From f508e30fcb8eb22175a609f6e24718aaa10623cb Mon Sep 17 00:00:00 2001
From: Chris Wanstrath <chris@ozmm.org>
Date: Thu, 19 Mar 2026 11:50:15 -0700
Subject: [PATCH] Merge setup and user macros into a single map and harden @def
 parsing

User macros already overwrote setup macros, so two separate maps were
unnecessary. Also prevents @def continuations from silently swallowing
command or directive lines, and stops expanding macros in teardown.
---
 src/cli/index.ts  | 10 ++++------
 src/parse.test.ts | 18 ++++++++++++++++++
 src/parse.ts      |  9 +++++++--
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/cli/index.ts b/src/cli/index.ts
index a682613..636c1e3 100755
--- a/src/cli/index.ts
+++ b/src/cli/index.ts
@@ -117,8 +117,7 @@ program
       const envVars: Record<string, string> = {}
       const setupEnvVars: Record<string, string> = {}
       const userEnvVars: Record<string, string> = {}
-      const setupMacros: Record<string, string> = {}
-      const userMacros: Record<string, string> = {}
+      const macros: Record<string, string> = {}
       const setupCommands: Command[] = []
       const teardownCommands: Command[] = [...parsed.teardownCommands]
       for (const d of parsed.directives) {
@@ -128,21 +127,20 @@ program
           const setupParsed = parseSetup(relative(cwd, setupPath), setupContent)
           for (const sd of setupParsed.directives) {
             if (sd.type === "env") setupEnvVars[sd.key] = sd.value
-            else if (sd.type === "def") setupMacros[sd.name] = sd.body
+            else if (sd.type === "def") macros[sd.name] = sd.body
           }
           setupCommands.push(...setupParsed.commands)
           teardownCommands.push(...setupParsed.teardownCommands)
         } else if (d.type === "env") {
           userEnvVars[d.key] = d.value
         } else if (d.type === "def") {
-          userMacros[d.name] = d.body
+          macros[d.name] = d.body
         }
       }
       Object.assign(envVars, setupEnvVars, userEnvVars)
       if (!("PORT" in userEnvVars) && !("PORT" in setupEnvVars)) {
         envVars["PORT"] = String(port)
       }
-      const macros: Record<string, string> = { ...setupMacros, ...userMacros }
       const expandMacro = (cmd: Command): Command => {
         const body = macros[cmd.command]
         return body !== undefined ? { ...cmd, command: body } : cmd
@@ -152,7 +150,7 @@ program
         commands: [
           ...setupCommands.map(expandMacro),
           ...parsed.commands.map(expandMacro),
-          ...teardownCommands.map(expandMacro),
+          ...teardownCommands,
         ],
       }
 
diff --git a/src/parse.test.ts b/src/parse.test.ts
index 0fa7c98..f7e3c96 100644
--- a/src/parse.test.ts
+++ b/src/parse.test.ts
@@ -218,6 +218,24 @@ describe("parse", () => {
     )
   })
 
+  test("@def with whitespace-only body throws", () => {
+    expect(() => parse("test.shout", "@def greet   \n$ echo hi\n")).toThrow(
+      "test.shout:1: @def requires a name and body",
+    )
+  })
+
+  test("@def continuation consuming $ line throws", () => {
+    expect(() => parse("test.shout", "@def foo echo a \\\n$ echo real\n")).toThrow(
+      "test.shout:2: @def continuation consumed a command or directive line",
+    )
+  })
+
+  test("@def continuation consuming @ directive throws", () => {
+    expect(() => parse("test.shout", "@def foo echo a \\\n@env PORT=3000\n")).toThrow(
+      "test.shout:2: @def continuation consumed a command or directive line",
+    )
+  })
+
   test("@def after first command is expected output", () => {
     const result = parse("test.shout", "$ cat file\n@def foo bar\n")
     expect(result.directives).toEqual([])
diff --git a/src/parse.ts b/src/parse.ts
index 413588e..802e2e6 100644
--- a/src/parse.ts
+++ b/src/parse.ts
@@ -81,11 +81,16 @@ function parseDefDirective(
   let extra = 0
 
   while (body.endsWith("\\") && i + extra + 1 < rawLines.length) {
-    body = body.slice(0, -1).trimEnd() + "\n" + rawLines[i + extra + 1]!.trim()
+    const next = rawLines[i + extra + 1]!
+    if (next.startsWith("$ ") || next.startsWith("@")) {
+      throw new Error(`${path}:${i + extra + 2}: @def continuation consumed a command or directive line`)
+    }
+    body = body.slice(0, -1).trimEnd() + "\n" + next.trim()
     extra++
   }
 
-  if (!body.trim()) {
+  body = body.trim()
+  if (!body) {
     throw new Error(`${path}:${i + 1}: @def requires a name and body`)
   }