pi update
This commit is contained in:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -8,3 +8,4 @@ pi/.pi/agent/usage-cache.json
|
|||||||
pi/.pi/agent/mcp-cache.json
|
pi/.pi/agent/mcp-cache.json
|
||||||
pi/.pi/agent/auth.json.current
|
pi/.pi/agent/auth.json.current
|
||||||
pi/.pi/agent/run-history.jsonl
|
pi/.pi/agent/run-history.jsonl
|
||||||
|
pi/.pi/context.md
|
||||||
|
|||||||
@@ -1,9 +1,9 @@
|
|||||||
---
|
---
|
||||||
name: explorer
|
name: explorer
|
||||||
description: Deep codebase and knowledge-base explorer using Claude Haiku with semantic search (QMD) and HDC-indexed context retrieval (opty). Use for thorough exploration, cross-cutting queries across docs and code, or when the local scout's Qwen model isn't cutting it.
|
description: Comprehensive codebase and knowledge-base explorer. Maps architecture, traces dependencies, synthesizes cross-cutting context with full code snippets and rationale. Use for deep refactoring, architectural decisions, or understanding complex subsystems.
|
||||||
tools: read, bash, write, mcp:qmd, mcp:opty
|
tools: read, bash, write, mcp:qmd, mcp:opty
|
||||||
model: anthropic/claude-haiku-4-5
|
model: anthropic/claude-haiku-4-5
|
||||||
output: context.md
|
output: /home/jonas/.pi/context.md
|
||||||
defaultProgress: true
|
defaultProgress: true
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -42,7 +42,11 @@ You are an explorer. Thoroughly investigate a codebase or knowledge base and ret
|
|||||||
- Medium: follow the most important cross-references, read critical sections
|
- Medium: follow the most important cross-references, read critical sections
|
||||||
- Thorough: trace all dependencies, check related files, synthesize a full picture
|
- Thorough: trace all dependencies, check related files, synthesize a full picture
|
||||||
|
|
||||||
## Output format (context.md)
|
## Output format (/home/jonas/.pi/context.md)
|
||||||
|
|
||||||
|
Write your report to **/home/jonas/.pi/context.md**. End your text response with:
|
||||||
|
|
||||||
|
> Report saved to: /home/jonas/.pi/context.md
|
||||||
|
|
||||||
# Exploration Context
|
# Exploration Context
|
||||||
|
|
||||||
|
|||||||
@@ -1,9 +1,16 @@
|
|||||||
/**
|
/**
|
||||||
* Claude Account Switch Extension
|
* Claude Account Switch Extension
|
||||||
*
|
*
|
||||||
* Switches between two Claude Pro accounts (personal and work).
|
* Switches between two Claude Pro accounts (personal and work) **without
|
||||||
* Tokens are saved from pi's own OAuth sessions (not Claude CLI),
|
* restarting pi**. Works by swapping auth.json files at the filesystem level,
|
||||||
* so token refresh works correctly.
|
* then reloading the auth storage and forcing an immediate token refresh to
|
||||||
|
* validate the switch.
|
||||||
|
*
|
||||||
|
* Why file-level swaps? Anthropic's OAuth rotates refresh tokens on every
|
||||||
|
* refresh. Calling authStorage.set() can appear to work, but the next
|
||||||
|
* getApiKey() call triggers refreshOAuthTokenWithLock(), which re-reads
|
||||||
|
* auth.json from disk — overwriting in-memory changes if persistence
|
||||||
|
* silently failed. Working at the file level avoids this entirely.
|
||||||
*
|
*
|
||||||
* Setup (one-time per account):
|
* Setup (one-time per account):
|
||||||
* 1. /login → authenticate with personal account
|
* 1. /login → authenticate with personal account
|
||||||
@@ -20,15 +27,17 @@ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
|||||||
import * as fs from "node:fs";
|
import * as fs from "node:fs";
|
||||||
import * as os from "node:os";
|
import * as os from "node:os";
|
||||||
import * as path from "node:path";
|
import * as path from "node:path";
|
||||||
|
import { execSync } from "node:child_process";
|
||||||
|
|
||||||
const HOME = os.homedir();
|
const HOME = os.homedir();
|
||||||
const AUTH_FILE = path.join(HOME, ".pi/agent/auth.json");
|
const AUTH_JSON = path.join(HOME, ".pi/agent/auth.json");
|
||||||
const AGENT_DIR = path.join(HOME, ".pi/agent");
|
|
||||||
const MARKER_FILE = path.join(HOME, ".pi/agent/auth.json.current");
|
const MARKER_FILE = path.join(HOME, ".pi/agent/auth.json.current");
|
||||||
const PROFILES_DIR = path.join(HOME, ".pi/agent/profiles");
|
const PROFILES_DIR = path.join(HOME, ".pi/agent/profiles");
|
||||||
|
|
||||||
type Account = "personal" | "work";
|
type Account = "personal" | "work";
|
||||||
|
|
||||||
|
// ── Profile helpers ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
function profilePath(account: Account): string {
|
function profilePath(account: Account): string {
|
||||||
return path.join(PROFILES_DIR, `auth-${account}.json`);
|
return path.join(PROFILES_DIR, `auth-${account}.json`);
|
||||||
}
|
}
|
||||||
@@ -37,17 +46,30 @@ function hasProfile(account: Account): boolean {
|
|||||||
return fs.existsSync(profilePath(account));
|
return fs.existsSync(profilePath(account));
|
||||||
}
|
}
|
||||||
|
|
||||||
function saveProfile(account: Account): void {
|
/**
|
||||||
|
* Save auth.json content directly to a profile file.
|
||||||
|
* This captures the exact on-disk state, including any tokens that were
|
||||||
|
* refreshed behind our back by the auth system.
|
||||||
|
*/
|
||||||
|
function saveCurrentAuthToProfile(account: Account): void {
|
||||||
fs.mkdirSync(PROFILES_DIR, { recursive: true });
|
fs.mkdirSync(PROFILES_DIR, { recursive: true });
|
||||||
fs.copyFileSync(AUTH_FILE, profilePath(account));
|
if (fs.existsSync(AUTH_JSON)) {
|
||||||
|
fs.copyFileSync(AUTH_JSON, profilePath(account));
|
||||||
fs.chmodSync(profilePath(account), 0o600);
|
fs.chmodSync(profilePath(account), 0o600);
|
||||||
fs.writeFileSync(MARKER_FILE, account);
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
function loadProfile(account: Account): void {
|
/**
|
||||||
fs.copyFileSync(profilePath(account), AUTH_FILE);
|
* Copy a profile file to auth.json. This is an atomic-ish swap that
|
||||||
fs.chmodSync(AUTH_FILE, 0o600);
|
* replaces the entire file rather than merging per-provider.
|
||||||
fs.writeFileSync(MARKER_FILE, account);
|
*/
|
||||||
|
function restoreProfileToAuth(account: Account): void {
|
||||||
|
fs.copyFileSync(profilePath(account), AUTH_JSON);
|
||||||
|
fs.chmodSync(AUTH_JSON, 0o600);
|
||||||
|
}
|
||||||
|
|
||||||
|
function setMarker(account: Account): void {
|
||||||
|
fs.writeFileSync(MARKER_FILE, account, "utf-8");
|
||||||
}
|
}
|
||||||
|
|
||||||
function getCurrentAccount(): Account | "unknown" {
|
function getCurrentAccount(): Account | "unknown" {
|
||||||
@@ -58,132 +80,223 @@ function getCurrentAccount(): Account | "unknown" {
|
|||||||
return "unknown";
|
return "unknown";
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Other session detection ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
function otherPiSessions(): number[] {
|
||||||
|
try {
|
||||||
|
const myPid = process.pid;
|
||||||
|
// Use a character class [c] trick so pgrep doesn't match its own process
|
||||||
|
const out = execSync("pgrep -f 'pi-[c]oding-agent' 2>/dev/null || true", {
|
||||||
|
encoding: "utf-8",
|
||||||
|
});
|
||||||
|
const pids = out
|
||||||
|
.trim()
|
||||||
|
.split("\n")
|
||||||
|
.map(Number)
|
||||||
|
.filter((p) => p && p !== myPid && !isNaN(p));
|
||||||
|
return pids;
|
||||||
|
} catch {
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function killOtherSessions(pids: number[]): number {
|
||||||
|
let killed = 0;
|
||||||
|
for (const pid of pids) {
|
||||||
|
try {
|
||||||
|
process.kill(pid);
|
||||||
|
killed++;
|
||||||
|
} catch {
|
||||||
|
// already dead or permission denied
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return killed;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── UI helpers ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
function statusLabel(account: Account | "unknown"): string {
|
function statusLabel(account: Account | "unknown"): string {
|
||||||
switch (account) {
|
switch (account) {
|
||||||
case "personal": return "🏠 personal";
|
case "personal":
|
||||||
case "work": return "💼 work";
|
return "🏠 personal";
|
||||||
default: return "❓ claude";
|
case "work":
|
||||||
|
return "💼 work";
|
||||||
|
default:
|
||||||
|
return "❓ claude";
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
// ── Extension ──────────────────────────────────────────────────────────────
|
||||||
* Sync auth.json into the active profile file whenever pi rotates tokens.
|
|
||||||
*
|
|
||||||
* Pi writes fresh OAuth tokens to auth.json (via direct write or atomic
|
|
||||||
* rename). Without this watcher the profile file would keep the original
|
|
||||||
* snapshot token, which Anthropic invalidates after its first use as a
|
|
||||||
* refresh-token rotation. The next /switch-claude would restore the dead
|
|
||||||
* refresh token and force a full re-login.
|
|
||||||
*
|
|
||||||
* We watch the agent directory (more reliable than watching the file
|
|
||||||
* directly across atomic renames) and copy auth.json → active profile
|
|
||||||
* on every change, debounced to avoid duplicate writes.
|
|
||||||
*/
|
|
||||||
function startProfileSyncer(getAccount: () => Account | "unknown"): fs.FSWatcher | null {
|
|
||||||
let debounceTimer: ReturnType<typeof setTimeout> | null = null;
|
|
||||||
|
|
||||||
const syncNow = () => {
|
|
||||||
const account = getAccount();
|
|
||||||
if (account === "unknown") return;
|
|
||||||
if (!fs.existsSync(AUTH_FILE)) return;
|
|
||||||
try {
|
|
||||||
const content = fs.readFileSync(AUTH_FILE, "utf8");
|
|
||||||
JSON.parse(content); // only sync valid JSON
|
|
||||||
const dest = profilePath(account);
|
|
||||||
fs.mkdirSync(PROFILES_DIR, { recursive: true });
|
|
||||||
fs.writeFileSync(dest, content, "utf8");
|
|
||||||
fs.chmodSync(dest, 0o600);
|
|
||||||
} catch {
|
|
||||||
// ignore transient read errors (file mid-write, etc.)
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
try {
|
|
||||||
return fs.watch(AGENT_DIR, { persistent: false }, (_event, filename) => {
|
|
||||||
if (filename !== "auth.json") return;
|
|
||||||
if (debounceTimer) clearTimeout(debounceTimer);
|
|
||||||
debounceTimer = setTimeout(syncNow, 200);
|
|
||||||
});
|
|
||||||
} catch {
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export default function (pi: ExtensionAPI) {
|
export default function (pi: ExtensionAPI) {
|
||||||
let currentAccount: Account | "unknown" = "unknown";
|
let currentAccount: Account | "unknown" = "unknown";
|
||||||
let syncer: fs.FSWatcher | null = null;
|
|
||||||
|
|
||||||
pi.on("session_start", async (_event, ctx) => {
|
pi.on("session_start", async (_event, ctx) => {
|
||||||
currentAccount = getCurrentAccount();
|
currentAccount = getCurrentAccount();
|
||||||
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
||||||
|
|
||||||
// Start watching auth.json so token refreshes are mirrored into
|
|
||||||
// the active profile file. Only one watcher is needed per session.
|
|
||||||
if (!syncer) {
|
|
||||||
syncer = startProfileSyncer(() => currentAccount);
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
pi.on("session_shutdown", async () => {
|
|
||||||
if (syncer) {
|
|
||||||
syncer.close();
|
|
||||||
syncer = null;
|
|
||||||
}
|
|
||||||
});
|
});
|
||||||
|
|
||||||
pi.registerCommand("switch-claude", {
|
pi.registerCommand("switch-claude", {
|
||||||
description: "Switch between personal (🏠) and work (💼) Claude accounts. Use 'save <name>' to save current login as a profile.",
|
description:
|
||||||
|
"Switch between personal (🏠) and work (💼) Claude accounts. Use 'save <name>' to save current login as a profile.",
|
||||||
handler: async (args, ctx) => {
|
handler: async (args, ctx) => {
|
||||||
|
const authStorage = ctx.modelRegistry.authStorage;
|
||||||
const trimmed = args?.trim() ?? "";
|
const trimmed = args?.trim() ?? "";
|
||||||
|
|
||||||
// Save current auth.json as a named profile
|
// ── Save current auth state as a named profile ──────────────────
|
||||||
if (trimmed.startsWith("save ")) {
|
if (trimmed.startsWith("save ")) {
|
||||||
const name = trimmed.slice(5).trim();
|
const name = trimmed.slice(5).trim();
|
||||||
if (name !== "personal" && name !== "work") {
|
if (name !== "personal" && name !== "work") {
|
||||||
ctx.ui.notify("Usage: /switch-claude save personal|work", "warning");
|
ctx.ui.notify(
|
||||||
return;
|
"Usage: /switch-claude save personal|work",
|
||||||
}
|
"warning",
|
||||||
saveProfile(name as Account);
|
);
|
||||||
currentAccount = name as Account;
|
|
||||||
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
|
||||||
ctx.ui.notify(`Saved current login as ${statusLabel(name as Account)} profile`, "info");
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Switch between profiles
|
if (!authStorage.has("anthropic")) {
|
||||||
|
ctx.ui.notify(
|
||||||
|
"No Anthropic credentials found. Run /login first.",
|
||||||
|
"warning",
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
saveCurrentAuthToProfile(name as Account);
|
||||||
|
currentAccount = name as Account;
|
||||||
|
setMarker(currentAccount);
|
||||||
|
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
||||||
|
ctx.ui.notify(
|
||||||
|
`Saved current login as ${statusLabel(name as Account)} profile`,
|
||||||
|
"info",
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Resolve target account (direct arg or interactive) ──────────
|
||||||
|
let newAccount: Account;
|
||||||
|
if (trimmed === "personal" || trimmed === "work") {
|
||||||
|
newAccount = trimmed;
|
||||||
|
} else if (trimmed === "") {
|
||||||
const personalLabel = `🏠 personal${currentAccount === "personal" ? " ← current" : ""}${!hasProfile("personal") ? " (no profile saved)" : ""}`;
|
const personalLabel = `🏠 personal${currentAccount === "personal" ? " ← current" : ""}${!hasProfile("personal") ? " (no profile saved)" : ""}`;
|
||||||
const workLabel = `💼 work${currentAccount === "work" ? " ← current" : ""}${!hasProfile("work") ? " (no profile saved)" : ""}`;
|
const workLabel = `💼 work${currentAccount === "work" ? " ← current" : ""}${!hasProfile("work") ? " (no profile saved)" : ""}`;
|
||||||
|
|
||||||
const choice = await ctx.ui.select("Switch Claude account:", [personalLabel, workLabel]);
|
const accountChoice = await ctx.ui.select(
|
||||||
if (choice === undefined) return;
|
"Switch Claude account:",
|
||||||
|
[personalLabel, workLabel],
|
||||||
const newAccount: Account = choice.startsWith("🏠") ? "personal" : "work";
|
);
|
||||||
|
if (accountChoice === undefined) return;
|
||||||
|
newAccount = accountChoice.startsWith("🏠") ? "personal" : "work";
|
||||||
|
} else {
|
||||||
|
ctx.ui.notify(
|
||||||
|
"Usage: /switch-claude [personal|work|save <name>]",
|
||||||
|
"warning",
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
if (newAccount === currentAccount) {
|
if (newAccount === currentAccount) {
|
||||||
ctx.ui.notify(`Already using ${statusLabel(newAccount)}`, "info");
|
ctx.ui.notify(
|
||||||
|
`Already using ${statusLabel(newAccount)}`,
|
||||||
|
"info",
|
||||||
|
);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Warn about other sessions ───────────────────────────────────
|
||||||
|
const otherPids = otherPiSessions();
|
||||||
|
if (otherPids.length > 0) {
|
||||||
|
const sessionChoice = await ctx.ui.select(
|
||||||
|
`⚠️ ${otherPids.length} other pi session(s) detected`,
|
||||||
|
[
|
||||||
|
"Continue anyway",
|
||||||
|
`Kill ${otherPids.length} other instance(s) and continue`,
|
||||||
|
"Cancel",
|
||||||
|
],
|
||||||
|
);
|
||||||
|
if (sessionChoice === undefined || sessionChoice.includes("Cancel"))
|
||||||
|
return;
|
||||||
|
|
||||||
|
if (sessionChoice.includes("Kill")) {
|
||||||
|
const killed = killOtherSessions(otherPids);
|
||||||
|
ctx.ui.notify(`Killed ${killed} pi session(s)`, "info");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if (!hasProfile(newAccount)) {
|
if (!hasProfile(newAccount)) {
|
||||||
ctx.ui.notify(
|
ctx.ui.notify(
|
||||||
`No profile saved for ${newAccount}. Run /login then /switch-claude save ${newAccount}`,
|
`No profile saved for ${newAccount}. Run /login then /switch-claude save ${newAccount}`,
|
||||||
"warning"
|
"warning",
|
||||||
);
|
);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Perform the switch ──────────────────────────────────────────
|
||||||
try {
|
try {
|
||||||
// Persist any token refreshes pi made since the last save so
|
// 1. Snapshot current auth.json → outgoing profile.
|
||||||
// we don't restore a stale refresh token when we come back.
|
// This captures any tokens that were silently refreshed
|
||||||
if (currentAccount !== "unknown") saveProfile(currentAccount);
|
// since the last save (the file is the source of truth,
|
||||||
loadProfile(newAccount);
|
// not the in-memory snapshot from getAll()).
|
||||||
|
if (currentAccount !== "unknown") {
|
||||||
|
saveCurrentAuthToProfile(currentAccount);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Copy incoming profile → auth.json (full file replace).
|
||||||
|
restoreProfileToAuth(newAccount);
|
||||||
|
|
||||||
|
// 3. Tell AuthStorage to re-read the file. This updates
|
||||||
|
// the in-memory credential cache from the new auth.json.
|
||||||
|
authStorage.reload();
|
||||||
|
|
||||||
|
// 4. Force an immediate token refresh to validate the switch.
|
||||||
|
// If the stored refresh token is stale, this will fail now
|
||||||
|
// rather than on the next chat message.
|
||||||
|
const apiKey = await authStorage.getApiKey("anthropic");
|
||||||
|
|
||||||
|
if (!apiKey) {
|
||||||
|
// Refresh failed → roll back to the previous account.
|
||||||
|
if (currentAccount !== "unknown") {
|
||||||
|
restoreProfileToAuth(currentAccount);
|
||||||
|
authStorage.reload();
|
||||||
|
}
|
||||||
|
ctx.ui.notify(
|
||||||
|
`❌ Switch failed: could not authenticate as ${newAccount}. ` +
|
||||||
|
`The saved refresh token may have expired. ` +
|
||||||
|
`Run /login then /switch-claude save ${newAccount} to re-save.`,
|
||||||
|
"error",
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 5. Success — the refresh worked, auth.json now has fresh
|
||||||
|
// tokens. Save them back to the profile so next switch
|
||||||
|
// has the latest refresh token.
|
||||||
|
saveCurrentAuthToProfile(newAccount);
|
||||||
|
|
||||||
currentAccount = newAccount;
|
currentAccount = newAccount;
|
||||||
|
setMarker(currentAccount);
|
||||||
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
|
||||||
pi.events.emit("claude-account:switched", { account: newAccount });
|
ctx.ui.notify(
|
||||||
ctx.ui.notify(`Switched to ${statusLabel(newAccount)} — restart pi to apply`, "info");
|
`Switched to ${statusLabel(newAccount)} ✓`,
|
||||||
|
"info",
|
||||||
|
);
|
||||||
} catch (e: unknown) {
|
} catch (e: unknown) {
|
||||||
|
// Something went wrong → try to roll back.
|
||||||
|
try {
|
||||||
|
if (currentAccount !== "unknown" && hasProfile(currentAccount)) {
|
||||||
|
restoreProfileToAuth(currentAccount);
|
||||||
|
authStorage.reload();
|
||||||
|
}
|
||||||
|
} catch {
|
||||||
|
// rollback failed too — nothing more we can do
|
||||||
|
}
|
||||||
|
|
||||||
const msg = e instanceof Error ? e.message : String(e);
|
const msg = e instanceof Error ? e.message : String(e);
|
||||||
ctx.ui.notify(`Failed to switch: ${msg}`, "error");
|
ctx.ui.notify(
|
||||||
|
`❌ Switch failed: ${msg}. Rolled back to ${statusLabel(currentAccount)}. ` +
|
||||||
|
`You may need to /login and /switch-claude save ${newAccount}.`,
|
||||||
|
"error",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -37,7 +37,17 @@ import {
|
|||||||
type UsageData,
|
type UsageData,
|
||||||
} from "./core";
|
} from "./core";
|
||||||
|
|
||||||
const CACHE_TTL_MS = 15 * 60 * 1000; // reuse cached data for 15 min
|
// 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; // 1 hour back-off after 429
|
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; // 1 hour back-off after 429
|
||||||
const STATUS_KEY = "usage-bars";
|
const STATUS_KEY = "usage-bars";
|
||||||
|
|
||||||
@@ -297,6 +307,17 @@ interface UsageState extends UsageByProvider {
|
|||||||
activeProvider: ProviderKey | null;
|
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) {
|
export default function (pi: ExtensionAPI) {
|
||||||
const endpoints = resolveUsageEndpoints();
|
const endpoints = resolveUsageEndpoints();
|
||||||
const state: UsageState = {
|
const state: UsageState = {
|
||||||
@@ -311,6 +332,8 @@ export default function (pi: ExtensionAPI) {
|
|||||||
|
|
||||||
let pollInFlight: Promise<void> | null = null;
|
let pollInFlight: Promise<void> | null = null;
|
||||||
let pollQueued = false;
|
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;
|
let ctx: any = null;
|
||||||
|
|
||||||
function renderPercent(theme: any, value: number): string {
|
function renderPercent(theme: any, value: number): string {
|
||||||
@@ -410,7 +433,7 @@ export default function (pi: ExtensionAPI) {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
async function runPoll() {
|
async function runPoll(options: PollOptions = {}) {
|
||||||
const auth = readAuth();
|
const auth = readAuth();
|
||||||
const active = state.activeProvider;
|
const active = state.activeProvider;
|
||||||
|
|
||||||
@@ -420,25 +443,21 @@ export default function (pi: ExtensionAPI) {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Shared disk cache check ---
|
|
||||||
// All pi sessions read and write the same cache file so that only one
|
|
||||||
// process hits the API per CACHE_TTL_MS window, no matter how many
|
|
||||||
// sessions are open at once.
|
|
||||||
const cache = readUsageCache();
|
const cache = readUsageCache();
|
||||||
const now = Date.now();
|
const now = Date.now();
|
||||||
|
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||||
|
|
||||||
// Respect cross-session rate-limit back-off written by any session.
|
// Respect cross-session rate-limit back-off written by any session.
|
||||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||||
if (now < blockedUntil) {
|
if (now < blockedUntil) {
|
||||||
// Use whatever data is in the cache so the bar still shows something.
|
|
||||||
if (cache?.data?.[active]) state[active] = cache.data[active]!;
|
if (cache?.data?.[active]) state[active] = cache.data[active]!;
|
||||||
state.lastPoll = now;
|
state.lastPoll = now;
|
||||||
updateStatus();
|
updateStatus();
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// If another session already polled recently, use their result.
|
// Use shared disk cache unless forceFresh is set (e.g. after account switch).
|
||||||
if (cache && now - cache.timestamp < CACHE_TTL_MS && cache.data?.[active]) {
|
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||||
state[active] = cache.data[active]!;
|
state[active] = cache.data[active]!;
|
||||||
state.lastPoll = now;
|
state.lastPoll = now;
|
||||||
updateStatus();
|
updateStatus();
|
||||||
@@ -446,13 +465,6 @@ export default function (pi: ExtensionAPI) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// --- Proactive token refresh ---
|
// --- Proactive token refresh ---
|
||||||
// Before hitting the API, check whether the stored access token is expired.
|
|
||||||
// This is the main cause of HTTP 401 errors: switching accounts via
|
|
||||||
// /switch-claude restores a profile whose access token has since expired
|
|
||||||
// (the refresh token is still valid). We use pi's own OAuth resolver so
|
|
||||||
// the new tokens are written back to auth.json and the profile stays in
|
|
||||||
// sync. This is safe at turn_start / session_start because pi hasn't made
|
|
||||||
// any Claude API calls yet, so there's no parallel refresh to conflict with.
|
|
||||||
const oauthId = providerToOAuthProviderId(active);
|
const oauthId = providerToOAuthProviderId(active);
|
||||||
let effectiveAuth = auth;
|
let effectiveAuth = auth;
|
||||||
if (oauthId && active !== "zai") {
|
if (oauthId && active !== "zai") {
|
||||||
@@ -506,16 +518,8 @@ export default function (pi: ExtensionAPI) {
|
|||||||
|
|
||||||
state[active] = result;
|
state[active] = result;
|
||||||
|
|
||||||
// Write result + rate-limit state to shared cache so other sessions
|
|
||||||
// don't need to re-hit the API within CACHE_TTL_MS.
|
|
||||||
//
|
|
||||||
// Error results (other than 429) are NOT cached: they should be retried
|
|
||||||
// on the next input instead of being replayed from cache for 15 minutes.
|
|
||||||
// The most common error is HTTP 401 (expired token after an account switch)
|
|
||||||
// which resolves on the very next poll once the token is refreshed above.
|
|
||||||
if (result.error) {
|
if (result.error) {
|
||||||
if (result.error === "HTTP 429") {
|
if (result.error === "HTTP 429") {
|
||||||
// Write rate-limit backoff but preserve the last good data in cache.
|
|
||||||
const nextCache: import("./core").UsageCache = {
|
const nextCache: import("./core").UsageCache = {
|
||||||
timestamp: cache?.timestamp ?? now,
|
timestamp: cache?.timestamp ?? now,
|
||||||
data: { ...(cache?.data ?? {}) },
|
data: { ...(cache?.data ?? {}) },
|
||||||
@@ -541,7 +545,7 @@ export default function (pi: ExtensionAPI) {
|
|||||||
updateStatus();
|
updateStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
async function poll() {
|
async function poll(options: PollOptions = {}) {
|
||||||
if (pollInFlight) {
|
if (pollInFlight) {
|
||||||
pollQueued = true;
|
pollQueued = true;
|
||||||
await pollInFlight;
|
await pollInFlight;
|
||||||
@@ -550,7 +554,7 @@ export default function (pi: ExtensionAPI) {
|
|||||||
|
|
||||||
do {
|
do {
|
||||||
pollQueued = false;
|
pollQueued = false;
|
||||||
pollInFlight = runPoll()
|
pollInFlight = runPoll(options)
|
||||||
.catch(() => {
|
.catch(() => {
|
||||||
// Never crash extension event handlers on transient polling errors.
|
// Never crash extension event handlers on transient polling errors.
|
||||||
})
|
})
|
||||||
@@ -562,6 +566,22 @@ export default function (pi: ExtensionAPI) {
|
|||||||
} while (pollQueued);
|
} while (pollQueued);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function startStreamingTimer() {
|
||||||
|
if (streamingTimer !== null) return; // already running
|
||||||
|
streamingTimer = setInterval(() => {
|
||||||
|
void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||||
|
}, STREAMING_POLL_INTERVAL_MS);
|
||||||
|
}
|
||||||
|
|
||||||
|
function stopStreamingTimer() {
|
||||||
|
if (streamingTimer !== null) {
|
||||||
|
clearInterval(streamingTimer);
|
||||||
|
streamingTimer = null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Session lifecycle ────────────────────────────────────────────────────
|
||||||
|
|
||||||
pi.on("session_start", async (_event, _ctx) => {
|
pi.on("session_start", async (_event, _ctx) => {
|
||||||
ctx = _ctx;
|
ctx = _ctx;
|
||||||
updateProviderFrom(_ctx.model);
|
updateProviderFrom(_ctx.model);
|
||||||
@@ -569,19 +589,13 @@ export default function (pi: ExtensionAPI) {
|
|||||||
});
|
});
|
||||||
|
|
||||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||||
|
stopStreamingTimer();
|
||||||
if (_ctx?.hasUI) {
|
if (_ctx?.hasUI) {
|
||||||
_ctx.ui.setStatus(STATUS_KEY, undefined);
|
_ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
// Refresh usage on every turn (like claude-pulse's UserPromptSubmit hook).
|
// ── Model change ─────────────────────────────────────────────────────────
|
||||||
// The disk cache means the API is only hit at most once per CACHE_TTL_MS
|
|
||||||
// regardless of how many turns or sessions are active.
|
|
||||||
pi.on("turn_start", async (_event, _ctx) => {
|
|
||||||
ctx = _ctx;
|
|
||||||
updateProviderFrom(_ctx.model);
|
|
||||||
await poll();
|
|
||||||
});
|
|
||||||
|
|
||||||
pi.on("model_select", async (event, _ctx) => {
|
pi.on("model_select", async (event, _ctx) => {
|
||||||
ctx = _ctx;
|
ctx = _ctx;
|
||||||
@@ -589,6 +603,58 @@ export default function (pi: ExtensionAPI) {
|
|||||||
if (changed) await poll();
|
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);
|
||||||
|
});
|
||||||
|
|
||||||
|
// ── 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();
|
||||||
|
});
|
||||||
|
|
||||||
|
// 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 },
|
||||||
|
};
|
||||||
|
delete nextCache.data.claude;
|
||||||
|
writeUsageCache(nextCache);
|
||||||
|
}
|
||||||
|
void poll({ forceFresh: true });
|
||||||
|
});
|
||||||
|
|
||||||
|
// ── /usage command ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
pi.registerCommand("usage", {
|
pi.registerCommand("usage", {
|
||||||
description: "Show API usage for all subscriptions",
|
description: "Show API usage for all subscriptions",
|
||||||
handler: async (_args, _ctx) => {
|
handler: async (_args, _ctx) => {
|
||||||
@@ -609,13 +675,8 @@ export default function (pi: ExtensionAPI) {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
} finally {
|
} finally {
|
||||||
await poll();
|
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
|
|
||||||
// Re-poll immediately when the Claude account is switched via /switch-claude
|
|
||||||
pi.events.on("claude-account:switched", () => {
|
|
||||||
void poll();
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user