diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 6ab69cdb6c..4e3068d5fb 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -115,6 +115,7 @@ vi.mock('../telemetry/loggers.js', () => ({ vi.mock('../telemetry/uiTelemetry.js', () => ({ uiTelemetryService: { setLastPromptTokenCount: vi.fn(), + setLastCachedContentTokenCount: vi.fn(), }, })); @@ -15604,6 +15605,116 @@ describe('GeminiChat', async () => { expect(compressSpy.mock.calls[0]?.[1].originalTokenCount).toBe(0); }); + it('invalidates a stale route count before fast-compression sizing', () => { + // compressFast is the third entrypoint alongside sendMessageStream + // and tryCompress: it reads the raw count field for its apiBaseline, + // so it must drop a pre-switch count before sizing the new route. + vi.mocked(mockConfig.getClearContextOnIdle).mockReturnValue({ + toolResultsThresholdMinutes: 30, + toolResultsNumToKeep: 1, + }); + const fastChat = new GeminiChat( + mockConfig, + config, + [ + { role: 'user', parts: [{ text: 'question' }] }, + { + role: 'model', + parts: [ + { text: 'reasoning '.repeat(100), thought: true }, + { text: 'answer' }, + ], + }, + ], + { + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + fastChat.setLastPromptTokenCount(691_000, false); + switchRoute('anthropic-model@beef1234'); + + const result = fastChat.compressFast(); + + expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); + // The stale 691_000 must not anchor sizing for the new route: the + // baseline falls back to the history-walk estimate. + expect(result.info.originalTokenCount).toBeLessThan(691_000); + expect(fastChat.getLastPromptTokenCount()).toBeLessThan(691_000); + }); + + it('zeroes the telemetry cached-content mirror when invalidating a foreign count', () => { + // The cached content count is written for the same foreign route's + // last response; leaving it up next to the zeroed prompt count gives + // /context an internally inconsistent capacity picture. + chat.setLastPromptTokenCount(691_000, false); + switchRoute('anthropic-model@beef1234'); + + expect(chat.getLastPromptTokenCount()).toBe(0); + expect( + uiTelemetryService.setLastCachedContentTokenCount, + ).toHaveBeenCalledWith(0); + }); + + it('restores the request route key when a failed hard-rescue rolls counts back', async () => { + // Hard-rescue only fires for non-exact sends, whose request route key + // can differ from the active route's. tryCompress re-stamps the key + // to the ACTIVE route mid-rescue; the rollback must restore the key + // alongside the counts, or the resurrected override-route count rides + // the active key past the next send's entry invalidation (#9454). + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + const rescueChat = new GeminiChat( + mockConfig, + config, + [ + { role: 'user', parts: [{ text: 'earlier turn' }] }, + { role: 'model', parts: [{ text: 'ack' }] }, + ], + { + recordAssistantTurn: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + // Authoritative count recorded by an earlier override-route turn. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'override-model@route', + ); + rescueChat.setLastPromptTokenCount(176_999, false); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + + 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, + }, + }); + + await expect( + rescueChat.sendMessageStream( + 'override-model', + { message: 'continue' }, + 'prompt-hard-rescue-route-key-restore', + ), + ).rejects.toThrow(/compression status: COMPRESSED/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + // The restored count belongs to the override route: an active-route + // read must invalidate it, not inherit it. + expect(rescueChat.getLastPromptTokenCount()).toBe(0); + }); + it('does not anchor an exact route output clamp on the active route count', async () => { // Authoritative count recorded by the ACTIVE route. An exact `\0` // route send targets a different serialization, so its output clamp @@ -15645,8 +15756,8 @@ describe('GeminiChat', async () => { vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({ resolveForModel, } as unknown as ReturnType); - vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation( - (model) => (model ? `${model}@route` : 'gemini-pro@test0001'), + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'gemini-pro@test0001', ); const selector = 'openai:vision-agent\0https://vision.example.com/v1\0'; diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 33b496a37c..cf6583189d 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -1988,13 +1988,21 @@ export class GeminiChat { * targets so a foreign count cannot anchor that request's decisions even * when the active route owns it — e.g. an exact `\0` route selector, or * a non-exact send whose `model` param overrides the active model. + * + * The telemetry mirror is display-only state: it is resynchronized here, + * at the next guarded read, so between a `/model` switch and the next + * chat touch the UI counters may briefly show the previous route's + * counts. Decision paths never read the mirror, only the route-aware + * chat getters above. */ - private invalidateTokenCountsIfRouteChanged( - targetRouteKey: string = this.currentRouteKey(), - ): void { + private invalidateTokenCountsIfRouteChanged(targetRouteKey?: string): void { if (this.lastPromptTokenCount === 0 && this.lastOutputTokenCount === 0) { return; } + // Resolve the active-route default only AFTER the zero-count fast path: + // computing a route identity (SHA-256 digest + config lookups) on every + // count read while both counts are 0 would defeat the guard above. + targetRouteKey ??= this.currentRouteKey(); if (this.tokenCountsRouteKey === targetRouteKey) { return; } @@ -2004,8 +2012,10 @@ export class GeminiChat { ); this.setLastPromptTokenCount(0); // Keep the telemetry mirror in sync, or the UI context counters - // and compression banners keep reading the foreign count. + // and compression banners keep reading the foreign count. The cached + // content count belongs to the same foreign route's last response. this.telemetryService?.setLastPromptTokenCount(0); + this.telemetryService?.setLastCachedContentTokenCount(0); } /** @@ -2393,15 +2403,15 @@ export class GeminiChat { if (exactRoute) { model = exactRoute.model; } - const requestRouteKey = exactRoute - ? this.config.getModelRouteIdentity( - model, - exactRoute.contentGeneratorConfig, - ) - : this.config.getModelRouteIdentity( - model, - this.config.getContentGeneratorConfig(), - ); + // Both arms are one call: for a non-exact send `exactRoute` is + // undefined, and `resolvedModelIdentity`'s second parameter defaults to + // `getContentGeneratorConfig()` — including when passed an explicit + // undefined. Keeping a single call site means a future change to how + // the request route is identified cannot drift between the arms. + const requestRouteKey = this.config.getModelRouteIdentity( + model, + exactRoute?.contentGeneratorConfig, + ); // Counts recorded for a route other than this request's target must not // anchor its admission/clamp/compression decisions (#9454). Comparing // against the REQUEST route — resolved above — keeps an exact `\0` @@ -2559,6 +2569,13 @@ export class GeminiChat { const lastPromptTokenCountBeforeHardRescue = this.lastPromptTokenCount; const lastPromptTokenCountWasEstimatedBeforeHardRescue = this.lastPromptTokenCountIsEstimated; + // tryCompress re-stamps tokenCountsRouteKey to the ACTIVE route (via + // setLastPromptTokenCount on the success path) even though this send + // targets the REQUEST route — and hard-rescue only fires for + // non-exact sends, whose request key can differ from the active one. + // Capture the key so the rollback below restores the resurrected + // count's original route attribution along with the count itself. + const tokenCountsRouteKeyBeforeHardRescue = this.tokenCountsRouteKey; const hardRescueFailureCountBeforeHardRescue = this.hardRescueFailureCount; if (shouldForceFromHard) { @@ -2635,6 +2652,7 @@ export class GeminiChat { this.lastPromptTokenCount = lastPromptTokenCountBeforeHardRescue; this.lastPromptTokenCountIsEstimated = lastPromptTokenCountWasEstimatedBeforeHardRescue; + this.tokenCountsRouteKey = tokenCountsRouteKeyBeforeHardRescue; this.telemetryService?.setLastPromptTokenCount( lastPromptTokenCountBeforeHardRescue, );