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

/**
 * 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.
 *
 * 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";

const HOME = os.homedir();
const AUTH_FILE = 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 PROFILES_DIR = path.join(HOME, ".pi/agent/profiles");

type Account = "personal" | "work";

function profilePath(account: Account): string {
	return path.join(PROFILES_DIR, `auth-${account}.json`);
}

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

function saveProfile(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);
}

function loadProfile(account: Account): void {
	fs.copyFileSync(profilePath(account), AUTH_FILE);
	fs.chmodSync(AUTH_FILE, 0o600);
	fs.writeFileSync(MARKER_FILE, account);
}

function getCurrentAccount(): Account | "unknown" {
	try {
		const marker = fs.readFileSync(MARKER_FILE, "utf8").trim();
		if (marker === "personal" || marker === "work") return marker;
	} catch {}
	return "unknown";
}

function statusLabel(account: Account | "unknown"): string {
	switch (account) {
		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;
	}
}

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.",
		handler: async (args, ctx) => {
			const trimmed = args?.trim() ?? "";

			// Save current auth.json 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;
				}
				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;
			}

			// 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)" : ""}`;

			const choice = await ctx.ui.select("Switch Claude account:", [personalLabel, workLabel]);
			if (choice === undefined) return;

			const newAccount: Account = choice.startsWith("🏠") ? "personal" : "work";

			if (newAccount === currentAccount) {
				ctx.ui.notify(`Already using ${statusLabel(newAccount)}`, "info");
				return;
			}

			if (!hasProfile(newAccount)) {
				ctx.ui.notify(
					`No profile saved for ${newAccount}. Run /login then /switch-claude save ${newAccount}`,
					"warning"
				);
				return;
			}

			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);
				currentAccount = newAccount;
				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");
			} catch (e: unknown) {
				const msg = e instanceof Error ? e.message : String(e);
				ctx.ui.notify(`Failed to switch: ${msg}`, "error");
			}
		},
	});
}