From 3edd4892194780e868249e8e1eda5e003f643665 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 16:37:41 -0400 Subject: [PATCH 01/21] Add Auto model routing extension Adds an "Auto" entry to /model that classifies each turn's complexity with the default model and routes to a configured model/effort tier, failing over to other configured models or higher tiers when one is unhealthy or out of usage. Health is tracked from observed provider responses and, best-effort, reconciled against real quota APIs (Anthropic, OpenAI Codex, Z.ai, Kimi Coding, OpenRouter) at session start and via /usage so state self-corrects across sessions. Co-Authored-By: Claude Sonnet 5 --- README.md | 42 ++- bun.lock | 3 + extensions/auto-router-classify.ts | 100 ++++++ extensions/auto-router-health.ts | 290 ++++++++++++++++ extensions/auto-router-quota.ts | 257 ++++++++++++++ extensions/auto-router-settings.ts | 237 +++++++++++++ extensions/auto-router.ts | 501 ++++++++++++++++++++++++++++ package.json | 8 +- tests/auto-router-extension.test.ts | 332 ++++++++++++++++++ tests/auto-router-health.test.ts | 141 ++++++++ tests/auto-router-quota.test.ts | 125 +++++++ tests/auto-router-settings.test.ts | 171 ++++++++++ 12 files changed, 2205 insertions(+), 2 deletions(-) create mode 100644 extensions/auto-router-classify.ts create mode 100644 extensions/auto-router-health.ts create mode 100644 extensions/auto-router-quota.ts create mode 100644 extensions/auto-router-settings.ts create mode 100644 extensions/auto-router.ts create mode 100644 tests/auto-router-extension.test.ts create mode 100644 tests/auto-router-health.test.ts create mode 100644 tests/auto-router-quota.test.ts create mode 100644 tests/auto-router-settings.test.ts diff --git a/README.md b/README.md index e91c44d..4765bb7 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,46 @@ 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. + +On every turn, Auto asks the `medium` tier's first healthy model (the "default model") to classify the turn as `low`, `medium`, `high`, or `xhigh`, 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. + +Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota API (currently Anthropic, OpenAI Codex, Z.ai, Kimi Coding, and OpenRouter) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. + +Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally; the footer (and Pi Web's model display) always shows the real underlying model and effort actually in use, plus a small `๐Ÿ”€` badge in the TUI footer while Auto is engaged. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. + +### 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) + ## 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 +169,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..635af18 --- /dev/null +++ b/extensions/auto-router-classify.ts @@ -0,0 +1,100 @@ +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[] = [ + "low", + "medium", + "high", + "xhigh", +]; +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: low, medium, high, or xhigh. + +- low: trivial, mechanical, or purely informational. One-line edits, formatting, simple 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. + +Reply with only the single word.`; + +export type ClassificationUsage = { + input: number; + output: number; + cost: number; +}; + +export type ClassificationResult = { + level: AutoRouterEffortLevel; + usage?: ClassificationUsage; +}; + +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(), + }, + ); + const reply = response.content + .filter( + (block): block is { type: "text"; text: string } => + block.type === "text", + ) + .map((block) => block.text) + .join("") + .trim() + .toLowerCase(); + const level = + VALID_LEVELS.find((candidate) => reply.includes(candidate)) ?? + 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 }; + } catch { + return { level: DEFAULT_LEVEL }; + } finally { + clearTimeout(timeout); + } +} diff --git a/extensions/auto-router-health.ts b/extensions/auto-router-health.ts new file mode 100644 index 0000000..523d40a --- /dev/null +++ b/extensions/auto-router-health.ts @@ -0,0 +1,290 @@ +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; +}; + +export type AutoRouterHealthState = Record; + +export type UsageDelta = { input: number; output: number; cost: number }; + +export type QuotaReconciliationResult = { + exhausted: boolean; + /** epoch ms the exhausted window resets, if known. */ + resetsAt?: number; +}; + +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, + } + : { + ...previous, + consecutiveFailures: 0, + cooldownUntil: undefined, + verifiedAt: now, + }; + 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, + }; + } + return state; +} + +/** 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 writeTimer: ReturnType | undefined; + + async load(): Promise { + try { + this.state = parseState( + JSON.parse(await readFile(statePath(), "utf8")), + ); + } catch { + this.state = {}; + } + } + + getState(): AutoRouterHealthState { + return this.state; + } + + 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(); + } + + 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(this.state, 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..42f48d6 --- /dev/null +++ b/extensions/auto-router-quota.ts @@ -0,0 +1,257 @@ +import { readFile } from "node:fs/promises"; +import { homedir } from "node:os"; +import { join } from "node:path"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; +import type { QuotaReconciliationResult } from "./auto-router-health.js"; + +const FETCH_TIMEOUT_MS = 15_000; +const EXHAUSTED_UTILIZATION_PERCENT = 99.5; + +type FetchResult = { ok: true; data: unknown } | { ok: false }; + +/** Injectable for tests; defaults to the real network/filesystem. */ +export type QuotaFetchDependencies = { + fetchImpl: typeof fetch; + readCodexAccountId: () => Promise; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +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; + } +} + +export const defaultQuotaFetchDependencies: QuotaFetchDependencies = { + fetchImpl: fetch, + readCodexAccountId: defaultReadCodexAccountId, +}; + +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; +} + +/** 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; + + for (const key of ["five_hour", "seven_day"] as const) { + const window = result.data[key]; + if (!isRecord(window)) continue; + const utilization = numeric(window.utilization); + if (utilization !== undefined && utilization >= EXHAUSTED_UTILIZATION_PERCENT) { + return { exhausted: true, resetsAt: parseDateish(window.resets_at) }; + } + } + return { exhausted: false }; +} + +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 { exhausted: true }; + } + + const rateLimit = result.data.rate_limit ?? result.data.rate_limits; + if (!isRecord(rateLimit)) return { exhausted: false }; + for (const window of [ + rateLimit.primary_window ?? rateLimit.primary ?? rateLimit.five_hour_limit ?? rateLimit.five_hour, + rateLimit.secondary_window ?? rateLimit.secondary ?? rateLimit.weekly_limit ?? rateLimit.weekly, + ]) { + if (!isRecord(window)) continue; + const percentLeft = numeric(window.percent_left) ?? numeric(window.remaining_percent); + if (percentLeft !== undefined && percentLeft <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { + exhausted: true, + resetsAt: parseDateish(window.reset_at ?? window.reset_time_ms), + }; + } + } + return { exhausted: false }; +} + +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; + + const nested = isRecord(result.data.data) ? result.data.data : result.data; + const limits = Array.isArray(nested.limits) ? nested.limits : []; + for (const entry of limits) { + if (!isRecord(entry) || entry.type !== "TOKENS_LIMIT") continue; + const percentage = numeric(entry.percentage); + if (percentage !== undefined && percentage >= EXHAUSTED_UTILIZATION_PERCENT) { + return { exhausted: true, resetsAt: parseDateish(entry.nextResetTime) }; + } + } + return { exhausted: false }; +} + +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)) { + const limit = numeric(weekly.limit); + const used = numeric(weekly.used); + if (limit !== undefined && used !== undefined && limit > 0 && used >= limit) { + return { exhausted: true, resetsAt: parseDateish(weekly.resetTime) }; + } + } + return { exhausted: false }; +} + +async function fetchOpenRouterQuota( + modelRegistry: ModelRegistry, + deps: QuotaFetchDependencies, +): Promise { + const apiKey = await modelRegistry.getApiKeyForProvider("openrouter"); + if (!apiKey) return undefined; + const result = await fetchJson( + "https://openrouter.ai/api/v1/key", + { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, + deps.fetchImpl, + ); + if (!result.ok || !isRecord(result.data)) return undefined; + + const keyData = isRecord(result.data.data) ? result.data.data : undefined; + if (!keyData) return { exhausted: false }; + const limit = numeric(keyData.limit); + const limitRemaining = numeric(keyData.limit_remaining); + if (limit !== undefined && limit > 0 && limitRemaining !== undefined && limitRemaining <= 0) { + return { exhausted: true }; + } + return { exhausted: false }; +} + +/** + * Best-effort real quota reconciliation, keyed by Pi provider id. Providers without a known + * fetcher (e.g. Minimax, arbitrary OpenAI-compatible custom providers) 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, + openrouter: fetchOpenRouterQuota, +}; + +/** 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..6b95e6d --- /dev/null +++ b/extensions/auto-router-settings.ts @@ -0,0 +1,237 @@ +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; +}; + +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 } = value; + if (typeof provider !== "string" || !provider.trim()) return undefined; + if (typeof id !== "string" || !id.trim()) return undefined; + return { provider: provider.trim(), id: id.trim() }; +} + +/** 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; + } +} + +/** Persist the `autoRouter` setting without dropping keys owned by Pi or other extensions. */ +export async function writeAutoRouterSettingsFile( + settingsPath: string, + settings: AutoRouterSettings, + 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)}`, + ); + } + } + root[SETTINGS_KEY] = settings; + const tempPath = join( + settingsDir, + `.settings.${process.pid}.${io.randomUUID()}.tmp`, + ); + try { + await io.writeFile(tempPath, `${JSON.stringify(root, 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 package-specific global `autoRouter` setting without dropping unknown keys. */ +export async function writeAutoRouterSettings( + settings: AutoRouterSettings, +): Promise { + await writeAutoRouterSettingsFile(settingsPath(), settings); +} + +/** + * 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..dbf212b --- /dev/null +++ b/extensions/auto-router.ts @@ -0,0 +1,501 @@ +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 ModelHealthEntry, + type ModelIdentity, + modelKey, +} from "./auto-router-health.js"; +import { reconcileProviderQuota } from "./auto-router-quota.js"; +import { + AUTO_ROUTER_EFFORT_ORDER, + type AutoRouterEffortLevel, + type AutoRouterModelRef, + type AutoRouterSettings, + allConfiguredModels, + 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 AUTO_STATUS_KEY = "auto-router"; +const FOOTER_KEY = "auto-router"; + +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; +} + +/** 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; +} + +function restoreAutoActive(ctx: ExtensionContext): boolean { + 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; + return isRecord(data) && data.enabled === true; + } + } + return false; +} + +function formatTier(tier: AutoRouterEffortLevel): string { + return tier; +} + +function statusLine(tier: AutoRouterEffortLevel, model: ModelIdentity): string { + return `๐Ÿ”€ auto โ†’ ${formatTier(tier)} ยท ${model.id}`; +} + +export default function autoRouter(pi: ExtensionAPI): void { + pi.registerProvider(AUTO_PROVIDER_ID, { + name: "Auto", + // Never actually dispatched to: `before_agent_start` always swaps to a real + // routed model before any request would be sent here. + baseUrl: "http://127.0.0.1:0", + apiKey: "auto-router", + api: "openai-completions", + models: [ + { + id: AUTO_MODEL_ID, + name: "Auto", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128_000, + maxTokens: 4096, + }, + ], + }); + + let currentSessionId: string | undefined; + let autoActive = false; + let routingInFlight = false; + let currentInFlightModel: ModelIdentity | undefined; + let healthStore = new AutoRouterHealthStore(); + + function publishFooter(tier: AutoRouterEffortLevel): void { + if (!currentSessionId) return; + pi.events.emit(FOOTER_CONTRIBUTION_EVENT, { + sessionId: currentSessionId, + key: FOOTER_KEY, + identitySuffix: (theme: Theme) => + theme.fg("accent", `๐Ÿ”€${formatTier(tier)}`), + } satisfies FooterContribution); + } + + function clearFooter(): void { + if (!currentSessionId) return; + pi.events.emit(FOOTER_CONTRIBUTION_EVENT, { + sessionId: currentSessionId, + key: FOOTER_KEY, + remove: true, + } satisfies FooterContribution); + } + + /** 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 } | 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 }; + } + } + // Last resort: nothing healthy anywhere. Use the first resolvable model in `tier`, + // or failing that the first resolvable model anywhere, rather than blocking the turn. + const fallbackRefs = + settings.efforts[tier]?.models ?? allConfiguredModels(settings); + const fallback = + resolveAvailableModels(ctx.modelRegistry, fallbackRefs)[0] ?? + resolveAvailableModels( + ctx.modelRegistry, + allConfiguredModels(settings), + )[0]; + 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 }; + } + return undefined; + } + + async function applyRouting( + pi: ExtensionAPI, + ctx: ExtensionContext, + model: Model, + tier: 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(tier); + } finally { + routingInFlight = false; + } + currentInFlightModel = { provider: model.provider, id: model.id }; + if (ctx.hasUI) ctx.ui.setStatus(AUTO_STATUS_KEY, statusLine(tier, model)); + publishFooter(tier); + } + + async function activateDefault( + pi: ExtensionAPI, + ctx: ExtensionContext, + ): Promise { + const settings = await readAutoRouterSettings(); + const picked = pickForTier(ctx, settings, "medium"); + 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; + } + await applyRouting(pi, ctx, picked.model, picked.tier); + } + + async function routeForPrompt( + pi: ExtensionAPI, + ctx: ExtensionContext, + prompt: string, + hasImages: boolean, + ): Promise { + const settings = await readAutoRouterSettings(); + 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 level: AutoRouterEffortLevel = "medium"; + if (classifierModel) { + const result = await classifyTurnComplexity( + ctx.modelRegistry, + classifierModel, + prompt, + hasImages, + ); + level = result.level; + if (result.usage) { + healthStore.recordSuccess(modelKey(classifierModel), result.usage); + } + } + + const tier = resolveEffortTier(settings, level); + 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; + } + await applyRouting(pi, ctx, picked.model, picked.tier); + } + + 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; + healthStore.applyQuotaResult(modelKey(model), result); + } + }), + ); + } + + pi.on("model_select", async (event, ctx) => { + if (routingInFlight) return; + if (event.model.provider === AUTO_PROVIDER_ID) { + autoActive = true; + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true }); + await activateDefault(pi, ctx); + return; + } + if (event.source !== "restore" && autoActive) { + autoActive = false; + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: false }); + currentInFlightModel = undefined; + if (ctx.hasUI) ctx.ui.setStatus(AUTO_STATUS_KEY, undefined); + clearFooter(); + } + }); + + pi.on("session_start", async (_event, ctx) => { + currentSessionId = ctx.sessionManager.getSessionId(); + routingInFlight = false; + currentInFlightModel = undefined; + healthStore = new AutoRouterHealthStore(); + await healthStore.load(); + autoActive = restoreAutoActive(ctx); + if ( + autoActive && + ctx.model && + ctx.model.provider !== AUTO_PROVIDER_ID && + ctx.thinkingLevel + ) { + if (ctx.hasUI) + ctx.ui.setStatus( + AUTO_STATUS_KEY, + statusLine(ctx.thinkingLevel, ctx.model), + ); + currentInFlightModel = { provider: ctx.model.provider, id: ctx.model.id }; + publishFooter(ctx.thinkingLevel); + } + const settings = await readAutoRouterSettings(); + void reconcileAllProviders(ctx.modelRegistry, settings); + }); + + pi.on("before_agent_start", async (event, ctx) => { + if (!autoActive) return; + await routeForPrompt(pi, ctx, event.prompt, Boolean(event.images?.length)); + }); + + pi.on("after_provider_response", (event) => { + if (!currentInFlightModel) return; + if (event.status >= 200 && event.status < 300) return; + healthStore.recordFailure( + modelKey(currentInFlightModel), + event.status, + event.headers, + ); + }); + + pi.on("message_end", (event) => { + if (!autoActive || !currentInFlightModel) return; + if (event.message.role !== "assistant") return; + const usage = isRecord(event.message) ? event.message.usage : undefined; + healthStore.recordSuccess(modelKey(currentInFlightModel), { + 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; + currentInFlightModel = undefined; + }); + + 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); + if (ctx.mode === "tui") { + await showUsageDashboard(rows, ctx); + } else if (ctx.hasUI) { + ctx.ui.notify(formatUsagePlainText(rows), "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`; +} + +function rowStatus(entry: ModelHealthEntry | undefined, now: number): string { + if (!entry) return "unused"; + if (entry.cooldownUntil && entry.cooldownUntil > now) { + const minutes = Math.max( + 1, + Math.round((entry.cooldownUntil - now) / 60_000), + ); + const cause = entry.lastError ? ` (${entry.lastError.status})` : ""; + return `cooldown${cause} ~${minutes}m`; + } + return "healthy"; +} + +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), + ); + const verified = entry?.verifiedAt ? "โœ“" : "~"; + return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${requests} req ยท ${tokens} tok ${verified}`; +} + +function formatUsagePlainText(rows: UsageRow[]): 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(" | ")}`, + ); + } + return lines.join("\n"); +} + +function formatUsageMarkdown(rows: UsageRow[]): string { + if (rows.length === 0) return "No models configured."; + const now = Date.now(); + const lines = [ + "| Tier | Model | Status | Req | Tokens | 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), + ); + const verified = entry?.verifiedAt ? " โœ“" : ""; + lines.push( + `| ${row.tier} | ${row.model.provider}/${row.model.id} | ${rowStatus(entry, now)}${verified} | ${entry?.totals.requests ?? 0} | ${tokens} | ${cost} |`, + ); + } + return lines.join("\n"); +} + +async function showUsageDashboard( + rows: UsageRow[], + 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), 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-extension.test.ts b/tests/auto-router-extension.test.ts new file mode 100644 index 0000000..eba3bb7 --- /dev/null +++ b/tests/auto-router-extension.test.ts @@ -0,0 +1,332 @@ +import { afterEach, 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, +} from "@earendil-works/pi-coding-agent"; +import autoRouter from "../extensions/auto-router.ts"; +import type { AutoRouterSettings } from "../extensions/auto-router-settings.ts"; + +const ENV_VAR = "PI_CODING_AGENT_DIR"; +let previousEnv: string | undefined; +let agentDir: string | undefined; + +beforeEach(async () => { + previousEnv = process.env[ENV_VAR]; + agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-agent-")); + process.env[ENV_VAR] = agentDir; +}); + +afterEach(async () => { + if (previousEnv === undefined) delete process.env[ENV_VAR]; + else process.env[ENV_VAR] = previousEnv; + if (agentDir) await rm(agentDir, { recursive: true, force: true }); + agentDir = undefined; +}); + +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; +} + +type FakeHandler = (event: unknown, ctx: unknown) => unknown; + +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 }>; +}; + +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 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: () => undefined, on: () => () => undefined }, + setModel: async (m: Model) => { + setModelCalls.push(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, + }; +} + +type FakeRegistryOptions = { + models: Model[]; + unavailable?: Model[]; + classify?: (prompt: string) => string; +}; + +function fakeModelRegistry({ models, unavailable = [], classify }: FakeRegistryOptions) { + const unavailableKeys = new Set(unavailable.map((m) => `${m.provider}/${m.id}`)); + return { + find: (provider: string, id: string) => models.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; + model?: Model; + thinkingLevel?: string; + entries?: unknown[]; +}) { + const notifications: Array<{ message: string; type?: string }> = []; + const statuses = new Map(); + return { + hasUI: true, + mode: "rpc", + model: options.model, + thinkingLevel: options.thinkingLevel, + modelRegistry: options.modelRegistry, + sessionManager: { + getSessionId: () => "session-1", + getEntries: () => options.entries ?? [], + }, + ui: { + notify: (message: string, type?: string) => notifications.push({ message, type }), + setStatus: (_key: string, text: string | undefined) => statuses.set(_key, text), + custom: async () => undefined, + }, + notifications, + statuses, + } as unknown as ExtensionContext & { notifications: typeof notifications; statuses: typeof statuses }; +} + +test("selecting Auto immediately routes to the first available medium-tier model", 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(); + autoRouter(fake.pi); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a, b] }) }); + + await fake.fire("session_start", {}, ctx); + await fake.fire( + "model_select", + { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, + ctx, + ); + + expect(fake.setModelCalls).toEqual([a]); + expect(fake.thinkingLevelCalls).toEqual(["medium"]); +}); + +test("before_agent_start classifies the turn and routes to the matching tier", 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(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ + models: [medium, high], + classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), + }); + const ctx = fakeCtx({ modelRegistry: registry }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); + fake.setModelCalls.length = 0; + fake.thinkingLevelCalls.length = 0; + + await fake.fire("before_agent_start", { prompt: "please refactor this multi-file module" }, ctx); + + expect(fake.setModelCalls).toEqual([high]); + expect(fake.thinkingLevelCalls).toEqual(["high"]); +}); + +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(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [medium, c, d], classify: () => "high" }); + const ctx = fakeCtx({ modelRegistry: registry }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); + fake.setModelCalls.length = 0; + + // 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(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [only] }); + const ctx = fakeCtx({ modelRegistry: registry }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); + fake.setModelCalls.length = 0; + + await fake.fire("after_provider_response", { status: 401, headers: {} }, ctx); + 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("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(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, b] }); + const ctx = fakeCtx({ modelRegistry: registry }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, 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("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(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [a, manual] }); + const ctx = fakeCtx({ modelRegistry: registry }); + + await fake.fire("session_start", {}, ctx); + await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, 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("/usage with no configured models notifies instead of throwing", async () => { + const fake = createFakePi(); + 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 restores an active Auto session without re-routing", async () => { + const already = model("prov", "already-selected"); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "already-selected" }] } } }); + + const fake = createFakePi(); + autoRouter(fake.pi); + const registry = fakeModelRegistry({ models: [already] }); + const ctx = fakeCtx({ + modelRegistry: registry, + model: already, + thinkingLevel: "medium", + entries: [{ type: "custom", customType: "vessup:auto-router:active", data: { enabled: true } }], + }); + + await fake.fire("session_start", {}, ctx); + + expect(fake.setModelCalls).toEqual([]); + expect(ctx.statuses.get("auto-router")).toContain("already-selected"); +}); diff --git a/tests/auto-router-health.test.ts b/tests/auto-router-health.test.ts new file mode 100644 index 0000000..023c262 --- /dev/null +++ b/tests/auto-router-health.test.ts @@ -0,0 +1,141 @@ +import { expect, test } from "bun:test"; +import { + applyFailure, + applyQuotaResult, + applySuccess, + type AutoRouterHealthState, + isHealthy, + modelKey, + parseRetryAfterMs, + pickHealthy, +} from "../extensions/auto-router-health.ts"; + +const NOW = 1_000_000_000_000; + +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); +}); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts new file mode 100644 index 0000000..4c6ad5e --- /dev/null +++ b/tests/auto-router-quota.test.ts @@ -0,0 +1,125 @@ +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 }: { accountId?: string } = { accountId: "acct-1" }, +): QuotaFetchDependencies { + return { + fetchImpl: (async (url: string | URL) => handler(url.toString())) as typeof fetch, + readCodexAccountId: async () => accountId, + }; +} + +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("minimax", 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({ exhausted: true, resetsAt: Date.parse(resetsAt) }); +}); + +test("reconcileProviderQuota(anthropic) reports headroom when utilization is low", 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({ exhausted: false }); +}); + +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({ exhausted: true }); +}); + +test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit window is depleted", 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({ exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z") }); +}); + +test("reconcileProviderQuota(zai) reports exhaustion from a TOKENS_LIMIT entry", async () => { + const deps = fakeDeps(() => + jsonResponse({ data: { limits: [{ type: "TOKENS_LIMIT", percentage: 100, nextResetTime: 4_102_444_800_000 }] } }), + ); + const result = await reconcileProviderQuota("zai", fakeRegistry("key"), deps); + expect(result).toEqual({ exhausted: true, resetsAt: 4_102_444_800_000 }); +}); + +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({ exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z") }); +}); + +test("reconcileProviderQuota(openrouter) reports exhaustion when no budget remains", async () => { + const deps = fakeDeps(() => jsonResponse({ data: { limit: 50, limit_remaining: 0 } })); + const result = await reconcileProviderQuota("openrouter", fakeRegistry("key"), deps); + expect(result).toEqual({ exhausted: true }); +}); + +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", + }; + 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..4edadaa --- /dev/null +++ b/tests/auto-router-settings.test.ts @@ -0,0 +1,171 @@ +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, + type AutoRouterSettings, + 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 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 }); +}); From 1a1e74a7495bdd0a29fbd156611e3c99b4a71e30 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 16:55:45 -0400 Subject: [PATCH 02/21] Expand classifier vocabulary, fix xhigh parsing, drop OpenRouter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The complexity classifier now covers the full minimal..max range instead of just low/medium/high/xhigh (all seven tiers were already configurable; only the classifier's own vocabulary was narrower). Fixes a substring-matching bug this surfaced: "xhigh" contains "high", so replies of "xhigh" were silently parsed as "high" โ€” now matched on word boundaries. Also drops the OpenRouter quota fetcher; it wasn't wanted and isn't part of the personal config this was built around. Co-Authored-By: Claude Sonnet 5 --- README.md | 4 +- extensions/auto-router-classify.ts | 12 +++-- extensions/auto-router-quota.ts | 24 --------- tests/auto-router-classify.test.ts | 84 ++++++++++++++++++++++++++++++ tests/auto-router-quota.test.ts | 6 --- 5 files changed, 95 insertions(+), 35 deletions(-) create mode 100644 tests/auto-router-classify.test.ts diff --git a/README.md b/README.md index 4765bb7..b0956c0 100644 --- a/README.md +++ b/README.md @@ -62,9 +62,9 @@ Configure it under a new `autoRouter` key in `~/.pi/agent/settings.json` (or `.p 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. -On every turn, Auto asks the `medium` tier's first healthy model (the "default model") to classify the turn as `low`, `medium`, `high`, or `xhigh`, 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. +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. -Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota API (currently Anthropic, OpenAI Codex, Z.ai, Kimi Coding, and OpenRouter) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. +Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota API (currently Anthropic, OpenAI Codex, Z.ai, and Kimi Coding) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. Providers without a known quota API (Minimax included โ€” it doesn't currently expose one) simply stay on router-observed data. Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally; the footer (and Pi Web's model display) always shows the real underlying model and effort actually in use, plus a small `๐Ÿ”€` badge in the TUI footer while Auto is engaged. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. diff --git a/extensions/auto-router-classify.ts b/extensions/auto-router-classify.ts index 635af18..84206a0 100644 --- a/extensions/auto-router-classify.ts +++ b/extensions/auto-router-classify.ts @@ -5,19 +5,23 @@ 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: low, medium, high, or xhigh. +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. -- low: trivial, mechanical, or purely informational. One-line edits, formatting, simple lookups, answering a quick factual question about the codebase. +- 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.`; @@ -81,8 +85,10 @@ export async function classifyTurnComplexity( .join("") .trim() .toLowerCase(); + // Word-boundary match, not plain substring: "high" is a substring of "xhigh", + // so a naive `.includes()` would mis-parse an "xhigh" reply as "high". const level = - VALID_LEVELS.find((candidate) => reply.includes(candidate)) ?? + VALID_LEVELS.find((candidate) => new RegExp(`\\b${candidate}\\b`).test(reply)) ?? DEFAULT_LEVEL; const usage = response.usage ? { diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 42f48d6..4df46d8 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -199,29 +199,6 @@ async function fetchKimiCodingQuota( return { exhausted: false }; } -async function fetchOpenRouterQuota( - modelRegistry: ModelRegistry, - deps: QuotaFetchDependencies, -): Promise { - const apiKey = await modelRegistry.getApiKeyForProvider("openrouter"); - if (!apiKey) return undefined; - const result = await fetchJson( - "https://openrouter.ai/api/v1/key", - { Authorization: `Bearer ${apiKey}`, Accept: "application/json" }, - deps.fetchImpl, - ); - if (!result.ok || !isRecord(result.data)) return undefined; - - const keyData = isRecord(result.data.data) ? result.data.data : undefined; - if (!keyData) return { exhausted: false }; - const limit = numeric(keyData.limit); - const limitRemaining = numeric(keyData.limit_remaining); - if (limit !== undefined && limit > 0 && limitRemaining !== undefined && limitRemaining <= 0) { - return { exhausted: true }; - } - return { exhausted: false }; -} - /** * Best-effort real quota reconciliation, keyed by Pi provider id. Providers without a known * fetcher (e.g. Minimax, arbitrary OpenAI-compatible custom providers) simply have no entry @@ -238,7 +215,6 @@ export const QUOTA_FETCHERS: Record< "openai-codex": fetchCodexQuota, zai: fetchZaiQuota, "kimi-coding": fetchKimiCodingQuota, - openrouter: fetchOpenRouterQuota, }; /** Reconcile one provider's real quota state. Never throws; returns `undefined` when there's no known fetcher, no credentials, or the request failed. */ diff --git a/tests/auto-router-classify.test.ts b/tests/auto-router-classify.test.ts new file mode 100644 index 0000000..5b8bd0b --- /dev/null +++ b/tests/auto-router-classify.test.ts @@ -0,0 +1,84 @@ +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 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", async () => { + const result = await classifyTurnComplexity(throwingRegistry(), MODEL, "do something", false); + expect(result.level).toBe("medium"); + expect(result.usage).toBeUndefined(); +}); + +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-quota.test.ts b/tests/auto-router-quota.test.ts index 4c6ad5e..492669c 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -101,12 +101,6 @@ test("reconcileProviderQuota(kimi-coding) reports exhaustion when used reaches t expect(result).toEqual({ exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z") }); }); -test("reconcileProviderQuota(openrouter) reports exhaustion when no budget remains", async () => { - const deps = fakeDeps(() => jsonResponse({ data: { limit: 50, limit_remaining: 0 } })); - const result = await reconcileProviderQuota("openrouter", fakeRegistry("key"), deps); - expect(result).toEqual({ exhausted: true }); -}); - test("reconcileProviderQuota degrades gracefully on network/HTTP failure", async () => { const deps = fakeDeps(() => jsonResponse({}, 500)); const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); From 51a543c9e547ca02249a42c73e246a5d2c8581d4 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 17:07:36 -0400 Subject: [PATCH 03/21] Add real Minimax quota reconciliation via the mmx CLI MiniMax has no documented HTTP quota endpoint, but its own mmx CLI does (confirmed via `mmx --verbose`: GET /v1/token_plan/remains with mmx's own OAuth session). Shell out to `mmx quota show --output json` rather than reading its private token cache, so mmx keeps owning token refresh/expiry; missing or logged-out CLI degrades gracefully to router-observed data like any other unsupported provider. Co-Authored-By: Claude Sonnet 5 --- README.md | 3 +- extensions/auto-router-quota.ts | 66 ++++++++++++++++++++++++++-- tests/auto-router-quota.test.ts | 77 ++++++++++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index b0956c0..438bcce 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ Each tier key is a Pi thinking level (`off`, `minimal`, `low`, `medium`, `high`, 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. -Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota API (currently Anthropic, OpenAI Codex, Z.ai, and Kimi Coding) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. Providers without a known quota API (Minimax included โ€” it doesn't currently expose one) simply stay on router-observed data. +Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota source (currently Anthropic, OpenAI Codex, Z.ai, and Kimi Coding via their HTTP APIs, plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no documented HTTP quota endpoint) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. Providers without a known quota source simply stay on router-observed data. Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally; the footer (and Pi Web's model display) always shows the real underlying model and effort actually in use, plus a small `๐Ÿ”€` badge in the TUI footer while Auto is engaged. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. @@ -72,6 +72,7 @@ Run `/usage` to see health/usage for every configured model, grouped by tier, in - 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 diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 4df46d8..cc37cc4 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -1,18 +1,24 @@ +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 }; -/** Injectable for tests; defaults to the real network/filesystem. */ +/** 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 { @@ -35,9 +41,22 @@ async function defaultReadCodexAccountId(): Promise { } } +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( @@ -199,10 +218,50 @@ async function fetchKimiCodingQuota( return { exhausted: false }; } +/** + * 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"); + const entries = isRecord(general) ? [general] : data.model_remains; + for (const entry of entries) { + if (!isRecord(entry)) continue; + // The API reports *remaining* percent (opposite convention from the other providers' + // *used* percent), and tracks a short rolling interval plus a weekly window separately. + const intervalRemaining = numeric(entry.current_interval_remaining_percent); + if (intervalRemaining !== undefined && intervalRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { exhausted: true, resetsAt: numeric(entry.end_time) }; + } + const weeklyRemaining = numeric(entry.current_weekly_remaining_percent); + if (weeklyRemaining !== undefined && weeklyRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { exhausted: true, resetsAt: numeric(entry.weekly_end_time) }; + } + } + return { exhausted: false }; +} + /** * Best-effort real quota reconciliation, keyed by Pi provider id. Providers without a known - * fetcher (e.g. Minimax, arbitrary OpenAI-compatible custom providers) simply have no entry - * here โ€” callers treat a missing/failed fetch as "no correction available", never as failure. + * 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, @@ -215,6 +274,7 @@ export const QUOTA_FETCHERS: Record< "openai-codex": fetchCodexQuota, zai: fetchZaiQuota, "kimi-coding": fetchKimiCodingQuota, + minimax: fetchMinimaxQuota, }; /** Reconcile one provider's real quota state. Never throws; returns `undefined` when there's no known fetcher, no credentials, or the request failed. */ diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index 492669c..dd67b9c 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -13,11 +13,17 @@ function fakeRegistry(apiKey: string | undefined): ModelRegistry { function fakeDeps( handler: (url: string) => Response, - { accountId }: { accountId?: string } = { accountId: "acct-1" }, + { + 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), }; } @@ -26,7 +32,11 @@ function jsonResponse(data: unknown, status = 200): Response { } test("reconcileProviderQuota returns undefined for providers without a known fetcher", async () => { - const result = await reconcileProviderQuota("minimax", fakeRegistry("key"), fakeDeps(() => jsonResponse({}))); + const result = await reconcileProviderQuota( + "totally-unsupported-provider", + fakeRegistry("key"), + fakeDeps(() => jsonResponse({})), + ); expect(result).toBeUndefined(); }); @@ -101,6 +111,68 @@ test("reconcileProviderQuota(kimi-coding) reports exhaustion when used reaches t expect(result).toEqual({ exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z") }); }); +test("reconcileProviderQuota(minimax) reports headroom from the mmx CLI's general bucket", async () => { + const deps = fakeDeps(() => jsonResponse({}), { + minimaxCli: async () => + JSON.stringify({ + model_remains: [ + { model_name: "general", 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({ exhausted: false }); +}); + +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({ exhausted: true, resetsAt: 4_102_444_800_000 }); +}); + +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({ exhausted: true, resetsAt: 4_102_444_800_000 }); +}); + +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 degrades gracefully on network/HTTP failure", async () => { const deps = fakeDeps(() => jsonResponse({}, 500)); const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); @@ -113,6 +185,7 @@ test("reconcileProviderQuota degrades gracefully when the fetcher throws", 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(); From 9367c08aab3848e7145a0af9cb15c3c539b36a43 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 17:21:18 -0400 Subject: [PATCH 04/21] Keep Auto visible when /model is scoped by enabledModels Pi's /model picker defaults to showing only the enabledModels-scoped list when that setting is non-empty, hiding everything else - including our own registered "auto" model - behind a manual Tab to "all". At session start, best-effort append an "auto/auto" pattern to enabledModels (only when scoping is already configured, and only if not already present) so Auto shows up by default instead of being invisible for anyone who has scoped their model list. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 + extensions/auto-router-settings.ts | 59 +++++++++++++++++++++++++++--- extensions/auto-router.ts | 2 + tests/auto-router-settings.test.ts | 48 ++++++++++++++++++++++++ 4 files changed, 106 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 438bcce..3c293fd 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,8 @@ Health is tracked from the router's own observed traffic (HTTP status codes, rat Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally; the footer (and Pi Web's model display) always shows the real underlying model and effort actually in use, plus a small `๐Ÿ”€` badge in the TUI footer while Auto is engaged. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" 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 Auto โ€” behind a manual Tab to "all". At session start, Auto best-effort appends its own `auto/auto` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so it shows up in the default scoped view too, without changing anything else about what's scoped. + ### Requirements - Pi 0.84.1 diff --git a/extensions/auto-router-settings.ts b/extensions/auto-router-settings.ts index 6b95e6d..d4d9dc2 100644 --- a/extensions/auto-router-settings.ts +++ b/extensions/auto-router-settings.ts @@ -114,10 +114,14 @@ export async function readAutoRouterSettings(): Promise { } } -/** Persist the `autoRouter` setting without dropping keys owned by Pi or other extensions. */ -export async function writeAutoRouterSettingsFile( +/** + * 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, - settings: AutoRouterSettings, + mutate: (root: Record) => Record | undefined, dependencies: Partial = {}, ): Promise { const io = { ...defaultWriteDependencies, ...dependencies }; @@ -148,13 +152,14 @@ export async function writeAutoRouterSettingsFile( ); } } - root[SETTINGS_KEY] = settings; + const next = mutate(root); + if (!next) return; const tempPath = join( settingsDir, `.settings.${process.pid}.${io.randomUUID()}.tmp`, ); try { - await io.writeFile(tempPath, `${JSON.stringify(root, null, 2)}\n`, { + await io.writeFile(tempPath, `${JSON.stringify(next, null, 2)}\n`, { mode: 0o600, flag: "wx", }); @@ -167,6 +172,19 @@ export async function writeAutoRouterSettingsFile( } } +/** 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, @@ -174,6 +192,37 @@ export async function writeAutoRouterSettings( await writeAutoRouterSettingsFile(settingsPath(), settings); } +/** Pattern that matches our registered virtual "Auto" model in `/model`'s scoping patterns. */ +export const AUTO_MODEL_SCOPE_PATTERN = "auto/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" + * entry) 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. diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index dbf212b..6a2d6ae 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -23,6 +23,7 @@ import { type AutoRouterModelRef, type AutoRouterSettings, allConfiguredModels, + ensureAutoModelScopedInGlobalSettings, escalationTiers, readAutoRouterSettings, resolveEffortTier, @@ -321,6 +322,7 @@ export default function autoRouter(pi: ExtensionAPI): void { } const settings = await readAutoRouterSettings(); void reconcileAllProviders(ctx.modelRegistry, settings); + void ensureAutoModelScopedInGlobalSettings().catch(() => undefined); }); pi.on("before_agent_start", async (event, ctx) => { diff --git a/tests/auto-router-settings.test.ts b/tests/auto-router-settings.test.ts index 4edadaa..fecc33b 100644 --- a/tests/auto-router-settings.test.ts +++ b/tests/auto-router-settings.test.ts @@ -4,7 +4,9 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { allConfiguredModels, + AUTO_MODEL_SCOPE_PATTERN, type AutoRouterSettings, + ensureAutoModelScoped, escalationTiers, parseAutoRouterSettings, resolveEffortTier, @@ -169,3 +171,49 @@ test("malformed settings reject without leaking the cross-process lock", async ( 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, + }); +}); From 08dfe655cdd0d805602496688a38d5bf6b165bbb Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 17:33:43 -0400 Subject: [PATCH 05/21] Keep /model showing Auto selected instead of the routed model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the real model stayed selected once Auto routed a turn, so reopening /model showed e.g. "MiniMax M3" instead of "Auto" - Auto's own selection only "stuck" for the very first turn, since pi.setModel just swaps ctx.model to whatever we route to and leaves it there. Now the real model is only swapped in for the duration of each turn: before_agent_start routes to it as before, and a new agent_settled handler swaps back to the inert Auto placeholder once the turn is fully done (including any retries/continuations), so /model shows Auto again between turns. Session restore mirrors this, reverting on resume if a session was interrupted mid-turn before agent_settled could fire. The footer badge (also reworked per feedback into a single "๐Ÿ”€ Auto ()" line instead of a separate status line) tracks the last-used tier independently, so it keeps showing useful information across the revert. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 +- extensions/auto-router.ts | 81 ++++++++------- tests/auto-router-extension.test.ts | 148 +++++++++++++++++++++------- 3 files changed, 151 insertions(+), 80 deletions(-) diff --git a/README.md b/README.md index 3c293fd..6820fcf 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ On every turn, Auto asks the `medium` tier's first healthy model (the "default m Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota source (currently Anthropic, OpenAI Codex, Z.ai, and Kimi Coding via their HTTP APIs, plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no documented HTTP quota endpoint) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. Providers without a known quota source simply stay on router-observed data. -Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally; the footer (and Pi Web's model display) always shows the real underlying model and effort actually in use, plus a small `๐Ÿ”€` badge in the TUI footer while Auto is engaged. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. +Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" 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 Auto โ€” behind a manual Tab to "all". At session start, Auto best-effort appends its own `auto/auto` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so it shows up in the default scoped view too, without changing anything else about what's scoped. diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 6a2d6ae..5cc774d 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -36,7 +36,6 @@ import { const AUTO_PROVIDER_ID = "auto"; const AUTO_MODEL_ID = "auto"; const AUTO_ACTIVE_ENTRY_TYPE = "vessup:auto-router:active"; -const AUTO_STATUS_KEY = "auto-router"; const FOOTER_KEY = "auto-router"; function numeric(value: unknown): number { @@ -79,8 +78,8 @@ function formatTier(tier: AutoRouterEffortLevel): string { return tier; } -function statusLine(tier: AutoRouterEffortLevel, model: ModelIdentity): string { - return `๐Ÿ”€ auto โ†’ ${formatTier(tier)} ยท ${model.id}`; +function footerBadge(tier: AutoRouterEffortLevel | undefined): string { + return tier ? `๐Ÿ”€ Auto (${formatTier(tier)})` : "๐Ÿ”€ Auto"; } export default function autoRouter(pi: ExtensionAPI): void { @@ -108,18 +107,33 @@ export default function autoRouter(pi: ExtensionAPI): void { let autoActive = false; let routingInFlight = false; let currentInFlightModel: ModelIdentity | undefined; + let lastKnownTier: AutoRouterEffortLevel | undefined; let healthStore = new AutoRouterHealthStore(); - function publishFooter(tier: AutoRouterEffortLevel): void { + function publishFooter(tier: AutoRouterEffortLevel | undefined): void { if (!currentSessionId) return; pi.events.emit(FOOTER_CONTRIBUTION_EVENT, { sessionId: currentSessionId, key: FOOTER_KEY, - identitySuffix: (theme: Theme) => - theme.fg("accent", `๐Ÿ”€${formatTier(tier)}`), + identitySuffix: (theme: Theme) => theme.fg("accent", footerBadge(tier)), } satisfies FooterContribution); } + /** Swap `ctx.model` back to the inert "auto" placeholder once a turn is fully done, so `/model` keeps showing Auto selected (not whichever real model just handled the turn) between turns. */ + async function revertToAutoPlaceholder( + pi: ExtensionAPI, + ctx: ExtensionContext, + ): Promise { + const placeholder = ctx.modelRegistry.find(AUTO_PROVIDER_ID, AUTO_MODEL_ID); + 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, { @@ -193,28 +207,10 @@ export default function autoRouter(pi: ExtensionAPI): void { routingInFlight = false; } currentInFlightModel = { provider: model.provider, id: model.id }; - if (ctx.hasUI) ctx.ui.setStatus(AUTO_STATUS_KEY, statusLine(tier, model)); + lastKnownTier = tier; publishFooter(tier); } - async function activateDefault( - pi: ExtensionAPI, - ctx: ExtensionContext, - ): Promise { - const settings = await readAutoRouterSettings(); - const picked = pickForTier(ctx, settings, "medium"); - 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; - } - await applyRouting(pi, ctx, picked.model, picked.tier); - } - async function routeForPrompt( pi: ExtensionAPI, ctx: ExtensionContext, @@ -282,19 +278,19 @@ export default function autoRouter(pi: ExtensionAPI): void { ); } - pi.on("model_select", async (event, ctx) => { + pi.on("model_select", (event, _ctx) => { if (routingInFlight) return; if (event.model.provider === AUTO_PROVIDER_ID) { autoActive = true; pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true }); - await activateDefault(pi, ctx); + publishFooter(lastKnownTier); return; } if (event.source !== "restore" && autoActive) { autoActive = false; pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: false }); currentInFlightModel = undefined; - if (ctx.hasUI) ctx.ui.setStatus(AUTO_STATUS_KEY, undefined); + lastKnownTier = undefined; clearFooter(); } }); @@ -303,28 +299,31 @@ export default function autoRouter(pi: ExtensionAPI): void { currentSessionId = ctx.sessionManager.getSessionId(); routingInFlight = false; currentInFlightModel = undefined; + lastKnownTier = undefined; healthStore = new AutoRouterHealthStore(); await healthStore.load(); autoActive = restoreAutoActive(ctx); - if ( - autoActive && - ctx.model && - ctx.model.provider !== AUTO_PROVIDER_ID && - ctx.thinkingLevel - ) { - if (ctx.hasUI) - ctx.ui.setStatus( - AUTO_STATUS_KEY, - statusLine(ctx.thinkingLevel, ctx.model), - ); - currentInFlightModel = { provider: ctx.model.provider, id: ctx.model.id }; - publishFooter(ctx.thinkingLevel); + 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. + currentInFlightModel = { provider: ctx.model.provider, id: ctx.model.id }; + lastKnownTier = ctx.thinkingLevel; + await revertToAutoPlaceholder(pi, ctx); + } + publishFooter(lastKnownTier); } 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) return; await routeForPrompt(pi, ctx, event.prompt, Boolean(event.images?.length)); diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index eba3bb7..77250e0 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -7,6 +7,7 @@ import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext, + Theme, } from "@earendil-works/pi-coding-agent"; import autoRouter from "../extensions/auto-router.ts"; import type { AutoRouterSettings } from "../extensions/auto-router-settings.ts"; @@ -40,7 +41,10 @@ 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; @@ -49,6 +53,8 @@ type FakePi = { setModelCalls: Model[]; thinkingLevelCalls: string[]; appendedEntries: Array<{ type: string; data: unknown }>; + footerEvents: unknown[]; + currentModel: ModelRef; }; function createFakePi(): FakePi { @@ -57,6 +63,8 @@ function createFakePi(): FakePi { 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, @@ -70,9 +78,10 @@ function createFakePi(): FakePi { appendEntry: (type: string, data: unknown) => { appendedEntries.push({ type, data }); }, - events: { emit: () => undefined, on: () => () => undefined }, + 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) => { @@ -95,9 +104,18 @@ function createFakePi(): FakePi { 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[]; @@ -105,9 +123,11 @@ type FakeRegistryOptions = { }; 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) => models.find((m) => m.provider === provider && m.id === id), + find: (provider: string, id: string) => 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 }> }> }) => { @@ -120,16 +140,19 @@ function fakeModelRegistry({ models, unavailable = [], classify }: FakeRegistryO function fakeCtx(options: { modelRegistry: ReturnType; + currentModel?: ModelRef; model?: Model; thinkingLevel?: string; entries?: unknown[]; }) { const notifications: Array<{ message: string; type?: string }> = []; - const statuses = new Map(); + const modelRef = options.currentModel ?? { value: options.model }; return { hasUI: true, mode: "rpc", - model: options.model, + get model() { + return modelRef.value; + }, thinkingLevel: options.thinkingLevel, modelRegistry: options.modelRegistry, sessionManager: { @@ -138,35 +161,34 @@ function fakeCtx(options: { }, ui: { notify: (message: string, type?: string) => notifications.push({ message, type }), - setStatus: (_key: string, text: string | undefined) => statuses.set(_key, text), + setStatus: () => undefined, custom: async () => undefined, }, notifications, - statuses, - } as unknown as ExtensionContext & { notifications: typeof notifications; statuses: typeof statuses }; + } 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); } -test("selecting Auto immediately routes to the first available medium-tier model", async () => { +test("selecting Auto marks it active without eagerly routing, showing a neutral footer badge", 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" }] } } }); + await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); autoRouter(fake.pi); - const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a, b] }) }); + const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a] }), currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire( - "model_select", - { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, - ctx, - ); + await selectAuto(fake, ctx); - expect(fake.setModelCalls).toEqual([a]); - expect(fake.thinkingLevelCalls).toEqual(["medium"]); + expect(fake.setModelCalls).toEqual([]); + expect(fake.thinkingLevelCalls).toEqual([]); + expect(lastFooterBadge(fake.footerEvents)).toBe("๐Ÿ”€ Auto"); }); -test("before_agent_start classifies the turn and routes to the matching tier", async () => { +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({ @@ -182,17 +204,25 @@ test("before_agent_start classifies the turn and routes to the matching tier", a models: [medium, high], classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), }); - const ctx = fakeCtx({ modelRegistry: registry }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); - fake.setModelCalls.length = 0; - fake.thinkingLevelCalls.length = 0; + 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("routing escalates to a higher tier when the classified tier is entirely unhealthy", async () => { @@ -210,11 +240,10 @@ test("routing escalates to a higher tier when the classified tier is entirely un const fake = createFakePi(); autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [medium, c, d], classify: () => "high" }); - const ctx = fakeCtx({ modelRegistry: registry }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); - fake.setModelCalls.length = 0; + 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); @@ -239,13 +268,15 @@ test("routing falls back to the classified tier's model as a last resort when no const fake = createFakePi(); autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [only] }); - const ctx = fakeCtx({ modelRegistry: registry }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); - fake.setModelCalls.length = 0; + 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]); @@ -269,10 +300,12 @@ test("a recorded failure fails over to the next configured model in the same tie const fake = createFakePi(); autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, b] }); - const ctx = fakeCtx({ modelRegistry: registry }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); + await selectAuto(fake, ctx); + + await fake.fire("before_agent_start", { prompt: "anything" }, ctx); expect(fake.setModelCalls).toEqual([a]); // model-a gets rate limited. @@ -291,10 +324,10 @@ test("manually picking a real model while Auto is active turns Auto off", async const fake = createFakePi(); autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, manual] }); - const ctx = fakeCtx({ modelRegistry: registry }); + const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); - await fake.fire("model_select", { model: { provider: "auto", id: "auto" }, previousModel: undefined, source: "set" }, ctx); + await selectAuto(fake, ctx); await fake.fire("model_select", { model: manual, previousModel: a, source: "set" }, ctx); fake.setModelCalls.length = 0; @@ -302,6 +335,23 @@ test("manually picking a real model while Auto is active turns Auto off", async 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(); + 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(); autoRouter(fake.pi); @@ -311,22 +361,44 @@ test("/usage with no configured models notifies instead of throwing", async () = expect(ctx.notifications.some((n) => n.message.includes("no configured models"))).toBe(true); }); -test("session_start restores an active Auto session without re-routing", async () => { +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(); autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [already] }); + fake.currentModel.value = AUTO_PLACEHOLDER; const ctx = fakeCtx({ modelRegistry: registry, - model: already, - thinkingLevel: "medium", + 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(ctx.statuses.get("auto-router")).toContain("already-selected"); + expect(lastFooterBadge(fake.footerEvents)).toBe("๐Ÿ”€ Auto"); +}); + +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(); + 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)"); }); From 51afc9a0486be9dee068102302817e319ad08974 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 18:10:24 -0400 Subject: [PATCH 06/21] Fix two real gaps that let an exhausted model look healthy Both found from a live bug report: Codex quota was hit mid-session, the turn errored out to the user, Auto never failed over, and /usage still showed the model as healthy with zero recorded requests. 1. message_end unconditionally recorded every assistant message as a success, including ones with stopReason "error" - a provider failure surfaced this way (HTTP 200, error only appears once Pi finalizes the message, so after_provider_response never sees a non-2xx status) was silently counted as a healthy request. Now checked via stopReason/errorMessage, with aborted (user-cancelled) turns correctly excluded from health tracking. 2. The Codex quota fetcher missed the account-wide authoritative signal entirely. Verified directly against the real exhausted account: `rate_limit.limit_reached`/`allowed` at the top level was true while a *healthy* per-model entry sat right next to it under additional_rate_limits for the model actually in use - only the account-wide flag reflects what's really blocking requests. Also added the missing `used_percent` field fallback the per-window percentage check was missing (present in the original pi-quotas reference this was adapted from, dropped in transcription). Verified both against the real account this was reported from, then used the fixed code directly to refresh the live persisted health state so /usage reflects it immediately rather than waiting for the next session. Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router-quota.ts | 20 +++++++- extensions/auto-router.ts | 25 +++++++++- tests/auto-router-extension.test.ts | 74 +++++++++++++++++++++++++++++ tests/auto-router-quota.test.ts | 41 +++++++++++++++- 4 files changed, 157 insertions(+), 3 deletions(-) diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index cc37cc4..20da662 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -153,13 +153,31 @@ async function fetchCodexQuota( const rateLimit = result.data.rate_limit ?? result.data.rate_limits; if (!isRecord(rateLimit)) return { exhausted: false }; + + // Codex reports this account-wide, authoritatively, right on the rate_limit object itself - + // check it before falling back to inferring exhaustion from individual 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 the specific model in use - the + // account-wide flag is the one that actually blocks every model under this provider.) + if (rateLimit.limit_reached === true || rateLimit.allowed === false) { + const window = rateLimit.primary_window ?? rateLimit.primary; + return { + exhausted: true, + resetsAt: isRecord(window) ? parseDateish(window.reset_at ?? window.reset_time_ms) : undefined, + }; + } + for (const window of [ rateLimit.primary_window ?? rateLimit.primary ?? rateLimit.five_hour_limit ?? rateLimit.five_hour, rateLimit.secondary_window ?? rateLimit.secondary ?? rateLimit.weekly_limit ?? rateLimit.weekly, ]) { if (!isRecord(window)) continue; + // The API has been observed reporting this three different ways: "percent left" fields + // (convert to used%) or a direct "used%" field. Check all three rather than assuming one. const percentLeft = numeric(window.percent_left) ?? numeric(window.remaining_percent); - if (percentLeft !== undefined && percentLeft <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + const usedPercent = percentLeft !== undefined ? 100 - percentLeft : numeric(window.used_percent); + if (usedPercent !== undefined && usedPercent >= EXHAUSTED_UTILIZATION_PERCENT) { return { exhausted: true, resetsAt: parseDateish(window.reset_at ?? window.reset_time_ms), diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 5cc774d..1b3bff1 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -46,6 +46,19 @@ 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, @@ -342,7 +355,17 @@ export default function autoRouter(pi: ExtensionAPI): void { pi.on("message_end", (event) => { if (!autoActive || !currentInFlightModel) return; if (event.message.role !== "assistant") return; - const usage = isRecord(event.message) ? event.message.usage : undefined; + const message = event.message; + if (message.stopReason === "aborted") return; // user-cancelled, not a provider health signal + if (message.stopReason === "error") { + healthStore.recordFailure( + modelKey(currentInFlightModel), + inferFailureStatus(message.errorMessage), + undefined, + ); + return; + } + const usage = isRecord(message) ? message.usage : undefined; healthStore.recordSuccess(modelKey(currentInFlightModel), { input: numeric(isRecord(usage) ? usage.input : undefined), output: numeric(isRecord(usage) ? usage.output : undefined), diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 77250e0..626c89e 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -316,6 +316,80 @@ test("a recorded failure fails over to the next configured model in the same tie 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(); + 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(); + 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("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"); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index dd67b9c..b25253b 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -89,7 +89,32 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion when the spend cap expect(result).toEqual({ exhausted: true }); }); -test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit window is depleted", async () => { +test("reconcileProviderQuota(openai-codex) reports exhaustion from the account-wide rate_limit.limit_reached flag", async () => { + // Shape verified against a real exhausted account: the account-wide flag was true while a + // per-model entry under additional_rate_limits for the model in active use was still healthy - + // the account-wide flag is what actually blocks every model under this provider. + const deps = fakeDeps(() => + jsonResponse({ + rate_limit: { + allowed: false, + limit_reached: true, + primary_window: { used_percent: 100, 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 } }, + }, + ], + spend_control: { reached: false }, + }), + ); + const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); + expect(result).toEqual({ exhausted: true, resetsAt: 1787197007 * 1000 }); +}); + +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" } } }), ); @@ -97,6 +122,20 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit expect(result).toEqual({ exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z") }); }); +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({ exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z") }); +}); + +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({ exhausted: false }); +}); + test("reconcileProviderQuota(zai) reports exhaustion from a TOKENS_LIMIT entry", async () => { const deps = fakeDeps(() => jsonResponse({ data: { limits: [{ type: "TOKENS_LIMIT", percentage: 100, nextResetTime: 4_102_444_800_000 }] } }), From 3ebe60a759427ff7466c63be03a7df6db7e8eccc Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 19:51:30 -0400 Subject: [PATCH 07/21] Show real per-model usage in /usage, and fix cooldown formatting Three more issues from the same live report: 1. /usage showed zero usage for a model the user actually had real usage on. Router-observed req/token/cost counters only ever see traffic Auto itself routed - they were the only usage figure shown, so anything used another way (manually, before installing Auto, etc.) looked unused even when verified quota data said otherwise. Quota reconciliation now carries a real "detail" string alongside its exhausted/resetsAt fields (e.g. "5% used", "interval 81% left, weekly 89% left") and /usage shows it as its own "verified usage" column/line, clearly separate from the router-observed counters. 2. "cooldown ~4557m" is unreadable. Cooldown duration now scales to minutes/hours/days instead of always minutes. 3. Codex reports quota two ways: an account-wide limit that actually gates every request, and per-model usage under additional_rate_limits for models recently used. Reconciliation results are now provider-wide with optional per-model overrides (matched by normalizing the provider's own model label), so Codex can report the real per-model percentage as detail while still using the account-wide flag - the one that actually blocks requests - to decide whether that model is marked unhealthy. Verified end-to-end against the real account this was reported against, then used the fixed code to refresh the live persisted health state directly. Co-Authored-By: Claude Sonnet 5 --- README.md | 6 +- extensions/auto-router-health.ts | 8 ++ extensions/auto-router-quota.ts | 188 ++++++++++++++++++++----------- extensions/auto-router.ts | 40 ++++--- tests/auto-router-quota.test.ts | 68 ++++++++--- 5 files changed, 219 insertions(+), 91 deletions(-) diff --git a/README.md b/README.md index 6820fcf..1a7cfdb 100644 --- a/README.md +++ b/README.md @@ -64,9 +64,11 @@ Each tier key is a Pi thinking level (`off`, `minimal`, `low`, `medium`, `high`, 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. -Health is tracked from the router's own observed traffic (HTTP status codes, rate-limit headers) and, best-effort, reconciled against real provider usage at session start and on `/usage` for providers with a known quota source (currently Anthropic, OpenAI Codex, Z.ai, and Kimi Coding via their HTTP APIs, plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no documented HTTP quota endpoint) โ€” this lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. Providers without a known quota source simply stay on router-observed data. +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 Auto itself sees a model fail. Separately, best-effort real quota reconciliation runs at session start and on `/usage`, for providers with a known quota source: Anthropic, OpenAI Codex, and Z.ai and Kimi Coding via their HTTP APIs, plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no documented HTTP quota endpoint. This is what lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures โ€” and, for Codex specifically, it's per-model where Codex reports it that way (its own account-wide limit still governs whether a model is actually blocked, since that's what's really stopping requests, but the displayed usage is the model's own). Providers without a known quota source simply stay on router-observed data. -Run `/usage` to see health/usage for every configured model, grouped by tier, in a bordered dashboard in the TUI or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. +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 (e.g. "5% used" or "interval 81% left, weekly 89% left"), and separately the request/token/cost totals *this router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). + +The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" 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 Auto โ€” behind a manual Tab to "all". At session start, Auto best-effort appends its own `auto/auto` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so it shows up in the default scoped view too, without changing anything else about what's scoped. diff --git a/extensions/auto-router-health.ts b/extensions/auto-router-health.ts index 523d40a..eb3c88d 100644 --- a/extensions/auto-router-health.ts +++ b/extensions/auto-router-health.ts @@ -28,6 +28,8 @@ export type ModelHealthEntry = { 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; @@ -38,6 +40,8 @@ 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 { @@ -160,12 +164,14 @@ export function applyQuotaResult( ? 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 }; } @@ -197,6 +203,8 @@ function parseState(value: unknown): AutoRouterHealthState { }, verifiedAt: typeof raw.verifiedAt === "number" ? raw.verifiedAt : undefined, + verifiedDetail: + typeof raw.verifiedDetail === "string" ? raw.verifiedDetail : undefined, }; } return state; diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 20da662..3b8469c 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -13,6 +13,18 @@ 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; @@ -25,6 +37,11 @@ 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"); @@ -80,6 +97,10 @@ 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; @@ -101,7 +122,7 @@ function isDirectAnthropicApiKey(token: string): boolean { async function fetchAnthropicQuota( modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, -): Promise { +): Promise { const token = await modelRegistry.getApiKeyForProvider("anthropic"); if (!token || isDirectAnthropicApiKey(token)) return undefined; const result = await fetchJson( @@ -115,21 +136,39 @@ async function fetchAnthropicQuota( ); if (!result.ok || !isRecord(result.data)) return undefined; - for (const key of ["five_hour", "seven_day"] as const) { + 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 && utilization >= EXHAUSTED_UTILIZATION_PERCENT) { - return { exhausted: true, resetsAt: parseDateish(window.resets_at) }; - } + 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 { exhausted: false }; + 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); } async function fetchCodexQuota( modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, -): Promise { +): Promise { const token = await modelRegistry.getApiKeyForProvider("openai-codex"); const accountId = await deps.readCodexAccountId(); if (!token || !accountId) return undefined; @@ -148,49 +187,63 @@ async function fetchCodexQuota( const spendControl = result.data.spend_control; if (isRecord(spendControl) && spendControl.reached === true) { - return { exhausted: true }; + return { default: { exhausted: true, detail: "spend cap reached" } }; } const rateLimit = result.data.rate_limit ?? result.data.rate_limits; - if (!isRecord(rateLimit)) return { exhausted: false }; + if (!isRecord(rateLimit)) return { default: { exhausted: false } }; - // Codex reports this account-wide, authoritatively, right on the rate_limit object itself - - // check it before falling back to inferring exhaustion from individual window percentages. + const accountWindow = isRecord(rateLimit.primary_window ?? rateLimit.primary) + ? ((rateLimit.primary_window ?? rateLimit.primary) as Record) + : undefined; + const accountUsedPercent = accountWindow ? windowUsedPercent(accountWindow) : undefined; + const accountDetail = accountUsedPercent !== undefined ? `account ${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 the specific model in use - the // account-wide flag is the one that actually blocks every model under this provider.) - if (rateLimit.limit_reached === true || rateLimit.allowed === false) { - const window = rateLimit.primary_window ?? rateLimit.primary; - return { - exhausted: true, - resetsAt: isRecord(window) ? parseDateish(window.reset_at ?? window.reset_time_ms) : undefined, + 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 }; + + // Per-model detail (and, when the account itself isn't blocking, per-model exhaustion) from + // additional_rate_limits, matched to configured model ids by normalized label. + 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 = + accountExhausted || + entryRateLimit.limit_reached === true || + entryRateLimit.allowed === false || + (usedPercent !== undefined && usedPercent >= EXHAUSTED_UTILIZATION_PERCENT); + perModel[normalizeModelId(entry.limit_name)] = { + exhausted: modelExhausted, + resetsAt: modelExhausted ? (accountResetsAt ?? (window ? windowResetsAt(window) : undefined)) : undefined, + detail: usedPercent !== undefined ? `${roundPercent(usedPercent)}% used` : accountDetail, }; } - for (const window of [ - rateLimit.primary_window ?? rateLimit.primary ?? rateLimit.five_hour_limit ?? rateLimit.five_hour, - rateLimit.secondary_window ?? rateLimit.secondary ?? rateLimit.weekly_limit ?? rateLimit.weekly, - ]) { - if (!isRecord(window)) continue; - // The API has been observed reporting this three different ways: "percent left" fields - // (convert to used%) or a direct "used%" field. Check all three rather than assuming one. - const percentLeft = numeric(window.percent_left) ?? numeric(window.remaining_percent); - const usedPercent = percentLeft !== undefined ? 100 - percentLeft : numeric(window.used_percent); - if (usedPercent !== undefined && usedPercent >= EXHAUSTED_UTILIZATION_PERCENT) { - return { - exhausted: true, - resetsAt: parseDateish(window.reset_at ?? window.reset_time_ms), - }; - } - } - return { exhausted: false }; + return { default: defaultResult, perModel: Object.keys(perModel).length > 0 ? perModel : undefined }; } async function fetchZaiQuota( modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, -): Promise { +): Promise { const apiKey = await modelRegistry.getApiKeyForProvider("zai"); if (!apiKey) return undefined; const result = await fetchJson( @@ -202,20 +255,26 @@ async function fetchZaiQuota( 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) || entry.type !== "TOKENS_LIMIT") continue; const percentage = numeric(entry.percentage); - if (percentage !== undefined && percentage >= EXHAUSTED_UTILIZATION_PERCENT) { - return { exhausted: true, resetsAt: parseDateish(entry.nextResetTime) }; - } + if (percentage === undefined) continue; + const label = entry.unit === 3 ? "hourly" : entry.unit === 6 ? "weekly" : "token"; + 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 { exhausted: false }; + return { default: { exhausted: false, detail } }; } async function fetchKimiCodingQuota( modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, -): Promise { +): Promise { const apiKey = await modelRegistry.getApiKeyForProvider("kimi-coding"); if (!apiKey) return undefined; const result = await fetchJson( @@ -226,14 +285,14 @@ async function fetchKimiCodingQuota( if (!result.ok || !isRecord(result.data)) return undefined; const weekly = result.data.usage; - if (isRecord(weekly)) { - const limit = numeric(weekly.limit); - const used = numeric(weekly.used); - if (limit !== undefined && used !== undefined && limit > 0 && used >= limit) { - return { exhausted: true, resetsAt: parseDateish(weekly.resetTime) }; - } + 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 { exhausted: false }; + return { default: { exhausted: false, detail } }; } /** @@ -247,7 +306,7 @@ async function fetchKimiCodingQuota( async function fetchMinimaxQuota( _modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, -): Promise { +): Promise { const stdout = await deps.runMinimaxCli(["quota", "show", "--output", "json"]); if (!stdout) return undefined; let data: unknown; @@ -259,21 +318,24 @@ async function fetchMinimaxQuota( if (!isRecord(data) || !Array.isArray(data.model_remains)) return undefined; const general = data.model_remains.find((entry) => isRecord(entry) && entry.model_name === "general"); - const entries = isRecord(general) ? [general] : data.model_remains; - for (const entry of entries) { - if (!isRecord(entry)) continue; - // The API reports *remaining* percent (opposite convention from the other providers' - // *used* percent), and tracks a short rolling interval plus a weekly window separately. - const intervalRemaining = numeric(entry.current_interval_remaining_percent); - if (intervalRemaining !== undefined && intervalRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { - return { exhausted: true, resetsAt: numeric(entry.end_time) }; - } - const weeklyRemaining = numeric(entry.current_weekly_remaining_percent); - if (weeklyRemaining !== undefined && weeklyRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { - return { exhausted: true, resetsAt: numeric(entry.weekly_end_time) }; - } + if (!isRecord(general)) return { default: { exhausted: false } }; + + // The API reports *remaining* percent (opposite convention from the other providers' *used* + // percent), and tracks a short rolling interval plus a weekly window separately. + const intervalRemaining = numeric(general.current_interval_remaining_percent); + const weeklyRemaining = numeric(general.current_weekly_remaining_percent); + const detailParts: string[] = []; + if (intervalRemaining !== undefined) detailParts.push(`interval ${roundPercent(intervalRemaining)}% left`); + if (weeklyRemaining !== undefined) detailParts.push(`weekly ${roundPercent(weeklyRemaining)}% left`); + const detail = detailParts.length > 0 ? detailParts.join(", ") : undefined; + + if (intervalRemaining !== undefined && intervalRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: numeric(general.end_time), detail } }; + } + if (weeklyRemaining !== undefined && weeklyRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { + return { default: { exhausted: true, resetsAt: numeric(general.weekly_end_time), detail } }; } - return { exhausted: false }; + return { default: { exhausted: false, detail } }; } /** @@ -286,7 +348,7 @@ export const QUOTA_FETCHERS: Record< ( modelRegistry: ModelRegistry, deps: QuotaFetchDependencies, - ) => Promise + ) => Promise > = { anthropic: fetchAnthropicQuota, "openai-codex": fetchCodexQuota, @@ -300,7 +362,7 @@ export async function reconcileProviderQuota( provider: string, modelRegistry: ModelRegistry, deps: QuotaFetchDependencies = defaultQuotaFetchDependencies, -): Promise { +): Promise { const fetcher = QUOTA_FETCHERS[provider]; if (!fetcher) return undefined; try { diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 1b3bff1..ac7b97f 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -16,7 +16,7 @@ import { type ModelIdentity, modelKey, } from "./auto-router-health.js"; -import { reconcileProviderQuota } from "./auto-router-quota.js"; +import { normalizeModelId, reconcileProviderQuota } from "./auto-router-quota.js"; import { AUTO_ROUTER_EFFORT_ORDER, type AutoRouterEffortLevel, @@ -285,7 +285,8 @@ export default function autoRouter(pi: ExtensionAPI): void { if (!result) return; for (const model of models) { if (model.provider !== provider) continue; - healthStore.applyQuotaResult(modelKey(model), result); + const specific = result.perModel?.[normalizeModelId(model.id)]; + healthStore.applyQuotaResult(modelKey(model), specific ?? result.default); } }), ); @@ -431,19 +432,30 @@ function formatTokenCount(count: number): string { 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 minutes = Math.max( - 1, - Math.round((entry.cooldownUntil - now) / 60_000), - ); const cause = entry.lastError ? ` (${entry.lastError.status})` : ""; - return `cooldown${cause} ~${minutes}m`; + 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 Auto itself routed. */ +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); @@ -451,8 +463,7 @@ function rowLine(row: UsageRow, now: number): string { const tokens = formatTokenCount( (entry?.totals.input ?? 0) + (entry?.totals.output ?? 0), ); - const verified = entry?.verifiedAt ? "โœ“" : "~"; - return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${requests} req ยท ${tokens} tok ${verified}`; + return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${verifiedUsageText(entry)} ยท ${requests} routed req ยท ${tokens} tok`; } function formatUsagePlainText(rows: UsageRow[]): string { @@ -477,8 +488,8 @@ function formatUsageMarkdown(rows: UsageRow[]): string { if (rows.length === 0) return "No models configured."; const now = Date.now(); const lines = [ - "| Tier | Model | Status | Req | Tokens | Cost |", - "|---|---|---|---|---|---|", + "| Tier | Model | Status | Verified usage | Routed req | Routed tokens | Routed cost |", + "|---|---|---|---|---|---|---|", ]; for (const row of rows) { const entry = row.entry; @@ -486,11 +497,14 @@ function formatUsageMarkdown(rows: UsageRow[]): string { const tokens = formatTokenCount( (entry?.totals.input ?? 0) + (entry?.totals.output ?? 0), ); - const verified = entry?.verifiedAt ? " โœ“" : ""; lines.push( - `| ${row.tier} | ${row.model.provider}/${row.model.id} | ${rowStatus(entry, now)}${verified} | ${entry?.totals.requests ?? 0} | ${tokens} | ${cost} |`, + `| ${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. Routed req/tokens/cost only count turns Auto itself routed to that model โ€” usage from other sessions, manual `/model` picks, or other tools isn't reflected there._", + ); return lines.join("\n"); } diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index b25253b..a5fda4f 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -68,13 +68,15 @@ test("reconcileProviderQuota(anthropic) reports exhaustion from an OAuth token's jsonResponse({ five_hour: { utilization: 100, resets_at: resetsAt } }), ); const result = await reconcileProviderQuota("anthropic", fakeRegistry("oauth-token"), deps); - expect(result).toEqual({ exhausted: true, resetsAt: Date.parse(resetsAt) }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse(resetsAt), detail: "5h 100% used" }, + }); }); -test("reconcileProviderQuota(anthropic) reports headroom when utilization is low", async () => { +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({ exhausted: false }); + expect(result).toEqual({ default: { exhausted: false, detail: "7d 20% used" } }); }); test("reconcileProviderQuota(openai-codex) returns undefined without a discoverable account id", async () => { @@ -86,7 +88,7 @@ test("reconcileProviderQuota(openai-codex) returns undefined without a discovera 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({ exhausted: true }); + expect(result).toEqual({ default: { exhausted: true, detail: "spend cap reached" } }); }); test("reconcileProviderQuota(openai-codex) reports exhaustion from the account-wide rate_limit.limit_reached flag", async () => { @@ -111,7 +113,12 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion from the account-w }), ); const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); - expect(result).toEqual({ exhausted: true, resetsAt: 1787197007 * 1000 }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "account 100% used" }, + // The per-model entry is *also* exhausted, since the account-wide flag blocks it too - + // but its detail reflects the model's own (much lower) usage, not the account's. + perModel: { gpt53codexspark: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "5% used" } }, + }); }); test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit window is depleted (percent_left)", async () => { @@ -119,7 +126,9 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit 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({ exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z") }); + 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 () => { @@ -127,13 +136,36 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion when a rate-limit 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({ exhausted: true, resetsAt: Date.parse("2030-06-01T00:00:00Z") }); + 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({ exhausted: false }); + 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 () => { @@ -141,13 +173,17 @@ test("reconcileProviderQuota(zai) reports exhaustion from a TOKENS_LIMIT entry", jsonResponse({ data: { limits: [{ type: "TOKENS_LIMIT", percentage: 100, nextResetTime: 4_102_444_800_000 }] } }), ); const result = await reconcileProviderQuota("zai", fakeRegistry("key"), deps); - expect(result).toEqual({ exhausted: true, resetsAt: 4_102_444_800_000 }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "token 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({ exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z") }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: Date.parse("2030-01-01T00:00:00Z"), detail: "100/100 this week" }, + }); }); test("reconcileProviderQuota(minimax) reports headroom from the mmx CLI's general bucket", async () => { @@ -161,7 +197,9 @@ test("reconcileProviderQuota(minimax) reports headroom from the mmx CLI's genera }), }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); - expect(result).toEqual({ exhausted: false }); + expect(result).toEqual({ + default: { exhausted: false, detail: "interval 84% left, weekly 89% left" }, + }); }); test("reconcileProviderQuota(minimax) reports exhaustion when the interval bucket is depleted", async () => { @@ -179,7 +217,9 @@ test("reconcileProviderQuota(minimax) reports exhaustion when the interval bucke }), }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); - expect(result).toEqual({ exhausted: true, resetsAt: 4_102_444_800_000 }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 0% left, weekly 50% left" }, + }); }); test("reconcileProviderQuota(minimax) reports exhaustion when only the weekly bucket is depleted", async () => { @@ -197,7 +237,9 @@ test("reconcileProviderQuota(minimax) reports exhaustion when only the weekly bu }), }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); - expect(result).toEqual({ exhausted: true, resetsAt: 4_102_444_800_000 }); + expect(result).toEqual({ + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 60% left, weekly 0.2% left" }, + }); }); test("reconcileProviderQuota(minimax) degrades gracefully when the CLI is missing or not logged in", async () => { From 14b005bb4651b0a9ae0c43b2afed40e16f5bd532 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 20:22:00 -0400 Subject: [PATCH 08/21] Add real OpenCode Go quota, fix Codex per-model independence Add opencode-go quota reconciliation: found a real, undocumented but clean JSON endpoint (GET /zen/go/v1/usage) that authenticates with the same API key Pi already uses for inference - no separate cookie/workspace-id setup needed, unlike pi-quotas' HTML-scraping approach for this provider. Verified directly against the real account. Also fixes a real regression the Codex per-model change just introduced: additional_rate_limits entries were being marked exhausted whenever the account-wide flag was set, even though the account-wide flag and a model's own entry are independent quota tracks - confirmed directly by the user, who could still use a model this was wrongly cooling down. A model with its own entry is now governed solely by that entry; only models without one fall back to the account-wide state. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 +- extensions/auto-router-quota.ts | 62 +++++++++++++++++++++++++++---- tests/auto-router-quota.test.ts | 66 +++++++++++++++++++++++++++++---- 3 files changed, 115 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 1a7cfdb..0950c40 100644 --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ Each tier key is a Pi thinking level (`off`, `minimal`, `low`, `medium`, `high`, 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. -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 Auto itself sees a model fail. Separately, best-effort real quota reconciliation runs at session start and on `/usage`, for providers with a known quota source: Anthropic, OpenAI Codex, and Z.ai and Kimi Coding via their HTTP APIs, plus Minimax via its `mmx` CLI (`mmx auth login`) since MiniMax has no documented HTTP quota endpoint. This is what lets the router self-correct for usage consumed outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures โ€” and, for Codex specifically, it's per-model where Codex reports it that way (its own account-wide limit still governs whether a model is actually blocked, since that's what's really stopping requests, but the displayed usage is the model's own). Providers without a known quota source simply stay on router-observed data. +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 Auto itself sees a model fail. 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 outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. 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 (e.g. "5% used" or "interval 81% left, weekly 89% left"), and separately the request/token/cost totals *this router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 3b8469c..f3e681f 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -204,8 +204,11 @@ async function fetchCodexQuota( // 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 the specific model in use - the - // account-wide flag is the one that actually blocks every model under this provider.) + // 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 || @@ -215,8 +218,9 @@ async function fetchCodexQuota( ? { exhausted: true, resetsAt: accountResetsAt, detail: accountDetail } : { exhausted: false, detail: accountDetail }; - // Per-model detail (and, when the account itself isn't blocking, per-model exhaustion) from - // additional_rate_limits, matched to configured model ids by normalized label. + // 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) { @@ -226,14 +230,13 @@ async function fetchCodexQuota( const window = isRecord(entryRateLimit.primary_window) ? entryRateLimit.primary_window : undefined; const usedPercent = window ? windowUsedPercent(window) : undefined; const modelExhausted = - accountExhausted || entryRateLimit.limit_reached === true || entryRateLimit.allowed === false || (usedPercent !== undefined && usedPercent >= EXHAUSTED_UTILIZATION_PERCENT); perModel[normalizeModelId(entry.limit_name)] = { exhausted: modelExhausted, - resetsAt: modelExhausted ? (accountResetsAt ?? (window ? windowResetsAt(window) : undefined)) : undefined, - detail: usedPercent !== undefined ? `${roundPercent(usedPercent)}% used` : accountDetail, + resetsAt: modelExhausted && window ? windowResetsAt(window) : undefined, + detail: usedPercent !== undefined ? `${roundPercent(usedPercent)}% used` : undefined, }; } @@ -295,6 +298,50 @@ async function fetchKimiCodingQuota( 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; + let mostUsed: { label: string; percent: number; window: Record } | undefined; + let blocked: { 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 = { window }; + } + const percent = numeric(window.percent); + if (percent === undefined) continue; + if (!mostUsed || percent > mostUsed.percent) mostUsed = { label, percent, window }; + } + if (!mostUsed && !blocked) return { default: { exhausted: false } }; + const detail = mostUsed ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` : 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 @@ -355,6 +402,7 @@ export const QUOTA_FETCHERS: Record< 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. */ diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index a5fda4f..c975e54 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -91,10 +91,13 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion when the spend cap expect(result).toEqual({ default: { exhausted: true, detail: "spend cap reached" } }); }); -test("reconcileProviderQuota(openai-codex) reports exhaustion from the account-wide rate_limit.limit_reached flag", async () => { - // Shape verified against a real exhausted account: the account-wide flag was true while a - // per-model entry under additional_rate_limits for the model in active use was still healthy - - // the account-wide flag is what actually blocks every model under this provider. +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: { @@ -114,10 +117,10 @@ test("reconcileProviderQuota(openai-codex) reports exhaustion from the account-w ); const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); expect(result).toEqual({ + // Applies to any configured model with no specific additional_rate_limits entry. default: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "account 100% used" }, - // The per-model entry is *also* exhausted, since the account-wide flag blocks it too - - // but its detail reflects the model's own (much lower) usage, not the account's. - perModel: { gpt53codexspark: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "5% used" } }, + // This model has its own entry, so it's unaffected by the account-wide flag. + perModel: { gpt53codexspark: { exhausted: false, detail: "5% used" } }, }); }); @@ -254,6 +257,55 @@ test("reconcileProviderQuota(minimax) degrades gracefully on unparseable CLI out expect(result).toBeUndefined(); }); +test("reconcileProviderQuota(opencode-go) reports headroom with the most-used window as detail", async () => { + const deps = fakeDeps(() => + jsonResponse({ + usage: { + rolling: { status: "ok", percent: 0, 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: "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) 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); From a4d5a035504aa59f3adf261b43cfc5c2899be444 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 21:05:18 -0400 Subject: [PATCH 09/21] Normalize quota detail to "% used" everywhere, fix Z.ai/GLM parsing Minimax's API reports *remaining* percent, the only one of the five providers that does - its detail strings read "X% left" while every other provider reads "X% used", which is exactly the kind of inconsistency that adds cognitive overhead when scanning /usage. Now converted to "used" like the rest. Also fixes Z.ai/GLM returning no usage at all: the real response for this account uses `type: "CREDIT_LIMIT"` entries (verified directly), not the `"TOKENS_LIMIT"` type the fetcher only checked for. Different plan tiers apparently report different type strings, but both carry the same `percentage` field, so the fetcher now keys off that directly instead of an incomplete type allowlist - and labels each window from its actual unit+count (e.g. "5h", "7d") instead of a fixed "token" label that ignored the count entirely. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 +- extensions/auto-router-quota.ts | 33 +++++++++++++++++++++++++++------ tests/auto-router-quota.test.ts | 31 ++++++++++++++++++++++++++----- 3 files changed, 54 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 0950c40..8103305 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ On every turn, Auto asks the `medium` tier's first healthy model (the "default m 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 Auto itself sees a model fail. 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 outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. 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 (e.g. "5% used" or "interval 81% left, weekly 89% left"), and separately the request/token/cost totals *this router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). +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 (e.g. "5% used" or "interval 16% used, weekly 11% used") โ€” and separately the request/token/cost totals *this router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index f3e681f..3aa27e3 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -243,6 +243,23 @@ async function fetchCodexQuota( 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, @@ -256,14 +273,18 @@ async function fetchZaiQuota( ); 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) || entry.type !== "TOKENS_LIMIT") continue; + if (!isRecord(entry)) continue; const percentage = numeric(entry.percentage); if (percentage === undefined) continue; - const label = entry.unit === 3 ? "hourly" : entry.unit === 6 ? "weekly" : "token"; + const label = zaiWindowLabel(entry.unit, entry.number); if (!mostUsed || percentage > mostUsed.percent) mostUsed = { label, percent: percentage, entry }; } if (!mostUsed) return { default: { exhausted: false } }; @@ -367,13 +388,13 @@ async function fetchMinimaxQuota( 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 the other providers' *used* - // percent), and tracks a short rolling interval plus a weekly window separately. + // The API reports *remaining* percent (opposite convention from every other provider here, + // which all report *used* percent) - convert so /usage reads consistently across providers. const intervalRemaining = numeric(general.current_interval_remaining_percent); const weeklyRemaining = numeric(general.current_weekly_remaining_percent); const detailParts: string[] = []; - if (intervalRemaining !== undefined) detailParts.push(`interval ${roundPercent(intervalRemaining)}% left`); - if (weeklyRemaining !== undefined) detailParts.push(`weekly ${roundPercent(weeklyRemaining)}% left`); + if (intervalRemaining !== undefined) detailParts.push(`interval ${roundPercent(100 - intervalRemaining)}% used`); + if (weeklyRemaining !== undefined) detailParts.push(`weekly ${roundPercent(100 - weeklyRemaining)}% used`); const detail = detailParts.length > 0 ? detailParts.join(", ") : undefined; if (intervalRemaining !== undefined && intervalRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index c975e54..419c6e0 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -173,14 +173,35 @@ test("reconcileProviderQuota(openai-codex) can report a model exhausted independ test("reconcileProviderQuota(zai) reports exhaustion from a TOKENS_LIMIT entry", async () => { const deps = fakeDeps(() => - jsonResponse({ data: { limits: [{ type: "TOKENS_LIMIT", percentage: 100, nextResetTime: 4_102_444_800_000 }] } }), + 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: "token 100% used" }, + 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); @@ -201,7 +222,7 @@ test("reconcileProviderQuota(minimax) reports headroom from the mmx CLI's genera }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); expect(result).toEqual({ - default: { exhausted: false, detail: "interval 84% left, weekly 89% left" }, + default: { exhausted: false, detail: "interval 16% used, weekly 11% used" }, }); }); @@ -221,7 +242,7 @@ test("reconcileProviderQuota(minimax) reports exhaustion when the interval bucke }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); expect(result).toEqual({ - default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 0% left, weekly 50% left" }, + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 100% used, weekly 50% used" }, }); }); @@ -241,7 +262,7 @@ test("reconcileProviderQuota(minimax) reports exhaustion when only the weekly bu }); const result = await reconcileProviderQuota("minimax", fakeRegistry(undefined), deps); expect(result).toEqual({ - default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 60% left, weekly 0.2% left" }, + default: { exhausted: true, resetsAt: 4_102_444_800_000, detail: "interval 40% used, weekly 99.8% used" }, }); }); From 738394611a19d4382671def5fa4d77dcb205062c Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 21:15:26 -0400 Subject: [PATCH 10/21] Label quota windows by their real duration, not placeholder words "account 100% used" and "interval 16% used" were both vague - the user correctly guessed both should read as real time windows (weekly and 5h respectively). Both APIs already carry the data needed to derive that properly instead of guessing or hardcoding it: - Codex's rate_limit windows carry their own limit_window_seconds (verified: 604800 = 7 days on the real account) - now labeled from that directly ("7d 100% used"), on both the account-wide and per-model detail strings. - Minimax's mmx CLI output carries start_time/end_time for the short window; its actual length (verified live: exactly 5 hours) is now computed from that instead of the placeholder word "interval". Falls back to "interval" only when those timestamps aren't present. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 +- extensions/auto-router-quota.ts | 39 +++++++++++++++++++++++++++++--- tests/auto-router-quota.test.ts | 40 +++++++++++++++++++++++++++------ 3 files changed, 70 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 8103305..5104978 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ On every turn, Auto asks the `medium` tier's first healthy model (the "default m 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 Auto itself sees a model fail. 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 outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. 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 (e.g. "5% used" or "interval 16% used, weekly 11% used") โ€” and separately the request/token/cost totals *this router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). +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 router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 3aa27e3..c54dba5 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -165,6 +165,27 @@ 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, @@ -197,7 +218,11 @@ async function fetchCodexQuota( ? ((rateLimit.primary_window ?? rateLimit.primary) as Record) : undefined; const accountUsedPercent = accountWindow ? windowUsedPercent(accountWindow) : undefined; - const accountDetail = accountUsedPercent !== undefined ? `account ${roundPercent(accountUsedPercent)}% used` : 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 @@ -233,10 +258,14 @@ async function fetchCodexQuota( 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 ? `${roundPercent(usedPercent)}% used` : undefined, + detail: + usedPercent !== undefined + ? `${modelLabel ? `${modelLabel} ` : ""}${roundPercent(usedPercent)}% used` + : undefined, }; } @@ -390,10 +419,14 @@ async function fetchMinimaxQuota( // 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(general.start_time, general.end_time); + const intervalLabel = intervalSeconds !== undefined ? secondsToLabel(intervalSeconds) : "interval"; const detailParts: string[] = []; - if (intervalRemaining !== undefined) detailParts.push(`interval ${roundPercent(100 - intervalRemaining)}% used`); + 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; diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index 419c6e0..04f5078 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -103,13 +103,17 @@ test("reconcileProviderQuota(openai-codex) applies the account-wide rate_limit.l rate_limit: { allowed: false, limit_reached: true, - primary_window: { used_percent: 100, reset_at: 1787197007 }, + 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 } }, + rate_limit: { + allowed: true, + limit_reached: false, + primary_window: { used_percent: 5, limit_window_seconds: 604_800 }, + }, }, ], spend_control: { reached: false }, @@ -117,10 +121,11 @@ test("reconcileProviderQuota(openai-codex) applies the account-wide rate_limit.l ); const result = await reconcileProviderQuota("openai-codex", fakeRegistry("token"), deps); expect(result).toEqual({ - // Applies to any configured model with no specific additional_rate_limits entry. - default: { exhausted: true, resetsAt: 1787197007 * 1000, detail: "account 100% used" }, + // 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: "5% used" } }, + perModel: { gpt53codexspark: { exhausted: false, detail: "7d 5% used" } }, }); }); @@ -210,17 +215,38 @@ test("reconcileProviderQuota(kimi-coding) reports exhaustion when used reaches t }); }); -test("reconcileProviderQuota(minimax) reports headroom from the mmx CLI's general bucket", async () => { +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", current_interval_remaining_percent: 84, current_weekly_remaining_percent: 89 }, + { + 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) 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" }, }); From da6ec2c7e1d71aa04e2f7a4efc096453540a28cd Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 21:39:38 -0400 Subject: [PATCH 11/21] Address PR review findings, track manually-selected models too CodeRabbit review fixes (verified each against current code first - one "Addressed" auto-label turned out to still be a live bug): - Auto could get stuck routing real requests at the inert "auto" placeholder (http://127.0.0.1:0) when nothing is configured, or becomes unconfigured mid-session, producing a bare connection error instead of a clear message. pickForTier now falls back to any authenticated model in the whole catalog as a last resort rather than ever leaving a turn pointed at the placeholder. - session_start replaced the healthStore instance instead of reusing it; a stale instance's pending debounced-save timer could still fire independently afterward and overwrite the freshly-reloaded state on disk. Now a single instance is reused and reloaded. - OpenCode Go's blocked-window detection used whichever window had the highest percentage for the reported detail, but the real reset time from whichever window was actually flagged blocked - a low blocked window next to a high healthy one produced a mismatched, misleading report. (This is the one CodeRabbit had auto-marked "Addressed" from an unrelated commit; still reproduces in current code, now actually fixed, with a regression test covering the mismatch scenario specifically.) - MiniMax reset timestamps went through a raw numeric() read instead of parseDateish's seconds-vs-milliseconds handling, unlike every other provider's reset time here. - peerDependenciesMeta "optional" findings for pi-ai/pi-coding-agent/ pi-tui: declined with an explanation rather than changed - this matches this repo's own pre-existing convention (session-footer.ts already imports runtime values from pi-tui the same way), since Pi itself provides these at runtime when loading extensions. - Two other findings were genuinely already fixed by earlier commits (Codex per-model resetsAt/account mixing, README wording); verified against current code and left as-is. Also: router-observed health/usage tracking was scoped to traffic Auto itself routed, so a model picked manually from /model - even one configured in autoRouter - showed zero usage in /usage regardless of real activity. message_end/after_provider_response now key off whichever model is actually active (ctx.model) and configured in autoRouter, not an internal "did Auto pick this" flag, so any turn against a configured model is tracked the same way. Simplified out the now-redundant currentInFlightModel tracking this replaces. /usage wording updated from "routed" to "observed" to match. Co-Authored-By: Claude Sonnet 5 --- README.md | 4 +- extensions/auto-router-quota.ts | 26 +++++++--- extensions/auto-router.ts | 73 ++++++++++++++++++----------- tests/auto-router-extension.test.ts | 68 +++++++++++++++++++++++++++ tests/auto-router-quota.test.ts | 18 +++++++ 5 files changed, 154 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 5104978..d88b4c3 100644 --- a/README.md +++ b/README.md @@ -64,9 +64,9 @@ Each tier key is a Pi thinking level (`off`, `minimal`, `low`, `medium`, `high`, 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. -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 Auto itself sees a model fail. 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 outside the current session (a different session, a manual `/model` pick, another tool) instead of only reacting to its own failures. 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. +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 router itself* has routed to that model this way โ€” the latter only reflects Auto's own traffic and will read zero for a model you've used through other means, which is exactly what verified usage is for. Shown as a bordered dashboard in the TUI, or a compact summary elsewhere (including Pi Web). +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). The `/model` picker's effort/thinking control is inert while Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index c54dba5..eec626c 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -370,22 +370,32 @@ async function fetchOpenCodeGoQuota( const usage = result.data.usage; let mostUsed: { label: string; percent: number; window: Record } | undefined; - let blocked: { window: Record } | undefined; + let blocked: { label: string; percent: number | undefined; window: Record } | undefined; for (const label of ["rolling", "weekly", "monthly"] as const) { const window = usage[label]; if (!isRecord(window)) continue; + const percent = numeric(window.percent); if (!blocked && typeof window.status === "string" && window.status !== "ok") { - blocked = { window }; + blocked = { label, percent, window }; } - const percent = numeric(window.percent); if (percent === undefined) continue; if (!mostUsed || percent > mostUsed.percent) mostUsed = { label, percent, window }; } if (!mostUsed && !blocked) return { default: { exhausted: false } }; - const detail = mostUsed ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` : undefined; if (blocked) { + // Report the blocked window's own detail, not whichever window happens to have the + // highest percentage - a blocked window can have a low percentage (e.g. a short rolling + // window resets rarely but hit its cap) while a healthy window has a higher one, and + // showing the wrong window's numbers next to the real reset time is actively misleading. + const detail = + blocked.percent !== undefined + ? `${blocked.label} ${roundPercent(blocked.percent)}% used` + : mostUsed + ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` + : undefined; return { default: { exhausted: true, resetsAt: parseDateish(blocked.window.resetsAt), detail } }; } + const detail = mostUsed ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` : undefined; if (mostUsed && mostUsed.percent >= EXHAUSTED_UTILIZATION_PERCENT) { return { default: { exhausted: true, resetsAt: parseDateish(mostUsed.window.resetsAt), detail } }; } @@ -430,11 +440,15 @@ async function fetchMinimaxQuota( 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: numeric(general.end_time), detail } }; + return { default: { exhausted: true, resetsAt: parseDateish(general.end_time), detail } }; } if (weeklyRemaining !== undefined && weeklyRemaining <= 100 - EXHAUSTED_UTILIZATION_PERCENT) { - return { default: { exhausted: true, resetsAt: numeric(general.weekly_end_time), detail } }; + return { default: { exhausted: true, resetsAt: parseDateish(general.weekly_end_time), detail } }; } return { default: { exhausted: false, detail } }; } diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index ac7b97f..a11cbe6 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -119,9 +119,8 @@ export default function autoRouter(pi: ExtensionAPI): void { let currentSessionId: string | undefined; let autoActive = false; let routingInFlight = false; - let currentInFlightModel: ModelIdentity | undefined; let lastKnownTier: AutoRouterEffortLevel | undefined; - let healthStore = new AutoRouterHealthStore(); + const healthStore = new AutoRouterHealthStore(); function publishFooter(tier: AutoRouterEffortLevel | undefined): void { if (!currentSessionId) return; @@ -194,6 +193,22 @@ export default function autoRouter(pi: ExtensionAPI): void { } return { model: fallback, tier }; } + // 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 }; + } return undefined; } @@ -219,7 +234,6 @@ export default function autoRouter(pi: ExtensionAPI): void { } finally { routingInFlight = false; } - currentInFlightModel = { provider: model.provider, id: model.id }; lastKnownTier = tier; publishFooter(tier); } @@ -303,7 +317,6 @@ export default function autoRouter(pi: ExtensionAPI): void { if (event.source !== "restore" && autoActive) { autoActive = false; pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: false }); - currentInFlightModel = undefined; lastKnownTier = undefined; clearFooter(); } @@ -312,16 +325,16 @@ export default function autoRouter(pi: ExtensionAPI): void { pi.on("session_start", async (_event, ctx) => { currentSessionId = ctx.sessionManager.getSessionId(); routingInFlight = false; - currentInFlightModel = undefined; lastKnownTier = undefined; - healthStore = new AutoRouterHealthStore(); + // 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(); autoActive = restoreAutoActive(ctx); 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. - currentInFlightModel = { provider: ctx.model.provider, id: ctx.model.id }; lastKnownTier = ctx.thinkingLevel; await revertToAutoPlaceholder(pi, ctx); } @@ -343,31 +356,38 @@ export default function autoRouter(pi: ExtensionAPI): void { await routeForPrompt(pi, ctx, event.prompt, Boolean(event.images?.length)); }); - pi.on("after_provider_response", (event) => { - if (!currentInFlightModel) return; - if (event.status >= 200 && event.status < 300) return; - healthStore.recordFailure( - modelKey(currentInFlightModel), - event.status, - event.headers, + // 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", (event) => { - if (!autoActive || !currentInFlightModel) return; + 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(currentInFlightModel), - inferFailureStatus(message.errorMessage), - undefined, - ); + healthStore.recordFailure(modelKey(model), inferFailureStatus(message.errorMessage), undefined); return; } const usage = isRecord(message) ? message.usage : undefined; - healthStore.recordSuccess(modelKey(currentInFlightModel), { + healthStore.recordSuccess(modelKey(model), { input: numeric(isRecord(usage) ? usage.input : undefined), output: numeric(isRecord(usage) ? usage.output : undefined), cost: numeric( @@ -380,7 +400,6 @@ export default function autoRouter(pi: ExtensionAPI): void { void healthStore.flush(); currentSessionId = undefined; autoActive = false; - currentInFlightModel = undefined; }); pi.registerCommand("usage", { @@ -450,7 +469,7 @@ function rowStatus(entry: ModelHealthEntry | undefined, now: number): string { 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 Auto itself routed. */ +/** 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"; @@ -463,7 +482,7 @@ function rowLine(row: UsageRow, now: number): string { const tokens = formatTokenCount( (entry?.totals.input ?? 0) + (entry?.totals.output ?? 0), ); - return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${verifiedUsageText(entry)} ยท ${requests} routed req ยท ${tokens} tok`; + return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${verifiedUsageText(entry)} ยท ${requests} req ยท ${tokens} tok`; } function formatUsagePlainText(rows: UsageRow[]): string { @@ -488,7 +507,7 @@ function formatUsageMarkdown(rows: UsageRow[]): string { if (rows.length === 0) return "No models configured."; const now = Date.now(); const lines = [ - "| Tier | Model | Status | Verified usage | Routed req | Routed tokens | Routed cost |", + "| Tier | Model | Status | Verified usage | Observed req | Observed tokens | Observed cost |", "|---|---|---|---|---|---|---|", ]; for (const row of rows) { @@ -503,7 +522,7 @@ function formatUsageMarkdown(rows: UsageRow[]): string { } lines.push( "", - "_Verified usage comes from the provider's own quota API, where available. Routed req/tokens/cost only count turns Auto itself routed to that model โ€” usage from other sessions, manual `/model` picks, or other tools isn't reflected there._", + "_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._", ); return lines.join("\n"); } diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 626c89e..ea82169 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -390,6 +390,74 @@ test("an aborted (user-cancelled) message does not count as a provider failure", 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(); + 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(); + 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(); + 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"); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index 04f5078..8f725e2 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -342,6 +342,24 @@ test("reconcileProviderQuota(opencode-go) reports exhaustion from a non-ok statu }); }); +test("reconcileProviderQuota(opencode-go) reports the blocked window's own detail, not a higher-percent healthy window's", 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) - the detail and + // resetsAt must describe the same (blocked) window, not whichever has the highest percentage. + 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" }, + }); +}); + test("reconcileProviderQuota(opencode-go) returns undefined without credentials, never calling fetch", async () => { let called = false; const deps = fakeDeps(() => { From ab82f2dd8f6c77cd447b96ffda9892f5209a39e1 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:22:47 -0400 Subject: [PATCH 12/21] Fix pre-existing typecheck errors blocking CI sessionId could be undefined when reused as a Map key in web/server/index.ts; web/client/app.tsx spread dnd-kit attributes after explicit role/tabIndex, silently discarding them (TS caught it as a duplicate-prop overwrite); and semantic-session.tsx passed optional cache/highlight values into APIs that require non-optional ones. Co-Authored-By: Claude Sonnet 5 --- web/client/app.tsx | 4 ++-- web/client/semantic-session.tsx | 8 ++++---- web/server/index.ts | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) 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; From 803c8fef87382462ed72323fe2d60b805d569bcb Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:35:38 -0400 Subject: [PATCH 13/21] Fix Auto never routing when it's the session's default model session_start only restored autoActive from a persisted session entry (written when Auto is explicitly picked via /model), so a brand-new session with defaultProvider/defaultModel set to "auto" in global settings had no entry to restore from - autoActive stayed false, and every turn's before_agent_start returned early, dispatching straight at the inert placeholder's dead baseUrl and failing with a bare "Connection error" x N. Now also treats ctx.model already being the auto placeholder at session start as active, which covers this case without disturbing the existing interrupted-mid-turn restore path. Also fixes a CodeRabbit finding: Minimax's start_time/end_time now go through parseDateish before durationSeconds, so the derived interval window label stays correct if the CLI ever reports epoch seconds instead of milliseconds (verified milliseconds only for this account so far, but nothing guarantees that universally). Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router-quota.ts | 5 ++++- extensions/auto-router.ts | 7 ++++++- tests/auto-router-extension.test.ts | 20 ++++++++++++++++++++ tests/auto-router-quota.test.ts | 21 +++++++++++++++++++++ 4 files changed, 51 insertions(+), 2 deletions(-) diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index eec626c..8a237fa 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -433,7 +433,10 @@ async function fetchMinimaxQuota( // 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(general.start_time, general.end_time); + 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`); diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index a11cbe6..c0465eb 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -330,7 +330,12 @@ export default function autoRouter(pi: ExtensionAPI): void { // 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(); - autoActive = restoreAutoActive(ctx); + // A session can arrive with Auto active two different ways: a persisted entry from an + // earlier explicit `/model` pick (`restoreAutoActive`), or `ctx.model` already being the + // placeholder because the user set `defaultProvider`/`defaultModel` to "auto" globally - a + // brand-new session in that case has no entries yet, so `restoreAutoActive` alone would + // miss it and leave every turn dispatching straight at the placeholder's dead URL. + autoActive = restoreAutoActive(ctx) || ctx.model?.provider === AUTO_PROVIDER_ID; if (autoActive) { if (ctx.model && ctx.model.provider !== AUTO_PROVIDER_ID) { // Restored mid-turn (e.g. an interrupted process, before agent_settled could diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index ea82169..3813a8d 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -523,6 +523,26 @@ test("session_start on a cleanly-idle Auto session leaves the placeholder select 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(); + 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("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" }] } } }); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index 8f725e2..93a8a05 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -237,6 +237,27 @@ test("reconcileProviderQuota(minimax) labels the short window from its real star }); }); +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 () => From 6b2d212948aeefc7f2ec7eb97b31500c56cfb976 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:38:32 -0400 Subject: [PATCH 14/21] Fix Informant build job for untracked web/dist The job's command still diffed a fresh build against a `web/dist` copied out of the checkout before building - a leftover from before c5da8a7 stopped tracking web/dist in git. Every checkout now starts with no web/dist at all, so `cp -R web/dist ...` fails immediately with "No such file or directory" before the actual build even runs. Since dist is generated, not committed, the job just needs to confirm the build succeeds. Co-Authored-By: Claude Sonnet 5 --- .informant/jobs/build.toml | 4 ---- 1 file changed, 4 deletions(-) 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 } From 0dd5206708229fae5d711cc0a15dc7cfa6cecda6 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:52:51 -0400 Subject: [PATCH 15/21] Fix classification silently downgrading to medium on a non-bare-word reply VALID_LEVELS.find(word-boundary test) picked the first match in the array's own fixed order (minimal, low, medium, high, xhigh, max), not the first one the model actually said. So a reply like "high complexity, more than a medium task" - a real answer of "high" with a "medium" comparison tacked on - matched "medium" first purely because it sorts earlier in that list, silently downgrading a hard task's routing tier. The classifier prompt asks for a single bare word, but nothing enforced it, and models don't always comply. Now parses with one alternation regex, which naturally returns whichever level word the model said *first in its own reply*, and add a small maxTokens cap so a rambling reply can't run on indefinitely. The existing "high" vs "xhigh" substring safety is preserved: `\b` can't match between two word characters, so "high" still can't match inside "xhigh". Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router-classify.ts | 15 ++++++++++----- tests/auto-router-classify.test.ts | 22 ++++++++++++++++++++++ 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/extensions/auto-router-classify.ts b/extensions/auto-router-classify.ts index 84206a0..00a0c98 100644 --- a/extensions/auto-router-classify.ts +++ b/extensions/auto-router-classify.ts @@ -74,6 +74,7 @@ export async function classifyTurnComplexity( reasoningEffort: "off", cacheRetention: "none", sessionId: uuidv7(), + maxTokens: 20, }, ); const reply = response.content @@ -85,11 +86,15 @@ export async function classifyTurnComplexity( .join("") .trim() .toLowerCase(); - // Word-boundary match, not plain substring: "high" is a substring of "xhigh", - // so a naive `.includes()` would mis-parse an "xhigh" reply as "high". - const level = - VALID_LEVELS.find((candidate) => new RegExp(`\\b${candidate}\\b`).test(reply)) ?? - DEFAULT_LEVEL; + // 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), diff --git a/tests/auto-router-classify.test.ts b/tests/auto-router-classify.test.ts index 5b8bd0b..c507adc 100644 --- a/tests/auto-router-classify.test.ts +++ b/tests/auto-router-classify.test.ts @@ -45,6 +45,28 @@ test("classifyTurnComplexity parses the reply case-insensitively with surroundin 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"); From 949ebc46475eb1bdf3826131e1f9e56943c80bed Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:09:14 -0400 Subject: [PATCH 16/21] Log what the classifier actually said, so a bad routing call can be verified The classification completion was a pure throwaway: its raw reply was parsed into a level and then discarded, so a routing decision that looked wrong (e.g. a hard task landing on medium) was impossible to actually verify afterward - only reasonable to guess about from reading the code. That's exactly the gap that made yesterday's "stayed on medium" report unresolvable from real evidence: the session transcript showed a clean route to a medium-tier model, but there was no way to tell whether the classifier had genuinely said medium or said something else that got misparsed. classifyTurnComplexity now returns its raw reply (or a reason string on failure/timeout) alongside the parsed level. auto-router-health.ts persists the last 20 routing decisions - prompt, raw reply, parsed level, resolved tier, and picked model - in the existing state file (now `{models, classifications}`, with back-compat parsing for the old flat format). /usage shows the last 5 under "Recent classifications" in both the TUI dashboard and the plain-text fallback, so a future misrouting report can be checked against what the classifier actually said instead of reasoned about from the code. Co-Authored-By: Claude Sonnet 5 --- README.md | 2 + extensions/auto-router-classify.ts | 14 +++- extensions/auto-router-health.ts | 99 ++++++++++++++++++++++++++++- extensions/auto-router.ts | 76 ++++++++++++++++++++-- tests/auto-router-classify.test.ts | 14 +++- tests/auto-router-extension.test.ts | 29 +++++++++ tests/auto-router-health.test.ts | 85 ++++++++++++++++++++++++- 7 files changed, 307 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index d88b4c3..915a7fe 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,8 @@ Health is tracked from two sources. Router-observed traffic (HTTP status codes, 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 Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" 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 Auto โ€” behind a manual Tab to "all". At session start, Auto best-effort appends its own `auto/auto` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so it shows up in the default scoped view too, without changing anything else about what's scoped. diff --git a/extensions/auto-router-classify.ts b/extensions/auto-router-classify.ts index 00a0c98..e8b9ef4 100644 --- a/extensions/auto-router-classify.ts +++ b/extensions/auto-router-classify.ts @@ -34,6 +34,13 @@ export type ClassificationUsage = { 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 { @@ -102,9 +109,10 @@ export async function classifyTurnComplexity( cost: numeric(response.usage.cost?.total), } : undefined; - return { level, usage }; - } catch { - return { level: DEFAULT_LEVEL }; + 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 index eb3c88d..40f053a 100644 --- a/extensions/auto-router-health.ts +++ b/extensions/auto-router-health.ts @@ -34,6 +34,31 @@ export type ModelHealthEntry = { 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. + */ +export type ClassificationLogEntry = { + timestamp: number; + prompt: string; + reply: string; + level: string; + tier: 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 = { @@ -210,18 +235,66 @@ function parseState(value: unknown): AutoRouterHealthState { 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.prompt !== "string" || 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; + entries.push({ + timestamp: raw.timestamp, + prompt: raw.prompt, + reply: raw.reply, + level: raw.level, + tier: 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 { - this.state = parseState( + const parsed = parsePersisted( JSON.parse(await readFile(statePath(), "utf8")), ); + this.state = parsed.models; + this.classifications = parsed.classifications; } catch { this.state = {}; + this.classifications = []; } } @@ -229,6 +302,10 @@ export class AutoRouterHealthStore { return this.state; } + getClassifications(): readonly ClassificationLogEntry[] { + return this.classifications; + } + getEntry(key: string): ModelHealthEntry | undefined { return this.state[key]; } @@ -268,6 +345,24 @@ export class AutoRouterHealthStore { this.scheduleSave(); } + recordClassification( + entry: Omit, + now: number = Date.now(), + ): void { + const full: ClassificationLogEntry = { + timestamp: now, + prompt: truncateForLog(entry.prompt), + reply: truncateForLog(entry.reply), + level: entry.level, + tier: entry.tier, + model: entry.model, + }; + this.classifications = [...this.classifications, full].slice( + -CLASSIFICATION_LOG_LIMIT, + ); + this.scheduleSave(); + } + private scheduleSave(): void { if (this.writeTimer) return; this.writeTimer = setTimeout(() => { @@ -287,7 +382,7 @@ export class AutoRouterHealthStore { try { await writeFile( tempPath, - `${JSON.stringify(this.state, null, 2)}\n`, + `${JSON.stringify({ models: this.state, classifications: this.classifications }, null, 2)}\n`, "utf8", ); await rename(tempPath, statePath()); diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index c0465eb..75e1ae3 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -12,6 +12,7 @@ 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, @@ -258,6 +259,7 @@ export default function autoRouter(pi: ExtensionAPI): void { : undefined; let level: AutoRouterEffortLevel = "medium"; + let classifierReply = "(no classifier available)"; if (classifierModel) { const result = await classifyTurnComplexity( ctx.modelRegistry, @@ -266,6 +268,7 @@ export default function autoRouter(pi: ExtensionAPI): void { hasImages, ); level = result.level; + classifierReply = result.reply; if (result.usage) { healthStore.recordSuccess(modelKey(classifierModel), result.usage); } @@ -282,6 +285,16 @@ export default function autoRouter(pi: ExtensionAPI): void { } 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({ + prompt, + reply: classifierReply, + level, + tier: picked.tier, + model: picked.model, + }); await applyRouting(pi, ctx, picked.model, picked.tier); } @@ -422,10 +435,11 @@ export default function autoRouter(pi: ExtensionAPI): void { } await reconcileAllProviders(ctx.modelRegistry, settings); const rows = buildUsageRows(settings, healthStore); + const classifications = healthStore.getClassifications(); if (ctx.mode === "tui") { - await showUsageDashboard(rows, ctx); + await showUsageDashboard(rows, classifications, ctx); } else if (ctx.hasUI) { - ctx.ui.notify(formatUsagePlainText(rows), "info"); + ctx.ui.notify(formatUsagePlainText(rows, classifications), "info"); } }, }); @@ -490,7 +504,28 @@ function rowLine(row: UsageRow, now: number): string { return `${row.model.provider}/${row.model.id} โ€” ${status} ยท ${verifiedUsageText(entry)} ยท ${requests} req ยท ${tokens} tok`; } -function formatUsagePlainText(rows: UsageRow[]): string { +/** 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`; +} + +function escapeTableCell(text: string): string { + return text.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(); @@ -505,10 +540,22 @@ function formatUsagePlainText(rows: UsageRow[]): string { `${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) { + lines.push( + ` ${formatRelativeTime(entry.timestamp, now)}: said "${entry.reply}" โ†’ ${entry.level}, routed to ${entry.tier} (${entry.model.provider}/${entry.model.id})`, + ); + } + } return lines.join("\n"); } -function formatUsageMarkdown(rows: UsageRow[]): string { +function formatUsageMarkdown( + rows: UsageRow[], + classifications: readonly ClassificationLogEntry[], +): string { if (rows.length === 0) return "No models configured."; const now = Date.now(); const lines = [ @@ -529,11 +576,28 @@ function formatUsageMarkdown(rows: UsageRow[]): string { "", "_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 | Model |", + "|---|---|---|---|---|", + ); + for (const entry of recent) { + lines.push( + `| ${formatRelativeTime(entry.timestamp, now)} | ${escapeTableCell(entry.reply)} | ${entry.level} | ${entry.tier} | ${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) => { @@ -545,7 +609,9 @@ async function showUsageDashboard( container.addChild( new Text(theme.fg("accent", theme.bold("Auto Router Usage")), 1, 0), ); - container.addChild(new Markdown(formatUsageMarkdown(rows), 1, 1, mdTheme)); + 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), ); diff --git a/tests/auto-router-classify.test.ts b/tests/auto-router-classify.test.ts index c507adc..6e417b8 100644 --- a/tests/auto-router-classify.test.ts +++ b/tests/auto-router-classify.test.ts @@ -72,10 +72,22 @@ test("classifyTurnComplexity falls back to medium on an unparseable reply", asyn expect(result.level).toBe("medium"); }); -test("classifyTurnComplexity falls back to medium when the provider call throws", async () => { +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 () => { diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 3813a8d..a4ffa43 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -225,6 +225,35 @@ test("before_agent_start routes to the classified tier, and the picker shows Aut expect(lastFooterBadge(fake.footerEvents)).toBe("๐Ÿ”€ Auto (high)"); }); +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(); + 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"); diff --git a/tests/auto-router-health.test.ts b/tests/auto-router-health.test.ts index 023c262..56aa9e3 100644 --- a/tests/auto-router-health.test.ts +++ b/tests/auto-router-health.test.ts @@ -1,8 +1,12 @@ -import { expect, test } from "bun:test"; +import { afterEach, 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 { applyFailure, applyQuotaResult, applySuccess, + AutoRouterHealthStore, type AutoRouterHealthState, isHealthy, modelKey, @@ -12,6 +16,23 @@ import { const NOW = 1_000_000_000_000; +const ENV_VAR = "PI_CODING_AGENT_DIR"; +let previousEnv: string | undefined; +let agentDir: string | undefined; + +beforeEach(async () => { + previousEnv = process.env[ENV_VAR]; + agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-health-")); + process.env[ENV_VAR] = agentDir; +}); + +afterEach(async () => { + if (previousEnv === undefined) delete process.env[ENV_VAR]; + else process.env[ENV_VAR] = previousEnv; + if (agentDir) await rm(agentDir, { recursive: true, force: true }); + agentDir = undefined; +}); + test("modelKey joins provider and id", () => { expect(modelKey({ provider: "openai", id: "gpt-5.3-codex" })).toBe("openai/gpt-5.3-codex"); }); @@ -139,3 +160,65 @@ test("a cooldown clears once its expiry has passed", () => { 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", async () => { + const store = new AutoRouterHealthStore(); + const model = { provider: "prov", id: "a" }; + store.recordSuccess(modelKey(model), { input: 10, output: 20, cost: 0.01 }); + store.recordClassification( + { prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", model }, + NOW, + ); + await store.flush(); + + 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, prompt: "do the thing", reply: "high complexity", level: "high", tier: "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( + { prompt: `turn ${i}`, reply: i === 24 ? longReply : `reply ${i}`, level: "medium", tier: "medium", model }, + NOW + i, + ); + } + const entries = store.getClassifications(); + expect(entries).toHaveLength(20); + expect(entries[0]?.prompt).toBe("turn 5"); + expect(entries.at(-1)?.reply.length).toBeLessThanOrEqual(201); + expect(entries.at(-1)?.reply.endsWith("โ€ฆ")).toBe(true); +}); From 4ffad9cfc3fc0baddf734ad59adce5d576443653 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:15:28 -0400 Subject: [PATCH 17/21] Make before_agent_start self-heal instead of trusting autoActive alone The previous fix (803c8fe) made session_start also detect Auto being active from ctx.model already being the placeholder, not just from a persisted session entry - but that only helps if ctx.model is already resolved to the settings default at the exact moment session_start fires. Whether that holds is an SDK timing assumption, and the connection-error report persisted after that fix landed and was pulled, meaning that assumption doesn't hold (or some other path still desyncs autoActive from reality) - the actual invariant needs to live somewhere it can't be defeated by a timing gap. before_agent_start is that place: it fires immediately before the turn actually dispatches, so ctx.model there is ground truth for what model is about to receive the request, no matter what happened earlier. It now checks that directly - if ctx.model is already the inert Auto placeholder, it self-heals (marks Auto active, persists the entry) and routes for real, regardless of what autoActive's own bookkeeping says. A request can no longer be sent against the placeholder as long as before_agent_start fires before dispatch, which is what it's documented to do. Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router.ts | 14 +++++++++++++- tests/auto-router-extension.test.ts | 28 ++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 75e1ae3..95b4895 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -370,7 +370,19 @@ export default function autoRouter(pi: ExtensionAPI): void { }); pi.on("before_agent_start", async (event, ctx) => { - if (!autoActive) return; + 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* the Auto placeholder, 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; + pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true }); + } await routeForPrompt(pi, ctx, event.prompt, Boolean(event.images?.length)); }); diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index a4ffa43..4397928 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -572,6 +572,34 @@ test("a brand-new session whose defaultModel is auto/auto routes on the first tu 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(); + 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" }] } } }); From 08cd972385966187da00b01317e497b21408ee5d Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:29:34 -0400 Subject: [PATCH 18/21] Show every OpenCode Go usage window, and fix a tier-mislabeling fallback bug OpenCode Go's quota fetcher tracked "mostUsed" (whichever of rolling/ weekly/monthly had the highest percent) but only ever put that single window in `detail`, silently dropping the others - a real rolling (short-window) usage figure could be sitting right there in the API response and never shown, exactly like a user just reported for kimi-k3 showing only "monthly". Now reports every window with a known percent together, e.g. "rolling 3% used, weekly 0% used, monthly 20% used", matching the multi-window pattern Minimax's fetcher already used. The blocked-window resetsAt/exhaustion behavior is unchanged. Separately, while checking real session logs for a reported "landed on the wrong tier" case, found (though couldn't confirm it was the actual cause, given other evidence pointed at a stale process) pickForTier's last-resort fallback could return a tier label that doesn't match where the picked model actually lives: if the resolved tier itself has no configured models (reachable if "medium" - the resolveEffortTier hardcoded floor - is left unconfigured) and the real fallback model comes from allConfiguredModels() spanning every tier, the old code still labeled it with the original (unconfigured) tier. That mislabeling would set the wrong thinking level for whatever model actually got picked. Fixed to search configured tiers in order and return the tier the picked model is actually configured under. Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router-quota.ts | 31 +++++++++++------------ extensions/auto-router.ts | 33 ++++++++++++++++-------- tests/auto-router-extension.test.ts | 39 +++++++++++++++++++++++++++++ tests/auto-router-quota.test.ts | 20 +++++++++------ 4 files changed, 89 insertions(+), 34 deletions(-) diff --git a/extensions/auto-router-quota.ts b/extensions/auto-router-quota.ts index 8a237fa..517eea3 100644 --- a/extensions/auto-router-quota.ts +++ b/extensions/auto-router-quota.ts @@ -369,33 +369,32 @@ async function fetchOpenCodeGoQuota( 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; percent: number | undefined; 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; - const percent = numeric(window.percent); if (!blocked && typeof window.status === "string" && window.status !== "ok") { - blocked = { label, percent, window }; + blocked = { label, window }; } + const percent = numeric(window.percent); if (percent === undefined) continue; - if (!mostUsed || percent > mostUsed.percent) mostUsed = { label, percent, window }; + const entry = { label, percent, window }; + windows.push(entry); + if (!mostUsed || percent > mostUsed.percent) mostUsed = entry; } - if (!mostUsed && !blocked) return { default: { exhausted: false } }; + 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) { - // Report the blocked window's own detail, not whichever window happens to have the - // highest percentage - a blocked window can have a low percentage (e.g. a short rolling - // window resets rarely but hit its cap) while a healthy window has a higher one, and - // showing the wrong window's numbers next to the real reset time is actively misleading. - const detail = - blocked.percent !== undefined - ? `${blocked.label} ${roundPercent(blocked.percent)}% used` - : mostUsed - ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` - : undefined; return { default: { exhausted: true, resetsAt: parseDateish(blocked.window.resetsAt), detail } }; } - const detail = mostUsed ? `${mostUsed.label} ${roundPercent(mostUsed.percent)}% used` : undefined; if (mostUsed && mostUsed.percent >= EXHAUSTED_UTILIZATION_PERCENT) { return { default: { exhausted: true, resetsAt: parseDateish(mostUsed.window.resetsAt), detail } }; } diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 95b4895..5e1c259 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -175,16 +175,27 @@ export default function autoRouter(pi: ExtensionAPI): void { if (model) return { model, tier: candidateTier }; } } - // Last resort: nothing healthy anywhere. Use the first resolvable model in `tier`, - // or failing that the first resolvable model anywhere, rather than blocking the turn. - const fallbackRefs = - settings.efforts[tier]?.models ?? allConfiguredModels(settings); - const fallback = - resolveAvailableModels(ctx.modelRegistry, fallbackRefs)[0] ?? - resolveAvailableModels( - ctx.modelRegistry, - allConfiguredModels(settings), - )[0]; + // 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; + 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; + break; + } + } + } if (fallback) { if (ctx.hasUI) { ctx.ui.notify( @@ -192,7 +203,7 @@ export default function autoRouter(pi: ExtensionAPI): void { "warning", ); } - return { model: fallback, tier }; + return { model: fallback, tier: 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 diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 4397928..9df4c45 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -312,6 +312,45 @@ test("routing falls back to the classified tier's model as a last resort when no 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(); + 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"); diff --git a/tests/auto-router-quota.test.ts b/tests/auto-router-quota.test.ts index 93a8a05..faf4cc7 100644 --- a/tests/auto-router-quota.test.ts +++ b/tests/auto-router-quota.test.ts @@ -325,18 +325,20 @@ test("reconcileProviderQuota(minimax) degrades gracefully on unparseable CLI out expect(result).toBeUndefined(); }); -test("reconcileProviderQuota(opencode-go) reports headroom with the most-used window as detail", async () => { +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: 0, resetsAt: "2030-01-01T00:00:00Z" }, + 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: "monthly 20% used" } }); + 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 () => { @@ -363,10 +365,10 @@ test("reconcileProviderQuota(opencode-go) reports exhaustion from a non-ok statu }); }); -test("reconcileProviderQuota(opencode-go) reports the blocked window's own detail, not a higher-percent healthy window's", async () => { +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) - the detail and - // resetsAt must describe the same (blocked) window, not whichever has the highest percentage. + // 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: { @@ -377,7 +379,11 @@ test("reconcileProviderQuota(opencode-go) reports the blocked window's own detai ); 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" }, + default: { + exhausted: true, + resetsAt: Date.parse("2030-01-01T00:00:00Z"), + detail: "rolling 10% used, weekly 90% used", + }, }); }); From 6bb1b6307a7d3f4d0e108c62bfbe8d5149449a52 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Mon, 17 Aug 2026 00:00:01 -0400 Subject: [PATCH 19/21] Let a model's thinking level be set independently of its routing tier A model's tier only ever decided two things: which classified- complexity bucket routes to it, and where it sits in the escalation order - but it was also blindly reused as the literal ThinkingLevel passed to pi.setThinkingLevel(), with no way to separate the two. A model that only performs well at its own maximum setting had no way to live under, say, "high" (participating in escalation and routing normally) while always actually running at "max". Model refs gain an optional `effort` field that overrides the dispatched thinking level; omitted, a model still just uses its own tier's name, unchanged from before. pickForTier now resolves and returns this alongside the tier itself (labeled correctly through every branch, including the last-resort fallbacks), and applyRouting dispatches at the resolved effort while the footer, health tracking, and classification log all keep using tier for what they're actually about (routing/grouping) - the classification log additionally records the applied effort for full transparency, and /usage shows both together as "high at max effort" when they differ. Co-Authored-By: Claude Sonnet 5 --- README.md | 10 ++++ extensions/auto-router-health.ts | 6 +++ extensions/auto-router-settings.ts | 16 +++++- extensions/auto-router.ts | 78 +++++++++++++++++++++-------- tests/auto-router-extension.test.ts | 31 ++++++++++++ tests/auto-router-health.test.ts | 21 ++++++-- tests/auto-router-settings.test.ts | 20 ++++++++ 7 files changed, 156 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 915a7fe..cfdfe55 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,16 @@ Configure it under a new `autoRouter` key in `~/.pi/agent/settings.json` (or `.p 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. 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. diff --git a/extensions/auto-router-health.ts b/extensions/auto-router-health.ts index 40f053a..3842649 100644 --- a/extensions/auto-router-health.ts +++ b/extensions/auto-router-health.ts @@ -49,6 +49,8 @@ export type ClassificationLogEntry = { 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; }; @@ -255,6 +257,9 @@ function parseClassifications(value: unknown): ClassificationLogEntry[] { 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 }, }); } @@ -355,6 +360,7 @@ export class AutoRouterHealthStore { reply: truncateForLog(entry.reply), level: entry.level, tier: entry.tier, + effort: entry.effort, model: entry.model, }; this.classifications = [...this.classifications, full].slice( diff --git a/extensions/auto-router-settings.ts b/extensions/auto-router-settings.ts index d4d9dc2..87e6ceb 100644 --- a/extensions/auto-router-settings.ts +++ b/extensions/auto-router-settings.ts @@ -59,6 +59,16 @@ 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 = { @@ -82,10 +92,12 @@ function isEffortLevel(value: string): value is AutoRouterEffortLevel { function parseModelRef(value: unknown): AutoRouterModelRef | undefined { if (!isRecord(value)) return undefined; - const { provider, id } = value; + const { provider, id, effort } = value; if (typeof provider !== "string" || !provider.trim()) return undefined; if (typeof id !== "string" || !id.trim()) return undefined; - return { provider: provider.trim(), id: id.trim() }; + 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. */ diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index 5e1c259..c6530d7 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -120,15 +120,17 @@ export default function autoRouter(pi: ExtensionAPI): void { let currentSessionId: string | undefined; let autoActive = false; let routingInFlight = false; - let lastKnownTier: AutoRouterEffortLevel | undefined; + /** 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(tier: AutoRouterEffortLevel | undefined): void { + 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(tier)), + identitySuffix: (theme: Theme) => theme.fg("accent", footerBadge(effort)), } satisfies FooterContribution); } @@ -156,12 +158,31 @@ export default function autoRouter(pi: ExtensionAPI): void { } 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 } | undefined { + ): + | { 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); @@ -172,7 +193,13 @@ export default function autoRouter(pi: ExtensionAPI): void { candidate.id === healthy.id && candidate.provider === healthy.provider, ); - if (model) return { model, tier: candidateTier }; + if (model) { + return { + model, + tier: candidateTier, + effort: resolveEffort(refs, model, candidateTier), + }; + } } } // Last resort: nothing healthy anywhere. Use the first resolvable model still configured @@ -184,6 +211,7 @@ export default function autoRouter(pi: ExtensionAPI): void { ? 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; @@ -192,6 +220,7 @@ export default function autoRouter(pi: ExtensionAPI): void { if (candidate) { fallback = candidate; fallbackTier = candidateTier; + fallbackRefs = refs; break; } } @@ -203,7 +232,11 @@ export default function autoRouter(pi: ExtensionAPI): void { "warning", ); } - return { model: fallback, tier: fallbackTier }; + 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 @@ -219,7 +252,7 @@ export default function autoRouter(pi: ExtensionAPI): void { "warning", ); } - return { model: anyModel, tier }; + return { model: anyModel, tier, effort: tier }; } return undefined; } @@ -228,7 +261,7 @@ export default function autoRouter(pi: ExtensionAPI): void { pi: ExtensionAPI, ctx: ExtensionContext, model: Model, - tier: AutoRouterEffortLevel, + effort: AutoRouterEffortLevel, ): Promise { routingInFlight = true; try { @@ -242,12 +275,12 @@ export default function autoRouter(pi: ExtensionAPI): void { } return; } - await pi.setThinkingLevel(tier); + await pi.setThinkingLevel(effort); } finally { routingInFlight = false; } - lastKnownTier = tier; - publishFooter(tier); + lastKnownEffort = effort; + publishFooter(effort); } async function routeForPrompt( @@ -304,9 +337,10 @@ export default function autoRouter(pi: ExtensionAPI): void { reply: classifierReply, level, tier: picked.tier, + effort: picked.effort, model: picked.model, }); - await applyRouting(pi, ctx, picked.model, picked.tier); + await applyRouting(pi, ctx, picked.model, picked.effort); } async function reconcileAllProviders( @@ -335,13 +369,13 @@ export default function autoRouter(pi: ExtensionAPI): void { if (event.model.provider === AUTO_PROVIDER_ID) { autoActive = true; pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true }); - publishFooter(lastKnownTier); + publishFooter(lastKnownEffort); return; } if (event.source !== "restore" && autoActive) { autoActive = false; pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: false }); - lastKnownTier = undefined; + lastKnownEffort = undefined; clearFooter(); } }); @@ -349,7 +383,7 @@ export default function autoRouter(pi: ExtensionAPI): void { pi.on("session_start", async (_event, ctx) => { currentSessionId = ctx.sessionManager.getSessionId(); routingInFlight = false; - lastKnownTier = undefined; + 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. @@ -364,10 +398,10 @@ export default function autoRouter(pi: ExtensionAPI): void { 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. - lastKnownTier = ctx.thinkingLevel; + lastKnownEffort = ctx.thinkingLevel; await revertToAutoPlaceholder(pi, ctx); } - publishFooter(lastKnownTier); + publishFooter(lastKnownEffort); } const settings = await readAutoRouterSettings(); void reconcileAllProviders(ctx.modelRegistry, settings); @@ -567,8 +601,10 @@ function formatUsagePlainText( 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} (${entry.model.provider}/${entry.model.id})`, + ` ${formatRelativeTime(entry.timestamp, now)}: said "${entry.reply}" โ†’ ${entry.level}, routed to ${entry.tier}${effortSuffix} (${entry.model.provider}/${entry.model.id})`, ); } } @@ -606,12 +642,12 @@ function formatUsageMarkdown( "### 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 | Model |", - "|---|---|---|---|---|", + "| 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.model.provider}/${entry.model.id} |`, + `| ${formatRelativeTime(entry.timestamp, now)} | ${escapeTableCell(entry.reply)} | ${entry.level} | ${entry.tier} | ${entry.effort} | ${entry.model.provider}/${entry.model.id} |`, ); } } diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 9df4c45..9874cc2 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -225,6 +225,37 @@ test("before_agent_start routes to the classified tier, and the picker shows Aut 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(); + 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"); diff --git a/tests/auto-router-health.test.ts b/tests/auto-router-health.test.ts index 56aa9e3..64610f1 100644 --- a/tests/auto-router-health.test.ts +++ b/tests/auto-router-health.test.ts @@ -171,7 +171,7 @@ test("AutoRouterHealthStore round-trips model health and classification log thro const model = { provider: "prov", id: "a" }; store.recordSuccess(modelKey(model), { input: 10, output: 20, cost: 0.01 }); store.recordClassification( - { prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", model }, + { prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", effort: "high", model }, NOW, ); await store.flush(); @@ -185,7 +185,15 @@ test("AutoRouterHealthStore round-trips model health and classification log thro cost: 0.01, }); expect(reloaded.getClassifications()).toEqual([ - { timestamp: NOW, prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", model }, + { + timestamp: NOW, + prompt: "do the thing", + reply: "high complexity", + level: "high", + tier: "high", + effort: "high", + model, + }, ]); }); @@ -212,7 +220,14 @@ test("AutoRouterHealthStore.recordClassification truncates long text and caps th const longReply = "x".repeat(500); for (let i = 0; i < 25; i++) { store.recordClassification( - { prompt: `turn ${i}`, reply: i === 24 ? longReply : `reply ${i}`, level: "medium", tier: "medium", model }, + { + prompt: `turn ${i}`, + reply: i === 24 ? longReply : `reply ${i}`, + level: "medium", + tier: "medium", + effort: "medium", + model, + }, NOW + i, ); } diff --git a/tests/auto-router-settings.test.ts b/tests/auto-router-settings.test.ts index fecc33b..6e7aa23 100644 --- a/tests/auto-router-settings.test.ts +++ b/tests/auto-router-settings.test.ts @@ -76,6 +76,26 @@ test("parseAutoRouterSettings drops malformed model refs and empty tiers", () => 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: {} }); From 28fc6aeb35c60c9bb724f1fd769684de27847d52 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Mon, 17 Aug 2026 00:30:42 -0400 Subject: [PATCH 20/21] Add a /model entry per configured tier: Auto (auto), Auto (max), etc. Previously the only way to skip classification for a turn was manual model selection, losing Auto's failover, escalation, and health tracking entirely. Now /model lists a separate Auto entry for the adaptive mode plus one per tier that has models configured - picking "Auto ()" pins every turn to route directly within that tier (still with normal failover/escalation/health tracking, just without asking a model to classify complexity first). Implementation: the extension factory is now async so it can read config before registering providers (a documented Pi pattern for "dynamically discovering available models"), registering one placeholder per configured tier alongside the existing adaptive one. Selection state gained a `pinnedTier` alongside `autoActive`, threaded through model_select, session_start's restore/self-heal path (both the persisted-entry and ctx.model-is-already-the-placeholder cases), before_agent_start's self-heal, and the revert-to-placeholder logic (which now reverts to whichever specific Auto entry was selected, not always the bare adaptive one). routeForPrompt skips classification entirely when pinned, using the pinned tier directly while still logging the routing decision for /usage. AUTO_MODEL_SCOPE_PATTERN also moved from the literal "auto/auto" to the glob "auto/*", since Pi matches enabledModels patterns with minimatch - one pattern now keeps every Auto entry visible under enabledModels scoping instead of just the adaptive one. Co-Authored-By: Claude Sonnet 5 --- README.md | 6 +- extensions/auto-router-settings.ts | 10 +- extensions/auto-router.ts | 174 ++++++++++++++++++++-------- tests/auto-router-extension.test.ts | 170 +++++++++++++++++++++++---- 4 files changed, 283 insertions(+), 77 deletions(-) diff --git a/README.md b/README.md index cfdfe55..fac55fd 100644 --- a/README.md +++ b/README.md @@ -74,15 +74,17 @@ 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 Auto is selected, since effort is chosen per turn internally. `/model` keeps showing "Auto" selected even after routing: the real model is only swapped in for the duration of each turn and swapped back to the inert Auto placeholder as soon as it settles, so reopening `/model` between turns still shows Auto, not whichever model last handled a turn. A `๐Ÿ”€ Auto ()` badge in the TUI footer tracks the most recently used tier regardless of which of the two is currently selected. Manually picking a different model from `/model` turns Auto off; reselecting "Auto" turns it back on. +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 Auto โ€” behind a manual Tab to "all". At session start, Auto best-effort appends its own `auto/auto` pattern to `enabledModels` (only when scoping is already configured, and only if it isn't already present) so it shows up in the default scoped view too, without changing anything else about what's scoped. +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 diff --git a/extensions/auto-router-settings.ts b/extensions/auto-router-settings.ts index 87e6ceb..1afff77 100644 --- a/extensions/auto-router-settings.ts +++ b/extensions/auto-router-settings.ts @@ -204,13 +204,17 @@ export async function writeAutoRouterSettings( await writeAutoRouterSettingsFile(settingsPath(), settings); } -/** Pattern that matches our registered virtual "Auto" model in `/model`'s scoping patterns. */ -export const AUTO_MODEL_SCOPE_PATTERN = "auto/auto"; +/** + * 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" - * entry) behind a manual Tab toggle to "all". If the user has scoping configured, make sure it + * 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. */ diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index c6530d7..cabce16 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -38,6 +38,21 @@ 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; @@ -73,7 +88,9 @@ function resolveAvailableModels( return models; } -function restoreAutoActive(ctx: ExtensionContext): boolean { +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]; @@ -82,10 +99,18 @@ function restoreAutoActive(ctx: ExtensionContext): boolean { entry.customType === AUTO_ACTIVE_ENTRY_TYPE ) { const data = entry.data; - return isRecord(data) && data.enabled === true; + 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 false; + return { active: false, pinnedTier: undefined }; } function formatTier(tier: AutoRouterEffortLevel): string { @@ -96,29 +121,44 @@ function footerBadge(tier: AutoRouterEffortLevel | undefined): string { return tier ? `๐Ÿ”€ Auto (${formatTier(tier)})` : "๐Ÿ”€ Auto"; } -export default function autoRouter(pi: ExtensionAPI): void { +/** 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", - // Never actually dispatched to: `before_agent_start` always swaps to a real - // routed model before any request would be sent here. baseUrl: "http://127.0.0.1:0", apiKey: "auto-router", api: "openai-completions", models: [ - { - id: AUTO_MODEL_ID, - name: "Auto", - reasoning: false, - input: ["text"], - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, - contextWindow: 128_000, - maxTokens: 4096, - }, + 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. */ @@ -134,12 +174,20 @@ export default function autoRouter(pi: ExtensionAPI): void { } satisfies FooterContribution); } - /** Swap `ctx.model` back to the inert "auto" placeholder once a turn is fully done, so `/model` keeps showing Auto selected (not whichever real model just handled the turn) between turns. */ + /** 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, AUTO_MODEL_ID); + const placeholder = ctx.modelRegistry.find( + AUTO_PROVIDER_ID, + currentAutoModelId(), + ); if (!placeholder) return; routingInFlight = true; try { @@ -290,35 +338,48 @@ export default function autoRouter(pi: ExtensionAPI): void { hasImages: boolean, ): Promise { const settings = await readAutoRouterSettings(); - 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 level: AutoRouterEffortLevel = "medium"; - let classifierReply = "(no classifier available)"; - if (classifierModel) { - const result = await classifyTurnComplexity( + 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, - classifierModel, - prompt, - hasImages, + settings.efforts.medium?.models ?? allConfiguredModels(settings), ); - level = result.level; - classifierReply = result.reply; - if (result.usage) { - healthStore.recordSuccess(modelKey(classifierModel), result.usage); + 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 tier = resolveEffortTier(settings, level); const picked = pickForTier(ctx, settings, tier); if (!picked) { if (ctx.hasUI) { @@ -368,12 +429,16 @@ export default function autoRouter(pi: ExtensionAPI): void { if (routingInFlight) return; if (event.model.provider === AUTO_PROVIDER_ID) { autoActive = true; - pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: true }); - publishFooter(lastKnownEffort); + 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(); @@ -389,11 +454,17 @@ export default function autoRouter(pi: ExtensionAPI): void { // 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 (`restoreAutoActive`), or `ctx.model` already being the - // placeholder because the user set `defaultProvider`/`defaultModel` to "auto" globally - a - // brand-new session in that case has no entries yet, so `restoreAutoActive` alone would - // miss it and leave every turn dispatching straight at the placeholder's dead URL. - autoActive = restoreAutoActive(ctx) || ctx.model?.provider === AUTO_PROVIDER_ID; + // 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 @@ -401,7 +472,7 @@ export default function autoRouter(pi: ExtensionAPI): void { lastKnownEffort = ctx.thinkingLevel; await revertToAutoPlaceholder(pi, ctx); } - publishFooter(lastKnownEffort); + publishFooter(pinnedTier ?? lastKnownEffort); } const settings = await readAutoRouterSettings(); void reconcileAllProviders(ctx.modelRegistry, settings); @@ -421,12 +492,13 @@ export default function autoRouter(pi: ExtensionAPI): void { // 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* the Auto placeholder, sending this turn as-is is - // guaranteed to fail with a bare connection error no matter why our own flag says + // 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; - pi.appendEntry(AUTO_ACTIVE_ENTRY_TYPE, { enabled: 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)); }); diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index 9874cc2..ecfb25a 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -127,7 +127,12 @@ function fakeModelRegistry({ models, unavailable = [], classify }: FakeRegistryO const allModels = [...models, AUTO_PLACEHOLDER]; const unavailableKeys = new Set(unavailable.map((m) => `${m.provider}/${m.id}`)); return { - find: (provider: string, id: string) => allModels.find((m) => m.provider === provider && m.id === id), + 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 }> }> }) => { @@ -172,12 +177,24 @@ 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(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a] }), currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); @@ -188,6 +205,117 @@ test("selecting Auto marks it active without eagerly routing, showing a neutral 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"); @@ -199,7 +327,7 @@ test("before_agent_start routes to the classified tier, and the picker shows Aut }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [medium, high], classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), @@ -236,7 +364,7 @@ test("a model's `effort` override sets its own thinking level, independent of th }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [high], classify: () => "high" }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -267,7 +395,7 @@ test("a routed turn's classification is logged and shows up in /usage, so a rout }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [medium, high], classify: (prompt) => (prompt.includes("refactor") ? "high" : "medium"), @@ -298,7 +426,7 @@ test("routing escalates to a higher tier when the classified tier is entirely un }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [medium, c, d], classify: () => "high" }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -326,7 +454,7 @@ test("routing falls back to the classified tier's model as a last resort when no await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "only-model" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [only] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -354,7 +482,7 @@ test("the last-resort fallback labels the model with the tier it actually belong await writeConfig({ efforts: { high: { models: [{ provider: "prov", id: "high-model" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [high], classify: () => "high" }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -397,7 +525,7 @@ test("a recorded failure fails over to the next configured model in the same tie }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, b] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -430,7 +558,7 @@ test("a message-level provider error (no distinct HTTP failure status) still fai }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, b] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -474,7 +602,7 @@ test("an aborted (user-cancelled) message does not count as a provider failure", }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, b] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -494,7 +622,7 @@ test("router-observed usage is tracked for a manually-selected configured model, await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -518,7 +646,7 @@ test("router-observed failures are tracked for a manually-selected configured mo await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -538,7 +666,7 @@ test("usage from a model that isn't configured anywhere in autoRouter is not tra await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, unrelated] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -563,7 +691,7 @@ test("manually picking a real model while Auto is active turns Auto off", async await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [a, manual] }); const ctx = fakeCtx({ modelRegistry: registry, currentModel: fake.currentModel }); @@ -582,7 +710,7 @@ test("deactivating Auto removes its footer badge", async () => { await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "model-a" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [a, manual] }), currentModel: fake.currentModel }); await fake.fire("session_start", {}, ctx); @@ -595,7 +723,7 @@ test("deactivating Auto removes its footer badge", async () => { test("/usage with no configured models notifies instead of throwing", async () => { const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const ctx = fakeCtx({ modelRegistry: fakeModelRegistry({ models: [] }) }); await fake.runCommand("usage", "", ctx); @@ -607,7 +735,7 @@ test("session_start on a cleanly-idle Auto session leaves the placeholder select await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "already-selected" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [already] }); fake.currentModel.value = AUTO_PLACEHOLDER; const ctx = fakeCtx({ @@ -627,7 +755,7 @@ test("a brand-new session whose defaultModel is auto/auto routes on the first tu await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "medium-model" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + 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 @@ -647,7 +775,7 @@ test("before_agent_start routes for real even if autoActive's own bookkeeping ne await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "medium-model" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + 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 @@ -675,7 +803,7 @@ test("session_start restored mid-turn (e.g. after a crash) reverts back to the A await writeConfig({ efforts: { medium: { models: [{ provider: "prov", id: "already-selected" }] } } }); const fake = createFakePi(); - autoRouter(fake.pi); + await autoRouter(fake.pi); const registry = fakeModelRegistry({ models: [already] }); fake.currentModel.value = already; const ctx = fakeCtx({ From 33c214bd04bdf3ee81bae713ff8df9504921a3a9 Mon Sep 17 00:00:00 2001 From: Ian Walter <122028+ianwalter@users.noreply.github.com> Date: Mon, 17 Aug 2026 00:44:10 -0400 Subject: [PATCH 21/21] Address CodeRabbit findings: drop persisted prompts, fix table-cell newlines Security/privacy (major): recordClassification was persisting each turn's raw prompt text into ~/.pi/agent/auto-router-state.json - a plaintext file - even though nothing in /usage ever displayed it. Prompts can contain source code, credentials, or personal data, so that was pure liability for zero benefit. Dropped `prompt` from ClassificationLogEntry and every place that wrote or read it; parseClassifications no longer requires or copies it back from old entries either, so a reload+resave scrubs it from disk. Correctness (minor): escapeTableCell only escaped `|`, not embedded line breaks. Every current caller already gets pre-collapsed text from truncateForLog, but that's a call-graph coincidence, not something the function itself enforced - a raw multi-line value would otherwise terminate its table row early and break the rest of the "Recent classifications" table. Now neutralizes both. Exported it and added a direct unit test, since nothing in the current call graph can actually exercise the newline path end-to-end to test it that way. Also fixed, found while updating the round-trip test for the prompt removal: AutoRouterHealthStore debounces its writes ~2s after the last record call, and neither auto-router-extension.test.ts nor auto-router-health.test.ts's afterEach accounted for that - a save scheduled by one test could fire after that test's own teardown had already restored PI_CODING_AGENT_DIR to its prior (usually unset) value, landing the write in the real global agent directory instead of the test's temp one. Confirmed this had already happened: the real ~/.pi/agent/auto-router-state.json got overwritten with test fixture data ("reply 5", provider "prov") partway through this session, which is what surfaced the bug. Fixed by never restoring the env var to anything but a temp dir for the life of the test run (only ever moving it to a new one) and deferring temp-dir cleanup to afterAll, so even a very late write can only ever land somewhere harmless. Co-Authored-By: Claude Sonnet 5 --- extensions/auto-router-health.ts | 12 +++++---- extensions/auto-router.ts | 11 ++++++--- tests/auto-router-extension.test.ts | 32 ++++++++++++++++-------- tests/auto-router-health.test.ts | 38 +++++++++++++++++------------ 4 files changed, 60 insertions(+), 33 deletions(-) diff --git a/extensions/auto-router-health.ts b/extensions/auto-router-health.ts index 3842649..5dd8d2d 100644 --- a/extensions/auto-router-health.ts +++ b/extensions/auto-router-health.ts @@ -41,11 +41,13 @@ 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. + * 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; - prompt: string; reply: string; level: string; tier: string; @@ -243,7 +245,7 @@ function parseClassifications(value: unknown): ClassificationLogEntry[] { for (const raw of value) { if (!isRecord(raw)) continue; if (typeof raw.timestamp !== "number") continue; - if (typeof raw.prompt !== "string" || typeof raw.reply !== "string") continue; + if (typeof raw.reply !== "string") continue; if (typeof raw.level !== "string" || typeof raw.tier !== "string") continue; if ( !isRecord(raw.model) || @@ -251,9 +253,10 @@ function parseClassifications(value: unknown): ClassificationLogEntry[] { 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, - prompt: raw.prompt, reply: raw.reply, level: raw.level, tier: raw.tier, @@ -356,7 +359,6 @@ export class AutoRouterHealthStore { ): void { const full: ClassificationLogEntry = { timestamp: now, - prompt: truncateForLog(entry.prompt), reply: truncateForLog(entry.reply), level: entry.level, tier: entry.tier, diff --git a/extensions/auto-router.ts b/extensions/auto-router.ts index cabce16..c1e25db 100644 --- a/extensions/auto-router.ts +++ b/extensions/auto-router.ts @@ -394,7 +394,6 @@ export default async function autoRouter(pi: ExtensionAPI): Promise { // 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({ - prompt, reply: classifierReply, level, tier: picked.tier, @@ -641,8 +640,14 @@ function formatRelativeTime(at: number, now: number): string { return delta < 60_000 ? "just now" : `${formatDuration(delta)} ago`; } -function escapeTableCell(text: string): string { - return text.replace(/\|/g, "\\|"); +/** + * 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( diff --git a/tests/auto-router-extension.test.ts b/tests/auto-router-extension.test.ts index ecfb25a..a2e5ce9 100644 --- a/tests/auto-router-extension.test.ts +++ b/tests/auto-router-extension.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, expect, test } from "bun:test"; +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"; @@ -9,24 +9,36 @@ import type { ExtensionContext, Theme, } from "@earendil-works/pi-coding-agent"; -import autoRouter from "../extensions/auto-router.ts"; +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 previousEnv: string | undefined; 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 () => { - previousEnv = process.env[ENV_VAR]; agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-agent-")); + usedDirs.push(agentDir); process.env[ENV_VAR] = agentDir; }); -afterEach(async () => { - if (previousEnv === undefined) delete process.env[ENV_VAR]; - else process.env[ENV_VAR] = previousEnv; - if (agentDir) await rm(agentDir, { recursive: true, force: true }); - agentDir = undefined; +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 { diff --git a/tests/auto-router-health.test.ts b/tests/auto-router-health.test.ts index 64610f1..1c38f42 100644 --- a/tests/auto-router-health.test.ts +++ b/tests/auto-router-health.test.ts @@ -1,5 +1,5 @@ -import { afterEach, beforeEach, expect, test } from "bun:test"; -import { mkdtemp, rm, writeFile } from "node:fs/promises"; +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 { @@ -17,20 +17,25 @@ import { const NOW = 1_000_000_000_000; const ENV_VAR = "PI_CODING_AGENT_DIR"; -let previousEnv: string | undefined; 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 () => { - previousEnv = process.env[ENV_VAR]; agentDir = await mkdtemp(join(tmpdir(), "pi-kit-auto-router-health-")); + usedDirs.push(agentDir); process.env[ENV_VAR] = agentDir; }); -afterEach(async () => { - if (previousEnv === undefined) delete process.env[ENV_VAR]; - else process.env[ENV_VAR] = previousEnv; - if (agentDir) await rm(agentDir, { recursive: true, force: true }); - agentDir = undefined; +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", () => { @@ -166,16 +171,21 @@ function statePath(): string { return join(agentDir, "auto-router-state.json"); } -test("AutoRouterHealthStore round-trips model health and classification log through flush/load", async () => { +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( - { prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", effort: "high", model }, + { 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({ @@ -187,7 +197,6 @@ test("AutoRouterHealthStore round-trips model health and classification log thro expect(reloaded.getClassifications()).toEqual([ { timestamp: NOW, - prompt: "do the thing", reply: "high complexity", level: "high", tier: "high", @@ -221,7 +230,6 @@ test("AutoRouterHealthStore.recordClassification truncates long text and caps th for (let i = 0; i < 25; i++) { store.recordClassification( { - prompt: `turn ${i}`, reply: i === 24 ? longReply : `reply ${i}`, level: "medium", tier: "medium", @@ -233,7 +241,7 @@ test("AutoRouterHealthStore.recordClassification truncates long text and caps th } const entries = store.getClassifications(); expect(entries).toHaveLength(20); - expect(entries[0]?.prompt).toBe("turn 5"); + expect(entries[0]?.reply).toBe("reply 5"); expect(entries.at(-1)?.reply.length).toBeLessThanOrEqual(201); expect(entries.at(-1)?.reply.endsWith("โ€ฆ")).toBe(true); });