From 42c12176a06cfdc1ad61494308b343b8daf19c2a Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Thu, 12 Mar 2026 13:59:16 -0700 Subject: [PATCH] fix(chatd): persist interrupted tool call steps instead of losing them (#23011) ## Problem When a chat is interrupted while tools are executing, the step content (text, reasoning, tool calls, and partial tool results) was being lost. Two gaps existed: 1. **During tool execution**: `executeTools` returns with error results for interrupted tools, but the subsequent `PersistStep(ctx, ...)` fails on the canceled context and returns `ErrInterrupted` without persisting anything. 2. **PersistStep race**: If the context is canceled between the post-tool interrupt check and the `PersistStep` call, the same loss occurs. This is inconsistent with how we handle stream interruptions (which properly flush and persist partial content via `persistInterruptedStep`) and how [coder/blink](https://github.com/coder/blink) handles interruptions (always inserting the response message regardless of execution phase). ## Fix Two changes in `chatloop.go`: - **Post-tool-execution interrupt check**: After `executeTools` returns, check if the context was interrupted and route through `persistInterruptedStep` (which uses `context.WithoutCancel` internally) to save the accumulated content. - **PersistStep fallback**: If `PersistStep` returns `ErrInterrupted`, retry via `persistInterruptedStep` so partial content is not lost. ## Tests - `TestRun_InterruptedDuringToolExecutionPersistsStep`: Verifies that when a tool is blocked and the chat is interrupted, the step (text + reasoning + tool call + tool error result) is persisted via the interrupt-safe path. - `TestRun_PersistStepInterruptedFallback`: Verifies that when `PersistStep` itself returns `ErrInterrupted`, the step is retried via the fallback path and content is saved. --- coderd/chatd/chatloop/chatloop.go | 24 +++- coderd/chatd/chatloop/chatloop_test.go | 176 +++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 3 deletions(-) diff --git a/coderd/chatd/chatloop/chatloop.go b/coderd/chatd/chatloop/chatloop.go index f3e3433e80..f493e38175 100644 --- a/coderd/chatd/chatloop/chatloop.go +++ b/coderd/chatd/chatloop/chatloop.go @@ -324,8 +324,20 @@ func Run(ctx context.Context, opts RunOptions) error { for _, tr := range toolResults { result.content = append(result.content, tr) } - } + // Check for interruption after tool execution. + // Tools that were canceled mid-flight produce error + // results via ctx cancellation. Persist the full + // step (assistant blocks + tool results) through + // the interrupt-safe path so nothing is lost. + if ctx.Err() != nil { + if errors.Is(context.Cause(ctx), ErrInterrupted) { + persistInterruptedStep(ctx, opts, &result) + return ErrInterrupted + } + return ctx.Err() + } + } // Extract context limit from provider metadata. contextLimit := extractContextLimit(result.providerMetadata) if !contextLimit.Valid && opts.ContextLimitFallback > 0 { @@ -334,15 +346,21 @@ func Run(ctx context.Context, opts RunOptions) error { Valid: true, } } - // Persist the step — errors propagate directly. + // Persist the step. If persistence fails because + // the chat was interrupted between the previous + // check and here, fall back to the interrupt-safe + // path so partial content is not lost. if err := opts.PersistStep(ctx, PersistedStep{ Content: result.content, Usage: result.usage, ContextLimit: contextLimit, }); err != nil { + if errors.Is(err, ErrInterrupted) { + persistInterruptedStep(ctx, opts, &result) + return ErrInterrupted + } return xerrors.Errorf("persist step: %w", err) } - lastUsage = result.usage lastProviderMetadata = result.providerMetadata diff --git a/coderd/chatd/chatloop/chatloop_test.go b/coderd/chatd/chatloop/chatloop_test.go index dd337354da..d715f558ca 100644 --- a/coderd/chatd/chatloop/chatloop_test.go +++ b/coderd/chatd/chatloop/chatloop_test.go @@ -588,3 +588,179 @@ func hasAnthropicEphemeralCacheControl(message fantasy.Message) bool { cacheOptions, ok := options.(*fantasyanthropic.ProviderCacheControlOptions) return ok && cacheOptions.CacheControl.Type == "ephemeral" } + +// TestRun_InterruptedDuringToolExecutionPersistsStep verifies that when +// tools are executing and the chat is interrupted, the accumulated step +// content (assistant blocks + tool results) is persisted via the +// interrupt-safe path rather than being lost. +func TestRun_InterruptedDuringToolExecutionPersistsStep(t *testing.T) { + t.Parallel() + + toolStarted := make(chan struct{}) + + // Model returns a completed tool call in the stream. + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "calling tool"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeReasoningStart, ID: "reason-1"}, + {Type: fantasy.StreamPartTypeReasoningDelta, ID: "reason-1", Delta: "let me think"}, + {Type: fantasy.StreamPartTypeReasoningEnd, ID: "reason-1"}, + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "slow_tool"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{"key":"value"}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "slow_tool", + ToolCallInput: `{"key":"value"}`, + }, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, + }), nil + }, + } + + // Tool that blocks until context is canceled, simulating + // a long-running operation interrupted by the user. + slowTool := fantasy.NewAgentTool( + "slow_tool", + "blocks until canceled", + func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + close(toolStarted) + <-ctx.Done() + return fantasy.ToolResponse{}, ctx.Err() + }, + ) + + ctx, cancel := context.WithCancelCause(context.Background()) + defer cancel(nil) + + go func() { + <-toolStarted + cancel(ErrInterrupted) + }() + + var persistedContent []fantasy.Content + persistedCtxErr := xerrors.New("unset") + + err := Run(ctx, RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "run the slow tool"), + }, + Tools: []fantasy.AgentTool{slowTool}, + MaxSteps: 3, + PersistStep: func(persistCtx context.Context, step PersistedStep) error { + persistedCtxErr = persistCtx.Err() + persistedContent = append([]fantasy.Content(nil), step.Content...) + return nil + }, + }) + require.ErrorIs(t, err, ErrInterrupted) + // persistInterruptedStep uses context.WithoutCancel, so the + // persist callback should see a non-canceled context. + require.NoError(t, persistedCtxErr) + require.NotEmpty(t, persistedContent) + + var ( + foundText bool + foundReasoning bool + foundToolCall bool + foundToolResult bool + ) + for _, block := range persistedContent { + if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { + if strings.Contains(text.Text, "calling tool") { + foundText = true + } + continue + } + if reasoning, ok := fantasy.AsContentType[fantasy.ReasoningContent](block); ok { + if strings.Contains(reasoning.Text, "let me think") { + foundReasoning = true + } + continue + } + if toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block); ok { + if toolCall.ToolCallID == "tc-1" && toolCall.ToolName == "slow_tool" { + foundToolCall = true + } + continue + } + if toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok { + if toolResult.ToolCallID == "tc-1" { + foundToolResult = true + } + } + } + require.True(t, foundText, "persisted content should include text from the stream") + require.True(t, foundReasoning, "persisted content should include reasoning from the stream") + require.True(t, foundToolCall, "persisted content should include the tool call") + require.True(t, foundToolResult, "persisted content should include the tool result (error from cancellation)") +} + +// TestRun_PersistStepInterruptedFallback verifies that when the normal +// PersistStep call returns ErrInterrupted (e.g., context canceled in a +// race), the step is retried via the interrupt-safe path. +func TestRun_PersistStepInterruptedFallback(t *testing.T) { + t.Parallel() + + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "hello world"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + }, + } + + var ( + mu sync.Mutex + persistCalls int + savedContent []fantasy.Content + ) + + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "hello"), + }, + MaxSteps: 1, + PersistStep: func(_ context.Context, step PersistedStep) error { + mu.Lock() + defer mu.Unlock() + persistCalls++ + if persistCalls == 1 { + // First call: simulate an interrupt race by + // returning ErrInterrupted without persisting. + return ErrInterrupted + } + // Second call (from persistInterruptedStep fallback): + // accept the content. + savedContent = append([]fantasy.Content(nil), step.Content...) + return nil + }, + }) + require.ErrorIs(t, err, ErrInterrupted) + + mu.Lock() + defer mu.Unlock() + require.Equal(t, 2, persistCalls, "PersistStep should be called twice: once normally (failing), once via fallback") + require.NotEmpty(t, savedContent) + + var foundText bool + for _, block := range savedContent { + if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { + if strings.Contains(text.Text, "hello world") { + foundText = true + } + } + } + require.True(t, foundText, "fallback should persist the text content") +}