diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 2b55cacc11f..2c12bdc0ce3 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -1632,6 +1632,11 @@ describe('GeminiChat', async () => { }); it('seeds inherited token count via setLastPromptTokenCount', async () => { + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + authType: AuthType.USE_GEMINI, + model: 'test-model', + contextWindowSize: 200_000, + }); const subagentChat = new GeminiChat(mockConfig, config, [ { role: 'user', parts: [{ text: 'inherited' }] }, { role: 'model', parts: [{ text: 'inherited reply' }] }, @@ -2382,6 +2387,17 @@ describe('GeminiChat', async () => { { role: 'user', parts: [{ text: 'summary' }] }, { role: 'model', parts: [{ text: 'ack' }] }, ]; + const recordChatCompression = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn: vi.fn(), + recordChatCompression, + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); const compressSpy = vi .spyOn(ChatCompressionService.prototype, 'compress') .mockResolvedValueOnce({ @@ -2399,10 +2415,10 @@ describe('GeminiChat', async () => { // Seed lastPromptTokenCount JUST under the 177K hard threshold; the // pending user message adds a handful of estimate-tokens that pushes // effective >= 177K, so the rescue must trigger. - chat.setLastPromptTokenCount(176_999); + chatWithRecording.setLastPromptTokenCount(176_999); const userMessage = 'this is the next user message'; - const stream = await chat.sendMessageStream( + const stream = await chatWithRecording.sendMessageStream( 'test-model', { message: userMessage }, 'prompt-id-hard-rescue-forces', @@ -2426,6 +2442,164 @@ describe('GeminiChat', async () => { (part) => part.text === userMessage, ), ).toBe(true); + expect(recordChatCompression).toHaveBeenCalledTimes(1); + const recordPayload = recordChatCompression.mock.calls[0][0]; + expect(recordPayload.info).toEqual( + expect.objectContaining({ + compressionStatus: CompressionStatus.COMPRESSED, + newTokenCount: 40_000, + }), + ); + expect(recordPayload.compressedHistory).toEqual([ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ]); + }); + + it('rejects before request serialization when oversized resumed history cannot be compressed', async () => { + const oversizedResumedHistory: Content[] = [ + { role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ]; + chat.setHistory(oversizedResumedHistory); + expect(chat.getLastPromptTokenCount()).toBe(0); + + const compressSpy = vi + .spyOn(ChatCompressionService.prototype, 'compress') + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 180_000, + newTokenCount: 180_000, + compressionStatus: + CompressionStatus.COMPRESSION_FAILED_EMPTY_SUMMARY, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue( + new Error('Invalid string length'), + ); + + await expect( + chat.sendMessageStream( + 'test-model', + { message: 'continue' }, + 'prompt-id-oversized-resume-guard', + ), + ).rejects.toThrow( + /compression status: COMPRESSION_FAILED_EMPTY_SUMMARY/i, + ); + + expect(compressSpy).toHaveBeenCalledTimes(1); + expect(compressSpy.mock.calls[0][1].force).toBe(true); + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + expect(chat.getLastPromptTokenCount()).toBe(0); + expect(chat.getHistory()).toHaveLength(2); + }); + + it('rejects before request serialization and restores history when hard-rescue compression is still oversized', async () => { + const originalHistory: Content[] = [ + { role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ]; + const recordChatCompression = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn: vi.fn(), + recordChatCompression, + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + chatWithRecording.setHistory(originalHistory); + chatWithRecording.setLastPromptTokenCount(176_999); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'still large summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 180_000, + newTokenCount: 177_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue( + new Error('Invalid string length'), + ); + + await expect( + chatWithRecording.sendMessageStream( + 'test-model', + { message: 'continue' }, + 'prompt-id-oversized-after-compression', + ), + ).rejects.toThrow(/compression status: COMPRESSED/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + expect(recordChatCompression).not.toHaveBeenCalled(); + expect(chatWithRecording.getLastPromptTokenCount()).toBe(176_999); + expect(chatWithRecording.getHistory()[0].parts?.[0].text).toBe( + originalHistory[0].parts?.[0].text, + ); + }); + + it('rejects when compressed history is below hard but the pending user message pushes it over', async () => { + const originalHistory: Content[] = [ + { role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ]; + const recordChatCompression = vi.fn(); + const chatWithRecording = new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn: vi.fn(), + recordChatCompression, + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + chatWithRecording.setHistory(originalHistory); + chatWithRecording.setLastPromptTokenCount(175_500); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 180_000, + newTokenCount: 176_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + makeStreamResponse('should not send'), + ); + + await expect( + chatWithRecording.sendMessageStream( + 'test-model', + { message: 'x'.repeat(8_000) }, + 'prompt-id-oversized-after-compression-and-user', + ), + ).rejects.toThrow(/Estimated prompt tokens: 178000; hard limit: 177000/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + expect(recordChatCompression).not.toHaveBeenCalled(); + expect(chatWithRecording.getLastPromptTokenCount()).toBe(175_500); + expect(chatWithRecording.getHistory()[0].parts?.[0].text).toBe( + originalHistory[0].parts?.[0].text, + ); }); it('forwards latched consecutiveFailures into hard-rescue (no pre-call reset); success recovers via the post-call branch', async () => { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 759a74b433a..dca02417fe0 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -106,6 +106,38 @@ function isCompressionFailureStatus(status: CompressionStatus): boolean { ); } +function shouldStopAfterHardRescue( + shouldForceFromHard: boolean, + hardLimit: number, + localPromptTokensAfterCompression: number, +): boolean { + return shouldForceFromHard && localPromptTokensAfterCompression >= hardLimit; +} + +function getHardRescueFailureMessage( + effectiveTokens: number, + hardLimit: number, + compressionInfo: ChatCompressionInfo, + localPromptTokensAfterCompression: number, +): string { + const compressionStatus = + CompressionStatus[compressionInfo.compressionStatus] ?? + String(compressionInfo.compressionStatus); + const tokenCount = + compressionInfo.compressionStatus === CompressionStatus.COMPRESSED + ? Math.max( + compressionInfo.newTokenCount, + localPromptTokensAfterCompression, + ) + : Math.max(effectiveTokens, localPromptTokensAfterCompression); + return ( + `Context is too large to send safely after automatic compression. ` + + `Estimated prompt tokens: ${tokenCount}; hard limit: ${hardLimit}; ` + + `compression status: ${compressionStatus}. ` + + `Start a new session or reduce the resumed history before continuing.` + ); +} + export enum StreamEventType { /** A regular content chunk from the API. */ CHUNK = 'chunk', @@ -157,6 +189,11 @@ interface TryCompressOptions { * `getHistory(true)` clone per send. (review #4168 R1.3 / R1.4) */ precomputedEffectiveTokens?: number; + /** + * Delay writing the compression checkpoint until the caller has run any + * post-compression guards that may roll the in-memory chat state back. + */ + deferChatCompressionRecord?: boolean; } const INVALID_CONTENT_RETRY_OPTIONS: ContentRetryOptions = { @@ -1325,9 +1362,11 @@ export class GeminiChat { * * Returns the compression info regardless of outcome. On a successful * compaction (`COMPRESSED`), this method has already mutated the chat's - * history, recorded the event to `chatRecordingService` (if wired), and - * updated both the per-chat token count and (when wired) the global - * telemetry singleton. + * history, recorded the event to `chatRecordingService` (if wired and + * unless `options.deferChatCompressionRecord` is set), and updated both + * the per-chat token count and (when wired) the global telemetry singleton. + * Deferred callers are responsible for recording after their own + * post-compression guards pass. */ async tryCompress( promptId: string, @@ -1352,10 +1391,12 @@ export class GeminiChat { }); if (info.compressionStatus === CompressionStatus.COMPRESSED && newHistory) { - this.chatRecordingService?.recordChatCompression({ - info, - compressedHistory: newHistory, - }); + if (!options?.deferChatCompressionRecord) { + this.chatRecordingService?.recordChatCompression({ + info, + compressedHistory: newHistory, + }); + } this.setHistory(newHistory); debugLogger.debug('[FILE_READ_CACHE] clear after auto tryCompress'); this.config.getFileReadCache().clear(); @@ -1525,9 +1566,13 @@ export class GeminiChat { imageTokenEstimate, ); const shouldForceFromHard = effectiveTokens >= hard; + const historyBeforeHardRescue = shouldForceFromHard + ? this.getHistoryShallow() + : undefined; + const lastPromptTokenCountBeforeHardRescue = this.lastPromptTokenCount; if (shouldForceFromHard) { debugLogger.warn( - `[compaction] hard-tier rescue triggered: effectiveTokens=${effectiveTokens}, hard=${hard}, consecutiveFailures=${this.consecutiveFailures}.`, + `[compaction] hard-tier rescue triggered: prompt_id=${prompt_id}, effectiveTokens=${effectiveTokens}, hard=${hard}, consecutiveFailures=${this.consecutiveFailures}.`, ); } @@ -1539,6 +1584,7 @@ export class GeminiChat { { pendingUserMessage: userContent, precomputedEffectiveTokens: effectiveTokens, + deferChatCompressionRecord: shouldForceFromHard, // Hard-rescue is force=true to bypass the cheap-gate breaker // but it's an AUTOMATIC trigger. Explicit trigger='auto' tells // the service to skip the manual-only orphan-strip that would @@ -1550,6 +1596,66 @@ export class GeminiChat { }, ); + const localPromptTokensAfterCompression = shouldForceFromHard + ? estimatePromptTokens( + this.lastPromptTokenCount > 0 ? [] : this.getHistoryShallow(true), + userContent, + this.lastPromptTokenCount, + imageTokenEstimate, + ) + : 0; + if ( + shouldStopAfterHardRescue( + shouldForceFromHard, + hard, + localPromptTokensAfterCompression, + ) + ) { + const message = getHardRescueFailureMessage( + effectiveTokens, + hard, + compressionInfo, + localPromptTokensAfterCompression, + ); + if ( + compressionInfo.compressionStatus === CompressionStatus.COMPRESSED && + historyBeforeHardRescue + ) { + // Hard-rescue compression mutates in-memory history before this + // guard can compare the compressed prompt size. If the compressed + // prompt is still too large to send, restore the pre-compression + // state. The JSONL compression checkpoint is intentionally not + // written because the send is about to be rejected. + this.setHistory(historyBeforeHardRescue); + this.lastPromptTokenCount = lastPromptTokenCountBeforeHardRescue; + this.telemetryService?.setLastPromptTokenCount( + lastPromptTokenCountBeforeHardRescue, + ); + } + const compressionStatus = + CompressionStatus[compressionInfo.compressionStatus] ?? + String(compressionInfo.compressionStatus); + debugLogger.warn( + `[compaction] hard-tier rescue stopped oversized prompt: ` + + `prompt_id=${prompt_id}, effectiveTokens=${effectiveTokens}, ` + + `hard=${hard}, localPromptTokensAfterCompression=` + + `${localPromptTokensAfterCompression}, compressionStatus=` + + `${compressionStatus}, newTokenCount=` + + `${compressionInfo.newTokenCount}, consecutiveFailures=` + + `${this.consecutiveFailures}. ${message}`, + ); + throw new Error(message); + } + if ( + shouldForceFromHard && + compressionInfo.compressionStatus === CompressionStatus.COMPRESSED + ) { + this.chatRecordingService?.recordChatCompression({ + info: compressionInfo, + compressedHistory: this.getHistoryShallow(), + }); + } + // Add user content to history ONCE before any attempts. this.history.push(userContent); userContentAdded = true;