lean_ctx/core/a2a/
relay.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3
4const DEFAULT_MAX_HOPS: usize = 5;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct RelayHop {
8    pub agent_id: String,
9    pub received_at: DateTime<Utc>,
10    pub forwarded_at: Option<DateTime<Utc>>,
11    pub processing_ms: Option<u64>,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct RelayChain {
16    pub hops: Vec<RelayHop>,
17    pub max_hops: usize,
18}
19
20impl RelayChain {
21    pub fn new(max_hops: usize) -> Self {
22        Self {
23            hops: Vec::new(),
24            max_hops,
25        }
26    }
27
28    pub fn add_hop(&mut self, hop: RelayHop) -> Result<(), RelayError> {
29        if self.hops.len() >= self.max_hops {
30            return Err(RelayError::MaxHopsExceeded(self.max_hops));
31        }
32        if self.contains_agent(&hop.agent_id) {
33            return Err(RelayError::CycleDetected(hop.agent_id));
34        }
35
36        self.hops.push(hop);
37        Ok(())
38    }
39
40    pub fn depth(&self) -> usize {
41        self.hops.len()
42    }
43
44    pub fn total_latency_ms(&self) -> u64 {
45        self.hops
46            .iter()
47            .filter_map(|hop| hop.processing_ms)
48            .fold(0, u64::saturating_add)
49    }
50
51    pub fn contains_agent(&self, agent_id: &str) -> bool {
52        self.hops.iter().any(|hop| hop.agent_id == agent_id)
53    }
54
55    pub fn origin(&self) -> Option<&str> {
56        self.hops.first().map(|hop| hop.agent_id.as_str())
57    }
58}
59
60impl Default for RelayChain {
61    fn default() -> Self {
62        Self::new(DEFAULT_MAX_HOPS)
63    }
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, thiserror::Error)]
67pub enum RelayError {
68    #[error("relay chain exceeded maximum of {0} hops")]
69    MaxHopsExceeded(usize),
70    #[error("relay cycle detected at agent {0}")]
71    CycleDetected(String),
72}
73
74#[cfg(test)]
75mod tests {
76    use chrono::{TimeZone, Utc};
77
78    use super::{RelayChain, RelayError, RelayHop};
79
80    fn hop(agent_id: &str, processing_ms: Option<u64>) -> RelayHop {
81        RelayHop {
82            agent_id: agent_id.to_owned(),
83            received_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
84            forwarded_at: None,
85            processing_ms,
86        }
87    }
88
89    #[test]
90    fn adds_hop_and_reports_chain_metadata() {
91        let mut chain = RelayChain::new(3);
92
93        chain.add_hop(hop("origin-agent", Some(4))).unwrap();
94
95        assert_eq!(chain.depth(), 1);
96        assert!(chain.contains_agent("origin-agent"));
97        assert_eq!(chain.origin(), Some("origin-agent"));
98    }
99
100    #[test]
101    fn rejects_cycle() {
102        let mut chain = RelayChain::new(3);
103        chain.add_hop(hop("relay-agent", None)).unwrap();
104
105        let error = chain.add_hop(hop("relay-agent", Some(1))).unwrap_err();
106
107        assert_eq!(error, RelayError::CycleDetected("relay-agent".to_owned()));
108        assert_eq!(chain.depth(), 1);
109    }
110
111    #[test]
112    fn rejects_hop_beyond_maximum() {
113        let mut chain = RelayChain::new(1);
114        chain.add_hop(hop("origin-agent", None)).unwrap();
115
116        let error = chain.add_hop(hop("next-agent", None)).unwrap_err();
117
118        assert_eq!(error, RelayError::MaxHopsExceeded(1));
119        assert_eq!(chain.depth(), 1);
120    }
121
122    #[test]
123    fn sums_available_processing_latency() {
124        let mut chain = RelayChain::new(3);
125        chain.add_hop(hop("origin-agent", Some(12))).unwrap();
126        chain.add_hop(hop("relay-agent", None)).unwrap();
127        chain.add_hop(hop("recipient-agent", Some(8))).unwrap();
128
129        assert_eq!(chain.total_latency_ms(), 20);
130    }
131
132    #[test]
133    fn default_chain_allows_five_hops() {
134        let mut chain = RelayChain::default();
135        for index in 0..5 {
136            chain
137                .add_hop(hop(&format!("agent-{index}"), Some(u64::MAX)))
138                .unwrap();
139        }
140
141        assert_eq!(chain.depth(), 5);
142        assert_eq!(chain.total_latency_ms(), u64::MAX);
143        assert_eq!(
144            chain.add_hop(hop("sixth-agent", None)),
145            Err(RelayError::MaxHopsExceeded(5))
146        );
147    }
148}
lean_ctx/core/a2a/relay.rs

lean_ctx/core/a2a/
relay.rs
