1use bytes::Bytes;
4use serde_json::json;
5
6use crate::{
7 codec::Codec,
8 protocol::{
9 agent::AgentCard,
10 error::A2AError,
11 operation::A2AOperation,
12 task::{Task, TaskListResponse},
13 },
14 service::response::A2AResponse,
15};
16
17#[derive(Debug, Clone)]
19pub struct JsonCodec;
20
21impl JsonCodec {
22 pub fn new() -> Self {
24 Self
25 }
26}
27
28impl Default for JsonCodec {
29 fn default() -> Self {
30 Self::new()
31 }
32}
33
34impl Codec for JsonCodec {
35 fn encode_request(&self, operation: &A2AOperation) -> Result<Bytes, A2AError> {
36 let json = match operation {
37 A2AOperation::SendMessage {
38 message,
39 stream,
40 context_id,
41 task_id,
42 } => {
43 let mut obj = json!({
44 "message": message,
45 "stream": stream,
46 });
47
48 if let Some(ctx_id) = context_id {
49 obj["contextId"] = json!(ctx_id);
50 }
51 if let Some(t_id) = task_id {
52 obj["taskId"] = json!(t_id);
53 }
54
55 obj
56 }
57 A2AOperation::CancelTask { task_id } => {
58 json!({
59 "taskId": task_id,
60 })
61 }
62 A2AOperation::RegisterWebhook { url, events, auth } => {
63 let mut obj = json!({
64 "url": url,
65 "events": events,
66 });
67
68 if let Some(auth_str) = auth {
69 obj["auth"] = json!(auth_str);
70 }
71
72 obj
73 }
74 _ => json!({}),
76 };
77
78 let bytes = serde_json::to_vec(&json)?;
79 Ok(Bytes::from(bytes))
80 }
81
82 fn decode_response(
83 &self,
84 body: &[u8],
85 operation: &A2AOperation,
86 ) -> Result<A2AResponse, A2AError> {
87 if body.is_empty() {
89 return Ok(A2AResponse::Empty);
90 }
91
92 match operation {
93 A2AOperation::SendMessage { .. } | A2AOperation::GetTask { .. } => {
94 let task: Task = serde_json::from_slice(body)?;
95 Ok(A2AResponse::Task(Box::new(task)))
96 }
97 A2AOperation::ListTasks { .. } => {
98 let list: TaskListResponse = serde_json::from_slice(body)?;
99 Ok(A2AResponse::TaskList {
100 tasks: list.tasks,
101 total: list.total,
102 next_token: list.next_token,
103 })
104 }
105 A2AOperation::DiscoverAgent => {
106 let card: AgentCard = serde_json::from_slice(body)?;
107 Ok(A2AResponse::AgentCard(Box::new(card)))
108 }
109 A2AOperation::CancelTask { .. } => {
110 let task: Task = serde_json::from_slice(body)?;
112 Ok(A2AResponse::Task(Box::new(task)))
113 }
114 A2AOperation::SubscribeTask { .. } => {
115 Ok(A2AResponse::Empty)
117 }
118 A2AOperation::RegisterWebhook { .. } => Ok(A2AResponse::Empty),
119 }
120 }
121
122 fn content_type(&self) -> &str {
123 "application/json"
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use serde_json::Value;
130
131 use super::*;
132 use crate::protocol::message::Message;
133
134 #[test]
135 fn test_encode_send_message() {
136 let codec = JsonCodec::new();
137 let message = Message::user("Hello");
138
139 let operation = A2AOperation::SendMessage {
140 message,
141 stream: false,
142 context_id: None,
143 task_id: None,
144 };
145
146 let bytes = codec.encode_request(&operation).unwrap();
147 assert!(!bytes.is_empty());
148
149 let json: Value = serde_json::from_slice(&bytes).unwrap();
151 assert!(json["message"].is_object());
152 assert_eq!(json["stream"], false);
153 }
154
155 #[test]
156 fn test_decode_task_response() {
157 let codec = JsonCodec::new();
158 let json = r#"{
159 "id": "task-123",
160 "status": "submitted",
161 "input": {
162 "role": "user",
163 "parts": [{"text": "Hello"}]
164 },
165 "createdAt": "2024-01-01T00:00:00Z"
166 }"#;
167
168 let operation = A2AOperation::GetTask {
169 task_id: "task-123".to_string(),
170 };
171
172 let response = codec.decode_response(json.as_bytes(), &operation).unwrap();
173
174 match response {
175 A2AResponse::Task(task) => {
176 assert_eq!(task.id, "task-123");
177 }
178 _ => panic!("Expected Task response"),
179 }
180 }
181
182 #[test]
183 fn test_content_type() {
184 let codec = JsonCodec::new();
185 assert_eq!(codec.content_type(), "application/json");
186 }
187}