jarvy 0.0.5

//! MCP Server Implementation
//!
//! Main server loop that handles JSON-RPC requests and routes them to appropriate handlers.

use crate::mcp::audit::AuditLog;
use crate::mcp::config::McpConfig;
use crate::mcp::error::{McpError, McpResult};
use crate::mcp::prompts;
use crate::mcp::resources;
use crate::mcp::safety::RateLimiter;
use crate::mcp::tools;
use crate::mcp::transport::{JsonRpcRequest, JsonRpcResponse, StdioTransport};
use crate::mcp::{PROTOCOL_VERSION, SERVER_NAME, SERVER_VERSION};
use serde::{Deserialize, Serialize};

/// MCP Server that handles JSON-RPC requests over stdio
pub struct McpServer {
    /// Server configuration
    config: McpConfig,
    /// Rate limiter for safety
    rate_limiter: RateLimiter,
    /// Audit logger
    audit_log: AuditLog,
    /// Client information (set during initialize)
    client_info: Option<ClientInfo>,
}

/// Information about the connected MCP client
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientInfo {
    /// Client name (e.g., "claude-desktop")
    pub name: String,
    /// Client version
    #[serde(default)]
    pub version: Option<String>,
}

/// MCP Initialize request parameters
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct InitializeParams {
    /// Protocol version the client supports
    #[allow(dead_code)] // Required by MCP protocol spec for future use
    protocol_version: String,
    /// Client capabilities
    #[serde(default)]
    #[allow(dead_code)] // Required by MCP protocol spec for future use
    capabilities: serde_json::Value,
    /// Client information
    #[serde(default)]
    client_info: Option<ClientInfo>,
}

/// MCP Initialize response
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct InitializeResult {
    /// Protocol version the server supports
    protocol_version: String,
    /// Server capabilities
    capabilities: ServerCapabilities,
    /// Server information
    server_info: ServerInfo,
}

/// Server capabilities
#[derive(Debug, Serialize)]
struct ServerCapabilities {
    /// Tool capabilities
    tools: ToolCapabilities,
    /// Resource capabilities
    resources: ResourceCapabilities,
    /// Prompt capabilities
    prompts: PromptCapabilities,
}

/// Tool capabilities
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct ToolCapabilities {
    /// Whether the server supports listing tools that changed
    list_changed: bool,
}

/// Resource capabilities
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct ResourceCapabilities {
    /// Whether resources can be subscribed to
    subscribe: bool,
    /// Whether the server supports listing resources that changed
    list_changed: bool,
}

/// Prompt capabilities
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct PromptCapabilities {
    /// Whether the server supports listing prompts that changed
    list_changed: bool,
}

/// Server information
#[derive(Debug, Serialize)]
struct ServerInfo {
    /// Server name
    name: String,
    /// Server version
    version: String,
}

impl McpServer {
    /// Create a new MCP server with the given configuration
    pub fn new(config: McpConfig) -> Self {
        let rate_limiter = RateLimiter::new(&config);
        let audit_log = AuditLog::new(&config).unwrap_or_else(|_| AuditLog::disabled());

        Self {
            config,
            rate_limiter,
            audit_log,
            client_info: None,
        }
    }

    /// Run the MCP server (blocking)
    pub fn run(mut self) -> McpResult<()> {
        let mut transport = StdioTransport::new();

        loop {
            match transport.read_request() {
                Ok(Some(request)) => {
                    let response = self.handle_request(&request);
                    transport.write_response(&response)?;
                }
                Ok(None) => {
                    // EOF - client disconnected
                    break;
                }
                Err(e) => {
                    // Parse error - send error response
                    let response = JsonRpcResponse::error(None, e);
                    transport.write_response(&response)?;
                }
            }
        }

        Ok(())
    }

    /// Handle a single JSON-RPC request
    fn handle_request(&mut self, request: &JsonRpcRequest) -> JsonRpcResponse {
        let result = match request.method.as_str() {
            // Lifecycle methods
            "initialize" => self.handle_initialize(request),
            "notifications/initialized" => {
                // Notification - no response needed, but we return empty for consistency
                return JsonRpcResponse::success(request.id.clone(), serde_json::json!({}));
            }
            "ping" => Ok(serde_json::json!({})),

            // Tool methods
            "tools/list" => self.handle_tools_list(),
            "tools/call" => self.handle_tools_call(request),

            // Resource methods
            "resources/list" => self.handle_resources_list(),
            "resources/read" => self.handle_resources_read(request),

            // Prompt methods
            "prompts/list" => self.handle_prompts_list(),
            "prompts/get" => self.handle_prompts_get(request),

            _ => Err(McpError::method_not_found(&request.method)),
        };

        match result {
            Ok(value) => JsonRpcResponse::success(request.id.clone(), value),
            Err(e) => JsonRpcResponse::error(request.id.clone(), e),
        }
    }

    /// Handle the initialize request
    fn handle_initialize(&mut self, request: &JsonRpcRequest) -> McpResult<serde_json::Value> {
        let params: InitializeParams = request
            .params
            .as_ref()
            .map(|p| serde_json::from_value(p.clone()))
            .transpose()?
            .unwrap_or(InitializeParams {
                protocol_version: PROTOCOL_VERSION.to_string(),
                capabilities: serde_json::json!({}),
                client_info: None,
            });

        // Store client info for audit logging
        self.client_info = params.client_info;

        let result = InitializeResult {
            protocol_version: PROTOCOL_VERSION.to_string(),
            capabilities: ServerCapabilities {
                tools: ToolCapabilities {
                    list_changed: false,
                },
                resources: ResourceCapabilities {
                    subscribe: false,
                    list_changed: false,
                },
                prompts: PromptCapabilities {
                    list_changed: false,
                },
            },
            server_info: ServerInfo {
                name: SERVER_NAME.to_string(),
                version: SERVER_VERSION.to_string(),
            },
        };

        Ok(serde_json::to_value(result)?)
    }

    /// Handle tools/list request
    fn handle_tools_list(&self) -> McpResult<serde_json::Value> {
        let tools_list = tools::list_tools();
        Ok(serde_json::json!({ "tools": tools_list }))
    }

    /// Handle tools/call request
    fn handle_tools_call(&self, request: &JsonRpcRequest) -> McpResult<serde_json::Value> {
        #[derive(Deserialize)]
        struct ToolsCallParams {
            name: String,
            #[serde(default)]
            arguments: Option<serde_json::Value>,
        }

        let params: ToolsCallParams = request
            .params
            .as_ref()
            .map(|p| serde_json::from_value(p.clone()))
            .transpose()?
            .ok_or_else(|| McpError::invalid_params("Missing params for tools/call"))?;

        let client_name = self.client_info.as_ref().map(|c| c.name.as_str());

        match params.name.as_str() {
            "jarvy_get_install_instructions" => {
                let result = tools::handle_get_install_instructions(params.arguments)?;
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_check_self" => {
                let result = tools::handle_check_self()?;
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_list_tools" => {
                let result = tools::handle_list_tools(params.arguments)?;
                self.audit_log
                    .log_list_tools(client_name, result["count"].as_u64().unwrap_or(0) as usize);
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_get_tool" => {
                let result = tools::handle_get_tool(params.arguments)?;
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_check_tool" => {
                self.rate_limiter.check_check_limit().inspect_err(|_e| {
                    self.audit_log.log_rate_limited(client_name, "check_tool");
                })?;

                let result = tools::handle_check_tool(params.arguments)?;
                self.audit_log.log_check_tool(
                    client_name,
                    result["name"].as_str().unwrap_or("unknown"),
                    result["installed"].as_bool().unwrap_or(false),
                    result["version"].as_str(),
                );
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_check_multiple" => {
                self.rate_limiter.check_check_limit().inspect_err(|_e| {
                    self.audit_log
                        .log_rate_limited(client_name, "check_multiple");
                })?;

                let result = tools::handle_check_multiple(params.arguments)?;
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            "jarvy_install_tool" => {
                let result = tools::handle_install_tool(
                    params.arguments,
                    &self.config,
                    &self.rate_limiter,
                    &self.audit_log,
                    client_name,
                )?;
                Ok(serde_json::json!({
                    "content": [{
                        "type": "text",
                        "text": serde_json::to_string_pretty(&result)?
                    }]
                }))
            }
            _ => Err(McpError::method_not_found(format!(
                "Unknown tool: {}",
                params.name
            ))),
        }
    }

    /// Handle resources/list request
    fn handle_resources_list(&self) -> McpResult<serde_json::Value> {
        let resources_list = resources::list_resources();
        Ok(serde_json::json!({ "resources": resources_list }))
    }

    /// Handle resources/read request
    fn handle_resources_read(&self, request: &JsonRpcRequest) -> McpResult<serde_json::Value> {
        #[derive(Deserialize)]
        struct ResourcesReadParams {
            uri: String,
        }

        let params: ResourcesReadParams = request
            .params
            .as_ref()
            .map(|p| serde_json::from_value(p.clone()))
            .transpose()?
            .ok_or_else(|| McpError::invalid_params("Missing params for resources/read"))?;

        let content = resources::read_resource(&params.uri)?;
        Ok(serde_json::json!({
            "contents": [{
                "uri": params.uri,
                "mimeType": "application/json",
                "text": content
            }]
        }))
    }

    /// Handle prompts/list request
    fn handle_prompts_list(&self) -> McpResult<serde_json::Value> {
        let prompts_list = prompts::list_prompts();
        Ok(serde_json::json!({ "prompts": prompts_list }))
    }

    /// Handle prompts/get request
    fn handle_prompts_get(&self, request: &JsonRpcRequest) -> McpResult<serde_json::Value> {
        #[derive(Deserialize)]
        struct PromptsGetParams {
            name: String,
            #[serde(default)]
            arguments: Option<serde_json::Value>,
        }

        let params: PromptsGetParams = request
            .params
            .as_ref()
            .map(|p| serde_json::from_value(p.clone()))
            .transpose()?
            .ok_or_else(|| McpError::invalid_params("Missing params for prompts/get"))?;

        let prompt = prompts::get_prompt(&params.name, params.arguments)?;
        Ok(prompt)
    }
}

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

    #[test]
    fn test_server_creation() {
        let config = McpConfig::default();
        let server = McpServer::new(config);
        assert!(server.client_info.is_none());
    }

    #[test]
    fn test_tools_list_response() {
        let config = McpConfig::default();
        let server = McpServer::new(config);
        let result = server.handle_tools_list().unwrap();
        assert!(result.get("tools").is_some());
    }

    #[test]
    fn test_resources_list_response() {
        let config = McpConfig::default();
        let server = McpServer::new(config);
        let result = server.handle_resources_list().unwrap();
        assert!(result.get("resources").is_some());
    }

    #[test]
    fn test_prompts_list_response() {
        let config = McpConfig::default();
        let server = McpServer::new(config);
        let result = server.handle_prompts_list().unwrap();
        assert!(result.get("prompts").is_some());
    }
}