klieo-a2a 0.4.0

//! `A2aServer` — listens on a NATS-style request/reply subject and
//! dispatches incoming JSON-RPC payloads to an [`A2aHandler`].
//!
//! # Transport contract
//!
//! - **Subject:** `<app_prefix>.a2a.<agent_id>.rpc`. Subscribe via
//!   [`klieo_core::Pubsub::subscribe`] with a fixed durable name.
//! - **Reply-to:** the inbound message MUST carry a `"Reply-To"` header
//!   naming the subject the client is listening on. The server publishes
//!   the [`JsonRpcResponse`] there.
//! - **Headers:** `Authorization`, `Content-Type`, `A2A-Version`,
//!   `A2A-Extensions` decoded via [`crate::envelope::A2aHeaders`].
//!
//! See module docs in `lib.rs` for why the server takes
//! `Arc<dyn Pubsub>` instead of `Arc<dyn RequestReply>`.

use crate::auth::{Authenticator, RequestContext};
use crate::envelope::{A2aHeaders, A2aMethod, JsonRpcError, JsonRpcRequest, JsonRpcResponse};
use crate::error::A2aError;
use crate::handler::A2aHandler;
use crate::types::{
    CancelTaskParams, DeletePushNotificationConfigParams, GetExtendedAgentCardParams,
    GetPushNotificationConfigParams, GetTaskParams, ListPushNotificationConfigsParams,
    ListTasksParams, PushNotificationConfigParams, SendMessageParams, SubscribeToTaskParams,
};
use bytes::Bytes;
use klieo_core::{DurableName, Headers, KvStore, Msg, Pubsub};
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::Value;
use std::sync::Arc;
use tokio_stream::StreamExt;
use tracing::{error, warn};

/// JSON-RPC error code for unauthenticated callers. Defined as
/// `-32001` ("Authentication required") in our local error space; the
/// JSON-RPC 2.0 spec reserves -32000..-32099 for server-defined errors.
const JSON_RPC_UNAUTHENTICATED: i32 = -32001;

/// JSON-RPC dispatcher that listens on a NATS subject.
pub struct A2aServer {
    app_prefix: String,
    agent_id: String,
    handler: Arc<dyn A2aHandler>,
    authenticator: Arc<dyn Authenticator>,
    pubsub: Arc<dyn Pubsub>,
    #[allow(dead_code)]
    kv: Arc<dyn KvStore>,
}

impl A2aServer {
    /// Build a new server.
    ///
    /// **Authentication is mandatory** since 0.2.0. Wire a real
    /// [`Authenticator`] (e.g.
    /// [`BearerTokenAuthenticator`](crate::auth::BearerTokenAuthenticator))
    /// for production deployments. For tests and demos that intentionally
    /// run unauthenticated, pass
    /// [`AllowAnonymous`](crate::auth::AllowAnonymous) explicitly — that
    /// makes the choice visible at the wiring site.
    pub fn new(
        app_prefix: String,
        agent_id: String,
        handler: Arc<dyn A2aHandler>,
        authenticator: Arc<dyn Authenticator>,
        pubsub: Arc<dyn Pubsub>,
        kv: Arc<dyn KvStore>,
    ) -> Self {
        Self {
            app_prefix,
            agent_id,
            handler,
            authenticator,
            pubsub,
            kv,
        }
    }

    /// Subject the server listens on.
    pub fn subject(&self) -> String {
        format!("{}.a2a.{}.rpc", self.app_prefix, self.agent_id)
    }

    /// Run the dispatch loop until the underlying subscription terminates.
    pub async fn run(self) -> Result<(), A2aError> {
        let subject = self.subject();
        let durable = DurableName::new(format!("a2a-server-{}", self.agent_id));
        let mut stream = self.pubsub.subscribe(&subject, durable).await?;
        while let Some(item) = stream.next().await {
            match item {
                Ok(msg) => {
                    self.handle_one(msg).await;
                }
                Err(e) => {
                    error!(error = %e, "a2a subscribe stream error; exiting loop");
                    return Err(A2aError::Bus(e));
                }
            }
        }
        Ok(())
    }

    async fn handle_one(&self, msg: Msg) {
        let reply_to = match msg.headers.get("Reply-To") {
            Some(s) => s.clone(),
            None => {
                warn!("a2a request missing Reply-To header; dropping");
                let _ = msg.ack.ack().await;
                return;
            }
        };
        // CWE-306: authenticate the inbound request BEFORE deserializing
        // or dispatching. Failures surface as JSON-RPC -32001 to the
        // caller; the request body never reaches the handler.
        let headers = A2aHeaders::decode_from(&msg.headers);
        let response = match self
            .authenticator
            .authenticate(&headers, &msg.payload)
            .await
        {
            Ok(identity) => {
                let ctx = RequestContext::new(headers, Some(identity));
                self.dispatch(&ctx, &msg.payload).await
            }
            Err(e) => {
                warn!(target: "security", error = %e, "a2a auth rejected");
                let req_id = serde_json::from_slice::<JsonRpcRequest>(&msg.payload)
                    .map(|r| r.id)
                    .unwrap_or(Value::Null);
                JsonRpcResponse {
                    jsonrpc: "2.0".into(),
                    id: req_id,
                    result: None,
                    error: Some(JsonRpcError {
                        code: JSON_RPC_UNAUTHENTICATED,
                        message: "Unauthenticated".into(),
                        data: None,
                    }),
                }
            }
        };
        let bytes = match serde_json::to_vec(&response) {
            Ok(b) => Bytes::from(b),
            Err(e) => {
                error!(error = %e, "encode response");
                let _ = msg.ack.ack().await;
                return;
            }
        };
        if let Err(e) = self.pubsub.publish(&reply_to, bytes, Headers::new()).await {
            error!(error = %e, "publish reply");
        }
        let _ = msg.ack.ack().await;
    }

    async fn dispatch(&self, ctx: &RequestContext, payload: &[u8]) -> JsonRpcResponse {
        let req: JsonRpcRequest = match serde_json::from_slice(payload) {
            Ok(r) => r,
            Err(e) => return A2aError::Json(e).to_json_rpc_error(Value::Null),
        };
        let method = match A2aMethod::from_str(&req.method) {
            Ok(m) => m,
            Err(e) => return e.to_json_rpc_error(req.id.clone()),
        };
        let id = req.id.clone();
        let params_raw = req.params;
        let result: Result<Value, A2aError> = match method {
            A2aMethod::SendMessage => {
                run_method::<SendMessageParams, _, _>(params_raw, |p| async move {
                    self.handler.send_message(ctx, p).await
                })
                .await
            }
            A2aMethod::SendStreamingMessage => {
                run_method::<SendMessageParams, _, _>(params_raw, |p| async move {
                    self.handler.send_streaming_message(ctx, p).await
                })
                .await
            }
            A2aMethod::GetTask => {
                run_method::<GetTaskParams, _, _>(params_raw, |p| async move {
                    self.handler.get_task(ctx, p).await
                })
                .await
            }
            A2aMethod::ListTasks => {
                run_method::<ListTasksParams, _, _>(params_raw, |p| async move {
                    self.handler.list_tasks(ctx, p).await
                })
                .await
            }
            A2aMethod::CancelTask => {
                run_method::<CancelTaskParams, _, _>(params_raw, |p| async move {
                    self.handler.cancel_task(ctx, p).await
                })
                .await
            }
            A2aMethod::SubscribeToTask => {
                run_method::<SubscribeToTaskParams, _, _>(params_raw, |p| async move {
                    self.handler.subscribe_to_task(ctx, p).await
                })
                .await
            }
            A2aMethod::CreateTaskPushNotificationConfig => {
                run_method::<PushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.handler
                        .create_task_push_notification_config(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::GetTaskPushNotificationConfig => {
                run_method::<GetPushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.handler.get_task_push_notification_config(ctx, p).await
                })
                .await
            }
            A2aMethod::ListTaskPushNotificationConfigs => {
                run_method::<ListPushNotificationConfigsParams, _, _>(params_raw, |p| async move {
                    self.handler
                        .list_task_push_notification_configs(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::DeleteTaskPushNotificationConfig => {
                run_method::<DeletePushNotificationConfigParams, _, _>(params_raw, |p| async move {
                    self.handler
                        .delete_task_push_notification_config(ctx, p)
                        .await
                })
                .await
            }
            A2aMethod::GetExtendedAgentCard => {
                run_method::<GetExtendedAgentCardParams, _, _>(params_raw, |p| async move {
                    self.handler.get_extended_agent_card(ctx, p).await
                })
                .await
            }
        };
        match result {
            Ok(value) => JsonRpcResponse {
                jsonrpc: "2.0".into(),
                id,
                result: Some(value),
                error: None,
            },
            Err(e) => e.to_json_rpc_error(id),
        }
    }
}

/// Helper: deserialise params from raw JSON, run the handler closure,
/// then serialise the typed result back to a JSON [`Value`].
async fn run_method<P, T, Fut>(raw: Value, run: impl FnOnce(P) -> Fut) -> Result<Value, A2aError>
where
    P: DeserializeOwned,
    T: Serialize,
    Fut: std::future::Future<Output = Result<T, A2aError>>,
{
    let params: P =
        serde_json::from_value(raw).map_err(|e| A2aError::InvalidParams(e.to_string()))?;
    let typed = run(params).await?;
    serde_json::to_value(typed).map_err(A2aError::from)
}