sagemath-podman-sandbox

import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { spawn } from "child_process";
import { z } from "zod";
import { configSchematics, globalConfigSchematics } from "./config";

const CONTAINER_USER = "sage";
const MAX_CODE_CHARS = 50_000;
const MAX_TIMEOUT_SECONDS = 600;
const NOFILE_LIMIT = "1024:1024";
const FSIZE_LIMIT_BYTES = "134217728:134217728"; // 128 MiB
const MEMORY_LIMITS = new Set(["256m", "512m", "1g", "2g", "4g", "8g"]);
const DEFAULT_SAGE_IMAGE = "docker.io/sagemath/sagemath:latest";
const ROOTLESS_CHECK_TIMEOUT_MS = 5_000;

type RootlessCheckResult = {
  ok: boolean;
  reported: "true" | "false" | "unknown";
  stdout?: string;
  stderr?: string;
  error?: string;
  exitCode?: number | null;
  timedOut?: boolean;
}

interface SandboxConfig {
  podmanExecutable: string;
  sageImage: string;
  defaultTimeoutSeconds: number;
  memoryLimit: string;
  cpuLimit: number;
  pidsLimit: number;
  outputLimitBytes: number;
}

function clampInt(value: number, min: number, max: number): number {
  if (!Number.isFinite(value)) {return min;}
  return Math.max(min, Math.min(max, Math.trunc(value)));
}

function clampNumber(value: number, min: number, max: number): number {
  if (!Number.isFinite(value)) {return min;}
  return Math.max(min, Math.min(max, value));
}

function normalizeExecutable(value: string): string {
  const trimmed = value.trim();
  return trimmed.length > 0 ? trimmed : "podman";
}

function normalizeImage(value: string): string {
  const trimmed = value.trim();
  if (trimmed.length === 0) {return DEFAULT_SAGE_IMAGE;}
  if (trimmed.startsWith("-")) {return DEFAULT_SAGE_IMAGE;}
  if (/[\s\x00-\x1F\x7F]/.test(trimmed)) {return DEFAULT_SAGE_IMAGE;}
  return trimmed;
}

function normalizeMemoryLimit(value: unknown): string {
  const s = String(value ?? "").trim().toLowerCase();
  if (MEMORY_LIMITS.has(s)) {return s;}
  return "256m";
}

function readSandboxConfig(ctl: ToolsProviderController): SandboxConfig {
  const pluginConfig = ctl.getPluginConfig(configSchematics);
  const globalConfig = ctl.getGlobalPluginConfig(globalConfigSchematics);
  
  const defaultTimeoutSeconds = clampInt(pluginConfig.get("defaultTimeoutSeconds"), 1, MAX_TIMEOUT_SECONDS,);
  const pidsLimit = clampInt(pluginConfig.get("pidsLimit"), 16, 512);
  const outputLimitKiB = clampInt(pluginConfig.get("outputLimitKiB"), 16, 2048);
  const cpuLimit = clampNumber(pluginConfig.get("cpuLimit"), 0.1, 4);
  
  return {
    podmanExecutable: normalizeExecutable(globalConfig.get("podmanExecutable")),
    sageImage: normalizeImage(globalConfig.get("sageImage")),
    defaultTimeoutSeconds,
    memoryLimit: normalizeMemoryLimit(pluginConfig.get("memoryLimit")),
    cpuLimit,
    pidsLimit,
    outputLimitBytes: outputLimitKiB * 1024,
  };
}

function safeContainerName(): string {
  const random = Math.random().toString(36).slice(2, 10);
  return `lms-sage-${Date.now()}-${random}`.replace(/[^a-zA-Z0-9_.-]/g, "-");
}

function utf8ByteLength(s: string): number {return Buffer.byteLength(s, "utf8");}

function appendLimited(current: string, chunk: Buffer, limitBytes: number) {
  const room = limitBytes - utf8ByteLength(current);
  if (room <= 0) {return { value: current, exceeded: true };}
  
  if (chunk.byteLength <= room) {return { value: current + chunk.toString("utf8"), exceeded: false };}
  
  return {
    value:
      current +
      chunk.subarray(0, Math.max(0, room)).toString("utf8") +
      "\n...[output truncated by sage podman sandbox]...",
    exceeded: true,
  };
}

function podmanBaseEnv(): NodeJS.ProcessEnv {
  // Do not pass secrets/API keys from LM Studio's environment into the Sage process.
  // PATH is retained only so the configured Podman executable can be found.
  return {
    PATH: process.env.PATH ?? "",
    HOME: process.env.HOME ?? "",
    XDG_RUNTIME_DIR: process.env.XDG_RUNTIME_DIR ?? "",
    DBUS_SESSION_BUS_ADDRESS: process.env.DBUS_SESSION_BUS_ADDRESS ?? "",
    TMPDIR: process.env.TMPDIR ?? "/tmp",
  };
}

function buildPodmanArgs(containerName: string, cfg: SandboxConfig): string[] {
  return [
    "run", "--rm", "-i", "--name", containerName,
    
    // No image/network surprise. Pull images manually before using the plugin.
    "--pull=never",
    "--network=none",
    
    // The container sees no host paths. Its root filesystem is read-only.
    "--read-only",
    "--tmpfs", "/tmp:rw,nosuid,nodev,size=256m,mode=1777",
    "--workdir", "/tmp",
    "--env", "HOME=/tmp/sage-home",
    "--env", "DOT_SAGE=/tmp/dot-sage",
    "--env", "SAGE_NUM_THREADS=1",
    "--env", "OPENBLAS_NUM_THREADS=1",
    "--env", "OMP_NUM_THREADS=1",
    "--env", "MKL_NUM_THREADS=1",
    "--env", "PYTHONNOUSERSITE=1",
    
    // Do not run Sage as root in the container.
    "--user", CONTAINER_USER,
    
    // Drop privilege and kernel attack surface as much as Podman allows.
    "--cap-drop=all",
    "--security-opt", "no-new-privileges",
    "--pids-limit", String(cfg.pidsLimit),
    
    // Resource controls. Rootless Podman requires a reasonably modern cgroups setup.
    "--memory", cfg.memoryLimit,
    "--cpus", String(cfg.cpuLimit),
    "--ulimit", `nofile=${NOFILE_LIMIT}`,
    "--ulimit", `nproc=${cfg.pidsLimit}:${cfg.pidsLimit}`,
    "--ulimit", `fsize=${FSIZE_LIMIT_BYTES}`,
    
    // Metadata only; useful when inspecting local containers.
    "--label", "lmstudio.tool=sage_podman_sandbox",
    "--label", `lmstudio.default_timeout_seconds=${cfg.defaultTimeoutSeconds}`,
    
    "--entrypoint", "/bin/bash",
    
    "--",
    cfg.sageImage,
    "-lc",
    [
      "set -euo pipefail",
      "umask 077",
      "mkdir -p /tmp/sage-home /tmp/dot-sage",
      "chmod 700 /tmp/sage-home /tmp/dot-sage",
      "cat > /tmp/user_code.sage",
      "/usr/bin/sage -q /tmp/user_code.sage",
    ].join("; "),
  ];
}

async function checkPodmanRootless(cfg: SandboxConfig): Promise<RootlessCheckResult> {
  const result = await runCommand(
    cfg.podmanExecutable,
    ["info", "--format", "{{.Host.Security.Rootless}}"],
    ROOTLESS_CHECK_TIMEOUT_MS,
  );
  
  const reportedText = result.stdout.trim().toLowerCase();
  
  if (!result.ok) {
    return {
      ok: false,
      reported: "unknown",
      stdout: result.stdout || undefined,
      stderr: result.stderr || undefined,
      error: result.error || result.stderr || `podman info exited with code ${result.exitCode}`,
      exitCode: result.exitCode,
      timedOut: result.timedOut,
    };
  }
  
  if (reportedText === "true") {
    return {
      ok: true,
      reported: "true",
      stdout: result.stdout,
      stderr: result.stderr || undefined,
      exitCode: result.exitCode,
      timedOut: result.timedOut,
    };
  }
  
  if (reportedText === "false") {
    return {
      ok: false,
      reported: "false",
      stdout: result.stdout,
      stderr: result.stderr || undefined,
      error: "Podman reports rootless=false.",
      exitCode: result.exitCode,
      timedOut: result.timedOut,
    };
  }
  
  return {
    ok: false,
    reported: "unknown",
    stdout: result.stdout || undefined,
    stderr: result.stderr || undefined,
    error: `Unexpected Podman rootless report: ${JSON.stringify(result.stdout)}`,
    exitCode: result.exitCode,
    timedOut: result.timedOut,
  };
}

function runCommand(command: string, args: string[], timeoutMs: number): Promise<{
  ok: boolean;
  exitCode: number | null;
  signal: NodeJS.Signals | null;
  stdout: string;
  stderr: string;
  error?: string;
  timedOut: boolean;
}> {
  return new Promise(resolve => {
    let stdout = "";
    let stderr = "";
    let timedOut = false;
    let settled = false;
    
    function finish(result: {
      ok: boolean;
      exitCode: number | null;
      signal: NodeJS.Signals | null;
      stdout: string;
      stderr: string;
      error?: string;
      timedOut: boolean;
    }) {
      if (settled) {return;}
      settled = true;
      clearTimeout(timer);
      resolve(result);
    }
    
    const child = spawn(command, args, {
      stdio: ["ignore", "pipe", "pipe"],
      env: podmanBaseEnv(),
    });
    
    const timer = setTimeout(() => {
      timedOut = true;
      child.kill("SIGKILL");
    }, timeoutMs);
    
    child.stdout.on("data", (chunk: Buffer) => {stdout += chunk.toString("utf8");});
    child.stderr.on("data", (chunk: Buffer) => {stderr += chunk.toString("utf8");});
    
    child.on("error", err => {
      finish({
        ok: false,
        exitCode: null,
        signal: null,
        stdout: stdout.trimEnd(),
        stderr: stderr.trimEnd(),
        error: err.message,
        timedOut,
      });
    });
    
    child.on("close", (code, signal) => {
      finish({
        ok: code === 0 && !timedOut,
        exitCode: code,
        signal,
        stdout: stdout.trimEnd(),
        stderr: stderr.trimEnd(),
        timedOut,
      });
    });
  });
}

async function forceRemoveContainer(cfg: SandboxConfig, containerName: string): Promise<{attempted: true; ok: boolean; error?: string}> {
  const result = await runCommand(cfg.podmanExecutable, ["rm", "-f", containerName], 5_000);
  if (result.ok) {return { attempted: true, ok: true };}
  return {attempted: true, ok: false,
          error: result.error || result.stderr || `podman rm exited with ${result.exitCode}`
         };
}

async function runSageInPodmanSandbox(code: string, timeoutSeconds: number, cfg: SandboxConfig) {
  const rootlessCheck = await checkPodmanRootless(cfg);
  
  if (!rootlessCheck.ok) {
    return {
      ok: false,
      error: "Refusing to run SageMath: Podman is not confirmed to be running rootless.",
      hint: "Run LM Studio and Podman as a normal user, not with sudo. Check manually with: podman info --format '{{.Host.Security.Rootless}}'. It should print true.",
      rootless_check: rootlessCheck,
      sandbox: sandboxSummary(cfg, timeoutSeconds),
    };
  }
  
  const containerName = safeContainerName();
  const args = buildPodmanArgs(containerName, cfg);
  
  let stdout = "";
  let stderr = "";
  let stdoutTruncated = false;
  let stderrTruncated = false;
  let timedOut = false;
  let outputLimitExceeded = false;
  
  const child = spawn(cfg.podmanExecutable, args, {stdio: ["pipe", "pipe", "pipe"], env: podmanBaseEnv()});
  
  child.stdout.on("data", (chunk: Buffer) => {
    const result = appendLimited(stdout, chunk, cfg.outputLimitBytes);
    stdout = result.value;
    stdoutTruncated ||= result.exceeded;
    outputLimitExceeded ||= result.exceeded;
    if (result.exceeded) {
      child.kill("SIGKILL");
      void forceRemoveContainer(cfg, containerName);
    }
  });
  
  child.stderr.on("data", (chunk: Buffer) => {
    const result = appendLimited(stderr, chunk, cfg.outputLimitBytes);
    stderr = result.value;
    stderrTruncated ||= result.exceeded;
    outputLimitExceeded ||= result.exceeded;
    if (result.exceeded) {
      child.kill("SIGKILL");
      void forceRemoveContainer(cfg, containerName);
    }
  });
  
  child.stdin.on("error", () => {}); // ignore or capture broken pipe
  
  let cleanupTimer: NodeJS.Timeout | undefined;
  
  const resultPromise = new Promise<{
    exitCode: number | null;
    signal: NodeJS.Signals | null;
    spawnError?: string;
  }>(resolve => {
    let settled = false;
    function finish(result: {
      exitCode: number | null;
      signal: NodeJS.Signals | null;
      spawnError?: string;
    }) {
      if (settled) return;
      settled = true;
      if (cleanupTimer !== undefined) {clearTimeout(cleanupTimer);}
      resolve(result);
    }
    child.on("error", err => {
      finish({
        exitCode: null,
        signal: null,
        spawnError: err.message,
      });
    });
    child.on("close", (code, signal) => {
      finish({
        exitCode: code,
        signal,
      });
    });
  });
  
  cleanupTimer = setTimeout(() => {
    timedOut = true;
    child.kill("SIGKILL");
    void forceRemoveContainer(cfg, containerName);
  }, timeoutSeconds * 1000); // setTimeout expects a number of milliseconds
  
  child.stdin.write(code, "utf8");
  child.stdin.end();
  
  const result = await resultPromise;
  
  if (result.spawnError) {
    return {
      ok: false,
      error: `Could not start Podman executable '${cfg.podmanExecutable}': ${result.spawnError}`,
      hint:
        "Install rootless Podman, pull the SageMath image manually, then configure this plugin in LM Studio's plugin settings UI.",
      sandbox: sandboxSummary(cfg, timeoutSeconds),
    };
  }
  
  return {
    ok: result.exitCode === 0 && !timedOut && !outputLimitExceeded,
    exit_code: result.exitCode,
    signal: result.signal,
    timed_out: timedOut,
    output_limit_exceeded: outputLimitExceeded,
    stdout: stdout.trimEnd(),
    stderr: stderr.trimEnd(),
    truncated: {
      stdout: stdoutTruncated,
      stderr: stderrTruncated,
      max_bytes_per_stream: cfg.outputLimitBytes,
    },
    sandbox: sandboxSummary(cfg, timeoutSeconds),
  };
}

function sandboxSummary(cfg: SandboxConfig, timeoutSeconds: number) {
  return {
    runtime: "podman",
    podman_executable: cfg.podmanExecutable,
    image: cfg.sageImage,
    timeout_seconds: timeoutSeconds,
    network: "none",
    host_mounts: "none",
    root_filesystem: "read-only",
    writable_filesystems: ["tmpfs:/tmp"],
    home_directory: "/tmp/sage-home",
    dot_sage: "/tmp/dot-sage",
    user: CONTAINER_USER,
    capabilities: "all dropped",
    no_new_privileges: true,
    memory_limit: cfg.memoryLimit,
    cpus: cfg.cpuLimit,
    pids_limit: cfg.pidsLimit,
    pull_policy: "never",
    output_limit_bytes_per_stream: cfg.outputLimitBytes,
    input_method: "stdin written to /tmp/user_code.sage",
    command: "/usr/bin/sage -q /tmp/user_code.sage",
    automatic_image_pull: false,
    warning: "This is a configuration summary, not an independent runtime verification.",
  };
}

async function podmanStatus(cfg: SandboxConfig) {
  const version = await runCommand(cfg.podmanExecutable, ["--version"], 5_000);
  const imageExists = await runCommand(cfg.podmanExecutable, ["image", "exists", cfg.sageImage], 5_000);
  const rootless = await checkPodmanRootless(cfg);
  
  return {
    ok: version.ok && imageExists.ok && rootless.ok,
    
    podman_version: version.stdout || version.stderr || version.error || null,
    rootless_reported_by_podman: rootless.reported,
    rootless_check: rootless,
    
    image_available_locally: imageExists.ok,
    image: cfg.sageImage,
    podman_executable: cfg.podmanExecutable,
    image_pull_command: [cfg.podmanExecutable, "pull", cfg.sageImage],
    
    notes: [
      "The plugin uses Podman for sandboxing.",
      "The image must already be pulled locally because execution uses --pull=never.",
      "Execution is refused unless Podman reports rootless=true.",
    ],
    
    diagnostics: {
      version_exit_code: version.exitCode,
      image_exists_exit_code: imageExists.exitCode,
      rootless_check_exit_code: rootless.exitCode,
      rootless_check_error: rootless.error,
    },
    
    sandbox: sandboxSummary(cfg, cfg.defaultTimeoutSeconds),
  };
}

export async function toolsProvider(ctl: ToolsProviderController) {
  const tools: Tool[] = [];
  
  const runSageTool = tool({
    name: "run_sagemath",
    description: text`
      Run SageMath code in a locked-down local Podman container sandbox and return stdout/stderr.
      
      Use this for algebra, number theory, symbolic math, matrices, exact arithmetic, and other SageMath tasks.
      
      Security boundaries: the Sage code runs through Podman, with no host directory mounts,
      no network, a read-only root filesystem, tmpfs-only writable directories, non-root container user,
      dropped Linux capabilities, no-new-privileges, and CPU/memory/process/time/output limits.
      The image is not pulled automatically; install it manually before use.
      Print values that should be returned.
    `,
    parameters: {
      code: z.string().max(MAX_CODE_CHARS),
      timeout_seconds: z.number().int().min(1).max(MAX_TIMEOUT_SECONDS).optional(),
    },
    implementation: async ({ code, timeout_seconds }) => {
      const cfg = readSandboxConfig(ctl);
      const timeout = clampInt(timeout_seconds ?? cfg.defaultTimeoutSeconds, 1, MAX_TIMEOUT_SECONDS);
      return await runSageInPodmanSandbox(code, timeout, cfg);
    },
  });
  
  const selfTestTool = tool({
    name: "sage_podman_sandbox_self_test",
    description: text`
      Verify that the SageMath Podman sandbox can start and execute a small calculation.
      This runs print(factor(x^4 - 1)) inside the same sandbox used by run_sagemath.
    `,
    parameters: {},
    implementation: async () => {
      const cfg = readSandboxConfig(ctl);
      const timeout = clampInt(Math.max(cfg.defaultTimeoutSeconds, 15), 1, MAX_TIMEOUT_SECONDS);
      return await runSageInPodmanSandbox("R.<x> = QQ[]\nprint(factor(x^4 - 1))\n", timeout, cfg);
    },
  });
  
  const statusTool = tool({
    name: "sage_podman_sandbox_status",
    description: text`
      Check whether Podman is reachable, whether the configured SageMath image exists locally,
      and report the active LM Studio GUI sandbox settings. This does not run Sage code.
    `,
    parameters: {},
    implementation: async () => {
      const cfg = readSandboxConfig(ctl);
      return await podmanStatus(cfg);
    },
  });
  
  tools.push(runSageTool, selfTestTool, statusTool);
  return tools;
}