Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/core/src/availability/testUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ export function createAvailabilityServiceMock(
markHealthy: vi.fn(),
markRetryOncePerTurn: vi.fn(),
consumeStickyAttempt: vi.fn(),
snapshot: vi.fn(),
snapshot: vi.fn().mockReturnValue({ available: true }),
resetTurn: vi.fn(),
selectFirstAvailable: vi.fn().mockReturnValue(selection),
};
Expand Down
35 changes: 35 additions & 0 deletions packages/core/src/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -863,6 +863,16 @@ describe('Server Config (config.ts)', () => {
expect(GeminiClient).toHaveBeenCalledWith(config);
});

it('should clear fallback overrides when refreshing auth', async () => {
const config = new Config(baseParams);
config.activateFallbackMode('fallback-model', 'failed-model');
expect(config.getFallbackOverride('failed-model')).toBe('fallback-model');

await config.refreshAuth(AuthType.USE_GEMINI);

expect(config.getFallbackOverride('failed-model')).toBeUndefined();
});

it('should pass Vertex AI routing settings when refreshing auth', async () => {
const vertexAiRouting = {
requestType: 'shared' as const,
Expand Down Expand Up @@ -1902,6 +1912,21 @@ describe('Server Config (config.ts)', () => {
);
});

it('clears fallback overrides when session changes', async () => {
const config = new Config({
...baseParams,
sessionId: 'session-one',
});
await config.initialize();

config.activateFallbackMode('fallback-model', 'failed-model');
expect(config.getFallbackOverride('failed-model')).toBe('fallback-model');

config.setSessionId('session-two');

expect(config.getFallbackOverride('failed-model')).toBeUndefined();
});

it('does not throw when changing sessions before the previous plans dir exists', async () => {
const config = new Config({
...baseParams,
Expand Down Expand Up @@ -2715,6 +2740,16 @@ describe('Config getHooks', () => {
expect(spy).toHaveBeenCalled();
});

it('should preserve fallback overrides when setting a new model', () => {
const config = new Config(baseParams);
config.activateFallbackMode('fallback-model', 'failed-model');
expect(config.getFallbackOverride('failed-model')).toBe('fallback-model');

config.setModel('new-model');

expect(config.getFallbackOverride('failed-model')).toBe('fallback-model');
});

it('should allow setting auto model from auto model and reset availability', () => {
const config = new Config({
cwd: '/tmp',
Expand Down
35 changes: 33 additions & 2 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -833,6 +833,7 @@ export class Config implements McpContext, AgentLoopContext {
private ideMode: boolean;

private _activeModel: string;
private fallbackOverrides = new Map<string, string>();
Comment thread
cocosheng-g marked this conversation as resolved.
private readonly maxSessionTurns: number;
private readonly listSessions: boolean;
private readonly deleteSession: string | undefined;
Expand Down Expand Up @@ -1567,6 +1568,8 @@ export class Config implements McpContext, AgentLoopContext {
) {
// Reset availability service when switching auth
this.modelAvailabilityService.reset();
this.fallbackOverrides.clear();
this.modelConfigService.clearRuntimeOverrides();

// Vertex and Genai have incompatible encryption and sending history with
// thoughtSignature from Genai to Vertex will fail, we need to strip them
Expand Down Expand Up @@ -1828,6 +1831,8 @@ export class Config implements McpContext, AgentLoopContext {
this._sessionId = sessionId;
this.storage.setSessionId(sessionId);
this.trackerService = undefined;
this.fallbackOverrides.clear();
this.modelConfigService.clearRuntimeOverrides();
this.approvedPlanPath = undefined;
this.topicState.reset();
this.skillManager.reset();
Expand Down Expand Up @@ -1923,14 +1928,40 @@ export class Config implements McpContext, AgentLoopContext {
this.modelAvailabilityService.reset();
}

activateFallbackMode(model: string): void {
this.setModel(model, true);
activateFallbackMode(model: string, failedModel?: string): void {
if (this.getActiveModel() !== model) {
this.setModel(model, true);
}
if (failedModel) {
// Chained fallback mitigation: If we already have overrides that point to the model
// that just failed, we need to update them to point to the new fallback model.
// e.g. A -> B, then B fails and we fallback to C. We must update A to point to C.
for (const [source, target] of this.fallbackOverrides.entries()) {
if (target === failedModel) {
this.fallbackOverrides.set(source, model);
this.modelConfigService.registerRuntimeModelOverride({
match: { model: source },
modelConfig: { model },
});
}
}

this.fallbackOverrides.set(failedModel, model);
this.modelConfigService.registerRuntimeModelOverride({
match: { model: failedModel },
modelConfig: { model },
});
}
Comment thread
cocosheng-g marked this conversation as resolved.
const authType = this.getContentGeneratorConfig()?.authType;
if (authType) {
logFlashFallback(this, new FlashFallbackEvent(authType));
}
}

getFallbackOverride(model: string): string | undefined {
return this.fallbackOverrides.get(model);
}

getActiveModel(): string {
return this._activeModel ?? this.model;
}
Expand Down
53 changes: 53 additions & 0 deletions packages/core/src/config/flashFallback.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -73,5 +73,58 @@ describe('Flash Model Fallback Configuration', () => {
expect.any(FlashFallbackEvent),
);
});

it('should set fallback override when failedModel is provided and register runtime override', () => {
config.activateFallbackMode(
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
);
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
expect(config.getFallbackOverride(DEFAULT_GEMINI_MODEL)).toBe(
DEFAULT_GEMINI_FLASH_MODEL,
);

// Verify it registers the runtime model override with ModelConfigService
expect(
config
.getModelConfigService()
.getResolvedConfig({ model: DEFAULT_GEMINI_MODEL }).model,
).toBe(DEFAULT_GEMINI_FLASH_MODEL);
});

it('should flatten override chains when a model that was previously a target fails', () => {
// 1. Initial fallback: A -> B
config.activateFallbackMode('model-B', 'model-A');
expect(config.getFallbackOverride('model-A')).toBe('model-B');
expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-A' })
.model,
).toBe('model-B');

// 2. Chained fallback: B fails, fallback to C
// This should update A -> C as well.
config.activateFallbackMode('model-C', 'model-B');

expect(config.getFallbackOverride('model-A')).toBe('model-C');
expect(config.getFallbackOverride('model-B')).toBe('model-C');

expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-A' })
.model,
).toBe('model-C');
expect(
config.getModelConfigService().getResolvedConfig({ model: 'model-B' })
.model,
).toBe('model-C');
});

it('should not reset availability service if model has not changed', () => {
const resetSpy = vi.spyOn(config.getModelAvailabilityService(), 'reset');
const currentModel = config.getActiveModel();

config.activateFallbackMode(currentModel);

expect(resetSpy).not.toHaveBeenCalled();
});
});
});
7 changes: 7 additions & 0 deletions packages/core/src/fallback/handler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ describe('handleFallback', () => {
expect(policyConfig.getFallbackModelHandler).not.toHaveBeenCalled();
expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith(
DEFAULT_GEMINI_FLASH_MODEL,
undefined,
);
} finally {
chainSpy.mockRestore();
Expand All @@ -207,6 +208,9 @@ describe('handleFallback', () => {
selectedModel: MOCK_PRO_MODEL,
skipped: [],
});
// Mock activeModel to be unavailable so the utility bypass heuristic is skipped
vi.mocked(availability.snapshot).mockReturnValue({ available: false });

policyHandler.mockResolvedValue('retry_once');

await handleFallback(
Expand Down Expand Up @@ -351,6 +355,8 @@ describe('handleFallback', () => {
vi.mocked(policyConfig.getModel).mockReturnValue(
DEFAULT_GEMINI_MODEL_AUTO,
);
// Mock activeModel to be unavailable so the utility bypass heuristic is skipped
vi.mocked(availability.snapshot).mockReturnValue({ available: false });

const result = await handleFallback(
policyConfig,
Expand Down Expand Up @@ -383,6 +389,7 @@ describe('handleFallback', () => {
expect(result).toBe(true);
expect(policyConfig.activateFallbackMode).toHaveBeenCalledWith(
FALLBACK_MODEL,
undefined,
);
// TODO: add logging expect statement
});
Expand Down
35 changes: 31 additions & 4 deletions packages/core/src/fallback/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,17 @@ export async function handleFallback(
return { service: availability, policy: failedPolicy };
};

const activeModel = config.getActiveModel();
let fallbackModel: string;

if (!candidates.length) {
if (
failedModel !== activeModel &&
availability.snapshot(activeModel).available
) {
applyAvailabilityTransition(getAvailabilityContext, failureKind);
return processIntent(config, 'retry_always', activeModel, failedModel);
}
fallbackModel = failedModel;
} else {
const selection = availability.selectFirstAvailable(
Expand All @@ -70,9 +79,21 @@ export async function handleFallback(
// failureKind is already declared and calculated above
const action = resolvePolicyAction(failureKind, selectedPolicy);

if (action === 'silent') {
if (
action === 'silent' ||
(fallbackModel === activeModel && failedModel !== activeModel)
) {
applyAvailabilityTransition(getAvailabilityContext, failureKind);
return processIntent(config, 'retry_always', fallbackModel);
// For standard auto-routing (silent), we only update the active model, so don't pass failedModel.
// For utility bypass, we want a hard runtime override, so pass failedModel.
const overrideFailedModel =
failedModel !== activeModel ? failedModel : undefined;
return processIntent(
config,
'retry_always',
fallbackModel,
overrideFailedModel,
);
}

// This will be used in the future when FallbackRecommendation is passed through UI
Expand Down Expand Up @@ -103,7 +124,12 @@ export async function handleFallback(
applyAvailabilityTransition(getAvailabilityContext, failureKind);
}

return await processIntent(config, intent, fallbackModel);
return await processIntent(
config,
intent,
fallbackModel,
failedModel !== activeModel ? failedModel : undefined,
);
} catch (handlerError) {
debugLogger.error('Fallback handler failed:', handlerError);
return null;
Expand Down Expand Up @@ -131,12 +157,13 @@ async function processIntent(
config: Config,
intent: FallbackIntent | null,
fallbackModel: string,
failedModel?: string,
): Promise<boolean> {
switch (intent) {
case 'retry_always':
// TODO(telemetry): Implement generic fallback event logging. Existing
// logFlashFallback is specific to a single Model.
config.activateFallbackMode(fallbackModel);
config.activateFallbackMode(fallbackModel, failedModel);
return true;

case 'retry_once':
Expand Down
25 changes: 25 additions & 0 deletions packages/core/src/services/modelConfigService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -668,6 +668,31 @@ describe('ModelConfigService', () => {
// Specificity should win over order
expect(resolved.generateContentConfig.temperature).toBe(0.1);
});

it('should clear runtime overrides', () => {
const config: ModelConfigServiceConfig = {
aliases: {},
overrides: [],
};
const service = new ModelConfigService(config);

service.registerRuntimeModelOverride({
match: { model: 'gemini-pro' },
modelConfig: { generateContentConfig: { temperature: 0.99 } },
});

expect(
service.getResolvedConfig({ model: 'gemini-pro' }).generateContentConfig
.temperature,
).toBe(0.99);

service.clearRuntimeOverrides();

expect(
service.getResolvedConfig({ model: 'gemini-pro' }).generateContentConfig
.temperature,
).toBeUndefined();
});
});

describe('custom aliases', () => {
Expand Down
4 changes: 4 additions & 0 deletions packages/core/src/services/modelConfigService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@
export interface ModelDefinition {
displayName?: string;
tier?: string; // 'pro' | 'flash' | 'flash-lite' | 'custom' | 'auto'
family?: string; // The gemini family, e.g. 'gemini-3' | 'gemini-2'

Check warning on line 68 in packages/core/src/services/modelConfigService.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-2". Please make sure this change is appropriate to submit.
isPreview?: boolean;
// Specifies whether the model should be visible in the dialog.
isVisible?: boolean;
Expand Down Expand Up @@ -344,6 +344,10 @@
this.runtimeOverrides.push(override);
}

clearRuntimeOverrides(): void {
this.runtimeOverrides.length = 0;
}

/**
* Resolves a model configuration by merging settings from aliases and applying overrides.
*
Expand Down
Loading