From 00cb68689e8a4db7e6e3eb7250b508249b8e9a83 Mon Sep 17 00:00:00 2001 From: Charles Packer Date: Fri, 16 Jan 2026 19:40:45 -0800 Subject: [PATCH] refactor: extract shared shell spawn logic into shellRunner.ts (#572) Co-authored-by: Letta --- src/tests/tools/run-shell-command.test.ts | 11 +- src/tools/impl/Bash.ts | 142 ++++------------------ src/tools/impl/RunShellCommandGemini.ts | 20 +-- src/tools/impl/Shell.ts | 111 ++++++----------- src/tools/impl/ShellCommand.ts | 9 +- src/tools/impl/shellRunner.ts | 135 ++++++++++++++++++++ src/tools/manager.ts | 12 +- 7 files changed, 225 insertions(+), 215 deletions(-) create mode 100644 src/tools/impl/shellRunner.ts diff --git a/src/tests/tools/run-shell-command.test.ts b/src/tests/tools/run-shell-command.test.ts index a9ce37b..117bedf 100644 --- a/src/tests/tools/run-shell-command.test.ts +++ b/src/tests/tools/run-shell-command.test.ts @@ -24,11 +24,10 @@ describe("RunShellCommand tool (Gemini)", () => { }); test("throws error when command is missing", async () => { - // Bash tool doesn't validate empty command, so skip this test - // or test that empty command still executes - const result = await run_shell_command({ - command: "", - } as Parameters[0]); - expect(result.message).toBeTruthy(); + await expect( + run_shell_command({ + command: "", + } as Parameters[0]), + ).rejects.toThrow(/non-empty string/); }); }); diff --git a/src/tools/impl/Bash.ts b/src/tools/impl/Bash.ts index 33c28d5..7e5d3d3 100644 --- a/src/tools/impl/Bash.ts +++ b/src/tools/impl/Bash.ts @@ -3,6 +3,7 @@ import { INTERRUPTED_BY_USER } from "../../constants"; import { backgroundProcesses, getNextBashId } from "./process_manager.js"; import { getShellEnv } from "./shellEnv.js"; import { buildShellLaunchers } from "./shellLaunchers.js"; +import { spawnWithLauncher } from "./shellRunner.js"; import { LIMITS, truncateByChars } from "./truncation.js"; import { validateRequiredParams } from "./validation.js"; @@ -26,123 +27,6 @@ function getBackgroundLauncher(command: string): string[] { return launchers[0] || []; } -/** - * Spawn a command with a specific launcher. - * Returns a promise that resolves with the output or rejects with an error. - */ -function spawnWithLauncher( - launcher: string[], - options: { - cwd: string; - env: NodeJS.ProcessEnv; - timeout: number; - signal?: AbortSignal; - onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; - }, -): Promise<{ stdout: string; stderr: string; exitCode: number | null }> { - return new Promise((resolve, reject) => { - const [executable, ...args] = launcher; - if (!executable) { - reject(new Error("Empty launcher")); - return; - } - - const childProcess = spawn(executable, args, { - cwd: options.cwd, - env: options.env, - shell: false, // Don't use another shell layer - stdio: ["ignore", "pipe", "pipe"], - }); - - const stdoutChunks: Buffer[] = []; - const stderrChunks: Buffer[] = []; - let timedOut = false; - let killTimer: ReturnType | null = null; - - const timeoutId = setTimeout(() => { - timedOut = true; - childProcess.kill("SIGTERM"); - }, options.timeout); - - const abortHandler = () => { - childProcess.kill("SIGTERM"); - if (!killTimer) { - killTimer = setTimeout(() => { - if (childProcess.exitCode === null && !childProcess.killed) { - childProcess.kill("SIGKILL"); - } - }, 2000); - } - }; - if (options.signal) { - options.signal.addEventListener("abort", abortHandler, { once: true }); - } - - childProcess.stdout?.on("data", (chunk: Buffer) => { - stdoutChunks.push(chunk); - options.onOutput?.(chunk.toString("utf8"), "stdout"); - }); - - childProcess.stderr?.on("data", (chunk: Buffer) => { - stderrChunks.push(chunk); - options.onOutput?.(chunk.toString("utf8"), "stderr"); - }); - - childProcess.on("error", (err) => { - clearTimeout(timeoutId); - if (killTimer) { - clearTimeout(killTimer); - killTimer = null; - } - if (options.signal) { - options.signal.removeEventListener("abort", abortHandler); - } - reject(err); - }); - - childProcess.on("close", (code) => { - clearTimeout(timeoutId); - if (killTimer) { - clearTimeout(killTimer); - killTimer = null; - } - if (options.signal) { - options.signal.removeEventListener("abort", abortHandler); - } - - const stdout = Buffer.concat(stdoutChunks).toString("utf8"); - const stderr = Buffer.concat(stderrChunks).toString("utf8"); - - if (timedOut) { - reject( - Object.assign(new Error("Command timed out"), { - killed: true, - signal: "SIGTERM", - stdout, - stderr, - code, - }), - ); - return; - } - - if (options.signal?.aborted) { - reject( - Object.assign(new Error("The operation was aborted"), { - name: "AbortError", - code: "ABORT_ERR", - stdout, - stderr, - }), - ); - return; - } - - resolve({ stdout, stderr, exitCode: code }); - }); - }); -} - /** * Execute a command using spawn with explicit shell. * This avoids the double-shell parsing that exec() does. @@ -164,7 +48,13 @@ export async function spawnCommand( if (process.platform !== "win32") { // On macOS, prefer zsh due to bash 3.2's HEREDOC bug with apostrophes const executable = process.platform === "darwin" ? "/bin/zsh" : "bash"; - return spawnWithLauncher([executable, "-c", command], options); + return spawnWithLauncher([executable, "-c", command], { + cwd: options.cwd, + env: options.env, + timeoutMs: options.timeout, + signal: options.signal, + onOutput: options.onOutput, + }); } // On Windows, use fallback logic to handle PowerShell ENOENT errors (PR #482) @@ -173,7 +63,13 @@ export async function spawnCommand( if (executable) { const newLauncher = [executable, ...launcherArgs.slice(0, -1), command]; try { - const result = await spawnWithLauncher(newLauncher, options); + const result = await spawnWithLauncher(newLauncher, { + cwd: options.cwd, + env: options.env, + timeoutMs: options.timeout, + signal: options.signal, + onOutput: options.onOutput, + }); return result; } catch (error) { const err = error as NodeJS.ErrnoException; @@ -195,7 +91,13 @@ export async function spawnCommand( for (const launcher of launchers) { try { - const result = await spawnWithLauncher(launcher, options); + const result = await spawnWithLauncher(launcher, { + cwd: options.cwd, + env: options.env, + timeoutMs: options.timeout, + signal: options.signal, + onOutput: options.onOutput, + }); cachedWorkingLauncher = launcher; return result; } catch (error) { diff --git a/src/tools/impl/RunShellCommandGemini.ts b/src/tools/impl/RunShellCommandGemini.ts index 5fd9af8..9984d76 100644 --- a/src/tools/impl/RunShellCommandGemini.ts +++ b/src/tools/impl/RunShellCommandGemini.ts @@ -3,26 +3,28 @@ * Uses Gemini's exact schema and description */ -import { bash } from "./Bash"; +import { shell_command } from "./ShellCommand"; interface RunShellCommandGeminiArgs { command: string; description?: string; dir_path?: string; + timeout_ms?: number; + signal?: AbortSignal; + onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; } export async function run_shell_command( args: RunShellCommandGeminiArgs, ): Promise<{ message: string }> { - // Adapt Gemini params to Letta Code's Bash tool - const lettaArgs = { + const result = await shell_command({ command: args.command, - description: args.description, - }; + workdir: args.dir_path, + timeout_ms: args.timeout_ms, + signal: args.signal, + onOutput: args.onOutput, + }); - const result = await bash(lettaArgs); - - // Bash returns { content: Array<{ type: string, text: string }>, status: string } - const message = result.content.map((item) => item.text).join("\n"); + const message = result.output.trim() || "(Command completed with no output)"; return { message }; } diff --git a/src/tools/impl/Shell.ts b/src/tools/impl/Shell.ts index 8069e39..dee0386 100644 --- a/src/tools/impl/Shell.ts +++ b/src/tools/impl/Shell.ts @@ -1,20 +1,17 @@ -import { spawn } from "node:child_process"; import * as path from "node:path"; import { getShellEnv } from "./shellEnv.js"; import { buildShellLaunchers } from "./shellLaunchers.js"; +import { ShellExecutionError, spawnWithLauncher } from "./shellRunner.js"; import { validateRequiredParams } from "./validation.js"; -export class ShellExecutionError extends Error { - code?: string; - executable?: string; -} - interface ShellArgs { command: string[]; workdir?: string; timeout_ms?: number; with_escalated_permissions?: boolean; justification?: string; + signal?: AbortSignal; + onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; } interface ShellResult { @@ -30,81 +27,39 @@ type SpawnContext = { cwd: string; env: NodeJS.ProcessEnv; timeout: number; + signal?: AbortSignal; + onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; }; -function runProcess(context: SpawnContext): Promise { - return new Promise((resolve, reject) => { - const { command, cwd, env, timeout } = context; - const [executable, ...execArgs] = command; - if (!executable) { - reject(new ShellExecutionError("Executable is required")); - return; - } +async function runProcess(context: SpawnContext): Promise { + const { stdout, stderr, exitCode } = await spawnWithLauncher( + context.command, + { + cwd: context.cwd, + env: context.env, + timeoutMs: context.timeout, + signal: context.signal, + onOutput: context.onOutput, + }, + ); - const stdoutChunks: Buffer[] = []; - const stderrChunks: Buffer[] = []; + const stdoutLines = stdout.split("\n").filter((line) => line.length > 0); + const stderrLines = stderr.split("\n").filter((line) => line.length > 0); + const output = [stdout, stderr].filter(Boolean).join("\n").trim(); - const child = spawn(executable, execArgs, { - cwd, - env, - stdio: ["ignore", "pipe", "pipe"], - }); + if (exitCode !== 0 && exitCode !== null) { + return { + output: output || `Command exited with code ${exitCode}`, + stdout: stdoutLines, + stderr: stderrLines, + }; + } - const timeoutId = setTimeout(() => { - child.kill("SIGKILL"); - reject(new Error(`Command timed out after ${timeout}ms`)); - }, timeout); - - child.stdout.on("data", (chunk: Buffer) => { - stdoutChunks.push(chunk); - }); - - child.stderr.on("data", (chunk: Buffer) => { - stderrChunks.push(chunk); - }); - - child.on("error", (err: NodeJS.ErrnoException) => { - clearTimeout(timeoutId); - const execError = new ShellExecutionError( - err?.code === "ENOENT" - ? `Executable not found: ${executable}` - : `Failed to execute command: ${err?.message || "unknown error"}`, - ); - execError.code = err?.code; - execError.executable = executable; - reject(execError); - }); - - child.on("close", (code: number | null) => { - clearTimeout(timeoutId); - - const stdoutText = Buffer.concat(stdoutChunks).toString("utf8"); - const stderrText = Buffer.concat(stderrChunks).toString("utf8"); - - const stdoutLines = stdoutText - .split("\n") - .filter((line) => line.length > 0); - const stderrLines = stderrText - .split("\n") - .filter((line) => line.length > 0); - - const output = [stdoutText, stderrText].filter(Boolean).join("\n").trim(); - - if (code !== 0 && code !== null) { - resolve({ - output: output || `Command exited with code ${code}`, - stdout: stdoutLines, - stderr: stderrLines, - }); - } else { - resolve({ - output, - stdout: stdoutLines, - stderr: stderrLines, - }); - } - }); - }); + return { + output, + stdout: stdoutLines, + stderr: stderrLines, + }; } /** @@ -115,7 +70,7 @@ function runProcess(context: SpawnContext): Promise { export async function shell(args: ShellArgs): Promise { validateRequiredParams(args, ["command"], "shell"); - const { command, workdir, timeout_ms } = args; + const { command, workdir, timeout_ms, signal, onOutput } = args; if (!Array.isArray(command) || command.length === 0) { throw new Error("command must be a non-empty array of strings"); } @@ -132,6 +87,8 @@ export async function shell(args: ShellArgs): Promise { cwd, env: getShellEnv(), timeout, + signal, + onOutput, }; try { diff --git a/src/tools/impl/ShellCommand.ts b/src/tools/impl/ShellCommand.ts index 8528e7e..e7f1857 100644 --- a/src/tools/impl/ShellCommand.ts +++ b/src/tools/impl/ShellCommand.ts @@ -1,5 +1,6 @@ -import { ShellExecutionError, shell } from "./Shell.js"; +import { shell } from "./Shell.js"; import { buildShellLaunchers } from "./shellLaunchers.js"; +import { ShellExecutionError } from "./shellRunner.js"; import { validateRequiredParams } from "./validation.js"; interface ShellCommandArgs { @@ -8,6 +9,8 @@ interface ShellCommandArgs { timeout_ms?: number; with_escalated_permissions?: boolean; justification?: string; + signal?: AbortSignal; + onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; } interface ShellCommandResult { @@ -31,6 +34,8 @@ export async function shell_command( timeout_ms, with_escalated_permissions, justification, + signal, + onOutput, } = args; const launchers = buildShellLaunchers(command); if (launchers.length === 0) { @@ -48,6 +53,8 @@ export async function shell_command( timeout_ms, with_escalated_permissions, justification, + signal, + onOutput, }); } catch (error) { if (error instanceof ShellExecutionError && error.code === "ENOENT") { diff --git a/src/tools/impl/shellRunner.ts b/src/tools/impl/shellRunner.ts new file mode 100644 index 0000000..40272b8 --- /dev/null +++ b/src/tools/impl/shellRunner.ts @@ -0,0 +1,135 @@ +import { spawn } from "node:child_process"; + +export class ShellExecutionError extends Error { + code?: string; + executable?: string; +} + +export type ShellSpawnOptions = { + cwd: string; + env: NodeJS.ProcessEnv; + timeoutMs: number; + signal?: AbortSignal; + onOutput?: (chunk: string, stream: "stdout" | "stderr") => void; +}; + +const ABORT_KILL_TIMEOUT_MS = 2000; + +/** + * Spawn a command with a specific launcher. + * Returns a promise that resolves with the output or rejects with an error. + */ +export function spawnWithLauncher( + launcher: string[], + options: ShellSpawnOptions, +): Promise<{ stdout: string; stderr: string; exitCode: number | null }> { + return new Promise((resolve, reject) => { + const [executable, ...args] = launcher; + if (!executable) { + reject(new ShellExecutionError("Executable is required")); + return; + } + + const childProcess = spawn(executable, args, { + cwd: options.cwd, + env: options.env, + shell: false, + stdio: ["ignore", "pipe", "pipe"], + }); + + const stdoutChunks: Buffer[] = []; + const stderrChunks: Buffer[] = []; + let timedOut = false; + let killTimer: ReturnType | null = null; + + const timeoutId = setTimeout(() => { + timedOut = true; + childProcess.kill("SIGTERM"); + }, options.timeoutMs); + + const abortHandler = () => { + childProcess.kill("SIGTERM"); + if (!killTimer) { + killTimer = setTimeout(() => { + if (childProcess.exitCode === null && !childProcess.killed) { + childProcess.kill("SIGKILL"); + } + }, ABORT_KILL_TIMEOUT_MS); + } + }; + if (options.signal) { + options.signal.addEventListener("abort", abortHandler, { once: true }); + } + + childProcess.stdout?.on("data", (chunk: Buffer) => { + stdoutChunks.push(chunk); + options.onOutput?.(chunk.toString("utf8"), "stdout"); + }); + + childProcess.stderr?.on("data", (chunk: Buffer) => { + stderrChunks.push(chunk); + options.onOutput?.(chunk.toString("utf8"), "stderr"); + }); + + childProcess.on("error", (err: NodeJS.ErrnoException) => { + clearTimeout(timeoutId); + if (killTimer) { + clearTimeout(killTimer); + killTimer = null; + } + if (options.signal) { + options.signal.removeEventListener("abort", abortHandler); + } + + const execError = new ShellExecutionError( + err?.code === "ENOENT" + ? `Executable not found: ${executable}` + : `Failed to execute command: ${err?.message || "unknown error"}`, + ); + execError.code = err?.code; + execError.executable = executable; + reject(execError); + }); + + childProcess.on("close", (code) => { + clearTimeout(timeoutId); + if (killTimer) { + clearTimeout(killTimer); + killTimer = null; + } + if (options.signal) { + options.signal.removeEventListener("abort", abortHandler); + } + + const stdout = Buffer.concat(stdoutChunks).toString("utf8"); + const stderr = Buffer.concat(stderrChunks).toString("utf8"); + + if (timedOut) { + reject( + Object.assign(new Error("Command timed out"), { + killed: true, + signal: "SIGTERM", + stdout, + stderr, + code, + }), + ); + return; + } + + if (options.signal?.aborted) { + reject( + Object.assign(new Error("The operation was aborted"), { + name: "AbortError", + code: "ABORT_ERR", + stdout, + stderr, + }), + ); + return; + } + + resolve({ stdout, stderr, exitCode: code }); + }); + }); +} diff --git a/src/tools/manager.ts b/src/tools/manager.ts index 6d92f51..1229e0d 100644 --- a/src/tools/manager.ts +++ b/src/tools/manager.ts @@ -5,6 +5,15 @@ import { telemetry } from "../telemetry"; import { TOOL_DEFINITIONS, type ToolName } from "./toolDefinitions"; export const TOOL_NAMES = Object.keys(TOOL_DEFINITIONS) as ToolName[]; +const STREAMING_SHELL_TOOLS = new Set([ + "Bash", + "shell_command", + "ShellCommand", + "shell", + "Shell", + "run_shell_command", + "RunShellCommand", +]); // Maps internal tool names to server/model-facing tool names // This allows us to have multiple implementations (e.g., write_file_gemini, Write from Anthropic) @@ -726,8 +735,7 @@ export async function executeTool( // Inject options for tools that support them without altering schemas let enhancedArgs = args; - // Inject abort signal and streaming callback for Bash tool - if (internalName === "Bash") { + if (STREAMING_SHELL_TOOLS.has(internalName)) { if (options?.signal) { enhancedArgs = { ...enhancedArgs, signal: options.signal }; }