From 56f9688b306b937968ce1576b5de27754f29ed5c Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:07:50 -0400 Subject: [PATCH] fix(core): stop a new user message fusing into an unanswered tool response (#28700) Co-authored-by: David Pierce --- packages/core/src/core/geminiChat.test.ts | 223 ++++++++++++++++++++++ packages/core/src/core/geminiChat.ts | 39 ++++ 2 files changed, 262 insertions(+) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 6443df5824..e46f6730e9 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -228,6 +228,16 @@ describe('GeminiChat', () => { // Disable 429 simulation for tests setSimulate429(false); + + // The mid-stream retry loop sleeps on a real timer (1s + 2s + 4s) between + // attempts, which exceeds the default 5s test timeout and silently killed + // every InvalidStreamError test before it reached its assertions. Run those + // delays instantly. + vi.spyOn(globalThis, 'setTimeout').mockImplementation(((fn: () => void) => { + fn(); + return 0; + }) as unknown as typeof globalThis.setTimeout); + // Reset history for each test by creating a new instance chat = new GeminiChat(mockConfig); mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); @@ -999,6 +1009,219 @@ describe('GeminiChat', () => { expect(lastTurn.content.parts?.[0]?.functionResponse).toBeDefined(); }); + it('should not fuse the next user message into a preserved tool-response turn', async () => { + // Regression: when a stream fails mid tool-loop the tool response is + // deliberately preserved (see the test above), which leaves history + // ending on a user turn. The user's next message was then coalesced into + // that same turn as [functionResponse, text]. The model reads the + // trailing text as a continuation of the tool result and completes the + // sentence instead of answering it. + chat.agentHistory.push({ + id: 'model-turn-1', + content: { + role: 'model', + parts: [{ functionCall: { name: 'test_tool', args: {} } }], + }, + }); + + // 1. Tool response goes back, model returns nothing -> InvalidStreamError. + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield { + candidates: [ + { content: { role: 'model', parts: [] }, finishReason: 'STOP' }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const failingStream = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + [ + { + functionResponse: { + name: 'test_tool', + response: { success: true }, + }, + }, + ], + 'prompt-id-fusion-setup', + new AbortController().signal, + LlmRole.MAIN, + ); + await expect( + (async () => { + for await (const _ of failingStream) { + // consume + } + })(), + ).rejects.toThrow(InvalidStreamError); + + // 2. The user types a brand new instruction. + let capturedContents: Content[] = []; + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async (req) => { + capturedContents = req.contents as Content[]; + return (async function* () { + yield { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + }, + ); + + const stream = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + 'are you done?', + 'prompt-id-fusion-check', + new AbortController().signal, + LlmRole.MAIN, + ); + for await (const _ of stream) { + // consume + } + + const fusedTurn = capturedContents.find( + (c) => + c.role === 'user' && + !!c.parts?.some((p) => !!p.functionResponse) && + !!c.parts?.some((p) => p.text?.includes('are you done?')), + ); + 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. + chat.agentHistory.push({ + id: 'model-turn-cancel', + content: { + role: 'model', + parts: [ + { functionCall: { id: 'c1', name: 'run_shell_command', args: {} } }, + ], + }, + }); + chat.addHistory({ + role: 'user', + parts: [ + { + functionResponse: { + id: 'c1', + name: 'run_shell_command', + response: { error: '[Operation Cancelled]' }, + }, + }, + ], + }); + + let capturedContents: Content[] = []; + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async (req) => { + capturedContents = req.contents as Content[]; + return (async function* () { + yield { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + }, + ); + + const stream = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + "you're querying local database, I meant nprd", + 'prompt-id-cancel-fusion', + new AbortController().signal, + LlmRole.MAIN, + ); + for await (const _ of stream) { + // consume + } + + const fusedCancelTurn = 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(); + }); + + 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([ + { role: 'user', parts: [{ text: 'run the tests' }] }, + { + role: 'model', + parts: [ + { functionCall: { id: 'c1', name: 'run_shell_command', args: {} } }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'c1', + name: 'run_shell_command', + response: { output: 'ok' }, + }, + }, + ], + }, + ]); + + let capturedContents: Content[] = []; + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async (req) => { + capturedContents = req.contents as Content[]; + return (async function* () { + yield { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + }, + ); + + const stream = await chat.sendMessageStream( + { model: 'gemini-2.0-flash' }, + 'are you done?', + 'prompt-id-resumed-fusion', + new AbortController().signal, + LlmRole.MAIN, + ); + for await (const _ of stream) { + // consume + } + + const fusedResumedTurn = 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(); + }); + it('should preserve mixed multimodal function responses during rollback when InvalidStreamError is thrown (regression)', async () => { // 1. Setup history ending with a model turn containing functionCall chat.agentHistory.push({ diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 93ea20ee2a..7dd419f7b5 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -108,6 +108,13 @@ 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.]'; + /** * Internal interface for parts that carry the magic 'callIndex' property * used during model response consolidation. @@ -408,6 +415,16 @@ 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 { model } = this.context.config.modelConfigService.getResolvedConfig(modelConfigKey); @@ -683,6 +700,28 @@ 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 { + const turns = this.agentHistory.get(); + const last = turns[turns.length - 1]; + if ( + last?.content.role !== 'user' || + !last.content.parts?.some((part) => !!part.functionResponse) + ) { + return; + } + this.agentHistory.push({ + id: randomUUID(), + content: { + role: 'model', + parts: [{ text: INTERRUPTED_RESPONSE_PLACEHOLDER }], + }, + }); + } + private extractBinaryInjections( parts: Part[] | undefined, ): Part[] | undefined {