feat: reset availabilityService on /auth (#14911)

This commit is contained in:
Adam Weidman
2025-12-10 13:53:12 -08:00
committed by GitHub
parent 91b15fc9dc
commit 8c83e1ea9a
3 changed files with 35 additions and 0 deletions

View File

@@ -123,6 +123,10 @@ export class ModelAvailabilityService {
}
}
reset() {
this.health.clear();
}
private setState(model: ModelId, nextState: HealthState) {
this.health.set(model, nextState);
}

View File

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

View File

@@ -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 {