diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index e57f16de64..ddf7fd666a 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -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 { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 6f7a11c791..c8aee1e35a 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -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; + 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; + 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, + ); + } + } + + yield chunk; // Yield every chunk to the UI immediately. } - - // 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. + } 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