codetether_agent/bus/
relay.rs1use super::{AgentBus, BusMessage};
10use crate::a2a::types::Part;
11use std::sync::Arc;
12use uuid::Uuid;
13
14#[derive(Debug, Clone)]
16pub struct RelayAgentProfile {
17 pub name: String,
18 pub capabilities: Vec<String>,
19}
20
21#[derive(Clone)]
23pub struct ProtocolRelayRuntime {
24 bus: Arc<AgentBus>,
25 relay_id: String,
26}
27
28impl ProtocolRelayRuntime {
29 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 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 pub fn relay_id(&self) -> &str {
47 &self.relay_id
48 }
49
50 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 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 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}