mcp-streamable-proxy 0.2.4

//! Session Manager with backend version tracking
//!
//! This module implements ProxyAwareSessionManager that integrates with
//! ProxyHandler's version control mechanism to automatically invalidate
//! sessions when the backend reconnects.
//!
//! # Architecture
//!
//! ```text
//! ProxyAwareSessionManager
//! ├── LocalSessionManager (rmcp 提供的基础实现)
//! ├── ProxyHandler (Arc, 访问 backend_version)
//! └── DashMap<SessionId, SessionMetadata> (跟踪 session 创建时的版本)
//!
//! 工作流程：
//! 1. create_session: 记录当前 backend_version
//! 2. resume: 检查版本是否匹配
//!    - 匹配 → 正常 resume
//!    - 不匹配 → 返回 NotFound，客户端重新创建 session
//! ```

use dashmap::DashMap;
use dashmap::mapref::entry::Entry;
use futures::Stream;
use rmcp::{
    model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
    transport::{
        WorkerTransport,
        common::server_side_http::ServerSseMessage,
        streamable_http_server::session::{
            SessionId, SessionManager,
            local::{LocalSessionManager, LocalSessionManagerError, LocalSessionWorker},
        },
    },
};
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, info, warn};

use super::proxy_handler::ProxyHandler;

/// Session 元数据：跟踪 session 创建时的后端版本
#[derive(Debug, Clone)]
struct SessionMetadata {
    backend_version: u64,
}

/// 感知代理状态的 SessionManager
///
/// 职责：
/// 1. 委托 LocalSessionManager 处理核心 session 逻辑
/// 2. 维护 session → backend_version 映射
/// 3. 在 resume 时检查版本一致性
/// 4. 版本不匹配时使 session 失效
pub struct ProxyAwareSessionManager {
    inner: LocalSessionManager,
    handler: Arc<ProxyHandler>,
    session_versions: DashMap<String, SessionMetadata>,
    max_sessions: Option<usize>,
}

impl ProxyAwareSessionManager {
    /// 默认 session keep_alive 超时: 30 分钟
    /// rmcp 默认 5 分钟太短, 对于代理场景, agent 可能长时间不发消息但仍需要保持 session
    const DEFAULT_SESSION_KEEP_ALIVE_SECS: u64 = 30 * 60;

    pub fn new(handler: Arc<ProxyHandler>) -> Self {
        Self::with_keep_alive(
            handler,
            Duration::from_secs(Self::DEFAULT_SESSION_KEEP_ALIVE_SECS),
        )
    }

    pub fn with_keep_alive(handler: Arc<ProxyHandler>, keep_alive: Duration) -> Self {
        info!(
            "[Session Manager] Create ProxyAwareSessionManager - MCP ID: {}, keep_alive: {}s",
            handler.mcp_id(),
            keep_alive.as_secs()
        );
        let mut inner = LocalSessionManager::default();
        inner.session_config.keep_alive = Some(keep_alive);
        Self {
            inner,
            handler,
            session_versions: DashMap::new(),
            max_sessions: None,
        }
    }

    /// Limit the number of concurrent sessions (`None` = unlimited).
    pub fn with_max_sessions(mut self, max_sessions: Option<usize>) -> Self {
        self.max_sessions = max_sessions;
        self
    }

    fn check_backend_version(&self, session_id: &SessionId) -> bool {
        // Per-session handlers own their own backend; template version is not authoritative.
        if self.handler.isolation() == crate::backend_connector::BackendIsolation::PerSession {
            return true;
        }
        let current_version = self.handler.get_backend_version();
        match self.session_versions.entry(session_id.to_string()) {
            Entry::Occupied(entry) if entry.get().backend_version != current_version => {
                warn!(
                    "[Session version mismatch] session_id={}, creation version={}, current version={}, MCP ID: {}",
                    session_id,
                    entry.get().backend_version,
                    current_version,
                    self.handler.mcp_id()
                );
                false
            }
            Entry::Occupied(_) | Entry::Vacant(_) => true,
        }
    }

    fn backend_ready_for_session_ops(&self) -> bool {
        match self.handler.isolation() {
            crate::backend_connector::BackendIsolation::PerSession => true,
            crate::backend_connector::BackendIsolation::Shared => {
                self.handler.is_backend_available()
            }
        }
    }
}

// Implement SessionManager trait
impl SessionManager for ProxyAwareSessionManager {
    type Error = LocalSessionManagerError;
    type Transport = WorkerTransport<LocalSessionWorker>;

    async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> {
        if let Some(max) = self.max_sessions {
            let current = self.session_versions.len();
            if current >= max {
                warn!(
                    "[Session limit] rejecting create_session: {} >= max {}, MCP ID: {}",
                    current,
                    max,
                    self.handler.mcp_id()
                );
                return Err(LocalSessionManagerError::SessionNotFound(
                    format!("session-limit-{max}").into(),
                ));
            }
        }

        let (session_id, transport) = self.inner.create_session().await?;

        let version = self.handler.get_backend_version();
        self.session_versions.insert(
            session_id.to_string(),
            SessionMetadata {
                backend_version: version,
            },
        );

        info!(
            "[SessionCreated] session_id={}, backend_version={}, MCP ID: {}",
            session_id,
            version,
            self.handler.mcp_id()
        );

        Ok((session_id, transport))
    }

    async fn initialize_session(
        &self,
        id: &SessionId,
        message: ClientJsonRpcMessage,
    ) -> Result<ServerJsonRpcMessage, Self::Error> {
        if !self.backend_ready_for_session_ops() {
            warn!(
                "[Session initialization failed] session_id={}, reason: backend is unavailable, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        if !self.check_backend_version(id) {
            warn!(
                "[Session initialization failed] session_id={}, reason: version mismatch, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        debug!(
            "[Session initialization] session_id={}, MCP ID: {}",
            id,
            self.handler.mcp_id()
        );
        self.inner.initialize_session(id, message).await
    }

    async fn has_session(&self, id: &SessionId) -> Result<bool, Self::Error> {
        if !self.check_backend_version(id) {
            return Ok(false);
        }
        self.inner.has_session(id).await
    }

    async fn close_session(&self, id: &SessionId) -> Result<(), Self::Error> {
        info!(
            "[SessionClosed] session_id={}, MCP ID: {}",
            id,
            self.handler.mcp_id()
        );
        self.session_versions.remove(id.as_ref());
        // Shared isolation: drop peers whose transport is already closed.
        self.handler.upstream_peers().reap_closed_peers();
        self.inner.close_session(id).await
    }

    async fn create_stream(
        &self,
        id: &SessionId,
        message: ClientJsonRpcMessage,
    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
        if !self.backend_ready_for_session_ops() {
            warn!(
                "[Stream creation failed] session_id={}, reason: backend is unavailable, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        if !self.check_backend_version(id) {
            warn!(
                "[Stream creation failed] session_id={}, reason: version mismatch, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        debug!(
            "[Stream creation] session_id={}, MCP ID: {}",
            id,
            self.handler.mcp_id()
        );
        self.inner.create_stream(id, message).await
    }

    async fn accept_message(
        &self,
        id: &SessionId,
        message: ClientJsonRpcMessage,
    ) -> Result<(), Self::Error> {
        if !self.backend_ready_for_session_ops() {
            warn!(
                "[Message rejected] session_id={}, reason: backend unavailable, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        if !self.check_backend_version(id) {
            warn!(
                "[Message rejected] session_id={}, reason: version mismatch, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        self.inner.accept_message(id, message).await
    }

    async fn create_standalone_stream(
        &self,
        id: &SessionId,
    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
        self.inner.create_standalone_stream(id).await
    }

    async fn resume(
        &self,
        id: &SessionId,
        last_event_id: String,
    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
        let current_version = self.handler.get_backend_version();
        // This entry guard is dropped before close_session().await below.
        let stale_version =
            remove_stale_session_version(&self.session_versions, id.as_ref(), current_version);
        if let Some(previous_version) = stale_version {
            warn!(
                "[Session recovery failed] session_id={}, reason: backend version change ({} -> {}), MCP ID: {}",
                id,
                previous_version,
                current_version,
                self.handler.mcp_id()
            );
            let _ = self.inner.close_session(id).await;
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        if !self.backend_ready_for_session_ops() {
            warn!(
                "[Session recovery failed] session_id={}, reason: backend is unavailable, MCP ID: {}",
                id,
                self.handler.mcp_id()
            );
            return Err(LocalSessionManagerError::SessionNotFound(id.clone()));
        }

        debug!(
            "[SessionResumed] session_id={}, last_event_id={}, MCP ID: {}",
            id,
            last_event_id,
            self.handler.mcp_id()
        );
        self.inner.resume(id, last_event_id).await
    }
}

fn remove_stale_session_version(
    versions: &DashMap<String, SessionMetadata>,
    session_id: &str,
    current_version: u64,
) -> Option<u64> {
    match versions.entry(session_id.to_owned()) {
        Entry::Occupied(entry) if entry.get().backend_version != current_version => {
            let previous = entry.get().backend_version;
            entry.remove();
            Some(previous)
        }
        Entry::Occupied(_) | Entry::Vacant(_) => None,
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Barrier};

    use super::*;

    #[test]
    fn stale_version_removal_does_not_remove_replacement() {
        let versions = DashMap::new();
        versions.insert("session".to_owned(), SessionMetadata { backend_version: 1 });

        assert_eq!(
            remove_stale_session_version(&versions, "session", 2),
            Some(1)
        );
        versions.insert("session".to_owned(), SessionMetadata { backend_version: 2 });

        assert_eq!(remove_stale_session_version(&versions, "session", 2), None);
        assert_eq!(
            versions
                .entry("session".to_owned())
                .or_insert(SessionMetadata { backend_version: 0 })
                .backend_version,
            2
        );
    }

    #[test]
    fn concurrent_version_change_uses_atomic_entry_operations() {
        let versions = Arc::new(DashMap::new());
        versions.insert("session".to_owned(), SessionMetadata { backend_version: 1 });
        let stale_removed = Arc::new(Barrier::new(2));
        let replacement_inserted = Arc::new(Barrier::new(2));

        let worker_versions = versions.clone();
        let worker_stale_removed = stale_removed.clone();
        let worker_replacement_inserted = replacement_inserted.clone();
        let worker = std::thread::spawn(move || {
            let removed = remove_stale_session_version(&worker_versions, "session", 2);
            worker_stale_removed.wait();
            worker_replacement_inserted.wait();
            let second_removal = remove_stale_session_version(&worker_versions, "session", 2);
            (removed, second_removal)
        });

        stale_removed.wait();
        versions.insert("session".to_owned(), SessionMetadata { backend_version: 2 });
        replacement_inserted.wait();

        let (removed, second_removal) = worker.join().expect("worker thread must finish");
        assert_eq!(removed, Some(1));
        assert_eq!(second_removal, None);
        assert_eq!(
            versions
                .entry("session".to_owned())
                .or_insert(SessionMetadata { backend_version: 0 })
                .backend_version,
            2
        );
    }
}