@tokenbuddy/tokenbuddy 1.0.9 → 1.0.11

package/src/prewarm-scheduler.test.ts

@@ -0,0 +1,367 @@
+import { ModelIndex } from "../src/model-index.js";
+import { PrewarmCache } from "../src/prewarm-cache.js";
+import { PrewarmScheduler, type ProbeResult, type SellerProber } from "../src/prewarm-scheduler.js";
+import type { RegistrySeller } from "../src/seller-catalog.js";
+
+interface FakeClock {
+  now: number;
+  advance: (ms: number) => void;
+}
+
+function makeClock(start = 1_000_000): FakeClock {
+  const clock = { now: start, advance: (ms: number) => { clock.now += ms; } };
+  return clock;
+}
+
+function makeSeller(overrides: Partial<RegistrySeller> & { id: string; models?: string[] }): RegistrySeller {
+  return {
+    id: overrides.id,
+    name: overrides.name ?? overrides.id,
+    url: overrides.url ?? `https://${overrides.id}.example.com`,
+    supportedProtocols: overrides.supportedProtocols ?? ["chat_completions"],
+    paymentMethods: overrides.paymentMethods ?? ["clawtip"],
+    models: overrides.models
+  };
+}
+
+function makeProberScript(script: Array<{ sellerId: string; ok?: boolean; latencyMs?: number; errorMessage?: string }>): SellerProber & { calls: string[] } {
+  const calls: string[] = [];
+  const remaining = script.slice();
+  const fn = (async (seller: RegistrySeller, _signal: AbortSignal): Promise<ProbeResult> => {
+    calls.push(seller.id);
+    const next = remaining.shift();
+    if (next && next.sellerId === seller.id) {
+      return {
+        ok: next.ok ?? true,
+        latencyMs: next.latencyMs ?? 100,
+        errorMessage: next.errorMessage,
+        httpStatus: next.ok === false ? 503 : 200
+      };
+    }
+    return { ok: true, latencyMs: 100, httpStatus: 200 };
+  }) as SellerProber & { calls: string[] };
+  fn.calls = calls;
+  return fn;
+}
+
+async function flushMicrotasks(times = 5): Promise<void> {
+  for (let i = 0; i < times; i += 1) {
+    await new Promise<void>((resolve) => setImmediate(resolve));
+  }
+}
+
+describe("PrewarmScheduler", () => {
+  test("warm task resolves with successful status when at least one seller probes ok", async () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["gpt-4o"] })]);
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([{ sellerId: "s1", ok: true, latencyMs: 200 }]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      // Disable idle loop so it does not race the test.
+      idleIntervalMs: 60_000,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const task = await scheduler.schedulePrewarm({ modelId: "gpt-4o", reason: "lazy" });
+    expect(task.status).toBe("succeeded");
+    expect(task.sellerIds).toEqual(["s1"]);
+    expect(prober.calls).toEqual(["s1"]);
+
+    const entry = cache.get("gpt-4o", "chat_completions", "clawtip");
+    expect(entry?.state).toBe("warm");
+    expect(entry?.candidates).toHaveLength(1);
+    expect(entry?.candidates[0].sellerId).toBe("s1");
+    expect(entry?.candidates[0].healthScore).toBeGreaterThan(0);
+  });
+
+  test("all-failed probe marks the cache entry stale and the task failed", async () => {
+    const index = new ModelIndex();
+    index.rebuild([makeSeller({ id: "s1", models: ["gpt-4o"] })]);
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([{ sellerId: "s1", ok: false, errorMessage: "503" }]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const task = await scheduler.schedulePrewarm({ modelId: "gpt-4o", reason: "lazy" });
+    expect(task.status).toBe("failed");
+    expect(task.errorMessage).toBe("all probes failed");
+
+    const entry = cache.get("gpt-4o", "chat_completions", "clawtip");
+    expect(entry?.state).toBe("stale");
+    expect(entry?.consecutiveWarmingFailures).toBe(1);
+  });
+
+  test("no matching sellers in the index marks the task failed and skips probing", async () => {
+    const index = new ModelIndex();
+    index.rebuild([]); // empty
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const task = await scheduler.schedulePrewarm({ modelId: "unknown", reason: "lazy" });
+    expect(task.status).toBe("failed");
+    expect(task.errorMessage).toBe("no sellers for model");
+    expect(prober.calls).toEqual([]);
+  });
+
+  test("concurrency caps in-flight probes to the configured value", async () => {
+    // Each task probes its sellers serially; concurrency is the cap on the
+    // number of *tasks* running in parallel. To exercise the cap we enqueue
+    // three independent (model, seller) pairs and verify the prober is
+    // never invoked more than `concurrency` times at once.
+    const sellers = [
+      makeSeller({ id: "s1", models: ["m1"] }),
+      makeSeller({ id: "s2", models: ["m2"] }),
+      makeSeller({ id: "s3", models: ["m3"] })
+    ];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+
+    let concurrent = 0;
+    let peak = 0;
+    const prober: SellerProber = async (_seller, _signal) => {
+      concurrent += 1;
+      peak = Math.max(peak, concurrent);
+      await new Promise<void>((resolve) => setTimeout(resolve, 20));
+      concurrent -= 1;
+      return { ok: true, latencyMs: 50, httpStatus: 200 };
+    };
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      concurrency: 2,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const [t1, t2, t3] = await Promise.all([
+      scheduler.schedulePrewarm({ modelId: "m1", reason: "lazy" }),
+      scheduler.schedulePrewarm({ modelId: "m2", reason: "lazy" }),
+      scheduler.schedulePrewarm({ modelId: "m3", reason: "lazy" })
+    ]);
+    expect([t1, t2, t3].map((t) => t.status)).toEqual(["succeeded", "succeeded", "succeeded"]);
+    expect(peak).toBeLessThanOrEqual(2);
+    expect(peak).toBe(2);
+  });
+
+  test("per-seller rate limit suppresses repeated probes within the minimum interval", async () => {
+    const sellers = [makeSeller({ id: "s1", models: ["m"] })];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+    const clock = makeClock();
+    const prober = makeProberScript([{ sellerId: "s1", ok: true, latencyMs: 50 }]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      perSellerMinIntervalMs: 30_000,
+      now: () => clock.now,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    // First probe at t=1_000_000 succeeds and updates the rate-limit ledger.
+    const first = await scheduler.schedulePrewarm({ modelId: "m", reason: "lazy" });
+    expect(first.status).toBe("succeeded");
+
+    // Second probe 5s later: seller is rate-limited, no new probe call.
+    clock.advance(5_000);
+    const second = await scheduler.schedulePrewarm({ modelId: "m", reason: "lazy" });
+    expect(second.status).toBe("succeeded"); // task itself still resolves
+    expect(prober.calls).toEqual(["s1"]); // prober was NOT called again
+
+    // After 30s have elapsed since the last probe, the seller can be probed again.
+    clock.advance(30_000);
+    const third = await scheduler.schedulePrewarm({ modelId: "m", reason: "lazy" });
+    expect(third.status).toBe("succeeded");
+    expect(prober.calls).toEqual(["s1", "s1"]);
+  });
+
+  test("global per-minute probe budget rate-limits excess tasks", async () => {
+    const sellers = [
+      makeSeller({ id: "s1", models: ["m1"] }),
+      makeSeller({ id: "s2", models: ["m2"] }),
+      makeSeller({ id: "s3", models: ["m3"] })
+    ];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([
+      { sellerId: "s1", ok: true, latencyMs: 10 },
+      { sellerId: "s2", ok: true, latencyMs: 10 },
+      { sellerId: "s3", ok: true, latencyMs: 10 }
+    ]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      maxPrewarmPerMinute: 2,
+      // Generous per-seller window so it does not interfere.
+      perSellerMinIntervalMs: 0,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const t1 = await scheduler.schedulePrewarm({ modelId: "m1", reason: "lazy" });
+    const t2 = await scheduler.schedulePrewarm({ modelId: "m2", reason: "lazy" });
+    const t3 = await scheduler.schedulePrewarm({ modelId: "m3", reason: "lazy" });
+    expect(t1.status).toBe("succeeded");
+    expect(t2.status).toBe("succeeded");
+    expect(t3.status).toBe("rate_limited");
+    expect(prober.calls).toEqual(["s1", "s2"]);
+    const stats = scheduler.stats();
+    expect(stats.totalRateLimited).toBe(1);
+  });
+
+  test("tickIdle enqueues prewarms only for entries that are expiring soon", async () => {
+    const index = new ModelIndex();
+    index.rebuild([
+      makeSeller({ id: "s1", models: ["m1"] }),
+      makeSeller({ id: "s2", models: ["m2"] })
+    ]);
+    const cache = new PrewarmCache({ defaultTtlMs: 1000 });
+    const prober = makeProberScript([
+      { sellerId: "s1", ok: true, latencyMs: 50 },
+      { sellerId: "s2", ok: true, latencyMs: 50 }
+    ]);
+    const clock = makeClock();
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      perSellerMinIntervalMs: 0,
+      now: () => clock.now,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    // Seed cache with two entries.
+    await scheduler.schedulePrewarm({ modelId: "m1", reason: "startup" });
+    await scheduler.schedulePrewarm({ modelId: "m2", reason: "startup" });
+    expect(cache.size()).toBe(2);
+
+    // Advance to t=950 (within 10% of 1000 TTL).
+    clock.advance(950);
+    const enqueued = scheduler.tickIdle();
+    expect(enqueued).toBe(2);
+    await flushMicrotasks();
+  });
+
+  test("stop() cancels queued tasks and prevents further dispatch", async () => {
+    const sellers = [makeSeller({ id: "s1", models: ["m"] })];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+    // Prober that observes its abort signal and rejects on abort. This is
+    // the contract real probers (e.g. `health-probe.ts`) must follow.
+    let proberStarted = false;
+    const prober: SellerProber = async (_seller, signal) => {
+      proberStarted = true;
+      await new Promise<void>((resolve, reject) => {
+        if (signal.aborted) {
+          reject(new Error("aborted"));
+          return;
+        }
+        signal.addEventListener("abort", () => reject(new Error("aborted")), { once: true });
+      });
+      return { ok: true, latencyMs: 10, httpStatus: 200 };
+    };
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      concurrency: 1,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    const task = scheduler.schedulePrewarm({ modelId: "m", reason: "lazy" });
+    // Let the dispatcher pick up the task and start the probe.
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    expect(proberStarted).toBe(true);
+
+    await scheduler.stop();
+    const result = await task;
+    expect(result.status).toBe("canceled");
+    expect(scheduler.stats().inFlight).toBe(0);
+  });
+
+  test("runStartupPrewarm honors startup jitter and processes every model", async () => {
+    const sellers = [makeSeller({ id: "s1", models: ["m1"] }), makeSeller({ id: "s2", models: ["m2"] })];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([
+      { sellerId: "s1", ok: true, latencyMs: 50 },
+      { sellerId: "s2", ok: true, latencyMs: 50 }
+    ]);
+    let sleepCalls = 0;
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      startupJitterMinMs: 100,
+      startupJitterMaxMs: 100,
+      sleep: async () => { sleepCalls += 1; },
+      perSellerMinIntervalMs: 0
+    });
+
+    await scheduler.runStartupPrewarm(["m1", "m2"]);
+    expect(sleepCalls).toBe(1); // single jitter wait before the sweep
+    expect(prober.calls.sort()).toEqual(["s1", "s2"]);
+  });
+
+  test("stats reports queue depth, in-flight, and counters", async () => {
+    const sellers = [makeSeller({ id: "s1", models: ["m"] })];
+    const index = new ModelIndex();
+    index.rebuild(sellers);
+    const cache = new PrewarmCache();
+    const prober = makeProberScript([{ sellerId: "s1", ok: true, latencyMs: 50 }]);
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober,
+      sleep: () => new Promise(() => undefined)
+    });
+
+    expect(scheduler.stats()).toMatchObject({
+      queueDepth: 0,
+      inFlight: 0,
+      totalScheduled: 0,
+      totalSucceeded: 0,
+      totalFailed: 0,
+      totalRateLimited: 0,
+      concurrency: 4,
+      maxPrewarmPerMinute: 30
+    });
+
+    await scheduler.schedulePrewarm({ modelId: "m", reason: "lazy" });
+    const stats = scheduler.stats();
+    expect(stats.totalSucceeded).toBe(1);
+    expect(stats.totalScheduled).toBe(1);
+  });
+
+  test("default options match the v1.2 design defaults", () => {
+    const index = new ModelIndex();
+    const cache = new PrewarmCache();
+    const scheduler = new PrewarmScheduler({
+      modelIndex: index,
+      cache,
+      prober: async () => ({ ok: true, latencyMs: 1, httpStatus: 200 })
+    });
+    const stats = scheduler.stats();
+    expect(stats.concurrency).toBe(4);
+    expect(stats.maxPrewarmPerMinute).toBe(30);
+  });
+});