diff --git a/codex-rs/rmcp-client/src/elicitation_client_service.rs b/codex-rs/rmcp-client/src/elicitation_client_service.rs index 70ac92582b..497ca0652d 100644 --- a/codex-rs/rmcp-client/src/elicitation_client_service.rs +++ b/codex-rs/rmcp-client/src/elicitation_client_service.rs @@ -268,24 +268,8 @@ impl Service for ElicitationClientService { if request.method == OPENAI_ELICITATION_METHOD && self.supports_openai_elicitation_form => { - let params = request - .params_as::() - .map_err(|err| { - rmcp::ErrorData::invalid_params(err.to_string(), /*data*/ None) - })? - .ok_or_else(|| { - rmcp::ErrorData::invalid_params("missing params", /*data*/ None) - })?; - let OpenAiElicitationRequestParams::Form(params) = params; let response = self - .create_elicitation( - Elicitation::OpenAiElicitationForm { - meta: params.meta, - message: params.message, - requested_schema: params.requested_schema, - }, - context, - ) + .create_elicitation(openai_elicitation_form(request)?, context) .await?; Ok(ClientResult::CustomResult(elicitation_response_result( response, @@ -371,6 +355,21 @@ fn openai_form_elicitation(request: CustomRequest) -> Result Result { + let params = request + .params_as::() + .map_err(|err| rmcp::ErrorData::invalid_params(err.to_string(), /*data*/ None))? + .ok_or_else(|| rmcp::ErrorData::invalid_params("missing params", /*data*/ None))?; + let OpenAiElicitationRequestParams::Form(params) = params; + Ok(Elicitation::OpenAiElicitationForm { + meta: params.meta, + message: params.message, + requested_schema: params.requested_schema, + }) +} + fn restore_context_meta( mut request: Elicitation, mut context_meta: RequestMetaObject, diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 56b278a580..0fd91d3867 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -27,6 +27,7 @@ mod rmcp_client; mod service_error; mod startup_error; mod stdio_server_launcher; +mod tool_input; mod user_verification; mod utils; mod www_authenticate; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 549e6d9d3f..004f55d0cf 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -846,7 +846,7 @@ impl RmcpClient { }); if modern_session { rmcp_params.meta = meta; - return service.call_tool(rmcp_params).await; + return crate::tool_input::call_tool(&service, rmcp_params).await; } let mut options = rmcp::service::PeerRequestOptions::no_options(); options.meta = meta; @@ -1662,6 +1662,10 @@ async fn create_oauth_transport_and_runtime( }) } +#[cfg(test)] +#[path = "tool_input_tests.rs"] +mod tool_input_tests; + #[cfg(test)] #[path = "user_verification_cancellation_tests.rs"] mod user_verification_cancellation_tests; diff --git a/codex-rs/rmcp-client/src/tool_input.rs b/codex-rs/rmcp-client/src/tool_input.rs new file mode 100644 index 0000000000..ffda73988d --- /dev/null +++ b/codex-rs/rmcp-client/src/tool_input.rs @@ -0,0 +1,208 @@ +//! Drives modern tool continuations, including the native verification extension. +//! Inputs use the existing 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::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::RequestId; +use rmcp::model::ServerRequest; +use rmcp::model::ServerResult; +use rmcp::service::PeerRequestOptions; +use rmcp::service::RequestContext; +use rmcp::service::RunningService; +use rmcp::service::Service; +use rmcp::service::ServiceError; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use serde::Deserialize; +use serde_json::Value; + +use crate::elicitation_client_service::ElicitationClientService; +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct ToolInput { + input_requests: Option>, + request_state: Option, +} + +pub(crate) async fn call_tool( + service: &RunningService, + mut params: CallToolRequestParams, +) -> Result { + let mut state_only_rounds = 0; + let mut submitted_verification_proof = false; + for round in 0..DEFAULT_MRTR_MAX_ROUNDS { + let response = async { + let handle = service + .peer() + .send_request_with_option( + ClientRequest::CallToolRequest(CallToolRequest::new(params.clone())), + PeerRequestOptions::no_options(), + ) + .await?; + let id = handle.id.clone(); + Ok::<_, ServiceError>((id, handle.await_response().await?)) + } + .await; + let (id, result) = response.map_err(|error| { + if submitted_verification_proof + && let ServiceError::TransportSend(transport) = &error + && !matches!( + transport + .error + .downcast_ref::>(), + Some(StreamableHttpError::AuthRequired(_)) + ) + { + // 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") + } 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"))?, + // RMCP's InputRequest union intentionally excludes custom methods. + ServerResult::CustomResult(result) + if result.0.get("resultType").and_then(Value::as_str) == Some("input_required") => + { + result.0 + } + _ => return Err(ServiceError::UnexpectedResponse), + }; + 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 requests = input.input_requests.unwrap_or_default(); + if requests.is_empty() && input.request_state.is_none() { + return Err(ServiceError::UnexpectedResponse); + } + // Parse the whole round before presenting any prompts. Only standard MCP + // inputs and supported OpenAI elicitation modes are accepted. + let requests = requests + .into_iter() + .map(|(key, value)| parse_input(value).map(|request| (key, request))) + .collect::, _>>()?; + if requests.is_empty() { + let millis = (50_u64 << state_only_rounds.min(/*other*/ 3)).min(/*other*/ 250); + tokio::time::sleep(Duration::from_millis(millis)).await; + state_only_rounds += 1; + } 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 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 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() + .handle_request(request, context) + .await + .map_err(ServiceError::McpError)?; + let result = serde_json::to_value(result) + .map_err(|_| invalid("invalid MCP tool 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 { + // RMCP exposes closure status but no awaitable closure notification. + // Include explicit cancellation and transport EOF; dropping the joined + // handlers releases both their UI callbacks and timeout pause guards. + while !service.is_closed() && !service.peer().is_transport_closed() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 50)).await; + } + } => return Err(ServiceError::TransportClosed), + responses = responses => responses?, + }; + params.input_responses = (!responses.is_empty()).then(|| { + responses + .into_iter() + .map(|(key, result, contains_proof)| { + submitted_verification_proof |= contains_proof; + (key, result) + }) + .collect() + }); + params.request_state = input.request_state; + } + Err(ServiceError::InputRequiredRoundsExceeded { + max_rounds: DEFAULT_MRTR_MAX_ROUNDS, + }) +} + +fn parse_input(value: Value) -> Result { + if let Ok(request) = serde_json::from_value::(value.clone()) { + return match request { + InputRequest::CreateMessage(request) => { + Ok(ServerRequest::CreateMessageRequest(request)) + } + InputRequest::Elicitation(request) => Ok(ServerRequest::ElicitRequest(request)), + InputRequest::ListRoots(request) => Ok(ServerRequest::ListRootsRequest(request)), + _ => Err(ServiceError::UnexpectedResponse), + }; + } + let request: ServerRequest = + serde_json::from_value(value).map_err(|_| invalid("invalid MCP tool input request"))?; + if let ServerRequest::CustomRequest(request) = &request + && request.method == "openai/elicitation/create" + { + match request + .params + .as_ref() + .and_then(|params| params.get("mode")) + .and_then(Value::as_str) + { + Some("form") => { + crate::elicitation_client_service::openai_elicitation_form(request.clone()) + .map_err(ServiceError::McpError)?; + } + Some(crate::user_verification::MODE) => { + crate::user_verification::parse_request(request.clone()) + .map_err(ServiceError::McpError)?; + } + _ => return Err(invalid("unsupported OpenAI elicitation mode")), + } + return Ok(ServerRequest::CustomRequest(request.clone())); + } + Err(invalid("unsupported MCP tool input request")) +} + +fn invalid(message: &'static str) -> ServiceError { + ServiceError::McpError(rmcp::ErrorData::invalid_request( + message, /*data*/ None, + )) +} diff --git a/codex-rs/rmcp-client/src/tool_input_tests.rs b/codex-rs/rmcp-client/src/tool_input_tests.rs new file mode 100644 index 0000000000..ebbdc962f6 --- /dev/null +++ b/codex-rs/rmcp-client/src/tool_input_tests.rs @@ -0,0 +1,101 @@ +//! Exercises pending MRTR input cleanup over a real duplex transport. + +use super::ElicitationClientService; +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; +use rmcp::model::ClientInfo; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::CustomResult; +use rmcp::model::ServerJsonRpcMessage; +use rmcp::model::ServerResult; +use rmcp::service::ServiceError; +use rmcp::service::serve_directly; +use rmcp::transport::IntoTransport; +use rmcp::transport::Transport; +use serde_json::json; +use std::time::Duration; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::time::timeout; + +#[tokio::test] +async fn connection_closure_releases_pending_inputs_and_timeout_pause() -> anyhow::Result<()> { + for mode in ["form", "openai/userVerification"] { + for close_transport in [false, true] { + let pause_state = ElicitationPauseState::new(); + let mut paused = pause_state.subscribe(); + let (route_tx, mut route_rx) = mpsc::unbounded_channel(); + let mut info = ClientInfo::default(); + info.capabilities.extensions = Some( + [( + OPENAI_ELICITATION_EXTENSION_ID.into(), + serde_json::Map::from_iter([ + ("userVerification".into(), json!({})), + ("form".into(), json!({})), + ]), + )] + .into_iter() + .collect(), + ); + let service = ElicitationClientService::new( + info, + Box::new(move |_, _| { + let (tx, rx) = oneshot::channel(); + route_tx.send(tx).unwrap(); + Box::pin(async move { Ok(rx.await?) }) + }), + pause_state, + ); + let (client_transport, server_transport) = + 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")); + tokio::pin!(call); + let request = tokio::select! { + result = &mut call => panic!("tool completed before input: {result:?}"), + request = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) => request?, + }; + let Some(ClientJsonRpcMessage::Request(request)) = request else { + anyhow::bail!("expected tool call"); + }; + let params = match mode { + "form" => json!({ + "mode": mode, "message": "Confirm", + "requestedSchema": {"type": "object", "properties": {}}, + }), + _ => json!({ + "mode": mode, "title": "Verify", "description": "", "challenge": "AQID", + }), + }; + server.send(ServerJsonRpcMessage::response( + ServerResult::CustomResult(CustomResult(json!({ + "resultType": "input_required", + "inputRequests": {"input": {"method": "openai/elicitation/create", "params": params}}, + }))), + request.id, + )).await?; + let reply = tokio::select! { + result = &mut call => panic!("tool completed before prompting: {result:?}"), + reply = timeout(Duration::from_secs(/*secs*/ 5), route_rx.recv()) => reply?.unwrap(), + }; + assert!(*paused.borrow()); + if close_transport { + drop(server); + } else { + client.cancellation_token().cancel(); + } + let result = timeout(Duration::from_secs(/*secs*/ 5), call).await?; + assert!( + matches!(result, Err(ServiceError::TransportClosed)), + "{result:?}" + ); + assert!(reply.is_closed()); + assert!(!*paused.borrow_and_update()); + } + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs index 6f6a8555f9..5d86e92219 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs @@ -34,6 +34,9 @@ use wiremock::matchers::path; const MODERN_VERSION: &str = "2026-07-28"; const OPAQUE_STATE: &str = " opaque/\u{2603}/=?base64?literal?=\n"; +#[path = "mcp_2026_mrtr/native_verification_tests.rs"] +mod native_verification; + fn discover_response(body: &Value) -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(json!({ "jsonrpc": "2.0", 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 new file mode 100644 index 0000000000..8f26c64986 --- /dev/null +++ b/codex-rs/rmcp-client/tests/mcp_2026_mrtr/native_verification_tests.rs @@ -0,0 +1,578 @@ +//! Exercises native verification through the HTTP MRTR tool-input envelope. + +use super::*; +use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; +use codex_rmcp_client::SendElicitation; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::time::timeout; + +fn native_request() -> Value { + json!({ + "method": "openai/elicitation/create", + "params": { + "mode": "openai/userVerification", "title": "Verify action", + "description": "Verify the requested operation", "challenge": "AQID" + } + }) +} + +fn rich_form_request() -> Value { + let mut request = elicitation_request("form"); + request["method"] = json!("openai/elicitation/create"); + request +} + +fn input_required(request: Value) -> Value { + json!({ + "resultType": "input_required", "content": [], "isError": false, + "inputRequests": {"verification": request}, "requestState": OPAQUE_STATE + }) +} + +fn capabilities() -> ClientCapabilities { + let mut capabilities = ClientCapabilities::default(); + capabilities.extensions = Some( + [( + OPENAI_ELICITATION_EXTENSION_ID.into(), + serde_json::Map::from_iter([("userVerification".into(), json!({}))]), + )] + .into_iter() + .collect(), + ); + capabilities +} + +async fn client( + server: &MockServer, + capabilities: ClientCapabilities, + handler: SendElicitation, +) -> anyhow::Result> { + let client = RmcpClient::new_streamable_http_client_with_protocol_mode( + "native-mrtr-test", + &format!("{}/mcp", server.uri()), + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + McpProtocolMode::V20260728, + ) + .await?; + client + .initialize( + InitializeRequestParams::new(capabilities, Implementation::new("test", "0.0.0")), + Some(Duration::from_secs(/*secs*/ 5)), + handler, + ) + .await?; + Ok(Arc::new(client)) +} + +async fn call(client: &RmcpClient) -> anyhow::Result { + Ok(serde_json::to_value( + client + .call_tool( + "verified_operation".into(), + Some(json!({"resource": "example"})), + Some(json!({"caller": "preserved"})), + Some(Duration::from_secs(/*secs*/ 5)), + ) + .await?, + )?) +} + +fn response(action: ElicitationAction, content: Value) -> ElicitationResponse { + ElicitationResponse { + action, + content: Some(content), + meta: Some(json!({"discard": true})), + } +} + +#[tokio::test] +async fn native_mrtr_returns_validated_proof_or_cancellation_in_content() -> anyhow::Result<()> { + let proof = json!({"credentialId": "test-credential", "signature": "BAUG"}); + for (action, content, expected) in [ + ( + ElicitationAction::Accept, + proof.clone(), + json!({"action": "accept", "content": proof}), + ), + ( + ElicitationAction::Cancel, + proof.clone(), + json!({"action": "cancel"}), + ), + ( + ElicitationAction::Decline, + proof, + json!({"action": "decline"}), + ), + ( + ElicitationAction::Accept, + json!({"signature": "invalid"}), + json!({"action": "cancel"}), + ), + ] { + let server = MockServer::start().await; + let final_result = + json!({"resultType": "complete", "content": [{"type": "text", "text": "done"}]}); + let final_response = final_result.clone(); + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") if body.pointer("/params/inputResponses").is_none() => { + let mut result = input_required(native_request()); + result["inputRequests"]["confirmation"] = elicitation_request("form"); + result["inputRequests"]["richConfirmation"] = rich_form_request(); + sse_result_response(&body, result) + } + Some("tools/call") => { + assert_eq!( + body["params"]["inputResponses"], + json!({ + "verification": expected, + "confirmation": {"action": "accept", "content": {"confirmed": true}}, + "richConfirmation": {"action": "accept", "content": {"confirmed": true}} + }) + ); + assert_eq!(body["params"]["requestState"], OPAQUE_STATE); + assert_eq!(body["params"]["name"], "verified_operation"); + assert_eq!(body["params"]["arguments"], json!({"resource": "example"})); + assert_eq!(body["params"]["_meta"]["caller"], "preserved"); + result_response(&body, final_response.clone()) + } + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 3) + .mount(&server) + .await; + let mut capabilities = capabilities(); + capabilities + .extensions + .as_mut() + .unwrap() + .get_mut(OPENAI_ELICITATION_EXTENSION_ID) + .unwrap() + .insert("form".into(), json!({})); + capabilities.elicitation = + Some(ElicitationCapability::new().with_form(FormElicitationCapability::new())); + let client = client( + &server, + capabilities, + Box::new(move |_, request| { + if let Elicitation::OpenAiElicitationForm { meta, message, requested_schema } = request { + let expected = rich_form_request(); + assert_eq!( + json!({"_meta": meta, "message": message, "requestedSchema": requested_schema}), + json!({ + "_meta": expected["params"]["_meta"], + "message": expected["params"]["message"], + "requestedSchema": expected["params"]["requestedSchema"] + }) + ); + return Box::pin(async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"confirmed": true})), + meta: None, + }) + }); + } + if let Elicitation::Mcp(ElicitRequestParams::FormElicitationParams { + meta, .. + }) = request + { + assert_eq!( + meta.and_then(|meta| meta.get("inputContext").cloned()), + Some(json!("server-context")) + ); + return Box::pin(async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"confirmed": true})), + meta: None, + }) + }); + } + assert_eq!( + request, + Elicitation::UserVerification { + title: "Verify action".into(), + description: "Verify the requested operation".into(), + challenge: "AQID".into(), + } + ); + let response = response(action.clone(), content.clone()); + Box::pin(async move { Ok(response) }) + }), + ) + .await?; + assert_eq!(call(&client).await?, final_result); + let calls: Vec<_> = server + .received_requests() + .await + .unwrap() + .into_iter() + .map(|request| request.body_json::().unwrap()) + .filter(|body| body["method"] == "tools/call") + .collect(); + assert_ne!(calls[0]["id"], calls[1]["id"]); + client.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn native_mrtr_bounds_repeated_input_requests() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") => result_response(&body, input_required(native_request())), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(1 + rmcp::model::DEFAULT_MRTR_MAX_ROUNDS as u64) + .mount(&server) + .await; + let client = client( + &server, + capabilities(), + Box::new(|_, _| { + Box::pin(async { + Ok(response( + ElicitationAction::Accept, + json!({"credentialId": "test", "signature": "BAUG"}), + )) + }) + }), + ) + .await?; + let error = call(&client).await.unwrap_err(); + assert!(error.to_string().contains("MRTR"), "{error}"); + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn native_mrtr_rejects_unadvertised_and_malformed_requests_without_prompting() +-> anyhow::Result<()> { + let mut malformed = native_request(); + malformed["params"]["challenge"] = json!("not base64!"); + for (capabilities, request) in [ + (ClientCapabilities::default(), native_request()), + (capabilities(), rich_form_request()), + (capabilities(), malformed), + ( + capabilities(), + json!({"method": "arbitrary/custom", "params": {}}), + ), + ] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |http: &Request| { + let body: Value = http.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") => result_response(&body, input_required(request.clone())), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 2) + .mount(&server) + .await; + let client = client( + &server, + capabilities, + Box::new(|_, _| panic!("must not prompt")), + ) + .await?; + assert!(call(&client).await.is_err()); + client.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn mrtr_rejects_unsupported_elicitation_modes_without_prompting() -> anyhow::Result<()> { + for mode in [None, Some(json!(null)), Some(json!(0)), Some(json!("url"))] { + let server = MockServer::start().await; + let mut input = native_request(); + if let Some(mode) = mode { + input["params"]["mode"] = mode; + } else { + input["params"].as_object_mut().unwrap().remove("mode"); + } + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") => result_response(&body, input_required(input.clone())), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 2) + .mount(&server) + .await; + let client = client( + &server, + capabilities(), + Box::new(|_, _| panic!("must not prompt")), + ) + .await?; + let error = call(&client).await.unwrap_err(); + assert_eq!( + codex_rmcp_client::mcp_error(&error), + Some(&rmcp::ErrorData::invalid_request( + "unsupported OpenAI elicitation mode", + /*data*/ None + )) + ); + client.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn mrtr_preserves_auth_challenges_after_input_responses() -> anyhow::Result<()> { + for (request, action, content) in [ + ( + elicitation_request("form"), + ElicitationAction::Accept, + json!({"confirmed": true}), + ), + ( + rich_form_request(), + ElicitationAction::Accept, + json!({"confirmed": true}), + ), + ( + native_request(), + ElicitationAction::Accept, + json!({"credentialId": "test", "signature": "BAUG"}), + ), + (native_request(), ElicitationAction::Cancel, Value::Null), + (native_request(), ElicitationAction::Decline, Value::Null), + ] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |http: &Request| { + let body: Value = http.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") if body.pointer("/params/inputResponses").is_none() => { + result_response(&body, input_required(request.clone())) + } + Some("tools/call") => ResponseTemplate::new(/*s*/ 401) + .insert_header("www-authenticate", "Bearer error=\"invalid_token\""), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 3) + .mount(&server) + .await; + let mut capabilities = capabilities(); + capabilities + .extensions + .as_mut() + .unwrap() + .get_mut(OPENAI_ELICITATION_EXTENSION_ID) + .unwrap() + .insert("form".into(), json!({})); + capabilities.elicitation = + Some(ElicitationCapability::new().with_form(FormElicitationCapability::new())); + let client = client( + &server, + capabilities, + Box::new(move |_, _| { + let response = response(action.clone(), content.clone()); + Box::pin(async move { Ok(response) }) + }), + ) + .await?; + assert_eq!( + call(&client).await?, + serde_json::to_value( + rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text( + "Authentication required", + )]) + .with_meta(Some(rmcp::model::MetaObject::from( + serde_json::Map::from_iter([( + "mcp/www_authenticate".into(), + json!(["Bearer error=\"invalid_token\""]), + )]) + ))) + )? + ); + client.shutdown().await; + server.verify().await; + } + Ok(()) +} + +#[tokio::test] +async fn native_mrtr_does_not_restart_after_continuation_session_expiry() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") if body.pointer("/params/inputResponses").is_none() => { + result_response(&body, input_required(native_request())) + } + Some("tools/call") => ResponseTemplate::new(/*s*/ 404), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 3) + .mount(&server) + .await; + let client = client( + &server, + capabilities(), + Box::new(|_, _| { + Box::pin(async { + Ok(response( + ElicitationAction::Accept, + json!({"credentialId": "test", "signature": "BAUG"}), + )) + }) + }), + ) + .await?; + let error = call(&client).await.unwrap_err(); + assert!( + error.to_string().contains("the tool was not restarted"), + "{error}" + ); + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn native_mrtr_does_not_restart_after_proof_then_state_only_round() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") if body.pointer("/params/requestState").is_none() => { + result_response(&body, input_required(native_request())) + } + Some("tools/call") if body.pointer("/params/inputResponses").is_some() => { + result_response( + &body, + json!({"resultType": "input_required", "requestState": "after-proof"}), + ) + } + Some("tools/call") => ResponseTemplate::new(/*s*/ 404), + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 4) + .mount(&server) + .await; + let client = client( + &server, + capabilities(), + Box::new(|_, _| { + Box::pin(async { + Ok(response( + ElicitationAction::Accept, + json!({"credentialId": "test", "signature": "BAUG"}), + )) + }) + }), + ) + .await?; + let error = call(&client).await.unwrap_err(); + assert!( + error.to_string().contains("the tool was not restarted"), + "{error}" + ); + client.shutdown().await; + server.verify().await; + Ok(()) +} + +#[tokio::test] +async fn native_mrtr_concurrent_prompts_have_independent_cancellation() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().unwrap(); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") if body.pointer("/params/inputResponses").is_none() => { + result_response(&body, input_required(native_request())) + } + Some("tools/call") => { + assert_eq!( + body["params"]["inputResponses"], + json!({"verification": {"action": "cancel"}}) + ); + result_response(&body, json!({"resultType": "complete", "content": []})) + } + other => panic!("unexpected request: {other:?}"), + } + }) + .expect(/*r*/ 4) + .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?) }) + }), + ) + .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(); + assert_ne!(first_id, second_id); + first_call.abort(); + assert!(first_call.await.unwrap_err().is_cancelled()); + timeout(Duration::from_secs(/*secs*/ 5), first_reply.closed()).await?; + assert!(!second_reply.is_closed()); + second_reply + .send(response(ElicitationAction::Cancel, Value::Null)) + .unwrap(); + assert_eq!( + timeout(Duration::from_secs(/*secs*/ 5), second_call).await???, + json!({"resultType": "complete", "content": []}) + ); + client.shutdown().await; + Ok(()) +}