neva 0.5.2

//! Tools and utilities for MCP Client Session

use crate::transport::http::ServiceUrl;
use once_cell::sync::OnceCell;
use std::sync::Arc;
use std::sync::RwLock;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;

/// Represents current MCP Session
pub(super) struct McpSession {
    /// Dual-mode protocol switch -- legacy peers get legacy headers.
    #[cfg(not(feature = "legacy-spec"))]
    peer_mode: crate::shared::PeerMode,
    initialized: Notify,
    sse_ready: Notify,
    url: Arc<str>,
    session_id: OnceCell<uuid::Uuid>,
    last_event_id: RwLock<Option<String>>,

    /// The reconnection delay the server last asked for with an SSE `retry:`
    /// field, in milliseconds; [`UNSTATED_RETRY`] while it has asked for none.
    ///
    /// The field is the server's, not a suggestion: a client that reconnects on
    /// its own schedule either hammers a server that asked for room or leaves a
    /// stream down long after it said to come back.
    retry_ms: std::sync::atomic::AtomicU64,
    cancellation_token: CancellationToken,

    /// Abort handles for in-flight requests whose reply is a long-lived stream
    /// (`subscriptions/listen`).
    ///
    /// A subscription is ended over HTTP by closing its response body -- that
    /// is the spec's cancellation mechanism on this transport, and the only one
    /// that works, since a `notifications/cancelled` travels on its own `POST`
    /// and carries no evidence of which stream it belongs to. Cancelling the
    /// token here drops the body without disturbing the rest of the session.
    #[cfg(not(feature = "legacy-spec"))]
    streams: dashmap::DashMap<crate::types::RequestId, CancellationToken>,
}

/// The stored `retry:` when the server has not stated one.
///
/// A sentinel rather than `0`, because `0` is a legal delay the server may
/// actually ask for; no server states a reconnection pause of 2^64-1 ms.
const UNSTATED_RETRY: u64 = u64::MAX;

impl McpSession {
    /// Creates a new [`McpSession`].
    ///
    /// The request URL is assembled once here (after all client configuration,
    /// including TLS, has settled in [`ServiceUrl`]) and cached as an
    /// [`Arc<str>`], so the per-request POST/GET paths borrow it directly
    /// instead of re-formatting the URL on every call.
    pub(super) fn new(
        url: ServiceUrl,
        token: CancellationToken,
        #[cfg(not(feature = "legacy-spec"))] peer_mode: crate::shared::PeerMode,
    ) -> Self {
        Self {
            #[cfg(not(feature = "legacy-spec"))]
            peer_mode,
            initialized: Notify::new(),
            sse_ready: Notify::new(),
            session_id: OnceCell::new(),
            last_event_id: RwLock::new(None),
            retry_ms: std::sync::atomic::AtomicU64::new(UNSTATED_RETRY),
            cancellation_token: token,
            url: Arc::from(url.to_url()),
            #[cfg(not(feature = "legacy-spec"))]
            streams: dashmap::DashMap::new(),
        }
    }

    /// Registers an abort handle for the in-flight request `id` and returns it.
    #[cfg(not(feature = "legacy-spec"))]
    pub(super) fn track_stream(&self, id: crate::types::RequestId) -> CancellationToken {
        let token = CancellationToken::new();
        self.streams.insert(id, token.clone());
        token
    }

    /// Drops the abort handle for `id`, once its reply has been read.
    #[cfg(not(feature = "legacy-spec"))]
    pub(super) fn untrack_stream(&self, id: &crate::types::RequestId) {
        self.streams.remove(id);
    }

    /// Aborts the in-flight request `id`, closing its response body.
    ///
    /// Returns whether there was one to abort.
    #[cfg(not(feature = "legacy-spec"))]
    pub(super) fn abort_stream(&self, id: &crate::types::RequestId) -> bool {
        match self.streams.remove(id) {
            Some((_, token)) => {
                token.cancel();
                true
            }
            None => false,
        }
    }

    /// Whether the connected peer negotiated the legacy (legacy)
    /// protocol via the dual-mode fallback.
    #[cfg(not(feature = "legacy-spec"))]
    pub(super) fn is_legacy(&self) -> bool {
        self.peer_mode.is_legacy()
    }

    /// Returns the pre-assembled request URL for this session.
    pub(super) fn url(&self) -> &str {
        &self.url
    }

    /// Returns the [`CancellationToken`] that can abort the whole session
    pub(super) fn cancellation_token(&self) -> CancellationToken {
        self.cancellation_token.clone()
    }

    /// Returns `true` if a Session ID has been specified
    pub(super) fn has_session_id(&self) -> bool {
        self.session_id.get().is_some()
    }

    /// Returns a reference to the current MCP Session ID
    pub(super) fn session_id(&self) -> Option<&uuid::Uuid> {
        self.session_id.get()
    }

    /// Sets the MCP Session ID
    pub(super) fn set_session_id(&self, id: uuid::Uuid) {
        if let Err(_err) = self.session_id.set(id) {
            #[cfg(feature = "tracing")]
            tracing::info!("MCP Session Id already set");
        }
    }

    /// Returns the last received SSE event ID, if any
    pub(super) fn last_event_id(&self) -> Option<String> {
        self.last_event_id.read().ok().and_then(|g| g.clone())
    }

    /// Updates the last received SSE event ID
    pub(super) fn set_last_event_id(&self, id: String) {
        if let Ok(mut guard) = self.last_event_id.write() {
            *guard = Some(id);
        }
    }

    /// Records an SSE `retry:` field as the reconnection delay to use from now
    /// on.
    ///
    /// `0` is a delay like any other -- a server asking to be reconnected
    /// immediately -- so it is stored rather than dropped. "Nothing was asked
    /// for" is a state of its own, and [`UNSTATED_RETRY`] is what says it;
    /// using `0` for both would silently turn a request for an immediate
    /// reconnect into the multi-second default.
    pub(super) fn set_retry(&self, ms: u64) {
        self.retry_ms
            .store(ms, std::sync::atomic::Ordering::Relaxed);
    }

    /// How long to wait before reconnecting a dropped stream: what the server
    /// last asked for, or `default` while it has asked for nothing.
    pub(super) fn retry_delay(&self, default: std::time::Duration) -> std::time::Duration {
        match self.retry_ms.load(std::sync::atomic::Ordering::Relaxed) {
            UNSTATED_RETRY => default,
            ms => std::time::Duration::from_millis(ms),
        }
    }

    /// Sends a signal that this MCP Session has been initialized
    #[inline]
    pub(super) fn notify_session_initialized(&self) {
        self.initialized.notify_one();
    }

    /// Sends a signal that the SSE-connection has been initialized
    #[inline]
    pub(super) fn notify_sse_initialized(&self) {
        self.sse_ready.notify_one();
    }

    /// Waits for MCP Session to be initialized
    #[inline]
    pub(super) async fn initialized(&self) {
        self.initialized.notified().await;
    }

    /// Waits for SSE connection to be initialized
    #[inline]
    pub(super) async fn sse_ready(&self) {
        self.sse_ready.notified().await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transport::http::HttpProto;
    use std::sync::Arc;
    use tokio::time::{Duration, timeout};
    use tokio_util::sync::CancellationToken;
    use uuid::Uuid;

    fn create_session() -> McpSession {
        let url = ServiceUrl {
            proto: HttpProto::Http,
            addr: "localhost".to_string(),
            endpoint: "init".to_string(),
        };
        let token = CancellationToken::new();
        McpSession::new(
            url,
            token,
            #[cfg(not(feature = "legacy-spec"))]
            Default::default(),
        )
    }

    #[tokio::test]
    async fn it_has_url() {
        let session = create_session();
        // The full URL is assembled once and cached: proto://addr + endpoint.
        assert_eq!(session.url(), "http://localhostinit");
    }

    #[tokio::test]
    async fn it_has_cancellable_and_synced_cancellation_token() {
        let session = create_session();
        let token = session.cancellation_token();
        token.cancel();
        assert!(token.is_cancelled());
    }

    #[tokio::test]
    async fn it_sets_and_gets_session_id() {
        let session = create_session();
        let id = Uuid::new_v4();
        assert!(!session.has_session_id());
        assert!(session.session_id().is_none());

        session.set_session_id(id);
        assert!(session.has_session_id());
        assert_eq!(session.session_id(), Some(&id));
    }

    #[test]
    fn it_returns_none_last_event_id_by_default() {
        let session = create_session();
        assert!(session.last_event_id().is_none());
    }

    #[test]
    fn it_sets_and_gets_last_event_id() {
        let session = create_session();
        session.set_last_event_id("abc-123".to_string());
        assert_eq!(session.last_event_id(), Some("abc-123".to_string()));
    }

    #[test]
    fn it_overwrites_last_event_id_on_each_set() {
        let session = create_session();
        session.set_last_event_id("first".to_string());
        session.set_last_event_id("second".to_string());
        assert_eq!(session.last_event_id(), Some("second".to_string()));
    }

    #[test]
    fn the_reconnect_delay_is_the_servers_to_state() {
        let default = std::time::Duration::from_secs(3);
        let session = create_session();
        assert_eq!(
            session.retry_delay(default),
            default,
            "a server that stated nothing gets the default"
        );

        session.set_retry(500);
        assert_eq!(
            session.retry_delay(default),
            std::time::Duration::from_millis(500)
        );

        // Zero is a delay the server is entitled to ask for -- come back at
        // once -- and it is not the same statement as having asked for nothing.
        session.set_retry(0);
        assert_eq!(
            session.retry_delay(default),
            std::time::Duration::ZERO,
            "a server asking for an immediate reconnect must get one"
        );

        session.set_retry(1200);
        assert_eq!(
            session.retry_delay(default),
            std::time::Duration::from_millis(1200),
            "the latest statement wins"
        );
    }

    #[tokio::test]
    async fn it_guarantees_session_id_cannot_be_overwritten() {
        let session = create_session();
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();

        session.set_session_id(id1);
        session.set_session_id(id2); // silently ignored

        assert_eq!(session.session_id(), Some(&id1));
        assert_ne!(session.session_id(), Some(&id2));
    }

    #[tokio::test]
    async fn it_notifies_and_initialized() {
        let session = Arc::new(create_session());

        let handle = tokio::spawn({
            let session = session.clone();
            async move {
                session.initialized().await;
            }
        });

        // Notify after a short_set_and_get delay
        tokio::time::sleep(Duration::from_millis(10)).await;
        session.notify_session_initialized();

        // Should complete within timeout
        assert!(timeout(Duration::from_secs(1), handle).await.is_ok());
    }
}