@jameslovespancakes/pi-plus 1.0.12 → 1.0.14

package/src/domains/subscriptions/providers/anthropic.ts

@@ -1,11 +1,14 @@
 import { randomUUID } from "node:crypto";
 import type { AccountContext, AccountProvider, ManagedAccount, RoutingMode } from "../../../core/accounts/registry.ts";
+import { anthropicAccountIdentity } from "../../../core/anthropic/identity.ts";
 import { authorize, exchange } from "../../../core/anthropic/oauth.ts";
+import { ensureAccessToken } from "../../../core/anthropic/quota.ts";
 import {
   getRoutingMode,
   loadAccounts,
   saveAccount,
   setRoutingMode,
+  type Account,
 } from "../../../core/anthropic/store.ts";
 
 /**
@@ -40,22 +43,45 @@ async function authenticate(ctx: AccountContext): Promise<{ access: string; refr
   return { access: result.access, refresh: result.refresh, expires: result.expires };
 }
 
+/** Resolves and persists identity once; token rotation does not change it. */
+async function stableIdentity(account: Account): Promise<string | undefined> {
+  if (account.identity) return account.identity;
+  let access = account.access;
+  if (!access || (typeof account.expires === "number" && account.expires <= Date.now())) {
+    try { access = await ensureAccessToken(account); } catch { return undefined; }
+  }
+  if (!access) return undefined;
+  const identity = await anthropicAccountIdentity(access);
+  if (!identity) return undefined;
+
+  const current = loadAccounts()?.accounts.find((candidate) => candidate.id === account.id);
+  if (current && !current.identity) saveAccount({ ...current, identity });
+  return identity;
+}
+
 export const anthropicAccounts: AccountProvider = {
   id: "anthropic",
   label: "Claude",
 
   async list(): Promise<ManagedAccount[]> {
-    const storage = loadAccounts();
-    return (storage?.accounts ?? [])
-      .filter((account) => account.type === "oauth")
-      .map((account) => ({
+    const accounts = (loadAccounts()?.accounts ?? []).filter((account) => account.type === "oauth");
+    const result: ManagedAccount[] = [];
+    // Bootstrap is deliberately serial: this is a one-time migration for old
+    // rows, and concurrent profile requests can be throttled independently.
+    for (const account of accounts) {
+      result.push({
         id: account.id,
         label: account.label ?? account.id.slice(0, 8),
         enabled: account.enabled !== false,
         expiresAt: typeof account.expires === "number" ? account.expires : undefined,
-      }));
+        identity: await stableIdentity(account),
+      });
+    }
+    return result;
   },
 
+  identify: anthropicAccountIdentity,
+
   async add(ctx, label): Promise<string | undefined> {
     const proceed = await ctx.ui.confirm(
       "Add Claude account",
@@ -67,6 +93,7 @@ export const anthropicAccounts: AccountProvider = {
     if (!result) return undefined;
 
     const now = Date.now();
+    const identity = await anthropicAccountIdentity(result.access);
     saveAccount({
       id: randomUUID(),
       label,
@@ -77,6 +104,7 @@ export const anthropicAccounts: AccountProvider = {
       expires: result.expires,
       addedAt: now,
       lastRefreshedAt: now,
+      identity,
       authLineageId: randomUUID(),
     });
     return label;
@@ -94,12 +122,16 @@ export const anthropicAccounts: AccountProvider = {
     const result = await authenticate(ctx);
     if (!result) return undefined;
 
+    const identity = await anthropicAccountIdentity(result.access);
     saveAccount({
       ...account,
       access: result.access,
       refresh: result.refresh,
       expires: result.expires,
       lastRefreshedAt: Date.now(),
+      lastRefreshError: undefined,
+      identity,
+      quota: undefined,
       authLineageId: randomUUID(),
     });
     return account.label ?? account.id;

package/src/domains/subscriptions/providers/codex.ts

@@ -1,10 +1,9 @@
-import { randomUUID } from "node:crypto";
-import type { ModelAuth, OAuthAuth, OAuthCredential } from "@earendil-works/pi-ai";
+import type { OAuthCredential } from "@earendil-works/pi-ai";
 import { openaiCodexProvider } from "@earendil-works/pi-ai/providers/openai-codex";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import type { AccountContext, AccountProvider, ManagedAccount, RoutingMode } from "../../../core/accounts/registry.ts";
-import { normalizeRoutingMode, selectRoutingCandidate, type AccountQuotaState } from "../../../core/accounts/routing.ts";
-import { authorizeCodex, exchangeCodex } from "../../../core/codex/oauth.ts";
+import type { RoutingMode } from "../../../core/accounts/registry.ts";
+import { normalizeRoutingMode, type AccountQuotaState } from "../../../core/accounts/routing.ts";
+import type { PooledOAuthAccount, PooledOAuthStore } from "../../../core/accounts/oauth-pool.ts";
 import { applyCodexQuotaHeaders } from "../../../core/codex/quota.ts";
 import {
   MAIN_ACCOUNT_ID,
@@ -14,170 +13,47 @@ import {
   saveCodexAccounts,
   type CodexAccount,
 } from "../../../core/codex/store.ts";
+import {
+  chooseCredential,
+  createPooledOAuthAdapter,
+  registerPooledOAuthProvider,
+  type PooledOAuthProviderSpec,
+} from "./oauth-pool.ts";
 
 /**
  * Codex adapter.
  *
- * Structurally the same as the Anthropic one, with two differences that come
- * from the provider rather than from us:
- *   - quota arrives only through response headers, so there is no poll;
- *   - every request must carry the account's `chatgpt-account-id`, so the
- *     account id is stored alongside the tokens rather than derived later.
+ * Codex shares the pooled serving path with every other subscription provider
+ * (one `routedAuth`, one bounded refresh, one quota observer). Only the two
+ * genuinely provider-specific pieces are supplied here:
+ *
+ *   - storage stays in `codex-accounts.json`, because it carries `accountId`
+ *     (which the usage bars read) and a rich quota snapshot that the generic
+ *     pool schema has no room for;
+ *   - quota arrives through `x-codex-*` response headers rather than the
+ *     generic `x-ratelimit-*` set, so there is no polling path at all.
+ *
+ * The ChatGPT account id does NOT need to be threaded through auth: the
+ * provider derives it from the access token's JWT claims and sets the
+ * `chatgpt-account-id` header itself, so swapping the token is self-consistent.
  */
 
-async function authenticate(ctx: AccountContext) {
-  const auth = await authorizeCodex();
-
-  try {
-    await ctx.openBrowser(auth.url);
-  } catch {
-    ctx.ui.notify(`Open this authorization URL:\n${auth.url}`, "warning");
-  }
-
-  const callback = await ctx.ui.input("Codex OAuth", "Paste the callback URL or authorization code");
-  if (!callback) return undefined;
-
-  return exchangeCodex(callback, auth.verifier, auth.redirectUri, auth.state);
-}
-
 function describePlan(account: CodexAccount): string {
-  return account.plan ? `${account.label ?? account.id.slice(0, 8)} (${account.plan})` : (account.label ?? account.id.slice(0, 8));
-}
-
-type RoutableCodexAccount = CodexAccount & { access: string; refresh: string; expires: number };
-
-const refreshes = new Map<string, Promise<RoutableCodexAccount>>();
-const codexLastUsed = new Map<string, number>();
-
-function isRoutableCodexAccount(account: CodexAccount): account is RoutableCodexAccount {
-  return account.enabled !== false
-    && typeof account.access === "string"
-    && typeof account.refresh === "string"
-    && typeof account.expires === "number";
-}
-
-export const codexAccounts: AccountProvider = {
-  id: "openai-codex",
-  label: "Codex",
-
-  async list(): Promise<ManagedAccount[]> {
-    return loadCodexAccounts().accounts.map((account) => ({
-      id: account.id,
-      label: describePlan(account),
-      enabled: account.enabled !== false,
-      expiresAt: typeof account.expires === "number" ? account.expires : undefined,
-    }));
-  },
-
-  async add(ctx, label): Promise<string | undefined> {
-    const proceed = await ctx.ui.confirm(
-      "Add Codex account",
-      "Sign in with a DIFFERENT ChatGPT account in the browser. Continue?",
-    );
-    if (!proceed) return undefined;
-
-    const tokens = await authenticate(ctx);
-    if (!tokens) return undefined;
-
-    const claims = claimsOf(tokens.access);
-
-    // Adding the same ChatGPT account twice gives routing two entries backed by
-    // one quota pool, which looks like capacity that does not exist.
-    const existing = loadCodexAccounts().accounts.find(
-      (a) => claims.accountId && a.accountId === claims.accountId,
-    );
-    if (existing) {
-      throw new Error(
-        `That ChatGPT account is already added as “${existing.label ?? existing.id.slice(0, 8)}”. `
-        + "Use the browser's account switcher to sign in as a different account.",
-      );
-    }
-
-    const now = Date.now();
-    saveCodexAccount({
-      id: randomUUID(),
-      label,
-      enabled: true,
-      access: tokens.access,
-      refresh: tokens.refresh,
-      expires: tokens.expires,
-      accountId: claims.accountId,
-      plan: claims.plan,
-      addedAt: now,
-      lastRefreshedAt: now,
-    });
-    return label;
-  },
-
-  async reauth(ctx, accountId): Promise<string | undefined> {
-    const account = loadCodexAccounts().accounts.find(
-      (a) => a.id === accountId || a.label === accountId,
-    );
-    if (!account) throw new Error(`Codex account “${accountId}” not found.`);
-
-    const tokens = await authenticate(ctx);
-    if (!tokens) return undefined;
-
-    const claims = claimsOf(tokens.access);
-    saveCodexAccount({
-      ...account,
-      access: tokens.access,
-      refresh: tokens.refresh,
-      expires: tokens.expires,
-      accountId: claims.accountId ?? account.accountId,
-      plan: claims.plan ?? account.plan,
-      lastRefreshedAt: Date.now(),
-    });
-    return account.label ?? account.id;
-  },
-
-  async setEnabled(accountId: string, enabled: boolean): Promise<void> {
-    const storage = loadCodexAccounts();
-    const account = storage.accounts.find((a) => a.id === accountId);
-    if (!account) throw new Error(`Account “${accountId}” not found.`);
-    account.enabled = enabled;
-    saveCodexAccounts(storage);
-  },
-
-  async rename(accountId: string, label: string): Promise<void> {
-    const storage = loadCodexAccounts();
-    const account = storage.accounts.find((a) => a.id === accountId);
-    if (!account) throw new Error(`Account “${accountId}” not found.`);
-    account.label = label;
-    saveCodexAccounts(storage);
-  },
-
-  routing: {
-    async get(): Promise<RoutingMode> {
-      return codexRoutingMode(loadCodexAccounts().routing?.mode);
-    },
-    async set(mode: RoutingMode): Promise<RoutingMode> {
-      const storage = loadCodexAccounts();
-      storage.routing = { ...(storage.routing ?? {}), mode };
-      saveCodexAccounts(storage);
-      return mode;
-    },
-    describe(mode: RoutingMode): string {
-      return mode === "quota-aware"
-        ? "Uses the Codex subscription with the most remaining quota."
-        : "Uses Codex subscriptions in order, moving on when one is exhausted.";
-    },
-  },
-};
-
-export function codexRoutingMode(value: string | undefined): RoutingMode {
-  return normalizeRoutingMode(value);
+  const name = account.label ?? account.id.slice(0, 8);
+  return account.plan ? `${name} (${account.plan})` : name;
 }
 
+/** Routing view of a stored Codex snapshot. */
 function quotaState(quota: CodexAccount["quota"], blockedUntil?: number): AccountQuotaState | undefined {
+  const blocked = blockedUntil && blockedUntil > Date.now() ? blockedUntil : undefined;
   if (!quota) {
-    return blockedUntil && blockedUntil > Date.now()
-      ? { remainingPercent: 0, checkedAt: Date.now(), blockedUntil }
-      : undefined;
+    return blocked ? { remainingPercent: 0, checkedAt: Date.now(), blockedUntil: blocked } : undefined;
   }
   const windows = [quota.five_hour, quota.seven_day, ...(quota.scoped ?? [])]
     .filter((window) => window && Number.isFinite(window.remainingPercent));
-  if (windows.length === 0) return undefined;
+  if (windows.length === 0) {
+    return blocked ? { remainingPercent: 0, checkedAt: quota.checkedAt ?? Date.now(), blockedUntil: blocked } : undefined;
+  }
 
   const remainingPercent = Math.min(...windows.map((window) => window!.remainingPercent!));
   const resetTimes = windows
@@ -186,89 +62,73 @@ function quotaState(quota: CodexAccount["quota"], blockedUntil?: number): Accoun
     .filter(Number.isFinite);
   const resetAt = resetTimes.length > 0 ? Math.min(...resetTimes) : undefined;
   return {
-    remainingPercent: blockedUntil && blockedUntil > Date.now() ? 0 : remainingPercent,
+    remainingPercent: blocked ? 0 : remainingPercent,
     resetAt,
     checkedAt: quota.checkedAt ?? Date.now(),
-    blockedUntil: blockedUntil && blockedUntil > Date.now()
-      ? blockedUntil
-      : remainingPercent <= 0 ? resetAt ?? Number.POSITIVE_INFINITY : undefined,
+    // An exhausted window with no reset time is already non-viable, so there is
+    // no sentinel here: `Infinity` would serialize to null and silently unblock.
+    blockedUntil: blocked ?? (remainingPercent <= 0 ? resetAt : undefined),
   };
 }
 
-export function chooseCodexCredential(primary: OAuthCredential): {
-  readonly id: string;
-  readonly credential: OAuthCredential;
-  readonly account?: RoutableCodexAccount;
-} {
-  const storage = loadCodexAccounts();
-  const mode = codexRoutingMode(storage.routing?.mode);
-  const candidates = [
-    {
-      id: MAIN_ACCOUNT_ID,
-      order: 0,
-      lastUsed: codexLastUsed.get(MAIN_ACCOUNT_ID) ?? storage.main?.lastUsed ?? 0,
-      quota: quotaState(storage.main?.quota, storage.main?.blockedUntil),
-      value: { id: MAIN_ACCOUNT_ID, credential: primary },
-    },
-    ...storage.accounts
-      .filter(isRoutableCodexAccount)
-      .map((account, index) => ({
-        id: account.id,
-        order: index + 1,
-        lastUsed: codexLastUsed.get(account.id) ?? account.lastUsed ?? 0,
+/** True when the account has a complete credential and can serve a request. */
+function isRoutable(account: CodexAccount): boolean {
+  return typeof account.access === "string"
+    && typeof account.refresh === "string"
+    && typeof account.expires === "number";
+}
+
+/**
+ * Adapts `codex-accounts.json` to the pooled store contract.
+ *
+ * `quota` and `identity` are derived views, so they are never written back:
+ * the pool would otherwise overwrite the rich snapshot the usage bars read
+ * with the flattened routing state.
+ */
+const CODEX_STORE: PooledOAuthStore = {
+  load() {
+    const storage = loadCodexAccounts();
+    return {
+      mode: normalizeRoutingMode(storage.routing?.mode),
+      accounts: storage.accounts.map((account) => ({
+        ...account,
+        type: "oauth" as const,
+        // Blanking access on an incomplete credential keeps it out of routing
+        // while still listing it under `/accounts`.
+        access: isRoutable(account) ? account.access! : "",
+        refresh: account.refresh ?? "",
+        expires: account.expires ?? 0,
+        label: account.label ?? "",
+        identity: account.accountId,
+        addedAt: account.addedAt ?? 0,
         quota: quotaState(account.quota, account.blockedUntil),
-        value: {
-          id: account.id,
-          credential: { ...account, type: "oauth" as const } satisfies OAuthCredential,
-          account,
-        },
       })),
-  ];
-  return selectRoutingCandidate(candidates, mode)?.value ?? { id: MAIN_ACCOUNT_ID, credential: primary };
-}
+    };
+  },
 
-async function freshCodexAccount(
-  oauth: OAuthAuth,
-  account: RoutableCodexAccount,
-  signal: AbortSignal,
-): Promise<RoutableCodexAccount> {
-  if (account.access && account.expires && account.expires > Date.now() + 60_000) return account;
-  const active = refreshes.get(account.id);
-  if (active) return active;
+  saveAccount(account: PooledOAuthAccount) {
+    const { quota: _quota, identity: _identity, type: _type, ...rest } = account as PooledOAuthAccount & CodexAccount;
+    const claims = claimsOf(rest.access);
+    saveCodexAccount({
+      ...rest,
+      accountId: rest.accountId ?? claims.accountId,
+      plan: rest.plan ?? claims.plan,
+    });
+  },
 
-  const credential: OAuthCredential = { ...account, type: "oauth" };
-  const refresh = oauth.refresh(credential, signal).then((updated) => {
-    const next: RoutableCodexAccount = { ...account, ...updated };
-    saveCodexAccount(next);
-    return next;
-  }).finally(() => refreshes.delete(account.id));
-  refreshes.set(account.id, refresh);
-  return refresh;
-}
+  saveMode(mode) {
+    const storage = loadCodexAccounts();
+    storage.routing = { ...(storage.routing ?? {}), mode };
+    saveCodexAccounts(storage);
+  },
 
-async function routedCodexAuth(
-  oauth: OAuthAuth,
-  primary: OAuthCredential,
-): Promise<ModelAuth> {
-  const selected = chooseCodexCredential(primary);
-  const account = selected.account
-    ? await freshCodexAccount(oauth, selected.account, new AbortController().signal)
-    : undefined;
-  const credential: OAuthCredential = account ? { ...account, type: "oauth" } : primary;
-  const now = Date.now();
-  codexLastUsed.set(selected.id, now);
-  const storage = loadCodexAccounts();
-  if (selected.id === MAIN_ACCOUNT_ID) {
-    if (now - (storage.main?.lastUsed ?? 0) >= 60_000) {
-      storage.main = { ...storage.main, lastUsed: now };
-      saveCodexAccounts(storage);
-    }
-  } else if (account && now - (account.lastUsed ?? 0) >= 60_000) {
-    saveCodexAccount({ ...account, lastUsed: now });
-  }
-  return oauth.toAuth(credential);
-}
+  primaryQuota() {
+    const storage = loadCodexAccounts();
+    return quotaState(storage.main?.quota, storage.main?.blockedUntil);
+  },
+};
 
+/** Persists a rate-limit block so routing skips the account until it clears. */
 function markCodexRateLimited(accountId: string, headers: Record<string, string>): void {
   const raw = Object.entries(headers).find(([key]) => key.toLowerCase() === "retry-after")?.[1];
   const seconds = Number(raw);
@@ -283,33 +143,39 @@ function markCodexRateLimited(accountId: string, headers: Record<string, string>
   if (account) saveCodexAccount({ ...account, blockedUntil });
 }
 
-export function registerCodexProvider(pi: ExtensionAPI): void {
-  const provider = openaiCodexProvider();
-  const oauth = provider.auth.oauth;
-  if (!oauth) return;
+export const CODEX_SPEC: PooledOAuthProviderSpec<"openai-codex-responses"> = {
+  id: "openai-codex",
+  label: "Codex",
+  createProvider: openaiCodexProvider,
+  store: CODEX_STORE,
+  addPrompt: "Sign in with a DIFFERENT ChatGPT account in the browser. Continue?",
+  describeAccount: (account) => describePlan(account as PooledOAuthAccount & CodexAccount),
+  // Two entries backed by one ChatGPT account would look like capacity that
+  // does not exist, so duplicates are matched on the token's account claim.
+  identityOf: (access) => claimsOf(access).accountId,
+  recordQuota(accountId, status, headers) {
+    // Never throws: swallows its own write failures.
+    applyCodexQuotaHeaders(accountId, headers);
+    if (status !== 429) return;
+    try {
+      markCodexRateLimited(accountId, headers);
+    } catch {
+      // Telemetry only — a failed block write must not fail the response.
+    }
+  },
+};
 
-  const withQuotaObserver = (options: any) => ({
-    ...options,
-    onResponse: async (response: { status: number; headers: Record<string, string> }, model: unknown) => {
-      const storage = loadCodexAccounts();
-      const account = storage.accounts.find((candidate) => candidate.access === options?.apiKey);
-      const accountId = account?.id ?? MAIN_ACCOUNT_ID;
-      applyCodexQuotaHeaders(accountId, response.headers);
-      if (response.status === 429) markCodexRateLimited(accountId, response.headers);
-      await options?.onResponse?.(response, model);
-    },
-  });
+export const codexAccounts = createPooledOAuthAdapter(CODEX_SPEC);
 
-  pi.registerProvider({
-    ...provider,
-    auth: {
-      ...provider.auth,
-      oauth: {
-        ...oauth,
-        toAuth: (primary) => routedCodexAuth(oauth, primary),
-      },
-    },
-    stream: (model: any, context: any, options: any) => provider.stream(model, context, withQuotaObserver(options)),
-    streamSimple: (model: any, context: any, options: any) => provider.streamSimple(model, context, withQuotaObserver(options)),
-  });
+export function codexRoutingMode(value: string | undefined): RoutingMode {
+  return normalizeRoutingMode(value);
+}
+
+/** Retained for callers and tests that select a Codex account directly. */
+export function chooseCodexCredential(primary: OAuthCredential) {
+  return chooseCredential(CODEX_SPEC, primary);
+}
+
+export function registerCodexProvider(pi: ExtensionAPI): void {
+  registerPooledOAuthProvider(pi, CODEX_SPEC);
 }

package/src/domains/subscriptions/providers/hosted.ts

@@ -2,13 +2,13 @@ import { kimiCodingProvider } from "@earendil-works/pi-ai/providers/kimi-coding"
 import { xaiProvider } from "@earendil-works/pi-ai/providers/xai";
 import { createPooledOAuthAdapter, type PooledOAuthProviderSpec } from "./oauth-pool.ts";
 
-export const KIMI_SPEC: PooledOAuthProviderSpec = {
+export const KIMI_SPEC: PooledOAuthProviderSpec<"anthropic-messages"> = {
   id: "kimi-coding",
   label: "Kimi",
   createProvider: kimiCodingProvider,
 };
 
-export const XAI_SPEC: PooledOAuthProviderSpec = {
+export const XAI_SPEC: PooledOAuthProviderSpec<"openai-responses"> = {
   id: "xai",
   label: "Grok",
   createProvider: xaiProvider,