fix: use Map keyed by callID to prevent race condition in command restoration

src/index.ts

@@ -39,7 +39,7 @@ export const SandboxPlugin: Plugin = async ({ directory, worktree }) => {
 
   if (!sandboxReady) return {}
 
-  let lastOriginalCommand: string | undefined
+  const originalCommands = new Map<string, string>()
 
   return {
     "tool.execute.before": async (input, output) => {
@@ -48,7 +48,7 @@ export const SandboxPlugin: Plugin = async ({ directory, worktree }) => {
       const command = output.args?.command
       if (typeof command !== "string" || !command) return
 
-      lastOriginalCommand = command
+      originalCommands.set(input.callID, command)
 
       try {
         output.args.command = await SandboxManager.wrapWithSandbox(command)
@@ -64,9 +64,10 @@ export const SandboxPlugin: Plugin = async ({ directory, worktree }) => {
       if (input.tool !== "bash") return
 
       // Restore original command so the UI shows it instead of the bwrap wrapper
-      if (lastOriginalCommand && input.args && typeof input.args.command === "string") {
-        input.args.command = lastOriginalCommand
-        lastOriginalCommand = undefined
+      const originalCommand = originalCommands.get(input.callID)
+      if (originalCommand && input.args && typeof input.args.command === "string") {
+        input.args.command = originalCommand
+        originalCommands.delete(input.callID)
       }
     },
   }

test/plugin.test.ts

@@ -138,4 +138,40 @@ describe("SandboxPlugin", () => {
     // Command should remain unchanged (fail open)
     expect(output.args.command).toBe("echo hello")
   })
+
+  test("restores correct command for concurrent bash calls", async () => {
+    if (process.platform === "win32") return
+
+    const hooks = await SandboxPlugin(makeCtx())
+
+    // Simulate two concurrent bash commands with different callIDs
+    const input1 = { tool: "bash", sessionID: "s1", callID: "c1" }
+    const output1 = { args: { command: "echo first" } }
+    const input2 = { tool: "bash", sessionID: "s1", callID: "c2" }
+    const output2 = { args: { command: "echo second" } }
+
+    // Both "before" hooks fire before either "after" (simulating concurrent execution)
+    await hooks["tool.execute.before"]?.(input1, output1)
+    await hooks["tool.execute.before"]?.(input2, output2)
+
+    // Now restore both - each should get its own original command
+    const afterInput1 = {
+      tool: "bash",
+      sessionID: "s1",
+      callID: "c1",
+      args: { command: output1.args.command },
+    }
+    const afterInput2 = {
+      tool: "bash",
+      sessionID: "s1",
+      callID: "c2",
+      args: { command: output2.args.command },
+    }
+
+    await hooks["tool.execute.after"]?.(afterInput1, {})
+    await hooks["tool.execute.after"]?.(afterInput2, {})
+
+    expect(afterInput1.args.command).toBe("echo first")
+    expect(afterInput2.args.command).toBe("echo second")
+  })
 })