1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use super::task::{Task, TaskMessage, TaskPart, TaskState, TaskStore};
5
6#[derive(Debug, Deserialize)]
7pub struct JsonRpcRequest {
8 pub jsonrpc: String,
9 pub id: Value,
10 pub method: String,
11 #[serde(default)]
12 pub params: Value,
13}
14
15#[derive(Debug, Serialize)]
16pub struct JsonRpcResponse {
17 pub jsonrpc: String,
18 pub id: Value,
19 #[serde(skip_serializing_if = "Option::is_none")]
20 pub result: Option<Value>,
21 #[serde(skip_serializing_if = "Option::is_none")]
22 pub error: Option<JsonRpcError>,
23}
24
25#[derive(Debug, Serialize)]
26pub struct JsonRpcError {
27 pub code: i32,
28 pub message: String,
29 #[serde(skip_serializing_if = "Option::is_none")]
30 pub data: Option<Value>,
31}
32
33impl JsonRpcResponse {
34 fn success(id: Value, result: Value) -> Self {
35 Self {
36 jsonrpc: "2.0".to_string(),
37 id,
38 result: Some(result),
39 error: None,
40 }
41 }
42
43 fn error(id: Value, code: i32, message: &str) -> Self {
44 Self {
45 jsonrpc: "2.0".to_string(),
46 id,
47 result: None,
48 error: Some(JsonRpcError {
49 code,
50 message: message.to_string(),
51 data: None,
52 }),
53 }
54 }
55}
56
57pub fn handle_a2a_jsonrpc(req: &JsonRpcRequest) -> JsonRpcResponse {
60 if req.jsonrpc != "2.0" {
61 return JsonRpcResponse::error(req.id.clone(), -32600, "invalid jsonrpc version");
62 }
63
64 match req.method.as_str() {
65 "tasks/send" => handle_send_message(req),
66 "tasks/get" => handle_get_task(req),
67 "tasks/cancel" => handle_cancel_task(req),
68 _ => JsonRpcResponse::error(
69 req.id.clone(),
70 -32601,
71 &format!("method not found: {}", req.method),
72 ),
73 }
74}
75
76fn handle_send_message(req: &JsonRpcRequest) -> JsonRpcResponse {
77 let params = &req.params;
78
79 let from_agent = params
80 .get("message")
81 .and_then(|m| m.get("role"))
82 .and_then(Value::as_str)
83 .unwrap_or("anonymous");
84 let to_agent = params
85 .get("to")
86 .and_then(Value::as_str)
87 .unwrap_or("lean-ctx");
88 let description = params
89 .get("message")
90 .and_then(|m| m.get("parts"))
91 .and_then(|p| p.as_array())
92 .and_then(|parts| parts.first())
93 .and_then(|p| p.get("text"))
94 .and_then(Value::as_str)
95 .unwrap_or("");
96
97 if description.is_empty() {
98 return JsonRpcResponse::error(req.id.clone(), -32602, "message text is required");
99 }
100
101 let mut store = TaskStore::load();
102
103 let task_id = if let Some(id) = params.get("id").and_then(Value::as_str) {
104 if let Some(task) = store.get_task_mut(id) {
105 let parts = extract_message_parts(params);
106 task.add_message(from_agent, parts);
107 if task.state == TaskState::InputRequired {
108 let _ = task.transition(TaskState::Working, Some("input received via A2A"));
109 }
110 id.to_string()
111 } else {
112 return JsonRpcResponse::error(req.id.clone(), -32602, "task not found");
113 }
114 } else {
115 store.create_task(from_agent, to_agent, description)
116 };
117
118 let _ = store.save();
119
120 let task = store.get_task(&task_id);
121 JsonRpcResponse::success(req.id.clone(), task_to_a2a_json(task))
122}
123
124fn handle_get_task(req: &JsonRpcRequest) -> JsonRpcResponse {
125 let Some(task_id) = req.params.get("id").and_then(Value::as_str) else {
126 return JsonRpcResponse::error(req.id.clone(), -32602, "id is required");
127 };
128
129 let store = TaskStore::load();
130 match store.get_task(task_id) {
131 Some(task) => JsonRpcResponse::success(req.id.clone(), task_to_a2a_json(Some(task))),
132 None => JsonRpcResponse::error(req.id.clone(), -32602, "task not found"),
133 }
134}
135
136fn handle_cancel_task(req: &JsonRpcRequest) -> JsonRpcResponse {
137 let Some(task_id) = req.params.get("id").and_then(Value::as_str) else {
138 return JsonRpcResponse::error(req.id.clone(), -32602, "id is required");
139 };
140
141 let mut store = TaskStore::load();
142 let Some(task) = store.get_task_mut(task_id) else {
143 return JsonRpcResponse::error(req.id.clone(), -32602, "task not found");
144 };
145
146 if let Err(e) = task.transition(TaskState::Canceled, Some("canceled via A2A")) {
147 return JsonRpcResponse::error(req.id.clone(), -32603, &e);
148 }
149 let _ = store.save();
150
151 let task = store.get_task(task_id);
152 JsonRpcResponse::success(req.id.clone(), task_to_a2a_json(task))
153}
154
155fn task_to_a2a_json(task: Option<&Task>) -> Value {
156 let Some(task) = task else {
157 return Value::Null;
158 };
159
160 let messages: Vec<Value> = task.messages.iter().map(message_to_a2a_json).collect();
161
162 let artifacts: Vec<Value> = task.artifacts.iter().map(part_to_a2a_json).collect();
163
164 let history: Vec<Value> = task
165 .history
166 .iter()
167 .map(|h| {
168 serde_json::json!({
169 "from": h.from.to_string(),
170 "to": h.to.to_string(),
171 "timestamp": h.timestamp.to_rfc3339(),
172 "reason": h.reason,
173 })
174 })
175 .collect();
176
177 serde_json::json!({
178 "id": task.id,
179 "status": {
180 "state": task.state.to_string(),
181 "timestamp": task.updated_at.to_rfc3339(),
182 },
183 "messages": messages,
184 "artifacts": artifacts,
185 "history": history,
186 "metadata": task.metadata,
187 })
188}
189
190fn message_to_a2a_json(m: &TaskMessage) -> Value {
191 let parts: Vec<Value> = m.parts.iter().map(part_to_a2a_json).collect();
192 serde_json::json!({
193 "role": m.role,
194 "parts": parts,
195 "timestamp": m.timestamp.to_rfc3339(),
196 })
197}
198
199fn part_to_a2a_json(p: &TaskPart) -> Value {
200 match p {
201 TaskPart::Text { text } => serde_json::json!({"type": "text", "text": text}),
202 TaskPart::Data { mime_type, data } => {
203 serde_json::json!({"type": "data", "mimeType": mime_type, "data": data})
204 }
205 TaskPart::File {
206 name,
207 mime_type,
208 data,
209 uri,
210 } => serde_json::json!({
211 "type": "file",
212 "file": {
213 "name": name,
214 "mimeType": mime_type,
215 "bytes": data,
216 "uri": uri,
217 }
218 }),
219 }
220}
221
222fn extract_message_parts(params: &Value) -> Vec<TaskPart> {
223 params
224 .get("message")
225 .and_then(|m| m.get("parts"))
226 .and_then(|p| p.as_array())
227 .map(|parts| {
228 parts
229 .iter()
230 .filter_map(|p| {
231 let ptype = p.get("type")?.as_str()?;
232 match ptype {
233 "text" => Some(TaskPart::Text {
234 text: p.get("text")?.as_str()?.to_string(),
235 }),
236 "data" => Some(TaskPart::Data {
237 mime_type: p
238 .get("mimeType")
239 .and_then(Value::as_str)
240 .unwrap_or("application/octet-stream")
241 .to_string(),
242 data: p.get("data")?.as_str()?.to_string(),
243 }),
244 _ => None,
245 }
246 })
247 .collect()
248 })
249 .unwrap_or_default()
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255
256 fn make_request(method: &str, params: Value) -> JsonRpcRequest {
257 JsonRpcRequest {
258 jsonrpc: "2.0".to_string(),
259 id: Value::Number(1.into()),
260 method: method.to_string(),
261 params,
262 }
263 }
264
265 #[test]
266 fn rejects_unknown_method() {
267 let req = make_request("tasks/unknown", serde_json::json!({}));
268 let resp = handle_a2a_jsonrpc(&req);
269 assert!(resp.error.is_some());
270 assert_eq!(resp.error.unwrap().code, -32601);
271 }
272
273 #[test]
274 fn rejects_missing_message_text() {
275 let req = make_request(
276 "tasks/send",
277 serde_json::json!({
278 "message": { "role": "user", "parts": [] }
279 }),
280 );
281 let resp = handle_a2a_jsonrpc(&req);
282 assert!(resp.error.is_some());
283 }
284
285 #[test]
286 fn send_creates_task() {
287 let req = make_request(
288 "tasks/send",
289 serde_json::json!({
290 "to": "lean-ctx",
291 "message": {
292 "role": "user",
293 "parts": [{"type": "text", "text": "Fix the auth bug"}]
294 }
295 }),
296 );
297 let resp = handle_a2a_jsonrpc(&req);
298 assert!(resp.result.is_some());
299 let result = resp.result.unwrap();
300 assert!(result.get("id").is_some());
301 assert_eq!(
302 result.get("status").unwrap().get("state").unwrap().as_str(),
303 Some("created")
304 );
305 }
306
307 #[test]
308 fn get_nonexistent_task_returns_error() {
309 let req = make_request(
310 "tasks/get",
311 serde_json::json!({"id": "nonexistent-task-id"}),
312 );
313 let resp = handle_a2a_jsonrpc(&req);
314 assert!(resp.error.is_some());
315 }
316}