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";

// ── Session-window helpers ─────────────────────────────────────────────────
//
// We store the actual `resets_at` timestamp returned by Claude's usage API
// (via the usage:update event) so the switch menu can show a live countdown
// to the next session reset rather than a guessed switchedAt + 5h window.

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

/** Persist the actual session-reset timestamp for an account. */
function saveSessionResetsAt(account: Account, resetsAt: number): void {
	try {
		fs.mkdirSync(PROFILES_DIR, { recursive: true });
		fs.writeFileSync(
			sessionStampPath(account),
			JSON.stringify({ resetsAt }, null, 2),
			{ mode: 0o600 },
		);
	} catch {}
}

/** Load the stored session-reset timestamp (ms epoch), or null. */
function loadSessionResetsAt(account: Account): number | null {
	try {
		const raw = fs.readFileSync(sessionStampPath(account), "utf-8");
		const { resetsAt } = JSON.parse(raw) as { resetsAt: number };
		if (typeof resetsAt === "number") return resetsAt;
	} catch {}
	return null;
}

/** Format milliseconds as a compact duration string. */
function formatDuration(ms: number): string {
	const totalSec = Math.ceil(ms / 1000);
	const h = Math.floor(totalSec / 3600);
	const m = Math.floor((totalSec % 3600) / 60);
	const s = totalSec % 60;
	if (h > 0) return `${h}h ${m}m`;
	if (m > 0) return `${m}m`;
	return `${s}s`;
}

/**
 * One-line session summary appended to each account option in the select menu.
 *
 * Window still running:  [resets in Xh Ym]
 * Window already passed: [0 (ready)]
 * Never recorded:        (empty)
 */
function sessionSummary(account: Account): string {
	const resetsAt = loadSessionResetsAt(account);
	if (resetsAt === null) return "";
	const remaining = resetsAt - Date.now();
	if (remaining <= 0) return "  [0 (ready)]";
	return `  [resets in ${formatDuration(remaining)}]`;
}

// ── 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.
 *
 * We parse + re-serialize the JSON to guard against corrupt auth.json
 * (e.g. trailing commas left by buggy serializers). If the file can't
 * be parsed, we skip the save rather than propagate bad data.
 */
function saveCurrentAuthToProfile(account: Account): boolean {
	fs.mkdirSync(PROFILES_DIR, { recursive: true });
	if (!fs.existsSync(AUTH_JSON)) return false;
	try {
		const raw = fs.readFileSync(AUTH_JSON, "utf-8");
		const parsed = JSON.parse(raw); // validates JSON
		const clean = JSON.stringify(parsed, null, 2);
		fs.writeFileSync(profilePath(account), clean, { mode: 0o600 });
		return true;
	} catch {
		// auth.json is missing or corrupt — don't propagate bad data
		return false;
	}
}

/**
 * Copy a profile file to auth.json. This is an atomic-ish swap that
 * replaces the entire file rather than merging per-provider.
 *
 * Like saveCurrentAuthToProfile, we round-trip through JSON.parse to
 * ensure we never write corrupt data to auth.json.
 */
function restoreProfileToAuth(account: Account): void {
	const raw = fs.readFileSync(profilePath(account), "utf-8");
	const parsed = JSON.parse(raw); // throws on corrupt profile
	const clean = JSON.stringify(parsed, null, 2);
	fs.writeFileSync(AUTH_JSON, clean, { mode: 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";

	// Whenever usage-bars fetches fresh data, save the real resets_at for the
	// current account so the switch menu shows an accurate live countdown.
	pi.events.on("usage:update", (event: unknown) => {
		const e = event as { sessionResetsAt?: number };
		if (currentAccount !== "unknown" && typeof e.sessionResetsAt === "number") {
			saveSessionResetsAt(currentAccount, e.sessionResetsAt);
		}
	});

	pi.on("session_start", async (_event, ctx) => {
		// Proper-lockfile creates auth.json.lock as a *directory* (atomic mkdir).
		// If a regular file exists at that path (e.g. left by an older pi version),
		// rmdir fails with ENOTDIR → lock acquisition throws → loadError is set →
		// credentials are never persisted after /login.  Delete the stale file and
		// reload so this session has working auth persistence.
		const lockPath = AUTH_JSON + ".lock";
		try {
			const stat = fs.statSync(lockPath);
			if (stat.isFile()) {
				fs.unlinkSync(lockPath);
				ctx.modelRegistry.authStorage.reload();
			}
		} catch {
			// lock doesn't exist or we can't stat it — nothing to fix
		}

		// Guard against corrupt auth.json (e.g. trailing commas from buggy
		// serializers).  Re-serialize to clean JSON and reload so the auth
		// system picks up valid credentials.
		try {
			const raw = fs.readFileSync(AUTH_JSON, "utf-8");
			const parsed = JSON.parse(raw);
			const clean = JSON.stringify(parsed, null, 2);
			if (clean !== raw) {
				fs.writeFileSync(AUTH_JSON, clean, { mode: 0o600 });
				ctx.modelRegistry.authStorage.reload();
			}
		} catch {
			// auth.json missing or unparseable — nothing we can fix here
		}

		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)" : ""}${sessionSummary("personal")}`;
				const workLabel = `󰃖  work${currentAccount === "work" ? "  ← current" : ""}${!hasProfile("work") ? "  (no profile saved)" : ""}${sessionSummary("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",
				);
				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));
				pi.events.emit("claude-account:switched", { account: newAccount });
				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",
				);
			}
		},
	});
}