a2a_rs/adapter/business/

request_processor.rs

//! A default request processor implementation

// This module is already conditionally compiled with #[cfg(feature = "server")] in mod.rs

use std::sync::Arc;

use async_trait::async_trait;

use crate::{
    application::{
        JSONRPCError, JSONRPCResponse,
        json_rpc::{
            self, A2ARequest, CancelTaskRequest, GetExtendedCardRequest,
            GetTaskPushNotificationRequest, GetTaskRequest, SendTaskRequest,
            SendTaskStreamingRequest, SetTaskPushNotificationRequest, TaskResubscriptionRequest,
        },
    },
    domain::A2AError,
    port::{AsyncMessageHandler, AsyncNotificationManager, AsyncTaskManager},
    services::server::{AgentInfoProvider, AsyncA2ARequestProcessor},
};

/// Default implementation of a request processor that routes requests to business handlers
#[derive(Clone)]
pub struct DefaultRequestProcessor<M, T, N, A = crate::adapter::SimpleAgentInfo>
where
    M: AsyncMessageHandler + Send + Sync + 'static,
    T: AsyncTaskManager + Send + Sync + 'static,
    N: AsyncNotificationManager + Send + Sync + 'static,
    A: AgentInfoProvider + Send + Sync + 'static,
{
    /// Message handler
    message_handler: Arc<M>,
    /// Task manager
    task_manager: Arc<T>,
    /// Notification manager
    notification_manager: Arc<N>,
    /// Agent info provider
    agent_info: Arc<A>,
}

impl<M, T, N, A> DefaultRequestProcessor<M, T, N, A>
where
    M: AsyncMessageHandler + Send + Sync + 'static,
    T: AsyncTaskManager + Send + Sync + 'static,
    N: AsyncNotificationManager + Send + Sync + 'static,
    A: AgentInfoProvider + Send + Sync + 'static,
{
    /// Create a new request processor with the given handlers
    pub fn new(
        message_handler: M,
        task_manager: T,
        notification_manager: N,
        agent_info: A,
    ) -> Self {
        Self {
            message_handler: Arc::new(message_handler),
            task_manager: Arc::new(task_manager),
            notification_manager: Arc::new(notification_manager),
            agent_info: Arc::new(agent_info),
        }
    }
}

impl<H, A> DefaultRequestProcessor<H, H, H, A>
where
    H: AsyncMessageHandler + AsyncTaskManager + AsyncNotificationManager + Send + Sync + 'static,
    A: AgentInfoProvider + Send + Sync + 'static,
{
    /// Create a new request processor with a single handler that implements all traits
    pub fn with_handler(handler: H, agent_info: A) -> Self {
        let handler_arc = Arc::new(handler);
        Self {
            message_handler: handler_arc.clone(),
            task_manager: handler_arc.clone(),
            notification_manager: handler_arc,
            agent_info: Arc::new(agent_info),
        }
    }
}

impl<M, T, N, A> DefaultRequestProcessor<M, T, N, A>
where
    M: AsyncMessageHandler + Send + Sync + 'static,
    T: AsyncTaskManager + Send + Sync + 'static,
    N: AsyncNotificationManager + Send + Sync + 'static,
    A: AgentInfoProvider + Send + Sync + 'static,
{
    /// Process a send task request
    async fn process_send_task(
        &self,
        request: &SendTaskRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let params = &request.params;
        let session_id = params.session_id.as_deref();

        tracing::info!(
            task_id = %params.id,
            message_id = %params.message.message_id,
            "🔄 DefaultRequestProcessor: About to call message_handler.process_message"
        );

        // Process the message through the handler
        // The handler is responsible for managing history
        let task = self
            .message_handler
            .process_message(&params.id, &params.message, session_id)
            .await?;

        tracing::info!(
            task_id = %params.id,
            "✅ DefaultRequestProcessor: Message handler returned successfully"
        );

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(task)?,
        ))
    }

    /// Process a get task request
    async fn process_get_task(
        &self,
        request: &GetTaskRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let params = &request.params;
        let task = self
            .task_manager
            .get_task(&params.id, params.history_length)
            .await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(task)?,
        ))
    }

    /// Process a cancel task request
    async fn process_cancel_task(
        &self,
        request: &CancelTaskRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let params = &request.params;
        let task = self.task_manager.cancel_task(&params.id).await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(task)?,
        ))
    }

    /// Process a set task push notification request
    async fn process_set_push_notification(
        &self,
        request: &SetTaskPushNotificationRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let config = self
            .notification_manager
            .set_task_notification(&request.params)
            .await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(config)?,
        ))
    }

    /// Process a get task push notification request
    async fn process_get_push_notification(
        &self,
        request: &GetTaskPushNotificationRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let params = &request.params;
        let config = self
            .notification_manager
            .get_task_notification(&params.id)
            .await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(config)?,
        ))
    }

    /// Process a task resubscription request
    async fn process_task_resubscription(
        &self,
        request: &TaskResubscriptionRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        // For resubscription, we return an initial success response,
        // and then the streaming updates are handled separately
        let params = &request.params;

        // Try to get the task, but don't fail if it doesn't exist
        // This allows clients to subscribe to tasks before they're created
        match self
            .task_manager
            .get_task(&params.id, params.history_length)
            .await
        {
            Ok(task) => {
                // Task exists, return it
                Ok(JSONRPCResponse::success(
                    request.id.clone(),
                    serde_json::to_value(task)?,
                ))
            }
            Err(A2AError::TaskNotFound(_)) => {
                // Task doesn't exist yet, return null result
                // The WebSocket server will still set up subscriptions
                // and send updates when the task is created
                Ok(JSONRPCResponse::success(
                    request.id.clone(),
                    serde_json::Value::Null,
                ))
            }
            Err(e) => {
                // Other errors should still be propagated
                Err(e)
            }
        }
    }

    /// Process a send task streaming request
    async fn process_send_task_streaming(
        &self,
        request: &SendTaskStreamingRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        // For streaming, we process the message and return an initial success response,
        // and then the streaming updates are handled separately
        let params = &request.params;
        let session_id = params.session_id.as_deref();

        // Process the message through the handler
        // The handler is responsible for managing history
        let task = self
            .message_handler
            .process_message(&params.id, &params.message, session_id)
            .await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(task)?,
        ))
    }

    /// Process a get extended card request (v0.3.0)
    async fn process_get_extended_card(
        &self,
        request: &GetExtendedCardRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        // Get the agent card from the agent info provider
        // For v0.3.0, this method should return extended information
        // that may only be available to authenticated clients.
        // Authentication checking should be handled by middleware.
        let card = self.agent_info.get_agent_card().await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(card)?,
        ))
    }

    // ===== v0.3.0 New Methods =====

    async fn process_list_tasks(
        &self,
        request: &crate::application::handlers::task::ListTasksRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let default_params = crate::domain::ListTasksParams::default();
        let params = request.params.as_ref().unwrap_or(&default_params);
        let result = self.task_manager.list_tasks_v3(params).await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(result)?,
        ))
    }

    async fn process_get_push_notification_config(
        &self,
        request: &crate::application::handlers::task::GetTaskPushNotificationConfigRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        if let Some(ref params) = request.params {
            let result = self
                .task_manager
                .get_push_notification_config(params)
                .await?;

            Ok(JSONRPCResponse::success(
                request.id.clone(),
                serde_json::to_value(result)?,
            ))
        } else {
            Err(A2AError::InvalidParams(
                "Missing params for get push notification config".to_string(),
            ))
        }
    }

    async fn process_list_push_notification_configs(
        &self,
        request: &crate::application::handlers::task::ListTaskPushNotificationConfigRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        let result = self
            .task_manager
            .list_push_notification_configs(&request.params)
            .await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(result)?,
        ))
    }

    async fn process_delete_push_notification_config(
        &self,
        request: &crate::application::handlers::task::DeleteTaskPushNotificationConfigRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        self.task_manager
            .delete_push_notification_config(&request.params)
            .await?;

        // Return null on success
        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::Value::Null,
        ))
    }

    async fn process_get_authenticated_extended_card(
        &self,
        request: &crate::application::handlers::agent::GetAuthenticatedExtendedCardRequest,
    ) -> Result<JSONRPCResponse, A2AError> {
        // Get the authenticated extended card from the agent info provider
        // Authentication checking should be handled by middleware before this point
        let card = self.agent_info.get_authenticated_extended_card().await?;

        Ok(JSONRPCResponse::success(
            request.id.clone(),
            serde_json::to_value(card)?,
        ))
    }
}

#[async_trait]
impl<M, T, N, A> AsyncA2ARequestProcessor for DefaultRequestProcessor<M, T, N, A>
where
    M: AsyncMessageHandler + Send + Sync + 'static,
    T: AsyncTaskManager + Send + Sync + 'static,
    N: AsyncNotificationManager + Send + Sync + 'static,
    A: AgentInfoProvider + Send + Sync + 'static,
{
    async fn process_raw_request(&self, request: &str) -> Result<String, A2AError> {
        // Parse the request
        let request = match json_rpc::parse_request(request) {
            Ok(req) => req,
            Err(e) => {
                // Return a JSON-RPC error response
                let error = JSONRPCError::from(e);
                let response = JSONRPCResponse::error(None, error);
                return Ok(serde_json::to_string(&response)?);
            }
        };

        // Process the request
        let response = match self.process_request(&request).await {
            Ok(resp) => resp,
            Err(e) => {
                // Return a JSON-RPC error response
                let error = JSONRPCError::from(e);
                let response = JSONRPCResponse::error(request.id().cloned(), error);
                return Ok(serde_json::to_string(&response)?);
            }
        };

        // Serialize the response
        Ok(serde_json::to_string(&response)?)
    }

    async fn process_request(&self, request: &A2ARequest) -> Result<JSONRPCResponse, A2AError> {
        match request {
            A2ARequest::SendTask(req) => self.process_send_task(req).await,
            A2ARequest::SendMessage(_req) => {
                // Convert MessageSendParams to TaskSendParams for backwards compatibility
                // TODO: Implement proper message handling
                Err(A2AError::UnsupportedOperation(
                    "Message sending not yet implemented".to_string(),
                ))
            }
            A2ARequest::GetTask(req) => self.process_get_task(req).await,
            A2ARequest::CancelTask(req) => self.process_cancel_task(req).await,
            A2ARequest::SetTaskPushNotification(req) => {
                self.process_set_push_notification(req).await
            }
            A2ARequest::GetTaskPushNotification(req) => {
                self.process_get_push_notification(req).await
            }
            A2ARequest::TaskResubscription(req) => self.process_task_resubscription(req).await,
            A2ARequest::SendTaskStreaming(req) => self.process_send_task_streaming(req).await,
            A2ARequest::SendMessageStreaming(_req) => {
                // Convert MessageSendParams to TaskSendParams for backwards compatibility
                // TODO: Implement proper message streaming
                Err(A2AError::UnsupportedOperation(
                    "Message streaming not yet implemented".to_string(),
                ))
            }
            A2ARequest::GetExtendedCard(req) => self.process_get_extended_card(req).await,
            // v0.3.0 new methods
            A2ARequest::ListTasks(req) => self.process_list_tasks(req).await,
            A2ARequest::GetTaskPushNotificationConfig(req) => {
                self.process_get_push_notification_config(req).await
            }
            A2ARequest::ListTaskPushNotificationConfigs(req) => {
                self.process_list_push_notification_configs(req).await
            }
            A2ARequest::DeleteTaskPushNotificationConfig(req) => {
                self.process_delete_push_notification_config(req).await
            }
            A2ARequest::GetAuthenticatedExtendedCard(req) => {
                self.process_get_authenticated_extended_card(req).await
            }
            A2ARequest::Generic(req) => {
                // Handle unknown method
                Err(A2AError::MethodNotFound(format!(
                    "Method '{}' not found",
                    req.method
                )))
            }
        }
    }
}