rework lua tool scheduler

1 files changed, 242 insertions, 107 deletions

diff --git a/src/lua_runtime.zig b/src/lua_runtime.zig
index fed56a8..2cf3cf9 100644
--- a/src/lua_runtime.zig
+++ b/src/lua_runtime.zig
@@ -63,15 +63,10 @@ pub const LuaRuntime = struct {
     /// `panto._record_result`. Allocated once at `create`; reused for
     /// every call. `0` until `installScheduler` runs.
     wrapper_ref: c_int = 0,
-    /// Registry ref to `require("luv").run`, the function we call to
-    /// tick libuv between coroutine resumes. `0` until
+    /// Registry ref to `require("luv").run`, the function the scheduler
+    /// calls to drive libuv to completion. `0` until
     /// `installScheduler` runs.
     uv_run_ref: c_int = 0,
-    /// Registry ref to `require("luv").loop_alive`, used by the
-    /// scheduler to detect handler coroutines that yielded without
-    /// arming any libuv work (a deadlock we surface rather than hang
-    /// on). `0` until `installScheduler` runs.
-    uv_loop_alive_ref: c_int = 0,
 
     /// Pointer to the in-flight batch, valid only for the duration of
     /// one `invoke_batch` call. The `panto._record_result` C function
@@ -134,9 +129,6 @@ pub const LuaRuntime = struct {
         if (self.uv_run_ref != 0) {
             c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, self.uv_run_ref);
         }
-        if (self.uv_loop_alive_ref != 0) {
-            c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, self.uv_loop_alive_ref);
-        }
 
         c.lua_close(self.L);
 
@@ -489,7 +481,12 @@ fn runBatch(
         if (r != 0) c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, r);
     };
 
-    var pending: usize = 0;
+    // Step 1: start every call's coroutine. A *synchronous* handler
+    // (one that never yields) runs to completion right here and
+    // records its result via the wrapper; `still_pending` is false.
+    // An *async* handler arms one or more libuv operations and yields
+    // — that armed work is what keeps the event loop alive in step 2.
+    var any_pending = false;
     for (calls, 0..) |call, i| {
         const handler_ref = self.handlers.get(call.tool_name) orelse {
             // Synthesize a recorded "err" result; don't even bother
@@ -503,74 +500,60 @@ fn runBatch(
         };
         const t = try startCoroutine(self, i, handler_ref, call.input, allocator);
         thread_refs[i] = t.thread_ref;
-        if (t.still_pending) pending += 1;
+        if (t.still_pending) any_pending = true;
     }
 
-    while (pending > 0) {
-        try driveUvOnce(self);
-        // Reap any coroutines that terminated during the tick.
-        var reaped: usize = 0;
-        for (thread_refs, 0..) |tref, i| {
-            if (tref == 0) continue;
-            // Push the thread, check status, pop.
-            _ = c.lua_rawgeti(self.L, lua_bridge.LUA_REGISTRYINDEX, tref);
-            const co: *c.lua_State = @ptrCast(c.lua_tothread(self.L, -1).?);
-            const status = c.lua_status(co);
-            c.lua_settop(self.L, c.lua_gettop(self.L) - 1);
-            if (status != c.LUA_YIELD) {
-                // Terminated (LUA_OK or error). The wrapper should
-                // have called `_record_result` already; if not, synthesize.
-                if (!slots[i].recorded) {
-                    slots[i] = .{
-                        .recorded = true,
-                        .ok = false,
-                        .err_msg = try allocator.dupe(
-                            u8,
-                            "panto: handler terminated without recording a result",
-                        ),
-                    };
-                }
-                c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, tref);
-                thread_refs[i] = 0;
-                reaped += 1;
-            }
-        }
-        if (reaped > 0) {
-            if (reaped > pending) {
-                // Defensive: keep the counter sane.
-                pending = 0;
-            } else {
-                pending -= reaped;
-            }
-            continue;
-        }
-        if (!try loopAlive(self)) {
-            // No libuv handles are pending, but we still have alive
-            // coroutines. They yielded without arranging to be woken.
-            // Mark them as failed and break.
-            //
-            // Note: we ask `uv.loop_alive` rather than relying on the
-            // return value of `uv.run("once")`, which is a boolean
-            // signalling whether `uv.stop()` was called — not a count
-            // of active handles. Conflating the two used to break
-            // multi-tick tools (e.g. `bash`) on their second tick.
+    // Step 2: drive the event loop to completion. The contract for
+    // async handlers is that the *only* way a handler coroutine may
+    // block is by yielding on a pending libuv operation whose callback
+    // will eventually resume it. Under that invariant, `uv.run`
+    // returns exactly when no active handles remain — which means
+    // every coroutine has either finished or violated the contract
+    // (yielded without keeping libuv work alive). Either way there is
+    // nothing left to wake, so a single run-to-completion pass is all
+    // we need; the per-tick reap/`loop_alive` dance the scheduler used
+    // to do is unnecessary.
+    if (any_pending) {
+        try driveUvToCompletion(self);
+    }
 
-            for (thread_refs, 0..) |tref, i| {
-                if (tref == 0) continue;
-                slots[i] = .{
-                    .recorded = true,
-                    .ok = false,
-                    .err_msg = try allocator.dupe(
-                        u8,
-                        "panto: handler yielded but no libuv handle is pending; " ++
-                            "did you forget to await with luv?",
-                    ),
-                };
-                c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, tref);
-                thread_refs[i] = 0;
-            }
-            break;
+    // Step 3: reap. Every coroutine should now be terminated. Any that
+    // is still suspended (`LUA_YIELD`) violated the contract — it
+    // yielded without arming libuv work that would resume it (or armed
+    // work, then tore it down without resolving). Surface that as a
+    // per-call error rather than hanging the process.
+    for (thread_refs, 0..) |tref, i| {
+        if (tref == 0) continue;
+        _ = c.lua_rawgeti(self.L, lua_bridge.LUA_REGISTRYINDEX, tref);
+        const co: *c.lua_State = @ptrCast(c.lua_tothread(self.L, -1).?);
+        const status = c.lua_status(co);
+        c.lua_settop(self.L, c.lua_gettop(self.L) - 1);
+        if (status == c.LUA_YIELD) {
+            slots[i] = .{
+                .recorded = true,
+                .ok = false,
+                .err_msg = try allocator.dupe(
+                    u8,
+                    "panto: handler is still suspended after the event loop " ++
+                        "drained; it yielded without a pending libuv operation " ++
+                        "to resume it (did a uv callback fail to resume the " ++
+                        "coroutine, or was its handle closed without resolving?)",
+                ),
+            };
+        } else if (!slots[i].recorded) {
+            // Terminated cleanly or with an error, but the wrapper
+            // never recorded — should not happen, but don't drop it.
+            slots[i] = .{
+                .recorded = true,
+                .ok = false,
+                .err_msg = try allocator.dupe(
+                    u8,
+                    "panto: handler terminated without recording a result",
+                ),
+            };
         }
+        c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, tref);
+        thread_refs[i] = 0;
     }
 
     // Translate slots into the libpanto-shaped results.
@@ -673,13 +656,15 @@ fn startCoroutine(
     };
 }
 
-/// Call `uv.run("once")`. The return value (a boolean meaning "was
-/// `uv.stop()` called with handles still alive?") is not useful to us
-/// — to decide whether to keep ticking we call `loopAlive` separately.
-fn driveUvOnce(self: *LuaRuntime) !void {
+/// Call `uv.run("default")`, which runs the event loop until there are
+/// no active handles or requests left. Under the async-tool contract
+/// (a handler may only block by yielding on a pending libuv operation
+/// whose callback resumes it), this returns exactly when every handler
+/// coroutine has finished. The boolean return value is discarded.
+fn driveUvToCompletion(self: *LuaRuntime) !void {
     const L = self.L;
     _ = c.lua_rawgeti(L, lua_bridge.LUA_REGISTRYINDEX, @intCast(self.uv_run_ref));
-    _ = c.lua_pushlstring(L, "once", 4);
+    _ = c.lua_pushlstring(L, "default", 7);
     if (c.lua_pcallk(L, 1, 1, 0, 0, null) != 0) {
         logTopAsError(L, "panto-lua: uv.run failed");
         c.lua_settop(L, c.lua_gettop(L) - 1);
@@ -689,23 +674,6 @@ fn driveUvOnce(self: *LuaRuntime) !void {
     c.lua_settop(L, c.lua_gettop(L) - 1);
 }
 
-/// Call `uv.loop_alive()`. Returns true iff libuv has any referenced
-/// active handles, requests, or closing handles. We use this — not
-/// the return value of `uv.run` — to detect coroutines that yielded
-/// without arranging to be woken.
-fn loopAlive(self: *LuaRuntime) !bool {
-    const L = self.L;
-    _ = c.lua_rawgeti(L, lua_bridge.LUA_REGISTRYINDEX, @intCast(self.uv_loop_alive_ref));
-    if (c.lua_pcallk(L, 0, 1, 0, 0, null) != 0) {
-        logTopAsError(L, "panto-lua: uv.loop_alive failed");
-        c.lua_settop(L, c.lua_gettop(L) - 1);
-        return error.UvRunFailed;
-    }
-    const alive = c.lua_toboolean(L, -1) != 0;
-    c.lua_settop(L, c.lua_gettop(L) - 1);
-    return alive;
-}
-
 /// Pre-scheduler fallback (used in unit tests and during early
 /// startup before `installScheduler` has run).
 fn runLegacySync(
@@ -895,32 +863,29 @@ fn installWrapperClosure(self: *LuaRuntime) !void {
     self.wrapper_ref = c.luaL_ref(L, lua_bridge.LUA_REGISTRYINDEX);
 }
 
-/// Cache `require("luv").run` and `require("luv").loop_alive` in the
-/// registry so the scheduler can invoke them cheaply per tick.
+/// Cache `require("luv").run` in the registry so the scheduler can
+/// invoke it cheaply without re-resolving `require` each batch.
 fn cacheUvRun(self: *LuaRuntime) !void {
     const L = self.L;
     const snippet =
         \\local uv = require("luv")
-        \\return uv.run, uv.loop_alive
+        \\return uv.run
     ;
     if (c.luaL_loadstring(L, snippet) != 0) {
         logTopAsError(L, "panto-lua: failed to compile luv lookup");
         c.lua_settop(L, c.lua_gettop(L) - 1);
         return RuntimeError.LuaInitFailed;
     }
-    if (c.lua_pcallk(L, 0, 2, 0, 0, null) != 0) {
+    if (c.lua_pcallk(L, 0, 1, 0, 0, null) != 0) {
         logTopAsError(L, "panto-lua: require('luv') failed (was the bootstrap successful?)");
         c.lua_settop(L, c.lua_gettop(L) - 1);
         return RuntimeError.LuaInitFailed;
     }
-    // Stack: [..., uv.run, uv.loop_alive]. luaL_ref pops the top.
-    if (c.lua_type(L, -1) != lua_bridge.T_FUNCTION or
-        c.lua_type(L, -2) != lua_bridge.T_FUNCTION)
-    {
-        c.lua_settop(L, c.lua_gettop(L) - 2);
+    // Stack: [..., uv.run]. luaL_ref pops the top.
+    if (c.lua_type(L, -1) != lua_bridge.T_FUNCTION) {
+        c.lua_settop(L, c.lua_gettop(L) - 1);
         return RuntimeError.LuaInitFailed;
     }
-    self.uv_loop_alive_ref = c.luaL_ref(L, lua_bridge.LUA_REGISTRYINDEX);
     self.uv_run_ref = c.luaL_ref(L, lua_bridge.LUA_REGISTRYINDEX);
 }
 
@@ -1328,6 +1293,176 @@ test "scheduler: yielding handler is resumed by libuv" {
     try testing.expectEqualStrings("awake", results[1].ok);
 }
 
+// Contract-violation guard: a handler that yields WITHOUT arming any
+// libuv work can never be resumed. Under the run-to-completion model,
+// `uv.run` drains immediately (nothing to wait on) and the scheduler
+// must surface the still-suspended coroutine as a per-call error
+// rather than hanging the process. A sibling well-behaved call in the
+// same batch must still succeed.
+test "scheduler: handler that yields without arming libuv work is surfaced as an error" {
+    var rt = try LuaRuntime.create(testing.allocator);
+    defer rt.deinit();
+    const luarocks_rt = bootstrapRealRuntime(rt) catch |e| return e;
+    defer luarocks_rt.deinit();
+
+    var tmp = testing.tmpDir(.{});
+    defer tmp.cleanup();
+    const source =
+        \\panto.register_tool {
+        \\  name = "abandon", description = "yields into the void",
+        \\  schema = { type = "object" },
+        \\  handler = function() coroutine.yield(); return "never" end,
+        \\}
+        \\panto.register_tool {
+        \\  name = "fine", description = "sync ok",
+        \\  schema = { type = "object" },
+        \\  handler = function() return "ok" end,
+        \\}
+    ;
+    const path = try writeTempScript(tmp.dir, "abandon.lua", source);
+    defer testing.allocator.free(path);
+    try rt.loadExtension(path, null);
+
+    var src = rt.toolSource();
+    const calls = [_]panto.ToolCall{
+        .{ .tool_name = "abandon", .input = "{}" },
+        .{ .tool_name = "fine", .input = "{}" },
+    };
+    var results: [2]panto.ToolCallResult = .{
+        .{ .err = error.SourceDroppedCall },
+        .{ .err = error.SourceDroppedCall },
+    };
+    try src.vtable.invoke_batch(src.ctx, &calls, &results, testing.allocator);
+    defer for (results) |r| switch (r) {
+        .ok => |b| testing.allocator.free(b),
+        .err => {},
+    };
+
+    try testing.expect(std.mem.startsWith(u8, results[0].ok, "panto-lua: tool 'abandon' failed:"));
+    try testing.expect(std.mem.indexOf(u8, results[0].ok, "still suspended") != null);
+    try testing.expectEqualStrings("ok", results[1].ok);
+}
+
+/// Absolute path to the repo's `agent/tools` directory, derived from
+/// the build-time repo root (the test runner's cwd is unreliable).
+/// Caller owns the returned slice. Skips the test if the directory
+/// isn't present.
+fn findAgentToolsDir() ![]const u8 {
+    const versions = @import("versions");
+    const tools = try std.fs.path.join(testing.allocator, &.{ versions.repo_root, "agent", "tools" });
+    errdefer testing.allocator.free(tools);
+    const probe = try std.fs.path.join(testing.allocator, &.{ tools, "shell.lua" });
+    defer testing.allocator.free(probe);
+    std.Io.Dir.cwd().access(testing.io, probe, .{}) catch {
+        testing.allocator.free(tools);
+        return error.SkipZigTest;
+    };
+    return tools;
+}
+
+fn bootstrapRealRuntime(rt: *LuaRuntime) !*@import("luarocks_runtime.zig").LuarocksRuntime {
+    const home_z = std.c.getenv("PANTO_HOME") orelse return error.SkipZigTest;
+    const panto_home_env = std.mem.sliceTo(home_z, 0);
+    const manifest = @import("manifest.zig");
+    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
+    const so_path = try std.fmt.bufPrint(
+        &path_buf,
+        "{s}/rocks/lua-{s}/lib/lua/{s}/luv.so",
+        .{ panto_home_env, manifest.lua_version, manifest.lua_short_version },
+    );
+    std.Io.Dir.cwd().access(testing.io, so_path, .{}) catch return error.SkipZigTest;
+
+    var env: std.process.Environ.Map = .init(testing.allocator);
+    defer env.deinit();
+    try env.put("PANTO_HOME", panto_home_env);
+    const luarocks_runtime = @import("luarocks_runtime.zig");
+    const luarocks_rt = try luarocks_runtime.bootstrap(
+        testing.allocator,
+        testing.io,
+        &env,
+        rt.L,
+        "/usr/bin/true",
+    );
+    try rt.installScheduler();
+    return luarocks_rt;
+}
+
+// Reproduction: two REAL `std.shell` calls in one batch. shell.lua
+// uses spawn + two pipes + a timeout timer (3+ libuv events per call)
+// and resumes its own coroutine from a libuv callback. This exercises
+// the exact C-resume topology that the synthetic `timer_say` test does
+// not (spawn, dual pipes, spill-file fs calls). Skipped without luv.
+test "scheduler: two real std.shell calls in one batch do not deadlock" {
+    var rt = try LuaRuntime.create(testing.allocator);
+    defer rt.deinit();
+    const luarocks_rt = bootstrapRealRuntime(rt) catch |e| return e;
+    defer luarocks_rt.deinit();
+
+    const tools_dir = try findAgentToolsDir();
+    defer testing.allocator.free(tools_dir);
+    const shell_path = try std.fs.path.join(testing.allocator, &.{ tools_dir, "shell.lua" });
+    defer testing.allocator.free(shell_path);
+    try rt.loadTool(shell_path, tools_dir);
+
+    var src = rt.toolSource();
+    const calls = [_]panto.ToolCall{
+        .{ .tool_name = "std.shell", .input = "{\"command\":\"echo first; sleep 0.05; echo first-done\"}" },
+        .{ .tool_name = "std.shell", .input = "{\"command\":\"echo second; sleep 0.05; echo second-done\"}" },
+    };
+    var results: [2]panto.ToolCallResult = .{
+        .{ .err = error.SourceDroppedCall },
+        .{ .err = error.SourceDroppedCall },
+    };
+    try src.vtable.invoke_batch(src.ctx, &calls, &results, testing.allocator);
+    defer for (results) |r| switch (r) {
+        .ok => |b| testing.allocator.free(b),
+        .err => {},
+    };
+
+    try testing.expect(std.mem.indexOf(u8, results[0].ok, "first-done") != null);
+    try testing.expect(std.mem.indexOf(u8, results[1].ok, "second-done") != null);
+}
+
+// Reproduction: two REAL `std.read` calls in one batch. read.lua uses
+// only synchronous `uv.fs_*` calls and never yields, so both should
+// complete during dispatch without entering the drive loop.
+test "scheduler: two real std.read calls in one batch do not deadlock" {
+    var rt = try LuaRuntime.create(testing.allocator);
+    defer rt.deinit();
+    const luarocks_rt = bootstrapRealRuntime(rt) catch |e| return e;
+    defer luarocks_rt.deinit();
+
+    const tools_dir = try findAgentToolsDir();
+    defer testing.allocator.free(tools_dir);
+    const read_path = try std.fs.path.join(testing.allocator, &.{ tools_dir, "read.lua" });
+    defer testing.allocator.free(read_path);
+    try rt.loadTool(read_path, tools_dir);
+
+    const in0 = try std.fmt.allocPrint(testing.allocator, "{{\"path\":\"{s}/shell.lua\"}}", .{tools_dir});
+    defer testing.allocator.free(in0);
+    const in1 = try std.fmt.allocPrint(testing.allocator, "{{\"path\":\"{s}/read.lua\"}}", .{tools_dir});
+    defer testing.allocator.free(in1);
+
+    var src = rt.toolSource();
+    const calls = [_]panto.ToolCall{
+        .{ .tool_name = "std.read", .input = in0 },
+        .{ .tool_name = "std.read", .input = in1 },
+    };
+    var results: [2]panto.ToolCallResult = .{
+        .{ .err = error.SourceDroppedCall },
+        .{ .err = error.SourceDroppedCall },
+    };
+
+    try src.vtable.invoke_batch(src.ctx, &calls, &results, testing.allocator);
+    defer for (results) |r| switch (r) {
+        .ok => |b| testing.allocator.free(b),
+        .err => {},
+    };
+
+    try testing.expect(results[0].ok.len > 0);
+    try testing.expect(results[1].ok.len > 0);
+}
+
 test "loadExtension: duplicate tool name from a second extension errors" {
     var tmp = testing.tmpDir(.{});
     defer tmp.cleanup();