diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index 18b209152d..dfde60b8c2 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -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>>; + pub(crate) struct RemoteFileSystem { client: LazyRemoteExecServerClient, + metadata_requests: Mutex>>, } 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 { + 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 { 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(()) } } diff --git a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs index 2f4f99ac45..cd0f5c1a2f 100644 --- a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs +++ b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs @@ -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), + 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, +) -> ( + String, + oneshot::Receiver>, + 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";