Skip to main content

netget/server/a2a/
mod.rs

1//! A2A 1.0 agent over HTTP/1.1: agent card, JSON-RPC dispatch, SSE streaming. Rust owns the
2//! envelope, ids and version rules; the handler owns every reply.
3pub mod actions;
4pub mod model;
5
6use crate::llm::action_helper::call_llm;
7use crate::llm::actions::protocol_trait::ActionResult;
8use crate::logging::emit::Log;
9use crate::protocol::{Event, SpawnContext};
10use crate::server::accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS};
11use crate::server::connection::ConnectionId;
12use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
13use anyhow::{bail, ensure, Context, Result};
14use bytes::Bytes;
15use http_body_util::{BodyExt, Full, Limited};
16use hyper::{
17    body::Incoming, header, server::conn::http1, service::service_fn, Method, Request, Response,
18    StatusCode,
19};
20use hyper_util::rt::{TokioIo, TokioTimer};
21use serde_json::{json, Value};
22use std::{convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
23
24pub const DEFAULT_NAME: &str = "NetGet Agent";
25pub const DEFAULT_DESCRIPTION: &str = "An A2A agent served by NetGet";
26pub const DEFAULT_STREAMING: bool = true;
27const HEADER_TIMEOUT: Duration = Duration::from_secs(30);
28const BODY_TIMEOUT: Duration = Duration::from_secs(30);
29
30struct Shared {
31    ctx: SpawnContext,
32    name: String,
33    description: String,
34    skills: Vec<Value>,
35    streaming: bool,
36    local: SocketAddr,
37}
38
39pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
40    let p = ctx.startup_params.as_ref();
41    let s = |k: &str, d: &str| -> Result<String> {
42        Ok(p.map(|p| p.get_optional_string(k))
43            .transpose()?
44            .flatten()
45            .unwrap_or_else(|| d.to_owned()))
46    };
47    let name = s("agent_name", DEFAULT_NAME)?;
48    let description = s("agent_description", DEFAULT_DESCRIPTION)?;
49    ensure!(
50        !name.is_empty() && name.len() <= 256 && description.len() <= 4096,
51        "agent_name/agent_description too long"
52    );
53    let skills = match p
54        .map(|p| p.get_optional_array("skills"))
55        .transpose()?
56        .flatten()
57    {
58        None => vec![
59            json!({"id": "chat", "name": "Chat", "description": "Converse with the agent", "tags": ["chat"]}),
60        ],
61        Some(list) => {
62            ensure!(list.len() <= 64, "at most 64 skills");
63            list.iter()
64                .map(|sk| {
65                    let id = sk["id"].as_str().context("each skill needs an id")?;
66                    let name = sk["name"].as_str().context("each skill needs a name")?;
67                    Ok(json!({"id": id, "name": name, "description": sk["description"].as_str().unwrap_or(""), "tags": sk.get("tags").cloned().unwrap_or(json!([])), "examples": sk.get("examples").cloned().unwrap_or(json!([]))}))
68                })
69                .collect::<Result<_>>()?
70        }
71    };
72    let streaming = p
73        .map(|p| p.get_optional_bool("streaming"))
74        .transpose()?
75        .flatten()
76        .unwrap_or(DEFAULT_STREAMING);
77    let listener = tokio::net::TcpListener::bind(ctx.legacy_listen_addr()).await?;
78    let local = listener.local_addr()?;
79    Log::new(Some(&ctx.status_tx)).info(format!(
80        "A2A agent '{name}' at http://{local}/ (card {})",
81        model::CARD_PATH
82    ));
83    let shared = Arc::new(Shared {
84        ctx: ctx.clone(),
85        name,
86        description,
87        skills,
88        streaming,
89        local,
90    });
91    let server_id = ctx.server_id;
92    let accept = tokio::spawn(async move {
93        let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
94        loop {
95            let (stream, peer, permit) = match accept_bounded(&listener, &limiter, b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", "A2A", Some(&shared.ctx.status_tx)).await {
96                Ok(v) => v,
97                Err(_) => break,
98            };
99            let id = ConnectionId::new(shared.ctx.state.get_next_unified_id().await);
100            let now = crate::utils::clock::Instant::now();
101            shared
102                .ctx
103                .state
104                .add_connection_to_server(
105                    server_id,
106                    ConnectionState {
107                        id,
108                        remote_addr: peer,
109                        local_addr: local,
110                        bytes_sent: 0,
111                        bytes_received: 0,
112                        packets_sent: 0,
113                        packets_received: 0,
114                        last_activity: now,
115                        status: ConnectionStatus::Active,
116                        status_changed_at: now,
117                        protocol_info: ProtocolConnectionInfo::empty(),
118                    },
119                )
120                .await;
121            let child = shared.clone();
122            shared
123                .ctx
124                .state
125                .spawn_server_task(server_id, async move {
126                    let _permit = permit;
127                    let svc_shared = child.clone();
128                    let service = service_fn(move |req| {
129                        let shared = svc_shared.clone();
130                        async move { Ok::<_, Infallible>(handle(&shared, id, req).await) }
131                    });
132                    let mut builder = http1::Builder::new();
133                    builder
134                        .timer(TokioTimer::new())
135                        .header_read_timeout(HEADER_TIMEOUT)
136                        .max_headers(64)
137                        .max_buf_size(32 * 1024);
138                    if let Err(e) = builder
139                        .serve_connection(TokioIo::new(stream), service)
140                        .await
141                    {
142                        Log::new(Some(&child.ctx.status_tx))
143                            .debug(format!("A2A connection {id}: {e}"));
144                    }
145                    child
146                        .ctx
147                        .state
148                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
149                        .await;
150                    let _ = child.ctx.status_tx.send("__UPDATE_UI__".into());
151                })
152                .await;
153        }
154    });
155    ctx.state.register_server_task(server_id, accept).await;
156    Ok(local)
157}
158
159fn respond(status: u16, content_type: &'static str, body: Vec<u8>) -> Response<Full<Bytes>> {
160    let mut r = Response::new(Full::new(Bytes::from(body)));
161    *r.status_mut() = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
162    r.headers_mut().insert(
163        header::CONTENT_TYPE,
164        header::HeaderValue::from_static(content_type),
165    );
166    r
167}
168fn json_response(value: &Value) -> Response<Full<Bytes>> {
169    respond(
170        200,
171        "application/json",
172        serde_json::to_vec(value).unwrap_or_default(),
173    )
174}
175
176fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
177    let summary = format!("A2A connection {id} method={operation} decision={decision}");
178    let log = Log::new(Some(&ctx.status_tx));
179    if decision.starts_with("fail_closed_") || decision == "model_silent" {
180        log.error(summary);
181    } else {
182        log.info(summary);
183    }
184}
185
186async fn decide(shared: &Shared, id: ConnectionId, event: Event) -> Result<Value> {
187    let ctx = &shared.ctx;
188    let result = match call_llm(
189        &ctx.llm_client,
190        &ctx.state,
191        ctx.server_id,
192        Some(id),
193        &event,
194        &actions::A2aProtocol,
195    )
196    .await
197    {
198        Ok(r) => r,
199        Err(e) => {
200            outcome(ctx, id, event.id(), "fail_closed_llm_error");
201            return Err(e);
202        }
203    };
204    if !result.failures.is_empty() {
205        outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
206        bail!("A2A handler supplied an invalid action");
207    }
208    let mut found = None;
209    let mut pending = result.protocol_results;
210    while let Some(r) = pending.pop() {
211        match r {
212            ActionResult::Custom { name, data } if name == "a2a_reply" => {
213                if found.replace(data).is_some() {
214                    outcome(ctx, id, event.id(), "fail_closed_invalid_reply");
215                    bail!("A2A handler supplied more than one reply");
216                }
217            }
218            ActionResult::Multiple(items) => pending.extend(items),
219            _ => {}
220        }
221    }
222    found
223        .context("A2A handler did not answer")
224        .inspect_err(|_| outcome(ctx, id, event.id(), "model_silent"))
225}
226
227async fn handle(
228    shared: &Shared,
229    id: ConnectionId,
230    req: Request<Incoming>,
231) -> Response<Full<Bytes>> {
232    let ctx = &shared.ctx;
233    ctx.state
234        .update_connection_stats(
235            ctx.server_id,
236            id,
237            Some(req.uri().to_string().len() as u64),
238            None,
239            Some(1),
240            None,
241        )
242        .await;
243    let response = route(shared, id, req).await;
244    let sent = hyper::body::Body::size_hint(response.body())
245        .exact()
246        .unwrap_or(0);
247    ctx.state
248        .update_connection_stats(ctx.server_id, id, None, Some(sent), None, Some(1))
249        .await;
250    response
251}
252
253async fn route(shared: &Shared, id: ConnectionId, req: Request<Incoming>) -> Response<Full<Bytes>> {
254    if req.method() == Method::GET && req.uri().path() == model::CARD_PATH {
255        let host = req
256            .headers()
257            .get(header::HOST)
258            .and_then(|h| h.to_str().ok())
259            .filter(|h| !h.is_empty() && h.len() <= 255 && !h.contains(['/', ' ']))
260            .map(str::to_owned)
261            .unwrap_or_else(|| shared.local.to_string());
262        let card = model::card(
263            &shared.name,
264            &shared.description,
265            "1.0.0",
266            &format!("http://{host}/"),
267            shared.streaming,
268            &shared.skills,
269        );
270        return json_response(&card);
271    }
272    if req.method() != Method::POST || req.uri().path() != "/" {
273        return respond(404, "text/plain", b"not found".to_vec());
274    }
275    let version = req
276        .headers()
277        .get(model::VERSION_HEADER)
278        .and_then(|v| v.to_str().ok())
279        .map(str::to_owned);
280    let body = match tokio::time::timeout(
281        BODY_TIMEOUT,
282        Limited::new(req.into_body(), model::MAX_BODY_BYTES).collect(),
283    )
284    .await
285    {
286        Ok(Ok(b)) => b.to_bytes(),
287        _ => {
288            return json_response(&model::rpc_error(
289                &Value::Null,
290                -32600,
291                "request body unreadable or over 1 MiB",
292            ))
293        }
294    };
295    let request: Value = match serde_json::from_slice(&body) {
296        Ok(v) if model::budget_ok(&v) => v,
297        _ => return json_response(&model::rpc_error(&Value::Null, -32700, "parse error")),
298    };
299    let rpc_id = request.get("id").cloned().unwrap_or(Value::Null);
300    if request["jsonrpc"] != "2.0"
301        || !matches!(rpc_id, Value::String(_) | Value::Number(_) | Value::Null)
302    {
303        return json_response(&model::rpc_error(
304            &rpc_id,
305            -32600,
306            "invalid JSON-RPC 2.0 request",
307        ));
308    }
309    let Some(method) = request["method"].as_str() else {
310        return json_response(&model::rpc_error(&rpc_id, -32600, "method is required"));
311    };
312    if !model::version_ok(version.as_deref()) {
313        return json_response(&model::rpc_error(
314            &rpc_id,
315            model::VERSION_NOT_SUPPORTED,
316            "A2A-Version 1.0 is required",
317        ));
318    }
319    let params = request.get("params").cloned().unwrap_or(json!({}));
320    if !params.is_object() {
321        return json_response(&model::rpc_error(
322            &rpc_id,
323            -32602,
324            "params must be an object",
325        ));
326    }
327    match method {
328        "SendMessage" | "SendStreamingMessage" => send(shared, id, method, &rpc_id, &params).await,
329        "GetTask" | "CancelTask" | "ListTasks" => {
330            task_request(shared, id, method, &rpc_id, &params).await
331        }
332        "SubscribeToTask"
333        | "CreateTaskPushNotificationConfig"
334        | "GetTaskPushNotificationConfig"
335        | "ListTaskPushNotificationConfigs"
336        | "DeleteTaskPushNotificationConfig" => json_response(&model::rpc_error(
337            &rpc_id,
338            if method == "SubscribeToTask" {
339                model::UNSUPPORTED_OPERATION
340            } else {
341                model::PUSH_NOT_SUPPORTED
342            },
343            "not supported by this agent",
344        )),
345        "GetExtendedAgentCard" => json_response(&model::rpc_error(
346            &rpc_id,
347            model::EXTENDED_CARD_NOT_CONFIGURED,
348            "no extended agent card",
349        )),
350        _ => json_response(&model::rpc_error(&rpc_id, -32601, "method not found")),
351    }
352}
353
354fn internal(rpc_id: &Value, streaming: bool) -> Response<Full<Bytes>> {
355    let err = model::rpc_error(rpc_id, -32603, "the agent cannot answer right now");
356    if streaming {
357        respond(200, "text/event-stream", model::sse(&err).into_bytes())
358    } else {
359        json_response(&err)
360    }
361}
362
363fn mapped_error(rpc_id: &Value, v: &Value) -> Value {
364    let code = v["error"]["code"].as_str().unwrap_or("");
365    let num = actions::ERROR_NAMES
366        .iter()
367        .find(|(n, _)| *n == code)
368        .map(|(_, c)| *c)
369        .unwrap_or(-32603);
370    model::rpc_error(rpc_id, num, v["error"]["message"].as_str().unwrap_or(code))
371}
372
373async fn send(
374    shared: &Shared,
375    id: ConnectionId,
376    method: &str,
377    rpc_id: &Value,
378    params: &Value,
379) -> Response<Full<Bytes>> {
380    let ctx = &shared.ctx;
381    let streaming = method == "SendStreamingMessage";
382    if streaming && !shared.streaming {
383        return json_response(&model::rpc_error(
384            rpc_id,
385            model::UNSUPPORTED_OPERATION,
386            "streaming is not offered",
387        ));
388    }
389    let incoming = match model::incoming_message(params) {
390        Ok(m) => m,
391        Err(e) => {
392            outcome(ctx, id, method, "protocol_refusal");
393            return json_response(&model::rpc_error(rpc_id, -32602, &e.to_string()));
394        }
395    };
396    let context_id = incoming.context_id.clone().unwrap_or_else(model::new_id);
397    let event = Event::new(
398        &actions::MESSAGE_EVENT,
399        json!({"method": method, "text": incoming.text, "parts": incoming.parts, "message_id": incoming.message_id, "context_id": context_id, "task_id": incoming.task_id}),
400    );
401    let reply = match decide(shared, id, event).await {
402        Ok(v) => v,
403        Err(_) => return internal(rpc_id, streaming),
404    };
405    let built = (|| -> Result<Value> {
406        actions::validate_reply(&reply)?;
407        if reply.get("error").is_some_and(|e| !e.is_null()) {
408            return Ok(json!({"error": mapped_error(rpc_id, &reply)}));
409        }
410        if reply.get("message").is_some_and(|m| !m.is_null()) {
411            return Ok(
412                json!({"message": model::agent_message(reply["message"]["text"].as_str(), reply["message"].get("parts"), &context_id, incoming.task_id.as_deref())?}),
413            );
414        }
415        let task = reply
416            .get("task")
417            .filter(|t| !t.is_null())
418            .context("SendMessage is answered with message, task or error")?;
419        let default_id = incoming.task_id.clone().unwrap_or_else(model::new_id);
420        Ok(json!({"task": model::task(task, &default_id, &context_id)?}))
421    })();
422    let built = match built {
423        Ok(b) => b,
424        Err(_) => {
425            outcome(ctx, id, method, "fail_closed_invalid_reply");
426            return internal(rpc_id, streaming);
427        }
428    };
429    if let Some(err) = built.get("error") {
430        outcome(ctx, id, method, "model_reject");
431        return if streaming {
432            respond(200, "text/event-stream", model::sse(err).into_bytes())
433        } else {
434            json_response(err)
435        };
436    }
437    outcome(ctx, id, method, "model_answer");
438    if !streaming {
439        return json_response(&model::rpc_result(rpc_id, built));
440    }
441    let events = match built.get("task") {
442        Some(task) => model::stream_events(rpc_id, task),
443        None => vec![model::rpc_result(rpc_id, built)],
444    };
445    respond(
446        200,
447        "text/event-stream",
448        events
449            .iter()
450            .map(model::sse)
451            .collect::<String>()
452            .into_bytes(),
453    )
454}
455
456async fn task_request(
457    shared: &Shared,
458    id: ConnectionId,
459    method: &str,
460    rpc_id: &Value,
461    params: &Value,
462) -> Response<Full<Bytes>> {
463    let ctx = &shared.ctx;
464    let task_id = params.get("id").and_then(Value::as_str).map(str::to_owned);
465    if method != "ListTasks"
466        && task_id
467            .as_deref()
468            .is_none_or(|t| t.is_empty() || t.len() > 256)
469    {
470        return json_response(&model::rpc_error(rpc_id, -32602, "params.id is required"));
471    }
472    let event = Event::new(
473        &actions::TASK_EVENT,
474        json!({"method": method, "task_id": task_id, "params": params}),
475    );
476    let reply = match decide(shared, id, event).await {
477        Ok(v) => v,
478        Err(_) => return internal(rpc_id, false),
479    };
480    let built = (|| -> Result<Value> {
481        actions::validate_reply(&reply)?;
482        if reply.get("error").is_some_and(|e| !e.is_null()) {
483            return Ok(mapped_error(rpc_id, &reply));
484        }
485        match method {
486            "ListTasks" => {
487                let list = reply
488                    .get("tasks")
489                    .and_then(Value::as_array)
490                    .context("ListTasks is answered with tasks or error")?;
491                let tasks: Vec<Value> = list
492                    .iter()
493                    .map(|t| {
494                        model::task(t, "unused", t["context_id"].as_str().unwrap_or("default"))
495                    })
496                    .collect::<Result<_>>()?;
497                Ok(model::rpc_result(
498                    rpc_id,
499                    json!({"tasks": tasks, "pageSize": tasks.len(), "totalSize": tasks.len(), "nextPageToken": ""}),
500                ))
501            }
502            _ => {
503                let spec = reply
504                    .get("task")
505                    .filter(|t| !t.is_null())
506                    .context("GetTask/CancelTask are answered with task or error")?;
507                let tid = task_id.clone().unwrap_or_default();
508                let task =
509                    model::task(spec, &tid, spec["context_id"].as_str().unwrap_or("default"))?;
510                ensure!(task["id"] == tid, "the answer describes a different task");
511                if method == "CancelTask" {
512                    ensure!(task["status"]["state"] == "TASK_STATE_CANCELED", "a successful CancelTask returns the task canceled; use task_not_cancelable otherwise");
513                }
514                Ok(model::rpc_result(rpc_id, task))
515            }
516        }
517    })();
518    match built {
519        Ok(v) => {
520            outcome(
521                ctx,
522                id,
523                method,
524                if v.get("error").is_some() {
525                    "model_reject"
526                } else {
527                    "model_answer"
528                },
529            );
530            json_response(&v)
531        }
532        Err(_) => {
533            outcome(ctx, id, method, "fail_closed_invalid_reply");
534            internal(rpc_id, false)
535        }
536    }
537}