Clean up MCP sessions on delete

Cargo.toml

@@ -41,7 +41,6 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tokio = { version = "1.48.0", features = ["full"] }
 axum = "0.8"
 tower-http = { version = "0.6.8", features = ["full"] }
-tower-layer = "0.3.3"
 tower = "0.5.3"
 http = "1.4.0"
 futures = { version = "0.3", features = ["std", "alloc"] }

crates/contextforge-gateway-rs-lib/Cargo.toml

@@ -18,7 +18,6 @@ tracing.workspace = true
 tokio.workspace = true
 axum.workspace = true
 tower-http.workspace = true
-tower-layer.workspace = true
 tower.workspace = true
 http.workspace = true
 futures.workspace = true

crates/contextforge-gateway-rs-lib/src/gateway/mcp_call_validator.rs

@@ -7,7 +7,10 @@ use rmcp::{
 };
 use tracing::info;
 
-use crate::layers::{session_id::SessionId, virtual_host_id::VirtualHostId};
+use crate::{
+    common::ContextForgeClaims,
+    layers::{session_id::SessionId, virtual_host_id::VirtualHostId},
+};
 
 pub struct AuthorizedCallValidator<'a> {
     call_name: &'a str,
@@ -18,10 +21,11 @@ impl<'a> AuthorizedCallValidator<'a> {
     pub fn new(call_name: &'a str, ctx: &'a RequestContext<RoleServer>) -> Self {
         Self { call_name, ctx }
     }
-    pub fn validate(self) -> Result<(&'a VirtualHost, &'a SessionId), ErrorData> {
+    pub fn validate(self) -> Result<(&'a VirtualHost, &'a SessionId, &'a ContextForgeClaims), ErrorData> {
         let maybe_parts = self.ctx.extensions.get::<Parts>();
         let maybe_session_id = maybe_parts.and_then(|parts| parts.extensions.get::<SessionId>());
         let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::<UserConfig>());
+        let maybe_claims = maybe_parts.and_then(|parts| parts.extensions.get::<ContextForgeClaims>());
 
         let maybe_virtual_host_id = maybe_parts.and_then(|parts| parts.extensions.get::<VirtualHostId>());
         info!(
@@ -61,7 +65,15 @@ impl<'a> AuthorizedCallValidator<'a> {
             });
         };
 
-        Ok((virtual_host, session_id))
+        let Some(claims) = maybe_claims else {
+            return Err(ErrorData {
+                code: ErrorCode::INTERNAL_ERROR,
+                message: "Routing problem... claims not found".into(),
+                data: None,
+            });
+        };
+
+        Ok((virtual_host, session_id, claims))
     }
 }
 
@@ -73,12 +85,13 @@ impl<'a> InitializeCallValidator<'a> {
     pub fn new(ctx: &'a RequestContext<RoleServer>) -> Self {
         Self { ctx }
     }
-    pub fn validate(self) -> Result<(&'a VirtualHost, &'a DownstreamSessionId), ErrorData> {
+    pub fn validate(self) -> Result<(&'a VirtualHost, &'a DownstreamSessionId, &'a ContextForgeClaims), ErrorData> {
         let maybe_parts = self.ctx.extensions.get::<Parts>();
 
         let maybe_downstream_session = self.ctx.extensions.get::<DownstreamSessionId>();
         let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::<UserConfig>());
         let maybe_virtual_host_id = maybe_parts.and_then(|parts| parts.extensions.get::<VirtualHostId>());
+        let maybe_claims = maybe_parts.and_then(|parts| parts.extensions.get::<ContextForgeClaims>());
         info!(
             "intialize user_config = {maybe_user_config:#?} downstream_session_id = {maybe_downstream_session:#?} virtual_host_id = {maybe_virtual_host_id:#?}"
         );
@@ -115,6 +128,14 @@ impl<'a> InitializeCallValidator<'a> {
             });
         };
 
-        Ok((virtual_host, downstream_session_id))
+        let Some(claims) = maybe_claims else {
+            return Err(ErrorData {
+                code: ErrorCode::INTERNAL_ERROR,
+                message: "Routing problem... claims not found".into(),
+                data: None,
+            });
+        };
+
+        Ok((virtual_host, downstream_session_id, claims))
     }
 }

crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs

@@ -35,6 +35,20 @@ use crate::{
     },
 };
 
+#[derive(Clone, Default)]
+pub struct BackendTransports(Arc<Mutex<HashMap<BackendTransportKey, BackendTransportService>>>);
+
+impl BackendTransports {
+    pub async fn remove_session(&self, principal: &str, session_id: &str) {
+        let mut transports = self.0.lock().await;
+        transports.retain(|key, _| key.principal != principal || key.session_id != session_id);
+    }
+
+    pub fn inner(&self) -> &Arc<Mutex<HashMap<BackendTransportKey, BackendTransportService>>> {
+        &self.0
+    }
+}
+
 #[derive(Clone, TypedBuilder)]
 #[builder(field_defaults(setter(prefix = "with_")))]
 pub struct McpService<T>
@@ -43,8 +57,8 @@ where
 {
     #[builder(default = Arc::new(Mutex::new(HashSet::new())))]
     subscriptions: Arc<Mutex<HashSet<String>>>,
-    #[builder(default = Arc::new(Mutex::new(HashMap::new())))]
-    transports: Arc<Mutex<HashMap<BackendTransportKey, BackendTransportService>>>,
+    #[builder(default = BackendTransports::default())]
+    transports: BackendTransports,
     #[builder(default = Arc::new(Mutex::new(LoggingLevel::Debug)))]
     log_level: Arc<Mutex<LoggingLevel>>,
     http_client: reqwest::Client,
@@ -55,6 +69,7 @@ where
 
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct BackendTransportKey {
+    principal: String,
     backend_name: String,
     session_id: String,
 }
@@ -80,15 +95,23 @@ pub struct BackendTransportService {
     pub(crate) service: Option<McpClientService>,
 }
 
-impl From<(&str, &str)> for BackendTransportKey {
-    fn from((backend_name, session_name): (&str, &str)) -> Self {
-        Self { backend_name: backend_name.to_owned(), session_id: session_name.to_owned() }
+impl From<(&str, &str, &str)> for BackendTransportKey {
+    fn from((backend_name, session_name, principal): (&str, &str, &str)) -> Self {
+        Self {
+            principal: principal.to_owned(),
+            backend_name: backend_name.to_owned(),
+            session_id: session_name.to_owned(),
+        }
     }
 }
 
-impl From<(&String, &SessionId)> for BackendTransportKey {
-    fn from((backend_name, session_name): (&String, &SessionId)) -> Self {
-        Self { backend_name: backend_name.to_owned(), session_id: session_name.value().to_owned() }
+impl From<(&String, &SessionId, &str)> for BackendTransportKey {
+    fn from((backend_name, session_name, principal): (&String, &SessionId, &str)) -> Self {
+        Self {
+            principal: principal.to_owned(),
+            backend_name: backend_name.to_owned(),
+            session_id: session_name.value().to_owned(),
+        }
     }
 }
 
@@ -108,10 +131,10 @@ where
         cx: RequestContext<RoleServer>,
     ) -> Result<InitializeResult, ErrorData> {
         let call_validator = InitializeCallValidator::new(&cx);
-        let (virtual_host, downstream_session_id) = call_validator.validate()?;
+        let (virtual_host, downstream_session_id, claims) = call_validator.validate()?;
         let session_mapping = if let Ok(maybe_session_mapping) = self
             .user_session_store
-            .get_session(&UserSession::new(String::new(), Arc::clone(&downstream_session_id.session_id)))
+            .get_session(&UserSession::new(claims.sub.clone(), Arc::clone(&downstream_session_id.session_id)))
             .await
         {
             maybe_session_mapping.unwrap_or_default()
@@ -184,18 +207,26 @@ where
             })
             .unzip();
 
-        let _ = self
+        if self
             .user_session_store
             .set_session(
-                &UserSession::new(String::new(), Arc::clone(&downstream_session_id.session_id)),
+                &UserSession::new(claims.sub.clone(), Arc::clone(&downstream_session_id.session_id)),
                 &session_mapping,
             )
-            .await;
+            .await
+            .is_err()
+        {
+            return Err(ErrorData {
+                code: ErrorCode::INTERNAL_ERROR,
+                message: "Internal problem... session store can't be written".into(),
+                data: None,
+            });
+        }
 
-        let mut transports = self.transports.lock().await;
+        let mut transports = self.transports.inner().lock().await;
         for (name, svc) in backend_services {
             transports
-                .entry(BackendTransportKey::from((name.as_str(), downstream_session_id.value())))
+                .entry(BackendTransportKey::from((name.as_str(), downstream_session_id.value(), claims.sub.as_str())))
                 .insert_entry(svc);
         }
         drop(transports);
@@ -215,9 +246,9 @@ where
         cx: RequestContext<RoleServer>,
     ) -> Result<ListToolsResult, ErrorData> {
         let mcp_call_validator = AuthorizedCallValidator::new("list_tools", &cx);
-        let (virtual_host, session_id) = mcp_call_validator.validate()?;
+        let (virtual_host, session_id, claims) = mcp_call_validator.validate()?;
 
-        let session_manager = SessionManager::new(virtual_host, session_id, &self.transports);
+        let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &self.transports);
         let backend_transports: Vec<_> = session_manager.borrow_transports().await;
 
         let list_tools_tasks = backend_transports
@@ -264,8 +295,8 @@ where
         cx: RequestContext<RoleServer>,
     ) -> Result<CallToolResult, ErrorData> {
         let mcp_call_validator = AuthorizedCallValidator::new("call_tool", &cx);
-        let (virtual_host, session_id) = mcp_call_validator.validate()?;
-        let session_manager = SessionManager::new(virtual_host, session_id, &self.transports);
+        let (virtual_host, session_id, claims) = mcp_call_validator.validate()?;
+        let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &self.transports);
 
         let backend_names = session_manager.get_backend_names();
 
@@ -355,9 +386,9 @@ where
         cx: RequestContext<RoleServer>,
     ) -> Result<ListResourcesResult, ErrorData> {
         let mcp_call_validator = AuthorizedCallValidator::new("list_resources", &cx);
-        let (virtual_host, session_id) = mcp_call_validator.validate()?;
+        let (virtual_host, session_id, claims) = mcp_call_validator.validate()?;
 
-        let session_manager = SessionManager::new(virtual_host, session_id, &self.transports);
+        let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &self.transports);
         let backend_transports: Vec<_> = session_manager.borrow_transports().await;
 
         let list_resources_tasks = backend_transports
@@ -404,8 +435,8 @@ where
         cx: RequestContext<RoleServer>,
     ) -> Result<ReadResourceResult, ErrorData> {
         let mcp_call_validator = AuthorizedCallValidator::new("read_resource", &cx);
-        let (virtual_host, session_id) = mcp_call_validator.validate()?;
-        let session_manager = SessionManager::new(virtual_host, session_id, &self.transports);
+        let (virtual_host, session_id, claims) = mcp_call_validator.validate()?;
+        let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &self.transports);
 
         let backend_names = session_manager.get_backend_names();
 

crates/contextforge-gateway-rs-lib/src/gateway/mod.rs

@@ -3,4 +3,5 @@ pub(crate) mod mcp_gateway;
 mod session_manager;
 mod session_store;
 
-pub use mcp_gateway::{LocalUserSessionStore, McpService};
+pub use mcp_gateway::{BackendTransports, LocalUserSessionStore, McpService};
+pub use session_store::{UserSession, UserSessionStore};

crates/contextforge-gateway-rs-lib/src/gateway/session_manager.rs

@@ -1,25 +1,24 @@
-use std::{collections::HashMap, sync::Arc};
-
 use contextforge_gateway_rs_apis::user_store::VirtualHost;
-use tokio::sync::Mutex;
 use tracing::{debug, info};
 
-use super::mcp_gateway::{BackendTransportKey, BackendTransportService, ServiceHolder};
+use super::mcp_gateway::{BackendTransportKey, BackendTransports, ServiceHolder};
 use crate::layers::session_id::SessionId;
 
 pub struct SessionManager<'a> {
     virtual_host: &'a VirtualHost,
     session_id: &'a SessionId,
-    transports: &'a Arc<Mutex<HashMap<BackendTransportKey, BackendTransportService>>>,
+    principal: &'a str,
+    transports: &'a BackendTransports,
 }
 
 impl<'a> SessionManager<'a> {
     pub fn new(
         virtual_host: &'a VirtualHost,
         session_id: &'a SessionId,
-        transports: &'a Arc<Mutex<HashMap<BackendTransportKey, BackendTransportService>>>,
+        principal: &'a str,
+        transports: &'a BackendTransports,
     ) -> Self {
-        Self { virtual_host, session_id, transports }
+        Self { virtual_host, session_id, principal, transports }
     }
 
     pub fn get_backend_names(&self) -> Vec<&str> {
@@ -28,12 +27,12 @@ impl<'a> SessionManager<'a> {
 
     pub async fn borrow_transports(&self) -> Vec<ServiceHolder> {
         let names: Vec<_> = self.virtual_host.backends.keys().cloned().collect();
-        let mut transports = self.transports.lock().await;
+        let mut transports = self.transports.inner().lock().await;
         names
             .into_iter()
             .filter_map(|name| {
                 transports
-                    .get_mut(&BackendTransportKey::from((&name, self.session_id)))
+                    .get_mut(&BackendTransportKey::from((&name, self.session_id, self.principal)))
                     .map(|b| ServiceHolder::new(name, b.service.clone()))
             })
             .collect()
@@ -42,20 +41,20 @@ impl<'a> SessionManager<'a> {
     // pub async fn return_transports(&self, backend_transports: impl Iterator<Item = ServiceHolder>) {
     //     let backend_transports = backend_transports.collect::<Vec<_>>();
     //     info!("Returning transports {:?} {backend_transports:?}", self.session_id);
-    //     let mut transports = self.transports.lock().await;
+    //     let mut transports = self.transports.inner().lock().await;
     //     for svc_holder in backend_transports {
     //         transports
-    //             .entry(BackendTransportKey::from((&svc_holder.name, self.session_id)))
+    //             .entry(BackendTransportKey::from((&svc_holder.name, self.session_id, self.principal)))
     //             .and_modify(|e| e.service = svc_holder.running_service);
     //     }
     // }
 
     pub async fn cleanup_backends(&self, reason: &'static str) {
         let names: Vec<_> = self.virtual_host.backends.keys().cloned().collect();
         info!("Cleaning up backends {:?}", self.session_id);
-        let mut transports = self.transports.lock().await;
+        let mut transports = self.transports.inner().lock().await;
         for name in names {
-            let key = BackendTransportKey::from((&name, self.session_id));
+            let key = BackendTransportKey::from((&name, self.session_id, self.principal));
             debug!("session_manager: removing transport for {key:?} {reason}");
             transports.remove(&key);
         }

crates/contextforge-gateway-rs-lib/src/gateway/session_store/local_session_store.rs

@@ -40,4 +40,9 @@ impl UserSessionStore for LocalUserSessionStore {
         self.cache.lock().await.insert(session_key.clone(), mapping.clone());
         Ok(())
     }
+
+    async fn remove_session<'a>(&self, session_key: &'a UserSession) -> Result<(), SessionStoreError> {
+        self.cache.lock().await.remove(session_key);
+        Ok(())
+    }
 }

crates/contextforge-gateway-rs-lib/src/gateway/session_store/mod.rs

@@ -77,4 +77,5 @@ pub trait UserSessionStore: Send + Sync {
         key: &'a UserSession,
         session_mapping: &'a SessionMapping,
     ) -> Result<(), SessionStoreError>;
+    async fn remove_session<'a>(&self, key: &'a UserSession) -> Result<(), SessionStoreError>;
 }

crates/contextforge-gateway-rs-lib/src/gateway/session_store/redis_session_store.rs

@@ -94,4 +94,18 @@ impl UserSessionStore for RedisUserSessionStore {
             return Err(SessionStoreError::CantWriteData);
         }
     }
+
+    async fn remove_session<'a>(&self, session_key: &'a UserSession) -> Result<(), SessionStoreError> {
+        let Ok(key) = rmp_serde::encode::to_vec::<UserSession>(session_key) else {
+            return Err(SessionStoreError::DataEncoding);
+        };
+
+        let mut connection = self.connection.clone();
+        if redis::cmd("DEL").arg(key).query_async::<()>(&mut connection).await.is_ok() {
+            self.cache.lock().await.remove(session_key);
+            Ok(())
+        } else {
+            Err(SessionStoreError::CantWriteData)
+        }
+    }
 }