diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index b0dd7a01f..1112ab433 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,14 +1,25 @@ # Changelog ## [Unreleased] + +### Added + +- Added usage tracking system with normalized schema for provider quota/limit endpoints +- Added Claude usage provider for 5-hour and 7-day quota windows +- Added GitHub Copilot usage provider for chat, completions, and premium requests +- Added Google Antigravity usage provider for model quota tracking +- Added Google Gemini CLI usage provider for tier-based quota monitoring +- Added OpenAI Codex usage provider for primary and secondary rate limit windows +- Added ZAI usage provider for token and request quota tracking + ### Fixed - Fixed API validation errors when sending empty user messages (resume with `.`) across all providers: - - Google Cloud Code Assist (google-shared.ts) - - OpenAI Responses API (openai-responses.ts) - - OpenAI Codex Responses API (openai-codex-responses.ts) - - Cursor (cursor.ts) - - Amazon Bedrock (amazon-bedrock.ts) +- Google Cloud Code Assist (google-shared.ts) +- OpenAI Responses API (openai-responses.ts) +- OpenAI Codex Responses API (openai-codex-responses.ts) +- Cursor (cursor.ts) +- Amazon Bedrock (amazon-bedrock.ts) - Clamped OpenAI Codex reasoning effort "minimal" to "low" for gpt-5.2 models to avoid API errors ## [6.9.69] - 2026-01-21 diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index c07928648..0861dd87a 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -5,11 +5,18 @@ export * from "./providers/anthropic"; export * from "./providers/cursor"; export * from "./providers/google"; export * from "./providers/google-gemini-cli"; +export * from "./providers/google-gemini-cli-usage"; export * from "./providers/google-vertex"; export * from "./providers/openai-completions"; export * from "./providers/openai-responses"; export * from "./stream"; export * from "./types"; +export * from "./usage"; +export * from "./usage/claude"; +export * from "./usage/github-copilot"; +export * from "./usage/google-antigravity"; +export * from "./usage/openai-codex"; +export * from "./usage/zai"; export * from "./utils/event-stream"; export * from "./utils/oauth/index"; export * from "./utils/overflow"; diff --git a/packages/ai/src/providers/google-gemini-cli-usage.ts b/packages/ai/src/providers/google-gemini-cli-usage.ts new file mode 100644 index 000000000..b320185b4 --- /dev/null +++ b/packages/ai/src/providers/google-gemini-cli-usage.ts @@ -0,0 +1,261 @@ +import type { + UsageAmount, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageWindow, +} from "../usage"; +import { refreshGoogleCloudToken } from "../utils/oauth/google-gemini-cli"; + +const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com"; +const CACHE_TTL_MS = 60_000; + +const GEMINI_CLI_HEADERS = { + "User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1", + "X-Goog-Api-Client": "gl-node/22.17.0", + "Client-Metadata": JSON.stringify({ + ideType: "IDE_UNSPECIFIED", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + }), +}; + +const GEMINI_TIER_MAP: Array<{ tier: string; models: string[] }> = [ + { + tier: "3-Flash", + models: ["gemini-3-flash-preview", "gemini-3-flash"], + }, + { + tier: "Flash", + models: ["gemini-2.5-flash", "gemini-2.5-flash-lite", "gemini-2.0-flash", "gemini-1.5-flash"], + }, + { + tier: "Pro", + models: ["gemini-2.5-pro", "gemini-3-pro-preview", "gemini-3-pro", "gemini-1.5-pro"], + }, +]; + +interface LoadCodeAssistResponse { + cloudaicompanionProject?: string | { id?: string }; + currentTier?: { id?: string; name?: string }; +} + +interface RetrieveUserQuotaResponse { + buckets?: Array<{ + modelId?: string; + remainingFraction?: number; + resetTime?: string; + }>; +} + +function getProjectId(payload: LoadCodeAssistResponse | undefined): string | undefined { + if (!payload) return undefined; + if (typeof payload.cloudaicompanionProject === "string") { + return payload.cloudaicompanionProject; + } + if (payload.cloudaicompanionProject && typeof payload.cloudaicompanionProject === "object") { + return payload.cloudaicompanionProject.id; + } + return undefined; +} + +function getModelTier(modelId: string): string | undefined { + for (const entry of GEMINI_TIER_MAP) { + if (entry.models.includes(modelId)) { + return entry.tier; + } + } + const normalized = modelId.toLowerCase(); + if (normalized.includes("flash")) return "Flash"; + if (normalized.includes("pro")) return "Pro"; + return undefined; +} + +function parseWindow(resetTime: string | undefined, now: number): UsageWindow | undefined { + if (!resetTime) return undefined; + const resetsAt = Date.parse(resetTime); + if (Number.isNaN(resetsAt)) return undefined; + return { + id: `reset-${resetsAt}`, + label: "Quota window", + resetsAt, + resetInMs: Math.max(0, resetsAt - now), + }; +} + +function buildAmount(remainingFraction: number | undefined): UsageAmount { + if (remainingFraction === undefined || !Number.isFinite(remainingFraction)) { + return { unit: "percent" }; + } + const remaining = Math.min(Math.max(remainingFraction, 0), 1); + const used = Math.min(Math.max(1 - remaining, 0), 1); + return { + unit: "percent", + used: Math.round(used * 1000) / 10, + remaining: Math.round(remaining * 1000) / 10, + limit: 100, + usedFraction: used, + remainingFraction: remaining, + }; +} + +async function resolveAccessToken(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + const { credential } = params; + if (credential.type !== "oauth") return undefined; + const now = ctx.now(); + if (credential.accessToken && (!credential.expiresAt || credential.expiresAt > now + 60_000)) { + return credential.accessToken; + } + if (!credential.refreshToken || !credential.projectId) return credential.accessToken; + try { + const refreshed = await refreshGoogleCloudToken(credential.refreshToken, credential.projectId); + return refreshed.access; + } catch (error) { + ctx.logger?.warn("Gemini CLI token refresh failed", { error: String(error) }); + return credential.accessToken; + } +} + +async function loadCodeAssist( + params: UsageFetchParams, + ctx: UsageFetchContext, + accessToken: string, + baseUrl: string, + projectId?: string, +): Promise { + const response = await ctx.fetch(`${baseUrl}/v1internal:loadCodeAssist`, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + ...GEMINI_CLI_HEADERS, + }, + body: JSON.stringify({ + ...(projectId ? { cloudaicompanionProject: projectId } : {}), + metadata: { + ideType: "IDE_UNSPECIFIED", + platform: "PLATFORM_UNSPECIFIED", + pluginType: "GEMINI", + }, + }), + signal: params.signal, + }); + + if (!response.ok) { + const errorText = await response.text(); + ctx.logger?.warn("Gemini CLI loadCodeAssist failed", { + status: response.status, + error: errorText, + }); + return undefined; + } + + return (await response.json()) as LoadCodeAssistResponse; +} + +async function fetchQuota( + params: UsageFetchParams, + ctx: UsageFetchContext, + accessToken: string, + baseUrl: string, + projectId?: string, +): Promise { + const response = await ctx.fetch(`${baseUrl}/v1internal:retrieveUserQuota`, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + ...GEMINI_CLI_HEADERS, + }, + body: JSON.stringify(projectId ? { project: projectId } : {}), + signal: params.signal, + }); + + if (!response.ok) { + const errorText = await response.text(); + ctx.logger?.warn("Gemini CLI retrieveUserQuota failed", { + status: response.status, + error: errorText, + }); + return undefined; + } + + return (await response.json()) as RetrieveUserQuotaResponse; +} + +export const googleGeminiCliUsageProvider: UsageProvider = { + id: "google-gemini-cli", + supports: ({ credential }) => credential.type === "oauth" && !!credential.accessToken, + async fetchUsage(params, ctx) { + const { credential } = params; + if (credential.type !== "oauth") { + return null; + } + const accessToken = await resolveAccessToken(params, ctx); + if (!accessToken) { + return null; + } + + const now = ctx.now(); + const baseUrl = (params.baseUrl?.trim() || DEFAULT_ENDPOINT).replace(/\/$/, ""); + const cacheKey = `usage:${params.provider}:${credential.accountId ?? credential.email ?? "default"}:${baseUrl}:${ + credential.projectId ?? "default" + }`; + const cached = await ctx.cache.get(cacheKey); + if (cached && cached.expiresAt > now) { + return cached.value; + } + + const loadResponse = await loadCodeAssist(params, ctx, accessToken, baseUrl, credential.projectId); + const projectId = credential.projectId ?? getProjectId(loadResponse); + const quotaResponse = await fetchQuota(params, ctx, accessToken, baseUrl, projectId); + if (!quotaResponse) { + const entry = { value: null, expiresAt: now + CACHE_TTL_MS }; + await ctx.cache.set(cacheKey, entry); + return null; + } + + const limits: UsageLimit[] = []; + const buckets = quotaResponse.buckets ?? []; + + buckets.forEach((bucket, index) => { + const modelId = bucket.modelId; + const window = parseWindow(bucket.resetTime, now); + const amount = buildAmount(bucket.remainingFraction); + const tier = modelId ? getModelTier(modelId) : undefined; + const label = modelId ? `Gemini ${modelId}` : "Gemini quota"; + const id = `${modelId ?? "unknown"}:${window?.id ?? index}`; + + limits.push({ + id, + label, + scope: { + provider: params.provider, + accountId: credential.accountId, + projectId, + modelId, + tier, + windowId: window?.id, + }, + window, + amount, + }); + }); + + const report: UsageReport = { + provider: params.provider, + fetchedAt: now, + limits, + metadata: { + currentTierId: loadResponse?.currentTier?.id, + currentTierName: loadResponse?.currentTier?.name, + }, + raw: quotaResponse, + }; + + await ctx.cache.set(cacheKey, { value: report, expiresAt: now + CACHE_TTL_MS }); + return report; + }, +}; diff --git a/packages/ai/src/usage.ts b/packages/ai/src/usage.ts new file mode 100644 index 000000000..8c424a472 --- /dev/null +++ b/packages/ai/src/usage.ts @@ -0,0 +1,133 @@ +/** + * Usage reporting types for provider quota/limit endpoints. + * + * Provides a normalized schema to represent multiple limit windows, model tiers, + * and shared quotas across providers. + */ + +import type { Provider } from "./types"; + +export type UsageUnit = "percent" | "tokens" | "requests" | "usd" | "minutes" | "bytes" | "unknown"; + +export type UsageStatus = "ok" | "warning" | "exhausted" | "unknown"; + +/** Time window for a limit (e.g. 5h, 7d, monthly). */ +export interface UsageWindow { + /** Stable identifier (e.g. "5h", "7d", "monthly"). */ + id: string; + /** Human label (e.g. "5 Hour", "7 Day"). */ + label: string; + /** Window duration in milliseconds, when known. */ + durationMs?: number; + /** Absolute reset timestamp in milliseconds since epoch. */ + resetsAt?: number; + /** Relative reset time in milliseconds, computed at fetch time. */ + resetInMs?: number; +} + +/** Quantitative usage data. */ +export interface UsageAmount { + /** Amount used in the given unit. */ + used?: number; + /** Maximum limit in the given unit. */ + limit?: number; + /** Remaining amount in the given unit. */ + remaining?: number; + /** Fraction used (0..1). */ + usedFraction?: number; + /** Fraction remaining (0..1). */ + remainingFraction?: number; + /** Unit for the amounts (percent, tokens, etc.). */ + unit: UsageUnit; +} + +/** Scope metadata describing what the limit applies to. */ +export interface UsageScope { + provider: Provider; + accountId?: string; + projectId?: string; + orgId?: string; + modelId?: string; + tier?: string; + windowId?: string; + shared?: boolean; +} + +/** Normalized limit entry for a single window or quota bucket. */ +export interface UsageLimit { + /** Stable identifier for this limit entry. */ + id: string; + /** Human label for display. */ + label: string; + scope: UsageScope; + window?: UsageWindow; + amount: UsageAmount; + status?: UsageStatus; + notes?: string[]; +} + +/** Aggregated usage report for a provider. */ +export interface UsageReport { + provider: Provider; + fetchedAt: number; + limits: UsageLimit[]; + metadata?: Record; + raw?: unknown; +} + +/** Cache entry for usage reports with absolute expiry. */ +export interface UsageCacheEntry { + value: UsageReport | null; + expiresAt: number; +} + +/** Dependency-injected cache store for usage responses. */ +export interface UsageCache { + get(key: string): UsageCacheEntry | undefined | Promise; + set(key: string, entry: UsageCacheEntry): void | Promise; + delete?(key: string): void | Promise; + cleanup?(): void | Promise; +} + +/** Optional logger for usage fetchers. */ +export interface UsageLogger { + debug(message: string, meta?: Record): void; + warn(message: string, meta?: Record): void; +} + +/** Credential bundle for usage endpoints. */ +export interface UsageCredential { + type: "api_key" | "oauth"; + apiKey?: string; + accessToken?: string; + refreshToken?: string; + expiresAt?: number; + accountId?: string; + projectId?: string; + email?: string; + enterpriseUrl?: string; + metadata?: Record; +} + +/** Parameters provided to a usage fetcher. */ +export interface UsageFetchParams { + provider: Provider; + credential: UsageCredential; + baseUrl?: string; + signal?: AbortSignal; +} + +/** Shared runtime utilities for fetchers. */ +export interface UsageFetchContext { + cache: UsageCache; + fetch: typeof fetch; + now: () => number; + logger?: UsageLogger; +} + +/** Provider implementation for fetching usage information. */ +export interface UsageProvider { + id: Provider; + fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise; + supports?(params: UsageFetchParams): boolean; +} diff --git a/packages/ai/src/usage/claude.ts b/packages/ai/src/usage/claude.ts new file mode 100644 index 000000000..6ea9d478c --- /dev/null +++ b/packages/ai/src/usage/claude.ts @@ -0,0 +1,321 @@ +import type { + UsageAmount, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, + UsageWindow, +} from "../usage"; + +const DEFAULT_ENDPOINT = "https://api.anthropic.com/api/oauth"; +const DEFAULT_CACHE_TTL_MS = 60_000; +const FIVE_HOURS_MS = 5 * 60 * 60 * 1000; +const SEVEN_DAYS_MS = 7 * 24 * 60 * 60 * 1000; +const MAX_RETRIES = 3; +const BASE_RETRY_DELAY_MS = 500; + +const CLAUDE_HEADERS = { + accept: "application/json, text/plain, */*", + "accept-encoding": "gzip, compress, deflate, br", + "anthropic-beta": "oauth-2025-04-20", + "content-type": "application/json", + "user-agent": "claude-code/2.0.20", + connection: "keep-alive", +} as const; + +function normalizeClaudeBaseUrl(baseUrl?: string): string { + if (!baseUrl || !baseUrl.trim()) return DEFAULT_ENDPOINT; + const trimmed = baseUrl.trim().replace(/\/+$/, ""); + const lower = trimmed.toLowerCase(); + if (lower.endsWith("/api/oauth")) return trimmed; + let url: URL; + try { + url = new URL(trimmed); + } catch { + return DEFAULT_ENDPOINT; + } + let path = url.pathname.replace(/\/+$/, ""); + if (path === "/") path = ""; + if (path.toLowerCase().endsWith("/v1")) { + path = path.slice(0, -3); + } + if (!path) return `${url.origin}/api/oauth`; + return `${url.origin}${path}/api/oauth`; +} + +interface ClaudeUsageBucket { + utilization?: number; + resets_at?: string; +} + +interface ParsedUsageBucket { + utilization?: number; + resetsAt?: number; +} + +interface ClaudeUsageResponse { + five_hour?: ClaudeUsageBucket | null; + seven_day?: ClaudeUsageBucket | null; + seven_day_opus?: ClaudeUsageBucket | null; + seven_day_sonnet?: ClaudeUsageBucket | null; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function toNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +function parseIsoTime(value: string | undefined): number | undefined { + if (!value) return undefined; + const parsed = Date.parse(value); + return Number.isFinite(parsed) ? parsed : undefined; +} + +function parseBucket(bucket: unknown): ParsedUsageBucket | undefined { + if (!isRecord(bucket)) return undefined; + const utilization = toNumber(bucket.utilization); + const resetsAt = parseIsoTime(typeof bucket.resets_at === "string" ? bucket.resets_at : undefined); + if (utilization === undefined && resetsAt === undefined) { + if ("utilization" in bucket || "resets_at" in bucket) { + return { utilization: 0, resetsAt: undefined }; + } + return undefined; + } + return { utilization, resetsAt }; +} + +function hasUsageData(payload: ClaudeUsageResponse): boolean { + return Boolean(payload.five_hour || payload.seven_day || payload.seven_day_opus || payload.seven_day_sonnet); +} + +async function fetchUsagePayload( + url: string, + headers: Record, + ctx: UsageFetchContext, + signal?: AbortSignal, +): Promise { + let lastPayload: ClaudeUsageResponse | null = null; + for (let attempt = 0; attempt < MAX_RETRIES; attempt++) { + try { + const response = await ctx.fetch(url, { headers, signal }); + if (!response.ok) { + ctx.logger?.warn("Claude usage fetch failed", { status: response.status, statusText: response.statusText }); + return null; + } + const payload = (await response.json()) as ClaudeUsageResponse; + lastPayload = payload; + if (payload && isRecord(payload) && hasUsageData(payload)) { + return payload; + } + } catch (error) { + ctx.logger?.warn("Claude usage fetch error", { error: String(error) }); + return null; + } + + if (attempt < MAX_RETRIES - 1) { + await Bun.sleep(BASE_RETRY_DELAY_MS * 2 ** attempt); + } + } + + return lastPayload; +} + +function buildUsageAmount(utilization: number | undefined): UsageAmount | undefined { + if (utilization === undefined) return undefined; + const clamped = Math.min(Math.max(utilization, 0), 100); + const usedFraction = clamped / 100; + return { + used: clamped, + limit: 100, + remaining: Math.max(0, 100 - clamped), + usedFraction, + remainingFraction: Math.max(0, 1 - usedFraction), + unit: "percent", + }; +} + +function buildUsageWindow( + id: string, + label: string, + durationMs: number, + resetsAt: number | undefined, + now: number, +): UsageWindow { + const resolvedResetAt = resetsAt ?? now + durationMs; + const resetInMs = Math.max(0, resolvedResetAt - now); + return { + id, + label, + durationMs, + resetsAt: resolvedResetAt, + resetInMs, + }; +} + +function buildUsageStatus(usedFraction: number | undefined): UsageStatus | undefined { + if (usedFraction === undefined) return undefined; + if (usedFraction >= 1) return "exhausted"; + if (usedFraction >= 0.9) return "warning"; + return "ok"; +} + +function buildUsageLimit(args: { + id: string; + label: string; + windowId: string; + windowLabel: string; + durationMs: number; + bucket: ParsedUsageBucket | undefined; + provider: "anthropic"; + tier?: string; + shared?: boolean; + now: number; +}): UsageLimit | null { + if (!args.bucket) return null; + const amount = buildUsageAmount(args.bucket.utilization); + if (!amount) return null; + const window = buildUsageWindow(args.windowId, args.windowLabel, args.durationMs, args.bucket.resetsAt, args.now); + return { + id: args.id, + label: args.label, + scope: { + provider: args.provider, + windowId: args.windowId, + tier: args.tier, + shared: args.shared, + }, + window, + amount, + status: buildUsageStatus(amount.usedFraction), + }; +} + +function buildCacheKey(params: UsageFetchParams): string { + const credential = params.credential; + const account = credential.accountId ?? credential.email ?? "unknown"; + const token = credential.accessToken ?? credential.refreshToken; + const fingerprint = token && typeof token === "string" ? Bun.hash(token).toString(16) : "anonymous"; + const baseUrl = params.baseUrl ?? DEFAULT_ENDPOINT; + return `usage:${params.provider}:${account}:${fingerprint}:${baseUrl}`; +} + +function resolveCacheExpiry(now: number, limits: UsageLimit[]): number { + const earliestReset = limits + .map((limit) => limit.window?.resetsAt) + .filter((value): value is number => typeof value === "number" && Number.isFinite(value)) + .reduce((min, value) => (min === undefined ? value : Math.min(min, value)), undefined as number | undefined); + const exhausted = limits.some((limit) => limit.status === "exhausted"); + if (earliestReset === undefined) return now + DEFAULT_CACHE_TTL_MS; + if (exhausted) return earliestReset; + return Math.min(now + DEFAULT_CACHE_TTL_MS, earliestReset); +} + +async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + if (params.provider !== "anthropic") return null; + const credential = params.credential; + if (credential.type !== "oauth" || !credential.accessToken) return null; + + const cacheKey = buildCacheKey(params); + const cachedEntry = await ctx.cache.get(cacheKey); + const now = ctx.now(); + if (cachedEntry && cachedEntry.expiresAt > now) { + return cachedEntry.value; + } + const cachedValue = cachedEntry?.value ?? null; + + const baseUrl = normalizeClaudeBaseUrl(params.baseUrl); + const url = `${baseUrl}/usage`; + const headers: Record = { + ...CLAUDE_HEADERS, + authorization: `Bearer ${credential.accessToken}`, + }; + + const payload = await fetchUsagePayload(url, headers, ctx, params.signal); + if (!payload || !isRecord(payload)) return cachedValue; + + const fiveHour = parseBucket(payload.five_hour); + const sevenDay = parseBucket(payload.seven_day); + const sevenDayOpus = parseBucket(payload.seven_day_opus); + const sevenDaySonnet = parseBucket(payload.seven_day_sonnet); + + const limits = [ + buildUsageLimit({ + id: "anthropic:5h", + label: "Claude 5 Hour", + windowId: "5h", + windowLabel: "5 Hour", + durationMs: FIVE_HOURS_MS, + bucket: fiveHour, + provider: "anthropic", + shared: true, + now, + }), + buildUsageLimit({ + id: "anthropic:7d", + label: "Claude 7 Day", + windowId: "7d", + windowLabel: "7 Day", + durationMs: SEVEN_DAYS_MS, + bucket: sevenDay, + provider: "anthropic", + shared: true, + now, + }), + buildUsageLimit({ + id: "anthropic:7d:opus", + label: "Claude 7 Day (Opus)", + windowId: "7d", + windowLabel: "7 Day", + durationMs: SEVEN_DAYS_MS, + bucket: sevenDayOpus, + provider: "anthropic", + tier: "opus", + now, + }), + buildUsageLimit({ + id: "anthropic:7d:sonnet", + label: "Claude 7 Day (Sonnet)", + windowId: "7d", + windowLabel: "7 Day", + durationMs: SEVEN_DAYS_MS, + bucket: sevenDaySonnet, + provider: "anthropic", + tier: "sonnet", + now, + }), + ].filter((limit): limit is UsageLimit => limit !== null); + + if (limits.length === 0) return cachedValue; + + const report: UsageReport = { + provider: params.provider, + fetchedAt: now, + limits, + metadata: { + accountId: credential.accountId, + email: credential.email, + endpoint: url, + }, + raw: payload, + }; + + const expiresAt = resolveCacheExpiry(now, limits); + await ctx.cache.set(cacheKey, { value: report, expiresAt }); + return report; +} + +export const claudeUsageProvider: UsageProvider = { + id: "anthropic", + fetchUsage: fetchClaudeUsage, + supports: (params) => params.provider === "anthropic" && params.credential.type === "oauth", +}; diff --git a/packages/ai/src/usage/github-copilot.ts b/packages/ai/src/usage/github-copilot.ts new file mode 100644 index 000000000..c98d792d8 --- /dev/null +++ b/packages/ai/src/usage/github-copilot.ts @@ -0,0 +1,502 @@ +/** + * GitHub Copilot usage provider. + * + * Normalizes Copilot quota usage into the shared UsageReport schema. + */ + +import type { + UsageAmount, + UsageCacheEntry, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, + UsageWindow, +} from "../usage"; + +const COPILOT_HEADERS = { + "User-Agent": "GitHubCopilotChat/0.35.0", + "Editor-Version": "vscode/1.107.0", + "Editor-Plugin-Version": "copilot-chat/0.35.0", + "Copilot-Integration-Id": "vscode-chat", +} as const; + +const DEFAULT_CACHE_TTL_MS = 60_000; +const MAX_CACHE_TTL_MS = 300_000; + +type CopilotQuotaDetail = { + entitlement: number; + overage_count: number; + overage_permitted: boolean; + percent_remaining: number; + quota_id: string; + quota_remaining: number; + remaining: number; + unlimited: boolean; +}; + +type CopilotQuotaSnapshots = { + chat?: CopilotQuotaDetail; + completions?: CopilotQuotaDetail; + premium_interactions?: CopilotQuotaDetail; +}; + +type CopilotUsageResponse = { + copilot_plan: string; + quota_reset_date: string; + quota_snapshots: CopilotQuotaSnapshots; +}; + +type CopilotTokenResponse = { + token: string; + expires_at: number; +}; + +type BillingUsageItem = { + product: string; + sku: string; + model?: string; + unitType: string; + grossQuantity: number; + netQuantity: number; + limit?: number; +}; + +type BillingUsageResponse = { + timePeriod: { year: number; month?: number }; + user: string; + usageItems: BillingUsageItem[]; +}; + +function toNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +function toBoolean(value: unknown): boolean | undefined { + return typeof value === "boolean" ? value : undefined; +} + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function resolveGitHubApiBaseUrl(params: UsageFetchParams): string { + const baseUrl = params.baseUrl?.replace(/\/$/, ""); + if (baseUrl && !baseUrl.includes("githubcopilot.com")) return baseUrl; + const enterpriseUrl = params.credential.enterpriseUrl?.trim(); + if (!enterpriseUrl) return "https://api.github.com"; + if (enterpriseUrl.startsWith("http://") || enterpriseUrl.startsWith("https://")) { + return enterpriseUrl.replace(/\/$/, ""); + } + if (enterpriseUrl.startsWith("api.")) { + return `https://${enterpriseUrl}`; + } + return `https://api.${enterpriseUrl}`; +} + +function resolveCopilotApiBaseUrl(params: UsageFetchParams): string { + if (params.baseUrl) return params.baseUrl.replace(/\/$/, ""); + const enterpriseUrl = params.credential.enterpriseUrl?.trim(); + if (enterpriseUrl) return `https://api.${enterpriseUrl}`; + return "https://api.individual.githubcopilot.com"; +} + +function buildCacheKey(params: UsageFetchParams): string { + const parts: string[] = [params.provider]; + const { credential } = params; + if (credential.accountId) parts.push(credential.accountId); + if (credential.email) parts.push(credential.email); + const token = + credential.apiKey || credential.accessToken || credential.refreshToken || credential.metadata?.username; + if (token && typeof token === "string") { + const fingerprint = Bun.hash(token).toString(16); + parts.push(fingerprint); + } + return parts.join(":"); +} + +function buildWindow(resetDate: string | undefined, now: number): UsageWindow | undefined { + if (!resetDate) return undefined; + const resetAt = Date.parse(resetDate); + if (!Number.isFinite(resetAt)) return undefined; + return { + id: "monthly", + label: "Monthly", + resetsAt: resetAt, + resetInMs: resetAt - now, + }; +} + +function buildAmount(used: number | undefined, limit: number | undefined, unit: UsageAmount["unit"]): UsageAmount { + const safeLimit = limit !== undefined && Number.isFinite(limit) ? limit : undefined; + const safeUsed = used !== undefined && Number.isFinite(used) ? used : undefined; + const remaining = safeLimit !== undefined && safeUsed !== undefined ? Math.max(0, safeLimit - safeUsed) : undefined; + const usedFraction = + safeLimit !== undefined && safeUsed !== undefined && safeLimit > 0 ? safeUsed / safeLimit : undefined; + const remainingFraction = + safeLimit !== undefined && remaining !== undefined && safeLimit > 0 ? remaining / safeLimit : undefined; + return { + used: safeUsed, + limit: safeLimit, + remaining, + usedFraction, + remainingFraction, + unit, + }; +} + +function deriveStatus(amount: UsageAmount, unlimited: boolean): UsageStatus { + if (unlimited) return "ok"; + if (amount.remainingFraction === undefined) return "unknown"; + if (amount.remainingFraction <= 0) return "exhausted"; + if (amount.remainingFraction <= 0.1) return "warning"; + return "ok"; +} + +function parseQuotaDetail(value: unknown): CopilotQuotaDetail | null { + if (!isRecord(value)) return null; + const entitlement = toNumber(value.entitlement); + const remaining = toNumber(value.remaining); + const percentRemaining = toNumber(value.percent_remaining); + const unlimited = toBoolean(value.unlimited); + if ( + entitlement === undefined || + remaining === undefined || + percentRemaining === undefined || + unlimited === undefined + ) { + return null; + } + const overageCount = toNumber(value.overage_count) ?? 0; + const overagePermitted = toBoolean(value.overage_permitted) ?? false; + const quotaId = typeof value.quota_id === "string" ? value.quota_id : ""; + const quotaRemaining = toNumber(value.quota_remaining) ?? remaining; + return { + entitlement, + overage_count: overageCount, + overage_permitted: overagePermitted, + percent_remaining: percentRemaining, + quota_id: quotaId, + quota_remaining: quotaRemaining, + remaining, + unlimited, + }; +} + +async function fetchJson(ctx: UsageFetchContext, url: string, init: RequestInit): Promise { + const response = await ctx.fetch(url, init); + if (!response.ok) { + const text = await response.text(); + throw new Error(`${response.status} ${response.statusText}: ${text}`); + } + return response.json(); +} + +async function resolveGitHubUsername( + ctx: UsageFetchContext, + baseUrl: string, + token: string, + signal?: AbortSignal, +): Promise { + try { + const data = await fetchJson(ctx, `${baseUrl}/user`, { + headers: { + Accept: "application/vnd.github+json", + Authorization: `Bearer ${token}`, + "X-GitHub-Api-Version": "2022-11-28", + }, + signal, + }); + if (!isRecord(data)) return undefined; + return typeof data.login === "string" ? data.login : undefined; + } catch { + return undefined; + } +} + +async function exchangeForCopilotToken( + ctx: UsageFetchContext, + baseUrl: string, + oauthToken: string, + signal?: AbortSignal, +): Promise { + try { + const data = await fetchJson(ctx, `${baseUrl}/copilot_internal/v2/token`, { + headers: { + Accept: "application/json", + Authorization: `Bearer ${oauthToken}`, + ...COPILOT_HEADERS, + }, + signal, + }); + + if (!isRecord(data)) return null; + const token = typeof data.token === "string" ? data.token : undefined; + const expiresAt = toNumber(data.expires_at); + if (!token || !expiresAt) return null; + return { token, expires_at: expiresAt }; + } catch { + return null; + } +} + +async function fetchInternalUsage( + ctx: UsageFetchContext, + baseUrl: string, + oauthToken: string, + accessToken: string | undefined, + expiresAt: number | undefined, + signal?: AbortSignal, +): Promise { + const requestWithToken = async (token: string, legacy: boolean) => { + const headers: Record = { + "Content-Type": "application/json", + Accept: "application/json", + Authorization: legacy ? `token ${token}` : `Bearer ${token}`, + ...COPILOT_HEADERS, + }; + const data = await fetchJson(ctx, `${baseUrl}/copilot_internal/user`, { headers, signal }); + if (!isRecord(data)) throw new Error("Invalid Copilot usage response"); + return data as CopilotUsageResponse; + }; + + const now = ctx.now(); + if (accessToken && expiresAt && accessToken !== oauthToken && expiresAt > now) { + try { + return await requestWithToken(accessToken, false); + } catch { + // Ignore and try other strategies. + } + } + + try { + return await requestWithToken(oauthToken, true); + } catch { + const exchanged = await exchangeForCopilotToken(ctx, baseUrl, oauthToken, signal); + if (!exchanged) throw new Error("Copilot usage token exchange failed"); + return requestWithToken(exchanged.token, false); + } +} + +async function fetchBillingUsage( + ctx: UsageFetchContext, + baseUrl: string, + username: string, + token: string, + signal?: AbortSignal, +): Promise { + const data = await fetchJson( + ctx, + `${baseUrl}/users/${encodeURIComponent(username)}/settings/billing/premium_request/usage`, + { + headers: { + Accept: "application/vnd.github+json", + Authorization: `Bearer ${token}`, + "X-GitHub-Api-Version": "2022-11-28", + }, + signal, + }, + ); + + if (!isRecord(data)) throw new Error("Invalid Copilot billing usage response"); + return data as BillingUsageResponse; +} + +function buildLimitFromQuota( + key: string, + label: string, + quota: CopilotQuotaDetail, + plan: string, + window: UsageWindow | undefined, +): UsageLimit { + const used = quota.unlimited ? undefined : Math.max(0, quota.entitlement - quota.remaining); + const limit = quota.unlimited ? undefined : quota.entitlement; + const amount = buildAmount(used, limit, "requests"); + const notes: string[] = []; + if (quota.unlimited) notes.push("Unlimited"); + if (quota.overage_count > 0) { + notes.push(`Overage requests: ${quota.overage_count}`); + } + return { + id: `copilot:${key}`, + label, + scope: { + provider: "github-copilot", + tier: plan, + windowId: window?.id, + }, + window, + amount, + status: deriveStatus(amount, quota.unlimited), + notes: notes.length > 0 ? notes : undefined, + }; +} + +function normalizeQuotaSnapshots( + data: CopilotUsageResponse, + now: number, +): { limits: UsageLimit[]; window?: UsageWindow } { + const window = buildWindow(data.quota_reset_date, now); + const snapshots = data.quota_snapshots ?? {}; + const limits: UsageLimit[] = []; + const premium = parseQuotaDetail(snapshots.premium_interactions); + if (premium) { + limits.push(buildLimitFromQuota("premium", "Premium Requests", premium, data.copilot_plan, window)); + } + const chat = parseQuotaDetail(snapshots.chat); + if (chat && !chat.unlimited) { + limits.push(buildLimitFromQuota("chat", "Chat Requests", chat, data.copilot_plan, window)); + } + const completions = parseQuotaDetail(snapshots.completions); + if (completions && !completions.unlimited) { + limits.push(buildLimitFromQuota("completions", "Completions", completions, data.copilot_plan, window)); + } + return { limits, window }; +} + +function normalizeBillingUsage(data: BillingUsageResponse): UsageLimit[] { + const limits: UsageLimit[] = []; + const periodLabel = data.timePeriod.month + ? `${data.timePeriod.year}-${String(data.timePeriod.month).padStart(2, "0")}` + : `${data.timePeriod.year}`; + const window: UsageWindow = { + id: "billing-period", + label: periodLabel, + }; + + const premiumItems = data.usageItems.filter( + (item) => item.sku === "Copilot Premium Request" || item.sku.includes("Premium"), + ); + const totalUsed = premiumItems.reduce((sum, item) => sum + item.grossQuantity, 0); + const totalLimit = premiumItems.reduce((sum, item) => sum + (item.limit ?? 0), 0) || undefined; + const totalAmount = buildAmount(totalUsed, totalLimit, "requests"); + limits.push({ + id: "copilot:premium", + label: "Premium Requests", + scope: { + provider: "github-copilot", + accountId: data.user, + windowId: window.id, + }, + window, + amount: totalAmount, + status: deriveStatus(totalAmount, false), + }); + + for (const item of data.usageItems) { + if (!item.model) continue; + if (item.grossQuantity <= 0) continue; + const amount = buildAmount(item.grossQuantity, item.limit, "requests"); + limits.push({ + id: `copilot:model:${item.model}`, + label: `Model ${item.model}`, + scope: { + provider: "github-copilot", + accountId: data.user, + modelId: item.model, + windowId: window.id, + }, + window, + amount, + status: deriveStatus(amount, false), + }); + } + + return limits; +} + +function resolveCacheTtl(now: number, report: UsageReport | null): UsageCacheEntry["expiresAt"] { + if (!report) return now + DEFAULT_CACHE_TTL_MS; + const resetInMs = report.limits + .map((limit) => limit.window?.resetInMs) + .find((value): value is number => typeof value === "number" && Number.isFinite(value)); + if (!resetInMs || resetInMs <= 0) return now + DEFAULT_CACHE_TTL_MS; + return now + Math.min(MAX_CACHE_TTL_MS, resetInMs); +} + +export const githubCopilotUsageProvider: UsageProvider = { + id: "github-copilot", + supports: ({ provider, credential }) => { + if (provider !== "github-copilot") return false; + if (credential.type === "oauth") { + return Boolean(credential.refreshToken || credential.accessToken); + } + return Boolean(credential.apiKey); + }, + fetchUsage: async (params, ctx) => { + if (!githubCopilotUsageProvider.supports?.(params)) return null; + const now = ctx.now(); + const cacheKey = buildCacheKey(params); + const cached = await ctx.cache.get(cacheKey); + if (cached && cached.expiresAt > now) return cached.value; + + const baseUrl = + params.credential.type === "api_key" ? resolveGitHubApiBaseUrl(params) : resolveCopilotApiBaseUrl(params); + let report: UsageReport | null = null; + + if (params.credential.type === "api_key") { + let username = + params.credential.accountId || params.credential.metadata?.username || params.credential.metadata?.user; + if ((!username || typeof username !== "string" || !username.trim()) && params.credential.apiKey) { + username = await resolveGitHubUsername(ctx, baseUrl, params.credential.apiKey, params.signal); + } + if (typeof username !== "string" || !username.trim()) { + ctx.logger?.warn("Copilot usage requires username for billing API", { provider: params.provider }); + } else if (params.credential.apiKey) { + try { + const billing = await fetchBillingUsage(ctx, baseUrl, username, params.credential.apiKey, params.signal); + report = { + provider: "github-copilot", + fetchedAt: now, + limits: normalizeBillingUsage(billing), + metadata: { + account: billing.user, + period: billing.timePeriod, + }, + }; + } catch (error) { + ctx.logger?.warn("Copilot usage fetch failed", { error: String(error) }); + } + } + } else { + const { refreshToken, accessToken, expiresAt } = params.credential; + if (!refreshToken && !accessToken) return null; + const oauthToken = refreshToken || accessToken; + if (!oauthToken) return null; + try { + const usage = await fetchInternalUsage( + ctx, + baseUrl, + oauthToken, + accessToken ?? undefined, + expiresAt ?? undefined, + params.signal, + ); + const normalized = normalizeQuotaSnapshots(usage, now); + report = { + provider: "github-copilot", + fetchedAt: now, + limits: normalized.limits, + metadata: { + plan: usage.copilot_plan, + quotaResetDate: usage.quota_reset_date, + }, + raw: usage, + }; + } catch (error) { + ctx.logger?.warn("Copilot usage fetch failed", { error: String(error) }); + } + } + + const expiresAt = resolveCacheTtl(now, report); + await ctx.cache.set(cacheKey, { value: report, expiresAt }); + return report; + }, +}; diff --git a/packages/ai/src/usage/google-antigravity.ts b/packages/ai/src/usage/google-antigravity.ts new file mode 100644 index 000000000..c9b78086b --- /dev/null +++ b/packages/ai/src/usage/google-antigravity.ts @@ -0,0 +1,218 @@ +import type { + UsageAmount, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, + UsageWindow, +} from "../usage"; +import { refreshAntigravityToken } from "../utils/oauth/google-antigravity"; + +interface AntigravityQuotaInfo { + remainingFraction?: number; + resetTime?: string; + tier?: string; + windowId?: string; + windowLabel?: string; +} + +interface AntigravityModelInfo { + displayName?: string; + quotaInfo?: AntigravityQuotaInfo | AntigravityQuotaInfo[]; + quotaInfos?: AntigravityQuotaInfo[]; + quotaInfoByTier?: Record; +} + +interface AntigravityUsageResponse { + models: Record; +} + +const DEFAULT_ENDPOINT = "https://daily-cloudcode-pa.googleapis.com"; +const FETCH_AVAILABLE_MODELS_PATH = "/v1internal:fetchAvailableModels"; +const USER_AGENT = "antigravity/1.11.9 windows/amd64"; +const DEFAULT_CACHE_TTL_MS = 60_000; + +function clampFraction(value: number | undefined): number | undefined { + if (value === undefined || !Number.isFinite(value)) return undefined; + if (value < 0) return 0; + if (value > 1) return 1; + return value; +} + +function getUsageStatus(remainingFraction: number | undefined): UsageStatus | undefined { + if (remainingFraction === undefined) return "unknown"; + if (remainingFraction <= 0) return "exhausted"; + if (remainingFraction <= 0.1) return "warning"; + return "ok"; +} + +function parseWindow(info: AntigravityQuotaInfo, nowMs: number): UsageWindow | undefined { + if (!info.resetTime) return undefined; + const resetAt = Date.parse(info.resetTime); + if (!Number.isFinite(resetAt)) return undefined; + return { + id: info.windowId ?? "default", + label: info.windowLabel ?? "Default", + resetsAt: resetAt, + resetInMs: Math.max(0, resetAt - nowMs), + }; +} + +function buildAmount(info: AntigravityQuotaInfo): UsageAmount { + const remainingFraction = clampFraction(info.remainingFraction); + const amount: UsageAmount = { unit: "percent" }; + if (remainingFraction === undefined) return amount; + const usedFraction = clampFraction(1 - remainingFraction); + amount.remainingFraction = remainingFraction; + amount.usedFraction = usedFraction; + amount.remaining = remainingFraction * 100; + amount.used = usedFraction !== undefined ? usedFraction * 100 : undefined; + amount.limit = 100; + return amount; +} + +function normalizeQuotaInfos(info: AntigravityModelInfo): AntigravityQuotaInfo[] { + const results: AntigravityQuotaInfo[] = []; + const addInfo = (value: AntigravityQuotaInfo, tier?: string) => { + results.push({ ...value, ...(tier ? { tier } : {}) }); + }; + const addArray = (values?: AntigravityQuotaInfo[]) => { + if (!values) return; + for (const value of values) addInfo(value); + }; + + if (Array.isArray(info.quotaInfo)) { + addArray(info.quotaInfo); + } else if (info.quotaInfo) { + addInfo(info.quotaInfo); + } + addArray(info.quotaInfos); + + if (info.quotaInfoByTier) { + for (const [tier, value] of Object.entries(info.quotaInfoByTier)) { + if (Array.isArray(value)) { + for (const entry of value) addInfo(entry, tier); + } else if (value) { + addInfo(value, tier); + } + } + } + + return results; +} + +function buildCacheKey(params: UsageFetchParams): string { + const credential = params.credential; + const accountPart = credential.accountId ?? credential.email ?? "unknown"; + const projectPart = credential.projectId ?? "unknown"; + return `usage:${params.provider}:${accountPart}:${projectPart}`; +} + +async function resolveAccessToken(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + const { credential } = params; + if (credential.accessToken && (!credential.expiresAt || credential.expiresAt > ctx.now() + 60_000)) { + return credential.accessToken; + } + if (!credential.refreshToken || !credential.projectId) return undefined; + try { + const refreshed = await refreshAntigravityToken(credential.refreshToken, credential.projectId); + return refreshed.access; + } catch (error) { + ctx.logger?.warn("Antigravity usage token refresh failed", { error: String(error) }); + return undefined; + } +} + +async function fetchAntigravityUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + const credential = params.credential; + if (!credential.projectId) return null; + + const cacheKey = buildCacheKey(params); + const cached = await ctx.cache.get(cacheKey); + const nowMs = ctx.now(); + if (cached && cached.expiresAt > nowMs) { + return cached.value; + } + + const accessToken = await resolveAccessToken(params, ctx); + if (!accessToken) return null; + + const baseUrl = params.baseUrl?.replace(/\/+$/, "") || DEFAULT_ENDPOINT; + const url = `${baseUrl}${FETCH_AVAILABLE_MODELS_PATH}`; + const response = await ctx.fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + "User-Agent": USER_AGENT, + }, + body: JSON.stringify({ project: credential.projectId }), + signal: params.signal, + }); + + if (!response.ok) { + ctx.logger?.warn("Antigravity usage fetch failed", { + status: response.status, + statusText: response.statusText, + }); + return null; + } + + const data = (await response.json()) as AntigravityUsageResponse; + const limits: UsageLimit[] = []; + let earliestReset: number | undefined; + + for (const [modelId, modelInfo] of Object.entries(data.models ?? {})) { + const quotaInfos = normalizeQuotaInfos(modelInfo); + for (const quotaInfo of quotaInfos) { + const amount = buildAmount(quotaInfo); + const window = parseWindow(quotaInfo, nowMs); + if (window?.resetsAt) { + earliestReset = earliestReset ? Math.min(earliestReset, window.resetsAt) : window.resetsAt; + } + const labelBase = modelInfo.displayName || modelId; + const label = quotaInfo.tier ? `${labelBase} (${quotaInfo.tier})` : labelBase; + const windowId = window?.id ?? "default"; + limits.push({ + id: `${modelId}:${quotaInfo.tier ?? "default"}:${windowId}`, + label, + scope: { + provider: params.provider, + accountId: credential.accountId, + projectId: credential.projectId, + modelId, + tier: quotaInfo.tier, + windowId, + }, + window, + amount, + status: getUsageStatus(amount.remainingFraction), + }); + } + } + + const report: UsageReport = { + provider: params.provider, + fetchedAt: nowMs, + limits, + metadata: { + endpoint: url, + projectId: credential.projectId, + }, + raw: data, + }; + + const expiresAt = earliestReset + ? Math.min(earliestReset, nowMs + DEFAULT_CACHE_TTL_MS) + : nowMs + DEFAULT_CACHE_TTL_MS; + await ctx.cache.set(cacheKey, { value: report, expiresAt }); + return report; +} + +export const antigravityUsageProvider: UsageProvider = { + id: "google-antigravity", + fetchUsage: fetchAntigravityUsage, + supports: (params) => params.provider === "google-antigravity", +}; diff --git a/packages/ai/src/usage/openai-codex.ts b/packages/ai/src/usage/openai-codex.ts new file mode 100644 index 000000000..ddfb6293d --- /dev/null +++ b/packages/ai/src/usage/openai-codex.ts @@ -0,0 +1,393 @@ +import { Buffer } from "node:buffer"; +import { CODEX_BASE_URL } from "../providers/openai-codex/constants"; +import type { + UsageAmount, + UsageCache, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageWindow, +} from "../usage"; + +const CODEX_USAGE_PATH = "wham/usage"; +const DEFAULT_CACHE_TTL_MS = 60_000; +const JWT_AUTH_CLAIM = "https://api.openai.com/auth"; +const JWT_PROFILE_CLAIM = "https://api.openai.com/profile"; + +interface CodexUsageWindowPayload { + used_percent?: number; + limit_window_seconds?: number; + reset_after_seconds?: number; + reset_at?: number; +} + +interface CodexUsageRateLimitPayload { + allowed?: boolean; + limit_reached?: boolean; + primary_window?: CodexUsageWindowPayload | null; + secondary_window?: CodexUsageWindowPayload | null; +} + +interface CodexUsagePayload { + plan_type?: string; + rate_limit?: CodexUsageRateLimitPayload | null; +} + +interface ParsedUsageWindow { + usedPercent?: number; + limitWindowSeconds?: number; + resetAfterSeconds?: number; + resetAt?: number; +} + +interface ParsedUsage { + planType?: string; + allowed?: boolean; + limitReached?: boolean; + primary?: ParsedUsageWindow; + secondary?: ParsedUsageWindow; + raw: CodexUsagePayload; +} + +interface JwtPayload { + [JWT_AUTH_CLAIM]?: { + chatgpt_account_id?: string; + }; + [JWT_PROFILE_CLAIM]?: { + email?: string; + }; +} + +const isRecord = (value: unknown): value is Record => + typeof value === "object" && value !== null && !Array.isArray(value); + +const toNumber = (value: unknown): number | undefined => { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string") { + const trimmed = value.trim(); + if (!trimmed) return undefined; + const parsed = Number(trimmed); + if (Number.isFinite(parsed)) return parsed; + } + return undefined; +}; + +const toBoolean = (value: unknown): boolean | undefined => { + if (typeof value === "boolean") return value; + return undefined; +}; + +function base64UrlDecode(input: string): string { + const base64 = input.replace(/-/g, "+").replace(/_/g, "/"); + const padLen = (4 - (base64.length % 4)) % 4; + const padded = base64 + "=".repeat(padLen); + return Buffer.from(padded, "base64").toString("utf8"); +} + +function parseJwt(token: string): JwtPayload | null { + const parts = token.split("."); + if (parts.length !== 3) return null; + try { + const payloadJson = base64UrlDecode(parts[1]); + return JSON.parse(payloadJson) as JwtPayload; + } catch { + return null; + } +} + +function extractAccountId(token: string | undefined): string | undefined { + if (!token) return undefined; + const payload = parseJwt(token); + return payload?.[JWT_AUTH_CLAIM]?.chatgpt_account_id ?? undefined; +} + +function extractEmail(token: string | undefined): string | undefined { + if (!token) return undefined; + const payload = parseJwt(token); + return payload?.[JWT_PROFILE_CLAIM]?.email ?? undefined; +} + +function parseUsageWindow(payload: unknown): ParsedUsageWindow | undefined { + if (!isRecord(payload)) return undefined; + const usedPercent = toNumber(payload.used_percent); + const limitWindowSeconds = toNumber(payload.limit_window_seconds); + const resetAfterSeconds = toNumber(payload.reset_after_seconds); + const resetAt = toNumber(payload.reset_at); + if ( + usedPercent === undefined && + limitWindowSeconds === undefined && + resetAfterSeconds === undefined && + resetAt === undefined + ) { + return undefined; + } + return { + usedPercent, + limitWindowSeconds, + resetAfterSeconds, + resetAt, + }; +} + +function parseUsagePayload(payload: unknown): ParsedUsage | null { + if (!isRecord(payload)) return null; + const planType = typeof payload.plan_type === "string" ? payload.plan_type : undefined; + const rateLimit = isRecord(payload.rate_limit) ? payload.rate_limit : undefined; + if (!rateLimit) return null; + const parsed: ParsedUsage = { + planType, + allowed: toBoolean(rateLimit.allowed), + limitReached: toBoolean(rateLimit.limit_reached), + primary: parseUsageWindow(rateLimit.primary_window), + secondary: parseUsageWindow(rateLimit.secondary_window), + raw: payload as CodexUsagePayload, + }; + if (!parsed.primary && !parsed.secondary && parsed.allowed === undefined && parsed.limitReached === undefined) { + return null; + } + return parsed; +} + +function normalizeCodexBaseUrl(baseUrl?: string): string { + const fallback = CODEX_BASE_URL; + const trimmed = baseUrl?.trim() ? baseUrl.trim() : fallback; + const base = trimmed.replace(/\/+$/, ""); + const lower = base.toLowerCase(); + if ( + (lower.startsWith("https://chatgpt.com") || lower.startsWith("https://chat.openai.com")) && + !lower.includes("/backend-api") + ) { + return `${base}/backend-api`; + } + return base; +} + +function buildCodexUsageUrl(baseUrl: string): string { + const normalized = baseUrl.endsWith("/") ? baseUrl : `${baseUrl}/`; + return `${normalized}${CODEX_USAGE_PATH}`; +} + +function formatWindowLabel(value: number, unit: "hour" | "day"): string { + const rounded = Math.round(value); + const suffix = rounded === 1 ? unit : `${unit}s`; + return `${rounded} ${suffix}`; +} + +function buildWindowLabel(seconds: number): { id: string; label: string } { + const daySeconds = 86_400; + if (seconds >= daySeconds) { + const days = Math.round(seconds / daySeconds); + return { id: `${days}d`, label: formatWindowLabel(days, "day") }; + } + const hours = Math.max(1, Math.round(seconds / 3600)); + return { id: `${hours}h`, label: formatWindowLabel(hours, "hour") }; +} + +function resolveResetTimes(window: ParsedUsageWindow, nowMs: number): Pick { + const resetAt = window.resetAt; + if (resetAt !== undefined) { + const resetAtMs = resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000; + if (Number.isFinite(resetAtMs)) { + return { resetsAt: resetAtMs, resetInMs: resetAtMs - nowMs }; + } + } + if (window.resetAfterSeconds !== undefined) { + const resetInMs = window.resetAfterSeconds * 1000; + return { resetsAt: nowMs + resetInMs, resetInMs }; + } + return {}; +} + +function buildUsageWindow(window: ParsedUsageWindow, key: string, nowMs: number): UsageWindow { + if (window.limitWindowSeconds !== undefined) { + const { id, label } = buildWindowLabel(window.limitWindowSeconds); + const durationMs = window.limitWindowSeconds * 1000; + return { id, label, durationMs, ...resolveResetTimes(window, nowMs) }; + } + const fallbackLabel = key === "primary" ? "Primary window" : "Secondary window"; + return { id: key, label: fallbackLabel, ...resolveResetTimes(window, nowMs) }; +} + +function buildUsageAmount(window: ParsedUsageWindow): UsageAmount { + const usedPercent = window.usedPercent; + if (usedPercent === undefined) { + return { unit: "percent" }; + } + const clamped = Math.min(Math.max(usedPercent, 0), 100); + const usedFraction = clamped / 100; + return { + used: clamped, + limit: 100, + remaining: Math.max(0, 100 - clamped), + usedFraction, + remainingFraction: Math.max(0, 1 - usedFraction), + unit: "percent", + }; +} + +function buildUsageStatus(usedFraction?: number, limitReached?: boolean): UsageLimit["status"] { + if (limitReached) return "exhausted"; + if (usedFraction === undefined) return "unknown"; + if (usedFraction >= 1) return "exhausted"; + if (usedFraction >= 0.9) return "warning"; + return "ok"; +} + +function buildUsageLimit(args: { + key: "primary" | "secondary"; + window: ParsedUsageWindow; + accountId?: string; + planType?: string; + limitReached?: boolean; + nowMs: number; +}): UsageLimit { + const usageWindow = buildUsageWindow(args.window, args.key, args.nowMs); + const amount = buildUsageAmount(args.window); + return { + id: `openai-codex:${args.key}`, + label: usageWindow.label, + scope: { + provider: "openai-codex", + accountId: args.accountId, + tier: args.planType, + windowId: usageWindow.id, + shared: true, + }, + window: usageWindow, + amount, + status: buildUsageStatus(amount.usedFraction, args.limitReached), + }; +} + +function resolveCacheExpiry(args: { report: UsageReport | null; nowMs: number }): number { + const { report, nowMs } = args; + if (!report) return nowMs + DEFAULT_CACHE_TTL_MS; + const exhausted = report.limits.some((limit) => limit.status === "exhausted"); + const resetCandidates = report.limits + .map((limit) => limit.window?.resetsAt) + .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); + const earliestReset = resetCandidates.length > 0 ? Math.min(...resetCandidates) : undefined; + if (exhausted && earliestReset) return earliestReset; + if (earliestReset) return Math.min(nowMs + DEFAULT_CACHE_TTL_MS, earliestReset); + return nowMs + DEFAULT_CACHE_TTL_MS; +} + +async function getCachedReport( + cache: UsageCache, + cacheKey: string, + nowMs: number, +): Promise { + const cached = await cache.get(cacheKey); + if (!cached) return undefined; + if (cached.expiresAt <= nowMs) return undefined; + return cached.value; +} + +async function setCachedReport( + cache: UsageCache, + cacheKey: string, + report: UsageReport | null, + expiresAt: number, +): Promise { + await cache.set(cacheKey, { value: report, expiresAt }); +} + +export const openaiCodexUsageProvider: UsageProvider = { + id: "openai-codex", + supports(params: UsageFetchParams): boolean { + return params.provider === "openai-codex" && params.credential.type === "oauth"; + }, + async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + if (params.provider !== "openai-codex") return null; + const { credential } = params; + if (credential.type !== "oauth") return null; + + const accessToken = credential.accessToken; + if (!accessToken) return null; + + const nowMs = ctx.now(); + if (credential.expiresAt !== undefined && credential.expiresAt <= nowMs) { + ctx.logger?.warn("Codex usage token expired", { provider: params.provider }); + return null; + } + + const baseUrl = normalizeCodexBaseUrl(params.baseUrl); + const accountId = credential.accountId ?? extractAccountId(accessToken); + const cacheKey = `usage:openai-codex:${accountId ?? "unknown"}:${baseUrl}`; + const cached = await getCachedReport(ctx.cache, cacheKey, nowMs); + if (cached !== undefined) return cached; + + const headers: Record = { + Authorization: `Bearer ${accessToken}`, + "User-Agent": "OpenCode-Status-Plugin/1.0", + }; + if (accountId) { + headers["ChatGPT-Account-Id"] = accountId; + } + + const url = buildCodexUsageUrl(baseUrl); + let payload: unknown; + try { + const response = await ctx.fetch(url, { headers, signal: params.signal }); + if (!response.ok) { + ctx.logger?.warn("Codex usage request failed", { status: response.status, provider: params.provider }); + return null; + } + payload = await response.json(); + } catch (error) { + ctx.logger?.warn("Codex usage request error", { provider: params.provider, error: String(error) }); + return null; + } + + const parsed = parseUsagePayload(payload); + if (!parsed) { + ctx.logger?.warn("Codex usage response invalid", { provider: params.provider }); + return null; + } + + const limits: UsageLimit[] = []; + if (parsed.primary) { + limits.push( + buildUsageLimit({ + key: "primary", + window: parsed.primary, + accountId, + planType: parsed.planType, + limitReached: parsed.limitReached, + nowMs, + }), + ); + } + if (parsed.secondary) { + limits.push( + buildUsageLimit({ + key: "secondary", + window: parsed.secondary, + accountId, + planType: parsed.planType, + limitReached: parsed.limitReached, + nowMs, + }), + ); + } + + const report: UsageReport = { + provider: "openai-codex", + fetchedAt: nowMs, + limits, + metadata: { + planType: parsed.planType, + allowed: parsed.allowed, + limitReached: parsed.limitReached, + email: credential.email ?? extractEmail(accessToken), + }, + raw: parsed.raw, + }; + + const expiresAt = resolveCacheExpiry({ report, nowMs }); + await setCachedReport(ctx.cache, cacheKey, report, expiresAt); + return report; + }, +}; diff --git a/packages/ai/src/usage/zai.ts b/packages/ai/src/usage/zai.ts new file mode 100644 index 000000000..47531b546 --- /dev/null +++ b/packages/ai/src/usage/zai.ts @@ -0,0 +1,292 @@ +import type { + UsageAmount, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, + UsageWindow, +} from "../usage"; + +const DEFAULT_ENDPOINT = "https://api.z.ai"; +const QUOTA_PATH = "/api/monitor/usage/quota/limit"; +const MODEL_USAGE_PATH = "/api/monitor/usage/model-usage"; +const DEFAULT_CACHE_TTL_MS = 60_000; +const SEVEN_DAYS_MS = 7 * 24 * 60 * 60 * 1000; + +function normalizeZaiBaseUrl(baseUrl?: string): string { + if (!baseUrl || !baseUrl.trim()) return DEFAULT_ENDPOINT; + try { + return new URL(baseUrl.trim()).origin; + } catch { + return DEFAULT_ENDPOINT; + } +} + +interface ZaiUsageLimitItem { + type?: string; + usage?: number; + currentValue?: number; + percentage?: number; + remaining?: number; + nextResetTime?: number; +} + +interface ZaiQuotaPayload { + success?: boolean; + code?: number; + msg?: string; + data?: { + limits?: ZaiUsageLimitItem[]; + }; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function toNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +function parseMillis(value: unknown): number | undefined { + const parsed = toNumber(value); + if (parsed === undefined) return undefined; + return parsed > 1_000_000_000_000 ? parsed : parsed * 1000; +} + +function parseLimitItem(value: unknown): ZaiUsageLimitItem | null { + if (!isRecord(value)) return null; + const type = typeof value.type === "string" ? value.type : undefined; + if (!type) return null; + return { + type, + usage: toNumber(value.usage), + currentValue: toNumber(value.currentValue), + percentage: toNumber(value.percentage), + remaining: toNumber(value.remaining), + nextResetTime: parseMillis(value.nextResetTime), + }; +} + +function buildUsageAmount(args: { + used: number | undefined; + limit: number | undefined; + remaining: number | undefined; + unit: UsageAmount["unit"]; + percentage?: number; +}): UsageAmount { + const usedFraction = + args.percentage !== undefined + ? Math.min(Math.max(args.percentage / 100, 0), 1) + : args.used !== undefined && args.limit !== undefined && args.limit > 0 + ? Math.min(args.used / args.limit, 1) + : undefined; + const remainingFraction = usedFraction !== undefined ? Math.max(1 - usedFraction, 0) : undefined; + return { + used: args.used, + limit: args.limit, + remaining: args.remaining, + usedFraction, + remainingFraction, + unit: args.unit, + }; +} + +function buildUsageWindow( + id: string, + label: string, + resetsAt: number | undefined, + now: number, +): UsageWindow | undefined { + if (!resetsAt) return { id, label }; + const resetInMs = Math.max(0, resetsAt - now); + return { + id, + label, + resetsAt, + resetInMs, + }; +} + +function getUsageStatus(usedFraction: number | undefined): UsageStatus | undefined { + if (usedFraction === undefined) return undefined; + if (usedFraction >= 1) return "exhausted"; + if (usedFraction >= 0.9) return "warning"; + return "ok"; +} + +function buildCacheKey(params: UsageFetchParams): string { + const credential = params.credential; + const account = credential.accountId ?? credential.email ?? "unknown"; + const token = credential.apiKey ?? credential.accessToken; + const fingerprint = token && typeof token === "string" ? Bun.hash(token).toString(16) : "anonymous"; + const baseUrl = params.baseUrl ?? DEFAULT_ENDPOINT; + return `usage:${params.provider}:${account}:${fingerprint}:${baseUrl}`; +} + +function resolveCacheExpiry(now: number, limits: UsageLimit[]): number { + const earliestReset = limits + .map((limit) => limit.window?.resetsAt) + .filter((value): value is number => typeof value === "number" && Number.isFinite(value)) + .reduce((min, value) => (min === undefined ? value : Math.min(min, value)), undefined as number | undefined); + if (!earliestReset) return now + DEFAULT_CACHE_TTL_MS; + return Math.min(earliestReset, now + DEFAULT_CACHE_TTL_MS); +} + +function formatDate(value: Date): string { + const pad = (input: number) => String(input).padStart(2, "0"); + return `${value.getFullYear()}-${pad(value.getMonth() + 1)}-${pad(value.getDate())}+${pad(value.getHours())}:${pad( + value.getMinutes(), + )}:${pad(value.getSeconds())}`; +} + +function buildModelUsageUrl(baseUrl: string, now: Date): string { + const start = new Date(now.getTime() - SEVEN_DAYS_MS); + const startTime = formatDate(start); + const endTime = formatDate(now); + return `${baseUrl}${MODEL_USAGE_PATH}?startTime=${encodeURIComponent(startTime)}&endTime=${encodeURIComponent(endTime)}`; +} + +async function fetchZaiUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + if (params.provider !== "zai") return null; + const credential = params.credential; + if (credential.type !== "api_key" || !credential.apiKey) return null; + + const cacheKey = buildCacheKey(params); + const cachedEntry = await ctx.cache.get(cacheKey); + const now = ctx.now(); + if (cachedEntry && cachedEntry.expiresAt > now) return cachedEntry.value; + + const baseUrl = normalizeZaiBaseUrl(params.baseUrl); + const url = `${baseUrl}${QUOTA_PATH}`; + const headers: Record = { + Authorization: credential.apiKey, + "Content-Type": "application/json", + "User-Agent": "OpenCode-Status-Plugin/1.0", + }; + + let payload: ZaiQuotaPayload | null = null; + try { + const response = await ctx.fetch(url, { + headers, + signal: params.signal, + }); + if (!response.ok) { + ctx.logger?.warn("ZAI usage fetch failed", { status: response.status, statusText: response.statusText }); + return null; + } + payload = (await response.json()) as ZaiQuotaPayload; + } catch (error) { + ctx.logger?.warn("ZAI usage fetch error", { error: String(error) }); + return null; + } + + if (!payload) return null; + if (payload.success !== true) { + ctx.logger?.warn("ZAI usage response invalid", { code: payload.code, message: payload.msg }); + return null; + } + + const limitsPayload = Array.isArray(payload.data?.limits) ? payload.data?.limits : []; + const limits: UsageLimit[] = []; + + for (const rawLimit of limitsPayload) { + const parsed = parseLimitItem(rawLimit); + if (!parsed) continue; + if (parsed.type === "TOKENS_LIMIT") { + const amount = buildUsageAmount({ + used: parsed.currentValue, + limit: parsed.usage, + remaining: parsed.remaining, + percentage: parsed.percentage, + unit: "tokens", + }); + const window = buildUsageWindow("quota", "Quota", parsed.nextResetTime, now); + limits.push({ + id: "zai:tokens", + label: "ZAI Token Quota", + scope: { + provider: params.provider, + windowId: window?.id ?? "quota", + shared: true, + }, + window, + amount, + status: getUsageStatus(amount.usedFraction), + }); + } + if (parsed.type === "TIME_LIMIT") { + const window = buildUsageWindow("quota", "Quota", undefined, now); + const amount = buildUsageAmount({ + used: parsed.currentValue, + limit: parsed.usage, + remaining: parsed.remaining, + percentage: parsed.percentage, + unit: "requests", + }); + limits.push({ + id: "zai:requests", + label: "ZAI Request Quota", + scope: { + provider: params.provider, + windowId: "quota", + shared: true, + }, + window, + amount, + status: getUsageStatus(amount.usedFraction), + }); + } + } + + if (limits.length === 0) return null; + + const report: UsageReport = { + provider: params.provider, + fetchedAt: now, + limits, + metadata: { + endpoint: url, + accountId: credential.accountId, + email: credential.email, + }, + raw: payload, + }; + + const expiresAt = resolveCacheExpiry(now, limits); + await ctx.cache.set(cacheKey, { value: report, expiresAt }); + + const modelUsageUrl = buildModelUsageUrl(baseUrl, new Date(now)); + try { + const response = await ctx.fetch(modelUsageUrl, { + headers, + signal: params.signal, + }); + if (response.ok) { + const modelUsagePayload = (await response.json()) as unknown; + if (isRecord(modelUsagePayload)) { + report.metadata = { + ...report.metadata, + modelUsage: modelUsagePayload, + }; + } + } + } catch (error) { + ctx.logger?.debug("ZAI model usage fetch failed", { error: String(error) }); + } + + return report; +} + +export const zaiUsageProvider: UsageProvider = { + id: "zai", + fetchUsage: fetchZaiUsage, + supports: (params) => params.provider === "zai" && params.credential.type === "api_key", +}; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d03df3c2e..6c9656763 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,9 +1,13 @@ # Changelog ## [Unreleased] - ### Added +- Added `/usage` command to display provider usage and limits +- Added support for multiple usage providers beyond Codex +- Added usage report caching with configurable TTL +- Added visual usage bars and account aggregation in usage display +- Added `fetchUsageReports()` method to agent session - Added `output()` function to read task/agent outputs by ID with support for multiple formats and queries - Added session file support to Python executor for accessing task outputs - Added support for jq-like queries when reading JSON outputs @@ -12,6 +16,9 @@ ### Changed +- Refactored usage tracking from Codex-specific to generic provider system +- Updated usage limit detection to work with multiple provider APIs +- Changed usage cache to use persistent storage instead of in-memory only - Limited diagnostic messages to 50 items to prevent overwhelming output when processing files with many issues - Changed `/dump` command to include complete agent context: system prompt, model config, available tools with schemas, and all message types (bash/python executions, custom messages, branch summaries, compaction summaries, file mentions) - Changed `/dump` format to use YAML instead of JSON for tool schemas and arguments (more readable) @@ -21,6 +28,7 @@ - Fixed external editor to work properly on Unix systems by correctly handling terminal I/O - Fixed external editor to show warning message when it fails to open instead of silently failing - Fixed find tool to properly handle no matches case without treating as error +- Fixed find tool to wait for fd exit so error messages no longer report exit null - Fixed read tool to properly handle no matches case without treating as error - Fixed orphaned Python kernel gateway processes not being killed on process exit diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index e2076ed29..294e100eb 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -13,22 +13,31 @@ * Modes use this class and add their own I/O layer on top. */ +import { existsSync, readFileSync } from "node:fs"; import type { Agent, AgentEvent, AgentMessage, AgentState, AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; -import type { AssistantMessage, ImageContent, Message, Model, TextContent, ToolCall, Usage } from "@oh-my-pi/pi-ai"; +import type { + AssistantMessage, + ImageContent, + Message, + Model, + TextContent, + ToolCall, + Usage, + UsageReport, +} from "@oh-my-pi/pi-ai"; import { isContextOverflow, modelsAreEqual, supportsXhigh } from "@oh-my-pi/pi-ai"; import { abortableSleep, logger } from "@oh-my-pi/pi-utils"; import { YAML } from "bun"; -import { existsSync, readFileSync } from "node:fs"; import type { Rule } from "../capability/rule"; import { getAgentDbPath } from "../config"; import { theme } from "../modes/interactive/theme/theme"; import ttsrInterruptTemplate from "../prompts/system/ttsr-interrupt.md" with { type: "text" }; import { type BashResult, executeBash as executeBashCommand } from "./bash-executor"; import { + type CompactionResult, calculateContextTokens, collectEntriesForBranchSummary, compact, - type CompactionResult, estimateTokens, generateBranchSummary, prepareCompaction, @@ -53,8 +62,8 @@ import { extractFileMentions, generateFileMentionMessages } from "./file-mention import type { HookCommandContext } from "./hooks/types"; import { type BashExecutionMessage, - bashExecutionToText, type BranchSummaryMessage, + bashExecutionToText, type CompactionSummaryMessage, type CustomMessage, type FileMentionMessage, @@ -64,7 +73,7 @@ import { } from "./messages"; import type { ModelRegistry } from "./model-registry"; import { parseModelString } from "./model-resolver"; -import { expandPromptTemplate, parseCommandArgs, type PromptTemplate, renderPromptTemplate } from "./prompt-templates"; +import { expandPromptTemplate, type PromptTemplate, parseCommandArgs, renderPromptTemplate } from "./prompt-templates"; import { executePython as executePythonCommand, type PythonResult } from "./python-executor"; import type { BranchSummaryEntry, CompactionEntry, NewSessionOptions, SessionManager } from "./session-manager"; import type { SettingsManager, SkillsSettings } from "./settings-manager"; @@ -3172,6 +3181,14 @@ export class AgentSession { }; } + async fetchUsageReports(): Promise { + const authStorage = this._modelRegistry.authStorage; + if (!authStorage.fetchUsageReports) return null; + return authStorage.fetchUsageReports({ + baseUrlResolver: (provider) => this._modelRegistry.getProviderBaseUrl?.(provider), + }); + } + /** * Estimate context tokens from messages, using the last assistant usage when available. */ diff --git a/packages/coding-agent/src/core/auth-storage.ts b/packages/coding-agent/src/core/auth-storage.ts index dc464a52d..c8f85ec98 100644 --- a/packages/coding-agent/src/core/auth-storage.ts +++ b/packages/coding-agent/src/core/auth-storage.ts @@ -3,10 +3,15 @@ * Handles loading, saving, and refreshing credentials from agent.db. */ +import { Buffer } from "node:buffer"; import { dirname, join } from "node:path"; import { + antigravityUsageProvider, + claudeUsageProvider, getEnvApiKey, getOAuthApiKey, + githubCopilotUsageProvider, + googleGeminiCliUsageProvider, loginAnthropic, loginAntigravity, loginCursor, @@ -16,6 +21,16 @@ import { type OAuthController, type OAuthCredentials, type OAuthProvider, + openaiCodexUsageProvider, + type Provider, + type UsageCache, + type UsageCacheEntry, + type UsageCredential, + type UsageLimit, + type UsageLogger, + type UsageProvider, + type UsageReport, + zaiUsageProvider, } from "@oh-my-pi/pi-ai"; import { logger } from "@oh-my-pi/pi-utils"; import { getAgentDbPath, getAuthPath } from "../config"; @@ -61,43 +76,64 @@ export interface SerializedAuthStorage { */ type StoredCredential = { id: number; credential: AuthCredential }; -/** Rate limit window from Codex usage API (primary or secondary quota). */ -type CodexUsageWindow = { - usedPercent?: number; - limitWindowSeconds?: number; - resetAt?: number; // Unix timestamp (seconds) +export type AuthStorageOptions = { + usageProviderResolver?: (provider: Provider) => UsageProvider | undefined; + usageCache?: UsageCache; + usageFetch?: typeof fetch; + usageNow?: () => number; + usageLogger?: UsageLogger; }; -/** Parsed usage data from Codex /wham/usage endpoint. */ -type CodexUsage = { - allowed?: boolean; - limitReached?: boolean; - primary?: CodexUsageWindow; - secondary?: CodexUsageWindow; -}; +const DEFAULT_USAGE_PROVIDERS: UsageProvider[] = [ + openaiCodexUsageProvider, + antigravityUsageProvider, + googleGeminiCliUsageProvider, + claudeUsageProvider, + zaiUsageProvider, + githubCopilotUsageProvider, +]; -/** Cached usage entry with TTL for avoiding redundant API calls. */ -type CodexUsageCacheEntry = { - fetchedAt: number; - expiresAt: number; - usage?: CodexUsage; -}; +const DEFAULT_USAGE_PROVIDER_MAP = new Map( + DEFAULT_USAGE_PROVIDERS.map((provider) => [provider.id, provider]), +); -function isRecord(value: unknown): value is Record { - return !!value && typeof value === "object" && !Array.isArray(value); +const USAGE_CACHE_PREFIX = "usage_cache:"; + +function resolveDefaultUsageProvider(provider: Provider): UsageProvider | undefined { + return DEFAULT_USAGE_PROVIDER_MAP.get(provider); } -function toNumber(value: unknown): number | undefined { - if (typeof value === "number" && Number.isFinite(value)) return value; - if (typeof value === "string" && value.trim()) { - const parsed = Number(value); - return Number.isFinite(parsed) ? parsed : undefined; +function parseUsageCacheEntry(raw: string): UsageCacheEntry | undefined { + try { + const parsed = JSON.parse(raw) as { value?: UsageReport | null; expiresAt?: unknown }; + const expiresAt = typeof parsed.expiresAt === "number" ? parsed.expiresAt : undefined; + if (!expiresAt || !Number.isFinite(expiresAt)) return undefined; + return { value: parsed.value ?? null, expiresAt }; + } catch { + return undefined; } - return undefined; } -function toBoolean(value: unknown): boolean | undefined { - return typeof value === "boolean" ? value : undefined; +class AuthStorageUsageCache implements UsageCache { + constructor(private storage: AgentStorage) {} + + get(key: string): UsageCacheEntry | undefined { + const raw = this.storage.getCache(`${USAGE_CACHE_PREFIX}${key}`); + if (!raw) return undefined; + const entry = parseUsageCacheEntry(raw); + if (!entry) return undefined; + if (entry.expiresAt <= Date.now()) return undefined; + return entry; + } + + set(key: string, entry: UsageCacheEntry): void { + const payload = JSON.stringify({ value: entry.value ?? null, expiresAt: entry.expiresAt }); + this.storage.setCache(`${USAGE_CACHE_PREFIX}${key}`, payload, Math.floor(entry.expiresAt / 1000)); + } + + cleanup(): void { + this.storage.cleanExpiredCache(); + } } /** @@ -105,14 +141,11 @@ function toBoolean(value: unknown): boolean | undefined { * Reads from SQLite and migrates legacy auth.json paths. */ export class AuthStorage { - private static readonly codexUsageCacheTtlMs = 60_000; // Cache usage data for 1 minute private static readonly defaultBackoffMs = 60_000; // Default backoff when no reset time available - private static readonly cacheCleanupIntervalMs = 300_000; // Clean expired cache every 5 minutes /** Provider -> credentials cache, populated from agent.db on reload(). */ private data: Map = new Map(); private storage: AgentStorage; - private lastCacheCleanup = 0; /** Resolved path to agent.db (derived from authPath or used directly if .db). */ private dbPath: string; private runtimeOverrides: Map = new Map(); @@ -122,8 +155,11 @@ export class AuthStorage { private sessionLastCredential: Map> = new Map(); /** Maps provider:type -> credentialIndex -> blockedUntilMs for temporary backoff. */ private credentialBackoff: Map> = new Map(); - /** Cached usage info for providers that expose usage endpoints. */ - private codexUsageCache: Map = new Map(); + private usageProviderResolver?: (provider: Provider) => UsageProvider | undefined; + private usageCache?: UsageCache; + private usageFetch: typeof fetch; + private usageNow: () => number; + private usageLogger?: UsageLogger; private fallbackResolver?: (provider: string) => string | undefined; /** @@ -133,16 +169,27 @@ export class AuthStorage { constructor( private authPath: string, private fallbackPaths: string[] = [], + options: AuthStorageOptions = {}, ) { this.dbPath = AuthStorage.resolveDbPath(authPath); this.storage = AgentStorage.open(this.dbPath); + this.usageProviderResolver = options.usageProviderResolver ?? resolveDefaultUsageProvider; + this.usageCache = options.usageCache ?? new AuthStorageUsageCache(this.storage); + this.usageFetch = options.usageFetch ?? fetch; + this.usageNow = options.usageNow ?? Date.now; + this.usageLogger = + options.usageLogger ?? + ({ + debug: (message, meta) => logger.debug(message, meta), + warn: (message, meta) => logger.warn(message, meta), + } satisfies UsageLogger); } /** * Create an in-memory AuthStorage instance from serialized data. * Used by subagent workers to bypass discovery and use parent's credentials. */ - static fromSerialized(data: SerializedAuthStorage): AuthStorage { + static fromSerialized(data: SerializedAuthStorage, options: AuthStorageOptions = {}): AuthStorage { const instance = Object.create(AuthStorage.prototype) as AuthStorage; const authPath = data.authPath ?? data.dbPath ?? getAuthPath(); instance.authPath = authPath; @@ -154,8 +201,16 @@ export class AuthStorage { instance.providerRoundRobinIndex = new Map(); instance.sessionLastCredential = new Map(); instance.credentialBackoff = new Map(); - instance.codexUsageCache = new Map(); - instance.lastCacheCleanup = 0; + instance.usageProviderResolver = options.usageProviderResolver ?? resolveDefaultUsageProvider; + instance.usageCache = options.usageCache ?? new AuthStorageUsageCache(instance.storage); + instance.usageFetch = options.usageFetch ?? fetch; + instance.usageNow = options.usageNow ?? Date.now; + instance.usageLogger = + options.usageLogger ?? + ({ + debug: (message, meta) => logger.debug(message, meta), + warn: (message, meta) => logger.warn(message, meta), + } satisfies UsageLogger); for (const [provider, creds] of Object.entries(data.credentials)) { instance.data.set( @@ -257,7 +312,15 @@ export class AuthStorage { list.push({ id: record.id, credential: record.credential }); grouped.set(record.provider, list); } - this.data = grouped; + + const dedupedGrouped = new Map(); + for (const [provider, entries] of grouped.entries()) { + const deduped = this.pruneDuplicateStoredCredentials(provider, entries); + if (deduped.length > 0) { + dedupedGrouped.set(provider, deduped); + } + } + this.data = dedupedGrouped; } /** @@ -283,6 +346,115 @@ export class AuthStorage { } } + private getOAuthIdentifiers(credential: OAuthCredential): string[] { + const identifiers: string[] = []; + const accountId = credential.accountId?.trim(); + if (accountId) identifiers.push(`account:${accountId}`); + const email = credential.email?.trim().toLowerCase(); + if (email) identifiers.push(`email:${email}`); + if (identifiers.length > 0) return identifiers; + const tokenIdentifiers = this.getOAuthIdentifiersFromToken(credential.access) ?? []; + for (const identifier of tokenIdentifiers) { + identifiers.push(identifier); + } + if (identifiers.length > 0) return identifiers; + const refreshIdentifiers = this.getOAuthIdentifiersFromToken(credential.refresh) ?? []; + for (const identifier of refreshIdentifiers) { + identifiers.push(identifier); + } + return identifiers; + } + + private getOAuthIdentifiersFromToken(token: string | undefined): string[] | undefined { + if (!token) return undefined; + const parts = token.split("."); + if (parts.length !== 3) return undefined; + const payloadRaw = parts[1]; + try { + const payload = JSON.parse( + Buffer.from(payloadRaw.replace(/-/g, "+").replace(/_/g, "/"), "base64").toString("utf8"), + ) as Record; + if (!payload || typeof payload !== "object") return undefined; + const identifiers: string[] = []; + const email = typeof payload.email === "string" ? payload.email.trim().toLowerCase() : undefined; + if (email) identifiers.push(`email:${email}`); + const accountId = + typeof payload.account_id === "string" + ? payload.account_id + : typeof payload.accountId === "string" + ? payload.accountId + : typeof payload.user_id === "string" + ? payload.user_id + : typeof payload.sub === "string" + ? payload.sub + : undefined; + const trimmedAccountId = accountId?.trim(); + if (trimmedAccountId) identifiers.push(`account:${trimmedAccountId}`); + return identifiers.length > 0 ? identifiers : undefined; + } catch { + return undefined; + } + } + + private dedupeOAuthCredentials(credentials: AuthCredential[]): AuthCredential[] { + const seen = new Set(); + const deduped: AuthCredential[] = []; + for (let index = credentials.length - 1; index >= 0; index -= 1) { + const credential = credentials[index]; + if (credential.type !== "oauth") { + deduped.push(credential); + continue; + } + const identifiers = this.getOAuthIdentifiers(credential); + if (identifiers.length === 0) { + deduped.push(credential); + continue; + } + if (identifiers.some((identifier) => seen.has(identifier))) { + continue; + } + for (const identifier of identifiers) { + seen.add(identifier); + } + deduped.push(credential); + } + return deduped.reverse(); + } + + private pruneDuplicateStoredCredentials(provider: string, entries: StoredCredential[]): StoredCredential[] { + const seen = new Set(); + const kept: StoredCredential[] = []; + const removed: StoredCredential[] = []; + for (let index = entries.length - 1; index >= 0; index -= 1) { + const entry = entries[index]; + const credential = entry.credential; + if (credential.type !== "oauth") { + kept.push(entry); + continue; + } + const identifiers = this.getOAuthIdentifiers(credential); + if (identifiers.length === 0) { + kept.push(entry); + continue; + } + if (identifiers.some((identifier) => seen.has(identifier))) { + removed.push(entry); + continue; + } + for (const identifier of identifiers) { + seen.add(identifier); + } + kept.push(entry); + } + if (removed.length > 0) { + for (const entry of removed) { + this.storage.deleteAuthCredential(entry.id); + } + this.resetProviderAssignments(provider); + } + return kept.reverse(); + } + /** Returns all credentials for a provider as an array */ private getCredentialsForProvider(provider: string): AuthCredential[] { return this.getStoredCredentials(provider).map((entry) => entry.credential); @@ -469,7 +641,8 @@ export class AuthStorage { */ async set(provider: string, credential: AuthCredentialEntry): Promise { const normalized = Array.isArray(credential) ? credential : [credential]; - const stored = this.storage.replaceAuthCredentialsForProvider(provider, normalized); + const deduped = this.dedupeOAuthCredentials(normalized); + const stored = this.storage.replaceAuthCredentialsForProvider(provider, deduped); this.setStoredCredentials( provider, stored.map((record) => ({ id: record.id, credential: record.credential })), @@ -611,202 +784,195 @@ export class AuthStorage { } // ───────────────────────────────────────────────────────────────────────────── - // Codex Usage API Integration - // Queries ChatGPT/Codex usage endpoints to detect rate limits before they occur. + // Usage API Integration + // Queries provider usage endpoints to detect rate limits before they occur. // ───────────────────────────────────────────────────────────────────────────── - /** Normalizes Codex base URL to include /backend-api path. */ - private normalizeCodexBaseUrl(baseUrl?: string): string { - const fallback = "https://chatgpt.com/backend-api"; - const trimmed = baseUrl?.trim() ? baseUrl.trim() : fallback; - const base = trimmed.replace(/\/+$/, ""); - const lower = base.toLowerCase(); - if ( - (lower.startsWith("https://chatgpt.com") || lower.startsWith("https://chat.openai.com")) && - !lower.includes("/backend-api") - ) { - return `${base}/backend-api`; - } - return base; - } - - private getCodexUsagePath(baseUrl: string): string { - return baseUrl.includes("/backend-api") ? "wham/usage" : "api/codex/usage"; - } - - private buildCodexUsageUrl(baseUrl: string, path: string): string { - const normalized = baseUrl.endsWith("/") ? baseUrl : `${baseUrl}/`; - return `${normalized}${path.replace(/^\/+/, "")}`; - } - - private getCodexUsageCacheKey(accountId: string, baseUrl: string): string { - return `${baseUrl}|${accountId}`; - } - - private extractCodexUsageWindow(window: unknown): CodexUsageWindow | undefined { - if (!isRecord(window)) return undefined; - const usedPercent = toNumber(window.used_percent); - const limitWindowSeconds = toNumber(window.limit_window_seconds); - const resetAt = toNumber(window.reset_at); - if (usedPercent === undefined && limitWindowSeconds === undefined && resetAt === undefined) return undefined; - return { usedPercent, limitWindowSeconds, resetAt }; - } - - private extractCodexUsage(payload: unknown): CodexUsage | undefined { - if (!isRecord(payload)) return undefined; - const rateLimit = isRecord(payload.rate_limit) ? payload.rate_limit : undefined; - if (!rateLimit) return undefined; - const primary = this.extractCodexUsageWindow(rateLimit.primary_window); - const secondary = this.extractCodexUsageWindow(rateLimit.secondary_window); - const usage: CodexUsage = { - allowed: toBoolean(rateLimit.allowed), - limitReached: toBoolean(rateLimit.limit_reached), - primary, - secondary, + private buildUsageCredential(credential: OAuthCredential): UsageCredential { + return { + type: "oauth", + accessToken: credential.access, + refreshToken: credential.refresh, + expiresAt: credential.expires, + accountId: credential.accountId, + projectId: credential.projectId, + email: credential.email, + enterpriseUrl: credential.enterpriseUrl, }; - if (!primary && !secondary && usage.allowed === undefined && usage.limitReached === undefined) return undefined; - return usage; } - /** Returns true if usage indicates rate limit has been reached. */ - private isCodexUsageLimitReached(usage: CodexUsage): boolean { - if (usage.allowed === false || usage.limitReached === true) return true; - if (usage.primary?.usedPercent !== undefined && usage.primary.usedPercent >= 100) return true; - if (usage.secondary?.usedPercent !== undefined && usage.secondary.usedPercent >= 100) return true; + private isUsageLimitExhausted(limit: UsageLimit): boolean { + if (limit.status === "exhausted") return true; + const amount = limit.amount; + if (amount.usedFraction !== undefined && amount.usedFraction >= 1) return true; + if (amount.remainingFraction !== undefined && amount.remainingFraction <= 0) return true; + if (amount.used !== undefined && amount.limit !== undefined && amount.used >= amount.limit) return true; + if (amount.remaining !== undefined && amount.remaining <= 0) return true; + if (amount.unit === "percent" && amount.used !== undefined && amount.used >= 100) return true; return false; } - /** Extracts the earliest reset timestamp from usage windows (in ms). */ - private getCodexResetAtMs(usage: CodexUsage): number | undefined { - const now = Date.now(); + /** Returns true if usage indicates rate limit has been reached. */ + private isUsageLimitReached(report: UsageReport): boolean { + return report.limits.some((limit) => this.isUsageLimitExhausted(limit)); + } + + /** Extracts the earliest reset timestamp from exhausted windows (in ms). */ + private getUsageResetAtMs(report: UsageReport, nowMs: number): number | undefined { const candidates: number[] = []; - const addCandidate = (value: number | undefined) => { - if (!value) return; - const ms = value > 1_000_000_000_000 ? value : value * 1000; - if (Number.isFinite(ms) && ms > now) { - candidates.push(ms); + for (const limit of report.limits) { + if (!this.isUsageLimitExhausted(limit)) continue; + const window = limit.window; + if (window?.resetsAt && window.resetsAt > nowMs) { + candidates.push(window.resetsAt); } - }; - const useAll = usage.limitReached === true || usage.allowed === false; - if (useAll) { - addCandidate(usage.primary?.resetAt); - addCandidate(usage.secondary?.resetAt); - } else { - if (usage.primary?.usedPercent !== undefined && usage.primary.usedPercent >= 100) { - addCandidate(usage.primary.resetAt); - } - if (usage.secondary?.usedPercent !== undefined && usage.secondary.usedPercent >= 100) { - addCandidate(usage.secondary.resetAt); + if (window?.resetInMs && window.resetInMs > 0) { + const resetAt = nowMs + window.resetInMs; + if (resetAt > nowMs) candidates.push(resetAt); } } if (candidates.length === 0) return undefined; return Math.min(...candidates); } - private getCodexUsageExpiryMs(usage: CodexUsage, nowMs: number): number { - const resetAtMs = this.getCodexResetAtMs(usage); - if (this.isCodexUsageLimitReached(usage)) { - if (resetAtMs) return resetAtMs; - return nowMs + AuthStorage.defaultBackoffMs; - } - const defaultExpiry = nowMs + AuthStorage.codexUsageCacheTtlMs; - if (!resetAtMs) return defaultExpiry; - return Math.min(defaultExpiry, resetAtMs); - } + private async getUsageReport( + provider: Provider, + credential: OAuthCredential, + options?: { baseUrl?: string }, + ): Promise { + const resolver = this.usageProviderResolver; + const cache = this.usageCache; + if (!resolver || !cache) return null; - /** Fetches usage data from Codex API. */ - private async fetchCodexUsage(credential: OAuthCredential, baseUrl?: string): Promise { - const accountId = credential.accountId; - if (!accountId) return undefined; + const providerImpl = resolver(provider); + if (!providerImpl) return null; - const normalizedBase = this.normalizeCodexBaseUrl(baseUrl); - const url = this.buildCodexUsageUrl(normalizedBase, this.getCodexUsagePath(normalizedBase)); - const headers = { - authorization: `Bearer ${credential.access}`, - "chatgpt-account-id": accountId, - "openai-beta": "responses=experimental", - originator: "codex_cli_rs", + const params = { + provider, + credential: this.buildUsageCredential(credential), + baseUrl: options?.baseUrl, }; - try { - const response = await fetch(url, { headers }); - if (!response.ok) { - logger.debug("AuthStorage codex usage fetch failed", { - status: response.status, - statusText: response.statusText, - }); - return undefined; - } + if (providerImpl.supports && !providerImpl.supports(params)) return null; - const payload = (await response.json()) as unknown; - return this.extractCodexUsage(payload); + try { + return await providerImpl.fetchUsage(params, { + cache, + fetch: this.usageFetch, + now: this.usageNow, + logger: this.usageLogger, + }); } catch (error) { - logger.debug("AuthStorage codex usage fetch error", { error: String(error) }); - return undefined; + logger.debug("AuthStorage usage fetch failed", { + provider, + error: String(error), + }); + return null; } } - /** Gets usage data with caching to avoid redundant API calls. */ - private async getCodexUsage(credential: OAuthCredential, baseUrl?: string): Promise { - const accountId = credential.accountId; - if (!accountId) return undefined; + async fetchUsageReports(options?: { + baseUrlResolver?: (provider: Provider) => string | undefined; + }): Promise { + const resolver = this.usageProviderResolver; + const cache = this.usageCache; + if (!resolver || !cache) return null; - const normalizedBase = this.normalizeCodexBaseUrl(baseUrl); - const cacheKey = this.getCodexUsageCacheKey(accountId, normalizedBase); - const now = Date.now(); + const tasks: Array> = []; + const providers = new Set([ + ...this.data.keys(), + ...DEFAULT_USAGE_PROVIDERS.map((provider) => provider.id), + ]); + for (const provider of providers) { + const providerImpl = resolver(provider as Provider); + if (!providerImpl) continue; + const baseUrl = options?.baseUrlResolver?.(provider as Provider); + let entries = this.getStoredCredentials(provider); + if (entries.length > 0) { + const dedupedEntries = this.pruneDuplicateStoredCredentials(provider, entries); + if (dedupedEntries.length !== entries.length) { + this.setStoredCredentials(provider, dedupedEntries); + } + entries = dedupedEntries; + } - if (now - this.lastCacheCleanup > AuthStorage.cacheCleanupIntervalMs) { - this.lastCacheCleanup = now; - this.storage.cleanExpiredCache(); - } + if (entries.length === 0) { + const runtimeKey = this.runtimeOverrides.get(provider); + const envKey = getEnvApiKey(provider); + const apiKey = runtimeKey ?? envKey; + if (!apiKey) { + continue; + } + const params = { + provider: provider as Provider, + credential: { type: "api_key", apiKey } satisfies UsageCredential, + baseUrl, + }; + if (providerImpl.supports && !providerImpl.supports(params)) { + continue; + } + tasks.push( + providerImpl + .fetchUsage(params, { + cache, + fetch: this.usageFetch, + now: this.usageNow, + logger: this.usageLogger, + }) + .catch((error) => { + logger.debug("AuthStorage usage fetch failed", { + provider, + error: String(error), + }); + return null; + }), + ); + continue; + } - // Check in-memory cache first (fastest) - const memCached = this.codexUsageCache.get(cacheKey); - if (memCached && memCached.expiresAt > now) { - return memCached.usage; - } + for (const entry of entries) { + const credential = entry.credential; + const usageCredential: UsageCredential = + credential.type === "api_key" + ? { type: "api_key", apiKey: credential.key } + : this.buildUsageCredential(credential); + const params = { + provider: provider as Provider, + credential: usageCredential, + baseUrl, + }; - // Check DB cache (survives restarts) - const dbCached = this.storage.getCache(`codex_usage:${cacheKey}`); - if (dbCached) { - try { - const parsed = JSON.parse(dbCached) as CodexUsage; - // Store in memory for faster subsequent access - this.codexUsageCache.set(cacheKey, { - fetchedAt: now, - expiresAt: now + AuthStorage.codexUsageCacheTtlMs, - usage: parsed, - }); - return parsed; - } catch { - // Invalid cache, continue to fetch + if (providerImpl.supports && !providerImpl.supports(params)) { + continue; + } + + tasks.push( + providerImpl + .fetchUsage(params, { + cache, + fetch: this.usageFetch, + now: this.usageNow, + logger: this.usageLogger, + }) + .catch((error) => { + logger.debug("AuthStorage usage fetch failed", { + provider, + error: String(error), + }); + return null; + }), + ); } } - // Fetch from API - const usage = await this.fetchCodexUsage(credential, normalizedBase); - if (usage) { - const expiresAt = this.getCodexUsageExpiryMs(usage, now); - this.codexUsageCache.set(cacheKey, { fetchedAt: now, expiresAt, usage }); - // Store in DB with 60s TTL - this.storage.setCache( - `codex_usage:${cacheKey}`, - JSON.stringify(usage), - Math.floor((now + AuthStorage.codexUsageCacheTtlMs) / 1000), - ); - return usage; - } - - this.codexUsageCache.set(cacheKey, { - fetchedAt: now, - expiresAt: now + AuthStorage.defaultBackoffMs, - }); - return undefined; + if (tasks.length === 0) return []; + const results = await Promise.all(tasks); + return results.filter((report): report is UsageReport => report !== null); } /** * Marks the current session's credential as temporarily blocked due to usage limits. - * Queries the Codex usage API to determine accurate reset time. + * Uses usage reports to determine accurate reset time when available. * Returns true if a credential was blocked, enabling automatic fallback to the next credential. */ async markUsageLimitReached( @@ -818,15 +984,15 @@ export class AuthStorage { if (!sessionCredential) return false; const providerKey = this.getProviderTypeKey(provider, sessionCredential.type); - const now = Date.now(); + const now = this.usageNow(); let blockedUntil = now + (options?.retryAfterMs ?? AuthStorage.defaultBackoffMs); if (provider === "openai-codex" && sessionCredential.type === "oauth") { const credential = this.getCredentialsForProvider(provider)[sessionCredential.index]; if (credential?.type === "oauth") { - const usage = await this.getCodexUsage(credential, options?.baseUrl); - if (usage) { - const resetAtMs = this.getCodexResetAtMs(usage); + const report = await this.getUsageReport(provider, credential, options); + if (report && this.isUsageLimitReached(report)) { + const resetAtMs = this.getUsageResetAtMs(report, this.usageNow()); if (resetAtMs && resetAtMs > blockedUntil) { blockedUntil = resetAtMs; } @@ -848,7 +1014,7 @@ export class AuthStorage { /** * Resolves an OAuth API key, trying credentials in priority order. - * Skips blocked credentials and checks usage limits for Codex accounts. + * Skips blocked credentials and checks usage limits for providers with usage data. * Falls back to earliest-unblocking credential if all are blocked. */ private async resolveOAuthApiKey( @@ -902,14 +1068,18 @@ export class AuthStorage { return undefined; } + let usage: UsageReport | null = null; + let usageChecked = false; + if (checkUsage) { - const usage = await this.getCodexUsage(selection.credential, options?.baseUrl); - if (usage && this.isCodexUsageLimitReached(usage)) { - const resetAtMs = this.getCodexResetAtMs(usage); + usage = await this.getUsageReport(provider, selection.credential, options); + usageChecked = true; + if (usage && this.isUsageLimitReached(usage)) { + const resetAtMs = this.getUsageResetAtMs(usage, this.usageNow()); this.markCredentialBlocked( providerKey, selection.index, - resetAtMs ?? Date.now() + AuthStorage.defaultBackoffMs, + resetAtMs ?? this.usageNow() + AuthStorage.defaultBackoffMs, ); return undefined; } @@ -927,13 +1097,16 @@ export class AuthStorage { this.replaceCredentialAt(provider, selection.index, updated); if (checkUsage) { - const usage = await this.getCodexUsage(updated, options?.baseUrl); - if (usage && this.isCodexUsageLimitReached(usage)) { - const resetAtMs = this.getCodexResetAtMs(usage); + const sameAccount = selection.credential.accountId === updated.accountId; + if (!usageChecked || !sameAccount) { + usage = await this.getUsageReport(provider, updated, options); + } + if (usage && this.isUsageLimitReached(usage)) { + const resetAtMs = this.getUsageResetAtMs(usage, this.usageNow()); this.markCredentialBlocked( providerKey, selection.index, - resetAtMs ?? Date.now() + AuthStorage.defaultBackoffMs, + resetAtMs ?? this.usageNow() + AuthStorage.defaultBackoffMs, ); return undefined; } diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 673648b88..37c3fbaf3 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -432,6 +432,13 @@ export class ModelRegistry { return this.models.find((m) => m.provider === provider && m.id === modelId); } + /** + * Get the base URL associated with a provider, if any model defines one. + */ + getProviderBaseUrl(provider: string): string | undefined { + return this.models.find((m) => m.provider === provider && m.baseUrl)?.baseUrl; + } + /** * Get API key for a model. */ diff --git a/packages/coding-agent/src/core/tools/find.ts b/packages/coding-agent/src/core/tools/find.ts index 20f213c6f..92d27e971 100644 --- a/packages/coding-agent/src/core/tools/find.ts +++ b/packages/coding-agent/src/core/tools/find.ts @@ -87,10 +87,22 @@ export async function runFd(fdPath: string, args: string[], signal?: AbortSignal throw err; } + let exitError: unknown; + try { + await child.exited; + } catch (err) { + exitError = err; + if (err instanceof ptree.Exception && err.aborted) { + throw new Error("Operation aborted"); + } + } + + const exitCode = child.exitCode ?? (exitError instanceof ptree.Exception ? exitError.exitCode : null); + return { stdout, stderr: child.peekStderr(), - exitCode: child.exitCode, + exitCode, }; } diff --git a/packages/coding-agent/src/modes/interactive/controllers/command-controller.ts b/packages/coding-agent/src/modes/interactive/controllers/command-controller.ts index 99dcbb0b2..10e9d3fc0 100644 --- a/packages/coding-agent/src/modes/interactive/controllers/command-controller.ts +++ b/packages/coding-agent/src/modes/interactive/controllers/command-controller.ts @@ -1,6 +1,7 @@ import { mkdir, rm } from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; +import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai"; import { Loader, Markdown, Spacer, Text, visibleWidth } from "@oh-my-pi/pi-tui"; import { $ } from "bun"; import { nanoid } from "nanoid"; @@ -284,6 +285,33 @@ export class CommandController { this.ctx.ui.requestRender(); } + async handleUsageCommand(reports?: UsageReport[] | null): Promise { + let usageReports = reports ?? null; + if (!usageReports) { + const provider = this.ctx.session as { fetchUsageReports?: () => Promise }; + if (!provider.fetchUsageReports) { + this.ctx.showWarning("Usage reporting is not configured for this session."); + return; + } + try { + usageReports = await provider.fetchUsageReports(); + } catch (error) { + this.ctx.showError(`Failed to fetch usage data: ${error instanceof Error ? error.message : String(error)}`); + return; + } + } + + if (!usageReports || usageReports.length === 0) { + this.ctx.showWarning("No usage data available."); + return; + } + + const output = renderUsageReports(usageReports, theme, Date.now()); + this.ctx.chatContainer.addChild(new Spacer(1)); + this.ctx.chatContainer.addChild(new Text(output, 1, 0)); + this.ctx.ui.requestRender(); + } + handleChangelogCommand(): void { const changelogPath = getChangelogPath(); const allEntries = parseChangelog(changelogPath); @@ -598,3 +626,307 @@ export class CommandController { await this.ctx.flushCompactionQueue({ willRetry: false }); } } + +const BAR_WIDTH = 24; +const COLUMN_WIDTH = BAR_WIDTH + 2; + +function formatProviderName(provider: string): string { + return provider + .split(/[-_]/g) + .map((part) => (part ? part[0].toUpperCase() + part.slice(1) : "")) + .join(" "); +} + +function formatNumber(value: number, maxFractionDigits = 1): string { + return new Intl.NumberFormat("en-US", { maximumFractionDigits: maxFractionDigits }).format(value); +} + +function formatUsedAccounts(value: number): string { + return `${value.toFixed(2)} used`; +} + +function formatDuration(ms: number): string { + const totalSeconds = Math.max(0, Math.round(ms / 1000)); + const minutes = Math.floor(totalSeconds / 60); + const seconds = totalSeconds % 60; + const hours = Math.floor(minutes / 60); + const mins = minutes % 60; + const days = Math.floor(hours / 24); + const hrs = hours % 24; + if (days > 0) return `${days}d ${hrs}h`; + if (hours > 0) return `${hours}h ${mins}m`; + if (minutes > 0) return `${minutes}m ${seconds}s`; + return `${seconds}s`; +} + +function formatDurationShort(ms: number): string { + const totalSeconds = Math.max(0, Math.round(ms / 1000)); + const minutes = Math.floor(totalSeconds / 60); + const hours = Math.floor(minutes / 60); + const mins = minutes % 60; + const days = Math.floor(hours / 24); + const hrs = hours % 24; + if (days > 0) return `${days}d${hrs > 0 ? ` ${hrs}h` : ""}`; + if (hours > 0) return `${hours}h${mins > 0 ? ` ${mins}m` : ""}`; + if (minutes > 0) return `${minutes}m`; + return `${totalSeconds}s`; +} + +function resolveFraction(limit: UsageLimit): number | undefined { + const amount = limit.amount; + if (amount.usedFraction !== undefined) return amount.usedFraction; + if (amount.used !== undefined && amount.limit !== undefined && amount.limit > 0) { + return amount.used / amount.limit; + } + if (amount.unit === "percent" && amount.used !== undefined) { + return amount.used / 100; + } + return undefined; +} + +function formatLimitTitle(limit: UsageLimit): string { + const tier = limit.scope.tier; + if (tier && !limit.label.toLowerCase().includes(tier.toLowerCase())) { + return `${limit.label} (${tier})`; + } + return limit.label; +} + +function formatWindowSuffix(label: string, windowLabel: string, uiTheme: typeof theme): string { + const normalizedLabel = label.toLowerCase(); + const normalizedWindow = windowLabel.toLowerCase(); + if (normalizedWindow === "quota window") return ""; + if (normalizedLabel.includes(normalizedWindow)) return ""; + return uiTheme.fg("dim", `(${windowLabel})`); +} + +function formatAccountLabel(limit: UsageLimit, report: UsageReport, index: number): string { + const email = (report.metadata?.email as string | undefined) ?? limit.scope.accountId; + if (email) return email; + const accountId = (report.metadata?.accountId as string | undefined) ?? limit.scope.accountId; + if (accountId) return accountId; + return `account ${index + 1}`; +} + +function formatResetShort(limit: UsageLimit, nowMs: number): string | undefined { + if (limit.window?.resetInMs !== undefined) { + return formatDurationShort(limit.window.resetInMs); + } + if (limit.window?.resetsAt !== undefined) { + return formatDurationShort(limit.window.resetsAt - nowMs); + } + return undefined; +} + +function formatAccountHeader(limit: UsageLimit, report: UsageReport, index: number, nowMs: number): string { + const label = formatAccountLabel(limit, report, index); + const reset = formatResetShort(limit, nowMs); + if (!reset) return label; + return `${label} (${reset})`; +} + +function padColumn(text: string, width: number): string { + const visible = visibleWidth(text); + if (visible >= width) return text; + return `${text}${" ".repeat(width - visible)}`; +} + +function resolveAggregateStatus(limits: UsageLimit[]): UsageLimit["status"] { + const hasOk = limits.some((limit) => limit.status === "ok"); + const hasWarning = limits.some((limit) => limit.status === "warning"); + const hasExhausted = limits.some((limit) => limit.status === "exhausted"); + if (!hasOk && !hasWarning && !hasExhausted) return "unknown"; + if (hasOk) { + return hasWarning || hasExhausted ? "warning" : "ok"; + } + if (hasWarning) return "warning"; + return "exhausted"; +} + +function isZeroUsage(limit: UsageLimit): boolean { + const amount = limit.amount; + if (amount.usedFraction !== undefined) return amount.usedFraction <= 0; + if (amount.used !== undefined) return amount.used <= 0; + if (amount.unit === "percent" && amount.used !== undefined) return amount.used <= 0; + if (amount.remainingFraction !== undefined) return amount.remainingFraction >= 1; + return false; +} + +function isZeroUsageGroup(limits: UsageLimit[]): boolean { + return limits.length > 0 && limits.every((limit) => isZeroUsage(limit)); +} + +function formatAggregateAmount(limits: UsageLimit[]): string { + const fractions = limits + .map((limit) => resolveFraction(limit)) + .filter((value): value is number => value !== undefined); + if (fractions.length === limits.length && fractions.length > 0) { + const sum = fractions.reduce((total, value) => total + value, 0); + const usedPct = Math.max(sum * 100, 0); + const remainingPct = Math.max(0, limits.length * 100 - usedPct); + const avgRemaining = limits.length > 0 ? remainingPct / limits.length : remainingPct; + return `${formatUsedAccounts(sum)} (${formatNumber(avgRemaining)}% left)`; + } + + const amounts = limits + .map((limit) => limit.amount) + .filter((amount) => amount.used !== undefined && amount.limit !== undefined && amount.limit > 0); + if (amounts.length === limits.length && amounts.length > 0) { + const totalUsed = amounts.reduce((sum, amount) => sum + (amount.used ?? 0), 0); + const totalLimit = amounts.reduce((sum, amount) => sum + (amount.limit ?? 0), 0); + const usedPct = totalLimit > 0 ? (totalUsed / totalLimit) * 100 : 0; + const remainingPct = Math.max(0, 100 - usedPct); + const usedAccounts = totalLimit > 0 ? (usedPct / 100) * limits.length : 0; + return `${formatUsedAccounts(usedAccounts)} (${formatNumber(remainingPct)}% left)`; + } + + return `Accounts: ${limits.length}`; +} + +function resolveResetRange(limits: UsageLimit[], nowMs: number): string | null { + const resets = limits + .map((limit) => limit.window?.resetInMs ?? undefined) + .filter((value): value is number => value !== undefined && Number.isFinite(value) && value > 0); + if (resets.length === 0) { + const absolute = limits + .map((limit) => limit.window?.resetsAt) + .filter((value): value is number => value !== undefined && Number.isFinite(value) && value > nowMs); + if (absolute.length === 0) return null; + const earliest = Math.min(...absolute); + return `resets at ${new Date(earliest).toLocaleString()}`; + } + const minReset = Math.min(...resets); + const maxReset = Math.max(...resets); + if (maxReset - minReset > 60_000) { + return `resets in ${formatDuration(minReset)}–${formatDuration(maxReset)}`; + } + return `resets in ${formatDuration(minReset)}`; +} + +function resolveStatusIcon(status: UsageLimit["status"], uiTheme: typeof theme): string { + if (status === "exhausted") return uiTheme.fg("error", uiTheme.status.error); + if (status === "warning") return uiTheme.fg("warning", uiTheme.status.warning); + if (status === "ok") return uiTheme.fg("success", uiTheme.status.success); + return uiTheme.fg("dim", uiTheme.status.pending); +} + +function resolveStatusColor(status: UsageLimit["status"]): "success" | "warning" | "error" | "dim" { + if (status === "exhausted") return "error"; + if (status === "warning") return "warning"; + if (status === "ok") return "success"; + return "dim"; +} + +function renderUsageBar(limit: UsageLimit, uiTheme: typeof theme): string { + const fraction = resolveFraction(limit); + if (fraction === undefined) { + return uiTheme.fg("dim", `[${"·".repeat(BAR_WIDTH)}]`); + } + const clamped = Math.min(Math.max(fraction, 0), 1); + const filled = Math.round(clamped * BAR_WIDTH); + const filledBar = "█".repeat(filled); + const emptyBar = "░".repeat(Math.max(0, BAR_WIDTH - filled)); + const color = resolveStatusColor(limit.status); + return `${uiTheme.fg("dim", "[")}${uiTheme.fg(color, filledBar)}${uiTheme.fg("dim", emptyBar)}${uiTheme.fg("dim", "]")}`; +} + +function renderUsageReports(reports: UsageReport[], uiTheme: typeof theme, nowMs: number): string { + const lines: string[] = []; + const latestFetchedAt = Math.max(...reports.map((report) => report.fetchedAt ?? 0)); + const headerSuffix = latestFetchedAt ? ` (${formatDuration(nowMs - latestFetchedAt)} ago)` : ""; + lines.push(uiTheme.bold(uiTheme.fg("accent", `Usage${headerSuffix}`))); + const grouped = new Map(); + for (const report of reports) { + const list = grouped.get(report.provider) ?? []; + list.push(report); + grouped.set(report.provider, list); + } + + for (const [provider, providerReports] of grouped.entries()) { + lines.push(""); + const providerName = formatProviderName(provider); + + const limitGroups = new Map< + string, + { label: string; windowLabel: string; limits: UsageLimit[]; reports: UsageReport[] } + >(); + for (const report of providerReports) { + for (const limit of report.limits) { + const windowId = limit.window?.id ?? limit.scope.windowId ?? "default"; + const key = `${formatLimitTitle(limit)}|${windowId}`; + const windowLabel = limit.window?.label ?? windowId; + const entry = limitGroups.get(key) ?? { + label: formatLimitTitle(limit), + windowLabel, + limits: [], + reports: [], + }; + entry.limits.push(limit); + entry.reports.push(report); + limitGroups.set(key, entry); + } + } + + const providerAllZero = isZeroUsageGroup(Array.from(limitGroups.values()).flatMap((group) => group.limits)); + if (providerAllZero) { + const providerTitle = `${resolveStatusIcon("ok", uiTheme)} ${uiTheme.fg("accent", `${providerName} (0%)`)}`; + lines.push(uiTheme.bold(providerTitle)); + continue; + } + + lines.push(uiTheme.bold(uiTheme.fg("accent", providerName))); + + for (const group of limitGroups.values()) { + const entries = group.limits.map((limit, index) => ({ + limit, + report: group.reports[index], + fraction: resolveFraction(limit), + index, + })); + entries.sort((a, b) => { + const aFraction = a.fraction ?? -1; + const bFraction = b.fraction ?? -1; + if (aFraction !== bFraction) return bFraction - aFraction; + return a.index - b.index; + }); + const sortedLimits = entries.map((entry) => entry.limit); + const sortedReports = entries.map((entry) => entry.report); + + const status = resolveAggregateStatus(sortedLimits); + const statusIcon = resolveStatusIcon(status, uiTheme); + if (isZeroUsageGroup(sortedLimits)) { + const resetText = resolveResetRange(sortedLimits, nowMs); + const resetSuffix = resetText ? ` | ${resetText}` : ""; + const windowSuffix = formatWindowSuffix(group.label, group.windowLabel, uiTheme); + lines.push( + `${statusIcon} ${uiTheme.bold(group.label)} ${windowSuffix} ${uiTheme.fg( + "dim", + `0%${resetSuffix}`, + )}`.trim(), + ); + continue; + } + + const windowSuffix = formatWindowSuffix(group.label, group.windowLabel, uiTheme); + lines.push(`${statusIcon} ${uiTheme.bold(group.label)} ${windowSuffix}`.trim()); + const accountLabels = sortedLimits.map((limit, index) => + padColumn(formatAccountHeader(limit, sortedReports[index], index, nowMs), COLUMN_WIDTH), + ); + lines.push(` ${accountLabels.join(" ")}`.trimEnd()); + const bars = sortedLimits.map((limit) => padColumn(renderUsageBar(limit, uiTheme), COLUMN_WIDTH)); + lines.push(` ${bars.join(" ")} ${formatAggregateAmount(sortedLimits)}`.trimEnd()); + const resetText = sortedLimits.length <= 1 ? resolveResetRange(sortedLimits, nowMs) : null; + if (resetText) { + lines.push(` ${uiTheme.fg("dim", resetText)}`.trimEnd()); + } + const notes = sortedLimits.flatMap((limit) => limit.notes ?? []); + if (notes.length > 0) { + lines.push(` ${uiTheme.fg("dim", notes.join(" • "))}`.trimEnd()); + } + } + + // No per-provider footer; global header shows last check. + } + + return lines.join("\n"); +} diff --git a/packages/coding-agent/src/modes/interactive/controllers/input-controller.ts b/packages/coding-agent/src/modes/interactive/controllers/input-controller.ts index 2303fc80e..7760cefe7 100644 --- a/packages/coding-agent/src/modes/interactive/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/interactive/controllers/input-controller.ts @@ -211,6 +211,11 @@ export class InputController { this.ctx.editor.setText(""); return; } + if (text === "/usage") { + await this.ctx.handleUsageCommand(); + this.ctx.editor.setText(""); + return; + } if (text === "/changelog") { this.ctx.handleChangelogCommand(); this.ctx.editor.setText(""); diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 705aec9cb..a6dfcdef3 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -5,7 +5,7 @@ import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; -import type { AssistantMessage, ImageContent, Message } from "@oh-my-pi/pi-ai"; +import type { AssistantMessage, ImageContent, Message, UsageReport } from "@oh-my-pi/pi-ai"; import type { Component, Loader, SlashCommand } from "@oh-my-pi/pi-tui"; import { CombinedAutocompleteProvider, @@ -187,6 +187,7 @@ export class InteractiveMode implements InteractiveModeContext { { name: "share", description: "Share session as a secret GitHub gist" }, { name: "copy", description: "Copy last agent message to clipboard" }, { name: "session", description: "Show session info and stats" }, + { name: "usage", description: "Show provider usage and limits" }, { name: "extensions", description: "Open Extension Control Center dashboard" }, { name: "status", description: "Alias for /extensions" }, { name: "changelog", description: "Show changelog entries" }, @@ -613,6 +614,10 @@ export class InteractiveMode implements InteractiveModeContext { this.commandController.handleSessionCommand(); } + handleUsageCommand(reports?: UsageReport[] | null): Promise { + return this.commandController.handleUsageCommand(reports); + } + handleChangelogCommand(): void { this.commandController.handleChangelogCommand(); } diff --git a/packages/coding-agent/src/modes/interactive/types.ts b/packages/coding-agent/src/modes/interactive/types.ts index d5846eab5..dbffb6048 100644 --- a/packages/coding-agent/src/modes/interactive/types.ts +++ b/packages/coding-agent/src/modes/interactive/types.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; -import type { AssistantMessage, ImageContent, Message } from "@oh-my-pi/pi-ai"; +import type { AssistantMessage, ImageContent, Message, UsageReport } from "@oh-my-pi/pi-ai"; import type { Component, Container, Loader, Spacer, Text, TUI } from "@oh-my-pi/pi-tui"; import type { AgentSession, AgentSessionEvent } from "../../core/agent-session"; import type { ExtensionUIContext } from "../../core/extensions/index"; @@ -134,6 +134,7 @@ export interface InteractiveModeContext { handleShareCommand(): Promise; handleCopyCommand(): Promise; handleSessionCommand(): void; + handleUsageCommand(reports?: UsageReport[] | null): Promise; handleChangelogCommand(): void; handleHotkeysCommand(): void; handleDumpCommand(): Promise; diff --git a/tsconfig.json b/tsconfig.json index 23f367414..93e4be9ed 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -4,6 +4,8 @@ "noEmit": true, "paths": { "*": ["./*"], + "@oh-my-pi/pi-ai": ["./packages/ai/src/index.ts"], + "@oh-my-pi/pi-ai/*": ["./packages/ai/src/*"], "@oh-my-pi/pi-agent-core": ["./packages/agent/src/index.ts"], "@oh-my-pi/pi-agent-core/*": ["./packages/agent/src/*"], "@oh-my-pi/pi-utils": ["./packages/pi-utils/src/index.ts"],