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;