libpanto fixes

- ensuring stream.next() sequences are interruptible
- Agent.run() takes any user message, not just chat text

2 files changed, 282 insertions, 73 deletions

diff --git a/libpanto/src/agent.zig b/libpanto/src/agent.zig
index ed1600f..e689924 100644
--- a/libpanto/src/agent.zig
+++ b/libpanto/src/agent.zig
@@ -200,11 +200,14 @@ fn toolErrorResult(
     return tool_mod.ResultParts.fromTextOwned(allocator, msg);
 }
 
-/// The user's submission that opens a turn. A struct (not a bare slice) so
-/// it can grow to carry file/image attachments alongside the chat text
-/// without changing `Agent.run`'s signature.
+/// The user's submission that opens a turn: an ordered list of content
+/// blocks (text, tool results, images, …), the same shape used to rebuild a
+/// persisted user message. Ownership of `blocks` transfers to the agent's
+/// conversation on `run`; the caller must not deinit them afterwards. A plain
+/// chat turn is one `.Text` block; a turn that resumes after the embedder
+/// handled tool calls is one or more `.ToolResult` blocks.
 pub const UserMessage = struct {
-    text: []const u8,
+    blocks: []const conversation.ContentBlock,
 };
 
 /// Outcome of a compaction attempt.
@@ -259,6 +262,14 @@ pub const Agent = struct {
     /// touched from the single agent-loop thread (retries are serial), so
     /// no synchronization is needed.
     _retry_prng: ?std.Random.DefaultPrng = null,
+    /// High-water mark of messages durably handed to the store: the count of
+    /// conversation messages already persisted. `flushPersist` appends only
+    /// `[_persisted_through .. coherent_end)` and advances this, so the store
+    /// stays current at every `Stream.next()` boundary without re-appending.
+    /// Initialized to the loaded history length (already persisted) and
+    /// rewound by `rewriteWithSummary` so a compaction re-persists its summary
+    /// + restated suffix.
+    _persisted_through: usize = 0,
 
     /// Construct an agent.
     ///
@@ -286,6 +297,9 @@ pub const Agent = struct {
             .conversation = maybe_conversation orelse conversation.Conversation.init(allocator),
             ._session = session,
         };
+        // Loaded history is already in the store; only messages produced from
+        // here on need persisting.
+        self._persisted_through = self.conversation.messages.items.len;
         return self;
     }
 
@@ -353,19 +367,11 @@ pub const Agent = struct {
         text: []const u8,
         mode: conversation.SystemMode,
     ) !void {
-        const start = self.conversation.messages.items.len;
         switch (mode) {
             .append => try self.conversation.addSystemMessage(text),
             .replace => try self.conversation.replaceSystemMessage(text),
         }
-        try turn_persist.persistTurn(
-            self._allocator,
-            &self._session,
-            &self.conversation,
-            start,
-            self.wireIdentity(),
-            &.{},
-        );
+        self.flushPersist();
     }
 
     /// Submit a user message and begin a turn, returning a resumable pull
@@ -390,49 +396,55 @@ pub const Agent = struct {
         self._auto_compacted = false;
 
         // Append + persist the user prompt up front (the dangling-prompt
-        // recovery guarantee).
-        const user_start = self.conversation.messages.items.len;
-        try addUserText(&self.conversation, message.text);
-        try turn_persist.persistTurn(
-            self._allocator,
-            &self._session,
-            &self.conversation,
-            user_start,
-            self.wireIdentity(),
-            &.{},
-        );
+        // recovery guarantee). `addMessage` adopts the blocks; persistence is
+        // brought current through the shared high-water flush.
+        try self.conversation.addMessage(.user, message.blocks, null);
+        self.flushPersist();
 
         const s = try self._allocator.create(Stream);
         s.* = Stream.init(self);
         return s;
     }
 
-    /// Persist the messages a turn produced. When the turn auto-compacted,
-    /// message indices shifted (the conversation was rewritten to
-    /// `[system..., summary, kept-suffix...]`), so persist the whole
-    /// post-compaction window instead of `[start..]`.
-    fn persistTurnTail(self: *Agent, start: usize) !void {
-        const id = self.wireIdentity();
-        if (self._auto_compacted) {
-            try turn_persist.persistCompaction(
-                self._allocator,
-                &self._session,
-                &self.conversation,
-                id,
-                &.{},
-            );
-        } else {
-            try turn_persist.persistTurn(
-                self._allocator,
-                &self._session,
-                &self.conversation,
-                start,
-                id,
-                &.{},
-            );
+    /// Bring the session store current with the conversation: append every
+    /// message in `[_persisted_through .. coherent_end)` as a single batch and
+    /// advance the high-water mark. `coherent_end` excludes a trailing
+    /// assistant message whose ToolUse blocks have no following results — a
+    /// turn interrupted between the assistant's tool call and dispatch. Leaving
+    /// it unpersisted keeps the log replayable and the live conversation in a
+    /// state `compact()` (or a tool-result `run`) can resume from. Called at
+    /// every `Stream.next()` boundary, so after any `next()` the store reflects
+    /// all committed, coherent messages.
+    pub fn flushPersist(self: *Agent) void {
+        const msgs = self.conversation.messages.items;
+        // An operation that rewrote the prefix (replace-system, cancel) can
+        // leave the mark past the new length; clamp before comparing.
+        if (self._persisted_through > msgs.len) self._persisted_through = msgs.len;
+        var end = msgs.len;
+        if (end > self._persisted_through and
+            msgs[end - 1].role == .assistant and
+            turn_persist.hasToolUseWithoutFollowingResults(&self.conversation, end - 1))
+        {
+            end -= 1;
         }
+        if (end <= self._persisted_through) return;
+        turn_persist.persistRange(
+            self._allocator,
+            &self._session,
+            &self.conversation,
+            self._persisted_through,
+            end,
+            self.wireIdentity(),
+            &.{},
+        ) catch |e| {
+            std.log.err("session: failed to persist turn: {t}", .{e});
+            return;
+        };
+        self._persisted_through = end;
     }
 
+    /// Persist the messages a turn produced. When the turn auto-compacted,
+    /// message indices shifted (the conversation was rewritten to
     fn hasToolUseBlock(msg: conversation.Message) bool {
         for (msg.content.items) |block| {
             if (block == .ToolUse) return true;
@@ -653,15 +665,9 @@ pub const Agent = struct {
             self._config.compaction.compaction_prompt orelse
             return error.NoCompactionPrompt;
         const res = try self._compactInPlace(system_prompt, extra_instructions);
-        if (res.compacted) {
-            try turn_persist.persistCompaction(
-                self._allocator,
-                &self._session,
-                &self.conversation,
-                self.wireIdentity(),
-                &.{},
-            );
-        }
+        // `rewriteWithSummary` rewound the high-water mark to the summary;
+        // flush appends the new compaction window (summary + restated suffix).
+        if (res.compacted) self.flushPersist();
         return res;
     }
 
@@ -752,6 +758,13 @@ pub const Agent = struct {
         for (conv.messages.items) |*m| m.deinit(alloc);
         conv.messages.deinit(alloc);
         conv.messages = rebuilt;
+
+        // The rewrite invalidated the old message indices. Rewind the
+        // persistence high-water mark to the inserted summary so the next
+        // `flushPersist` re-appends the new compaction window (summary +
+        // restated kept suffix) as fresh entries.
+        self._persisted_through =
+            conversation.latestCompactionIndex(conv.messages.items) orelse 0;
     }
 
     /// Run a single compaction provider call against a throwaway
@@ -1122,10 +1135,8 @@ pub const Stream = struct {
     state: State,
     /// The active provider response, when in `.streaming`.
     _response: ?provider_mod.ProviderStream = null,
-    /// First message index of this turn (for persistence).
+    /// First message index of this turn (the boundary `cancel` rolls back to).
     _start: usize,
-    /// Set once the turn's tail has been persisted (on terminal or deinit).
-    _persisted: bool = false,
     /// A terminal error to surface once any already-queued events (e.g.
     /// `provider_retry` notices pushed before the failing attempt) have been
     /// drained. `next()` yields the queue first, then this error.
@@ -1173,7 +1184,9 @@ pub const Stream = struct {
     pub fn deinit(self: *Stream) void {
         // Persist whatever the turn committed, on every exit path — including
         // dropping the stream mid-turn after some messages were committed.
-        self.persistTail();
+        // `next()` already flushes at every boundary, so this is usually a
+        // no-op; it catches a stream dropped without a final `next()`.
+        self._agent.flushPersist();
         if (self._response) |ps| ps.deinit();
         if (self._retry_message) |m| self._agent._allocator.free(m);
         self._queue.deinit();
@@ -1197,17 +1210,11 @@ pub const Stream = struct {
             var msg = conv.messages.pop().?;
             msg.deinit(conv.allocator);
         }
+        // Rolled the conversation back; clamp the persistence mark to match.
+        self._agent.flushPersist();
         self.state = .done;
     }
 
-    fn persistTail(self: *Stream) void {
-        if (self._persisted) return;
-        self._persisted = true;
-        self._agent.persistTurnTail(self._start) catch |e| {
-            std.log.err("session: failed to persist turn: {t}", .{e});
-        };
-    }
-
     /// Reset a `.failed` stream back to `.turn_start` so the caller can resume
     /// `next()` after changing the agent config (e.g. an embedder that catches
     /// a terminal `ProviderBadRequest`, rewrites the provider config, and
@@ -1230,6 +1237,11 @@ pub const Stream = struct {
     /// Pull the next event, or null past the terminal. See the contract
     /// above.
     pub fn next(self: *Stream) !?Event {
+        // Bring persistence current on every return path (events, terminal,
+        // errors). This is the interruptibility guarantee: after any `next()`
+        // the store reflects all committed, coherent messages, so the caller
+        // can `break` and immediately `compact()` or resume with a new turn.
+        defer self._agent.flushPersist();
         // Always drain queued events first; they borrow decode/conversation
         // state valid until this call returns.
         if (self._queue.pop()) |ev| return ev;
@@ -1349,7 +1361,6 @@ pub const Stream = struct {
 
                     if (!Agent.hasToolUseBlock(last)) {
                         self.state = .done;
-                        self.persistTail();
                         return .turn_complete;
                     }
 
@@ -1571,11 +1582,21 @@ const testing = std.testing;
 /// `submitUserMessage` + `runStep`: it returns the same terminal error a
 /// turn would raise.
 fn drainTurn(agent: *Agent, text: []const u8) !void {
-    var s = try agent.run(.{ .text = text });
+    var s = try runUserText(agent, text);
     defer s.deinit();
     while (try s.next()) |_| {}
 }
 
+/// Test helper: open a turn from a single user `.Text` block, mirroring the
+/// old text-only `run(.{ .text = ... })`. The block is adopted by the agent's
+/// conversation.
+fn runUserText(agent: *Agent, text: []const u8) !*Stream {
+    var blocks = [_]conversation.ContentBlock{
+        .{ .Text = try conversation.textualBlockFromSlice(agent._allocator, text) },
+    };
+    return agent.run(.{ .blocks = &blocks });
+}
+
 /// Test helper: the items of a ToolResultBlock's first text part.
 fn trText(tr: conversation.ToolResultBlock) []const u8 {
     for (tr.parts.items) |p| {
@@ -2102,6 +2123,78 @@ test "agent persists user, assistant, and tool-result messages of a turn" {
     }
 }
 
+test "interruption: persistence stays current and excludes a dangling tool call" {
+    const allocator = testing.allocator;
+
+    const scripted = [_]StubProvider.ScriptedTurn{
+        .{ .blocks = &.{
+            .{ .ToolUse = .{ .id = "tc_1", .name = "echo", .input = "hello" } },
+        } },
+        .{ .blocks = &.{
+            .{ .Text = "ok" },
+        } },
+    };
+    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
+    var threaded: std.Io.Threaded = .init(allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+    var h = TestHarness.init(allocator);
+    defer h.deinit();
+    try h.registry.register(try EchoTool.create(allocator, "echo", "ECHO:"));
+    h.activate();
+
+    var cap = CapturingStore.init(allocator);
+    defer cap.deinit();
+    const agent = try Agent.init(allocator, io, &h.config, cap.store().create(), null);
+    defer agent.deinit();
+    h.seedInto(agent);
+    agent._open_stream_fn = stub.install();
+
+    var s = try runUserText(agent, "call a tool");
+
+    // `run` persisted the user prompt up front, and nothing else yet.
+    try testing.expectEqual(@as(usize, 1), cap.roles.items.len);
+    try testing.expectEqual(conversation.MessageRole.user, cap.roles.items[0]);
+
+    // Pull until the assistant's tool-call message completes, then break — the
+    // embedder's interruption point (e.g. to handle the tool calls itself).
+    while (try s.next()) |ev| {
+        if (ev == .message_complete) break;
+    }
+
+    // The tool-call message is live in the conversation...
+    const msgs = agent.conversation.messages.items;
+    try testing.expectEqual(conversation.MessageRole.assistant, msgs[msgs.len - 1].role);
+    try testing.expect(Agent.hasToolUseBlock(msgs[msgs.len - 1]));
+    // ...but it is a dangling tool call (no results), so the store still holds
+    // only the coherent prefix: the interruption guarantee.
+    try testing.expectEqual(@as(usize, 1), cap.roles.items.len);
+
+    // Dropping the stream mid-turn must not persist the dangling call either.
+    s.deinit();
+    try testing.expectEqual(@as(usize, 1), cap.roles.items.len);
+
+    // The conversation is coherent enough to resume from: feeding the tool
+    // result as a fresh user turn resolves the dangling call, and now both the
+    // assistant tool-call and the user result become persistable.
+    var parts: std.ArrayList(conversation.ResultPartStored) = .empty;
+    try parts.append(allocator, .{ .text = try conversation.textualBlockFromSlice(allocator, "42") });
+    var s2 = try agent.run(.{ .blocks = &.{
+        .{ .ToolResult = .{
+            .tool_use_id = try allocator.dupe(u8, "tc_1"),
+            .parts = parts,
+            .is_error = false,
+        } },
+    } });
+    defer s2.deinit();
+    // user prompt, assistant(ToolUse), user(ToolResult) are all persisted once
+    // the dangling call is resolved by the new user turn.
+    try testing.expectEqual(@as(usize, 3), cap.roles.items.len);
+    try testing.expectEqual(conversation.MessageRole.user, cap.roles.items[0]);
+    try testing.expectEqual(conversation.MessageRole.assistant, cap.roles.items[1]);
+    try testing.expectEqual(conversation.MessageRole.user, cap.roles.items[2]);
+}
+
 test "agent runs a turn against NullStore without persisting or erroring" {
     const allocator = testing.allocator;
 
@@ -2455,7 +2548,7 @@ test "Stream emits tool_dispatch_start before running tools" {
     h.seedInto(agent);
     agent._open_stream_fn = stub.install();
 
-    var s = try agent.run(.{ .text = "call a tool" });
+    var s = try runUserText(agent, "call a tool");
     defer s.deinit();
 
     const first = (try s.next()).?;
@@ -2917,6 +3010,106 @@ test "compact: summarizes prefix, keeps suffix, system survives" {
     );
 }
 
+test "interruption: resuming next() after a break dispatches the pending tools" {
+    const allocator = testing.allocator;
+
+    const scripted = [_]StubProvider.ScriptedTurn{
+        .{ .blocks = &.{
+            .{ .ToolUse = .{ .id = "tc_1", .name = "echo", .input = "hello" } },
+        } },
+        .{ .blocks = &.{
+            .{ .Text = "ok" },
+        } },
+    };
+    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
+    var threaded: std.Io.Threaded = .init(allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+    var h = TestHarness.init(allocator);
+    defer h.deinit();
+    try h.registry.register(try EchoTool.create(allocator, "echo", "ECHO:"));
+    h.activate();
+
+    var cap = CapturingStore.init(allocator);
+    defer cap.deinit();
+    const agent = try Agent.init(allocator, io, &h.config, cap.store().create(), null);
+    defer agent.deinit();
+    h.seedInto(agent);
+    agent._open_stream_fn = stub.install();
+
+    var s = try runUserText(agent, "call a tool");
+    defer s.deinit();
+
+    // Break right after the assistant's tool-call message completes — the
+    // stream is paused, NOT closed. The tool has not been dispatched.
+    while (try s.next()) |ev| {
+        if (ev == .message_complete) break;
+    }
+    try testing.expectEqual(@as(usize, 1), cap.roles.items.len); // only the prompt
+
+    // Resuming `next()` on the same stream picks up where it left off and runs
+    // the agent's registered-tool machinery: no new user turn required.
+    var saw_dispatch = false;
+    while (try s.next()) |ev| {
+        if (ev == .tool_dispatch_complete) saw_dispatch = true;
+    }
+    try testing.expect(saw_dispatch);
+
+    // The full turn is now persisted: prompt, assistant(ToolUse),
+    // user(ToolResult), assistant(text).
+    try testing.expectEqual(@as(usize, 4), cap.roles.items.len);
+    try testing.expectEqual(conversation.MessageRole.user, cap.roles.items[2]);
+}
+
+test "compact: tolerates a trailing dangling tool call (interrupted turn)" {
+    const allocator = testing.allocator;
+
+    const scripted = [_]StubProvider.ScriptedTurn{
+        .{ .blocks = &.{.{ .Text = "SUMMARY OF EARLIER" }} },
+    };
+    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
+    var threaded: std.Io.Threaded = .init(allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+    var h = TestHarness.init(allocator);
+    defer h.deinit();
+    h.activate();
+    h.config.compaction = .{ .keep_verbatim = 10 };
+    var ns = null_store_mod.NullStore.init(allocator);
+    const agent = try Agent.init(allocator, io, &h.config, ns.store().create(), null);
+    defer agent.deinit();
+    h.seedInto(agent);
+    agent._open_stream_fn = stub.install();
+
+    const conv = &agent.conversation;
+    try conv.addSystemMessage("you are helpful");
+    try addUserText(conv, "first question here with several words");
+    try conv.addAssistantMessage(&.{
+        .{ .Text = try conversation.textualBlockFromSlice(allocator, "first answer with several words") },
+    }, null);
+    // A fresh turn whose assistant message ends in a tool call with no results
+    // yet — exactly the state a break right after `message_complete` leaves.
+    try addUserText(conv, "second recent question");
+    try conv.addAssistantMessage(&.{
+        .{ .ToolUse = .{
+            .id = try allocator.dupe(u8, "tc_1"),
+            .name = try allocator.dupe(u8, "echo"),
+            .input = try conversation.textualBlockFromSlice(allocator, "hello"),
+        } },
+    }, null);
+
+    // Compaction must not choke on the dangling tail: it summarizes the older
+    // turn and keeps the interrupted turn (incl. the dangling call) verbatim.
+    const res = try agent._compactInPlace("Summarize the conversation.", null);
+    try testing.expect(res.compacted);
+
+    // The dangling tool call survives as the tail, intact and resumable.
+    const last = conv.messages.items[conv.messages.items.len - 1];
+    try testing.expectEqual(conversation.MessageRole.assistant, last.role);
+    try testing.expect(Agent.hasToolUseBlock(last));
+    try testing.expectEqualStrings("tc_1", last.content.items[0].ToolUse.id);
+}
+
 test "compact: restated suffix usage reconstructs a fresh cumulative chain" {
     const allocator = testing.allocator;
 
@@ -3175,7 +3368,7 @@ const RetryRecorder = struct {
 /// `provider_retry` event into `rr` and discarding the rest. Returns the
 /// same terminal error the turn would raise.
 fn drainTurnRecording(agent: *Agent, text: []const u8, rr: *RetryRecorder) !void {
-    var s = try agent.run(.{ .text = text });
+    var s = try runUserText(agent, text);
     defer s.deinit();
     while (try s.next()) |ev| {
         if (ev == .provider_retry) try rr.append(ev.provider_retry);
diff --git a/libpanto/src/turn_persist.zig b/libpanto/src/turn_persist.zig
index 1dacfec..6d260c8 100644
--- a/libpanto/src/turn_persist.zig
+++ b/libpanto/src/turn_persist.zig
@@ -46,12 +46,28 @@ pub fn persistTurn(
     identity: WireIdentity,
     tools: []const ToolDecl,
 ) !void {
+    return persistRange(alloc, session, conv, start_index, conv.messages.items.len, identity, tools);
+}
+
+/// Persist conversation messages in `[start_index, end_index)`. Like
+/// `persistTurn` but with an explicit upper bound, so an incremental flush can
+/// exclude a trailing not-yet-coherent message (a dangling tool call awaiting
+/// its results) while still committing everything before it.
+pub fn persistRange(
+    alloc: Allocator,
+    session: *Session,
+    conv: *conversation.Conversation,
+    start_index: usize,
+    end_index: usize,
+    identity: WireIdentity,
+    tools: []const ToolDecl,
+) !void {
     var batch: std.ArrayList(PersistentMessage) = .empty;
     defer batch.deinit(alloc);
 
     const all_messages = conv.messages.items;
     var i = start_index;
-    while (i < all_messages.len) : (i += 1) {
+    while (i < end_index) : (i += 1) {
         const msg = &all_messages[i];
         if (msg.role == .assistant and hasToolUseWithoutFollowingResults(conv, i)) {
             continue;