feat(model-availability): introduce ModelPolicy and PolicyCatalog (#13751)

Author: adamfweidman

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

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<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[];

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

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