diff --git a/packages/cli/src/ui/commands/contextCommand.ts b/packages/cli/src/ui/commands/contextCommand.ts index 9d574d520c0..905deb601b4 100644 --- a/packages/cli/src/ui/commands/contextCommand.ts +++ b/packages/cli/src/ui/commands/contextCommand.ts @@ -326,7 +326,7 @@ export async function collectContextData( // single render — that resolves the moment any send happens. // // TODO: plumb the chat history into collectContextData and use - // estimatePromptTokens(history, undefined, 0, imageTokenEstimate) here + // estimatePromptTokens(history, undefined, 0, 0, imageTokenEstimate) here // for same-source-of-truth as the cheap-gate. Defer because Config // doesn't expose the active chat instance today. const tierTokens = isEstimated ? rawOverhead : apiTotalTokens; diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 5b022f38751..4cecabd4403 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -32,7 +32,7 @@ import { } from './contentGenerator.js'; import { BaseLlmClient } from './baseLlmClient.js'; import { buildAgentContentGeneratorConfig } from '../models/content-generator-config.js'; -import { type GeminiChat } from './geminiChat.js'; +import { GeminiChat } from './geminiChat.js'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../config/config.js'; import { @@ -598,6 +598,47 @@ describe('Gemini Client (client.ts)', () => { expect(resumedClient.getChat().getLastPromptTokenCount()).toBe(123_456); }); + it('seeds resumed chat with previous response output token count', async () => { + const seedResumeTokenCountsSpy = vi.spyOn( + GeminiChat.prototype, + 'seedResumeTokenCounts', + ); + vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ + conversation: { + sessionId: 'resumed-session-id', + projectHash: 'project-hash', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + messages: [ + { + uuid: 'assistant-1', + parentUuid: null, + sessionId: 'resumed-session-id', + timestamp: new Date(0).toISOString(), + type: 'assistant', + cwd: '/test/project', + version: '1.0.0', + message: { role: 'model', parts: [{ text: 'done' }] }, + usageMetadata: { + promptTokenCount: 200, + candidatesTokenCount: 60, + thoughtsTokenCount: 20, + totalTokenCount: 280, + }, + }, + ], + }, + filePath: '/test/session.jsonl', + lastCompletedUuid: null, + }); + + const resumedClient = new GeminiClient(mockConfig); + await resumedClient.initialize(); + + expect(resumedClient.getChat().getLastPromptTokenCount()).toBe(200); + expect(seedResumeTokenCountsSpy).toHaveBeenCalledWith(200, 80); + }); + it('seeds recently completed tools from resumed history', async () => { vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ conversation: { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index fbe56153536..45a132ec122 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -290,7 +290,7 @@ export class GeminiClient { // Check if we're resuming from a previous session const resumedSessionData = this.config.getResumedSessionData(); if (resumedSessionData) { - replayUiTelemetryFromConversation( + const resumeTokenCounts = replayUiTelemetryFromConversation( resumedSessionData.conversation, this.config.getSessionId(), ); @@ -304,9 +304,17 @@ export class GeminiClient { resumedHistory, sessionStartSource ?? SessionStartSource.Resume, ); - this.getChat().setLastPromptTokenCount( - uiTelemetryService.getLastPromptTokenCount(), - ); + const chat = this.getChat(); + if (resumeTokenCounts) { + chat.seedResumeTokenCounts( + resumeTokenCounts.promptTokenCount, + resumeTokenCounts.outputTokenCount, + ); + } else { + chat.setLastPromptTokenCount( + uiTelemetryService.getLastPromptTokenCount(), + ); + } // Restore attribution state from the last snapshot in the session this.restoreAttributionFromSession(resumedSessionData.conversation); diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 017065765c8..71bc82b460d 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -8,13 +8,10 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import type { Content, GenerateContentConfig, - Part, -} from '@google/genai'; -import { - ApiError, - FinishReason, GenerateContentResponse, + Part, } from '@google/genai'; +import { ApiError } from '@google/genai'; import { AuthType, type ContentGenerator } from '../core/contentGenerator.js'; import { GeminiChat, @@ -32,6 +29,10 @@ import { ChatCompressionService, MAX_CONSECUTIVE_FAILURES, } from '../services/chatCompressionService.js'; +import { + estimateContentTokens, + estimatePromptTokens, +} from '../services/tokenEstimation.js'; import { SYSTEM_REMINDER_OPEN } from '../utils/environmentContext.js'; import { SessionStartSource } from '../hooks/types.js'; @@ -80,24 +81,6 @@ const { mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({ mockLogContentRetryFailure: vi.fn(), })); -const { mockDebugLoggerWarn } = vi.hoisted(() => ({ - mockDebugLoggerWarn: vi.fn(), -})); - -vi.mock('../utils/debugLogger.js', async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createDebugLogger: vi.fn(() => ({ - debug: vi.fn(), - info: vi.fn(), - warn: mockDebugLoggerWarn, - error: vi.fn(), - })), - }; -}); - vi.mock('../telemetry/loggers.js', () => ({ logContentRetry: mockLogContentRetry, logContentRetryFailure: mockLogContentRetryFailure, @@ -124,6 +107,24 @@ vi.mock('../services/sleepInhibitor.js', () => ({ acquireSleepInhibitor: mockAcquireSleepInhibitor, })); +const { mockDebugLoggerWarn } = vi.hoisted(() => ({ + mockDebugLoggerWarn: vi.fn(), +})); + +vi.mock('../utils/debugLogger.js', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createDebugLogger: () => ({ + debug: vi.fn(), + info: vi.fn(), + warn: mockDebugLoggerWarn, + error: vi.fn(), + }), + }; +}); + describe('GeminiChat', async () => { let mockContentGenerator: ContentGenerator; let chat: GeminiChat; @@ -452,238 +453,6 @@ describe('GeminiChat', async () => { expect(modelTurn?.parts![0]!.functionCall).toBeDefined(); }); - it('suffixes cross-turn reused functionCall ids before yielding and recording history', async () => { - chat = new GeminiChat( - mockConfig, - config, - [ - { role: 'user', parts: [{ text: 'first' }] }, - { - role: 'model', - parts: [ - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: { file_path: 'a.ts' }, - }, - }, - ], - }, - { - role: 'user', - parts: [ - { - functionResponse: { - id: 'dup_id_0001', - name: 'read_file', - response: { output: 'A' }, - }, - }, - ], - }, - ], - undefined, - uiTelemetryService, - ); - - vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( - (async function* () { - yield { - functionCalls: [ - { - id: 'dup_id_0001', - name: 'read_file', - args: { file_path: 'b.ts' }, - }, - ], - candidates: [ - { - content: { - role: 'model', - parts: [ - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: { file_path: 'b.ts' }, - }, - }, - ], - }, - finishReason: 'STOP', - }, - ], - } as unknown as GenerateContentResponse; - })(), - ); - - const stream = await chat.sendMessageStream( - 'test-model', - { message: 'second' }, - 'prompt-id-dup-tool-call', - ); - const events: StreamEvent[] = []; - for await (const event of stream) { - events.push(event); - } - - const chunk = events.find((event) => event.type === StreamEventType.CHUNK) - ?.value as GenerateContentResponse | undefined; - expect(chunk?.functionCalls?.map((call) => call.id)).toEqual([ - 'dup_id_0001__qwen_dup_2', - ]); - expect( - chunk?.candidates?.[0]?.content?.parts?.map( - (part) => part.functionCall?.id, - ), - ).toEqual(['dup_id_0001__qwen_dup_2']); - - const history = chat.getHistory(); - expect(history.at(-1)?.parts?.[0]?.functionCall?.id).toBe( - 'dup_id_0001__qwen_dup_2', - ); - }); - - it('normalizes ids visible through the real GenerateContentResponse functionCalls getter', async () => { - chat = new GeminiChat( - mockConfig, - config, - [ - { - role: 'model', - parts: [ - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: { file_path: 'a.ts' }, - }, - }, - ], - }, - { - role: 'user', - parts: [ - { - functionResponse: { - id: 'dup_id_0001', - name: 'read_file', - response: { output: 'A' }, - }, - }, - ], - }, - ], - undefined, - uiTelemetryService, - ); - const response = new GenerateContentResponse(); - response.candidates = [ - { - content: { - role: 'model', - parts: [ - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: { file_path: 'b.ts' }, - }, - }, - ], - }, - finishReason: FinishReason.STOP, - }, - ]; - vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( - (async function* () { - yield response; - })(), - ); - - const stream = await chat.sendMessageStream( - 'test-model', - { message: 'second' }, - 'prompt-id-real-response-getter', - ); - const events: StreamEvent[] = []; - for await (const event of stream) { - events.push(event); - } - - const chunk = events.find((event) => event.type === StreamEventType.CHUNK) - ?.value as GenerateContentResponse | undefined; - expect(chunk?.functionCalls?.map((call) => call.id)).toEqual([ - 'dup_id_0001__qwen_dup_2', - ]); - }); - - it('drops same-turn replayed functionCall ids before yielding and recording history', async () => { - vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( - (async function* () { - yield { - functionCalls: [ - { id: 'dup_id_0001', name: 'read_file', args: {} }, - { id: 'dup_id_0001', name: 'read_file', args: {} }, - ], - candidates: [ - { - content: { - role: 'model', - parts: [ - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: {}, - }, - }, - { - functionCall: { - id: 'dup_id_0001', - name: 'read_file', - args: {}, - }, - }, - ], - }, - finishReason: 'STOP', - }, - ], - } as unknown as GenerateContentResponse; - })(), - ); - - const stream = await chat.sendMessageStream( - 'test-model', - { message: 'run once' }, - 'prompt-id-same-turn-dup-tool-call', - ); - const events: StreamEvent[] = []; - for await (const event of stream) { - events.push(event); - } - - const chunk = events.find((event) => event.type === StreamEventType.CHUNK) - ?.value as GenerateContentResponse | undefined; - expect(chunk?.functionCalls?.map((call) => call.id)).toEqual([ - 'dup_id_0001', - ]); - expect( - chunk?.candidates?.[0]?.content?.parts?.map( - (part) => part.functionCall?.id, - ), - ).toEqual(['dup_id_0001']); - - const functionCallIds = chat - .getHistory() - .at(-1) - ?.parts?.map((part) => part.functionCall?.id) - .filter((id): id is string => Boolean(id)); - expect(functionCallIds).toEqual(['dup_id_0001']); - }); - it('should fail if the stream ends with an empty part and has no finishReason', async () => { vi.useFakeTimers(); try { @@ -2054,7 +1823,10 @@ describe('GeminiChat', async () => { }); describe('auto-compression integration', () => { - function makeStreamResponse(text = 'ok') { + function makeStreamResponse( + text = 'ok', + usageMetadata?: GenerateContentResponse['usageMetadata'], + ) { return (async function* () { yield { candidates: [ @@ -2065,6 +1837,7 @@ describe('GeminiChat', async () => { safetyRatings: [], }, ], + usageMetadata, text: () => text, } as unknown as GenerateContentResponse; })(); @@ -2885,7 +2658,10 @@ describe('GeminiChat', async () => { // 2) call tryCompress with force=true (so MAX_CONSECUTIVE_FAILURES does // not gate the only attempt that can save the next round-trip). describe('sendMessageStream hard-tier rescue', () => { - function makeStreamResponse(text = 'ok') { + function makeStreamResponse( + text = 'ok', + usageMetadata?: GenerateContentResponse['usageMetadata'], + ) { return (async function* () { yield { candidates: [ @@ -2896,6 +2672,7 @@ describe('GeminiChat', async () => { safetyRatings: [], }, ], + usageMetadata, text: () => text, } as unknown as GenerateContentResponse; })(); @@ -3082,133 +2859,520 @@ describe('GeminiChat', async () => { ); }); - it('stops hard-rescue after repeated compressed results are still oversized', async () => { + 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: vi.fn(), + recordChatCompression, } as unknown as ConstructorParameters[3], uiTelemetryService, ); chatWithRecording.setHistory(originalHistory); - chatWithRecording.setLastPromptTokenCount(176_999); + chatWithRecording.setLastPromptTokenCount(175_500); - const compressSpy = vi - .spyOn(ChatCompressionService.prototype, 'compress') - .mockResolvedValue({ - newHistory: [ - { role: 'user', parts: [{ text: 'still large summary' }] }, - { role: 'model', parts: [{ text: 'ack' }] }, - ], - info: { - originalTokenCount: 180_000, - newTokenCount: 177_000, + 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('does not treat the image token estimate as output tokens after hard-rescue compression', 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_500); + + const compressSpy = 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('sent after compression'), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'x'.repeat(3_000) }, + 'prompt-id-hard-rescue-image-estimate-slot', + ); + for await (const _ of stream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(1); + expect(compressSpy.mock.calls[0][1].force).toBe(true); + expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( + 1, + ); + expect(recordChatCompression).toHaveBeenCalledTimes(1); + expect(chatWithRecording.getLastPromptTokenCount()).toBe(176_000); + }); + + it('includes previous response output tokens in the hard-tier estimate', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 50_000, + newTokenCount: 50_000, + compressionStatus: CompressionStatus.NOOP, + }, + }) + .mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 176_000, + newTokenCount: 40_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 176_000, + candidatesTokenCount: 1_500, + totalTokenCount: 177_500, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('after rescue')); + + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'prime the token counters' }, + 'prompt-prime-candidates', + ); + for await (const _ of firstStream) { + /* consume */ + } + + const rescueStream = await chat.sendMessageStream( + 'test-model', + { message: 'small follow-up' }, + 'prompt-hard-rescue-candidates', + ); + for await (const _ of rescueStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[0][1].force).toBe(false); + expect(compressSpy.mock.calls[1][1].force).toBe(true); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeGreaterThanOrEqual(177_000); + }); + + it('does not double-count output tokens when prompt count falls back to total token count', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy.mockResolvedValue({ + newHistory: null, + info: { + originalTokenCount: 176_000, + newTokenCount: 176_000, + compressionStatus: CompressionStatus.NOOP, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + candidatesTokenCount: 1_500, + totalTokenCount: 176_000, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('second')); + + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'prime fallback token counters' }, + 'prompt-prime-total-token-fallback', + ); + for await (const _ of firstStream) { + /* consume */ + } + + const secondStream = await chat.sendMessageStream( + 'test-model', + { message: 'small follow-up' }, + 'prompt-total-token-fallback-follow-up', + ); + for await (const _ of secondStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[1][1].force).toBe(false); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeLessThan(177_000); + }); + + it('includes previous response thought tokens in the hard-tier estimate', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 50_000, + newTokenCount: 50_000, + compressionStatus: CompressionStatus.NOOP, + }, + }) + .mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 176_000, + newTokenCount: 40_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 176_000, + candidatesTokenCount: 500, + thoughtsTokenCount: 1_000, + totalTokenCount: 177_500, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('after rescue')); + + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'prime thought token counters' }, + 'prompt-prime-thought-tokens', + ); + for await (const _ of firstStream) { + /* consume */ + } + + const rescueStream = await chat.sendMessageStream( + 'test-model', + { message: 'small follow-up' }, + 'prompt-hard-rescue-thought-tokens', + ); + for await (const _ of rescueStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[1][1].force).toBe(true); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeGreaterThanOrEqual(177_000); + }); + + it('includes disjoint candidate and thought tokens when total token count is unavailable', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 50_000, + newTokenCount: 50_000, + compressionStatus: CompressionStatus.NOOP, + }, + }) + .mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 176_000, + newTokenCount: 40_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 176_000, + candidatesTokenCount: 1_200, + thoughtsTokenCount: 300, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('after rescue')); + + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'prime disjoint output token counters' }, + 'prompt-prime-disjoint-output-tokens', + ); + for await (const _ of firstStream) { + /* consume */ + } + + const rescueStream = await chat.sendMessageStream( + 'test-model', + { message: 'small follow-up' }, + 'prompt-hard-rescue-disjoint-output-tokens', + ); + for await (const _ of rescueStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[1][1].force).toBe(true); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeGreaterThanOrEqual(177_000); + }); + + it('does not double-count OpenAI-compatible reasoning tokens already included in candidates', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy.mockResolvedValue({ + newHistory: null, + info: { + originalTokenCount: 176_400, + newTokenCount: 176_400, + compressionStatus: CompressionStatus.NOOP, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 175_400, + candidatesTokenCount: 1_000, + thoughtsTokenCount: 500, + totalTokenCount: 176_400, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('second')); + + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'prime OpenAI-compatible reasoning token counters' }, + 'prompt-prime-openai-reasoning-tokens', + ); + for await (const _ of firstStream) { + /* consume */ + } + + const secondStream = await chat.sendMessageStream( + 'test-model', + { message: 'small follow-up' }, + 'prompt-openai-reasoning-follow-up', + ); + for await (const _ of secondStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[1][1].force).toBe(false); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeLessThan(177_000); + }); + + it('resets previous response output tokens when seeding last prompt tokens externally', async () => { + const compressSpy = vi + .spyOn(ChatCompressionService.prototype, 'compress') + .mockResolvedValue({ + newHistory: null, + info: { + originalTokenCount: 176_000, + newTokenCount: 176_000, + compressionStatus: CompressionStatus.NOOP, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 10_000, + candidatesTokenCount: 5_000, + totalTokenCount: 15_000, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('second')); + + const firstStream = await chat.sendMessageStream( + 'test-model', + { message: 'collect candidates' }, + 'prompt-collect-candidates', + ); + for await (const _ of firstStream) { + /* consume */ + } + + chat.setLastPromptTokenCount(176_000); + const secondStream = await chat.sendMessageStream( + 'test-model', + { message: 'seeded follow-up' }, + 'prompt-seeded-after-candidates', + ); + for await (const _ of secondStream) { + /* consume */ + } + + expect(compressSpy).toHaveBeenCalledTimes(2); + expect(compressSpy.mock.calls[1][1].force).toBe(false); + expect( + compressSpy.mock.calls[1][1].precomputedEffectiveTokens, + ).toBeLessThan(177_000); + }); + + it('resets previous response output tokens after successful compression', async () => { + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 50_000, + newTokenCount: 50_000, + compressionStatus: CompressionStatus.NOOP, + }, + }) + .mockResolvedValueOnce({ + newHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + info: { + originalTokenCount: 176_000, + newTokenCount: 40_000, compressionStatus: CompressionStatus.COMPRESSED, }, + }) + .mockResolvedValueOnce({ + newHistory: null, + info: { + originalTokenCount: 40_000, + newTokenCount: 40_000, + compressionStatus: CompressionStatus.NOOP, + }, }); - vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( - makeStreamResponse('after bounded compressed hard-rescue'), - ); - - for (let i = 0; i < MAX_CONSECUTIVE_FAILURES; i++) { - await expect( - chatWithRecording.sendMessageStream( - 'test-model', - { message: `still-oversized-after-compression-${i}` }, - `prompt-hard-rescue-compressed-bound-${i}`, - ), - ).rejects.toThrow(/compression status: COMPRESSED/i); - } + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + makeStreamResponse('first', { + promptTokenCount: 176_000, + candidatesTokenCount: 100_000, + totalTokenCount: 276_000, + }), + ) + .mockResolvedValueOnce(makeStreamResponse('after compression')) + .mockResolvedValueOnce(makeStreamResponse('after reset')); - const callsBeforeBound = compressSpy.mock.calls.length; - const stream = await chatWithRecording.sendMessageStream( + chat.setLastPromptTokenCount(50_000); + const firstStream = await chat.sendMessageStream( 'test-model', - { message: 'send after bounded compressed hard-rescue' }, - 'prompt-hard-rescue-after-compressed-bound', + { message: 'prime output tokens' }, + 'prompt-prime-compression-reset', ); - for await (const _ of stream) { + for await (const _ of firstStream) { /* consume */ } - expect(compressSpy).toHaveBeenCalledTimes(callsBeforeBound); - expect(callsBeforeBound).toBe(MAX_CONSECUTIVE_FAILURES); - expect(compressSpy.mock.calls.map(([, opts]) => opts.force)).toEqual( - Array(MAX_CONSECUTIVE_FAILURES).fill(true), - ); - expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( - 1, - ); - expect(mockDebugLoggerWarn).toHaveBeenCalledWith( - expect.stringContaining('hardRescueFailureCount=1'), - ); - expect(mockDebugLoggerWarn).toHaveBeenCalledWith( - expect.stringContaining('hard-tier rescue skipped'), - ); - expect(mockDebugLoggerWarn).toHaveBeenCalledWith( - expect.stringContaining( - 'prompt_id=prompt-hard-rescue-after-compressed-bound', - ), - ); - }); - - 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, + const rescueStream = await chat.sendMessageStream( + 'test-model', + { message: 'trigger compression' }, + 'prompt-compression-reset-rescue', ); - chatWithRecording.setHistory(originalHistory); - chatWithRecording.setLastPromptTokenCount(175_500); + for await (const _ of rescueStream) { + /* consume */ + } - 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'), + const followUpStream = await chat.sendMessageStream( + 'test-model', + { message: 'after compression reset' }, + 'prompt-after-compression-reset', ); + for await (const _ of followUpStream) { + /* consume */ + } - 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, - ); + expect(compressSpy).toHaveBeenCalledTimes(3); + expect(compressSpy.mock.calls[1][1].force).toBe(true); + expect( + compressSpy.mock.calls[2][1].precomputedEffectiveTokens, + ).toBeLessThan(100_000); }); it('stops pre-send hard-rescue after repeated failed hard-tier compactions', async () => { @@ -3408,6 +3572,50 @@ describe('GeminiChat', async () => { ); }); + it('does not replace token counters when usage reports zero prompt tokens', async () => { + vi.spyOn(ChatCompressionService.prototype, 'compress').mockResolvedValue({ + newHistory: null, + info: { + originalTokenCount: 123_456, + newTokenCount: 123_456, + compressionStatus: CompressionStatus.NOOP, + }, + }); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + makeStreamResponse('zero prompt count', { + promptTokenCount: 0, + totalTokenCount: 5000, + }), + ); + + chat.setLastPromptTokenCount(123_456); + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'zero prompt count should not reseed' }, + 'prompt-zero-count-no-reseed', + ); + for await (const _ of stream) { + /* consume */ + } + + expect(chat.getLastPromptTokenCount()).toBe(123_456); + }); + + it('ignores previous response output tokens when the prompt token count is zero', () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'history question' }] }, + { role: 'model', parts: [{ text: 'history answer' }] }, + ]; + const userMessage: Content = { + role: 'user', + parts: [{ text: 'follow-up question' }], + }; + + expect(estimatePromptTokens(history, userMessage, 0, 9999)).toBe( + estimateContentTokens([...history, userMessage]), + ); + }); + it('forwards latched consecutiveFailures into hard-rescue (no pre-call reset); success recovers via the post-call branch', async () => { // Hard-rescue uses force=true, which already bypasses the // chatCompressionService breaker (the `!force` check in compress's @@ -4579,75 +4787,6 @@ describe('GeminiChat', async () => { } }); - it('should pass configured retry error codes into streamed retry diagnostics', async () => { - vi.useFakeTimers(); - - try { - vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ - authType: AuthType.USE_OPENAI, - model: 'test-model', - retryErrorCodes: [4999], - }); - const providerThrottle = Object.assign( - new StreamContentError('Provider-specific throttle'), - { status: 4999 }, - ); - - vi.mocked(mockContentGenerator.generateContentStream) - .mockResolvedValueOnce( - (async function* () { - throw providerThrottle; - - yield {} as GenerateContentResponse; - })(), - ) - .mockResolvedValueOnce( - (async function* () { - yield { - candidates: [ - { - content: { - parts: [{ text: 'Success after custom code retry' }], - }, - finishReason: 'STOP', - }, - ], - } as unknown as GenerateContentResponse; - })(), - ); - - const stream = await chat.sendMessageStream( - 'test-model', - { message: 'test' }, - 'prompt-id-custom-retry-code', - ); - - const iterator = stream[Symbol.asyncIterator](); - const first = await iterator.next(); - expect(first.value.type).toBe(StreamEventType.RETRY); - - const secondPromise = iterator.next(); - await vi.advanceTimersByTimeAsync(60_000); - await secondPromise; - - for (;;) { - const next = await iterator.next(); - if (next.done) break; - } - - expect(mockDebugLoggerWarn).toHaveBeenCalledWith( - 'Rate limit retry scheduled', - expect.objectContaining({ - classificationDiagnosis: 'retryable', - classificationReason: 'rate-limit', - errorKind: 'provider', - }), - ); - } finally { - vi.useRealTimers(); - } - }); - it('should retry immediately when skipDelay is called during rate-limit wait', async () => { vi.useFakeTimers(); @@ -6485,37 +6624,6 @@ describe('GeminiChat', async () => { ); }); - it('should not re-escalate when the request already uses the escalated output limit', async () => { - vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( - makeStream([makeChunk([{ text: 'still partial' }], 'MAX_TOKENS')]), - ); - - const stream = await chat.sendMessageStream( - 'gemini-3-pro', - { - message: 'continue the agent task', - config: { maxOutputTokens: 65_536 }, - }, - 'prompt-sticky-escalation', - ); - - const events: StreamEvent[] = []; - for await (const event of stream) { - events.push(event); - } - - expect( - events.some( - (event) => - event.type === StreamEventType.RETRY && - event.maxOutputTokensEscalated !== undefined, - ), - ).toBe(false); - expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes( - 1, - ); - }); - it('should coalesce overlapping recovery continuation text', async () => { const streams = [ makeStream([makeChunk([{ text: 'discarded initial' }], 'MAX_TOKENS')]), @@ -8090,138 +8198,4 @@ describe('GeminiChat', async () => { expect(compressSpy.mock.calls[3][1].consecutiveFailures).toBe(0); }); }); - - describe('compressFast', () => { - const userMsg = (text: string): Content => ({ - role: 'user' as const, - parts: [{ text }], - }); - const modelMsg = (text: string): Content => ({ - role: 'model' as const, - parts: [{ text }], - }); - const modelMsgWithThinking = ( - text: string | null, - thinking: string, - ): Content => ({ - role: 'model' as const, - parts: [{ thought: true, text: thinking }, ...(text ? [{ text }] : [])], - }); - const toolCall = (name: string): Content => ({ - role: 'model' as const, - parts: [{ functionCall: { name, args: {} } }], - }); - const toolResult = (name: string, output: string): Content => ({ - role: 'user' as const, - parts: [{ functionResponse: { name, response: { output } } }], - }); - - beforeEach(() => { - (mockConfig as unknown as Record)[ - 'getClearContextOnIdle' - ] = vi.fn().mockReturnValue({ - toolResultsThresholdMinutes: 60, - toolResultsNumToKeep: 5, - }); - }); - - it('strips thinking from model turns', () => { - chat.setHistory([ - userMsg('hello'), - modelMsgWithThinking('response text', 'internal reasoning'), - userMsg('next'), - modelMsg('plain reply'), - ]); - chat.setLastPromptTokenCount(1000); - - const result = chat.compressFast(); - - expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); - // Thinking parts should be stripped from the first model message - const history = chat.getHistory(); - const firstModel = history.find((c) => c.role === 'model'); - expect(firstModel?.parts).toEqual([{ text: 'response text' }]); - }); - - it('NOOP when nothing is compressible', () => { - chat.setHistory([ - userMsg('hello'), - modelMsg('hi'), - userMsg('how are you'), - modelMsg('good'), - ]); - chat.setLastPromptTokenCount(1000); - - const result = chat.compressFast(); - - expect(result.info.compressionStatus).toBe(CompressionStatus.NOOP); - }); - - it('clears old tool results via microcompaction', () => { - const history: Content[] = []; - // Create many tool calls so keepRecent=5 kicks in - for (let i = 0; i < 8; i++) { - history.push(toolCall('read_file')); - history.push( - toolResult('read_file', `content for file ${i} `.repeat(50)), - ); - } - history.push(userMsg('final')); - history.push(modelMsg('done')); - chat.setHistory(history); - chat.setLastPromptTokenCount(5000); - - const result = chat.compressFast(); - - expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); - expect(result.microcompactMeta).toBeDefined(); - expect(result.microcompactMeta!.toolsCleared).toBeGreaterThan(0); - }); - - it('adjusts lastPromptTokenCount by estimated delta on COMPRESSED', () => { - const history: Content[] = []; - for (let i = 0; i < 5; i++) { - history.push( - userMsg(`question ${i}`), - modelMsgWithThinking( - `response ${i}`, - `very long internal reasoning for turn ${i} `.repeat(100), - ), - ); - } - chat.setHistory(history); - const apiBaseline = 50000; - chat.setLastPromptTokenCount(apiBaseline); - - const result = chat.compressFast(); - - expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); - expect(result.info.newTokenCount).toBeLessThan(apiBaseline); - expect(result.info.newTokenCount).toBeGreaterThan(0); - expect(chat.getLastPromptTokenCount()).toBe(result.info.newTokenCount); - }); - - it('falls back to estimateContentTokens when lastPromptTokenCount is 0', () => { - const history: Content[] = []; - for (let i = 0; i < 5; i++) { - history.push( - userMsg(`question ${i}`), - modelMsgWithThinking( - `response ${i}`, - `very long internal reasoning for turn ${i} `.repeat(100), - ), - ); - } - chat.setHistory(history); - chat.setLastPromptTokenCount(0); - - const result = chat.compressFast(); - - expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); - expect(result.info.originalTokenCount).toBeGreaterThan(0); - expect(result.info.newTokenCount).toBeLessThan( - result.info.originalTokenCount, - ); - }); - }); }); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 1bf532bcaec..54ba4eb0a64 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -59,6 +59,7 @@ import { resolveSlimmingConfig } from '../services/compactionInputSlimming.js'; import { estimateContentTokens, estimatePromptTokens, + getUsageOutputTokenCountForPromptEstimate, } from '../services/tokenEstimation.js'; import { microcompactHistory, @@ -1363,6 +1364,14 @@ export class GeminiChat { */ private lastPromptTokenCount = 0; + /** + * Per-chat output-token count from the previous model response. The + * previous response is appended to local history after `promptTokenCount` + * was reported, so steady-state prompt estimates add this value to avoid + * under-counting the next request near the hard compaction threshold. + */ + private lastOutputTokenCount = 0; + /** * Number of consecutive auto-compaction failures for this chat. The * cheap-gate NOOPs once this reaches MAX_CONSECUTIVE_FAILURES (default 3) @@ -1467,10 +1476,32 @@ export class GeminiChat { * history (forks, subagents, speculation). Without this, the auto-compress * threshold check sees `0` and refuses to compress — so the first API call * can 400 from oversized history. Callers pass the parent chat's - * `getLastPromptTokenCount()` here. + * `getLastPromptTokenCount()` here. This also clears any remembered + * previous-response output token count because the seeded prompt count + * comes from a different chat instance and should not inherit this chat's + * last response size. */ setLastPromptTokenCount(count: number): void { this.lastPromptTokenCount = count; + this.lastOutputTokenCount = 0; + } + + /** + * Seed the restored prompt and previous-response output token counts in one + * step. Resume restores chat history plus both counters from the same + * assistant usage record, so callers must avoid the normal + * setLastPromptTokenCount() clearing behavior. + */ + seedResumeTokenCounts( + promptTokenCount: number, + outputTokenCount: number, + ): void { + this.lastPromptTokenCount = Number.isFinite(promptTokenCount) + ? Math.max(0, promptTokenCount) + : 0; + this.lastOutputTokenCount = Number.isFinite(outputTokenCount) + ? Math.max(0, outputTokenCount) + : 0; } /** @@ -1519,6 +1550,7 @@ export class GeminiChat { this.config.getFileReadCache().clear(); clearDetailedSpanState(); this.lastPromptTokenCount = info.newTokenCount; + this.lastOutputTokenCount = 0; this.telemetryService?.setLastPromptTokenCount(info.newTokenCount); // Reset the consecutive-failure counter on success so a forced /compress // (or any successful compaction) recovers a chat whose breaker had @@ -1733,9 +1765,10 @@ export class GeminiChat { this.config.getChatCompression(), ).imageTokenEstimate; // When lastPromptTokenCount > 0, estimatePromptTokens uses the - // API-authoritative count + a tiny estimate of just the new user - // message — it does NOT touch the history at all in that branch, so - // skip the costly `getHistory(true)` clone on the steady-state path. + // API-authoritative previous prompt count + the previous response's + // output token count + a tiny estimate of just the new user message. + // It does NOT touch the history at all in that branch, so skip the + // costly `getHistory(true)` clone on the steady-state path. // The lastPromptTokenCount=0 branch (first send after --continue // restore / subagent inheritance) walks history with a char/4 // heuristic that can under-count by ~15-20K tokens; the reactive @@ -1747,6 +1780,7 @@ export class GeminiChat { this.lastPromptTokenCount > 0 ? [] : this.getHistoryShallow(true), userContent, this.lastPromptTokenCount, + this.lastOutputTokenCount, imageTokenEstimate, ); const isHardTier = effectiveTokens >= hard; @@ -1798,6 +1832,7 @@ export class GeminiChat { this.lastPromptTokenCount > 0 ? [] : this.getHistoryShallow(true), userContent, this.lastPromptTokenCount, + this.lastOutputTokenCount, imageTokenEstimate, ) : 0; @@ -2877,6 +2912,7 @@ export class GeminiChat { totalTokenCount: number; candidatesTokenCount: number; cachedContentTokenCount: number; + thoughtsTokenCount: number; } | undefined; @@ -2924,6 +2960,14 @@ export class GeminiChat { // Coerce hostile-provider values (NaN / Infinity / negative) to 0 // so the compaction gate arithmetic stays well-defined; see // `coerceUsageCount` for the failure modes this guards against. + const hasUsablePromptTokenCount = + typeof usageMetadata.promptTokenCount === 'number' && + Number.isFinite(usageMetadata.promptTokenCount) && + usageMetadata.promptTokenCount >= 0; + const hasUsableTotalTokenCount = + typeof usageMetadata.totalTokenCount === 'number' && + Number.isFinite(usageMetadata.totalTokenCount) && + usageMetadata.totalTokenCount >= 0; const promptTokenCount = coerceUsageCount( usageMetadata.promptTokenCount, 'promptTokenCount', @@ -2940,6 +2984,10 @@ export class GeminiChat { usageMetadata.cachedContentTokenCount, 'cachedContentTokenCount', ); + const thoughtsTokenCount = coerceUsageCount( + usageMetadata.thoughtsTokenCount, + 'thoughtsTokenCount', + ); // Stash coerced values so recordAssistantTurn can reuse them // without re-calling coerceUsageCount inline. coercedUsage = { @@ -2947,12 +2995,23 @@ export class GeminiChat { totalTokenCount, candidatesTokenCount, cachedContentTokenCount, + thoughtsTokenCount, }; - const lastPromptTokenCount = promptTokenCount || totalTokenCount; + const lastPromptTokenCount = hasUsablePromptTokenCount + ? promptTokenCount + : totalTokenCount; if (lastPromptTokenCount) { // Always update the per-chat counter so this chat (including // subagents) can make its own compaction decisions. this.lastPromptTokenCount = lastPromptTokenCount; + this.lastOutputTokenCount = hasUsablePromptTokenCount + ? getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount, + ...(hasUsableTotalTokenCount ? { totalTokenCount } : {}), + candidatesTokenCount, + thoughtsTokenCount, + }) + : 0; // Mirror to the global telemetry only when wired — subagents // pass `telemetryService=undefined` to keep their context usage // out of the main session's UI counters. diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index 97ede7deec1..f955abc0885 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -356,6 +356,10 @@ export class ChatCompressionService { chat.getHistoryShallow(true), pendingUserMessage, originalTokenCount, + // lastOutputTokenCount is unavailable here. The common + // sendMessageStream path passes precomputedEffectiveTokens, + // which already includes the chat's previous output tokens. + 0, slimmingConfig.imageTokenEstimate, ) : originalTokenCount; diff --git a/packages/core/src/services/sessionService.test.ts b/packages/core/src/services/sessionService.test.ts index 971a43fe376..d670b9fbae6 100644 --- a/packages/core/src/services/sessionService.test.ts +++ b/packages/core/src/services/sessionService.test.ts @@ -22,6 +22,7 @@ import { SessionService, buildApiHistoryFromConversation, getResumePromptTokenCount, + getResumeTokenCounts, type ConversationRecord, } from './sessionService.js'; import { CompressionStatus } from '../core/turn.js'; @@ -1295,6 +1296,9 @@ describe('SessionService', () => { makeConversation([compressionRecord, assistant]), ), ).toBe(450); + expect( + getResumeTokenCounts(makeConversation([compressionRecord, assistant])), + ).toEqual({ promptTokenCount: 450, outputTokenCount: 0 }); }); it('should prefer promptTokenCount over totalTokenCount when both are present', () => { @@ -1310,6 +1314,26 @@ describe('SessionService', () => { makeConversation([compressionRecord, assistant]), ), ).toBe(200); + expect( + getResumeTokenCounts(makeConversation([compressionRecord, assistant])), + ).toEqual({ promptTokenCount: 200, outputTokenCount: 250 }); + }); + + it('should restore disjoint candidate and thought output tokens when total is unavailable', () => { + const assistant: ChatRecord = { + ...baseRecord, + uuid: 'a1', + parentUuid: 'comp', + type: 'assistant', + usageMetadata: { + promptTokenCount: 200, + candidatesTokenCount: 40, + thoughtsTokenCount: 60, + }, + }; + expect( + getResumeTokenCounts(makeConversation([compressionRecord, assistant])), + ).toEqual({ promptTokenCount: 200, outputTokenCount: 100 }); }); it('should fall back to compression when latest assistant has zero usage', () => { @@ -1325,6 +1349,9 @@ describe('SessionService', () => { makeConversation([compressionRecord, assistant]), ), ).toBe(300); + expect( + getResumeTokenCounts(makeConversation([compressionRecord, assistant])), + ).toEqual({ promptTokenCount: 300, outputTokenCount: 0 }); }); }); diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 998510308a0..93a8a966ded 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -33,6 +33,7 @@ import { readLastJsonStringFieldSync, readLastJsonStringFieldsSync, } from '../utils/sessionStorageUtils.js'; +import { getUsageOutputTokenCountForPromptEstimate } from './tokenEstimation.js'; const debugLogger = createDebugLogger('SESSION'); @@ -1462,7 +1463,7 @@ export function buildApiHistoryFromConversation( export function replayUiTelemetryFromConversation( conversation: ConversationRecord, sessionId?: string, -): void { +): ResumeTokenCounts | undefined { if (sessionId) { uiTelemetryService.resetSession(sessionId); } else { @@ -1482,10 +1483,18 @@ export function replayUiTelemetryFromConversation( } } - const resumePromptTokens = getResumePromptTokenCount(conversation); - if (resumePromptTokens !== undefined) { - uiTelemetryService.setLastPromptTokenCount(resumePromptTokens); + const resumeTokenCounts = getResumeTokenCounts(conversation); + if (resumeTokenCounts !== undefined) { + uiTelemetryService.setLastPromptTokenCount( + resumeTokenCounts.promptTokenCount, + ); } + return resumeTokenCounts; +} + +export interface ResumeTokenCounts { + promptTokenCount: number; + outputTokenCount: number; } /** @@ -1497,6 +1506,18 @@ export function replayUiTelemetryFromConversation( export function getResumePromptTokenCount( conversation: ConversationRecord, ): number | undefined { + return getResumeTokenCounts(conversation)?.promptTokenCount; +} + +/** + * Returns the prompt and previous-response output token counts used to seed a + * resumed chat. The prompt count restores the context anchor; the output + * count preserves the output tokens appended after that prompt count was + * reported, matching steady-state prompt estimation on the next send. + */ +export function getResumeTokenCounts( + conversation: ConversationRecord, +): ResumeTokenCounts | undefined { for (let i = conversation.messages.length - 1; i >= 0; i--) { const record = conversation.messages[i]; @@ -1504,7 +1525,10 @@ export function getResumePromptTokenCount( const usage = record.usageMetadata; const candidate = usage?.promptTokenCount ?? usage?.totalTokenCount; if (candidate) { - return candidate; + return { + promptTokenCount: candidate, + outputTokenCount: getUsageOutputTokenCountForPromptEstimate(usage), + }; } } @@ -1513,7 +1537,10 @@ export function getResumePromptTokenCount( | ChatCompressionRecordPayload | undefined; if (payload?.info) { - return payload.info.newTokenCount; + return { + promptTokenCount: payload.info.newTokenCount, + outputTokenCount: 0, + }; } } } diff --git a/packages/core/src/services/tokenEstimation.test.ts b/packages/core/src/services/tokenEstimation.test.ts index b853ffc3d10..3d1b45d0405 100644 --- a/packages/core/src/services/tokenEstimation.test.ts +++ b/packages/core/src/services/tokenEstimation.test.ts @@ -9,6 +9,7 @@ import type { Content } from '@google/genai'; import { estimateContentTokens, estimatePromptTokens, + getUsageOutputTokenCountForPromptEstimate, } from './tokenEstimation.js'; const textContent = (text: string): Content => ({ @@ -84,8 +85,79 @@ describe('estimatePromptTokens', () => { expect(estimatePromptTokens(history, user, 5000)).toBe(5000 + userEst); }); + it('includes the previous turn candidate tokens in the steady-state estimate', () => { + const userEst = estimateContentTokens([user]); + expect(estimatePromptTokens(history, user, 5000, 1200)).toBe( + 5000 + 1200 + userEst, + ); + }); + + it('keeps custom image-token estimates as the fifth argument', () => { + const imageUser: Content = { + role: 'user', + parts: [{ inlineData: { mimeType: 'image/png', data: 'xxx' } }], + }; + + expect(estimatePromptTokens(history, imageUser, 5000, 1200, 1600)).toBe( + 5000 + 1200 + 1600, + ); + }); + it('falls back to full estimate when lastPromptTokenCount is 0', () => { const fullEst = estimateContentTokens([...history, user]); expect(estimatePromptTokens(history, user, 0)).toBe(fullEst); }); }); + +describe('getUsageOutputTokenCountForPromptEstimate', () => { + it('uses totalTokenCount when available to avoid candidate/thought overlap ambiguity', () => { + expect( + getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount: 100, + totalTokenCount: 180, + candidatesTokenCount: 70, + thoughtsTokenCount: 50, + }), + ).toBe(80); + }); + + it('does not double-count thoughts that appear included in candidates', () => { + expect( + getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount: 100, + candidatesTokenCount: 150, + thoughtsTokenCount: 120, + }), + ).toBe(150); + }); + + it('adds thoughts when they exceed candidates and are likely disjoint', () => { + expect( + getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount: 100, + candidatesTokenCount: 50, + thoughtsTokenCount: 120, + }), + ).toBe(170); + }); + + it('adds equal candidate and thought counts because equality does not prove overlap', () => { + expect( + getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount: 100, + candidatesTokenCount: 80, + thoughtsTokenCount: 80, + }), + ).toBe(160); + }); + + it('clamps negative disjoint output token counts to zero', () => { + expect( + getUsageOutputTokenCountForPromptEstimate({ + promptTokenCount: 100, + candidatesTokenCount: -10, + thoughtsTokenCount: -5, + }), + ).toBe(0); + }); +}); diff --git a/packages/core/src/services/tokenEstimation.ts b/packages/core/src/services/tokenEstimation.ts index 2946bc0fab8..1bdb11503d6 100644 --- a/packages/core/src/services/tokenEstimation.ts +++ b/packages/core/src/services/tokenEstimation.ts @@ -4,7 +4,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Content } from '@google/genai'; +import type { + Content, + GenerateContentResponseUsageMetadata, +} from '@google/genai'; import { DEFAULT_IMAGE_TOKEN_ESTIMATE, TOKEN_TO_CHAR_RATIO, @@ -49,8 +52,10 @@ export function estimateContentTokens( * Compute an effective prompt-token count for the auto-compaction gate. * * `lastPromptTokenCount` (from the previous turn's usage metadata) lacks - * two things: the current user message, and any initial value on the - * very first send. This helper closes both gaps via local estimation. + * three things: the current user message, the previous model response that + * was appended to local history after that prompt count was reported, and + * any initial value on the very first send. This helper closes those gaps via + * local estimation plus `lastOutputTokenCount` when available. * * WARNING: like estimateContentTokens, this is a conservative lower * bound. Use it to TRIGGER earlier, never to SKIP — the fallback path @@ -61,11 +66,13 @@ export function estimatePromptTokens( history: Content[], userMessage: Content, lastPromptTokenCount: number, + lastOutputTokenCount: number = 0, imageTokenEstimate: number = DEFAULT_IMAGE_TOKEN_ESTIMATE, ): number { if (lastPromptTokenCount > 0) { return ( lastPromptTokenCount + + lastOutputTokenCount + estimateContentTokens([userMessage], imageTokenEstimate) ); } @@ -76,3 +83,22 @@ export function estimatePromptTokens( // misses for that reason. See review #4168 R3.3. return estimateContentTokens([...history, userMessage], imageTokenEstimate); } + +export function getUsageOutputTokenCountForPromptEstimate( + usage: GenerateContentResponseUsageMetadata | undefined, +): number { + if (usage?.promptTokenCount === undefined) { + return 0; + } + if (usage.totalTokenCount !== undefined) { + return Math.max(0, usage.totalTokenCount - usage.promptTokenCount); + } + const candidates = Math.max(0, usage.candidatesTokenCount ?? 0); + const thoughts = Math.max(0, usage.thoughtsTokenCount ?? 0); + // Some OpenAI-compatible providers include reasoning tokens inside + // candidatesTokenCount when totalTokenCount is unavailable. If candidates + // strictly dominates thoughts, treat thoughts as potentially overlapping; + // otherwise add the larger reasoning-only count so long-thinking responses + // still advance the steady-state prompt estimate. + return candidates > thoughts ? candidates : candidates + thoughts; +}