use super::notification::TaskNotification;
use super::{IncomingMessage, OutputChunk};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
const MAX_OUTPUT_CHUNK_BYTES: usize = 64 * 1024;
pub fn session_key(repo: &str, task_id: &str) -> String {
if task_id.starts_with("internal:") {
task_id.to_string()
} else {
format!("{repo}:{task_id}")
}
}
pub fn parse_session_key(key: &str) -> Option<(&str, &str)> {
if key.starts_with("internal:") {
None
} else {
let (repo, task_id) = key.split_once(':')?;
Some((repo, task_id))
}
}
#[derive(Debug, Clone)]
pub struct SessionBinding {
pub tmux_session: String,
pub connected_threads: Vec<String>,
pub output_tx: broadcast::Sender<OutputChunk>,
}
pub fn conversation_key(channel: &str, thread_id: &str, topic_id: Option<&str>) -> String {
match topic_id {
Some(tid) if !tid.is_empty() => format!("{channel}:{thread_id}|{tid}"),
_ => format!("{channel}:{thread_id}"),
}
}
pub fn parse_conversation_key(key: &str) -> Option<(&str, &str, Option<&str>)> {
let (channel, rest) = key.split_once(':')?;
let (thread_id, topic_id) = match rest.split_once('|') {
Some((tid, topic)) => (tid, Some(topic)),
None => (rest, None),
};
Some((channel, thread_id, topic_id))
}
pub struct Transport {
bindings: Arc<RwLock<HashMap<String, SessionBinding>>>,
thread_to_task: Arc<RwLock<HashMap<String, String>>>,
notification_tx: broadcast::Sender<TaskNotification>,
}
impl Transport {
pub fn new() -> Self {
let (notification_tx, _) = broadcast::channel(64);
Self {
bindings: Arc::new(RwLock::new(HashMap::new())),
thread_to_task: Arc::new(RwLock::new(HashMap::new())),
notification_tx,
}
}
pub async fn bind(
&self,
repo: &str,
task_id: &str,
tmux_session: &str,
channel: &str,
thread_id: &str,
topic_id: Option<&str>,
) {
let key = conversation_key(channel, thread_id, topic_id);
let skey = session_key(repo, task_id);
{
let mut bindings = self.bindings.write().await;
let binding = bindings.entry(skey.clone()).or_insert_with(|| {
let (tx, _) = broadcast::channel(256);
SessionBinding {
tmux_session: tmux_session.to_string(),
connected_threads: Vec::new(),
output_tx: tx,
}
});
binding.tmux_session = tmux_session.to_string();
if !binding.connected_threads.contains(&key) {
binding.connected_threads.push(key.clone());
}
} self.thread_to_task.write().await.insert(key, skey);
}
pub async fn subscribe(
&self,
repo: &str,
task_id: &str,
) -> Option<broadcast::Receiver<OutputChunk>> {
let skey = session_key(repo, task_id);
let bindings = self.bindings.read().await;
bindings.get(&skey).map(|b| b.output_tx.subscribe())
}
pub async fn push_output(&self, repo: &str, task_id: &str, chunk: OutputChunk) {
let skey = session_key(repo, task_id);
if chunk.content.is_empty() {
let bindings = self.bindings.read().await;
if let Some(binding) = bindings.get(&skey) {
let _ = binding.output_tx.send(chunk.clone());
}
return;
}
let parts = split_chunks(&chunk.content, MAX_OUTPUT_CHUNK_BYTES);
let last_index = parts.len().saturating_sub(1);
for (idx, part) in parts.into_iter().enumerate() {
let part_chunk = OutputChunk {
content: part,
is_final: chunk.is_final && idx == last_index,
};
let bindings = self.bindings.read().await;
if let Some(binding) = bindings.get(&skey) {
let _ = binding.output_tx.send(part_chunk.clone());
}
}
}
pub async fn get_binding(&self, repo: &str, task_id: &str) -> Option<SessionBinding> {
let skey = session_key(repo, task_id);
self.bindings.read().await.get(&skey).cloned()
}
pub async fn route(&self, msg: &IncomingMessage) -> MessageRoute {
let key = conversation_key(&msg.channel, &msg.thread_id, msg.topic_id.as_deref());
if let Some(session_key) = self.thread_to_task.read().await.get(&key) {
return MessageRoute::TaskSession {
session_key: session_key.clone(),
};
}
let body = msg.body.trim();
if body.starts_with('/') || body.starts_with("orch ") {
return MessageRoute::Command {
raw: body.to_string(),
};
}
MessageRoute::NewTask
}
pub async fn unbind(&self, repo: &str, task_id: &str) {
let skey = session_key(repo, task_id);
let connected_threads = {
let mut bindings = self.bindings.write().await;
bindings
.remove(&skey)
.map(|b| b.connected_threads)
.unwrap_or_default()
};
if !connected_threads.is_empty() {
let mut t2t = self.thread_to_task.write().await;
for key in &connected_threads {
t2t.remove(key);
}
}
}
pub async fn active_sessions(&self) -> Vec<SessionBinding> {
self.bindings.read().await.values().cloned().collect()
}
pub fn push_notification(&self, notification: TaskNotification) {
let _ = self.notification_tx.send(notification);
}
pub fn subscribe_notifications(&self) -> broadcast::Receiver<TaskNotification> {
self.notification_tx.subscribe()
}
}
fn split_chunks(content: &str, max_bytes: usize) -> Vec<String> {
if content.len() <= max_bytes {
return vec![content.to_string()];
}
let mut chunks = Vec::new();
let mut start = 0;
let total = content.len();
while start < total {
let mut end = (start + max_bytes).min(total);
while end > start && !content.is_char_boundary(end) {
end -= 1;
}
if end == start {
if let Some(ch) = content[start..].chars().next() {
end = start + ch.len_utf8();
} else {
break;
}
}
chunks.push(content[start..end].to_string());
start = end;
}
chunks
}
#[derive(Debug)]
pub enum MessageRoute {
TaskSession { session_key: String },
Command { raw: String },
ControlSession,
NewTask,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn push_notification_reaches_subscriber() {
let transport = Transport::new();
let mut rx = transport.subscribe_notifications();
let notification = TaskNotification {
task_id: "42".to_string(),
title: "Test task".to_string(),
status: "done".to_string(),
agent: "claude".to_string(),
duration_seconds: 60.0,
summary: "Completed successfully".to_string(),
repo: None,
notify_target: None,
pr_number: None,
};
transport.push_notification(notification.clone());
let received = rx.recv().await.unwrap();
assert_eq!(received.task_id, "42");
assert_eq!(received.status, "done");
assert_eq!(received.agent, "claude");
assert_eq!(received.summary, "Completed successfully");
}
#[tokio::test]
async fn push_notification_multiple_subscribers() {
let transport = Transport::new();
let mut rx1 = transport.subscribe_notifications();
let mut rx2 = transport.subscribe_notifications();
transport.push_notification(TaskNotification {
task_id: "1".to_string(),
title: "Task".to_string(),
status: "done".to_string(),
agent: "codex".to_string(),
duration_seconds: 10.0,
summary: "Done".to_string(),
repo: None,
notify_target: None,
pr_number: None,
});
let n1 = rx1.recv().await.unwrap();
let n2 = rx2.recv().await.unwrap();
assert_eq!(n1.task_id, "1");
assert_eq!(n2.task_id, "1");
}
#[test]
fn split_chunks_multibyte_at_boundary_does_not_panic() {
let content = "x".repeat(97) + "😀"; let chunks = split_chunks(&content, 100);
assert_eq!(chunks.join(""), content);
for c in &chunks {
assert!(std::str::from_utf8(c.as_bytes()).is_ok());
}
}
#[test]
fn split_chunks_only_emojis() {
let content = "😀".repeat(50); let chunks = split_chunks(&content, 10); assert_eq!(chunks.join(""), content);
for c in &chunks {
assert!(std::str::from_utf8(c.as_bytes()).is_ok());
}
}
#[test]
fn push_notification_no_subscribers_does_not_panic() {
let transport = Transport::new();
transport.push_notification(TaskNotification {
task_id: "1".to_string(),
title: "Task".to_string(),
status: "done".to_string(),
agent: "claude".to_string(),
duration_seconds: 0.0,
summary: "Done".to_string(),
repo: None,
notify_target: None,
pr_number: None,
});
}
#[tokio::test]
async fn bind_and_route_to_session() {
let transport = Transport::new();
transport
.bind(
"owner/repo",
"42",
"orch-myproject-42",
"telegram",
"12345",
None,
)
.await;
let msg = IncomingMessage {
channel: "telegram".to_string(),
id: "msg1".to_string(),
thread_id: "12345".to_string(),
author: "user".to_string(),
body: "hello".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: None,
};
match transport.route(&msg).await {
MessageRoute::TaskSession { session_key } => {
assert_eq!(session_key, "owner/repo:42");
}
_ => panic!("expected TaskSession"),
}
}
#[tokio::test]
async fn route_command() {
let transport = Transport::new();
let msg = IncomingMessage {
channel: "telegram".to_string(),
id: "msg1".to_string(),
thread_id: "99".to_string(),
author: "user".to_string(),
body: "/status".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: None,
};
match transport.route(&msg).await {
MessageRoute::Command { raw } => assert_eq!(raw, "/status"),
_ => panic!("expected Command"),
}
}
#[test]
fn session_key_external_is_prefixed_with_repo() {
assert_eq!(session_key("owner/repo", "42"), "owner/repo:42");
}
#[test]
fn session_key_internal_is_unchanged() {
assert_eq!(session_key("owner/repo", "internal:99"), "internal:99");
}
#[test]
fn session_key_same_id_different_repos_are_unique() {
let k1 = session_key("owner/repo-a", "42");
let k2 = session_key("owner/repo-b", "42");
assert_ne!(k1, k2);
assert_eq!(k1, "owner/repo-a:42");
assert_eq!(k2, "owner/repo-b:42");
}
#[tokio::test]
async fn same_external_id_different_repos_do_not_collide() {
let transport = Transport::new();
transport
.bind(
"owner/repo-a",
"42",
"orch-repo-a-42",
"telegram",
"111",
None,
)
.await;
transport
.bind(
"owner/repo-b",
"42",
"orch-repo-b-42",
"telegram",
"222",
None,
)
.await;
let binding_a = transport.get_binding("owner/repo-a", "42").await.unwrap();
let binding_b = transport.get_binding("owner/repo-b", "42").await.unwrap();
assert_eq!(binding_a.tmux_session, "orch-repo-a-42");
assert_eq!(binding_b.tmux_session, "orch-repo-b-42");
assert_ne!(binding_a.tmux_session, binding_b.tmux_session);
}
#[tokio::test]
async fn push_output_isolated_across_repos_with_same_task_id() {
let transport = Transport::new();
transport
.bind("owner/repo-a", "42", "orch-a-42", "cli", "stream-a", None)
.await;
let mut rx_a = transport.subscribe("owner/repo-a", "42").await.unwrap();
transport
.bind("owner/repo-b", "42", "orch-b-42", "cli", "stream-b", None)
.await;
let mut rx_b = transport.subscribe("owner/repo-b", "42").await.unwrap();
transport
.push_output(
"owner/repo-a",
"42",
OutputChunk {
content: "hello from repo-a".to_string(),
is_final: false,
},
)
.await;
let chunk_a = tokio::time::timeout(std::time::Duration::from_millis(100), rx_a.recv())
.await
.unwrap()
.unwrap();
assert_eq!(chunk_a.content, "hello from repo-a");
let result_b =
tokio::time::timeout(std::time::Duration::from_millis(100), rx_b.recv()).await;
assert!(
result_b.is_err(),
"repo-b should not receive repo-a's output"
);
}
#[test]
fn conversation_key_without_topic() {
let key = conversation_key("telegram", "12345", None);
assert_eq!(key, "telegram:12345");
}
#[test]
fn conversation_key_with_topic() {
let key = conversation_key("telegram", "12345", Some("678"));
assert_eq!(key, "telegram:12345|678");
}
#[test]
fn conversation_key_empty_topic_treated_as_none() {
let key_empty = conversation_key("discord", "abc", Some(""));
let key_none = conversation_key("discord", "abc", None);
assert_eq!(key_empty, key_none);
}
#[test]
fn parse_key_without_topic() {
let (channel, thread_id, topic_id) =
parse_conversation_key("telegram:12345").expect("valid key");
assert_eq!(channel, "telegram");
assert_eq!(thread_id, "12345");
assert!(topic_id.is_none());
}
#[test]
fn parse_key_with_topic() {
let (channel, thread_id, topic_id) =
parse_conversation_key("telegram:12345|678").expect("valid key");
assert_eq!(channel, "telegram");
assert_eq!(thread_id, "12345");
assert_eq!(topic_id, Some("678"));
}
#[test]
fn parse_key_malformed_returns_none() {
assert!(parse_conversation_key("no-colon-here").is_none());
}
#[test]
fn conversation_key_round_trips() {
for (ch, tid, topic) in [
("telegram", "111", Some("222")),
("discord", "abc", None),
("slack", "x", Some("y")),
] {
let key = conversation_key(ch, tid, topic);
let (ch2, tid2, topic2) = parse_conversation_key(&key).expect("round-trip valid");
assert_eq!(ch2, ch);
assert_eq!(tid2, tid);
assert_eq!(topic2, topic);
}
}
#[tokio::test]
async fn different_topics_do_not_collide() {
let transport = Transport::new();
transport
.bind(
"owner/repo",
"10",
"orch-proj-10",
"telegram",
"999",
Some("100"),
)
.await;
transport
.bind(
"owner/repo",
"20",
"orch-proj-20",
"telegram",
"999",
Some("200"),
)
.await;
let msg_topic_100 = IncomingMessage {
channel: "telegram".to_string(),
id: "m1".to_string(),
thread_id: "999".to_string(),
author: "user".to_string(),
body: "reply for task 10".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: Some("100".to_string()),
};
let msg_topic_200 = IncomingMessage {
channel: "telegram".to_string(),
id: "m2".to_string(),
thread_id: "999".to_string(),
author: "user".to_string(),
body: "reply for task 20".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: Some("200".to_string()),
};
match transport.route(&msg_topic_100).await {
MessageRoute::TaskSession { session_key } => {
assert_eq!(session_key, "owner/repo:10");
}
other => panic!("expected TaskSession for task 10, got {other:?}"),
}
match transport.route(&msg_topic_200).await {
MessageRoute::TaskSession { session_key } => {
assert_eq!(session_key, "owner/repo:20");
}
other => panic!("expected TaskSession for task 20, got {other:?}"),
}
}
#[tokio::test]
async fn unbind_removes_binding_and_reverse_lookup() {
let transport = Transport::new();
transport
.bind("owner/repo", "99", "orch-proj-99", "telegram", "555", None)
.await;
assert!(transport.get_binding("owner/repo", "99").await.is_some());
transport.unbind("owner/repo", "99").await;
assert!(transport.get_binding("owner/repo", "99").await.is_none());
let msg = IncomingMessage {
channel: "telegram".to_string(),
id: "m1".to_string(),
thread_id: "555".to_string(),
author: "user".to_string(),
body: "hello".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: None,
};
match transport.route(&msg).await {
MessageRoute::NewTask => {} other => panic!("expected NewTask after unbind, got {other:?}"),
}
}
#[tokio::test]
async fn unbind_nonexistent_does_not_panic() {
let transport = Transport::new();
transport.unbind("owner/repo", "nonexistent").await;
}
#[tokio::test]
async fn wrong_topic_does_not_route_to_bound_task() {
let transport = Transport::new();
transport
.bind(
"owner/repo",
"5",
"orch-proj-5",
"telegram",
"42",
Some("10"),
)
.await;
let msg = IncomingMessage {
channel: "telegram".to_string(),
id: "m1".to_string(),
thread_id: "42".to_string(),
author: "user".to_string(),
body: "hello from a different topic".to_string(),
timestamp: chrono::Utc::now(),
metadata: serde_json::json!({}),
topic_id: Some("99".to_string()),
};
match transport.route(&msg).await {
MessageRoute::NewTask => {} MessageRoute::TaskSession { session_key } => {
panic!("should not have routed to session {session_key}")
}
other => panic!("unexpected route result: {other:?}"),
}
}
}