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