1use futures::{SinkExt, StreamExt};
4use serde::{Deserialize, Serialize};
5use tokio::net::TcpListener;
6use tokio::sync::broadcast;
7use tokio_tungstenite::tungstenite::handshake::server::Request;
8
9#[derive(Clone, Debug, Default)]
11pub struct EventFilter {
12 pub task_id: Option<String>,
13 pub repo: Option<String>,
14}
15
16impl EventFilter {
17 fn from_request(request: &Request) -> Self {
19 let uri = request.uri();
20 let Some(query) = uri.query() else {
21 return Self::default();
22 };
23
24 let mut task_id = None;
25 let mut repo = None;
26
27 for pair in query.split('&') {
28 let Some((key, value)) = pair.split_once('=') else {
29 continue;
30 };
31 let decoded = urlencoding::decode(value).unwrap_or(std::borrow::Cow::Borrowed(value));
32 match key {
33 "task_id" => task_id = Some(decoded.into_owned()),
34 "repo" => repo = Some(decoded.into_owned()),
35 _ => {}
36 }
37 }
38
39 Self { task_id, repo }
40 }
41
42 fn matches(&self, event: &TaskEvent) -> bool {
44 if let Some(ref task_id) = self.task_id {
45 if event.task_id != *task_id {
46 return false;
47 }
48 }
49 if let Some(ref repo) = self.repo {
50 if event.repo != *repo {
51 return false;
52 }
53 }
54 true
55 }
56}
57
58#[derive(Clone, Debug, Serialize, Deserialize)]
64pub struct ControlMessage {
65 pub kind: String,
67 pub message: String,
69 pub missed: u64,
71}
72
73#[derive(Clone, Debug, Serialize, Deserialize)]
75pub struct TaskEvent {
76 pub task_id: String,
77 pub repo: String,
78 pub old_status: String,
79 pub new_status: String,
80 pub agent: Option<String>,
81 pub model: Option<String>,
82 pub pr_number: Option<String>,
83 pub branch: Option<String>,
84 pub review_context: Option<String>,
85 pub error: Option<String>,
86 pub timestamp: String,
87 pub title: Option<String>,
88 pub summary: Option<String>,
89 pub duration_seconds: Option<f64>,
90}
91
92pub struct EventBus {
94 tx: broadcast::Sender<TaskEvent>,
95}
96
97impl EventBus {
98 pub fn new(capacity: usize) -> Self {
99 let (tx, _) = broadcast::channel(capacity);
100 Self { tx }
101 }
102
103 pub fn sender(&self) -> broadcast::Sender<TaskEvent> {
105 self.tx.clone()
106 }
107
108 pub fn subscribe(&self) -> broadcast::Receiver<TaskEvent> {
110 self.tx.subscribe()
111 }
112
113 #[cfg_attr(not(test), allow(dead_code))]
115 pub fn publish(&self, event: TaskEvent) -> usize {
116 self.tx.send(event).unwrap_or(0)
117 }
118
119 pub async fn start_ws_server(&self) -> anyhow::Result<u16> {
122 let (listener, port) = select_available_listener()?;
125 listener.set_nonblocking(true)?;
126 let listener = TcpListener::from_std(listener)?;
127
128 let port_path = ws_port_path().await?;
130 if let Some(parent) = port_path.parent() {
131 tokio::fs::create_dir_all(parent).await?;
132 }
133 tokio::fs::write(&port_path, port.to_string()).await?;
134
135 if let Ok(path) = crate::home::service_version_path() {
138 if let Some(parent) = path.parent() {
139 let _ = tokio::fs::create_dir_all(parent).await;
140 }
141 let pid = std::process::id();
142 let content = format!("{}\t{}", pid, env!("ORCH_VERSION"));
143 let _ = tokio::fs::write(path, content).await;
144 }
145
146 let tx = self.tx.clone();
147 tokio::spawn(async move {
148 loop {
149 match listener.accept().await {
150 Ok((stream, _addr)) => {
151 let mut rx = tx.subscribe();
152 tokio::spawn(async move {
153 let mut filter = EventFilter::default();
154 #[allow(clippy::result_large_err)]
155 let callback = |req: &Request, resp| {
156 filter = EventFilter::from_request(req);
157 tracing::debug!(?filter, "ws client connected with filter");
158 Ok(resp)
159 };
160 let Ok(ws) =
161 tokio_tungstenite::accept_hdr_async(stream, callback).await
162 else {
163 return;
164 };
165 let (mut write, mut read) = ws.split();
166 loop {
167 tokio::select! {
168 event = rx.recv() => {
169 match event {
170 Ok(e) => {
171 if !filter.matches(&e) {
172 continue;
173 }
174 let json = serde_json::to_string(&e)
175 .unwrap_or_default();
176 let msg = tokio_tungstenite::tungstenite::Message::Text(json);
177 if write.send(msg).await.is_err() {
178 break;
179 }
180 }
181 Err(broadcast::error::RecvError::Lagged(n)) => {
182 tracing::warn!(
183 missed = n,
184 "event bus ws client lagged — {n} events dropped"
185 );
186 let ctrl = ControlMessage {
187 kind: "lagged".to_string(),
188 message: format!(
189 "event bus lagged: {n} events were dropped — resync from store"
190 ),
191 missed: n,
192 };
193 if let Ok(json) = serde_json::to_string(&ctrl) {
194 let msg = tokio_tungstenite::tungstenite::Message::Text(json);
195 if write.send(msg).await.is_err() {
196 break;
197 }
198 }
199 }
200 Err(_) => break,
201 }
202 }
203 msg = read.next() => {
204 if msg.is_none() {
206 break;
207 }
208 }
209 }
210 }
211 });
212 }
213 Err(e) => {
214 tracing::warn!(?e, "ws accept failed");
215 }
216 }
217 }
218 });
219
220 tracing::info!(port, "event bus websocket server started on 127.0.0.1");
221 Ok(port)
222 }
223}
224
225pub fn select_available_listener() -> anyhow::Result<(std::net::TcpListener, u16)> {
229 let hostname = std::env::var("HOSTNAME").unwrap_or_else(|_| "orch".to_string());
230 let hash = hostname
231 .bytes()
232 .fold(0u32, |acc, b| acc.wrapping_add(b as u32));
233 let range = 65535u32 - 49152;
234 let start = 49152u16 + (hash % range) as u16;
235
236 for offset in 0..1000u16 {
237 let port = start.wrapping_add(offset);
238 if port < 49152 {
239 continue;
240 }
241 if let Ok(listener) = std::net::TcpListener::bind(("127.0.0.1", port)) {
242 return Ok((listener, port));
243 }
244 }
245 anyhow::bail!("no available port found in range 49152-65535")
246}
247
248pub async fn cleanup_port_file() {
250 if let Ok(path) = ws_port_path().await {
251 let _ = tokio::fs::remove_file(path).await;
252 }
253}
254
255pub fn cleanup_version_file() {
257 if let Ok(path) = crate::home::service_version_path() {
258 let owned = std::fs::read_to_string(&path)
259 .ok()
260 .and_then(|s| s.split('\t').next()?.parse::<u32>().ok())
261 .map(|pid| pid == std::process::id())
262 .unwrap_or(true); if owned {
264 let _ = std::fs::remove_file(path);
265 }
266 }
267}
268
269async fn ws_port_path() -> anyhow::Result<std::path::PathBuf> {
270 if cfg!(test) {
271 Ok(std::env::temp_dir().join("orch-test-state").join("ws.port"))
272 } else {
273 Ok(crate::home::state_dir()?.join("ws.port"))
274 }
275}
276
277#[cfg(test)]
278mod tests {
279 use super::*;
280
281 #[test]
282 fn control_message_serializes_to_json() {
283 let ctrl = ControlMessage {
284 kind: "lagged".to_string(),
285 message: "event bus lagged: 5 events were dropped".to_string(),
286 missed: 5,
287 };
288 let json = serde_json::to_string(&ctrl).unwrap();
289 assert!(json.contains("\"kind\":\"lagged\""));
290 assert!(json.contains("\"missed\":5"));
291 }
292
293 #[test]
294 fn control_message_round_trips() {
295 let ctrl = ControlMessage {
296 kind: "lagged".to_string(),
297 message: "test".to_string(),
298 missed: 42,
299 };
300 let json = serde_json::to_string(&ctrl).unwrap();
301 let decoded: ControlMessage = serde_json::from_str(&json).unwrap();
302 assert_eq!(decoded.kind, "lagged");
303 assert_eq!(decoded.missed, 42);
304 }
305
306 #[test]
307 fn task_event_serializes_to_json() {
308 let event = TaskEvent {
309 task_id: "123".to_string(),
310 repo: "owner/repo".to_string(),
311 old_status: "new".to_string(),
312 new_status: "routed".to_string(),
313 agent: Some("claude".to_string()),
314 model: None,
315 pr_number: None,
316 branch: None,
317 review_context: None,
318 error: None,
319 timestamp: "2026-03-23T12:00:00Z".to_string(),
320 title: None,
321 summary: None,
322 duration_seconds: None,
323 };
324 let json = serde_json::to_string(&event).unwrap();
325 assert!(json.contains("\"task_id\":\"123\""));
326 assert!(json.contains("\"new_status\":\"routed\""));
327 }
328
329 #[test]
330 fn event_bus_send_receive() {
331 let bus = EventBus::new(256);
332 let mut rx = bus.subscribe();
333 let event = TaskEvent {
334 task_id: "1".to_string(),
335 repo: "r".to_string(),
336 old_status: "new".to_string(),
337 new_status: "routed".to_string(),
338 agent: None,
339 model: None,
340 pr_number: None,
341 branch: None,
342 review_context: None,
343 error: None,
344 timestamp: "2026-03-23T12:00:00Z".to_string(),
345 title: None,
346 summary: None,
347 duration_seconds: None,
348 };
349 bus.publish(event.clone());
350 let received = rx.try_recv().unwrap();
351 assert_eq!(received.task_id, "1");
352 assert_eq!(received.new_status, "routed");
353 }
354
355 #[test]
356 fn select_port_finds_available_port() {
357 let (listener, port) = select_available_listener().unwrap();
358 assert!(port >= 49152);
359 let second_bind = std::net::TcpListener::bind(("127.0.0.1", port));
361 assert!(second_bind.is_err());
362 drop(listener);
363 }
364
365 fn make_event(task_id: &str, repo: &str) -> TaskEvent {
366 TaskEvent {
367 task_id: task_id.to_string(),
368 repo: repo.to_string(),
369 old_status: "new".to_string(),
370 new_status: "routed".to_string(),
371 agent: None,
372 model: None,
373 pr_number: None,
374 branch: None,
375 review_context: None,
376 error: None,
377 timestamp: "2026-03-23T12:00:00Z".to_string(),
378 title: None,
379 summary: None,
380 duration_seconds: None,
381 }
382 }
383
384 #[test]
385 fn event_filter_empty_matches_all() {
386 let filter = EventFilter::default();
387 assert!(filter.matches(&make_event("1", "a/b")));
388 assert!(filter.matches(&make_event("999", "x/y")));
389 }
390
391 #[test]
392 fn event_filter_task_id_filters() {
393 let filter = EventFilter {
394 task_id: Some("42".to_string()),
395 repo: None,
396 };
397 assert!(filter.matches(&make_event("42", "a/b")));
398 assert!(!filter.matches(&make_event("43", "a/b")));
399 }
400
401 #[test]
402 fn event_filter_repo_filters() {
403 let filter = EventFilter {
404 task_id: None,
405 repo: Some("owner/repo".to_string()),
406 };
407 assert!(filter.matches(&make_event("1", "owner/repo")));
408 assert!(!filter.matches(&make_event("1", "other/repo")));
409 }
410
411 #[test]
412 fn event_filter_repo_exact_match() {
413 let filter = EventFilter {
414 task_id: None,
415 repo: Some("owner".to_string()),
416 };
417 assert!(filter.matches(&make_event("1", "owner")));
418 assert!(!filter.matches(&make_event("1", "owner/repo")));
419 assert!(!filter.matches(&make_event("1", "someone/else")));
420 }
421
422 #[test]
423 fn event_filter_combined() {
424 let filter = EventFilter {
425 task_id: Some("5".to_string()),
426 repo: Some("a/b".to_string()),
427 };
428 assert!(filter.matches(&make_event("5", "a/b")));
429 assert!(!filter.matches(&make_event("6", "a/b")));
430 assert!(!filter.matches(&make_event("5", "c/d")));
431 }
432
433 #[tokio::test]
434 async fn ws_server_broadcasts_events() {
435 let bus = EventBus::new(256);
436 let port = bus.start_ws_server().await.unwrap();
437
438 let url = format!("ws://127.0.0.1:{}", port);
440 let (ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
441 let (_write, mut read) = ws.split();
442
443 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
445
446 let mut event_with_agent = make_event("test-1", "owner/repo");
448 event_with_agent.agent = Some("claude".to_string());
449 bus.publish(event_with_agent);
450
451 let msg = tokio::time::timeout(
453 std::time::Duration::from_secs(2),
454 StreamExt::next(&mut read),
455 )
456 .await
457 .unwrap()
458 .unwrap()
459 .unwrap();
460
461 if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
462 let received: TaskEvent = serde_json::from_str(&text).unwrap();
463 assert_eq!(received.task_id, "test-1");
464 assert_eq!(received.new_status, "routed");
465 assert_eq!(received.agent, Some("claude".to_string()));
466 } else {
467 panic!("expected text message");
468 }
469
470 cleanup_port_file().await;
472 }
473
474 #[tokio::test]
475 async fn ws_filter_by_task_id() {
476 let bus = EventBus::new(256);
477 let port = bus.start_ws_server().await.unwrap();
478
479 let url = format!("ws://127.0.0.1:{port}/events?task_id=42");
481 let (ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
482 let (_write, mut read) = ws.split();
483
484 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
485
486 bus.publish(make_event("42", "owner/repo"));
488 bus.publish(make_event("99", "owner/repo"));
490
491 let msg = tokio::time::timeout(
492 std::time::Duration::from_secs(2),
493 StreamExt::next(&mut read),
494 )
495 .await
496 .unwrap()
497 .unwrap()
498 .unwrap();
499
500 if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
501 let received: TaskEvent = serde_json::from_str(&text).unwrap();
502 assert_eq!(received.task_id, "42");
503 } else {
504 panic!("expected text message");
505 }
506
507 let result = tokio::time::timeout(
509 std::time::Duration::from_millis(300),
510 StreamExt::next(&mut read),
511 )
512 .await;
513 assert!(result.is_err(), "expected no more messages (filtered)");
514
515 cleanup_port_file().await;
516 }
517
518 #[tokio::test]
519 async fn ws_filter_by_repo() {
520 let bus = EventBus::new(256);
521 let port = bus.start_ws_server().await.unwrap();
522
523 let url = format!(
525 "ws://127.0.0.1:{port}/events?repo={}",
526 urlencoding::encode("gabrielkoerich/orch")
527 );
528 let (ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
529 let (_write, mut read) = ws.split();
530
531 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
532
533 bus.publish(make_event("1", "gabrielkoerich/orch"));
535 bus.publish(make_event("2", "other/repo"));
537
538 let msg = tokio::time::timeout(
539 std::time::Duration::from_secs(2),
540 StreamExt::next(&mut read),
541 )
542 .await
543 .unwrap()
544 .unwrap()
545 .unwrap();
546
547 if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
548 let received: TaskEvent = serde_json::from_str(&text).unwrap();
549 assert_eq!(received.task_id, "1");
550 assert_eq!(received.repo, "gabrielkoerich/orch");
551 } else {
552 panic!("expected text message");
553 }
554
555 let result = tokio::time::timeout(
557 std::time::Duration::from_millis(300),
558 StreamExt::next(&mut read),
559 )
560 .await;
561 assert!(result.is_err(), "expected no more messages (filtered)");
562
563 cleanup_port_file().await;
564 }
565
566 #[tokio::test]
567 async fn ws_unfiltered_client_receives_all() {
568 let bus = EventBus::new(256);
569 let port = bus.start_ws_server().await.unwrap();
570
571 let url = format!("ws://127.0.0.1:{port}/events");
573 let (ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
574 let (_write, mut read) = ws.split();
575
576 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
577
578 bus.publish(make_event("1", "a/b"));
579 bus.publish(make_event("2", "c/d"));
580 bus.publish(make_event("3", "e/f"));
581
582 let mut received_ids = Vec::new();
583 for _ in 0..3 {
584 let msg = tokio::time::timeout(
585 std::time::Duration::from_secs(2),
586 StreamExt::next(&mut read),
587 )
588 .await
589 .unwrap()
590 .unwrap()
591 .unwrap();
592
593 if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
594 let event: TaskEvent = serde_json::from_str(&text).unwrap();
595 received_ids.push(event.task_id);
596 }
597 }
598
599 assert_eq!(received_ids, vec!["1", "2", "3"]);
600
601 cleanup_port_file().await;
602 }
603
604 #[tokio::test]
607 async fn ws_server_sends_lag_control_message() {
608 let bus = EventBus::new(1);
611 let port = match bus.start_ws_server().await {
612 Ok(p) => p,
613 Err(_) => {
614 return;
616 }
617 };
618
619 let url = format!("ws://127.0.0.1:{port}/events");
621 let (ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap();
622 let (_write, mut read) = ws.split();
623
624 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
626
627 for i in 0..20 {
630 bus.publish(make_event(&i.to_string(), "owner/repo"));
631 }
632
633 let mut got_lag = false;
635 for _ in 0..30 {
636 let result = tokio::time::timeout(
637 std::time::Duration::from_secs(2),
638 StreamExt::next(&mut read),
639 )
640 .await;
641
642 let Ok(Some(Ok(msg))) = result else { break };
643
644 if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
645 if let Ok(ctrl) = serde_json::from_str::<ControlMessage>(&text) {
646 if ctrl.kind == "lagged" {
647 got_lag = true;
648 assert!(ctrl.missed > 0, "missed count must be > 0");
649 break;
650 }
651 }
652 }
653 }
654
655 assert!(
656 got_lag,
657 "expected a lagged ControlMessage when bus capacity is exceeded"
658 );
659
660 cleanup_port_file().await;
661 }
662}