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