Skip to main content

lean_ctx/core/a2a/
message.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3
4pub use lean_ctx_ocla::{MessagePriority, PrivacyLevel};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct A2AMessage {
8    pub id: String,
9    pub from_agent: String,
10    pub to_agent: Option<String>,
11    pub task_id: Option<String>,
12    pub category: MessageCategory,
13    pub priority: MessagePriority,
14    pub privacy: PrivacyLevel,
15    pub content: String,
16    pub metadata: std::collections::HashMap<String, String>,
17    #[serde(default)]
18    pub project_root: Option<String>,
19    pub timestamp: DateTime<Utc>,
20    pub read_by: Vec<String>,
21    pub expires_at: Option<DateTime<Utc>>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
25pub enum MessageCategory {
26    TaskDelegation,
27    TaskUpdate,
28    TaskResult,
29    ContextShare,
30    Question,
31    Answer,
32    Notification,
33    Handoff,
34}
35
36impl MessageCategory {
37    pub fn parse_str(s: &str) -> Self {
38        match s.to_lowercase().as_str() {
39            "task_delegation" | "delegation" => Self::TaskDelegation,
40            "task_update" | "update" => Self::TaskUpdate,
41            "task_result" | "result" => Self::TaskResult,
42            "context_share" | "share" => Self::ContextShare,
43            "question" => Self::Question,
44            "answer" => Self::Answer,
45            "handoff" => Self::Handoff,
46            _ => Self::Notification,
47        }
48    }
49}
50
51impl std::fmt::Display for MessageCategory {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        match self {
54            Self::TaskDelegation => write!(f, "task_delegation"),
55            Self::TaskUpdate => write!(f, "task_update"),
56            Self::TaskResult => write!(f, "task_result"),
57            Self::ContextShare => write!(f, "context_share"),
58            Self::Question => write!(f, "question"),
59            Self::Answer => write!(f, "answer"),
60            Self::Notification => write!(f, "notification"),
61            Self::Handoff => write!(f, "handoff"),
62        }
63    }
64}
65
66impl A2AMessage {
67    pub fn new(from: &str, to: Option<&str>, category: MessageCategory, content: &str) -> Self {
68        Self {
69            id: generate_msg_id(),
70            from_agent: from.to_string(),
71            to_agent: to.map(std::string::ToString::to_string),
72            task_id: None,
73            category,
74            priority: MessagePriority::Normal,
75            privacy: PrivacyLevel::Team,
76            content: content.to_string(),
77            metadata: std::collections::HashMap::new(),
78            project_root: None,
79            timestamp: Utc::now(),
80            read_by: vec![from.to_string()],
81            expires_at: None,
82        }
83    }
84
85    pub fn with_task(mut self, task_id: &str) -> Self {
86        self.task_id = Some(task_id.to_string());
87        self
88    }
89
90    pub fn with_priority(mut self, priority: MessagePriority) -> Self {
91        self.priority = priority;
92        self
93    }
94
95    pub fn with_privacy(mut self, privacy: PrivacyLevel) -> Self {
96        self.privacy = privacy;
97        self
98    }
99
100    pub fn with_ttl_hours(mut self, hours: u64) -> Self {
101        self.expires_at = Some(Utc::now() + chrono::Duration::hours(hours as i64));
102        self
103    }
104
105    pub fn is_expired(&self) -> bool {
106        self.expires_at.is_some_and(|exp| Utc::now() > exp)
107    }
108
109    pub fn is_visible_to(&self, agent_id: &str) -> bool {
110        if self.is_expired() {
111            return false;
112        }
113        let is_sender = self.from_agent == agent_id;
114        let is_recipient = self.to_agent.as_deref().is_some_and(|t| t == agent_id);
115        self.privacy.allows_access(is_sender, is_recipient)
116    }
117
118    pub fn mark_read(&mut self, agent_id: &str) {
119        if !self.read_by.contains(&agent_id.to_string()) {
120            self.read_by.push(agent_id.to_string());
121        }
122    }
123}
124
125fn generate_msg_id() -> String {
126    use std::time::{SystemTime, UNIX_EPOCH};
127    let ts = SystemTime::now()
128        .duration_since(UNIX_EPOCH)
129        .unwrap_or_default()
130        .as_nanos();
131    let rand: u64 = (ts as u64).wrapping_mul(6364136223846793005);
132    format!("msg-{:x}-{:08x}", ts % 0xFFFF_FFFF, rand as u32)
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn privacy_access_control() {
141        assert!(PrivacyLevel::Public.allows_access(false, false));
142        assert!(PrivacyLevel::Team.allows_access(false, false));
143        assert!(!PrivacyLevel::Private.allows_access(false, false));
144        assert!(PrivacyLevel::Private.allows_access(true, false));
145        assert!(PrivacyLevel::Private.allows_access(false, true));
146    }
147
148    #[test]
149    fn message_visibility() {
150        let msg = A2AMessage::new(
151            "agent-a",
152            Some("agent-b"),
153            MessageCategory::Notification,
154            "hello",
155        )
156        .with_privacy(PrivacyLevel::Private);
157
158        assert!(msg.is_visible_to("agent-a"));
159        assert!(msg.is_visible_to("agent-b"));
160        assert!(!msg.is_visible_to("agent-c"));
161    }
162
163    #[test]
164    fn broadcast_visibility() {
165        let msg = A2AMessage::new("agent-a", None, MessageCategory::Notification, "hey all");
166        assert!(msg.is_visible_to("agent-a"));
167        assert!(msg.is_visible_to("agent-x"));
168    }
169
170    #[test]
171    fn message_expiry() {
172        let mut msg = A2AMessage::new("a", None, MessageCategory::Notification, "tmp");
173        msg.expires_at = Some(Utc::now() - chrono::Duration::hours(1));
174        assert!(msg.is_expired());
175        assert!(!msg.is_visible_to("a"));
176    }
177}