Coalesce concurrent remote metadata requests (#36184)

## Why Concurrent callers can request metadata for the same remote path, resulting in duplicate `fs/getMetadata` RPCs. ## What changed - Share an unsandboxed metadata request with concurrent callers for the same path while the RPC is in flight. - Remove completed and failed requests so later callers fetch fresh metadata, and clear in-flight entries after filesystem mutations, including failed mutations. - Keep requests for different paths separate and bypass sharing for sandboxed requests. ## Testing Add WebSocket-backed tests covering shared results and errors, retry after cancellation or failure, mutation invalidation, and path and sandbox isolation. GitOrigin-RevId: f91e51b94db3fb984597aec8d302b8715778121c
2026-10-18 17:18:52 +00:00 · 2026-07-30 16:35:20 +00:00
parent 0dcad0c972
commit 483559cc75
2 changed files with 505 additions and 13 deletions
@@ -1,7 +1,12 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
 use base64::Engine as _;
 use base64::engine::general_purpose::STANDARD;
 use codex_utils_path_uri::PathUri;
 use tokio::io;
+use tokio::sync::Mutex;
+use tokio::sync::OnceCell;
 use tracing::trace;

 use crate::CopyOptions;
@@ -35,14 +40,20 @@ const NOT_FOUND_ERROR_CODE: i64 = -32004;
 #[path = "remote_file_stream.rs"]
 mod file_stream;

+type InFlightMetadataRequest = OnceCell<Result<FileMetadata, Arc<io::Error>>>;
+
 pub(crate) struct RemoteFileSystem {
    client: LazyRemoteExecServerClient,
+    metadata_requests: Mutex<HashMap<PathUri, Arc<InFlightMetadataRequest>>>,
 }

 impl RemoteFileSystem {
    pub(crate) fn new(client: LazyRemoteExecServerClient) -> Self {
        trace!("remote fs new");
-        Self { client }
+        Self {
+            client,
+            metadata_requests: Mutex::new(HashMap::new()),
+        }
    }

    async fn canonicalize(
@@ -108,14 +119,15 @@ impl RemoteFileSystem {
    ) -> FileSystemResult<()> {
        trace!("remote fs write_file");
        let client = self.client.get().await.map_err(map_remote_error)?;
-        client
+        let result = client
            .fs_write_file(FsWriteFileParams {
                path: path.clone(),
                data_base64: STANDARD.encode(contents),
                sandbox: remote_sandbox_context(sandbox),
            })
-            .await
-            .map_err(map_remote_error)?;
+            .await;
+        self.metadata_requests.lock().await.clear();
+        result.map_err(map_remote_error)?;
        Ok(())
    }

@@ -127,21 +139,61 @@ impl RemoteFileSystem {
    ) -> FileSystemResult<()> {
        trace!("remote fs create_directory");
        let client = self.client.get().await.map_err(map_remote_error)?;
-        client
+        let result = client
            .fs_create_directory(FsCreateDirectoryParams {
                path: path.clone(),
                recursive: Some(options.recursive),
                sandbox: remote_sandbox_context(sandbox),
            })
-            .await
-            .map_err(map_remote_error)?;
+            .await;
+        self.metadata_requests.lock().await.clear();
+        result.map_err(map_remote_error)?;
        Ok(())
    }

+    /// Shares identical unsandboxed metadata requests only while their RPC remains in flight.
    async fn get_metadata(
        &self,
        path: &PathUri,
        sandbox: Option<&FileSystemSandboxContext>,
+    ) -> FileSystemResult<FileMetadata> {
+        if sandbox.is_some() {
+            return self.get_metadata_uncached(path, sandbox).await;
+        }
+
+        let request = {
+            let mut requests = self.metadata_requests.lock().await;
+            requests.retain(|_, in_flight| Arc::strong_count(in_flight) > 1);
+            Arc::clone(requests.entry(path.clone()).or_default())
+        };
+        let result = match request
+            .get_or_init(|| async {
+                self.get_metadata_uncached(path, /*sandbox*/ None)
+                    .await
+                    .map_err(Arc::new)
+            })
+            .await
+        {
+            Ok(metadata) => Ok(metadata.clone()),
+            Err(error) => Err(io::Error::new(error.kind(), error.to_string())),
+        };
+
+        let mut requests = self.metadata_requests.lock().await;
+        if requests
+            .get(path)
+            .is_some_and(|in_flight| Arc::ptr_eq(in_flight, &request))
+        {
+            requests.remove(path);
+        }
+
+        result
+    }
+
+    /// Sends a fresh metadata request without sharing another caller's sandbox or result.
+    async fn get_metadata_uncached(
+        &self,
+        path: &PathUri,
+        sandbox: Option<&FileSystemSandboxContext>,
    ) -> FileSystemResult<FileMetadata> {
        trace!("remote fs get_metadata");
        let client = self.client.get().await.map_err(map_remote_error)?;
@@ -226,15 +278,16 @@ impl RemoteFileSystem {
    ) -> FileSystemResult<()> {
        trace!("remote fs remove");
        let client = self.client.get().await.map_err(map_remote_error)?;
-        client
+        let result = client
            .fs_remove(FsRemoveParams {
                path: path.clone(),
                recursive: Some(options.recursive),
                force: Some(options.force),
                sandbox: remote_sandbox_context(sandbox),
            })
-            .await
-            .map_err(map_remote_error)?;
+            .await;
+        self.metadata_requests.lock().await.clear();
+        result.map_err(map_remote_error)?;
        Ok(())
    }

@@ -247,15 +300,16 @@ impl RemoteFileSystem {
    ) -> FileSystemResult<()> {
        trace!("remote fs copy");
        let client = self.client.get().await.map_err(map_remote_error)?;
-        client
+        let result = client
            .fs_copy(FsCopyParams {
                source_path: source_path.clone(),
                destination_path: destination_path.clone(),
                recursive: options.recursive,
                sandbox: remote_sandbox_context(sandbox),
            })
-            .await
-            .map_err(map_remote_error)?;
+            .await;
+        self.metadata_requests.lock().await.clear();
+        result.map_err(map_remote_error)?;
        Ok(())
    }
 }
@@ -1,5 +1,7 @@
 #![allow(clippy::expect_used)]

+use codex_exec_server_protocol::JSONRPCError;
+use codex_exec_server_protocol::JSONRPCErrorError;
 use codex_exec_server_protocol::JSONRPCMessage;
 use codex_exec_server_protocol::JSONRPCResponse;
 use codex_http_client::HttpClientFactory;
@@ -27,7 +29,14 @@ use tokio_tungstenite::tungstenite::Message;
 use super::*;
 use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
 use crate::client_api::ExecServerTransportParams;
+use crate::protocol::FS_COPY_METHOD;
+use crate::protocol::FS_CREATE_DIRECTORY_METHOD;
+use crate::protocol::FS_GET_METADATA_METHOD;
 use crate::protocol::FS_READ_FILE_METHOD;
+use crate::protocol::FS_REMOVE_METHOD;
+use crate::protocol::FS_WRITE_FILE_METHOD;
+use crate::protocol::FsGetMetadataParams;
+use crate::protocol::FsGetMetadataResponse;
 use crate::protocol::FsReadFileParams;
 use crate::protocol::FsReadFileResponse;
 use crate::protocol::INITIALIZE_METHOD;
@@ -86,6 +95,335 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve
    server.await.expect("recording server should succeed");
 }

+#[tokio::test]
+async fn concurrent_remote_metadata_requests_share_only_in_flight_results() {
+    let (abandoned_request_tx, abandoned_request_rx) = oneshot::channel();
+    let (websocket_url, captured_params, server) = record_metadata_params(vec![
+        MetadataResponse::Immediate(Ok(metadata_response(/*size*/ 42))),
+        MetadataResponse::Abandoned(abandoned_request_tx),
+        MetadataResponse::Immediate(Ok(metadata_response(/*size*/ 43))),
+    ])
+    .await;
+    let file_system = std::sync::Arc::new(RemoteFileSystem::new(LazyRemoteExecServerClient::new(
+        ExecServerTransportParams::websocket_url(
+            websocket_url,
+            DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+        ),
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )));
+    let path = PathUri::parse("file:///workspace/project/AGENTS.md").expect("valid path URI");
+
+    let (first, second) = tokio::join!(
+        file_system.get_metadata(&path, /*sandbox*/ None),
+        file_system.get_metadata(&path, /*sandbox*/ None),
+    );
+    let expected = FileMetadata {
+        is_directory: false,
+        is_file: true,
+        is_symlink: false,
+        size: 42,
+        created_at_ms: 10,
+        modified_at_ms: 20,
+    };
+    assert_eq!(first.expect("first metadata request"), expected);
+    assert_eq!(second.expect("second metadata request"), expected);
+    assert!(file_system.metadata_requests.lock().await.is_empty());
+
+    let initializing_file_system = std::sync::Arc::clone(&file_system);
+    let initializing_path = path.clone();
+    let initializer = tokio::spawn(async move {
+        initializing_file_system
+            .get_metadata(&initializing_path, /*sandbox*/ None)
+            .await
+    });
+    abandoned_request_rx
+        .await
+        .expect("server should receive the abandoned metadata request");
+
+    let follower = file_system.get_metadata(&path, /*sandbox*/ None);
+    tokio::pin!(follower);
+    assert!(futures::poll!(follower.as_mut()).is_pending());
+    assert_eq!(
+        file_system
+            .metadata_requests
+            .lock()
+            .await
+            .get(&path)
+            .map(std::sync::Arc::strong_count),
+        Some(3)
+    );
+
+    initializer.abort();
+    assert!(
+        initializer
+            .await
+            .expect_err("metadata initializer should be aborted")
+            .is_cancelled()
+    );
+
+    let mut refreshed = expected;
+    refreshed.size = 43;
+    assert_eq!(
+        follower
+            .await
+            .expect("waiting metadata request should retry the abandoned initializer"),
+        refreshed
+    );
+    assert!(file_system.metadata_requests.lock().await.is_empty());
+    assert_eq!(
+        captured_params.await.expect("captured metadata requests"),
+        vec![
+            FsGetMetadataParams {
+                path: path.clone(),
+                sandbox: None,
+            };
+            3
+        ]
+    );
+    server.await.expect("metadata server should succeed");
+}
+
+#[tokio::test]
+async fn concurrent_remote_metadata_errors_are_shared_but_retried() {
+    let (websocket_url, captured_params, server) = record_metadata_params(vec![
+        MetadataResponse::Immediate(Err(JSONRPCErrorError {
+            code: NOT_FOUND_ERROR_CODE,
+            data: None,
+            message: "metadata not found".to_string(),
+        })),
+        MetadataResponse::Immediate(Ok(metadata_response(/*size*/ 42))),
+    ])
+    .await;
+    let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new(
+        ExecServerTransportParams::websocket_url(
+            websocket_url,
+            DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+        ),
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    ));
+    let path = PathUri::parse("file:///workspace/project/AGENTS.md").expect("valid path URI");
+
+    let (first, second) = tokio::join!(
+        file_system.get_metadata(&path, /*sandbox*/ None),
+        file_system.get_metadata(&path, /*sandbox*/ None),
+    );
+    assert_eq!(
+        first.expect_err("first metadata error").kind(),
+        std::io::ErrorKind::NotFound
+    );
+    assert_eq!(
+        second.expect_err("shared metadata error").kind(),
+        std::io::ErrorKind::NotFound
+    );
+    assert_eq!(
+        file_system
+            .get_metadata(&path, /*sandbox*/ None)
+            .await
+            .expect("failed metadata request should be retried")
+            .size,
+        42
+    );
+    assert_eq!(
+        captured_params.await.expect("captured metadata requests"),
+        vec![
+            FsGetMetadataParams {
+                path: path.clone(),
+                sandbox: None,
+            },
+            FsGetMetadataParams {
+                path,
+                sandbox: None,
+            },
+        ]
+    );
+    server.await.expect("metadata server should succeed");
+}
+
+#[tokio::test]
+async fn remote_metadata_starts_fresh_after_intervening_filesystem_mutation() {
+    let path = PathUri::parse("file:///workspace/project/AGENTS.md").expect("valid path URI");
+    let source_path =
+        PathUri::parse("file:///workspace/project/source.md").expect("valid source path URI");
+
+    for (mutation, mutation_response) in [
+        (MetadataMutation::Write, Ok(())),
+        (MetadataMutation::CreateDirectory, Ok(())),
+        (MetadataMutation::Remove, Ok(())),
+        (MetadataMutation::Copy, Ok(())),
+        (
+            MetadataMutation::Copy,
+            Err(JSONRPCErrorError {
+                code: INVALID_REQUEST_ERROR_CODE,
+                data: None,
+                message: "mutation partially failed".to_string(),
+            }),
+        ),
+    ] {
+        let mutation_fails = mutation_response.is_err();
+        let (websocket_url, captured_params, server) = record_metadata_params(vec![
+            MetadataResponse::Immediate(Ok(metadata_response(/*size*/ 42))),
+            MetadataResponse::Mutation(mutation_response),
+            MetadataResponse::Immediate(Ok(metadata_response(/*size*/ 43))),
+        ])
+        .await;
+        let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new(
+            ExecServerTransportParams::websocket_url(
+                websocket_url,
+                DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+            ),
+            HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ));
+        file_system
+            .client
+            .get()
+            .await
+            .expect("remote filesystem client should connect");
+
+        let stale_request = file_system.get_metadata(&path, /*sandbox*/ None);
+        tokio::pin!(stale_request);
+        assert!(futures::poll!(stale_request.as_mut()).is_pending());
+
+        let result = match mutation {
+            MetadataMutation::Write => {
+                file_system
+                    .write_file(&path, b"updated".to_vec(), /*sandbox*/ None)
+                    .await
+            }
+            MetadataMutation::CreateDirectory => {
+                file_system
+                    .create_directory(
+                        &path,
+                        CreateDirectoryOptions { recursive: true },
+                        /*sandbox*/ None,
+                    )
+                    .await
+            }
+            MetadataMutation::Remove => {
+                file_system
+                    .remove(
+                        &path,
+                        RemoveOptions {
+                            recursive: true,
+                            force: true,
+                        },
+                        /*sandbox*/ None,
+                    )
+                    .await
+            }
+            MetadataMutation::Copy => {
+                file_system
+                    .copy(
+                        &source_path,
+                        &path,
+                        CopyOptions { recursive: true },
+                        /*sandbox*/ None,
+                    )
+                    .await
+            }
+        };
+        if mutation_fails {
+            assert_eq!(
+                result
+                    .expect_err("remote filesystem mutation should fail")
+                    .kind(),
+                std::io::ErrorKind::InvalidInput
+            );
+        } else {
+            result.expect("remote filesystem mutation should succeed");
+        }
+
+        let (refreshed, stale) = tokio::join!(
+            file_system.get_metadata(&path, /*sandbox*/ None),
+            stale_request.as_mut(),
+        );
+        assert_eq!(stale.expect("original metadata request").size, 42);
+        assert_eq!(refreshed.expect("post-mutation metadata request").size, 43);
+        assert_eq!(
+            captured_params.await.expect("captured metadata requests"),
+            vec![
+                FsGetMetadataParams {
+                    path: path.clone(),
+                    sandbox: None,
+                };
+                2
+            ]
+        );
+        server.await.expect("metadata server should succeed");
+    }
+}
+
+#[tokio::test]
+async fn remote_metadata_requests_do_not_cross_path_or_sandbox_boundaries() {
+    let metadata = metadata_response(/*size*/ 42);
+    let (websocket_url, captured_params, server) = record_metadata_params(vec![
+        MetadataResponse::Immediate(Ok(metadata.clone())),
+        MetadataResponse::Immediate(Ok(metadata.clone())),
+        MetadataResponse::Immediate(Ok(metadata.clone())),
+        MetadataResponse::Immediate(Ok(metadata)),
+    ])
+    .await;
+    let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new(
+        ExecServerTransportParams::websocket_url(
+            websocket_url,
+            DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
+        ),
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    ));
+    let first_path = PathUri::parse("file:///workspace/project/AGENTS.md").expect("valid path URI");
+    let second_path = PathUri::parse("file:///workspace/project/SKILL.md").expect("valid path URI");
+    let sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
+        PermissionProfile::from_runtime_permissions(
+            &FileSystemSandboxPolicy::restricted(vec![FileSystemSandboxEntry {
+                path: FileSystemPath::Special {
+                    value: FileSystemSpecialPath::project_roots(/*subpath*/ None),
+                },
+                access: FileSystemAccessMode::Write,
+                missing_path_behavior: None,
+            }]),
+            NetworkSandboxPolicy::Restricted,
+        ),
+        non_native_cwd(),
+    );
+
+    let (first, second) = tokio::join!(
+        file_system.get_metadata(&first_path, /*sandbox*/ None),
+        file_system.get_metadata(&second_path, /*sandbox*/ None),
+    );
+    first.expect("metadata for first path");
+    second.expect("metadata for second path");
+
+    let (first, second) = tokio::join!(
+        file_system.get_metadata(&first_path, Some(&sandbox)),
+        file_system.get_metadata(&first_path, Some(&sandbox)),
+    );
+    first.expect("first sandboxed metadata request");
+    second.expect("second sandboxed metadata request");
+
+    let captured_params = captured_params.await.expect("captured metadata requests");
+    assert_eq!(
+        captured_params
+            .iter()
+            .filter(|params| params.path == first_path && params.sandbox.is_none())
+            .count(),
+        1
+    );
+    assert_eq!(
+        captured_params
+            .iter()
+            .filter(|params| params.path == second_path && params.sandbox.is_none())
+            .count(),
+        1
+    );
+    assert_eq!(
+        captured_params
+            .iter()
+            .filter(|params| params.path == first_path && params.sandbox.as_ref() == Some(&sandbox))
+            .count(),
+        2
+    );
+    server.await.expect("metadata server should succeed");
+}
+
 async fn record_read_file_params(
    expected_requests: usize,
 ) -> (
@@ -137,6 +475,106 @@ async fn record_read_file_params(
    (websocket_url, captured_params_rx, server)
 }

+enum MetadataResponse {
+    Immediate(Result<FsGetMetadataResponse, JSONRPCErrorError>),
+    Abandoned(oneshot::Sender<()>),
+    Mutation(Result<(), JSONRPCErrorError>),
+}
+
+fn metadata_response(size: u64) -> FsGetMetadataResponse {
+    FsGetMetadataResponse {
+        is_directory: false,
+        is_file: true,
+        is_symlink: false,
+        size,
+        created_at_ms: 10,
+        modified_at_ms: 20,
+    }
+}
+
+#[derive(Clone, Copy)]
+enum MetadataMutation {
+    Write,
+    CreateDirectory,
+    Remove,
+    Copy,
+}
+
+async fn record_metadata_params(
+    responses: Vec<MetadataResponse>,
+) -> (
+    String,
+    oneshot::Receiver<Vec<FsGetMetadataParams>>,
+    tokio::task::JoinHandle<()>,
+) {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("listener should bind");
+    let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address"));
+    let (captured_params_tx, captured_params_rx) = oneshot::channel();
+    let server = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.expect("listener should accept");
+        let mut websocket = accept_async(stream)
+            .await
+            .expect("websocket handshake should succeed");
+        complete_websocket_initialize(&mut websocket).await;
+
+        let mut captured_params = Vec::with_capacity(responses.len());
+        for response in responses {
+            let request = match read_jsonrpc_websocket(&mut websocket).await {
+                JSONRPCMessage::Request(request) => request,
+                other => panic!("expected filesystem request, got {other:?}"),
+            };
+            if matches!(response, MetadataResponse::Mutation(_)) {
+                assert!(matches!(
+                    request.method.as_str(),
+                    FS_WRITE_FILE_METHOD
+                        | FS_CREATE_DIRECTORY_METHOD
+                        | FS_REMOVE_METHOD
+                        | FS_COPY_METHOD
+                ));
+            } else {
+                assert_eq!(request.method, FS_GET_METADATA_METHOD);
+                let params: FsGetMetadataParams = serde_json::from_value(
+                    request.params.expect("fs/getMetadata params should exist"),
+                )
+                .expect("fs/getMetadata params should deserialize");
+                captured_params.push(params);
+            }
+            let response = match response {
+                MetadataResponse::Immediate(Ok(metadata)) => {
+                    JSONRPCMessage::Response(JSONRPCResponse {
+                        id: request.id,
+                        result: serde_json::to_value(metadata)
+                            .expect("fs/getMetadata response should serialize"),
+                    })
+                }
+                MetadataResponse::Immediate(Err(error))
+                | MetadataResponse::Mutation(Err(error)) => JSONRPCMessage::Error(JSONRPCError {
+                    error,
+                    id: request.id,
+                }),
+                MetadataResponse::Mutation(Ok(())) => JSONRPCMessage::Response(JSONRPCResponse {
+                    id: request.id,
+                    result: serde_json::json!({}),
+                }),
+                MetadataResponse::Abandoned(received_tx) => {
+                    received_tx
+                        .send(())
+                        .expect("abandoned request observer should stay open");
+                    continue;
+                }
+            };
+            write_jsonrpc_websocket(&mut websocket, response).await;
+        }
+        captured_params_tx
+            .send(captured_params)
+            .expect("captured params receiver should stay open");
+    });
+
+    (websocket_url, captured_params_rx, server)
+}
+
 fn non_native_cwd() -> PathUri {
    #[cfg(unix)]
    let uri = "file://server/share/checkout";