@llblab/pi-codex-usage 0.10.1 → 0.12.0

package/lib/extension.ts

@@ -0,0 +1,25 @@
+/** Domain: extension composition. Owns: provider registration, Fast bridge and status wiring. Excludes: shared command ownership. */
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { registerFastProvider } from "@llblab/pi-command-fast";
+import { applyFastToRequest, isFastEnabled, isFastEligibleModel, toggleFast } from "./fast.ts";
+import { createCodexUsageStatus } from "./status.ts";
+
+export default function codexUsage(pi: ExtensionAPI) {
+  let releaseFast: (() => void) | undefined;
+  pi.on("session_start", (_event, ctx) => {
+    releaseFast = registerFastProvider(pi, ctx, {
+      provider: "openai-codex",
+      toggle(commandCtx) {
+        toggleFast(commandCtx.model!.id);
+        status.fastChanged(commandCtx, commandCtx.model!);
+      },
+    });
+  });
+  pi.on("session_shutdown", () => { releaseFast?.(); releaseFast = undefined; });
+  const status = createCodexUsageStatus(pi);
+  pi.on("before_provider_request", (event, ctx) => {
+    const model = ctx.model;
+    if (!isFastEligibleModel(model)) return;
+    return applyFastToRequest(event.payload, model.id, isFastEnabled(model.id));
+  });
+}

package/lib/fast.ts

@@ -0,0 +1,23 @@
+/** Domain: Codex Fast. Owns: provider semantics and wire adaptation. Excludes: command arbitration and JSONC editing. */
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { isModelOverrideValue, toggleModelOverrideValue } from "@llblab/pi-command-fast";
+import { isOpenAICodexModel } from "./usage.ts";
+
+const target = (id: string) => ({ provider: "openai-codex", id, property: "serviceTier", enabledValue: "priority" });
+
+export function isFastEligibleModel(model: ExtensionContext["model"]): model is NonNullable<ExtensionContext["model"]> {
+  return !!model && isOpenAICodexModel(model);
+}
+export function isFastEnabled(modelId: string, path?: string): boolean {
+  return isModelOverrideValue(target(modelId), path);
+}
+export function toggleFast(modelId: string, path?: string): boolean {
+  return toggleModelOverrideValue(target(modelId), path);
+}
+
+export function applyFastToRequest(payload: unknown, modelId: string, enabled: boolean): unknown {
+  if (!enabled || !payload || typeof payload !== "object" || Array.isArray(payload)) return undefined;
+  const body = payload as Record<string, unknown>;
+  if (body.model !== modelId || "service_tier" in body) return undefined;
+  return { ...body, service_tier: "priority" };
+}

package/lib/query.ts

@@ -0,0 +1,368 @@
+/** Domain: quota transport. Owns: Pi auth, Codex app-server fallback and errors. Excludes: shared leadership and UI. */
+import { type ChildProcessWithoutNullStreams, spawn } from "node:child_process";
+import { createInterface } from "node:readline";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { assertObject, errorMessage, normalizeAppServerResponse, normalizeBackendPayload, parseJsonObject, selectUsageSnapshot, type AppServerRateLimitResponse, type CodexUsageModel, type CodexUsageReport, type RateLimitStatusPayload, type UsageQueryError } from "./usage.ts";
+const CODEX_PROVIDER_ID = "openai-codex";
+const CODEX_USAGE_LIMIT_ID = "codex";
+const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
+const MAX_ERROR_BODY_CHARS = 600;
+type PiModel = NonNullable<ExtensionContext["model"]>;
+type QueryUsageOptions = { timeoutMs: number };
+export type QueryUsageResult = { ok: true; report: CodexUsageReport } | { ok: false; errors: UsageQueryError[] };
+type RpcResponse = {
+  id?: unknown;
+  result?: unknown;
+  error?: { message?: unknown; code?: unknown };
+};
+
+type PendingRpc = {
+  resolve: (value: unknown) => void;
+  reject: (error: Error) => void;
+};
+
+export async function queryUsage(
+  ctx: ExtensionContext,
+  options: Pick<QueryUsageOptions, "timeoutMs">,
+  model: CodexUsageModel | undefined,
+  mayQuery: () => boolean,
+): Promise<QueryUsageResult | undefined> {
+  const errors: UsageQueryError[] = [];
+  const sources = ["pi-auth", "codex-app-server"] as const;
+
+  for (const source of sources) {
+    if (!mayQuery()) return undefined;
+    try {
+      const report =
+        source === "pi-auth"
+          ? await queryViaPiAuth(ctx, options.timeoutMs, mayQuery)
+          : await queryViaCodexAppServer(options.timeoutMs, mayQuery);
+      if (!report) return undefined;
+      if (
+        selectUsageSnapshot(report, CODEX_USAGE_LIMIT_ID) ||
+        report.credits
+      ) {
+        return { ok: true, report };
+      }
+      errors.push({
+        source,
+        message: `${source} returned no displayable codex rate-limit windows`,
+      });
+    } catch (cause) {
+      errors.push({ source, message: errorMessage(cause), cause });
+    }
+  }
+
+  return { ok: false, errors };
+}
+
+async function queryViaPiAuth(
+  ctx: ExtensionContext,
+  timeoutMs: number,
+  mayQuery: () => boolean,
+): Promise<CodexUsageReport | undefined> {
+  const auth = await resolvePiCodexAuth(ctx);
+  if (!auth) {
+    throw new Error(
+      "No Pi OpenAI Codex subscription auth was available. Use a Pi OpenAI Codex model or run /login for OpenAI ChatGPT Plus/Pro (Codex).",
+    );
+  }
+
+  if (!mayQuery()) return undefined;
+  const response = await fetchWithTimeout(
+    CODEX_USAGE_URL,
+    { headers: auth.headers },
+    timeoutMs,
+  );
+  const text = await response.text();
+  if (!response.ok) {
+    throw new Error(
+      `Codex usage endpoint returned ${response.status} ${response.statusText}: ${redactErrorBody(text)}`,
+    );
+  }
+
+  const payload = parseJsonObject(text, "Codex usage endpoint response");
+  return normalizeBackendPayload(
+    payload as RateLimitStatusPayload,
+    Date.now(),
+    "pi-auth",
+  );
+}
+
+async function resolvePiCodexAuth(
+  ctx: ExtensionContext,
+): Promise<{ headers: Record<string, string> } | undefined> {
+  const models = codexAuthCandidateModels(ctx);
+  const errors: string[] = [];
+
+  for (const model of models) {
+    const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+    if (!auth.ok) {
+      errors.push(auth.error);
+      continue;
+    }
+
+    const headers: Record<string, string> = {};
+    for (const [name, value] of Object.entries(auth.headers ?? {})) {
+      if (value !== null) headers[name] = value;
+    }
+    if (!hasHeader(headers, "Authorization") && auth.apiKey) {
+      headers.Authorization = `Bearer ${auth.apiKey}`;
+    }
+    if (!hasHeader(headers, "User-Agent")) {
+      headers["User-Agent"] = "pi-codex-usage";
+    }
+    if (hasHeader(headers, "Authorization")) {
+      return { headers };
+    }
+  }
+
+  if (errors.length > 0) {
+    throw new Error(errors.join("; "));
+  }
+  return undefined;
+}
+
+function codexAuthCandidateModels(ctx: ExtensionContext): PiModel[] {
+  const candidates: PiModel[] = [];
+  const seen = new Set<string>();
+  const add = (model: PiModel | undefined) => {
+    if (!model || model.provider !== CODEX_PROVIDER_ID) return;
+    const key = `${model.provider}/${model.id}`;
+    if (seen.has(key)) return;
+    seen.add(key);
+    candidates.push(model);
+  };
+
+  add(ctx.model);
+  for (const model of ctx.modelRegistry.getAvailable()) add(model);
+  for (const model of ctx.modelRegistry.getAll()) add(model);
+  return candidates;
+}
+
+async function fetchWithTimeout(
+  url: string,
+  init: RequestInit,
+  timeoutMs: number,
+): Promise<Response> {
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    return await fetch(url, { ...init, signal: controller.signal });
+  } catch (error) {
+    if (controller.signal.aborted) {
+      throw new Error(
+        `Timed out after ${Math.round(timeoutMs / 1000)}s while fetching Codex usage.`,
+      );
+    }
+    throw error;
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+async function queryViaCodexAppServer(
+  timeoutMs: number,
+  mayQuery: () => boolean,
+): Promise<CodexUsageReport | undefined> {
+  const client = new CodexAppServerClient(timeoutMs);
+  try {
+    await client.start();
+    if (!mayQuery()) return undefined;
+    await client.request("initialize", {
+      clientInfo: {
+        name: "pi_codex_usage",
+        title: "Pi Codex Usage",
+        version: "0.1.0",
+      },
+      capabilities: {
+        experimentalApi: false,
+        requestAttestation: false,
+        optOutNotificationMethods: [],
+      },
+    });
+    client.notify("initialized");
+    if (!mayQuery()) return undefined;
+    const result = await client.request("account/rateLimits/read", undefined);
+    return normalizeAppServerResponse(
+      assertObject(
+        result,
+        "account/rateLimits/read result",
+      ) as AppServerRateLimitResponse,
+      Date.now(),
+    );
+  } finally {
+    client.dispose();
+  }
+}
+
+class CodexAppServerClient {
+  private child?: ChildProcessWithoutNullStreams;
+  private nextId = 1;
+  private stderr = "";
+  private readonly pending = new Map<number, PendingRpc>();
+  private startPromise?: Promise<void>;
+  private exitError?: Error;
+  private readonly timeoutMs: number;
+
+  constructor(timeoutMs: number) {
+    this.timeoutMs = timeoutMs;
+  }
+
+  start(): Promise<void> {
+    if (this.startPromise) return this.startPromise;
+
+    this.startPromise = new Promise((resolve, reject) => {
+      const child = spawn("codex", ["app-server", "--listen", "stdio://"], {
+        stdio: ["pipe", "pipe", "pipe"],
+      });
+      this.child = child;
+
+      const startupTimeout = setTimeout(() => {
+        reject(
+          new Error(
+            `Timed out after ${Math.round(this.timeoutMs / 1000)}s starting codex app-server.`,
+          ),
+        );
+      }, this.timeoutMs);
+
+      child.once("spawn", () => {
+        clearTimeout(startupTimeout);
+        resolve();
+      });
+
+      child.once("error", (error) => {
+        clearTimeout(startupTimeout);
+        reject(new Error(`Failed to start codex app-server: ${error.message}`));
+        this.rejectAll(error);
+      });
+
+      child.once("exit", (code, signal) => {
+        const suffix = this.stderr
+          ? ` stderr: ${redactErrorBody(this.stderr)}`
+          : "";
+        this.exitError = new Error(
+          `codex app-server exited before completing the request (code ${code ?? "unknown"}, signal ${signal ?? "none"}).${suffix}`,
+        );
+        this.rejectAll(this.exitError);
+      });
+
+      child.stderr.setEncoding("utf8");
+      child.stderr.on("data", (chunk: string) => {
+        this.stderr = truncateEnd(this.stderr + chunk, MAX_ERROR_BODY_CHARS);
+      });
+
+      const lines = createInterface({ input: child.stdout });
+      lines.on("line", (line) => this.handleLine(line));
+    });
+
+    return this.startPromise;
+  }
+
+  request(method: string, params: unknown): Promise<unknown> {
+    const child = this.child;
+    if (!child?.stdin.writable) {
+      throw new Error("codex app-server is not running.");
+    }
+    if (this.exitError) throw this.exitError;
+
+    const id = this.nextId++;
+    const payload =
+      params === undefined ? { method, id } : { method, id, params };
+    const response = new Promise<unknown>((resolve, reject) => {
+      const timeout = setTimeout(() => {
+        this.pending.delete(id);
+        reject(
+          new Error(
+            `Timed out after ${Math.round(this.timeoutMs / 1000)}s waiting for ${method}.`,
+          ),
+        );
+      }, this.timeoutMs);
+
+      this.pending.set(id, {
+        resolve: (value) => {
+          clearTimeout(timeout);
+          resolve(value);
+        },
+        reject: (error) => {
+          clearTimeout(timeout);
+          reject(error);
+        },
+      });
+    });
+
+    child.stdin.write(`${JSON.stringify(payload)}\n`);
+    return response;
+  }
+
+  notify(method: string): void {
+    const child = this.child;
+    if (!child?.stdin.writable) return;
+    child.stdin.write(`${JSON.stringify({ method })}\n`);
+  }
+
+  dispose(): void {
+    for (const [id, pending] of this.pending) {
+      pending.reject(new Error(`codex app-server request ${id} cancelled.`));
+    }
+    this.pending.clear();
+
+    const child = this.child;
+    if (!child) return;
+    child.stdin.end();
+    if (!child.killed) child.kill();
+    this.child = undefined;
+  }
+
+  private handleLine(line: string): void {
+    let parsed: RpcResponse;
+    try {
+      parsed = JSON.parse(line) as RpcResponse;
+    } catch {
+      return;
+    }
+
+    if (typeof parsed.id !== "number") return;
+    const pending = this.pending.get(parsed.id);
+    if (!pending) return;
+    this.pending.delete(parsed.id);
+
+    if (parsed.error) {
+      const message =
+        typeof parsed.error.message === "string"
+          ? parsed.error.message
+          : "unknown error";
+      pending.reject(new Error(`codex app-server request failed: ${message}`));
+      return;
+    }
+
+    pending.resolve(parsed.result);
+  }
+
+  private rejectAll(error: Error): void {
+    for (const pending of this.pending.values()) pending.reject(error);
+    this.pending.clear();
+  }
+}
+
+
+function hasHeader(headers: Record<string, string>, name: string): boolean {
+  return Object.keys(headers).some(
+    (key) => key.toLowerCase() === name.toLowerCase(),
+  );
+}
+
+function redactErrorBody(body: string): string {
+  return truncateEnd(
+    body
+      .replace(/Bearer\s+[A-Za-z0-9._~+/=-]+/gi, "Bearer <redacted>")
+      .replace(/"access_token"\s*:\s*"[^"]+"/gi, '"access_token":"<redacted>"')
+      .trim(),
+    MAX_ERROR_BODY_CHARS,
+  );
+}
+
+function truncateEnd(value: string, maxChars: number): string {
+  if (value.length <= maxChars) return value;
+  return `${value.slice(0, maxChars - 1)}…`;
+}