bamboo-engine 2026.10.10

//! 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 and is cleared only after the durable history barrier is published.

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

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

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

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

/// 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,
}

/// 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. Silently evicting it would make late 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.content.is_empty())
                .cloned()
                .collect(),
            terminal: state.terminal.clone(),
            history_committed: state.history_committed,
        };
        (receiver, snapshot)
    }

    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();
        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,
            });
        }
        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();
        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,
                });
                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.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 },
        });
    }

    /// Clear replay only after durable history has actually been published.
    pub fn history_committed(&self) {
        let mut state = self.lock();
        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);
    }
}