Skip to main content

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
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    fn env(message_id: &str) -> Envelope {
54        Envelope {
55            macp_version: "1.0".into(),
56            mode: "macp.mode.decision.v1".into(),
57            message_type: "Proposal".into(),
58            message_id: message_id.into(),
59            session_id: "s1".into(),
60            sender: "agent://sender".into(),
61            timestamp_unix_ms: 1,
62            payload: vec![],
63        }
64    }
65
66    #[test]
67    fn subscribe_then_publish_round_trip() {
68        let bus = SessionStreamBus::default();
69        let mut rx = bus.subscribe("s1");
70        bus.publish("s1", env("m1"));
71        let envelope = rx.try_recv().expect("stream event");
72        assert_eq!(envelope.message_id, "m1");
73    }
74}