bamboo_engine/runtime/execution/

visible_messages.rs

//! Safe, runner-owned visible-message replay for late subscribers.
//!
//! This state is deliberately independent from the full-fidelity `AgentEvent`
//! stream. It retains only visible assistant text plus minimal reconciliation
//! controls. Ordinary replay clears at the durable history barrier; a selected
//! execution can retain the same buffer for its finite runner lifetime.

use std::sync::{Mutex, MutexGuard};

use chrono::{DateTime, Utc};
use serde::Serialize;
use tokio::sync::broadcast;
use uuid::Uuid;

/// One visible assistant message retained until durable history is confirmed.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct VisibleAssistantMessage {
    pub id: String,
    pub content: String,
    pub created_at: DateTime<Utc>,
    #[serde(skip)]
    round_count: u32,
    #[serde(skip)]
    committed: bool,
}

/// Atomic point-in-time replay for a message-only subscriber.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct VisibleMessageSnapshot {
    pub version: u64,
    pub messages: Vec<VisibleAssistantMessage>,
    pub terminal: Option<String>,
    pub history_committed: bool,
}

/// Host-captured identity for one receipt-bound execution's visible output.
/// This is read data, never an execution capability or a transcript proof.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SelectedRunOutputIdentity {
    pub session_id: String,
    pub run_id: String,
    pub session_created_at: DateTime<Utc>,
    pub receipt: bamboo_domain::SessionInboxReceipt,
}

/// A replacement snapshot of visible text, not a full event/completion replay.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SelectedRunOutputSnapshot {
    pub identity: SelectedRunOutputIdentity,
    pub output: VisibleMessageSnapshot,
}

/// Safe realtime event emitted by [`VisibleMessageStream`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VisibleMessageEvent {
    pub version: u64,
    pub kind: VisibleMessageEventKind,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VisibleMessageEventKind {
    Started {
        message_id: String,
        created_at: DateTime<Utc>,
    },
    Delta {
        message_id: String,
        offset: usize,
        content: String,
        created_at: DateTime<Utc>,
    },
    Discarded {
        message_id: String,
    },
    Terminal {
        reason: String,
    },
    HistoryCommitted,
}

#[derive(Debug, Default)]
struct VisibleMessageReplay {
    version: u64,
    round_count: u32,
    active_message_id: Option<String>,
    messages: Vec<VisibleAssistantMessage>,
    terminal: Option<String>,
    history_committed: bool,
    selected_output: Option<SelectedRunOutputIdentity>,
}

/// Per-run message-only replay and broadcast state.
///
/// The replay is intentionally not byte-capped: it contains only visible text
/// generated by the current run and remains authoritative until durable history
/// is acknowledged. Selected output may retain it until the existing runner
/// is replaced/evicted. Silently evicting it would make subscribers lose text.
#[derive(Debug)]
pub struct VisibleMessageStream {
    state: Mutex<VisibleMessageReplay>,
    sender: broadcast::Sender<VisibleMessageEvent>,
}

impl Default for VisibleMessageStream {
    fn default() -> Self {
        Self::new()
    }
}

impl VisibleMessageStream {
    const CHANNEL_CAPACITY: usize = 1000;

    pub fn new() -> Self {
        let (sender, _) = broadcast::channel(Self::CHANNEL_CAPACITY);
        Self {
            state: Mutex::new(VisibleMessageReplay::default()),
            sender,
        }
    }

    fn lock(&self) -> MutexGuard<'_, VisibleMessageReplay> {
        self.state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    fn send_locked(&self, event: VisibleMessageEvent) {
        let _ = self.sender.send(event);
    }

    /// Subscribe before taking the snapshot. Concurrent events are either in the
    /// snapshot or queued with a greater version, never neither.
    pub fn subscribe_with_snapshot(
        &self,
    ) -> (
        broadcast::Receiver<VisibleMessageEvent>,
        VisibleMessageSnapshot,
    ) {
        let receiver = self.sender.subscribe();
        let state = self.lock();
        let snapshot = VisibleMessageSnapshot {
            version: state.version,
            messages: state
                .messages
                .iter()
                .filter(|message| !message.committed && !message.content.is_empty())
                .cloned()
                .collect(),
            terminal: state.terminal.clone(),
            history_committed: state.history_committed,
        };
        (receiver, snapshot)
    }

    /// Called only from a live selected reservation with its actual Root writer,
    /// before any producer starts. The same buffer remains owned by this runner;
    /// replacement/eviction/restart never reconstructs it from session history.
    pub(crate) fn retain_selected_output(
        &self,
        owner: &bamboo_domain::RootActorRuntimeWrite,
        receipt: bamboo_domain::SessionInboxReceipt,
    ) -> std::io::Result<()> {
        let identity = SelectedRunOutputIdentity {
            session_id: owner.fence.actor_id.clone(),
            run_id: owner.fence.run_id.clone(),
            session_created_at: owner.created_at,
            receipt,
        };
        let mut state = self.lock();
        if let Some(current) = &state.selected_output {
            return if current == &identity {
                Ok(())
            } else {
                Err(std::io::Error::other(
                    "Selected output identity cannot change",
                ))
            };
        }
        if state.version != 0 || identity.receipt.generation == 0 {
            return Err(std::io::Error::other(
                "Selected output must bind before publication",
            ));
        }
        state.selected_output = Some(identity);
        Ok(())
    }

    /// Read the same visible-message buffer atomically, including text already
    /// acknowledged by history. Ordinary WS snapshots still hide that prefix.
    /// Only the owning runner's lifetime retains these bytes; no new store/TTL.
    pub fn selected_output_snapshot(&self) -> Option<SelectedRunOutputSnapshot> {
        let state = self.lock();
        Some(SelectedRunOutputSnapshot {
            identity: state.selected_output.clone()?,
            output: VisibleMessageSnapshot {
                version: state.version,
                messages: state
                    .messages
                    .iter()
                    .filter(|m| !m.content.is_empty())
                    .cloned()
                    .collect(),
                terminal: state.terminal.clone(),
                history_committed: state.history_committed,
            },
        })
    }

    pub fn begin_round(&self, round_count: u32) {
        self.lock().round_count = round_count;
    }

    /// Start one provider response. Previous segments remain replayable until an
    /// explicit discard or the durable-history barrier; a later attempt in the
    /// same round does not prove the earlier segment was rolled back.
    pub fn start(&self, message_id: String, created_at: DateTime<Utc>) {
        let mut state = self.lock();
        // A post-barrier producer reusing an id starts the same replacement
        // value an ordinary cleared stream would have; never append old text.
        if let Some(previous) = state
            .messages
            .iter()
            .position(|m| m.id == message_id && m.committed)
        {
            state.messages.remove(previous);
        }
        if !state
            .messages
            .iter()
            .any(|message| message.id == message_id)
        {
            let round_count = state.round_count;
            state.messages.push(VisibleAssistantMessage {
                id: message_id.clone(),
                content: String::new(),
                created_at,
                round_count,
                committed: false,
            });
        }
        state.active_message_id = Some(message_id.clone());
        state.history_committed = false;
        state.version = state.version.saturating_add(1);
        self.send_locked(VisibleMessageEvent {
            version: state.version,
            kind: VisibleMessageEventKind::Started {
                message_id,
                created_at,
            },
        });
    }

    /// Append visible assistant text. Legacy producers without an explicit start
    /// receive a host-minted temporary identity and reconcile on history commit.
    pub fn append(&self, content: String) {
        if content.is_empty() {
            return;
        }
        let mut state = self.lock();
        state.history_committed = false;
        let message_id = match state.active_message_id.clone() {
            Some(message_id) => message_id,
            None => {
                let message_id = Uuid::new_v4().to_string();
                let created_at = Utc::now();
                let round_count = state.round_count;
                state.messages.push(VisibleAssistantMessage {
                    id: message_id.clone(),
                    content: String::new(),
                    created_at,
                    round_count,
                    committed: false,
                });
                state.active_message_id = Some(message_id.clone());
                state.version = state.version.saturating_add(1);
                self.send_locked(VisibleMessageEvent {
                    version: state.version,
                    kind: VisibleMessageEventKind::Started {
                        message_id: message_id.clone(),
                        created_at,
                    },
                });
                message_id
            }
        };

        let Some(index) = state
            .messages
            .iter()
            .position(|message| message.id == message_id)
        else {
            return;
        };
        let offset = state.messages[index].content.len();
        let created_at = state.messages[index].created_at;
        state.messages[index].content.push_str(&content);
        state.version = state.version.saturating_add(1);
        self.send_locked(VisibleMessageEvent {
            version: state.version,
            kind: VisibleMessageEventKind::Delta {
                message_id,
                offset,
                content,
                created_at,
            },
        });
    }

    pub fn discard(&self, message_id: &str) {
        let mut state = self.lock();
        let previous_len = state.messages.len();
        state.messages.retain(|message| message.id != message_id);
        if state.messages.len() == previous_len {
            return;
        }
        if state.active_message_id.as_deref() == Some(message_id) {
            state.active_message_id = None;
        }
        state.history_committed = false;
        state.version = state.version.saturating_add(1);
        self.send_locked(VisibleMessageEvent {
            version: state.version,
            kind: VisibleMessageEventKind::Discarded {
                message_id: message_id.to_string(),
            },
        });
    }

    pub fn mark_terminal(&self, reason: impl Into<String>) {
        let mut state = self.lock();
        let reason = reason.into();
        if state.terminal.as_deref() == Some(reason.as_str()) {
            return;
        }
        state.terminal = Some(reason.clone());
        state.version = state.version.saturating_add(1);
        self.send_locked(VisibleMessageEvent {
            version: state.version,
            kind: VisibleMessageEventKind::Terminal { reason },
        });
    }

    /// Reconcile ordinary replay only after durable history publication.
    /// Selected output retains the same bytes until its runner is reclaimed.
    pub fn history_committed(&self) {
        let mut state = self.lock();
        if state.selected_output.is_some() {
            for message in &mut state.messages {
                message.committed = true;
            }
        } else {
            state.messages.clear();
        }
        state.active_message_id = None;
        state.history_committed = true;
        state.version = state.version.saturating_add(1);
        self.send_locked(VisibleMessageEvent {
            version: state.version,
            kind: VisibleMessageEventKind::HistoryCommitted,
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn subscribe_snapshot_and_live_tail_have_no_overlap_after_version_filter() {
        let stream = VisibleMessageStream::new();
        stream.begin_round(0);
        let created_at = Utc::now();
        stream.start("message-1".into(), created_at);
        stream.append("before".into());

        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
        assert_eq!(snapshot.messages[0].content, "before");
        stream.append(" after".into());

        let event = receiver.try_recv().expect("live delta");
        assert!(event.version > snapshot.version);
        assert!(matches!(
            event.kind,
            VisibleMessageEventKind::Delta { offset: 6, ref content, .. } if content == " after"
        ));
    }

    #[test]
    fn explicit_retry_rollback_discards_previous_attempt() {
        let stream = VisibleMessageStream::new();
        stream.begin_round(2);
        stream.start("attempt-1".into(), Utc::now());
        stream.append("partial".into());
        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
        assert_eq!(snapshot.messages.len(), 1);

        stream.discard("attempt-1");
        stream.start("attempt-2".into(), Utc::now());
        let discarded = receiver.try_recv().expect("discarded event");
        let started = receiver.try_recv().expect("replacement start");
        assert!(matches!(
            discarded.kind,
            VisibleMessageEventKind::Discarded { ref message_id } if message_id == "attempt-1"
        ));
        assert!(matches!(
            started.kind,
            VisibleMessageEventKind::Started { ref message_id, .. } if message_id == "attempt-2"
        ));
    }

    #[test]
    fn multi_round_messages_survive_until_history_commit() {
        let stream = VisibleMessageStream::new();
        stream.begin_round(1);
        stream.start("round-1".into(), Utc::now());
        stream.append("first".into());
        stream.begin_round(2);
        stream.start("round-2".into(), Utc::now());
        stream.append("second".into());

        let (_, snapshot) = stream.subscribe_with_snapshot();
        assert_eq!(snapshot.messages.len(), 2);
        assert_eq!(snapshot.messages[0].round_count, 1);
        assert_eq!(snapshot.messages[0].content, "first");
        assert_eq!(snapshot.messages[1].round_count, 2);
        assert_eq!(snapshot.messages[1].content, "second");
    }

    #[test]
    fn failed_and_cancelled_terminal_reasons_remain_in_late_snapshots() {
        for reason in ["error", "cancelled"] {
            let stream = VisibleMessageStream::new();
            stream.start("partial".into(), Utc::now());
            stream.append("visible prefix".into());
            stream.mark_terminal(reason);

            let (_, snapshot) = stream.subscribe_with_snapshot();
            assert_eq!(snapshot.terminal.as_deref(), Some(reason));
            assert_eq!(snapshot.messages[0].content, "visible prefix");
        }
    }

    #[test]
    fn history_commit_clears_text_only_after_control_is_emitted() {
        let stream = VisibleMessageStream::new();
        stream.start("message-1".into(), Utc::now());
        stream.append("visible".into());
        let (mut receiver, snapshot) = stream.subscribe_with_snapshot();
        assert_eq!(snapshot.messages.len(), 1);

        stream.history_committed();
        let committed = receiver.try_recv().expect("history commit control");
        assert!(matches!(
            committed.kind,
            VisibleMessageEventKind::HistoryCommitted
        ));
        let (_, settled) = stream.subscribe_with_snapshot();
        assert!(settled.messages.is_empty());
        assert!(settled.history_committed);
    }

    fn selected_owner() -> bamboo_domain::RootActorRuntimeWrite {
        bamboo_domain::RootActorRuntimeWrite {
            fence: bamboo_domain::ActorActivationFence {
                schema_version: 1,
                actor_id: "output-root".into(),
                activation_id: "activation".into(),
                attempt: 1,
                run_id: "output-run".into(),
                lease_owner: "host".into(),
                lease_epoch: 1,
            },
            created_at: Utc::now(),
        }
    }
    fn selected_receipt() -> bamboo_domain::SessionInboxReceipt {
        bamboo_domain::SessionInboxReceipt {
            id: bamboo_domain::SessionMessageId::parse("selected-input").unwrap(),
            generation: 7,
        }
    }

    #[test]
    fn run_output_retains_same_bytes_after_history_without_changing_ws_replay() {
        let stream = VisibleMessageStream::new();
        let owner = selected_owner();
        stream
            .retain_selected_output(&owner, selected_receipt())
            .unwrap();
        stream.start("visible".into(), Utc::now());
        stream.append("界\r\nfinal".into());
        stream.mark_terminal("complete");
        let before = stream.selected_output_snapshot().unwrap();
        stream.history_committed();
        let replay = stream.selected_output_snapshot().unwrap();
        assert_eq!(replay.identity.session_id, owner.fence.actor_id);
        assert_eq!(replay.identity.run_id, owner.fence.run_id);
        assert_eq!(replay.identity.session_created_at, owner.created_at);
        assert_eq!(replay.identity.receipt, selected_receipt());
        assert_eq!(
            serde_json::to_value(&replay.output.messages).unwrap(),
            serde_json::to_value(&before.output.messages).unwrap()
        );
        assert!(replay.output.history_committed && replay.output.version > before.output.version);
        assert_eq!(replay.output.terminal.as_deref(), Some("complete"));
        let (_, legacy) = stream.subscribe_with_snapshot();
        assert!(legacy.messages.is_empty());
        assert!(legacy.history_committed);
        let wire = serde_json::to_value(&replay.output).unwrap();
        assert_eq!(wire["messages"][0].as_object().unwrap().len(), 3);
        assert!(wire["messages"][0].get("committed").is_none());
    }

    #[test]
    fn run_output_rejects_rebinding_or_late_opt_in_and_defaults_to_no_receipt() {
        let stream = VisibleMessageStream::new();
        assert!(stream.selected_output_snapshot().is_none());
        let owner = selected_owner();
        stream
            .retain_selected_output(&owner, selected_receipt())
            .unwrap();
        stream
            .retain_selected_output(&owner, selected_receipt())
            .unwrap();
        let mut foreign = owner.clone();
        foreign.fence.run_id = "successor".into();
        assert!(stream
            .retain_selected_output(&foreign, selected_receipt())
            .is_err());
        let mut foreign = owner.clone();
        foreign.created_at += chrono::Duration::seconds(1);
        assert!(stream
            .retain_selected_output(&foreign, selected_receipt())
            .is_err());
        let mut receipt = selected_receipt();
        receipt.generation += 1;
        assert!(stream.retain_selected_output(&owner, receipt).is_err());
        let late = VisibleMessageStream::new();
        late.append("unattributed prefix".into());
        assert!(late
            .retain_selected_output(&owner, selected_receipt())
            .is_err());
        assert!(late.selected_output_snapshot().is_none());
    }

    #[test]
    fn run_output_later_segments_do_not_restore_committed_text_to_ws() {
        let stream = VisibleMessageStream::new();
        stream
            .retain_selected_output(&selected_owner(), selected_receipt())
            .unwrap();
        stream.start("old".into(), Utc::now());
        stream.append("committed prefix".into());
        stream.history_committed();
        stream.start("next".into(), Utc::now());
        stream.append("new segment".into());
        let (_, legacy) = stream.subscribe_with_snapshot();
        assert_eq!(legacy.messages.len(), 1);
        assert_eq!(legacy.messages[0].id, "next");
        let output = stream.selected_output_snapshot().unwrap().output;
        assert_eq!(output.messages.len(), 2);
        assert!(!output.history_committed);
        stream.discard("next");
        let (_, legacy) = stream.subscribe_with_snapshot();
        assert!(legacy.messages.is_empty());
        assert_eq!(
            stream.selected_output_snapshot().unwrap().output.messages[0].content,
            "committed prefix"
        );
    }

    #[test]
    fn run_output_reused_id_and_legacy_delta_never_append_committed_bytes() {
        let stream = VisibleMessageStream::new();
        stream
            .retain_selected_output(&selected_owner(), selected_receipt())
            .unwrap();
        stream.start("same".into(), Utc::now());
        stream.append("old".into());
        stream.history_committed();
        stream.start("same".into(), Utc::now());
        stream.append("replacement".into());
        let snapshot = stream.selected_output_snapshot().unwrap().output;
        assert_eq!(snapshot.messages.len(), 1);
        assert_eq!(snapshot.messages[0].content, "replacement");
        stream.history_committed();
        stream.append("legacy new segment".into());
        let snapshot = stream.selected_output_snapshot().unwrap().output;
        assert_eq!(snapshot.messages.len(), 2);
        assert!(!snapshot.history_committed);
        stream.history_committed();
        stream.discard("same");
        assert!(
            !stream
                .selected_output_snapshot()
                .unwrap()
                .output
                .history_committed
        );
    }

    #[test]
    fn run_output_error_and_cancellation_keep_prefix_without_inventing_commit() {
        for reason in ["error", "cancelled"] {
            let stream = VisibleMessageStream::new();
            stream
                .retain_selected_output(&selected_owner(), selected_receipt())
                .unwrap();
            stream.append("partial visible output".into());
            stream.mark_terminal(reason);
            let snapshot = stream.selected_output_snapshot().unwrap().output;
            assert_eq!(snapshot.messages[0].content, "partial visible output");
            assert_eq!(snapshot.terminal.as_deref(), Some(reason));
            assert!(!snapshot.history_committed);
            stream.history_committed();
            let snapshot = stream.selected_output_snapshot().unwrap().output;
            assert_eq!(snapshot.messages[0].content, "partial visible output");
            assert!(snapshot.history_committed);
        }
    }
}