rework lua tool scheduler

1 files changed, 105 insertions, 74 deletions

diff --git a/agent/tools/shell.lua b/agent/tools/shell.lua
index a9e9de3..815a8b9 100644
--- a/agent/tools/shell.lua
+++ b/agent/tools/shell.lua
@@ -33,6 +33,42 @@ local DEFAULT_TIMEOUT_S = 300   -- 5 minutes
 local SIGKILL_GRACE_MS = 250
 
 -- ---------------------------------------------------------------------------
+-- Coroutine-synchronous coordination.
+--
+-- panto runs every tool handler in its own coroutine and drives a
+-- single `uv.run()` to completion; the only legal way for a handler to
+-- block is to yield on pending libuv work whose callback(s) resume it.
+-- The shell tool waits on THREE independent events before it can build
+-- its result: stdout EOF, stderr EOF, and process exit. `await_n`
+-- captures that "resume once, after N signals" pattern.
+--
+-- `arm(signal)` is called immediately with a `signal` function; it must
+-- arrange for `signal` to be called exactly `n` times across its libuv
+-- callbacks. After the n-th signal, the coroutine is resumed exactly
+-- once. `signal` is safe to call from any libuv callback.
+local function await_n(n, arm)
+    local co = assert(coroutine.running(),
+        "shell: await_n must run inside a tool handler coroutine")
+    local remaining, resumed = n, false
+    local function signal()
+        if remaining <= 0 then return end
+        remaining = remaining - 1
+        if remaining > 0 or resumed then return end
+        resumed = true
+        local ok, err = coroutine.resume(co)
+        if not ok then
+            -- This runs inside a libuv callback, not the handler's
+            -- own frame; re-raising here would unwind the event loop.
+            -- Log instead; the scheduler surfaces the unfinished
+            -- coroutine as a tool error on reap.
+            io.stderr:write("shell tool: resume failed: " .. tostring(err) .. "\n")
+        end
+    end
+    arm(signal)
+    coroutine.yield()
+end
+
+-- ---------------------------------------------------------------------------
 -- Spill-file path resolution. We mirror Zig's `panto_home.resolveHome`:
 -- PANTO_HOME wins, then XDG_DATA_HOME/panto, then HOME/.local/share/panto.
 -- The shell-output subdirectory is created lazily.
@@ -117,8 +153,7 @@ return {
         local timeout_s = input.timeout or DEFAULT_TIMEOUT_S
         local timeout_ms = timeout_s * 1000
 
-        local co = coroutine.running()
-        if not co then
+        if not coroutine.running() then
             -- Should never happen — panto always runs handlers in a
             -- coroutine — but be defensive rather than silently hang.
             return "Error: shell tool requires a coroutine scheduler " ..
@@ -148,6 +183,14 @@ return {
         local spill_error = nil    -- if opening/writing the spill file fails
         local total_written_to_spill = 0
 
+        -- NOTE: the spill helpers below use SYNCHRONOUS `uv.fs_*`
+        -- calls on purpose. They run inside `append_chunk`, which runs
+        -- inside the stdout/stderr read callbacks — i.e. in event-loop
+        -- context, NOT in the handler coroutine's own frame. You cannot
+        -- `coroutine.yield` (and thus cannot `await`) from a libuv
+        -- callback, so these must be the blocking variants. The brief
+        -- per-write stall is acceptable; spilling only happens once
+        -- output exceeds the in-memory cap.
         local function open_spill_and_seed()
             -- Opens the spill file and writes every chunk currently
             -- live in the FIFO. Called exactly once, the moment we
@@ -241,15 +284,10 @@ return {
             total_bytes = total_bytes + #data
         end
 
-        -- Coordination state. We resume the coroutine exactly once,
-        -- after `pending_events` hits zero: stdout EOF, stderr EOF,
-        -- on_exit fired = 3 events.
-        local pending_events = 3
-        local resumed = false
+        -- Per-call outcome state, filled by the libuv callbacks below.
         local exit_code = nil
         local exit_signal = nil
         local timed_out = false
-        local spawn_error = nil
         local process_handle = nil
         local timeout_timer = nil
         local kill_timer = nil
@@ -264,97 +302,90 @@ return {
             return table.concat(out)
         end
 
-        local function try_resume()
-            if pending_events > 0 then return end
-            if resumed then return end
-            resumed = true
-            -- Cancel any still-pending timers. They may have been
-            -- nil'd out already inside their own callback.
-            if timeout_timer and not timeout_timer:is_closing() then
-                timeout_timer:stop(); timeout_timer:close()
-            end
-            if kill_timer and not kill_timer:is_closing() then
-                kill_timer:stop(); kill_timer:close()
-            end
-            local ok, err = coroutine.resume(co)
-            if not ok then
-                -- The handler shouldn't be raising here; this is the
-                -- resume scheduler, not user code. Log via stderr;
-                -- the parent will surface the failure as a tool error.
-                io.stderr:write("shell tool: resume failed: " .. tostring(err) .. "\n")
-            end
-        end
-
-        local function on_read(err, data)
-            if err then
-                -- Treat read errors as EOF — luv signals EOF by
-                -- calling with `data == nil`, so an actual `err` means
-                -- something genuinely broke. Record it once.
-                append_chunk("[read error: " .. tostring(err) .. "]\n")
-                pending_events = pending_events - 1
-                try_resume()
-                return
-            end
-            if data == nil then
-                -- EOF on this pipe.
-                pending_events = pending_events - 1
-                try_resume()
-                return
-            end
-            append_chunk(data)
-        end
-
+        -- Spawn first so a spawn failure can return *before* we commit
+        -- to awaiting (a failed spawn fires no events to resume us).
         local stdio = { nil, stdout_pipe, stderr_pipe }
         local opts = { args = { "-c", command }, stdio = stdio }
         if cwd then opts.cwd = cwd end
 
-        local handle, pid_or_err, name = uv.spawn("/bin/sh", opts, function(code, signal)
+        local on_exit_signal   -- set by await_n before any event fires
+        local handle, pid_or_err, name = uv.spawn("/bin/sh", opts, function(code, sig)
             exit_code = code
-            exit_signal = signal
-            pending_events = pending_events - 1
-            -- The on-exit callback is responsible for closing the
-            -- process handle, per luv convention.
+            exit_signal = sig
+            -- The on-exit callback owns closing the process handle.
             if process_handle and not process_handle:is_closing() then
                 process_handle:close()
             end
-            try_resume()
+            on_exit_signal()
         end)
 
         if not handle then
-            -- Spawn failed before we ever started — `pid_or_err` is
-            -- the error message. Clean up pipes and return.
-            spawn_error = tostring(pid_or_err) ..
+            -- Spawn failed before any event could fire. Clean up the
+            -- pipes and return without awaiting.
+            local spawn_error = tostring(pid_or_err) ..
                 (name and (" (" .. name .. ")") or "")
             stdout_pipe:close(); stderr_pipe:close()
             return "Error: failed to spawn shell: " .. spawn_error
         end
         process_handle = handle
 
-        stdout_pipe:read_start(on_read)
-        stderr_pipe:read_start(on_read)
+        -- `await_n(3, ...)` resumes this coroutine exactly once, after
+        -- all three of {stdout EOF, stderr EOF, process exit} signal.
+        await_n(3, function(signal)
+            -- The exit callback was armed above (before this awaiter
+            -- existed); wire its signal now.
+            on_exit_signal = signal
 
-        -- Timeout timer. If it fires, we send SIGTERM and arm a
-        -- second short timer to follow up with SIGKILL.
-        timeout_timer = uv.new_timer()
-        timeout_timer:start(timeout_ms, 0, function()
-            timed_out = true
-            if process_handle and not process_handle:is_closing() then
-                process_handle:kill("sigterm")
+            -- stdout/stderr each signal once, on EOF or read error.
+            -- luv delivers stream EOF as a nil `data` with no error.
+            local function on_read(err, data)
+                if err then
+                    append_chunk("[read error: " .. tostring(err) .. "]\n")
+                    signal()
+                elseif data == nil then
+                    signal()
+                else
+                    append_chunk(data)
+                end
             end
-            timeout_timer:stop(); timeout_timer:close(); timeout_timer = nil
-            kill_timer = uv.new_timer()
-            kill_timer:start(SIGKILL_GRACE_MS, 0, function()
+            stdout_pipe:read_start(on_read)
+            stderr_pipe:read_start(on_read)
+
+            -- Timeout timer. If it fires, send SIGTERM and arm a short
+            -- follow-up timer for SIGKILL. These timers do NOT signal;
+            -- they cause the child to exit, which fires the exit/EOF
+            -- callbacks that do. We close them here so they stop
+            -- holding the event loop open.
+            timeout_timer = uv.new_timer()
+            timeout_timer:start(timeout_ms, 0, function()
+                timed_out = true
                 if process_handle and not process_handle:is_closing() then
-                    process_handle:kill("sigkill")
+                    process_handle:kill("sigterm")
                 end
-                kill_timer:stop(); kill_timer:close(); kill_timer = nil
+                timeout_timer:stop(); timeout_timer:close(); timeout_timer = nil
+                kill_timer = uv.new_timer()
+                kill_timer:start(SIGKILL_GRACE_MS, 0, function()
+                    if process_handle and not process_handle:is_closing() then
+                        process_handle:kill("sigkill")
+                    end
+                    kill_timer:stop(); kill_timer:close(); kill_timer = nil
+                end)
             end)
         end)
 
-        -- Wait for stdout EOF, stderr EOF, and on_exit.
-        coroutine.yield()
+        if spawn_error then
+            return "Error: failed to spawn shell: " .. spawn_error
+        end
 
-        -- Close pipes (idempotent if already closing).
+        -- Resumed: all three events fired. Cancel any still-pending
+        -- timers (the command may have finished before the timeout) so
+        -- they don't keep the loop alive, then close pipes/spill.
+        if timeout_timer and not timeout_timer:is_closing() then
+            timeout_timer:stop(); timeout_timer:close()
+        end
+        if kill_timer and not kill_timer:is_closing() then
+            kill_timer:stop(); kill_timer:close()
+        end
         if not stdout_pipe:is_closing() then stdout_pipe:close() end
         if not stderr_pipe:is_closing() then stderr_pipe:close() end
         if spill_handle then uv.fs_close(spill_handle) end