systemprompt-mcp 0.32.2

//! Database-backed MCP session manager.
//!
//! [`DatabaseSessionHandler`] implements the rmcp `SessionManager` trait (see
//! the `session_manager_impl` submodule), wrapping rmcp's in-memory
//! `LocalSessionManager` while mirroring session lifecycle (create, activity,
//! close) into the `mcp_sessions` table for cross-restart visibility.
//! [`DatabaseSessionManagerError`] models the local, database, and
//! reconnect-signalling failure cases; database persistence is best-effort and
//! never fails an in-memory operation.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

mod session_manager_impl;
mod session_store;

pub use session_store::PostgresSessionStore;

use std::fmt;
use std::sync::Arc;

use crate::repository::McpSessionRepository;
use rmcp::transport::streamable_http_server::session::SessionId;
use rmcp::transport::streamable_http_server::session::local::{
    LocalSessionManager, LocalSessionManagerError,
};

#[derive(Debug)]
pub enum DatabaseSessionManagerError {
    Local(LocalSessionManagerError),
    Database(crate::error::McpDomainError),
    SessionNotFound(String),
    SessionExpired(String),
    SessionNeedsReconnect(String),
}

impl fmt::Display for DatabaseSessionManagerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Local(e) => write!(f, "Local session error: {e}"),
            Self::Database(e) => write!(f, "Database error: {e}"),
            Self::SessionNotFound(id) => write!(f, "Session not found: {id}"),
            Self::SessionExpired(id) => write!(f, "Session expired: {id}"),
            Self::SessionNeedsReconnect(id) => write!(f, "Session needs reconnect: {id}"),
        }
    }
}

impl std::error::Error for DatabaseSessionManagerError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Local(e) => Some(e),
            Self::Database(e) => Some(e),
            _ => None,
        }
    }
}

impl From<LocalSessionManagerError> for DatabaseSessionManagerError {
    fn from(e: LocalSessionManagerError) -> Self {
        Self::Local(e)
    }
}

pub struct DatabaseSessionHandler {
    local_manager: LocalSessionManager,
    repository: Arc<McpSessionRepository>,
}

impl fmt::Debug for DatabaseSessionHandler {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("DatabaseSessionHandler")
            .field("local_manager", &self.local_manager)
            .field("repository", &self.repository)
            .finish()
    }
}

impl DatabaseSessionHandler {
    pub fn new(repository: Arc<McpSessionRepository>) -> Self {
        Self::with_timeouts(repository, crate::SessionTimeouts::default())
    }

    pub fn with_timeouts(
        repository: Arc<McpSessionRepository>,
        timeouts: crate::SessionTimeouts,
    ) -> Self {
        let mut local_manager = LocalSessionManager::default();
        let cfg = &mut local_manager.session_config;
        cfg.init_timeout = timeouts.init.or(cfg.init_timeout);
        cfg.keep_alive = timeouts.keep_alive.or(cfg.keep_alive);
        Self {
            local_manager,
            repository,
        }
    }

    async fn persist_create(&self, session_id: &SessionId) {
        if let Err(e) = self
            .repository
            .create(
                &systemprompt_identifiers::SessionId::new(session_id.as_ref()),
                None,
                None,
            )
            .await
        {
            tracing::warn!(
                session_id = %session_id,
                error = %e,
                "Failed to persist session creation to database"
            );
        }
    }

    async fn persist_close(&self, session_id: &SessionId) {
        if let Err(e) = self
            .repository
            .close(&systemprompt_identifiers::SessionId::new(
                session_id.as_ref(),
            ))
            .await
        {
            tracing::warn!(
                session_id = %session_id,
                error = %e,
                "Failed to persist session close to database"
            );
        }
    }

    pub(crate) async fn update_activity(&self, session_id: &SessionId) {
        if let Err(e) = self
            .repository
            .update_activity(&systemprompt_identifiers::SessionId::new(
                session_id.as_ref(),
            ))
            .await
        {
            tracing::debug!(
                session_id = %session_id,
                error = %e,
                "Failed to update session activity"
            );
        }
    }

    async fn check_db_session(&self, session_id: &SessionId) -> Option<bool> {
        match self
            .repository
            .find_active(&systemprompt_identifiers::SessionId::new(
                session_id.as_ref(),
            ))
            .await
        {
            Ok(Some(_)) => Some(true),
            Ok(None) => Some(false),
            Err(e) => {
                tracing::warn!(
                    session_id = %session_id,
                    error = %e,
                    "Failed to check session in database"
                );
                None
            },
        }
    }
}