diff --git a/src/index.ts b/src/index.ts index c34f32c..c9f7c36 100644 --- a/src/index.ts +++ b/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() 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) } }, } diff --git a/test/plugin.test.ts b/test/plugin.test.ts index c0012af..f980f0d 100644 --- a/test/plugin.test.ts +++ b/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") + }) })