diff --git a/packages/core/src/availability/modelAvailabilityService.ts b/packages/core/src/availability/modelAvailabilityService.ts index 0a08c28655c..f804c5dacd5 100644 --- a/packages/core/src/availability/modelAvailabilityService.ts +++ b/packages/core/src/availability/modelAvailabilityService.ts @@ -14,6 +14,8 @@ export type UnavailabilityReason = | TurnUnavailabilityReason | 'unknown'; +export type ModelHealthStatus = 'terminal' | 'sticky_retry'; + type HealthState = | { status: 'terminal'; reason: TerminalUnavailabilityReason } | { diff --git a/packages/core/src/availability/modelPolicy.ts b/packages/core/src/availability/modelPolicy.ts new file mode 100644 index 00000000000..296047edb6b --- /dev/null +++ b/packages/core/src/availability/modelPolicy.ts @@ -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>; + +/** + * What state (e.g. Terminal, Sticky Retry) to set a model after failed API call. + */ +export type ModelPolicyStateMap = Partial< + Record +>; + +/** + * 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[]; diff --git a/packages/core/src/availability/policyCatalog.test.ts b/packages/core/src/availability/policyCatalog.test.ts new file mode 100644 index 00000000000..092123ff65a --- /dev/null +++ b/packages/core/src/availability/policyCatalog.test.ts @@ -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'); + }); +}); diff --git a/packages/core/src/availability/policyCatalog.ts b/packages/core/src/availability/policyCatalog.ts new file mode 100644 index 00000000000..74b810d7de7 --- /dev/null +++ b/packages/core/src/availability/policyCatalog.ts @@ -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 & { + 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); +}