From bb20bd8d70d32d459f8c6802d9715b9f775d3c17 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 13:04:10 +0100 Subject: [PATCH 1/7] feat: added pagination support for tools, resources and prompts inline with mcp specs Signed-off-by: Lang-Akshay --- .../src/gateway/list_aggregation.rs | 236 ++++++++++++------ .../src/gateway/mcp_service/prompts.rs | 26 +- .../src/gateway/mcp_service/resources.rs | 52 +++- .../src/gateway/mcp_service/tools.rs | 29 ++- .../tests/gateway_pagination.rs | 177 +++++++++++++ .../tests/support/mod.rs | 1 + .../tests/support/paginating_mock.rs | 52 ++++ docs/book/src/mcp-routing-semantics.md | 24 +- 8 files changed, 490 insertions(+), 107 deletions(-) create mode 100644 crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs create mode 100755 crates/contextforge-gateway-rs-lib/tests/support/paginating_mock.rs diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs index 2612f9e..cf44220 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs @@ -1,7 +1,12 @@ +use std::{collections::HashMap, future::Future}; + use contextforge_gateway_rs_apis::user_store::VirtualHost; -use rmcp::model::{ - ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, Prompt, Resource, - ResourceTemplate, Tool, +use rmcp::{ + ErrorData, + model::{ + ErrorCode, ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, + Prompt, Resource, ResourceTemplate, Tool, + }, }; use tracing::{info, warn}; @@ -10,6 +15,33 @@ use super::{ identifier_routing::{exposed_tool_name, prefixed_name}, }; +/// Per-backend cursor state encoded as the gateway's opaque cursor token. +/// Only backends that still have pages appear in `backends`; absent == exhausted. +#[derive(Debug, Default, serde::Serialize, serde::Deserialize)] +pub(super) struct GatewayCursor { + pub(super) backends: HashMap, +} + +/// Decode an incoming gateway cursor (raw JSON, opaque to MCP clients). +/// `None` means first page; an undecodable value returns `-32602 Invalid params`. +pub(super) fn decode_gateway_cursor(raw: Option<&str>) -> Result { + let Some(raw) = raw else { return Ok(GatewayCursor::default()) }; + serde_json::from_str(raw) + .map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None)) +} + +/// Build the next gateway cursor from backends that returned a `next_cursor`. +/// Returns `None` (= no more pages) when all backends are exhausted. +fn encode_next_cursor(backends: HashMap) -> Option { + if backends.is_empty() { + return None; + } + Some( + serde_json::to_string(&GatewayCursor { backends }) + .expect("GatewayCursor is always serializable"), + ) +} + /// Fans a paginated list request out to every connected backend concurrently, logs each response, /// and returns the `(backend_name, result)` pairs that succeeded. pub(super) async fn fan_out_list( @@ -19,8 +51,8 @@ pub(super) async fn fan_out_list( call: F, ) -> Vec<(String, R)> where - F: Fn(McpClientService) -> Fut, - Fut: std::future::Future>, + F: Fn(String, McpClientService) -> Fut, + Fut: Future>, C: Fn(&R) -> usize, E: std::fmt::Debug, { @@ -28,7 +60,7 @@ where let call = &call; async move { let response = match service_holder.running_service { - Some(service) => Some(call(service).await), + Some(service) => Some(call(service_holder.name.clone(), service).await), None => None, }; (service_holder.name, response) @@ -61,126 +93,136 @@ fn log_backend_response( } } -pub(super) fn merge_tools(tools: Vec<(String, ListToolsResult)>, virtual_host: &VirtualHost) -> Vec { - let mut tools = tools - .into_iter() - .flat_map(|(backend_name, result)| { - result - .tools - .into_iter() - .map(|mut tool| { - tool.name = exposed_tool_name(virtual_host, &backend_name, &tool.name).into(); - tool - }) - .collect::>() - }) - .collect::>(); - tools.sort_unstable_by(|tool, other| tool.name.cmp(&other.name)); - tools +pub(super) fn merge_tools( + tools: Vec<(String, ListToolsResult)>, + virtual_host: &VirtualHost, +) -> (Vec, Option) { + let mut next_backends = HashMap::new(); + let mut merged = Vec::new(); + for (backend_name, result) in tools { + if let Some(c) = result.next_cursor { + next_backends.insert(backend_name.clone(), c); + } + for mut tool in result.tools { + tool.name = exposed_tool_name(virtual_host, &backend_name, &tool.name).into(); + merged.push(tool); + } + } + merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + (merged, encode_next_cursor(next_backends)) } pub(super) fn merge_resources( resources: Vec<(String, ListResourcesResult)>, namespace_identifiers: bool, -) -> Vec { - let mut resources = resources - .into_iter() - .flat_map(|(backend_name, result)| { - result - .resources - .into_iter() - .map(|mut resource| { - if namespace_identifiers { - resource.name = prefixed_name(&backend_name, &resource.name); - resource.uri = prefixed_name(&backend_name, &resource.uri); - } - resource - }) - .collect::>() - }) - .collect::>(); - resources.sort_unstable_by(|resource, other| resource.name.cmp(&other.name)); - resources +) -> (Vec, Option) { + let mut next_backends = HashMap::new(); + let mut merged = Vec::new(); + for (backend_name, result) in resources { + if let Some(c) = result.next_cursor { + next_backends.insert(backend_name.clone(), c); + } + for mut resource in result.resources { + if namespace_identifiers { + resource.name = prefixed_name(&backend_name, &resource.name); + resource.uri = prefixed_name(&backend_name, &resource.uri); + } + merged.push(resource); + } + } + merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + (merged, encode_next_cursor(next_backends)) } pub(super) fn merge_resource_templates( templates: Vec<(String, ListResourceTemplatesResult)>, namespace_identifiers: bool, -) -> Vec { - let mut templates = templates - .into_iter() - .flat_map(|(backend_name, result)| { - result.resource_templates.into_iter().map(move |mut template| { - if namespace_identifiers { - template.name = prefixed_name(&backend_name, &template.name); - template.uri_template = prefixed_name(&backend_name, &template.uri_template); - } - template - }) - }) - .collect::>(); - templates.sort_unstable_by(|template, other| template.name.cmp(&other.name)); - templates +) -> (Vec, Option) { + let mut next_backends = HashMap::new(); + let mut merged = Vec::new(); + for (backend_name, result) in templates { + if let Some(c) = result.next_cursor { + next_backends.insert(backend_name.clone(), c); + } + for mut template in result.resource_templates { + if namespace_identifiers { + template.name = prefixed_name(&backend_name, &template.name); + template.uri_template = prefixed_name(&backend_name, &template.uri_template); + } + merged.push(template); + } + } + merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + (merged, encode_next_cursor(next_backends)) } -pub(super) fn merge_prompts(prompts: Vec<(String, ListPromptsResult)>, namespace_identifiers: bool) -> Vec { - let mut prompts = prompts - .into_iter() - .flat_map(|(backend_name, result)| { - result.prompts.into_iter().map(move |mut prompt| { - if namespace_identifiers { - prompt.name = prefixed_name(&backend_name, &prompt.name); - } - prompt - }) - }) - .collect::>(); - prompts.sort_unstable_by(|prompt, other| prompt.name.cmp(&other.name)); - prompts +pub(super) fn merge_prompts( + prompts: Vec<(String, ListPromptsResult)>, + namespace_identifiers: bool, +) -> (Vec, Option) { + let mut next_backends = HashMap::new(); + let mut merged = Vec::new(); + for (backend_name, result) in prompts { + if let Some(c) = result.next_cursor { + next_backends.insert(backend_name.clone(), c); + } + for mut prompt in result.prompts { + if namespace_identifiers { + prompt.name = prefixed_name(&backend_name, &prompt.name); + } + merged.push(prompt); + } + } + merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + (merged, encode_next_cursor(next_backends)) } #[cfg(test)] mod tests { use super::*; - #[test] - fn single_backend_listings_preserve_identifiers() { + fn test_virtual_host(backend_id: &str) -> VirtualHost { let config_json = serde_json::json!({ "backends": { - "backend-id": { + backend_id: { "name": "backend", "url": "http://upstream:9000/mcp", "transport": "STREAMABLEHTTP", "passthrough_headers": [], - "allowed_tool_names": ["test_simple_text"], + "allowed_tool_names": [], "allowed_resource_names": [], "allowed_prompt_names": [] } } }); - let virtual_host: VirtualHost = serde_json::from_value(config_json).expect("valid virtual host"); - let tools = merge_tools( + serde_json::from_value(config_json).expect("valid virtual host") + } + + #[test] + fn single_backend_listings_preserve_identifiers() { + let virtual_host = test_virtual_host("backend-id"); + let (tools, _) = merge_tools( vec![( "backend-id".to_owned(), ListToolsResult::with_all_items(vec![Tool::new("test_simple_text", "", serde_json::Map::new())]), )], &virtual_host, ); - let prompts = merge_prompts( + let (prompts, _) = merge_prompts( vec![( "backend-id".to_owned(), ListPromptsResult::with_all_items(vec![Prompt::new("test_prompt", None::, None)]), )], false, ); - let resources = merge_resources( + let (resources, _) = merge_resources( vec![( "backend-id".to_owned(), ListResourcesResult::with_all_items(vec![Resource::new("test://resource", "test_resource")]), )], false, ); - let templates = merge_resource_templates( + let (templates, _) = merge_resource_templates( vec![( "backend-id".to_owned(), ListResourceTemplatesResult::with_all_items(vec![ResourceTemplate::new( @@ -198,4 +240,40 @@ mod tests { assert_eq!("test_template", templates[0].name); assert_eq!("test://template/{id}/data", templates[0].uri_template); } + + #[test] + fn backend_next_cursor_is_preserved_in_gateway_cursor() { + let virtual_host = test_virtual_host("b1"); + let mut result = + ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); + result.next_cursor = Some("backend-page2".to_owned()); + + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host); + let raw = next_cursor.expect("should have a next cursor"); + + let cursor: GatewayCursor = serde_json::from_str(&raw).expect("valid JSON"); + assert_eq!(cursor.backends.get("b1").map(String::as_str), Some("backend-page2")); + } + + #[test] + fn exhausted_backends_produce_no_next_cursor() { + let virtual_host = test_virtual_host("b1"); + // next_cursor is None → backend is exhausted + let result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); + + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host); + assert!(next_cursor.is_none(), "no cursor when all backends exhausted"); + } + + #[test] + fn invalid_cursor_returns_invalid_params_error() { + let err = decode_gateway_cursor(Some("not-json!")).unwrap_err(); + assert_eq!(err.code, ErrorCode::INVALID_PARAMS); + } + + #[test] + fn none_cursor_returns_default_empty_gateway_cursor() { + let cursor = decode_gateway_cursor(None).expect("None is a valid first-page indicator"); + assert!(cursor.backends.is_empty()); + } } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index 3139d11..87557a0 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -8,7 +8,7 @@ use tracing::info; use super::McpService; use crate::gateway::{ identifier_routing::{backend_forward_error, route_identifier_to_backend}, - list_aggregation::{fan_out_list, merge_prompts}, + list_aggregation::{decode_gateway_cursor, fan_out_list, merge_prompts}, mcp_call_validator::AuthorizedCallValidator, session_manager::SessionManager, session_store::UserSessionStore, @@ -27,20 +27,34 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); - let backend_transports: Vec<_> = session_manager.borrow_transports().await; + let all_transports: Vec<_> = session_manager.borrow_transports().await; + + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { + all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() + } else { + all_transports + }; let responses = fan_out_list( backend_transports, "list_prompts", |response: &ListPromptsResult| response.prompts.len(), - |service| { - let request = request.clone(); - async move { service.list_prompts(request).await } + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service + .list_prompts(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) + .await + } }, ) .await; - Ok(ListPromptsResult::with_all_items(merge_prompts(responses, namespace_identifiers))) + let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers); + let mut result = ListPromptsResult::with_all_items(prompts); + result.next_cursor = next_cursor; + Ok(result) } pub(super) async fn get_prompt( diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index 70cb30e..44e8c4d 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -11,7 +11,7 @@ use tracing::info; use super::McpService; use crate::gateway::{ identifier_routing::{backend_forward_error, route_identifier_to_backend}, - list_aggregation::{fan_out_list, merge_resource_templates, merge_resources}, + list_aggregation::{decode_gateway_cursor, fan_out_list, merge_resource_templates, merge_resources}, mcp_call_validator::AuthorizedCallValidator, session_manager::SessionManager, session_store::UserSessionStore, @@ -30,20 +30,34 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); - let backend_transports: Vec<_> = session_manager.borrow_transports().await; + let all_transports: Vec<_> = session_manager.borrow_transports().await; + + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { + all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() + } else { + all_transports + }; let responses = fan_out_list( backend_transports, "list_resources", |response: &ListResourcesResult| response.resources.len(), - |service| { - let request = request.clone(); - async move { service.list_resources(request).await } + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service + .list_resources(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) + .await + } }, ) .await; - Ok(ListResourcesResult::with_all_items(merge_resources(responses, namespace_identifiers))) + let (resources, next_cursor) = merge_resources(responses, namespace_identifiers); + let mut result = ListResourcesResult::with_all_items(resources); + result.next_cursor = next_cursor; + Ok(result) } pub(super) async fn read_resource( @@ -89,20 +103,36 @@ where let namespace_identifiers = virtual_host.backends.len() > 1; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); - let backend_transports: Vec<_> = session_manager.borrow_transports().await; + let all_transports: Vec<_> = session_manager.borrow_transports().await; + + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { + all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() + } else { + all_transports + }; let responses = fan_out_list( backend_transports, "list_resource_templates", |response: &ListResourceTemplatesResult| response.resource_templates.len(), - |service| { - let request = request.clone(); - async move { service.list_resource_templates(request).await } + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service + .list_resource_templates( + cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))), + ) + .await + } }, ) .await; - Ok(ListResourceTemplatesResult::with_all_items(merge_resource_templates(responses, namespace_identifiers))) + let (resource_templates, next_cursor) = merge_resource_templates(responses, namespace_identifiers); + let mut result = ListResourceTemplatesResult::with_all_items(resource_templates); + result.next_cursor = next_cursor; + Ok(result) } #[expect(deprecated, reason = "temporary RMCP v3 compatibility; subscriptions/listen migration is deferred")] diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs index 50715b0..612a117 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs @@ -10,7 +10,7 @@ use super::McpService; use crate::gateway::{ backend_client::call_backend_tool, identifier_routing::{backend_forward_error, resolve_backend, resolve_tool_route}, - list_aggregation::{fan_out_list, merge_tools}, + list_aggregation::{decode_gateway_cursor, fan_out_list, merge_tools}, mcp_call_validator::AuthorizedCallValidator, session_manager::SessionManager, session_store::UserSessionStore, @@ -27,20 +27,37 @@ where let mcp_call_validator = AuthorizedCallValidator::new("list_tools", &cx); let (virtual_host, session_id, claims) = mcp_call_validator.validate()?; let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); - let backend_transports: Vec<_> = session_manager.borrow_transports().await; + let all_transports: Vec<_> = session_manager.borrow_transports().await; + + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + // On resume, skip backends already exhausted in the prior page. + // ponytail: topology changes between pages silently drop removed backends; + // add a cursor version field if reconfiguration stability matters. + let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { + all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() + } else { + all_transports + }; let responses = fan_out_list( backend_transports, "list_tools", |response: &ListToolsResult| response.tools.len(), - |service| { - let request = request.clone(); - async move { service.list_tools(request).await } + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service + .list_tools(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) + .await + } }, ) .await; - Ok(ListToolsResult::with_all_items(merge_tools(responses, virtual_host))) + let (tools, next_cursor) = merge_tools(responses, virtual_host); + let mut result = ListToolsResult::with_all_items(tools); + result.next_cursor = next_cursor; + Ok(result) } pub(super) async fn call_tool( diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs new file mode 100644 index 0000000..7b1dbc9 --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs @@ -0,0 +1,177 @@ +mod support; + +use std::{collections::HashMap, sync::Arc}; + +use contextforge_gateway_rs_apis::{ + User, + user_store::{BackendMCPGateway, Transport, UserConfig, VirtualHost}, +}; +use contextforge_gateway_rs_lib::{Config, Gateway, Result, UserConfigStore, UserConfigStoreType}; +use rmcp::{ + model::PaginatedRequestParams, + transport::{ + StreamableHttpServerConfig, StreamableHttpService, + streamable_http_server::session::local::LocalSessionManager, + }, +}; +use tracing::warn; + +use support::{ + MemoryUserConfigStore, TEST_USER_ID, connect_client, create_client, create_ports, paginating_mock, + plaintext_config, +}; + +/// Build a single-backend `BackendMCPGateway` pointed at `port`. +fn paginating_backend(port: u16) -> BackendMCPGateway { + BackendMCPGateway { + name: format!("backend-{port}"), + url: format!("http://127.0.0.1:{port}/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 backend_id(port: u16) -> String { + format!("00000000-0000-0000-0000-{port:012}") +} + +/// Start an axum MCP server at `port` serving a `PaginatingServer`. +async fn serve_paginating_backend(port: u16) { + let service = StreamableHttpService::new( + || Ok(paginating_mock::PaginatingServer), + LocalSessionManager::default().into(), + StreamableHttpServerConfig::default(), + ); + let router = axum::Router::new().route_service("/mcp", service); + let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")) + .await + .expect("bind backend"); + axum::serve(listener, router).await.expect("backend server"); +} + +/// Boot the gateway with the given config and user config; return the gateway URL. +async fn start_gateway(config: Config, virtual_host_id: &str, user_config: UserConfig) -> String { + let store = MemoryUserConfigStore::default(); + store.set_config(&User::new(TEST_USER_ID), &user_config).await.expect("set config"); + + let address = config.address.clone().expect("address required"); + let gateway_url = + format!("http://{address}/contextforge-rs/servers/{virtual_host_id}/mcp"); + + let gateway = Gateway::builder() + .with_config(config) + .with_session_manager(Arc::new(LocalSessionManager::default())) + .with_user_config_store_type(UserConfigStoreType::Test(Arc::new(store))) + .build(); + + tokio::spawn(async move { + let res = gateway.run_gateway().await; + warn!("Gateway exited {res:?}"); + }); + + gateway_url +} + +/// A paginating backend returns tools across two pages; the gateway must expose +/// all of them to the client without any items being silently dropped. +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[test_log::test] +async fn single_backend_pagination_all_tools_reachable() -> Result<()> { + let ports = create_ports(2); + let (backend_port, gateway_port) = (ports[0], ports[1]); + let config = plaintext_config(gateway_port); + + let virtual_host_id = "22222222-2222-2222-2222-222222222222"; + let backends = HashMap::from([(backend_id(backend_port), paginating_backend(backend_port))]); + let user_config = + UserConfig { virtual_hosts: HashMap::from([(virtual_host_id.to_owned(), VirtualHost { backends })]) }; + + tokio::spawn(serve_paginating_backend(backend_port)); + let gateway_url = start_gateway(config, virtual_host_id, user_config).await; + + let svc = connect_client(gateway_url, create_client(TEST_USER_ID)).await?; + + // Page 1 + let page1 = svc.list_tools(None).await.expect("page 1"); + let page1_names: Vec<&str> = page1.tools.iter().map(|t| t.name.as_ref()).collect(); + assert!(page1.next_cursor.is_some(), "page 1 must carry a next_cursor"); + assert_eq!(page1_names, ["tool_alpha", "tool_beta"]); + + // Page 2 + let cursor = + page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); + let page2 = svc.list_tools(cursor).await.expect("page 2"); + let page2_names: Vec<&str> = page2.tools.iter().map(|t| t.name.as_ref()).collect(); + assert!(page2.next_cursor.is_none(), "page 2 must be the final page"); + assert_eq!(page2_names, ["tool_gamma"]); + + // All tools reachable with no duplication + let mut all_names = page1_names.clone(); + all_names.extend_from_slice(&page2_names); + all_names.sort_unstable(); + assert_eq!(all_names, paginating_mock::PaginatingServer::all_tool_names()); + + Ok(()) +} + +/// When one backend exhausts its pages, it must be excluded from the resume +/// request. Without the filter, the exhausted backend would be re-queried and +/// its tools would appear in every subsequent page as duplicates. +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[test_log::test] +async fn multi_backend_exhausted_backend_not_requeried() -> Result<()> { + // Backend A: PaginatingServer (2 pages: 2 tools + 1 tool) + // Backend B: another PaginatingServer (same 2 pages, different backend ID) + // + // With 2 backends, tool names get the backend-ID prefix. + // Page 1: 2 tools from A page 1 + 2 tools from B page 1 = 4 total + // Page 2: 1 tool from A page 2 + 1 tool from B page 2 = 2 total + // If either backend were re-queried, its page-1 tools would reappear. + let ports = create_ports(3); + let (port_a, port_b, gateway_port) = (ports[0], ports[1], ports[2]); + let config = plaintext_config(gateway_port); + + let virtual_host_id = "33333333-3333-3333-3333-333333333333"; + let backends = HashMap::from([ + (backend_id(port_a), paginating_backend(port_a)), + (backend_id(port_b), paginating_backend(port_b)), + ]); + let user_config = + UserConfig { virtual_hosts: HashMap::from([(virtual_host_id.to_owned(), VirtualHost { backends })]) }; + + tokio::spawn(serve_paginating_backend(port_a)); + tokio::spawn(serve_paginating_backend(port_b)); + let gateway_url = start_gateway(config, virtual_host_id, user_config).await; + + let svc = connect_client(gateway_url, create_client(TEST_USER_ID)).await?; + + // Page 1: both backends contribute their first page (2 tools each) + let page1 = svc.list_tools(None).await.expect("page 1"); + assert!(page1.next_cursor.is_some(), "page 1 must carry a next_cursor"); + assert_eq!(page1.tools.len(), 4, "page 1 should have 2 tools from each backend"); + + // Page 2: both backends contribute their second page (1 tool each) + let cursor = + page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); + let page2 = svc.list_tools(cursor).await.expect("page 2"); + assert!(page2.next_cursor.is_none(), "page 2 must be the final page"); + assert_eq!(page2.tools.len(), 2, "page 2 should have 1 tool from each backend"); + + // Union has 6 unique tools, no duplicates + let mut all_names: Vec<_> = page1.tools.iter().chain(page2.tools.iter()).map(|t| t.name.clone()).collect(); + all_names.sort_unstable(); + all_names.dedup(); + assert_eq!( + all_names.len(), + page1.tools.len() + page2.tools.len(), + "no duplicate tools across pages" + ); + + Ok(()) +} diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs index 37db231..54f6898 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs @@ -4,6 +4,7 @@ mod auth; mod client; mod list_tools_gateway; pub(crate) mod mock_counter; +pub(crate) mod paginating_mock; mod plugin; mod plugin_gateway; mod runtime; diff --git a/crates/contextforge-gateway-rs-lib/tests/support/paginating_mock.rs b/crates/contextforge-gateway-rs-lib/tests/support/paginating_mock.rs new file mode 100755 index 0000000..0e8b206 --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/tests/support/paginating_mock.rs @@ -0,0 +1,52 @@ +#![allow(dead_code, reason = "shared test fixture used by separate integration test targets")] + +use rmcp::{ErrorData as McpError, RoleServer, ServerHandler, model::*, service::RequestContext}; + +/// Backend cursor the mock uses to signal "page 2 available". +const PAGE2_CURSOR: &str = "page2"; + +/// A minimal backend that returns tools across two pages. +/// +/// Page 1: `tool_alpha`, `tool_beta` → `next_cursor = Some("page2")` +/// Page 2: `tool_gamma` → `next_cursor = None` +pub struct PaginatingServer; + +impl PaginatingServer { + pub fn page1_tools() -> Vec { + vec![ + Tool::new("tool_alpha", "Alpha tool", serde_json::Map::new()), + Tool::new("tool_beta", "Beta tool", serde_json::Map::new()), + ] + } + + pub fn page2_tools() -> Vec { + vec![Tool::new("tool_gamma", "Gamma tool", serde_json::Map::new())] + } + + /// Sorted union of all tool names across both pages. + pub fn all_tool_names() -> &'static [&'static str] { + &["tool_alpha", "tool_beta", "tool_gamma"] + } +} + +impl ServerHandler for PaginatingServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::from_build_env()) + .with_protocol_version(ProtocolVersion::V_2024_11_05) + } + + async fn list_tools( + &self, + request: Option, + _: RequestContext, + ) -> Result { + if request.as_ref().and_then(|r| r.cursor.as_deref()) == Some(PAGE2_CURSOR) { + Ok(ListToolsResult::with_all_items(Self::page2_tools())) + } else { + let mut result = ListToolsResult::with_all_items(Self::page1_tools()); + result.next_cursor = Some(PAGE2_CURSOR.to_owned()); + Ok(result) + } + } +} diff --git a/docs/book/src/mcp-routing-semantics.md b/docs/book/src/mcp-routing-semantics.md index 5b2923b..079d5ad 100644 --- a/docs/book/src/mcp-routing-semantics.md +++ b/docs/book/src/mcp-routing-semantics.md @@ -86,12 +86,26 @@ service for the principal and downstream session. Missing backends fail the call. Duplicate matches are treated as invalid session state and trigger backend cleanup. -## Known Gaps +## Federated Pagination + +All four list operations support cursor-based pagination across backends. + +The gateway wraps per-backend cursors inside its own opaque token (JSON serialised, +treated as an opaque string by MCP clients per spec). On the first request (no cursor) +every backend is queried. On a resume request the gateway decodes its cursor to +recover per-backend positions, skips backends that have already been exhausted, and +forwards each remaining backend its own cursor. Results from the current page are +merged and sorted. A new gateway cursor is emitted when at least one backend still +has more pages; the cursor is omitted once all backends are exhausted. -Pagination is not complete. `list_tools`, `list_resources`, `list_prompts`, -and `list_resource_templates` currently perform one backend call and return a -merged response with no downstream cursor. Full parity needs to gather all -backend pages or define a merged cursor strategy. +An undecodable cursor returns `-32602 Invalid params` (MCP spec §Pagination). + +**Known limitation:** if the backend set changes between pages (reconfiguration), the +cursor for a removed backend is silently dropped and its remaining items are not +returned. Add a cursor version field if stability under live reconfiguration is +required. + +## Known Gaps Streaming/SSE behavior is also still a tracked design area. The target is to stream downstream as backend chunks arrive while preserving plugin behavior, From 73700bfb4d17b9dd08b2bdeb7142cf83c3521883 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 13:17:03 +0100 Subject: [PATCH 2/7] fixup: fmt and clippy Signed-off-by: Lang-Akshay --- .../src/gateway/list_aggregation.rs | 15 ++++------- .../src/gateway/mcp_service/prompts.rs | 27 +++++++++---------- .../src/gateway/mcp_service/resources.rs | 8 ++---- .../src/gateway/mcp_service/tools.rs | 27 +++++++++---------- .../tests/gateway_pagination.rs | 27 ++++++------------- 5 files changed, 41 insertions(+), 63 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs index cf44220..3ebcdb7 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs @@ -4,8 +4,8 @@ use contextforge_gateway_rs_apis::user_store::VirtualHost; use rmcp::{ ErrorData, model::{ - ErrorCode, ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, - Prompt, Resource, ResourceTemplate, Tool, + ErrorCode, ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, Prompt, + Resource, ResourceTemplate, Tool, }, }; use tracing::{info, warn}; @@ -26,8 +26,7 @@ pub(super) struct GatewayCursor { /// `None` means first page; an undecodable value returns `-32602 Invalid params`. pub(super) fn decode_gateway_cursor(raw: Option<&str>) -> Result { let Some(raw) = raw else { return Ok(GatewayCursor::default()) }; - serde_json::from_str(raw) - .map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None)) + serde_json::from_str(raw).map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None)) } /// Build the next gateway cursor from backends that returned a `next_cursor`. @@ -36,10 +35,7 @@ fn encode_next_cursor(backends: HashMap) -> Option { if backends.is_empty() { return None; } - Some( - serde_json::to_string(&GatewayCursor { backends }) - .expect("GatewayCursor is always serializable"), - ) + Some(serde_json::to_string(&GatewayCursor { backends }).expect("GatewayCursor is always serializable")) } /// Fans a paginated list request out to every connected backend concurrently, logs each response, @@ -244,8 +240,7 @@ mod tests { #[test] fn backend_next_cursor_is_preserved_in_gateway_cursor() { let virtual_host = test_virtual_host("b1"); - let mut result = - ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); + let mut result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); result.next_cursor = Some("backend-page2".to_owned()); let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host); diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index 87557a0..220b1b8 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -36,20 +36,19 @@ where all_transports }; - let responses = fan_out_list( - backend_transports, - "list_prompts", - |response: &ListPromptsResult| response.prompts.len(), - |name, service| { - let cursor = gateway_cursor.backends.get(&name).cloned(); - async move { - service - .list_prompts(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) - .await - } - }, - ) - .await; + let responses = + fan_out_list( + backend_transports, + "list_prompts", + |response: &ListPromptsResult| response.prompts.len(), + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service.list_prompts(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await + } + }, + ) + .await; let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers); let mut result = ListPromptsResult::with_all_items(prompts); diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index 44e8c4d..c5bfe8a 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -46,9 +46,7 @@ where |name, service| { let cursor = gateway_cursor.backends.get(&name).cloned(); async move { - service - .list_resources(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) - .await + service.list_resources(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await } }, ) @@ -120,9 +118,7 @@ where let cursor = gateway_cursor.backends.get(&name).cloned(); async move { service - .list_resource_templates( - cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))), - ) + .list_resource_templates(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) .await } }, diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs index 612a117..1fadd9d 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs @@ -39,20 +39,19 @@ where all_transports }; - let responses = fan_out_list( - backend_transports, - "list_tools", - |response: &ListToolsResult| response.tools.len(), - |name, service| { - let cursor = gateway_cursor.backends.get(&name).cloned(); - async move { - service - .list_tools(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) - .await - } - }, - ) - .await; + let responses = + fan_out_list( + backend_transports, + "list_tools", + |response: &ListToolsResult| response.tools.len(), + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + async move { + service.list_tools(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await + } + }, + ) + .await; let (tools, next_cursor) = merge_tools(responses, virtual_host); let mut result = ListToolsResult::with_all_items(tools); diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs index 7b1dbc9..59a0071 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs @@ -10,15 +10,13 @@ use contextforge_gateway_rs_lib::{Config, Gateway, Result, UserConfigStore, User use rmcp::{ model::PaginatedRequestParams, transport::{ - StreamableHttpServerConfig, StreamableHttpService, - streamable_http_server::session::local::LocalSessionManager, + StreamableHttpServerConfig, StreamableHttpService, streamable_http_server::session::local::LocalSessionManager, }, }; use tracing::warn; use support::{ - MemoryUserConfigStore, TEST_USER_ID, connect_client, create_client, create_ports, paginating_mock, - plaintext_config, + MemoryUserConfigStore, TEST_USER_ID, connect_client, create_client, create_ports, paginating_mock, plaintext_config, }; /// Build a single-backend `BackendMCPGateway` pointed at `port`. @@ -49,9 +47,7 @@ async fn serve_paginating_backend(port: u16) { StreamableHttpServerConfig::default(), ); let router = axum::Router::new().route_service("/mcp", service); - let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")) - .await - .expect("bind backend"); + let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")).await.expect("bind backend"); axum::serve(listener, router).await.expect("backend server"); } @@ -60,9 +56,8 @@ async fn start_gateway(config: Config, virtual_host_id: &str, user_config: UserC let store = MemoryUserConfigStore::default(); store.set_config(&User::new(TEST_USER_ID), &user_config).await.expect("set config"); - let address = config.address.clone().expect("address required"); - let gateway_url = - format!("http://{address}/contextforge-rs/servers/{virtual_host_id}/mcp"); + let address = config.address.expect("address required"); + let gateway_url = format!("http://{address}/contextforge-rs/servers/{virtual_host_id}/mcp"); let gateway = Gateway::builder() .with_config(config) @@ -104,8 +99,7 @@ async fn single_backend_pagination_all_tools_reachable() -> Result<()> { assert_eq!(page1_names, ["tool_alpha", "tool_beta"]); // Page 2 - let cursor = - page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); + let cursor = page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); let page2 = svc.list_tools(cursor).await.expect("page 2"); let page2_names: Vec<&str> = page2.tools.iter().map(|t| t.name.as_ref()).collect(); assert!(page2.next_cursor.is_none(), "page 2 must be the final page"); @@ -157,8 +151,7 @@ async fn multi_backend_exhausted_backend_not_requeried() -> Result<()> { assert_eq!(page1.tools.len(), 4, "page 1 should have 2 tools from each backend"); // Page 2: both backends contribute their second page (1 tool each) - let cursor = - page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); + let cursor = page1.next_cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c))); let page2 = svc.list_tools(cursor).await.expect("page 2"); assert!(page2.next_cursor.is_none(), "page 2 must be the final page"); assert_eq!(page2.tools.len(), 2, "page 2 should have 1 tool from each backend"); @@ -167,11 +160,7 @@ async fn multi_backend_exhausted_backend_not_requeried() -> Result<()> { let mut all_names: Vec<_> = page1.tools.iter().chain(page2.tools.iter()).map(|t| t.name.clone()).collect(); all_names.sort_unstable(); all_names.dedup(); - assert_eq!( - all_names.len(), - page1.tools.len() + page2.tools.len(), - "no duplicate tools across pages" - ); + assert_eq!(all_names.len(), page1.tools.len() + page2.tools.len(), "no duplicate tools across pages"); Ok(()) } From 4573499c38fb9e208b53a064a4c7c17edafdff4d Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 13:25:15 +0100 Subject: [PATCH 3/7] feat: updated book to capture new request-flow Signed-off-by: Lang-Akshay --- docs/book/src/mcp-method-reference.md | 2 +- docs/book/src/request-flow.md | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/book/src/mcp-method-reference.md b/docs/book/src/mcp-method-reference.md index 061e36f..d378a85 100644 --- a/docs/book/src/mcp-method-reference.md +++ b/docs/book/src/mcp-method-reference.md @@ -33,7 +33,7 @@ share one fanout path: | Identifiers | A single backend preserves upstream identifiers; multiple backends use the backend map-key prefix. Explicit control-plane tool aliases are returned exactly as configured. Resource templates apply the same single-versus-multi rule to both names and URI templates. | | Ordering | Merged output is sorted by name. | | Failures | Failed or unavailable backends are logged and omitted from the merged result. | -| Pagination | One backend call per request and no downstream cursor; see [Known Gaps](mcp-routing-semantics.md#known-gaps). | +| Pagination | Cursor-based, per the [MCP spec §Pagination](https://modelcontextprotocol.io/specification/2026-07-28/server/utilities/pagination). The gateway cursor is an opaque JSON token encoding each active backend's position. On the first request (no cursor) all backends are queried; on resume only backends with remaining pages are queried. An undecodable cursor returns `−32602 Invalid params`. See [Federated Pagination](mcp-routing-semantics.md#federated-pagination). | ## Routed Targeted Methods diff --git a/docs/book/src/request-flow.md b/docs/book/src/request-flow.md index 09abf09..f4320b1 100644 --- a/docs/book/src/request-flow.md +++ b/docs/book/src/request-flow.md @@ -160,7 +160,7 @@ Current routed method families: | Method family | Flow | | --- | --- | -| `list_tools`, `list_resources`, `list_prompts`, `list_resource_templates` | Borrow all configured backend services, call every available backend concurrently with `fan_out_list`, preserve identifiers for one backend or namespace them for multiple backends, sort merged output, and return one list. Explicit tool aliases are preserved exactly. | +| `list_tools`, `list_resources`, `list_prompts`, `list_resource_templates` | Decode the incoming gateway cursor (if present) to recover per-backend positions. On the first request (no cursor) all configured backend services are queried; on a resume request only backends that still have pages are queried. Call each selected backend concurrently via `fan_out_list`, passing its own per-backend cursor. Preserve identifiers for a single backend or namespace them for multiple backends. Sort the merged page output. Encode a new gateway cursor when at least one backend returned a `next_cursor`; omit it when all are exhausted. Explicit tool aliases are preserved exactly. An undecodable cursor returns `−32602 Invalid params`. | | `call_tool` | Resolve an exact tool alias first; otherwise preserve the name for one backend or split `{backend_name}-{tool_name}` for multiple backends. Resolve one backend, run the optional tool hooks, track progress, call the backend, and return the result. | | `read_resource`, `subscribe`, `unsubscribe`, `get_prompt` | Select the only backend and forward the identifier unchanged, or split and strip the prefix for a multi-backend host, then call the resolved backend. | | `complete` | Apply the same conditional routing to the prompt name or resource URI in `ref`, then return the selected backend's completion result. | @@ -193,8 +193,12 @@ backend fanout or identifier routing. ## Response Path Backend responses return to `McpService` first. `call_tool` may run response -plugin hooks before returning. List calls merge backend output, preserving -single-backend identifiers and namespacing multi-backend identifiers as needed. +plugin hooks before returning. + +List calls decode the incoming gateway cursor, fan out to the active backends, +merge the current page of output (preserving single-backend identifiers and +namespacing multi-backend identifiers as needed), and encode a new gateway +cursor when more pages remain across any backend. The HTTP response then unwinds through `virtual_host_config_layer`, `user_config_store_layer`, `session_id_layer`, `claims_layer`, From c782021c97cf08bc89a3dcd93fc2a4ef184fc1a9 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 13:52:07 +0100 Subject: [PATCH 4/7] feat: preserve request._meta and fixed backend readiness race Signed-off-by: Lang-Akshay --- .../src/gateway/mcp_service/prompts.rs | 11 ++++++++- .../src/gateway/mcp_service/resources.rs | 24 +++++++++++++++---- .../src/gateway/mcp_service/tools.rs | 11 ++++++++- .../tests/gateway_pagination.rs | 20 +++++++++++----- 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index 220b1b8..a0da384 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -43,8 +43,17 @@ where |response: &ListPromptsResult| response.prompts.len(), |name, service| { let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); async move { - service.list_prompts(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + } + None => req, + }; + service.list_prompts(backend_req).await } }, ) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index c5bfe8a..0f65a81 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -45,8 +45,17 @@ where |response: &ListResourcesResult| response.resources.len(), |name, service| { let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); async move { - service.list_resources(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + } + None => req, + }; + service.list_resources(backend_req).await } }, ) @@ -116,10 +125,17 @@ where |response: &ListResourceTemplatesResult| response.resource_templates.len(), |name, service| { let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); async move { - service - .list_resource_templates(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))) - .await + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + } + None => req, + }; + service.list_resource_templates(backend_req).await } }, ) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs index 1fadd9d..baa1c84 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs @@ -46,8 +46,17 @@ where |response: &ListToolsResult| response.tools.len(), |name, service| { let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); async move { - service.list_tools(cursor.map(|c| PaginatedRequestParams::default().with_cursor(Some(c)))).await + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + } + None => req, + }; + service.list_tools(backend_req).await } }, ) diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs index 59a0071..4f4679d 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_pagination.rs @@ -39,15 +39,20 @@ fn backend_id(port: u16) -> String { format!("00000000-0000-0000-0000-{port:012}") } -/// Start an axum MCP server at `port` serving a `PaginatingServer`. -async fn serve_paginating_backend(port: u16) { +/// Bind the TCP port for a backend; returns the ready listener. +/// Call this *before* `tokio::spawn` so the port is reserved before the test proceeds. +async fn bind_backend_port(port: u16) -> tokio::net::TcpListener { + tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")).await.expect("bind backend") +} + +/// Start an axum MCP server on an already-bound listener serving a `PaginatingServer`. +async fn serve_paginating_backend(listener: tokio::net::TcpListener) { let service = StreamableHttpService::new( || Ok(paginating_mock::PaginatingServer), LocalSessionManager::default().into(), StreamableHttpServerConfig::default(), ); let router = axum::Router::new().route_service("/mcp", service); - let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")).await.expect("bind backend"); axum::serve(listener, router).await.expect("backend server"); } @@ -87,7 +92,8 @@ async fn single_backend_pagination_all_tools_reachable() -> Result<()> { let user_config = UserConfig { virtual_hosts: HashMap::from([(virtual_host_id.to_owned(), VirtualHost { backends })]) }; - tokio::spawn(serve_paginating_backend(backend_port)); + let backend_listener = bind_backend_port(backend_port).await; + tokio::spawn(serve_paginating_backend(backend_listener)); let gateway_url = start_gateway(config, virtual_host_id, user_config).await; let svc = connect_client(gateway_url, create_client(TEST_USER_ID)).await?; @@ -139,8 +145,10 @@ async fn multi_backend_exhausted_backend_not_requeried() -> Result<()> { let user_config = UserConfig { virtual_hosts: HashMap::from([(virtual_host_id.to_owned(), VirtualHost { backends })]) }; - tokio::spawn(serve_paginating_backend(port_a)); - tokio::spawn(serve_paginating_backend(port_b)); + let listener_a = bind_backend_port(port_a).await; + let listener_b = bind_backend_port(port_b).await; + tokio::spawn(serve_paginating_backend(listener_a)); + tokio::spawn(serve_paginating_backend(listener_b)); let gateway_url = start_gateway(config, virtual_host_id, user_config).await; let svc = connect_client(gateway_url, create_client(TEST_USER_ID)).await?; From cee3ca881c4ae826665fa5fd51422214704d0e21 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 15:06:35 +0100 Subject: [PATCH 5/7] fixup: fmt Signed-off-by: Lang-Akshay --- .../src/gateway/mcp_service/prompts.rs | 43 +++++++++---------- .../src/gateway/mcp_service/resources.rs | 4 +- .../src/gateway/mcp_service/tools.rs | 43 +++++++++---------- 3 files changed, 44 insertions(+), 46 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index a0da384..fcc13f6 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -36,28 +36,27 @@ where all_transports }; - let responses = - fan_out_list( - backend_transports, - "list_prompts", - |response: &ListPromptsResult| response.prompts.len(), - |name, service| { - let cursor = gateway_cursor.backends.get(&name).cloned(); - let req = request.clone(); - async move { - let backend_req = match cursor { - Some(c) => { - let mut r = req.unwrap_or_default(); - r.cursor = Some(c); - Some(r) - } - None => req, - }; - service.list_prompts(backend_req).await - } - }, - ) - .await; + let responses = fan_out_list( + backend_transports, + "list_prompts", + |response: &ListPromptsResult| response.prompts.len(), + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); + async move { + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + }, + None => req, + }; + service.list_prompts(backend_req).await + } + }, + ) + .await; let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers); let mut result = ListPromptsResult::with_all_items(prompts); diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index 0f65a81..8c98974 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -52,7 +52,7 @@ where let mut r = req.unwrap_or_default(); r.cursor = Some(c); Some(r) - } + }, None => req, }; service.list_resources(backend_req).await @@ -132,7 +132,7 @@ where let mut r = req.unwrap_or_default(); r.cursor = Some(c); Some(r) - } + }, None => req, }; service.list_resource_templates(backend_req).await diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs index baa1c84..349fc49 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs @@ -39,28 +39,27 @@ where all_transports }; - let responses = - fan_out_list( - backend_transports, - "list_tools", - |response: &ListToolsResult| response.tools.len(), - |name, service| { - let cursor = gateway_cursor.backends.get(&name).cloned(); - let req = request.clone(); - async move { - let backend_req = match cursor { - Some(c) => { - let mut r = req.unwrap_or_default(); - r.cursor = Some(c); - Some(r) - } - None => req, - }; - service.list_tools(backend_req).await - } - }, - ) - .await; + let responses = fan_out_list( + backend_transports, + "list_tools", + |response: &ListToolsResult| response.tools.len(), + |name, service| { + let cursor = gateway_cursor.backends.get(&name).cloned(); + let req = request.clone(); + async move { + let backend_req = match cursor { + Some(c) => { + let mut r = req.unwrap_or_default(); + r.cursor = Some(c); + Some(r) + }, + None => req, + }; + service.list_tools(backend_req).await + } + }, + ) + .await; let (tools, next_cursor) = merge_tools(responses, virtual_host); let mut result = ListToolsResult::with_all_items(tools); From b9159a18260d05611f0c0b342db547f9e8954bd0 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 16:10:28 +0100 Subject: [PATCH 6/7] feat: resolved backend timeout during resume page and gateway cursor casrries no operation discriminator Signed-off-by: Lang-Akshay --- .../src/gateway/list_aggregation.rs | 194 +++++++++++++++--- .../src/gateway/mcp_service/prompts.rs | 4 +- .../src/gateway/mcp_service/resources.rs | 8 +- .../src/gateway/mcp_service/tools.rs | 4 +- 4 files changed, 170 insertions(+), 40 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs index 3ebcdb7..8f41912 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, future::Future}; +use std::{collections::{HashMap, HashSet}, future::Future}; use contextforge_gateway_rs_apis::user_store::VirtualHost; use rmcp::{ @@ -16,26 +16,50 @@ use super::{ }; /// Per-backend cursor state encoded as the gateway's opaque cursor token. -/// Only backends that still have pages appear in `backends`; absent == exhausted. -#[derive(Debug, Default, serde::Serialize, serde::Deserialize)] +/// +/// `op` names the list operation that issued this cursor (e.g. `"list_tools"`); a cursor +/// presented to a different operation is rejected with `-32602`. +/// +/// Only backends that still have pages **or that failed mid-page** appear in `backends`. +/// Backends that returned `next_cursor = None` are absent (truly exhausted). +#[derive(Debug, serde::Serialize, serde::Deserialize)] pub(super) struct GatewayCursor { + pub(super) op: String, pub(super) backends: HashMap, } +impl Default for GatewayCursor { + /// Returns the first-page sentinel. `op` is empty; it is never validated because + /// `decode_gateway_cursor` returns this value only when `raw` is `None`. + fn default() -> Self { + Self { op: String::new(), backends: HashMap::new() } + } +} + /// Decode an incoming gateway cursor (raw JSON, opaque to MCP clients). /// `None` means first page; an undecodable value returns `-32602 Invalid params`. -pub(super) fn decode_gateway_cursor(raw: Option<&str>) -> Result { +/// `expected_op` must match the `op` field stored in the cursor, preventing a cursor +/// issued by one list operation from being accepted by a different one. +pub(super) fn decode_gateway_cursor(raw: Option<&str>, expected_op: &str) -> Result { let Some(raw) = raw else { return Ok(GatewayCursor::default()) }; - serde_json::from_str(raw).map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None)) + let cursor: GatewayCursor = serde_json::from_str(raw) + .map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None))?; + if cursor.op != expected_op { + return Err(ErrorData::new(ErrorCode::INVALID_PARAMS, "cursor operation mismatch", None)); + } + Ok(cursor) } -/// Build the next gateway cursor from backends that returned a `next_cursor`. +/// Build the next gateway cursor from backends that still have pages. /// Returns `None` (= no more pages) when all backends are exhausted. -fn encode_next_cursor(backends: HashMap) -> Option { +fn encode_next_cursor(op: &str, backends: HashMap) -> Option { if backends.is_empty() { return None; } - Some(serde_json::to_string(&GatewayCursor { backends }).expect("GatewayCursor is always serializable")) + Some( + serde_json::to_string(&GatewayCursor { op: op.to_owned(), backends }) + .expect("GatewayCursor is always serializable"), + ) } /// Fans a paginated list request out to every connected backend concurrently, logs each response, @@ -89,35 +113,65 @@ fn log_backend_response( } } +/// Merges tool listings from multiple backends into a single sorted list. +/// +/// `op` is the list operation name used to tag the outgoing cursor. +/// `incoming_cursor` is used to re-emit the prior cursor position for any backend +/// that was expected to continue but failed on this page (transient failure ≠ exhaustion). pub(super) fn merge_tools( tools: Vec<(String, ListToolsResult)>, virtual_host: &VirtualHost, + incoming_cursor: &GatewayCursor, + op: &str, ) -> (Vec, Option) { - let mut next_backends = HashMap::new(); + let mut next_backends: HashMap = HashMap::new(); + let mut succeeded: HashSet<&str> = HashSet::new(); let mut merged = Vec::new(); - for (backend_name, result) in tools { - if let Some(c) = result.next_cursor { - next_backends.insert(backend_name.clone(), c); + for (backend_name, result) in &tools { + succeeded.insert(backend_name.as_str()); + if let Some(c) = &result.next_cursor { + next_backends.insert(backend_name.clone(), c.clone()); } + } + // Preserve cursor for backends that were expected to continue but failed this page. + for (backend_name, prior_cursor) in &incoming_cursor.backends { + if !succeeded.contains(backend_name.as_str()) { + next_backends.entry(backend_name.clone()).or_insert_with(|| prior_cursor.clone()); + } + } + for (backend_name, result) in tools { for mut tool in result.tools { tool.name = exposed_tool_name(virtual_host, &backend_name, &tool.name).into(); merged.push(tool); } } merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); - (merged, encode_next_cursor(next_backends)) + (merged, encode_next_cursor(op, next_backends)) } +/// Merges resource listings from multiple backends into a single sorted list. +/// See `merge_tools` for the `incoming_cursor` / `op` contract. pub(super) fn merge_resources( resources: Vec<(String, ListResourcesResult)>, namespace_identifiers: bool, + incoming_cursor: &GatewayCursor, + op: &str, ) -> (Vec, Option) { - let mut next_backends = HashMap::new(); + let mut next_backends: HashMap = HashMap::new(); + let mut succeeded: HashSet<&str> = HashSet::new(); let mut merged = Vec::new(); - for (backend_name, result) in resources { - if let Some(c) = result.next_cursor { - next_backends.insert(backend_name.clone(), c); + for (backend_name, result) in &resources { + succeeded.insert(backend_name.as_str()); + if let Some(c) = &result.next_cursor { + next_backends.insert(backend_name.clone(), c.clone()); + } + } + for (backend_name, prior_cursor) in &incoming_cursor.backends { + if !succeeded.contains(backend_name.as_str()) { + next_backends.entry(backend_name.clone()).or_insert_with(|| prior_cursor.clone()); } + } + for (backend_name, result) in resources { for mut resource in result.resources { if namespace_identifiers { resource.name = prefixed_name(&backend_name, &resource.name); @@ -127,19 +181,32 @@ pub(super) fn merge_resources( } } merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); - (merged, encode_next_cursor(next_backends)) + (merged, encode_next_cursor(op, next_backends)) } +/// Merges resource-template listings from multiple backends into a single sorted list. +/// See `merge_tools` for the `incoming_cursor` / `op` contract. pub(super) fn merge_resource_templates( templates: Vec<(String, ListResourceTemplatesResult)>, namespace_identifiers: bool, + incoming_cursor: &GatewayCursor, + op: &str, ) -> (Vec, Option) { - let mut next_backends = HashMap::new(); + let mut next_backends: HashMap = HashMap::new(); + let mut succeeded: HashSet<&str> = HashSet::new(); let mut merged = Vec::new(); - for (backend_name, result) in templates { - if let Some(c) = result.next_cursor { - next_backends.insert(backend_name.clone(), c); + for (backend_name, result) in &templates { + succeeded.insert(backend_name.as_str()); + if let Some(c) = &result.next_cursor { + next_backends.insert(backend_name.clone(), c.clone()); } + } + for (backend_name, prior_cursor) in &incoming_cursor.backends { + if !succeeded.contains(backend_name.as_str()) { + next_backends.entry(backend_name.clone()).or_insert_with(|| prior_cursor.clone()); + } + } + for (backend_name, result) in templates { for mut template in result.resource_templates { if namespace_identifiers { template.name = prefixed_name(&backend_name, &template.name); @@ -149,19 +216,32 @@ pub(super) fn merge_resource_templates( } } merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); - (merged, encode_next_cursor(next_backends)) + (merged, encode_next_cursor(op, next_backends)) } +/// Merges prompt listings from multiple backends into a single sorted list. +/// See `merge_tools` for the `incoming_cursor` / `op` contract. pub(super) fn merge_prompts( prompts: Vec<(String, ListPromptsResult)>, namespace_identifiers: bool, + incoming_cursor: &GatewayCursor, + op: &str, ) -> (Vec, Option) { - let mut next_backends = HashMap::new(); + let mut next_backends: HashMap = HashMap::new(); + let mut succeeded: HashSet<&str> = HashSet::new(); let mut merged = Vec::new(); - for (backend_name, result) in prompts { - if let Some(c) = result.next_cursor { - next_backends.insert(backend_name.clone(), c); + for (backend_name, result) in &prompts { + succeeded.insert(backend_name.as_str()); + if let Some(c) = &result.next_cursor { + next_backends.insert(backend_name.clone(), c.clone()); + } + } + for (backend_name, prior_cursor) in &incoming_cursor.backends { + if !succeeded.contains(backend_name.as_str()) { + next_backends.entry(backend_name.clone()).or_insert_with(|| prior_cursor.clone()); } + } + for (backend_name, result) in prompts { for mut prompt in result.prompts { if namespace_identifiers { prompt.name = prefixed_name(&backend_name, &prompt.name); @@ -170,7 +250,7 @@ pub(super) fn merge_prompts( } } merged.sort_unstable_by(|a, b| a.name.cmp(&b.name)); - (merged, encode_next_cursor(next_backends)) + (merged, encode_next_cursor(op, next_backends)) } #[cfg(test)] @@ -203,6 +283,8 @@ mod tests { ListToolsResult::with_all_items(vec![Tool::new("test_simple_text", "", serde_json::Map::new())]), )], &virtual_host, + &GatewayCursor::default(), + "list_tools", ); let (prompts, _) = merge_prompts( vec![( @@ -210,6 +292,8 @@ mod tests { ListPromptsResult::with_all_items(vec![Prompt::new("test_prompt", None::, None)]), )], false, + &GatewayCursor::default(), + "list_prompts", ); let (resources, _) = merge_resources( vec![( @@ -217,6 +301,8 @@ mod tests { ListResourcesResult::with_all_items(vec![Resource::new("test://resource", "test_resource")]), )], false, + &GatewayCursor::default(), + "list_resources", ); let (templates, _) = merge_resource_templates( vec![( @@ -227,6 +313,8 @@ mod tests { )]), )], false, + &GatewayCursor::default(), + "list_resource_templates", ); assert_eq!("test_simple_text", tools[0].name); @@ -243,10 +331,11 @@ mod tests { let mut result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); result.next_cursor = Some("backend-page2".to_owned()); - let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host); + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); let raw = next_cursor.expect("should have a next cursor"); let cursor: GatewayCursor = serde_json::from_str(&raw).expect("valid JSON"); + assert_eq!(cursor.op, "list_tools"); assert_eq!(cursor.backends.get("b1").map(String::as_str), Some("backend-page2")); } @@ -256,19 +345,60 @@ mod tests { // next_cursor is None → backend is exhausted let result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); - let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host); + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); assert!(next_cursor.is_none(), "no cursor when all backends exhausted"); } #[test] fn invalid_cursor_returns_invalid_params_error() { - let err = decode_gateway_cursor(Some("not-json!")).unwrap_err(); + let err = decode_gateway_cursor(Some("not-json!"), "list_tools").unwrap_err(); assert_eq!(err.code, ErrorCode::INVALID_PARAMS); } #[test] fn none_cursor_returns_default_empty_gateway_cursor() { - let cursor = decode_gateway_cursor(None).expect("None is a valid first-page indicator"); + let cursor = decode_gateway_cursor(None, "list_tools").expect("None is a valid first-page indicator"); assert!(cursor.backends.is_empty()); } + + // Finding 1: a backend that fails mid-page keeps its prior cursor position so the + // client can reach its remaining items on the next page. + #[test] + fn failed_backend_cursor_preserved_across_page() { + let virtual_host = test_virtual_host("b1"); + // Simulate an incoming cursor where both b1 and b2 had more pages. + let incoming = GatewayCursor { + op: "list_tools".to_owned(), + backends: [ + ("b1".to_owned(), "b1-page2".to_owned()), + ("b2".to_owned(), "b2-page2".to_owned()), + ] + .into(), + }; + // b1 succeeds and is done (no next_cursor); b2 fails (absent from results). + let b1_result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), b1_result)], &virtual_host, &incoming, "list_tools"); + + let raw = next_cursor.expect("b2 failure must produce a next cursor to retry"); + let cursor: GatewayCursor = serde_json::from_str(&raw).expect("valid JSON"); + assert!(cursor.backends.get("b1").is_none(), "b1 exhausted itself — must not appear"); + assert_eq!(cursor.backends.get("b2").map(String::as_str), Some("b2-page2"), "b2 prior cursor preserved"); + } + + // Finding 2: a cursor issued by one list operation must be rejected when presented + // to a different operation, returning -32602. + #[test] + fn cross_operation_cursor_rejected() { + let virtual_host = test_virtual_host("b1"); + let mut result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); + result.next_cursor = Some("page2".to_owned()); + let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); + let raw = next_cursor.expect("cursor present"); + + // Presenting a list_tools cursor to list_prompts must fail. + let err = decode_gateway_cursor(Some(&raw), "list_prompts").unwrap_err(); + assert_eq!(err.code, ErrorCode::INVALID_PARAMS); + // Presenting it to its own operation succeeds. + assert!(decode_gateway_cursor(Some(&raw), "list_tools").is_ok()); + } } diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs index fcc13f6..14c0654 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/prompts.rs @@ -29,7 +29,7 @@ where let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); let all_transports: Vec<_> = session_manager.borrow_transports().await; - let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_prompts")?; let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() } else { @@ -58,7 +58,7 @@ where ) .await; - let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers); + let (prompts, next_cursor) = merge_prompts(responses, namespace_identifiers, &gateway_cursor, "list_prompts"); let mut result = ListPromptsResult::with_all_items(prompts); result.next_cursor = next_cursor; Ok(result) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index 8c98974..e1d9ee7 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -32,7 +32,7 @@ where let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); let all_transports: Vec<_> = session_manager.borrow_transports().await; - let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_resources")?; let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() } else { @@ -61,7 +61,7 @@ where ) .await; - let (resources, next_cursor) = merge_resources(responses, namespace_identifiers); + let (resources, next_cursor) = merge_resources(responses, namespace_identifiers, &gateway_cursor, "list_resources"); let mut result = ListResourcesResult::with_all_items(resources); result.next_cursor = next_cursor; Ok(result) @@ -112,7 +112,7 @@ where let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); let all_transports: Vec<_> = session_manager.borrow_transports().await; - let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_resource_templates")?; let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() } else { @@ -141,7 +141,7 @@ where ) .await; - let (resource_templates, next_cursor) = merge_resource_templates(responses, namespace_identifiers); + let (resource_templates, next_cursor) = merge_resource_templates(responses, namespace_identifiers, &gateway_cursor, "list_resource_templates"); let mut result = ListResourceTemplatesResult::with_all_items(resource_templates); result.next_cursor = next_cursor; Ok(result) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs index 349fc49..903eb43 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/tools.rs @@ -29,7 +29,7 @@ where let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); let all_transports: Vec<_> = session_manager.borrow_transports().await; - let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()))?; + let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_tools")?; // On resume, skip backends already exhausted in the prior page. // ponytail: topology changes between pages silently drop removed backends; // add a cursor version field if reconfiguration stability matters. @@ -61,7 +61,7 @@ where ) .await; - let (tools, next_cursor) = merge_tools(responses, virtual_host); + let (tools, next_cursor) = merge_tools(responses, virtual_host, &gateway_cursor, "list_tools"); let mut result = ListToolsResult::with_all_items(tools); result.next_cursor = next_cursor; Ok(result) From f1243d20551236c1fed3457b69b607a42e5d7530 Mon Sep 17 00:00:00 2001 From: Lang-Akshay Date: Thu, 30 Jul 2026 16:16:10 +0100 Subject: [PATCH 7/7] fixup: fmt and clippy Signed-off-by: Lang-Akshay --- .../src/gateway/list_aggregation.rs | 26 ++++++++++--------- .../src/gateway/mcp_service/resources.rs | 6 +++-- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs index 8f41912..f9945f1 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/list_aggregation.rs @@ -1,4 +1,7 @@ -use std::{collections::{HashMap, HashSet}, future::Future}; +use std::{ + collections::{HashMap, HashSet}, + future::Future, +}; use contextforge_gateway_rs_apis::user_store::VirtualHost; use rmcp::{ @@ -42,8 +45,8 @@ impl Default for GatewayCursor { /// issued by one list operation from being accepted by a different one. pub(super) fn decode_gateway_cursor(raw: Option<&str>, expected_op: &str) -> Result { let Some(raw) = raw else { return Ok(GatewayCursor::default()) }; - let cursor: GatewayCursor = serde_json::from_str(raw) - .map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None))?; + let cursor: GatewayCursor = + serde_json::from_str(raw).map_err(|_| ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid cursor", None))?; if cursor.op != expected_op { return Err(ErrorData::new(ErrorCode::INVALID_PARAMS, "cursor operation mismatch", None)); } @@ -331,7 +334,8 @@ mod tests { let mut result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); result.next_cursor = Some("backend-page2".to_owned()); - let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); + let (_, next_cursor) = + merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); let raw = next_cursor.expect("should have a next cursor"); let cursor: GatewayCursor = serde_json::from_str(&raw).expect("valid JSON"); @@ -345,7 +349,8 @@ mod tests { // next_cursor is None → backend is exhausted let result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); - let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); + let (_, next_cursor) = + merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); assert!(next_cursor.is_none(), "no cursor when all backends exhausted"); } @@ -369,11 +374,7 @@ mod tests { // Simulate an incoming cursor where both b1 and b2 had more pages. let incoming = GatewayCursor { op: "list_tools".to_owned(), - backends: [ - ("b1".to_owned(), "b1-page2".to_owned()), - ("b2".to_owned(), "b2-page2".to_owned()), - ] - .into(), + backends: [("b1".to_owned(), "b1-page2".to_owned()), ("b2".to_owned(), "b2-page2".to_owned())].into(), }; // b1 succeeds and is done (no next_cursor); b2 fails (absent from results). let b1_result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); @@ -381,7 +382,7 @@ mod tests { let raw = next_cursor.expect("b2 failure must produce a next cursor to retry"); let cursor: GatewayCursor = serde_json::from_str(&raw).expect("valid JSON"); - assert!(cursor.backends.get("b1").is_none(), "b1 exhausted itself — must not appear"); + assert!(!cursor.backends.contains_key("b1"), "b1 exhausted itself — must not appear"); assert_eq!(cursor.backends.get("b2").map(String::as_str), Some("b2-page2"), "b2 prior cursor preserved"); } @@ -392,7 +393,8 @@ mod tests { let virtual_host = test_virtual_host("b1"); let mut result = ListToolsResult::with_all_items(vec![Tool::new("t1", "", serde_json::Map::new())]); result.next_cursor = Some("page2".to_owned()); - let (_, next_cursor) = merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); + let (_, next_cursor) = + merge_tools(vec![("b1".to_owned(), result)], &virtual_host, &GatewayCursor::default(), "list_tools"); let raw = next_cursor.expect("cursor present"); // Presenting a list_tools cursor to list_prompts must fail. diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs index e1d9ee7..d75dea0 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_service/resources.rs @@ -112,7 +112,8 @@ where let session_manager = SessionManager::new(virtual_host, session_id, claims.sub.as_str(), &mcp_service.transports); let all_transports: Vec<_> = session_manager.borrow_transports().await; - let gateway_cursor = decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_resource_templates")?; + let gateway_cursor = + decode_gateway_cursor(request.as_ref().and_then(|r| r.cursor.as_deref()), "list_resource_templates")?; let backend_transports: Vec<_> = if request.as_ref().and_then(|r| r.cursor.as_ref()).is_some() { all_transports.into_iter().filter(|b| gateway_cursor.backends.contains_key(&b.name)).collect() } else { @@ -141,7 +142,8 @@ where ) .await; - let (resource_templates, next_cursor) = merge_resource_templates(responses, namespace_identifiers, &gateway_cursor, "list_resource_templates"); + let (resource_templates, next_cursor) = + merge_resource_templates(responses, namespace_identifiers, &gateway_cursor, "list_resource_templates"); let mut result = ListResourceTemplatesResult::with_all_items(resource_templates); result.next_cursor = next_cursor; Ok(result)