diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index dd1c4e2cd56..b5f82bcbff6 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -53,6 +53,17 @@ const MockedGeminiClientClass = vi.hoisted(() => this.addHistory = vi.fn(); this.consumePendingMemoryTaskPromises = vi.fn().mockReturnValue([]); this.recordCompletedToolCall = vi.fn(); + // Default to the fast-path accessor returning an empty Set so the + // dedup dispatcher in `handleCompletedTools` takes the + // `getHistoryFunctionResponseIds` branch by default (matching + // production). Tests that need a non-empty dedup set override + // this. Without exposing the method at all, the dispatcher would + // fall through to the `structuredClone(getHistory())` slow path + // and any regression in the fast path would silently route + // production onto the expensive branch while CI stays green. + this.getHistoryFunctionResponseIds = vi + .fn() + .mockReturnValue(new Set()); this.getChatRecordingService = vi.fn().mockReturnValue({ recordThought: vi.fn(), initialize: vi.fn(), @@ -939,6 +950,659 @@ describe('useGeminiStream', () => { }); }); + it('drops a late tool result whose callId is already paired in chat.history (Race A dedup)', async () => { + // Race A repro: the chat-internal repair pass already synthesized a + // functionResponse for this callId on the Retry push (because the + // partial-tool_use turn was orphan when Ctrl+Y landed). The live + // scheduler's late real result must NOT also be submitted, otherwise + // the wire payload would carry two functionResponse parts for the + // same callId and the second one would land as an orphan tool_result. + // The dedup MUST run regardless of `isResponding`, because the + // scheduler's `onAllToolCallsComplete` is single-shot and would + // otherwise leave the tool stuck in `completed-but-not-submitted`. + const lateRealResult = { + request: { + callId: 'call_race_A', + name: 'read_file', + args: { path: '/tmp/x.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-race-a', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call_race_A', + responseParts: [ + { + functionResponse: { + id: 'call_race_A', + name: 'read_file', + response: { output: 'real file contents' }, + }, + }, + ], + resultDisplay: undefined, + error: undefined, + errorType: undefined, + }, + tool: { + name: 'read_file', + displayName: 'ReadFile', + description: 'Read a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'read /tmp/x.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + // Simulate the chat-internal repair pass having already planted a + // synthetic functionResponse for the same callId on the previous + // (Retry) push. The dedup dispatcher consults + // `getHistoryFunctionResponseIds` first; we override the default + // empty-Set mock to return the matching callId so the fast path + // is what production code exercises in this test (instead of + // falling through to the structuredClone slow path). + client.getHistoryFunctionResponseIds = vi + .fn() + .mockReturnValue(new Set(['call_race_A'])); + client.getHistory = vi.fn().mockReturnValue([ + { role: 'user', parts: [{ text: 'open /tmp/x.txt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_race_A', + name: 'read_file', + args: { path: '/tmp/x.txt' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { text: 'retry' }, + { + functionResponse: { + id: 'call_race_A', + name: 'read_file', + response: { + error: 'Tool execution result was not recorded', + }, + }, + }, + ], + }, + ]); + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([lateRealResult]); + } + }); + + await waitFor(() => { + // The dedup hit must `markToolsAsSubmitted` so the UI/scheduler is + // unblocked even though we drop the real result on the wire. + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['call_race_A']); + }); + + // The deduped tool DID run locally — `recordCompletedToolCall` must + // still fire so toolCallCount / skillsModifiedInSession reflect it, + // even though the wire-side submission is dropped. Regression guard: + // an earlier version filtered deduped tools out of `geminiTools` + // without recording, skipping the metric increment. + expect(client.recordCompletedToolCall).toHaveBeenCalledWith('read_file', { + path: '/tmp/x.txt', + }); + + // No follow-up submission: the synthetic in history already closes + // the tool_use ↔ tool_result pair. + expect(mockSendMessageStream).not.toHaveBeenCalled(); + }); + + it('skips recordCompletedToolCall for deduped CANCELLED tools (telemetry parity)', async () => { + // A deduped tool with status='cancelled' never actually produced + // model-visible output — counting it via `recordCompletedToolCall` + // (which increments toolCallCount and can flip + // skillsModifiedInSession on a skill-write path) would inflate the + // metric for a call that never ran end-to-end. Dedup must skip + // BOTH client-initiated (already skipped) AND cancelled tools, + // while still calling `markToolsAsSubmitted` so the scheduler + // unblocks. + const cancelledDedupedTool = { + request: { + callId: 'call_dedup_cancelled', + name: 'write_file', + args: { path: '/tmp/cancelled.txt', content: 'x' }, + isClientInitiated: false, + prompt_id: 'prompt-dedup-cancel', + }, + status: 'cancelled', + responseSubmittedToGemini: false, + response: { + callId: 'call_dedup_cancelled', + responseParts: [ + { + functionResponse: { + id: 'call_dedup_cancelled', + name: 'write_file', + response: { error: 'cancelled' }, + }, + }, + ], + resultDisplay: undefined, + error: undefined, + errorType: undefined, + }, + tool: { + name: 'write_file', + displayName: 'WriteFile', + description: 'Write a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'cancelled write', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCancelledToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + // Pre-paired in history: dedup will fire for this callId. Wire + // the fast-path accessor so the dispatcher takes the + // `getHistoryFunctionResponseIds` branch (matches production + // path; see the default mock comment in MockedGeminiClientClass). + client.getHistoryFunctionResponseIds = vi + .fn() + .mockReturnValue(new Set(['call_dedup_cancelled'])); + client.getHistory = vi.fn().mockReturnValue([ + { role: 'user', parts: [{ text: 'cancelled write' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_dedup_cancelled', + name: 'write_file', + args: { path: '/tmp/cancelled.txt', content: 'x' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_dedup_cancelled', + name: 'write_file', + response: { error: 'synthetic' }, + }, + }, + ], + }, + ]); + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([cancelledDedupedTool]); + } + }); + + // Scheduler still gets unblocked. + await waitFor(() => { + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ + 'call_dedup_cancelled', + ]); + }); + + // Telemetry NOT incremented — the cancelled filter held. + expect(client.recordCompletedToolCall).not.toHaveBeenCalled(); + }); + + it('runs Race A dedup BEFORE the isResponding early-return (regression guard)', async () => { + // The dedup block in handleCompletedTools is intentionally placed + // ABOVE the `if (isResponding) return;` early-return: the scheduler's + // `onAllToolCallsComplete` is single-shot per batch, so if the dedup + // sat below the guard a tool whose result was already paired in + // history would be left in `completed-but-not-submitted` forever + // whenever the late completion lands while the next stream is still + // in flight (isResponding=true). This test holds a stream open to + // pin isResponding=true, then asserts `markToolsAsSubmitted` still + // fires for the deduped callId. A future refactor that moves the + // dedup below the guard would silently break this and pass every + // other test. + const lateRealResult = { + request: { + callId: 'call_race_A_responding', + name: 'read_file', + args: { path: '/tmp/y.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-race-a-responding', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call_race_A_responding', + responseParts: [ + { + functionResponse: { + id: 'call_race_A_responding', + name: 'read_file', + response: { output: 'real file contents' }, + }, + }, + ], + resultDisplay: undefined, + error: undefined, + errorType: undefined, + }, + tool: { + name: 'read_file', + displayName: 'ReadFile', + description: 'Read a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'read /tmp/y.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + // Wire the fast-path accessor so the dispatcher takes the + // `getHistoryFunctionResponseIds` branch (matches production + // path; see the default mock comment in MockedGeminiClientClass). + client.getHistoryFunctionResponseIds = vi + .fn() + .mockReturnValue(new Set(['call_race_A_responding'])); + client.getHistory = vi.fn().mockReturnValue([ + { role: 'user', parts: [{ text: 'open /tmp/y.txt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_race_A_responding', + name: 'read_file', + args: { path: '/tmp/y.txt' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_race_A_responding', + name: 'read_file', + response: { + error: 'Tool execution result was not recorded', + }, + }, + }, + { text: 'next prompt' }, + ], + }, + ]); + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + // Held stream: never yields, never returns. Pins isResponding=true. + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + // Intentionally yield-less: holds the stream open without producing + // chunks so isResponding stays true while we trigger onComplete. + // eslint-disable-next-line require-yield + const heldStream = (async function* () { + await holdStream; + })(); + mockSendMessageStream.mockReturnValue(heldStream); + + const { result } = renderHook(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + // Kick the stream so submitQuery flips isResponding=true and parks + // on the first `await` inside the held async generator. + act(() => { + void result.current.submitQuery('next prompt'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + }); + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + + // Now fire the deduped completion while isResponding=true. + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([lateRealResult]); + } + }); + + // The dedup MUST still fire — markToolsAsSubmitted called with the + // deduped callId — even though the early-return on isResponding + // would otherwise skip every later branch. + await waitFor(() => { + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ + 'call_race_A_responding', + ]); + }); + + // No additional sendMessageStream: the held one is still the only + // call. The dedup path does NOT submit a new request. + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + + // Release the held stream so the test exits cleanly. + releaseStream(); + }); + + it('handles a mixed batch (one deduped + one non-deduped) without double-counting telemetry', async () => { + // The dedup filter on `geminiTools` (`!historyCallIdsWithResponse.has(callId)`) + // is the only thing preventing double `recordCompletedToolCall` + // for tools whose late real result lands AFTER the orphan-tool_use + // repair already planted a synthetic. Existing dedup tests supply + // ONLY deduped tools, so a regression that removed that filter + // would silently inflate `toolCallCount` (and flip + // `skillsModifiedInSession` for the SAME skill-write callId twice) + // without breaking any current test. + // + // Mixed-batch repro: scheduler completes two tools in the same + // batch — one whose callId already has a fr in history (deduped), + // one whose callId is fresh (must reach sendMessageStream). Pin: + // (a) markToolsAsSubmitted called with BOTH callIds, + // (b) recordCompletedToolCall fires once per non-deduped tool, + // NOT twice for the deduped one, + // (c) sendMessageStream IS called (the non-deduped tool's real + // result must reach the wire). + const dedupedTool = { + request: { + callId: 'call_mixed_deduped', + name: 'read_file', + args: { path: '/tmp/d.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-mixed', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call_mixed_deduped', + responseParts: [ + { + functionResponse: { + id: 'call_mixed_deduped', + name: 'read_file', + response: { output: 'late real for deduped' }, + }, + }, + ], + resultDisplay: undefined, + error: undefined, + errorType: undefined, + }, + tool: { + name: 'read_file', + displayName: 'ReadFile', + description: 'Read a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'read /tmp/d.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const freshTool = { + request: { + callId: 'call_mixed_fresh', + name: 'read_file', + args: { path: '/tmp/f.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-mixed', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call_mixed_fresh', + responseParts: [ + { + functionResponse: { + id: 'call_mixed_fresh', + name: 'read_file', + response: { output: 'real for fresh' }, + }, + }, + ], + resultDisplay: undefined, + error: undefined, + errorType: undefined, + }, + tool: { + name: 'read_file', + displayName: 'ReadFile', + description: 'Read a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'read /tmp/f.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + // Wire BOTH the fast-path accessor (`getHistoryFunctionResponseIds`) + // and the legacy `getHistory()` fallback. Wiring the fast path + // is the actual point of this test: production code prefers + // `getHistoryFunctionResponseIds` to skip the multi-millisecond + // `structuredClone` cost on long sessions, and an earlier + // version of this test only mocked `getHistory()` so the slow + // path was always the one exercised. We assert below that the + // fast path was the only one called — a regression that drops + // the fast-path branch from the dispatcher would silently + // re-route every batch onto the slow clone path with no test + // failure. + client.getHistoryFunctionResponseIds = vi + .fn() + .mockReturnValue(new Set(['call_mixed_deduped'])); + client.getHistory = vi.fn().mockReturnValue([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_mixed_deduped', + name: 'read_file', + args: { path: '/tmp/d.txt' }, + }, + }, + { + functionCall: { + id: 'call_mixed_fresh', + name: 'read_file', + args: { path: '/tmp/f.txt' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_mixed_deduped', + name: 'read_file', + response: { error: 'Tool execution result was not recorded' }, + }, + }, + ], + }, + ]); + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([dedupedTool, freshTool]); + } + }); + + await waitFor(() => { + // (a) Both callIds were marked submitted somewhere across the + // dedup pass (deduped) and the post-isResponding flow (fresh). + const allMarked = mockMarkToolsAsSubmitted.mock.calls.flatMap( + (call) => call[0] as string[], + ); + expect(allMarked).toContain('call_mixed_deduped'); + expect(allMarked).toContain('call_mixed_fresh'); + }); + + // (b) recordCompletedToolCall fires EXACTLY once per tool (deduped + // gets one call from the dedup-loop; fresh gets one from the + // geminiTools loop). The filter is what prevents the double + // record on the deduped callId. + const recordedCallIds = ( + client.recordCompletedToolCall as unknown as ReturnType + ).mock.calls.map((call) => (call[1] as { path: string }).path); + expect(recordedCallIds.filter((p) => p === '/tmp/d.txt').length).toBe(1); + expect(recordedCallIds.filter((p) => p === '/tmp/f.txt').length).toBe(1); + + // (c) The fresh tool's real result reaches sendMessageStream — + // dedup didn't accidentally suppress it. + expect(mockSendMessageStream).toHaveBeenCalled(); + + // (d) Fast-path was taken: `getHistoryFunctionResponseIds` was + // called for the dedup pass, and the cloning `getHistory()` + // fallback was NOT used by the dedup. (Other call sites in the + // hook may still call getHistory for their own purposes; we + // pin only that the dedup itself did not re-clone.) A future + // refactor that drops the fast-path branch from the dispatcher + // would re-route the dedup pass onto the structuredClone path + // and break this assertion — exactly the regression the + // accessor was added to prevent. + expect(client.getHistoryFunctionResponseIds).toHaveBeenCalled(); + expect(client.getHistory).not.toHaveBeenCalled(); + }); + it('should not flicker streaming state to Idle between tool completion and submission', async () => { const toolCallResponseParts: PartListUnion = [ { text: 'tool 1 final response' }, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 4e50b1f4b45..d6f937ee8a7 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -2037,10 +2037,6 @@ export const useGeminiStream = ( const handleCompletedTools = useCallback( async (completedToolCallsFromScheduler: TrackedToolCall[]) => { - if (isResponding) { - return; - } - const completedAndReadyToSubmitTools = completedToolCallsFromScheduler.filter( ( @@ -2063,9 +2059,87 @@ export const useGeminiStream = ( }, ); + // History-based dedup MUST run before the `isResponding` early-return. + // If a synthetic `functionResponse` for this callId is already in + // chat.history (planted on session-load by + // `client.repairOrphanedToolUseTurnsInHistory` or on every + // `chat.sendMessageStream` push by the inline repair pass), the + // in-flight scheduler result must be marked submitted NOW — + // `useReactToolScheduler.allToolCallsCompleteHandler` is single-shot + // per batch, so a later isResponding=true early-return would leave + // the tool stuck in `completed-but-not-submitted` forever (Race A + // surfaced in PR #4176 review). The real result is dropped on the + // wire — same trade-off upstream Claude Code makes when its + // `StreamingToolExecutor.discard()` follows a + // `yieldMissingToolResultBlocks` synthesis (`query.ts:733` + `:984`). + // Walk raw history WITHOUT cloning — `geminiClient.getHistory()` + // returns `structuredClone(this.history)`, which on long sessions + // (200+ entries with sizable tool outputs) costs several ms on + // the React UI thread and visibly stalls streaming when the + // dedup pass runs on every tool-completion batch. + // `getHistoryFunctionResponseIds` walks history in place and + // returns only the id Set this dispatcher needs. The + // GeminiClient implementation is mandatory — production and + // test mocks both expose it. Skip the dedup pass entirely if + // the client is missing (only happens in unit tests that + // construct a hook without a client). + const historyCallIdsWithResponse: Set = geminiClient + ? geminiClient.getHistoryFunctionResponseIds() + : new Set(); + const dedupedTools = completedAndReadyToSubmitTools.filter((tc) => + historyCallIdsWithResponse.has(tc.request.callId), + ); + const dedupedCallIds = dedupedTools.map((tc) => tc.request.callId); + if (dedupedCallIds.length > 0) { + debugLogger.warn( + `[REPAIR] Dropping ${dedupedCallIds.length} late tool result(s) ` + + `whose callId already has a functionResponse in history: ` + + `${dedupedCallIds.join(', ')}`, + ); + // Even though the wire-side submission is dropped, the tool DID + // run locally — `toolCallCount` and `skillsModifiedInSession` + // must reflect that. Without this, deduped skill-write tools + // (e.g. write_file under a project SKILLS path) would silently + // skip the `skillsModifiedInSession` flip that gates the + // skills-reload prompt at end-of-turn. Mirrors the + // `recordCompletedToolCall` loop below over `geminiTools` — + // filter to the same shape (non-client-initiated) so client + // tools (which the original loop also skipped) stay skipped. + // + // Cancelled tools are also skipped: `dedupedTools` includes + // anything in a terminal state (success | error | cancelled), + // but cancelled means the tool never actually ran end-to-end — + // the `allToolsCancelled` branch below would have surfaced + // them via `addHistory + reportCancelled` rather than the + // completed-call metric, and the metric should match. Without + // this filter, a deduped + cancelled tool would inflate + // `toolCallCount` for a call that never produced a result + // (and could also flip `skillsModifiedInSession` for a + // never-executed skill-write). + for (const tc of dedupedTools) { + if (tc.request.isClientInitiated) continue; + if (tc.status === 'cancelled') continue; + geminiClient?.recordCompletedToolCall( + tc.request.name, + tc.request.args as Record, + ); + } + markToolsAsSubmitted(dedupedCallIds); + } + + if (isResponding) { + return; + } + // Finalize any client-initiated tools as soon as they are done. + // Skip ones whose callId already lives in chat history with a + // matching `functionResponse` — the dedup block above already + // called `markToolsAsSubmitted` for those, and re-dispatching + // the same callIds here would queue an extra React render. const clientTools = completedAndReadyToSubmitTools.filter( - (t) => t.request.isClientInitiated, + (t) => + t.request.isClientInitiated && + !historyCallIdsWithResponse.has(t.request.callId), ); if (clientTools.length > 0) { markToolsAsSubmitted(clientTools.map((t) => t.request.callId)); @@ -2089,7 +2163,9 @@ export const useGeminiStream = ( } const geminiTools = completedAndReadyToSubmitTools.filter( - (t) => !t.request.isClientInitiated, + (t) => + !t.request.isClientInitiated && + !historyCallIdsWithResponse.has(t.request.callId), ); for (const toolCall of geminiTools) { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 95b2e7c2540..318514f00d8 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -1004,6 +1004,80 @@ describe('Gemini Client (client.ts)', () => { }); }); + describe('startChat — repair orphan tool_use on resume', () => { + it('synthesizes a functionResponse for a transcript ending in a dangling model[functionCall]', async () => { + // --resume of a session that crashed (OOM / SIGKILL / process exit) + // between the partial-tool_use push in `processStreamResponse` and + // the React scheduler's `submitQuery(ToolResult)`. The persisted + // JSONL ends with `model[functionCall]` and no matching user + // `functionResponse`. Without the repair pass running at session + // load, the first API call after `--resume` would 400 with + // "tool_use_id ... must have a corresponding tool_use block in + // the previous message" — exactly the wedge this PR is supposed + // to escape. Covers the only resume-time integration point for + // the repair, so a future reorder/removal of the call in + // `startChat()` regresses this test. + await client.startChat([ + { + role: 'user', + parts: [{ text: 'open /tmp/crash.txt' }], + }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_crash_resume', + name: 'read_file', + args: { path: '/tmp/crash.txt' }, + }, + } as never, + ], + }, + ]); + + const history = client.getHistory(); + // startChat prepends a mocked env-context user/model pair, then + // appends the supplied extraHistory; the repair pass must then + // splice a synthetic user[functionResponse] AFTER the dangling + // model[fc]. Locate the dangling model entry by its callId and + // verify the immediately-following entry carries the synthetic. + const danglingIdx = history.findIndex( + (h) => + h.role === 'model' && + h.parts?.some((p) => p.functionCall?.id === 'call_crash_resume'), + ); + expect(danglingIdx).toBeGreaterThanOrEqual(0); + const userAfter = history[danglingIdx + 1]; + expect(userAfter?.role).toBe('user'); + const fr = userAfter?.parts!.find((p) => p.functionResponse); + expect(fr?.functionResponse?.id).toBe('call_crash_resume'); + expect(fr?.functionResponse?.name).toBe('read_file'); + expect( + (fr?.functionResponse?.response as { error?: string })?.error, + ).toMatch(/interrupted/i); + }); + + it('is a no-op when the resumed transcript has no dangling tool_use', async () => { + // Happy resume path: don't inject a synthetic functionResponse + // into a transcript whose tool_use pairing is already valid (or, + // as here, has no tool_use at all). Defends against a future + // regression where the repair pass starts spuriously injecting on + // perfectly-formed history. + await client.startChat([ + { role: 'user', parts: [{ text: 'q' }] }, + { role: 'model', parts: [{ text: 'plain text reply' }] }, + ]); + + const history = client.getHistory(); + // No functionResponse anywhere — repair did nothing. + const hasAnyFunctionResponse = history.some((h) => + h.parts?.some((p) => p.functionResponse), + ); + expect(hasAnyFunctionResponse).toBe(false); + }); + }); + describe('setTools — system instruction refresh', () => { // Regression coverage for the progressive-MCP wiring bug: when MCP // discovery completes AFTER startChat() (the new default), `setTools()` @@ -1464,6 +1538,7 @@ describe('Gemini Client (client.ts)', () => { getHistory: vi.fn().mockReturnValue([]), getHistoryLength, stripOrphanedUserEntriesFromHistory, + repairOrphanedToolUseTurns: vi.fn().mockReturnValue({ injected: [] }), } as unknown as GeminiChat; mockTurnRunFn.mockReturnValue( (async function* () { @@ -4772,6 +4847,7 @@ Other open files: getHistoryLength: vi.fn().mockReturnValueOnce(3).mockReturnValue(2), setHistory: vi.fn(), stripOrphanedUserEntriesFromHistory: vi.fn(), + repairOrphanedToolUseTurns: vi.fn().mockReturnValue({ injected: [] }), }; client['chat'] = mockChat as GeminiChat; @@ -4804,6 +4880,7 @@ Other open files: getHistoryLength: vi.fn().mockReturnValue(0), setHistory: vi.fn(), stripOrphanedUserEntriesFromHistory: vi.fn(), + repairOrphanedToolUseTurns: vi.fn().mockReturnValue({ injected: [] }), }; client['chat'] = mockChat as GeminiChat; diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 14cc4421f3e..2a35fc2ff44 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -355,6 +355,21 @@ export class GeminiClient { return undefined; } + /** + * Walk-only accessor for the set of `functionResponse.id` strings in + * raw history. Callers that only need the dedup id set (notably + * `useGeminiStream.handleCompletedTools`) MUST prefer this over + * {@link getHistory}, which deep-clones the entire conversation via + * `structuredClone` on every call. On long sessions with sizable + * tool outputs the clone is a multi-millisecond hit on the React UI + * thread; running it on every tool-completion batch caused visible + * frame drops during streaming. See + * `GeminiChat.getHistoryFunctionResponseIds` for the implementation. + */ + getHistoryFunctionResponseIds(): Set { + return this.getChat().getHistoryFunctionResponseIds(); + } + /** * Pop orphaned trailing user entries from the in-memory chat history. * Used by: @@ -394,6 +409,62 @@ export class GeminiClient { this.forceFullIdeContext = true; } + /** + * Synthesize a `functionResponse` for every dangling `model[functionCall]` + * in chat history whose corresponding tool_result never landed. Inverse of + * {@link stripOrphanedUserEntriesFromHistory}, which only handles trailing + * `user` entries. + * + * This `GeminiClient` method is the resume-path entry point — called once + * from {@link startChat} after the transcript loads, covering `--resume` + * of a session that crashed between a partial-tool_use push and the + * tool's eventual completion. + * + * The other two coverage points (Retry submit path after + * `stripOrphanedUserEntriesFromHistory`, and the defensive pass at the + * start of every UserQuery / Cron send) live one layer down inside + * `GeminiChat.sendMessageStream` and call the standalone + * `repairOrphanedToolUseTurns(history)` function directly — they don't + * route through this wrapper. Anyone tracing the repair-pass coupling + * between the client and chat layers should follow that path + * separately rather than expect everything to funnel through here. + * + * Synthesizes an `error` `functionResponse`. The React tool scheduler + * (`useGeminiStream.handleCompletedTools`) MUST dedupe by `callId` against + * the live history before submitting its own `tool_result` — otherwise a + * late real result lands as a second `user[tool_result]` block (orphan + * because the synthetic already consumed the matching `tool_use`). + */ + repairOrphanedToolUseTurnsInHistory(reason?: string): { + injected: Array<{ callId: string; name: string }>; + droppedDuplicates: Array<{ callId: string; name: string }>; + } { + const result = this.getChat().repairOrphanedToolUseTurns(reason); + if (result.injected.length > 0) { + debugLogger.warn( + `[REPAIR] Synthesized ${result.injected.length} functionResponse(s) ` + + `for dangling tool_use(s): ${result.injected + .map((e) => `${e.name}(${e.callId})`) + .join(', ')}`, + ); + } + if (result.droppedDuplicates.length > 0) { + // Surface the duplicate-cleanup pass so investigators tracing + // a dedup-drop log have a breadcrumb pointing back to the + // repair function. Without this a duplicate-only repair (no + // synthesis, no hoist) leaves zero diagnostic trail and a + // future callId-collision bug would silently delete the + // wrong fr. + debugLogger.warn( + `[REPAIR] Dropped ${result.droppedDuplicates.length} duplicate ` + + `functionResponse(s) for callId(s): ${result.droppedDuplicates + .map((e) => `${e.name}(${e.callId})`) + .join(', ')}`, + ); + } + return result; + } + setHistory(history: Content[]) { this.getChat().setHistory(history); // Replacing history wholesale drops any prior read_file tool @@ -754,6 +825,21 @@ export class GeminiClient { uiTelemetryService, ); + // Repair any dangling `model[functionCall]` whose `functionResponse` + // never made it back into the transcript before we wrote the JSONL. + // The common cause is a process crash / OOM / SIGKILL between the + // partial-tool_use push (see `processStreamResponse`) and the React + // scheduler's tool_result submission. Without this pass, the first + // API call on a resumed session would 400 with the same + // `tool_use_id ... corresponding tool_use` error this whole + // subsystem is trying to escape. (Belt-and-suspenders: the same + // helper runs again inside `chat.sendMessageStream` after the user + // content is pushed, so a dangling left here by setHistory / + // compaction reordering is also caught — but doing it here keeps + // any pre-send code reading `chat.history` from seeing a malformed + // shape.) + this.repairOrphanedToolUseTurnsInHistory(); + const sessionStartAdditionalContext = await this.fireSessionStartHook(sessionStartSource); this.lastSessionStartContext = sessionStartAdditionalContext; @@ -1138,6 +1224,13 @@ export class GeminiClient { if (messageType === SendMessageType.Retry) { this.stripOrphanedUserEntriesFromHistory(); + // The matching dangling-`functionCall` repair runs inside + // `chat.sendMessageStream` AFTER the user content is pushed, so any + // tool_result the user is supplying (Retry of a ToolResult + // submission, lastPrompt === fr parts) closes the pair via the real + // `functionResponse` before we synthesize an error one. Doing the + // repair here would happen pre-push and race against the user + // content's own pairing — see PR #4176 review for the corner. } // Fire UserPromptSubmit hook through MessageBus (only if hooks are enabled) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 5f3caa976b3..b434aa08724 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -9,6 +9,7 @@ import type { Content, GenerateContentConfig, GenerateContentResponse, + Part, } from '@google/genai'; import { ApiError } from '@google/genai'; import { AuthType, type ContentGenerator } from '../core/contentGenerator.js'; @@ -597,6 +598,145 @@ describe('GeminiChat', async () => { 'This is the visible text that should not be lost.', ); }); + + it('synthesizes a functionResponse for a dangling tool_use before sending', async () => { + // End-to-end: when sendMessageStream is invoked on a chat whose + // history carries a dangling `model[functionCall]` (typical state + // after a Ctrl+Y race or a crash-resume on a partial-tool_use + // turn), the inline repair pass closes the pair against the + // just-pushed user content so the wire payload doesn't 400 with + // "tool_use_id ... corresponding tool_use". + chat.setHistory([ + { role: 'user', parts: [{ text: 'first message' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_dangling_for_send', + name: 'read_file', + args: { path: '/tmp/x' }, + }, + }, + ], + }, + ]); + + const ackStream = (async function* () { + yield { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + ackStream, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'next user prompt after a stream-error-mid-tool_use' }, + 'prompt-send-repair', + ); + for await (const _ of stream) { + /* drain */ + } + + const history = chat.getHistory(); + // The dangling fc should now be followed by a user turn that + // carries both the user-supplied text AND the synthetic fr that + // closes the pair. + const userTurn = history[2]!; + expect(userTurn.role).toBe('user'); + const fr = userTurn.parts!.find((p) => p.functionResponse); + expect(fr?.functionResponse?.id).toBe('call_dangling_for_send'); + expect(fr?.functionResponse?.name).toBe('read_file'); + expect( + (fr?.functionResponse?.response as { error?: string })?.error, + ).toMatch(/interrupted/i); + // The user's own text part is still present. + expect( + userTurn.parts!.some( + (p) => + p.text === 'next user prompt after a stream-error-mid-tool_use', + ), + ).toBe(true); + // tool_result block must come BEFORE the text — Anthropic- + // compatible backends reject a user message whose first content + // block isn't the tool_result answering the immediately preceding + // tool_use. Mirrors upstream Claude Code's `hoistToolResults`. + expect(userTurn.parts![0]!.functionResponse?.id).toBe( + 'call_dangling_for_send', + ); + }); + + it('does NOT synthesize when the user supplies a matching tool_result', async () => { + // Retry-of-ToolResult case (lastPrompt is a functionResponse Part + // array): the user-supplied tool_result must close the pair before + // the inline repair pass sees it, so no synthetic error is + // injected. Otherwise the wire payload would carry two + // functionResponse parts for the same callId — the real one and a + // bogus synthetic. + chat.setHistory([ + { role: 'user', parts: [{ text: 'do the read' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_retry_real_fr', + name: 'read_file', + args: { path: '/tmp/y' }, + }, + }, + ], + }, + ]); + + const ackStream = (async function* () { + yield { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ack' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + ackStream, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { + message: { + functionResponse: { + id: 'call_retry_real_fr', + name: 'read_file', + response: { output: 'real-tool-output' }, + }, + }, + }, + 'prompt-retry-real-fr', + ); + for await (const _ of stream) { + /* drain */ + } + + const userTurn = chat.getHistory()[2]!; + const frParts = userTurn.parts!.filter((p) => p.functionResponse); + // Exactly ONE functionResponse — the real one. No synthetic. + expect(frParts.length).toBe(1); + expect(frParts[0]!.functionResponse?.id).toBe('call_retry_real_fr'); + expect( + (frParts[0]!.functionResponse?.response as { output?: string })?.output, + ).toBe('real-tool-output'); + }); + it('should throw an error when a tool call is followed by an empty stream response', async () => { vi.useFakeTimers(); try { @@ -699,6 +839,186 @@ 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('preserves thinking parts alongside tool_use when stream throws mid-tool', async () => { + // Covers reasoning-mode providers (DeepSeek, Claude 4.6+) where the + // assistant turn carries both a thinking block and a tool_use. The + // partial-history push must keep the thinking part so DeepSeek's + // `injectThinkingOnToolUseTurns` converter pass sees an existing + // block on the replayed turn and does not pre-pend a synthetic one + // (which would discard the model's original reasoning text). + mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall()); + const networkError = new Error('SSE timeout'); + const streamWithThinkingAndTool = (async function* () { + yield { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'planning the read', thought: true }], + }, + }, + ], + } as unknown as GenerateContentResponse; + yield { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_thinking_tool_use', + name: 'read_file', + args: { path: '/tmp/a.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw networkError; + })(); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + streamWithThinkingAndTool, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'read /tmp/a.txt' }, + 'prompt-thinking-tool-weak-network', + ); + await expect( + (async () => { + for await (const _ of stream) { + /* drain */ + } + })(), + ).rejects.toBe(networkError); + + const history = chat.getHistory(); + expect(history.length).toBe(2); + const modelTurn = history[1]!; + expect(modelTurn.role).toBe('model'); + const parts = modelTurn.parts!; + // The thinking part must come before the functionCall — Anthropic + // requires thinking blocks first in the assistant content array. + expect(parts[0]!.thought).toBe(true); + expect(parts[0]!.text).toBe('planning the read'); + const functionCallPart = parts.find((p) => p.functionCall); + expect(functionCallPart?.functionCall?.id).toBe('call_thinking_tool_use'); + }); + + 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 { @@ -1945,46 +2265,203 @@ describe('GeminiChat', async () => { }); }); - describe('getHistoryTail', () => { - it('returns only the requested recent entries as a deep copy', () => { - const oldContent: Content = { role: 'user', parts: [{ text: 'old' }] }; - const recentContent: Content = { - role: 'model', - parts: [{ text: 'recent' }], - }; - chat.addHistory(oldContent); - chat.addHistory(recentContent); + describe('getHistoryFunctionResponseIds', () => { + // Walk-only accessor used by `useGeminiStream.handleCompletedTools` + // for the dedup pass. The whole point of this method is to avoid + // the multi-millisecond `structuredClone` hit that + // `getHistory()` pays on long sessions when only the id Set is + // needed. Pin the contract: returned Set contains every fr id + // present in user turns (including duplicates collapsed to one + // Set entry), and ignores parts that aren't functionResponses + // and turns that aren't user. + it('returns an empty Set for empty history', () => { + expect(chat.getHistoryFunctionResponseIds()).toEqual(new Set()); + }); - const tail = chat.getHistoryTail(1); + it('collects fr ids from user turns and ignores non-fr parts', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'go' }] }, + { + role: 'model', + parts: [ + { functionCall: { id: 'cid_a', name: 'read_file', args: {} } }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_a', + name: 'read_file', + response: { output: 'a' }, + }, + }, + { text: 'follow up' }, + ], + }, + ]); - expect(tail).toEqual([recentContent]); - expect(tail[0]).not.toBe(recentContent); - tail[0]!.parts![0]!.text = 'mutated'; - expect(chat.getHistory()[1]!.parts![0]!.text).toBe('recent'); + expect(chat.getHistoryFunctionResponseIds()).toEqual(new Set(['cid_a'])); }); - it('returns an empty tail for non-positive counts', () => { - chat.addHistory({ role: 'user', parts: [{ text: 'a' }] }); - expect(chat.getHistoryTail(0)).toEqual([]); - expect(chat.getHistoryTail(-1)).toEqual([]); + it('skips functionCall parts in model turns (only user[fr] counts)', () => { + // Defensive: a regression that walks all turns instead of just + // user turns would pull in `functionCall.id`s and double-count. + chat.setHistory([ + { + role: 'model', + parts: [ + { functionCall: { id: 'cid_model', name: 'read_file', args: {} } }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_user', + name: 'read_file', + response: { output: 'u' }, + }, + }, + ], + }, + ]); + + const ids = chat.getHistoryFunctionResponseIds(); + expect(ids).toEqual(new Set(['cid_user'])); + expect(ids.has('cid_model')).toBe(false); }); - }); - describe('getHistoryShallow', () => { - it('copies containers without structured-cloning large part payloads', () => { - const payload = { output: 'x'.repeat(128 * 1024) }; - const content: Content = { - role: 'user', - parts: [ - { - functionResponse: { - id: 'call-1', - name: 'read_file', - response: payload, + it('collapses duplicate fr ids across multiple user turns to one Set entry', () => { + // Same id echoed twice in different user turns: dedup callers + // only need to know "is this id paired anywhere", not the + // count, so a Set is sufficient and natural. + chat.setHistory([ + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_dup', + name: 'read_file', + response: { output: '1' }, + }, }, - }, - ], - }; + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_dup', + name: 'read_file', + response: { output: '2' }, + }, + }, + ], + }, + ]); + + const ids = chat.getHistoryFunctionResponseIds(); + expect(ids.size).toBe(1); + expect(ids.has('cid_dup')).toBe(true); + }); + + it('handles entries with no parts and parts with no functionResponse', () => { + // Defensive against malformed history (missing parts, parts + // with neither text nor fr): must not crash. + chat.setHistory([ + { role: 'user', parts: undefined as unknown as Part[] }, + { role: 'user', parts: [] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_ok', + name: 'read_file', + response: { output: 'ok' }, + }, + }, + ], + }, + ]); + + expect(chat.getHistoryFunctionResponseIds()).toEqual(new Set(['cid_ok'])); + }); + + it('does not deep-clone history (returns a fresh Set, not aliased to internal state)', () => { + // The whole reason this method exists is to avoid the + // structuredClone in getHistory(). Mutating the returned Set + // must not bleed into the next call. + chat.setHistory([ + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_immut', + name: 'read_file', + response: { output: 'v' }, + }, + }, + ], + }, + ]); + + const first = chat.getHistoryFunctionResponseIds(); + first.add('cid_FAKE'); + first.delete('cid_immut'); + + const second = chat.getHistoryFunctionResponseIds(); + expect(second.has('cid_immut')).toBe(true); + expect(second.has('cid_FAKE')).toBe(false); + }); + }); + + describe('getHistoryTail', () => { + it('returns only the requested recent entries as a deep copy', () => { + const oldContent: Content = { role: 'user', parts: [{ text: 'old' }] }; + const recentContent: Content = { + role: 'model', + parts: [{ text: 'recent' }], + }; + chat.addHistory(oldContent); + chat.addHistory(recentContent); + + const tail = chat.getHistoryTail(1); + + expect(tail).toEqual([recentContent]); + expect(tail[0]).not.toBe(recentContent); + tail[0]!.parts![0]!.text = 'mutated'; + expect(chat.getHistory()[1]!.parts![0]!.text).toBe('recent'); + }); + + it('returns an empty tail for non-positive counts', () => { + chat.addHistory({ role: 'user', parts: [{ text: 'a' }] }); + expect(chat.getHistoryTail(0)).toEqual([]); + expect(chat.getHistoryTail(-1)).toEqual([]); + }); + }); + + describe('getHistoryShallow', () => { + it('copies containers without structured-cloning large part payloads', () => { + const payload = { output: 'x'.repeat(128 * 1024) }; + const content: Content = { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call-1', + name: 'read_file', + response: payload, + }, + }, + ], + }; chat.addHistory(content); const structuredCloneSpy = vi .spyOn(globalThis, 'structuredClone') @@ -2261,6 +2738,360 @@ describe('GeminiChat', async () => { } }); + it('rolls back the partial assistant turn when a retryable error fires after a tool_use chunk', async () => { + // Regression for a stream attempt that yields a `functionCall` + // (which triggers the partial-history push in + // `processStreamResponse`), then throws a retryable error (e.g. + // a TPM 429 `StreamContentError`). The outer retry loop must + // drop the partial before issuing the + // retry — otherwise the retry's response lands as a SECOND + // consecutive `model` entry and the failed-attempt `tool_use` + // becomes orphan on the wire (invalid alternation + + // tool_use_id-with-no-matching-tool_use 400). + vi.useFakeTimers(); + try { + const tpmError = new StreamContentError( + '{"error":{"code":"429","message":"Throttling: TPM(1/1)"}}', + ); + const failingStream = (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_failed_retry_attempt', + name: 'read_file', + args: { path: '/tmp/a.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw tpmError; + })(); + const successStream = (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Success after retry' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(failingStream) + .mockResolvedValueOnce(successStream); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-rollback-on-retry', + ); + const iterator = stream[Symbol.asyncIterator](); + // Advance through the rate-limit RETRY + delay, drain all events. + for (;;) { + const next = iterator.next(); + await vi.advanceTimersByTimeAsync(60_000); + const r = await next; + if (r.done) break; + } + + const history = chat.getHistory(); + // History must NOT contain the failed attempt's partial + // model[functionCall]. Expected shape: [user, model(success + // text)] — exactly two entries, alternation intact. + expect(history.length).toBe(2); + expect(history[0]!.role).toBe('user'); + expect(history[1]!.role).toBe('model'); + const successText = history[1]!.parts!.find((p) => p.text)?.text; + expect(successText).toBe('Success after retry'); + // Defensively: NO functionCall anywhere in history. + expect(history.some((h) => h.parts?.some((p) => p.functionCall))).toBe( + false, + ); + } finally { + vi.useRealTimers(); + } + }); + + it('rolls back the partial assistant turn when an InvalidStreamError fires after a tool_use chunk on the transient-stream retry budget', async () => { + // Counterpart to the rate-limit rollback above. The + // transient-stream retry budget (NO_FINISH_REASON / + // NO_RESPONSE_TEXT) has its own popPartialIfPushed call site — + // separate from the rate-limit branch the existing test + // covers. Without a regression test, that call could be + // accidentally removed and the rate-limit test would still + // pass while a stale partial silently rode the retry. + vi.useFakeTimers(); + try { + const failingStream = (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_transient_retry_partial', + name: 'read_file', + args: { path: '/tmp/t.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + // Mid-tool_use cut without a finish reason — the transient- + // stream retry budget catches this and retries with delay. + throw new InvalidStreamError( + 'Model stream ended without a finish reason.', + 'NO_FINISH_REASON', + ); + })(); + const successStream = (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Recovered on retry' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(failingStream) + .mockResolvedValueOnce(successStream); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-rollback-transient', + ); + const iterator = stream[Symbol.asyncIterator](); + // Advance through the transient-retry delay (initial 2000 ms). + for (;;) { + const next = iterator.next(); + await vi.advanceTimersByTimeAsync(5_000); + const r = await next; + if (r.done) break; + } + + const history = chat.getHistory(); + // Final shape must be clean: [user, model(success text)]. + // The failed attempt's partial functionCall must NOT survive. + expect(history.length).toBe(2); + expect(history[0]!.role).toBe('user'); + expect(history[1]!.role).toBe('model'); + expect(history[1]!.parts!.find((p) => p.text)?.text).toBe( + 'Recovered on retry', + ); + expect(history.some((h) => h.parts?.some((p) => p.functionCall))).toBe( + false, + ); + } finally { + vi.useRealTimers(); + } + }); + + // NOTE: no test for the InvalidStreamError content-retry branch + // (geminiChat.ts ~line 1399). Verified unreachable for that error + // class: `isTransientStreamError` and `isContentError` are the + // same predicate (`error instanceof InvalidStreamError`), so the + // transient branch above always either `continue`s or `break`s + // before control reaches the content branch. The + // `popPartialIfPushed()` call there is preserved as + // defense-in-depth for a future error class that should diverge + // the predicates; see the comment block at that call site for + // the full analysis. + + it('rolls back the chat-recording entry too when the retry succeeds', async () => { + // The in-memory rollback test above asserts `this.history` ends + // clean after a retry-success. This test asserts the same about + // chat-recording JSONL: the failed attempt's `recordAssistantTurn` + // call must NOT have been flushed, so `--resume` won't rehydrate + // a model[functionCall] turn the live session correctly discarded. + // Without the deferred-flush stash + popPartialIfPushed clear, + // `recordAssistantTurn` was called twice (once for the partial, + // once for the success) and only the in-memory pop fixed live + // history; the durable transcript stayed corrupt. + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + + const tpmError = new StreamContentError( + '{"error":{"code":"429","message":"Throttling: TPM(1/1)"}}', + ); + const failingStream = (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_failed_retry_recording', + name: 'read_file', + args: { path: '/tmp/a.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw tpmError; + })(); + const successStream = (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Success after retry' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(failingStream) + .mockResolvedValueOnce(successStream); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-recording-rollback', + ); + const iterator = stream[Symbol.asyncIterator](); + for (;;) { + const next = iterator.next(); + await vi.advanceTimersByTimeAsync(60_000); + const r = await next; + if (r.done) break; + } + + // Exactly one recording: the successful retry's text turn. + // The failed attempt's partial functionCall must have been + // discarded by `popPartialIfPushed` clearing the deferred-flush + // stash, never reaching the JSONL. + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + const recordedMessage = recordAssistantTurn.mock.calls[0]![0] + ?.message as Array<{ text?: string; functionCall?: unknown }>; + const recordedText = recordedMessage.find((p) => p.text)?.text; + expect(recordedText).toBe('Success after retry'); + // No functionCall part anywhere in the recorded turn. + expect(recordedMessage.some((p) => p.functionCall)).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + it('flushes the chat-recording entry on the unretryable break path (kept partial → durable JSONL)', async () => { + // Counterpart to the rollback test: when the retry budget is + // exhausted (or the error is unretryable from the start), the + // partial assistant turn IS kept in `this.history` — and the + // chat-recording JSONL must match. Without the deferred-flush + // path firing at the rethrow site, the JSONL silently drops a + // partial that's still in live history, and the orphan-tool_use + // repair pass at session-load has no dangling functionCall to + // close → `--resume` first send 400s with the very wedge the + // repair was supposed to escape. + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + + // Unretryable: a non-rate-limit, non-InvalidStream error after + // a tool_use chunk lands. The catch block falls through to + // `break` with the partial kept in memory. + const failingStream = (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_unretryable_kept', + name: 'read_file', + args: { path: '/tmp/k.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw new Error('synthetic unretryable mid-stream failure'); + })(); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce(failingStream); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-recording-flush-on-break', + ); + const iterator = stream[Symbol.asyncIterator](); + await expect( + (async () => { + for (;;) { + const r = await iterator.next(); + if (r.done) return; + } + })(), + ).rejects.toThrow(/synthetic unretryable/); + + // In-memory: partial is kept (the wedge-recovery contract that + // the rest of this PR's machinery relies on). + const history = chatWithRecording.getHistory(); + const lastModelTurn = history.findLast((h) => h.role === 'model'); + expect( + lastModelTurn?.parts?.some( + (p) => p.functionCall?.id === 'call_unretryable_kept', + ), + ).toBe(true); + + // JSONL: must contain the same partial turn so `--resume` sees + // a transcript that matches live history. Exactly one record + // (no success retry happened on this path). + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + const recordedMessage = recordAssistantTurn.mock.calls[0]![0] + ?.message as Array<{ functionCall?: { id?: string } }>; + expect( + recordedMessage.some( + (p) => p.functionCall?.id === 'call_unretryable_kept', + ), + ).toBe(true); + } finally { + vi.useRealTimers(); + } + }); + it('should retry on TPM throttling StreamContentError with initial delay', async () => { vi.useFakeTimers(); @@ -3330,6 +4161,784 @@ describe('GeminiChat', async () => { }); }); + describe('partial-push marker invariants on history mutation', () => { + // The whole partial-push lifecycle relies on the invariant + // "every history-mutation method clears the partial-push markers" + // — six sites enforce it (clearHistory, addHistory, setHistory, + // truncateHistory, stripThoughtsFromHistory, + // stripOrphanedUserEntriesFromHistory). If any site forgets, a + // stale `pendingPartialAssistantTurnIndex` could line up with an + // unrelated model turn in the post-mutation history and cause + // `popPartialIfPushed` to splice the WRONG entry — silently losing + // a real assistant response. + // + // The markers are ephemeral within a single sendMessageStream + // call: the `finally` block flushes the deferred JSONL record + // and calls `clearPendingPartialState()` before the generator + // unwinds. So we can't observe non-null markers after a real + // mid-stream error completes — by that point the lifecycle has + // already cleared them. Instead, we plant the markers directly + // via the same private-field assignment the production code uses, + // then call each mutation method and verify both fields are reset + // in lockstep. This pins the invariant against future refactors + // that drop a `clearPendingPartialState()` call from one site + // while the other five still pass. + type PrivateFields = { + pendingPartialAssistantTurnIndex: number | null; + pendingPartialAssistantRecord: unknown; + }; + function plantMarkers(c: GeminiChat): void { + const internal = c as unknown as PrivateFields; + internal.pendingPartialAssistantTurnIndex = 0; + internal.pendingPartialAssistantRecord = { + model: 'test-model', + message: [{ functionCall: { id: 'call_test', name: 't', args: {} } }], + }; + } + function markers(c: GeminiChat): { + idx: number | null; + record: unknown; + } { + const internal = c as unknown as PrivateFields; + return { + idx: internal.pendingPartialAssistantTurnIndex, + record: internal.pendingPartialAssistantRecord, + }; + } + + it('clearHistory() clears the partial-push markers', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.clearHistory(); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + + it('addHistory() clears the partial-push markers (violation path)', () => { + // addHistory is documented to be called between sends, NOT + // mid-send. Calling it with markers active is a violation — + // the implementation logs a warn so the offending caller is + // visible in diagnostics, then clears the markers. + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.addHistory({ role: 'user', parts: [{ text: 'between sends' }] }); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + + it('setHistory() clears the partial-push markers', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.setHistory([{ role: 'user', parts: [{ text: 'replacement' }] }]); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + + it('truncateHistory() clears the partial-push markers', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.truncateHistory(1); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + + it('stripThoughtsFromHistory() clears the partial-push markers', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.stripThoughtsFromHistory(); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + + it('stripOrphanedUserEntriesFromHistory() clears the partial-push markers', () => { + // History tail is a model turn — strip is a no-op on history, + // but the marker reset must still fire so all six mutation + // sites stay uniform. + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [{ functionCall: { id: 'x', name: 't', args: {} } }], + }, + ]); + plantMarkers(chat); + expect(markers(chat).idx).toBe(0); + + chat.stripOrphanedUserEntriesFromHistory(); + + expect(markers(chat).idx).toBeNull(); + expect(markers(chat).record).toBeNull(); + }); + }); + + describe('repairOrphanedToolUseTurns', () => { + // Verifies the inverse-of-strip pass: every `model[functionCall]` + // without a matching `user[functionResponse]` in the next turn gets + // a synthesized error functionResponse. This closes the + // tool_use ↔ tool_result wire invariant for the residual races + // (`--resume` of a crashed session, Ctrl+Y before in-flight tool + // finishes, scheduler abort before submitQuery, manual JSONL edits). + + it('injects a synthetic functionResponse for a trailing tool_use (Race B/C)', () => { + // --resume of a session that crashed after the partial-tool_use push + // in `processStreamResponse` but before the scheduler submitted the + // tool_result. First API call would 400 without repair. + chat.setHistory([ + { role: 'user', parts: [{ text: 'open /tmp/a.txt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_crash_A', + name: 'read_file', + args: { path: '/tmp/a.txt' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([ + { callId: 'call_crash_A', name: 'read_file' }, + ]); + const history = chat.getHistory(); + expect(history.length).toBe(3); + expect(history[2]!.role).toBe('user'); + const fr = history[2]!.parts![0]!.functionResponse; + expect(fr?.id).toBe('call_crash_A'); + expect(fr?.name).toBe('read_file'); + expect((fr?.response as { error?: string })?.error).toMatch( + /interrupted/i, + ); + }); + + it('hoists synthetic functionResponse to the front of an existing user turn (Race A)', () => { + // Ctrl+Y race: the user retried while the in-flight tool was still + // running. `stripOrphanedUserEntriesFromHistory` leaves the + // model[functionCall] in place (trailing entry is model), then the + // Retry pushes a fresh user turn with the user prompt. Repair must + // splice the synthetic response onto that user turn so it sits + // immediately after the model[tool_use] — NOT create a stray + // synthetic user turn between them. Crucially the synthetic + // functionResponse must come BEFORE the text part: Anthropic- + // compatible backends require tool_result blocks to be first in + // the user message (mirrors upstream Claude Code's + // `hoistToolResults`). Otherwise the wire payload re-triggers the + // "tool_use_id ... must have a corresponding tool_use block in the + // previous message" 400 this PR is supposed to escape. + chat.setHistory([ + { role: 'user', parts: [{ text: 'open /tmp/a.txt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_race_A', + name: 'read_file', + args: { path: '/tmp/a.txt' }, + }, + }, + ], + }, + { role: 'user', parts: [{ text: 'retry prompt' }] }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected.map((e) => e.callId)).toEqual(['call_race_A']); + const history = chat.getHistory(); + expect(history.length).toBe(3); + expect(history[2]!.role).toBe('user'); + expect(history[2]!.parts!.length).toBe(2); + // synthetic fr FIRST, user text AFTER. + expect(history[2]!.parts![0]!.functionResponse?.id).toBe('call_race_A'); + expect(history[2]!.parts![1]).toEqual({ text: 'retry prompt' }); + }); + + it('hoists synthetic functionResponse AFTER pre-existing real ones (parallel partial submit)', () => { + // Parallel tool_use with one real functionResponse already in the + // user turn — synthetic for the missing callId must slot in + // between the real fr and any non-fr parts so the user message + // shape stays `[real_fr, synthetic_fr, text]` (every tool_result + // before any other content, preserving the real-fr order). + chat.setHistory([ + { role: 'user', parts: [{ text: 'batch read' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'call_A', name: 'read_file', args: {} }, + }, + { + functionCall: { id: 'call_B', name: 'read_file', args: {} }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_A', + name: 'read_file', + response: { output: 'a' }, + }, + }, + { text: 'retry prompt' }, + ], + }, + ]); + + chat.repairOrphanedToolUseTurns(); + const parts = chat.getHistory()[2]!.parts!; + expect(parts.length).toBe(3); + expect(parts[0]!.functionResponse?.id).toBe('call_A'); + expect(parts[1]!.functionResponse?.id).toBe('call_B'); + expect(parts[2]).toEqual({ text: 'retry prompt' }); + }); + + it('handles parallel tool_use turns with only some responses present', () => { + // Common shape after #4176's partial-history push: the stream + // emitted multiple `content_block_stop`s for parallel tool_uses, + // but the React scheduler only submitted some before the user hit + // Ctrl+Y. The Retry path's repair must close every missing pair — + // the present `functionResponse` for A must NOT be duplicated. + chat.setHistory([ + { role: 'user', parts: [{ text: 'batch read' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_A', + name: 'read_file', + args: { path: '/a' }, + }, + }, + { + functionCall: { + id: 'call_B', + name: 'read_file', + args: { path: '/b' }, + }, + }, + { + functionCall: { + id: 'call_C', + name: 'read_file', + args: { path: '/c' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_A', + name: 'read_file', + response: { output: 'a-content' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + const injectedIds = result.injected.map((e) => e.callId); + expect(injectedIds.sort()).toEqual(['call_B', 'call_C']); + const history = chat.getHistory(); + // Same shape — synthetics merge into the existing user turn. + expect(history.length).toBe(3); + const fr = history[2]!.parts!.map((p) => p.functionResponse?.id); + expect(fr).toEqual(['call_A', 'call_B', 'call_C']); + // The pre-existing `call_A` response is untouched (real result kept). + expect( + ( + history[2]!.parts![0]!.functionResponse?.response as { + output?: string; + } + )?.output, + ).toBe('a-content'); + }); + + it('is a no-op when every tool_use already has a matching response', () => { + // Happy path: don't churn history when the invariant already holds. + const happy = [ + { role: 'user' as const, parts: [{ text: 'q' }] }, + { + role: 'model' as const, + parts: [ + { + functionCall: { + id: 'call_ok', + name: 'read_file', + args: {}, + }, + }, + ], + }, + { + role: 'user' as const, + parts: [ + { + functionResponse: { + id: 'call_ok', + name: 'read_file', + response: { output: 'fine' }, + }, + }, + ], + }, + ]; + chat.setHistory(structuredClone(happy)); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([]); + expect(chat.getHistory()).toEqual(happy); + }); + + it('repairs multiple non-adjacent dangling tool_uses across history', () => { + // Stress case for the forward-walk algorithm: dangling turn near the + // start AND another near the end. Both should be repaired and the + // outer loop must not re-scan synthetic user turns it just inserted. + chat.setHistory([ + { + role: 'model', + parts: [ + { + functionCall: { + id: 'early_orphan', + name: 'glob', + args: {}, + }, + }, + ], + }, + { role: 'user', parts: [{ text: 'second user prompt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'late_orphan', + name: 'read_file', + args: { path: '/x' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + const injectedIds = result.injected.map((e) => e.callId); + expect(injectedIds.sort()).toEqual(['early_orphan', 'late_orphan']); + const history = chat.getHistory(); + // early_orphan got the synthetic spliced into the existing user turn + // between the two model entries; late_orphan got a brand-new + // trailing user turn appended after the second model entry. + expect(history.length).toBe(4); + expect(history[0]!.role).toBe('model'); + expect(history[1]!.role).toBe('user'); + expect( + history[1]!.parts!.some( + (p) => p.functionResponse?.id === 'early_orphan', + ), + ).toBe(true); + expect(history[2]!.role).toBe('model'); + expect(history[3]!.role).toBe('user'); + expect(history[3]!.parts![0]!.functionResponse?.id).toBe('late_orphan'); + }); + + it('ignores model turns with no functionCall parts', () => { + const plain = [ + { role: 'user' as const, parts: [{ text: 'hi' }] }, + { role: 'model' as const, parts: [{ text: 'hello' }] }, + ]; + chat.setHistory(structuredClone(plain)); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([]); + expect(chat.getHistory()).toEqual(plain); + }); + + it('uses caller-provided reason text', () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'q' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'cid', name: 'read_file', args: {} }, + }, + ], + }, + ]); + + chat.repairOrphanedToolUseTurns('custom reason'); + + const fr = chat.getHistory()[2]!.parts![0]!.functionResponse; + expect((fr?.response as { error?: string })?.error).toBe('custom reason'); + }); + + it('hoists the real functionResponse from a non-adjacent later user turn into the adjacent one', () => { + // Regression for the shape + // `[user, model[fc], user[text], user[fr_real]]` — arises when + // the user aborts a long-running tool, types a follow-up text + // turn, and the React scheduler's late submitQuery then appends + // the real tool_result as a SEPARATE user entry. + // + // Forward scanning alone prevents the *synthesis* duplicate, + // but the wire layout is still + // `model[tool_use] → user[text] → user[tool_result]`, which + // Anthropic-compatible backends reject because the tool_result + // is not at the head of the IMMEDIATELY following user message. + // The repair must MOVE the real fr from history[3] into + // history[2] (before the text part) so the wire format becomes + // `model[tool_use] → user[tool_result, text]`. + chat.setHistory([ + { role: 'user', parts: [{ text: 'open /tmp/long.txt' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_nonadjacent_real', + name: 'read_file', + args: { path: '/tmp/long.txt' }, + }, + }, + ], + }, + { role: 'user', parts: [{ text: 'never mind, do something else' }] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call_nonadjacent_real', + name: 'read_file', + response: { output: 'real file contents' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + // No synthesis (the fr is real, just relocated) — `injected` + // stays empty so the React scheduler dedup doesn't see it as a + // synthesized callId. + expect(result.injected).toEqual([]); + const history = chat.getHistory(); + // History is now 3 entries: the source turn for the hoisted fr + // had only the one fr part, so it becomes empty and is removed. + expect(history.length).toBe(3); + // Real fr now at the head of the immediate next user turn, + // before the text part, satisfying the wire-format invariant. + expect(history[2]!.parts![0]!.functionResponse?.id).toBe( + 'call_nonadjacent_real', + ); + expect(history[2]!.parts![0]!.functionResponse?.response).toEqual({ + output: 'real file contents', + }); + expect(history[2]!.parts![1]).toEqual({ + text: 'never mind, do something else', + }); + }); + + it('synthesizes missing fr AND hoists real fr in a parallel tool_use mismatch', () => { + // Counterpart to the hoist case: when the real fr only covers + // SOME callIds in a parallel tool_use, and the real one is in a + // non-adjacent later user turn, BOTH fix-ups apply on the same + // model turn — synthesize the missing callId AND hoist the real + // fr from the non-adjacent location into the adjacent turn. + chat.setHistory([ + { role: 'user', parts: [{ text: 'fan out two reads' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'cid_a', name: 'read_file', args: {} }, + }, + { + functionCall: { id: 'cid_b', name: 'read_file', args: {} }, + }, + ], + }, + { role: 'user', parts: [{ text: 'follow up' }] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_a', + name: 'read_file', + response: { output: 'real for a' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + // cid_b synthesized (no real fr anywhere). cid_a is hoisted, not + // synthesized — `injected` only contains the synthetic. + expect(result.injected).toEqual([{ callId: 'cid_b', name: 'read_file' }]); + const history = chat.getHistory(); + // The non-adjacent turn that held cid_a's real fr is now empty + // and removed → 3 entries instead of the original 4. + expect(history.length).toBe(3); + // Adjacent user turn now leads with the synthesized fr_b, then + // the hoisted real fr_a, then the text. Both tool_results sit + // at the head, satisfying the Anthropic wire-format invariant. + const adjacentParts = history[2]!.parts!; + expect(adjacentParts[0]!.functionResponse?.id).toBe('cid_b'); + expect( + (adjacentParts[0]!.functionResponse?.response as { error?: string }) + ?.error, + ).toBeDefined(); + expect(adjacentParts[1]!.functionResponse?.id).toBe('cid_a'); + expect(adjacentParts[1]!.functionResponse?.response).toEqual({ + output: 'real for a', + }); + expect(adjacentParts[2]).toEqual({ text: 'follow up' }); + }); + + it('hoists real fr but preserves the source user turn when it carries other content', () => { + // Edge case for the hoist path: if the source turn for the real + // fr ALSO carries text (or any non-fr part), removing the fr + // alone must NOT delete the turn — the remaining text is the + // user's real message and must be preserved at its original + // position. Confirms the empty-turn cleanup only deletes turns + // whose parts list goes to zero after the splice. + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'cid_mix', name: 'read_file', args: {} }, + }, + ], + }, + { role: 'user', parts: [{ text: 'never mind' }] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_mix', + name: 'read_file', + response: { output: 'data' }, + }, + }, + { text: 'thanks anyway' }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([]); + const history = chat.getHistory(); + // The source turn lost its fr but kept its trailing text, so + // history is still 4 entries — the source turn survives as a + // text-only user message. + expect(history.length).toBe(4); + expect(history[2]!.parts![0]!.functionResponse?.id).toBe('cid_mix'); + expect(history[2]!.parts![1]).toEqual({ text: 'never mind' }); + expect(history[3]!.parts).toEqual([{ text: 'thanks anyway' }]); + }); + + it('drops duplicate functionResponse entries for the same callId across user turns', () => { + // Critical regression: when the same callId is echoed back more + // than once (e.g. the React scheduler retries the late submitQuery + // after the orphan repair already planted one, or two parallel + // late-submit paths land), hoisting only the first leaves the + // duplicate behind. The wire payload then serializes + // `model[tool_use] -> user[tool_result] -> user[tool_result]` + // and Anthropic-compatible backends reject the trailing block as + // an orphan, re-wedging the session. The repair MUST hoist one + // canonical fr into the adjacent turn AND delete every duplicate. + chat.setHistory([ + { role: 'user', parts: [{ text: 'open file' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'cid_dup', name: 'read_file', args: {} }, + }, + ], + }, + { role: 'user', parts: [{ text: 'never mind' }] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_dup', + name: 'read_file', + response: { output: 'data' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_dup', + name: 'read_file', + response: { output: 'data' }, + }, + }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([]); + const history = chat.getHistory(); + // 5 → 3: both source turns held only the duplicate fr, so both + // are removed; the canonical fr is hoisted into history[2] and + // sits at the head before the text part. + expect(history.length).toBe(3); + expect(history[2]!.parts![0]!.functionResponse?.id).toBe('cid_dup'); + expect(history[2]!.parts![1]).toEqual({ text: 'never mind' }); + // No fr for cid_dup remains anywhere AFTER the adjacent turn. + const trailingHasDup = history + .slice(3) + .some((entry) => + (entry.parts ?? []).some( + (part) => part.functionResponse?.id === 'cid_dup', + ), + ); + expect(trailingHasDup).toBe(false); + }); + + it('drops duplicate fr even when the canonical copy is already in the adjacent turn', () => { + // Variant of the duplicate case where the FIRST fr lands in the + // immediate next user turn (no hoist needed) but a second + // duplicate copy is in a later user turn. The hoist branch is + // skipped, but duplicate cleanup must still fire — otherwise the + // wire payload still has two `tool_result` blocks for the same id. + chat.setHistory([ + { role: 'user', parts: [{ text: 'kick off' }] }, + { + role: 'model', + parts: [ + { + functionCall: { id: 'cid_adj_dup', name: 'read_file', args: {} }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_adj_dup', + name: 'read_file', + response: { output: 'real' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'cid_adj_dup', + name: 'read_file', + response: { output: 'real' }, + }, + }, + { text: 'follow up' }, + ], + }, + ]); + + const result = chat.repairOrphanedToolUseTurns(); + + expect(result.injected).toEqual([]); + const history = chat.getHistory(); + // The source duplicate turn loses its fr but keeps its text part + // → 4 entries preserved, but the duplicate fr is gone. + expect(history.length).toBe(4); + expect(history[2]!.parts![0]!.functionResponse?.id).toBe('cid_adj_dup'); + expect(history[2]!.parts!.length).toBe(1); + expect(history[3]!.parts).toEqual([{ text: 'follow up' }]); + // The model[fc] is followed by exactly one fr for that id across + // all subsequent user turns. + const allFrIds = history + .slice(2) + .flatMap((entry) => + (entry.parts ?? []).map((p) => p.functionResponse?.id), + ) + .filter((id): id is string => Boolean(id)); + expect(allFrIds).toEqual(['cid_adj_dup']); + }); + }); + describe('output token recovery', () => { function makeChunk( parts: Array<{ text?: string; functionCall?: unknown }>, @@ -3533,6 +5142,104 @@ describe('GeminiChat', async () => { expect(lastEntry.parts!.length).toBeGreaterThan(0); }); + it('should pop both the partial model turn AND the recovery user message when recovery throws after a functionCall', async () => { + // Critical regression for the recovery catch's pop ordering. + // When the recovery stream yields a `functionCall` chunk and + // then throws, `processStreamResponse` pushes a partial `model` + // turn into history BEFORE re-throwing — so by the time the + // recovery catch runs, the trailing entries are + // [..., user(OUTPUT_RECOVERY_MESSAGE), model(partial fc)] + // The naive "if last is user, pop" check would no-op here (last + // is now `model`), leaving the OUTPUT_RECOVERY_MESSAGE control + // prompt stranded as a real user turn. The catch must pop the + // partial model turn FIRST, then the recovery user turn, and + // clear the partial-push markers so the outer `finally` JSONL + // flush doesn't resurrect the partial we just deleted. + const streams = [ + // Initial: text + MAX_TOKENS → triggers escalation. + makeStream([makeChunk([{ text: 'initial' }], 'MAX_TOKENS')]), + // Escalated: text + MAX_TOKENS → triggers recovery iteration 1. + makeStream([makeChunk([{ text: 'escalated' }], 'MAX_TOKENS')]), + // Recovery iter 1: yields functionCall chunk, then throws. + // processStreamResponse pushes a partial model turn before + // re-throwing the synthetic error. + (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_recovery_throw', + name: 'read_file', + args: { path: '/tmp/r.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw new Error('synthetic recovery mid-tool_use cut'); + })(), + ]; + let callIndex = 0; + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async () => streams[callIndex++]!, + ); + + const stream = await chat.sendMessageStream( + 'gemini-3-pro', + { message: 'recovery throws after functionCall' }, + 'prompt-recovery-fc-throw', + ); + + // Consume; the catch swallows the error and emits a synthetic + // STOP chunk so the consumer sees a clean termination. + for await (const _ of stream) { + /* consume */ + } + + const history = chat.getHistory(); + + // OUTPUT_RECOVERY_MESSAGE must NOT appear anywhere in history. + // The pop-ordering bug strands it as a real user turn that then + // pollutes durable history and biases later turns. + const flattened = JSON.stringify(history); + expect(flattened).not.toContain('Output token limit hit'); + expect(flattened).not.toContain('Resume directly'); + + // The partial model[functionCall] from the recovery throw must + // also be popped — leaving it would create a dangling tool_use + // that the inline repair on the next sendMessageStream would + // synthesize an `error` functionResponse for, and the React + // scheduler's late real result would be dropped by the + // history-based dedup. Symptom: model sees an "execution result + // was not recorded" error for a tool that actually succeeded. + const stillHasPartialFc = history.some((entry) => + (entry.parts ?? []).some( + (part) => part.functionCall?.id === 'call_recovery_throw', + ), + ); + expect(stillHasPartialFc).toBe(false); + + // Roles must strictly alternate (no consecutive same-role) so + // providers don't reject the next turn. + for (let i = 1; i < history.length; i++) { + expect(history[i]!.role).not.toBe(history[i - 1]!.role); + } + + // History tail should be the escalated model response (text: + // 'escalated'), preserved as the user-visible answer. + const lastEntry = history[history.length - 1]!; + expect(lastEntry.role).toBe('model'); + const lastModelText = (lastEntry.parts ?? []) + .map((p) => ('text' in p ? ((p as { text?: string }).text ?? '') : '')) + .join(''); + expect(lastModelText).toContain('escalated'); + }); + it('should stop recovery mid-loop when a later iteration emits functionCall', async () => { // Covers the cross-iteration guard: iter 1 returns plain text (recovery // proceeds), iter 2 returns a functionCall (recovery must break before @@ -3627,6 +5334,112 @@ describe('GeminiChat', async () => { .join(''); expect(mergedText).toBe('BCD'); }); + + it('flushes the JSONL record when escalated stream throws mid-tool_use', async () => { + // Critical regression for the max-tokens escalation path: + // 1) initial stream succeeds with text + MAX_TOKENS → triggers + // escalation, no partial set, deferred record clean. + // 2) escalated stream throws AFTER yielding a functionCall chunk + // → processStreamResponse pushes a partial model[fc] into + // `this.history` and stashes a NEW `pendingPartialAssistantRecord`. + // 3) The throw escapes through the for-await on the escalated + // stream, propagates past the (now-passed) retry loop, and + // lands in the outer `finally` block. + // + // BEFORE the fix: the flush only ran BEFORE the escalation block, + // so the new record set in step 2 was never appended to JSONL — + // live history disagreed with disk; `--resume` rehydrated a + // truncated transcript and `repairOrphanedToolUseTurnsInHistory` + // had nothing to repair, leaving the React scheduler's late real + // result as a permanent orphan. + // + // AFTER the fix: the flush is in `finally`, so the record lands + // on disk regardless of which stream raised. + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + + // Stream 1: text + MAX_TOKENS (success, triggers escalation). + // Stream 2: yields a functionCall chunk THEN throws — simulates a + // mid-tool_use stream cut on the escalated request. + const streams = [ + makeStream([makeChunk([{ text: 'partial answer' }], 'MAX_TOKENS')]), + (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_escalation_throw', + name: 'read_file', + args: { path: '/tmp/escalated.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + throw new Error('synthetic mid-tool_use cut on escalated stream'); + })(), + ]; + let callIndex = 0; + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async () => streams[callIndex++]!, + ); + + const stream = await chatWithRecording.sendMessageStream( + 'gemini-3-pro', + { message: 'kick off' }, + 'prompt-escalation-flush', + ); + + // Consume the stream and expect the synthetic mid-tool_use error + // to escape (escalation errors do not retry). + await expect( + (async () => { + for await (const _ of stream) { + /* consume */ + } + })(), + ).rejects.toThrow(/synthetic mid-tool_use cut/); + + // In-memory: the partial functionCall pushed by the escalated + // processStreamResponse must be in history. + const history = chatWithRecording.getHistory(); + const partialModel = history.findLast((h) => h.role === 'model'); + expect( + partialModel?.parts?.some( + (p) => p.functionCall?.id === 'call_escalation_throw', + ), + ).toBe(true); + + // JSONL: at least one record must mention the partial functionCall + // (the escalation throw flushed it). Without the finally-block + // flush, this assertion would fail and the durable transcript + // would silently lose a tool_use that's still live in memory. + const recordedHasPartial = recordAssistantTurn.mock.calls.some((call) => { + const message = ( + call[0] as { + message?: Array<{ functionCall?: { id?: string } }>; + } + )?.message; + return message?.some( + (p) => p.functionCall?.id === 'call_escalation_throw', + ); + }); + expect(recordedHasPartial).toBe(true); + }); }); describe('redactStructuredOutputArgsForRecording', () => { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 2655acd61cb..ce6e5106b0f 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -385,6 +385,324 @@ export class InvalidStreamError extends Error { } } +/** + * Default error text used when a synthesized `functionResponse` has to stand + * in for a real tool result that never made it back into history (e.g. the + * process crashed between the partial-tool_use push and tool completion, or + * the user hit Ctrl+Y before the in-flight tool finished and the scheduler's + * `onAllToolCallsComplete` was a single-shot that already fired into an + * `isResponding` early-return). + */ +const ORPHAN_TOOL_USE_REPAIR_REASON = + 'Tool execution result was not recorded — likely interrupted by network ' + + 'failure, abort, or process exit. Treat as failure and retry if needed.'; + +/* + * ============================================================================ + * Partial-tool_use repair subsystem — canonical design note. + * ============================================================================ + * + * Every comment block elsewhere in this file that mentions one of the + * concepts below points back here. Per-site comments should be one or two + * lines stating WHAT the local code does; the WHY lives here. + * + * --- The wedge ---------------------------------------------------------- + * + * Anthropic-compatible backends (Anthropic, DeepSeek, …) reject a request + * whose `user[tool_result]` blocks are not at the HEAD of the user message + * immediately following the `model[tool_use]` they answer: + * + * "tool_use_id ... must have a corresponding tool_use block in the + * previous message" + * + * Without a matching pair the session is unrecoverable — `stripOrphanedUser + * EntriesFromHistory` only strips trailing user entries, so a lost tool_use + * cannot be resurrected and the next send 400s repeatedly. + * + * --- The race classes that produce dangling tool_uses -------------------- + * + * Race A (Ctrl+Y mid-flight): user retries before the in-flight tool + * finishes. The scheduler's `onAllToolCallsComplete` is single-shot + * per batch and would otherwise leave the tool stuck in + * `completed-but-not-submitted` forever. + * Race B (process crash / OOM mid-flight): the JSONL transcript captures + * the dangling `model[fc]` and `--resume` rehydrates it. + * Race C (network drop between `content_block_stop` of a tool_use and + * the terminal `message_stop`): `processStreamResponse` re-throws + * after we have already yielded a `functionCall` chunk, so the React + * scheduler is on its way to submit a real `functionResponse` while + * in-memory history has no matching `model[fc]`. + * + * --- The two-layer fix --------------------------------------------------- + * + * (1) Persist the partial assistant turn at the failure point in + * `processStreamResponse` (`this.history.push({role: 'model', parts: + * [...]})` plus the `pendingPartialAssistantTurnIndex` / + * `pendingPartialAssistantRecord` markers) so the matching + * `model[fc]` is on disk and in memory when the late `user[fr]` + * arrives. + * (2) Repair any remaining dangling `model[fc]` whose + * `user[fr]` never landed (`repairOrphanedToolUseTurns`): + * - SYNTHESIZE an `error` fr for ids with no matching response; + * - HOIST the real fr into the immediately-adjacent user turn + * when it landed in a non-adjacent later turn; + * - DROP duplicate fr copies for the same id. + * Then `useGeminiStream.handleCompletedTools` dedupes the + * scheduler's late real result against `chat.history` so the + * synthetic and the real result never collide on the wire. + * + * --- Partial-push marker lifecycle --------------------------------------- + * + * Set together on (streamError + hasToolCall + hasContent) inside + * `processStreamResponse`. Cleared together by `popPartialIfPushed` on a + * retryable error rollback, or flushed together to JSONL by the outer + * `finally` after the retry loop exits. Defense-in-depth: every + * history-mutation method (clearHistory / addHistory / setHistory / + * truncateHistory / stripThoughtsFromHistory / + * stripOrphanedUserEntriesFromHistory) resets both markers in lockstep so + * a stale index can't shift onto an unrelated model turn and cause + * `popPartialIfPushed` to splice the wrong entry. Any single-field reset + * is a bug. + * ============================================================================ + */ + +/** + * Walk `history` left-to-right and close every dangling + * tool_use ↔ tool_result pair. For each `model[functionCall]`: + * - SYNTHESIZE an `error` `functionResponse` for ids with no match; + * - HOIST a real fr from a non-adjacent later user turn into the + * adjacent one; + * - drop duplicate fr copies for the same id. + * + * Mutates `history` in place. Returns the synthesized (callId, name) + * pairs so the React scheduler's dedup can drop late real results for + * those ids; hoisted ids are NOT returned (the real fr is still in + * history, scheduler dedup handles them naturally). See the canonical + * note above `ORPHAN_TOOL_USE_REPAIR_REASON`. qwen-code analogue of + * upstream Claude Code's `yieldMissingToolResultBlocks`. + */ +/** Location of a `functionResponse` part within `history`. */ +interface FrLocation { + turnIdx: number; + partIdx: number; + part: Part; +} + +/** + * Output of the scan phase for a single `model[functionCall]` turn at + * `modelIdx`. `expected` maps each `functionCall.id` to its tool name, + * `matched` maps that same id to ALL locations of matching + * `functionResponse` parts across the consecutive user turns that + * follow, and `scanEnd` is one past the last user turn visited. + */ +interface ScanResult { + modelIdx: number; + expected: Map; + matched: Map; + scanEnd: number; +} + +/** Decision-phase output: exact mutations the next phase will apply. */ +interface RepairPlan { + modelIdx: number; + scanEnd: number; + synthesizeIds: Array<[string, string]>; + hoistedParts: Part[]; + removalTargets: Array<{ turnIdx: number; partIdx: number }>; + droppedDuplicates: Array<{ callId: string; name: string }>; +} + +/** + * SCAN — collect every `functionCall.id → name` from the model turn at + * `modelIdx` and EVERY `functionResponse.id → location` from the + * consecutive user turns that follow. Pure read. Storing all locations + * (not just the first) is what lets the decision phase drop duplicates. + */ +function scanModelTurn(history: Content[], modelIdx: number): ScanResult { + const expected = new Map(); + for (const part of history[modelIdx]?.parts ?? []) { + const fc = part.functionCall; + if (fc?.id) expected.set(fc.id, fc.name ?? 'unknown'); + } + + const matched = new Map(); + let scanIdx = modelIdx + 1; + while (scanIdx < history.length && history[scanIdx]?.role === 'user') { + const parts = history[scanIdx].parts ?? []; + for (let pIdx = 0; pIdx < parts.length; pIdx++) { + const part = parts[pIdx]; + const id = part.functionResponse?.id; + if (id) { + const list = matched.get(id); + if (list) list.push({ turnIdx: scanIdx, partIdx: pIdx, part }); + else matched.set(id, [{ turnIdx: scanIdx, partIdx: pIdx, part }]); + } + } + scanIdx++; + } + + return { modelIdx, expected, matched, scanEnd: scanIdx }; +} + +/** + * DECISION — classify each expected id: no match → SYNTHESIZE; first + * match adjacent → SKIP relocation; first match non-adjacent → HOIST. + * Every duplicate beyond the first is always dropped. Pure compute. + */ +function planRepair(scan: ScanResult): RepairPlan { + const synthesizeIds: Array<[string, string]> = []; + const hoistedParts: Part[] = []; + const removalTargets: Array<{ turnIdx: number; partIdx: number }> = []; + const droppedDuplicates: Array<{ callId: string; name: string }> = []; + + const adjacentIdx = scan.modelIdx + 1; + for (const [id, name] of scan.expected) { + const locations = scan.matched.get(id); + if (!locations || locations.length === 0) { + synthesizeIds.push([id, name]); + continue; + } + // First copy is the canonical survivor — payloads should be + // identical for the same callId; if they differ, the wire is + // already corrupt and the backend rejects regardless. + const survivor = locations[0]!; + if (survivor.turnIdx !== adjacentIdx) { + hoistedParts.push(survivor.part); + removalTargets.push({ + turnIdx: survivor.turnIdx, + partIdx: survivor.partIdx, + }); + } + for (let k = 1; k < locations.length; k++) { + removalTargets.push({ + turnIdx: locations[k]!.turnIdx, + partIdx: locations[k]!.partIdx, + }); + droppedDuplicates.push({ callId: id, name }); + } + } + + return { + modelIdx: scan.modelIdx, + scanEnd: scan.scanEnd, + synthesizeIds, + hoistedParts, + removalTargets, + droppedDuplicates, + }; +} + +/** + * MUTATION — apply the plan to `history` in place. Returns the count + * of new user turns inserted ahead of `modelIdx + 1` (0 or 1) so the + * outer loop can advance its cursor. + * + * Order: (1) splice removal targets desc-by-desc, (2) drop empty user + * turns in `[modelIdx + 2, scanEnd)`, (3) HEAD-insert at the adjacent + * user turn OR splice a new user turn between. The HEAD insert is + * load-bearing (mirrors upstream `hoistToolResults`) — see the + * canonical note for why tail-append re-triggers the wedge. + */ +function applyRepair( + history: Content[], + plan: RepairPlan, + reason: string, +): { insertedBefore: number } { + if (plan.synthesizeIds.length === 0 && plan.removalTargets.length === 0) { + return { insertedBefore: 0 }; + } + + const syntheticParts: Part[] = plan.synthesizeIds.map(([callId, name]) => ({ + functionResponse: { id: callId, name, response: { error: reason } }, + })); + const partsToInject: Part[] = [...syntheticParts, ...plan.hoistedParts]; + + // (1) Splice removal targets, descending so indices stay valid. + const removals = [...plan.removalTargets].sort((a, b) => { + if (a.turnIdx !== b.turnIdx) return b.turnIdx - a.turnIdx; + return b.partIdx - a.partIdx; + }); + for (const loc of removals) { + const turnParts = history[loc.turnIdx].parts; + if (turnParts) turnParts.splice(loc.partIdx, 1); + } + + // (2) Drop now-empty user turns within [modelIdx + 2, scanEnd). + // Preserve the adjacent turn even if empty — we'll rewrite it + // below. + const adjacentIdx = plan.modelIdx + 1; + for (let j = plan.scanEnd - 1; j > adjacentIdx; j--) { + if (history[j]?.role === 'user' && (history[j].parts?.length ?? 0) === 0) { + history.splice(j, 1); + } + } + + // (3) Place new parts at the head of the adjacent user turn, OR + // insert a fresh user turn between this model turn and whatever + // follows. + const next = history[adjacentIdx]; + if (next?.role === 'user') { + const existing = next.parts ?? []; + const firstNonFr = existing.findIndex((part) => !part.functionResponse); + const insertAt = firstNonFr === -1 ? existing.length : firstNonFr; + next.parts = [ + ...existing.slice(0, insertAt), + ...partsToInject, + ...existing.slice(insertAt), + ]; + return { insertedBefore: 0 }; + } + history.splice(adjacentIdx, 0, { role: 'user', parts: partsToInject }); + return { insertedBefore: 1 }; +} + +/** + * Forward-walk `history`, planning and applying the repair for each + * `model[functionCall]` turn in turn. Iteration is index-based and the + * cursor advances by the count of user turns inserted ahead of it so + * a freshly-injected turn isn't re-visited. + * + * Splitting scan / decision / mutation into separate functions keeps + * each phase auditable in isolation — index drift can only happen in + * `applyRepair`, the only function that mutates `history`. + */ +export function repairOrphanedToolUseTurns( + history: Content[], + reason: string = ORPHAN_TOOL_USE_REPAIR_REASON, +): { + injected: Array<{ callId: string; name: string }>; + droppedDuplicates: Array<{ callId: string; name: string }>; +} { + const injected: Array<{ callId: string; name: string }> = []; + const droppedDuplicates: Array<{ callId: string; name: string }> = []; + + for (let i = 0; i < history.length; i++) { + if (history[i].role !== 'model') continue; + + const scan = scanModelTurn(history, i); + if (scan.expected.size === 0) continue; + + const plan = planRepair(scan); + if (plan.synthesizeIds.length === 0 && plan.removalTargets.length === 0) { + continue; + } + + const { insertedBefore } = applyRepair(history, plan, reason); + // Only synthesized ids feed `injected` — hoisted ids reference real + // frs that were ALREADY in history before this pass (just + // relocated), so the scheduler's dedup naturally handles them. + for (const [callId, name] of plan.synthesizeIds) { + injected.push({ callId, name }); + } + droppedDuplicates.push(...plan.droppedDuplicates); + // Advance past any freshly-inserted user turn so the outer loop + // doesn't revisit it. Keeps the walk linear-time. + i += insertedBefore; + } + + return { injected, droppedDuplicates }; +} + /** * Chat session that enables sending messages to the model with previous * conversation context. @@ -443,6 +761,27 @@ export class GeminiChat { */ private hasFailedCompressionAttempt = false; + /** + * Partial-push markers — index of the in-memory `model[partial fc]` + * and the matching deferred JSONL record. See the canonical note + * above `ORPHAN_TOOL_USE_REPAIR_REASON` for the lifecycle and the + * wedge they prevent. + */ + private pendingPartialAssistantTurnIndex: number | null = null; + private pendingPartialAssistantRecord: + | Parameters[0] + | null = null; + + /** + * Reset both partial-push markers in lockstep. Every history-mutation + * site uses this — single-field resets are a bug because the fields + * are always paired by lifecycle. + */ + private clearPendingPartialState(): void { + this.pendingPartialAssistantTurnIndex = null; + this.pendingPartialAssistantRecord = null; + } + /** * Creates a new GeminiChat instance. * @@ -617,6 +956,16 @@ export class GeminiChat { }); this.sendPromise = streamDonePromise; + // Clear any partial-push marker left over from a prior unretryable + // break path — the marker is per-send; carrying it across sends + // would let the next send's retry catch wrongly pop a now-valid + // model entry sitting at the stale index. The deferred-record + // stash gets the same per-send reset for the same reason: a + // leftover from a prior unretryable break would otherwise get + // appended to JSONL by THIS send's retry-loop flush, attaching + // someone else's failed turn to this conversation. + this.clearPendingPartialState(); + let compressionInfo: ChatCompressionInfo; let requestContents: Content[]; let userContentAdded = false; @@ -637,6 +986,33 @@ export class GeminiChat { // Add user content to history ONCE before any attempts. this.history.push(userContent); userContentAdded = true; + // Per-send orphan repair (belt-and-suspenders alongside the + // startChat load-time pass). Runs AFTER user content lands so a + // user-supplied tool_result closes the pair before we synthesize + // anything. Logs are tagged so investigators can distinguish this + // pass from the session-load pass and from the React scheduler's + // dedup-drop. See the canonical note above + // `ORPHAN_TOOL_USE_REPAIR_REASON`. + const inlineRepair = repairOrphanedToolUseTurns(this.history); + if (inlineRepair.injected.length > 0) { + debugLogger.warn( + `[REPAIR] sendMessageStream inline pass synthesized ` + + `${inlineRepair.injected.length} functionResponse(s): ` + + inlineRepair.injected + .map((entry) => `${entry.name}(${entry.callId})`) + .join(', '), + ); + } + if (inlineRepair.droppedDuplicates.length > 0) { + debugLogger.warn( + `[REPAIR] sendMessageStream inline pass dropped ` + + `${inlineRepair.droppedDuplicates.length} duplicate ` + + `functionResponse(s): ` + + inlineRepair.droppedDuplicates + .map((entry) => `${entry.name}(${entry.callId})`) + .join(', '), + ); + } requestContents = this.getRequestHistory(); } catch (error) { if (userContentAdded) { @@ -722,12 +1098,61 @@ export class GeminiChat { } catch (error) { lastError = error; + // If `processStreamResponse` persisted a partial assistant turn + // (mid-stream error after a `functionCall` was already + // yielded), every retry-and-continue path below must drop + // that turn first; otherwise the retry's response lands as + // a second consecutive model turn with an orphan tool_use + // (the wedge — see the canonical note above + // `ORPHAN_TOOL_USE_REPAIR_REASON`). Paths that `break` + // (unretryable) keep the partial. + const popPartialIfPushed = () => { + const idx = self.pendingPartialAssistantTurnIndex; + if (idx === null) return; + if ( + self.history.length > idx && + self.history[idx]?.role === 'model' + ) { + self.history.splice(idx, 1); + } else { + // Marker was set but the entry it pointed at is gone or + // is no longer a `model` turn. Today this can't happen: + // every history-mutation path (clearHistory, addHistory, + // setHistory, truncateHistory, stripThoughtsFromHistory, + // stripOrphanedUserEntriesFromHistory) calls + // clearPendingPartialState() in lockstep, so the marker + // is null whenever the index basis is invalidated. + // Logging the mismatch makes the invariant observable — + // without this, a future caller that mutates history + // without resetting the marker would silently leave a + // stale partial in `this.history` (popPartialIfPushed + // skipping the splice) AND the field-level invariant + // that "marker non-null ⇒ a real partial sits at idx" + // would be quietly violated. With the warn, anyone + // investigating a stale-partial wedge sees a log line + // pointing straight at the offending caller. + debugLogger.warn( + `[PARTIAL_POP] Splice skipped: idx=${idx}, ` + + `historyLength=${self.history.length}, ` + + `roleAtIdx=${self.history[idx]?.role ?? 'undefined'}`, + ); + } + // Drop both markers in lockstep — the deferred chat- + // recording record must be discarded alongside the + // in-memory splice so the JSONL transcript also drops the + // failed attempt. See the field-level comment on + // `pendingPartialAssistantRecord` for the failure mode + // this prevents. + self.clearPendingPartialState(); + }; + // Handle rate-limit / throttling errors returned as stream content. // These arrive as StreamContentError with finish_reason="error_finish" // from the pipeline, containing the throttling message in the content. // Covers TPM throttling, GLM rate limits, and other provider throttling. const isRateLimit = isRateLimitError(error, extraRetryErrorCodes); if (isRateLimit && rateLimitRetryCount < maxRateLimitRetries) { + popPartialIfPushed(); rateLimitRetryCount++; const delayMs = getRateLimitRetryDelayMs(rateLimitRetryCount, { ...RATE_LIMIT_RETRY_OPTIONS, @@ -802,6 +1227,11 @@ export class GeminiChat { reactiveInfo.compressionStatus === CompressionStatus.COMPRESSED ) { + // No-op today: tryCompress's setHistory has already + // cleared the marker. Kept for uniformity with the + // other retry branches in case a future in-place + // tryCompress stops resetting it. + popPartialIfPushed(); requestContents = self.getRequestHistory(); debugLogger.info( `Reactive compression succeeded: ` + @@ -858,6 +1288,7 @@ export class GeminiChat { isTransientStreamError && invalidStreamRetryCount < INVALID_STREAM_RETRY_CONFIG.maxRetries ) { + popPartialIfPushed(); invalidStreamRetryCount++; const delayMs = INVALID_STREAM_RETRY_CONFIG.initialDelayMs * @@ -887,10 +1318,20 @@ export class GeminiChat { break; } - // Other content validation errors (e.g. NO_FINISH_REASON). + // Currently unreachable for `InvalidStreamError`. The + // `isContentError` predicate is identical to + // `isTransientStreamError` (`error instanceof InvalidStreamError`), + // and the transient branch above already either continued or + // broke for that class. The branch is preserved as + // defense-in-depth: a future error class that should consume + // its own content-retry budget but NOT the transient one + // could be threaded through here without re-deriving the + // popPartialIfPushed sequence. No reachable test path until + // the predicates diverge. const isContentError = error instanceof InvalidStreamError; if (isContentError) { if (attempt < INVALID_CONTENT_RETRY_OPTIONS.maxAttempts - 1) { + popPartialIfPushed(); logContentRetry( self.config, new ContentRetryEvent( @@ -1025,9 +1466,33 @@ export class GeminiChat { // coalesced back into the preceding model entry after the loop. successfulRecoveries++; } catch (recoveryError) { - // If a recovery attempt fails (e.g., empty response, network - // error), stop recovering and let the partial output stand. - // Pop the dangling recovery message to keep history valid. + // Pop the partial `model[fc]` FIRST (if processStreamResponse + // pushed one before re-throwing), THEN the recovery user + // turn. Reversed order would strand `OUTPUT_RECOVERY_MESSAGE` + // as a real user turn. Index-checked pop mirrors + // `popPartialIfPushed` above — see the design note above + // `ORPHAN_TOOL_USE_REPAIR_REASON` for the wedge mechanism + // and the partial-push marker lifecycle. + const expectedIdx = self.pendingPartialAssistantTurnIndex; + const lastIdx = self.history.length - 1; + if ( + expectedIdx !== null && + self.history.length > 0 && + self.history[lastIdx]?.role === 'model' + ) { + if (expectedIdx !== lastIdx) { + debugLogger.warn( + `[RECOVERY_POP] Marker/last-index mismatch: ` + + `marker=${expectedIdx}, lastIdx=${lastIdx}, ` + + `historyLength=${self.history.length}. Popping ` + + `last entry as best-effort rollback — investigate ` + + `any history mutation between processStreamResponse's ` + + `partial push and this catch.`, + ); + } + self.history.pop(); + self.clearPendingPartialState(); + } if ( self.history.length > 0 && self.history[self.history.length - 1].role === 'user' @@ -1083,6 +1548,28 @@ export class GeminiChat { } } finally { streamDoneResolver!(); + // Flush any deferred partial-tool_use record. Covers both the + // post-retry-loop unretryable break AND the max-tokens + // escalation throw (the escalated processStreamResponse can + // set a new record that escapes the retry-loop catch). + // Recording-service errors are logged at error level (sustained + // failure = monitoring signal) and swallowed — propagating + // would mask the real send outcome. + if (self.pendingPartialAssistantRecord) { + try { + self.chatRecordingService?.recordAssistantTurn( + self.pendingPartialAssistantRecord, + ); + } catch (recordErr) { + debugLogger.error( + '[PARTIAL_FLUSH] Failed to persist deferred JSONL record: ' + + (recordErr instanceof Error + ? recordErr.message + : String(recordErr)), + ); + } + self.clearPendingPartialState(); + } } })(); } @@ -1243,11 +1730,38 @@ export class GeminiChat { return this.history.length; } + /** + * Set of `functionResponse.id` strings in user turns. Walk-only, + * no clone — `useGeminiStream.handleCompletedTools` calls this per + * tool-completion batch, so {@link getHistory}'s `structuredClone` + * would stall the UI on long sessions. + */ + getHistoryFunctionResponseIds(): Set { + const ids = new Set(); + for (const entry of this.history) { + if (entry.role !== 'user') continue; + for (const part of entry.parts ?? []) { + const id = part.functionResponse?.id; + if (id) ids.add(id); + } + } + return ids; + } + /** * Clears the chat history. */ clearHistory(): void { this.history = []; + // Any pending partial-push state points into the now-empty history; + // resetting prevents `popPartialIfPushed` from splicing whatever + // shows up at that index in a future send (defense-in-depth — the + // helper also bounds-checks, but a stale marker that happens to + // line up with a real model turn could otherwise pop the wrong + // entry). The deferred-record stash is dropped for the same reason: + // a later flush would append a turn that doesn't match the (now- + // empty) live history. + this.clearPendingPartialState(); } /** @@ -1255,20 +1769,56 @@ export class GeminiChat { */ addHistory(content: Content): void { this.history.push(content); + // addHistory only runs between sends, so the partial-push marker + // should already be cleared. If it is not, a new caller is + // violating that invariant — surface it at error level so the + // offending stack is visible. See the design note above + // `ORPHAN_TOOL_USE_REPAIR_REASON` for the marker lifecycle. + if ( + this.pendingPartialAssistantTurnIndex !== null || + this.pendingPartialAssistantRecord !== null + ) { + debugLogger.error( + '[INVARIANT_VIOLATION] addHistory called while a partial-push ' + + 'marker is active — clearing it.', + ); + } + this.clearPendingPartialState(); } setHistory(history: Content[]): void { this.history = history; + // History replacement (compression, /clear, --resume reload) wipes + // the index basis the partial-push marker was captured against. The + // marker MUST be cleared — otherwise `popPartialIfPushed` could find + // a model turn at the stale index in the replacement history and + // splice an entry that has nothing to do with the original partial + // push, corrupting the conversation. Drop the paired deferred-record + // stash too: its referent (the model turn at the old index) is gone. + this.clearPendingPartialState(); } truncateHistory(keepCount: number): void { this.history = this.history.slice(0, keepCount); + // Truncation can drop the entry the partial-push marker points at, + // or leave it valid but shift the meaning of nearby indices. Reset + // both fields rather than try to fix them up — they're per-send and + // ephemeral, so losing them across a truncate is safe (the + // sendMessageStream that pushed them has already finished or will + // start fresh on the next call). + this.clearPendingPartialState(); } stripThoughtsFromHistory(): void { this.history = this.history .map(stripThoughtPartsFromContent) .filter((content): content is Content => content !== null); + // Filter+map replaces `this.history` with a new array, so any pending + // partial-push marker is now indexed against an array that no longer + // exists. Clear it for the same reason setHistory does — and drop + // the paired deferred-record stash so a later flush can't land a + // turn that doesn't exist in live history. + this.clearPendingPartialState(); } /** @@ -1283,6 +1833,29 @@ export class GeminiChat { ) { this.history.pop(); } + // Today this is safe even without the reset — only trailing user + // entries are popped, which can't shift the index of an earlier + // `model` partial. But every other history-mutation method now + // clears the partial-push state in lockstep + // (clearHistory/addHistory/setHistory/truncateHistory/ + // stripThoughtsFromHistory), so omitting it here would be a silent + // exception to the uniform invariant: a future caller invoking + // this method between the deferred JSONL flush and the next + // `sendMessageStream` would otherwise leave a stale marker that + // happens to line up with whatever model entry is at that index + // in the meanwhile. + this.clearPendingPartialState(); + } + + /** + * Instance wrapper around the free-function {@link repairOrphanedToolUseTurns}. + * See the canonical note above `ORPHAN_TOOL_USE_REPAIR_REASON`. + */ + repairOrphanedToolUseTurns(reason?: string): { + injected: Array<{ callId: string; name: string }>; + droppedDuplicates: Array<{ callId: string; name: string }>; + } { + return repairOrphanedToolUseTurns(this.history, reason); } setTools(tools: Tool[]): void { @@ -1334,48 +1907,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; @@ -1421,11 +2008,25 @@ export class GeminiChat { .join('') .trim(); - // Record assistant turn with raw Content and metadata - if (thoughtContentPart || contentText || hasToolCall || usageMetadata) { + // Record assistant turn with raw Content and metadata. Gate matches + // the in-memory `this.history.push` decision below so chat-recording + // JSONL never carries a partial turn we deliberately dropped from + // history: on `--resume` the transcript-load path would otherwise + // re-inject a model turn the in-session run intentionally discarded + // (text-only mid-stream errors, where the Retry re-issues the user + // prompt — a stale partial-text record would bias the resumed + // conversation or surface as duplicate output). + const willPersistToHistory = + streamError === null || + (hasToolCall && + (thoughtContentPart || consolidatedHistoryParts.length > 0)); + if ( + willPersistToHistory && + (thoughtContentPart || contentText || hasToolCall || usageMetadata) + ) { const contextWindowSize = this.config.getContentGeneratorConfig()?.contextWindowSize; - this.chatRecordingService?.recordAssistantTurn({ + const recordArgs = { model, message: [ ...(thoughtContentPart ? [thoughtContentPart] : []), @@ -1443,7 +2044,84 @@ export class GeminiChat { ], tokens: usageMetadata, contextWindowSize, - }); + }; + if (streamError !== null) { + // Stream-error + tool-use partial: defer the JSONL append until + // the outer retry loop decides whether to roll back this attempt. + // If the same send retries successfully, popPartialIfPushed clears + // this stash and the failed attempt never lands on disk; if the + // retry path doesn't apply (unretryable break), the stash is + // flushed at the rethrow site so JSONL stays aligned with the + // partial that survives in-memory. Without this, retry-success + // leaves a failed `model[functionCall]` durable in JSONL and + // `--resume` rehydrates a turn the live session correctly + // discarded. + this.pendingPartialAssistantRecord = recordArgs; + } else { + this.chatRecordingService?.recordAssistantTurn(recordArgs); + } + } + + // Mid-stream failure recovery (Race C in the canonical note above + // `ORPHAN_TOOL_USE_REPAIR_REASON`): if the upstream stream threw + // AFTER a `functionCall` chunk was already yielded — typical on + // weak networks: SSE cut between a tool_use `content_block_stop` + // and the terminal `message_stop` — we persist the partial + // assistant turn so the React scheduler's incoming + // `user[functionResponse]` has a matching `model[tool_use]` to + // pair with. + // + // 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) { + // Reuse the `willPersistToHistory` gate from the recordAssistantTurn + // block above instead of re-deriving it. When `streamError !== null`, + // `willPersistToHistory` reduces to exactly the original expression + // `hasToolCall && (thoughtContentPart || consolidatedHistoryParts.length > 0)`; + // sharing the single binding eliminates drift risk if one gate is + // tightened without the other and the JSONL recording silently + // desyncs from in-memory history. + if (willPersistToHistory) { + this.history.push({ + role: 'model', + parts: [ + ...(thoughtContentPart ? [thoughtContentPart] : []), + ...consolidatedHistoryParts, + ], + }); + // Track the pushed turn so the outer sendMessageStream retry loop + // can roll it back if it decides to retry the same send. Without + // this, a successful retry would leave the failed attempt's + // partial `model[functionCall]` as a stale leading model turn in + // front of the retry's real response. + this.pendingPartialAssistantTurnIndex = this.history.length - 1; + // Trace the push event so the lifecycle is observable end-to-end: + // dedup in `useGeminiStream.handleCompletedTools` already logs + // `[REPAIR] Dropping ...`, and `repairOrphanedToolUseTurnsInHistory` + // logs `[REPAIR] Synthesized ...`. Without a corresponding + // `[PARTIAL_PUSH]` line here, an investigator looking at a + // stale-partial wedge sees the downstream symptom but has no + // anchor for when/why the partial originated. + debugLogger.warn( + '[PARTIAL_PUSH] Persisting partial assistant turn for ' + + 'mid-stream error recovery (will be rolled back if retry ' + + 'succeeds, kept if break is unretryable). ' + + `pendingIndex=${this.pendingPartialAssistantTurnIndex} ` + + `callIds=${consolidatedHistoryParts + .map((p) => p.functionCall?.id) + .filter((id): id is string => Boolean(id)) + .join(',')} ` + + `error=${ + streamError instanceof Error + ? streamError.message + : String(streamError) + }`, + ); + } + throw streamError; } // Stream validation logic: A stream is considered successful if: