diff options
Diffstat (limited to 'libpanto/src/agent.zig')
| -rw-r--r-- | libpanto/src/agent.zig | 772 |
1 files changed, 623 insertions, 149 deletions
diff --git a/libpanto/src/agent.zig b/libpanto/src/agent.zig index 3e240d4..55396d1 100644 --- a/libpanto/src/agent.zig +++ b/libpanto/src/agent.zig @@ -1,35 +1,56 @@ //! 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. +//! On each turn, after the provider streams an assistant message, the +//! agent inspects it for ToolUse blocks. If any are present, the agent: +//! +//! 1. Groups them by their *owning registration* in the registry — a +//! single `Tool` is its own group; every `ToolSource`-backed tool +//! whose name maps to the same source forms one group. +//! 2. Spawns one concurrent task per group via `std.Io.Group`. +//! A single-`Tool` group runs the tool's `invoke` once; a +//! `ToolSource` group calls the source's `invoke_batch` with all +//! of its calls at once. We use `Group.concurrent` (not `async`) +//! because tool invocations may block on I/O and we need real +//! concurrency, not just expressed asynchrony. +//! 3. Awaits the group. ToolResult blocks are assembled in the +//! *original* call order (i.e. the order the LLM emitted them). +//! 4. Appends a user message containing the ToolResult blocks back +//! into the conversation and loops. +//! +//! The "thread-safe" promise for single `Tool` registrations is +//! unchanged. For `ToolSource`-backed tools, the source's runtime +//! receives all of its calls on one thread per turn, so it can keep a +//! single-threaded interpreter (Lua, Python, ...) without further +//! synchronization. const std = @import("std"); const Allocator = std.mem.Allocator; -const Thread = std.Thread; +const Io = std.Io; const provider_mod = @import("provider.zig"); const conversation = @import("conversation.zig"); const tool_mod = @import("tool.zig"); +const tool_source_mod = @import("tool_source.zig"); const tool_registry_mod = @import("tool_registry.zig"); pub const Tool = tool_mod.Tool; +pub const ToolSource = tool_source_mod.ToolSource; pub const ToolRegistry = tool_registry_mod.ToolRegistry; +const Entry = tool_registry_mod.Entry; + pub const Agent = struct { provider: provider_mod.Provider, allocator: Allocator, + io: Io, registry: ToolRegistry, - pub fn init(allocator: Allocator, prov: provider_mod.Provider) Agent { + pub fn init(allocator: Allocator, io: Io, prov: provider_mod.Provider) Agent { return .{ .provider = prov, .allocator = allocator, + .io = io, .registry = ToolRegistry.init(allocator), }; } @@ -39,24 +60,23 @@ pub const Agent = struct { self.provider.deinit(); } - /// Register a tool. The agent's registry takes ownership. + /// Register a single tool. The agent's registry takes ownership. pub fn registerTool(self: *Agent, tool: Tool) !void { try self.registry.register(tool); } - /// Remove a tool by name. No-op if not registered. + /// Register a tool source. The agent's registry takes ownership. + pub fn registerToolSource(self: *Agent, src: ToolSource) !void { + try self.registry.registerSource(src); + } + + /// Remove a tool by name. No-op if not registered or if the name + /// belongs to a source. 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, @@ -68,23 +88,17 @@ pub const Agent = struct { 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. + // Defense-in-depth: a provider that silently committed an + // empty assistant message means the turn made no observable + // progress. Surface it instead of looping 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; @@ -92,92 +106,97 @@ pub const Agent = struct { 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. + /// Dispatch every ToolUse block in `assistant_msg`. Groups by owning + /// registration; one OS thread per group; results assembled in the + /// original call order. fn dispatchToolCalls( self: *Agent, conv: *conversation.Conversation, assistant_msg: conversation.Message, ) !void { - // Count tool uses for sizing. - var n: usize = 0; + // Build the flat call list (in original order) and group calls + // by owning registration. + var calls: std.array_list.Managed(FlatCall) = .init(self.allocator); + defer calls.deinit(); + for (assistant_msg.content.items) |block| { - if (block == .ToolUse) n += 1; + if (block != .ToolUse) continue; + const tu = block.ToolUse; + const entry = self.registry.lookup(tu.name) orelse { + // Unknown tool: abort the turn with a clear error. + return error.UnknownTool; + }; + try calls.append(.{ + .tool_use_id = tu.id, + .tool_name = tu.name, + .input = tu.input.items, + .entry = entry.entry, + .result = null, + .err = null, + }); } - std.debug.assert(n > 0); - - const tasks = try self.allocator.alloc(ToolCallTask, n); - defer self.allocator.free(tasks); + std.debug.assert(calls.items.len > 0); - // 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; - } + // Partition into groups. A group's `kind` determines how it + // runs; the `member_indices` are positions into `calls` (the + // original call order) so we can write back results without + // re-ordering. + var groups: std.array_list.Managed(Group) = .init(self.allocator); + defer { + for (groups.items) |*g| g.deinit(self.allocator); + groups.deinit(); } + try buildGroups(self.allocator, calls.items, &groups); - // 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; + // Spawn one concurrent task per group via `std.Io.Group`. + // Single-tool groups run the tool's vtable; source groups run + // the source's `invoke_batch`. We use `concurrent` rather than + // `async` because tool work may block on I/O — under a + // single-threaded `Io` `async` would deadlock; `concurrent` + // forces real concurrency (or `error.ConcurrencyUnavailable`). + var task_group: Io.Group = .init; + // `cancel` is idempotent with `await`; if anything below this + // point errors before we successfully `await`, this releases + // the group's resources. + defer task_group.cancel(self.io); 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 (calls.items) |*c| { + if (c.result) |r| self.allocator.free(r); } } - for (tasks, 0..) |*task, idx| { - threads[idx] = try Thread.spawn(.{}, runToolTask, .{task}); - spawned += 1; + for (groups.items) |*g| { + try task_group.concurrent(self.io, runGroup, .{ self, g, calls.items }); } - for (threads[0..spawned]) |t| t.join(); - joined = true; + // `error.Canceled` here means cancellation propagated into this + // dispatch from above; surface it like any other error. + try task_group.await(self.io); - // Build the user ToolResult message. From here on we own all - // result byte slices; transfer them into ToolResultBlocks. + // Assemble ToolResult blocks in original call order. If any + // call errored, prefer to abort the turn — but only after the + // standard errdefer above has freed remaining results. 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); + try content.ensureTotalCapacity(self.allocator, calls.items.len); - // 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| { + for (calls.items) |*c| { + if (c.err) |e| { first_err = e; continue; } - const result_bytes = task.result.?; - task.result = null; // ownership transferred below + const result_bytes = c.result orelse { + // Internal error: every successful call should have left + // bytes behind. Treat as MissingToolResult. + first_err = error.MissingToolResult; + continue; + }; + c.result = null; // ownership transferred below - const id_copy = try self.allocator.dupe(u8, task.tool_use_id); + const id_copy = try self.allocator.dupe(u8, c.tool_use_id); errdefer self.allocator.free(id_copy); var content_buf: conversation.TextualBlock = .empty; @@ -193,7 +212,6 @@ pub const Agent = struct { 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, @@ -201,31 +219,158 @@ pub const Agent = struct { } }; -/// Per-tool-call work item passed into a worker thread. -const ToolCallTask = struct { - agent: *Agent, +/// One ToolUse, as flattened into the agent's dispatch list. `result` +/// and `err` are filled in by the worker; exactly one is non-null on +/// successful task completion. +const FlatCall = struct { tool_use_id: []const u8, // borrowed from assistant_msg tool_name: []const u8, // borrowed from assistant_msg input: []const u8, // borrowed from assistant_msg + entry: Entry, - /// Owned result bytes from `Tool.invoke`. Allocated with - /// `agent.allocator`. Transferred into a ToolResultBlock on success. + /// Owned result bytes from `Tool.invoke` or `ToolSource.invoke_batch`. + /// Allocated with the agent's allocator. Transferred into a + /// ToolResultBlock on success. result: ?[]u8, - /// If non-null, the tool failed and the turn must abort. + /// If non-null, the call 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; +/// One dispatch group. Either a single Tool invocation, or a batch of +/// calls headed to one ToolSource. +const Group = union(enum) { + single: SingleGroup, + source: SourceGroup, + + pub const SingleGroup = struct { + tool: Tool, + /// Index into the flat calls array. + call_index: usize, + }; + + pub const SourceGroup = struct { + source: *ToolSource, + /// Indices into the flat calls array. Owned by the group. + member_indices: []usize, + }; + + fn deinit(self: *Group, allocator: Allocator) void { + switch (self.*) { + .single => {}, + .source => |sg| allocator.free(sg.member_indices), + } + } +}; + +/// Partition the flat call list into groups. Order of groups is +/// arbitrary; order within a `source` group preserves the original +/// call order so that batch results can be written back positionally. +fn buildGroups( + allocator: Allocator, + calls: []const FlatCall, + out: *std.array_list.Managed(Group), +) !void { + // Map from source pointer to the index of its group in `out`. + // Buffers per source, accumulated then frozen into slices. + var pending: std.AutoHashMap(*ToolSource, std.array_list.Managed(usize)) = + .init(allocator); + defer { + var it = pending.valueIterator(); + while (it.next()) |l| l.deinit(); + pending.deinit(); + } + + for (calls, 0..) |c, i| { + switch (c.entry) { + .single => |t| try out.append(.{ .single = .{ .tool = t, .call_index = i } }), + .source => |sr| { + const gop = try pending.getOrPut(sr.source); + if (!gop.found_existing) { + gop.value_ptr.* = std.array_list.Managed(usize).init(allocator); + } + try gop.value_ptr.append(i); + }, + } + } + + // Freeze each pending list into a source-group entry. We move + // ownership of the indices into `Group.source.member_indices`. + var pit = pending.iterator(); + while (pit.next()) |entry| { + const src = entry.key_ptr.*; + const indices = try entry.value_ptr.toOwnedSlice(); + try out.append(.{ .source = .{ .source = src, .member_indices = indices } }); + } +} + +/// Worker entry point. Runs one group to completion, populating +/// `calls[i].result` or `calls[i].err` for each member call. +/// +/// Return type is `void`, which coerces to `Io.Cancelable!void` as +/// required by `Group.concurrent`. Tool errors are reported via +/// `FlatCall.err`, not by returning from this function. +fn runGroup(agent: *Agent, group: *Group, calls: []FlatCall) void { + switch (group.*) { + .single => |sg| { + const i = sg.call_index; + const c = &calls[i]; + const out = sg.tool.vtable.invoke(sg.tool.ctx, c.input, agent.allocator) catch |e| { + c.err = e; + return; + }; + c.result = out; + }, + .source => |sg| runSourceGroup(agent, sg, calls), + } +} + +fn runSourceGroup(agent: *Agent, sg: Group.SourceGroup, calls: []FlatCall) void { + const n = sg.member_indices.len; + + const batch_calls = agent.allocator.alloc(tool_source_mod.Call, n) catch |e| { + for (sg.member_indices) |i| calls[i].err = e; + return; + }; + defer agent.allocator.free(batch_calls); + + const batch_results = agent.allocator.alloc(tool_source_mod.CallResult, n) catch |e| { + for (sg.member_indices) |i| calls[i].err = e; return; }; - const out = tool.vtable.invoke(tool.ctx, task.input, task.agent.allocator) catch |e| { - task.err = e; + defer agent.allocator.free(batch_results); + + for (sg.member_indices, 0..) |idx, j| { + batch_calls[j] = .{ + .tool_name = calls[idx].tool_name, + .input = calls[idx].input, + }; + batch_results[j] = .{ .err = error.SourceDroppedCall }; + } + + sg.source.vtable.invoke_batch( + sg.source.ctx, + batch_calls, + batch_results, + agent.allocator, + ) catch |e| { + // Whole-batch failure: free any partial successes the source + // already wrote, then mark every member as failed. + for (batch_results) |r| switch (r) { + .ok => |b| agent.allocator.free(b), + .err => {}, + }; + for (sg.member_indices) |i| calls[i].err = e; return; }; - task.result = out; + + // Per-call success/error. + for (sg.member_indices, 0..) |i, j| { + switch (batch_results[j]) { + .ok => |b| calls[i].result = b, + .err => |e| calls[i].err = e, + } + } } // ----------------------------------------------------------------------------- @@ -234,18 +379,12 @@ fn runToolTask(task: *ToolCallTask) void { 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, }; @@ -314,7 +453,6 @@ const StubProvider = struct { fn vtDeinit(_: *anyopaque) void {} }; -/// Simple in-test tool: returns `prefix ++ input`. Used in dispatch tests. const EchoTool = struct { prefix_owned: []u8, name_owned: []u8, @@ -326,9 +464,11 @@ const EchoTool = struct { errdefer allocator.free(self.name_owned); self.prefix_owned = try allocator.dupe(u8, prefix); return .{ - .name = self.name_owned, - .description = "echo", - .schema_json = "{}", + .decl = .{ + .name = self.name_owned, + .description = "echo", + .schema_json = "{}", + }, .ctx = self, .vtable = &vt, }; @@ -349,14 +489,6 @@ const EchoTool = struct { } }; -/// 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, @@ -375,9 +507,11 @@ const BarrierTool = struct { self.name_owned = try allocator.dupe(u8, name); self.barrier = barrier; return .{ - .name = self.name_owned, - .description = "barrier", - .schema_json = "{}", + .decl = .{ + .name = self.name_owned, + .description = "barrier", + .schema_json = "{}", + }, .ctx = self, .vtable = &vt, }; @@ -392,9 +526,6 @@ const BarrierTool = struct { 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; @@ -410,7 +541,6 @@ const BarrierTool = struct { } }; -/// Tool whose invoke always errors. Used to verify the turn aborts. const FailingTool = struct { name_owned: []u8, @@ -419,9 +549,11 @@ const FailingTool = struct { errdefer allocator.destroy(self); self.name_owned = try allocator.dupe(u8, name); return .{ - .name = self.name_owned, - .description = "fails", - .schema_json = "{}", + .decl = .{ + .name = self.name_owned, + .description = "fails", + .schema_json = "{}", + }, .ctx = self, .vtable = &vt, }; @@ -461,9 +593,189 @@ const NoopReceiver = struct { fn noop6(_: *anyopaque, _: anyerror) void {} }; +/// A configurable ToolSource for testing the grouped-dispatch path. +/// Stores every batch it receives so tests can assert "calls X and Y +/// arrived in the same batch on the same thread". +const TestSource = struct { + name_owned: []u8, + decls: []tool_source_mod.ToolDecl, + decl_strings: std.array_list.Managed([]u8), + /// Sequence of (thread_id, [tool_name; n]) per batch received. + /// Only mutated inside `invoke_batch`. Because libpanto guarantees + /// at most one outstanding `invoke_batch` per source at any time + /// (one batch per turn per source), no synchronization is needed. + batches: std.array_list.Managed(Batch), + allocator: Allocator, + + const Batch = struct { + thread_id: u64, + names: std.array_list.Managed([]u8), + }; + + 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 strings = std.array_list.Managed([]u8).init(allocator); + errdefer { + for (strings.items) |s| allocator.free(s); + strings.deinit(); + } + + const name_owned = try allocator.dupe(u8, source_name); + try strings.append(name_owned); + + const decls = try allocator.alloc(tool_source_mod.ToolDecl, tool_names.len); + errdefer allocator.free(decls); + for (tool_names, 0..) |tn, i| { + const n = try allocator.dupe(u8, tn); + try strings.append(n); + const d = try allocator.dupe(u8, "test src tool"); + try strings.append(d); + const s = try allocator.dupe(u8, "{}"); + try strings.append(s); + decls[i] = .{ .name = n, .description = d, .schema_json = s }; + } + + self.* = .{ + .name_owned = name_owned, + .decls = decls, + .decl_strings = strings, + .batches = std.array_list.Managed(Batch).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)); + var batch: Batch = .{ + .thread_id = std.Thread.getCurrentId(), + .names = std.array_list.Managed([]u8).init(self.allocator), + }; + for (calls) |c| { + const copy = try self.allocator.dupe(u8, c.tool_name); + try batch.names.append(copy); + } + try self.batches.append(batch); + + for (calls, 0..) |c, i| { + results[i] = .{ + .ok = std.fmt.allocPrint( + allocator, + "{s}->{s}", + .{ c.tool_name, c.input }, + ) catch |e| { + results[i] = .{ .err = e }; + continue; + }, + }; + } + } + + fn deinitSrc(ctx: *anyopaque, _: Allocator) void { + const self: *TestSource = @ptrCast(@alignCast(ctx)); + for (self.decl_strings.items) |s| self.allocator.free(s); + self.decl_strings.deinit(); + for (self.batches.items) |*b| { + for (b.names.items) |n| self.allocator.free(n); + b.names.deinit(); + } + self.batches.deinit(); + self.allocator.free(self.decls); + self.allocator.destroy(self); + } +}; + +/// A source that always fails the whole batch by returning an error +/// from invoke_batch (rather than recording per-call errors). Used to +/// verify libpanto's whole-batch-failure path. +const FailingSource = struct { + name_owned: []u8, + decls: []tool_source_mod.ToolDecl, + decl_strings: std.array_list.Managed([]u8), + allocator: Allocator, + + fn create(allocator: Allocator, source_name: []const u8, tool_names: []const []const u8) !ToolSource { + const self = try allocator.create(FailingSource); + errdefer allocator.destroy(self); + + var strings = std.array_list.Managed([]u8).init(allocator); + errdefer { + for (strings.items) |s| allocator.free(s); + strings.deinit(); + } + + const name_owned = try allocator.dupe(u8, source_name); + try strings.append(name_owned); + + const decls = try allocator.alloc(tool_source_mod.ToolDecl, tool_names.len); + errdefer allocator.free(decls); + for (tool_names, 0..) |tn, i| { + const n = try allocator.dupe(u8, tn); + try strings.append(n); + const d = try allocator.dupe(u8, "fails"); + try strings.append(d); + const s = try allocator.dupe(u8, "{}"); + try strings.append(s); + decls[i] = .{ .name = n, .description = d, .schema_json = s }; + } + + self.* = .{ + .name_owned = name_owned, + .decls = decls, + .decl_strings = strings, + .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( + _: *anyopaque, + _: []const tool_source_mod.Call, + _: []tool_source_mod.CallResult, + _: Allocator, + ) anyerror!void { + return error.SourceExploded; + } + + fn deinitSrc(ctx: *anyopaque, _: Allocator) void { + const self: *FailingSource = @ptrCast(@alignCast(ctx)); + for (self.decl_strings.items) |s| self.allocator.free(s); + self.decl_strings.deinit(); + self.allocator.free(self.decls); + self.allocator.destroy(self); + } +}; + test "registerTool and lookup via registry" { var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} }; - var agent = Agent.init(testing.allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(testing.allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(testing.allocator, io, stub.provider()); defer agent.deinit(); try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "ECHO:")); @@ -473,7 +785,10 @@ test "registerTool and lookup via registry" { test "duplicate registerTool returns error" { var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} }; - var agent = Agent.init(testing.allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(testing.allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(testing.allocator, io, stub.provider()); defer agent.deinit(); try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "A:")); @@ -495,7 +810,10 @@ test "runStep dispatches a tool call and loops to a final text turn" { } }, }; var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; - var agent = Agent.init(allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); try agent.registerTool(try EchoTool.create(allocator, "echo", "ECHO:")); @@ -507,7 +825,6 @@ test "runStep dispatches a tool call and loops to a final text turn" { 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); @@ -527,10 +844,6 @@ test "runStep dispatches a tool call and loops to a final text turn" { 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{ @@ -544,7 +857,10 @@ test "runStep dispatches multiple tool calls in parallel" { } }, }; var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; - var agent = Agent.init(allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); try agent.registerTool(try BarrierTool.create(allocator, "barrierA", &barrier)); @@ -558,14 +874,12 @@ test "runStep dispatches multiple tool calls in parallel" { 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); @@ -580,11 +894,13 @@ test "runStep propagates tool errors and aborts the turn" { .{ .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()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); try agent.registerTool(try FailingTool.create(allocator, "boom")); @@ -596,8 +912,6 @@ test "runStep propagates tool errors and aborts the turn" { 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); } @@ -610,7 +924,10 @@ test "runStep errors UnknownTool when the model calls something unregistered" { } }, }; var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; - var agent = Agent.init(allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); var conv = conversation.Conversation.init(allocator); @@ -628,7 +945,10 @@ test "runStep with no tool calls returns after one provider step" { .{ .blocks = &.{.{ .Text = "hi" }} }, }; var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; - var agent = Agent.init(allocator, stub.provider()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); var conv = conversation.Conversation.init(allocator); @@ -643,17 +963,16 @@ test "runStep with no tool calls returns after one provider step" { } 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()); + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); defer agent.deinit(); var conv = conversation.Conversation.init(allocator); @@ -663,3 +982,158 @@ test "runStep surfaces EmptyAssistantResponse when provider commits an empty mes var recv = NoopReceiver.make(); try testing.expectError(error.EmptyAssistantResponse, agent.runStep(&conv, &recv)); } + +// ------------ ToolSource tests ------------ + +test "runStep delivers all source-backed calls in one batch on one thread" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "a", .name = "lua_x", .input = "1" } }, + .{ .ToolUse = .{ .id = "b", .name = "lua_y", .input = "2" } }, + .{ .ToolUse = .{ .id = "c", .name = "lua_x", .input = "3" } }, + } }, + .{ .blocks = &.{.{ .Text = "done" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); + defer agent.deinit(); + + try agent.registerToolSource(try TestSource.create(allocator, "panto-lua", &.{ "lua_x", "lua_y" })); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("go"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + // Locate the source and inspect its observed batches. + const view = agent.registry.lookup("lua_x") orelse return error.NotFound; + const src_ptr = view.entry.source.source; + const test_src: *TestSource = @ptrCast(@alignCast(src_ptr.ctx)); + + try testing.expectEqual(@as(usize, 1), test_src.batches.items.len); + const b = test_src.batches.items[0]; + try testing.expectEqual(@as(usize, 3), b.names.items.len); + try testing.expectEqualStrings("lua_x", b.names.items[0]); + try testing.expectEqualStrings("lua_y", b.names.items[1]); + try testing.expectEqualStrings("lua_x", b.names.items[2]); + + // ToolResults arrived in the original call 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("lua_x->1", tr_msg.content.items[0].ToolResult.content.items); + try testing.expectEqualStrings("b", tr_msg.content.items[1].ToolResult.tool_use_id); + try testing.expectEqualStrings("lua_y->2", tr_msg.content.items[1].ToolResult.content.items); + try testing.expectEqualStrings("c", tr_msg.content.items[2].ToolResult.tool_use_id); + try testing.expectEqualStrings("lua_x->3", tr_msg.content.items[2].ToolResult.content.items); +} + +test "runStep: distinct sources run on distinct threads in parallel" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "a", .name = "src_a_t", .input = "" } }, + .{ .ToolUse = .{ .id = "b", .name = "src_b_t", .input = "" } }, + } }, + .{ .blocks = &.{.{ .Text = "done" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); + defer agent.deinit(); + + try agent.registerToolSource(try TestSource.create(allocator, "src_a", &.{"src_a_t"})); + try agent.registerToolSource(try TestSource.create(allocator, "src_b", &.{"src_b_t"})); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("go"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + const view_a = agent.registry.lookup("src_a_t") orelse return error.NotFound; + const view_b = agent.registry.lookup("src_b_t") orelse return error.NotFound; + const sa: *TestSource = @ptrCast(@alignCast(view_a.entry.source.source.ctx)); + const sb: *TestSource = @ptrCast(@alignCast(view_b.entry.source.source.ctx)); + + try testing.expectEqual(@as(usize, 1), sa.batches.items.len); + try testing.expectEqual(@as(usize, 1), sb.batches.items.len); + // The two sources ran on distinct OS threads. + try testing.expect(sa.batches.items[0].thread_id != sb.batches.items[0].thread_id); +} + +test "runStep: source whole-batch error aborts the turn" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "a", .name = "fa", .input = "" } }, + .{ .ToolUse = .{ .id = "b", .name = "fb", .input = "" } }, + } }, + .{ .blocks = &.{.{ .Text = "never" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); + defer agent.deinit(); + + try agent.registerToolSource(try FailingSource.create(allocator, "fs", &.{ "fa", "fb" })); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("kaboom"); + + var recv = NoopReceiver.make(); + try testing.expectError(error.SourceExploded, agent.runStep(&conv, &recv)); + + // Conversation stops at user + assistant(tool_use). No ToolResult appended. + try testing.expectEqual(@as(usize, 2), conv.messages.items.len); +} + +test "runStep: mixed single Tools and source-backed tools coexist in one turn" { + const allocator = testing.allocator; + + const scripted = [_]StubProvider.ScriptedTurn{ + .{ .blocks = &.{ + .{ .ToolUse = .{ .id = "a", .name = "single", .input = "X" } }, + .{ .ToolUse = .{ .id = "b", .name = "src_t1", .input = "Y" } }, + .{ .ToolUse = .{ .id = "c", .name = "src_t2", .input = "Z" } }, + } }, + .{ .blocks = &.{.{ .Text = "done" }} }, + }; + var stub = StubProvider{ .allocator = allocator, .scripted = &scripted }; + var threaded: std.Io.Threaded = .init(allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + var agent = Agent.init(allocator, io, stub.provider()); + defer agent.deinit(); + + try agent.registerTool(try EchoTool.create(allocator, "single", "S:")); + try agent.registerToolSource(try TestSource.create(allocator, "src", &.{ "src_t1", "src_t2" })); + + var conv = conversation.Conversation.init(allocator); + defer conv.deinit(); + try conv.addUserMessage("go"); + + var recv = NoopReceiver.make(); + try agent.runStep(&conv, &recv); + + const tr_msg = conv.messages.items[2]; + try testing.expectEqual(@as(usize, 3), tr_msg.content.items.len); + try testing.expectEqualStrings("S:X", tr_msg.content.items[0].ToolResult.content.items); + try testing.expectEqualStrings("src_t1->Y", tr_msg.content.items[1].ToolResult.content.items); + try testing.expectEqualStrings("src_t2->Z", tr_msg.content.items[2].ToolResult.content.items); +} |