diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index b575eb01a5..7916d6c967 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -211,6 +211,14 @@ impl ElicitationRequestManager { let server_name = server_name.clone(); let authority = authority.clone(); async move { + // Activate only once the typed app-server and UI path is available. + if matches!(&elicitation, Elicitation::UserVerification { .. }) { + return Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + }); + } if router.auto_deny() { return Ok(ElicitationResponse { action: ElicitationAction::Decline, @@ -336,7 +344,9 @@ impl ElicitationRequestManager { .get("properties") .and_then(Value::as_object) .is_some_and(|properties| !properties.is_empty()), - Elicitation::Mcp(_) | Elicitation::OpenAiForm { .. } => false, + Elicitation::Mcp(_) + | Elicitation::OpenAiForm { .. } + | Elicitation::UserVerification { .. } => false, }; if !should_surface_form_in_full_access { @@ -374,6 +384,13 @@ impl ElicitationRequestManager { ); let routed_request_id = RequestId::String(public_request_id.clone().into()); let request = match elicitation { + Elicitation::UserVerification { .. } => { + return Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + }); + } Elicitation::Mcp(rmcp::model::ElicitRequestParams::FormElicitationParams { meta, message, @@ -491,7 +508,8 @@ fn can_auto_accept_elicitation(elicitation: &Elicitation) -> bool { } Elicitation::Mcp(_) | Elicitation::OpenAiForm { .. } - | Elicitation::OpenAiElicitationForm { .. } => false, + | Elicitation::OpenAiElicitationForm { .. } + | Elicitation::UserVerification { .. } => false, } } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index bac8c641cb..1444c1e64c 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1023,7 +1023,9 @@ fn guardian_elicitation_review_request( "guardian MCP elicitation review does not support this elicitation mode", ); } - Elicitation::OpenAiForm { .. } | Elicitation::OpenAiElicitationForm { .. } => { + Elicitation::OpenAiForm { .. } + | Elicitation::OpenAiElicitationForm { .. } + | Elicitation::UserVerification { .. } => { return GuardianElicitationReview::NotRequested; } }; diff --git a/codex-rs/core/tests/suite/mcp_user_verification.rs b/codex-rs/core/tests/suite/mcp_user_verification.rs new file mode 100644 index 0000000000..47354a2339 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_user_verification.rs @@ -0,0 +1,172 @@ +//! Before activation, verification never prompts or prevents the enclosing tool/turn completing. + +use anyhow::Result; +use codex_core::StartThreadOptions; +use codex_core::TurnInputRequest; +use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::collections::HashMap; +use test_case::test_case; + +const SERVER: &str = r#" +import json +import sys + +def send(message): + print(json.dumps({"jsonrpc": "2.0", **message}), flush=True) + +extensions = {} +pending = None +for line in sys.stdin: + request = json.loads(line) + method = request.get("method") + if method == "initialize": + extensions = request["params"]["capabilities"].get("extensions", {}) + result = {"protocolVersion": request["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "verification-test", "version": "1"}} + elif method == "tools/list": + result = {"tools": [{"name": "verify", "inputSchema": {"type": "object", "properties": {}}, + "annotations": {"readOnlyHint": True}}]} + elif method == "tools/call": + pending = request["id"] + send({"id": "verification", "method": "openai/elicitation/create", "params": { + "mode": "openai/userVerification", "title": "Approve", "description": "", "challenge": "AQID"}}) + continue + elif method is None and request.get("id") == "verification": + send({"id": pending, "result": {"content": [{"type": "text", "text": json.dumps({ + "extensions": extensions, "response": request.get("result", request.get("error"))})}]}}) + continue + elif method == "resources/list": + result = {"resources": []} + elif method == "resources/templates/list": + result = {"resourceTemplates": []} + elif "id" not in request: + continue + else: + result = {} + send({"id": request["id"], "result": result}) +"#; + +#[derive(Clone, Copy)] +enum CapabilitySource { + HostProjection, + ExplicitSession, +} + +#[test_case(CapabilitySource::HostProjection; "host_projection_does_not_advertise_verification")] +#[test_case(CapabilitySource::ExplicitSession; "explicit_session_cancels_without_ui")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_user_verification_before_activation(source: CapabilitySource) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "the MCP fixture requires a host Python interpreter"); + let server = responses::start_mock_server().await; + let test = test_codex().build_with_auto_env(&server).await?; + let declarations = HashMap::from([( + OPENAI_ELICITATION_EXTENSION_ID.to_string(), + json!({"userVerification": {}}), + )]); + let client_mcp_extensions = match source { + CapabilitySource::HostProjection => codex_mcp::client_mcp_extensions( + Some(&declarations), + /*legacy_openai_form_elicitation*/ false, + ), + // Exercise the manager's fail-closed route even if a host directly supplies + // this capability before the typed transport/UI activation stage lands. + CapabilitySource::ExplicitSession => ClientMcpExtensions::new(declarations), + }; + let mut config = test.config.clone(); + config.mcp_servers.set(serde_json::from_value(json!({ + "verification": { + "command": if cfg!(windows) { "python" } else { "python3" }, + "args": ["-u", "-c", SERVER], "default_tools_approval_mode": "approve" + } + }))?)?; + let thread = test + .thread_manager + .start_thread(StartThreadOptions { + client_mcp_extensions, + environments: Some(vec![test.executor_environment().selection().clone()]), + ..StartThreadOptions::new(config) + }) + .await? + .thread; + wait_for_mcp_server(&thread, "verification").await?; + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "verify-call", + "mcp__verification", + "verify", + "{}", + ), + responses::ev_completed("tool-response"), + ]), + ) + .await; + let follow_up = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message("done-message", "done"), + responses::ev_completed("done-response"), + ]), + ) + .await; + thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Run verification".into(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&thread, |event| { + assert!( + !matches!(event, EventMsg::ElicitationRequest(_)), + "verification must not prompt before activation" + ); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let output = follow_up + .single_request() + .function_call_output("verify-call"); + let echoed: Value = serde_json::from_str( + output["output"][1]["text"] + .as_str() + .expect("MCP wire response"), + )?; + match source { + CapabilitySource::HostProjection => { + assert!( + echoed["extensions"][OPENAI_ELICITATION_EXTENSION_ID] + .get("userVerification") + .is_none() + ); + assert_eq!( + echoed["response"], + json!({"code": -32601, "message": "openai/elicitation/create"}) + ); + } + CapabilitySource::ExplicitSession => { + assert_eq!( + echoed["extensions"][OPENAI_ELICITATION_EXTENSION_ID], + json!({"userVerification": {}}) + ); + assert_eq!(echoed["response"], json!({"action": "cancel"})); + } + } + thread.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 2040f6ba3b..1c6a629926 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -103,6 +103,7 @@ mod mcp_startup_refresh_http_proxy; mod mcp_tool_cache; mod mcp_tool_exposure; mod mcp_turn_metadata; +mod mcp_user_verification; mod model_overrides; mod model_runtime_selectors; mod model_switching; diff --git a/codex-rs/rmcp-client/src/elicitation_client_service.rs b/codex-rs/rmcp-client/src/elicitation_client_service.rs index 12cad74bd3..1eff73f788 100644 --- a/codex-rs/rmcp-client/src/elicitation_client_service.rs +++ b/codex-rs/rmcp-client/src/elicitation_client_service.rs @@ -1,4 +1,8 @@ +use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; @@ -11,6 +15,7 @@ use rmcp::model::ElicitResult; use rmcp::model::ElicitationAction; use rmcp::model::MetaObject; use rmcp::model::ProtocolVersion; +use rmcp::model::RequestId; use rmcp::model::RequestMetaObject; use rmcp::model::RequestParamsMeta; use rmcp::model::ServerNotification; @@ -22,6 +27,7 @@ use serde::Deserialize; use serde::Serialize; use serde_json::Map; use serde_json::Value; +use tokio::sync::oneshot; use crate::logging_client_handler::LoggingClientHandler; use crate::rmcp_client::Elicitation; @@ -55,8 +61,36 @@ pub(crate) struct ElicitationClientService { handler: LoggingClientHandler, supports_openai_form: bool, supports_openai_elicitation_form: bool, + supports_user_verification: bool, send_elicitation: Arc, pause_state: ElicitationPauseState, + pending_verifications: Arc>, +} + +// A notification handler can run before its request handler. Never evict an early +// cancellation: after saturation, cancel new verifications for this connection. +const MAX_EARLY_CANCELLATIONS: usize = 1024; + +#[derive(Default)] +struct VerificationCancellations { + pending: HashMap>, + early: HashSet, + saturated: bool, +} + +struct PendingVerification { + request_id: RequestId, + cancellations: Arc>, +} + +impl Drop for PendingVerification { + fn drop(&mut self) { + self.cancellations + .lock() + .unwrap_or_else(PoisonError::into_inner) + .pending + .remove(&self.request_id); + } } impl ElicitationClientService { @@ -78,6 +112,13 @@ impl ElicitationClientService { .and_then(|settings| settings.get("form")) .is_some_and(Value::is_object); let send_elicitation = Arc::new(send_elicitation); + let supports_user_verification = client_info + .capabilities + .extensions + .as_ref() + .and_then(|extensions| extensions.get(OPENAI_ELICITATION_EXTENSION_ID)) + .and_then(|settings| settings.get("userVerification")) + .is_some_and(Value::is_object); Self { handler: LoggingClientHandler::new( client_info, @@ -85,8 +126,10 @@ impl ElicitationClientService { ), supports_openai_form, supports_openai_elicitation_form, + supports_user_verification, send_elicitation, pause_state, + pending_verifications: Arc::default(), } } @@ -95,12 +138,55 @@ impl ElicitationClientService { request: Elicitation, context: RequestContext, ) -> Result { - let RequestContext { id, meta, .. } = context; + let RequestContext { id, meta, ct, .. } = context; let request = restore_context_meta(request, meta); + let user_verification = matches!(&request, Elicitation::UserVerification { .. }); + let (cancel_tx, cancel_rx) = oneshot::channel(); + let _pending = if user_verification { + let mut cancellations = self + .pending_verifications + .lock() + .unwrap_or_else(PoisonError::into_inner); + if cancellations.saturated || cancellations.early.remove(&id) { + return Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + }); + } + cancellations.pending.insert(id.clone(), cancel_tx); + Some(PendingVerification { + request_id: id.clone(), + cancellations: Arc::clone(&self.pending_verifications), + }) + } else { + None + }; let _pause = self.pause_state.enter(); - (self.send_elicitation)(id, request) - .await - .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None)) + let response = tokio::select! { + biased; + _ = ct.cancelled(), if user_verification => { + return Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + }); + } + _ = cancel_rx, if user_verification => { + return Ok(ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + }); + } + response = (self.send_elicitation)(id, request) => response, + } + .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?; + Ok(if user_verification { + crate::user_verification::validate_response(response) + } else { + response + }) } } @@ -153,6 +239,23 @@ impl Service for ElicitationClientService { )?)) } } + ServerRequest::CustomRequest(request) + if request.method == OPENAI_ELICITATION_METHOD + && self.supports_user_verification + && request + .params + .as_ref() + .and_then(|params| params.get("mode")) + .and_then(Value::as_str) + == Some(crate::user_verification::MODE) => + { + let response = self + .create_elicitation(crate::user_verification::parse_request(request)?, context) + .await?; + Ok(ClientResult::CustomResult(elicitation_response_result( + response, + )?)) + } ServerRequest::CustomRequest(request) if request.method == OPENAI_FORM_METHOD && self.supports_openai_form => { @@ -190,6 +293,15 @@ impl Service for ElicitationClientService { response, )?)) } + ServerRequest::CustomRequest(request) + if request.method == OPENAI_ELICITATION_METHOD + && self.supports_user_verification => + { + Err(rmcp::ErrorData::invalid_params( + "invalid elicitation mode", + /*data*/ None, + )) + } request => { >::handle_request( &self.handler, @@ -206,6 +318,28 @@ impl Service for ElicitationClientService { notification: ServerNotification, context: NotificationContext, ) -> Result<(), rmcp::ErrorData> { + if let ServerNotification::CancelledNotification(cancelled) = ¬ification + && self.supports_user_verification + && let Some(request_id) = cancelled.params.request_id.as_ref() + { + let mut cancellations = self + .pending_verifications + .lock() + .unwrap_or_else(PoisonError::into_inner); + if let Some(cancel) = cancellations.pending.remove(request_id) { + let _ = cancel.send(()); + } else if !cancellations.saturated && !cancellations.early.contains(request_id) { + // Bound both the number and size of retained request IDs. + if cancellations.early.len() == MAX_EARLY_CANCELLATIONS + || matches!(request_id, RequestId::String(id) if id.len() > 1024) + { + cancellations.saturated = true; + cancellations.early.clear(); + } else { + cancellations.early.insert(request_id.clone()); + } + } + } >::handle_notification( &self.handler, notification, @@ -251,6 +385,7 @@ fn restore_context_meta( } match &mut request { + Elicitation::UserVerification { .. } => {} Elicitation::Mcp(request) => request .meta_mut() .get_or_insert_with(RequestMetaObject::new) @@ -562,3 +697,7 @@ mod tests { RequestMetaObject::from(map) } } + +#[cfg(test)] +#[path = "user_verification_dispatch_tests.rs"] +mod user_verification_dispatch_tests; diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 4c41f0823a..7792e53806 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -26,6 +26,7 @@ mod rmcp_client; mod service_error; mod startup_error; mod stdio_server_launcher; +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 31998b31b4..048e893bae 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -194,7 +194,7 @@ pub(crate) struct ElicitationPauseState { } impl ElicitationPauseState { - fn new() -> Self { + pub(crate) fn new() -> Self { let (paused, _rx) = watch::channel(false); Self { active_count: Arc::new(AtomicUsize::new(0)), @@ -314,6 +314,11 @@ pub enum Elicitation { message: String, requested_schema: serde_json::Value, }, + UserVerification { + title: String, + description: String, + challenge: String, + }, } impl Elicitation { @@ -323,6 +328,7 @@ impl Elicitation { Self::OpenAiForm { meta, .. } | Self::OpenAiElicitationForm { meta, .. } => { meta.as_ref().and_then(serde_json::Value::as_object) } + Self::UserVerification { .. } => None, } } } @@ -1634,6 +1640,10 @@ async fn create_oauth_transport_and_runtime( }) } +#[cfg(test)] +#[path = "user_verification_cancellation_tests.rs"] +mod user_verification_cancellation_tests; + #[cfg(test)] mod tests { use std::time::Duration; diff --git a/codex-rs/rmcp-client/src/user_verification.rs b/codex-rs/rmcp-client/src/user_verification.rs new file mode 100644 index 0000000000..fb93b42ff6 --- /dev/null +++ b/codex-rs/rmcp-client/src/user_verification.rs @@ -0,0 +1,91 @@ +//! Validation for the user-verification elicitation extension. + +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use rmcp::model::CustomRequest; +use rmcp::model::ElicitationAction; +use serde::Deserialize; + +use crate::rmcp_client::Elicitation; +use crate::rmcp_client::ElicitationResponse; + +pub(crate) const MODE: &str = "openai/userVerification"; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct RequestParams { + mode: String, + title: String, + description: String, + challenge: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct Proof<'a> { + credential_id: &'a str, + signature: &'a str, +} + +pub(crate) fn parse_request(request: CustomRequest) -> Result { + let params = request + .params_as::() + .ok() + .flatten() + .ok_or_else(invalid_request)?; + if params.mode != MODE + || params.title.is_empty() + || params.title.len() > 256 + || params.description.len() > 4096 + || !valid_bytes(¶ms.challenge, /*max_decoded_bytes*/ 4096) + { + return Err(invalid_request()); + } + Ok(Elicitation::UserVerification { + title: params.title, + description: params.description, + challenge: params.challenge, + }) +} + +/// Accept only a bounded proof, and never return proof material for cancellation or rejection. +pub(crate) fn validate_response(mut response: ElicitationResponse) -> ElicitationResponse { + response.meta = None; + match response.action { + ElicitationAction::Accept => { + let proof = response + .content + .as_ref() + .and_then(|content| Proof::deserialize(content).ok()); + if proof.is_some_and(|proof| { + !proof.credential_id.is_empty() + && proof.credential_id.len() <= 1024 + && valid_bytes(proof.signature, /*max_decoded_bytes*/ 128) + }) { + return response; + } + tracing::warn!("user-verification acceptance omitted a valid proof; cancelling"); + response.action = ElicitationAction::Cancel; + } + ElicitationAction::Decline | ElicitationAction::Cancel => {} + _ => response.action = ElicitationAction::Cancel, + } + response.content = None; + response +} + +fn valid_bytes(encoded: &str, max_decoded_bytes: usize) -> bool { + !encoded.is_empty() + && encoded.len() <= max_decoded_bytes.div_ceil(3) * 4 + && URL_SAFE_NO_PAD + .decode(encoded) + .is_ok_and(|bytes| !bytes.is_empty() && bytes.len() <= max_decoded_bytes) +} + +fn invalid_request() -> rmcp::ErrorData { + rmcp::ErrorData::invalid_params("invalid user-verification request", /*data*/ None) +} + +#[cfg(test)] +#[path = "user_verification_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs b/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs new file mode 100644 index 0000000000..5f69fcebf6 --- /dev/null +++ b/codex-rs/rmcp-client/src/user_verification_cancellation_tests.rs @@ -0,0 +1,192 @@ +//! Ensures MCP service cancellation releases the pending verification route and tool timeout pause. + +use super::ElicitationClientService; +use super::ElicitationPauseState; +use super::ElicitationResponse; +use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID; +use rmcp::RoleServer; +use rmcp::model::CancelledNotification; +use rmcp::model::CancelledNotificationParam; +use rmcp::model::ClientInfo; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::CustomRequest; +use rmcp::model::ElicitationAction; +use rmcp::model::RequestId; +use rmcp::model::ServerJsonRpcMessage; +use rmcp::model::ServerNotification; +use rmcp::model::ServerRequest; +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 user_verification_service_cancellation_drops_pending_response() -> anyhow::Result<()> { + let pause_state = ElicitationPauseState::new(); + let mut paused = pause_state.subscribe(); + let (route_tx, mut route_rx) = mpsc::unbounded_channel(); + let mut client_info = ClientInfo::default(); + client_info.capabilities.extensions = Some( + [( + OPENAI_ELICITATION_EXTENSION_ID.to_string(), + serde_json::Map::from_iter([("userVerification".to_string(), json!({}))]), + )] + .into_iter() + .collect(), + ); + let service = ElicitationClientService::new( + client_info, + Box::new(move |_, _| { + let (response_tx, response_rx) = oneshot::channel(); + route_tx + .send(response_tx) + .expect("observe pending verification"); + Box::pin(async move { Ok(response_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); + server + .send(ServerJsonRpcMessage::request( + ServerRequest::CustomRequest(CustomRequest::new( + "openai/elicitation/create", + Some(json!({ + "mode": "openai/userVerification", + "title": "Approve", + "description": "", + "challenge": "AQID", + })), + )), + RequestId::Number(1), + )) + .await?; + let mut response_tx = timeout(Duration::from_secs(/*secs*/ 5), route_rx.recv()) + .await? + .expect("verification was routed to the UI"); + assert!(*paused.borrow()); + + client.cancel().await?; + + timeout(Duration::from_secs(/*secs*/ 5), response_tx.closed()).await?; + timeout( + Duration::from_secs(/*secs*/ 5), + paused.wait_for(|paused| !*paused), + ) + .await??; + Ok(()) +} + +#[tokio::test] +async fn cancelling_one_verification_leaves_the_mcp_connection_and_other_requests_alive() +-> anyhow::Result<()> { + let pause_state = ElicitationPauseState::new(); + let mut paused = pause_state.subscribe(); + let (route_tx, mut route_rx) = mpsc::unbounded_channel(); + let mut client_info = ClientInfo::default(); + client_info.capabilities.extensions = Some( + [( + OPENAI_ELICITATION_EXTENSION_ID.to_string(), + serde_json::Map::from_iter([("userVerification".to_string(), json!({}))]), + )] + .into_iter() + .collect(), + ); + let service = ElicitationClientService::new( + client_info, + Box::new(move |id, _| { + let (response_tx, response_rx) = oneshot::channel(); + route_tx + .send((id, response_tx)) + .expect("observe verification"); + Box::pin(async move { Ok(response_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); + + for id in [1, 2] { + server + .send(ServerJsonRpcMessage::request( + ServerRequest::CustomRequest(CustomRequest::new( + "openai/elicitation/create", + Some(json!({ + "mode": "openai/userVerification", + "title": "Approve", + "description": "", + "challenge": "AQID", + })), + )), + RequestId::Number(id), + )) + .await?; + } + let (first_id, mut first) = timeout(Duration::from_secs(/*secs*/ 5), route_rx.recv()) + .await? + .expect("first verification reached UI"); + let (second_id, second) = timeout(Duration::from_secs(/*secs*/ 5), route_rx.recv()) + .await? + .expect("second verification reached UI"); + assert_ne!(first_id, second_id); + assert!(*paused.borrow()); + + server + .send(ServerJsonRpcMessage::notification( + ServerNotification::CancelledNotification(CancelledNotification::new( + CancelledNotificationParam::new(Some(first_id.clone()), None), + )), + )) + .await?; + timeout(Duration::from_secs(/*secs*/ 5), first.closed()).await?; + assert!(!second.is_closed()); + assert!(*paused.borrow()); + let cancel = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) + .await? + .expect("cancelled verification still sends a response"); + let ClientJsonRpcMessage::Response(cancel) = cancel else { + anyhow::bail!("expected a cancellation response"); + }; + assert_eq!(cancel.id, first_id); + assert_eq!(serde_json::to_value(cancel.result)?["action"], "cancel"); + + for request_id in [first_id, RequestId::Number(999)] { + server + .send(ServerJsonRpcMessage::notification( + ServerNotification::CancelledNotification(CancelledNotification::new( + CancelledNotificationParam::new(Some(request_id), None), + )), + )) + .await?; + } + assert!(!second.is_closed()); + + second + .send(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"credentialId": "AQID", "signature": "BAUG"})), + meta: None, + }) + .expect("second verification remains pending"); + let accepted = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) + .await? + .expect("second request survived cancellation"); + let ClientJsonRpcMessage::Response(accepted) = accepted else { + anyhow::bail!("expected the second verification response"); + }; + assert_eq!(accepted.id, second_id); + assert_eq!(serde_json::to_value(accepted.result)?["action"], "accept"); + timeout( + Duration::from_secs(/*secs*/ 5), + paused.wait_for(|paused| !*paused), + ) + .await??; + client.cancel().await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/user_verification_dispatch_tests.rs b/codex-rs/rmcp-client/src/user_verification_dispatch_tests.rs new file mode 100644 index 0000000000..ea1ac31610 --- /dev/null +++ b/codex-rs/rmcp-client/src/user_verification_dispatch_tests.rs @@ -0,0 +1,164 @@ +use super::*; +use pretty_assertions::assert_eq; +use rmcp::RoleServer; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::ServerJsonRpcMessage; +use rmcp::service::serve_directly; +use rmcp::transport::IntoTransport; +use rmcp::transport::Transport; +use serde_json::json; +use std::time::Duration; +use tokio::time::timeout; + +fn service() -> ElicitationClientService { + let mut info = ClientInfo::default(); + info.capabilities.extensions = Some( + [( + OPENAI_ELICITATION_EXTENSION_ID.into(), + Map::from_iter([("userVerification".into(), json!({}))]), + )] + .into_iter() + .collect(), + ); + ElicitationClientService::new( + info, + Box::new(|_, _| panic!("cancelled or malformed verification must not reach the UI")), + ElicitationPauseState::new(), + ) +} + +#[tokio::test] +async fn user_verification_remembers_cancellation_before_request_handler_runs() -> anyhow::Result<()> +{ + let service = service(); + let (client_transport, server_transport) = tokio::io::duplex(/*max_buf_size*/ 4096); + let client = serve_directly(service.clone(), client_transport, /*peer_info*/ None); + let mut server = IntoTransport::::into_transport(server_transport); + server + .send(serde_json::from_value::(json!({ + "jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 1} + }))?) + .await?; + // Force the scheduling order that RMCP's independently spawned handlers allow. + timeout(Duration::from_secs(/*secs*/ 5), async { + loop { + if service + .pending_verifications + .lock() + .unwrap() + .early + .contains(&RequestId::Number(1)) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await?; + server.send(serde_json::from_value::(json!({ + "jsonrpc": "2.0", "id": 1, "method": OPENAI_ELICITATION_METHOD, + "params": {"mode": crate::user_verification::MODE, "title": "Approve", "description": "", "challenge": "AQID"} + }))?).await?; + let response = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) + .await? + .unwrap(); + assert_eq!( + serde_json::to_value(response)?, + json!({ + "jsonrpc": "2.0", "id": 1, "result": {"action": "cancel"} + }) + ); + assert!( + service + .pending_verifications + .lock() + .unwrap() + .early + .is_empty() + ); + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn user_verification_early_cancellation_storage_fails_closed_at_capacity() +-> anyhow::Result<()> { + let service = service(); + let (client_transport, server_transport) = tokio::io::duplex(/*max_buf_size*/ 4096); + let client = serve_directly(service.clone(), client_transport, /*peer_info*/ None); + let mut server = IntoTransport::::into_transport(server_transport); + for id in 0..=MAX_EARLY_CANCELLATIONS { + server + .send(serde_json::from_value::(json!({ + "jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": id} + }))?) + .await?; + } + timeout(Duration::from_secs(/*secs*/ 5), async { + while !service.pending_verifications.lock().unwrap().saturated { + tokio::task::yield_now().await; + } + }) + .await?; + assert!( + service + .pending_verifications + .lock() + .unwrap() + .early + .is_empty() + ); + let response = service.handle_request( + ServerRequest::CustomRequest(CustomRequest::new(OPENAI_ELICITATION_METHOD, Some(json!({ + "mode": crate::user_verification::MODE, "title": "Approve", "description": "", "challenge": "AQID" + })))), + RequestContext::new(RequestId::Number(5000), client.peer().clone()), + ).await?; + assert_eq!(serde_json::to_value(response)?, json!({"action": "cancel"})); + // Ordinary requests still work after verification fails closed. + server + .send(serde_json::from_value::(json!({ + "jsonrpc": "2.0", "id": 5001, "method": "ping" + }))?) + .await?; + let response = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) + .await? + .unwrap(); + assert_eq!( + serde_json::to_value(response)?, + json!({"jsonrpc": "2.0", "id": 5001, "result": {}}) + ); + client.cancel().await?; + Ok(()) +} + +#[tokio::test] +async fn user_verification_dispatch_rejects_malformed_modes_as_invalid_params() -> anyhow::Result<()> +{ + 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); + for params in [json!({}), json!({"mode": 7}), json!({"mode": "unknown"})] { + server + .send(ServerJsonRpcMessage::request( + ServerRequest::CustomRequest(CustomRequest::new( + OPENAI_ELICITATION_METHOD, + Some(params), + )), + RequestId::Number(1), + )) + .await?; + let response = timeout(Duration::from_secs(/*secs*/ 5), server.receive()) + .await? + .unwrap(); + let ClientJsonRpcMessage::Error(response) = response else { + anyhow::bail!("expected invalid params"); + }; + assert_eq!( + response.error, + rmcp::ErrorData::invalid_params("invalid elicitation mode", /*data*/ None) + ); + } + client.cancel().await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/user_verification_tests.rs b/codex-rs/rmcp-client/src/user_verification_tests.rs new file mode 100644 index 0000000000..1465f041dc --- /dev/null +++ b/codex-rs/rmcp-client/src/user_verification_tests.rs @@ -0,0 +1,112 @@ +use super::*; +use pretty_assertions::assert_eq; +use serde_json::json; + +#[test] +fn user_verification_request_preserves_signed_bytes_and_display_text() { + assert_eq!( + parse_request(CustomRequest::new( + "openai/elicitation/create", + Some(json!({ + "mode": MODE, + "title": "Approve purchase", + "description": "Pay $200 to Example Store", + "challenge": "AQID", + })), + )) + .unwrap(), + Elicitation::UserVerification { + title: "Approve purchase".to_string(), + description: "Pay $200 to Example Store".to_string(), + challenge: "AQID".to_string(), + } + ); +} + +#[test] +fn user_verification_rejects_invalid_or_unbounded_requests_without_echoing_values() { + for (field, value) in [ + ("mode", json!("unknown")), + ("title", json!("")), + ("title", json!("x".repeat(257))), + ("description", json!("x".repeat(4097))), + ("challenge", json!("")), + ("challenge", json!("not base64url!")), + ("challenge", json!(URL_SAFE_NO_PAD.encode(vec![0; 4097]))), + ("unexpected", json!("secret")), + ] { + let mut params = json!({ + "mode": MODE, "title": "Approve", "description": "", "challenge": "AQID" + }); + params[field] = value; + assert_eq!( + parse_request(CustomRequest::new( + "openai/elicitation/create", + Some(params) + )), + Err(invalid_request()), + ); + } +} + +#[test] +fn user_verification_acceptance_returns_proof_only_in_content() { + let content = Some(json!({"credentialId": "credential:用户/1=", "signature": "BAUG"})); + assert_eq!( + validate_response(ElicitationResponse { + action: ElicitationAction::Accept, + content: content.clone(), + meta: Some(json!({"untrusted": "ignored"})), + }), + ElicitationResponse { + action: ElicitationAction::Accept, + content, + meta: None, + } + ); +} + +#[test] +fn user_verification_cancels_acceptance_without_a_valid_bounded_proof() { + for content in [ + None, + Some(json!({})), + Some(json!({"credentialId": "AQID", "signature": ""})), + Some(json!({"credentialId": "AQID", "signature": "BAUG="})), + Some(json!({"credentialId": "é".repeat(513), "signature": "BAUG"})), + Some(json!({"credentialId": "", "signature": "BAUG"})), + Some(json!({"credentialId": "AQID", "signature": URL_SAFE_NO_PAD.encode(vec![0; 129])})), + Some(json!({"credentialId": "AQID", "signature": "BAUG", "extra": true})), + ] { + assert_eq!( + validate_response(ElicitationResponse { + action: ElicitationAction::Accept, + content, + meta: None, + }), + ElicitationResponse { + action: ElicitationAction::Cancel, + content: None, + meta: None, + } + ); + } +} + +#[test] +fn user_verification_decline_and_cancel_discard_proof_material() { + for action in [ElicitationAction::Decline, ElicitationAction::Cancel] { + assert_eq!( + validate_response(ElicitationResponse { + action: action.clone(), + content: Some(json!({"credentialId": "AQID", "signature": "BAUG"})), + meta: Some(json!({"secret": "discarded"})), + }), + ElicitationResponse { + action, + content: None, + meta: None + } + ); + } +}