diff options
Diffstat (limited to 'src/main.zig')
| -rw-r--r-- | src/main.zig | 369 |
1 files changed, 361 insertions, 8 deletions
diff --git a/src/main.zig b/src/main.zig index 6654c95..3bc24e3 100644 --- a/src/main.zig +++ b/src/main.zig @@ -8,6 +8,8 @@ const panto_home = @import("panto_home.zig"); const luarocks_runtime = @import("luarocks_runtime.zig"); const self_exe = @import("self_exe.zig"); const subcommand = @import("subcommand.zig"); +const session_paths = @import("session_paths.zig"); +const models_toml = @import("models_toml.zig"); // Shorthand alias for the Lua C API. The bridge module owns the actual // `@cImport`; we re-use it here so the smoke check uses identical types. @@ -31,6 +33,7 @@ test { _ = luarocks_runtime; _ = self_exe; _ = subcommand; + _ = models_toml; } const Receiver = panto.provider.Receiver; @@ -40,14 +43,37 @@ const MessageRole = panto.conversation.MessageRole; /// Receiver that prints streaming deltas to stdout. Thinking blocks are /// dimmed with ANSI escape codes; text blocks render plain. +/// +/// Also captures one `?Usage` per assistant response during a turn. In +/// a tool-using turn the agent loop drives multiple streamStep calls; +/// each one ends with `onMessageComplete(msg, usage)`. We append the +/// usage — including `null` when the wire didn't deliver any — to +/// `per_message_usage` in order so `persistTurn` can pair each captured +/// usage with its corresponding assistant message. const CLIReceiver = struct { stdout: *std.Io.Writer, file: *std.Io.File.Writer, + allocator: std.mem.Allocator, + + /// One slot per assistant message completed during the current + /// turn, in completion order. `null` means the provider did not + /// report usage on the wire for that message (typical of + /// OpenAI-compatible proxies that ignore `stream_options.include_usage`). + per_message_usage: std.ArrayList(?panto.session_manager.Usage) = .empty, pub fn receiver(self: *CLIReceiver) Receiver { return .{ .ptr = self, .vtable = &vtable }; } + /// Reset usage state at the start of each turn. + pub fn beginTurn(self: *CLIReceiver) void { + self.per_message_usage.clearRetainingCapacity(); + } + + pub fn deinit(self: *CLIReceiver) void { + self.per_message_usage.deinit(self.allocator); + } + const vtable: ReceiverVTable = .{ .onMessageStart = onMessageStart, .onBlockStart = onBlockStart, @@ -121,9 +147,19 @@ const CLIReceiver = struct { try self.file.flush(); } - fn onMessageComplete(ptr: *anyopaque, message: panto.conversation.Message) anyerror!void { - _ = message; + fn onMessageComplete( + ptr: *anyopaque, + message: panto.conversation.Message, + usage: ?panto.session_manager.Usage, + ) anyerror!void { const self: *CLIReceiver = @ptrCast(@alignCast(ptr)); + // Only assistant messages come through streaming. The receiver + // contract says onMessageComplete fires exactly once per + // streamStep, with role=.assistant; record the usage slot in + // turn order regardless of whether the wire actually had usage. + if (message.role == .assistant) { + try self.per_message_usage.append(self.allocator, usage); + } try self.stdout.writeAll("\n"); try self.file.flush(); } @@ -261,6 +297,10 @@ pub fn main(init: std.process.Init) !void { const config = try loadConfig(init.environ_map); + // Parse the agent-mode flags. Currently only `--resume [<id>]`. + const cli_flags = try parseAgentFlags(alloc, init.minimal.args); + defer cli_flags.deinit(alloc); + var stdout_buffer: [4096]u8 = undefined; var stdout_file = std.Io.File.stdout().writer(io, &stdout_buffer); const stdout = &stdout_file.interface; @@ -269,9 +309,62 @@ pub fn main(init: std.process.Init) !void { var stdin_file = std.Io.File.stdin().reader(io, &stdin_buffer); const stdin = &stdin_file.interface; + // Resolve where this project's sessions live. + var cwd_buf: [std.fs.max_path_bytes]u8 = undefined; + const cwd_n = try std.process.currentPath(io, &cwd_buf); + const cwd = cwd_buf[0..cwd_n]; + const session_dir = try session_paths.sessionDirForCwd(alloc, init.environ_map, cwd); + defer alloc.free(session_dir); + + // Load the user's models.toml — a missing file is fine (empty + // registry). Cost lookups against an empty registry return null, + // and the display layer will format that as "unknown." + const models_toml_path = try models_toml.configPath(alloc, init.environ_map); + defer alloc.free(models_toml_path); + var pricing_registry = try models_toml.loadFromPath(alloc, io, models_toml_path); + defer pricing_registry.deinit(); + std.log.debug("models.toml: {d} entries from {s}", .{ pricing_registry.count(), models_toml_path }); + + const banner_model_initial: []const u8 = switch (config) { + inline else => |c| c.model, + }; + const banner_provider_initial: []const u8 = @tagName(config); + + // Create or resume the session. Resume failures (missing/ambiguous id) + // are user errors — print a tidy message and exit 1 rather than + // printing a Zig stack trace. + var session_mgr = openSession( + alloc, + io, + session_dir, + cwd, + cli_flags, + stdout, + &stdout_file, + ) catch |err| switch (err) { + error.SessionNotFound, error.AmbiguousSessionId => std.process.exit(1), + else => return err, + }; + defer session_mgr.deinit(); + var conv = panto.conversation.Conversation.init(alloc); defer conv.deinit(); - try conv.addSystemMessage("You are a helpful assistant."); + + if (session_mgr.getEntries().len > 0) { + // Resumed an existing session — rebuild the conversation from the + // log. The system prompt is part of the log. + conv.deinit(); + conv = try session_mgr.rebuildConversation(); + try stdout.print( + "resumed session {s} ({d} entries)\n", + .{ session_mgr.getSessionId()[0..@min(8, session_mgr.getSessionId().len)], session_mgr.getEntries().len }, + ); + } else { + // Fresh session — install the default system prompt and record it. + const system_text = "You are a helpful assistant."; + try conv.addSystemMessage(system_text); + try appendSystemToSession(alloc, &session_mgr, system_text); + } const prov = try panto.provider.Provider.init(alloc, io, config); var agent = panto.agent.Agent.init(alloc, io, prov); @@ -325,17 +418,22 @@ pub fn main(init: std.process.Init) !void { try agent.registerToolSource(rt.toolSource()); } - const banner_model: []const u8 = switch (config) { - inline else => |c| c.model, - }; const banner_base: []const u8 = switch (config) { inline else => |c| c.base_url, }; - try stdout.print("panto — {s}: {s} @ {s}\n", .{ @tagName(config), banner_model, banner_base }); + try stdout.print( + "panto — {s}: {s} @ {s}\n", + .{ banner_provider_initial, banner_model_initial, banner_base }, + ); try stdout.print("> ", .{}); try stdout_file.flush(); - var cli_recv = CLIReceiver{ .stdout = stdout, .file = &stdout_file }; + var cli_recv = CLIReceiver{ + .stdout = stdout, + .file = &stdout_file, + .allocator = alloc, + }; + defer cli_recv.deinit(); var recv = cli_recv.receiver(); while (true) { @@ -356,12 +454,267 @@ pub fn main(init: std.process.Init) !void { } try conv.addUserMessage(line); + try appendUserPromptToSession( + alloc, + &session_mgr, + line, + banner_provider_initial, + banner_model_initial, + ); + const entries_before_step = conv.messages.items.len; + + cli_recv.beginTurn(); agent.runStep(&conv, &recv) catch |err| { try stdout.print("\n[error: {s}]\n", .{@errorName(err)}); }; + // Persist whatever new entries the agent produced this turn. The + // agent loop may have appended: + // - assistant message(s) (one per provider response) + // - user messages containing ToolResult blocks (one per tool round) + // Each assistant message gets paired with the Usage that the + // receiver captured at its onMessageComplete time (or null if + // the provider didn't emit usage that round). + try persistTurn( + alloc, + &session_mgr, + &conv, + entries_before_step, + banner_provider_initial, + banner_model_initial, + cli_recv.per_message_usage.items, + ); + try stdout.writeAll("\n> "); try stdout_file.flush(); } } + +// ----------------------------------------------------------------------------- +// CLI flag parsing +// ----------------------------------------------------------------------------- + +const AgentFlags = struct { + /// `--resume` without an id: resume the most recent session. + /// `--resume <id>`: resume the session whose id has this prefix. + /// Not present: start a new session. + resume_kind: ResumeKind = .none, + resume_id: ?[]const u8 = null, // owned + + pub fn deinit(self: AgentFlags, alloc: std.mem.Allocator) void { + if (self.resume_id) |id| alloc.free(id); + } +}; + +const ResumeKind = enum { none, most_recent, by_id }; + +fn parseAgentFlags(alloc: std.mem.Allocator, args: std.process.Args) !AgentFlags { + var flags: AgentFlags = .{}; + errdefer flags.deinit(alloc); + + var it = args.iterate(); + defer it.deinit(); + _ = it.next(); // argv[0] + + while (it.next()) |a| { + if (std.mem.eql(u8, a, "--resume")) { + // Peek at the next arg. If it exists and doesn't start with `-`, + // treat it as a session id (or prefix). + const next = it.next(); + if (next) |id| { + if (id.len > 0 and id[0] != '-') { + flags.resume_kind = .by_id; + flags.resume_id = try alloc.dupe(u8, id); + continue; + } else { + // Not an id; rewind by treating it as a separate flag. + // The Args API doesn't support rewind, so handle inline. + flags.resume_kind = .most_recent; + if (std.mem.eql(u8, id, "--resume")) { + // back-to-back --resume; second resets, fine. + continue; + } + // Otherwise, fall through and process this token as a flag. + // (Currently we don't have other flags; ignore unknowns.) + std.log.warn("panto: ignoring unknown argument '{s}'", .{id}); + continue; + } + } + flags.resume_kind = .most_recent; + continue; + } + // Future agent-mode flags would land here. Unknown args are tolerated + // (the user might be passing something we don't recognize yet). + } + return flags; +} + +// ----------------------------------------------------------------------------- +// Session bootstrap +// ----------------------------------------------------------------------------- + +fn openSession( + alloc: std.mem.Allocator, + io: std.Io, + session_dir: []const u8, + cwd: []const u8, + flags: AgentFlags, + stdout: *std.Io.Writer, + stdout_file: *std.Io.File.Writer, +) !panto.session_manager.SessionManager { + switch (flags.resume_kind) { + .none => return try panto.session_manager.SessionManager.init( + alloc, + io, + session_dir, + cwd, + ), + .most_recent => { + const path_opt = panto.session_manager.findMostRecentSession(alloc, io, session_dir) catch null; + if (path_opt) |path| { + defer alloc.free(path); + return try panto.session_manager.SessionManager.open(alloc, io, path); + } + try stdout.print("no sessions to resume; starting fresh.\n", .{}); + try stdout_file.flush(); + return try panto.session_manager.SessionManager.init( + alloc, + io, + session_dir, + cwd, + ); + }, + .by_id => { + const id = flags.resume_id.?; + const path = panto.session_manager.resolveSessionId(alloc, io, session_dir, id) catch |err| switch (err) { + error.SessionNotFound => { + try stdout.print("error: no session matching '{s}' in {s}\n", .{ id, session_dir }); + try stdout_file.flush(); + return err; + }, + error.AmbiguousSessionId => { + try stdout.print("error: session id '{s}' is ambiguous\n", .{id}); + try stdout_file.flush(); + return err; + }, + else => return err, + }; + defer alloc.free(path); + return try panto.session_manager.SessionManager.open(alloc, io, path); + }, + } +} + +// ----------------------------------------------------------------------------- +// Session append helpers — bridge in-memory ContentBlocks to on-disk entries. +// ----------------------------------------------------------------------------- + +fn appendSystemToSession( + alloc: std.mem.Allocator, + mgr: *panto.session_manager.SessionManager, + text: []const u8, +) !void { + const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, 1); + blocks[0] = .{ .text = .{ .text = try alloc.dupe(u8, text) } }; + _ = try mgr.appendMessage( + .{ .role = .system, .content = blocks }, + null, + null, + ); +} + +fn appendUserPromptToSession( + alloc: std.mem.Allocator, + mgr: *panto.session_manager.SessionManager, + text: []const u8, + provider: []const u8, + model: []const u8, +) !void { + const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, 1); + blocks[0] = .{ .text = .{ .text = try alloc.dupe(u8, text) } }; + _ = try mgr.appendMessage( + .{ .role = .user, .content = blocks }, + provider, + model, + ); +} + +/// After the agent loop has driven a turn to completion, persist every +/// new in-memory message at index `>= start_index` to the session log. +/// +/// Each in-memory message is mapped to disk: +/// - assistant → assistant entry with provider/model/stop_reason metadata +/// and Usage (from `per_message_usage`, one slot per assistant message in +/// completion order). +/// - user (with ToolResult blocks) → user entry stamped with provider/model. +/// +/// `stop_reason` is recorded as `"stop"` for now; surfacing the real wire +/// value requires plumbing it through the Receiver vtable (future work, +/// separate from token plumbing). +fn persistTurn( + alloc: std.mem.Allocator, + mgr: *panto.session_manager.SessionManager, + conv: *const panto.conversation.Conversation, + start_index: usize, + provider: []const u8, + model: []const u8, + per_message_usage: []const ?panto.session_manager.Usage, +) !void { + var i = start_index; + var assistant_seen: usize = 0; + while (i < conv.messages.items.len) : (i += 1) { + const msg = conv.messages.items[i]; + const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, msg.content.items.len); + var allocated: usize = 0; + errdefer { + for (blocks[0..allocated]) |b| b.deinit(alloc); + alloc.free(blocks); + } + for (msg.content.items) |block| { + blocks[allocated] = try panto.session.contentBlockToDisk(alloc, block); + allocated += 1; + } + switch (msg.role) { + .system => { + // Mid-turn system messages aren't a thing in pantograph today. + // Treat as harmless and persist verbatim, sans stamps. + _ = try mgr.appendMessage( + .{ .role = .system, .content = blocks }, + null, + null, + ); + }, + .user => { + _ = try mgr.appendMessage( + .{ .role = .user, .content = blocks }, + provider, + model, + ); + }, + .assistant => { + // Pair this assistant message with its usage, if the + // receiver captured one for this position. (A turn + // ending in a stream error can have fewer usage entries + // than assistant messages; treat as null and move on.) + const usage: ?panto.session_manager.Usage = if (assistant_seen < per_message_usage.len) + per_message_usage[assistant_seen] + else + null; + assistant_seen += 1; + _ = try mgr.appendMessage( + .{ + .role = .assistant, + .content = blocks, + .provider = try alloc.dupe(u8, provider), + .model = try alloc.dupe(u8, model), + .stop_reason = try alloc.dupe(u8, "stop"), + .usage = usage, + }, + null, + null, + ); + }, + } + } +} |