fix: isolate downstream subscription registrations

Signed-off-by: Pratik Gandhi <gandhipratik203@gmail.com>

Author: gandhipratik203

crates/contextforge-data-plane-lib/src/gateway/downstream_subscriptions.rs

@@ -1,6 +1,9 @@
 use std::{
     collections::HashMap,
-    sync::{Arc, Mutex},
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicU64, Ordering},
+    },
 };
 
 use rmcp::{
@@ -13,6 +16,7 @@ use crate::layers::request_context::GatewayRequestContext;
 #[derive(Clone, Default)]
 pub(crate) struct DownstreamSubscriptionRegistry {
     inner: Arc<Mutex<HashMap<DownstreamSubscriptionKey, SubscriptionSink>>>,
+    next_registration_id: Arc<AtomicU64>,
 }
 
 impl DownstreamSubscriptionRegistry {
@@ -22,7 +26,8 @@ impl DownstreamSubscriptionRegistry {
         filter: &SubscriptionFilter,
         sink: &SubscriptionSink,
     ) -> DownstreamSubscriptionGuard {
-        let keys = subscription_keys(context, filter, sink.id());
+        let registration_id = self.next_registration_id.fetch_add(1, Ordering::Relaxed);
+        let keys = subscription_keys(context, filter, sink.id(), registration_id);
         let mut subscriptions = self.inner.lock().expect("downstream subscription registry lock poisoned");
         for key in &keys {
             subscriptions.insert(key.clone(), sink.clone());
@@ -54,6 +59,7 @@ pub(crate) struct DownstreamSubscriptionKey {
     principal: String,
     virtual_host_id: String,
     subscription_id: RequestId,
+    registration_id: u64,
     notification: DownstreamSubscriptionNotification,
 }
 
@@ -69,22 +75,39 @@ pub(super) fn subscription_keys(
     context: &GatewayRequestContext,
     filter: &SubscriptionFilter,
     subscription_id: &RequestId,
+    registration_id: u64,
 ) -> Vec<DownstreamSubscriptionKey> {
     let mut keys = Vec::new();
     if filter.tools_list_changed == Some(true) {
-        keys.push(subscription_key(context, subscription_id, DownstreamSubscriptionNotification::ToolsListChanged));
+        keys.push(subscription_key(
+            context,
+            subscription_id,
+            registration_id,
+            DownstreamSubscriptionNotification::ToolsListChanged,
+        ));
     }
     if filter.prompts_list_changed == Some(true) {
-        keys.push(subscription_key(context, subscription_id, DownstreamSubscriptionNotification::PromptsListChanged));
+        keys.push(subscription_key(
+            context,
+            subscription_id,
+            registration_id,
+            DownstreamSubscriptionNotification::PromptsListChanged,
+        ));
     }
     if filter.resources_list_changed == Some(true) {
-        keys.push(subscription_key(context, subscription_id, DownstreamSubscriptionNotification::ResourcesListChanged));
+        keys.push(subscription_key(
+            context,
+            subscription_id,
+            registration_id,
+            DownstreamSubscriptionNotification::ResourcesListChanged,
+        ));
     }
     if let Some(uris) = &filter.resource_subscriptions {
         keys.extend(uris.iter().map(|uri| {
             subscription_key(
                 context,
                 subscription_id,
+                registration_id,
                 DownstreamSubscriptionNotification::ResourceUpdated { uri: uri.clone() },
             )
         }));
@@ -95,12 +118,91 @@ pub(super) fn subscription_keys(
 fn subscription_key(
     context: &GatewayRequestContext,
     subscription_id: &RequestId,
+    registration_id: u64,
     notification: DownstreamSubscriptionNotification,
 ) -> DownstreamSubscriptionKey {
     DownstreamSubscriptionKey {
         principal: context.principal().to_owned(),
         virtual_host_id: context.virtual_host_id().to_owned(),
         subscription_id: subscription_id.clone(),
+        registration_id,
         notification,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use contextforge_data_plane_apis::user_store::{BackendMCPGateway, Transport, VirtualHost};
+    use rmcp::model::RequestId;
+
+    use super::*;
+
+    #[test]
+    fn keys_include_subscription_id_and_notification_kind() {
+        let gateway_context = GatewayRequestContext::new(&test_claims(), &test_virtual_host_id(), &test_virtual_host());
+        let filter = SubscriptionFilter::builder().tools_list_changed().resource_subscription("memo://known").build();
+
+        let keys = subscription_keys(&gateway_context, &filter, &RequestId::Number(7), 9);
+
+        assert_eq!(2, keys.len());
+        assert!(keys.iter().all(|key| key.subscription_id == RequestId::Number(7)));
+        assert!(keys.iter().all(|key| key.registration_id == 9));
+    }
+
+    #[test]
+    fn registration_id_is_part_of_key_identity() {
+        let gateway_context = GatewayRequestContext::new(&test_claims(), &test_virtual_host_id(), &test_virtual_host());
+        let filter = SubscriptionFilter::builder().tools_list_changed().build();
+
+        let first = subscription_keys(&gateway_context, &filter, &RequestId::Number(7), 0);
+        let second = subscription_keys(&gateway_context, &filter, &RequestId::Number(7), 1);
+
+        assert_ne!(first, second);
+    }
+
+    fn test_virtual_host() -> VirtualHost {
+        VirtualHost {
+            backends: HashMap::from([(
+                "backend-one".to_owned(),
+                BackendMCPGateway {
+                    name: "backend-one".to_owned(),
+                    url: "http://127.0.0.1:9999/mcp".parse().expect("valid URL"),
+                    transport: Transport::default(),
+                    passthrough_headers: Vec::new(),
+                    add_headers: HashMap::new(),
+                    remove_headers: Vec::new(),
+                    allowed_tool_names: Vec::new(),
+                    tool_name_aliases: HashMap::new(),
+                    allowed_resource_names: Vec::new(),
+                    allowed_prompt_names: Vec::new(),
+                },
+            )]),
+        }
+    }
+
+    fn test_claims() -> crate::common::ContextForgeClaims {
+        crate::common::ContextForgeClaims {
+            sub: "test-principal".to_owned(),
+            jti: "test-jti".to_owned(),
+            token_use: None,
+            iat: None,
+            iss: "test-issuer".to_owned(),
+            aud: "test-audience".to_owned(),
+            exp: 1,
+            teams: None,
+            user: crate::common::User::builder()
+                .email("test@example.com".to_owned())
+                .full_name(None)
+                .is_admin(false)
+                .auth_provider("test".to_owned())
+                .build(),
+            scopes: None,
+        }
+    }
+
+    fn test_virtual_host_id() -> crate::layers::virtual_host_id::VirtualHostId {
+        crate::layers::virtual_host_id::VirtualHostId::new("test-vhost".to_owned())
+    }
+}

crates/contextforge-data-plane-lib/src/gateway/mcp_service.rs

@@ -376,8 +376,12 @@ mod tests {
         let virtual_host = virtual_host_with_backends(&["backend-one"]);
         let gateway_context = GatewayRequestContext::new(&test_claims(), &test_virtual_host_id(), &virtual_host);
         let filter = SubscriptionFilter::builder().tools_list_changed().resource_subscription("memo://known").build();
-        let keys =
-            super::super::downstream_subscriptions::subscription_keys(&gateway_context, &filter, &RequestId::Number(7));
+        let keys = super::super::downstream_subscriptions::subscription_keys(
+            &gateway_context,
+            &filter,
+            &RequestId::Number(7),
+            0,
+        );
 
         assert_eq!(2, keys.len());
     }

crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs

@@ -31,26 +31,14 @@ pub async fn virtual_host_config_layer(request: http::Request<axum::body::Body>,
             let gateway_context = GatewayRequestContext::new(claims, virtual_host_id, virtual_host);
             return scope_gateway_request_context(gateway_context, next.run(request)).await;
         }
-    }
 
-    if let (Some(virtual_host_id), Some(user_config)) = (virtual_host_id, user_config)
-        && !has_virtual_host(user_config, virtual_host_id)
-    {
         let virtual_host_id = virtual_host_id.value();
-        let virtual_hosts = user_config.virtual_hosts.len();
-        debug!(
-            "virtual_host_config_layer - virtual host config missing virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}"
-        );
-        return server_not_found_response();
+        debug!("virtual_host_config_layer - claims missing virtual_host_id = {virtual_host_id}");
     }
 
     next.run(request).await
 }
 
-fn has_virtual_host(user_config: &UserConfig, virtual_host_id: &VirtualHostId) -> bool {
-    user_config.virtual_hosts.contains_key(virtual_host_id.value())
-}
-
 fn server_not_found_response() -> Response {
     Response::builder()
         .status(StatusCode::NOT_FOUND)