diff options
Diffstat (limited to 'src/tool_registry.zig')
| -rw-r--r-- | src/tool_registry.zig | 663 |
1 files changed, 663 insertions, 0 deletions
diff --git a/src/tool_registry.zig b/src/tool_registry.zig new file mode 100644 index 0000000..430cda5 --- /dev/null +++ b/src/tool_registry.zig @@ -0,0 +1,663 @@ +//! 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); +} |