diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 1fb02f99df..7346e10162 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -258,10 +258,10 @@ describe('GeminiChat', async () => { } function streamResponse( - response: GenerateContentResponse, + ...responses: GenerateContentResponse[] ): AsyncGenerator { return (async function* () { - yield response; + yield* responses; })(); } @@ -6490,7 +6490,7 @@ describe('GeminiChat', async () => { expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); }); - it('tries the next fallback when a fallback emits only preparation metadata', async () => { + it('continues fallback after usage and preparation metadata', async () => { vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ authType: AuthType.USE_GEMINI, model: 'test-model', @@ -6547,7 +6547,14 @@ describe('GeminiChat', async () => { ); vi.mocked( mockContentGenerator.generateContentStream, - ).mockRejectedValueOnce(capacityError); + ).mockResolvedValueOnce( + (async function* () { + yield { + usageMetadata: { promptTokenCount: 10, totalTokenCount: 10 }, + } as GenerateContentResponse; + throw capacityError; + })(), + ); const preparationResponse = { candidates: [{ content: { parts: [] } }], } as unknown as GenerateContentResponse; @@ -6596,6 +6603,13 @@ describe('GeminiChat', async () => { expect( events.filter((event) => event.type === StreamEventType.MODEL_FALLBACK), ).toHaveLength(2); + expect( + events.some( + (event) => + event.type === StreamEventType.CHUNK && + event.value.usageMetadata?.promptTokenCount === 10, + ), + ).toBe(true); expect( events.filter((event) => event.type === StreamEventType.MODEL_FALLBACK), ).toEqual([ @@ -9274,6 +9288,533 @@ describe('GeminiChat', async () => { ]); }); + it.each([ + { + name: 'array with a different first argument key', + leakedJson: JSON.stringify([ + { + file_path: 'a.ts', + prompt: 'Create the node.', + name: 'create_node', + subagent_type: 'general-purpose', + run_in_background: true, + }, + { + name: 'read_ref', + prompt: 'Read the reference.', + subagent_type: 'general-purpose', + run_in_background: true, + }, + ]), + trailingText: '', + finishWithContent: false, + }, + { + name: 'single object with trailing prose', + leakedJson: JSON.stringify({ command: 'ls', name: 'run_shell_command' }), + trailingText: 'Let me continue.', + finishWithContent: true, + }, + ])( + 'retries a JSON tool protocol leak: $name', + async ({ leakedJson, trailingText, finishWithContent }) => { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + const leakedText = + leakedJson + '\n\n\n' + trailingText; + const leakedResponses = [ + { + candidates: [ + { + content: { + parts: [ + { text: '...', thought: true }, + { text: leakedText.slice(0, 40) }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse, + { + candidates: [ + { + content: { parts: [{ text: leakedText.slice(40) }] }, + ...(finishWithContent ? { finishReason: 'STOP' as const } : {}), + }, + ], + } as unknown as GenerateContentResponse, + ]; + if (!finishWithContent) { + leakedResponses.push({ + candidates: [{ finishReason: 'STOP' }], + } as unknown as GenerateContentResponse); + } + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(streamResponse(...leakedResponses)) + .mockResolvedValueOnce( + streamResponse(stopResponse([{ text: 'Successful final response' }])), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-tool-protocol-leak', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 2, + ); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + true, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(emittedParts).toEqual([{ text: 'Successful final response' }]); + expect(chatWithRecording.getLastModelMessageText()).toBe( + 'Successful final response', + ); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordAssistantTurn.mock.calls[0]?.[0].message).toEqual([ + { text: 'Successful final response' }, + ]); + }, + ); + + it.each([ + [JSON.stringify([{ name: 'example', value: 1 }]), 8, 1], + ['[1,2,3]', 2, 2], + ])( + 'preserves an ordinary leading JSON array: %s', + async (response, splitAt, expectedTextChunks) => { + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce( + streamResponse( + { + candidates: [ + { + content: { + parts: [ + { text: 'thinking', thought: true }, + { text: response.slice(0, splitAt) }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse, + stopResponse([{ text: response.slice(splitAt) }]), + ), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-array-literal', + ); + const events: StreamEvent[] = []; + const streamedTextChunks: string[] = []; + for await (const event of stream) { + events.push(event); + if (event.type === StreamEventType.CHUNK) { + const text = + event.value.candidates?.[0]?.content?.parts + ?.filter((part) => !part.thought) + .map((part) => part.text ?? '') + .join('') ?? ''; + if (text) streamedTextChunks.push(text); + } + } + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 1, + ); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + false, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(streamedTextChunks).toHaveLength(expectedTextChunks); + expect(emittedParts.find((part) => part.thought)?.text).toBe('thinking'); + expect(streamedTextChunks.join('')).toBe(response); + expect(chat.getLastModelMessageText()).toBe(response); + }, + ); + + it('releases buffered JSON through a finish-only chunk without leaked tags', async () => { + const response = JSON.stringify([{ name: 'example', value: 1 }]); + const splitAt = 12; + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValueOnce( + streamResponse( + { + candidates: [ + { content: { parts: [{ text: response.slice(0, splitAt) }] } }, + ], + } as unknown as GenerateContentResponse, + { + candidates: [ + { content: { parts: [{ text: response.slice(splitAt) }] } }, + ], + } as unknown as GenerateContentResponse, + { + candidates: [{ finishReason: 'STOP' }], + } as unknown as GenerateContentResponse, + ), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-finish-only', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + false, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + const emittedText = emittedParts + .filter((part) => !part.thought) + .map((part) => part.text ?? '') + .join(''); + expect(emittedText).toBe(response); + expect(chat.getLastModelMessageText()).toBe(response); + expect(chat.getHistory().at(-1)?.parts).toEqual(emittedParts); + }); + + it.each(['preparation', 'function call'] as const)( + 'retries when a %s interrupts a partial JSON protocol leak', + async (middleChunkType) => { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + const leakedText = + JSON.stringify([{ name: 'read_file', file_path: 'a.ts' }]) + + '\n\n\n'; + const splitAt = 12; + let middleResponse: GenerateContentResponse; + if (middleChunkType === 'preparation') { + middleResponse = { + candidates: [{ content: { parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(middleResponse, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + } else { + middleResponse = { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + } + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + streamResponse( + { + candidates: [ + { + content: { parts: [{ text: leakedText.slice(0, splitAt) }] }, + }, + ], + } as unknown as GenerateContentResponse, + middleResponse, + stopResponse([{ text: leakedText.slice(splitAt) }]), + ), + ) + .mockResolvedValueOnce( + streamResponse(stopResponse([{ text: 'Successful final response' }])), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-interrupted-json-tool-protocol-leak', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 2, + ); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + true, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(emittedParts).toEqual([{ text: 'Successful final response' }]); + expect(chatWithRecording.getLastModelMessageText()).toBe( + 'Successful final response', + ); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordAssistantTurn.mock.calls[0]?.[0].message).toEqual([ + { text: 'Successful final response' }, + ]); + }, + ); + + it.each([true, false])( + 'preserves leading JSON when a tool call ends without a finish reason (tool call first: %s)', + async (toolCallFirst) => { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + const response = JSON.stringify([{ name: 'example', value: 1 }]); + const splitAt = 12; + const functionCallPart = { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }; + const functionCallResponse = { + candidates: [{ content: { parts: [functionCallPart] } }], + } as unknown as GenerateContentResponse; + const textResponses = [ + response.slice(0, splitAt), + response.slice(splitAt), + ].map( + (text) => + ({ + candidates: [{ content: { parts: [{ text }] } }], + }) as unknown as GenerateContentResponse, + ); + const responses = toolCallFirst + ? [functionCallResponse, ...textResponses] + : [textResponses[0]!, functionCallResponse, textResponses[1]!]; + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce(streamResponse(...responses)); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + `prompt-id-json-tool-call-no-finish-${toolCallFirst}`, + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 1, + ); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + false, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(chatWithRecording.getHistory().at(-1)?.parts).toEqual( + emittedParts, + ); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + const recordedParts = recordAssistantTurn.mock.calls[0]?.[0] + .message as Part[]; + for (const parts of [emittedParts, recordedParts]) { + expect(parts.filter((part) => part.functionCall)).toEqual([ + functionCallPart, + ]); + expect( + parts + .filter((part) => part.text) + .map((part) => part.text) + .join(''), + ).toBe(response); + } + }, + ); + + it.each([false, true])( + 'keeps leading JSON before a later structured tool call (preparation: %s)', + async (withPreparation) => { + const response = JSON.stringify([{ name: 'example', value: 1 }]); + const preparationResponse = { + candidates: [{ content: { parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + const usageResponse = { + usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 2 }, + } as GenerateContentResponse; + const responses = [ + { + candidates: [ + { + content: { parts: [{ text: response }] }, + }, + ], + } as unknown as GenerateContentResponse, + usageResponse, + ]; + if (withPreparation) responses.push(preparationResponse); + responses.push( + stopResponse([ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ]), + ); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce(streamResponse(...responses)); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-before-tool-call', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + const emittedPreparations = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => getToolCallPreparations(event.value)); + expect( + events.some( + (event) => + event.type === StreamEventType.CHUNK && + event.value.usageMetadata?.promptTokenCount === 10 && + event.value.usageMetadata.candidatesTokenCount === 2, + ), + ).toBe(true); + expect(emittedPreparations).toEqual( + withPreparation ? [{ callId: 'call-1', toolName: 'read_file' }] : [], + ); + for (const parts of [ + emittedParts, + chat.getHistory().at(-1)?.parts ?? [], + ]) { + expect(parts.findIndex((part) => part.text === response)).toBe(0); + expect(parts.findIndex((part) => part.functionCall)).toBe(1); + } + }, + ); + + it('does not retry after a structured tool call has already been emitted', async () => { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + const leakedText = + JSON.stringify([{ name: 'read_file', file_path: 'a.ts' }]) + + '\n\n\n'; + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + streamResponse( + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse, + stopResponse([{ text: leakedText }]), + ), + ) + .mockResolvedValueOnce( + streamResponse(stopResponse([{ text: 'Unexpected retry response' }])), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-tool-call-before-json-protocol-leak', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1); + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + false, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(emittedParts).toEqual([ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ]); + expect(chatWithRecording.getHistory().at(-1)?.parts).toEqual(emittedParts); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordAssistantTurn.mock.calls[0]?.[0].message).toEqual( + emittedParts, + ); + }); + it('does not reject normal HTML or protocol tag names in prose', async () => { const response = '
Title
' + @@ -9307,6 +9848,146 @@ describe('GeminiChat', async () => { expect(chat.getLastModelMessageText()).toBe(response); }); + it('does not reject closing protocol tags inside a JSON string', async () => { + const response = JSON.stringify({ + example: '} text', + }); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValueOnce( + streamResponse(stopResponse([{ text: response }])), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-protocol-literal', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(events.some((event) => event.type === StreamEventType.RETRY)).toBe( + false, + ); + expect(chat.getLastModelMessageText()).toBe(response); + }); + + it('retries leaked JSON before a structured tool call', async () => { + vi.useFakeTimers(); + try { + const leakedText = + JSON.stringify([{ name: 'read_file', file_path: 'a.ts' }]) + + '\n\n\n'; + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + streamResponse( + { + candidates: [{ content: { parts: [{ text: leakedText }] } }], + } as unknown as GenerateContentResponse, + stopResponse([ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ]), + ), + ) + .mockResolvedValueOnce( + streamResponse(stopResponse([{ text: 'Successful final response' }])), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-leak-before-tool-call', + ); + const events: StreamEvent[] = []; + const iterator = stream[Symbol.asyncIterator](); + for (;;) { + const next = iterator.next(); + await vi.advanceTimersByTimeAsync(5_000); + const result = await next; + if (result.done) break; + events.push(result.value); + } + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 2, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(emittedParts).toEqual([{ text: 'Successful final response' }]); + expect(chat.getHistory().at(-1)?.parts).toEqual(emittedParts); + } finally { + vi.useRealTimers(); + } + }); + + it('retries leaked JSON without a finish reason via the post-stream leak guard', async () => { + vi.useFakeTimers(); + try { + const leakedText = + JSON.stringify([{ name: 'read_file', file_path: 'a.ts' }]) + + '\n\n\n\n'; + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + streamResponse( + { + candidates: [{ content: { parts: [{ text: leakedText }] } }], + } as unknown as GenerateContentResponse, + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.ts' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse, + ), + ) + .mockResolvedValueOnce( + streamResponse(stopResponse([{ text: 'Successful final response' }])), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-json-leak-no-finish-reason', + ); + const events: StreamEvent[] = []; + const iterator = stream[Symbol.asyncIterator](); + for (;;) { + const next = iterator.next(); + await vi.advanceTimersByTimeAsync(5_000); + const result = await next; + if (result.done) break; + events.push(result.value); + } + + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 2, + ); + const emittedParts = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => event.value.candidates?.[0]?.content?.parts ?? []); + expect(emittedParts).toEqual([{ text: 'Successful final response' }]); + expect(chat.getHistory().at(-1)?.parts).toEqual(emittedParts); + } finally { + vi.useRealTimers(); + } + }); + it('retries a protocol-tagged turn even when the leaked attempt also contains a tool call', async () => { vi.useFakeTimers(); try { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 639784d42a..a67f23e2ba 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -134,15 +134,14 @@ const debugLogger = createDebugLogger('QWEN_CODE_CHAT'); // must stay in sync. const GEMINI_EMPTY_CONTENT_PLACEHOLDER = '(empty content)'; -function isToolCallPreparationOnly(response: GenerateContentResponse): boolean { - if (getToolCallPreparations(response).length === 0) return false; - - const hasCandidateOutput = response.candidates?.some( - (candidate) => - Boolean(candidate.finishReason) || - (candidate.content?.parts?.length ?? 0) > 0, +function hasCandidateOutput(response: GenerateContentResponse): boolean { + return Boolean( + response.candidates?.some( + (candidate) => + Boolean(candidate.finishReason) || + (candidate.content?.parts?.length ?? 0) > 0, + ), ); - return !hasCandidateOutput && !response.usageMetadata; } /** @@ -1184,9 +1183,29 @@ const PROTOCOL_TAG_PREFIXES = [ '\s*<\/function>/iy; + +function hasLeakedToolCallTags(text: string): boolean { + let inString = false; + let escaped = false; + for (let i = 0; i < text.length; i++) { + const char = text[i]; + if (inString) { + if (escaped) escaped = false; + else if (char === '\\') escaped = true; + else if (char === '"') inString = false; + } else if (char === '"') { + inString = true; + } else if (char === '}' || char === ']') { + LEAKED_TOOL_CALL_TAGS.lastIndex = i; + if (LEAKED_TOOL_CALL_TAGS.test(text)) return true; + } + } + return false; +} class LeadingProtocolTagLeakDetector { - private state: 'detecting' | 'clean' | 'leaked' = 'detecting'; + private state: 'detecting' | 'json' | 'clean' | 'leaked' = 'detecting'; private buffer = ''; accept(text: string): string { @@ -1194,6 +1213,7 @@ class LeadingProtocolTagLeakDetector { if (this.state === 'leaked') return ''; this.buffer += text; + if (this.state === 'json') return ''; const candidate = this.buffer.trimStart().toLowerCase(); if (!candidate) return ''; if (PROTOCOL_TAG_PREFIXES.some((prefix) => prefix.startsWith(candidate))) { @@ -1210,14 +1230,31 @@ class LeadingProtocolTagLeakDetector { return ''; } } + if (candidate.startsWith('{')) { + this.state = 'json'; + return ''; + } + if (candidate.startsWith('[')) { + const normalized = candidate.replace(/\s/g, ''); + if (normalized === '[') return ''; + if (normalized.startsWith('[{')) { + this.state = 'json'; + return ''; + } + } - this.state = 'clean'; - const output = this.buffer; - this.buffer = ''; - return output; + return this.release(); } finish(): string { + if (this.state === 'json') { + if (hasLeakedToolCallTags(this.buffer)) { + this.state = 'leaked'; + this.buffer = ''; + return ''; + } + return this.release(); + } if (this.state !== 'detecting') return ''; const candidate = this.buffer.trimStart().toLowerCase(); if ( @@ -1228,8 +1265,12 @@ class LeadingProtocolTagLeakDetector { this.buffer = ''; return ''; } - this.state = 'clean'; + return this.release(); + } + + private release(): string { const output = this.buffer; + this.state = 'clean'; this.buffer = ''; return output; } @@ -2513,7 +2554,7 @@ export class GeminiChat { lastFinishReason = undefined; for await (const chunk of stream) { - if (!isToolCallPreparationOnly(chunk)) { + if (hasCandidateOutput(chunk)) { streamYieldedChunk = true; streamYieldedAnyChunk = true; } @@ -3350,7 +3391,7 @@ export class GeminiChat { )) { const emittedUserVisibleOutput = event.type !== StreamEventType.CHUNK || - !isToolCallPreparationOnly(event.value); + hasCandidateOutput(event.value); if (emittedUserVisibleOutput) { currentFallbackYieldedAnyChunk = true; fallbackStreamYieldedAnyChunk = true; @@ -4063,6 +4104,25 @@ export class GeminiChat { let hasToolCall = false; let hasFinishReason = false; const protocolTagDetector = new LeadingProtocolTagLeakDetector(); + let pendingProtocolParts: Part[] = []; + const takePendingProtocolParts = (): Part[] => { + const parts = pendingProtocolParts; + pendingProtocolParts = []; + const released: Part[] = []; + for (const part of parts) { + const previous = released.at(-1); + if ( + previous && + isValidNonThoughtTextPart(previous) && + isValidNonThoughtTextPart(part) + ) { + previous.text! += part.text!; + } else { + released.push(isValidNonThoughtTextPart(part) ? { ...part } : part); + } + } + return released; + }; let protocolTextWasSuppressed = false; const currentUserTurn = this.history[this.history.length - 1]; const isToolResultContinuation = @@ -4102,28 +4162,64 @@ export class GeminiChat { if (isValidResponse(chunk)) { const candidate = chunk.candidates?.[0]; - const content = candidate?.content; + let content = candidate?.content; + if (candidate?.finishReason && !content?.parts) { + protocolTagDetector.finish(); + if (protocolTagDetector.leaked) { + pendingProtocolParts = []; + } else { + const parts = takePendingProtocolParts(); + if (parts.length > 0) { + content = { + ...content, + role: content?.role ?? 'model', + parts, + }; + candidate.content = content; + } + } + } if (content?.parts) { - content.parts = content.parts.flatMap((part) => { + const outputParts: Part[] = []; + for (const part of content.parts) { if ( isToolResultContinuation && !part.thought && part.text?.trim() === GEMINI_EMPTY_CONTENT_PLACEHOLDER ) { - return []; + continue; + } + if (typeof part.text !== 'string' || part.thought) { + if ( + pendingProtocolParts.length > 0 || + protocolTagDetector.leaked + ) { + pendingProtocolParts.push(part); + } else { + outputParts.push(part); + } + continue; } - if (typeof part.text !== 'string' || part.thought) return [part]; const text = protocolTagDetector.accept(part.text); - if (text) return [{ ...part, text }]; + if (text) { + if (pendingProtocolParts.length > 0) { + outputParts.push(...takePendingProtocolParts(), part); + } else { + outputParts.push({ ...part, text }); + } + continue; + } + pendingProtocolParts.push(...outputParts.splice(0), part); protocolTextWasSuppressed ||= part.text.length > 0; - const { text: _text, ...rest } = part; - return Object.values(rest).some((value) => value !== undefined) - ? [rest] - : []; - }); + } + content.parts = outputParts; if (candidate?.finishReason) { - const text = protocolTagDetector.finish(); - if (text) content.parts.push({ text }); + protocolTagDetector.finish(); + if (protocolTagDetector.leaked) { + pendingProtocolParts = []; + } else { + content.parts.push(...takePendingProtocolParts()); + } } content.parts = normalizeModelToolCallIds( content.parts, @@ -4226,7 +4322,12 @@ export class GeminiChat { } } - if (!protocolTextWasSuppressed || !protocolTagDetector.blockingOutput) { + if ( + !chunk.candidates?.length || + preparations.length > 0 || + !protocolTextWasSuppressed || + !protocolTagDetector.blockingOutput + ) { yield chunk; } } @@ -4234,6 +4335,32 @@ export class GeminiChat { streamError = e; } + if ( + streamError === null && + pendingProtocolParts.length > 0 && + (hasToolCall || + pendingProtocolParts.some((part) => part.functionCall !== undefined)) + ) { + protocolTagDetector.finish(); + if (protocolTagDetector.leaked) { + pendingProtocolParts = []; + } else { + const parts = normalizeModelToolCallIds( + takePendingProtocolParts(), + usedToolCallIds, + rawToolCallIdsInCurrentTurn, + reservedToolCallIds, + ); + const chunk = { + candidates: [{ content: { role: 'model', parts } }], + } as GenerateContentResponse; + syncFunctionCallsField(chunk, parts); + hasToolCall ||= parts.some((part) => part.functionCall); + allModelParts.push(...parts); + yield chunk; + } + } + let thoughtContentPart: Part | undefined; const thoughtText = allModelParts .filter((part) => part.thought) @@ -4348,7 +4475,7 @@ export class GeminiChat { } } - if (streamError === null && protocolTagDetector.leaked) { + if (streamError === null && protocolTagDetector.leaked && !hasToolCall) { throw new InvalidStreamError( 'Model response started with leaked protocol tags.', 'PROTOCOL_TAG_LEAK',