strip ANSI color sequences from tool output for display

1 files changed, 90 insertions, 2 deletions

diff --git a/src/tui_components.zig b/src/tui_components.zig
index 48f9377..a9648ee 100644
--- a/src/tui_components.zig
+++ b/src/tui_components.zig
@@ -196,13 +196,69 @@ fn wrapParagraph(text: []const u8, width: usize, out: *std.ArrayList([]const u8)
     try out.append(alloc, text[line_start..]);
 }
 
+/// Append `text` to `out` with terminal escape sequences removed. Tool output
+/// frequently carries SGR colour codes and other ANSI control sequences; the
+/// render component styles its own background, so raw escapes corrupt the box.
+/// Recognised forms:
+///   - CSI:  ESC `[` ... final byte in 0x40..0x7E   (e.g. SGR colours)
+///   - OSC:  ESC `]` ... terminated by BEL (0x07) or ST (ESC `\`)
+///   - Other escapes: ESC, optional intermediate bytes (0x20..0x2F), final byte
+///     (e.g. charset select `ESC ( B`)
+/// A bare trailing ESC (truncated sequence) is dropped. All non-escape bytes,
+/// including newlines and UTF-8, pass through unchanged.
+fn stripAnsi(text: []const u8, out: *std.ArrayList(u8), alloc: std.mem.Allocator) !void {
+    var i: usize = 0;
+    while (i < text.len) {
+        const c = text[i];
+        if (c != 0x1B) {
+            try out.append(alloc, c);
+            i += 1;
+            continue;
+        }
+        // ESC seen. Look at the next byte to classify the sequence.
+        if (i + 1 >= text.len) break; // bare trailing ESC: drop it.
+        const kind = text[i + 1];
+        if (kind == '[') {
+            // CSI: consume until a final byte in 0x40..0x7E (inclusive).
+            var j = i + 2;
+            while (j < text.len and !(text[j] >= 0x40 and text[j] <= 0x7E)) j += 1;
+            i = if (j < text.len) j + 1 else text.len;
+        } else if (kind == ']') {
+            // OSC: consume until BEL or ST (ESC `\`).
+            var j = i + 2;
+            while (j < text.len) {
+                if (text[j] == 0x07) {
+                    j += 1;
+                    break;
+                }
+                if (text[j] == 0x1B and j + 1 < text.len and text[j + 1] == '\\') {
+                    j += 2;
+                    break;
+                }
+                j += 1;
+            }
+            i = j;
+        } else {
+            // Other escapes (e.g. charset select `ESC ( B`): consume any
+            // intermediate bytes (0x20..0x2F) then one final byte.
+            var j = i + 1;
+            while (j < text.len and text[j] >= 0x20 and text[j] <= 0x2F) j += 1;
+            i = if (j < text.len) j + 1 else text.len;
+        }
+    }
+}
+
 /// Split `buffer` on newlines into paragraphs and wrap each to `width`,
 /// appending all produced lines to `out`. A trailing newline produces a final
 /// empty line (so a freshly-typed "\n" shows a blank row). An empty buffer
 /// produces no lines.
 fn wrapBuffer(buffer: []const u8, width: usize, out: *std.ArrayList([]const u8), alloc: std.mem.Allocator) !void {
     if (buffer.len == 0) return;
-    var it = std.mem.splitScalar(u8, buffer, '\n');
+    // A single trailing newline is a line *terminator*, not an empty final
+    // line — strip it so `"a\nb\n"` wraps to two lines, not three.
+    const trimmed = if (buffer[buffer.len - 1] == '\n') buffer[0 .. buffer.len - 1] else buffer;
+    if (trimmed.len == 0) return;
+    var it = std.mem.splitScalar(u8, trimmed, '\n');
     while (it.next()) |para| {
         try wrapParagraph(para, width, out, alloc);
     }
@@ -1576,7 +1632,9 @@ pub const ToolUse = struct {
     pub fn setOutput(self: *ToolUse, value: []const u8) !void {
         if (self.output == null) self.output = .empty;
         self.output.?.clearRetainingCapacity();
-        try self.output.?.appendSlice(self.alloc, value);
+        // Strip ANSI escapes once on ingest so colour codes from commands
+        // don't corrupt the component's own background styling.
+        try stripAnsi(value, &self.output.?, self.alloc);
         self.cache.markDirty();
     }
 
@@ -2518,6 +2576,36 @@ test "ToolUse: short output is shown whole even when collapsed" {
     try testing.expect(seen_only and seen_two);
 }
 
+test "ToolUse: ANSI escapes are stripped from output on ingest" {
+    var t = ToolUse.init(testing.allocator);
+    defer t.deinit();
+    try t.setName("grep");
+    try t.setInput("{}");
+    // SGR colour around "match", an OSC title, and a charset-select escape.
+    try t.setOutput("\x1b[1;31mmatch\x1b[0m here \x1b]0;title\x07\x1b(Bdone");
+    // The stored output must be free of escapes but keep the text.
+    try testing.expectEqualStrings("match here done", t.output.?.items);
+}
+
+test "stripAnsi: handles CSI, OSC, plain text, and truncated escapes" {
+    const a = testing.allocator;
+    const cases = [_]struct { in: []const u8, want: []const u8 }{
+        .{ .in = "plain", .want = "plain" },
+        .{ .in = "\x1b[31mred\x1b[0m", .want = "red" },
+        .{ .in = "a\x1b[1mb\nc", .want = "ab\nc" },
+        .{ .in = "x\x1b]0;t\x07y", .want = "xy" },
+        .{ .in = "x\x1b]0;t\x1b\\y", .want = "xy" },
+        .{ .in = "trailing\x1b", .want = "trailing" },
+        .{ .in = "\x1b(Bplain", .want = "plain" },
+    };
+    for (cases) |c| {
+        var out: std.ArrayList(u8) = .empty;
+        defer out.deinit(a);
+        try stripAnsi(c.in, &out, a);
+        try testing.expectEqualStrings(c.want, out.items);
+    }
+}
+
 test "ToolUse: collapse/expand is a length change with a cache-derived firstLineChanged" {
     // Expanding/collapsing changes the rendered LINE COUNT (plan §3.3). A
     // collapse toggle is a structural change (the whole output region shifts),