bamboo_engine/runtime/execution/
visible_messages.rs

1//! Safe, runner-owned visible-message replay for late subscribers.
2//!
3//! This state is deliberately independent from the full-fidelity `AgentEvent`
4//! stream. It retains only visible assistant text plus minimal reconciliation
5//! controls and is cleared only after the durable history barrier is published.
6
7use std::sync::{Mutex, MutexGuard};
8
9use chrono::{DateTime, Utc};
10use tokio::sync::broadcast;
11use uuid::Uuid;
12
13/// One visible assistant message retained until durable history is confirmed.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct VisibleAssistantMessage {
16    pub id: String,
17    pub content: String,
18    pub created_at: DateTime<Utc>,
19    round_count: u32,
20}
21
22/// Atomic point-in-time replay for a message-only subscriber.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct VisibleMessageSnapshot {
25    pub version: u64,
26    pub messages: Vec<VisibleAssistantMessage>,
27    pub terminal: Option<String>,
28    pub history_committed: bool,
29}
30
31/// Safe realtime event emitted by [`VisibleMessageStream`].
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct VisibleMessageEvent {
34    pub version: u64,
35    pub kind: VisibleMessageEventKind,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub enum VisibleMessageEventKind {
40    Started {
41        message_id: String,
42        created_at: DateTime<Utc>,
43    },
44    Delta {
45        message_id: String,
46        offset: usize,
47        content: String,
48        created_at: DateTime<Utc>,
49    },
50    Discarded {
51        message_id: String,
52    },
53    Terminal {
54        reason: String,
55    },
56    HistoryCommitted,
57}
58
59#[derive(Debug, Default)]
60struct VisibleMessageReplay {
61    version: u64,
62    round_count: u32,
63    active_message_id: Option<String>,
64    messages: Vec<VisibleAssistantMessage>,
65    terminal: Option<String>,
66    history_committed: bool,
67}
68
69/// Per-run message-only replay and broadcast state.
70///
71/// The replay is intentionally not byte-capped: it contains only visible text
72/// generated by the current run and remains authoritative until durable history
73/// is acknowledged. Silently evicting it would make late subscribers lose text.
74#[derive(Debug)]
75pub struct VisibleMessageStream {
76    state: Mutex<VisibleMessageReplay>,
77    sender: broadcast::Sender<VisibleMessageEvent>,
78}
79
80impl Default for VisibleMessageStream {
81    fn default() -> Self {
82        Self::new()
83    }
84}
85
86impl VisibleMessageStream {
87    const CHANNEL_CAPACITY: usize = 1000;
88
89    pub fn new() -> Self {
90        let (sender, _) = broadcast::channel(Self::CHANNEL_CAPACITY);
91        Self {
92            state: Mutex::new(VisibleMessageReplay::default()),
93            sender,
94        }
95    }
96
97    fn lock(&self) -> MutexGuard<'_, VisibleMessageReplay> {
98        self.state
99            .lock()
100            .unwrap_or_else(std::sync::PoisonError::into_inner)
101    }
102
103    fn send_locked(&self, event: VisibleMessageEvent) {
104        let _ = self.sender.send(event);
105    }
106
107    /// Subscribe before taking the snapshot. Concurrent events are either in the
108    /// snapshot or queued with a greater version, never neither.
109    pub fn subscribe_with_snapshot(
110        &self,
111    ) -> (
112        broadcast::Receiver<VisibleMessageEvent>,
113        VisibleMessageSnapshot,
114    ) {
115        let receiver = self.sender.subscribe();
116        let state = self.lock();
117        let snapshot = VisibleMessageSnapshot {
118            version: state.version,
119            messages: state
120                .messages
121                .iter()
122                .filter(|message| !message.content.is_empty())
123                .cloned()
124                .collect(),
125            terminal: state.terminal.clone(),
126            history_committed: state.history_committed,
127        };
128        (receiver, snapshot)
129    }
130
131    pub fn begin_round(&self, round_count: u32) {
132        self.lock().round_count = round_count;
133    }
134
135    /// Start one provider response. Previous segments remain replayable until an
136    /// explicit discard or the durable-history barrier; a later attempt in the
137    /// same round does not prove the earlier segment was rolled back.
138    pub fn start(&self, message_id: String, created_at: DateTime<Utc>) {
139        let mut state = self.lock();
140        if !state
141            .messages
142            .iter()
143            .any(|message| message.id == message_id)
144        {
145            let round_count = state.round_count;
146            state.messages.push(VisibleAssistantMessage {
147                id: message_id.clone(),
148                content: String::new(),
149                created_at,
150                round_count,
151            });
152        }
153        state.active_message_id = Some(message_id.clone());
154        state.history_committed = false;
155        state.version = state.version.saturating_add(1);
156        self.send_locked(VisibleMessageEvent {
157            version: state.version,
158            kind: VisibleMessageEventKind::Started {
159                message_id,
160                created_at,
161            },
162        });
163    }
164
165    /// Append visible assistant text. Legacy producers without an explicit start
166    /// receive a host-minted temporary identity and reconcile on history commit.
167    pub fn append(&self, content: String) {
168        if content.is_empty() {
169            return;
170        }
171        let mut state = self.lock();
172        let message_id = match state.active_message_id.clone() {
173            Some(message_id) => message_id,
174            None => {
175                let message_id = Uuid::new_v4().to_string();
176                let created_at = Utc::now();
177                let round_count = state.round_count;
178                state.messages.push(VisibleAssistantMessage {
179                    id: message_id.clone(),
180                    content: String::new(),
181                    created_at,
182                    round_count,
183                });
184                state.active_message_id = Some(message_id.clone());
185                state.version = state.version.saturating_add(1);
186                self.send_locked(VisibleMessageEvent {
187                    version: state.version,
188                    kind: VisibleMessageEventKind::Started {
189                        message_id: message_id.clone(),
190                        created_at,
191                    },
192                });
193                message_id
194            }
195        };
196
197        let Some(index) = state
198            .messages
199            .iter()
200            .position(|message| message.id == message_id)
201        else {
202            return;
203        };
204        let offset = state.messages[index].content.len();
205        let created_at = state.messages[index].created_at;
206        state.messages[index].content.push_str(&content);
207        state.version = state.version.saturating_add(1);
208        self.send_locked(VisibleMessageEvent {
209            version: state.version,
210            kind: VisibleMessageEventKind::Delta {
211                message_id,
212                offset,
213                content,
214                created_at,
215            },
216        });
217    }
218
219    pub fn discard(&self, message_id: &str) {
220        let mut state = self.lock();
221        let previous_len = state.messages.len();
222        state.messages.retain(|message| message.id != message_id);
223        if state.messages.len() == previous_len {
224            return;
225        }
226        if state.active_message_id.as_deref() == Some(message_id) {
227            state.active_message_id = None;
228        }
229        state.version = state.version.saturating_add(1);
230        self.send_locked(VisibleMessageEvent {
231            version: state.version,
232            kind: VisibleMessageEventKind::Discarded {
233                message_id: message_id.to_string(),
234            },
235        });
236    }
237
238    pub fn mark_terminal(&self, reason: impl Into<String>) {
239        let mut state = self.lock();
240        let reason = reason.into();
241        if state.terminal.as_deref() == Some(reason.as_str()) {
242            return;
243        }
244        state.terminal = Some(reason.clone());
245        state.version = state.version.saturating_add(1);
246        self.send_locked(VisibleMessageEvent {
247            version: state.version,
248            kind: VisibleMessageEventKind::Terminal { reason },
249        });
250    }
251
252    /// Clear replay only after durable history has actually been published.
253    pub fn history_committed(&self) {
254        let mut state = self.lock();
255        state.messages.clear();
256        state.active_message_id = None;
257        state.history_committed = true;
258        state.version = state.version.saturating_add(1);
259        self.send_locked(VisibleMessageEvent {
260            version: state.version,
261            kind: VisibleMessageEventKind::HistoryCommitted,
262        });
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn subscribe_snapshot_and_live_tail_have_no_overlap_after_version_filter() {
272        let stream = VisibleMessageStream::new();
273        stream.begin_round(0);
274        let created_at = Utc::now();
275        stream.start("message-1".into(), created_at);
276        stream.append("before".into());
277
278        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
279        assert_eq!(snapshot.messages[0].content, "before");
280        stream.append(" after".into());
281
282        let event = receiver.try_recv().expect("live delta");
283        assert!(event.version > snapshot.version);
284        assert!(matches!(
285            event.kind,
286            VisibleMessageEventKind::Delta { offset: 6, ref content, .. } if content == " after"
287        ));
288    }
289
290    #[test]
291    fn explicit_retry_rollback_discards_previous_attempt() {
292        let stream = VisibleMessageStream::new();
293        stream.begin_round(2);
294        stream.start("attempt-1".into(), Utc::now());
295        stream.append("partial".into());
296        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
297        assert_eq!(snapshot.messages.len(), 1);
298
299        stream.discard("attempt-1");
300        stream.start("attempt-2".into(), Utc::now());
301        let discarded = receiver.try_recv().expect("discarded event");
302        let started = receiver.try_recv().expect("replacement start");
303        assert!(matches!(
304            discarded.kind,
305            VisibleMessageEventKind::Discarded { ref message_id } if message_id == "attempt-1"
306        ));
307        assert!(matches!(
308            started.kind,
309            VisibleMessageEventKind::Started { ref message_id, .. } if message_id == "attempt-2"
310        ));
311    }
312
313    #[test]
314    fn multi_round_messages_survive_until_history_commit() {
315        let stream = VisibleMessageStream::new();
316        stream.begin_round(1);
317        stream.start("round-1".into(), Utc::now());
318        stream.append("first".into());
319        stream.begin_round(2);
320        stream.start("round-2".into(), Utc::now());
321        stream.append("second".into());
322
323        let (_, snapshot) = stream.subscribe_with_snapshot();
324        assert_eq!(snapshot.messages.len(), 2);
325        assert_eq!(snapshot.messages[0].round_count, 1);
326        assert_eq!(snapshot.messages[0].content, "first");
327        assert_eq!(snapshot.messages[1].round_count, 2);
328        assert_eq!(snapshot.messages[1].content, "second");
329    }
330
331    #[test]
332    fn failed_and_cancelled_terminal_reasons_remain_in_late_snapshots() {
333        for reason in ["error", "cancelled"] {
334            let stream = VisibleMessageStream::new();
335            stream.start("partial".into(), Utc::now());
336            stream.append("visible prefix".into());
337            stream.mark_terminal(reason);
338
339            let (_, snapshot) = stream.subscribe_with_snapshot();
340            assert_eq!(snapshot.terminal.as_deref(), Some(reason));
341            assert_eq!(snapshot.messages[0].content, "visible prefix");
342        }
343    }
344
345    #[test]
346    fn history_commit_clears_text_only_after_control_is_emitted() {
347        let stream = VisibleMessageStream::new();
348        stream.start("message-1".into(), Utc::now());
349        stream.append("visible".into());
350        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
351        assert_eq!(snapshot.messages.len(), 1);
352
353        stream.history_committed();
354        let committed = receiver.try_recv().expect("history commit control");
355        assert!(matches!(
356            committed.kind,
357            VisibleMessageEventKind::HistoryCommitted
358        ));
359        let (_, settled) = stream.subscribe_with_snapshot();
360        assert!(settled.messages.is_empty());
361        assert!(settled.history_committed);
362    }
363}
bamboo_engine/runtime/execution/visible_messages.rs

bamboo_engine/runtime/execution/
visible_messages.rs
