Vapora/crates/vapora-mcp-server/src/handlers.rs

use std::sync::Arc;

use serde_json::json;

use crate::{
    backend_client::BackendClient,
    error::{McpError, Result},
    protocol::{JsonRpcError, JsonRpcRequest, JsonRpcResponse, ToolCallParams},
    registry::ToolRegistry,
};

pub struct RequestHandler {
    backend: Arc<BackendClient>,
    registry: Arc<ToolRegistry>,
}

impl RequestHandler {
    pub fn new(backend: Arc<BackendClient>, registry: Arc<ToolRegistry>) -> Self {
        Self { backend, registry }
    }

    pub async fn handle(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
        if request.jsonrpc != "2.0" {
            return JsonRpcResponse {
                jsonrpc: "2.0".to_string(),
                id: request.id.clone(),
                result: None,
                error: Some(JsonRpcError {
                    code: -32600,
                    message: "Invalid Request: jsonrpc must be 2.0".to_string(),
                    data: None,
                }),
            };
        }

        let result = match request.method.as_str() {
            "initialize" => self.handle_initialize(request).await,
            "tools/list" => self.handle_tools_list(request).await,
            "tools/call" => self.handle_tools_call(request).await,
            "resources/list" => self.handle_resources_list(request).await,
            "prompts/list" => self.handle_prompts_list(request).await,
            method => Err(McpError::MethodNotFound(method.to_string())),
        };

        match result {
            Ok(result_value) => JsonRpcResponse {
                jsonrpc: "2.0".to_string(),
                id: request.id.clone(),
                result: Some(result_value),
                error: None,
            },
            Err(e) => {
                let error_obj = e.to_json_rpc_error();
                JsonRpcResponse {
                    jsonrpc: "2.0".to_string(),
                    id: request.id.clone(),
                    result: None,
                    error: Some(JsonRpcError {
                        code: error_obj["error"]["code"].as_i64().unwrap_or(-32603) as i32,
                        message: error_obj["error"]["message"]
                            .as_str()
                            .unwrap_or("Unknown error")
                            .to_string(),
                        data: None,
                    }),
                }
            }
        }
    }

    async fn handle_initialize(&self, _request: &JsonRpcRequest) -> Result<serde_json::Value> {
        Ok(json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": { "list_changed": false },
                "resources": { "subscribe": false, "list_changed": false },
                "prompts": { "list_changed": false }
            },
            "serverInfo": {
                "name": "vapora-mcp-server",
                "version": env!("CARGO_PKG_VERSION"),
            }
        }))
    }

    async fn handle_tools_list(&self, _request: &JsonRpcRequest) -> Result<serde_json::Value> {
        let tools: Vec<_> = self
            .registry
            .list()
            .into_iter()
            .map(|t| {
                json!({
                    "name": t.name,
                    "description": t.description,
                    "inputSchema": {
                        "type": t.input_schema.schema_type,
                        "properties": t.input_schema.properties,
                        "required": t.input_schema.required,
                    }
                })
            })
            .collect();

        Ok(json!({ "tools": tools }))
    }

    async fn handle_tools_call(&self, request: &JsonRpcRequest) -> Result<serde_json::Value> {
        let params: ToolCallParams =
            serde_json::from_value(request.params.clone()).map_err(|e| {
                McpError::InvalidParams {
                    method: "tools/call".to_string(),
                    reason: e.to_string(),
                }
            })?;

        let result = match params.name.as_str() {
            "kanban_create_task" => {
                let project_id = params
                    .arguments
                    .get("project_id")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "kanban_create_task".to_string(),
                        reason: "Missing project_id".to_string(),
                    })?;

                let title = params
                    .arguments
                    .get("title")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "kanban_create_task".to_string(),
                        reason: "Missing title".to_string(),
                    })?;

                self.backend
                    .create_task(project_id, title.to_string())
                    .await?
            }
            "kanban_update_task" => {
                let task_id = params
                    .arguments
                    .get("task_id")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "kanban_update_task".to_string(),
                        reason: "Missing task_id".to_string(),
                    })?;

                let status = params
                    .arguments
                    .get("status")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "kanban_update_task".to_string(),
                        reason: "Missing status".to_string(),
                    })?;

                self.backend.update_task(task_id, status).await?
            }
            "get_project_summary" => {
                let project_id = params
                    .arguments
                    .get("project_id")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "get_project_summary".to_string(),
                        reason: "Missing project_id".to_string(),
                    })?;

                self.backend.get_project_summary(project_id).await?
            }
            "list_agents" => self.backend.list_agents().await?,
            "get_agent_capabilities" => {
                let agent_id = params
                    .arguments
                    .get("agent_id")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "get_agent_capabilities".to_string(),
                        reason: "Missing agent_id".to_string(),
                    })?;

                self.backend.get_agent_capabilities(agent_id).await?
            }
            "assign_task_to_agent" => {
                let role = params
                    .arguments
                    .get("role")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "assign_task_to_agent".to_string(),
                        reason: "Missing role".to_string(),
                    })?;

                let title = params
                    .arguments
                    .get("task_title")
                    .and_then(|v| v.as_str())
                    .ok_or(McpError::InvalidParams {
                        method: "assign_task_to_agent".to_string(),
                        reason: "Missing task_title".to_string(),
                    })?;

                let description = params
                    .arguments
                    .get("task_description")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");

                self.backend
                    .assign_task_to_agent(role, title.to_string(), description.to_string())
                    .await?
            }
            tool_name => {
                return Err(McpError::MethodNotFound(format!("Tool: {}", tool_name)));
            }
        };

        Ok(json!({
            "content": [{
                "type": "text",
                "text": serde_json::to_string(&result)?
            }],
            "isError": false
        }))
    }

    async fn handle_resources_list(&self, _request: &JsonRpcRequest) -> Result<serde_json::Value> {
        Ok(json!({
            "resources": [
                { "uri": "vapora://projects", "name": "Projects", "description": "All projects" },
                { "uri": "vapora://tasks", "name": "Tasks", "description": "All tasks" },
                { "uri": "vapora://agents", "name": "Agents", "description": "Agent registry" },
            ]
        }))
    }

    async fn handle_prompts_list(&self, _request: &JsonRpcRequest) -> Result<serde_json::Value> {
        Ok(json!({
            "prompts": [
                { "name": "analyze_task", "description": "Analyze task and suggest improvements" },
                { "name": "code_review_prompt", "description": "Generate code review feedback" },
            ]
        }))
    }
}