mur_core/a2a/

server.rs

//! A2A Server — expose Commander capabilities as an A2A-compatible agent.

use super::protocol::*;
use crate::types::Workflow;
use crate::workflow::parser;
use anyhow::Result;
use chrono::Utc;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

/// Maximum number of tasks stored in the task map before evicting terminal tasks.
const MAX_TASKS: usize = 1000;

/// A2A server that handles incoming task requests.
pub struct A2aServer {
    /// Active tasks keyed by task ID.
    tasks: Arc<RwLock<HashMap<String, TaskResponse>>>,
    /// Agent configuration.
    config: A2aServerConfig,
    /// Cached agent card (regenerated on demand).
    agent_card_cache: RwLock<Option<AgentCard>>,
}

/// Configuration for the A2A server.
pub struct A2aServerConfig {
    /// Agent name.
    pub name: String,
    /// Agent description.
    pub description: String,
    /// Base URL where this agent is reachable.
    pub base_url: String,
    /// Agent version.
    pub version: String,
}

impl Default for A2aServerConfig {
    fn default() -> Self {
        Self {
            name: "MUR Commander".into(),
            description: "Autonomous workflow execution agent".into(),
            base_url: "http://localhost:3939".into(),
            version: env!("CARGO_PKG_VERSION").into(),
        }
    }
}

impl A2aServer {
    /// Create a new A2A server.
    pub fn new(config: A2aServerConfig) -> Self {
        Self {
            tasks: Arc::new(RwLock::new(HashMap::new())),
            config,
            agent_card_cache: RwLock::new(None),
        }
    }

    /// Return the cached agent card, generating it on first call.
    /// Use [`invalidate_agent_card`] to force regeneration.
    pub fn agent_card(&self) -> AgentCard {
        // Fast path: return cached card if present.
        {
            let cache = self.agent_card_cache.read().unwrap_or_else(|e| e.into_inner());
            if let Some(ref card) = *cache {
                return card.clone();
            }
        }

        // Slow path: generate and cache.
        let card = self.generate_agent_card();
        {
            let mut cache = self.agent_card_cache.write().unwrap_or_else(|e| e.into_inner());
            *cache = Some(card.clone());
        }
        card
    }

    /// Invalidate the cached agent card so the next call to [`agent_card`]
    /// regenerates it (e.g. after workflows change on disk).
    pub fn invalidate_agent_card(&self) {
        let mut cache = self.agent_card_cache.write().unwrap_or_else(|e| e.into_inner());
        *cache = None;
    }

    /// Generate a fresh agent card from disk workflows.
    fn generate_agent_card(&self) -> AgentCard {
        let workflows = parser::load_all_workflows().unwrap_or_default();
        let skills = Self::workflows_to_skills(&workflows);

        AgentCard {
            name: self.config.name.clone(),
            description: self.config.description.clone(),
            url: format!("{}/a2a", self.config.base_url),
            version: self.config.version.clone(),
            protocol_version: "0.1".into(),
            capabilities: AgentCapabilities {
                streaming: true,
                push_notifications: false,
                state_management: true,
            },
            skills,
            authentication: None,
        }
    }

    /// Handle an incoming JSON-RPC request.
    pub fn handle_request(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
        match request.method.as_str() {
            methods::TASKS_SEND => self.handle_task_send(request),
            methods::TASKS_GET => self.handle_task_get(request),
            methods::TASKS_CANCEL => self.handle_task_cancel(request),
            _ => JsonRpcResponse::error(
                request.id.clone(),
                error_codes::METHOD_NOT_FOUND,
                &format!("Unknown method: {}", request.method),
            ),
        }
    }

    /// Handle `tasks/send` — create and start a task.
    fn handle_task_send(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
        let params = match &request.params {
            Some(p) => p,
            None => {
                return JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INVALID_PARAMS,
                    "Missing params",
                );
            }
        };

        let task_req: TaskRequest = match serde_json::from_value(params.clone()) {
            Ok(t) => t,
            Err(e) => {
                return JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INVALID_PARAMS,
                    &format!("Invalid task request: {}", e),
                );
            }
        };

        // Create the task response with "submitted" status
        let task_response = TaskResponse {
            id: task_req.id.clone(),
            status: TaskStatus {
                state: TaskState::Submitted,
                message: Some(TaskMessage {
                    role: MessageRole::Agent,
                    parts: vec![MessagePart::Text {
                        text: "Task received and queued for execution".into(),
                    }],
                }),
            },
            artifacts: vec![],
            history: task_req.messages.clone(),
            metadata: task_req.metadata.clone(),
        };

        // Store the task, evicting oldest terminal tasks if at capacity.
        match self.tasks.write() {
            Ok(mut tasks) => {
                if tasks.len() >= MAX_TASKS {
                    Self::evict_terminal_tasks(&mut tasks);
                }
                tasks.insert(task_req.id.clone(), task_response.clone());
            }
            Err(_) => {
                return JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INTERNAL_ERROR,
                    "Internal error: task store unavailable",
                );
            }
        }

        match serde_json::to_value(&task_response) {
            Ok(v) => JsonRpcResponse::success(request.id.clone(), v),
            Err(e) => JsonRpcResponse::error(
                request.id.clone(),
                error_codes::INTERNAL_ERROR,
                &format!("Serialization error: {}", e),
            ),
        }
    }

    /// Handle `tasks/get` — get the status of a task.
    fn handle_task_get(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
        let task_id = request
            .params
            .as_ref()
            .and_then(|p| p.get("id"))
            .and_then(|v| v.as_str());

        let task_id = match task_id {
            Some(id) => id,
            None => {
                return JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INVALID_PARAMS,
                    "Missing task id",
                );
            }
        };

        let tasks = self.tasks.read().unwrap_or_else(|e| e.into_inner());
        match tasks.get(task_id) {
            Some(task) => match serde_json::to_value(task) {
                Ok(v) => JsonRpcResponse::success(request.id.clone(), v),
                Err(e) => JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INTERNAL_ERROR,
                    &format!("Serialization error: {}", e),
                ),
            },
            None => JsonRpcResponse::error(
                request.id.clone(),
                error_codes::TASK_NOT_FOUND,
                &format!("Task '{}' not found", task_id),
            ),
        }
    }

    /// Handle `tasks/cancel` — cancel a running task.
    fn handle_task_cancel(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
        let task_id = request
            .params
            .as_ref()
            .and_then(|p| p.get("id"))
            .and_then(|v| v.as_str());

        let task_id = match task_id {
            Some(id) => id,
            None => {
                return JsonRpcResponse::error(
                    request.id.clone(),
                    error_codes::INVALID_PARAMS,
                    "Missing task id",
                );
            }
        };

        let mut tasks = self.tasks.write().unwrap_or_else(|e| e.into_inner());
        match tasks.get_mut(task_id) {
            Some(task) => {
                // Only tasks in Submitted or Working states can be canceled.
                match task.status.state {
                    TaskState::Submitted | TaskState::Working => {}
                    _ => {
                        return JsonRpcResponse::error(
                            request.id.clone(),
                            error_codes::INVALID_PARAMS,
                            &format!(
                                "Task '{}' cannot be canceled in state {:?}",
                                task_id, task.status.state
                            ),
                        );
                    }
                }

                task.status = TaskStatus {
                    state: TaskState::Canceled,
                    message: Some(TaskMessage {
                        role: MessageRole::Agent,
                        parts: vec![MessagePart::Text {
                            text: "Task canceled".into(),
                        }],
                    }),
                };

                let update = TaskStatusUpdate {
                    id: task_id.to_string(),
                    status: task.status.clone(),
                    final_update: true,
                    timestamp: Utc::now(),
                };

                match serde_json::to_value(&update) {
                    Ok(v) => JsonRpcResponse::success(request.id.clone(), v),
                    Err(e) => JsonRpcResponse::error(
                        request.id.clone(),
                        error_codes::INTERNAL_ERROR,
                        &format!("Serialization error: {}", e),
                    ),
                }
            }
            None => JsonRpcResponse::error(
                request.id.clone(),
                error_codes::TASK_NOT_FOUND,
                &format!("Task '{}' not found", task_id),
            ),
        }
    }

    /// Update the status of a task (called by workflow runner).
    pub fn update_task_status(
        &self,
        task_id: &str,
        state: TaskState,
        message: Option<String>,
        artifacts: Option<Vec<TaskArtifact>>,
    ) -> Result<()> {
        let mut tasks = self
            .tasks
            .write()
            .map_err(|_| anyhow::anyhow!("Lock poisoned"))?;

        let task = tasks
            .get_mut(task_id)
            .ok_or_else(|| anyhow::anyhow!("Task not found: {}", task_id))?;

        task.status = TaskStatus {
            state,
            message: message.map(|text| TaskMessage {
                role: MessageRole::Agent,
                parts: vec![MessagePart::Text { text }],
            }),
        };

        if let Some(arts) = artifacts {
            task.artifacts = arts;
        }

        Ok(())
    }

    /// List all active tasks.
    pub fn list_tasks(&self) -> Vec<TaskResponse> {
        let tasks = self.tasks.read().unwrap_or_else(|e| e.into_inner());
        tasks.values().cloned().collect()
    }

    /// Evict terminal (Completed, Failed, Canceled) tasks to make room.
    /// Removes all terminal tasks; if none exist the map is left as-is.
    fn evict_terminal_tasks(tasks: &mut HashMap<String, TaskResponse>) {
        let terminal_ids: Vec<String> = tasks
            .iter()
            .filter(|(_, t)| matches!(
                t.status.state,
                TaskState::Completed | TaskState::Failed | TaskState::Canceled
            ))
            .map(|(id, _)| id.clone())
            .collect();

        for id in terminal_ids {
            tasks.remove(&id);
        }
    }

    /// Convert workflows to A2A skills.
    fn workflows_to_skills(workflows: &[Workflow]) -> Vec<AgentSkill> {
        workflows
            .iter()
            .map(|w| AgentSkill {
                id: w.id.clone(),
                name: w.name.clone(),
                description: w.description.clone(),
                tags: vec!["workflow".into()],
                input_schema: Some(serde_json::json!({
                    "type": "object",
                    "properties": {
                        "variables": {
                            "type": "object",
                            "description": "Workflow variables",
                            "properties": w.variables.iter().map(|(k, v)| {
                                (k.clone(), serde_json::json!({
                                    "type": "string",
                                    "default": v,
                                }))
                            }).collect::<HashMap<_, _>>(),
                        },
                        "shadow": {
                            "type": "boolean",
                            "default": false,
                            "description": "Run in dry-run mode",
                        },
                    },
                })),
                output_schema: None,
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_server() -> A2aServer {
        A2aServer::new(A2aServerConfig::default())
    }

    #[test]
    fn test_agent_card_generation() {
        let server = test_server();
        let card = server.agent_card();
        assert_eq!(card.name, "MUR Commander");
        assert!(card.capabilities.streaming);
        assert_eq!(card.protocol_version, "0.1");
    }

    #[test]
    fn test_handle_unknown_method() {
        let server = test_server();
        let req = JsonRpcRequest::new("unknown/method", None, serde_json::json!(1));
        let resp = server.handle_request(&req);
        assert!(resp.error.is_some());
        assert_eq!(resp.error.unwrap().code, error_codes::METHOD_NOT_FOUND);
    }

    #[test]
    fn test_handle_task_send() {
        let server = test_server();
        let task_req = TaskRequest {
            id: "test-task-1".into(),
            skill_id: None,
            messages: vec![TaskMessage {
                role: MessageRole::User,
                parts: vec![MessagePart::Text {
                    text: "Hello".into(),
                }],
            }],
            metadata: HashMap::new(),
        };
        let req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(1),
        );
        let resp = server.handle_request(&req);
        assert!(resp.error.is_none());
        assert!(resp.result.is_some());
    }

    #[test]
    fn test_handle_task_get() {
        let server = test_server();

        // Send a task first
        let task_req = TaskRequest {
            id: "test-task-2".into(),
            skill_id: None,
            messages: vec![],
            metadata: HashMap::new(),
        };
        let send_req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(1),
        );
        server.handle_request(&send_req);

        // Get the task
        let get_req = JsonRpcRequest::new(
            methods::TASKS_GET,
            Some(serde_json::json!({"id": "test-task-2"})),
            serde_json::json!(2),
        );
        let resp = server.handle_request(&get_req);
        assert!(resp.error.is_none());
    }

    #[test]
    fn test_handle_task_get_not_found() {
        let server = test_server();
        let req = JsonRpcRequest::new(
            methods::TASKS_GET,
            Some(serde_json::json!({"id": "nonexistent"})),
            serde_json::json!(1),
        );
        let resp = server.handle_request(&req);
        assert!(resp.error.is_some());
        assert_eq!(resp.error.unwrap().code, error_codes::TASK_NOT_FOUND);
    }

    #[test]
    fn test_handle_task_cancel() {
        let server = test_server();

        // Send a task first
        let task_req = TaskRequest {
            id: "cancel-me".into(),
            skill_id: None,
            messages: vec![],
            metadata: HashMap::new(),
        };
        let send_req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(1),
        );
        server.handle_request(&send_req);

        // Cancel it
        let cancel_req = JsonRpcRequest::new(
            methods::TASKS_CANCEL,
            Some(serde_json::json!({"id": "cancel-me"})),
            serde_json::json!(2),
        );
        let resp = server.handle_request(&cancel_req);
        assert!(resp.error.is_none());
    }

    #[test]
    fn test_update_task_status() {
        let server = test_server();

        // Send a task first
        let task_req = TaskRequest {
            id: "update-me".into(),
            skill_id: None,
            messages: vec![],
            metadata: HashMap::new(),
        };
        let send_req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(1),
        );
        server.handle_request(&send_req);

        // Update status
        server
            .update_task_status(
                "update-me",
                TaskState::Working,
                Some("Processing...".into()),
                None,
            )
            .unwrap();

        let tasks = server.list_tasks();
        assert_eq!(tasks.len(), 1);
        assert_eq!(tasks[0].status.state, TaskState::Working);
    }

    #[test]
    fn test_cancel_already_completed_task() {
        let server = test_server();

        // Send a task and mark it completed
        let task_req = TaskRequest {
            id: "done-task".into(),
            skill_id: None,
            messages: vec![],
            metadata: HashMap::new(),
        };
        let send_req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(1),
        );
        server.handle_request(&send_req);
        server
            .update_task_status("done-task", TaskState::Completed, None, None)
            .unwrap();

        // Attempt to cancel — should fail
        let cancel_req = JsonRpcRequest::new(
            methods::TASKS_CANCEL,
            Some(serde_json::json!({"id": "done-task"})),
            serde_json::json!(2),
        );
        let resp = server.handle_request(&cancel_req);
        assert!(resp.error.is_some());
        assert_eq!(resp.error.unwrap().code, error_codes::INVALID_PARAMS);
    }

    #[test]
    fn test_task_eviction() {
        let server = test_server();

        // Fill up to MAX_TASKS with completed tasks
        for i in 0..MAX_TASKS {
            let task_req = TaskRequest {
                id: format!("task-{}", i),
                skill_id: None,
                messages: vec![],
                metadata: HashMap::new(),
            };
            let send_req = JsonRpcRequest::new(
                methods::TASKS_SEND,
                Some(serde_json::to_value(&task_req).unwrap()),
                serde_json::json!(i),
            );
            server.handle_request(&send_req);
            server
                .update_task_status(
                    &format!("task-{}", i),
                    TaskState::Completed,
                    None,
                    None,
                )
                .unwrap();
        }

        assert_eq!(server.list_tasks().len(), MAX_TASKS);

        // Adding one more should trigger eviction of terminal tasks
        let task_req = TaskRequest {
            id: "new-task".into(),
            skill_id: None,
            messages: vec![],
            metadata: HashMap::new(),
        };
        let send_req = JsonRpcRequest::new(
            methods::TASKS_SEND,
            Some(serde_json::to_value(&task_req).unwrap()),
            serde_json::json!(9999),
        );
        server.handle_request(&send_req);

        // All old completed tasks evicted, only the new one remains
        assert_eq!(server.list_tasks().len(), 1);
    }
}