wenshao · May 21, 2026 · May 15, 2026 · May 15, 2026 · May 15, 2026 · May 15, 2026
@@ -699,6 +699,116 @@ describe('GeminiChat', async () => {
       ).resolves.not.toThrow();
     });
 
+    it('persists partial assistant turn when stream throws after a tool_use chunk', async () => {
+      // Weak-network scenario: Anthropic-compatible providers emit the
+      // `functionCall` part on `content_block_stop`; the SSE may then drop
+      // before `message_stop`. The yielded chunk is enough for `Turn.run`
+      // to queue a `ToolCallRequest`, the tool scheduler will eventually
+      // submit a `functionResponse` user turn — without a matching
+      // tool_use in history, the next request body shows
+      // `user → user[tool_result]` and DeepSeek/Anthropic rejects with
+      // "tool_use_id ... must have a corresponding tool_use block in the
+      // previous message". `processStreamResponse` must persist the
+      // partial model turn before re-throwing so the pairing is intact.
+      mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
+      const networkError = new Error('SSE connection reset by peer');
+      const streamThatThrowsAfterToolCall = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: {
+                role: 'model',
+                parts: [
+                  {
+                    functionCall: {
+                      id: 'call_00_CeJrKJB0PSmXUZTCWHET7332',
+                      name: 'read_file',
+                      args: { path: '/tmp/x.txt' },
+                    },
+                  },
+                ],
+              },
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+        throw networkError;
+      })();
+      vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
+        streamThatThrowsAfterToolCall,
+      );
+
+      const stream = await chat.sendMessageStream(
+        'test-model',
+        { message: 'open /tmp/x.txt please' },
+        'prompt-weak-network-tool',
+      );
+      await expect(
+        (async () => {
+          for await (const _ of stream) {
+            /* drain */
+          }
+        })(),
+      ).rejects.toBe(networkError);
+
+      const history = chat.getHistory();
+      expect(history.length).toBe(2);
+      expect(history[0]!.role).toBe('user');
+      const modelTurn = history[1]!;
+      expect(modelTurn.role).toBe('model');
+      expect(modelTurn.parts).toBeDefined();
+      const functionCallPart = modelTurn.parts!.find((p) => p.functionCall);
+      expect(functionCallPart?.functionCall?.id).toBe(
+        'call_00_CeJrKJB0PSmXUZTCWHET7332',
+      );
+      expect(functionCallPart?.functionCall?.name).toBe('read_file');
+    });
+
+    it('does NOT persist partial assistant turn when stream throws before any tool_use chunk', async () => {
+      // Plain-text partial responses are deliberately dropped on stream
+      // error: the Retry path pops the trailing user prompt and re-issues
+      // it, so a stale partial-text model turn between them would bias
+      // the retry or surface as duplicate output. Only tool_use turns
+      // need the partial-history bridge to preserve the tool_use →
+      // tool_result invariant — text alone has no such invariant.
+      mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
+      const networkError = new Error('connection reset');
+      const streamThatThrowsAfterText = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: {
+                role: 'model',
+                parts: [{ text: 'partial reply that will be lost' }],
+              },
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+        throw networkError;
+      })();
+      vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
+        streamThatThrowsAfterText,
+      );
+
+      const stream = await chat.sendMessageStream(
+        'test-model',
+        { message: 'hello' },
+        'prompt-weak-network-text',
+      );
+      await expect(
+        (async () => {
+          for await (const _ of stream) {
+            /* drain */
+          }
+        })(),
+      ).rejects.toBe(networkError);
+
+      const history = chat.getHistory();
+      // Only the user turn is in history — the partial-text model turn is
+      // intentionally not persisted.
+      expect(history.length).toBe(1);
+      expect(history[0]!.role).toBe('user');
+    });
+
     it('should throw InvalidStreamError when no tool call and no finish reason', async () => {
       vi.useFakeTimers();
       try {

@@ -1257,48 +1257,62 @@ export class GeminiChat {
 
     let hasToolCall = false;
     let hasFinishReason = false;
+    // Captured if the upstream stream throws mid-iteration (typical on weak
+    // networks: SSE drops between `content_block_stop` of a tool_use and the
+    // terminal `message_stop`). We still build / record / push a partial
+    // assistant turn below before re-throwing — see the dedicated branch in
+    // the post-loop block for why this is needed to keep tool_use/tool_result
+    // pairing intact across the failure.
+    let streamError: unknown = null;
 
-    for await (const chunk of streamResponse) {
-      // Use ||= to avoid later usage-only chunks (no candidates) overwriting
-      // a finishReason that was already seen in an earlier chunk.
-      hasFinishReason ||=
-        chunk?.candidates?.some((candidate) => candidate.finishReason) ?? false;
-
-      if (isValidResponse(chunk)) {
-        const content = chunk.candidates?.[0]?.content;
-        if (content?.parts) {
-          if (content.parts.some((part) => part.functionCall)) {
-            hasToolCall = true;
-          }
+    try {
+      for await (const chunk of streamResponse) {
+        // Use ||= to avoid later usage-only chunks (no candidates) overwriting
+        // a finishReason that was already seen in an earlier chunk.
+        hasFinishReason ||=
+          chunk?.candidates?.some((candidate) => candidate.finishReason) ??
+          false;
+
+        if (isValidResponse(chunk)) {
+          const content = chunk.candidates?.[0]?.content;
+          if (content?.parts) {
+            if (content.parts.some((part) => part.functionCall)) {
+              hasToolCall = true;
+            }
 
-          // Collect all parts for recording
-          allModelParts.push(...content.parts);
+            // Collect all parts for recording
+            allModelParts.push(...content.parts);
+          }
         }
-      }
 
-      // Collect token usage for consolidated recording
-      if (chunk.usageMetadata) {
-        usageMetadata = chunk.usageMetadata;
-        // Context usage tracks prompt size; output isn't in history yet.
-        const lastPromptTokenCount =
-          usageMetadata.promptTokenCount || usageMetadata.totalTokenCount;
-        if (lastPromptTokenCount) {
-          // Always update the per-chat counter so this chat (including
-          // subagents) can make its own compaction decisions.
-          this.lastPromptTokenCount = lastPromptTokenCount;
-          // Mirror to the global telemetry only when wired — subagents
-          // pass `telemetryService=undefined` to keep their context usage
-          // out of the main session's UI counters.
-          this.telemetryService?.setLastPromptTokenCount(lastPromptTokenCount);
-        }
-        if (usageMetadata.cachedContentTokenCount && this.telemetryService) {
-          this.telemetryService.setLastCachedContentTokenCount(
-            usageMetadata.cachedContentTokenCount,
-          );
+        // Collect token usage for consolidated recording
+        if (chunk.usageMetadata) {
+          usageMetadata = chunk.usageMetadata;
+          // Context usage tracks prompt size; output isn't in history yet.
+          const lastPromptTokenCount =
+            usageMetadata.promptTokenCount || usageMetadata.totalTokenCount;
+          if (lastPromptTokenCount) {
+            // Always update the per-chat counter so this chat (including
+            // subagents) can make its own compaction decisions.
+            this.lastPromptTokenCount = lastPromptTokenCount;
+            // Mirror to the global telemetry only when wired — subagents
+            // pass `telemetryService=undefined` to keep their context usage
+            // out of the main session's UI counters.
+            this.telemetryService?.setLastPromptTokenCount(
+              lastPromptTokenCount,
+            );
+          }
+          if (usageMetadata.cachedContentTokenCount && this.telemetryService) {
+            this.telemetryService.setLastCachedContentTokenCount(
+              usageMetadata.cachedContentTokenCount,
+            );
+          }
         }
-      }
 
-      yield chunk; // Yield every chunk to the UI immediately.
+        yield chunk; // Yield every chunk to the UI immediately.
+      }
+    } catch (e) {
+      streamError = e;
     }
 
     let thoughtContentPart: Part | undefined;
@@ -1369,6 +1383,46 @@ export class GeminiChat {
       });
     }
 
+    // Mid-stream failure recovery: if the upstream stream threw (typical on
+    // weak networks — SSE cut between a tool_use `content_block_stop` and
+    // the terminal `message_stop`) AND any `functionCall` chunk was already
+    // yielded to consumers, we must persist the partial assistant turn here.
+    //
+    // The content generator (Anthropic / OpenAI) emits a `functionCall` part
+    // only at the end of a tool_use block. Once yielded, `Turn.run` registers
+    // a `ToolCallRequest` event, the React tool scheduler queues the call,
+    // and `handleCompletedTools` will fire `submitQuery(..., ToolResult)` —
+    // pushing a user message with `functionResponse` into history — even
+    // though the parent stream errored. Without preserving the matching
+    // tool_use on the model side, the next request body would have
+    // `user → user[tool_result]` with no tool_use in between, and the
+    // Anthropic-compatible API (DeepSeek, Anthropic, etc.) rejects with
+    //   "tool_use_id ... must have a corresponding tool_use block in the
+    //    previous message"
+    // — an unrecoverable state because Ctrl+Y's `stripOrphanedUserEntries`
+    // only strips trailing user entries; the lost tool_use can't be
+    // resurrected.
+    //
+    // Plain-text partial turns (no functionCall yielded) are deliberately
+    // NOT persisted — the Retry path pops the trailing user prompt and
+    // re-issues it; a stale partial-text model turn between them would
+    // either bias the retry or surface as a duplicate.
+    if (streamError !== null) {
+      if (
+        hasToolCall &&
+        (thoughtContentPart || consolidatedHistoryParts.length > 0)
+      ) {
+        this.history.push({
+          role: 'model',
+          parts: [
+            ...(thoughtContentPart ? [thoughtContentPart] : []),
+            ...consolidatedHistoryParts,
+          ],
+        });
+      }
+      throw streamError;
+    }
+
     // Stream validation logic: A stream is considered successful if:
     // 1. There's a tool call (tool calls can end without explicit finish reasons), OR
     // 2. There's a finish reason AND we have non-empty response text or thought text