Add Lock + Confirm/Cancel semantics to MessageQueue

Add Confirm and Cancel methods to the MessageQueue interface so that
persistent queue implementations can use a transactional dequeue pattern:
Dequeue locks a message (in-flight), Confirm permanently removes it
after sess.AddMessage succeeds, Cancel releases it back to the queue
on failure.

This prevents message loss when the process crashes or the context is
cancelled between dequeue and session persistence. The in-memory
implementation treats Confirm/Cancel as no-ops since the message is
already consumed from the channel on Dequeue.

The agent loop now calls Confirm after successfully adding a follow-up
message to the session. Drain (used for steer messages) auto-confirms
all messages in a batch.

Author: trungutt

pkg/runtime/loop.go

@@ -427,9 +427,14 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				// a new turn — the model sees them as fresh input, not a
 				// mid-stream interruption. Each follow-up gets a full
 				// undivided agent turn.
+				//
+				// Dequeue locks the message; Confirm is called after
+				// AddMessage succeeds so persistent queue implementations
+				// can safely re-queue on failure.
 				if followUp, ok := r.DequeueFollowUp(ctx); ok {
 					userMsg := session.UserMessage(followUp.Content, followUp.MultiContent...)
 					sess.AddMessage(userMsg)
+					_ = r.followUpQueue.Confirm(ctx)
 					events <- UserMessage(followUp.Content, sess.ID, followUp.MultiContent, len(sess.Messages)-1)
 					r.compactIfNeeded(ctx, sess, a, m, contextLimit, messageCountBeforeTools, events)
 					continue // re-enter the loop for a new turn

pkg/runtime/runtime.go

@@ -93,18 +93,35 @@ type QueuedMessage struct {
 // is called from API handlers while Dequeue/Drain are called from the agent
 // loop goroutine.
 //
+// Dequeue uses a Lock + Confirm/Cancel pattern: Dequeue locks the next
+// message (making it invisible to subsequent Dequeue calls), Confirm
+// permanently removes it after the message has been successfully added to
+// the session, and Cancel releases it back to the queue if processing
+// fails. This prevents message loss in persistent queue implementations
+// when the process crashes or the context is cancelled between dequeue
+// and session persistence.
+//
 // The default implementation is NewInMemoryMessageQueue. Callers that need
 // durable or distributed storage can provide their own implementation
 // via the WithSteerQueue or WithFollowUpQueue options.
 type MessageQueue interface {
 	// Enqueue adds a message to the queue. Returns false if the queue is
 	// full or the context is cancelled.
 	Enqueue(ctx context.Context, msg QueuedMessage) bool
-	// Dequeue removes and returns the next message from the queue.
-	// Returns the message and true, or a zero value and false if the
-	// queue is empty. Must not block.
+	// Dequeue locks and returns the next message from the queue. The
+	// message is invisible to subsequent Dequeue calls until Confirm or
+	// Cancel is called. Returns the message and true, or a zero value
+	// and false if the queue is empty. Must not block.
 	Dequeue(ctx context.Context) (QueuedMessage, bool)
-	// Drain returns all pending messages and removes them from the queue.
+	// Confirm permanently removes the most recently dequeued message.
+	// Must be called after the message has been successfully persisted
+	// to the session. For in-memory queues this is a no-op.
+	Confirm(ctx context.Context) error
+	// Cancel releases the most recently dequeued message back to the
+	// queue. For in-memory queues this is a no-op (the message was
+	// already consumed from the channel).
+	Cancel(ctx context.Context) error
+	// Drain locks, returns, and auto-confirms all pending messages.
 	// Must not block — if the queue is empty it returns nil.
 	Drain(ctx context.Context) []QueuedMessage
 	// Len returns the current number of messages in the queue.
@@ -148,6 +165,14 @@ func (q *inMemoryMessageQueue) Dequeue(_ context.Context) (QueuedMessage, bool)
 	}
 }
 
+// Confirm is a no-op for in-memory queues — the message was already
+// removed from the channel on Dequeue.
+func (q *inMemoryMessageQueue) Confirm(_ context.Context) error { return nil }
+
+// Cancel is a no-op for in-memory queues — the message cannot be put
+// back into a buffered channel without risking deadlock.
+func (q *inMemoryMessageQueue) Cancel(_ context.Context) error { return nil }
+
 func (q *inMemoryMessageQueue) Drain(_ context.Context) []QueuedMessage {
 	var msgs []QueuedMessage
 	for {