mcp_streamable_proxy/

session_manager.rs

//! 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 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;
        }
        if let Some(meta) = self.session_versions.get(session_id.as_ref()) {
            let current_version = self.handler.get_backend_version();
            if meta.backend_version != current_version {
                warn!(
                    "[Session version mismatch] session_id={}, creation version={}, current version={}, MCP ID: {}",
                    session_id,
                    meta.backend_version,
                    current_version,
                    self.handler.mcp_id()
                );
                return false;
            }
        }
        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> {
        // 关键：检查后端版本
        if let Some(meta) = self.session_versions.get(id.as_ref()) {
            let current_version = self.handler.get_backend_version();
            if meta.backend_version != current_version {
                warn!(
                    "[Session recovery failed] session_id={}, reason: backend version change ({} -> {}), MCP ID: {}",
                    id,
                    meta.backend_version,
                    current_version,
                    self.handler.mcp_id()
                );

                // 清理失效 session
                drop(meta); // 释放 DashMap 的读锁
                self.session_versions.remove(id.as_ref());
                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
    }
}