Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions _context/wiki/architecture.md
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,8 @@ Order is invariant: auth/config before backend selection; request plugins before
| User config | `RedisUserConfigStore` (LRU + Redis) | Request-path consumed; control-plane authored |
| Request identity / VirtualHostId | Request extensions | One HTTP request |
| Downstream session id | RMCP + `SessionId` extension | MCP session |
| Backend RMCP services | `BackendTransports` map | Local process, per principal/backend/session |
| Backend RMCP services (initialize, list ops) | `BackendTransports` map | Local process, per principal/backend/session |
| Backend RMCP services (call_tool) | Per-request connection | Single HTTP request |
| Local user session mapping | `LocalUserSessionStore` | Local LRU, 50k entries, 1 hour |
| Plugin manager | `CpexRuntimeRegistry` | Process, reloadable |

Expand All @@ -121,7 +122,7 @@ In multi-runtime mode, the first thread initializes the optional CPEX plugin run

| State | Lock | Contention profile |
| --- | --- | --- |
| `BackendTransports` map | `Arc<tokio::sync::Mutex<HashMap<...>>>` | Locked briefly on initialize insert, per-call borrow, and cleanup. Borrowing clones `Arc<RunningService>` handles so the lock is not held across backend calls. |
| `BackendTransports` map | `Arc<tokio::sync::Mutex<HashMap<...>>>` | Locked briefly on initialize insert, list-op borrow, and cleanup. Borrowing clones `Arc<RunningService>` handles so the lock is not held across backend calls. `call_tool` bypasses this map entirely. |
| Subscription set | `Arc<tokio::sync::Mutex<HashSet<String>>>` | Local `subscribe`/`unsubscribe` only. |
| User config LRU cache | `Arc<tokio::sync::Mutex<LruCache>>` inside `RedisUserConfigStore` | One lock per config lookup on the hot path; misses add a Redis round trip. |
| User session LRU cache | Same pattern in `LocalUserSessionStore` | Initialize and delete paths. |
Expand All @@ -141,7 +142,8 @@ The binary sets `tikv_jemallocator` as the global allocator. jemalloc holds up b

- `initialize` opens one backend transport per configured backend concurrently (`futures::future::join_all`); a failed backend degrades that backend only.
- List methods fan out to all connected backends concurrently and merge.
- Targeted calls resolve exactly one backend service handle.
- Targeted calls (except `call_tool`) resolve exactly one backend service handle from `BackendTransports`.
- `call_tool` creates a fresh per-request backend connection via `connect_backend_for_request`, runs pre/post plugin hooks, executes the call, then explicitly closes the connection before returning.
- `call_tool` watches the downstream cancellation token and forwards a cancel to the backend if the client gives up first; backend progress notifications are forwarded downstream while the call is in flight.

## Startup And Response Flow
Expand Down
4 changes: 3 additions & 1 deletion _context/wiki/routing.md
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ This is **local process state only**. Implications:
- Gateway restart → all sessions lost → clients must re-run `initialize`.
- Multi-runtime mode (`--single-runtime false`): each runtime thread has its own `BackendTransports` with no cross-thread affinity.

**Exception: `call_tool` uses per-request backend lifecycle.** Each tool call creates a fresh backend connection, executes the call with plugin hooks, then closes the connection. This bypasses `BackendTransports` entirely and does not require session affinity for tool calls specifically (though other MCP methods still do).


```mermaid
sequenceDiagram
Expand Down Expand Up @@ -130,7 +132,7 @@ If RMCP rejects the delete, local state is untouched.
| `list_resources` | List | Same as list_tools. |
| `list_prompts` | List | Same as list_tools. |
| `list_resource_templates` | List | Same — both name and URI template get prefixed for multi-backend. |
| `call_tool` | Targeted | Resolves alias → single/multi-backend fallback. Runs pre/post plugin hooks. Forwards downstream cancellation to backend. Tracks backend progress tokens: RMCP assigns a new token per backend request; the gateway maps each backend token to the downstream token. Request enqueue and mapping publication are serialized against progress lookup so an immediate backend notification cannot overtake registration. When the notification matches an in-flight token, the gateway restores the downstream token and forwards it to the client. |
| `call_tool` | Targeted | **Per-request backend lifecycle:** creates fresh connection via `connect_backend_for_request`, runs pre-hook, executes call, runs post-hook, closes connection. Resolves alias → single/multi-backend fallback. Forwards downstream cancellation to backend. Tracks backend progress tokens: RMCP assigns a new token per backend request; the gateway maps each backend token to the downstream token. Request enqueue and mapping publication are serialized against progress lookup so an immediate backend notification cannot overtake registration. When the notification matches an in-flight token, the gateway restores the downstream token and forwards it to the client. Does not use session-backed `BackendTransports`. |
| `read_resource` | Targeted | Single-backend: URI unchanged. Multi-backend: strips prefix. |
| `subscribe` / `unsubscribe` | Targeted | Same resource-URI routing; forwards/stops resource-update notifications. |
| `get_prompt` | Targeted | Single-backend: name unchanged. Multi-backend: strips prefix. Runs pre/post prompt hooks around the backend call: the pre hook may rewrite arguments or deny, the post hook may rewrite or reject the rendered messages. |
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,62 @@ impl<'a> AuthorizedCallValidator<'a> {
pub fn new(call_name: &'a str, ctx: &'a RequestContext<RoleServer>) -> Self {
Self { call_name, ctx }
}

pub fn validate_stateless(self) -> Result<(&'a VirtualHost, &'a ContextForgeClaims), ErrorData> {
Comment thread
lucarlig marked this conversation as resolved.
let maybe_parts = self.ctx.extensions.get::<Parts>();
let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::<UserConfig>());
let maybe_claims = maybe_parts.and_then(|parts| parts.extensions.get::<ContextForgeClaims>());
let maybe_virtual_host_id = maybe_parts.and_then(|parts| parts.extensions.get::<VirtualHostId>());
let call_name = self.call_name;
let has_user_config = maybe_user_config.is_some();
let virtual_hosts = maybe_user_config.map_or(0, |user_config| user_config.virtual_hosts.len());
let has_claims = maybe_claims.is_some();
let virtual_host_id = maybe_virtual_host_id.map_or("<missing>", |id| id.value().as_str());
debug!(
"AuthorizedCallValidator::validate - mcp call validation call_name = {call_name} has_user_config = {has_user_config} virtual_hosts = {virtual_hosts} has_claims = {has_claims} virtual_host_id = {virtual_host_id}"
);

let Some(user_config) = maybe_user_config else {
return Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: "Routing problem... user config not found".into(),
data: None,
});
};

let Some(virtual_host_id) = maybe_virtual_host_id else {
return Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: "Routing problem... virtual host not known".into(),
data: None,
});
};

let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.value()) else {
let call_name = self.call_name;
let virtual_host_id = virtual_host_id.value();
let virtual_hosts = user_config.virtual_hosts.len();
debug!(
"AuthorizedCallValidator::validate - mcp virtual host config missing call_name = {call_name} virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}"
);
return Err(ErrorData {
code: ErrorCode::RESOURCE_NOT_FOUND,
message: "No configuration".into(),
data: None,
});
};

let Some(claims) = maybe_claims else {
return Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: "Routing problem... claims not found".into(),
data: None,
});
};

Ok((virtual_host, claims))
}

pub fn validate(self) -> Result<(&'a VirtualHost, &'a SessionId, &'a ContextForgeClaims), ErrorData> {
let maybe_parts = self.ctx.extensions.get::<Parts>();
let maybe_session_id = maybe_parts.and_then(|parts| parts.extensions.get::<SessionId>());
Expand Down Expand Up @@ -120,7 +176,7 @@ impl<'a> InitializeCallValidator<'a> {
let Some(virtual_host_id) = maybe_virtual_host_id else {
return Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: "Routing problem... virutal host not known".into(),
message: "Routing problem... virtual host not known".into(),
data: None,
});
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@ use rmcp::{
};
use typed_builder::TypedBuilder;

use super::{backend_transports::BackendTransports, session_store::UserSessionStore};
use crate::gateway::UserSessionStore;

use super::backend_transports::BackendTransports;

#[derive(Clone, TypedBuilder)]
#[builder(field_defaults(setter(prefix = "with_")))]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,12 @@ use std::{collections::HashMap, sync::Arc};
use contextforge_data_plane_apis::user_store::BackendMCPGateway;
use http::request::Parts;
use rmcp::{
ErrorData, RoleClient, RoleServer, ServiceExt,
model::{ErrorCode, Implementation, InitializeRequestParams, InitializeResult, ServerCapabilities},
ClientLifecycleMode, ErrorData, RoleClient, RoleServer, ServiceExt,
model::{
ClientCapabilities, ErrorCode, Implementation, InitializeRequestParams, InitializeResult, ProtocolVersion,
ServerCapabilities,
},
service::serve_client_with_lifecycle_and_ct,
service::{RequestContext, RunningService},
transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig},
};
Expand Down Expand Up @@ -193,20 +197,68 @@ fn merge_and_build_capabilities(server_capabilities: Vec<(String, Option<ServerC
merged
}

pub(super) async fn connect_backend_for_request<T>(
Comment thread
lucarlig marked this conversation as resolved.
mcp_service: &McpService<T>,
backend_name: &str,
backend: &BackendMCPGateway,
namespace_identifiers: bool,
cx: &RequestContext<RoleServer>,
) -> Result<RunningService<RoleClient, GatewayBackendClient>, ErrorData>
where
T: UserSessionStore + Send + Sync + 'static,
{
let mut headers = HashMap::new();
let downstream_headers = cx.extensions.get::<Parts>().map(|parts| &parts.headers);

if let Some(host) = backend.url.host_str()
&& backend.url.scheme() == "https"
{
let authority = if let Some(port) = backend.url.port() { format!("{host}:{port}") } else { host.to_owned() };
if let Ok(value) = http::HeaderValue::from_str(&authority) {
headers.insert(http::header::HOST, value);
} else {
warn!("connect_backend_for_request - invalid backend host backend_name = {backend_name}");
}
}

apply_header_config(&mut headers, backend, downstream_headers);
crate::telemetry::inject_current_context(&mut headers);

let config = StreamableHttpClientTransportConfig::with_uri(backend.url.to_string()).custom_headers(headers);
let transport = StreamableHttpClientTransport::with_client(mcp_service.http_client.clone(), config);
let client_info = InitializeRequestParams::new(
ClientCapabilities::default(),
Implementation::new("contextforge-data-plane", env!("CARGO_PKG_VERSION")),
)
.with_protocol_version(ProtocolVersion::V_2026_07_28);
Comment thread
lucarlig marked this conversation as resolved.

let backend_client = GatewayBackendClient::new(
backend_name.to_owned(),
namespace_identifiers,
client_info,
mcp_service.plugin_runtime.clone(),
);

serve_client_with_lifecycle_and_ct(
backend_client,
transport,
ClientLifecycleMode::Discover { preferred_versions: vec![ProtocolVersion::V_2026_07_28] },
cx.ct.clone(),
)
.await
.map_err(|error| {
warn!(
"connect_backend_for_request - backend connection failed backend_name = {backend_name} error = {error:?}"
);
ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: "Routing problem... backend unavailable".into(),
data: None,
}
})
}

/// Apply a backend's header config to the upstream header map.
///
/// Order: passthrough (copy named headers from the downstream request) -> add
/// (inject/override static headers) -> remove (strip named headers).
///
/// Protected headers are silently skipped in every phase:
/// - Gateway-managed: `Host` (set from backend URL before this runs)
/// - Body-framing: `Content-Length`, `Content-Type` (gateway owns framing)
/// - Hop-by-hop (RFC 7230 §6.1): `Connection`, `Keep-Alive`, `Proxy-Authenticate`,
/// `Proxy-Authorization`, `TE`, `Trailer`, `Trailers`, `Transfer-Encoding`, `Upgrade`
/// - Non-standard hop-by-hop: `Proxy-Connection`
/// - RMCP transport-reserved: `Mcp-Session-Id`, `Accept`, `Last-Event-Id`
///
/// ponytail: single-value per name; a repeated downstream header keeps its first value.
fn apply_header_config(
headers: &mut HashMap<http::HeaderName, http::HeaderValue>,
backend: &BackendMCPGateway,
Expand Down
26 changes: 16 additions & 10 deletions crates/contextforge-data-plane-lib/src/gateway/mcp_service/tools.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,15 @@ use rmcp::{
model::{CallToolRequestParams, CallToolResponse, ErrorCode, ListToolsResult, PaginatedRequestParams},
service::RequestContext,
};
use tracing::info;
use tracing::{info, warn};

use super::McpService;
use crate::gateway::{
backend_client::call_backend_tool,
identifier_routing::{backend_forward_error, resolve_backend, resolve_tool_route},
identifier_routing::{backend_forward_error, resolve_tool_route},
list_aggregation::{decode_gateway_cursor, fan_out_list, merge_tools},
mcp_call_validator::AuthorizedCallValidator,
mcp_service::initialization::connect_backend_for_request,
session_manager::SessionManager,
session_store::UserSessionStore,
};
Expand Down Expand Up @@ -76,11 +77,8 @@ where
T: UserSessionStore + Send + Sync + 'static,
{
let mcp_call_validator = AuthorizedCallValidator::new("call_tool", &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_names = session_manager.get_backend_names();

let (virtual_host, _claims) = mcp_call_validator.validate_stateless()?;
Comment thread
cafalchio marked this conversation as resolved.
let backend_names: Vec<&str> = virtual_host.backends.keys().map(String::as_str).collect();
let Some((backend_name, tool_name)) = resolve_tool_route(virtual_host, &request.name, &backend_names) else {
return Err(ErrorData {
code: ErrorCode::INVALID_PARAMS,
Expand All @@ -90,14 +88,19 @@ where
};
let backend_name = backend_name.to_owned();
let tool_name = tool_name.to_owned();

let (service_name, backend_service) = resolve_backend(&session_manager, "call_tool", &backend_name).await?;

let backend = virtual_host.backends.get(&backend_name).ok_or_else(|| ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: "Routing problem... backend not found".into(),
data: None,
})?;
let service_name = backend_name.clone();
let pre_result = if let Some(plugin_runtime) = &mcp_service.plugin_runtime {
plugin_runtime.before_tool_call(&request, &tool_name, &service_name).await?
} else {
ToolPreCallResult::unchanged()
};
let mut backend_service =
connect_backend_for_request(mcp_service, &backend_name, backend, virtual_host.backends.len() > 1, &cx).await?;
let post_state = pre_result.state;
let mut routed_request = request;
pre_result.arguments.apply_to_request(&mut routed_request, &tool_name);
Expand All @@ -118,6 +121,9 @@ where
let backend_progress_token = handle.progress_token.clone();
let response = call_backend_tool(handle, cx.ct.clone()).await;
backend_service.service().stop_tracking_tool_call(&backend_progress_token).await;
if let Err(error) = backend_service.close().await {
warn!("call_tool: backend cleanup failed backend_name = {service_name} error = {error:?}");
}

let response = response.map_err(|error| backend_forward_error("call_tool", &service_name, &error))?;
let response = match (&mcp_service.plugin_runtime, post_state) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ mod tests {
use crate::layers::virtual_host_id::{VirtualHostId, extract_virtual_host_id};

#[test]
fn test_virutal_host_extractor() {
fn test_virtual_host_extractor() {
assert_eq!(None, extract_virtual_host_id("/mcp/servers"));
assert_eq!(None, extract_virtual_host_id("/servers"));
assert_eq!(None, extract_virtual_host_id("/servers/12345_abcd-efgh/mcp/dkfjk"));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use support::{

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[test_log::test]
#[ignore = "2026-07-28 protocol transition"]
Comment thread
cafalchio marked this conversation as resolved.
async fn plaintext_completes_prompt_argument_through_prefixed_backend() -> Result<()> {
let gateway_port = create_ports(1)[0];
let user = TEST_USER_ID;
Expand All @@ -28,6 +29,7 @@ async fn plaintext_completes_prompt_argument_through_prefixed_backend() -> Resul

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[test_log::test]
#[ignore = "2026-07-28 protocol transition"]
Comment thread
cafalchio marked this conversation as resolved.
async fn plaintext_completes_resource_argument_through_prefixed_backend() -> Result<()> {
let gateway_port = create_ports(1)[0];
let user = TEST_USER_ID;
Expand Down
Loading