1use super::notification::TaskNotification;
30use super::{IncomingMessage, OutputChunk};
31use std::collections::HashMap;
32use std::sync::Arc;
33use tokio::sync::{broadcast, RwLock};
34
35const MAX_OUTPUT_CHUNK_BYTES: usize = 64 * 1024;
36
37pub fn session_key(repo: &str, task_id: &str) -> String {
44 if task_id.starts_with("internal:") {
45 task_id.to_string()
46 } else {
47 format!("{repo}:{task_id}")
48 }
49}
50
51pub fn parse_session_key(key: &str) -> Option<(&str, &str)> {
55 if key.starts_with("internal:") {
56 None
59 } else {
60 let (repo, task_id) = key.split_once(':')?;
63 Some((repo, task_id))
64 }
65}
66
67#[derive(Debug, Clone)]
69pub struct SessionBinding {
70 pub tmux_session: String,
72 pub connected_threads: Vec<String>,
81 pub output_tx: broadcast::Sender<OutputChunk>,
83}
84
85pub fn conversation_key(channel: &str, thread_id: &str, topic_id: Option<&str>) -> String {
96 match topic_id {
97 Some(tid) if !tid.is_empty() => format!("{channel}:{thread_id}|{tid}"),
98 _ => format!("{channel}:{thread_id}"),
99 }
100}
101
102pub fn parse_conversation_key(key: &str) -> Option<(&str, &str, Option<&str>)> {
107 let (channel, rest) = key.split_once(':')?;
108 let (thread_id, topic_id) = match rest.split_once('|') {
109 Some((tid, topic)) => (tid, Some(topic)),
110 None => (rest, None),
111 };
112 Some((channel, thread_id, topic_id))
113}
114
115pub struct Transport {
117 bindings: Arc<RwLock<HashMap<String, SessionBinding>>>,
119 thread_to_task: Arc<RwLock<HashMap<String, String>>>,
121 notification_tx: broadcast::Sender<TaskNotification>,
123}
124
125impl Transport {
126 pub fn new() -> Self {
127 let (notification_tx, _) = broadcast::channel(64);
128 Self {
129 bindings: Arc::new(RwLock::new(HashMap::new())),
130 thread_to_task: Arc::new(RwLock::new(HashMap::new())),
131 notification_tx,
132 }
133 }
134
135 pub async fn bind(
144 &self,
145 repo: &str,
146 task_id: &str,
147 tmux_session: &str,
148 channel: &str,
149 thread_id: &str,
150 topic_id: Option<&str>,
151 ) {
152 let key = conversation_key(channel, thread_id, topic_id);
153 let skey = session_key(repo, task_id);
154 {
156 let mut bindings = self.bindings.write().await;
157 let binding = bindings.entry(skey.clone()).or_insert_with(|| {
158 let (tx, _) = broadcast::channel(256);
159 SessionBinding {
160 tmux_session: tmux_session.to_string(),
161 connected_threads: Vec::new(),
162 output_tx: tx,
163 }
164 });
165 binding.tmux_session = tmux_session.to_string();
167 if !binding.connected_threads.contains(&key) {
168 binding.connected_threads.push(key.clone());
169 }
170 } self.thread_to_task.write().await.insert(key, skey);
173 }
174
175 pub async fn subscribe(
177 &self,
178 repo: &str,
179 task_id: &str,
180 ) -> Option<broadcast::Receiver<OutputChunk>> {
181 let skey = session_key(repo, task_id);
182 let bindings = self.bindings.read().await;
183 bindings.get(&skey).map(|b| b.output_tx.subscribe())
184 }
185
186 pub async fn push_output(&self, repo: &str, task_id: &str, chunk: OutputChunk) {
188 let skey = session_key(repo, task_id);
189 if chunk.content.is_empty() {
190 let bindings = self.bindings.read().await;
191 if let Some(binding) = bindings.get(&skey) {
192 let _ = binding.output_tx.send(chunk.clone());
193 }
194 return;
195 }
196
197 let parts = split_chunks(&chunk.content, MAX_OUTPUT_CHUNK_BYTES);
198 let last_index = parts.len().saturating_sub(1);
199
200 for (idx, part) in parts.into_iter().enumerate() {
201 let part_chunk = OutputChunk {
202 content: part,
203 is_final: chunk.is_final && idx == last_index,
204 };
205 let bindings = self.bindings.read().await;
206 if let Some(binding) = bindings.get(&skey) {
207 let _ = binding.output_tx.send(part_chunk.clone());
208 }
209 }
210 }
211
212 pub async fn get_binding(&self, repo: &str, task_id: &str) -> Option<SessionBinding> {
214 let skey = session_key(repo, task_id);
215 self.bindings.read().await.get(&skey).cloned()
216 }
217
218 pub async fn route(&self, msg: &IncomingMessage) -> MessageRoute {
221 let key = conversation_key(&msg.channel, &msg.thread_id, msg.topic_id.as_deref());
222
223 if let Some(session_key) = self.thread_to_task.read().await.get(&key) {
225 return MessageRoute::TaskSession {
226 session_key: session_key.clone(),
227 };
228 }
229
230 let body = msg.body.trim();
232 if body.starts_with('/') || body.starts_with("orch ") {
233 return MessageRoute::Command {
234 raw: body.to_string(),
235 };
236 }
237
238 MessageRoute::NewTask
240 }
241
242 pub async fn unbind(&self, repo: &str, task_id: &str) {
248 let skey = session_key(repo, task_id);
249
250 let connected_threads = {
252 let mut bindings = self.bindings.write().await;
253 bindings
254 .remove(&skey)
255 .map(|b| b.connected_threads)
256 .unwrap_or_default()
257 };
258
259 if !connected_threads.is_empty() {
261 let mut t2t = self.thread_to_task.write().await;
262 for key in &connected_threads {
263 t2t.remove(key);
264 }
265 }
266 }
267
268 pub async fn active_sessions(&self) -> Vec<SessionBinding> {
270 self.bindings.read().await.values().cloned().collect()
271 }
272
273 pub fn push_notification(&self, notification: TaskNotification) {
275 let _ = self.notification_tx.send(notification);
277 }
278
279 pub fn subscribe_notifications(&self) -> broadcast::Receiver<TaskNotification> {
281 self.notification_tx.subscribe()
282 }
283}
284
285fn split_chunks(content: &str, max_bytes: usize) -> Vec<String> {
286 if content.len() <= max_bytes {
287 return vec![content.to_string()];
288 }
289
290 let mut chunks = Vec::new();
291 let mut start = 0;
292 let total = content.len();
293
294 while start < total {
295 let mut end = (start + max_bytes).min(total);
296 while end > start && !content.is_char_boundary(end) {
297 end -= 1;
298 }
299 if end == start {
300 if let Some(ch) = content[start..].chars().next() {
302 end = start + ch.len_utf8();
303 } else {
304 break;
305 }
306 }
307 chunks.push(content[start..end].to_string());
308 start = end;
309 }
310
311 chunks
312}
313
314#[derive(Debug)]
316pub enum MessageRoute {
317 TaskSession { session_key: String },
320 Command { raw: String },
322 ControlSession,
324 NewTask,
326}
327
328#[cfg(test)]
329mod tests {
330 use super::*;
331
332 #[tokio::test]
333 async fn push_notification_reaches_subscriber() {
334 let transport = Transport::new();
335 let mut rx = transport.subscribe_notifications();
336
337 let notification = TaskNotification {
338 task_id: "42".to_string(),
339 title: "Test task".to_string(),
340 status: "done".to_string(),
341 agent: "claude".to_string(),
342 duration_seconds: 60.0,
343 summary: "Completed successfully".to_string(),
344 repo: None,
345 notify_target: None,
346 pr_number: None,
347 };
348
349 transport.push_notification(notification.clone());
350
351 let received = rx.recv().await.unwrap();
352 assert_eq!(received.task_id, "42");
353 assert_eq!(received.status, "done");
354 assert_eq!(received.agent, "claude");
355 assert_eq!(received.summary, "Completed successfully");
356 }
357
358 #[tokio::test]
359 async fn push_notification_multiple_subscribers() {
360 let transport = Transport::new();
361 let mut rx1 = transport.subscribe_notifications();
362 let mut rx2 = transport.subscribe_notifications();
363
364 transport.push_notification(TaskNotification {
365 task_id: "1".to_string(),
366 title: "Task".to_string(),
367 status: "done".to_string(),
368 agent: "codex".to_string(),
369 duration_seconds: 10.0,
370 summary: "Done".to_string(),
371 repo: None,
372 notify_target: None,
373 pr_number: None,
374 });
375
376 let n1 = rx1.recv().await.unwrap();
377 let n2 = rx2.recv().await.unwrap();
378 assert_eq!(n1.task_id, "1");
379 assert_eq!(n2.task_id, "1");
380 }
381
382 #[test]
383 fn split_chunks_multibyte_at_boundary_does_not_panic() {
384 let content = "x".repeat(97) + "😀"; let chunks = split_chunks(&content, 100);
388 assert_eq!(chunks.join(""), content);
389 for c in &chunks {
390 assert!(std::str::from_utf8(c.as_bytes()).is_ok());
391 }
392 }
393
394 #[test]
395 fn split_chunks_only_emojis() {
396 let content = "😀".repeat(50); let chunks = split_chunks(&content, 10); assert_eq!(chunks.join(""), content);
400 for c in &chunks {
401 assert!(std::str::from_utf8(c.as_bytes()).is_ok());
402 }
403 }
404
405 #[test]
406 fn push_notification_no_subscribers_does_not_panic() {
407 let transport = Transport::new();
408 transport.push_notification(TaskNotification {
410 task_id: "1".to_string(),
411 title: "Task".to_string(),
412 status: "done".to_string(),
413 agent: "claude".to_string(),
414 duration_seconds: 0.0,
415 summary: "Done".to_string(),
416 repo: None,
417 notify_target: None,
418 pr_number: None,
419 });
420 }
421
422 #[tokio::test]
423 async fn bind_and_route_to_session() {
424 let transport = Transport::new();
425 transport
426 .bind(
427 "owner/repo",
428 "42",
429 "orch-myproject-42",
430 "telegram",
431 "12345",
432 None,
433 )
434 .await;
435
436 let msg = IncomingMessage {
437 channel: "telegram".to_string(),
438 id: "msg1".to_string(),
439 thread_id: "12345".to_string(),
440 author: "user".to_string(),
441 body: "hello".to_string(),
442 timestamp: chrono::Utc::now(),
443 metadata: serde_json::json!({}),
444 topic_id: None,
445 };
446
447 match transport.route(&msg).await {
448 MessageRoute::TaskSession { session_key } => {
449 assert_eq!(session_key, "owner/repo:42");
450 }
451 _ => panic!("expected TaskSession"),
452 }
453 }
454
455 #[tokio::test]
456 async fn route_command() {
457 let transport = Transport::new();
458
459 let msg = IncomingMessage {
460 channel: "telegram".to_string(),
461 id: "msg1".to_string(),
462 thread_id: "99".to_string(),
463 author: "user".to_string(),
464 body: "/status".to_string(),
465 timestamp: chrono::Utc::now(),
466 metadata: serde_json::json!({}),
467 topic_id: None,
468 };
469
470 match transport.route(&msg).await {
471 MessageRoute::Command { raw } => assert_eq!(raw, "/status"),
472 _ => panic!("expected Command"),
473 }
474 }
475
476 #[test]
479 fn session_key_external_is_prefixed_with_repo() {
480 assert_eq!(session_key("owner/repo", "42"), "owner/repo:42");
481 }
482
483 #[test]
484 fn session_key_internal_is_unchanged() {
485 assert_eq!(session_key("owner/repo", "internal:99"), "internal:99");
486 }
487
488 #[test]
489 fn session_key_same_id_different_repos_are_unique() {
490 let k1 = session_key("owner/repo-a", "42");
491 let k2 = session_key("owner/repo-b", "42");
492 assert_ne!(k1, k2);
493 assert_eq!(k1, "owner/repo-a:42");
494 assert_eq!(k2, "owner/repo-b:42");
495 }
496
497 #[tokio::test]
501 async fn same_external_id_different_repos_do_not_collide() {
502 let transport = Transport::new();
503
504 transport
506 .bind(
507 "owner/repo-a",
508 "42",
509 "orch-repo-a-42",
510 "telegram",
511 "111",
512 None,
513 )
514 .await;
515 transport
517 .bind(
518 "owner/repo-b",
519 "42",
520 "orch-repo-b-42",
521 "telegram",
522 "222",
523 None,
524 )
525 .await;
526
527 let binding_a = transport.get_binding("owner/repo-a", "42").await.unwrap();
529 let binding_b = transport.get_binding("owner/repo-b", "42").await.unwrap();
530 assert_eq!(binding_a.tmux_session, "orch-repo-a-42");
531 assert_eq!(binding_b.tmux_session, "orch-repo-b-42");
532 assert_ne!(binding_a.tmux_session, binding_b.tmux_session);
533 }
534
535 #[tokio::test]
538 async fn push_output_isolated_across_repos_with_same_task_id() {
539 let transport = Transport::new();
540
541 transport
543 .bind("owner/repo-a", "42", "orch-a-42", "cli", "stream-a", None)
544 .await;
545 let mut rx_a = transport.subscribe("owner/repo-a", "42").await.unwrap();
546
547 transport
549 .bind("owner/repo-b", "42", "orch-b-42", "cli", "stream-b", None)
550 .await;
551 let mut rx_b = transport.subscribe("owner/repo-b", "42").await.unwrap();
552
553 transport
555 .push_output(
556 "owner/repo-a",
557 "42",
558 OutputChunk {
559 content: "hello from repo-a".to_string(),
560 is_final: false,
561 },
562 )
563 .await;
564
565 let chunk_a = tokio::time::timeout(std::time::Duration::from_millis(100), rx_a.recv())
567 .await
568 .unwrap()
569 .unwrap();
570 assert_eq!(chunk_a.content, "hello from repo-a");
571
572 let result_b =
574 tokio::time::timeout(std::time::Duration::from_millis(100), rx_b.recv()).await;
575 assert!(
576 result_b.is_err(),
577 "repo-b should not receive repo-a's output"
578 );
579 }
580
581 #[test]
584 fn conversation_key_without_topic() {
585 let key = conversation_key("telegram", "12345", None);
586 assert_eq!(key, "telegram:12345");
587 }
588
589 #[test]
590 fn conversation_key_with_topic() {
591 let key = conversation_key("telegram", "12345", Some("678"));
592 assert_eq!(key, "telegram:12345|678");
593 }
594
595 #[test]
596 fn conversation_key_empty_topic_treated_as_none() {
597 let key_empty = conversation_key("discord", "abc", Some(""));
599 let key_none = conversation_key("discord", "abc", None);
600 assert_eq!(key_empty, key_none);
601 }
602
603 #[test]
606 fn parse_key_without_topic() {
607 let (channel, thread_id, topic_id) =
608 parse_conversation_key("telegram:12345").expect("valid key");
609 assert_eq!(channel, "telegram");
610 assert_eq!(thread_id, "12345");
611 assert!(topic_id.is_none());
612 }
613
614 #[test]
615 fn parse_key_with_topic() {
616 let (channel, thread_id, topic_id) =
617 parse_conversation_key("telegram:12345|678").expect("valid key");
618 assert_eq!(channel, "telegram");
619 assert_eq!(thread_id, "12345");
620 assert_eq!(topic_id, Some("678"));
621 }
622
623 #[test]
624 fn parse_key_malformed_returns_none() {
625 assert!(parse_conversation_key("no-colon-here").is_none());
626 }
627
628 #[test]
629 fn conversation_key_round_trips() {
630 for (ch, tid, topic) in [
632 ("telegram", "111", Some("222")),
633 ("discord", "abc", None),
634 ("slack", "x", Some("y")),
635 ] {
636 let key = conversation_key(ch, tid, topic);
637 let (ch2, tid2, topic2) = parse_conversation_key(&key).expect("round-trip valid");
638 assert_eq!(ch2, ch);
639 assert_eq!(tid2, tid);
640 assert_eq!(topic2, topic);
641 }
642 }
643
644 #[tokio::test]
649 async fn different_topics_do_not_collide() {
650 let transport = Transport::new();
651
652 transport
654 .bind(
655 "owner/repo",
656 "10",
657 "orch-proj-10",
658 "telegram",
659 "999",
660 Some("100"),
661 )
662 .await;
663 transport
665 .bind(
666 "owner/repo",
667 "20",
668 "orch-proj-20",
669 "telegram",
670 "999",
671 Some("200"),
672 )
673 .await;
674
675 let msg_topic_100 = IncomingMessage {
676 channel: "telegram".to_string(),
677 id: "m1".to_string(),
678 thread_id: "999".to_string(),
679 author: "user".to_string(),
680 body: "reply for task 10".to_string(),
681 timestamp: chrono::Utc::now(),
682 metadata: serde_json::json!({}),
683 topic_id: Some("100".to_string()),
684 };
685
686 let msg_topic_200 = IncomingMessage {
687 channel: "telegram".to_string(),
688 id: "m2".to_string(),
689 thread_id: "999".to_string(),
690 author: "user".to_string(),
691 body: "reply for task 20".to_string(),
692 timestamp: chrono::Utc::now(),
693 metadata: serde_json::json!({}),
694 topic_id: Some("200".to_string()),
695 };
696
697 match transport.route(&msg_topic_100).await {
698 MessageRoute::TaskSession { session_key } => {
699 assert_eq!(session_key, "owner/repo:10");
700 }
701 other => panic!("expected TaskSession for task 10, got {other:?}"),
702 }
703
704 match transport.route(&msg_topic_200).await {
705 MessageRoute::TaskSession { session_key } => {
706 assert_eq!(session_key, "owner/repo:20");
707 }
708 other => panic!("expected TaskSession for task 20, got {other:?}"),
709 }
710 }
711
712 #[tokio::test]
716 async fn unbind_removes_binding_and_reverse_lookup() {
717 let transport = Transport::new();
718 transport
719 .bind("owner/repo", "99", "orch-proj-99", "telegram", "555", None)
720 .await;
721
722 assert!(transport.get_binding("owner/repo", "99").await.is_some());
724
725 transport.unbind("owner/repo", "99").await;
727
728 assert!(transport.get_binding("owner/repo", "99").await.is_none());
730
731 let msg = IncomingMessage {
733 channel: "telegram".to_string(),
734 id: "m1".to_string(),
735 thread_id: "555".to_string(),
736 author: "user".to_string(),
737 body: "hello".to_string(),
738 timestamp: chrono::Utc::now(),
739 metadata: serde_json::json!({}),
740 topic_id: None,
741 };
742 match transport.route(&msg).await {
743 MessageRoute::NewTask => {} other => panic!("expected NewTask after unbind, got {other:?}"),
745 }
746 }
747
748 #[tokio::test]
750 async fn unbind_nonexistent_does_not_panic() {
751 let transport = Transport::new();
752 transport.unbind("owner/repo", "nonexistent").await;
754 }
755
756 #[tokio::test]
759 async fn wrong_topic_does_not_route_to_bound_task() {
760 let transport = Transport::new();
761
762 transport
764 .bind(
765 "owner/repo",
766 "5",
767 "orch-proj-5",
768 "telegram",
769 "42",
770 Some("10"),
771 )
772 .await;
773
774 let msg = IncomingMessage {
776 channel: "telegram".to_string(),
777 id: "m1".to_string(),
778 thread_id: "42".to_string(),
779 author: "user".to_string(),
780 body: "hello from a different topic".to_string(),
781 timestamp: chrono::Utc::now(),
782 metadata: serde_json::json!({}),
783 topic_id: Some("99".to_string()),
784 };
785
786 match transport.route(&msg).await {
787 MessageRoute::NewTask => {} MessageRoute::TaskSession { session_key } => {
789 panic!("should not have routed to session {session_key}")
790 }
791 other => panic!("unexpected route result: {other:?}"),
792 }
793 }
794}