path: root/libpanto/src/tool_registry.zig

//! Registry of `Tool`s owned by an `Agent`.
//!
//! Tools are keyed by name. The registry takes ownership: on `unregister`
//! or `deinit`, it calls the tool's `vtable.deinit`.
//!
//! Iteration is not synchronized — callers must avoid mutating the registry
//! during iteration. In the current agent loop this is naturally true: the
//! provider iterates the registry 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 = tool_mod.Tool;

pub const ToolRegistry = struct {
    tools: std.StringHashMap(Tool),
    allocator: Allocator,

    pub fn init(allocator: Allocator) ToolRegistry {
        return .{
            .tools = std.StringHashMap(Tool).init(allocator),
            .allocator = allocator,
        };
    }

    /// Tear down the registry. Each remaining tool's `vtable.deinit` is
    /// invoked.
    pub fn deinit(self: *ToolRegistry) void {
        var it = self.tools.iterator();
        while (it.next()) |entry| {
            const tool = entry.value_ptr.*;
            tool.vtable.deinit(tool.ctx, self.allocator);
        }
        self.tools.deinit();
    }

    /// Register a tool. The registry takes ownership.
    ///
    /// Returns `error.DuplicateTool` if a tool with the same name is
    /// already registered — in that 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 {
        const gop = try self.tools.getOrPut(tool.name);
        if (gop.found_existing) return error.DuplicateTool;
        gop.value_ptr.* = tool;
    }

    /// Remove a tool by name. Calls the tool's `vtable.deinit`. No-op if
    /// the name is not registered.
    pub fn unregister(self: *ToolRegistry, name: []const u8) void {
        if (self.tools.fetchRemove(name)) |kv| {
            kv.value.vtable.deinit(kv.value.ctx, self.allocator);
        }
    }

    /// Look up a tool by name. The returned pointer is invalidated by any
    /// subsequent register/unregister call.
    pub fn lookup(self: *const ToolRegistry, name: []const u8) ?*const Tool {
        return self.tools.getPtr(name);
    }

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

    /// Iterate registered tools. Caller must not mutate the registry during
    /// iteration.
    pub fn iterator(self: *const ToolRegistry) Iterator {
        return .{ .inner = self.tools.iterator() };
    }

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

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

// -----------------------------------------------------------------------------
// 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 .{
            .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![]u8 {
        const self: *TestTool = @ptrCast(@alignCast(ctx));
        self.invocations += 1;
        return try allocator.dupe(u8, 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);
    }
};

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").?.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.name, "a")) saw_a = true;
        if (std.mem.eql(u8, t.name, "b")) saw_b = true;
        if (std.mem.eql(u8, t.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();
}