//! 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::future::Future;
use std::pin::Pin;
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::managers::cancellation::{CancellationManager, CancelledNotificationParams};
use crate::managers::completion::CompletionManager;
use crate::managers::subscription::SubscriptionManager;
use crate::protocol::capabilities::ServerCapabilities;
use crate::protocol::errors::McpError;
use crate::protocol::methods::McpMethod;
use crate::protocol::types::{JsonRpcRequest, JsonRpcResponse};
use crate::registry::prompts::PromptManager;
use crate::registry::resources::ResourceManager;
use crate::registry::tools::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::session_store::{NullSessionStore, SessionStore};
use crate::server::visibility::{Environment, VisibilityContext};
use crate::transport::traits::{IncomingMessage, JsonRpcNotification, Transport};

/// Main MCP server
#[derive(Clone)]
pub struct Server {
    /// Server name
    pub(crate) name: String,

    /// Server version
    pub(crate) version: String,

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

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

    /// Active sessions by session ID
    pub(crate) sessions: Arc<DashMap<String, Session>>,

    /// Middleware chain for request processing
    pub(crate) middleware: MiddlewareChain,

    /// Tool registry
    pub(crate) tool_registry: ToolRegistry,

    /// Resource manager
    pub(crate) resource_manager: ResourceManager,

    /// Prompt manager
    pub(crate) prompt_manager: PromptManager,

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

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

    /// Global logger for server notifications
    pub(crate) logger: crate::logging::McpLogger,

    /// Request multiplexer for server→client requests
    pub(crate) multiplexer: Arc<RequestMultiplexer>,

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

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

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

    /// Cancellation manager for tracking cancellable operations
    pub(crate) cancellation_manager: CancellationManager,

    /// Subscription manager for resource subscriptions
    pub(crate) subscription_manager: SubscriptionManager,

    /// Completion manager for argument value suggestions
    pub(crate) completion_manager: CompletionManager,

    /// Optional environment for visibility checks
    pub(crate) environment: Option<Arc<dyn Environment>>,

    /// Hook called after client sends notifications/initialized (client is fully ready).
    /// Receives the session ID and an optional ClientRequester for bidirectional requests.
    pub(crate) on_initialized_hook: Option<
        Arc<
            dyn Fn(String, Option<ClientRequester>) -> Pin<Box<dyn Future<Output = ()> + Send>>
                + Send
                + Sync,
        >,
    >,

    /// Optional session store for persistence across restarts
    session_store: Arc<dyn SessionStore>,

    /// Optional callback invoked when a new session is created.
    /// Use this to pre-initialize sessions (e.g., after hot-swap re-exec).
    session_initializer: Option<Arc<dyn Fn(&mut Session) + Send + Sync>>,
}

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,
        // wired to the notification channel for notifications/tasks/status (2025-11-25)
        let task_store = Arc::new(crate::managers::task::TaskStore::with_notifications(
            std::time::Duration::from_secs(300),
            std::time::Duration::from_secs(5),
            notification_tx.clone(),
        ));

        // 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: Arc::new(DashMap::new()),
            middleware: MiddlewareChain::new(),
            tool_registry: ToolRegistry::with_notifications(notification_tx.clone()),
            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,
            on_initialized_hook: None,
            session_store: Arc::new(NullSessionStore),
            session_initializer: None,
        }
    }

    /// Register a hook called after the client sends `notifications/initialized`.
    ///
    /// At this point the client is fully ready and will respond to server→client requests
    /// such as `roots/list` and `sampling/createMessage`.
    ///
    /// # Example
    /// ```rust,ignore
    /// server.on_initialized(|session_id, requester| async move {
    ///     if let Some(req) = requester {
    ///         if let Ok(roots) = req.request_roots(None).await {
    ///             tracing::info!(session = %session_id, roots = ?roots, "Client roots");
    ///         }
    ///     }
    /// });
    /// ```
    pub fn on_initialized<F, Fut>(&mut self, hook: F) -> &mut Self
    where
        F: Fn(String, Option<ClientRequester>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.on_initialized_hook = Some(Arc::new(move |session_id, requester| {
            Box::pin(hook(session_id, requester))
        }));
        self
    }

    /// Set the initialized hook directly (used by ServerBuilder).
    pub(crate) fn set_initialized_hook(
        &mut self,
        hook: Arc<
            dyn Fn(String, Option<ClientRequester>) -> Pin<Box<dyn Future<Output = ()> + Send>>
                + Send
                + Sync,
        >,
    ) {
        self.on_initialized_hook = Some(hook);
    }

    /// 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);
    }

    /// Set a session store for persistence across restarts.
    pub fn set_session_store(&mut self, store: Arc<dyn SessionStore>) {
        self.session_store = store;
    }

    /// Set a callback invoked when a new session is created.
    ///
    /// Use this for hot-swap scenarios where the server re-execs and needs
    /// to pre-initialize sessions so clients can continue without re-handshaking.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// server.set_session_initializer(|session| {
    ///     session.initialize(
    ///         Implementation::new("hot-swap", "0.0"),
    ///         Default::default(),
    ///         "2025-06-18".to_string(),
    ///     );
    /// });
    /// ```
    pub fn set_session_initializer(&mut self, f: impl Fn(&mut Session) + Send + Sync + 'static) {
        self.session_initializer = Some(Arc::new(f));
    }

    pub(crate) 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 a sender for the server's notification channel.
    ///
    /// Use this to wire background tasks (outside the normal tool execution path)
    /// so they can emit `notifications/progress`, `notifications/tasks/status`, etc.
    pub fn notification_sender(&self) -> mpsc::UnboundedSender<JsonRpcNotification> {
        self.notification_tx.clone()
    }

    /// 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
    }

    /// 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);
        let removed = self.sessions.remove(session_id).map(|(_, s)| s);
        if removed.is_some() {
            self.session_store.delete(session_id);
        }
        removed
    }

    /// 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.supports_form_elicitation(),
        ))
    }

    /// 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.supports_form_elicitation() {
                    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;

        // Channel for spawned handlers to send responses back to the write loop.
        // This breaks the deadlock: a tool handler can make server→client requests
        // (e.g. roots/list) while the event loop continues reading responses.
        let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel::<(
            crate::protocol::types::JsonRpcResponse,
            bool,
        )>();

        loop {
            tokio::select! {
                // Deliver responses produced by spawned request handlers
                Some((response, is_request)) = response_rx.recv() => {
                    if is_request {
                        if let Err(e) = transport.write_message(response).await {
                            tracing::error!(error = %e, "Failed to write response");
                            break;
                        }
                    }
                }

                // 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();

                            // Spawn the handler so the event loop stays free to route
                            // server→client responses (e.g. roots/list) while this runs.
                            let method = request.method.clone();
                            let server = self.clone();
                            let session_id = session_id.clone();
                            let tx = response_tx.clone();
                            tokio::spawn(async move {
                                tracing::debug!(%method, "Handling request");
                                let response = server.handle_request(&session_id, request).await;
                                let _ = tx.send((response, !is_notification));
                            });
                        }
                        Ok(IncomingMessage::Response(response)) => {
                            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 {
        // Fast path: session already exists in the map
        let session = if let Some(existing) = self.sessions.get(session_id) {
            existing.clone()
        } else {
            // Slow path: create a new session, load from store, run initializer
            let store = Arc::clone(&self.session_store);
            let initializer = self.session_initializer.clone();
            let notification_tx = self.notification_tx.clone();

            self.sessions
                .entry(session_id.to_string())
                .or_insert_with(|| {
                    let mut session = Session::with_id(session_id);
                    session.set_notification_channel(notification_tx.clone());
                    session.set_store(Arc::clone(&store));

                    if let Some(snapshot) = store.load(session_id) {
                        session.apply_snapshot(&snapshot);
                    }

                    if let Some(initializer) = initializer.as_ref() {
                        initializer(&mut session);
                    }

                    session
                })
                .clone()
        };

        // Keep a reference for post-request flush
        let session_ref = 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) => {
                session_ref.flush_if_dirty();
                return error_response(request.id, err.to_jsonrpc());
            }
        };

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

        let response = 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,
            McpMethod::NotificationsInitialized => self.handle_notifications_initialized(ctx).await,
            McpMethod::NotificationsProgress => self.handle_notifications_progress(ctx).await,
            McpMethod::NotificationsRootsListChanged => {
                self.handle_notifications_roots_list_changed(ctx).await
            }
            // Server→client only notifications — if client somehow sends these, ignore silently
            McpMethod::NotificationsMessage
            | McpMethod::NotificationsResourcesListChanged
            | McpMethod::NotificationsResourcesUpdated
            | McpMethod::NotificationsPromptsListChanged
            | McpMethod::NotificationsToolsListChanged
            | McpMethod::NotificationsTasksStatus
            | McpMethod::NotificationsElicitationComplete => self.handle_notification_ignored(ctx),
            // Unknown method: error for requests, silent ignore for notifications (per spec)
            _ => {
                if request.is_notification() {
                    self.handle_notification_ignored(ctx)
                } else {
                    error_response(
                        request.id,
                        McpError::method_not_found(&request.method).to_jsonrpc(),
                    )
                }
            }
        };

        // Flush any dirty state accumulated during the request
        session_ref.flush_if_dirty();

        response
    }

    /// 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 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!({}))
    }

    /// Handle `notifications/initialized` — client signals it is fully ready.
    ///
    /// This is the correct moment to fire server→client requests (roots/list, sampling)
    /// because the client's message loop is running and will respond.
    async fn handle_notifications_initialized(&self, ctx: RequestContext) -> JsonRpcResponse {
        let session_id = ctx.session.id.clone();
        tracing::debug!(session = %session_id, "Client initialized");

        if let Some(hook) = &self.on_initialized_hook {
            let hook = hook.clone();
            let requester = self.create_client_requester(&session_id);
            tokio::spawn(async move {
                hook(session_id, requester).await;
            });
        }

        // Notifications never get a response
        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle `notifications/progress` — client sends progress update for a server-initiated
    /// request (e.g. roots/list with _meta.progressToken).
    ///
    /// Per spec, `notifications/progress` is bidirectional. When the server sends a request
    /// to the client with `_meta.progressToken`, the client may respond with this notification.
    async fn handle_notifications_progress(&self, ctx: RequestContext) -> JsonRpcResponse {
        if let Some(params) = ctx.params() {
            let token = params
                .get("progressToken")
                .and_then(|v| v.as_str())
                .unwrap_or("<unknown>");
            let progress = params
                .get("progress")
                .and_then(|v| v.as_f64())
                .unwrap_or(0.0);
            let total = params.get("total").and_then(|v| v.as_f64());
            let message = params.get("message").and_then(|v| v.as_str());

            match (total, message) {
                (Some(t), Some(m)) => tracing::debug!(
                    token = %token,
                    progress = %progress,
                    total = %t,
                    message = %m,
                    "Received progress notification from client"
                ),
                (Some(t), None) => tracing::debug!(
                    token = %token,
                    progress = %progress,
                    total = %t,
                    "Received progress notification from client"
                ),
                (None, Some(m)) => tracing::debug!(
                    token = %token,
                    progress = %progress,
                    message = %m,
                    "Received progress notification from client"
                ),
                (None, None) => tracing::debug!(
                    token = %token,
                    progress = %progress,
                    "Received progress notification from client"
                ),
            }
        } else {
            tracing::debug!("Received notifications/progress with no params");
        }

        // Notifications never get a response
        success_response(ctx.request.id, serde_json::json!({}))
    }

    /// Handle `notifications/roots/list_changed` — client signals its workspace roots changed.
    ///
    /// The server should re-query roots if it maintains a cache. Per spec, this is a
    /// hint that a previous roots/list result may be stale.
    async fn handle_notifications_roots_list_changed(
        &self,
        ctx: RequestContext,
    ) -> JsonRpcResponse {
        if let Err(err) = require_initialization(&ctx) {
            // Notifications don't get error responses; just log and return empty
            tracing::warn!(error = %err, "notifications/roots/list_changed before initialization");
            return success_response(ctx.request.id, serde_json::json!({}));
        }

        tracing::debug!(
            session = %ctx.session.id,
            "Client roots changed; server should re-query roots/list if caching"
        );

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

    /// Silently ignore a notification — returns empty success without logging an error.
    ///
    /// Used for:
    /// - Server→client-only notifications that the client incorrectly sends back
    /// - Unknown notifications (spec §3.3: receivers MUST silently ignore unknown notifications)
    fn handle_notification_ignored(&self, ctx: RequestContext) -> JsonRpcResponse {
        tracing::trace!(method = %ctx.request.method, "Ignoring notification (no-op)");
        success_response(ctx.request.id, serde_json::json!({}))
    }
}

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

    // Helper: build a notification request (id: None)
    fn notification(method: &str, params: Option<serde_json::Value>) -> JsonRpcRequest {
        JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: None,
            method: method.to_string(),
            params,
        }
    }

    #[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_unknown_notification_ignored() {
        let server = Server::new("test-server", "1.0.0");

        let response = server
            .handle_request(
                "test-session",
                notification("notifications/custom/foo", None),
            )
            .await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());
        assert_eq!(response.result.unwrap(), serde_json::json!({}));
    }

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

        let response = server
            .handle_request(
                "test-session",
                notification(
                    "notifications/progress",
                    Some(serde_json::json!({
                        "progressToken": "t1",
                        "progress": 5,
                        "total": 10,
                        "message": "halfway"
                    })),
                ),
            )
            .await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());
        assert_eq!(response.result.unwrap(), serde_json::json!({}));
    }

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

        let response = server
            .handle_request("test-session", notification("notifications/progress", None))
            .await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());
        assert_eq!(response.result.unwrap(), serde_json::json!({}));
    }

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

        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"
                }
            })),
        };
        let _ = server.handle_request("test-session", init_request).await;

        let response = server
            .handle_request(
                "test-session",
                notification("notifications/roots/list_changed", None),
            )
            .await;

        assert!(response.result.is_some());
        assert!(response.error.is_none());
        assert_eq!(response.result.unwrap(), serde_json::json!({}));
    }

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

        for method in [
            "notifications/message",
            "notifications/resources/list_changed",
            "notifications/resources/updated",
            "notifications/prompts/list_changed",
            "notifications/tools/list_changed",
            "notifications/tasks/status",
            "notifications/elicitation/complete",
        ] {
            let response = server
                .handle_request("test-session", notification(method, None))
                .await;

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

    #[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);
    }

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

        // Drain any pre-existing notifications before the state transition we care about.
        {
            let mut rx = server.notification_rx.lock().await;
            while rx.try_recv().is_ok() {}
        }

        assert!(server.tool_registry().disable_tool("test_task"));

        let notification = {
            let mut rx = server.notification_rx.lock().await;
            tokio::time::timeout(Duration::from_millis(250), rx.recv())
                .await
                .expect("timed out waiting for tools/list_changed notification")
                .expect("notification channel closed unexpectedly")
        };

        assert_eq!(notification.method, "notifications/tools/list_changed");
        assert!(notification.params.is_none());
    }

    // ========================================================================
    // 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;
    }

    struct HiddenTestPrompt;

    impl crate::registry::prompts::Prompt for HiddenTestPrompt {
        fn name(&self) -> &str {
            "hidden_prompt"
        }

        fn get<'a>(
            &'a self,
            _ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::prompts::PromptFuture<'a> {
            Box::pin(async move {
                Ok(crate::registry::prompts::GetPromptResult {
                    description: Some("Hidden test prompt".to_string()),
                    messages: vec![crate::content::prompt::PromptMessage::assistant("hello")],
                })
            })
        }
    }

    struct HiddenTestResource;

    impl crate::registry::resources::Resource for HiddenTestResource {
        fn uri(&self) -> &str {
            "test://hidden-resource"
        }

        fn name(&self) -> &str {
            "hidden_resource"
        }

        fn read<'a>(
            &'a self,
            _ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::resources::ResourceReadFuture<'a> {
            Box::pin(async move { Ok(vec![self.text_content("secret")]) })
        }
    }

    struct InvisibleTestPrompt;

    impl crate::registry::prompts::Prompt for InvisibleTestPrompt {
        fn name(&self) -> &str {
            "invisible_prompt"
        }

        fn is_visible(&self, _ctx: &crate::server::visibility::VisibilityContext) -> bool {
            false
        }

        fn get<'a>(
            &'a self,
            _ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::prompts::PromptFuture<'a> {
            Box::pin(async move {
                Ok(crate::registry::prompts::GetPromptResult {
                    description: Some("Invisible test prompt".to_string()),
                    messages: vec![crate::content::prompt::PromptMessage::assistant("hello")],
                })
            })
        }
    }

    struct InvisibleTestResource;

    impl crate::registry::resources::Resource for InvisibleTestResource {
        fn uri(&self) -> &str {
            "test://invisible-resource"
        }

        fn name(&self) -> &str {
            "invisible_resource"
        }

        fn is_visible(&self, _ctx: &crate::server::visibility::VisibilityContext) -> bool {
            false
        }

        fn read<'a>(
            &'a self,
            _ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::resources::ResourceReadFuture<'a> {
            Box::pin(async move { Ok(vec![self.text_content("secret")]) })
        }
    }

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

        let session = server.get_session("test-session").unwrap();
        session.hide_prompt("hidden_prompt", false);

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "prompts/get".to_string(),
            params: Some(serde_json::json!({
                "name": "hidden_prompt",
                "arguments": {}
            })),
        };

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

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

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

        let session = server.get_session("test-session").unwrap();
        session.hide_resource("test://hidden-resource", false);

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

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

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

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

        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(2.into())),
            method: "prompts/get".to_string(),
            params: Some(serde_json::json!({
                "name": "invisible_prompt",
                "arguments": {}
            })),
        };

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

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

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

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

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

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

    /// Test tool that completes immediately
    struct TestTaskTool;

    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),
            })
        }

        fn execute<'a>(
            &'a self,
            ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::tools::ToolFuture<'a> {
            Box::pin(async move {
                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;

    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),
            })
        }

        fn execute<'a>(
            &'a self,
            _ctx: crate::prelude::ExecutionContext<'a>,
        ) -> crate::registry::tools::ToolFuture<'a> {
            Box::pin(async move {
                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());
    }

    #[tokio::test]
    async fn test_client_requester_with_url_only_elicitation_capability() {
        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": {
                    "elicitation": {
                        "url": {}
                    }
                },
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            })),
        };
        server.handle_request("test-session", request).await;

        let requester = server.create_client_requester("test-session").unwrap();
        assert!(!requester.supports_elicitation());

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

        assert!(matches!(
            result,
            Err(MultiplexerError::UnsupportedCapability(cap)) if cap == "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;

        impl CompletionProvider for TestProvider {
            fn complete_prompt<'a>(
                &'a self,
                prompt_name: &'a str,
                argument_name: &'a str,
                argument_value: &'a str,
            ) -> crate::managers::completion::CompletionFuture<'a> {
                Box::pin(async move {
                    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;

        impl CompletionProvider for FileProvider {
            fn complete_resource<'a>(
                &'a self,
                _uri: &'a str,
                _argument_name: &'a str,
                argument_value: &'a str,
            ) -> crate::managers::completion::CompletionFuture<'a> {
                Box::pin(async move {
                    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
    }
}

#[cfg(test)]
mod bidirectional_tests {
    use super::*;
    use crate::prelude::ExecutionContext;
    use crate::registry::tools::{Tool, ToolFuture, ToolOutput};
    use crate::transport::{InProcessTransport, IncomingMessage, Transport};

    struct GetRootsTool;

    impl Tool for GetRootsTool {
        fn name(&self) -> &str {
            "get_roots"
        }
        fn description(&self) -> Option<&str> {
            Some("Returns client roots")
        }
        fn input_schema(&self) -> Value {
            serde_json::json!({"type": "object", "properties": {}})
        }
        fn execute<'a>(&'a self, ctx: ExecutionContext<'a>) -> ToolFuture<'a> {
            Box::pin(async move {
                let requester = ctx.client_requester().ok_or_else(|| {
                    crate::registry::tools::ToolError::Execution("no requester".into())
                })?;
                let roots = tokio::time::timeout(
                    tokio::time::Duration::from_secs(5),
                    requester.request_roots(None),
                )
                .await
                .map_err(|_| crate::registry::tools::ToolError::Execution("timeout".into()))?
                .map_err(|e| crate::registry::tools::ToolError::Execution(e.to_string()))?;
                Ok(ToolOutput::text(format!("roots={}", roots.len())))
            })
        }
    }

    #[tokio::test]
    async fn tool_can_call_roots_list_without_deadlock() {
        let server = Arc::new(Server::new("test", "1.0"));
        server.tool_registry().register(GetRootsTool);

        let (server_transport, mut client_transport) = InProcessTransport::pair();

        let server_task = tokio::spawn({
            let server = server.clone();
            async move {
                server
                    .run(server_transport)
                    .await
                    .expect("server.run should succeed");
            }
        });

        client_transport
            .send_request(crate::server::multiplexer::JsonRpcClientRequest::new(
                "1",
                "initialize",
                Some(serde_json::json!({
                    "protocolVersion": "2025-11-25",
                    "capabilities": {"roots": {}},
                    "clientInfo": {"name": "test", "version": "1.0"}
                })),
            ))
            .await
            .unwrap();

        let initialize_response = match tokio::time::timeout(
            tokio::time::Duration::from_secs(5),
            client_transport.read_incoming(),
        )
        .await
        .expect("timed out waiting for initialize response")
        .unwrap()
        {
            IncomingMessage::Response(response) => response,
            IncomingMessage::Request(request) => {
                panic!(
                    "expected initialize response, got request {}",
                    request.method
                )
            }
        };

        assert_eq!(initialize_response.id, serde_json::json!("1"));

        client_transport
            .send_request(crate::server::multiplexer::JsonRpcClientRequest::new(
                "2",
                "tools/call",
                Some(serde_json::json!({
                    "name": "get_roots",
                    "arguments": {}
                })),
            ))
            .await
            .unwrap();

        let mut saw_roots_list = false;
        let mut tool_result_text = None;

        for _ in 0..4 {
            match tokio::time::timeout(
                tokio::time::Duration::from_secs(5),
                client_transport.read_incoming(),
            )
            .await
            .expect("timed out waiting for bidirectional exchange")
            .unwrap()
            {
                IncomingMessage::Request(request) if request.method == "roots/list" => {
                    saw_roots_list = true;
                    client_transport
                        .write_message(JsonRpcResponse::success(
                            request.id.expect("roots/list must carry an id"),
                            serde_json::json!({
                                "roots": [{"uri": "file:///workspace"}]
                            }),
                        ))
                        .await
                        .unwrap();
                }
                IncomingMessage::Response(response) if response.id == serde_json::json!("2") => {
                    let text = response
                        .result
                        .as_ref()
                        .and_then(|result| result.get("content"))
                        .and_then(|content| content.as_array())
                        .and_then(|items| items.first())
                        .and_then(|item| item.get("text"))
                        .and_then(|text| text.as_str())
                        .map(str::to_owned);
                    tool_result_text = text;
                    break;
                }
                IncomingMessage::Request(request) => {
                    panic!("unexpected request during test: {}", request.method)
                }
                IncomingMessage::Response(response) => {
                    panic!("unexpected response during test: {:?}", response.id)
                }
            }
        }

        assert!(saw_roots_list, "expected server to request roots/list");
        assert_eq!(tool_result_text.as_deref(), Some("roots=1"));

        client_transport.shutdown().await.unwrap();

        tokio::time::timeout(tokio::time::Duration::from_secs(5), server_task)
            .await
            .expect("run() timed out — deadlock present in bidirectional handling")
            .expect("server task failed");
    }
}
