diff --git a/.gitignore b/.gitignore index 30d48f0..b2eccb1 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ pi/.pi/agent/usage-cache.json pi/.pi/agent/mcp-cache.json pi/.pi/agent/auth.json.current pi/.pi/agent/run-history.jsonl +pi/.pi/context.md diff --git a/pi/.pi/agent/agents/explorer.md b/pi/.pi/agent/agents/explorer.md index eeca4d5..f0f7b9a 100644 --- a/pi/.pi/agent/agents/explorer.md +++ b/pi/.pi/agent/agents/explorer.md @@ -1,9 +1,9 @@ --- 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 model: anthropic/claude-haiku-4-5 -output: context.md +output: /home/jonas/.pi/context.md 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 - 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 diff --git a/pi/.pi/agent/extensions/claude-account-switch.ts b/pi/.pi/agent/extensions/claude-account-switch.ts index abe33bb..d7e30d6 100644 --- a/pi/.pi/agent/extensions/claude-account-switch.ts +++ b/pi/.pi/agent/extensions/claude-account-switch.ts @@ -1,9 +1,16 @@ /** * Claude Account Switch Extension * - * Switches between two Claude Pro accounts (personal and work). - * Tokens are saved from pi's own OAuth sessions (not Claude CLI), - * so token refresh works correctly. + * Switches between two Claude Pro accounts (personal and work) **without + * restarting pi**. Works by swapping auth.json files at the filesystem level, + * 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): * 1. /login → authenticate with personal account @@ -12,23 +19,25 @@ * 4. /switch-claude save work * * Usage: - * /switch-claude — pick account interactively - * /switch-claude save — save current pi login as a named profile + * /switch-claude — pick account interactively + * /switch-claude save — save current pi login as a named profile */ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; +import { execSync } from "node:child_process"; const HOME = os.homedir(); -const AUTH_FILE = path.join(HOME, ".pi/agent/auth.json"); -const AGENT_DIR = path.join(HOME, ".pi/agent"); +const AUTH_JSON = path.join(HOME, ".pi/agent/auth.json"); const MARKER_FILE = path.join(HOME, ".pi/agent/auth.json.current"); const PROFILES_DIR = path.join(HOME, ".pi/agent/profiles"); type Account = "personal" | "work"; +// ── Profile helpers ──────────────────────────────────────────────────────── + function profilePath(account: Account): string { return path.join(PROFILES_DIR, `auth-${account}.json`); } @@ -37,17 +46,30 @@ function hasProfile(account: Account): boolean { 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.copyFileSync(AUTH_FILE, profilePath(account)); - fs.chmodSync(profilePath(account), 0o600); - fs.writeFileSync(MARKER_FILE, account); + if (fs.existsSync(AUTH_JSON)) { + fs.copyFileSync(AUTH_JSON, profilePath(account)); + fs.chmodSync(profilePath(account), 0o600); + } } -function loadProfile(account: Account): void { - fs.copyFileSync(profilePath(account), AUTH_FILE); - fs.chmodSync(AUTH_FILE, 0o600); - fs.writeFileSync(MARKER_FILE, account); +/** + * Copy a profile file to auth.json. This is an atomic-ish swap that + * replaces the entire file rather than merging per-provider. + */ +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" { @@ -58,132 +80,223 @@ function getCurrentAccount(): Account | "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 { switch (account) { - case "personal": return "🏠 personal"; - case "work": return "💼 work"; - default: return "❓ claude"; + case "personal": + return "🏠 personal"; + case "work": + return "💼 work"; + default: + return "❓ claude"; } } -/** - * 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 | 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; - } -} +// ── Extension ────────────────────────────────────────────────────────────── export default function (pi: ExtensionAPI) { let currentAccount: Account | "unknown" = "unknown"; - let syncer: fs.FSWatcher | null = null; pi.on("session_start", async (_event, ctx) => { currentAccount = getCurrentAccount(); 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", { - description: "Switch between personal (🏠) and work (💼) Claude accounts. Use 'save ' to save current login as a profile.", + description: + "Switch between personal (🏠) and work (💼) Claude accounts. Use 'save ' to save current login as a profile.", handler: async (args, ctx) => { + const authStorage = ctx.modelRegistry.authStorage; const trimmed = args?.trim() ?? ""; - // Save current auth.json as a named profile + // ── Save current auth state as a named profile ────────────────── if (trimmed.startsWith("save ")) { const name = trimmed.slice(5).trim(); if (name !== "personal" && name !== "work") { - ctx.ui.notify("Usage: /switch-claude save personal|work", "warning"); + ctx.ui.notify( + "Usage: /switch-claude save personal|work", + "warning", + ); return; } - saveProfile(name as Account); + + 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"); + ctx.ui.notify( + `Saved current login as ${statusLabel(name as Account)} profile`, + "info", + ); return; } - // Switch between profiles - const personalLabel = `🏠 personal${currentAccount === "personal" ? " ← current" : ""}${!hasProfile("personal") ? " (no profile saved)" : ""}`; - const workLabel = `💼 work${currentAccount === "work" ? " ← current" : ""}${!hasProfile("work") ? " (no profile saved)" : ""}`; + // ── 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 workLabel = `💼 work${currentAccount === "work" ? " ← current" : ""}${!hasProfile("work") ? " (no profile saved)" : ""}`; - const choice = await ctx.ui.select("Switch Claude account:", [personalLabel, workLabel]); - if (choice === undefined) return; - - const newAccount: Account = choice.startsWith("🏠") ? "personal" : "work"; + const accountChoice = await ctx.ui.select( + "Switch Claude account:", + [personalLabel, workLabel], + ); + if (accountChoice === undefined) return; + newAccount = accountChoice.startsWith("🏠") ? "personal" : "work"; + } else { + ctx.ui.notify( + "Usage: /switch-claude [personal|work|save ]", + "warning", + ); + return; + } if (newAccount === currentAccount) { - ctx.ui.notify(`Already using ${statusLabel(newAccount)}`, "info"); + ctx.ui.notify( + `Already using ${statusLabel(newAccount)}`, + "info", + ); 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)) { ctx.ui.notify( `No profile saved for ${newAccount}. Run /login then /switch-claude save ${newAccount}`, - "warning" + "warning", ); return; } + // ── Perform the switch ────────────────────────────────────────── try { - // Persist any token refreshes pi made since the last save so - // we don't restore a stale refresh token when we come back. - if (currentAccount !== "unknown") saveProfile(currentAccount); - loadProfile(newAccount); + // 1. Snapshot current auth.json → outgoing profile. + // This captures any tokens that were silently refreshed + // since the last save (the file is the source of truth, + // 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; + setMarker(currentAccount); ctx.ui.setStatus("claude-account", statusLabel(currentAccount)); - pi.events.emit("claude-account:switched", { account: newAccount }); - ctx.ui.notify(`Switched to ${statusLabel(newAccount)} — restart pi to apply`, "info"); + ctx.ui.notify( + `Switched to ${statusLabel(newAccount)} ✓`, + "info", + ); } 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); - 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", + ); } }, }); diff --git a/pi/.pi/agent/extensions/usage-bars/index.ts b/pi/.pi/agent/extensions/usage-bars/index.ts index 4f9acb1..e8558eb 100644 --- a/pi/.pi/agent/extensions/usage-bars/index.ts +++ b/pi/.pi/agent/extensions/usage-bars/index.ts @@ -37,7 +37,17 @@ import { type UsageData, } 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 STATUS_KEY = "usage-bars"; @@ -297,6 +307,17 @@ interface UsageState extends UsageByProvider { 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 = { @@ -311,6 +332,8 @@ export default function (pi: ExtensionAPI) { let pollInFlight: Promise | null = null; let pollQueued = false; + /** Timer running during an active agent loop to refresh usage periodically. */ + let streamingTimer: ReturnType | null = null; let ctx: any = null; function renderPercent(theme: any, value: number): string { @@ -410,7 +433,7 @@ export default function (pi: ExtensionAPI) { return false; } - async function runPoll() { + async function runPoll(options: PollOptions = {}) { const auth = readAuth(); const active = state.activeProvider; @@ -420,25 +443,21 @@ export default function (pi: ExtensionAPI) { 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 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) { - // Use whatever data is in the cache so the bar still shows something. if (cache?.data?.[active]) state[active] = cache.data[active]!; state.lastPoll = now; updateStatus(); return; } - // If another session already polled recently, use their result. - if (cache && now - cache.timestamp < CACHE_TTL_MS && cache.data?.[active]) { + // 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(); @@ -446,13 +465,6 @@ export default function (pi: ExtensionAPI) { } // --- 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); let effectiveAuth = auth; if (oauthId && active !== "zai") { @@ -506,16 +518,8 @@ export default function (pi: ExtensionAPI) { 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 === "HTTP 429") { - // Write rate-limit backoff but preserve the last good data in cache. const nextCache: import("./core").UsageCache = { timestamp: cache?.timestamp ?? now, data: { ...(cache?.data ?? {}) }, @@ -541,7 +545,7 @@ export default function (pi: ExtensionAPI) { updateStatus(); } - async function poll() { + async function poll(options: PollOptions = {}) { if (pollInFlight) { pollQueued = true; await pollInFlight; @@ -550,7 +554,7 @@ export default function (pi: ExtensionAPI) { do { pollQueued = false; - pollInFlight = runPoll() + pollInFlight = runPoll(options) .catch(() => { // Never crash extension event handlers on transient polling errors. }) @@ -562,6 +566,22 @@ export default function (pi: ExtensionAPI) { } 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) => { ctx = _ctx; updateProviderFrom(_ctx.model); @@ -569,19 +589,13 @@ export default function (pi: ExtensionAPI) { }); pi.on("session_shutdown", async (_event, _ctx) => { + stopStreamingTimer(); if (_ctx?.hasUI) { _ctx.ui.setStatus(STATUS_KEY, undefined); } }); - // Refresh usage on every turn (like claude-pulse's UserPromptSubmit hook). - // 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(); - }); + // ── Model change ───────────────────────────────────────────────────────── pi.on("model_select", async (event, _ctx) => { ctx = _ctx; @@ -589,6 +603,58 @@ export default function (pi: ExtensionAPI) { 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", { description: "Show API usage for all subscriptions", handler: async (_args, _ctx) => { @@ -609,13 +675,8 @@ export default function (pi: ExtensionAPI) { }); } } 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(); - }); }