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.
2026-09-22 05:05:20 +08:00 · 2026-03-12 16:59:16 -04:00
parent 072e9a212f
commit 42c12176a0
2 changed files with 197 additions and 3 deletions
@@ -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

@@ -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")
+}