python-tools

import { spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { mkdir, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";

import { resolvePythonCommand } from "./pythonResolver";
import {
  ensurePythonFile,
  validateArgs,
  validatePathInput,
  validatePythonFilePath,
  validateWindowTitle,
  validateWorkingDirectory,
} from "./safePaths";

export interface LaunchInteractiveRequest {
  filePath: string;
  args?: string[];
  cwd?: string;
  windowTitle?: string;
}

export interface LaunchInteractiveResult {
  launched: true;
  pythonExecutableUsed: string;
  filePath: string;
  args: string[];
  cwd: string;
  launcherPath: string;
}

export async function writeTemporaryPythonFile(code: string): Promise<string> {
  validateCode(code);
  const tempDir = path.join(tmpdir(), "lmstudio-python-tools");
  const scriptPath = path.join(tempDir, `interactive-${randomUUID()}.py`);

  await mkdir(tempDir, { recursive: true });
  await writeFile(scriptPath, code, "utf8");
  return scriptPath;
}

export async function launchPythonFileInteractive(
  request: LaunchInteractiveRequest,
): Promise<LaunchInteractiveResult> {
  const resolvedFilePath = await ensurePythonFile(request.filePath);
  return await launchValidatedPythonFileInteractive(request, resolvedFilePath);
}

export async function launchTemporaryPythonFileInteractive(
  request: LaunchInteractiveRequest,
): Promise<LaunchInteractiveResult> {
  const resolvedFilePath = await validatePythonFilePath(request.filePath);
  return await launchValidatedPythonFileInteractive(request, resolvedFilePath);
}

async function launchValidatedPythonFileInteractive(
  request: LaunchInteractiveRequest,
  resolvedFilePath: string,
): Promise<LaunchInteractiveResult> {
  const resolvedArgs = validateArgs(request.args);
  const resolvedCwd = (await validateWorkingDirectory(request.cwd)) ?? path.dirname(resolvedFilePath);
  const resolvedWindowTitle = validateWindowTitle(request.windowTitle);
  const pythonCommand = await resolvePythonCommand();

  const launcherPath = await createWindowsLauncher({
    pythonExecutable: pythonCommand.pythonExecutableUsed,
    scriptPath: resolvedFilePath,
    args: resolvedArgs,
    cwd: resolvedCwd,
    windowTitle: resolvedWindowTitle,
  });
  await launchWindowsTerminal(launcherPath);

  return {
    launched: true,
    pythonExecutableUsed: pythonCommand.pythonExecutableUsed,
    filePath: resolvedFilePath,
    args: resolvedArgs,
    cwd: resolvedCwd,
    launcherPath,
  };
}

function validateCode(code: string): void {
  if (code.trim().length === 0) {
    throw new Error("Python code cannot be empty.");
  }

  if (code.length > 50_000) {
    throw new Error("Python code must be 50000 characters or fewer.");
  }
}

async function createWindowsLauncher(input: {
  pythonExecutable: string;
  scriptPath: string;
  args: string[];
  cwd: string;
  windowTitle: string;
}): Promise<string> {
  if (process.platform !== "win32") {
    throw new Error("Interactive launch is currently supported on Windows only.");
  }

  validatePathInput(input.pythonExecutable, "pythonExecutable");
  validatePathInput(input.scriptPath, "scriptPath");
  validatePathInput(input.cwd, "cwd");

  const tempDir = path.join(tmpdir(), "lmstudio-python-tools");
  const launcherPath = path.join(tempDir, `launch-${randomUUID()}.cmd`);
  const argsForCmd = input.args.map(quoteForCmd).join(" ");
  const runCommand = `${quoteForCmd(input.pythonExecutable)} ${quoteForCmd(input.scriptPath)}${argsForCmd.length > 0 ? ` ${argsForCmd}` : ""}`;
  const content = [
    "@echo off",
    `title ${input.windowTitle}`,
    `cd /d ${quoteForCmd(input.cwd)}`,
    "echo.",
    `echo Python executable: ${quoteForCmd(input.pythonExecutable)}`,
    `echo Script: ${quoteForCmd(input.scriptPath)}`,
    "echo.",
    runCommand,
    "set EXIT_CODE=%ERRORLEVEL%",
    "echo.",
    "echo Python exited with code %EXIT_CODE%.",
    "pause",
    "",
  ].join("\r\n");

  await mkdir(tempDir, { recursive: true });
  await writeFile(launcherPath, content, "utf8");
  return launcherPath;
}

async function launchWindowsTerminal(launcherPath: string): Promise<void> {
  if (process.platform !== "win32") {
    throw new Error("Interactive launch is currently supported on Windows only.");
  }

  validatePathInput(launcherPath, "launcherPath");

  try {
    await spawnDetached("wt.exe", ["cmd.exe", "/k", launcherPath]);
    return;
  } catch (error) {
    if (!(error instanceof Error) || !error.message.includes("ENOENT")) {
      throw error;
    }
  }

  const wrapperPath = await createCmdFallbackWrapper(launcherPath);
  await spawnDetached("cmd.exe", ["/c", wrapperPath]);
}

async function createCmdFallbackWrapper(launcherPath: string): Promise<string> {
  const tempDir = path.join(tmpdir(), "lmstudio-python-tools");
  const wrapperPath = path.join(tempDir, `open-${randomUUID()}.cmd`);
  const content = [
    "@echo off",
    `start "" cmd.exe /k ${quoteForCmd(launcherPath)}`,
    "",
  ].join("\r\n");

  await mkdir(tempDir, { recursive: true });
  await writeFile(wrapperPath, content, "utf8");
  return wrapperPath;
}

function quoteForCmd(value: string): string {
  if (value.includes("\n") || value.includes("\r")) {
    throw new Error("Command value must not contain newline characters.");
  }

  return `"${value.replace(/%/g, "%%").replace(/"/g, "\"\"")}"`;
}

function spawnDetached(command: string, args: string[]): Promise<void> {
  return new Promise((resolve, reject) => {
    const child = spawn(command, args, {
      shell: false,
      detached: true,
      stdio: "ignore",
      windowsHide: false,
    });

    child.once("error", (error) => {
      reject(new Error(`Failed to launch interactive Python window with ${command}: ${error.message}`));
    });

    child.once("spawn", () => {
      child.unref();
      resolve();
    });
  });
}