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

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

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

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

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

@@ -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 {
+            serde_json::to_value(CapabilityRootsDiscoverParams { roots })?,
-                roots: vec![CapabilityRootDiscoverRequest {
-                    id: id.to_string(),
-                    path,
-                    sandbox,
-                }],
-            })?,
         )
         .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)?;
+    Ok(serde_json::from_value(result)?)
-    let [discovery] = response.roots.as_slice() else {
-        anyhow::bail!("expected exactly one discovered root");
-    };
-    Ok(discovery.clone())
 }
 
 async fn initialize(server: &mut common::exec_server::ExecServerHarness) -> anyhow::Result<()> {

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

@@ -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)]