From 064265ec4fe0bc40875f257031dc4e55a49d51c5 Mon Sep 17 00:00:00 2001 From: KittiphonKamnuan Date: Mon, 22 Jun 2026 23:23:38 +0700 Subject: [PATCH 1/2] fix(core): drop late tool calls after SIGINT cancellation Closes #28091. A SIGINT delivered after the stream consumer has already started a turn could still result in a delayed provider tool-call chunk being executed locally: the side-effect ran, and the result was submitted back to the model after the user had cancelled. The root cause is that the abort signal was only checked at the top of each `for await` iteration on the model stream. A chunk that arrived after cancellation but in the same iteration as the abort would still materialize tool calls, push them onto the scheduler, and run. Defense-in-depth fix at three layers: 1. `Turn.run` (`packages/core/src/core/turn.ts`): re-check `signal.aborted` between processing text/thought parts and yielding `ToolCallRequest` events. If the user cancelled mid-chunk, yield `UserCancelled` and stop instead of forwarding the tool call. 2. Non-interactive CLI driver (`packages/cli/src/nonInteractiveCli.ts`): re-check `signal.aborted` after the stream loop exits and before calling `scheduler.schedule(...)`, so a late tool call that survived the stream layer is not handed to the scheduler. 3. Scheduler (`packages/core/src/scheduler/scheduler.ts`): in `_startBatch`, short-circuit via `cancelAllQueued('Operation cancelled')` if the caller's signal is already aborted, before validating + enqueuing. This mirrors the existing pattern in `_processNextItem` and is the final backstop against late tool calls reaching the executor. Each layer has a regression test for the late-after-cancel ordering. --- packages/cli/src/nonInteractiveCli.ts | 7 +++ packages/core/src/core/turn.test.ts | 48 +++++++++++++++++++ packages/core/src/core/turn.ts | 12 ++++- packages/core/src/scheduler/scheduler.test.ts | 15 ++++++ packages/core/src/scheduler/scheduler.ts | 12 +++++ 5 files changed, 93 insertions(+), 1 deletion(-) diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 00b48be9540..45c3027232e 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -451,6 +451,13 @@ export async function runNonInteractive( } if (toolCallRequests.length > 0) { + // Re-check the abort signal before scheduling any tool calls. + // A SIGINT can fire after the stream loop has already buffered + // a tool-call event; without this guard the side effect would + // execute even though the user has cancelled. + if (abortController.signal.aborted) { + handleCancellationError(config); + } textOutput.ensureTrailingNewline(); const completedToolCalls = await scheduler.schedule( toolCallRequests, diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index cd769abb0dd..65972bd5ae6 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -237,6 +237,54 @@ describe('Turn', () => { expect(turn.getDebugResponses().length).toBe(1); }); + it('should not yield tool_call_request when signal is aborted before a function-call chunk is materialized', async () => { + // Regression test for #28091: a SIGINT can fire after the top-of-loop + // abort check but before we materialize tool calls from a chunk that + // arrived after cancellation. Without the extra guard the side effect + // runs after the user cancelled. + const abortController = new AbortController(); + const mockResponseStream = (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'First part' }] } }], + } as GenerateContentResponse, + }; + abortController.abort(); + yield { + type: StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'late-call', + name: 'late_after_cancel', + args: {}, + isClientInitiated: false, + }, + ], + } as unknown as GenerateContentResponse, + }; + })(); + mockSendMessageStream.mockResolvedValue(mockResponseStream); + + const events = []; + const reqParts: Part[] = [{ text: 'Test late tool call' }]; + for await (const event of turn.run( + { model: 'gemini' }, + reqParts, + abortController.signal, + )) { + events.push(event); + } + + expect(events).toEqual([ + { type: GeminiEventType.Content, value: 'First part' }, + { type: GeminiEventType.UserCancelled }, + ]); + // No tool call should have been queued for the scheduler. + expect(turn.pendingToolCalls).toEqual([]); + }); + it('should yield InvalidStream event if sendMessageStream throws InvalidStreamError', async () => { const error = new InvalidStreamError( 'Test invalid stream', diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 74771c4478e..c37113a512d 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -364,8 +364,18 @@ export class Turn { yield { type: GeminiEventType.Content, value: text, traceId }; } - // Handle function calls (requesting tool execution) + // Handle function calls (requesting tool execution). + // + // Re-check the abort signal here: a SIGINT can fire between the + // top-of-loop check above and the moment we materialize tool calls + // from a chunk that arrived after cancellation. Without this guard + // a delayed function-call chunk is yielded to the scheduler and + // the side effect runs after the user cancelled. const functionCalls = resp.functionCalls ?? []; + if (functionCalls.length > 0 && signal?.aborted) { + yield { type: GeminiEventType.UserCancelled }; + return; + } for (const fnCall of functionCalls) { const event = this.handlePendingFunctionCall(fnCall, traceId); if (event) { diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 9b5935cdbeb..89706164586 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -612,6 +612,21 @@ describe('Scheduler (Orchestrator)', () => { expect(mockStateManager.dequeue).not.toHaveBeenCalled(); // Loop broke }); + it('should not enqueue or validate tool calls when scheduled with an already-aborted signal (regression #28091)', async () => { + // If a delayed tool-call chunk reaches the scheduler after the user + // cancelled, we must not invoke the tool registry / validators or + // enqueue the call — the late side effect would run before the queue + // processor's own abort check kicked in. + abortController.abort(); + + await scheduler.schedule(req1, signal); + + expect(mockStateManager.enqueue).not.toHaveBeenCalled(); + expect(mockStateManager.cancelAllQueued).toHaveBeenCalledWith( + 'Operation cancelled', + ); + }); + it('cancelAll() should cancel active call and clear queue', () => { const activeCall: ValidatingToolCall = { status: CoreToolCallStatus.Validating, diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 76529c14d36..49527f72ac4 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -306,6 +306,18 @@ export class Scheduler { this.isProcessing = true; this.isCancelling = false; this.state.clearBatch(); + + // Guard against late-arriving requests scheduled after cancellation: + // if the caller's signal is already aborted, short-circuit instead of + // validating + enqueuing tool calls that would then execute their + // local side effects before the queue processor sees the abort. + // Match the cancellation pattern used by `_processNextItem`. + if (signal.aborted) { + this.state.cancelAllQueued('Operation cancelled'); + this.isProcessing = false; + return this.state.completedBatch; + } + const currentApprovalMode = this.config.getApprovalMode(); // Sort requests to ensure Topic changes happen before actions in the same batch. From e0efeb8cfe2c2a6f5d173c1c106f8fe687237747 Mon Sep 17 00:00:00 2001 From: KittiphonKamnuan Date: Tue, 23 Jun 2026 08:29:32 +0700 Subject: [PATCH 2/2] fix(scheduler): drain request queue when aborting in _startBatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to gemini-code-assist review on #28096. The previous fix short-circuited `_startBatch` with an early `return` *outside* the `try { ... } finally { ... }`. That bypassed the `finally` block, so `_processNextInRequestQueue()` never ran when the signal was already aborted — any concurrent batches queued in `requestQueue` would hang indefinitely instead of being drained / rejected. Move the abort guard inside the `try` so the existing `finally` (`isProcessing = false`, `clearBatch()`, `_processNextInRequestQueue()`) runs on the early-return path too. Add a regression test that schedules a second batch after the aborted one to confirm the scheduler still makes progress. --- packages/core/src/scheduler/scheduler.test.ts | 18 +++++++++ packages/core/src/scheduler/scheduler.ts | 37 ++++++++++--------- 2 files changed, 37 insertions(+), 18 deletions(-) diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 89706164586..4ba7304ceb0 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -627,6 +627,24 @@ describe('Scheduler (Orchestrator)', () => { ); }); + it('should not leak queued batches when scheduled with an already-aborted signal (regression #28091)', async () => { + // The aborted-signal short-circuit must still drain the request + // queue via the `finally` block — otherwise a follow-up batch + // queued behind it would never resolve and hang the caller. + abortController.abort(); + + // First batch is rejected by the abort guard. + await scheduler.schedule(req1, signal); + + // A second batch scheduled with a fresh, non-aborted signal must + // still make progress (i.e. the scheduler must not be stuck in the + // "busy" state after the early-return path). + const fresh = new AbortController(); + await expect( + scheduler.schedule(req2, fresh.signal), + ).resolves.toBeDefined(); + }); + it('cancelAll() should cancel active call and clear queue', () => { const activeCall: ValidatingToolCall = { status: CoreToolCallStatus.Validating, diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 49527f72ac4..53d338e32e1 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -307,27 +307,28 @@ export class Scheduler { this.isCancelling = false; this.state.clearBatch(); - // Guard against late-arriving requests scheduled after cancellation: - // if the caller's signal is already aborted, short-circuit instead of - // validating + enqueuing tool calls that would then execute their - // local side effects before the queue processor sees the abort. - // Match the cancellation pattern used by `_processNextItem`. - if (signal.aborted) { - this.state.cancelAllQueued('Operation cancelled'); - this.isProcessing = false; - return this.state.completedBatch; - } + try { + // Guard against late-arriving requests scheduled after cancellation: + // if the caller's signal is already aborted, short-circuit instead of + // validating + enqueuing tool calls that would then execute their + // local side effects before the queue processor sees the abort. + // Match the cancellation pattern used by `_processNextItem`. + // Inside the `try` so the `finally` still drains the request queue — + // otherwise concurrent batches waiting in `requestQueue` would hang. + if (signal.aborted) { + this.state.cancelAllQueued('Operation cancelled'); + return this.state.completedBatch; + } - const currentApprovalMode = this.config.getApprovalMode(); + const currentApprovalMode = this.config.getApprovalMode(); - // Sort requests to ensure Topic changes happen before actions in the same batch. - const sortedRequests = [...requests].sort((a, b) => { - if (a.name === UPDATE_TOPIC_TOOL_NAME) return -1; - if (b.name === UPDATE_TOPIC_TOOL_NAME) return 1; - return 0; - }); + // Sort requests to ensure Topic changes happen before actions in the same batch. + const sortedRequests = [...requests].sort((a, b) => { + if (a.name === UPDATE_TOPIC_TOOL_NAME) return -1; + if (b.name === UPDATE_TOPIC_TOOL_NAME) return 1; + return 0; + }); - try { const toolRegistry = this.context.toolRegistry; const newCalls: ToolCall[] = sortedRequests.map((request) => { const enrichedRequest: ToolCallRequestInfo = {