Skip to main content

codetether_agent/bus/
relay.rs

1//! Protocol-first relay runtime for multi-agent handoff flows.
2//!
3//! This layer keeps orchestration transport-agnostic while using the local
4//! `AgentBus` as the default protocol transport. It provides:
5//! - Agent lifecycle registration/deregistration
6//! - Message handoff publication with correlation ids
7//! - Relay topic emission for observability (`relay.{relay_id}`)
8
9use super::{AgentBus, BusMessage};
10use crate::a2a::types::Part;
11use std::sync::Arc;
12use uuid::Uuid;
13
14/// Profile for a relay-participating agent.
15#[derive(Debug, Clone)]
16pub struct RelayAgentProfile {
17    pub name: String,
18    pub capabilities: Vec<String>,
19}
20
21/// Protocol runtime used by auto-chat and other relay-style orchestrators.
22#[derive(Clone)]
23pub struct ProtocolRelayRuntime {
24    bus: Arc<AgentBus>,
25    relay_id: String,
26}
27
28impl ProtocolRelayRuntime {
29    /// Build a relay runtime with an auto-generated relay id.
30    pub fn new(bus: Arc<AgentBus>) -> Self {
31        Self {
32            bus,
33            relay_id: format!("relay-{}", &Uuid::new_v4().to_string()[..8]),
34        }
35    }
36
37    /// Build a relay runtime with an explicit id.
38    pub fn with_relay_id(bus: Arc<AgentBus>, relay_id: impl Into<String>) -> Self {
39        Self {
40            bus,
41            relay_id: relay_id.into(),
42        }
43    }
44
45    /// The relay id used for correlation and topic publication.
46    pub fn relay_id(&self) -> &str {
47        &self.relay_id
48    }
49
50    /// Register relay agents on the protocol bus.
51    pub fn register_agents(&self, agents: &[RelayAgentProfile]) {
52        for agent in agents {
53            let mut capabilities = agent.capabilities.clone();
54            if !capabilities.iter().any(|c| c == "relay") {
55                capabilities.push("relay".to_string());
56            }
57            capabilities.push(format!("relay:{}", self.relay_id));
58
59            let handle = self.bus.handle(agent.name.clone());
60            handle.announce_ready(capabilities);
61        }
62    }
63
64    /// Deregister relay agents from the protocol bus.
65    pub fn shutdown_agents(&self, agent_ids: &[String]) {
66        for agent_id in agent_ids {
67            let handle = self.bus.handle(agent_id.clone());
68            handle.announce_shutdown();
69        }
70    }
71
72    /// Send one protocol handoff from `from` to `to`.
73    ///
74    /// Publishes to both `agent.{to}` and `relay.{relay_id}` so downstream
75    /// observers (TUI bus log, metrics) can trace the conversation flow.
76    pub fn send_handoff(&self, from: &str, to: &str, text: &str) {
77        let payload = BusMessage::AgentMessage {
78            from: from.to_string(),
79            to: to.to_string(),
80            parts: vec![Part::Text {
81                text: text.to_string(),
82            }],
83        };
84
85        let correlation = Some(format!("{}:{}", self.relay_id, Uuid::new_v4()));
86        let handle = self.bus.handle(from.to_string());
87
88        handle.send_with_correlation(format!("agent.{to}"), payload.clone(), correlation.clone());
89        handle.send_with_correlation(format!("relay.{}", self.relay_id), payload, correlation);
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use tokio::time::{Duration, timeout};
97
98    #[tokio::test]
99    async fn test_register_and_shutdown_agents() {
100        let bus = AgentBus::new().into_arc();
101        let relay = ProtocolRelayRuntime::with_relay_id(bus.clone(), "relay-test");
102
103        let agents = vec![
104            RelayAgentProfile {
105                name: "auto-planner".to_string(),
106                capabilities: vec!["planning".to_string()],
107            },
108            RelayAgentProfile {
109                name: "auto-coder".to_string(),
110                capabilities: vec!["coding".to_string()],
111            },
112        ];
113
114        relay.register_agents(&agents);
115
116        assert!(bus.registry.get("auto-planner").is_some());
117        assert!(bus.registry.get("auto-coder").is_some());
118
119        relay.shutdown_agents(&["auto-planner".to_string(), "auto-coder".to_string()]);
120
121        assert!(bus.registry.get("auto-planner").is_none());
122        assert!(bus.registry.get("auto-coder").is_none());
123    }
124
125    #[tokio::test]
126    async fn test_send_handoff_emits_agent_and_relay_topics() {
127        let bus = AgentBus::new().into_arc();
128        let relay = ProtocolRelayRuntime::with_relay_id(bus.clone(), "relay-test");
129        let mut observer = bus.handle("observer");
130
131        relay.send_handoff("auto-planner", "auto-coder", "handoff payload");
132
133        let first = timeout(Duration::from_millis(200), observer.recv())
134            .await
135            .expect("first envelope timeout")
136            .expect("first envelope missing");
137        let second = timeout(Duration::from_millis(200), observer.recv())
138            .await
139            .expect("second envelope timeout")
140            .expect("second envelope missing");
141
142        let topics = [first.topic, second.topic];
143        assert!(topics.iter().any(|t| t == "agent.auto-coder"));
144        assert!(topics.iter().any(|t| t == "relay.relay-test"));
145    }
146}