velocia 0.3.1

velocia – production-ready AI agent framework using ADK-Rust, A2A protocol, and AWS DynamoDB
//! `A2AGatewayClient` – high-level client for sending messages to A2A agents.
//!
//! Mirrors Python's `A2AGatewayClient` from `a2a/client/app_client.py`.

use std::collections::HashMap;

use uuid::Uuid;

use crate::a2a::types::{AgentCard, Message, Part, Role, Task};
use crate::error::{AgentKitError, Result};

pub struct A2AGatewayClient {
    card: AgentCard,
    http: reqwest::Client,
    #[allow(dead_code)]
    owns_client: bool,
}

impl A2AGatewayClient {
    /// Create a client from a pre-fetched agent card.
    pub fn from_card(card: AgentCard) -> Self {
        Self {
            card,
            http: reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(60))
                .build()
                .expect("reqwest client"),
            owns_client: true,
        }
    }

    /// Discover an agent card from its URL and build a client.
    pub async fn from_card_url(
        agent_card_url: &str,
        headers: Option<HashMap<String, String>>,
    ) -> Result<Self> {
        let http = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(60))
            .build()
            .map_err(|e| AgentKitError::A2aClient(e.to_string()))?;

        let mut req = http.get(agent_card_url);
        if let Some(hdrs) = headers {
            for (k, v) in hdrs {
                req = req.header(k, v);
            }
        }

        let card = req
            .send()
            .await?
            .error_for_status()
            .map_err(|e| AgentKitError::A2aClient(e.to_string()))?
            .json::<AgentCard>()
            .await?;

        Ok(Self { card, http, owns_client: true })
    }

    pub fn supports_streaming(&self) -> bool {
        self.card.capabilities.streaming
    }

    /// Send a text message and return the resulting task.
    pub async fn send_message(
        &self,
        content: &str,
        context_id: Option<String>,
        propagation_headers: Option<HashMap<String, String>>,
    ) -> Result<Task> {
        let message = Message {
            kind: "message".to_string(),
            role: Role::User,
            parts: vec![Part::Text { text: content.to_string(), metadata: None }],
            message_id: Uuid::new_v4().to_string(),
            context_id,
            task_id: None,
            metadata: None,
        };

        let mut req = self.http.post(&self.card.url).json(&message);

        if let Some(hdrs) = propagation_headers {
            for (k, v) in hdrs {
                req = req.header(k, v);
            }
        }

        let task = req
            .send()
            .await?
            .error_for_status()
            .map_err(|e| AgentKitError::A2aClient(e.to_string()))?
            .json::<Task>()
            .await?;

        Ok(task)
    }

    /// Build propagation headers that forward tracing context to the remote.
    pub fn build_propagation_context(
        &self,
        incoming_headers: &HashMap<String, String>,
    ) -> HashMap<String, String> {
        incoming_headers
            .iter()
            .filter(|(k, _)| k.starts_with("x-velocia-"))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect()
    }

    pub fn agent_card(&self) -> &AgentCard {
        &self.card
    }
}