basic lua tools

1 files changed, 341 insertions, 0 deletions

diff --git a/src/lua_tool.zig b/src/lua_tool.zig
new file mode 100644
index 0000000..3064d2b
--- /dev/null
+++ b/src/lua_tool.zig
@@ -0,0 +1,341 @@
+//! `LuaTool` — adapts a Lua-defined tool to libpanto's `Tool` interface.
+//!
+//! Every call to `invoke` opens a fresh `lua_State`, re-runs the extension
+//! script (which calls `panto.register_tool(...)` and registers its handler
+//! into the registrations table), locates the handler by name, runs it
+//! under `xpcall` with a `debug.traceback` error handler, and tears the
+//! state down before returning. No state is reused across calls.
+//!
+//! This is the simplest possible model: each `invoke` is hermetic. The
+//! consequence is a few milliseconds of Lua startup per call, which is
+//! invisible next to LLM round-trip latency. A pool can be added later
+//! behind the same `LuaTool` interface without changing libpanto.
+
+const std = @import("std");
+const panto = @import("panto");
+const lua_bridge = @import("lua_bridge.zig");
+
+const c = lua_bridge.c;
+const Allocator = std.mem.Allocator;
+
+/// A LuaTool owns all the strings the `Tool` interface borrows (name,
+/// description, schema_json), plus the path to the script it dispatches
+/// to. `vtable.deinit` frees everything including the LuaTool itself.
+pub const LuaTool = struct {
+    allocator: Allocator,
+    // Owned, NUL-terminated for `luaL_loadfilex`.
+    script_path_z: [:0]u8,
+    name_owned: []u8,
+    description_owned: []u8,
+    schema_owned: []u8,
+
+    /// Build a LuaTool from a harvested registration plus the script path
+    /// it came from. All strings are copied into freshly-allocated bytes
+    /// owned by this LuaTool — the source slices can be freed after this
+    /// returns.
+    pub fn create(
+        allocator: Allocator,
+        script_path: []const u8,
+        name: []const u8,
+        description: []const u8,
+        schema_json: []const u8,
+    ) !panto.Tool {
+        const self = try allocator.create(LuaTool);
+        errdefer allocator.destroy(self);
+
+        self.* = .{
+            .allocator = allocator,
+            .script_path_z = try allocator.dupeZ(u8, script_path),
+            .name_owned = try allocator.dupe(u8, name),
+            .description_owned = try allocator.dupe(u8, description),
+            .schema_owned = try allocator.dupe(u8, schema_json),
+        };
+
+        return panto.Tool{
+            .name = self.name_owned,
+            .description = self.description_owned,
+            .schema_json = self.schema_owned,
+            .ctx = self,
+            .vtable = &vtable,
+        };
+    }
+
+    fn freeAll(self: *LuaTool) void {
+        const a = self.allocator;
+        a.free(self.script_path_z);
+        a.free(self.name_owned);
+        a.free(self.description_owned);
+        a.free(self.schema_owned);
+        a.destroy(self);
+    }
+};
+
+const vtable: panto.Tool.VTable = .{
+    .invoke = invoke,
+    .deinit = deinitTool,
+};
+
+fn invoke(
+    ctx: *anyopaque,
+    input: []const u8,
+    allocator: Allocator,
+) anyerror![]u8 {
+    const self: *LuaTool = @ptrCast(@alignCast(ctx));
+    return runLuaHandler(self, input, allocator);
+}
+
+fn deinitTool(ctx: *anyopaque, _: Allocator) void {
+    const self: *LuaTool = @ptrCast(@alignCast(ctx));
+    self.freeAll();
+}
+
+/// Open a fresh Lua state, re-load the script (running `panto.register_tool`),
+/// push the handler for `self.name_owned`, push the parsed input, run under
+/// `xpcall`, and return the result bytes.
+fn runLuaHandler(
+    self: *LuaTool,
+    input: []const u8,
+    out_allocator: Allocator,
+) anyerror![]u8 {
+    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
+    defer c.lua_close(L);
+    c.luaL_openlibs(L);
+    lua_bridge.install(L);
+
+    // Run the script so register_tool fires.
+    lua_bridge.loadFile(L, self.script_path_z) catch |err| {
+        logTopAsError(L, "lua: failed to load extension");
+        return err;
+    };
+
+    // Push a traceback handler at the bottom of the upcoming call frame.
+    _ = c.lua_getglobal(L, "debug");
+    _ = c.lua_getfield(L, -1, "traceback");
+    // Replace the `debug` slot with its `traceback` field.
+    c.lua_copy(L, -1, -2);
+    c.lua_settop(L, c.lua_gettop(L) - 1);
+    const errfunc_idx = c.lua_gettop(L);
+
+    try lua_bridge.pushHandler(L, self.name_owned);
+
+    // Parse the LLM's JSON input into a Lua table and push as argument.
+    var arena_state = std.heap.ArenaAllocator.init(out_allocator);
+    defer arena_state.deinit();
+    try lua_bridge.pushJsonAsLua(L, arena_state.allocator(), input);
+
+    // Call handler(input). 1 arg, 1 return value, traceback at errfunc_idx.
+    const rc = c.lua_pcallk(L, 1, 1, errfunc_idx, 0, null);
+    if (rc != 0) {
+        logTopAsError(L, "lua: handler crashed");
+        return error.LuaHandlerCrashed;
+    }
+
+    return lua_bridge.readHandlerResult(L, -1, out_allocator);
+}
+
+/// Best-effort: read the top of the Lua stack as a string and log it under
+/// the given prefix. Always leaves the stack as we found it.
+///
+/// In test builds we log at `warn` instead of `err` so the test runner
+/// doesn't treat expected-failure paths (e.g. a crash-protection test) as
+/// overall test failures. End users still see warnings in normal runs.
+fn logTopAsError(L: *c.lua_State, prefix: []const u8) void {
+    var len: usize = 0;
+    const msg = c.lua_tolstring(L, -1, &len);
+    const is_test = @import("builtin").is_test;
+    if (msg != null) {
+        if (is_test) {
+            std.log.warn("{s}: {s}", .{ prefix, msg[0..len] });
+        } else {
+            std.log.err("{s}: {s}", .{ prefix, msg[0..len] });
+        }
+    } else {
+        if (is_test) {
+            std.log.warn("{s} (no error message)", .{prefix});
+        } else {
+            std.log.err("{s} (no error message)", .{prefix});
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Standalone discovery helper
+// ---------------------------------------------------------------------------
+
+/// Open a *throwaway* Lua state, run `script_path`, harvest every
+/// `panto.register_tool` call into a slice of `LuaTool`s registered with
+/// the given registry. The state is closed before returning; only the
+/// metadata (name, description, schema) survives. Each tool rebuilds its
+/// own state on every invocation.
+///
+/// Returns the number of tools registered. On any failure, the caller
+/// should treat the extension as not loaded — partial-success cleanup is
+/// the caller's responsibility (typically: surface the error and abort).
+pub fn loadExtension(
+    allocator: Allocator,
+    registry: *panto.ToolRegistry,
+    script_path: []const u8,
+) !usize {
+    const path_z = try allocator.dupeZ(u8, script_path);
+    defer allocator.free(path_z);
+
+    const L = c.luaL_newstate() orelse return error.LuaInitFailed;
+    defer c.lua_close(L);
+    c.luaL_openlibs(L);
+    lua_bridge.install(L);
+
+    lua_bridge.loadFile(L, path_z) catch |err| {
+        logTopAsError(L, "lua: failed to load extension");
+        return err;
+    };
+
+    var arena_state = std.heap.ArenaAllocator.init(allocator);
+    defer arena_state.deinit();
+
+    const regs = try lua_bridge.harvestRegistrations(L, arena_state.allocator());
+    for (regs) |r| {
+        const tool = try LuaTool.create(
+            allocator,
+            script_path,
+            r.name,
+            r.description,
+            r.schema_json,
+        );
+        // If registration fails (e.g. duplicate name), free the tool we just
+        // built and propagate the error.
+        registry.register(tool) catch |err| {
+            tool.vtable.deinit(tool.ctx, allocator);
+            return err;
+        };
+    }
+    return regs.len;
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+const testing = std.testing;
+const Io = std.Io;
+
+fn writeTempScript(dir: Io.Dir, name: []const u8, source: []const u8) ![]const u8 {
+    try dir.writeFile(testing.io, .{ .sub_path = name, .data = source });
+    // Construct an absolute path so luaL_loadfilex finds it regardless of cwd.
+    var buf: [std.fs.max_path_bytes]u8 = undefined;
+    const n = try dir.realPathFile(testing.io, name, &buf);
+    return testing.allocator.dupe(u8, buf[0..n]);
+}
+
+test "loadExtension registers tools and invoke runs the handler" {
+    var tmp = testing.tmpDir(.{});
+    defer tmp.cleanup();
+
+    const source =
+        \\panto.register_tool {
+        \\  name = "greet", description = "Says hi.",
+        \\  schema = { type = "object", properties = { name = { type = "string" } } },
+        \\  handler = function(input) return "hi, " .. input.name end,
+        \\}
+    ;
+    const path = try writeTempScript(tmp.dir, "greet.lua", source);
+    defer testing.allocator.free(path);
+
+    var registry = panto.ToolRegistry.init(testing.allocator);
+    defer registry.deinit();
+
+    const n = try loadExtension(testing.allocator, &registry, path);
+    try testing.expectEqual(@as(usize, 1), n);
+
+    const tool = registry.lookup("greet") orelse return error.NotRegistered;
+    try testing.expectEqualStrings("greet", tool.name);
+    try testing.expectEqualStrings("Says hi.", tool.description);
+    try testing.expect(std.mem.indexOf(u8, tool.schema_json, "\"object\"") != null);
+
+    // Invoke through the vtable.
+    const result = try tool.vtable.invoke(tool.ctx, "{\"name\":\"travis\"}", testing.allocator);
+    defer testing.allocator.free(result);
+    try testing.expectEqualStrings("hi, travis", result);
+}
+
+test "invoke surfaces handler crashes as LuaHandlerCrashed" {
+    var tmp = testing.tmpDir(.{});
+    defer tmp.cleanup();
+
+    const source =
+        \\panto.register_tool {
+        \\  name = "boom", description = "crashes",
+        \\  schema = { type = "object" },
+        \\  handler = function(input) error("kaboom") end,
+        \\}
+    ;
+    const path = try writeTempScript(tmp.dir, "boom.lua", source);
+    defer testing.allocator.free(path);
+
+    var registry = panto.ToolRegistry.init(testing.allocator);
+    defer registry.deinit();
+
+    _ = try loadExtension(testing.allocator, &registry, path);
+    const tool = registry.lookup("boom") orelse return error.NotRegistered;
+
+    const result = tool.vtable.invoke(tool.ctx, "{}", testing.allocator);
+    try testing.expectError(error.LuaHandlerCrashed, result);
+}
+
+test "concurrent invoke: each call gets its own Lua state" {
+    var tmp = testing.tmpDir(.{});
+    defer tmp.cleanup();
+
+    // The handler returns the pointer-as-hex of `_G`, which differs between
+    // distinct Lua states. If two threads share a state, two of the four
+    // calls will return the same string.
+    const source =
+        \\panto.register_tool {
+        \\  name = "whoami", description = "state id",
+        \\  schema = { type = "object" },
+        \\  handler = function(input) return tostring(_G) end,
+        \\}
+    ;
+    const path = try writeTempScript(tmp.dir, "whoami.lua", source);
+    defer testing.allocator.free(path);
+
+    var registry = panto.ToolRegistry.init(testing.allocator);
+    defer registry.deinit();
+    _ = try loadExtension(testing.allocator, &registry, path);
+    const tool = registry.lookup("whoami") orelse return error.NotRegistered;
+
+    const Worker = struct {
+        tool: *const panto.Tool,
+        out: *[]u8,
+        err: *?anyerror,
+
+        fn run(self: @This()) void {
+            const r = self.tool.vtable.invoke(self.tool.ctx, "{}", testing.allocator) catch |e| {
+                self.err.* = e;
+                return;
+            };
+            self.out.* = r;
+        }
+    };
+
+    var results: [4][]u8 = .{ undefined, undefined, undefined, undefined };
+    var errs: [4]?anyerror = .{ null, null, null, null };
+    var threads: [4]std.Thread = undefined;
+    for (&threads, 0..) |*t, i| {
+        t.* = try std.Thread.spawn(.{}, Worker.run, .{Worker{
+            .tool = tool,
+            .out = &results[i],
+            .err = &errs[i],
+        }});
+    }
+    for (&threads) |t| t.join();
+    defer for (results) |r| if (r.len != 0) testing.allocator.free(r);
+
+    for (errs) |e| try testing.expectEqual(@as(?anyerror, null), e);
+
+    // All four "_G" identifiers should be distinct, proving distinct states.
+    for (0..4) |i| {
+        for ((i + 1)..4) |j| {
+            try testing.expect(!std.mem.eql(u8, results[i], results[j]));
+        }
+    }
+}