feat: reset availabilityService on /auth (#14911)

Author: adamfweidman

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

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

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 {