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}