Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
26 changes: 26 additions & 0 deletions packages/core/src/availability/modelPolicy.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/

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

export type FallbackAction = 'silent' | 'prompt';

export type FailureKind = 'terminal' | 'transient' | 'not_found' | 'unknown';

export type ModelPolicyActionMap = Partial<Record<FailureKind, FallbackAction>>;

export type ModelPolicyStateMap = Partial<
Record<FailureKind, ModelHealthStatus>
>;

export interface ModelPolicy {
model: ModelId;
actions: ModelPolicyActionMap;
stateTransitions: ModelPolicyStateMap;
isLastResort?: boolean;
}

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