fix(core): preserve route-scoped token guards (#9454)

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>

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();
+    });
   });
 
   /**

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

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

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,
+              );
+            }
           }
         }