Ignore heredoc openers inside arithmetic and comments

A << inside ((/$(( is the left-shift operator and a << after a top-level # is documentation, so neither the splitter nor the approval display may open a heredoc there and swallow the following chain. Track arithmetic depth and skip #-to-EOL comments in both.

Author: TheGreatAxios

src/permission/command.test.ts

@@ -193,3 +193,69 @@ describe("splitChainedCommand heredoc boundaries", () => {
     ).toEqual(["cat <<OUTER\nfoo <<INNER\nOUTER", "echo done"]);
   });
 });
+
+describe("splitChainedCommand lexical context (arithmetic and comments)", () => {
+  // Inside `((` / `$((` the `<<` token is the left-shift operator, never a
+  // heredoc opener — the chain after it must still split.
+  test("never opens a heredoc inside arithmetic expansion", () => {
+    expect(splitChainedCommand("echo $((a<<1))")).toEqual(["echo $((a<<1))"]);
+    expect(splitChainedCommand("echo $((a << 1)) && echo done")).toEqual([
+      "echo $((a << 1))",
+      "echo done",
+    ]);
+  });
+
+  test("never opens a heredoc inside a (( )) arithmetic command", () => {
+    expect(splitChainedCommand("((x = a << 1)) && echo done")).toEqual([
+      "x = a << 1",
+      "echo done",
+    ]);
+  });
+
+  // A bare `( ... )` subshell is not arithmetic: a heredoc inside it is real.
+  test("still opens a heredoc inside a bare-paren subshell", () => {
+    const command = "(cat <<EOF\nbody\nEOF) && echo done";
+    expect(splitChainedCommand(command)).toEqual([command]);
+  });
+
+  // A top-level `#` starts a comment through end of line: a `<<` down there
+  // documents rather than opens, so the next line still splits.
+  test("never opens a heredoc from a #-to-EOL comment", () => {
+    expect(splitChainedCommand("# example: cat <<EOF\necho hi")).toEqual([
+      "# example: cat <<EOF",
+      "echo hi",
+    ]);
+    expect(splitChainedCommand("echo hi # tail <<EOF\n&& echo done")).toEqual([
+      "echo hi # tail <<EOF",
+      "echo done",
+    ]);
+  });
+
+  // Comment text never touches arithmetic depth: an unbalanced `((` inside
+  // a `#` comment must not poison later lines, so a genuine heredoc after
+  // the comment still opens and the following chain still splits.
+  test("never counts comment parens toward arithmetic depth", () => {
+    const command = "# (( \ncat <<EOF\nbody\nEOF\n&& echo done";
+    expect(splitChainedCommand(command)).toEqual([
+      "# ((",
+      "cat <<EOF\nbody\nEOF",
+      "echo done",
+    ]);
+  });
+
+  // Chain operators after `#` still split, so a dangerous command hiding
+  // behind a comment still surfaces as its own approval subject.
+  test("still splits chain operators after a # comment", () => {
+    expect(splitChainedCommand("# note && rm -rf /")).toEqual([
+      "# note",
+      "rm -rf /",
+    ]);
+  });
+
+  // A `#` line inside a genuine heredoc body stays payload: the marker still
+  // closes and the following chain still splits.
+  test("keeps a # line inside a heredoc body as payload", () => {
+    const command = "cat <<EOF\n# payload\nEOF\necho done";
+    expect(splitChainedCommand(command)).toEqual([command]);
+  });
+});

src/shell/command-segments.ts

@@ -7,14 +7,21 @@
 // Parentheses group: operators inside a subshell or command substitution never
 // split, and a segment that is exactly one `( ... )` group is unwrapped and its
 // inner chain split recursively — so `(cd a && b)` yields `cd a` and `b`, not
-// the fragment `(cd a`.
+// the fragment `(cd a`. `<<` inside `(( ... ))` / `$(( ... ))` arithmetic is the
+// left-shift operator and a top-level `#` starts a comment — neither opens a
+// heredoc (see isArithmeticOpener / isCommentStart).
 export function splitChainedCommand(command: string): string[] {
   const segments: string[] = [];
   let current = "";
   let quote: '"' | "'" | "`" | null = null;
   let heredocMarker: string | null = null;
   let heredocStripTabs = false;
   let parenDepth = 0;
+  let arithDepth = 0;
+  let commentToEOL = false;
+  // Inside a top-level `#`-to-EOL comment: suppresses only the `<<` heredoc
+  // opener below. Chain operators after `#` still split, so
+  // `# note && rm -rf /` surfaces `rm -rf /` as its own segment.
 
   const push = (): void => {
     const trimmed = current.trim();
@@ -72,7 +79,20 @@ export function splitChainedCommand(command: string): string[] {
     }
 
     // Detect heredoc redirect: << or <<-
-    if (ch === "<" && command[i + 1] === "<") {
+    // A top-level `#` starts a comment through end of line: a `<<` down there
+    // (e.g. `# example: cat <<EOF`) documents rather than opens. Only the
+    // opener is suppressed — the comment text flows through the normal scan
+    // below, so chain operators after `#` still split.
+    if (ch === "\n") commentToEOL = false;
+    if (!commentToEOL && arithDepth === 0 && isCommentStart(command, i)) {
+      commentToEOL = true;
+    }
+    if (
+      !commentToEOL &&
+      arithDepth === 0 &&
+      ch === "<" &&
+      command[i + 1] === "<"
+    ) {
       const opener = parseHeredocOpener(command, i);
       if (opener !== null) {
         current += command.slice(i, opener.lineEnd);
@@ -83,12 +103,17 @@ export function splitChainedCommand(command: string): string[] {
       }
     }
 
-    if (ch === "(") {
+    if (ch === "(" && !commentToEOL) {
+      // `((` / `$((` opens arithmetic, where `<<` shifts instead of opening a
+      // heredoc (see the `<<` guard above). Bare `(` subshells still detect
+      // heredocs — e.g. `(cat <<EOF ...)` is genuine.
+      if (isArithmeticOpener(command, i)) arithDepth++;
       parenDepth++;
       current += ch;
       continue;
     }
-    if (ch === ")") {
+    if (ch === ")" && !commentToEOL) {
+      if (isArithmeticCloser(command, i) && arithDepth > 0) arithDepth--;
       if (parenDepth > 0) parenDepth--;
       current += ch;
       continue;
@@ -145,6 +170,38 @@ export function splitChainedCommand(command: string): string[] {
   return segments;
 }
 
+// Whether text[i] opens an arithmetic context (`((` or `$((`): inside it `<<`
+// is the left-shift operator, never a heredoc opener. Keyed on the doubled
+// paren — a bare `( ... )` subshell can still contain a genuine heredoc.
+// Deliberately not a full arithmetic evaluator: callers only track depth.
+export function isArithmeticOpener(text: string, i: number): boolean {
+  return text[i] === "(" && text[i + 1] === "(";
+}
+
+// Whether text[i] closes one arithmetic-context level (`))`).
+export function isArithmeticCloser(text: string, i: number): boolean {
+  return text[i] === ")" && text[i + 1] === ")";
+}
+
+// Whether text[i] starts a `#`-to-EOL comment: at the very start of the input
+// or right after whitespace, a newline, or a command separator (`;`, `&`,
+// `|`, `(`). A `#` glued to a word (`foo#bar`, `$#`, `${a#b}`) is data.
+export function isCommentStart(text: string, i: number): boolean {
+  if (text[i] !== "#") return false;
+  if (i === 0) return true;
+  const prev = text[i - 1] as string;
+  return (
+    prev === " " ||
+    prev === "\t" ||
+    prev === "\r" ||
+    prev === "\n" ||
+    prev === ";" ||
+    prev === "&" ||
+    prev === "|" ||
+    prev === "("
+  );
+}
+
 // Parses a heredoc opener (`<<` or `<<-`) starting at `command[i]` (which must
 // be the first "<"). Returns the terminating marker text, the exclusive end
 // index of the line that opened the heredoc, and whether the opener was `<<-`

src/tui/command-display.test.ts

@@ -23,6 +23,11 @@ test("display segments exactly match authorization segments", () => {
     "cmd1 && \\\ncmd2",
     "(cd packages/shared && bunx tsc --noEmit 2>&1 | tail -3)",
     "echo start && (cd apps/web && bun test) && echo done",
+    "echo $((a << 1)) && echo done",
+    "((x = a << 1)) && echo done",
+    "# example: cat <<EOF\necho hi",
+    "# (( \ncat <<EOF\nbody\nEOF\n&& echo done",
+    "cat <<EOF\n# payload\nEOF\necho done",
   ];
 
   for (const command of commands) {
@@ -76,6 +81,40 @@ test("a here-string never opens a pending heredoc", () => {
   ]);
 });
 
+test("a << inside arithmetic never opens a pending heredoc", () => {
+  expect(groupChainSegmentsForDisplay("echo $((a << 1)) && echo done")).toEqual(
+    ["echo $((a << 1))", "echo done"],
+  );
+  // With no pending heredoc, a later line is ordinary text, never body.
+  expect(verbatimCommandLines("echo $((a<<1))\nEOF\necho done")).toEqual([
+    { text: "echo $((a<<1))", isComment: false },
+    { text: "EOF", isComment: false },
+    { text: "echo done", isComment: false },
+  ]);
+});
+
+test("a << inside a comment documents rather than opens", () => {
+  expect(verbatimCommandLines("# example: cat <<EOF\necho hi")).toEqual([
+    { text: "# example: cat <<EOF", isComment: true },
+    { text: "echo hi", isComment: false },
+  ]);
+  expect(groupChainSegmentsForDisplay("# c <<EOF")).toEqual(["# c <<EOF"]);
+});
+
+test("comment parens never suppress a later heredoc on the display", () => {
+  // Mirrors the splitter pin: `# ((` is comment text, so the heredoc opens
+  // here exactly as it does for authorization and `&& echo done` separates.
+  expect(
+    verbatimCommandLines("# (( \ncat <<EOF\nbody\nEOF\n&& echo done"),
+  ).toEqual([
+    { text: "# (( ", isComment: true },
+    { text: "cat <<EOF", isComment: false },
+    { text: "body", isComment: false },
+    { text: "EOF", isComment: false },
+    { text: "&& echo done", isComment: false },
+  ]);
+});
+
 test("top-level newlines become verbatim lines; quoted newlines stay marked inline", () => {
   expect(verbatimCommandLines('echo "a\nb"\necho two')).toEqual([
     { text: 'echo "a↵b"', isComment: false },

src/tui/command-display.ts

@@ -1,4 +1,7 @@
 import {
+  isArithmeticCloser,
+  isArithmeticOpener,
+  isCommentStart,
   isHeredocTerminator,
   splitChainedCommand,
 } from "../shell/command-segments.js";
@@ -66,6 +69,9 @@ export function verbatimCommandLines(text: string): VerbatimLine[] {
   let heredocStripTabs = false;
   let heredocPending: { marker: string; stripTabs: boolean } | null = null;
   let continued = false;
+  // Arithmetic depth (`((` / `$((`): `<<` inside is the shift operator and
+  // `#`-to-EOL comments never open a heredoc — mirrors the splitter.
+  let arithDepth = 0;
 
   const push = (): void => {
     const isComment =
@@ -135,7 +141,25 @@ export function verbatimCommandLines(text: string): VerbatimLine[] {
       continue;
     }
 
-    if (ch === "<" && normalized[i + 1] === "<" && heredocPending === null) {
+    if (isArithmeticOpener(normalized, i)) arithDepth++;
+    else if (isArithmeticCloser(normalized, i) && arithDepth > 0) arithDepth--;
+
+    // A top-level `#` comment runs to end of line: `<<` inside it documents
+    // rather than opens. The line still renders whole (see push's isComment).
+    if (arithDepth === 0 && isCommentStart(normalized, i)) {
+      let j = i;
+      while (j < normalized.length && normalized[j] !== "\n") j++;
+      current += normalized.slice(i, j);
+      i = j - 1;
+      continue;
+    }
+
+    if (
+      arithDepth === 0 &&
+      ch === "<" &&
+      normalized[i + 1] === "<" &&
+      heredocPending === null
+    ) {
       const opener = parseHeredocMarker(normalized, i);
       if (opener !== null) heredocPending = opener;
     }
@@ -232,6 +256,9 @@ function segmentWords(segment: string): string[] {
   let heredocMarker: string | null = null;
   let heredocStripTabs = false;
   let heredocPending: { marker: string; stripTabs: boolean } | null = null;
+  // Arithmetic depth (`((` / `$((`): `<<` inside shifts, never opens —
+  // mirrors the splitter (keyed on arithmetic, NOT on bare parens).
+  let arithDepth = 0;
 
   const push = (): void => {
     if (current.length > 0) words.push(current);
@@ -275,7 +302,10 @@ function segmentWords(segment: string): string[] {
       continue;
     }
 
-    if (ch === "<" && segment[i + 1] === "<") {
+    if (isArithmeticOpener(segment, i)) arithDepth++;
+    else if (isArithmeticCloser(segment, i) && arithDepth > 0) arithDepth--;
+
+    if (arithDepth === 0 && ch === "<" && segment[i + 1] === "<") {
       const opener = parseHeredocMarker(segment, i);
       if (opener !== null) {
         push();