tool call resiliency

1 files changed, 131 insertions, 29 deletions

diff --git a/libpanto/src/session_manager.zig b/libpanto/src/session_manager.zig
index 8b503f2..9e8e2c8 100644
--- a/libpanto/src/session_manager.zig
+++ b/libpanto/src/session_manager.zig
@@ -355,10 +355,11 @@ pub const SessionManager = struct {
             try truncateFileTo(io, path_copy, valid_bytes);
         }
 
-        // Run format migration. Currently a no-op for version 1, but the
-        // hook exists so future versions can rewrite the file.
         var migrated_header = header;
-        const did_migrate = migrate(allocator, &migrated_header, &entries);
+        var did_migrate = migrate(allocator, &migrated_header, &entries);
+        if (elideDanglingToolUses(allocator, &entries)) {
+            did_migrate = true;
+        }
         // We didn't reassign by_id during migration; rebuild if needed.
         if (did_migrate) {
             by_id.clearRetainingCapacity();
@@ -567,16 +568,11 @@ pub const SessionManager = struct {
 
     // ---------- Persistence ----------
 
-    /// Write the header + all currently-buffered entries + `final_entry`
+    /// Write the header + all currently-buffered entries + `new_entries`
     /// to the file as a single batch. Creates the directory and file.
-    /// Called only when `flushed = false` and we just got an assistant
-    /// message.
-    fn flushBuffered(self: *SessionManager, final_entry: SessionEntry) !void {
-        // Ensure the directory exists.
+    fn flushBufferedMany(self: *SessionManager, new_entries: []const SessionEntry) !void {
         try mkdirP(self.io, self.session_dir);
 
-        // Open (create exclusively isn't required; if a stale file with the
-        // same UUIDv7 existed we'd just overwrite — UUIDv7s are unique).
         const file = try Io.Dir.cwd().createFile(self.io, self.session_file, .{
             .truncate = true,
             .read = false,
@@ -584,7 +580,6 @@ pub const SessionManager = struct {
         defer file.close(self.io);
 
         var offset: u64 = 0;
-        // Header
         const header_line = try session_mod.serializeHeader(self.allocator, self.header);
         defer self.allocator.free(header_line);
         try file.writePositionalAll(self.io, header_line, offset);
@@ -592,7 +587,6 @@ pub const SessionManager = struct {
         try file.writePositionalAll(self.io, "\n", offset);
         offset += 1;
 
-        // Buffered entries.
         for (self.entries.items) |e| {
             const line = try session_mod.serializeEntry(self.allocator, e);
             defer self.allocator.free(line);
@@ -602,38 +596,105 @@ pub const SessionManager = struct {
             offset += 1;
         }
 
-        // Final (the assistant message we just got).
-        const final_line = try session_mod.serializeEntry(self.allocator, final_entry);
-        defer self.allocator.free(final_line);
-        try file.writePositionalAll(self.io, final_line, offset);
-        offset += final_line.len;
-        try file.writePositionalAll(self.io, "\n", offset);
-        offset += 1;
+        for (new_entries) |entry| {
+            const line = try session_mod.serializeEntry(self.allocator, entry);
+            defer self.allocator.free(line);
+            try file.writePositionalAll(self.io, line, offset);
+            offset += line.len;
+            try file.writePositionalAll(self.io, "\n", offset);
+            offset += 1;
+        }
 
-        // Best-effort fsync — the directory itself doesn't need syncing
-        // for our purposes (we don't unlink/rename); the file body does.
         file.sync(self.io) catch {};
-
         self.flushed = true;
         self.written_bytes = offset;
     }
 
+    fn flushBuffered(self: *SessionManager, final_entry: SessionEntry) !void {
+        try self.flushBufferedMany(&.{final_entry});
+    }
+
     /// Append a single line for `entry` to the open session file. Caller
     /// must have already verified `flushed`.
     fn persistEntry(self: *SessionManager, entry: SessionEntry) !void {
+        try self.persistEntries(&.{entry});
+    }
+
+    pub fn appendMessagesAtomic(
+        self: *SessionManager,
+        messages: []DiskMessage,
+        providers: []const ?[]const u8,
+        models: []const ?[]const u8,
+    ) !void {
+        std.debug.assert(messages.len == providers.len);
+        std.debug.assert(messages.len == models.len);
+        if (messages.len == 0) return;
+
+        const base_len = self.entries.items.len;
+        try self.entries.ensureUnusedCapacity(self.allocator, messages.len);
+        try self.by_id.ensureUnusedCapacity(@intCast(messages.len));
+
+        var entries = try self.allocator.alloc(SessionEntry, messages.len);
+        defer self.allocator.free(entries);
+        var built: usize = 0;
+        errdefer {
+            for (entries[0..built]) |*e| e.deinit(self.allocator);
+        }
+
+        var prev_leaf = self.leaf_id;
+        for (messages, 0..) |msg, i| {
+            const msg_local = msg;
+            const id_buf = try self.newEntryId();
+            errdefer self.allocator.free(id_buf);
+            const timestamp = try isoTimestamp(self.allocator, self.io);
+            errdefer self.allocator.free(timestamp);
+            const parent_id_copy: ?[]const u8 = if (prev_leaf) |l| try self.allocator.dupe(u8, l) else null;
+            errdefer if (parent_id_copy) |p| self.allocator.free(p);
+            const provider_copy: ?[]const u8 = if (providers[i]) |p| try self.allocator.dupe(u8, p) else null;
+            errdefer if (provider_copy) |p| self.allocator.free(p);
+            const model_copy: ?[]const u8 = if (models[i]) |m| try self.allocator.dupe(u8, m) else null;
+            errdefer if (model_copy) |m| self.allocator.free(m);
+            entries[i] = .{ .message = .{
+                .base = .{ .id = id_buf, .parent_id = parent_id_copy, .timestamp = timestamp },
+                .provider = provider_copy,
+                .model = model_copy,
+                .message = msg_local,
+            } };
+            built += 1;
+            prev_leaf = entries[i].base().id;
+        }
+
+        if (self.flushed) {
+            try self.persistEntries(entries);
+        } else {
+            try self.flushBufferedMany(entries);
+        }
+
+        for (entries, 0..) |entry, i| {
+            self.entries.appendAssumeCapacity(entry);
+            self.by_id.putAssumeCapacity(entry.base().id, base_len + i);
+        }
+        self.leaf_id = entries[entries.len - 1].base().id;
+        built = 0;
+    }
+
+    fn persistEntries(self: *SessionManager, entries: []const SessionEntry) !void {
         const file = try Io.Dir.cwd().openFile(self.io, self.session_file, .{
             .mode = .write_only,
         });
         defer file.close(self.io);
 
-        const line = try session_mod.serializeEntry(self.allocator, entry);
-        defer self.allocator.free(line);
-
-        try file.writePositionalAll(self.io, line, self.written_bytes);
-        self.written_bytes += line.len;
-        try file.writePositionalAll(self.io, "\n", self.written_bytes);
-        self.written_bytes += 1;
+        var offset = self.written_bytes;
+        for (entries) |entry| {
+            const line = try session_mod.serializeEntry(self.allocator, entry);
+            defer self.allocator.free(line);
+            try file.writePositionalAll(self.io, line, offset);
+            offset += line.len;
+            try file.writePositionalAll(self.io, "\n", offset);
+            offset += 1;
+        }
         file.sync(self.io) catch {};
+        self.written_bytes = offset;
     }
 
     // =============================================================================
@@ -711,6 +772,47 @@ fn migrate(
     return false;
 }
 
+fn elideDanglingToolUses(allocator: Allocator, entries: *std.ArrayList(SessionEntry)) bool {
+    var needed: std.StringHashMap(void) = .init(allocator);
+    defer needed.deinit();
+    var changed = false;
+
+    var i = entries.items.len;
+    while (i > 0) {
+        i -= 1;
+        const entry = &entries.items[i];
+        if (entry.* != .message) continue;
+        const msg = &entry.message.message;
+
+        if (msg.role == .user) {
+            for (msg.content) |block| {
+                if (block == .tool_result) {
+                    needed.put(block.tool_result.tool_use_id, {}) catch {};
+                }
+            }
+            continue;
+        }
+
+        if (msg.role != .assistant) continue;
+        var kept: std.ArrayList(DiskContentBlock) = .empty;
+        defer kept.deinit(allocator);
+        var removed = false;
+        for (msg.content) |block| {
+            if (block == .tool_use and !needed.contains(block.tool_use.id)) {
+                block.deinit(allocator);
+                removed = true;
+                continue;
+            }
+            kept.append(allocator, block) catch unreachable;
+        }
+        if (!removed) continue;
+        allocator.free(msg.content);
+        msg.content = kept.toOwnedSlice(allocator) catch unreachable;
+        changed = true;
+    }
+    return changed;
+}
+
 // =============================================================================
 // File utilities
 // =============================================================================