path: root/subagents/luatool.lua

-- subagents/luatool.lua
--
-- The `subagents.lua` model-facing tool: run a transient, model-authored Lua
-- workflow without writing a definition to disk. The source must evaluate to
-- `subagents.workflow(function(ctx, input) ... end)`; the tool then executes it
-- with the tool's `prompt` as the workflow input and formats the terminal
-- results for the calling model.
--
-- The source is loaded in text mode only (`load(source, chunkname, "t", env)`)
-- against a restricted `_ENV`. That environment holds a safe slice of the
-- standard library plus `subagents.workflow`; it has no `os`, `io`, `debug`,
-- `package`, `require`, `load`, `dofile`, `coroutine`, `setmetatable`, or
-- `getmetatable`, and `print` is a no-op so generated code cannot scribble on
-- the TUI. `string`, `table`, `math`, and `utf8` are shallow copies, so a guest
-- that reassigns `table.insert` only breaks itself, and `string.dump` is
-- removed from the copy.
--
-- `pcall`/`xpcall` are deliberately absent: they would let a guest catch the
-- instruction-budget error and spin again in fresh 10M-instruction chunks. A
-- guest has no need to recover from its own errors — the handler reports them —
-- and with no `pcall`, `coroutine`, or metatable access left, nothing in the
-- environment can trap the hook error before it reaches the host.
--
-- Known, accepted gaps in that sandbox:
--
-- * The real string metatable is still reachable through any string literal
--   (`("").dump`), so `string.dump` is obtainable. Without `load` there is no
--   way to turn bytecode back into a running function, so this is noise rather
--   than an escape.
-- * `ctx:agent` returns handles the guest can read and scribble on. Nothing
--   reachable from one starts work: the running job — which owns the child's
--   agent and could start turns outside the cap — lives in a private side table
--   in workflow.lua, as do the job cap, the counter, and the profile set, so a
--   guest holding `ctx` and its handles cannot raise its own cap or reach the
--   host.
--
-- Runaway generated Lua is bounded two ways: `max_jobs = 32` caps how many
-- children one transient workflow may start, and a debug count hook is armed
-- for the duration of the guest callback and disarmed as soon as it returns.
-- The hook lives here rather than in workflow.lua because the guest has no
-- `debug` library but the host does; workflow.execute only exposes the
-- on_resume/on_yield seam the hook needs. The guest runs on the tool handler's
-- own coroutine (an await parks and resumes exactly that coroutine when a child
-- settles), so the budget covers the whole run rather than one slice; awaiting
-- a child executes no instructions, so only real spinning trips it.

local workflow = require("subagents.workflow")
local run = require("subagents.run")

local MAX_JOBS = 32
local INSTRUCTION_BUDGET = 10000000
local CHUNK_NAME = "subagents.lua"

local M = {}

M.max_jobs = MAX_JOBS

local function shallow_copy(source, skip)
    local copy = {}
    for key, value in pairs(source) do
        if key ~= skip then
            copy[key] = value
        end
    end
    return copy
end

-- Build a fresh restricted environment per call: the guest may mutate anything
-- it can reach, so nothing here is shared between invocations.
local function build_env()
    local env = {
        assert = assert,
        error = error,
        ipairs = ipairs,
        next = next,
        pairs = pairs,
        select = select,
        tonumber = tonumber,
        tostring = tostring,
        type = type,
        string = shallow_copy(string, "dump"),
        table = shallow_copy(table),
        math = shallow_copy(math),
        utf8 = shallow_copy(utf8),
        print = function() end,
        subagents = { workflow = workflow.workflow },
    }
    env._G = env
    return env
end

M.build_env = build_env

local function budget_hook()
    error("instruction budget exceeded", 0)
end

-- subagents.run's block formatter expects string output; a structured worker's
-- output is a decoded table, so it is re-encoded first.
local function format_one(result)
    if type(result.output) == "table" then
        local flattened = {}
        for key, value in pairs(result) do
            flattened[key] = value
        end
        flattened.output = workflow.output_text(result)
        return run.format_result(flattened)
    end
    return run.format_result(result)
end

-- Format whatever the workflow callback returned. Result-shaped tables (the
-- common case: one settled result, or an array of them) render as the same
-- plain "key: value" block subagents.run uses; anything else is encoded
-- compactly so the model still sees it.
local function format_return(value)
    if value == nil then
        return "The workflow returned no value."
    end
    if type(value) ~= "table" then
        return tostring(value)
    end
    if value.status ~= nil then
        return format_one(value)
    end
    local blocks, count = {}, 0
    for index, entry in ipairs(value) do
        if type(entry) ~= "table" or entry.status == nil then
            blocks = nil
            break
        end
        blocks[index] = format_one(entry)
        count = count + 1
    end
    if blocks and count > 0 then
        return table.concat(blocks, "\n\n")
    end
    return workflow.json_encode(value)
end

-- Tool handler for `subagents.lua`. `profiles` is the discovered profile set
-- from activation; when omitted the workflow API discovers it lazily.
function M.handle(input, profiles)
    if type(input) ~= "table" then
        return "Error: expected an input object"
    end
    if type(input.prompt) ~= "string" or input.prompt == "" then
        return "Error: prompt is required and must be a non-empty string"
    end
    if type(input.source) ~= "string" or input.source == "" then
        return "Error: source is required and must be a non-empty string"
    end

    local chunk, load_err = load(input.source, CHUNK_NAME, "t", build_env())
    if not chunk then
        return "Error: source did not compile: " .. tostring(load_err)
    end

    local built_ok, built = pcall(chunk)
    if not built_ok then
        return "Error: source failed to run: " .. tostring(built)
    end
    if not workflow.is_workflow(built) then
        return "Error: source must return subagents.workflow(function(ctx, input) ... end)"
    end

    local armed = nil
    local ran_ok, result = pcall(workflow.execute, built, input.prompt, {
        max_jobs = MAX_JOBS,
        profiles = profiles,
        on_resume = function(co)
            armed = co
            debug.sethook(co, budget_hook, "", INSTRUCTION_BUDGET)
        end,
        on_yield = function(co)
            debug.sethook(co)
            armed = nil
        end,
    })
    if armed then
        debug.sethook(armed)
    end

    if not ran_ok then
        return "Error: " .. tostring(result)
    end
    return format_return(result)
end

return M