From 7d1492ffabc9a8d734d729aa5d184823f143de8b Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 23 Oct 2026 23:39:13 +0000 Subject: [PATCH] Isolate MCP inputs and timeouts for callers sharing a connection (#57196) ## Why Concurrent requests on a shared MCP connection need to route prompts to their own callers and pause only their own timeouts while awaiting input. ## What changed - Add `call_tool_with_elicitation` and `read_resource_with_elicitation` with per-request callbacks and timeout pause state, requiring MCP protocol `2026-07-28`. - Share continuation handling between tool calls and resource reads, and assign globally unique input IDs across requests and connections. - Add `enable_request_cancellation` to extend cancellation on dropped response futures to catalog, resource, and custom requests. ## Testing Add or extend tests for concurrent form and URL input routing, unique prompt IDs, independent native verification cancellation and timeouts, resource continuations, and remote cancellation when resource or custom request futures are dropped. GitOrigin-RevId: ad5a3cd0e289726828a694d00c20870e2d1108c6 --- .../src/elicitation_client_service.rs | 13 + codex-rs/rmcp-client/src/lib.rs | 4 +- .../rmcp-client/src/request_cancellation.rs | 55 +++ .../src/{tool_input.rs => request_input.rs} | 128 +++++-- codex-rs/rmcp-client/src/rmcp_client.rs | 342 +++++++++++++++--- codex-rs/rmcp-client/src/tool_input_tests.rs | 8 +- codex-rs/rmcp-client/src/tool_request.rs | 49 --- .../user_verification_cancellation_tests.rs | 77 ++++ codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs | 102 ++++-- .../native_verification_tests.rs | 79 +++- 10 files changed, 668 insertions(+), 189 deletions(-) create mode 100644 codex-rs/rmcp-client/src/request_cancellation.rs rename codex-rs/rmcp-client/src/{tool_input.rs => request_input.rs} (65%) delete mode 100644 codex-rs/rmcp-client/src/tool_request.rs diff --git a/codex-rs/rmcp-client/src/elicitation_client_service.rs b/codex-rs/rmcp-client/src/elicitation_client_service.rs index ef2fbacd92..8a54bc9367 100644 --- a/codex-rs/rmcp-client/src/elicitation_client_service.rs +++ b/codex-rs/rmcp-client/src/elicitation_client_service.rs @@ -94,6 +94,19 @@ impl Drop for PendingVerification { } impl ElicitationClientService { + /// Retains protocol validation and cancellation dispatch while isolating the caller's inputs. + pub(crate) fn for_request( + &self, + send_elicitation: Arc, + pause_state: ElicitationPauseState, + ) -> Self { + Self { + send_elicitation, + pause_state, + ..self.clone() + } + } + pub(crate) fn new( client_info: ClientConfig, send_elicitation: SendElicitation, diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index a46b300771..0cdfed8ca7 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -27,13 +27,13 @@ mod ordered_byte_transport; mod perform_oauth_login; mod program_resolver; mod protocol_mode; +mod request_cancellation; +mod request_input; mod rmcp_client; mod service_error; mod startup_error; mod stdio_server_launcher; -mod tool_input; mod tool_list_changes; -mod tool_request; mod trace_context; mod user_verification; mod utils; diff --git a/codex-rs/rmcp-client/src/request_cancellation.rs b/codex-rs/rmcp-client/src/request_cancellation.rs new file mode 100644 index 0000000000..8f0d9cadc5 --- /dev/null +++ b/codex-rs/rmcp-client/src/request_cancellation.rs @@ -0,0 +1,55 @@ +//! Cancel remote MCP work when its local response future is dropped. + +use rmcp::Peer; +use rmcp::RoleClient; +use rmcp::model::CancelledNotificationParam; +use rmcp::model::ClientRequest; +use rmcp::model::RequestId; +use rmcp::model::ServerResult; +use rmcp::service::PeerRequestOptions; +use rmcp::service::ServiceError; + +/// Whether dropping a caller's future should cancel its request on the peer. +#[derive(Clone, Copy)] +pub(crate) enum RequestCancellation { + Disabled, + OnDrop, +} + +impl RequestCancellation { + /// Send a request and tie its remote lifetime to the caller when requested. + pub(crate) async fn send( + self, + peer: &Peer, + request: ClientRequest, + options: PeerRequestOptions, + ) -> Result { + let handle = peer.send_request_with_option(request, options).await?; + let mut cancellation = CancelOnDrop(match self { + Self::Disabled => None, + Self::OnDrop => Some((peer.clone(), handle.id.clone())), + }); + let result = handle.await_response().await; + cancellation.0 = None; + result + } +} + +struct CancelOnDrop(Option<(Peer, RequestId)>); + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + if let Some((peer, id)) = self.0.take() + && let Ok(runtime) = tokio::runtime::Handle::try_current() + { + runtime.spawn(async move { + let _ = peer + .notify_cancelled(CancelledNotificationParam::new( + Some(id), + Some("MCP response future dropped".into()), + )) + .await; + }); + } + } +} diff --git a/codex-rs/rmcp-client/src/tool_input.rs b/codex-rs/rmcp-client/src/request_input.rs similarity index 65% rename from codex-rs/rmcp-client/src/tool_input.rs rename to codex-rs/rmcp-client/src/request_input.rs index e06a8aaa26..9b39fe9667 100644 --- a/codex-rs/rmcp-client/src/tool_input.rs +++ b/codex-rs/rmcp-client/src/request_input.rs @@ -1,17 +1,24 @@ -//! Drives modern tool continuations, including the native verification extension. -//! Inputs use the existing client service; proofs never bypass its validation. +//! Drives modern request continuations, including the native verification extension. +//! Inputs use the caller's client service; proofs never bypass its validation. //! Session recovery cannot replay a submitted proof; connection closure drops pending inputs. use std::collections::BTreeMap; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::time::Duration; use rmcp::RoleClient; +use rmcp::model::CallToolRequest; use rmcp::model::CallToolRequestParams; use rmcp::model::CallToolResult; +use rmcp::model::ClientRequest; use rmcp::model::DEFAULT_MRTR_MAX_ROUNDS; use rmcp::model::GetExtensions; use rmcp::model::GetMeta; use rmcp::model::InputRequest; +use rmcp::model::ReadResourceRequest; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; use rmcp::model::ServerRequest; use rmcp::model::ServerResult; @@ -26,28 +33,71 @@ use serde_json::Value; use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::request_cancellation::RequestCancellation; + +static NEXT_INPUT_ID: AtomicU64 = AtomicU64::new(0); #[derive(Deserialize)] #[serde(rename_all = "camelCase")] -struct ToolInput { +struct RequiredInput { input_requests: Option>, request_state: Option, } pub(crate) async fn call_tool( service: &RunningService, - mut params: CallToolRequestParams, + input_service: &ElicitationClientService, + params: CallToolRequestParams, ) -> Result { + match request( + service, + input_service, + ClientRequest::CallToolRequest(CallToolRequest::new(params)), + RequestCancellation::OnDrop, + ) + .await? + { + ServerResult::CallToolResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } +} + +pub(crate) async fn read_resource( + service: &RunningService, + input_service: &ElicitationClientService, + params: ReadResourceRequestParams, + cancellation: RequestCancellation, +) -> Result { + match request( + service, + input_service, + ClientRequest::ReadResourceRequest(ReadResourceRequest::new(params)), + cancellation, + ) + .await? + { + ServerResult::ReadResourceResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } +} + +async fn request( + service: &RunningService, + input_service: &ElicitationClientService, + mut request: ClientRequest, + cancellation: RequestCancellation, +) -> Result { let mut state_only_rounds = 0; let mut submitted_verification_proof = false; for round in 0..DEFAULT_MRTR_MAX_ROUNDS { - let response = crate::tool_request::call_tool( - service.peer(), - params.clone(), - PeerRequestOptions::no_options(), - ) - .await; - let (id, result) = response.map_err(|error| { + let response = cancellation + .send( + service.peer(), + request.clone(), + PeerRequestOptions::no_options(), + ) + .await; + let result = response.map_err(|error| { if submitted_verification_proof && let ServiceError::TransportSend(transport) = &error && !matches!( @@ -60,15 +110,18 @@ pub(crate) async fn call_tool( // Session recovery must not replay the original operation after a // continuation (which may already have consumed a signed proof). // Auth challenges already return without replaying the operation. - invalid("MCP tool continuation failed; the tool was not restarted") + invalid("MCP continuation failed; the operation was not restarted") } else { error } })?; let input = match result { - ServerResult::CallToolResult(result) => return Ok(result), - ServerResult::InputRequiredResult(result) => serde_json::to_value(result) - .map_err(|_| invalid("invalid MCP tool input request"))?, + result @ (ServerResult::CallToolResult(_) | ServerResult::ReadResourceResult(_)) => { + return Ok(result); + } + ServerResult::InputRequiredResult(result) => { + serde_json::to_value(result).map_err(|_| invalid("invalid MCP input request"))? + } // RMCP's InputRequest union intentionally excludes custom methods. ServerResult::CustomResult(result) if result.0.get("resultType").and_then(Value::as_str) == Some("input_required") => @@ -80,8 +133,8 @@ pub(crate) async fn call_tool( if round + 1 == DEFAULT_MRTR_MAX_ROUNDS { break; } - let input: ToolInput = - serde_json::from_value(input).map_err(|_| invalid("invalid MCP tool input request"))?; + let input: RequiredInput = + serde_json::from_value(input).map_err(|_| invalid("invalid MCP input request"))?; let requests = input.input_requests.unwrap_or_default(); if requests.is_empty() && input.request_state.is_none() { return Err(ServiceError::UnexpectedResponse); @@ -99,35 +152,30 @@ pub(crate) async fn call_tool( } else { state_only_rounds = 0; } - let responses = futures::future::try_join_all(requests.into_iter().enumerate().map( - |(index, (key, mut request))| { - // Native prompts need distinct UI/cancellation ownership when - // concurrent tool calls use the same server-assigned input key. + let responses = + futures::future::try_join_all(requests.into_iter().map(|(key, mut request)| { + // Prompts outlive individual transport requests, including across + // session recovery, so their IDs must not reuse a peer's counter. let native_verification = matches!(&request, ServerRequest::CustomRequest(request) if request.params.as_ref().and_then(|params| params.get("mode")) .and_then(Value::as_str) == Some(crate::user_verification::MODE)); - let request_id = if native_verification { - RequestId::String(format!("tool-input/{id}/{index}").into()) - } else { - RequestId::String(key.clone().into()) - }; + let id = NEXT_INPUT_ID.fetch_add(1, Ordering::Relaxed); + let request_id = RequestId::String(format!("mcp-input/{id}").into()); let mut context = RequestContext::new(request_id, service.peer().clone()); context.meta = std::mem::take(request.get_meta_mut()); context.extensions = std::mem::take(request.extensions_mut()); async move { - let result = service - .service() + let result = input_service .handle_request(request, context) .await .map_err(ServiceError::McpError)?; let result = serde_json::to_value(result) - .map_err(|_| invalid("invalid MCP tool input response"))?; + .map_err(|_| invalid("invalid MCP input response"))?; let contains_proof = native_verification && result.get("action").and_then(Value::as_str) == Some("accept"); Ok::<_, ServiceError>((key, result, contains_proof)) } - }, - )); + })); let responses = tokio::select! { biased; _ = async { @@ -140,7 +188,7 @@ pub(crate) async fn call_tool( } => return Err(ServiceError::TransportClosed), responses = responses => responses?, }; - params.input_responses = (!responses.is_empty()).then(|| { + let input_responses = (!responses.is_empty()).then(|| { responses .into_iter() .map(|(key, result, contains_proof)| { @@ -149,7 +197,17 @@ pub(crate) async fn call_tool( }) .collect() }); - params.request_state = input.request_state; + match &mut request { + ClientRequest::CallToolRequest(request) => { + request.params.input_responses = input_responses; + request.params.request_state = input.request_state; + } + ClientRequest::ReadResourceRequest(request) => { + request.params.input_responses = input_responses; + request.params.request_state = input.request_state; + } + _ => return Err(ServiceError::UnexpectedResponse), + } } Err(ServiceError::InputRequiredRoundsExceeded { max_rounds: DEFAULT_MRTR_MAX_ROUNDS, @@ -168,7 +226,7 @@ fn parse_input(value: Value) -> Result { }; } let request: ServerRequest = - serde_json::from_value(value).map_err(|_| invalid("invalid MCP tool input request"))?; + serde_json::from_value(value).map_err(|_| invalid("invalid MCP input request"))?; if let ServerRequest::CustomRequest(request) = &request && request.method == "openai/elicitation/create" { @@ -190,7 +248,7 @@ fn parse_input(value: Value) -> Result { } return Ok(ServerRequest::CustomRequest(request.clone())); } - Err(invalid("unsupported MCP tool input request")) + Err(invalid("unsupported MCP input request")) } fn invalid(message: &'static str) -> ServiceError { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index d5cdf26aa1..ed3a144d5d 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -39,8 +40,11 @@ use rmcp::model::ElicitResult; use rmcp::model::ElicitationAction; use rmcp::model::Extensions; use rmcp::model::InitializeRequestParams; +use rmcp::model::ListResourceTemplatesRequest; use rmcp::model::ListResourceTemplatesResult; +use rmcp::model::ListResourcesRequest; use rmcp::model::ListResourcesResult; +use rmcp::model::ListToolsRequest; use rmcp::model::ListToolsResult; use rmcp::model::MetaObject; use rmcp::model::PaginatedRequestParams; @@ -97,6 +101,7 @@ use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::oauth_refresh_mode::McpOAuthRefreshMode; use crate::ordered_byte_transport::OrderedByteTransport; use crate::protocol_mode::McpProtocolMode; +use crate::request_cancellation::RequestCancellation; use crate::startup_error::is_authentication_required_error; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; @@ -414,9 +419,23 @@ pub struct RmcpClient { session_recovery_lock: Semaphore, elicitation_pause_state: ElicitationPauseState, tool_list_revision: Arc, + cancel_requests_on_drop: AtomicBool, } impl RmcpClient { + /// Extend tool-call cancellation to catalog, resource, and custom requests on a shared connection. + pub fn enable_request_cancellation(&self) { + self.cancel_requests_on_drop.store(true, Ordering::Relaxed); + } + + fn request_cancellation(&self) -> RequestCancellation { + if self.cancel_requests_on_drop.load(Ordering::Relaxed) { + RequestCancellation::OnDrop + } else { + RequestCancellation::Disabled + } + } + /// Tool catalog invalidation sequence, ordered before subsequent responses. pub fn tool_list_revision(&self) -> crate::ToolListRevision { crate::ToolListRevision(Arc::clone(&self.tool_list_revision)) @@ -460,6 +479,7 @@ impl RmcpClient { session_recovery_lock: Semaphore::new(/*permits*/ 1), elicitation_pause_state: ElicitationPauseState::new(), tool_list_revision: Arc::default(), + cancel_requests_on_drop: AtomicBool::new(false), }) } @@ -535,6 +555,7 @@ impl RmcpClient { session_recovery_lock: Semaphore::new(/*permits*/ 1), elicitation_pause_state: ElicitationPauseState::new(), tool_list_revision: Arc::default(), + cancel_requests_on_drop: AtomicBool::new(false), }) } @@ -643,6 +664,7 @@ impl RmcpClient { session_recovery_lock: Semaphore::new(/*permits*/ 1), elicitation_pause_state: ElicitationPauseState::new(), tool_list_revision: Arc::default(), + cancel_requests_on_drop: AtomicBool::new(false), }) } @@ -717,10 +739,31 @@ impl RmcpClient { if self.requires_read_only_tools { self.apply_read_only_tools_meta(&mut params.get_or_insert_default().meta); } + let cancellation = self.request_cancellation(); let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); - async move { service.list_tools(params).await }.boxed() + async move { + if matches!(cancellation, RequestCancellation::Disabled) { + return service.list_tools(params).await; + } + let result = cancellation + .send( + service.peer(), + ClientRequest::ListToolsRequest(ListToolsRequest { + method: Default::default(), + params, + extensions: Default::default(), + }), + rmcp::service::PeerRequestOptions::no_options(), + ) + .await?; + match result { + ServerResult::ListToolsResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } + } + .boxed() }) .await?; self.persist_oauth_tokens().await; @@ -774,10 +817,31 @@ impl RmcpClient { ) -> Result { let params = crate::trace_context::traced_pagination(params); self.refresh_oauth_if_needed().await?; + let cancellation = self.request_cancellation(); let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); - async move { service.list_resources(params).await }.boxed() + async move { + if matches!(cancellation, RequestCancellation::Disabled) { + return service.list_resources(params).await; + } + let result = cancellation + .send( + service.peer(), + ClientRequest::ListResourcesRequest(ListResourcesRequest { + method: Default::default(), + params, + extensions: Default::default(), + }), + rmcp::service::PeerRequestOptions::no_options(), + ) + .await?; + match result { + ServerResult::ListResourcesResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } + } + .boxed() }) .await?; self.persist_oauth_tokens().await; @@ -791,10 +855,33 @@ impl RmcpClient { ) -> Result { let params = crate::trace_context::traced_pagination(params); self.refresh_oauth_if_needed().await?; + let cancellation = self.request_cancellation(); let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); - async move { service.list_resource_templates(params).await }.boxed() + async move { + if matches!(cancellation, RequestCancellation::Disabled) { + return service.list_resource_templates(params).await; + } + let result = cancellation + .send( + service.peer(), + ClientRequest::ListResourceTemplatesRequest( + ListResourceTemplatesRequest { + method: Default::default(), + params, + extensions: Default::default(), + }, + ), + rmcp::service::PeerRequestOptions::no_options(), + ) + .await?; + match result { + ServerResult::ListResourceTemplatesResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } + } + .boxed() }) .await?; self.persist_oauth_tokens().await; @@ -809,11 +896,21 @@ impl RmcpClient { let mut params = params; params.meta = crate::trace_context::with_current_trace(params.meta); self.refresh_oauth_if_needed().await?; + let cancellation = self.request_cancellation(); let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); async move { + if matches!(cancellation, RequestCancellation::OnDrop) { + return crate::request_input::read_resource( + &service, + service.service(), + params, + cancellation, + ) + .await; + } let modern_session = requested_modern && service.peer().peer_info().is_some_and(|info| { info.protocol_version == ProtocolVersion::V_2026_07_28 @@ -831,6 +928,52 @@ impl RmcpClient { Ok(result) } + /// Reads a resource with inputs and timeout pauses owned by the calling chat. + pub async fn read_resource_with_elicitation( + &self, + mut params: ReadResourceRequestParams, + timeout: Option, + send_elicitation: Arc, + ) -> Result { + params.meta = crate::trace_context::with_current_trace(params.meta); + self.refresh_oauth_if_needed().await?; + let cancellation = self.request_cancellation(); + let pause_state = ElicitationPauseState::new(); + let result = self + .run_service_operation_with_pause_state( + "resources/read", + timeout, + pause_state.clone(), + move |service| { + let params = params.clone(); + let input_service = service + .service() + .for_request(Arc::clone(&send_elicitation), pause_state.clone()); + async move { + if !service.peer().peer_info().is_some_and(|info| { + info.protocol_version == ProtocolVersion::V_2026_07_28 + }) { + return Err(ServiceError::McpError(rmcp::ErrorData::invalid_request( + "request-scoped MCP inputs require protocol 2026-07-28", + /*data*/ None, + ))); + } + crate::request_input::read_resource( + &service, + &input_service, + params, + cancellation, + ) + .await + } + .boxed() + }, + ) + .await?; + self.persist_oauth_tokens().await; + Ok(result) + } + /// Builds metadata against the session that sends each attempt, including a recovery retry. /// The resolver must be safe to invoke again and must retain the call's captured authority. pub async fn call_tool( @@ -840,6 +983,50 @@ impl RmcpClient { meta_for_session: F, timeout: Option, ) -> Result + where + F: Fn(&rmcp::model::ServerCapabilities) -> Result> + Send + Sync, + { + self.call_tool_inner( + name, + arguments, + meta_for_session, + timeout, + /*send_elicitation*/ None, + ) + .await + } + + /// Routes modern tool inputs and timeout pauses to this caller on a shared connection. + /// Legacy connections are rejected because their callbacks belong to connection startup. + pub async fn call_tool_with_elicitation( + &self, + name: String, + arguments: Option, + meta_for_session: F, + timeout: Option, + send_elicitation: Arc, + ) -> Result + where + F: Fn(&rmcp::model::ServerCapabilities) -> Result> + Send + Sync, + { + self.call_tool_inner( + name, + arguments, + meta_for_session, + timeout, + Some(send_elicitation), + ) + .await + } + + async fn call_tool_inner( + &self, + name: String, + arguments: Option, + meta_for_session: F, + timeout: Option, + send_elicitation: Option>, + ) -> Result where F: Fn(&rmcp::model::ServerCapabilities) -> Result> + Send + Sync, { @@ -877,55 +1064,85 @@ impl RmcpClient { rmcp_params.arguments = arguments; let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; let meta_for_session = &meta_for_session; + let pause_state = if send_elicitation.is_some() { + ElicitationPauseState::new() + } else { + self.elicitation_pause_state.clone() + }; match self - .run_service_operation("tools/call", timeout, move |service| { - let mut rmcp_params = rmcp_params.clone(); - async move { - // Resolve from this service, not a cached startup capability or another - // service observed before recovery. Local policy failures never send a call. - let meta = (|| { - let info = service - .peer() - .peer_info() - .context("MCP session has no negotiated server capabilities")?; - let mut meta = match meta_for_session(&info.capabilities)? { - Some(Value::Object(map)) => Some(RequestMetaObject::from(map)), - Some(other) => { - return Err(anyhow!( - "MCP tool request _meta must be a JSON object, got {other}" - )); - } - None => None, + .run_service_operation_with_pause_state( + "tools/call", + timeout, + pause_state.clone(), + move |service| { + let mut rmcp_params = rmcp_params.clone(); + let send_elicitation = send_elicitation.clone(); + let pause_state = pause_state.clone(); + async move { + // Resolve from this service, not a cached startup capability or another + // service observed before recovery. Local policy failures never send a call. + let meta = (|| { + let info = service + .peer() + .peer_info() + .context("MCP session has no negotiated server capabilities")?; + let mut meta = match meta_for_session(&info.capabilities)? { + Some(Value::Object(map)) => Some(RequestMetaObject::from(map)), + Some(other) => { + return Err(anyhow!( + "MCP tool request _meta must be a JSON object, got {other}" + )); + } + None => None, + }; + self.apply_read_only_tools_meta(&mut meta); + Ok(crate::trace_context::with_current_trace(meta)) + })(); + let meta = match meta { + Ok(meta) => meta, + Err(error) => return Ok(Err(error)), }; - self.apply_read_only_tools_meta(&mut meta); - Ok(crate::trace_context::with_current_trace(meta)) - })(); - let meta = match meta { - Ok(meta) => meta, - Err(error) => return Ok(Err(error)), - }; - let modern_session = requested_modern - && service.peer().peer_info().is_some_and(|info| { - info.protocol_version == ProtocolVersion::V_2026_07_28 - }); - if modern_session { - rmcp_params.meta = meta; - return crate::tool_input::call_tool(&service, rmcp_params) + let modern_session = requested_modern + && service.peer().peer_info().is_some_and(|info| { + info.protocol_version == ProtocolVersion::V_2026_07_28 + }); + if modern_session { + rmcp_params.meta = meta; + let input_service = send_elicitation + .map(|sender| service.service().for_request(sender, pause_state)); + return crate::request_input::call_tool( + &service, + input_service.as_ref().unwrap_or_else(|| service.service()), + rmcp_params, + ) .await .map(Ok); - } - let mut options = rmcp::service::PeerRequestOptions::no_options(); - options.meta = meta; - let (_, result) = - crate::tool_request::call_tool(service.peer(), rmcp_params, options) + } + if send_elicitation.is_some() { + return Err(ServiceError::McpError(rmcp::ErrorData::invalid_request( + "request-scoped MCP inputs require protocol 2026-07-28", + /*data*/ None, + ))); + } + let mut options = rmcp::service::PeerRequestOptions::no_options(); + options.meta = meta; + let result = RequestCancellation::OnDrop + .send( + service.peer(), + ClientRequest::CallToolRequest(rmcp::model::CallToolRequest::new( + rmcp_params, + )), + options, + ) .await?; - match result { - ServerResult::CallToolResult(result) => Ok(Ok(result)), - _ => Err(ServiceError::UnexpectedResponse), + match result { + ServerResult::CallToolResult(result) => Ok(Ok(result)), + _ => Err(rmcp::service::ServiceError::UnexpectedResponse), + } } - } - .boxed() - }) + .boxed() + }, + ) .await { Ok(Ok(result)) => { @@ -1005,13 +1222,18 @@ impl RmcpClient { timeout: Option, ) -> Result { self.refresh_oauth_if_needed().await?; + let cancellation = self.request_cancellation(); let response = self .run_service_operation("requests/custom", timeout, move |service| { let params = params.clone(); async move { let request = crate::trace_context::traced_custom_request(method, params); - service - .send_request(ClientRequest::CustomRequest(request)) + cancellation + .send( + service.peer(), + ClientRequest::CustomRequest(request), + rmcp::service::PeerRequestOptions::no_options(), + ) .await } .boxed() @@ -1444,6 +1666,26 @@ impl RmcpClient { timeout: Option, operation: F, ) -> Result + where + F: Fn(Arc>) -> Fut, + Fut: std::future::Future>, + { + self.run_service_operation_with_pause_state( + label, + timeout, + self.elicitation_pause_state.clone(), + operation, + ) + .await + } + + async fn run_service_operation_with_pause_state( + &self, + label: &str, + timeout: Option, + pause_state: ElicitationPauseState, + operation: F, + ) -> Result where F: Fn(Arc>) -> Fut, Fut: std::future::Future>, @@ -1453,7 +1695,7 @@ impl RmcpClient { Arc::clone(&service), label, timeout, - self.elicitation_pause_state.clone(), + pause_state.clone(), &operation, ) .await @@ -1466,7 +1708,7 @@ impl RmcpClient { recovered_service, label, timeout, - self.elicitation_pause_state.clone(), + pause_state, &operation, ) .await diff --git a/codex-rs/rmcp-client/src/tool_input_tests.rs b/codex-rs/rmcp-client/src/tool_input_tests.rs index 306de871e0..5d45b46f3f 100644 --- a/codex-rs/rmcp-client/src/tool_input_tests.rs +++ b/codex-rs/rmcp-client/src/tool_input_tests.rs @@ -1,8 +1,8 @@ //! Exercises pending MRTR input cleanup over a real duplex transport. use super::ElicitationClientService; +use crate::request_input::call_tool; use crate::rmcp_client::ElicitationPauseState; -use crate::tool_input::call_tool; use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; use rmcp::RoleServer; use rmcp::model::CallToolRequestParams; @@ -53,7 +53,11 @@ async fn connection_closure_releases_pending_inputs_and_timeout_pause() -> anyho tokio::io::duplex(/*max_buf_size*/ 4096); let client = serve_directly(service, client_transport, /*peer_info*/ None); let mut server = IntoTransport::::into_transport(server_transport); - let call = call_tool(&client, CallToolRequestParams::new("test")); + let call = call_tool( + &client, + client.service(), + CallToolRequestParams::new("test"), + ); tokio::pin!(call); let request = tokio::select! { result = &mut call => panic!("tool completed before input: {result:?}"), diff --git a/codex-rs/rmcp-client/src/tool_request.rs b/codex-rs/rmcp-client/src/tool_request.rs deleted file mode 100644 index 70be1846fb..0000000000 --- a/codex-rs/rmcp-client/src/tool_request.rs +++ /dev/null @@ -1,49 +0,0 @@ -//! Cancel remote tool work when its local response future is dropped. - -use rmcp::Peer; -use rmcp::RoleClient; -use rmcp::model::CallToolRequest; -use rmcp::model::CallToolRequestParams; -use rmcp::model::CancelledNotificationParam; -use rmcp::model::ClientRequest; -use rmcp::model::RequestId; -use rmcp::model::ServerResult; -use rmcp::service::PeerRequestOptions; -use rmcp::service::ServiceError; - -pub(crate) async fn call_tool( - peer: &Peer, - params: CallToolRequestParams, - options: PeerRequestOptions, -) -> Result<(RequestId, ServerResult), ServiceError> { - let handle = peer - .send_request_with_option( - ClientRequest::CallToolRequest(CallToolRequest::new(params)), - options, - ) - .await?; - let id = handle.id.clone(); - let mut cancellation = CancelOnDrop(Some((peer.clone(), id.clone()))); - let result = handle.await_response().await; - cancellation.0 = None; - result.map(|result| (id, result)) -} - -struct CancelOnDrop(Option<(Peer, RequestId)>); - -impl Drop for CancelOnDrop { - fn drop(&mut self) { - if let Some((peer, id)) = self.0.take() - && let Ok(runtime) = tokio::runtime::Handle::try_current() - { - runtime.spawn(async move { - let _ = peer - .notify_cancelled(CancelledNotificationParam::new( - Some(id), - Some("tool call response future dropped".into()), - )) - .await; - }); - } - } -} diff --git a/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs b/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs index 30a21bc930..aab81ce435 100644 --- a/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs +++ b/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs @@ -5,7 +5,10 @@ use super::ElicitationPauseState; use super::ElicitationResponse; use super::RmcpClient; use crate::InProcessTransportFactory; +use crate::McpProtocolMode; use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; +use futures::FutureExt; +use futures::TryFutureExt; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use rmcp::RoleServer; @@ -19,9 +22,11 @@ use rmcp::model::ClientJsonRpcMessage; use rmcp::model::ClientNotification; use rmcp::model::ClientRequest; use rmcp::model::CustomRequest; +use rmcp::model::DiscoverResult; use rmcp::model::ElicitationAction; use rmcp::model::InitializeResult; use rmcp::model::ProtocolVersion; +use rmcp::model::ReadResourceRequestParams; use rmcp::model::RequestId; use rmcp::model::ServerJsonRpcMessage; use rmcp::model::ServerNotification; @@ -149,6 +154,78 @@ async fn dropping_tool_call_cancels_remote_request() -> anyhow::Result<()> { .await? } +#[tokio::test] +async fn dropping_shared_resource_and_custom_requests_cancels_the_pending_request() +-> anyhow::Result<()> { + let (servers, mut server_rx) = mpsc::unbounded_channel(); + let mut client = + RmcpClient::new_in_process_client(Arc::new(ElicitationTestTransport { servers })).await?; + client.protocol_mode = McpProtocolMode::V20260728; + let mut server = + IntoTransport::::into_transport(server_rx.recv().await.unwrap()); + let (initialized, discovery) = tokio::join!( + client.initialize( + ClientConfig::default(), + Some(Duration::from_secs(5)), + Box::new(|_, _| panic!("unexpected startup elicitation")), + ), + async { + let Some(ClientJsonRpcMessage::Request(request)) = server.receive().await else { + anyhow::bail!("expected discovery"); + }; + assert!(matches!(request.request, ClientRequest::DiscoverRequest(_))); + let result = DiscoverResult::from_server_info( + vec![ProtocolVersion::V_2026_07_28], + InitializeResult::default(), + ); + server + .send(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(result), + request.id, + )) + .await?; + anyhow::Ok(()) + }, + ); + discovery?; + initialized?; + client.enable_request_cancellation(); + for mut pending in [ + client + .read_resource_with_elicitation( + ReadResourceRequestParams::new("test://pending"), + /*timeout*/ None, + Arc::new(Box::new(|_, _| panic!("unexpected input"))), + ) + .map_ok(|_| ()) + .boxed(), + client + .send_custom_request("test/pending", /*params*/ None) + .map_ok(|_| ()) + .boxed(), + ] { + let request = tokio::select! { + result = &mut pending => panic!("request completed before a response: {result:?}"), + message = timeout(Duration::from_secs(5), server.receive()) => message?, + }; + let Some(ClientJsonRpcMessage::Request(request)) = request else { + anyhow::bail!("expected pending request"); + }; + drop(pending); + let Some(ClientJsonRpcMessage::Notification(notification)) = + timeout(Duration::from_secs(5), server.receive()).await? + else { + anyhow::bail!("expected cancellation"); + }; + let ClientNotification::CancelledNotification(cancelled) = notification.notification else { + anyhow::bail!("expected cancellation"); + }; + assert_eq!(cancelled.params.request_id, Some(request.id)); + } + client.shutdown().await; + Ok(()) +} + #[traced_test] #[tokio::test] async fn recovered_connections_accept_elicitations_with_previously_cancelled_ids() diff --git a/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs index 5890e3f753..91bf1bca0e 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; use std::time::Duration; @@ -10,6 +11,7 @@ use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::McpProtocolMode; use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::SendElicitation; use futures::FutureExt; use pretty_assertions::assert_eq; use rmcp::model::ClientCapabilities; @@ -133,10 +135,9 @@ async fn create_client( ) .with_protocol_version(ProtocolVersion::V_2025_06_18), Some(Duration::from_secs(5)), - Box::new(move |request_id, request| { + Box::new(move |_, request| { let elicitation_modes = Arc::clone(&elicitation_modes); async move { - assert_eq!(request_id.to_string(), "confirmation"); let Elicitation::Mcp(request) = request else { anyhow::bail!("MRTR input must be a standard MCP elicitation"); }; @@ -219,12 +220,10 @@ async fn modern_sse_input_required_preserves_response_metadata() -> anyhow::Resu } #[tokio::test] -async fn modern_tool_mrtr_drives_form_and_url_elicitation_and_preserves_metadata() +async fn modern_tool_mrtr_routes_concurrent_inputs_to_their_callers_and_preserves_metadata() -> anyhow::Result<()> { for mode in ["form", "url"] { let server = MockServer::start().await; - let calls = Arc::new(Mutex::new(Vec::::new())); - let recorded = Arc::clone(&calls); Mock::given(method("POST")) .and(path("/mcp")) @@ -233,11 +232,7 @@ async fn modern_tool_mrtr_drives_form_and_url_elicitation_and_preserves_metadata match body["method"].as_str() { Some("server/discover") => discover_response(&body), Some("tools/call") => { - recorded.lock().expect("requests lock").push(body.clone()); - assert_eq!( - body.pointer("/params/_meta/requestContext"), - Some(&json!("caller-context")) - ); + let caller = body["params"]["_meta"]["requestContext"].clone(); assert_eq!( body.pointer("/params/_meta/io.modelcontextprotocol~1protocolVersion"), Some(&json!(MODERN_VERSION)) @@ -269,7 +264,7 @@ async fn modern_tool_mrtr_drives_form_and_url_elicitation_and_preserves_metadata body.pointer("/params/inputResponses/confirmation"), Some(&json!({ "action": "accept", - "content": {"confirmed": true}, + "content": {"confirmed": true, "caller": caller}, "_meta": {"clientContext": "preserved"}, })) ); @@ -285,35 +280,65 @@ async fn modern_tool_mrtr_drives_form_and_url_elicitation_and_preserves_metadata other => panic!("unexpected MRTR request: {other:?}"), } }) - .expect(3) + .expect(8) .mount(&server) .await; - let elicitation_modes = Arc::new(Mutex::new(Vec::new())); - let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; - let result = client - .call_tool( + let client = create_client(&server, Arc::default()).await?; + let other_client = create_client(&server, Arc::default()).await?; + let inputs = Arc::new(Mutex::new(Vec::new())); + let callers = [ + (&client, "first"), + (&client, "second"), + (&other_client, "other connection"), + ]; + let results = futures::future::try_join_all(callers.map(|(client, caller)| { + let inputs = Arc::clone(&inputs); + let sender: SendElicitation = Box::new(move |id, request| { + inputs.lock().unwrap().push((caller, id)); + let Elicitation::Mcp(request) = request else { + panic!("expected standard MCP input"); + }; + let meta = match request { + ElicitRequestParams::FormElicitationParams { meta, .. } + | ElicitRequestParams::UrlElicitationParams { meta, .. } => meta, + _ => panic!("expected form or URL input"), + }; + assert_eq!( + meta.and_then(|meta| meta.get("inputContext").cloned()), + Some(json!("server-context")) + ); + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"confirmed": true, "caller": caller})), + meta: Some(json!({"clientContext": "preserved"})), + })) + .boxed() + }); + client.call_tool_with_elicitation( "confirm".into(), Some(json!({"mode": mode})), - |_| Ok(Some(json!({"requestContext": "caller-context"}))), + move |_| Ok(Some(json!({"requestContext": caller}))), Some(Duration::from_secs(5)), + Arc::new(sender), ) - .await?; - assert_eq!( - result.content[0].as_text().map(|text| text.text.as_str()), - Some("MRTR completed") - ); - assert_eq!( - *elicitation_modes.lock().expect("elicitation lock"), - vec![mode] - ); + })) + .await?; + for result in results { + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some("MRTR completed") + ); + } { - let calls = calls.lock().expect("requests lock"); - assert_eq!(calls.len(), 2); - assert_ne!(calls[0]["id"], calls[1]["id"]); + let inputs = inputs.lock().unwrap(); + assert_eq!(inputs.len(), 3); + let ids = inputs.iter().map(|(_, id)| id).collect::>(); + assert_eq!(ids.len(), inputs.len()); } client.shutdown().await; + other_client.shutdown().await; } Ok(()) } @@ -548,11 +573,26 @@ async fn modern_resource_mrtr_retries_with_opaque_request_state() -> anyhow::Res .await; let elicitation_modes = Arc::new(Mutex::new(Vec::new())); - let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let client = create_client(&server, Arc::default()).await?; + let observed_modes = Arc::clone(&elicitation_modes); + let sender: SendElicitation = Box::new(move |_, request| { + assert!(matches!( + request, + Elicitation::Mcp(ElicitRequestParams::FormElicitationParams { .. }) + )); + observed_modes.lock().unwrap().push("form"); + futures::future::ready(Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"confirmed": true})), + meta: Some(json!({"clientContext": "preserved"})), + })) + .boxed() + }); let result = client - .read_resource( + .read_resource_with_elicitation( ReadResourceRequestParams::new("memo://requires-confirmation"), Some(Duration::from_secs(5)), + Arc::new(sender), ) .await?; assert_eq!(result.contents.len(), 1); diff --git a/codex-rs/rmcp-client/tests/mcp_2026_mrtr/native_verification_tests.rs b/codex-rs/rmcp-client/tests/mcp_2026_mrtr/native_verification_tests.rs index 705867acef..de78ad7b5c 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_mrtr/native_verification_tests.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_mrtr/native_verification_tests.rs @@ -465,7 +465,9 @@ async fn native_mrtr_does_not_restart_after_continuation_session_expiry() -> any .await?; let error = call(&client).await.unwrap_err(); assert!( - error.to_string().contains("the tool was not restarted"), + error + .to_string() + .contains("the operation was not restarted"), "{error}" ); client.shutdown().await; @@ -512,7 +514,9 @@ async fn native_mrtr_does_not_restart_after_proof_then_state_only_round() -> any .await?; let error = call(&client).await.unwrap_err(); assert!( - error.to_string().contains("the tool was not restarted"), + error + .to_string() + .contains("the operation was not restarted"), "{error}" ); client.shutdown().await; @@ -521,7 +525,7 @@ async fn native_mrtr_does_not_restart_after_proof_then_state_only_round() -> any } #[tokio::test] -async fn native_mrtr_concurrent_prompts_have_independent_cancellation() -> anyhow::Result<()> { +async fn native_mrtr_shared_calls_isolate_inputs_cancellation_and_timeouts() -> anyhow::Result<()> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/mcp")) @@ -529,6 +533,10 @@ async fn native_mrtr_concurrent_prompts_have_independent_cancellation() -> anyho let body: Value = request.body_json().unwrap(); match body["method"].as_str() { Some("server/discover") => discover_response(&body), + Some("tools/call") if body["params"]["name"] == "slow" => { + result_response(&body, json!({"resultType": "complete", "content": []})) + .set_delay(Duration::from_secs(1)) + } Some("tools/call") if body.pointer("/params/inputResponses").is_none() => { result_response(&body, input_required(native_request())) } @@ -542,31 +550,62 @@ async fn native_mrtr_concurrent_prompts_have_independent_cancellation() -> anyho other => panic!("unexpected request: {other:?}"), } }) - .expect(/*r*/ 4) + .expect(/*r*/ 5) .mount(&server) .await; let (tx, mut rx) = mpsc::unbounded_channel(); let client = client( &server, capabilities(), - Box::new(move |id, _| { - let (reply_tx, reply_rx) = oneshot::channel(); - tx.send((id, reply_tx)).unwrap(); - Box::pin(async move { Ok(reply_rx.await?) }) - }), + Box::new(|_, _| panic!("shared input must use the calling chat's callback")), ) .await?; - let first_client = Arc::clone(&client); - let first_call = tokio::spawn(async move { call(&first_client).await }); - let (first_id, mut first_reply) = timeout(Duration::from_secs(/*secs*/ 5), rx.recv()) - .await? - .unwrap(); - let second_client = Arc::clone(&client); - let second_call = tokio::spawn(async move { call(&second_client).await }); - let (second_id, second_reply) = timeout(Duration::from_secs(/*secs*/ 5), rx.recv()) - .await? - .unwrap(); + let launch = |caller| { + let client = Arc::clone(&client); + let tx = tx.clone(); + tokio::spawn(async move { + let sender: SendElicitation = Box::new(move |id, _| { + let (reply_tx, reply_rx) = oneshot::channel(); + tx.send((caller, id, reply_tx)).unwrap(); + Box::pin(async move { Ok(reply_rx.await?) }) + }); + client + .call_tool_with_elicitation( + "verified_operation".into(), + /*arguments*/ None, + |_| Ok(None), + Some(Duration::from_secs(5)), + Arc::new(sender), + ) + .await + }) + }; + let first_call = launch("first"); + let (first_caller, first_id, mut first_reply) = + timeout(Duration::from_secs(/*secs*/ 5), rx.recv()) + .await? + .unwrap(); + let second_call = launch("second"); + let (second_caller, second_id, second_reply) = + timeout(Duration::from_secs(/*secs*/ 5), rx.recv()) + .await? + .unwrap(); assert_ne!(first_id, second_id); + assert_eq!((first_caller, second_caller), ("first", "second")); + let timeout_error = client + .call_tool_with_elicitation( + "slow".into(), + /*arguments*/ None, + |_| Ok(None), + Some(Duration::from_millis(100)), + Arc::new(Box::new(|_, _| panic!("slow call must not ask for input"))), + ) + .await + .unwrap_err(); + assert!( + timeout_error.to_string().contains("timed out"), + "{timeout_error}" + ); first_call.abort(); assert!(first_call.await.unwrap_err().is_cancelled()); timeout(Duration::from_secs(/*secs*/ 5), first_reply.closed()).await?; @@ -575,7 +614,7 @@ async fn native_mrtr_concurrent_prompts_have_independent_cancellation() -> anyho .send(response(ElicitationAction::Cancel, Value::Null)) .unwrap(); assert_eq!( - timeout(Duration::from_secs(/*secs*/ 5), second_call).await???, + serde_json::to_value(timeout(Duration::from_secs(/*secs*/ 5), second_call).await???)?, json!({"resultType": "complete", "content": []}) ); client.shutdown().await;