Merge #2: Add --color flag to preserve ANSI escape sequences (with debounce fix)

src/cli.ts

@@ -64,8 +64,9 @@ program
   .command("snapshot")
   .description("Return the current rendered screen snapshot (requires: tui-use use <id> first)")
   .option("--format <fmt>", "Output format: pretty, json (default: pretty)", "pretty")
+  .option("--color", "Preserve ANSI color/style escape sequences in output lines")
   .action(async (opts) => {
-    const res = await sendRequest({ type: "snapshot" });
+    const res = await sendRequest({ type: "snapshot", color: opts.color ?? false });
     handleResponse(res, (r) => printScreen(r as ScreenResponse, opts.format as "pretty" | "json"));
   });
 
@@ -82,6 +83,7 @@ program
   .option("--text <pattern>", "Wait until screen contains text (substring or regex)")
   .option("--debounce <ms>", "Idle time after last change before resolving, in ms (default: 100)", "100")
   .option("--format <fmt>", "Output format: pretty, json (default: pretty)", "pretty")
+  .option("--color", "Preserve ANSI color/style escape sequences in output lines")
   .action(async (duration: string | undefined, opts) => {
     // Parse duration (position arg) or use default
     let timeoutMs = 3000;
@@ -97,6 +99,7 @@ program
       timeout_ms: timeoutMs,
       text: opts.text,
       debounce_ms: isNaN(debounceMs) ? 100 : debounceMs,
+      color: opts.color ?? false,
     });
     handleResponse(res, (r) => printScreen(r as ScreenResponse, opts.format as "pretty" | "json"));
   });

src/daemon.ts

@@ -121,7 +121,7 @@ async function handleRequest(req: Request): Promise<Response> {
       if (!session) {
         return { type: "error", message: `Session not found: ${currentSession}` };
       }
-      const { lines, cursor, changed, highlights, title, is_fullscreen } = session.snapshot();
+      const { lines, cursor, changed, highlights, title, is_fullscreen } = session.snapshot({ color: (req as SnapshotRequest).color });
       return {
         type: "snapshot",
         session_id: currentSession,
@@ -146,7 +146,8 @@ async function handleRequest(req: Request): Promise<Response> {
       if (!session) {
         return { type: "error", message: `Session not found: ${currentSession}` };
       }
-      const { lines, cursor, changed, highlights, title, is_fullscreen } = await session.wait((req as WaitRequest).timeout_ms ?? 3000, (req as WaitRequest).text, (req as WaitRequest).debounce_ms ?? 100);
+      const waitReq = req as WaitRequest;
+      const { lines, cursor, changed, highlights, title, is_fullscreen } = await session.wait(waitReq.timeout_ms ?? 3000, waitReq.text, waitReq.debounce_ms ?? 100, { color: waitReq.color });
       return {
         type: "wait",
         session_id: currentSession,

src/protocol.ts

@@ -35,13 +35,15 @@ export type ScreenFormat = "text" | "lines" | "numbered" | "pretty";
 
 export interface SnapshotRequest {
   type: "snapshot";
+  color?: boolean;       // when true, lines contain ANSI escape sequences
 }
 
 export interface WaitRequest {
   type: "wait";
   timeout_ms?: number;   // default 3000
   text?: string;         // substring or regex — block until screen contains text
   debounce_ms?: number;  // idle time after last change before resolving (default 100ms)
+  color?: boolean;       // when true, lines contain ANSI escape sequences
 }
 
 /** Send literal text to the PTY. Supports \n \r \t escape sequences. */

src/session.ts

@@ -11,6 +11,128 @@ import { Terminal } from "@xterm/headless";
 import { SessionInfo } from "./protocol";
 import { extractHighlights, Highlight } from "./highlights";
 
+// ---- ANSI re-encoding helpers ----
+
+interface CellStyle {
+  fg: number; fgMode: "default" | "palette" | "rgb";
+  bg: number; bgMode: "default" | "palette" | "rgb";
+  bold: boolean; dim: boolean; italic: boolean;
+  underline: boolean; blink: boolean; inverse: boolean; strikethrough: boolean;
+}
+
+function getCellStyle(cell: { getFgColor(): number; getBgColor(): number; isFgDefault(): boolean; isFgPalette(): boolean; isFgRGB(): boolean; isBgDefault(): boolean; isBgPalette(): boolean; isBgRGB(): boolean; isBold(): number; isDim(): number; isItalic(): number; isUnderline(): number; isBlink(): number; isInverse(): number; isStrikethrough(): number }): CellStyle {
+  return {
+    fg: cell.getFgColor(),
+    fgMode: cell.isFgDefault() ? "default" : cell.isFgPalette() ? "palette" : "rgb",
+    bg: cell.getBgColor(),
+    bgMode: cell.isBgDefault() ? "default" : cell.isBgPalette() ? "palette" : "rgb",
+    bold: cell.isBold() !== 0,
+    dim: cell.isDim() !== 0,
+    italic: cell.isItalic() !== 0,
+    underline: cell.isUnderline() !== 0,
+    blink: cell.isBlink() !== 0,
+    inverse: cell.isInverse() !== 0,
+    strikethrough: cell.isStrikethrough() !== 0,
+  };
+}
+
+function stylesEqual(a: CellStyle, b: CellStyle): boolean {
+  return a.fg === b.fg && a.fgMode === b.fgMode &&
+    a.bg === b.bg && a.bgMode === b.bgMode &&
+    a.bold === b.bold && a.dim === b.dim && a.italic === b.italic &&
+    a.underline === b.underline && a.blink === b.blink &&
+    a.inverse === b.inverse && a.strikethrough === b.strikethrough;
+}
+
+function isDefaultStyle(s: CellStyle): boolean {
+  return s.fgMode === "default" && s.bgMode === "default" &&
+    !s.bold && !s.dim && !s.italic && !s.underline && !s.blink && !s.inverse && !s.strikethrough;
+}
+
+function styleToAnsi(s: CellStyle): string {
+  const codes: string[] = [];
+  if (s.bold) codes.push("1");
+  if (s.dim) codes.push("2");
+  if (s.italic) codes.push("3");
+  if (s.underline) codes.push("4");
+  if (s.blink) codes.push("5");
+  if (s.inverse) codes.push("7");
+  if (s.strikethrough) codes.push("9");
+  if (s.fgMode === "palette") {
+    if (s.fg < 8) codes.push(String(30 + s.fg));
+    else if (s.fg < 16) codes.push(String(90 + s.fg - 8));
+    else codes.push(`38;5;${s.fg}`);
+  } else if (s.fgMode === "rgb") {
+    codes.push(`38;2;${(s.fg >> 16) & 0xff};${(s.fg >> 8) & 0xff};${s.fg & 0xff}`);
+  }
+  if (s.bgMode === "palette") {
+    if (s.bg < 8) codes.push(String(40 + s.bg));
+    else if (s.bg < 16) codes.push(String(100 + s.bg - 8));
+    else codes.push(`48;5;${s.bg}`);
+  } else if (s.bgMode === "rgb") {
+    codes.push(`48;2;${(s.bg >> 16) & 0xff};${(s.bg >> 8) & 0xff};${s.bg & 0xff}`);
+  }
+  if (codes.length === 0) return "";
+  return `\x1b[${codes.join(";")}m`;
+}
+
+/** Minimal interface for an xterm buffer line (avoids importing concrete types). */
+interface ColorBufferLine {
+  getCell(x: number): { getChars(): string; getWidth(): number; getFgColor(): number; getBgColor(): number; isFgDefault(): boolean; isFgPalette(): boolean; isFgRGB(): boolean; isBgDefault(): boolean; isBgPalette(): boolean; isBgRGB(): boolean; isBold(): number; isDim(): number; isItalic(): number; isUnderline(): number; isBlink(): number; isInverse(): number; isStrikethrough(): number } | undefined;
+  length: number;
+}
+
+/** Render a buffer line with ANSI escape sequences preserved. */
+function renderLineWithColor(line: ColorBufferLine): string {
+  // First pass: find the rightmost cell with non-default style or non-space content.
+  // This lets us trim trailing default-styled whitespace without breaking styled spans.
+  let lastStyledOrContent = -1;
+  for (let x = 0; x < line.length; x++) {
+    const cell = line.getCell(x);
+    if (!cell) break;
+    if (cell.getWidth() === 0) continue;
+    const chars = cell.getChars();
+    if (!isDefaultStyle(getCellStyle(cell)) || (chars !== "" && chars !== " ")) {
+      lastStyledOrContent = x;
+    }
+  }
+
+  if (lastStyledOrContent < 0) return ""; // entirely empty/default line
+
+  let result = "";
+  let currentStyle: CellStyle | null = null;
+
+  for (let x = 0; x <= lastStyledOrContent; x++) {
+    const cell = line.getCell(x);
+    if (!cell) break;
+    if (cell.getWidth() === 0) continue; // wide char continuation cell
+
+    const chars = cell.getChars();
+    const style = getCellStyle(cell);
+
+    if (!currentStyle || !stylesEqual(currentStyle, style)) {
+      // Style changed — always reset before switching
+      if (currentStyle !== null && !isDefaultStyle(currentStyle)) {
+        result += "\x1b[0m";
+      }
+      if (!isDefaultStyle(style)) {
+        result += styleToAnsi(style);
+      }
+      currentStyle = style;
+    }
+
+    // Empty cells (width=1, no chars) are spaces — preserve them for background color
+    result += chars || " ";
+  }
+
+  // Reset at end of line if we had styling
+  if (currentStyle && !isDefaultStyle(currentStyle)) {
+    result += "\x1b[0m";
+  }
+
+  return result;
+}
+
 // ---- Pure helper functions (exported for testing) ----
 
 /** Extract title from a raw title string (or undefined). */
@@ -172,25 +294,45 @@ export class Session {
    * Trailing empty lines and per-line trailing spaces are removed.
    * Updates lastSnapshot for change detection.
    */
-  snapshot(): { lines: string[]; cursor: { x: number; y: number }; changed: boolean; highlights: Highlight[]; title: string; is_fullscreen: boolean } {
+  snapshot(options?: { color?: boolean }): { lines: string[]; cursor: { x: number; y: number }; changed: boolean; highlights: Highlight[]; title: string; is_fullscreen: boolean } {
     const buf = this.terminal.buffer.active;
-    const lines: string[] = [];
-    // Start from viewportY to support scrolling through buffer history
+    const useColor = options?.color ?? false;
+    const plainLines: string[] = [];
     const startY = buf.viewportY;
     for (let i = 0; i < this.terminal.rows; i++) {
-      lines.push((buf.getLine(startY + i)?.translateToString(true) ?? "").trimEnd());
+      plainLines.push((buf.getLine(startY + i)?.translateToString(true) ?? "").trimEnd());
     }
     // Remove trailing empty lines
-    while (lines.length > 0 && lines[lines.length - 1] === "") {
-      lines.pop();
+    while (plainLines.length > 0 && plainLines[plainLines.length - 1] === "") {
+      plainLines.pop();
     }
     // Remove leading empty lines (TUI apps like fzf render from bottom)
-    while (lines.length > 0 && lines[0] === "") {
-      lines.shift();
+    while (plainLines.length > 0 && plainLines[0] === "") {
+      plainLines.shift();
+    }
+    // Change detection always uses plain text
+    const plainScreen = plainLines.join("\n");
+    const changed = plainScreen !== this.lastSnapshot;
+    this.lastSnapshot = plainScreen;
+
+    // Build color lines if requested
+    let lines: string[];
+    if (useColor) {
+      lines = [];
+      for (let i = 0; i < this.terminal.rows; i++) {
+        const bufLine = buf.getLine(startY + i);
+        lines.push(bufLine ? renderLineWithColor(bufLine) : "");
+      }
+      // Trim trailing empty lines (match plain text trimming)
+      while (lines.length > 0 && lines[lines.length - 1] === "") {
+        lines.pop();
+      }
+      while (lines.length > 0 && lines[0] === "") {
+        lines.shift();
+      }
+    } else {
+      lines = plainLines;
     }
-    const screen = lines.join("\n");
-    const changed = screen !== this.lastSnapshot;
-    this.lastSnapshot = screen;
     const highlights = extractHighlights(buf, this.terminal.rows, startY);
     return {
       lines,
@@ -209,7 +351,8 @@ export class Session {
   async wait(
     timeoutMs: number = 3000,
     text?: string,
-    debounceMs: number = 100
+    debounceMs: number = 100,
+    options?: { color?: boolean }
   ): Promise<{ lines: string[]; cursor: { x: number; y: number }; changed: boolean; highlights: ReturnType<typeof extractHighlights>; title: string; is_fullscreen: boolean }> {
     const beforeScreen = this.lastSnapshot;
     const beforeTitle = this._title;
@@ -269,7 +412,7 @@ export class Session {
     });
 
     this.changeListeners = this.changeListeners.filter((l) => typeof l === "function");
-    return this.snapshot();
+    return this.snapshot(options);
   }
 
   kill(): void {