macp_runtime/
stream_bus.rs

1use crate::pb::Envelope;
2use std::collections::HashMap;
3use std::sync::Mutex;
4use tokio::sync::broadcast;
5
6const DEFAULT_SESSION_STREAM_CAPACITY: usize = 256;
7
8pub struct SessionStreamBus {
9    channels: Mutex<HashMap<String, broadcast::Sender<Envelope>>>,
10    capacity: usize,
11}
12
13impl Default for SessionStreamBus {
14    fn default() -> Self {
15        Self::new(DEFAULT_SESSION_STREAM_CAPACITY)
16    }
17}
18
19impl SessionStreamBus {
20    pub fn new(capacity: usize) -> Self {
21        Self {
22            channels: Mutex::new(HashMap::new()),
23            capacity,
24        }
25    }
26
27    pub fn subscribe(&self, session_id: &str) -> broadcast::Receiver<Envelope> {
28        let mut guard = self.channels.lock().unwrap_or_else(|e| e.into_inner());
29        guard
30            .entry(session_id.to_string())
31            .or_insert_with(|| {
32                let (sender, _receiver) = broadcast::channel(self.capacity);
33                sender
34            })
35            .subscribe()
36    }
37
38    pub fn publish(&self, session_id: &str, envelope: Envelope) {
39        let sender = {
40            let guard = self.channels.lock().unwrap_or_else(|e| e.into_inner());
41            guard.get(session_id).cloned()
42        };
43        if let Some(sender) = sender {
44            let _ = sender.send(envelope);
45        }
46    }
47
48    /// Remove a session's broadcast channel if it has no live receivers.
49    /// Channels are created lazily on subscribe and were previously never
50    /// removed — every session ever streamed pinned a map entry (and a
51    /// lagging receiver pins up to `capacity` buffered envelopes) for the
52    /// process lifetime. Called from session eviction; a channel with active
53    /// receivers is left in place (`false`) and retried on the next sweep.
54    pub fn remove_if_unused(&self, session_id: &str) -> bool {
55        let mut guard = self.channels.lock().unwrap_or_else(|e| e.into_inner());
56        match guard.get(session_id) {
57            Some(sender) if sender.receiver_count() == 0 => {
58                guard.remove(session_id);
59                true
60            }
61            Some(_) => false,
62            None => true,
63        }
64    }
65
66    /// Number of live channels (observability / tests).
67    pub fn channel_count(&self) -> usize {
68        self.channels
69            .lock()
70            .unwrap_or_else(|e| e.into_inner())
71            .len()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    fn env(message_id: &str) -> Envelope {
80        Envelope {
81            macp_version: "1.0".into(),
82            mode: "macp.mode.decision.v1".into(),
83            message_type: "Proposal".into(),
84            message_id: message_id.into(),
85            session_id: "s1".into(),
86            sender: "agent://sender".into(),
87            timestamp_unix_ms: 1,
88            payload: vec![],
89        }
90    }
91
92    #[test]
93    fn subscribe_then_publish_round_trip() {
94        let bus = SessionStreamBus::default();
95        let mut rx = bus.subscribe("s1");
96        bus.publish("s1", env("m1"));
97        let envelope = rx.try_recv().expect("stream event");
98        assert_eq!(envelope.message_id, "m1");
99    }
100
101    #[test]
102    fn remove_if_unused_respects_live_receivers() {
103        let bus = SessionStreamBus::default();
104        let rx = bus.subscribe("s1");
105        assert_eq!(bus.channel_count(), 1);
106
107        // Live receiver: channel must survive eviction attempts.
108        assert!(!bus.remove_if_unused("s1"));
109        assert_eq!(bus.channel_count(), 1);
110
111        // Receiver dropped: channel is removable.
112        drop(rx);
113        assert!(bus.remove_if_unused("s1"));
114        assert_eq!(bus.channel_count(), 0);
115
116        // Unknown session is trivially "removed".
117        assert!(bus.remove_if_unused("nope"));
118    }
119}
macp_runtime/stream_bus.rs

macp_runtime/
stream_bus.rs
