path: root/subagents/jobs.lua

-- Start child agent jobs, bound how many run at once, and drain their events.
--
-- One module owns the session-wide concurrency bound and the event pump, so
-- the policy that decides *what* to start (subagents/spawn.lua) never touches
-- luv, and the callers that wait for a result (subagents/run.lua, the workflow
-- API) never touch a job.
--
-- Threading and ownership: everything here runs on panto's Lua owner thread.
-- A job's pump runs on its own thread inside the binding; it buffers events on
-- the job and writes one byte to the wake pipe handed to it here. That byte is
-- an edge trigger, never a count, so every wake drains the pipe, then drains
-- `job:next_event()` to exhaustion, then checks `job:result()`.
--
-- Waiting parks the CALLING coroutine — the tool handler's — and a poll
-- callback resumes exactly that coroutine. subagents/workflow.lua's rule that
-- a workflow callback must never run on a nested coroutine follows from this:
-- the parked thread is the one the resume goes to. Where there is no coroutine
-- to park (a plain script) or a job with no wake pipe, the identical drain
-- runs in a loop instead — same gate, same settle bookkeeping, only the wait
-- differs.
--
-- A settled result is read and cached the moment it appears, because
-- `job:close()` frees it. Jobs are otherwise left open until close_all() ends
-- the turn, so a result can still be read after its coroutine resumed.

local ok_uv, uv = pcall(require, "luv")

local READ_CHUNK = 4096

local M = {}

-- The session-wide bound. A field, not a constant, so a caller (or a spec) can
-- lower it without reaching into the queue.
M.MAX_CONCURRENT = 4

local handle_mt = {}
handle_mt.__index = handle_mt
handle_mt.__name = "subagents.job"

-- Every handle that has not been closed, in start order; the queue is the
-- subset waiting for a slot; `running` counts the started-but-unsettled ones.
local live = {}
local queued = {}
local waiters = {}
local running = 0
local pumping = false

-- ---------------------------------------------------------------------------
-- Wake pipes
-- ---------------------------------------------------------------------------

local function open_pipe()
    if not ok_uv or type(uv.pipe) ~= "function" then
        return nil
    end
    local ok, pair = pcall(uv.pipe, { nonblock = true }, { nonblock = true })
    if not ok or type(pair) ~= "table" then
        return nil
    end
    return pair
end

-- Close the poll before the fds: the pump has already exited by the time a job
-- settles, so nothing can be mid-write when the read end goes away.
local function close_pipe(handle)
    if handle.poll then
        pcall(handle.poll.stop, handle.poll)
        pcall(handle.poll.close, handle.poll)
        handle.poll = nil
    end
    if handle.fds then
        pcall(uv.fs_close, handle.fds.read)
        pcall(uv.fs_close, handle.fds.write)
        handle.fds = nil
    end
end

local function drain_pipe(handle)
    if not handle.fds then
        return
    end
    while true do
        local data = uv.fs_read(handle.fds.read, READ_CHUNK, -1)
        if type(data) ~= "string" or #data < READ_CHUNK then
            return
        end
    end
end

-- ---------------------------------------------------------------------------
-- Gate, drain, settle
-- ---------------------------------------------------------------------------

local drain
local pump_queue
local wake

local function settle(handle, raw)
    if handle.settled ~= nil then
        return
    end
    local result = raw
    if handle.spec.shape then
        local ok, shaped = pcall(handle.spec.shape, raw)
        if ok and type(shaped) == "table" then
            result = shaped
        elseif not ok then
            result = { status = "failed", error = tostring(shaped), resumable = false }
        end
    end
    handle.settled = result
    if handle.state == "running" then
        running = running - 1
    end
    handle.state = "settled"
    close_pipe(handle)
    pump_queue()
end

local function launch(handle)
    handle.fds = open_pipe()
    local ok, job, err = pcall(handle.spec.build, handle.fds and handle.fds.write or nil)
    if not ok then
        close_pipe(handle)
        return false, tostring(job)
    end
    if not job then
        close_pipe(handle)
        return false, err and tostring(err) or "the child could not be started"
    end

    handle.job = job
    handle.state = "running"
    running = running + 1

    if handle.fds and ok_uv then
        local armed, poll = pcall(uv.new_poll, handle.fds.read)
        if armed and poll then
            handle.poll = poll
            poll:start("r", function()
                drain(handle)
                wake()
            end)
        end
    end
    return true
end

-- Start queued jobs while the gate has room. Reentrant: a job that settles the
-- instant it starts calls back in here, and the outer loop keeps going.
function pump_queue()
    if pumping then
        return
    end
    pumping = true
    while running < M.MAX_CONCURRENT do
        local handle = table.remove(queued, 1)
        if handle == nil then
            break
        end
        if handle.state == "queued" then
            local ok, err = launch(handle)
            if not ok then
                settle(handle, { status = "failed", error = tostring(err) })
            end
        end
    end
    pumping = false
end

-- Drain one job: the wake pipe, then every buffered event, then the result.
-- Returns true when anything moved, which is what the fallback wait loop uses
-- to decide whether it is spinning.
function drain(handle)
    local job = handle.job
    if job == nil or handle.settled ~= nil then
        return false
    end
    drain_pipe(handle)

    local moved = false
    local on_event = handle.spec.on_event
    local event = job:next_event()
    while event ~= nil do
        moved = true
        if on_event then
            pcall(on_event, event)
        end
        event = job:next_event()
    end

    local raw = job:result()
    if raw ~= nil then
        settle(handle, raw)
        return true
    end
    return moved
end

local function drain_all()
    local moved = false
    local index = 1
    while index <= #live do
        if drain(live[index]) then
            moved = true
        end
        index = index + 1
    end
    return moved
end

-- ---------------------------------------------------------------------------
-- Handles
-- ---------------------------------------------------------------------------

function handle_mt:result()
    return self.settled
end

-- A queued child settles without its build ever running: no user message is
-- appended, so nothing on disk claims a turn that never happened.
function handle_mt:cancel()
    if self.settled ~= nil then
        return
    end
    if self.state == "queued" then
        for index, other in ipairs(queued) do
            if other == self then
                table.remove(queued, index)
                break
            end
        end
        settle(self, { status = "cancelled", error = "cancelled before the child started" })
        return
    end
    if self.job then
        pcall(self.job.request_cancel, self.job)
    end
end

-- start(startspec) -> handle | nil, err
--
-- startspec = {
--   build     = function(wake_fd) -> job | nil, err   -- calls agent:run_async
--   label     = string?,  id = string?,  one_shot = boolean?
--   on_event  = function(event)?   -- one call per drained run_async event
--   shape     = function(raw) -> result?  -- maps the settled run_async result
--                                            onto the caller's result table
-- }
--
-- Over the bound the job is queued and `build` is not called yet.
function M.start(spec)
    if type(spec) ~= "table" or type(spec.build) ~= "function" then
        return nil, "jobs.start needs a build function"
    end

    local handle = setmetatable({
        spec = spec,
        label = spec.label,
        id = spec.id,
        one_shot = spec.one_shot == true,
        state = "queued",
    }, handle_mt)
    live[#live + 1] = handle

    if running >= M.MAX_CONCURRENT then
        queued[#queued + 1] = handle
        return handle
    end

    local ok, err = launch(handle)
    if not ok then
        for index, other in ipairs(live) do
            if other == handle then
                table.remove(live, index)
                break
            end
        end
        return nil, err
    end
    -- Deliberately not drained here: starting a child must not settle it, and
    -- the first wake byte is already in the pipe by the time anyone waits.
    return handle
end

-- True while a child with this id has a turn in flight.
function M.active(id)
    if id == nil then
        return false
    end
    for _, handle in ipairs(live) do
        if handle.id == id and handle.settled == nil then
            return true
        end
    end
    return false
end

-- ---------------------------------------------------------------------------
-- Awaiting
-- ---------------------------------------------------------------------------

-- "all" is satisfied only when every handle has settled; "first" as soon as one
-- has. Both answer in input order, and "first" hands the rest back by identity
-- so a caller can await them again.
local function collect(handles, mode)
    if mode == "first" then
        local results, remaining = {}, {}
        for _, handle in ipairs(handles) do
            if handle.settled ~= nil then
                results[#results + 1] = handle.settled
            else
                remaining[#remaining + 1] = handle
            end
        end
        if #results == 0 and #remaining > 0 then
            return nil
        end
        return results, remaining
    end

    local results = {}
    for index, handle in ipairs(handles) do
        if handle.settled == nil then
            return nil
        end
        results[index] = handle.settled
    end
    return results, {}
end

-- Resume every parked await whose condition now holds. Never resumes the
-- running coroutine: a coroutine that is executing cannot also be parked.
function wake()
    local self_co = coroutine.running()
    local index = 1
    while index <= #waiters do
        local waiter = waiters[index]
        local results, remaining = collect(waiter.handles, waiter.mode)
        if results ~= nil and waiter.co ~= self_co and coroutine.status(waiter.co) == "suspended" then
            table.remove(waiters, index)
            coroutine.resume(waiter.co, results, remaining)
        else
            index = index + 1
        end
    end
end

-- Parking is only safe when something will wake us: a started job with no wake
-- pipe is drained by asking it directly instead.
local function can_park(handles)
    if not coroutine.isyieldable() then
        return false
    end
    for _, handle in ipairs(handles) do
        if handle.state == "running" and handle.poll == nil then
            return false
        end
    end
    return true
end

-- await(handles, mode) -> results, remaining
function M.await(handles, mode)
    mode = mode or "all"
    if type(handles) ~= "table" then
        error("jobs.await expects an array of job handles", 2)
    end
    if mode ~= "all" and mode ~= "first" then
        error("jobs.await mode must be \"all\" or \"first\"", 2)
    end

    while true do
        local moved = drain_all()
        local results, remaining = collect(handles, mode)
        if results ~= nil then
            return results, remaining
        end
        wake()

        if can_park(handles) then
            waiters[#waiters + 1] = { co = coroutine.running(), handles = handles, mode = mode }
            local woken, rest = coroutine.yield()
            if woken ~= nil then
                return woken, rest
            end
        elseif not moved and ok_uv then
            -- Nothing to drain and nothing to park on: yield the core rather
            -- than burn it while the pump threads work.
            uv.sleep(1)
        end
    end
end

-- ---------------------------------------------------------------------------
-- Turn lifecycle
-- ---------------------------------------------------------------------------

-- The turn was interrupted: ask every child to stop. Cancellation is a request,
-- not a settle — each job still reports its own cancelled result.
function M.cancel_all()
    for index = #live, 1, -1 do
        live[index]:cancel()
    end
end

-- The turn is over: join every pump and drop the state. Requests go out first
-- so the joins overlap instead of running one child's teardown at a time.
function M.close_all()
    for _, handle in ipairs(live) do
        if handle.job and handle.settled == nil then
            pcall(handle.job.request_cancel, handle.job)
        end
    end
    for _, handle in ipairs(live) do
        if handle.job then
            pcall(handle.job.close, handle.job)
            handle.job = nil
        end
        close_pipe(handle)
        handle.state = "closed"
    end
    live, queued, waiters = {}, {}, {}
    running = 0
end

return M