diff --git a/.informant/jobs/build.toml b/.informant/jobs/build.toml index 84a2d53..0d317f3 100644 --- a/.informant/jobs/build.toml +++ b/.informant/jobs/build.toml @@ -2,11 +2,7 @@ name = "build" needs = ["test", "typecheck"] command = """ set -eu -expected_dist="$(mktemp -d)" -trap 'rm -rf "$expected_dist"' EXIT -cp -R web/dist "$expected_dist/dist" bun run build -diff -ru "$expected_dist/dist" web/dist """ timeout_minutes = 15 container = { cpu = 2, memory_mb = 4096 } diff --git a/README.md b/README.md index e91c44d..fac55fd 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,65 @@ It uses the GitHub CLI to resolve the pull request and check status for the chec The subagent extension independently contributes its token use and status to `extensions/session-footer.ts`, the package's generic composable footer. When subagents are involved, a third footer line shows their aggregate status. With an empty editor, press Option+Down (Alt+Down) to select that line and Enter to open the manager; `/subagents` opens it directly. The manager shows individual status and transcripts and supports model, effort, messaging, and termination controls. Run `/subagents-cleanup` to stop and remove every retained subagent. +## Auto model routing + +`extensions/auto-router.ts` adds an "Auto" entry to `/model`. Selecting it routes each turn to a model/reasoning-effort pair chosen from your own configured lists, based on the turn's classified complexity, and fails over to other configured models or tiers when one is unhealthy or out of usage. + +Configure it under a new `autoRouter` key in `~/.pi/agent/settings.json` (or `.pi/settings.json` for a project override): + +```json +{ + "autoRouter": { + "efforts": { + "medium": { + "models": [ + { "provider": "anthropic", "id": "claude-sonnet-4-5" }, + { "provider": "openai", "id": "gpt-5.3-codex" } + ] + }, + "high": { + "models": [{ "provider": "anthropic", "id": "claude-opus-4-7" }] + }, + "xhigh": { + "models": [{ "provider": "openai", "id": "gpt-5.6-sol" }] + } + } + } +} +``` + +Each tier key is a Pi thinking level (`off`, `minimal`, `low`, `medium`, `high`, `xhigh`, `max`); `medium` is the default/anchor tier. Each tier holds an ordered list of `{ provider, id }` model references — the first is preferred, later entries are failover within that tier. + +Which tier a model is listed under only decides *when it's used* (which classified-complexity bucket routes to it, and where it sits in the escalation order) — it doesn't have to be the reasoning effort that model is actually dispatched at. Add `"effort"` to a model reference to pin its own thinking level independent of its tier, e.g. a model that only performs well at its own maximum setting can still live under `high` (so moderately-hard tasks reach it and it takes part in escalation normally) while always running at `max`: + +```json +"high": { + "models": [{ "provider": "opencode-go", "id": "kimi-k3", "effort": "max" }] +} +``` + +Omit `effort` and a model just uses its tier's own name, as before. + +On every turn, Auto asks the `medium` tier's first healthy model (the "default model") to classify the turn as `minimal`, `low`, `medium`, `high`, `xhigh`, or `max`, then routes to the resolved tier: if the classified level has no configured models, it steps toward `medium` until it finds one (a classified `low` with nothing configured there falls back to `medium`). Within that tier it picks the first model that isn't in a failure/rate-limit cooldown; if every model in the tier is unhealthy, it escalates to the next *higher* configured tier; if nothing anywhere is healthy, it uses the first configured model anyway rather than blocking the turn, with a warning. + +`/model` shows a separate entry per configured tier — "Auto (auto)" for the classify-every-turn behavior above, plus "Auto (medium)", "Auto (high)", and so on for each tier that has at least one model configured (tiers with nothing configured don't get an entry). Picking a specific one pins Auto to that tier: every turn skips classification and routes directly within it — still with the same failover, escalation, and health tracking as the adaptive mode, just without asking a model to judge complexity first. This list is fixed at startup from whatever's configured then, so adding a new tier to `autoRouter` needs a Pi restart before its "Auto (\)" entry shows up. + +Health is tracked from two sources. Router-observed traffic (HTTP status codes, rate-limit headers, and message-level provider errors that never surface as a bad HTTP status) sets an immediate cooldown the moment any turn against a configured model fails — whether Auto routed there itself or you picked it manually from `/model`; a model configured in `autoRouter` is tracked the same way either way. Separately, best-effort real quota reconciliation runs at session start and on `/usage`, for providers with a known quota source: Anthropic, OpenAI Codex, Z.ai, Kimi Coding, and OpenCode Go via their HTTP APIs (using the same credentials Pi already has for each), plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no HTTP quota endpoint of its own. This is what lets the router self-correct for usage consumed truly outside its view (a different session or machine, another tool, or before Auto was set up) instead of only reacting to its own observations. Codex specifically reports quota per-model for models it meters individually (its own `additional_rate_limits` entries) — those are independent of its account-wide limit in both directions, so a model with its own entry is neither blocked by, nor shielded by, the account-wide state; only models without one fall back to it. Providers without a known quota source simply stay on router-observed data. + +Run `/usage` to see health and usage for every configured model, grouped by tier. Each row shows its cooldown status if any, the real "verified usage" reported by the provider's own quota API when available — always normalized to "X% used" regardless of how the provider itself reports it, with each window labeled by its real duration rather than a vague placeholder where the provider's response makes that derivable (e.g. "7d 5% used", "5h 16% used, weekly 11% used") — and separately the request/token/cost totals *this Pi installation* has observed for that model, whether Auto routed there or it was picked manually. The latter still won't reflect usage from other sessions/machines/tools or from before Auto started tracking, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). + +`/usage` also shows the last several routing decisions under "Recent classifications" — what the classifier's raw reply actually was, the level it parsed to, and the tier/model it routed to. The classification call itself is otherwise a throwaway completion whose result would normally vanish the moment it's parsed, so if a turn ever looks under- or over-routed, this is what to check first rather than guessing from the code. + +The `/model` picker's effort/thinking control is inert while any Auto entry is selected, since effort is chosen per turn (or fixed to the pinned tier) internally. `/model` keeps showing whichever Auto entry you picked selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to that same inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows "Auto (auto)" or "Auto (high)" (whichever you picked), not whichever model last handled a turn. A `🔀 Auto ()` badge in the TUI footer tracks the most recently applied thinking level regardless of which Auto entry is currently selected. Manually picking a real (non-Auto) model from `/model` turns Auto off; reselecting any Auto entry turns it back on. + +If you've scoped `/model` with `enabledModels` (or `--models`), Pi's picker defaults to showing only that scoped list, hiding everything else — including every Auto entry — behind a manual Tab to "all". At session start, Auto best-effort appends an `auto/*` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so every Auto entry shows up in the default scoped view too, without changing anything else about what's scoped. + +### Requirements + +- Pi 0.84.1 +- Network access from the machine running Pi, for the optional quota reconciliation calls (never required — routing and `/usage` work fully offline from router-observed data alone) +- For Minimax quota reconciliation specifically: MiniMax's own `mmx` CLI on `PATH`, logged in via `mmx auth login`. Without it, Minimax models just stay on router-observed data like any other unsupported provider. + ## Worktrees Run `/worktree ` to create `/.pi/worktrees/`, run the optional `.pi/worktrees/setup.sh`, and move the active conversation into a replacement session rooted in the managed checkout. The backward-compatible default creates or reuses local branch ``; a missing branch starts at the selected checkout's `HEAD`. @@ -129,5 +188,5 @@ bun install --frozen-lockfile bun run check bun test bun run webBuild -pi -e ./extensions/session-footer.ts -e ./extensions/pr-footer.ts -e ./extensions/subagents.ts -e ./extensions/worktree.ts -e ./extensions/web-sessions.ts +pi -e ./extensions/session-footer.ts -e ./extensions/pr-footer.ts -e ./extensions/subagents.ts -e ./extensions/worktree.ts -e ./extensions/web-sessions.ts -e ./extensions/auto-router.ts ``` diff --git a/bun.lock b/bun.lock index ba4da97..82f360c 100644 --- a/bun.lock +++ b/bun.lock @@ -15,6 +15,7 @@ }, "devDependencies": { "@biomejs/biome": "2.3.7", + "@earendil-works/pi-ai": "0.84.1", "@earendil-works/pi-coding-agent": "0.84.1", "@earendil-works/pi-tui": "0.84.1", "@tailwindcss/vite": "4.1.18", @@ -37,11 +38,13 @@ "vite": "7.3.6", }, "peerDependencies": { + "@earendil-works/pi-ai": "*", "@earendil-works/pi-coding-agent": "0.84.1", "@earendil-works/pi-tui": "*", "typebox": "*", }, "optionalPeers": [ + "@earendil-works/pi-ai", "@earendil-works/pi-coding-agent", "@earendil-works/pi-tui", "typebox", diff --git a/extensions/auto-router-classify.ts b/extensions/auto-router-classify.ts new file mode 100644 index 0000000..e8b9ef4 --- /dev/null +++ b/extensions/auto-router-classify.ts @@ -0,0 +1,119 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import { uuidv7 } from "@earendil-works/pi-ai"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; +import type { AutoRouterEffortLevel } from "./auto-router-settings.js"; + +const CLASSIFY_TIMEOUT_MS = 15_000; +const VALID_LEVELS: readonly AutoRouterEffortLevel[] = [ + "minimal", + "low", + "medium", + "high", + "xhigh", + "max", +]; +const DEFAULT_LEVEL: AutoRouterEffortLevel = "medium"; + +const SYSTEM_PROMPT = `You triage the complexity of a single upcoming coding-agent turn so it can be routed to an appropriately capable model. Reply with exactly one word, lowercase, no punctuation: minimal, low, medium, high, xhigh, or max. + +- minimal: rote, no real reasoning needed. A one-word answer, a pure formatting pass, a trivial rename, echoing back something already known. +- low: simple and mechanical, but not entirely rote. One-line edits, small lookups, answering a quick factual question about the codebase. +- medium: a typical coding task. Implementing a small-to-moderate feature, fixing a well-understood bug, writing straightforward tests. This is the default for ordinary work. +- high: meaningfully harder. Multi-file refactors, tricky or intermittent bugs, non-obvious architectural changes, tasks that require holding a lot of context at once. +- xhigh: very hard, high-stakes, or open-ended. Large-scope redesigns, subtle correctness/security-critical work, or reasoning-heavy problems where getting it wrong is costly. +- max: the hardest, rarest cases. Deep multi-step reasoning under real stakes — major system redesigns, subtle distributed-systems or security bugs, decisions with significant real-world consequences. + +Reply with only the single word.`; + +export type ClassificationUsage = { + input: number; + output: number; + cost: number; +}; + +export type ClassificationResult = { + level: AutoRouterEffortLevel; + usage?: ClassificationUsage; + /** + * The classifier's raw reply (trimmed/lowercased), or a bracketed placeholder when the call + * itself failed. Callers should log this alongside `level` - otherwise there's no way to tell + * apart "the model genuinely said medium" from "parsing picked the wrong word out of a messy + * reply" after the fact, since the model call itself is never persisted anywhere else. + */ + reply: string; +}; + +function numeric(value: unknown): number { + return typeof value === "number" && Number.isFinite(value) ? value : 0; +} + +/** + * Classify a turn's complexity using the given (default/medium-tier) model. Never throws and + * never blocks indefinitely: a bounded timeout, a provider error, or an unparseable reply all + * fall back to `medium` so classification can never stall or break the user's turn. + */ +export async function classifyTurnComplexity( + modelRegistry: ModelRegistry, + model: Model, + prompt: string, + hasImages: boolean, +): Promise { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), CLASSIFY_TIMEOUT_MS); + try { + const text = hasImages + ? `${prompt}\n\n(This turn also includes attached images.)` + : prompt; + const response = await modelRegistry.complete( + model, + { + systemPrompt: SYSTEM_PROMPT, + messages: [ + { + role: "user", + content: [{ type: "text", text }], + timestamp: Date.now(), + }, + ], + }, + { + signal: controller.signal, + reasoningEffort: "off", + cacheRetention: "none", + sessionId: uuidv7(), + maxTokens: 20, + }, + ); + const reply = response.content + .filter( + (block): block is { type: "text"; text: string } => + block.type === "text", + ) + .map((block) => block.text) + .join("") + .trim() + .toLowerCase(); + // Match whichever valid level word appears *first in the reply text*, not the first one + // in VALID_LEVELS' own order - a naive `VALID_LEVELS.find(word-boundary test)` would let an + // earlier-in-that-list word like "medium" win over a later one like "high" even when "high" + // is the word the model actually led with (e.g. "high complexity, more than a medium task"), + // silently downgrading the classification. `\b(...)\b` as one alternation also keeps the + // existing "high" vs "xhigh" substring safety: `\b` can't match between two word characters, + // so "high" never matches inside "xhigh" regardless of alternation order. + const match = reply.match(new RegExp(`\\b(${VALID_LEVELS.join("|")})\\b`)); + const level = (match?.[1] as AutoRouterEffortLevel | undefined) ?? DEFAULT_LEVEL; + const usage = response.usage + ? { + input: numeric(response.usage.input), + output: numeric(response.usage.output), + cost: numeric(response.usage.cost?.total), + } + : undefined; + return { level, usage, reply: reply || "(empty reply)" }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { level: DEFAULT_LEVEL, reply: `(classification failed: ${message})` }; + } finally { + clearTimeout(timeout); + } +} diff --git a/extensions/auto-router-health.ts b/extensions/auto-router-health.ts new file mode 100644 index 0000000..5dd8d2d --- /dev/null +++ b/extensions/auto-router-health.ts @@ -0,0 +1,401 @@ +import { randomUUID } from "node:crypto"; +import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises"; +import { dirname, join } from "node:path"; +import { getAgentDir } from "@earendil-works/pi-coding-agent"; + +/** Structural — matches both `AutoRouterModelRef` and the SDK's `Model`. */ +export type ModelIdentity = { provider: string; id: string }; + +const SAVE_DEBOUNCE_MS = 2_000; + +/** Resolved at call time (not module load) so it honors a `PI_CODING_AGENT_DIR` override set after import. */ +function statePath(): string { + return join(getAgentDir(), "auto-router-state.json"); +} + +const RATE_LIMIT_BASE_COOLDOWN_MS = 5 * 60_000; +const RATE_LIMIT_MAX_COOLDOWN_MS = 60 * 60_000; +const SERVER_ERROR_FAILURE_THRESHOLD = 3; +const SERVER_ERROR_COOLDOWN_MS = 2 * 60_000; +const AUTH_ERROR_COOLDOWN_MS = 30 * 60_000; +const GENERIC_FAILURE_THRESHOLD = 5; +const GENERIC_COOLDOWN_MS = 5 * 60_000; + +export type ModelHealthEntry = { + consecutiveFailures: number; + cooldownUntil?: number; + lastError?: { status: number; at: number }; + totals: { requests: number; input: number; output: number; cost: number }; + /** epoch ms of the last successful real quota-API reconciliation, if any. */ + verifiedAt?: number; + /** Human-readable real usage from the last quota-API reconciliation, e.g. "62% used (7d)". */ + verifiedDetail?: string; +}; + +export type AutoRouterHealthState = Record; + +const CLASSIFICATION_LOG_LIMIT = 20; +const CLASSIFICATION_LOG_TEXT_LIMIT = 200; + +/** + * A single routing decision, kept so `/usage` can show what the classifier actually said - the + * classification call itself is otherwise a throwaway completion whose result is discarded + * after parsing, so without this there's no way to tell apart "the model genuinely said medium" + * from "the reply parsed wrong" after the fact. Deliberately excludes the user's prompt text: + * `/usage` never displays it, and this file is plaintext on disk, so persisting it would be + * pure liability - prompts can contain source code, credentials, or personal data - for no + * actual benefit. + */ +export type ClassificationLogEntry = { + timestamp: number; + reply: string; + level: string; + tier: string; + /** The thinking level actually applied - a model's own `effort` override, or `tier` when it has none. */ + effort: string; + model: ModelIdentity; +}; + +function truncateForLog(text: string): string { + const collapsed = text.replace(/\s+/g, " ").trim(); + return collapsed.length > CLASSIFICATION_LOG_TEXT_LIMIT + ? `${collapsed.slice(0, CLASSIFICATION_LOG_TEXT_LIMIT)}…` + : collapsed; +} + +export type UsageDelta = { input: number; output: number; cost: number }; + +export type QuotaReconciliationResult = { + exhausted: boolean; + /** epoch ms the exhausted window resets, if known. */ + resetsAt?: number; + /** Human-readable real usage summary from the provider, e.g. "62% used (7d)". Shown in /usage when present. */ + detail?: string; +}; + +export function modelKey(model: ModelIdentity): string { + return `${model.provider}/${model.id}`; +} + +function defaultEntry(): ModelHealthEntry { + return { + consecutiveFailures: 0, + totals: { requests: 0, input: 0, output: 0, cost: 0 }, + }; +} + +/** Parse a `retry-after` header value (seconds, or an HTTP date) into a millisecond delay from `now`. */ +export function parseRetryAfterMs( + headers: Record | undefined, + now: number, +): number | undefined { + const raw = headers?.["retry-after"] ?? headers?.["Retry-After"]; + if (!raw) return undefined; + const seconds = Number(raw); + if (Number.isFinite(seconds) && seconds >= 0) return seconds * 1000; + const at = Date.parse(raw); + return Number.isFinite(at) ? Math.max(0, at - now) : undefined; +} + +export function isHealthy( + state: AutoRouterHealthState, + key: string, + now: number, +): boolean { + const entry = state[key]; + return !entry?.cooldownUntil || entry.cooldownUntil <= now; +} + +export function pickHealthy( + state: AutoRouterHealthState, + models: ModelIdentity[], + now: number, +): ModelIdentity | undefined { + return models.find((model) => isHealthy(state, modelKey(model), now)); +} + +export function applySuccess( + state: AutoRouterHealthState, + key: string, + usage: UsageDelta, +): AutoRouterHealthState { + const previous = state[key] ?? defaultEntry(); + const entry: ModelHealthEntry = { + ...previous, + consecutiveFailures: 0, + cooldownUntil: undefined, + totals: { + requests: previous.totals.requests + 1, + input: previous.totals.input + usage.input, + output: previous.totals.output + usage.output, + cost: previous.totals.cost + usage.cost, + }, + }; + return { ...state, [key]: entry }; +} + +/** Apply a provider HTTP response outcome. `status` in `[200,300)` is treated as success with no usage delta (see `applySuccess` for usage accounting). */ +export function applyFailure( + state: AutoRouterHealthState, + key: string, + status: number, + headers: Record | undefined, + now: number, +): AutoRouterHealthState { + const previous = state[key] ?? defaultEntry(); + const consecutiveFailures = previous.consecutiveFailures + 1; + let cooldownUntil = previous.cooldownUntil; + + if (status === 429) { + const retryAfterMs = parseRetryAfterMs(headers, now); + const backoffExponent = Math.min(consecutiveFailures, 5) - 1; + const backoffMs = Math.min( + RATE_LIMIT_BASE_COOLDOWN_MS * 2 ** backoffExponent, + RATE_LIMIT_MAX_COOLDOWN_MS, + ); + cooldownUntil = now + (retryAfterMs ?? backoffMs); + } else if (status === 401 || status === 403) { + cooldownUntil = now + AUTH_ERROR_COOLDOWN_MS; + } else if (status >= 500) { + if (consecutiveFailures >= SERVER_ERROR_FAILURE_THRESHOLD) { + cooldownUntil = now + SERVER_ERROR_COOLDOWN_MS; + } + } else if (consecutiveFailures >= GENERIC_FAILURE_THRESHOLD) { + cooldownUntil = now + GENERIC_COOLDOWN_MS; + } + + const entry: ModelHealthEntry = { + ...previous, + consecutiveFailures, + cooldownUntil, + lastError: { status, at: now }, + }; + return { ...state, [key]: entry }; +} + +/** + * Merge a real quota-API reconciliation result. Real data always wins over locally-inferred + * state: an exhausted window sets `cooldownUntil` even if the router never saw a 429 itself, + * and confirmed headroom clears any existing cooldown/failure count outright. + */ +export function applyQuotaResult( + state: AutoRouterHealthState, + key: string, + result: QuotaReconciliationResult, + now: number, +): AutoRouterHealthState { + const previous = state[key] ?? defaultEntry(); + const entry: ModelHealthEntry = result.exhausted + ? { + ...previous, + cooldownUntil: + result.resetsAt && result.resetsAt > now + ? result.resetsAt + : now + GENERIC_COOLDOWN_MS, + verifiedAt: now, + verifiedDetail: result.detail, + } + : { + ...previous, + consecutiveFailures: 0, + cooldownUntil: undefined, + verifiedAt: now, + verifiedDetail: result.detail, + }; + return { ...state, [key]: entry }; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function parseState(value: unknown): AutoRouterHealthState { + if (!isRecord(value)) return {}; + const state: AutoRouterHealthState = {}; + for (const [key, raw] of Object.entries(value)) { + if (!isRecord(raw) || !isRecord(raw.totals)) continue; + state[key] = { + consecutiveFailures: Number(raw.consecutiveFailures) || 0, + cooldownUntil: + typeof raw.cooldownUntil === "number" ? raw.cooldownUntil : undefined, + lastError: isRecord(raw.lastError) + ? { + status: Number(raw.lastError.status) || 0, + at: Number(raw.lastError.at) || 0, + } + : undefined, + totals: { + requests: Number(raw.totals.requests) || 0, + input: Number(raw.totals.input) || 0, + output: Number(raw.totals.output) || 0, + cost: Number(raw.totals.cost) || 0, + }, + verifiedAt: + typeof raw.verifiedAt === "number" ? raw.verifiedAt : undefined, + verifiedDetail: + typeof raw.verifiedDetail === "string" ? raw.verifiedDetail : undefined, + }; + } + return state; +} + +function parseClassifications(value: unknown): ClassificationLogEntry[] { + if (!Array.isArray(value)) return []; + const entries: ClassificationLogEntry[] = []; + for (const raw of value) { + if (!isRecord(raw)) continue; + if (typeof raw.timestamp !== "number") continue; + if (typeof raw.reply !== "string") continue; + if (typeof raw.level !== "string" || typeof raw.tier !== "string") continue; + if ( + !isRecord(raw.model) || + typeof raw.model.provider !== "string" || + typeof raw.model.id !== "string" + ) + continue; + // Any `prompt` field from an entry logged before this was dropped is intentionally not + // read back here, so a reload+resave (e.g. the next classification) scrubs it from disk. + entries.push({ + timestamp: raw.timestamp, + reply: raw.reply, + level: raw.level, + tier: raw.tier, + // Back-compat: entries logged before the `effort` field existed default to the tier name, + // matching what actually ran for them at the time. + effort: typeof raw.effort === "string" ? raw.effort : raw.tier, + model: { provider: raw.model.provider, id: raw.model.id }, + }); + } + return entries; +} + +/** + * Handles both the current `{models, classifications}` shape and the flat `Record` shape every persisted file had before classification logging existed. + */ +function parsePersisted(value: unknown): { + models: AutoRouterHealthState; + classifications: ClassificationLogEntry[]; +} { + if (!isRecord(value)) return { models: {}, classifications: [] }; + if (isRecord(value.models)) { + return { + models: parseState(value.models), + classifications: parseClassifications(value.classifications), + }; + } + return { models: parseState(value), classifications: [] }; +} + +/** Debounced, best-effort persistence for router-observed health/usage state, shared across concurrent Pi processes on a last-write-wins basis (telemetry, not correctness-critical config). */ +export class AutoRouterHealthStore { + private state: AutoRouterHealthState = {}; + private classifications: ClassificationLogEntry[] = []; + private writeTimer: ReturnType | undefined; + + async load(): Promise { + try { + const parsed = parsePersisted( + JSON.parse(await readFile(statePath(), "utf8")), + ); + this.state = parsed.models; + this.classifications = parsed.classifications; + } catch { + this.state = {}; + this.classifications = []; + } + } + + getState(): AutoRouterHealthState { + return this.state; + } + + getClassifications(): readonly ClassificationLogEntry[] { + return this.classifications; + } + + getEntry(key: string): ModelHealthEntry | undefined { + return this.state[key]; + } + + isHealthy(key: string, now: number = Date.now()): boolean { + return isHealthy(this.state, key, now); + } + + pickHealthy( + models: ModelIdentity[], + now: number = Date.now(), + ): ModelIdentity | undefined { + return pickHealthy(this.state, models, now); + } + + recordSuccess(key: string, usage: UsageDelta): void { + this.state = applySuccess(this.state, key, usage); + this.scheduleSave(); + } + + recordFailure( + key: string, + status: number, + headers: Record | undefined, + now: number = Date.now(), + ): void { + this.state = applyFailure(this.state, key, status, headers, now); + this.scheduleSave(); + } + + applyQuotaResult( + key: string, + result: QuotaReconciliationResult, + now: number = Date.now(), + ): void { + this.state = applyQuotaResult(this.state, key, result, now); + this.scheduleSave(); + } + + recordClassification( + entry: Omit, + now: number = Date.now(), + ): void { + const full: ClassificationLogEntry = { + timestamp: now, + reply: truncateForLog(entry.reply), + level: entry.level, + tier: entry.tier, + effort: entry.effort, + model: entry.model, + }; + this.classifications = [...this.classifications, full].slice( + -CLASSIFICATION_LOG_LIMIT, + ); + this.scheduleSave(); + } + + private scheduleSave(): void { + if (this.writeTimer) return; + this.writeTimer = setTimeout(() => { + this.writeTimer = undefined; + void this.flush(); + }, SAVE_DEBOUNCE_MS); + this.writeTimer.unref?.(); + } + + async flush(): Promise { + const dir = dirname(statePath()); + await mkdir(dir, { recursive: true }); + const tempPath = join( + dir, + `.auto-router-state.${process.pid}.${randomUUID()}.tmp`, + ); + try { + await writeFile( + tempPath, + `${JSON.stringify({ models: this.state, classifications: this.classifications }, null, 2)}\n`, + "utf8", + ); + await rename(tempPath, statePath()); + } finally { + await rm(tempPath, { force: true }).catch(() => undefined); + } + } +} diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts new file mode 100644 index 0000000..517eea3 --- /dev/null +++ b/extensions/auto-router-quota.ts @@ -0,0 +1,491 @@ +import { execFile } from "node:child_process"; +import { readFile } from "node:fs/promises"; +import { homedir } from "node:os"; +import { join } from "node:path"; +import { promisify } from "node:util"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; +import type { QuotaReconciliationResult } from "./auto-router-health.js"; + +const execFileAsync = promisify(execFile); + +const FETCH_TIMEOUT_MS = 15_000; +const EXHAUSTED_UTILIZATION_PERCENT = 99.5; + +type FetchResult = { ok: true; data: unknown } | { ok: false }; + +/** + * A provider's reconciliation result. `default` applies to every configured model under the + * provider; `perModel` (keyed by `normalizeModelId(modelId)`) overrides it for models the + * provider reports on individually — e.g. Codex's per-model usage alongside its account-wide + * limit. Falling back to `default` for anything not in `perModel` keeps this correct even for + * models the provider doesn't break out individually. + */ +export type ProviderQuotaResult = { + default: QuotaReconciliationResult; + perModel?: Record; +}; + +/** Injectable for tests; defaults to the real network/filesystem/CLI. */ +export type QuotaFetchDependencies = { + fetchImpl: typeof fetch; + readCodexAccountId: () => Promise; + /** Runs the `mmx` CLI (MiniMax's own tool) and returns stdout, or `undefined` if it's missing, not logged in, or fails. */ + runMinimaxCli: (args: string[]) => Promise; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +/** Loosely matches a provider's own model label (e.g. "GPT-5.3-Codex-Spark") to a configured model id (e.g. "gpt-5.3-codex-spark"). */ +export function normalizeModelId(id: string): string { + return id.toLowerCase().replace(/[^a-z0-9]/g, ""); +} + +async function defaultReadCodexAccountId(): Promise { + try { + const authPath = join(homedir(), ".codex", "auth.json"); + const data: unknown = JSON.parse(await readFile(authPath, "utf8")); + if (!isRecord(data) || !isRecord(data.tokens)) return undefined; + const { account_id: accountId, accountId: camelAccountId } = data.tokens as Record; + return typeof accountId === "string" + ? accountId + : typeof camelAccountId === "string" + ? camelAccountId + : undefined; + } catch { + return undefined; + } +} + +async function defaultRunMinimaxCli(args: string[]): Promise { + try { + const { stdout } = await execFileAsync("mmx", args, { + timeout: FETCH_TIMEOUT_MS, + encoding: "utf8", + }); + return stdout; + } catch { + return undefined; + } +} + +export const defaultQuotaFetchDependencies: QuotaFetchDependencies = { + fetchImpl: fetch, + readCodexAccountId: defaultReadCodexAccountId, + runMinimaxCli: defaultRunMinimaxCli, +}; + +async function fetchJson( + url: string, + headers: Record, + fetchImpl: typeof fetch, +): Promise { + try { + const response = await fetchImpl(url, { + headers, + signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), + }); + if (!response.ok) return { ok: false }; + return { ok: true, data: await response.json() }; + } catch { + return { ok: false }; + } +} + +function numeric(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +function roundPercent(value: number): number { + return Math.round(value * 10) / 10; +} + +/** Accepts epoch seconds, epoch milliseconds, or an ISO date string. */ +function parseDateish(value: unknown): number | undefined { + if (typeof value === "number") return value > 10 ** 11 ? value : value * 1000; + if (typeof value === "string") { + const parsed = Date.parse(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +/** + * A raw Anthropic API key (`sk-ant-...`) has no subscription usage window to report; only + * `pi /login` OAuth subscription credentials do. + */ +function isDirectAnthropicApiKey(token: string): boolean { + return token.startsWith("sk-ant-"); +} + +async function fetchAnthropicQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const token = await modelRegistry.getApiKeyForProvider("anthropic"); + if (!token || isDirectAnthropicApiKey(token)) return undefined; + const result = await fetchJson( + "https://api.anthropic.com/api/oauth/usage", + { + Authorization: `Bearer ${token}`, + "anthropic-beta": "oauth-2025-04-20", + Accept: "application/json", + }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data)) return undefined; + + let mostUsed: { label: string; percent: number; window: Record } | undefined; + for (const [key, label] of [["five_hour", "5h"], ["seven_day", "7d"]] as const) { + const window = result.data[key]; + if (!isRecord(window)) continue; + const utilization = numeric(window.utilization); + if (utilization === undefined) continue; + if (!mostUsed || utilization > mostUsed.percent) mostUsed = { label, percent: utilization, window }; + } + if (!mostUsed) return { default: { exhausted: false } }; + const detail = `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used`; + if (mostUsed.percent >= EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: parseDateish(mostUsed.window.resets_at), detail } }; + } + return { default: { exhausted: false, detail } }; +} + +/** + * The API has been observed reporting a window's usage three different ways: "percent left" + * fields (converted to used%) or a direct "used%" field. Check all three rather than assuming one. + */ +function windowUsedPercent(window: Record): number | undefined { + const percentLeft = numeric(window.percent_left) ?? numeric(window.remaining_percent); + return percentLeft !== undefined ? 100 - percentLeft : numeric(window.used_percent); +} + +function windowResetsAt(window: Record): number | undefined { + return parseDateish(window.reset_at ?? window.reset_time_ms); +} + +/** e.g. 604800 -> "7d", 3600 -> "1h". Codex's own `limit_window_seconds` field, so this is the window's real duration, not a guess. */ +function secondsToLabel(seconds: number): string { + if (seconds % 86_400 === 0) return `${seconds / 86_400}d`; + if (seconds % 3600 === 0) return `${seconds / 3600}h`; + if (seconds % 60 === 0) return `${seconds / 60}m`; + return `${seconds}s`; +} + +/** Window length from a pair of epoch-ms timestamps, rounded to the nearest second. */ +function durationSeconds(startMs: unknown, endMs: unknown): number | undefined { + const start = numeric(startMs); + const end = numeric(endMs); + if (start === undefined || end === undefined || end <= start) return undefined; + return Math.round((end - start) / 1000); +} + +function windowLabel(window: Record): string | undefined { + const seconds = numeric(window.limit_window_seconds); + return seconds !== undefined && seconds > 0 ? secondsToLabel(seconds) : undefined; +} + +async function fetchCodexQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const token = await modelRegistry.getApiKeyForProvider("openai-codex"); + const accountId = await deps.readCodexAccountId(); + if (!token || !accountId) return undefined; + const result = await fetchJson( + "https://chatgpt.com/backend-api/wham/usage", + { + Authorization: `Bearer ${token}`, + "ChatGPT-Account-Id": accountId, + Accept: "application/json", + Origin: "https://chatgpt.com", + Referer: "https://chatgpt.com/", + }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data)) return undefined; + + const spendControl = result.data.spend_control; + if (isRecord(spendControl) && spendControl.reached === true) { + return { default: { exhausted: true, detail: "spend cap reached" } }; + } + + const rateLimit = result.data.rate_limit ?? result.data.rate_limits; + if (!isRecord(rateLimit)) return { default: { exhausted: false } }; + + const accountWindow = isRecord(rateLimit.primary_window ?? rateLimit.primary) + ? ((rateLimit.primary_window ?? rateLimit.primary) as Record) + : undefined; + const accountUsedPercent = accountWindow ? windowUsedPercent(accountWindow) : undefined; + // Codex's own `limit_window_seconds` gives the real window duration (e.g. 604800 = 7 days / + // weekly) rather than a vague "account" placeholder. + const accountLabel = accountWindow ? (windowLabel(accountWindow) ?? "account") : "account"; + const accountDetail = + accountUsedPercent !== undefined ? `${accountLabel} ${roundPercent(accountUsedPercent)}% used` : undefined; + const accountResetsAt = accountWindow ? windowResetsAt(accountWindow) : undefined; + + // Codex reports exhaustion account-wide, authoritatively, right on the rate_limit object + // itself - check it before falling back to inferring exhaustion from window percentages. + // (Verified directly against a real exhausted account: `{"allowed":false,"limit_reached":true, + // "primary_window":{"used_percent":100,...}}` at the top level, alongside a *healthy* + // per-model entry under `additional_rate_limits` for a specific model — and that model kept + // working normally. So the account-wide flag applies only to the "default" bucket, i.e. + // whichever configured models aren't separately metered below; a model with its own + // additional_rate_limits entry has an independent quota track that the account-wide flag + // does not override in either direction.) + const accountExhausted = + rateLimit.limit_reached === true || + rateLimit.allowed === false || + (accountUsedPercent !== undefined && accountUsedPercent >= EXHAUSTED_UTILIZATION_PERCENT); + + const defaultResult: QuotaReconciliationResult = accountExhausted + ? { exhausted: true, resetsAt: accountResetsAt, detail: accountDetail } + : { exhausted: false, detail: accountDetail }; + + // Independently-metered models from additional_rate_limits, matched to configured model ids + // by normalized label. Each one's own status is authoritative for that model — it neither + // inherits the account-wide block nor is protected by the account being otherwise healthy. + const perModel: Record = {}; + const additional = Array.isArray(result.data.additional_rate_limits) ? result.data.additional_rate_limits : []; + for (const entry of additional) { + if (!isRecord(entry) || typeof entry.limit_name !== "string") continue; + const entryRateLimit = entry.rate_limit; + if (!isRecord(entryRateLimit)) continue; + const window = isRecord(entryRateLimit.primary_window) ? entryRateLimit.primary_window : undefined; + const usedPercent = window ? windowUsedPercent(window) : undefined; + const modelExhausted = + entryRateLimit.limit_reached === true || + entryRateLimit.allowed === false || + (usedPercent !== undefined && usedPercent >= EXHAUSTED_UTILIZATION_PERCENT); + const modelLabel = window ? windowLabel(window) : undefined; + perModel[normalizeModelId(entry.limit_name)] = { + exhausted: modelExhausted, + resetsAt: modelExhausted && window ? windowResetsAt(window) : undefined, + detail: + usedPercent !== undefined + ? `${modelLabel ? `${modelLabel} ` : ""}${roundPercent(usedPercent)}% used` + : undefined, + }; + } + + return { default: defaultResult, perModel: Object.keys(perModel).length > 0 ? perModel : undefined }; +} + +/** `unit` is a time-unit enum (3=hour, 4=day, 6=week, 5=month observed); `number` is the count, e.g. unit 3 + number 5 = a 5-hour window. */ +function zaiWindowLabel(unit: unknown, count: unknown): string { + const n = numeric(count) ?? 1; + switch (unit) { + case 3: + return `${n}h`; + case 4: + return `${n}d`; + case 6: + return `${n * 7}d`; + case 5: + return "monthly"; + default: + return "usage"; + } +} + +async function fetchZaiQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const apiKey = await modelRegistry.getApiKeyForProvider("zai"); + if (!apiKey) return undefined; + const result = await fetchJson( + "https://api.z.ai/api/monitor/usage/quota/limit", + { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data)) return undefined; + + // Different plan tiers report different `type` values (observed: "TOKENS_LIMIT" on some + // accounts, "CREDIT_LIMIT" — with usage/currentValue/remaining alongside it — on others, + // e.g. a "pro" plan). Both carry a real `percentage` field with the same meaning, so key off + // that directly rather than an allowlist of type strings that isn't fully known. + const nested = isRecord(result.data.data) ? result.data.data : result.data; + const limits = Array.isArray(nested.limits) ? nested.limits : []; + let mostUsed: { label: string; percent: number; entry: Record } | undefined; + for (const entry of limits) { + if (!isRecord(entry)) continue; + const percentage = numeric(entry.percentage); + if (percentage === undefined) continue; + const label = zaiWindowLabel(entry.unit, entry.number); + if (!mostUsed || percentage > mostUsed.percent) mostUsed = { label, percent: percentage, entry }; + } + if (!mostUsed) return { default: { exhausted: false } }; + const detail = `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used`; + if (mostUsed.percent >= EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: parseDateish(mostUsed.entry.nextResetTime), detail } }; + } + return { default: { exhausted: false, detail } }; +} + +async function fetchKimiCodingQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const apiKey = await modelRegistry.getApiKeyForProvider("kimi-coding"); + if (!apiKey) return undefined; + const result = await fetchJson( + "https://api.kimi.com/coding/v1/usages", + { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data)) return undefined; + + const weekly = result.data.usage; + if (!isRecord(weekly)) return { default: { exhausted: false } }; + const limit = numeric(weekly.limit); + const used = numeric(weekly.used); + const detail = limit !== undefined && used !== undefined ? `${used}/${limit} this week` : undefined; + if (limit !== undefined && used !== undefined && limit > 0 && used >= limit) { + return { default: { exhausted: true, resetsAt: parseDateish(weekly.resetTime), detail } }; + } + return { default: { exhausted: false, detail } }; +} + +/** + * `GET /zen/go/v1/usage` with the same API key Pi already uses for inference — a real, clean + * JSON endpoint, not documented anywhere but discovered directly (no scraping, no separate + * cookie/workspace-id setup needed, unlike pi-quotas' HTML-scraping approach for this provider). + * Verified shape: `{"usage":{"rolling":{"status":"ok","percent":0,"resetsAt":"..."}, + * "weekly":{...},"monthly":{...}}}`. + */ +async function fetchOpenCodeGoQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const apiKey = await modelRegistry.getApiKeyForProvider("opencode-go"); + if (!apiKey) return undefined; + const result = await fetchJson( + "https://opencode.ai/zen/go/v1/usage", + { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data) || !isRecord(result.data.usage)) return undefined; + + const usage = result.data.usage; + const windows: { label: string; percent: number; window: Record }[] = []; + let mostUsed: { label: string; percent: number; window: Record } | undefined; + let blocked: { label: string; window: Record } | undefined; + for (const label of ["rolling", "weekly", "monthly"] as const) { + const window = usage[label]; + if (!isRecord(window)) continue; + if (!blocked && typeof window.status === "string" && window.status !== "ok") { + blocked = { label, window }; + } + const percent = numeric(window.percent); + if (percent === undefined) continue; + const entry = { label, percent, window }; + windows.push(entry); + if (!mostUsed || percent > mostUsed.percent) mostUsed = entry; + } + if (windows.length === 0 && !blocked) return { default: { exhausted: false } }; + // Report every window's usage together, not just whichever has the highest percentage - + // "rolling" (a short window, e.g. hourly) and the longer weekly/monthly windows are + // independent limits, and showing only one silently hides real usage against the others. + const detail = + windows.length > 0 + ? windows.map((w) => `${w.label} ${roundPercent(w.percent)}% used`).join(", ") + : undefined; + if (blocked) { + return { default: { exhausted: true, resetsAt: parseDateish(blocked.window.resetsAt), detail } }; + } + if (mostUsed && mostUsed.percent >= EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: parseDateish(mostUsed.window.resetsAt), detail } }; + } + return { default: { exhausted: false, detail } }; +} + +/** + * MiniMax has no documented HTTP quota endpoint, but its own `mmx` CLI does — `mmx --verbose` + * shows it calling `GET https://api.minimax.io/v1/token_plan/remains` with its own OAuth + * session (`mmx auth login`), separate from whatever credential Pi itself uses for inference. + * Shelling out to the CLI (rather than reading its private token cache directly) lets `mmx` + * own token refresh/expiry, and only degrades — never breaks — when `mmx` isn't installed or + * isn't logged in. + */ +async function fetchMinimaxQuota( + _modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const stdout = await deps.runMinimaxCli(["quota", "show", "--output", "json"]); + if (!stdout) return undefined; + let data: unknown; + try { + data = JSON.parse(stdout); + } catch { + return undefined; + } + if (!isRecord(data) || !Array.isArray(data.model_remains)) return undefined; + + const general = data.model_remains.find((entry) => isRecord(entry) && entry.model_name === "general"); + if (!isRecord(general)) return { default: { exhausted: false } }; + + // The API reports *remaining* percent (opposite convention from every other provider here, + // which all report *used* percent) - convert so /usage reads consistently across providers. + // The short window's actual length (observed 5h) comes from start_time/end_time rather than + // being hardcoded, since nothing else in the response names it. + const intervalRemaining = numeric(general.current_interval_remaining_percent); + const weeklyRemaining = numeric(general.current_weekly_remaining_percent); + const intervalSeconds = durationSeconds( + parseDateish(general.start_time), + parseDateish(general.end_time), + ); + const intervalLabel = intervalSeconds !== undefined ? secondsToLabel(intervalSeconds) : "interval"; + const detailParts: string[] = []; + if (intervalRemaining !== undefined) detailParts.push(`${intervalLabel} ${roundPercent(100 - intervalRemaining)}% used`); + if (weeklyRemaining !== undefined) detailParts.push(`weekly ${roundPercent(100 - weeklyRemaining)}% used`); + const detail = detailParts.length > 0 ? detailParts.join(", ") : undefined; + + // Through parseDateish, not a raw numeric() read: this API has been observed returning + // milliseconds, but nothing guarantees every account/response does, and parseDateish's + // seconds-vs-milliseconds heuristic is what every other provider's reset time already goes + // through here. + if (intervalRemaining !== undefined && intervalRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: parseDateish(general.end_time), detail } }; + } + if (weeklyRemaining !== undefined && weeklyRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: parseDateish(general.weekly_end_time), detail } }; + } + return { default: { exhausted: false, detail } }; +} + +/** + * Best-effort real quota reconciliation, keyed by Pi provider id. Providers without a known + * fetcher (arbitrary OpenAI-compatible custom providers, mostly) simply have no entry here — + * callers treat a missing/failed fetch as "no correction available", never as failure. + */ +export const QUOTA_FETCHERS: Record< + string, + ( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, + ) => Promise +> = { + anthropic: fetchAnthropicQuota, + "openai-codex": fetchCodexQuota, + zai: fetchZaiQuota, + "kimi-coding": fetchKimiCodingQuota, + minimax: fetchMinimaxQuota, + "opencode-go": fetchOpenCodeGoQuota, +}; + +/** Reconcile one provider's real quota state. Never throws; returns `undefined` when there's no known fetcher, no credentials, or the request failed. */ +export async function reconcileProviderQuota( + provider: string, + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies = defaultQuotaFetchDependencies, +): Promise { + const fetcher = QUOTA_FETCHERS[provider]; + if (!fetcher) return undefined; + try { + return await fetcher(modelRegistry, deps); + } catch { + return undefined; + } +} diff --git a/extensions/auto-router-settings.ts b/extensions/auto-router-settings.ts new file mode 100644 index 0000000..1afff77 --- /dev/null +++ b/extensions/auto-router-settings.ts @@ -0,0 +1,302 @@ +import { randomUUID } from "node:crypto"; +import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises"; +import { createRequire } from "node:module"; +import { dirname, join } from "node:path"; +import { getAgentDir } from "@earendil-works/pi-coding-agent"; + +const SETTINGS_KEY = "autoRouter"; + +/** Resolved at call time (not module load) so it honors a `PI_CODING_AGENT_DIR` override set after import. */ +function settingsPath(): string { + return join(getAgentDir(), "settings.json"); +} + +type ReleaseLock = () => Promise; +type LockSettingsFile = ( + path: string, + options: { + realpath: boolean; + retries: { retries: number; minTimeout: number; maxTimeout: number }; + }, +) => Promise; +const lockfile = createRequire(import.meta.url)("proper-lockfile") as { + lock: LockSettingsFile; +}; + +type SettingsWriteDependencies = { + lock: LockSettingsFile; + mkdir: typeof mkdir; + readFile: typeof readFile; + writeFile: typeof writeFile; + rename: typeof rename; + rm: typeof rm; + randomUUID: typeof randomUUID; +}; + +const defaultWriteDependencies: SettingsWriteDependencies = { + lock: lockfile.lock, + mkdir, + readFile, + writeFile, + rename, + rm, + randomUUID, +}; + +/** Pi thinking levels, ordered from lightest to heaviest. `medium` is the conventional default/anchor tier. */ +export const AUTO_ROUTER_EFFORT_ORDER = [ + "off", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max", +] as const; + +export type AutoRouterEffortLevel = (typeof AUTO_ROUTER_EFFORT_ORDER)[number]; + +export type AutoRouterModelRef = { + provider: string; + id: string; + /** + * Thinking level to actually apply when this model is used, if different from the tier it's + * listed under. The tier a model is configured in only controls which classified-complexity + * bucket routes to it and where it sits in the escalation order - it's a routing decision, not + * a claim about what reasoning effort suits that specific model. A model that only performs + * well at its own max setting, for instance, can be listed under `high` (so moderately-hard + * tasks reach it and it participates in escalation the same way) while still always being + * dispatched at `max` effort. Defaults to the tier's own name when omitted. + */ + effort?: AutoRouterEffortLevel; +}; + +export type AutoRouterTierConfig = { + /** Ordered list of models for this effort tier; first is preferred, later entries are failover. */ + models: AutoRouterModelRef[]; +}; + +export type AutoRouterSettings = { + efforts: Partial>; +}; + +export const EMPTY_AUTO_ROUTER_SETTINGS: AutoRouterSettings = { efforts: {} }; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function isEffortLevel(value: string): value is AutoRouterEffortLevel { + return (AUTO_ROUTER_EFFORT_ORDER as readonly string[]).includes(value); +} + +function parseModelRef(value: unknown): AutoRouterModelRef | undefined { + if (!isRecord(value)) return undefined; + const { provider, id, effort } = value; + if (typeof provider !== "string" || !provider.trim()) return undefined; + if (typeof id !== "string" || !id.trim()) return undefined; + const ref: AutoRouterModelRef = { provider: provider.trim(), id: id.trim() }; + if (typeof effort === "string" && isEffortLevel(effort)) ref.effort = effort; + return ref; +} + +/** Parse the raw `autoRouter` settings value, silently dropping malformed entries rather than throwing on hand-edited config. */ +export function parseAutoRouterSettings(value: unknown): AutoRouterSettings { + const efforts: AutoRouterSettings["efforts"] = {}; + if (!isRecord(value) || !isRecord(value.efforts)) return { efforts }; + for (const [key, tier] of Object.entries(value.efforts)) { + if (!isEffortLevel(key) || !isRecord(tier) || !Array.isArray(tier.models)) + continue; + const models = tier.models + .map(parseModelRef) + .filter((model): model is AutoRouterModelRef => model !== undefined); + if (models.length > 0) efforts[key] = { models }; + } + return { efforts }; +} + +/** Read the shared `autoRouter` global setting. Returns an empty config (Auto effectively unconfigured) on any read/parse failure. */ +export async function readAutoRouterSettings(): Promise { + try { + const root: unknown = JSON.parse(await readFile(settingsPath(), "utf8")); + if (!isRecord(root)) return EMPTY_AUTO_ROUTER_SETTINGS; + return parseAutoRouterSettings(root[SETTINGS_KEY]); + } catch { + return EMPTY_AUTO_ROUTER_SETTINGS; + } +} + +/** + * Lock, read, and rewrite the shared settings file. `mutate` receives the current parsed root + * (empty object on first write) and returns the new root to persist, or `undefined` to skip + * the write entirely (e.g. no change needed) while still safely releasing the lock. + */ +async function withLockedSettingsFile( + settingsPath: string, + mutate: (root: Record) => Record | undefined, + dependencies: Partial = {}, +): Promise { + const io = { ...defaultWriteDependencies, ...dependencies }; + const settingsDir = dirname(settingsPath); + await io.mkdir(settingsDir, { recursive: true }); + const release = await io.lock(settingsPath, { + realpath: false, + retries: { retries: 9, minTimeout: 20, maxTimeout: 20 }, + }); + try { + let root: Record = {}; + try { + const parsed: unknown = JSON.parse( + await io.readFile(settingsPath, "utf8"), + ); + if (isRecord(parsed)) root = parsed; + } catch (error) { + if ( + error && + typeof error === "object" && + "code" in error && + error.code === "ENOENT" + ) { + // First global setting write. + } else { + throw new Error( + `Could not read ${settingsPath}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } + const next = mutate(root); + if (!next) return; + const tempPath = join( + settingsDir, + `.settings.${process.pid}.${io.randomUUID()}.tmp`, + ); + try { + await io.writeFile(tempPath, `${JSON.stringify(next, null, 2)}\n`, { + mode: 0o600, + flag: "wx", + }); + await io.rename(tempPath, settingsPath); + } finally { + await io.rm(tempPath, { force: true }).catch(() => undefined); + } + } finally { + await release(); + } +} + +/** Persist the `autoRouter` setting without dropping keys owned by Pi or other extensions. */ +export async function writeAutoRouterSettingsFile( + settingsPath: string, + settings: AutoRouterSettings, + dependencies: Partial = {}, +): Promise { + await withLockedSettingsFile( + settingsPath, + (root) => ({ ...root, [SETTINGS_KEY]: settings }), + dependencies, + ); +} + +/** Persist the package-specific global `autoRouter` setting without dropping unknown keys. */ +export async function writeAutoRouterSettings( + settings: AutoRouterSettings, +): Promise { + await writeAutoRouterSettingsFile(settingsPath(), settings); +} + +/** + * Pattern that matches every virtual "Auto" model in `/model`'s scoping patterns: the plain + * adaptive one and every "Auto ()" pinned-tier entry, both under the `auto` provider. + * `enabledModels` patterns are matched with minimatch, so a single glob covers all of them. + */ +export const AUTO_MODEL_SCOPE_PATTERN = "auto/*"; + +/** + * Best-effort: Pi's `/model` picker defaults to showing only `enabledModels`-scoped models + * when that setting is non-empty, hiding everything else (including our own registered "Auto" + * entries) behind a manual Tab toggle to "all". If the user has scoping configured, make sure it + * includes a pattern matching Auto so it's visible by default. No-op when scoping isn't + * configured at all (everything is already visible) or already includes a matching pattern. + */ +export async function ensureAutoModelScoped( + settingsPath: string, + dependencies: Partial = {}, +): Promise { + await withLockedSettingsFile( + settingsPath, + (root) => { + const current = root.enabledModels; + if (!Array.isArray(current) || current.length === 0) return undefined; + if (current.includes(AUTO_MODEL_SCOPE_PATTERN)) return undefined; + return { ...root, enabledModels: [...current, AUTO_MODEL_SCOPE_PATTERN] }; + }, + dependencies, + ); +} + +/** `ensureAutoModelScoped` against the real global settings file. */ +export async function ensureAutoModelScopedInGlobalSettings(): Promise { + await ensureAutoModelScoped(settingsPath()); +} + +/** + * Resolve a classified effort level to the concrete tier to route to: walk the ordered + * level list from `level` toward `medium`, using the first tier with configured models. + * `medium` is the last resort even if unconfigured (callers should validate config has + * at least one tier before routing). + */ +export function resolveEffortTier( + settings: AutoRouterSettings, + level: AutoRouterEffortLevel, +): AutoRouterEffortLevel { + const levelIndex = AUTO_ROUTER_EFFORT_ORDER.indexOf(level); + const mediumIndex = AUTO_ROUTER_EFFORT_ORDER.indexOf("medium"); + const step = levelIndex <= mediumIndex ? 1 : -1; + for ( + let index = levelIndex; + step > 0 ? index <= mediumIndex : index >= mediumIndex; + index += step + ) { + const candidate = AUTO_ROUTER_EFFORT_ORDER[index]; + if (settings.efforts[candidate]?.models.length) return candidate; + } + return "medium"; +} + +/** + * Tiers to try, in escalation order, when every model in `fromTier` is unhealthy: strictly + * upward through configured tiers only (`fromTier` excluded), e.g. medium -> high -> xhigh -> max. + */ +export function escalationTiers( + settings: AutoRouterSettings, + fromTier: AutoRouterEffortLevel, +): AutoRouterEffortLevel[] { + const fromIndex = AUTO_ROUTER_EFFORT_ORDER.indexOf(fromTier); + const tiers: AutoRouterEffortLevel[] = []; + for ( + let index = fromIndex + 1; + index < AUTO_ROUTER_EFFORT_ORDER.length; + index++ + ) { + const candidate = AUTO_ROUTER_EFFORT_ORDER[index]; + if (settings.efforts[candidate]?.models.length) tiers.push(candidate); + } + return tiers; +} + +/** All models referenced anywhere in the config, in tier order then list order, deduplicated by provider/id. */ +export function allConfiguredModels( + settings: AutoRouterSettings, +): AutoRouterModelRef[] { + const seen = new Set(); + const models: AutoRouterModelRef[] = []; + for (const level of AUTO_ROUTER_EFFORT_ORDER) { + for (const model of settings.efforts[level]?.models ?? []) { + const key = `${model.provider}/${model.id}`; + if (seen.has(key)) continue; + seen.add(key); + models.push(model); + } + } + return models; +} diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts new file mode 100644 index 0000000..c1e25db --- /dev/null +++ b/extensions/auto-router.ts @@ -0,0 +1,765 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import { + DynamicBorder, + type ExtensionAPI, + type ExtensionCommandContext, + type ExtensionContext, + getMarkdownTheme, + type ModelRegistry, + type Theme, +} from "@earendil-works/pi-coding-agent"; +import { Container, Markdown, matchesKey, Text } from "@earendil-works/pi-tui"; +import { classifyTurnComplexity } from "./auto-router-classify.js"; +import { + AutoRouterHealthStore, + type ClassificationLogEntry, + type ModelHealthEntry, + type ModelIdentity, + modelKey, +} from "./auto-router-health.js"; +import { normalizeModelId, reconcileProviderQuota } from "./auto-router-quota.js"; +import { + AUTO_ROUTER_EFFORT_ORDER, + type AutoRouterEffortLevel, + type AutoRouterModelRef, + type AutoRouterSettings, + allConfiguredModels, + ensureAutoModelScopedInGlobalSettings, + escalationTiers, + readAutoRouterSettings, + resolveEffortTier, +} from "./auto-router-settings.js"; +import { + FOOTER_CONTRIBUTION_EVENT, + type FooterContribution, +} from "./footer-events.js"; + +const AUTO_PROVIDER_ID = "auto"; +const AUTO_MODEL_ID = "auto"; +const AUTO_ACTIVE_ENTRY_TYPE = "vessup:auto-router:active"; +const FOOTER_KEY = "auto-router"; +const PINNED_MODEL_PREFIX = `${AUTO_MODEL_ID}-`; + +/** Model id for the `/model` entry that pins Auto to `tier` instead of classifying each turn. */ +function pinnedModelId(tier: AutoRouterEffortLevel): string { + return `${PINNED_MODEL_PREFIX}${tier}`; +} + +/** The reverse of `pinnedModelId`: which tier (if any) a selected auto-provider model id pins to. `undefined` for the plain adaptive "auto" id. */ +function tierFromModelId(id: string): AutoRouterEffortLevel | undefined { + if (!id.startsWith(PINNED_MODEL_PREFIX)) return undefined; + const candidate = id.slice(PINNED_MODEL_PREFIX.length); + return (AUTO_ROUTER_EFFORT_ORDER as readonly string[]).includes(candidate) + ? (candidate as AutoRouterEffortLevel) + : undefined; +} + +function numeric(value: unknown): number { + return typeof value === "number" && Number.isFinite(value) ? value : 0; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +/** + * A `stopReason: "error"` assistant message carries no HTTP status — only a human-readable + * `errorMessage` (Pi's own normalized wording, e.g. "Codex error: The usage limit has been + * reached"). By the time an extension sees this, Pi's own agent-level retry has already given + * up against this exact model/provider, so any such error is treated as at least as serious as + * a rate limit (immediate cooldown) rather than requiring several occurrences first. + */ +function inferFailureStatus(errorMessage: string | undefined): number { + const text = (errorMessage ?? "").toLowerCase(); + if (/unauthoriz|authentication|invalid api key|forbidden/.test(text)) return 401; + return 429; +} + +/** Resolve config model refs to real, currently-usable `Model` objects (auth configured), preserving order. */ +function resolveAvailableModels( + modelRegistry: ModelRegistry, + refs: AutoRouterModelRef[], +): Model[] { + const models: Model[] = []; + for (const ref of refs) { + const model = modelRegistry.find(ref.provider, ref.id); + if (model && modelRegistry.hasConfiguredAuth(model)) models.push(model); + } + return models; +} + +type AutoState = { active: boolean; pinnedTier: AutoRouterEffortLevel | undefined }; + +function restoreAutoState(ctx: ExtensionContext): AutoState { + const entries = ctx.sessionManager.getEntries(); + for (let index = entries.length - 1; index >= 0; index--) { + const entry = entries[index]; + if ( + entry.type === "custom" && + entry.customType === AUTO_ACTIVE_ENTRY_TYPE + ) { + const data = entry.data; + if (!isRecord(data) || data.enabled !== true) { + return { active: false, pinnedTier: undefined }; + } + const pinnedTier = + typeof data.pinnedTier === "string" && + (AUTO_ROUTER_EFFORT_ORDER as readonly string[]).includes(data.pinnedTier) + ? (data.pinnedTier as AutoRouterEffortLevel) + : undefined; + return { active: true, pinnedTier }; + } + } + return { active: false, pinnedTier: undefined }; +} + +function formatTier(tier: AutoRouterEffortLevel): string { + return tier; +} + +function footerBadge(tier: AutoRouterEffortLevel | undefined): string { + return tier ? `🔀 Auto (${formatTier(tier)})` : "🔀 Auto"; +} + +/** A registered-but-inert `/model` entry: never actually dispatched to, since `before_agent_start` always swaps in a real routed model first. */ +function placeholderModel(id: string, name: string) { + return { + id, + name, + reasoning: false, + input: ["text"] as ("text" | "image")[], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128_000, + maxTokens: 4096, + }; +} + +export default async function autoRouter(pi: ExtensionAPI): Promise { + // Read once at startup (not re-read later) to decide which per-tier "Auto ()" picker + // entries to register - config changes after this need a process restart to pick up new + // tiers, same as any other extension code change. + const initialSettings = await readAutoRouterSettings(); + const pinnedTierModels = AUTO_ROUTER_EFFORT_ORDER.filter( + (tier) => (initialSettings.efforts[tier]?.models.length ?? 0) > 0, + ).map((tier) => placeholderModel(pinnedModelId(tier), `Auto (${tier})`)); + + pi.registerProvider(AUTO_PROVIDER_ID, { + name: "Auto", + baseUrl: "http://127.0.0.1:0", + apiKey: "auto-router", + api: "openai-completions", + models: [ + placeholderModel(AUTO_MODEL_ID, "Auto (auto)"), + ...pinnedTierModels, + ], + }); + + let currentSessionId: string | undefined; + let autoActive = false; + /** Set when the user picked a specific "Auto ()" entry rather than plain "Auto": + * every turn routes within that tier directly, skipping classification entirely. */ + let pinnedTier: AutoRouterEffortLevel | undefined; + let routingInFlight = false; + /** The last thinking level actually applied to a routed model - not necessarily the tier + * name it's classified under, since a model's `effort` override can differ from its tier. */ + let lastKnownEffort: AutoRouterEffortLevel | undefined; + const healthStore = new AutoRouterHealthStore(); + + function publishFooter(effort: AutoRouterEffortLevel | undefined): void { + if (!currentSessionId) return; + pi.events.emit(FOOTER_CONTRIBUTION_EVENT, { + sessionId: currentSessionId, + key: FOOTER_KEY, + identitySuffix: (theme: Theme) => theme.fg("accent", footerBadge(effort)), + } satisfies FooterContribution); + } + + /** The specific auto-provider `/model` id currently selected: the pinned tier's, or the plain adaptive one. */ + function currentAutoModelId(): string { + return pinnedTier ? pinnedModelId(pinnedTier) : AUTO_MODEL_ID; + } + + /** Swap `ctx.model` back to the inert Auto placeholder (whichever one is selected - adaptive or a pinned tier) once a turn is fully done, so `/model` keeps showing it selected instead of whichever real model just handled the turn. */ + async function revertToAutoPlaceholder( + pi: ExtensionAPI, + ctx: ExtensionContext, + ): Promise { + const placeholder = ctx.modelRegistry.find( + AUTO_PROVIDER_ID, + currentAutoModelId(), + ); + if (!placeholder) return; + routingInFlight = true; + try { + await pi.setModel(placeholder); + } finally { + routingInFlight = false; + } + } + + function clearFooter(): void { + if (!currentSessionId) return; + pi.events.emit(FOOTER_CONTRIBUTION_EVENT, { + sessionId: currentSessionId, + key: FOOTER_KEY, + remove: true, + } satisfies FooterContribution); + } + + /** + * The thinking level to actually dispatch a model at: its own configured `effort` override, + * or the tier name itself when it doesn't have one. `refs` is whichever tier's config list + * `model` was resolved from, so the matching entry (and its override, if any) can be found. + */ + function resolveEffort( + refs: AutoRouterModelRef[], + model: ModelIdentity, + tier: AutoRouterEffortLevel, + ): AutoRouterEffortLevel { + const ref = refs.find( + (candidate) => + candidate.provider === model.provider && candidate.id === model.id, + ); + return ref?.effort ?? tier; + } + + /** Pick the best available (resolved + healthy) model for `tier`, escalating to higher configured tiers when everything in `tier` is unhealthy, then falling back to the first available model anywhere as a last resort. */ + function pickForTier( + ctx: ExtensionContext, + settings: AutoRouterSettings, + tier: AutoRouterEffortLevel, + ): + | { model: Model; tier: AutoRouterEffortLevel; effort: AutoRouterEffortLevel } + | undefined { + for (const candidateTier of [tier, ...escalationTiers(settings, tier)]) { + const refs = settings.efforts[candidateTier]?.models ?? []; + const available = resolveAvailableModels(ctx.modelRegistry, refs); + const healthy = healthStore.pickHealthy(available); + if (healthy) { + const model = available.find( + (candidate) => + candidate.id === healthy.id && + candidate.provider === healthy.provider, + ); + if (model) { + return { + model, + tier: candidateTier, + effort: resolveEffort(refs, model, candidateTier), + }; + } + } + } + // Last resort: nothing healthy anywhere. Use the first resolvable model still configured + // for `tier` itself, or failing that the first resolvable model in any other configured + // tier - labeled with whichever tier the picked model actually belongs to, not blindly + // `tier`, since mislabeling it would apply the wrong thinking level to the model in use. + const ownRefs = settings.efforts[tier]?.models; + let fallback = ownRefs + ? resolveAvailableModels(ctx.modelRegistry, ownRefs)[0] + : undefined; + let fallbackTier = tier; + let fallbackRefs = ownRefs ?? []; + if (!fallback) { + for (const candidateTier of AUTO_ROUTER_EFFORT_ORDER) { + const refs = settings.efforts[candidateTier]?.models; + if (!refs) continue; + const candidate = resolveAvailableModels(ctx.modelRegistry, refs)[0]; + if (candidate) { + fallback = candidate; + fallbackTier = candidateTier; + fallbackRefs = refs; + break; + } + } + } + if (fallback) { + if (ctx.hasUI) { + ctx.ui.notify( + `Auto: every configured ${formatTier(tier)} model looks unavailable; using ${fallback.provider}/${fallback.id} anyway. Check /usage.`, + "warning", + ); + } + return { + model: fallback, + tier: fallbackTier, + effort: resolveEffort(fallbackRefs, fallback, fallbackTier), + }; + } + // Nothing configured for Auto at all (or resolvable) - the "auto" placeholder has no real + // backend, so leaving it selected here would send the actual request to it and fail with a + // connection error instead of a clear message. Fall back to any authenticated model in the + // whole catalog rather than ever letting a turn run against the placeholder. + const anyModel = ctx.modelRegistry + .getAvailable() + .find((candidate) => candidate.provider !== AUTO_PROVIDER_ID); + if (anyModel) { + if (ctx.hasUI) { + ctx.ui.notify( + `Auto has no configured models. Add an \`autoRouter\` entry to ~/.pi/agent/settings.json — falling back to ${anyModel.provider}/${anyModel.id} for now.`, + "warning", + ); + } + return { model: anyModel, tier, effort: tier }; + } + return undefined; + } + + async function applyRouting( + pi: ExtensionAPI, + ctx: ExtensionContext, + model: Model, + effort: AutoRouterEffortLevel, + ): Promise { + routingInFlight = true; + try { + const success = await pi.setModel(model); + if (!success) { + if (ctx.hasUI) { + ctx.ui.notify( + `Auto: no credentials configured for ${model.provider}/${model.id}`, + "warning", + ); + } + return; + } + await pi.setThinkingLevel(effort); + } finally { + routingInFlight = false; + } + lastKnownEffort = effort; + publishFooter(effort); + } + + async function routeForPrompt( + pi: ExtensionAPI, + ctx: ExtensionContext, + prompt: string, + hasImages: boolean, + ): Promise { + const settings = await readAutoRouterSettings(); + + let tier: AutoRouterEffortLevel; + let level: string; + let classifierReply: string; + if (pinnedTier) { + // The user picked "Auto ()" specifically to skip classification and always route + // within that tier - pickForTier's own escalation/fallback still applies if it's unhealthy. + tier = pinnedTier; + level = "(pinned)"; + classifierReply = `pinned to ${pinnedTier} - not classified`; + } else { + const classifierPool = resolveAvailableModels( + ctx.modelRegistry, + settings.efforts.medium?.models ?? allConfiguredModels(settings), + ); + const classifierRef = healthStore.pickHealthy(classifierPool); + const classifierModel = classifierRef + ? classifierPool.find( + (m) => + m.id === classifierRef.id && m.provider === classifierRef.provider, + ) + : undefined; + + let classifiedLevel: AutoRouterEffortLevel = "medium"; + classifierReply = "(no classifier available)"; + if (classifierModel) { + const result = await classifyTurnComplexity( + ctx.modelRegistry, + classifierModel, + prompt, + hasImages, + ); + classifiedLevel = result.level; + classifierReply = result.reply; + if (result.usage) { + healthStore.recordSuccess(modelKey(classifierModel), result.usage); + } + } + level = classifiedLevel; + tier = resolveEffortTier(settings, classifiedLevel); + } + + const picked = pickForTier(ctx, settings, tier); + if (!picked) { + if (ctx.hasUI) { + ctx.ui.notify( + "Auto has no configured models yet. Add an `autoRouter` entry to ~/.pi/agent/settings.json.", + "warning", + ); + } + return; + } + // Persisted so `/usage` can show what the classifier actually said - the classify call + // itself is otherwise a throwaway completion whose result is discarded after parsing, which + // made a prior misrouting report impossible to actually verify against real evidence. + healthStore.recordClassification({ + reply: classifierReply, + level, + tier: picked.tier, + effort: picked.effort, + model: picked.model, + }); + await applyRouting(pi, ctx, picked.model, picked.effort); + } + + async function reconcileAllProviders( + modelRegistry: ModelRegistry, + settings: AutoRouterSettings, + ): Promise { + const models = allConfiguredModels(settings); + const providers = Array.from( + new Set(models.map((model) => model.provider)), + ); + await Promise.all( + providers.map(async (provider) => { + const result = await reconcileProviderQuota(provider, modelRegistry); + if (!result) return; + for (const model of models) { + if (model.provider !== provider) continue; + const specific = result.perModel?.[normalizeModelId(model.id)]; + healthStore.applyQuotaResult(modelKey(model), specific ?? result.default); + } + }), + ); + } + + pi.on("model_select", (event, _ctx) => { + if (routingInFlight) return; + if (event.model.provider === AUTO_PROVIDER_ID) { + autoActive = true; + pinnedTier = tierFromModelId(event.model.id); + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true, pinnedTier }); + // A pinned tier is worth showing immediately (the user explicitly locked to it), even + // before any turn has routed; the adaptive "auto" entry still waits for a real turn. + publishFooter(pinnedTier ?? lastKnownEffort); + return; + } + if (event.source !== "restore" && autoActive) { + autoActive = false; + pinnedTier = undefined; + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: false }); + lastKnownEffort = undefined; + clearFooter(); + } + }); + + pi.on("session_start", async (_event, ctx) => { + currentSessionId = ctx.sessionManager.getSessionId(); + routingInFlight = false; + lastKnownEffort = undefined; + // Reuse the single instance rather than replacing it: a stale instance's pending + // debounced-save timer would otherwise still fire independently and could overwrite + // this reload's freshly-loaded state on disk with the old in-memory data. + await healthStore.load(); + // A session can arrive with Auto active two different ways: a persisted entry from an + // earlier explicit `/model` pick (`restoreAutoState`), or `ctx.model` already being one of + // Auto's placeholders because the user set `defaultProvider`/`defaultModel` to "auto" (or a + // pinned "auto-" id) globally - a brand-new session in that case has no entries yet, + // so `restoreAutoState` alone would miss it and leave every turn dispatching straight at + // the placeholder's dead URL. + const restored = restoreAutoState(ctx); + const modelIsAuto = ctx.model?.provider === AUTO_PROVIDER_ID; + autoActive = restored.active || modelIsAuto; + pinnedTier = + restored.pinnedTier ?? + (modelIsAuto && ctx.model ? tierFromModelId(ctx.model.id) : undefined); + if (autoActive) { + if (ctx.model && ctx.model.provider !== AUTO_PROVIDER_ID) { + // Restored mid-turn (e.g. an interrupted process, before agent_settled could + // revert it). Normalize back to the placeholder so /model shows Auto again. + lastKnownEffort = ctx.thinkingLevel; + await revertToAutoPlaceholder(pi, ctx); + } + publishFooter(pinnedTier ?? lastKnownEffort); + } + const settings = await readAutoRouterSettings(); + void reconcileAllProviders(ctx.modelRegistry, settings); + void ensureAutoModelScopedInGlobalSettings().catch(() => undefined); + }); + + pi.on("agent_settled", async (_event, ctx) => { + if (!autoActive) return; + if (!ctx.model || ctx.model.provider === AUTO_PROVIDER_ID) return; + await revertToAutoPlaceholder(pi, ctx); + }); + + pi.on("before_agent_start", async (event, ctx) => { + if (!autoActive) { + // `autoActive` is bookkeeping derived from model_select/session_start events, and every + // path that's supposed to keep it in sync with reality is a separate thing to get right - + // exactly the kind of thing that's easy to miss one case of (as happened with a brand-new + // session whose defaultModel is "auto" in settings). But whether a request is about to be + // sent against the inert placeholder is directly observable right here, right before + // dispatch: if `ctx.model` already *is* one of Auto's placeholders, sending this turn as-is + // is guaranteed to fail with a bare connection error no matter why our own flag says + // inactive. So treat that as authoritative and self-heal instead of trusting the flag. + if (ctx.model?.provider !== AUTO_PROVIDER_ID) return; + autoActive = true; + pinnedTier = tierFromModelId(ctx.model.id); + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true, pinnedTier }); + } + await routeForPrompt(pi, ctx, event.prompt, Boolean(event.images?.length)); + }); + + // Health/usage tracking isn't limited to turns Auto itself routed: any turn against a model + // that's *configured* somewhere in autoRouter (picked manually from /model, or left over from + // before Auto was engaged) is just as real a signal for future routing decisions and /usage, + // so it's tracked the same way regardless of who selected the model. + async function trackedModel(model: ModelIdentity | undefined): Promise { + if (!model || model.provider === AUTO_PROVIDER_ID) return undefined; + const settings = await readAutoRouterSettings(); + const configured = allConfiguredModels(settings).some( + (candidate) => candidate.provider === model.provider && candidate.id === model.id, + ); + return configured ? model : undefined; + } + + pi.on("after_provider_response", async (event, ctx) => { + if (event.status >= 200 && event.status < 300) return; + const model = await trackedModel(ctx.model); + if (!model) return; + healthStore.recordFailure(modelKey(model), event.status, event.headers); + }); + + pi.on("message_end", async (event, ctx) => { + if (event.message.role !== "assistant") return; + const message = event.message; + if (message.stopReason === "aborted") return; // user-cancelled, not a provider health signal + const model = await trackedModel(ctx.model); + if (!model) return; + if (message.stopReason === "error") { + healthStore.recordFailure(modelKey(model), inferFailureStatus(message.errorMessage), undefined); + return; + } + const usage = isRecord(message) ? message.usage : undefined; + healthStore.recordSuccess(modelKey(model), { + input: numeric(isRecord(usage) ? usage.input : undefined), + output: numeric(isRecord(usage) ? usage.output : undefined), + cost: numeric( + isRecord(usage) && isRecord(usage.cost) ? usage.cost.total : undefined, + ), + }); + }); + + pi.on("session_shutdown", () => { + void healthStore.flush(); + currentSessionId = undefined; + autoActive = false; + }); + + pi.registerCommand("usage", { + description: "Show Auto router health/usage for every configured model", + handler: async (_args, ctx: ExtensionCommandContext) => { + const settings = await readAutoRouterSettings(); + if (Object.keys(settings.efforts).length === 0) { + if (ctx.hasUI) { + ctx.ui.notify( + "Auto has no configured models. Add an `autoRouter` entry to ~/.pi/agent/settings.json.", + "info", + ); + } + return; + } + await reconcileAllProviders(ctx.modelRegistry, settings); + const rows = buildUsageRows(settings, healthStore); + const classifications = healthStore.getClassifications(); + if (ctx.mode === "tui") { + await showUsageDashboard(rows, classifications, ctx); + } else if (ctx.hasUI) { + ctx.ui.notify(formatUsagePlainText(rows, classifications), "info"); + } + }, + }); +} + +type UsageRow = { + tier: AutoRouterEffortLevel; + model: AutoRouterModelRef; + entry: ModelHealthEntry | undefined; +}; + +function buildUsageRows( + settings: AutoRouterSettings, + healthStore: AutoRouterHealthStore, +): UsageRow[] { + const rows: UsageRow[] = []; + for (const tier of AUTO_ROUTER_EFFORT_ORDER) { + for (const model of settings.efforts[tier]?.models ?? []) { + rows.push({ tier, model, entry: healthStore.getEntry(modelKey(model)) }); + } + } + return rows; +} + +function formatTokenCount(count: number): string { + if (count < 1_000) return String(count); + if (count < 1_000_000) return `${(count / 1_000).toFixed(1)}k`; + return `${(count / 1_000_000).toFixed(1)}M`; +} + +/** `4557m` is meaningless at a glance; scale to hours/days once a cooldown is that long. */ +function formatDuration(ms: number): string { + const minutes = Math.max(1, Math.round(ms / 60_000)); + if (minutes < 60) return `${minutes}m`; + const hours = Math.round(minutes / 60); + if (hours < 48) return `${hours}h`; + return `${Math.round(hours / 24)}d`; +} + +function rowStatus(entry: ModelHealthEntry | undefined, now: number): string { + if (!entry) return "unused"; + if (entry.cooldownUntil && entry.cooldownUntil > now) { + const cause = entry.lastError ? ` (${entry.lastError.status})` : ""; + return `cooldown${cause} ~${formatDuration(entry.cooldownUntil - now)}`; + } + return "healthy"; +} + +/** Real usage from the provider's own quota API, when reconciliation has run for this model — separate from (and often more accurate than) this router's own request/token counters, which only see traffic this Pi installation itself made against the model. */ +function verifiedUsageText(entry: ModelHealthEntry | undefined): string { + if (!entry?.verifiedAt) return "—"; + return entry.verifiedDetail ?? "verified, no detail"; +} + +function rowLine(row: UsageRow, now: number): string { + const entry = row.entry; + const status = rowStatus(entry, now); + const requests = entry?.totals.requests ?? 0; + const tokens = formatTokenCount( + (entry?.totals.input ?? 0) + (entry?.totals.output ?? 0), + ); + return `${row.model.provider}/${row.model.id} — ${status} · ${verifiedUsageText(entry)} · ${requests} req · ${tokens} tok`; +} + +/** How many recent classification decisions `/usage` shows, most recent first. */ +const CLASSIFICATION_DISPLAY_LIMIT = 5; + +function formatRelativeTime(at: number, now: number): string { + const delta = now - at; + return delta < 60_000 ? "just now" : `${formatDuration(delta)} ago`; +} + +/** + * Neutralizes both pipes and line breaks for a Markdown table cell. Every current caller + * already gets pre-collapsed text from `truncateForLog`, but this stays a defense of its own + * rather than relying on that - a raw multi-line value would otherwise terminate the row early + * and break the rest of the table, not just that one cell. + */ +export function escapeTableCell(text: string): string { + return text.replace(/\s*[\r\n]+\s*/g, " ").replace(/\|/g, "\\|"); +} + +function recentClassifications( + entries: readonly ClassificationLogEntry[], +): ClassificationLogEntry[] { + return entries.slice(-CLASSIFICATION_DISPLAY_LIMIT).reverse(); +} + +function formatUsagePlainText( + rows: UsageRow[], + classifications: readonly ClassificationLogEntry[], +): string { + if (rows.length === 0) return "Auto: no models configured."; + const now = Date.now(); + const byTier = new Map(); + for (const row of rows) { + const list = byTier.get(row.tier) ?? []; + list.push(row); + byTier.set(row.tier, list); + } + const lines: string[] = []; + for (const [tier, tierRows] of byTier) { + lines.push( + `${tier}: ${tierRows.map((row) => rowLine(row, now)).join(" | ")}`, + ); + } + const recent = recentClassifications(classifications); + if (recent.length > 0) { + lines.push("", "Recent classifications:"); + for (const entry of recent) { + const effortSuffix = + entry.effort !== entry.tier ? ` at ${entry.effort} effort` : ""; + lines.push( + ` ${formatRelativeTime(entry.timestamp, now)}: said "${entry.reply}" → ${entry.level}, routed to ${entry.tier}${effortSuffix} (${entry.model.provider}/${entry.model.id})`, + ); + } + } + return lines.join("\n"); +} + +function formatUsageMarkdown( + rows: UsageRow[], + classifications: readonly ClassificationLogEntry[], +): string { + if (rows.length === 0) return "No models configured."; + const now = Date.now(); + const lines = [ + "| Tier | Model | Status | Verified usage | Observed req | Observed tokens | Observed cost |", + "|---|---|---|---|---|---|---|", + ]; + for (const row of rows) { + const entry = row.entry; + const cost = entry ? `$${entry.totals.cost.toFixed(3)}` : "$0.000"; + const tokens = formatTokenCount( + (entry?.totals.input ?? 0) + (entry?.totals.output ?? 0), + ); + lines.push( + `| ${row.tier} | ${row.model.provider}/${row.model.id} | ${rowStatus(entry, now)} | ${verifiedUsageText(entry)} | ${entry?.totals.requests ?? 0} | ${tokens} | ${cost} |`, + ); + } + lines.push( + "", + "_Verified usage comes from the provider's own quota API, where available. Observed req/tokens/cost count every turn this Pi installation has run against that model since Auto started tracking it — whether Auto routed there or it was picked manually from `/model` — but not usage from other sessions/machines/tools or from before that; that's exactly what verified usage is for._", + ); + const recent = recentClassifications(classifications); + if (recent.length > 0) { + lines.push( + "", + "### Recent classifications", + "_What the classifier actually said, so a routing decision that looks wrong can be checked against real evidence instead of guessed at._", + "", + "| When | Said | Level | Tier used | Effort applied | Model |", + "|---|---|---|---|---|---|", + ); + for (const entry of recent) { + lines.push( + `| ${formatRelativeTime(entry.timestamp, now)} | ${escapeTableCell(entry.reply)} | ${entry.level} | ${entry.tier} | ${entry.effort} | ${entry.model.provider}/${entry.model.id} |`, + ); + } + } + return lines.join("\n"); +} + +async function showUsageDashboard( + rows: UsageRow[], + classifications: readonly ClassificationLogEntry[], + ctx: ExtensionCommandContext, +): Promise { + await ctx.ui.custom((_tui, theme, _kb, done) => { + const container = new Container(); + const border = new DynamicBorder((s: string) => theme.fg("accent", s)); + const mdTheme = getMarkdownTheme(); + + container.addChild(border); + container.addChild( + new Text(theme.fg("accent", theme.bold("Auto Router Usage")), 1, 0), + ); + container.addChild( + new Markdown(formatUsageMarkdown(rows, classifications), 1, 1, mdTheme), + ); + container.addChild( + new Text(theme.fg("dim", "Press Enter or Esc to close"), 1, 0), + ); + container.addChild(border); + + return { + render: (width: number) => container.render(width), + invalidate: () => container.invalidate(), + handleInput: (data: string) => { + if (matchesKey(data, "enter") || matchesKey(data, "escape")) + done(undefined); + }, + }; + }); +} diff --git a/package.json b/package.json index fcf7988..178073e 100644 --- a/package.json +++ b/package.json @@ -29,7 +29,8 @@ "./extensions/subagents.ts", "./extensions/worktree.ts", "./extensions/web-sessions.ts", - "./extensions/delete-session.ts" + "./extensions/delete-session.ts", + "./extensions/auto-router.ts" ], "prompts": [ "./prompts/address-pr.md" @@ -39,11 +40,15 @@ ] }, "peerDependencies": { + "@earendil-works/pi-ai": "*", "@earendil-works/pi-coding-agent": "0.84.1", "@earendil-works/pi-tui": "*", "typebox": "*" }, "peerDependenciesMeta": { + "@earendil-works/pi-ai": { + "optional": true + }, "@earendil-works/pi-coding-agent": { "optional": true }, @@ -56,6 +61,7 @@ }, "devDependencies": { "@biomejs/biome": "2.3.7", + "@earendil-works/pi-ai": "0.84.1", "@earendil-works/pi-coding-agent": "0.84.1", "@earendil-works/pi-tui": "0.84.1", "@tailwindcss/vite": "4.1.18", diff --git a/tests/auto-router-classify.test.ts b/tests/auto-router-classify.test.ts new file mode 100644 index 0000000..6e417b8 --- /dev/null +++ b/tests/auto-router-classify.test.ts @@ -0,0 +1,118 @@ +import { expect, test } from "bun:test"; +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; +import { classifyTurnComplexity } from "../extensions/auto-router-classify.ts"; +import type { AutoRouterEffortLevel } from "../extensions/auto-router-settings.ts"; + +const MODEL = { provider: "prov", id: "classifier" } as unknown as Model; + +function registryReplying( + text: string, + usage?: { input: number; output: number; cost: { total: number } }, +): ModelRegistry { + return { + complete: async () => ({ + content: [{ type: "text", text }], + usage, + }), + } as unknown as ModelRegistry; +} + +function throwingRegistry(): ModelRegistry { + return { + complete: async () => { + throw new Error("provider down"); + }, + } as unknown as ModelRegistry; +} + +const ALL_LEVELS: AutoRouterEffortLevel[] = ["minimal", "low", "medium", "high", "xhigh", "max"]; + +for (const level of ALL_LEVELS) { + test(`classifyTurnComplexity parses a bare "${level}" reply`, async () => { + const result = await classifyTurnComplexity(registryReplying(level), MODEL, "do something", false); + expect(result.level).toBe(level); + }); +} + +test("classifyTurnComplexity parses the reply case-insensitively with surrounding text", async () => { + const result = await classifyTurnComplexity( + registryReplying(" I'd say MAX. \n"), + MODEL, + "do something", + false, + ); + expect(result.level).toBe("max"); +}); + +test("classifyTurnComplexity picks the level word the model actually led with, not whichever word sorts earliest in the level list", async () => { + // "medium" sorts before "high" in the internal level list, but the model's stated verdict + // here is "high" - a naive "first match in list order" parse would wrongly return "medium". + const result = await classifyTurnComplexity( + registryReplying("high complexity, more than a medium task"), + MODEL, + "do something", + false, + ); + expect(result.level).toBe("high"); +}); + +test("classifyTurnComplexity picks the level word the model actually led with, even when a later caveat mentions an earlier-sorting level", async () => { + const result = await classifyTurnComplexity( + registryReplying("medium at first glance, though parts could be high"), + MODEL, + "do something", + false, + ); + expect(result.level).toBe("medium"); +}); + +test("classifyTurnComplexity falls back to medium on an unparseable reply", async () => { + const result = await classifyTurnComplexity(registryReplying("uh, tricky one"), MODEL, "do something", false); + expect(result.level).toBe("medium"); +}); + +test("classifyTurnComplexity falls back to medium when the provider call throws, and records why in reply", async () => { + const result = await classifyTurnComplexity(throwingRegistry(), MODEL, "do something", false); + expect(result.level).toBe("medium"); + expect(result.usage).toBeUndefined(); + expect(result.reply).toContain("provider down"); +}); + +test("classifyTurnComplexity surfaces the raw reply text alongside the parsed level, for diagnosing mismatches later", async () => { + const result = await classifyTurnComplexity( + registryReplying("high complexity, more than a medium task"), + MODEL, + "do something", + false, + ); + expect(result.level).toBe("high"); + expect(result.reply).toBe("high complexity, more than a medium task"); +}); + +test("classifyTurnComplexity surfaces usage from the response when present", async () => { + const result = await classifyTurnComplexity( + registryReplying("low", { input: 42, output: 7, cost: { total: 0.002 } }), + MODEL, + "do something", + false, + ); + expect(result.level).toBe("low"); + expect(result.usage).toEqual({ input: 42, output: 7, cost: 0.002 }); +}); + +test("classifyTurnComplexity notes attached images in the classification prompt", async () => { + let capturedText: string | undefined; + const registry = { + complete: async ( + _model: Model, + context: { messages: Array<{ content: Array<{ text?: string }> }> }, + ) => { + capturedText = context.messages[0]?.content?.[0]?.text; + return { content: [{ type: "text", text: "medium" }], usage: undefined }; + }, + } as unknown as ModelRegistry; + + await classifyTurnComplexity(registry, MODEL, "describe this screenshot", true); + expect(capturedText).toContain("attached images"); +}); diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts new file mode 100644 index 0000000..a2e5ce9 --- /dev/null +++ b/tests/auto-router-extension.test.ts @@ -0,0 +1,833 @@ +import { afterAll, beforeEach, expect, test } from "bun:test"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { + ExtensionAPI, + ExtensionCommandContext, + ExtensionContext, + Theme, +} from "@earendil-works/pi-coding-agent"; +import autoRouter, { escapeTableCell } from "../extensions/auto-router.ts"; +import type { AutoRouterSettings } from "../extensions/auto-router-settings.ts"; + +test("escapeTableCell neutralizes both pipes and line breaks, so one bad reply can't break the rest of the table", () => { + expect(escapeTableCell("a | b")).toBe("a \\| b"); + expect(escapeTableCell("line one\nline two")).toBe("line one line two"); + expect(escapeTableCell("windows\r\nstyle")).toBe("windows style"); + expect(escapeTableCell("multi\n\n\nblank\nlines")).toBe("multi blank lines"); +}); + +const ENV_VAR = "PI_CODING_AGENT_DIR"; +let agentDir: string | undefined; +const usedDirs: string[] = []; + +// The extension's internal AutoRouterHealthStore debounces its writes (~2s after the last +// record call), so a save scheduled by one test can fire well after that test's own teardown - +// if `afterEach` restored PI_CODING_AGENT_DIR to its prior (usually unset) value in the +// meantime, that late write would land in the real global agent directory instead of a test's +// temp one. So the env var is never restored to anything other than a temp dir for the whole +// run - only ever moved to a new one - and every temp dir used stays on disk until all tests +// finish, so even a very late write can only ever land somewhere harmless. +beforeEach(async () => { + agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-agent-")); + usedDirs.push(agentDir); + process.env[ENV_VAR] = agentDir; +}); + +afterAll(async () => { + delete process.env[ENV_VAR]; + await Promise.all(usedDirs.map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +async function writeConfig(settings: AutoRouterSettings): Promise { + if (!agentDir) throw new Error("agentDir not set"); + await writeFile( + join(agentDir, "settings.json"), + JSON.stringify({ autoRouter: settings }), + ); +} + +function model(provider: string, id: string): Model { + return { provider, id } as unknown as Model; +} + +const AUTO_PLACEHOLDER = model("auto", "auto"); + +type FakeHandler = (event: unknown, ctx: unknown) => unknown; +type ModelRef = { value: Model | undefined }; + +type FakePi = { + pi: ExtensionAPI; + fire: (event: string, payload: unknown, ctx: unknown) => Promise; + runCommand: (name: string, args: string, ctx: unknown) => Promise; + setModelCalls: Model[]; + thinkingLevelCalls: string[]; + appendedEntries: Array<{ type: string; data: unknown }>; + footerEvents: unknown[]; + currentModel: ModelRef; +}; + +function createFakePi(): FakePi { + const handlers = new Map(); + const commands = new Map Promise }>(); + const setModelCalls: Model[] = []; + const thinkingLevelCalls: string[] = []; + const appendedEntries: Array<{ type: string; data: unknown }> = []; + const footerEvents: unknown[] = []; + const currentModel: ModelRef = { value: undefined }; + + const pi = { + registerProvider: () => undefined, + on: (event: string, handler: FakeHandler) => { + handlers.set(event, handler); + return () => handlers.delete(event); + }, + registerCommand: (name: string, options: { handler: (args: string, ctx: unknown) => Promise }) => { + commands.set(name, options); + }, + appendEntry: (type: string, data: unknown) => { + appendedEntries.push({ type, data }); + }, + events: { emit: (_event: string, value: unknown) => footerEvents.push(value), on: () => () => undefined }, + setModel: async (m: Model) => { + setModelCalls.push(m); + currentModel.value = m; + return true; + }, + setThinkingLevel: async (level: string) => { + thinkingLevelCalls.push(level); + }, + } as unknown as ExtensionAPI; + + return { + pi, + fire: async (event, payload, ctx) => { + const handler = handlers.get(event); + if (!handler) throw new Error(`no handler registered for ${event}`); + return handler(payload, ctx); + }, + runCommand: async (name, args, ctx) => { + const command = commands.get(name); + if (!command) throw new Error(`no command registered: ${name}`); + await command.handler(args, ctx as ExtensionCommandContext); + }, + setModelCalls, + thinkingLevelCalls, + appendedEntries, + footerEvents, + currentModel, + }; +} + +const FAKE_THEME = { fg: (_kind: string, text: string) => text } as unknown as Theme; + +function lastFooterBadge(footerEvents: unknown[]): string | undefined { + const last = footerEvents.at(-1) as { identitySuffix?: (theme: Theme) => string | undefined } | undefined; + return last?.identitySuffix?.(FAKE_THEME); +} + +type FakeRegistryOptions = { + models: Model[]; + unavailable?: Model[]; + classify?: (prompt: string) => string; +}; + +function fakeModelRegistry({ models, unavailable = [], classify }: FakeRegistryOptions) { + // Real Pi always has our registered "auto" placeholder findable, same as any other provider. + const allModels = [...models, AUTO_PLACEHOLDER]; + const unavailableKeys = new Set(unavailable.map((m) => `${m.provider}/${m.id}`)); + return { + find: (provider: string, id: string) => + // Real Pi also has every registered "auto-" pinned placeholder findable, the same + // way it has the bare "auto" one - synthesize rather than requiring every test to list them. + provider === "auto" && id.startsWith("auto-") + ? model(provider, id) + : allModels.find((m) => m.provider === provider && m.id === id), + hasConfiguredAuth: (m: Model) => !unavailableKeys.has(`${m.provider}/${m.id}`), + getApiKeyForProvider: async () => undefined, + complete: async (_model: Model, context: { messages: Array<{ content: Array<{ text?: string }> }> }) => { + const prompt = context.messages[0]?.content?.[0]?.text ?? ""; + const level = classify ? classify(prompt) : "medium"; + return { content: [{ type: "text", text: level }], usage: undefined }; + }, + }; +} + +function fakeCtx(options: { + modelRegistry: ReturnType; + currentModel?: ModelRef; + model?: Model; + thinkingLevel?: string; + entries?: unknown[]; +}) { + const notifications: Array<{ message: string; type?: string }> = []; + const modelRef = options.currentModel ?? { value: options.model }; + return { + hasUI: true, + mode: "rpc", + get model() { + return modelRef.value; + }, + thinkingLevel: options.thinkingLevel, + modelRegistry: options.modelRegistry, + sessionManager: { + getSessionId: () => "session-1", + getEntries: () => options.entries ?? [], + }, + ui: { + notify: (message: string, type?: string) => notifications.push({ message, type }), + setStatus: () => undefined, + custom: async () => undefined, + }, + notifications, + } as unknown as ExtensionContext & { notifications: typeof notifications }; +} + +function selectAuto(fake: FakePi, ctx: unknown): Promise { + return fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); +} + +function pinnedPlaceholder(tier: string): Model { + return model("auto", `auto-${tier}`); +} + +function selectPinned(fake: FakePi, ctx: unknown, tier: string): Promise { + return fake.fire( + "model_select", + { model: { provider: "auto", id: `auto-${tier}` }, previousModel: undefined, source: "set" }, + ctx, + ); +} + +test("selecting Auto marks it active without eagerly routing, showing a neutral footer badge", async () => { + const a = model("prov", "model-a"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a] }), currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + expect(fake.setModelCalls).toEqual([]); + expect(fake.thinkingLevelCalls).toEqual([]); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto"); +}); + +test("selecting a pinned Auto () entry shows that tier in the footer immediately, before any turn runs", async () => { + const a = model("prov", "model-a"); + await writeConfig({ efforts: { high: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a] }), currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectPinned(fake, ctx, "high"); + + expect(fake.setModelCalls).toEqual([]); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (high)"); +}); + +test("a pinned Auto () entry routes directly within that tier, skipping classification entirely", async () => { + const medium = model("prov", "medium-model"); + const high = model("prov", "high-model"); + let classifyCalls = 0; + await writeConfig({ + efforts: { + medium: { models: [{ provider: "prov", id: "medium-model" }] }, + high: { models: [{ provider: "prov", id: "high-model" }] }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ + models: [medium, high], + // If classification ran, it would say "medium" - proving the pinned tier (high) wins + // regardless, and that this classifier was never actually asked. + classify: () => { + classifyCalls++; + return "medium"; + }, + }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectPinned(fake, ctx, "high"); + await fake.fire("before_agent_start", { prompt: "anything, complexity doesn't matter here" }, ctx); + + expect(fake.setModelCalls).toEqual([high]); + expect(fake.thinkingLevelCalls).toEqual(["high"]); + expect(classifyCalls).toBe(0); +}); + +test("after a turn settles, /model reverts to the same pinned entry that was selected, not the adaptive one", async () => { + const high = model("prov", "high-model"); + await writeConfig({ efforts: { high: { models: [{ provider: "prov", id: "high-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [high] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectPinned(fake, ctx, "high"); + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + await fake.fire("agent_settled", {}, ctx); + + expect(ctx.model).toEqual(pinnedPlaceholder("high")); +}); + +test("session_start restores a pinned tier from a persisted entry and reverts to that specific placeholder", async () => { + const already = model("prov", "already-selected"); + await writeConfig({ efforts: { xhigh: { models: [{ provider: "prov", id: "already-selected" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [already] }); + fake.currentModel.value = already; + const ctx = fakeCtx({ + modelRegistry: registry, + currentModel: fake.currentModel, + thinkingLevel: "xhigh", + entries: [ + { + type: "custom", + customType: "vessup:auto-router:active", + data: { enabled: true, pinnedTier: "xhigh" }, + }, + ], + }); + + await fake.fire("session_start", {}, ctx); + + expect(fake.setModelCalls).toEqual([pinnedPlaceholder("xhigh")]); + expect(ctx.model).toEqual(pinnedPlaceholder("xhigh")); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (xhigh)"); +}); + +test("before_agent_start self-heals into the correct pinned tier when ctx.model is already that pinned placeholder", async () => { + const max = model("prov", "max-model"); + await writeConfig({ efforts: { max: { models: [{ provider: "prov", id: "max-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [max] }); + // Simulates defaultModel/defaultProvider being set to a pinned "auto-max" entry globally - + // no model_select, no session entries, exactly like the earlier defaultModel=auto bug. + fake.currentModel.value = pinnedPlaceholder("max"); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel, entries: [] }); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([max]); + expect(fake.thinkingLevelCalls).toEqual(["max"]); +}); + +test("before_agent_start routes to the classified tier, and the picker shows Auto again once the turn settles", async () => { + const medium = model("prov", "medium-model"); + const high = model("prov", "high-model"); + await writeConfig({ + efforts: { + medium: { models: [{ provider: "prov", id: "medium-model" }] }, + high: { models: [{ provider: "prov", id: "high-model" }] }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ + models: [medium, high], + classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), + }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "please refactor this multi-file module" }, ctx); + expect(fake.setModelCalls).toEqual([high]); + expect(fake.thinkingLevelCalls).toEqual(["high"]); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (high)"); + // Mid-turn, /model would show the real routed model, not "Auto". + expect(ctx.model).toEqual(high); + + await fake.fire("agent_settled", {}, ctx); + + // Once the turn settles, /model shows Auto selected again... + expect(ctx.model).toEqual(AUTO_PLACEHOLDER); + expect(fake.setModelCalls.at(-1)).toEqual(AUTO_PLACEHOLDER); + // ...while the footer badge keeps reflecting the last real routing. + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (high)"); +}); + +test("a model's `effort` override sets its own thinking level, independent of the tier that routed to it", async () => { + const high = model("opencode-go", "kimi-k3"); + await writeConfig({ + efforts: { + high: { + models: [{ provider: "opencode-go", id: "kimi-k3", effort: "max" }], + }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [high], classify: () => "high" }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([high]); + // Routed via the "high" tier (that's what got classified and what /usage groups it under), + // but dispatched at "max" thinking level per the model's own override. + expect(fake.thinkingLevelCalls).toEqual(["max"]); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (max)"); + + await fake.runCommand("usage", "", ctx); + const notified = ctx.notifications.at(-1)?.message ?? ""; + expect(notified).toContain("high"); // still grouped under its configured tier + expect(notified).toContain("at max effort"); // classification log shows the real applied effort +}); + +test("a routed turn's classification is logged and shows up in /usage, so a routing decision can be checked against what the classifier actually said", async () => { + const medium = model("prov", "medium-model"); + const high = model("prov", "high-model"); + await writeConfig({ + efforts: { + medium: { models: [{ provider: "prov", id: "medium-model" }] }, + high: { models: [{ provider: "prov", id: "high-model" }] }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ + models: [medium, high], + classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), + }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + await fake.fire("before_agent_start", { prompt: "please refactor this multi-file module" }, ctx); + + await fake.runCommand("usage", "", ctx); + const notified = ctx.notifications.at(-1)?.message ?? ""; + expect(notified).toContain("Recent classifications"); + expect(notified).toContain("high"); + expect(notified).toContain("prov/high-model"); +}); + +test("routing escalates to a higher tier when the classified tier is entirely unhealthy", async () => { + const c = model("prov", "high-model"); + const d = model("prov", "xhigh-model"); + const medium = model("prov", "medium-model"); + await writeConfig({ + efforts: { + medium: { models: [{ provider: "prov", id: "medium-model" }] }, + high: { models: [{ provider: "prov", id: "high-model" }] }, + xhigh: { models: [{ provider: "prov", id: "xhigh-model" }] }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [medium, c, d], classify: () => "high" }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + // First turn classifies "high" and routes to the high-tier model, making it current. + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([c]); + + // That model then takes a hard auth failure. + await fake.fire("after_provider_response", { status: 401, headers: {} }, ctx); + fake.setModelCalls.length = 0; + fake.thinkingLevelCalls.length = 0; + + // Same classification again should now escalate past the unhealthy high tier. + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([d]); + expect(fake.thinkingLevelCalls).toEqual(["xhigh"]); +}); + +test("routing falls back to the classified tier's model as a last resort when nothing anywhere is healthy, and warns", async () => { + const only = model("prov", "only-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "only-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [only] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + await fake.fire("after_provider_response", { status: 401, headers: {} }, ctx); + fake.setModelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([only]); + expect(ctx.notifications.some((n) => n.type === "warning" && n.message.includes("unavailable"))).toBe(true); +}); + +test("the last-resort fallback labels the model with the tier it actually belongs to, not the (unconfigured) resolved tier", async () => { + // Only "high" is configured; "medium" has nothing. resolveEffortTier falls back to the + // literal "medium" when nothing between the classified level and medium has models, so + // pickForTier can be called with a tier that itself has zero configured entries. When that + // happens and the only real fallback model comes from a *different* tier (high), the + // returned tier must reflect that model's real tier - otherwise the wrong thinking level + // gets applied to it. + const high = model("prov", "high-model"); + await writeConfig({ efforts: { high: { models: [{ provider: "prov", id: "high-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [high], classify: () => "high" }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + // First turn: classified "high", routes normally to the only configured model. + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([high]); + expect(fake.thinkingLevelCalls).toEqual(["high"]); + + // It then fails, which also makes it unhealthy as the classifier - so the next turn's level + // defaults straight to "medium" without even running classification, and with nothing above + // it healthy either (high is now the only tier and it just failed), forcing the cross-tier + // last-resort fallback. + await fake.fire("after_provider_response", { status: 401, headers: {} }, ctx); + fake.setModelCalls.length = 0; + fake.thinkingLevelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([high]); + // Must be labeled "high" (where the model actually lives), not "medium" (the resolved-but- + // unconfigured tier that triggered the fallback). + expect(fake.thinkingLevelCalls).toEqual(["high"]); +}); + +test("a recorded failure fails over to the next configured model in the same tier", async () => { + const a = model("prov", "model-a"); + const b = model("prov", "model-b"); + await writeConfig({ + efforts: { + medium: { + models: [ + { provider: "prov", id: "model-a" }, + { provider: "prov", id: "model-b" }, + ], + }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, b] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([a]); + + // model-a gets rate limited. + await fake.fire("after_provider_response", { status: 429, headers: {} }, ctx); + fake.setModelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([b]); +}); + +test("a message-level provider error (no distinct HTTP failure status) still fails over on the next turn", async () => { + const a = model("prov", "model-a"); + const b = model("prov", "model-b"); + await writeConfig({ + efforts: { + medium: { + models: [ + { provider: "prov", id: "model-a" }, + { provider: "prov", id: "model-b" }, + ], + }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, b] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([a]); + + // The HTTP response came back 200, so after_provider_response never fires as a failure - + // the error only shows up once Pi finalizes the assistant message. + await fake.fire( + "message_end", + { + message: { + role: "assistant", + stopReason: "error", + errorMessage: "Codex error: The usage limit has been reached", + }, + }, + ctx, + ); + fake.setModelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([b]); +}); + +test("an aborted (user-cancelled) message does not count as a provider failure", async () => { + const a = model("prov", "model-a"); + const b = model("prov", "model-b"); + await writeConfig({ + efforts: { + medium: { + models: [ + { provider: "prov", id: "model-a" }, + { provider: "prov", id: "model-b" }, + ], + }, + }, + }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, b] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + await fake.fire("message_end", { message: { role: "assistant", stopReason: "aborted" } }, ctx); + fake.setModelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([a]); +}); + +test("router-observed usage is tracked for a manually-selected configured model, not just ones Auto routed to", async () => { + const a = model("prov", "model-a"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + // The user picks model-a straight from /model - Auto is never engaged. + fake.currentModel.value = a; + + await fake.fire( + "message_end", + { message: { role: "assistant", stopReason: "stop", usage: { input: 10, output: 5, cost: { total: 0.01 } } } }, + ctx, + ); + + await fake.runCommand("usage", "", ctx); + const notified = ctx.notifications.at(-1)?.message ?? ""; + expect(notified).toContain("1 req"); +}); + +test("router-observed failures are tracked for a manually-selected configured model too", async () => { + const a = model("prov", "model-a"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + fake.currentModel.value = a; + + await fake.fire("after_provider_response", { status: 429, headers: {} }, ctx); + + await fake.runCommand("usage", "", ctx); + const notified = ctx.notifications.at(-1)?.message ?? ""; + expect(notified).toContain("cooldown"); +}); + +test("usage from a model that isn't configured anywhere in autoRouter is not tracked", async () => { + const a = model("prov", "model-a"); + const unrelated = model("other", "unrelated-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, unrelated] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + fake.currentModel.value = unrelated; + + await fake.fire( + "message_end", + { message: { role: "assistant", stopReason: "stop", usage: { input: 10, output: 5, cost: { total: 0.01 } } } }, + ctx, + ); + + await fake.runCommand("usage", "", ctx); + const notified = ctx.notifications.at(-1)?.message ?? ""; + expect(notified).toContain("prov/model-a"); + expect(notified).not.toContain("unrelated-model"); +}); + +test("manually picking a real model while Auto is active turns Auto off", async () => { + const a = model("prov", "model-a"); + const manual = model("prov", "manual-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, manual] }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + await fake.fire("model_select", { model: manual, previousModel: a, source: "set" }, ctx); + fake.setModelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + expect(fake.setModelCalls).toEqual([]); +}); + +test("deactivating Auto removes its footer badge", async () => { + const a = model("prov", "model-a"); + const manual = model("prov", "manual-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a, manual] }), currentModel: fake.currentModel }); + + await fake.fire("session_start", {}, ctx); + await selectAuto(fake, ctx); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto"); + + await fake.fire("model_select", { model: manual, previousModel: a, source: "set" }, ctx); + expect(fake.footerEvents.at(-1)).toMatchObject({ remove: true }); +}); + +test("/usage with no configured models notifies instead of throwing", async () => { + const fake = createFakePi(); + await autoRouter(fake.pi); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [] }) }); + + await fake.runCommand("usage", "", ctx); + expect(ctx.notifications.some((n) => n.message.includes("no configured models"))).toBe(true); +}); + +test("session_start on a cleanly-idle Auto session leaves the placeholder selected", async () => { + const already = model("prov", "already-selected"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "already-selected" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [already] }); + fake.currentModel.value = AUTO_PLACEHOLDER; + const ctx = fakeCtx({ + modelRegistry: registry, + currentModel: fake.currentModel, + entries: [{ type: "custom", customType: "vessup:auto-router:active", data: { enabled: true } }], + }); + + await fake.fire("session_start", {}, ctx); + + expect(fake.setModelCalls).toEqual([]); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto"); +}); + +test("a brand-new session whose defaultModel is auto/auto routes on the first turn with no prior /model pick or session entries", async () => { + const medium = model("prov", "medium-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "medium-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [medium] }); + // No `model_select` for Auto ever fired, and no session entries exist - this is what a + // brand-new session looks like when `defaultProvider`/`defaultModel` are set to "auto" in + // global settings rather than picked interactively. + fake.currentModel.value = AUTO_PLACEHOLDER; + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel, entries: [] }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([medium]); + expect(ctx.model).toEqual(medium); +}); + +test("before_agent_start routes for real even if autoActive's own bookkeeping never saw the placeholder become active", async () => { + const medium = model("prov", "medium-model"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "medium-model" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [medium] }); + fake.currentModel.value = AUTO_PLACEHOLDER; + // Deliberately no `session_start` and no `model_select` at all here - autoActive is stuck at + // its initial `false`. Whatever caused that desync (a timing race, a code path we haven't + // accounted for), `ctx.model` already being the inert placeholder right before dispatch is + // itself proof a real request is about to fail, and that must be enough to route for real. + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel, entries: [] }); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); + + expect(fake.setModelCalls).toEqual([medium]); + expect(ctx.model).toEqual(medium); + expect(fake.appendedEntries).toContainEqual({ + type: "vessup:auto-router:active", + data: { enabled: true }, + }); + + // And the self-heal sticks: the next turn's placeholder-revert and routing both still work. + await fake.fire("agent_settled", {}, ctx); + expect(ctx.model).toEqual(AUTO_PLACEHOLDER); +}); + +test("session_start restored mid-turn (e.g. after a crash) reverts back to the Auto placeholder", async () => { + const already = model("prov", "already-selected"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "already-selected" }] } } }); + + const fake = createFakePi(); + await autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [already] }); + fake.currentModel.value = already; + const ctx = fakeCtx({ + modelRegistry: registry, + currentModel: fake.currentModel, + thinkingLevel: "medium", + entries: [{ type: "custom", customType: "vessup:auto-router:active", data: { enabled: true } }], + }); + + await fake.fire("session_start", {}, ctx); + + expect(fake.setModelCalls).toEqual([AUTO_PLACEHOLDER]); + expect(ctx.model).toEqual(AUTO_PLACEHOLDER); + expect(lastFooterBadge(fake.footerEvents)).toBe("🔀 Auto (medium)"); +}); diff --git a/tests/auto-router-health.test.ts b/tests/auto-router-health.test.ts new file mode 100644 index 0000000..1c38f42 --- /dev/null +++ b/tests/auto-router-health.test.ts @@ -0,0 +1,247 @@ +import { afterAll, beforeEach, expect, test } from "bun:test"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + applyFailure, + applyQuotaResult, + applySuccess, + AutoRouterHealthStore, + type AutoRouterHealthState, + isHealthy, + modelKey, + parseRetryAfterMs, + pickHealthy, +} from "../extensions/auto-router-health.ts"; + +const NOW = 1_000_000_000_000; + +const ENV_VAR = "PI_CODING_AGENT_DIR"; +let agentDir: string | undefined; +const usedDirs: string[] = []; + +// AutoRouterHealthStore debounces its writes (~2s after the last record call), so a save +// scheduled by one test can fire well after that test's own teardown - if `afterEach` restored +// PI_CODING_AGENT_DIR to its prior (usually unset) value in the meantime, that late write would +// land in the real global agent directory instead of a test's temp one. So the env var is never +// restored to anything other than a temp dir for the whole run - only ever moved to a new one - +// and every temp dir used stays on disk until all tests finish, so even a very late write can +// only ever land somewhere harmless. +beforeEach(async () => { + agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-health-")); + usedDirs.push(agentDir); + process.env[ENV_VAR] = agentDir; +}); + +afterAll(async () => { + delete process.env[ENV_VAR]; + await Promise.all(usedDirs.map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +test("modelKey joins provider and id", () => { + expect(modelKey({ provider: "openai", id: "gpt-5.3-codex" })).toBe("openai/gpt-5.3-codex"); +}); + +test("parseRetryAfterMs reads a numeric seconds header", () => { + expect(parseRetryAfterMs({ "retry-after": "30" }, NOW)).toBe(30_000); +}); + +test("parseRetryAfterMs reads an HTTP-date header", () => { + const future = new Date(NOW + 60_000).toUTCString(); + expect(parseRetryAfterMs({ "retry-after": future }, NOW)).toBeCloseTo(60_000, -2); +}); + +test("parseRetryAfterMs returns undefined when the header is missing or unparseable", () => { + expect(parseRetryAfterMs(undefined, NOW)).toBeUndefined(); + expect(parseRetryAfterMs({}, NOW)).toBeUndefined(); + expect(parseRetryAfterMs({ "retry-after": "not-a-value" }, NOW)).toBeUndefined(); +}); + +test("applySuccess resets failure state and accumulates usage totals", () => { + let state: AutoRouterHealthState = {}; + state = applyFailure(state, "m", 500, undefined, NOW); + state = applyFailure(state, "m", 500, undefined, NOW); + state = applyFailure(state, "m", 500, undefined, NOW); + expect(state.m.cooldownUntil).toBeDefined(); + + state = applySuccess(state, "m", { input: 10, output: 5, cost: 0.01 }); + expect(state.m.consecutiveFailures).toBe(0); + expect(state.m.cooldownUntil).toBeUndefined(); + expect(state.m.totals).toEqual({ requests: 1, input: 10, output: 5, cost: 0.01 }); + + state = applySuccess(state, "m", { input: 3, output: 2, cost: 0.001 }); + expect(state.m.totals).toEqual({ requests: 2, input: 13, output: 7, cost: 0.011 }); +}); + +test("applyFailure on 429 honors retry-after when present", () => { + const state = applyFailure({}, "m", 429, { "retry-after": "120" }, NOW); + expect(state.m.cooldownUntil).toBe(NOW + 120_000); +}); + +test("applyFailure on 429 backs off exponentially without retry-after, capped", () => { + let state: AutoRouterHealthState = {}; + const cooldowns: number[] = []; + for (let i = 0; i < 8; i++) { + state = applyFailure(state, "m", 429, undefined, NOW); + cooldowns.push((state.m.cooldownUntil ?? NOW) - NOW); + } + // Strictly increasing until it hits the cap, then flat. + for (let i = 1; i < cooldowns.length; i++) { + expect(cooldowns[i]).toBeGreaterThanOrEqual(cooldowns[i - 1]); + } + expect(Math.max(...cooldowns)).toBe(60 * 60_000); +}); + +test("applyFailure on 401/403 sets a fixed auth cooldown immediately", () => { + const state401 = applyFailure({}, "m", 401, undefined, NOW); + expect(state401.m.cooldownUntil).toBe(NOW + 30 * 60_000); + const state403 = applyFailure({}, "m", 403, undefined, NOW); + expect(state403.m.cooldownUntil).toBe(NOW + 30 * 60_000); +}); + +test("applyFailure on 5xx only cools down after crossing the consecutive-failure threshold", () => { + let state: AutoRouterHealthState = {}; + state = applyFailure(state, "m", 500, undefined, NOW); + expect(state.m.cooldownUntil).toBeUndefined(); + state = applyFailure(state, "m", 500, undefined, NOW); + expect(state.m.cooldownUntil).toBeUndefined(); + state = applyFailure(state, "m", 500, undefined, NOW); + expect(state.m.cooldownUntil).toBe(NOW + 2 * 60_000); +}); + +test("applyFailure on other 4xx codes eventually applies a generic cooldown", () => { + let state: AutoRouterHealthState = {}; + for (let i = 0; i < 4; i++) { + state = applyFailure(state, "m", 400, undefined, NOW); + expect(state.m.cooldownUntil).toBeUndefined(); + } + state = applyFailure(state, "m", 400, undefined, NOW); + expect(state.m.cooldownUntil).toBe(NOW + 5 * 60_000); +}); + +test("applyQuotaResult sets a cooldown to the real reset time even without a local 429", () => { + const state = applyQuotaResult({}, "m", { exhausted: true, resetsAt: NOW + 90_000 }, NOW); + expect(state.m.cooldownUntil).toBe(NOW + 90_000); + expect(state.m.verifiedAt).toBe(NOW); +}); + +test("applyQuotaResult without a reset time falls back to a short default cooldown", () => { + const state = applyQuotaResult({}, "m", { exhausted: true }, NOW); + expect(state.m.cooldownUntil).toBe(NOW + 5 * 60_000); +}); + +test("applyQuotaResult with confirmed headroom clears a stale local cooldown outright", () => { + let state: AutoRouterHealthState = applyFailure({}, "m", 429, undefined, NOW); + expect(state.m.cooldownUntil).toBeDefined(); + state = applyQuotaResult(state, "m", { exhausted: false }, NOW + 1); + expect(state.m.cooldownUntil).toBeUndefined(); + expect(state.m.consecutiveFailures).toBe(0); + expect(state.m.verifiedAt).toBe(NOW + 1); +}); + +test("isHealthy/pickHealthy skip models in cooldown and pick the first healthy candidate", () => { + const a = { provider: "p", id: "a" }; + const b = { provider: "p", id: "b" }; + let state: AutoRouterHealthState = {}; + state = applyFailure(state, modelKey(a), 401, undefined, NOW); + + expect(isHealthy(state, modelKey(a), NOW)).toBe(false); + expect(isHealthy(state, modelKey(b), NOW)).toBe(true); + expect(pickHealthy(state, [a, b], NOW)).toEqual(b); +}); + +test("pickHealthy returns undefined once every candidate is in cooldown", () => { + const a = { provider: "p", id: "a" }; + const b = { provider: "p", id: "b" }; + let state: AutoRouterHealthState = {}; + state = applyFailure(state, modelKey(a), 401, undefined, NOW); + state = applyFailure(state, modelKey(b), 401, undefined, NOW); + expect(pickHealthy(state, [a, b], NOW)).toBeUndefined(); +}); + +test("a cooldown clears once its expiry has passed", () => { + const a = { provider: "p", id: "a" }; + const state = applyFailure({}, modelKey(a), 401, undefined, NOW); + expect(isHealthy(state, modelKey(a), NOW + 30 * 60_000 - 1)).toBe(false); + expect(isHealthy(state, modelKey(a), NOW + 30 * 60_000 + 1)).toBe(true); +}); + +function statePath(): string { + if (!agentDir) throw new Error("agentDir not set"); + return join(agentDir, "auto-router-state.json"); +} + +test("AutoRouterHealthStore round-trips model health and classification log through flush/load, without ever persisting prompt content", async () => { + const store = new AutoRouterHealthStore(); + const model = { provider: "prov", id: "a" }; + store.recordSuccess(modelKey(model), { input: 10, output: 20, cost: 0.01 }); + store.recordClassification( + { reply: "high complexity", level: "high", tier: "high", effort: "high", model }, + NOW, + ); + await store.flush(); + + // The prompt itself (which can contain source code, credentials, or personal data) must + // never reach the plaintext state file on disk - only routing metadata does. + const onDisk = await readFile(statePath(), "utf8"); + expect(onDisk).not.toContain("prompt"); + + const reloaded = new AutoRouterHealthStore(); + await reloaded.load(); + expect(reloaded.getEntry(modelKey(model))?.totals).toEqual({ + requests: 1, + input: 10, + output: 20, + cost: 0.01, + }); + expect(reloaded.getClassifications()).toEqual([ + { + timestamp: NOW, + reply: "high complexity", + level: "high", + tier: "high", + effort: "high", + model, + }, + ]); +}); + +test("AutoRouterHealthStore.load reads a pre-classification-log file (flat model-keyed record) as model health with an empty log", async () => { + const model = { provider: "prov", id: "a" }; + await writeFile( + statePath(), + JSON.stringify({ + [modelKey(model)]: { + consecutiveFailures: 0, + totals: { requests: 3, input: 1, output: 2, cost: 0.001 }, + }, + }), + ); + const store = new AutoRouterHealthStore(); + await store.load(); + expect(store.getEntry(modelKey(model))?.totals.requests).toBe(3); + expect(store.getClassifications()).toEqual([]); +}); + +test("AutoRouterHealthStore.recordClassification truncates long text and caps the log at the most recent 20 entries", async () => { + const store = new AutoRouterHealthStore(); + const model = { provider: "prov", id: "a" }; + const longReply = "x".repeat(500); + for (let i = 0; i < 25; i++) { + store.recordClassification( + { + reply: i === 24 ? longReply : `reply ${i}`, + level: "medium", + tier: "medium", + effort: "medium", + model, + }, + NOW + i, + ); + } + const entries = store.getClassifications(); + expect(entries).toHaveLength(20); + expect(entries[0]?.reply).toBe("reply 5"); + expect(entries.at(-1)?.reply.length).toBeLessThanOrEqual(201); + expect(entries.at(-1)?.reply.endsWith("…")).toBe(true); +}); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts new file mode 100644 index 0000000..faf4cc7 --- /dev/null +++ b/tests/auto-router-quota.test.ts @@ -0,0 +1,417 @@ +import { expect, test } from "bun:test"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; +import { + type QuotaFetchDependencies, + reconcileProviderQuota, +} from "../extensions/auto-router-quota.ts"; + +function fakeRegistry(apiKey: string | undefined): ModelRegistry { + return { + getApiKeyForProvider: async () => apiKey, + } as unknown as ModelRegistry; +} + +function fakeDeps( + handler: (url: string) => Response, + { + accountId, + minimaxCli, + }: { accountId?: string; minimaxCli?: (args: string[]) => Promise } = { + accountId: "acct-1", + }, +): QuotaFetchDependencies { + return { + fetchImpl: (async (url: string | URL) => handler(url.toString())) as typeof fetch, + readCodexAccountId: async () => accountId, + runMinimaxCli: minimaxCli ?? (async () => undefined), + }; +} + +function jsonResponse(data: unknown, status = 200): Response { + return new Response(JSON.stringify(data), { status }); +} + +test("reconcileProviderQuota returns undefined for providers without a known fetcher", async () => { + const result = await reconcileProviderQuota( + "totally-unsupported-provider", + fakeRegistry("key"), + fakeDeps(() => jsonResponse({})), + ); + expect(result).toBeUndefined(); +}); + +test("reconcileProviderQuota returns undefined without credentials, never calling fetch", async () => { + let called = false; + const deps = fakeDeps(() => { + called = true; + return jsonResponse({}); + }); + const result = await reconcileProviderQuota("anthropic", fakeRegistry(undefined), deps); + expect(result).toBeUndefined(); + expect(called).toBe(false); +}); + +test("reconcileProviderQuota(anthropic) skips a raw API key without calling fetch", async () => { + let called = false; + const deps = fakeDeps(() => { + called = true; + return jsonResponse({}); + }); + const result = await reconcileProviderQuota("anthropic", fakeRegistry("sk-ant-api03-abc"), deps); + expect(result).toBeUndefined(); + expect(called).toBe(false); +}); + +test("reconcileProviderQuota(anthropic) reports exhaustion from an OAuth token's usage window", async () => { + const resetsAt = "2030-01-01T00:00:00Z"; + const deps = fakeDeps(() => + jsonResponse({ five_hour: { utilization: 100, resets_at: resetsAt } }), + ); + const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse(resetsAt), detail: "5h 100% used" }, + }); +}); + +test("reconcileProviderQuota(anthropic) reports headroom, with the most-used window as detail", async () => { + const deps = fakeDeps(() => jsonResponse({ five_hour: { utilization: 10 }, seven_day: { utilization: 20 } })); + const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); + expect(result).toEqual({ default: { exhausted: false, detail: "7d 20% used" } }); +}); + +test("reconcileProviderQuota(openai-codex) returns undefined without a discoverable account id", async () => { + const deps = fakeDeps(() => jsonResponse({}), { accountId: undefined }); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toBeUndefined(); +}); + +test("reconcileProviderQuota(openai-codex) reports exhaustion when the spend cap is reached", async () => { + const deps = fakeDeps(() => jsonResponse({ spend_control: { reached: true } })); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ default: { exhausted: true, detail: "spend cap reached" } }); +}); + +test("reconcileProviderQuota(openai-codex) applies the account-wide rate_limit.limit_reached flag only to models without their own additional_rate_limits entry", async () => { + // Shape verified against a real exhausted account, *and* the model-specific behavior verified + // directly by the user: the account-wide flag was true, a per-model entry under + // additional_rate_limits for the model actually in use was healthy, and that model kept + // working normally despite the account-wide flag. The two are independent quota tracks - + // the account-wide flag governs the "default" bucket (models with no specific entry below), + // not every model under the provider. + const deps = fakeDeps(() => + jsonResponse({ + rate_limit: { + allowed: false, + limit_reached: true, + primary_window: { used_percent: 100, limit_window_seconds: 604_800, reset_at: 1787197007 }, + secondary_window: null, + }, + additional_rate_limits: [ + { + limit_name: "GPT-5.3-Codex-Spark", + rate_limit: { + allowed: true, + limit_reached: false, + primary_window: { used_percent: 5, limit_window_seconds: 604_800 }, + }, + }, + ], + spend_control: { reached: false }, + }), + ); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ + // Applies to any configured model with no specific additional_rate_limits entry. The + // window's own limit_window_seconds (604800 = 7 days) labels it, not a vague "account". + default: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "7d 100% used" }, + // This model has its own entry, so it's unaffected by the account-wide flag. + perModel: { gpt53codexspark: { exhausted: false, detail: "7d 5% used" } }, + }); +}); + +test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit window is depleted (percent_left)", async () => { + const deps = fakeDeps(() => + jsonResponse({ rate_limit: { primary_window: { percent_left: 0, reset_at: "2030-06-01T00:00:00Z" } } }), + ); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z"), detail: "account 100% used" }, + }); +}); + +test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit window is depleted (used_percent)", async () => { + const deps = fakeDeps(() => + jsonResponse({ rate_limit: { primary_window: { used_percent: 100, reset_at: "2030-06-01T00:00:00Z" } } }), + ); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z"), detail: "account 100% used" }, + }); +}); + +test("reconcileProviderQuota(openai-codex) reports headroom when used_percent is low", async () => { + const deps = fakeDeps(() => jsonResponse({ rate_limit: { primary_window: { used_percent: 12 } } })); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ default: { exhausted: false, detail: "account 12% used" } }); +}); + +test("reconcileProviderQuota(openai-codex) can report a model exhausted independently of a healthy account", async () => { + const deps = fakeDeps(() => + jsonResponse({ + rate_limit: { allowed: true, limit_reached: false, primary_window: { used_percent: 10 } }, + additional_rate_limits: [ + { + limit_name: "GPT-5.6-Sol", + rate_limit: { allowed: false, limit_reached: true, primary_window: { used_percent: 100, reset_at: "2030-06-01T00:00:00Z" } }, + }, + ], + }), + ); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result?.default).toEqual({ exhausted: false, detail: "account 10% used" }); + expect(result?.perModel?.gpt56sol).toEqual({ + exhausted: true, + resetsAt: Date.parse("2030-06-01T00:00:00Z"), + detail: "100% used", + }); +}); + +test("reconcileProviderQuota(zai) reports exhaustion from a TOKENS_LIMIT entry", async () => { + const deps = fakeDeps(() => + jsonResponse({ + data: { limits: [{ type: "TOKENS_LIMIT", unit: 3, number: 5, percentage: 100, nextResetTime: 4_102_444_800_000 }] }, + }), + ); + const result = await reconcileProviderQuota("zai", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "5h 100% used" }, + }); +}); + +test("reconcileProviderQuota(zai) also reports usage from a CREDIT_LIMIT entry (a different plan tier's shape)", async () => { + // Verified directly against a real "pro" plan account: entries carry `type: "CREDIT_LIMIT"` + // (with absolute usage/currentValue/remaining alongside it) rather than "TOKENS_LIMIT", but + // the same `percentage` field either way. + const deps = fakeDeps(() => + jsonResponse({ + data: { + limits: [ + { type: "CREDIT_LIMIT", unit: 3, number: 5, usage: 12000, currentValue: 12023, remaining: 0, percentage: 100 }, + { type: "CREDIT_LIMIT", unit: 6, number: 1, usage: 60000, currentValue: 12023, remaining: 47976, percentage: 20 }, + ], + }, + }), + ); + const result = await reconcileProviderQuota("zai", fakeRegistry("key"), deps); + // The 5h window (100%) is more used than the 7d window (20%), so it wins as "most used". + expect(result).toEqual({ default: { exhausted: true, detail: "5h 100% used" } }); +}); + +test("reconcileProviderQuota(kimi-coding) reports exhaustion when used reaches the weekly limit", async () => { + const deps = fakeDeps(() => jsonResponse({ usage: { limit: 100, used: 100, resetTime: "2030-01-01T00:00:00Z" } })); + const result = await reconcileProviderQuota("kimi-coding", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z"), detail: "100/100 this week" }, + }); +}); + +test("reconcileProviderQuota(minimax) labels the short window from its real start/end times (observed 5h), not a vague 'interval'", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { + model_name: "general", + start_time: 1_000_000_000_000, + end_time: 1_000_000_000_000 + 5 * 60 * 60 * 1000, + current_interval_remaining_percent: 84, + current_weekly_remaining_percent: 89, + }, + { model_name: "video", current_interval_remaining_percent: 100, current_weekly_remaining_percent: 100 }, + ], + }), + }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toEqual({ + default: { exhausted: false, detail: "5h 16% used, weekly 11% used" }, + }); +}); + +test("reconcileProviderQuota(minimax) still derives the right window duration when start/end times are epoch seconds, not milliseconds", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { + model_name: "general", + start_time: 1_700_000_000, + end_time: 1_700_000_000 + 5 * 60 * 60, + current_interval_remaining_percent: 84, + current_weekly_remaining_percent: 89, + }, + ], + }), + }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toEqual({ + default: { exhausted: false, detail: "5h 16% used, weekly 11% used" }, + }); +}); + +test("reconcileProviderQuota(minimax) falls back to 'interval' when start/end times aren't present", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { model_name: "general", current_interval_remaining_percent: 84, current_weekly_remaining_percent: 89 }, + ], + }), + }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toEqual({ + default: { exhausted: false, detail: "interval 16% used, weekly 11% used" }, + }); +}); + +test("reconcileProviderQuota(minimax) reports exhaustion when the interval bucket is depleted", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { + model_name: "general", + current_interval_remaining_percent: 0, + current_weekly_remaining_percent: 50, + end_time: 4_102_444_800_000, + }, + ], + }), + }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 100% used, weekly 50% used" }, + }); +}); + +test("reconcileProviderQuota(minimax) reports exhaustion when only the weekly bucket is depleted", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { + model_name: "general", + current_interval_remaining_percent: 60, + current_weekly_remaining_percent: 0.2, + weekly_end_time: 4_102_444_800_000, + }, + ], + }), + }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 40% used, weekly 99.8% used" }, + }); +}); + +test("reconcileProviderQuota(minimax) degrades gracefully when the CLI is missing or not logged in", async () => { + const deps = fakeDeps(() => jsonResponse({}), { minimaxCli: async () => undefined }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toBeUndefined(); +}); + +test("reconcileProviderQuota(minimax) degrades gracefully on unparseable CLI output", async () => { + const deps = fakeDeps(() => jsonResponse({}), { minimaxCli: async () => "not json" }); + const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); + expect(result).toBeUndefined(); +}); + +test("reconcileProviderQuota(opencode-go) reports every window's usage together, not just whichever has the highest percentage", async () => { + const deps = fakeDeps(() => + jsonResponse({ + usage: { + rolling: { status: "ok", percent: 3, resetsAt: "2030-01-01T00:00:00Z" }, + weekly: { status: "ok", percent: 0, resetsAt: "2030-01-08T00:00:00Z" }, + monthly: { status: "ok", percent: 20, resetsAt: "2030-02-01T00:00:00Z" }, + }, + }), + ); + const result = await reconcileProviderQuota("opencode-go", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { exhausted: false, detail: "rolling 3% used, weekly 0% used, monthly 20% used" }, + }); +}); + +test("reconcileProviderQuota(opencode-go) reports exhaustion when a window's percent crosses the threshold", async () => { + const deps = fakeDeps(() => + jsonResponse({ + usage: { monthly: { status: "ok", percent: 100, resetsAt: "2030-02-01T00:00:00Z" } }, + }), + ); + const result = await reconcileProviderQuota("opencode-go", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-02-01T00:00:00Z"), detail: "monthly 100% used" }, + }); +}); + +test("reconcileProviderQuota(opencode-go) reports exhaustion from a non-ok status even at low percent", async () => { + const deps = fakeDeps(() => + jsonResponse({ + usage: { rolling: { status: "limited", percent: 10, resetsAt: "2030-01-01T00:00:00Z" } }, + }), + ); + const result = await reconcileProviderQuota("opencode-go", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z"), detail: "rolling 10% used" }, + }); +}); + +test("reconcileProviderQuota(opencode-go) uses the blocked window's own resetsAt even when a healthy window has a higher percentage, while still showing both in detail", async () => { + // A blocked window can have a lower percentage than a healthy one (e.g. a short window that + // resets rarely hit its own cap while a longer window is nowhere near full) - resetsAt must + // describe the blocked window specifically, even though detail now shows every window. + const deps = fakeDeps(() => + jsonResponse({ + usage: { + rolling: { status: "limited", percent: 10, resetsAt: "2030-01-01T00:00:00Z" }, + weekly: { status: "ok", percent: 90, resetsAt: "2030-01-08T00:00:00Z" }, + }, + }), + ); + const result = await reconcileProviderQuota("opencode-go", fakeRegistry("key"), deps); + expect(result).toEqual({ + default: { + exhausted: true, + resetsAt: Date.parse("2030-01-01T00:00:00Z"), + detail: "rolling 10% used, weekly 90% used", + }, + }); +}); + +test("reconcileProviderQuota(opencode-go) returns undefined without credentials, never calling fetch", async () => { + let called = false; + const deps = fakeDeps(() => { + called = true; + return jsonResponse({}); + }); + const result = await reconcileProviderQuota("opencode-go", fakeRegistry(undefined), deps); + expect(result).toBeUndefined(); + expect(called).toBe(false); +}); + +test("reconcileProviderQuota degrades gracefully on network/HTTP failure", async () => { + const deps = fakeDeps(() => jsonResponse({}, 500)); + const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); + expect(result).toBeUndefined(); +}); + +test("reconcileProviderQuota degrades gracefully when the fetcher throws", async () => { + const deps: QuotaFetchDependencies = { + fetchImpl: (async () => { + throw new Error("network down"); + }) as typeof fetch, + readCodexAccountId: async () => "acct-1", + runMinimaxCli: async () => undefined, + }; + const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); + expect(result).toBeUndefined(); +}); diff --git a/tests/auto-router-settings.test.ts b/tests/auto-router-settings.test.ts new file mode 100644 index 0000000..6e7aa23 --- /dev/null +++ b/tests/auto-router-settings.test.ts @@ -0,0 +1,239 @@ +import { afterEach, expect, test } from "bun:test"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + allConfiguredModels, + AUTO_MODEL_SCOPE_PATTERN, + type AutoRouterSettings, + ensureAutoModelScoped, + escalationTiers, + parseAutoRouterSettings, + resolveEffortTier, + writeAutoRouterSettingsFile, +} from "../extensions/auto-router-settings.ts"; + +let directory: string | undefined; + +afterEach(async () => { + if (directory) await rm(directory, { recursive: true, force: true }); + directory = undefined; +}); + +async function settingsPath(): Promise { + directory = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-settings-")); + return join(directory, "settings.json"); +} + +async function readSettings(path: string): Promise> { + return JSON.parse(await readFile(path, "utf8")) as Record; +} + +const EXAMPLE: AutoRouterSettings = { + efforts: { + medium: { + models: [ + { provider: "minimax", id: "m3" }, + { provider: "openai", id: "gpt-5.3-codex" }, + ], + }, + high: { + models: [ + { provider: "moonshot", id: "kimi-k3-max" }, + { provider: "zhipu", id: "glm-5.3-max" }, + ], + }, + xhigh: { models: [{ provider: "openai", id: "sol-5.6-xhigh" }] }, + }, +}; + +test("parseAutoRouterSettings accepts a well-formed config", () => { + expect(parseAutoRouterSettings(EXAMPLE)).toEqual(EXAMPLE); +}); + +test("parseAutoRouterSettings drops unknown effort levels", () => { + const parsed = parseAutoRouterSettings({ + efforts: { + medium: { models: [{ provider: "a", id: "b" }] }, + turbo: { models: [{ provider: "a", id: "b" }] }, + }, + }); + expect(Object.keys(parsed.efforts)).toEqual(["medium"]); +}); + +test("parseAutoRouterSettings drops malformed model refs and empty tiers", () => { + const parsed = parseAutoRouterSettings({ + efforts: { + medium: { + models: [{ provider: "a", id: "b" }, { provider: "a" }, { id: "b" }, "nope", 42], + }, + high: { models: [] }, + xhigh: { models: [{ provider: "", id: "b" }] }, + }, + }); + expect(parsed.efforts.medium?.models).toEqual([{ provider: "a", id: "b" }]); + expect(parsed.efforts.high).toBeUndefined(); + expect(parsed.efforts.xhigh).toBeUndefined(); +}); + +test("parseAutoRouterSettings accepts a per-model effort override, independent of the tier it's listed under", () => { + const parsed = parseAutoRouterSettings({ + efforts: { + high: { models: [{ provider: "opencode-go", id: "kimi-k3", effort: "max" }] }, + }, + }); + expect(parsed.efforts.high?.models).toEqual([ + { provider: "opencode-go", id: "kimi-k3", effort: "max" }, + ]); +}); + +test("parseAutoRouterSettings drops an invalid effort override rather than the whole model entry", () => { + const parsed = parseAutoRouterSettings({ + efforts: { + high: { models: [{ provider: "a", id: "b", effort: "turbo" }] }, + }, + }); + expect(parsed.efforts.high?.models).toEqual([{ provider: "a", id: "b" }]); +}); + +test("parseAutoRouterSettings tolerates garbage input", () => { + expect(parseAutoRouterSettings(null)).toEqual({ efforts: {} }); + expect(parseAutoRouterSettings("nope")).toEqual({ efforts: {} }); + expect(parseAutoRouterSettings({})).toEqual({ efforts: {} }); +}); + +test("resolveEffortTier returns the level itself when configured", () => { + expect(resolveEffortTier(EXAMPLE, "high")).toBe("high"); + expect(resolveEffortTier(EXAMPLE, "xhigh")).toBe("xhigh"); +}); + +test("resolveEffortTier falls back toward medium when the classified level is unconfigured", () => { + // low is unset; steps up to medium, matching the request's own example. + expect(resolveEffortTier(EXAMPLE, "low")).toBe("medium"); + expect(resolveEffortTier(EXAMPLE, "minimal")).toBe("medium"); + expect(resolveEffortTier(EXAMPLE, "off")).toBe("medium"); +}); + +test("resolveEffortTier never overshoots past medium to a configured higher tier", () => { + const onlyXhigh: AutoRouterSettings = { + efforts: { xhigh: { models: [{ provider: "a", id: "b" }] } }, + }; + // high is unconfigured; walking toward medium never considers xhigh even though it exists. + expect(resolveEffortTier(onlyXhigh, "high")).toBe("medium"); +}); + +test("escalationTiers walks strictly upward through configured tiers only", () => { + expect(escalationTiers(EXAMPLE, "medium")).toEqual(["high", "xhigh"]); + expect(escalationTiers(EXAMPLE, "high")).toEqual(["xhigh"]); + expect(escalationTiers(EXAMPLE, "xhigh")).toEqual([]); +}); + +test("escalationTiers skips tiers below fromTier even if configured", () => { + const settings: AutoRouterSettings = { + efforts: { + medium: { models: [{ provider: "a", id: "b" }] }, + xhigh: { models: [{ provider: "c", id: "d" }] }, + }, + }; + expect(escalationTiers(settings, "high")).toEqual(["xhigh"]); +}); + +test("allConfiguredModels lists every model in tier order, deduplicated", () => { + const settings: AutoRouterSettings = { + efforts: { + medium: { models: [{ provider: "a", id: "b" }] }, + high: { + models: [ + { provider: "a", id: "b" }, + { provider: "c", id: "d" }, + ], + }, + }, + }; + expect(allConfiguredModels(settings)).toEqual([ + { provider: "a", id: "b" }, + { provider: "c", id: "d" }, + ]); +}); + +test("auto router settings write preserves unrelated Pi and extension keys", async () => { + const path = await settingsPath(); + await writeFile( + path, + JSON.stringify({ theme: "dark", web: { tailscale: { enabled: true } } }), + ); + await writeAutoRouterSettingsFile(path, EXAMPLE); + expect(await readSettings(path)).toEqual({ + theme: "dark", + web: { tailscale: { enabled: true } }, + autoRouter: EXAMPLE, + }); +}); + +test("concurrent auto router settings updates remain valid and retain unrelated keys", async () => { + const path = await settingsPath(); + await writeFile(path, JSON.stringify({ theme: "light" })); + const other: AutoRouterSettings = { efforts: { medium: { models: [{ provider: "x", id: "y" }] } } }; + await Promise.all([ + writeAutoRouterSettingsFile(path, EXAMPLE), + writeAutoRouterSettingsFile(path, other), + ]); + const result = (await readSettings(path)) as { theme?: string; autoRouter?: AutoRouterSettings }; + expect(result.theme).toBe("light"); + expect([EXAMPLE, other]).toContainEqual(result.autoRouter); +}); + +test("malformed settings reject without leaking the cross-process lock", async () => { + const path = await settingsPath(); + await writeFile(path, "{broken"); + await expect(writeAutoRouterSettingsFile(path, EXAMPLE)).rejects.toThrow("Could not read"); + await writeFile(path, JSON.stringify({ recovered: true })); + await writeAutoRouterSettingsFile(path, EXAMPLE); + expect(await readSettings(path)).toEqual({ recovered: true, autoRouter: EXAMPLE }); +}); + +test("ensureAutoModelScoped adds Auto's pattern when scoping is configured", async () => { + const path = await settingsPath(); + await writeFile(path, JSON.stringify({ enabledModels: ["claude-*", "gpt-4o"] })); + await ensureAutoModelScoped(path); + expect(await readSettings(path)).toEqual({ + enabledModels: ["claude-*", "gpt-4o", AUTO_MODEL_SCOPE_PATTERN], + }); +}); + +test("ensureAutoModelScoped is a no-op when Auto's pattern is already present", async () => { + const path = await settingsPath(); + await writeFile(path, JSON.stringify({ enabledModels: ["claude-*", AUTO_MODEL_SCOPE_PATTERN] })); + await ensureAutoModelScoped(path); + expect(await readSettings(path)).toEqual({ + enabledModels: ["claude-*", AUTO_MODEL_SCOPE_PATTERN], + }); +}); + +test("ensureAutoModelScoped is a no-op when no scoping is configured", async () => { + const path = await settingsPath(); + await writeFile(path, JSON.stringify({ theme: "dark" })); + await ensureAutoModelScoped(path); + expect(await readSettings(path)).toEqual({ theme: "dark" }); +}); + +test("ensureAutoModelScoped is a no-op when enabledModels is an empty array", async () => { + const path = await settingsPath(); + await writeFile(path, JSON.stringify({ enabledModels: [] })); + await ensureAutoModelScoped(path); + expect(await readSettings(path)).toEqual({ enabledModels: [] }); +}); + +test("ensureAutoModelScoped preserves unrelated keys", async () => { + const path = await settingsPath(); + await writeFile( + path, + JSON.stringify({ theme: "dark", enabledModels: ["gpt-4o"], autoRouter: EXAMPLE }), + ); + await ensureAutoModelScoped(path); + expect(await readSettings(path)).toEqual({ + theme: "dark", + enabledModels: ["gpt-4o", AUTO_MODEL_SCOPE_PATTERN], + autoRouter: EXAMPLE, + }); +}); diff --git a/web/client/app.tsx b/web/client/app.tsx index 037add0..1f8d3a3 100644 --- a/web/client/app.tsx +++ b/web/client/app.tsx @@ -1166,10 +1166,10 @@ function SessionListItem({
{ event.preventDefault(); event.stopPropagation(); diff --git a/web/client/semantic-session.tsx b/web/client/semantic-session.tsx index 6bd7d0a..0db8f47 100644 --- a/web/client/semantic-session.tsx +++ b/web/client/semantic-session.tsx @@ -661,12 +661,12 @@ function TokenDetails({ session }: { session: WebSession }) { {(usage?.cacheRead ?? 0) > 0 && ( - Cache read {formatTokenCount(usage?.cacheRead)} + Cache read {formatTokenCount(usage?.cacheRead ?? 0)} )} {(usage?.cacheWrite ?? 0) > 0 && ( - Cache write {formatTokenCount(usage?.cacheWrite)} + Cache write {formatTokenCount(usage?.cacheWrite ?? 0)} )} @@ -1229,9 +1229,9 @@ function ChangedLine({ row, language }: { row: DiffRow; language?: string }) { {pieces.map((piece, index) => { const highlighted = row.kind === "added" - ? piece.added + ? (piece.added ?? false) : row.kind === "removed" - ? piece.removed + ? (piece.removed ?? false) : false; const hidden = row.kind === "added" diff --git a/web/server/index.ts b/web/server/index.ts index b3019d9..1fe556f 100644 --- a/web/server/index.ts +++ b/web/server/index.ts @@ -1395,7 +1395,7 @@ async function recoverStagedSourceSessionDeletions(): Promise { ); }) : undefined; - if (!replacement) { + if (!sourceId || !replacement) { if (!existsSync(staged.source)) renameSync(staged.tombstone, staged.source); continue;