diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 80d37e7fbdf..30d3353e51b 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -50,6 +50,9 @@ import { LlmRole } from '../telemetry/types.js'; import { BINARY_INJECTION_KEY } from '../utils/generateContentResponseUtilities.js'; import type { ResumedSessionData } from '../services/chatRecordingTypes.js'; +const INTERRUPTED_RESPONSE_BOUNDARY = + 'A previous model response was interrupted after a tool result. The content that follows is a new user message. Respond to that message directly.'; + // Mock fs module to prevent actual file system operations during tests const mockFileSystem = new Map(); @@ -1108,10 +1111,7 @@ describe('GeminiChat', () => { expect(fusedTurn).toBeUndefined(); }); - it('should not fuse the next user message into a cancelled tool response', async () => { - // Same defect reached by a different trigger: cancelling a tool call - // records its response via addHistory then returns without submitting, - // leaving history on an unanswered user turn just like a stream failure. + it('should separate a new user message from a cancelled tool response without a model placeholder', async () => { chat.agentHistory.push({ id: 'model-turn-cancel', content: { @@ -1153,7 +1153,10 @@ describe('GeminiChat', () => { const stream = await chat.sendMessageStream( { model: 'gemini-2.0-flash' }, - "you're querying local database, I meant nprd", + [ + { text: "you're querying local database, I meant nprd" }, + { text: 'check the target before retrying' }, + ], 'prompt-id-cancel-fusion', new AbortController().signal, LlmRole.MAIN, @@ -1162,21 +1165,34 @@ describe('GeminiChat', () => { // consume } - const fusedCancelTurn = capturedContents.find( + const repairedTurn = capturedContents.find( (c) => c.role === 'user' && !!c.parts?.some((p) => !!p.functionResponse) && !!c.parts?.some((p) => p.text?.includes('I meant nprd')), ); - expect(fusedCancelTurn).toBeUndefined(); + expect(repairedTurn?.parts).toEqual([ + { + functionResponse: { + id: 'c1', + name: 'run_shell_command', + response: { error: '[Operation Cancelled]' }, + }, + }, + { text: INTERRUPTED_RESPONSE_BOUNDARY }, + { text: "you're querying local database, I meant nprd" }, + { text: 'check the target before retrying' }, + ]); + + expect(JSON.stringify(chat.agentHistory.get())).not.toContain( + INTERRUPTED_RESPONSE_BOUNDARY, + ); }); - it('should close a dangling tool response restored from a resumed session', async () => { - // The guard runs when a new user message arrives rather than when the - // turn fails, so it does not depend on a placeholder having been - // persisted. A session resumed from disk that ends on an unanswered tool - // response is repaired on the next message just the same. - chat.setHistory([ + it('should add a transient boundary to a context-managed history override after a resumed dangling tool response', async () => { + vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(true); + + const history: Content[] = [ { role: 'user', parts: [{ text: 'run the tests' }] }, { role: 'model', @@ -1196,7 +1212,8 @@ describe('GeminiChat', () => { }, ], }, - ]); + ]; + chat.setHistory(history); let capturedContents: Content[] = []; vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( @@ -1221,18 +1238,30 @@ describe('GeminiChat', () => { 'prompt-id-resumed-fusion', new AbortController().signal, LlmRole.MAIN, + undefined, + [...history, { role: 'user', parts: [{ text: 'are you done?' }] }], ); for await (const _ of stream) { // consume } - const fusedResumedTurn = capturedContents.find( + const repairedTurn = capturedContents.find( (c) => c.role === 'user' && !!c.parts?.some((p) => !!p.functionResponse) && !!c.parts?.some((p) => p.text?.includes('are you done?')), ); - expect(fusedResumedTurn).toBeUndefined(); + expect(repairedTurn?.parts).toEqual([ + { + functionResponse: { + id: 'c1', + name: 'run_shell_command', + response: { output: 'ok' }, + }, + }, + { text: INTERRUPTED_RESPONSE_BOUNDARY }, + { text: 'are you done?' }, + ]); }); it('should preserve mixed multimodal function responses during rollback when InvalidStreamError is thrown (regression)', async () => { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 1fe4a305275..fab72a14c15 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -109,12 +109,21 @@ const MID_STREAM_RETRY_OPTIONS: MidStreamRetryOptions = { export const SYNTHETIC_THOUGHT_SIGNATURE = 'skip_thought_signature_validator'; -/** - * Stands in for a model turn that never arrived because the stream failed - * after a tool response was already committed to history. - */ -export const INTERRUPTED_RESPONSE_PLACEHOLDER = - '[The previous response was interrupted before it completed.]'; +const INTERRUPTED_RESPONSE_BOUNDARY = + 'A previous model response was interrupted after a tool result. The content that follows is a new user message. Respond to that message directly.'; + +function insertInterruptionBoundary( + parts: readonly Part[] | undefined, + newMessagePartCount: number, +): Part[] { + const existingParts = parts ?? []; + const boundaryIndex = Math.max(0, existingParts.length - newMessagePartCount); + return [ + ...existingParts.slice(0, boundaryIndex), + { text: INTERRUPTED_RESPONSE_BOUNDARY }, + ...existingParts.slice(boundaryIndex), + ]; +} /** * Internal interface for parts that carry the magic 'callIndex' property @@ -511,14 +520,8 @@ export class GeminiChat { let userContent = createUserContent(message); const isOriginalFunctionResponse = isFunctionResponse(userContent); - // A turn can end leaving history on an unanswered tool response: a stream - // error after the response was committed, or a cancelled tool call. Close - // it before recording a genuinely new user message, otherwise the two user - // turns are coalesced into one and the model continues the trailing text - // instead of answering it. - if (!isOriginalFunctionResponse) { - this.closeUnansweredToolResponseTurn(); - } + const needsInterruptionBoundary = + !isOriginalFunctionResponse && this.hasUnansweredToolResponseTurn(); const { model } = this.context.config.modelConfigService.getResolvedConfig(modelConfigKey); @@ -638,7 +641,22 @@ export class GeminiChat { } } - const requestHistory = this.getHistoryTurns(true); + let requestHistory = this.getHistoryTurns(true); + let requestApiHistoryOverride = apiHistoryOverride; + if (needsInterruptionBoundary) { + const newMessagePartCount = userContent.parts?.length ?? 0; + if (requestApiHistoryOverride) { + requestApiHistoryOverride = this.addInterruptionBoundaryToContents( + requestApiHistoryOverride, + newMessagePartCount, + ); + } else { + requestHistory = this.addInterruptionBoundaryToHistory( + requestHistory, + newMessagePartCount, + ); + } + } const streamWithRetries = async function* ( this: GeminiChat, @@ -670,7 +688,7 @@ export class GeminiChat { prompt_id, signal, role, - apiHistoryOverride, + requestApiHistoryOverride, isOriginalFunctionResponse, ); isConnectionPhase = false; @@ -824,26 +842,55 @@ export class GeminiChat { return streamWithRetries.call(this); } - /** - * Appends a closing model turn when history ends with an unanswered tool - * response, so the next user message stays a turn of its own. - */ - private closeUnansweredToolResponseTurn(): void { + private hasUnansweredToolResponseTurn(): boolean { const turns = this.agentHistory.get(); - const last = turns[turns.length - 1]; - if ( - last?.content.role !== 'user' || - !last.content.parts?.some((part) => !!part.functionResponse) - ) { - return; + const last = turns.at(-1); + return !!( + last?.content.role === 'user' && + last.content.parts?.some((part) => !!part.functionResponse) + ); + } + + private addInterruptionBoundaryToHistory( + history: readonly HistoryTurn[], + newMessagePartCount: number, + ): HistoryTurn[] { + const last = history.at(-1); + if (last?.content.role !== 'user') { + return [...history]; } - this.agentHistory.push({ - id: randomUUID(), - content: { - role: 'model', - parts: [{ text: INTERRUPTED_RESPONSE_PLACEHOLDER }], + + return [ + ...history.slice(0, -1), + { + ...last, + content: { + ...last.content, + parts: insertInterruptionBoundary( + last.content.parts, + newMessagePartCount, + ), + }, }, - }); + ]; + } + + private addInterruptionBoundaryToContents( + contents: readonly Content[], + newMessagePartCount: number, + ): Content[] { + const last = contents.at(-1); + if (last?.role !== 'user') { + return [...contents]; + } + + return [ + ...contents.slice(0, -1), + { + ...last, + parts: insertInterruptionBoundary(last.parts, newMessagePartCount), + }, + ]; } private extractBinaryInjections(