Skip to main content

vtcode_webmcp/
event_hub.rs

1use crate::error::{Result, WebmcpError};
2use std::collections::{HashMap, VecDeque};
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, Mutex};
5use tokio::sync::mpsc;
6use vtcode_exec_events::VersionedThreadEvent;
7
8pub(crate) const MAX_EVENT_BYTES: usize = 512 * 1024;
9const MAX_REPLAY_BYTES: usize = 8 * 1024 * 1024;
10const MAX_REPLAY_CAPACITY: usize = 4096;
11const MAX_SUBSCRIBER_CAPACITY: usize = 1024;
12const MAX_SUBSCRIBERS: usize = 1024;
13const MAX_SUBSCRIBER_BYTES: usize = 8 * 1024 * 1024;
14
15/// Event hub retention and subscriber queue limits.
16#[derive(Debug, Clone, Copy)]
17pub struct EventHubConfig {
18    /// Number of events retained for reconnect replay.
19    pub replay_capacity: usize,
20    /// Number of events buffered per connected browser.
21    pub subscriber_capacity: usize,
22}
23
24impl Default for EventHubConfig {
25    fn default() -> Self {
26        Self { replay_capacity: 256, subscriber_capacity: 64 }
27    }
28}
29
30#[derive(Debug, Clone)]
31struct HubEvent {
32    sequence: u64,
33    event: VersionedThreadEvent,
34    size_bytes: usize,
35}
36
37#[derive(Debug)]
38struct HubState {
39    next_sequence: u64,
40    replay: VecDeque<HubEvent>,
41    replay_bytes: usize,
42    subscribers: HashMap<u64, Subscriber>,
43    next_subscriber_id: u64,
44}
45
46#[derive(Debug)]
47struct Subscriber {
48    sender: mpsc::Sender<HubEvent>,
49    queued_bytes: Arc<AtomicUsize>,
50}
51
52/// Bounded event hub that retains canonical VT Code runtime events.
53#[derive(Clone)]
54pub struct WebmcpEventHub {
55    config: EventHubConfig,
56    max_event_bytes: usize,
57    state: Arc<Mutex<HubState>>,
58}
59
60/// A sequenced event returned during replay or live subscription.
61#[derive(Debug, Clone)]
62pub struct SequencedThreadEvent {
63    /// Monotonic bridge sequence number.
64    pub sequence: u64,
65    /// Canonical versioned runtime event.
66    pub event: VersionedThreadEvent,
67}
68
69/// A browser event subscription with a replay prefix.
70pub struct EventHubSubscription {
71    replay: Vec<SequencedThreadEvent>,
72    receiver: mpsc::Receiver<HubEvent>,
73    state: Arc<Mutex<HubState>>,
74    subscriber_id: u64,
75    queued_bytes: Arc<AtomicUsize>,
76}
77
78impl WebmcpEventHub {
79    /// Creates a bounded event hub.
80    pub fn new(config: EventHubConfig) -> Result<Self> {
81        Self::new_with_max_event_bytes(config, MAX_EVENT_BYTES)
82    }
83
84    /// Creates a bounded event hub with a smaller event payload limit.
85    pub fn new_with_max_event_bytes(config: EventHubConfig, max_event_bytes: usize) -> Result<Self> {
86        if config.replay_capacity == 0
87            || config.subscriber_capacity == 0
88            || config.replay_capacity > MAX_REPLAY_CAPACITY
89            || config.subscriber_capacity > MAX_SUBSCRIBER_CAPACITY
90            || max_event_bytes == 0
91            || max_event_bytes > MAX_EVENT_BYTES
92        {
93            return Err(WebmcpError::InvalidRequest(
94                "event hub capacities are outside the supported bounds".to_string(),
95            ));
96        }
97        Ok(Self {
98            config,
99            max_event_bytes,
100            state: Arc::new(Mutex::new(HubState {
101                next_sequence: 1,
102                replay: VecDeque::with_capacity(config.replay_capacity),
103                replay_bytes: 0,
104                subscribers: HashMap::new(),
105                next_subscriber_id: 1,
106            })),
107        })
108    }
109
110    /// Publishes a canonical runtime event and returns its bridge sequence.
111    pub fn publish(&self, event: VersionedThreadEvent) -> Result<u64> {
112        let event_bytes = serde_json::to_vec(&event)?;
113        if event_bytes.len() > self.max_event_bytes {
114            return Err(WebmcpError::LimitExceeded);
115        }
116        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
117        let sequence = state.next_sequence;
118        state.next_sequence = state.next_sequence.checked_add(1).ok_or(WebmcpError::LimitExceeded)?;
119        let hub_event = HubEvent { sequence, event, size_bytes: event_bytes.len() };
120        state.replay.push_back(hub_event.clone());
121        state.replay_bytes = state.replay_bytes.saturating_add(hub_event.size_bytes);
122        while state.replay.len() > self.config.replay_capacity || state.replay_bytes > MAX_REPLAY_BYTES {
123            if let Some(removed) = state.replay.pop_front() {
124                state.replay_bytes = state.replay_bytes.saturating_sub(removed.size_bytes);
125            }
126        }
127
128        let mut slow_subscribers = Vec::new();
129        for (subscriber_id, subscriber) in &state.subscribers {
130            let queued_bytes = subscriber.queued_bytes.load(Ordering::Relaxed);
131            if queued_bytes.saturating_add(hub_event.size_bytes) > MAX_SUBSCRIBER_BYTES {
132                slow_subscribers.push(*subscriber_id);
133                continue;
134            }
135            let _ = subscriber.queued_bytes.fetch_add(hub_event.size_bytes, Ordering::Relaxed);
136            if subscriber.sender.try_send(hub_event.clone()).is_err() {
137                let _ = subscriber.queued_bytes.fetch_sub(hub_event.size_bytes, Ordering::Relaxed);
138                slow_subscribers.push(*subscriber_id);
139            }
140        }
141        for subscriber_id in slow_subscribers {
142            drop(state.subscribers.remove(&subscriber_id));
143        }
144        Ok(sequence)
145    }
146
147    /// Subscribes after a sequence, returning retained events before live events.
148    pub fn subscribe(&self, after_sequence: Option<u64>) -> Result<EventHubSubscription> {
149        let hub_state = Arc::clone(&self.state);
150        let mut state = hub_state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
151        let replay = match after_sequence {
152            None => Vec::new(),
153            Some(requested) => {
154                if let Some(oldest) = state.replay.front().map(|event| event.sequence)
155                    && requested.saturating_add(1) < oldest
156                {
157                    return Err(WebmcpError::SequenceGap { requested, oldest });
158                }
159                state
160                    .replay
161                    .iter()
162                    .filter(|event| event.sequence > requested)
163                    .map(|event| SequencedThreadEvent {
164                        sequence: event.sequence,
165                        event: event.event.clone(),
166                    })
167                    .collect()
168            }
169        };
170        let (sender, receiver) = mpsc::channel(self.config.subscriber_capacity);
171        if state.subscribers.len() >= MAX_SUBSCRIBERS {
172            return Err(WebmcpError::LimitExceeded);
173        }
174        let subscriber_id = state.next_subscriber_id;
175        state.next_subscriber_id = state.next_subscriber_id.checked_add(1).ok_or(WebmcpError::LimitExceeded)?;
176        let queued_bytes = Arc::new(AtomicUsize::new(0));
177        drop(
178            state
179                .subscribers
180                .insert(subscriber_id, Subscriber { sender, queued_bytes: Arc::clone(&queued_bytes) }),
181        );
182        drop(state);
183        Ok(EventHubSubscription {
184            replay,
185            receiver,
186            state: hub_state,
187            subscriber_id,
188            queued_bytes,
189        })
190    }
191
192    /// Returns the latest assigned sequence, or zero before the first event.
193    pub fn latest_sequence(&self) -> u64 {
194        let state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
195        state.next_sequence.saturating_sub(1)
196    }
197}
198
199impl EventHubSubscription {
200    /// Returns retained replay events in sequence order.
201    pub fn replay(&self) -> &[SequencedThreadEvent] {
202        &self.replay
203    }
204
205    /// Waits for the next live event. A closed receiver means the client was
206    /// removed because it could not keep up. The subscription owns the state
207    /// needed to receive retained and live events after the hub is dropped.
208    pub async fn recv(&mut self) -> Option<SequencedThreadEvent> {
209        let event = self.receiver.recv().await?;
210        let _ = self.queued_bytes.fetch_sub(event.size_bytes, Ordering::Relaxed);
211        Some(SequencedThreadEvent { sequence: event.sequence, event: event.event })
212    }
213}
214
215impl Drop for EventHubSubscription {
216    fn drop(&mut self) {
217        let mut state = self.state.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
218        drop(state.subscribers.remove(&self.subscriber_id));
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use vtcode_exec_events::{ThreadEvent, ThreadStartedEvent, VersionedThreadEvent};
226
227    fn event(id: &str) -> ThreadEvent {
228        ThreadEvent::ThreadStarted(ThreadStartedEvent { thread_id: id.to_string() })
229    }
230
231    #[tokio::test]
232    async fn replays_events_and_reports_old_sequence_gaps() {
233        let hub = WebmcpEventHub::new(EventHubConfig { replay_capacity: 2, subscriber_capacity: 2 }).expect("hub");
234        let _ = hub.publish(VersionedThreadEvent::new(event("one"))).expect("publish");
235        let _ = hub.publish(VersionedThreadEvent::new(event("two"))).expect("publish");
236        let _ = hub.publish(VersionedThreadEvent::new(event("three"))).expect("publish");
237
238        assert!(hub.subscribe(None).expect("fresh subscription").replay().is_empty());
239        assert!(matches!(hub.subscribe(Some(0)), Err(WebmcpError::SequenceGap { requested: 0, oldest: 2 })));
240        let subscription = hub.subscribe(Some(1)).expect("replay");
241        assert_eq!(subscription.replay().len(), 2);
242        assert_eq!(subscription.replay()[0].sequence, 2);
243    }
244
245    #[tokio::test]
246    async fn slow_subscriber_is_closed_instead_of_dropping_silently() {
247        let hub = WebmcpEventHub::new(EventHubConfig { replay_capacity: 4, subscriber_capacity: 1 }).expect("hub");
248        let mut subscription = hub.subscribe(None).expect("subscription");
249        let _ = hub.publish(VersionedThreadEvent::new(event("one"))).expect("publish");
250        let _ = hub.publish(VersionedThreadEvent::new(event("two"))).expect("publish");
251        assert_eq!(subscription.recv().await.expect("first event").sequence, 1);
252        assert!(subscription.recv().await.is_none());
253    }
254
255    #[test]
256    fn rejects_unbounded_queue_configuration() {
257        assert!(
258            WebmcpEventHub::new(EventHubConfig {
259                replay_capacity: MAX_REPLAY_CAPACITY + 1,
260                subscriber_capacity: 1
261            })
262            .is_err()
263        );
264        assert!(
265            WebmcpEventHub::new(EventHubConfig {
266                replay_capacity: 1,
267                subscriber_capacity: MAX_SUBSCRIBER_CAPACITY + 1
268            })
269            .is_err()
270        );
271    }
272
273    #[test]
274    fn rejects_subscriber_overflow() {
275        let hub = WebmcpEventHub::new(EventHubConfig::default()).expect("hub");
276        let mut subscriptions = Vec::with_capacity(MAX_SUBSCRIBERS);
277        for _ in 0..MAX_SUBSCRIBERS {
278            subscriptions.push(hub.subscribe(None).expect("subscriber"));
279        }
280        assert!(matches!(hub.subscribe(None), Err(WebmcpError::LimitExceeded)));
281    }
282}