Skip to main content

a3s_code_core/tools/
notification.rs

1//! Task notification format for multi-agent coordination
2//!
3//! Provides serialization of tasks into XML and JSON formats for inter-agent
4//! communication, similar to Claude Code's coordinator mode.
5
6use serde::{Deserialize, Serialize};
7use std::fmt;
8
9/// Task notification envelope
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct TaskNotification {
12    /// Unique task identifier
13    pub id: String,
14    /// Task type (e.g., "explore", "implement", "review")
15    pub task_type: String,
16    /// Task prompt/description
17    pub prompt: String,
18    /// Optional parent task ID for dependency tracking
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub parent_id: Option<String>,
21    /// Priority level (1-5, 5 being highest)
22    #[serde(default)]
23    pub priority: u8,
24    /// Additional metadata
25    #[serde(default)]
26    pub metadata: serde_json::Value,
27}
28
29impl TaskNotification {
30    pub fn new(id: String, task_type: String, prompt: String) -> Self {
31        Self {
32            id,
33            task_type,
34            prompt,
35            parent_id: None,
36            priority: 3,
37            metadata: serde_json::json!({}),
38        }
39    }
40
41    pub fn with_priority(mut self, priority: u8) -> Self {
42        self.priority = priority.min(5);
43        self
44    }
45
46    pub fn with_parent(mut self, parent_id: String) -> Self {
47        self.parent_id = Some(parent_id);
48        self
49    }
50
51    pub fn with_metadata(mut self, key: &str, value: serde_json::Value) -> Self {
52        self.metadata[key] = value;
53        self
54    }
55}
56
57/// Notification format serializer trait
58///
59/// Allows customizing how task notifications are serialized for different
60/// communication protocols (XML for Claude Code compatibility, JSON for REST).
61pub trait NotificationFormat: Send + Sync {
62    /// Serialize a task notification
63    fn serialize(&self, task: &TaskNotification) -> Result<String, NotificationError>;
64
65    /// Deserialize a task notification
66    fn deserialize(&self, data: &str) -> Result<TaskNotification, NotificationError>;
67
68    /// Format name (e.g., "xml", "json")
69    fn name(&self) -> &str;
70}
71
72/// Notification serialization error
73#[derive(Debug)]
74pub struct NotificationError {
75    message: String,
76}
77
78impl fmt::Display for NotificationError {
79    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80        write!(f, "{}", self.message)
81    }
82}
83
84impl std::error::Error for NotificationError {}
85
86impl NotificationError {
87    pub fn new(msg: impl Into<String>) -> Self {
88        Self {
89            message: msg.into(),
90        }
91    }
92}
93
94/// JSON notification format (default)
95pub struct JsonNotificationFormat;
96
97impl JsonNotificationFormat {
98    pub fn new() -> Self {
99        Self
100    }
101}
102
103impl Default for JsonNotificationFormat {
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109impl NotificationFormat for JsonNotificationFormat {
110    fn serialize(&self, task: &TaskNotification) -> Result<String, NotificationError> {
111        serde_json::to_string_pretty(task).map_err(|e| NotificationError::new(e.to_string()))
112    }
113
114    fn deserialize(&self, data: &str) -> Result<TaskNotification, NotificationError> {
115        serde_json::from_str(data).map_err(|e| NotificationError::new(e.to_string()))
116    }
117
118    fn name(&self) -> &str {
119        "json"
120    }
121}
122
123/// XML notification format (Claude Code compatible)
124pub struct XmlNotificationFormat;
125
126impl XmlNotificationFormat {
127    pub fn new() -> Self {
128        Self
129    }
130}
131
132impl Default for XmlNotificationFormat {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138impl NotificationFormat for XmlNotificationFormat {
139    fn serialize(&self, task: &TaskNotification) -> Result<String, NotificationError> {
140        let mut xml = format!(
141            r#"<task>
142  <id>{}</id>
143  <type>{}</type>
144  <prompt><![CDATA[{}]]></prompt>"#,
145            escape_xml(&task.id),
146            escape_xml(&task.task_type),
147            escape_xml(&task.prompt)
148        );
149
150        if let Some(ref parent) = task.parent_id {
151            xml.push_str(&format!(
152                "\n  <parent_id>{}</parent_id>",
153                escape_xml(parent)
154            ));
155        }
156
157        xml.push_str(&format!("\n  <priority>{}</priority>", task.priority));
158
159        // Serialize metadata as child elements
160        if !task.metadata.is_null() {
161            xml.push_str("\n  <metadata>");
162            if let Some(obj) = task.metadata.as_object() {
163                for (k, v) in obj {
164                    xml.push_str(&format!(
165                        "\n    <{}>{}</{}>",
166                        escape_xml(k),
167                        escape_xml(&v.to_string()),
168                        escape_xml(k)
169                    ));
170                }
171            }
172            xml.push_str("\n  </metadata>");
173        }
174
175        xml.push_str("\n</task>");
176        Ok(xml)
177    }
178
179    fn deserialize(&self, data: &str) -> Result<TaskNotification, NotificationError> {
180        // Simple XML parsing for task elements
181        let id =
182            extract_xml_value(data, "id").ok_or_else(|| NotificationError::new("Missing id"))?;
183        let task_type = extract_xml_value(data, "type")
184            .ok_or_else(|| NotificationError::new("Missing type"))?;
185        let prompt = extract_xml_value(data, "prompt")
186            .ok_or_else(|| NotificationError::new("Missing prompt"))?;
187        let parent_id = extract_xml_value(data, "parent_id");
188        let priority = extract_xml_value(data, "priority")
189            .and_then(|s| s.parse().ok())
190            .unwrap_or(3);
191
192        Ok(TaskNotification {
193            id,
194            task_type,
195            prompt,
196            parent_id,
197            priority,
198            metadata: serde_json::json!({}),
199        })
200    }
201
202    fn name(&self) -> &str {
203        "xml"
204    }
205}
206
207/// Escape special XML characters
208fn escape_xml(s: &str) -> String {
209    s.replace('&', "&amp;")
210        .replace('<', "&lt;")
211        .replace('>', "&gt;")
212        .replace('"', "&quot;")
213        .replace('\'', "&apos;")
214}
215
216/// Extract value from XML element (simple regex-based parser)
217fn extract_xml_value(xml: &str, tag: &str) -> Option<String> {
218    let pattern = format!(r#"<{tag}><!\[CDATA\[([\s\S]*?)\]\]></{tag}>"#, tag = tag);
219    if let Ok(re) = regex::Regex::new(&pattern) {
220        if let Some(cap) = re.captures(xml) {
221            return Some(cap.get(1).unwrap().as_str().to_string());
222        }
223    }
224
225    // Try without CDATA
226    let pattern = format!(r"<{tag}>([^<]*)</{tag}>", tag = tag);
227    if let Ok(re) = regex::Regex::new(&pattern) {
228        if let Some(cap) = re.captures(xml) {
229            return Some(cap.get(1).unwrap().as_str().to_string());
230        }
231    }
232
233    None
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    #[test]
241    fn test_json_roundtrip() {
242        let task = TaskNotification::new(
243            "task-1".to_string(),
244            "explore".to_string(),
245            "Find all files".to_string(),
246        )
247        .with_priority(5);
248
249        let format = JsonNotificationFormat::new();
250        let serialized = format.serialize(&task).unwrap();
251        let deserialized = format.deserialize(&serialized).unwrap();
252
253        assert_eq!(deserialized.id, "task-1");
254        assert_eq!(deserialized.task_type, "explore");
255        assert_eq!(deserialized.prompt, "Find all files");
256        assert_eq!(deserialized.priority, 5);
257    }
258
259    #[test]
260    fn test_xml_roundtrip() {
261        let task = TaskNotification::new(
262            "task-2".to_string(),
263            "implement".to_string(),
264            "Add feature".to_string(),
265        )
266        .with_parent("task-1".to_string());
267
268        let format = XmlNotificationFormat::new();
269        let serialized = format.serialize(&task).unwrap();
270        assert!(serialized.contains("<task>"));
271        assert!(serialized.contains("<id>task-2</id>"));
272
273        let deserialized = format.deserialize(&serialized).unwrap();
274        assert_eq!(deserialized.id, "task-2");
275        assert_eq!(deserialized.parent_id, Some("task-1".to_string()));
276    }
277
278    #[test]
279    fn test_xml_escape_cdata() {
280        let task = TaskNotification::new(
281            "task-3".to_string(),
282            "review".to_string(),
283            "Check <input> & \"output\"".to_string(),
284        );
285
286        let format = XmlNotificationFormat::new();
287        let serialized = format.serialize(&task).unwrap();
288        assert!(serialized.contains("&lt;input&gt;"));
289        assert!(serialized.contains("&amp;"));
290    }
291
292    #[test]
293    fn test_task_notification_builder() {
294        let task = TaskNotification::new(
295            "test-id".to_string(),
296            "test-type".to_string(),
297            "test prompt".to_string(),
298        )
299        .with_priority(4)
300        .with_parent("parent-id".to_string())
301        .with_metadata("key", serde_json::json!("value"));
302
303        assert_eq!(task.id, "test-id");
304        assert_eq!(task.task_type, "test-type");
305        assert_eq!(task.prompt, "test prompt");
306        assert_eq!(task.priority, 4);
307        assert_eq!(task.parent_id, Some("parent-id".to_string()));
308        assert_eq!(task.metadata["key"], "value");
309    }
310
311    #[test]
312    fn test_task_notification_priority_clamped() {
313        let task =
314            TaskNotification::new("id".to_string(), "type".to_string(), "prompt".to_string())
315                .with_priority(10); // Should be clamped to 5
316
317        assert_eq!(task.priority, 5);
318    }
319
320    #[test]
321    fn test_json_notification_format_name() {
322        let format = JsonNotificationFormat::new();
323        assert_eq!(format.name(), "json");
324    }
325
326    #[test]
327    fn test_xml_notification_format_name() {
328        let format = XmlNotificationFormat::new();
329        assert_eq!(format.name(), "xml");
330    }
331
332    #[test]
333    fn test_json_serialize_error() {
334        let format = JsonNotificationFormat::new();
335        // Invalid task with circular reference would fail
336        let task =
337            TaskNotification::new("id".to_string(), "type".to_string(), "prompt".to_string());
338        let result = format.serialize(&task);
339        assert!(result.is_ok());
340    }
341
342    #[test]
343    fn test_json_deserialize_invalid() {
344        let format = JsonNotificationFormat::new();
345        let result = format.deserialize("not json");
346        assert!(result.is_err());
347    }
348
349    #[test]
350    fn test_xml_deserialize_invalid() {
351        let format = XmlNotificationFormat::new();
352        let result = format.deserialize("<task><id>only id</task>");
353        assert!(result.is_err());
354    }
355
356    #[test]
357    fn test_xml_deserialize_missing_fields() {
358        let format = XmlNotificationFormat::new();
359        let result = format.deserialize("<task><id>test</id></task>");
360        assert!(result.is_err()); // Missing type and prompt
361    }
362
363    #[test]
364    fn test_xml_serialize_with_metadata() {
365        let task =
366            TaskNotification::new("id".to_string(), "type".to_string(), "prompt".to_string())
367                .with_metadata("extra", serde_json::json!({"nested": true}));
368
369        let format = XmlNotificationFormat::new();
370        let serialized = format.serialize(&task).unwrap();
371        assert!(serialized.contains("<extra>"));
372    }
373
374    #[test]
375    fn test_notification_error_display() {
376        let err = NotificationError::new("test error");
377        assert_eq!(format!("{}", err), "test error");
378    }
379
380    #[test]
381    fn test_notification_error_fromserde() {
382        let err = NotificationError::new("json error");
383        let result: Result<String, _> = Err(err);
384        assert!(result.is_err());
385    }
386
387    #[test]
388    fn test_escape_xml() {
389        assert_eq!(escape_xml("<>&\"'"), "&lt;&gt;&amp;&quot;&apos;");
390        assert_eq!(escape_xml("normal"), "normal");
391        assert_eq!(escape_xml(""), "");
392    }
393
394    #[test]
395    fn test_extract_xml_value_cdata() {
396        let xml = r#"<test><![CDATA[content here]]></test>"#;
397        let value = extract_xml_value(xml, "test");
398        assert_eq!(value, Some("content here".to_string()));
399    }
400
401    #[test]
402    fn test_extract_xml_value_simple() {
403        let xml = "<test>simple content</test>";
404        let value = extract_xml_value(xml, "test");
405        assert_eq!(value, Some("simple content".to_string()));
406    }
407
408    #[test]
409    fn test_extract_xml_value_not_found() {
410        let xml = "<other>content</other>";
411        let value = extract_xml_value(xml, "test");
412        assert!(value.is_none());
413    }
414
415    #[test]
416    fn test_task_notification_default_priority() {
417        let task =
418            TaskNotification::new("id".to_string(), "type".to_string(), "prompt".to_string());
419        assert_eq!(task.priority, 3); // Default priority
420    }
421
422    #[test]
423    fn test_task_notification_with_metadata_multiple() {
424        let task =
425            TaskNotification::new("id".to_string(), "type".to_string(), "prompt".to_string())
426                .with_metadata("key1", serde_json::json!("value1"))
427                .with_metadata("key2", serde_json::json!(42));
428
429        assert_eq!(task.metadata["key1"], "value1");
430        assert_eq!(task.metadata["key2"], 42);
431    }
432}