From 5b379b744513efcac6db0aa18dadae7de3881ac7 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 23 Jun 2026 10:38:22 +0100 Subject: [PATCH 1/7] gateway: route MCP resource subscribe/unsubscribe through backends Replace the local in-memory subscription set with real backend routing: subscribe/unsubscribe now validate the request, split the namespaced resource URI, resolve the single owning backend, and forward the call with the prefix stripped, returning backend and routing errors instead of reporting local-only success. Mirrors read_resource/get_prompt and drops the now-unused subscriptions HashSet field. The mock backend gains subscribe/unsubscribe handlers that only accept its own backend-local URIs, and a new integration test proves a namespaced subscribe/unsubscribe round-trips to the backend and that an unrouted URI fails with a gateway routing error. Implements IBM/mcp-context-forge#5253 Signed-off-by: lucarlig --- .../src/gateway/mcp_gateway.rs | 73 ++++++++++++---- .../tests/gateway_subscriptions.rs | 83 +++++++++++++++++++ .../tests/support/mock_counter.rs | 26 ++++++ 3 files changed, 164 insertions(+), 18 deletions(-) create mode 100644 crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index 0ae9c9dc..d4ad3ae9 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -1,7 +1,4 @@ -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; +use std::{collections::HashMap, sync::Arc}; use contextforge_gateway_rs_apis::user_store::UserConfig; use contextforge_gateway_rs_cpex::{GatewayPluginRuntimeHandle, ToolPreCallResult}; @@ -57,8 +54,6 @@ pub struct McpService where T: UserSessionStore, { - #[builder(default = Arc::new(Mutex::new(HashSet::new())))] - subscriptions: Arc>>, #[builder(default = BackendTransports::default())] transports: BackendTransports, http_client: reqwest::Client, @@ -418,13 +413,34 @@ where request: SubscribeRequestParams, cx: RequestContext, ) -> Result<(), ErrorData> { - let maybe_parts = cx.extensions.get::(); - let maybe_session = maybe_parts.and_then(|parts| parts.extensions.get::()); - let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::()); - info!("subscribe user_config = {maybe_user_config:#?} session_id = {maybe_session:#?}"); + let mcp_call_validator = AuthorizedCallValidator::new("subscribe", &cx); + 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 mut subs = self.subscriptions.lock().await; - subs.insert(request.uri.clone()); + let backend_names = session_manager.get_backend_names(); + + let Some((backend_name, resource_uri)) = split_prefixed_name(&request.uri, &backend_names) else { + return Err(ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: "Routing problem... wrong resource name".into(), + data: None, + }); + }; + let resource_uri = resource_uri.to_owned(); + + let (service_name, service) = resolve_backend(&session_manager, "subscribe", backend_name).await?; + + let mut routed_request = request; + routed_request.uri = resource_uri; + service.subscribe(routed_request).await.map_err(|error| { + warn!("subscribe: backend {service_name} {error:?}"); + ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: "Routing problem... got no responses from backends".into(), + data: None, + } + })?; + info!("subscribe: backend {service_name} completed"); Ok(()) } @@ -433,13 +449,34 @@ where request: UnsubscribeRequestParams, cx: RequestContext, ) -> Result<(), ErrorData> { - let maybe_parts = cx.extensions.get::(); - let maybe_session = maybe_parts.and_then(|parts| parts.extensions.get::()); - let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::()); - info!("unsubscribe user_config = {maybe_user_config:#?} session_id = {maybe_session:#?}"); + let mcp_call_validator = AuthorizedCallValidator::new("unsubscribe", &cx); + 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(); + + let Some((backend_name, resource_uri)) = split_prefixed_name(&request.uri, &backend_names) else { + return Err(ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: "Routing problem... wrong resource name".into(), + data: None, + }); + }; + let resource_uri = resource_uri.to_owned(); + + let (service_name, service) = resolve_backend(&session_manager, "unsubscribe", backend_name).await?; - let mut subs = self.subscriptions.lock().await; - subs.remove(request.uri.as_str()); + let mut routed_request = request; + routed_request.uri = resource_uri; + service.unsubscribe(routed_request).await.map_err(|error| { + warn!("unsubscribe: backend {service_name} {error:?}"); + ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: "Routing problem... got no responses from backends".into(), + data: None, + } + })?; + info!("unsubscribe: backend {service_name} completed"); Ok(()) } diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs new file mode 100644 index 00000000..7f0519dd --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs @@ -0,0 +1,83 @@ +mod support; + +use contextforge_gateway_rs_lib::{Config, Result, UpstreamConnectionMode}; +use rmcp::model::{SubscribeRequestParams, UnsubscribeRequestParams}; +use tracing::info; + +use support::{ + ListToolsGatewaySettings, connect_client, create_client, create_gateway_with_four_counters, create_ports, +}; + +fn plaintext_config(gateway_port: u16) -> Config { + Config { + address: Some(format!("127.0.0.1:{gateway_port}").parse().expect("This should work")), + token_verification_public_key: Some("../../assets/jwt.key.pub".into()), + upstream_connection_mode: Some(UpstreamConnectionMode::PlainTextOrTls), + ..Default::default() + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[test_log::test] +async fn plaintext_subscribes_and_unsubscribes_through_prefixed_backend() -> Result<()> { + let gateway_port = create_ports(1)[0]; + let user = "admin@example.com"; + let Ok(ListToolsGatewaySettings { handle, gateway_url, .. }) = + create_gateway_with_four_counters(user, plaintext_config(gateway_port)).await + else { + panic!("Invalid configuration "); + }; + + let client = create_client(user); + let maybe_passed = assert_subscribe_roundtrip(gateway_url, client).await; + + handle.abort(); + maybe_passed +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[test_log::test] +async fn plaintext_subscribe_to_unrouted_resource_errors() -> Result<()> { + let gateway_port = create_ports(1)[0]; + let user = "admin@example.com"; + let Ok(ListToolsGatewaySettings { handle, gateway_url, .. }) = + create_gateway_with_four_counters(user, plaintext_config(gateway_port)).await + else { + panic!("Invalid configuration "); + }; + + let client = create_client(user); + let maybe_passed = assert_unrouted_subscribe_errors(gateway_url, client).await; + + handle.abort(); + maybe_passed +} + +async fn assert_subscribe_roundtrip(gateway_url: String, client: reqwest::Client) -> Result<()> { + info!("Sending request to {gateway_url}"); + let running_service = connect_client(gateway_url, client).await?; + + // Pick a namespaced resource URI the gateway federated from a backend. + let resources = running_service.list_resources(None).await?; + let resource = resources.resources.first().ok_or("expected at least one federated resource")?; + let uri = resource.uri.clone(); + + // The mock backend only accepts its own backend-local URIs, so success proves the gateway + // routed to a single backend and stripped the namespace prefix before forwarding. + running_service.subscribe(SubscribeRequestParams::new(uri.clone())).await?; + running_service.unsubscribe(UnsubscribeRequestParams::new(uri)).await?; + + Ok(()) +} + +async fn assert_unrouted_subscribe_errors(gateway_url: String, client: reqwest::Client) -> Result<()> { + let running_service = connect_client(gateway_url, client).await?; + + // No backend namespace prefix => no route, so the gateway must reject it. + let result = running_service.subscribe(SubscribeRequestParams::new("unrouted://resource")).await; + if result.is_ok() { + return Err("expected a routing error for an unrouted resource URI".into()); + } + + Ok(()) +} diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs index b482132d..a29a1495 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs @@ -259,6 +259,26 @@ impl ServerHandler for Counter { Ok(CompleteResult::new(CompletionInfo::new(values).map_err(|e| McpError::internal_error(e, None))?)) } + async fn subscribe(&self, request: SubscribeRequestParams, _: RequestContext) -> Result<(), McpError> { + if is_known_resource_uri(&request.uri) { + Ok(()) + } else { + Err(McpError::resource_not_found("resource_not_found", Some(json!({ "uri": request.uri })))) + } + } + + async fn unsubscribe( + &self, + request: UnsubscribeRequestParams, + _: RequestContext, + ) -> Result<(), McpError> { + if is_known_resource_uri(&request.uri) { + Ok(()) + } else { + Err(McpError::resource_not_found("resource_not_found", Some(json!({ "uri": request.uri })))) + } + } + async fn list_resource_templates( &self, _request: Option, @@ -300,3 +320,9 @@ impl ServerHandler for Counter { Ok(self.get_info()) } } + +/// The backend-local resource URIs this mock owns; subscribe/unsubscribe only succeed for these, +/// so a successful gateway call proves the namespace prefix was stripped before forwarding. +fn is_known_resource_uri(uri: &str) -> bool { + matches!(uri, "str:////Users/to/some/path/" | "memo://insights") +} From 10e612381cc4c1dbef7c4dff6052212aa9c01a8b Mon Sep 17 00:00:00 2001 From: lucarlig Date: Wed, 24 Jun 2026 09:23:27 +0100 Subject: [PATCH 2/7] gateway: advertise resources.subscribe capability Now that subscribe/unsubscribe route to backends, declare the resources.subscribe server capability so spec-compliant clients discover and use the feature. Signed-off-by: lucarlig --- .../src/gateway/mcp_gateway.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index d4ad3ae9..f448fe32 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -682,7 +682,13 @@ async fn resolve_backend( } fn merge_capabilities(_server_capabilities: Vec<(String, Option)>) -> ServerCapabilities { - ServerCapabilities::builder().enable_completions().enable_prompts().enable_resources().enable_tools().build() + ServerCapabilities::builder() + .enable_completions() + .enable_prompts() + .enable_resources() + .enable_resources_subscribe() + .enable_tools() + .build() } fn log_list_backend_response( From c58ff13126394817750c227b4d98b229be3f8b29 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 7 Jul 2026 10:09:33 +0100 Subject: [PATCH 3/7] gateway: forward resource update notifications Signed-off-by: lucarlig --- .../src/gateway/backend_client.rs | 46 +++++- .../src/gateway/mcp_gateway.rs | 17 +- .../tests/gateway_subscriptions.rs | 148 ++++++++++++++++-- .../tests/support/mock_counter.rs | 17 +- 4 files changed, 206 insertions(+), 22 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs index 0208cd00..a46652e6 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs @@ -5,7 +5,7 @@ use rmcp::{ ClientHandler, Peer, RoleClient, RoleServer, model::{ CallToolRequestParams, CallToolResult, ClientRequest, InitializeRequestParams, Meta, ProgressNotificationParam, - ProgressToken, Request, ServerResult, + ProgressToken, Request, ResourceUpdatedNotificationParam, ServerResult, }, serde::{Serialize, de::DeserializeOwned}, service::{NotificationContext, PeerRequestOptions, ServiceError}, @@ -16,9 +16,11 @@ use tracing::{debug, warn}; #[derive(Clone)] pub(crate) struct GatewayBackendClient { + backend_name: String, initialize_request: InitializeRequestParams, plugin_runtime: Option, in_flight_calls: Arc>>>, + resource_subscriptions: Arc>>>, } #[derive(Debug)] @@ -31,10 +33,17 @@ struct InFlightToolCall { impl GatewayBackendClient { pub(crate) fn new( + backend_name: String, initialize_request: InitializeRequestParams, plugin_runtime: Option, ) -> Self { - Self { initialize_request, plugin_runtime, in_flight_calls: Arc::default() } + Self { + backend_name, + initialize_request, + plugin_runtime, + in_flight_calls: Arc::default(), + resource_subscriptions: Arc::default(), + } } pub(crate) async fn track_tool_call( @@ -70,6 +79,23 @@ impl GatewayBackendClient { calls.get(progress_token).cloned() } + pub(crate) async fn track_resource_subscription(&self, resource_uri: String, downstream: Peer) { + debug!("track_resource_subscription backend {} uri {resource_uri}", self.backend_name); + let mut subscriptions = self.resource_subscriptions.lock().await; + subscriptions.insert(resource_uri, downstream); + } + + pub(crate) async fn stop_tracking_resource_subscription(&self, resource_uri: &str) { + debug!("stop_tracking_resource_subscription backend {} uri {resource_uri}", self.backend_name); + let mut subscriptions = self.resource_subscriptions.lock().await; + subscriptions.remove(resource_uri); + } + + async fn resource_subscription(&self, resource_uri: &str) -> Option> { + let subscriptions = self.resource_subscriptions.lock().await; + subscriptions.get(resource_uri).cloned() + } + async fn stream_event_post_hook(&self, call: &InFlightToolCall, event: T) -> Option where T: Serialize + DeserializeOwned, @@ -116,6 +142,22 @@ impl ClientHandler for GatewayBackendClient { warn!("call_tool: unable to forward backend progress notification downstream: {error:?}"); } } + + async fn on_resource_updated( + &self, + mut params: ResourceUpdatedNotificationParam, + _context: NotificationContext, + ) { + let Some(downstream) = self.resource_subscription(¶ms.uri).await else { + debug!("resource_updated: dropping backend notification for unsubscribed uri {}", params.uri); + return; + }; + + params.uri = format!("{}-{}", self.backend_name, params.uri); + if let Err(error) = downstream.notify_resource_updated(params).await { + warn!("resource_updated: unable to forward backend notification downstream: {error:?}"); + } + } } /// Calls the tool on the backend, keeping the downstream progress token on diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index f448fe32..9c07d175 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -146,7 +146,8 @@ where .iter() .map(|(name, backend)| { let client = self.http_client.clone(); - let backend_client = GatewayBackendClient::new(request.clone(), self.plugin_runtime.clone()); + let backend_client = + GatewayBackendClient::new(name.clone(), request.clone(), self.plugin_runtime.clone()); let backend_url = backend.url.clone(); let downstream_session_id = downstream_session_id.clone(); @@ -432,14 +433,18 @@ where let mut routed_request = request; routed_request.uri = resource_uri; - service.subscribe(routed_request).await.map_err(|error| { + let tracked_uri = routed_request.uri.clone(); + service.service().track_resource_subscription(tracked_uri.clone(), cx.peer.clone()).await; + + if let Err(error) = service.subscribe(routed_request).await { + service.service().stop_tracking_resource_subscription(&tracked_uri).await; warn!("subscribe: backend {service_name} {error:?}"); - ErrorData { + return Err(ErrorData { code: ErrorCode::INTERNAL_ERROR, message: "Routing problem... got no responses from backends".into(), data: None, - } - })?; + }); + } info!("subscribe: backend {service_name} completed"); Ok(()) } @@ -467,6 +472,7 @@ where let (service_name, service) = resolve_backend(&session_manager, "unsubscribe", backend_name).await?; let mut routed_request = request; + let tracked_uri = resource_uri.clone(); routed_request.uri = resource_uri; service.unsubscribe(routed_request).await.map_err(|error| { warn!("unsubscribe: backend {service_name} {error:?}"); @@ -476,6 +482,7 @@ where data: None, } })?; + service.service().stop_tracking_resource_subscription(&tracked_uri).await; info!("unsubscribe: backend {service_name} completed"); Ok(()) } diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs index 7f0519dd..952492f9 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs @@ -1,13 +1,56 @@ mod support; +use std::{ + collections::HashSet, + sync::{Arc, Mutex as StdMutex}, + time::{Duration, Instant}, +}; + use contextforge_gateway_rs_lib::{Config, Result, UpstreamConnectionMode}; -use rmcp::model::{SubscribeRequestParams, UnsubscribeRequestParams}; -use tracing::info; +use futures::future::try_join_all; +use rmcp::{ + ClientHandler, ServiceExt, + model::{ + ClientCapabilities, Implementation, InitializeRequestParams, ResourceUpdatedNotificationParam, + SubscribeRequestParams, UnsubscribeRequestParams, + }, + service::{NotificationContext, RoleClient, RunningService}, + transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig}, +}; +use tracing::warn; use support::{ ListToolsGatewaySettings, connect_client, create_client, create_gateway_with_four_counters, create_ports, }; +const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); +const EXPECTED_UPDATES_PER_BACKEND: usize = 4; + +type Recorded = Arc>>; + +#[derive(Clone, Default)] +struct RecordingClient { + resource_updates: Recorded, +} + +impl ClientHandler for RecordingClient { + fn get_info(&self) -> InitializeRequestParams { + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("resource-update-recording-test-client", "0.1.0"), + ) + } + + async fn on_resource_updated( + &self, + params: ResourceUpdatedNotificationParam, + _context: NotificationContext, + ) { + self.resource_updates.lock().expect("resource update lock poisoned").push(params); + } +} + fn plaintext_config(gateway_port: u16) -> Config { Config { address: Some(format!("127.0.0.1:{gateway_port}").parse().expect("This should work")), @@ -19,7 +62,7 @@ fn plaintext_config(gateway_port: u16) -> Config { #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[test_log::test] -async fn plaintext_subscribes_and_unsubscribes_through_prefixed_backend() -> Result<()> { +async fn plaintext_subscribes_and_unsubscribes_through_two_prefixed_backends() -> Result<()> { let gateway_port = create_ports(1)[0]; let user = "admin@example.com"; let Ok(ListToolsGatewaySettings { handle, gateway_url, .. }) = @@ -29,7 +72,7 @@ async fn plaintext_subscribes_and_unsubscribes_through_prefixed_backend() -> Res }; let client = create_client(user); - let maybe_passed = assert_subscribe_roundtrip(gateway_url, client).await; + let maybe_passed = assert_two_backend_subscribe_roundtrips(gateway_url, client).await; handle.abort(); maybe_passed @@ -53,19 +96,36 @@ async fn plaintext_subscribe_to_unrouted_resource_errors() -> Result<()> { maybe_passed } -async fn assert_subscribe_roundtrip(gateway_url: String, client: reqwest::Client) -> Result<()> { - info!("Sending request to {gateway_url}"); - let running_service = connect_client(gateway_url, client).await?; +async fn assert_two_backend_subscribe_roundtrips(gateway_url: String, client: reqwest::Client) -> Result<()> { + let recording_client = RecordingClient::default(); + let resource_updates = Arc::clone(&recording_client.resource_updates); + let running_service = connect_recording_client(gateway_url, client, recording_client).await?; - // Pick a namespaced resource URI the gateway federated from a backend. let resources = running_service.list_resources(None).await?; - let resource = resources.resources.first().ok_or("expected at least one federated resource")?; - let uri = resource.uri.clone(); + let mut selected_backends = HashSet::new(); + let mut selected_uris = Vec::new(); + + for resource in resources.resources { + let backend_name = mock_backend_name(&resource.uri) + .ok_or_else(|| format!("expected mock backend-prefixed URI, got {}", resource.uri))?; + if selected_backends.insert(backend_name) { + selected_uris.push(resource.uri.clone()); + } + if selected_uris.len() == 2 { + break; + } + } + + if selected_uris.len() != 2 { + return Err(format!("expected resources from at least two backends, got {}", selected_uris.len()).into()); + } - // The mock backend only accepts its own backend-local URIs, so success proves the gateway - // routed to a single backend and stripped the namespace prefix before forwarding. - running_service.subscribe(SubscribeRequestParams::new(uri.clone())).await?; - running_service.unsubscribe(UnsubscribeRequestParams::new(uri)).await?; + try_join_all(selected_uris.iter().map(|uri| running_service.subscribe(SubscribeRequestParams::new(uri.clone())))) + .await?; + wait_for_resource_updates(&resource_updates, &selected_uris, EXPECTED_UPDATES_PER_BACKEND).await?; + for uri in selected_uris { + running_service.unsubscribe(UnsubscribeRequestParams::new(uri)).await?; + } Ok(()) } @@ -81,3 +141,63 @@ async fn assert_unrouted_subscribe_errors(gateway_url: String, client: reqwest:: Ok(()) } + +async fn connect_recording_client( + gateway_url: String, + client: reqwest::Client, + recording_client: RecordingClient, +) -> Result> { + let deadline = Instant::now() + CLIENT_CONNECT_TIMEOUT; + loop { + let config = StreamableHttpClientTransportConfig::with_uri(gateway_url.clone()); + let transport = StreamableHttpClientTransport::with_client(client.clone(), config); + + match recording_client.clone().serve(transport).await { + Ok(running_service) => return Ok(running_service), + Err(error) if Instant::now() < deadline => { + warn!("No Service {error:?}"); + tokio::time::sleep(TEST_POLL_INTERVAL).await; + }, + Err(error) => { + warn!("No Service {error:?}"); + return Err("Couldn't get a service".into()); + }, + } + } +} + +async fn wait_for_resource_updates( + resource_updates: &StdMutex>, + expected_uris: &[String], + expected_count_per_uri: usize, +) -> Result<()> { + let deadline = Instant::now() + CLIENT_CONNECT_TIMEOUT; + + loop { + let counts = { + let updates = resource_updates.lock().expect("resource update lock poisoned"); + expected_uris + .iter() + .map(|uri| { + let count = updates.iter().filter(|update| update.uri == *uri).count(); + (uri.clone(), count) + }) + .collect::>() + }; + + if counts.iter().all(|(_, count)| *count >= expected_count_per_uri) { + return Ok(()); + } + if Instant::now() >= deadline { + return Err( + format!("expected {expected_count_per_uri} resource updates per URI, got counts {counts:?}").into() + ); + } + + tokio::time::sleep(TEST_POLL_INTERVAL).await; + } +} + +fn mock_backend_name(uri: &str) -> Option { + uri.strip_suffix("-str:////Users/to/some/path/").or_else(|| uri.strip_suffix("-memo://insights")).map(str::to_owned) +} diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs index a29a1495..f6ba2b45 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs @@ -259,8 +259,23 @@ impl ServerHandler for Counter { Ok(CompleteResult::new(CompletionInfo::new(values).map_err(|e| McpError::internal_error(e, None))?)) } - async fn subscribe(&self, request: SubscribeRequestParams, _: RequestContext) -> Result<(), McpError> { + async fn subscribe( + &self, + request: SubscribeRequestParams, + context: RequestContext, + ) -> Result<(), McpError> { if is_known_resource_uri(&request.uri) { + let uri = request.uri.clone(); + let peer = context.peer; + tokio::spawn(async move { + for _ in 0..4 { + if let Err(error) = + peer.notify_resource_updated(ResourceUpdatedNotificationParam::new(uri.clone())).await + { + tracing::warn!("mock_counter: failed to send resource update notification: {error:?}"); + } + } + }); Ok(()) } else { Err(McpError::resource_not_found("resource_not_found", Some(json!({ "uri": request.uri })))) From fbb712b80faa8d65c7b09b3d90da1ff800565519 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 7 Jul 2026 10:15:34 +0100 Subject: [PATCH 4/7] gateway: remove unused imports Signed-off-by: lucarlig --- crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index 9c07d175..c0a4f125 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -1,8 +1,6 @@ use std::{collections::HashMap, sync::Arc}; -use contextforge_gateway_rs_apis::user_store::UserConfig; use contextforge_gateway_rs_cpex::{GatewayPluginRuntimeHandle, ToolPreCallResult}; -use http::request::Parts; use itertools::Itertools; use rmcp::{ ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, From 5daf33d10a8f0f8cffb3b56853befc929f5a8b5f Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 7 Jul 2026 13:56:15 +0100 Subject: [PATCH 5/7] gateway: share prefixed-name routing across handlers call_tool, read_resource, get_prompt, complete, subscribe, and unsubscribe all repeated the same split-prefix/resolve-backend preamble; extract it into route_prefixed_name. Signed-off-by: lucarlig --- .../src/gateway/mcp_gateway.rs | 129 +++++++----------- 1 file changed, 51 insertions(+), 78 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index c0a4f125..812f00b8 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -268,22 +268,12 @@ where 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(); - - let Some((backend_name, tool_name)) = split_prefixed_name(&request.name, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong tool name".into(), - data: None, - }); - }; - let backend_name = backend_name.to_owned(); - let tool_name = tool_name.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "call_tool", &backend_name).await?; + let (service_name, service, tool_name) = + route_prefixed_name(&session_manager, "call_tool", &request.name, "Routing problem... wrong tool name") + .await?; let pre_result = if let Some(plugin_runtime) = &self.plugin_runtime { - plugin_runtime.before_tool_call(&request, &tool_name, &backend_name).await? + plugin_runtime.before_tool_call(&request, &tool_name, &service_name).await? } else { ToolPreCallResult::unchanged() }; @@ -351,18 +341,13 @@ where 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(); - - let Some((backend_name, resource_uri)) = split_prefixed_name(&request.uri, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong resource name".into(), - data: None, - }); - }; - let resource_uri = resource_uri.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "read_resource", backend_name).await?; + let (service_name, service, resource_uri) = route_prefixed_name( + &session_manager, + "read_resource", + &request.uri, + "Routing problem... wrong resource name", + ) + .await?; let mut routed_request = request; routed_request.uri = resource_uri; @@ -416,22 +401,13 @@ where 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(); - - let Some((backend_name, resource_uri)) = split_prefixed_name(&request.uri, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong resource name".into(), - data: None, - }); - }; - let resource_uri = resource_uri.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "subscribe", backend_name).await?; + let (service_name, service, resource_uri) = + route_prefixed_name(&session_manager, "subscribe", &request.uri, "Routing problem... wrong resource name") + .await?; + let tracked_uri = resource_uri.clone(); let mut routed_request = request; routed_request.uri = resource_uri; - let tracked_uri = routed_request.uri.clone(); service.service().track_resource_subscription(tracked_uri.clone(), cx.peer.clone()).await; if let Err(error) = service.subscribe(routed_request).await { @@ -456,21 +432,16 @@ where 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(); - - let Some((backend_name, resource_uri)) = split_prefixed_name(&request.uri, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong resource name".into(), - data: None, - }); - }; - let resource_uri = resource_uri.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "unsubscribe", backend_name).await?; + let (service_name, service, resource_uri) = route_prefixed_name( + &session_manager, + "unsubscribe", + &request.uri, + "Routing problem... wrong resource name", + ) + .await?; - let mut routed_request = request; let tracked_uri = resource_uri.clone(); + let mut routed_request = request; routed_request.uri = resource_uri; service.unsubscribe(routed_request).await.map_err(|error| { warn!("unsubscribe: backend {service_name} {error:?}"); @@ -519,19 +490,9 @@ where 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(); - - let Some((backend_name, prompt_name)) = split_prefixed_name(&request.name, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong prompt name".into(), - data: None, - }); - }; - let backend_name = backend_name.to_owned(); - let prompt_name = prompt_name.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "get_prompt", &backend_name).await?; + let (service_name, service, prompt_name) = + route_prefixed_name(&session_manager, "get_prompt", &request.name, "Routing problem... wrong prompt name") + .await?; let mut routed_request = request; routed_request.name = prompt_name; @@ -556,25 +517,19 @@ where 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(); - // The reference carries a namespaced prompt name or resource URI; route on that. let namespaced = match &request.r#ref { Reference::Prompt(prompt) => prompt.name.as_str(), Reference::Resource(resource) => resource.uri.as_str(), }; - let Some((backend_name, stripped)) = split_prefixed_name(namespaced, &backend_names) else { - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... wrong completion reference".into(), - data: None, - }); - }; - let backend_name = backend_name.to_owned(); - let stripped = stripped.to_owned(); - - let (service_name, service) = resolve_backend(&session_manager, "complete", &backend_name).await?; + let (service_name, service, stripped) = route_prefixed_name( + &session_manager, + "complete", + namespaced, + "Routing problem... wrong completion reference", + ) + .await?; let mut routed_request = request; match &mut routed_request.r#ref { @@ -641,6 +596,24 @@ where .collect() } +/// Routes a namespaced `{backend}-{rest}` name: splits it against the session's backend names +/// and resolves the owning backend, returning `(backend_name, service, rest)`. Shared by tool, +/// resource, prompt, and completion routing. +async fn route_prefixed_name( + session_manager: &SessionManager<'_>, + op: &str, + namespaced: &str, + no_route_message: &'static str, +) -> Result<(String, McpClientService, String), ErrorData> { + let backend_names = session_manager.get_backend_names(); + let Some((backend_name, rest)) = split_prefixed_name(namespaced, &backend_names) else { + return Err(ErrorData { code: ErrorCode::INTERNAL_ERROR, message: no_route_message.into(), data: None }); + }; + let rest = rest.to_owned(); + let (backend_name, service) = resolve_backend(session_manager, op, backend_name).await?; + Ok((backend_name, service, rest)) +} + /// Resolves the single connected backend named `backend_name` and takes its running service. /// Shared by tool, resource, and prompt routing so they reject duplicate or missing backends /// the same way; a duplicate match means the session is invalid, so it is cleaned up. From ce47364a00b79e7dbe763162b01aa90f6ec62935 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 7 Jul 2026 13:56:15 +0100 Subject: [PATCH 6/7] gateway tests: assert updates stop after unsubscribe The mock backend now notifies on an interval and deliberately keeps notifying after unsubscribe, so the round-trip test can assert the gateway itself stops forwarding updates for unsubscribed URIs. This also decouples the test's update threshold from a hardcoded mock notification count. Signed-off-by: lucarlig --- .../tests/gateway_subscriptions.rs | 28 +++++++++++++++++-- .../tests/support/mock_counter.rs | 19 +++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs index 952492f9..4172434d 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs @@ -21,11 +21,13 @@ use tracing::warn; use support::{ ListToolsGatewaySettings, connect_client, create_client, create_gateway_with_four_counters, create_ports, + mock_counter::RESOURCE_UPDATE_NOTIFY_INTERVAL, }; const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); -const EXPECTED_UPDATES_PER_BACKEND: usize = 4; +/// The mocks notify continuously, so this is just the threshold proving delivery works. +const MIN_UPDATES_PER_BACKEND: usize = 4; type Recorded = Arc>>; @@ -122,11 +124,33 @@ async fn assert_two_backend_subscribe_roundtrips(gateway_url: String, client: re try_join_all(selected_uris.iter().map(|uri| running_service.subscribe(SubscribeRequestParams::new(uri.clone())))) .await?; - wait_for_resource_updates(&resource_updates, &selected_uris, EXPECTED_UPDATES_PER_BACKEND).await?; + wait_for_resource_updates(&resource_updates, &selected_uris, MIN_UPDATES_PER_BACKEND).await?; for uri in selected_uris { running_service.unsubscribe(UnsubscribeRequestParams::new(uri)).await?; } + assert_no_more_resource_updates(&resource_updates).await +} + +/// The mock backends keep notifying after unsubscribe, so any update recorded after the quiet +/// window starts would mean the gateway kept forwarding for an unsubscribed URI. +async fn assert_no_more_resource_updates( + resource_updates: &StdMutex>, +) -> Result<()> { + // Let updates the gateway forwarded before the unsubscribe finish arriving. + tokio::time::sleep(RESOURCE_UPDATE_NOTIFY_INTERVAL * 5).await; + let count_after_drain = resource_updates.lock().expect("resource update lock poisoned").len(); + + tokio::time::sleep(RESOURCE_UPDATE_NOTIFY_INTERVAL * 10).await; + let count_after_quiet = resource_updates.lock().expect("resource update lock poisoned").len(); + + if count_after_quiet != count_after_drain { + return Err(format!( + "expected no resource updates after unsubscribe, got {} new", + count_after_quiet - count_after_drain + ) + .into()); + } Ok(()) } diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs index f6ba2b45..2a855ce8 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs @@ -1,7 +1,7 @@ #![allow(clippy::pedantic)] #![allow(dead_code)] -use std::{any::Any, sync::Arc}; +use std::{any::Any, sync::Arc, time::Duration}; use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, @@ -267,13 +267,14 @@ impl ServerHandler for Counter { if is_known_resource_uri(&request.uri) { let uri = request.uri.clone(); let peer = context.peer; + // Keeps notifying even after unsubscribe (a rude backend): the gateway must stop + // forwarding updates for unsubscribed URIs itself, and tests assert exactly that. tokio::spawn(async move { - for _ in 0..4 { - if let Err(error) = - peer.notify_resource_updated(ResourceUpdatedNotificationParam::new(uri.clone())).await - { - tracing::warn!("mock_counter: failed to send resource update notification: {error:?}"); + for _ in 0..MAX_RESOURCE_UPDATE_NOTIFICATIONS { + if peer.notify_resource_updated(ResourceUpdatedNotificationParam::new(uri.clone())).await.is_err() { + break; } + tokio::time::sleep(RESOURCE_UPDATE_NOTIFY_INTERVAL).await; } }); Ok(()) @@ -341,3 +342,9 @@ impl ServerHandler for Counter { fn is_known_resource_uri(uri: &str) -> bool { matches!(uri, "str:////Users/to/some/path/" | "memo://insights") } + +/// Interval between the resource-update notifications sent after a subscribe is accepted. +pub const RESOURCE_UPDATE_NOTIFY_INTERVAL: Duration = Duration::from_millis(10); + +/// Safety cap so notify loops can't outlive a hung test run. +const MAX_RESOURCE_UPDATE_NOTIFICATIONS: usize = 1000; From 207e338e56ea46071c1e42d889217953a3a2175f Mon Sep 17 00:00:00 2001 From: lucarlig Date: Thu, 9 Jul 2026 10:38:03 +0100 Subject: [PATCH 7/7] gateway: consolidate prefix join, backend errors, and test helpers Add prefixed_name as the inverse of split_prefixed_name and route all namespace joins through it, share one backend_forward_error helper across the forwarding handlers, make subscription tracking take &str, and reuse the support connect/config helpers instead of per-test copies. Signed-off-by: lucarlig --- .../src/gateway/backend_client.rs | 8 +- .../src/gateway/mcp_gateway.rs | 104 ++++++++---------- .../tests/gateway_completions.rs | 12 +- .../tests/gateway_subscriptions.rs | 61 +++------- .../tests/support/client.rs | 19 +++- .../tests/support/list_tools_gateway.rs | 14 ++- .../tests/support/mock_counter.rs | 4 +- .../tests/support/mod.rs | 7 +- 8 files changed, 104 insertions(+), 125 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs index a46652e6..247cfe7d 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs @@ -14,6 +14,8 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; +use super::mcp_gateway::prefixed_name; + #[derive(Clone)] pub(crate) struct GatewayBackendClient { backend_name: String, @@ -79,10 +81,10 @@ impl GatewayBackendClient { calls.get(progress_token).cloned() } - pub(crate) async fn track_resource_subscription(&self, resource_uri: String, downstream: Peer) { + pub(crate) async fn track_resource_subscription(&self, resource_uri: &str, downstream: Peer) { debug!("track_resource_subscription backend {} uri {resource_uri}", self.backend_name); let mut subscriptions = self.resource_subscriptions.lock().await; - subscriptions.insert(resource_uri, downstream); + subscriptions.insert(resource_uri.to_owned(), downstream); } pub(crate) async fn stop_tracking_resource_subscription(&self, resource_uri: &str) { @@ -153,7 +155,7 @@ impl ClientHandler for GatewayBackendClient { return; }; - params.uri = format!("{}-{}", self.backend_name, params.uri); + params.uri = prefixed_name(&self.backend_name, ¶ms.uri); if let Err(error) = downstream.notify_resource_updated(params).await { warn!("resource_updated: unable to forward backend notification downstream: {error:?}"); } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index 812f00b8..f5a21f1c 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -289,14 +289,7 @@ where let response = call_backend_tool(service.peer(), routed_request, progress_token.clone(), cx.ct.clone()).await; service.service().stop_tracking_tool_call(progress_token.clone()).await; - let response = response.map_err(|error| { - warn!("call_tool: backend {service_name} {error:?}"); - ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - } - })?; + let response = response.map_err(|error| backend_forward_error("call_tool", &service_name, &error))?; let response = match (&self.plugin_runtime, post_state) { (Some(plugin_runtime), Some(post_state)) => { plugin_runtime.after_tool_call(&tool_name, response, Some(post_state)).await? @@ -351,14 +344,10 @@ where let mut routed_request = request; routed_request.uri = resource_uri; - let response = service.read_resource(routed_request).await.map_err(|error| { - warn!("read_resource: backend {service_name} {error:?}"); - ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - } - })?; + let response = service + .read_resource(routed_request) + .await + .map_err(|error| backend_forward_error("read_resource", &service_name, &error))?; info!("read_resource: backend {service_name} returned {} contents", response.contents.len()); Ok(response) } @@ -405,19 +394,13 @@ where route_prefixed_name(&session_manager, "subscribe", &request.uri, "Routing problem... wrong resource name") .await?; - let tracked_uri = resource_uri.clone(); let mut routed_request = request; - routed_request.uri = resource_uri; - service.service().track_resource_subscription(tracked_uri.clone(), cx.peer.clone()).await; + routed_request.uri = resource_uri.clone(); + service.service().track_resource_subscription(&resource_uri, cx.peer.clone()).await; if let Err(error) = service.subscribe(routed_request).await { - service.service().stop_tracking_resource_subscription(&tracked_uri).await; - warn!("subscribe: backend {service_name} {error:?}"); - return Err(ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - }); + service.service().stop_tracking_resource_subscription(&resource_uri).await; + return Err(backend_forward_error("subscribe", &service_name, &error)); } info!("subscribe: backend {service_name} completed"); Ok(()) @@ -440,18 +423,13 @@ where ) .await?; - let tracked_uri = resource_uri.clone(); let mut routed_request = request; - routed_request.uri = resource_uri; - service.unsubscribe(routed_request).await.map_err(|error| { - warn!("unsubscribe: backend {service_name} {error:?}"); - ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - } - })?; - service.service().stop_tracking_resource_subscription(&tracked_uri).await; + routed_request.uri = resource_uri.clone(); + service + .unsubscribe(routed_request) + .await + .map_err(|error| backend_forward_error("unsubscribe", &service_name, &error))?; + service.service().stop_tracking_resource_subscription(&resource_uri).await; info!("unsubscribe: backend {service_name} completed"); Ok(()) } @@ -496,14 +474,10 @@ where let mut routed_request = request; routed_request.name = prompt_name; - let response = service.get_prompt(routed_request).await.map_err(|_| { - warn!("get_prompt: backend {service_name} returned an error"); - ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - } - })?; + let response = service + .get_prompt(routed_request) + .await + .map_err(|error| backend_forward_error("get_prompt", &service_name, &error))?; info!("get_prompt: backend {service_name} returned {} messages", response.messages.len()); Ok(response) } @@ -536,14 +510,10 @@ where Reference::Prompt(prompt) => prompt.name = stripped, Reference::Resource(resource) => resource.uri = stripped, } - let response = service.complete(routed_request).await.map_err(|error| { - warn!("complete: backend {service_name} {error:?}"); - ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: "Routing problem... got no responses from backends".into(), - data: None, - } - })?; + let response = service + .complete(routed_request) + .await + .map_err(|error| backend_forward_error("complete", &service_name, &error))?; info!("complete: backend {service_name} returned {} values", response.completion.values.len()); Ok(response) } @@ -558,6 +528,22 @@ fn split_prefixed_name<'a, N: AsRef>(name: &'a str, backend_names: &'a [N]) }) } +/// Joins a backend name and a backend-local name into the namespaced `{backend}-{rest}` form. +/// Inverse of [`split_prefixed_name`]; together they own the naming convention. +pub(crate) fn prefixed_name(backend_name: &str, rest: &str) -> String { + format!("{backend_name}-{rest}") +} + +/// Logs a backend forwarding failure and maps it to the routing error every handler returns. +fn backend_forward_error(op: &str, backend_name: &str, error: &impl std::fmt::Debug) -> ErrorData { + warn!("{op}: backend {backend_name} {error:?}"); + ErrorData { + code: ErrorCode::INTERNAL_ERROR, + message: "Routing problem... got no responses from backends".into(), + data: None, + } +} + /// Fans a paginated list request out to every connected backend concurrently, logs each response, /// and returns the `(backend_name, result)` pairs that succeeded. async fn fan_out_list( @@ -690,7 +676,7 @@ fn merge_tools(tools: Vec<(String, ListToolsResult)>) -> Vec { .tools .into_iter() .map(|mut t| { - t.name = format!("{backend_name}-{}", t.name).into(); + t.name = prefixed_name(&backend_name, &t.name).into(); t }) .collect::>() @@ -707,8 +693,8 @@ fn merge_resources(resources: Vec<(String, ListResourcesResult)>) -> Vec>() @@ -722,8 +708,8 @@ fn merge_resource_templates(templates: Vec<(String, ListResourceTemplatesResult) .into_iter() .flat_map(|(backend_name, result)| { result.resource_templates.into_iter().map(move |mut template| { - template.name = format!("{backend_name}-{}", template.name); - template.uri_template = format!("{backend_name}-{}", template.uri_template); + template.name = prefixed_name(&backend_name, &template.name); + template.uri_template = prefixed_name(&backend_name, &template.uri_template); template }) }) @@ -736,7 +722,7 @@ fn merge_prompts(prompts: Vec<(String, ListPromptsResult)>) -> Vec { .into_iter() .flat_map(|(backend_name, result)| { result.prompts.into_iter().map(move |mut prompt| { - prompt.name = format!("{backend_name}-{}", prompt.name); + prompt.name = prefixed_name(&backend_name, &prompt.name); prompt }) }) diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_completions.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_completions.rs index 4b41f589..fbc0464b 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_completions.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_completions.rs @@ -1,21 +1,13 @@ mod support; -use contextforge_gateway_rs_lib::{Config, Result, UpstreamConnectionMode}; +use contextforge_gateway_rs_lib::Result; use tracing::info; use support::{ ListToolsGatewaySettings, connect_client, create_client, create_gateway_with_four_counters, create_ports, + plaintext_config, }; -fn plaintext_config(gateway_port: u16) -> Config { - Config { - address: Some(format!("127.0.0.1:{gateway_port}").parse().expect("This should work")), - token_verification_public_key: Some("../../assets/jwt.key.pub".into()), - upstream_connection_mode: Some(UpstreamConnectionMode::PlainTextOrTls), - ..Default::default() - } -} - #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[test_log::test] async fn plaintext_completes_prompt_argument_through_prefixed_backend() -> Result<()> { diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs index 4172434d..b63dddce 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_subscriptions.rs @@ -3,29 +3,27 @@ mod support; use std::{ collections::HashSet, sync::{Arc, Mutex as StdMutex}, - time::{Duration, Instant}, + time::Instant, }; -use contextforge_gateway_rs_lib::{Config, Result, UpstreamConnectionMode}; +use contextforge_gateway_rs_lib::Result; use futures::future::try_join_all; use rmcp::{ - ClientHandler, ServiceExt, + ClientHandler, model::{ ClientCapabilities, Implementation, InitializeRequestParams, ResourceUpdatedNotificationParam, SubscribeRequestParams, UnsubscribeRequestParams, }, - service::{NotificationContext, RoleClient, RunningService}, - transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig}, + service::{NotificationContext, RoleClient}, }; -use tracing::warn; use support::{ - ListToolsGatewaySettings, connect_client, create_client, create_gateway_with_four_counters, create_ports, - mock_counter::RESOURCE_UPDATE_NOTIFY_INTERVAL, + CLIENT_CONNECT_TIMEOUT, ListToolsGatewaySettings, TEST_POLL_INTERVAL, connect_client, connect_client_with_handler, + create_client, create_gateway_with_four_counters, create_ports, + mock_counter::{KNOWN_RESOURCE_URIS, RESOURCE_UPDATE_NOTIFY_INTERVAL}, + plaintext_config, }; -const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); -const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); /// The mocks notify continuously, so this is just the threshold proving delivery works. const MIN_UPDATES_PER_BACKEND: usize = 4; @@ -53,15 +51,6 @@ impl ClientHandler for RecordingClient { } } -fn plaintext_config(gateway_port: u16) -> Config { - Config { - address: Some(format!("127.0.0.1:{gateway_port}").parse().expect("This should work")), - token_verification_public_key: Some("../../assets/jwt.key.pub".into()), - upstream_connection_mode: Some(UpstreamConnectionMode::PlainTextOrTls), - ..Default::default() - } -} - #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[test_log::test] async fn plaintext_subscribes_and_unsubscribes_through_two_prefixed_backends() -> Result<()> { @@ -101,7 +90,7 @@ async fn plaintext_subscribe_to_unrouted_resource_errors() -> Result<()> { async fn assert_two_backend_subscribe_roundtrips(gateway_url: String, client: reqwest::Client) -> Result<()> { let recording_client = RecordingClient::default(); let resource_updates = Arc::clone(&recording_client.resource_updates); - let running_service = connect_recording_client(gateway_url, client, recording_client).await?; + let running_service = connect_client_with_handler(gateway_url, client, recording_client).await?; let resources = running_service.list_resources(None).await?; let mut selected_backends = HashSet::new(); @@ -166,30 +155,6 @@ async fn assert_unrouted_subscribe_errors(gateway_url: String, client: reqwest:: Ok(()) } -async fn connect_recording_client( - gateway_url: String, - client: reqwest::Client, - recording_client: RecordingClient, -) -> Result> { - let deadline = Instant::now() + CLIENT_CONNECT_TIMEOUT; - loop { - let config = StreamableHttpClientTransportConfig::with_uri(gateway_url.clone()); - let transport = StreamableHttpClientTransport::with_client(client.clone(), config); - - match recording_client.clone().serve(transport).await { - Ok(running_service) => return Ok(running_service), - Err(error) if Instant::now() < deadline => { - warn!("No Service {error:?}"); - tokio::time::sleep(TEST_POLL_INTERVAL).await; - }, - Err(error) => { - warn!("No Service {error:?}"); - return Err("Couldn't get a service".into()); - }, - } - } -} - async fn wait_for_resource_updates( resource_updates: &StdMutex>, expected_uris: &[String], @@ -222,6 +187,12 @@ async fn wait_for_resource_updates( } } +/// Extracts the backend prefix from a namespaced mock resource URI; the suffixes are the +/// backend-local URIs the mock owns, so this stays in lockstep with the mock's list. fn mock_backend_name(uri: &str) -> Option { - uri.strip_suffix("-str:////Users/to/some/path/").or_else(|| uri.strip_suffix("-memo://insights")).map(str::to_owned) + KNOWN_RESOURCE_URIS + .iter() + .find_map(|known| uri.strip_suffix(known)) + .and_then(|prefix| prefix.strip_suffix('-')) + .map(str::to_owned) } diff --git a/crates/contextforge-gateway-rs-lib/tests/support/client.rs b/crates/contextforge-gateway-rs-lib/tests/support/client.rs index 282cba85..135e84a2 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/client.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/client.rs @@ -11,8 +11,8 @@ use tracing::warn; use super::auth::token; -const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); -const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); +pub(crate) const CLIENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2); +pub(crate) const TEST_POLL_INTERVAL: Duration = Duration::from_millis(20); pub(crate) fn create_client(user: &str) -> reqwest::Client { reqwest::Client::builder().default_headers(auth_headers(user)).build().expect("This should work") @@ -41,13 +41,24 @@ pub(crate) async fn connect_client( gateway_url: String, client: reqwest::Client, ) -> Result> { + connect_client_with_handler(gateway_url, client, InitializeRequestParams::default()).await +} + +/// Connects any `ClientHandler` to the gateway, retrying until `CLIENT_CONNECT_TIMEOUT`. +pub(crate) async fn connect_client_with_handler( + gateway_url: String, + client: reqwest::Client, + handler: H, +) -> Result> +where + H: rmcp::ClientHandler + Clone, +{ let deadline = Instant::now() + CLIENT_CONNECT_TIMEOUT; loop { let config = StreamableHttpClientTransportConfig::with_uri(gateway_url.clone()); let transport = StreamableHttpClientTransport::with_client(client.clone(), config); - let request = InitializeRequestParams::default(); - match request.serve(transport).await { + match handler.clone().serve(transport).await { Ok(running_service) => return Ok(running_service), Err(error) if Instant::now() < deadline => { warn!("No Service {error:?}"); diff --git a/crates/contextforge-gateway-rs-lib/tests/support/list_tools_gateway.rs b/crates/contextforge-gateway-rs-lib/tests/support/list_tools_gateway.rs index b759e443..15182952 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/list_tools_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/list_tools_gateway.rs @@ -4,7 +4,9 @@ use contextforge_gateway_rs_apis::{ User, user_store::{BackendMCPGateway, Transport, UserConfig, VirtualHost}, }; -use contextforge_gateway_rs_lib::{Config, Gateway, Result, UserConfigStore, UserConfigStoreType}; +use contextforge_gateway_rs_lib::{ + Config, Gateway, Result, UpstreamConnectionMode, UserConfigStore, UserConfigStoreType, +}; use futures::{FutureExt, future::BoxFuture}; use rmcp::transport::{ StreamableHttpServerConfig, StreamableHttpService, streamable_http_server::session::local::LocalSessionManager, @@ -28,6 +30,16 @@ pub(crate) struct ListToolsGatewaySettings { pub(crate) expected_resource_template_uris: Vec, } +/// Gateway config for plaintext-upstream tests, shared by the integration test binaries. +pub(crate) fn plaintext_config(gateway_port: u16) -> Config { + Config { + address: Some(format!("127.0.0.1:{gateway_port}").parse().expect("This should work")), + token_verification_public_key: Some("../../assets/jwt.key.pub".into()), + upstream_connection_mode: Some(UpstreamConnectionMode::PlainTextOrTls), + ..Default::default() + } +} + pub(crate) fn create_ports(ports: usize) -> Vec { (0..ports).map(|_| openport::pick_random_unused_port().expect("Expecting to find port")).collect() } diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs index 2a855ce8..ff24089f 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mock_counter.rs @@ -339,8 +339,10 @@ impl ServerHandler for Counter { /// The backend-local resource URIs this mock owns; subscribe/unsubscribe only succeed for these, /// so a successful gateway call proves the namespace prefix was stripped before forwarding. +pub const KNOWN_RESOURCE_URIS: [&str; 2] = ["str:////Users/to/some/path/", "memo://insights"]; + fn is_known_resource_uri(uri: &str) -> bool { - matches!(uri, "str:////Users/to/some/path/" | "memo://insights") + KNOWN_RESOURCE_URIS.contains(&uri) } /// Interval between the resource-update notifications sent after a subscribe is accepted. diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs index da6cac71..ba8ef9dc 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs @@ -11,10 +11,13 @@ mod tool; mod user_config_store; pub(crate) use auth::token; -pub(crate) use client::{connect_client, create_client, create_tls_client}; +pub(crate) use client::{ + CLIENT_CONNECT_TIMEOUT, TEST_POLL_INTERVAL, connect_client, connect_client_with_handler, create_client, + create_tls_client, +}; pub(crate) use list_tools_gateway::{ ListToolsGatewaySettings, create_gateway_with_four_counters, create_ports, - create_tls_gateway_with_four_tls_counters, + create_tls_gateway_with_four_tls_counters, plaintext_config, }; pub(crate) use plugin::{ POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, TestPlugin, TestPluginFactory,