pi config update

This commit is contained in:
Jonas H
2026-03-19 07:58:49 +01:00
parent a3c9183485
commit 871caa5adc
24 changed files with 6198 additions and 555 deletions

View File

@@ -1,9 +1,12 @@
/**
* Usage Extension - Minimal API usage indicator for pi
*
* Shows Codex (OpenAI), Anthropic (Claude), Z.AI, and optionally
* Google Gemini CLI / Antigravity usage as color-coded percentages
* in the footer status bar.
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
* via two channels:
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
*
* Rendering / footer layout is handled by the separate footer-display extension.
*/
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
@@ -37,20 +40,39 @@ import {
type UsageData,
} from "./core";
// Disk cache TTL for idle/background reads (session start, etc.)
const CACHE_TTL_MS = 15 * 60 * 1000;
// Shorter TTL for event-driven polls (after prompt submit / after turn end).
// With the shared disk cache, only one pi instance per ACTIVE_CACHE_TTL_MS window
// will actually hit the API — the rest will read from the cached result.
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
// How often to re-poll while the model is actively streaming / running tools.
// Combined with the shared disk cache this means at most one HTTP request per
// STREAMING_POLL_INTERVAL_MS regardless of how many pi sessions are open.
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; // 1 hour back-off after 429
const STATUS_KEY = "usage-bars";
// ---------------------------------------------------------------------------
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
// ---------------------------------------------------------------------------
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
const BRAILLE_EMPTY = "\u28C0";
const BAR_WIDTH = 5;
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
const v = clampPercent(value);
const levels = BRAILLE_GRADIENT.length - 1;
const totalSteps = width * levels;
const filledSteps = Math.round((v / 100) * totalSteps);
const full = Math.floor(filledSteps / levels);
const partial = filledSteps % levels;
const empty = width - full - (partial ? 1 : 0);
const color = colorForPercent(v);
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
}
function renderBrailleBarWide(theme: any, value: number): string {
return renderBrailleBar(theme, value, 12);
}
const PROVIDER_LABELS: Record<ProviderKey, string> = {
codex: "Codex",
claude: "Claude",
@@ -59,6 +81,9 @@ const PROVIDER_LABELS: Record<ProviderKey, string> = {
antigravity: "Antigravity",
};
// ---------------------------------------------------------------------------
// /usage command popup
// ---------------------------------------------------------------------------
interface SubscriptionItem {
name: string;
provider: ProviderKey;
@@ -81,14 +106,8 @@ class UsageSelectorComponent extends Container implements Focusable {
private fetchAllFn: () => Promise<UsageByProvider>;
private _focused = false;
get focused(): boolean {
return this._focused;
}
set focused(value: boolean) {
this._focused = value;
this.searchInput.focused = value;
}
get focused(): boolean { return this._focused; }
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
constructor(
tui: any,
@@ -106,19 +125,15 @@ class UsageSelectorComponent extends Container implements Focusable {
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
this.addChild(new Spacer(1));
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
this.addChild(this.hintText);
this.addChild(new Spacer(1));
this.searchInput = new Input();
this.addChild(this.searchInput);
this.addChild(new Spacer(1));
this.listContainer = new Container();
this.addChild(this.listContainer);
this.addChild(new Spacer(1));
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
this.fetchAllFn()
@@ -149,7 +164,6 @@ class UsageSelectorComponent extends Container implements Focusable {
{ key: "gemini", name: "Gemini" },
{ key: "antigravity", name: "Antigravity" },
];
this.allItems = [];
for (const p of providers) {
if (results[p.key] !== null) {
@@ -161,7 +175,6 @@ class UsageSelectorComponent extends Container implements Focusable {
});
}
}
this.filteredItems = this.allItems;
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
}
@@ -178,23 +191,12 @@ class UsageSelectorComponent extends Container implements Focusable {
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
}
private renderBar(pct: number, width = 16): string {
const value = clampPercent(pct);
const filled = Math.round((value / 100) * width);
const color = colorForPercent(value);
const full = "█".repeat(Math.max(0, filled));
const empty = "░".repeat(Math.max(0, width - filled));
return this.theme.fg(color, full) + this.theme.fg("dim", empty);
}
private renderItem(item: SubscriptionItem, isSelected: boolean) {
const t = this.theme;
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
const indent = " ";
if (!item.data) {
@@ -204,69 +206,45 @@ class UsageSelectorComponent extends Container implements Focusable {
} else {
const session = clampPercent(item.data.session);
const weekly = clampPercent(item.data.weekly);
const sessionReset = item.data.sessionResetsIn
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`)
: "";
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
const weeklyReset = item.data.weeklyResetsIn
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`)
: "";
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
this.listContainer.addChild(
new Text(
indent +
t.fg("muted", "Session ") +
this.renderBar(session) +
" " +
t.fg(colorForPercent(session), `${session}%`.padStart(4)) +
sessionReset,
0,
0,
),
);
this.listContainer.addChild(
new Text(
indent +
t.fg("muted", "Weekly ") +
this.renderBar(weekly) +
" " +
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) +
weeklyReset,
0,
0,
),
);
this.listContainer.addChild(new Text(
indent + t.fg("muted", "Session ") +
renderBrailleBarWide(t, session) + " " +
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
0, 0,
));
this.listContainer.addChild(new Text(
indent + t.fg("muted", "Weekly ") +
renderBrailleBarWide(t, weekly) + " " +
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
0, 0,
));
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
this.listContainer.addChild(
new Text(
indent +
t.fg("muted", "Extra ") +
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
0,
0,
),
);
this.listContainer.addChild(new Text(
indent + t.fg("muted", "Extra ") +
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
0, 0,
));
}
}
this.listContainer.addChild(new Spacer(1));
}
private updateList() {
this.listContainer.clear();
if (this.loading) {
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
return;
}
if (this.filteredItems.length === 0) {
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
return;
}
for (let i = 0; i < this.filteredItems.length; i++) {
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
}
@@ -274,91 +252,58 @@ class UsageSelectorComponent extends Container implements Focusable {
handleInput(keyData: string): void {
const kb = getEditorKeybindings();
if (kb.matches(keyData, "selectUp")) {
if (this.filteredItems.length === 0) return;
this.selectedIndex =
this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
this.updateList();
return;
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
this.updateList(); return;
}
if (kb.matches(keyData, "selectDown")) {
if (this.filteredItems.length === 0) return;
this.selectedIndex =
this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
this.updateList();
return;
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
this.updateList(); return;
}
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
this.onCancelCallback();
return;
this.onCancelCallback(); return;
}
this.searchInput.handleInput(keyData);
this.filterItems(this.searchInput.getValue());
this.updateList();
}
}
// ---------------------------------------------------------------------------
// Extension state
// ---------------------------------------------------------------------------
interface UsageState extends UsageByProvider {
lastPoll: number;
activeProvider: ProviderKey | null;
}
interface PollOptions {
/** Override the disk-cache TTL for this poll (default: CACHE_TTL_MS). */
cacheTtl?: number;
/**
* Skip the shared disk-cache TTL check entirely and always fetch from the
* API. Used after an account switch where the cached data belongs to a
* different account.
*/
forceFresh?: boolean;
}
export default function (pi: ExtensionAPI) {
const endpoints = resolveUsageEndpoints();
const state: UsageState = {
codex: null,
claude: null,
zai: null,
gemini: null,
antigravity: null,
lastPoll: 0,
activeProvider: null,
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
lastPoll: 0, activeProvider: null,
};
let pollInFlight: Promise<void> | null = null;
let pollQueued = false;
/** Timer running during an active agent loop to refresh usage periodically. */
let streamingTimer: ReturnType<typeof setInterval> | null = null;
let ctx: any = null;
function renderPercent(theme: any, value: number): string {
const v = clampPercent(value);
return theme.fg(colorForPercent(v), `${v}%`);
}
function renderBar(theme: any, value: number): string {
const v = clampPercent(value);
const width = 8;
const filled = Math.round((v / 100) * width);
const full = "█".repeat(Math.max(0, Math.min(width, filled)));
const empty = "░".repeat(Math.max(0, width - filled));
return theme.fg(colorForPercent(v), full) + theme.fg("dim", empty);
}
function pickDataForProvider(provider: ProviderKey | null): UsageData | null {
if (!provider) return null;
return state[provider];
}
// ---------------------------------------------------------------------------
// Status update
// ---------------------------------------------------------------------------
function updateStatus() {
const active = state.activeProvider;
const data = pickDataForProvider(active);
const data = active ? state[active] : null;
// Always emit event for other extensions (e.g. footer-display)
if (data && !data.error) {
pi.events.emit("usage:update", {
session: data.session,
@@ -370,6 +315,8 @@ export default function (pi: ExtensionAPI) {
if (!ctx?.hasUI) return;
const theme = ctx.ui.theme;
if (!active) {
ctx.ui.setStatus(STATUS_KEY, undefined);
return;
@@ -381,128 +328,92 @@ export default function (pi: ExtensionAPI) {
return;
}
const theme = ctx.ui.theme;
const label = PROVIDER_LABELS[active];
if (!data) {
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", `${label} usage: loading…`));
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
return;
}
if (data.error) {
const cache = readUsageCache();
const blockedUntil = active ? (cache?.rateLimitedUntil?.[active] ?? 0) : 0;
const backoffNote = blockedUntil > Date.now()
? ` retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m`
: "";
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${label} usage unavailable (${data.error}${backoffNote})`));
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
const note = blockedUntil > Date.now()
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
return;
}
const session = clampPercent(data.session);
const weekly = clampPercent(data.weekly);
const sessionReset = data.sessionResetsIn ? theme.fg("dim", `${data.sessionResetsIn}`) : "";
const weeklyReset = data.weeklyResetsIn ? theme.fg("dim", `${data.weeklyResetsIn}`) : "";
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
const status =
theme.fg("dim", `${label} `) +
theme.fg("muted", "S ") +
renderBar(theme, session) +
" " +
renderPercent(theme, session) +
sessionReset +
theme.fg("muted", " W ") +
renderBar(theme, weekly) +
" " +
renderPercent(theme, weekly) +
weeklyReset;
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
ctx.ui.setStatus(STATUS_KEY, status);
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
}
function updateProviderFrom(modelLike: any): boolean {
const previous = state.activeProvider;
state.activeProvider = detectProvider(modelLike);
if (previous !== state.activeProvider) {
updateStatus();
return true;
}
if (previous !== state.activeProvider) { updateStatus(); return true; }
return false;
}
// ---------------------------------------------------------------------------
// Polling
// ---------------------------------------------------------------------------
async function runPoll(options: PollOptions = {}) {
const auth = readAuth();
const active = state.activeProvider;
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
state.lastPoll = Date.now();
updateStatus();
return;
state.lastPoll = Date.now(); updateStatus(); return;
}
const cache = readUsageCache();
const now = Date.now();
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
// Respect cross-session rate-limit back-off written by any session.
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
if (now < blockedUntil) {
if (cache?.data?.[active]) state[active] = cache.data[active]!;
state.lastPoll = now;
updateStatus();
return;
state.lastPoll = now; updateStatus(); return;
}
// Use shared disk cache unless forceFresh is set (e.g. after account switch).
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
state[active] = cache.data[active]!;
state.lastPoll = now;
updateStatus();
return;
state.lastPoll = now; updateStatus(); return;
}
// --- Proactive token refresh ---
const oauthId = providerToOAuthProviderId(active);
let effectiveAuth = auth;
if (oauthId && active !== "zai") {
const creds = auth[oauthId as keyof typeof auth] as
| { access?: string; refresh?: string; expires?: number }
| undefined;
| { access?: string; refresh?: string; expires?: number } | undefined;
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
const tokenExpiredOrMissing =
!creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
if (tokenExpiredOrMissing && creds?.refresh) {
try {
const refreshed = await ensureFreshAuthForProviders([oauthId as OAuthProviderId], {
auth,
persist: true,
});
const refreshed = await ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
if (refreshed.auth) effectiveAuth = refreshed.auth;
} catch {
// Ignore refresh errors — fall through with existing auth
}
} catch {}
}
}
let result: UsageData;
if (active === "codex") {
const access = effectiveAuth["openai-codex"]?.access;
result = access
? await fetchCodexUsage(access)
result = access ? await fetchCodexUsage(access)
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
} else if (active === "claude") {
const access = effectiveAuth.anthropic?.access;
result = access
? await fetchClaudeUsage(access)
result = access ? await fetchClaudeUsage(access)
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
} else if (active === "zai") {
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
result = token
? await fetchZaiUsage(token, { endpoints })
result = token ? await fetchZaiUsage(token, { endpoints })
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
} else if (active === "gemini") {
const creds = effectiveAuth["google-gemini-cli"];
@@ -523,14 +434,10 @@ export default function (pi: ExtensionAPI) {
const nextCache: import("./core").UsageCache = {
timestamp: cache?.timestamp ?? now,
data: { ...(cache?.data ?? {}) },
rateLimitedUntil: {
...(cache?.rateLimitedUntil ?? {}),
[active]: now + RATE_LIMITED_BACKOFF_MS,
},
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
};
writeUsageCache(nextCache);
}
// All other errors: don't update cache — next turn will retry from scratch.
} else {
const nextCache: import("./core").UsageCache = {
timestamp: now,
@@ -546,41 +453,24 @@ export default function (pi: ExtensionAPI) {
}
async function poll(options: PollOptions = {}) {
if (pollInFlight) {
pollQueued = true;
await pollInFlight;
return;
}
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
do {
pollQueued = false;
pollInFlight = runPoll(options)
.catch(() => {
// Never crash extension event handlers on transient polling errors.
})
.finally(() => {
pollInFlight = null;
});
pollInFlight = runPoll(options).catch(() => {}).finally(() => { pollInFlight = null; });
await pollInFlight;
} while (pollQueued);
}
function startStreamingTimer() {
if (streamingTimer !== null) return; // already running
streamingTimer = setInterval(() => {
void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
}, STREAMING_POLL_INTERVAL_MS);
if (streamingTimer !== null) return;
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
}
function stopStreamingTimer() {
if (streamingTimer !== null) {
clearInterval(streamingTimer);
streamingTimer = null;
}
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
}
// ── Session lifecycle ────────────────────────────────────────────────────
// ── Lifecycle ────────────────────────────────────────────────────────────
pi.on("session_start", async (_event, _ctx) => {
ctx = _ctx;
@@ -590,63 +480,34 @@ export default function (pi: ExtensionAPI) {
pi.on("session_shutdown", async (_event, _ctx) => {
stopStreamingTimer();
if (_ctx?.hasUI) {
_ctx.ui.setStatus(STATUS_KEY, undefined);
}
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
});
// ── Model change ─────────────────────────────────────────────────────────
pi.on("model_select", async (event, _ctx) => {
ctx = _ctx;
const changed = updateProviderFrom(event.model ?? _ctx.model);
if (changed) await poll();
});
// Keep provider detection up-to-date across turns (cheap, no API call unless
// the provider changed).
pi.on("turn_start", (_event, _ctx) => {
ctx = _ctx;
updateProviderFrom(_ctx.model);
});
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
// ── Agent loop ───────────────────────────────────────────────────────────
// Poll when the user submits a prompt — captures usage right before the
// new turn (useful to see current state before tokens are consumed).
pi.on("before_agent_start", async (_event, _ctx) => {
ctx = _ctx;
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
});
// Start interval timer so usage stays fresh during long-running agents
// (lots of tool calls, extended thinking, etc.).
pi.on("agent_start", (_event, _ctx) => {
ctx = _ctx;
startStreamingTimer();
});
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
// When the agent finishes, stop the timer and do a final fresh poll to
// capture the usage that was just consumed.
pi.on("agent_end", async (_event, _ctx) => {
ctx = _ctx;
stopStreamingTimer();
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
});
// ── Account switch ───────────────────────────────────────────────────────
// Re-poll immediately when the Claude account is switched via /switch-claude.
// We must invalidate the claude entry in the shared disk cache first —
// otherwise poll() would serve the previous account's data which is still
// within CACHE_TTL_MS.
pi.events.on("claude-account:switched", () => {
const cache = readUsageCache();
if (cache?.data?.claude) {
const nextCache: import("./core").UsageCache = {
...cache,
data: { ...cache.data },
};
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
delete nextCache.data.claude;
writeUsageCache(nextCache);
}
@@ -660,18 +521,14 @@ export default function (pi: ExtensionAPI) {
handler: async (_args, _ctx) => {
ctx = _ctx;
updateProviderFrom(_ctx.model);
try {
if (_ctx?.hasUI) {
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
const selector = new UsageSelectorComponent(
tui,
theme,
state.activeProvider,
return new UsageSelectorComponent(
tui, theme, state.activeProvider,
() => fetchAllUsages({ endpoints }),
() => done(),
);
return selector;
});
}
} finally {