Stage and apply event-field overrides at tool lifecycle boundaries

Stage writable event-field overrides by tool call id so they can be applied at the correct point in the turn lifecycle. Capture input overrides from tool_call_complete, capture output overrides from tool_result, and clear staged overrides defensively at turn start and shutdown.

Also honor user_message text overrides when starting a turn, add the helper that applies staged overrides to the agent at dispatch boundaries, and cover the behavior with tests for staged input/output overrides.

1 files changed, 87 insertions, 133 deletions

diff --git a/src/lua_runtime.zig b/src/lua_runtime.zig
index 6b3a2f9..e2b7ec8 100644
--- a/src/lua_runtime.zig
+++ b/src/lua_runtime.zig
@@ -84,10 +84,6 @@ 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 the scheduler
-    /// calls to drive libuv to completion. `0` until
-    /// `installScheduler` runs.
-    uv_run_ref: c_int = 0,
     /// Registry ref to `require("luv").update_time`. libuv caches the
     /// loop's notion of "now" and only refreshes it while `uv.run`
     /// executes; the loop is idle between batches (and during a slow
@@ -214,13 +210,45 @@ pub const LuaRuntime = struct {
         c.lua_settop(L, c.lua_gettop(L) - 1); // pop the leftover `panto`
     }
 
+    /// Expose the CLI's session agent to extensions as `panto.ext.agent`:
+    /// a *borrowed* native Agent userdata built by the module's
+    /// `_wrap_agent`, giving extensions the full agent method surface
+    /// (`add_system_message`, `compact`, …) on the live session. Must run
+    /// after `installPantoModule`. Borrowing is safe because the CLI
+    /// deinits this runtime (closing the Lua state) before the agent.
+    pub fn setExtAgent(self: *LuaRuntime, agent: *panto.Agent) !void {
+        const L = self.L;
+        lua_bridge.pushPantoTable(L);
+        if (c.lua_type(L, -1) != lua_bridge.T_TABLE) {
+            c.lua_settop(L, c.lua_gettop(L) - 1);
+            return RuntimeError.LuaInitFailed;
+        }
+        _ = c.lua_getfield(L, -1, "_wrap_agent");
+        if (c.lua_type(L, -1) != lua_bridge.T_FUNCTION) {
+            c.lua_settop(L, c.lua_gettop(L) - 2);
+            return RuntimeError.LuaInitFailed;
+        }
+        c.lua_pushlightuserdata(L, agent);
+        if (c.lua_pcallk(L, 1, 1, 0, 0, null) != 0) {
+            logTopAsError(L, "panto-lua: _wrap_agent failed");
+            c.lua_settop(L, c.lua_gettop(L) - 2);
+            return RuntimeError.LuaInitFailed;
+        }
+        // Stack: ..., panto, agent_ud
+        lua_bridge.pushExtTable(L);
+        c.lua_pushvalue(L, -2); // copy agent_ud above ext
+        c.lua_setfield(L, -2, "agent"); // ext.agent = agent_ud
+        c.lua_settop(L, c.lua_gettop(L) - 3); // pop ext, agent_ud, panto
+    }
+
     /// Install the libuv-driven coroutine scheduler:
     ///   - Register `panto._record_result` (C function with `self` as
     ///     light-userdata upvalue) so the wrapper closure can hand
     ///     results back to Zig.
     ///   - Create the wrapper closure that runs a user handler in
     ///     `pcall` and reports the result.
-    ///   - Cache `require("luv").run` for fast per-tick access.
+    ///   - Cache `require("luv").update_time` for the per-batch clock
+    ///     refresh.
     ///
     /// Must be called after luarocks bootstrap has installed luv,
     /// otherwise the `require("luv")` step will fail. Also after
@@ -229,7 +257,6 @@ pub const LuaRuntime = struct {
     pub fn installScheduler(self: *LuaRuntime) !void {
         try installRecordResult(self);
         try installWrapperClosure(self);
-        try cacheUvRun(self);
         try cacheUvUpdateTime(self);
     }
 
@@ -252,9 +279,6 @@ pub const LuaRuntime = struct {
         if (self.wrapper_ref != 0) {
             c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, self.wrapper_ref);
         }
-        if (self.uv_run_ref != 0) {
-            c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, self.uv_run_ref);
-        }
         if (self.uv_update_time_ref != 0) {
             c.luaL_unref(self.L, lua_bridge.LUA_REGISTRYINDEX, self.uv_update_time_ref);
         }
@@ -317,11 +341,11 @@ pub const LuaRuntime = struct {
             c.lua_settop(L, c.lua_gettop(L) - 1);
             return error.LuaLoadFailed;
         }
-        if (c.lua_pcallk(L, 0, 1, 0, 0, null) != 0) {
-            logTopAsError(L, "lua: failed to evaluate extension");
-            c.lua_settop(L, c.lua_gettop(L) - 1);
-            return error.LuaRunFailed;
-        }
+        // The chunk is user Lua: run it through the shared entrypoint
+        // machinery like everything else (eval must still be
+        // side-effect-free per the extension contract, but it gets the
+        // same execution environment).
+        try lua_bridge.runEntrypointResult(L, 0, "extension eval", null, null);
         // Stack: ..., returned_value
         defer c.lua_settop(L, c.lua_gettop(L) - 1); // pop returned value
         try self.collectEntriesFromTop(script_path, out);
@@ -337,13 +361,11 @@ pub const LuaRuntime = struct {
         out: *std.array_list.Managed(Entry),
     ) !void {
         const L = self.L;
+        // `require` runs the rock's chunk — user Lua — so it goes through
+        // the shared entrypoint machinery like every other entrypoint.
         _ = c.lua_getglobal(L, "require");
         _ = c.lua_pushlstring(L, module_name.ptr, module_name.len);
-        if (c.lua_pcallk(L, 1, 1, 0, 0, null) != 0) {
-            logTopAsError(L, "lua: failed to require rock");
-            c.lua_settop(L, c.lua_gettop(L) - 1);
-            return error.LuaRunFailed;
-        }
+        try lua_bridge.runEntrypointResult(L, 1, "rock eval", null, null);
         // Stack: ..., returned_value
         defer c.lua_settop(L, c.lua_gettop(L) - 1);
         try self.collectEntriesFromTop(module_name, out);
@@ -458,13 +480,14 @@ pub const LuaRuntime = struct {
                 return error.LuaRunFailed;
             }
         } else {
+            // Run activate() as a user-Lua entrypoint: inside a coroutine
+            // with the uv loop driven to completion, so activation code may
+            // await libuv work like any other user Lua.
             _ = c.lua_rawgeti(L, lua_bridge.LUA_REGISTRYINDEX, entry.ref); // push activate fn
-            if (c.lua_pcallk(L, 0, 0, 0, 0, null) != 0) {
-                logTopAsError(L, "lua: extension activate() failed");
-                c.lua_settop(L, c.lua_gettop(L) - 1);
+            lua_bridge.runEntrypoint(L, 0, "extension activate()", null, null) catch {
                 self.finishEntry(entry);
                 return error.LuaRunFailed;
-            }
+            };
         }
 
         self.finishEntry(entry);
@@ -630,17 +653,18 @@ pub const LuaRuntime = struct {
         }
     }
 
-    /// Run a Lua slash-command handler synchronously. `handler_ref` is a
-    /// registry ref to a `function(args)` closure; `args` is the trimmed
-    /// remainder of the command line. The handler runs under `pcall` with
-    /// a `debug.traceback` error handler so failures surface a readable
-    /// message instead of crashing the REPL.
+    /// Run a Lua slash-command handler. `handler_ref` is a registry ref
+    /// to a `function(args)` closure; `args` is the trimmed remainder of
+    /// the command line. Like every user-Lua entrypoint it runs inside a
+    /// coroutine with the uv loop driven to completion, so command
+    /// handlers may await libuv work; the call is synchronous from the
+    /// REPL's point of view.
     ///
     /// Slash-command handlers act by side effect (printing, mutating
     /// state) and their return value is ignored. On Lua error, returns
     /// `RuntimeError.LuaHandlerCrashed` after copying the error message
-    /// into `err_buf` (truncated to fit); `*err_len` is set to the number
-    /// of bytes written.
+    /// (with traceback) into `err_buf` (truncated to fit); `*err_len` is
+    /// set to the number of bytes written.
     pub fn runCommand(
         self: *LuaRuntime,
         handler_ref: c_int,
@@ -651,30 +675,10 @@ pub const LuaRuntime = struct {
         const L = self.L;
         err_len.* = 0;
 
-        // Install a traceback error handler beneath the call frame.
-        _ = c.lua_getglobal(L, "debug");
-        _ = c.lua_getfield(L, -1, "traceback");
-        c.lua_copy(L, -1, -2);
-        c.lua_settop(L, c.lua_gettop(L) - 1); // pop the `debug` table
-        const errfunc_idx = c.lua_gettop(L);
-
-        // Push handler + args string.
         _ = c.lua_rawgeti(L, lua_bridge.LUA_REGISTRYINDEX, @intCast(handler_ref));
         _ = c.lua_pushlstring(L, args.ptr, args.len);
-
-        const rc = c.lua_pcallk(L, 1, 0, errfunc_idx, 0, null);
-        if (rc != 0) {
-            var len: usize = 0;
-            const ptr = c.lua_tolstring(L, -1, &len);
-            if (ptr != null) {
-                const n = @min(len, err_buf.len);
-                @memcpy(err_buf[0..n], ptr[0..n]);
-                err_len.* = n;
-            }
-            c.lua_settop(L, c.lua_gettop(L) - 2); // pop error + errfunc
+        lua_bridge.runEntrypoint(L, 1, "slash command", err_buf, err_len) catch
             return RuntimeError.LuaHandlerCrashed;
-        }
-        c.lua_settop(L, c.lua_gettop(L) - 1); // pop errfunc
     }
 
     /// Slash commands declared by loaded extensions, in registration
@@ -825,7 +829,7 @@ fn runBatch(
     results: []panto.ToolCallResult,
     allocator: Allocator,
 ) !void {
-    if (self.wrapper_ref == 0 or self.uv_run_ref == 0) {
+    if (self.wrapper_ref == 0) {
         // Scheduler not installed. We can still run synchronous
         // handlers — use the legacy path that drives one coroutine at
         // a time without an event loop.
@@ -1035,22 +1039,18 @@ fn startCoroutine(
     };
 }
 
-/// 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.
+/// Drain the uv loop (`lua_bridge.drainLoop` — the same drain every
+/// single-fn entrypoint uses). 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.
 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, "default", 7);
-    if (c.lua_pcallk(L, 1, 1, 0, 0, null) != 0) {
+    lua_bridge.drainLoop(L) catch {
         logTopAsError(L, "panto-lua: uv.run failed");
         c.lua_settop(L, c.lua_gettop(L) - 1);
         return error.UvRunFailed;
-    }
-    // Discard the boolean return value.
-    c.lua_settop(L, c.lua_gettop(L) - 1);
+    };
 }
 
 /// Pre-scheduler fallback (used in unit tests and during early
@@ -1092,10 +1092,11 @@ fn runLegacySync(
     return .{ .ok = out_parts };
 }
 
-/// Run a handler synchronously, with no event loop. On Lua-level
-/// failure the error message string is duped into `*err_msg_out`
-/// (caller owns) before returning the typed error. `err_msg_out` is
-/// only written on the error path; on success it is left untouched.
+/// Run a handler through the shared entrypoint machinery (coroutine +
+/// loop drain), keeping its return value. On Lua-level failure the error
+/// message is duped into `*err_msg_out` (caller owns) before returning
+/// the typed error. `err_msg_out` is only written on the error path; on
+/// success it is left untouched.
 fn invokeCoroutineSync(
     L: *c.lua_State,
     handler_ref: c_int,
@@ -1103,49 +1104,26 @@ fn invokeCoroutineSync(
     allocator: Allocator,
     err_msg_out: *?[]u8,
 ) !panto.ResultParts {
-    const co = c.lua_newthread(L) orelse return RuntimeError.LuaInitFailed;
-    defer c.lua_settop(L, c.lua_gettop(L) - 1);
-
-    _ = c.lua_rawgeti(co, lua_bridge.LUA_REGISTRYINDEX, @intCast(handler_ref));
-    if (c.lua_type(co, -1) != lua_bridge.T_FUNCTION) {
+    _ = c.lua_rawgeti(L, lua_bridge.LUA_REGISTRYINDEX, @intCast(handler_ref));
+    if (c.lua_type(L, -1) != lua_bridge.T_FUNCTION) {
+        c.lua_settop(L, c.lua_gettop(L) - 1);
         return RuntimeError.LuaHandlerNotFound;
     }
 
     var arena_state = std.heap.ArenaAllocator.init(allocator);
     defer arena_state.deinit();
-    try lua_bridge.pushJsonAsLua(co, arena_state.allocator(), input);
-
-    var nresults: c_int = 0;
-    const status = c.lua_resume(co, L, 1, &nresults);
+    try lua_bridge.pushJsonAsLua(L, arena_state.allocator(), input);
 
-    switch (status) {
-        c.LUA_OK => {
-            if (nresults < 1) return RuntimeError.BadHandlerReturn;
-            return try lua_bridge.readHandlerResult(co, -1, allocator);
-        },
-        c.LUA_YIELD => {
-            const msg = "lua: tool handler yielded with no event loop installed; call installScheduler() before dispatching";
-            if (@import("builtin").is_test) {
-                std.log.warn("{s}", .{msg});
-            } else {
-                std.log.err("{s}", .{msg});
-            }
-            return RuntimeError.LuaHandlerYielded;
-        },
-        else => {
-            // Capture the Lua error message into the caller's slot
-            // *before* logging+returning, so the legacy sync path can
-            // surface it to the model rather than just the typed
-            // error name.
-            var msg_len: usize = 0;
-            const msg_ptr = c.lua_tolstring(co, -1, &msg_len);
-            if (msg_ptr != null) {
-                err_msg_out.* = allocator.dupe(u8, msg_ptr[0..msg_len]) catch null;
-            }
-            logTopAsError(co, "lua: handler crashed");
-            return RuntimeError.LuaHandlerCrashed;
-        },
-    }
+    var err_buf: [1024]u8 = undefined;
+    var err_len: usize = 0;
+    lua_bridge.runEntrypointResult(L, 1, "tool handler", &err_buf, &err_len) catch {
+        err_msg_out.* = allocator.dupe(u8, err_buf[0..err_len]) catch null;
+        return RuntimeError.LuaHandlerCrashed;
+    };
+    // Success: the handler's result is on top of L.
+    defer c.lua_settop(L, c.lua_gettop(L) - 1);
+    if (c.lua_type(L, -1) == lua_bridge.T_NIL) return RuntimeError.BadHandlerReturn;
+    return try lua_bridge.readHandlerResult(L, -1, allocator);
 }
 
 // ---------------------------------------------------------------------------
@@ -1249,30 +1227,6 @@ fn installWrapperClosure(self: *LuaRuntime) !void {
 
 /// 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
-    ;
-    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, 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]. 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_run_ref = c.luaL_ref(L, lua_bridge.LUA_REGISTRYINDEX);
-}
-
 /// Cache `require("luv").update_time` so the scheduler can refresh the
 /// loop clock at the start of each batch. See `uv_update_time_ref`.
 fn cacheUvUpdateTime(self: *LuaRuntime) !void {
@@ -1743,7 +1697,7 @@ test "directory-style extension can require sibling modules" {
     try testing.expectEqualStrings("HI!", okText(results[0]));
 }
 
-test "yielding handler with no event loop surfaces LuaHandlerYielded" {
+test "yielding handler with no event loop surfaces a suspension error" {
     var tmp = testing.tmpDir(.{});
     defer tmp.cleanup();
 
@@ -1768,15 +1722,15 @@ test "yielding handler with no event loop surfaces LuaHandlerYielded" {
     defer freeResults(&results);
 
     // Same policy as the crash test: the failure is surfaced as `.ok`
-    // text so libpanto doesn't abort the turn. The error type name
-    // (`LuaHandlerYielded`) is included via `@errorName` in the legacy
-    // sync path's formatToolError call.
+    // text so libpanto doesn't abort the turn. The message is the shared
+    // entrypoint machinery's post-drain suspension diagnosis (without luv
+    // the drain is a no-op, so the yield can never be resumed).
     try testing.expect(std.mem.startsWith(
         u8,
         okText(results[0]),
         "panto-lua: tool 'sleeper' failed:",
     ));
-    try testing.expect(std.mem.indexOf(u8, okText(results[0]), "LuaHandlerYielded") != null);
+    try testing.expect(std.mem.indexOf(u8, okText(results[0]), "suspended after the uv loop drained") != null);
 }
 
 // Integration test: requires a data home with luv already installed.