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. Ordinary replay clears at the durable history barrier; a selected
6//! execution can retain the same buffer for its finite runner lifetime.
7
8use std::sync::{Mutex, MutexGuard};
9
10use chrono::{DateTime, Utc};
11use serde::Serialize;
12use tokio::sync::broadcast;
13use uuid::Uuid;
14
15/// One visible assistant message retained until durable history is confirmed.
16#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
17pub struct VisibleAssistantMessage {
18    pub id: String,
19    pub content: String,
20    pub created_at: DateTime<Utc>,
21    #[serde(skip)]
22    round_count: u32,
23    #[serde(skip)]
24    committed: bool,
25}
26
27/// Atomic point-in-time replay for a message-only subscriber.
28#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
29pub struct VisibleMessageSnapshot {
30    pub version: u64,
31    pub messages: Vec<VisibleAssistantMessage>,
32    pub terminal: Option<String>,
33    pub history_committed: bool,
34}
35
36/// Host-captured identity for one receipt-bound execution's visible output.
37/// This is read data, never an execution capability or a transcript proof.
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct SelectedRunOutputIdentity {
40    pub session_id: String,
41    pub run_id: String,
42    pub session_created_at: DateTime<Utc>,
43    pub receipt: bamboo_domain::SessionInboxReceipt,
44}
45
46/// A replacement snapshot of visible text, not a full event/completion replay.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct SelectedRunOutputSnapshot {
49    pub identity: SelectedRunOutputIdentity,
50    pub output: VisibleMessageSnapshot,
51}
52
53/// Safe realtime event emitted by [`VisibleMessageStream`].
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct VisibleMessageEvent {
56    pub version: u64,
57    pub kind: VisibleMessageEventKind,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum VisibleMessageEventKind {
62    Started {
63        message_id: String,
64        created_at: DateTime<Utc>,
65    },
66    Delta {
67        message_id: String,
68        offset: usize,
69        content: String,
70        created_at: DateTime<Utc>,
71    },
72    Discarded {
73        message_id: String,
74    },
75    Terminal {
76        reason: String,
77    },
78    HistoryCommitted,
79}
80
81#[derive(Debug, Default)]
82struct VisibleMessageReplay {
83    version: u64,
84    round_count: u32,
85    active_message_id: Option<String>,
86    messages: Vec<VisibleAssistantMessage>,
87    terminal: Option<String>,
88    history_committed: bool,
89    selected_output: Option<SelectedRunOutputIdentity>,
90}
91
92/// Per-run message-only replay and broadcast state.
93///
94/// The replay is intentionally not byte-capped: it contains only visible text
95/// generated by the current run and remains authoritative until durable history
96/// is acknowledged. Selected output may retain it until the existing runner
97/// is replaced/evicted. Silently evicting it would make subscribers lose text.
98#[derive(Debug)]
99pub struct VisibleMessageStream {
100    state: Mutex<VisibleMessageReplay>,
101    sender: broadcast::Sender<VisibleMessageEvent>,
102}
103
104impl Default for VisibleMessageStream {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110impl VisibleMessageStream {
111    const CHANNEL_CAPACITY: usize = 1000;
112
113    pub fn new() -> Self {
114        let (sender, _) = broadcast::channel(Self::CHANNEL_CAPACITY);
115        Self {
116            state: Mutex::new(VisibleMessageReplay::default()),
117            sender,
118        }
119    }
120
121    fn lock(&self) -> MutexGuard<'_, VisibleMessageReplay> {
122        self.state
123            .lock()
124            .unwrap_or_else(std::sync::PoisonError::into_inner)
125    }
126
127    fn send_locked(&self, event: VisibleMessageEvent) {
128        let _ = self.sender.send(event);
129    }
130
131    /// Subscribe before taking the snapshot. Concurrent events are either in the
132    /// snapshot or queued with a greater version, never neither.
133    pub fn subscribe_with_snapshot(
134        &self,
135    ) -> (
136        broadcast::Receiver<VisibleMessageEvent>,
137        VisibleMessageSnapshot,
138    ) {
139        let receiver = self.sender.subscribe();
140        let state = self.lock();
141        let snapshot = VisibleMessageSnapshot {
142            version: state.version,
143            messages: state
144                .messages
145                .iter()
146                .filter(|message| !message.committed && !message.content.is_empty())
147                .cloned()
148                .collect(),
149            terminal: state.terminal.clone(),
150            history_committed: state.history_committed,
151        };
152        (receiver, snapshot)
153    }
154
155    /// Called only from a live selected reservation with its actual Root writer,
156    /// before any producer starts. The same buffer remains owned by this runner;
157    /// replacement/eviction/restart never reconstructs it from session history.
158    pub(crate) fn retain_selected_output(
159        &self,
160        owner: &bamboo_domain::RootActorRuntimeWrite,
161        receipt: bamboo_domain::SessionInboxReceipt,
162    ) -> std::io::Result<()> {
163        let identity = SelectedRunOutputIdentity {
164            session_id: owner.fence.actor_id.clone(),
165            run_id: owner.fence.run_id.clone(),
166            session_created_at: owner.created_at,
167            receipt,
168        };
169        let mut state = self.lock();
170        if let Some(current) = &state.selected_output {
171            return if current == &identity {
172                Ok(())
173            } else {
174                Err(std::io::Error::other(
175                    "Selected output identity cannot change",
176                ))
177            };
178        }
179        if state.version != 0 || identity.receipt.generation == 0 {
180            return Err(std::io::Error::other(
181                "Selected output must bind before publication",
182            ));
183        }
184        state.selected_output = Some(identity);
185        Ok(())
186    }
187
188    /// Read the same visible-message buffer atomically, including text already
189    /// acknowledged by history. Ordinary WS snapshots still hide that prefix.
190    /// Only the owning runner's lifetime retains these bytes; no new store/TTL.
191    pub fn selected_output_snapshot(&self) -> Option<SelectedRunOutputSnapshot> {
192        let state = self.lock();
193        Some(SelectedRunOutputSnapshot {
194            identity: state.selected_output.clone()?,
195            output: VisibleMessageSnapshot {
196                version: state.version,
197                messages: state
198                    .messages
199                    .iter()
200                    .filter(|m| !m.content.is_empty())
201                    .cloned()
202                    .collect(),
203                terminal: state.terminal.clone(),
204                history_committed: state.history_committed,
205            },
206        })
207    }
208
209    pub fn begin_round(&self, round_count: u32) {
210        self.lock().round_count = round_count;
211    }
212
213    /// Start one provider response. Previous segments remain replayable until an
214    /// explicit discard or the durable-history barrier; a later attempt in the
215    /// same round does not prove the earlier segment was rolled back.
216    pub fn start(&self, message_id: String, created_at: DateTime<Utc>) {
217        let mut state = self.lock();
218        // A post-barrier producer reusing an id starts the same replacement
219        // value an ordinary cleared stream would have; never append old text.
220        if let Some(previous) = state
221            .messages
222            .iter()
223            .position(|m| m.id == message_id && m.committed)
224        {
225            state.messages.remove(previous);
226        }
227        if !state
228            .messages
229            .iter()
230            .any(|message| message.id == message_id)
231        {
232            let round_count = state.round_count;
233            state.messages.push(VisibleAssistantMessage {
234                id: message_id.clone(),
235                content: String::new(),
236                created_at,
237                round_count,
238                committed: false,
239            });
240        }
241        state.active_message_id = Some(message_id.clone());
242        state.history_committed = false;
243        state.version = state.version.saturating_add(1);
244        self.send_locked(VisibleMessageEvent {
245            version: state.version,
246            kind: VisibleMessageEventKind::Started {
247                message_id,
248                created_at,
249            },
250        });
251    }
252
253    /// Append visible assistant text. Legacy producers without an explicit start
254    /// receive a host-minted temporary identity and reconcile on history commit.
255    pub fn append(&self, content: String) {
256        if content.is_empty() {
257            return;
258        }
259        let mut state = self.lock();
260        state.history_committed = false;
261        let message_id = match state.active_message_id.clone() {
262            Some(message_id) => message_id,
263            None => {
264                let message_id = Uuid::new_v4().to_string();
265                let created_at = Utc::now();
266                let round_count = state.round_count;
267                state.messages.push(VisibleAssistantMessage {
268                    id: message_id.clone(),
269                    content: String::new(),
270                    created_at,
271                    round_count,
272                    committed: false,
273                });
274                state.active_message_id = Some(message_id.clone());
275                state.version = state.version.saturating_add(1);
276                self.send_locked(VisibleMessageEvent {
277                    version: state.version,
278                    kind: VisibleMessageEventKind::Started {
279                        message_id: message_id.clone(),
280                        created_at,
281                    },
282                });
283                message_id
284            }
285        };
286
287        let Some(index) = state
288            .messages
289            .iter()
290            .position(|message| message.id == message_id)
291        else {
292            return;
293        };
294        let offset = state.messages[index].content.len();
295        let created_at = state.messages[index].created_at;
296        state.messages[index].content.push_str(&content);
297        state.version = state.version.saturating_add(1);
298        self.send_locked(VisibleMessageEvent {
299            version: state.version,
300            kind: VisibleMessageEventKind::Delta {
301                message_id,
302                offset,
303                content,
304                created_at,
305            },
306        });
307    }
308
309    pub fn discard(&self, message_id: &str) {
310        let mut state = self.lock();
311        let previous_len = state.messages.len();
312        state.messages.retain(|message| message.id != message_id);
313        if state.messages.len() == previous_len {
314            return;
315        }
316        if state.active_message_id.as_deref() == Some(message_id) {
317            state.active_message_id = None;
318        }
319        state.history_committed = false;
320        state.version = state.version.saturating_add(1);
321        self.send_locked(VisibleMessageEvent {
322            version: state.version,
323            kind: VisibleMessageEventKind::Discarded {
324                message_id: message_id.to_string(),
325            },
326        });
327    }
328
329    pub fn mark_terminal(&self, reason: impl Into<String>) {
330        let mut state = self.lock();
331        let reason = reason.into();
332        if state.terminal.as_deref() == Some(reason.as_str()) {
333            return;
334        }
335        state.terminal = Some(reason.clone());
336        state.version = state.version.saturating_add(1);
337        self.send_locked(VisibleMessageEvent {
338            version: state.version,
339            kind: VisibleMessageEventKind::Terminal { reason },
340        });
341    }
342
343    /// Reconcile ordinary replay only after durable history publication.
344    /// Selected output retains the same bytes until its runner is reclaimed.
345    pub fn history_committed(&self) {
346        let mut state = self.lock();
347        if state.selected_output.is_some() {
348            for message in &mut state.messages {
349                message.committed = true;
350            }
351        } else {
352            state.messages.clear();
353        }
354        state.active_message_id = None;
355        state.history_committed = true;
356        state.version = state.version.saturating_add(1);
357        self.send_locked(VisibleMessageEvent {
358            version: state.version,
359            kind: VisibleMessageEventKind::HistoryCommitted,
360        });
361    }
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367
368    #[test]
369    fn subscribe_snapshot_and_live_tail_have_no_overlap_after_version_filter() {
370        let stream = VisibleMessageStream::new();
371        stream.begin_round(0);
372        let created_at = Utc::now();
373        stream.start("message-1".into(), created_at);
374        stream.append("before".into());
375
376        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
377        assert_eq!(snapshot.messages[0].content, "before");
378        stream.append(" after".into());
379
380        let event = receiver.try_recv().expect("live delta");
381        assert!(event.version > snapshot.version);
382        assert!(matches!(
383            event.kind,
384            VisibleMessageEventKind::Delta { offset: 6, ref content, .. } if content == " after"
385        ));
386    }
387
388    #[test]
389    fn explicit_retry_rollback_discards_previous_attempt() {
390        let stream = VisibleMessageStream::new();
391        stream.begin_round(2);
392        stream.start("attempt-1".into(), Utc::now());
393        stream.append("partial".into());
394        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
395        assert_eq!(snapshot.messages.len(), 1);
396
397        stream.discard("attempt-1");
398        stream.start("attempt-2".into(), Utc::now());
399        let discarded = receiver.try_recv().expect("discarded event");
400        let started = receiver.try_recv().expect("replacement start");
401        assert!(matches!(
402            discarded.kind,
403            VisibleMessageEventKind::Discarded { ref message_id } if message_id == "attempt-1"
404        ));
405        assert!(matches!(
406            started.kind,
407            VisibleMessageEventKind::Started { ref message_id, .. } if message_id == "attempt-2"
408        ));
409    }
410
411    #[test]
412    fn multi_round_messages_survive_until_history_commit() {
413        let stream = VisibleMessageStream::new();
414        stream.begin_round(1);
415        stream.start("round-1".into(), Utc::now());
416        stream.append("first".into());
417        stream.begin_round(2);
418        stream.start("round-2".into(), Utc::now());
419        stream.append("second".into());
420
421        let (_, snapshot) = stream.subscribe_with_snapshot();
422        assert_eq!(snapshot.messages.len(), 2);
423        assert_eq!(snapshot.messages[0].round_count, 1);
424        assert_eq!(snapshot.messages[0].content, "first");
425        assert_eq!(snapshot.messages[1].round_count, 2);
426        assert_eq!(snapshot.messages[1].content, "second");
427    }
428
429    #[test]
430    fn failed_and_cancelled_terminal_reasons_remain_in_late_snapshots() {
431        for reason in ["error", "cancelled"] {
432            let stream = VisibleMessageStream::new();
433            stream.start("partial".into(), Utc::now());
434            stream.append("visible prefix".into());
435            stream.mark_terminal(reason);
436
437            let (_, snapshot) = stream.subscribe_with_snapshot();
438            assert_eq!(snapshot.terminal.as_deref(), Some(reason));
439            assert_eq!(snapshot.messages[0].content, "visible prefix");
440        }
441    }
442
443    #[test]
444    fn history_commit_clears_text_only_after_control_is_emitted() {
445        let stream = VisibleMessageStream::new();
446        stream.start("message-1".into(), Utc::now());
447        stream.append("visible".into());
448        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
449        assert_eq!(snapshot.messages.len(), 1);
450
451        stream.history_committed();
452        let committed = receiver.try_recv().expect("history commit control");
453        assert!(matches!(
454            committed.kind,
455            VisibleMessageEventKind::HistoryCommitted
456        ));
457        let (_, settled) = stream.subscribe_with_snapshot();
458        assert!(settled.messages.is_empty());
459        assert!(settled.history_committed);
460    }
461
462    fn selected_owner() -> bamboo_domain::RootActorRuntimeWrite {
463        bamboo_domain::RootActorRuntimeWrite {
464            fence: bamboo_domain::ActorActivationFence {
465                schema_version: 1,
466                actor_id: "output-root".into(),
467                activation_id: "activation".into(),
468                attempt: 1,
469                run_id: "output-run".into(),
470                lease_owner: "host".into(),
471                lease_epoch: 1,
472            },
473            created_at: Utc::now(),
474            conversation_generation: bamboo_domain::Session::new("output-root", "model")
475                .conversation_generation_observation(),
476        }
477    }
478    fn selected_receipt() -> bamboo_domain::SessionInboxReceipt {
479        bamboo_domain::SessionInboxReceipt {
480            id: bamboo_domain::SessionMessageId::parse("selected-input").unwrap(),
481            generation: 7,
482        }
483    }
484
485    #[test]
486    fn run_output_retains_same_bytes_after_history_without_changing_ws_replay() {
487        let stream = VisibleMessageStream::new();
488        let owner = selected_owner();
489        stream
490            .retain_selected_output(&owner, selected_receipt())
491            .unwrap();
492        stream.start("visible".into(), Utc::now());
493        stream.append("界\r\nfinal".into());
494        stream.mark_terminal("complete");
495        let before = stream.selected_output_snapshot().unwrap();
496        stream.history_committed();
497        let replay = stream.selected_output_snapshot().unwrap();
498        assert_eq!(replay.identity.session_id, owner.fence.actor_id);
499        assert_eq!(replay.identity.run_id, owner.fence.run_id);
500        assert_eq!(replay.identity.session_created_at, owner.created_at);
501        assert_eq!(replay.identity.receipt, selected_receipt());
502        assert_eq!(
503            serde_json::to_value(&replay.output.messages).unwrap(),
504            serde_json::to_value(&before.output.messages).unwrap()
505        );
506        assert!(replay.output.history_committed && replay.output.version > before.output.version);
507        assert_eq!(replay.output.terminal.as_deref(), Some("complete"));
508        let (_, legacy) = stream.subscribe_with_snapshot();
509        assert!(legacy.messages.is_empty());
510        assert!(legacy.history_committed);
511        let wire = serde_json::to_value(&replay.output).unwrap();
512        assert_eq!(wire["messages"][0].as_object().unwrap().len(), 3);
513        assert!(wire["messages"][0].get("committed").is_none());
514    }
515
516    #[test]
517    fn run_output_rejects_rebinding_or_late_opt_in_and_defaults_to_no_receipt() {
518        let stream = VisibleMessageStream::new();
519        assert!(stream.selected_output_snapshot().is_none());
520        let owner = selected_owner();
521        stream
522            .retain_selected_output(&owner, selected_receipt())
523            .unwrap();
524        stream
525            .retain_selected_output(&owner, selected_receipt())
526            .unwrap();
527        let mut foreign = owner.clone();
528        foreign.fence.run_id = "successor".into();
529        assert!(stream
530            .retain_selected_output(&foreign, selected_receipt())
531            .is_err());
532        let mut foreign = owner.clone();
533        foreign.created_at += chrono::Duration::seconds(1);
534        assert!(stream
535            .retain_selected_output(&foreign, selected_receipt())
536            .is_err());
537        let mut receipt = selected_receipt();
538        receipt.generation += 1;
539        assert!(stream.retain_selected_output(&owner, receipt).is_err());
540        let late = VisibleMessageStream::new();
541        late.append("unattributed prefix".into());
542        assert!(late
543            .retain_selected_output(&owner, selected_receipt())
544            .is_err());
545        assert!(late.selected_output_snapshot().is_none());
546    }
547
548    #[test]
549    fn run_output_later_segments_do_not_restore_committed_text_to_ws() {
550        let stream = VisibleMessageStream::new();
551        stream
552            .retain_selected_output(&selected_owner(), selected_receipt())
553            .unwrap();
554        stream.start("old".into(), Utc::now());
555        stream.append("committed prefix".into());
556        stream.history_committed();
557        stream.start("next".into(), Utc::now());
558        stream.append("new segment".into());
559        let (_, legacy) = stream.subscribe_with_snapshot();
560        assert_eq!(legacy.messages.len(), 1);
561        assert_eq!(legacy.messages[0].id, "next");
562        let output = stream.selected_output_snapshot().unwrap().output;
563        assert_eq!(output.messages.len(), 2);
564        assert!(!output.history_committed);
565        stream.discard("next");
566        let (_, legacy) = stream.subscribe_with_snapshot();
567        assert!(legacy.messages.is_empty());
568        assert_eq!(
569            stream.selected_output_snapshot().unwrap().output.messages[0].content,
570            "committed prefix"
571        );
572    }
573
574    #[test]
575    fn run_output_reused_id_and_legacy_delta_never_append_committed_bytes() {
576        let stream = VisibleMessageStream::new();
577        stream
578            .retain_selected_output(&selected_owner(), selected_receipt())
579            .unwrap();
580        stream.start("same".into(), Utc::now());
581        stream.append("old".into());
582        stream.history_committed();
583        stream.start("same".into(), Utc::now());
584        stream.append("replacement".into());
585        let snapshot = stream.selected_output_snapshot().unwrap().output;
586        assert_eq!(snapshot.messages.len(), 1);
587        assert_eq!(snapshot.messages[0].content, "replacement");
588        stream.history_committed();
589        stream.append("legacy new segment".into());
590        let snapshot = stream.selected_output_snapshot().unwrap().output;
591        assert_eq!(snapshot.messages.len(), 2);
592        assert!(!snapshot.history_committed);
593        stream.history_committed();
594        stream.discard("same");
595        assert!(
596            !stream
597                .selected_output_snapshot()
598                .unwrap()
599                .output
600                .history_committed
601        );
602    }
603
604    #[test]
605    fn run_output_error_and_cancellation_keep_prefix_without_inventing_commit() {
606        for reason in ["error", "cancelled"] {
607            let stream = VisibleMessageStream::new();
608            stream
609                .retain_selected_output(&selected_owner(), selected_receipt())
610                .unwrap();
611            stream.append("partial visible output".into());
612            stream.mark_terminal(reason);
613            let snapshot = stream.selected_output_snapshot().unwrap().output;
614            assert_eq!(snapshot.messages[0].content, "partial visible output");
615            assert_eq!(snapshot.terminal.as_deref(), Some(reason));
616            assert!(!snapshot.history_committed);
617            stream.history_committed();
618            let snapshot = stream.selected_output_snapshot().unwrap().output;
619            assert_eq!(snapshot.messages[0].content, "partial visible output");
620            assert!(snapshot.history_committed);
621        }
622    }
623}
bamboo_engine/runtime/execution/visible_messages.rs

bamboo_engine/runtime/execution/
visible_messages.rs
