qmd/test/embedding-autofallback.test.ts

/**
 * embedding-autofallback.test.ts - Tests for AutoFallbackEmbeddingProvider.
 */

import { describe, test, expect } from "vitest";
import {
  AutoFallbackEmbeddingProvider,
  type AutoFallbackProviderConfig,
} from "../src/embedding/autofallback.js";
import { CircuitOpenError } from "../src/embedding/openai.js";
import type {
  EmbeddingProvider,
  ProviderEmbedOptions,
  ProviderEmbedding,
  ProviderHealth,
  ProviderKind,
} from "../src/embedding/provider.js";

// ─────────────────────────── Test fakes ──────────────────────────────────────

class FakeProvider implements EmbeddingProvider {
  readonly kind: ProviderKind;
  readonly modelId: string;
  readonly dim: number;
  embedCalls = 0;
  embedBatchCalls = 0;
  healthcheckCalls = 0;
  disposed = false;
  /** Override behavior for next N calls */
  nextThrows: Array<Error | null> = [];
  /** Always-throw mode */
  alwaysThrows: Error | null = null;
  /** Health response */
  healthResponse: ProviderHealth | null = null;

  constructor(kind: ProviderKind, modelId: string, dim = 4) {
    this.kind = kind;
    this.modelId = modelId;
    this.dim = dim;
  }

  getModelId(): string {
    return this.modelId;
  }
  getDimensions(): number | undefined {
    return this.dim;
  }

  async healthcheck(): Promise<ProviderHealth> {
    this.healthcheckCalls++;
    if (this.healthResponse) return this.healthResponse;
    return { ok: true, model: this.modelId, dimensions: this.dim };
  }

  async embed(text: string, _options?: ProviderEmbedOptions): Promise<ProviderEmbedding | null> {
    this.embedCalls++;
    this.maybeThrow();
    return { embedding: this.fakeEmbed(text), model: this.modelId };
  }

  async embedBatch(texts: string[], _options?: ProviderEmbedOptions): Promise<(ProviderEmbedding | null)[]> {
    this.embedBatchCalls++;
    this.maybeThrow();
    return texts.map((t) => ({ embedding: this.fakeEmbed(t), model: this.modelId }));
  }

  async dispose(): Promise<void> {
    this.disposed = true;
  }

  private maybeThrow(): void {
    if (this.alwaysThrows) throw this.alwaysThrows;
    const next = this.nextThrows.shift();
    if (next) throw next;
  }

  private fakeEmbed(text: string): number[] {
    return Array.from({ length: this.dim }, (_, i) => (text.length + i) * 0.01);
  }
}

function buildAutoFallback(opts: Partial<AutoFallbackProviderConfig> = {}): {
  af: AutoFallbackEmbeddingProvider;
  primary: FakeProvider;
  fallback: FakeProvider;
  warns: string[];
  setNow: (n: number) => void;
} {
  const primary = new FakeProvider("openai", "embeddinggemma");
  const fallback = new FakeProvider("local", "embeddinggemma");
  const warns: string[] = [];
  let now = 1_000_000;
  const af = new AutoFallbackEmbeddingProvider({
    primary,
    fallback,
    failureStreakThreshold: opts.failureStreakThreshold ?? 3,
    cooldownMs: opts.cooldownMs ?? 60_000,
    warn: (m) => warns.push(m),
    now: () => now,
    ...opts,
  });
  return { af, primary, fallback, warns, setNow: (n) => (now = n) };
}

// ─────────────────────────── Construction ────────────────────────────────────

describe("AutoFallbackEmbeddingProvider — construction", () => {
  test("requires primary", () => {
    expect(
      () =>
        new AutoFallbackEmbeddingProvider({
          // @ts-expect-error testing runtime guard
          primary: undefined,
          fallback: new FakeProvider("local", "x"),
        }),
    ).toThrow(/primary is required/);
  });

  test("requires fallback", () => {
    expect(
      () =>
        new AutoFallbackEmbeddingProvider({
          primary: new FakeProvider("openai", "x"),
          // @ts-expect-error testing runtime guard
          fallback: undefined,
        }),
    ).toThrow(/fallback is required/);
  });

  test("rejects identical primary and fallback", () => {
    const same = new FakeProvider("openai", "x");
    expect(
      () =>
        new AutoFallbackEmbeddingProvider({
          primary: same,
          fallback: same,
        }),
    ).toThrow(/must differ/);
  });

  test("inherits primary's kind", () => {
    const { af } = buildAutoFallback();
    expect(af.kind).toBe("openai");
  });
});

// ─────────────────────────── Happy path ──────────────────────────────────────

describe("AutoFallbackEmbeddingProvider — happy path", () => {
  test("primary succeeds → fallback never called", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    const r = await af.embed("hello");
    expect(r).not.toBeNull();
    expect(primary.embedCalls).toBe(1);
    expect(fallback.embedCalls).toBe(0);
    expect(af.getRoutingState()).toBe("primary");
  });

  test("primary embedBatch succeeds → fallback untouched", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    const out = await af.embedBatch(["a", "b"]);
    expect(out.length).toBe(2);
    expect(primary.embedBatchCalls).toBe(1);
    expect(fallback.embedBatchCalls).toBe(0);
  });

  test("getModelId / getDimensions delegate to primary", () => {
    const { af, primary } = buildAutoFallback();
    expect(af.getModelId()).toBe(primary.getModelId());
    expect(af.getDimensions()).toBe(primary.getDimensions());
  });
});

// ─────────────────────────── Circuit-open fallback ───────────────────────────

describe("AutoFallbackEmbeddingProvider — CircuitOpenError handling", () => {
  test("primary throws CircuitOpenError → fallback served + cooldown opens", async () => {
    const { af, primary, fallback, warns } = buildAutoFallback();
    primary.nextThrows.push(new CircuitOpenError());
    const r = await af.embed("hello");
    expect(r).not.toBeNull();
    expect(r!.embedding.length).toBe(4); // came from fallback
    expect(primary.embedCalls).toBe(1);
    expect(fallback.embedCalls).toBe(1);
    expect(af.getRoutingState()).toBe("fallback");
    expect(warns.some((w) => w.includes("CircuitOpenError"))).toBe(true);
  });

  test("during cooldown subsequent calls skip primary entirely", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    primary.nextThrows.push(new CircuitOpenError());
    await af.embed("first");
    expect(primary.embedCalls).toBe(1);
    expect(fallback.embedCalls).toBe(1);
    // Subsequent call within cooldown
    await af.embed("second");
    expect(primary.embedCalls).toBe(1); // unchanged
    expect(fallback.embedCalls).toBe(2);
  });

  test("after cooldown expires, primary is retried", async () => {
    const { af, primary, fallback, setNow } = buildAutoFallback({ cooldownMs: 5000 });
    primary.nextThrows.push(new CircuitOpenError());
    await af.embed("a");
    expect(af.getRoutingState()).toBe("fallback");

    setNow(1_000_000 + 5_001);
    expect(af.getRoutingState()).toBe("primary");

    // Next call reaches primary again
    await af.embed("b");
    expect(primary.embedCalls).toBe(2);
    expect(fallback.embedCalls).toBe(1);
  });

  test("WARN fired only once per transition (not per call during cooldown)", async () => {
    const { af, primary, warns } = buildAutoFallback();
    primary.nextThrows.push(new CircuitOpenError());
    await af.embed("a");
    await af.embed("b");
    await af.embed("c");
    const fallbackWarns = warns.filter((w) => w.includes("falling back"));
    expect(fallbackWarns.length).toBe(1);
  });
});

// ─────────────────────────── Failure-streak threshold ────────────────────────

describe("AutoFallbackEmbeddingProvider — failure streak", () => {
  test("non-CircuitOpen errors below threshold → no cooldown", async () => {
    const { af, primary, fallback } = buildAutoFallback({ failureStreakThreshold: 3 });
    primary.nextThrows.push(new Error("transient"));
    const r = await af.embed("a");
    expect(r).not.toBeNull(); // fallback served it
    expect(af.getRoutingState()).toBe("primary");
    expect(primary.embedCalls).toBe(1);
    expect(fallback.embedCalls).toBe(1);
  });

  test("threshold consecutive failures → cooldown opens", async () => {
    const { af, primary, fallback } = buildAutoFallback({ failureStreakThreshold: 3 });
    for (let i = 0; i < 3; i++) {
      primary.nextThrows.push(new Error(`err ${i}`));
    }
    await af.embed("a");
    await af.embed("b");
    await af.embed("c");
    expect(af.getRoutingState()).toBe("fallback");
    expect(primary.embedCalls).toBe(3);
    expect(fallback.embedCalls).toBe(3);
  });

  test("a single primary success resets the streak", async () => {
    const { af, primary } = buildAutoFallback({ failureStreakThreshold: 3 });
    primary.nextThrows.push(new Error("e1"));
    primary.nextThrows.push(new Error("e2"));
    await af.embed("a");
    await af.embed("b");
    // Now success
    await af.embed("c");
    // Streak reset; another two failures shouldn't trip cooldown yet
    primary.nextThrows.push(new Error("e3"));
    primary.nextThrows.push(new Error("e4"));
    await af.embed("d");
    await af.embed("e");
    expect(af.getRoutingState()).toBe("primary");
  });
});

// ─────────────────────────── Recovery transition ─────────────────────────────

describe("AutoFallbackEmbeddingProvider — recovery transitions", () => {
  test("recovery WARN fires when primary call succeeds after fallback", async () => {
    const { af, primary, warns, setNow } = buildAutoFallback({ cooldownMs: 5000 });
    primary.nextThrows.push(new CircuitOpenError());
    await af.embed("a");
    setNow(1_000_000 + 5_001);
    await af.embed("b"); // primary succeeds
    const recoveryWarns = warns.filter((w) => w.includes("recovered"));
    expect(recoveryWarns.length).toBe(1);
  });

  test("reset() clears state + transitions back to primary", async () => {
    const { af, primary } = buildAutoFallback({ cooldownMs: 60_000 });
    primary.nextThrows.push(new CircuitOpenError());
    await af.embed("a");
    expect(af.getRoutingState()).toBe("fallback");
    af.reset();
    expect(af.getRoutingState()).toBe("primary");
    await af.embed("b");
    expect(primary.embedCalls).toBe(2);
  });
});

// ─────────────────────────── Both fail ───────────────────────────────────────

describe("AutoFallbackEmbeddingProvider — both providers fail", () => {
  test("primary throws + fallback throws → embedBatch returns nulls", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    primary.alwaysThrows = new Error("primary down");
    fallback.alwaysThrows = new Error("local broken");
    const r = await af.embedBatch(["a", "b"]);
    expect(r).toEqual([null, null]);
  });

  test("primary throws + fallback throws → embed propagates fallback error", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    primary.alwaysThrows = new Error("primary down");
    fallback.alwaysThrows = new Error("local broken");
    await expect(af.embed("a")).rejects.toThrow(/local broken/);
  });
});

// ─────────────────────────── Healthcheck ─────────────────────────────────────

describe("AutoFallbackEmbeddingProvider — healthcheck", () => {
  test("primary healthy → returns primary health", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    const h = await af.healthcheck();
    expect(h.ok).toBe(true);
    expect(primary.healthcheckCalls).toBe(1);
    expect(fallback.healthcheckCalls).toBe(0);
  });

  test("primary unhealthy → fallback checked + reported", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    primary.healthResponse = { ok: false, model: "primary-model", detail: "down" };
    fallback.healthResponse = { ok: true, model: "local-model", detail: "fine" };
    const h = await af.healthcheck();
    expect(h.ok).toBe(true);
    expect(primary.healthcheckCalls).toBe(1);
    expect(fallback.healthcheckCalls).toBe(1);
    expect(h.detail).toContain("primary");
    expect(h.detail).toContain("fallback");
  });

  test("both unhealthy → ok=false", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    primary.healthResponse = { ok: false, model: "p", detail: "down" };
    fallback.healthResponse = { ok: false, model: "f", detail: "down" };
    const h = await af.healthcheck();
    expect(h.ok).toBe(false);
  });
});

// ─────────────────────────── dispose ─────────────────────────────────────────

describe("AutoFallbackEmbeddingProvider — dispose", () => {
  test("dispose cascades to both providers", async () => {
    const { af, primary, fallback } = buildAutoFallback();
    await af.dispose();
    expect(primary.disposed).toBe(true);
    expect(fallback.disposed).toBe(true);
  });
});