diff --git a/README.md b/README.md index 434f0f8..edd30bd 100644 --- a/README.md +++ b/README.md @@ -14,8 +14,8 @@ Sure, writing LLM agent loops is a solved problem and it's completely unnecessar - **Type-safe Comptime Tool Registration**: Register standard Zig functions as agent tools. The compiler checks that your function parameters match the tool's descriptor structure using Zig's compile-time reflection. - **ACP Server Integration**: Built-in JSON-RPC stdio server supporting the Agent Connection Protocol (ACP) for external client/IDE integration (`--acp`). -- **Context-Aware Tools & State Management**: Support for custom user-data contexts (`Tool.initWithContext`), session state tracking, and dynamic context injections across turns. -- **Built-in Agent Tools**: Standardized built-in tools (such as `agent.Tool.BuiltIn.Todo` for planning and tracking multi-step execution). +- **Context-Aware Tools & State Management**: Support for stateful tools with intrusive VTables (`Tool.wrapExecute`), session state tracking, and dynamic context injections across turns. +- **Built-in Agent Tools**: Standardized built-in tools (such as `agent.Tool.BuiltIn.Todo` for planning and tracking multi-step execution, and `agent.Tool.BuiltIn.Shell` for executing shell commands with multi-shell support: `sh`, `bash`, `zsh`, `cmd`, `powershell`, `pwsh`, and host `default`). - **Incremental Streaming & Response Accumulator**: Built-in support for streaming model responses chunk-by-chunk with ANSI terminal markdown formatting. - **Asynchronous & Concurrent Tool Execution**: Built on top of Zig's native non-blocking I/O event loop (`std.Io`). Tools run concurrently via `Io.Select`. - **Async Zero-Dependency HTTP Client**: Direct communication with LLM backends asynchronously using Zig's native `std.http.Client`. @@ -32,7 +32,7 @@ fn getWeather(allocator: std.mem.Allocator, zip_code: i64) ![]const u8 { } // Register it with the Agent -const weather_tool = Tool.init(.{ +const weather_tool = Tool.fromFn(.{ .name = "get_weather", .description = "Get the current weather for a given zip code.", .parameters = &.{ @@ -46,7 +46,7 @@ const weather_tool = Tool.init(.{ }, getWeather); ``` -You can also attach custom context to tools using `Tool.initWithContext` or include built-in tools like `agent.Tool.BuiltIn.Todo`. +You can also create stateful tools using `Tool.wrapExecute` and `Tool.VTable` or include built-in tools like `agent.Tool.BuiltIn.Todo` and `agent.Tool.BuiltIn.Shell.default`. ## Library Usage @@ -55,18 +55,25 @@ To use COMA in your own project, initialize a provider and the agent loop: ```zig var http_client: std.http.Client = .{ .allocator = allocator, .io = io }; var gemini_client = try provider.Gemini.init(allocator, &http_client, api_key); -var client = gemini_client.provider(); +const client = gemini_client.provider(); + +const tools = [_]*const Tool{ + &weather_tool, + &agent.Tool.BuiltIn.Todo, + &agent.Tool.BuiltIn.Shell.default, +}; const session_config: types.SessionConfig = .{ .model = selected_model, .system_prompt = "You are a helpful assistant.", - .tools = &[_]Tool{ weather_tool, agent.Tool.BuiltIn.Todo }, + .tools = &tools, }; var session = try Session.init(allocator, io, client, session_config); defer session.deinit(); -const result = try session.executeTurn(.{ .prompt = "What is the weather in 90210?" }); +var result = try session.executeTurn(.{ .prompt = "What is the weather in 90210?" }); +defer result.deinit(); ``` ## Project Structure diff --git a/src/acp/Config.zig b/src/acp/Config.zig index bba7b3b..0c6fe95 100644 --- a/src/acp/Config.zig +++ b/src/acp/Config.zig @@ -6,5 +6,5 @@ const SessionConfig = @import("agent").types.SessionConfig; const Config = @This(); -provider: Provider, +provider: *Provider, default_session_config: SessionConfig, diff --git a/src/acp/Server.zig b/src/acp/Server.zig index fad73e6..e669cac 100644 --- a/src/acp/Server.zig +++ b/src/acp/Server.zig @@ -30,7 +30,7 @@ pub const AcpProtocolError = error{ /// Internal context payload passed to session streaming callbacks. const ServerSessionContext = struct { - session_state: *SessionStorage.SessionState, + session_state: *const SessionStorage.SessionState, json_rpc_writer: *JsonRpcWriter, allocator: Allocator, io: Io, @@ -140,13 +140,13 @@ fn checkClientRequestValid(request: client_api.ClientRequest) AcpProtocolError!v /// Callback handler for streaming turn updates, converting agent streaming chunks into JSON-RPC notifications. fn handleTurnUpdate(ctx: ?*anyopaque, chunk: agent.types.StreamingChunk) void { - const stream_ctx: *ServerSessionContext = @ptrCast(@alignCast(ctx)); + const stream_ctx: *const ServerSessionContext = @ptrCast(@alignCast(ctx)); const notification = converter.streamingChunkToNotification(stream_ctx.allocator, stream_ctx.session_state.id, chunk) catch return orelse return; stream_ctx.json_rpc_writer.writeJsonObject(stream_ctx.io, notification, .{ .use_headers = false }) catch {}; } /// Formats and writes a JSON-RPC error response using the provided JSON-RPC writer. -fn sendError(self: *Server, writer: *JsonRpcWriter, id: shared_api.RequestId, code: agent_api.JsonRpcErrorCode, message: []const u8) !void { +fn sendError(self: *const Server, writer: *JsonRpcWriter, id: shared_api.RequestId, code: agent_api.JsonRpcErrorCode, message: []const u8) !void { try writer.writeJsonObject(self.io, agent_api.AgentErrorResponse{ .id = id, .@"error" = .{ @@ -157,7 +157,7 @@ fn sendError(self: *Server, writer: *JsonRpcWriter, id: shared_api.RequestId, co } /// Handles `initialize` request negotiation and writes JSON-RPC `initialize` response. -fn handleInitialize(self: *Server, writer: *JsonRpcWriter, id: shared_api.RequestId, params: client_api.InitializeRequest) !void { +fn handleInitialize(self: *const Server, writer: *JsonRpcWriter, id: shared_api.RequestId, params: client_api.InitializeRequest) !void { _ = params; const reply: agent_api.AgentResponse = .{ .id = id, @@ -202,7 +202,7 @@ fn handleSessionNew(self: *Server, writer: *JsonRpcWriter, id: shared_api.Reques /// Wrapper for `handleSessionPrompt` that allows it to be called from the main async loop. /// /// Takes ownership of `parse_result`, deinit-ing it upon completion. -fn handleSessionPromptFireAndForget(self: *Server, writer: *JsonRpcWriter, parse_result: std.json.Parsed(client_api.ClientRequest)) Io.Cancelable!void { +fn handleSessionPromptFireAndForget(self: *const Server, writer: *JsonRpcWriter, parse_result: std.json.Parsed(client_api.ClientRequest)) Io.Cancelable!void { defer parse_result.deinit(); const client_request = parse_result.value; self.handleSessionPrompt(writer, client_request.id, client_request.params.session_prompt) catch |err| { @@ -215,7 +215,7 @@ fn handleSessionPromptFireAndForget(self: *Server, writer: *JsonRpcWriter, parse /// TODO(razza): This method currently assume that only one prompt will ever be active for a given session at a time. /// This gurantee does not exist. We should either buffer the prompts or block/lock waiting to acquire exclusive access /// to the session. -fn handleSessionPrompt(self: *Server, writer: *JsonRpcWriter, id: shared_api.RequestId, params: client_api.PromptRequest) !void { +fn handleSessionPrompt(self: *const Server, writer: *JsonRpcWriter, id: shared_api.RequestId, params: client_api.PromptRequest) !void { const prompt_blocks = params.prompt; if (prompt_blocks.len == 0) { try self.sendError(writer, id, .invalid_params, "Prompt array cannot be empty"); @@ -339,9 +339,9 @@ test "Server prompt handling - multiple items in prompt array" { defer mock_provider.deinit(); const prov = mock_provider.provider(); - const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; diff --git a/src/acp/client_api.zig b/src/acp/client_api.zig index fa59506..84aaa63 100644 --- a/src/acp/client_api.zig +++ b/src/acp/client_api.zig @@ -155,9 +155,9 @@ pub const NewSessionRequest = struct { /// These expand the session's filesystem scope without changing `cwd`, which /// remains the base for relative paths. When omitted or empty, no /// additional roots are activated for the new session. - additionalDirectories: ?[][]const u8 = null, + additionalDirectories: ?[]const []const u8 = null, /// List of MCP (Model Context Protocol) servers the agent should connect to. - mcpServers: []McpServer, + mcpServers: []const McpServer, }; /// MCP server transport types. @@ -217,9 +217,9 @@ pub const McpServerStdio = struct { /// Absolute path to the MCP server executable. command: []const u8, /// Command-line arguments to pass to the MCP server. - args: [][]const u8, + args: []const []const u8, /// Environment variables to set when launching the MCP server. - env: []EnvVariable, + env: []const EnvVariable, }; /// An environment variable to set when launching an MCP server. @@ -252,7 +252,7 @@ pub const PromptRequest = struct { /// When available, [`ContentBlock::Resource`] is preferred /// as it avoids extra round-trips and allows the message to include /// pieces of context from sources the agent may not have access to. - prompt: []shared_api.ContentBlock, + prompt: []const shared_api.ContentBlock, }; test "json parse initialize ClientRequest" { diff --git a/src/agent/Session.zig b/src/agent/Session.zig index f5f84e1..c9a067b 100644 --- a/src/agent/Session.zig +++ b/src/agent/Session.zig @@ -19,10 +19,10 @@ const Session = @This(); allocator: Allocator, io: Io, -provider: Provider, -tools: []const Tool, +provider: *Provider, +tools: []const *const Tool, session_config: llm.types.SessionConfig, -prev_continuation: ?llm.types.StepContinuation, +prev_continuation: ?*llm.types.StepContinuation, session_state: SessionState, turn_mutex: Io.Mutex, @@ -35,7 +35,7 @@ pub const SessionError = error{TurnAlreadyRunning} || ToolError || Provider.Prov /// `io` is the I/O context to use for operations. /// `provider` is a Provider interface implementation that determines the LLM provider to use. /// `config` contains the configuration for the Session. -pub fn init(allocator: Allocator, io: Io, provider: Provider, config: types.SessionConfig) !Session { +pub fn init(allocator: Allocator, io: Io, provider: *Provider, config: types.SessionConfig) !Session { const descriptors = try allocator.alloc(llm.types.Tool, config.tools.len); errdefer allocator.free(descriptors); for (config.tools, 0..) |tool, i| { @@ -60,7 +60,7 @@ pub fn init(allocator: Allocator, io: Io, provider: Provider, config: types.Sess /// Deinitializes the Session, releasing any accumulated session history and internal resources. pub fn deinit(self: *Session) void { - if (self.prev_continuation) |*ls| { + if (self.prev_continuation) |ls| { ls.deinit(); self.prev_continuation = null; } @@ -112,7 +112,7 @@ fn streamingCallbackProxy(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) voi } fn executeToolCall(self: *Session, tool_call: llm.types.ToolCall) ToolError!llm.types.ToolResult { - const tool = for (self.tools) |*t| { + const tool = for (self.tools) |t| { if (std.mem.eql(u8, t.descriptor.name, tool_call.name)) { break t; } @@ -172,10 +172,10 @@ fn executeTurnInternal(self: *Session, turn: types.Turn, callback_context: ?*Str ); next_steps.clearRetainingCapacity(); - var step_result = step_outcome.result; + const step_result = step_outcome.result; const step_continuation = step_outcome.continuation; - if (self.prev_continuation) |*old_continuation| old_continuation.deinit(); + if (self.prev_continuation) |old_continuation| old_continuation.deinit(); self.prev_continuation = step_continuation; if (step_result.tool_calls.len > 0) { @@ -198,9 +198,7 @@ fn executeTurnInternal(self: *Session, turn: types.Turn, callback_context: ?*Str tool_futures.async(.result, executeToolCall, .{ self, tool_call }); } for (0..step_result.tool_calls.len) |_| { - const tool_result_wrapper = tool_futures.await() catch |err| switch (err) { - error.Canceled => unreachable, - }; + const tool_result_wrapper = try tool_futures.await(); var tool_result = try tool_result_wrapper.result; intermediate_results.append(allocator, .{ .tool_result = tool_result }) catch |err| { @@ -224,6 +222,7 @@ fn executeTurnInternal(self: *Session, turn: types.Turn, callback_context: ?*Str } const testing = @import("testing"); +const BlockingTool = @import("testing/BlockingTool.zig"); const MockToolImpl = struct { pub fn execute(allocator: std.mem.Allocator, val: i64) ![]const u8 { @@ -249,9 +248,9 @@ test "Session.executeTurn - no tool calls" { }); defer session.deinit(); - const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -288,9 +287,9 @@ test "Session.executeTurnStreaming - no tool calls" { ); defer session.deinit(); - const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -338,8 +337,8 @@ test "Session.executeTurn - executes tool call and runs again" { }, }; - const tool = Tool.init(tool_desc, MockToolImpl.execute); - const tools = &[_]Tool{tool}; + const tool = Tool.fromFn(tool_desc, MockToolImpl.execute); + const tools = &[_]*const Tool{&tool}; const mock_model = llm.types.Model{ .id = "mock-model", @@ -364,22 +363,22 @@ test "Session.executeTurn - executes tool call and runs again" { .{ .id = "call-id-123", .name = "mock_tool", - .arguments = @constCast(&args), + .arguments = &args, }, }; - const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); const continuation1 = testing.MockProvider.stepContinuation(); const final_outputs = [_]llm.types.ModelOutput{ .{ .text = "Final output after tool" }, }; - const result2 = testing.MockProvider.stepResult(&final_outputs, &.{}, &.{}); + var result2 = testing.MockProvider.stepResult(&final_outputs, &.{}, &.{}); const continuation2 = testing.MockProvider.stepContinuation(); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - llm.types.StepOutcome{ .result = result1, .continuation = continuation1 }, - llm.types.StepOutcome{ .result = result2, .continuation = continuation2 }, + llm.types.StepOutcome{ .result = &result1, .continuation = continuation1 }, + llm.types.StepOutcome{ .result = &result2, .continuation = continuation2 }, }; mock_provider.execute_step_results = &outcomes; @@ -440,9 +439,9 @@ test "Session.executeTurnStreaming - model chunks streaming" { const chunks_list = [_][]const llm.types.StreamingChunk{&chunks}; mock_provider.execute_step_streaming_chunks = &chunks_list; - const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello world!" }}, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello world!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -493,15 +492,15 @@ test "Session.executeTurnStreaming - with tool calls" { .{ .id = "call-id-123", .name = "mock_tool", - .arguments = @constCast(&args_buf), + .arguments = &args_buf, }, }; - const step_result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls_buf); - const step_result2 = testing.MockProvider.stepResult(&.{.{ .text = "Tool executed!" }}, &.{}, &.{}); + var step_result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls_buf); + var step_result2 = testing.MockProvider.stepResult(&.{.{ .text = "Tool executed!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = step_result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result2, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -517,8 +516,8 @@ test "Session.executeTurnStreaming - with tool calls" { }, }, }; - const tool = Tool.init(tool_desc, MockToolImpl.execute); - const tools = &[_]Tool{tool}; + const tool = Tool.fromFn(tool_desc, MockToolImpl.execute); + const tools = &[_]*const Tool{&tool}; var session = try Session.init( allocator, @@ -609,8 +608,8 @@ test "Session.executeTurn - tool call error cleanup" { return error.ArgumentTypeMismatch; } }; - const tool = Tool.init(tool_desc, error_tool_impl.execute); - const tools = &[_]Tool{tool}; + const tool = Tool.fromFn(tool_desc, error_tool_impl.execute); + const tools = &[_]*const Tool{&tool}; var session = try Session.init( allocator, @@ -630,10 +629,10 @@ test "Session.executeTurn - tool call error cleanup" { .arguments = &.{}, }, }; - const step_result = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var step_result = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); const step_continuation = testing.MockProvider.stepContinuation(); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = step_continuation }, + .{ .result = &step_result, .continuation = step_continuation }, }; mock_provider.execute_step_results = &outcomes; @@ -675,8 +674,8 @@ test "Session.executeTurn - tool receives ToolCallContext" { return try std.fmt.allocPrint(alloc, "State value is {d}", .{obj.value}); } }; - const tool = Tool.init(tool_desc, state_tool_impl.execute); - const tools = &[_]Tool{tool}; + const tool = Tool.fromFn(tool_desc, state_tool_impl.execute); + const tools = &[_]*const Tool{&tool}; var session = try Session.init( allocator, @@ -696,15 +695,15 @@ test "Session.executeTurn - tool receives ToolCallContext" { .{ .id = "call-id-123", .name = "state_tool", - .arguments = @constCast(&args), + .arguments = &args, }, }; - const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); - const result2 = testing.MockProvider.stepResult(&.{.{ .text = "Done" }}, &.{}, &.{}); + var result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var result2 = testing.MockProvider.stepResult(&.{.{ .text = "Done" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result2, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -755,19 +754,19 @@ test "Session.executeTurn prepends injected context from session_state" { return ""; } }; - const tool = Tool.init(tool_desc, tool_impl.run); + const tool = Tool.fromFn(tool_desc, tool_impl.run); var session = try Session.init(allocator, io, prov, .{ .model = .{ .id = "mock-model", .display_name = "Mock Model" }, - .tools = &.{tool}, + .tools = &.{&tool}, }); defer session.deinit(); - try session.session_state.setInjectedContext(&session.tools[0], "Injected tool info"); + try session.session_state.setInjectedContext(session.tools[0], "Injected tool info"); - const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -801,8 +800,8 @@ test "Session.executeTurn prepends injected context added during tool execution" } }; - const tool = Tool.init(dummy_tool_desc, ContextToolImpl.run); - const tools = &[_]Tool{tool}; + const tool = Tool.fromFn(dummy_tool_desc, ContextToolImpl.run); + const tools = &[_]*const Tool{&tool}; var session: Session = try .init( allocator, @@ -823,13 +822,13 @@ test "Session.executeTurn prepends injected context added during tool execution" }, }; - const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); - const result2 = testing.MockProvider.stepResult(&.{.{ .text = "Final turn output" }}, &.{}, &.{}); - const result3 = testing.MockProvider.stepResult(&.{.{ .text = "Second turn output" }}, &.{}, &.{}); + var result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var result2 = testing.MockProvider.stepResult(&.{.{ .text = "Final turn output" }}, &.{}, &.{}); + var result3 = testing.MockProvider.stepResult(&.{.{ .text = "Second turn output" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = result2, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = result3, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result3, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -859,8 +858,8 @@ test "Session.executeTurn - error.TurnAlreadyRunning when turn is in progress" { defer mock_provider.deinit(); const prov = mock_provider.provider(); - var blocking_tool = testing.BlockingTool.init(io); - const tool = blocking_tool.toolWithContext(Tool); + var blocking_state: BlockingTool.State = .{}; + var blocking_tool = BlockingTool.init(&blocking_state, io); var session = try Session.init( allocator, @@ -868,7 +867,7 @@ test "Session.executeTurn - error.TurnAlreadyRunning when turn is in progress" { prov, .{ .model = .{ .id = "mock-model", .display_name = "Mock Model" }, - .tools = &.{tool}, + .tools = &.{&blocking_tool.tool}, }, ); defer session.deinit(); @@ -876,15 +875,15 @@ test "Session.executeTurn - error.TurnAlreadyRunning when turn is in progress" { const tool_calls = [_]llm.types.ToolCall{ .{ .id = "call-blocking-1", - .name = testing.BlockingTool.descriptor.name, + .name = BlockingTool.descriptor.name, .arguments = &.{}, }, }; - const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); - const result2 = testing.MockProvider.stepResult(&.{.{ .text = "Finished blocking turn" }}, &.{}, &.{}); + var result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var result2 = testing.MockProvider.stepResult(&.{.{ .text = "Finished blocking turn" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result2, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -928,8 +927,8 @@ test "Session.executeTurnStreaming - error.TurnAlreadyRunning when turn is in pr defer mock_provider.deinit(); const prov = mock_provider.provider(); - var blocking_tool = testing.BlockingTool.init(io); - const tool = blocking_tool.toolWithContext(Tool); + var blocking_state: BlockingTool.State = .{}; + var blocking_tool = BlockingTool.init(&blocking_state, io); var session = try Session.init( allocator, @@ -937,7 +936,7 @@ test "Session.executeTurnStreaming - error.TurnAlreadyRunning when turn is in pr prov, .{ .model = .{ .id = "mock-model", .display_name = "Mock Model" }, - .tools = &.{tool}, + .tools = &.{&blocking_tool.tool}, }, ); defer session.deinit(); @@ -945,15 +944,15 @@ test "Session.executeTurnStreaming - error.TurnAlreadyRunning when turn is in pr const tool_calls = [_]llm.types.ToolCall{ .{ .id = "call-blocking-1", - .name = testing.BlockingTool.descriptor.name, + .name = BlockingTool.descriptor.name, .arguments = &.{}, }, }; - const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); - const result2 = testing.MockProvider.stepResult(&.{.{ .text = "Finished blocking turn" }}, &.{}, &.{}); + var result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + var result2 = testing.MockProvider.stepResult(&.{.{ .text = "Finished blocking turn" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &result2, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; @@ -1012,12 +1011,12 @@ test "Session.executeTurn - sequential turns execute successfully" { ); defer session.deinit(); - const step_result1 = testing.MockProvider.stepResult(&.{.{ .text = "Turn 1 output" }}, &.{}, &.{}); - const step_result2 = testing.MockProvider.stepResult(&.{.{ .text = "Turn 2 output" }}, &.{}, &.{}); + var step_result1 = testing.MockProvider.stepResult(&.{.{ .text = "Turn 1 output" }}, &.{}, &.{}); + var step_result2 = testing.MockProvider.stepResult(&.{.{ .text = "Turn 2 output" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result1, .continuation = testing.MockProvider.stepContinuation() }, - .{ .result = step_result2, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result1, .continuation = testing.MockProvider.stepContinuation() }, + .{ .result = &step_result2, .continuation = testing.MockProvider.stepContinuation() }, }; mock_provider.execute_step_results = &outcomes; diff --git a/src/agent/SessionState.zig b/src/agent/SessionState.zig index 59e2383..c01ee28 100644 --- a/src/agent/SessionState.zig +++ b/src/agent/SessionState.zig @@ -242,7 +242,7 @@ test setInjectedContext { return ""; } }; - const tool = Tool.init(tool_desc, tool_impl.run); + const tool = Tool.fromFn(tool_desc, tool_impl.run); try state.setInjectedContext(&tool, "System context info"); try std.testing.expectEqualStrings("System context info", state.getInjectedContext(&tool).?); @@ -267,7 +267,7 @@ test getInjectedContext { return ""; } }; - const tool = Tool.init(tool_desc, tool_impl.run); + const tool = Tool.fromFn(tool_desc, tool_impl.run); try std.testing.expectEqual(null, state.getInjectedContext(&tool)); @@ -290,7 +290,7 @@ test clearInjectedContext { return ""; } }; - const tool = Tool.init(tool_desc, tool_impl.run); + const tool = Tool.fromFn(tool_desc, tool_impl.run); // Clearing non-existent entry is a safe no-op state.clearInjectedContext(&tool); @@ -323,7 +323,7 @@ test getInjectedContextString { return ""; } }; - const tool = Tool.init(tool_desc, tool_impl.run); + const tool = Tool.fromFn(tool_desc, tool_impl.run); try state.setInjectedContext(&tool, "Use carefully."); const single_str = (try state.getInjectedContextString(std.testing.allocator)).?; diff --git a/src/agent/Tool.zig b/src/agent/Tool.zig index 340584c..0993c3a 100644 --- a/src/agent/Tool.zig +++ b/src/agent/Tool.zig @@ -21,14 +21,55 @@ pub const CallError = error{ /// An argument was provided with a name that does not match the expected name; /// either an unexpected name or the order does not match. ArgumentMismatch, -} || std.mem.Allocator.Error; +} || std.mem.Allocator.Error || std.Io.Cancelable; +/// Intrusive interface representing a tool an Agent can call. +/// +/// Implementations must embed `Tool` directly as a field in their enclosing struct and implement +/// the methods specified in `Tool.VTable`. +/// +/// **Pointer Stability & Memory Safety:** +/// Implementations recover the outer struct reference using `@fieldParentPtr("tool", self)` inside +/// VTable function callbacks. Because this relies on fixed relative memory offsets, the enclosing struct +/// instance MUST NOT be copied by value or relocated in memory once a pointer to its embedded `Tool` +/// has been passed to a Session or another consumer. +/// +/// **Thread Safety & Mutable State:** +/// Tools receive immutable `*const Tool` references during execution. Execution methods can be called +/// concurrently by parallel turns or simultaneous tool invocations within a single turn (e.g. via `Io.Select`). +/// Implementations should avoid casting away const-ness with `@constCast`, as doing so obscures mutability +/// and can cause undefined behavior if the tool instance is declared `const` (e.g. in read-only memory). +/// +/// To manage mutable state cleanly without `@constCast`: +/// 1. **Session-scoped state:** Use `ToolCallContext.getOrInitState` to store state inside `SessionState`. +/// This isolates state to the active session (avoiding cross-session leakage). Note that multiple tool +/// calls within a single turn can execute concurrently, so any shared state accessed simultaneously +/// still requires appropriate synchronization (e.g., atomics or a mutex). +/// 2. **Tool-scoped mutable objects (Clients, Pools, Atomics):** Store a pointer to the mutable object +/// (e.g., `client: *HttpClient`, `state: *SharedState`). Because Zig's `const` is shallow, accessing +/// a pointer field through `*const Tool` yields a mutable pointer (`*T`) with zero `@constCast`. const Tool = @This(); -const ToolExecuteFn = *const fn (allocator: Allocator, io: Io, call_ctx: ToolCallContext, ctx: ?*anyopaque, args: []const Argument) CallError![]const u8; + +pub const ToolExecuteFn = *const fn ( + tool: *const Tool, + allocator: Allocator, + io: Io, + call_ctx: ToolCallContext, + args: []const Argument, +) CallError![]const u8; +pub const ToolDeinitFn = *const fn (tool: *Tool) void; descriptor: llm.types.Tool, -execute_fn: ToolExecuteFn, -ctx: ?*anyopaque, +vtable: *const VTable, + +/// The virtual function table defining the method references used during tool execution. +pub const VTable = struct { + /// Function to call when executing the tool. + execute: ToolExecuteFn, + + /// Frees the resources associated with the tool. + deinit: ?ToolDeinitFn = null, +}; /// Executes the tool with the given arguments. /// @@ -43,43 +84,39 @@ ctx: ?*anyopaque, /// for freeing the `ToolResult` by calling `deinit()`. pub fn execute(self: *const Tool, allocator: Allocator, io: Io, session_state: *SessionState, id: []const u8, args: []const Argument) CallError!ToolResult { const call_ctx: ToolCallContext = .{ .tool = self, .session_state = session_state }; - const result = try self.execute_fn(allocator, io, call_ctx, self.ctx, args); + const result = try self.vtable.execute(self, allocator, io, call_ctx, args); errdefer allocator.free(result); return ToolResult.initTakingResultOwnership(allocator, self.descriptor.name, id, result); } -/// Creates a Tool from a descriptor and a function. -/// Tool calls will delegate to the provided function when called. -/// -/// `descriptor` is the tool's descriptor defining the structure of the tool for the LLM. -/// `execute_fn` is the function to be called when the tool is executed. -/// -/// The function arguments must match the descriptor parameters and be in the same order. -/// The function should also take an allocator as an argument which will be used, at least, -/// to allocate its result. The allocator can be in any position provided the other arguments -/// are still in the same order as the parameters in the descriptor. +/// Wraps an execution function into a `ToolExecuteFn` suitable for a `Tool.VTable`. /// -/// Additionally, the function can optionally accept an Io struct which represents the IO to use -/// during the tool call, and/or a ToolCallContext struct which provides access to a session-scoped -/// state store and the ability to define arbitrary context messages to inject before each turn. -/// The function MUST be thread-safe as multiple calls to the tool function can occur in parallel. -/// -/// The `execute_fn` should return the result of the tool call as a string and transfer -/// ownership of the memory to the caller. The result will be passed to the LLM as the -/// result of the tool call. -pub fn init(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) Tool { - if (comptime expectsContext(execute_fn)) { - @compileError("Tool function '" ++ descriptor.name ++ "' expects a context parameter. Use Tool.initWithContext instead."); - } - return initInternal(descriptor, execute_fn, null); +/// Validates at compile time that the function parameters match the tool descriptor. +/// Stateful tools that take `*const Tool` should use this function to populate +/// their `VTable.execute`. +pub fn wrapExecute(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) ToolExecuteFn { + return comptime blk: { + for (descriptor.parameters, 0..) |p1, i| { + for (descriptor.parameters, 0..) |p2, j| { + if (i != j and std.mem.eql(u8, p1.name, p2.name)) { + @compileError("Duplicate parameter name: " ++ p1.name); + } + } + } + + const res = makeExecuteFn(descriptor, execute_fn); + break :blk switch (res) { + .ok => |f| f, + .err => |e| @compileError(e.msg), + }; + }; } -/// Creates a Tool from a descriptor, a function, and a provided context object. +/// Creates a Tool from a descriptor and a standalone function. /// Tool calls will delegate to the provided function when called. /// /// `descriptor` is the tool's descriptor defining the structure of the tool for the LLM. /// `execute_fn` is the function to be called when the tool is executed. -/// `ctx` is a pointer that will be passed to the `execute_fn`. /// /// The function arguments must match the descriptor parameters and be in the same order. /// The function should also take an allocator as an argument which will be used, at least, @@ -87,56 +124,38 @@ pub fn init(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) T /// are still in the same order as the parameters in the descriptor. /// /// Additionally, the function can optionally accept an Io struct which represents the IO to use -/// during the tool call, and/or a ToolCallContext struct. +/// during the tool call and/or a ToolCallContext struct. +/// The function MUST NOT accept a `*const Tool` or `*Tool` parameter. For tools with tool-level +/// state, define a custom struct embedding `tool: Tool`, use `Tool.wrapExecute` in a `Tool.VTable`, +/// and use `@fieldParentPtr`. +/// The function MUST be thread-safe as multiple calls to the tool function can occur in parallel. /// /// The `execute_fn` should return the result of the tool call as a string and transfer /// ownership of the memory to the caller. The result will be passed to the LLM as the /// result of the tool call. -pub fn initWithContext(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype, ctx: anytype) Tool { - if (comptime !expectsContext(execute_fn)) { - @compileError("Tool function '" ++ descriptor.name ++ "' does not take a context parameter. Use Tool.init instead."); +pub fn fromFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) Tool { + if (comptime takesToolParam(@TypeOf(execute_fn))) { + @compileError("Tool.fromFn function '" ++ descriptor.name ++ "' cannot take a *const Tool or *Tool parameter. Construct a Tool with a VTable and Tool.wrapExecute instead."); } - if (comptime !isContextType(@TypeOf(ctx))) { - @compileError("Context argument to Tool.initWithContext must be a pointer, got: " ++ @typeName(@TypeOf(ctx))); - } - return initInternal(descriptor, execute_fn, ctx); -} -fn initInternal(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype, ctx: anytype) Tool { - const CtxType: ?type = comptime if (@TypeOf(ctx) == @TypeOf(null)) null else @TypeOf(ctx); - - var value = comptime blk: { - for (descriptor.parameters, 0..) |p1, i| { - for (descriptor.parameters, 0..) |p2, j| { - if (i != j and std.mem.eql(u8, p1.name, p2.name)) { - @compileError("Duplicate parameter name: " ++ p1.name); - } - } - } + return comptime blk: { + const vtable: VTable = .{ + .execute = wrapExecute(descriptor, execute_fn), + .deinit = null, + }; - const res = makeExecuteFn(descriptor, execute_fn, CtxType); break :blk Tool{ .descriptor = descriptor, - .execute_fn = switch (res) { - .ok => |f| f, - .err => |e| @compileError(e.msg), - }, - .ctx = null, + .vtable = &vtable, }; }; +} - value.ctx = blk: { - if (comptime CtxType == null) break :blk null; - - if (comptime @typeInfo(CtxType.?) == .optional) { - const ptr = ctx orelse break :blk null; - break :blk @ptrCast(@alignCast(@constCast(ptr))); - } - - break :blk @ptrCast(@alignCast(@constCast(ctx))); - }; - - return value; +/// Frees the resources associated with the tool. +pub fn deinit(self: *Tool) void { + if (self.vtable.deinit) |deinit_fn| { + deinit_fn(self); + } } const ValidationError = error{ @@ -155,6 +174,25 @@ const ValidationResult = union(enum) { }, }; +fn takesToolParam(comptime FnType: type) bool { + const fn_info = switch (@typeInfo(FnType)) { + .@"fn" => |info| info, + .pointer => |ptr_info| switch (@typeInfo(ptr_info.child)) { + .@"fn" => |info| info, + else => return false, + }, + else => return false, + }; + inline for (fn_info.params) |param| { + if (param.type) |T| { + if (T == *const Tool or T == *Tool) { + return true; + } + } + } + return false; +} + fn expectedTagForType(comptime T: type) ValidationError!struct { tag: std.meta.Tag(Argument.Value), required: bool } { var InputT = T; var is_optional = false; @@ -211,74 +249,7 @@ inline fn findArgument(args: []const Argument, name: []const u8) ?*const Argumen return null; } -/// Returns true if type `T` represents a valid context pointer type. -/// A context pointer type is a single-item pointer (e.g. `*MyCtx`, `*const MyCtx`, `*anyopaque`) -/// or an optional single-item pointer (e.g. `?*MyCtx`). -fn isContextType(comptime T: type) bool { - if (T == ToolCallContext) return false; - return switch (@typeInfo(T)) { - .pointer => |ptr_info| ptr_info.size == .one, - .optional => |opt_info| switch (@typeInfo(opt_info.child)) { - .pointer => |ptr_info| ptr_info.size == .one, - else => false, - }, - else => false, - }; -} - -/// Returns true if the provided context type `ProvidedCtx` can be passed to a tool parameter of type `ParamType`. -/// Handles non-const to const pointer conversions, optional pointer unwrapping, and opaque pointer compatibility. -inline fn isContextTypeCompatible(comptime ProvidedCtx: type, comptime ParamType: type) bool { - if (ParamType == *anyopaque or ParamType == ?*anyopaque or ParamType == *const anyopaque or ParamType == ?*const anyopaque) return true; - if (ProvidedCtx == ParamType) return true; - - const UnwrappedProvided = switch (@typeInfo(ProvidedCtx)) { - .optional => |opt| opt.child, - else => ProvidedCtx, - }; - const UnwrappedParam = switch (@typeInfo(ParamType)) { - .optional => |opt| opt.child, - else => ParamType, - }; - - if (UnwrappedProvided == UnwrappedParam) return true; - - if (@typeInfo(UnwrappedProvided) == .pointer and @typeInfo(UnwrappedParam) == .pointer) { - const prov_ptr = @typeInfo(UnwrappedProvided).pointer; - const param_ptr = @typeInfo(UnwrappedParam).pointer; - - if (prov_ptr.child == anyopaque or param_ptr.child == anyopaque) return true; - - if (prov_ptr.child == param_ptr.child) { - if (param_ptr.is_const) return true; - return !prov_ptr.is_const; - } - } - - return false; -} - -/// Returns true if the function or function pointer `execute_fn` has at least one context parameter. -fn expectsContext(comptime execute_fn: anytype) bool { - const FnType = @TypeOf(execute_fn); - const fn_info = switch (@typeInfo(FnType)) { - .@"fn" => |info| info, - .pointer => |ptr_info| switch (@typeInfo(ptr_info.child)) { - .@"fn" => |info| info, - else => return false, - }, - else => return false, - }; - - inline for (fn_info.params) |param| { - if (param.type) |T| { - if (comptime isContextType(T)) return true; - } - } - return false; -} - -fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype, comptime CtxType: ?type) ValidationResult { +fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) ValidationResult { const FnType = @TypeOf(execute_fn); const fn_info = switch (@typeInfo(FnType)) { .@"fn" => |info| info, @@ -303,17 +274,13 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty for (fn_info.params, 0..) |fn_param, i| { result_types[i] = fn_param.type.?; - if (fn_param.type.? == Allocator or fn_param.type.? == Io or fn_param.type.? == ToolCallContext) { + if (fn_param.type.? == Allocator or fn_param.type.? == Io or fn_param.type.? == ToolCallContext or fn_param.type.? == *const Tool) { // Environment parameter - } else if (isContextType(fn_param.type.?)) { - if (CtxType) |ProvidedCtxType| { - if (!isContextTypeCompatible(ProvidedCtxType, fn_param.type.?)) { - return .{ .err = .{ - .code = ValidationError.ArgumentTypeMismatch, - .msg = "Tool function '" ++ descriptor.name ++ "' expects context parameter of type '" ++ @typeName(fn_param.type.?) ++ "', but got context of type '" ++ @typeName(ProvidedCtxType) ++ "'", - } }; - } - } + } else if (fn_param.type.? == *Tool) { + return .{ .err = .{ + .code = ValidationError.UnsupportedType, + .msg = "Tool execution function parameter cannot be '*Tool'. Use '*const Tool' instead, as tools receive immutable references during execution.", + } }; } else { if (param_idx >= descriptor_params.len) { return .{ .err = .{ .code = ValidationError.ArgumentCountMismatch, .msg = "More arguments in function than descriptor." } }; @@ -341,7 +308,7 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty const TupleType = @Tuple(&types); return .{ .ok = struct { - pub fn call(allocator: Allocator, io: Io, call_ctx: ToolCallContext, ctx: ?*anyopaque, input_args: []const Argument) CallError![]const u8 { + pub fn call(tool: *const Tool, allocator: Allocator, io: Io, call_ctx: ToolCallContext, input_args: []const Argument) CallError![]const u8 { var args: TupleType = undefined; comptime var descriptor_idx: usize = 0; inline for (0..fn_info.params.len) |func_idx| { @@ -352,12 +319,8 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty args[func_idx] = io; } else if (comptime T == ToolCallContext) { args[func_idx] = call_ctx; - } else if (comptime (CtxType != null and isContextTypeCompatible(CtxType.?, T))) { - if (comptime @typeInfo(T) == .optional) { - args[func_idx] = if (ctx) |c| @ptrCast(@alignCast(c)) else null; - } else { - args[func_idx] = @ptrCast(@alignCast(ctx.?)); - } + } else if (comptime T == *const Tool) { + args[func_idx] = tool; } else { const curr_descriptor = comptime descriptor.parameters[descriptor_idx]; descriptor_idx += 1; @@ -386,7 +349,18 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty }; } -test init { +test takesToolParam { + const S = struct { + fn fConstTool(_: *const Tool) void {} + fn fTool(_: *Tool) void {} + fn fStateless(_: Allocator, _: i64) void {} + }; + try std.testing.expect(takesToolParam(@TypeOf(S.fConstTool))); + try std.testing.expect(takesToolParam(@TypeOf(S.fTool))); + try std.testing.expect(!takesToolParam(@TypeOf(S.fStateless))); +} + +test fromFn { const allocator = std.testing.allocator; const io = std.testing.io; var session_state = SessionState.init(allocator); @@ -415,7 +389,7 @@ test init { return try std.fmt.allocPrint(allocator, "{d}{s}", .{ arg1, arg2 }); } }; - const tool = comptime init(tool_descriptor, tool_impl.example_function); + const tool = comptime fromFn(tool_descriptor, tool_impl.example_function); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .integer = 12 } }, @@ -430,7 +404,7 @@ test init { try std.testing.expectEqualStrings("12hello", result.result); } -test initWithContext { +test "stateful tool with @fieldParentPtr" { const allocator = std.testing.allocator; const io = std.testing.io; var session_state = SessionState.init(allocator); @@ -454,22 +428,33 @@ test initWithContext { }, }, }; - const CtxStruct = struct { prefix: []const u8 }; - var ctx_obj: CtxStruct = .{ .prefix = "ctx-" }; + const StatefulTool = struct { + tool: Tool, + prefix: []const u8, - const tool_impl = struct { - pub fn example_context_function(_: Allocator, arg1: i64, arg2: []const u8, ctx: *CtxStruct) ![]const u8 { - return try std.fmt.allocPrint(allocator, "{s}{d}{s}", .{ ctx.prefix, arg1, arg2 }); + const vtable: Tool.VTable = .{ + .execute = Tool.wrapExecute(tool_descriptor, run), + }; + + fn run(tool: *const Tool, alloc: Allocator, arg1: i64, arg2: []const u8) ![]const u8 { + const self: *const @This() = @fieldParentPtr("tool", tool); + return try std.fmt.allocPrint(alloc, "{s}{d}{s}", .{ self.prefix, arg1, arg2 }); } }; - const tool = initWithContext(tool_descriptor, tool_impl.example_context_function, &ctx_obj); + const stateful_tool = StatefulTool{ + .tool = .{ + .descriptor = tool_descriptor, + .vtable = &StatefulTool.vtable, + }, + .prefix = "ctx-", + }; const args = [_]Argument{ .{ .name = "arg1", .value = .{ .integer = 12 } }, .{ .name = "arg2", .value = .{ .string = "hello" } }, }; - var result = try tool.execute(allocator, io, &session_state, "123", &args); + var result = try stateful_tool.tool.execute(allocator, io, &session_state, "123", &args); defer result.deinit(); try std.testing.expectEqualStrings("example_context_function", result.tool_name); @@ -483,7 +468,7 @@ test "makeExecuteFn - ExpectedFunctionOrPointer" { .description = "desc", .parameters = &.{}, }; - const res = comptime makeExecuteFn(desc, 42, null); + const res = comptime makeExecuteFn(desc, 42); try std.testing.expectEqual(ValidationError.ExpectedFunctionOrPointer, res.err.code); } @@ -505,7 +490,7 @@ test "makeExecuteFn - ArgumentTypeMismatch" { return arg1; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected i64 but got []const u8", res.err.msg); } @@ -522,7 +507,7 @@ test "makeExecuteFn - ArgumentCountMismatch (too many arguments)" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); try std.testing.expectEqualStrings("More arguments in function than descriptor.", res.err.msg); } @@ -545,7 +530,7 @@ test "makeExecuteFn - ArgumentCountMismatch (too few arguments)" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); try std.testing.expectEqualStrings("Fewer arguments in function than descriptor.", res.err.msg); } @@ -569,7 +554,7 @@ test "makeExecuteFn - ParamTypeArrayNotSupported" { return arg1; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.UnsupportedType, res.err.code); try std.testing.expectEqualStrings("Failed to resolve type from descriptor (tool: test_tool, param: arg1)", res.err.msg); } @@ -593,7 +578,7 @@ test "makeExecuteFn - ReturnTypeMismatch" { return 10; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ReturnTypeMismatch, res.err.code); try std.testing.expectEqualStrings("Function return type must be a string.", res.err.msg); } @@ -617,7 +602,7 @@ test "makeExecuteFn - argument optionals OK" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expect(res == .ok); } @@ -639,7 +624,7 @@ test "makeExecuteFn - argument optional in descriptor, required in fn" { return required; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected ?[]const u8 but got []const u8", res.err.msg); } @@ -663,7 +648,7 @@ test "makeExecuteFn - argument required in descriptor, optional in fn" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run, null); + const res = comptime makeExecuteFn(desc, Impl.run); try std.testing.expectEqual(ValidationError.ArgumentTypeMismatch, res.err.code); try std.testing.expectEqualStrings("Argument type mismatch in tool test_tool for argument arg1: expected []const u8 but got ?[]const u8", res.err.msg); } @@ -691,7 +676,7 @@ test execute { return try allocator.dupe(u8, arg1); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{ .{ .name = "arg1", @@ -729,7 +714,7 @@ test "execute - unknown argument" { return try allocator.dupe(u8, arg1); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{ .{ .name = "unknown", @@ -762,7 +747,7 @@ test "execute - extra argument ignored" { return try allocator.dupe(u8, arg1); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{ .{ .name = "arg1", @@ -801,7 +786,7 @@ test "execute - missing required argument" { return try allocator.dupe(u8, arg1); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{}; try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, &session_state, "id", args)); } @@ -830,7 +815,7 @@ test "execute - missing optional argument" { return allocator.dupe(u8, "default"); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{}; var result = try tool.execute(testing_allocator, io, &session_state, "id", args); defer result.deinit(); @@ -864,7 +849,7 @@ test "execute - optional argument provided" { return allocator.dupe(u8, "default"); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{ .{ .name = "arg1", .value = .{ .string = "provided" } }, }; @@ -899,7 +884,7 @@ test "execute - argument type mismatch" { return try allocator.dupe(u8, arg1); } }; - const tool = Tool.init(desc, Impl.run); + const tool = Tool.fromFn(desc, Impl.run); const args: []const Argument = &.{ .{ .name = "arg1", @@ -932,7 +917,7 @@ test "execute - no Allocator parameter" { return allocator.dupe(u8, arg1); } }; - const tool = comptime init(tool_descriptor, tool_impl.no_allocator_func); + const tool = comptime fromFn(tool_descriptor, tool_impl.no_allocator_func); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .string = "hello" } }, @@ -975,7 +960,7 @@ test "execute - Allocator as middle/last parameter" { return try std.fmt.allocPrint(tool_allocator, "middle-{d}-{s}", .{ arg1, arg2 }); } }; - const tool = comptime init(tool_descriptor, tool_impl.middle_last_allocator_func); + const tool = comptime fromFn(tool_descriptor, tool_impl.middle_last_allocator_func); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .integer = 99 } }, @@ -1014,7 +999,7 @@ test "execute - multiple Allocator parameters" { return try std.fmt.allocPrint(alloc1, "multi-{d}", .{arg1}); } }; - const tool = comptime init(tool_descriptor, tool_impl.multi_allocator_func); + const tool = comptime fromFn(tool_descriptor, tool_impl.multi_allocator_func); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .integer = 7 } }, @@ -1052,7 +1037,7 @@ test "execute - Io parameter" { return try std.fmt.allocPrint(tool_allocator, "io-{s}", .{arg1}); } }; - const tool = comptime init(tool_descriptor, tool_impl.io_func); + const tool = comptime fromFn(tool_descriptor, tool_impl.io_func); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .string = "test" } }, @@ -1091,7 +1076,7 @@ test "execute - multiple Io parameters" { return try std.fmt.allocPrint(tool_allocator, "multi-io-{d}", .{arg1}); } }; - const tool = comptime init(tool_descriptor, tool_impl.multi_io_func); + const tool = comptime fromFn(tool_descriptor, tool_impl.multi_io_func); const args = [_]Argument{ .{ .name = "arg1", .value = .{ .integer = 77 } }, @@ -1105,234 +1090,201 @@ test "execute - multiple Io parameters" { try std.testing.expectEqualStrings("multi-io-77", result.result); } -test expectsContext { - const ImplNoCtx = struct { - pub fn run(_: Allocator, _: i64) ![]const u8 { - return ""; - } +test "execute - inject *const Tool with wrapExecute" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(allocator); + defer session_state.deinit(); + + const desc: llm.types.Tool = .{ + .name = "tool_ptr_test", + .description = "desc", + .parameters = &.{}, }; - const ImplWithCtx = struct { - pub fn run(_: Allocator, _: i64, _: *anyopaque) ![]const u8 { - return ""; + const Impl = struct { + fn run(tool: *const Tool, alloc: Allocator) ![]const u8 { + return try alloc.dupe(u8, tool.descriptor.name); } }; - const CtxStruct = struct { val: u32 }; - const ImplWithTypedCtx = struct { - pub fn run(_: Allocator, _: i64, _: *CtxStruct) ![]const u8 { - return ""; - } + const vtable: Tool.VTable = .{ + .execute = Tool.wrapExecute(desc, Impl.run), }; - const ImplWithCallContext = struct { - pub fn run(_: Allocator, _: ToolCallContext, _: i64) ![]const u8 { - return ""; - } + const tool = Tool{ + .descriptor = desc, + .vtable = &vtable, }; - try std.testing.expect(!expectsContext(ImplNoCtx.run)); - try std.testing.expect(expectsContext(ImplWithCtx.run)); - try std.testing.expect(expectsContext(ImplWithTypedCtx.run)); - try std.testing.expect(!expectsContext(ImplWithCallContext.run)); + var res = try tool.execute(allocator, io, &session_state, "id", &.{}); + defer res.deinit(); + try std.testing.expectEqualStrings("tool_ptr_test", res.result); } -test "execute - has context" { +test "execute with ToolCallContext" { const allocator = std.testing.allocator; const io = std.testing.io; var session_state = SessionState.init(allocator); defer session_state.deinit(); + const CounterState = struct { + count: i64, + fn initFn(self: *@This(), _: Allocator) !void { + self.count = 0; + } + }; + const tool_descriptor: llm.types.Tool = .{ - .name = "ctx_func", - .description = "Takes context and arguments", + .name = "counter_tool", + .description = "Increments counter in SessionState", .parameters = &.{ .{ - .name = "arg1", - .description = "The first argument", - .type = .string, + .name = "increment", + .type = .integer, .required = true, + .description = "amount to increment", }, }, }; - const CtxStruct = struct { prefix: []const u8 }; - var ctx_obj: CtxStruct = .{ .prefix = "ctx-hello-" }; const tool_impl = struct { - pub fn ctx_func(tool_allocator: Allocator, arg1: []const u8, ctx: *CtxStruct) ![]const u8 { - return try std.fmt.allocPrint(tool_allocator, "{s}{s}", .{ ctx.prefix, arg1 }); + pub fn increment(alloc: Allocator, call_ctx: ToolCallContext, inc: i64) CallError![]const u8 { + const counter = call_ctx.getOrInitState(CounterState, CounterState.initFn) catch return error.OutOfMemory; + counter.count += inc; + return try std.fmt.allocPrint(alloc, "Count: {d}", .{counter.count}); + } + pub fn read_const(alloc: Allocator, call_ctx: ToolCallContext, inc: i64) CallError![]const u8 { + const count = if (call_ctx.getState(CounterState)) |c| c.count else 0; + return try std.fmt.allocPrint(alloc, "ConstCount: {d}", .{count + inc}); } }; - const tool = initWithContext(tool_descriptor, tool_impl.ctx_func, &ctx_obj); + const args = [_]Argument{.{ .name = "increment", .value = .{ .integer = 5 } }}; - const args = [_]Argument{ - .{ .name = "arg1", .value = .{ .string = "world" } }, - }; + const tool_inc = fromFn(tool_descriptor, tool_impl.increment); + var res1 = try tool_inc.execute(allocator, io, &session_state, "id1", &args); + defer res1.deinit(); + try std.testing.expectEqualStrings("Count: 5", res1.result); - var result = try tool.execute(allocator, io, &session_state, "ctx-id", &args); - defer result.deinit(); + var res2 = try tool_inc.execute(allocator, io, &session_state, "id2", &args); + defer res2.deinit(); + try std.testing.expectEqualStrings("Count: 10", res2.result); - try std.testing.expectEqualStrings("ctx_func", result.tool_name); - try std.testing.expectEqualStrings("ctx-id", result.id); - try std.testing.expectEqualStrings("ctx-hello-world", result.result); + const tool_read = fromFn(tool_descriptor, tool_impl.read_const); + var res3 = try tool_read.execute(allocator, io, &session_state, "id3", &args); + defer res3.deinit(); + try std.testing.expectEqualStrings("ConstCount: 15", res3.result); } -test "execute - const and optional context pointers" { +test "custom VTable implementation" { const allocator = std.testing.allocator; const io = std.testing.io; var session_state = SessionState.init(allocator); defer session_state.deinit(); - const tool_descriptor: llm.types.Tool = .{ - .name = "const_ctx_func", - .description = "Takes const and optional context pointers", - .parameters = &.{ - .{ - .name = "val", - .type = .integer, - .required = true, - .description = "val", - }, - }, + const desc: llm.types.Tool = .{ + .name = "custom_vtable_tool", + .description = "desc", + .parameters = &.{}, }; - const CtxStruct = struct { factor: i64 }; - const const_ctx_obj: CtxStruct = .{ .factor = 10 }; - var mut_ctx_obj: CtxStruct = .{ .factor = 10 }; + const CustomTool = struct { + tool: Tool, + calls: *std.atomic.Value(usize), + deinited: *bool, - const tool_impl = struct { - pub fn const_ctx_func(tool_allocator: Allocator, val: i64, ctx: *const CtxStruct) ![]const u8 { - return try std.fmt.allocPrint(tool_allocator, "{d}", .{val * ctx.factor}); + const vtable: Tool.VTable = .{ + .execute = executeFn, + .deinit = deinitFn, + }; + + fn initTool(self: *@This(), calls: *std.atomic.Value(usize), deinited: *bool) void { + self.* = .{ + .tool = .{ + .descriptor = desc, + .vtable = &vtable, + }, + .calls = calls, + .deinited = deinited, + }; } - pub fn optional_ctx_func(tool_allocator: Allocator, val: i64, ctx: ?*const CtxStruct) ![]const u8 { - const factor = if (ctx) |c| c.factor else 1; - return try std.fmt.allocPrint(tool_allocator, "{d}", .{val * factor}); + + fn executeFn(tool: *const Tool, alloc: Allocator, _: Io, _: ToolCallContext, _: []const Argument) Tool.CallError![]const u8 { + const self: *const @This() = @fieldParentPtr("tool", tool); + _ = self.calls.fetchAdd(1, .monotonic); + return try alloc.dupe(u8, "custom vtable called"); } - }; - const args = [_]Argument{.{ .name = "val", .value = .{ .integer = 5 } }}; + fn deinitFn(tool: *Tool) void { + const self: *const @This() = @fieldParentPtr("tool", tool); + self.deinited.* = true; + } + }; - // const -> const - const tool_const1 = initWithContext(tool_descriptor, tool_impl.const_ctx_func, &const_ctx_obj); - var res_const1 = try tool_const1.execute(allocator, io, &session_state, "id1", &args); - defer res_const1.deinit(); - try std.testing.expectEqualStrings("50", res_const1.result); + var calls = std.atomic.Value(usize).init(0); + var deinited = false; + var custom: CustomTool = undefined; + custom.initTool(&calls, &deinited); - // non-const -> const - const tool_const2 = initWithContext(tool_descriptor, tool_impl.const_ctx_func, &mut_ctx_obj); - var res_const2 = try tool_const2.execute(allocator, io, &session_state, "id2", &args); - defer res_const2.deinit(); - try std.testing.expectEqualStrings("50", res_const2.result); + var res = try custom.tool.execute(allocator, io, &session_state, "id_vtable", &.{}); + defer res.deinit(); - // const -> optional const - const tool_opt1 = initWithContext(tool_descriptor, tool_impl.optional_ctx_func, &const_ctx_obj); - var res_opt1 = try tool_opt1.execute(allocator, io, &session_state, "id3", &args); - defer res_opt1.deinit(); - try std.testing.expectEqualStrings("50", res_opt1.result); + try std.testing.expectEqual(@as(usize, 1), calls.load(.monotonic)); + try std.testing.expectEqualStrings("custom vtable called", res.result); - // non-const -> optional const - const tool_opt2 = initWithContext(tool_descriptor, tool_impl.optional_ctx_func, &mut_ctx_obj); - var res_opt2 = try tool_opt2.execute(allocator, io, &session_state, "id4", &args); - defer res_opt2.deinit(); - try std.testing.expectEqualStrings("50", res_opt2.result); + try std.testing.expect(!deinited); + custom.tool.deinit(); + try std.testing.expect(deinited); } -test "execute - multiple context pointers" { - const allocator = std.testing.allocator; - const io = std.testing.io; - var session_state = SessionState.init(allocator); - defer session_state.deinit(); - - const tool_descriptor: llm.types.Tool = .{ - .name = "multi_ctx_func", - .description = "Takes multiple context parameters of the same type", - .parameters = &.{ - .{ - .name = "val", - .type = .integer, - .required = true, - .description = "val", - }, - }, +test "makeExecuteFn rejects *Tool parameter" { + const desc: llm.types.Tool = .{ + .name = "mut_tool_test", + .description = "desc", + .parameters = &.{}, }; + const Impl = struct { + fn run(_: *Tool) ![]const u8 { + return ""; + } + }; + const res = comptime makeExecuteFn(desc, Impl.run); + try std.testing.expectEqual(ValidationError.UnsupportedType, res.err.code); + try std.testing.expectEqualStrings("Tool execution function parameter cannot be '*Tool'. Use '*const Tool' instead, as tools receive immutable references during execution.", res.err.msg); +} - const CtxStruct = struct { factor: i64 }; - var ctx_obj: CtxStruct = .{ .factor = 7 }; - +test "Tool.deinit on Tool.fromFn is safe no-op" { + const descriptor: llm.types.Tool = .{ + .name = "test_init_deinit", + .description = "test init deinit", + .parameters = &.{}, + }; const tool_impl = struct { - pub fn multi_ctx_func(tool_allocator: Allocator, ctx1: *CtxStruct, val: i64, ctx2: *CtxStruct) ![]const u8 { - return try std.fmt.allocPrint(tool_allocator, "{d}", .{val + ctx1.factor + ctx2.factor}); + pub fn run(_: Allocator) ![]const u8 { + return ""; } }; - const tool = initWithContext(tool_descriptor, tool_impl.multi_ctx_func, &ctx_obj); - const args = [_]Argument{.{ .name = "val", .value = .{ .integer = 10 } }}; - - var res = try tool.execute(allocator, io, &session_state, "id_multi", &args); - defer res.deinit(); - try std.testing.expectEqualStrings("24", res.result); -} - -test isContextTypeCompatible { - const Foo = struct { val: u32 }; - const Bar = struct { val: u32 }; - try std.testing.expect(isContextTypeCompatible(*Foo, *Foo)); - try std.testing.expect(isContextTypeCompatible(*Foo, *const Foo)); - try std.testing.expect(isContextTypeCompatible(*Foo, ?*Foo)); - try std.testing.expect(isContextTypeCompatible(*Foo, *anyopaque)); - try std.testing.expect(isContextTypeCompatible(*anyopaque, *Foo)); - try std.testing.expect(!isContextTypeCompatible(*Foo, *Bar)); - try std.testing.expect(!isContextTypeCompatible(*const Foo, *Foo)); + var tool = Tool.fromFn(descriptor, tool_impl.run); + tool.deinit(); } -test "execute with ToolCallContext" { +test "Tool propagates error.Canceled on execution" { const allocator = std.testing.allocator; const io = std.testing.io; var session_state = SessionState.init(allocator); defer session_state.deinit(); - const CounterState = struct { - count: i64, - fn initFn(self: *@This(), _: Allocator) !void { - self.count = 0; - } - }; - - const tool_descriptor: llm.types.Tool = .{ - .name = "counter_tool", - .description = "Increments counter in SessionState", - .parameters = &.{ - .{ - .name = "increment", - .type = .integer, - .required = true, - .description = "amount to increment", - }, - }, + const descriptor: llm.types.Tool = .{ + .name = "test_canceled_tool", + .description = "tool returning error.Canceled", + .parameters = &.{}, }; - const tool_impl = struct { - pub fn increment(alloc: Allocator, call_ctx: ToolCallContext, inc: i64) CallError![]const u8 { - const counter = call_ctx.getOrInitState(CounterState, CounterState.initFn) catch return error.OutOfMemory; - counter.count += inc; - return try std.fmt.allocPrint(alloc, "Count: {d}", .{counter.count}); - } - pub fn read_const(alloc: Allocator, call_ctx: ToolCallContext, inc: i64) CallError![]const u8 { - const count = if (call_ctx.getState(CounterState)) |c| c.count else 0; - return try std.fmt.allocPrint(alloc, "ConstCount: {d}", .{count + inc}); + pub fn run(_: Allocator) (Allocator.Error || std.Io.Cancelable)![]const u8 { + return error.Canceled; } }; - const args = [_]Argument{.{ .name = "increment", .value = .{ .integer = 5 } }}; - - const tool_inc = init(tool_descriptor, tool_impl.increment); - var res1 = try tool_inc.execute(allocator, io, &session_state, "id1", &args); - defer res1.deinit(); - try std.testing.expectEqualStrings("Count: 5", res1.result); - - var res2 = try tool_inc.execute(allocator, io, &session_state, "id2", &args); - defer res2.deinit(); - try std.testing.expectEqualStrings("Count: 10", res2.result); + var tool = Tool.fromFn(descriptor, tool_impl.run); + defer tool.deinit(); - const tool_read = init(tool_descriptor, tool_impl.read_const); - var res3 = try tool_read.execute(allocator, io, &session_state, "id3", &args); - defer res3.deinit(); - try std.testing.expectEqualStrings("ConstCount: 15", res3.result); + try std.testing.expectError(error.Canceled, tool.execute(allocator, io, &session_state, "id_cancel", &.{})); } diff --git a/src/agent/ToolCallContext.zig b/src/agent/ToolCallContext.zig index eea08e9..00e7204 100644 --- a/src/agent/ToolCallContext.zig +++ b/src/agent/ToolCallContext.zig @@ -52,7 +52,7 @@ test setInjectedContext { var session_state = SessionState.init(std.testing.allocator); defer session_state.deinit(); - const tool = Tool.init(.{ + const tool = Tool.fromFn(.{ .name = "dummy_tool", .description = "A dummy tool", .parameters = &.{}, @@ -75,7 +75,7 @@ test clearInjectedContext { var session_state = SessionState.init(std.testing.allocator); defer session_state.deinit(); - const tool = Tool.init(.{ + const tool = Tool.fromFn(.{ .name = "dummy_tool", .description = "A dummy tool", .parameters = &.{}, @@ -101,7 +101,7 @@ test getInjectedContext { var session_state = SessionState.init(std.testing.allocator); defer session_state.deinit(); - const tool = Tool.init(.{ + const tool = Tool.fromFn(.{ .name = "dummy_tool", .description = "A dummy tool", .parameters = &.{}, @@ -126,7 +126,7 @@ test getOrInitState { var session_state = SessionState.init(std.testing.allocator); defer session_state.deinit(); - const tool = Tool.init(.{ + const tool = Tool.fromFn(.{ .name = "dummy_tool", .description = "A dummy tool", .parameters = &.{}, @@ -156,7 +156,7 @@ test getState { var session_state = SessionState.init(std.testing.allocator); defer session_state.deinit(); - const tool = Tool.init(.{ + const tool = Tool.fromFn(.{ .name = "dummy_tool", .description = "A dummy tool", .parameters = &.{}, diff --git a/src/agent/root.zig b/src/agent/root.zig index 3385a70..640e2b3 100644 --- a/src/agent/root.zig +++ b/src/agent/root.zig @@ -4,8 +4,10 @@ pub const Session = @import("Session.zig"); pub const ToolCallContext = @import("ToolCallContext.zig"); pub const Tool = @import("Tool.zig"); pub const types = @import("types.zig"); +pub const testing = @import("testing/root.zig"); test { std.testing.refAllDecls(@This()); std.testing.refAllDecls(@import("tool/root.zig")); + std.testing.refAllDecls(@import("testing/root.zig")); } diff --git a/src/agent/testing/BlockingTool.zig b/src/agent/testing/BlockingTool.zig new file mode 100644 index 0000000..a371e3f --- /dev/null +++ b/src/agent/testing/BlockingTool.zig @@ -0,0 +1,80 @@ +//! A test tool implementation that blocks until `should_stop` is set to true +//! or the underlying I/O context is canceled. + +const std = @import("std"); +const Io = std.Io; +const llm = @import("llm"); +const Tool = @import("../Tool.zig"); + +const BlockingTool = @This(); + +/// Full descriptor for the blocking test tool. +pub const descriptor: llm.types.Tool = .{ + .name = "blocking_tool", + .description = "A test tool that blocks until signaled to stop", + .parameters = &.{}, +}; + +const vtable: Tool.VTable = .{ + .execute = Tool.wrapExecute(descriptor, execute), +}; + +pub const State = struct { + started: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), + should_stop: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), + + pub fn isStarted(self: *const State) bool { + return self.started.load(.monotonic); + } + + pub fn waitUntilStarted(self: *const State) void { + var attempts: usize = 0; + while (!self.isStarted() and attempts < 1_000_000) : (attempts += 1) { + std.atomic.spinLoopHint(); + } + } + + pub fn stop(self: *State) void { + self.should_stop.store(true, .monotonic); + } +}; + +tool: Tool, +state: *State, +io: Io, + +pub fn init(state: *State, io: Io) BlockingTool { + return .{ + .tool = .{ + .descriptor = descriptor, + .vtable = &vtable, + }, + .state = state, + .io = io, + }; +} + +/// Execution function using @fieldParentPtr. +pub fn execute(tool_ptr: *const Tool, alloc: std.mem.Allocator) Tool.CallError![]const u8 { + const self: *const BlockingTool = @fieldParentPtr("tool", tool_ptr); + self.state.started.store(true, .monotonic); + while (!self.state.should_stop.load(.monotonic)) { + try self.io.checkCancel(); + } + return try alloc.dupe(u8, "blocking tool done"); +} + +/// Returns true if the tool execution function has been entered. +pub fn isStarted(self: *const BlockingTool) bool { + return self.state.isStarted(); +} + +/// Blocks until the tool execution function is entered. +pub fn waitUntilStarted(self: *const BlockingTool) void { + self.state.waitUntilStarted(); +} + +/// Signals the tool execution to stop blocking and exit. +pub fn stop(self: *const BlockingTool) void { + self.state.stop(); +} diff --git a/src/agent/testing/root.zig b/src/agent/testing/root.zig new file mode 100644 index 0000000..1f0b6f5 --- /dev/null +++ b/src/agent/testing/root.zig @@ -0,0 +1,7 @@ +pub const BlockingTool = @import("BlockingTool.zig"); + +const std = @import("std"); + +test { + std.testing.refAllDecls(@This()); +} diff --git a/src/agent/tool/root.zig b/src/agent/tool/root.zig index cbee524..df44a16 100644 --- a/src/agent/tool/root.zig +++ b/src/agent/tool/root.zig @@ -3,7 +3,10 @@ const std = @import("std"); /// A simple todo list tool for managing tasks. pub const Todo = @import("todo.zig").Tool; +/// Shell tool utilities and pre-configured shell tools. +pub const Shell = @import("shell.zig"); + test { std.testing.refAllDecls(@This()); + std.testing.refAllDecls(@import("shell.zig")); } - diff --git a/src/agent/tool/shell.zig b/src/agent/tool/shell.zig new file mode 100644 index 0000000..95f3ad8 --- /dev/null +++ b/src/agent/tool/shell.zig @@ -0,0 +1,684 @@ +//! Agent tool for executing shell commands. + +const std = @import("std"); +const builtin = @import("builtin"); +const agent = @import("../root.zig"); +const SessionState = @import("../SessionState.zig"); + +pub const MAX_STREAM_LINES: usize = 2000; +pub const MAX_STREAM_BYTES: usize = 64 * 1024; // 64 KiB display limit per stream +pub const MAX_PROCESS_STREAM_BYTES: usize = 10 * 1024 * 1024; // 10 MiB limit per process stream +pub const DEFAULT_TIMEOUT_SECS: i64 = 30; + +/// Structured result of a shell execution. +const ShellResult = struct { + stdout: []const u8, + stderr: []const u8, + exit_code: ?u32 = null, + error_message: ?[]const u8 = null, +}; + +/// Constructs an argv builder function for POSIX shells (` -c `). +pub fn makePosixArgv(comptime shell_binary: []const u8) fn (std.mem.Allocator, []const u8) std.mem.Allocator.Error![]const []const u8 { + return struct { + fn make(allocator: std.mem.Allocator, command: []const u8) std.mem.Allocator.Error![]const []const u8 { + const argv = try allocator.alloc([]const u8, 3); + argv[0] = shell_binary; + argv[1] = "-c"; + argv[2] = command; + return argv; + } + }.make; +} + +/// Constructs an argv builder function for Windows Command Prompt (`cmd.exe /c `). +pub fn makeCmdArgv(allocator: std.mem.Allocator, command: []const u8) std.mem.Allocator.Error![]const []const u8 { + const argv = try allocator.alloc([]const u8, 3); + argv[0] = "cmd.exe"; + argv[1] = "/c"; + argv[2] = command; + return argv; +} + +/// Constructs an argv builder function for PowerShell (` -NoProfile -NonInteractive -Command `). +pub fn makePowershellArgv(comptime pwsh_binary: []const u8) fn (std.mem.Allocator, []const u8) std.mem.Allocator.Error![]const []const u8 { + return struct { + fn make(allocator: std.mem.Allocator, command: []const u8) std.mem.Allocator.Error![]const []const u8 { + const argv = try allocator.alloc([]const u8, 5); + argv[0] = pwsh_binary; + argv[1] = "-NoProfile"; + argv[2] = "-NonInteractive"; + argv[3] = "-Command"; + argv[4] = command; + return argv; + } + }.make; +} + +/// Creates an agent `Tool` configured for the specified shell name and argv builder function. +pub fn makeTool( + comptime shell_name: []const u8, + comptime make_argv_fn: fn (allocator: std.mem.Allocator, command: []const u8) std.mem.Allocator.Error![]const []const u8, +) agent.Tool { + const description = std.fmt.comptimePrint( + \\Execute a shell command in the current dir. + \\Commands run under `{s}` - write command strings in that shell's syntax. + \\Execution is non-interactive; stdin is ignored. + \\Returns an object with stdout and stderr as separate fields. + \\The output of each stream is limited to up to {d} lines ({d} KiB); longer outputs will be saved to a temporary file. + , .{ shell_name, MAX_STREAM_LINES, MAX_STREAM_BYTES / 1024 }); + + const S = struct { + fn execute(allocator: std.mem.Allocator, io: std.Io, command: []const u8, timeout_secs: ?i64) (std.mem.Allocator.Error || std.Io.Cancelable)![]const u8 { + const argv = try make_argv_fn(allocator, command); + defer allocator.free(argv); + + const effective_timeout_secs: i64 = if (timeout_secs) |t| (if (t > 0) t else DEFAULT_TIMEOUT_SECS) else DEFAULT_TIMEOUT_SECS; + const dur: std.Io.Duration = .fromSeconds(effective_timeout_secs); + const timeout: std.Io.Timeout = (std.Io.Timeout{ .duration = .{ .raw = dur, .clock = .awake } }).toDeadline(io); + + const run_res = std.process.run(allocator, io, .{ + .argv = argv, + .timeout = timeout, + .stdout_limit = std.Io.Limit.limited(MAX_PROCESS_STREAM_BYTES), + .stderr_limit = std.Io.Limit.limited(MAX_PROCESS_STREAM_BYTES), + }); + + if (run_res) |res| { + defer { + allocator.free(res.stdout); + allocator.free(res.stderr); + } + + const processed_stdout = try processStream(allocator, io, "stdout", res.stdout); + defer allocator.free(processed_stdout); + + const processed_stderr = try processStream(allocator, io, "stderr", res.stderr); + defer allocator.free(processed_stderr); + + const exit_code: ?u32 = switch (res.term) { + .exited => |code| code, + else => null, + }; + + const term_error_msg: ?[]const u8 = switch (res.term) { + .exited => null, + .signal => |sig| try std.fmt.allocPrint(allocator, "Process terminated by signal {d}", .{@intFromEnum(sig)}), + .stopped => |sig| try std.fmt.allocPrint(allocator, "Process stopped by signal {d}", .{sig}), + .unknown => |val| try std.fmt.allocPrint(allocator, "Process terminated abnormally (status {d})", .{val}), + }; + defer if (term_error_msg) |msg| allocator.free(msg); + + const shell_res: ShellResult = .{ + .stdout = processed_stdout, + .stderr = processed_stderr, + .exit_code = exit_code, + .error_message = term_error_msg, + }; + + return try formatJson(allocator, shell_res); + } else |err| { + const error_msg = switch (err) { + error.OutOfMemory, error.Canceled => |e| return e, + error.Timeout => try std.fmt.allocPrint(allocator, "Command timed out after {d} seconds", .{effective_timeout_secs}), + error.StreamTooLong => try std.fmt.allocPrint(allocator, "Command output exceeded maximum process stream limit ({d} bytes)", .{MAX_PROCESS_STREAM_BYTES}), + else => try std.fmt.allocPrint(allocator, "Failed to execute command: {}", .{err}), + }; + defer allocator.free(error_msg); + + const shell_res: ShellResult = .{ + .stdout = "", + .stderr = "", + .exit_code = null, + .error_message = error_msg, + }; + + return try formatJson(allocator, shell_res); + } + } + }; + + return agent.Tool.fromFn(.{ + .name = "shell", + .description = description, + .parameters = &.{ + .{ + .name = "command", + .description = "The command to execute.", + .type = .string, + .required = true, + }, + .{ + .name = "timeout_secs", + .description = std.fmt.comptimePrint( + "The number of seconds to wait for the command to complete before terminating. If omitted, defaults to {d} seconds.", + .{DEFAULT_TIMEOUT_SECS}, + ), + .type = .integer, + }, + }, + }, S.execute); +} + +/// Pre-configured shell tools. +pub const sh = makeTool("sh", makePosixArgv("sh")); +pub const bash = makeTool("bash", makePosixArgv("bash")); +pub const zsh = makeTool("zsh", makePosixArgv("zsh")); +pub const cmd = makeTool("cmd", makeCmdArgv); +pub const powershell = makeTool("powershell", makePowershellArgv("powershell.exe")); +pub const pwsh = makeTool("pwsh", makePowershellArgv("pwsh")); + +/// Default shell tool based on the host operating system. +pub const default = if (builtin.os.tag == .windows) cmd else sh; + +/// Replaces invalid UTF-8 byte sequences with Unicode replacement character U+FFFD. +fn sanitizeUtf8(allocator: std.mem.Allocator, input: []const u8) std.mem.Allocator.Error![]const u8 { + if (std.unicode.utf8ValidateSlice(input)) { + return try allocator.dupe(u8, input); + } + + var out: std.ArrayList(u8) = try .initCapacity(allocator, input.len); + defer out.deinit(allocator); + + var i: usize = 0; + while (i < input.len) { + const cp_len = std.unicode.utf8ByteSequenceLength(input[i]) catch { + try out.appendSlice(allocator, "\xEF\xBF\xBD"); + i += 1; + continue; + }; + + if (i + cp_len > input.len) { + try out.appendSlice(allocator, "\xEF\xBF\xBD"); + break; + } + + if (std.unicode.utf8Decode(input[i .. i + cp_len])) |_| { + try out.appendSlice(allocator, input[i .. i + cp_len]); + i += cp_len; + } else |_| { + try out.appendSlice(allocator, "\xEF\xBF\xBD"); + i += 1; + } + } + + return try out.toOwnedSlice(allocator); +} + +fn getTempDirAlloc(allocator: std.mem.Allocator) std.mem.Allocator.Error![]const u8 { + if (builtin.os.tag == .windows) { + const env: std.process.Environ = .{ .block = .global }; + inline for (.{ "TEMP", "TMP" }) |var_name| { + if (env.getAlloc(allocator, var_name)) |val| { + if (val.len > 0) return val; + allocator.free(val); + } else |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + else => {}, + } + } + if (env.getAlloc(allocator, "USERPROFILE")) |profile| { + defer allocator.free(profile); + if (profile.len > 0) { + return try std.fs.path.join(allocator, &.{ profile, "AppData", "Local", "Temp" }); + } + } else |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + else => {}, + } + return try allocator.dupe(u8, "C:\\Windows\\Temp"); + } + return try allocator.dupe(u8, "/tmp"); +} + +/// Creates a temporary file with the given prefix and writes the provided content into it. +/// Returns the allocated path of the created temporary file, which the caller is responsible for freeing. +fn writeTempFile(allocator: std.mem.Allocator, io: std.Io, prefix: []const u8, content: []const u8) ![]const u8 { + var rand_bytes: [8]u8 = undefined; + io.random(&rand_bytes); + const rand_val = std.mem.readInt(u64, &rand_bytes, .little); + const temp_dir = try getTempDirAlloc(allocator); + defer allocator.free(temp_dir); + + const file_name = try std.fmt.allocPrint(allocator, "{s}_{x}.txt", .{ prefix, rand_val }); + defer allocator.free(file_name); + + const temp_path = try std.fs.path.join(allocator, &.{ temp_dir, file_name }); + errdefer allocator.free(temp_path); + + const file = try std.Io.Dir.createFileAbsolute(io, temp_path, .{}); + defer file.close(io); + + var offset: u64 = 0; + while (offset < content.len) { + const written = try file.writePositional(io, &.{content[offset..]}, offset); + if (written == 0) return error.WriteFailed; + offset += written; + } + + return temp_path; +} + +const CutoffReason = enum { lines, bytes }; +const CutoffResult = struct { + cutoff: usize, + reason: CutoffReason, +}; + +fn processStream(allocator: std.mem.Allocator, io: std.Io, stream_name: []const u8, content: []const u8) (std.mem.Allocator.Error || std.Io.Cancelable)![]const u8 { + const maybe_cutoff = findCutoff(content, MAX_STREAM_LINES, MAX_STREAM_BYTES); + if (maybe_cutoff == null) { + return try sanitizeUtf8(allocator, content); + } + + const cutoff_res = maybe_cutoff.?; + const truncated_raw = content[0..cutoff_res.cutoff]; + const sanitized_truncated = try sanitizeUtf8(allocator, truncated_raw); + defer allocator.free(sanitized_truncated); + + const prefix = if (std.mem.eql(u8, stream_name, "stdout")) "coma_shell_stdout" else "coma_shell_stderr"; + const temp_file_res = writeTempFile(allocator, io, prefix, content); + + const reason_str = switch (cutoff_res.reason) { + .lines => try std.fmt.allocPrint(allocator, "{d} lines", .{MAX_STREAM_LINES}), + .bytes => try std.fmt.allocPrint(allocator, "{d} KiB", .{MAX_STREAM_BYTES / 1024}), + }; + defer allocator.free(reason_str); + + if (temp_file_res) |temp_path| { + defer allocator.free(temp_path); + return try std.fmt.allocPrint( + allocator, + "{s}\n\n[Output truncated to {s}. Full output saved to {s}]", + .{ sanitized_truncated, reason_str, temp_path }, + ); + } else |err| switch (err) { + error.OutOfMemory, error.Canceled => |e| return e, + else => return try std.fmt.allocPrint( + allocator, + "{s}\n\n[Output truncated to {s}. Failed to save overflow file: {}]", + .{ sanitized_truncated, reason_str, err }, + ), + } +} + +fn findCutoff(content: []const u8, max_lines: usize, max_bytes: usize) ?CutoffResult { + var line_count: usize = 0; + for (content, 0..) |byte, idx| { + if (idx >= max_bytes) { + var cut = max_bytes; + while (cut > 0 and cut < content.len and (content[cut] & 0xC0) == 0x80) : (cut -= 1) {} + return .{ .cutoff = cut, .reason = .bytes }; + } + if (byte == '\n') { + line_count += 1; + if (line_count >= max_lines and idx + 1 < content.len) { + return .{ .cutoff = idx + 1, .reason = .lines }; + } + } + } + return null; +} + +fn formatJson(allocator: std.mem.Allocator, result: ShellResult) std.mem.Allocator.Error![]const u8 { + var write_buffer: std.Io.Writer.Allocating = .init(allocator); + defer write_buffer.deinit(); + var stringifier = std.json.Stringify{ + .writer = &write_buffer.writer, + .options = .{}, + }; + // std.json.Stringify writes to std.Io.Writer.Allocating, which only fails if memory allocation fails. + stringifier.write(result) catch return error.OutOfMemory; + return try write_buffer.toOwnedSlice(); +} + +const llm = @import("llm"); + +test "Shell tool descriptor metadata" { + const tool = default; + try std.testing.expectEqualStrings("shell", tool.descriptor.name); + try std.testing.expect(std.mem.indexOf(u8, tool.descriptor.description, "2000 lines (64 KiB)") != null); + try std.testing.expect(std.mem.indexOf(u8, tool.descriptor.description, "non-interactive") != null); + try std.testing.expectEqual(2, tool.descriptor.parameters.len); + try std.testing.expectEqualStrings("command", tool.descriptor.parameters[0].name); + try std.testing.expectEqual(llm.types.Tool.Param.Type.string, tool.descriptor.parameters[0].type); + try std.testing.expect(tool.descriptor.parameters[0].required); + try std.testing.expectEqualStrings("timeout_secs", tool.descriptor.parameters[1].name); + try std.testing.expectEqual(llm.types.Tool.Param.Type.integer, tool.descriptor.parameters[1].type); + try std.testing.expect(!tool.descriptor.parameters[1].required); + try std.testing.expect(std.mem.indexOf(u8, tool.descriptor.parameters[1].description, "defaults to 30 seconds") != null); +} + +test "Shell tool stdout execution" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "echo hello coma" }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_1", args); + defer result.deinit(); + + try std.testing.expectEqualStrings("shell", result.tool_name); + try std.testing.expectEqualStrings("call_1", result.id); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("hello coma\n", parsed.value.stdout); + try std.testing.expectEqualStrings("", parsed.value.stderr); + try std.testing.expectEqual(0, parsed.value.exit_code.?); + try std.testing.expectEqual(null, parsed.value.error_message); +} + +test "Shell tool stderr and non-zero exit code" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "echo error output >&2; exit 42" }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_2", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("", parsed.value.stdout); + try std.testing.expectEqualStrings("error output\n", parsed.value.stderr); + try std.testing.expectEqual(42, parsed.value.exit_code.?); + try std.testing.expectEqual(null, parsed.value.error_message); +} + +test "Shell tool timeout triggered" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "sleep 2" }, + }, + .{ + .name = "timeout_secs", + .value = .{ .integer = 1 }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_3", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expect(parsed.value.error_message != null); + try std.testing.expectEqualStrings("Command timed out after 1 seconds", parsed.value.error_message.?); + try std.testing.expectEqual(null, parsed.value.exit_code); +} + +test "Shell tool line cutoff and temp file overflow" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + // Generate 2005 lines of output + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "i=1; while [ $i -le 2005 ]; do echo line $i; i=$((i+1)); done" }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_4", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(0, parsed.value.exit_code.?); + try std.testing.expect(std.mem.indexOf(u8, parsed.value.stdout, "[Output truncated to 2000 lines. Full output saved to ") != null); + + // Extract the temp file path from stdout notice + const prefix = "Full output saved to "; + if (std.mem.indexOf(u8, parsed.value.stdout, prefix)) |idx| { + const start = idx + prefix.len; + if (std.mem.indexOf(u8, parsed.value.stdout[start..], "]")) |end_rel| { + const temp_path = parsed.value.stdout[start .. start + end_rel]; + // Verify file exists and clean it up + std.Io.Dir.deleteFileAbsolute(io, temp_path) catch {}; + } + } +} + +test "Shell tool missing required argument" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = default; + const empty_args: []const llm.types.Argument = &.{}; + try std.testing.expectError( + agent.Tool.CallError.RequiredArgumentMissing, + tool.execute(testing_allocator, io, &session_state, "call_err", empty_args), + ); +} + +test "Argv builders" { + const allocator = std.testing.allocator; + + const posix_builder = makePosixArgv("custom_sh"); + const posix_argv = try posix_builder(allocator, "echo 1"); + defer allocator.free(posix_argv); + try std.testing.expectEqual(3, posix_argv.len); + try std.testing.expectEqualStrings("custom_sh", posix_argv[0]); + try std.testing.expectEqualStrings("-c", posix_argv[1]); + try std.testing.expectEqualStrings("echo 1", posix_argv[2]); + + const cmd_argv = try makeCmdArgv(allocator, "dir"); + defer allocator.free(cmd_argv); + try std.testing.expectEqual(3, cmd_argv.len); + try std.testing.expectEqualStrings("cmd.exe", cmd_argv[0]); + try std.testing.expectEqualStrings("/c", cmd_argv[1]); + try std.testing.expectEqualStrings("dir", cmd_argv[2]); + + const pwsh_builder = makePowershellArgv("pwsh"); + const pwsh_argv = try pwsh_builder(allocator, "Get-Process"); + defer allocator.free(pwsh_argv); + try std.testing.expectEqual(5, pwsh_argv.len); + try std.testing.expectEqualStrings("pwsh", pwsh_argv[0]); + try std.testing.expectEqualStrings("-NoProfile", pwsh_argv[1]); + try std.testing.expectEqualStrings("-NonInteractive", pwsh_argv[2]); + try std.testing.expectEqualStrings("-Command", pwsh_argv[3]); + try std.testing.expectEqualStrings("Get-Process", pwsh_argv[4]); +} + +test "Custom makeTool" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const custom_tool = makeTool("sh", makePosixArgv("sh")); + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "printf 'custom tool output'" }, + }, + }; + + var result = try custom_tool.execute(testing_allocator, io, &session_state, "call_custom", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("custom tool output", parsed.value.stdout); + try std.testing.expectEqual(0, parsed.value.exit_code.?); +} + +test "findCutoff helper" { + const text_few_lines = "line 1\nline 2\nline 3\n"; + try std.testing.expectEqual(null, findCutoff(text_few_lines, 3, 100)); + try std.testing.expectEqual(null, findCutoff(text_few_lines, 5, 100)); + + // Cutoff by lines + const line_cutoff = findCutoff(text_few_lines, 2, 100); + try std.testing.expect(line_cutoff != null); + try std.testing.expectEqual(CutoffReason.lines, line_cutoff.?.reason); + try std.testing.expectEqualStrings("line 1\nline 2\n", text_few_lines[0..line_cutoff.?.cutoff]); + + // Cutoff by bytes + const byte_cutoff = findCutoff("abcdefghijklmnop", 10, 5); + try std.testing.expect(byte_cutoff != null); + try std.testing.expectEqual(CutoffReason.bytes, byte_cutoff.?.reason); + try std.testing.expectEqualStrings("abcde", "abcdefghijklmnop"[0..byte_cutoff.?.cutoff]); + + // Cutoff does not split multi-byte UTF-8 sequence + // '€' is 3 bytes: 0xE2 0x82 0xAC + const utf8_text = "ab€cd"; // 2 + 3 + 2 = 7 bytes + const cut_mid_char = findCutoff(utf8_text, 10, 4); // index 4 is in the middle of '€' + try std.testing.expect(cut_mid_char != null); + try std.testing.expectEqual(CutoffReason.bytes, cut_mid_char.?.reason); + try std.testing.expectEqualStrings("ab", utf8_text[0..cut_mid_char.?.cutoff]); +} + +test "sanitizeUtf8 helper" { + const allocator = std.testing.allocator; + + const valid = "Hello, world! 123"; + const s_valid = try sanitizeUtf8(allocator, valid); + defer allocator.free(s_valid); + try std.testing.expectEqualStrings(valid, s_valid); + + const invalid = "bad \xFF\xFE bytes"; + const s_invalid = try sanitizeUtf8(allocator, invalid); + defer allocator.free(s_invalid); + try std.testing.expectEqualStrings("bad \xEF\xBF\xBD\xEF\xBF\xBD bytes", s_invalid); +} + +test "Shell tool chatty command total timeout" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + // Command prints periodically for ~3s total; each output interval (0.2s) is well below the 1s timeout. + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "for i in 1 2 3 4 5 6 7 8 9 10; do echo $i; sleep 0.2; done" }, + }, + .{ + .name = "timeout_secs", + .value = .{ .integer = 1 }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_chatty_timeout", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expect(parsed.value.error_message != null); + try std.testing.expectEqualStrings("Command timed out after 1 seconds", parsed.value.error_message.?); + try std.testing.expectEqual(null, parsed.value.exit_code); +} + +test "Shell tool single long line byte cap truncation" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + // Generate a single 70 KiB string without newlines + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "printf '%*s' 70000 '' | tr ' ' 'a'" }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_byte_cap", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(0, parsed.value.exit_code.?); + try std.testing.expect(std.mem.indexOf(u8, parsed.value.stdout, "[Output truncated to 64 KiB. Full output saved to ") != null); + + // Clean up temporary file + const prefix = "Full output saved to "; + if (std.mem.indexOf(u8, parsed.value.stdout, prefix)) |idx| { + const start = idx + prefix.len; + if (std.mem.indexOf(u8, parsed.value.stdout[start..], "]")) |end_rel| { + const temp_path = parsed.value.stdout[start .. start + end_rel]; + std.Io.Dir.deleteFileAbsolute(io, temp_path) catch {}; + } + } +} + +test "Shell tool non-UTF8 output sanitization" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const tool = sh; + const args: []const llm.types.Argument = &.{ + .{ + .name = "command", + .value = .{ .string = "printf 'hello \\377\\376 world'" }, + }, + }; + + var result = try tool.execute(testing_allocator, io, &session_state, "call_non_utf8", args); + defer result.deinit(); + + const parsed = try std.json.parseFromSlice(ShellResult, testing_allocator, result.result, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(0, parsed.value.exit_code.?); + try std.testing.expectEqualStrings("hello \u{FFFD}\u{FFFD} world", parsed.value.stdout); +} + +test "writeTempFile creates and writes file content" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + + const sample_content = "This is a test temporary file content.\n"; + const temp_path = try writeTempFile(testing_allocator, io, "coma_test_file", sample_content); + defer testing_allocator.free(temp_path); + defer std.Io.Dir.deleteFileAbsolute(io, temp_path) catch {}; + + const file = try std.Io.Dir.openFileAbsolute(io, temp_path, .{ .mode = .read_only }); + defer file.close(io); + + var read_buf: [100]u8 = undefined; + const slices = [_][]u8{&read_buf}; + const bytes_read = try file.readPositional(io, &slices, 0); + try std.testing.expectEqualStrings(sample_content, read_buf[0..bytes_read]); +} diff --git a/src/agent/tool/todo.zig b/src/agent/tool/todo.zig index a9ff243..0151a14 100644 --- a/src/agent/tool/todo.zig +++ b/src/agent/tool/todo.zig @@ -6,7 +6,7 @@ const agent = @import("../root.zig"); const SessionState = @import("../SessionState.zig"); /// `write_todo` tool for creating, updating, or clearing the persistent TODO list. -pub const Tool = agent.Tool.init(.{ +pub const Tool = agent.Tool.fromFn(.{ .name = "write_todo", .description = \\Overwrite the entire TODO content. @@ -16,7 +16,8 @@ pub const Tool = agent.Tool.init(.{ \\ - Task tracking and progress updates \\ - Important notes and reminders \\ - \\Writing an empty string clears the entire todo list. + \\Writing an empty string clears the entire todo list. Consider clearing + \\the todo list when starting a new task. \\ \\WARNING: This operation completely replaces the existing content. \\Always include all content you want to keep, not just the changes. diff --git a/src/agent/types.zig b/src/agent/types.zig index 45af69c..67eb128 100644 --- a/src/agent/types.zig +++ b/src/agent/types.zig @@ -9,7 +9,7 @@ pub const SessionConfig = struct { /// The LLM model to be used by the Session. model: llm.types.Model, /// A list of executable tools available for the Session to use. - tools: []const Tool = &.{}, + tools: []const *const Tool = &.{}, /// Optional system prompt for the Session. system_prompt: ?[]const u8 = null, }; @@ -27,14 +27,14 @@ pub const Turn = struct { /// the model. pub const IntermediateStepResult = union(enum) { /// The result of an LLM generation step. - step_result: llm.types.StepResult, + step_result: *llm.types.StepResult, /// The result of executing a tool call requested by the model. tool_result: llm.types.ToolResult, /// Frees the resources associated with the intermediate step result. pub fn deinit(self: *IntermediateStepResult) void { switch (self.*) { - .step_result => |*step_result| step_result.deinit(), + .step_result => |step_result| step_result.deinit(), .tool_result => |*tool_result| tool_result.deinit(), } self.* = undefined; @@ -49,7 +49,7 @@ pub const TurnResult = struct { /// The history of intermediate steps (model thoughts, tool calls, and results) executed during the turn. intermediate_steps: []IntermediateStepResult, /// The final step result that completed the turn (typically the model's final text output). - final_step: llm.types.StepResult, + final_step: *llm.types.StepResult, /// Deinitializes the TurnResult and frees all associated memory. /// diff --git a/src/llm/Provider.tests.zig b/src/llm/Provider.tests.zig index cc29ea6..396549e 100644 --- a/src/llm/Provider.tests.zig +++ b/src/llm/Provider.tests.zig @@ -5,9 +5,9 @@ const testing = @import("testing"); test "Provider.listModels delegates to VTable" { const allocator = std.testing.allocator; var mock_impl = testing.MockProvider{}; - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); - var result = try prov.listModels(allocator); + const result = try prov.listModels(allocator); defer result.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.list_models_calls); @@ -18,7 +18,7 @@ test "Provider.executeStep delegates to VTable" { const allocator = std.testing.allocator; var mock_impl = testing.MockProvider{}; defer mock_impl.deinit(); - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); const model = llm.types.Model{ .id = "test-model-id", @@ -31,7 +31,7 @@ test "Provider.executeStep delegates to VTable" { const last_step_continuation = testing.MockProvider.stepContinuation(); - var outcome = try prov.executeStep(allocator, session_config, input_steps, last_step_continuation); + const outcome = try prov.executeStep(allocator, session_config, input_steps, last_step_continuation); defer outcome.result.deinit(); defer outcome.continuation.deinit(); @@ -39,12 +39,12 @@ test "Provider.executeStep delegates to VTable" { try std.testing.expectEqual(allocator, mock_impl.last_allocator.?); try std.testing.expectEqualStrings("test-model-id", mock_impl.last_session_config.?.model.id); try std.testing.expectEqualStrings("hello", mock_impl.last_input_steps.?[0].prompt); - try std.testing.expectEqual(last_step_continuation.ptr, mock_impl.last_previous_step.?.ptr); + try std.testing.expectEqual(last_step_continuation, mock_impl.last_previous_step.?); } test "Provider.deinit delegates to VTable" { var mock_impl = testing.MockProvider{}; - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); prov.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.deinit_calls); @@ -59,10 +59,11 @@ test "Provider.listModels returns custom success and error" { const dummy_models = &[_]llm.types.Model{ .{ .id = "custom-id", .display_name = "Custom Name" }, }; - mock_impl.list_models_result = testing.MockProvider.listModelsResult(dummy_models); + var expected_result = testing.MockProvider.listModelsResult(dummy_models); + mock_impl.list_models_result = &expected_result; - var prov = mock_impl.provider(); - var result = try prov.listModels(allocator); + const prov = mock_impl.provider(); + const result = try prov.listModels(allocator); defer result.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_impl.list_models_calls); @@ -74,7 +75,7 @@ test "Provider.listModels returns custom success and error" { var mock_impl = testing.MockProvider{}; mock_impl.list_models_result = error.HttpRequestFailed; - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); try std.testing.expectError(error.HttpRequestFailed, prov.listModels(allocator)); } } @@ -90,15 +91,15 @@ test "Provider.executeStep returns custom success and error" { const dummy_outputs = &[_]llm.types.ModelOutput{ .{ .text = "custom-output" }, }; - const step_result = testing.MockProvider.stepResult(dummy_outputs, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(dummy_outputs, &.{}, &.{}); const step_continuation = testing.MockProvider.stepContinuation(); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = step_continuation }, + .{ .result = &step_result, .continuation = step_continuation }, }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); - var outcome = try prov.executeStep(allocator, session_config, &.{}, null); + const prov = mock_impl.provider(); + const outcome = try prov.executeStep(allocator, session_config, &.{}, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); @@ -114,7 +115,7 @@ test "Provider.executeStep returns custom success and error" { }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); try std.testing.expectError(error.HttpRequestFailed, prov.executeStep(allocator, session_config, &.{}, null)); } } @@ -123,7 +124,7 @@ test "Provider.executeStepStreaming delegates to VTable" { const allocator = std.testing.allocator; var mock_impl = testing.MockProvider{}; defer mock_impl.deinit(); - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); const model = llm.types.Model{ .id = "test-model-id", @@ -143,7 +144,7 @@ test "Provider.executeStepStreaming delegates to VTable" { } }; - var outcome = try prov.executeStepStreaming(allocator, session_config, input_steps, prev_continuation, CallbackState.callback, null); + const outcome = try prov.executeStepStreaming(allocator, session_config, input_steps, prev_continuation, CallbackState.callback, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); @@ -152,7 +153,7 @@ test "Provider.executeStepStreaming delegates to VTable" { try std.testing.expectEqual(allocator, mock_impl.last_allocator.?); try std.testing.expectEqualStrings("test-model-id", mock_impl.last_session_config.?.model.id); try std.testing.expectEqualStrings("hello", mock_impl.last_input_steps.?[0].prompt); - try std.testing.expectEqual(prev_continuation.ptr, mock_impl.last_previous_step.?.ptr); + try std.testing.expectEqual(prev_continuation, mock_impl.last_previous_step.?); } test "Provider.executeStepStreaming returns custom success and error" { @@ -173,15 +174,15 @@ test "Provider.executeStepStreaming returns custom success and error" { const dummy_outputs = &[_]llm.types.ModelOutput{ .{ .text = "custom-output" }, }; - const step_result = testing.MockProvider.stepResult(dummy_outputs, &.{}, &.{}); + var step_result = testing.MockProvider.stepResult(dummy_outputs, &.{}, &.{}); const step_continuation = testing.MockProvider.stepContinuation(); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ - .{ .result = step_result, .continuation = step_continuation }, + .{ .result = &step_result, .continuation = step_continuation }, }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); - var outcome = try prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null); + const prov = mock_impl.provider(); + const outcome = try prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); @@ -197,7 +198,7 @@ test "Provider.executeStepStreaming returns custom success and error" { }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); + const prov = mock_impl.provider(); try std.testing.expectError(error.HttpRequestFailed, prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null)); } } diff --git a/src/llm/Provider.zig b/src/llm/Provider.zig index 0d6b3b1..7f8d35f 100644 --- a/src/llm/Provider.zig +++ b/src/llm/Provider.zig @@ -11,10 +11,17 @@ const StreamingCallback = types.StreamingCallback; const StepContinuation = types.StepContinuation; const StepOutcome = types.StepOutcome; +/// Intrusive interface for an LLM provider. +/// +/// Concrete providers (such as Gemini) implement this interface by embedding `Provider` as a field +/// (e.g. `interface: Provider = .{ .vtable = &vtable }`) and implementing the methods in `Provider.VTable`. +/// +/// **Pointer Stability & Memory Safety:** +/// Provider method implementations recover the enclosing provider struct using +/// `@fieldParentPtr("interface", provider)`. The enclosing provider instance MUST NOT be copied by +/// value or moved in memory once its pointer or interface reference is taken. const Provider = @This(); -/// Opaque pointer to the provider implementation's context. -ptr: *anyopaque, /// Pointer to the provider implementation's virtual function table. vtable: *const VTable, @@ -28,38 +35,38 @@ pub const ProviderError = error{ /// The virtual function table defining the methods that an LLM provider must implement. pub const VTable = struct { /// Lists all models available from the provider. - list_models: *const fn (ptr: *anyopaque, allocator: Allocator) ProviderError!ListModelsResult, + list_models: *const fn (provider: *Provider, allocator: Allocator) ProviderError!*ListModelsResult, /// Executes a single interaction step with the LLM. execute_step: *const fn ( - ptr: *anyopaque, + provider: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, ) ProviderError!StepOutcome, /// Executes a single interaction step with the LLM, streaming the response. execute_step_streaming: *const fn ( - ptr: *anyopaque, + provider: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, callback: StreamingCallback, callback_context: ?*anyopaque, ) ProviderError!StepOutcome, /// Frees the resources associated with the provider. - deinit: *const fn (ptr: *anyopaque) void, + deinit: *const fn (provider: *Provider) void, }; /// Lists all models available from the current provider. /// /// `allocator` is used to allocate the structures and strings in the returned `ListModelsResult`. /// The caller **MUST** call `deinit()` on the returned `ListModelsResult` to free the allocated memory. -pub fn listModels(provider: *Provider, allocator: Allocator) ProviderError!ListModelsResult { - return provider.vtable.list_models(provider.ptr, allocator); +pub fn listModels(provider: *Provider, allocator: Allocator) ProviderError!*ListModelsResult { + return provider.vtable.list_models(provider, allocator); } /// Executes a single step of interaction with the LLM provider. @@ -76,9 +83,9 @@ pub fn executeStep( allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, ) ProviderError!StepOutcome { - return provider.vtable.execute_step(provider.ptr, allocator, session_config, input, previous_step); + return provider.vtable.execute_step(provider, allocator, session_config, input, previous_step); } /// Executes a single step of interaction with the LLM provider, streaming chunks back to the callback. @@ -100,14 +107,14 @@ pub fn executeStepStreaming( allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, callback: StreamingCallback, callback_context: ?*anyopaque, ) ProviderError!StepOutcome { - return provider.vtable.execute_step_streaming(provider.ptr, allocator, session_config, input, previous_step, callback, callback_context); + return provider.vtable.execute_step_streaming(provider, allocator, session_config, input, previous_step, callback, callback_context); } /// Frees the resources associated with the provider. pub fn deinit(provider: *Provider) void { - provider.vtable.deinit(provider.ptr); + provider.vtable.deinit(provider); } diff --git a/src/llm/types.zig b/src/llm/types.zig index 06775c4..cf85c07 100644 --- a/src/llm/types.zig +++ b/src/llm/types.zig @@ -58,17 +58,15 @@ pub const ListModelsResult = struct { /// The list of available models. models: []const Model, - ptr: *anyopaque, vtable: *const VTable, pub const VTable = struct { - deinit: *const fn (ptr: *anyopaque) void, + deinit: *const fn (self: *ListModelsResult) void, }; /// Frees the resources associated with the ListModelsResult. pub fn deinit(self: *ListModelsResult) void { - self.vtable.deinit(self.ptr); - self.* = undefined; + self.vtable.deinit(self); } }; @@ -102,7 +100,7 @@ pub const ToolCall = struct { /// The name of the tool to be called. name: []const u8, /// The arguments provided for the tool call. - arguments: []Argument, + arguments: []const Argument, }; /// The result of executing a step (interaction) with the LLM. @@ -115,42 +113,38 @@ pub const StepResult = struct { /// A list of tool calls requested by the model. tool_calls: []const ToolCall, - ptr: *anyopaque, vtable: *const VTable, pub const VTable = struct { - deinit: *const fn (ptr: *anyopaque) void, + deinit: *const fn (self: *StepResult) void, }; /// Frees the resources associated with the StepResult. pub fn deinit(self: *StepResult) void { - self.vtable.deinit(self.ptr); - self.* = undefined; + self.vtable.deinit(self); } }; /// Represents information required to continue an interaction with an LLM. pub const StepContinuation = struct { - ptr: *anyopaque, vtable: *const VTable, pub const VTable = struct { - deinit: *const fn (ptr: *anyopaque) void, + deinit: *const fn (self: *StepContinuation) void, }; /// Frees the resources associated with the StepContinuation. pub fn deinit(self: *StepContinuation) void { - self.vtable.deinit(self.ptr); - self.* = undefined; + self.vtable.deinit(self); } }; /// Represents the outcome of executing a single interaction step with the LLM. pub const StepOutcome = struct { /// The result of the step execution containing generated outputs, thoughts, and tool calls. - result: StepResult, + result: *StepResult, /// The continuation required to maintain conversation state for subsequent steps. - continuation: StepContinuation, + continuation: *StepContinuation, }; /// Represents a single input step to the LLM. diff --git a/src/main.zig b/src/main.zig index 7487070..02e2e67 100644 --- a/src/main.zig +++ b/src/main.zig @@ -13,7 +13,7 @@ const MarkdownRendering = @import("MarkdownRendering.zig"); /// Loads the GEMINI_API_KEY from the environment variables. /// If not found, it attempts to read it from a `.env` file in the current working directory. /// Returns an allocated string containing the API key, or `error.ApiKeyMissing` if not found. -fn loadApiKey(allocator: std.mem.Allocator, io: std.Io, environ_map: *std.process.Environ.Map) ![]const u8 { +fn loadApiKey(allocator: std.mem.Allocator, io: std.Io, environ_map: *const std.process.Environ.Map) ![]const u8 { if (environ_map.get("GEMINI_API_KEY")) |env_val| { return try allocator.dupe(u8, env_val); } @@ -22,7 +22,7 @@ fn loadApiKey(allocator: std.mem.Allocator, io: std.Io, environ_map: *std.proces return try allocator.dupe(u8, env_val); } - var file = std.Io.Dir.openFile(.cwd(), io, ".env", .{}) catch |err| switch (err) { + const file = std.Io.Dir.openFile(.cwd(), io, ".env", .{}) catch |err| switch (err) { error.FileNotFound => { std.debug.print("Error: GEMINI_API_KEY and GOOGLE_API_KEY environment variable is not set, and no .env file was found.\n", .{}); return error.ApiKeyMissing; @@ -79,17 +79,46 @@ fn executeTypescript(allocator: std.mem.Allocator, io: std.Io, code: []const u8) return result.stdout; } -fn getWeather(allocator: std.mem.Allocator, zip_code: i64, ctx: *WeatherToolCtx) ![]const u8 { - const result_str = if (zip_code == 7302) - try allocator.dupe(u8, ctx.weather_str) - else - try std.fmt.allocPrint(allocator, "Error: Weather data is only available for zip code 07302. Requested: {}", .{zip_code}); +const WeatherTool = struct { + tool: Tool, + weather_str: []const u8, - return result_str; -} + pub const descriptor: llm.types.Tool = .{ + .name = "get_weather", + .description = "Get the current weather for a given zip code.", + .parameters = &.{ + .{ + .name = "zip_code", + .type = .integer, + .required = true, + .description = "The 5-digit zip code to get the weather for.", + }, + }, + }; -const WeatherToolCtx = struct { - weather_str: []const u8, + const vtable: Tool.VTable = .{ + .execute = Tool.wrapExecute(descriptor, getWeather), + }; + + pub fn init(weather_str: []const u8) WeatherTool { + return .{ + .tool = .{ + .descriptor = descriptor, + .vtable = &vtable, + }, + .weather_str = weather_str, + }; + } + + fn getWeather(tool: *const Tool, allocator: std.mem.Allocator, zip_code: i64) ![]const u8 { + const self: *const WeatherTool = @fieldParentPtr("tool", tool); + const result_str = if (zip_code == 7302) + try allocator.dupe(u8, self.weather_str) + else + try std.fmt.allocPrint(allocator, "Error: Weather data is only available for zip code 07302. Requested: {}", .{zip_code}); + + return result_str; + } }; /// The main entry point of the application. @@ -110,10 +139,10 @@ pub fn main(init: std.process.Init) !void { var http_client: std.http.Client = .{ .allocator = allocator, .io = io }; defer http_client.deinit(); var gemini_client: provider.Gemini = try .init(allocator, &http_client, api_key); - var client = gemini_client.provider(); + const client = gemini_client.provider(); defer client.deinit(); - var models_list = try client.listModels(allocator); + const models_list = try client.listModels(allocator); defer models_list.deinit(); const models = models_list.models; var selected_model: ?llm.types.Model = null; @@ -124,37 +153,31 @@ pub fn main(init: std.process.Init) !void { } } else unreachable; - var weather_ctx: WeatherToolCtx = .{ .weather_str = "Weather report for 07302: Sunny, 72°F, Humidity 50%, Wind 5 mph" }; - - const tools = &[_]Tool{ - Tool.init(.{ - .name = "execute_typescript", - .description = "Executes typescript code and returns the output printed to stdout. Takes a single string argument.", - .parameters = &.{ - .{ - .name = "code", - .type = .string, - .required = true, - .description = "The typescript code to execute.", - }, - }, - }, executeTypescript), - Tool.initWithContext(.{ - .name = "get_weather", - .description = "Get the current weather for a given zip code.", - .parameters = &.{ - .{ - .name = "zip_code", - .type = .integer, - .required = true, - .description = "The 5-digit zip code to get the weather for.", - }, + var weather_tool = WeatherTool.init("Weather report for 07302: Sunny, 72°F, Humidity 50%, Wind 5 mph"); + + var typescript_tool = Tool.fromFn(.{ + .name = "execute_typescript", + .description = "Executes typescript code and returns the output printed to stdout. Takes a single string argument.", + .parameters = &.{ + .{ + .name = "code", + .type = .string, + .required = true, + .description = "The typescript code to execute.", }, - }, getWeather, &weather_ctx), - agent.Tool.BuiltIn.Todo, + }, + }, executeTypescript); + + const tools = [_]*const Tool{ + &typescript_tool, + &weather_tool.tool, + &agent.Tool.BuiltIn.Todo, + &agent.Tool.BuiltIn.Shell.default, }; + defer typescript_tool.deinit(); + defer weather_tool.tool.deinit(); - const session_config: types.SessionConfig = .{ .model = selected_model.?, .tools = tools, .system_prompt = "You're a helpful agent. The user can ask you questions and you can use your tools to answer them. When receiving a new request from the user, plan out how you will address the request and document the steps on your Todo list. Keep the todo list updated as you progress." }; + const session_config: types.SessionConfig = .{ .model = selected_model.?, .tools = &tools, .system_prompt = "You're a helpful agent. The user can ask you questions and you can use your tools to answer them. When receiving a new request from the user, plan out how you will address the request and document the steps on your Todo list. Keep the todo list updated as you progress and remove completed todo steps if they are no longer relevant." }; const args = try init.minimal.args.toSlice(allocator); defer allocator.free(args); diff --git a/src/provider/google/api.zig b/src/provider/google/api.zig index a4778f7..6a943a6 100644 --- a/src/provider/google/api.zig +++ b/src/provider/google/api.zig @@ -242,7 +242,7 @@ pub const FunctionArgument = struct { name: []const u8, value: llm.types.Argument.Value, - pub fn parseFromJsonObject(allocator: std.mem.Allocator, source: std.json.Value) ![]FunctionArgument { + pub fn parseFromJsonObject(allocator: std.mem.Allocator, source: std.json.Value) ![]const FunctionArgument { var arguments: std.ArrayList(FunctionArgument) = .empty; defer arguments.deinit(allocator); var argument_iterator = source.object.iterator(); @@ -264,7 +264,7 @@ pub const FunctionArgument = struct { pub const FunctionCall = struct { id: []const u8, name: []const u8, - arguments: []FunctionArgument, + arguments: []const FunctionArgument, pub fn jsonParseFromValue(allocator: std.mem.Allocator, source: std.json.Value, options: std.json.ParseOptions) !FunctionCall { const id = try std.json.innerParseFromValue([]const u8, allocator, source.object.get("id") orelse return error.MissingField, options); @@ -303,8 +303,8 @@ pub const StepType = enum { /// Represents a single step in a Gemini interaction, such as a thought, output, or function call. pub const Step = union(StepType) { - thought: []Content, - model_output: []Content, + thought: []const Content, + model_output: []const Content, function_call: FunctionCall, pub fn jsonParse(allocator: std.mem.Allocator, source: anytype, options: std.json.ParseOptions) !Step { @@ -319,7 +319,7 @@ pub const Step = union(StepType) { .thought => { if (source.object.get("summary")) |summary| { return Step{ - .thought = try std.json.innerParseFromValue([]Content, allocator, summary, options), + .thought = try std.json.innerParseFromValue([]const Content, allocator, summary, options), }; } else { return Step{ .thought = &.{} }; @@ -328,7 +328,7 @@ pub const Step = union(StepType) { .model_output => { if (source.object.get("content")) |content| { return Step{ - .model_output = try std.json.innerParseFromValue([]Content, allocator, content, options), + .model_output = try std.json.innerParseFromValue([]const Content, allocator, content, options), }; } else { return Step{ .model_output = &.{} }; @@ -344,7 +344,7 @@ pub const Step = union(StepType) { /// The full response from the Gemini API representing an interaction. pub const Interaction = struct { id: []const u8, - steps: []Step, + steps: []const Step, }; pub const StreamingInteraction = struct { diff --git a/src/provider/google/converter.zig b/src/provider/google/converter.zig index db2c650..df526fd 100644 --- a/src/provider/google/converter.zig +++ b/src/provider/google/converter.zig @@ -194,7 +194,7 @@ pub fn toArgument(arg: api.FunctionArgument) llm.types.Argument { /// Converts a slice of Gemini API function arguments to a generic argument slice. /// Caller owns the returned slice, but not the individual strings. -pub fn toArguments(allocator: Allocator, args: []const api.FunctionArgument) ![]llm.types.Argument { +pub fn toArguments(allocator: Allocator, args: []const api.FunctionArgument) ![]const llm.types.Argument { const result = try allocator.alloc(llm.types.Argument, args.len); for (args, 0..) |arg, i| { result[i] = toArgument(arg); @@ -204,7 +204,7 @@ pub fn toArguments(allocator: Allocator, args: []const api.FunctionArgument) ![] /// Converts a slice of Gemini API function arguments to a generic argument slice, duplicating all strings. /// Caller owns the returned slice and the individual strings. -pub fn dupeArguments(allocator: Allocator, args: []const api.FunctionArgument) ![]llm.types.Argument { +pub fn dupeArguments(allocator: Allocator, args: []const api.FunctionArgument) ![]const llm.types.Argument { const result = try allocator.alloc(llm.types.Argument, args.len); var initialized: usize = 0; errdefer { @@ -251,7 +251,7 @@ pub fn toModel(model: api.GeminiModel) llm.types.Model { /// Converts a slice of Gemini API models to a generic model slice. /// Caller owns the returned slice. -pub fn toModels(allocator: Allocator, models: []const api.GeminiModel) ![]llm.types.Model { +pub fn toModels(allocator: Allocator, models: []const api.GeminiModel) ![]const llm.types.Model { const result = try allocator.alloc(llm.types.Model, models.len); for (models, 0..) |model, i| { result[i] = toModel(model); diff --git a/src/provider/google/provider.zig b/src/provider/google/provider.zig index f211779..9be9b37 100644 --- a/src/provider/google/provider.zig +++ b/src/provider/google/provider.zig @@ -21,6 +21,7 @@ pub const Gemini = MakeProvider(*std.http.Client); /// Returns a type representing the provider instance. fn MakeProvider(comptime ClientType: type) type { return struct { + interface: Provider = .{ .vtable = &vtable }, /// Memory allocator used for copying the API key and other internal allocations. allocator: Allocator, /// The JSON client wrapper used to handle low-level HTTP requests and responses. @@ -30,6 +31,12 @@ fn MakeProvider(comptime ClientType: type) type { api_key: []const u8, const Self = @This(); + const vtable: Provider.VTable = .{ + .list_models = listModels, + .execute_step = executeStep, + .execute_step_streaming = executeStepStreaming, + .deinit = deinitVtable, + }; /// Initializes a new instance of the Gemini provider. /// @@ -52,34 +59,31 @@ fn MakeProvider(comptime ClientType: type) type { } /// Returns the generic `llm.Provider` interface for this Gemini instance. - pub fn provider(self: *Self) Provider { - return .{ .ptr = self, .vtable = &.{ - .list_models = listModels, - .execute_step = executeStep, - .execute_step_streaming = executeStepStreaming, - .deinit = deinit, - } }; + pub fn provider(self: *Self) *Provider { + return &self.interface; } /// Frees the resources associated with the provider, including the copied API key. - /// - /// `ctx` is an opaque pointer to the `Self` instance. - pub fn deinit(ctx: *anyopaque) void { - const self: *Self = @ptrCast(@alignCast(ctx)); + pub fn deinit(self: *Self) void { self.allocator.free(self.api_key); self.* = undefined; } + fn deinitVtable(prov: *Provider) void { + const self: *Self = @fieldParentPtr("interface", prov); + self.deinit(); + } + /// Lists all available Gemini models from the remote API. /// - /// `ctx` is an opaque pointer to the `Self` instance. + /// `prov` is a pointer to the `llm.Provider` interface. /// `allocator` is used to allocate internal structures. /// /// The caller **MUST** call `deinit()` on the returned `ListModelsResult` to free the allocated model details. /// /// Returns a `ListModelsResult` structure mapping the remote models, or a `ProviderError` on failure. - pub fn listModels(ctx: *anyopaque, allocator: Allocator) !llm.types.ListModelsResult { - const self: *Self = @ptrCast(@alignCast(ctx)); + pub fn listModels(prov: *Provider, allocator: Allocator) ProviderError!*llm.types.ListModelsResult { + const self: *Self = @fieldParentPtr("interface", prov); const str_uri = try makeUri(allocator, .{ .path = &.{"models"}, .api_key = self.api_key }); defer allocator.free(str_uri); @@ -103,21 +107,22 @@ fn MakeProvider(comptime ClientType: type) type { /// /// Returns the structured URI and `CreateInteractionRequest` payload. fn buildCreateInteractionRequest( - self: *Self, + self: *const Self, arena_allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, - previous_step: ?llm.types.StepContinuation, + previous_step: ?*const llm.types.StepContinuation, stream: bool, ) !struct { uri: std.Uri, payload: api.CreateInteractionRequest } { - const previous_gemini_step: ?*gemini_types.StepContinuation = if (previous_step) |step| @ptrCast(@alignCast(step.ptr)) else null; + const previous_gemini_step: ?*const gemini_types.StepContinuation = if (previous_step) |step| @fieldParentPtr("interface", step) else null; - var tools: []api.Tool = &.{}; + var tools: []const api.Tool = &.{}; if (session_config.tools.len > 0) { - tools = try arena_allocator.alloc(api.Tool, session_config.tools.len); + const tools_buf = try arena_allocator.alloc(api.Tool, session_config.tools.len); for (session_config.tools, 0..) |tool, i| { - tools[i] = try converter.toGoogleTool(arena_allocator, tool); + tools_buf[i] = try converter.toGoogleTool(arena_allocator, tool); } + tools = tools_buf; } const str_uri = try makeUri(arena_allocator, .{ .path = &.{"interactions"}, .api_key = self.api_key }); @@ -142,7 +147,7 @@ fn MakeProvider(comptime ClientType: type) type { /// Executes a single interaction step using the Gemini API, waiting for the full response. /// - /// `ctx` is an opaque pointer to the `Self` instance. + /// `prov` is a pointer to the `llm.Provider` interface. /// `allocator` is used to allocate the structures and strings in the returned `StepResult`. /// `session_config` provides the model ID and active tools configuration. /// `input` contains the array of steps to send as context/prompts. @@ -152,13 +157,13 @@ fn MakeProvider(comptime ClientType: type) type { /// /// Returns the `StepResult` detailing thoughts, outputs, and tool calls, or a `ProviderError` on failure. pub fn executeStep( - ctx: *anyopaque, + prov: *Provider, allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, - previous_step: ?llm.types.StepContinuation, - ) !llm.types.StepOutcome { - const self: *Self = @ptrCast(@alignCast(ctx)); + previous_step: ?*const llm.types.StepContinuation, + ) ProviderError!llm.types.StepOutcome { + const self: *Self = @fieldParentPtr("interface", prov); var arena = std.heap.ArenaAllocator.init(allocator); defer arena.deinit(); @@ -171,7 +176,7 @@ fn MakeProvider(comptime ClientType: type) type { /// Executes a single interaction step using the Gemini API, streaming chunks of the response back. /// - /// `ctx` is an opaque pointer to the `Self` instance. + /// `prov` is a pointer to the `llm.Provider` interface. /// `allocator` is used to allocate internal streaming state and structures in the final returned `StepResult`. /// `session_config` provides the model ID and active tools configuration. /// `input` contains the array of steps to send as context/prompts. @@ -186,15 +191,15 @@ fn MakeProvider(comptime ClientType: type) type { /// /// Returns the accumulated `StepResult` once the stream completes, or a `ProviderError` on failure. pub fn executeStepStreaming( - ctx: *anyopaque, + prov: *Provider, allocator: Allocator, session_config: llm.types.SessionConfig, input: []const llm.types.Step, - previous_step: ?llm.types.StepContinuation, + previous_step: ?*const llm.types.StepContinuation, callback: llm.types.StreamingCallback, callback_context: ?*anyopaque, - ) !llm.types.StepOutcome { - const self: *Self = @ptrCast(@alignCast(ctx)); + ) ProviderError!llm.types.StepOutcome { + const self: *Self = @fieldParentPtr("interface", prov); var arena = std.heap.ArenaAllocator.init(allocator); defer arena.deinit(); @@ -376,7 +381,7 @@ test "Gemini initialization and deinitialization" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); p.deinit(); } @@ -413,10 +418,10 @@ test "Gemini.listModels success" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); - var result = try p.listModels(allocator); + const result = try p.listModels(allocator); defer result.deinit(); try std.testing.expectEqual(1, result.models.len); @@ -443,7 +448,7 @@ test "Gemini.listModels HTTP failure" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); try std.testing.expectError(error.HttpRequestFailed, p.listModels(allocator)); @@ -499,7 +504,7 @@ test "Gemini.executeStep success" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -509,7 +514,7 @@ test "Gemini.executeStep success" { .{ .prompt = "Hello" }, }; - var outcome = try p.executeStep(allocator, config, input, null); + const outcome = try p.executeStep(allocator, config, input, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); const result = outcome.result; @@ -520,7 +525,7 @@ test "Gemini.executeStep success" { try std.testing.expectEqual(1, result.thoughts.len); try std.testing.expectEqualStrings("Thinking...", result.thoughts[0].text); - const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + const gemini_continuation: *const gemini_types.StepContinuation = @fieldParentPtr("interface", continuation); try std.testing.expectEqualStrings("interaction_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } @@ -594,7 +599,7 @@ test "Gemini.executeStep with previous step" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -605,7 +610,7 @@ test "Gemini.executeStep with previous step" { const input1 = &[_]llm.types.Step{ .{ .prompt = "Hello" }, }; - var outcome1 = try p.executeStep(allocator, config, input1, null); + const outcome1 = try p.executeStep(allocator, config, input1, null); defer outcome1.result.deinit(); defer outcome1.continuation.deinit(); const result1 = outcome1.result; @@ -617,7 +622,7 @@ test "Gemini.executeStep with previous step" { const input2 = &[_]llm.types.Step{ .{ .prompt = "Next prompt" }, }; - var outcome2 = try p.executeStep(allocator, config, input2, continuation1); + const outcome2 = try p.executeStep(allocator, config, input2, continuation1); defer outcome2.result.deinit(); defer outcome2.continuation.deinit(); const result2 = outcome2.result; @@ -668,7 +673,7 @@ test "Gemini.executeStep with tools" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const param = llm.types.Tool.Param{ @@ -690,7 +695,7 @@ test "Gemini.executeStep with tools" { .{ .prompt = "Hello" }, }; - var outcome = try p.executeStep(allocator, config, input, null); + const outcome = try p.executeStep(allocator, config, input, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); const result = outcome.result; @@ -718,7 +723,7 @@ test "Gemini.executeStep HTTP failure" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -796,7 +801,7 @@ test "Gemini.executeStepStreaming success" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -854,7 +859,7 @@ test "Gemini.executeStepStreaming success" { defer ctx.model_output_text.deinit(); defer ctx.thought_text.deinit(); - var outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); + const outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); defer outcome.result.deinit(); defer outcome.continuation.deinit(); const result = outcome.result; @@ -880,7 +885,7 @@ test "Gemini.executeStepStreaming success" { try std.testing.expectEqualStrings("location", result.tool_calls[0].arguments[0].name); try std.testing.expectEqualStrings("San Francisco", result.tool_calls[0].arguments[0].value.string); - const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + const gemini_continuation: *const gemini_types.StepContinuation = @fieldParentPtr("interface", continuation); try std.testing.expectEqualStrings("interaction_streaming_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } @@ -958,7 +963,7 @@ test "Gemini.executeStepStreaming with CRLF line endings" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -1016,7 +1021,7 @@ test "Gemini.executeStepStreaming with CRLF line endings" { defer ctx.model_output_text.deinit(); defer ctx.thought_text.deinit(); - var outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); + const outcome = try p.executeStepStreaming(allocator, config, input, null, Context.callback, &ctx); defer outcome.result.deinit(); defer outcome.continuation.deinit(); const result = outcome.result; @@ -1041,7 +1046,7 @@ test "Gemini.executeStepStreaming with CRLF line endings" { try std.testing.expectEqualStrings("location", result.tool_calls[0].arguments[0].name); try std.testing.expectEqualStrings("San Francisco", result.tool_calls[0].arguments[0].value.string); - const gemini_continuation: *gemini_types.StepContinuation = @ptrCast(@alignCast(continuation.ptr)); + const gemini_continuation: *const gemini_types.StepContinuation = @fieldParentPtr("interface", continuation); try std.testing.expectEqualStrings("interaction_streaming_123", gemini_continuation.interaction_id); try std.testing.expectEqual(1, call_counts[0]); } @@ -1082,20 +1087,19 @@ test "Gemini.executeStep returns function call" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, }; - var outcome = try p.executeStep(allocator, config, &.{}, null); + const outcome = try p.executeStep(allocator, config, &.{}, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); const result = outcome.result; - const continuation = outcome.continuation; - - try std.testing.expectEqualStrings("interaction_fc", (@as(*gemini_types.StepContinuation, @ptrCast(@alignCast(continuation.ptr)))).interaction_id); + const gemini_continuation: *const gemini_types.StepContinuation = @fieldParentPtr("interface", outcome.continuation); + try std.testing.expectEqualStrings("interaction_fc", gemini_continuation.interaction_id); try std.testing.expectEqual(1, result.tool_calls.len); try std.testing.expectEqualStrings("call_abc", result.tool_calls[0].id); try std.testing.expectEqualStrings("get_weather", result.tool_calls[0].name); @@ -1135,7 +1139,7 @@ test "Gemini.executeStepStreaming multiple interaction_created" { }, }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, @@ -1174,7 +1178,7 @@ test "Gemini.executeStepStreaming duplicate step start" { }, }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, @@ -1212,7 +1216,7 @@ test "Gemini.executeStepStreaming delta for non-existent step index" { }, }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, @@ -1251,7 +1255,7 @@ test "Gemini.executeStepStreaming mismatched interaction completed ID" { }, }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, @@ -1289,16 +1293,16 @@ test "Gemini.executeStepStreaming fallback interaction ID to unknown" { }, }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, }; - var outcome = try p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null); + const outcome = try p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); - const continuation = outcome.continuation; - try std.testing.expectEqualStrings("unknown", (@as(*gemini_types.StepContinuation, @ptrCast(@alignCast(continuation.ptr)))).interaction_id); + const gemini_continuation: *const gemini_types.StepContinuation = @fieldParentPtr("interface", outcome.continuation); + try std.testing.expectEqualStrings("unknown", gemini_continuation.interaction_id); } test "ListModelsResult.init OOM" { @@ -1357,7 +1361,7 @@ test "Gemini.executeStepStreaming malformed stream payload" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -1418,7 +1422,7 @@ test "Gemini.executeStep with system prompt" { }; var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); - var p = prov.provider(); + const p = prov.provider(); defer p.deinit(); const config = llm.types.SessionConfig{ @@ -1429,7 +1433,7 @@ test "Gemini.executeStep with system prompt" { .{ .prompt = "Hello" }, }; - var outcome = try p.executeStep(allocator, config, input, null); + const outcome = try p.executeStep(allocator, config, input, null); defer outcome.result.deinit(); defer outcome.continuation.deinit(); diff --git a/src/provider/google/types.zig b/src/provider/google/types.zig index 2f18bc6..75d9d5f 100644 --- a/src/provider/google/types.zig +++ b/src/provider/google/types.zig @@ -8,36 +8,38 @@ const Allocator = std.mem.Allocator; /// A Gemini-specific implementation of the `llm.types.ListModelsResult` interface. /// Represents the result of a request to list Gemini models. pub const ListModelsResult = struct { - models: []const llm.types.Model, - + interface: llm.types.ListModelsResult, allocator: Allocator, parsed_response: std.json.Parsed(api.ListModelsResponse), + const vtable = llm.types.ListModelsResult.VTable{ + .deinit = deinit, + }; + /// Initializes a new ListModelsResult by parsing the raw API response and returns a - /// reference to this struct via the `llm.types.ListModelsResult` interface. - pub fn init(allocator: Allocator, parsed_response: std.json.Parsed(api.ListModelsResponse)) !llm.types.ListModelsResult { + /// pointer to the `llm.types.ListModelsResult` interface. + pub fn init(allocator: Allocator, parsed_response: std.json.Parsed(api.ListModelsResponse)) !*llm.types.ListModelsResult { const self = try allocator.create(ListModelsResult); errdefer allocator.destroy(self); const models = try converter.toModels(allocator, parsed_response.value.models); errdefer allocator.free(models); self.* = .{ + .interface = .{ + .models = models, + .vtable = &vtable, + }, .allocator = allocator, .parsed_response = parsed_response, - .models = models, }; - return .{ - .ptr = self, - .vtable = &.{ .deinit = deinit }, - .models = models, - }; + return &self.interface; } /// Frees resources associated with the result. - pub fn deinit(ctx: *anyopaque) void { - const self: *ListModelsResult = @ptrCast(@alignCast(ctx)); + pub fn deinit(res: *llm.types.ListModelsResult) void { + const self: *ListModelsResult = @fieldParentPtr("interface", res); const allocator = self.allocator; - allocator.free(self.models); + allocator.free(self.interface.models); self.parsed_response.deinit(); self.* = undefined; allocator.destroy(self); @@ -47,6 +49,7 @@ pub const ListModelsResult = struct { /// A Gemini-specific implementation of the `llm.types.StepResult` interface. /// Represents the result of an interaction step with the Gemini API. pub const StepResult = struct { + interface: llm.types.StepResult, model_output: std.ArrayList(llm.types.ModelOutput), thoughts: std.ArrayList(llm.types.Thought), tool_calls: std.ArrayList(llm.types.ToolCall), @@ -54,9 +57,13 @@ pub const StepResult = struct { allocator: std.mem.Allocator, parsed_response: std.json.Parsed(api.Interaction), + const vtable = llm.types.StepResult.VTable{ + .deinit = deinit, + }; + /// Initializes a new StepResult by parsing the raw interaction response and returns a - /// reference to this struct via the `llm.types.StepResult` interface. - pub fn init(allocator: std.mem.Allocator, parsed_response: std.json.Parsed(api.Interaction)) !llm.types.StepResult { + /// pointer to the `llm.types.StepResult` interface. + pub fn init(allocator: std.mem.Allocator, parsed_response: std.json.Parsed(api.Interaction)) !*llm.types.StepResult { var model_output_list: std.ArrayList(llm.types.ModelOutput) = .empty; errdefer model_output_list.deinit(allocator); var thoughts_list: std.ArrayList(llm.types.Thought) = .empty; @@ -104,6 +111,12 @@ pub const StepResult = struct { const self = try allocator.create(StepResult); errdefer allocator.destroy(self); self.* = .{ + .interface = .{ + .model_output = model_output_list.items, + .thoughts = thoughts_list.items, + .tool_calls = tool_calls_list.items, + .vtable = &vtable, + }, .model_output = model_output_list, .thoughts = thoughts_list, .tool_calls = tool_calls_list, @@ -111,18 +124,12 @@ pub const StepResult = struct { .parsed_response = parsed_response, }; - return .{ - .ptr = self, - .vtable = &.{ .deinit = deinit }, - .model_output = model_output_list.items, - .thoughts = thoughts_list.items, - .tool_calls = tool_calls_list.items, - }; + return &self.interface; } /// Frees resources associated with the result. - pub fn deinit(ctx: *anyopaque) void { - const self: *StepResult = @ptrCast(@alignCast(ctx)); + pub fn deinit(res: *llm.types.StepResult) void { + const self: *StepResult = @fieldParentPtr("interface", res); const allocator = self.allocator; self.model_output.deinit(allocator); self.thoughts.deinit(allocator); @@ -138,43 +145,41 @@ pub const StepResult = struct { /// A Gemini-specific implementation of the `llm.types.StepResult` interface for streaming requests. pub const StreamingStepResult = struct { - model_output: []const llm.types.ModelOutput, - thoughts: []const llm.types.Thought, - tool_calls: []const llm.types.ToolCall, - + interface: llm.types.StepResult, arena: std.heap.ArenaAllocator, allocator: std.mem.Allocator, - /// Initializes a new StreamingStepResult and returns a reference to it via `llm.types.StepResult`. + const vtable = llm.types.StepResult.VTable{ + .deinit = deinit, + }; + + /// Initializes a new StreamingStepResult and returns a pointer to the `llm.types.StepResult` interface. pub fn init( allocator: std.mem.Allocator, arena: std.heap.ArenaAllocator, model_output: []const llm.types.ModelOutput, thoughts: []const llm.types.Thought, tool_calls: []const llm.types.ToolCall, - ) !llm.types.StepResult { + ) !*llm.types.StepResult { const self = try allocator.create(StreamingStepResult); errdefer allocator.destroy(self); self.* = .{ - .model_output = model_output, - .thoughts = thoughts, - .tool_calls = tool_calls, + .interface = .{ + .model_output = model_output, + .thoughts = thoughts, + .tool_calls = tool_calls, + .vtable = &vtable, + }, .arena = arena, .allocator = allocator, }; - return .{ - .ptr = self, - .vtable = &.{ .deinit = deinit }, - .model_output = model_output, - .thoughts = thoughts, - .tool_calls = tool_calls, - }; + return &self.interface; } /// Frees the resources associated with the StreamingStepResult. - pub fn deinit(ctx: *anyopaque) void { - const self: *StreamingStepResult = @ptrCast(@alignCast(ctx)); + pub fn deinit(res: *llm.types.StepResult) void { + const self: *StreamingStepResult = @fieldParentPtr("interface", res); const allocator = self.allocator; self.arena.deinit(); self.* = undefined; @@ -183,33 +188,37 @@ pub const StreamingStepResult = struct { }; pub const StepContinuation = struct { + interface: llm.types.StepContinuation, interaction_id: []const u8, - allocator: std.mem.Allocator, - /// Initializes a new StepContinuation and returns a reference to it via `llm.types.StepContinuation`. + const vtable = llm.types.StepContinuation.VTable{ + .deinit = deinit, + }; + + /// Initializes a new StepContinuation and returns a pointer to the `llm.types.StepContinuation` interface. /// /// This function creates a copy of the interaction_id and stores it in the new StepContinuation. /// The StepContinuation will deallocate the interaction_id when it is deinitialized. - pub fn init(allocator: std.mem.Allocator, interaction_id: []const u8) !llm.types.StepContinuation { + pub fn init(allocator: std.mem.Allocator, interaction_id: []const u8) !*llm.types.StepContinuation { const self = try allocator.create(StepContinuation); errdefer allocator.destroy(self); const dupe_interaction_id = try allocator.dupe(u8, interaction_id); errdefer allocator.free(dupe_interaction_id); self.* = .{ + .interface = .{ + .vtable = &vtable, + }, .interaction_id = dupe_interaction_id, .allocator = allocator, }; - return .{ - .ptr = self, - .vtable = &.{ .deinit = deinit }, - }; + return &self.interface; } /// Frees the resources associated with the StepContinuation. - pub fn deinit(ctx: *anyopaque) void { - const self: *StepContinuation = @ptrCast(@alignCast(ctx)); + pub fn deinit(cont: *llm.types.StepContinuation) void { + const self: *StepContinuation = @fieldParentPtr("interface", cont); const allocator = self.allocator; allocator.free(self.interaction_id); self.* = undefined; diff --git a/src/testing/BlockingTool.zig b/src/testing/BlockingTool.zig deleted file mode 100644 index b61f7a6..0000000 --- a/src/testing/BlockingTool.zig +++ /dev/null @@ -1,55 +0,0 @@ -//! A test tool implementation that blocks until `should_stop` is set to true -//! or the underlying I/O context is canceled. - -const std = @import("std"); -const Io = std.Io; -const llm = @import("llm"); - -const BlockingTool = @This(); - -/// Full descriptor for the blocking test tool. -pub const descriptor: llm.types.Tool = .{ - .name = "blocking_tool", - .description = "A test tool that blocks until signaled to stop", - .parameters = &.{}, -}; - -started: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), -should_stop: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), -io: Io, - -pub fn init(io: Io) BlockingTool { - return .{ .io = io }; -} - -/// Creates an agent Tool instance bound to this BlockingTool context. -pub fn toolWithContext(self: *BlockingTool, comptime Tool: type) Tool { - return Tool.initWithContext(descriptor, execute, self); -} - -/// Execution function compatible with agent.Tool.initWithContext. -pub fn execute(self: *BlockingTool, alloc: std.mem.Allocator) error{ ArgumentMismatch, ArgumentTypeMismatch, OutOfMemory, RequiredArgumentMissing }![]const u8 { - self.started.store(true, .monotonic); - while (!self.should_stop.load(.monotonic)) { - self.io.checkCancel() catch return error.OutOfMemory; - } - return alloc.dupe(u8, "blocking tool done") catch return error.OutOfMemory; -} - -/// Returns true if the tool execution function has been entered. -pub fn isStarted(self: *const BlockingTool) bool { - return self.started.load(.monotonic); -} - -/// Blocks until the tool execution function is entered. -pub fn waitUntilStarted(self: *BlockingTool) void { - var attempts: usize = 0; - while (!self.isStarted() and attempts < 1_000_000) : (attempts += 1) { - std.atomic.spinLoopHint(); - } -} - -/// Signals the tool execution to stop blocking and exit. -pub fn stop(self: *BlockingTool) void { - self.should_stop.store(true, .monotonic); -} diff --git a/src/testing/MockProvider.zig b/src/testing/MockProvider.zig index 0896b42..16a78fb 100644 --- a/src/testing/MockProvider.zig +++ b/src/testing/MockProvider.zig @@ -15,6 +15,8 @@ const StepOutcome = types.StepOutcome; const MockProvider = @This(); +interface: Provider = .{ .vtable = &vtable }, + /// Tracks the number of times `listModels` was called. list_models_calls: usize = 0, /// Tracks the number of times `executeStep` was called. @@ -31,10 +33,10 @@ last_session_config: ?SessionConfig = null, /// Stores the input steps from the last `executeStep` call. last_input_steps: ?[]const Step = null, /// Stores the previous step from the last `executeStep` call. -last_previous_step: ?StepContinuation = null, +last_previous_step: ?*const StepContinuation = null, /// Optional fixed result to return from `listModels`. -list_models_result: ?(Provider.ProviderError!ListModelsResult) = null, +list_models_result: ?(Provider.ProviderError!*ListModelsResult) = null, /// Optional sequence of results/outcomes to return from successive `executeStep` or `executeStepStreaming` calls. execute_step_results: ?[]const (Provider.ProviderError!StepOutcome) = null, /// Whether to loop back to the start of execute_step_results when more calls are made than outcomes available. @@ -50,46 +52,62 @@ const vtable = Provider.VTable{ }; /// Returns the generic `llm.Provider` interface for this mock instance. -pub fn provider(self: *MockProvider) Provider { - return .{ - .ptr = self, - .vtable = &vtable, - }; +pub fn provider(self: *MockProvider) *Provider { + return &self.interface; } /// A dummy virtual table for `ListModelsResult` used by the mock. pub const mock_list_models_vtable = ListModelsResult.VTable{ - .deinit = dummyDeinit, + .deinit = dummyDeinitListModels, }; /// A dummy virtual table for `StepResult` used by the mock. pub const mock_step_vtable = StepResult.VTable{ - .deinit = dummyDeinit, + .deinit = dummyDeinitStepResult, }; /// A dummy virtual table for `StepContinuation` used by the mock. pub const mock_continuation_vtable = StepContinuation.VTable{ - .deinit = dummyDeinit, + .deinit = dummyDeinitStepContinuation, +}; + +var default_step_result: StepResult = .{ + .model_output = &.{}, + .thoughts = &.{}, + .tool_calls = &.{}, + .vtable = &mock_step_vtable, +}; + +var default_step_continuation: StepContinuation = .{ + .vtable = &mock_continuation_vtable, +}; + +var default_list_models_result: ListModelsResult = .{ + .models = &.{}, + .vtable = &mock_list_models_vtable, }; -/// A no-op implementation of `deinit`. -fn dummyDeinit(ptr: *anyopaque) void { - _ = ptr; +fn dummyDeinitListModels(self: *ListModelsResult) void { + _ = self; +} + +fn dummyDeinitStepResult(self: *StepResult) void { + _ = self; +} + +fn dummyDeinitStepContinuation(self: *StepContinuation) void { + _ = self; } /// Mock implementation of `listModels`. -fn list_models(ptr: *anyopaque, allocator: Allocator) Provider.ProviderError!ListModelsResult { - const self: *MockProvider = @ptrCast(@alignCast(ptr)); +fn list_models(prov: *Provider, allocator: Allocator) Provider.ProviderError!*ListModelsResult { + const self: *MockProvider = @fieldParentPtr("interface", prov); self.list_models_calls += 1; self.last_allocator = allocator; if (self.list_models_result) |res| { return res; } - return ListModelsResult{ - .models = &.{}, - .ptr = ptr, - .vtable = &mock_list_models_vtable, - }; + return &default_list_models_result; } fn clearLastInputSteps(self: *MockProvider) void { @@ -129,13 +147,13 @@ fn recordInputSteps(self: *MockProvider, allocator: Allocator, input: []const St /// Mock implementation of `executeStep`. fn execute_step( - ptr: *anyopaque, + prov: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, ) Provider.ProviderError!StepOutcome { - const self: *MockProvider = @ptrCast(@alignCast(ptr)); + const self: *MockProvider = @fieldParentPtr("interface", prov); self.execute_step_calls += 1; self.last_allocator = allocator; self.last_session_config = session_config; @@ -155,28 +173,22 @@ fn execute_step( } } return StepOutcome{ - .result = StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = ptr, - .vtable = &mock_step_vtable, - }, - .continuation = StepContinuation{ .ptr = ptr, .vtable = &mock_continuation_vtable }, + .result = &default_step_result, + .continuation = &default_step_continuation, }; } /// Mock implementation of `executeStepStreaming`. fn execute_step_streaming( - ptr: *anyopaque, + prov: *Provider, allocator: Allocator, session_config: SessionConfig, input: []const Step, - previous_step: ?StepContinuation, + previous_step: ?*const StepContinuation, callback: types.StreamingCallback, callback_context: ?*anyopaque, ) Provider.ProviderError!StepOutcome { - const self: *MockProvider = @ptrCast(@alignCast(ptr)); + const self: *MockProvider = @fieldParentPtr("interface", prov); self.execute_step_streaming_calls += 1; self.last_allocator = allocator; self.last_session_config = session_config; @@ -209,21 +221,15 @@ fn execute_step_streaming( } } return StepOutcome{ - .result = StepResult{ - .model_output = &.{}, - .thoughts = &.{}, - .tool_calls = &.{}, - .ptr = ptr, - .vtable = &mock_step_vtable, - }, - .continuation = StepContinuation{ .ptr = ptr, .vtable = &mock_continuation_vtable }, + .result = &default_step_result, + .continuation = &default_step_continuation, }; } /// Mock implementation of `deinit`. -fn deinitVtable(ctx: *anyopaque) void { - const self: *MockProvider = @ptrCast(@alignCast(ctx)); - deinit(self); +fn deinitVtable(prov: *Provider) void { + const self: *MockProvider = @fieldParentPtr("interface", prov); + self.deinit(); } pub fn deinit(self: *MockProvider) void { @@ -241,24 +247,19 @@ pub fn stepResult( .model_output = model_output, .thoughts = thoughts, .tool_calls = tool_calls, - .ptr = @constCast(&mock_step_vtable), .vtable = &mock_step_vtable, }; } /// Helper constructor to create a mock `StepContinuation` with standard mock vtable. -pub fn stepContinuation() StepContinuation { - return .{ - .ptr = @constCast(&mock_continuation_vtable), - .vtable = &mock_continuation_vtable, - }; +pub fn stepContinuation() *StepContinuation { + return &default_step_continuation; } /// Helper constructor to create a mock `ListModelsResult` with standard mock vtable. pub fn listModelsResult(models: []const types.Model) ListModelsResult { return .{ .models = models, - .ptr = @constCast(&mock_list_models_vtable), .vtable = &mock_list_models_vtable, }; } diff --git a/src/testing/root.zig b/src/testing/root.zig index 5cffd34..3e1342f 100644 --- a/src/testing/root.zig +++ b/src/testing/root.zig @@ -1,5 +1,4 @@ pub const BlockingReader = @import("BlockingReader.zig"); -pub const BlockingTool = @import("BlockingTool.zig"); pub const MockHttpClient = @import("MockHttpClient.zig"); pub const MockProvider = @import("MockProvider.zig");