From db6ef39d5e019e1356ad22c250dde9200236785a Mon Sep 17 00:00:00 2001 From: stcl-bsm <2659616637@qq.com> Date: Sun, 13 Sep 2026 18:54:10 +0800 Subject: [PATCH] Migrate peerDependencies and imports from @mariozechner/* to @earendil-works/* --- extensions/usage/core.ts | 1394 ++++++++++++++++++------------------- extensions/usage/index.ts | 592 ++++++++-------- package.json | 78 +-- 3 files changed, 1032 insertions(+), 1032 deletions(-) diff --git a/extensions/usage/core.ts b/extensions/usage/core.ts index 5a0def0..ee29288 100644 --- a/extensions/usage/core.ts +++ b/extensions/usage/core.ts @@ -1,697 +1,697 @@ -import * as fs from "node:fs"; -import * as path from "node:path"; -import * as os from "node:os"; - -export type ProviderKey = "codex" | "claude" | "zai" | "gemini" | "antigravity"; -export type OAuthProviderId = "openai-codex" | "anthropic" | "google-gemini-cli" | "google-antigravity"; - -export interface AuthData { - "openai-codex"?: { access?: string; refresh?: string; expires?: number }; - anthropic?: { access?: string; refresh?: string; expires?: number }; - zai?: { key?: string; access?: string; refresh?: string; expires?: number }; - "google-gemini-cli"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; - "google-antigravity"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; -} - -export interface UsageData { - session: number; - weekly: number; - sessionResetsIn?: string; - weeklyResetsIn?: string; - extraSpend?: number; - extraLimit?: number; - error?: string; -} - -export type UsageByProvider = Record; - -export interface UsageEndpoints { - zai: string; - gemini: string; - antigravity: string; - googleLoadCodeAssistEndpoints: string[]; -} - -export interface FetchResponseLike { - ok: boolean; - status: number; - json(): Promise; -} - -export type FetchLike = (input: string, init?: RequestInit) => Promise; - -export interface RequestConfig { - fetchFn?: FetchLike; - timeoutMs?: number; -} - -export interface FetchConfig extends RequestConfig { - endpoints?: UsageEndpoints; - env?: NodeJS.ProcessEnv; -} - -export interface OAuthApiKeyResult { - newCredentials: Record; - apiKey: string; -} - -export type OAuthApiKeyResolver = ( - providerId: OAuthProviderId, - credentials: Record>, -) => Promise; - -export interface EnsureFreshAuthConfig { - auth?: AuthData | null; - authFile?: string; - oauthResolver?: OAuthApiKeyResolver; - nowMs?: number; - persist?: boolean; -} - -export interface FreshAuthResult { - auth: AuthData | null; - changed: boolean; - refreshErrors: Partial>; -} - -export interface FetchAllUsagesConfig extends FetchConfig, EnsureFreshAuthConfig { - auth?: AuthData | null; - authFile?: string; -} - -const DEFAULT_FETCH_TIMEOUT_MS = 12_000; -const TOKEN_REFRESH_SKEW_MS = 60_000; - -export const DEFAULT_AUTH_FILE = path.join(os.homedir(), ".pi", "agent", "auth.json"); -export const DEFAULT_ZAI_USAGE_ENDPOINT = "https://api.z.ai/api/monitor/usage/quota/limit"; -export const GOOGLE_QUOTA_ENDPOINT = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"; -export const GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS = [ - "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", - "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist", -]; - -export function resolveUsageEndpoints(): UsageEndpoints { - return { - zai: DEFAULT_ZAI_USAGE_ENDPOINT, - gemini: GOOGLE_QUOTA_ENDPOINT, - antigravity: GOOGLE_QUOTA_ENDPOINT, - googleLoadCodeAssistEndpoints: GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS, - }; -} - -function toErrorMessage(error: unknown): string { - if (error instanceof Error) { - if (error.name === "AbortError") return "request timeout"; - return error.message || String(error); - } - return String(error); -} - -function asObject(value: unknown): Record | null { - if (!value || typeof value !== "object") return null; - return value as Record; -} - -function normalizeUsagePair(session: number, weekly: number): { session: number; weekly: number } { - const clean = (v: number) => { - if (!Number.isFinite(v)) return 0; - return Number(v.toFixed(2)); - }; - return { session: clean(session), weekly: clean(weekly) }; -} - -async function requestJson(url: string, init: RequestInit, config: RequestConfig = {}): Promise<{ ok: true; data: any } | { ok: false; error: string }> { - const fetchFn = config.fetchFn ?? ((fetch as unknown) as FetchLike); - const timeoutMs = config.timeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS; - const controller = new AbortController(); - const timeout = timeoutMs > 0 ? setTimeout(() => controller.abort(), timeoutMs) : null; - - try { - const response = await fetchFn(url, { ...init, signal: controller.signal }); - if (!response.ok) return { ok: false, error: `HTTP ${response.status}` }; - - try { - const data = await response.json(); - return { ok: true, data }; - } catch { - return { ok: false, error: "invalid JSON response" }; - } - } catch (error) { - return { ok: false, error: toErrorMessage(error) }; - } finally { - if (timeout) clearTimeout(timeout); - } -} - -export function formatDuration(seconds: number): string { - if (!Number.isFinite(seconds) || seconds <= 0) return "now"; - const d = Math.floor(seconds / 86400); - const h = Math.floor((seconds % 86400) / 3600); - const m = Math.floor((seconds % 3600) / 60); - if (d > 0 && h > 0) return `${d}d ${h}h`; - if (d > 0) return `${d}d`; - if (h > 0 && m > 0) return `${h}h ${m}m`; - if (h > 0) return `${h}h`; - if (m > 0) return `${m}m`; - return "<1m"; -} - -export function formatResetsAt(isoDate: string, nowMs = Date.now()): string { - const resetTime = new Date(isoDate).getTime(); - if (!Number.isFinite(resetTime)) return ""; - const diffSeconds = Math.max(0, (resetTime - nowMs) / 1000); - return formatDuration(diffSeconds); -} - -export function readAuth(authFile = DEFAULT_AUTH_FILE): AuthData | null { - try { - const parsed = JSON.parse(fs.readFileSync(authFile, "utf-8")); - return asObject(parsed) as AuthData; - } catch { - return null; - } -} - -export function writeAuth(auth: AuthData, authFile = DEFAULT_AUTH_FILE): boolean { - try { - const dir = path.dirname(authFile); - if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true }); - - const tmpPath = `${authFile}.tmp-${process.pid}-${Date.now()}`; - fs.writeFileSync(tmpPath, JSON.stringify(auth, null, 2)); - fs.renameSync(tmpPath, authFile); - return true; - } catch { - return false; - } -} - -let cachedOAuthResolver: OAuthApiKeyResolver | null = null; - -async function getDefaultOAuthResolver(): Promise { - if (cachedOAuthResolver) return cachedOAuthResolver; - - const mod = await import("@mariozechner/pi-ai"); - if (typeof (mod as any).getOAuthApiKey !== "function") { - throw new Error("oauth resolver unavailable"); - } - - cachedOAuthResolver = (providerId, credentials) => - (mod as any).getOAuthApiKey(providerId, credentials) as Promise; - - return cachedOAuthResolver; -} - -function isCredentialExpired(creds: { expires?: number } | undefined, nowMs: number): boolean { - if (!creds) return false; - if (typeof creds.expires !== "number") return false; - return nowMs + TOKEN_REFRESH_SKEW_MS >= creds.expires; -} - -export async function ensureFreshAuthForProviders( - providerIds: OAuthProviderId[], - config: EnsureFreshAuthConfig = {}, -): Promise { - const authFile = config.authFile ?? DEFAULT_AUTH_FILE; - const auth = config.auth ?? readAuth(authFile); - if (!auth) { - return { auth: null, changed: false, refreshErrors: {} }; - } - - const nowMs = config.nowMs ?? Date.now(); - const uniqueProviders = Array.from(new Set(providerIds)); - const nextAuth: AuthData = { ...auth }; - const refreshErrors: Partial> = {}; - - let changed = false; - - for (const providerId of uniqueProviders) { - const creds = (nextAuth as any)[providerId] as { access?: string; refresh?: string; expires?: number } | undefined; - if (!creds?.refresh) continue; - - const needsRefresh = !creds.access || isCredentialExpired(creds, nowMs); - if (!needsRefresh) continue; - - try { - const resolver = config.oauthResolver ?? (await getDefaultOAuthResolver()); - const resolved = await resolver(providerId, nextAuth as any); - if (!resolved?.newCredentials) { - refreshErrors[providerId] = "missing OAuth credentials"; - continue; - } - - (nextAuth as any)[providerId] = { - ...(nextAuth as any)[providerId], - ...resolved.newCredentials, - }; - changed = true; - } catch (error) { - refreshErrors[providerId] = toErrorMessage(error); - } - } - - if (changed && config.persist !== false) { - writeAuth(nextAuth, authFile); - } - - return { auth: nextAuth, changed, refreshErrors }; -} - -export function readPercentCandidate(value: unknown): number | null { - if (typeof value !== "number" || !Number.isFinite(value)) return null; - - if (value >= 0 && value <= 1) { - if (Number.isInteger(value)) return value; - return value * 100; - } - - if (value >= 0 && value <= 100) return value; - return null; -} - -export function readLimitPercent(limit: any): number | null { - const direct = [ - limit?.percentage, - limit?.utilization, - limit?.used_percent, - limit?.usedPercent, - limit?.usagePercent, - limit?.usage_percent, - ] - .map(readPercentCandidate) - .find((v) => v != null); - - if (direct != null) return direct; - - const current = typeof limit?.currentValue === "number" ? limit.currentValue : null; - const remaining = typeof limit?.remaining === "number" ? limit.remaining : null; - - if (current != null && remaining != null && current + remaining > 0) { - return (current / (current + remaining)) * 100; - } - - return null; -} - -export function extractUsageFromPayload(data: any): { session: number; weekly: number } | null { - const limitArrays = [data?.data?.limits, data?.limits, data?.quota?.limits, data?.data?.quota?.limits]; - const limits = limitArrays.find((arr) => Array.isArray(arr)) as any[] | undefined; - - if (limits) { - const byType = (types: string[]) => - limits.find((l) => { - const t = String(l?.type || "").toUpperCase(); - return types.some((x) => t === x); - }); - - const sessionLimit = byType(["TIME_LIMIT", "SESSION_LIMIT", "REQUEST_LIMIT", "RPM_LIMIT", "RPD_LIMIT"]); - const weeklyLimit = byType(["TOKENS_LIMIT", "TOKEN_LIMIT", "WEEK_LIMIT", "WEEKLY_LIMIT", "TPM_LIMIT", "DAILY_LIMIT"]); - - const s = readLimitPercent(sessionLimit); - const w = readLimitPercent(weeklyLimit); - if (s != null && w != null) return normalizeUsagePair(s, w); - } - - const sessionCandidates = [ - data?.session, - data?.sessionPercent, - data?.session_percent, - data?.five_hour?.utilization, - data?.rate_limit?.primary_window?.used_percent, - data?.limits?.session?.utilization, - data?.usage?.session, - data?.data?.session, - data?.data?.sessionPercent, - data?.data?.session_percent, - data?.data?.usage?.session, - data?.quota?.session?.percentage, - data?.data?.quota?.session?.percentage, - ]; - - const weeklyCandidates = [ - data?.weekly, - data?.weeklyPercent, - data?.weekly_percent, - data?.seven_day?.utilization, - data?.rate_limit?.secondary_window?.used_percent, - data?.limits?.weekly?.utilization, - data?.usage?.weekly, - data?.data?.weekly, - data?.data?.weeklyPercent, - data?.data?.weekly_percent, - data?.data?.usage?.weekly, - data?.quota?.weekly?.percentage, - data?.data?.quota?.weekly?.percentage, - data?.quota?.daily?.percentage, - data?.data?.quota?.daily?.percentage, - ]; - - const session = sessionCandidates.map(readPercentCandidate).find((v) => v != null); - const weekly = weeklyCandidates.map(readPercentCandidate).find((v) => v != null); - - if (session == null || weekly == null) return null; - return normalizeUsagePair(session, weekly); -} - -export function googleMetadata(projectId?: string) { - return { - ideType: "IDE_UNSPECIFIED", - platform: "PLATFORM_UNSPECIFIED", - pluginType: "GEMINI", - ...(projectId ? { duetProject: projectId } : {}), - }; -} - -export function googleHeaders(token: string, projectId?: string) { - return { - Authorization: `Bearer ${token}`, - "Content-Type": "application/json", - "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1", - "X-Goog-Api-Client": "gl-node/22.17.0", - "Client-Metadata": JSON.stringify(googleMetadata(projectId)), - }; -} - -export async function discoverGoogleProjectId(token: string, config: FetchConfig = {}): Promise { - const env = config.env ?? process.env; - const envProjectId = env.GOOGLE_CLOUD_PROJECT || env.GOOGLE_CLOUD_PROJECT_ID; - if (envProjectId) return envProjectId; - - const endpoints = config.endpoints ?? resolveUsageEndpoints(); - - for (const endpoint of endpoints.googleLoadCodeAssistEndpoints) { - const result = await requestJson( - endpoint, - { - method: "POST", - headers: googleHeaders(token), - body: JSON.stringify({ metadata: googleMetadata() }), - }, - config, - ); - - if (!result.ok) continue; - - const data = result.data; - if (typeof data?.cloudaicompanionProject === "string" && data.cloudaicompanionProject) { - return data.cloudaicompanionProject; - } - if (data?.cloudaicompanionProject && typeof data.cloudaicompanionProject === "object") { - const id = data.cloudaicompanionProject.id; - if (typeof id === "string" && id) return id; - } - } - - return undefined; -} - -export function usedPercentFromRemainingFraction(value: unknown): number | null { - if (typeof value !== "number" || !Number.isFinite(value)) return null; - const remaining = Math.max(0, Math.min(1, value)); - return (1 - remaining) * 100; -} - -export function pickMostUsedBucket(buckets: any[]): any | null { - let best: any | null = null; - let bestUsed = -1; - for (const bucket of buckets) { - const used = usedPercentFromRemainingFraction(bucket?.remainingFraction); - if (used == null) continue; - if (used > bestUsed) { - bestUsed = used; - best = bucket; - } - } - return best; -} - -export function parseGoogleQuotaBuckets(data: any, provider: "gemini" | "antigravity"): { session: number; weekly: number } | null { - const allBuckets = Array.isArray(data?.buckets) ? data.buckets : []; - if (!allBuckets.length) return null; - - const requestBuckets = allBuckets.filter((b: any) => String(b?.tokenType || "").toUpperCase() === "REQUESTS"); - const buckets = requestBuckets.length ? requestBuckets : allBuckets; - - const modelId = (b: any) => String(b?.modelId || "").toLowerCase(); - const claudeNonThinking = buckets.filter((b: any) => modelId(b).includes("claude") && !modelId(b).includes("thinking")); - const geminiPro = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("pro")); - const geminiFlash = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("flash")); - - const primaryBucket = - provider === "antigravity" - ? pickMostUsedBucket(claudeNonThinking) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets) - : pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets); - - const secondaryBucket = pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(buckets); - - const session = usedPercentFromRemainingFraction(primaryBucket?.remainingFraction); - const weekly = usedPercentFromRemainingFraction(secondaryBucket?.remainingFraction); - - if (session == null || weekly == null) return null; - return normalizeUsagePair(session, weekly); -} - -export async function fetchCodexUsage(token: string, config: RequestConfig = {}): Promise { - const result = await requestJson( - "https://chatgpt.com/backend-api/wham/usage", - { headers: { Authorization: `Bearer ${token}` } }, - config, - ); - - if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; - - const primary = result.data?.rate_limit?.primary_window; - const secondary = result.data?.rate_limit?.secondary_window; - - return { - session: readPercentCandidate(primary?.used_percent) ?? 0, - weekly: readPercentCandidate(secondary?.used_percent) ?? 0, - sessionResetsIn: typeof primary?.reset_after_seconds === "number" ? formatDuration(primary.reset_after_seconds) : undefined, - weeklyResetsIn: typeof secondary?.reset_after_seconds === "number" ? formatDuration(secondary.reset_after_seconds) : undefined, - }; -} - -export async function fetchClaudeUsage(token: string, config: RequestConfig = {}): Promise { - const result = await requestJson( - "https://api.anthropic.com/api/oauth/usage", - { - headers: { - Authorization: `Bearer ${token}`, - "anthropic-beta": "oauth-2025-04-20", - }, - }, - config, - ); - - if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; - - const data = result.data; - const usage: UsageData = { - session: readPercentCandidate(data?.five_hour?.utilization) ?? 0, - weekly: readPercentCandidate(data?.seven_day?.utilization) ?? 0, - sessionResetsIn: data?.five_hour?.resets_at ? formatResetsAt(data.five_hour.resets_at) : undefined, - weeklyResetsIn: data?.seven_day?.resets_at ? formatResetsAt(data.seven_day.resets_at) : undefined, - }; - - if (data?.extra_usage?.is_enabled) { - usage.extraSpend = typeof data.extra_usage.used_credits === "number" ? data.extra_usage.used_credits : undefined; - usage.extraLimit = typeof data.extra_usage.monthly_limit === "number" ? data.extra_usage.monthly_limit : undefined; - } - - return usage; -} - -export async function fetchZaiUsage(token: string, config: FetchConfig = {}): Promise { - const endpoint = (config.endpoints ?? resolveUsageEndpoints()).zai; - if (!endpoint) return { session: 0, weekly: 0, error: "usage endpoint unavailable" }; - - const result = await requestJson( - endpoint, - { headers: { Authorization: `Bearer ${token}` } }, - config, - ); - - if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; - - const parsed = extractUsageFromPayload(result.data); - if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" }; - return parsed; -} - -export async function fetchGoogleUsage( - token: string, - endpoint: string, - projectId: string | undefined, - provider: "gemini" | "antigravity", - config: FetchConfig = {}, -): Promise { - if (!endpoint) return { session: 0, weekly: 0, error: "configure endpoint" }; - - const discoveredProjectId = projectId || (await discoverGoogleProjectId(token, config)); - if (!discoveredProjectId) { - return { session: 0, weekly: 0, error: "missing projectId (try /login again)" }; - } - - const result = await requestJson( - endpoint, - { - method: "POST", - headers: googleHeaders(token, discoveredProjectId), - body: JSON.stringify({ project: discoveredProjectId }), - }, - config, - ); - - if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; - - const quota = parseGoogleQuotaBuckets(result.data, provider); - if (quota) return quota; - - const parsed = extractUsageFromPayload(result.data); - if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" }; - return parsed; -} - -export function detectProvider( - model: { provider?: string; id?: string; name?: string; api?: string } | string | undefined | null, -): ProviderKey | null { - if (!model) return null; - if (typeof model === "string") return null; - - const provider = (model.provider || "").toLowerCase(); - - if (provider === "openai-codex") return "codex"; - if (provider === "anthropic") return "claude"; - if (provider === "zai") return "zai"; - if (provider === "google-gemini-cli") return "gemini"; - if (provider === "google-antigravity") return "antigravity"; - - return null; -} - -export function providerToOAuthProviderId(active: ProviderKey | null): OAuthProviderId | null { - if (active === "codex") return "openai-codex"; - if (active === "claude") return "anthropic"; - if (active === "gemini") return "google-gemini-cli"; - if (active === "antigravity") return "google-antigravity"; - return null; -} - -export function canShowForProvider(active: ProviderKey | null, auth: AuthData | null, endpoints: UsageEndpoints): boolean { - if (!active || !auth) return false; - if (active === "codex") return !!(auth["openai-codex"]?.access || auth["openai-codex"]?.refresh); - if (active === "claude") return !!(auth.anthropic?.access || auth.anthropic?.refresh); - if (active === "zai") return !!(auth.zai?.access || auth.zai?.key) && !!endpoints.zai; - if (active === "gemini") { - return !!(auth["google-gemini-cli"]?.access || auth["google-gemini-cli"]?.refresh) && !!endpoints.gemini; - } - if (active === "antigravity") { - return !!(auth["google-antigravity"]?.access || auth["google-antigravity"]?.refresh) && !!endpoints.antigravity; - } - return false; -} - -export function clampPercent(value: number): number { - if (!Number.isFinite(value)) return 0; - return Math.max(0, Math.min(100, Math.round(value))); -} - -export function colorForPercent(value: number): "success" | "warning" | "error" { - if (value >= 90) return "error"; - if (value >= 70) return "warning"; - return "success"; -} - -export async function fetchAllUsages(config: FetchAllUsagesConfig = {}): Promise { - const authFile = config.authFile ?? DEFAULT_AUTH_FILE; - const auth = config.auth ?? readAuth(authFile); - const endpoints = config.endpoints ?? resolveUsageEndpoints(); - - const results: UsageByProvider = { - codex: null, - claude: null, - zai: null, - gemini: null, - antigravity: null, - }; - - if (!auth) return results; - - const oauthProviders: OAuthProviderId[] = [ - "openai-codex", - "anthropic", - "google-gemini-cli", - "google-antigravity", - ]; - - const refreshed = await ensureFreshAuthForProviders(oauthProviders, { - ...config, - auth, - authFile, - }); - - const authData = refreshed.auth ?? auth; - - const refreshError = (providerId: OAuthProviderId): string | null => { - const error = refreshed.refreshErrors[providerId]; - return error ? `auth refresh failed (${error})` : null; - }; - - const tasks: Promise[] = []; - const assign = (provider: ProviderKey, task: Promise) => { - tasks.push( - task - .then((usage) => { - results[provider] = usage; - }) - .catch((error) => { - results[provider] = { session: 0, weekly: 0, error: toErrorMessage(error) }; - }), - ); - }; - - if (authData["openai-codex"]?.access) { - const err = refreshError("openai-codex"); - if (err) results.codex = { session: 0, weekly: 0, error: err }; - else assign("codex", fetchCodexUsage(authData["openai-codex"].access, config)); - } - - if (authData.anthropic?.access) { - const err = refreshError("anthropic"); - if (err) results.claude = { session: 0, weekly: 0, error: err }; - else assign("claude", fetchClaudeUsage(authData.anthropic.access, config)); - } - - if (authData.zai?.access || authData.zai?.key) { - assign("zai", fetchZaiUsage(authData.zai.access || authData.zai.key!, { ...config, endpoints })); - } - - if (authData["google-gemini-cli"]?.access) { - const err = refreshError("google-gemini-cli"); - if (err) { - results.gemini = { session: 0, weekly: 0, error: err }; - } else { - const creds = authData["google-gemini-cli"]; - assign( - "gemini", - fetchGoogleUsage(creds.access!, endpoints.gemini, creds.projectId, "gemini", { ...config, endpoints }), - ); - } - } - - if (authData["google-antigravity"]?.access) { - const err = refreshError("google-antigravity"); - if (err) { - results.antigravity = { session: 0, weekly: 0, error: err }; - } else { - const creds = authData["google-antigravity"]; - assign( - "antigravity", - fetchGoogleUsage(creds.access!, endpoints.antigravity, creds.projectId, "antigravity", { ...config, endpoints }), - ); - } - } - - await Promise.all(tasks); - return results; -} +import * as fs from "node:fs"; +import * as path from "node:path"; +import * as os from "node:os"; + +export type ProviderKey = "codex" | "claude" | "zai" | "gemini" | "antigravity"; +export type OAuthProviderId = "openai-codex" | "anthropic" | "google-gemini-cli" | "google-antigravity"; + +export interface AuthData { + "openai-codex"?: { access?: string; refresh?: string; expires?: number }; + anthropic?: { access?: string; refresh?: string; expires?: number }; + zai?: { key?: string; access?: string; refresh?: string; expires?: number }; + "google-gemini-cli"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; + "google-antigravity"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; +} + +export interface UsageData { + session: number; + weekly: number; + sessionResetsIn?: string; + weeklyResetsIn?: string; + extraSpend?: number; + extraLimit?: number; + error?: string; +} + +export type UsageByProvider = Record; + +export interface UsageEndpoints { + zai: string; + gemini: string; + antigravity: string; + googleLoadCodeAssistEndpoints: string[]; +} + +export interface FetchResponseLike { + ok: boolean; + status: number; + json(): Promise; +} + +export type FetchLike = (input: string, init?: RequestInit) => Promise; + +export interface RequestConfig { + fetchFn?: FetchLike; + timeoutMs?: number; +} + +export interface FetchConfig extends RequestConfig { + endpoints?: UsageEndpoints; + env?: NodeJS.ProcessEnv; +} + +export interface OAuthApiKeyResult { + newCredentials: Record; + apiKey: string; +} + +export type OAuthApiKeyResolver = ( + providerId: OAuthProviderId, + credentials: Record>, +) => Promise; + +export interface EnsureFreshAuthConfig { + auth?: AuthData | null; + authFile?: string; + oauthResolver?: OAuthApiKeyResolver; + nowMs?: number; + persist?: boolean; +} + +export interface FreshAuthResult { + auth: AuthData | null; + changed: boolean; + refreshErrors: Partial>; +} + +export interface FetchAllUsagesConfig extends FetchConfig, EnsureFreshAuthConfig { + auth?: AuthData | null; + authFile?: string; +} + +const DEFAULT_FETCH_TIMEOUT_MS = 12_000; +const TOKEN_REFRESH_SKEW_MS = 60_000; + +export const DEFAULT_AUTH_FILE = path.join(os.homedir(), ".pi", "agent", "auth.json"); +export const DEFAULT_ZAI_USAGE_ENDPOINT = "https://api.z.ai/api/monitor/usage/quota/limit"; +export const GOOGLE_QUOTA_ENDPOINT = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"; +export const GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS = [ + "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", + "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist", +]; + +export function resolveUsageEndpoints(): UsageEndpoints { + return { + zai: DEFAULT_ZAI_USAGE_ENDPOINT, + gemini: GOOGLE_QUOTA_ENDPOINT, + antigravity: GOOGLE_QUOTA_ENDPOINT, + googleLoadCodeAssistEndpoints: GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS, + }; +} + +function toErrorMessage(error: unknown): string { + if (error instanceof Error) { + if (error.name === "AbortError") return "request timeout"; + return error.message || String(error); + } + return String(error); +} + +function asObject(value: unknown): Record | null { + if (!value || typeof value !== "object") return null; + return value as Record; +} + +function normalizeUsagePair(session: number, weekly: number): { session: number; weekly: number } { + const clean = (v: number) => { + if (!Number.isFinite(v)) return 0; + return Number(v.toFixed(2)); + }; + return { session: clean(session), weekly: clean(weekly) }; +} + +async function requestJson(url: string, init: RequestInit, config: RequestConfig = {}): Promise<{ ok: true; data: any } | { ok: false; error: string }> { + const fetchFn = config.fetchFn ?? ((fetch as unknown) as FetchLike); + const timeoutMs = config.timeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS; + const controller = new AbortController(); + const timeout = timeoutMs > 0 ? setTimeout(() => controller.abort(), timeoutMs) : null; + + try { + const response = await fetchFn(url, { ...init, signal: controller.signal }); + if (!response.ok) return { ok: false, error: `HTTP ${response.status}` }; + + try { + const data = await response.json(); + return { ok: true, data }; + } catch { + return { ok: false, error: "invalid JSON response" }; + } + } catch (error) { + return { ok: false, error: toErrorMessage(error) }; + } finally { + if (timeout) clearTimeout(timeout); + } +} + +export function formatDuration(seconds: number): string { + if (!Number.isFinite(seconds) || seconds <= 0) return "now"; + const d = Math.floor(seconds / 86400); + const h = Math.floor((seconds % 86400) / 3600); + const m = Math.floor((seconds % 3600) / 60); + if (d > 0 && h > 0) return `${d}d ${h}h`; + if (d > 0) return `${d}d`; + if (h > 0 && m > 0) return `${h}h ${m}m`; + if (h > 0) return `${h}h`; + if (m > 0) return `${m}m`; + return "<1m"; +} + +export function formatResetsAt(isoDate: string, nowMs = Date.now()): string { + const resetTime = new Date(isoDate).getTime(); + if (!Number.isFinite(resetTime)) return ""; + const diffSeconds = Math.max(0, (resetTime - nowMs) / 1000); + return formatDuration(diffSeconds); +} + +export function readAuth(authFile = DEFAULT_AUTH_FILE): AuthData | null { + try { + const parsed = JSON.parse(fs.readFileSync(authFile, "utf-8")); + return asObject(parsed) as AuthData; + } catch { + return null; + } +} + +export function writeAuth(auth: AuthData, authFile = DEFAULT_AUTH_FILE): boolean { + try { + const dir = path.dirname(authFile); + if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true }); + + const tmpPath = `${authFile}.tmp-${process.pid}-${Date.now()}`; + fs.writeFileSync(tmpPath, JSON.stringify(auth, null, 2)); + fs.renameSync(tmpPath, authFile); + return true; + } catch { + return false; + } +} + +let cachedOAuthResolver: OAuthApiKeyResolver | null = null; + +async function getDefaultOAuthResolver(): Promise { + if (cachedOAuthResolver) return cachedOAuthResolver; + + const mod = await import("@earendil-works/pi-ai"); + if (typeof (mod as any).getOAuthApiKey !== "function") { + throw new Error("oauth resolver unavailable"); + } + + cachedOAuthResolver = (providerId, credentials) => + (mod as any).getOAuthApiKey(providerId, credentials) as Promise; + + return cachedOAuthResolver; +} + +function isCredentialExpired(creds: { expires?: number } | undefined, nowMs: number): boolean { + if (!creds) return false; + if (typeof creds.expires !== "number") return false; + return nowMs + TOKEN_REFRESH_SKEW_MS >= creds.expires; +} + +export async function ensureFreshAuthForProviders( + providerIds: OAuthProviderId[], + config: EnsureFreshAuthConfig = {}, +): Promise { + const authFile = config.authFile ?? DEFAULT_AUTH_FILE; + const auth = config.auth ?? readAuth(authFile); + if (!auth) { + return { auth: null, changed: false, refreshErrors: {} }; + } + + const nowMs = config.nowMs ?? Date.now(); + const uniqueProviders = Array.from(new Set(providerIds)); + const nextAuth: AuthData = { ...auth }; + const refreshErrors: Partial> = {}; + + let changed = false; + + for (const providerId of uniqueProviders) { + const creds = (nextAuth as any)[providerId] as { access?: string; refresh?: string; expires?: number } | undefined; + if (!creds?.refresh) continue; + + const needsRefresh = !creds.access || isCredentialExpired(creds, nowMs); + if (!needsRefresh) continue; + + try { + const resolver = config.oauthResolver ?? (await getDefaultOAuthResolver()); + const resolved = await resolver(providerId, nextAuth as any); + if (!resolved?.newCredentials) { + refreshErrors[providerId] = "missing OAuth credentials"; + continue; + } + + (nextAuth as any)[providerId] = { + ...(nextAuth as any)[providerId], + ...resolved.newCredentials, + }; + changed = true; + } catch (error) { + refreshErrors[providerId] = toErrorMessage(error); + } + } + + if (changed && config.persist !== false) { + writeAuth(nextAuth, authFile); + } + + return { auth: nextAuth, changed, refreshErrors }; +} + +export function readPercentCandidate(value: unknown): number | null { + if (typeof value !== "number" || !Number.isFinite(value)) return null; + + if (value >= 0 && value <= 1) { + if (Number.isInteger(value)) return value; + return value * 100; + } + + if (value >= 0 && value <= 100) return value; + return null; +} + +export function readLimitPercent(limit: any): number | null { + const direct = [ + limit?.percentage, + limit?.utilization, + limit?.used_percent, + limit?.usedPercent, + limit?.usagePercent, + limit?.usage_percent, + ] + .map(readPercentCandidate) + .find((v) => v != null); + + if (direct != null) return direct; + + const current = typeof limit?.currentValue === "number" ? limit.currentValue : null; + const remaining = typeof limit?.remaining === "number" ? limit.remaining : null; + + if (current != null && remaining != null && current + remaining > 0) { + return (current / (current + remaining)) * 100; + } + + return null; +} + +export function extractUsageFromPayload(data: any): { session: number; weekly: number } | null { + const limitArrays = [data?.data?.limits, data?.limits, data?.quota?.limits, data?.data?.quota?.limits]; + const limits = limitArrays.find((arr) => Array.isArray(arr)) as any[] | undefined; + + if (limits) { + const byType = (types: string[]) => + limits.find((l) => { + const t = String(l?.type || "").toUpperCase(); + return types.some((x) => t === x); + }); + + const sessionLimit = byType(["TIME_LIMIT", "SESSION_LIMIT", "REQUEST_LIMIT", "RPM_LIMIT", "RPD_LIMIT"]); + const weeklyLimit = byType(["TOKENS_LIMIT", "TOKEN_LIMIT", "WEEK_LIMIT", "WEEKLY_LIMIT", "TPM_LIMIT", "DAILY_LIMIT"]); + + const s = readLimitPercent(sessionLimit); + const w = readLimitPercent(weeklyLimit); + if (s != null && w != null) return normalizeUsagePair(s, w); + } + + const sessionCandidates = [ + data?.session, + data?.sessionPercent, + data?.session_percent, + data?.five_hour?.utilization, + data?.rate_limit?.primary_window?.used_percent, + data?.limits?.session?.utilization, + data?.usage?.session, + data?.data?.session, + data?.data?.sessionPercent, + data?.data?.session_percent, + data?.data?.usage?.session, + data?.quota?.session?.percentage, + data?.data?.quota?.session?.percentage, + ]; + + const weeklyCandidates = [ + data?.weekly, + data?.weeklyPercent, + data?.weekly_percent, + data?.seven_day?.utilization, + data?.rate_limit?.secondary_window?.used_percent, + data?.limits?.weekly?.utilization, + data?.usage?.weekly, + data?.data?.weekly, + data?.data?.weeklyPercent, + data?.data?.weekly_percent, + data?.data?.usage?.weekly, + data?.quota?.weekly?.percentage, + data?.data?.quota?.weekly?.percentage, + data?.quota?.daily?.percentage, + data?.data?.quota?.daily?.percentage, + ]; + + const session = sessionCandidates.map(readPercentCandidate).find((v) => v != null); + const weekly = weeklyCandidates.map(readPercentCandidate).find((v) => v != null); + + if (session == null || weekly == null) return null; + return normalizeUsagePair(session, weekly); +} + +export function googleMetadata(projectId?: string) { + return { + ideType: "IDE_UNSPECIFIED", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + ...(projectId ? { duetProject: projectId } : {}), + }; +} + +export function googleHeaders(token: string, projectId?: string) { + return { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1", + "X-Goog-Api-Client": "gl-node/22.17.0", + "Client-Metadata": JSON.stringify(googleMetadata(projectId)), + }; +} + +export async function discoverGoogleProjectId(token: string, config: FetchConfig = {}): Promise { + const env = config.env ?? process.env; + const envProjectId = env.GOOGLE_CLOUD_PROJECT || env.GOOGLE_CLOUD_PROJECT_ID; + if (envProjectId) return envProjectId; + + const endpoints = config.endpoints ?? resolveUsageEndpoints(); + + for (const endpoint of endpoints.googleLoadCodeAssistEndpoints) { + const result = await requestJson( + endpoint, + { + method: "POST", + headers: googleHeaders(token), + body: JSON.stringify({ metadata: googleMetadata() }), + }, + config, + ); + + if (!result.ok) continue; + + const data = result.data; + if (typeof data?.cloudaicompanionProject === "string" && data.cloudaicompanionProject) { + return data.cloudaicompanionProject; + } + if (data?.cloudaicompanionProject && typeof data.cloudaicompanionProject === "object") { + const id = data.cloudaicompanionProject.id; + if (typeof id === "string" && id) return id; + } + } + + return undefined; +} + +export function usedPercentFromRemainingFraction(value: unknown): number | null { + if (typeof value !== "number" || !Number.isFinite(value)) return null; + const remaining = Math.max(0, Math.min(1, value)); + return (1 - remaining) * 100; +} + +export function pickMostUsedBucket(buckets: any[]): any | null { + let best: any | null = null; + let bestUsed = -1; + for (const bucket of buckets) { + const used = usedPercentFromRemainingFraction(bucket?.remainingFraction); + if (used == null) continue; + if (used > bestUsed) { + bestUsed = used; + best = bucket; + } + } + return best; +} + +export function parseGoogleQuotaBuckets(data: any, provider: "gemini" | "antigravity"): { session: number; weekly: number } | null { + const allBuckets = Array.isArray(data?.buckets) ? data.buckets : []; + if (!allBuckets.length) return null; + + const requestBuckets = allBuckets.filter((b: any) => String(b?.tokenType || "").toUpperCase() === "REQUESTS"); + const buckets = requestBuckets.length ? requestBuckets : allBuckets; + + const modelId = (b: any) => String(b?.modelId || "").toLowerCase(); + const claudeNonThinking = buckets.filter((b: any) => modelId(b).includes("claude") && !modelId(b).includes("thinking")); + const geminiPro = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("pro")); + const geminiFlash = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("flash")); + + const primaryBucket = + provider === "antigravity" + ? pickMostUsedBucket(claudeNonThinking) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets) + : pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets); + + const secondaryBucket = pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(buckets); + + const session = usedPercentFromRemainingFraction(primaryBucket?.remainingFraction); + const weekly = usedPercentFromRemainingFraction(secondaryBucket?.remainingFraction); + + if (session == null || weekly == null) return null; + return normalizeUsagePair(session, weekly); +} + +export async function fetchCodexUsage(token: string, config: RequestConfig = {}): Promise { + const result = await requestJson( + "https://chatgpt.com/backend-api/wham/usage", + { headers: { Authorization: `Bearer ${token}` } }, + config, + ); + + if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; + + const primary = result.data?.rate_limit?.primary_window; + const secondary = result.data?.rate_limit?.secondary_window; + + return { + session: readPercentCandidate(primary?.used_percent) ?? 0, + weekly: readPercentCandidate(secondary?.used_percent) ?? 0, + sessionResetsIn: typeof primary?.reset_after_seconds === "number" ? formatDuration(primary.reset_after_seconds) : undefined, + weeklyResetsIn: typeof secondary?.reset_after_seconds === "number" ? formatDuration(secondary.reset_after_seconds) : undefined, + }; +} + +export async function fetchClaudeUsage(token: string, config: RequestConfig = {}): Promise { + const result = await requestJson( + "https://api.anthropic.com/api/oauth/usage", + { + headers: { + Authorization: `Bearer ${token}`, + "anthropic-beta": "oauth-2025-04-20", + }, + }, + config, + ); + + if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; + + const data = result.data; + const usage: UsageData = { + session: readPercentCandidate(data?.five_hour?.utilization) ?? 0, + weekly: readPercentCandidate(data?.seven_day?.utilization) ?? 0, + sessionResetsIn: data?.five_hour?.resets_at ? formatResetsAt(data.five_hour.resets_at) : undefined, + weeklyResetsIn: data?.seven_day?.resets_at ? formatResetsAt(data.seven_day.resets_at) : undefined, + }; + + if (data?.extra_usage?.is_enabled) { + usage.extraSpend = typeof data.extra_usage.used_credits === "number" ? data.extra_usage.used_credits : undefined; + usage.extraLimit = typeof data.extra_usage.monthly_limit === "number" ? data.extra_usage.monthly_limit : undefined; + } + + return usage; +} + +export async function fetchZaiUsage(token: string, config: FetchConfig = {}): Promise { + const endpoint = (config.endpoints ?? resolveUsageEndpoints()).zai; + if (!endpoint) return { session: 0, weekly: 0, error: "usage endpoint unavailable" }; + + const result = await requestJson( + endpoint, + { headers: { Authorization: `Bearer ${token}` } }, + config, + ); + + if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; + + const parsed = extractUsageFromPayload(result.data); + if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" }; + return parsed; +} + +export async function fetchGoogleUsage( + token: string, + endpoint: string, + projectId: string | undefined, + provider: "gemini" | "antigravity", + config: FetchConfig = {}, +): Promise { + if (!endpoint) return { session: 0, weekly: 0, error: "configure endpoint" }; + + const discoveredProjectId = projectId || (await discoverGoogleProjectId(token, config)); + if (!discoveredProjectId) { + return { session: 0, weekly: 0, error: "missing projectId (try /login again)" }; + } + + const result = await requestJson( + endpoint, + { + method: "POST", + headers: googleHeaders(token, discoveredProjectId), + body: JSON.stringify({ project: discoveredProjectId }), + }, + config, + ); + + if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; + + const quota = parseGoogleQuotaBuckets(result.data, provider); + if (quota) return quota; + + const parsed = extractUsageFromPayload(result.data); + if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" }; + return parsed; +} + +export function detectProvider( + model: { provider?: string; id?: string; name?: string; api?: string } | string | undefined | null, +): ProviderKey | null { + if (!model) return null; + if (typeof model === "string") return null; + + const provider = (model.provider || "").toLowerCase(); + + if (provider === "openai-codex") return "codex"; + if (provider === "anthropic") return "claude"; + if (provider === "zai") return "zai"; + if (provider === "google-gemini-cli") return "gemini"; + if (provider === "google-antigravity") return "antigravity"; + + return null; +} + +export function providerToOAuthProviderId(active: ProviderKey | null): OAuthProviderId | null { + if (active === "codex") return "openai-codex"; + if (active === "claude") return "anthropic"; + if (active === "gemini") return "google-gemini-cli"; + if (active === "antigravity") return "google-antigravity"; + return null; +} + +export function canShowForProvider(active: ProviderKey | null, auth: AuthData | null, endpoints: UsageEndpoints): boolean { + if (!active || !auth) return false; + if (active === "codex") return !!(auth["openai-codex"]?.access || auth["openai-codex"]?.refresh); + if (active === "claude") return !!(auth.anthropic?.access || auth.anthropic?.refresh); + if (active === "zai") return !!(auth.zai?.access || auth.zai?.key) && !!endpoints.zai; + if (active === "gemini") { + return !!(auth["google-gemini-cli"]?.access || auth["google-gemini-cli"]?.refresh) && !!endpoints.gemini; + } + if (active === "antigravity") { + return !!(auth["google-antigravity"]?.access || auth["google-antigravity"]?.refresh) && !!endpoints.antigravity; + } + return false; +} + +export function clampPercent(value: number): number { + if (!Number.isFinite(value)) return 0; + return Math.max(0, Math.min(100, Math.round(value))); +} + +export function colorForPercent(value: number): "success" | "warning" | "error" { + if (value >= 90) return "error"; + if (value >= 70) return "warning"; + return "success"; +} + +export async function fetchAllUsages(config: FetchAllUsagesConfig = {}): Promise { + const authFile = config.authFile ?? DEFAULT_AUTH_FILE; + const auth = config.auth ?? readAuth(authFile); + const endpoints = config.endpoints ?? resolveUsageEndpoints(); + + const results: UsageByProvider = { + codex: null, + claude: null, + zai: null, + gemini: null, + antigravity: null, + }; + + if (!auth) return results; + + const oauthProviders: OAuthProviderId[] = [ + "openai-codex", + "anthropic", + "google-gemini-cli", + "google-antigravity", + ]; + + const refreshed = await ensureFreshAuthForProviders(oauthProviders, { + ...config, + auth, + authFile, + }); + + const authData = refreshed.auth ?? auth; + + const refreshError = (providerId: OAuthProviderId): string | null => { + const error = refreshed.refreshErrors[providerId]; + return error ? `auth refresh failed (${error})` : null; + }; + + const tasks: Promise[] = []; + const assign = (provider: ProviderKey, task: Promise) => { + tasks.push( + task + .then((usage) => { + results[provider] = usage; + }) + .catch((error) => { + results[provider] = { session: 0, weekly: 0, error: toErrorMessage(error) }; + }), + ); + }; + + if (authData["openai-codex"]?.access) { + const err = refreshError("openai-codex"); + if (err) results.codex = { session: 0, weekly: 0, error: err }; + else assign("codex", fetchCodexUsage(authData["openai-codex"].access, config)); + } + + if (authData.anthropic?.access) { + const err = refreshError("anthropic"); + if (err) results.claude = { session: 0, weekly: 0, error: err }; + else assign("claude", fetchClaudeUsage(authData.anthropic.access, config)); + } + + if (authData.zai?.access || authData.zai?.key) { + assign("zai", fetchZaiUsage(authData.zai.access || authData.zai.key!, { ...config, endpoints })); + } + + if (authData["google-gemini-cli"]?.access) { + const err = refreshError("google-gemini-cli"); + if (err) { + results.gemini = { session: 0, weekly: 0, error: err }; + } else { + const creds = authData["google-gemini-cli"]; + assign( + "gemini", + fetchGoogleUsage(creds.access!, endpoints.gemini, creds.projectId, "gemini", { ...config, endpoints }), + ); + } + } + + if (authData["google-antigravity"]?.access) { + const err = refreshError("google-antigravity"); + if (err) { + results.antigravity = { session: 0, weekly: 0, error: err }; + } else { + const creds = authData["google-antigravity"]; + assign( + "antigravity", + fetchGoogleUsage(creds.access!, endpoints.antigravity, creds.projectId, "antigravity", { ...config, endpoints }), + ); + } + } + + await Promise.all(tasks); + return results; +} diff --git a/extensions/usage/index.ts b/extensions/usage/index.ts index b047aa8..97fd9ac 100644 --- a/extensions/usage/index.ts +++ b/extensions/usage/index.ts @@ -1,296 +1,296 @@ -/** - * Simple Usage Extension for pi - * - * Provides a single /usage command that shows the current provider's - * daily (session) and weekly limits in a TUI panel. - */ - -import { AuthStorage, DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; -import { matchesKey, Container, Spacer, Text, type Focusable } from "@mariozechner/pi-tui"; -import { - clampPercent, - colorForPercent, - detectProvider, - fetchClaudeUsage, - fetchCodexUsage, - fetchGoogleUsage, - fetchZaiUsage, - providerToOAuthProviderId, - resolveUsageEndpoints, - type ProviderKey, - type UsageData, - type UsageEndpoints, -} from "./core.js"; - -// ── Provider labels ────────────────────────────────────────────── - -const PROVIDER_LABELS: Record = { - codex: "Codex", - claude: "Claude", - zai: "Z.AI", - gemini: "Gemini", - antigravity: "Antigravity", -}; - -// ── Self-managing usage panel ──────────────────────────────────── - -class UsagePanelComponent extends Container implements Focusable { - private _focused = false; - private tui: any; - private theme: any; - private onDone: () => void; - - // Mutable content area (everything between the two borders) - private contentContainer = new Container(); - - get focused() { - return this._focused; - } - set focused(v: boolean) { - this._focused = v; - } - - constructor( - tui: any, - theme: any, - provider: ProviderKey, - fetchPromise: Promise, - onDone: () => void, - ) { - super(); - this.tui = tui; - this.theme = theme; - this.onDone = onDone; - - // Static structure: top border → content → bottom border - this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); - this.addChild(this.contentContainer); - this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); - - // Show loading state - this.renderLoading(provider); - - // Kick off fetch – when it resolves, rebuild content in place - fetchPromise - .then((data) => this.renderResult(provider, data)) - .catch((err) => - this.renderResult(provider, { - session: 0, - weekly: 0, - error: String(err?.message ?? err), - }), - ); - } - - // ── Renderers ──────────────────────────────────────────────── - - private renderLoading(provider: ProviderKey) { - this.contentContainer.clear(); - const t = this.theme; - this.contentContainer.addChild(new Spacer(1)); - this.contentContainer.addChild( - new Text( - " " + t.fg("accent", t.bold(PROVIDER_LABELS[provider] ?? provider)) + " Usage", - 0, - 0, - ), - ); - this.contentContainer.addChild(new Spacer(1)); - this.contentContainer.addChild( - new Text(" " + t.fg("muted", "Fetching usage data…"), 0, 0), - ); - this.contentContainer.addChild(new Spacer(1)); - } - - private renderResult(provider: ProviderKey, data: UsageData | null) { - this.contentContainer.clear(); - const t = this.theme; - const label = PROVIDER_LABELS[provider] ?? provider; - - this.contentContainer.addChild(new Spacer(1)); - this.contentContainer.addChild( - new Text(" " + t.fg("accent", t.bold(label)) + " Usage", 0, 0), - ); - this.contentContainer.addChild(new Spacer(1)); - - if (!data) { - this.contentContainer.addChild( - new Text(" " + t.fg("dim", "No credentials found for this provider."), 0, 0), - ); - } else if (data.error) { - this.contentContainer.addChild( - new Text(" " + t.fg("error", `Error: ${data.error}`), 0, 0), - ); - } else { - const session = clampPercent(data.session); - const sessionReset = data.sessionResetsIn - ? t.fg("dim", ` resets in ${data.sessionResetsIn}`) - : ""; - this.contentContainer.addChild( - new Text( - " " + - t.fg("muted", "Daily ") + - renderBar(t, session) + - " " + - t.fg(colorForPercent(session), `${session}%`.padStart(4)) + - sessionReset, - 0, - 0, - ), - ); - - const weekly = clampPercent(data.weekly); - const weeklyReset = data.weeklyResetsIn - ? t.fg("dim", ` resets in ${data.weeklyResetsIn}`) - : ""; - this.contentContainer.addChild( - new Text( - " " + - t.fg("muted", "Weekly ") + - renderBar(t, weekly) + - " " + - t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + - weeklyReset, - 0, - 0, - ), - ); - - if (typeof data.extraSpend === "number" && typeof data.extraLimit === "number") { - this.contentContainer.addChild( - new Text( - " " + - t.fg("muted", "Extra ") + - t.fg("dim", `$${data.extraSpend.toFixed(2)} / $${data.extraLimit}`), - 0, - 0, - ), - ); - } - } - - this.contentContainer.addChild(new Spacer(1)); - this.contentContainer.addChild( - new Text(" " + t.fg("dim", "Press Enter or Escape to close"), 0, 0), - ); - this.contentContainer.addChild(new Spacer(1)); - - // ← This is the key fix: trigger a TUI re-render after content changes - this.tui.requestRender(); - } - - // ── Input ──────────────────────────────────────────────────── - - handleInput(keyData: string): void { - // Use matchesKey for raw key matching – no dependency on keybinding config - if (matchesKey(keyData, "return") || matchesKey(keyData, "escape")) { - this.onDone(); - } - } -} - -// ── Helpers ────────────────────────────────────────────────────── - -function renderBar(theme: any, value: number, width = 20): string { - const v = clampPercent(value); - const filled = Math.round((v / 100) * width); - const full = "█".repeat(Math.max(0, Math.min(width, filled))); - const empty = "░".repeat(Math.max(0, width - filled)); - return theme.fg(colorForPercent(v), full) + theme.fg("dim", empty); -} - -async function getAccessToken(providerId: string): Promise { - try { - const auth = AuthStorage.create(); - return (await auth.getApiKey(providerId)) ?? null; - } catch { - return null; - } -} - -async function fetchProviderUsage( - provider: ProviderKey, - endpoints: UsageEndpoints, -): Promise { - const oauthId = providerToOAuthProviderId(provider); - - switch (provider) { - case "codex": { - const access = oauthId ? await getAccessToken(oauthId) : null; - return access - ? fetchCodexUsage(access) - : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; - } - case "claude": { - const access = oauthId ? await getAccessToken(oauthId) : null; - return access - ? fetchClaudeUsage(access) - : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; - } - case "zai": { - const token = process.env.ZAI_API_KEY || (oauthId ? await getAccessToken(oauthId) : null); - return token - ? fetchZaiUsage(token, { endpoints }) - : { session: 0, weekly: 0, error: "missing token (set ZAI_API_KEY or try /login again)" }; - } - case "gemini": { - const access = oauthId ? await getAccessToken(oauthId) : null; - // AuthStorage manages projectId internally via auth.json - return access - ? fetchGoogleUsage(access, endpoints.gemini, undefined, "gemini", { endpoints }) - : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; - } - case "antigravity": { - const access = oauthId ? await getAccessToken(oauthId) : null; - return access - ? fetchGoogleUsage(access, endpoints.antigravity, undefined, "antigravity", { endpoints }) - : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; - } - default: - return null; - } -} - -// ── Extension entry point ──────────────────────────────────────── - -export default function (pi: ExtensionAPI) { - const endpoints = resolveUsageEndpoints(); - - pi.registerCommand("usage", { - description: "Show current provider's daily & weekly usage limits", - handler: async (_args, ctx) => { - if (!ctx?.hasUI) return; - - const provider = detectProvider(ctx.model); - if (!provider) { - ctx.ui.notify( - "Cannot detect current provider – usage data unavailable", - "warning", - ); - return; - } - - const auth = AuthStorage.create(); - const oauthId = providerToOAuthProviderId(provider); - - // Z.AI uses ZAI_API_KEY env var, not OAuth - if (provider === "zai" && !process.env.ZAI_API_KEY && !(oauthId && auth.hasAuth(oauthId))) { - ctx.ui.notify("No ZAI_API_KEY found (set env var or run /login)", "warning"); - return; - } - - if (provider !== "zai" && oauthId && !auth.hasAuth(oauthId)) { - ctx.ui.notify( - `No credentials found for ${PROVIDER_LABELS[provider] ?? provider}`, - "warning", - ); - return; - } - - await ctx.ui.custom((tui, theme, _keybindings, done) => { - const fetchPromise = fetchProviderUsage(provider, endpoints); - return new UsagePanelComponent(tui, theme, provider, fetchPromise, () => done()); - }); - }, - }); -} +/** + * Simple Usage Extension for pi + * + * Provides a single /usage command that shows the current provider's + * daily (session) and weekly limits in a TUI panel. + */ + +import { AuthStorage, DynamicBorder, type ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { matchesKey, Container, Spacer, Text, type Focusable } from "@earendil-works/pi-tui"; +import { + clampPercent, + colorForPercent, + detectProvider, + fetchClaudeUsage, + fetchCodexUsage, + fetchGoogleUsage, + fetchZaiUsage, + providerToOAuthProviderId, + resolveUsageEndpoints, + type ProviderKey, + type UsageData, + type UsageEndpoints, +} from "./core.js"; + +// ── Provider labels ────────────────────────────────────────────── + +const PROVIDER_LABELS: Record = { + codex: "Codex", + claude: "Claude", + zai: "Z.AI", + gemini: "Gemini", + antigravity: "Antigravity", +}; + +// ── Self-managing usage panel ──────────────────────────────────── + +class UsagePanelComponent extends Container implements Focusable { + private _focused = false; + private tui: any; + private theme: any; + private onDone: () => void; + + // Mutable content area (everything between the two borders) + private contentContainer = new Container(); + + get focused() { + return this._focused; + } + set focused(v: boolean) { + this._focused = v; + } + + constructor( + tui: any, + theme: any, + provider: ProviderKey, + fetchPromise: Promise, + onDone: () => void, + ) { + super(); + this.tui = tui; + this.theme = theme; + this.onDone = onDone; + + // Static structure: top border → content → bottom border + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + this.addChild(this.contentContainer); + this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s))); + + // Show loading state + this.renderLoading(provider); + + // Kick off fetch – when it resolves, rebuild content in place + fetchPromise + .then((data) => this.renderResult(provider, data)) + .catch((err) => + this.renderResult(provider, { + session: 0, + weekly: 0, + error: String(err?.message ?? err), + }), + ); + } + + // ── Renderers ──────────────────────────────────────────────── + + private renderLoading(provider: ProviderKey) { + this.contentContainer.clear(); + const t = this.theme; + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild( + new Text( + " " + t.fg("accent", t.bold(PROVIDER_LABELS[provider] ?? provider)) + " Usage", + 0, + 0, + ), + ); + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild( + new Text(" " + t.fg("muted", "Fetching usage data…"), 0, 0), + ); + this.contentContainer.addChild(new Spacer(1)); + } + + private renderResult(provider: ProviderKey, data: UsageData | null) { + this.contentContainer.clear(); + const t = this.theme; + const label = PROVIDER_LABELS[provider] ?? provider; + + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild( + new Text(" " + t.fg("accent", t.bold(label)) + " Usage", 0, 0), + ); + this.contentContainer.addChild(new Spacer(1)); + + if (!data) { + this.contentContainer.addChild( + new Text(" " + t.fg("dim", "No credentials found for this provider."), 0, 0), + ); + } else if (data.error) { + this.contentContainer.addChild( + new Text(" " + t.fg("error", `Error: ${data.error}`), 0, 0), + ); + } else { + const session = clampPercent(data.session); + const sessionReset = data.sessionResetsIn + ? t.fg("dim", ` resets in ${data.sessionResetsIn}`) + : ""; + this.contentContainer.addChild( + new Text( + " " + + t.fg("muted", "Daily ") + + renderBar(t, session) + + " " + + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + + sessionReset, + 0, + 0, + ), + ); + + const weekly = clampPercent(data.weekly); + const weeklyReset = data.weeklyResetsIn + ? t.fg("dim", ` resets in ${data.weeklyResetsIn}`) + : ""; + this.contentContainer.addChild( + new Text( + " " + + t.fg("muted", "Weekly ") + + renderBar(t, weekly) + + " " + + t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + + weeklyReset, + 0, + 0, + ), + ); + + if (typeof data.extraSpend === "number" && typeof data.extraLimit === "number") { + this.contentContainer.addChild( + new Text( + " " + + t.fg("muted", "Extra ") + + t.fg("dim", `$${data.extraSpend.toFixed(2)} / $${data.extraLimit}`), + 0, + 0, + ), + ); + } + } + + this.contentContainer.addChild(new Spacer(1)); + this.contentContainer.addChild( + new Text(" " + t.fg("dim", "Press Enter or Escape to close"), 0, 0), + ); + this.contentContainer.addChild(new Spacer(1)); + + // ← This is the key fix: trigger a TUI re-render after content changes + this.tui.requestRender(); + } + + // ── Input ──────────────────────────────────────────────────── + + handleInput(keyData: string): void { + // Use matchesKey for raw key matching – no dependency on keybinding config + if (matchesKey(keyData, "return") || matchesKey(keyData, "escape")) { + this.onDone(); + } + } +} + +// ── Helpers ────────────────────────────────────────────────────── + +function renderBar(theme: any, value: number, width = 20): string { + const v = clampPercent(value); + const filled = Math.round((v / 100) * width); + const full = "█".repeat(Math.max(0, Math.min(width, filled))); + const empty = "░".repeat(Math.max(0, width - filled)); + return theme.fg(colorForPercent(v), full) + theme.fg("dim", empty); +} + +async function getAccessToken(providerId: string): Promise { + try { + const auth = AuthStorage.create(); + return (await auth.getApiKey(providerId)) ?? null; + } catch { + return null; + } +} + +async function fetchProviderUsage( + provider: ProviderKey, + endpoints: UsageEndpoints, +): Promise { + const oauthId = providerToOAuthProviderId(provider); + + switch (provider) { + case "codex": { + const access = oauthId ? await getAccessToken(oauthId) : null; + return access + ? fetchCodexUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + case "claude": { + const access = oauthId ? await getAccessToken(oauthId) : null; + return access + ? fetchClaudeUsage(access) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + case "zai": { + const token = process.env.ZAI_API_KEY || (oauthId ? await getAccessToken(oauthId) : null); + return token + ? fetchZaiUsage(token, { endpoints }) + : { session: 0, weekly: 0, error: "missing token (set ZAI_API_KEY or try /login again)" }; + } + case "gemini": { + const access = oauthId ? await getAccessToken(oauthId) : null; + // AuthStorage manages projectId internally via auth.json + return access + ? fetchGoogleUsage(access, endpoints.gemini, undefined, "gemini", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + case "antigravity": { + const access = oauthId ? await getAccessToken(oauthId) : null; + return access + ? fetchGoogleUsage(access, endpoints.antigravity, undefined, "antigravity", { endpoints }) + : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; + } + default: + return null; + } +} + +// ── Extension entry point ──────────────────────────────────────── + +export default function (pi: ExtensionAPI) { + const endpoints = resolveUsageEndpoints(); + + pi.registerCommand("usage", { + description: "Show current provider's daily & weekly usage limits", + handler: async (_args, ctx) => { + if (!ctx?.hasUI) return; + + const provider = detectProvider(ctx.model); + if (!provider) { + ctx.ui.notify( + "Cannot detect current provider – usage data unavailable", + "warning", + ); + return; + } + + const auth = AuthStorage.create(); + const oauthId = providerToOAuthProviderId(provider); + + // Z.AI uses ZAI_API_KEY env var, not OAuth + if (provider === "zai" && !process.env.ZAI_API_KEY && !(oauthId && auth.hasAuth(oauthId))) { + ctx.ui.notify("No ZAI_API_KEY found (set env var or run /login)", "warning"); + return; + } + + if (provider !== "zai" && oauthId && !auth.hasAuth(oauthId)) { + ctx.ui.notify( + `No credentials found for ${PROVIDER_LABELS[provider] ?? provider}`, + "warning", + ); + return; + } + + await ctx.ui.custom((tui, theme, _keybindings, done) => { + const fetchPromise = fetchProviderUsage(provider, endpoints); + return new UsagePanelComponent(tui, theme, provider, fetchPromise, () => done()); + }); + }, + }); +} diff --git a/package.json b/package.json index cd09517..4fda677 100644 --- a/package.json +++ b/package.json @@ -1,39 +1,39 @@ -{ - "name": "pi-usage", - "version": "0.2.1", - "description": "/usage command for pi – shows current provider's daily & weekly limits", - "keywords": [ - "pi-package", - "pi", - "extension", - "usage" - ], - "license": "MIT", - "repository": { - "type": "git", - "url": "git+https://github.com/iefnaf/pi-usage.git" - }, - "bugs": { - "url": "https://github.com/iefnaf/pi-usage/issues" - }, - "homepage": "https://github.com/iefnaf/pi-usage#readme", - "type": "module", - "scripts": { - "test": "bun test" - }, - "files": [ - "extensions", - "README.md", - "LICENSE" - ], - "peerDependencies": { - "@mariozechner/pi-ai": "*", - "@mariozechner/pi-coding-agent": "*", - "@mariozechner/pi-tui": "*" - }, - "pi": { - "extensions": [ - "./extensions/usage/index.ts" - ] - } -} +{ + "name": "pi-usage", + "version": "0.2.1", + "description": "/usage command for pi – shows current provider's daily & weekly limits", + "keywords": [ + "pi-package", + "pi", + "extension", + "usage" + ], + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/iefnaf/pi-usage.git" + }, + "bugs": { + "url": "https://github.com/iefnaf/pi-usage/issues" + }, + "homepage": "https://github.com/iefnaf/pi-usage#readme", + "type": "module", + "scripts": { + "test": "bun test" + }, + "files": [ + "extensions", + "README.md", + "LICENSE" + ], + "peerDependencies": { + "@earendil-works/pi-ai": "*", + "@earendil-works/pi-coding-agent": "*", + "@earendil-works/pi-tui": "*" + }, + "pi": { + "extensions": [ + "./extensions/usage/index.ts" + ] + } +}