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 pub timestamp: DateTime<Utc>,
72 pub read_by: Vec<String>,
73 pub expires_at: Option<DateTime<Utc>>,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
77pub enum MessageCategory {
78 TaskDelegation,
79 TaskUpdate,
80 TaskResult,
81 ContextShare,
82 Question,
83 Answer,
84 Notification,
85 Handoff,
86}
87
88impl MessageCategory {
89 pub fn parse_str(s: &str) -> Self {
90 match s.to_lowercase().as_str() {
91 "task_delegation" | "delegation" => Self::TaskDelegation,
92 "task_update" | "update" => Self::TaskUpdate,
93 "task_result" | "result" => Self::TaskResult,
94 "context_share" | "share" => Self::ContextShare,
95 "question" => Self::Question,
96 "answer" => Self::Answer,
97 "handoff" => Self::Handoff,
98 _ => Self::Notification,
99 }
100 }
101}
102
103impl std::fmt::Display for MessageCategory {
104 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105 match self {
106 Self::TaskDelegation => write!(f, "task_delegation"),
107 Self::TaskUpdate => write!(f, "task_update"),
108 Self::TaskResult => write!(f, "task_result"),
109 Self::ContextShare => write!(f, "context_share"),
110 Self::Question => write!(f, "question"),
111 Self::Answer => write!(f, "answer"),
112 Self::Notification => write!(f, "notification"),
113 Self::Handoff => write!(f, "handoff"),
114 }
115 }
116}
117
118impl A2AMessage {
119 pub fn new(from: &str, to: Option<&str>, category: MessageCategory, content: &str) -> Self {
120 Self {
121 id: generate_msg_id(),
122 from_agent: from.to_string(),
123 to_agent: to.map(std::string::ToString::to_string),
124 task_id: None,
125 category,
126 priority: MessagePriority::Normal,
127 privacy: PrivacyLevel::Team,
128 content: content.to_string(),
129 metadata: std::collections::HashMap::new(),
130 timestamp: Utc::now(),
131 read_by: vec![from.to_string()],
132 expires_at: None,
133 }
134 }
135
136 pub fn with_task(mut self, task_id: &str) -> Self {
137 self.task_id = Some(task_id.to_string());
138 self
139 }
140
141 pub fn with_priority(mut self, priority: MessagePriority) -> Self {
142 self.priority = priority;
143 self
144 }
145
146 pub fn with_privacy(mut self, privacy: PrivacyLevel) -> Self {
147 self.privacy = privacy;
148 self
149 }
150
151 pub fn with_ttl_hours(mut self, hours: u64) -> Self {
152 self.expires_at = Some(Utc::now() + chrono::Duration::hours(hours as i64));
153 self
154 }
155
156 pub fn is_expired(&self) -> bool {
157 self.expires_at.is_some_and(|exp| Utc::now() > exp)
158 }
159
160 pub fn is_visible_to(&self, agent_id: &str) -> bool {
161 if self.is_expired() {
162 return false;
163 }
164 let is_sender = self.from_agent == agent_id;
165 let is_recipient = self.to_agent.as_ref().is_none_or(|t| t == agent_id);
166 self.privacy.allows_access(is_sender, is_recipient)
167 }
168
169 pub fn mark_read(&mut self, agent_id: &str) {
170 if !self.read_by.contains(&agent_id.to_string()) {
171 self.read_by.push(agent_id.to_string());
172 }
173 }
174}
175
176fn generate_msg_id() -> String {
177 use std::time::{SystemTime, UNIX_EPOCH};
178 let ts = SystemTime::now()
179 .duration_since(UNIX_EPOCH)
180 .unwrap_or_default()
181 .as_nanos();
182 let rand: u64 = (ts as u64).wrapping_mul(6364136223846793005);
183 format!("msg-{:x}-{:08x}", ts % 0xFFFF_FFFF, rand as u32)
184}
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189
190 #[test]
191 fn privacy_access_control() {
192 assert!(PrivacyLevel::Public.allows_access(false, false));
193 assert!(PrivacyLevel::Team.allows_access(false, false));
194 assert!(!PrivacyLevel::Private.allows_access(false, false));
195 assert!(PrivacyLevel::Private.allows_access(true, false));
196 assert!(PrivacyLevel::Private.allows_access(false, true));
197 }
198
199 #[test]
200 fn message_visibility() {
201 let msg = A2AMessage::new(
202 "agent-a",
203 Some("agent-b"),
204 MessageCategory::Notification,
205 "hello",
206 )
207 .with_privacy(PrivacyLevel::Private);
208
209 assert!(msg.is_visible_to("agent-a"));
210 assert!(msg.is_visible_to("agent-b"));
211 assert!(!msg.is_visible_to("agent-c"));
212 }
213
214 #[test]
215 fn broadcast_visibility() {
216 let msg = A2AMessage::new("agent-a", None, MessageCategory::Notification, "hey all");
217 assert!(msg.is_visible_to("agent-a"));
218 assert!(msg.is_visible_to("agent-x"));
219 }
220
221 #[test]
222 fn message_expiry() {
223 let mut msg = A2AMessage::new("a", None, MessageCategory::Notification, "tmp");
224 msg.expires_at = Some(Utc::now() - chrono::Duration::hours(1));
225 assert!(msg.is_expired());
226 assert!(!msg.is_visible_to("a"));
227 }
228}