path: root/libpanto/src/tool_registry.zig

//! Registry of tools owned by an `Agent`.
//!
//! Two kinds of registration coexist:
//!
//!   - A single `Tool`: a thread-safe, self-contained handler. The
//!     registry holds one entry keyed by `tool.decl.name`.
//!   - A `ToolSource`: a batch-dispatched runtime that owns many tools.
//!     The registry holds one entry per declared tool, all pointing back
//!     at the same source (different `tool_index` per entry).
//!
//! Iteration yields the per-tool metadata as a uniform `ToolView` so
//! callers (chiefly: provider request serializers) don't need to know
//! which flavor of registration each tool came from.
//!
//! Iteration is not synchronized — callers must avoid mutating the
//! registry during iteration. In the current agent loop this is naturally
//! true: the provider iterates once at request-build time, and tool
//! registration only happens at agent setup.

const std = @import("std");
const Allocator = std.mem.Allocator;
const tool_mod = @import("tool.zig");
const tool_source_mod = @import("tool_source.zig");

const Tool = tool_mod.Tool;
const ToolSource = tool_source_mod.ToolSource;
const ToolDecl = tool_source_mod.ToolDecl;

// ===========================================================================
// Wire-name encoding
// ===========================================================================
//
// Internally, tool names use dots for namespacing (`std.read`), which our
// glob-based allow/deny policies rely on. But the OpenAI and Anthropic
// tool-name grammars forbid dots: both require `^[a-zA-Z0-9_-]{1,128}$`.
//
// So names are translated at the wire boundary only: `.` <-> `__`. The
// mapping is a clean bijection because a literal `__` is forbidden in an
// internal name (enforced by `validateName` at registration). Everything
// inside libpanto keeps speaking dots; only the serializers (via
// `toolsForLLM`) and inbound dispatch (via `lookupLLM`) cross the boundary.

/// The largest a wire (LLM-facing) tool name may be, per the provider
/// grammars. We validate the *encoded* length against this so an encoded
/// name is always acceptable to both providers.
pub const max_wire_name_len = 128;

pub const NameError = error{
    /// Name is empty or its encoded form exceeds `max_wire_name_len`.
    NameTooLong,
    /// Name contains a literal `__` (reserved as the encoded form of `.`)
    /// or a character outside `[a-zA-Z0-9_.-]`.
    InvalidNameChar,
};

/// Validate an internal tool name. Permits `[a-zA-Z0-9_.-]` but forbids a
/// literal `__` (which would collide with an encoded `.`), and requires
/// the encoded form to be 1..=`max_wire_name_len` bytes. Each `.` expands
/// to two bytes when encoded, so the cap is checked against that.
pub fn validateName(name: []const u8) NameError!void {
    if (name.len == 0) return error.NameTooLong;
    var encoded_len: usize = 0;
    for (name, 0..) |ch, i| {
        const ok = (ch >= 'a' and ch <= 'z') or (ch >= 'A' and ch <= 'Z') or
            (ch >= '0' and ch <= '9') or ch == '_' or ch == '-' or ch == '.';
        if (!ok) return error.InvalidNameChar;
        // Reject a literal double underscore: it is reserved for `.`.
        if (ch == '_' and i + 1 < name.len and name[i + 1] == '_') return error.InvalidNameChar;
        encoded_len += if (ch == '.') 2 else 1;
    }
    if (encoded_len > max_wire_name_len) return error.NameTooLong;
}

/// Encode an internal name for the wire: `.` -> `__`. Writes into `buf`
/// (which must be at least `max_wire_name_len` bytes) and returns the
/// written slice. Names that passed `validateName` always fit.
pub fn encodeName(buf: []u8, name: []const u8) []const u8 {
    var w: usize = 0;
    for (name) |ch| {
        if (ch == '.') {
            buf[w] = '_';
            buf[w + 1] = '_';
            w += 2;
        } else {
            buf[w] = ch;
            w += 1;
        }
    }
    return buf[0..w];
}

/// Decode a wire name back to internal form: `__` -> `.`. Writes into
/// `buf` (at least `wire.len` bytes) and returns the written slice. The
/// decode is unambiguous because internal names never contain `__`.
pub fn decodeName(buf: []u8, wire: []const u8) []u8 {
    var r: usize = 0;
    var w: usize = 0;
    while (r < wire.len) {
        if (wire[r] == '_' and r + 1 < wire.len and wire[r + 1] == '_') {
            buf[w] = '.';
            w += 1;
            r += 2;
        } else {
            buf[w] = wire[r];
            w += 1;
            r += 1;
        }
    }
    return buf[0..w];
}

/// Tagged registry value. The registry stores one of these per *tool
/// name*. ToolSources expand to one entry per declared tool, each with a
/// distinct `tool_index`.
pub const Entry = union(enum) {
    single: Tool,
    source: SourceRef,

    pub const SourceRef = struct {
        source: *ToolSource,
        /// Index into `source.tools`.
        tool_index: usize,
    };
};

/// Read-only view of a tool's metadata, uniform across `Tool` and
/// `ToolSource` registrations. Returned by registry iteration and
/// lookup.
pub const ToolView = struct {
    decl: ToolDecl,
    /// Which entry this view came from. Carries enough information to
    /// dispatch the call (single Tool vs source-backed).
    entry: Entry,

    pub fn name(self: ToolView) []const u8 {
        return self.decl.name;
    }
};

pub const ToolRegistry = struct {
    /// Per-tool-name entries.
    entries: std.StringHashMap(Entry),
    /// Heap-allocated sources, kept in a list so `deinit` can tear each
    /// down exactly once even though many entries reference a single
    /// source.
    sources: std.array_list.Managed(*ToolSource),
    allocator: Allocator,

    pub fn init(allocator: Allocator) ToolRegistry {
        return .{
            .entries = std.StringHashMap(Entry).init(allocator),
            .sources = std.array_list.Managed(*ToolSource).init(allocator),
            .allocator = allocator,
        };
    }

    /// Tear down the registry. Each single `Tool`'s `vtable.deinit` is
    /// invoked once. Each `ToolSource`'s `vtable.deinit` is invoked once
    /// (not once per declared tool).
    pub fn deinit(self: *ToolRegistry) void {
        var it = self.entries.iterator();
        while (it.next()) |entry| {
            switch (entry.value_ptr.*) {
                .single => |t| t.vtable.deinit(t.ctx, self.allocator),
                .source => {},
            }
        }
        self.entries.deinit();

        for (self.sources.items) |src| {
            src.vtable.deinit(src.ctx, self.allocator);
            self.allocator.destroy(src);
        }
        self.sources.deinit();
    }

    /// Register a single tool. The registry takes ownership.
    ///
    /// Returns `error.DuplicateTool` if a tool with the same name is
    /// already registered (whether from a single Tool or from a source).
    /// In the duplicate case the caller's tool is NOT taken over; the
    /// caller is responsible for tearing it down.
    pub fn register(self: *ToolRegistry, tool: Tool) !void {
        try validateName(tool.decl.name);
        const gop = try self.entries.getOrPut(tool.decl.name);
        if (gop.found_existing) return error.DuplicateTool;
        gop.value_ptr.* = .{ .single = tool };
    }

    /// Register a tool source. The registry takes ownership of `src` —
    /// it is heap-copied into the registry's source list and freed at
    /// deinit.
    ///
    /// Returns `error.DuplicateTool` if any of the source's declared
    /// tools collides with an existing registration. On collision the
    /// source is NOT taken over (caller still owns it and must tear it
    /// down) and any tools that *had* been inserted before the collision
    /// are rolled back.
    pub fn registerSource(self: *ToolRegistry, src: ToolSource) !void {
        // First pass: validate names and check for any collision before
        // committing anything.
        for (src.tools) |decl| {
            try validateName(decl.name);
            if (self.entries.contains(decl.name)) return error.DuplicateTool;
        }

        // Allocate the persistent heap copy of the source. From this
        // point forward, on any failure we must free the allocation and
        // roll back any entries we inserted.
        const heap = try self.allocator.create(ToolSource);
        errdefer self.allocator.destroy(heap);
        heap.* = src;

        var inserted: usize = 0;
        errdefer {
            // Roll back any inserts we made before the failure.
            for (src.tools[0..inserted]) |decl| {
                _ = self.entries.remove(decl.name);
            }
        }

        for (src.tools, 0..) |decl, i| {
            const gop = try self.entries.getOrPut(decl.name);
            if (gop.found_existing) return error.DuplicateTool;
            gop.value_ptr.* = .{ .source = .{ .source = heap, .tool_index = i } };
            inserted = i + 1;
        }

        try self.sources.append(heap);
    }

    /// Remove a single-tool registration by name. Calls the tool's
    /// `vtable.deinit`. No-op if the name is not registered or if it
    /// belongs to a source (sources are removed as a unit; not yet
    /// exposed).
    pub fn unregister(self: *ToolRegistry, name: []const u8) void {
        const entry_ptr = self.entries.getPtr(name) orelse return;
        switch (entry_ptr.*) {
            .single => |t| {
                _ = self.entries.remove(name);
                t.vtable.deinit(t.ctx, self.allocator);
            },
            .source => {}, // ignore — sources tear down at registry deinit
        }
    }

    /// Look up a tool by name. Returns a uniform `ToolView`. Pointer
    /// invariants are the same as `std.StringHashMap.getPtr`: invalidated
    /// by subsequent register/unregister calls.
    pub fn lookup(self: *const ToolRegistry, name: []const u8) ?ToolView {
        const entry = self.entries.get(name) orelse return null;
        return makeView(entry);
    }

    pub fn count(self: *const ToolRegistry) usize {
        return self.entries.count();
    }

    pub fn iterator(self: *const ToolRegistry) Iterator {
        return .{ .inner = self.entries.iterator() };
    }

    pub const Iterator = struct {
        inner: std.StringHashMap(Entry).Iterator,

        pub fn next(self: *Iterator) ?ToolView {
            const entry = self.inner.next() orelse return null;
            return makeView(entry.value_ptr.*);
        }
    };

    /// Iterate tools with their names **wire-encoded** (`.` -> `__`) for
    /// the LLM. The yielded `ToolView.decl.name` borrows the iterator's
    /// internal buffer and is only valid until the next `next()` call;
    /// serializers consume it immediately, so this is safe. Description
    /// and schema are unchanged.
    pub fn toolsForLLM(self: *const ToolRegistry) LLMIterator {
        return .{ .inner = self.entries.iterator() };
    }

    pub const LLMIterator = struct {
        inner: std.StringHashMap(Entry).Iterator,
        name_buf: [max_wire_name_len]u8 = undefined,

        pub fn next(self: *LLMIterator) ?ToolView {
            const entry = self.inner.next() orelse return null;
            var view = makeView(entry.value_ptr.*);
            view.decl.name = encodeName(&self.name_buf, view.decl.name);
            return view;
        }
    };

    fn makeView(entry: Entry) ToolView {
        return switch (entry) {
            .single => |t| .{ .decl = t.decl, .entry = entry },
            .source => |sr| .{ .decl = sr.source.tools[sr.tool_index], .entry = entry },
        };
    }
};

// -----------------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------------

const testing = std.testing;

/// A trivial in-test Tool implementation backed by a single owned counter
/// allocation. Used to verify ownership/deinit behavior.
const TestTool = struct {
    invocations: u32 = 0,
    name_owned: []u8,
    desc_owned: []u8,
    schema_owned: []u8,

    fn create(allocator: Allocator, name: []const u8) !Tool {
        const self = try allocator.create(TestTool);
        errdefer allocator.destroy(self);

        const name_owned = try allocator.dupe(u8, name);
        errdefer allocator.free(name_owned);
        const desc_owned = try allocator.dupe(u8, "test tool");
        errdefer allocator.free(desc_owned);
        const schema_owned = try allocator.dupe(u8, "{}");
        errdefer allocator.free(schema_owned);

        self.* = .{
            .name_owned = name_owned,
            .desc_owned = desc_owned,
            .schema_owned = schema_owned,
        };
        return .{
            .decl = .{
                .name = self.name_owned,
                .description = self.desc_owned,
                .schema_json = self.schema_owned,
            },
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: Tool.VTable = .{
        .invoke = invoke,
        .deinit = deinit,
    };

    fn invoke(ctx: *anyopaque, input: []const u8, allocator: Allocator) anyerror!tool_mod.ResultParts {
        const self: *TestTool = @ptrCast(@alignCast(ctx));
        self.invocations += 1;
        return tool_mod.ResultParts.fromText(allocator, input);
    }

    fn deinit(ctx: *anyopaque, allocator: Allocator) void {
        const self: *TestTool = @ptrCast(@alignCast(ctx));
        allocator.free(self.name_owned);
        allocator.free(self.desc_owned);
        allocator.free(self.schema_owned);
        allocator.destroy(self);
    }
};

/// A minimal source backing N tools. Each tool name maps to a configured
/// response prefix; invoke_batch returns "<prefix>:<input>" for each
/// call. Tracks the batch sizes it was called with for inspection.
const TestSource = struct {
    name_owned: []u8,
    decls: []ToolDecl,
    /// Allocations backing every `decl`'s strings. Freed at deinit.
    allocations: std.array_list.Managed([]u8),
    batch_sizes: std.array_list.Managed(usize),
    allocator: Allocator,

    fn create(
        allocator: Allocator,
        source_name: []const u8,
        tool_names: []const []const u8,
    ) !ToolSource {
        const self = try allocator.create(TestSource);
        errdefer allocator.destroy(self);

        var allocations = std.array_list.Managed([]u8).init(allocator);
        errdefer {
            for (allocations.items) |s| allocator.free(s);
            allocations.deinit();
        }

        const name_owned = try allocator.dupe(u8, source_name);
        try allocations.append(name_owned);

        const decls = try allocator.alloc(ToolDecl, tool_names.len);
        errdefer allocator.free(decls);

        for (tool_names, 0..) |tn, i| {
            const n = try allocator.dupe(u8, tn);
            try allocations.append(n);
            const d = try allocator.dupe(u8, "test src tool");
            try allocations.append(d);
            const s = try allocator.dupe(u8, "{}");
            try allocations.append(s);
            decls[i] = .{ .name = n, .description = d, .schema_json = s };
        }

        self.* = .{
            .name_owned = name_owned,
            .decls = decls,
            .allocations = allocations,
            .batch_sizes = std.array_list.Managed(usize).init(allocator),
            .allocator = allocator,
        };

        return ToolSource{
            .name = self.name_owned,
            .tools = self.decls,
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: ToolSource.VTable = .{
        .invoke_batch = invokeBatch,
        .deinit = deinitSrc,
    };

    fn invokeBatch(
        ctx: *anyopaque,
        calls: []const tool_source_mod.Call,
        results: []tool_source_mod.CallResult,
        allocator: Allocator,
    ) anyerror!void {
        const self: *TestSource = @ptrCast(@alignCast(ctx));
        try self.batch_sizes.append(calls.len);
        for (calls, 0..) |call, i| {
            const buf = std.fmt.allocPrint(
                allocator,
                "{s}:{s}",
                .{ call.tool_name, call.input },
            ) catch |e| {
                results[i] = .{ .err = e };
                continue;
            };
            results[i] = .{
                .ok = tool_mod.ResultParts.fromTextOwned(allocator, buf) catch |e| {
                    results[i] = .{ .err = e };
                    continue;
                },
            };
        }
    }

    fn deinitSrc(ctx: *anyopaque, _: Allocator) void {
        const self: *TestSource = @ptrCast(@alignCast(ctx));
        for (self.allocations.items) |s| self.allocator.free(s);
        self.allocations.deinit();
        self.batch_sizes.deinit();
        self.allocator.free(self.decls);
        self.allocator.destroy(self);
    }
};

test "register, lookup, count" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "echo"));
    try reg.register(try TestTool.create(allocator, "ls"));

    try testing.expectEqual(@as(usize, 2), reg.count());
    try testing.expect(reg.lookup("echo") != null);
    try testing.expect(reg.lookup("ls") != null);
    try testing.expect(reg.lookup("missing") == null);
    try testing.expectEqualStrings("echo", reg.lookup("echo").?.decl.name);
}

test "duplicate registration returns error and leaves original in place" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "echo"));

    // The second tool isn't taken over on duplicate; tear it down ourselves.
    var dup = try TestTool.create(allocator, "echo");
    try testing.expectError(error.DuplicateTool, reg.register(dup));
    dup.vtable.deinit(dup.ctx, allocator);

    try testing.expectEqual(@as(usize, 1), reg.count());
}

test "unregister calls deinit and removes" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "tmp"));
    try testing.expectEqual(@as(usize, 1), reg.count());

    reg.unregister("tmp");
    try testing.expectEqual(@as(usize, 0), reg.count());
    try testing.expect(reg.lookup("tmp") == null);

    // No-op on missing.
    reg.unregister("never_existed");
}

test "iterator visits every tool" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "a"));
    try reg.register(try TestTool.create(allocator, "b"));
    try reg.register(try TestTool.create(allocator, "c"));

    var saw_a = false;
    var saw_b = false;
    var saw_c = false;

    var it = reg.iterator();
    while (it.next()) |t| {
        if (std.mem.eql(u8, t.decl.name, "a")) saw_a = true;
        if (std.mem.eql(u8, t.decl.name, "b")) saw_b = true;
        if (std.mem.eql(u8, t.decl.name, "c")) saw_c = true;
    }
    try testing.expect(saw_a and saw_b and saw_c);
}

test "deinit frees all remaining tools" {
    // If this leaks, the testing allocator will catch it.
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    try reg.register(try TestTool.create(allocator, "x"));
    try reg.register(try TestTool.create(allocator, "y"));
    reg.deinit();
}

test "registerSource exposes every declared tool by name" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    const src = try TestSource.create(allocator, "panto-lua", &.{ "alpha", "beta", "gamma" });
    try reg.registerSource(src);

    try testing.expectEqual(@as(usize, 3), reg.count());
    const v = reg.lookup("beta") orelse return error.NotFound;
    try testing.expectEqualStrings("beta", v.decl.name);
    try testing.expect(v.entry == .source);
}

test "registerSource: collision with existing single tool aborts and rolls back" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "shared"));

    // Build a source that includes the colliding name. We must tear it
    // down ourselves on failure.
    var src = try TestSource.create(allocator, "src", &.{ "first", "shared", "third" });
    try testing.expectError(error.DuplicateTool, reg.registerSource(src));
    src.vtable.deinit(src.ctx, allocator);

    // No partial state from the source remains.
    try testing.expectEqual(@as(usize, 1), reg.count());
    try testing.expect(reg.lookup("first") == null);
    try testing.expect(reg.lookup("third") == null);
}

test "registerSource: collision between two sources" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.registerSource(try TestSource.create(allocator, "a", &.{ "foo", "bar" }));

    var s = try TestSource.create(allocator, "b", &.{ "baz", "foo" });
    try testing.expectError(error.DuplicateTool, reg.registerSource(s));
    s.vtable.deinit(s.ctx, allocator);

    try testing.expectEqual(@as(usize, 2), reg.count());
}

test "source view exposes per-tool metadata uniformly" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.registerSource(try TestSource.create(allocator, "lua", &.{ "x", "y" }));
    try reg.register(try TestTool.create(allocator, "z"));

    try testing.expectEqual(@as(usize, 3), reg.count());

    // Every entry has the canonical fields populated.
    var it = reg.iterator();
    var n: usize = 0;
    while (it.next()) |v| : (n += 1) {
        try testing.expect(v.decl.name.len > 0);
        try testing.expect(v.decl.description.len > 0);
        try testing.expect(v.decl.schema_json.len > 0);
    }
    try testing.expectEqual(@as(usize, 3), n);
}

// --- wire-name encoding ---

test "validateName: accepts dotted names, rejects literal __ and bad chars" {
    try validateName("std.read");
    try validateName("pkg.read_file");
    try validateName("a-b_c.d");
    try testing.expectError(error.InvalidNameChar, validateName("std__read"));
    try testing.expectError(error.InvalidNameChar, validateName("has space"));
    try testing.expectError(error.InvalidNameChar, validateName("slash/name"));
    try testing.expectError(error.NameTooLong, validateName(""));
    // 64 dots -> 128 encoded bytes: OK; 65 -> 130: too long.
    try validateName("." ** 64);
    try testing.expectError(error.NameTooLong, validateName("." ** 65));
}

test "encode/decode: dots <-> double underscores, bijective" {
    var buf: [max_wire_name_len]u8 = undefined;
    var buf2: [max_wire_name_len]u8 = undefined;

    const cases = [_][]const u8{ "std.read", "pkg.read_file", "a.b.c", "plain", "a-b" };
    inline for (cases) |internal| {
        const wire = encodeName(&buf, internal);
        try testing.expect(std.mem.indexOf(u8, wire, ".") == null);
        const back = decodeName(&buf2, wire);
        try testing.expectEqualStrings(internal, back);
    }

    // Spot-check the exact wire form and the read_file distinction.
    try testing.expectEqualStrings("std__read", encodeName(&buf, "std.read"));
    try testing.expectEqualStrings("pkg__read_file", encodeName(&buf, "pkg.read_file"));
    try testing.expectEqualStrings("pkg__read__file", encodeName(&buf, "pkg.read.file"));
}

test "register rejects names with literal double underscore" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    var bad = try TestTool.create(allocator, "std__read");
    try testing.expectError(error.InvalidNameChar, reg.register(bad));
    // Registration refused ownership; tear the tool down ourselves.
    bad.vtable.deinit(bad.ctx, allocator);
}

test "toolsForLLM yields wire-encoded names; iterator keeps dotted names" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();
    try reg.register(try TestTool.create(allocator, "std.read"));

    var llm = reg.toolsForLLM();
    const v = llm.next().?;
    try testing.expectEqualStrings("std__read", v.decl.name);
    try testing.expect(llm.next() == null);

    // The internal iterator is unchanged.
    var it = reg.iterator();
    try testing.expectEqualStrings("std.read", it.next().?.decl.name);
}