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