cmux.ts

/**
 * cmux — Push pi agent state into the cmux sidebar.
 *
 * Hooks into pi lifecycle events and fires cmux CLI commands to update
 * sidebar status keys, progress, and notifications. Fire-and-forget —
 * errors are silently ignored so cmux issues never affect pi.
 *
 * No-op when CMUX_SOCKET_PATH is not set (i.e. not running inside cmux).
 *
 * Source: https://github.com/HazAT/pi-config/tree/main/extensions/cmux
 */

import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";

const CMUX_SOCKET = process.env.CMUX_SOCKET_PATH;

const GREEN = "#22C55E";
const AMBER = "#F59E0B";
const PURPLE = "#8B5CF6";
const BLUE = "#3B82F6";
const GRAY = "#6B7280";

/** Status keys owned by this extension — cleared on shutdown. */
const STATUS_KEYS = [
	"pi_state",
	"pi_model",
	"pi_thinking",
	"pi_tokens",
	"pi_cost",
	"pi_tool",
];

/** Format a token count into a compact human-readable string. */
function formatTokens(n: number): string {
	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
	if (n >= 10_000) return `${Math.round(n / 1000)}k`;
	if (n >= 1_000) return `${(n / 1000).toFixed(1)}k`;
	return String(n);
}

/** Format a cost value as a dollar string. */
function formatCost(n: number): string {
	return `$${n.toFixed(2)}`;
}

/** Strip common prefixes/suffixes from model IDs for display. */
function shortModel(id: string): string {
	return id.replace(/^claude-/, "").replace(/-\d{8}$/, "");
}

/** cmux sidebar integration — pushes agent state to cmux when running inside it. */
export default function cmuxExtension(pi: ExtensionAPI): void {
	if (!CMUX_SOCKET) return;

	let sessionCost = 0;
	let hasUI = false;

	function run(...args: string[]): void {
		if (!hasUI) return;
		pi.exec("cmux", args, { timeout: 2000 }).catch(() => {});
	}

	function setStatus(key: string, value: string, icon: string, color: string): void {
		run("set-status", key, value, "--icon", icon, "--color", color);
	}

	function clearStatus(key: string): void {
		run("clear-status", key);
	}

	/** Update the token-count status if context usage is available. */
	function syncTokenStatus(ctx: ExtensionContext): void {
		const usage = ctx.getContextUsage();
		if (usage && usage.tokens != null && usage.tokens > 0) {
			setStatus("pi_tokens", formatTokens(usage.tokens), "number", BLUE);
		}
	}

	// --- Session lifecycle ---

	pi.on("session_start", async (_event, ctx) => {
		hasUI = ctx.hasUI;
		if (!hasUI) return;

		// Reconstruct session cost from existing entries
		sessionCost = 0;
		for (const entry of ctx.sessionManager.getBranch()) {
			if (
				entry.type === "message" &&
				entry.message.role === "assistant" &&
				(entry.message as any).usage?.cost?.total
			) {
				sessionCost += (entry.message as any).usage.cost.total;
			}
		}

		setStatus("pi_state", "Idle", "checkmark.circle", GREEN);

		if (ctx.model?.id) {
			setStatus("pi_model", shortModel(ctx.model.id), "brain", PURPLE);
		}

		const thinking = pi.getThinkingLevel();
		if (thinking && thinking !== "off") {
			setStatus("pi_thinking", thinking, "sparkles", AMBER);
		}

		if (sessionCost > 0) {
			setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
		}

		syncTokenStatus(ctx);
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		if (!ctx.hasUI) return;
		for (const key of STATUS_KEYS) {
			clearStatus(key);
		}
	});

	// --- Agent working state ---

	pi.on("agent_start", async (_event, ctx) => {
		if (!ctx.hasUI) return;
		setStatus("pi_state", "Working", "arrow.circlepath", AMBER);
	});

	pi.on("agent_end", async (_event, ctx) => {
		if (!ctx.hasUI) return;
		setStatus("pi_state", "Idle", "checkmark.circle", GREEN);
		clearStatus("pi_tool");

		syncTokenStatus(ctx);

		if (sessionCost > 0) {
			setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
		}

		// Notify user — empty body triggers blue ring + tab highlight
		run("notify", "--title", "Needs attention");
	});

	// --- Turn tracking (tokens + cost) ---

	pi.on("turn_end", async (event, ctx) => {
		if (!ctx.hasUI) return;

		const msg = event.message;
		if (msg?.role === "assistant" && (msg as any).usage?.cost?.total) {
			sessionCost += (msg as any).usage.cost.total;
			setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
		}

		syncTokenStatus(ctx);
	});

	// --- Model / thinking changes ---

	pi.on("model_select", async (event, ctx) => {
		if (!ctx.hasUI) return;
		setStatus("pi_model", shortModel(event.model.id), "brain", PURPLE);
		const thinking = pi.getThinkingLevel();
		setStatus(
			"pi_thinking",
			thinking === "off" ? "off" : thinking,
			"sparkles",
			thinking === "off" ? GRAY : AMBER,
		);
	});

	// --- Tool execution tracking ---

	pi.on("tool_execution_start", async (event, ctx) => {
		if (!ctx.hasUI) return;
		setStatus("pi_tool", event.toolName, "wrench", GRAY);
	});

	pi.on("tool_execution_end", async (_event, ctx) => {
		if (!ctx.hasUI) return;
		clearStatus("pi_tool");
	});
}