fix: change codex/sandbox-state/update from a notification to a request (#8142)

Historically, `accept_elicitation_for_prompt_rule()` was flaky because
we were using a notification to update the sandbox followed by a `shell`
tool request that we expected to be subject to the new sandbox config,
but because [rmcp](https://crates.io/crates/rmcp) MCP servers delegate
each incoming message to a new Tokio task, messages are not guaranteed
to be processed in order, so sometimes the `shell` tool call would run
before the notification was processed.

Prior to this PR, we relied on a generous `sleep()` between the
notification and the request to reduce the change of the test flaking
out.

This PR implements a proper fix, which is to use a _request_ instead of
a notification for the sandbox update so that we can wait for the
response to the sandbox request before sending the request to the
`shell` tool call. Previously, `rmcp` did not support custom requests,
but I fixed that in
https://github.com/modelcontextprotocol/rust-sdk/pull/590, which made it
into the `0.12.0` release (see #8288).

This PR updates `shell-tool-mcp` to expect
`"codex/sandbox-state/update"` as a _request_ instead of a notification
and sends the appropriate ack. Note this behavior is tied to our custom
`codex/sandbox-state` capability, which Codex honors as an MCP client,
which is why `core/src/mcp_connection_manager.rs` had to be updated as
part of this PR, as well.

This PR also updates the docs at `shell-tool-mcp/README.md`.
This commit is contained in:
Michael Bolin
2025-12-18 15:32:01 -08:00
committed by GitHub
parent 358a5baba0
commit 46baedd7cb
7 changed files with 98 additions and 62 deletions

View File

@@ -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 {