From c257394a826eb8d60254020992207ac4fc9ddaf9 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 13:05:10 +0800 Subject: [PATCH 1/8] fix(core): sanitize standalone closing thinking tags --- packages/core/src/core/geminiChat.test.ts | 87 ++++++- .../openaiContentGenerator/converter.test.ts | 246 ++++++++++++++++++ .../core/openaiContentGenerator/converter.ts | 139 +++++++++- .../openaiContentGenerator/pipeline.test.ts | 129 +++++++++ .../core/openaiContentGenerator/pipeline.ts | 38 ++- .../src/core/openaiContentGenerator/types.ts | 9 + packages/core/src/telemetry/constants.ts | 2 + packages/core/src/telemetry/index.ts | 2 + packages/core/src/telemetry/loggers.test.ts | 40 +++ packages/core/src/telemetry/loggers.ts | 21 ++ .../telemetry/qwen-logger/qwen-logger.test.ts | 34 +++ .../src/telemetry/qwen-logger/qwen-logger.ts | 17 ++ packages/core/src/telemetry/types.ts | 29 +++ 13 files changed, 779 insertions(+), 14 deletions(-) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 10e5d41aa57..100f0319616 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -81,14 +81,20 @@ vi.mock('../utils/retry.js', async (importOriginal) => { }; }); -const { mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({ +const { + mockLogContentRetry, + mockLogContentRetryFailure, + mockLogProtocolTagSanitized, +} = vi.hoisted(() => ({ mockLogContentRetry: vi.fn(), mockLogContentRetryFailure: vi.fn(), + mockLogProtocolTagSanitized: vi.fn(), })); vi.mock('../telemetry/loggers.js', () => ({ logContentRetry: mockLogContentRetry, logContentRetryFailure: mockLogContentRetryFailure, + logProtocolTagSanitized: mockLogProtocolTagSanitized, // Real ChatCompressionService.compress() calls logChatCompression on // every attempt; the R3.4 integration test exercises that path, so the // mock has to expose it (no-op). @@ -1885,6 +1891,85 @@ describe('GeminiChat', async () => { }, ); + it('sanitizes a standalone closing thinking tag without retrying valid tool calls', async () => { + const create = vi.fn().mockImplementation(async () => + (async function* () { + yield { + id: 'sanitized-protocol-tag', + created: 1, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + reasoning_content: 'hidden reasoning', + content: '\n\n', + tool_calls: [ + { + index: 0, + id: 'call_read', + type: 'function', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }, + finish_reason: 'tool_calls', + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk; + })(), + ); + const provider = { + buildClient: () => + ({ chat: { completions: { create } } }) as unknown as OpenAI, + buildRequest: (request: OpenAI.Chat.ChatCompletionCreateParams) => + request, + buildHeaders: () => ({}), + getDefaultGenerationConfig: () => ({}), + } as OpenAICompatibleProvider; + const generator = new OpenAIContentGenerator( + { model: 'test-model', authType: AuthType.USE_OPENAI }, + mockConfig, + provider, + ); + vi.mocked(mockConfig.getContentGenerator).mockReturnValue(generator); + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + model: 'test-model', + authType: AuthType.USE_OPENAI, + }); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-sanitized-protocol-tag', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + const parts = events.flatMap((event) => + event.type === StreamEventType.CHUNK + ? (event.value.candidates?.[0]?.content?.parts ?? []) + : [], + ); + + expect(create).toHaveBeenCalledTimes(1); + expect(mockLogContentRetry).not.toHaveBeenCalled(); + expect(mockLogProtocolTagSanitized).toHaveBeenCalledTimes(1); + expect(mockLogProtocolTagSanitized).toHaveBeenCalledWith( + mockConfig, + expect.objectContaining({ + model: 'test-model', + prompt_id: 'prompt-id-sanitized-protocol-tag', + response_id: 'sanitized-protocol-tag', + tag_name: 'think', + tool_call_count: 1, + }), + ); + expect(parts).toContainEqual({ + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }); + expect(parts.some((part) => part.text?.includes(''))).toBe(false); + }); + it('falls back to coerced totalTokenCount when promptTokenCount is hostile', async () => { const response = (async function* () { yield { diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 79e9d427106..16e643715a5 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -451,6 +451,252 @@ describe('OpenAIContentConverter', () => { ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); }); + it('sanitizes a standalone closing thinking tag when complete tool calls are available', () => { + const stream = withStreamParser(); + const reasoning = converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + const leakedTag = converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content: '\n\n\n', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + const finish = converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ); + + expect(reasoning.candidates?.[0]?.content?.parts).toEqual([ + { thought: true, text: 'Let me check.' }, + ]); + expect(leakedTag.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + + it('sanitizes a standalone closing thinking tag alias', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content: ' ', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ); + + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'thinking', + toolCallCount: 1, + }); + }); + + it('sanitizes a standalone closing thinking tag split across chunks', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + const firstHalf = converter.convertOpenAIChunkToGemini( + streamChunk('tag-start', { content: '\n\n', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + const finish = converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ); + + expect(firstHalf.candidates?.[0]?.content?.parts).toEqual([]); + expect(secondHalf.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + + it('releases a split tag-like prefix when it becomes ordinary text', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Explain the syntax.' }), + stream, + ); + const prefix = converter.convertOpenAIChunkToGemini( + streamChunk('prefix', { content: ' { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('tag', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + + for (let i = 0; i < 1_000; i++) { + converter.convertOpenAIChunkToGemini( + streamChunk(`whitespace-${i}`, { content: ' ' }), + stream, + ); + } + + expect(stream.pendingThinkingTagCandidate).toEqual({ + text: '', + closingTagName: 'think', + }); + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + + it('rejects a standalone closing thinking tag without a complete tool call', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + const leakedTag = converter.convertOpenAIChunkToGemini( + streamChunk('content', { content: '' }), + stream, + ); + + expect(leakedTag.candidates?.[0]?.content?.parts).toEqual([]); + expect(() => + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'stop'), + stream, + ), + ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + + it('rejects visible content after a deferred closing thinking tag', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('content', { content: '' }), + stream, + ); + + expect(() => + converter.convertOpenAIChunkToGemini( + streamChunk('content-after-tag', { content: 'unexpected' }), + stream, + ), + ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + + it('rejects a standalone closing thinking tag with an incomplete tool call', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { + name: 'read_file', + arguments: '{"path":', + }, + }, + ], + }), + stream, + ); + + expect(() => + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ), + ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + it('rejects a closing tag split after a visible line break', () => { const stream = withStreamParser(); converter.convertOpenAIChunkToGemini( diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index b1c7a81f1c6..5a1c9262b6c 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1092,6 +1092,22 @@ const CLOSING_THINKING_TAG_PATTERN = /\n[^\S\r\n]*<\/think(?:ing)?[^\S\r\n]*>/i; const LEADING_CLOSING_THINKING_TAG_PATTERN = /^[^\S\r\n]*<\/think(?:ing)?[^\S\r\n]*>/i; const LEADING_THINKING_TAG_PATTERN = /^\s*<\/?think(?:ing)?\s*>/i; +const STANDALONE_CLOSING_THINKING_TAG_PATTERN = + /^\s*<\/(think|thinking)\s*>\s*$/i; +const STANDALONE_OPENING_THINKING_TAG_PATTERN = + /^\s*<(think|thinking)\s*>\s*$/i; +const MAX_THINKING_TAG_CANDIDATE_LENGTH = 128; + +function canBeStandaloneThinkingTagPrefix(text: string): boolean { + const candidate = text.trimStart().toLowerCase(); + if (!candidate) return true; + + return [' { + if (tag.startsWith(candidate)) return true; + if (!candidate.startsWith(tag)) return false; + return /^\s*(?:>\s*)?$/.test(candidate.slice(tag.length)); + }); +} /** * Convert OpenAI response to Gemini format. @@ -1379,11 +1395,86 @@ export function convertOpenAIChunkToGemini( } } - const visibleText = parts - .map((part) => - part.thought !== true && typeof part.text === 'string' ? part.text : '', - ) - .join(''); + const getVisibleText = (part: Part): string => + part.thought !== true && typeof part.text === 'string' ? part.text : ''; + let visibleText = parts.map(getVisibleText).join(''); + + const pendingTagCandidate = requestContext.pendingThinkingTagCandidate; + const combinedCandidateText = + (pendingTagCandidate?.text ?? '') + visibleText; + const canStartTagCandidate = + requestContext.hasStructuredReasoningContent === true && + requestContext.hasVisibleContent !== true && + /\S/.test(visibleText) && + canBeStandaloneThinkingTagPrefix(combinedCandidateText); + + if (pendingTagCandidate || canStartTagCandidate) { + const closingTag = STANDALONE_CLOSING_THINKING_TAG_PATTERN.exec( + combinedCandidateText, + )?.[1]?.toLowerCase(); + const closingTagName = + closingTag === 'think' || closingTag === 'thinking' + ? closingTag + : undefined; + const openingTag = STANDALONE_OPENING_THINKING_TAG_PATTERN.test( + combinedCandidateText, + ); + const isPossibleTag = canBeStandaloneThinkingTagPrefix( + combinedCandidateText, + ); + + if (openingTag) { + requestContext.pendingThinkingTagCandidate = undefined; + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + + if (pendingTagCandidate?.closingTagName && !closingTagName) { + requestContext.pendingThinkingTagCandidate = undefined; + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + + if (isPossibleTag) { + if ( + !closingTagName && + combinedCandidateText.length > MAX_THINKING_TAG_CANDIDATE_LENGTH + ) { + requestContext.pendingThinkingTagCandidate = undefined; + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + requestContext.pendingThinkingTagCandidate = closingTagName + ? { text: ``, closingTagName } + : { text: combinedCandidateText }; + parts = parts.filter((part) => !getVisibleText(part)); + visibleText = ''; + + if (choice.finish_reason && !closingTagName) { + requestContext.pendingThinkingTagCandidate = undefined; + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + } else if (pendingTagCandidate) { + parts = parts.filter((part) => !getVisibleText(part)); + parts.push({ text: combinedCandidateText }); + visibleText = combinedCandidateText; + requestContext.pendingThinkingTagCandidate = undefined; + } + } + const leakedThinkingTag = requestContext.hasStructuredReasoningContent === true && ((requestContext.hasVisibleContent !== true && @@ -1392,6 +1483,7 @@ export function convertOpenAIChunkToGemini( (CLOSING_THINKING_TAG_PATTERN.test(visibleText) || (requestContext.atVisibleLineStart === true && LEADING_CLOSING_THINKING_TAG_PATTERN.test(visibleText))))); + if (/\S/.test(visibleText)) { requestContext.hasVisibleContent = true; } @@ -1414,6 +1506,34 @@ export function convertOpenAIChunkToGemini( const completedToolCalls = choice.finish_reason ? toolCallParser.getCompletedToolCalls() : []; + // Some providers report "stop" or "tool_calls" for JSON cut off by the + // token limit, so validate the parser state independently of finish_reason. + const toolCallsTruncated = choice.finish_reason + ? toolCallParser.hasIncompleteToolCalls() + : false; + + if ( + choice.finish_reason && + requestContext.pendingThinkingTagCandidate?.closingTagName + ) { + if ( + completedToolCalls.length === 0 || + toolCallWithoutName || + toolCallsTruncated + ) { + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + requestContext.protocolTagSanitized = { + tagName: requestContext.pendingThinkingTagCandidate.closingTagName, + toolCallCount: completedToolCalls.length, + }; + requestContext.pendingThinkingTagCandidate = undefined; + } + if ( choice.finish_reason && (toolCallWithoutName || @@ -1430,7 +1550,8 @@ export function convertOpenAIChunkToGemini( const shouldHoldParts = !choice.finish_reason && (toolCallWithoutName || - requestContext.hasThinkingTagInReasoning === true); + requestContext.hasThinkingTagInReasoning === true || + requestContext.pendingThinkingTagCandidate !== undefined); if (shouldHoldParts) { (requestContext.pendingUntrustedResponseParts ??= []).push(...parts); parts.length = 0; @@ -1440,13 +1561,7 @@ export function convertOpenAIChunkToGemini( } // Only emit function calls when streaming is complete (finish_reason is present) - let toolCallsTruncated = false; if (choice.finish_reason) { - // Detect truncation the provider may not report correctly. - // Some providers (e.g. DashScope/Qwen) send "stop" or "tool_calls" - // even when output was cut off mid-JSON due to max_tokens. - toolCallsTruncated = toolCallParser.hasIncompleteToolCalls(); - for (const toolCall of completedToolCalls) { if (toolCall.name) { parts.push({ diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 7e5230d6e1c..9c1b0684410 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -27,6 +27,7 @@ import { MAX_STREAM_IDLE_TIMEOUT_MS, QWEN_STREAM_IDLE_TIMEOUT_MS_ENV, } from './constants.js'; +import { logProtocolTagSanitized } from '../../telemetry/loggers.js'; // Mock dependencies vi.mock('./converter.js', () => ({ @@ -38,6 +39,9 @@ vi.mock('./converter.js', () => ({ }, })); vi.mock('openai'); +vi.mock('../../telemetry/loggers.js', () => ({ + logProtocolTagSanitized: vi.fn(), +})); describe('ContentGenerationPipeline', () => { let pipeline: ContentGenerationPipeline; @@ -1652,6 +1656,131 @@ describe('ContentGenerationPipeline', () => { expect(results[0]).toBe(mockValidResponse); }); + it('logs a protocol tag sanitization marker exactly once', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const finishChunk = { + id: 'response-id', + choices: [{ delta: {}, finish_reason: 'tool_calls' }], + } as OpenAI.Chat.ChatCompletionChunk; + const usageChunk = { + id: 'response-id', + choices: [], + usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 }, + } as unknown as OpenAI.Chat.ChatCompletionChunk; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield finishChunk; + yield usageChunk; + }, + }; + const finishResponse = new GenerateContentResponse(); + finishResponse.responseId = 'response-id'; + finishResponse.candidates = [ + { + content: { + parts: [ + { + functionCall: { id: 'call-1', name: 'read_file', args: {} }, + }, + ], + role: 'model', + }, + finishReason: FinishReason.STOP, + }, + ]; + const usageResponse = new GenerateContentResponse(); + usageResponse.responseId = 'response-id'; + usageResponse.candidates = []; + usageResponse.usageMetadata = { + promptTokenCount: 1, + candidatesTokenCount: 2, + totalTokenCount: 3, + }; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock) + .mockImplementationOnce((_chunk, context) => { + context.protocolTagSanitized = { + tagName: 'think', + toolCallCount: 1, + }; + return finishResponse; + }) + .mockReturnValueOnce(usageResponse); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + for await (const _ of resultGenerator) { + // Consume the stream so the trailing usage chunk is processed. + } + + expect(logProtocolTagSanitized).toHaveBeenCalledTimes(1); + expect(logProtocolTagSanitized).toHaveBeenCalledWith( + mockCliConfig, + expect.objectContaining({ + model: 'test-model', + prompt_id: 'test-prompt-id', + response_id: 'response-id', + tag_name: 'think', + handling: 'suppress_standalone_closing_tag', + tool_call_count: 1, + }), + ); + }); + + it('rejects an unresolved thinking-tag candidate at clean stream EOF', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'response-id', + choices: [{ delta: { content: '' }, finish_reason: null }], + } as OpenAI.Chat.ChatCompletionChunk; + }, + }; + const emptyResponse = new GenerateContentResponse(); + emptyResponse.candidates = [ + { content: { parts: [], role: 'model' }, index: 0 }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.pendingThinkingTagCandidate = { + text: '', + closingTagName: 'think', + }; + return emptyResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + + await expect(async () => { + for await (const _ of resultGenerator) { + // Consume until EOF validation runs. + } + }).rejects.toMatchObject({ type: 'PROTOCOL_TAG_LEAK' }); + expect(logProtocolTagSanitized).not.toHaveBeenCalled(); + }); + it('should handle streaming errors and reset tool calls', async () => { // Arrange const request: GenerateContentParameters = { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index f0048e5b8f9..1839397d8e6 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -28,6 +28,8 @@ import { } from './constants.js'; import { createDebugLogger } from '../../utils/debugLogger.js'; import { InvalidStreamError } from '../invalid-stream-error.js'; +import { logProtocolTagSanitized } from '../../telemetry/loggers.js'; +import { ProtocolTagSanitizedEvent } from '../../telemetry/types.js'; const debugLogger = createDebugLogger('OPENAI_PIPELINE'); @@ -499,6 +501,27 @@ export class ContentGenerationPipeline { context, ); + const sanitization = context.protocolTagSanitized; + if (sanitization) { + context.protocolTagSanitized = undefined; + const event = new ProtocolTagSanitizedEvent({ + model: context.model, + promptId: context.userPromptId, + responseId: response.responseId, + tagName: sanitization.tagName, + toolCallCount: sanitization.toolCallCount, + }); + debugLogger.warn('Sanitized a model protocol tag', { + model: event.model, + promptId: event.prompt_id, + responseId: event.response_id, + tagName: event.tag_name, + handling: event.handling, + toolCallCount: event.tool_call_count, + }); + logProtocolTagSanitized(this.config.cliConfig, event); + } + // Stage 2b: Filter empty responses to avoid downstream issues if ( response.candidates?.[0]?.content?.parts?.length === 0 && @@ -552,6 +575,13 @@ export class ContentGenerationPipeline { } } + if (context.pendingThinkingTagCandidate) { + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + // Stage 2d: If there's still a pending finish response at the end // (e.g. no usage chunk arrived after the finish chunk), yield it. if (pendingFinishResponse && !finishYielded) { @@ -941,7 +971,11 @@ export class ContentGenerationPipeline { context: RequestContext, ) => Promise, ): Promise { - const context = this.createRequestContext(request, isStreaming); + const context = this.createRequestContext( + request, + isStreaming, + userPromptId, + ); try { const openaiRequest = await this.buildRequest( @@ -983,6 +1017,7 @@ export class ContentGenerationPipeline { private createRequestContext( request: GenerateContentParameters, isStreaming: boolean, + userPromptId: string, ): RequestContext { const effectiveModel = request.model || this.contentGeneratorConfig.model; const providerOverrides = @@ -999,6 +1034,7 @@ export class ContentGenerationPipeline { return { model: effectiveModel, + userPromptId, modalities: this.contentGeneratorConfig.modalities ?? {}, startTime: Date.now(), splitToolMedia: diff --git a/packages/core/src/core/openaiContentGenerator/types.ts b/packages/core/src/core/openaiContentGenerator/types.ts index 3e638cd4da3..3134ab77867 100644 --- a/packages/core/src/core/openaiContentGenerator/types.ts +++ b/packages/core/src/core/openaiContentGenerator/types.ts @@ -39,6 +39,7 @@ export interface StreamingTextDeltaState { export interface RequestContext { model: string; + userPromptId?: string; modalities: InputModalities; startTime: number; toolCallParser?: StreamingToolCallParser; @@ -84,6 +85,14 @@ export interface RequestContext { hasThinkingTagInReasoning?: boolean; hasVisibleContent?: boolean; atVisibleLineStart?: boolean; + pendingThinkingTagCandidate?: { + text: string; + closingTagName?: 'think' | 'thinking'; + }; + protocolTagSanitized?: { + tagName: 'think' | 'thinking'; + toolCallCount: number; + }; } export interface ErrorHandler { diff --git a/packages/core/src/telemetry/constants.ts b/packages/core/src/telemetry/constants.ts index 52124e48550..b589e29d46e 100644 --- a/packages/core/src/telemetry/constants.ts +++ b/packages/core/src/telemetry/constants.ts @@ -29,6 +29,8 @@ export const EVENT_INVALID_CHUNK = 'qwen-code.chat.invalid_chunk'; export const EVENT_CONTENT_RETRY = 'qwen-code.chat.content_retry'; export const EVENT_CONTENT_RETRY_FAILURE = 'qwen-code.chat.content_retry_failure'; +export const EVENT_PROTOCOL_TAG_SANITIZED = + 'qwen-code.chat.protocol_tag_sanitized'; // Phase 4b — HTTP-status retry telemetry emitted by `retryWithBackoff` for // 429 / 5xx errors at LLM call sites. Distinct from EVENT_CONTENT_RETRY, // which is fired by geminiChat for InvalidStreamError retries on a separate diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts index 8aee18b21f2..abb93b7a0ab 100644 --- a/packages/core/src/telemetry/index.ts +++ b/packages/core/src/telemetry/index.ts @@ -62,6 +62,7 @@ export { logMemoryExtract, logMemoryDream, logMemoryRecall, + logProtocolTagSanitized, } from './loggers.js'; export type { SlashCommandEvent, ChatCompressionEvent } from './types.js'; export { @@ -91,6 +92,7 @@ export { MemoryExtractEvent, MemoryDreamEvent, MemoryRecallEvent, + ProtocolTagSanitizedEvent, } from './types.js'; export { makeSlashCommandEvent, makeChatCompressionEvent } from './types.js'; export type { diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 04176bd2b23..0c55a16a621 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -39,6 +39,7 @@ import { EVENT_EXTENSION_INSTALL, EVENT_EXTENSION_UNINSTALL, EVENT_TOOL_OUTPUT_TRUNCATED, + EVENT_PROTOCOL_TAG_SANITIZED, } from './constants.js'; import { logApiRequest, @@ -60,6 +61,7 @@ import { logHookCall, logApiError, logApiRetry, + logProtocolTagSanitized, } from './loggers.js'; import * as metrics from './metrics.js'; import { apiActivityTracker } from './api-activity-tracker.js'; @@ -87,6 +89,7 @@ import { HookCallEvent, ApiErrorEvent, ApiRetryEvent, + ProtocolTagSanitizedEvent, } from './types.js'; import { FileOperation } from './metrics.js'; import type { @@ -160,6 +163,43 @@ describe('loggers', () => { }); }); + describe('logProtocolTagSanitized', () => { + it('emits a privacy-safe handled event to QwenLogger and OpenTelemetry', () => { + const config = makeFakeConfig({ sessionId: 'test-session-id' }); + vi.spyOn(QwenLogger.prototype, 'logProtocolTagSanitizedEvent'); + const event = new ProtocolTagSanitizedEvent({ + model: 'test-model', + promptId: 'prompt-id', + responseId: 'response-id', + tagName: 'think', + toolCallCount: 2, + }); + + logProtocolTagSanitized(config, event); + + expect( + QwenLogger.prototype.logProtocolTagSanitizedEvent, + ).toHaveBeenCalledWith(event); + expect(mockLogger.emit).toHaveBeenCalledWith({ + body: 'Suppressed a standalone closing think tag and preserved 2 tool call(s).', + attributes: { + 'session.id': 'test-session-id', + 'event.name': EVENT_PROTOCOL_TAG_SANITIZED, + 'event.timestamp': '2025-01-01T00:00:00.000Z', + model: 'test-model', + prompt_id: 'prompt-id', + response_id: 'response-id', + tag_name: 'think', + handling: 'suppress_standalone_closing_tag', + tool_call_count: 2, + }, + }); + expect(JSON.stringify(mockLogger.emit.mock.calls[0])).not.toMatch( + /response_text|reasoning|tool_name|arguments/, + ); + }); + }); + describe('logCliConfiguration', () => { it('should log the cli configuration', () => { const mockConfig = { diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts index a89300edfcd..292d4b99e8a 100644 --- a/packages/core/src/telemetry/loggers.ts +++ b/packages/core/src/telemetry/loggers.ts @@ -30,6 +30,7 @@ import { EVENT_CHAT_COMPRESSION, EVENT_CONTENT_RETRY, EVENT_CONTENT_RETRY_FAILURE, + EVENT_PROTOCOL_TAG_SANITIZED, EVENT_API_RETRY, EVENT_FILE_OPERATION, EVENT_RIPGREP_FALLBACK, @@ -98,6 +99,7 @@ import type { ChatCompressionEvent, ContentRetryEvent, ContentRetryFailureEvent, + ProtocolTagSanitizedEvent, ApiRetryEvent, RipgrepFallbackEvent, ToolOutputTruncatedEvent, @@ -762,6 +764,25 @@ export function logContentRetry( recordContentRetry(config); } +export function logProtocolTagSanitized( + config: Config, + event: ProtocolTagSanitizedEvent, +): void { + QwenLogger.getInstance(config)?.logProtocolTagSanitizedEvent(event); + if (!isTelemetrySdkInitialized()) return; + + const attributes: LogAttributes = { + ...getCommonAttributes(config), + ...event, + 'event.name': EVENT_PROTOCOL_TAG_SANITIZED, + }; + + logs.getLogger(SERVICE_NAME).emit({ + body: `Suppressed a standalone closing ${event.tag_name} tag and preserved ${event.tool_call_count} tool call(s).`, + attributes, + }); +} + export function logContentRetryFailure( config: Config, event: ContentRetryFailureEvent, diff --git a/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts b/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts index 8cc3fc09b64..47e0ec1ac3a 100644 --- a/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts +++ b/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts @@ -25,6 +25,7 @@ import { IdeConnectionType, HookCallEvent, SkillLaunchEvent, + ProtocolTagSanitizedEvent, } from '../types.js'; import type { RumEvent, RumPayload } from './event-types.js'; @@ -359,6 +360,39 @@ describe('QwenLogger', () => { }); describe('event handlers', () => { + it('logs protocol tag sanitization without model content', () => { + const logger = QwenLogger.getInstance(mockConfig)!; + const enqueueSpy = vi.spyOn(logger, 'enqueueLogEvent'); + const event = new ProtocolTagSanitizedEvent({ + model: 'test-model', + promptId: 'prompt-id', + responseId: 'response-id', + tagName: 'thinking', + toolCallCount: 3, + }); + + logger.logProtocolTagSanitizedEvent(event); + + expect(enqueueSpy).toHaveBeenCalledWith( + expect.objectContaining({ + event_type: 'action', + type: 'misc', + name: 'protocol_tag_sanitized', + properties: { + model: 'test-model', + prompt_id: 'prompt-id', + response_id: 'response-id', + tag_name: 'thinking', + handling: 'suppress_standalone_closing_tag', + tool_call_count: 3, + }, + }), + ); + expect(JSON.stringify(enqueueSpy.mock.calls[0])).not.toMatch( + /response_text|reasoning|tool_name|arguments/, + ); + }); + it('should log IDE connection events', () => { const logger = QwenLogger.getInstance(mockConfig)!; const enqueueSpy = vi.spyOn(logger, 'enqueueLogEvent'); diff --git a/packages/core/src/telemetry/qwen-logger/qwen-logger.ts b/packages/core/src/telemetry/qwen-logger/qwen-logger.ts index 15f9d11bcc0..5b7cd8a4f0f 100644 --- a/packages/core/src/telemetry/qwen-logger/qwen-logger.ts +++ b/packages/core/src/telemetry/qwen-logger/qwen-logger.ts @@ -31,6 +31,7 @@ import type { ChatCompressionEvent, InvalidChunkEvent, ContentRetryEvent, + ProtocolTagSanitizedEvent, ApiRetryEvent, ContentRetryFailureEvent, ConversationFinishedEvent, @@ -961,6 +962,22 @@ export class QwenLogger { this.flushIfNeeded(); } + logProtocolTagSanitizedEvent(event: ProtocolTagSanitizedEvent): void { + const rumEvent = this.createActionEvent('misc', 'protocol_tag_sanitized', { + properties: { + model: event.model, + prompt_id: event.prompt_id ?? '', + response_id: event.response_id ?? '', + tag_name: event.tag_name, + handling: event.handling, + tool_call_count: event.tool_call_count, + }, + }); + + this.enqueueLogEvent(rumEvent); + this.flushIfNeeded(); + } + // Phase 4b — HTTP-status retry from retryWithBackoff (429/5xx). Distinct from // logContentRetryEvent which is fired by geminiChat's content-recovery loop. logApiRetryEvent(event: ApiRetryEvent): void { diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index dea9cea23c0..ba385d17e52 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -671,6 +671,34 @@ export class ContentRetryEvent implements BaseTelemetryEvent { } } +export class ProtocolTagSanitizedEvent implements BaseTelemetryEvent { + 'event.name': 'protocol_tag_sanitized'; + 'event.timestamp': string; + model: string; + prompt_id?: string; + response_id?: string; + tag_name: 'think' | 'thinking'; + handling: 'suppress_standalone_closing_tag'; + tool_call_count: number; + + constructor(opts: { + model: string; + promptId?: string; + responseId?: string; + tagName: 'think' | 'thinking'; + toolCallCount: number; + }) { + this['event.name'] = 'protocol_tag_sanitized'; + this['event.timestamp'] = new Date().toISOString(); + this.model = opts.model; + this.prompt_id = opts.promptId; + this.response_id = opts.responseId; + this.tag_name = opts.tagName; + this.handling = 'suppress_standalone_closing_tag'; + this.tool_call_count = opts.toolCallCount; + } +} + /** * Phase 4b — HTTP-status retry telemetry. Emitted by `retryWithBackoff` (via * the `onRetry` callback opt-in) for HTTP 429 / 5xx retries at LLM call sites. @@ -1075,6 +1103,7 @@ export type TelemetryEvent = | FileOperationEvent | InvalidChunkEvent | ContentRetryEvent + | ProtocolTagSanitizedEvent | ContentRetryFailureEvent | ApiRetryEvent | SubagentExecutionEvent From 516c817e0c9f0abf6bca6cef293fff58222f8bb6 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 13:23:27 +0800 Subject: [PATCH 2/8] fix(core): validate sanitized tool calls --- .../openaiContentGenerator/converter.test.ts | 35 ++++++++++++ .../core/openaiContentGenerator/converter.ts | 53 ++++++------------- .../core/openaiContentGenerator/pipeline.ts | 12 ++--- .../streamingToolCallParser.test.ts | 18 +++++++ .../streamingToolCallParser.ts | 16 ++++++ .../src/core/openaiContentGenerator/types.ts | 1 - packages/core/src/telemetry/index.ts | 2 - 7 files changed, 90 insertions(+), 47 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 16e643715a5..29cfa905881 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -697,6 +697,41 @@ describe('OpenAIContentConverter', () => { expect(stream.protocolTagSanitized).toBeUndefined(); }); + it.each(['{bad}', 'null', '[]', '42'])( + 'rejects a standalone closing thinking tag with invalid tool arguments %s', + (toolArguments) => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { + name: 'read_file', + arguments: toolArguments, + }, + }, + ], + }), + stream, + ); + + expect(() => + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, 'tool_calls'), + stream, + ), + ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(stream.protocolTagSanitized).toBeUndefined(); + }, + ); + it('rejects a closing tag split after a visible line break', () => { const stream = withStreamParser(); converter.convertOpenAIChunkToGemini( diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index 5a1c9262b6c..1ef3c98c2fc 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1109,6 +1109,15 @@ function canBeStandaloneThinkingTagPrefix(text: string): boolean { }); } +function throwProtocolTagLeak(requestContext: RequestContext): never { + requestContext.pendingThinkingTagCandidate = undefined; + requestContext.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); +} + /** * Convert OpenAI response to Gemini format. */ @@ -1424,21 +1433,11 @@ export function convertOpenAIChunkToGemini( ); if (openingTag) { - requestContext.pendingThinkingTagCandidate = undefined; - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } if (pendingTagCandidate?.closingTagName && !closingTagName) { - requestContext.pendingThinkingTagCandidate = undefined; - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } if (isPossibleTag) { @@ -1446,12 +1445,7 @@ export function convertOpenAIChunkToGemini( !closingTagName && combinedCandidateText.length > MAX_THINKING_TAG_CANDIDATE_LENGTH ) { - requestContext.pendingThinkingTagCandidate = undefined; - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } requestContext.pendingThinkingTagCandidate = closingTagName ? { text: ``, closingTagName } @@ -1460,12 +1454,7 @@ export function convertOpenAIChunkToGemini( visibleText = ''; if (choice.finish_reason && !closingTagName) { - requestContext.pendingThinkingTagCandidate = undefined; - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } } else if (pendingTagCandidate) { parts = parts.filter((part) => !getVisibleText(part)); @@ -1495,11 +1484,7 @@ export function convertOpenAIChunkToGemini( /^[^\S\r\n]*$/.test(lineSuffix); } if (leakedThinkingTag) { - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } const toolCallWithoutName = toolCallParser.hasNamelessToolCall(); @@ -1511,7 +1496,6 @@ export function convertOpenAIChunkToGemini( const toolCallsTruncated = choice.finish_reason ? toolCallParser.hasIncompleteToolCalls() : false; - if ( choice.finish_reason && requestContext.pendingThinkingTagCandidate?.closingTagName @@ -1519,13 +1503,10 @@ export function convertOpenAIChunkToGemini( if ( completedToolCalls.length === 0 || toolCallWithoutName || - toolCallsTruncated + toolCallsTruncated || + toolCallParser.hasInvalidToolCallArguments() ) { - requestContext.pendingUntrustedResponseParts = undefined; - throw new InvalidStreamError( - 'Model response leaked thinking tags.', - 'PROTOCOL_TAG_LEAK', - ); + throwProtocolTagLeak(requestContext); } requestContext.protocolTagSanitized = { tagName: requestContext.pendingThinkingTagCandidate.closingTagName, diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index 1839397d8e6..1af443c4f42 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -443,6 +443,7 @@ export class ContentGenerationPipeline { guarded, context, request, + userPromptId, ); async function* drainThenCleanup(): AsyncGenerator { try { @@ -469,6 +470,7 @@ export class ContentGenerationPipeline { stream: AsyncIterable, context: RequestContext, request: GenerateContentParameters, + userPromptId: string, ): AsyncGenerator { const collectedGeminiResponses: GenerateContentResponse[] = []; @@ -506,7 +508,7 @@ export class ContentGenerationPipeline { context.protocolTagSanitized = undefined; const event = new ProtocolTagSanitizedEvent({ model: context.model, - promptId: context.userPromptId, + promptId: userPromptId, responseId: response.responseId, tagName: sanitization.tagName, toolCallCount: sanitization.toolCallCount, @@ -971,11 +973,7 @@ export class ContentGenerationPipeline { context: RequestContext, ) => Promise, ): Promise { - const context = this.createRequestContext( - request, - isStreaming, - userPromptId, - ); + const context = this.createRequestContext(request, isStreaming); try { const openaiRequest = await this.buildRequest( @@ -1017,7 +1015,6 @@ export class ContentGenerationPipeline { private createRequestContext( request: GenerateContentParameters, isStreaming: boolean, - userPromptId: string, ): RequestContext { const effectiveModel = request.model || this.contentGeneratorConfig.model; const providerOverrides = @@ -1034,7 +1031,6 @@ export class ContentGenerationPipeline { return { model: effectiveModel, - userPromptId, modalities: this.contentGeneratorConfig.modalities ?? {}, startTime: Date.now(), splitToolMedia: diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts index f9485a2d9d5..a90e2af73d7 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts @@ -1057,4 +1057,22 @@ describe('StreamingToolCallParser', () => { expect(parser.getState(0).depth).toBe(1); }); }); + + describe('hasInvalidToolCallArguments', () => { + it('accepts empty and object arguments', () => { + parser.addChunk(0, '', 'call_empty', 'list_sessions'); + parser.addChunk(1, '{"path":"a.ts"}', 'call_object', 'read_file'); + + expect(parser.hasInvalidToolCallArguments()).toBe(false); + }); + + it.each(['{bad}', 'null', '[]', '42'])( + 'rejects invalid or non-object arguments %s', + (toolArguments) => { + parser.addChunk(0, toolArguments, 'call_1', 'read_file'); + + expect(parser.hasInvalidToolCallArguments()).toBe(true); + }, + ); + }); }); diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index 2c962397b89..39f21635d5f 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -276,6 +276,22 @@ export class StreamingToolCallParser { return this.namelessToolCallIndices.size > 0; } + hasInvalidToolCallArguments(): boolean { + for (const [index, buffer] of this.buffers.entries()) { + if (!this.toolCallMeta.get(index)?.name || !buffer.trim()) continue; + + try { + const args: unknown = JSON.parse(buffer); + if (typeof args !== 'object' || args === null || Array.isArray(args)) { + return true; + } + } catch { + return true; + } + } + return false; + } + /** * Gets all completed tool calls that are ready to be emitted * diff --git a/packages/core/src/core/openaiContentGenerator/types.ts b/packages/core/src/core/openaiContentGenerator/types.ts index 3134ab77867..4278b081977 100644 --- a/packages/core/src/core/openaiContentGenerator/types.ts +++ b/packages/core/src/core/openaiContentGenerator/types.ts @@ -39,7 +39,6 @@ export interface StreamingTextDeltaState { export interface RequestContext { model: string; - userPromptId?: string; modalities: InputModalities; startTime: number; toolCallParser?: StreamingToolCallParser; diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts index abb93b7a0ab..8aee18b21f2 100644 --- a/packages/core/src/telemetry/index.ts +++ b/packages/core/src/telemetry/index.ts @@ -62,7 +62,6 @@ export { logMemoryExtract, logMemoryDream, logMemoryRecall, - logProtocolTagSanitized, } from './loggers.js'; export type { SlashCommandEvent, ChatCompressionEvent } from './types.js'; export { @@ -92,7 +91,6 @@ export { MemoryExtractEvent, MemoryDreamEvent, MemoryRecallEvent, - ProtocolTagSanitizedEvent, } from './types.js'; export { makeSlashCommandEvent, makeChatCompressionEvent } from './types.js'; export type { From 62e133ad42dcc00b717edfd33558b66444e7fafb Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 14:04:12 +0800 Subject: [PATCH 3/8] refactor(core): simplify protocol tag sanitization --- .../openaiContentGenerator/converter.test.ts | 260 ++++++------------ .../openaiContentGenerator/pipeline.test.ts | 80 ------ .../core/openaiContentGenerator/pipeline.ts | 1 - .../streamingToolCallParser.test.ts | 24 +- packages/core/src/telemetry/loggers.test.ts | 1 - .../telemetry/qwen-logger/qwen-logger.test.ts | 4 - .../src/telemetry/qwen-logger/qwen-logger.ts | 1 - packages/core/src/telemetry/types.ts | 2 - 8 files changed, 87 insertions(+), 286 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 29cfa905881..21479a40d8a 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -108,6 +108,40 @@ describe('OpenAIContentConverter', () => { choices: [{ index: 0, delta, finish_reason: finishReason }], }) as unknown as OpenAI.Chat.ChatCompletionChunk; + const emitReasoning = (stream: RequestContext) => + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { reasoning_content: 'Let me check.' }), + stream, + ); + + const emitToolCall = ( + stream: RequestContext, + content: string, + toolArguments = '{}', + ) => + converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content, + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: toolArguments }, + }, + ], + }), + stream, + ); + + const finishStream = ( + stream: RequestContext, + finishReason = 'tool_calls', + ) => + converter.convertOpenAIChunkToGemini( + streamChunk('finish', {}, finishReason), + stream, + ); + it('creates fresh parser instances', () => { const ctx1 = new StreamingToolCallParser(); const ctx2 = new StreamingToolCallParser(); @@ -451,102 +485,42 @@ describe('OpenAIContentConverter', () => { ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); }); - it('sanitizes a standalone closing thinking tag when complete tool calls are available', () => { - const stream = withStreamParser(); - const reasoning = converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); - const leakedTag = converter.convertOpenAIChunkToGemini( - streamChunk('tool-call', { - content: '\n\n\n', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { name: 'read_file', arguments: '{}' }, - }, - ], - }), - stream, - ); - const finish = converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ); - - expect(reasoning.candidates?.[0]?.content?.parts).toEqual([ - { thought: true, text: 'Let me check.' }, - ]); - expect(leakedTag.candidates?.[0]?.content?.parts).toEqual([]); - expect(finish.candidates?.[0]?.content?.parts).toEqual([ - { - functionCall: { id: 'call_read', name: 'read_file', args: {} }, - }, - ]); - expect(stream.protocolTagSanitized).toEqual({ - tagName: 'think', - toolCallCount: 1, - }); - }); - - it('sanitizes a standalone closing thinking tag alias', () => { - const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('tool-call', { - content: ' ', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { name: 'read_file', arguments: '{}' }, - }, - ], - }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ); + it.each([ + ['think', '\n\n\n'], + ['thinking', ' '], + ] as const)( + 'sanitizes a standalone closing %s tag with complete tool calls', + (tagName, tag) => { + const stream = withStreamParser(); + const reasoning = emitReasoning(stream); + const leakedTag = emitToolCall(stream, tag); + const finish = finishStream(stream); - expect(stream.protocolTagSanitized).toEqual({ - tagName: 'thinking', - toolCallCount: 1, - }); - }); + expect(reasoning.candidates?.[0]?.content?.parts).toEqual([ + { thought: true, text: 'Let me check.' }, + ]); + expect(leakedTag.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName, + toolCallCount: 1, + }); + }, + ); it('sanitizes a standalone closing thinking tag split across chunks', () => { const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); + emitReasoning(stream); const firstHalf = converter.convertOpenAIChunkToGemini( streamChunk('tag-start', { content: '\n\n', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { name: 'read_file', arguments: '{}' }, - }, - ], - }), - stream, - ); - const finish = converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ); + const secondHalf = emitToolCall(stream, 'nk>\n'); + const finish = finishStream(stream); expect(firstHalf.candidates?.[0]?.content?.parts).toEqual([]); expect(secondHalf.candidates?.[0]?.content?.parts).toEqual([]); @@ -585,23 +559,8 @@ describe('OpenAIContentConverter', () => { it('does not accumulate whitespace after a complete closing tag candidate', () => { const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('tag', { - content: '', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { name: 'read_file', arguments: '{}' }, - }, - ], - }), - stream, - ); + emitReasoning(stream); + emitToolCall(stream, ''); for (let i = 0; i < 1_000; i++) { converter.convertOpenAIChunkToGemini( @@ -614,10 +573,7 @@ describe('OpenAIContentConverter', () => { text: '', closingTagName: 'think', }); - converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ); + finishStream(stream); expect(stream.protocolTagSanitized).toEqual({ tagName: 'think', toolCallCount: 1, @@ -626,31 +582,22 @@ describe('OpenAIContentConverter', () => { it('rejects a standalone closing thinking tag without a complete tool call', () => { const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); + emitReasoning(stream); const leakedTag = converter.convertOpenAIChunkToGemini( streamChunk('content', { content: '' }), stream, ); expect(leakedTag.candidates?.[0]?.content?.parts).toEqual([]); - expect(() => - converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'stop'), - stream, - ), - ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(() => finishStream(stream, 'stop')).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); expect(stream.protocolTagSanitized).toBeUndefined(); }); it('rejects visible content after a deferred closing thinking tag', () => { const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); + emitReasoning(stream); converter.convertOpenAIChunkToGemini( streamChunk('content', { content: '' }), stream, @@ -665,69 +612,16 @@ describe('OpenAIContentConverter', () => { expect(stream.protocolTagSanitized).toBeUndefined(); }); - it('rejects a standalone closing thinking tag with an incomplete tool call', () => { - const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('tool-call', { - content: '', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { - name: 'read_file', - arguments: '{"path":', - }, - }, - ], - }), - stream, - ); - - expect(() => - converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ), - ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); - expect(stream.protocolTagSanitized).toBeUndefined(); - }); - - it.each(['{bad}', 'null', '[]', '42'])( - 'rejects a standalone closing thinking tag with invalid tool arguments %s', + it.each(['{"path":', '{bad}', 'null', '[]', '42'])( + 'rejects a standalone closing thinking tag with unsafe tool arguments %s', (toolArguments) => { const stream = withStreamParser(); - converter.convertOpenAIChunkToGemini( - streamChunk('reasoning', { reasoning_content: 'Let me check.' }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('tool-call', { - content: '', - tool_calls: [ - { - index: 0, - id: 'call_read', - function: { - name: 'read_file', - arguments: toolArguments, - }, - }, - ], - }), - stream, - ); + emitReasoning(stream); + emitToolCall(stream, '', toolArguments); - expect(() => - converter.convertOpenAIChunkToGemini( - streamChunk('finish', {}, 'tool_calls'), - stream, - ), - ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); expect(stream.protocolTagSanitized).toBeUndefined(); }, ); diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 9c1b0684410..03086877303 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -1656,86 +1656,6 @@ describe('ContentGenerationPipeline', () => { expect(results[0]).toBe(mockValidResponse); }); - it('logs a protocol tag sanitization marker exactly once', async () => { - const request: GenerateContentParameters = { - model: 'test-model', - contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], - }; - const finishChunk = { - id: 'response-id', - choices: [{ delta: {}, finish_reason: 'tool_calls' }], - } as OpenAI.Chat.ChatCompletionChunk; - const usageChunk = { - id: 'response-id', - choices: [], - usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 }, - } as unknown as OpenAI.Chat.ChatCompletionChunk; - const mockStream = { - async *[Symbol.asyncIterator]() { - yield finishChunk; - yield usageChunk; - }, - }; - const finishResponse = new GenerateContentResponse(); - finishResponse.responseId = 'response-id'; - finishResponse.candidates = [ - { - content: { - parts: [ - { - functionCall: { id: 'call-1', name: 'read_file', args: {} }, - }, - ], - role: 'model', - }, - finishReason: FinishReason.STOP, - }, - ]; - const usageResponse = new GenerateContentResponse(); - usageResponse.responseId = 'response-id'; - usageResponse.candidates = []; - usageResponse.usageMetadata = { - promptTokenCount: 1, - candidatesTokenCount: 2, - totalTokenCount: 3, - }; - - (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); - (mockConverter.convertOpenAIChunkToGemini as Mock) - .mockImplementationOnce((_chunk, context) => { - context.protocolTagSanitized = { - tagName: 'think', - toolCallCount: 1, - }; - return finishResponse; - }) - .mockReturnValueOnce(usageResponse); - (mockClient.chat.completions.create as Mock).mockResolvedValue( - mockStream, - ); - - const resultGenerator = await pipeline.executeStream( - request, - 'test-prompt-id', - ); - for await (const _ of resultGenerator) { - // Consume the stream so the trailing usage chunk is processed. - } - - expect(logProtocolTagSanitized).toHaveBeenCalledTimes(1); - expect(logProtocolTagSanitized).toHaveBeenCalledWith( - mockCliConfig, - expect.objectContaining({ - model: 'test-model', - prompt_id: 'test-prompt-id', - response_id: 'response-id', - tag_name: 'think', - handling: 'suppress_standalone_closing_tag', - tool_call_count: 1, - }), - ); - }); - it('rejects an unresolved thinking-tag candidate at clean stream EOF', async () => { const request: GenerateContentParameters = { model: 'test-model', diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index 1af443c4f42..e9b6b9259d2 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -518,7 +518,6 @@ export class ContentGenerationPipeline { promptId: event.prompt_id, responseId: event.response_id, tagName: event.tag_name, - handling: event.handling, toolCallCount: event.tool_call_count, }); logProtocolTagSanitized(this.config.cliConfig, event); diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts index a90e2af73d7..630145f567d 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts @@ -1059,20 +1059,16 @@ describe('StreamingToolCallParser', () => { }); describe('hasInvalidToolCallArguments', () => { - it('accepts empty and object arguments', () => { - parser.addChunk(0, '', 'call_empty', 'list_sessions'); - parser.addChunk(1, '{"path":"a.ts"}', 'call_object', 'read_file'); - - expect(parser.hasInvalidToolCallArguments()).toBe(false); + it.each([ + ['', false], + ['{"path":"a.ts"}', false], + ['{bad}', true], + ['null', true], + ['[]', true], + ['42', true], + ])('validates %s', (toolArguments, invalid) => { + parser.addChunk(0, toolArguments, 'call_1', 'read_file'); + expect(parser.hasInvalidToolCallArguments()).toBe(invalid); }); - - it.each(['{bad}', 'null', '[]', '42'])( - 'rejects invalid or non-object arguments %s', - (toolArguments) => { - parser.addChunk(0, toolArguments, 'call_1', 'read_file'); - - expect(parser.hasInvalidToolCallArguments()).toBe(true); - }, - ); }); }); diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 0c55a16a621..65201327908 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -190,7 +190,6 @@ describe('loggers', () => { prompt_id: 'prompt-id', response_id: 'response-id', tag_name: 'think', - handling: 'suppress_standalone_closing_tag', tool_call_count: 2, }, }); diff --git a/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts b/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts index 47e0ec1ac3a..819dd70bb51 100644 --- a/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts +++ b/packages/core/src/telemetry/qwen-logger/qwen-logger.test.ts @@ -383,14 +383,10 @@ describe('QwenLogger', () => { prompt_id: 'prompt-id', response_id: 'response-id', tag_name: 'thinking', - handling: 'suppress_standalone_closing_tag', tool_call_count: 3, }, }), ); - expect(JSON.stringify(enqueueSpy.mock.calls[0])).not.toMatch( - /response_text|reasoning|tool_name|arguments/, - ); }); it('should log IDE connection events', () => { diff --git a/packages/core/src/telemetry/qwen-logger/qwen-logger.ts b/packages/core/src/telemetry/qwen-logger/qwen-logger.ts index 5b7cd8a4f0f..aec004b8817 100644 --- a/packages/core/src/telemetry/qwen-logger/qwen-logger.ts +++ b/packages/core/src/telemetry/qwen-logger/qwen-logger.ts @@ -969,7 +969,6 @@ export class QwenLogger { prompt_id: event.prompt_id ?? '', response_id: event.response_id ?? '', tag_name: event.tag_name, - handling: event.handling, tool_call_count: event.tool_call_count, }, }); diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index ba385d17e52..881d84ccdd0 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -678,7 +678,6 @@ export class ProtocolTagSanitizedEvent implements BaseTelemetryEvent { prompt_id?: string; response_id?: string; tag_name: 'think' | 'thinking'; - handling: 'suppress_standalone_closing_tag'; tool_call_count: number; constructor(opts: { @@ -694,7 +693,6 @@ export class ProtocolTagSanitizedEvent implements BaseTelemetryEvent { this.prompt_id = opts.promptId; this.response_id = opts.responseId; this.tag_name = opts.tagName; - this.handling = 'suppress_standalone_closing_tag'; this.tool_call_count = opts.toolCallCount; } } From cab608fcf099c74487564176e700350c7accefcd Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 14:33:41 +0800 Subject: [PATCH 4/8] fix(core): restrict protocol tag recovery finish --- packages/core/src/core/geminiChat.test.ts | 20 ++++++++++++++++++- .../openaiContentGenerator/converter.test.ts | 14 +++++++++++++ .../core/openaiContentGenerator/converter.ts | 1 + 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 100f0319616..c3d3c3ed553 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -1892,6 +1892,17 @@ describe('GeminiChat', async () => { ); it('sanitizes a standalone closing thinking tag without retrying valid tool calls', async () => { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); const create = vi.fn().mockImplementation(async () => (async function* () { yield { @@ -1938,7 +1949,7 @@ describe('GeminiChat', async () => { authType: AuthType.USE_OPENAI, }); - const stream = await chat.sendMessageStream( + const stream = await chatWithRecording.sendMessageStream( 'test-model', { message: 'test' }, 'prompt-id-sanitized-protocol-tag', @@ -1968,6 +1979,13 @@ describe('GeminiChat', async () => { functionCall: { id: 'call_read', name: 'read_file', args: {} }, }); expect(parts.some((part) => part.text?.includes(''))).toBe(false); + expect(JSON.stringify(chatWithRecording.getHistory())).not.toContain( + '', + ); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect( + JSON.stringify(recordAssistantTurn.mock.calls[0]?.[0].message), + ).not.toContain(''); }); it('falls back to coerced totalTokenCount when promptTokenCount is hostile', async () => { diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 21479a40d8a..63d0297cb20 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -595,6 +595,20 @@ describe('OpenAIContentConverter', () => { expect(stream.protocolTagSanitized).toBeUndefined(); }); + it.each(['stop', 'length', 'content_filter', 'unknown'])( + 'rejects recovery when the stream finishes with %s', + (finishReason) => { + const stream = withStreamParser(); + emitReasoning(stream); + emitToolCall(stream, ''); + + expect(() => finishStream(stream, finishReason)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); + expect(stream.protocolTagSanitized).toBeUndefined(); + }, + ); + it('rejects visible content after a deferred closing thinking tag', () => { const stream = withStreamParser(); emitReasoning(stream); diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index 1ef3c98c2fc..fb6b441ca90 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1501,6 +1501,7 @@ export function convertOpenAIChunkToGemini( requestContext.pendingThinkingTagCandidate?.closingTagName ) { if ( + choice.finish_reason !== 'tool_calls' || completedToolCalls.length === 0 || toolCallWithoutName || toolCallsTruncated || From 5d7a00dfe64ac0ad0f568bc14bb3d22492513af4 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 20:58:37 +0800 Subject: [PATCH 5/8] fix(core): close protocol tag recovery gaps --- .../openaiContentGenerator/converter.test.ts | 140 ++++++++++++++++++ .../core/openaiContentGenerator/converter.ts | 9 +- .../openaiContentGenerator/pipeline.test.ts | 52 +++++++ .../core/openaiContentGenerator/pipeline.ts | 44 ++++-- .../streamingToolCallParser.ts | 20 +-- 5 files changed, 240 insertions(+), 25 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 63d0297cb20..cccd2b365a9 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -467,6 +467,29 @@ describe('OpenAIContentConverter', () => { ).toThrowError(expect.objectContaining({ type: 'MALFORMED_TOOL_CALL' })); }); + it('rejects a protocol-tag recovery with a whitespace-only function name', () => { + const stream = withStreamParser(); + emitReasoning(stream); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-call', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_blank', + function: { name: ' ', arguments: '{}' }, + }, + ], + }), + stream, + ); + + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + it('rejects the recorded cross-channel thinking-tag leak', () => { const stream = withStreamParser(); const reasoning = converter.convertOpenAIChunkToGemini( @@ -535,6 +558,123 @@ describe('OpenAIContentConverter', () => { }); }); + it('sanitizes a standalone closing thinking tag with multiple tool calls', () => { + const stream = withStreamParser(); + emitReasoning(stream); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-calls', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + { + index: 1, + id: 'call_list', + function: { name: 'list_directory', arguments: '{}' }, + }, + ], + }), + stream, + ); + const finish = finishStream(stream); + + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + { + functionCall: { id: 'call_list', name: 'list_directory', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 2, + }); + }); + + it('rejects a protocol-tag recovery when a new id relabels nameless arguments', () => { + const stream = withStreamParser(); + emitReasoning(stream); + converter.convertOpenAIChunkToGemini( + streamChunk('old-arguments', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_old', + function: { arguments: '{"path":"old.txt"}' }, + }, + ], + }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('new-name', { + tool_calls: [ + { + index: 0, + id: 'call_new', + function: { name: 'read_file' }, + }, + ], + }), + stream, + ); + + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + + it('buffers leading whitespace before a split standalone closing tag', () => { + const stream = withStreamParser(); + emitReasoning(stream); + const whitespace = converter.convertOpenAIChunkToGemini( + streamChunk('whitespace', { content: '\n' }), + stream, + ); + const tag = emitToolCall(stream, ''); + const finish = finishStream(stream); + + expect(whitespace.candidates?.[0]?.content?.parts).toEqual([]); + expect(tag.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + + it('ignores an exact cumulative replay of a deferred closing tag', () => { + const stream = withStreamParser(); + emitReasoning(stream); + emitToolCall(stream, ''); + const replay = converter.convertOpenAIChunkToGemini( + streamChunk('replay', { content: '' }), + stream, + ); + const finish = finishStream(stream); + + expect(replay.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + it('releases a split tag-like prefix when it becomes ordinary text', () => { const stream = withStreamParser(); converter.convertOpenAIChunkToGemini( diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index fb6b441ca90..d852b94c404 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1409,12 +1409,19 @@ export function convertOpenAIChunkToGemini( let visibleText = parts.map(getVisibleText).join(''); const pendingTagCandidate = requestContext.pendingThinkingTagCandidate; + if ( + pendingTagCandidate?.closingTagName && + visibleText.trim() === pendingTagCandidate.text + ) { + parts = parts.filter((part) => !getVisibleText(part)); + visibleText = ''; + } const combinedCandidateText = (pendingTagCandidate?.text ?? '') + visibleText; const canStartTagCandidate = requestContext.hasStructuredReasoningContent === true && requestContext.hasVisibleContent !== true && - /\S/.test(visibleText) && + visibleText.length > 0 && canBeStandaloneThinkingTagPrefix(combinedCandidateText); if (pendingTagCandidate || canStartTagCandidate) { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 03086877303..10a45003bd5 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -1701,6 +1701,58 @@ describe('ContentGenerationPipeline', () => { expect(logProtocolTagSanitized).not.toHaveBeenCalled(); }); + it('does not log protocol-tag sanitization before a held finish is yielded', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const streamError = new Error('stream failed after finish'); + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'finish-chunk', + choices: [{ delta: {}, finish_reason: 'tool_calls' }], + } as OpenAI.Chat.ChatCompletionChunk; + throw streamError; + }, + }; + const finishResponse = new GenerateContentResponse(); + finishResponse.responseId = 'finish-response'; + finishResponse.candidates = [ + { + content: { parts: [{ functionCall: { name: 'read_file' } }] }, + finishReason: FinishReason.STOP, + index: 0, + }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.protocolTagSanitized = { + tagName: 'think', + toolCallCount: 1, + }; + return finishResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + + await expect(async () => { + for await (const _ of resultGenerator) { + // Consume until the stream error after the held finish. + } + }).rejects.toThrow(streamError); + expect(logProtocolTagSanitized).not.toHaveBeenCalled(); + }); + it('should handle streaming errors and reset tool calls', async () => { // Arrange const request: GenerateContentParameters = { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index e9b6b9259d2..231d80274d5 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -483,6 +483,30 @@ export class ContentGenerationPipeline { // function-call parts from the finish chunk). let pendingFinishResponse: GenerateContentResponse | null = null; let finishYielded = false; + let pendingProtocolTagSanitized: + | NonNullable + | undefined; + const logPendingProtocolTagSanitized = ( + response: GenerateContentResponse, + ) => { + if (!pendingProtocolTagSanitized) return; + const event = new ProtocolTagSanitizedEvent({ + model: context.model, + promptId: userPromptId, + responseId: response.responseId, + tagName: pendingProtocolTagSanitized.tagName, + toolCallCount: pendingProtocolTagSanitized.toolCallCount, + }); + pendingProtocolTagSanitized = undefined; + debugLogger.warn('Sanitized a model protocol tag', { + model: event.model, + promptId: event.prompt_id, + responseId: event.response_id, + tagName: event.tag_name, + toolCallCount: event.tool_call_count, + }); + logProtocolTagSanitized(this.config.cliConfig, event); + }; try { // Stage 2a: Convert and yield each chunk while preserving original @@ -505,22 +529,8 @@ export class ContentGenerationPipeline { const sanitization = context.protocolTagSanitized; if (sanitization) { + pendingProtocolTagSanitized = sanitization; context.protocolTagSanitized = undefined; - const event = new ProtocolTagSanitizedEvent({ - model: context.model, - promptId: userPromptId, - responseId: response.responseId, - tagName: sanitization.tagName, - toolCallCount: sanitization.toolCallCount, - }); - debugLogger.warn('Sanitized a model protocol tag', { - model: event.model, - promptId: event.prompt_id, - responseId: event.response_id, - tagName: event.tag_name, - toolCallCount: event.tool_call_count, - }); - logProtocolTagSanitized(this.config.cliConfig, event); } // Stage 2b: Filter empty responses to avoid downstream issues @@ -529,6 +539,7 @@ export class ContentGenerationPipeline { !response.candidates?.[0]?.finishReason && !response.usageMetadata ) { + pendingProtocolTagSanitized = undefined; continue; } @@ -566,11 +577,13 @@ export class ContentGenerationPipeline { if (shouldYield) { // If we have a pending finish response, yield it instead if (pendingFinishResponse) { + logPendingProtocolTagSanitized(pendingFinishResponse); yield pendingFinishResponse; finishYielded = true; // Keep pendingFinishResponse alive so late-arriving usage // metadata can still be merged (see finishYielded block above). } else { + logPendingProtocolTagSanitized(response); yield response; } } @@ -586,6 +599,7 @@ export class ContentGenerationPipeline { // Stage 2d: If there's still a pending finish response at the end // (e.g. no usage chunk arrived after the finish chunk), yield it. if (pendingFinishResponse && !finishYielded) { + logPendingProtocolTagSanitized(pendingFinishResponse); yield pendingFinishResponse; } } catch (error) { diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index 39f21635d5f..c9013becd63 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -72,7 +72,8 @@ export class StreamingToolCallParser { id?: string, name?: string, ): ToolCallParseResult { - if (!id && !name && !chunk.trim()) { + const validName = name?.trim() ? name : undefined; + if (!id && !validName && !chunk.trim()) { const depth = this.depths.get(index) ?? 0; const inString = this.inStrings.get(index) ?? false; if (!this.buffers.has(index) || (depth === 0 && !inString)) { @@ -82,7 +83,7 @@ export class StreamingToolCallParser { let actualIndex = index; const isKnownId = Boolean(id && this.idToIndexMap.has(id)); - const isNameOnlyDelta = Boolean(name && chunk.length === 0); + const isNameOnlyDelta = Boolean(validName && chunk.length === 0); // Handle tool call ID mapping for collision detection if (id) { @@ -102,10 +103,9 @@ export class StreamingToolCallParser { // is signaled by the name metadata, not the buffer: an empty // buffer with a name is a complete no-argument call. if ( - existingMeta?.name && - existingDepth === 0 && existingMeta?.id && - existingMeta.id !== id + existingMeta.id !== id && + (!existingMeta.name || existingDepth === 0) ) { let existingComplete = true; if (existingBuffer.trim()) { @@ -119,7 +119,9 @@ export class StreamingToolCallParser { if (existingComplete) { // We have a complete tool call with a different ID at this index // Find a new index for this tool call - actualIndex = this.findNextAvailableIndex(); + while (this.buffers.has(actualIndex)) { + actualIndex++; + } } } } @@ -165,9 +167,9 @@ export class StreamingToolCallParser { const currentBuffer = this.buffers.get(actualIndex)!; const currentDepth = this.depths.get(actualIndex)!; const meta = this.toolCallMeta.get(actualIndex)!; - if (chunk.length === 0 && (id || name)) { + if (chunk.length === 0 && (id || validName)) { if (id) meta.id = id; - if (name && !meta.name) meta.name = name; + if (validName && !meta.name) meta.name = validName; if (!meta.name && meta.id) { this.namelessToolCallIndices.add(actualIndex); } else { @@ -192,7 +194,7 @@ export class StreamingToolCallParser { // Update metadata if (id) meta.id = id; - if (name) meta.name = name; + if (validName) meta.name = validName; // Get current state for the actual index const currentInString = this.inStrings.get(actualIndex)!; From 3ec7226213ea38bc7b41db070a2d96ed136b969f Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 21:14:02 +0800 Subject: [PATCH 6/8] fix(core): preserve protocol recovery semantics --- .../openaiContentGenerator/converter.test.ts | 96 ++++++----- .../core/openaiContentGenerator/converter.ts | 7 +- .../openaiContentGenerator/pipeline.test.ts | 156 ++++++++++++++++++ .../core/openaiContentGenerator/pipeline.ts | 15 +- .../streamingToolCallParser.ts | 10 ++ 5 files changed, 243 insertions(+), 41 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index cccd2b365a9..f5b225c52b0 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -595,40 +595,46 @@ describe('OpenAIContentConverter', () => { }); }); - it('rejects a protocol-tag recovery when a new id relabels nameless arguments', () => { - const stream = withStreamParser(); - emitReasoning(stream); - converter.convertOpenAIChunkToGemini( - streamChunk('old-arguments', { - content: '', - tool_calls: [ - { - index: 0, - id: 'call_old', - function: { arguments: '{"path":"old.txt"}' }, - }, - ], - }), - stream, - ); - converter.convertOpenAIChunkToGemini( - streamChunk('new-name', { - tool_calls: [ - { - index: 0, - id: 'call_new', - function: { name: 'read_file' }, - }, - ], - }), - stream, - ); + it.each([ + ['complete', '{"path":"old.txt"}', ''], + ['incomplete', '{"path":"old.txt"', '}'], + ] as const)( + 'rejects a protocol-tag recovery when a new id relabels %s nameless arguments', + (_state, oldArguments, newArguments) => { + const stream = withStreamParser(); + emitReasoning(stream); + converter.convertOpenAIChunkToGemini( + streamChunk('old-arguments', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_old', + function: { arguments: oldArguments }, + }, + ], + }), + stream, + ); + converter.convertOpenAIChunkToGemini( + streamChunk('new-name', { + tool_calls: [ + { + index: 0, + id: 'call_new', + function: { name: 'read_file', arguments: newArguments }, + }, + ], + }), + stream, + ); - expect(() => finishStream(stream)).toThrowError( - expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), - ); - expect(stream.protocolTagSanitized).toBeUndefined(); - }); + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); + expect(stream.protocolTagSanitized).toBeUndefined(); + }, + ); it('buffers leading whitespace before a split standalone closing tag', () => { const stream = withStreamParser(); @@ -656,14 +662,28 @@ describe('OpenAIContentConverter', () => { it('ignores an exact cumulative replay of a deferred closing tag', () => { const stream = withStreamParser(); emitReasoning(stream); - emitToolCall(stream, ''); - const replay = converter.convertOpenAIChunkToGemini( - streamChunk('replay', { content: '' }), + converter.convertOpenAIChunkToGemini( + streamChunk('tag', { content: '' }), + stream, + ); + const finish = converter.convertOpenAIChunkToGemini( + streamChunk( + 'finish', + { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }, + 'tool_calls', + ), stream, ); - const finish = finishStream(stream); - expect(replay.candidates?.[0]?.content?.parts).toEqual([]); expect(finish.candidates?.[0]?.content?.parts).toEqual([ { functionCall: { id: 'call_read', name: 'read_file', args: {} }, diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index d852b94c404..8f065be6811 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1409,9 +1409,13 @@ export function convertOpenAIChunkToGemini( let visibleText = parts.map(getVisibleText).join(''); const pendingTagCandidate = requestContext.pendingThinkingTagCandidate; + const replayedClosingTag = + STANDALONE_CLOSING_THINKING_TAG_PATTERN.exec( + visibleText, + )?.[1]?.toLowerCase(); if ( pendingTagCandidate?.closingTagName && - visibleText.trim() === pendingTagCandidate.text + pendingTagCandidate.closingTagName === replayedClosingTag ) { parts = parts.filter((part) => !getVisibleText(part)); visibleText = ''; @@ -1511,6 +1515,7 @@ export function convertOpenAIChunkToGemini( choice.finish_reason !== 'tool_calls' || completedToolCalls.length === 0 || toolCallWithoutName || + toolCallParser.hasConflictingToolCallIdentity() || toolCallsTruncated || toolCallParser.hasInvalidToolCallArguments() ) { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 10a45003bd5..564ecf198ad 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -1753,6 +1753,162 @@ describe('ContentGenerationPipeline', () => { expect(logProtocolTagSanitized).not.toHaveBeenCalled(); }); + it('logs only the accepted finish after duplicate and empty trailing chunks', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const chunks = ['finish-1', 'finish-2', 'trailing-empty'].map( + (id) => + ({ + id, + choices: [{ delta: {}, finish_reason: null }], + }) as OpenAI.Chat.ChatCompletionChunk, + ); + const mockStream = { + async *[Symbol.asyncIterator]() { + yield* chunks; + }, + }; + const makeFinishResponse = (responseId: string, callId: string) => { + const response = new GenerateContentResponse(); + response.responseId = responseId; + response.candidates = [ + { + content: { + parts: [{ functionCall: { id: callId, name: 'read_file' } }], + }, + finishReason: FinishReason.STOP, + index: 0, + }, + ]; + return response; + }; + const firstFinish = makeFinishResponse('finish-1', 'call-1'); + const secondFinish = makeFinishResponse('finish-2', 'call-2'); + const emptyResponse = new GenerateContentResponse(); + emptyResponse.candidates = [ + { content: { parts: [], role: 'model' }, index: 0 }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (chunk, context) => { + if (chunk.id === 'finish-1') { + context.protocolTagSanitized = { + tagName: 'think', + toolCallCount: 1, + }; + return firstFinish; + } + if (chunk.id === 'finish-2') { + context.protocolTagSanitized = { + tagName: 'thinking', + toolCallCount: 2, + }; + return secondFinish; + } + return emptyResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + const results = []; + for await (const result of resultGenerator) results.push(result); + + expect(results).toEqual([firstFinish]); + expect(logProtocolTagSanitized).toHaveBeenCalledTimes(1); + expect(logProtocolTagSanitized).toHaveBeenCalledWith( + mockCliConfig, + expect.objectContaining({ + response_id: 'finish-1', + tag_name: 'think', + tool_call_count: 1, + }), + ); + }); + + it.each(['transport error', 'explicit abort'] as const)( + 'handles a pending closing tag on %s', + async (termination) => { + const abortController = new AbortController(); + const streamError = new Error( + termination === 'explicit abort' ? 'Aborted' : 'socket reset', + ) as Error & { code?: string }; + if (termination === 'explicit abort') { + streamError.name = 'AbortError'; + } else { + streamError.code = 'ECONNRESET'; + } + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + config: { abortSignal: abortController.signal }, + }; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'pending-tag', + choices: [{ delta: {}, finish_reason: null }], + } as OpenAI.Chat.ChatCompletionChunk; + if (termination === 'explicit abort') abortController.abort(); + throw streamError; + }, + }; + const reasoningResponse = new GenerateContentResponse(); + reasoningResponse.candidates = [ + { + content: { + parts: [{ thought: true, text: 'reasoning' }], + role: 'model', + }, + index: 0, + }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue( + [], + ); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.pendingThinkingTagCandidate = { + text: '', + closingTagName: 'think', + }; + return reasoningResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + const results = []; + let caught: unknown; + try { + for await (const result of resultGenerator) results.push(result); + } catch (error) { + caught = error; + } + + expect(results).toEqual([reasoningResponse]); + if (termination === 'explicit abort') { + expect(caught).toBe(streamError); + } else { + expect(caught).toMatchObject({ type: 'PROTOCOL_TAG_LEAK' }); + } + }, + ); + it('should handle streaming errors and reset tool calls', async () => { // Arrange const request: GenerateContentParameters = { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index 231d80274d5..e15b5717a72 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -529,7 +529,7 @@ export class ContentGenerationPipeline { const sanitization = context.protocolTagSanitized; if (sanitization) { - pendingProtocolTagSanitized = sanitization; + pendingProtocolTagSanitized ??= sanitization; context.protocolTagSanitized = undefined; } @@ -539,7 +539,6 @@ export class ContentGenerationPipeline { !response.candidates?.[0]?.finishReason && !response.usageMetadata ) { - pendingProtocolTagSanitized = undefined; continue; } @@ -607,6 +606,18 @@ export class ContentGenerationPipeline { throw error; } + if ( + context.pendingThinkingTagCandidate?.closingTagName && + request.config?.abortSignal?.aborted !== true + ) { + context.pendingThinkingTagCandidate = undefined; + context.pendingUntrustedResponseParts = undefined; + throw new InvalidStreamError( + 'Model response leaked thinking tags.', + 'PROTOCOL_TAG_LEAK', + ); + } + // Re-throw StreamContentError directly so it can be handled by // the caller's retry logic (e.g., TPM throttling retry in sendMessageStream) if (error instanceof StreamContentError) { diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index c9013becd63..5996233f8d0 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -49,6 +49,7 @@ export class StreamingToolCallParser { private idToIndexMap: Map = new Map(); /** Counter for generating new indices when collisions occur */ private nextAvailableIndex: number = 0; + private conflictingToolCallIdentity = false; /** * Processes a new chunk of tool call data and attempts to parse complete JSON objects @@ -99,6 +100,10 @@ export class StreamingToolCallParser { const existingDepth = this.depths.get(index)!; const existingMeta = this.toolCallMeta.get(index); + if (existingMeta?.id && existingMeta.id !== id) { + this.conflictingToolCallIdentity = true; + } + // Check if we have a complete tool call at this index. Occupancy // is signaled by the name metadata, not the buffer: an empty // buffer with a name is a complete no-argument call. @@ -278,6 +283,10 @@ export class StreamingToolCallParser { return this.namelessToolCallIndices.size > 0; } + hasConflictingToolCallIdentity(): boolean { + return this.conflictingToolCallIdentity; + } + hasInvalidToolCallArguments(): boolean { for (const [index, buffer] of this.buffers.entries()) { if (!this.toolCallMeta.get(index)?.name || !buffer.trim()) continue; @@ -496,6 +505,7 @@ export class StreamingToolCallParser { this.namelessToolCallIndices.clear(); this.idToIndexMap.clear(); this.nextAvailableIndex = 0; + this.conflictingToolCallIdentity = false; } /** From 61d67a15ac0c0d5b44c66f25b8ac103c6e87463c Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 14 Jul 2026 23:47:50 +0800 Subject: [PATCH 7/8] fix(core): close protocol recovery review gaps --- .../openaiContentGenerator/converter.test.ts | 55 ++++++++++ .../core/openaiContentGenerator/converter.ts | 13 ++- .../openaiContentGenerator/pipeline.test.ts | 100 ++++++++++++++++++ .../core/openaiContentGenerator/pipeline.ts | 21 ++-- .../streamingToolCallParser.test.ts | 34 ++++++ .../streamingToolCallParser.ts | 44 ++++---- 6 files changed, 239 insertions(+), 28 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index f5b225c52b0..29efc2e65d0 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -595,6 +595,47 @@ describe('OpenAIContentConverter', () => { }); }); + it('recovers complete same-index tool calls with distinct IDs', () => { + const stream = withStreamParser(); + emitReasoning(stream); + converter.convertOpenAIChunkToGemini( + streamChunk('tool-calls', { + content: '', + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + { + index: 0, + id: 'call_list', + function: { name: 'list_directory', arguments: '{}' }, + }, + ], + }), + stream, + ); + const finish = finishStream(stream); + + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + { + functionCall: { + id: 'call_list', + name: 'list_directory', + args: {}, + }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 2, + }); + }); + it.each([ ['complete', '{"path":"old.txt"}', ''], ['incomplete', '{"path":"old.txt"', '}'], @@ -659,6 +700,20 @@ describe('OpenAIContentConverter', () => { }); }); + it('emits trailing whitespace when the stream finishes', () => { + const stream = withStreamParser(); + emitReasoning(stream); + const whitespace = converter.convertOpenAIChunkToGemini( + streamChunk('whitespace', { content: ' \n' }), + stream, + ); + const finish = finishStream(stream, 'stop'); + + expect(whitespace.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([{ text: ' \n' }]); + expect(stream.pendingThinkingTagCandidate).toBeUndefined(); + }); + it('ignores an exact cumulative replay of a deferred closing tag', () => { const stream = withStreamParser(); emitReasoning(stream); diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index 8f065be6811..ee48184d656 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1442,6 +1442,10 @@ export function convertOpenAIChunkToGemini( const isPossibleTag = canBeStandaloneThinkingTagPrefix( combinedCandidateText, ); + const finishedWhitespaceCandidate = + Boolean(choice.finish_reason) && + !closingTagName && + !/\S/.test(combinedCandidateText); if (openingTag) { throwProtocolTagLeak(requestContext); @@ -1451,7 +1455,14 @@ export function convertOpenAIChunkToGemini( throwProtocolTagLeak(requestContext); } - if (isPossibleTag) { + if (finishedWhitespaceCandidate) { + parts = parts.filter((part) => !getVisibleText(part)); + if (combinedCandidateText) { + parts.push({ text: combinedCandidateText }); + } + visibleText = combinedCandidateText; + requestContext.pendingThinkingTagCandidate = undefined; + } else if (isPossibleTag) { if ( !closingTagName && combinedCandidateText.length > MAX_THINKING_TAG_CANDIDATE_LENGTH diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 564ecf198ad..9b1e1eb06f5 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -1701,6 +1701,45 @@ describe('ContentGenerationPipeline', () => { expect(logProtocolTagSanitized).not.toHaveBeenCalled(); }); + it('allows a whitespace-only tag candidate at clean stream EOF', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'response-id', + choices: [{ delta: { content: ' ' }, finish_reason: null }], + } as OpenAI.Chat.ChatCompletionChunk; + }, + }; + const emptyResponse = new GenerateContentResponse(); + emptyResponse.candidates = [ + { content: { parts: [], role: 'model' }, index: 0 }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.pendingThinkingTagCandidate = { text: ' ' }; + return emptyResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + const results = []; + for await (const result of resultGenerator) results.push(result); + + expect(results).toEqual([]); + }); + it('does not log protocol-tag sanitization before a held finish is yielded', async () => { const request: GenerateContentParameters = { model: 'test-model', @@ -1909,6 +1948,67 @@ describe('ContentGenerationPipeline', () => { }, ); + it('preserves a StreamContentError while a closing tag is pending', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'pending-tag', + object: 'chat.completion.chunk', + created: Date.now(), + model: 'test-model', + choices: [{ index: 0, delta: {}, finish_reason: null }], + } as OpenAI.Chat.ChatCompletionChunk; + yield { + id: 'error', + object: 'chat.completion.chunk', + created: Date.now(), + model: 'test-model', + choices: [ + { + index: 0, + delta: { content: 'Throttling: TPM(1/1)' }, + finish_reason: 'error_finish', + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk; + }, + }; + const emptyResponse = new GenerateContentResponse(); + emptyResponse.candidates = [ + { content: { parts: [], role: 'model' }, index: 0 }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.pendingThinkingTagCandidate = { + text: '', + closingTagName: 'think', + }; + return emptyResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + + await expect(async () => { + for await (const _ of resultGenerator) { + // Consume until the provider error is raised. + } + }).rejects.toThrow(StreamContentError); + expect(mockErrorHandler.handle).not.toHaveBeenCalled(); + }); + it('should handle streaming errors and reset tool calls', async () => { // Arrange const request: GenerateContentParameters = { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index e15b5717a72..ade6deedc9e 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -588,7 +588,14 @@ export class ContentGenerationPipeline { } } - if (context.pendingThinkingTagCandidate) { + if ( + context.pendingThinkingTagCandidate && + !context.pendingThinkingTagCandidate.closingTagName && + !/\S/.test(context.pendingThinkingTagCandidate.text) + ) { + context.pendingThinkingTagCandidate = undefined; + context.pendingUntrustedResponseParts = undefined; + } else if (context.pendingThinkingTagCandidate) { throw new InvalidStreamError( 'Model response leaked thinking tags.', 'PROTOCOL_TAG_LEAK', @@ -606,6 +613,12 @@ export class ContentGenerationPipeline { throw error; } + // Re-throw StreamContentError directly so it can be handled by + // the caller's retry logic (e.g., TPM throttling retry in sendMessageStream) + if (error instanceof StreamContentError) { + throw redactProxyError(error); + } + if ( context.pendingThinkingTagCandidate?.closingTagName && request.config?.abortSignal?.aborted !== true @@ -618,12 +631,6 @@ export class ContentGenerationPipeline { ); } - // Re-throw StreamContentError directly so it can be handled by - // the caller's retry logic (e.g., TPM throttling retry in sendMessageStream) - if (error instanceof StreamContentError) { - throw redactProxyError(error); - } - // Bypass handleError: it strips `code` from timeout errors, which would // prevent classifyRetryError from recognizing retryable ETIMEDOUT. if (error instanceof StreamInactivityTimeoutError) { diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts index 630145f567d..181c09c3f88 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts @@ -767,6 +767,27 @@ describe('StreamingToolCallParser', () => { ]); }); + it('should normalize a tool call name before storing it', () => { + parser.addChunk(0, '{}', 'call_1', ' read_file '); + + expect(parser.getCompletedToolCalls()).toEqual([ + { + id: 'call_1', + name: 'read_file', + args: {}, + index: 0, + }, + ]); + }); + + it('should preserve the first non-empty name for a tool call ID', () => { + parser.addChunk(0, '{"file_path":', 'call_1', 'read_file'); + parser.addChunk(0, '"a.ts"}', 'call_1', 'shell'); + + expect(parser.getCompletedToolCalls()[0]?.name).toBe('read_file'); + expect(parser.hasConflictingToolCallIdentity()).toBe(true); + }); + it('should detect index collision and find new index', () => { // First complete tool call at index 0 parser.addChunk(0, '{"param1": "value1"}', 'call_1', 'function1'); @@ -789,6 +810,19 @@ describe('StreamingToolCallParser', () => { expect(call2).toBeDefined(); expect(call1?.args).toEqual({ param1: 'value1' }); expect(call2?.args).toEqual({ param2: 'value2' }); + expect(parser.hasConflictingToolCallIdentity()).toBe(false); + }); + + it('should reject unsafe provider indices', () => { + const result = parser.addChunk( + Number.MAX_SAFE_INTEGER + 1, + '{}', + 'call_1', + 'read_file', + ); + + expect(result.error?.message).toContain('Invalid tool call index'); + expect(parser.hasConflictingToolCallIdentity()).toBe(true); }); it('should handle continuation chunks without ID correctly', () => { diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index 5996233f8d0..b5c5e76f9ef 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -73,7 +73,7 @@ export class StreamingToolCallParser { id?: string, name?: string, ): ToolCallParseResult { - const validName = name?.trim() ? name : undefined; + const validName = name?.trim() || undefined; if (!id && !validName && !chunk.trim()) { const depth = this.depths.get(index) ?? 0; const inString = this.inStrings.get(index) ?? false; @@ -81,6 +81,13 @@ export class StreamingToolCallParser { return { complete: false }; } } + if (!Number.isSafeInteger(index) || index < 0) { + this.conflictingToolCallIdentity = true; + return { + complete: false, + error: new Error(`Invalid tool call index: ${index}`), + }; + } let actualIndex = index; const isKnownId = Boolean(id && this.idToIndexMap.has(id)); @@ -101,32 +108,22 @@ export class StreamingToolCallParser { const existingMeta = this.toolCallMeta.get(index); if (existingMeta?.id && existingMeta.id !== id) { - this.conflictingToolCallIdentity = true; - } - - // Check if we have a complete tool call at this index. Occupancy - // is signaled by the name metadata, not the buffer: an empty - // buffer with a name is a complete no-argument call. - if ( - existingMeta?.id && - existingMeta.id !== id && - (!existingMeta.name || existingDepth === 0) - ) { - let existingComplete = true; - if (existingBuffer.trim()) { + let existingComplete = existingDepth === 0; + if (existingComplete && existingBuffer.trim()) { try { JSON.parse(existingBuffer); } catch { - // Existing buffer is not complete JSON, we can reuse this index existingComplete = false; } } if (existingComplete) { - // We have a complete tool call with a different ID at this index - // Find a new index for this tool call - while (this.buffers.has(actualIndex)) { - actualIndex++; + actualIndex = 0; + while (this.buffers.has(actualIndex)) actualIndex += 1; + if (!existingMeta.name) { + this.conflictingToolCallIdentity = true; } + } else { + this.conflictingToolCallIdentity = true; } } } @@ -198,8 +195,15 @@ export class StreamingToolCallParser { } // Update metadata + const identityChanged = Boolean(id && meta.id && meta.id !== id); if (id) meta.id = id; - if (validName) meta.name = validName; + if (validName) { + if (!identityChanged && meta.name && meta.name !== validName) { + this.conflictingToolCallIdentity = true; + } else { + meta.name = validName; + } + } // Get current state for the actual index const currentInString = this.inStrings.get(actualIndex)!; From 8771e4835791226b187ee3da9e39fcc4e740b203 Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Wed, 15 Jul 2026 10:53:03 +0800 Subject: [PATCH 8/8] fix(core): harden protocol tag recovery --- .../openaiContentGenerator/converter.test.ts | 108 ++++++++++- .../core/openaiContentGenerator/converter.ts | 16 +- .../openaiContentGenerator/pipeline.test.ts | 172 ++++++++++++++++++ .../core/openaiContentGenerator/pipeline.ts | 58 +++++- .../streamingToolCallParser.test.ts | 9 +- .../streamingToolCallParser.ts | 23 ++- 6 files changed, 358 insertions(+), 28 deletions(-) diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index adbf87b74cc..73364464590 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -509,6 +509,22 @@ describe('OpenAIContentConverter', () => { ).toThrowError(expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' })); }); + it('rejects closing-tag recovery after a tag leaked in reasoning', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('reasoning', { + reasoning_content: 'Let me check', + }), + stream, + ); + emitToolCall(stream, ''); + + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'PROTOCOL_TAG_LEAK' }), + ); + expect(stream.protocolTagSanitized).toBeUndefined(); + }); + it.each([ ['think', '\n\n\n'], ['thinking', ' '], @@ -701,6 +717,23 @@ describe('OpenAIContentConverter', () => { }); }); + it('buffers long whitespace without treating it as a protocol tag', () => { + const stream = withStreamParser(); + const whitespace = ' '.repeat(129); + emitReasoning(stream); + + const pending = converter.convertOpenAIChunkToGemini( + streamChunk('whitespace', { content: whitespace }), + stream, + ); + const finish = finishStream(stream, 'stop'); + + expect(pending.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { text: whitespace }, + ]); + }); + it('emits trailing whitespace when the stream finishes', () => { const stream = withStreamParser(); emitReasoning(stream); @@ -751,6 +784,79 @@ describe('OpenAIContentConverter', () => { }); }); + it('ignores cumulative replays of an incomplete closing tag', () => { + const stream = withStreamParser(); + emitReasoning(stream); + const first = converter.convertOpenAIChunkToGemini( + streamChunk('tag-prefix', { content: ''); + const finish = finishStream(stream); + + expect(first.candidates?.[0]?.content?.parts).toEqual([]); + expect(replay.candidates?.[0]?.content?.parts).toEqual([]); + expect(tag.candidates?.[0]?.content?.parts).toEqual([]); + expect(finish.candidates?.[0]?.content?.parts).toEqual([ + { + functionCall: { id: 'call_read', name: 'read_file', args: {} }, + }, + ]); + expect(stream.protocolTagSanitized).toEqual({ + tagName: 'think', + toolCallCount: 1, + }); + }); + + it('rejects an invalid tool-call index on a stop finish', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('invalid-tool-call', { + tool_calls: [ + { + index: Number.MAX_SAFE_INTEGER + 1, + id: 'call_invalid', + function: { name: 'read_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + + expect(() => finishStream(stream, 'stop')).toThrowError( + expect.objectContaining({ type: 'MALFORMED_TOOL_CALL' }), + ); + }); + + it('rejects valid tool calls accompanied by an invalid index', () => { + const stream = withStreamParser(); + converter.convertOpenAIChunkToGemini( + streamChunk('mixed-tool-calls', { + tool_calls: [ + { + index: 0, + id: 'call_read', + function: { name: 'read_file', arguments: '{}' }, + }, + { + index: Number.MAX_SAFE_INTEGER + 1, + id: 'call_invalid', + function: { name: 'write_file', arguments: '{}' }, + }, + ], + }), + stream, + ); + + expect(() => finishStream(stream)).toThrowError( + expect.objectContaining({ type: 'MALFORMED_TOOL_CALL' }), + ); + }); + it('releases a split tag-like prefix when it becomes ordinary text', () => { const stream = withStreamParser(); converter.convertOpenAIChunkToGemini( @@ -842,7 +948,7 @@ describe('OpenAIContentConverter', () => { expect(stream.protocolTagSanitized).toBeUndefined(); }); - it.each(['{"path":', '{bad}', 'null', '[]', '42'])( + it.each(['{"path":', '{bad}', 'null', '[]', '42', ' '])( 'rejects a standalone closing thinking tag with unsafe tool arguments %s', (toolArguments) => { const stream = withStreamParser(); diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index 86ea9c7c588..cd48cd605e2 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -1422,13 +1422,18 @@ export function convertOpenAIChunkToGemini( let visibleText = parts.map(getVisibleText).join(''); const pendingTagCandidate = requestContext.pendingThinkingTagCandidate; + const replayedTagPrefix = + !pendingTagCandidate?.closingTagName && + /\S/.test(pendingTagCandidate?.text ?? '') && + pendingTagCandidate?.text === visibleText; const replayedClosingTag = STANDALONE_CLOSING_THINKING_TAG_PATTERN.exec( visibleText, )?.[1]?.toLowerCase(); if ( - pendingTagCandidate?.closingTagName && - pendingTagCandidate.closingTagName === replayedClosingTag + replayedTagPrefix || + (pendingTagCandidate?.closingTagName && + pendingTagCandidate.closingTagName === replayedClosingTag) ) { parts = parts.filter((part) => !getVisibleText(part)); visibleText = ''; @@ -1478,7 +1483,8 @@ export function convertOpenAIChunkToGemini( } else if (isPossibleTag) { if ( !closingTagName && - combinedCandidateText.length > MAX_THINKING_TAG_CANDIDATE_LENGTH + combinedCandidateText.trimStart().length > + MAX_THINKING_TAG_CANDIDATE_LENGTH ) { throwProtocolTagLeak(requestContext); } @@ -1536,6 +1542,7 @@ export function convertOpenAIChunkToGemini( requestContext.pendingThinkingTagCandidate?.closingTagName ) { if ( + requestContext.hasThinkingTagInReasoning === true || choice.finish_reason !== 'tool_calls' || completedToolCalls.length === 0 || toolCallWithoutName || @@ -1554,7 +1561,8 @@ export function convertOpenAIChunkToGemini( if ( choice.finish_reason && - (toolCallWithoutName || + (toolCallParser.hasInvalidToolCallIndex() || + toolCallWithoutName || (choice.finish_reason === 'tool_calls' && completedToolCalls.length === 0)) ) { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 028df5f6a1e..4c522f2bceb 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -1741,6 +1741,51 @@ describe('ContentGenerationPipeline', () => { expect(results).toEqual([]); }); + it('flushes held response parts for a whitespace-only candidate at clean EOF', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + id: 'response-id', + choices: [{ delta: { content: ' ' }, finish_reason: null }], + } as OpenAI.Chat.ChatCompletionChunk; + }, + }; + const emptyResponse = new GenerateContentResponse(); + emptyResponse.candidates = [ + { content: { parts: [], role: 'model' }, index: 0 }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (_chunk, context) => { + context.pendingThinkingTagCandidate = { text: ' ' }; + context.pendingUntrustedResponseParts = [ + { thought: true, text: 'reasoning' }, + ]; + return emptyResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + const results = []; + for await (const result of resultGenerator) results.push(result); + + expect(results).toHaveLength(1); + expect(results[0]?.candidates?.[0]?.content?.parts).toEqual([ + { thought: true, text: 'reasoning' }, + ]); + }); + it('does not log protocol-tag sanitization before a held finish is yielded', async () => { const request: GenerateContentParameters = { model: 'test-model', @@ -1874,6 +1919,133 @@ describe('ContentGenerationPipeline', () => { ); }); + it('does not attribute sanitization from a discarded duplicate finish', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const chunks = ['finish-1', 'finish-2', 'usage'].map( + (id) => + ({ + id, + choices: [{ delta: {}, finish_reason: null }], + }) as OpenAI.Chat.ChatCompletionChunk, + ); + const mockStream = { + async *[Symbol.asyncIterator]() { + yield* chunks; + }, + }; + const makeFinishResponse = (responseId: string) => { + const response = new GenerateContentResponse(); + response.responseId = responseId; + response.candidates = [ + { + content: { parts: [{ functionCall: { name: 'read_file' } }] }, + finishReason: FinishReason.STOP, + index: 0, + }, + ]; + return response; + }; + const firstFinish = makeFinishResponse('finish-1'); + const secondFinish = makeFinishResponse('finish-2'); + const usageResponse = new GenerateContentResponse(); + usageResponse.usageMetadata = { totalTokenCount: 1 }; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (chunk, context) => { + if (chunk.id === 'finish-1') return firstFinish; + if (chunk.id === 'finish-2') { + context.protocolTagSanitized = { + tagName: 'think', + toolCallCount: 1, + }; + return secondFinish; + } + return usageResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + for await (const _ of resultGenerator) { + // Consume the merged finish response. + } + + expect(logProtocolTagSanitized).not.toHaveBeenCalled(); + }); + + it('rejects visible content after a sanitized finish', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const chunks = ['finish', 'trailing-content'].map( + (id) => + ({ + id, + choices: [{ delta: {}, finish_reason: null }], + }) as OpenAI.Chat.ChatCompletionChunk, + ); + const mockStream = { + async *[Symbol.asyncIterator]() { + yield* chunks; + }, + }; + const finishResponse = new GenerateContentResponse(); + finishResponse.responseId = 'finish'; + finishResponse.candidates = [ + { + content: { parts: [{ functionCall: { name: 'read_file' } }] }, + finishReason: FinishReason.STOP, + index: 0, + }, + ]; + const trailingResponse = new GenerateContentResponse(); + trailingResponse.candidates = [ + { + content: { parts: [{ text: 'unexpected' }], role: 'model' }, + index: 0, + }, + ]; + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockImplementation( + (chunk, context) => { + if (chunk.id === 'finish') { + context.protocolTagSanitized = { + tagName: 'think', + toolCallCount: 1, + }; + return finishResponse; + } + return trailingResponse; + }, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + + await expect(async () => { + for await (const _ of resultGenerator) { + // Consume until trailing content validation runs. + } + }).rejects.toMatchObject({ type: 'PROTOCOL_TAG_LEAK' }); + expect(logProtocolTagSanitized).not.toHaveBeenCalled(); + }); + it.each(['transport error', 'explicit abort'] as const)( 'handles a pending closing tag on %s', async (termination) => { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index 86ac7a09ada..d41dd77605a 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -484,21 +484,23 @@ export class ContentGenerationPipeline { // function-call parts from the finish chunk). let pendingFinishResponse: GenerateContentResponse | null = null; let finishYielded = false; - let pendingProtocolTagSanitized: + let pendingFinishProtocolTagSanitized: | NonNullable | undefined; const logPendingProtocolTagSanitized = ( response: GenerateContentResponse, + sanitization: + | NonNullable + | undefined, ) => { - if (!pendingProtocolTagSanitized) return; + if (!sanitization) return; const event = new ProtocolTagSanitizedEvent({ model: context.model, promptId: userPromptId, responseId: response.responseId, - tagName: pendingProtocolTagSanitized.tagName, - toolCallCount: pendingProtocolTagSanitized.toolCallCount, + tagName: sanitization.tagName, + toolCallCount: sanitization.toolCallCount, }); - pendingProtocolTagSanitized = undefined; debugLogger.warn('Sanitized a model protocol tag', { model: event.model, promptId: event.prompt_id, @@ -530,7 +532,6 @@ export class ContentGenerationPipeline { const sanitization = context.protocolTagSanitized; if (sanitization) { - pendingProtocolTagSanitized ??= sanitization; context.protocolTagSanitized = undefined; } @@ -545,6 +546,20 @@ export class ContentGenerationPipeline { continue; } + if ( + pendingFinishProtocolTagSanitized && + pendingFinishResponse && + !response.candidates?.[0]?.finishReason && + response.candidates?.some( + (candidate) => (candidate.content?.parts?.length ?? 0) > 0, + ) + ) { + throw new InvalidStreamError( + 'Model response continued after a finish reason.', + 'PROTOCOL_TAG_LEAK', + ); + } + // Stage 2c: Handle chunk merging for providers that send // finishReason and usageMetadata in separate chunks. // Once the merged finish response has been yielded, skip @@ -568,6 +583,14 @@ export class ContentGenerationPipeline { continue; } + if ( + !pendingFinishResponse && + response.candidates?.[0]?.finishReason && + sanitization + ) { + pendingFinishProtocolTagSanitized = sanitization; + } + const shouldYield = this.handleChunkMerging( response, collectedGeminiResponses, @@ -579,13 +602,16 @@ export class ContentGenerationPipeline { if (shouldYield) { // If we have a pending finish response, yield it instead if (pendingFinishResponse) { - logPendingProtocolTagSanitized(pendingFinishResponse); + logPendingProtocolTagSanitized( + pendingFinishResponse, + pendingFinishProtocolTagSanitized, + ); yield pendingFinishResponse; finishYielded = true; // Keep pendingFinishResponse alive so late-arriving usage // metadata can still be merged (see finishYielded block above). } else { - logPendingProtocolTagSanitized(response); + logPendingProtocolTagSanitized(response, sanitization); yield response; } } @@ -596,8 +622,19 @@ export class ContentGenerationPipeline { !context.pendingThinkingTagCandidate.closingTagName && !/\S/.test(context.pendingThinkingTagCandidate.text) ) { + const pendingParts = context.pendingUntrustedResponseParts; context.pendingThinkingTagCandidate = undefined; context.pendingUntrustedResponseParts = undefined; + if (pendingParts?.length) { + const response = new GenerateContentResponse(); + response.candidates = [ + { + content: { parts: pendingParts, role: 'model' }, + index: 0, + }, + ]; + yield response; + } } else if (context.pendingThinkingTagCandidate) { throw new InvalidStreamError( 'Model response leaked thinking tags.', @@ -608,7 +645,10 @@ export class ContentGenerationPipeline { // Stage 2d: If there's still a pending finish response at the end // (e.g. no usage chunk arrived after the finish chunk), yield it. if (pendingFinishResponse && !finishYielded) { - logPendingProtocolTagSanitized(pendingFinishResponse); + logPendingProtocolTagSanitized( + pendingFinishResponse, + pendingFinishProtocolTagSanitized, + ); yield pendingFinishResponse; } } catch (error) { diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts index 2a26bbf1330..b9522dd9d62 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts @@ -814,14 +814,10 @@ describe('StreamingToolCallParser', () => { }); it('should reject unsafe provider indices', () => { - const result = parser.addChunk( - Number.MAX_SAFE_INTEGER + 1, - '{}', - 'call_1', - 'read_file', - ); + const result = parser.addChunk(Number.MAX_SAFE_INTEGER + 1, ' '); expect(result.error?.message).toContain('Invalid tool call index'); + expect(parser.hasInvalidToolCallIndex()).toBe(true); expect(parser.hasConflictingToolCallIdentity()).toBe(true); }); @@ -1133,6 +1129,7 @@ describe('StreamingToolCallParser', () => { describe('hasInvalidToolCallArguments', () => { it.each([ ['', false], + [' ', true], ['{"path":"a.ts"}', false], ['{bad}', true], ['null', true], diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index b8a90642094..57f53ce2375 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -54,6 +54,7 @@ export class StreamingToolCallParser { /** Counter for generating new indices when collisions occur */ private nextAvailableIndex: number = 0; private conflictingToolCallIdentity = false; + private invalidToolCallIndex = false; /** * Processes a new chunk of tool call data and attempts to parse complete JSON objects @@ -78,20 +79,21 @@ export class StreamingToolCallParser { name?: string, ): ToolCallParseResult { const validName = name?.trim() || undefined; - if (!id && !validName && !chunk.trim()) { - const depth = this.depths.get(index) ?? 0; - const inString = this.inStrings.get(index) ?? false; - if (!this.buffers.has(index) || (depth === 0 && !inString)) { - return { complete: false }; - } - } if (!Number.isSafeInteger(index) || index < 0) { this.conflictingToolCallIdentity = true; + this.invalidToolCallIndex = true; return { complete: false, error: new Error(`Invalid tool call index: ${index}`), }; } + if (!id && !validName && !chunk.trim()) { + const depth = this.depths.get(index) ?? 0; + const inString = this.inStrings.get(index) ?? false; + if (!this.buffers.has(index) || (depth === 0 && !inString)) { + return { complete: false }; + } + } let actualIndex = index; const isKnownId = Boolean(id && this.idToIndexMap.has(id)); @@ -326,9 +328,13 @@ export class StreamingToolCallParser { return this.conflictingToolCallIdentity; } + hasInvalidToolCallIndex(): boolean { + return this.invalidToolCallIndex; + } + hasInvalidToolCallArguments(): boolean { for (const [index, buffer] of this.buffers.entries()) { - if (!this.toolCallMeta.get(index)?.name || !buffer.trim()) continue; + if (!this.toolCallMeta.get(index)?.name || buffer.length === 0) continue; try { const args: unknown = JSON.parse(buffer); @@ -551,6 +557,7 @@ export class StreamingToolCallParser { this.pendingIndexRemaps.clear(); this.nextAvailableIndex = 0; this.conflictingToolCallIdentity = false; + this.invalidToolCallIndex = false; } /**