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

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@<sha-prefix>` 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);
+  });
+});

packages/core/src/config/config.ts

@@ -4366,12 +4366,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')
@@ -4381,6 +4385,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

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),
@@ -3018,6 +3019,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
@@ -3043,6 +3045,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<Config['getBaseLlmClient']>);
+      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);
+    });
   });
 
   /**
@@ -9137,6 +9303,7 @@ Other open files:
       const mockChat: Partial<GeminiChat> = {
         addHistory: vi.fn(),
         getHistory: vi.fn().mockReturnValue([]),
+        getLastPromptTokenCount: vi.fn().mockReturnValue(9999),
       };
       client['chat'] = mockChat as GeminiChat;
 

packages/core/src/core/client.ts

@@ -3130,10 +3130,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 {
@@ -3225,9 +3240,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