diff options
Diffstat (limited to 'libpanto/src/agent.zig')
| -rw-r--r-- | libpanto/src/agent.zig | 658 |
1 files changed, 650 insertions, 8 deletions
diff --git a/libpanto/src/agent.zig b/libpanto/src/agent.zig index 7ebc058..3e240d4 100644 --- a/libpanto/src/agent.zig +++ b/libpanto/src/agent.zig @@ -1,23 +1,665 @@ +//! The Agent owns the conversation-driving loop: provider streaming + +//! tool dispatch. +//! +//! In phase 1/2 this was a thin pass-through to the provider. In phase 3 +//! it grows the tool-call loop: after each provider streaming step, the +//! agent inspects the assistant message for ToolUse blocks, dispatches +//! the registered handlers (in parallel when there are multiple), and +//! appends a user message containing the ToolResult blocks back into the +//! conversation. The loop continues until a turn arrives with no ToolUse +//! blocks. + const std = @import("std"); -const provider = @import("provider.zig"); +const Allocator = std.mem.Allocator; +const Thread = std.Thread; + +const provider_mod = @import("provider.zig"); const conversation = @import("conversation.zig"); +const tool_mod = @import("tool.zig"); +const tool_registry_mod = @import("tool_registry.zig"); + +pub const Tool = tool_mod.Tool; +pub const ToolRegistry = tool_registry_mod.ToolRegistry; pub const Agent = struct { - provider: provider.Provider, - allocator: std.mem.Allocator, + provider: provider_mod.Provider, + allocator: Allocator, + registry: ToolRegistry, - pub fn init(allocator: std.mem.Allocator, prov: provider.Provider) Agent { + pub fn init(allocator: Allocator, prov: provider_mod.Provider) Agent { return .{ .provider = prov, .allocator = allocator, + .registry = ToolRegistry.init(allocator), }; } - pub fn runStep(self: Agent, conv: *conversation.Conversation, receiver: *provider.Receiver) !void { - try self.provider.streamStep(conv, receiver); + pub fn deinit(self: *Agent) void { + self.registry.deinit(); + self.provider.deinit(); + } + + /// Register a tool. The agent's registry takes ownership. + pub fn registerTool(self: *Agent, tool: Tool) !void { + try self.registry.register(tool); } - pub fn deinit(self: Agent) void { - self.provider.deinit(); + /// Remove a tool by name. No-op if not registered. + pub fn unregisterTool(self: *Agent, name: []const u8) void { + self.registry.unregister(name); + } + + /// Drive the conversation forward until the model stops calling tools. + /// + /// A single `runStep` invocation may call the provider multiple times + /// if the model chains tool calls. Each provider call streams a new + /// assistant message into `conv`; if that message contains ToolUse + /// blocks the agent dispatches them concurrently, appends a user + /// message of ToolResult blocks, and loops. The loop terminates when + /// the provider's most recent response has no ToolUse blocks. + pub fn runStep( + self: *Agent, + conv: *conversation.Conversation, + receiver: *provider_mod.Receiver, + ) !void { + while (true) { + try self.provider.streamStep(conv, &self.registry, receiver); + + const last = conv.messages.items[conv.messages.items.len - 1]; + std.debug.assert(last.role == .assistant); + + // Defense-in-depth: if the provider committed an assistant + // message with zero content blocks, something went wrong + // upstream that wasn't surfaced as a provider error (e.g. a + // mid-stream provider error that an older codepath swallowed, + // or a model that genuinely returned nothing). Either way the + // turn made no observable progress — surface it instead of + // silently dropping back to the prompt. + if (last.content.items.len == 0) return error.EmptyAssistantResponse; + + if (!hasToolUseBlock(last)) return; + + try self.dispatchToolCalls(conv, last); + // Loop: feed the ToolResult message back to the provider. + } + } + + /// Returns true if the message contains at least one ToolUse block. + fn hasToolUseBlock(msg: conversation.Message) bool { + for (msg.content.items) |block| { + if (block == .ToolUse) return true; + } + return false; + } + + /// Dispatch every ToolUse block in `assistant_msg` concurrently, then + /// append a single user Message containing all ToolResult blocks to + /// `conv` in the same order the tool calls appeared. + fn dispatchToolCalls( + self: *Agent, + conv: *conversation.Conversation, + assistant_msg: conversation.Message, + ) !void { + // Count tool uses for sizing. + var n: usize = 0; + for (assistant_msg.content.items) |block| { + if (block == .ToolUse) n += 1; + } + std.debug.assert(n > 0); + + const tasks = try self.allocator.alloc(ToolCallTask, n); + defer self.allocator.free(tasks); + + // Populate tasks. We borrow ID/name slices from the conversation — + // the assistant message stays in `conv` throughout dispatch, so + // these slices remain valid until we copy them into the new + // ToolResultBlock. + { + var i: usize = 0; + for (assistant_msg.content.items) |block| { + if (block != .ToolUse) continue; + const tu = block.ToolUse; + tasks[i] = .{ + .agent = self, + .tool_use_id = tu.id, + .tool_name = tu.name, + .input = tu.input.items, + .result = null, + .err = null, + }; + i += 1; + } + } + + // Spawn one thread per tool call. `std.Thread.spawn` is cheap + // (sub-millisecond on Linux/macOS) compared to typical tool + // latency, and `Tool.invoke` is contractually thread-safe, so we + // fan out without a pool. + const threads = try self.allocator.alloc(Thread, n); + defer self.allocator.free(threads); + + var spawned: usize = 0; + var joined = false; + errdefer { + // Join any in-flight threads so they don't outlive `tasks`. + if (!joined) for (threads[0..spawned]) |t| t.join(); + for (tasks) |*task| { + if (task.result) |r| self.allocator.free(r); + } + } + + for (tasks, 0..) |*task, idx| { + threads[idx] = try Thread.spawn(.{}, runToolTask, .{task}); + spawned += 1; + } + for (threads[0..spawned]) |t| t.join(); + joined = true; + + // Build the user ToolResult message. From here on we own all + // result byte slices; transfer them into ToolResultBlocks. + var content: std.ArrayList(conversation.ContentBlock) = .empty; + errdefer { + for (content.items) |*b| b.deinit(self.allocator); + content.deinit(self.allocator); + } + try content.ensureTotalCapacity(self.allocator, n); + + // If any task failed, prefer to abort the turn — but first move + // every successful result into a block so it gets freed by the + // standard cleanup path, and free errored ones eagerly (there + // are none to move). The errdefer above handles teardown. + var first_err: ?anyerror = null; + for (tasks) |*task| { + if (task.err) |e| { + first_err = e; + continue; + } + const result_bytes = task.result.?; + task.result = null; // ownership transferred below + + const id_copy = try self.allocator.dupe(u8, task.tool_use_id); + errdefer self.allocator.free(id_copy); + + var content_buf: conversation.TextualBlock = .empty; + errdefer content_buf.deinit(self.allocator); + try content_buf.appendSlice(self.allocator, result_bytes); + self.allocator.free(result_bytes); + + content.appendAssumeCapacity(.{ .ToolResult = .{ + .tool_use_id = id_copy, + .content = content_buf, + } }); + } + + if (first_err) |e| return e; + + // Wrap the ToolResult blocks into a user Message and append. + try conv.messages.append(self.allocator, .{ + .role = .user, + .content = content, + }); + } +}; + +/// Per-tool-call work item passed into a worker thread. +const ToolCallTask = struct { + agent: *Agent, + tool_use_id: []const u8, // borrowed from assistant_msg + tool_name: []const u8, // borrowed from assistant_msg + input: []const u8, // borrowed from assistant_msg + + /// Owned result bytes from `Tool.invoke`. Allocated with + /// `agent.allocator`. Transferred into a ToolResultBlock on success. + result: ?[]u8, + + /// If non-null, the tool failed and the turn must abort. + err: ?anyerror, +}; + +fn runToolTask(task: *ToolCallTask) void { + const tool = task.agent.registry.lookup(task.tool_name) orelse { + task.err = error.UnknownTool; + return; + }; + const out = tool.vtable.invoke(tool.ctx, task.input, task.agent.allocator) catch |e| { + task.err = e; + return; + }; + task.result = out; +} + +// ----------------------------------------------------------------------------- +// Tests +// ----------------------------------------------------------------------------- + +const testing = std.testing; + +/// A stub Provider that, on each call to `streamStep`, appends a +/// pre-canned assistant message to the conversation. Used to drive the +/// agent's tool-call loop without any HTTP plumbing. +const StubProvider = struct { + allocator: Allocator, + scripted: []const ScriptedTurn, + next: usize = 0, + + const ScriptedTurn = struct { + /// Blocks to append as the next assistant message. The producer + /// owns these — the stub clones them per turn so the conversation + /// can take ownership. + blocks: []const TestBlock, + }; + + const TestBlock = union(enum) { + Text: []const u8, + ToolUse: struct { + id: []const u8, + name: []const u8, + input: []const u8, + }, + }; + + fn provider(self: *StubProvider) provider_mod.Provider { + return .{ .ptr = self, .vtable = &vt }; + } + + const vt: provider_mod.ProviderVTable = .{ + .streamStep = vtStreamStep, + .deinit = vtDeinit, + }; + + fn vtStreamStep( + ptr: *anyopaque, + conv: *conversation.Conversation, + _: *const ToolRegistry, + _: *provider_mod.Receiver, + ) anyerror!void { + const self: *StubProvider = @ptrCast(@alignCast(ptr)); + if (self.next >= self.scripted.len) return error.NoMoreScriptedTurns; + const turn = self.scripted[self.next]; + self.next += 1; + + var blocks: std.ArrayList(conversation.ContentBlock) = .empty; + errdefer { + for (blocks.items) |*b| b.deinit(self.allocator); + blocks.deinit(self.allocator); + } + for (turn.blocks) |tb| { + switch (tb) { + .Text => |s| { + try blocks.append(self.allocator, .{ + .Text = try conversation.textualBlockFromSlice(self.allocator, s), + }); + }, + .ToolUse => |tu| { + const id = try self.allocator.dupe(u8, tu.id); + errdefer self.allocator.free(id); + const name = try self.allocator.dupe(u8, tu.name); + errdefer self.allocator.free(name); + var input_buf: conversation.TextualBlock = .empty; + errdefer input_buf.deinit(self.allocator); + try input_buf.appendSlice(self.allocator, tu.input); + try blocks.append(self.allocator, .{ .ToolUse = .{ + .id = id, + .name = name, + .input = input_buf, + } }); + }, + } + } + const moved = try blocks.toOwnedSlice(self.allocator); + defer self.allocator.free(moved); + try conv.addAssistantMessage(moved); + } + + fn vtDeinit(_: *anyopaque) void {} +}; + +/// Simple in-test tool: returns `prefix ++ input`. Used in dispatch tests. +const EchoTool = struct { + prefix_owned: []u8, + name_owned: []u8, + + fn create(allocator: Allocator, name: []const u8, prefix: []const u8) !Tool { + const self = try allocator.create(EchoTool); + errdefer allocator.destroy(self); + self.name_owned = try allocator.dupe(u8, name); + errdefer allocator.free(self.name_owned); + self.prefix_owned = try allocator.dupe(u8, prefix); + return .{ + .name = self.name_owned, + .description = "echo", + .schema_json = "{}", + .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: *EchoTool = @ptrCast(@alignCast(ctx)); + return try std.fmt.allocPrint(allocator, "{s}{s}", .{ self.prefix_owned, input }); + } + + fn deinit(ctx: *anyopaque, allocator: Allocator) void { + const self: *EchoTool = @ptrCast(@alignCast(ctx)); + allocator.free(self.name_owned); + allocator.free(self.prefix_owned); + allocator.destroy(self); + } +}; + +/// Tool that records the thread it ran on, then participates in a +/// rendezvous: every invocation must reach the barrier before any can +/// return. If dispatch is sequential, the first invocation would deadlock +/// (only one tool runs at a time, never reaching the threshold) — so this +/// test only passes when invocations run truly concurrently. +/// +/// The barrier is bounded by a spin-with-yield with a wall-time ceiling +/// of 5 seconds; failure to reach quorum surfaces as an `error.BarrierTimeout`. +const BarrierTool = struct { + name_owned: []u8, + barrier: *Barrier, + + const Barrier = struct { + target: u32, + arrived: std.atomic.Value(u32) = .init(0), + thread_ids: [4]std.atomic.Value(u64) = .{ + .init(0), .init(0), .init(0), .init(0), + }, + }; + + fn create(allocator: Allocator, name: []const u8, barrier: *Barrier) !Tool { + const self = try allocator.create(BarrierTool); + errdefer allocator.destroy(self); + self.name_owned = try allocator.dupe(u8, name); + self.barrier = barrier; + return .{ + .name = self.name_owned, + .description = "barrier", + .schema_json = "{}", + .ctx = self, + .vtable = &vt, + }; + } + + const vt: Tool.VTable = .{ .invoke = invoke, .deinit = deinit }; + + fn invoke(ctx: *anyopaque, _: []const u8, allocator: Allocator) anyerror![]u8 { + const self: *BarrierTool = @ptrCast(@alignCast(ctx)); + const arrived = self.barrier.arrived.fetchAdd(1, .acq_rel); + if (arrived < self.barrier.thread_ids.len) { + self.barrier.thread_ids[arrived].store(std.Thread.getCurrentId(), .release); + } + + // Spin-with-yield until everyone has arrived. ~5s ceiling at the + // typical yield granularity is plenty for a 3-way barrier; on a + // truly single-threaded dispatch this loop never resolves. + var i: usize = 0; + while (self.barrier.arrived.load(.acquire) < self.barrier.target) : (i += 1) { + if (i > 50_000) return error.BarrierTimeout; + std.Thread.yield() catch {}; + } + return try allocator.dupe(u8, "done"); + } + + fn deinit(ctx: *anyopaque, allocator: Allocator) void { + const self: *BarrierTool = @ptrCast(@alignCast(ctx)); + allocator.free(self.name_owned); + allocator.destroy(self); + } +}; + +/// Tool whose invoke always errors. Used to verify the turn aborts. +const FailingTool = struct { + name_owned: []u8, + + fn create(allocator: Allocator, name: []const u8) !Tool { + const self = try allocator.create(FailingTool); + errdefer allocator.destroy(self); + self.name_owned = try allocator.dupe(u8, name); + return .{ + .name = self.name_owned, + .description = "fails", + .schema_json = "{}", + .ctx = self, + .vtable = &vt, + }; + } + + const vt: Tool.VTable = .{ .invoke = invoke, .deinit = deinit }; + + fn invoke(_: *anyopaque, _: []const u8, _: Allocator) anyerror![]u8 { + return error.ToolExploded; + } + + fn deinit(ctx: *anyopaque, allocator: Allocator) void { + const self: *FailingTool = @ptrCast(@alignCast(ctx)); + allocator.free(self.name_owned); + allocator.destroy(self); + } +}; + +const NoopReceiver = struct { + fn make() provider_mod.Receiver { + return .{ .ptr = @constCast(@ptrCast(&dummy)), .vtable = &vt }; } + var dummy: u8 = 0; + const vt: provider_mod.ReceiverVTable = .{ + .onMessageStart = noop1, + .onBlockStart = noop2, + .onContentDelta = noop3, + .onBlockComplete = noop4, + .onMessageComplete = noop5, + .onError = noop6, + }; + fn noop1(_: *anyopaque, _: conversation.MessageRole) anyerror!void {} + fn noop2(_: *anyopaque, _: provider_mod.ContentBlockType, _: usize, _: ?provider_mod.BlockMeta) anyerror!void {} + fn noop3(_: *anyopaque, _: usize, _: []const u8) anyerror!void {} + fn noop4(_: *anyopaque, _: usize, _: conversation.ContentBlock) anyerror!void {} + fn noop5(_: *anyopaque, _: conversation.Message) anyerror!void {} + fn noop6(_: *anyopaque, _: anyerror) void {} }; + +test "registerTool and lookup via registry" { + var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} }; + var agent = Agent.init(testing.allocator, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "ECHO:")); + try testing.expectEqual(@as(usize, 1), agent.registry.count()); + try testing.expect(agent.registry.lookup("echo") != null); +} + +test "duplicate registerTool returns error" { + var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} }; + var agent = Agent.init(testing.allocator, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "A:")); + + var dup = try EchoTool.create(testing.allocator, "echo", "B:"); + try testing.expectError(error.DuplicateTool, agent.registerTool(dup)); + dup.vtable.deinit(dup.ctx, testing.allocator); +} + +test "runStep dispatches a tool call and loops to a final text turn" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "tc_1", .name = "echo", .input = "hello" } }, + } }, + .{ .blocks = &.{ + .{ .Text = "ok" }, + } }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try EchoTool.create(allocator, "echo", "ECHO:")); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("call a tool"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + // user, assistant(tool_use), user(tool_result), assistant(text) + try testing.expectEqual(@as(usize, 4), conv.messages.items.len); + + try testing.expectEqual(conversation.MessageRole.assistant, conv.messages.items[1].role); + try testing.expectEqual(@as(usize, 1), conv.messages.items[1].content.items.len); + try testing.expectEqualStrings("tc_1", conv.messages.items[1].content.items[0].ToolUse.id); + + try testing.expectEqual(conversation.MessageRole.user, conv.messages.items[2].role); + try testing.expectEqual(@as(usize, 1), conv.messages.items[2].content.items.len); + const tr = conv.messages.items[2].content.items[0].ToolResult; + try testing.expectEqualStrings("tc_1", tr.tool_use_id); + try testing.expectEqualStrings("ECHO:hello", tr.content.items); + + try testing.expectEqual(conversation.MessageRole.assistant, conv.messages.items[3].role); + try testing.expectEqualStrings("ok", conv.messages.items[3].content.items[0].Text.items); +} + +test "runStep dispatches multiple tool calls in parallel" { + const allocator = testing.allocator; + + // Use a barrier: each tool must wait until all three have arrived + // before returning. If dispatch were sequential, the first tool + // would hit its iteration ceiling and `error.BarrierTimeout`. Reaching + // the barrier proves all three ran concurrently. + var barrier: BarrierTool.Barrier = .{ .target = 3 }; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "a", .name = "barrierA", .input = "" } }, + .{ .ToolUse = .{ .id = "b", .name = "barrierB", .input = "" } }, + .{ .ToolUse = .{ .id = "c", .name = "barrierC", .input = "" } }, + } }, + .{ .blocks = &.{ + .{ .Text = "done" }, + } }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try BarrierTool.create(allocator, "barrierA", &barrier)); + try agent.registerTool(try BarrierTool.create(allocator, "barrierB", &barrier)); + try agent.registerTool(try BarrierTool.create(allocator, "barrierC", &barrier)); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("go"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + // Each tool produced one ToolResult, in original order. + const tr_msg = conv.messages.items[2]; + try testing.expectEqual(@as(usize, 3), tr_msg.content.items.len); + try testing.expectEqualStrings("a", tr_msg.content.items[0].ToolResult.tool_use_id); + try testing.expectEqualStrings("b", tr_msg.content.items[1].ToolResult.tool_use_id); + try testing.expectEqualStrings("c", tr_msg.content.items[2].ToolResult.tool_use_id); + + // And the three calls happened on three distinct threads. + const t0 = barrier.thread_ids[0].load(.acquire); + const t1 = barrier.thread_ids[1].load(.acquire); + const t2 = barrier.thread_ids[2].load(.acquire); + try testing.expect(t0 != 0 and t1 != 0 and t2 != 0); + try testing.expect(t0 != t1 and t1 != t2 and t0 != t2); +} + +test "runStep propagates tool errors and aborts the turn" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "x", .name = "boom", .input = "" } }, + } }, + // Second turn should never run. + .{ .blocks = &.{.{ .Text = "should-not-see" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try FailingTool.create(allocator, "boom")); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("break it"); + + var recv = NoopReceiver.make(); + try testing.expectError(error.ToolExploded, agent.runStep(&conv, &recv)); + + // Conversation has user + assistant(tool_use). No ToolResult message + // was appended because the dispatch errored before append. + try testing.expectEqual(@as(usize, 2), conv.messages.items.len); +} + +test "runStep errors UnknownTool when the model calls something unregistered" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "z", .name = "ghost", .input = "" } }, + } }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("call a ghost"); + + var recv = NoopReceiver.make(); + try testing.expectError(error.UnknownTool, agent.runStep(&conv, &recv)); +} + +test "runStep with no tool calls returns after one provider step" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{.{ .Text = "hi" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("hello"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + try testing.expectEqual(@as(usize, 2), conv.messages.items.len); + try testing.expectEqualStrings("hi", conv.messages.items[1].content.items[0].Text.items); +} + +test "runStep surfaces EmptyAssistantResponse when provider commits an empty message" { + // Mirrors the real-world failure mode where a provider silently ends the + // turn with no content blocks — e.g. a mid-stream error that an older + // codepath swallowed. The agent must surface the failure so the user + // doesn't see the prompt come back with no explanation. + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var agent = Agent.init(allocator, stub.provider()); + defer agent.deinit(); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("hi"); + + var recv = NoopReceiver.make(); + try testing.expectError(error.EmptyAssistantResponse, agent.runStep(&conv, &recv)); +} |