pmcp 2.4.0

//! Cancellation token infrastructure and the canonical
//! [`RequestHandlerExtra`] request-context struct handed to every tool,
//! prompt, and resource handler.
//!
//! [`RequestHandlerExtra`] carries two cross-cutting request-scoped surfaces:
//!
//! - `extensions: http::Extensions` — typed-key typemap for cross-middleware
//!   state transfer. Insert/retrieve typed values via
//!   [`RequestHandlerExtra::extensions_mut`] / [`RequestHandlerExtra::extensions`].
//! - `peer: Option<Arc<dyn PeerHandle>>` — server-to-client back-channel
//!   exposing `sample()`, `list_roots()`, and `progress_notify()` from inside
//!   tool/prompt/resource handlers. Non-wasm only. `None` when the enclosing
//!   [`crate::server::core::ServerCore`] was not configured with a
//!   [`crate::server::server_request_dispatcher::ServerRequestDispatcher`]
//!   (tests, custom integrations).
//!
//! # Semver posture
//!
//! [`RequestHandlerExtra`] is now `#[non_exhaustive]`. This is a **breaking
//! change** for any downstream crate that constructed `RequestHandlerExtra`
//! with a positional struct literal. It is **not breaking** for code that
//! uses [`RequestHandlerExtra::new`], [`RequestHandlerExtra::default`], or
//! the `.with_*(...)` builder chain. Migration path: switch positional
//! literals to
//! `RequestHandlerExtra::new(request_id, cancellation_token)
//!     .with_auth_context(...)
//!     .with_peer(...)`.
//!
//! # Known limitation: session-id plumbing
//!
//! The `session_id` field on [`RequestHandlerExtra`] is not populated at
//! dispatch time in v2.2 — [`crate::server::auth::AuthContext`] does not
//! carry a `session_id` and `ProtocolHandler::handle_request` does not thread
//! one. Peer session isolation is therefore enforced at the
//! [`crate::server::Server`] level: each `Server` instance owns one
//! dispatcher bound to one transport, so cross-session confusion requires
//! cross-process access which is out of the current threat model. Follow-on
//! work can plumb `session_id` through `ProtocolHandler::handle_request` when
//! rmcp-parity for session-scoped auth becomes a scheduled phase goal.
//!
//! # Usage
//!
//! ```rust,no_run
//! use pmcp::RequestHandlerExtra;
//!
//! #[derive(Clone)]
//! struct RequestContext { user_id: u64 }
//!
//! // Middleware populates cross-cutting state before the handler runs.
//! let mut extra = RequestHandlerExtra::default();
//! extra.extensions_mut().insert(RequestContext { user_id: 42 });
//!
//! // Handler retrieves the typed value.
//! let ctx = extra.extensions().get::<RequestContext>().cloned();
//! assert!(ctx.is_some());
//! ```
//!
//! For an end-to-end runnable demonstration see
//! `examples/s42_handler_extensions.rs` (extensions) and
//! `examples/s43_handler_peer_sample.rs` (peer.sample from inside a
//! `ToolHandler`).

use crate::error::Result;
use crate::server::progress::ProgressReporter;
use crate::types::{CancelledNotification, Notification};
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(not(target_arch = "wasm32"))]
use tokio::sync::RwLock;
#[cfg(not(target_arch = "wasm32"))]
use tokio_util::sync::CancellationToken;

/// Manages cancellation tokens for requests.
pub struct CancellationManager {
    tokens: Arc<RwLock<HashMap<String, CancellationToken>>>,
    notification_sender: Option<Arc<dyn Fn(Notification) + Send + Sync>>,
}

impl CancellationManager {
    /// Create a new cancellation manager.
    pub fn new() -> Self {
        Self {
            tokens: Arc::new(RwLock::new(HashMap::new())),
            notification_sender: None,
        }
    }

    /// Set the notification sender.
    pub fn set_notification_sender(&mut self, sender: Arc<dyn Fn(Notification) + Send + Sync>) {
        self.notification_sender = Some(sender);
    }

    /// Create a cancellation token for a request.
    pub async fn create_token(&self, request_id: String) -> CancellationToken {
        let token = CancellationToken::new();
        let mut tokens = self.tokens.write().await;
        tokens.insert(request_id, token.clone());
        token
    }

    /// Cancel a request by ID.
    pub async fn cancel_request(&self, request_id: String, reason: Option<String>) -> Result<()> {
        let token = {
            let mut tokens = self.tokens.write().await;
            tokens.remove(&request_id)
        };

        if let Some(token) = token {
            // Cancel the token
            token.cancel();

            // Send cancellation notification
            if let Some(sender) = &self.notification_sender {
                let notification =
                    Notification::Client(crate::types::ClientNotification::Cancelled(
                        CancelledNotification::new(crate::types::RequestId::String(
                            request_id.clone(),
                        ))
                        .with_reason(reason.unwrap_or_else(|| "Cancelled by server".to_string())),
                    ));
                sender(notification);
            }
        }

        Ok(())
    }

    /// Remove a completed request's token.
    pub async fn remove_token(&self, request_id: &str) {
        let mut tokens = self.tokens.write().await;
        tokens.remove(request_id);
    }

    /// Check if a request is cancelled.
    pub async fn is_cancelled(&self, request_id: &str) -> bool {
        let tokens = self.tokens.read().await;
        tokens
            .get(request_id)
            .is_some_and(tokio_util::sync::CancellationToken::is_cancelled)
    }

    /// Get the cancellation token for a request.
    pub async fn get_token(&self, request_id: &str) -> Option<CancellationToken> {
        let tokens = self.tokens.read().await;
        tokens.get(request_id).cloned()
    }

    /// Clear all cancellation tokens.
    pub async fn clear(&self) {
        let mut tokens = self.tokens.write().await;
        // Cancel all active tokens
        for token in tokens.values() {
            token.cancel();
        }
        tokens.clear();
    }
}

impl Default for CancellationManager {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for CancellationManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CancellationManager")
            .field(
                "active_tokens",
                &self.tokens.try_read().map_or(0, |t| t.len()),
            )
            .finish()
    }
}

/// Extra context passed to request handlers.
#[derive(Clone)]
#[non_exhaustive]
pub struct RequestHandlerExtra {
    /// Cancellation token for the request
    pub cancellation_token: CancellationToken,
    /// Request ID
    pub request_id: String,
    /// Session ID
    pub session_id: Option<String>,
    /// Authentication info
    pub auth_info: Option<crate::types::auth::AuthInfo>,
    /// Validated authentication context (if auth is enabled)
    pub auth_context: Option<crate::server::auth::AuthContext>,
    /// Custom metadata for middleware (e.g., OAuth tokens, session data)
    ///
    /// **Security Note**: Metadata may contain sensitive values like OAuth tokens.
    /// The Debug implementation redacts these values to prevent accidental logging.
    pub metadata: HashMap<String, String>,
    /// Optional progress reporter for this request
    #[allow(dead_code)]
    pub progress_reporter: Option<Arc<dyn ProgressReporter>>,
    /// Task augmentation request from the client (MCP Tasks).
    ///
    /// When `Some`, the client supports async task polling and requested
    /// task-augmented behavior for this call. The tool handler can check
    /// this to decide between a sync path (await and return results) or
    /// an async path (create task, return immediately).
    ///
    /// When `None`, the client does not support tasks or did not request
    /// task mode — the tool should return results synchronously.
    pub task_request: Option<serde_json::Value>,
    /// Typed request-scoped state for middleware→handler transfer.
    ///
    /// Inserting values requires `T: Clone + Send + Sync + 'static`. Debug prints type names only,
    /// not values, making this safe for logging. Cloning `RequestHandlerExtra` clones the entire
    /// extensions map — prefer `Arc<T>` for large values.
    pub extensions: http::Extensions,
    /// Server-to-client back-channel for `sample` / `list_roots` / `progress_notify`.
    ///
    /// `None` when the request originated from a code path (wasm dispatch, unit-test
    /// fixture, any `ServerCore::new(...)` without
    /// [`crate::server::core::ServerCore::with_server_request_dispatcher`]) that did
    /// not configure a `ServerRequestDispatcher`. Tool/prompt/resource handlers
    /// running under a fully-configured `ServerCore` receive a populated peer handle
    /// that routes `sample` / `list_roots` / `progress_notify` back to the originating
    /// client.
    ///
    /// # Session isolation
    /// Each `Server` instance owns its own dispatcher, so peer handles are scoped to
    /// the enclosing `Server` — cross-session confusion requires cross-process access.
    #[cfg(not(target_arch = "wasm32"))]
    pub peer: Option<Arc<dyn crate::shared::peer::PeerHandle>>,
}

impl RequestHandlerExtra {
    /// Create new handler extra context.
    pub fn new(request_id: String, cancellation_token: CancellationToken) -> Self {
        Self {
            cancellation_token,
            request_id,
            session_id: None,
            auth_info: None,
            auth_context: None,
            metadata: HashMap::new(),
            progress_reporter: None,
            task_request: None,
            extensions: http::Extensions::new(),
            #[cfg(not(target_arch = "wasm32"))]
            peer: None,
        }
    }

    /// Set the session ID.
    pub fn with_session_id(mut self, session_id: Option<String>) -> Self {
        self.session_id = session_id;
        self
    }

    /// Set the auth info.
    pub fn with_auth_info(mut self, auth_info: Option<crate::types::auth::AuthInfo>) -> Self {
        self.auth_info = auth_info;
        self
    }

    /// Set the auth context.
    pub fn with_auth_context(
        mut self,
        auth_context: Option<crate::server::auth::AuthContext>,
    ) -> Self {
        self.auth_context = auth_context;
        self
    }

    /// Attach a progress reporter.
    pub fn with_progress_reporter(
        mut self,
        progress_reporter: Option<Arc<dyn ProgressReporter>>,
    ) -> Self {
        self.progress_reporter = progress_reporter;
        self
    }

    /// Set the task request from the client's `tools/call` params.
    ///
    /// When present, the tool handler knows the client supports task-augmented
    /// responses and can choose an async path (create task, return immediately)
    /// instead of awaiting the full result.
    pub fn with_task_request(mut self, task_request: Option<serde_json::Value>) -> Self {
        self.task_request = task_request;
        self
    }

    /// Returns a reference to the typed extensions map.
    ///
    /// # Example
    /// ```rust,no_run
    /// use pmcp::RequestHandlerExtra;
    /// let extra = RequestHandlerExtra::default();
    /// let _: Option<&String> = extra.extensions().get::<String>();
    /// ```
    pub fn extensions(&self) -> &http::Extensions {
        &self.extensions
    }

    /// Returns a mutable reference to the typed extensions map.
    ///
    /// # Example
    /// ```rust,no_run
    /// use pmcp::RequestHandlerExtra;
    /// let mut extra = RequestHandlerExtra::default();
    /// extra.extensions_mut().insert(42u64);
    /// ```
    pub fn extensions_mut(&mut self) -> &mut http::Extensions {
        &mut self.extensions
    }

    /// Attach a peer handle for server-to-client RPCs.
    ///
    /// When set, tool/prompt/resource handlers can invoke `extra.peer().unwrap().sample(...)`
    /// (or `list_roots` / `progress_notify`) to make outbound requests back to the
    /// originating client. The peer handle is populated by the enclosing
    /// `ServerCore` at each dispatch site when a
    /// [`crate::server::server_request_dispatcher::ServerRequestDispatcher`]
    /// has been attached — tests or ad-hoc constructions that skip the dispatcher
    /// leave this as `None`, and handlers should treat `None` as "no client
    /// available for back-channel".
    ///
    /// # Example
    /// ```rust,no_run
    /// # #[cfg(not(target_arch = "wasm32"))]
    /// # {
    /// use pmcp::RequestHandlerExtra;
    /// use std::sync::Arc;
    /// # fn build_peer() -> Arc<dyn pmcp::PeerHandle> { unimplemented!() }
    /// let extra = RequestHandlerExtra::default().with_peer(build_peer());
    /// # }
    /// ```
    #[cfg(not(target_arch = "wasm32"))]
    pub fn with_peer(mut self, peer: Arc<dyn crate::shared::peer::PeerHandle>) -> Self {
        self.peer = Some(peer);
        self
    }

    /// Returns the peer handle, if one was attached.
    ///
    /// Handlers should treat `None` as "no client back-channel available" and
    /// skip any sample/list_roots-dependent code paths gracefully.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn peer(&self) -> Option<&Arc<dyn crate::shared::peer::PeerHandle>> {
        self.peer.as_ref()
    }

    /// Returns `true` if the client requested task-augmented behavior.
    pub fn is_task_request(&self) -> bool {
        self.task_request.is_some()
    }

    /// Get the auth context if available.
    pub fn auth_context(&self) -> Option<&crate::server::auth::AuthContext> {
        self.auth_context.as_ref()
    }

    /// Get metadata value by key.
    ///
    /// Metadata is typically set by middleware (e.g., OAuth token injection).
    pub fn get_metadata(&self, key: &str) -> Option<&String> {
        self.metadata.get(key)
    }

    /// Set metadata value.
    ///
    /// This is typically used by middleware to inject data for tools to consume.
    pub fn set_metadata(&mut self, key: String, value: String) {
        self.metadata.insert(key, value);
    }

    /// Check if the request has been cancelled.
    pub fn is_cancelled(&self) -> bool {
        self.cancellation_token.is_cancelled()
    }

    /// Wait for cancellation.
    pub async fn cancelled(&self) {
        self.cancellation_token.cancelled().await;
    }

    /// Report progress if a reporter is available.
    pub async fn report_progress(
        &self,
        progress: f64,
        total: Option<f64>,
        message: Option<String>,
    ) -> crate::Result<()> {
        if let Some(rep) = &self.progress_reporter {
            rep.report_progress(progress, total, message).await
        } else {
            Ok(())
        }
    }

    /// Report percentage progress (0-100) if available.
    pub async fn report_percent(&self, percent: f64, message: Option<String>) -> crate::Result<()> {
        if let Some(rep) = &self.progress_reporter {
            rep.report_percent(percent, message).await
        } else {
            Ok(())
        }
    }

    /// Report count-based progress if available.
    pub async fn report_count(
        &self,
        current: usize,
        total: usize,
        message: Option<String>,
    ) -> crate::Result<()> {
        if let Some(rep) = &self.progress_reporter {
            rep.report_count(current, total, message).await
        } else {
            Ok(())
        }
    }
}

impl Default for RequestHandlerExtra {
    /// Create a default `RequestHandlerExtra` for testing and simple tool invocations.
    ///
    /// Uses a generated UUID as `request_id` and an uncancellable (never-cancelled)
    /// `CancellationToken`. Not suitable for production use where cancellation
    /// tracking is needed.
    fn default() -> Self {
        Self {
            cancellation_token: CancellationToken::new(),
            request_id: uuid::Uuid::new_v4().to_string(),
            session_id: None,
            auth_info: None,
            auth_context: None,
            metadata: HashMap::new(),
            progress_reporter: None,
            task_request: None,
            extensions: http::Extensions::new(),
            #[cfg(not(target_arch = "wasm32"))]
            peer: None,
        }
    }
}

impl std::fmt::Debug for RequestHandlerExtra {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // List of sensitive metadata keys that should be redacted
        const SENSITIVE_KEYS: &[&str] = &[
            "oauth_token",
            "access_token",
            "refresh_token",
            "api_key",
            "secret",
            "password",
            "bearer_token",
            "auth_token",
        ];

        // Create a redacted version of metadata
        let redacted_metadata: HashMap<String, String> = self
            .metadata
            .iter()
            .map(|(k, v)| {
                let is_sensitive = SENSITIVE_KEYS
                    .iter()
                    .any(|sensitive| k.to_lowercase().contains(sensitive));
                if is_sensitive {
                    (k.clone(), "[REDACTED]".to_string())
                } else {
                    (k.clone(), v.clone())
                }
            })
            .collect();

        let mut debug = f.debug_struct("RequestHandlerExtra");
        debug
            .field("cancellation_token", &self.cancellation_token)
            .field("request_id", &self.request_id)
            .field("session_id", &self.session_id)
            .field("auth_info", &self.auth_info)
            .field("auth_context", &self.auth_context)
            .field("metadata", &redacted_metadata)
            .field("task_request", &self.task_request.is_some())
            .field("extensions", &self.extensions);
        #[cfg(not(target_arch = "wasm32"))]
        debug.field("peer", &self.peer.as_ref().map(|_| "Arc<dyn PeerHandle>"));
        debug.finish()
    }
}

impl CancellationManager {
    /// Cancel a request silently (no notification sent to the client).
    pub async fn cancel_request_silent(&self, request_id: String) -> Result<()> {
        let token = {
            let mut tokens = self.tokens.write().await;
            tokens.remove(&request_id)
        };
        if let Some(token) = token {
            token.cancel();
        }
        Ok(())
    }
}

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

    #[tokio::test]
    async fn test_create_and_cancel_token() {
        let manager = CancellationManager::new();

        // Create a token
        let token = manager.create_token("test-request".to_string()).await;
        assert!(!token.is_cancelled());

        // Cancel the request
        manager
            .cancel_request("test-request".to_string(), None)
            .await
            .unwrap();

        // Token should be cancelled
        assert!(token.is_cancelled());

        // Token should be removed from manager
        assert!(manager.get_token("test-request").await.is_none());
    }

    #[tokio::test]
    async fn test_cancel_with_reason() {
        let manager = CancellationManager::new();

        // Set up notification tracking
        let notifications = Arc::new(RwLock::new(Vec::new()));
        let notifications_clone = notifications.clone();

        let mut manager = manager;
        manager.set_notification_sender(Arc::new(move |notif| {
            let notifications = notifications_clone.clone();
            tokio::spawn(async move {
                notifications.write().await.push(notif);
            });
        }));

        // Create and cancel with reason
        let _token = manager.create_token("test-request".to_string()).await;
        manager
            .cancel_request("test-request".to_string(), Some("Test reason".to_string()))
            .await
            .unwrap();

        // Give notification time to be sent
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        // Check notification was sent
        let notifs = notifications.read().await;
        assert_eq!(notifs.len(), 1);

        if let Notification::Client(crate::types::ClientNotification::Cancelled(cancelled)) =
            &notifs[0]
        {
            assert_eq!(
                cancelled.request_id,
                crate::types::RequestId::String("test-request".to_string())
            );
            assert_eq!(cancelled.reason, Some("Test reason".to_string()));
        } else {
            panic!("Expected Cancelled notification");
        }
    }

    #[tokio::test]
    async fn test_remove_token() {
        let manager = CancellationManager::new();

        // Create a token
        let token = manager.create_token("test-request".to_string()).await;
        assert!(manager.get_token("test-request").await.is_some());

        // Remove the token
        manager.remove_token("test-request").await;
        assert!(manager.get_token("test-request").await.is_none());

        // Token should still be valid (not cancelled)
        assert!(!token.is_cancelled());
    }

    #[tokio::test]
    async fn test_clear_all_tokens() {
        let manager = CancellationManager::new();

        // Create multiple tokens
        let token1 = manager.create_token("request1".to_string()).await;
        let token2 = manager.create_token("request2".to_string()).await;
        let token3 = manager.create_token("request3".to_string()).await;

        // Clear all tokens
        manager.clear().await;

        // All tokens should be cancelled
        assert!(token1.is_cancelled());
        assert!(token2.is_cancelled());
        assert!(token3.is_cancelled());

        // Manager should have no tokens
        assert!(manager.get_token("request1").await.is_none());
        assert!(manager.get_token("request2").await.is_none());
        assert!(manager.get_token("request3").await.is_none());
    }

    #[tokio::test]
    async fn test_request_handler_extra() {
        let token = CancellationToken::new();
        let extra = RequestHandlerExtra::new("test-req".to_string(), token.clone())
            .with_session_id(Some("session-123".to_string()));

        assert_eq!(extra.request_id, "test-req");
        assert_eq!(extra.session_id, Some("session-123".to_string()));
        assert!(!extra.is_cancelled());

        // Cancel the token
        token.cancel();
        assert!(extra.is_cancelled());
    }

    #[tokio::test]
    async fn test_metadata_redaction_in_debug() {
        let token = CancellationToken::new();
        let mut extra = RequestHandlerExtra::new("test-req".to_string(), token);

        // Add sensitive and non-sensitive metadata
        extra.set_metadata("oauth_token".to_string(), "secret-token-123".to_string());
        extra.set_metadata("access_token".to_string(), "bearer-xyz".to_string());
        extra.set_metadata("user_id".to_string(), "user-456".to_string());
        extra.set_metadata("request_count".to_string(), "42".to_string());

        // Get debug representation
        let debug_output = format!("{:?}", extra);

        // Verify sensitive values are redacted
        assert!(
            debug_output.contains("[REDACTED]"),
            "Expected redacted values in: {}",
            debug_output
        );
        assert!(
            !debug_output.contains("secret-token-123"),
            "OAuth token should be redacted: {}",
            debug_output
        );
        assert!(
            !debug_output.contains("bearer-xyz"),
            "Access token should be redacted: {}",
            debug_output
        );

        // Verify non-sensitive values are not redacted
        assert!(
            debug_output.contains("user-456"),
            "Non-sensitive metadata should not be redacted: {}",
            debug_output
        );
        assert!(
            debug_output.contains("42"),
            "Non-sensitive metadata should not be redacted: {}",
            debug_output
        );
    }

    #[tokio::test]
    async fn test_extensions_default_empty() {
        let extra = RequestHandlerExtra::default();
        assert!(extra.extensions().get::<String>().is_none());
    }

    #[tokio::test]
    async fn test_extensions_insert_overwrite_returns_old() {
        let mut extra = RequestHandlerExtra::default();
        assert_eq!(extra.extensions_mut().insert(42u64), None);
        assert_eq!(extra.extensions_mut().insert(99u64), Some(42u64));
        assert_eq!(extra.extensions().get::<u64>(), Some(&99u64));
    }

    #[tokio::test]
    async fn test_debug_extensions_prints_type_names_only() {
        let mut extra = RequestHandlerExtra::default();
        extra
            .extensions_mut()
            .insert("SECRET_VALUE_DO_NOT_LEAK".to_string());
        let debug_out = format!("{:?}", extra);
        // http::Extensions Debug prints type names, not field values
        assert!(!debug_out.contains("SECRET_VALUE_DO_NOT_LEAK"));
    }
}