diff --git a/src/agent/Agent.zig b/src/agent/Agent.zig index 308538f2..447351b3 100644 --- a/src/agent/Agent.zig +++ b/src/agent/Agent.zig @@ -672,24 +672,31 @@ fn pruneMessages(self: *Self) void { const tail_start = zenai.provider.safeTruncationStart(msgs, msgs.len - prune_keep) orelse return; + // Dupe the kept tail into a scratch slice in the new arena first. Only + // mutate self.messages once every dupe has succeeded — otherwise a + // partial failure would leave self.messages.items[1..] pointing into + // the freed `new_arena`, since dupeMessage already wrote into it. var new_arena: std.heap.ArenaAllocator = .init(self.allocator); - const na = new_arena.allocator(); + const duped = dupeMessages(new_arena.allocator(), msgs[tail_start..]) orelse { + new_arena.deinit(); + return; + }; - // System prompt content points outside the arena, so only the tail needs copying. - var i: usize = 0; - for (msgs[tail_start..]) |msg| { - msgs[1 + i] = dupeMessage(na, msg) orelse { - new_arena.deinit(); - return; - }; - i += 1; - } - - self.messages.shrinkRetainingCapacity(1 + i); + // System prompt at index 0 lives outside the arena and is preserved. + @memcpy(self.messages.items[1..][0..duped.len], duped); + self.messages.shrinkRetainingCapacity(1 + duped.len); self.message_arena.deinit(); self.message_arena = new_arena; } +fn dupeMessages(arena: std.mem.Allocator, msgs: []const zenai.provider.Message) ?[]zenai.provider.Message { + const out = arena.alloc(zenai.provider.Message, msgs.len) catch return null; + for (msgs, 0..) |msg, i| { + out[i] = dupeMessage(arena, msg) orelse return null; + } + return out; +} + fn dupeMessage(alloc: std.mem.Allocator, msg: zenai.provider.Message) ?zenai.provider.Message { return .{ .role = msg.role, @@ -1253,3 +1260,36 @@ test "formatReplacement: multiple commands produce multi-line replacement" { replacement.new_text, ); } + +test "dupeMessages: happy path" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + + const src = [_]zenai.provider.Message{ + .{ .role = .user, .content = "hello" }, + .{ .role = .assistant, .content = "world" }, + }; + + const out = dupeMessages(arena.allocator(), &src) orelse return error.UnexpectedNull; + try std.testing.expectEqual(@as(usize, 2), out.len); + try std.testing.expectEqualStrings("hello", out[0].content.?); + try std.testing.expectEqualStrings("world", out[1].content.?); + try std.testing.expect(out[0].content.?.ptr != src[0].content.?.ptr); +} + +test "dupeMessages: returns null on mid-iteration alloc failure" { + // The contract pruneMessages depends on: on any partial failure, + // dupeMessages returns null without mutating its inputs. pruneMessages + // can then deinit the scratch arena and leave self.messages untouched. + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + + var failing = std.testing.FailingAllocator.init(arena.allocator(), .{ .fail_index = 2 }); + const src = [_]zenai.provider.Message{ + .{ .role = .user, .content = "hello" }, + .{ .role = .assistant, .content = "world" }, + .{ .role = .user, .content = "third" }, + }; + try std.testing.expect(dupeMessages(failing.allocator(), &src) == null); + try std.testing.expect(failing.has_induced_failure); +}