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: 2 additions & 0 deletions packages/core/src/availability/modelAvailabilityService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ export type UnavailabilityReason =
| TurnUnavailabilityReason
| 'unknown';

export type ModelHealthStatus = 'terminal' | 'sticky_retry';

type HealthState =
| { status: 'terminal'; reason: TerminalUnavailabilityReason }
| {
Expand Down
51 changes: 51 additions & 0 deletions packages/core/src/availability/modelPolicy.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import type { ModelHealthStatus, ModelId } from './modelAvailabilityService.js';

/**
* Whether to prompt the user or fallback silently on a model API failure.
*/
export type FallbackAction = 'silent' | 'prompt';

/**
* Type of possible errors from model API failures.
*/
export type FailureKind = 'terminal' | 'transient' | 'not_found' | 'unknown';

/**
* Map from model API failure reason to user interaction.
*/
export type ModelPolicyActionMap = Partial<Record<FailureKind, FallbackAction>>;

/**
* What state (e.g. Terminal, Sticky Retry) to set a model after failed API call.
*/
export type ModelPolicyStateMap = Partial<
Record<FailureKind, ModelHealthStatus>
>;

/**
* Defines the policy for a single model in the availability chain.
*
* This includes:
* - Which model this policy applies to.
* - What actions to take (prompt vs silent fallback) for different failure kinds.
* - How the model's health status should transition upon failure.
* - Whether this model is considered a "last resort" (i.e. use if all models are unavailable).
*/
export interface ModelPolicy {
model: ModelId;
actions: ModelPolicyActionMap;
stateTransitions: ModelPolicyStateMap;
isLastResort?: boolean;
}

/**
* A chain of model policies defining the priority and fallback behavior.
* The first model in the chain is the primary model.
*/
export type ModelPolicyChain = ModelPolicy[];
93 changes: 93 additions & 0 deletions packages/core/src/availability/policyCatalog.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import { describe, it, expect } from 'vitest';
import {
createDefaultPolicy,
getModelPolicyChain,
validateModelPolicyChain,
} from './policyCatalog.js';
import {
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL,
} from '../config/models.js';

describe('policyCatalog', () => {
it('returns preview chain when preview enabled', () => {
const chain = getModelPolicyChain({ previewEnabled: true });
expect(chain[0]?.model).toBe(PREVIEW_GEMINI_MODEL);
expect(chain).toHaveLength(3);
});

it('returns default chain when preview disabled', () => {
const chain = getModelPolicyChain({ previewEnabled: false });
expect(chain[0]?.model).toBe(DEFAULT_GEMINI_MODEL);
expect(chain).toHaveLength(2);
});

it('marks preview transients as sticky retries', () => {
const [previewPolicy] = getModelPolicyChain({ previewEnabled: true });
expect(previewPolicy.model).toBe(PREVIEW_GEMINI_MODEL);
expect(previewPolicy.stateTransitions.transient).toBe('sticky_retry');
});

it('applies default actions and state transitions for unspecified kinds', () => {
const [previewPolicy] = getModelPolicyChain({ previewEnabled: true });
expect(previewPolicy.stateTransitions.not_found).toBe('terminal');
expect(previewPolicy.stateTransitions.unknown).toBe('terminal');
expect(previewPolicy.actions.unknown).toBe('prompt');
});

it('clones policy maps so edits do not leak between calls', () => {
const firstCall = getModelPolicyChain({ previewEnabled: false });
firstCall[0]!.actions.terminal = 'silent';
const secondCall = getModelPolicyChain({ previewEnabled: false });
expect(secondCall[0]!.actions.terminal).toBe('prompt');
});

it('passes when there is exactly one last-resort policy', () => {
const validChain = [
createDefaultPolicy('test-model'),
{ ...createDefaultPolicy('last-resort'), isLastResort: true },
];
expect(() => validateModelPolicyChain(validChain)).not.toThrow();
});

it('fails when no policies are marked last-resort', () => {
const chain = [
createDefaultPolicy('model-a'),
createDefaultPolicy('model-b'),
];
expect(() => validateModelPolicyChain(chain)).toThrow(
'must include an `isLastResort`',
);
});

it('fails when a single-model chain is not last-resort', () => {
const chain = [createDefaultPolicy('lonely-model')];
expect(() => validateModelPolicyChain(chain)).toThrow(
'must include an `isLastResort`',
);
});

it('fails when multiple policies are marked last-resort', () => {
const chain = [
{ ...createDefaultPolicy('model-a'), isLastResort: true },
{ ...createDefaultPolicy('model-b'), isLastResort: true },
];
expect(() => validateModelPolicyChain(chain)).toThrow(
'must only have one `isLastResort`',
);
});

it('createDefaultPolicy seeds default actions and states', () => {
const policy = createDefaultPolicy('custom');
expect(policy.actions.terminal).toBe('prompt');
expect(policy.actions.unknown).toBe('prompt');
expect(policy.stateTransitions.terminal).toBe('terminal');
expect(policy.stateTransitions.unknown).toBe('terminal');
});
});
118 changes: 118 additions & 0 deletions packages/core/src/availability/policyCatalog.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

import type {
ModelPolicy,
ModelPolicyActionMap,
ModelPolicyChain,
ModelPolicyStateMap,
} from './modelPolicy.js';
import {
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL,
} from '../config/models.js';
import type { UserTierId } from '../code_assist/types.js';

// actions and stateTransitions are optional when defining ModelPolicy
type PolicyConfig = Omit<ModelPolicy, 'actions' | 'stateTransitions'> & {
actions?: ModelPolicyActionMap;
stateTransitions?: ModelPolicyStateMap;
};

export interface ModelPolicyOptions {
previewEnabled: boolean;
userTier?: UserTierId;
}

const DEFAULT_ACTIONS: ModelPolicyActionMap = {
terminal: 'prompt',
transient: 'prompt',
not_found: 'prompt',
unknown: 'prompt',
};

const DEFAULT_STATE: ModelPolicyStateMap = {
terminal: 'terminal',
transient: 'terminal',
not_found: 'terminal',
unknown: 'terminal',
};

const DEFAULT_CHAIN: ModelPolicyChain = [
definePolicy({ model: DEFAULT_GEMINI_MODEL }),
definePolicy({ model: DEFAULT_GEMINI_FLASH_MODEL, isLastResort: true }),
];

const PREVIEW_CHAIN: ModelPolicyChain = [
definePolicy({
model: PREVIEW_GEMINI_MODEL,
stateTransitions: { transient: 'sticky_retry' },
}),
definePolicy({ model: DEFAULT_GEMINI_MODEL }),
definePolicy({ model: DEFAULT_GEMINI_FLASH_MODEL, isLastResort: true }),
];

/**
* Returns the default ordered model policy chain for the user.
*/
export function getModelPolicyChain(
options: ModelPolicyOptions,
): ModelPolicyChain {
if (options.previewEnabled) {
return cloneChain(PREVIEW_CHAIN);
}

return cloneChain(DEFAULT_CHAIN);
}

/**
* Provides a default policy scaffold for models not present in the catalog.
*/
export function createDefaultPolicy(model: string): ModelPolicy {
return definePolicy({ model });
}

export function validateModelPolicyChain(chain: ModelPolicyChain): void {
if (chain.length === 0) {
throw new Error('Model policy chain must include at least one model.');
}
const lastResortCount = chain.filter((policy) => policy.isLastResort).length;
if (lastResortCount === 0) {
throw new Error('Model policy chain must include an `isLastResort` model.');
}
if (lastResortCount > 1) {
throw new Error('Model policy chain must only have one `isLastResort`.');
}
}

/**
* Helper to define a ModelPolicy with default actions and state transitions.
* Ensures every policy is a fresh instance to avoid shared state.
*/
function definePolicy(config: PolicyConfig): ModelPolicy {
return {
model: config.model,
isLastResort: config.isLastResort,
actions: { ...DEFAULT_ACTIONS, ...(config.actions ?? {}) },
stateTransitions: {
...DEFAULT_STATE,
...(config.stateTransitions ?? {}),
},
};
}

function clonePolicy(policy: ModelPolicy): ModelPolicy {
return {
...policy,
actions: { ...policy.actions },
stateTransitions: { ...policy.stateTransitions },
};
}

function cloneChain(chain: ModelPolicyChain): ModelPolicyChain {
return chain.map(clonePolicy);
}