From 8c83e1ea9a60ffd7c06192f52a4c0f259db1b090 Mon Sep 17 00:00:00 2001
From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com>
Date: Wed, 10 Dec 2025 13:53:12 -0800
Subject: [PATCH] feat: reset availabilityService on /auth (#14911)

---
 .../availability/modelAvailabilityService.ts  |  4 +++
 packages/core/src/config/config.test.ts       | 27 +++++++++++++++++++
 packages/core/src/config/config.ts            |  4 +++
 3 files changed, 35 insertions(+)

diff --git a/packages/core/src/availability/modelAvailabilityService.ts b/packages/core/src/availability/modelAvailabilityService.ts
index ac32089490..8c424c0d78 100644
--- a/packages/core/src/availability/modelAvailabilityService.ts
+++ b/packages/core/src/availability/modelAvailabilityService.ts
@@ -123,6 +123,10 @@ export class ModelAvailabilityService {
     }
   }
 
+  reset() {
+    this.health.clear();
+  }
+
   private setState(model: ModelId, nextState: HealthState) {
     this.health.set(model, nextState);
   }
diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
index ea2f8b40ac..2c51151967 100644
--- a/packages/core/src/config/config.test.ts
+++ b/packages/core/src/config/config.test.ts
@@ -355,6 +355,21 @@ describe('Server Config (config.ts)', () => {
       expect(config.isInFallbackMode()).toBe(false);
     });
 
+    it('should reset model availability status', async () => {
+      const config = new Config(baseParams);
+      const service = config.getModelAvailabilityService();
+      const spy = vi.spyOn(service, 'reset');
+
+      vi.mocked(createContentGeneratorConfig).mockImplementation(
+        async (_: Config, authType: AuthType | undefined) =>
+          ({ authType }) as unknown as ContentGeneratorConfig,
+      );
+
+      await config.refreshAuth(AuthType.USE_GEMINI);
+
+      expect(spy).toHaveBeenCalled();
+    });
+
     it('should strip thoughts when switching from GenAI to Vertex', async () => {
       const config = new Config(baseParams);
 
@@ -1516,6 +1531,9 @@ describe('Config getHooks', () => {
   describe('setModel', () => {
     it('should allow setting a pro (any) model and disable fallback mode', () => {
       const config = new Config(baseParams);
+      const service = config.getModelAvailabilityService();
+      const spy = vi.spyOn(service, 'reset');
+
       config.setFallbackMode(true);
       expect(config.isInFallbackMode()).toBe(true);
 
@@ -1525,10 +1543,14 @@ describe('Config getHooks', () => {
       expect(config.getModel()).toBe(proModel);
       expect(config.isInFallbackMode()).toBe(false);
       expect(mockCoreEvents.emitModelChanged).toHaveBeenCalledWith(proModel);
+      expect(spy).toHaveBeenCalled();
     });
 
     it('should allow setting auto model from non-auto model and disable fallback mode', () => {
       const config = new Config(baseParams);
+      const service = config.getModelAvailabilityService();
+      const spy = vi.spyOn(service, 'reset');
+
       config.setFallbackMode(true);
       expect(config.isInFallbackMode()).toBe(true);
 
@@ -1537,6 +1559,7 @@ describe('Config getHooks', () => {
       expect(config.getModel()).toBe('auto');
       expect(config.isInFallbackMode()).toBe(false);
       expect(mockCoreEvents.emitModelChanged).toHaveBeenCalledWith('auto');
+      expect(spy).toHaveBeenCalled();
     });
 
     it('should allow setting auto model from auto model if it is in the fallback mode', () => {
@@ -1548,6 +1571,9 @@ describe('Config getHooks', () => {
         model: 'auto',
         usageStatisticsEnabled: false,
       });
+      const service = config.getModelAvailabilityService();
+      const spy = vi.spyOn(service, 'reset');
+
       config.setFallbackMode(true);
       expect(config.isInFallbackMode()).toBe(true);
 
@@ -1556,6 +1582,7 @@ describe('Config getHooks', () => {
       expect(config.getModel()).toBe('auto');
       expect(config.isInFallbackMode()).toBe(false);
       expect(mockCoreEvents.emitModelChanged).toHaveBeenCalledWith('auto');
+      expect(spy).toHaveBeenCalled();
     });
   });
 });
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 4c0c26647e..36c53bb431 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -703,6 +703,9 @@ export class Config {
   }
 
   async refreshAuth(authMethod: AuthType) {
+    // Reset availability service when switching auth
+    this.modelAvailabilityService.reset();
+
     // Vertex and Genai have incompatible encryption and sending history with
     // thoughtSignature from Genai to Vertex will fail, we need to strip them
     if (
@@ -828,6 +831,7 @@ export class Config {
       coreEvents.emitModelChanged(newModel);
     }
     this.setFallbackMode(false);
+    this.modelAvailabilityService.reset();
   }
 
   getActiveModel(): string {