klieo_a2a/

handler.rs

//! `A2aHandler` trait + `EchoHandler` test fixture.
//!
//! The trait has 11 methods — one per A2A v1.0 spec §9.4 op — each
//! defaulted to return [`A2aError::MethodNotFound`]. Implementors
//! override only the methods their agent supports; the dispatcher
//! cleanly returns -32601 for everything else.

use crate::auth::RequestContext;
use crate::error::A2aError;
use crate::server::TaskEventStream;
use crate::types::{
    AgentCard, CancelTaskParams, DeletePushNotificationConfigParams, GetExtendedAgentCardParams,
    GetPushNotificationConfigParams, GetTaskParams, ListPushNotificationConfigsParams,
    ListPushNotificationConfigsResult, ListTasksParams, ListTasksResult, PushNotificationConfig,
    PushNotificationConfigParams, SendMessageParams, SendMessageResult, SubscribeToTaskParams,
    Task,
};
use async_trait::async_trait;

#[cfg(any(test, feature = "test-fixtures"))]
use crate::types::{AgentCapabilities, Message, Role, TaskStatus};
#[cfg(any(test, feature = "test-fixtures"))]
use std::collections::HashMap;
#[cfg(any(test, feature = "test-fixtures"))]
use std::sync::Mutex;
#[cfg(any(test, feature = "test-fixtures"))]
use uuid::Uuid;

/// Server-side handler implementing the 11 A2A v1.0 operations.
///
/// Default implementations return [`A2aError::MethodNotFound`]. Override
/// only the methods your agent supports; `A2aServer` translates the
/// error into a JSON-RPC -32601 response automatically.
#[async_trait]
pub trait A2aHandler: Send + Sync {
    /// Spec §9.4.1. `ctx.caller` is the verified identity (or `None`
    /// only when `AllowAnonymous` was wired). Implementations should
    /// reject `None` for mutating methods.
    async fn send_message(
        &self,
        _ctx: &RequestContext,
        _params: SendMessageParams,
    ) -> Result<SendMessageResult, A2aError> {
        Err(A2aError::MethodNotFound("SendMessage".into()))
    }
    /// Spec §9.4.2. Streaming variant — v0.0.2. Return a
    /// [`TaskEventStream`] to override the dispatcher's default
    /// broadcast-derived stream; return [`A2aError::MethodNotFound`]
    /// to fall back to it.
    async fn send_streaming_message(
        &self,
        _ctx: &RequestContext,
        _params: SendMessageParams,
    ) -> Result<TaskEventStream, A2aError> {
        Err(A2aError::MethodNotFound("SendStreamingMessage".into()))
    }
    /// Spec §9.4.3.
    async fn get_task(
        &self,
        _ctx: &RequestContext,
        _params: GetTaskParams,
    ) -> Result<Task, A2aError> {
        Err(A2aError::MethodNotFound("GetTask".into()))
    }
    /// Spec §9.4.4.
    async fn list_tasks(
        &self,
        _ctx: &RequestContext,
        _params: ListTasksParams,
    ) -> Result<ListTasksResult, A2aError> {
        Err(A2aError::MethodNotFound("ListTasks".into()))
    }
    /// Spec §9.4.5.
    async fn cancel_task(
        &self,
        _ctx: &RequestContext,
        _params: CancelTaskParams,
    ) -> Result<Task, A2aError> {
        Err(A2aError::MethodNotFound("CancelTask".into()))
    }
    /// Spec §9.4.6. v0.0.2. Return a [`TaskEventStream`] to override
    /// the dispatcher's default broadcast-derived stream; return
    /// [`A2aError::MethodNotFound`] to fall back to it.
    async fn subscribe_to_task(
        &self,
        _ctx: &RequestContext,
        _params: SubscribeToTaskParams,
    ) -> Result<TaskEventStream, A2aError> {
        Err(A2aError::MethodNotFound("SubscribeToTask".into()))
    }
    /// Spec §9.4.7. v0.0.2.
    async fn create_task_push_notification_config(
        &self,
        _ctx: &RequestContext,
        _params: PushNotificationConfigParams,
    ) -> Result<PushNotificationConfig, A2aError> {
        Err(A2aError::MethodNotFound(
            "CreateTaskPushNotificationConfig".into(),
        ))
    }
    /// Spec §9.4.8. v0.0.2.
    async fn get_task_push_notification_config(
        &self,
        _ctx: &RequestContext,
        _params: GetPushNotificationConfigParams,
    ) -> Result<PushNotificationConfig, A2aError> {
        Err(A2aError::MethodNotFound(
            "GetTaskPushNotificationConfig".into(),
        ))
    }
    /// Spec §9.4.9. v0.0.2.
    async fn list_task_push_notification_configs(
        &self,
        _ctx: &RequestContext,
        _params: ListPushNotificationConfigsParams,
    ) -> Result<ListPushNotificationConfigsResult, A2aError> {
        Err(A2aError::MethodNotFound(
            "ListTaskPushNotificationConfigs".into(),
        ))
    }
    /// Spec §9.4.10. v0.0.2.
    async fn delete_task_push_notification_config(
        &self,
        _ctx: &RequestContext,
        _params: DeletePushNotificationConfigParams,
    ) -> Result<(), A2aError> {
        Err(A2aError::MethodNotFound(
            "DeleteTaskPushNotificationConfig".into(),
        ))
    }
    /// Spec §9.4.11.
    async fn get_extended_agent_card(
        &self,
        _ctx: &RequestContext,
        _params: GetExtendedAgentCardParams,
    ) -> Result<AgentCard, A2aError> {
        Err(A2aError::MethodNotFound("GetExtendedAgentCard".into()))
    }

    /// The PUBLIC agent card served over
    /// `GET /.well-known/agent-card.json` (A2A spec §5, IANA
    /// registration in §12 of the spec's well-known-URI appendix).
    ///
    /// Deliberately a SEPARATE method from [`Self::get_extended_agent_card`]:
    /// the well-known endpoint is unauthenticated BY PROTOCOL DESIGN (a
    /// client fetches it before it knows which auth scheme to use), so
    /// `ctx.caller` is always `None` here — implementations MUST NOT
    /// return anything in this card that isn't safe to hand to an
    /// anonymous caller. Richer, auth-gated detail belongs in
    /// [`Self::get_extended_agent_card`] instead, which only runs after
    /// the dispatcher's `Authenticator` has verified the caller.
    async fn get_agent_card(&self, _ctx: &RequestContext) -> Result<AgentCard, A2aError> {
        Err(A2aError::MethodNotFound("GetAgentCard".into()))
    }

    /// Returns true if all streaming invocations against this
    /// handler are safe to re-execute from the original payload
    /// — i.e. the result is deterministic OR duplicate execution
    /// has no business effect.
    ///
    /// Default false (safe). Mark true ONLY when the handler is
    /// provably stateless / read-only or returns the same result
    /// for the same input.
    ///
    /// Cluster 0.24: when true, the follower replica detecting
    /// an orphaned stream re-invokes from the cached payload
    /// instead of writing a terminal "leader died" frame. See
    /// ADR-024.
    fn is_idempotent(&self) -> bool {
        false
    }
}

/// **TEST FIXTURE — DO NOT WIRE INTO PRODUCTION.**
///
/// `EchoHandler` is an unauthenticated, in-memory `A2aHandler` used by
/// the crate's own integration tests, the conformance suite, and the
/// `cargo klieo a2a-conformance` CLI. It accepts every inbound
/// `SendMessage` regardless of caller identity, advertises empty
/// `securitySchemes` on its agent card, and stores tasks in process
/// memory only.
///
/// Wiring this handler into an `A2aServer` outside a test context is a
/// security bug — any caller reachable on the transport can submit
/// tasks that are recorded as `Completed`.
///
/// Available only when the `test-fixtures` Cargo feature is enabled
/// (or in `cfg(test)` builds of `klieo-a2a` itself).
#[cfg(any(test, feature = "test-fixtures"))]
#[derive(Default)]
pub struct EchoHandler {
    tasks: Mutex<HashMap<String, Task>>,
}

#[cfg(any(test, feature = "test-fixtures"))]
#[async_trait]
impl A2aHandler for EchoHandler {
    async fn send_message(
        &self,
        _ctx: &RequestContext,
        params: SendMessageParams,
    ) -> Result<SendMessageResult, A2aError> {
        let task_id = format!("task-{}", Uuid::new_v4());
        let context_id = params
            .message
            .contextId
            .clone()
            .unwrap_or_else(|| format!("ctx-{}", Uuid::new_v4()));
        let response = Message {
            messageId: format!("msg-{}", Uuid::new_v4()),
            contextId: Some(context_id.clone()),
            taskId: Some(task_id.clone()),
            role: Role::Agent,
            parts: params.message.parts.clone(),
            metadata: None,
            extensions: vec![],
            referenceTaskIds: vec![],
        };
        let task = Task {
            id: task_id.clone(),
            contextId: context_id,
            status: TaskStatus::Completed,
            artifacts: vec![],
            history: vec![params.message, response],
            metadata: None,
        };
        self.tasks
            .lock()
            .map_err(|_| A2aError::Server("mutex poisoned".into()))?
            .insert(task_id.clone(), task.clone());
        Ok(SendMessageResult::Task(task))
    }

    async fn get_task(
        &self,
        _ctx: &RequestContext,
        params: GetTaskParams,
    ) -> Result<Task, A2aError> {
        self.tasks
            .lock()
            .map_err(|_| A2aError::Server("mutex poisoned".into()))?
            .get(&params.id)
            .cloned()
            .ok_or(A2aError::TaskNotFound(params.id))
    }

    async fn list_tasks(
        &self,
        _ctx: &RequestContext,
        params: ListTasksParams,
    ) -> Result<ListTasksResult, A2aError> {
        let guard = self
            .tasks
            .lock()
            .map_err(|_| A2aError::Server("mutex poisoned".into()))?;
        let tasks: Vec<Task> = guard
            .values()
            .filter(|t| match &params.contextId {
                Some(ctx) => &t.contextId == ctx,
                None => true,
            })
            .cloned()
            .collect();
        Ok(ListTasksResult {
            tasks,
            nextCursor: None,
        })
    }

    async fn cancel_task(
        &self,
        _ctx: &RequestContext,
        params: CancelTaskParams,
    ) -> Result<Task, A2aError> {
        let mut guard = self
            .tasks
            .lock()
            .map_err(|_| A2aError::Server("mutex poisoned".into()))?;
        let task = guard
            .get_mut(&params.id)
            .ok_or_else(|| A2aError::TaskNotFound(params.id.clone()))?;
        task.status = TaskStatus::Canceled;
        Ok(task.clone())
    }

    async fn get_extended_agent_card(
        &self,
        _ctx: &RequestContext,
        _params: GetExtendedAgentCardParams,
    ) -> Result<AgentCard, A2aError> {
        Ok(Self::static_card())
    }

    async fn get_agent_card(&self, _ctx: &RequestContext) -> Result<AgentCard, A2aError> {
        Ok(Self::static_card())
    }
}

#[cfg(any(test, feature = "test-fixtures"))]
impl EchoHandler {
    /// Shared card body for [`A2aHandler::get_extended_agent_card`] and
    /// [`A2aHandler::get_agent_card`]. `EchoHandler` has no auth-gated
    /// extra detail to withhold from anonymous callers, so both methods
    /// return the identical, fully-public card.
    fn static_card() -> AgentCard {
        AgentCard {
            id: "echo".into(),
            name: "Echo Agent".into(),
            description: Some("Echoes incoming messages as completed tasks.".into()),
            provider: None,
            capabilities: AgentCapabilities {
                streaming: false,
                pushNotifications: false,
                stateTransitionHistory: false,
            },
            skills: vec![],
            interfaces: vec![],
            securitySchemes: serde_json::Map::new(),
            security: vec![],
            extensions: vec![],
            signature: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::envelope::A2aHeaders;
    use crate::types::{
        CancelTaskParams, GetExtendedAgentCardParams, GetTaskParams, Message, Part, Role,
        SendMessageParams, TaskStatus,
    };

    struct AlwaysDefault;
    #[async_trait::async_trait]
    impl A2aHandler for AlwaysDefault {}

    fn anon_ctx() -> RequestContext {
        RequestContext::new(
            A2aHeaders::decode_from(&klieo_core::Headers::new()),
            Some(klieo_auth_common::Identity::anonymous()),
        )
    }

    #[tokio::test]
    async fn default_handler_returns_method_not_found() {
        let h = AlwaysDefault;
        let err = h
            .send_message(
                &anon_ctx(),
                SendMessageParams {
                    message: Message {
                        messageId: "m".into(),
                        contextId: None,
                        taskId: None,
                        role: Role::User,
                        parts: vec![],
                        metadata: None,
                        extensions: vec![],
                        referenceTaskIds: vec![],
                    },
                    configuration: None,
                },
            )
            .await
            .unwrap_err();
        assert!(matches!(err, A2aError::MethodNotFound(_)));
    }

    #[tokio::test]
    async fn echo_handler_send_message_returns_completed_task() {
        let h = EchoHandler::default();
        let msg = Message {
            messageId: "m-1".into(),
            contextId: Some("c-1".into()),
            taskId: None,
            role: Role::User,
            parts: vec![Part::Text {
                content: "hi".into(),
                metadata: None,
                mediaType: None,
                filename: None,
            }],
            metadata: None,
            extensions: vec![],
            referenceTaskIds: vec![],
        };
        let result = h
            .send_message(
                &anon_ctx(),
                SendMessageParams {
                    message: msg,
                    configuration: None,
                },
            )
            .await
            .unwrap();
        match result {
            crate::types::SendMessageResult::Task(t) => {
                assert!(matches!(t.status, TaskStatus::Completed));
                assert_eq!(t.history.len(), 2);
            }
            _ => panic!("expected Task"),
        }
    }

    #[tokio::test]
    async fn echo_handler_get_task_recovers_stored_task() {
        let h = EchoHandler::default();
        let msg = Message {
            messageId: "m-2".into(),
            contextId: Some("c-2".into()),
            taskId: None,
            role: Role::User,
            parts: vec![],
            metadata: None,
            extensions: vec![],
            referenceTaskIds: vec![],
        };
        let result = h
            .send_message(
                &anon_ctx(),
                SendMessageParams {
                    message: msg,
                    configuration: None,
                },
            )
            .await
            .unwrap();
        let task_id = match result {
            crate::types::SendMessageResult::Task(t) => t.id,
            _ => panic!(),
        };
        let recovered = h
            .get_task(
                &anon_ctx(),
                GetTaskParams {
                    id: task_id.clone(),
                    historyLength: None,
                },
            )
            .await
            .unwrap();
        assert_eq!(recovered.id, task_id);
    }

    #[tokio::test]
    async fn echo_handler_cancel_task_marks_canceled() {
        let h = EchoHandler::default();
        let msg = Message {
            messageId: "m-3".into(),
            contextId: Some("c-3".into()),
            taskId: None,
            role: Role::User,
            parts: vec![],
            metadata: None,
            extensions: vec![],
            referenceTaskIds: vec![],
        };
        let task_id = match h
            .send_message(
                &anon_ctx(),
                SendMessageParams {
                    message: msg,
                    configuration: None,
                },
            )
            .await
            .unwrap()
        {
            crate::types::SendMessageResult::Task(t) => t.id,
            _ => panic!(),
        };
        let cancelled = h
            .cancel_task(&anon_ctx(), CancelTaskParams { id: task_id })
            .await
            .unwrap();
        assert!(matches!(cancelled.status, TaskStatus::Canceled));
    }

    #[test]
    fn echo_handler_default_is_not_idempotent() {
        let h = EchoHandler::default();
        assert!(!h.is_idempotent());
    }

    #[test]
    fn always_default_handler_is_not_idempotent() {
        let h = AlwaysDefault;
        assert!(!h.is_idempotent());
    }

    #[tokio::test]
    async fn echo_handler_extended_agent_card_returns_static() {
        let h = EchoHandler::default();
        let card = h
            .get_extended_agent_card(&anon_ctx(), GetExtendedAgentCardParams::default())
            .await
            .unwrap();
        assert_eq!(card.id, "echo");
    }

    #[tokio::test]
    async fn echo_handler_get_agent_card_returns_static() {
        let h = EchoHandler::default();
        let card = h.get_agent_card(&anon_ctx()).await.unwrap();
        assert_eq!(card.id, "echo");
    }

    #[tokio::test]
    async fn default_handler_get_agent_card_returns_method_not_found() {
        let h = AlwaysDefault;
        let err = h.get_agent_card(&anon_ctx()).await.unwrap_err();
        assert!(matches!(err, A2aError::MethodNotFound(_)));
    }
}