Skip to content
Closed
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
53 changes: 11 additions & 42 deletions crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,18 @@ use contextforge_gateway_rs_cpex::{GatewayPluginRuntimeHandle, RuntimeHookState}
use rmcp::{
ClientHandler, Peer, RoleClient, RoleServer,
model::{
CallToolRequestParams, CallToolResult, ClientRequest, InitializeRequestParams, LoggingMessageNotificationParam,
Meta, ProgressNotificationParam, ProgressToken, Request, ServerResult,
CallToolRequestParams, CallToolResult, ClientRequest, InitializeRequestParams, Meta, ProgressNotificationParam,
ProgressToken, Request, ServerResult,
},
serde::{Serialize, de::DeserializeOwned},
service::{NotificationContext, PeerRequestOptions, ServiceError},
};
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};

/// Client handler for a backend MCP connection. Progress and logging
/// notifications streamed by the backend are run through the tool post hooks
/// of the in-flight tool call they belong to and forwarded to the downstream
/// peer of the gateway session.
/// Client handler for a backend MCP connection. Progress notifications streamed
/// by the backend are run through the tool post hooks of the in-flight tool call
/// they belong to and forwarded to the downstream peer of the gateway session.
#[derive(Clone)]
pub(crate) struct GatewayBackendClient {
initialize_request: InitializeRequestParams,
Expand All @@ -30,11 +29,9 @@ pub(crate) struct GatewayBackendClient {

/// Tool calls currently in flight on this backend session. Progress
/// notifications are routed strictly by the downstream progress token, as MCP
/// requires backends to echo the token the client sent; logging notifications
/// carry no token and are attributed to the most recently started call.
/// requires backends to echo the token the client sent.
#[derive(Default)]
struct InFlightCalls {
ordered: Vec<Arc<InFlightToolCall>>,
by_progress_token: HashMap<ProgressToken, Arc<InFlightToolCall>>,
}

Expand All @@ -53,13 +50,11 @@ pub(crate) struct InFlightToolCallGuard {

impl Drop for InFlightToolCallGuard {
fn drop(&mut self) {
if let Ok(mut calls) = self.calls.lock() {
calls.ordered.retain(|call| !Arc::ptr_eq(call, &self.call));
if let Some(token) = &self.call.progress_token
&& calls.by_progress_token.get(token).is_some_and(|call| Arc::ptr_eq(call, &self.call))
{
calls.by_progress_token.remove(token);
}
if let Ok(mut calls) = self.calls.lock()
&& let Some(token) = &self.call.progress_token
&& calls.by_progress_token.get(token).is_some_and(|call| Arc::ptr_eq(call, &self.call))
{
calls.by_progress_token.remove(token);
}
}
}
Expand All @@ -81,7 +76,6 @@ impl GatewayBackendClient {
) -> InFlightToolCallGuard {
let call = Arc::new(InFlightToolCall { progress_token, tool_name, post_state });
let mut calls = self.in_flight_calls.lock().expect("in-flight tool call lock poisoned");
calls.ordered.push(Arc::clone(&call));
if let Some(token) = &call.progress_token {
calls.by_progress_token.entry(token.clone()).or_insert_with(|| Arc::clone(&call));
}
Expand All @@ -94,10 +88,6 @@ impl GatewayBackendClient {
calls.by_progress_token.get(progress_token).cloned()
}

fn latest_call(&self) -> Option<Arc<InFlightToolCall>> {
self.in_flight_calls.lock().expect("in-flight tool call lock poisoned").ordered.last().cloned()
}

async fn stream_event_post_hook<T>(&self, call: &InFlightToolCall, event: T) -> Option<T>
where
T: Serialize + DeserializeOwned,
Expand Down Expand Up @@ -145,27 +135,6 @@ impl ClientHandler for GatewayBackendClient {
warn!("call_tool: unable to forward backend progress notification downstream: {error:?}");
}
}

#[expect(deprecated, reason = "logging forwarding is kept until the SEP-2577 removal lands in MCP")]
async fn on_logging_message(
&self,
message: LoggingMessageNotificationParam,
_context: NotificationContext<RoleClient>,
) {
// MCP does not correlate logging messages with requests, so they are
// attributed to the most recent in-flight tool call. Messages with no
// in-flight call are dropped so nothing bypasses the post hooks.
let Some(call) = self.latest_call() else {
debug!("call_tool: dropping backend logging notification without an in-flight tool call");
return;
};
let Some(message) = self.stream_event_post_hook(&call, message).await else {
return;
};
if let Err(error) = self.downstream.notify_logging_message(message).await {
warn!("call_tool: unable to forward backend logging notification downstream: {error:?}");
}
}
}

/// Calls the tool on the backend, keeping the downstream progress token on
Expand Down
23 changes: 5 additions & 18 deletions crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,10 @@ use rmcp::{
model::{
AnnotateAble, CallToolRequestParams, CallToolResult, CompleteRequestParams, CompleteResult, CompletionInfo,
ErrorCode, GetPromptRequestParams, GetPromptResult, Implementation, InitializeRequestParams, InitializeResult,
ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, LoggingLevel,
PaginatedRequestParams, Prompt, PromptArgument, PromptMessage, PromptMessageContent, PromptMessageRole,
RawImageContent, RawResourceTemplate, ReadResourceRequestParams, ReadResourceResult, Reference, Resource,
ServerCapabilities, SetLevelRequestParams, SubscribeRequestParams, Tool, UnsubscribeRequestParams,
ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams,
Prompt, PromptArgument, PromptMessage, PromptMessageContent, PromptMessageRole, RawImageContent,
RawResourceTemplate, ReadResourceRequestParams, ReadResourceResult, Reference, Resource, ServerCapabilities,
SubscribeRequestParams, Tool, UnsubscribeRequestParams,
},
service::{RequestContext, RunningService},
transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig},
Expand Down Expand Up @@ -62,8 +62,6 @@ where
subscriptions: Arc<Mutex<HashSet<String>>>,
#[builder(default = BackendTransports::default())]
transports: BackendTransports,
#[builder(default = Arc::new(Mutex::new(LoggingLevel::Debug)))]
log_level: Arc<Mutex<LoggingLevel>>,
http_client: reqwest::Client,
user_session_store: T,
#[builder(default)]
Expand Down Expand Up @@ -694,16 +692,6 @@ where
};
Ok(CompleteResult::new(CompletionInfo::new(values).map_err(|e| ErrorData::internal_error(e, None))?))
}

async fn set_level(&self, request: SetLevelRequestParams, cx: RequestContext<RoleServer>) -> Result<(), ErrorData> {
let maybe_parts = cx.extensions.get::<Parts>();
let maybe_session = maybe_parts.and_then(|parts| parts.extensions.get::<SessionId>());
let maybe_user_config = maybe_parts.and_then(|parts| parts.extensions.get::<UserConfig>());
info!("set_level user_config = {maybe_user_config:#?} session_id = {maybe_session:#?}");
let mut level = self.log_level.lock().await;
*level = request.level;
Ok(())
}
}

#[derive(Debug, PartialEq, PartialOrd, Ord, Eq)]
Expand Down Expand Up @@ -737,9 +725,8 @@ const TEST_IMAGE_DATA: &str =
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==";
// Small base64-encoded WAV (silence)

#[expect(deprecated, reason = "logging capability is kept until the SEP-2577 removal lands in MCP")]
fn merge_capabilities(_server_capabilities: Vec<(String, Option<ServerCapabilities>)>) -> ServerCapabilities {
ServerCapabilities::builder().enable_prompts().enable_resources().enable_tools().enable_logging().build()
ServerCapabilities::builder().enable_prompts().enable_resources().enable_tools().build()
}

fn merge_tools(tools: Vec<(String, ListToolsResult)>) -> Vec<Tool> {
Expand Down
49 changes: 14 additions & 35 deletions crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,7 @@ use rmcp::{
ClientHandler,
model::{
CallToolRequestParams, CallToolResult, ClientCapabilities, ClientRequest, ErrorCode, Implementation,
InitializeRequestParams, LoggingMessageNotificationParam, Meta, NumberOrString, ProgressNotificationParam,
ProgressToken, Request, ServerResult,
InitializeRequestParams, Meta, NumberOrString, ProgressNotificationParam, ProgressToken, Request, ServerResult,
},
service::{NotificationContext, PeerRequestOptions, RoleClient},
};
Expand All @@ -27,7 +26,6 @@ type Recorded<T> = Arc<StdMutex<Vec<T>>>;
#[derive(Clone, Default)]
struct RecordingClient {
progress: Recorded<ProgressNotificationParam>,
messages: Recorded<LoggingMessageNotificationParam>,
}

impl ClientHandler for RecordingClient {
Expand All @@ -41,33 +39,23 @@ impl ClientHandler for RecordingClient {
async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext<RoleClient>) {
self.progress.lock().expect("progress lock poisoned").push(params);
}

async fn on_logging_message(
&self,
params: LoggingMessageNotificationParam,
_context: NotificationContext<RoleClient>,
) {
self.messages.lock().expect("messages lock poisoned").push(params);
}
}

async fn call_progress_sum(
gateway: &RunningGateway,
user: &str,
) -> (CallToolResult, Recorded<ProgressNotificationParam>, Recorded<LoggingMessageNotificationParam>) {
let (result, progress, messages) = send_progress_sum(gateway, user).await;
) -> (CallToolResult, Recorded<ProgressNotificationParam>) {
let (result, progress) = send_progress_sum(gateway, user).await;
wait_for_event_count(&progress, 4).await;
wait_for_event_count(&messages, 4).await;
(result, progress, messages)
(result, progress)
}

async fn send_progress_sum(
gateway: &RunningGateway,
user: &str,
) -> (CallToolResult, Recorded<ProgressNotificationParam>, Recorded<LoggingMessageNotificationParam>) {
) -> (CallToolResult, Recorded<ProgressNotificationParam>) {
let client = RecordingClient::default();
let progress = Arc::clone(&client.progress);
let messages = Arc::clone(&client.messages);
let service = gateway.connect_with_handler(user, client).await;
let request = CallToolRequestParams::new(format!("{}-progress_sum", gateway.backend_name));
let mut options = PeerRequestOptions::no_options();
Expand All @@ -81,7 +69,7 @@ async fn send_progress_sum(
else {
panic!("expected call tool result");
};
(result, progress, messages)
(result, progress)
}

async fn wait_for_event_count<T>(events: &StdMutex<Vec<T>>, expected: usize) {
Expand Down Expand Up @@ -232,7 +220,6 @@ async fn concurrent_progress_calls_forward_each_token_without_plugins() {
let gateway = start_gateway("admin@example.com", false, Arc::new(CpexRuntimeRegistry::default())).await;
let client = RecordingClient::default();
let progress = Arc::clone(&client.progress);
let messages = Arc::clone(&client.messages);
let service = gateway.connect_with_handler("admin@example.com", client).await;
let request = CallToolRequestParams::new(format!("{}-progress_sum", gateway.backend_name));

Expand Down Expand Up @@ -266,7 +253,6 @@ async fn concurrent_progress_calls_forward_each_token_without_plugins() {
assert_eq!("completed 4 packages", text(&second));

wait_for_event_count(&progress, 8).await;
wait_for_event_count(&messages, 8).await;
let progress = progress.lock().expect("progress lock poisoned");
let first_count = progress
.iter()
Expand Down Expand Up @@ -384,41 +370,35 @@ async fn post_hook_receives_backend_result_and_modifies_client_result() {
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn post_hook_can_modify_stream_progress_and_message_notifications() {
async fn post_hook_can_modify_stream_progress_notifications() {
let plugin =
Arc::new(TestPlugin::new("post-stream", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_stream_event_rewrite());
let observations = plugin.observations();
let runtime = runtime_with_post(plugin).await;

let gateway = start_gateway("admin@example.com", true, runtime).await;
let (result, progress, messages) = call_progress_sum(&gateway, "admin@example.com").await;
let (result, progress) = call_progress_sum(&gateway, "admin@example.com").await;

assert_eq!("completed 4 packages", text(&result));
let progress = progress.lock().expect("progress lock poisoned");
assert_eq!(Some("plugin:package 4/4"), progress.last().and_then(|notification| notification.message.as_deref()));
let messages = messages.lock().expect("messages lock poisoned");
assert_eq!(
Some("message"),
messages.last().and_then(|notification| notification.data.get("plugin")).and_then(Value::as_str)
);

let observations = observations.lock().expect("observations lock poisoned");
assert_eq!(9, observations.post_calls);
assert_eq!(5, observations.post_calls);
let first_id = observations.post_tool_call_ids.first().expect("post call id");
assert!(observations.post_tool_call_ids.iter().all(|id| id == first_id));
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn json_response_mode_forwards_backend_progress_and_message_notifications() {
async fn json_response_mode_forwards_backend_progress_notifications() {
let plugin = Arc::new(TestPlugin::new("post", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_post_rewrite());
let runtime = runtime_with_post(plugin).await;

let gateway = start_gateway_with_json_backend_responses("admin@example.com", true, runtime).await;
let (result, progress, messages) = call_progress_sum(&gateway, "admin@example.com").await;
let (result, progress) = call_progress_sum(&gateway, "admin@example.com").await;

assert_eq!("post:completed 4 packages", text(&result));
assert_eq!(4, progress.lock().expect("progress lock poisoned").len());
assert_eq!(4, messages.lock().expect("messages lock poisoned").len());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
Expand All @@ -429,19 +409,18 @@ async fn post_hook_deny_drops_stream_notifications_without_failing_call() {
let runtime = runtime_with_post(plugin).await;

let gateway = start_gateway("admin@example.com", true, runtime).await;
let (result, progress, messages) = send_progress_sum(&gateway, "admin@example.com").await;
let (result, progress) = send_progress_sum(&gateway, "admin@example.com").await;

assert_eq!("completed 4 packages", text(&result));
for _ in 0..50 {
if observations.lock().expect("observations lock poisoned").post_calls >= 9 {
if observations.lock().expect("observations lock poisoned").post_calls >= 5 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
let observations = observations.lock().expect("observations lock poisoned");
assert_eq!(9, observations.post_calls);
assert_eq!(5, observations.post_calls);
assert!(progress.lock().expect("progress lock poisoned").is_empty());
assert!(messages.lock().expect("messages lock poisoned").is_empty());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
Expand Down
26 changes: 8 additions & 18 deletions crates/contextforge-gateway-rs-lib/tests/support/plugin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use cpex_core::{
hooks::{Extensions, HookHandler, PluginResult, TypedHandlerAdapter, types::cmf_hook_names},
plugin::{Plugin, PluginConfig},
};
use rmcp::model::{CallToolResult, Content, LoggingMessageNotificationParam, ProgressNotificationParam};
use rmcp::model::{CallToolResult, Content, ProgressNotificationParam};
use serde_json::{Value, json};

use super::tool::text;
Expand Down Expand Up @@ -215,21 +215,12 @@ impl HookHandler<CmfHook> for TestPlugin {
let mut modified = payload.clone();
if let Some(ContentPart::ToolResult { content }) =
modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::ToolResult { .. }))
{
if let Ok(mut progress) =
&& let Ok(mut progress) =
serde_json::from_value::<ProgressNotificationParam>(content.content.clone())
{
progress.message = progress.message.map(|message| format!("plugin:{message}"));
content.content = serde_json::to_value(progress).expect("progress serializes");
return PluginResult::modify_payload(modified);
}
if let Ok(mut message) =
serde_json::from_value::<LoggingMessageNotificationParam>(content.content.clone())
{
message.data = json!({ "plugin": "message", "original": message.data });
content.content = serde_json::to_value(message).expect("message serializes");
return PluginResult::modify_payload(modified);
}
{
progress.message = progress.message.map(|message| format!("plugin:{message}"));
content.content = serde_json::to_value(progress).expect("progress serializes");
return PluginResult::modify_payload(modified);
}
PluginResult::allow()
},
Expand Down Expand Up @@ -292,9 +283,8 @@ impl HookHandler<CmfHook> for TestPlugin {
}
}

/// Stream events (progress and logging notifications) run through the same
/// post hook as tool results; result-rewriting behaviors must leave them
/// untouched.
/// Stream events run through the same post hook as tool results;
/// result-rewriting behaviors must leave them untouched.
fn is_tool_result_content(content: &Value) -> bool {
serde_json::from_value::<CallToolResult>(content.clone()).is_ok()
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,7 @@ use rmcp::{
ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt,
model::{
CallToolRequestParams, CallToolResult, Content, ErrorCode, Implementation, InitializeRequestParams,
InitializeResult, LoggingLevel, LoggingMessageNotificationParam, NumberOrString, ProgressNotificationParam,
ProgressToken, ServerCapabilities,
InitializeResult, NumberOrString, ProgressNotificationParam, ProgressToken, ServerCapabilities,
},
service::{RequestContext, Service},
transport::{
Expand Down Expand Up @@ -63,7 +62,6 @@ impl ServerHandler for TestBackend {
.with_server_info(Implementation::new("test-backend", "0.1.0")))
}

#[expect(deprecated, reason = "logging notifications are exercised until the SEP-2577 removal lands in MCP")]
async fn call_tool(
&self,
request: CallToolRequestParams,
Expand Down Expand Up @@ -94,16 +92,6 @@ impl ServerHandler for TestBackend {
"progress_sum" => {
if let Some(progress_token) = cx.meta.get_progress_token() {
for package in 1..=4 {
cx.peer
.notify_logging_message(LoggingMessageNotificationParam {
level: LoggingLevel::Info,
logger: Some("test-backend".to_owned()),
data: serde_json::json!({ "package": package, "state": "started" }),
})
.await
.map_err(|error| {
ErrorData::internal_error(format!("message notification failed: {error}"), None)
})?;
cx.peer
.notify_progress(
ProgressNotificationParam::new(progress_token.clone(), f64::from(package))
Expand Down