opencode-ufr 0.2.9 → 0.2.11

package/src/daemon/catalog-source.ts

@@ -40,6 +40,15 @@ export async function loadUfrModels(o: {
         redirect: "manual",
       })
       const type = res.headers.get("content-type") ?? ""
+      // The 3xx check must come first: real redirects carry text/html, so the HTML
+      // branch below would otherwise swallow them (and could even misfire isVpnPage
+      // on a redirect page). Same ordering as callUpstream.
+      if (res.status >= 300 && res.status < 400) {
+        // Not proof of the VPN wall — cancel the body (like callUpstream does) so
+        // the connection is not left hanging on a redirect we will not follow.
+        await res.body?.cancel()
+        throw new Error(`UFR answered with a redirect (HTTP ${res.status}) instead of JSON`)
+      }
       if (type.includes("text/html")) {
         const html = await res.text()
         throw new Error(isVpnPage(html) ? VPN_MESSAGE : `HTML instead of JSON (HTTP ${res.status})`)

package/src/daemon/daemon.ts

@@ -2,6 +2,7 @@ import { randomBytes } from "node:crypto"
 import { mkdir, open, rm, stat } from "node:fs/promises"
 import { dirname } from "node:path"
 import { loadConfig, saveConfig } from "../shared/config"
+import { readDaemonInfo } from "../shared/daemon-client"
 import { readJson, readText, writeFileAtomic } from "../shared/fs"
 import type { Paths } from "../shared/paths"
 import type { SecretStore } from "../shared/secrets"
@@ -44,6 +45,8 @@ export type StatusJson = {
   vpn: { mode: string; detail: string } | null
   spendToday: Record<string, number>
   dailyBudgetUsd: number
+  /** Rolling-window throughput — what /v1/_status dashboards display as req/s and tok/s. */
+  rates: { windowMs: number; requests: number; reqPerSec: number; tokensInPerSec: number; tokensOutPerSec: number }
 }
 
 export type DaemonOptions = {
@@ -116,8 +119,11 @@ export async function acquireLock(lockFile: string, o: AcquireLockOptions = {}):
     heldNonces.add(nonce)
     try {
       const fh = await open(lockFile, "wx")
-      await fh.writeFile(`${process.pid} ${nonce}`)
-      await fh.close()
+      try {
+        await fh.writeFile(`${process.pid} ${nonce}`)
+      } finally {
+        await fh.close() // a failed writeFile must not leak the fd
+      }
       nonceByLock.set(lockFile, nonce)
       return true
     } catch (e) {
@@ -173,8 +179,8 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
   const p = o.paths
   const config = await loadConfig(p.configFile)
   const confirmDaemon = async (pid: number): Promise<boolean> => {
-    const saved = (await readJson(p.daemonFile)) as { pid?: number; port?: number } | null
-    if (!saved || saved.pid !== pid || typeof saved.port !== "number") return false
+    const saved = await readDaemonInfo(p)
+    if (!saved || saved.pid !== pid) return false
     try {
       const res = await fetch(`http://127.0.0.1:${saved.port}/health`, { signal: AbortSignal.timeout(2_000) })
       if (!res.ok) return false
@@ -290,11 +296,18 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
       probing = true
       try {
         for (const id of unknown.slice(0, 2)) {
+          if (stopping) break
+          // Probes share UFR's rate-limit bucket with real traffic: reserve a
+          // key slot per probe and keep the admission (the router keeps a
+          // completed call's too), so probes cannot oversubscribe the bucket.
+          const reserved = keys.acquire()
+          if (reserved.kind === "none") break // pool exhausted — the next refresh probes again
+          if (reserved.waitMs > 0) await sleep(reserved.waitMs)
           if (stopping) break
           const result = await probeContextLimit({
             model: id,
             baseUrl: config.upstream.baseUrl,
-            key: keyInfos[0]!.secret,
+            key: reserved.secret,
             transport,
             log,
           })
@@ -308,8 +321,10 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
           }
         }
         if (unknown.length > 0) {
-          const mf2 = await loadBundledModels({ bundledPath: o.bundledModelsPath })
-          catalog = buildCatalog(ufrModels, mf2.file, { allowPaid: config.allowPaid, probes: probeStore })
+          // Rebuild from the latest raw inputs: a newer catalog refresh may
+          // have replaced the catalog while these probes were in flight — the
+          // captured arguments are stale by now.
+          catalog = buildCatalog(lastUfrModels, lastModelsFile, { allowPaid: config.allowPaid, probes: probeStore })
           catalogInfo = { ...catalogInfo, loadedAt: now() }
         }
       } finally {
@@ -323,6 +338,8 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
         cachePath: p.ufrModelsCache, fetch: (u, i) => transport.fetch(u, i), log })
       catalog = buildCatalog(ufr.models, mf.file, { allowPaid: config.allowPaid, probes: probeStore })
       catalogInfo = { source: mf.source, ufrSource: ufr.source, loadedAt: now() }
+      lastUfrModels = ufr.models
+      lastModelsFile = mf.file
       if (ufr.error) Object.assign(reach, { ok: false, message: ufr.error, at: now() })
       else if (reach.ok !== true) Object.assign(reach, { ok: true, message: "", at: now() })
       for (const w of catalog.warnings) log(w)
@@ -331,6 +348,13 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
       else clearCatalogRetry()
       void scheduleProbes(ufr.models, mf.file).catch(() => {}) // must not outlive a shutdown
     }
+    // Raw inputs of the catalog, for scripts that probe every model (GET /v1/_catalog).
+    // lastModelsFile is assigned by refreshCatalog() right below, before any
+    // reader runs: scheduleProbes only fires from refreshCatalog (after the
+    // assignment) and catalogData is only served once the server has started,
+    // which happens after this awaited call.
+    let lastUfrModels: UfrModel[] = []
+    let lastModelsFile: ModelsFile
     await refreshCatalog()
 
     const router = new Router({
@@ -349,6 +373,19 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
     const startedAt = now()
     let lastActivity = now()
     let port = 0
+    /** Throughput over the last 60 s, from completed requests (a running stream's tokens land here when it ends). */
+    const ratesSnapshot = (): StatusJson["rates"] => {
+      const windowMs = 60_000
+      const r = db.rates(now() - windowMs)
+      const sec = windowMs / 1000
+      return {
+        windowMs,
+        requests: r.requests,
+        reqPerSec: Math.round((r.requests / sec) * 100) / 100,
+        tokensInPerSec: Math.round(r.promptTokens / sec),
+        tokensOutPerSec: Math.round(r.completionTokens / sec),
+      }
+    }
     const status = (): StatusJson => ({
       version: VERSION,
       pid: process.pid,
@@ -364,27 +401,40 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
       vpn: vpn ? { mode: vpn.status.mode, detail: vpn.status.detail } : null,
       spendToday: db.spendByKeySince(startOfLocalDay(now())),
       dailyBudgetUsd: config.dailyBudgetUsd,
+      rates: ratesSnapshot(),
     })
 
     let server: ReturnType<typeof startServer> | null = null
+    /** How long in-flight requests get to finish before connections and the db are force-closed. */
+    const drainGraceMs = 10_000
     const stop = async () => {
       if (stopping) return
       stopping = true
       for (const t of timers) clearInterval(t)
       clearCatalogRetry()
       try {
-        db.setKv("breakers", JSON.stringify(breakers.snapshot()))
-      } catch {
-        // db already closed
+        // Drain first: in-flight requests still write stats.record() through the
+        // db handle and owe their clients a response that server.stop(true)
+        // would reset. Idle shutdowns reach this loop with nothing in flight.
+        const deadline = now() + drainGraceMs
+        while (router.inFlight > 0 && now() < deadline) await sleep(50)
+        try {
+          db.setKv("breakers", JSON.stringify(breakers.snapshot()))
+        } catch {
+          // db already closed
+        }
+        server?.stop()
+        db.close()
+        await vpn?.stop()
+      } finally {
+        // Cleanup always runs — even when vpn.stop() or anything above throws —
+        // so no daemon-file/lock-file is left behind and onStopped still fires.
+        await rm(p.daemonFile, { force: true }).catch(() => {})
+        await rm(p.lockFile, { force: true }).catch(() => {})
+        releaseLockNonce(p.lockFile)
+        log("stopped")
+        o.onStopped?.()
       }
-      server?.stop()
-      db.close()
-      await vpn?.stop()
-      await rm(p.daemonFile, { force: true })
-      await rm(p.lockFile, { force: true })
-      releaseLockNonce(p.lockFile)
-      log("stopped")
-      o.onStopped?.()
     }
 
     const fixed = o.port !== undefined || config.port !== null
@@ -398,6 +448,8 @@ export async function startDaemon(o: DaemonOptions): Promise<RunningDaemon> {
           router,
           models: () => listModels(catalog),
           status,
+          draining: () => stopping,
+          catalogData: () => ({ ufr: lastUfrModels, file: lastModelsFile }),
           onActivity: () => {
             lastActivity = now()
           },

package/src/daemon/keypool.ts

@@ -5,10 +5,15 @@ export type Acquired = { kind: "ok"; alias: string; secret: string; waitMs: numb
 export type NotAcquired = { kind: "none"; reason: "no_keys" | "all_tried" | "exhausted"; retryAfterMs: number }
 export type KeySnapshot = { alias: string; used: number; cap: number; blockedForMs: number; invalid: boolean }
 
+/** A 401 can be transient (portal hiccup, short-lived token rotation): the key
+ *  is skipped for this long, then retried — instead of being dead until restart. */
+export const INVALID_KEY_TTL_MS = 10 * 60_000
+
 type Slot = KeyInfo & {
   window: SlidingWindow
   blockedUntil: number
-  invalid: boolean
+  /** Until when the key is considered invalid (0 = valid); a transient 401 expires. */
+  invalidUntil: number
   recent429: { at: number; model: string }[]
 }
 
@@ -24,14 +29,14 @@ export class KeyPool {
       ...k,
       window: new SlidingWindow({ cap: o.cap, windowMs: o.windowMs, maxWaitMs: o.maxWaitMs, now: o.now }),
       blockedUntil: 0,
-      invalid: false,
+      invalidUntil: 0,
       recent429: [],
     }))
   }
 
-  /** Number of keys not marked invalid. */
+  /** Number of keys not currently marked invalid. */
   get size(): number {
-    return this.slots.filter((s) => !s.invalid).length
+    return this.slots.filter((s) => s.invalidUntil <= this.o.now()).length
   }
 
   private find(alias: string): Slot | undefined {
@@ -44,7 +49,7 @@ export class KeyPool {
 
   acquire(exclude: ReadonlySet<string> = new Set()): Acquired | NotAcquired {
     const now = this.o.now()
-    const valid = this.slots.filter((s) => !s.invalid)
+    const valid = this.slots.filter((s) => s.invalidUntil <= now)
     if (valid.length === 0) return { kind: "none", reason: "no_keys", retryAfterMs: 0 }
     const usable = valid.filter((s) => !exclude.has(s.alias))
     if (usable.length === 0) return { kind: "none", reason: "all_tried", retryAfterMs: 0 }
@@ -93,7 +98,7 @@ export class KeyPool {
 
   onInvalid(alias: string): void {
     const s = this.find(alias)
-    if (s) s.invalid = true
+    if (s) s.invalidUntil = this.o.now() + INVALID_KEY_TTL_MS
   }
 
   snapshot(): KeySnapshot[] {
@@ -103,7 +108,7 @@ export class KeyPool {
       used: s.window.inWindow(),
       cap: this.o.cap,
       blockedForMs: Math.max(0, s.blockedUntil - now),
-      invalid: s.invalid,
+      invalid: s.invalidUntil > now,
     }))
   }
 }

package/src/daemon/main.ts

@@ -7,7 +7,13 @@ const paths = resolvePaths()
 const log = createLogger(paths.logFile)
 try {
   const d = await startDaemon({ paths, secrets: new KeyringStore(), log, onStopped: () => process.exit(0) })
-  for (const sig of ["SIGINT", "SIGTERM"] as const) process.on(sig, () => void d.stop())
+  let stopping = false
+  for (const sig of ["SIGINT", "SIGTERM"] as const)
+    process.on(sig, () => {
+      if (stopping) process.exit(1) // second signal: stop() is hanging — get out now
+      stopping = true
+      void d.stop()
+    })
 } catch (e) {
   if (e instanceof AlreadyRunningError) {
     log("another daemon is already running — exiting")

package/src/daemon/probe-all.ts

@@ -13,13 +13,23 @@
  */
 
 import type { ModelsFile } from "../shared/models-file"
-import type { FetchLike } from "./catalog-source"
-import { parseUfrModels } from "./catalog"
+import type { UfrModel } from "./catalog"
 import { fillerForTokens, parseLimitFromBody } from "./probe"
 
 /** No UFR model is known to exceed this; probing higher only risks an accepted (paid) rung. */
 export const PROBE_CEILING = 1_600_000
 
+/**
+ * Sub-1% deltas are noise, not context changes: some backends name the
+ * input-only limit (context minus reserved output tokens), so the same model
+ * can answer 262,144 locally and 261,144 from CI (observed 2026-10-05).
+ * Without a tolerance the weekly run would flip-flop such values forever.
+ */
+export const PROBE_TOLERANCE = 0.01
+
+export const withinTolerance = (probed: number, known: number): boolean =>
+  Math.abs(probed - known) / known < PROBE_TOLERANCE
+
 export type ProbeHow =
   | "error-named" // the server named the exact limit — the trustworthy result
   | "accepted-floor" // every rung up to `probed` was accepted — `probed` is a lower bound only
@@ -37,13 +47,23 @@ export type ProbeRow = {
   detail?: string
 }
 
+/**
+ * One model call. Through the gateway's relay (POST /v1/_relay) this gets the
+ * central key rotation and soft rate limiting for free — the call waits for a
+ * key slot instead of being told to back off, and UFR's status and body come
+ * back verbatim (a context probe needs the 400 body that names the limit).
+ */
+export type ProbeCallResult = { status: number; body: string; retryAfterMs?: number }
+export type ProbeCall = (body: Record<string, unknown>) => Promise<ProbeCallResult>
+
 export type ProbeAllDeps = {
-  baseUrl: string
-  key: string
-  fetch: FetchLike
+  call: ProbeCall
+  /** UFR's model list — the caller's job to fetch (gateway catalog or direct). */
+  ufr: UfrModel[]
   file: ModelsFile // known values + exclude list
   includePaid?: boolean
-  /** Pause between models — the per-key bucket is 20/min across ALL models. */
+  /** Pause between models for direct callers that pace themselves. The gateway
+   *  relay paces on its own — pass 0 there and the key pool decides. */
   paceMs?: number
   log: (m: string) => void
   now?: () => number
@@ -72,60 +92,58 @@ async function probeOne(o: ProbeAllDeps, id: string, known: number | null): Prom
   let last = { status: 0, body: "" }
   for (const size of sizesFor(known)) {
     for (let tries = 0; ; tries++) {
-      let res: Response
+      let res: ProbeCallResult
       try {
-        res = await o.fetch(`${o.baseUrl}/chat/completions`, {
-          method: "POST",
-          headers: { Authorization: `Bearer ${o.key}`, "Content-Type": "application/json" },
-          body: JSON.stringify({
-            model: id,
-            max_tokens: 1,
-            messages: [{ role: "user", content: fillerForTokens(size) + "\n\nReply with exactly: OK" }],
-          }),
-          signal: AbortSignal.timeout(180_000),
-          redirect: "manual",
+        res = await o.call({
+          model: id,
+          max_tokens: 1,
+          messages: [{ role: "user", content: fillerForTokens(size) + "\n\nReply with exactly: OK" }],
         })
       } catch (e) {
         return { id, tier: "free", known, probed: floor.value || null,
           how: floor.value ? "accepted-floor" : "http-error", detail: `transport: ${(e as Error).message}` }
       }
-      if (res.ok) {
-        const j = (await res.json().catch(() => null)) as { usage?: { prompt_tokens?: number } } | null
+      if (res.status >= 200 && res.status < 300) {
+        const j = parseJson(res.body) as { usage?: { prompt_tokens?: number } } | null
         floor.value = j?.usage?.prompt_tokens ?? size
         break // accepted — try the next rung
       }
-      const text = await res.text().catch(() => "")
       if (res.status === 429 && tries < 2) {
-        o.log(`  ${id}: 429 at ${size} tokens — backing off 20 s`)
-        await doSleep(20_000)
-        continue // bucket, not a context answer — retry the same rung
+        // 429 through the relay means walled, pool-capped or budget: all rare.
+        // Respect the gateway's retry-after when it names one (capped so a
+        // 30-minute ladder doesn't stall a whole probe run).
+        const wait = Math.min(res.retryAfterMs ?? 20_000, 120_000)
+        o.log(`  ${id}: 429 at ${size} tokens — backing off ${Math.round(wait / 1000)} s`)
+        await doSleep(wait)
+        continue // not a context answer — retry the same rung
       }
-      last = { status: res.status, body: text }
+      last = { status: res.status, body: res.body }
       if (res.status === 400 || res.status === 413) {
-        const named = parseLimitFromBody(text)
+        const named = parseLimitFromBody(res.body)
         if (named) return { id, tier: "free", known, probed: named, how: "error-named" }
         return { id, tier: "free", known, probed: floor.value || null,
-          how: floor.value ? "accepted-floor" : "rejected-unnamed", detail: text.slice(0, 200) }
+          how: floor.value ? "accepted-floor" : "rejected-unnamed", detail: res.body.slice(0, 200) }
       }
       return { id, tier: "free", known, probed: floor.value || null,
-        how: floor.value ? "accepted-floor" : "http-error", detail: `HTTP ${res.status}: ${text.slice(0, 200)}` }
+        how: floor.value ? "accepted-floor" : "http-error", detail: `HTTP ${res.status}: ${res.body.slice(0, 200)}` }
     }
   }
   return { id, tier: "free", known, probed: floor.value || null, how: "accepted-floor",
     detail: last.status ? `last: HTTP ${last.status}` : undefined }
 }
 
+function parseJson(text: string): unknown {
+  try {
+    return JSON.parse(text)
+  } catch {
+    return null
+  }
+}
+
 export async function probeAllContexts(o: ProbeAllDeps): Promise<ProbeRow[]> {
-  const res = await o.fetch(`${o.baseUrl}/models`, {
-    headers: { Authorization: `Bearer ${o.key}` },
-    signal: AbortSignal.timeout(20_000),
-    redirect: "manual",
-  })
-  if (!res.ok) throw new Error(`UFR /api/models: HTTP ${res.status}`)
-  const ufr = parseUfrModels(await res.json())
   const excluded = new Set(o.file.exclude)
   const rows: ProbeRow[] = []
-  for (const m of ufr) {
+  for (const m of o.ufr) {
     if (excluded.has(m.id)) {
       rows.push({ id: m.id, tier: m.tier, known: o.file.models[m.id]?.context ?? null, probed: null, how: "skipped-hidden" })
       continue
@@ -135,7 +153,7 @@ export async function probeAllContexts(o: ProbeAllDeps): Promise<ProbeRow[]> {
       continue
     }
     rows.push(await probeOne(o, m.id, o.file.models[m.id]?.context ?? null))
-    await (o.sleep ?? sleep)(o.paceMs ?? 3_000) // stay under the 20/min key bucket
+    if ((o.paceMs ?? 0) > 0) await (o.sleep ?? sleep)(o.paceMs!) // direct mode only: the relay paces itself
   }
   return rows
 }
@@ -152,7 +170,11 @@ export function formatProbeReport(rows: ProbeRow[]): string {
     const context = r.probed === null ? "—" : r.probed.toLocaleString("en-US")
     const status =
       r.how === "error-named" && r.probed !== null
-        ? r.known === null ? "NEW — needs models.json entry" : r.probed === r.known ? "ok" : `MISMATCH (was ${r.known.toLocaleString("en-US")})`
+        ? r.known === null
+          ? "NEW — needs models.json entry"
+          : withinTolerance(r.probed, r.known)
+            ? r.probed === r.known ? "ok" : `ok (±${Math.abs(r.probed - r.known).toLocaleString("en-US")})`
+            : `MISMATCH (was ${r.known.toLocaleString("en-US")})`
         : r.how
     if (status.startsWith("NEW") || status.startsWith("MISMATCH")) mismatches++
     lines.push(pad(r.id, 38) + pad(context, 12) + pad(r.how, 17) + pad(r.known === null ? "—" : r.known.toLocaleString("en-US"), 12) + status)
@@ -180,7 +202,7 @@ export function applyProbeResults(
     if (r.how !== "error-named" || r.probed === null) continue
     const e = (out.models[r.id] ??= {})
     const from = e.context ?? null
-    if (from === r.probed) continue
+    if (from === r.probed || (from !== null && withinTolerance(r.probed, from))) continue
     e.context = r.probed
     const probeNote = `context ${r.probed} probed live ${date} (probe-all-contexts)`
     e.note = e.note?.includes("probed live") ? probeNote : e.note ? `${e.note} | ${probeNote}` : probeNote