dotfiles/pi/.pi/agent/extensions/claude-account-switch.ts

/**
 * Claude Account Switch Extension
 *
 * 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
 *   2. /switch-claude save personal
 *   3. /login  →  authenticate with work account
 *   4. /switch-claude save work
 *
 * Usage:
 *   /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_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`);
}

function hasProfile(account: Account): boolean {
	return fs.existsSync(profilePath(account));
}

/**
 * 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 });
	if (fs.existsSync(AUTH_JSON)) {
		fs.copyFileSync(AUTH_JSON, profilePath(account));
		fs.chmodSync(profilePath(account), 0o600);
	}
}

/**
 * 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" {
	try {
		const marker = fs.readFileSync(MARKER_FILE, "utf8").trim();
		if (marker === "personal" || marker === "work") return marker;
	} catch {}
	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";
	}
}

// ── Extension ──────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	let currentAccount: Account | "unknown" = "unknown";

	pi.on("session_start", async (_event, ctx) => {
		currentAccount = getCurrentAccount();
		ctx.ui.setStatus("claude-account", statusLabel(currentAccount));
	});

	pi.registerCommand("switch-claude", {
		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 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",
					);
					return;
				}

				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 workLabel = `💼  work${currentAccount === "work" ? "  ← current" : ""}${!hasProfile("work") ? "  (no profile saved)" : ""}`;

				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",
				);
				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",
				);
				return;
			}

			// ── Perform the switch ──────────────────────────────────────────
			try {
				// 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));
				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(
					`❌ Switch failed: ${msg}. Rolled back to ${statusLabel(currentAccount)}. ` +
						`You may need to /login and /switch-claude save ${newAccount}.`,
					"error",
				);
			}
		},
	});
}