@jameslovespancakes/pi-plus 1.0.0

package/src/domains/subscriptions/provider.ts

@@ -0,0 +1,219 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { authorize, exchange, refreshToken } from "../../core/anthropic/oauth.ts";
+import {
+  billingHeader, clientIdentityHeaders, firstUserText, prependPromptBlock, signRequestBody, splitSystemPrompt,
+} from "../../core/anthropic/client-identity.ts";
+import { ANTHROPIC_MODELS } from "../../core/anthropic/models.ts";
+import { applyQuotaHeaders, refreshAllQuota } from "../../core/anthropic/quota.ts";
+import { applyCodexQuotaHeaders } from "../../core/codex/quota.ts";
+import { loadCodexAccounts } from "../../core/codex/store.ts";
+import {
+  MAIN_ACCOUNT_ID, familyForModel, selectAccount, type Candidate,
+} from "../../core/anthropic/routing.ts";
+import { getRoutingMode, loadAccounts, saveAccount } from "../../core/anthropic/store.ts";
+
+/**
+ * The Anthropic provider.
+ *
+ * Rather than reimplementing the Messages API, this registers pi's built-in
+ * `anthropic-messages` implementation and supplies only the credential. pi
+ * resolves auth per request, so `getApiKey` is the hook where account routing
+ * happens: it returns the token of whichever account the router picked.
+ *
+ * Consequences of that choice, stated plainly:
+ *   - `getApiKey` is synchronous, so selection reads the CACHED quota snapshot.
+ *     A background timer keeps it fresh instead of polling inline.
+ *   - A mid-stream 429 is handled by pi's retry rather than by migrating the
+ *     in-flight request to another account. The next request routes elsewhere
+ *     once the failure is recorded, which is a real behavioural difference from
+ *     the vendored provider.
+ *
+ * What this buys: pi's streaming, tool conversion, cost accounting and retry
+ * logic are reused unchanged, instead of a 1,500-line reimplementation whose
+ * edge cases we could not see.
+ */
+
+let lastSelected: { id: string; at: number } | undefined;
+
+/** Which account served the most recent request, for the UI. */
+export function lastRoutedAccount(): { id: string; at: number } | undefined {
+  return lastSelected;
+}
+
+/**
+ * Chooses an account and returns its access token.
+ *
+ * `primary` is pi's own credential, which participates as `main` at order 0.
+ * Falls back to the primary token whenever routing has nothing better, so a
+ * single-account setup behaves exactly as it did before any of this existed.
+ */
+// `_sessionId` is accepted for signature compatibility with pi's per-request
+// hook; stickiness is derived from stored assignments rather than the id.
+export function routeAccessToken(primary: string, modelId?: string, _sessionId?: string): string {
+  const storage = loadAccounts();
+  if (!storage) return primary;
+
+  const mode = getRoutingMode(storage);
+  const family = familyForModel(modelId);
+
+  const candidates: Candidate[] = [
+    { id: MAIN_ACCOUNT_ID, access: primary, quota: storage.main?.quota as any, order: 0 },
+    ...storage.accounts
+      .filter((a) => a.enabled !== false && a.type === "oauth" && a.access)
+      .map((a, index) => ({ id: a.id, access: a.access, quota: a.quota, order: index + 1, account: a })),
+  ];
+
+  // Nothing to choose between: keep pi's own credential.
+  if (candidates.length <= 1) return primary;
+
+  const picked = selectAccount({ candidates, family, modelId, mode });
+  if (!picked) {
+    // Everything is quota-blocked. Returning the primary lets Anthropic issue
+    // the authoritative 429 rather than inventing a local failure.
+    return primary;
+  }
+
+  lastSelected = { id: picked.candidate.id, at: Date.now() };
+  if (picked.candidate.account) {
+    saveAccount({ ...picked.candidate.account, lastUsed: Date.now() });
+  }
+  return picked.candidate.access ?? primary;
+}
+
+async function login(callbacks: any) {
+  const auth = await authorize("max");
+  callbacks.onAuth({ url: auth.url });
+  const pasted = await callbacks.onPrompt({ message: "Paste the Claude OAuth callback URL or code:" });
+  const result = await exchange(pasted, auth.verifier, auth.redirectUri, auth.state);
+  if (result.type !== "success") throw new Error(`Anthropic OAuth failed: ${result.reason}`);
+  return { access: result.access, refresh: result.refresh, expires: result.expires };
+}
+
+export function registerAnthropicProvider(pi: ExtensionAPI): void {
+  pi.registerProvider("anthropic", {
+    name: "Anthropic",
+    baseUrl: "https://api.anthropic.com",
+    api: "anthropic-messages",
+    models: ANTHROPIC_MODELS,
+    // See core/anthropic/client-identity.ts. Without these, Anthropic bills
+    // requests to extra usage rather than plan limits. Delete that import and
+    // this line to opt out; nothing else depends on it.
+    headers: clientIdentityHeaders(),
+    oauth: {
+      name: "Anthropic Claude Pro/Max",
+      isSubscription: true,
+      login,
+      refreshToken: async (credentials: any) => {
+        const refreshed = await refreshToken({ refreshToken: credentials.refresh });
+        return { access: refreshed.access, refresh: refreshed.refresh, expires: refreshed.expires };
+      },
+      // The routing hook. pi calls this per request.
+      getApiKey: (credentials: any) => routeAccessToken(credentials.access),
+    },
+  });
+
+  /**
+   * Reshapes the request the way the endpoint requires, then signs it.
+   *
+   * Three things have to happen, in order:
+   *   1. pi's system prompt is split. The documentation paragraph cannot live in
+   *      `system` at all, or Anthropic answers 400 "Third-party apps now draw
+   *      from your extra usage"; it moves into the first user message instead.
+   *   2. The billing header goes first in `system`, alongside the Claude Code
+   *      identity line.
+   *   3. The body is signed, which must happen last because the checksum covers
+   *      the canonicalised body.
+   *
+   * See core/anthropic/client-identity.ts for what all of this is and the
+   * caveats around it.
+   */
+  pi.on("before_provider_request", async (event: any) => {
+    const payload = event?.payload;
+    if (!payload || typeof payload !== "object") return undefined;
+
+    const serialized = typeof payload.system === "string"
+      ? payload.system
+      : (payload.system ?? []).map((b: any) => b?.text ?? "").join("\n\n");
+    const split = splitSystemPrompt(serialized);
+
+    const messages = structuredClone(payload.messages ?? []);
+    prependPromptBlock(messages, split.messageText);
+
+    const system = [
+      { type: "text", text: billingHeader(firstUserText(messages)) },
+      { type: "text", text: "You are Claude Code, Anthropic's official CLI for Claude." },
+      ...(split.systemText ? [{ type: "text", text: split.systemText }] : []),
+    ];
+
+    const signed = await signRequestBody(JSON.stringify({ ...payload, system, messages }));
+    return JSON.parse(signed);
+  });
+
+  /**
+   * Free quota refresh for whichever account just served a request.
+   *
+   * Anthropic returns the same utilisation numbers as the usage endpoint in
+   * `anthropic-ratelimit-unified-*` response headers, so the active account
+   * stays current without spending a request on asking.
+   *
+   * `lastRoutedAccount()` is set by `getApiKey` immediately before the call,
+   * which is what lets the reply be attributed to the right account.
+   */
+  pi.on("after_provider_response", (event: any) => {
+    const routed = lastRoutedAccount();
+    // `main` lives in pi's auth.json, not in storage.accounts, so there is no
+    // per-account record to update for it.
+    if (!routed || routed.id === MAIN_ACCOUNT_ID) return;
+    try {
+      applyQuotaHeaders(routed.id, event?.headers);
+    } catch {
+      // Never let bookkeeping disturb a response.
+    }
+  });
+
+  /**
+   * The same trick for Codex, which reports usage as `x-codex-*` headers.
+   *
+   * Codex has no usage endpoint, so headers are the ONLY quota source: an
+   * account that is not serving traffic keeps whatever snapshot it last had.
+   * Attribution is by `chatgpt-account-id` on the request, since Codex
+   * requests carry the account explicitly rather than only in the token.
+   */
+  pi.on("after_provider_response", (event: any) => {
+    const headers = event?.headers;
+    if (!headers) return;
+    const hasCodexQuota = Object.keys(headers).some((k) => k.toLowerCase().startsWith("x-codex-"));
+    if (!hasCodexQuota) return;
+    try {
+      const sent = event?.request?.headers ?? {};
+      const accountId = sent["chatgpt-account-id"] ?? sent["Chatgpt-Account-Id"];
+      const match = loadCodexAccounts().accounts.find(
+        (a) => (accountId ? a.accountId === accountId : false) || a.enabled !== false,
+      );
+      if (match) applyCodexQuotaHeaders(match.id, headers);
+    } catch {
+      // Bookkeeping only.
+    }
+  });
+
+  /**
+   * Idle accounts still need polling, because routing compares accounts and
+   * response headers only ever refresh the one that served traffic.
+   *
+   * Triggered by sending a message, not by a timer, and `refreshAllQuota`
+   * skips accounts whose snapshot is still fresh. So the poll rate is at most
+   * once per QUOTA_FRESH_MS however fast messages are sent, and an idle
+   * session costs nothing at all. The previous 5 minute interval ran forever
+   * regardless of activity and was being answered with 429s.
+   *
+   * `input` is the per-message hook; `turn_start` would fire once per LLM
+   * response, which is many times per message in a tool-use loop.
+   */
+  pi.on("input", async () => {
+    void refreshAllQuota().catch(() => {});
+  });
+
+  pi.on("session_start", async () => {
+    void refreshAllQuota().catch(() => {});
+  });
+}

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

@@ -0,0 +1,149 @@
+import { randomUUID } from "node:crypto";
+import type { AccountContext, AccountProvider, ManagedAccount, RoutingMode } from "../../../core/accounts/registry.ts";
+import { authorize, exchange } from "../../../core/anthropic/oauth.ts";
+import {
+  getRoutingMode,
+  loadAccounts,
+  saveAccount,
+  setRoutingMode,
+  type RoutingMode as StoreRoutingMode,
+} from "../../../core/anthropic/store.ts";
+
+/**
+ * Anthropic adapter.
+ *
+ * pi itself holds one Anthropic credential; the additional accounts live in
+ * our own store and tokens are swapped per request. Everything
+ * provider-specific lives here so `/accounts` and `/routing` stay generic.
+ *
+ * This used to read through `vendor/anthropic.ts`, which loaded the cortexkit
+ * package. That package was removed when the provider was extracted, so every
+ * call here threw and the account list silently came back empty. It now uses
+ * 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");
+
+  try {
+    await ctx.openBrowser(auth.url);
+  } catch {
+    ctx.ui.notify(`Open this authorization URL:\n${auth.url}`, "warning");
+  }
+
+  const callback = await ctx.ui.input("Claude OAuth", "Paste the callback URL or authorization code");
+  if (!callback) return undefined;
+
+  const result = await exchange(callback, auth.verifier, auth.redirectUri, auth.state);
+  if (result.type !== "success") {
+    throw new Error(`Claude OAuth exchange failed: ${result.reason}`);
+  }
+  return { access: result.access, refresh: result.refresh, expires: result.expires };
+}
+
+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) => ({
+        id: account.id,
+        label: account.label ?? account.id.slice(0, 8),
+        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 Claude account",
+      "Use a different authorized Anthropic account in the browser. Continue?",
+    );
+    if (!proceed) return undefined;
+
+    const result = await authenticate(ctx);
+    if (!result) return undefined;
+
+    const now = Date.now();
+    saveAccount({
+      id: randomUUID(),
+      label,
+      type: "oauth",
+      enabled: true,
+      access: result.access,
+      refresh: result.refresh,
+      expires: result.expires,
+      addedAt: now,
+      lastRefreshedAt: now,
+      authLineageId: randomUUID(),
+    });
+    return label;
+  },
+
+  async reauth(ctx, accountId): Promise<string | undefined> {
+    const storage = loadAccounts();
+    const account = storage?.accounts.find(
+      (candidate) => candidate.id === accountId || candidate.label === accountId,
+    );
+    if (!account || account.type !== "oauth") {
+      throw new Error(`OAuth account “${accountId}” not found.`);
+    }
+
+    const result = await authenticate(ctx);
+    if (!result) return undefined;
+
+    saveAccount({
+      ...account,
+      access: result.access,
+      refresh: result.refresh,
+      expires: result.expires,
+      lastRefreshedAt: Date.now(),
+      authLineageId: randomUUID(),
+    });
+    return account.label ?? account.id;
+  },
+
+  async setEnabled(accountId: string, enabled: boolean): Promise<void> {
+    const storage = loadAccounts();
+    const account = storage?.accounts.find((candidate) => candidate.id === accountId);
+    if (!account) throw new Error(`Account “${accountId}” not found.`);
+    // Credentials are preserved; only the eligibility flag changes, so a
+    // disabled account can be re-enabled without another OAuth round trip.
+    saveAccount({ ...account, enabled });
+  },
+
+  async rename(accountId: string, label: string): Promise<void> {
+    const account = loadAccounts()?.accounts.find((candidate) => candidate.id === accountId);
+    if (!account) throw new Error(`Account “${accountId}” not found.`);
+    saveAccount({ ...account, label });
+  },
+
+  routing: {
+    async get(): Promise<RoutingMode> {
+      return fromStoreMode(getRoutingMode(loadAccounts()));
+    },
+    async set(mode: RoutingMode): Promise<RoutingMode> {
+      setRoutingMode(MODE_TO_STORE[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.";
+    },
+  },
+};

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

@@ -0,0 +1,148 @@
+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 {
+  claimsOf,
+  loadCodexAccounts,
+  saveCodexAccount,
+  saveCodexAccounts,
+  type CodexAccount,
+} from "../../../core/codex/store.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.
+ */
+
+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));
+}
+
+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 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.";
+    },
+  },
+};

package/src/domains/subscriptions/routing.ts

@@ -0,0 +1,72 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { accountProvider, routableProviders, type RoutingMode } from "../../core/accounts/registry.ts";
+
+/**
+ * `/routing [standard|optimal] [provider]`
+ *
+ *   standard  main account first, fall back only when exhausted
+ *   optimal   balance across accounts by remaining quota and time to reset
+ *
+ * With no provider named, the mode is applied to every provider that supports
+ * routing. With no mode, the current modes are reported.
+ */
+
+const MODES: RoutingMode[] = ["standard", "optimal"];
+
+function isMode(value: string): value is RoutingMode {
+  return (MODES as string[]).includes(value);
+}
+
+export function registerRoutingCommands(pi: ExtensionAPI): void {
+  pi.registerCommand("routing", {
+    description: "Account routing: standard (main first) or optimal (quota balanced)",
+    getArgumentCompletions: (prefix) =>
+      MODES.filter((mode) => mode.startsWith(prefix)).map((mode) => ({
+        value: mode,
+        label: mode === "optimal" ? "optimal: balance by remaining quota" : "standard: main account first",
+      })),
+    handler: async (args, ctx) => {
+      const [first, second] = args.trim().toLowerCase().split(/\s+/).filter(Boolean);
+      const providers = second ? [accountProvider(second)].filter((entry) => !!entry) : routableProviders();
+
+      if (providers.length === 0) {
+        ctx.ui.notify(
+          second ? `“${second}” does not support account routing.` : "No providers support account routing.",
+          "warning",
+        );
+        return;
+      }
+
+      // No mode given: report current state.
+      if (!first) {
+        const lines: string[] = [];
+        for (const provider of providers) {
+          try {
+            const mode = await provider.routing!.get();
+            lines.push(`${provider.label} (${provider.id}): ${mode}`, `  ${provider.routing!.describe(mode)}`);
+          } catch (error) {
+            lines.push(`${provider.label}: ${error instanceof Error ? error.message : String(error)}`);
+          }
+        }
+        ctx.ui.notify(lines.join("\n"), "info");
+        return;
+      }
+
+      if (!isMode(first)) {
+        ctx.ui.notify(`Usage: /routing [standard|optimal] [provider]`, "warning");
+        return;
+      }
+
+      const results: string[] = [];
+      for (const provider of providers) {
+        try {
+          const mode = await provider.routing!.set(first);
+          results.push(`${provider.label}: ${mode}`, `  ${provider.routing!.describe(mode)}`);
+        } catch (error) {
+          results.push(`${provider.label}: failed, ${error instanceof Error ? error.message : String(error)}`);
+        }
+      }
+      ctx.ui.notify(results.join("\n"), "info");
+    },
+  });
+}