diff options
Diffstat (limited to 'subagents/jobs.lua')
| -rw-r--r-- | subagents/jobs.lua | 427 |
1 files changed, 427 insertions, 0 deletions
diff --git a/subagents/jobs.lua b/subagents/jobs.lua new file mode 100644 index 0000000..a8cde6d --- /dev/null +++ b/subagents/jobs.lua @@ -0,0 +1,427 @@ +-- 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 |