Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
12 kB · 333 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334import { spawn } from "node:child_process";import { readFileSync } from "node:fs";import { join } from "node:path";import { collectUsage, type FetchLike, type UsageCommandRequest, type UsageCommandRunner, type UsageCredential, type UsageLimit, type UsageReport, type UsageResponse, type UsageStatus, usageFeed,} from "@earendil-works/pi-ai";import { getAgentDir } from "../../config.js";import type { AuthCredential, AuthStorage } from "../../core/auth-storage.js";import { credentialIdentity } from "../../core/auth-storage.js";import type { ContextTreeNode } from "../../core/context-tree.js";import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "../../core/provider-display-names.js";import { formatContextTree } from "./components/context-tree-format.js";import { theme } from "./theme/theme.js";
const BAR_WIDTH = 24;
export type { FetchLike, UsageLimit, UsageReport };
export interface SkippedProvider { provider: string; reason: string;}
export interface FetchUsageResult { reports: UsageReport[]; skipped: SkippedProvider[];}
/** * A console session captured out of band (alibaba's quota only answers to a browser * session, because its API keys are inference-only). Providers whose feed accepts a * `session` credential read one of these when it exists; without it they are reported as * skipped, with the reason, rather than as a provider with no limits. */export function sessionPathFor(provider: string): string { return join(getAgentDir(), `${provider}-session.json`);}
function readSession(provider: string): Record<string, unknown> | undefined { try { const parsed = JSON.parse(readFileSync(sessionPathFor(provider), "utf8")); return parsed && typeof parsed === "object" ? (parsed as Record<string, unknown>) : undefined; } catch { return undefined; }}
function accountFields(credential: AuthCredential): { accountId?: string; email?: string } { const accountId = (credential as { accountId?: unknown }).accountId; const email = (credential as { email?: unknown }).email; return { ...(typeof accountId === "string" ? { accountId } : {}), ...(typeof email === "string" ? { email } : {}), };}
/** Resolve every stored credential into the payload its provider's feed accepts. */async function credentialEntries( authStorage: AuthStorage,): Promise<{ entries: { provider: string; credential: UsageCredential }[]; unavailable: UsageReport[] }> { const entries: { provider: string; credential: UsageCredential }[] = []; const unavailable: UsageReport[] = [];
for (const provider of authStorage.list()) { const displayName = BUILT_IN_PROVIDER_DISPLAY_NAMES[provider] ?? provider; for (const credential of authStorage.getCredentialList(provider)) { const account = accountFields(credential); const identity = credentialIdentity(credential);
if (credential.type === "oauth") { let resolved: Awaited<ReturnType<AuthStorage["getApiKeyForCredential"]>>; try { resolved = await authStorage.getApiKeyForCredential(provider, identity); } catch { resolved = undefined; } const resolvedCredential = resolved?.credential; const access = resolvedCredential?.type === "oauth" ? resolvedCredential.access : credential.access; if (typeof access !== "string" || access === "") { unavailable.push({ provider, providerName: displayName, fetchedAt: Date.now(), ...account, limits: [], error: "OAuth token unavailable; run /login", }); continue; } entries.push({ provider, credential: { kind: "oauth", access, ...(typeof credential.refresh === "string" ? { refresh: credential.refresh } : {}), ...account, extra: { ...credential }, }, }); continue; }
let apiKey: string | undefined; try { apiKey = (await authStorage.getApiKeyForCredential(provider, identity))?.apiKey; } catch { apiKey = undefined; } if (apiKey) entries.push({ provider, credential: { kind: "apiKey", apiKey, ...account } }); } }
// Providers whose feed accepts a session get one when it exists, and providers whose feed accepts a // CLI get the chance to use it. Both come from the feed's declared credential kinds, so the bridge // never assumes which provider wants what. for (const provider of authStorage.list()) { const feed = usageFeed(provider); if (!feed) continue; if (feed.credentialKinds.includes("session")) { const session = readSession(provider); if (session) entries.push({ provider, credential: { kind: "session", session } }); } if (feed.credentialKinds.includes("cli") && !entries.some((entry) => entry.provider === provider)) { const region = process.env.PRIME_USAGE_CLI_REGION; entries.push({ provider, credential: { kind: "cli", ...(region ? { config: { region } } : {}) } }); } }
return { entries, unavailable };}
/** * Commands a feed may ask the host to run. A feed is trusted code (a shipped provider or a * registered extension), but the *name* of the binary is still data crossing a boundary, so the host * runs only what it recognises rather than whatever a feed names. */const ALLOWED_USAGE_COMMANDS = new Set(["bl"]);
const COMMAND_OUTPUT_LIMIT_BYTES = 64 * 1024;const DEFAULT_COMMAND_TIMEOUT_MS = 15_000;
/** Environment a provider CLI legitimately needs: PATH resolution, config discovery, locale, proxies. */const COMMAND_ENV_ALLOWLIST = [ "PATH", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "TZ", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "all_proxy", "no_proxy",];
function commandEnvironment(): NodeJS.ProcessEnv { const environment: NodeJS.ProcessEnv = {}; for (const key of COMMAND_ENV_ALLOWLIST) { const value = process.env[key]; if (value !== undefined) environment[key] = value; } return environment;}
/** * Answer a feed's `command` request by running it. Exit status is the response status and stdout is * the body, matching what an HTTP response looks like to a feed — so a provider can treat "the CLI is * missing" and "the CLI failed" as data and route around them. */export const runUsageCommand: UsageCommandRunner = (request: UsageCommandRequest): Promise<UsageResponse> => new Promise((resolve) => { if (!ALLOWED_USAGE_COMMANDS.has(request.command)) { resolve({ status: 126, body: `the host will not run ${request.command}` }); return; } const child = spawn(request.command, request.args, { env: commandEnvironment(), stdio: ["pipe", "pipe", "pipe"], }); let stdout = ""; let settled = false; const finish = (response: UsageResponse) => { if (settled) return; settled = true; clearTimeout(timer); resolve(response); }; const timer = setTimeout(() => { child.kill("SIGKILL"); finish({ status: 124, body: `${request.command} timed out` }); }, request.timeoutMs ?? DEFAULT_COMMAND_TIMEOUT_MS);
child.stdout?.on("data", (chunk: Buffer) => { if (stdout.length < COMMAND_OUTPUT_LIMIT_BYTES) stdout += chunk.toString("utf8"); }); if (request.stdin) child.stdin?.end(request.stdin); else child.stdin?.end(); child.on("error", (error) => finish({ status: 127, body: error.message })); child.on("close", (code) => finish({ status: code ?? 1, body: stdout })); });
/** Fetch usage for every stored credential, across every provider with a feed. */export async function fetchUsageReports( authStorage: AuthStorage, options: { fetchImpl?: FetchLike; signal?: AbortSignal; nowMs?: number } = {},): Promise<FetchUsageResult> { const { entries, unavailable } = await credentialEntries(authStorage); // The host runs commands for feeds that prefer a provider's own CLI: `bl` for alibaba's token plan, // which is a supported interface and needs no captured browser session. const collected = await collectUsage(entries, { runCommand: runUsageCommand, ...options }); return { reports: [...collected.reports, ...unavailable], skipped: collected.skipped, };}
function accountLabel(report: UsageReport, index: number): string { if (report.email) return report.email; if (report.accountId) return report.accountId; return `account ${index + 1}`;}
function formatDuration(ms: number): string { const seconds = Math.max(0, Math.round(ms / 1000)); if (seconds < 60) return `${seconds}s`; const minutes = Math.round(seconds / 60); if (minutes < 60) return `${minutes}m`; const hours = Math.round(minutes / 60); if (hours < 48) return `${hours}h`; return `${Math.round(hours / 24)}d`;}
function renderBar(usedPercent: number | undefined): string { if (usedPercent === undefined) return `[${"·".repeat(BAR_WIDTH)}] unknown`; const fraction = Math.max(0, Math.min(1, usedPercent / 100)); const filled = Math.round(fraction * BAR_WIDTH); return `[${"█".repeat(filled)}${"░".repeat(BAR_WIDTH - filled)}] ${Math.round(usedPercent)}%`;}
/** Exhaustion is a colour, not a number: the bar is already the number. */function statusColour(status: UsageStatus): "error" | "warning" | "dim" | "text" { if (status === "exhausted") return "error"; if (status === "warning") return "warning"; if (status === "unknown") return "dim"; return "text";}
function renderLimit(limit: UsageLimit, nowMs: number): string[] { const used = limit.usedPercent === undefined ? "no percentage reported" : `${limit.usedPercent.toFixed(1)}% used`; const reset = limit.resetsAt && limit.resetsAt > nowMs ? ` · resets in ${formatDuration(limit.resetsAt - nowMs)}` : ""; const detail = limit.detail ? ` · ${limit.detail}` : ""; const head = ` ${limit.label}${limit.windowLabel ? ` (${limit.windowLabel})` : ""}: ${used}${reset}${detail}`; const bar = limit.usedPercent === undefined ? "" : `\n ${renderBar(limit.usedPercent)}`; return [theme.fg(statusColour(limit.status), head + bar)];}
function formatProviderHeading(report: UsageReport): string { return report.providerName ?? BUILT_IN_PROVIDER_DISPLAY_NAMES[report.provider] ?? report.provider;}
function formatSkippedProviderName(provider: string): string { return BUILT_IN_PROVIDER_DISPLAY_NAMES[provider] ?? provider;}
export function renderUsageReportText( result: { reports: readonly UsageReport[]; skipped?: readonly SkippedProvider[] }, contextTree: ContextTreeNode | undefined, width: number, options: { loading?: boolean; fetchedAt?: number; includeSession?: boolean } = {},): string { const reports = result.reports; const skipped = result.skipped ?? []; const nowMs = Date.now(); const lines = [theme.bold(theme.fg("accent", "Usage"))]; if (options.fetchedAt) lines[0] += theme.fg("dim", ` (${formatDuration(nowMs - options.fetchedAt)} ago)`); if (options.loading) lines.push(theme.fg("dim", "Refreshing provider quotas in the background…"));
let currentProvider: string | undefined; reports.forEach((report, index) => { if (report.provider !== currentProvider) { currentProvider = report.provider; lines.push(""); lines.push(theme.bold(theme.fg("accent", formatProviderHeading(report)))); } lines.push(""); lines.push( ` ${theme.fg("accent", accountLabel(report, index))}${report.planType ? ` · ${report.planType}` : ""}`, ); if (report.error) { lines.push(` ${theme.fg("warning", `Unable to fetch quota: ${report.error}`)}`); return; } for (const limit of report.limits) lines.push(...renderLimit(limit, nowMs)); if (report.limits.length === 0 && !report.balance) lines.push(theme.fg("dim", " Nothing reported")); if (report.balance) lines.push(` ${theme.fg("accent", "Balance")}: ${report.balance}`); for (const note of report.notes ?? []) lines.push(theme.fg("dim", ` ${note}`)); });
if (skipped.length > 0) { lines.push(""); lines.push(theme.fg("dim", "No quota feed")); for (const entry of skipped) { lines.push(theme.fg("dim", ` ${formatSkippedProviderName(entry.provider)}: ${entry.reason}`)); } } if (reports.length === 0 && skipped.length === 0) { lines.push(theme.fg("dim", "No provider accounts to report.")); }
if (options.includeSession !== false) { if (contextTree) { lines.push("", theme.bold(theme.fg("accent", "This session"))); lines.push(formatContextTree(contextTree, Math.max(60, width))); } else { lines.push("", theme.fg("dim", "Session usage will appear when a context snapshot is available.")); } } return lines.join("\n");}