Skip to main content

a2a/
event.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// Copyright A2A Contributors (https://github.com/a2aproject)
3// SPDX-License-Identifier: Apache-2.0
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashMap;
7
8use crate::types::*;
9
10// ---------------------------------------------------------------------------
11// StreamResponse (field-presence union — 4 variants)
12// ---------------------------------------------------------------------------
13
14/// A streaming event. Uses field-presence serialization for wire compatibility.
15#[derive(Debug, Clone, PartialEq)]
16#[cfg_attr(feature = "fuzzing", derive(arbitrary::Arbitrary))]
17pub enum StreamResponse {
18    Task(Task),
19    Message(Message),
20    StatusUpdate(TaskStatusUpdateEvent),
21    ArtifactUpdate(TaskArtifactUpdateEvent),
22}
23
24impl Serialize for StreamResponse {
25    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
26        use serde::ser::SerializeMap;
27        let mut map = serializer.serialize_map(Some(1))?;
28        match self {
29            StreamResponse::Task(t) => map.serialize_entry("task", t)?,
30            StreamResponse::Message(m) => map.serialize_entry("message", m)?,
31            StreamResponse::StatusUpdate(s) => map.serialize_entry("statusUpdate", s)?,
32            StreamResponse::ArtifactUpdate(a) => map.serialize_entry("artifactUpdate", a)?,
33        }
34        map.end()
35    }
36}
37
38impl<'de> Deserialize<'de> for StreamResponse {
39    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
40        let raw: HashMap<String, Value> = HashMap::deserialize(deserializer)?;
41        if let Some(v) = raw.get("message") {
42            Ok(StreamResponse::Message(
43                serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?,
44            ))
45        } else if let Some(v) = raw.get("task") {
46            Ok(StreamResponse::Task(
47                serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?,
48            ))
49        } else if let Some(v) = raw.get("statusUpdate") {
50            Ok(StreamResponse::StatusUpdate(
51                serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?,
52            ))
53        } else if let Some(v) = raw.get("artifactUpdate") {
54            Ok(StreamResponse::ArtifactUpdate(
55                serde_json::from_value(v.clone()).map_err(serde::de::Error::custom)?,
56            ))
57        } else {
58            Err(serde::de::Error::custom("unknown StreamResponse variant"))
59        }
60    }
61}
62
63// ---------------------------------------------------------------------------
64// TaskStatusUpdateEvent
65// ---------------------------------------------------------------------------
66
67/// Event: a task's status has changed.
68#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
69#[cfg_attr(feature = "fuzzing", derive(arbitrary::Arbitrary))]
70#[serde(rename_all = "camelCase")]
71pub struct TaskStatusUpdateEvent {
72    pub task_id: TaskId,
73    pub context_id: String,
74    pub status: TaskStatus,
75
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    #[cfg_attr(feature = "fuzzing", arbitrary(default))]
78    pub metadata: Option<HashMap<String, Value>>,
79}
80
81// ---------------------------------------------------------------------------
82// TaskArtifactUpdateEvent
83// ---------------------------------------------------------------------------
84
85/// Event: an artifact has been generated or updated.
86#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
87#[cfg_attr(feature = "fuzzing", derive(arbitrary::Arbitrary))]
88#[serde(rename_all = "camelCase")]
89pub struct TaskArtifactUpdateEvent {
90    pub task_id: TaskId,
91    pub context_id: String,
92    pub artifact: Artifact,
93
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub append: Option<bool>,
96
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub last_chunk: Option<bool>,
99
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    #[cfg_attr(feature = "fuzzing", arbitrary(default))]
102    pub metadata: Option<HashMap<String, Value>>,
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn test_stream_response_status_update_serde() {
111        let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
112            task_id: "t1".into(),
113            context_id: "c1".into(),
114            status: TaskStatus {
115                state: TaskState::Working,
116                message: None,
117                timestamp: None,
118            },
119            metadata: None,
120        });
121        let json = serde_json::to_string(&event).unwrap();
122        let v: Value = serde_json::from_str(&json).unwrap();
123        assert!(v.get("statusUpdate").is_some());
124
125        let back: StreamResponse = serde_json::from_str(&json).unwrap();
126        assert!(matches!(back, StreamResponse::StatusUpdate(_)));
127    }
128
129    #[test]
130    fn test_stream_response_task_serde() {
131        let task = Task {
132            id: "t1".into(),
133            context_id: "c1".into(),
134            status: TaskStatus {
135                state: TaskState::Completed,
136                message: None,
137                timestamp: None,
138            },
139            artifacts: None,
140            history: None,
141            metadata: None,
142        };
143        let event = StreamResponse::Task(task.clone());
144        let json = serde_json::to_string(&event).unwrap();
145        let v: Value = serde_json::from_str(&json).unwrap();
146        assert!(v.get("task").is_some());
147        let back: StreamResponse = serde_json::from_str(&json).unwrap();
148        assert!(matches!(back, StreamResponse::Task(_)));
149    }
150
151    #[test]
152    fn test_stream_response_message_serde() {
153        let msg = Message::new(Role::Agent, vec![Part::text("hello")]);
154        let event = StreamResponse::Message(msg);
155        let json = serde_json::to_string(&event).unwrap();
156        let v: Value = serde_json::from_str(&json).unwrap();
157        assert!(v.get("message").is_some());
158        let back: StreamResponse = serde_json::from_str(&json).unwrap();
159        assert!(matches!(back, StreamResponse::Message(_)));
160    }
161
162    #[test]
163    fn test_stream_response_artifact_update_serde() {
164        let event = StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
165            task_id: "t1".into(),
166            context_id: "c1".into(),
167            artifact: Artifact {
168                artifact_id: "a1".into(),
169                name: None,
170                description: None,
171                parts: vec![],
172                metadata: None,
173                extensions: None,
174            },
175            append: Some(true),
176            last_chunk: Some(false),
177            metadata: None,
178        });
179        let json = serde_json::to_string(&event).unwrap();
180        let v: Value = serde_json::from_str(&json).unwrap();
181        assert!(v.get("artifactUpdate").is_some());
182        let back: StreamResponse = serde_json::from_str(&json).unwrap();
183        assert!(matches!(back, StreamResponse::ArtifactUpdate(_)));
184    }
185
186    #[test]
187    fn test_stream_response_unknown_variant() {
188        let json = r#"{"unknown": {}}"#;
189        let result = serde_json::from_str::<StreamResponse>(json);
190        assert!(result.is_err());
191    }
192
193    #[test]
194    fn test_task_status_update_event_with_metadata() {
195        let mut meta = HashMap::new();
196        meta.insert("key".to_string(), Value::String("val".to_string()));
197        let event = TaskStatusUpdateEvent {
198            task_id: "t1".into(),
199            context_id: "c1".into(),
200            status: TaskStatus {
201                state: TaskState::Working,
202                message: Some(Message::new(Role::Agent, vec![])),
203                timestamp: None,
204            },
205            metadata: Some(meta),
206        };
207        let json = serde_json::to_string(&event).unwrap();
208        let back: TaskStatusUpdateEvent = serde_json::from_str(&json).unwrap();
209        assert!(back.metadata.is_some());
210        assert_eq!(back.status.state, TaskState::Working);
211    }
212
213    #[test]
214    fn test_task_artifact_update_event_full() {
215        let event = TaskArtifactUpdateEvent {
216            task_id: "t1".into(),
217            context_id: "c1".into(),
218            artifact: Artifact {
219                artifact_id: "a1".into(),
220                name: Some("file.txt".into()),
221                description: Some("A file".into()),
222                parts: vec![Part::text("content")],
223                metadata: None,
224                extensions: None,
225            },
226            append: None,
227            last_chunk: Some(true),
228            metadata: None,
229        };
230        let json = serde_json::to_string(&event).unwrap();
231        let back: TaskArtifactUpdateEvent = serde_json::from_str(&json).unwrap();
232        assert_eq!(back.last_chunk, Some(true));
233        assert!(back.append.is_none());
234    }
235}