Skip to main content

lean_ctx/core/
a2a_transport.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5const MAX_ENVELOPE_BYTES: usize = 2_000_000;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct AgentIdentityV1 {
9    pub agent_id: String,
10    pub agent_type: String,
11    pub daemon_fingerprint: String,
12    pub capabilities: Vec<String>,
13}
14
15impl AgentIdentityV1 {
16    pub fn from_current(agent_id: &str, agent_type: &str) -> Self {
17        Self {
18            agent_id: agent_id.to_string(),
19            agent_type: agent_type.to_string(),
20            daemon_fingerprint: compute_daemon_fingerprint(),
21            capabilities: default_capabilities(),
22        }
23    }
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct TransportEnvelopeV1 {
28    pub format_version: u32,
29    pub sent_at: DateTime<Utc>,
30    pub sender: AgentIdentityV1,
31    pub recipient: Option<String>,
32    pub content_type: TransportContentType,
33    pub payload_json: String,
34    pub signature: Option<String>,
35    pub metadata: HashMap<String, String>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
39#[serde(rename_all = "snake_case")]
40pub enum TransportContentType {
41    HandoffBundle,
42    ContextPackage,
43    A2AMessage,
44    A2ATask,
45}
46
47impl TransportEnvelopeV1 {
48    pub fn new(
49        sender: AgentIdentityV1,
50        recipient: Option<&str>,
51        content_type: TransportContentType,
52        payload_json: String,
53    ) -> Self {
54        Self {
55            format_version: 1,
56            sent_at: Utc::now(),
57            sender,
58            recipient: recipient.map(std::string::ToString::to_string),
59            content_type,
60            payload_json,
61            signature: None,
62            metadata: HashMap::new(),
63        }
64    }
65
66    pub fn sign(&mut self, secret: &[u8]) {
67        use hmac::{Hmac, Mac};
68        use sha2::Sha256;
69
70        let header = format!(
71            "{}:{}:{}:{}",
72            self.format_version,
73            self.sender.agent_id,
74            self.content_type_str(),
75            self.payload_json.len()
76        );
77        let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("HMAC accepts any key length");
78        mac.update(header.as_bytes());
79        mac.update(b"\0");
80        mac.update(self.payload_json.as_bytes());
81        let result = mac.finalize().into_bytes();
82        let mut sig = String::with_capacity(result.len() * 2);
83        for b in &result {
84            use std::fmt::Write;
85            let _ = write!(sig, "{b:02x}");
86        }
87        self.signature = Some(sig);
88    }
89
90    pub fn verify_signature(&self, secret: &[u8]) -> bool {
91        let Some(ref sig) = self.signature else {
92            return false;
93        };
94        use hmac::{Hmac, Mac};
95        use sha2::Sha256;
96
97        let header = format!(
98            "{}:{}:{}:{}",
99            self.format_version,
100            self.sender.agent_id,
101            self.content_type_str(),
102            self.payload_json.len()
103        );
104        let mut mac = Hmac::<Sha256>::new_from_slice(secret).expect("HMAC accepts any key length");
105        mac.update(header.as_bytes());
106        mac.update(b"\0");
107        mac.update(self.payload_json.as_bytes());
108
109        let expected: Vec<u8> = (0..sig.len())
110            .step_by(2)
111            .filter_map(|i| u8::from_str_radix(&sig[i..i + 2], 16).ok())
112            .collect();
113        if expected.len() != sig.len() / 2 {
114            return false;
115        }
116        mac.verify_slice(&expected).is_ok()
117    }
118
119    fn content_type_str(&self) -> &str {
120        match self.content_type {
121            TransportContentType::HandoffBundle => "handoff_bundle",
122            TransportContentType::ContextPackage => "context_package",
123            TransportContentType::A2AMessage => "a2a_message",
124            TransportContentType::A2ATask => "a2a_task",
125        }
126    }
127}
128
129pub fn serialize_envelope(envelope: &TransportEnvelopeV1) -> Result<String, String> {
130    let json = serde_json::to_string_pretty(envelope).map_err(|e| e.to_string())?;
131    if json.len() > MAX_ENVELOPE_BYTES {
132        return Err(format!(
133            "envelope too large ({} bytes, max {})",
134            json.len(),
135            MAX_ENVELOPE_BYTES
136        ));
137    }
138    Ok(json)
139}
140
141pub fn parse_envelope(json: &str) -> Result<TransportEnvelopeV1, String> {
142    if json.len() > MAX_ENVELOPE_BYTES {
143        return Err(format!(
144            "envelope too large ({} bytes, max {})",
145            json.len(),
146            MAX_ENVELOPE_BYTES
147        ));
148    }
149    let env: TransportEnvelopeV1 = serde_json::from_str(json).map_err(|e| e.to_string())?;
150    if env.format_version != 1 {
151        return Err(format!(
152            "unsupported format_version {} (expected 1)",
153            env.format_version
154        ));
155    }
156    Ok(env)
157}
158
159fn compute_daemon_fingerprint() -> String {
160    use sha2::{Digest, Sha256};
161    let mut hasher = Sha256::new();
162    hasher.update(env!("CARGO_PKG_VERSION").as_bytes());
163    if let Ok(exe) = std::env::current_exe() {
164        hasher.update(exe.to_string_lossy().as_bytes());
165    }
166    format!("{:x}", hasher.finalize())[..16].to_string()
167}
168
169fn default_capabilities() -> Vec<String> {
170    vec![
171        "context_compression".to_string(),
172        "knowledge_graph".to_string(),
173        "shared_sessions".to_string(),
174        "a2a_messaging".to_string(),
175    ]
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    fn test_sender() -> AgentIdentityV1 {
183        AgentIdentityV1 {
184            agent_id: "test-agent".to_string(),
185            agent_type: "cursor".to_string(),
186            daemon_fingerprint: "abcd1234".to_string(),
187            capabilities: vec!["context_compression".to_string()],
188        }
189    }
190
191    #[test]
192    fn envelope_roundtrip() {
193        let env = TransportEnvelopeV1::new(
194            test_sender(),
195            Some("target-agent"),
196            TransportContentType::A2AMessage,
197            r#"{"hello":"world"}"#.to_string(),
198        );
199        let json = serialize_envelope(&env).unwrap();
200        let parsed = parse_envelope(&json).unwrap();
201        assert_eq!(parsed.format_version, 1);
202        assert_eq!(parsed.sender.agent_id, "test-agent");
203        assert_eq!(parsed.recipient, Some("target-agent".to_string()));
204        assert_eq!(parsed.content_type, TransportContentType::A2AMessage);
205    }
206
207    #[test]
208    fn hmac_sign_verify() {
209        let secret = b"test-secret-key";
210        let mut env = TransportEnvelopeV1::new(
211            test_sender(),
212            None,
213            TransportContentType::HandoffBundle,
214            "payload".to_string(),
215        );
216        assert!(!env.verify_signature(secret));
217
218        env.sign(secret);
219        assert!(env.signature.is_some());
220        assert!(env.verify_signature(secret));
221        assert!(!env.verify_signature(b"wrong-key"));
222    }
223
224    #[test]
225    fn rejects_oversized_envelope() {
226        let big = "x".repeat(MAX_ENVELOPE_BYTES + 1);
227        assert!(parse_envelope(&big).is_err());
228    }
229
230    #[test]
231    fn rejects_wrong_version() {
232        let json = r#"{"format_version":99,"sent_at":"2026-01-01T00:00:00Z","sender":{"agent_id":"a","agent_type":"b","daemon_fingerprint":"c","capabilities":[]},"recipient":null,"content_type":"a2a_message","payload_json":"{}","signature":null,"metadata":{}}"#;
233        assert!(parse_envelope(json).is_err());
234    }
235}