agent-orchestrator/agor-pty/clients/ts/pty-client.ts

/**
 * agor-pty TypeScript IPC client.
 * Connects to the agor-ptyd daemon via Unix socket.
 * Works with both Bun (Electrobun) and Node.js (Tauri sidecar).
 */

import { connect, type Socket } from "net";
import { readFileSync } from "fs";
import { join } from "path";
import { EventEmitter } from "events";

// ── IPC Protocol Types ──────────────────────────────────────────

export interface SessionInfo {
  id: string;
  pid: number;
  shell: string;
  cwd: string;
  cols: number;
  rows: number;
  created_at: number;
  alive: boolean;
}

export type DaemonEvent =
  | { type: "auth_result"; ok: boolean }
  | { type: "session_created"; session_id: string; pid: number }
  /** data is base64-encoded bytes from the PTY. */
  | { type: "session_output"; session_id: string; data: string }
  | { type: "session_closed"; session_id: string; exit_code: number | null }
  | { type: "session_list"; sessions: SessionInfo[] }
  | { type: "pong" }
  | { type: "error"; message: string };

// ── Client ──────────────────────────────────────────────────────

export class PtyClient extends EventEmitter {
  private socket: Socket | null = null;
  private buffer = "";
  private authenticated = false;
  private socketPath: string;
  private tokenPath: string;

  constructor(socketDir?: string) {
    super();
    const dir =
      socketDir ??
      (process.env.XDG_RUNTIME_DIR
        ? join(process.env.XDG_RUNTIME_DIR, "agor")
        : join(
            process.env.HOME ?? "/tmp",
            ".local",
            "share",
            "agor",
            "run"
          ));
    this.socketPath = join(dir, "ptyd.sock");
    this.tokenPath = join(dir, "ptyd.token");
  }

  /** Connect to daemon and authenticate. */
  async connect(): Promise<void> {
    return new Promise((resolve, reject) => {
      let token: string;
      try {
        token = readFileSync(this.tokenPath, "utf-8").trim();
      } catch {
        reject(new Error(`Cannot read token at ${this.tokenPath}. Is agor-ptyd running?`));
        return;
      }

      this.socket = connect(this.socketPath);

      this.socket.on("connect", () => {
        this.send({ type: "auth", token });
      });

      this.socket.on("data", (chunk: Buffer) => {
        this.buffer += chunk.toString("utf-8");
        let newlineIdx: number;
        while ((newlineIdx = this.buffer.indexOf("\n")) !== -1) {
          const line = this.buffer.slice(0, newlineIdx);
          this.buffer = this.buffer.slice(newlineIdx + 1);
          try {
            const msg = JSON.parse(line) as DaemonEvent;
            if (!this.authenticated && msg.type === "auth_result") {
              if (msg.ok) {
                this.authenticated = true;
                resolve();
              } else {
                reject(new Error("Authentication failed"));
              }
            } else {
              this.emit("message", msg);
              this.emit(msg.type, msg);
            }
          } catch {
            // Ignore malformed lines
          }
        }
      });

      this.socket.on("error", (err) => {
        if (!this.authenticated) reject(err);
        this.emit("error", err);
      });

      this.socket.on("close", () => {
        this.authenticated = false;
        this.emit("close");
      });
    });
  }

  /** Create a new PTY session. */
  createSession(opts: {
    id: string;
    shell?: string;
    cwd?: string;
    env?: Record<string, string>;
    cols?: number;
    rows?: number;
  }): void {
    this.send({
      type: "create_session",
      id: opts.id,
      shell: opts.shell ?? null,
      cwd: opts.cwd ?? null,
      env: opts.env ?? null,
      cols: opts.cols ?? 80,
      rows: opts.rows ?? 24,
    });
  }

  /** Write input to a session's PTY. data is encoded as base64 for transport. */
  writeInput(sessionId: string, data: string | Uint8Array): void {
    const bytes =
      typeof data === "string"
        ? new TextEncoder().encode(data)
        : data;
    // Daemon expects base64 string per protocol.rs WriteInput { data: String }
    const b64 = btoa(String.fromCharCode(...bytes));
    this.send({ type: "write_input", session_id: sessionId, data: b64 });
  }

  /** Resize a session's terminal. */
  resize(sessionId: string, cols: number, rows: number): void {
    this.send({ type: "resize", session_id: sessionId, cols, rows });
  }

  /** Subscribe to a session's output. */
  subscribe(sessionId: string): void {
    this.send({ type: "subscribe", session_id: sessionId });
  }

  /** Unsubscribe from a session's output. */
  unsubscribe(sessionId: string): void {
    this.send({ type: "unsubscribe", session_id: sessionId });
  }

  /** Close/kill a session. */
  closeSession(sessionId: string): void {
    this.send({ type: "close_session", session_id: sessionId });
  }

  /** List all active sessions. */
  listSessions(): void {
    this.send({ type: "list_sessions" });
  }

  /** Ping the daemon. */
  ping(): void {
    this.send({ type: "ping" });
  }

  /** Disconnect from daemon (sessions stay alive). */
  disconnect(): void {
    this.socket?.end();
    this.socket = null;
    this.authenticated = false;
  }

  /** Check if connected and authenticated. */
  get isConnected(): boolean {
    return this.authenticated && this.socket !== null && !this.socket.destroyed;
  }

  private send(msg: Record<string, unknown>): void {
    if (!this.socket || this.socket.destroyed) {
      throw new Error("Not connected to daemon");
    }
    this.socket.write(JSON.stringify(msg) + "\n");
  }
}

// ── Convenience: auto-connect + session helpers ─────────────────

/**
 * Connect to daemon, create a session, and return output as an async iterator.
 * Usage:
 *   for await (const chunk of ptySession("my-shell", { cols: 120, rows: 40 })) {
 *     terminal.write(new Uint8Array(chunk));
 *   }
 */
/**
 * Connect to daemon, create a session, and return output as an async iterator.
 * Each yielded chunk is a Uint8Array of raw PTY bytes (decoded from base64).
 * Usage:
 *   for await (const chunk of ptySession("my-shell", { cols: 120, rows: 40 })) {
 *     terminal.write(chunk);
 *   }
 */
export async function* ptySession(
  sessionId: string,
  opts?: { shell?: string; cwd?: string; cols?: number; rows?: number; socketDir?: string }
): AsyncGenerator<Uint8Array, void, void> {
  const client = new PtyClient(opts?.socketDir);
  await client.connect();

  client.createSession({
    id: sessionId,
    shell: opts?.shell,
    cwd: opts?.cwd,
    cols: opts?.cols ?? 80,
    rows: opts?.rows ?? 24,
  });
  client.subscribe(sessionId);

  try {
    while (client.isConnected) {
      const msg: DaemonEvent = await new Promise((resolve) => {
        client.once("message", resolve);
      });

      if (msg.type === "session_output" && msg.session_id === sessionId) {
        // Decode base64 → raw bytes
        const binary = atob(msg.data);
        const bytes = new Uint8Array(binary.length);
        for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i);
        yield bytes;
      } else if (msg.type === "session_closed" && msg.session_id === sessionId) {
        break;
      }
    }
  } finally {
    client.disconnect();
  }
}