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,45 @@ 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);
+    });
   });
 
   /**
@@ -9137,6 +9178,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

@@ -3132,8 +3132,7 @@ export class GeminiClient {
       // forceFullIdeContext flip).
       const sessionTokenLimit = this.config.getSessionTokenLimit();
       if (sessionTokenLimit > 0) {
-        const lastPromptTokenCount =
-          uiTelemetryService.getLastPromptTokenCount();
+        const lastPromptTokenCount = this.getChat().getLastPromptTokenCount();
         if (lastPromptTokenCount > sessionTokenLimit) {
           this.cancelPendingMemoryPrefetch('no_safe_delivery_point');
           yield {