path: root/src/lua_tool.zig

//! `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,
    /// Optional directory to prepend to `package.path` so the script can
    /// `require` sibling files. Owned, NUL-terminated. Null for single-file
    /// extensions.
    package_root_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,
        package_root: ?[]const u8,
    ) !panto.Tool {
        const self = try allocator.create(LuaTool);
        errdefer allocator.destroy(self);

        const script_path_z = try allocator.dupeZ(u8, script_path);
        errdefer allocator.free(script_path_z);
        const name_owned = try allocator.dupe(u8, name);
        errdefer allocator.free(name_owned);
        const description_owned = try allocator.dupe(u8, description);
        errdefer allocator.free(description_owned);
        const schema_owned = try allocator.dupe(u8, schema_json);
        errdefer allocator.free(schema_owned);
        const package_root_z: ?[:0]u8 = if (package_root) |p|
            try allocator.dupeZ(u8, p)
        else
            null;

        self.* = .{
            .allocator = allocator,
            .script_path_z = script_path_z,
            .package_root_z = package_root_z,
            .name_owned = name_owned,
            .description_owned = description_owned,
            .schema_owned = schema_owned,
        };

        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);
        if (self.package_root_z) |p| a.free(p);
        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);

    if (self.package_root_z) |root| {
        try prependPackagePath(L, root);
    }

    // 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 {
    return loadExtensionImpl(allocator, registry, script_path, null);
}

/// Like `loadExtension`, but extends `package.path` with `<package_root>/?.lua`
/// and `<package_root>/?/init.lua` so the script can `require` sibling files.
/// `package_root` should be the directory containing the script.
pub fn loadExtensionWithPackageRoot(
    allocator: Allocator,
    registry: *panto.ToolRegistry,
    script_path: []const u8,
    package_root: []const u8,
) !usize {
    return loadExtensionImpl(allocator, registry, script_path, package_root);
}

fn loadExtensionImpl(
    allocator: Allocator,
    registry: *panto.ToolRegistry,
    script_path: []const u8,
    package_root: ?[]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);

    // Configure require() for directory-style extensions before running.
    if (package_root) |root| {
        const root_z = try allocator.dupeZ(u8, root);
        defer allocator.free(root_z);
        try prependPackagePath(L, root_z);
    }

    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,
            package_root,
        );
        // 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;
}

/// Prepend `<root>/?.lua` and `<root>/?/init.lua` to `package.path` so that
/// `require("foo")` resolves against files next to the extension's
/// `init.lua`. We prepend (not replace) so the standard library remains
/// available.
///
/// Stack effect: net zero.
fn prependPackagePath(L: *c.lua_State, root: [:0]const u8) !void {
    // Run a tiny Lua snippet rather than fiddle with the stack manually:
    // string concatenation is so much nicer in Lua.
    const snippet =
        \\local root = ...
        \\package.path = root .. "/?.lua;" .. root .. "/?/init.lua;" .. package.path
    ;
    if (c.luaL_loadstring(L, snippet) != 0) {
        logTopAsError(L, "lua: package.path loader failed to compile");
        return error.LuaPackagePathLoadFailed;
    }
    _ = c.lua_pushlstring(L, root.ptr, root.len);
    if (c.lua_pcallk(L, 1, 0, 0, 0, null) != 0) {
        logTopAsError(L, "lua: package.path setup failed");
        return error.LuaPackagePathSetupFailed;
    }
}

// ---------------------------------------------------------------------------
// 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]));
        }
    }
}