takomi 2.1.25 → 2.1.26

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,4 +1,4 @@
1
- import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
1
+ import { chmodSync, existsSync, mkdirSync, readFileSync, renameSync, rmSync, statSync, writeFileSync } from "node:fs";
2
2
  import { homedir } from "node:os";
3
3
  import { dirname, join } from "node:path";
4
4
  import type { RouterConfig, RouterModelConfig, RouterUpstreamConfig } from "./types.ts";
@@ -96,6 +96,16 @@ const DEFAULT_UPSTREAMS: RouterUpstreamConfig[] = [
96
96
  oauthProviderId: "openai-codex",
97
97
  enabled: true,
98
98
  modelIds: ["gpt-5.1", "gpt-5.4", "gpt-5.4-mini", "gpt-5.5"],
99
+ usageProbe: {
100
+ enabled: true,
101
+ endpoints: [
102
+ "/wham/usage",
103
+ "/codex/usage",
104
+ "/codex/limits",
105
+ "/codex/rate_limits",
106
+ "/codex/subscription",
107
+ ],
108
+ },
99
109
  },
100
110
  ];
101
111
 
@@ -110,23 +120,71 @@ export const DEFAULT_CONFIG: RouterConfig = {
110
120
  upstreams: DEFAULT_UPSTREAMS,
111
121
  };
112
122
 
113
- function applySecurePermissions(path: string, mode: number) {
123
+ function applySecurePermissions(filePath: string, mode: number) {
114
124
  try {
115
- chmodSync(path, mode);
125
+ chmodSync(filePath, mode);
116
126
  } catch {
117
127
  // Best effort only. Windows commonly ignores POSIX chmod semantics.
118
128
  }
119
129
  }
120
130
 
121
- export function ensureDirectory(path: string) {
122
- mkdirSync(path, { recursive: true, mode: 0o700 });
123
- applySecurePermissions(path, 0o700);
131
+ export function ensureDirectory(filePath: string) {
132
+ mkdirSync(filePath, { recursive: true, mode: 0o700 });
133
+ applySecurePermissions(filePath, 0o700);
124
134
  }
125
135
 
126
- export function writeJsonFile(path: string, value: unknown, secure = true) {
127
- ensureDirectory(dirname(path));
128
- writeFileSync(path, `${JSON.stringify(value, null, 2)}\n`, { encoding: "utf8", mode: secure ? 0o600 : 0o644 });
129
- applySecurePermissions(path, secure ? 0o600 : 0o644);
136
+ function sleepSync(ms: number): void {
137
+ Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms);
138
+ }
139
+
140
+ const HELD_JSON_LOCKS = new Set<string>();
141
+
142
+ export function withJsonFileLock<T>(filePath: string, fn: () => T): T {
143
+ const lockPath = `${filePath}.lock`;
144
+ if (HELD_JSON_LOCKS.has(lockPath)) return fn();
145
+ const started = Date.now();
146
+ while (true) {
147
+ try {
148
+ mkdirSync(lockPath, { mode: 0o700 });
149
+ break;
150
+ } catch {
151
+ try {
152
+ const ageMs = Date.now() - statSync(lockPath).mtimeMs;
153
+ if (ageMs > 30_000) rmSync(lockPath, { recursive: true, force: true });
154
+ } catch {
155
+ // The lock disappeared between mkdir attempts.
156
+ }
157
+ if (Date.now() - started > 10_000) {
158
+ throw new Error(`Timed out waiting for oauth-router JSON lock: ${lockPath}`);
159
+ }
160
+ sleepSync(50);
161
+ }
162
+ }
163
+
164
+ HELD_JSON_LOCKS.add(lockPath);
165
+ try {
166
+ return fn();
167
+ } finally {
168
+ HELD_JSON_LOCKS.delete(lockPath);
169
+ rmSync(lockPath, { recursive: true, force: true });
170
+ }
171
+ }
172
+
173
+ export function writeJsonFile(filePath: string, value: unknown, secure = true) {
174
+ ensureDirectory(dirname(filePath));
175
+ withJsonFileLock(filePath, () => {
176
+ const mode = secure ? 0o600 : 0o644;
177
+ const tempPath = `${filePath}.${process.pid}.${Date.now()}.tmp`;
178
+ try {
179
+ writeFileSync(tempPath, `${JSON.stringify(value, null, 2)}\n`, { encoding: "utf8", mode });
180
+ applySecurePermissions(tempPath, mode);
181
+ renameSync(tempPath, filePath);
182
+ applySecurePermissions(filePath, mode);
183
+ } catch (error) {
184
+ rmSync(tempPath, { force: true });
185
+ throw error;
186
+ }
187
+ });
130
188
  }
131
189
 
132
190
  function deepClone<T>(value: T): T {
@@ -160,7 +218,12 @@ function mergeUpstreamConfigs(candidateUpstreams: RouterUpstreamConfig[] | undef
160
218
  }
161
219
 
162
220
  const modelIds = Array.from(new Set([...(previous.modelIds ?? []), ...(upstream.modelIds ?? [])]));
163
- merged.set(upstream.id, { ...previous, ...deepClone(upstream), id: upstream.id, modelIds });
221
+ const usageProbe = {
222
+ ...(previous.usageProbe ?? {}),
223
+ ...(upstream.usageProbe ?? {}),
224
+ endpoints: Array.from(new Set([...(previous.usageProbe?.endpoints ?? []), ...(upstream.usageProbe?.endpoints ?? [])])),
225
+ };
226
+ merged.set(upstream.id, { ...previous, ...deepClone(upstream), id: upstream.id, modelIds, usageProbe });
164
227
  }
165
228
  return Array.from(merged.values());
166
229
  }
@@ -1,9 +1,10 @@
1
1
  import { spawn } from "node:child_process";
2
+ import { Buffer } from "node:buffer";
2
3
  import { randomUUID } from "node:crypto";
3
- import type { OAuthCredentials } from "@mariozechner/pi-ai";
4
- import { getOAuthProvider } from "@mariozechner/pi-ai/oauth";
4
+ import type { OAuthCredentials, OAuthSelectPrompt } from "@earendil-works/pi-ai/oauth";
5
+ import { getOAuthProvider } from "@earendil-works/pi-ai/oauth";
5
6
  import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
6
- import type { RouterUpstreamConfig, StoredRouterAccount } from "./types.ts";
7
+ import type { RouterProviderQuotaWindow, RouterProviderUsageSnapshot, RouterUpstreamConfig, StoredRouterAccount } from "./types.ts";
7
8
 
8
9
  function now() {
9
10
  return Date.now();
@@ -51,6 +52,17 @@ async function promptRequired(ctx: ExtensionContext, message: string, placeholde
51
52
  return response;
52
53
  }
53
54
 
55
+ async function selectOAuthOption(ctx: ExtensionContext, prompt: OAuthSelectPrompt): Promise<string | undefined> {
56
+ if (!prompt.options.length) return undefined;
57
+ if (!ctx.hasUI) return prompt.options[0]?.id;
58
+
59
+ const labels = prompt.options.map((option) => `${option.id} — ${option.label}`);
60
+ const choice = await ctx.ui.select(prompt.message, labels);
61
+ if (!choice) return undefined;
62
+ const id = choice.split(" — ")[0]?.trim();
63
+ return prompt.options.find((option) => option.id === id)?.id;
64
+ }
65
+
54
66
  export async function createAccountFromUpstream(
55
67
  upstream: RouterUpstreamConfig,
56
68
  label: string,
@@ -91,12 +103,21 @@ export async function createAccountFromUpstream(
91
103
  ctx.ui.notify(`${provider.name}: ${info.instructions ?? "Finish login in your browser."}`, "info");
92
104
  ctx.ui.notify(info.url, "info");
93
105
  },
106
+ onDeviceCode(info) {
107
+ openUrlInBrowser(info.verificationUri);
108
+ ctx.ui.notify(`${provider.name}: device code ${info.userCode}`, "info");
109
+ ctx.ui.notify(`Open ${info.verificationUri} and enter code ${info.userCode}`, "info");
110
+ },
94
111
  onPrompt(prompt) {
95
112
  return promptRequired(ctx, prompt.message, prompt.placeholder);
96
113
  },
97
114
  onProgress(message) {
98
115
  ctx.ui.notify(message, "info");
99
116
  },
117
+ onSelect(prompt) {
118
+ return selectOAuthOption(ctx, prompt);
119
+ },
120
+ signal: ctx.signal,
100
121
  });
101
122
 
102
123
  const normalized = normalizeCredentials(credentials);
@@ -126,6 +147,260 @@ function toCredentials(account: StoredRouterAccount): OAuthCredentials {
126
147
  };
127
148
  }
128
149
 
150
+ function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
151
+ const [, payload] = token.split(".");
152
+ if (!payload) return undefined;
153
+
154
+ try {
155
+ const normalized = payload.replace(/-/g, "+").replace(/_/g, "/");
156
+ const padded = normalized.padEnd(normalized.length + ((4 - (normalized.length % 4)) % 4), "=");
157
+ return JSON.parse(Buffer.from(padded, "base64").toString("utf8")) as Record<string, unknown>;
158
+ } catch {
159
+ return undefined;
160
+ }
161
+ }
162
+
163
+ function getStringClaim(claims: Record<string, unknown>, key: string): string | undefined {
164
+ const value = claims[key];
165
+ return typeof value === "string" && value ? value : undefined;
166
+ }
167
+
168
+ function getAudience(claims: Record<string, unknown>): string | string[] | undefined {
169
+ const audience = claims.aud;
170
+ if (typeof audience === "string") return audience;
171
+ if (Array.isArray(audience) && audience.every((item) => typeof item === "string")) return audience as string[];
172
+ return undefined;
173
+ }
174
+
175
+ function getOpenAIAuthClaim(claims: Record<string, unknown>): Record<string, unknown> | undefined {
176
+ const value = claims["https://api.openai.com/auth"];
177
+ return value && typeof value === "object" && !Array.isArray(value) ? (value as Record<string, unknown>) : undefined;
178
+ }
179
+
180
+ export function inspectAccountToken(account: StoredRouterAccount): RouterProviderUsageSnapshot {
181
+ const claims = decodeJwtPayload(account.access);
182
+ if (!claims) {
183
+ return {
184
+ fetchedAt: now(),
185
+ source: "unavailable",
186
+ accountId: typeof account.meta?.accountId === "string" ? account.meta.accountId : undefined,
187
+ expires: account.expires,
188
+ message: "Access token is not a readable JWT or exposes no local claims. Provider-side usage needs an authenticated usage endpoint.",
189
+ };
190
+ }
191
+
192
+ const openaiAuth = getOpenAIAuthClaim(claims);
193
+ const exp = typeof claims.exp === "number" && Number.isFinite(claims.exp) ? claims.exp * 1000 : account.expires;
194
+ const accountId =
195
+ (typeof account.meta?.accountId === "string" ? account.meta.accountId : undefined) ??
196
+ (openaiAuth && getStringClaim(openaiAuth, "chatgpt_account_id")) ??
197
+ getStringClaim(claims, "account_id");
198
+ const planType =
199
+ (typeof account.meta?.planType === "string" ? account.meta.planType : undefined) ??
200
+ (openaiAuth && (getStringClaim(openaiAuth, "plan_type") ?? getStringClaim(openaiAuth, "planType"))) ??
201
+ getStringClaim(claims, "plan_type") ??
202
+ getStringClaim(claims, "planType");
203
+
204
+ return {
205
+ fetchedAt: now(),
206
+ source: "token-claims",
207
+ accountId,
208
+ planType,
209
+ email: getStringClaim(claims, "email"),
210
+ subject: getStringClaim(claims, "sub"),
211
+ issuer: getStringClaim(claims, "iss"),
212
+ audience: getAudience(claims),
213
+ expires: exp,
214
+ claimKeys: Object.keys(claims).sort(),
215
+ message: "Token claims expose identity/expiry metadata only. 5h and weekly quota windows are not present in this token snapshot; local router-observed usage is shown separately.",
216
+ };
217
+ }
218
+
219
+ function firstNumber(...values: unknown[]): number | undefined {
220
+ for (const value of values) {
221
+ if (typeof value === "number" && Number.isFinite(value)) return value;
222
+ if (typeof value === "string" && value.trim() && Number.isFinite(Number(value))) return Number(value);
223
+ }
224
+ return undefined;
225
+ }
226
+
227
+ function firstString(...values: unknown[]): string | undefined {
228
+ for (const value of values) {
229
+ if (typeof value === "string" && value.trim()) return value.trim();
230
+ }
231
+ return undefined;
232
+ }
233
+
234
+ function isRecord(value: unknown): value is Record<string, unknown> {
235
+ return Boolean(value) && typeof value === "object" && !Array.isArray(value);
236
+ }
237
+
238
+ function walkRecords(value: unknown, visit: (record: Record<string, unknown>) => void, depth = 0) {
239
+ if (depth > 8) return;
240
+ if (Array.isArray(value)) {
241
+ for (const item of value) walkRecords(item, visit, depth + 1);
242
+ return;
243
+ }
244
+ if (!isRecord(value)) return;
245
+ visit(value);
246
+ for (const item of Object.values(value)) walkRecords(item, visit, depth + 1);
247
+ }
248
+
249
+ function normalizeResetAt(record: Record<string, unknown>): number | undefined {
250
+ const seconds = firstNumber(record.reset_after_seconds, record.resetAfterSeconds, record.resets_in_seconds, record.resetsInSeconds);
251
+ if (seconds !== undefined) return now() + seconds * 1000;
252
+ const raw = firstNumber(record.reset_at, record.resetAt, record.resets_at, record.resetsAt, record.next_reset_at, record.nextResetAt);
253
+ if (raw === undefined) return undefined;
254
+ return raw < 10_000_000_000 ? raw * 1000 : raw;
255
+ }
256
+
257
+ function quotaFromRecord(label: string, record: Record<string, unknown>): RouterProviderQuotaWindow | undefined {
258
+ const limit = firstNumber(record.limit, record.cap, record.total, record.max, record.quota);
259
+ const remaining = firstNumber(record.remaining, record.available, record.left, record.remaining_messages, record.remainingMessages);
260
+ const used = firstNumber(record.used, record.consumed, record.current, record.used_messages, record.usedMessages);
261
+ const usedPercent = firstNumber(record.used_percent, record.usedPercent);
262
+ const percentRemaining = firstNumber(record.percent_remaining, record.percentRemaining, record.remaining_percent, record.remainingPercent) ?? (usedPercent !== undefined ? 100 - Math.max(0, Math.min(100, usedPercent)) : undefined);
263
+ const resetAt = normalizeResetAt(record);
264
+
265
+ if (limit === undefined && remaining === undefined && used === undefined && percentRemaining === undefined && resetAt === undefined) {
266
+ return undefined;
267
+ }
268
+
269
+ return {
270
+ label,
271
+ used,
272
+ limit,
273
+ remaining,
274
+ percentRemaining: percentRemaining !== undefined ? percentRemaining : limit && remaining !== undefined ? (remaining / limit) * 100 : undefined,
275
+ resetAt,
276
+ };
277
+ }
278
+
279
+ function mergeQuota(previous: RouterProviderQuotaWindow | undefined, next: RouterProviderQuotaWindow | undefined): RouterProviderQuotaWindow | undefined {
280
+ if (!previous) return next;
281
+ if (!next) return previous;
282
+ return { ...previous, ...next };
283
+ }
284
+
285
+ function extractProviderUsage(json: unknown): Pick<RouterProviderUsageSnapshot, "planType" | "fiveHour" | "weekly"> {
286
+ let planType: string | undefined;
287
+ let fiveHour: RouterProviderQuotaWindow | undefined;
288
+ let weekly: RouterProviderQuotaWindow | undefined;
289
+
290
+ walkRecords(json, (record) => {
291
+ planType ??= firstString(record.plan_type, record.planType, record.plan, record.subscription_plan, record.subscriptionPlan, record.account_plan, record.accountPlan);
292
+
293
+ const rateLimit = isRecord(record.rate_limit) ? record.rate_limit : undefined;
294
+ if (rateLimit) {
295
+ const primary = isRecord(rateLimit.primary_window) ? rateLimit.primary_window : isRecord(rateLimit.primaryWindow) ? rateLimit.primaryWindow : undefined;
296
+ const secondary = isRecord(rateLimit.secondary_window) ? rateLimit.secondary_window : isRecord(rateLimit.secondaryWindow) ? rateLimit.secondaryWindow : undefined;
297
+ fiveHour = mergeQuota(fiveHour, primary ? quotaFromRecord("5h", primary) : undefined);
298
+ weekly = mergeQuota(weekly, secondary ? quotaFromRecord("weekly", secondary) : undefined);
299
+ }
300
+
301
+ const name = [record.name, record.label, record.bucket, record.window, record.period, record.type, record.key, record.id]
302
+ .map((value) => (typeof value === "string" ? value.toLowerCase() : ""))
303
+ .join(" ");
304
+
305
+ const directFive = isRecord(record.five_hour) ? record.five_hour : isRecord(record.fiveHour) ? record.fiveHour : isRecord(record["5h"]) ? record["5h"] : undefined;
306
+ const directWeekly = isRecord(record.weekly) ? record.weekly : isRecord(record.week) ? record.week : isRecord(record["7d"]) ? record["7d"] : undefined;
307
+ fiveHour = mergeQuota(fiveHour, directFive ? quotaFromRecord("5h", directFive) : undefined);
308
+ weekly = mergeQuota(weekly, directWeekly ? quotaFromRecord("weekly", directWeekly) : undefined);
309
+
310
+ if (/5\s*h|five.?hour|primary.?window/.test(name)) {
311
+ fiveHour = mergeQuota(fiveHour, quotaFromRecord("5h", record));
312
+ }
313
+ if (/week|weekly|7\s*d|secondary.?window/.test(name)) {
314
+ weekly = mergeQuota(weekly, quotaFromRecord("weekly", record));
315
+ }
316
+ });
317
+
318
+ return { planType, fiveHour, weekly };
319
+ }
320
+
321
+ function resolveProbeUrl(baseUrl: string, endpoint: string): string {
322
+ if (/^https?:\/\//i.test(endpoint)) return endpoint;
323
+ return `${baseUrl.replace(/\/+$/, "")}/${endpoint.replace(/^\/+/, "")}`;
324
+ }
325
+
326
+ function collectRateLimitHeaders(headers: Headers): Record<string, string> {
327
+ const result: Record<string, string> = {};
328
+ for (const [key, value] of headers.entries()) {
329
+ if (/rate.?limit|reset|remaining|quota/i.test(key)) result[key] = value;
330
+ }
331
+ return result;
332
+ }
333
+
334
+ export async function refreshProviderUsageSnapshot(account: StoredRouterAccount, upstream: RouterUpstreamConfig): Promise<RouterProviderUsageSnapshot> {
335
+ const base = inspectAccountToken(account);
336
+ if (account.provider !== "openai-codex" || upstream.usageProbe?.enabled === false) return base;
337
+
338
+ const accountId = base.accountId;
339
+ const endpoints = Array.from(new Set(["/wham/usage", ...(upstream.usageProbe?.endpoints ?? [])]));
340
+ if (!accountId || endpoints.length === 0) {
341
+ return { ...base, message: `${base.message ?? "Token inspected."} No provider usage probe endpoints are configured.` };
342
+ }
343
+
344
+ let lastStatus: number | undefined;
345
+ let lastEndpoint: string | undefined;
346
+ let lastHeaders: Record<string, string> | undefined;
347
+ const errors: string[] = [];
348
+
349
+ for (const endpoint of endpoints) {
350
+ const url = resolveProbeUrl(upstream.baseUrl, endpoint);
351
+ lastEndpoint = url;
352
+ try {
353
+ const response = await fetch(url, {
354
+ method: "GET",
355
+ headers: {
356
+ Authorization: `Bearer ${account.access}`,
357
+ "ChatGPT-Account-Id": accountId,
358
+ "chatgpt-account-id": accountId,
359
+ Originator: "codex_cli_rs",
360
+ originator: "codex_cli_rs",
361
+ "User-Agent": "codex_cli_rs/0.133.0 (Windows; x86_64) pi/oauth-router",
362
+ accept: "application/json",
363
+ },
364
+ });
365
+ lastStatus = response.status;
366
+ lastHeaders = collectRateLimitHeaders(response.headers);
367
+ const text = await response.text();
368
+ const json = text ? JSON.parse(text) : undefined;
369
+ const extracted = extractProviderUsage(json);
370
+ const hasQuota = Boolean(extracted.fiveHour || extracted.weekly || extracted.planType);
371
+ if (response.ok && hasQuota) {
372
+ return {
373
+ ...base,
374
+ source: extracted.fiveHour || extracted.weekly ? "provider" : "token-claims",
375
+ fetchedAt: now(),
376
+ planType: extracted.planType ?? base.planType,
377
+ fiveHour: extracted.fiveHour,
378
+ weekly: extracted.weekly,
379
+ endpoint: url,
380
+ status: response.status,
381
+ rateLimitHeaders: lastHeaders,
382
+ message: extracted.fiveHour || extracted.weekly
383
+ ? "Provider-side quota metadata was extracted from an authenticated ChatGPT/Codex endpoint."
384
+ : "Provider endpoint responded, but no 5h/weekly quota windows were found; token identity metadata is shown.",
385
+ };
386
+ }
387
+ errors.push(`${endpoint}: HTTP ${response.status}${hasQuota ? " partial metadata only" : " no quota fields"}`);
388
+ } catch (error) {
389
+ const message = error instanceof Error ? error.message : String(error);
390
+ errors.push(`${endpoint}: ${message}`);
391
+ }
392
+ }
393
+
394
+ return {
395
+ ...base,
396
+ fetchedAt: now(),
397
+ endpoint: lastEndpoint,
398
+ status: lastStatus,
399
+ rateLimitHeaders: lastHeaders,
400
+ message: `Provider quota probe did not find 5h/weekly windows. Tried ${endpoints.length} endpoint(s): ${errors.slice(0, 4).join("; ")}${errors.length > 4 ? "; …" : ""}`,
401
+ };
402
+ }
403
+
129
404
  export async function refreshAccountCredentials(account: StoredRouterAccount): Promise<StoredRouterAccount> {
130
405
  if (account.provider === "api-key") return account;
131
406
 
@@ -1,5 +1,5 @@
1
1
  import { existsSync, readFileSync } from "node:fs";
2
- import { CREDENTIALS_PATH, writeJsonFile } from "./config.ts";
2
+ import { CREDENTIALS_PATH, withJsonFileLock, writeJsonFile } from "./config.ts";
3
3
  import type { RouterCredentialStore, StoredRouterAccount } from "./types.ts";
4
4
 
5
5
  const EMPTY_STORE: RouterCredentialStore = {
@@ -73,6 +73,14 @@ export class RouterAccountStore {
73
73
  writeJsonFile(CREDENTIALS_PATH, this.data, true);
74
74
  }
75
75
 
76
+ private mutate(fn: () => void) {
77
+ withJsonFileLock(CREDENTIALS_PATH, () => {
78
+ this.data = this.load();
79
+ fn();
80
+ this.save();
81
+ });
82
+ }
83
+
76
84
  reload() {
77
85
  this.data = this.load();
78
86
  }
@@ -86,21 +94,33 @@ export class RouterAccountStore {
86
94
  }
87
95
 
88
96
  add(account: StoredRouterAccount) {
89
- this.data.accounts.push(normalizeAccount(account));
90
- this.save();
97
+ this.mutate(() => {
98
+ this.data.accounts.push(normalizeAccount(account));
99
+ });
91
100
  }
92
101
 
93
102
  update(account: StoredRouterAccount) {
94
- const index = this.data.accounts.findIndex((item) => item.id === account.id);
95
- if (index === -1) throw new Error(`Unknown account: ${account.id}`);
96
- this.data.accounts[index] = normalizeAccount(account);
97
- this.save();
103
+ this.mutate(() => {
104
+ const index = this.data.accounts.findIndex((item) => item.id === account.id);
105
+ if (index === -1) throw new Error(`Unknown account: ${account.id}`);
106
+ this.data.accounts[index] = normalizeAccount(account);
107
+ });
98
108
  }
99
109
 
100
110
  remove(id: string) {
101
- const next = this.data.accounts.filter((account) => account.id !== id);
102
- this.data.accounts = next;
103
- this.save();
111
+ this.mutate(() => {
112
+ this.data.accounts = this.data.accounts.filter((account) => account.id !== id);
113
+ });
114
+ }
115
+
116
+ rename(id: string, label: string) {
117
+ const account = this.get(id);
118
+ if (!account) throw new Error(`Unknown account: ${id}`);
119
+ const nextLabel = label.trim();
120
+ if (!nextLabel) throw new Error("Account label cannot be empty");
121
+ account.label = nextLabel;
122
+ account.updatedAt = Date.now();
123
+ this.update(account);
104
124
  }
105
125
 
106
126
  setEnabled(id: string, enabled: boolean) {
@@ -13,7 +13,7 @@ import {
13
13
  } from "@mariozechner/pi-ai";
14
14
  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
15
15
  import { loadRouterConfig } from "./config.ts";
16
- import { refreshAccountCredentials, getApiKeyForAccount } from "./oauth-flow.ts";
16
+ import { refreshAccountCredentials, getApiKeyForAccount, refreshProviderUsageSnapshot } from "./oauth-flow.ts";
17
17
  import { RouterAccountStore } from "./oauth-store.ts";
18
18
  import { chooseEligibleAccount } from "./policies.ts";
19
19
  import { RouterStateStore } from "./state.ts";
@@ -82,6 +82,12 @@ function createErrorMessage({ model, message, stopReason = "error" }: RouterErro
82
82
  };
83
83
  }
84
84
 
85
+ function isAbortLike(message?: string, signal?: AbortSignal, stopReason?: string): boolean {
86
+ if (signal?.aborted || stopReason === "aborted") return true;
87
+ const lower = (message ?? "").toLowerCase();
88
+ return lower.includes("abort") || lower.includes("cancelled") || lower.includes("canceled");
89
+ }
90
+
85
91
  function classifyFailure(
86
92
  status: number | undefined,
87
93
  headers: Record<string, string>,
@@ -182,6 +188,10 @@ export class RouterRuntime {
182
188
  this.state.pruneAccountIds(this.accounts.list().map((account) => account.id));
183
189
  }
184
190
 
191
+ renameAccount(id: string, label: string) {
192
+ this.accounts.rename(id, label);
193
+ }
194
+
185
195
  setEnabled(id: string, enabled: boolean) {
186
196
  this.accounts.setEnabled(id, enabled);
187
197
  if (enabled) this.state.clearHealth(id);
@@ -191,6 +201,38 @@ export class RouterRuntime {
191
201
  this.accounts.setWeight(id, weight);
192
202
  }
193
203
 
204
+ clearAccountHealth(id: string) {
205
+ this.state.clearHealth(id);
206
+ }
207
+
208
+ getUsageSummaries() {
209
+ this.reloadConfig();
210
+ return this.accounts.list().map((account) => this.state.getUsageSummary(account.id));
211
+ }
212
+
213
+ getUsageSummary(id: string) {
214
+ this.reloadConfig();
215
+ if (!this.accounts.get(id)) throw new Error(`Unknown account: ${id}`);
216
+ return this.state.getUsageSummary(id);
217
+ }
218
+
219
+ async refreshUsageSnapshot(id: string) {
220
+ this.reloadConfig();
221
+ const account = this.accounts.get(id);
222
+ if (!account) throw new Error(`Unknown account: ${id}`);
223
+ const upstream = this.getUpstream(account.upstreamId);
224
+ if (!upstream) throw new Error(`Unknown upstream for account ${id}: ${account.upstreamId}`);
225
+ const snapshot = await refreshProviderUsageSnapshot(account, upstream);
226
+ this.state.setProviderUsage(id, snapshot);
227
+ return snapshot;
228
+ }
229
+
230
+ resetUsage(id: string) {
231
+ this.reloadConfig();
232
+ if (!this.accounts.get(id)) throw new Error(`Unknown account: ${id}`);
233
+ this.state.resetUsage(id);
234
+ }
235
+
194
236
  setPolicy(policy: RoutingPolicyName) {
195
237
  this.state.setPolicy(policy);
196
238
  }
@@ -372,7 +414,17 @@ export class RouterRuntime {
372
414
 
373
415
  for await (const event of inner) {
374
416
  if (event.type === "error") {
375
- const failure = classifyFailure(responseStatus, responseHeaders, event.error.errorMessage || "Upstream request failed", this.config);
417
+ const message = event.error.errorMessage || "Upstream request failed";
418
+ if (isAbortLike(message, options?.signal, event.error.stopReason)) {
419
+ outerStream.push({
420
+ ...event,
421
+ reason: "aborted",
422
+ error: { ...event.error, stopReason: "aborted", errorMessage: message },
423
+ });
424
+ return { completed: true, emittedMeaningfulOutput };
425
+ }
426
+
427
+ const failure = classifyFailure(responseStatus, responseHeaders, message, this.config);
376
428
  this.markFailure(selection.account.id, failure);
377
429
 
378
430
  if (!emittedMeaningfulOutput) {
@@ -400,6 +452,7 @@ export class RouterRuntime {
400
452
  for (const pending of buffered) outerStream.push(pending);
401
453
  buffered.length = 0;
402
454
  }
455
+ this.state.recordUsage(selection.account.id, selection.modelConfig.id, event.message.usage, responseStatus ?? 200);
403
456
  this.state.markSuccess(selection.account.id, responseStatus ?? 200);
404
457
  return { completed: true, emittedMeaningfulOutput };
405
458
  }
@@ -418,6 +471,7 @@ export class RouterRuntime {
418
471
 
419
472
  stream(model: Model<Api>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream {
420
473
  this.reloadConfig();
474
+ this.state.clearAbortHealth(this.accounts.list().map((account) => account.id));
421
475
  const outer = createAssistantMessageEventStream();
422
476
 
423
477
  (async () => {
@@ -427,6 +481,12 @@ export class RouterRuntime {
427
481
 
428
482
  try {
429
483
  while (true) {
484
+ if (options?.signal?.aborted) {
485
+ outer.push({ type: "error", reason: "aborted", error: createErrorMessage({ model, message: "Request was aborted", stopReason: "aborted" }) });
486
+ outer.end();
487
+ return;
488
+ }
489
+
430
490
  const eligible = this.getEligibleAccounts(model.id, tried);
431
491
  const policy = this.getPolicy();
432
492
  const cursor = this.state.getCursor(policy);
@@ -450,6 +510,11 @@ export class RouterRuntime {
450
510
  selection = await this.prepareSelection(picked.selected, model, options);
451
511
  } catch (error) {
452
512
  const message = error instanceof Error ? error.message : String(error);
513
+ if (isAbortLike(message, options?.signal)) {
514
+ outer.push({ type: "error", reason: "aborted", error: createErrorMessage({ model, message, stopReason: "aborted" }) });
515
+ outer.end();
516
+ return;
517
+ }
453
518
  lastFailure = { kind: "auth", status: 401, message };
454
519
  this.markFailure(picked.selected.account.id, lastFailure);
455
520
  continue;
@@ -471,7 +536,8 @@ export class RouterRuntime {
471
536
  }
472
537
  } catch (error) {
473
538
  const message = error instanceof Error ? error.message : String(error);
474
- outer.push({ type: "error", reason: "error", error: createErrorMessage({ model, message }) });
539
+ const aborted = isAbortLike(message, options?.signal);
540
+ outer.push({ type: "error", reason: aborted ? "aborted" : "error", error: createErrorMessage({ model, message, stopReason: aborted ? "aborted" : "error" }) });
475
541
  outer.end();
476
542
  }
477
543
  })();
@@ -20,10 +20,6 @@ REQUIRED = [
20
20
  ROOT / "config.ts",
21
21
  ROOT / "types.ts",
22
22
  ROOT / "README.md",
23
- ROOT / "docs" / "Project_Requirements.md",
24
- ROOT / "docs" / "Coding_Guidelines.md",
25
- ROOT / "docs" / "Builder_Prompt.md",
26
- ROOT / "docs" / "issues" / "FR-001.md",
27
23
  ]
28
24
  EXPECTED_MODELS = ["oauth-router", "gpt-4o", "gpt-4.1", "o4-mini", "gpt-5.4"]
29
25
  PI_CANDIDATES = [