openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Batch sandboxed capability root discovery (#40443)

Browse Source 
## What changed

- Run capability discovery for roots with the same sandbox context in one sandbox helper invocation.
- Keep roots with different sandbox permissions isolated, and fall back to per-root discovery if the batched invocation fails.

## Testing

- Add coverage for single-launch batching, fallback after a sandbox launch failure, and roots with different permission profiles.

GitOrigin-RevId: a6e185a8edfd028fc3639b7c6c3cc7da342b6fe3
This commit is contained in:
jif
2026-08-24 17:39:16 +00:00
committed by copyberry
parent a12c4f4c01
commit eb15245d82
7 changed files with 373 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/fs_helper.rs
+26
View File
@@ -5,6 +5,8 @@ use serde::Deserialize;
use serde::Serialize;
use tokio::io;
use crate::CapabilityRootsDiscoverParams;
use crate::CapabilityRootsDiscoverResponse;
use crate::CopyOptions;
use crate::CreateDirectoryOptions;
use crate::ExecutorFileSystem;
@@ -13,6 +15,7 @@ use crate::ReadFileOptions;
use crate::RemoveOptions;
use crate::WriteFileOptions;
use crate::local_file_system::DirectFileSystem;
use crate::protocol::CAPABILITY_ROOTS_DISCOVER_METHOD;
use crate::protocol::FS_CANONICALIZE_METHOD;
use crate::protocol::FS_COPY_METHOD;
use crate::protocol::FS_CREATE_DIRECTORY_METHOD;
@@ -51,6 +54,8 @@ pub const CODEX_FS_HELPER_ARG1: &str = "--codex-run-as-fs-helper";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "operation", content = "params")]
pub(crate) enum FsHelperRequest {
#[serde(rename = "capabilityRoots/discoverV1")]
DiscoverCapabilityRoots(CapabilityRootsDiscoverParams),
#[serde(rename = "fs/open")]
Open(FsReadFileParams),
#[serde(rename = "fs/readFile")]
@@ -94,6 +99,8 @@ pub(crate) struct FsHelperOpenResponse {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "operation", content = "response")]
pub(crate) enum FsHelperPayload {
#[serde(rename = "capabilityRoots/discoverV1")]
DiscoverCapabilityRoots(CapabilityRootsDiscoverResponse),
#[serde(rename = "fs/open")]
Open(FsHelperOpenResponse),
#[serde(rename = "fs/readFile")]
@@ -119,6 +126,7 @@ pub(crate) enum FsHelperPayload {
impl FsHelperPayload {
fn operation(&self) -> &'static str {
match self {
Self::DiscoverCapabilityRoots(_) => CAPABILITY_ROOTS_DISCOVER_METHOD,
Self::Open(_) => FS_OPEN_METHOD,
Self::ReadFile(_) => FS_READ_FILE_METHOD,
Self::WriteFile(_) => FS_WRITE_FILE_METHOD,
@@ -132,6 +140,18 @@ impl FsHelperPayload {
}
}
pub(crate) fn expect_capability_roots_discover(
self,
) -> Result<CapabilityRootsDiscoverResponse, JSONRPCErrorError> {
match self {
Self::DiscoverCapabilityRoots(response) => Ok(response),
other => Err(unexpected_response(
CAPABILITY_ROOTS_DISCOVER_METHOD,
other.operation(),
)),
}
}
pub(crate) fn expect_read_file(self) -> Result<FsReadFileResponse, JSONRPCErrorError> {
match self {
Self::ReadFile(response) => Ok(response),
@@ -223,6 +243,12 @@ pub(crate) async fn run_direct_request(
) -> Result<FsHelperPayload, JSONRPCErrorError> {
let file_system = DirectFileSystem;
match request {
FsHelperRequest::DiscoverCapabilityRoots(params) => {
let response = crate::discover_capability_roots(&file_system, params)
.await
.map_err(|error| invalid_request(error.to_string()))?;
Ok(FsHelperPayload::DiscoverCapabilityRoots(response))
}
FsHelperRequest::Open(_) => Err(invalid_request(
"opening a file requires descriptor handoff".to_string(),
)),
codex-rs/exec-server/src/local_file_system.rs
+1 -1
View File
@@ -83,7 +83,7 @@ impl LocalFileSystem {
}
}
fn sandboxed(&self) -> io::Result<&SandboxedFileSystem> {
pub(crate) fn sandboxed(&self) -> io::Result<&SandboxedFileSystem> {
self.sandboxed.as_ref().ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
codex-rs/exec-server/src/sandboxed_file_system.rs
+18
View File
@@ -5,6 +5,8 @@ use codex_utils_path_uri::PathUri;
use tokio::io;
use tokio_util::io::ReaderStream;
use crate::CapabilityRootsDiscoverParams;
use crate::CapabilityRootsDiscoverResponse;
use crate::CopyOptions;
use crate::CreateDirectoryOptions;
use crate::ExecServerRuntimePaths;
@@ -41,6 +43,22 @@ pub struct SandboxedFileSystem {
}
impl SandboxedFileSystem {
#[tracing::instrument(
name = "capability_roots.discover_v1",
skip_all,
fields(root_count = params.roots.len())
)]
pub(crate) async fn discover_capability_roots(
&self,
params: CapabilityRootsDiscoverParams,
sandbox: &FileSystemSandboxContext,
) -> FileSystemResult<CapabilityRootsDiscoverResponse> {
self.run_sandboxed(sandbox, FsHelperRequest::DiscoverCapabilityRoots(params))
.await?
.expect_capability_roots_discover()
.map_err(map_sandbox_error)
}
pub(crate) async fn open_file_for_read(
&self,
path: &PathUri,
codex-rs/exec-server/src/server/file_system_handler.rs
+37
View File
@@ -3,6 +3,7 @@ use std::io;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_exec_server_protocol::JSONRPCErrorError;
use codex_protocol::config_types::WindowsSandboxLevel;
use crate::CapabilityRootsDiscoverParams;
use crate::CapabilityRootsDiscoverResponse;
@@ -74,6 +75,42 @@ impl FileSystemHandler {
&self,
params: CapabilityRootsDiscoverParams,
) -> Result<CapabilityRootsDiscoverResponse, JSONRPCErrorError> {
let sandbox = params
.roots
.first()
.and_then(|root| root.sandbox.as_ref())
.filter(|sandbox| {
sandbox.should_run_in_sandbox()
&& (!cfg!(target_os = "windows")
|| sandbox.windows_sandbox_level != WindowsSandboxLevel::Disabled)
&& params
.roots
.iter()
.all(|root| root.sandbox.as_ref() == Some(*sandbox))
})
.cloned();
if let Some(sandbox) = sandbox {
let mut batched_params = params.clone();
for root in &mut batched_params.roots {
root.sandbox = None;
}
let result = match self.file_system.sandboxed() {
Ok(file_system) => {
file_system
.discover_capability_roots(batched_params, &sandbox)
.await
}
Err(error) => Err(error),
};
match result {
Ok(response) => return Ok(response),
Err(error) => {
tracing::warn!(%error, "batched capability discovery failed; retrying roots separately");
}
}
}
crate::discover_capability_roots(&self.file_system, params)
.await
.map_err(|error| invalid_request(error.to_string()))
codex-rs/exec-server/tests/capability_discovery.rs
+224 -12
View File
@@ -1,5 +1,10 @@
mod common;
#[cfg(target_os = "linux")]
#[path = "common/fake_bwrap.rs"]
mod fake_bwrap;
#[cfg(target_os = "linux")]
use anyhow::Context as _;
use codex_exec_server::CAPABILITY_ROOTS_DISCOVER_METHOD;
use codex_exec_server::CapabilityRootDiscovery;
use codex_exec_server::CapabilityRootsDiscoverParams;
@@ -24,6 +29,10 @@ use codex_protocol::permissions::NetworkSandboxPolicy;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
use common::exec_server::exec_server;
#[cfg(target_os = "linux")]
use common::exec_server::exec_server_with_env;
#[cfg(target_os = "linux")]
use fake_bwrap::write_fake_bwrap;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -179,6 +188,200 @@ async fn discovers_cursor_plugin_without_reading_default_mcp_for_inline_servers(
Ok(())
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sandboxed_discovery_batches_roots_without_combining_different_permissions()
-> anyhow::Result<()> {
let workspace = tempfile::tempdir()?;
let first_root = workspace.path().join("first");
let second_root = workspace.path().join("second");
write_file(
&first_root.join("skills/first/SKILL.md"),
"---\nname: first\ndescription: First skill.\n---\n",
)?;
write_file(
&second_root.join("skills/second/SKILL.md"),
"---\nname: second\ndescription: Second skill.\n---\n",
)?;
let workspace_uri = PathUri::from_host_native_path(workspace.path())?;
let first_uri = PathUri::from_host_native_path(&first_root)?;
let second_uri = PathUri::from_host_native_path(&second_root)?;
let read_workspace = FileSystemSandboxEntry::new(
AbsolutePathBuf::from_absolute_path(workspace.path())?.into(),
FileSystemAccessMode::Read,
);
let policy = FileSystemSandboxPolicy::restricted(vec![read_workspace]);
let shared_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
PermissionProfile::from_runtime_permissions(&policy, NetworkSandboxPolicy::Restricted),
workspace_uri,
);
#[cfg(target_os = "linux")]
let fake_bwrap_directory = tempfile::tempdir()?;
#[cfg(target_os = "linux")]
let (mut server, fake_bwrap) = {
let fake_bin_dir = fake_bwrap_directory.path().to_path_buf();
let fake_bwrap = write_fake_bwrap(&fake_bin_dir)?;
let mut path_entries = vec![fake_bin_dir];
if let Some(path) = std::env::var_os("PATH") {
path_entries.extend(std::env::split_paths(&path));
}
let helper_path = std::env::join_paths(path_entries)?;
(
exec_server_with_env([("PATH", helper_path.as_os_str())], &[]).await?,
fake_bwrap,
)
};
#[cfg(not(target_os = "linux"))]
let mut server = exec_server().await?;
initialize(&mut server).await?;
let response = discover_roots(
&mut server,
vec![
CapabilityRootDiscoverRequest {
id: "first".to_string(),
path: first_uri.clone(),
sandbox: Some(shared_sandbox.clone()),
},
CapabilityRootDiscoverRequest {
id: "second".to_string(),
path: second_uri.clone(),
sandbox: Some(shared_sandbox.clone()),
},
],
)
.await?;
assert_eq!(
response
.roots
.into_iter()
.map(|root| (
root.id,
root.path,
root.skills
.into_iter()
.map(|skill| skill.instructions.path)
.collect::<Vec<_>>(),
root.error,
))
.collect::<Vec<_>>(),
vec![
(
"first".to_string(),
first_uri.clone(),
vec![first_uri.join("skills/first/SKILL.md")?],
None,
),
(
"second".to_string(),
second_uri.clone(),
vec![second_uri.join("skills/second/SKILL.md")?],
None,
),
]
);
#[cfg(target_os = "linux")]
{
let launch_log = fake_bwrap.with_file_name("bwrap.log");
let launch_count = std::fs::read_to_string(&launch_log)
.with_context(|| format!("expected fake bwrap launch log at {}", launch_log.display()))?
.lines()
.count();
assert_eq!(launch_count, 1);
std::fs::write(fake_bwrap.with_file_name("bwrap.fail-once"), "")?;
let fallback = discover_roots(
&mut server,
vec![
CapabilityRootDiscoverRequest {
id: "fallback-first".to_string(),
path: first_uri.clone(),
sandbox: Some(shared_sandbox.clone()),
},
CapabilityRootDiscoverRequest {
id: "fallback-second".to_string(),
path: second_uri.clone(),
sandbox: Some(shared_sandbox.clone()),
},
],
)
.await?;
assert_eq!(
fallback
.roots
.into_iter()
.map(|root| (root.id, root.skills.len(), root.error))
.collect::<Vec<_>>(),
vec![
("fallback-first".to_string(), 1, None),
("fallback-second".to_string(), 1, None),
]
);
assert!(std::fs::read_to_string(&launch_log)?.lines().count() > 2);
server.shutdown().await?;
server = exec_server().await?;
initialize(&mut server).await?;
}
let read_first_root = FileSystemSandboxEntry::new(
AbsolutePathBuf::from_absolute_path(&first_root)?.into(),
FileSystemAccessMode::Read,
);
let first_policy = FileSystemSandboxPolicy::restricted(vec![read_first_root]);
let first_only_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
PermissionProfile::from_runtime_permissions(
&first_policy,
NetworkSandboxPolicy::Restricted,
),
first_uri.clone(),
);
let read_second_root = FileSystemSandboxEntry::new(
AbsolutePathBuf::from_absolute_path(&second_root)?.into(),
FileSystemAccessMode::Read,
);
let second_policy = FileSystemSandboxPolicy::restricted(vec![read_second_root]);
let second_only_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
PermissionProfile::from_runtime_permissions(
&second_policy,
NetworkSandboxPolicy::Restricted,
),
second_uri.clone(),
);
let response = discover_roots(
&mut server,
vec![
CapabilityRootDiscoverRequest {
id: "first-isolated".to_string(),
path: first_uri,
sandbox: Some(first_only_sandbox),
},
CapabilityRootDiscoverRequest {
id: "second-isolated".to_string(),
path: second_uri,
sandbox: Some(second_only_sandbox),
},
],
)
.await?;
assert_eq!(
response
.roots
.into_iter()
.map(|root| (root.id, root.skills.len(), root.error))
.collect::<Vec<_>>(),
vec![
("first-isolated".to_string(), 1, None),
("second-isolated".to_string(), 1, None),
]
);
server.shutdown().await?;
Ok(())
}
#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sandboxed_discovery_follows_only_permitted_external_symlinks() -> anyhow::Result<()> {
@@ -277,16 +480,29 @@ async fn discover_root_with_sandbox(
path: PathUri,
sandbox: Option<FileSystemSandboxContext>,
) -> anyhow::Result<CapabilityRootDiscovery> {
let response = discover_roots(
server,
vec![CapabilityRootDiscoverRequest {
id: id.to_string(),
path,
sandbox,
}],
)
.await?;
let [discovery] = response.roots.as_slice() else {
anyhow::bail!("expected exactly one discovered root");
};
Ok(discovery.clone())
}
async fn discover_roots(
server: &mut common::exec_server::ExecServerHarness,
roots: Vec<CapabilityRootDiscoverRequest>,
) -> anyhow::Result<CapabilityRootsDiscoverResponse> {
let request_id = server
.send_request(
CAPABILITY_ROOTS_DISCOVER_METHOD,
serde_json::to_value(CapabilityRootsDiscoverParams {
roots: vec![CapabilityRootDiscoverRequest {
id: id.to_string(),
path,
sandbox,
}],
})?,
serde_json::to_value(CapabilityRootsDiscoverParams { roots })?,
)
.await?;
let response = server.next_event().await?;
@@ -294,11 +510,7 @@ async fn discover_root_with_sandbox(
anyhow::bail!("expected discovery response, received {response:?}");
};
assert_eq!(id, request_id);
let response: CapabilityRootsDiscoverResponse = serde_json::from_value(result)?;
let [discovery] = response.roots.as_slice() else {
anyhow::bail!("expected exactly one discovered root");
};
Ok(discovery.clone())
Ok(serde_json::from_value(result)?)
}
async fn initialize(server: &mut common::exec_server::ExecServerHarness) -> anyhow::Result<()> {
codex-rs/exec-server/tests/common/fake_bwrap.rs
+62
View File
@@ -0,0 +1,62 @@
use std::os::unix::fs::PermissionsExt;
use std::path::Path;
use std::path::PathBuf;
pub(crate) fn write_fake_bwrap(bin_dir: &Path) -> anyhow::Result<PathBuf> {
std::fs::create_dir_all(bin_dir)?;
let fake_bwrap = bin_dir.join("bwrap");
std::fs::write(
&fake_bwrap,
r#"#!/bin/bash
set -euo pipefail
for arg in "$@"; do
if [[ "${arg}" == "--help" ]]; then
echo "Usage: bwrap --argv0 --perms --as-pid-1"
exit 0
fi
done
args=("$@")
argv0=""
command_start=-1
for i in "${!args[@]}"; do
if [[ "${args[$i]}" == "--argv0" && $((i + 1)) -lt ${#args[@]} ]]; then
argv0="${args[$((i + 1))]}"
fi
if [[ "${args[$i]}" == "--" ]]; then
command_start=$((i + 1))
break
fi
done
if [[ "${command_start}" -lt 0 || "${command_start}" -ge "${#args[@]}" ]]; then
echo "fake bwrap did not find an inner command" >&2
exit 125
fi
cmd=("${args[@]:$command_start}")
case "${cmd[0]}" in
/usr/bin/true|/bin/true|true)
exec "${cmd[@]}"
;;
esac
printf '%s\n' "$*" >> "${0}.log"
if [[ -f "${0}.fail-once" ]]; then
rm -f "${0}.fail-once"
echo "forced fake bwrap failure" >&2
exit 125
fi
if [[ -n "${argv0}" ]]; then
exec -a "${argv0}" "${cmd[@]}"
fi
exec "${cmd[@]}"
"#,
)?;
let mut permissions = std::fs::metadata(&fake_bwrap)?.permissions();
permissions.set_mode(0o755);
std::fs::set_permissions(&fake_bwrap, permissions)?;
Ok(fake_bwrap)
}
codex-rs/exec-server/tests/file_system_unix.rs
+5 -49
View File
@@ -2,6 +2,9 @@
#![allow(clippy::expect_used)]
mod common;
#[cfg(target_os = "linux")]
#[path = "common/fake_bwrap.rs"]
mod fake_bwrap;
#[path = "file_system/shared.rs"]
mod shared;
@@ -55,6 +58,8 @@ use tokio::time::timeout;
#[cfg(target_os = "linux")]
use crate::common::exec_server::exec_server_with_env;
#[cfg(target_os = "linux")]
use crate::fake_bwrap::write_fake_bwrap;
use crate::support::FileSystemImplementation;
use crate::support::create_file_system_context;
@@ -117,55 +122,6 @@ fn create_directory_symlink(target: &Path, alias: &Path) -> Result<()> {
Ok(())
}
#[cfg(target_os = "linux")]
fn write_fake_bwrap(bin_dir: &Path) -> Result<PathBuf> {
std::fs::create_dir_all(bin_dir)?;
let fake_bwrap = bin_dir.join("bwrap");
std::fs::write(
&fake_bwrap,
r#"#!/bin/bash
set -euo pipefail
for arg in "$@"; do
if [[ "${arg}" == "--help" ]]; then
echo "Usage: bwrap --argv0 --perms --as-pid-1"
exit 0
fi
done
printf '%s\n' "$*" >> "${0}.log"
args=("$@")
argv0=""
command_start=-1
for i in "${!args[@]}"; do
if [[ "${args[$i]}" == "--argv0" && $((i + 1)) -lt ${#args[@]} ]]; then
argv0="${args[$((i + 1))]}"
fi
if [[ "${args[$i]}" == "--" ]]; then
command_start=$((i + 1))
break
fi
done
if [[ "${command_start}" -lt 0 || "${command_start}" -ge "${#args[@]}" ]]; then
echo "fake bwrap did not find an inner command" >&2
exit 125
fi
cmd=("${args[@]:$command_start}")
if [[ -n "${argv0}" ]]; then
exec -a "${argv0}" "${cmd[@]}"
fi
exec "${cmd[@]}"
"#,
)?;
let mut permissions = std::fs::metadata(&fake_bwrap)?.permissions();
permissions.set_mode(0o755);
std::fs::set_permissions(&fake_bwrap, permissions)?;
Ok(fake_bwrap)
}
#[test_case(FileSystemImplementation::Local ; "local")]
#[test_case(FileSystemImplementation::Remote ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]