diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 43a762bbd92..ab98e093ad6 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -1323,9 +1323,9 @@ describe('useGeminiStream', () => { expect(client.recordCompletedToolCall).not.toHaveBeenCalled(); }); - it('runs Race A dedup BEFORE the isResponding early-return (regression guard)', async () => { + it('runs Race A dedup BEFORE the active-stream early-return (regression guard)', async () => { // The dedup block in handleCompletedTools is intentionally placed - // ABOVE the `if (isResponding) return;` early-return: the scheduler's + // ABOVE the active-stream 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 @@ -1472,8 +1472,8 @@ describe('useGeminiStream', () => { }); // The dedup MUST still fire — markToolsAsSubmitted called with the - // deduped callId — even though the early-return on isResponding - // would otherwise skip every later branch. + // deduped callId — even though the active-stream guard would + // otherwise skip every later branch. await waitFor(() => { expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ 'call_race_A_responding', @@ -1488,6 +1488,238 @@ describe('useGeminiStream', () => { releaseStream(); }); + it('submits a fast tool result after the stream ended but before React replaces the callback', async () => { + const responseParts: Part[] = [ + { + functionResponse: { + id: 'call_fast_after_stream', + name: 'read_file', + response: { error: 'ENOENT: missing file' }, + }, + }, + ]; + const fastFailedTool = { + request: { + callId: 'call_fast_after_stream', + name: 'read_file', + args: { path: '/tmp/missing.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-fast-after-stream', + }, + status: 'error', + responseSubmittedToGemini: false, + response: { + callId: 'call_fast_after_stream', + responseParts, + resultDisplay: undefined, + error: new Error('ENOENT: missing file'), + errorType: ToolErrorType.UNHANDLED_EXCEPTION, + }, + tool: { + name: 'read_file', + displayName: 'ReadFile', + description: 'Read a file', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => 'read /tmp/missing.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + // 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, + ), + ); + + let submitPromise: Promise | undefined; + act(() => { + submitPromise = result.current.submitQuery('edit the missing file'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + }); + + // Save the callback from the render where React state still says + // "responding". The scheduler can call this stale closure if a tool + // finishes immediately after the stream returns. + const staleOnComplete = capturedOnComplete; + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + + releaseStream(); + await act(async () => { + await submitPromise; + }); + + const staleCompletedOnComplete = staleOnComplete as + | ((completedTools: TrackedCompletedToolCall[]) => Promise) + | null; + await act(async () => { + await staleCompletedOnComplete?.([fastFailedTool]); + }); + + await waitFor(() => { + expect(mockSendMessageStream).toHaveBeenCalledTimes(2); + }); + expect(mockSendMessageStream).toHaveBeenNthCalledWith( + 2, + responseParts, + expect.any(AbortSignal), + 'prompt-fast-after-stream', + expect.objectContaining({ type: SendMessageType.ToolResult }), + ); + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ + 'call_fast_after_stream', + ]); + }); + + it('drops a fast tool result after cancellation even if the stale callback runs later', async () => { + const responseParts: Part[] = [ + { + functionResponse: { + id: 'call_fast_after_cancel', + name: 'read_file', + response: { output: 'secret file contents' }, + }, + }, + ]; + const fastToolAfterCancel = { + request: { + callId: 'call_fast_after_cancel', + name: 'read_file', + args: { path: '/tmp/secret.txt' }, + isClientInitiated: false, + prompt_id: 'prompt-fast-after-cancel', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'call_fast_after_cancel', + responseParts, + 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/secret.txt', + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall; + + const client = new MockedGeminiClientClass(mockConfig); + let capturedOnComplete: + | ((completedTools: TrackedCompletedToolCall[]) => Promise) + | null = null; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + // 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, + ), + ); + + let submitPromise: Promise | undefined; + act(() => { + submitPromise = result.current.submitQuery('read the file'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + }); + + const staleOnComplete = capturedOnComplete; + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + + act(() => { + result.current.cancelOngoingRequest(); + }); + releaseStream(); + await act(async () => { + await submitPromise; + }); + + await act(async () => { + await staleOnComplete?.([fastToolAfterCancel]); + }); + + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith([ + 'call_fast_after_cancel', + ]); + }); + 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` diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index bdbbe387140..218aaef7b3f 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -347,6 +347,8 @@ export const useGeminiStream = ( const lastPromptErroredRef = useRef(false); const dualOutput = useDualOutput(); const [isResponding, setIsResponding] = useState(false); + // React state can lag by one render; this tracks the actual stream lifetime. + const activeModelStreamsRef = useRef(0); const [thought, setThought] = useState(null); // Hold the latest history in a ref so handleCompletedTools can read it // without depending on `history` (which would recreate the tool scheduler @@ -1855,6 +1857,7 @@ export const useGeminiStream = ( logUserRetry(config, new UserRetryEvent(prompt_id)); } + activeModelStreamsRef.current += 1; setIsResponding(true); setInitError(null); // Entering "requesting" phase — no content yet for this API call. @@ -1969,7 +1972,13 @@ export const useGeminiStream = ( } } finally { submitPromptOnCompleteRef.current = null; - setIsResponding(false); + activeModelStreamsRef.current = Math.max( + 0, + activeModelStreamsRef.current - 1, + ); + if (activeModelStreamsRef.current === 0) { + setIsResponding(false); + } isSubmittingQueryRef.current = false; } }); @@ -2114,14 +2123,14 @@ export const useGeminiStream = ( }, ); - // History-based dedup MUST run before the `isResponding` early-return. + // History-based dedup MUST run before the active-stream 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 + // per batch, so a later active-stream 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 @@ -2182,7 +2191,7 @@ export const useGeminiStream = ( markToolsAsSubmitted(dedupedCallIds); } - if (isResponding) { + if (activeModelStreamsRef.current > 0) { return; } @@ -2234,6 +2243,16 @@ export const useGeminiStream = ( return; } + if ( + turnCancelledRef.current || + abortControllerRef.current?.signal.aborted + ) { + markToolsAsSubmitted( + geminiTools.map((toolCall) => toolCall.request.callId), + ); + return; + } + // If all the tools were cancelled, don't submit a response to Gemini. const allToolsCancelled = geminiTools.every( (tc) => tc.status === 'cancelled', @@ -2409,7 +2428,6 @@ export const useGeminiStream = ( submitQuery(responsesToSend, SendMessageType.ToolResult, prompt_ids[0]); }, [ - isResponding, submitQuery, markToolsAsSubmitted, geminiClient,