fix: change codex/sandbox-state/update from a notification to a request

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;

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")]

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<rmcp::RoleServer>,
-    ) {
-        let rmcp::model::CustomNotification { method, params, .. } = notification;
-        if method == MCP_SANDBOX_STATE_NOTIFICATION
-            && let Some(params) = params
-        {
-            match serde_json::from_value::<SandboxState>(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<rmcp::RoleServer>,
+    ) -> Result<CustomResult, McpError> {
+        let CustomRequest { method, params, .. } = request;
+        if method != MCP_SANDBOX_STATE_METHOD {
+            return Err(McpError::method_not_found::<CodexSandboxStateUpdateMethod>());
         }
+
+        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::<SandboxState>(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!({})))
     }
 }
 

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<P, S>(
     sandbox_cwd: P,
     codex_linux_sandbox_exe: Option<PathBuf>,
     service: &RunningService<RoleClient, S>,
-) -> anyhow::Result<()>
+) -> anyhow::Result<ServerResult>
 where
     P: AsRef<Path>,
     S: Service<RoleClient> + 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<P, S>(
     writable_folder: P,
     codex_linux_sandbox_exe: Option<PathBuf>,
     service: &RunningService<RoleClient, S>,
-) -> anyhow::Result<()>
+) -> anyhow::Result<ServerResult>
 where
     P: AsRef<Path>,
     S: Service<RoleClient> + 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<S>(
+async fn send_sandbox_state_update<S>(
     sandbox_state: SandboxState,
     service: &RunningService<RoleClient, S>,
-) -> anyhow::Result<()>
+) -> anyhow::Result<ServerResult>
 where
     S: Service<RoleClient> + 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 {

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.

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<serde_json::Value>,
     ) -> Result<()> {
         let service: Arc<RunningService<RoleClient, LoggingClientHandler>> = 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<serde_json::Value>,
+    ) -> Result<ServerResult> {
+        let service: Arc<RunningService<RoleClient, LoggingClientHandler>> = self.service().await?;
+        let response = service
+            .send_request(ClientRequest::CustomRequest(CustomRequest::new(
+                method, params,
+            )))
+            .await?;
+        Ok(response)
+    }
+
     async fn service(&self) -> Result<Arc<RunningService<RoleClient, LoggingClientHandler>>> {
         let guard = self.state.lock().await;
         match &*guard {

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