klieo-a2a 3.4.0

//! Typed client for calling an A2A agent over [`klieo_core::Pubsub`].
//!
//! # Why a typed client?
//!
//! Calling an A2A agent by hand is a six-step dance: subscribe a reply
//! inbox, build a [`JsonRpcRequest`], stamp `Reply-To` + `A2A-Version`
//! headers, publish to the agent's rpc subject, await one reply with a
//! timeout, then decode the [`JsonRpcResponse`]. [`A2aClient`] collapses
//! that to one typed call per A2A method and maps a JSON-RPC error
//! envelope onto a typed [`A2aError`].
//!
//! # Transport
//!
//! The [`crate::server::A2aServer`] replies by publishing on the subject
//! named in the request's `Reply-To` header (it does **not** use
//! [`klieo_core::RequestReply`] correlation). So this client wraps
//! `Arc<dyn Pubsub>`: it subscribes ONE reply inbox per client instance,
//! advertises it in `Reply-To`, and awaits answers there.
//!
//! # Why one reply consumer per instance (not per call)
//!
//! On the NATS backend, [`Pubsub::subscribe`] idempotently creates a
//! *durable* JetStream pull consumer keyed by the [`DurableName`], and NATS
//! never auto-deletes it. A per-call durable name therefore leaks one
//! consumer per RPC — unbounded under load (the in-memory bus ignores the
//! name, so the leak only bites in production). The trait exposes no
//! ephemeral-subscribe and no unsubscribe, so each client instead holds ONE
//! fixed reply subject + durable name and serialises its in-flight requests
//! so the single shared consumer is never read by two calls at once.

use crate::envelope::{JsonRpcError, JsonRpcRequest, JsonRpcResponse};
use crate::error::A2aError;
use crate::types::{CancelTaskParams, GetTaskParams, Message, SendMessageResult, Task};
use bytes::Bytes;
use klieo_core::{DurableName, Headers, Pubsub};
use serde::de::DeserializeOwned;
use serde_json::{json, Value};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio_stream::StreamExt;

/// JSON-RPC protocol version stamped on every request.
const JSONRPC_VERSION: &str = "2.0";
/// A2A protocol version advertised in the `A2A-Version` header.
const A2A_VERSION: &str = "1.0";

/// Process-global sequence handing each `A2aClient` a distinct reply
/// subject + durable name. Names are unique within a process — the
/// realistic case — so this avoids a uuid/rand dependency.
static CLIENT_INSTANCE_SEQ: AtomicU64 = AtomicU64::new(0);

/// Typed caller for a remote A2A agent over a [`Pubsub`] transport.
///
/// One client can address many agents — each method takes the target
/// `agent_id`. Cloning the `Arc<dyn Pubsub>` is cheap; share a single
/// `A2aClient` across tasks (it is `Send + Sync`).
///
/// ponytail: one durable reply consumer per client, serialized to one
/// in-flight request via `in_flight`. For concurrent pipelining, construct
/// multiple `A2aClient`s (each = one consumer) or add an id-demux reader.
/// The per-instance id comes from a per-process counter; a cross-process
/// restart can reuse a name — same ceiling as the server's
/// `a2a-server-{agent_id}` durable.
pub struct A2aClient {
    transport: Arc<dyn Pubsub>,
    app_prefix: String,
    reply_subject: String,
    durable: DurableName,
    next_id: AtomicU64,
    in_flight: tokio::sync::Mutex<()>,
}

impl A2aClient {
    /// Build a client publishing under `app_prefix` (matching the
    /// `A2aServer`'s prefix so subjects line up).
    pub fn new(transport: Arc<dyn Pubsub>, app_prefix: impl Into<String>) -> Self {
        let app_prefix = app_prefix.into();
        let instance = CLIENT_INSTANCE_SEQ.fetch_add(1, Ordering::Relaxed);
        Self {
            reply_subject: format!("{app_prefix}.a2a.reply.{instance}"),
            durable: DurableName::new(format!("a2a-client-{instance}")),
            app_prefix,
            transport,
            next_id: AtomicU64::new(1),
            in_flight: tokio::sync::Mutex::new(()),
        }
    }

    /// Send a message to `agent_id` and decode the agent's
    /// [`SendMessageResult`] (a completed `Message` or a long-running
    /// `Task`). Errors if the peer returns a JSON-RPC error, the reply
    /// times out, or the transport fails.
    pub async fn send_message(
        &self,
        agent_id: &str,
        message: Message,
        timeout: Duration,
    ) -> Result<SendMessageResult, A2aError> {
        let params = json!({ "message": message });
        let response = self
            .request_one(agent_id, "SendMessage", params, timeout)
            .await?;
        decode_result(response)
    }

    /// Fetch a task by id from `agent_id`. Errors with
    /// [`A2aError::Rpc`] carrying the peer's code when the task is unknown
    /// or access is denied.
    pub async fn get_task(
        &self,
        agent_id: &str,
        params: GetTaskParams,
        timeout: Duration,
    ) -> Result<Task, A2aError> {
        let raw = serde_json::to_value(params)?;
        let response = self.request_one(agent_id, "GetTask", raw, timeout).await?;
        decode_result(response)
    }

    /// Cancel a task on `agent_id`, returning the task in its post-cancel
    /// state. Errors with [`A2aError::Rpc`] when the task is unknown or
    /// access is denied.
    pub async fn cancel_task(
        &self,
        agent_id: &str,
        params: CancelTaskParams,
        timeout: Duration,
    ) -> Result<Task, A2aError> {
        let raw = serde_json::to_value(params)?;
        let response = self
            .request_one(agent_id, "CancelTask", raw, timeout)
            .await?;
        decode_result(response)
    }

    async fn request_one(
        &self,
        agent_id: &str,
        method: &str,
        params: Value,
        timeout: Duration,
    ) -> Result<JsonRpcResponse, A2aError> {
        // Serialize so the single shared reply consumer is never read by
        // two calls at once (a durable pull consumer load-balances across
        // concurrent readers, which would let one call steal another's
        // reply).
        let _guard = self.in_flight.lock().await;
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let mut reply_stream = self
            .transport
            .subscribe(&self.reply_subject, self.durable.clone())
            .await?;
        let request = JsonRpcRequest {
            jsonrpc: JSONRPC_VERSION.to_string(),
            id: json!(id),
            method: method.to_string(),
            params,
        };
        self.publish_request(agent_id, &request).await?;
        let reply = await_reply(&mut reply_stream, timeout, id).await?;
        check_for_error(reply)
    }

    async fn publish_request(
        &self,
        agent_id: &str,
        request: &JsonRpcRequest,
    ) -> Result<(), A2aError> {
        let subject = format!("{}.a2a.{agent_id}.rpc", self.app_prefix);
        let payload = Bytes::from(serde_json::to_vec(request)?);
        let mut headers = Headers::new();
        headers.insert("Reply-To".into(), self.reply_subject.clone());
        headers.insert("A2A-Version".into(), A2A_VERSION.into());
        self.transport.publish(&subject, payload, headers).await?;
        Ok(())
    }
}

/// Await the reply correlated to request `id` on `stream` within `timeout`,
/// skipping any reply whose JSON-RPC id does not match (a shared per-instance
/// inbox may carry a stale answer from an earlier call). The whole skip-loop
/// runs under one `timeout` budget so a flood of mismatched replies cannot
/// extend the deadline. Maps a deadline overrun to [`A2aError::Timeout`] and
/// a closed stream to [`A2aError::Server`] — never silently returns another
/// call's reply.
async fn await_reply(
    stream: &mut klieo_core::MsgStream,
    timeout: Duration,
    id: u64,
) -> Result<JsonRpcResponse, A2aError> {
    let matched = tokio::time::timeout(timeout, async {
        loop {
            let msg = match stream.next().await {
                Some(Ok(msg)) => msg,
                Some(Err(bus_err)) => return Err(A2aError::Bus(bus_err)),
                None => return Err(A2aError::Server("reply stream closed before reply".into())),
            };
            let response: JsonRpcResponse = serde_json::from_slice(&msg.payload)?;
            if response.id == json!(id) {
                return Ok(response);
            }
        }
    })
    .await;
    matched.unwrap_or_else(|_| Err(A2aError::Timeout(timeout)))
}

/// Turn a JSON-RPC `error` envelope into [`A2aError::Rpc`]; otherwise pass
/// the response through for result decoding.
fn check_for_error(response: JsonRpcResponse) -> Result<JsonRpcResponse, A2aError> {
    if let Some(JsonRpcError { code, message, .. }) = response.error {
        return Err(A2aError::Rpc { code, message });
    }
    Ok(response)
}

/// A missing `result` on an error-free response is a malformed reply.
fn decode_result<T: DeserializeOwned>(response: JsonRpcResponse) -> Result<T, A2aError> {
    let result = response
        .result
        .ok_or_else(|| A2aError::Server("reply had neither result nor error".into()))?;
    Ok(serde_json::from_value(result)?)
}