Bound in-flight file opens in exec-server (#50162)

## Why

The per-connection limit of 128 file read handles was checked only after opening the file, leaving pending opens outside the limit.

## What changed

Reserve a semaphore permit before polling the file-open future, counting pending opens and registered handles against the same limit. Release capacity when an open fails or is cancelled, or when a handle closes, even if a read still retains the file.

Use synchronous handle-table locks that never cross an await, and recheck duplicate IDs after opening to preserve the winning handle during concurrent opens.

## Testing

Add tests for stalled-open progress, pending-open admission limits, permit cleanup after failure or cancellation, capacity reuse while a file remains in use, and concurrent duplicate IDs.

GitOrigin-RevId: caabf38c74dda5951f2fc75234a7149daaa01263

codex-rs/exec-server/src/file_handle.rs

@@ -1,10 +1,18 @@
+//! Connection-scoped file handles with bounded positional reads.
+//! Semaphore permits bound open handles and in-flight opens; table locks never cross an await.
+
 use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::fs::File;
+use std::future::Future;
 use std::io;
 use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::MutexGuard;
 
 use codex_file_system::FILE_READ_CHUNK_SIZE;
-use tokio::sync::Mutex;
+use tokio::sync::OwnedSemaphorePermit;
+use tokio::sync::Semaphore;
 
 const MAX_OPEN_FILES: usize = 128;
 
@@ -14,32 +22,51 @@ pub(crate) struct FileReadBlock {
     pub(crate) eof: bool,
 }
 
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub(crate) struct FileHandleManager {
-    handles: Arc<Mutex<HashMap<String, Arc<File>>>>,
+    handles: Arc<Mutex<HashMap<String, FileHandleEntry>>>,
+    slots: Arc<Semaphore>,
+}
+
+impl Default for FileHandleManager {
+    fn default() -> Self {
+        Self {
+            handles: Arc::default(),
+            slots: Arc::new(Semaphore::new(MAX_OPEN_FILES)),
+        }
+    }
 }
 
 impl FileHandleManager {
     pub(crate) async fn open(
         &self,
         handle_id: String,
-        file: tokio::fs::File,
+        open_file: impl Future<Output = io::Result<tokio::fs::File>>,
     ) -> io::Result<String> {
-        let file = Arc::new(file.into_std().await);
-        let mut handles = self.handles.lock().await;
-        if handles.contains_key(&handle_id) {
+        if self.lock_handles().contains_key(&handle_id) {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidInput,
                 format!("file read handle `{handle_id}` already exists"),
             ));
         }
-        if handles.len() >= MAX_OPEN_FILES {
-            return Err(io::Error::new(
+        let permit = Arc::clone(&self.slots).try_acquire_owned().map_err(|_| {
+            io::Error::new(
                 io::ErrorKind::InvalidInput,
                 format!("at most {MAX_OPEN_FILES} file reads may be open per connection"),
+            )
+        })?;
+        let file = Arc::new(open_file.await?.into_std().await);
+        let mut handles = self.lock_handles();
+        let Entry::Vacant(entry) = handles.entry(handle_id.clone()) else {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("file read handle `{handle_id}` already exists"),
             ));
-        }
-        handles.insert(handle_id.clone(), file);
+        };
+        entry.insert(FileHandleEntry {
+            file,
+            _permit: permit,
+        });
         Ok(handle_id)
     }
 
@@ -50,13 +77,7 @@ impl FileHandleManager {
         len: usize,
     ) -> io::Result<FileReadBlock> {
         validate_read_block_len(len)?;
-        let file = {
-            let handles = self.handles.lock().await;
-            handles
-                .get(handle_id)
-                .cloned()
-                .ok_or_else(|| unknown_handle_error(handle_id))?
-        };
+        let file = self.get(handle_id)?;
         let result =
             match tokio::task::spawn_blocking(move || read_block_at(&file, offset, len)).await {
                 Ok(result) => result,
@@ -65,18 +86,37 @@ impl FileHandleManager {
                 ))),
             };
         if result.is_err() {
-            self.close(handle_id).await;
+            self.close(handle_id);
         }
         result
     }
 
-    pub(crate) async fn close(&self, handle_id: &str) {
-        self.handles.lock().await.remove(handle_id);
+    fn get(&self, handle_id: &str) -> io::Result<Arc<File>> {
+        self.lock_handles()
+            .get(handle_id)
+            .map(|entry| Arc::clone(&entry.file))
+            .ok_or_else(|| unknown_handle_error(handle_id))
     }
 
-    pub(crate) async fn close_all(&self) {
-        self.handles.lock().await.clear();
+    pub(crate) fn close(&self, handle_id: &str) {
+        self.lock_handles().remove(handle_id);
     }
+
+    pub(crate) fn close_all(&self) {
+        self.lock_handles().clear();
+    }
+
+    fn lock_handles(&self) -> MutexGuard<'_, HashMap<String, FileHandleEntry>> {
+        self.handles
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+    }
+}
+
+struct FileHandleEntry {
+    file: Arc<File>,
+    // Closing an entry releases capacity even if a read still holds the file.
+    _permit: OwnedSemaphorePermit,
 }
 
 fn read_block_at(file: &File, offset: u64, len: usize) -> io::Result<FileReadBlock> {
@@ -126,3 +166,7 @@ fn unknown_handle_error(handle_id: &str) -> io::Error {
         format!("unknown file read handle `{handle_id}`"),
     )
 }
+
+#[cfg(test)]
+#[path = "file_handle_tests.rs"]
+mod tests;

codex-rs/exec-server/src/file_handle_tests.rs

@@ -0,0 +1,134 @@
+//! Exercises permit cleanup and progress without relying on I/O timing.
+
+use std::future::pending;
+use std::future::ready;
+use std::io;
+use std::io::Write;
+use std::task::Poll;
+
+use anyhow::Result;
+use pretty_assertions::assert_eq;
+use tokio::fs::File;
+use tokio::sync::oneshot;
+
+use super::FileHandleManager;
+use super::FileReadBlock;
+use super::MAX_OPEN_FILES;
+
+/// A stalled open must let unrelated handles open and read without waiting for it.
+#[tokio::test]
+async fn stalled_open_does_not_block_other_handles() -> Result<()> {
+    let manager = FileHandleManager::default();
+    let mut stalled = Box::pin(manager.open("stalled".to_string(), pending()));
+    assert!(futures::poll!(stalled.as_mut()).is_pending());
+
+    let mut other = Box::pin(manager.open(
+        "other".to_string(),
+        ready(Ok(File::from_std(tempfile::tempfile()?))),
+    ));
+    let Poll::Ready(result) = futures::poll!(other.as_mut()) else {
+        anyhow::bail!("an unrelated open was blocked by the stalled open");
+    };
+    assert_eq!(result?, "other");
+    manager.read_block("other", /*offset*/ 0, /*len*/ 1).await?;
+    Ok(())
+}
+
+/// Pending opens consume capacity before polling another file open.
+#[tokio::test]
+async fn pending_opens_reserve_capacity() -> Result<()> {
+    let manager = FileHandleManager::default();
+    let mut opens = Vec::new();
+    for index in 0..MAX_OPEN_FILES {
+        let mut open = Box::pin(manager.open(index.to_string(), pending()));
+        assert!(futures::poll!(open.as_mut()).is_pending());
+        opens.push(open);
+    }
+    let mut rejected = Box::pin(manager.open("overflow".to_string(), pending()));
+    let Poll::Ready(Err(error)) = futures::poll!(rejected.as_mut()) else {
+        anyhow::bail!("admission must fail before polling the file open");
+    };
+    assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+    Ok(())
+}
+
+/// Failed and cancelled opens must return their permits for reuse.
+#[test_case::test_case(true; "cancelled")]
+#[test_case::test_case(false; "failed")]
+#[tokio::test]
+async fn unfinished_opens_release_capacity(cancelled: bool) -> Result<()> {
+    let manager = FileHandleManager::default();
+    for _ in 0..=MAX_OPEN_FILES {
+        let mut open = Box::pin(manager.open("reused".to_string(), async {
+            if cancelled {
+                pending::<()>().await;
+            }
+            Err(io::Error::other("open failed"))
+        }));
+        let result = futures::poll!(open.as_mut());
+        if cancelled {
+            assert!(result.is_pending());
+        } else {
+            let Poll::Ready(Err(error)) = result else {
+                anyhow::bail!("the open error must reach the caller");
+            };
+            assert_eq!(error.kind(), io::ErrorKind::Other);
+        }
+        drop(open);
+    }
+    Ok(())
+}
+
+/// Closing an entry frees capacity even while an operation retains its file descriptor.
+#[tokio::test]
+async fn close_releases_capacity_while_the_file_is_still_in_use() -> Result<()> {
+    let manager = FileHandleManager::default();
+    let file = tempfile::tempfile()?;
+    for index in 0..MAX_OPEN_FILES {
+        manager
+            .open(
+                index.to_string(),
+                ready(Ok(File::from_std(file.try_clone()?))),
+            )
+            .await?;
+    }
+    let retained = manager.get("0")?;
+    manager.close("0");
+    manager
+        .open("next".to_string(), ready(Ok(File::from_std(file))))
+        .await?;
+    drop(retained);
+    Ok(())
+}
+
+/// Concurrent opens of one ID must reject the loser without changing the winning file.
+#[tokio::test]
+async fn concurrent_opens_reject_duplicate_ids_without_replacing_the_file() -> Result<()> {
+    let manager = FileHandleManager::default();
+    let (send, receive) = oneshot::channel();
+    let mut late = Box::pin(manager.open("shared".to_string(), async {
+        receive.await.map_err(io::Error::other)
+    }));
+    assert!(futures::poll!(late.as_mut()).is_pending());
+
+    let mut file = tempfile::tempfile()?;
+    file.write_all(b"first")?;
+    manager
+        .open("shared".to_string(), ready(Ok(File::from_std(file))))
+        .await?;
+    send.send(File::from_std(tempfile::tempfile()?))
+        .map_err(|_| anyhow::anyhow!("the pending open dropped its receiver"))?;
+
+    let error = late.await.expect_err("the duplicate open must fail");
+    assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+    assert_eq!(
+        manager
+            .read_block("shared", /*offset*/ 0, /*len*/ 6)
+            .await?,
+        FileReadBlock {
+            bytes: b"first".to_vec(),
+            eof: true,
+        }
+    );
+    Ok(())
+}

codex-rs/exec-server/src/local_file_system.rs

@@ -118,7 +118,7 @@ impl LocalFileSystem {
 }
 
 impl LocalFileSystem {
-    pub(crate) async fn open_file_for_read(
+    pub(crate) async fn open_file(
         &self,
         path: &PathUri,
         sandbox: Option<&FileSystemSandboxContext>,
@@ -127,11 +127,9 @@ impl LocalFileSystem {
             sandbox.validate_file_system_paths_for_current_host()?;
         }
         if sandbox.is_some_and(FileSystemSandboxContext::should_read_from_sandbox) {
-            return self.sandboxed()?.open_file_for_read(path, sandbox).await;
+            return self.sandboxed()?.open_file(path, sandbox).await;
         }
-        self.unsandboxed
-            .open_file_for_read(path, /*sandbox*/ None)
-            .await
+        self.unsandboxed.open_file(path, /*sandbox*/ None).await
     }
 
     async fn canonicalize(
@@ -350,15 +348,13 @@ impl ExecutorFileSystem for LocalFileSystem {
 }
 
 impl UnsandboxedFileSystem {
-    async fn open_file_for_read(
+    async fn open_file(
         &self,
         path: &PathUri,
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<tokio::fs::File> {
         reject_platform_sandbox_context(sandbox)?;
-        self.file_system
-            .open_file_for_read(path, /*sandbox*/ None)
-            .await
+        self.file_system.open_file(path, /*sandbox*/ None).await
     }
 
     async fn canonicalize(
@@ -581,7 +577,7 @@ impl ExecutorFileSystem for UnsandboxedFileSystem {
 }
 
 impl DirectFileSystem {
-    async fn open_file_for_read(
+    async fn open_file(
         &self,
         path: &PathUri,
         sandbox: Option<&FileSystemSandboxContext>,
@@ -619,7 +615,7 @@ impl DirectFileSystem {
         path: &PathUri,
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<FileSystemReadStream> {
-        let file = self.open_file_for_read(path, sandbox).await?;
+        let file = self.open_file(path, sandbox).await?;
         Ok(FileSystemReadStream::new(ReaderStream::with_capacity(
             file,
             FILE_READ_CHUNK_SIZE,

codex-rs/exec-server/src/sandboxed_file_system.rs

@@ -80,7 +80,7 @@ impl SandboxedFileSystem {
             .map_err(map_sandbox_error)
     }
 
-    pub(crate) async fn open_file_for_read(
+    pub(crate) async fn open_file(
         &self,
         path: &PathUri,
         sandbox: Option<&FileSystemSandboxContext>,
@@ -367,7 +367,7 @@ impl ExecutorFileSystem for SandboxedFileSystem {
         sandbox: Option<&'a FileSystemSandboxContext>,
     ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
         Box::pin(async move {
-            let file = self.open_file_for_read(path, sandbox).await?;
+            let file = self.open_file(path, sandbox).await?;
             Ok(FileSystemReadStream::new(ReaderStream::with_capacity(
                 file,
                 FILE_READ_CHUNK_SIZE,

codex-rs/exec-server/src/server/file_system_handler.rs

@@ -70,7 +70,7 @@ impl FileSystemHandler {
     }
 
     pub(crate) async fn shutdown(&self) {
-        self.file_handles.close_all().await;
+        self.file_handles.close_all();
     }
 
     pub(crate) async fn discover_capability_roots(
@@ -131,14 +131,13 @@ impl FileSystemHandler {
                 "exec-server does not support writable file streams".to_string(),
             ));
         }
-        let file = self
-            .file_system
-            .open_file_for_read(&params.path, params.sandbox.as_ref())
-            .await
-            .map_err(map_fs_error)?;
         let handle_id = self
             .file_handles
-            .open(params.handle_id, file)
+            .open(
+                params.handle_id,
+                self.file_system
+                    .open_file(&params.path, params.sandbox.as_ref()),
+            )
             .await
             .map_err(map_fs_error)?;
         Ok(FsOpenResponse { handle_id })
@@ -176,7 +175,7 @@ impl FileSystemHandler {
         params: FsCloseParams,
     ) -> Result<FsCloseResponse, JSONRPCErrorError> {
         validate_file_handle_id(&params.handle_id)?;
-        self.file_handles.close(&params.handle_id).await;
+        self.file_handles.close(&params.handle_id);
         Ok(FsCloseResponse {})
     }