diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 73c94bd34ea..f7b1be37a5f 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -1628,6 +1628,89 @@ describe('GeminiChat', () => { expect(chat.agentHistory.length).toBe(initialHistoryLength); }); + it('should roll back the entire multi-turn request including function responses when a continuation stream is aborted/cancelled', async () => { + const initialHistoryLength = chat.agentHistory.length; + const abortController = new AbortController(); + + // 1. Send the first message of the prompt. This will succeed and register prompt-id-multi-turn-abort. + const streamFirst = (async function* () { + yield { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'model first response' }], + }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + streamFirst, + ); + + const s1 = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + 'user original prompt', + 'prompt-id-multi-turn-abort', + new AbortController().signal, + LlmRole.MAIN, + ); + for await (const _ of s1) { + // consume the stream + } + + // Expect history to contain: user, model + expect(chat.agentHistory.length).toBe(initialHistoryLength + 2); + + // 2. Send a continuation (functionResponse), which is cancelled mid-stream. + const streamSecond = (async function* () { + yield { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'partial model response before abort' }], + }, + }, + ], + } as unknown as GenerateContentResponse; + abortController.abort(); + throw new Error('User aborted a continuation stream.'); + })(); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + streamSecond, + ); + + const s2 = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + [ + { + functionResponse: { + id: 'call_id_1', + name: 'my_tool', + response: { result: 'success' }, + }, + }, + ], + 'prompt-id-multi-turn-abort', + abortController.signal, + LlmRole.MAIN, + ); + + await expect( + (async () => { + for await (const _ of s2) { + // consume the stream to trigger abort + } + })(), + ).rejects.toThrow(); + + // Verify history has been rolled back entirely to initialHistoryLength (before the original prompt started)! + expect(chat.agentHistory.length).toBe(initialHistoryLength); + }); + it('should roll back the un-responded user turn from history when an ApiError is thrown', async () => { const initialHistoryLength = chat.agentHistory.length; diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index b561285635f..332a10a1c2c 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -38,6 +38,7 @@ import { import { hasCycleInSchema } from '../tools/tools.js'; import type { StructuredError } from './turn.js'; import type { CompletedToolCall } from '../scheduler/types.js'; +import { isAbortError } from '../utils/errors.js'; import { logContentRetry, logContentRetryFailure, @@ -296,6 +297,9 @@ export class GeminiChat { private lastPromptTokenCount: number; private callCounter = 0; agentHistory: AgentChatHistory; + private lastPromptId?: string; + private promptOriginalHistoryLength?: number; + private promptOriginalTokenCount?: number; constructor( readonly context: AgentLoopContext, @@ -407,6 +411,17 @@ export class GeminiChat { const historyLengthBefore = this.agentHistory.length; const baselinePromptTokenCount = this.lastPromptTokenCount; + if (this.lastPromptId && this.lastPromptId !== prompt_id) { + this.promptOriginalHistoryLength = undefined; + this.promptOriginalTokenCount = undefined; + } + this.lastPromptId = prompt_id; + + if (this.promptOriginalHistoryLength === undefined) { + this.promptOriginalHistoryLength = historyLengthBefore; + this.promptOriginalTokenCount = baselinePromptTokenCount; + } + let streamDoneResolver: () => void; const streamDonePromise = new Promise((resolve) => { streamDoneResolver = resolve; @@ -684,7 +699,26 @@ export class GeminiChat { } } } catch (error) { - if (!isOriginalFunctionResponse) { + const isAborted = + signal?.aborted || + isAbortError(error) || + (error instanceof Error && + (error.name === 'CanceledError' || + error.name === 'FatalCancellationError')); + const originalLength = this.promptOriginalHistoryLength; + const originalTokenCount = this.promptOriginalTokenCount; + if (isAborted && originalLength !== undefined) { + this.agentHistory.rollback(originalLength); + this.chatRecordingService.updateMessagesFromHistory( + this.agentHistory.get(), + ); + if (originalTokenCount !== undefined) { + this.lastPromptTokenCount = originalTokenCount; + } + this.promptOriginalHistoryLength = undefined; + this.promptOriginalTokenCount = undefined; + this.lastPromptId = undefined; + } else if (!isOriginalFunctionResponse) { this.agentHistory.rollback(historyLengthBefore); this.chatRecordingService.updateMessagesFromHistory( this.agentHistory.get(),