diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f78c19328f0..4eeb1746bc9 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 3213b22b71a..6c0b48b1bd5 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 3fec7e4dd95..620d332e71e 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!({}))) } } diff --git a/codex-rs/exec-server/tests/common/lib.rs b/codex-rs/exec-server/tests/common/lib.rs index 99587a2ad5e..c2202a168a8 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,10 +7,12 @@ 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::model::ServerResult; use rmcp::service::RunningService; use rmcp::transport::ConfigureCommandExt; use rmcp::transport::TokioChildProcess; @@ -82,7 +84,7 @@ pub async fn notify_readable_sandbox( sandbox_cwd: P, codex_linux_sandbox_exe: Option, service: &RunningService, -) -> anyhow::Result<()> +) -> anyhow::Result where P: AsRef, S: Service + ClientHandler, @@ -92,14 +94,14 @@ 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( writable_folder: P, codex_linux_sandbox_exe: Option, service: &RunningService, -) -> anyhow::Result<()> +) -> anyhow::Result where P: AsRef, S: Service + ClientHandler, @@ -119,24 +121,23 @@ 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<()> +) -> 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()) + let response = service + .send_request(ClientRequest::CustomRequest(CustomRequest::new( + MCP_SANDBOX_STATE_METHOD, + Some(serde_json::to_value(sandbox_state)?), + ))) .await?; - Ok(()) + Ok(response) } pub struct InteractiveClient { diff --git a/codex-rs/exec-server/tests/suite/accept_elicitation.rs b/codex-rs/exec-server/tests/suite/accept_elicitation.rs index b703eaf4a70..81283a91d53 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; @@ -19,6 +18,8 @@ use rmcp::ServiceExt; use rmcp::model::CallToolRequestParam; use rmcp::model::CallToolResult; use rmcp::model::CreateElicitationRequestParam; +use rmcp::model::EmptyResult; +use rmcp::model::ServerResult; use rmcp::model::object; use serde_json::json; use std::os::unix::fs::PermissionsExt; @@ -82,19 +83,11 @@ prefix_rule( } else { None }; - 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; + let response = + notify_readable_sandbox(&project_root_path, codex_linux_sandbox_exe, &service).await?; + let ServerResult::EmptyResult(EmptyResult {}) = response else { + panic!("expected EmptyResult from sandbox state notification but found: {response:?}"); + }; // Call the shell tool and verify that an elicitation was created and // auto-approved. diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index cd92cd08c40..b977389eab0 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -26,13 +26,16 @@ 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::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 +373,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 +383,20 @@ 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?; + Ok(response) + } + async fn service(&self) -> Result>> { let guard = self.state.lock().await; match &*guard { diff --git a/shell-tool-mcp/README.md b/shell-tool-mcp/README.md index 16a8492656e..ccfd0bcfbad 100644 --- a/shell-tool-mcp/README.md +++ b/shell-tool-mcp/README.md @@ -65,10 +65,11 @@ This MCP server is designed to be used with [Codex](https://developers.openai.co } ``` -This capability means the MCP server honors notifications like the following to update the sandbox policy the MCP server uses when spawning Bash: +This capability means the MCP server honors requests like the following to update the sandbox policy the MCP server uses when spawning Bash: ```json { + "id": "req-42", "method": "codex/sandbox-state/update", "params": { "sandboxPolicy": { @@ -82,7 +83,16 @@ This capability means the MCP server honors notifications like the following to } ``` -The Codex harness (used by the CLI and the VS Code extension) sends such notifications to MCP servers that declare the `codex/sandbox-state` capability. +Once the server has processed the update, it sends an empty response to acknowledge the request: + +```json +{ + "id": "req-42", + "result": {} +} +``` + +The Codex harness (used by the CLI and the VS Code extension) sends such requests to MCP servers that declare the `codex/sandbox-state` capability. ## Package Contents