From eaf60149e5ca7500f54f7186de1e3785f5101bd7 Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Wed, 19 Nov 2025 12:33:29 -0500 Subject: [PATCH 1/3] feat: add modelAvailabilityService --- .../modelAvailabilityService.test.ts | 105 ++++++++++++++++ .../availability/modelAvailabilityService.ts | 118 ++++++++++++++++++ 2 files changed, 223 insertions(+) create mode 100644 packages/core/src/availability/modelAvailabilityService.test.ts create mode 100644 packages/core/src/availability/modelAvailabilityService.ts diff --git a/packages/core/src/availability/modelAvailabilityService.test.ts b/packages/core/src/availability/modelAvailabilityService.test.ts new file mode 100644 index 00000000000..cf0ebe233dd --- /dev/null +++ b/packages/core/src/availability/modelAvailabilityService.test.ts @@ -0,0 +1,105 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi, beforeEach } from 'vitest'; +import { ModelAvailabilityService } from './modelAvailabilityService.js'; + +describe('ModelAvailabilityService', () => { + let service: ModelAvailabilityService; + const model = 'test-model'; + + beforeEach(() => { + service = new ModelAvailabilityService(); + vi.useRealTimers(); + }); + + it('returns available snapshot when no state recorded', () => { + expect(service.snapshot(model)).toEqual({ available: true }); + }); + + it('tracks retry-once-per-turn failures', () => { + service.markRetryOncePerTurn(model); + expect(service.snapshot(model)).toEqual({ available: true }); + + service.consumeStickyAttempt(model); + expect(service.snapshot(model)).toEqual({ + available: false, + reason: 'retry_once_per_turn', + }); + + service.resetTurn(); + expect(service.snapshot(model)).toEqual({ available: true }); + }); + + it('tracks terminal failures', () => { + service.markTerminal(model, 'quota'); + expect(service.snapshot(model)).toEqual({ + available: false, + reason: 'quota', + }); + }); + + it('selects models respecting terminal and sticky states', () => { + const stickyModel = 'stick-model'; + const healthyModel = 'healthy-model'; + + service.markTerminal(model, 'capacity'); + service.markRetryOncePerTurn(stickyModel); + + const first = service.selectFirstAvailable([ + model, + stickyModel, + healthyModel, + ]); + expect(first).toEqual({ + selected: stickyModel, + attempts: 1, + skipped: [ + { + model, + reason: 'capacity', + }, + ], + }); + + service.consumeStickyAttempt(stickyModel); + const second = service.selectFirstAvailable([ + model, + stickyModel, + healthyModel, + ]); + expect(second).toEqual({ + selected: healthyModel, + skipped: [ + { + model, + reason: 'capacity', + }, + { + model: stickyModel, + reason: 'retry_once_per_turn', + }, + ], + }); + + service.resetTurn(); + const third = service.selectFirstAvailable([ + model, + stickyModel, + healthyModel, + ]); + expect(third).toEqual({ + selected: stickyModel, + attempts: 1, + skipped: [ + { + model, + reason: 'capacity', + }, + ], + }); + }); +}); diff --git a/packages/core/src/availability/modelAvailabilityService.ts b/packages/core/src/availability/modelAvailabilityService.ts new file mode 100644 index 00000000000..447b26828bc --- /dev/null +++ b/packages/core/src/availability/modelAvailabilityService.ts @@ -0,0 +1,118 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export type ModelId = string; + +export type TerminalAvailabilityReason = 'quota' | 'capacity'; +export type TurnAvailabilityReason = 'retry_once_per_turn'; + +export type AvailabilityReason = + | TerminalAvailabilityReason + | TurnAvailabilityReason + | 'unknown'; + +type HealthState = + | { status: 'terminal'; reason: TerminalAvailabilityReason } + | { status: 'sticky_retry'; reason: TurnAvailabilityReason }; + +export interface ModelAvailabilitySnapshot { + available: boolean; + reason?: AvailabilityReason; +} + +export interface ModelSelectionResult { + selected: ModelId | null; + attempts?: number; + skipped: Array<{ + model: ModelId; + reason: AvailabilityReason; + }>; +} + +export class ModelAvailabilityService { + private readonly health = new Map(); + private readonly consumedSticky = new Set(); + + markTerminal(model: ModelId, reason: TerminalAvailabilityReason) { + this.setState(model, { + status: 'terminal', + reason, + }); + } + + markHealthy(model: ModelId) { + this.clearState(model); + } + + markRetryOncePerTurn(model: ModelId) { + // Only reset consumption if we are not already in the sticky_retry state. + // This prevents infinite loops if the model fails repeatedly in the same turn. + const currentState = this.health.get(model); + if (currentState?.status !== 'sticky_retry') { + this.consumedSticky.delete(model); + } + + this.setState(model, { + status: 'sticky_retry', + reason: 'retry_once_per_turn', + }); + } + + consumeStickyAttempt(model: ModelId) { + if (this.health.get(model)?.status === 'sticky_retry') { + this.consumedSticky.add(model); + } + } + + snapshot(model: ModelId): ModelAvailabilitySnapshot { + const state = this.health.get(model); + if (!state) { + return { available: true }; + } + if (state.status === 'terminal') { + return { available: false, reason: state.reason }; + } + const available = !this.consumedSticky.has(model); + return available + ? { available: true } + : { available: false, reason: state.reason }; + } + + selectFirstAvailable(models: ModelId[]): ModelSelectionResult { + const skipped: ModelSelectionResult['skipped'] = []; + for (const model of models) { + const state = this.health.get(model); + if (!state) { + return { selected: model, skipped }; + } + if (state.status === 'terminal') { + skipped.push({ model, reason: state.reason }); + continue; + } + if (state.status === 'sticky_retry') { + if (this.consumedSticky.has(model)) { + skipped.push({ model, reason: state.reason }); + continue; + } + return { selected: model, attempts: 1, skipped }; + } + } + return { selected: null, skipped }; + } + + resetTurn() { + this.consumedSticky.clear(); + } + + private setState(model: ModelId, nextState: HealthState) { + this.health.set(model, nextState); + } + + private clearState(model: ModelId) { + this.health.delete(model); + this.consumedSticky.delete(model); + } +} From 7ab0f74b792215c8b342b082ac40ee51cf3a46da Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Wed, 19 Nov 2025 13:04:48 -0500 Subject: [PATCH 2/3] cleanup --- .../modelAvailabilityService.test.ts | 9 ++++ .../availability/modelAvailabilityService.ts | 46 +++++++++---------- 2 files changed, 31 insertions(+), 24 deletions(-) diff --git a/packages/core/src/availability/modelAvailabilityService.test.ts b/packages/core/src/availability/modelAvailabilityService.test.ts index cf0ebe233dd..26634467e02 100644 --- a/packages/core/src/availability/modelAvailabilityService.test.ts +++ b/packages/core/src/availability/modelAvailabilityService.test.ts @@ -42,6 +42,15 @@ describe('ModelAvailabilityService', () => { }); }); + it('does not override terminal failure with sticky failure', () => { + service.markTerminal(model, 'quota'); + service.markRetryOncePerTurn(model); + expect(service.snapshot(model)).toEqual({ + available: false, + reason: 'quota', + }); + }); + it('selects models respecting terminal and sticky states', () => { const stickyModel = 'stick-model'; const healthyModel = 'healthy-model'; diff --git a/packages/core/src/availability/modelAvailabilityService.ts b/packages/core/src/availability/modelAvailabilityService.ts index 447b26828bc..f5e5192692c 100644 --- a/packages/core/src/availability/modelAvailabilityService.ts +++ b/packages/core/src/availability/modelAvailabilityService.ts @@ -48,9 +48,14 @@ export class ModelAvailabilityService { } markRetryOncePerTurn(model: ModelId) { + const currentState = this.health.get(model); + // Do not override a terminal failure with a transient one. + if (currentState?.status === 'terminal') { + return; + } + // Only reset consumption if we are not already in the sticky_retry state. // This prevents infinite loops if the model fails repeatedly in the same turn. - const currentState = this.health.get(model); if (currentState?.status !== 'sticky_retry') { this.consumedSticky.delete(model); } @@ -69,35 +74,28 @@ export class ModelAvailabilityService { snapshot(model: ModelId): ModelAvailabilitySnapshot { const state = this.health.get(model); - if (!state) { - return { available: true }; - } - if (state.status === 'terminal') { - return { available: false, reason: state.reason }; + const isUnavailable = + state?.status === 'terminal' || + (state?.status === 'sticky_retry' && this.consumedSticky.has(model)); + + if (isUnavailable) { + return { available: false, reason: state!.reason }; } - const available = !this.consumedSticky.has(model); - return available - ? { available: true } - : { available: false, reason: state.reason }; + return { available: true }; } selectFirstAvailable(models: ModelId[]): ModelSelectionResult { const skipped: ModelSelectionResult['skipped'] = []; + for (const model of models) { - const state = this.health.get(model); - if (!state) { - return { selected: model, skipped }; - } - if (state.status === 'terminal') { - skipped.push({ model, reason: state.reason }); - continue; - } - if (state.status === 'sticky_retry') { - if (this.consumedSticky.has(model)) { - skipped.push({ model, reason: state.reason }); - continue; - } - return { selected: model, attempts: 1, skipped }; + const snapshot = this.snapshot(model); + if (snapshot.available) { + const state = this.health.get(model); + // A sticky model is being attempted, so note that. + const attempts = state?.status === 'sticky_retry' ? 1 : undefined; + return { selected: model, skipped, attempts }; + } else { + skipped.push({ model, reason: snapshot.reason ?? 'unknown' }); } } return { selected: null, skipped }; From 99392de8d50276db659e2038e539aeff6d8d5426 Mon Sep 17 00:00:00 2001 From: Adam Weidman Date: Fri, 21 Nov 2025 10:59:51 -0500 Subject: [PATCH 3/3] move consumed into healthstate --- .../modelAvailabilityService.test.ts | 51 ++++++++++++++++ .../availability/modelAvailabilityService.ts | 59 ++++++++++++------- 2 files changed, 88 insertions(+), 22 deletions(-) diff --git a/packages/core/src/availability/modelAvailabilityService.test.ts b/packages/core/src/availability/modelAvailabilityService.test.ts index 26634467e02..bddb3b49469 100644 --- a/packages/core/src/availability/modelAvailabilityService.test.ts +++ b/packages/core/src/availability/modelAvailabilityService.test.ts @@ -111,4 +111,55 @@ describe('ModelAvailabilityService', () => { ], }); }); + + it('preserves consumed state when marking retry-once-per-turn again', () => { + service.markRetryOncePerTurn(model); + service.consumeStickyAttempt(model); + + // It is currently consumed + expect(service.snapshot(model).available).toBe(false); + + // Marking it again should not reset the consumed flag + service.markRetryOncePerTurn(model); + expect(service.snapshot(model).available).toBe(false); + }); + + it('clears consumed state when marked healthy', () => { + service.markRetryOncePerTurn(model); + service.consumeStickyAttempt(model); + expect(service.snapshot(model).available).toBe(false); + + service.markHealthy(model); + expect(service.snapshot(model).available).toBe(true); + + // If we mark it sticky again, it should be fresh (not consumed) + service.markRetryOncePerTurn(model); + expect(service.snapshot(model).available).toBe(true); + }); + + it('resetTurn resets consumed state for multiple sticky models', () => { + const model2 = 'model-2'; + service.markRetryOncePerTurn(model); + service.markRetryOncePerTurn(model2); + + service.consumeStickyAttempt(model); + service.consumeStickyAttempt(model2); + + expect(service.snapshot(model).available).toBe(false); + expect(service.snapshot(model2).available).toBe(false); + + service.resetTurn(); + + expect(service.snapshot(model).available).toBe(true); + expect(service.snapshot(model2).available).toBe(true); + }); + + it('resetTurn does not affect terminal models', () => { + service.markTerminal(model, 'quota'); + service.resetTurn(); + expect(service.snapshot(model)).toEqual({ + available: false, + reason: 'quota', + }); + }); }); diff --git a/packages/core/src/availability/modelAvailabilityService.ts b/packages/core/src/availability/modelAvailabilityService.ts index f5e5192692c..0a08c28655c 100644 --- a/packages/core/src/availability/modelAvailabilityService.ts +++ b/packages/core/src/availability/modelAvailabilityService.ts @@ -6,21 +6,25 @@ export type ModelId = string; -export type TerminalAvailabilityReason = 'quota' | 'capacity'; -export type TurnAvailabilityReason = 'retry_once_per_turn'; +export type TerminalUnavailabilityReason = 'quota' | 'capacity'; +export type TurnUnavailabilityReason = 'retry_once_per_turn'; -export type AvailabilityReason = - | TerminalAvailabilityReason - | TurnAvailabilityReason +export type UnavailabilityReason = + | TerminalUnavailabilityReason + | TurnUnavailabilityReason | 'unknown'; type HealthState = - | { status: 'terminal'; reason: TerminalAvailabilityReason } - | { status: 'sticky_retry'; reason: TurnAvailabilityReason }; + | { status: 'terminal'; reason: TerminalUnavailabilityReason } + | { + status: 'sticky_retry'; + reason: TurnUnavailabilityReason; + consumed: boolean; + }; export interface ModelAvailabilitySnapshot { available: boolean; - reason?: AvailabilityReason; + reason?: UnavailabilityReason; } export interface ModelSelectionResult { @@ -28,15 +32,14 @@ export interface ModelSelectionResult { attempts?: number; skipped: Array<{ model: ModelId; - reason: AvailabilityReason; + reason: UnavailabilityReason; }>; } export class ModelAvailabilityService { private readonly health = new Map(); - private readonly consumedSticky = new Set(); - markTerminal(model: ModelId, reason: TerminalAvailabilityReason) { + markTerminal(model: ModelId, reason: TerminalUnavailabilityReason) { this.setState(model, { status: 'terminal', reason, @@ -56,31 +59,40 @@ export class ModelAvailabilityService { // Only reset consumption if we are not already in the sticky_retry state. // This prevents infinite loops if the model fails repeatedly in the same turn. - if (currentState?.status !== 'sticky_retry') { - this.consumedSticky.delete(model); + let consumed = false; + if (currentState?.status === 'sticky_retry') { + consumed = currentState.consumed; } this.setState(model, { status: 'sticky_retry', reason: 'retry_once_per_turn', + consumed, }); } consumeStickyAttempt(model: ModelId) { - if (this.health.get(model)?.status === 'sticky_retry') { - this.consumedSticky.add(model); + const state = this.health.get(model); + if (state?.status === 'sticky_retry') { + this.setState(model, { ...state, consumed: true }); } } snapshot(model: ModelId): ModelAvailabilitySnapshot { const state = this.health.get(model); - const isUnavailable = - state?.status === 'terminal' || - (state?.status === 'sticky_retry' && this.consumedSticky.has(model)); - if (isUnavailable) { - return { available: false, reason: state!.reason }; + if (!state) { + return { available: true }; + } + + if (state.status === 'terminal') { + return { available: false, reason: state.reason }; + } + + if (state.status === 'sticky_retry' && state.consumed) { + return { available: false, reason: state.reason }; } + return { available: true }; } @@ -102,7 +114,11 @@ export class ModelAvailabilityService { } resetTurn() { - this.consumedSticky.clear(); + for (const [model, state] of this.health.entries()) { + if (state.status === 'sticky_retry') { + this.setState(model, { ...state, consumed: false }); + } + } } private setState(model: ModelId, nextState: HealthState) { @@ -111,6 +127,5 @@ export class ModelAvailabilityService { private clearState(model: ModelId) { this.health.delete(model); - this.consumedSticky.delete(model); } }