diff --git a/Cargo.lock b/Cargo.lock index bde5f59..92224ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -607,6 +607,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tracing", diff --git a/crates/contextforge-gateway-rs-cpex/src/handle.rs b/crates/contextforge-gateway-rs-cpex/src/handle.rs index 3c3231c..02500da 100644 --- a/crates/contextforge-gateway-rs-cpex/src/handle.rs +++ b/crates/contextforge-gateway-rs-cpex/src/handle.rs @@ -13,7 +13,8 @@ use cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult, ErrorCode, ProgressNotificationParam}, + model::{CallToolRequestParams, CallToolResult, ErrorCode}, + serde::{Serialize, de::DeserializeOwned}, }; use tokio::task::JoinHandle; @@ -271,17 +272,20 @@ impl GatewayPluginRuntimeHandle { } } - /// Runs the tool post hooks over a streamed progress notification. Returns - /// `None` when a plugin denies the notification. - pub async fn after_progress_notification( + /// Runs the tool post hooks over a streamed tool event (progress or logging + /// notification). Returns `None` when a plugin denies the event. + pub async fn after_stream_event( &self, tool_name: &str, - progress: ProgressNotificationParam, + event: T, state: Option, - ) -> Result, ErrorData> { + ) -> Result, ErrorData> + where + T: Serialize + DeserializeOwned, + { match state.and_then(|state| state.downcast::().ok()) { - Some(state) => state.runtime.after_progress_notification(tool_name, progress, state.state.clone()).await, - None => Ok(Some(progress)), + Some(state) => state.runtime.after_tool_event(tool_name, event, state.state.clone()).await, + None => Ok(Some(event)), } } } @@ -853,8 +857,7 @@ mod tests { let observations = plugin.observations(); let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; - let event = - runtime.handle().after_progress_notification("sum", progress_event(), None).await.expect("event passes"); + let event = runtime.handle().after_stream_event("sum", progress_event(), None).await.expect("event passes"); assert_eq!(Some("step 1/2"), event.expect("event is kept").message.as_deref()); assert_eq!(0, observations.lock().expect("observations lock poisoned").post_calls); @@ -868,11 +871,8 @@ mod tests { let runtime = runtime_with_plugin(&plugin, plugin_config(&[Arc::clone(&plugin)])).await; let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); - let event = runtime - .handle() - .after_progress_notification("sum", progress_event(), pre.state) - .await - .expect("event passes"); + let event = + runtime.handle().after_stream_event("sum", progress_event(), pre.state).await.expect("event passes"); assert_eq!(Some("plugin:step 1/2"), event.expect("event is kept").message.as_deref()); assert_eq!(1, observations.lock().expect("observations lock poisoned").post_calls); @@ -886,7 +886,7 @@ mod tests { let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); let event = runtime .handle() - .after_progress_notification("sum", progress_event(), pre.state) + .after_stream_event("sum", progress_event(), pre.state) .await .expect("deny drops the event"); @@ -902,7 +902,7 @@ mod tests { let pre = runtime.before_tool_call(&sum_request(1, 2), "sum", "backend").await.expect("pre state is created"); let error = runtime .handle() - .after_progress_notification("sum", progress_event(), pre.state) + .after_stream_event("sum", progress_event(), pre.state) .await .expect_err("invalid rewrite is rejected"); diff --git a/crates/contextforge-gateway-rs-cpex/src/pipeline.rs b/crates/contextforge-gateway-rs-cpex/src/pipeline.rs index 429c3bd..a56be02 100644 --- a/crates/contextforge-gateway-rs-cpex/src/pipeline.rs +++ b/crates/contextforge-gateway-rs-cpex/src/pipeline.rs @@ -2,7 +2,8 @@ use cpex_core::cmf::MessagePayload; use cpex_core::executor::PipelineResult; use rmcp::{ ErrorData, - model::{CallToolResult, ErrorCode, ProgressNotificationParam}, + model::{CallToolResult, ErrorCode}, + serde::de::DeserializeOwned, }; use tracing::warn; @@ -45,23 +46,23 @@ pub(crate) fn effective_post_result(original: CallToolResult, result: &PipelineR } } -pub(crate) fn effective_post_progress( - original: ProgressNotificationParam, - result: &PipelineResult, -) -> Result { +pub(crate) fn effective_post_json(original: T, result: &PipelineResult) -> Result +where + T: DeserializeOwned, +{ let Some(payload) = modified_message_payload(result) else { return Ok(original); }; let Some(content) = tool_result_content(payload) else { return Err(ErrorData { code: ErrorCode::INVALID_PARAMS, - message: "Plugin modified progress notification payload without a tool result".into(), + message: "Plugin modified stream event payload without a tool result".into(), data: None, }); }; - serde_json::from_value(content).map_err(|error| ErrorData { + serde_json::from_value::(content).map_err(|error| ErrorData { code: ErrorCode::INVALID_PARAMS, - message: format!("Plugin modified progress notification payload with invalid JSON: {error}").into(), + message: format!("Plugin modified stream event payload with invalid JSON: {error}").into(), data: None, }) } diff --git a/crates/contextforge-gateway-rs-cpex/src/runtime.rs b/crates/contextforge-gateway-rs-cpex/src/runtime.rs index 3755800..d5b140e 100644 --- a/crates/contextforge-gateway-rs-cpex/src/runtime.rs +++ b/crates/contextforge-gateway-rs-cpex/src/runtime.rs @@ -14,7 +14,8 @@ use cpex_core::{ }; use rmcp::{ ErrorData, - model::{CallToolRequestParams, CallToolResult, ProgressNotificationParam}, + model::{CallToolRequestParams, CallToolResult}, + serde::{Serialize, de::DeserializeOwned}, }; use tokio::sync::Mutex; @@ -23,7 +24,7 @@ use crate::{ error::GatewayPluginRuntimeError, hooks::{RuntimeHookState, ToolArgumentsUpdate, ToolPreCallResult}, pipeline::{ - effective_post_progress, effective_post_result, effective_pre_args, log_pipeline_errors, plugin_denied_error, + effective_post_json, effective_post_result, effective_pre_args, log_pipeline_errors, plugin_denied_error, }, }; @@ -145,24 +146,6 @@ impl GatewayPluginRuntime { result } - /// Runs the tool post pipeline for one event of an in-flight call: the - /// payload is built from the call's `tool_call_id`, the call's context table - /// seeds the pipeline, and the resulting context table is carried back so - /// later events in the same call observe it. - async fn run_tool_post( - &self, - state: &SharedToolCallState, - build_payload: impl FnOnce(&str) -> MessagePayload, - ) -> PipelineResult { - let mut state = state.lock().await; - let payload = build_payload(&state.tool_call_id); - let post_result = self.invoke_tool_post(payload, Some(state.context_table.clone())).await; - if !post_result.is_denied() { - state.context_table = post_result.context_table.clone(); - } - post_result - } - pub(crate) async fn before_tool_call( &self, request: &CallToolRequestParams, @@ -196,37 +179,53 @@ impl GatewayPluginRuntime { return Ok(response); } - let Some(state) = state.and_then(|state| state.downcast::().ok()) else { - return Ok(response); - }; + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(response) }; - let post_result = self.run_tool_post(&state, |id| tool_result_payload(tool_name, &response, id)).await; + let mut state = state.lock().await; + let post_result = self + .invoke_tool_post( + tool_result_payload(tool_name, &response, &state.tool_call_id), + Some(state.context_table.clone()), + ) + .await; if post_result.is_denied() { return Err(plugin_denied_error(post_result)); } + + state.context_table = post_result.context_table.clone(); Ok(effective_post_result(response, &post_result)) } - pub(crate) async fn after_progress_notification( + pub(crate) async fn after_tool_event( &self, tool_name: &str, - progress: ProgressNotificationParam, + event: T, state: Option, - ) -> Result, ErrorData> { + ) -> Result, ErrorData> + where + T: Serialize + DeserializeOwned, + { if !self.has_post_hook { - return Ok(Some(progress)); + return Ok(Some(event)); } - let Some(state) = state.and_then(|state| state.downcast::().ok()) else { - return Ok(Some(progress)); - }; + let state = state.and_then(|state| state.downcast::().ok()); + let Some(state) = state else { return Ok(Some(event)) }; - let content = serde_json::to_value(&progress).unwrap_or(serde_json::Value::Null); - let post_result = - self.run_tool_post(&state, |id| tool_json_result_payload(tool_name, content, false, id)).await; + let content = serde_json::to_value(&event).unwrap_or(serde_json::Value::Null); + let mut state = state.lock().await; + let post_result = self + .invoke_tool_post( + tool_json_result_payload(tool_name, content, false, &state.tool_call_id), + Some(state.context_table.clone()), + ) + .await; if post_result.is_denied() { return Ok(None); } - Ok(Some(effective_post_progress(progress, &post_result)?)) + + state.context_table = post_result.context_table.clone(); + Ok(Some(effective_post_json(event, &post_result)?)) } } diff --git a/crates/contextforge-gateway-rs-lib/Cargo.toml b/crates/contextforge-gateway-rs-lib/Cargo.toml index c735046..af259f9 100644 --- a/crates/contextforge-gateway-rs-lib/Cargo.toml +++ b/crates/contextforge-gateway-rs-lib/Cargo.toml @@ -16,6 +16,7 @@ serde.workspace = true serde_json.workspace = true tracing.workspace = true tokio.workspace = true +tokio-util = "0.7" axum.workspace = true axum-otel-metrics.workspace = true tower-http.workspace = true diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs new file mode 100644 index 0000000..0208cd0 --- /dev/null +++ b/crates/contextforge-gateway-rs-lib/src/gateway/backend_client.rs @@ -0,0 +1,154 @@ +use std::{collections::HashMap, sync::Arc}; + +use contextforge_gateway_rs_cpex::{GatewayPluginRuntimeHandle, RuntimeHookState}; +use rmcp::{ + ClientHandler, Peer, RoleClient, RoleServer, + model::{ + CallToolRequestParams, CallToolResult, ClientRequest, InitializeRequestParams, Meta, ProgressNotificationParam, + ProgressToken, Request, ServerResult, + }, + serde::{Serialize, de::DeserializeOwned}, + service::{NotificationContext, PeerRequestOptions, ServiceError}, +}; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +#[derive(Clone)] +pub(crate) struct GatewayBackendClient { + initialize_request: InitializeRequestParams, + plugin_runtime: Option, + in_flight_calls: Arc>>>, +} + +#[derive(Debug)] +struct InFlightToolCall { + progress_token: Option, + tool_name: String, + post_state: Option, + downstream: Peer, +} + +impl GatewayBackendClient { + pub(crate) fn new( + initialize_request: InitializeRequestParams, + plugin_runtime: Option, + ) -> Self { + Self { initialize_request, plugin_runtime, in_flight_calls: Arc::default() } + } + + pub(crate) async fn track_tool_call( + &self, + tool_name: String, + downstream: Peer, + progress_token: Option, + post_state: Option, + ) { + debug!("track_tool_call {tool_name} {progress_token:?} {post_state:?}"); + let call = Arc::new(InFlightToolCall { progress_token, tool_name, post_state, downstream }); + let mut calls = self.in_flight_calls.lock().await; + + if let Some(token) = &call.progress_token { + calls.entry(token.clone()).or_insert_with(|| Arc::clone(&call)); + } + drop(calls); + } + + pub(crate) async fn stop_tracking_tool_call(&self, progress_token: Option) { + debug!("stop_tracking_tool_call {progress_token:?}"); + + let mut calls = self.in_flight_calls.lock().await; + + if let Some(token) = &progress_token { + calls.remove(token); + } + drop(calls); + } + + async fn progress_call(&self, progress_token: &ProgressToken) -> Option> { + let calls = self.in_flight_calls.lock().await; + calls.get(progress_token).cloned() + } + + async fn stream_event_post_hook(&self, call: &InFlightToolCall, event: T) -> Option + where + T: Serialize + DeserializeOwned, + { + let Some(plugin_runtime) = &self.plugin_runtime else { + return Some(event); + }; + match plugin_runtime.after_stream_event(&call.tool_name, event, call.post_state.clone()).await { + Ok(event) => event, + Err(error) => { + warn!("call_tool: plugin rejected backend notification: {error:?}"); + None + }, + } + } +} + +impl std::fmt::Debug for GatewayBackendClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GatewayBackendClient") + .field("initialize_request", &self.initialize_request) + .finish_non_exhaustive() + } +} + +impl ClientHandler for GatewayBackendClient { + fn get_info(&self) -> InitializeRequestParams { + self.initialize_request.clone() + } + + async fn on_progress(&self, progress: ProgressNotificationParam, _context: NotificationContext) { + let Some(call) = self.progress_call(&progress.progress_token).await else { + debug!( + "call_tool: dropping backend progress notification with unknown token {:?}", + progress.progress_token + ); + return; + }; + debug!("Processing Progress Notification {progress:?} {call:?}"); + let Some(progress) = self.stream_event_post_hook(&call, progress).await else { + return; + }; + if let Err(error) = call.downstream.notify_progress(progress).await { + warn!("call_tool: unable to forward backend progress notification downstream: {error:?}"); + } + } +} + +/// Calls the tool on the backend, keeping the downstream progress token on +/// the request (`Peer::call_tool` would stamp an auto-generated token over it +/// at serialization) and relaying a downstream cancellation to the backend. +pub(crate) async fn call_backend_tool( + peer: &Peer, + request: CallToolRequestParams, + progress_token: Option, + cancellation: CancellationToken, +) -> Result { + let mut options = PeerRequestOptions::no_options(); + if let Some(progress_token) = progress_token { + let mut meta = Meta::new(); + meta.set_progress_token(progress_token); + options.meta = Some(meta); + } + let mut handle = + peer.send_cancellable_request(ClientRequest::CallToolRequest(Request::new(request)), options).await?; + let response = tokio::select! { + response = &mut handle.rx => Some(response), + () = cancellation.cancelled() => None, + }; + + let Some(response) = response else { + let reason = "tool call cancelled by the downstream client".to_owned(); + if let Err(error) = handle.cancel(Some(reason.clone())).await { + warn!("call_tool: unable to relay cancellation to the backend: {error:?}"); + } + return Err(ServiceError::Cancelled { reason: Some(reason) }); + }; + match response.map_err(|_| ServiceError::TransportClosed)?? { + ServerResult::CallToolResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } +} diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs index 18b9764..93825ad 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs @@ -24,7 +24,10 @@ use tokio::sync::Mutex; use tracing::{debug, info, warn}; use typed_builder::TypedBuilder; -use super::mcp_call_validator::AuthorizedCallValidator; +use super::{ + backend_client::{GatewayBackendClient, call_backend_tool}, + mcp_call_validator::AuthorizedCallValidator, +}; pub use crate::gateway::session_store::LocalUserSessionStore; use crate::{ SessionId, @@ -72,7 +75,7 @@ pub struct BackendTransportKey { session_id: String, } -type McpClientService = Arc>; +type McpClientService = Arc>; #[derive(Debug)] pub struct ServiceHolder { @@ -149,7 +152,7 @@ where .iter() .map(|(name, backend)| { let client = self.http_client.clone(); - let request = request.clone(); + let backend_client = GatewayBackendClient::new(request.clone(), self.plugin_runtime.clone()); let backend_url = backend.url.clone(); let downstream_session_id = downstream_session_id.clone(); @@ -172,7 +175,7 @@ where let config = StreamableHttpClientTransportConfig::with_uri(backend_url.to_string()) .custom_headers(headers); let transport = StreamableHttpClientTransport::with_client(client, config); - let maybe_running_service = request.serve(transport).await; + let maybe_running_service = backend_client.serve(transport).await; if let Ok(running_service) = maybe_running_service { info!("initialize: intialized for {downstream_session_id:?} {name:?}"); (name, Some(running_service)) @@ -183,7 +186,7 @@ where }) }).collect(); - let initialization_results: Vec<(&String, Option>)> = + let initialization_results: Vec<(&String, Option>)> = futures::future::join_all(tasks).await; let (capabilities, backend_services): (Vec<_>, Vec<_>) = initialization_results @@ -359,7 +362,14 @@ where pre_result.arguments.apply_to_request(&mut routed_request, &tool_name); let service_name = target_service.name.clone(); - let response = service.call_tool(routed_request).await; + let progress_token = cx.meta.get_progress_token(); + service + .service() + .track_tool_call(tool_name.clone(), cx.peer.clone(), progress_token.clone(), post_state.clone()) + .await; + let response = call_backend_tool(service.peer(), routed_request, progress_token.clone(), cx.ct.clone()).await; + service.service().stop_tracking_tool_call(progress_token.clone()).await; + let response = response.map_err(|error| { warn!("call_tool: backend {service_name} {error:?}"); ErrorData { diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/mod.rs b/crates/contextforge-gateway-rs-lib/src/gateway/mod.rs index 5339216..978229a 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/mod.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/mod.rs @@ -1,3 +1,4 @@ +mod backend_client; mod mcp_call_validator; pub(crate) mod mcp_gateway; mod session_manager; diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs index f2bc5f5..2feb68c 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs @@ -1,17 +1,311 @@ mod support; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; +use contextforge_gateway_rs_cpex::CpexRuntimeRegistry; use cpex_core::cmf::Role; use cpex_core::hooks::types::cmf_hook_names; -use rmcp::model::ErrorCode; +use rmcp::{ + ClientHandler, + model::{ + CallToolRequestParams, CallToolResult, ClientCapabilities, ClientRequest, ErrorCode, Implementation, + InitializeRequestParams, Meta, NumberOrString, ProgressNotificationParam, ProgressToken, Request, ServerResult, + }, + service::{NotificationContext, PeerRequestOptions, RequestHandle, RoleClient, RunningService}, +}; use serde_json::Value; use support::{ - POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, TestPlugin, error_code, - runtime_with_post, runtime_with_pre, runtime_with_pre_and_post, start_gateway, sum_request, text, + POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, RunningGateway, TestPlugin, + error_code, runtime_with_post, runtime_with_pre, runtime_with_pre_and_post, start_gateway, + start_gateway_with_json_backend_responses, sum_request, text, token, }; +type Recorded = Arc>>; + +#[derive(Clone, Default)] +struct RecordingClient { + progress: Recorded, +} + +impl ClientHandler for RecordingClient { + fn get_info(&self) -> InitializeRequestParams { + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("recording-test-client", "0.1.0"), + ) + } + + async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext) { + self.progress.lock().expect("progress lock poisoned").push(params); + } +} + +async fn call_progress_sum( + gateway: &RunningGateway, + user: &str, +) -> (CallToolResult, Recorded) { + let (result, progress) = send_progress_sum(gateway, user).await; + wait_for_event_count(&progress, 4).await; + (result, progress) +} + +async fn send_progress_sum( + gateway: &RunningGateway, + user: &str, +) -> (CallToolResult, Recorded) { + let client = RecordingClient::default(); + let progress = Arc::clone(&client.progress); + let service = gateway.connect_with_handler(user, client).await; + let token = ProgressToken(NumberOrString::String("package-progress".into())); + let handle = send_progress_call(&service, &format!("{}-progress_sum", gateway.backend_name), token).await; + + let ServerResult::CallToolResult(result) = handle.await_response().await.expect("progress_sum call succeeds") + else { + panic!("expected call tool result"); + }; + (result, progress) +} + +/// Sends a `tools/call` carrying `progress_token` and returns the in-flight +/// request handle without awaiting it. +async fn send_progress_call( + service: &RunningService, + tool_name: &str, + progress_token: ProgressToken, +) -> RequestHandle { + let request = CallToolRequestParams::new(tool_name.to_owned()); + let mut options = PeerRequestOptions::no_options(); + options.meta = Some(Meta::with_progress_token(progress_token)); + service + .send_cancellable_request(ClientRequest::CallToolRequest(Request::new(request)), options) + .await + .expect("progress request is sent") +} + +async fn wait_for_event_count(events: &StdMutex>, expected: usize) { + for _ in 0..50 { + if events.lock().expect("events lock poisoned").len() >= expected { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + panic!("timed out waiting for {expected} recorded events"); +} + +fn raw_mcp_request( + client: &reqwest::Client, + gateway: &RunningGateway, + user: &str, + session_id: Option<&str>, + body: &Value, +) -> reqwest::RequestBuilder { + let mut request = client + .post(gateway.gateway_url()) + .bearer_auth(token(user)) + .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::ACCEPT, "application/json, text/event-stream") + .json(body); + if let Some(session_id) = session_id { + request = request.header("Mcp-Session-Id", session_id).header("MCP-Protocol-Version", "2025-11-25"); + } + request +} + +fn raw_tool_call(tool_name: &str, request_id: i64, progress_token: i64) -> Value { + serde_json::json!({ + "method": "tools/call", + "params": { + "name": tool_name, + "arguments": {}, + "_meta": { "progressToken": progress_token } + }, + "jsonrpc": "2.0", + "id": request_id + }) +} + +fn sse_data_values(body: &str) -> Vec { + let values = body + .lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .filter(|data| !data.is_empty()) + .map(|data| serde_json::from_str(data).expect("SSE data is JSON")) + .collect::>(); + if !values.is_empty() { + return values; + } + let body = body.trim(); + if body.is_empty() { Vec::new() } else { vec![serde_json::from_str(body).expect("JSON response body")] } +} + +fn assert_raw_progress_stream(body: &str, response_id: i64, progress_token: i64) { + let messages = sse_data_values(body); + let progress = messages + .iter() + .filter(|message| message.get("method").and_then(Value::as_str) == Some("notifications/progress")) + .collect::>(); + assert_eq!(4, progress.len(), "unexpected progress events in body: {body}"); + assert!( + progress + .iter() + .all(|message| message.pointer("/params/progressToken").and_then(Value::as_i64) == Some(progress_token)), + "progress events with foreign tokens in body: {body}" + ); + let result = messages + .iter() + .find(|message| message.get("id").and_then(Value::as_i64) == Some(response_id)) + .unwrap_or_else(|| panic!("missing response id {response_id} in body: {body}")); + assert_eq!(Some("completed 4 packages"), result.pointer("/result/content/0/text").and_then(Value::as_str)); +} + +async fn start_raw_mcp_session(client: &reqwest::Client, gateway: &RunningGateway, user: &str) -> String { + let initialize = raw_mcp_request( + client, + gateway, + user, + None, + &serde_json::json!({ + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "clientInfo": { "name": "raw-test-client", "version": "0.1.0" } + }, + "jsonrpc": "2.0", + "id": 0 + }), + ) + .send() + .await + .expect("initialize request is sent"); + assert!(initialize.status().is_success(), "initialize failed: {initialize:?}"); + let session_id = initialize + .headers() + .get("mcp-session-id") + .expect("initialize response has MCP session id") + .to_str() + .expect("MCP session id is valid") + .to_owned(); + let _initialize_body = initialize.text().await.expect("initialize body is read"); + + let initialized = raw_mcp_request( + client, + gateway, + user, + Some(&session_id), + &serde_json::json!({ "method": "notifications/initialized", "jsonrpc": "2.0" }), + ) + .send() + .await + .expect("initialized notification is sent"); + assert!(initialized.status().is_success(), "initialized notification failed: {initialized:?}"); + let _initialized_body = initialized.text().await.expect("initialized body is read"); + + session_id +} + +async fn read_concurrent_raw_progress_streams( + first: reqwest::RequestBuilder, + second: reqwest::RequestBuilder, +) -> (String, String) { + let (first, second) = + tokio::time::timeout(std::time::Duration::from_secs(3), async { tokio::join!(first.send(), second.send()) }) + .await + .expect("both raw progress requests receive response headers"); + let first = first.expect("first raw progress request succeeds"); + let second = second.expect("second raw progress request succeeds"); + assert!(first.status().is_success(), "first raw progress request failed: {first:?}"); + assert!(second.status().is_success(), "second raw progress request failed: {second:?}"); + + let (first_body, second_body) = + tokio::time::timeout(std::time::Duration::from_secs(3), async { tokio::join!(first.text(), second.text()) }) + .await + .expect("both raw progress streams complete"); + (first_body.expect("first raw progress body is read"), second_body.expect("second raw progress body is read")) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +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 service = gateway.connect_with_handler("admin@example.com", client).await; + let tool_name = format!("{}-progress_sum", gateway.backend_name); + + let first = send_progress_call(&service, &tool_name, ProgressToken(NumberOrString::Number(1))).await; + let second = send_progress_call(&service, &tool_name, ProgressToken(NumberOrString::Number(2))).await; + + let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(3), async { + tokio::join!(first.await_response(), second.await_response()) + }) + .await + .expect("both concurrent progress_sum calls complete"); + + let ServerResult::CallToolResult(first) = first.expect("first progress_sum call succeeds") else { + panic!("expected first call tool result"); + }; + let ServerResult::CallToolResult(second) = second.expect("second progress_sum call succeeds") else { + panic!("expected second call tool result"); + }; + assert_eq!("completed 4 packages", text(&first)); + assert_eq!("completed 4 packages", text(&second)); + + wait_for_event_count(&progress, 8).await; + let progress = progress.lock().expect("progress lock poisoned"); + let first_count = progress + .iter() + .filter(|notification| notification.progress_token == ProgressToken(NumberOrString::Number(1))) + .count(); + let second_count = progress + .iter() + .filter(|notification| notification.progress_token == ProgressToken(NumberOrString::Number(2))) + .count(); + assert_eq!(4, first_count); + assert_eq!(4, second_count); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn raw_streamable_http_concurrent_progress_calls_complete_without_plugins() { + let gateway = start_gateway("admin@example.com", false, Arc::new(CpexRuntimeRegistry::default())).await; + let client = reqwest::Client::new(); + let session_id = start_raw_mcp_session(&client, &gateway, "admin@example.com").await; + + let tool_name = format!("{}-progress_sum", gateway.backend_name); + let first = + raw_mcp_request(&client, &gateway, "admin@example.com", Some(&session_id), &raw_tool_call(&tool_name, 2, 1)); + let second = + raw_mcp_request(&client, &gateway, "admin@example.com", Some(&session_id), &raw_tool_call(&tool_name, 3, 2)); + let (first_body, second_body) = read_concurrent_raw_progress_streams(first, second).await; + + assert_raw_progress_stream(&first_body, 2, 1); + assert_raw_progress_stream(&second_body, 3, 2); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn backend_generated_progress_tokens_are_dropped() { + let gateway = start_gateway("admin@example.com", false, Arc::new(CpexRuntimeRegistry::default())).await; + let client = RecordingClient::default(); + let progress = Arc::clone(&client.progress); + let service = gateway.connect_with_handler("admin@example.com", client).await; + + let tool_name = format!("{}-progress_counter_tokens", gateway.backend_name); + let handle = send_progress_call(&service, &tool_name, ProgressToken(NumberOrString::Number(10))).await; + + let ServerResult::CallToolResult(result) = + handle.await_response().await.expect("progress_counter_tokens call succeeds") + else { + panic!("expected call tool result"); + }; + assert_eq!("completed 4 packages", text(&result)); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert!( + progress.lock().expect("progress lock poisoned").is_empty(), + "backend-generated progress tokens must not be forwarded" + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn disabled_runtime_does_not_invoke_registered_plugin() { let pre_plugin = @@ -70,6 +364,81 @@ async fn post_hook_receives_backend_result_and_modifies_client_result() { assert_eq!(Some("3".to_owned()), observations.post_result_text); } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +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) = 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 observations = observations.lock().expect("observations lock poisoned"); + // four progress notifications plus the final tool result + 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_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) = 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()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn post_hook_deny_drops_progress_notifications_without_failing_call() { + let plugin = + Arc::new(TestPlugin::new("post-stream-deny", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_stream_event_deny()); + let observations = plugin.observations(); + let runtime = runtime_with_post(plugin).await; + + let gateway = start_gateway("admin@example.com", true, runtime).await; + let (result, progress) = send_progress_sum(&gateway, "admin@example.com").await; + + assert_eq!("completed 4 packages", text(&result)); + // four denied progress notifications plus the final tool result + for _ in 0..50 { + 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!(5, observations.post_calls); + assert!(progress.lock().expect("progress lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn downstream_cancellation_is_relayed_to_backend() { + let gateway = start_gateway("admin@example.com", true, Arc::new(CpexRuntimeRegistry::default())).await; + let service = gateway.connect("admin@example.com").await; + + let request = CallToolRequestParams::new(format!("{}-wait_for_cancellation", gateway.backend_name)); + let handle = service + .send_cancellable_request( + ClientRequest::CallToolRequest(Request::new(request)), + PeerRequestOptions::no_options(), + ) + .await + .expect("wait_for_cancellation request is sent"); + wait_for_event_count(&gateway.backend_state.calls, 1).await; + + handle.cancel(Some("client gave up".to_owned())).await.expect("cancellation is sent"); + wait_for_event_count(&gateway.backend_state.cancellations, 1).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn post_hook_can_return_raw_cmf_result_content() { let plugin = Arc::new(TestPlugin::new("post", vec![cmf_hook_names::TOOL_POST_INVOKE]).with_raw_post_rewrite()); diff --git a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs index 7d0226a..ca80e50 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/mod.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/mod.rs @@ -17,7 +17,7 @@ pub(crate) use list_tools_gateway::{ pub(crate) use plugin::{ POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, TestPlugin, TestPluginFactory, }; -pub(crate) use plugin_gateway::start_gateway; +pub(crate) use plugin_gateway::{RunningGateway, start_gateway, start_gateway_with_json_backend_responses}; pub(crate) use runtime::{runtime_with_post, runtime_with_pre, runtime_with_pre_and_post}; pub(crate) use tool::{error_code, sum_request, text}; pub(crate) use user_config_store::MemoryUserConfigStore; diff --git a/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs b/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs index ca7d8d6..b4bca58 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/plugin.rs @@ -12,8 +12,8 @@ use cpex_core::{ hooks::{Extensions, HookHandler, PluginResult, TypedHandlerAdapter, types::cmf_hook_names}, plugin::{Plugin, PluginConfig}, }; -use rmcp::model::{CallToolResult, Content}; -use serde_json::json; +use rmcp::model::{CallToolResult, Content, ProgressNotificationParam}; +use serde_json::{Value, json}; use super::tool::text; @@ -33,7 +33,7 @@ pub(crate) struct Observations { pub(crate) pre_payload_role: Option, pub(crate) pre_tool_call_id: Option, pub(crate) post_payload_name: Option, - pub(crate) post_tool_call_id: Option, + pub(crate) post_tool_call_ids: Vec, pub(crate) post_result_text: Option, } @@ -53,6 +53,8 @@ pub(crate) enum PostBehavior { Allow, Rewrite, RewriteRaw, + RewriteStreamEvents, + DenyStreamEvents, Deny, RequireContext, } @@ -103,6 +105,16 @@ impl TestPlugin { self } + pub(crate) fn with_stream_event_rewrite(mut self) -> Self { + self.post_behavior = PostBehavior::RewriteStreamEvents; + self + } + + pub(crate) fn with_stream_event_deny(mut self) -> Self { + self.post_behavior = PostBehavior::DenyStreamEvents; + self + } + pub(crate) fn with_pre_deny(mut self) -> Self { self.pre_behavior = PreBehavior::Deny; self @@ -154,7 +166,7 @@ impl HookHandler for TestPlugin { observations.post_calls += 1; if let Some(result) = payload.message.get_tool_results().first() { observations.post_payload_name = Some(result.tool_name.clone()); - observations.post_tool_call_id = Some(result.tool_call_id.clone()); + observations.post_tool_call_ids.push(result.tool_call_id.clone()); } observations.post_result_text = Some(cmf_result_text(payload)); } else { @@ -177,6 +189,9 @@ impl HookHandler for TestPlugin { if let Some(ContentPart::ToolResult { content }) = modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::ToolResult { .. })) { + if !is_tool_result_content(&content.content) { + return PluginResult::allow(); + } content.content = serde_json::to_value(CallToolResult::success(vec![Content::text(format!( "post:{result_text}" ))])) @@ -189,10 +204,38 @@ impl HookHandler for TestPlugin { if let Some(ContentPart::ToolResult { content }) = modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::ToolResult { .. })) { + if !is_tool_result_content(&content.content) { + return PluginResult::allow(); + } content.content = json!("raw-post"); } PluginResult::modify_payload(modified) }, + PostBehavior::RewriteStreamEvents => { + let mut modified = payload.clone(); + if let Some(ContentPart::ToolResult { content }) = + modified.message.content.iter_mut().find(|part| matches!(part, ContentPart::ToolResult { .. })) + && let Ok(mut progress) = + serde_json::from_value::(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); + } + PluginResult::allow() + }, + PostBehavior::DenyStreamEvents => { + let is_stream_event = payload + .message + .get_tool_results() + .first() + .is_some_and(|result| !is_tool_result_content(&result.content)); + if is_stream_event { + PluginResult::deny(PluginViolation::new("stream_denied", "stream denied")) + } else { + PluginResult::allow() + } + }, PostBehavior::Deny => PluginResult::deny( PluginViolation::new("post_denied", "post denied") .with_proto_error_code(i64::from(POST_DENY_ERROR_CODE)), @@ -240,6 +283,12 @@ impl HookHandler for TestPlugin { } } +/// Progress notifications 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::(content.clone()).is_ok() +} + fn cmf_result_text(payload: &MessagePayload) -> String { payload .message diff --git a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs index a3b5c97..d330313 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs @@ -13,12 +13,12 @@ use contextforge_gateway_rs_lib::{Config, Gateway, UpstreamConnectionMode, UserC use futures::FutureExt; use http::{HeaderMap, HeaderValue}; use rmcp::{ - ErrorData, RoleServer, ServerHandler, ServiceExt, + ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, model::{ CallToolRequestParams, CallToolResult, Content, ErrorCode, Implementation, InitializeRequestParams, - InitializeResult, ServerCapabilities, + InitializeResult, NumberOrString, ProgressNotificationParam, ProgressToken, ServerCapabilities, }, - service::RequestContext, + service::{RequestContext, Service}, transport::{ StreamableHttpClientTransport, StreamableHttpServerConfig, StreamableHttpService, streamable_http_client::StreamableHttpClientTransportConfig, @@ -44,6 +44,7 @@ pub(crate) struct BackendObservation { #[derive(Clone, Default)] pub(crate) struct BackendState { pub(crate) calls: Arc>>, + pub(crate) cancellations: Arc>>, } #[derive(Clone)] @@ -64,7 +65,7 @@ impl ServerHandler for TestBackend { async fn call_tool( &self, request: CallToolRequestParams, - _cx: RequestContext, + cx: RequestContext, ) -> Result { self.state .calls @@ -88,6 +89,52 @@ impl ServerHandler for TestBackend { .ok_or_else(|| ErrorData::invalid_params("sum requires numeric b", None))?; Ok(CallToolResult::success(vec![Content::text((a + b).to_string())])) }, + "progress_sum" => { + if let Some(progress_token) = cx.meta.get_progress_token() { + for package in 1..=4 { + cx.peer + .notify_progress( + ProgressNotificationParam::new(progress_token.clone(), f64::from(package)) + .with_total(4.0) + .with_message(format!("package {package}/4")), + ) + .await + .map_err(|error| { + ErrorData::internal_error(format!("progress notification failed: {error}"), None) + })?; + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + Ok(CallToolResult::success(vec![Content::text("completed 4 packages")])) + }, + "progress_counter_tokens" => { + for package in 1..=4i32 { + cx.peer + .notify_progress( + ProgressNotificationParam::new( + ProgressToken(NumberOrString::Number(i64::from(package))), + f64::from(package), + ) + .with_total(4.0) + .with_message(format!("package {package}/4")), + ) + .await + .map_err(|error| { + ErrorData::internal_error(format!("progress notification failed: {error}"), None) + })?; + tokio::time::sleep(Duration::from_millis(10)).await; + } + Ok(CallToolResult::success(vec![Content::text("completed 4 packages")])) + }, + "wait_for_cancellation" => { + cx.ct.cancelled().await; + self.state + .cancellations + .lock() + .expect("backend cancellations lock poisoned") + .push(request.name.to_string()); + Ok(CallToolResult::success(vec![Content::text("cancelled")])) + }, _ => Err(ErrorData { code: ErrorCode::METHOD_NOT_FOUND, message: format!("unknown tool {}", request.name).into(), @@ -105,10 +152,25 @@ pub(crate) struct RunningGateway { } impl RunningGateway { + pub(crate) fn gateway_url(&self) -> &str { + &self.gateway_url + } + pub(crate) async fn connect( &self, user: &str, ) -> rmcp::service::RunningService { + self.connect_with_handler(user, InitializeRequestParams::default()).await + } + + pub(crate) async fn connect_with_handler( + &self, + user: &str, + handler: S, + ) -> rmcp::service::RunningService + where + S: Service + Send + Sync + Clone + 'static, + { let deadline = Instant::now() + CLIENT_CONNECT_TIMEOUT; loop { let mut headers = HeaderMap::new(); @@ -121,7 +183,7 @@ impl RunningGateway { client, StreamableHttpClientTransportConfig::with_uri(self.gateway_url.clone()), ); - match InitializeRequestParams::default().serve(transport).await { + match handler.clone().serve(transport).await { Ok(service) => return service, Err(error) if Instant::now() < deadline => { let _ = error; @@ -146,13 +208,22 @@ pub(crate) async fn start_gateway( runtime_plugins_enabled: bool, plugin_runtime: Arc, ) -> RunningGateway { - start_gateway_with_runtime(user, runtime_plugins_enabled, plugin_runtime).await + start_gateway_with_runtime(user, runtime_plugins_enabled, plugin_runtime, false).await +} + +pub(crate) async fn start_gateway_with_json_backend_responses( + user: &str, + runtime_plugins_enabled: bool, + plugin_runtime: Arc, +) -> RunningGateway { + start_gateway_with_runtime(user, runtime_plugins_enabled, plugin_runtime, true).await } async fn start_gateway_with_runtime( user: &str, runtime_plugins_enabled: bool, plugin_runtime: Arc, + json_backend_responses: bool, ) -> RunningGateway { let port_lock = Arc::clone(GATEWAY_PORT_LOCK.get_or_init(|| Arc::new(TokioMutex::new(())))); let port_guard = port_lock.lock().await; @@ -169,7 +240,7 @@ async fn start_gateway_with_runtime( move || Ok(TestBackend { state: backend_state.clone() }) }, LocalSessionManager::default().into(), - StreamableHttpServerConfig::default(), + StreamableHttpServerConfig::default().with_json_response(json_backend_responses), ); let backend_router = axum::Router::new().route_service("/mcp", backend_service);