vtcode_webmcp/
event_hub.rs1use 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#[derive(Debug, Clone, Copy)]
17pub struct EventHubConfig {
18 pub replay_capacity: usize,
20 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#[derive(Clone)]
54pub struct WebmcpEventHub {
55 config: EventHubConfig,
56 max_event_bytes: usize,
57 state: Arc<Mutex<HubState>>,
58}
59
60#[derive(Debug, Clone)]
62pub struct SequencedThreadEvent {
63 pub sequence: u64,
65 pub event: VersionedThreadEvent,
67}
68
69pub 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 pub fn new(config: EventHubConfig) -> Result<Self> {
81 Self::new_with_max_event_bytes(config, MAX_EVENT_BYTES)
82 }
83
84 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 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 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 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 pub fn replay(&self) -> &[SequencedThreadEvent] {
202 &self.replay
203 }
204
205 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}