1use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashMap;
7
8use crate::types::*;
9
10#[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#[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#[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}