@jameslovespancakes/pi-plus 1.0.11 → 1.0.13

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

@@ -1,12 +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 RoutingMode as StoreRoutingMode,
+  type Account,
 } from "../../../core/anthropic/store.ts";
 
 /**
@@ -22,16 +24,6 @@ import {
  * the extracted store directly.
  */
 
-/** The registry's vocabulary maps onto the store's routing modes. */
-const MODE_TO_STORE: Record<RoutingMode, StoreRoutingMode> = {
-  standard: "main-first",
-  optimal: "sticky-balanced",
-};
-
-function fromStoreMode(mode: string): RoutingMode {
-  return mode === "sticky-balanced" ? "optimal" : "standard";
-}
-
 async function authenticate(ctx: AccountContext): Promise<{ access: string; refresh: string; expires: number } | undefined> {
   const auth = await authorize("max");
 
@@ -51,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",
@@ -78,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,
@@ -88,6 +104,7 @@ export const anthropicAccounts: AccountProvider = {
       expires: result.expires,
       addedAt: now,
       lastRefreshedAt: now,
+      identity,
       authLineageId: randomUUID(),
     });
     return label;
@@ -105,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;
@@ -133,17 +154,16 @@ export const anthropicAccounts: AccountProvider = {
 
   routing: {
     async get(): Promise<RoutingMode> {
-      return fromStoreMode(getRoutingMode(loadAccounts()));
+      return getRoutingMode(loadAccounts());
     },
     async set(mode: RoutingMode): Promise<RoutingMode> {
-      setRoutingMode(MODE_TO_STORE[mode]);
+      setRoutingMode(mode);
       return mode;
     },
     describe(mode: RoutingMode): string {
-      return mode === "optimal"
-        ? "Balances sessions across subscriptions using remaining quota and time until reset. "
-          + "Keeps session caches sticky; migrates on confirmed exhaustion."
-        : "Main subscription first, then eligible fallbacks.";
+      return mode === "quota-aware"
+        ? "Uses the Claude subscription with the most remaining quota."
+        : "Uses Claude subscriptions in order, moving on when one is exhausted.";
     },
   },
 };

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

@@ -1,148 +1,181 @@
-import { randomUUID } from "node:crypto";
-import type { AccountContext, AccountProvider, ManagedAccount, RoutingMode } from "../../../core/accounts/registry.ts";
-import { authorizeCodex, exchangeCodex } from "../../../core/codex/oauth.ts";
+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 { 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,
   claimsOf,
   loadCodexAccounts,
   saveCodexAccount,
   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();
+function describePlan(account: CodexAccount): string {
+  const name = account.label ?? account.id.slice(0, 8);
+  return account.plan ? `${name} (${account.plan})` : name;
+}
 
-  try {
-    await ctx.openBrowser(auth.url);
-  } catch {
-    ctx.ui.notify(`Open this authorization URL:\n${auth.url}`, "warning");
+/** 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 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 blocked ? { remainingPercent: 0, checkedAt: quota.checkedAt ?? Date.now(), blockedUntil: blocked } : undefined;
   }
 
-  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);
+  const remainingPercent = Math.min(...windows.map((window) => window!.remainingPercent!));
+  const resetTimes = windows
+    .filter((window) => (window!.remainingPercent ?? 0) <= 0 && window!.resetsAt)
+    .map((window) => Date.parse(window!.resetsAt!))
+    .filter(Number.isFinite);
+  const resetAt = resetTimes.length > 0 ? Math.min(...resetTimes) : undefined;
+  return {
+    remainingPercent: blocked ? 0 : remainingPercent,
+    resetAt,
+    checkedAt: quota.checkedAt ?? Date.now(),
+    // 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),
+  };
 }
 
-function describePlan(account: CodexAccount): string {
-  return account.plan ? `${account.label ?? account.id.slice(0, 8)} (${account.plan})` : (account.label ?? account.id.slice(0, 8));
+/** 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";
 }
 
-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;
+/**
+ * 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),
+      })),
+    };
   },
 
-  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);
+  saveAccount(account: PooledOAuthAccount) {
+    const { quota: _quota, identity: _identity, type: _type, ...rest } = account as PooledOAuthAccount & CodexAccount;
+    const claims = claimsOf(rest.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(),
+      ...rest,
+      accountId: rest.accountId ?? claims.accountId,
+      plan: rest.plan ?? claims.plan,
     });
-    return account.label ?? account.id;
   },
 
-  async setEnabled(accountId: string, enabled: boolean): Promise<void> {
+  saveMode(mode) {
     const storage = loadCodexAccounts();
-    const account = storage.accounts.find((a) => a.id === accountId);
-    if (!account) throw new Error(`Account “${accountId}” not found.`);
-    account.enabled = enabled;
+    storage.routing = { ...(storage.routing ?? {}), mode };
     saveCodexAccounts(storage);
   },
 
-  async rename(accountId: string, label: string): Promise<void> {
+  primaryQuota() {
     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);
+    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);
+  const blockedUntil = Date.now() + (Number.isFinite(seconds) && seconds >= 0 ? seconds * 1_000 : 60_000);
+  const storage = loadCodexAccounts();
+  if (accountId === MAIN_ACCOUNT_ID) {
+    storage.main = { ...storage.main, blockedUntil };
+    saveCodexAccounts(storage);
+    return;
+  }
+  const account = storage.accounts.find((candidate) => candidate.id === accountId);
+  if (account) saveCodexAccount({ ...account, blockedUntil });
+}
 
-  routing: {
-    async get(): Promise<RoutingMode> {
-      return loadCodexAccounts().routing?.mode === "optimal" ? "optimal" : "standard";
-    },
-    async set(mode: RoutingMode): Promise<RoutingMode> {
-      const storage = loadCodexAccounts();
-      storage.routing = { ...(storage.routing ?? {}), mode };
-      saveCodexAccounts(storage);
-      return mode;
-    },
-    describe(mode: RoutingMode): string {
-      return mode === "optimal"
-        ? "Balances across Codex subscriptions using remaining quota and time until reset."
-        : "Main Codex subscription first, then eligible fallbacks.";
-    },
+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.
+    }
   },
 };
+
+export const codexAccounts = createPooledOAuthAdapter(CODEX_SPEC);
+
+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

@@ -0,0 +1,18 @@
+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<"anthropic-messages"> = {
+  id: "kimi-coding",
+  label: "Kimi",
+  createProvider: kimiCodingProvider,
+};
+
+export const XAI_SPEC: PooledOAuthProviderSpec<"openai-responses"> = {
+  id: "xai",
+  label: "Grok",
+  createProvider: xaiProvider,
+};
+
+export const kimiAccounts = createPooledOAuthAdapter(KIMI_SPEC);
+export const xaiAccounts = createPooledOAuthAdapter(XAI_SPEC);