maestro

/**
 * @file codeTools.ts
 * Code execution tools: JavaScript (Deno), Python, shell commands, terminal, tests.
 * Exports originalRunJs/PyImplementation for use by the sub-agent dispatcher.
 */

import { text, tool, type Tool } from "@lmstudio/sdk";
import { spawn } from "child_process";
import { rm, writeFile } from "fs/promises";
import { join } from "path";
import { z } from "zod";
import { createSafeToolImplementation, type ToolContext } from "./shared";
import { findLMStudioHome } from "./findLMStudioHome";

function getDenoPath() {
  const lmstudioHome = findLMStudioHome();
  const utilPath = join(lmstudioHome, ".internal", "utils");
  return join(utilPath, process.platform === "win32" ? "deno.exe" : "deno");
}

export interface CodeToolsResult {
  tools: Tool[];
  /** Exposed for sub-agent tool dispatch. */
  originalRunJavascript: (params: { javascript: string; timeout_seconds?: number }) => Promise<{ stdout: string; stderr: string }>;
  originalRunPython: (params: { python: string; timeout_seconds?: number }) => Promise<{ stdout: string; stderr: string }>;
}

export function createCodeTools(
  ctx: ToolContext,
  config: { allowJavascript: boolean; allowPython: boolean; allowShell: boolean; allowTerminal: boolean },
): CodeToolsResult {
  const tools: Tool[] = [];
  const MAX_OUTPUT = 4000;

  const originalRunJavascript = async ({ javascript, timeout_seconds }: { javascript: string; timeout_seconds?: number }) => {
    const scriptFileName = `temp_script_${Date.now()}.ts`;
    const scriptFilePath = join(ctx.cwd, scriptFileName);
    try {
      await writeFile(scriptFilePath, javascript, "utf-8");
      const childProcess = spawn(getDenoPath(), ["run", "--allow-read=.", "--allow-write=.", "--no-prompt", "--deny-net", "--deny-env", "--deny-sys", "--deny-run", "--deny-ffi", scriptFilePath], {
        cwd: ctx.cwd, timeout: (timeout_seconds ?? 5) * 1000, stdio: "pipe", env: { NO_COLOR: "true" },
      });
      let stdout = "", stderr = "";
      childProcess.stdout.setEncoding("utf-8"); childProcess.stderr.setEncoding("utf-8");
      childProcess.stdout.on("data", d => stdout += d);
      childProcess.stderr.on("data", d => stderr += d);
      await new Promise<void>((resolve, reject) => {
        childProcess.on("close", code => code === 0 ? resolve() : reject(new Error(`Process exited with code ${code}. Stderr: ${stderr}`)));
        childProcess.on("error", reject);
      });
      const outJs = stdout.trim(), errJs = stderr.trim();
      return {
        stdout: outJs.length > MAX_OUTPUT ? outJs.substring(0, MAX_OUTPUT) + `\n... (truncated, ${outJs.length} chars total)` : outJs,
        stderr: errJs.length > MAX_OUTPUT ? errJs.substring(0, MAX_OUTPUT) + `\n... (truncated)` : errJs,
      };
    } finally {
      await rm(scriptFilePath, { force: true }).catch(() => {});
    }
  };

  tools.push(tool({
    name: "run_javascript",
    description: text`
      Run a JavaScript code snippet using deno. You cannot import external modules but you have
      read/write access to the current working directory.
      Pass the code you wish to run as a string in the 'javascript' parameter.
      By default, the code will timeout in 5 seconds. You can extend this timeout by setting the
      'timeout_seconds' parameter to a higher value in seconds, up to a maximum of 60 seconds.
      You will get the stdout and stderr output of the code execution, thus please print the output
      you wish to return using 'console.log' or 'console.error'.
    `,
    parameters: { javascript: z.string(), timeout_seconds: z.number().min(0.1).max(60).optional() },
    implementation: createSafeToolImplementation(originalRunJavascript, config.allowJavascript, "run_javascript"),
  }));

  const originalRunPython = async ({ python, timeout_seconds }: { python: string; timeout_seconds?: number }) => {
    const scriptFileName = `temp_script_${Date.now()}.py`;
    const scriptFilePath = join(ctx.cwd, scriptFileName);
    try {
      await writeFile(scriptFilePath, python, "utf-8");
      const childProcess = spawn("python", [scriptFilePath], {
        cwd: ctx.cwd, timeout: (timeout_seconds ?? 5) * 1000, stdio: "pipe",
      });
      let stdout = "", stderr = "";
      childProcess.stdout.setEncoding("utf-8"); childProcess.stderr.setEncoding("utf-8");
      childProcess.stdout.on("data", d => stdout += d);
      childProcess.stderr.on("data", d => stderr += d);
      await new Promise<void>((resolve, reject) => {
        childProcess.on("close", code => code === 0 ? resolve() : reject(new Error(`Process exited with code ${code}. Stderr: ${stderr}`)));
        childProcess.on("error", reject);
      });
      const outStr = stdout.trim(), errStr = stderr.trim();
      return {
        stdout: outStr.length > MAX_OUTPUT ? outStr.substring(0, MAX_OUTPUT) + `\n... (truncated, ${outStr.length} chars total)` : outStr,
        stderr: errStr.length > MAX_OUTPUT ? errStr.substring(0, MAX_OUTPUT) + `\n... (truncated)` : errStr,
      };
    } finally {
      await rm(scriptFilePath, { force: true }).catch(() => {});
    }
  };

  tools.push(tool({
    name: "run_python",
    description: text`
      Run a Python code snippet. You cannot import external modules but you have
      read/write access to the current working directory.
      Pass the code you wish to run as a string in the 'python' parameter.
      By default, the code will timeout in 5 seconds. You can extend this timeout by setting the
      'timeout_seconds' parameter to a higher value in seconds, up to a maximum of 60 seconds.
      You will get the stdout and stderr output of the code execution, thus please print the output
      you wish to return using 'print()'.
    `,
    parameters: { python: z.string(), timeout_seconds: z.number().min(0.1).max(60).optional() },
    implementation: createSafeToolImplementation(originalRunPython, config.allowPython, "run_python"),
  }));

  const originalExecuteCommand = async ({ command, input, timeout_seconds }: { command: string; input?: string; timeout_seconds?: number }) => {
    const childProcess = spawn(command, [], {
      cwd: ctx.cwd, shell: true, timeout: (timeout_seconds ?? 5) * 1000, stdio: "pipe",
    });
    if (input) { childProcess.stdin.write(input); childProcess.stdin.end(); } else { childProcess.stdin.end(); }
    let stdout = "", stderr = "";
    childProcess.stdout.setEncoding("utf-8"); childProcess.stderr.setEncoding("utf-8");
    childProcess.stdout.on("data", d => stdout += d);
    childProcess.stderr.on("data", d => stderr += d);
    await new Promise<void>((resolve, reject) => {
      childProcess.on("close", code => code === 0 ? resolve() : reject(new Error(`Process exited with code ${code}. Stderr: ${stderr}`)));
      childProcess.on("error", reject);
    });
    const outStr = stdout.trim();
    const errStr = stderr.trim();
    return {
      stdout: outStr.length > MAX_OUTPUT ? outStr.substring(0, MAX_OUTPUT) + `\n... (truncated, ${outStr.length} chars total)` : outStr,
      stderr: errStr.length > MAX_OUTPUT ? errStr.substring(0, MAX_OUTPUT) + `\n... (truncated)` : errStr,
    };
  };

  tools.push(tool({
    name: "execute_command",
    description: text`
      Execute a shell command in the current working directory.
      Returns the stdout and stderr output of the command.
      You can optionally provide input to be piped to the command's stdin.
      IMPORTANT: The host operating system is '${process.platform}'.
      If the OS is 'win32' (Windows), do NOT use 'bash' or 'sh' commands unless you are certain WSL is available.
      Instead, use standard Windows 'cmd' or 'powershell' syntax.
    `,
    parameters: {
      command: z.string(),
      input: z.string().optional().describe("Input text to pipe to the command's stdin."),
      timeout_seconds: z.number().min(0.1).max(60).optional().describe("Timeout in seconds (default: 5, max: 60)"),
    },
    implementation: createSafeToolImplementation(originalExecuteCommand, config.allowShell, "execute_command"),
  }));

  const originalRunInTerminal = async ({ command }: { command: string }) => {
    if (process.platform === "win32") {
      const escapedDir = ctx.cwd.replace(/"/g, '""');
      const escapedCmd = command.replace(/"/g, '""');
      const child = spawn("cmd.exe", ["/c", `start "" /D "${escapedDir}" cmd.exe /k "${escapedCmd}"`], { detached: true, stdio: "ignore", windowsHide: false });
      child.unref();
    } else if (process.platform === "darwin") {
      const safeCmd = command.replace(/\\/g, "\\\\").replace(/"/g, '\\"');
      const safeCwd = ctx.cwd.replace(/\\/g, "\\\\").replace(/"/g, '\\"');
      const child = spawn("osascript", ["-e", `tell application "Terminal"\ndo script "cd \\"${safeCwd}\\" && ${safeCmd}"\nactivate\nend tell`], { detached: true, stdio: "ignore" });
      child.unref();
    } else {
      const safeCwd = ctx.cwd.replace(/'/g, "'\\''");
      const safeCmd = command.replace(/'/g, "'\\''");
      const bashScript = `cd '${safeCwd}' && ${safeCmd}; bash`;
      const child = spawn("x-terminal-emulator", ["-e", "bash", "-c", bashScript], { detached: true, stdio: "ignore" });
      child.on("error", () => {
        const child2 = spawn("gnome-terminal", ["--", "bash", "-c", bashScript], { detached: true, stdio: "ignore" });
        child2.unref();
      });
      child.unref();
    }
    return { success: true, message: "Terminal window launched. Please check your taskbar." };
  };

  tools.push(tool({
    name: "run_in_terminal",
    description: text`Launch a command in a new, separate interactive terminal window.`,
    parameters: { command: z.string() },
    implementation: createSafeToolImplementation(originalRunInTerminal, config.allowTerminal, "run_in_terminal"),
  }));

  tools.push(tool({
    name: "run_test_command",
    description: "Execute a test command (like 'npm test') and return the results.",
    parameters: { command: z.string().describe("The test command to run (e.g., 'npm test', 'pytest').") },
    implementation: async ({ command }) => {
      return new Promise((resolve) => {
        const parts = command.split(" ");
        const TIMEOUT_MS = 120_000;
        const child = spawn(parts[0], parts.slice(1), { cwd: ctx.cwd, shell: true, env: { ...process.env, CI: 'true' } });
        let stdout = "", stderr = "", timedOut = false;
        const timer = setTimeout(() => { timedOut = true; child.kill("SIGTERM"); }, TIMEOUT_MS);
        child.stdout.on("data", d => stdout += d.toString());
        child.stderr.on("data", d => stderr += d.toString());
        child.on("close", code => {
          clearTimeout(timer);
          const out = stdout.trim();
          const err = stderr.trim();
          resolve({
            command, exit_code: code, passed: code === 0,
            stdout: out.length > MAX_OUTPUT ? out.substring(0, MAX_OUTPUT) + `\n... (truncated, ${out.length} chars total)` : out,
            stderr: err.length > MAX_OUTPUT ? err.substring(0, MAX_OUTPUT) + `\n... (truncated)` : err,
            ...(timedOut ? { warning: `Process killed after ${TIMEOUT_MS / 1000}s timeout` } : {}),
          });
        });
        child.on("error", err => { clearTimeout(timer); resolve({ command, error: err.message, passed: false }); });
      });
    },
  }));

  return { tools, originalRunJavascript, originalRunPython };
}