From eb15245d820d094cc0e5a950bc052bd7c2c5f925 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 24 Aug 2026 17:23:01 +0000 Subject: [PATCH] Batch sandboxed capability root discovery (#40443) ## 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 --- codex-rs/exec-server/src/fs_helper.rs | 26 ++ codex-rs/exec-server/src/local_file_system.rs | 2 +- .../exec-server/src/sandboxed_file_system.rs | 18 ++ .../src/server/file_system_handler.rs | 37 +++ .../exec-server/tests/capability_discovery.rs | 236 +++++++++++++++++- .../exec-server/tests/common/fake_bwrap.rs | 62 +++++ .../exec-server/tests/file_system_unix.rs | 54 +--- 7 files changed, 373 insertions(+), 62 deletions(-) create mode 100644 codex-rs/exec-server/tests/common/fake_bwrap.rs diff --git a/codex-rs/exec-server/src/fs_helper.rs b/codex-rs/exec-server/src/fs_helper.rs index e519215faf..ad6eacf93c 100644 --- a/codex-rs/exec-server/src/fs_helper.rs +++ b/codex-rs/exec-server/src/fs_helper.rs @@ -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 { + 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 { match self { Self::ReadFile(response) => Ok(response), @@ -223,6 +243,12 @@ pub(crate) async fn run_direct_request( ) -> Result { 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(), )), diff --git a/codex-rs/exec-server/src/local_file_system.rs b/codex-rs/exec-server/src/local_file_system.rs index 7f33b5a697..66a4876d45 100644 --- a/codex-rs/exec-server/src/local_file_system.rs +++ b/codex-rs/exec-server/src/local_file_system.rs @@ -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, diff --git a/codex-rs/exec-server/src/sandboxed_file_system.rs b/codex-rs/exec-server/src/sandboxed_file_system.rs index 571deda1bb..61c3c06126 100644 --- a/codex-rs/exec-server/src/sandboxed_file_system.rs +++ b/codex-rs/exec-server/src/sandboxed_file_system.rs @@ -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 { + 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, diff --git a/codex-rs/exec-server/src/server/file_system_handler.rs b/codex-rs/exec-server/src/server/file_system_handler.rs index fd5eaaea2d..be6802e98c 100644 --- a/codex-rs/exec-server/src/server/file_system_handler.rs +++ b/codex-rs/exec-server/src/server/file_system_handler.rs @@ -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 { + 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())) diff --git a/codex-rs/exec-server/tests/capability_discovery.rs b/codex-rs/exec-server/tests/capability_discovery.rs index c970c91b78..a547c978ba 100644 --- a/codex-rs/exec-server/tests/capability_discovery.rs +++ b/codex-rs/exec-server/tests/capability_discovery.rs @@ -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::>(), + root.error, + )) + .collect::>(), + 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![ + ("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![ + ("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, ) -> anyhow::Result { + 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, +) -> anyhow::Result { 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<()> { diff --git a/codex-rs/exec-server/tests/common/fake_bwrap.rs b/codex-rs/exec-server/tests/common/fake_bwrap.rs new file mode 100644 index 0000000000..ddd2936480 --- /dev/null +++ b/codex-rs/exec-server/tests/common/fake_bwrap.rs @@ -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 { + 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) +} diff --git a/codex-rs/exec-server/tests/file_system_unix.rs b/codex-rs/exec-server/tests/file_system_unix.rs index e27c8ed1be..4b5615611f 100644 --- a/codex-rs/exec-server/tests/file_system_unix.rs +++ b/codex-rs/exec-server/tests/file_system_unix.rs @@ -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 { - 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)]