fix(core): invalidate token counts recorded for a switched model route

packages/core/src/config/config.ts

@@ -4381,6 +4381,16 @@ 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(): string {
+    return this.resolvedModelIdentity();
+  }
+
   /**
    * Returns the configured fast model selector when it resolves to an available
    * model. Bare selectors stay bare and authType-qualified selectors keep their

packages/core/src/core/geminiChat.test.ts

@@ -180,6 +180,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'),
@@ -15469,6 +15470,71 @@ 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('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);
+    });
+  });
+
   // 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

packages/core/src/core/geminiChat.ts

@@ -1845,6 +1845,15 @@ 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;
+
   /**
    * Number of consecutive auto-compaction failures for this chat. The
    * cheap-gate NOOPs once this reaches MAX_CONSECUTIVE_FAILURES (default 3)
@@ -1955,6 +1964,46 @@ 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?.() ?? '';
+  }
+
+  /**
+   * Drop token counts recorded for a different route (model, auth type, or
+   * endpoint). `/model` switches rebuild the content generator but keep
+   * this chat instance, so without this reset the previous route's
+   * API-authoritative counts would anchor admission, output clamping, and
+   * compression decisions for a different serialization (#9454). After
+   * invalidation those decisions fall back to the history-walk estimate,
+   * with reactive overflow recovery as the safety net.
+   */
+  private invalidateTokenCountsIfRouteChanged(): void {
+    if (this.lastPromptTokenCount === 0 && this.lastOutputTokenCount === 0) {
+      return;
+    }
+    const currentRoute = this.currentRouteKey();
+    if (this.tokenCountsRouteKey === currentRoute) {
+      return;
+    }
+    debugLogger.debug(
+      `[token-counts] route changed; invalidating counts recorded for ` +
+        `${this.tokenCountsRouteKey ?? 'unknown'} (now ${currentRoute})`,
+    );
+    this.lastPromptTokenCount = 0;
+    this.lastPromptTokenCountIsEstimated = false;
+    this.lastOutputTokenCount = 0;
+    this.tokenCountsRouteKey = currentRoute;
+    // Keep the telemetry mirror in sync, or the session token-limit gate
+    // and compression banners keep reading the foreign count.
+    this.telemetryService?.setLastPromptTokenCount(0);
+  }
+
   /**
    * Most recent prompt-token count reported by the model for *this* chat,
    * mirroring the value in {@link UiTelemetryService} for the main session.
@@ -1963,11 +2012,13 @@ export class GeminiChat {
    * of whether the global telemetry is updated.
    */
   getLastPromptTokenCount(): number {
+    this.invalidateTokenCountsIfRouteChanged();
     return this.lastPromptTokenCount;
   }
 
   /** Previous model-response tokens used by the next prompt estimate. */
   getLastOutputTokenCount(): number {
+    this.invalidateTokenCountsIfRouteChanged();
     return this.lastOutputTokenCount;
   }
 
@@ -2042,9 +2093,11 @@ export class GeminiChat {
     this.lastPromptTokenCount = count;
     this.lastPromptTokenCountIsEstimated = isEstimated;
     this.lastOutputTokenCount = 0;
+    this.tokenCountsRouteKey = this.currentRouteKey();
   }
 
   isLastPromptTokenCountEstimated(): boolean {
+    this.invalidateTokenCountsIfRouteChanged();
     return this.lastPromptTokenCountIsEstimated;
   }
 
@@ -2072,6 +2125,11 @@ 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();
   }
 
   /**
@@ -2091,6 +2149,9 @@ export class GeminiChat {
     signal?: AbortSignal,
     options?: TryCompressOptions,
   ): Promise<ChatCompressionInfo> {
+    // Counts from a pre-switch route must not anchor compression admission
+    // or sizing for the active route (#9454).
+    this.invalidateTokenCountsIfRouteChanged();
     const originalTokenCountIsEstimated =
       options?.originalTokenCountOverride === undefined &&
       this.promptCountIsEstimateDerived();
@@ -2175,6 +2236,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.invalidateTokenCountsIfRouteChanged();
     // 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.
@@ -2245,6 +2309,7 @@ export class GeminiChat {
     this.setHistory(newHistory);
     this.lastPromptTokenCount = adjustedTokenCount;
     this.lastPromptTokenCountIsEstimated = true;
+    this.tokenCountsRouteKey = this.currentRouteKey();
     this.telemetryService?.setLastPromptTokenCount(adjustedTokenCount);
     this.consecutiveFailures = 0;
 
@@ -2314,6 +2379,9 @@ export class GeminiChat {
     goalContext?: GoalTurnPermit,
     options?: GeminiChatSendOptions,
   ): Promise<AsyncGenerator<StreamEvent>> {
+    // Counts recorded for a pre-switch route must not anchor this send's
+    // admission/clamp/compression decisions for the active route (#9454).
+    this.invalidateTokenCountsIfRouteChanged();
     const turnGoalContext = goalContext ? { ...goalContext } : undefined;
     const fullTurnRoute = model.endsWith('\0');
     const exactRoute = fullTurnRoute
@@ -4813,6 +4881,9 @@ export class GeminiChat {
                   thoughtsTokenCount,
                 })
               : 0;
+            // Attribute these counts to the route that reported them so a
+            // later model switch invalidates them (#9454).
+            this.tokenCountsRouteKey = this.currentRouteKey();
             // 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.