@jameslovespancakes/pi-plus 1.0.11 → 1.0.12

package/src/domains/subscriptions/provider.ts

@@ -5,50 +5,36 @@ import {
 } 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.
+ * Anthropic uses pi's Messages API client with per-request account routing.
+ * Routing reads cached quota; a 429 affects the next request, not its stream.
  */
 
+/**
+ * The hook is global, so require both Anthropic's payload shape and a Claude
+ * model. This prevents `system` injection into other provider requests.
+ */
+function isAnthropicMessagesPayload(payload: any): boolean {
+  if ("instructions" in payload || "input" in payload) return false;
+  if (!Array.isArray(payload.messages)) return false;
+  const model = typeof payload.model === "string" ? payload.model.toLowerCase() : "";
+  return model.startsWith("claude") || ANTHROPIC_MODELS.some((m: any) => m.id === payload.model);
+}
+
 let lastSelected: { id: string; at: number } | undefined;
+const accountLastUsed = new Map<string, number>();
 
 /** 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.
+/** Routes to a pooled account, falling back to pi's primary credential. */
 export function routeAccessToken(primary: string, modelId?: string, _sessionId?: string): string {
   const storage = loadAccounts();
   if (!storage) return primary;
@@ -57,10 +43,23 @@ export function routeAccessToken(primary: string, modelId?: string, _sessionId?:
   const family = familyForModel(modelId);
 
   const candidates: Candidate[] = [
-    { id: MAIN_ACCOUNT_ID, access: primary, quota: storage.main?.quota as any, order: 0 },
+    {
+      id: MAIN_ACCOUNT_ID,
+      access: primary,
+      quota: storage.main?.quota as any,
+      order: 0,
+      lastUsed: accountLastUsed.get(MAIN_ACCOUNT_ID) ?? Number(storage.main?.lastUsed ?? 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 })),
+      .map((a, index) => ({
+        id: a.id,
+        access: a.access,
+        quota: a.quota,
+        order: index + 1,
+        lastUsed: accountLastUsed.get(a.id) ?? a.lastUsed ?? 0,
+        account: a,
+      })),
   ];
 
   // Nothing to choose between: keep pi's own credential.
@@ -75,10 +74,9 @@ export function routeAccessToken(primary: string, modelId?: string, _sessionId?:
 
   const now = Date.now();
   lastSelected = { id: picked.candidate.id, at: now };
+  accountLastUsed.set(picked.candidate.id, now);
   if (picked.candidate.account) {
-    // `lastUsed` only orders the footer's account list, so second-level
-    // precision buys nothing. Persisting it on every request rewrote two
-    // credential files per call; once a minute is indistinguishable in the UI.
+    // Minute precision avoids a credential write on every request.
     const previous = picked.candidate.account.lastUsed ?? 0;
     if (now - previous > 60_000) {
       saveAccount({ ...picked.candidate.account, lastUsed: now });
@@ -102,9 +100,7 @@ export function registerAnthropicProvider(pi: ExtensionAPI): void {
     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.
+    // Required for subscription billing. See client-identity.ts.
     headers: clientIdentityHeaders(),
     oauth: {
       name: "Anthropic Claude Pro/Max",
@@ -114,29 +110,15 @@ export function registerAnthropicProvider(pi: ExtensionAPI): void {
         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.
-   */
+  /** Rebuilds the Anthropic prompt, then signs the final body. */
   pi.on("before_provider_request", async (event: any) => {
     const payload = event?.payload;
     if (!payload || typeof payload !== "object") return undefined;
+    if (!isAnthropicMessagesPayload(payload)) return undefined;
 
     const serialized = typeof payload.system === "string"
       ? payload.system
@@ -156,20 +138,10 @@ export function registerAnthropicProvider(pi: ExtensionAPI): void {
     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.
-   */
+  /** Updates the routed account from response quota headers. */
   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.
+    // The primary account lives in pi's auth store.
     if (!routed || routed.id === MAIN_ACCOUNT_ID) return;
     try {
       applyQuotaHeaders(routed.id, event?.headers);
@@ -178,44 +150,7 @@ export function registerAnthropicProvider(pi: ExtensionAPI): void {
     }
   });
 
-  /**
-   * 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.
-   */
+  /** Refreshes stale idle-account quota when a message starts. */
   pi.on("input", async () => {
     void refreshAllQuota().catch(() => {});
   });

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

@@ -6,7 +6,6 @@ import {
   loadAccounts,
   saveAccount,
   setRoutingMode,
-  type RoutingMode as StoreRoutingMode,
 } from "../../../core/anthropic/store.ts";
 
 /**
@@ -22,16 +21,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");
 
@@ -133,17 +122,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,7 +1,13 @@
 import { randomUUID } from "node:crypto";
+import type { ModelAuth, OAuthAuth, 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 { applyCodexQuotaHeaders } from "../../../core/codex/quota.ts";
 import {
+  MAIN_ACCOUNT_ID,
   claimsOf,
   loadCodexAccounts,
   saveCodexAccount,
@@ -38,6 +44,18 @@ 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",
@@ -131,7 +149,7 @@ export const codexAccounts: AccountProvider = {
 
   routing: {
     async get(): Promise<RoutingMode> {
-      return loadCodexAccounts().routing?.mode === "optimal" ? "optimal" : "standard";
+      return codexRoutingMode(loadCodexAccounts().routing?.mode);
     },
     async set(mode: RoutingMode): Promise<RoutingMode> {
       const storage = loadCodexAccounts();
@@ -140,9 +158,158 @@ export const codexAccounts: AccountProvider = {
       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.";
+      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);
+}
+
+function quotaState(quota: CodexAccount["quota"], blockedUntil?: number): AccountQuotaState | undefined {
+  if (!quota) {
+    return blockedUntil && blockedUntil > Date.now()
+      ? { remainingPercent: 0, checkedAt: Date.now(), blockedUntil }
+      : undefined;
+  }
+  const windows = [quota.five_hour, quota.seven_day, ...(quota.scoped ?? [])]
+    .filter((window) => window && Number.isFinite(window.remainingPercent));
+  if (windows.length === 0) return undefined;
+
+  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: blockedUntil && blockedUntil > Date.now() ? 0 : remainingPercent,
+    resetAt,
+    checkedAt: quota.checkedAt ?? Date.now(),
+    blockedUntil: blockedUntil && blockedUntil > Date.now()
+      ? blockedUntil
+      : remainingPercent <= 0 ? resetAt ?? Number.POSITIVE_INFINITY : 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,
+        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;
+
+  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;
+}
+
+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);
+}
+
+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 });
+}
+
+export function registerCodexProvider(pi: ExtensionAPI): void {
+  const provider = openaiCodexProvider();
+  const oauth = provider.auth.oauth;
+  if (!oauth) return;
+
+  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);
+    },
+  });
+
+  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)),
+  });
+}

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