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

This commit is contained in:
Michael Bolin
2025-12-16 14:39:00 -08:00
parent a3b137d093
commit 41a0c52ca1
8 changed files with 129 additions and 77 deletions
+42 -12
View File
@@ -2204,6 +2204,16 @@ dependencies = [
"darling_macro 0.21.3",
]
[[package]]
name = "darling"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d"
dependencies = [
"darling_core 0.23.0",
"darling_macro 0.23.0",
]
[[package]]
name = "darling_core"
version = "0.20.11"
@@ -2232,6 +2242,19 @@ dependencies = [
"syn 2.0.104",
]
[[package]]
name = "darling_core"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0"
dependencies = [
"ident_case",
"proc-macro2",
"quote",
"strsim 0.11.1",
"syn 2.0.104",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
@@ -2254,6 +2277,17 @@ dependencies = [
"syn 2.0.104",
]
[[package]]
name = "darling_macro"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d"
dependencies = [
"darling_core 0.23.0",
"quote",
"syn 2.0.104",
]
[[package]]
name = "dbus"
version = "0.9.9"
@@ -3365,7 +3399,7 @@ dependencies = [
"libc",
"percent-encoding",
"pin-project-lite",
"socket2 0.5.10",
"socket2 0.6.1",
"system-configuration",
"tokio",
"tower-service",
@@ -5051,9 +5085,9 @@ dependencies = [
[[package]]
name = "process-wrap"
version = "8.2.1"
version = "9.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a3ef4f2f0422f23a82ec9f628ea2acd12871c81a9362b02c43c1aa86acfc3ba1"
checksum = "5e5fd83ab7fa55fd06f5e665e3fc52b8bca451c0486b8ea60ad649cd1c10a5da"
dependencies = [
"futures",
"indexmap 2.12.0",
@@ -5151,7 +5185,7 @@ dependencies = [
"quinn-udp",
"rustc-hash",
"rustls",
"socket2 0.5.10",
"socket2 0.6.1",
"thiserror 2.0.17",
"tokio",
"tracing",
@@ -5188,7 +5222,7 @@ dependencies = [
"cfg_aliases 0.2.1",
"libc",
"once_cell",
"socket2 0.5.10",
"socket2 0.6.1",
"tracing",
"windows-sys 0.60.2",
]
@@ -5463,9 +5497,7 @@ dependencies = [
[[package]]
name = "rmcp"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "38b18323edc657390a6ed4d7a9110b0dec2dc3ed128eb2a123edfbafabdbddc5"
version = "0.11.0"
dependencies = [
"async-trait",
"base64",
@@ -5498,11 +5530,9 @@ dependencies = [
[[package]]
name = "rmcp-macros"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c75d0a62676bf8c8003c4e3c348e2ceb6a7b3e48323681aaf177fdccdac2ce50"
version = "0.11.0"
dependencies = [
"darling 0.21.3",
"darling 0.23.0",
"proc-macro2",
"quote",
"serde_json",
+3 -3
View File
@@ -162,7 +162,6 @@ opentelemetry-appender-tracing = "0.30.0"
opentelemetry-otlp = "0.30.0"
opentelemetry-semantic-conventions = "0.30.0"
opentelemetry_sdk = "0.30.0"
tracing-opentelemetry = "0.31.0"
os_info = "3.12.0"
owo-colors = "4.2.0"
path-absolutize = "3.1.1"
@@ -177,7 +176,7 @@ ratatui-macros = "0.6.0"
regex = "1.12.2"
regex-lite = "0.1.7"
reqwest = "0.12"
rmcp = { version = "0.10.0", default-features = false }
rmcp = { version = "0.11.0", default-features = false }
schemars = "0.8.22"
seccompiler = "0.5.0"
sentry = "0.46.0"
@@ -211,6 +210,7 @@ toml_edit = "0.23.5"
tonic = "0.13.1"
tracing = "0.1.43"
tracing-appender = "0.2.3"
tracing-opentelemetry = "0.31.0"
tracing-subscriber = "0.3.20"
tracing-test = "0.2.5"
tree-sitter = "0.25.10"
@@ -296,4 +296,4 @@ crossterm = { git = "https://github.com/nornagon/crossterm", branch = "nornagon/
ratatui = { git = "https://github.com/nornagon/ratatui", branch = "nornagon-v0.29.0-patch" }
# Uncomment to debug local changes.
# rmcp = { path = "../../rust-sdk/crates/rmcp" }
rmcp = { path = "../../rust-sdk/crates/rmcp" }
+1 -1
View File
@@ -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;
+9 -6
View File
@@ -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")]
+36 -23
View File
@@ -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::CustomClientNotification,
_context: rmcp::service::NotificationContext<rmcp::RoleServer>,
) {
let rmcp::model::CustomClientNotification { 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!({ "status": "ok" })))
}
}
+10 -10
View File
@@ -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::CustomClientNotification;
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<P, S>(
@@ -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<S>(
async fn send_sandbox_state_update<S>(
sandbox_state: SandboxState,
service: &RunningService<RoleClient, S>,
) -> anyhow::Result<()>
where
S: Service<RoleClient> + ClientHandler,
{
let sandbox_state_notification = CustomClientNotification::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(())
}
@@ -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 {
+28 -9
View File
@@ -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::CustomClientNotification;
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,19 +374,34 @@ impl RmcpClient {
params: Option<serde_json::Value>,
) -> Result<()> {
let service: Arc<RunningService<RoleClient, LoggingClientHandler>> = self.service().await?;
service.service();
service
.send_notification(ClientNotification::CustomClientNotification(
CustomClientNotification {
method: method.to_string(),
params,
extensions: Extensions::new(),
},
))
.send_notification(ClientNotification::CustomNotification(CustomNotification {
method: method.to_string(),
params,
extensions: Extensions::new(),
}))
.await?;
Ok(())
}
pub async fn send_custom_request(
&self,
method: &str,
params: Option<serde_json::Value>,
) -> Result<CustomResult> {
let service: Arc<RunningService<RoleClient, LoggingClientHandler>> = 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<Arc<RunningService<RoleClient, LoggingClientHandler>>> {
let guard = self.state.lock().await;
match &*guard {