pi config update
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user