marmar-context / web / context / server / model-runtime.ts
model-runtime.ts
Raw
// Local preview admission control. There is no unbounded waiting queue.
export class ModelError extends Error {
  constructor(
    message: string,
    public status = 503,
    public retryAfter?: number,
  ) {
    super(message);
  }
}
export const MODEL_CONCURRENCY = 4;
const active = new Set<AbortController>();
let cooldownUntil = 0;
export const modelLoad = () => ({
  active: active.size,
  limit: MODEL_CONCURRENCY,
});
export function cancelModelRequests() {
  for (const controller of active) controller.abort();
  cooldownUntil = 0;
}
export function retryAfterMs(value: string | null, now = Date.now()): number {
  if (!value) return 0;
  const seconds = Number(value);
  const ms = Number.isFinite(seconds)
    ? seconds * 1000
    : Date.parse(value) - now;
  return Number.isFinite(ms) ? Math.max(0, Math.min(3_600_000, ms)) : 0;
}
export function providerCooldown(ms: number) {
  cooldownUntil = Math.max(cooldownUntil, Date.now() + ms);
}
export function acquireModel(signal?: AbortSignal) {
  if (signal?.aborted)
    throw new ModelError(
      "Extraction cancelled. No model result was used.",
      499,
    );
  if (cooldownUntil > Date.now())
    throw new ModelError(
      "The model provider is rate limited. Please retry later; no model result was used.",
      503,
      Math.ceil((cooldownUntil - Date.now()) / 1000),
    );
  if (active.size >= MODEL_CONCURRENCY)
    throw new ModelError(
      "All four model slots are busy. Please retry shortly; your source notes are unchanged.",
      503,
      2,
    );
  const controller = new AbortController();
  const abort = () => controller.abort();
  signal?.addEventListener("abort", abort, { once: true });
  active.add(controller);
  return {
    signal: controller.signal,
    release() {
      signal?.removeEventListener("abort", abort);
      active.delete(controller);
    },
  };
}