packages/core/src/availability/modelAvailabilityService.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, expect, it, vi, beforeEach, afterEach } 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 retry with custom attempts', () => {
    service.markRetryOncePerTurn(model, 3);
    const selection = service.selectFirstAvailable([model]);
    expect(selection.attempts).toBe(3);
  });

  it('tracks terminal failures', () => {
    service.markTerminal(model, 'quota');
    expect(service.snapshot(model)).toEqual({
      available: false,
      reason: 'quota',
    });
  });

  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';

    service.markTerminal(model, 'capacity');
    service.markRetryOncePerTurn(stickyModel);

    const first = service.selectFirstAvailable([
      model,
      stickyModel,
      healthyModel,
    ]);
    expect(first).toEqual({
      selectedModel: stickyModel,
      attempts: 1,
      skipped: [
        {
          model,
          reason: 'capacity',
        },
      ],
    });

    service.consumeStickyAttempt(stickyModel);
    const second = service.selectFirstAvailable([
      model,
      stickyModel,
      healthyModel,
    ]);
    expect(second).toEqual({
      selectedModel: healthyModel,
      skipped: [
        {
          model,
          reason: 'capacity',
        },
        {
          model: stickyModel,
          reason: 'retry_once_per_turn',
        },
      ],
    });

    service.resetTurn();
    const third = service.selectFirstAvailable([
      model,
      stickyModel,
      healthyModel,
    ]);
    expect(third).toEqual({
      selectedModel: stickyModel,
      attempts: 1,
      skipped: [
        {
          model,
          reason: 'capacity',
        },
      ],
    });
  });

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

  describe('prefix normalization', () => {
    it('treats prefixed and non-prefixed models as identical when marking terminal', () => {
      service.markTerminal('models/gemini-3.1-pro-preview', 'quota');

      // Checking the non-prefixed version should show it as unavailable
      expect(service.snapshot('gemini-3.1-pro-preview')).toEqual({
        available: false,
        reason: 'quota',
      });

      // Checking the prefixed version should also show it as unavailable
      expect(service.snapshot('models/gemini-3.1-pro-preview')).toEqual({
        available: false,
        reason: 'quota',
      });
    });

    it('treats prefixed and non-prefixed models as identical when selecting', () => {
      service.markTerminal('gemini-3-flash-preview', 'quota');

      // Attempting to select the prefixed version should skip it because the base is exhausted
      const result = service.selectFirstAvailable([
        'models/gemini-3-flash-preview',
        'gemini-3.1-pro-preview',
      ]);

      expect(result.selectedModel).toBe('gemini-3.1-pro-preview');
      expect(result.skipped).toEqual([
        { model: 'gemini-3-flash-preview', reason: 'quota' },
      ]);
    });

    it('treats prefixed and non-prefixed models as identical when marking healthy', () => {
      service.markTerminal('gemini-3-flash-preview', 'quota');
      service.markHealthy('models/gemini-3-flash-preview');

      expect(service.snapshot('gemini-3-flash-preview')).toEqual({
        available: true,
      });
    });
  });

  describe('Capacity TTL expiration', () => {
    beforeEach(() => {
      vi.useFakeTimers();
    });

    afterEach(() => {
      vi.useRealTimers();
    });

    it('recovers terminal status for capacity after the TTL expires', () => {
      service.markTerminal(model, 'capacity');

      // Before TTL expiration, model is unavailable
      expect(service.snapshot(model)).toEqual({
        available: false,
        reason: 'capacity',
      });

      // Fast forward 29 seconds (default is 30s)
      vi.advanceTimersByTime(29000);
      expect(service.snapshot(model)).toEqual({
        available: false,
        reason: 'capacity',
      });

      // Fast forward past 30 seconds
      vi.advanceTimersByTime(1001);
      expect(service.snapshot(model)).toEqual({
        available: true,
      });
    });

    it('allows custom TTL parameter in snapshot', () => {
      service.markTerminal(model, 'capacity');

      // Check with a custom 5s TTL
      expect(service.snapshot(model, 5000)).toEqual({
        available: false,
        reason: 'capacity',
      });

      // Advance 6 seconds
      vi.advanceTimersByTime(6000);
      expect(service.snapshot(model, 5000)).toEqual({
        available: true,
      });
    });

    it('does not expire terminal status for non-capacity reasons (e.g. quota)', () => {
      service.markTerminal(model, 'quota');

      // Even after advance, quota remains terminal
      vi.advanceTimersByTime(60000);
      expect(service.snapshot(model)).toEqual({
        available: false,
        reason: 'quota',
      });
    });
  });
});