diff --git a/README.md b/README.md index 231dffb..e9ec47c 100644 --- a/README.md +++ b/README.md @@ -52,15 +52,15 @@ 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 agent_config: types.AgentConfig = .{ +const session_config: types.SessionConfig = .{ .model = selected_model, .tools = &[_]Tool{ weather_tool }, }; -var agent = try Agent.init(allocator, io, client, agent_config); -defer agent.deinit(); +var session = try Session.init(allocator, io, client, session_config); +defer session.deinit(); -const result = try agent.executeTurn(.{ .prompt = "What is the weather in 90210?" }); +const result = try session.executeTurn(.{ .prompt = "What is the weather in 90210?" }); ``` ## Project Structure diff --git a/build.zig b/build.zig index 6d187e8..c10c9f7 100644 --- a/build.zig +++ b/build.zig @@ -56,6 +56,17 @@ pub fn build(b: *std.Build) void { }, }); + const acp = b.addModule("acp", .{ + .root_source_file = b.path("src/acp/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "agent", .module = agent }, + .{ .name = "llm", .module = llm }, + .{ .name = "testing", .module = testing }, + }, + }); + // This creates a module, which represents a collection of source files alongside // some compilation options, such as optimization mode and linked system libraries. // Zig modules are the preferred way of making Zig code available to consumers. @@ -78,6 +89,7 @@ pub fn build(b: *std.Build) void { .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, .{ .name = "agent", .module = agent }, + .{ .name = "acp", .module = acp }, }, }); @@ -122,6 +134,7 @@ pub fn build(b: *std.Build) void { .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, .{ .name = "agent", .module = agent }, + .{ .name = "acp", .module = acp }, }, }), }); @@ -180,6 +193,7 @@ pub fn build(b: *std.Build) void { "kcov-out/suite_3", "kcov-out/suite_4", "kcov-out/suite_5", + "kcov-out/suite_6", }); for (coverage_test_suites, 0..) |suite, i| { @@ -222,7 +236,7 @@ fn createTestSuites( b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, -) [6]*std.Build.Step.Compile { +) [7]*std.Build.Step.Compile { const llm = b.createModule(.{ .root_source_file = b.path("src/llm/root.zig"), .target = target, @@ -259,6 +273,17 @@ fn createTestSuites( }, }); + const acp = b.createModule(.{ + .root_source_file = b.path("src/acp/root.zig"), + .target = target, + .optimize = optimize, + .imports = &.{ + .{ .name = "agent", .module = agent }, + .{ .name = "llm", .module = llm }, + .{ .name = "testing", .module = testing }, + }, + }); + const coma = b.createModule(.{ .root_source_file = b.path("src/root.zig"), .target = target, @@ -267,6 +292,7 @@ fn createTestSuites( .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, .{ .name = "agent", .module = agent }, + .{ .name = "acp", .module = acp }, }, }); @@ -279,6 +305,7 @@ fn createTestSuites( .{ .name = "llm", .module = llm }, .{ .name = "provider", .module = provider }, .{ .name = "agent", .module = agent }, + .{ .name = "acp", .module = acp }, }, }); @@ -299,5 +326,6 @@ fn createTestSuites( b.addTest(.{ .root_module = agent }), b.addTest(.{ .root_module = llm_test_module }), b.addTest(.{ .root_module = testing }), + b.addTest(.{ .root_module = acp }), }; } diff --git a/src/acp/Config.zig b/src/acp/Config.zig new file mode 100644 index 0000000..bba7b3b --- /dev/null +++ b/src/acp/Config.zig @@ -0,0 +1,10 @@ +//! Configuration of an ACP server. + +const std = @import("std"); +const Provider = @import("llm").Provider; +const SessionConfig = @import("agent").types.SessionConfig; + +const Config = @This(); + +provider: Provider, +default_session_config: SessionConfig, diff --git a/src/acp/Server.zig b/src/acp/Server.zig new file mode 100644 index 0000000..e5f9484 --- /dev/null +++ b/src/acp/Server.zig @@ -0,0 +1,315 @@ +//! Agent Communication Protocol (ACP) server implementation. +//! +//! Handles JSON-RPC 2.0 requests from client applications over standard input/output +//! or streaming I/O interfaces, dispatching initialization, session creation, and turn execution. + +const std = @import("std"); +const agent = @import("agent"); +const llm = @import("llm"); +const agent_api = @import("agent_api.zig"); +const client_api = @import("client_api.zig"); +const shared_api = @import("shared_api.zig"); +const converter = @import("converter.zig"); +const JsonRpcReader = @import("json_rpc/JsonRpcReader.zig"); +const JsonRpcWriter = @import("json_rpc/JsonRpcWriter.zig"); +const SessionStorage = @import("SessionStorage.zig"); + +/// ACP server configuration. +pub const Config = @import("Config.zig"); + +const Io = std.Io; +const Allocator = std.mem.Allocator; + +/// Protocol errors encountered when decoding or validating ACP JSON-RPC requests. +pub const AcpProtocolError = error{ + InvalidJsonRpcVersion, + MissingId, +} || std.json.Error; + +/// Internal context payload passed to session streaming callbacks. +const ServerSessionContext = struct { + session_state: *SessionStorage.SessionState, + json_rpc_writer: *JsonRpcWriter, + allocator: Allocator, +}; + +/// ACP JSON-RPC Server instance managing reader/writer loops and session state. +const Server = @This(); + +allocator: Allocator, +io: Io, +input_reader: *Io.Reader, +output_writer: *Io.Writer, +sessions: SessionStorage, + +/// Initializes a new ACP `Server` with the provided allocator, I/O context, reader, and writer. +pub fn init(allocator: Allocator, io: Io, input_reader: *Io.Reader, output_writer: *Io.Writer) Server { + return .{ + .allocator = allocator, + .io = io, + .input_reader = input_reader, + .output_writer = output_writer, + .sessions = .init(allocator), + }; +} + +/// Deinitializes the server and frees all tracked session resources. +pub fn deinit(self: *Server) void { + self.sessions.deinit(); +} + +/// 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 notification = converter.streamingChunkToNotification(stream_ctx.allocator, stream_ctx.session_state.id, chunk) catch return orelse return; + stream_ctx.json_rpc_writer.writeJsonObject(notification, .{ .use_headers = false }) catch {}; +} + +/// Formats and writes a JSON-RPC error response to the output writer. +fn sendError(self: *Server, allocator: Allocator, id: shared_api.RequestId, code: agent_api.JsonRpcErrorCode, message: []const u8) !void { + var writer = JsonRpcWriter.init(allocator, self.output_writer); + defer writer.deinit(); + try writer.writeJsonObject(agent_api.AgentErrorResponse{ + .id = id, + .@"error" = .{ + .code = code, + .message = message, + }, + }, .{}); +} + +/// Runs the main server request handling loop. +/// +/// Continuously reads JSON-RPC requests from `input_reader`, validates them, +/// and processes supported methods (`initialize`, `session/new`, `session/prompt`). +/// +/// This method blocks. To cancel, request cancelation via `Io.cancel`. +pub fn run(self: *Server, acp_config: Config) !void { + var arena = std.heap.ArenaAllocator.init(self.allocator); + defer arena.deinit(); + const arena_allocator = arena.allocator(); + + while (true) { + try self.io.checkCancel(); + _ = arena.reset(.retain_capacity); + + var json_rpc_reader = JsonRpcReader.init(arena_allocator, self.input_reader); + defer json_rpc_reader.deinit(); + + const parse_result = json_rpc_reader.readJsonObject(client_api.ClientRequest) catch |err| { + if (err == error.EndOfStream) return; + try self.sendError(arena_allocator, .null, .parse_error, "Parse error"); + continue; + }; + + const client_request = parse_result.value; + checkClientRequestValid(client_request) catch |err| { + const msg = switch (err) { + AcpProtocolError.InvalidJsonRpcVersion => "Invalid JSON-RPC version (must be 2.0)", + AcpProtocolError.MissingId => "Missing request ID", + else => "Invalid request", + }; + try self.sendError(arena_allocator, client_request.id, .invalid_request, msg); + continue; + }; + + switch (client_request.method) { + .initialize => { + const reply: agent_api.AgentResponse = .{ + .id = client_request.id, + .result = .{ + .initialize = .{ + .protocolVersion = 1, + .agentCapabilities = null, + .agentInfo = null, + .authMethods = {}, + }, + }, + }; + var json_rpc_writer = JsonRpcWriter.init(arena_allocator, self.output_writer); + defer json_rpc_writer.deinit(); + + try json_rpc_writer.writeJsonObject(reply, .{}); + }, + .session_new => { + const session_state = self.sessions.createSession(.{ + self.allocator, + self.io, + acp_config.provider, + acp_config.default_session_config, + }) catch { + try self.sendError(arena_allocator, client_request.id, .internal_error, "Failed to create session"); + continue; + }; + + const reply: agent_api.AgentResponse = .{ + .id = client_request.id, + .result = .{ + .session_new = .{ + .sessionId = session_state.id, + }, + }, + }; + var json_rpc_writer = JsonRpcWriter.init(arena_allocator, self.output_writer); + defer json_rpc_writer.deinit(); + + try json_rpc_writer.writeJsonObject(reply, .{}); + }, + .session_prompt => { + const prompt_blocks = client_request.params.session_prompt.prompt; + if (prompt_blocks.len == 0) { + try self.sendError(arena_allocator, client_request.id, .invalid_params, "Prompt array cannot be empty"); + continue; + } + + const session = self.sessions.getSession(client_request.params.session_prompt.sessionId) catch |err| { + const code: agent_api.JsonRpcErrorCode = if (err == error.SessionNotFound) .session_not_found else .internal_error; + const msg = if (err == error.SessionNotFound) "Session not found" else "Session retrieval error"; + try self.sendError(arena_allocator, client_request.id, code, msg); + continue; + }; + var json_rpc_writer = JsonRpcWriter.init(arena_allocator, self.output_writer); + defer json_rpc_writer.deinit(); + + var combined_prompt: std.ArrayList(u8) = .empty; + defer combined_prompt.deinit(arena_allocator); + for (prompt_blocks) |block| { + switch (block) { + .text => |txt| try combined_prompt.appendSlice(arena_allocator, txt), + } + } + + var ctx: ServerSessionContext = .{ + .session_state = session, + .json_rpc_writer = &json_rpc_writer, + .allocator = arena_allocator, + }; + + _ = session.session.executeTurnStreaming(.{ .prompt = combined_prompt.items }, handleTurnUpdate, &ctx) catch { + try self.sendError(arena_allocator, client_request.id, .internal_error, "Failed to execute turn"); + continue; + }; + + const reply: agent_api.AgentResponse = .{ + .id = client_request.id, + .result = .{ + .session_prompt = .{ + .stopReason = agent_api.StopReason.end_turn, + }, + }, + }; + + try json_rpc_writer.writeJsonObject(reply, .{}); + }, + .unknown => { + try self.sendError(arena_allocator, client_request.id, .method_not_found, "Method not found"); + }, + } + } +} + +/// Validates basic ACP JSON-RPC request structure (protocol version and request ID presence). +fn checkClientRequestValid(request: client_api.ClientRequest) AcpProtocolError!void { + if (!std.mem.eql(u8, request.jsonrpc, "2.0")) return AcpProtocolError.InvalidJsonRpcVersion; + if (request.id == .null) return AcpProtocolError.MissingId; +} + +test "Server error handling - malformed JSON and recovery" { + const allocator = std.testing.allocator; + + const input = + \\{ malformed json + \\{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":1}} + ; + var reader_buf = std.Io.Reader.fixed(input); + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var server = Server.init(allocator, std.testing.io, &reader_buf, &buffer.writer); + defer server.deinit(); + + try server.run(.{ .provider = undefined, .default_session_config = undefined }); + + const output = buffer.written(); + try std.testing.expect(std.mem.indexOf(u8, output, "-32700") != null); + try std.testing.expect(std.mem.indexOf(u8, output, "\"protocolVersion\":1") != null); +} + +test "Server error handling - invalid session ID" { + const allocator = std.testing.allocator; + + const input = + \\{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{"sessionId":"nonexistent","prompt":[{"type":"text","text":"hello"}]}} + ; + var reader_buf = std.Io.Reader.fixed(input); + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var server = Server.init(allocator, std.testing.io, &reader_buf, &buffer.writer); + defer server.deinit(); + + try server.run(.{ .provider = undefined, .default_session_config = undefined }); + + const output = buffer.written(); + try std.testing.expect(std.mem.indexOf(u8, output, "-32001") != null); + try std.testing.expect(std.mem.indexOf(u8, output, "Session not found") != null); +} + +test "Server error handling - empty prompt array" { + const allocator = std.testing.allocator; + + const input = + \\{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{"sessionId":"s1","prompt":[]}} + ; + var reader_buf = std.Io.Reader.fixed(input); + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var server = Server.init(allocator, std.testing.io, &reader_buf, &buffer.writer); + defer server.deinit(); + + try server.run(.{ .provider = undefined, .default_session_config = undefined }); + + const output = buffer.written(); + try std.testing.expect(std.mem.indexOf(u8, output, "-32602") != null); + try std.testing.expect(std.mem.indexOf(u8, output, "Prompt array cannot be empty") != null); +} + +test "Server prompt handling - multiple items in prompt array" { + const testing = @import("testing"); + const allocator = std.testing.allocator; + + var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); + const prov = mock_provider.provider(); + + const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); + const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ + .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + }; + mock_provider.execute_step_results = &outcomes; + + const input = + \\{"jsonrpc":"2.0","id":1,"method":"session/new","params":{"cwd":"/tmp","mcpServers":[]}} + \\{"jsonrpc":"2.0","id":2,"method":"session/prompt","params":{"sessionId":"session_0","prompt":[{"type":"text","text":"Hello "},{"type":"text","text":"world!"}]}} + ; + var reader_buf = std.Io.Reader.fixed(input); + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var server = Server.init(allocator, std.testing.io, &reader_buf, &buffer.writer); + defer server.deinit(); + + try server.run(.{ + .provider = prov, + .default_session_config = .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + }, + }); + + const output = buffer.written(); + try std.testing.expect(std.mem.indexOf(u8, output, "stopReason") != null); + try std.testing.expectEqual(@as(usize, 1), mock_provider.last_input_steps.?.len); + try std.testing.expectEqualStrings("Hello world!", mock_provider.last_input_steps.?[0].prompt); +} + diff --git a/src/acp/SessionStorage.zig b/src/acp/SessionStorage.zig new file mode 100644 index 0000000..389bab7 --- /dev/null +++ b/src/acp/SessionStorage.zig @@ -0,0 +1,119 @@ +//! Manages the storage and tracking of sessions in an ACP server. + +const std = @import("std"); +const agent = @import("agent"); + +const Allocator = std.mem.Allocator; +const Io = std.Io; + +/// ACP session storage container. +const SessionStorage = @This(); + +allocator: Allocator, +sessions: std.StringHashMapUnmanaged(*SessionState), +session_counter: u64, + +/// Holds state for an active ACP session, including its unique ID and underlying `agent.Session`. +pub const SessionState = struct { + id: []const u8, + session: agent.Session, +}; + +/// Initializes an empty `SessionStorage` instance. +pub fn init(allocator: Allocator) SessionStorage { + return .{ .allocator = allocator, .sessions = .{}, .session_counter = 0 }; +} + +/// Frees all stored session states, IDs, and internal map memory. +pub fn deinit(self: *SessionStorage) void { + var it = self.sessions.valueIterator(); + while (it.next()) |state_ptr| { + const state = state_ptr.*; + state.session.deinit(); + self.allocator.free(state.id); + self.allocator.destroy(state); + } + self.sessions.deinit(self.allocator); +} + +/// Tuple type representing the argument types required by `agent.Session.init`. +pub const SessionInitArgs = std.meta.ArgsTuple(@TypeOf(agent.Session.init)); + +/// Creates and stores a new `SessionState` with an auto-generated session ID. +pub fn createSession(self: *SessionStorage, args: SessionInitArgs) !*SessionState { + const session_id = try std.fmt.allocPrint(self.allocator, "session_{}", .{self.session_counter}); + errdefer self.allocator.free(session_id); + + const session_state = try self.allocator.create(SessionState); + errdefer self.allocator.destroy(session_state); + + var session = try @call(.auto, agent.Session.init, args); + errdefer session.deinit(); + session_state.* = .{ .id = session_id, .session = session }; + + try self.sessions.put(self.allocator, session_id, session_state); + + self.session_counter += 1; + return session_state; +} + +/// Retrieves a pointer to a `SessionState` by its session ID. +/// +/// Returns `error.SessionNotFound` if no session matches `id`. +pub fn getSession(self: *const SessionStorage, id: []const u8) !*SessionState { + return self.sessions.get(id) orelse return error.SessionNotFound; +} + +const testing = @import("testing"); + +test createSession { + const allocator = std.testing.allocator; + const io = std.testing.io; + + var mock_provider: testing.MockProvider = .{}; + + var session_storage = init(allocator); + defer session_storage.deinit(); + const session_state = try session_storage.createSession(.{ + allocator, io, mock_provider.provider(), + .{ + .model = .{ + .id = "mock-model", + .display_name = "Mock Model", + }, + .tools = &.{}, + }, + }); + const session = session_state.session; + try std.testing.expectEqualStrings("session_0", session_state.id); + try std.testing.expectEqual(mock_provider.provider(), session.provider); +} + +test getSession { + const allocator = std.testing.allocator; + const io = std.testing.io; + + var mock_provider: testing.MockProvider = .{}; + + var session_storage = init(allocator); + defer session_storage.deinit(); + + const session_state = try session_storage.createSession(.{ + allocator, io, mock_provider.provider(), + .{ + .model = .{ + .id = "mock-model", + .display_name = "Mock Model", + }, + .tools = &.{}, + }, + }); + const session = &session_state.session; + try std.testing.expectEqual(mock_provider.provider(), session.provider); + + const retrieved_session_state = try session_storage.getSession("session_0"); + try std.testing.expectEqual(&retrieved_session_state.session, session); + + const session_not_found = session_storage.getSession("session_1"); + try std.testing.expectError(error.SessionNotFound, session_not_found); +} diff --git a/src/acp/agent_api.zig b/src/acp/agent_api.zig new file mode 100644 index 0000000..5ce0a04 --- /dev/null +++ b/src/acp/agent_api.zig @@ -0,0 +1,644 @@ +//! Collection of API types sent by servers in ACP implementations. +//! +//! Contains primarily response objects. + +const std = @import("std"); +const shared_api = @import("shared_api.zig"); + +/// A JSON-RPC response sent by the agent to the client. +pub const AgentResponse = struct { + /// The ID of the request this response answers. + id: shared_api.RequestId, + /// Method-specific response data. + result: AgentResponseResult, + + /// Stringify the response as a JSON object. + /// + /// Prepends `"jsonrpc": "2.0"` to the response. + /// + /// See https://www.jsonrpc.org/specification#response_object + pub fn jsonStringify(self: AgentResponse, jw: anytype) !void { + try jw.beginObject(); + try jw.objectField("jsonrpc"); + try jw.write("2.0"); + try shared_api.jsonStringifyFields(self, jw); + try jw.endObject(); + } +}; + +/// JSON-RPC 2.0 error codes used in agent error responses. +/// +/// See https://www.jsonrpc.org/specification#error_object +pub const JsonRpcErrorCode = enum(i32) { + /// Invalid JSON was received by the server. An error occurred on the server while parsing the JSON text. + parse_error = -32700, + /// The JSON sent is not a valid Request object. + invalid_request = -32600, + /// The method does not exist / is not available. + method_not_found = -32601, + /// Invalid method parameter(s). + invalid_params = -32602, + /// Internal JSON-RPC error. + internal_error = -32603, + /// The requested session ID was not found. + session_not_found = -32001, + + pub fn jsonStringify(self: JsonRpcErrorCode, jw: anytype) !void { + try jw.write(@intFromEnum(self)); + } +}; + +/// Represents a JSON-RPC 2.0 error payload. +/// +/// See https://www.jsonrpc.org/specification#error_object +pub const JsonRpcError = struct { + /// Error code indicating the error type that occurred. + code: JsonRpcErrorCode, + /// A short description of the error. + message: []const u8, +}; + +/// A JSON-RPC error response sent by the agent to the client. +/// +/// See https://www.jsonrpc.org/specification#error_object +pub const AgentErrorResponse = struct { + /// The ID of the request this error response answers (or null if parsing failed). + id: shared_api.RequestId, + /// Error payload describing the failure. + @"error": JsonRpcError, + + /// Stringify the response as a JSON object. + /// + /// Prepends `"jsonrpc": "2.0"` to the response. + /// + /// See https://www.jsonrpc.org/specification#response_object + pub fn jsonStringify(self: AgentErrorResponse, jw: anytype) !void { + try jw.beginObject(); + try jw.objectField("jsonrpc"); + try jw.write("2.0"); + try shared_api.jsonStringifyFields(self, jw); + try jw.endObject(); + } +}; + +/// Union of all possible response data types. +/// +/// The filled member is dependent on the method type the agent is responding to. +pub const AgentResponseResult = union(enum) { + initialize: InitializeResponse, + session_new: NewSessionResponse, + session_prompt: PromptResponse, + + pub fn jsonStringify(self: AgentResponseResult, jw: anytype) !void { + switch (self) { + inline else => |payload| { + try jw.write(payload); + }, + } + } +}; + +/// Response from creating a new session. +/// +/// See protocol docs: [Creating a Session](https://agentclientprotocol.com/protocol/session-setup#creating-a-session) +pub const NewSessionResponse = struct { + /// Unique identifier for the created session. + /// + /// Used in all subsequent requests for this conversation. + sessionId: shared_api.SessionId, +}; + +/// Response to the `initialize` method. +/// +/// Contains the negotiated protocol version and agent capabilities. +/// +/// See protocol docs: [Initialization](https://agentclientprotocol.com/protocol/initialization) +pub const InitializeResponse = struct { + /// The protocol version the client specified if supported by the agent, + /// or the latest protocol version supported by the agent. + /// + /// The client should disconnect if it does not support this version. + protocolVersion: shared_api.ProtocolVersion, + /// Capabilities supported by the agent. + agentCapabilities: ?AgentCapabilities, + /// Authentication methods supported by the agent. + /// + /// This is currently not implemented. + authMethods: void, + /// Information about the Agent name and version sent to the client. + agentInfo: ?shared_api.Implementation, +}; + +/// Capabilities supported by the agent. +/// +/// These are advertised during initialization to inform the client about +/// available features and content types. +/// +/// See protocol docs: [Agent Capabilities](https://agentclientprotocol.com/protocol/initialization#agent-capabilities) +pub const AgentCapabilities = struct { + /// Whether the agent supports `session/load`. + loadSession: ?bool, + /// Prompt capabilities supported by the agent. + promptCapabilities: ?PromptCapabilities, + /// MCP capabilities supported by the agent. + mcpCapabilities: ?McpCapabilities, + /// Session lifecycle and prompt capabilities advertised by the agent. + sessionCapabilities: ?SessionCapabilities, + /// Authentication-related capabilities supported by the agent. + auth: ?AgentAuthCapabilities, +}; + +/// Prompt capabilities supported by the agent in `session/prompt` requests. +/// +/// Baseline agent functionality requires support for [`ContentBlock::Text`] +/// and [`ContentBlock::ResourceLink`] in prompt requests. +/// +/// Other variants must be explicitly opted in to. +/// Capabilities for different types of content in prompt requests. +/// +/// Indicates which content types beyond the baseline (text and resource links) +/// the agent can process. +/// +/// See protocol docs: [Prompt Capabilities](https://agentclientprotocol.com/protocol/initialization#prompt-capabilities) +pub const PromptCapabilities = struct { + /// Whether the agent supports `ContentBlock::Image` in prompt requests. + image: ?bool, + /// Whether the agent supports `ContentBlock::Audio` in prompt requests. + audio: ?bool, + /// Whether the agent supports embedded context in`session/prompt` requests. + /// + /// When enabled, the Client is allowed to include [`ContentBlock::Resource`] + /// in prompt requests for pieces of content that are referenced in the message. + embeddedContext: ?bool, +}; + +/// MCP capabilities supported by the agent. +pub const McpCapabilities = struct { + /// Whether the agent supports MCP HTTP servers. + http: ?bool, + /// Whether the agent supports MCP SSE servers. + sse: ?bool, +}; + +/// Session capabilities supported by the agent. +/// +/// As a baseline, all Agents **MUST** support `session/new`, `session/prompt`, `session/cancel`, and `session/update`. +/// +/// Optionally, they **MAY** support other session methods and notifications by specifying additional capabilities. +/// +/// Note: `session/load` is still handled by the top-level `load_session` capability. This will be unified in future versions of the protocol. +/// +/// See protocol docs: [Session Capabilities](https://agentclientprotocol.com/protocol/initialization#session-capabilities) +pub const SessionCapabilities = struct {}; + +/// Authentication-related capabilities supported by the agent. +pub const AgentAuthCapabilities = struct {}; + +pub const StopReason = enum { + end_turn, + max_tokens, + max_turn_requests, + refusal, + cancelled, +}; + +pub const PromptResponse = struct { + stopReason: StopReason, +}; + +/// Possible set of notification methods. +pub const AgentNotificationMethod = enum { + session_update, + + pub fn jsonStringify(self: AgentNotificationMethod, jw: anytype) !void { + try jw.write(shared_api.stringifyEnum(self)); + } +}; + +/// A JSON-RPC notification object. +pub const AgentNotification = struct { + /// The notification method name. + method: AgentNotificationMethod, + /// Method-specific notification parameters. + params: AgentNotificationParams, +}; + +/// Union for method-specific notification parameters; split by method type. +pub const AgentNotificationParams = union(AgentNotificationMethod) { + session_update: SessionNotificationParams, + + pub fn jsonStringify(self: AgentNotificationParams, jw: anytype) !void { + switch (self) { + inline else => |payload| { + try jw.write(payload); + }, + } + } +}; + +/// Notification containing a session update from the agent. +/// +/// Used to stream real-time progress and results during prompt processing. +/// +/// See protocol docs: [Agent Reports Output](https://agentclientprotocol.com/protocol/prompt-turn#3-agent-reports-output) +pub const SessionNotificationParams = struct { + /// The ID of the session that is receiving an update. + sessionId: shared_api.SessionId, + /// The update to send to the session + update: SessionUpdate, +}; + +/// Different types of updates that can be sent during session processing. +/// +/// These updates provide real-time feedback about the agent's progress. +/// +/// See protocol docs: [Agent Reports Output](https://agentclientprotocol.com/protocol/prompt-turn#3-agent-reports-output) +pub const SessionUpdate = union(enum) { + /// Agent message output chunk. + agent_message_chunk: ContentChunk, + /// Agent thinking chunk. + agent_thought_chunk: ContentChunk, + /// Tool call update chunk. + tool_call_update: ToolCallUpdate, + + pub fn jsonStringify(self: SessionUpdate, jw: anytype) !void { + switch (self) { + inline else => |payload| { + try jw.beginObject(); + try jw.objectField("sessionUpdate"); + try jw.write(@tagName(self)); + try shared_api.jsonStringifyFields(payload, jw); + try jw.endObject(); + }, + } + } +}; + +/// A streamed item of content. +pub const ContentChunk = struct { + content: shared_api.ContentBlock, + + // TODO(razza): Do we need messageId? +}; + +/// Execution status of a tool call. +/// +/// Tool calls progress through different statuses during their lifecycle. +/// +/// See protocol docs: [Status](https://agentclientprotocol.com/protocol/tool-calls#status) +pub const ToolCallStatus = enum { + pending, + in_progress, + completed, + failed, +}; + +/// An update to an existing tool call. +/// +/// Used to report progress and results as tools execute. All fields except +/// the tool call ID are optional - only changed fields need to be included. +/// +/// See protocol docs: [Updating](https://agentclientprotocol.com/protocol/tool-calls#updating) +pub const ToolCallUpdate = struct { + /// The ID of the tool call being updated. + toolCallId: shared_api.ToolCallId, + /// Human-readable title describing what the tool is doing. + title: ?[]const u8, + /// Programmatic name of the tool being invoked. + name: ?[]const u8, + /// Execution status of the tool call. + status: ?ToolCallStatus, + /// Update the raw input. + rawInput: ?[]const u8, + /// Update the raw output. + rawOutput: ?[]const u8, +}; + +/// Serializes a value into a newly allocated JSON string. +/// +/// Caller owns the returned slice and must free it using `allocator`. +/// Note: This helper is provided for unit tests to verify JSON output. +fn stringify(allocator: std.mem.Allocator, value: anytype) ![]u8 { + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var stringifier = std.json.Stringify{ + .writer = &buffer.writer, + .options = .{}, + }; + try stringifier.write(value); + return allocator.dupe(u8, buffer.written()); +} + +test "json stringify AgentResponse initialize" { + const allocator = std.testing.allocator; + + const response: AgentResponse = .{ + .id = .{ .integer = 1 }, + .result = .{ + .initialize = .{ + .protocolVersion = 1, + .agentCapabilities = .{ + .loadSession = true, + .promptCapabilities = .{ + .image = true, + .audio = false, + .embeddedContext = true, + }, + .mcpCapabilities = .{ + .http = true, + .sse = false, + }, + .sessionCapabilities = .{}, + .auth = null, + }, + .authMethods = {}, + .agentInfo = .{ + .name = "coma", + .title = "Coma ACP Agent", + .version = "0.1.0", + }, + }, + }, + }; + + const json_str = try stringify(allocator, response); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(@as(i64, 1), parsed.value.object.get("id").?.integer); + + const result = parsed.value.object.get("result").?.object; + try std.testing.expectEqual(@as(i64, 1), result.get("protocolVersion").?.integer); + + const agent_info = result.get("agentInfo").?.object; + try std.testing.expectEqualStrings("coma", agent_info.get("name").?.string); + try std.testing.expectEqualStrings("Coma ACP Agent", agent_info.get("title").?.string); + try std.testing.expectEqualStrings("0.1.0", agent_info.get("version").?.string); + + const caps = result.get("agentCapabilities").?.object; + try std.testing.expect(caps.get("loadSession").?.bool); + + const prompt_caps = caps.get("promptCapabilities").?.object; + try std.testing.expect(prompt_caps.get("image").?.bool); + try std.testing.expect(!prompt_caps.get("audio").?.bool); + try std.testing.expect(prompt_caps.get("embeddedContext").?.bool); + + const mcp_caps = caps.get("mcpCapabilities").?.object; + try std.testing.expect(mcp_caps.get("http").?.bool); + try std.testing.expect(!mcp_caps.get("sse").?.bool); +} + +test "json stringify AgentResponse session_new" { + const allocator = std.testing.allocator; + + const response: AgentResponse = .{ + .id = .{ .string = "req-123" }, + .result = .{ + .session_new = .{ + .sessionId = "sess-456", + }, + }, + }; + + const json_str = try stringify(allocator, response); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("req-123", parsed.value.object.get("id").?.string); + + const result = parsed.value.object.get("result").?.object; + try std.testing.expectEqualStrings("sess-456", result.get("sessionId").?.string); +} + +test "json stringify InitializeResponse optional fields" { + const allocator = std.testing.allocator; + + const response: AgentResponse = .{ + .id = .{ .integer = 42 }, + .result = .{ + .initialize = .{ + .protocolVersion = 1, + .agentCapabilities = null, + .authMethods = {}, + .agentInfo = null, + }, + }, + }; + + const json_str = try stringify(allocator, response); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(@as(i64, 42), parsed.value.object.get("id").?.integer); + + const result = parsed.value.object.get("result").?.object; + try std.testing.expectEqual(@as(i64, 1), result.get("protocolVersion").?.integer); + try std.testing.expect(result.get("agentCapabilities").? == .null); + try std.testing.expect(result.get("agentInfo").? == .null); +} + +test "json stringify AgentCapabilities minimal vs full" { + const allocator = std.testing.allocator; + + // Minimal capabilities (all optionals null) + { + const min_caps: AgentCapabilities = .{ + .loadSession = null, + .promptCapabilities = null, + .mcpCapabilities = null, + .sessionCapabilities = null, + .auth = null, + }; + + const json_str = try stringify(allocator, min_caps); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expect(parsed.value.object.get("loadSession").? == .null); + try std.testing.expect(parsed.value.object.get("promptCapabilities").? == .null); + try std.testing.expect(parsed.value.object.get("mcpCapabilities").? == .null); + } + + // Full capabilities + { + const full_caps: AgentCapabilities = .{ + .loadSession = true, + .promptCapabilities = .{ + .image = true, + .audio = null, + .embeddedContext = false, + }, + .mcpCapabilities = .{ + .http = false, + .sse = true, + }, + .sessionCapabilities = .{}, + .auth = .{}, + }; + + const json_str = try stringify(allocator, full_caps); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expect(parsed.value.object.get("loadSession").?.bool); + + const prompt_caps = parsed.value.object.get("promptCapabilities").?.object; + try std.testing.expect(prompt_caps.get("image").?.bool); + try std.testing.expect(prompt_caps.get("audio").? == .null); + try std.testing.expect(!prompt_caps.get("embeddedContext").?.bool); + + const mcp_caps = parsed.value.object.get("mcpCapabilities").?.object; + try std.testing.expect(!mcp_caps.get("http").?.bool); + try std.testing.expect(mcp_caps.get("sse").?.bool); + } +} + +test "json stringify AgentResponse session_prompt" { + const allocator = std.testing.allocator; + + const response: AgentResponse = .{ + .id = .{ .integer = 10 }, + .result = .{ + .session_prompt = .{ + .stopReason = .end_turn, + }, + }, + }; + + const json_str = try stringify(allocator, response); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(@as(i64, 10), parsed.value.object.get("id").?.integer); + + const result = parsed.value.object.get("result").?.object; + try std.testing.expectEqualStrings("end_turn", result.get("stopReason").?.string); +} + +test "json stringify AgentNotification session_update message chunk" { + const allocator = std.testing.allocator; + + const notification: AgentNotification = .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = "sess-789", + .update = .{ + .agent_message_chunk = .{ + .content = .{ + .text = "Hello world chunk", + }, + }, + }, + }, + }, + }; + + const json_str = try stringify(allocator, notification); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("session/update", parsed.value.object.get("method").?.string); + + const params = parsed.value.object.get("params").?.object; + try std.testing.expectEqualStrings("sess-789", params.get("sessionId").?.string); + + const update = params.get("update").?.object; + try std.testing.expectEqualStrings("agent_message_chunk", update.get("sessionUpdate").?.string); + + const content = update.get("content").?.object; + try std.testing.expectEqualStrings("text", content.get("type").?.string); + try std.testing.expectEqualStrings("Hello world chunk", content.get("text").?.string); +} + +test "json stringify AgentNotification session_update thought chunk" { + const allocator = std.testing.allocator; + + const notification: AgentNotification = .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = "sess-789", + .update = .{ + .agent_thought_chunk = .{ + .content = .{ + .text = "Thinking...", + }, + }, + }, + }, + }, + }; + + const json_str = try stringify(allocator, notification); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("session/update", parsed.value.object.get("method").?.string); + + const params = parsed.value.object.get("params").?.object; + const update = params.get("update").?.object; + try std.testing.expectEqualStrings("agent_thought_chunk", update.get("sessionUpdate").?.string); + + const content = update.get("content").?.object; + try std.testing.expectEqualStrings("Thinking...", content.get("text").?.string); +} + +test "json stringify AgentNotification session_update tool_call_update" { + const allocator = std.testing.allocator; + + const notification: AgentNotification = .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = "sess-123", + .update = .{ + .tool_call_update = .{ + .toolCallId = "tc-456", + .title = "Reading file", + .name = "read_file", + .status = .in_progress, + .rawInput = "{\"path\":\"/foo/bar\"}", + .rawOutput = null, + }, + }, + }, + }, + }; + + const json_str = try stringify(allocator, notification); + defer allocator.free(json_str); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("session/update", parsed.value.object.get("method").?.string); + + const params = parsed.value.object.get("params").?.object; + try std.testing.expectEqualStrings("sess-123", params.get("sessionId").?.string); + + const update = params.get("update").?.object; + try std.testing.expectEqualStrings("tool_call_update", update.get("sessionUpdate").?.string); + try std.testing.expectEqualStrings("tc-456", update.get("toolCallId").?.string); + try std.testing.expectEqualStrings("Reading file", update.get("title").?.string); + try std.testing.expectEqualStrings("read_file", update.get("name").?.string); + try std.testing.expectEqualStrings("in_progress", update.get("status").?.string); + try std.testing.expectEqualStrings("{\"path\":\"/foo/bar\"}", update.get("rawInput").?.string); +} diff --git a/src/acp/client_api.zig b/src/acp/client_api.zig new file mode 100644 index 0000000..fa59506 --- /dev/null +++ b/src/acp/client_api.zig @@ -0,0 +1,445 @@ +//! Collection of API types sent by clients in the Agent Client Protocol (ACP). +//! +//! Contains primarily requests and request parameters. + +const std = @import("std"); +const shared_api = @import("shared_api.zig"); + +const Allocator = std.mem.Allocator; + +/// Valid types of request methods that can be invoked by a client. +pub const RequestMethod = enum { + /// Initialize method used to negotiate protocol version and capabilities. + /// + /// Before a Session can be created or any other methods can be called, + /// Clients MUST initialize the connection by calling this method. + initialize, + /// New Session method (mapped to "session/new") used to create a new session. + session_new, + /// Session Prompt method (mapped to "session/prompt") used to send a prompt to the agent. + session_prompt, + /// The request method is not recognized. + unknown, + + pub fn jsonParse(allocator: Allocator, source: anytype, options: std.json.ParseOptions) !RequestMethod { + const json_value = try std.json.innerParse(std.json.Value, allocator, source, options); + return jsonParseFromValue(allocator, json_value, options); + } + + pub fn jsonParseFromValue(_: Allocator, source: std.json.Value, _: std.json.ParseOptions) !RequestMethod { + return shared_api.parseEnumWithMappedFallback(RequestMethod, .unknown, source); + } +}; + +/// A JSON-RPC request sent by the client to the agent. +/// +/// See protocol docs: [Requests](https://agentclientprotocol.com/protocol/overview) +pub const ClientRequest = struct { + /// The JSON-RPC protocol version. + jsonrpc: []const u8, + /// The request id used to correlate the matching response. + id: shared_api.RequestId, + /// The method to be invoked. + method: RequestMethod, + /// The parameters for the method invocation. + params: ClientRequestParams, + + pub fn jsonParse(allocator: Allocator, source: anytype, options: std.json.ParseOptions) !ClientRequest { + const json_value = try std.json.innerParse(std.json.Value, allocator, source, options); + return jsonParseFromValue(allocator, json_value, options); + } + + pub fn jsonParseFromValue(allocator: Allocator, source: std.json.Value, options: std.json.ParseOptions) !ClientRequest { + if (source != .object) return error.UnexpectedToken; + const jsonrpc = try std.json.innerParseFromValue([]const u8, allocator, source.object.get("jsonrpc") orelse return error.MissingField, options); + const id = try std.json.innerParseFromValue(shared_api.RequestId, allocator, source.object.get("id") orelse return error.MissingField, options); + const method = try std.json.innerParseFromValue(RequestMethod, allocator, source.object.get("method") orelse return error.MissingField, options); + const params: ClientRequestParams = switch (method) { + .initialize => .{ .initialize = try std.json.innerParseFromValue(InitializeRequest, allocator, source.object.get("params") orelse return error.MissingField, options) }, + .session_new => .{ .session_new = try std.json.innerParseFromValue(NewSessionRequest, allocator, source.object.get("params") orelse return error.MissingField, options) }, + .session_prompt => .{ .session_prompt = try std.json.innerParseFromValue(PromptRequest, allocator, source.object.get("params") orelse return error.MissingField, options) }, + .unknown => .{ .unknown = {} }, + }; + return ClientRequest{ + .jsonrpc = jsonrpc, + .id = id, + .method = method, + .params = params, + }; + } +}; + +/// Union of all possible request parameters. +/// +/// The filled member is dependent on the method type. +pub const ClientRequestParams = union(RequestMethod) { + initialize: InitializeRequest, + session_new: NewSessionRequest, + session_prompt: PromptRequest, + unknown, +}; + +/// Request parameters for the initialize method. +/// +/// Sent by the client to establish a connection and negotiate capabilities. +/// +/// See protocol docs: [Initialization](https://agentclientprotocol.com/protocol/initialization) +pub const InitializeRequest = struct { + /// The latest protocol version supported by the client. + protocolVersion: shared_api.ProtocolVersion, + /// Capabilities supported by the client. + clientCapabilities: ?ClientCapabilities = null, + /// Information about the Client name and version sent to the Agent. + /// Note: in future versions of the protocol, this will be required. + clientInfo: ?shared_api.Implementation = null, +}; + +/// Describes capabilities supported by the client. +/// +/// See protocol docs: [Client Capabilities](https://agentclientprotocol.com/protocol/initialization#client-capabilities) +pub const ClientCapabilities = struct { + /// File system capabilities supported by the client. + /// + /// Determines which file operations the agent can request from + /// the client. + fs: ?FileSystemCapabilities = null, + + /// Whether the client supports all `terminal*` methods. + terminal: bool = false, + + /// Session-related capabilities supported by the client. + /// + /// Optional. Omitted or `null` both mean the client does not advertise any + /// session-related extensions. + session: ?ClientSessionCapabilities = null, +}; + +/// File system capabilities that a client may support. +/// +/// See protocol docs: [FileSystem](https://agentclientprotocol.com/protocol/initialization#filesystem) +pub const FileSystemCapabilities = struct { + /// Whether the Client supports `fs/read_text_file` requests. + readTextFile: bool = false, + /// Whether the Client supports `fs/write_text_file` requests. + writeTextFile: bool = false, +}; + +/// Session-related capabilities supported by the client. +pub const ClientSessionCapabilities = struct { + /// Config option capabilities supported by the client. + /// + /// Omitted or `null` both mean the client does not advertise support for any + /// config option extensions. + configOptions: ?SessionConfigOptionsCapabilities = null, +}; + +/// Session configuration option capabilities supported by the client. +pub const SessionConfigOptionsCapabilities = struct { + /// Whether the client supports boolean session configuration options. + /// + /// Optional. Omitted or `null` both mean the client does not advertise support. + /// Supplying `{}` means agents may include `type: "boolean"` entries in + /// `configOptions`, and the client may send `session/set_config_option` + /// requests with `type: "boolean"` and a boolean `value`. + boolean: bool = false, +}; + +/// Request parameters for creating a new session. +/// +/// See protocol docs: [Session Setup](https://agentclientprotocol.com/protocol/session-setup#creating-a-session) +pub const NewSessionRequest = struct { + /// The working directory for this session. Must be an absolute path. + cwd: []const u8, + /// Additional workspace roots for this session. Each path must be absolute. + /// + /// 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, + /// List of MCP (Model Context Protocol) servers the agent should connect to. + mcpServers: []McpServer, +}; + +/// MCP server transport types. +pub const McpServerType = enum { + /// HTTP-based MCP server. + http, + /// Server-Sent Events (SSE) based MCP server. + sse, + /// ACP-based MCP server. + acp, + /// Stdio-based MCP server (default). + stdio, +}; + +/// Configuration for connecting to an MCP (Model Context Protocol) server. +/// +/// MCP servers provide tools and context that the agent can use when +/// processing prompts. +/// +/// The default MCP Server type is `stdio`. +/// +/// See protocol docs: [MCP Servers](https://agentclientprotocol.com/protocol/session-setup#mcp-servers) +pub const McpServer = union(McpServerType) { + /// HTTP-based MCP server. + http: void, + /// Server-Sent Events (SSE) based MCP server. + sse: void, + /// ACP-based MCP server. + acp: void, + /// Stdio-based MCP server (default). + stdio: McpServerStdio, + + pub fn jsonParse(allocator: Allocator, source: anytype, options: std.json.ParseOptions) !McpServer { + const json_value = try std.json.innerParse(std.json.Value, allocator, source, options); + return jsonParseFromValue(allocator, json_value, options); + } + + pub fn jsonParseFromValue(allocator: Allocator, source: std.json.Value, options: std.json.ParseOptions) !McpServer { + if (source != .object) return error.UnexpectedToken; + const server_type = if (source.object.get("type")) |type_val| + try std.json.innerParseFromValue(McpServerType, allocator, type_val, options) + else + McpServerType.stdio; + return switch (server_type) { + .http => .{ .http = {} }, + .sse => .{ .sse = {} }, + .acp => .{ .acp = {} }, + .stdio => .{ .stdio = try std.json.innerParseFromValue(McpServerStdio, allocator, source, options) }, + }; + } +}; + +/// Stdio transport configuration for an MCP server. +pub const McpServerStdio = struct { + /// Human-readable name identifying this MCP server. + name: []const u8, + /// Absolute path to the MCP server executable. + command: []const u8, + /// Command-line arguments to pass to the MCP server. + args: [][]const u8, + /// Environment variables to set when launching the MCP server. + env: []EnvVariable, +}; + +/// An environment variable to set when launching an MCP server. +pub const EnvVariable = struct { + /// The name of the environment variable. + name: []const u8, + /// The value of the environment variable. + value: []const u8, +}; + +/// Request parameters for sending a user prompt to the agent. +/// +/// Contains the user's message and any additional context. +/// +/// See protocol docs: [User Message](https://agentclientprotocol.com/protocol/prompt-turn#1-user-message) +pub const PromptRequest = struct { + /// The ID of the session to send this user message to + sessionId: shared_api.SessionId, + + /// The blocks of content that compose the user's message. + /// + /// As a baseline, the Agent MUST support [`ContentBlock::Text`] and [`ContentBlock::ResourceLink`], + /// while other variants are optionally enabled via [`PromptCapabilities`]. + /// + /// The Client MUST adapt its interface according to [`PromptCapabilities`]. + /// + /// The client MAY include referenced pieces of context as either + /// [`ContentBlock::Resource`] or [`ContentBlock::ResourceLink`]. + /// + /// 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, +}; + +test "json parse initialize ClientRequest" { + const allocator = std.testing.allocator; + const json_str = + \\{ + \\ "jsonrpc": "2.0", + \\ "id": 42, + \\ "method": "initialize", + \\ "params": { + \\ "protocolVersion": 1, + \\ "clientCapabilities": { + \\ "fs": { + \\ "readTextFile": true, + \\ "writeTextFile": false + \\ }, + \\ "terminal": true + \\ }, + \\ "clientInfo": { + \\ "name": "test-client", + \\ "version": "1.0.0" + \\ } + \\ } + \\} + ; + + const parsed = try std.json.parseFromSlice(ClientRequest, allocator, json_str, .{}); + defer parsed.deinit(); + + const request = parsed.value; + try std.testing.expectEqualStrings("2.0", request.jsonrpc); + try std.testing.expectEqual(RequestMethod.initialize, request.method); + try std.testing.expectEqual(shared_api.RequestId{ .integer = 42 }, request.id); + + const init_params = request.params.initialize; + try std.testing.expectEqual(@as(shared_api.ProtocolVersion, 1), init_params.protocolVersion); + + const capabilities = init_params.clientCapabilities.?; + try std.testing.expect(capabilities.terminal); + const fs = capabilities.fs.?; + try std.testing.expect(fs.readTextFile); + try std.testing.expect(!fs.writeTextFile); + + const client_info = init_params.clientInfo.?; + try std.testing.expectEqualStrings("test-client", client_info.name); + try std.testing.expectEqualStrings("1.0.0", client_info.version); + try std.testing.expect(client_info.title == null); +} + +test "json parse RequestMethod mapping" { + const allocator = std.testing.allocator; + + { + const parsed = try std.json.parseFromSlice(RequestMethod, allocator, "\"initialize\"", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestMethod.initialize, parsed.value); + } + + { + const parsed = try std.json.parseFromSlice(RequestMethod, allocator, "\"session/new\"", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestMethod.session_new, parsed.value); + } + + { + const parsed = try std.json.parseFromSlice(RequestMethod, allocator, "\"session/prompt\"", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestMethod.session_prompt, parsed.value); + } + + { + const parsed = try std.json.parseFromSlice(RequestMethod, allocator, "\"some/unknown/method\"", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestMethod.unknown, parsed.value); + } +} + +test "json parse McpServer stdio" { + const allocator = std.testing.allocator; + const json_str = + \\{ + \\ "name": "filesystem", + \\ "command": "/path/to/mcp-server", + \\ "args": ["--stdio"], + \\ "env": [] + \\} + ; + + const parsed = try std.json.parseFromSlice(McpServer, allocator, json_str, .{}); + defer parsed.deinit(); + + const server = parsed.value; + try std.testing.expectEqual(McpServerType.stdio, @as(McpServerType, server)); + try std.testing.expectEqualStrings("filesystem", server.stdio.name); + try std.testing.expectEqualStrings("/path/to/mcp-server", server.stdio.command); + try std.testing.expectEqual(1, server.stdio.args.len); + try std.testing.expectEqualStrings("--stdio", server.stdio.args[0]); + try std.testing.expectEqual(0, server.stdio.env.len); +} + +test "json parse McpServer http" { + const allocator = std.testing.allocator; + const json_str = + \\{ + \\ "type": "http" + \\} + ; + + const parsed = try std.json.parseFromSlice(McpServer, allocator, json_str, .{}); + defer parsed.deinit(); + + const server = parsed.value; + try std.testing.expectEqual(McpServerType.http, @as(McpServerType, server)); +} + +test "json parse session/new ClientRequest with McpServer" { + const allocator = std.testing.allocator; + const json_str = + \\{ + \\ "jsonrpc": "2.0", + \\ "id": 1, + \\ "method": "session/new", + \\ "params": { + \\ "cwd": "/home/user/project", + \\ "mcpServers": [ + \\ { + \\ "name": "filesystem", + \\ "command": "/path/to/mcp-server", + \\ "args": ["--stdio"], + \\ "env": [] + \\ } + \\ ] + \\ } + \\} + ; + + const parsed = try std.json.parseFromSlice(ClientRequest, allocator, json_str, .{}); + defer parsed.deinit(); + + const request = parsed.value; + try std.testing.expectEqualStrings("2.0", request.jsonrpc); + try std.testing.expectEqual(RequestMethod.session_new, request.method); + try std.testing.expectEqual(shared_api.RequestId{ .integer = 1 }, request.id); + + const session_params = request.params.session_new; + try std.testing.expectEqualStrings("/home/user/project", session_params.cwd); + try std.testing.expectEqual(1, session_params.mcpServers.len); + + const mcp_server = session_params.mcpServers[0]; + try std.testing.expectEqual(McpServerType.stdio, @as(McpServerType, mcp_server)); + try std.testing.expectEqualStrings("filesystem", mcp_server.stdio.name); + try std.testing.expectEqualStrings("/path/to/mcp-server", mcp_server.stdio.command); + try std.testing.expectEqual(1, mcp_server.stdio.args.len); + try std.testing.expectEqualStrings("--stdio", mcp_server.stdio.args[0]); + try std.testing.expectEqual(0, mcp_server.stdio.env.len); +} + +test "json parse session/prompt ClientRequest with PromptRequest" { + const allocator = std.testing.allocator; + const json_str = + \\{ + \\ "jsonrpc": "2.0", + \\ "id": 3, + \\ "method": "session/prompt", + \\ "params": { + \\ "sessionId": "session-123", + \\ "prompt": [ + \\ { + \\ "type": "text", + \\ "text": "Hello ACP agent!" + \\ } + \\ ] + \\ } + \\} + ; + + const parsed = try std.json.parseFromSlice(ClientRequest, allocator, json_str, .{}); + defer parsed.deinit(); + + const request = parsed.value; + try std.testing.expectEqualStrings("2.0", request.jsonrpc); + try std.testing.expectEqual(RequestMethod.session_prompt, request.method); + try std.testing.expectEqual(shared_api.RequestId{ .integer = 3 }, request.id); + + const prompt_params = request.params.session_prompt; + try std.testing.expectEqualStrings("session-123", prompt_params.sessionId); + try std.testing.expectEqual(1, prompt_params.prompt.len); + try std.testing.expectEqual(shared_api.ContentType.text, @as(shared_api.ContentType, prompt_params.prompt[0])); + try std.testing.expectEqualStrings("Hello ACP agent!", prompt_params.prompt[0].text); +} diff --git a/src/acp/converter.zig b/src/acp/converter.zig new file mode 100644 index 0000000..c4d4d68 --- /dev/null +++ b/src/acp/converter.zig @@ -0,0 +1,373 @@ +//! Methods to convert between internal COMA data types and ACP API types. + +const std = @import("std"); +const agent = @import("agent"); +const llm = @import("llm"); +const agent_api = @import("agent_api.zig"); +const shared_api = @import("shared_api.zig"); + +const Allocator = std.mem.Allocator; + +/// Extracts the `Delta` payload from a `StreamingChunk` if it is a step delta event. +/// +/// - `chunk`: The streaming chunk from the agent layer. +pub fn extractDelta(chunk: agent.types.StreamingChunk) ?llm.types.Delta { + if (chunk != .model_chunk) return null; + if (chunk.model_chunk.event != .step_event) return null; + if (chunk.model_chunk.event.step_event.event != .delta) return null; + return chunk.model_chunk.event.step_event.event.delta; +} + +/// Converts a `ModelOutput` into a `ContentChunk`. +/// +/// - `model_output`: The model output payload containing generated text. +pub fn modelOutputToContentChunk(model_output: llm.types.ModelOutput) agent_api.ContentChunk { + return .{ + .content = .{ .text = model_output.text }, + }; +} + +/// Converts a `Thought` into a `ContentChunk`. +/// +/// - `thought`: The thought payload containing reasoning text. +pub fn thoughtToContentChunk(thought: llm.types.Thought) agent_api.ContentChunk { + return .{ + .content = .{ .text = thought.text }, + }; +} + +/// Constructs an `AgentNotification` for an agent message chunk update. +/// +/// - `session_id`: The ID of the active session receiving the update. +/// - `model_output`: The model output payload to include in the message chunk notification. +fn agentMessageChunk(session_id: shared_api.SessionId, model_output: llm.types.ModelOutput) agent_api.AgentNotification { + return .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = session_id, + .update = .{ + .agent_message_chunk = modelOutputToContentChunk(model_output), + }, + }, + }, + }; +} + +/// Constructs an `AgentNotification` for an agent thought chunk update. +/// +/// - `session_id`: The ID of the active session receiving the update. +/// - `thought`: The thought payload to include in the thought chunk notification. +fn agentThoughtChunk(session_id: shared_api.SessionId, thought: llm.types.Thought) agent_api.AgentNotification { + return .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = session_id, + .update = .{ + .agent_thought_chunk = thoughtToContentChunk(thought), + }, + }, + }, + }; +} + +/// Constructs an `AgentNotification` for a tool call update chunk. +/// +/// - `allocator`: Allocator used to JSON-serialize the tool call arguments into `rawInput`. +/// - `session_id`: The ID of the active session receiving the update. +/// - `tool_call_delta`: The tool call delta payload containing tool ID, name, and arguments. +fn agentToolCallUpdate(allocator: Allocator, session_id: shared_api.SessionId, tool_call_delta: llm.types.ToolCallDelta) !agent_api.AgentNotification { + var write_buffer: std.Io.Writer.Allocating = .init(allocator); + defer write_buffer.deinit(); + var stringifier = std.json.Stringify{ + .writer = &write_buffer.writer, + .options = .{}, + }; + try stringifier.write(tool_call_delta.arguments); + return .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = session_id, + .update = .{ + .tool_call_update = .{ + .toolCallId = tool_call_delta.id, + .title = tool_call_delta.name, + .name = tool_call_delta.name, + .status = .in_progress, + .rawInput = try write_buffer.toOwnedSlice(), + .rawOutput = null, + }, + }, + }, + }, + }; +} + +/// Constructs an `AgentNotification` for a tool execution result update. +/// +/// - `session_id`: The ID of the active session receiving the update. +/// - `tool_result`: The tool result payload containing tool ID, name, and result string. +fn agentToolResult(session_id: shared_api.SessionId, tool_result: llm.types.ToolResult) agent_api.AgentNotification { + return .{ + .method = .session_update, + .params = .{ + .session_update = .{ + .sessionId = session_id, + .update = .{ + .tool_call_update = .{ + .toolCallId = tool_result.id, + .title = tool_result.tool_name, + .name = tool_result.tool_name, + .status = .completed, + .rawInput = null, + .rawOutput = tool_result.result, + }, + }, + }, + }, + }; +} + +/// Converts a `StreamingChunk` into an optional `AgentNotification`. +/// +/// - `allocator`: Allocator to use when producing this chunk. +/// - `session_id`: The ID of the active session receiving the update. +/// - `chunk`: The streaming chunk from the agent layer to convert. +pub fn streamingChunkToNotification(allocator: Allocator, session_id: shared_api.SessionId, chunk: agent.types.StreamingChunk) !?agent_api.AgentNotification { + switch (chunk) { + .model_chunk => { + const delta = extractDelta(chunk) orelse return null; + return switch (delta) { + .model_output => |model_output| agentMessageChunk(session_id, model_output), + .thought => |thought| agentThoughtChunk(session_id, thought), + .tool_call => |tool_call| try agentToolCallUpdate(allocator, session_id, tool_call), + }; + }, + .tool_result => |tool_result| { + return agentToolResult(session_id, tool_result); + }, + } +} + +test "extractDelta extracts model output, thought, and tool call deltas" { + const chunk_output = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .model_output = .{ .text = "hello" }, + }, + }, + }, + }, + }, + }; + const delta_output = extractDelta(chunk_output); + try std.testing.expect(delta_output != null); + try std.testing.expectEqualStrings("hello", delta_output.?.model_output.text); + + const chunk_thought = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .thought = .{ .text = "thinking" }, + }, + }, + }, + }, + }, + }; + const delta_thought = extractDelta(chunk_thought); + try std.testing.expect(delta_thought != null); + try std.testing.expectEqualStrings("thinking", delta_thought.?.thought.text); + + const chunk_tool_call = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .tool_call = .{ + .id = "call-1", + .name = "read_file", + .arguments = &.{}, + }, + }, + }, + }, + }, + }, + }; + const delta_tool_call = extractDelta(chunk_tool_call); + try std.testing.expect(delta_tool_call != null); + try std.testing.expectEqualStrings("call-1", delta_tool_call.?.tool_call.id); + try std.testing.expectEqualStrings("read_file", delta_tool_call.?.tool_call.name); + + const chunk_other = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .interaction_created, + }, + }; + try std.testing.expectEqual(@as(?llm.types.Delta, null), extractDelta(chunk_other)); +} + +test "modelOutputToContentChunk and thoughtToContentChunk" { + const content_output = modelOutputToContentChunk(.{ .text = "test output" }); + try std.testing.expectEqualStrings("test output", content_output.content.text); + + const content_thought = thoughtToContentChunk(.{ .text = "test thought" }); + try std.testing.expectEqualStrings("test thought", content_thought.content.text); +} + +test "agentMessageChunk and agentThoughtChunk build expected notifications" { + const session_id: shared_api.SessionId = "session-42"; + + const msg_notif = agentMessageChunk(session_id, .{ .text = "message delta" }); + try std.testing.expectEqual(agent_api.AgentNotificationMethod.session_update, msg_notif.method); + try std.testing.expectEqualStrings(session_id, msg_notif.params.session_update.sessionId); + try std.testing.expectEqualStrings("message delta", msg_notif.params.session_update.update.agent_message_chunk.content.text); + + const thought_notif = agentThoughtChunk(session_id, .{ .text = "thought delta" }); + try std.testing.expectEqual(agent_api.AgentNotificationMethod.session_update, thought_notif.method); + try std.testing.expectEqualStrings(session_id, thought_notif.params.session_update.sessionId); + try std.testing.expectEqualStrings("thought delta", thought_notif.params.session_update.update.agent_thought_chunk.content.text); +} + +test "agentToolCallUpdate and agentToolResult build expected notifications" { + const session_id: shared_api.SessionId = "session-42"; + const allocator = std.testing.allocator; + + const tool_call_notif = try agentToolCallUpdate(allocator, session_id, .{ + .id = "tc-1", + .name = "list_dir", + .arguments = &.{}, + }); + defer allocator.free(tool_call_notif.params.session_update.update.tool_call_update.rawInput.?); + + try std.testing.expectEqual(agent_api.AgentNotificationMethod.session_update, tool_call_notif.method); + try std.testing.expectEqualStrings(session_id, tool_call_notif.params.session_update.sessionId); + + const tc_update = tool_call_notif.params.session_update.update.tool_call_update; + try std.testing.expectEqualStrings("tc-1", tc_update.toolCallId); + try std.testing.expectEqualStrings("list_dir", tc_update.name.?); + try std.testing.expectEqual(agent_api.ToolCallStatus.in_progress, tc_update.status.?); + try std.testing.expectEqualStrings("[]", tc_update.rawInput.?); + try std.testing.expect(tc_update.rawOutput == null); + + const tool_result_notif = agentToolResult(session_id, .{ + .id = "tc-1", + .tool_name = "list_dir", + .result = "file1.txt\nfile2.txt", + .allocator = allocator, + }); + const tr_update = tool_result_notif.params.session_update.update.tool_call_update; + try std.testing.expectEqualStrings("tc-1", tr_update.toolCallId); + try std.testing.expectEqualStrings("list_dir", tr_update.name.?); + try std.testing.expectEqual(agent_api.ToolCallStatus.completed, tr_update.status.?); + try std.testing.expect(tr_update.rawInput == null); + try std.testing.expectEqualStrings("file1.txt\nfile2.txt", tr_update.rawOutput.?); +} + +test "streamingChunkToNotification converts streaming chunks correctly" { + const session_id: shared_api.SessionId = "session-100"; + + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + + // 1. Model output chunk + const chunk_msg = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .model_output = .{ .text = "hello notification" }, + }, + }, + }, + }, + }, + }; + + const notif_msg = try streamingChunkToNotification(arena.allocator(), session_id, chunk_msg); + try std.testing.expect(notif_msg != null); + try std.testing.expectEqualStrings("hello notification", notif_msg.?.params.session_update.update.agent_message_chunk.content.text); + + // 2. Thought chunk + const chunk_thought = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .thought = .{ .text = "deep thought" }, + }, + }, + }, + }, + }, + }; + + const notif_thought = try streamingChunkToNotification(arena.allocator(), session_id, chunk_thought); + try std.testing.expect(notif_thought != null); + try std.testing.expectEqualStrings("deep thought", notif_thought.?.params.session_update.update.agent_thought_chunk.content.text); + + // 3. Tool call chunk + const chunk_tool_call = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .{ + .step_event = .{ + .index = 0, + .event = .{ + .delta = .{ + .tool_call = .{ + .id = "tc-99", + .name = "grep_search", + .arguments = &.{}, + }, + }, + }, + }, + }, + }, + }; + + const notif_tc = try streamingChunkToNotification(arena.allocator(), session_id, chunk_tool_call); + try std.testing.expect(notif_tc != null); + try std.testing.expectEqualStrings("tc-99", notif_tc.?.params.session_update.update.tool_call_update.toolCallId); + try std.testing.expectEqual(agent_api.ToolCallStatus.in_progress, notif_tc.?.params.session_update.update.tool_call_update.status.?); + + // 4. Tool result chunk + const chunk_tool_result = agent.types.StreamingChunk{ + .tool_result = .{ + .id = "tc-99", + .tool_name = "grep_search", + .result = "match found", + .allocator = arena.allocator(), + }, + }; + + const notif_tr = try streamingChunkToNotification(arena.allocator(), session_id, chunk_tool_result); + try std.testing.expect(notif_tr != null); + try std.testing.expectEqualStrings("tc-99", notif_tr.?.params.session_update.update.tool_call_update.toolCallId); + try std.testing.expectEqual(agent_api.ToolCallStatus.completed, notif_tr.?.params.session_update.update.tool_call_update.status.?); + try std.testing.expectEqualStrings("match found", notif_tr.?.params.session_update.update.tool_call_update.rawOutput.?); + + // 5. Non-delta event returns null + const chunk_other = agent.types.StreamingChunk{ + .model_chunk = .{ + .event = .interaction_created, + }, + }; + try std.testing.expectEqual(null, try streamingChunkToNotification(arena.allocator(), session_id, chunk_other)); +} diff --git a/src/acp/json_rpc/JsonRpcReader.zig b/src/acp/json_rpc/JsonRpcReader.zig new file mode 100644 index 0000000..3bdc99e --- /dev/null +++ b/src/acp/json_rpc/JsonRpcReader.zig @@ -0,0 +1,273 @@ +//! Takes a `*std.Io.Reader` that contains a [JSON-RPC](https://www.jsonrpc.org/specification) stream. +//! +//! Reads individual JSON-RPC messages from the stream and returns them as byte slices. +//! Supports both streams prefixed with `Content-Length` headers and streams without headers. +//! +//! This is NOT threadsafe. + +const std = @import("std"); + +const Allocator = std.mem.Allocator; +const Io = std.Io; + +const JsonRpcReader = @This(); + +allocator: Allocator, +reader: *Io.Reader, +read_buffer: std.Io.Writer.Allocating, + +/// Initializes a new `JsonRpcReader`. +/// +/// The returned `JsonRpcReader` must be deinitialized when no longer needed by calling `deinit()`. +pub fn init(allocator: Allocator, reader: *Io.Reader) JsonRpcReader { + return .{ + .allocator = allocator, + .reader = reader, + .read_buffer = .init(allocator), + }; +} + +/// Frees the resources associated with the `JsonRpcReader`. +pub fn deinit(self: *JsonRpcReader) void { + self.read_buffer.deinit(); +} + +/// Reads the next raw JSON-RPC message from the stream and returns it as a raw byte slice. +/// +/// The returned slice must be freed by the caller. +pub fn readRawMessage(self: *JsonRpcReader) ![]const u8 { + return self.allocator.dupe(u8, try self.readRawMessageInternal()); +} + +/// Reads the next JSON-RPC message from the stream and parses it as an ObjectType. +/// +/// The returned `Parsed(ObjectType)` must be deinitialized when no longer needed by called `deinit()`. +pub fn readJsonObject(self: *JsonRpcReader, ObjectType: type) !std.json.Parsed(ObjectType) { + const raw = try self.readRawMessageInternal(); + return try std.json.parseFromSlice(ObjectType, self.allocator, raw, .{ .ignore_unknown_fields = true, .allocate = .alloc_always }); +} + +fn streamLine(self: *JsonRpcReader) ![]const u8 { + const line_length = self.reader.streamDelimiter(&self.read_buffer.writer, '\n'); + defer self.read_buffer.clearRetainingCapacity(); + if (line_length) |_| { + self.reader.toss(1); + } else |err| switch (err) { + error.EndOfStream => if (self.read_buffer.written().len == 0) return err, // Expected if there are no more messages in the stream. + else => return err, + } + var line = self.read_buffer.written(); + if (line.len > 0 and line[line.len - 1] == '\r') { + line = line[0 .. line.len - 1]; + } + return line; +} + +fn readRawMessageInternal(self: *JsonRpcReader) ![]const u8 { + const line = try self.streamLine(); + + if (line.len == 0) return error.UnexpectedEndOfInput; + + // TODO: We assume that the first header (in LSP-style JSON-RPC) is "Content-Length". + // This may not always be true. "Content-Type" is another common header and the + // spec does not specify ordering. + if (std.ascii.startsWithIgnoreCase(line, "Content-Length:")) { + const parts = std.mem.trim(u8, line["Content-Length:".len..], " "); + const content_length = try std.fmt.parseInt(usize, parts, 10); + + // The protocol specifies that all headers are followed by a blank line. + // Therefore, we consume all content until a blank line. + while (true) { + const expected_empty = try self.streamLine(); + if (expected_empty.len == 0) { + break; + } + } + + try self.reader.streamExact(&self.read_buffer.writer, content_length); + defer self.read_buffer.clearRetainingCapacity(); + return self.read_buffer.written(); + } else { + return line; + } +} + +test readJsonObject { + const TestObjectParams = struct { + version: []const u8, + }; + + const TestObject = struct { + jsonrpc: []const u8, + method: []const u8, + params: TestObjectParams, + id: u64, + }; + + // test reading newline delimited messages + { + const input = + \\{"jsonrpc":"2.0","method":"initialize","params":{"version":"2.0"},"id":1} + \\{"jsonrpc":"2.0","method":"initialize","params":{"version":"2.0"},"id":2} + ; + + const allocator = std.testing.allocator; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const json_msg1 = try reader.readJsonObject(TestObject); + defer json_msg1.deinit(); + const msg1 = json_msg1.value; + + const json_msg2 = try reader.readJsonObject(TestObject); + defer json_msg2.deinit(); + const msg2 = json_msg2.value; + + try std.testing.expectEqualStrings("2.0", msg1.jsonrpc); + try std.testing.expectEqualStrings("initialize", msg1.method); + try std.testing.expectEqual(1, msg1.id); + try std.testing.expectEqualStrings("2.0", msg1.params.version); + + try std.testing.expectEqualStrings("2.0", msg2.jsonrpc); + try std.testing.expectEqualStrings("initialize", msg2.method); + try std.testing.expectEqual(2, msg2.id); + try std.testing.expectEqualStrings("2.0", msg2.params.version); + + try std.testing.expectError(error.EndOfStream, reader.readJsonObject(TestObject)); + } + + // test reading messages with headers + { + const input = + \\Content-Length: 73 + \\ + \\{"jsonrpc":"2.0","method":"initialize","params":{"version":"2.0"},"id":1}Content-Length: 73 + \\ + \\{"jsonrpc":"2.0","method":"initialize","params":{"version":"2.0"},"id":2} + ; + + const allocator = std.testing.allocator; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const json_msg1 = try reader.readJsonObject(TestObject); + defer json_msg1.deinit(); + const msg1 = json_msg1.value; + + const json_msg2 = try reader.readJsonObject(TestObject); + defer json_msg2.deinit(); + const msg2 = json_msg2.value; + + try std.testing.expectEqualStrings("2.0", msg1.jsonrpc); + try std.testing.expectEqualStrings("initialize", msg1.method); + try std.testing.expectEqual(1, msg1.id); + try std.testing.expectEqualStrings("2.0", msg1.params.version); + + try std.testing.expectEqualStrings("2.0", msg2.jsonrpc); + try std.testing.expectEqualStrings("initialize", msg2.method); + try std.testing.expectEqual(2, msg2.id); + try std.testing.expectEqualStrings("2.0", msg2.params.version); + + try std.testing.expectError(error.EndOfStream, reader.readJsonObject(TestObject)); + } +} + +test readRawMessage { + const allocator = std.testing.allocator; + + // test reading newline delimited messages + { + const input = + \\{"jsonrpc":"2.0","id":1,"method":"initialize"} + \\{"jsonrpc":"2.0","id":2,"method":"initialize"} + ; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const msg1 = try reader.readRawMessage(); + defer allocator.free(msg1); + const msg2 = try reader.readRawMessage(); + defer allocator.free(msg2); + + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}", msg1); + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"initialize\"}", msg2); + try std.testing.expectError(error.EndOfStream, reader.readRawMessage()); + } + + // test reading messages with headers + { + const input = + \\Content-Length: 46 + \\ + \\{"jsonrpc":"2.0","id":1,"method":"initialize"}Content-Length: 46 + \\ + \\{"jsonrpc":"2.0","id":2,"method":"initialize"} + ; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const msg1 = try reader.readRawMessage(); + defer allocator.free(msg1); + const msg2 = try reader.readRawMessage(); + defer allocator.free(msg2); + + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}", msg1); + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"initialize\"}", msg2); + try std.testing.expectError(error.EndOfStream, reader.readRawMessage()); + } + + // test reading messages with lower-case header + { + const input = + \\content-length: 46 + \\ + \\{"jsonrpc":"2.0","id":1,"method":"initialize"} + ; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const msg1 = try reader.readRawMessage(); + defer allocator.free(msg1); + + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}", msg1); + try std.testing.expectError(error.EndOfStream, reader.readRawMessage()); + } +} + +test "readRawMessage - delimiter before EOF" { + const allocator = std.testing.allocator; + + const input = + \\{"jsonrpc":"2.0","id":1,"method":"initialize"} + \\ + ; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + const msg1 = try reader.readRawMessage(); + defer allocator.free(msg1); + + try std.testing.expectEqualStrings("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}", msg1); + try std.testing.expectError(error.EndOfStream, reader.readRawMessage()); +} + +test "readRawMessage - empty line" { + const allocator = std.testing.allocator; + + const input = + \\ + \\ + ; + var r = std.Io.Reader.fixed(input); + var reader: JsonRpcReader = .init(allocator, &r); + defer reader.deinit(); + + try std.testing.expectError(error.UnexpectedEndOfInput, reader.readRawMessage()); + try std.testing.expectError(error.EndOfStream, reader.readRawMessage()); +} diff --git a/src/acp/json_rpc/JsonRpcWriter.zig b/src/acp/json_rpc/JsonRpcWriter.zig new file mode 100644 index 0000000..a9c3110 --- /dev/null +++ b/src/acp/json_rpc/JsonRpcWriter.zig @@ -0,0 +1,183 @@ +//! Takes a `*std.Io.Writer` and supports writing objects to that stream via the JSON-RPC format. +//! +//! This is NOT threadsafe. + +const std = @import("std"); + +const Allocator = std.mem.Allocator; +const Io = std.Io; + +const JsonRpcWriter = @This(); + +allocator: Allocator, +writer: *Io.Writer, +write_buffer: std.Io.Writer.Allocating, + +/// Options for writing JSON-RPC messages. +pub const Options = struct { + /// Whether to prefix the message with a "Content-Length" header. + /// + /// Defaults to `false`. + use_headers: bool = false, +}; + +/// Initializes a new `JsonRpcWriter`. +/// +/// The returned `JsonRpcWriter` must be deinitialized when no longer needed by calling `deinit()`. +pub fn init(allocator: Allocator, writer: *Io.Writer) JsonRpcWriter { + return .{ + .allocator = allocator, + .writer = writer, + .write_buffer = .init(allocator), + }; +} + +/// Frees the resources associated with the `JsonRpcWriter`. +pub fn deinit(self: *JsonRpcWriter) void { + self.write_buffer.deinit(); +} + +/// Writes a JSON-RPC raw message (payload) directly to the stream. +pub fn writeRawMessage(self: *JsonRpcWriter, payload: []const u8, options: Options) !void { + if (options.use_headers) { + var header_buf: [64]u8 = undefined; + const header = try std.fmt.bufPrint(&header_buf, "Content-Length: {d}\r\n\r\n", .{payload.len}); + _ = try self.writer.write(header); + _ = try self.writer.write(payload); + } else { + _ = try self.writer.write(payload); + _ = try self.writer.write("\n"); + } + _ = try self.writer.flush(); +} + +/// Serializes the given value to JSON and writes it to the stream. +pub fn writeJsonObject(self: *JsonRpcWriter, value: anytype, options: Options) !void { + self.write_buffer.clearRetainingCapacity(); + var stringifier = std.json.Stringify{ + .writer = &self.write_buffer.writer, + .options = .{}, + }; + try stringifier.write(value); + + try self.writeRawMessage(self.write_buffer.written(), options); +} + +test writeJsonObject { + const TestObject = struct { + jsonrpc: []const u8, + method: []const u8, + id: u64, + }; + + const allocator = std.testing.allocator; + + // Test writing with headers + { + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var writer = JsonRpcWriter.init(allocator, &buffer.writer); + defer writer.deinit(); + + const msg = TestObject{ + .jsonrpc = "2.0", + .method = "initialize", + .id = 1, + }; + + try writer.writeJsonObject(msg, .{ .use_headers = true }); + + const expected = "Content-Length: 46\r\n\r\n{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":1}"; + try std.testing.expectEqualStrings(expected, buffer.written()); + } + + // Test writing without headers + { + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var writer = JsonRpcWriter.init(allocator, &buffer.writer); + defer writer.deinit(); + + const msg = TestObject{ + .jsonrpc = "2.0", + .method = "initialize", + .id = 2, + }; + + try writer.writeJsonObject(msg, .{ .use_headers = false }); + + const expected = "{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":2}\n"; + try std.testing.expectEqualStrings(expected, buffer.written()); + } +} + +test writeRawMessage { + const allocator = std.testing.allocator; + + // Test writing raw message with headers + { + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var writer = JsonRpcWriter.init(allocator, &buffer.writer); + defer writer.deinit(); + + const payload = "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}"; + try writer.writeRawMessage(payload, .{ .use_headers = true }); + + const expected = "Content-Length: 46\r\n\r\n{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\"}"; + try std.testing.expectEqualStrings(expected, buffer.written()); + } + + // Test writing raw message without headers + { + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var writer = JsonRpcWriter.init(allocator, &buffer.writer); + defer writer.deinit(); + + const payload = "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"initialize\"}"; + try writer.writeRawMessage(payload, .{ .use_headers = false }); + + const expected = "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"initialize\"}\n"; + try std.testing.expectEqualStrings(expected, buffer.written()); + } +} + +test "writeJsonObject - consecutive calls with headers" { + const TestObject = struct { + jsonrpc: []const u8, + method: []const u8, + id: u64, + }; + + const allocator = std.testing.allocator; + + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var writer = JsonRpcWriter.init(allocator, &buffer.writer); + defer writer.deinit(); + + const msg1 = TestObject{ + .jsonrpc = "2.0", + .method = "initialize", + .id = 1, + }; + try writer.writeJsonObject(msg1, .{ .use_headers = true }); + + const msg2 = TestObject{ + .jsonrpc = "2.0", + .method = "initialized", + .id = 2, + }; + try writer.writeJsonObject(msg2, .{ .use_headers = true }); + + const expected = + "Content-Length: 46\r\n\r\n{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":1}" ++ + "Content-Length: 47\r\n\r\n{\"jsonrpc\":\"2.0\",\"method\":\"initialized\",\"id\":2}"; + try std.testing.expectEqualStrings(expected, buffer.written()); +} diff --git a/src/acp/root.zig b/src/acp/root.zig new file mode 100644 index 0000000..b9efa00 --- /dev/null +++ b/src/acp/root.zig @@ -0,0 +1,7 @@ +const std = @import("std"); + +pub const Server = @import("Server.zig"); + +test { + std.testing.refAllDecls(@This()); +} diff --git a/src/acp/shared_api.zig b/src/acp/shared_api.zig new file mode 100644 index 0000000..c1a4b3c --- /dev/null +++ b/src/acp/shared_api.zig @@ -0,0 +1,416 @@ +//! Contains structs and functionality shared by both ACP client (requests) and +//! ACP servers (responses). + +const std = @import("std"); + +const Allocator = std.mem.Allocator; + +/// Parses a JSON string value into an enum tag, dynamically mapping enum tag names +/// containing underscores to wire names containing slashes (e.g. mapping `session_new` +/// to `"session/new"`). +/// +/// Unrecognized or invalid tags will return the specified `fallback` enum value. +/// String mapping is resolved entirely at compile-time and has zero runtime allocation cost. +pub fn parseEnumWithMappedFallback( + comptime T: type, + comptime fallback: T, + source: std.json.Value, +) !T { + if (source != .string) return error.UnexpectedToken; + const s = source.string; + + inline for (std.meta.fields(T)) |f| { + if (!std.mem.eql(u8, f.name, @tagName(fallback))) { + const wire_name = comptime blk: { + var buf: [f.name.len]u8 = undefined; + for (f.name, 0..) |char, i| { + buf[i] = if (char == '_') '/' else char; + } + const final_buf = buf; + break :blk &final_buf; + }; + + if (std.mem.eql(u8, s, wire_name)) { + return @field(T, f.name); + } + } + } + return fallback; +} + +/// Converts an enum tag value into a string slice, dynamically mapping enum tag names +/// containing underscores to wire names containing slashes (e.g. mapping `.session_new` +/// to `"session/new"`). +/// +/// String mapping is resolved entirely at compile-time and has zero runtime allocation cost. +pub fn stringifyEnum(value: anytype) []const u8 { + comptime { + const T = @TypeOf(value); + if (@typeInfo(T) != .@"enum") { + @compileError("stringifyEnum expects an enum value, found " ++ @typeName(T)); + } + } + + switch (value) { + inline else => |tag| { + return comptime blk: { + const name = @tagName(tag); + var buf: [name.len]u8 = undefined; + for (name, 0..) |char, i| { + buf[i] = if (char == '_') '/' else char; + } + const final_buf = buf; + break :blk &final_buf; + }; + }, + } +} + +/// A generic JSON stringifier that iterates over the fields of a struct. +/// It skips fields that are optional and have a null value. +/// +/// Useful when implementing a custom json stringifier that writes additional fields before the object fields. +/// But still needs the object fields. +/// TODO(razza): Move to a general location? This is re-used in provider/google/api. +pub fn jsonStringifyFields(object: anytype, jw: anytype) !void { + const info = @typeInfo(@TypeOf(object)); + if (info != .@"struct" and info != .@"union") { + @compileError("jsonStringifyFields only supports struct and union types"); + } + + inline for (std.meta.fields(@TypeOf(object))) |field| { + const value = @field(object, field.name); + if (@typeInfo(field.type) != .optional or value != null) { + try jw.objectField(field.name); + try jw.write(value); + } + } +} + +/// JSON RPC Request ID +/// +/// An identifier established by the Client that MUST contain a string, an integer, or null. +/// If it is not included, the request is assumed to be a notification. +/// +/// See the [JSON RPC spec](https://www.jsonrpc.org/specification). +pub const RequestId = union(enum) { + integer: i64, + string: []const u8, + null, + + pub fn jsonStringify(self: RequestId, jw: anytype) !void { + switch (self) { + .null => try jw.write(null), + inline else => |payload| { + try jw.write(payload); + }, + } + } + + pub fn jsonParse(allocator: Allocator, source: anytype, options: std.json.ParseOptions) !RequestId { + const json_value = try std.json.innerParse(std.json.Value, allocator, source, options); + return jsonParseFromValue(allocator, json_value, options); + } + + pub fn jsonParseFromValue(_: Allocator, source: std.json.Value, _: std.json.ParseOptions) !RequestId { + switch (source) { + .integer => |i| return .{ .integer = i }, + .string => |s| return .{ .string = s }, + .null => return .{ .null = {} }, + else => return error.InvalidEnumTag, + } + } +}; + +/// Metadata about the implementation of the client or agent. +/// +/// Describes the name and version of an ACP implementation, with an optional +/// title for display to the user. +pub const Implementation = struct { + /// Name identifying the ACP implementation intended for programmatic or + /// logical use, but can be used as a display name if `title` is not present. + name: []const u8, + /// Intended for UI and end-user contexts — optimized to be human-readable + /// and easily understood. + /// + /// If not provided, the `name` should be used for display. + title: ?[]const u8 = null, + /// Version of the implementation. Can be displayed to the user or used + /// for debugging or metrics purposes. (e.g. "1.0.0"). + version: []const u8, +}; + +/// Protocol version identifier. +/// +/// In JSON this is a Number but it is unlikely to be fractional so using u32 here. +/// +/// This version is only bumped for breaking changes. +/// Non-breaking changes should be introduced via capabilities. +pub const ProtocolVersion = u32; + +/// A unique identifier for a conversation session between a client and agent. +/// +/// Sessions maintain their own context, conversation history, and state, +/// allowing multiple independent interactions with the same agent. +/// +/// See protocol docs: [Session ID](https://agentclientprotocol.com/protocol/session-setup#session-id) +pub const SessionId = []const u8; + +/// Types of content in a `ContentBlock`. +/// +/// See protocol docs: [Content](https://agentclientprotocol.com/protocol/v1/content) +pub const ContentType = enum { + text, +}; + +/// Unique identifier for a tool call within a session. +/// +/// See protocol docs: [Tool Calls](https://agentclientprotocol.com/protocol/tool-calls) +pub const ToolCallId = []const u8; + +/// Content blocks represent displayable information in the Agent Client Protocol. +/// +/// They provide a structured way to handle various types of user-facing content—whether +/// it's text from language models, images for analysis, or embedded resources for context. +/// +/// Content blocks appear in: +/// - User prompts sent via `session/prompt` +/// - Language model output streamed through `session/update` notifications +/// - Progress updates and results from tool calls +/// +/// This structure is compatible with the Model Context Protocol (MCP), enabling +/// agents to seamlessly forward content from MCP tool outputs without transformation. +/// +/// See protocol docs: [Content](https://agentclientprotocol.com/protocol/v1/content) +pub const ContentBlock = union(ContentType) { + /// Text content. + text: []const u8, + + pub fn jsonParse(allocator: Allocator, source: anytype, options: std.json.ParseOptions) !ContentBlock { + const json_value = try std.json.innerParse(std.json.Value, allocator, source, options); + return jsonParseFromValue(allocator, json_value, options); + } + + pub fn jsonParseFromValue(allocator: Allocator, source: std.json.Value, options: std.json.ParseOptions) !ContentBlock { + if (source != .object) return error.UnexpectedToken; + const content_type = try std.json.innerParseFromValue(ContentType, allocator, source.object.get("type") orelse return error.MissingField, options); + return switch (content_type) { + .text => .{ .text = try std.json.innerParseFromValue([]const u8, allocator, source.object.get("text") orelse return error.MissingField, options) }, + }; + } + + pub fn jsonStringify(self: ContentBlock, jw: anytype) !void { + try jw.beginObject(); + try jw.objectField("type"); + try jw.write(@tagName(self)); + try jsonStringifyFields(self, jw); + try jw.endObject(); + } +}; + +test "RequestId json parsing - integer" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(RequestId, allocator, "42", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestId{ .integer = 42 }, parsed.value); +} + +test "RequestId json parsing - string" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(RequestId, allocator, "\"abc\"", .{}); + defer parsed.deinit(); + try std.testing.expect(parsed.value == .string); + try std.testing.expectEqualStrings("abc", parsed.value.string); +} + +test "RequestId json parsing - null" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(RequestId, allocator, "null", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(RequestId.null, parsed.value); +} + +test "Implementation json parsing - without title" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(Implementation, allocator, + \\{"name": "test", "version": "1.0.0"} + , .{}); + defer parsed.deinit(); + try std.testing.expectEqualStrings("test", parsed.value.name); + try std.testing.expect(parsed.value.title == null); + try std.testing.expectEqualStrings("1.0.0", parsed.value.version); +} + +test "Implementation json parsing - with title" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(Implementation, allocator, + \\{"name": "test", "title": "My Title", "version": "1.0.0"} + , .{}); + defer parsed.deinit(); + try std.testing.expectEqualStrings("test", parsed.value.name); + try std.testing.expectEqualStrings("My Title", parsed.value.title.?); + try std.testing.expectEqualStrings("1.0.0", parsed.value.version); +} + +test "ProtocolVersion json parsing" { + const allocator = std.testing.allocator; + const parsed = try std.json.parseFromSlice(ProtocolVersion, allocator, "1", .{}); + defer parsed.deinit(); + try std.testing.expectEqual(@as(ProtocolVersion, 1), parsed.value); +} + +test "ContentBlock json parsing - text" { + const allocator = std.testing.allocator; + const json_str = + \\{"type": "text", "text": "hello world"} + ; + const parsed = try std.json.parseFromSlice(ContentBlock, allocator, json_str, .{}); + defer parsed.deinit(); + + try std.testing.expectEqual(ContentType.text, @as(ContentType, parsed.value)); + try std.testing.expectEqualStrings("hello world", parsed.value.text); +} + +test "ContentBlock json parsing - invalid" { + const allocator = std.testing.allocator; + + // Missing text field + try std.testing.expectError( + error.MissingField, + std.json.parseFromSlice(ContentBlock, allocator, "{\"type\": \"text\"}", .{}), + ); + + // Missing type field + try std.testing.expectError( + error.MissingField, + std.json.parseFromSlice(ContentBlock, allocator, "{\"text\": \"hello\"}", .{}), + ); + + // Invalid type tag + try std.testing.expectError( + error.InvalidEnumTag, + std.json.parseFromSlice(ContentBlock, allocator, "{\"type\": \"invalid\", \"text\": \"hello\"}", .{}), + ); + + // Non-object token + try std.testing.expectError( + error.UnexpectedToken, + std.json.parseFromSlice(ContentBlock, allocator, "\"not an object\"", .{}), + ); +} + +test "ContentBlock json stringify - text" { + const allocator = std.testing.allocator; + const block = ContentBlock{ .text = "hello world" }; + + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var stringifier = std.json.Stringify{ + .writer = &buffer.writer, + .options = .{}, + }; + try stringifier.write(block); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, buffer.written(), .{}); + defer parsed.deinit(); + + try std.testing.expectEqualStrings("text", parsed.value.object.get("type").?.string); + try std.testing.expectEqualStrings("hello world", parsed.value.object.get("text").?.string); +} + +test parseEnumWithMappedFallback { + const TestEnum = enum { + simple, + under_score, + other_long_name, + unknown, + }; + + // Test simple tag + { + const val = std.json.Value{ .string = "simple" }; + const result = try parseEnumWithMappedFallback(TestEnum, .unknown, val); + try std.testing.expectEqual(TestEnum.simple, result); + } + + // Test tag containing underscore (mapped to slash) + { + const val = std.json.Value{ .string = "under/score" }; + const result = try parseEnumWithMappedFallback(TestEnum, .unknown, val); + try std.testing.expectEqual(TestEnum.under_score, result); + } + + // Test tag containing multiple underscores (mapped to multiple slashes) + { + const val = std.json.Value{ .string = "other/long/name" }; + const result = try parseEnumWithMappedFallback(TestEnum, .unknown, val); + try std.testing.expectEqual(TestEnum.other_long_name, result); + } + + // Test unknown fallback + { + const val = std.json.Value{ .string = "nonexistent/name" }; + const result = try parseEnumWithMappedFallback(TestEnum, .unknown, val); + try std.testing.expectEqual(TestEnum.unknown, result); + } + + // Test non-string value error + { + const val = std.json.Value{ .integer = 42 }; + try std.testing.expectError(error.UnexpectedToken, parseEnumWithMappedFallback(TestEnum, .unknown, val)); + } +} + +test stringifyEnum { + const TestEnum = enum { + simple, + under_score, + other_long_name, + unknown, + }; + + try std.testing.expectEqualStrings("simple", stringifyEnum(TestEnum.simple)); + try std.testing.expectEqualStrings("under/score", stringifyEnum(TestEnum.under_score)); + try std.testing.expectEqualStrings("other/long/name", stringifyEnum(TestEnum.other_long_name)); + try std.testing.expectEqualStrings("unknown", stringifyEnum(TestEnum.unknown)); +} + +test jsonStringifyFields { + const allocator = std.testing.allocator; + + const TestStruct = struct { + foo: []const u8, + bar: ?i32 = null, + baz: ?bool = null, + }; + + const obj = TestStruct{ + .foo = "hello", + .bar = 42, + .baz = null, + }; + + var buffer = std.Io.Writer.Allocating.init(allocator); + defer buffer.deinit(); + + var stringifier = std.json.Stringify{ + .writer = &buffer.writer, + .options = .{}, + }; + + try stringifier.beginObject(); + try stringifier.objectField("extra"); + try stringifier.write("prefix"); + try jsonStringifyFields(obj, &stringifier); + try stringifier.endObject(); + + const parsed = try std.json.parseFromSlice(std.json.Value, allocator, buffer.written(), .{}); + defer parsed.deinit(); + + const map = parsed.value.object; + try std.testing.expectEqualStrings("prefix", map.get("extra").?.string); + try std.testing.expectEqualStrings("hello", map.get("foo").?.string); + try std.testing.expectEqual(@as(i64, 42), map.get("bar").?.integer); + try std.testing.expect(map.get("baz") == null); +} diff --git a/src/agent/Agent.zig b/src/agent/Session.zig similarity index 60% rename from src/agent/Agent.zig rename to src/agent/Session.zig index dff4c76..0819b64 100644 --- a/src/agent/Agent.zig +++ b/src/agent/Session.zig @@ -1,14 +1,22 @@ +//! Represents a single agentic session with the specified provider. +//! +//! A session is a single context window with an agent model as well +//! as the configuration associated with that context window (tools, +//! system prompt, etc). + const std = @import("std"); const llm = @import("llm"); -const Tool = @import("./Tool.zig"); -const types = @import("./types.zig"); +const Tool = @import("Tool.zig"); +const types = @import("types.zig"); +const SessionState = @import("SessionState.zig"); +const ToolCallContext = @import("ToolCallContext.zig"); const Io = std.Io; const Allocator = std.mem.Allocator; const Provider = llm.Provider; const Future = std.Io.Future; -const Agent = @This(); +const Session = @This(); allocator: Allocator, io: Io, @@ -16,24 +24,25 @@ provider: Provider, tools: []const Tool, session_config: llm.types.SessionConfig, prev_continuation: ?llm.types.StepContinuation, +session_state: SessionState, const ToolError = error{ToolNotFound} || Tool.CallError; -pub const AgentError = ToolError || Provider.ProviderError; +pub const SessionError = ToolError || Provider.ProviderError; -/// Initializes a new Agent instance. +/// Initializes a new Session instance. /// /// `allocator` is used for all internal dynamic memory allocations. /// `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 Agent. -pub fn init(allocator: Allocator, io: Io, provider: Provider, config: types.AgentConfig) !Agent { +/// `config` contains the configuration for the 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| { descriptors[i] = tool.descriptor; } - return Agent{ + return Session{ .allocator = allocator, .io = io, .provider = provider, @@ -41,34 +50,37 @@ pub fn init(allocator: Allocator, io: Io, provider: Provider, config: types.Agen .session_config = .{ .model = config.model, .tools = descriptors, + .system_prompt = config.system_prompt, }, .prev_continuation = null, + .session_state = .init(allocator), }; } -/// Deinitializes the Agent, releasing any accumulated session history and internal resources. -pub fn deinit(self: *Agent) void { +/// Deinitializes the Session, releasing any accumulated session history and internal resources. +pub fn deinit(self: *Session) void { if (self.prev_continuation) |*ls| { ls.deinit(); self.prev_continuation = null; } + self.session_state.deinit(); self.allocator.free(self.session_config.tools); } -/// Executes a single non-streaming turn of the agent. +/// Executes a single non-streaming turn in the agentic session. /// -/// The agent sends the turn's prompt to the LLM, handles any tool calls recommended by the model +/// The session sends the turn's prompt to the LLM, handles any tool calls recommended by the model /// sequentially/concurrently, and returns a `TurnResult` containing the final output and history /// when the model is finished thinking and using tools. /// /// `turn` contains the prompt to send to the LLM. /// /// The caller is responsible for deinitializing the returned `TurnResult` by calling deinit() on it. -pub fn executeTurn(self: *Agent, turn: types.Turn) AgentError!types.TurnResult { +pub fn executeTurn(self: *Session, turn: types.Turn) SessionError!types.TurnResult { return self.executeTurnInternal(turn, null); } -/// Executes a single turn of the agent while streaming progress back via a callback. +/// Executes a single turn in the agentic session while streaming progress back via a callback. /// /// Model chunks and tool results are streamed back via `callback`. /// Like `executeTurn`, this handles intermediate tool executions, and returns a final `TurnResult`. @@ -79,11 +91,11 @@ pub fn executeTurn(self: *Agent, turn: types.Turn) AgentError!types.TurnResult { /// /// The caller is responsible for deinitializing the returned `TurnResult` by calling deinit() on it. pub fn executeTurnStreaming( - self: *Agent, + self: *Session, turn: types.Turn, callback: types.StreamingCallback, callback_context: ?*anyopaque, -) AgentError!types.TurnResult { +) SessionError!types.TurnResult { var agent_streaming_ctx: StreamingContext = .{ .callback = callback, .context = callback_context }; return self.executeTurnInternal(turn, &agent_streaming_ctx); } @@ -98,8 +110,8 @@ fn streamingCallbackProxy(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) voi streaming_ctx.callback(streaming_ctx.context, .{ .model_chunk = chunk }); } -fn executeToolCall(self: *Agent, tool_call: llm.types.ToolCall) ToolError!llm.types.ToolResult { - const tool = for (self.tools) |t| { +fn executeToolCall(self: *Session, tool_call: llm.types.ToolCall) ToolError!llm.types.ToolResult { + const tool = for (self.tools) |*t| { if (std.mem.eql(u8, t.descriptor.name, tool_call.name)) { break t; } @@ -107,14 +119,20 @@ fn executeToolCall(self: *Agent, tool_call: llm.types.ToolCall) ToolError!llm.ty return ToolError.ToolNotFound; }; - return try tool.execute(self.allocator, self.io, tool_call.id, tool_call.arguments); + return try tool.execute(self.allocator, self.io, &self.session_state, tool_call.id, tool_call.arguments); } -fn executeTurnInternal(self: *Agent, turn: types.Turn, callback_context: ?*StreamingContext) AgentError!types.TurnResult { +fn executeTurnInternal(self: *Session, turn: types.Turn, callback_context: ?*StreamingContext) SessionError!types.TurnResult { var next_steps: std.ArrayList(llm.types.Step) = .empty; const allocator = self.allocator; const io = self.io; defer next_steps.deinit(allocator); + + const maybe_context = try self.session_state.getInjectedContextString(allocator); + defer if (maybe_context) |context_str| allocator.free(context_str); + if (maybe_context) |context_str| { + try next_steps.append(allocator, .{ .prompt = context_str }); + } try next_steps.append(allocator, .{ .prompt = turn.prompt }); var intermediate_results: std.ArrayList(types.IntermediateStepResult) = .empty; @@ -130,6 +148,7 @@ fn executeTurnInternal(self: *Agent, turn: types.Turn, callback_context: ?*Strea } while (true) { + const step_outcome = if (callback_context) |cb| try self.provider.executeStepStreaming( allocator, @@ -207,10 +226,11 @@ const MockToolImpl = struct { } }; -test "Agent.executeTurn - no tool calls" { +test "Session.executeTurn - no tool calls" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); const mock_model = llm.types.Model{ @@ -218,11 +238,11 @@ test "Agent.executeTurn - no tool calls" { .display_name = "Mock Model", }; - var agent = try Agent.init(allocator, io, prov, .{ + var session = try Session.init(allocator, io, prov, .{ .model = mock_model, .tools = &.{}, }); - defer agent.deinit(); + defer session.deinit(); const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ @@ -230,20 +250,21 @@ test "Agent.executeTurn - no tool calls" { }; mock_provider.execute_step_results = &outcomes; - const turn = types.Turn{ .prompt = "Hi agent" }; + const turn = types.Turn{ .prompt = "Hi session" }; - var result = try agent.executeTurn(turn); + var result = try session.executeTurn(turn); defer result.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_provider.execute_step_calls); - try std.testing.expect(agent.prev_continuation != null); + try std.testing.expect(session.prev_continuation != null); try std.testing.expectEqualStrings("Hello user!", result.final_step.model_output[0].text); } -test "Agent.executeTurnStreaming - no tool calls" { +test "Session.executeTurnStreaming - no tool calls" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); const mock_model = llm.types.Model{ @@ -251,7 +272,7 @@ test "Agent.executeTurnStreaming - no tool calls" { .display_name = "Mock Model", }; - var agent: Agent = try .init( + var session: Session = try .init( allocator, io, prov, @@ -260,7 +281,7 @@ test "Agent.executeTurnStreaming - no tool calls" { .tools = &.{}, }, ); - defer agent.deinit(); + defer session.deinit(); const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Hello user!" }}, &.{}, &.{}); const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ @@ -268,7 +289,7 @@ test "Agent.executeTurnStreaming - no tool calls" { }; mock_provider.execute_step_results = &outcomes; - const turn = types.Turn{ .prompt = "Hi agent" }; + const turn = types.Turn{ .prompt = "Hi session" }; const DummyContext = struct { called: bool = false, @@ -283,19 +304,20 @@ test "Agent.executeTurnStreaming - no tool calls" { } }.cb; - var result = try agent.executeTurnStreaming(turn, callback, &dummy_ctx); + var result = try session.executeTurnStreaming(turn, callback, &dummy_ctx); defer result.deinit(); try std.testing.expectEqual(@as(usize, 1), mock_provider.execute_step_streaming_calls); try std.testing.expectEqual(@as(usize, 0), mock_provider.execute_step_calls); - try std.testing.expect(agent.prev_continuation != null); + try std.testing.expect(session.prev_continuation != null); try std.testing.expectEqualStrings("Hello user!", result.final_step.model_output[0].text); } -test "Agent.executeTurn - executes tool call and runs again" { +test "Session.executeTurn - executes tool call and runs again" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); const tool_desc = llm.types.Tool{ @@ -319,7 +341,7 @@ test "Agent.executeTurn - executes tool call and runs again" { .display_name = "Mock Model", }; - var agent: Agent = try .init( + var session: Session = try .init( allocator, io, prov, @@ -328,7 +350,7 @@ test "Agent.executeTurn - executes tool call and runs again" { .tools = tools, }, ); - defer agent.deinit(); + defer session.deinit(); const args = [_]llm.types.Argument{ .{ .name = "val", .value = .{ .integer = 42 } }, @@ -356,24 +378,24 @@ test "Agent.executeTurn - executes tool call and runs again" { }; mock_provider.execute_step_results = &outcomes; + const turn = types.Turn{ .prompt = "Hi session, run mock_tool" }; - const turn = types.Turn{ .prompt = "Hi agent, run mock_tool" }; - - var result = try agent.executeTurn(turn); + var result = try session.executeTurn(turn); defer result.deinit(); try std.testing.expectEqual(@as(usize, 2), mock_provider.execute_step_calls); - try std.testing.expect(agent.prev_continuation != null); + try std.testing.expect(session.prev_continuation != null); try std.testing.expectEqualStrings("Final output after tool", result.final_step.model_output[0].text); } -test "Agent.executeTurnStreaming - model chunks streaming" { +test "Session.executeTurnStreaming - model chunks streaming" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); - var agent = try Agent.init( + var session = try Session.init( allocator, io, prov, @@ -382,7 +404,7 @@ test "Agent.executeTurnStreaming - model chunks streaming" { .tools = &.{}, }, ); - defer agent.deinit(); + defer session.deinit(); const chunk1 = llm.types.StreamingChunk{ .event = .{ @@ -410,7 +432,7 @@ test "Agent.executeTurnStreaming - model chunks streaming" { }; const chunks = [_]llm.types.StreamingChunk{ chunk1, chunk2 }; - const chunks_list = [_][]const llm.types.StreamingChunk{ &chunks }; + 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!" }}, &.{}, &.{}); @@ -441,8 +463,8 @@ test "Agent.executeTurnStreaming - model chunks streaming" { var cb_state = CallbackState.init(allocator); defer cb_state.deinit(); - const turn = types.Turn{ .prompt = "Hi agent" }; - var result = try agent.executeTurnStreaming(turn, CallbackState.cb, &cb_state); + const turn = types.Turn{ .prompt = "Hi session" }; + var result = try session.executeTurnStreaming(turn, CallbackState.cb, &cb_state); defer result.deinit(); try std.testing.expectEqual(@as(usize, 2), cb_state.chunks.items.len); @@ -451,11 +473,12 @@ test "Agent.executeTurnStreaming - model chunks streaming" { try std.testing.expectEqualStrings("Hello world!", result.final_step.model_output[0].text); } -test "Agent.executeTurnStreaming - with tool calls" { +test "Session.executeTurnStreaming - with tool calls" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); const args_buf = [_]llm.types.Argument{ @@ -492,7 +515,7 @@ test "Agent.executeTurnStreaming - with tool calls" { const tool = Tool.init(tool_desc, MockToolImpl.execute); const tools = &[_]Tool{tool}; - var agent = try Agent.init( + var session = try Session.init( allocator, io, prov, @@ -501,7 +524,7 @@ test "Agent.executeTurnStreaming - with tool calls" { .tools = tools, }, ); - defer agent.deinit(); + defer session.deinit(); const CallbackState = struct { const Self = @This(); @@ -525,8 +548,8 @@ test "Agent.executeTurnStreaming - with tool calls" { var cb_state = CallbackState.init(allocator); defer cb_state.deinit(); - const turn = types.Turn{ .prompt = "Hi agent" }; - var result = try agent.executeTurnStreaming(turn, CallbackState.cb, &cb_state); + const turn = types.Turn{ .prompt = "Hi session" }; + var result = try session.executeTurnStreaming(turn, CallbackState.cb, &cb_state); defer result.deinit(); try std.testing.expectEqual(@as(usize, 1), cb_state.chunks.items.len); @@ -536,13 +559,14 @@ test "Agent.executeTurnStreaming - with tool calls" { try std.testing.expectEqualStrings("Tool result for 42", cb_state.chunks.items[0].tool_result.result); } -test "Agent.executeToolCall - tool not found" { +test "Session.executeToolCall - tool not found" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); - var agent = try Agent.init( + var session = try Session.init( allocator, io, prov, @@ -551,7 +575,7 @@ test "Agent.executeToolCall - tool not found" { .tools = &.{}, }, ); - defer agent.deinit(); + defer session.deinit(); const tool_call = llm.types.ToolCall{ .id = "call-id", @@ -559,13 +583,14 @@ test "Agent.executeToolCall - tool not found" { .arguments = &.{}, }; - try std.testing.expectError(error.ToolNotFound, agent.executeToolCall(tool_call)); + try std.testing.expectError(error.ToolNotFound, session.executeToolCall(tool_call)); } -test "Agent.executeTurn - tool call error cleanup" { +test "Session.executeTurn - tool call error cleanup" { const allocator = std.testing.allocator; const io = std.testing.io; var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); const prov = mock_provider.provider(); const tool_desc = llm.types.Tool{ @@ -582,7 +607,7 @@ test "Agent.executeTurn - tool call error cleanup" { const tool = Tool.init(tool_desc, error_tool_impl.execute); const tools = &[_]Tool{tool}; - var agent = try Agent.init( + var session = try Session.init( allocator, io, prov, @@ -591,7 +616,7 @@ test "Agent.executeTurn - tool call error cleanup" { .tools = tools, }, ); - defer agent.deinit(); + defer session.deinit(); const tool_calls = [_]llm.types.ToolCall{ .{ @@ -607,7 +632,217 @@ test "Agent.executeTurn - tool call error cleanup" { }; mock_provider.execute_step_results = &outcomes; - const turn = types.Turn{ .prompt = "Run error_tool" }; - try std.testing.expectError(error.ArgumentTypeMismatch, agent.executeTurn(turn)); + try std.testing.expectError(error.ArgumentTypeMismatch, session.executeTurn(turn)); +} + +test "Session.executeTurn - tool receives ToolCallContext" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); + const prov = mock_provider.provider(); + + const StateObj = struct { + value: i64, + fn initFn(self: *@This(), _: Allocator) !void { + self.value = 0; + } + }; + + const tool_desc = llm.types.Tool{ + .name = "state_tool", + .description = "A tool that interacts with ToolCallContext", + .parameters = &.{ + .{ + .name = "val", + .description = "integer value", + .type = .integer, + .required = true, + }, + }, + }; + + const state_tool_impl = struct { + fn execute(alloc: Allocator, call_ctx: ToolCallContext, val: i64) Tool.CallError![]const u8 { + const obj = call_ctx.getOrInitState(StateObj, StateObj.initFn) catch return error.OutOfMemory; + obj.value += val; + 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}; + + var session = try Session.init( + allocator, + io, + prov, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = tools, + }, + ); + defer session.deinit(); + + const args = [_]llm.types.Argument{ + .{ .name = "val", .value = .{ .integer = 50 } }, + }; + const tool_calls = [_]llm.types.ToolCall{ + .{ + .id = "call-id-123", + .name = "state_tool", + .arguments = @constCast(&args), + }, + }; + + const result1 = testing.MockProvider.stepResult(&.{}, &.{}, &tool_calls); + const 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() }, + }; + mock_provider.execute_step_results = &outcomes; + + const turn = types.Turn{ .prompt = "Run state_tool" }; + var turn_res = try session.executeTurn(turn); + defer turn_res.deinit(); + + const state_obj = session.session_state.getState(StateObj); + try std.testing.expect(state_obj != null); + try std.testing.expectEqual(@as(i64, 50), state_obj.?.value); +} + +test "Session.init forwards system_prompt to session_config" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider: testing.MockProvider = .{}; + defer mock_provider.deinit(); + const prov = mock_provider.provider(); + + var session = try Session.init( + allocator, + io, + prov, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .system_prompt = "Custom system prompt", + }, + ); + defer session.deinit(); + + try std.testing.expectEqualStrings("Custom system prompt", session.session_config.system_prompt.?); +} + +test "Session.executeTurn prepends injected context from session_state" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); + const prov = mock_provider.provider(); + + const tool_desc = llm.types.Tool{ + .name = "ctx_tool", + .description = "Tool with context", + .parameters = &.{}, + }; + const tool_impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const tool = Tool.init(tool_desc, tool_impl.run); + + var session = try Session.init(allocator, io, prov, .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = &.{tool}, + }); + defer session.deinit(); + + try session.session_state.setInjectedContext(&session.tools[0], "Injected tool info"); + + const step_result = testing.MockProvider.stepResult(&.{.{ .text = "Response text" }}, &.{}, &.{}); + const outcomes = [_](llm.Provider.ProviderError!llm.types.StepOutcome){ + .{ .result = step_result, .continuation = testing.MockProvider.stepContinuation() }, + }; + mock_provider.execute_step_results = &outcomes; + + const turn = types.Turn{ .prompt = "Hello with context" }; + var result = try session.executeTurn(turn); + defer result.deinit(); + + try std.testing.expectEqual(2, mock_provider.last_input_steps.?.len); + const expected_ctx = "[TOOL_CONTEXT: ctx_tool]\nInjected tool info\n[/TOOL_CONTEXT]\n\n"; + try std.testing.expectEqualStrings(expected_ctx, mock_provider.last_input_steps.?[0].prompt); + try std.testing.expectEqualStrings("Hello with context", mock_provider.last_input_steps.?[1].prompt); +} + +test "Session.executeTurn prepends injected context added during tool execution" { + const allocator = std.testing.allocator; + const io = std.testing.io; + var mock_provider = testing.MockProvider{}; + defer mock_provider.deinit(); + const prov = mock_provider.provider(); + + const dummy_tool_desc = llm.types.Tool{ + .name = "context_setter_tool", + .description = "Tool that sets injected context", + .parameters = &.{}, + }; + + const ContextToolImpl = struct { + pub fn run(allocator_arg: std.mem.Allocator, call_ctx: ToolCallContext) ![]const u8 { + try call_ctx.setInjectedContext("New context from tool execution"); + return try allocator_arg.dupe(u8, "Tool completed successfully"); + } + }; + + const tool = Tool.init(dummy_tool_desc, ContextToolImpl.run); + const tools = &[_]Tool{tool}; + + var session: Session = try .init( + allocator, + io, + prov, + .{ + .model = .{ .id = "mock-model", .display_name = "Mock Model" }, + .tools = tools, + }, + ); + defer session.deinit(); + + const tool_calls = [_]llm.types.ToolCall{ + .{ + .id = "call-id-999", + .name = "context_setter_tool", + .arguments = &.{}, + }, + }; + + 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" }}, &.{}, &.{}); + 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() }, + }; + mock_provider.execute_step_results = &outcomes; + + const turn = types.Turn{ .prompt = "Run tool and get context" }; + var turn_res = try session.executeTurn(turn); + defer turn_res.deinit(); + + // On step 2 of turn 1 (tool response), tool_result is sent directly without prepending prompt + try std.testing.expectEqual(@as(usize, 1), mock_provider.last_input_steps.?.len); + try std.testing.expectEqualStrings("Tool completed successfully", mock_provider.last_input_steps.?[0].tool_result.result); + + // On turn 2, injected context set by tool execution is prepended ahead of turn prompt + const turn2 = types.Turn{ .prompt = "Follow up prompt" }; + var turn2_res = try session.executeTurn(turn2); + defer turn2_res.deinit(); + + try std.testing.expectEqual(@as(usize, 2), mock_provider.last_input_steps.?.len); + const expected_ctx = "[TOOL_CONTEXT: context_setter_tool]\nNew context from tool execution\n[/TOOL_CONTEXT]\n\n"; + try std.testing.expectEqualStrings(expected_ctx, mock_provider.last_input_steps.?[0].prompt); + try std.testing.expectEqualStrings("Follow up prompt", mock_provider.last_input_steps.?[1].prompt); } diff --git a/src/agent/SessionState.zig b/src/agent/SessionState.zig new file mode 100644 index 0000000..59e2383 --- /dev/null +++ b/src/agent/SessionState.zig @@ -0,0 +1,333 @@ +//! Manages typed, session-scoped state objects. Within a given session there is one instance +//! of each `type`. + +const std = @import("std"); +const Tool = @import("Tool.zig"); + +const Allocator = std.mem.Allocator; +const ArenaAllocator = std.heap.ArenaAllocator; + +const SessionState = @This(); + +allocator: Allocator, +arena_allocator: ArenaAllocator, +state_store: std.AutoHashMapUnmanaged(usize, *anyopaque), +context_injections: std.AutoHashMapUnmanaged(*const Tool, []const u8), + +/// Initializes a new `SessionState` instance. +/// +/// `allocator` to be used to store session state objects and internal memory. +pub fn init(allocator: Allocator) SessionState { + return .{ + .allocator = allocator, + .arena_allocator = .init(allocator), + .state_store = .{}, + .context_injections = .{}, + }; +} + +/// Deinitializes the `SessionState`, releasing the state store map and all stored state +/// objects allocated in the internal arena allocator. +pub fn deinit(self: *SessionState) void { + self.state_store.deinit(self.allocator); + var it = self.context_injections.valueIterator(); + while (it.next()) |v| { + self.allocator.free(v.*); + } + self.context_injections.deinit(self.allocator); + self.arena_allocator.deinit(); + self.* = undefined; +} + +/// Retrieves a pointer to the state object of type `T`, initializing it first if it does not already exist. +/// +/// If an instance of `T` does not exist in the store, memory for `T` is allocated using an internal allocator. +/// This memory is managed by this struct. It need not be freed by the caller. +/// If `constructor` is provided, it is invoked with the pointer to the newly allocated `T` and an `Allocator` +/// (from the session's internal arena) to initialize its fields. +/// +/// Returns a pointer to the existing or newly initialized instance of `T`. +pub fn getOrInitState(self: *SessionState, comptime T: type, comptime constructor: ?fn (*T, Allocator) anyerror!void) !*T { + const obj_allocator = self.arena_allocator.allocator(); + const type_id = typeId(T); + const store_result = try self.state_store.getOrPut(self.allocator, type_id); + + if (store_result.found_existing) { + return @ptrCast(@alignCast(store_result.value_ptr.*)); + } + + const alloc_ptr = try obj_allocator.create(T); + errdefer { + obj_allocator.destroy(alloc_ptr); + _ = self.state_store.remove(type_id); + } + + if (constructor) |c| { + try c(alloc_ptr, obj_allocator); + } + + store_result.value_ptr.* = alloc_ptr; + return alloc_ptr; +} + +/// Returns a pointer to the state object of type `T` if it exists in the store, or `null` if it has not been initialized. +pub fn getState(self: *const SessionState, comptime T: type) ?*T { + const ptr = self.state_store.get(typeId(T)) orelse return null; + return @ptrCast(@alignCast(ptr)); +} + +/// Generates a unique numeric identifier for a given type `T` using the memory address of a static variable. +fn typeId(comptime T: type) usize { + const Container = struct { + comptime { + _ = T; + } + var id: u8 = 0; + }; + return @intFromPtr(&Container.id); +} + +/// Associates additional prompt context string with a specific `Tool`. +/// +/// This context will be injected ahead of the user prompt or tool result each turn. +/// This call copies the context-slice internally and so it has no lifetime requirements +/// beyond this call. +/// +/// If a context association already exists for `tool`, it is overwritten. +pub fn setInjectedContext(self: *SessionState, tool: *const Tool, context: []const u8) !void { + const context_copy = try self.allocator.dupe(u8, context); + errdefer self.allocator.free(context_copy); + self.clearInjectedContext(tool); + try self.context_injections.put(self.allocator, tool, context_copy); +} + +/// Removes any injected context string associated with `tool`. +/// +/// If no context association exists for `tool`, this function is a no-op. +pub fn clearInjectedContext(self: *SessionState, tool: *const Tool) void { + const existing = self.context_injections.fetchRemove(tool); + if (existing) |entry| { + self.allocator.free(entry.value); + } +} + +/// Returns the injected context string associated with `tool`, or `null` if none has been set. +pub fn getInjectedContext(self: *const SessionState, tool: *const Tool) ?[]const u8 { + return self.context_injections.get(tool); +} + +/// Formats and combines all registered tool context injections into a single string to be prepended +/// to the LLM system prompt, or `null` if no context injections are present. +/// +/// The resulting string is dynamically allocated using `allocator`. The caller owns the returned memory. +pub fn getInjectedContextString(self: *const SessionState, allocator: Allocator) !?[]const u8 { + if (self.context_injections.count() == 0) return null; + + var list: std.ArrayList(u8) = .empty; + defer list.deinit(allocator); + + var iterator = self.context_injections.iterator(); + while (iterator.next()) |entry| { + const tool = entry.key_ptr.*; + const context = entry.value_ptr.*; + try list.appendSlice(allocator, "[TOOL_CONTEXT: "); + try list.appendSlice(allocator, tool.descriptor.name); + try list.appendSlice(allocator, "]\n"); + try list.appendSlice(allocator, context); + try list.appendSlice(allocator, "\n[/TOOL_CONTEXT]\n\n"); + } + + return try list.toOwnedSlice(allocator); +} + +test getOrInitState { + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const DummyState = struct { + count: u32, + buffer: []u8, + + fn initFn(self: *@This(), allocator: Allocator) !void { + self.count = 100; + self.buffer = try allocator.alloc(u8, 16); + @memset(self.buffer, 0xAB); + } + }; + + // First call initializes + const ptr1 = try state.getOrInitState(DummyState, DummyState.initFn); + try std.testing.expectEqual(100, ptr1.count); + try std.testing.expectEqual(16, ptr1.buffer.len); + try std.testing.expectEqual(0xAB, ptr1.buffer[0]); + + // Second call returns cached pointer + const ptr2 = try state.getOrInitState(DummyState, DummyState.initFn); + try std.testing.expectEqual(ptr1, ptr2); +} + +test "getOrInit without constructor" { + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const Config = struct { + value: i32 = 42, + }; + + const ptr1 = try state.getOrInitState(Config, null); + ptr1.value = 123; + + const ptr2 = try state.getOrInitState(Config, null); + try std.testing.expectEqual(123, ptr2.value); + try std.testing.expectEqual(ptr1, ptr2); +} + +test "getOrInit constructor error cleanup" { + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const FailingState = struct { + fn initFail(_: *@This(), _: Allocator) !void { + return error.InitializationFailed; + } + }; + + try std.testing.expectError(error.InitializationFailed, state.getOrInitState(FailingState, FailingState.initFail)); + try std.testing.expectEqual(null, state.getState(FailingState)); +} + +test getState { + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const Config = struct { + value: i32 = 42, + }; + + try std.testing.expectEqual(null, state.getState(Config)); + + const ptr1 = try state.getOrInitState(Config, null); + ptr1.value = 100; + + const ptr2 = state.getState(Config); + try std.testing.expect(ptr2 != null); + try std.testing.expectEqual(ptr1, ptr2.?); + try std.testing.expectEqual(100, ptr2.?.value); +} + +test typeId { + const id_u32_1 = typeId(u32); + const id_u32_2 = typeId(u32); + const id_u64 = typeId(u64); + const id_str = typeId([]const u8); + + try std.testing.expectEqual(id_u32_1, id_u32_2); + try std.testing.expect(id_u32_1 != id_u64); + try std.testing.expect(id_u32_1 != id_str); + try std.testing.expect(id_u64 != id_str); +} + +test setInjectedContext { + const llm = @import("llm"); + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const tool_desc: llm.types.Tool = .{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }; + const tool_impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const tool = Tool.init(tool_desc, tool_impl.run); + + try state.setInjectedContext(&tool, "System context info"); + try std.testing.expectEqualStrings("System context info", state.getInjectedContext(&tool).?); + + // Overwrite existing context + try state.setInjectedContext(&tool, "Updated context info"); + try std.testing.expectEqualStrings("Updated context info", state.getInjectedContext(&tool).?); +} + +test getInjectedContext { + const llm = @import("llm"); + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const tool_desc: llm.types.Tool = .{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }; + const tool_impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const tool = Tool.init(tool_desc, tool_impl.run); + + try std.testing.expectEqual(null, state.getInjectedContext(&tool)); + + try state.setInjectedContext(&tool, "System context info"); + try std.testing.expectEqualStrings("System context info", state.getInjectedContext(&tool).?); +} + +test clearInjectedContext { + const llm = @import("llm"); + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + const tool_desc: llm.types.Tool = .{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }; + const tool_impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const tool = Tool.init(tool_desc, tool_impl.run); + + // Clearing non-existent entry is a safe no-op + state.clearInjectedContext(&tool); + try std.testing.expectEqual(null, state.getInjectedContext(&tool)); + + try state.setInjectedContext(&tool, "Some context"); + try std.testing.expectEqualStrings("Some context", state.getInjectedContext(&tool).?); + + state.clearInjectedContext(&tool); + try std.testing.expectEqual(null, state.getInjectedContext(&tool)); +} + +test getInjectedContextString { + const llm = @import("llm"); + var state = SessionState.init(std.testing.allocator); + defer state.deinit(); + + // Empty state returns null + const empty_str = try state.getInjectedContextString(std.testing.allocator); + try std.testing.expectEqual(null, empty_str); + + // Single tool context injection + const tool_desc: llm.types.Tool = .{ + .name = "test_tool", + .description = "Test tool description", + .parameters = &.{}, + }; + const tool_impl = struct { + pub fn run() ![]const u8 { + return ""; + } + }; + const tool = Tool.init(tool_desc, tool_impl.run); + + try state.setInjectedContext(&tool, "Use carefully."); + const single_str = (try state.getInjectedContextString(std.testing.allocator)).?; + defer std.testing.allocator.free(single_str); + const expected_single = "[TOOL_CONTEXT: test_tool]\nUse carefully.\n[/TOOL_CONTEXT]\n\n"; + try std.testing.expectEqualStrings(expected_single, single_str); +} diff --git a/src/agent/Tool.zig b/src/agent/Tool.zig index 6533d81..5a2863e 100644 --- a/src/agent/Tool.zig +++ b/src/agent/Tool.zig @@ -1,11 +1,17 @@ +//! The definition of a Tool an Agent can call. + const std = @import("std"); const llm = @import("llm"); +const SessionState = @import("SessionState.zig"); +const ToolCallContext = @import("ToolCallContext.zig"); const Allocator = std.mem.Allocator; const Io = std.Io; const Argument = llm.types.Argument; const ToolResult = llm.types.ToolResult; +pub const BuiltIn = @import("tool/root.zig"); + /// Errors that can occur when executing or calling a tool. pub const CallError = error{ /// A required argument was not provided. @@ -18,23 +24,26 @@ pub const CallError = error{ } || std.mem.Allocator.Error; const Tool = @This(); -const ToolExecuteFn = *const fn (allocator: Allocator, io: Io, args: []const Argument) CallError![]const u8; +const ToolExecuteFn = *const fn (allocator: Allocator, io: Io, call_ctx: ToolCallContext, ctx: ?*anyopaque, args: []const Argument) CallError![]const u8; descriptor: llm.types.Tool, execute_fn: ToolExecuteFn, +ctx: ?*anyopaque, /// Executes the tool with the given arguments. /// /// `allocator` is used to allocate memory for the result. /// `io` is the IO to use for the tool call. +/// `session_state` is the session-scoped state object store to use for the tool call. /// `id` is the identifier of the tool call. /// `args` is the list of arguments to pass to the tool function. All required arguments must be present. /// Order is not relevant. Unexpected arguments are ignored. /// /// Returns a `ToolResult` containing the result of the tool call. The caller is responsible /// for freeing the `ToolResult` by calling `deinit()`. -pub fn execute(self: *const Tool, allocator: Allocator, io: Io, id: []const u8, args: []const Argument) CallError!ToolResult { - const result = try self.execute_fn(allocator, io, args); +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); errdefer allocator.free(result); return ToolResult.initTakingResultOwnership(allocator, self.descriptor.name, id, result); } @@ -51,13 +60,51 @@ pub fn execute(self: *const Tool, allocator: Allocator, io: Io, id: []const u8, /// 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. +/// during the tool call, and/or a ToolCallContext struct which provides access to a session-scoped +/// state store and the ability to consiste /// /// 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 { - return comptime blk: { + if (comptime expectsContext(execute_fn)) { + @compileError("Tool function '" ++ descriptor.name ++ "' expects a context parameter. Use Tool.initWithContext instead."); + } + return initInternal(descriptor, execute_fn, null); +} + +/// Creates a Tool from a descriptor, a function, and a provided context object. +/// 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, +/// 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. +/// +/// Additionally, the function can optionally accept an Io struct which represents the IO to use +/// during the tool call, and/or a ToolCallContext struct. +/// +/// 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."); + } + 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)) { @@ -66,15 +113,29 @@ pub fn init(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) T } } - const res = makeExecuteFn(descriptor, execute_fn); + 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, }; }; + + 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; } const ValidationError = error{ @@ -149,7 +210,74 @@ inline fn findArgument(args: []const Argument, name: []const u8) ?*const Argumen return null; } -fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anytype) ValidationResult { +/// 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 { const FnType = @TypeOf(execute_fn); const fn_info = switch (@typeInfo(FnType)) { .@"fn" => |info| info, @@ -174,7 +302,18 @@ 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 and fn_param.type.? != Io) { + if (fn_param.type.? == Allocator or fn_param.type.? == Io or fn_param.type.? == ToolCallContext) { + // 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 (param_idx >= descriptor_params.len) { return .{ .err = .{ .code = ValidationError.ArgumentCountMismatch, .msg = "More arguments in function than descriptor." } }; } @@ -201,17 +340,25 @@ 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, input_args: []const Argument) CallError![]const u8 { + pub fn call(allocator: Allocator, io: Io, call_ctx: ToolCallContext, ctx: ?*anyopaque, input_args: []const Argument) CallError![]const u8 { var args: TupleType = undefined; - var descriptor_idx: usize = 0; + comptime var descriptor_idx: usize = 0; inline for (0..fn_info.params.len) |func_idx| { const T = types[func_idx]; - if (T == Allocator) { + if (comptime T == Allocator) { args[func_idx] = allocator; - } else if (T == Io) { + } else if (comptime T == Io) { 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 { - const curr_descriptor = descriptor.parameters[descriptor_idx]; + const curr_descriptor = comptime descriptor.parameters[descriptor_idx]; descriptor_idx += 1; if (findArgument(input_args, curr_descriptor.name)) |argument| { const expected_tag = comptime try expectedTagForType(T); @@ -225,8 +372,8 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty .boolean => argument.value.boolean, }; } else { - if (curr_descriptor.required) return CallError.RequiredArgumentMissing; - if (@typeInfo(@TypeOf(args[func_idx])) != .optional) unreachable; + if (comptime curr_descriptor.required) return CallError.RequiredArgumentMissing; + if (comptime @typeInfo(@TypeOf(args[func_idx])) != .optional) unreachable; args[func_idx] = null; } } @@ -241,6 +388,8 @@ fn makeExecuteFn(comptime descriptor: llm.types.Tool, comptime execute_fn: anyty test init { 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 = "example_function", @@ -272,7 +421,7 @@ test init { .{ .name = "arg2", .value = .{ .string = "hello" } }, }; - var result = try tool.execute(allocator, io, "123", &args); + var result = try tool.execute(allocator, io, &session_state, "123", &args); defer result.deinit(); try std.testing.expectEqualStrings("example_function", result.tool_name); @@ -280,13 +429,60 @@ test init { try std.testing.expectEqualStrings("12hello", result.result); } +test initWithContext { + 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 = "example_context_function", + .description = "An example function that takes context and two arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument to the function", + .type = .integer, + .required = true, + }, + .{ + .name = "arg2", + .description = "The second argument to the function", + .type = .string, + .required = true, + }, + }, + }; + const CtxStruct = struct { prefix: []const u8 }; + var ctx_obj: CtxStruct = .{ .prefix = "ctx-" }; + + 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 tool = initWithContext(tool_descriptor, tool_impl.example_context_function, &ctx_obj); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .integer = 12 } }, + .{ .name = "arg2", .value = .{ .string = "hello" } }, + }; + + var result = try tool.execute(allocator, io, &session_state, "123", &args); + defer result.deinit(); + + try std.testing.expectEqualStrings("example_context_function", result.tool_name); + try std.testing.expectEqualStrings("123", result.id); + try std.testing.expectEqualStrings("ctx-12hello", result.result); +} + test "makeExecuteFn - ExpectedFunctionOrPointer" { const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", .parameters = &.{}, }; - const res = comptime makeExecuteFn(desc, 42); + const res = comptime makeExecuteFn(desc, 42, null); try std.testing.expectEqual(ValidationError.ExpectedFunctionOrPointer, res.err.code); } @@ -308,7 +504,7 @@ test "makeExecuteFn - ArgumentTypeMismatch" { return arg1; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); 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); } @@ -325,7 +521,7 @@ test "makeExecuteFn - ArgumentCountMismatch (too many arguments)" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); try std.testing.expectEqualStrings("More arguments in function than descriptor.", res.err.msg); } @@ -348,7 +544,7 @@ test "makeExecuteFn - ArgumentCountMismatch (too few arguments)" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); try std.testing.expectEqual(ValidationError.ArgumentCountMismatch, res.err.code); try std.testing.expectEqualStrings("Fewer arguments in function than descriptor.", res.err.msg); } @@ -372,7 +568,7 @@ test "makeExecuteFn - ParamTypeArrayNotSupported" { return arg1; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); 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); } @@ -396,7 +592,7 @@ test "makeExecuteFn - ReturnTypeMismatch" { return 10; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); try std.testing.expectEqual(ValidationError.ReturnTypeMismatch, res.err.code); try std.testing.expectEqualStrings("Function return type must be a string.", res.err.msg); } @@ -420,7 +616,7 @@ test "makeExecuteFn - argument optionals OK" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); try std.testing.expect(res == .ok); } @@ -442,7 +638,7 @@ test "makeExecuteFn - argument optional in descriptor, required in fn" { return required; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); 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); } @@ -466,7 +662,7 @@ test "makeExecuteFn - argument required in descriptor, optional in fn" { return ""; } }; - const res = comptime makeExecuteFn(desc, Impl.run); + const res = comptime makeExecuteFn(desc, Impl.run, null); 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); } @@ -474,6 +670,9 @@ test "makeExecuteFn - argument required in descriptor, optional in fn" { test execute { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -498,7 +697,7 @@ test execute { .value = .{ .string = "value" }, }, }; - var result = try tool.execute(testing_allocator, io, "id", args); + var result = try tool.execute(testing_allocator, io, &session_state, "id", args); defer result.deinit(); try std.testing.expectEqualStrings("test_tool", result.tool_name); @@ -509,6 +708,9 @@ test execute { test "execute - unknown argument" { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -533,12 +735,15 @@ test "execute - unknown argument" { .value = .{ .string = "value" }, }, }; - try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, "id", args)); + try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, &session_state, "id", args)); } test "execute - extra argument ignored" { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -567,7 +772,7 @@ test "execute - extra argument ignored" { .value = .{ .string = "value2" }, }, }; - var result = try tool.execute(testing_allocator, io, "id", args); + var result = try tool.execute(testing_allocator, io, &session_state, "id", args); defer result.deinit(); try std.testing.expectEqualStrings("value1", result.result); } @@ -575,6 +780,9 @@ test "execute - extra argument ignored" { test "execute - 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 desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -594,12 +802,15 @@ test "execute - missing required argument" { }; const tool = Tool.init(desc, Impl.run); const args: []const Argument = &.{}; - try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, "id", args)); + try std.testing.expectError(CallError.RequiredArgumentMissing, tool.execute(testing_allocator, io, &session_state, "id", args)); } test "execute - missing optional argument" { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -620,7 +831,7 @@ test "execute - missing optional argument" { }; const tool = Tool.init(desc, Impl.run); const args: []const Argument = &.{}; - var result = try tool.execute(testing_allocator, io, "id", args); + var result = try tool.execute(testing_allocator, io, &session_state, "id", args); defer result.deinit(); try std.testing.expectEqualStrings("test_tool", result.tool_name); @@ -631,6 +842,9 @@ test "execute - missing optional argument" { test "execute - optional argument provided" { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -653,7 +867,7 @@ test "execute - optional argument provided" { const args: []const Argument = &.{ .{ .name = "arg1", .value = .{ .string = "provided" } }, }; - var result = try tool.execute(testing_allocator, io, "id", args); + var result = try tool.execute(testing_allocator, io, &session_state, "id", args); defer result.deinit(); try std.testing.expectEqualStrings("test_tool", result.tool_name); @@ -664,6 +878,9 @@ test "execute - optional argument provided" { test "execute - argument type mismatch" { const testing_allocator = std.testing.allocator; const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + const desc: llm.types.Tool = .{ .name = "test_tool", .description = "desc", @@ -688,12 +905,14 @@ test "execute - argument type mismatch" { .value = .{ .integer = 10 }, }, }; - try std.testing.expectError(CallError.ArgumentTypeMismatch, tool.execute(testing_allocator, io, "id", args)); + try std.testing.expectError(CallError.ArgumentTypeMismatch, tool.execute(testing_allocator, io, &session_state, "id", args)); } test "execute - no Allocator parameter" { 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 = "no_allocator_func", @@ -718,7 +937,7 @@ test "execute - no Allocator parameter" { .{ .name = "arg1", .value = .{ .string = "hello" } }, }; - var result = try tool.execute(allocator, io, "abc", &args); + var result = try tool.execute(allocator, io, &session_state, "abc", &args); defer result.deinit(); try std.testing.expectEqualStrings("no_allocator_func", result.tool_name); @@ -729,6 +948,8 @@ test "execute - no Allocator parameter" { test "execute - Allocator as middle/last parameter" { 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 = "middle_last_allocator_func", @@ -760,7 +981,7 @@ test "execute - Allocator as middle/last parameter" { .{ .name = "arg2", .value = .{ .string = "test" } }, }; - var result = try tool.execute(allocator, io, "xyz", &args); + var result = try tool.execute(allocator, io, &session_state, "xyz", &args); defer result.deinit(); try std.testing.expectEqualStrings("middle_last_allocator_func", result.tool_name); @@ -771,6 +992,8 @@ test "execute - Allocator as middle/last parameter" { test "execute - multiple Allocator parameters" { 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_allocator_func", @@ -796,7 +1019,7 @@ test "execute - multiple Allocator parameters" { .{ .name = "arg1", .value = .{ .integer = 7 } }, }; - var result = try tool.execute(allocator, io, "multi", &args); + var result = try tool.execute(allocator, io, &session_state, "multi", &args); defer result.deinit(); try std.testing.expectEqualStrings("multi_allocator_func", result.tool_name); @@ -807,6 +1030,8 @@ test "execute - multiple Allocator parameters" { test "execute - Io parameter" { 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 = "io_func", @@ -832,7 +1057,7 @@ test "execute - Io parameter" { .{ .name = "arg1", .value = .{ .string = "test" } }, }; - var result = try tool.execute(allocator, io, "io-test", &args); + var result = try tool.execute(allocator, io, &session_state, "io-test", &args); defer result.deinit(); try std.testing.expectEqualStrings("io_func", result.tool_name); @@ -843,6 +1068,8 @@ test "execute - Io parameter" { test "execute - multiple Io parameters" { 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_io_func", @@ -869,10 +1096,242 @@ test "execute - multiple Io parameters" { .{ .name = "arg1", .value = .{ .integer = 77 } }, }; - var result = try tool.execute(allocator, io, "multi-io", &args); + var result = try tool.execute(allocator, io, &session_state, "multi-io", &args); defer result.deinit(); try std.testing.expectEqualStrings("multi_io_func", result.tool_name); try std.testing.expectEqualStrings("multi-io", result.id); try std.testing.expectEqualStrings("multi-io-77", result.result); } + +test expectsContext { + const ImplNoCtx = struct { + pub fn run(_: Allocator, _: i64) ![]const u8 { + return ""; + } + }; + const ImplWithCtx = struct { + pub fn run(_: Allocator, _: i64, _: *anyopaque) ![]const u8 { + return ""; + } + }; + const CtxStruct = struct { val: u32 }; + const ImplWithTypedCtx = struct { + pub fn run(_: Allocator, _: i64, _: *CtxStruct) ![]const u8 { + return ""; + } + }; + const ImplWithCallContext = struct { + pub fn run(_: Allocator, _: ToolCallContext, _: i64) ![]const u8 { + return ""; + } + }; + 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)); +} + +test "execute - has context" { + 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 = "ctx_func", + .description = "Takes context and arguments", + .parameters = &.{ + .{ + .name = "arg1", + .description = "The first argument", + .type = .string, + .required = true, + }, + }, + }; + 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 }); + } + }; + + const tool = initWithContext(tool_descriptor, tool_impl.ctx_func, &ctx_obj); + + const args = [_]Argument{ + .{ .name = "arg1", .value = .{ .string = "world" } }, + }; + + var result = try tool.execute(allocator, io, &session_state, "ctx-id", &args); + defer result.deinit(); + + 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); +} + +test "execute - const and optional 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 = "const_ctx_func", + .description = "Takes const and optional context pointers", + .parameters = &.{ + .{ + .name = "val", + .type = .integer, + .required = true, + .description = "val", + }, + }, + }; + + const CtxStruct = struct { factor: i64 }; + const const_ctx_obj: CtxStruct = .{ .factor = 10 }; + var mut_ctx_obj: CtxStruct = .{ .factor = 10 }; + + 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}); + } + 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}); + } + }; + + const args = [_]Argument{.{ .name = "val", .value = .{ .integer = 5 } }}; + + // 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); + + // 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); + + // 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); + + // 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); +} + +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", + }, + }, + }; + + const CtxStruct = struct { factor: i64 }; + var ctx_obj: CtxStruct = .{ .factor = 7 }; + + 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}); + } + }; + + 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)); +} + +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 = "counter_tool", + .description = "Increments counter in SessionState", + .parameters = &.{ + .{ + .name = "increment", + .type = .integer, + .required = true, + .description = "amount to increment", + }, + }, + }; + + 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}); + } + }; + + 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); + + 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); +} diff --git a/src/agent/ToolCallContext.zig b/src/agent/ToolCallContext.zig new file mode 100644 index 0000000..eea08e9 --- /dev/null +++ b/src/agent/ToolCallContext.zig @@ -0,0 +1,186 @@ +//! Execution context provided to a `Tool` during call execution. +//! Encapsulates both the specific `Tool` being executed and the `SessionState`. + +const std = @import("std"); +const SessionState = @import("SessionState.zig"); +const Tool = @import("Tool.zig"); + +const Allocator = std.mem.Allocator; + +const ToolCallContext = @This(); + +tool: *const Tool, +session_state: *SessionState, + +/// Associates additional prompt context string with the calling tool. +/// +/// This context will be injected ahead of the user prompt or tool result each turn. +/// This call copies the context-slice internally and so it has no lifetime requirements +/// beyond this call. +/// +/// If a context association already exists for the calling tool, it is overwritten. +pub fn setInjectedContext(self: ToolCallContext, context: []const u8) !void { + return self.session_state.setInjectedContext(self.tool, context); +} + +/// Removes any injected context string associated with the calling tool. +/// +/// If no context association exists for the calling tool, this function is a no-op. +pub fn clearInjectedContext(self: ToolCallContext) void { + self.session_state.clearInjectedContext(self.tool); +} + +/// Returns the injected context string associated with the calling tool, or `null` if none has been set. +pub fn getInjectedContext(self: ToolCallContext) ?[]const u8 { + return self.session_state.getInjectedContext(self.tool); +} + +/// Retrieves a pointer to the state object of type `T`, initializing it first if it does not already exist. +/// +/// If an instance of `T` does not exist in the store, memory for `T` is allocated using an internal allocator. +/// Memory is managed by `SessionState`. +pub fn getOrInitState(self: ToolCallContext, comptime T: type, comptime constructor: ?fn (*T, Allocator) anyerror!void) !*T { + return self.session_state.getOrInitState(T, constructor); +} + +/// Returns a pointer to the state object of type `T` if it exists in the store, or `null` if it has not been initialized. +pub fn getState(self: ToolCallContext, comptime T: type) ?*T { + return self.session_state.getState(T); +} + +test setInjectedContext { + var session_state = SessionState.init(std.testing.allocator); + defer session_state.deinit(); + + const tool = Tool.init(.{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }, struct { + fn run() ![]const u8 { + return ""; + } + }.run); + + const call_ctx: ToolCallContext = .{ + .tool = &tool, + .session_state = &session_state, + }; + + try call_ctx.setInjectedContext("System context info"); + try std.testing.expectEqualStrings("System context info", session_state.getInjectedContext(&tool).?); +} + +test clearInjectedContext { + var session_state = SessionState.init(std.testing.allocator); + defer session_state.deinit(); + + const tool = Tool.init(.{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }, struct { + fn run() ![]const u8 { + return ""; + } + }.run); + + const call_ctx: ToolCallContext = .{ + .tool = &tool, + .session_state = &session_state, + }; + + try call_ctx.setInjectedContext("System context info"); + try std.testing.expectEqualStrings("System context info", call_ctx.getInjectedContext().?); + + call_ctx.clearInjectedContext(); + try std.testing.expectEqual(null, call_ctx.getInjectedContext()); +} + +test getInjectedContext { + var session_state = SessionState.init(std.testing.allocator); + defer session_state.deinit(); + + const tool = Tool.init(.{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }, struct { + fn run() ![]const u8 { + return ""; + } + }.run); + + const call_ctx: ToolCallContext = .{ + .tool = &tool, + .session_state = &session_state, + }; + + try std.testing.expectEqual(null, call_ctx.getInjectedContext()); + + try call_ctx.setInjectedContext("System context info"); + try std.testing.expectEqualStrings("System context info", call_ctx.getInjectedContext().?); +} + +test getOrInitState { + var session_state = SessionState.init(std.testing.allocator); + defer session_state.deinit(); + + const tool = Tool.init(.{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }, struct { + fn run() ![]const u8 { + return ""; + } + }.run); + + const call_ctx: ToolCallContext = .{ + .tool = &tool, + .session_state = &session_state, + }; + + const DummyState = struct { + count: u32, + fn initFn(self: *@This(), _: Allocator) !void { + self.count = 100; + } + }; + + const ptr = try call_ctx.getOrInitState(DummyState, DummyState.initFn); + try std.testing.expectEqual(100, ptr.count); +} + +test getState { + var session_state = SessionState.init(std.testing.allocator); + defer session_state.deinit(); + + const tool = Tool.init(.{ + .name = "dummy_tool", + .description = "A dummy tool", + .parameters = &.{}, + }, struct { + fn run() ![]const u8 { + return ""; + } + }.run); + + const call_ctx: ToolCallContext = .{ + .tool = &tool, + .session_state = &session_state, + }; + + const Config = struct { + value: i32 = 42, + }; + + try std.testing.expectEqual(null, call_ctx.getState(Config)); + + const ptr1 = try call_ctx.getOrInitState(Config, null); + ptr1.value = 100; + + const ptr2 = call_ctx.getState(Config); + try std.testing.expect(ptr2 != null); + try std.testing.expectEqual(100, ptr2.?.value); +} diff --git a/src/agent/root.zig b/src/agent/root.zig index 72fc1fd..3385a70 100644 --- a/src/agent/root.zig +++ b/src/agent/root.zig @@ -1,9 +1,11 @@ const std = @import("std"); -pub const Agent = @import("Agent.zig"); +pub const Session = @import("Session.zig"); +pub const ToolCallContext = @import("ToolCallContext.zig"); pub const Tool = @import("Tool.zig"); pub const types = @import("types.zig"); test { std.testing.refAllDecls(@This()); + std.testing.refAllDecls(@import("tool/root.zig")); } diff --git a/src/agent/tool/root.zig b/src/agent/tool/root.zig new file mode 100644 index 0000000..cbee524 --- /dev/null +++ b/src/agent/tool/root.zig @@ -0,0 +1,9 @@ +const std = @import("std"); + +/// A simple todo list tool for managing tasks. +pub const Todo = @import("todo.zig").Tool; + +test { + std.testing.refAllDecls(@This()); +} + diff --git a/src/agent/tool/todo.zig b/src/agent/tool/todo.zig new file mode 100644 index 0000000..a9ff243 --- /dev/null +++ b/src/agent/tool/todo.zig @@ -0,0 +1,148 @@ +//! Agent tool for managing and injecting persistent TODO content into session context. + +const std = @import("std"); +const llm = @import("llm"); +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(.{ + .name = "write_todo", + .description = + \\Overwrite the entire TODO content. + \\ + \\The content persists across conversation turns and compaction and + \\is injected frequently into the context window. Use this for: + \\ - Task tracking and progress updates + \\ - Important notes and reminders + \\ + \\Writing an empty string clears the entire todo list. + \\ + \\WARNING: This operation completely replaces the existing content. + \\Always include all content you want to keep, not just the changes. + , + .parameters = &.{ + .{ + .name = "content", + .description = "The todo list in markdown format. To clear the todo list, pass an empty string.", + .type = .string, + .required = true, + }, + }, +}, execute); + +/// Executes the tool call to set or clear the session's injected TODO context. +fn execute(allocator: std.mem.Allocator, ctx: agent.ToolCallContext, content: []const u8) ![]const u8 { + if (content.len == 0) { + ctx.clearInjectedContext(); + return try allocator.dupe(u8, "Todo list cleared."); + } else { + try ctx.setInjectedContext(content); + return try allocator.dupe(u8, "Todo list updated."); + } +} + +test "Todo tool descriptor metadata" { + try std.testing.expectEqualStrings("write_todo", Tool.descriptor.name); + try std.testing.expectEqual(1, Tool.descriptor.parameters.len); + try std.testing.expectEqualStrings("content", 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); +} + +test "Todo tool update content" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const args: []const llm.types.Argument = &.{ + .{ + .name = "content", + .value = .{ .string = "- [ ] Write tests\n- [ ] Run tests" }, + }, + }; + + var result = try Tool.execute(testing_allocator, io, &session_state, "call_1", args); + defer result.deinit(); + + try std.testing.expectEqualStrings("write_todo", result.tool_name); + try std.testing.expectEqualStrings("call_1", result.id); + try std.testing.expectEqualStrings("Todo list updated.", result.result); + + const injected = session_state.getInjectedContext(&Tool); + try std.testing.expect(injected != null); + try std.testing.expectEqualStrings("- [ ] Write tests\n- [ ] Run tests", injected.?); +} + +test "Todo tool clear content with empty string" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + // First update the todo list + const update_args: []const llm.types.Argument = &.{ + .{ + .name = "content", + .value = .{ .string = "- [ ] Task to clear" }, + }, + }; + var update_res = try Tool.execute(testing_allocator, io, &session_state, "call_1", update_args); + update_res.deinit(); + + try std.testing.expectEqualStrings("- [ ] Task to clear", session_state.getInjectedContext(&Tool).?); + + // Now clear it + const clear_args: []const llm.types.Argument = &.{ + .{ + .name = "content", + .value = .{ .string = "" }, + }, + }; + var clear_res = try Tool.execute(testing_allocator, io, &session_state, "call_2", clear_args); + defer clear_res.deinit(); + + try std.testing.expectEqualStrings("Todo list cleared.", clear_res.result); + try std.testing.expectEqual(null, session_state.getInjectedContext(&Tool)); +} + +test "Todo tool overwrite existing content" { + const testing_allocator = std.testing.allocator; + const io = std.testing.io; + var session_state = SessionState.init(testing_allocator); + defer session_state.deinit(); + + const args1: []const llm.types.Argument = &.{ + .{ + .name = "content", + .value = .{ .string = "Initial todo list" }, + }, + }; + var res1 = try Tool.execute(testing_allocator, io, &session_state, "call_1", args1); + res1.deinit(); + + try std.testing.expectEqualStrings("Initial todo list", session_state.getInjectedContext(&Tool).?); + + const args2: []const llm.types.Argument = &.{ + .{ + .name = "content", + .value = .{ .string = "Overwritten todo list" }, + }, + }; + var res2 = try Tool.execute(testing_allocator, io, &session_state, "call_2", args2); + defer res2.deinit(); + + try std.testing.expectEqualStrings("Todo list updated.", res2.result); + try std.testing.expectEqualStrings("Overwritten todo list", session_state.getInjectedContext(&Tool).?); +} + +test "Todo 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 empty_args: []const llm.types.Argument = &.{}; + try std.testing.expectError(agent.Tool.CallError.RequiredArgumentMissing, Tool.execute(testing_allocator, io, &session_state, "call_1", empty_args)); +} diff --git a/src/agent/types.zig b/src/agent/types.zig index 9323595..45af69c 100644 --- a/src/agent/types.zig +++ b/src/agent/types.zig @@ -4,21 +4,23 @@ const Tool = @import("Tool.zig"); const Allocator = std.mem.Allocator; -/// Configuration for initializing an Agent. -pub const AgentConfig = struct { - /// The LLM model to be used by the Agent. +/// Configuration for initializing a Session. +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 Agent to use. + /// A list of executable tools available for the Session to use. tools: []const Tool = &.{}, + /// Optional system prompt for the Session. + system_prompt: ?[]const u8 = null, }; -/// Represents the input to start a single execution turn for the Agent. +/// Represents the input to start a single execution turn in a Session. pub const Turn = struct { /// The input text prompt or message for this turn. prompt: []const u8, }; -/// The result of an intermediate step executed during an agent turn. +/// The result of an intermediate step executed during a turn. /// /// An intermediate step is either a response from the model that required further /// processing (e.g. a tool call) or the result of executing a tool call requested by @@ -39,7 +41,7 @@ pub const IntermediateStepResult = union(enum) { } }; -/// The final result of an agent turn, including all intermediate steps taken and the final LLM step result. +/// The final result of a turn, including all intermediate steps taken and the final LLM step result. /// Memory must be freed using the `deinit` method. pub const TurnResult = struct { /// The allocator used to allocate resources in this TurnResult. diff --git a/src/llm/Provider.tests.zig b/src/llm/Provider.tests.zig index fec28c3..cc29ea6 100644 --- a/src/llm/Provider.tests.zig +++ b/src/llm/Provider.tests.zig @@ -17,16 +17,14 @@ test "Provider.listModels delegates to VTable" { 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 model = llm.types.Model{ .id = "test-model-id", .display_name = "Test Model", }; - const session_config = llm.types.SessionConfig{ - .model = model, - .tools = &.{}, - }; + const session_config = llm.types.SessionConfig{ .model = model }; const input_steps = &[_]llm.types.Step{ .{ .prompt = "hello" }, }; @@ -40,7 +38,7 @@ test "Provider.executeStep delegates to VTable" { try std.testing.expectEqual(@as(usize, 1), mock_impl.execute_step_calls); 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.?[0].prompt); + 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); } @@ -84,7 +82,7 @@ test "Provider.listModels returns custom success and error" { test "Provider.executeStep returns custom success and error" { const allocator = std.testing.allocator; const model = llm.types.Model{ .id = "id", .display_name = "name" }; - const session_config = llm.types.SessionConfig{ .model = model, .tools = &.{} }; + const session_config = llm.types.SessionConfig{ .model = model }; // Test custom success { @@ -99,7 +97,6 @@ test "Provider.executeStep returns custom success and error" { }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); var outcome = try prov.executeStep(allocator, session_config, &.{}, null); defer outcome.result.deinit(); @@ -125,16 +122,14 @@ test "Provider.executeStep returns custom success and error" { 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 model = llm.types.Model{ .id = "test-model-id", .display_name = "Test Model", }; - const session_config = llm.types.SessionConfig{ - .model = model, - .tools = &.{}, - }; + const session_config = llm.types.SessionConfig{ .model = model }; const input_steps = &[_]llm.types.Step{ .{ .prompt = "hello" }, }; @@ -156,14 +151,14 @@ test "Provider.executeStepStreaming delegates to VTable" { try std.testing.expectEqual(@as(usize, 0), mock_impl.execute_step_calls); 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.?[0].prompt); + 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); } test "Provider.executeStepStreaming returns custom success and error" { const allocator = std.testing.allocator; const model = llm.types.Model{ .id = "id", .display_name = "name" }; - const session_config = llm.types.SessionConfig{ .model = model, .tools = &.{} }; + const session_config = llm.types.SessionConfig{ .model = model }; const CallbackState = struct { fn callback(ctx: ?*anyopaque, chunk: llm.types.StreamingChunk) void { @@ -185,7 +180,6 @@ test "Provider.executeStepStreaming returns custom success and error" { }; mock_impl.execute_step_results = &outcomes; - var prov = mock_impl.provider(); var outcome = try prov.executeStepStreaming(allocator, session_config, &.{}, null, CallbackState.callback, null); defer outcome.result.deinit(); diff --git a/src/llm/Provider.zig b/src/llm/Provider.zig index ddb5749..0d6b3b1 100644 --- a/src/llm/Provider.zig +++ b/src/llm/Provider.zig @@ -1,3 +1,5 @@ +//! An interface for an LLM provider. + const std = @import("std"); const Allocator = std.mem.Allocator; const types = @import("types.zig"); @@ -9,7 +11,6 @@ const StreamingCallback = types.StreamingCallback; const StepContinuation = types.StepContinuation; const StepOutcome = types.StepOutcome; -/// An interface for an LLM provider. const Provider = @This(); /// Opaque pointer to the provider implementation's context. diff --git a/src/llm/types.zig b/src/llm/types.zig index 15fac0d..06775c4 100644 --- a/src/llm/types.zig +++ b/src/llm/types.zig @@ -7,7 +7,9 @@ pub const SessionConfig = struct { /// The model to be used for the session. model: Model, /// A list of tools available for the model to use during the session. - tools: []const Tool, + tools: []const Tool = &.{}, + /// The system prompt to be used for the session. + system_prompt: ?[]const u8 = null, }; /// Represents an LLM model provided by the backend. diff --git a/src/main.zig b/src/main.zig index 31818a9..be1a5a3 100644 --- a/src/main.zig +++ b/src/main.zig @@ -1,10 +1,11 @@ const std = @import("std"); const provider = @import("provider"); const llm = @import("llm"); -const agent_pkg = @import("agent"); -const Agent = agent_pkg.Agent; -const Tool = agent_pkg.Tool; -const types = agent_pkg.types; +const agent = @import("agent"); +const Session = agent.Session; +const Tool = agent.Tool; +const types = agent.types; +const acp_pkg = @import("acp"); const coma = @import("coma"); @@ -118,7 +119,7 @@ fn printMarkdown(stream_ctx: *StreamContext, text: []const u8) void { } } -fn streamCallback(ctx: ?*anyopaque, agent_chunk: agent_pkg.types.StreamingChunk) void { +fn streamCallback(ctx: ?*anyopaque, agent_chunk: types.StreamingChunk) void { const stream_ctx: *StreamContext = @ptrCast(@alignCast(ctx)); switch (agent_chunk) { .model_chunk => |chunk| { @@ -252,15 +253,19 @@ fn executeTypescript(allocator: std.mem.Allocator, io: std.Io, code: []const u8) return result.stdout; } -fn getWeather(allocator: std.mem.Allocator, zip_code: i64) ![]const u8 { +fn getWeather(allocator: std.mem.Allocator, zip_code: i64, ctx: *WeatherToolCtx) ![]const u8 { const result_str = if (zip_code == 7302) - try allocator.dupe(u8, "Weather report for 07302: Sunny, 72°F, Humidity 50%, Wind 5 mph") + 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}); return result_str; } +const WeatherToolCtx = struct { + weather_str: []const u8, +}; + /// The main entry point of the application. /// Currently used for testing. pub fn main(init: std.process.Init) !void { @@ -293,6 +298,8 @@ 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", @@ -306,7 +313,7 @@ pub fn main(init: std.process.Init) !void { }, }, }, executeTypescript), - Tool.init(.{ + Tool.initWithContext(.{ .name = "get_weather", .description = "Get the current weather for a given zip code.", .parameters = &.{ @@ -317,16 +324,43 @@ pub fn main(init: std.process.Init) !void { .description = "The 5-digit zip code to get the weather for.", }, }, - }, getWeather), + }, getWeather, &weather_ctx), + agent.Tool.BuiltIn.Todo, }; - const agent_config: types.AgentConfig = .{ - .model = selected_model.?, - .tools = tools, - }; + 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 args = try init.minimal.args.toSlice(allocator); + defer allocator.free(args); + + var run_acp = false; + for (args) |arg| { + if (std.mem.eql(u8, arg, "--acp") or std.mem.eql(u8, arg, "acp")) { + run_acp = true; + break; + } + } + + if (run_acp) { + var stdin_buffer: [1024]u8 = undefined; + var stdin_reader = std.Io.File.stdin().reader(io, &stdin_buffer); + var stdout_buffer: [1024]u8 = undefined; + var stdout_writer = std.Io.File.stdout().writer(io, &stdout_buffer); + + const acp_config = acp_pkg.Server.Config{ + .provider = gemini_client.provider(), + .default_session_config = session_config, + }; + var server = acp_pkg.Server.init(allocator, io, &stdin_reader.interface, &stdout_writer.interface); + defer server.deinit(); + + std.debug.print("ACP Server: starting standard input/output loop...\n", .{}); + try server.run(acp_config); + return; + } - var agent: Agent = try .init(allocator, io, client, agent_config); - defer agent.deinit(); + var session: Session = try .init(allocator, io, client, session_config); + defer session.deinit(); std.debug.print( \\{s}============================================================================ @@ -349,7 +383,7 @@ pub fn main(init: std.process.Init) !void { const turn = types.Turn{ .prompt = user_input }; var stream_ctx = StreamContext{ .allocator = allocator }; - var result = agent.executeTurnStreaming(turn, streamCallback, &stream_ctx) catch |err| { + var result = session.executeTurnStreaming(turn, streamCallback, &stream_ctx) catch |err| { std.debug.print("Error during execution: {}\n", .{err}); continue; }; diff --git a/src/provider/google/api.zig b/src/provider/google/api.zig index dfe64ee..a4778f7 100644 --- a/src/provider/google/api.zig +++ b/src/provider/google/api.zig @@ -6,7 +6,13 @@ const llm = @import("llm"); /// /// Useful when implementing a custom json stringifier that writes additional fields before the object fields. /// But still needs the object fields. +/// TODO(razza): Move to a general location? This is reused in acp/shared_api fn jsonStringifyFields(object: anytype, jw: anytype) !void { + const info = @typeInfo(@TypeOf(object)); + if (info != .@"struct") { + @compileError("jsonStringifyFields only supports struct types"); + } + inline for (std.meta.fields(@TypeOf(object))) |field| { const value = @field(object, field.name); if (@typeInfo(field.type) == .optional and value == null) { @@ -214,6 +220,7 @@ pub const CreateInteractionRequest = struct { model: []const u8, input: []const CreateInteractionRequest.Step, + system_instruction: ?[]const u8, previous_interaction_id: ?[]const u8 = null, generation_config: GenerationConfig = .{}, tools: []const Tool, diff --git a/src/provider/google/provider.zig b/src/provider/google/provider.zig index ca1beaf..f211779 100644 --- a/src/provider/google/provider.zig +++ b/src/provider/google/provider.zig @@ -134,6 +134,7 @@ fn MakeProvider(comptime ClientType: type) type { .tools = tools, .previous_interaction_id = if (previous_gemini_step) |step| step.interaction_id else null, .stream = stream, + .system_instruction = session_config.system_prompt, }; return .{ .uri = uri, .payload = request_payload }; @@ -226,7 +227,7 @@ fn MakeProvider(comptime ClientType: type) type { _ = reader.streamDelimiter(&line_writer.writer, '\n') catch |err| { if (err == error.EndOfStream) break else return ProviderError.BadResponse; }; - _ = reader.toss(1); + reader.toss(1); const line = std.mem.trimEnd(u8, line_writer.written(), "\r"); if (!std.mem.startsWith(u8, line, "data: ") or @@ -503,7 +504,6 @@ test "Gemini.executeStep success" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; const input = &[_]llm.types.Step{ .{ .prompt = "Hello" }, @@ -599,7 +599,6 @@ test "Gemini.executeStep with previous step" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; // First call @@ -724,7 +723,6 @@ test "Gemini.executeStep HTTP failure" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; const input = &[_]llm.types.Step{ .{ .prompt = "Hello" }, @@ -736,7 +734,7 @@ test "Gemini.executeStep HTTP failure" { test "Gemini.executeStepStreaming success" { const allocator = std.testing.allocator; - const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; + const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; const response_body = \\event: interaction.created \\data: {"event_type":"interaction.created","interaction":{"id":"interaction_streaming_123"}} @@ -803,7 +801,6 @@ test "Gemini.executeStepStreaming success" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; const input = &[_]llm.types.Step{ .{ .prompt = "Hello" }, @@ -891,7 +888,7 @@ test "Gemini.executeStepStreaming success" { test "Gemini.executeStepStreaming with CRLF line endings" { const allocator = std.testing.allocator; - const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; + const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; const response_body_lf = \\event: interaction.created \\data: {"event_type":"interaction.created","interaction":{"id":"interaction_streaming_123"}} @@ -966,7 +963,6 @@ test "Gemini.executeStepStreaming with CRLF line endings" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; const input = &[_]llm.types.Step{ .{ .prompt = "Hello" }, @@ -1091,7 +1087,6 @@ test "Gemini.executeStep returns function call" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; var outcome = try p.executeStep(allocator, config, &.{}, null); @@ -1133,7 +1128,7 @@ test "Gemini.executeStepStreaming multiple interaction_created" { .expected_path = "/v1beta/interactions", .expected_query = "key=TEST_API_KEY", .expected_method = .POST, - .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", + .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", .response_status = .ok, .response_body = payload, }, @@ -1144,7 +1139,6 @@ test "Gemini.executeStepStreaming multiple interaction_created" { defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null)); } @@ -1173,7 +1167,7 @@ test "Gemini.executeStepStreaming duplicate step start" { .expected_path = "/v1beta/interactions", .expected_query = "key=TEST_API_KEY", .expected_method = .POST, - .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", + .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", .response_status = .ok, .response_body = payload, }, @@ -1184,7 +1178,6 @@ test "Gemini.executeStepStreaming duplicate step start" { defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null)); } @@ -1212,7 +1205,7 @@ test "Gemini.executeStepStreaming delta for non-existent step index" { .expected_path = "/v1beta/interactions", .expected_query = "key=TEST_API_KEY", .expected_method = .POST, - .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", + .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", .response_status = .ok, .response_body = payload, }, @@ -1223,7 +1216,6 @@ test "Gemini.executeStepStreaming delta for non-existent step index" { defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null)); } @@ -1252,7 +1244,7 @@ test "Gemini.executeStepStreaming mismatched interaction completed ID" { .expected_path = "/v1beta/interactions", .expected_query = "key=TEST_API_KEY", .expected_method = .POST, - .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", + .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", .response_status = .ok, .response_body = payload, }, @@ -1263,7 +1255,6 @@ test "Gemini.executeStepStreaming mismatched interaction completed ID" { defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null)); } @@ -1291,7 +1282,7 @@ test "Gemini.executeStepStreaming fallback interaction ID to unknown" { .expected_path = "/v1beta/interactions", .expected_query = "key=TEST_API_KEY", .expected_method = .POST, - .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", + .expected_payload = "{\"model\":\"gemini-model\",\"input\":[],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}", .response_status = .ok, .response_body = payload, }, @@ -1302,7 +1293,6 @@ test "Gemini.executeStepStreaming fallback interaction ID to unknown" { defer p.deinit(); const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-model", .display_name = "Gemini Model" }, - .tools = &.{}, }; var outcome = try p.executeStepStreaming(allocator, config, &.{}, null, CallbackState.callback, null); defer outcome.result.deinit(); @@ -1340,7 +1330,7 @@ test "StepContinuation.init OOM" { test "Gemini.executeStepStreaming malformed stream payload" { const allocator = std.testing.allocator; - const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; + const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"system_instruction\":null,\"previous_interaction_id\":null,\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":true}"; const response_body = \\event: step.delta \\data: {invalid_json @@ -1372,7 +1362,6 @@ test "Gemini.executeStepStreaming malformed stream payload" { const config = llm.types.SessionConfig{ .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, - .tools = &.{}, }; const input = &[_]llm.types.Step{ .{ .prompt = "Hello" }, @@ -1387,3 +1376,62 @@ test "Gemini.executeStepStreaming malformed stream payload" { try std.testing.expectError(error.BadResponse, p.executeStepStreaming(allocator, config, input, null, CallbackState.callback, null)); } + +test "Gemini.executeStep with system prompt" { + const allocator = std.testing.allocator; + + const expected_payload = "{\"model\":\"gemini-2.0-flash\",\"input\":[{\"type\":\"user_input\",\"content\":[{\"type\":\"text\",\"text\":\"Hello\"}]}],\"system_instruction\":\"You are a helpful assistant.\",\"generation_config\":{\"thinking_summaries\":\"auto\"},\"tools\":[],\"stream\":false}"; + const response_json = + \\{ + \\ "id": "interaction_sp", + \\ "steps": [ + \\ { + \\ "type": "model_output", + \\ "content": [ + \\ { + \\ "type": "text", + \\ "text": "Hello user!" + \\ } + \\ ] + \\ } + \\ ] + \\} + ; + + var call_counts = [_]usize{0}; + const mock_client = testing.MockHttpClient{ + .allocator = allocator, + .expectations = &.{ + .{ + .expected_scheme = "https", + .expected_host = "generativelanguage.googleapis.com", + .expected_path = "/v1beta/interactions", + .expected_query = "key=TEST_API_KEY", + .expected_method = .POST, + .expected_payload = expected_payload, + .response_status = .ok, + .response_body = response_json, + }, + }, + .sequential = false, + .call_counts = &call_counts, + }; + + var prov = try MakeProvider(testing.MockHttpClient).init(allocator, mock_client, "TEST_API_KEY"); + var p = prov.provider(); + defer p.deinit(); + + const config = llm.types.SessionConfig{ + .model = .{ .id = "gemini-2.0-flash", .display_name = "Gemini 2.0 Flash" }, + .system_prompt = "You are a helpful assistant.", + }; + const input = &[_]llm.types.Step{ + .{ .prompt = "Hello" }, + }; + + var outcome = try p.executeStep(allocator, config, input, null); + defer outcome.result.deinit(); + defer outcome.continuation.deinit(); + + try std.testing.expectEqual(1, call_counts[0]); +} diff --git a/src/provider/json_client.zig b/src/provider/json_client.zig index b299617..ebc6182 100644 --- a/src/provider/json_client.zig +++ b/src/provider/json_client.zig @@ -8,6 +8,10 @@ pub const JsonHttpClient = MakeJsonClient(*std.http.Client); pub const MockJsonClient = MakeJsonClient(testing.MockHttpClient); /// Generates a JSON HTTP client wrapper parameterized by the underlying HTTP client type. +/// +/// The resulting type is as thread-safe as the underlying client type. Therefore, for JsonClients +/// created against `std.http.Client` individual connects and requests are created in a thread-safe +/// manner but each request, itself, is not thread-safe. pub fn MakeJsonClient(comptime ClientType: type) type { return struct { http_client: ClientType, diff --git a/src/testing/MockHttpClient.zig b/src/testing/MockHttpClient.zig index 0f2e86b..abd72c9 100644 --- a/src/testing/MockHttpClient.zig +++ b/src/testing/MockHttpClient.zig @@ -1,6 +1,7 @@ +//! A mock HTTP client (`std.http.Client`) used for testing. + const std = @import("std"); -/// A mock HTTP client used for testing. const MockHttpClient = @This(); /// Represents an expected HTTP request and its corresponding mock response. diff --git a/src/testing/MockProvider.zig b/src/testing/MockProvider.zig index 77256fb..0896b42 100644 --- a/src/testing/MockProvider.zig +++ b/src/testing/MockProvider.zig @@ -1,3 +1,5 @@ +//! A mock implementation of the `llm.Provider` interface for testing. + const std = @import("std"); const Allocator = std.mem.Allocator; const llm = @import("llm"); @@ -11,7 +13,6 @@ const SessionConfig = types.SessionConfig; const Step = types.Step; const StepOutcome = types.StepOutcome; -/// A mock implementation of the `llm.Provider` interface for testing. const MockProvider = @This(); /// Tracks the number of times `listModels` was called. @@ -28,7 +29,7 @@ last_allocator: ?Allocator = null, /// Stores the session configuration from the last `executeStep` call. last_session_config: ?SessionConfig = null, /// Stores the input steps from the last `executeStep` call. -last_input: ?[]const Step = null, +last_input_steps: ?[]const Step = null, /// Stores the previous step from the last `executeStep` call. last_previous_step: ?StepContinuation = null, @@ -41,12 +42,11 @@ execute_step_results_loop: bool = true, /// Optional sequence of streaming chunks to emit during successive `executeStepStreaming` calls. execute_step_streaming_chunks: ?[]const []const types.StreamingChunk = null, - const vtable = Provider.VTable{ .list_models = MockProvider.list_models, .execute_step = MockProvider.execute_step, .execute_step_streaming = MockProvider.execute_step_streaming, - .deinit = MockProvider.deinit, + .deinit = MockProvider.deinitVtable, }; /// Returns the generic `llm.Provider` interface for this mock instance. @@ -92,6 +92,41 @@ fn list_models(ptr: *anyopaque, allocator: Allocator) Provider.ProviderError!Lis }; } +fn clearLastInputSteps(self: *MockProvider) void { + if (self.last_input_steps) |steps| { + if (self.last_allocator) |alloc| { + for (steps) |step| { + switch (step) { + .prompt => |p| alloc.free(p), + .tool_result => {}, + } + } + alloc.free(steps); + } + self.last_input_steps = null; + } +} + +fn recordInputSteps(self: *MockProvider, allocator: Allocator, input: []const Step) void { + self.clearLastInputSteps(); + var list: std.ArrayList(Step) = .empty; + for (input) |step| { + switch (step) { + .prompt => |p| { + const p_copy = allocator.dupe(u8, p) catch return; + list.append(allocator, .{ .prompt = p_copy }) catch { + allocator.free(p_copy); + return; + }; + }, + .tool_result => |tr| { + list.append(allocator, .{ .tool_result = tr }) catch return; + }, + } + } + self.last_input_steps = list.toOwnedSlice(allocator) catch return; +} + /// Mock implementation of `executeStep`. fn execute_step( ptr: *anyopaque, @@ -104,7 +139,7 @@ fn execute_step( self.execute_step_calls += 1; self.last_allocator = allocator; self.last_session_config = session_config; - self.last_input = input; + self.recordInputSteps(allocator, input); self.last_previous_step = previous_step; if (self.execute_step_results) |results| { if (results.len == 0) { @@ -145,7 +180,7 @@ fn execute_step_streaming( self.execute_step_streaming_calls += 1; self.last_allocator = allocator; self.last_session_config = session_config; - self.last_input = input; + self.recordInputSteps(allocator, input); self.last_previous_step = previous_step; if (self.execute_step_streaming_chunks) |chunks_list| { if (chunks_list.len > 0) { @@ -186,9 +221,14 @@ fn execute_step_streaming( } /// Mock implementation of `deinit`. -fn deinit(ptr: *anyopaque) void { - const self: *MockProvider = @ptrCast(@alignCast(ptr)); +fn deinitVtable(ctx: *anyopaque) void { + const self: *MockProvider = @ptrCast(@alignCast(ctx)); + deinit(self); +} + +pub fn deinit(self: *MockProvider) void { self.deinit_calls += 1; + self.clearLastInputSteps(); } /// Helper constructor to create a mock `StepResult` with standard mock vtable.