@tokenbuddy/tokenbuddy 1.0.9 → 1.0.11

package/src/prewarm-scheduler.ts

@@ -0,0 +1,581 @@
+import { createModuleLogger } from "@tokenbuddy/logging";
+import type { RegistrySeller } from "./seller-catalog.js";
+import type { ModelIndex } from "./model-index.js";
+import type { PrewarmCache, PrewarmCandidate } from "./prewarm-cache.js";
+
+const logger = createModuleLogger("tb-proxyd:prewarm-scheduler");
+
+export type PrewarmReason = "startup" | "lazy" | "idle" | "explicit";
+
+export interface ProbeResult {
+  ok: boolean;
+  latencyMs: number;
+  httpStatus?: number;
+  errorMessage?: string;
+}
+
+/**
+ * The probe function used by the scheduler. Decoupled so the scheduler can
+ * be unit-tested without spinning up HTTP servers. The default
+ * implementation in `health-probe.ts` (PR-2/PR-3) calls
+ * `GET <seller.url>/healthz` with a 3s `AbortSignal.timeout`. Probers must
+ * observe the provided `AbortSignal` and reject when it aborts so the
+ * scheduler can short-circuit in-flight probes on `stop()`.
+ */
+export type SellerProber = (seller: RegistrySeller, signal: AbortSignal) => Promise<ProbeResult>;
+
+export interface PrewarmSchedulerOptions {
+  modelIndex: ModelIndex;
+  cache: PrewarmCache;
+  prober: SellerProber;
+  // Limits (defaults match buyer-driven-fallback-design.md §18.6).
+  concurrency?: number;
+  perSellerMinIntervalMs?: number;
+  maxPrewarmPerMinute?: number;
+  // Idle loop cadence; the scheduler can also be driven externally
+  // (PR-2.1 wires `tickIdle` into the existing registry-loop heartbeat).
+  idleIntervalMs?: number;
+  // Startup jitter (5-10s by default per §18.5.1).
+  startupJitterMinMs?: number;
+  startupJitterMaxMs?: number;
+  // Hooks for testing; defaults to Node's setTimeout / setImmediate.
+  sleep?: (ms: number, signal?: AbortSignal) => Promise<void>;
+  random?: () => number;
+  now?: () => number;
+  // Optional filter applied to every probe (e.g. preferred protocol).
+  protocol?: string;
+  paymentMethod?: string;
+}
+
+interface PrewarmTask {
+  id: number;
+  modelId: string;
+  reason: PrewarmReason;
+  protocol: string;
+  paymentMethod: string;
+  enqueuedAt: number;
+  sellerIds: string[];
+  startedAt?: number;
+  completedAt?: number;
+  status: "queued" | "running" | "succeeded" | "failed" | "canceled" | "rate_limited";
+  errorMessage?: string;
+}
+
+export interface PrewarmSchedulerStats {
+  queueDepth: number;
+  inFlight: number;
+  totalScheduled: number;
+  totalSucceeded: number;
+  totalFailed: number;
+  totalRateLimited: number;
+  recentProbesInLastMinute: number;
+  concurrency: number;
+  maxPrewarmPerMinute: number;
+}
+
+/**
+ * Background scheduler that warms up sellers for a (model, protocol,
+ * payment) triple on demand. The scheduler owns:
+ *   - queue management with bounded concurrency (default 4)
+ *   - per-seller rate limiting (default 30s between probes to the same
+ *     seller, even across different models)
+ *   - global rate limiting (default 30 probes/minute)
+ *   - jitter on startup and between probes to avoid thundering herds
+ *
+ * The scheduler does NOT own HTTP I/O; that lives in the injected
+ * `prober` so tests can swap in a deterministic stub.
+ */
+export class PrewarmScheduler {
+  private readonly modelIndex: ModelIndex;
+  private readonly cache: PrewarmCache;
+  private readonly prober: SellerProber;
+
+  private readonly concurrency: number;
+  private readonly perSellerMinIntervalMs: number;
+  private readonly maxPrewarmPerMinute: number;
+  private readonly idleIntervalMs: number;
+  private readonly startupJitterMinMs: number;
+  private readonly startupJitterMaxMs: number;
+  private readonly sleep: (ms: number, signal?: AbortSignal) => Promise<void>;
+  private readonly random: () => number;
+  private readonly now: () => number;
+  private readonly protocol: string | undefined;
+  private readonly paymentMethod: string | undefined;
+
+  private readonly queue: PrewarmTask[] = [];
+  private inFlight = 0;
+  private recentProbes: number[] = [];
+  private lastProbeAtBySeller = new Map<string, number>();
+  private nextTaskId = 1;
+
+  private totalScheduled = 0;
+  private totalSucceeded = 0;
+  private totalFailed = 0;
+  private totalRateLimited = 0;
+
+  private abortController: AbortController | null = null;
+  private idleLoopPromise: Promise<void> | null = null;
+
+  constructor(options: PrewarmSchedulerOptions) {
+    this.modelIndex = options.modelIndex;
+    this.cache = options.cache;
+    this.prober = options.prober;
+    this.concurrency = options.concurrency ?? 4;
+    this.perSellerMinIntervalMs = options.perSellerMinIntervalMs ?? 30_000;
+    this.maxPrewarmPerMinute = options.maxPrewarmPerMinute ?? 30;
+    this.idleIntervalMs = options.idleIntervalMs ?? 60_000;
+    this.startupJitterMinMs = options.startupJitterMinMs ?? 5_000;
+    this.startupJitterMaxMs = options.startupJitterMaxMs ?? 10_000;
+    this.sleep = options.sleep ?? defaultSleep;
+    this.random = options.random ?? Math.random;
+    this.now = options.now ?? Date.now;
+    this.protocol = options.protocol;
+    this.paymentMethod = options.paymentMethod;
+  }
+
+  /**
+   * Start the background idle loop. Safe to call once per scheduler
+   * instance; subsequent calls are no-ops. The idle loop probes any cached
+   * entry whose TTL is within 10% of expiry (`isExpiringSoon`).
+   */
+  start(): void {
+    if (this.abortController) {
+      return;
+    }
+    this.abortController = new AbortController();
+    this.idleLoopPromise = this.runIdleLoop(this.abortController.signal);
+  }
+
+  /**
+   * Cancel the idle loop and any pending tasks. Existing `inFlight` probes
+   * are not aborted (the prober owns its own timeout) but will not be
+   * dispatched to the cache.
+   */
+  async stop(): Promise<void> {
+    if (!this.abortController) {
+      return;
+    }
+    this.abortController.abort();
+    this.abortController = null;
+    if (this.idleLoopPromise) {
+      await this.idleLoopPromise.catch(() => undefined);
+      this.idleLoopPromise = null;
+    }
+    // Mark queued tasks as canceled so callers awaiting them can short-circuit.
+    for (const task of this.queue) {
+      task.status = "canceled";
+    }
+    this.queue.length = 0;
+  }
+
+  /**
+   * Enqueue a prewarm for a (model, protocol, payment) triple. The
+   * `reason` controls how aggressively the scheduler resolves candidates
+   * (e.g. `startup` defers; `lazy` waits on the returned promise). The
+   * returned promise resolves with the final task status once the queue
+   * drains or the scheduler is stopped.
+   */
+  schedulePrewarm(input: {
+    modelId: string;
+    reason: PrewarmReason;
+    protocol?: string;
+    paymentMethod?: string;
+    blockOnFirst?: boolean;
+  }): Promise<PrewarmTask> {
+    const protocol = input.protocol ?? this.protocol ?? "chat_completions";
+    const paymentMethod = input.paymentMethod ?? this.paymentMethod ?? "clawtip";
+    const task: PrewarmTask = {
+      id: this.nextTaskId++,
+      modelId: input.modelId,
+      reason: input.reason,
+      protocol,
+      paymentMethod,
+      enqueuedAt: this.now(),
+      sellerIds: [],
+      status: "queued"
+    };
+    this.queue.push(task);
+    this.totalScheduled += 1;
+    logger.info("prewarm.scheduled", "prewarm task enqueued", {
+      taskId: task.id,
+      modelId: task.modelId,
+      reason: task.reason,
+      protocol,
+      paymentMethod,
+      queueDepth: this.queue.length
+    });
+
+    // Fire-and-forget dispatch; the awaiter observes `task.status` via
+    // `taskResolved` rather than blocking the queue.
+    this.dispatch().catch((err) => {
+      logger.error("prewarm.dispatch.unexpected", "dispatcher threw unexpectedly", {
+        errorMessage: err instanceof Error ? err.message : String(err)
+      });
+    });
+
+    return new Promise<PrewarmTask>((resolve) => {
+      const check = () => {
+        if (task.status === "succeeded" || task.status === "failed" || task.status === "canceled" || task.status === "rate_limited") {
+          resolve(task);
+        } else {
+          setImmediate(check);
+        }
+      };
+      check();
+    });
+  }
+
+  /**
+   * Run a one-shot sweep that probes every focus-set model. Used by the
+   * `tb doctor --prewarm` explicit trigger and by the startup hook after
+   * the configured jitter window. Resolves once every scheduled task has
+   * reached a terminal state.
+   */
+  async runStartupPrewarm(modelIds: string[]): Promise<void> {
+    await this.sleep(this.jitterMs(), this.abortController?.signal);
+    if (this.abortController?.signal.aborted) {
+      return;
+    }
+    const tasks = modelIds.map((modelId) =>
+      this.schedulePrewarm({ modelId, reason: "startup" })
+    );
+    await Promise.all(tasks);
+  }
+
+  /**
+   * Force a sweep of any cache key whose TTL is about to expire. Returns
+   * the number of tasks that were enqueued. Intended to be called from
+   * the registry loop's heartbeat (replaces the v1 "all sellers" probe
+   * cycle with "only the ones we are about to forget").
+   */
+  tickIdle(): number {
+    const expiring = this.cache
+      .snapshot()
+      .filter((entry) => this.cache.isExpiringSoon(entry.modelId, entry.protocol, entry.paymentMethod, 60_000));
+    if (expiring.length === 0) {
+      return 0;
+    }
+    for (const entry of expiring) {
+      this.schedulePrewarm({
+        modelId: entry.modelId,
+        protocol: entry.protocol,
+        paymentMethod: entry.paymentMethod,
+        reason: "idle"
+      });
+    }
+    return expiring.length;
+  }
+
+  stats(): PrewarmSchedulerStats {
+    const now = this.now();
+    const cutoff = now - 60_000;
+    const recentProbesInLastMinute = this.recentProbes.filter((ts) => ts >= cutoff).length;
+    return {
+      queueDepth: this.queue.length,
+      inFlight: this.inFlight,
+      totalScheduled: this.totalScheduled,
+      totalSucceeded: this.totalSucceeded,
+      totalFailed: this.totalFailed,
+      totalRateLimited: this.totalRateLimited,
+      recentProbesInLastMinute,
+      concurrency: this.concurrency,
+      maxPrewarmPerMinute: this.maxPrewarmPerMinute
+    };
+  }
+
+  private jitterMs(): number {
+    const span = Math.max(0, this.startupJitterMaxMs - this.startupJitterMinMs);
+    return this.startupJitterMinMs + Math.floor(this.random() * span);
+  }
+
+  private async runIdleLoop(signal: AbortSignal): Promise<void> {
+    while (!signal.aborted) {
+      try {
+        await this.sleep(this.idleIntervalMs, signal);
+      } catch {
+        return;
+      }
+      if (signal.aborted) {
+        return;
+      }
+      try {
+        this.tickIdle();
+      } catch (err) {
+        logger.error("prewarm.idle.failed", "idle tick threw unexpectedly", {
+          errorMessage: err instanceof Error ? err.message : String(err)
+        });
+      }
+    }
+  }
+
+  private async dispatch(): Promise<void> {
+    while (this.queue.length > 0 && this.inFlight < this.concurrency) {
+      const task = this.queue.shift();
+      if (!task) {
+        break;
+      }
+      if (task.status === "canceled") {
+        continue;
+      }
+      if (this.isOverBudget()) {
+        task.status = "rate_limited";
+        this.totalRateLimited += 1;
+        logger.warn("prewarm.rate_limited", "global per-minute probe budget exhausted", {
+          taskId: task.id,
+          modelId: task.modelId,
+          recentProbes: this.recentProbesInLastMinute()
+        });
+        continue;
+      }
+      const sellers = this.modelIndex.sellersFor(task.modelId, {
+        protocol: task.protocol,
+        paymentMethod: task.paymentMethod
+      });
+      task.sellerIds = sellers.map((s) => s.id);
+      if (sellers.length === 0) {
+        task.status = "failed";
+        task.errorMessage = "no sellers for model";
+        task.completedAt = this.now();
+        this.totalFailed += 1;
+        logger.warn("prewarm.no_sellers", "no registry sellers match model", {
+          taskId: task.id,
+          modelId: task.modelId,
+          protocol: task.protocol,
+          paymentMethod: task.paymentMethod
+        });
+        continue;
+      }
+      // Ensure an abort controller exists so `stop()` works even when the
+      // caller never invoked `start()`. `start()` is otherwise responsible
+      // for the idle loop; dispatch only borrows the controller for
+      // short-lived abort propagation.
+      if (!this.abortController) {
+        this.abortController = new AbortController();
+      }
+      this.inFlight += 1;
+      task.status = "running";
+      task.startedAt = this.now();
+      // Capture the abort signal so an in-flight task can still observe
+      // `stop()` even after the controller reference is cleared.
+      const signal = this.abortController.signal;
+      // Run async without awaiting; the loop continues to dispatch.
+      void this.runTask(task, sellers, signal);
+    }
+  }
+
+  private async runTask(task: PrewarmTask, sellers: RegistrySeller[], signal: AbortSignal | undefined): Promise<void> {
+    const begin = this.cache.beginWarming(task.modelId, task.protocol, task.paymentMethod);
+    const candidates: PrewarmCandidate[] = [];
+    let anyOk = false;
+    let probedAny = false;
+
+    for (const seller of sellers) {
+      if (signal?.aborted) {
+        task.status = "canceled";
+        task.completedAt = this.now();
+        this.inFlight -= 1;
+        this.dispatch().catch(() => undefined);
+        return;
+      }
+      if (this.isSellerRateLimited(seller.id)) {
+        logger.debug("prewarm.seller_skipped", "seller probe skipped due to per-seller rate limit", {
+          taskId: task.id,
+          sellerId: seller.id
+        });
+        continue;
+      }
+      probedAny = true;
+      this.recordProbeAttempt();
+      const probeSignal = composeProbeSignal(signal);
+      let result: ProbeResult;
+      try {
+        result = await this.prober(seller, probeSignal.signal);
+      } catch (err) {
+        // Prober rejected (typically because of `stop()` aborting the probe
+        // signal). Treat the rejection as a canceled run and exit early.
+        if (signal?.aborted) {
+          task.status = "canceled";
+          task.errorMessage = err instanceof Error ? err.message : String(err);
+          task.completedAt = this.now();
+          this.inFlight -= 1;
+          this.dispatch().catch(() => undefined);
+          return;
+        }
+        // An unexpected prober error is recorded as a per-seller failure
+        // and the loop continues with the next seller.
+        logger.error("prewarm.probe.threw", "seller prober threw unexpectedly", {
+          taskId: task.id,
+          sellerId: seller.id,
+          modelId: task.modelId,
+          errorMessage: err instanceof Error ? err.message : String(err)
+        });
+        candidates.push({
+          sellerId: seller.id,
+          url: seller.url,
+          healthScore: 0,
+          lastSuccessAt: 0,
+          lastFailAt: this.now(),
+          avgLatencyMs: 0
+        });
+        continue;
+      }
+      if (signal?.aborted) {
+        task.status = "canceled";
+        task.completedAt = this.now();
+        this.lastProbeAtBySeller.set(seller.id, this.now());
+        this.inFlight -= 1;
+        this.dispatch().catch(() => undefined);
+        return;
+      }
+      this.lastProbeAtBySeller.set(seller.id, this.now());
+      if (result.ok) {
+        anyOk = true;
+        candidates.push({
+          sellerId: seller.id,
+          url: seller.url,
+          healthScore: scoreFromLatency(result.latencyMs),
+          lastSuccessAt: this.now(),
+          lastFailAt: 0,
+          avgLatencyMs: result.latencyMs
+        });
+        logger.info("prewarm.succeeded", "seller probe succeeded", {
+          taskId: task.id,
+          sellerId: seller.id,
+          modelId: task.modelId,
+          latencyMs: result.latencyMs,
+          httpStatus: result.httpStatus
+        });
+      } else {
+        candidates.push({
+          sellerId: seller.id,
+          url: seller.url,
+          healthScore: 0,
+          lastSuccessAt: 0,
+          lastFailAt: this.now(),
+          avgLatencyMs: result.latencyMs
+        });
+        logger.warn("prewarm.failed", "seller probe failed", {
+          taskId: task.id,
+          sellerId: seller.id,
+          modelId: task.modelId,
+          errorMessage: result.errorMessage,
+          httpStatus: result.httpStatus
+        });
+      }
+    }
+
+    if (!probedAny) {
+      // Every seller was rate-limited; the task is a no-op. Preserve the
+      // prior cache entry untouched and report the task as a no-op success
+      // so callers do not see a transient failure.
+      task.status = "succeeded";
+      task.completedAt = this.now();
+      this.totalSucceeded += 1;
+      this.inFlight -= 1;
+      this.dispatch().catch(() => undefined);
+      return;
+    }
+
+    // Filter out hard failures so the cache only stores reachable sellers.
+    const viable = candidates.filter((c) => c.healthScore > 0);
+    this.cache.commitWarm({
+      modelId: task.modelId,
+      protocol: task.protocol,
+      paymentMethod: task.paymentMethod,
+      candidates: viable
+    });
+
+    if (!anyOk) {
+      // No seller responded; record failure for the (model, protocol, payment)
+      // entry so the scheduler can back off.
+      this.cache.recordFailure(task.modelId, task.protocol, task.paymentMethod, "all probes failed");
+      task.status = "failed";
+      task.errorMessage = "all probes failed";
+      this.totalFailed += 1;
+    } else {
+      task.status = "succeeded";
+      this.totalSucceeded += 1;
+    }
+    task.completedAt = this.now();
+
+    // Free a slot and keep dispatching.
+    this.inFlight -= 1;
+    this.dispatch().catch(() => undefined);
+  }
+
+  private isOverBudget(): boolean {
+    return this.recentProbesInLastMinute() >= this.maxPrewarmPerMinute;
+  }
+
+  private recentProbesInLastMinute(): number {
+    const cutoff = this.now() - 60_000;
+    while (this.recentProbes.length > 0 && this.recentProbes[0] < cutoff) {
+      this.recentProbes.shift();
+    }
+    return this.recentProbes.length;
+  }
+
+  private recordProbeAttempt(): void {
+    this.recentProbes.push(this.now());
+  }
+
+  private isSellerRateLimited(sellerId: string): boolean {
+    const last = this.lastProbeAtBySeller.get(sellerId);
+    if (last === undefined) {
+      return false;
+    }
+    return this.now() - last < this.perSellerMinIntervalMs;
+  }
+}
+
+function defaultSleep(ms: number, signal?: AbortSignal): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(new Error("aborted"));
+      return;
+    }
+    const timer = setTimeout(() => {
+      cleanup();
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      cleanup();
+      reject(new Error("aborted"));
+    };
+    const cleanup = () => {
+      clearTimeout(timer);
+      signal?.removeEventListener("abort", onAbort);
+    };
+    signal?.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
+/**
+ * Build a per-probe abort signal that mirrors the scheduler's overall abort
+ * signal. Probers receive this scoped signal so aborting the scheduler
+ * propagates into any in-flight HTTP request (typically wired through
+ * `fetch(..., { signal })`).
+ */
+function composeProbeSignal(parent: AbortSignal | undefined): { signal: AbortSignal; abort(reason?: unknown): void } {
+  const controller = new AbortController();
+  if (parent) {
+    if (parent.aborted) {
+      controller.abort(parent.reason);
+    } else {
+      parent.addEventListener("abort", () => controller.abort(parent.reason), { once: true });
+    }
+  }
+  return controller;
+}
+
+function scoreFromLatency(latencyMs: number): number {
+  if (!Number.isFinite(latencyMs) || latencyMs < 0) {
+    return 30;
+  }
+  if (latencyMs <= 100) return 100;
+  if (latencyMs <= 300) return 90;
+  if (latencyMs <= 800) return 75;
+  if (latencyMs <= 1500) return 60;
+  if (latencyMs <= 3000) return 40;
+  return 20;
+}

package/src/provider-install.ts

@@ -469,7 +469,12 @@ function opencodeConfig(home: string, proxyUrl: string, config: ProviderRuntimeC
     : {};
   providers.tokenbuddy = {
     name: "TokenBuddy",
-    npm: "@ai-sdk/openai",
+    // v1.0.10 起默认走 OpenAI Responses API 协议（/v1/responses）。
+    // 验证：之前默认 @ai-sdk/openai（chat completions）也能 work，但 Responses API
+    // 才是上游（code.shoestravel.xin 等）原生支持的 SSE 事件链，type 字段更标准
+    // （response.created / response.output_text.delta / response.completed），
+    // 让 buyer 端 SseUsageExtractor 能稳定 parse usage 字段。
+    npm: "@ai-sdk/openai-responses",
     options: {
       apiKey: PROXY_ACCESS_TOKEN_PLACEHOLDER,
       baseURL: openAiBaseUrl(proxyUrl),
@@ -483,6 +488,9 @@ function opencodeConfig(home: string, proxyUrl: string, config: ProviderRuntimeC
     },
   };
   current.provider = providers;
+  // 写顶层 model / small_model，让 opencode 默认走 tokenbuddy 而不是残留的 openai/qwen-plus 死链
+  current.model = `tokenbuddy/${model}`;
+  current.small_model = `tokenbuddy/${model}`;
   return [makeChange("opencode", configPath, "configure OpenCode provider for TokenBuddy proxy", jsonContent(current))];
 }