rework lua tool scheduler

4 files changed, 334 insertions, 114 deletions

diff --git a/agent/tools/edit.lua b/agent/tools/edit.lua
index 7ad39c2..ec96aa6 100644
--- a/agent/tools/edit.lua
+++ b/agent/tools/edit.lua
@@ -10,6 +10,82 @@
 -- All matches are evaluated against the original file, not
 -- progressively against the result of earlier entries.
 
+local uv = require("luv")
+
+-- ---------------------------------------------------------------------------
+-- Coroutine-synchronous libuv wrappers. panto runs each tool handler in
+-- its own coroutine and drives a single `uv.run()` to completion; a
+-- handler may only block by yielding on a pending libuv op whose
+-- callback resumes it exactly once. `await` encapsulates that: arm one
+-- op, yield, return the callback's `(err, value)` results.
+-- ---------------------------------------------------------------------------
+local function await(arm)
+    local co = assert(coroutine.running(),
+        "edit: await must run inside a tool handler coroutine")
+    local res, fired = nil, false
+    arm(function(...)
+        assert(not fired, "edit: libuv callback fired twice")
+        fired = true
+        res = table.pack(...)
+        local ok, err = coroutine.resume(co)
+        if not ok then error(err, 0) end
+    end)
+    coroutine.yield()
+    return table.unpack(res, 1, res.n)
+end
+
+local function fs_open(path, flags, mode)
+    return await(function(resolve) uv.fs_open(path, flags, mode, resolve) end)
+end
+local function fs_fstat(fd)
+    return await(function(resolve) uv.fs_fstat(fd, resolve) end)
+end
+local function fs_read(fd, size, offset)
+    return await(function(resolve) uv.fs_read(fd, size, offset, resolve) end)
+end
+local function fs_write(fd, data, offset)
+    return await(function(resolve) uv.fs_write(fd, data, offset, resolve) end)
+end
+local function fs_close(fd)
+    return await(function(resolve) uv.fs_close(fd, resolve) end)
+end
+
+-- Read an entire file into a string. Returns (content, nil) or
+-- (nil, err). EOF on `fs_read` is the empty string, not nil.
+local function read_all(path)
+    local err, fd = fs_open(path, "r", 0)
+    if not fd then return nil, err end
+    local serr, st = fs_fstat(fd)
+    if not st then fs_close(fd); return nil, serr end
+    local parts, offset = {}, 0
+    while true do
+        local rerr, data = fs_read(fd, 64 * 1024, offset)
+        if rerr then fs_close(fd); return nil, rerr end
+        if data == nil or #data == 0 then break end
+        parts[#parts + 1] = data
+        offset = offset + #data
+    end
+    fs_close(fd)
+    return table.concat(parts), nil
+end
+
+-- Overwrite a file with `content` (O_WRONLY|O_CREAT|O_TRUNC, 0644).
+-- Returns (true, nil) or (nil, err).
+local function write_all(path, content)
+    local err, fd = fs_open(path, "w", tonumber("644", 8))
+    if not fd then return nil, err end
+    -- A single `fs_write` may short-write; loop until all bytes land.
+    local offset = 0
+    while offset < #content do
+        local werr, n = fs_write(fd, content:sub(offset + 1), offset)
+        if not n then fs_close(fd); return nil, werr end
+        offset = offset + n
+    end
+    -- Empty content still needs the truncate that O_TRUNC gave us.
+    fs_close(fd)
+    return true, nil
+end
+
 return {
     name = "std.edit",
     description = "Apply exact-text replacements to a file. Each entry's `old` must match exactly once in the original file (not progressively). Edits whose ranges overlap, match zero times, or match multiple times reject the entire call — no partial edits, file untouched. Keep `old` minimal but unique; widen with surrounding context if needed.",
@@ -45,14 +121,9 @@ return {
         end
 
         -- Load the file.
-        local f, open_err = io.open(path, "rb")
-        if not f then
-            return "Error: " .. (open_err or ("could not open " .. path))
-        end
-        local original = f:read("a")
-        f:close()
+        local original, read_err = read_all(path)
         if not original then
-            return "Error: failed to read " .. path
+            return "Error: " .. (read_err or ("could not read " .. path))
         end
 
         -- Validate every entry first, collecting the byte-range each
@@ -170,19 +241,10 @@ return {
         local new_content = table.concat(out)
 
         -- Write back.
-        local wf, werr = io.open(path, "wb")
-        if not wf then
-            return "Error: file read OK but failed to reopen for write: " ..
-                (werr or "?")
-        end
-        local ok, write_err = wf:write(new_content)
+        local ok, write_err = write_all(path, new_content)
         if not ok then
-            wf:close()
-            return "Error: write failed: " .. tostring(write_err)
-        end
-        local close_ok, close_err = wf:close()
-        if not close_ok then
-            return "Error: close failed: " .. tostring(close_err)
+            return "Error: file read OK but write failed: " ..
+                tostring(write_err)
         end
 
         return string.format(
diff --git a/agent/tools/read.lua b/agent/tools/read.lua
index d345f99..f232890 100644
--- a/agent/tools/read.lua
+++ b/agent/tools/read.lua
@@ -19,18 +19,98 @@ local MAX_BYTES = 50 * 1024   -- 50 KB
 local MAX_LINES = 2000
 local READ_CHUNK_SIZE = MAX_BYTES
 
+-- ---------------------------------------------------------------------------
+-- Coroutine-synchronous libuv wrappers.
+--
+-- panto runs every tool handler inside its own Lua coroutine and drives
+-- a single `uv.run()` to completion. The contract for an async handler
+-- is: the ONLY way it may block is to yield on a pending libuv request
+-- whose callback resumes it exactly once. These helpers encapsulate
+-- that pattern — each fires one async libuv op, yields, and returns the
+-- callback's results once resumed. They look synchronous to the caller
+-- but cooperate with sibling tool calls sharing the event loop.
+--
+-- INVARIANTS each wrapper upholds (and every async tool must too):
+--   * exactly one libuv op armed per yield;
+--   * the callback resumes the coroutine exactly once;
+--   * `coroutine.resume` errors are re-raised so a crash in the resumed
+--     handler isn't silently swallowed inside the uv callback.
+-- ---------------------------------------------------------------------------
+
+-- Run a single callback-style libuv fs operation and block the
+-- coroutine until its callback fires. `arm` receives a `resolve`
+-- function and must start exactly one libuv op that calls it.
+local function await(arm)
+    local co = assert(coroutine.running(),
+        "read: await must run inside a tool handler coroutine")
+    local res, fired = nil, false
+    arm(function(...)
+        assert(not fired, "read: libuv callback fired twice")
+        fired = true
+        res = table.pack(...)
+        local ok, err = coroutine.resume(co)
+        if not ok then error(err, 0) end
+    end)
+    coroutine.yield()
+    return table.unpack(res, 1, res.n)
+end
+
+local function fs_stat(path)
+    return await(function(resolve) uv.fs_stat(path, resolve) end)
+end
+
+local function fs_open(path, flags, mode)
+    return await(function(resolve) uv.fs_open(path, flags, mode, resolve) end)
+end
+
+local function fs_read(fd, size, offset)
+    return await(function(resolve) uv.fs_read(fd, size, offset, resolve) end)
+end
+
+local function fs_close(fd)
+    return await(function(resolve) uv.fs_close(fd, resolve) end)
+end
+
+-- `uv.fs_opendir` takes the entries-per-readdir batch hint as its
+-- SECOND positional arg, before the callback. We pass it explicitly so
+-- each `fs_readdir` returns up to `n` entries per round-trip.
+local READDIR_BATCH = 256
+local function fs_opendir(path)
+    return await(function(resolve) uv.fs_opendir(path, resolve, READDIR_BATCH) end)
+end
+local function fs_readdir(dir)
+    return await(function(resolve) uv.fs_readdir(dir, resolve) end)
+end
+local function fs_closedir(dir)
+    return await(function(resolve) uv.fs_closedir(dir, resolve) end)
+end
+
 local function scandir_sorted(path)
-    local req, err = uv.fs_scandir(path)
-    if not req then
+    -- Use the streaming opendir/readdir/closedir API rather than
+    -- `fs_scandir` + `fs_scandir_next`: the latter's iterator advances
+    -- with synchronous `readdir(3)` syscalls, which would block the
+    -- shared event loop on large directories. opendir/readdir route
+    -- every batch through libuv's thread pool, keeping us cooperative.
+    local err, dir = fs_opendir(path)
+    if not dir then
         return nil, err
     end
 
     local entries = {}
     while true do
-        local name, typ = uv.fs_scandir_next(req)
-        if not name then break end
-        entries[#entries + 1] = { name = name, typ = typ }
+        local rerr, batch = fs_readdir(dir)
+        if rerr then
+            fs_closedir(dir)
+            return nil, rerr
+        end
+        -- A nil/empty batch signals end of the directory stream.
+        if batch == nil or #batch == 0 then break end
+        for _, e in ipairs(batch) do
+            -- async readdir yields { name = ..., type = ... }.
+            entries[#entries + 1] = { name = e.name, typ = e.type }
+        end
     end
+    fs_closedir(dir)
 
     table.sort(entries, function(a, b)
         return a.name < b.name
@@ -136,7 +216,7 @@ return {
             return string.format("Error: limit (%d) must be >= 1.", limit)
         end
 
-        local stat, stat_err = uv.fs_stat(path)
+        local stat_err, stat = fs_stat(path)
         if not stat then
             return "Error: " .. (stat_err or ("could not stat " .. path))
         end
@@ -144,7 +224,7 @@ return {
             return render_directory_listing(path, offset, limit)
         end
 
-        local fd, open_err = uv.fs_open(path, "r", 0)
+        local open_err, fd = fs_open(path, "r", 0)
         if not fd then
             return "Error: " .. (open_err or ("could not open " .. path))
         end
@@ -190,12 +270,16 @@ return {
         end
 
         while not done do
-            local data, err = uv.fs_read(fd, READ_CHUNK_SIZE, file_offset)
+            local err, data = fs_read(fd, READ_CHUNK_SIZE, file_offset)
             if err then
                 read_err = err
                 break
             end
-            if data == nil then
+            -- luv's `fs_read` signals EOF by returning the empty
+            -- string, NOT nil (nil only accompanies an error). Treat
+            -- both nil and "" as EOF; otherwise the loop re-reads at
+            -- the same offset forever, pinning a core at 100%.
+            if data == nil or #data == 0 then
                 if #carry > 0 then
                     process_line(carry)
                     carry = ""
@@ -226,7 +310,7 @@ return {
             end
         end
 
-        uv.fs_close(fd)
+        fs_close(fd)
 
         if read_err then
             return "Error: read failed: " .. tostring(read_err)
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
diff --git a/agent/tools/write.lua b/agent/tools/write.lua
index 01bf3b2..42cb85a 100644
--- a/agent/tools/write.lua
+++ b/agent/tools/write.lua
@@ -4,8 +4,46 @@
 
 local uv = require("luv")
 
+-- ---------------------------------------------------------------------------
+-- Coroutine-synchronous libuv wrappers. panto runs each tool handler in
+-- its own coroutine and drives a single `uv.run()` to completion; a
+-- handler may only block by yielding on a pending libuv op whose
+-- callback resumes it exactly once. `await` encapsulates that.
+-- ---------------------------------------------------------------------------
+local function await(arm)
+    local co = assert(coroutine.running(),
+        "write: await must run inside a tool handler coroutine")
+    local res, fired = nil, false
+    arm(function(...)
+        assert(not fired, "write: libuv callback fired twice")
+        fired = true
+        res = table.pack(...)
+        local ok, err = coroutine.resume(co)
+        if not ok then error(err, 0) end
+    end)
+    coroutine.yield()
+    return table.unpack(res, 1, res.n)
+end
+
+-- Async `fs_mkdir`. luv's async callback is `(err, success)`, and on
+-- EEXIST it reports the error via a string code in `err`; we detect
+-- that textually since the async form doesn't surface the errno name
+-- separately the way the sync form's third return does.
+local function fs_mkdir(path, mode)
+    return await(function(resolve) uv.fs_mkdir(path, mode, resolve) end)
+end
+local function fs_open(path, flags, mode)
+    return await(function(resolve) uv.fs_open(path, flags, mode, resolve) end)
+end
+local function fs_write(fd, data, offset)
+    return await(function(resolve) uv.fs_write(fd, data, offset, resolve) end)
+end
+local function fs_close(fd)
+    return await(function(resolve) uv.fs_close(fd, resolve) end)
+end
+
 -- mkdir -p equivalent: walk the path components and create each.
--- Ignores EEXIST so the call is idempotent.
+-- Ignores "already exists" so the call is idempotent.
 local function mkdir_p(path)
     if path == "" or path == "." or path == "/" then return true end
     local parts = {}
@@ -15,8 +53,8 @@ local function mkdir_p(path)
     local cur = (path:sub(1, 1) == "/") and "" or "."
     for _, piece in ipairs(parts) do
         cur = cur .. "/" .. piece
-        local ok, err, name = uv.fs_mkdir(cur, tonumber("755", 8))
-        if not ok and name ~= "EEXIST" then
+        local err = fs_mkdir(cur, tonumber("755", 8))
+        if err and not tostring(err):find("EEXIST") then
             return nil, err
         end
     end
@@ -60,19 +98,24 @@ return {
             end
         end
 
-        local f, open_err = io.open(path, "wb")
-        if not f then
+        local open_err, fd = fs_open(path, "w", tonumber("644", 8))
+        if not fd then
             return "Error: " .. (open_err or ("could not open " .. path .. " for writing"))
         end
 
-        local ok, write_err = f:write(content)
-        if not ok then
-            f:close()
-            return "Error: write failed: " .. tostring(write_err)
+        -- A single `fs_write` may short-write; loop until all bytes land.
+        local offset = 0
+        while offset < #content do
+            local werr, n = fs_write(fd, content:sub(offset + 1), offset)
+            if not n then
+                fs_close(fd)
+                return "Error: write failed: " .. tostring(werr)
+            end
+            offset = offset + n
         end
 
-        local close_ok, close_err = f:close()
-        if not close_ok then
+        local close_err = fs_close(fd)
+        if close_err then
             return "Error: close failed: " .. tostring(close_err)
         end