diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 72b89f54add..cb0df1f80bf 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -238,6 +238,25 @@ describe('GeminiChat', async () => { return collecting; } + function streamResponse( + response: GenerateContentResponse, + ): AsyncGenerator { + return (async function* () { + yield response; + })(); + } + + function stopResponse(parts: Part[]): GenerateContentResponse { + return { + candidates: [ + { + content: { parts }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + } + describe('system instruction helpers', () => { it('replaces prior session-start context instead of appending indefinitely', () => { const isolatedChat = new GeminiChat( @@ -4607,12 +4626,20 @@ describe('GeminiChat', async () => { ); await expectStreamExhaustion(stream); - // Should be called 3 times (1 initial + 2 transient retries) + // Should be called 5 times (1 initial + 4 transient retries) expect( mockContentGenerator.generateContentStream, - ).toHaveBeenCalledTimes(3); - expect(mockLogContentRetry).toHaveBeenCalledTimes(2); + ).toHaveBeenCalledTimes(5); + expect(mockLogContentRetry).toHaveBeenCalledTimes(4); expect(mockLogContentRetryFailure).toHaveBeenCalledTimes(1); + expect(mockLogContentRetryFailure).toHaveBeenCalledWith( + mockConfig, + expect.objectContaining({ + total_attempts: 5, + final_error_type: 'NO_FINISH_REASON', + model: 'test-model', + }), + ); // History should still contain the user message. const history = chat.getHistory(); @@ -4626,9 +4653,237 @@ describe('GeminiChat', async () => { } }); - it('should retry usage-only empty streams and succeed on a later attempt', async () => { + it('should recover after four consecutive invalid streams', async () => { + vi.useFakeTimers(); + try { + let callCount = 0; + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => { + callCount++; + if (callCount <= 4) { + return (async function* () { + yield { + candidates: [ + { + content: { parts: [] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + } + + return (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Recovered response' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + }); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-four-invalid-streams', + ); + const events = await collectStreamWithFakeTimers(stream, 25_000); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(5); + expect(mockLogContentRetry).toHaveBeenCalledTimes(4); + for (const [index, retryDelayMs] of [ + 2000, 4000, 6000, 8000, + ].entries()) { + expect(mockLogContentRetry).toHaveBeenNthCalledWith( + index + 1, + mockConfig, + expect.objectContaining({ + attempt_number: index, + error_type: 'NO_RESPONSE_TEXT', + retry_delay_ms: retryDelayMs, + model: 'test-model', + }), + ); + } + expect(mockLogContentRetryFailure).not.toHaveBeenCalled(); + expect( + events.some( + (event) => + event.type === StreamEventType.CHUNK && + event.value.candidates?.[0]?.content?.parts?.[0]?.text === + 'Recovered response', + ), + ).toBe(true); + expect(chat.getHistory()).toEqual([ + { role: 'user', parts: [{ text: 'test' }] }, + { role: 'model', parts: [{ text: 'Recovered response' }] }, + ]); + } finally { + vi.useRealTimers(); + } + }); + + it('should keep protocol tag leak retries at the existing budget', async () => { + vi.useFakeTimers(); + try { + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => + (async function* () { + yield { + candidates: [ + { + content: { + parts: [ + { + text: 'hiddenleaked', + }, + ], + }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-protocol-leak-budget', + ); + await expectStreamExhaustion(stream); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(3); + expect(mockLogContentRetry).toHaveBeenCalledTimes(2); + expect(mockLogContentRetryFailure).toHaveBeenCalledWith( + mockConfig, + expect.objectContaining({ + total_attempts: 3, + final_error_type: 'PROTOCOL_TAG_LEAK', + model: 'test-model', + }), + ); + } finally { + vi.useRealTimers(); + } + }); + + it('keeps invalid stream retry budgets independent across error types', async () => { + vi.useFakeTimers(); + try { + let callCount = 0; + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => { + callCount++; + if (callCount <= 2) { + return streamResponse(stopResponse([])); + } + if (callCount === 3) { + return streamResponse( + stopResponse([ + { + text: 'hiddenleaked', + }, + ]), + ); + } + + return streamResponse(stopResponse([{ text: 'Recovered response' }])); + }); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-mixed-invalid-streams', + ); + const events = await collectStreamWithFakeTimers(stream, 15_000); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(4); + expect(mockLogContentRetry).toHaveBeenCalledTimes(3); + expect(mockLogContentRetry).toHaveBeenLastCalledWith( + mockConfig, + expect.objectContaining({ + attempt_number: 0, + error_type: 'PROTOCOL_TAG_LEAK', + retry_delay_ms: 2000, + model: 'test-model', + }), + ); + expect( + events.some( + (event) => + event.type === StreamEventType.CHUNK && + event.value.candidates?.[0]?.content?.parts?.[0]?.text === + 'Recovered response', + ), + ).toBe(true); + } finally { + vi.useRealTimers(); + } + }); + + it('surfaces an abort fired during the invalid-stream retry delay without retrying again', async () => { vi.useFakeTimers(); try { + const abortController = new AbortController(); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => streamResponse(stopResponse([]))); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test', config: { abortSignal: abortController.signal } }, + 'prompt-id-invalid-stream-abort-delay', + ); + + const iterator = stream[Symbol.asyncIterator](); + let next = await iterator.next(); + while (!next.done && next.value.type !== StreamEventType.RETRY) { + next = await iterator.next(); + } + if (next.done) { + throw new Error('Expected invalid stream retry event.'); + } + expect(next.value.type).toBe(StreamEventType.RETRY); + + const nextPromise = iterator.next(); + abortController.abort(); + await expect(nextPromise).rejects.toThrow(); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it('should retry usage-only empty streams without recording failed attempts', async () => { + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); vi.mocked(mockContentGenerator.generateContentStream) .mockImplementationOnce(async () => (async function* () { @@ -4656,7 +4911,7 @@ describe('GeminiChat', async () => { })(), ); - const stream = await chat.sendMessageStream( + const stream = await chatWithRecording.sendMessageStream( 'test-model', { message: 'test' }, 'prompt-id-empty-usage-retry', @@ -4667,6 +4922,10 @@ describe('GeminiChat', async () => { mockContentGenerator.generateContentStream, ).toHaveBeenCalledTimes(2); expect(mockLogContentRetry).toHaveBeenCalledTimes(1); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordAssistantTurn.mock.calls[0]?.[0].message).toEqual([ + { text: 'Recovered after empty stream' }, + ]); expect( events.some( (e) => @@ -8335,6 +8594,25 @@ describe('GeminiChat', async () => { })(); } + function invalidStream( + type: 'NO_FINISH_REASON' | 'PROTOCOL_TAG_LEAK', + ): AsyncGenerator { + return { + [Symbol.asyncIterator]() { + return this; + }, + async next() { + throw new InvalidStreamError('Invalid continuation stream.', type); + }, + async return() { + return { done: true, value: undefined }; + }, + async throw(error?: unknown) { + throw error; + }, + } as AsyncGenerator; + } + it('re-clamps maxOutputTokens on each recovery send as the prompt grows (window invariant)', async () => { // The #5950 shape at recovery time: 131,072 window, 71,349 prompt, // 64K-ceiling model. Initial clamp grants 49,722. The response @@ -9620,6 +9898,96 @@ describe('GeminiChat', async () => { ).toBe(true); }); + it('keeps protocol tag leak budget during output continuation', async () => { + vi.useFakeTimers(); + try { + const streams = [ + makeStream([makeChunk([{ text: 'initial' }], 'MAX_TOKENS')]), + makeStream([makeChunk([{ text: 'escalated' }], 'MAX_TOKENS')]), + invalidStream('PROTOCOL_TAG_LEAK'), + invalidStream('PROTOCOL_TAG_LEAK'), + invalidStream('PROTOCOL_TAG_LEAK'), + invalidStream('PROTOCOL_TAG_LEAK'), + invalidStream('PROTOCOL_TAG_LEAK'), + ]; + let callIndex = 0; + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => streams[callIndex++]!); + + const stream = await chat.sendMessageStream( + 'gemini-pro', + { message: 'essay' }, + 'prompt-recovery-protocol-leak-budget', + ); + + await collectStreamWithFakeTimers(stream, 35_000); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(5); + expect(mockLogContentRetry).toHaveBeenCalledTimes(2); + expect(mockLogContentRetry).toHaveBeenLastCalledWith( + mockConfig, + expect.objectContaining({ + attempt_number: 1, + error_type: 'PROTOCOL_TAG_LEAK', + }), + ); + } finally { + vi.useRealTimers(); + } + }); + + it('keeps continuation retry budgets independent across error types', async () => { + vi.useFakeTimers(); + try { + const streams = [ + makeStream([makeChunk([{ text: 'initial' }], 'MAX_TOKENS')]), + makeStream([makeChunk([{ text: 'escalated' }], 'MAX_TOKENS')]), + invalidStream('NO_FINISH_REASON'), + invalidStream('NO_FINISH_REASON'), + invalidStream('PROTOCOL_TAG_LEAK'), + makeStream([makeChunk([{ text: ' recovered' }], 'STOP')]), + ]; + let callIndex = 0; + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockImplementation(async () => streams[callIndex++]!); + + const stream = await chat.sendMessageStream( + 'gemini-pro', + { message: 'essay' }, + 'prompt-recovery-mixed-invalid-streams', + ); + + const events = await collectStreamWithFakeTimers(stream, 15_000); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(6); + expect(mockLogContentRetry).toHaveBeenCalledTimes(3); + expect(mockLogContentRetry).toHaveBeenLastCalledWith( + mockConfig, + expect.objectContaining({ + attempt_number: 0, + error_type: 'PROTOCOL_TAG_LEAK', + retry_delay_ms: 2000, + }), + ); + expect( + events.some( + (e) => + e.type === StreamEventType.CHUNK && + e.value.candidates?.[0]?.content?.parts?.[0]?.text === + ' recovered', + ), + ).toBe(true); + } finally { + vi.useRealTimers(); + } + }); + it('should cap recovery attempts at MAX_OUTPUT_RECOVERY_ATTEMPTS (3)', async () => { // Every stream returns MAX_TOKENS with text (no functionCall). vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index e3230e8f73b..1069ad96d93 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -342,7 +342,10 @@ const INVALID_CONTENT_RETRY_OPTIONS: ContentRetryOptions = { // reason. All are retried with an independent budget (similar to rate-limit // retries) so they do not consume each other's retry budgets. const INVALID_STREAM_RETRY_CONFIG = { - maxRetries: 2, + transientMaxRetries: 4, + // Protocol-tag leaks are model-output validation failures, not the + // provider-side empty/truncated streams covered by issue #6670. + protocolTagLeakMaxRetries: 2, initialDelayMs: 2000, }; @@ -2220,7 +2223,10 @@ export class GeminiChat { let lastError: unknown = new Error('Request failed after all retries.'); let rateLimitRetryCount = 0; - let invalidStreamRetryCount = 0; + let transientInvalidStreamRetryCount = 0; + let protocolTagLeakRetryCount = 0; + const totalInvalidStreamRetryCount = () => + transientInvalidStreamRetryCount + protocolTagLeakRetryCount; let transportStreamRetryCount = 0; let reactiveCompressionAttempted = false; let suppressNextRetryEvent = false; @@ -2257,7 +2263,7 @@ export class GeminiChat { } else if ( attempt > 0 || rateLimitRetryCount > 0 || - invalidStreamRetryCount > 0 || + totalInvalidStreamRetryCount() > 0 || transportStreamRetryCount > 0 ) { yield { type: StreamEventType.RETRY }; @@ -2505,28 +2511,40 @@ export class GeminiChat { break; } - // Transient stream anomalies (NO_FINISH_REASON / NO_RESPONSE_TEXT): - // independent retry budget, similar to rate-limit handling. - // Does NOT consume the content retry budget. - const isTransientStreamError = error instanceof InvalidStreamError; + // Invalid stream responses use an independent retry budget and do + // not consume the content retry budget. + const isInvalidStreamError = error instanceof InvalidStreamError; + const maxInvalidStreamRetries = + isInvalidStreamError && error.type === 'PROTOCOL_TAG_LEAK' + ? INVALID_STREAM_RETRY_CONFIG.protocolTagLeakMaxRetries + : INVALID_STREAM_RETRY_CONFIG.transientMaxRetries; + const invalidStreamRetryCount = + isInvalidStreamError && error.type === 'PROTOCOL_TAG_LEAK' + ? protocolTagLeakRetryCount + : transientInvalidStreamRetryCount; if ( - isTransientStreamError && - invalidStreamRetryCount < INVALID_STREAM_RETRY_CONFIG.maxRetries + isInvalidStreamError && + invalidStreamRetryCount < maxInvalidStreamRetries ) { self.popPendingPartialAssistantTurn(); - invalidStreamRetryCount++; + const nextInvalidStreamRetryCount = invalidStreamRetryCount + 1; + if (error.type === 'PROTOCOL_TAG_LEAK') { + protocolTagLeakRetryCount = nextInvalidStreamRetryCount; + } else { + transientInvalidStreamRetryCount = nextInvalidStreamRetryCount; + } const delayMs = INVALID_STREAM_RETRY_CONFIG.initialDelayMs * - invalidStreamRetryCount; + nextInvalidStreamRetryCount; debugLogger.warn( `Invalid stream [${(error as InvalidStreamError).type}] ` + - `(retry ${invalidStreamRetryCount}/${INVALID_STREAM_RETRY_CONFIG.maxRetries}). ` + + `(retry ${nextInvalidStreamRetryCount}/${maxInvalidStreamRetries}). ` + `Waiting ${delayMs / 1000}s before retrying...`, ); logContentRetry( self.config, new ContentRetryEvent( - invalidStreamRetryCount - 1, + nextInvalidStreamRetryCount - 1, (error as InvalidStreamError).type, delayMs, model, @@ -2538,14 +2556,14 @@ export class GeminiChat { await delay(delayMs, params.config?.abortSignal).promise; continue; } - // Transient budget exhausted — stop immediately. - if (isTransientStreamError) { + // Invalid-stream budget exhausted — stop immediately. + if (isInvalidStreamError) { break; } // Currently unreachable for `InvalidStreamError`. The // `isContentError` predicate is identical to - // `isTransientStreamError` (`error instanceof InvalidStreamError`), + // `isInvalidStreamError` (`error instanceof InvalidStreamError`), // and the transient branch above already either continued or // broke for that class. The branch is preserved as // defense-in-depth: a future error class that should consume @@ -2636,7 +2654,8 @@ export class GeminiChat { type: StreamEventType.RETRY, }, ): AsyncGenerator { - let retryCount = 0; + let transientRetryCount = 0; + let protocolTagLeakRetryCount = 0; for (;;) { const attemptState = buildAttempt(); try { @@ -2654,22 +2673,36 @@ export class GeminiChat { if (!(error instanceof InvalidStreamError)) throw error; attemptState.rollback(); - if (retryCount >= INVALID_STREAM_RETRY_CONFIG.maxRetries) { + const maxContinuationRetries = + error.type === 'PROTOCOL_TAG_LEAK' + ? INVALID_STREAM_RETRY_CONFIG.protocolTagLeakMaxRetries + : INVALID_STREAM_RETRY_CONFIG.transientMaxRetries; + const continuationRetryCount = + error.type === 'PROTOCOL_TAG_LEAK' + ? protocolTagLeakRetryCount + : transientRetryCount; + if (continuationRetryCount >= maxContinuationRetries) { throw error; } - retryCount++; + const nextContinuationRetryCount = continuationRetryCount + 1; + if (error.type === 'PROTOCOL_TAG_LEAK') { + protocolTagLeakRetryCount = nextContinuationRetryCount; + } else { + transientRetryCount = nextContinuationRetryCount; + } const delayMs = - INVALID_STREAM_RETRY_CONFIG.initialDelayMs * retryCount; + INVALID_STREAM_RETRY_CONFIG.initialDelayMs * + nextContinuationRetryCount; debugLogger.warn( `Invalid stream [${error.type}] during output continuation ` + - `(retry ${retryCount}/${INVALID_STREAM_RETRY_CONFIG.maxRetries}). ` + + `(retry ${nextContinuationRetryCount}/${maxContinuationRetries}). ` + `Waiting ${delayMs / 1000}s before retrying...`, ); logContentRetry( self.config, new ContentRetryEvent( - retryCount - 1, + nextContinuationRetryCount - 1, error.type, delayMs, model, @@ -2917,7 +2950,7 @@ export class GeminiChat { if (lastError) { if (lastError instanceof InvalidStreamError) { - const totalAttempts = invalidStreamRetryCount + 1; + const totalAttempts = totalInvalidStreamRetryCount() + 1; logContentRetryFailure( self.config, new ContentRetryFailureEvent( @@ -3799,6 +3832,29 @@ export class GeminiChat { ); } + // Stream validation logic: A stream is considered successful if: + // 1. There's a tool call (tool calls can end without explicit finish reasons), OR + // 2. There's a finish reason AND we have non-empty response text or thought text + // + // Note: Thoughts-only responses are valid for models that use thinking modes. + const hasAnyContent = contentText || thoughtText; + if ( + streamError === null && + !hasToolCall && + (!hasFinishReason || !hasAnyContent) + ) { + if (!hasFinishReason) { + throw new InvalidStreamError( + 'Model stream ended without a finish reason.', + 'NO_FINISH_REASON', + ); + } + throw new InvalidStreamError( + 'Model stream ended with empty response text.', + 'NO_RESPONSE_TEXT', + ); + } + // Record assistant turn with raw Content and metadata. Gate matches // the in-memory `this.history.push` decision below so chat-recording // JSONL never carries a partial turn we deliberately dropped from @@ -3917,31 +3973,6 @@ export class GeminiChat { throw streamError; } - // Stream validation logic: A stream is considered successful if: - // 1. There's a tool call (tool calls can end without explicit finish reasons), OR - // 2. There's a finish reason AND we have non-empty response text or thought text - // - // We throw an error only when there's no tool call AND: - // - No finish reason, OR - // - Empty response text (e.g., no actual content and no thoughts) - // - // Note: Thoughts-only responses are valid for models that use thinking modes - // These models may send only reasoning content without explicit text output. - const hasAnyContent = contentText || thoughtText; - if (!hasToolCall && (!hasFinishReason || !hasAnyContent)) { - if (!hasFinishReason) { - throw new InvalidStreamError( - 'Model stream ended without a finish reason.', - 'NO_FINISH_REASON', - ); - } else { - throw new InvalidStreamError( - 'Model stream ended with empty response text.', - 'NO_RESPONSE_TEXT', - ); - } - } - this.history.push({ role: 'model', parts: [