@ai-devkit/agent-manager 0.37.0 → 0.39.0

package/src/capacity/{zai.ts → providers/zai.ts}

@@ -1,16 +1,11 @@
-import { readFile } from "node:fs/promises";
-import { homedir } from "node:os";
-import { join } from "node:path";
 import type {
   Availability,
   CapacityLimitType,
-  CapacityReport,
   CapacityWindow,
-} from "./types.js";
+  ProviderCapacitySnapshot,
+} from "../types.js";
 
 const ZAI_QUOTA_URL = "https://api.z.ai/api/monitor/usage/quota/limit";
-const MISSING_API_KEY_MESSAGE =
-  "z.ai API key not found; set Z_AI_API_KEY or configure Pi provider zai";
 const LIMIT_METADATA = {
   TOKENS_LIMIT: { id: "tokens", label: "Tokens" },
   CREDIT_LIMIT: { id: "credit", label: "Credit" },
@@ -24,13 +19,7 @@ const UNIT_METADATA = new Map<number, { minutes: number; label: string }>([
 ]);
 type UnknownRecord = Record<string, unknown>;
 
-export type ZaiCredentialOptions = {
-  env?: NodeJS.ProcessEnv;
-  readFile?: (path: string, encoding: BufferEncoding) => Promise<string>;
-};
-
-export type ZaiProbeOptions = ZaiCredentialOptions & {
-  checkedAt: string;
+export type ZaiRequestOptions = {
   fetch?: typeof globalThis.fetch;
   timeoutMs?: number;
 };
@@ -41,10 +30,6 @@ function record(value: unknown): UnknownRecord | null {
     : null;
 }
 
-function nonEmptyText(value: unknown): string | null {
-  return typeof value === "string" && value.trim() ? value.trim() : null;
-}
-
 function optionalNumber(value: unknown, field: string): number | null {
   if (value === undefined || value === null) return null;
   if (typeof value !== "number" || !Number.isFinite(value)) {
@@ -53,10 +38,6 @@ function optionalNumber(value: unknown, field: string): number | null {
   return value;
 }
 
-function clampPercent(value: number): number {
-  return Math.max(0, Math.min(100, value));
-}
-
 function integer(value: unknown): number | null {
   return typeof value === "number" && Number.isInteger(value) ? value : null;
 }
@@ -65,50 +46,33 @@ function isCapacityLimitType(value: string): value is CapacityLimitType {
   return Object.hasOwn(LIMIT_METADATA, value);
 }
 
-function durationMinutes(
-  unit: number,
-  number: number,
-  type: CapacityLimitType,
-): number | null {
-  if (type === "TIME_LIMIT" && unit === 5 && number === 1) return 30 * 24 * 60;
+function durationMinutes(unit: number, count: number, type: CapacityLimitType): number | null {
+  if (type === "TIME_LIMIT" && unit === 5 && count === 1) return 30 * 24 * 60;
   const metadata = UNIT_METADATA.get(unit);
-  return metadata === undefined || number <= 0
-    ? null
-    : metadata.minutes * number;
+  return metadata === undefined || count <= 0 ? null : metadata.minutes * count;
 }
 
-function windowLabel(
-  type: CapacityLimitType,
-  unit: number,
-  number: number,
-): string {
-  if (type === "TIME_LIMIT" && unit === 5 && number === 1)
-    return "MCP · monthly";
+function windowLabel(type: CapacityLimitType, unit: number, count: number): string {
+  if (type === "TIME_LIMIT" && unit === 5 && count === 1) return "MCP · monthly";
   const kind = LIMIT_METADATA[type].label;
   const unitLabel = UNIT_METADATA.get(unit)?.label;
   if (!unitLabel) return kind;
-  return `${kind} · ${number} ${unitLabel}${number === 1 ? "" : "s"}`;
+  return `${kind} · ${count} ${unitLabel}${count === 1 ? "" : "s"}`;
 }
 
 function resetTime(value: unknown): string | null {
   const milliseconds = optionalNumber(value, "limit nextResetTime");
   if (milliseconds === null) return null;
   const date = new Date(milliseconds);
-  if (Number.isNaN(date.getTime()))
-    throw new Error("Invalid z.ai quota limit nextResetTime");
+  if (Number.isNaN(date.getTime())) throw new Error("Invalid z.ai quota limit nextResetTime");
   return date.toISOString();
 }
 
 function availability(windows: CapacityWindow[]): Availability {
   if (windows.length === 0) return "unknown";
-  if (
-    windows.some(
-      (window) => window.usedPercent !== null && window.usedPercent < 100,
-    )
-  ) {
-    return "yes";
-  }
-  return "no";
+  return windows.some((window) => window.usedPercent !== null && window.usedPercent < 100)
+    ? "yes"
+    : "no";
 }
 
 function parseLimit(value: unknown, index: number): CapacityWindow | null {
@@ -120,29 +84,25 @@ function parseLimit(value: unknown, index: number): CapacityWindow | null {
     throw new Error("Invalid z.ai quota limit entry");
   }
   const serverPercent = optionalNumber(limit.percentage, "limit percentage");
-  if (serverPercent === null)
-    throw new Error("Invalid z.ai quota limit percentage");
+  if (serverPercent === null) throw new Error("Invalid z.ai quota limit percentage");
   if (!isCapacityLimitType(limit.type)) return null;
 
   const total = optionalNumber(limit.usage, "limit usage");
   const current = optionalNumber(limit.currentValue, "limit currentValue");
   const remaining = optionalNumber(limit.remaining, "limit remaining");
-  let consumed =
-    total !== null && remaining !== null ? total - remaining : null;
+  let consumed = total !== null && remaining !== null ? total - remaining : null;
   if (current !== null) consumed = Math.max(consumed ?? current, current);
-
   let usedPercent = serverPercent;
   if (total !== null && total > 0 && consumed !== null) {
     usedPercent = (consumed / total) * 100;
   }
   const metadata = LIMIT_METADATA[limit.type];
-
   return {
     id: `zai:${metadata.id}:${index + 1}`,
     label: windowLabel(limit.type, unit, count),
     limitType: limit.type,
     durationMinutes: durationMinutes(unit, count, limit.type),
-    usedPercent: clampPercent(usedPercent),
+    usedPercent: Math.max(0, Math.min(100, usedPercent)),
     resetsAt: resetTime(limit.nextResetTime),
     total,
     current,
@@ -150,55 +110,16 @@ function parseLimit(value: unknown, index: number): CapacityWindow | null {
   };
 }
 
-export async function resolveZaiApiKey(
-  options: ZaiCredentialOptions = {},
-): Promise<string> {
-  const env = options.env ?? process.env;
-  const environmentKey = nonEmptyText(env.Z_AI_API_KEY);
-  if (environmentKey) return environmentKey;
-
-  const authPath = join(env.HOME || homedir(), ".pi", "agent", "auth.json");
-  let parsed: UnknownRecord;
-  try {
-    const contents = await (options.readFile ?? readFile)(authPath, "utf8");
-    const root = record(JSON.parse(contents));
-    if (!root) throw new Error("invalid");
-    parsed = root;
-  } catch (error) {
-    if (record(error)?.code === "ENOENT") {
-      throw new Error(MISSING_API_KEY_MESSAGE);
-    }
-    throw new Error("z.ai Pi auth file is malformed");
-  }
-
-  const credential = record(parsed.zai);
-  if (!credential) {
-    throw new Error(MISSING_API_KEY_MESSAGE);
-  }
-  const key =
-    credential.type === "api_key" ? nonEmptyText(credential.key) : null;
-  if (!key) throw new Error("z.ai Pi auth file has a malformed zai credential");
-  return key;
-}
-
-export function parseZaiQuota(raw: unknown, checkedAt: string): CapacityReport {
+export function parseZaiCapacity(raw: unknown): ProviderCapacitySnapshot {
   const root = record(raw);
   const data = record(root?.data);
-  if (
-    root?.success !== true ||
-    root.code !== 200 ||
-    !data ||
-    !Array.isArray(data.limits)
-  ) {
+  if (root?.success !== true || root.code !== 200 || !data || !Array.isArray(data.limits)) {
     throw new Error("Invalid z.ai quota response");
   }
   const windows = data.limits
     .map(parseLimit)
     .filter((window): window is CapacityWindow => window !== null);
   return {
-    harness: "pi",
-    provider: "zai",
-    generatedAt: checkedAt,
     authenticated: true,
     available: availability(windows),
     windows,
@@ -206,32 +127,29 @@ export function parseZaiQuota(raw: unknown, checkedAt: string): CapacityReport {
   };
 }
 
-export async function probeZaiCapacity(
-  options: ZaiProbeOptions,
-): Promise<CapacityReport> {
-  const key = await resolveZaiApiKey(options);
-  const controller = new AbortController();
-  const timer = setTimeout(() => controller.abort(), options.timeoutMs ?? 5000);
+export async function fetchZaiCapacity(
+  key: string,
+  options: ZaiRequestOptions = {},
+): Promise<ProviderCapacitySnapshot> {
   let response: Response;
   try {
     response = await (options.fetch ?? globalThis.fetch)(ZAI_QUOTA_URL, {
       method: "GET",
       headers: { Authorization: `Bearer ${key}` },
-      signal: controller.signal,
+      signal: AbortSignal.timeout(options.timeoutMs ?? 5000),
     });
   } catch {
     throw new Error("z.ai quota request failed");
-  } finally {
-    clearTimeout(timer);
   }
   if (response.status !== 200) {
     throw new Error(`z.ai quota request failed: HTTP ${response.status}`);
   }
-  let raw: unknown;
   try {
-    raw = await response.json();
-  } catch {
-    throw new Error("z.ai quota response is not valid JSON");
+    return parseZaiCapacity(await response.json());
+  } catch (error) {
+    if (error instanceof SyntaxError) {
+      throw new Error("z.ai quota response is not valid JSON");
+    }
+    throw error;
   }
-  return parseZaiQuota(raw, options.checkedAt);
 }

package/src/capacity/sources/claude.ts

@@ -0,0 +1,30 @@
+import {
+  fetchAnthropicCapacity,
+  type AnthropicCapacityRequestOptions,
+} from "../providers/anthropic.js";
+import type { CapacityReport, ProviderCapacitySnapshot } from "../types.js";
+import {
+  resolveClaudeAnthropicToken,
+  type ClaudeCredentialOptions,
+} from "../../harnesses/claude/credentials.js";
+
+export type ClaudeCapacityOptions = ClaudeCredentialOptions & AnthropicCapacityRequestOptions;
+
+export type ClaudeProbeOptions = ClaudeCapacityOptions & {
+  checkedAt: string;
+};
+
+function claudeReport(snapshot: ProviderCapacitySnapshot, checkedAt: string): CapacityReport {
+  return {
+    harness: "claude",
+    provider: "anthropic",
+    generatedAt: checkedAt,
+    ...snapshot,
+  };
+}
+
+export async function probeClaudeCapacity(options: ClaudeProbeOptions): Promise<CapacityReport> {
+  const token = await resolveClaudeAnthropicToken(options);
+  const snapshot = await fetchAnthropicCapacity(token, options);
+  return claudeReport(snapshot, options.checkedAt);
+}

package/src/capacity/{codex.ts → sources/codex.ts}

@@ -1,10 +1,16 @@
 import { spawn } from "node:child_process";
-import { readFile } from "node:fs/promises";
-import { join } from "node:path";
-import { parseWhamUsage, resetTime, safeIdentifier } from "./wham.js";
-import type { CapacityReport, CapacityWindow } from "./types.js";
-
-export { toRateWindow } from "./wham.js";
+import {
+  fetchOpenAiPatAccountId,
+  fetchOpenAiSubscriptionCapacity,
+  parseWhamUsage,
+  resetTime,
+  safeIdentifier,
+} from "../providers/openai-subscription.js";
+import type { CapacityReport, CapacityWindow } from "../types.js";
+import {
+  resolveCodexCredentials,
+  type CodexCredentialOptions,
+} from "../../harnesses/codex/credentials.js";
 
 type CodexUsageSource = "pat" | "oauth" | "cli";
 type UsageSnapshot = {
@@ -18,19 +24,16 @@ type RpcMessage = { id?: number; method: string; params?: UnknownRecord };
 type CliResponses = { rateLimits: unknown; account: unknown };
 type CodexRpc = (messages: RpcMessage[]) => Promise<CliResponses>;
 
-export const CODEX_APP_SERVER_ARGS = ["-s", "read-only", "-a", "untrusted", "app-server"] as const;
-
-type CodexProbeOptions = {
+export type CodexProbeOptions = CodexCredentialOptions & {
   installed: boolean;
   checkedAt: string;
-  readFile?: (path: string, encoding: BufferEncoding) => Promise<string>;
   fetch?: typeof globalThis.fetch;
   rpc?: CodexRpc;
   timeoutMs?: number;
-  env?: NodeJS.ProcessEnv;
-  now?: () => Date;
 };
 
+export const CODEX_APP_SERVER_ARGS = ["-s", "read-only", "-a", "untrusted", "app-server"] as const;
+
 function record(value: unknown): UnknownRecord | null {
   return value !== null && typeof value === "object" && !Array.isArray(value)
     ? (value as UnknownRecord)
@@ -45,11 +48,6 @@ function nonEmptyText(value: unknown): string | null {
   return typeof value === "string" && value.length > 0 ? value : null;
 }
 
-export function resolveCodexAuthPath(env: NodeJS.ProcessEnv = process.env): string {
-  const root = env.CODEX_HOME || join(env.HOME || "", ".codex");
-  return join(root, "auth.json");
-}
-
 export function parseUsage(raw: unknown, source: "pat" | "oauth"): UsageSnapshot {
   const { windows, creditsRemaining } = parseWhamUsage(raw);
   return { windows, creditsRemaining, source };
@@ -77,14 +75,15 @@ function cliSnapshotWindows(value: unknown, fallbackId: string): CapacityWindow[
   ].filter((item): item is CapacityWindow => item !== null);
 }
 
-export function parseCliUsage(raw: unknown): UsageSnapshot {
+function parseCliUsage(raw: unknown): UsageSnapshot {
   const response = record(raw) ?? {};
   const primary = record(response.rateLimits);
   const windows = primary ? cliSnapshotWindows(primary, "codex") : [];
   const buckets = record(response.rateLimitsByLimitId);
   if (buckets) {
-    for (const [id, snapshot] of Object.entries(buckets))
+    for (const [id, snapshot] of Object.entries(buckets)) {
       windows.push(...cliSnapshotWindows(snapshot, id));
+    }
   }
   const unique = [...new Map(windows.map((window) => [window.id, window])).values()];
   return { windows: unique, creditsRemaining: null, source: "cli" };
@@ -92,7 +91,7 @@ export function parseCliUsage(raw: unknown): UsageSnapshot {
 
 function capacityFromSnapshot(
   snapshot: UsageSnapshot,
-  context: CodexProbeOptions,
+  options: CodexProbeOptions,
   raw?: unknown,
 ): CapacityReport {
   const hasUsage = snapshot.windows.some((window) => window.usedPercent !== null);
@@ -103,7 +102,7 @@ function capacityFromSnapshot(
   return {
     harness: "codex",
     provider: "openai",
-    generatedAt: context.checkedAt,
+    generatedAt: options.checkedAt,
     authenticated: true,
     available: reached ? "no" : hasUsage ? "yes" : "unknown",
     windows: snapshot.windows,
@@ -111,70 +110,9 @@ function capacityFromSnapshot(
   };
 }
 
-function jwtExpiry(token: string): number | null {
-  const part = token.split(".")[1];
-  if (!part) return null;
-  try {
-    return finiteNumber(record(JSON.parse(Buffer.from(part, "base64url").toString("utf8")))?.exp);
-  } catch {
-    return null;
-  }
-}
-
-function staleOAuth(tokens: UnknownRecord, token: string, now: Date): boolean {
-  const metadata = tokens.expires_at ?? tokens.expiresAt ?? tokens.expiry;
-  let expiry: number | null = finiteNumber(metadata);
-  if (typeof metadata === "string") {
-    const parsed = Date.parse(metadata);
-    expiry = Number.isNaN(parsed) ? null : parsed / 1000;
-  }
-  expiry ??= jwtExpiry(token);
-  return expiry !== null && expiry <= now.getTime() / 1000;
-}
-
-async function fetchJson(
-  fetcher: typeof globalThis.fetch,
-  url: string,
-  init: RequestInit,
-  timeoutMs: number,
-): Promise<unknown> {
-  const controller = new AbortController();
-  const timer = setTimeout(() => controller.abort(), timeoutMs);
-  try {
-    const response = await fetcher(url, { ...init, signal: controller.signal });
-    if (!response.ok) throw new Error(response.status === 401 ? "unauthorized" : "request failed");
-    return await response.json();
-  } finally {
-    clearTimeout(timer);
-  }
-}
-
-async function apiSnapshot(
-  token: string,
-  accountId: string,
-  source: "pat" | "oauth",
-  options: CodexProbeOptions,
-): Promise<UsageSnapshot> {
-  const fetcher = options.fetch ?? globalThis.fetch;
-  const raw = await fetchJson(
-    fetcher,
-    "https://chatgpt.com/backend-api/wham/usage",
-    {
-      headers: {
-        Authorization: `Bearer ${token}`,
-        "ChatGPT-Account-Id": accountId,
-      },
-    },
-    options.timeoutMs ?? 5000,
-  );
-  return parseUsage(raw, source);
-}
-
 function appServerRpc(messages: RpcMessage[], timeoutMs = 5000): Promise<CliResponses> {
   return new Promise((resolve, reject) => {
-    const child = spawn("codex", CODEX_APP_SERVER_ARGS, {
-      stdio: ["pipe", "pipe", "ignore"],
-    });
+    const child = spawn("codex", CODEX_APP_SERVER_ARGS, { stdio: ["pipe", "pipe", "ignore"] });
     const results: Partial<CliResponses> = {};
     let buffer = "";
     let settled = false;
@@ -235,12 +173,8 @@ export function codexUnavailableReport(checkedAt: string): CapacityReport {
   };
 }
 
-function unavailable(options: CodexProbeOptions): CapacityReport {
-  return codexUnavailableReport(options.checkedAt);
-}
-
 async function cliFallback(options: CodexProbeOptions): Promise<CapacityReport> {
-  if (!options.installed) return unavailable(options);
+  if (!options.installed) return codexUnavailableReport(options.checkedAt);
   const messages: RpcMessage[] = [
     {
       id: 1,
@@ -273,57 +207,53 @@ async function cliFallback(options: CodexProbeOptions): Promise<CapacityReport>
     }
     return result;
   } catch {
-    return unavailable(options);
+    return codexUnavailableReport(options.checkedAt);
   }
 }
 
+async function subscriptionSnapshot(
+  access: string,
+  accountId: string,
+  source: "pat" | "oauth",
+  options: CodexProbeOptions,
+): Promise<UsageSnapshot> {
+  const snapshot = await fetchOpenAiSubscriptionCapacity({ access, accountId }, options);
+  if (snapshot.authenticated === false) throw new Error("unauthorized");
+  return {
+    windows: snapshot.windows,
+    creditsRemaining: snapshot.creditsRemaining,
+    source,
+  };
+}
+
 export async function probeCodexCapacity(options: CodexProbeOptions): Promise<CapacityReport> {
-  let parsed: UnknownRecord | null = null;
-  try {
-    const contents = await (options.readFile ?? readFile)(
-      resolveCodexAuthPath(options.env),
-      "utf8",
-    );
-    parsed = record(JSON.parse(contents));
-  } catch {
-    return cliFallback(options);
-  }
+  const credentials = await resolveCodexCredentials(options);
+  if (!credentials) return cliFallback(options);
 
-  const auth = parsed ?? {};
-  const pat = nonEmptyText(auth.personal_access_token);
-  if (pat) {
+  if (credentials.personalAccessToken) {
     try {
-      const fetcher = options.fetch ?? globalThis.fetch;
-      const whoami = record(
-        await fetchJson(
-          fetcher,
-          "https://auth.openai.com/api/accounts/v1/user-auth-credential/whoami",
-          { headers: { Authorization: `Bearer ${pat}` } },
-          options.timeoutMs ?? 5000,
-        ),
+      const accountId = await fetchOpenAiPatAccountId(credentials.personalAccessToken, options);
+      const snapshot = await subscriptionSnapshot(
+        credentials.personalAccessToken,
+        accountId,
+        "pat",
+        options,
       );
-      const accountId = nonEmptyText(whoami?.chatgpt_account_id);
-      if (!accountId) throw new Error("account unavailable");
-      return capacityFromSnapshot(await apiSnapshot(pat, accountId, "pat", options), options);
+      return capacityFromSnapshot(snapshot, options);
     } catch {
       // Continue to a separately available OAuth credential before using the CLI.
     }
   }
 
-  const tokens = record(auth.tokens);
-  const accessToken = nonEmptyText(tokens?.access_token);
-  const accountId = nonEmptyText(tokens?.account_id);
-  if (
-    tokens &&
-    accessToken &&
-    accountId &&
-    !staleOAuth(tokens, accessToken, (options.now ?? (() => new Date()))())
-  ) {
+  if (credentials.oauth) {
     try {
-      return capacityFromSnapshot(
-        await apiSnapshot(accessToken, accountId, "oauth", options),
+      const snapshot = await subscriptionSnapshot(
+        credentials.oauth.access,
+        credentials.oauth.accountId,
+        "oauth",
         options,
       );
+      return capacityFromSnapshot(snapshot, options);
     } catch {
       return cliFallback(options);
     }

package/src/capacity/sources/pi.ts

@@ -0,0 +1,131 @@
+import { fetchAnthropicCapacity } from "../providers/anthropic.js";
+import { fetchOpenAiPlatformCapacity } from "../providers/openai-platform.js";
+import { fetchOpenAiSubscriptionCapacity } from "../providers/openai-subscription.js";
+import { fetchZaiCapacity } from "../providers/zai.js";
+import type { CapacityReport, ProviderCapacitySnapshot } from "../types.js";
+import {
+  resolvePiAnthropicCredential,
+  resolvePiOpenAiCredential,
+  resolvePiZaiApiKey,
+  type PiCredentialOptions,
+  type PiOpenAiCredential,
+} from "../../harnesses/pi/credentials.js";
+
+export type PiAnthropicCapacityOptions = PiCredentialOptions & {
+  now?: () => Date;
+  fetch?: typeof globalThis.fetch;
+  timeoutMs?: number;
+};
+
+export type PiAnthropicProbeOptions = PiAnthropicCapacityOptions & {
+  checkedAt: string;
+};
+
+export type PiOpenAiCapacityOptions = PiCredentialOptions & {
+  now?: () => Date;
+  fetch?: typeof globalThis.fetch;
+  timeoutMs?: number;
+};
+
+export type PiOpenAiProbeOptions = PiOpenAiCapacityOptions & {
+  checkedAt: string;
+};
+
+export type PiZaiCapacityOptions = PiCredentialOptions & {
+  now?: () => Date;
+  fetch?: typeof globalThis.fetch;
+  timeoutMs?: number;
+};
+
+export type PiZaiProbeOptions = PiZaiCapacityOptions & {
+  checkedAt: string;
+};
+
+export async function probePiAnthropicCapacity(
+  options: PiAnthropicProbeOptions,
+): Promise<CapacityReport> {
+  const credential = await resolvePiAnthropicCredential(options);
+  const now = options.now?.() ?? new Date();
+  if (credential.expiresMs !== null && credential.expiresMs <= now.getTime()) {
+    return {
+      harness: "pi",
+      provider: "anthropic",
+      generatedAt: options.checkedAt,
+      authenticated: true,
+      available: "unknown",
+      windows: [],
+      creditsRemaining: null,
+    };
+  }
+  const snapshot = await fetchAnthropicCapacity(credential.access, options);
+  return {
+    harness: "pi",
+    provider: "anthropic",
+    generatedAt: options.checkedAt,
+    ...snapshot,
+  };
+}
+
+function piOpenAiReport(snapshot: ProviderCapacitySnapshot, checkedAt: string): CapacityReport {
+  return {
+    harness: "pi",
+    provider: "openai",
+    generatedAt: checkedAt,
+    ...snapshot,
+  };
+}
+
+function jwtExpiryMs(token: string): number | null {
+  const part = token.split(".")[1];
+  if (!part) return null;
+  try {
+    const payload = JSON.parse(Buffer.from(part, "base64url").toString("utf8")) as unknown;
+    if (!payload || typeof payload !== "object" || Array.isArray(payload)) return null;
+    const exp = (payload as Record<string, unknown>).exp;
+    return typeof exp === "number" && Number.isFinite(exp) ? exp * 1000 : null;
+  } catch {
+    return null;
+  }
+}
+
+function oauthExpired(
+  credential: Extract<PiOpenAiCredential, { kind: "oauth" }>,
+  nowMs: number,
+): boolean {
+  const expiryMs = credential.expiresMs ?? jwtExpiryMs(credential.access);
+  return expiryMs !== null && expiryMs <= nowMs;
+}
+
+export async function probeOpenAiCapacity(options: PiOpenAiProbeOptions): Promise<CapacityReport> {
+  const credential = await resolvePiOpenAiCredential(options);
+  if (credential.kind === "platform") {
+    const snapshot = await fetchOpenAiPlatformCapacity(credential.key, options);
+    return piOpenAiReport(snapshot, options.checkedAt);
+  }
+
+  const nowMs = Date.parse(options.checkedAt);
+  if (Number.isNaN(nowMs) || oauthExpired(credential, nowMs)) {
+    return piOpenAiReport(
+      {
+        authenticated: true,
+        available: "unknown",
+        windows: [],
+        creditsRemaining: null,
+      },
+      options.checkedAt,
+    );
+  }
+  const snapshot = await fetchOpenAiSubscriptionCapacity(credential, options);
+  return piOpenAiReport(snapshot, options.checkedAt);
+}
+
+export async function probeZaiCapacity(options: PiZaiProbeOptions): Promise<CapacityReport> {
+  const key = await resolvePiZaiApiKey(options);
+  const snapshot = await fetchZaiCapacity(key, options);
+  return {
+    harness: "pi",
+    provider: "zai",
+    generatedAt: options.checkedAt,
+    ...snapshot,
+  };
+}

package/src/capacity/types.ts

@@ -24,3 +24,6 @@ export interface CapacityReport {
   windows: CapacityWindow[];
   creditsRemaining: number | null;
 }
+
+/** Provider quota data before a harness/source identity is attached. */
+export type ProviderCapacitySnapshot = Omit<CapacityReport, "harness" | "provider" | "generatedAt">;