From 9d9b558c4c591d3127e67ae54f683dcbb86a11a9 Mon Sep 17 00:00:00 2001 From: jack Date: Wed, 5 Aug 2026 01:00:46 +0800 Subject: [PATCH 1/2] fix(runner): preserve structured turn history --- internal/command/acp.go | 6 +- internal/command/interactive.go | 28 +-- internal/runner/cancel_backfill_test.go | 131 +++++++++++- internal/runner/history_result_test.go | 271 ++++++++++++++++++++++++ internal/runner/runner.go | 150 +++++++++---- internal/session/session.go | 13 +- internal/web/chat.go | 6 +- 7 files changed, 538 insertions(+), 67 deletions(-) create mode 100644 internal/runner/history_result_test.go diff --git a/internal/command/acp.go b/internal/command/acp.go index 62a7ddab..a5e77552 100644 --- a/internal/command/acp.go +++ b/internal/command/acp.go @@ -917,11 +917,11 @@ func (a *acpAgent) Prompt(ctx context.Context, params acp.PromptRequest) (acp.Pr // Reset per-turn approval-reviewer state (denial circuit breaker) at the // start of each user turn. sess.approvalState.OnTurnStart() - resp := runner.Run(promptCtx, sess.ag, history, sess.h, sess.rec, sess.todoStore, sess.env.GoalStore, sess.tracer, sess.tokenUsage) + result := runner.Run(promptCtx, sess.ag, history, sess.h, sess.rec, sess.todoStore, sess.env.GoalStore, sess.tracer, sess.tokenUsage) sess.mu.Lock() - if resp != "" { - sess.history = append(sess.history, &schema.Message{Role: schema.Assistant, Content: resp}) + if len(result.Messages) > 0 { + sess.history = append(sess.history, result.Messages...) } sess.cancel = nil sess.mu.Unlock() diff --git a/internal/command/interactive.go b/internal/command/interactive.go index cd5d0a5a..a7964af9 100644 --- a/internal/command/interactive.go +++ b/internal/command/interactive.go @@ -615,12 +615,12 @@ func (s *interactiveState) handlePrompt(userPrompt string) { s.history = append(s.history, schema.UserMessage(userPrompt)) s.history = agent.DrainBgNotifications(s.bgManager, s.history) s.approvalState.OnTurnStart() // reset the per-turn reviewer denial breaker - resp := runner.Run(runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) - if resp != "" { - s.history = append(s.history, &schema.Message{Role: schema.Assistant, Content: resp}) + result := runner.Run(runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) + if len(result.Messages) > 0 { + s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) - s.handlePlanCompletion(resp) + s.handlePlanCompletion(result.Response) } func (s *interactiveState) handlePlanCompletion(resp string) { @@ -657,12 +657,12 @@ func (s *interactiveState) handlePlanCompletion(resp string) { s.rec.RecordUser(revisePrompt) } s.history = append(s.history, schema.UserMessage(revisePrompt)) - newResp := runner.Run(s.runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) - if newResp != "" { - s.history = append(s.history, &schema.Message{Role: schema.Assistant, Content: newResp}) + result := runner.Run(s.runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) + if len(result.Messages) > 0 { + s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) - s.handlePlanCompletion(newResp) + s.handlePlanCompletion(result.Response) return } @@ -687,9 +687,9 @@ func (s *interactiveState) handlePlanCompletion(resp string) { s.rec.RecordUser(execPrompt) } s.history = append(s.history, schema.UserMessage(execPrompt)) - execResp := runner.Run(s.ctx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) - if execResp != "" { - s.history = append(s.history, &schema.Message{Role: schema.Assistant, Content: execResp}) + result := runner.Run(s.ctx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) + if len(result.Messages) > 0 { + s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) @@ -995,12 +995,12 @@ func (s *interactiveState) runEventLoop(initialHistory []adk.Message, initialRes s.runCtx = runCtx s.agentRunning.Store(true) s.history = append(s.history, schema.UserMessage(prompt)) - resp := runner.Run(runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) + result := runner.Run(runCtx, s.ag, s.history, s.h, s.rec, s.env.TodoStore, s.env.GoalStore, s.langfuseTracer, s.agentTokenUsage) runCancel() s.runCtx = nil s.agentRunning.Store(false) - if resp != "" { - s.history = append(s.history, &schema.Message{Role: schema.Assistant, Content: resp}) + if len(result.Messages) > 0 { + s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) } diff --git a/internal/runner/cancel_backfill_test.go b/internal/runner/cancel_backfill_test.go index 77882224..b3af2ae5 100644 --- a/internal/runner/cancel_backfill_test.go +++ b/internal/runner/cancel_backfill_test.go @@ -48,6 +48,43 @@ func (m *cancelModel) WithTools([]*schema.ToolInfo) (einomodel.ToolCallingChatMo return m, nil } +// partialStreamCancelModel emits a visible text prefix and an incomplete tool +// call, then returns the context error after the handler cancels the turn. +type partialStreamCancelModel struct{} + +func (m *partialStreamCancelModel) WithTools([]*schema.ToolInfo) (einomodel.ToolCallingChatModel, error) { + return m, nil +} + +func (m *partialStreamCancelModel) Generate(context.Context, []*schema.Message, ...einomodel.Option) (*schema.Message, error) { + return nil, errors.New("Generate is not used: streaming is enabled") +} + +func (m *partialStreamCancelModel) Stream(ctx context.Context, _ []*schema.Message, _ ...einomodel.Option) (*schema.StreamReader[*schema.Message], error) { + reader, writer := schema.Pipe[*schema.Message](1) + go func() { + defer writer.Close() + index := 0 + if writer.Send(&schema.Message{ + Role: schema.Assistant, + Content: "partial", + ToolCalls: []schema.ToolCall{{ + Index: &index, + ID: "call-partial-1", + Function: schema.FunctionCall{ + Name: "block", + Arguments: `{"note":"`, + }, + }}, + }, nil) { + return + } + <-ctx.Done() + writer.Send(nil, ctx.Err()) + }() + return reader, nil +} + func (m *cancelModel) Generate(context.Context, []*schema.Message, ...einomodel.Option) (*schema.Message, error) { return nil, errors.New("Generate is not used: streaming is enabled") } @@ -78,6 +115,23 @@ type cancelRecordingHandler struct { doneErr error } +type cancelOnTextHandler struct { + stubHandler + cancel context.CancelFunc + toolCalls int + toolResults int + doneErr error +} + +func (h *cancelOnTextHandler) OnAgentText(string) { h.cancel() } +func (h *cancelOnTextHandler) OnToolCall(handler.ToolCallEvent) { + h.toolCalls++ +} +func (h *cancelOnTextHandler) OnToolResult(handler.ToolResultEvent) { + h.toolResults++ +} +func (h *cancelOnTextHandler) OnAgentDone(err error) { h.doneErr = err } + func (h *cancelRecordingHandler) OnToolCall(handler.ToolCallEvent) { h.cancel() } func (h *cancelRecordingHandler) OnToolResult(ev handler.ToolResultEvent) { @@ -127,15 +181,20 @@ func TestRunInnerCancellationBackfillsToolResult(t *testing.T) { } h := &cancelRecordingHandler{cancel: cancel, done: make(chan struct{})} - finished := make(chan bool, 1) + type runOutcome struct { + result RunResult + done bool + } + finished := make(chan runOutcome, 1) go func() { - _, done := runInner(ctx, ag, []adk.Message{schema.UserMessage("go")}, h, rec) - finished <- done + result, done := runInner(ctx, ag, []adk.Message{schema.UserMessage("go")}, h, rec) + finished <- runOutcome{result: result, done: done} }() + var outcome runOutcome select { - case done := <-finished: - if !done { + case outcome = <-finished: + if !outcome.done { t.Errorf("runInner done = false, want true after cancellation") } case <-time.After(10 * time.Second): @@ -161,6 +220,15 @@ func TestRunInnerCancellationBackfillsToolResult(t *testing.T) { if h.results[0].ToolCallID != "call-block-1" { t.Errorf("result ToolCallID = %q, want call-block-1", h.results[0].ToolCallID) } + if len(outcome.result.Messages) != 2 { + t.Fatalf("live history messages = %d, want assistant tool call + backfill result", len(outcome.result.Messages)) + } + if outcome.result.Messages[0].Role != schema.Assistant || len(outcome.result.Messages[0].ToolCalls) != 1 { + t.Fatalf("live history first message = %#v, want assistant tool call", outcome.result.Messages[0]) + } + if outcome.result.Messages[1].Role != schema.Tool || outcome.result.Messages[1].ToolCallID != "call-block-1" { + t.Fatalf("live history second message = %#v, want matching backfill result", outcome.result.Messages[1]) + } // The session on disk pairs the recorded call with a recorded result. entries, err := session.LoadSession(rec.UUID()) @@ -199,3 +267,56 @@ func TestRunInnerCancellationBackfillsToolResult(t *testing.T) { } } } + +func TestRunInnerCancellationDropsIncompleteStreamingToolCall(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ag, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "partial-cancel-test", + Description: "partial-cancel-test", + Instruction: "test", + Model: &partialStreamCancelModel{}, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{newBlockingTool()}, + }}, + MaxIterations: 5, + }) + if err != nil { + t.Fatalf("create agent: %v", err) + } + + rec, err := session.NewRecorder("partial-cancel-test", "test", "test") + if err != nil { + t.Fatalf("create recorder: %v", err) + } + h := &cancelOnTextHandler{cancel: cancel} + result, done := runInner(ctx, ag, []adk.Message{schema.UserMessage("go")}, h, rec) + + if !done { + t.Fatal("runInner done = false, want true after stream cancellation") + } + if !errors.Is(h.doneErr, context.Canceled) { + t.Fatalf("OnAgentDone err = %v, want context.Canceled", h.doneErr) + } + if h.toolCalls != 0 || h.toolResults != 0 { + t.Fatalf("handler saw calls/results = %d/%d, want 0/0 for incomplete call", h.toolCalls, h.toolResults) + } + if result.Response != "partial" || len(result.Messages) != 1 { + t.Fatalf("result = %#v, want one partial text message", result) + } + if result.Messages[0].Role != schema.Assistant || result.Messages[0].Content != "partial" || len(result.Messages[0].ToolCalls) != 0 { + t.Fatalf("result message = %#v, want text-only assistant", result.Messages[0]) + } + + entries, err := session.LoadSession(rec.UUID()) + if err != nil { + t.Fatalf("load session: %v", err) + } + for _, entry := range entries { + if entry.Type == session.EntryToolCall || entry.Type == session.EntryToolResult { + t.Fatalf("session persisted incomplete tool entry: %#v", entry) + } + } +} diff --git a/internal/runner/history_result_test.go b/internal/runner/history_result_test.go new file mode 100644 index 00000000..af6c7fcf --- /dev/null +++ b/internal/runner/history_result_test.go @@ -0,0 +1,271 @@ +package runner + +import ( + "context" + "errors" + "reflect" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + einomodel "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + + "github.com/cnjack/jcode/internal/session" +) + +type historyModel struct{} + +func (m *historyModel) WithTools([]*schema.ToolInfo) (einomodel.ToolCallingChatModel, error) { + return m, nil +} + +func (m *historyModel) Generate(context.Context, []*schema.Message, ...einomodel.Option) (*schema.Message, error) { + return nil, errors.New("Generate is not used: streaming is enabled") +} + +func (m *historyModel) Stream(_ context.Context, input []*schema.Message, _ ...einomodel.Option) (*schema.StreamReader[*schema.Message], error) { + if last := input[len(input)-1]; last.Role == schema.Tool { + return schema.StreamReaderFromArray([]*schema.Message{{ + Role: schema.Assistant, + Content: "installed", + }}), nil + } + return schema.StreamReaderFromArray([]*schema.Message{{ + Role: schema.Assistant, + Content: "checking", + ToolCalls: []schema.ToolCall{{ + ID: "call-check-1", + Function: schema.FunctionCall{Name: "check", Arguments: `{}`}, + }}, + }}), nil +} + +type historyTool struct{} + +func (historyTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "check", + Desc: "returns an installation result", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}), + }, nil +} + +func (historyTool) InvokableRun(context.Context, string, ...tool.Option) (string, error) { + return "uv 0.12.1", nil +} + +type parallelHistoryModel struct{} + +func (m *parallelHistoryModel) WithTools([]*schema.ToolInfo) (einomodel.ToolCallingChatModel, error) { + return m, nil +} + +func (m *parallelHistoryModel) Generate(context.Context, []*schema.Message, ...einomodel.Option) (*schema.Message, error) { + return nil, errors.New("Generate is not used: streaming is enabled") +} + +func (m *parallelHistoryModel) Stream(_ context.Context, input []*schema.Message, _ ...einomodel.Option) (*schema.StreamReader[*schema.Message], error) { + if last := input[len(input)-1]; last.Role == schema.Tool { + return schema.StreamReaderFromArray([]*schema.Message{{Role: schema.Assistant, Content: "done"}}), nil + } + first, second := 0, 1 + return schema.StreamReaderFromArray([]*schema.Message{{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {Index: &first, ID: "call-a", Function: schema.FunctionCall{Name: "check_a", Arguments: `{}`}}, + {Index: &second, ID: "call-b", Function: schema.FunctionCall{Name: "check_b", Arguments: `{}`}}, + }, + }}), nil +} + +type namedHistoryTool struct { + name string + output string +} + +func (t namedHistoryTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "returns a fixed result", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}), + }, nil +} + +func (t namedHistoryTool) InvokableRun(context.Context, string, ...tool.Option) (string, error) { + return t.output, nil +} + +func TestRunReturnsStructuredMessagesMatchingSessionReplay(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx := context.Background() + ag := newHistoryAgent(ctx, t) + + rec, err := session.NewRecorder("history-test", "test", "test") + if err != nil { + t.Fatalf("create recorder: %v", err) + } + rec.RecordUser("install uv") + + result := Run( + ctx, + ag, + []adk.Message{schema.UserMessage("install uv")}, + stubHandler{}, + rec, + nil, + nil, + nil, + nil, + ) + + if result.Response != "checkinginstalled" { + t.Fatalf("Response = %q, want %q", result.Response, "checkinginstalled") + } + assertStructuredTurn(t, result.Messages) + + entries, err := session.LoadSession(rec.UUID()) + if err != nil { + t.Fatalf("load session: %v", err) + } + replayed := session.ReconstructState(entries).History + if len(replayed) != len(result.Messages)+1 { + t.Fatalf("replayed messages = %d, want user + %d turn messages", len(replayed), len(result.Messages)) + } + assertMessagesEqual(t, replayed[1:], result.Messages) +} + +func TestRunStructuredMessagesDoNotDependOnRecorder(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx := context.Background() + result := Run( + ctx, + newHistoryAgent(ctx, t), + []adk.Message{schema.UserMessage("install uv")}, + stubHandler{}, + nil, + nil, + nil, + nil, + nil, + ) + + if result.Response != "checkinginstalled" { + t.Fatalf("Response = %q, want %q", result.Response, "checkinginstalled") + } + assertStructuredTurn(t, result.Messages) +} + +func TestRunPreservesParallelToolCallBatchInLiveAndReplayHistory(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx := context.Background() + ag, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "parallel-history-test", + Description: "parallel-history-test", + Instruction: "test", + Model: ¶llelHistoryModel{}, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{ + namedHistoryTool{name: "check_a", output: "a"}, + namedHistoryTool{name: "check_b", output: strings.Repeat("line\n", 3000)}, + }, + }}, + MaxIterations: 5, + }) + if err != nil { + t.Fatalf("create agent: %v", err) + } + rec, err := session.NewRecorder("parallel-history-test", "test", "test") + if err != nil { + t.Fatalf("create recorder: %v", err) + } + rec.RecordUser("check both") + + result := Run(ctx, ag, []adk.Message{schema.UserMessage("check both")}, stubHandler{}, rec, nil, nil, nil, nil) + if len(result.Messages) != 4 { + t.Fatalf("turn messages = %d, want assistant + 2 tools + assistant", len(result.Messages)) + } + if result.Messages[0].Role != schema.Assistant || len(result.Messages[0].ToolCalls) != 2 { + t.Fatalf("first message = %#v, want assistant with 2 tool calls", result.Messages[0]) + } + wantIDs := map[string]bool{"call-a": true, "call-b": true} + for _, message := range result.Messages[1:3] { + if message.Role != schema.Tool || !wantIDs[message.ToolCallID] { + t.Fatalf("parallel tool message = %#v, want call-a or call-b result", message) + } + delete(wantIDs, message.ToolCallID) + } + if len(wantIDs) != 0 { + t.Fatalf("missing tool results for %v", wantIDs) + } + var sawTruncated bool + for _, message := range result.Messages[1:3] { + if strings.Contains(message.Content, "truncated") { + sawTruncated = true + } + } + if !sawTruncated { + t.Fatal("large tool result was not normalized to the persisted truncated form") + } + if result.Messages[3].Role != schema.Assistant || result.Messages[3].Content != "done" { + t.Fatalf("last message = %#v, want final assistant", result.Messages[3]) + } + + entries, err := session.LoadSession(rec.UUID()) + if err != nil { + t.Fatalf("load session: %v", err) + } + replayed := session.ReconstructState(entries).History + assertMessagesEqual(t, replayed[1:], result.Messages) +} + +func newHistoryAgent(ctx context.Context, t *testing.T) *adk.ChatModelAgent { + t.Helper() + ag, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "history-test", + Description: "history-test", + Instruction: "test", + Model: &historyModel{}, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{historyTool{}}, + }}, + MaxIterations: 5, + }) + if err != nil { + t.Fatalf("create agent: %v", err) + } + return ag +} + +func assertStructuredTurn(t *testing.T, messages []adk.Message) { + t.Helper() + if len(messages) != 3 { + t.Fatalf("turn messages = %d, want assistant + tool + assistant", len(messages)) + } + if messages[0].Role != schema.Assistant || messages[0].Content != "checking" || len(messages[0].ToolCalls) != 1 { + t.Fatalf("first message = %#v, want assistant text with tool call", messages[0]) + } + if messages[0].ToolCalls[0].ID != "call-check-1" { + t.Errorf("tool call ID = %q, want call-check-1", messages[0].ToolCalls[0].ID) + } + if messages[1].Role != schema.Tool || messages[1].ToolCallID != "call-check-1" || messages[1].Content != "uv 0.12.1" { + t.Fatalf("tool message = %#v, want matching tool result", messages[1]) + } + if messages[2].Role != schema.Assistant || messages[2].Content != "installed" { + t.Fatalf("last message = %#v, want final assistant response", messages[2]) + } +} + +func assertMessagesEqual(t *testing.T, got, want []adk.Message) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("message lengths differ: got %d want %d", len(got), len(want)) + } + for i := range want { + if !reflect.DeepEqual(got[i], want[i]) { + t.Errorf("message[%d] mismatch:\n got %#v\n want %#v", i, got[i], want[i]) + } + } +} diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 44e2267a..86846139 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -23,6 +23,17 @@ import ( "github.com/cnjack/jcode/internal/usage" ) +// RunResult is the agent-generated, persistable output of one user turn. +// Response preserves the flattened assistant text used by UI-only consumers, +// while Messages carries the structured assistant/tool transcript that must be +// appended to live conversation history for the next turn. Internal +// continuation prompts are deliberately excluded from Messages, matching the +// session transcript reconstructed after a restart. +type RunResult struct { + Response string + Messages []adk.Message +} + // Run executes the agent for a single turn, wrapping the response with a // Langfuse trace when a tracer is present, enforcing todo-completion guards, // and sending token-usage updates to the handler when done. @@ -36,11 +47,11 @@ func Run( goalStore *tools.GoalStore, tracer *telemetry.LangfuseTracer, tokenUsage *internalmodel.TokenUsage, -) (response string) { +) (result RunResult) { if tracer != nil { ctx = tracer.WithNewTrace(ctx, "coding_agent", messages) defer func() { - tracer.EndTrace(ctx, response) + tracer.EndTrace(ctx, result.Response) }() } if rec != nil { @@ -102,15 +113,17 @@ func Run( } h.OnAgentStart() - resp, done := runInner(ctx, ag, messages, h, rec) + // Continuations need the exact structured output of the previous lap, but + // must not mutate the caller's backing slice while assembling that context. + runMessages := append([]adk.Message(nil), messages...) + lap, done := runInner(ctx, ag, runMessages, h, rec) + result.Response = lap.Response + result.Messages = append(result.Messages, lap.Messages...) if done { // runInner already signaled completion (cancellation or a real error). - return resp + return result } - // pending is the assistant text produced since messages was last extended; - // each continuation appends only this delta, never the accumulated resp, - // so earlier turns are not duplicated into the context. - pending := resp + runMessages = append(runMessages, lap.Messages...) // Unified continuation pipeline. Three mechanisms can keep the agent going // after it stops calling tools — incomplete-todo guard, active-goal guard, and @@ -181,13 +194,16 @@ continuationLoop: } h.OnAgentText(banner) - messages = append(messages, &schema.Message{Role: schema.Assistant, Content: pending}) - messages = append(messages, schema.UserMessage(reason)) - extra, done := runInner(ctx, ag, messages, h, rec) - resp += extra - pending = extra + // Continuation prompts are internal control messages. They are included in + // the next lap's context, but not returned as user-visible turn history or + // persisted as user messages in the session transcript. + runMessages = append(runMessages, schema.UserMessage(reason)) + extra, done := runInner(ctx, ag, runMessages, h, rec) + result.Response += extra.Response + result.Messages = append(result.Messages, extra.Messages...) + runMessages = append(runMessages, extra.Messages...) if done { - return resp + return result } } @@ -202,7 +218,7 @@ continuationLoop: // This turn's token delta is persisted by the deferred recordUsageTurn, which // also covers the early-return (cancel/error) paths above. h.OnAgentDone(nil) - return resp + return result } // continuationSource picks which mechanism drives the next continuation lap, @@ -259,13 +275,18 @@ func runInner( messages []adk.Message, h handler.AgentEventHandler, rec *session.Recorder, -) (string, bool) { +) (RunResult, bool) { input := &adk.AgentInput{ Messages: messages, EnableStreaming: true, } - var assistantText strings.Builder + var result RunResult + var responseText strings.Builder + finish := func(done bool) (RunResult, bool) { + result.Response = responseText.String() + return result, done + } // toolStarts records when each tool call was announced so results can // carry a call→result latency. The event loop below is the only reader @@ -288,9 +309,15 @@ func runInner( } applyApprovalOutcome(&ev, meter) h.OnToolResult(ev) + historyOutput := output if rec != nil { - rec.RecordToolResult(name, output, toolCallID, err, ev.Denied, ev.Duration) + historyOutput = rec.RecordToolResult(name, output, toolCallID, err, ev.Denied, ev.Duration) } + result.Messages = append(result.Messages, schema.ToolMessage( + historyOutput, + toolCallID, + schema.WithToolName(name), + )) } // drainDanglingToolResults backfills a result for every announced tool // call that never produced one (user stop, fatal tool-node failure). The @@ -321,7 +348,7 @@ func runInner( config.Logger().Printf("[runner] context cancelled, stopping iteration") drainDanglingToolResults() h.OnAgentDone(ctx.Err()) - return assistantText.String(), true + return finish(true) default: } @@ -342,7 +369,7 @@ func runInner( config.Logger().Printf("[runner] event error during cancellation: %v", event.Err) drainDanglingToolResults() h.OnAgentDone(ctx.Err()) - return assistantText.String(), true + return finish(true) } // Log the provider's raw payload, hand the frontends a sentence a // human can act on. This is the single choke point for model errors, @@ -351,7 +378,7 @@ func runInner( config.Logger().Printf("[runner] event error: %v", event.Err) drainDanglingToolResults() h.OnAgentDone(internalmodel.WrapFriendly(event.Err, "", "")) - return assistantText.String(), true + return finish(true) } if event.Output == nil || event.Output.MessageOutput == nil { config.Logger().Printf("[runner] event #%d: nil output", eventCount) @@ -405,10 +432,13 @@ func runInner( } if mo.IsStreaming { + var messageText strings.Builder + var streamErr error // Accumulate streaming tool call names, args, and IDs across chunks. type pendingTC struct { name string id string + typ string args strings.Builder } pending := make(map[int]*pendingTC) @@ -418,6 +448,7 @@ func runInner( break } if err != nil { + streamErr = err break } if chunk == nil { @@ -429,7 +460,7 @@ func runInner( idx = *tc.Index } if tc.Function.Name != "" { - p := &pendingTC{name: tc.Function.Name, id: tc.ID} + p := &pendingTC{name: tc.Function.Name, id: tc.ID, typ: tc.Type} p.args.WriteString(tc.Function.Arguments) pending[idx] = p } else if p, ok := pending[idx]; ok { @@ -437,17 +468,39 @@ func runInner( } } if chunk.Content != "" { - assistantText.WriteString(chunk.Content) + messageText.WriteString(chunk.Content) + responseText.WriteString(chunk.Content) h.OnAgentText(chunk.Content) } } + if streamErr != nil { + // A failed stream may contain only a prefix of tool-call arguments. + // Preserve text already shown to the user, but never persist or + // expose incomplete calls as executable conversation history. + if messageText.Len() > 0 { + if rec != nil { + rec.RecordAssistant(messageText.String()) + } + result.Messages = append(result.Messages, &schema.Message{ + Role: schema.Assistant, + Content: messageText.String(), + }) + } + if ctx.Err() != nil { + config.Logger().Printf("[runner] assistant stream cancelled: %v", streamErr) + h.OnAgentDone(ctx.Err()) + return finish(true) + } + config.Logger().Printf("[runner] assistant stream error: %v", streamErr) + h.OnAgentDone(internalmodel.WrapFriendly(streamErr, "", "")) + return finish(true) + } // Flush assistant text at the end of each assistant message so the // session file preserves the true message/tool interleaving. Without // this, the whole run accumulates into a single assistant entry and // replay collapses all surrounding tool calls into one big group. - if rec != nil && assistantText.Len() > 0 { - rec.RecordAssistant(assistantText.String()) - assistantText.Reset() + if rec != nil && messageText.Len() > 0 { + rec.RecordAssistant(messageText.String()) } // Notify and record accumulated tool calls in index order. // All tool calls from this assistant message form one batch. @@ -456,6 +509,25 @@ func runInner( indices = append(indices, idx) } sort.Ints(indices) + var toolCalls []schema.ToolCall + if len(indices) > 0 { + toolCalls = make([]schema.ToolCall, 0, len(indices)) + } + for _, idx := range indices { + p := pending[idx] + toolCalls = append(toolCalls, schema.ToolCall{ + ID: p.id, + Type: p.typ, + Function: schema.FunctionCall{Name: p.name, Arguments: p.args.String()}, + }) + } + if messageText.Len() > 0 || len(toolCalls) > 0 { + result.Messages = append(result.Messages, &schema.Message{ + Role: schema.Assistant, + Content: messageText.String(), + ToolCalls: toolCalls, + }) + } if len(indices) > 0 { batchID := nextBatchID() startedAt := time.Now() @@ -477,6 +549,17 @@ func runInner( } } } else if mo.Message != nil { + if mo.Message.Content != "" || len(mo.Message.ToolCalls) > 0 { + var toolCalls []schema.ToolCall + if len(mo.Message.ToolCalls) > 0 { + toolCalls = append(toolCalls, mo.Message.ToolCalls...) + } + result.Messages = append(result.Messages, &schema.Message{ + Role: schema.Assistant, + Content: mo.Message.Content, + ToolCalls: toolCalls, + }) + } if len(mo.Message.ToolCalls) > 0 { batchID := nextBatchID() startedAt := time.Now() @@ -498,26 +581,19 @@ func runInner( } } if mo.Message.Content != "" { - assistantText.WriteString(mo.Message.Content) + responseText.WriteString(mo.Message.Content) h.OnAgentText(mo.Message.Content) } // Flush non-streaming assistant text immediately so each assistant // message is a distinct session entry with its surrounding tool calls. - if rec != nil && assistantText.Len() > 0 { - rec.RecordAssistant(assistantText.String()) - assistantText.Reset() + if rec != nil && mo.Message.Content != "" { + rec.RecordAssistant(mo.Message.Content) } } } - // Final safety flush for any remaining text (e.g. returns above that skip - // the per-message flush, or trailing content after the last tool batch). - if rec != nil && assistantText.Len() > 0 { - rec.RecordAssistant(assistantText.String()) - } - // Clean completion: Run emits the single final OnAgentDone(nil). - return assistantText.String(), false + return finish(false) } // buildTokenUsage snapshots a tracker into a handler.TokenUsage. TotalTokens is diff --git a/internal/session/session.go b/internal/session/session.go index cf4c30f1..1a7fa6c8 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -577,11 +577,13 @@ func (r *Recorder) RecordToolCall(name, args, toolCallID, batchID string, batchI }) } -// RecordToolResult appends a tool-result entry. denied marks a user-rejected -// approval; duration is the approval-wait-adjusted execution latency (0 when -// unknown). Large outputs are automatically truncated (head+tail preserved) -// and the full content is saved to an overflow file on disk. -func (r *Recorder) RecordToolResult(name, output, toolCallID string, err error, denied bool, duration time.Duration) { +// RecordToolResult appends a tool-result entry and returns the exact output +// stored in the transcript. denied marks a user-rejected approval; duration is +// the approval-wait-adjusted execution latency (0 when unknown). Large outputs +// are automatically truncated (head+tail preserved) and the full content is +// saved to an overflow file on disk. Callers should use the returned output in +// live model history so live and replayed sessions have the same context. +func (r *Recorder) RecordToolResult(name, output, toolCallID string, err error, denied bool, duration time.Duration) string { errStr := "" if err != nil { errStr = err.Error() @@ -591,6 +593,7 @@ func (r *Recorder) RecordToolResult(name, output, toolCallID string, err error, Type: EntryToolResult, Name: name, Output: output, ToolCallID: toolCallID, Error: errStr, Denied: denied, DurationMs: duration.Milliseconds(), }) + return output } // RecordToolObservation appends metadata-only progressive-disclosure evidence. diff --git a/internal/web/chat.go b/internal/web/chat.go index 8695eea6..801a047f 100644 --- a/internal/web/chat.go +++ b/internal/web/chat.go @@ -312,10 +312,10 @@ func (s *Server) submitMessage(eng *Engine, message, mode, source, sessionID str // Inject the hook dispatcher so PreToolUse/PostToolUse/Stop hooks run on the // Web surface too (parity with the TUI); reloaded per turn for hot-apply. hookCtx := hooks.WithDispatcher(runCtx, hooks.NewSessionDispatcher(config.ConfigDir(), eng.env.Pwd(), recorder.UUID(), config.Logger().Printf)) - resp := runner.Run(hookCtx, agent, history, eng.eventHandler, recorder, eng.todoStore, eng.env.GoalStore, s.tracer, eng.tokenUsage) - if resp != "" { + result := runner.Run(hookCtx, agent, history, eng.eventHandler, recorder, eng.todoStore, eng.env.GoalStore, s.tracer, eng.tokenUsage) + if len(result.Messages) > 0 { eng.emu.Lock() - eng.history = append(eng.history, &schema.Message{Role: schema.Assistant, Content: resp}) + eng.history = append(eng.history, result.Messages...) eng.emu.Unlock() } }() From b576d4114be079a740716c7428187701614a9574 Mon Sep 17 00:00:00 2001 From: jack Date: Wed, 5 Aug 2026 01:57:51 +0800 Subject: [PATCH 2/2] fix(runner): harden interrupted turn history --- internal/command/interactive.go | 9 +- internal/command/interactive_plan_test.go | 28 +++ internal/runner/cancel_backfill_test.go | 27 ++- internal/runner/history_result_test.go | 8 +- internal/runner/review_feedback_test.go | 135 ++++++++++++++ internal/runner/runner.go | 208 ++++++++++++++++++---- internal/session/session.go | 31 +++- 7 files changed, 403 insertions(+), 43 deletions(-) create mode 100644 internal/command/interactive_plan_test.go create mode 100644 internal/runner/review_feedback_test.go diff --git a/internal/command/interactive.go b/internal/command/interactive.go index a7964af9..66953261 100644 --- a/internal/command/interactive.go +++ b/internal/command/interactive.go @@ -620,13 +620,14 @@ func (s *interactiveState) handlePrompt(userPrompt string) { s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) - s.handlePlanCompletion(result.Response) + s.handlePlanCompletion(result) } -func (s *interactiveState) handlePlanCompletion(resp string) { - if s.agentMode != tui.ModePlanning || resp == "" { +func (s *interactiveState) handlePlanCompletion(planResult runner.RunResult) { + if s.agentMode != tui.ModePlanning || planResult.Response == "" || planResult.Err != nil { return } + resp := planResult.Response s.planStore.Submit("Plan", resp) config.Logger().Printf("[plan] plan submitted for review (%d chars)", len(resp)) @@ -662,7 +663,7 @@ func (s *interactiveState) handlePlanCompletion(resp string) { s.history = append(s.history, result.Messages...) } s.history = agent.SyncSummarization(s.summCapture, s.history, s.rec) - s.handlePlanCompletion(result.Response) + s.handlePlanCompletion(result) return } diff --git a/internal/command/interactive_plan_test.go b/internal/command/interactive_plan_test.go new file mode 100644 index 00000000..8120c395 --- /dev/null +++ b/internal/command/interactive_plan_test.go @@ -0,0 +1,28 @@ +package command + +import ( + "context" + "testing" + "time" + + "github.com/cnjack/jcode/internal/runner" + "github.com/cnjack/jcode/internal/tui" +) + +func TestHandlePlanCompletionSkipsInterruptedRun(t *testing.T) { + state := &interactiveState{agentMode: tui.ModePlanning} + returned := make(chan struct{}) + go func() { + state.handlePlanCompletion(runner.RunResult{ + Response: "partial plan", + Err: context.Canceled, + }) + close(returned) + }() + + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("interrupted plan entered the approval flow") + } +} diff --git a/internal/runner/cancel_backfill_test.go b/internal/runner/cancel_backfill_test.go index b3af2ae5..fe9fb896 100644 --- a/internal/runner/cancel_backfill_test.go +++ b/internal/runner/cancel_backfill_test.go @@ -211,6 +211,9 @@ func TestRunInnerCancellationBackfillsToolResult(t *testing.T) { if !errors.Is(h.doneErr, context.Canceled) { t.Errorf("OnAgentDone err = %v, want context.Canceled", h.doneErr) } + if !errors.Is(outcome.result.Err, context.Canceled) { + t.Errorf("RunResult.Err = %v, want context.Canceled", outcome.result.Err) + } // The announced call got exactly one result (drain backfill or a folded // framework result — either satisfies the invariant). @@ -252,6 +255,7 @@ func TestRunInnerCancellationBackfillsToolResult(t *testing.T) { } state := session.ReconstructState(entries) + assertMessagesEqual(t, state.History, outcome.result.Messages) for i, m := range state.History { if m.Role != schema.Assistant || len(m.ToolCalls) == 0 { continue @@ -292,7 +296,24 @@ func TestRunInnerCancellationDropsIncompleteStreamingToolCall(t *testing.T) { t.Fatalf("create recorder: %v", err) } h := &cancelOnTextHandler{cancel: cancel} - result, done := runInner(ctx, ag, []adk.Message{schema.UserMessage("go")}, h, rec) + type runOutcome struct { + result RunResult + done bool + } + finished := make(chan runOutcome, 1) + go func() { + result, done := runInner(ctx, ag, []adk.Message{schema.UserMessage("go")}, h, rec) + finished <- runOutcome{result: result, done: done} + }() + + var result RunResult + var done bool + select { + case outcome := <-finished: + result, done = outcome.result, outcome.done + case <-time.After(10 * time.Second): + t.Fatal("runInner did not return after stream cancellation") + } if !done { t.Fatal("runInner done = false, want true after stream cancellation") @@ -300,6 +321,9 @@ func TestRunInnerCancellationDropsIncompleteStreamingToolCall(t *testing.T) { if !errors.Is(h.doneErr, context.Canceled) { t.Fatalf("OnAgentDone err = %v, want context.Canceled", h.doneErr) } + if !errors.Is(result.Err, context.Canceled) { + t.Fatalf("RunResult.Err = %v, want context.Canceled", result.Err) + } if h.toolCalls != 0 || h.toolResults != 0 { t.Fatalf("handler saw calls/results = %d/%d, want 0/0 for incomplete call", h.toolCalls, h.toolResults) } @@ -319,4 +343,5 @@ func TestRunInnerCancellationDropsIncompleteStreamingToolCall(t *testing.T) { t.Fatalf("session persisted incomplete tool entry: %#v", entry) } } + assertMessagesEqual(t, session.ReconstructState(entries).History, result.Messages) } diff --git a/internal/runner/history_result_test.go b/internal/runner/history_result_test.go index af6c7fcf..7f56cfa6 100644 --- a/internal/runner/history_result_test.go +++ b/internal/runner/history_result_test.go @@ -190,7 +190,13 @@ func TestRunPreservesParallelToolCallBatchInLiveAndReplayHistory(t *testing.T) { if result.Messages[0].Role != schema.Assistant || len(result.Messages[0].ToolCalls) != 2 { t.Fatalf("first message = %#v, want assistant with 2 tool calls", result.Messages[0]) } - wantIDs := map[string]bool{"call-a": true, "call-b": true} + wantIDs := make(map[string]bool, len(result.Messages[0].ToolCalls)) + for _, toolCall := range result.Messages[0].ToolCalls { + if toolCall.ID == "" || wantIDs[toolCall.ID] { + t.Fatalf("assistant tool calls = %#v, want two unique IDs", result.Messages[0].ToolCalls) + } + wantIDs[toolCall.ID] = true + } for _, message := range result.Messages[1:3] { if message.Role != schema.Tool || !wantIDs[message.ToolCallID] { t.Fatalf("parallel tool message = %#v, want call-a or call-b result", message) diff --git a/internal/runner/review_feedback_test.go b/internal/runner/review_feedback_test.go new file mode 100644 index 00000000..7de16d39 --- /dev/null +++ b/internal/runner/review_feedback_test.go @@ -0,0 +1,135 @@ +package runner + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + + "github.com/cnjack/jcode/internal/handler" + "github.com/cnjack/jcode/internal/session" +) + +var errFirstToolChunk = errors.New("tool stream failed before first chunk") + +type firstRecvErrorTool struct{} + +func (firstRecvErrorTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "check", + Desc: "fails before emitting its first result chunk", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}), + }, nil +} + +func (firstRecvErrorTool) StreamableRun(context.Context, string, ...tool.Option) (*schema.StreamReader[string], error) { + reader, writer := schema.Pipe[string](1) + writer.Send("", errFirstToolChunk) + writer.Close() + return reader, nil +} + +type resultCaptureHandler struct { + stubHandler + results []handler.ToolResultEvent + doneErr error +} + +func (h *resultCaptureHandler) OnToolResult(event handler.ToolResultEvent) { + h.results = append(h.results, event) +} + +func (h *resultCaptureHandler) OnAgentDone(err error) { h.doneErr = err } + +type closeRecorderOnResultHandler struct { + resultCaptureHandler + recorder *session.Recorder +} + +func (h *closeRecorderOnResultHandler) OnToolResult(event handler.ToolResultEvent) { + h.resultCaptureHandler.OnToolResult(event) + h.recorder.Close() +} + +func TestRunPairsFirstReceiveToolStreamFailureWithAnnouncedCall(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx := context.Background() + ag, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "stream-error-test", + Description: "stream-error-test", + Instruction: "test", + Model: &historyModel{}, + ToolsConfig: adk.ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{firstRecvErrorTool{}}, + }}, + MaxIterations: 5, + }) + if err != nil { + t.Fatalf("create agent: %v", err) + } + rec, err := session.NewRecorder("stream-error-test", "test", "test") + if err != nil { + t.Fatalf("create recorder: %v", err) + } + rec.RecordUser("check") + h := &resultCaptureHandler{} + + result := Run(ctx, ag, []adk.Message{schema.UserMessage("check")}, h, rec, nil, nil, nil, nil) + if result.Err == nil { + t.Fatal("RunResult.Err = nil, want tool stream failure") + } + if len(h.results) != 1 || h.results[0].ToolCallID != "call-check-1" { + t.Fatalf("tool results = %#v, want one result for call-check-1", h.results) + } + if len(result.Messages) != 2 || result.Messages[1].Role != schema.Tool || result.Messages[1].ToolCallID != "call-check-1" { + t.Fatalf("result messages = %#v, want paired assistant call and tool failure", result.Messages) + } + + entries, err := session.LoadSession(rec.UUID()) + if err != nil { + t.Fatalf("load session: %v", err) + } + assertMessagesEqual(t, session.ReconstructState(entries).History[1:], result.Messages) +} + +func TestRunSurfacesToolResultPersistenceFailureWithoutUsingUnstoredOutput(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + ctx := context.Background() + rec, err := session.NewRecorder("persist-error-test", "test", "test") + if err != nil { + t.Fatalf("create recorder: %v", err) + } + rec.RecordUser("check") + h := &closeRecorderOnResultHandler{recorder: rec} + + result := Run(ctx, newHistoryAgent(ctx, t), []adk.Message{schema.UserMessage("check")}, h, rec, nil, nil, nil, nil) + if result.Err == nil { + t.Fatal("RunResult.Err = nil, want recorder persistence failure") + } + if h.doneErr == nil { + t.Fatal("OnAgentDone err = nil, want recorder persistence failure") + } + for _, want := range []string{rec.UUID() + ".json", "call-check-1"} { + if !strings.Contains(result.Err.Error(), want) { + t.Errorf("RunResult.Err = %q, want it to contain %q", result.Err, want) + } + } + if len(result.Messages) != 2 { + t.Fatalf("result messages = %d, want assistant call + interrupted result", len(result.Messages)) + } + toolMessage := result.Messages[1] + if toolMessage.Role != schema.Tool || toolMessage.ToolCallID != "call-check-1" || toolMessage.Content != session.InterruptedToolOutput { + t.Fatalf("tool message = %#v, want replay-equivalent interrupted result", toolMessage) + } + + entries, err := session.LoadSession(rec.UUID()) + if err != nil { + t.Fatalf("load session: %v", err) + } + assertMessagesEqual(t, session.ReconstructState(entries).History[1:], result.Messages) +} diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 86846139..2f7b3958 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -2,6 +2,7 @@ package runner import ( "context" + "errors" "fmt" "io" "sort" @@ -26,12 +27,15 @@ import ( // RunResult is the agent-generated, persistable output of one user turn. // Response preserves the flattened assistant text used by UI-only consumers, // while Messages carries the structured assistant/tool transcript that must be -// appended to live conversation history for the next turn. Internal -// continuation prompts are deliberately excluded from Messages, matching the -// session transcript reconstructed after a restart. +// appended to live conversation history for the next turn. Err reports an +// interrupted or failed turn so callers do not treat a partial response as a +// completed artifact such as an approval-ready plan. Internal continuation +// prompts are deliberately excluded from Messages, matching the session +// transcript reconstructed after a restart. type RunResult struct { Response string Messages []adk.Message + Err error } // Run executes the agent for a single turn, wrapping the response with a @@ -119,6 +123,7 @@ func Run( lap, done := runInner(ctx, ag, runMessages, h, rec) result.Response = lap.Response result.Messages = append(result.Messages, lap.Messages...) + result.Err = lap.Err if done { // runInner already signaled completion (cancellation or a real error). return result @@ -148,7 +153,9 @@ continuationLoop: select { case <-ctx.Done(): config.Logger().Printf("[runner] continuation cancelled") - break continuationLoop + result.Err = ctx.Err() + h.OnAgentDone(result.Err) + return result default: } if goalStore != nil && tokenUsage != nil { @@ -201,6 +208,7 @@ continuationLoop: extra, done := runInner(ctx, ag, runMessages, h, rec) result.Response += extra.Response result.Messages = append(result.Messages, extra.Messages...) + result.Err = extra.Err runMessages = append(runMessages, extra.Messages...) if done { return result @@ -283,8 +291,9 @@ func runInner( var result RunResult var responseText strings.Builder - finish := func(done bool) (RunResult, bool) { + finish := func(done bool, runErr error) (RunResult, bool) { result.Response = responseText.String() + result.Err = runErr return result, done } @@ -292,16 +301,18 @@ func runInner( // carry a call→result latency. The event loop below is the only reader // and writer (single goroutine), so no locking is needed. type toolStart struct { - at time.Time - name string + at time.Time + name string + order int } toolStarts := make(map[string]toolStart) + nextToolOrder := 0 // meter carries approval wait/denied outcomes from the approval path (see // Run) so the emitted Duration is pure execution time and denied calls are // flagged. emitToolResult also owns session recording so the persisted // entry matches the emitted event exactly (denied + adjusted duration). meter := approvalMeterFrom(ctx) - emitToolResult := func(name, output, toolCallID string, err error) { + emitToolResult := func(name, output, toolCallID string, err error) error { ev := handler.ToolResultEvent{Name: name, Output: output, ToolCallID: toolCallID, Err: err} if started, ok := toolStarts[toolCallID]; ok { ev.Duration = time.Since(started.at) @@ -311,13 +322,68 @@ func runInner( h.OnToolResult(ev) historyOutput := output if rec != nil { - historyOutput = rec.RecordToolResult(name, output, toolCallID, err, ev.Denied, ev.Duration) + var persistErr error + historyOutput, persistErr = rec.RecordToolResult(name, output, toolCallID, err, ev.Denied, ev.Duration) + if persistErr != nil { + // ReconstructState repairs an on-disk call without a result using + // this same placeholder. Keep live history valid and replay-equivalent, + // but stop the turn instead of claiming the real output was durable. + result.Messages = append(result.Messages, schema.ToolMessage( + session.InterruptedToolOutput, + toolCallID, + schema.WithToolName(name), + )) + return persistErr + } } result.Messages = append(result.Messages, schema.ToolMessage( historyOutput, toolCallID, schema.WithToolName(name), )) + return nil + } + + // A streaming tool event carries the call ID only in its chunks. If the + // first receive fails, defer pairing the failure until the remaining + // announced calls identify the only possible ID. This also handles parallel + // same-name calls without assigning a failure to whichever map entry happens + // to be visited first. + pendingToolFailures := make(map[string][]error) + matchingToolCallIDs := func(name string) []string { + ids := make([]string, 0) + for id, started := range toolStarts { + if started.name == name { + ids = append(ids, id) + } + } + sort.Slice(ids, func(i, j int) bool { + return toolStarts[ids[i]].order < toolStarts[ids[j]].order + }) + return ids + } + resolvePendingToolFailures := func(name string, force bool) error { + failures := pendingToolFailures[name] + if len(failures) == 0 { + return nil + } + ids := matchingToolCallIDs(name) + if !force && len(ids) > len(failures) { + return nil + } + count := min(len(ids), len(failures)) + var firstErr error + for i := 0; i < count; i++ { + if persistErr := emitToolResult(name, "", ids[i], failures[i]); persistErr != nil && firstErr == nil { + firstErr = persistErr + } + } + if count == len(failures) { + delete(pendingToolFailures, name) + } else { + pendingToolFailures[name] = failures[count:] + } + return firstErr } // drainDanglingToolResults backfills a result for every announced tool // call that never produced one (user stop, fatal tool-node failure). The @@ -327,10 +393,33 @@ func runInner( // backfill carries no error: an interrupted call is not a failed call, and // a non-nil Err would paint every front-end's tool row red (raw // "context.Canceled" text) right next to the calm stop notice. - drainDanglingToolResults := func() { - for id, started := range toolStarts { - emitToolResult(started.name, session.InterruptedToolOutput, id, nil) + drainDanglingToolResults := func() error { + var firstErr error + names := make([]string, 0, len(pendingToolFailures)) + for name := range pendingToolFailures { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + if persistErr := resolvePendingToolFailures(name, true); persistErr != nil && firstErr == nil { + firstErr = persistErr + } + } + + ids := make([]string, 0, len(toolStarts)) + for id := range toolStarts { + ids = append(ids, id) } + sort.Slice(ids, func(i, j int) bool { + return toolStarts[ids[i]].order < toolStarts[ids[j]].order + }) + for _, id := range ids { + started := toolStarts[id] + if persistErr := emitToolResult(started.name, session.InterruptedToolOutput, id, nil); persistErr != nil && firstErr == nil { + firstErr = persistErr + } + } + return firstErr } config.Logger().Printf("[runner] runInner start, messages=%d", len(messages)) @@ -346,9 +435,12 @@ func runInner( // calm "Stopped". runInner owns this OnAgentDone (returns done=true), // so Run does not emit a second one. config.Logger().Printf("[runner] context cancelled, stopping iteration") - drainDanglingToolResults() - h.OnAgentDone(ctx.Err()) - return finish(true) + runErr := ctx.Err() + if persistErr := drainDanglingToolResults(); persistErr != nil { + runErr = errors.Join(runErr, persistErr) + } + h.OnAgentDone(runErr) + return finish(true, runErr) default: } @@ -367,18 +459,24 @@ func runInner( // "[NodeRunError] context canceled"); report the clean context // error instead of the noisy wrapped one. config.Logger().Printf("[runner] event error during cancellation: %v", event.Err) - drainDanglingToolResults() - h.OnAgentDone(ctx.Err()) - return finish(true) + runErr := ctx.Err() + if persistErr := drainDanglingToolResults(); persistErr != nil { + runErr = errors.Join(runErr, persistErr) + } + h.OnAgentDone(runErr) + return finish(true, runErr) } // Log the provider's raw payload, hand the frontends a sentence a // human can act on. This is the single choke point for model errors, // so wrapping here fixes the display in the TUI, the web UI and ACP // at once — and stops the next frontend from having to remember. config.Logger().Printf("[runner] event error: %v", event.Err) - drainDanglingToolResults() - h.OnAgentDone(internalmodel.WrapFriendly(event.Err, "", "")) - return finish(true) + runErr := internalmodel.WrapFriendly(event.Err, "", "") + if persistErr := drainDanglingToolResults(); persistErr != nil { + runErr = errors.Join(runErr, persistErr) + } + h.OnAgentDone(runErr) + return finish(true, runErr) } if event.Output == nil || event.Output.MessageOutput == nil { config.Logger().Printf("[runner] event #%d: nil output", eventCount) @@ -392,13 +490,24 @@ func runInner( toolName := mo.ToolName if !mo.IsStreaming && mo.Message != nil { output := toolMessageText(mo.Message) - emitToolResult(toolName, output, mo.Message.ToolCallID, nil) + persistErr := emitToolResult(toolName, output, mo.Message.ToolCallID, nil) + if pendingErr := resolvePendingToolFailures(toolName, false); pendingErr != nil { + persistErr = errors.Join(persistErr, pendingErr) + } + if persistErr != nil { + if drainErr := drainDanglingToolResults(); drainErr != nil { + persistErr = errors.Join(persistErr, drainErr) + } + h.OnAgentDone(persistErr) + return finish(true, persistErr) + } if toolName == "todowrite" || toolName == "todoread" { h.OnTodoUpdate() } } else if mo.IsStreaming { var sb strings.Builder var toolErr error + var persistErr error var toolCallID string for { chunk, err := mo.MessageStream.Recv() @@ -407,7 +516,11 @@ func runInner( } if err != nil { toolErr = err - emitToolResult(toolName, "", toolCallID, err) + if toolCallID == "" { + pendingToolFailures[toolName] = append(pendingToolFailures[toolName], err) + } else if writeErr := emitToolResult(toolName, "", toolCallID, err); writeErr != nil { + persistErr = errors.Join(persistErr, writeErr) + } break } if chunk != nil { @@ -418,11 +531,27 @@ func runInner( } } if toolErr == nil { - emitToolResult(toolName, sb.String(), toolCallID, nil) + if writeErr := emitToolResult(toolName, sb.String(), toolCallID, nil); writeErr != nil { + persistErr = errors.Join(persistErr, writeErr) + } if toolName == "todowrite" || toolName == "todoread" { h.OnTodoUpdate() } } + if pendingErr := resolvePendingToolFailures(toolName, false); pendingErr != nil { + persistErr = errors.Join(persistErr, pendingErr) + } + // A tool execution error is model-visible and normally non-fatal; + // only a recorder failure (joined above) terminates the runner. The + // pending first-receive failure remains paired and persisted, then the + // ADK decides whether the agent can continue. + if persistErr != nil { + if drainErr := drainDanglingToolResults(); drainErr != nil { + persistErr = errors.Join(persistErr, drainErr) + } + h.OnAgentDone(persistErr) + return finish(true, persistErr) + } } continue } @@ -488,12 +617,20 @@ func runInner( } if ctx.Err() != nil { config.Logger().Printf("[runner] assistant stream cancelled: %v", streamErr) - h.OnAgentDone(ctx.Err()) - return finish(true) + runErr := ctx.Err() + if persistErr := drainDanglingToolResults(); persistErr != nil { + runErr = errors.Join(runErr, persistErr) + } + h.OnAgentDone(runErr) + return finish(true, runErr) } config.Logger().Printf("[runner] assistant stream error: %v", streamErr) - h.OnAgentDone(internalmodel.WrapFriendly(streamErr, "", "")) - return finish(true) + runErr := internalmodel.WrapFriendly(streamErr, "", "") + if persistErr := drainDanglingToolResults(); persistErr != nil { + runErr = errors.Join(runErr, persistErr) + } + h.OnAgentDone(runErr) + return finish(true, runErr) } // Flush assistant text at the end of each assistant message so the // session file preserves the true message/tool interleaving. Without @@ -533,7 +670,8 @@ func runInner( startedAt := time.Now() for i, idx := range indices { p := pending[idx] - toolStarts[p.id] = toolStart{at: startedAt, name: p.name} + toolStarts[p.id] = toolStart{at: startedAt, name: p.name, order: nextToolOrder} + nextToolOrder++ h.OnToolCall(handler.ToolCallEvent{ Name: p.name, Args: p.args.String(), @@ -565,7 +703,8 @@ func runInner( startedAt := time.Now() size := len(mo.Message.ToolCalls) for i, tc := range mo.Message.ToolCalls { - toolStarts[tc.ID] = toolStart{at: startedAt, name: tc.Function.Name} + toolStarts[tc.ID] = toolStart{at: startedAt, name: tc.Function.Name, order: nextToolOrder} + nextToolOrder++ h.OnToolCall(handler.ToolCallEvent{ Name: tc.Function.Name, Args: tc.Function.Arguments, @@ -592,8 +731,13 @@ func runInner( } } - // Clean completion: Run emits the single final OnAgentDone(nil). - return finish(false) + // Clean completion: pair any framework-abandoned tool streams before Run + // emits the single final OnAgentDone(nil). + if persistErr := drainDanglingToolResults(); persistErr != nil { + h.OnAgentDone(persistErr) + return finish(true, persistErr) + } + return finish(false, nil) } // buildTokenUsage snapshots a tracker into a handler.TokenUsage. TotalTokens is diff --git a/internal/session/session.go b/internal/session/session.go index 1a7fa6c8..038bcff4 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -582,18 +582,39 @@ func (r *Recorder) RecordToolCall(name, args, toolCallID, batchID string, batchI // the approval-wait-adjusted execution latency (0 when unknown). Large outputs // are automatically truncated (head+tail preserved) and the full content is // saved to an overflow file on disk. Callers should use the returned output in -// live model history so live and replayed sessions have the same context. -func (r *Recorder) RecordToolResult(name, output, toolCallID string, err error, denied bool, duration time.Duration) string { +// live model history so live and replayed sessions have the same context. A +// persistence failure is returned so callers do not treat an unstored result as +// durable conversation history. +func (r *Recorder) RecordToolResult(name, output, toolCallID string, err error, denied bool, duration time.Duration) (string, error) { errStr := "" if err != nil { errStr = err.Error() } output = TruncateToolOutput(output, r.uuid, toolCallID) - _ = r.writeEntry(Entry{ + if writeErr := r.writeEntry(Entry{ Type: EntryToolResult, Name: name, Output: output, ToolCallID: toolCallID, Error: errStr, Denied: denied, DurationMs: duration.Milliseconds(), - }) - return output + }); writeErr != nil { + return output, fmt.Errorf("record tool result in session %s for tool call %q: %w", + r.sessionFilePathForError(), toolCallID, writeErr) + } + return output, nil +} + +// sessionFilePathForError returns the recorder's intended transcript path for +// diagnostics. It is best-effort because resolving the config directory may be +// the same operation that caused the write to fail. +func (r *Recorder) sessionFilePathForError() string { + dir, err := config.SessionsDir() + if err != nil { + return r.UUID() + } + r.mu.Lock() + defer r.mu.Unlock() + if r.agentID != "" { + return filepath.Join(dir, r.customDir, "subagents", "agent-"+r.agentID+".jsonl") + } + return filepath.Join(dir, r.uuid+".json") } // RecordToolObservation appends metadata-only progressive-disclosure evidence.