diff --git a/README.md b/README.md index 5e92d20..3d67854 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # pi-usage -A lightweight pi extension that provides a single `/usage` command to display the **current provider's** daily and weekly usage limits. +A lightweight pi extension that provides a single `/usage` command to display the **current provider's** usage limits (daily & weekly, plus monthly for OpenCode Go). ## Install @@ -16,7 +16,7 @@ In pi, type: /usage ``` -A panel appears showing the current provider's daily and weekly limits with progress bars. Press Enter or Escape to close. +A panel appears showing the current provider's usage limits with progress bars (Rolling/Weekly/Monthly for OpenCode Go, Daily/Weekly otherwise). Press Enter or Escape to close. ## Supported Providers @@ -27,6 +27,7 @@ A panel appears showing the current provider's daily and weekly limits with prog | Z.AI | API key | | Gemini CLI | OAuth (`/login`) | | Antigravity (Google) | OAuth (`/login`) | +| OpenCode Go | API key (`/connect` → OpenCode Go) | ## Credits diff --git a/extensions/usage/core.ts b/extensions/usage/core.ts index 5a0def0..0989b07 100644 --- a/extensions/usage/core.ts +++ b/extensions/usage/core.ts @@ -2,7 +2,7 @@ import * as fs from "node:fs"; import * as path from "node:path"; import * as os from "node:os"; -export type ProviderKey = "codex" | "claude" | "zai" | "gemini" | "antigravity"; +export type ProviderKey = "codex" | "claude" | "zai" | "gemini" | "antigravity" | "opencode-go"; export type OAuthProviderId = "openai-codex" | "anthropic" | "google-gemini-cli" | "google-antigravity"; export interface AuthData { @@ -11,6 +11,7 @@ export interface AuthData { zai?: { key?: string; access?: string; refresh?: string; expires?: number }; "google-gemini-cli"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; "google-antigravity"?: { access?: string; refresh?: string; projectId?: string; expires?: number }; + "opencode-go"?: { key?: string; access?: string; type?: string }; } export interface UsageData { @@ -18,6 +19,8 @@ export interface UsageData { weekly: number; sessionResetsIn?: string; weeklyResetsIn?: string; + monthly?: number; + monthlyResetsIn?: string; extraSpend?: number; extraLimit?: number; error?: string; @@ -29,6 +32,7 @@ export interface UsageEndpoints { zai: string; gemini: string; antigravity: string; + opencodeGo: string; googleLoadCodeAssistEndpoints: string[]; } @@ -84,6 +88,7 @@ const TOKEN_REFRESH_SKEW_MS = 60_000; export const DEFAULT_AUTH_FILE = path.join(os.homedir(), ".pi", "agent", "auth.json"); export const DEFAULT_ZAI_USAGE_ENDPOINT = "https://api.z.ai/api/monitor/usage/quota/limit"; +export const DEFAULT_OPENCODE_GO_USAGE_ENDPOINT = "https://opencode.ai/zen/go/v1/usage"; export const GOOGLE_QUOTA_ENDPOINT = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"; export const GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS = [ "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", @@ -95,6 +100,7 @@ export function resolveUsageEndpoints(): UsageEndpoints { zai: DEFAULT_ZAI_USAGE_ENDPOINT, gemini: GOOGLE_QUOTA_ENDPOINT, antigravity: GOOGLE_QUOTA_ENDPOINT, + opencodeGo: DEFAULT_OPENCODE_GO_USAGE_ENDPOINT, googleLoadCodeAssistEndpoints: GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS, }; } @@ -552,6 +558,48 @@ export async function fetchGoogleUsage( return parsed; } +export async function fetchOpencodeGoUsage(apiKey: string, config: FetchConfig = {}): Promise { + const endpoint = (config.endpoints ?? resolveUsageEndpoints()).opencodeGo || DEFAULT_OPENCODE_GO_USAGE_ENDPOINT; + if (!endpoint) return { session: 0, weekly: 0, error: "usage endpoint unavailable" }; + + const result = await requestJson( + endpoint, + { headers: { Authorization: `Bearer ${apiKey}` } }, + config, + ); + + if (!result.ok) return { session: 0, weekly: 0, error: (result as { ok: false; error: string }).error }; + + const usage = result.data?.usage; + if (!usage) return { session: 0, weekly: 0, error: "unrecognized response shape" }; + + const rolling = usage.rolling ?? usage.rolling5h ?? usage.session ?? usage.daily; + const weekly = usage.weekly ?? usage.weeklyUsage ?? usage.seven_day; + const monthly = usage.monthly ?? usage.monthlyUsage; + + const session = readPercentCandidate(rolling?.percent ?? rolling?.usagePercent ?? rolling); + const weeklyPct = readPercentCandidate(weekly?.percent ?? weekly?.usagePercent ?? weekly); + const monthlyPct = readPercentCandidate(monthly?.percent ?? monthly?.usagePercent ?? monthly); + + if (session == null || weeklyPct == null) { + return { session: 0, weekly: 0, error: "unrecognized response shape" }; + } + + const out: UsageData = { + ...normalizeUsagePair(session, weeklyPct), + }; + const rollingReset = rolling?.resetsAt ?? rolling?.resetAt ?? rolling?.resets_at; + const weeklyReset = weekly?.resetsAt ?? weekly?.resetAt ?? weekly?.resets_at; + const monthlyReset = monthly?.resetsAt ?? monthly?.resetAt ?? monthly?.resets_at; + if (typeof rollingReset === "string" && rollingReset) out.sessionResetsIn = formatResetsAt(rollingReset); + if (typeof weeklyReset === "string" && weeklyReset) out.weeklyResetsIn = formatResetsAt(weeklyReset); + if (monthlyPct != null) { + out.monthly = clampPercent(monthlyPct); + if (typeof monthlyReset === "string" && monthlyReset) out.monthlyResetsIn = formatResetsAt(monthlyReset); + } + return out; +} + export function detectProvider( model: { provider?: string; id?: string; name?: string; api?: string } | string | undefined | null, ): ProviderKey | null { @@ -565,6 +613,7 @@ export function detectProvider( if (provider === "zai") return "zai"; if (provider === "google-gemini-cli") return "gemini"; if (provider === "google-antigravity") return "antigravity"; + if (provider === "opencode-go" || provider === "opencode") return "opencode-go"; return null; } @@ -581,6 +630,8 @@ export function canShowForProvider(active: ProviderKey | null, auth: AuthData | if (!active || !auth) return false; if (active === "codex") return !!(auth["openai-codex"]?.access || auth["openai-codex"]?.refresh); if (active === "claude") return !!(auth.anthropic?.access || auth.anthropic?.refresh); + if (active === "opencode-go") + return !!((auth as AuthData)["opencode-go"]?.key || (auth as AuthData)["opencode-go"]?.access) && !!endpoints.opencodeGo; if (active === "zai") return !!(auth.zai?.access || auth.zai?.key) && !!endpoints.zai; if (active === "gemini") { return !!(auth["google-gemini-cli"]?.access || auth["google-gemini-cli"]?.refresh) && !!endpoints.gemini; @@ -613,6 +664,7 @@ export async function fetchAllUsages(config: FetchAllUsagesConfig = {}): Promise zai: null, gemini: null, antigravity: null, + "opencode-go": null, }; if (!auth) return results; @@ -666,6 +718,12 @@ export async function fetchAllUsages(config: FetchAllUsagesConfig = {}): Promise assign("zai", fetchZaiUsage(authData.zai.access || authData.zai.key!, { ...config, endpoints })); } + const goCreds = (authData as AuthData)["opencode-go"]; + const goKey = goCreds?.access || goCreds?.key; + if (goKey) { + assign("opencode-go", fetchOpencodeGoUsage(goKey, { ...config, endpoints })); + } + if (authData["google-gemini-cli"]?.access) { const err = refreshError("google-gemini-cli"); if (err) { diff --git a/extensions/usage/index.ts b/extensions/usage/index.ts index b047aa8..bb2a61b 100644 --- a/extensions/usage/index.ts +++ b/extensions/usage/index.ts @@ -14,6 +14,7 @@ import { fetchClaudeUsage, fetchCodexUsage, fetchGoogleUsage, + fetchOpencodeGoUsage, fetchZaiUsage, providerToOAuthProviderId, resolveUsageEndpoints, @@ -30,6 +31,7 @@ const PROVIDER_LABELS: Record = { zai: "Z.AI", gemini: "Gemini", antigravity: "Antigravity", + "opencode-go": "OpenCode Go", }; // ── Self-managing usage panel ──────────────────────────────────── @@ -122,6 +124,8 @@ class UsagePanelComponent extends Container implements Focusable { new Text(" " + t.fg("error", `Error: ${data.error}`), 0, 0), ); } else { + const isGo = provider === "opencode-go"; + const sessionLabel = isGo ? "Rolling " : "Daily "; const session = clampPercent(data.session); const sessionReset = data.sessionResetsIn ? t.fg("dim", ` resets in ${data.sessionResetsIn}`) @@ -129,7 +133,7 @@ class UsagePanelComponent extends Container implements Focusable { this.contentContainer.addChild( new Text( " " + - t.fg("muted", "Daily ") + + t.fg("muted", sessionLabel) + renderBar(t, session) + " " + t.fg(colorForPercent(session), `${session}%`.padStart(4)) + @@ -156,6 +160,25 @@ class UsagePanelComponent extends Container implements Focusable { ), ); + if (typeof data.monthly === "number") { + const monthly = clampPercent(data.monthly); + const monthlyReset = data.monthlyResetsIn + ? t.fg("dim", ` resets in ${data.monthlyResetsIn}`) + : ""; + this.contentContainer.addChild( + new Text( + " " + + t.fg("muted", "Monthly ") + + renderBar(t, monthly) + + " " + + t.fg(colorForPercent(monthly), `${monthly}%`.padStart(4)) + + monthlyReset, + 0, + 0, + ), + ); + } + if (typeof data.extraSpend === "number" && typeof data.extraLimit === "number") { this.contentContainer.addChild( new Text( @@ -246,6 +269,12 @@ async function fetchProviderUsage( ? fetchGoogleUsage(access, endpoints.antigravity, undefined, "antigravity", { endpoints }) : { session: 0, weekly: 0, error: "missing access token (try /login again)" }; } + case "opencode-go": { + const apiKey = await getAccessToken("opencode-go"); + return apiKey + ? fetchOpencodeGoUsage(apiKey, { endpoints }) + : { session: 0, weekly: 0, error: "missing API key (connect OpenCode Go to get a key)" }; + } default: return null; } @@ -279,7 +308,13 @@ export default function (pi: ExtensionAPI) { return; } - if (provider !== "zai" && oauthId && !auth.hasAuth(oauthId)) { + // OpenCode Go uses an API key (sk-...) stored as "opencode-go", not OAuth + if (provider === "opencode-go" && !auth.hasAuth("opencode-go")) { + ctx.ui.notify("No OpenCode Go API key found (connect OpenCode Go to get a key)", "warning"); + return; + } + + if (provider !== "zai" && provider !== "opencode-go" && oauthId && !auth.hasAuth(oauthId)) { ctx.ui.notify( `No credentials found for ${PROVIDER_LABELS[provider] ?? provider}`, "warning", diff --git a/tests/usage-core.test.ts b/tests/usage-core.test.ts index 82390b2..2c7ad83 100644 --- a/tests/usage-core.test.ts +++ b/tests/usage-core.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "bun:test"; import { detectProvider, fetchCodexUsage, + fetchOpencodeGoUsage, type FetchLike, type FetchResponseLike, } from "../extensions/usage/core.ts"; @@ -73,3 +74,50 @@ describe("usage core Codex support", () => { expect(badJson.error).toBe("invalid JSON response"); }); }); + +describe("usage core OpenCode Go support", () => { + it("detects opencode-go models as OpenCode Go", () => { + expect(detectProvider({ provider: "opencode-go", id: "kimi-k2.7-code" })).toBe("opencode-go"); + expect(detectProvider({ provider: "opencode", id: "glm-5.2" })).toBe("opencode-go"); + }); + + it("fetches rolling, weekly and monthly usage from Zen Go endpoint", async () => { + const calls: Array<{ url: string; authorization: string }> = []; + const fetchFn: FetchLike = async (url, init) => { + calls.push({ + url, + authorization: String((init?.headers as any)?.Authorization ?? ""), + }); + + return jsonResponse(200, { + usage: { + rolling: { status: "ok", percent: 5, resetsAt: new Date(Date.now() + 3600_000).toISOString() }, + weekly: { status: "ok", percent: 27, resetsAt: new Date(Date.now() + 86400_000).toISOString() }, + monthly: { status: "ok", percent: 13, resetsAt: new Date(Date.now() + 86400_000 * 30).toISOString() }, + }, + }); + }; + + const usage = await fetchOpencodeGoUsage("sk-test", { fetchFn }); + + expect(calls).toEqual([ + { + url: "https://opencode.ai/zen/go/v1/usage", + authorization: "Bearer sk-test", + }, + ]); + expect(usage.session).toBe(5); + expect(usage.weekly).toBe(27); + expect(usage.monthly).toBe(13); + expect(typeof usage.sessionResetsIn).toBe("string"); + expect(typeof usage.weeklyResetsIn).toBe("string"); + expect(typeof usage.monthlyResetsIn).toBe("string"); + }); + + it("returns explicit OpenCode Go errors for HTTP failures", async () => { + const badHttp = await fetchOpencodeGoUsage("sk-test", { + fetchFn: async () => jsonResponse(401, {}), + }); + expect(badHttp.error).toBe("HTTP 401"); + }); +});