From f515e339816a150d15433c2ed4541e33f97d1a7d Mon Sep 17 00:00:00 2001 From: "jinjing.zzj" Date: Thu, 20 Aug 2026 23:26:21 +0800 Subject: [PATCH] fix(core): preserve route-scoped token guards (#9454) Co-authored-by: Qwen-Coder --- packages/core/src/core/client.test.ts | 29 +++++++++++++++++++ packages/core/src/core/client.ts | 8 +++--- packages/core/src/core/geminiChat.test.ts | 35 +++++++++++++++++++++++ packages/core/src/core/geminiChat.ts | 14 ++++----- 4 files changed, 75 insertions(+), 11 deletions(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 8e78abaf2a..38f4a4aea1 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -3084,6 +3084,35 @@ describe('Gemini Client (client.ts)', () => { ); 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(); + }); }); /** diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 7e3364973d..4138c20317 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -3130,9 +3130,12 @@ 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) { - const lastPromptTokenCount = this.getChat().getLastPromptTokenCount(); + const requestRouteKey = this.config.getModelRouteIdentity(model); + const lastPromptTokenCount = + this.getChat().getLastPromptTokenCount(requestRouteKey); if (lastPromptTokenCount > sessionTokenLimit) { this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); yield { @@ -3224,9 +3227,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 4e3068d5fb..daa1b8f6a1 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -15656,6 +15656,41 @@ describe('GeminiChat', async () => { ).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('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 diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index cf6583189d..a1e65c70dd 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -2025,8 +2025,8 @@ export class GeminiChat { * count for compaction decisions, so this is always populated regardless * of whether the global telemetry is updated. */ - getLastPromptTokenCount(): number { - this.invalidateTokenCountsIfRouteChanged(); + getLastPromptTokenCount(targetRouteKey?: string): number { + this.invalidateTokenCountsIfRouteChanged(targetRouteKey); return this.lastPromptTokenCount; } @@ -4938,11 +4938,11 @@ export class GeminiChat { this.telemetryService?.setLastPromptTokenCount( lastPromptTokenCount, ); - } - if (cachedContentTokenCount && this.telemetryService) { - this.telemetryService.setLastCachedContentTokenCount( - cachedContentTokenCount, - ); + if (cachedContentTokenCount && this.telemetryService) { + this.telemetryService.setLastCachedContentTokenCount( + cachedContentTokenCount, + ); + } } }