Skip to content

fix(cli): submit fast tool results after stream end #5071

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wenshao merged 1 commit into from
Jun 13, 2026
+259 −9
Merged

fix(cli): submit fast tool results after stream end #5071

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
240 changes: 236 additions & 4 deletions packages/cli/src/ui/hooks/useGeminiStream.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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',
Expand All @@ -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<void>)
| null = null;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});

let releaseStream!: () => void;
const holdStream = new Promise<void>((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<unknown> | 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<void>)
| 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<void>)
| null = null;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});

let releaseStream!: () => void;
const holdStream = new Promise<void>((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<unknown> | 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`
Expand Down
28 changes: 23 additions & 5 deletions packages/cli/src/ui/hooks/useGeminiStream.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,8 @@ export const useGeminiStream = (
const lastPromptErroredRef = useRef(false);
const dualOutput = useDualOutput();
const [isResponding, setIsResponding] = useState<boolean>(false);
// React state can lag by one render; this tracks the actual stream lifetime.
const activeModelStreamsRef = useRef(0);
const [thought, setThought] = useState<ThoughtSummary | null>(null);
// Hold the latest history in a ref so handleCompletedTools can read it
// without depending on `history` (which would recreate the tool scheduler
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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;
}
});
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -2182,7 +2191,7 @@ export const useGeminiStream = (
markToolsAsSubmitted(dedupedCallIds);
}

if (isResponding) {
if (activeModelStreamsRef.current > 0) {
return;
}

Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -2409,7 +2428,6 @@ export const useGeminiStream = (
submitQuery(responsesToSend, SendMessageType.ToolResult, prompt_ids[0]);
},
[

@wenshao wenshao Jun 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Critical] Missing cancellation guard: removing isResponding from the dependency array eliminates the stale-closure that accidentally protected against post-cancellation tool submissions. Now activeModelStreamsRef.current always reads 0 after stream abort, so handleCompletedTools proceeds to submitQuery at line 2418 — sending tool results (potentially containing file contents or shell output) to a new API call the user explicitly tried to cancel via Ctrl+C.

Scenario: Ctrl+C → stream aborts → counter→0 in finally → tool completes → activeModelStreamsRef.current > 0 guard passes → submitQuery fires unwanted API call. The submitQuery function resets turnCancelledRef.current = false, so the cancel state is lost.

Before this PR, the stale isResponding === true closure returned early. Now the ref is always current, making this path deterministic.

Fix: Add a cancellation check before submitQuery at line 2418:

if (turnCancelledRef.current || abortControllerRef.current?.signal.aborted) {
  markToolsAsSubmitted(callIdsToMarkAsSubmitted);
  return;
}

— qwen3.7-max via Qwen Code /review

isResponding,
submitQuery,
markToolsAsSubmitted,
geminiClient,
Expand Down
Loading