Plugin for the Maki LLM harness that adds a `python` tool that executes unsandboxed Python code

python_tool.lua

-- Plugin for the Maki LLM harness that adds a `python` tool that executes unsandboxed Python code.

local truncate = require("maki.truncate")
local ToolView = require("maki.tool_view")
local output_limits = require("maki.output_limits")
local partial = require("maki.partial")

local DEFAULT_TIMEOUT_SECS = 120
local SEPARATOR = "──────"
local NO_OUTPUT = "(no output)"

local function build_env(ctx)
  local env = { GIT_TERMINAL_PROMPT = "0" }
  local config_dir = maki.env.config_dir()
  if config_dir then
    local pythonlib = config_dir .. "/pythonlib"
    local existing = maki.uv.os_getenv("PYTHONPATH") or ""
    if existing ~= "" then
      env.PYTHONPATH = pythonlib .. ":" .. existing
    else
      env.PYTHONPATH = pythonlib
    end
  end
  return env
end

local function create_python_view(code, ctx)
  local tol = ctx:tool_output_lines()
  local buf = maki.ui.buf()
  local view = ToolView.new(buf, {
    max_lines = (tol and tol.python) or 30,
    keep = "tail",
    max_line_bytes = output_limits.DEFAULT_MAX_LINE_BYTES,
  })
  local header_lines = {}
  local highlighted = maki.ui.highlight(code, "py")
  if highlighted then
    for _, line in ipairs(highlighted) do
      header_lines[#header_lines + 1] = line
    end
  else
    header_lines[#header_lines + 1] = code
  end
  header_lines[#header_lines + 1] = { { SEPARATOR, "dim" } }
  view:set_header(header_lines)
  buf:on("click", function()
    view:toggle()
  end)
  return buf, view
end

local description = [[Execute Python code in a fresh python3 process.

Use for system commands, builds, tests, and ad-hoc scripting where you would otherwise use bash.

- Available modules include the full Python stdlib and anything in ~/.config/maki/pythonlib/.
- Use subprocess to run commands: `subprocess.run(["git", "status"], check=True)`.
- Output is captured and returned to the model.
- Output truncated beyond 2000 lines or 50KB.
- Non-zero exit codes are reported as errors.
]]

maki.api.register_prompt_hint({
  slot = "tool_usage",
  content = "Use the `python` tool for scripting instead of bash. Each call is a fresh python3 process.",
})

local opts = maki.api.register_options(output_limits.extend({
  timeout_secs = {
    default = 120,
    min = 5,
    desc = "Kill Python code after this many seconds. A call's `timeout` param overrides it.",
  },
}))

maki.api.register_tool({
  name = "python",
  kind = "execute",
  description = description,
  schema = {
    type = "object",
    required = { "code" },
    additionalProperties = false,
    properties = {
      code = {
        type = "string",
        description = "Python code to execute. Runs in a fresh python3 process.",
      },
      timeout = {
        type = "integer",
        description = "Execution timeout in seconds (default 120)",
      },
      workdir = {
        type = "string",
        description = "Working directory (default: cwd)",
      },
    },
  },
  permission = "run",
  permission_scopes = function(input)
    if not input.code or input.code:match("^%s*$") then
      return nil
    end
    return { scopes = { input.code }, force_prompt = false }
  end,

  header = function(input)
    local lines = select(2, input.code:gsub("\n", "\n")) + 1
    local s = lines .. " lines of python"
    if input.timeout then
      local buf = maki.ui.buf()
      buf:line({ { s }, { " (" .. maki.ui.humantime(input.timeout) .. " timeout)", "dim" } })
      return buf
    end
    return s
  end,

  restore = function(input, output, is_error, ctx)
    local buf, view = create_python_view(input.code, ctx)
    local timeout_secs = output:match("^tool python timed out after (%d+)s$")
    if timeout_secs then
      view:append({ { "Timed out after " .. timeout_secs .. "s", "dim" } })
    elseif is_error then
      local body, code = output:match("^(.-)\nExit code: (%d+)$")
      if body then
        view:append_text(body)
        view:append({ { "Exit code: " .. code, "dim" } })
      else
        view:append_text(output)
      end
    else
      if output == "Exit code: 0" or output == "" then
        view:clear()
        view:append({ { "No output", "dim" } })
      else
        view:append_text(output)
      end
    end
    view:finish()
    return buf
  end,

  handler = function(input, ctx)
    if not input.code or input.code:match("^%s*$") then
      return { llm_output = "error: code is required", is_error = true }
    end

    local timeout_secs = input.timeout or opts.timeout_secs
    local max_lines, max_bytes = output_limits.resolve(opts, ctx)

    ctx:set_deadline(timeout_secs)

    local buf, view = create_python_view(input.code, ctx)

    local output_parts = {}
    local has_output = false
    local finished = false

    local function append_output(line)
      if not has_output then
        has_output = true
        view:clear()
      end
      if #output_parts > 0 then
        output_parts[#output_parts + 1] = "\n"
      end
      output_parts[#output_parts + 1] = line
      view:append(line)
    end

    local function finish(exit_code)
      if finished then
        return
      end
      finished = true
      local output = table.concat(output_parts)
      output = truncate(output, max_lines, max_bytes)

      local is_error = exit_code ~= 0
      local llm_output
      if exit_code == 0 then
        llm_output = output == "" and "Exit code: 0" or output
      else
        if output == "" then
          llm_output = "Exit code: " .. exit_code
        else
          llm_output = output .. "\nExit code: " .. exit_code
        end
      end

      if output == "" then
        view:clear()
        view:append({ { "No output", "dim" } })
      end

      if is_error then
        view:append({ { "Exit code: " .. exit_code, "dim" } })
      end
      view:finish()

      ctx:finish({ llm_output = llm_output, is_error = is_error, body = buf })
    end

    view:append({ { "Waiting for output...", "dim" } })

    maki.fn.jobstart({ "python3", "-c", input.code }, {
      cwd = input.workdir,
      env = build_env(ctx),
      on_stdout = function(_, line)
        append_output(line)
      end,
      on_stderr = function(_, line)
        append_output(line)
      end,
      on_exit = function(_, code)
        finish(code)
      end,
    })

    maki.async.on_cancel(function(reason)
      if finished then
        return
      end
      finished = true
      local out = truncate(table.concat(output_parts), max_lines, max_bytes)
      ctx:finish(partial.cut(view, out, reason, timeout_secs))
    end)

    return nil
  end,
})