mcp_host/server/

core.rs

//! Server core implementation
//!
//! Main MCP server with method handlers and session management.
//!
//! # Bidirectional Communication
//!
//! The server supports bidirectional communication:
//! - Client→Server: tools/call, resources/read, prompts/get, etc.
//! - Server→Client: roots/list, sampling/createMessage, etc.
//!
//! Server-initiated requests are handled via the [`RequestMultiplexer`].

use std::sync::Arc;
use std::time::Duration;

use arc_swap::ArcSwap;
use dashmap::DashMap;
use serde_json::Value;
use tokio::sync::mpsc;

use crate::content::types::Content;
use crate::managers::cancellation::{CancellationManager, CancelledNotificationParams};
use crate::managers::completion::CompletionManager;
use crate::managers::subscription::{SubscribeRequest, SubscriptionManager, UnsubscribeRequest};
use crate::protocol::capabilities::{InitializeRequest, InitializeResult, ServerCapabilities};
use crate::protocol::errors::{ErrorType, McpError, codes};
use crate::protocol::methods::McpMethod;
use crate::protocol::types::{Implementation, JsonRpcRequest, JsonRpcResponse};
use crate::protocol::version;
use crate::registry::prompts::PromptManager;
use crate::registry::resources::ResourceManager;
use crate::registry::tools::{ToolError, ToolRegistry};
use crate::server::handler::{
    RequestContext, error_response, require_initialization, success_response,
};
use crate::server::middleware::MiddlewareChain;
use crate::server::multiplexer::{
    ClientRequester, CreateMessageParams, CreateMessageResult, JsonRpcClientRequest,
    ListRootsResult, MultiplexerError, RequestMultiplexer, Root,
};
use crate::server::session::Session;
use crate::server::visibility::{Environment, VisibilityContext};
use crate::transport::traits::{IncomingMessage, JsonRpcNotification, Transport};

fn tool_error_to_mcp_error(tool_name: &str, err: ToolError) -> McpError {
    match err {
        ToolError::NotFound(_) => McpError::not_found(
            codes::TOOL_NOT_FOUND,
            format!("Tool not found: {}", tool_name),
        ),
        ToolError::InvalidArguments(message) => {
            McpError::validation(codes::INVALID_TOOL_ARGS, message)
        }
        ToolError::Execution(message) | ToolError::Internal(message) => {
            McpError::internal(codes::TOOL_EXECUTION_FAILED, message)
        }
        ToolError::CircuitOpen { tool, message } => McpError::builder(
            ErrorType::CircuitOpen,
            codes::CIRCUIT_BREAKER_OPEN,
        )
        .message(message)
        .detail("tool", tool)
        .build(),
    }
}

/// Main MCP server
pub struct Server {
    /// Server name
    name: String,

    /// Server version
    version: String,

    /// Server instructions for LLMs (atomically swappable)
    instructions: Arc<ArcSwap<Option<String>>>,

    /// Server capabilities (atomically swappable for dynamic updates)
    capabilities: Arc<ArcSwap<ServerCapabilities>>,

    /// Active sessions by session ID
    sessions: DashMap<String, Session>,

    /// Middleware chain for request processing
    middleware: MiddlewareChain,

    /// Tool registry
    tool_registry: ToolRegistry,

    /// Resource manager
    resource_manager: ResourceManager,

    /// Prompt manager
    prompt_manager: PromptManager,

    /// Notification sender (for background tasks to send notifications)
    notification_tx: mpsc::UnboundedSender<JsonRpcNotification>,

    /// Notification receiver (internal, used by run loop)
    notification_rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<JsonRpcNotification>>>,

    /// Global logger for server notifications
    logger: crate::logging::McpLogger,

    /// Request multiplexer for server→client requests
    multiplexer: Arc<RequestMultiplexer>,

    /// Channel to send server→client requests to the transport
    request_tx: mpsc::UnboundedSender<JsonRpcClientRequest>,

    /// Receiver for server→client requests (used by run loop)
    request_rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<JsonRpcClientRequest>>>,

    /// Task store for async task execution
    task_store: Arc<crate::managers::task::TaskStore>,

    /// Cancellation manager for tracking cancellable operations
    cancellation_manager: CancellationManager,

    /// Subscription manager for resource subscriptions
    subscription_manager: SubscriptionManager,

    /// Completion manager for argument value suggestions
    completion_manager: CompletionManager,

    /// Optional environment for visibility checks
    environment: Option<Arc<dyn Environment>>,
}

impl Server {
    /// Create new server
    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
        let (notification_tx, notification_rx) = mpsc::unbounded_channel();
        let (request_tx, request_rx) = mpsc::unbounded_channel();

        // Create task store with 5-minute TTL and 5-second poll interval
        let task_store = Arc::new(crate::managers::task::TaskStore::new(
            std::time::Duration::from_secs(300),
            std::time::Duration::from_secs(5),
        ));

        // Spawn cleanup task if in tokio runtime context (runs every minute)
        if tokio::runtime::Handle::try_current().is_ok() {
            task_store
                .clone()
                .spawn_cleanup_task(std::time::Duration::from_secs(60));
        }

        // Default capabilities with tasks support
        let default_caps = ServerCapabilities {
            tasks: Some(crate::protocol::capabilities::TasksCapability {
                list: Some(crate::protocol::capabilities::EmptyObject {}),
                cancel: Some(crate::protocol::capabilities::EmptyObject {}),
                requests: Some(crate::protocol::capabilities::TasksRequestsCapability {
                    tools: Some(crate::protocol::capabilities::TasksToolsCapability {
                        call: Some(crate::protocol::capabilities::EmptyObject {}),
                    }),
                    ..Default::default()
                }),
            }),
            ..Default::default()
        };

        // Create logger with default config
        let logger = crate::logging::McpLogger::new(notification_tx.clone(), "mcp-server");

        Self {
            name: name.into(),
            version: version.into(),
            instructions: Arc::new(ArcSwap::new(Arc::new(None))),
            capabilities: Arc::new(ArcSwap::new(Arc::new(default_caps))),
            sessions: DashMap::new(),
            middleware: MiddlewareChain::new(),
            tool_registry: ToolRegistry::new(),
            resource_manager: ResourceManager::new(),
            prompt_manager: PromptManager::new(),
            notification_tx: notification_tx.clone(),
            notification_rx: Arc::new(tokio::sync::Mutex::new(notification_rx)),
            logger,
            multiplexer: Arc::new(RequestMultiplexer::new()),
            request_tx,
            request_rx: Arc::new(tokio::sync::Mutex::new(request_rx)),
            task_store,
            cancellation_manager: CancellationManager::new(),
            subscription_manager: SubscriptionManager::with_notifications(notification_tx),
            completion_manager: CompletionManager::new(),
            environment: None,
        }
    }

    /// Get server name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get server version
    pub fn version(&self) -> &str {
        &self.version
    }

    /// Get current capabilities
    pub fn capabilities(&self) -> Arc<ServerCapabilities> {
        self.capabilities.load_full()
    }

    /// Update server capabilities
    pub fn set_capabilities(&self, capabilities: ServerCapabilities) {
        self.capabilities.store(Arc::new(capabilities));
    }

    /// Get current instructions
    pub fn instructions(&self) -> Option<String> {
        (**self.instructions.load()).clone()
    }

    /// Set server instructions for LLMs
    pub fn set_instructions(&self, instructions: Option<String>) {
        self.instructions.store(Arc::new(instructions));
    }

    /// Set environment for visibility checks
    pub fn set_environment(&mut self, environment: Arc<dyn Environment>) {
        self.environment = Some(environment);
    }

    fn visibility_context<'a>(&'a self, session: &'a Session) -> VisibilityContext<'a> {
        match self.environment.as_deref() {
            Some(env) => VisibilityContext::with_environment(session, env),
            None => VisibilityContext::new(session),
        }
    }

    /// Add middleware to the chain
    pub fn add_middleware(&mut self, middleware: crate::server::middleware::MiddlewareFn) {
        self.middleware.add(middleware);
    }

    /// Get tool registry
    pub fn tool_registry(&self) -> &ToolRegistry {
        &self.tool_registry
    }

    /// Get resource manager
    pub fn resource_manager(&self) -> &ResourceManager {
        &self.resource_manager
    }

    /// Get prompt manager
    pub fn prompt_manager(&self) -> &PromptManager {
        &self.prompt_manager
    }

    /// Get subscription manager for resource subscriptions
    pub fn subscription_manager(&self) -> &SubscriptionManager {
        &self.subscription_manager
    }

    /// Get completion manager for argument value suggestions
    pub fn completion_manager(&self) -> &CompletionManager {
        &self.completion_manager
    }

    /// Get global logger
    pub fn logger(&self) -> &crate::logging::McpLogger {
        &self.logger
    }

    /// Get notification sender (for background tasks)
    pub fn notification_sender(&self) -> mpsc::UnboundedSender<JsonRpcNotification> {
        self.notification_tx.clone()
    }

    /// Send a notification to the client
    pub fn send_notification(
        &self,
        method: impl Into<String>,
        params: Option<Value>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let notification = JsonRpcNotification::new(method, params);
        self.notification_tx.send(notification)?;
        Ok(())
    }

    /// Get session by ID
    pub fn get_session(&self, session_id: &str) -> Option<Session> {
        self.sessions.get(session_id).map(|s| s.clone())
    }

    /// Remove session
    ///
    /// Also cleans up any resource subscriptions for the session.
    pub fn remove_session(&self, session_id: &str) -> Option<Session> {
        // Clean up subscriptions before removing session
        self.subscription_manager.remove_session(session_id);
        self.sessions.remove(session_id).map(|(_, s)| s)
    }

    /// Notify subscribers that a resource has been updated
    ///
    /// Sends `notifications/resources/updated` to all sessions subscribed to the URI.
    /// Call this when a resource's content changes.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // After updating a file resource
    /// server.notify_resource_updated("file:///path/to/file.txt");
    /// ```
    pub fn notify_resource_updated(&self, uri: &str) {
        self.subscription_manager.notify_resource_updated(uri);
    }

    /// Notify subscribers for multiple resources (batch operation)
    ///
    /// Convenience method for notifying about multiple resource updates at once.
    pub fn notify_resources_updated(&self, uris: &[&str]) {
        self.subscription_manager.notify_resources_updated(uris);
    }

    /// Get the request multiplexer (for advanced use cases)
    pub fn multiplexer(&self) -> Arc<RequestMultiplexer> {
        self.multiplexer.clone()
    }

    /// Get the cancellation manager
    ///
    /// Use this to register cancellable operations and check cancellation status.
    pub fn cancellation_manager(&self) -> &CancellationManager {
        &self.cancellation_manager
    }

    /// Create a client requester for the given session
    ///
    /// The client requester allows tools to make server→client requests
    /// like roots/list, sampling/createMessage, and elicitation/create.
    pub fn create_client_requester(&self, session_id: &str) -> Option<ClientRequester> {
        let session = self.get_session(session_id)?;
        let caps = session.capabilities.as_ref()?;

        Some(ClientRequester::new(
            self.request_tx.clone(),
            self.multiplexer.clone(),
            caps.roots.is_some(),
            caps.sampling.is_some(),
            caps.elicitation.is_some(),
        ))
    }

    /// Request workspace roots from the client
    ///
    /// Sends a `roots/list` request to the client and waits for the response.
    /// The client must have the `roots` capability advertised.
    ///
    /// # Arguments
    ///
    /// * `session_id` - The session to check for roots capability
    /// * `timeout` - Optional timeout (defaults to 30 seconds)
    ///
    /// # Returns
    ///
    /// List of workspace roots, or an error if:
    /// - Client doesn't support roots capability
    /// - Request times out
    /// - Client returns an error
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let roots = server.request_roots("session-123", None).await?;
    /// for root in roots {
    ///     println!("Root: {} ({})", root.name.unwrap_or_default(), root.uri);
    /// }
    /// ```
    pub async fn request_roots(
        &self,
        session_id: &str,
        timeout: Option<Duration>,
    ) -> Result<Vec<Root>, MultiplexerError> {
        // Check if client supports roots
        if let Some(session) = self.get_session(session_id) {
            if let Some(caps) = &session.capabilities {
                if caps.roots.is_none() {
                    return Err(MultiplexerError::UnsupportedCapability("roots".to_string()));
                }
            } else {
                return Err(MultiplexerError::UnsupportedCapability("roots".to_string()));
            }
        } else {
            return Err(MultiplexerError::Transport("session not found".to_string()));
        }

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("roots/list");

        // Build and send the request
        let request = JsonRpcClientRequest::new(&id, "roots/list", Some(serde_json::json!({})));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let list_result: ListRootsResult = serde_json::from_value(result)?;
        Ok(list_result.roots)
    }

    /// Request an LLM completion from the client
    ///
    /// Sends a `sampling/createMessage` request to the client.
    /// The client must have the `sampling` capability advertised.
    ///
    /// # Arguments
    ///
    /// * `session_id` - The session to check for sampling capability
    /// * `params` - The sampling parameters
    /// * `timeout` - Optional timeout (defaults to 30 seconds)
    ///
    /// # Returns
    ///
    /// The completion result, or an error if:
    /// - Client doesn't support sampling capability
    /// - Request times out
    /// - Client returns an error
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use mcp_host::server::multiplexer::{CreateMessageParams, SamplingContent, SamplingMessage};
    ///
    /// let params = CreateMessageParams {
    ///     messages: vec![SamplingMessage {
    ///         role: "user".to_string(),
    ///         content: SamplingContent::Text { text: "Hello!".to_string() },
    ///     }],
    ///     max_tokens: 1000,
    ///     ..Default::default()
    /// };
    ///
    /// let result = server.request_sampling("session-123", params, None).await?;
    /// println!("Response: {:?}", result.content);
    /// ```
    pub async fn request_sampling(
        &self,
        session_id: &str,
        params: CreateMessageParams,
        timeout: Option<Duration>,
    ) -> Result<CreateMessageResult, MultiplexerError> {
        // Check if client supports sampling
        if let Some(session) = self.get_session(session_id) {
            if let Some(caps) = &session.capabilities {
                if caps.sampling.is_none() {
                    return Err(MultiplexerError::UnsupportedCapability(
                        "sampling".to_string(),
                    ));
                }
            } else {
                return Err(MultiplexerError::UnsupportedCapability(
                    "sampling".to_string(),
                ));
            }
        } else {
            return Err(MultiplexerError::Transport("session not found".to_string()));
        }

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("sampling/createMessage");

        // Build and send the request
        let params_value = serde_json::to_value(&params)?;
        let request = JsonRpcClientRequest::new(&id, "sampling/createMessage", Some(params_value));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let create_result: CreateMessageResult = serde_json::from_value(result)?;
        Ok(create_result)
    }

    /// Request structured user input from the client
    ///
    /// Sends an `elicitation/create` request to the client.
    /// The client must have the `elicitation` capability advertised.
    ///
    /// # Arguments
    ///
    /// * `session_id` - The session to check for elicitation capability
    /// * `message` - The message to show the user
    /// * `requested_schema` - JSON Schema for the structured input
    /// * `timeout` - Optional timeout (defaults to 30 seconds)
    ///
    /// # Returns
    ///
    /// The elicitation result (action + optional content), or an error if:
    /// - Client doesn't support elicitation capability
    /// - Request times out
    /// - Client returns an error
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use mcp_host::protocol::elicitation::ElicitationSchema;
    /// use mcp_host::protocol::types::ElicitationAction;
    ///
    /// let schema = ElicitationSchema::builder()
    ///     .required_email("email")
    ///     .required_integer("age", 0, 150)
    ///     .optional_bool("newsletter", false)
    ///     .build_unchecked();
    ///
    /// let result = server.request_elicitation(
    ///     "session-123",
    ///     "Please provide your information",
    ///     serde_json::to_value(&schema).unwrap(),
    ///     None,
    /// ).await?;
    ///
    /// match result.action {
    ///     ElicitationAction::Accept => {
    ///         if let Some(content) = result.content {
    ///             let email = content["email"].as_str().unwrap();
    ///             let age = content["age"].as_i64().unwrap();
    ///             println!("User: {} (age {})", email, age);
    ///         }
    ///     }
    ///     ElicitationAction::Decline => println!("User declined"),
    ///     ElicitationAction::Cancel => println!("User cancelled"),
    /// }
    /// ```
    pub async fn request_elicitation(
        &self,
        session_id: &str,
        message: impl Into<String>,
        requested_schema: Value,
        timeout: Option<Duration>,
    ) -> Result<crate::protocol::types::CreateElicitationResult, MultiplexerError> {
        // Check if client supports elicitation
        if let Some(session) = self.get_session(session_id) {
            if let Some(caps) = &session.capabilities {
                if caps.elicitation.is_none() {
                    return Err(MultiplexerError::UnsupportedCapability(
                        "elicitation".to_string(),
                    ));
                }
            } else {
                return Err(MultiplexerError::UnsupportedCapability(
                    "elicitation".to_string(),
                ));
            }
        } else {
            return Err(MultiplexerError::Transport("session not found".to_string()));
        }

        // Create pending request
        let (id, rx) = self.multiplexer.create_pending("elicitation/create");

        // Build and send the request
        let params = serde_json::json!({
            "message": message.into(),
            "requestedSchema": requested_schema,
        });
        let request = JsonRpcClientRequest::new(&id, "elicitation/create", Some(params));

        self.request_tx
            .send(request)
            .map_err(|e| MultiplexerError::Transport(e.to_string()))?;

        // Wait for response with timeout
        let timeout = timeout.unwrap_or(self.multiplexer.default_timeout());
        let result = tokio::time::timeout(timeout, rx)
            .await
            .map_err(|_| MultiplexerError::Timeout(timeout))?
            .map_err(|_| MultiplexerError::ChannelClosed)??;

        // Parse the result
        let elicitation_result: crate::protocol::types::CreateElicitationResult =
            serde_json::from_value(result)?;
        Ok(elicitation_result)
    }

    /// Run server with the given transport
    ///
    /// This is the main event loop that reads requests from the transport,
    /// processes them, and writes responses back.
    ///
    /// Supports bidirectional communication:
    /// - Incoming client requests are handled and responses sent
    /// - Server-initiated requests are sent via channels and responses routed back
    pub async fn run<T: Transport>(
        &self,
        mut transport: T,
    ) -> Result<(), Box<dyn std::error::Error>> {
        // Generate a unique session ID per MCP spec (globally unique, cryptographically secure)
        let session_id = uuid::Uuid::new_v4().to_string();

        // Get receivers
        let mut notification_rx = self.notification_rx.lock().await;
        let mut request_rx = self.request_rx.lock().await;

        loop {
            tokio::select! {
                // Handle outgoing notifications
                Some(notification) = notification_rx.recv() => {
                    tracing::debug!(method = %notification.method, "Sending notification");
                    if let Err(e) = transport.send_notification(notification).await {
                        tracing::error!(error = %e, "Failed to send notification");
                    }
                }

                // Handle outgoing server→client requests
                Some(request) = request_rx.recv() => {
                    tracing::debug!(method = %request.method, id = %request.id, "Sending request to client");
                    if let Err(e) = transport.send_request(request).await {
                        tracing::error!(error = %e, "Failed to send request to client");
                    }
                }

                // Handle incoming messages (requests or responses)
                result = transport.read_incoming() => {
                    match result {
                        Ok(IncomingMessage::Request(request)) => {
                            // Check if this is a notification (no ID) or a request (has ID)
                            let is_notification = request.id.is_none();

                            // Handle client request
                            let response = self.handle_request(&session_id, request).await;

                            // Only write response if this was NOT a notification
                            // Notifications must NOT receive responses per JSON-RPC spec
                            if !is_notification
                                && let Err(e) = transport.write_message(response).await {
                                    tracing::error!(error = %e, "Failed to write message");
                                    break;
                                }
                        }
                        Ok(IncomingMessage::Response(response)) => {
                            // Route response to pending server-initiated request
                            if !self.multiplexer.route_response(&response) {
                                tracing::warn!(
                                    id = ?response.id,
                                    "Received response for unknown request ID"
                                );
                            }
                        }
                        Err(crate::transport::traits::TransportError::Closed) => {
                            tracing::info!("Transport closed, shutting down");
                            break;
                        }
                        Err(e) => {
                            tracing::error!(error = %e, "Failed to read message");
                            continue;
                        }
                    }
                }
            }
        }

        // Cancel any pending requests
        self.multiplexer.cancel_all();

        // Signal cancellation to all tracked operations
        self.cancellation_manager.cancel_all();

        // Shutdown transport gracefully
        transport.shutdown().await?;

        Ok(())
    }

    /// Handle incoming JSON-RPC request
    pub async fn handle_request(
        &self,
        session_id: &str,
        request: JsonRpcRequest,
    ) -> JsonRpcResponse {
        // Get or create session
        let session = self
            .sessions
            .entry(session_id.to_string())
            .or_insert_with(|| {
                let mut session = Session::with_id(session_id);
                session.set_notification_channel(self.notification_tx.clone());
                session
            })
            .clone();

        // Create request context
        let ctx = RequestContext::new(session, request.clone());

        // Process through middleware chain
        let ctx = match self.middleware.process(ctx) {
            Ok(ctx) => ctx,
            Err(err) => return error_response(request.id, err.to_jsonrpc()),
        };

        // Route to appropriate handler
        let method = McpMethod::from(request.method.clone());

        match method {
            McpMethod::Initialize => self.handle_initialize(ctx).await,
            McpMethod::Ping => self.handle_ping(ctx).await,
            McpMethod::LoggingSetLevel => self.handle_logging_set_level(ctx).await,
            McpMethod::ToolsList => self.handle_tools_list(ctx).await,
            McpMethod::ToolsCall => self.handle_tools_call(ctx).await,
            McpMethod::ResourcesList => self.handle_resources_list(ctx).await,
            McpMethod::ResourcesTemplatesList => self.handle_resources_templates_list(ctx).await,
            McpMethod::ResourcesRead => self.handle_resources_read(ctx).await,
            McpMethod::ResourcesSubscribe => self.handle_resources_subscribe(ctx).await,
            McpMethod::ResourcesUnsubscribe => self.handle_resources_unsubscribe(ctx).await,
            McpMethod::PromptsList => self.handle_prompts_list(ctx).await,
            McpMethod::PromptsGet => self.handle_prompts_get(ctx).await,
            McpMethod::RootsList => self.handle_roots_list(ctx).await,
            McpMethod::SamplingCreateMessage => self.handle_sampling_create_message(ctx).await,
            McpMethod::ElicitationCreate => self.handle_elicitation_create(ctx).await,
            McpMethod::TasksGet => self.handle_tasks_get(ctx).await,
            McpMethod::TasksResult => self.handle_tasks_result(ctx).await,
            McpMethod::TasksList => self.handle_tasks_list(ctx).await,
            McpMethod::TasksCancel => self.handle_tasks_cancel(ctx).await,
            McpMethod::CompletionComplete => self.handle_completion_complete(ctx).await,
            McpMethod::NotificationsCancelled => self.handle_notifications_cancelled(ctx).await,
            _ => error_response(
                request.id,
                McpError::method_not_found(&request.method).to_jsonrpc(),
            ),
        }
    }

    /// Handle initialize request
    async fn handle_initialize(&self, ctx: RequestContext) -> JsonRpcResponse {
        // Parse initialize params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let req: InitializeRequest = match serde_json::from_value(params) {
            Ok(req) => req,
            Err(_) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Invalid initialize parameters")
                        .to_jsonrpc(),
                );
            }
        };

        // Negotiate protocol version
        let protocol_version = match version::negotiate_protocol_version(&req.protocol_version) {
            Ok(version) => version,
            Err(supported_versions) => {
                tracing::warn!(
                    client = %req.client_info.name,
                    requested = %req.protocol_version,
                    supported = ?supported_versions,
                    "Unsupported protocol version"
                );
                return error_response(
                    ctx.request.id,
                    McpError::builder(ErrorType::Validation, "unsupported_protocol_version")
                        .message("Unsupported protocol version")
                        .detail(
                            "supported",
                            serde_json::to_value(&supported_versions).unwrap(),
                        )
                        .detail("requested", req.protocol_version.clone())
                        .build()
                        .to_jsonrpc(),
                );
            }
        };

        // Log client connection
        tracing::info!(
            client = %req.client_info.name,
            version = %req.client_info.version,
            protocol = %protocol_version,
            "Client connected"
        );

        // Initialize session
        if let Some(mut session) = self.sessions.get_mut(&ctx.session.id) {
            session.initialize(req.client_info, req.capabilities, protocol_version.clone());
        }

        // Build response
        let result = InitializeResult {
            protocol_version,
            capabilities: (**self.capabilities.load()).clone(),
            server_info: Implementation {
                name: self.name.clone(),
                version: self.version.clone(),
            },
            instructions: self.instructions(),
        };

        success_response(
            ctx.request.id,
            serde_json::to_value(result).expect("Failed to serialize InitializeResult"),
        )
    }

    /// Handle ping request
    async fn handle_ping(&self, ctx: RequestContext) -> JsonRpcResponse {
        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle logging/setLevel request
    async fn handle_logging_set_level(&self, ctx: RequestContext) -> JsonRpcResponse {
        use crate::logging::LogLevel;
        use crate::protocol::types::SetLevelRequest;

        // Parse request params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let req: SetLevelRequest = match serde_json::from_value(params) {
            Ok(req) => req,
            Err(_) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Invalid setLevel parameters")
                        .to_jsonrpc(),
                );
            }
        };

        // Parse log level
        let level = match req.level.parse::<LogLevel>() {
            Ok(level) => level,
            Err(_) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation(
                        "invalid_level",
                        format!(
                            "Invalid log level '{}'. Valid levels: debug, info, notice, warning, error, critical, alert, emergency",
                            req.level
                        ),
                    )
                    .to_jsonrpc(),
                )
            }
        };

        // Update logger min level
        self.logger.set_min_level(level);

        tracing::debug!(level = %req.level, "Log level updated");

        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle tools/list request (with pagination support)
    async fn handle_tools_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse cursor if present
        let cursor = ctx
            .params()
            .and_then(|p| p.get("cursor"))
            .and_then(|c| c.as_str());

        let visibility_ctx = self.visibility_context(&ctx.session);
        let all_tools = self
            .tool_registry
            .list_for_session(&ctx.session, &visibility_ctx);

        // Apply pagination
        let result = crate::utils::paginate(&all_tools, cursor, crate::utils::DEFAULT_PAGE_SIZE);

        // Build response (only include nextCursor if present)
        let mut response = serde_json::json!({"tools": result.items});
        if let Some(next) = result.next_cursor {
            response["nextCursor"] = serde_json::Value::String(next);
        }
        success_response(ctx.request.id, response)
    }

    /// Handle tools/call request
    async fn handle_tools_call(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse tool call params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let tool_name = match params.get("name").and_then(|v| v.as_str()) {
            Some(name) => name,
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing 'name' field").to_jsonrpc(),
                );
            }
        };

        let tool_params = params.get("arguments").cloned().unwrap_or(Value::Null);

        // Check for task-augmented execution
        let task_meta: Option<crate::protocol::types::TaskMetadata> = params
            .get("task")
            .and_then(|t| serde_json::from_value(t.clone()).ok());

        if let Some(task_metadata) = task_meta {
            // Task-augmented execution
            return self
                .handle_task_augmented_tool_call(ctx, tool_name, tool_params, task_metadata)
                .await;
        }

        // Regular synchronous execution
        let client_requester = self.create_client_requester(&ctx.session.id);
        let visibility_ctx = self.visibility_context(&ctx.session);

        match self
            .tool_registry
            .call_for_session(
                tool_name,
                tool_params,
                &ctx.session,
                &self.logger,
                &visibility_ctx,
                client_requester,
            )
            .await
        {
            Ok(output) => {
                let result = match output {
                    crate::registry::tools::ToolOutput::Content(items) => {
                        let content_values: Vec<Value> =
                            items.iter().map(|c| c.to_value()).collect();
                        serde_json::json!({"content": content_values})
                    }
                    crate::registry::tools::ToolOutput::Structured(value) => {
                        serde_json::json!({"structuredContent": value})
                    }
                };
                success_response(ctx.request.id, result)
            }
            Err(e) => error_response(ctx.request.id, tool_error_to_mcp_error(tool_name, e).to_jsonrpc()),
        }
    }

    /// Handle task-augmented tool call
    async fn handle_task_augmented_tool_call(
        &self,
        ctx: RequestContext,
        tool_name: &str,
        tool_params: Value,
        task_metadata: crate::protocol::types::TaskMetadata,
    ) -> JsonRpcResponse {
        // Create task
        let (task, _result_rx) =
            self.task_store
                .create_task(&ctx.session.id, ctx.request.clone(), task_metadata.ttl);

        let task_id = task.task_id.clone();

        // Spawn background execution
        let task_store = self.task_store.clone();
        let tool_registry = self.tool_registry.clone();
        let logger = self.logger.clone();
        let session = ctx.session.clone();
        let environment = self.environment.clone();
        let client_requester = self.create_client_requester(&ctx.session.id);
        let tool_name = tool_name.to_string();

        tokio::spawn(async move {
            // Execute tool
            let visibility_ctx = match environment.as_deref() {
                Some(env) => VisibilityContext::with_environment(&session, env),
                None => VisibilityContext::new(&session),
            };

            match tool_registry
                .call_for_session(
                    &tool_name,
                    tool_params,
                    &session,
                    &logger,
                    &visibility_ctx,
                    client_requester,
                )
                .await
            {
                Ok(output) => {
                    // Success - store result
                    let result = match output {
                        crate::registry::tools::ToolOutput::Content(items) => {
                            let content_values: Vec<Value> =
                                items.iter().map(|c| c.to_value()).collect();
                            serde_json::json!({"content": content_values})
                        }
                        crate::registry::tools::ToolOutput::Structured(value) => {
                            serde_json::json!({"structuredContent": value})
                        }
                    };

                    let _ = task_store
                        .update_status(
                            &task_id,
                            crate::protocol::types::TaskStatus::Completed,
                            None,
                        )
                        .await;
                    let _ = task_store.store_result(&task_id, result).await;
                }
                Err(e) => {
                    // Failure - store error
                    let error_message = e.to_string();
                    let _ = task_store
                        .update_status(
                            &task_id,
                            crate::protocol::types::TaskStatus::Failed,
                            Some(error_message.clone()),
                        )
                        .await;

                    // Store error as result (for tasks/result)
                    let error_result = serde_json::json!({
                        "content": [{
                            "type": "text",
                            "text": error_message
                        }],
                        "isError": true
                    });
                    let _ = task_store.store_result(&task_id, error_result).await;
                }
            }
        });

        // Return immediately with task info
        success_response(
            ctx.request.id,
            serde_json::to_value(crate::protocol::types::CreateTaskResult { task }).unwrap(),
        )
    }

    /// Handle resources/list request (with pagination support)
    async fn handle_resources_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse cursor if present
        let cursor = ctx
            .params()
            .and_then(|p| p.get("cursor"))
            .and_then(|c| c.as_str());

        let visibility_ctx = self.visibility_context(&ctx.session);
        let all_resources = self
            .resource_manager
            .list_for_session(&ctx.session, &visibility_ctx);

        // Apply pagination
        let result = crate::utils::paginate(&all_resources, cursor, crate::utils::DEFAULT_PAGE_SIZE);

        // Build response (only include nextCursor if present)
        let mut response = serde_json::json!({"resources": result.items});
        if let Some(next) = result.next_cursor {
            response["nextCursor"] = serde_json::Value::String(next);
        }
        success_response(ctx.request.id, response)
    }

    /// Handle resources/templates/list request (with pagination support)
    async fn handle_resources_templates_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse cursor if present
        let cursor = ctx
            .params()
            .and_then(|p| p.get("cursor"))
            .and_then(|c| c.as_str());

        let visibility_ctx = self.visibility_context(&ctx.session);
        let all_templates = self
            .resource_manager
            .list_templates_for_session(&ctx.session, &visibility_ctx);

        // Apply pagination
        let result = crate::utils::paginate(&all_templates, cursor, crate::utils::DEFAULT_PAGE_SIZE);

        // Build response (only include nextCursor if present)
        let mut response = serde_json::json!({"resourceTemplates": result.items});
        if let Some(next) = result.next_cursor {
            response["nextCursor"] = serde_json::Value::String(next);
        }
        success_response(ctx.request.id, response)
    }

    /// Handle resources/read request
    async fn handle_resources_read(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse resource read params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let uri = match params.get("uri").and_then(|v| v.as_str()) {
            Some(uri) => uri,
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing 'uri' field").to_jsonrpc(),
                );
            }
        };

        // Read resource - ResourceContent already includes uri/mimeType
        match self
            .resource_manager
            .read(
                uri,
                std::collections::HashMap::new(),
                &ctx.session,
                &self.logger,
            )
            .await
        {
            Ok(contents) => {
                // ResourceContent.to_value() already produces MCP-compliant JSON with uri/mimeType
                let content_values: Vec<Value> = contents.iter().map(|c| c.to_value()).collect();
                success_response(
                    ctx.request.id,
                    serde_json::json!({"contents": content_values}),
                )
            }
            Err(e) => error_response(
                ctx.request.id,
                McpError::internal("resource_read_failed", e.to_string()).to_jsonrpc(),
            ),
        }
    }

    /// Handle resources/subscribe request
    ///
    /// Subscribes the current session to receive notifications when a resource changes.
    /// When the resource is updated, the server will send a `notifications/resources/updated`
    /// notification to this session.
    async fn handle_resources_subscribe(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Check if resources.subscribe capability is advertised
        let caps = self.capabilities.load();
        if caps.resources.is_none()
            || caps.resources.as_ref().and_then(|r| r.subscribe) != Some(true)
        {
            return error_response(
                ctx.request.id,
                McpError::validation(
                    "capability_not_supported",
                    "Resource subscriptions are not enabled on this server",
                )
                .to_jsonrpc(),
            );
        }

        // Parse subscribe params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let req: SubscribeRequest = match serde_json::from_value(params) {
            Ok(req) => req,
            Err(_) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing or invalid 'uri' field")
                        .to_jsonrpc(),
                );
            }
        };

        // Subscribe the session to the resource URI
        self.subscription_manager.subscribe(&ctx.session.id, &req.uri);

        tracing::debug!(
            session_id = %ctx.session.id,
            uri = %req.uri,
            "Session subscribed to resource"
        );

        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle resources/unsubscribe request
    ///
    /// Unsubscribes the current session from receiving notifications for a resource.
    async fn handle_resources_unsubscribe(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse unsubscribe params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let req: UnsubscribeRequest = match serde_json::from_value(params) {
            Ok(req) => req,
            Err(_) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing or invalid 'uri' field")
                        .to_jsonrpc(),
                );
            }
        };

        // Unsubscribe the session from the resource URI
        self.subscription_manager
            .unsubscribe(&ctx.session.id, &req.uri);

        tracing::debug!(
            session_id = %ctx.session.id,
            uri = %req.uri,
            "Session unsubscribed from resource"
        );

        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle prompts/list request (with pagination support)
    async fn handle_prompts_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse cursor if present
        let cursor = ctx
            .params()
            .and_then(|p| p.get("cursor"))
            .and_then(|c| c.as_str());

        let visibility_ctx = self.visibility_context(&ctx.session);
        let all_prompts = self
            .prompt_manager
            .list_for_session(&ctx.session, &visibility_ctx);

        // Apply pagination
        let result = crate::utils::paginate(&all_prompts, cursor, crate::utils::DEFAULT_PAGE_SIZE);

        // Build response (only include nextCursor if present)
        let mut response = serde_json::json!({"prompts": result.items});
        if let Some(next) = result.next_cursor {
            response["nextCursor"] = serde_json::Value::String(next);
        }
        success_response(ctx.request.id, response)
    }

    /// Handle prompts/get request
    async fn handle_prompts_get(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse prompt get params
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let prompt_name = match params.get("name").and_then(|v| v.as_str()) {
            Some(name) => name,
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing 'name' field").to_jsonrpc(),
                );
            }
        };

        let prompt_params = params.get("arguments").cloned().unwrap_or(Value::Null);

        // Call prompt
        match self
            .prompt_manager
            .call(prompt_name, prompt_params, &ctx.session, &self.logger)
            .await
        {
            Ok(result) => success_response(
                ctx.request.id,
                serde_json::to_value(result).expect("Failed to serialize prompt result"),
            ),
            Err(e) => error_response(
                ctx.request.id,
                McpError::internal("prompt_get_failed", e.to_string()).to_jsonrpc(),
            ),
        }
    }

    /// Handle roots/list request
    ///
    /// NOTE: This is a CLIENT capability - the server requests roots FROM the client.
    /// This handler would only be used if mcphost-rs is acting as a client.
    /// For server→client root requests, use `request_roots()` instead.
    async fn handle_roots_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // If mcphost-rs is acting as a server (normal case), we don't have roots to provide
        // The client provides roots TO us, not the other way around
        // Return empty list for now
        use crate::protocol::types::ListRootsResult;

        let result = ListRootsResult { roots: vec![] };

        success_response(
            ctx.request.id,
            serde_json::to_value(result).expect("Failed to serialize roots list"),
        )
    }

    /// Handle sampling/createMessage request
    ///
    /// NOTE: This is a CLIENT capability - the server requests LLM completions FROM the client.
    /// This is for when the MCP server needs the client to generate LLM responses.
    /// This is NOT for the server to generate responses itself.
    async fn handle_sampling_create_message(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Servers don't generate LLM responses, clients do
        // This would only be used if mcphost-rs acts as a client
        error_response(
            ctx.request.id,
            McpError::not_implemented(
                "sampling/createMessage is a client capability. Use ClientRequester.create_message() for server→client requests."
            ).to_jsonrpc(),
        )
    }

    /// Handle elicitation/create request
    ///
    /// NOTE: This is a CLIENT capability - the server requests user input FROM the client.
    /// This is for when the MCP server needs the client to prompt the user for structured input.
    async fn handle_elicitation_create(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Servers don't prompt users directly, clients do
        // This would only be used if mcphost-rs acts as a client
        error_response(
            ctx.request.id,
            McpError::not_implemented(
                "elicitation/create is a client capability. Use ClientRequester.create_elicitation() for server→client requests."
            ).to_jsonrpc(),
        )
    }

    /// Handle tasks/get request
    async fn handle_tasks_get(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        let params: crate::protocol::types::GetTaskParams = match ctx.params() {
            Some(p) => match serde_json::from_value(p.clone()) {
                Ok(params) => params,
                Err(_) => {
                    return error_response(
                        ctx.request.id,
                        McpError::validation("invalid_params", "Missing or invalid taskId")
                            .to_jsonrpc(),
                    );
                }
            },
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing taskId parameter").to_jsonrpc(),
                );
            }
        };

        match self
            .task_store
            .get_task_for_session(&params.task_id, &ctx.session.id)
            .await
        {
            Some(task) => success_response(ctx.request.id, serde_json::to_value(task).unwrap()),
            None => error_response(
                ctx.request.id,
                McpError::validation("invalid_params", "Task not found").to_jsonrpc(),
            ),
        }
    }

    /// Handle tasks/result request (blocks until terminal state)
    async fn handle_tasks_result(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        let params: crate::protocol::types::GetTaskParams = match ctx.params() {
            Some(p) => match serde_json::from_value(p.clone()) {
                Ok(params) => params,
                Err(_) => {
                    return error_response(
                        ctx.request.id,
                        McpError::validation("invalid_params", "Missing or invalid taskId")
                            .to_jsonrpc(),
                    );
                }
            },
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing taskId parameter").to_jsonrpc(),
                );
            }
        };

        // Verify session owns task
        if self
            .task_store
            .get_task_for_session(&params.task_id, &ctx.session.id)
            .await
            .is_none()
        {
            return error_response(
                ctx.request.id,
                McpError::validation("invalid_params", "Task not found").to_jsonrpc(),
            );
        }

        // Wait for result (5-minute timeout)
        match self
            .task_store
            .wait_for_result(&params.task_id, std::time::Duration::from_secs(300))
            .await
        {
            Ok(result) => success_response(ctx.request.id, result),
            Err(e) => error_response(
                ctx.request.id,
                McpError::internal("task_error", e.to_string()).to_jsonrpc(),
            ),
        }
    }

    /// Handle tasks/list request
    async fn handle_tasks_list(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Parse cursor if present
        let cursor = ctx
            .params()
            .and_then(|p| p.get("cursor"))
            .and_then(|c| c.as_str());

        let (tasks, next_cursor) = self
            .task_store
            .list_tasks(&ctx.session.id, cursor, 100)
            .await;

        success_response(
            ctx.request.id,
            serde_json::json!({
                "tasks": tasks,
                "nextCursor": next_cursor,
            }),
        )
    }

    /// Handle tasks/cancel request
    async fn handle_tasks_cancel(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        let params: crate::protocol::types::CancelTaskParams = match ctx.params() {
            Some(p) => match serde_json::from_value(p.clone()) {
                Ok(params) => params,
                Err(_) => {
                    return error_response(
                        ctx.request.id,
                        McpError::validation("invalid_params", "Missing or invalid taskId")
                            .to_jsonrpc(),
                    );
                }
            },
            None => {
                return error_response(
                    ctx.request.id,
                    McpError::validation("invalid_params", "Missing taskId parameter").to_jsonrpc(),
                );
            }
        };

        match self
            .task_store
            .cancel_task(&params.task_id, &ctx.session.id)
            .await
        {
            Ok(task) => success_response(ctx.request.id, serde_json::to_value(task).unwrap()),
            Err(e) => {
                let error_msg = match e {
                    crate::managers::task::TaskError::NotFound(_) => {
                        McpError::validation("invalid_params", "Task not found")
                    }
                    crate::managers::task::TaskError::AlreadyTerminal(status) => {
                        McpError::validation(
                            "invalid_params",
                            format!(
                                "Cannot cancel task: already in terminal status '{:?}'",
                                status
                            ),
                        )
                    }
                    _ => McpError::internal("task_error", e.to_string()),
                };
                error_response(ctx.request.id, error_msg.to_jsonrpc())
            }
        }
    }

    /// Handle completion/complete request
    ///
    /// Provides argument value completion suggestions for prompts and resources.
    /// This allows clients to get autocomplete suggestions while typing arguments.
    async fn handle_completion_complete(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            return error_response(ctx.request.id, err.to_jsonrpc());
        }

        // Check if server has completion capability
        let caps = self.capabilities.load();
        if caps.completion.is_none() {
            return error_response(
                ctx.request.id,
                McpError::not_implemented("Completion capability not enabled").to_jsonrpc(),
            );
        }

        // Parse completion request
        let params = ctx.params().cloned().unwrap_or(Value::Null);
        let request: crate::protocol::types::CompleteRequest = match serde_json::from_value(params)
        {
            Ok(req) => req,
            Err(e) => {
                return error_response(
                    ctx.request.id,
                    McpError::validation(
                        "invalid_params",
                        format!("Invalid completion request: {}", e),
                    )
                    .to_jsonrpc(),
                );
            }
        };

        // Execute completion
        match self.completion_manager.complete(request).await {
            Ok(result) => success_response(
                ctx.request.id,
                serde_json::to_value(result).expect("Failed to serialize completion result"),
            ),
            Err(e) => error_response(
                ctx.request.id,
                McpError::internal("completion_failed", e.to_string()).to_jsonrpc(),
            ),
        }
    }

    /// Handle notifications/cancelled notification
    ///
    /// This is a notification (no response expected). It signals that the client
    /// wants to cancel a previously submitted request. We look up the request ID
    /// and signal cancellation to any registered handlers.
    async fn handle_notifications_cancelled(&self, ctx: RequestContext) -> JsonRpcResponse {
        // Parse cancellation params
        let params: CancelledNotificationParams = match ctx.params() {
            Some(p) => match serde_json::from_value(p.clone()) {
                Ok(params) => params,
                Err(e) => {
                    tracing::warn!(error = %e, "Invalid notifications/cancelled params");
                    // Notifications don't get error responses, just return empty
                    return success_response(ctx.request.id, serde_json::json!({}));
                }
            },
            None => {
                tracing::warn!("notifications/cancelled missing params");
                return success_response(ctx.request.id, serde_json::json!({}));
            }
        };

        let request_id = params.request_id_string();
        tracing::debug!(
            request_id = %request_id,
            reason = ?params.reason,
            "Received cancellation request"
        );

        // Signal cancellation to any registered handlers
        self.cancellation_manager.cancel(&request_id, params.reason);

        // Notifications don't get responses, but we return an empty success
        // in case the run loop tries to write it (it won't for notifications)
        success_response(ctx.request.id, serde_json::json!({}))
    }
}

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

    #[tokio::test]
    async fn test_server_creation() {
        let server = Server::new("test-server", "1.0.0");
        assert_eq!(server.name(), "test-server");
        assert_eq!(server.version(), "1.0.0");
    }

    #[tokio::test]
    async fn test_ping() {
        let server = Server::new("test-server", "1.0.0");

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "ping".to_string(),
            params: None,
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());
    }

    #[tokio::test]
    async fn test_initialize() {
        let server = Server::new("test-server", "1.0.0");

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());

        // Check session was initialized
        let session = server.get_session("test-session").unwrap();
        assert!(session.is_initialized());
        assert_eq!(session.client_info.unwrap().name, "test-client");
    }

    #[tokio::test]
    async fn test_method_not_found() {
        let server = Server::new("test-server", "1.0.0");

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "unknown/method".to_string(),
            params: None,
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_none());
        assert!(response.error.is_some());
        assert_eq!(response.error.unwrap().code, -32601);
    }

    #[tokio::test]
    async fn test_requires_initialization() {
        let server = Server::new("test-server", "1.0.0");

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "tools/list".to_string(),
            params: None,
        };

        // Should fail without initialization
        let response = server.handle_request("test-session", request.clone()).await;
        assert!(response.error.is_some());

        // Initialize session
        let init_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", init_request).await;

        // Should succeed after initialization
        let response = server.handle_request("test-session", request).await;
        assert!(response.result.is_some());
    }

    #[tokio::test]
    async fn test_session_management() {
        let server = Server::new("test-server", "1.0.0");

        // Create session
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "ping".to_string(),
            params: None,
        };
        server.handle_request("session-1", request).await;

        // Session should exist
        assert!(server.get_session("session-1").is_some());

        // Remove session
        let removed = server.remove_session("session-1");
        assert!(removed.is_some());

        // Session should no longer exist
        assert!(server.get_session("session-1").is_none());
    }

    #[tokio::test]
    async fn test_capabilities_update() {
        let server = Server::new("test-server", "1.0.0");

        let caps = ServerCapabilities {
            tools: Some(crate::protocol::capabilities::ToolsCapability {
                list_changed: Some(true),
            }),
            ..Default::default()
        };

        server.set_capabilities(caps.clone());

        let loaded_caps = server.capabilities();
        assert_eq!(loaded_caps.tools, caps.tools);
    }

    // ========================================================================
    // Task Integration Tests
    // ========================================================================

    /// Helper: Initialize a test session
    async fn init_test_session(server: &Server, session_id: &str) {
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {
                    "tasks": {
                        "list": {},
                        "cancel": {},
                        "requests": {
                            "tools": {
                                "call": {}
                            }
                        }
                    }
                },
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };

        server.handle_request(session_id, request).await;
    }

    /// Test tool that completes immediately
    struct TestTaskTool;

    #[async_trait::async_trait]
    impl crate::registry::tools::Tool for TestTaskTool {
        fn name(&self) -> &str {
            "test_task"
        }

        fn description(&self) -> Option<&str> {
            Some("Test tool for task execution")
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({
                "type": "object",
                "properties": {
                    "message": {"type": "string"}
                }
            })
        }

        fn execution(&self) -> Option<crate::protocol::types::ToolExecution> {
            Some(crate::protocol::types::ToolExecution {
                task_support: Some(crate::protocol::types::TaskSupport::Optional),
            })
        }

        async fn execute(
            &self,
            ctx: crate::prelude::ExecutionContext<'_>,
        ) -> Result<crate::registry::tools::ToolOutput, crate::registry::tools::ToolError> {
            let msg = ctx
                .params
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("default");

            Ok(crate::registry::tools::ToolOutput::text(format!(
                "Processed: {}",
                msg
            )))
        }
    }

    /// Test tool that takes time to complete
    struct SlowTestTool;

    #[async_trait::async_trait]
    impl crate::registry::tools::Tool for SlowTestTool {
        fn name(&self) -> &str {
            "slow_test"
        }

        fn description(&self) -> Option<&str> {
            Some("Slow test tool")
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({"type": "object"})
        }

        fn execution(&self) -> Option<crate::protocol::types::ToolExecution> {
            Some(crate::protocol::types::ToolExecution {
                task_support: Some(crate::protocol::types::TaskSupport::Optional),
            })
        }

        async fn execute(
            &self,
            _ctx: crate::prelude::ExecutionContext<'_>,
        ) -> Result<crate::registry::tools::ToolOutput, crate::registry::tools::ToolError> {
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
            Ok(crate::registry::tools::ToolOutput::text(
                "Slow operation complete",
            ))
        }
    }

    #[tokio::test]
    async fn test_task_augmented_tool_call() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(TestTaskTool);

        init_test_session(&server, "test-session").await;

        // Call tool with task metadata
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "test_task",
                "arguments": {"message": "hello"},
                "task": {"ttl": 60000}
            })),
        };

        let response = server.handle_request("test-session", request).await;

        // Should get CreateTaskResult with task metadata
        assert!(response.result.is_some());
        assert!(response.error.is_none());

        let result = response.result.unwrap();
        assert!(result.get("task").is_some());

        let task = result.get("task").unwrap();
        assert!(task.get("taskId").is_some());
        assert_eq!(task.get("status").unwrap().as_str().unwrap(), "working");
        assert!(task.get("createdAt").is_some());
        assert_eq!(task.get("ttl").unwrap().as_u64().unwrap(), 60000);
    }

    #[tokio::test]
    async fn test_task_get_status() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(SlowTestTool);

        init_test_session(&server, "test-session").await;

        // Create task
        let create_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "slow_test",
                "arguments": {},
                "task": {"ttl": 60000}
            })),
        };

        let create_response = server.handle_request("test-session", create_request).await;
        let task_id = create_response.result.unwrap()["task"]["taskId"]
            .as_str()
            .unwrap()
            .to_string();

        // Get task status immediately (should be working)
        let get_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "tasks/get".to_string(),
            params: Some(serde_json::json!({"taskId": task_id})),
        };

        let get_response = server.handle_request("test-session", get_request).await;

        assert!(get_response.result.is_some());
        let result = get_response.result.unwrap();
        let status = result["status"].as_str().unwrap();
        assert!(status == "working" || status == "completed");
    }

    #[tokio::test]
    async fn test_task_result_blocking() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(SlowTestTool);

        init_test_session(&server, "test-session").await;

        // Create task
        let create_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "slow_test",
                "arguments": {},
                "task": {"ttl": 60000}
            })),
        };

        let create_response = server.handle_request("test-session", create_request).await;
        let task_id = create_response.result.unwrap()["task"]["taskId"]
            .as_str()
            .unwrap()
            .to_string();

        // Get result (should block until complete)
        let result_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "tasks/result".to_string(),
            params: Some(serde_json::json!({"taskId": task_id})),
        };

        let result_response = server.handle_request("test-session", result_request).await;

        assert!(result_response.result.is_some());
        assert!(result_response.error.is_none());

        // Should have tool result
        let result = result_response.result.unwrap();
        assert!(result.get("content").is_some());
    }

    #[tokio::test]
    async fn test_task_cancel() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(SlowTestTool);

        init_test_session(&server, "test-session").await;

        // Create task
        let create_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "slow_test",
                "arguments": {},
                "task": {"ttl": 60000}
            })),
        };

        let create_response = server.handle_request("test-session", create_request).await;
        let task_id = create_response.result.unwrap()["task"]["taskId"]
            .as_str()
            .unwrap()
            .to_string();

        // Cancel immediately
        let cancel_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "tasks/cancel".to_string(),
            params: Some(serde_json::json!({"taskId": task_id})),
        };

        let cancel_response = server.handle_request("test-session", cancel_request).await;

        // Should succeed if still working
        if cancel_response.result.is_some() {
            let result = cancel_response.result.unwrap();
            let status = result["status"].as_str().unwrap();
            assert_eq!(status, "cancelled");
        }
        // Or fail if already completed (timing-dependent)
    }

    #[tokio::test]
    async fn test_task_list() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(TestTaskTool);

        init_test_session(&server, "test-session").await;

        // Create multiple tasks
        for i in 0..3 {
            let request = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: Some(Value::Number((i + 2).into())),
                method: "tools/call".to_string(),
                params: Some(serde_json::json!({
                    "name": "test_task",
                    "arguments": {"message": format!("task-{}", i)},
                    "task": {"ttl": 60000}
                })),
            };
            server.handle_request("test-session", request).await;
        }

        // List tasks
        let list_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(10.into())),
            method: "tasks/list".to_string(),
            params: None,
        };

        let list_response = server.handle_request("test-session", list_request).await;

        assert!(list_response.result.is_some());
        let result = list_response.result.unwrap();
        let tasks = result["tasks"].as_array().unwrap();
        assert!(tasks.len() >= 3);
    }

    #[tokio::test]
    async fn test_task_session_isolation() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(TestTaskTool);

        init_test_session(&server, "session-1").await;
        init_test_session(&server, "session-2").await;

        // Create task in session-1
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "test_task",
                "arguments": {"message": "private"},
                "task": {"ttl": 60000}
            })),
        };

        let response = server.handle_request("session-1", request).await;
        let task_id = response.result.unwrap()["task"]["taskId"]
            .as_str()
            .unwrap()
            .to_string();

        // Try to access from session-2 (should fail)
        let get_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "tasks/get".to_string(),
            params: Some(serde_json::json!({"taskId": task_id})),
        };

        let get_response = server.handle_request("session-2", get_request).await;

        // Should return error (task not found for this session)
        assert!(get_response.error.is_some());
    }

    #[tokio::test]
    async fn test_task_not_found() {
        let server = Server::new("test-server", "1.0.0");
        init_test_session(&server, "test-session").await;

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tasks/get".to_string(),
            params: Some(serde_json::json!({"taskId": "nonexistent-task-id"})),
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.error.is_some());
        assert_eq!(response.error.unwrap().code, -32602);
    }

    #[tokio::test]
    async fn test_task_double_cancel() {
        let server = Server::new("test-server", "1.0.0");
        server.tool_registry().register(SlowTestTool);

        init_test_session(&server, "test-session").await;

        // Create task
        let create_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "tools/call".to_string(),
            params: Some(serde_json::json!({
                "name": "slow_test",
                "arguments": {},
                "task": {"ttl": 60000}
            })),
        };

        let create_response = server.handle_request("test-session", create_request).await;
        let task_id = create_response.result.unwrap()["task"]["taskId"]
            .as_str()
            .unwrap()
            .to_string();

        // Cancel first time
        let cancel_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "tasks/cancel".to_string(),
            params: Some(serde_json::json!({"taskId": task_id.clone()})),
        };

        let _ = server
            .handle_request("test-session", cancel_request.clone())
            .await;

        // Wait a bit to ensure cancellation completes
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;

        // Try to cancel again (should fail - already terminal)
        let cancel_request2 = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(4.into())),
            method: "tasks/cancel".to_string(),
            params: Some(serde_json::json!({"taskId": task_id})),
        };

        let cancel_response2 = server.handle_request("test-session", cancel_request2).await;

        // Second cancel should fail (already in terminal state)
        assert!(cancel_response2.error.is_some());
    }

    // ========================================================================
    // Elicitation Tests
    // ========================================================================

    #[tokio::test]
    async fn test_request_elicitation_no_capability() {
        let server = Server::new("test-server", "1.0.0");

        // Initialize session WITHOUT elicitation capability
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {
                    // No elicitation capability
                },
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", request).await;

        // Try to request elicitation
        let result = server
            .request_elicitation(
                "test-session",
                "Please provide your email",
                serde_json::json!({"type": "object", "properties": {"email": {"type": "string"}}}),
                None,
            )
            .await;

        // Should fail with UnsupportedCapability
        assert!(matches!(
            result,
            Err(MultiplexerError::UnsupportedCapability(cap)) if cap == "elicitation"
        ));
    }

    #[tokio::test]
    async fn test_request_elicitation_session_not_found() {
        let server = Server::new("test-server", "1.0.0");

        // Try to request elicitation for non-existent session
        let result = server
            .request_elicitation(
                "nonexistent-session",
                "Please provide your email",
                serde_json::json!({"type": "object", "properties": {"email": {"type": "string"}}}),
                None,
            )
            .await;

        // Should fail with Transport error
        assert!(matches!(result, Err(MultiplexerError::Transport(_))));
    }

    #[tokio::test]
    async fn test_client_requester_with_elicitation_capability() {
        let server = Server::new("test-server", "1.0.0");

        // Initialize session WITH elicitation capability
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {
                    "elicitation": {}
                },
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", request).await;

        // Create client requester and verify elicitation is supported
        let requester = server.create_client_requester("test-session");
        assert!(requester.is_some());
        let requester = requester.unwrap();
        assert!(requester.supports_elicitation());
    }

    // ========================================================================
    // Resource Subscription Tests
    // ========================================================================

    /// Helper: Initialize a test session with resource subscription capability
    async fn init_session_with_subscriptions(server: &Server, session_id: &str) {
        // First set server capabilities to support subscriptions
        server.set_capabilities(ServerCapabilities {
            resources: Some(crate::protocol::capabilities::ResourcesCapability {
                subscribe: Some(true),
                list_changed: Some(true),
                list_templates: Some(true),
            }),
            ..Default::default()
        });

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request(session_id, request).await;
    }

    #[tokio::test]
    async fn test_resources_subscribe() {
        let server = Server::new("test-server", "1.0.0");
        init_session_with_subscriptions(&server, "test-session").await;

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "resources/subscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());

        // Verify subscription was recorded
        assert!(server
            .subscription_manager()
            .is_subscribed("test-session", "file:///test.txt"));
    }

    #[tokio::test]
    async fn test_resources_unsubscribe() {
        let server = Server::new("test-server", "1.0.0");
        init_session_with_subscriptions(&server, "test-session").await;

        // First subscribe
        let subscribe_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "resources/subscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };
        server
            .handle_request("test-session", subscribe_request)
            .await;

        // Now unsubscribe
        let unsubscribe_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(3.into())),
            method: "resources/unsubscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };
        let response = server
            .handle_request("test-session", unsubscribe_request)
            .await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());

        // Verify subscription was removed
        assert!(!server
            .subscription_manager()
            .is_subscribed("test-session", "file:///test.txt"));
    }

    #[tokio::test]
    async fn test_resources_subscribe_without_capability() {
        let server = Server::new("test-server", "1.0.0");

        // Initialize WITHOUT subscription capability
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", request).await;

        // Try to subscribe (should fail)
        let subscribe_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "resources/subscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };
        let response = server
            .handle_request("test-session", subscribe_request)
            .await;

        assert!(response.error.is_some());
    }

    #[tokio::test]
    async fn test_session_removal_cleans_subscriptions() {
        let server = Server::new("test-server", "1.0.0");
        init_session_with_subscriptions(&server, "test-session").await;

        // Subscribe to multiple resources
        for uri in &["file:///a.txt", "file:///b.txt", "file:///c.txt"] {
            let request = JsonRpcRequest {
                jsonrpc: "2.0".to_string(),
                id: Some(Value::Number(2.into())),
                method: "resources/subscribe".to_string(),
                params: Some(serde_json::json!({"uri": uri})),
            };
            server.handle_request("test-session", request).await;
        }

        // Verify subscriptions exist
        assert_eq!(
            server
                .subscription_manager()
                .get_session_subscriptions("test-session")
                .len(),
            3
        );

        // Remove session
        server.remove_session("test-session");

        // Verify subscriptions were cleaned up
        assert_eq!(
            server
                .subscription_manager()
                .get_session_subscriptions("test-session")
                .len(),
            0
        );
    }

    #[tokio::test]
    async fn test_notify_resource_updated() {
        let server = Server::new("test-server", "1.0.0");
        init_session_with_subscriptions(&server, "test-session").await;

        // Subscribe
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "resources/subscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };
        server.handle_request("test-session", request).await;

        // Trigger update notification (this won't actually send anywhere in test,
        // but verifies the method exists and doesn't panic)
        server.notify_resource_updated("file:///test.txt");

        // The notification would be sent to the notification channel
        // In a real scenario, the transport would receive it
    }

    #[tokio::test]
    async fn test_subscription_requires_initialization() {
        let server = Server::new("test-server", "1.0.0");

        // Try to subscribe without initializing
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "resources/subscribe".to_string(),
            params: Some(serde_json::json!({"uri": "file:///test.txt"})),
        };

        let response = server.handle_request("test-session", request).await;

        // Should fail - not initialized
        assert!(response.error.is_some());
    }

    // ========================================================================
    // Completion Tests
    // ========================================================================

    /// Helper: Initialize a test session with completion capability
    async fn init_session_with_completion(server: &Server, session_id: &str) {
        // Set server capabilities to support completion
        server.set_capabilities(ServerCapabilities {
            completion: Some(crate::protocol::capabilities::CompletionCapability {}),
            ..Default::default()
        });

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request(session_id, request).await;
    }

    #[tokio::test]
    async fn test_completion_complete_prompt() {
        use crate::managers::completion::CompletionProvider;
        use crate::protocol::types::CompletionValue;
        use std::sync::Arc;

        let server = Server::new("test-server", "1.0.0");

        // Create a test completion provider
        struct TestProvider;

        #[async_trait::async_trait]
        impl CompletionProvider for TestProvider {
            async fn complete_prompt(
                &self,
                prompt_name: &str,
                argument_name: &str,
                argument_value: &str,
            ) -> Vec<CompletionValue> {
                if prompt_name == "greet" && argument_name == "name" {
                    vec!["Alice", "Amy", "Bob"]
                        .into_iter()
                        .filter(|v| v.to_lowercase().starts_with(&argument_value.to_lowercase()))
                        .map(CompletionValue::new)
                        .collect()
                } else {
                    vec![]
                }
            }
        }

        // Register provider
        server
            .completion_manager()
            .register_prompt_provider("greet", Arc::new(TestProvider));

        init_session_with_completion(&server, "test-session").await;

        // Request completions
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "completion/complete".to_string(),
            params: Some(serde_json::json!({
                "ref": {
                    "type": "ref/prompt",
                    "name": "greet"
                },
                "argument": {
                    "name": "name",
                    "value": "A"
                }
            })),
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());

        let result = response.result.unwrap();
        let values = result["completion"]["values"].as_array().unwrap();
        assert_eq!(values.len(), 2); // Alice, Amy
    }

    #[tokio::test]
    async fn test_completion_complete_no_capability() {
        let server = Server::new("test-server", "1.0.0");

        // Initialize WITHOUT completion capability (default)
        let init_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "initialize".to_string(),
            params: Some(serde_json::json!({
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", init_request).await;

        // Request completions should fail
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "completion/complete".to_string(),
            params: Some(serde_json::json!({
                "ref": {
                    "type": "ref/prompt",
                    "name": "greet"
                },
                "argument": {
                    "name": "name",
                    "value": "A"
                }
            })),
        };

        let response = server.handle_request("test-session", request).await;

        // Should fail - completion capability not enabled
        assert!(response.error.is_some());
    }

    #[tokio::test]
    async fn test_completion_complete_not_initialized() {
        let server = Server::new("test-server", "1.0.0");

        // Request completions without initialization
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "completion/complete".to_string(),
            params: Some(serde_json::json!({
                "ref": {
                    "type": "ref/prompt",
                    "name": "greet"
                },
                "argument": {
                    "name": "name",
                    "value": "A"
                }
            })),
        };

        let response = server.handle_request("test-session", request).await;

        // Should fail - not initialized
        assert!(response.error.is_some());
    }

    #[tokio::test]
    async fn test_completion_complete_invalid_params() {
        let server = Server::new("test-server", "1.0.0");
        init_session_with_completion(&server, "test-session").await;

        // Request completions with invalid params
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "completion/complete".to_string(),
            params: Some(serde_json::json!({
                "invalid": "params"
            })),
        };

        let response = server.handle_request("test-session", request).await;

        // Should fail - invalid params
        assert!(response.error.is_some());
        let error = response.error.unwrap();
        assert_eq!(error.code, -32602); // Invalid params
    }

    #[tokio::test]
    async fn test_completion_complete_resource() {
        use crate::managers::completion::CompletionProvider;
        use crate::protocol::types::CompletionValue;
        use std::sync::Arc;

        let server = Server::new("test-server", "1.0.0");

        // Create a test completion provider for resources
        struct FileProvider;

        #[async_trait::async_trait]
        impl CompletionProvider for FileProvider {
            async fn complete_resource(
                &self,
                _uri: &str,
                _argument_name: &str,
                argument_value: &str,
            ) -> Vec<CompletionValue> {
                vec!["src/main.rs", "src/lib.rs", "tests/test.rs"]
                    .into_iter()
                    .filter(|v| v.starts_with(argument_value))
                    .map(CompletionValue::new)
                    .collect()
            }
        }

        // Register provider
        server
            .completion_manager()
            .register_resource_provider("file:///", Arc::new(FileProvider));

        init_session_with_completion(&server, "test-session").await;

        // Request completions
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "completion/complete".to_string(),
            params: Some(serde_json::json!({
                "ref": {
                    "type": "ref/resource",
                    "uri": "file:///project"
                },
                "argument": {
                    "name": "path",
                    "value": "src/"
                }
            })),
        };

        let response = server.handle_request("test-session", request).await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());

        let result = response.result.unwrap();
        let values = result["completion"]["values"].as_array().unwrap();
        assert_eq!(values.len(), 2); // src/main.rs, src/lib.rs
    }
}