jonas/dotfiles
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pi update

Browse Source
This commit is contained in:
Jonas H
2026-03-12 09:12:31 +01:00
parent 870dc6ac58
commit de49d03182
4 changed files with 317 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

303
pi/.pi/agent/extensions/claude-account-switch.ts
View File

@@ -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 <name> — save current pi login as a named profile
* /switch-claude — pick account interactively
* /switch-claude save <name> — 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<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;
}
}
// ── 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 <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) => {
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 <name>]",
"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",
);
}
},
});