cmux-remote/src/bridge/cmux-client.ts

import { execFile } from 'node:child_process';
import { readFile, unlink } from 'node:fs/promises';
import { createConnection } from 'node:net';
import { promisify } from 'node:util';
import { homedir, tmpdir } from 'node:os';
import { join } from 'node:path';
import type { WorkspaceInfo, PaneInfo, SurfaceInfo } from '../protocol/messages.js';

const execFileAsync = promisify(execFile);

const MAX_CONCURRENT = 5;
const CACHE_TTL_MS = 2000;

interface CacheEntry<T> {
  data: T;
  timestamp: number;
}

export class CmuxClient {
  private cmuxPath: string;
  private running = 0;
  private queue: Array<() => void> = [];
  private cache = new Map<string, CacheEntry<unknown>>();
  private lastGood = new Map<string, string[]>();
  private shuttingDown = false;

  constructor(cmuxPath = 'cmux') {
    this.cmuxPath = cmuxPath;
  }

  shutdown(): void {
    this.shuttingDown = true;
  }

  private async throttle(): Promise<void> {
    if (this.running < MAX_CONCURRENT) {
      this.running++;
      return;
    }
    return new Promise((resolve) => {
      this.queue.push(() => {
        this.running++;
        resolve();
      });
    });
  }

  private release(): void {
    this.running--;
    const next = this.queue.shift();
    if (next) next();
  }

  private getCached<T>(key: string): T | null {
    const entry = this.cache.get(key);
    if (entry && Date.now() - entry.timestamp < CACHE_TTL_MS) {
      return entry.data as T;
    }
    this.cache.delete(key);
    return null;
  }

  private setCache<T>(key: string, data: T): void {
    this.cache.set(key, { data, timestamp: Date.now() });
  }

  async exec(args: string[]): Promise<string> {
    if (this.shuttingDown) return '';
    await this.throttle();
    try {
      const { stdout } = await execFileAsync(this.cmuxPath, args, {
        timeout: 10000,
        maxBuffer: 1024 * 1024,
      });
      return stdout;
    } catch (err) {
      if (this.shuttingDown) return '';
      throw err;
    } finally {
      this.release();
    }
  }

  async listWorkspaces(): Promise<WorkspaceInfo[]> {
    const cached = this.getCached<WorkspaceInfo[]>('workspaces');
    if (cached) return cached;

    const output = await this.exec(['tree', '--all']);
    const workspaces = this.parseTree(output);
    this.setCache('workspaces', workspaces);
    return workspaces;
  }

  parseTree(output: string): WorkspaceInfo[] {
    const workspaces: WorkspaceInfo[] = [];
    let currentWs: WorkspaceInfo | null = null;
    let currentPane: PaneInfo | null = null;

    for (const line of output.split('\n')) {
      // Strip tree-drawing characters and leading whitespace
      const stripped = line.replace(/^[\s│├└─┬┤┼┌┐┘┴]+/g, '').trim();
      if (!stripped) continue;

      // Match workspace lines: "workspace workspace:1 "name" [flags]"
      const wsMatch = stripped.match(/^workspace\s+(workspace:\d+)\s*(.*)/);
      if (wsMatch) {
        // Extract quoted name if present
        const nameMatch = wsMatch[2].match(/"([^"]+)"/);
        currentWs = {
          id: wsMatch[1],
          ref: wsMatch[1],
          name: nameMatch?.[1] || undefined,
          panes: [],
        };
        workspaces.push(currentWs);
        currentPane = null;
        continue;
      }

      // Match pane lines: "pane pane:2 [flags]"
      const paneMatch = stripped.match(/^pane\s+(pane:\d+)/);
      if (paneMatch && currentWs) {
        currentPane = {
          id: paneMatch[1],
          ref: paneMatch[1],
          surfaces: [],
        };
        currentWs.panes.push(currentPane);
        continue;
      }

      // Match surface lines: "surface surface:2 [terminal] "title" [flags]"
      const surfaceMatch = stripped.match(/^surface\s+(surface:\d+)\s*(?:\[(\w+)\])?\s*(.*)/);
      if (surfaceMatch && currentPane) {
        const titleMatch = surfaceMatch[3].match(/"([^"]+)"/);
        currentPane.surfaces.push({
          id: surfaceMatch[1],
          ref: surfaceMatch[1],
          type: surfaceMatch[2] || 'terminal',
          title: titleMatch?.[1] || undefined,
        });
      }
    }

    return workspaces;
  }

  async listPaneSurfaces(workspace: string, pane?: string): Promise<SurfaceInfo[]> {
    const cacheKey = `surfaces:${workspace}:${pane || 'all'}`;
    const cached = this.getCached<SurfaceInfo[]>(cacheKey);
    if (cached) return cached;

    const args = ['list-pane-surfaces', '--workspace', workspace];
    if (pane) args.push('--pane', pane);

    const output = await this.exec(args);
    const surfaces = this.parseSurfaces(output);
    this.setCache(cacheKey, surfaces);
    return surfaces;
  }

  parseSurfaces(output: string): SurfaceInfo[] {
    const surfaces: SurfaceInfo[] = [];
    for (const line of output.split('\n')) {
      const trimmed = line.trim();
      if (!trimmed) continue;
      const match = trimmed.match(/^(surface:\d+)\s*(?:\[(\w+)\])?\s*(.*)/);
      if (match) {
        surfaces.push({
          id: match[1],
          ref: match[1],
          type: match[2] || 'terminal',
          title: match[3]?.trim() || undefined,
        });
      }
    }
    return surfaces;
  }

  async readScreen(
    workspace: string,
    surface: string,
    opts?: { scrollback?: boolean; lines?: number }
  ): Promise<string> {
    const args = ['read-screen', '--workspace', workspace, '--surface', surface];
    if (opts?.scrollback) args.push('--scrollback');
    if (opts?.lines) args.push('--lines', String(opts.lines));
    return this.exec(args);
  }

  private get socketPath(): string {
    return process.env.CMUX_SOCKET_PATH ||
      join(homedir(), 'Library', 'Application Support', 'cmux', 'cmux.sock');
  }

  private socketQuery<T>(method: string, params: Record<string, unknown>): Promise<T> {
    return new Promise((resolve, reject) => {
      const id = Math.random().toString(36).slice(2);
      const conn = createConnection(this.socketPath, () => {
        conn.write(JSON.stringify({ id, method, params }) + '\n');
      });
      const timeout = setTimeout(() => {
        conn.destroy();
        reject(new Error(`Socket timed out after 10s (method: ${method})`));
      }, 10000);
      conn.on('data', (data) => {
        clearTimeout(timeout);
        conn.destroy();
        try {
          const msg = JSON.parse(data.toString());
          if (msg.ok) resolve(msg.result as T);
          else reject(new Error(msg.error?.message || 'socket error'));
        } catch {
          reject(new Error('Invalid response from cmux socket'));
        }
      });
      conn.on('error', (err) => {
        clearTimeout(timeout);
        reject(err);
      });
    });
  }

  private sendViaSocket(method: string, params: Record<string, unknown>): Promise<void> {
    return this.socketQuery<void>(method, params).then(() => {});
  }

  private async _getSurfaceRefsForProcess(processName: string, cacheKey: string): Promise<string[]> {
    const cached = this.getCached<string[]>(cacheKey);
    if (cached) return cached;

    // Step 1: find all matching process PIDs and extract their CMUX_SURFACE_ID
    const uuids = new Set<string>();
    try {
      const { stdout: psListOut } = await execFileAsync('ps', ['-axo', 'pid=,args='], { timeout: 5000, maxBuffer: 10 * 1024 * 1024 });
      const pids: string[] = [];
      for (const line of psListOut.split('\n')) {
        const parts = line.trim().split(/\s+/);
        if (parts.length < 2) continue;
        // Check all args tokens: handles both native binaries (/usr/bin/claude)
        // and Node.js wrappers (node /usr/local/bin/claude ...).
        const matched = parts.slice(1).some(
          p => p === processName || p.endsWith(`/${processName}`)
        );
        if (matched) pids.push(parts[0]);
      }
      for (const pid of pids) {
        try {
          const { stdout: psOut } = await execFileAsync('ps', ['-E', '-p', pid], { timeout: 5000 });
          const match = psOut.match(/CMUX_SURFACE_ID=([A-F0-9-]+)/i);
          if (match) uuids.add(match[1]);
        } catch { /* skip */ }
      }
    } catch { /* no matching processes */ }

    if (uuids.size === 0) {
      // Don't cache empty — transient ps failure; return last known good result.
      return this.lastGood.get(cacheKey) ?? [];
    }

    // Step 2: get ALL workspaces via cmux tree --all (covers every window, not just
    // the active one that workspace.list socket returns).
    const allWorkspaces = await this.listWorkspaces();

    // Step 3: for each workspace, get surfaces and match UUID → ref
    const refs: string[] = [];
    for (const ws of allWorkspaces) {
      try {
        const surfResult = await this.socketQuery<{ surfaces: Array<{ id: string; ref: string }> }>('surface.list', { workspace_id: ws.ref });
        for (const s of surfResult.surfaces) {
          if (uuids.has(s.id)) refs.push(s.ref);
        }
      } catch { /* skip workspace */ }
    }

    if (refs.length > 0) this.lastGood.set(cacheKey, refs);
    // Only cache non-empty results; empty may mean transient socket failure.
    if (refs.length > 0) this.setCache(cacheKey, refs);
    return refs.length > 0 ? refs : (this.lastGood.get(cacheKey) ?? []);
  }

  async getClaudeCodeSurfaceRefs(): Promise<string[]> {
    return this._getSurfaceRefsForProcess('claude', 'claude-surfaces');
  }

  async getCodexSurfaceRefs(): Promise<string[]> {
    return this._getSurfaceRefsForProcess('codex', 'codex-surfaces');
  }

  async sendText(workspace: string, surface: string, text: string): Promise<void> {
    await this.sendViaSocket('surface.send_text', { workspace_id: workspace, surface_id: surface, text });
  }

  async sendKey(workspace: string, surface: string, key: string): Promise<void> {
    await this.sendViaSocket('surface.send_key', { workspace_id: workspace, surface_id: surface, key });
  }


  async browserScreenshot(workspace: string, surface: string): Promise<{ data: string; mime: string }> {
    const tmpFile = join(tmpdir(), `cmux-screenshot-${surface}-${Date.now()}.jpg`);
    try {
      await this.exec(['browser', '--surface', surface, 'screenshot', '--out', tmpFile]);
      const data = await readFile(tmpFile);
      return { data: data.toString('base64'), mime: 'image/jpeg' };
    } finally {
      await unlink(tmpFile).catch(() => {});
    }
  }

  async capturePaneScrollback(
    workspace: string,
    surface: string,
    lines?: number
  ): Promise<string> {
    const args = ['capture-pane', '--workspace', workspace, '--surface', surface, '--scrollback'];
    if (lines) args.push('--lines', String(lines));
    return this.exec(args);
  }

  clearCache(): void {
    this.cache.clear();
  }
}