diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f78c19328f..4eeb1746bc 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -35,7 +35,7 @@ pub mod mcp; mod mcp_connection_manager; pub mod openai_models; pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY; -pub use mcp_connection_manager::MCP_SANDBOX_STATE_NOTIFICATION; +pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD; pub use mcp_connection_manager::SandboxState; mod mcp_tool_call; mod message_history; diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 3213b22b71..0797ceede5 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -184,17 +184,20 @@ struct ManagedClient { } impl ManagedClient { + /// Returns once the server has ack'd the sandbox state update. async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> { if !self.server_supports_sandbox_state_capability { return Ok(()); } - self.client - .send_custom_notification( - MCP_SANDBOX_STATE_NOTIFICATION, + let _response = self + .client + .send_custom_request( + MCP_SANDBOX_STATE_METHOD, Some(serde_json::to_value(sandbox_state)?), ) - .await + .await; + Ok(()) } } @@ -253,9 +256,9 @@ impl AsyncManagedClient { pub const MCP_SANDBOX_STATE_CAPABILITY: &str = "codex/sandbox-state"; -/// Custom MCP notification for sandbox state updates. +/// Custom MCP request to push sandbox state updates. /// When used, the `params` field of the notification is [`SandboxState`]. -pub const MCP_SANDBOX_STATE_NOTIFICATION: &str = "codex/sandbox-state/update"; +pub const MCP_SANDBOX_STATE_METHOD: &str = "codex/sandbox-state/update"; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/codex-rs/exec-server/src/posix/mcp.rs b/codex-rs/exec-server/src/posix/mcp.rs index 3fec7e4dd9..d07774a355 100644 --- a/codex-rs/exec-server/src/posix/mcp.rs +++ b/codex-rs/exec-server/src/posix/mcp.rs @@ -5,7 +5,7 @@ use std::time::Duration; use anyhow::Context as _; use anyhow::Result; use codex_core::MCP_SANDBOX_STATE_CAPABILITY; -use codex_core::MCP_SANDBOX_STATE_NOTIFICATION; +use codex_core::MCP_SANDBOX_STATE_METHOD; use codex_core::SandboxState; use codex_core::protocol::SandboxPolicy; use codex_execpolicy::Policy; @@ -15,6 +15,8 @@ use rmcp::ServerHandler; use rmcp::ServiceExt; use rmcp::handler::server::router::tool::ToolRouter; use rmcp::handler::server::wrapper::Parameters; +use rmcp::model::CustomRequest; +use rmcp::model::CustomResult; use rmcp::model::*; use rmcp::schemars; use rmcp::service::RequestContext; @@ -23,8 +25,8 @@ use rmcp::tool; use rmcp::tool_handler; use rmcp::tool_router; use rmcp::transport::stdio; +use serde_json::json; use tokio::sync::RwLock; -use tracing::debug; use crate::posix::escalate_server::EscalateServer; use crate::posix::escalate_server::{self}; @@ -146,6 +148,13 @@ impl ExecTool { } } +#[derive(Default)] +pub struct CodexSandboxStateUpdateMethod; + +impl rmcp::model::ConstString for CodexSandboxStateUpdateMethod { + const VALUE: &'static str = MCP_SANDBOX_STATE_METHOD; +} + #[tool_handler] impl ServerHandler for ExecTool { fn get_info(&self) -> ServerInfo { @@ -181,29 +190,33 @@ impl ServerHandler for ExecTool { Ok(self.get_info()) } - async fn on_custom_notification( + async fn on_custom_request( &self, - notification: rmcp::model::CustomNotification, - _context: rmcp::service::NotificationContext, - ) { - let rmcp::model::CustomNotification { method, params, .. } = notification; - if method == MCP_SANDBOX_STATE_NOTIFICATION - && let Some(params) = params - { - match serde_json::from_value::(params) { - Ok(sandbox_state) => { - debug!( - ?sandbox_state.sandbox_policy, - "received sandbox state notification" - ); - let mut state = self.sandbox_state.write().await; - *state = Some(sandbox_state); - } - Err(err) => { - tracing::warn!(?err, "failed to deserialize sandbox state notification"); - } - } + request: CustomRequest, + _context: rmcp::service::RequestContext, + ) -> Result { + let CustomRequest { method, params, .. } = request; + if method != MCP_SANDBOX_STATE_METHOD { + return Err(McpError::method_not_found::()); } + + let Some(params) = params else { + return Err(McpError::invalid_params( + "missing params for sandbox state request".to_string(), + None, + )); + }; + + let Ok(sandbox_state) = serde_json::from_value::(params.clone()) else { + return Err(McpError::invalid_params( + "failed to deserialize sandbox state".to_string(), + Some(params), + )); + }; + + *self.sandbox_state.write().await = Some(sandbox_state); + + Ok(CustomResult::new(json!({ "status": "ok" }))) } } diff --git a/codex-rs/exec-server/tests/common/lib.rs b/codex-rs/exec-server/tests/common/lib.rs index 99587a2ad5..82f98c0012 100644 --- a/codex-rs/exec-server/tests/common/lib.rs +++ b/codex-rs/exec-server/tests/common/lib.rs @@ -1,4 +1,4 @@ -use codex_core::MCP_SANDBOX_STATE_NOTIFICATION; +use codex_core::MCP_SANDBOX_STATE_METHOD; use codex_core::SandboxState; use codex_core::protocol::SandboxPolicy; use rmcp::ClientHandler; @@ -7,9 +7,10 @@ use rmcp::RoleClient; use rmcp::Service; use rmcp::model::ClientCapabilities; use rmcp::model::ClientInfo; +use rmcp::model::ClientRequest; use rmcp::model::CreateElicitationRequestParam; use rmcp::model::CreateElicitationResult; -use rmcp::model::CustomNotification; +use rmcp::model::CustomRequest; use rmcp::model::ElicitationAction; use rmcp::service::RunningService; use rmcp::transport::ConfigureCommandExt; @@ -92,7 +93,7 @@ where codex_linux_sandbox_exe, sandbox_cwd: sandbox_cwd.as_ref().to_path_buf(), }; - send_sandbox_notification(sandbox_state, service).await + send_sandbox_state_update(sandbox_state, service).await } pub async fn notify_writable_sandbox_only_one_folder( @@ -119,22 +120,21 @@ where codex_linux_sandbox_exe, sandbox_cwd: writable_folder.as_ref().to_path_buf(), }; - send_sandbox_notification(sandbox_state, service).await + send_sandbox_state_update(sandbox_state, service).await } -async fn send_sandbox_notification( +async fn send_sandbox_state_update( sandbox_state: SandboxState, service: &RunningService, ) -> anyhow::Result<()> where S: Service + ClientHandler, { - let sandbox_state_notification = CustomNotification::new( - MCP_SANDBOX_STATE_NOTIFICATION, - Some(serde_json::to_value(sandbox_state)?), - ); service - .send_notification(sandbox_state_notification.into()) + .send_request(ClientRequest::CustomRequest(CustomRequest::new( + MCP_SANDBOX_STATE_METHOD, + Some(serde_json::to_value(sandbox_state)?), + ))) .await?; Ok(()) } diff --git a/codex-rs/exec-server/tests/suite/accept_elicitation.rs b/codex-rs/exec-server/tests/suite/accept_elicitation.rs index b703eaf4a7..5f9b53a523 100644 --- a/codex-rs/exec-server/tests/suite/accept_elicitation.rs +++ b/codex-rs/exec-server/tests/suite/accept_elicitation.rs @@ -3,7 +3,6 @@ use std::borrow::Cow; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; -use std::time::Duration; use anyhow::Context; use anyhow::Result; @@ -84,18 +83,6 @@ prefix_rule( }; notify_readable_sandbox(&project_root_path, codex_linux_sandbox_exe, &service).await?; - // TODO(mbolin): Remove this hack to remove flakiness when possible. - // As noted in the commentary on https://github.com/openai/codex/pull/7832, - // an rmcp server does not process messages serially: it takes messages off - // the queue and immediately dispatches them to handlers, which may complete - // out of order. The proper fix is to replace our custom notification with a - // custom request where we wait for the response before proceeding. However, - // rmcp does not currently support custom requests, so as a temporary - // workaround we just wait a bit to increase the probability the server has - // processed the notification. Assuming we can upstream rmcp support for - // custom requests, we will remove this once the functionality is available. - tokio::time::sleep(Duration::from_secs(4)).await; - // Call the shell tool and verify that an elicitation was created and // auto-approved. let CallToolResult { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index cd92cd08c4..6a87f379c6 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -26,13 +26,17 @@ use mcp_types::RequestId; use reqwest::header::HeaderMap; use rmcp::model::CallToolRequestParam; use rmcp::model::ClientNotification; +use rmcp::model::ClientRequest; use rmcp::model::CreateElicitationRequestParam; use rmcp::model::CreateElicitationResult; use rmcp::model::CustomNotification; +use rmcp::model::CustomRequest; +use rmcp::model::CustomResult; use rmcp::model::Extensions; use rmcp::model::InitializeRequestParam; use rmcp::model::PaginatedRequestParam; use rmcp::model::ReadResourceRequestParam; +use rmcp::model::ServerResult; use rmcp::service::RoleClient; use rmcp::service::RunningService; use rmcp::service::{self}; @@ -370,7 +374,6 @@ impl RmcpClient { params: Option, ) -> Result<()> { let service: Arc> = self.service().await?; - service.service(); service .send_notification(ClientNotification::CustomNotification(CustomNotification { method: method.to_string(), @@ -381,6 +384,24 @@ impl RmcpClient { Ok(()) } + pub async fn send_custom_request( + &self, + method: &str, + params: Option, + ) -> Result { + let service: Arc> = self.service().await?; + let response = service + .send_request(ClientRequest::CustomRequest(CustomRequest::new( + method, params, + ))) + .await?; + + match response { + ServerResult::CustomResult(resp) => Ok(resp), + _ => Err(anyhow!("custom request error: {response:?}")), + } + } + async fn service(&self) -> Result>> { let guard = self.state.lock().await; match &*guard {