mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: declare server capability in shell-tool-mcp (#7112)
This introduces a new feature to Codex when it operates as an MCP
_client_ where if an MCP _server_ replies that it has an entry named
`"codex/sandbox-state"` in its _server capabilities_, then Codex will
send it an MCP notification with the following structure:
```json
{
"method": "codex/sandbox-state/update",
"params": {
"sandboxPolicy": {
"type": "workspace-write",
"network-access": false,
"exclude-tmpdir-env-var": false
"exclude-slash-tmp": false
},
"codexLinuxSandboxExe": null,
"sandboxCwd": "/Users/mbolin/code/codex2"
}
}
```
or with whatever values are appropriate for the initial `sandboxPolicy`.
**NOTE:** Codex _should_ continue to send the MCP server notifications
of the same format if these things change over the lifetime of the
thread, but that isn't wired up yet.
The result is that `shell-tool-mcp` can consume these values so that
when it calls `codex_core::exec::process_exec_tool_call()` in
`codex-rs/exec-server/src/posix/escalate_server.rs`, it is now sure to
call it with the correct values (whereas previously we relied on
hardcoded values).
While I would argue this is a supported use case within the MCP
protocol, the `rmcp` crate that we are using today does not support
custom notifications. As such, I had to patch it and I submitted it for
review, so hopefully it will be accepted in some form:
https://github.com/modelcontextprotocol/rust-sdk/pull/556
To test out this change from end-to-end:
- I ran `cargo build` in `~/code/codex2/codex-rs/exec-server`
- I built the fork of Bash in `~/code/bash/bash`
- I added the following to my `~/.codex/config.toml`:
```toml
# Use with `codex --disable shell_tool`.
[mcp_servers.execshell]
args = ["--bash", "/Users/mbolin/code/bash/bash"]
command = "/Users/mbolin/code/codex2/codex-rs/target/debug/codex-exec-mcp-server"
```
- From `~/code/codex2/codex-rs`, I ran `just codex --disable shell_tool`
- When the TUI started up, I verified that the sandbox mode is
`workspace-write`
- I ran `/mcp` to verify that the shell tool from the MCP is there:
<img width="1387" height="1400" alt="image"
src="https://github.com/user-attachments/assets/1a8addcc-5005-4e16-b59f-95cfd06fd4ab"
/>
- Then I asked it:
> what is the output of `gh issue list`
because this should be auto-approved with our existing dummy policy:
https://github.com/openai/codex/blob/af63e6eccc35783f1bf4dca3c61adb090efb6b8a/codex-rs/exec-server/src/posix.rs#L157-L164
And it worked:
<img width="1387" height="1400" alt="image"
src="https://github.com/user-attachments/assets/7568d2f7-80da-4d68-86d0-c265a6f5e6c1"
/>
This commit is contained in:
committed by
GitHub
Unverified
parent
af63e6eccc
commit
c6f68c9df8
Generated
+5
-6
@@ -5097,10 +5097,10 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rmcp"
|
||||
version = "0.8.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca"
|
||||
version = "0.9.0"
|
||||
source = "git+https://github.com/bolinfest/rust-sdk?branch=pr556#4d9cc16f4c76c84486344f542ed9a3e9364019ba"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64",
|
||||
"bytes",
|
||||
"chrono",
|
||||
@@ -5131,9 +5131,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rmcp-macros"
|
||||
version = "0.8.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "01263441d3f8635c628e33856c468b96ebbce1af2d3699ea712ca71432d4ee7a"
|
||||
version = "0.9.0"
|
||||
source = "git+https://github.com/bolinfest/rust-sdk?branch=pr556#4d9cc16f4c76c84486344f542ed9a3e9364019ba"
|
||||
dependencies = [
|
||||
"darling 0.21.3",
|
||||
"proc-macro2",
|
||||
|
||||
+5
-8
@@ -108,8 +108,8 @@ async-trait = "0.1.89"
|
||||
axum = { version = "0.8", default-features = false }
|
||||
base64 = "0.22.1"
|
||||
bytes = "1.10.1"
|
||||
chrono = "0.4.42"
|
||||
chardetng = "0.1.17"
|
||||
chrono = "0.4.42"
|
||||
clap = "4"
|
||||
clap_complete = "4"
|
||||
color-eyre = "0.6.3"
|
||||
@@ -120,9 +120,9 @@ diffy = "0.4.2"
|
||||
dirs = "6"
|
||||
dotenvy = "0.15.7"
|
||||
dunce = "1.0.4"
|
||||
encoding_rs = "0.8.35"
|
||||
env-flags = "0.1.1"
|
||||
env_logger = "0.11.5"
|
||||
encoding_rs = "0.8.35"
|
||||
escargot = "0.5"
|
||||
eventsource-stream = "0.2.3"
|
||||
futures = { version = "0.3", default-features = false }
|
||||
@@ -167,7 +167,7 @@ ratatui-macros = "0.6.0"
|
||||
regex-lite = "0.1.7"
|
||||
regex = "1.11.1"
|
||||
reqwest = "0.12"
|
||||
rmcp = { version = "0.8.5", default-features = false }
|
||||
rmcp = { version = "0.9.0", default-features = false }
|
||||
schemars = "0.8.22"
|
||||
seccompiler = "0.5.0"
|
||||
serde = "1"
|
||||
@@ -261,11 +261,7 @@ unwrap_used = "deny"
|
||||
# cargo-shear cannot see the platform-specific openssl-sys usage, so we
|
||||
# silence the false positive here instead of deleting a real dependency.
|
||||
[workspace.metadata.cargo-shear]
|
||||
ignored = [
|
||||
"icu_provider",
|
||||
"openssl-sys",
|
||||
"codex-utils-readiness",
|
||||
]
|
||||
ignored = ["icu_provider", "openssl-sys", "codex-utils-readiness"]
|
||||
|
||||
[profile.release]
|
||||
lto = "fat"
|
||||
@@ -286,6 +282,7 @@ opt-level = 0
|
||||
# ratatui = { path = "../../ratatui" }
|
||||
crossterm = { git = "https://github.com/nornagon/crossterm", branch = "nornagon/color-query" }
|
||||
ratatui = { git = "https://github.com/nornagon/ratatui", branch = "nornagon-v0.29.0-patch" }
|
||||
rmcp = { git = "https://github.com/bolinfest/rust-sdk", branch = "pr556" }
|
||||
|
||||
# Uncomment to debug local changes.
|
||||
# rmcp = { path = "../../rust-sdk/crates/rmcp" }
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
|
||||
use crate::AuthManager;
|
||||
use crate::SandboxState;
|
||||
use crate::client_common::REVIEW_PROMPT;
|
||||
use crate::compact;
|
||||
use crate::compact::run_inline_auto_compact_task;
|
||||
@@ -614,6 +615,22 @@ impl Session {
|
||||
)
|
||||
.await;
|
||||
|
||||
let sandbox_state = SandboxState {
|
||||
sandbox_policy: session_configuration.sandbox_policy.clone(),
|
||||
codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(),
|
||||
sandbox_cwd: session_configuration.cwd.clone(),
|
||||
};
|
||||
if let Err(e) = sess
|
||||
.services
|
||||
.mcp_connection_manager
|
||||
.read()
|
||||
.await
|
||||
.notify_sandbox_state_change(&sandbox_state)
|
||||
.await
|
||||
{
|
||||
tracing::error!("Failed to notify sandbox state change: {e}");
|
||||
}
|
||||
|
||||
// record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted.
|
||||
sess.record_initial_history(initial_history).await;
|
||||
|
||||
|
||||
@@ -32,6 +32,9 @@ pub mod git_info;
|
||||
pub mod landlock;
|
||||
pub mod mcp;
|
||||
mod mcp_connection_manager;
|
||||
pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY;
|
||||
pub use mcp_connection_manager::MCP_SANDBOX_STATE_NOTIFICATION;
|
||||
pub use mcp_connection_manager::SandboxState;
|
||||
mod mcp_tool_call;
|
||||
mod message_history;
|
||||
mod model_provider_info;
|
||||
|
||||
@@ -10,6 +10,7 @@ use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::env;
|
||||
use std::ffi::OsString;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
@@ -28,6 +29,7 @@ use codex_protocol::protocol::McpStartupCompleteEvent;
|
||||
use codex_protocol::protocol::McpStartupFailure;
|
||||
use codex_protocol::protocol::McpStartupStatus;
|
||||
use codex_protocol::protocol::McpStartupUpdateEvent;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_rmcp_client::ElicitationResponse;
|
||||
use codex_rmcp_client::OAuthCredentialsStoreMode;
|
||||
use codex_rmcp_client::RmcpClient;
|
||||
@@ -48,6 +50,8 @@ use mcp_types::Resource;
|
||||
use mcp_types::ResourceTemplate;
|
||||
use mcp_types::Tool;
|
||||
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::json;
|
||||
use sha1::Digest;
|
||||
use sha1::Sha1;
|
||||
@@ -174,6 +178,7 @@ struct ManagedClient {
|
||||
tools: Vec<ToolInfo>,
|
||||
tool_filter: ToolFilter,
|
||||
tool_timeout: Option<Duration>,
|
||||
server_supports_sandbox_state_capability: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -222,6 +227,35 @@ impl AsyncManagedClient {
|
||||
async fn client(&self) -> Result<ManagedClient, StartupOutcomeError> {
|
||||
self.client.clone().await
|
||||
}
|
||||
|
||||
async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> {
|
||||
let managed = self.client().await?;
|
||||
if !managed.server_supports_sandbox_state_capability {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
managed
|
||||
.client
|
||||
.send_custom_notification(
|
||||
MCP_SANDBOX_STATE_NOTIFICATION,
|
||||
Some(serde_json::to_value(sandbox_state)?),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub const MCP_SANDBOX_STATE_CAPABILITY: &str = "codex/sandbox-state";
|
||||
|
||||
/// Custom MCP notification for sandbox state updates.
|
||||
/// When used, the `params` field of the notification is [`SandboxState`].
|
||||
pub const MCP_SANDBOX_STATE_NOTIFICATION: &str = "codex/sandbox-state/update";
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SandboxState {
|
||||
pub sandbox_policy: SandboxPolicy,
|
||||
pub codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
pub sandbox_cwd: PathBuf,
|
||||
}
|
||||
|
||||
/// A thin wrapper around a set of running [`RmcpClient`] instances.
|
||||
@@ -567,6 +601,34 @@ impl McpConnectionManager {
|
||||
.get(tool_name)
|
||||
.map(|tool| (tool.server_name.clone(), tool.tool_name.clone()))
|
||||
}
|
||||
|
||||
pub async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> {
|
||||
let mut join_set = JoinSet::new();
|
||||
|
||||
for async_managed_client in self.clients.values() {
|
||||
let sandbox_state = sandbox_state.clone();
|
||||
let async_managed_client = async_managed_client.clone();
|
||||
join_set.spawn(async move {
|
||||
async_managed_client
|
||||
.notify_sandbox_state_change(&sandbox_state)
|
||||
.await
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(join_res) = join_set.join_next().await {
|
||||
match join_res {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(err)) => {
|
||||
warn!("Failed to notify sandbox state change to MCP server: {err:#}");
|
||||
}
|
||||
Err(err) => {
|
||||
warn!("Task panic when notifying sandbox state change to MCP server: {err:#}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn emit_update(
|
||||
@@ -700,7 +762,7 @@ async fn start_server_task(
|
||||
|
||||
let send_elicitation = elicitation_requests.make_sender(server_name.clone(), tx_event);
|
||||
|
||||
client
|
||||
let initialize_result = client
|
||||
.initialize(params, startup_timeout, send_elicitation)
|
||||
.await
|
||||
.map_err(StartupOutcomeError::from)?;
|
||||
@@ -709,11 +771,19 @@ async fn start_server_task(
|
||||
.await
|
||||
.map_err(StartupOutcomeError::from)?;
|
||||
|
||||
let server_supports_sandbox_state_capability = initialize_result
|
||||
.capabilities
|
||||
.experimental
|
||||
.as_ref()
|
||||
.and_then(|exp| exp.get(MCP_SANDBOX_STATE_CAPABILITY))
|
||||
.is_some();
|
||||
|
||||
let managed = ManagedClient {
|
||||
client: Arc::clone(&client),
|
||||
tools,
|
||||
tool_timeout: Some(tool_timeout),
|
||||
tool_filter,
|
||||
server_supports_sandbox_state_capability,
|
||||
};
|
||||
|
||||
Ok(managed)
|
||||
|
||||
@@ -8,8 +8,8 @@ use std::time::Duration;
|
||||
use anyhow::Context as _;
|
||||
use path_absolutize::Absolutize as _;
|
||||
|
||||
use codex_core::SandboxState;
|
||||
use codex_core::exec::process_exec_tool_call;
|
||||
use codex_core::protocol::SandboxPolicy;
|
||||
use tokio::process::Command;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
@@ -48,6 +48,7 @@ impl EscalateServer {
|
||||
&self,
|
||||
params: ExecParams,
|
||||
cancel_rx: CancellationToken,
|
||||
sandbox_state: &SandboxState,
|
||||
) -> anyhow::Result<ExecResult> {
|
||||
let (escalate_server, escalate_client) = AsyncDatagramSocket::pair()?;
|
||||
let client_socket = escalate_client.into_inner();
|
||||
@@ -64,12 +65,6 @@ impl EscalateServer {
|
||||
self.execve_wrapper.to_string_lossy().to_string(),
|
||||
);
|
||||
|
||||
// TODO: use the sandbox policy and cwd from the calling client.
|
||||
// Note that sandbox_cwd is ignored for ReadOnly, but needs to be legit
|
||||
// for `SandboxPolicy::WorkspaceWrite`.
|
||||
let sandbox_policy = SandboxPolicy::ReadOnly;
|
||||
let sandbox_cwd = PathBuf::from("/__NONEXISTENT__");
|
||||
|
||||
let ExecParams {
|
||||
command,
|
||||
workdir,
|
||||
@@ -94,9 +89,9 @@ impl EscalateServer {
|
||||
justification: None,
|
||||
arg0: None,
|
||||
},
|
||||
&sandbox_policy,
|
||||
&sandbox_cwd,
|
||||
&None,
|
||||
&sandbox_state.sandbox_policy,
|
||||
&sandbox_state.sandbox_cwd,
|
||||
&sandbox_state.codex_linux_sandbox_exe,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
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::SandboxState;
|
||||
use codex_core::protocol::SandboxPolicy;
|
||||
use rmcp::ErrorData as McpError;
|
||||
use rmcp::RoleServer;
|
||||
use rmcp::ServerHandler;
|
||||
@@ -17,6 +22,8 @@ use rmcp::tool;
|
||||
use rmcp::tool_handler;
|
||||
use rmcp::tool_router;
|
||||
use rmcp::transport::stdio;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::posix::escalate_server::EscalateServer;
|
||||
use crate::posix::escalate_server::{self};
|
||||
@@ -27,6 +34,8 @@ use crate::posix::stopwatch::Stopwatch;
|
||||
/// Path to our patched bash.
|
||||
const CODEX_BASH_PATH_ENV_VAR: &str = "CODEX_BASH_PATH";
|
||||
|
||||
const SANDBOX_STATE_CAPABILITY_VERSION: &str = "1.0.0";
|
||||
|
||||
pub(crate) fn get_bash_path() -> Result<PathBuf> {
|
||||
std::env::var(CODEX_BASH_PATH_ENV_VAR)
|
||||
.map(PathBuf::from)
|
||||
@@ -70,6 +79,7 @@ pub struct ExecTool {
|
||||
bash_path: PathBuf,
|
||||
execve_wrapper: PathBuf,
|
||||
policy: ExecPolicy,
|
||||
sandbox_state: Arc<RwLock<Option<SandboxState>>>,
|
||||
}
|
||||
|
||||
#[tool_router]
|
||||
@@ -80,6 +90,7 @@ impl ExecTool {
|
||||
bash_path,
|
||||
execve_wrapper,
|
||||
policy,
|
||||
sandbox_state: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,13 +108,24 @@ impl ExecTool {
|
||||
);
|
||||
let stopwatch = Stopwatch::new(effective_timeout);
|
||||
let cancel_token = stopwatch.cancellation_token();
|
||||
let sandbox_state =
|
||||
self.sandbox_state
|
||||
.read()
|
||||
.await
|
||||
.clone()
|
||||
.unwrap_or_else(|| SandboxState {
|
||||
sandbox_policy: SandboxPolicy::ReadOnly,
|
||||
codex_linux_sandbox_exe: None,
|
||||
sandbox_cwd: PathBuf::from(¶ms.workdir),
|
||||
});
|
||||
let escalate_server = EscalateServer::new(
|
||||
self.bash_path.clone(),
|
||||
self.execve_wrapper.clone(),
|
||||
McpEscalationPolicy::new(self.policy, context, stopwatch.clone()),
|
||||
);
|
||||
|
||||
let result = escalate_server
|
||||
.exec(params, cancel_token)
|
||||
.exec(params, cancel_token, &sandbox_state)
|
||||
.await
|
||||
.map_err(|e| McpError::internal_error(e.to_string(), None))?;
|
||||
Ok(CallToolResult::success(vec![Content::json(
|
||||
@@ -115,9 +137,22 @@ impl ExecTool {
|
||||
#[tool_handler]
|
||||
impl ServerHandler for ExecTool {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
let mut experimental_capabilities = ExperimentalCapabilities::new();
|
||||
let mut sandbox_state_capability = JsonObject::new();
|
||||
sandbox_state_capability.insert(
|
||||
"version".to_string(),
|
||||
serde_json::Value::String(SANDBOX_STATE_CAPABILITY_VERSION.to_string()),
|
||||
);
|
||||
experimental_capabilities.insert(
|
||||
MCP_SANDBOX_STATE_CAPABILITY.to_string(),
|
||||
sandbox_state_capability,
|
||||
);
|
||||
ServerInfo {
|
||||
protocol_version: ProtocolVersion::V_2025_06_18,
|
||||
capabilities: ServerCapabilities::builder().enable_tools().build(),
|
||||
capabilities: ServerCapabilities::builder()
|
||||
.enable_tools()
|
||||
.enable_experimental_with(experimental_capabilities)
|
||||
.build(),
|
||||
server_info: Implementation::from_build_env(),
|
||||
instructions: Some(
|
||||
"This server provides a tool to execute shell commands and return their output."
|
||||
@@ -133,6 +168,31 @@ impl ServerHandler for ExecTool {
|
||||
) -> Result<InitializeResult, McpError> {
|
||||
Ok(self.get_info())
|
||||
}
|
||||
|
||||
async fn on_custom_notification(
|
||||
&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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn serve(
|
||||
|
||||
@@ -25,8 +25,11 @@ use mcp_types::ReadResourceResult;
|
||||
use mcp_types::RequestId;
|
||||
use reqwest::header::HeaderMap;
|
||||
use rmcp::model::CallToolRequestParam;
|
||||
use rmcp::model::ClientNotification;
|
||||
use rmcp::model::CreateElicitationRequestParam;
|
||||
use rmcp::model::CreateElicitationResult;
|
||||
use rmcp::model::CustomClientNotification;
|
||||
use rmcp::model::Extensions;
|
||||
use rmcp::model::InitializeRequestParam;
|
||||
use rmcp::model::PaginatedRequestParam;
|
||||
use rmcp::model::ReadResourceRequestParam;
|
||||
@@ -361,6 +364,25 @@ impl RmcpClient {
|
||||
Ok(converted)
|
||||
}
|
||||
|
||||
pub async fn send_custom_notification(
|
||||
&self,
|
||||
method: &str,
|
||||
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(),
|
||||
},
|
||||
))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn service(&self) -> Result<Arc<RunningService<RoleClient, LoggingClientHandler>>> {
|
||||
let guard = self.state.lock().await;
|
||||
match &*guard {
|
||||
|
||||
Reference in New Issue
Block a user