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. Source can start one with
-- `subagents.workflow(function(ctx) ... end)`, which immediately returns a
-- workflow id, or inspect a prior run through `subagents.workflows[id]` and
-- return any model-visible value. Optional inline agent profiles are overlaid
-- for workflows started by this call only; they are never persisted or added
-- to discovery.
--
-- 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, a `json` codec, and `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
-- on each background workflow coroutine for its full execution. Top-level
-- source evaluation gets the same budget before it can schedule anything.

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

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

-- JSON for the guest. `decode` returns `nil, message` instead of raising:
-- the sandbox has no `pcall`, so a raising decoder would make one malformed
-- agent line kill the whole workflow.
local function guest_json()
    return {
        decode = function(text)
            if type(text) ~= "string" then
                return nil, "expected a string"
            end
            local ok, value = pcall(workflow.json_decode, text)
            if not ok then
                return nil, tostring(value)
            end
            return value
        end,
        encode = function(value)
            return workflow.json_encode(value)
        end,
    }
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(schedule)
    schedule = schedule or workflow.workflow
    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,
        json = guest_json(),
        subagents = {
            workflow = schedule,
            workflows = workflow.workflows,
        },
    }
    env._G = env
    return env
end

M.build_env = build_env
M.guest_json = guest_json

-- Overlay workflow-local profiles without mutating the activation-time set.
-- Inline names intentionally win, matching normal layered profile precedence
-- while keeping their lifetime bounded to this one execute() call.
local function workflow_profiles(inline, discovered)
    if inline == nil then
        return discovered, nil
    end
    if type(inline) ~= "table" then
        return nil, "agents must be an array"
    end

    local count = 0
    for key in pairs(inline) do
        if type(key) ~= "number" then
            return nil, "agents must be an array"
        end
        count = count + 1
    end
    if count ~= #inline then
        return nil, "agents must be a dense array"
    end

    local base = spawn.profiles(discovered)
    local by_name = {}
    for name, profile in pairs(base.by_name or {}) do
        by_name[name] = profile
    end
    local seen = {}
    for index, raw in ipairs(inline) do
        if type(raw) ~= "table" then
            return nil, string.format("agents[%d] must be an object", index)
        end
        local name = raw.name
        if type(name) ~= "string" or name:match("^%s*$") then
            return nil, string.format("agents[%d].name must be a non-empty string", index)
        end
        if seen[name] then
            return nil, string.format("duplicate inline agent '%s'", name)
        end
        local system_prompt = raw.system_prompt
        if type(system_prompt) ~= "string" or system_prompt:match("^%s*$") then
            return nil, string.format("agents[%d].system_prompt must be a non-empty string", index)
        end
        if raw.description ~= nil and type(raw.description) ~= "string" then
            return nil, string.format("agents[%d].description must be a string when given", index)
        end
        seen[name] = true
        by_name[name] = {
            name = name,
            description = raw.description or "",
            body = system_prompt,
            layer = "workflow",
        }
    end

    local list = {}
    for _, profile in pairs(by_name) do
        list[#list + 1] = profile
    end
    table.sort(list, function(a, b) return a.name < b.name end)
    return { list = list, by_name = by_name, warnings = base.warnings or {} }, nil
end

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, context)
    if type(input) ~= "table" then
        return "Error: expected an input object"
    end
    if type(input.source) ~= "string" or input.source == "" then
        return "Error: source is required and must be a non-empty string"
    end

    local profiles_for_run, profiles_err = workflow_profiles(input.agents, profiles)
    if not profiles_for_run then
        return "Error: " .. profiles_err
    end

    local function on_resume(co)
        debug.sethook(co, budget_hook, "", INSTRUCTION_BUDGET)
    end
    local function on_yield(co)
        debug.sethook(co)
    end
    local function schedule(fn)
        return workflow.start(workflow.workflow(fn), {
            max_jobs = MAX_JOBS,
            profiles = profiles_for_run,
            tool_call_id = type(context) == "table" and context.tool_call_id or nil,
            on_resume = on_resume,
            on_yield = on_yield,
        })
    end

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

    local co = coroutine.running()
    if co then debug.sethook(co, budget_hook, "", INSTRUCTION_BUDGET) end
    local built_ok, built = pcall(chunk)
    if co then debug.sethook(co) end
    if not built_ok then
        return "Error: source failed to run: " .. tostring(built)
    end
    return format_return(built)
end

return M