1  /**
  2   * cmux — Push pi agent state into the cmux sidebar.
  3   *
  4   * Hooks into pi lifecycle events and fires cmux CLI commands to update
  5   * sidebar status keys, progress, and notifications. Fire-and-forget —
  6   * errors are silently ignored so cmux issues never affect pi.
  7   *
  8   * No-op when CMUX_SOCKET_PATH is not set (i.e. not running inside cmux).
  9   */
 10  
 11  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 12  
 13  const CMUX_SOCKET = process.env.CMUX_SOCKET_PATH;
 14  // Track the workspace that spawned this Pi process; `cmux current-workspace`
 15  // follows the focused workspace and can point at a different one.
 16  const CMUX_WORKSPACE_ID = process.env.CMUX_WORKSPACE_ID;
 17  
 18  // Colors
 19  const GREEN = "#22C55E";
 20  const AMBER = "#F59E0B";
 21  const PURPLE = "#8B5CF6";
 22  const BLUE = "#3B82F6";
 23  const GRAY = "#6B7280";
 24  
 25  // Status keys we own — cleared on shutdown
 26  const STATUS_KEYS = [
 27    "pi_state",
 28    "pi_model",
 29    "pi_thinking",
 30    "pi_tokens",
 31    "pi_cost",
 32    "pi_tool",
 33  ];
 34  
 35  function formatTokens(n: number): string {
 36    if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
 37    if (n >= 10_000) return `${Math.round(n / 1000)}k`;
 38    if (n >= 1_000) return `${(n / 1000).toFixed(1)}k`;
 39    return String(n);
 40  }
 41  
 42  function formatCost(n: number): string {
 43    return `$${n.toFixed(2)}`;
 44  }
 45  
 46  function shortModel(id: string): string {
 47    // Strip common prefixes: "claude-" etc., keep it readable
 48    return id
 49      .replace(/^claude-/, "")
 50      .replace(/-\d{8}$/, "");
 51  }
 52  
 53  function countMatches(text: string, pattern: RegExp): number {
 54    return [...text.matchAll(pattern)].length;
 55  }
 56  
 57  export default function (pi: ExtensionAPI) {
 58    if (!CMUX_SOCKET) return;
 59  
 60    let sessionCost = 0;
 61    let hasUI = false;
 62  
 63    // Fire-and-forget cmux CLI call
 64    function run(...args: string[]) {
 65      if (!hasUI) return;
 66      pi.exec("cmux", args, { timeout: 2000 }).catch(() => {});
 67    }
 68  
 69    function setStatus(key: string, value: string, icon: string, color: string) {
 70      run("set-status", key, value, "--icon", icon, "--color", color);
 71    }
 72  
 73    function clearStatus(key: string) {
 74      run("clear-status", key);
 75    }
 76  
 77    async function syncWorkspaceName() {
 78      if (!hasUI || !CMUX_WORKSPACE_ID) return;
 79  
 80      const sessionName = pi.getSessionName()?.trim();
 81      if (!sessionName) return;
 82  
 83      const treeResult = await pi.exec(
 84        "cmux",
 85        ["tree", "--workspace", CMUX_WORKSPACE_ID],
 86        { timeout: 2000 }
 87      ).catch(() => undefined);
 88      const tree = treeResult?.stdout;
 89      if (!tree || treeResult?.code !== 0) return;
 90  
 91      const paneCount = countMatches(tree, /\bpane pane:/g);
 92      const surfaceCount = countMatches(tree, /\bsurface surface:/g);
 93      if (paneCount !== 1 || surfaceCount !== 1) return;
 94  
 95      const workspaceMatch = tree.match(/workspace\s+workspace:\S+\s+"([^"]*)"/);
 96      const currentTitle = workspaceMatch?.[1]?.trim();
 97      if (currentTitle === sessionName) return;
 98  
 99      await pi.exec(
100        "cmux",
101        ["rename-workspace", "--workspace", CMUX_WORKSPACE_ID, sessionName],
102        { timeout: 2000 }
103      ).catch(() => {});
104    }
105  
106    // --- Session lifecycle ---
107  
108    pi.on("session_start", async (_event, ctx) => {
109      hasUI = ctx.hasUI;
110      if (!hasUI) return;
111  
112      await syncWorkspaceName();
113  
114      // Reconstruct session cost from existing entries
115      sessionCost = 0;
116      for (const entry of ctx.sessionManager.getBranch()) {
117        if (
118          entry.type === "message" &&
119          entry.message.role === "assistant" &&
120          (entry.message as any).usage?.cost?.total
121        ) {
122          sessionCost += (entry.message as any).usage.cost.total;
123        }
124      }
125  
126      // Set initial sidebar state
127      setStatus("pi_state", "Idle", "checkmark.circle", GREEN);
128  
129      if (ctx.model?.id) {
130        setStatus("pi_model", shortModel(ctx.model.id), "brain", PURPLE);
131      }
132  
133      const thinking = pi.getThinkingLevel();
134      if (thinking && thinking !== "off") {
135        setStatus("pi_thinking", thinking, "sparkles", AMBER);
136      }
137  
138      if (sessionCost > 0) {
139        setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
140      }
141  
142      const usage = ctx.getContextUsage();
143      if (usage && usage.tokens != null && usage.tokens > 0) {
144        setStatus("pi_tokens", formatTokens(usage.tokens), "number", BLUE);
145      }
146    });
147  
148    pi.on("session_shutdown", async (_event, ctx) => {
149      if (!ctx.hasUI) return;
150      for (const key of STATUS_KEYS) {
151        clearStatus(key);
152      }
153    });
154  
155    // --- Agent working state ---
156  
157    pi.on("agent_start", async (_event, ctx) => {
158      if (!ctx.hasUI) return;
159      setStatus("pi_state", "Working", "arrow.circlepath", AMBER);
160    });
161  
162    pi.on("agent_end", async (_event, ctx) => {
163      if (!ctx.hasUI) return;
164      setStatus("pi_state", "Idle", "checkmark.circle", GREEN);
165      clearStatus("pi_tool");
166      await syncWorkspaceName();
167  
168      // Update final token count
169      const usage = ctx.getContextUsage();
170      if (usage && usage.tokens != null && usage.tokens > 0) {
171        setStatus("pi_tokens", formatTokens(usage.tokens), "number", BLUE);
172      }
173  
174      // Update cost
175      if (sessionCost > 0) {
176        setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
177      }
178  
179      // Notify user that agent needs attention — use empty body so it
180      // triggers the blue ring and tab highlight without leaving persistent
181      // text in the sidebar.
182      run("notify", "--title", "Needs attention");
183    });
184  
185    // --- Turn tracking (tokens + cost) ---
186  
187    pi.on("turn_end", async (event, ctx) => {
188      if (!ctx.hasUI) return;
189  
190      // Accumulate cost from the assistant message
191      const msg = event.message;
192      if (msg?.role === "assistant" && (msg as any).usage?.cost?.total) {
193        sessionCost += (msg as any).usage.cost.total;
194        setStatus("pi_cost", formatCost(sessionCost), "dollarsign.circle", GREEN);
195      }
196  
197      // Update token count
198      const usage = ctx.getContextUsage();
199      if (usage && usage.tokens != null && usage.tokens > 0) {
200        setStatus("pi_tokens", formatTokens(usage.tokens), "number", BLUE);
201      }
202    });
203  
204    // --- Model / thinking changes ---
205  
206    pi.on("model_select", async (event, ctx) => {
207      if (!ctx.hasUI) return;
208      setStatus("pi_model", shortModel(event.model.id), "brain", PURPLE);
209      const thinking = pi.getThinkingLevel();
210      setStatus(
211        "pi_thinking",
212        thinking === "off" ? "off" : thinking,
213        "sparkles",
214        thinking === "off" ? GRAY : AMBER
215      );
216    });
217  
218    // --- Tool execution tracking ---
219  
220    pi.on("tool_execution_start", async (event, ctx) => {
221      if (!ctx.hasUI) return;
222      setStatus("pi_tool", event.toolName, "wrench", GRAY);
223    });
224  
225    pi.on("tool_execution_end", async (_event, ctx) => {
226      if (!ctx.hasUI) return;
227      clearStatus("pi_tool");
228    });
229  }