From 3fd059368b3654d120faac792d734cd240fd0d69 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=93=E8=89=AF?= <1204183885@qq.com> Date: Mon, 24 Aug 2026 05:00:16 +0000 Subject: [PATCH] fix(core): invalidate token counts recorded for a switched model route (#9506) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(core): invalidate token counts recorded for a switched model route `/model` switches rebuild the content generator but keep the same GeminiChat instance, so API-reported prompt/output token counts from the previous route survived and anchored admission, output clamping, and compression decisions for a different serialization (#9454). Attribute the counts to the route that produced them (Config.getModelRouteIdentity) and invalidate them on a route change so all safety decisions fall back to the history-walk estimate. Closes #9454 * fix(core): keep token counts scoped to producing route * fix(core): close round-2 route-scoping findings (#9454) - C1: client.test.ts session-token-limit gate test now seeds the chat mock's getLastPromptTokenCount (the gate's new source) instead of relying on the telemetry stub alone. - C2: goal-turn-integration.test.ts passes the routeKey positional added to processStreamResponse and widens the local cast. - S1: fix stale rationale comment — no session token-limit gate reads the telemetry mirror anymore; UI context counters and compression banners do. - S2: pin tryCompress's entry invalidation with a focused test; manual /compress reaches it without sendMessageStream's entry reset. - S3: direct config tests for getModelRouteIdentity — call stability, model@ shape, and the guard keeping the registry baseUrl out of non-active model identities. - R2-3: sendMessageStream now resolves the request route first and invalidates counts against IT (deriving the key from the actual model param on the non-exact branch), so an active-route count can no longer anchor an exact `\0` route's output clamp; regression test added. * fix(core): close round-3 route-scoping findings (#9454) - R3-1: hard-rescue rollback now restores tokenCountsRouteKey alongside the counts; tryCompress re-stamps the key to the active route mid-rescue, which left the resurrected request-route count riding the active key past the next entry invalidation (regression test: active-route read after failed rescue must not inherit the override count). - R3-2(2): invalidation zeroes the telemetry cached-content mirror together with the prompt mirror so /context stops rendering a foreign cached count beside a zeroed prompt count; the mirror is documented as best-effort display state between a switch and the next guarded read (R3-2(1)). - R3-4: collapse the requestRouteKey ternary into one getModelRouteIdentity call (the non-exact arm passed exactly the default parameter value). - R3-5: resolve the active-route default lazily after the zero-count fast path instead of eagerly in the default parameter. - R1-3: add the missing compressFast route-invalidation test (third entrypoint; mutation-verified). * fix(core): preserve route-scoped token guards (#9454) Co-authored-by: Qwen-Coder * fix(core): resolve full-turn route selectors at the session token limit gate (#9454) * fix(core): stamp fallback-served token counts under the request route (#9454) * fix(core): retain route-scoped token counts across route crossings (#9454) * fix(core): keep route-scoped token counts consistent through compression (#9454) * fix(core): keep rescued output counts and compression stamps route-scoped (#9454) --------- Co-authored-by: yiliang114 Co-authored-by: Qwen-Coder --- .../src/config/config-session-env.test.ts | 103 ++ packages/core/src/config/config.ts | 27 +- packages/core/src/core/client.test.ts | 169 +++- packages/core/src/core/client.ts | 20 +- packages/core/src/core/geminiChat.test.ts | 906 ++++++++++++++++++ packages/core/src/core/geminiChat.ts | 337 ++++++- .../src/core/goal-turn-integration.test.ts | 3 + 7 files changed, 1549 insertions(+), 16 deletions(-) diff --git a/packages/core/src/config/config-session-env.test.ts b/packages/core/src/config/config-session-env.test.ts index 134aa841a2..3ae5226296 100644 --- a/packages/core/src/config/config-session-env.test.ts +++ b/packages/core/src/config/config-session-env.test.ts @@ -436,3 +436,106 @@ describe('Config provider-qualified model identity', () => { expect(process.env['QWEN_CODE_MODEL_IDENTITY']).toBe(claimed); }); }); + +describe('Config.getModelRouteIdentity (#9454 route key)', () => { + it('returns an identical identity across repeated calls for one configuration', async () => { + const { Config } = await import('./config.js'); + const { resolveContentGeneratorConfigWithSources, AuthType } = await import( + '../core/contentGenerator.js' + ); + vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({ + config: { + model: 'steady-model', + apiKey: 'k', + baseUrl: 'https://provider-a.example/v1', + } as ContentGeneratorConfig, + sources: {}, + }); + const config = new Config({ ...baseParams }); + await config.refreshAuth(AuthType.USE_GEMINI); + + // Route-scoped caches (e.g. GeminiChat token counts) compare these + // strings for equality — the value must not drift between calls. + const first = config.getModelRouteIdentity(); + expect(config.getModelRouteIdentity()).toBe(first); + expect(config.getModelRouteIdentity()).toBe(first); + expect(first).toMatch(/^steady-model@[0-9a-f]{8}$/); + }); + + it('keeps the `model@` digest shape for explicit route queries', async () => { + const { Config } = await import('./config.js'); + const { resolveContentGeneratorConfigWithSources, AuthType } = await import( + '../core/contentGenerator.js' + ); + vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({ + config: { + model: 'active-model', + apiKey: 'k', + baseUrl: 'https://provider-a.example/v1', + } as ContentGeneratorConfig, + sources: {}, + }); + const config = new Config({ ...baseParams }); + await config.refreshAuth(AuthType.USE_GEMINI); + + // The readable model id stays the prefix; the discriminator is exactly + // eight hex characters, stable for one (auth type, endpoint) pair. + const explicit = config.getModelRouteIdentity('route-model', { + model: 'route-model', + authType: 'openai', + baseUrl: 'https://route.example/v1', + } as ContentGeneratorConfig); + expect(explicit).toMatch(/^route-model@[0-9a-f]{8}$/); + expect(config.getModelRouteIdentity('route-model', { + model: 'route-model', + authType: 'openai', + baseUrl: 'https://route.example/v1', + } as ContentGeneratorConfig)).toBe(explicit); + }); + + it('does not mix the registry base URL into a non-active model identity', async () => { + // The registry-baseUrl fallback qualifies the ACTIVE model's route when + // its own generator config carries no baseUrl. A foreign model queried + // with its own configuration must not pick that fallback up — hashing + // the registry endpoint into another model's identity would invalidate + // its route-scoped state on unrelated registry changes. + const { Config } = await import('./config.js'); + const { resolveContentGeneratorConfigWithSources, AuthType } = await import( + '../core/contentGenerator.js' + ); + vi.mocked(resolveContentGeneratorConfigWithSources).mockReturnValue({ + config: { + model: 'active-model', + apiKey: 'k', + // No baseUrl of its own: the active model falls back to the + // registry base URL below. + } as ContentGeneratorConfig, + sources: {}, + }); + const config = new Config({ ...baseParams }); + await config.refreshAuth(AuthType.USE_GEMINI); + const registrySpy = vi.spyOn(config, 'getCurrentModelRegistryBaseUrl'); + + const foreignGeneratorConfig = { + model: 'foreign-model', + authType: 'openai', + } as ContentGeneratorConfig; + + registrySpy.mockReturnValue('https://registry.example/v1'); + const activeWithRegistry = config.getModelRouteIdentity(); + const foreignWithRegistry = + config.getModelRouteIdentity('foreign-model', foreignGeneratorConfig); + + registrySpy.mockReturnValue(null); + const activeWithoutRegistry = config.getModelRouteIdentity(); + const foreignWithoutRegistry = + config.getModelRouteIdentity('foreign-model', foreignGeneratorConfig); + + // The fallback is load-bearing for the ACTIVE model… + expect(activeWithRegistry).toMatch(/^active-model@[0-9a-f]{8}$/); + expect(activeWithRegistry).not.toBe(activeWithoutRegistry); + // …but the foreign model's identity ignores the registry entirely. + expect(foreignWithRegistry).toMatch(/^foreign-model@[0-9a-f]{8}$/); + expect(foreignWithRegistry).toBe(foreignWithoutRegistry); + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index fd84704c4b..75ce85f662 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -4392,12 +4392,16 @@ export class Config { * configurations without publishing where they point. The bare id stays the * readable half, so a mismatch still names the model a human recognises. */ - private resolvedModelIdentity(): string { - const model = this.getModel(); - const authType = this.getContentGeneratorConfig()?.authType ?? ''; + private resolvedModelIdentity( + model = this.getModel(), + generatorConfig = this.getContentGeneratorConfig(), + ): string { + const authType = generatorConfig?.authType ?? ''; const baseUrl = - this.getContentGeneratorConfig()?.baseUrl ?? - this.getCurrentModelRegistryBaseUrl() ?? + generatorConfig?.baseUrl ?? + (model === this.getModel() + ? this.getCurrentModelRegistryBaseUrl() + : undefined) ?? ''; if (authType === '' && baseUrl === '') return model; const digest = createHash('sha256') @@ -4407,6 +4411,19 @@ export class Config { return `${model}@${digest}`; } + /** + * Identity of the currently active model route for consumers that cache + * route-specific state and must invalidate it when a model/auth/endpoint + * switch swaps the content generator — e.g. GeminiChat's API-reported + * token counts (#9454). Same identity ⇒ same serialization target. + */ + getModelRouteIdentity( + model?: string, + generatorConfig?: ContentGeneratorConfig, + ): string { + return this.resolvedModelIdentity(model, generatorConfig); + } + /** * Returns the configured fast model selector when it resolves to an available * model. Bare selectors stay bare and authType-qualified selectors keep their diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index c35f8d7a96..799063d774 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -619,7 +619,7 @@ describe('Gemini Client (client.ts)', () => { toolResultsThresholdMinutes: 60, toolResultsNumToKeep: 5, }), - getSessionTokenLimit: vi.fn().mockReturnValue(32000), + getSessionTokenLimit: vi.fn().mockReturnValue(0), getNoBrowser: vi.fn().mockReturnValue(false), getUsageStatisticsEnabled: vi.fn().mockReturnValue(true), getTelemetryIncludeSensitiveSpanAttributes: vi @@ -654,6 +654,7 @@ describe('Gemini Client (client.ts)', () => { .mockReturnValue('/test/project/root/.gemini/projects/test-project'), }, getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator), + getModelRouteIdentity: vi.fn().mockReturnValue('test-route'), getEffectiveInputModalities: vi.fn().mockReturnValue({}), getBaseLlmClient: vi.fn(), getSkipLoopDetection: vi.fn().mockReturnValue(false), @@ -3147,6 +3148,7 @@ describe('Gemini Client (client.ts)', () => { addHistory, getHistory: vi.fn().mockReturnValue([]), getHistoryLength: vi.fn().mockReturnValue(1), + getLastPromptTokenCount: vi.fn().mockReturnValue(101), // Send is skipped, so the push counter never advances → restore. getUserContentPushCount: vi.fn().mockReturnValue(0), stripOrphanedUserEntriesFromHistory: vi @@ -3172,6 +3174,170 @@ describe('Gemini Client (client.ts)', () => { expect(mockTurnRunFn).not.toHaveBeenCalled(); expect(addHistory).toHaveBeenCalledWith(retryEntry); }); + + it('invalidates a foreign route count before the session limit gate', async () => { + let route = 'route-a'; + let telemetryCount = 691_000; + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation( + () => route, + ); + vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100_000); + vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockImplementation( + () => telemetryCount, + ); + vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockImplementation( + (count) => { + telemetryCount = count; + }, + ); + client.getChat().setLastPromptTokenCount(telemetryCount); + route = 'route-b'; + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'response' }; + })(), + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'new route' }], + new AbortController().signal, + 'prompt-route-switch', + ), + ); + + expect(events).not.toContainEqual( + expect.objectContaining({ + type: GeminiEventType.SessionTokenLimitExceeded, + }), + ); + expect(telemetryCount).toBe(0); + }); + + it('applies the session limit to the requested override route', async () => { + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'override-model@route', + ); + vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100); + client.getChat().setLastPromptTokenCount(101); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active-model@route', + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'override route' }], + new AbortController().signal, + 'prompt-override-limit', + { + type: SendMessageType.UserQuery, + modelOverride: 'override-model', + }, + ), + ); + + expect(events).toContainEqual({ + type: GeminiEventType.SessionTokenLimitExceeded, + value: expect.objectContaining({ currentTokens: 101, limit: 100 }), + }); + expect(mockTurnRunFn).not.toHaveBeenCalled(); + }); + + it('applies the session limit to a resolved full-turn route selector', async () => { + // The vision-bridge full-turn selector `${id}\0${baseUrl}\0` arrives as + // modelOverride. GeminiChat.sendMessageStream resolves it and stamps + // counts under the RESOLVED route's identity, so the gate must resolve + // the selector before keying — the raw selector key (always containing + // a NUL) can never match a stamped count (#9454). + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'vision-agent@route', + ); + vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100); + client.getChat().setLastPromptTokenCount(101); + const resolveForModel = vi.fn().mockResolvedValue({ + model: 'vision-agent', + contentGeneratorConfig: undefined, + }); + vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({ + resolveForModel, + } as unknown as ReturnType); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active-model@route', + ); + + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'vision route' }], + new AbortController().signal, + 'prompt-selector-limit', + { + type: SendMessageType.UserQuery, + modelOverride: 'openai:vision-agent\0https://vision.example/v1\0', + }, + ), + ); + + expect(resolveForModel).toHaveBeenCalledWith( + 'openai:vision-agent\0https://vision.example/v1', + { failClosed: true }, + ); + expect(events).toContainEqual({ + type: GeminiEventType.SessionTokenLimitExceeded, + value: expect.objectContaining({ currentTokens: 101, limit: 100 }), + }); + expect(mockTurnRunFn).not.toHaveBeenCalled(); + }); + + it('keeps the session limit enforced when turns alternate routes (#9506)', async () => { + // Counts are retained per route (#9506): an intervening turn on + // another route must not destroy the count the gate later reads for + // the original route. Pre-fix, the foreign-route gate read zeroed + // the only slot, so the returning turn read 0 and was admitted + // regardless of size — steady alternation disabled the limit. + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model === 'route-x' ? 'route-x@route' : 'route-a', + ); + vi.mocked(mockConfig.getSessionTokenLimit).mockReturnValue(100); + // Route A's last response stamped an over-limit count. + client.getChat().setLastPromptTokenCount(101); + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'response' }; + })(), + ); + + // Intervening turn on route X: no counts recorded for X yet, so the + // gate admits it. + const foreignEvents = await fromAsync( + client.sendMessageStream( + [{ text: 'foreign route turn' }], + new AbortController().signal, + 'prompt-alternate-foreign', + { type: SendMessageType.UserQuery, modelOverride: 'route-x' }, + ), + ); + expect(foreignEvents).not.toContainEqual( + expect.objectContaining({ + type: GeminiEventType.SessionTokenLimitExceeded, + }), + ); + + // Returning to route A must still trip the gate with the retained + // over-limit count — the alternation must not have zeroed it. + const events = await fromAsync( + client.sendMessageStream( + [{ text: 'back on route a' }], + new AbortController().signal, + 'prompt-alternate-return', + { type: SendMessageType.UserQuery }, + ), + ); + expect(events).toContainEqual({ + type: GeminiEventType.SessionTokenLimitExceeded, + value: expect.objectContaining({ currentTokens: 101, limit: 100 }), + }); + expect(mockTurnRunFn).toHaveBeenCalledTimes(1); + }); }); /** @@ -9266,6 +9432,7 @@ Other open files: const mockChat: Partial = { addHistory: vi.fn(), getHistory: vi.fn().mockReturnValue([]), + getLastPromptTokenCount: vi.fn().mockReturnValue(9999), }; client['chat'] = mockChat as GeminiChat; diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index a4f8018515..cd135b9687 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -3142,10 +3142,25 @@ export class GeminiClient { // via the `compressed → ChatCompressed` bridge in turn.ts. Manual /compress // still calls tryCompressChat directly for the full reset (env refresh + // forceFullIdeContext flip). + const model = options?.modelOverride ?? this.config.getModel(); const sessionTokenLimit = this.config.getSessionTokenLimit(); if (sessionTokenLimit > 0) { + // An exact `\0` full-turn route selector resolves to its route before + // GeminiChat.sendMessageStream stamps counts under it, so the gate + // must key the resolved route too — the raw selector key can never + // match a stamped count. Mirrors the resolution at the top of + // GeminiChat.sendMessageStream (#9454). + const exactRoute = model.endsWith('\0') + ? await this.config + .getBaseLlmClient() + .resolveForModel(model.slice(0, -1), { failClosed: true }) + : undefined; + const requestRouteKey = this.config.getModelRouteIdentity( + exactRoute ? exactRoute.model : model, + exactRoute?.contentGeneratorConfig, + ); const lastPromptTokenCount = - uiTelemetryService.getLastPromptTokenCount(); + this.getChat().getLastPromptTokenCount(requestRouteKey); if (lastPromptTokenCount > sessionTokenLimit) { this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); yield { @@ -3237,9 +3252,6 @@ export class GeminiClient { const turn = new Turn(this.getChat(), prompt_id, goalPermit); - // Determine the model to use for this turn - const model = options?.modelOverride ?? this.config.getModel(); - // Assemble the outgoing request. IDE context is merged into the // user prompt's first text part, then on UserQuery / Cron turns // the system reminders block is prepended in front of everything diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 1152224fab..56a4ca6f4b 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(), }, })); @@ -180,6 +181,7 @@ describe('GeminiChat', async () => { model: 'test-model', }), getModel: vi.fn().mockReturnValue('gemini-pro'), + getModelRouteIdentity: vi.fn().mockReturnValue('gemini-pro@test0001'), setModel: vi.fn(), getProjectRoot: vi.fn().mockReturnValue('/test/project/root'), getTargetDir: vi.fn().mockReturnValue('/test/project/root'), @@ -6986,6 +6988,7 @@ describe('GeminiChat', async () => { finishReason: 'STOP', }, ], + usageMetadata: { promptTokenCount: 99_999 }, } as unknown as GenerateContentResponse; })(), ); @@ -7010,6 +7013,9 @@ 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.getEffectiveInputModalities).mockReturnValue({ pdf: true, }); @@ -7062,6 +7068,7 @@ describe('GeminiChat', async () => { for await (const _ of stream) { /* consume */ } + expect(chat.getLastPromptTokenCount()).toBe(0); expect(resolveForModel).toHaveBeenCalledOnce(); expect(resolveForModel).toHaveBeenCalledWith(routeSelector, { @@ -7223,6 +7230,9 @@ 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', + ); const capacityError = Object.assign( new StreamContentError( @@ -7261,6 +7271,7 @@ describe('GeminiChat', async () => { finishReason: 'STOP', }, ], + usageMetadata: { promptTokenCount: 99_999 }, } as unknown as GenerateContentResponse; })(), ); @@ -7328,6 +7339,7 @@ describe('GeminiChat', async () => { ).toContain('fallback-image'); expect(fallbackAGenerateContentStream).toHaveBeenCalledTimes(1); expect(fallbackBGenerateContentStream).toHaveBeenCalledTimes(1); + expect(chat.getLastPromptTokenCount()).toBe(0); expect( events.some( (event) => @@ -7338,6 +7350,88 @@ describe('GeminiChat', async () => { ).toBe(true); }); + it('stamps fallback-served counts under the request route key (#9454)', async () => { + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + authType: AuthType.USE_GEMINI, + model: 'test-model', + maxRetries: 0, + }); + vi.mocked(mockConfig.getModelFallbacks).mockReturnValue(['fallback-b']); + + const fallbackBGenerateContentStream = vi.fn(); + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: { + generateContent: vi.fn(), + generateContentStream: fallbackBGenerateContentStream, + countTokens: vi.fn(), + embedContent: vi.fn(), + batchEmbedContents: vi.fn(), + useSummarizedThinking: vi.fn().mockReturnValue(false), + } as unknown as ContentGenerator, + contentGeneratorConfig: { modalities: {} }, + retryAuthType: AuthType.USE_GEMINI, + retryErrorCodes: undefined, + model: 'fallback-b', + }); + vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({ + resolveForModel, + } as unknown as ReturnType); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'gemini-pro@test0001', + ); + + const capacityError = Object.assign( + new StreamContentError( + '{"error":{"code":"429","message":"Throttling: TPM(1/1)"}}', + ), + { status: 429 }, + ); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce( + (async function* () { + yield { + usageMetadata: { promptTokenCount: 10, totalTokenCount: 10 }, + } as GenerateContentResponse; + throw capacityError; + })(), + ); + fallbackBGenerateContentStream.mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'fallback-b ok' }], + }, + finishReason: 'STOP', + }, + ], + usageMetadata: { promptTokenCount: 99_999 }, + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: [{ text: 'test' }] }, + 'prompt-fallback-route-stamp', + ); + for await (const _ of stream) { + /* consume */ + } + + // The session-token-limit gate in Client reads the count keyed by the + // REQUEST route. A fallback serves on behalf of the same request (the + // session model never changes), so its count must survive that keyed + // read instead of being invalidated as a foreign route's (#9454). + expect(chat.getLastPromptTokenCount('test-model@route')).toBe(99_999); + // The count still belongs to the serving turn's request route: a read + // for a different route invalidates it as before. + expect(chat.getLastPromptTokenCount('other-model@route')).toBe(0); + }); + it('skips a fallback alias that resolves to the current model', async () => { vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ authType: AuthType.USE_GEMINI, @@ -15558,6 +15652,818 @@ describe('GeminiChat', async () => { }); }); + // Route-scoped token counts (#9454): API-reported prompt/output token + // counts describe the serialization of the route (model + auth type + + // endpoint) that produced them. A /model switch rebuilds the content + // generator but keeps this GeminiChat instance, so counts recorded for the + // previous route must be invalidated — otherwise they anchor admission, + // clamp, and compression decisions for a different serialization. + describe('route-scoped token counts (#9454)', () => { + const switchRoute = (routeKey: string) => { + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue(routeKey); + }; + + it('invalidates API-reported counts when the model route changes', () => { + // Count reported by the pre-switch route (authoritative, not estimated). + chat.setLastPromptTokenCount(691_000, false); + expect(chat.getLastPromptTokenCount()).toBe(691_000); + + // Simulate /model switching to a different route; the same chat + // instance survives with its history. + switchRoute('anthropic-model@beef1234'); + + // The stale count must not size requests for the new route: safety + // decisions fall back to the history-walk estimate (count 0). + expect(chat.getLastPromptTokenCount()).toBe(0); + expect(chat.getLastOutputTokenCount()).toBe(0); + expect(chat.isLastPromptTokenCountEstimated()).toBe(false); + // The telemetry mirror must drop the stale count too, or the session + // token-limit gate and compression banners keep using it. + expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith( + 0, + ); + }); + + it('keeps counts authoritative while the route is unchanged', () => { + chat.setLastPromptTokenCount(50_000, false); + + // Repeated reads on the same route keep the API-authoritative count. + expect(chat.getLastPromptTokenCount()).toBe(50_000); + expect(chat.getLastOutputTokenCount()).toBe(0); + expect(chat.isLastPromptTokenCountEstimated()).toBe(false); + expect(chat.getLastPromptTokenCount()).toBe(50_000); + }); + + it('keeps a foreign count intact across a keyless display read (#9506)', () => { + // Counts stamped under one route key (e.g. the vision bridge's + // full-turn selector route) must survive a keyless read: /context + // calls the getters with no argument, which defaults to the ACTIVE + // route key and used to zero the only slot before the + // session-token-limit gate's keyed read got to it. + chat.setLastPromptTokenCount(500_000, false); + switchRoute('other-active@route'); + + // The foreign count must not leak to the active route... + expect(chat.getLastPromptTokenCount()).toBe(0); + expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith( + 0, + ); + // ...and the crossing must not have destroyed it: the gate's keyed + // read for the original route restores the exact API-reported value. + expect(chat.getLastPromptTokenCount('gemini-pro@test0001')).toBe(500_000); + expect(chat.getLastPromptTokenCount()).toBe(0); + }); + + it('restores retained counts when the route switches back (#9506)', () => { + chat.seedResumeTokenCounts(321, 45, true); + switchRoute('anthropic-model@beef1234'); + expect(chat.getLastPromptTokenCount()).toBe(0); + expect(chat.getLastOutputTokenCount()).toBe(0); + + // A turn returning to the original route reads its exact retained + // counts — prompt, previous-output, and provenance — instead of the + // destructive zero a foreign touch used to leave behind. + switchRoute('gemini-pro@test0001'); + expect(chat.getLastPromptTokenCount()).toBe(321); + expect(chat.getLastOutputTokenCount()).toBe(45); + expect(chat.isLastPromptTokenCountEstimated()).toBe(true); + }); + + it('invalidates seeded resume counts after a later route change', () => { + chat.seedResumeTokenCounts(321, 45, false); + expect(chat.getLastPromptTokenCount()).toBe(321); + expect(chat.getLastOutputTokenCount()).toBe(45); + + switchRoute('other-model@1234abcd'); + + expect(chat.getLastPromptTokenCount()).toBe(0); + expect(chat.getLastOutputTokenCount()).toBe(0); + }); + + it('accepts counts recorded on the new route after a switch', () => { + chat.setLastPromptTokenCount(691_000, false); + switchRoute('anthropic-model@beef1234'); + expect(chat.getLastPromptTokenCount()).toBe(0); + + // First response on the new route re-establishes authoritative counts. + chat.setLastPromptTokenCount(120_000, false); + expect(chat.getLastPromptTokenCount()).toBe(120_000); + expect(chat.isLastPromptTokenCountEstimated()).toBe(false); + }); + + it('invalidates a stale count before sending on the new route', async () => { + chat.setLastPromptTokenCount(691_000, false); + switchRoute('anthropic-model@beef1234'); + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async () => { + expect( + uiTelemetryService.setLastPromptTokenCount, + ).toHaveBeenCalledWith(0); + return (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(); + }, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'new route' }, + 'prompt-route-switch', + ); + for await (const _ of stream) { + /* consume */ + } + }); + + it('invalidates a stale route count before manual compression sizing', async () => { + // Authoritative count recorded by the pre-switch route. Manual + // /compress reaches tryCompress without sendMessageStream's entry + // invalidation, and tryCompress reads the count field directly, so it + // must drop the foreign count itself before admission/sizing. + chat.setLastPromptTokenCount(691_000, false); + switchRoute('anthropic-model@beef1234'); + + const compressSpy = vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ); + compressSpy.mockResolvedValue({ + newHistory: null, + info: { + originalTokenCount: 0, + newTokenCount: 0, + compressionStatus: CompressionStatus.NOOP, + }, + }); + + await chat.tryCompress('prompt-manual-compress', true); + + // History is empty, so the estimate path sizes the attempt at 0 — the + // stale 691_000 must not have anchored the compression decision. + expect(compressSpy).toHaveBeenCalledTimes(1); + 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('does not mirror cached content without a route-stamped prompt count', async () => { + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'cached' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + promptTokenCount: 0, + totalTokenCount: 0, + cachedContentTokenCount: 42, + }, + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'cached-only' }, + 'prompt-cached-only', + ); + for await (const _ of stream) { + /* consume */ + } + switchRoute('anthropic-model@beef1234'); + expect(chat.getLastPromptTokenCount()).toBe(0); + + expect( + uiTelemetryService.setLastCachedContentTokenCount, + ).not.toHaveBeenCalledWith(42); + }); + + it('mirrors cached content alongside a route-stamped prompt count', async () => { + // The cached-content mirror's only non-zero production write lives + // inside the prompt-count guard; deleting it must not leave the suite + // green. Consumed without a route switch, a cached-content response + // must reach the /context cached-tokens line (#9454). + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'cached' }] }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + promptTokenCount: 100, + totalTokenCount: 100, + cachedContentTokenCount: 42, + }, + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'cached-happy' }, + 'prompt-cached-happy', + ); + for await (const _ of stream) { + /* consume */ + } + + expect(chat.getLastPromptTokenCount()).toBe(100); + expect( + uiTelemetryService.setLastCachedContentTokenCount, + ).toHaveBeenCalledWith(42); + }); + + 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('restores the retention map when a failed hard-rescue rolls counts back (#9506)', async () => { + // The rescue's own compression consumes retained map entries + // mid-flight (the service's keyless getter reads adopt the active + // route) and a successful compression clears the map outright. The + // rollback must restore the pre-rescue snapshot, or the resurrected + // route's over-limit count survives nowhere and its next + // session-token-limit gate read passes with 0. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'active@route', + ); + const rescueChat = new GeminiChat( + mockConfig, + config, + [{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }], + { + recordAssistantTurn: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + rescueChat.setLastPromptTokenCount(190_000, false); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockImplementationOnce(async (chatToCompress) => { + // Mirror the real service's unconditional keyless reads: they + // adopt the ACTIVE route, consuming its retained entry mid-rescue. + chatToCompress.getLastPromptTokenCount(); + chatToCompress.isLastPromptTokenCountEstimated(); + return { + newHistory: [ + { role: 'user', parts: [{ text: 'still large summary' }] }, + ], + info: { + originalTokenCount: 180_000, + newTokenCount: 178_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }; + }); + + await expect( + rescueChat.sendMessageStream( + 'override-model', + { message: 'continue' }, + 'prompt-rescue-retention-map-restore', + ), + ).rejects.toThrow(/compression status: COMPRESSED/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + // The active route's over-limit count survived the failed rescue: + // its next keyed read must still see it, not 0. + expect(rescueChat.getLastPromptTokenCount('active@route')).toBe(190_000); + expect(rescueChat.getLastPromptTokenCount('override-model@route')).toBe( + 0, + ); + }); + + it('restores the output token count when a failed hard-rescue rolls counts back (#9506)', async () => { + // The rescue's COMPRESSED stamp zeroes lastOutputTokenCount via + // setLastPromptTokenCount. The rollback restores the resurrected + // prompt count, its provenance, the route key and the retention map + // — it must restore the output half of the pair too, or the next + // turn's additive prompt estimate (prompt + output + new content) + // under-counts by the last response's size. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'override-model@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 pair recorded by an earlier override-route turn. + rescueChat.seedResumeTokenCounts(170_000, 8_000, 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-output-restore', + ), + ).rejects.toThrow(/compression status: COMPRESSED/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + // Reading through the override route (the resurrected slot's key) + // must return the full pre-rescue pair, output half included. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'override-model@route', + ); + expect(rescueChat.getLastPromptTokenCount()).toBe(170_000); + expect(rescueChat.getLastOutputTokenCount()).toBe(8_000); + }); + + it('re-adopts the request route after the compression service flips the slots (#9506)', async () => { + // ChatCompressionService.compress reads the KEYLESS count getters, + // which adopt the ACTIVE route. On a non-exact override send whose + // request route differs, that flips the slots back to the active + // route's retained counts mid-rescue. When the summarization side + // query then fails, the post-rescue stop check must size from the + // honest history-walk estimate, not the flipped foreign count. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'active@route', + ); + const rescueChat = new GeminiChat( + mockConfig, + config, + [{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }], + { + recordAssistantTurn: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + rescueChat.setLastPromptTokenCount(150_000, false); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockImplementationOnce(async (chatToCompress) => { + chatToCompress.getLastPromptTokenCount(); + chatToCompress.isLastPromptTokenCountEstimated(); + return { + newHistory: null, + info: { + originalTokenCount: 180_000, + newTokenCount: 0, + compressionStatus: + CompressionStatus.COMPRESSION_FAILED_EMPTY_SUMMARY, + }, + }; + }); + + await expect( + rescueChat.sendMessageStream( + 'override-model', + { message: 'continue' }, + 'prompt-rescue-flip-re-adopt', + ), + ).rejects.toThrow(/Context is too large to send safely/i); + + expect(mockContentGenerator.generateContentStream).not.toHaveBeenCalled(); + // The active route's retained count survived the failed rescue. + expect(rescueChat.getLastPromptTokenCount('active@route')).toBe(150_000); + }); + + it('retains a foreign-keyed slot occupant when the usage stamp re-keys (#9506)', async () => { + // Mid-send compression can leave the slots keyed to the ACTIVE route + // when the response's usage report arrives for the REQUEST route. + // The stamp must retain the displaced occupant before overwriting + // it, or the active route's count is destroyed and its next keyed + // read returns 0 — bypassing the session token limit. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'active@route', + ); + const stampChat = new GeminiChat( + mockConfig, + config, + [{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }], + { + recordAssistantTurn: vi.fn(), + // Successful hard-rescue compression records after the + // post-compression guard passes (deferred recording). + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + stampChat.setLastPromptTokenCount(150_000, false); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockImplementationOnce(async (chatToCompress) => { + chatToCompress.getLastPromptTokenCount(); + chatToCompress.isLastPromptTokenCountEstimated(); + return { + newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }], + info: { + originalTokenCount: 180_000, + newTokenCount: 60_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }; + }); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'ok' }], role: 'model' }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 61_000, + totalTokenCount: 62_000, + }, + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await stampChat.sendMessageStream( + 'override-model', + { message: 'continue' }, + 'prompt-stamp-retains-occupant', + ); + for await (const _ of stream) { + /* consume */ + } + + // The request route's fresh API report occupies the slots... + expect(stampChat.getLastPromptTokenCount('override-model@route')).toBe( + 61_000, + ); + // ...and the active route's post-compression count was retained, + // not destroyed by the re-keying stamp. + expect(stampChat.getLastPromptTokenCount('active@route')).toBe(60_000); + }); + + it('stamps the compressed count under the request route when the send ends without usage (#9506)', async () => { + // In-send compression runs for the REQUEST route, but + // setLastPromptTokenCount re-keys the fresh count to the ACTIVE + // route. If the request then ends without a usage report (abort, + // 400 — the reactive-overflow path exists for exactly those), the + // request route never stamps a count of its own, and its next + // session-token-limit gate read passes with 0 even though the + // shared compressed history's exact measure is on record. + vi.mocked(mockConfig.getModelRouteIdentity).mockReturnValue( + 'active@route', + ); + const stampChat = new GeminiChat( + mockConfig, + config, + [{ role: 'user', parts: [{ text: 'x'.repeat(720_000) }] }], + { + recordAssistantTurn: vi.fn(), + // Successful hard-rescue compression records after the + // post-compression guard passes (deferred recording). + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + stampChat.setLastPromptTokenCount(150_000, false); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'active@route', + ); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockImplementationOnce(async (chatToCompress) => { + chatToCompress.getLastPromptTokenCount(); + chatToCompress.isLastPromptTokenCountEstimated(); + return { + newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }], + info: { + originalTokenCount: 180_000, + newTokenCount: 60_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }; + }); + vi.mocked( + mockContentGenerator.generateContentStream, + ).mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'ok' }], role: 'model' }, + finishReason: 'STOP', + index: 0, + }, + ], + // No usageMetadata: the request route never stamps a count of + // its own, so the compression stamp must be readable under it. + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await stampChat.sendMessageStream( + 'override-model', + { message: 'continue' }, + 'prompt-compression-stamps-request-route', + ); + for await (const _ of stream) { + /* consume */ + } + + // The request route's keyed read sees the compressed history's + // count instead of passing the gate with 0... + expect(stampChat.getLastPromptTokenCount('override-model@route')).toBe( + 60_000, + ); + // ...and the active route still sees it through the retained entry, + // because the compressed history is shared by every route. + expect(stampChat.getLastPromptTokenCount('active@route')).toBe(60_000); + }); + + it('drops stale retained counts when a successful compression rewrites the history (#9506)', async () => { + // Compression rewrites the shared history every retained entry + // sizes. Retained pre-compression counts must not survive the + // success path, or a later keyed read adopts one and the session- + // token-limit gate blocks a prompt that fits the compressed history. + chat.setLastPromptTokenCount(691_000, false); + switchRoute('override@route'); + // The crossing retains the original route's count under its own key. + expect(chat.getLastPromptTokenCount()).toBe(0); + + vi.spyOn( + ChatCompressionService.prototype, + 'compress', + ).mockResolvedValueOnce({ + newHistory: [{ role: 'user', parts: [{ text: 'summary' }] }], + info: { + originalTokenCount: 691_000, + newTokenCount: 50_000, + compressionStatus: CompressionStatus.COMPRESSED, + }, + }); + + await chat.tryCompress('prompt-compression-drops-retained', true); + + expect(chat.getLastPromptTokenCount()).toBe(50_000); + expect(chat.getLastPromptTokenCount('gemini-pro@test0001')).toBe(0); + }); + + it('drops all retained counts when fast compression rewrites the history (#9506)', () => { + // compressFast rewrites the same shared history the other routes' + // retained entries size; clearing only the active route's entry + // would leave stale pre-compression counts adoptable by later keyed + // reads. + 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); + // Cross routes so the count is retained under the original key. + switchRoute('other-route@fast'); + expect(fastChat.getLastPromptTokenCount()).toBe(0); + + const result = fastChat.compressFast(); + + expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); + // The retained pre-compression entry did not survive the rewrite. + expect(fastChat.getLastPromptTokenCount('gemini-pro@test0001')).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 + // must not read this count — the entry invalidation has to compare + // against the request's route, not the active one. + chat.setLastPromptTokenCount(691_000, false); + + const routeGenerateContentStream = vi.fn().mockResolvedValue( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'ok' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + const routeGenerator = { + ...mockContentGenerator, + generateContentStream: routeGenerateContentStream, + } as ContentGenerator; + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: routeGenerator, + contentGeneratorConfig: { + model: 'vision-agent', + authType: AuthType.USE_OPENAI, + maxRetries: 0, + // Large enough that the zeroed-count estimate path (history walk + // + ESTIMATE_CLAMP_OVERHEAD_PAD + clamp margin) still leaves room + // for the full explicit ceiling below. + contextWindowSize: 64_000, + modalities: {}, + }, + retryAuthType: AuthType.USE_OPENAI, + model: 'vision-agent', + }); + vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({ + resolveForModel, + } as unknown as ReturnType); + vi.mocked(mockConfig.getModelRouteIdentity).mockImplementation((model) => + model ? `${model}@route` : 'gemini-pro@test0001', + ); + + const selector = 'openai:vision-agent\0https://vision.example.com/v1\0'; + const stream = await chat.sendMessageStream( + selector, + { + message: 'clamp probe', + config: { maxOutputTokens: 8_000 }, + }, + 'prompt-exact-route-clamp', + ); + for await (const _ of stream) { + /* consume */ + } + + // With the foreign count zeroed the estimate leaves room for the full + // 8_000 ceiling. Had the active route's 691_000 anchored the clamp, + // the request would have been floored at MIN_CLAMPED_OUTPUT_TOKENS. + const routeRequest = routeGenerateContentStream.mock.calls[0]?.[0] as { + config?: { maxOutputTokens?: number }; + }; + expect(routeRequest.config?.maxOutputTokens).toBe(8_000); + }); + }); + // The circuit breaker is the three-strike replacement for the old // single-shot hasFailedCompressionAttempt lock. After // MAX_CONSECUTIVE_FAILURES failures the chat stops trying to auto-compact diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 3325e3fe17..e0763a2b55 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -469,6 +469,14 @@ interface TryCompressOptions { precomputedEffectiveTokens?: number; /** Per-request overrides needed to preserve the main request cache prefix. */ requestGenerationConfig?: GenerateContentConfig; + /** + * Route the enclosing send targets. The entry adoption compares against + * this instead of the active route, so an in-send compression never + * re-adopts counts the active route retained while the request targets + * another one (#9506). Omitted by between-sends callers (manual + * `/compress`), which compress the active route's state. + */ + requestRouteKey?: string; /** * Delay writing the compression checkpoint until the caller has run any * post-compression guards that may roll the in-memory chat state back. @@ -525,6 +533,14 @@ const TRANSPORT_STREAM_RETRY_CONFIG = { */ const ESTIMATE_CLAMP_OVERHEAD_PAD = 20_000; +/** + * Cap on how many routes' token counts are retained while their route is + * not the one owning the chat's count slots (#9506). Route identities are + * bounded by the session's model routes, so this only guards pathological + * selector churn; eviction is FIFO. + */ +const MAX_RETAINED_ROUTE_COUNTS = 8; + /** * Max recovery attempts when the escalated response is also truncated. * Each attempt keeps the partial response in history and injects a recovery @@ -1856,6 +1872,35 @@ export class GeminiChat { */ private lastOutputTokenCount = 0; + /** + * Route identity (model + auth type + endpoint; see + * Config.getModelRouteIdentity) of the content generator that produced + * the counts above. API-reported sizes are wire-specific: one route's + * count cannot size another route's serialization (#9454). Undefined + * until the first count is recorded. + */ + private tokenCountsRouteKey: string | undefined = undefined; + + /** + * Token counts retained for routes other than the one currently owning + * the slots above, keyed by route identity (#9506). Crossing routes + * retains the current slots here and adopts the target's entry back + * instead of destroying the value: API-reported sizes are per-route + * state that a later turn on the same route still needs — most + * critically the session-token-limit gate, whose keyed read would + * otherwise see 0 after any foreign-route touch between turns. + * Invariant: never holds an entry for {@link tokenCountsRouteKey}. + */ + private readonly tokenCountsByRouteKey = new Map< + string, + { + promptTokenCount: number; + promptTokenCountIsEstimated: boolean; + outputTokenCount: number; + cachedContentTokenCount: number; + } + >(); + /** * Number of consecutive auto-compaction failures for this chat. The * cheap-gate NOOPs once this reaches MAX_CONSECUTIVE_FAILURES (default 3) @@ -1966,6 +2011,134 @@ export class GeminiChat { this.manualPlanExitNoticesEnabled = true; } + /** + * Identity of the currently active model route. Optional chaining keeps + * partial Config test mocks (`{} as Config`) from throwing on count + * reads/writes; a missing identity degrades to one stable key, i.e. no + * route-change invalidation. + */ + private currentRouteKey(): string { + return this.config.getModelRouteIdentity?.() ?? ''; + } + + /** + * Make the single-slot token counters describe the route identified by + * `targetRouteKey` (default: the active route). Counts recorded for a + * different route must not anchor admission, output clamping, or + * compression decisions for this one (`/model` switches rebuild the + * content generator but keep this chat instance; #9454). + * + * The crossing is NON-DESTRUCTIVE (#9506): the current slots are + * retained in {@link tokenCountsByRouteKey} under their own route key, + * and the target's retained entry — if any — is adopted back into the + * slots. Zeroing a foreign count outright let any foreign-route touch + * between two turns destroy the value before the session-token-limit + * gate (the only yield site of `SessionTokenLimitExceeded`) could read + * it back keyed by its request route. With retention, a route with no + * counts of its own still falls back to the history-walk estimate + * (slots 0), with reactive overflow recovery as the safety net, while a + * turn returning to a route that has counts reads the exact + * API-reported values. + * + * Defaults to comparing against the ACTIVE route (lazy reads on the + * getters). Send paths pass the route the upcoming request actually + * 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 + * (adopted or zeroed alongside the slots), 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 adoptTokenCountsForRoute(targetRouteKey?: string): void { + if ( + this.lastPromptTokenCount === 0 && + this.lastOutputTokenCount === 0 && + this.tokenCountsByRouteKey.size === 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 (and nothing is retained) would + // defeat the guard above. + targetRouteKey ??= this.currentRouteKey(); + if (this.tokenCountsRouteKey === targetRouteKey) { + return; + } + const retained = this.tokenCountsByRouteKey.get(targetRouteKey); + if (retained) { + this.tokenCountsByRouteKey.delete(targetRouteKey); + this.retainCurrentTokenCounts(); + debugLogger.debug( + `[token-counts] restoring retained counts for route ${targetRouteKey}`, + ); + this.lastPromptTokenCount = retained.promptTokenCount; + this.lastPromptTokenCountIsEstimated = + retained.promptTokenCountIsEstimated; + this.lastOutputTokenCount = retained.outputTokenCount; + this.tokenCountsRouteKey = targetRouteKey; + this.telemetryService?.setLastPromptTokenCount(retained.promptTokenCount); + this.telemetryService?.setLastCachedContentTokenCount( + retained.cachedContentTokenCount, + ); + return; + } + debugLogger.debug( + `[token-counts] route changed; retaining counts recorded for ` + + `${this.tokenCountsRouteKey ?? 'unknown'} (now ${targetRouteKey})`, + ); + this.retainCurrentTokenCounts(); + // Raw assignment on purpose: setLastPromptTokenCount would re-attribute + // the zero slot to the ACTIVE route. The slot is attributed to the + // TARGET route instead so it can never collide with the just-retained + // entry (retained under the evicted slot's key, which differs from the + // target) — a colliding key would make the next keyed read for the + // retained route early-return the zero slot without consulting the map. + this.lastPromptTokenCount = 0; + this.lastPromptTokenCountIsEstimated = false; + this.lastOutputTokenCount = 0; + this.tokenCountsRouteKey = targetRouteKey; + // Keep the telemetry mirror in sync, or the UI context counters + // 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); + } + + /** + * Save the current slots into {@link tokenCountsByRouteKey} under their + * owning route key so a later read keyed back to that route restores the + * exact API-reported values. Zero slots carry nothing worth retaining; + * the telemetry mirror still holds the owning route's cached-content + * count at this point, so it is captured here too. + */ + private retainCurrentTokenCounts(): void { + if ( + this.tokenCountsRouteKey === undefined || + (this.lastPromptTokenCount === 0 && this.lastOutputTokenCount === 0) + ) { + return; + } + if (this.tokenCountsByRouteKey.size >= MAX_RETAINED_ROUTE_COUNTS) { + const oldestKey = this.tokenCountsByRouteKey.keys().next().value; + if (oldestKey !== undefined) { + this.tokenCountsByRouteKey.delete(oldestKey); + } + } + this.tokenCountsByRouteKey.set(this.tokenCountsRouteKey, { + promptTokenCount: this.lastPromptTokenCount, + promptTokenCountIsEstimated: this.lastPromptTokenCountIsEstimated, + outputTokenCount: this.lastOutputTokenCount, + // Optional chaining keeps partial telemetry test mocks from throwing + // (same convention as currentRouteKey's Config lookups). + cachedContentTokenCount: + this.telemetryService?.getLastCachedContentTokenCount?.() ?? 0, + }); + } + /** * Most recent prompt-token count reported by the model for *this* chat, * mirroring the value in {@link UiTelemetryService} for the main session. @@ -1973,12 +2146,14 @@ export class GeminiChat { * count for compaction decisions, so this is always populated regardless * of whether the global telemetry is updated. */ - getLastPromptTokenCount(): number { + getLastPromptTokenCount(targetRouteKey?: string): number { + this.adoptTokenCountsForRoute(targetRouteKey); return this.lastPromptTokenCount; } /** Previous model-response tokens used by the next prompt estimate. */ getLastOutputTokenCount(): number { + this.adoptTokenCountsForRoute(); return this.lastOutputTokenCount; } @@ -2053,9 +2228,16 @@ export class GeminiChat { this.lastPromptTokenCount = count; this.lastPromptTokenCountIsEstimated = isEstimated; this.lastOutputTokenCount = 0; + this.tokenCountsRouteKey = this.currentRouteKey(); + // A fresh count supersedes anything this route retained while another + // route owned the slots. Without the delete this writer alone among the + // count writers would leave an entry for tokenCountsRouteKey behind, + // breaking the map's documented invariant (#9506). + this.tokenCountsByRouteKey.delete(this.tokenCountsRouteKey); } isLastPromptTokenCountEstimated(): boolean { + this.adoptTokenCountsForRoute(); return this.lastPromptTokenCountIsEstimated; } @@ -2083,6 +2265,14 @@ export class GeminiChat { this.lastOutputTokenCount = Number.isFinite(outputTokenCount) ? Math.max(0, outputTokenCount) : 0; + // Attribute the seeded counts to the active route so a model switch + // after resume invalidates them like any API-reported count. (Detecting + // a route that already differed at save time requires persisting route + // identity in the session transcript; tracked as a follow-up to #9454.) + this.tokenCountsRouteKey = this.currentRouteKey(); + // A fresh seed supersedes any count this route retained while another + // route owned the slots (#9506). + this.tokenCountsByRouteKey.delete(this.tokenCountsRouteKey); } /** @@ -2102,6 +2292,11 @@ export class GeminiChat { signal?: AbortSignal, options?: TryCompressOptions, ): Promise { + // Counts from a pre-switch route must not anchor compression admission + // or sizing for this route (#9454). In-send callers pass the request + // route so the adoption never re-adopts the active route's retained + // counts mid-send (#9506). + this.adoptTokenCountsForRoute(options?.requestRouteKey); const originalTokenCountIsEstimated = options?.originalTokenCountOverride === undefined && this.promptCountIsEstimateDerived(); @@ -2135,6 +2330,14 @@ export class GeminiChat { signal, }); + // ChatCompressionService reads the keyless count getters, which adopt + // the ACTIVE route — flipping the slots away from the request route + // adopted above whenever the two differ (non-exact override sends). + // Re-adopt the request route so neither the COMPRESSED stamp below nor + // the caller's post-compression sizing anchors on the flipped + // attribution (#9506). + this.adoptTokenCountsForRoute(options?.requestRouteKey); + if (info.compressionStatus === CompressionStatus.COMPRESSED && newHistory) { // ChatCompressionService owns provenance. Keep a conservative fallback // for older/custom implementations that omit the field, but preserve an @@ -2149,10 +2352,35 @@ export class GeminiChat { this.setHistory(newHistory); debugLogger.debug('[FILE_READ_CACHE] clear after auto tryCompress'); this.config.getFileReadCache().clear(); + // Compression rewrote the shared history every retained entry sizes, + // so ALL retained counts are stale — not just the current route's. + // Drop them, or a later keyed read adopts a pre-compression count and + // the session-token-limit gate blocks a prompt that fits the + // compressed history (#9506). + this.tokenCountsByRouteKey.clear(); this.setLastPromptTokenCount( info.newTokenCount, info.newTokenCountIsEstimated, ); + // setLastPromptTokenCount re-keyed the fresh count to the ACTIVE + // route, but in-send callers compress for the REQUEST route: the + // session-token-limit gate reads by that key (client.ts's sole + // SessionTokenLimitExceeded yield site), and a request that ends + // without a usage report (abort, 400 — the reactive-overflow path + // exists for exactly those) never stamps a count of its own. Re-key + // the fresh count to the request route, retaining it under the + // active key first: the compressed history is shared, so the count + // must anchor BOTH routes' next gate reads (#9506). + if ( + options?.requestRouteKey && + this.tokenCountsRouteKey !== options.requestRouteKey + ) { + this.retainCurrentTokenCounts(); + this.tokenCountsRouteKey = options.requestRouteKey; + // Same invariant as the other count writers: the fresh count + // supersedes anything the request route retained. + this.tokenCountsByRouteKey.delete(options.requestRouteKey); + } 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 @@ -2186,6 +2414,9 @@ export class GeminiChat { info: ChatCompressionInfo; microcompactMeta?: MicrocompactMeta; } { + // A pre-switch route's count must not anchor fast-compression sizing + // for the active route (#9454). + this.adoptTokenCountsForRoute(); // Use the same estimator on both sides so the NOOP gate compares // apples to apples. The API-authoritative lastPromptTokenCount is // then adjusted by the estimated delta — never replaced wholesale. @@ -2256,6 +2487,11 @@ export class GeminiChat { this.setHistory(newHistory); this.lastPromptTokenCount = adjustedTokenCount; this.lastPromptTokenCountIsEstimated = true; + this.tokenCountsRouteKey = this.currentRouteKey(); + // Fast compression rewrote the shared history every retained entry + // sizes, so ALL retained counts are stale — the other routes' entries + // describe the same pre-compression history (#9506). + this.tokenCountsByRouteKey.clear(); this.telemetryService?.setLastPromptTokenCount(adjustedTokenCount); this.consecutiveFailures = 0; @@ -2335,6 +2571,23 @@ export class GeminiChat { if (exactRoute) { model = exactRoute.model; } + // 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` + // route's decisions off the active route's counts, and a differing + // `model` param gets its own identity instead of borrowing the active + // route's. The crossing retains the current counts under their own + // route key so a later turn back on that route restores them (#9506). + this.adoptTokenCountsForRoute(requestRouteKey); const requestModalities = exactRoute?.contentGeneratorConfig.modalities ?? this.config.getEffectiveInputModalities(); @@ -2485,6 +2738,28 @@ export class GeminiChat { const lastPromptTokenCountBeforeHardRescue = this.lastPromptTokenCount; const lastPromptTokenCountWasEstimatedBeforeHardRescue = this.lastPromptTokenCountIsEstimated; + // The rescue's COMPRESSED stamp zeroes lastOutputTokenCount (via + // setLastPromptTokenCount), so the rollback below must restore the + // output half of the resurrected count pair alongside the prompt + // half, or the next turn's additive prompt estimate under-counts by + // the last response's size (#9506). + const lastOutputTokenCountBeforeHardRescue = this.lastOutputTokenCount; + // 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; + // Snapshot the retention map too: the rescue's compression consumes + // retained entries mid-flight (ChatCompressionService's keyless getter + // reads adopt the active route, deleting-and-consuming its entry) and + // a successful compression clears the map outright. Without the + // snapshot the rollback would restore the slots but not the map, + // leaving the resurrected route's count nowhere (#9506). + const retainedTokenCountsBeforeHardRescue = new Map( + this.tokenCountsByRouteKey, + ); const hardRescueFailureCountBeforeHardRescue = this.hardRescueFailureCount; if (shouldForceFromHard) { @@ -2512,6 +2787,7 @@ export class GeminiChat { pendingUserMessage: userContent, precomputedEffectiveTokens: effectiveTokens, requestGenerationConfig: params.config, + requestRouteKey, deferChatCompressionRecord: shouldForceFromHard, // Hard-rescue is force=true to bypass the cheap-gate breaker // but it remains a semantically AUTOMATIC trigger. Tag the @@ -2561,6 +2837,21 @@ export class GeminiChat { this.lastPromptTokenCount = lastPromptTokenCountBeforeHardRescue; this.lastPromptTokenCountIsEstimated = lastPromptTokenCountWasEstimatedBeforeHardRescue; + this.lastOutputTokenCount = lastOutputTokenCountBeforeHardRescue; + this.tokenCountsRouteKey = tokenCountsRouteKeyBeforeHardRescue; + // Restore the retention map alongside the slots: the rescue's + // compression consumed/cleared entries mid-flight, and without + // the restore the resurrected route's count would survive + // nowhere — its next gate read would pass with 0 (#9506). The + // snapshot predates the rescue, so it already satisfies the + // invariant (no entry for the resurrected slot key). + this.tokenCountsByRouteKey.clear(); + for (const [ + retainedRouteKey, + retainedCounts, + ] of retainedTokenCountsBeforeHardRescue) { + this.tokenCountsByRouteKey.set(retainedRouteKey, retainedCounts); + } this.telemetryService?.setLastPromptTokenCount( lastPromptTokenCountBeforeHardRescue, ); @@ -2914,6 +3205,7 @@ export class GeminiChat { params, prompt_id, requestOverrides, + requestRouteKey, turnGoalContext, // Captured by value, so the attempt records exactly the prefix // `buildAttemptContents()` just asked the model to resume from, @@ -3213,6 +3505,7 @@ export class GeminiChat { originalTokenCountOverride: reactiveOriginalTokenCount, precomputedEffectiveTokens: reactiveOriginalTokenCount, requestGenerationConfig: params.config, + requestRouteKey, trigger: 'auto', }, ); @@ -3446,6 +3739,7 @@ export class GeminiChat { attemptState.params, prompt_id, requestOverrides, + requestRouteKey, turnGoalContext, ); for await (const chunk of stream) { @@ -3857,6 +4151,14 @@ export class GeminiChat { currentUserContent, fallbackModalities ?? {}, ); + // Stamp the fallback-served counts under the REQUEST route + // key: a fallback serves on behalf of the same session + // request (the session model never changes), and the + // session-token-limit gate in Client reads the count keyed + // by the request route. Attributing the count to the + // fallback's own route would make every later gate read + // invalidate it, silently disabling the limit for any + // session ever served through fallback (#9454). for await (const event of self.makeFallbackStream( resolvedFallbackModel, fallbackRequestContents, @@ -3865,6 +4167,7 @@ export class GeminiChat { fallbackGenerator, fallbackRetryAuthType, fallbackRetryErrorCodes, + requestRouteKey, turnGoalContext, )) { const emittedUserVisibleOutput = @@ -4019,6 +4322,7 @@ export class GeminiChat { retryAuthType?: string; retryErrorCodes?: readonly number[]; }, + routeKey = this.currentRouteKey(), goalContext?: GoalTurnPermit, transportContinuationPrefix?: string, ): Promise> { @@ -4100,6 +4404,7 @@ export class GeminiChat { return this.processStreamResponse( model, rejectDegradedPlaceholderResponse(streamResponse), + routeKey, goalContext, transportContinuationPrefix, ); @@ -4113,6 +4418,7 @@ export class GeminiChat { contentGenerator: ContentGenerator, retryAuthType?: string, retryErrorCodes?: readonly number[], + routeKey?: string, goalContext?: GoalTurnPermit, ): AsyncGenerator { const stream = await this.makeApiCallAndProcessStream( @@ -4121,6 +4427,7 @@ export class GeminiChat { params, prompt_id, { contentGenerator, retryAuthType, retryErrorCodes }, + routeKey, goalContext, ); @@ -4614,6 +4921,7 @@ export class GeminiChat { private async *processStreamResponse( model: string, streamResponse: AsyncGenerator, + routeKey: string, goalContext?: GoalTurnPermit, transportContinuationPrefix?: string, ): AsyncGenerator { @@ -4820,6 +5128,17 @@ export class GeminiChat { if (lastPromptTokenCount) { // Always update the per-chat counter so this chat (including // subagents) can make its own compaction decisions. + // Retain whatever route's counts currently occupy the slots + // before overwriting them: a foreign-keyed slot holds another + // route's state that its next keyed read still needs — mid-send + // compression can leave the slots keyed to the active route + // even though this report comes from the request route (#9506). + if ( + this.tokenCountsRouteKey !== undefined && + this.tokenCountsRouteKey !== routeKey + ) { + this.retainCurrentTokenCounts(); + } this.lastPromptTokenCount = lastPromptTokenCount; this.lastPromptTokenCountIsEstimated = false; this.lastOutputTokenCount = hasUsablePromptTokenCount @@ -4830,17 +5149,23 @@ export class GeminiChat { thoughtsTokenCount, }) : 0; + // Attribute these counts to the route that reported them so a + // later model switch invalidates them (#9454). + this.tokenCountsRouteKey = routeKey; + // A fresh API report supersedes anything retained for this + // route while another route owned the slots (#9506). + this.tokenCountsByRouteKey.delete(routeKey); // 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. this.telemetryService?.setLastPromptTokenCount( lastPromptTokenCount, ); - } - if (cachedContentTokenCount && this.telemetryService) { - this.telemetryService.setLastCachedContentTokenCount( - cachedContentTokenCount, - ); + if (cachedContentTokenCount && this.telemetryService) { + this.telemetryService.setLastCachedContentTokenCount( + cachedContentTokenCount, + ); + } } } diff --git a/packages/core/src/core/goal-turn-integration.test.ts b/packages/core/src/core/goal-turn-integration.test.ts index a1c4122d5b..39b325d8a8 100644 --- a/packages/core/src/core/goal-turn-integration.test.ts +++ b/packages/core/src/core/goal-turn-integration.test.ts @@ -82,6 +82,7 @@ describe('Goal turn evidence propagation', () => { processStreamResponse: ( model: string, stream: AsyncGenerator, + routeKey: string, goalContext?: GoalTurnPermit, ) => AsyncGenerator; pendingPartialAssistantRecord: @@ -102,6 +103,7 @@ describe('Goal turn evidence propagation', () => { for await (const _ of internal.processStreamResponse( 'test-model', normalStream, + 'test-route', permit, )) { // Consume the persisted normal assistant attempt. @@ -136,6 +138,7 @@ describe('Goal turn evidence propagation', () => { for await (const _ of internal.processStreamResponse( 'test-model', partialStream, + 'test-route', permit, )) { // Consume until the deferred partial attempt is staged.