From 6920411925cbb6dd33a022f4422abee864919383 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 19 Jun 2026 11:59:23 +0200 Subject: [PATCH 1/7] Batch skill discovery filesystem reads --- codex-rs/core-skills/src/loader.rs | 365 ++++++++++++------ codex-rs/utils/plugins/src/lib.rs | 1 + .../utils/plugins/src/plugin_namespace.rs | 90 +++-- 3 files changed, 316 insertions(+), 140 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 2a3a3189be9b..51bf448396f3 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -13,14 +13,16 @@ use codex_config::default_project_root_markers; use codex_config::merge_toml_values; use codex_config::project_root_markers_from_config; use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FileSystemOperation; +use codex_exec_server::FileSystemOperationOutput; use codex_exec_server::LOCAL_FS; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBufGuard; use codex_utils_path_uri::PathUri; +use codex_utils_plugins::PluginNamespaceResolver; use codex_utils_plugins::PluginSkillRoot; -use codex_utils_plugins::plugin_namespace_for_skill_path; use dirs::home_dir; use serde::Deserialize; use std::collections::HashSet; @@ -68,6 +70,12 @@ struct LoadedSkillMetadata { policy: Option, } +struct PreloadedSkillFile { + contents: io::Result, + metadata_contents: io::Result, + resolved_path: AbsolutePathBuf, +} + #[derive(Debug, Default, Deserialize)] struct Interface { display_name: Option, @@ -523,99 +531,168 @@ async fn discover_skills_under_root( let mut queue: VecDeque<(AbsolutePathBuf, usize)> = VecDeque::from([(root.clone(), 0)]); let mut truncated_by_dir_limit = false; - - while let Some((dir, depth)) = queue.pop_front() { - let dir_uri = PathUri::from_abs_path(&dir); - let entries = match fs.read_directory(&dir_uri, /*sandbox*/ None).await { - Ok(entries) => entries, - Err(e) => { - error!("failed to read skills dir {}: {e:#}", dir.display()); - continue; + let mut skill_paths = Vec::new(); + + while !queue.is_empty() { + let dirs = queue.drain(..).collect::>(); + let directory_results = match fs + .execute_batch( + dirs.iter() + .map(|(dir, _)| FileSystemOperation::ReadDirectory { + path: PathUri::from_abs_path(dir), + }) + .collect(), + /*sandbox*/ None, + ) + .await + { + Ok(results) => results, + Err(err) => { + error!("failed to batch-read skills directories: {err:#}"); + break; } }; - for entry in entries { - let file_name = entry.file_name; - if file_name.starts_with('.') { - continue; - } - - let path = dir.join(&file_name); - let path_uri = PathUri::from_abs_path(&path); - let metadata = match fs.get_metadata(&path_uri, /*sandbox*/ None).await { - Ok(metadata) => metadata, - Err(e) => { - error!("failed to stat skills path {}: {e:#}", path.display()); + let mut candidates = Vec::new(); + for ((dir, depth), result) in dirs.into_iter().zip(directory_results) { + let entries = match result { + Ok(FileSystemOperationOutput::ReadDirectory(entries)) => entries, + Ok(_) => { + error!("filesystem returned the wrong result for {}", dir.display()); + continue; + } + Err(err) => { + error!("failed to read skills dir {}: {err:#}", dir.display()); continue; } }; + candidates.extend(entries.into_iter().filter_map(|entry| { + let file_name = entry.file_name; + if file_name.starts_with('.') { + return None; + } + Some((dir.join(&file_name), depth, file_name)) + })); + } - if metadata.is_symlink { - if !follow_symlinks { + let metadata_results = match fs + .execute_batch( + candidates + .iter() + .map(|(path, _, _)| FileSystemOperation::GetMetadata { + path: PathUri::from_abs_path(path), + }) + .collect(), + /*sandbox*/ None, + ) + .await + { + Ok(results) => results, + Err(err) => { + error!("failed to batch-stat skills paths: {err:#}"); + break; + } + }; + + let mut directories = Vec::new(); + for ((path, depth, file_name), result) in candidates.into_iter().zip(metadata_results) { + let metadata = match result { + Ok(FileSystemOperationOutput::GetMetadata(metadata)) => metadata, + Ok(_) => { + error!( + "filesystem returned the wrong result for {}", + path.display() + ); continue; } - match fs.read_directory(&path_uri, /*sandbox*/ None).await { - Ok(_) => { - let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; - enqueue_dir( - &mut queue, - &mut visited_dirs, - &mut truncated_by_dir_limit, - resolved_dir, - depth + 1, - ); - } - Err(err) - if matches!( - err.kind(), - io::ErrorKind::NotADirectory | io::ErrorKind::NotFound - ) => {} - Err(err) => { - error!( - "failed to read skills symlink dir {}: {err:#}", - path.display() - ); - } + Err(err) => { + error!("failed to stat skills path {}: {err:#}", path.display()); + continue; } - continue; + }; + if metadata.is_symlink { + if follow_symlinks && metadata.is_directory { + directories.push((path, depth + 1)); + } + } else if metadata.is_directory { + directories.push((path, depth + 1)); + } else if metadata.is_file && file_name == SKILLS_FILENAME { + skill_paths.push(path); } + } - if metadata.is_directory { - let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; - enqueue_dir( - &mut queue, - &mut visited_dirs, - &mut truncated_by_dir_limit, - resolved_dir, - depth + 1, - ); - continue; + let canonical_results = match fs + .execute_batch( + directories + .iter() + .map(|(path, _)| FileSystemOperation::Canonicalize { + path: PathUri::from_abs_path(path), + }) + .collect(), + /*sandbox*/ None, + ) + .await + { + Ok(results) => results, + Err(err) => { + error!("failed to batch-canonicalize skills directories: {err:#}"); + break; } - - if metadata.is_file && file_name == SKILLS_FILENAME { - match parse_skill_file( - fs, - &path, - scope, - plugin_id, - plugin_namespace, - plugin_root.as_ref(), - ) - .await - { - Ok(skill) => { - outcome.skills.push(skill); - } - Err(err) => { - if scope != SkillScope::System { - outcome.errors.push(SkillError { - path: path.clone(), - message: err.to_string(), - }); - } - } + }; + for ((path, depth), result) in directories.into_iter().zip(canonical_results) { + let resolved_dir = match result { + Ok(FileSystemOperationOutput::Canonicalize(resolved_path)) => { + resolved_path.to_abs_path().unwrap_or_else(|_| path.clone()) } - } + Ok(_) => { + error!( + "filesystem returned the wrong result for {}", + path.display() + ); + path + } + Err(_) => path, + }; + enqueue_dir( + &mut queue, + &mut visited_dirs, + &mut truncated_by_dir_limit, + resolved_dir, + depth, + ); + } + } + + let mut plugin_namespace_resolver = PluginNamespaceResolver::default(); + if plugin_namespace.is_none() { + plugin_namespace_resolver.prime(fs, &skill_paths).await; + } + let preloaded_skills = match preload_skill_files(fs, &skill_paths).await { + Ok(preloaded_skills) => preloaded_skills, + Err(err) => { + error!("failed to batch-read skill files: {err:#}"); + Vec::new() + } + }; + for (path, preloaded) in skill_paths.into_iter().zip(preloaded_skills) { + match parse_skill_file( + fs, + &path, + scope, + plugin_id, + plugin_namespace, + plugin_root.as_ref(), + &mut plugin_namespace_resolver, + preloaded, + ) + .await + { + Ok(skill) => outcome.skills.push(skill), + Err(err) if scope != SkillScope::System => outcome.errors.push(SkillError { + path, + message: err.to_string(), + }), + Err(_) => {} } } @@ -628,6 +705,79 @@ async fn discover_skills_under_root( } } +async fn preload_skill_files( + fs: &dyn ExecutorFileSystem, + skill_paths: &[AbsolutePathBuf], +) -> io::Result> { + let operations = skill_paths + .iter() + .flat_map(|path| { + let metadata_path = path + .parent() + .map(|parent| { + parent + .join(SKILLS_METADATA_DIR) + .join(SKILLS_METADATA_FILENAME) + }) + .unwrap_or_else(|| path.clone()); + [ + FileSystemOperation::ReadFile { + path: PathUri::from_abs_path(path), + }, + FileSystemOperation::ReadFile { + path: PathUri::from_abs_path(&metadata_path), + }, + FileSystemOperation::Canonicalize { + path: PathUri::from_abs_path(path), + }, + ] + }) + .collect(); + let mut results = fs + .execute_batch(operations, /*sandbox*/ None) + .await? + .into_iter(); + + Ok(skill_paths + .iter() + .map(|path| { + let contents = decode_text_operation_result(results.next(), "read skill file"); + let metadata_contents = + decode_text_operation_result(results.next(), "read skill metadata"); + let resolved_path = match results.next() { + Some(Ok(FileSystemOperationOutput::Canonicalize(resolved_path))) => { + resolved_path.to_abs_path().unwrap_or_else(|_| path.clone()) + } + Some(Ok(_)) | Some(Err(_)) | None => path.clone(), + }; + PreloadedSkillFile { + contents, + metadata_contents, + resolved_path, + } + }) + .collect()) +} + +fn decode_text_operation_result( + result: Option>, + operation: &str, +) -> io::Result { + match result { + Some(Ok(FileSystemOperationOutput::ReadFile(contents))) => String::from_utf8(contents) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)), + Some(Ok(_)) => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("filesystem returned the wrong result for {operation}"), + )), + Some(Err(error)) => Err(error), + None => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("filesystem returned no result for {operation}"), + )), + } +} + async fn parse_skill_file( fs: &dyn ExecutorFileSystem, path: &AbsolutePathBuf, @@ -635,12 +785,15 @@ async fn parse_skill_file( plugin_id: Option<&str>, plugin_namespace: Option<&str>, plugin_root: Option<&AbsolutePathBuf>, + plugin_namespace_resolver: &mut PluginNamespaceResolver, + preloaded: PreloadedSkillFile, ) -> Result { - let path_uri = PathUri::from_abs_path(path); - let contents = fs - .read_file_text(&path_uri, /*sandbox*/ None) - .await - .map_err(SkillParseError::Read)?; + let PreloadedSkillFile { + contents, + metadata_contents, + resolved_path, + } = preloaded; + let contents = contents.map_err(SkillParseError::Read)?; let frontmatter = extract_frontmatter(&contents).ok_or(SkillParseError::MissingFrontmatter)?; @@ -664,7 +817,14 @@ async fn parse_skill_file( .map(sanitize_single_line) .filter(|value| !value.is_empty()) .unwrap_or_else(|| default_skill_name(path)); - let name = namespaced_skill_name(fs, path, &base_name, plugin_namespace).await; + let name = namespaced_skill_name( + fs, + path, + &base_name, + plugin_namespace, + plugin_namespace_resolver, + ) + .await; let description = parsed .description .as_deref() @@ -680,7 +840,7 @@ async fn parse_skill_file( interface, dependencies, policy, - } = load_skill_metadata(fs, path, plugin_root).await; + } = load_skill_metadata(path, plugin_root, metadata_contents); validate_len(&base_name, MAX_NAME_LEN, "name")?; validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name")?; @@ -693,8 +853,6 @@ async fn parse_skill_file( )?; } - let resolved_path = canonicalize_for_skill_identity(fs, path).await; - Ok(SkillMetadata { name, description, @@ -725,20 +883,22 @@ async fn namespaced_skill_name( path: &AbsolutePathBuf, base_name: &str, plugin_namespace: Option<&str>, + plugin_namespace_resolver: &mut PluginNamespaceResolver, ) -> String { if let Some(plugin_namespace) = plugin_namespace { return format!("{plugin_namespace}:{base_name}"); } - plugin_namespace_for_skill_path(fs, path) + plugin_namespace_resolver + .resolve(fs, path) .await .map(|namespace| format!("{namespace}:{base_name}")) .unwrap_or_else(|| base_name.to_string()) } -async fn load_skill_metadata( - fs: &dyn ExecutorFileSystem, +fn load_skill_metadata( skill_path: &AbsolutePathBuf, plugin_root: Option<&AbsolutePathBuf>, + contents: io::Result, ) -> LoadedSkillMetadata { // Fail open: optional metadata should not block loading SKILL.md. let Some(skill_dir) = skill_path.parent() else { @@ -747,28 +907,11 @@ async fn load_skill_metadata( let metadata_path = skill_dir .join(SKILLS_METADATA_DIR) .join(SKILLS_METADATA_FILENAME); - let metadata_path_uri = PathUri::from_abs_path(&metadata_path); - match fs.get_metadata(&metadata_path_uri, /*sandbox*/ None).await { - Ok(metadata) if metadata.is_file => {} - Ok(_) => return LoadedSkillMetadata::default(), + let contents = match contents { + Ok(contents) => contents, Err(error) if error.kind() == io::ErrorKind::NotFound => { return LoadedSkillMetadata::default(); } - Err(error) => { - tracing::warn!( - "ignoring {path}: failed to stat {label}: {error}", - path = metadata_path.display(), - label = SKILLS_METADATA_FILENAME - ); - return LoadedSkillMetadata::default(); - } - } - - let contents = match fs - .read_file_text(&metadata_path_uri, /*sandbox*/ None) - .await - { - Ok(contents) => contents, Err(error) => { tracing::warn!( "ignoring {path}: failed to read {label}: {error}", diff --git a/codex-rs/utils/plugins/src/lib.rs b/codex-rs/utils/plugins/src/lib.rs index 203b1f22e41d..238d11f7cb41 100644 --- a/codex-rs/utils/plugins/src/lib.rs +++ b/codex-rs/utils/plugins/src/lib.rs @@ -8,6 +8,7 @@ pub mod mention_syntax; pub mod plugin_namespace; pub use plugin_namespace::DISCOVERABLE_PLUGIN_MANIFEST_PATHS; +pub use plugin_namespace::PluginNamespaceResolver; pub use plugin_namespace::find_plugin_manifest_path; pub use plugin_namespace::plugin_namespace_for_skill_path; diff --git a/codex-rs/utils/plugins/src/plugin_namespace.rs b/codex-rs/utils/plugins/src/plugin_namespace.rs index b29ff5c5d684..7ed49148daf1 100644 --- a/codex-rs/utils/plugins/src/plugin_namespace.rs +++ b/codex-rs/utils/plugins/src/plugin_namespace.rs @@ -1,8 +1,12 @@ //! Resolve plugin namespace from skill file paths by walking ancestors for `plugin.json`. use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FileSystemOperation; +use codex_exec_server::FileSystemOperationOutput; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; +use std::collections::HashMap; +use std::collections::HashSet; use std::path::Path; use std::path::PathBuf; @@ -24,29 +28,8 @@ struct RawPluginManifestName { name: String, } -async fn plugin_manifest_name( - fs: &dyn ExecutorFileSystem, - plugin_root: &AbsolutePathBuf, -) -> Option { - let mut manifest_path = None; - for relative_path in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { - let candidate = plugin_root.join(relative_path); - let candidate_uri = PathUri::from_abs_path(&candidate); - match fs.get_metadata(&candidate_uri, /*sandbox*/ None).await { - Ok(metadata) if metadata.is_file => { - manifest_path = Some(candidate); - break; - } - Ok(_) | Err(_) => {} - } - } - let manifest_path = manifest_path?; - let manifest_path_uri = PathUri::from_abs_path(&manifest_path); - let contents = fs - .read_file_text(&manifest_path_uri, /*sandbox*/ None) - .await - .ok()?; - let RawPluginManifestName { name: raw_name } = serde_json::from_str(&contents).ok()?; +fn plugin_manifest_name(plugin_root: &AbsolutePathBuf, contents: &[u8]) -> Option { + let RawPluginManifestName { name: raw_name } = serde_json::from_slice(contents).ok()?; Some( plugin_root .file_name() @@ -57,18 +40,67 @@ async fn plugin_manifest_name( ) } +/// Resolves plugin namespaces while caching and batching ancestor manifest probes. +#[derive(Default)] +pub struct PluginNamespaceResolver { + manifests_by_root: HashMap>, +} + +impl PluginNamespaceResolver { + pub async fn prime(&mut self, fs: &dyn ExecutorFileSystem, paths: &[AbsolutePathBuf]) { + let mut seen = HashSet::new(); + let unresolved_roots = paths + .iter() + .flat_map(AbsolutePathBuf::ancestors) + .filter(|ancestor| !self.manifests_by_root.contains_key(ancestor)) + .filter(|ancestor| seen.insert(ancestor.clone())) + .collect::>(); + let operations = unresolved_roots + .iter() + .flat_map(|root| { + DISCOVERABLE_PLUGIN_MANIFEST_PATHS + .iter() + .map(|relative_path| FileSystemOperation::ReadFile { + path: PathUri::from_abs_path(&root.join(relative_path)), + }) + }) + .collect(); + let mut results = match fs.execute_batch(operations, /*sandbox*/ None).await { + Ok(results) => results.into_iter(), + Err(_) => Vec::new().into_iter(), + }; + for root in unresolved_roots { + let mut name = None; + for _ in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { + if let Some(Ok(FileSystemOperationOutput::ReadFile(contents))) = results.next() + && name.is_none() + { + name = plugin_manifest_name(&root, &contents); + } + } + self.manifests_by_root.insert(root, name); + } + } + + pub async fn resolve( + &mut self, + fs: &dyn ExecutorFileSystem, + path: &AbsolutePathBuf, + ) -> Option { + self.prime(fs, std::slice::from_ref(path)).await; + + path.ancestors() + .find_map(|ancestor| self.manifests_by_root.get(&ancestor).cloned().flatten()) + } +} + /// Returns the plugin manifest `name` for the nearest ancestor of `path` that contains a valid /// plugin manifest (same `name` rules as full manifest loading in codex-core). pub async fn plugin_namespace_for_skill_path( fs: &dyn ExecutorFileSystem, path: &AbsolutePathBuf, ) -> Option { - for ancestor in path.ancestors() { - if let Some(name) = plugin_manifest_name(fs, &ancestor).await { - return Some(name); - } - } - None + PluginNamespaceResolver::default().resolve(fs, path).await } #[cfg(test)] From f3c0ccae9611f8a1759dd7918408dfed33470f99 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 19 Jun 2026 13:28:54 +0200 Subject: [PATCH 2/7] Preserve plugin manifest precedence --- codex-rs/core-skills/src/loader.rs | 61 +++------- .../utils/plugins/src/plugin_namespace.rs | 111 +++++++++++++----- 2 files changed, 98 insertions(+), 74 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 51bf448396f3..58641d4bd882 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -26,7 +26,6 @@ use codex_utils_plugins::PluginSkillRoot; use dirs::home_dir; use serde::Deserialize; use std::collections::HashSet; -use std::collections::VecDeque; use std::error::Error; use std::fmt; use std::io; @@ -502,7 +501,7 @@ async fn discover_skills_under_root( } fn enqueue_dir( - queue: &mut VecDeque<(AbsolutePathBuf, usize)>, + queue: &mut Vec<(AbsolutePathBuf, usize)>, visited_dirs: &mut HashSet, truncated_by_dir_limit: &mut bool, path: AbsolutePathBuf, @@ -516,7 +515,7 @@ async fn discover_skills_under_root( return; } if visited_dirs.insert(path.clone()) { - queue.push_back((path, depth)); + queue.push((path, depth)); } } @@ -529,12 +528,12 @@ async fn discover_skills_under_root( let mut visited_dirs: HashSet = HashSet::new(); visited_dirs.insert(root.clone()); - let mut queue: VecDeque<(AbsolutePathBuf, usize)> = VecDeque::from([(root.clone(), 0)]); + let mut queue = vec![(root.clone(), 0)]; let mut truncated_by_dir_limit = false; let mut skill_paths = Vec::new(); while !queue.is_empty() { - let dirs = queue.drain(..).collect::>(); + let dirs = std::mem::take(&mut queue); let directory_results = match fs .execute_batch( dirs.iter() @@ -663,10 +662,11 @@ async fn discover_skills_under_root( } } - let mut plugin_namespace_resolver = PluginNamespaceResolver::default(); - if plugin_namespace.is_none() { - plugin_namespace_resolver.prime(fs, &skill_paths).await; - } + let plugin_namespace_resolver = if plugin_namespace.is_none() { + Some(PluginNamespaceResolver::load(fs, &skill_paths).await) + } else { + None + }; let preloaded_skills = match preload_skill_files(fs, &skill_paths).await { Ok(preloaded_skills) => preloaded_skills, Err(err) => { @@ -675,18 +675,19 @@ async fn discover_skills_under_root( } }; for (path, preloaded) in skill_paths.into_iter().zip(preloaded_skills) { + let plugin_namespace = plugin_namespace.or_else(|| { + plugin_namespace_resolver + .as_ref() + .and_then(|resolver| resolver.resolve(&path)) + }); match parse_skill_file( - fs, &path, scope, plugin_id, plugin_namespace, plugin_root.as_ref(), - &mut plugin_namespace_resolver, preloaded, - ) - .await - { + ) { Ok(skill) => outcome.skills.push(skill), Err(err) if scope != SkillScope::System => outcome.errors.push(SkillError { path, @@ -778,14 +779,12 @@ fn decode_text_operation_result( } } -async fn parse_skill_file( - fs: &dyn ExecutorFileSystem, +fn parse_skill_file( path: &AbsolutePathBuf, scope: SkillScope, plugin_id: Option<&str>, plugin_namespace: Option<&str>, plugin_root: Option<&AbsolutePathBuf>, - plugin_namespace_resolver: &mut PluginNamespaceResolver, preloaded: PreloadedSkillFile, ) -> Result { let PreloadedSkillFile { @@ -817,14 +816,9 @@ async fn parse_skill_file( .map(sanitize_single_line) .filter(|value| !value.is_empty()) .unwrap_or_else(|| default_skill_name(path)); - let name = namespaced_skill_name( - fs, - path, - &base_name, - plugin_namespace, - plugin_namespace_resolver, - ) - .await; + let name = plugin_namespace + .map(|namespace| format!("{namespace}:{base_name}")) + .unwrap_or_else(|| base_name.clone()); let description = parsed .description .as_deref() @@ -878,23 +872,6 @@ fn default_skill_name(path: &AbsolutePathBuf) -> String { .unwrap_or_else(|| "skill".to_string()) } -async fn namespaced_skill_name( - fs: &dyn ExecutorFileSystem, - path: &AbsolutePathBuf, - base_name: &str, - plugin_namespace: Option<&str>, - plugin_namespace_resolver: &mut PluginNamespaceResolver, -) -> String { - if let Some(plugin_namespace) = plugin_namespace { - return format!("{plugin_namespace}:{base_name}"); - } - plugin_namespace_resolver - .resolve(fs, path) - .await - .map(|namespace| format!("{namespace}:{base_name}")) - .unwrap_or_else(|| base_name.to_string()) -} - fn load_skill_metadata( skill_path: &AbsolutePathBuf, plugin_root: Option<&AbsolutePathBuf>, diff --git a/codex-rs/utils/plugins/src/plugin_namespace.rs b/codex-rs/utils/plugins/src/plugin_namespace.rs index 7ed49148daf1..c829d4c26d37 100644 --- a/codex-rs/utils/plugins/src/plugin_namespace.rs +++ b/codex-rs/utils/plugins/src/plugin_namespace.rs @@ -41,56 +41,71 @@ fn plugin_manifest_name(plugin_root: &AbsolutePathBuf, contents: &[u8]) -> Optio } /// Resolves plugin namespaces while caching and batching ancestor manifest probes. -#[derive(Default)] pub struct PluginNamespaceResolver { - manifests_by_root: HashMap>, + manifests_by_root: HashMap, } impl PluginNamespaceResolver { - pub async fn prime(&mut self, fs: &dyn ExecutorFileSystem, paths: &[AbsolutePathBuf]) { + /// Loads plugin manifest names for every ancestor of `paths`. + pub async fn load(fs: &dyn ExecutorFileSystem, paths: &[AbsolutePathBuf]) -> Self { let mut seen = HashSet::new(); - let unresolved_roots = paths + let roots = paths .iter() .flat_map(AbsolutePathBuf::ancestors) - .filter(|ancestor| !self.manifests_by_root.contains_key(ancestor)) .filter(|ancestor| seen.insert(ancestor.clone())) .collect::>(); - let operations = unresolved_roots + let operations = roots .iter() .flat_map(|root| { DISCOVERABLE_PLUGIN_MANIFEST_PATHS .iter() - .map(|relative_path| FileSystemOperation::ReadFile { - path: PathUri::from_abs_path(&root.join(relative_path)), + .flat_map(|relative_path| { + let path = PathUri::from_abs_path(&root.join(relative_path)); + [ + FileSystemOperation::GetMetadata { path: path.clone() }, + FileSystemOperation::ReadFile { path }, + ] }) }) .collect(); - let mut results = match fs.execute_batch(operations, /*sandbox*/ None).await { - Ok(results) => results.into_iter(), - Err(_) => Vec::new().into_iter(), - }; - for root in unresolved_roots { - let mut name = None; - for _ in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { - if let Some(Ok(FileSystemOperationOutput::ReadFile(contents))) = results.next() - && name.is_none() - { - name = plugin_manifest_name(&root, &contents); + let mut results = fs + .execute_batch(operations, /*sandbox*/ None) + .await + .unwrap_or_default() + .into_iter(); + let manifests_by_root = roots + .into_iter() + .filter_map(|root| { + let mut name = None; + let mut selected = false; + for _ in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { + let metadata_result = results.next(); + let contents_result = results.next(); + let is_file = matches!( + metadata_result, + Some(Ok(FileSystemOperationOutput::GetMetadata(metadata))) + if metadata.is_file + ); + if !selected && is_file { + selected = true; + if let Some(Ok(FileSystemOperationOutput::ReadFile(contents))) = + contents_result + { + name = plugin_manifest_name(&root, &contents); + } + } } - } - self.manifests_by_root.insert(root, name); - } - } + name.map(|name| (root, name)) + }) + .collect(); - pub async fn resolve( - &mut self, - fs: &dyn ExecutorFileSystem, - path: &AbsolutePathBuf, - ) -> Option { - self.prime(fs, std::slice::from_ref(path)).await; + Self { manifests_by_root } + } + /// Returns the nearest preloaded plugin namespace for `path`. + pub fn resolve(&self, path: &AbsolutePathBuf) -> Option<&str> { path.ancestors() - .find_map(|ancestor| self.manifests_by_root.get(&ancestor).cloned().flatten()) + .find_map(|ancestor| self.manifests_by_root.get(&ancestor).map(String::as_str)) } } @@ -100,7 +115,10 @@ pub async fn plugin_namespace_for_skill_path( fs: &dyn ExecutorFileSystem, path: &AbsolutePathBuf, ) -> Option { - PluginNamespaceResolver::default().resolve(fs, path).await + PluginNamespaceResolver::load(fs, std::slice::from_ref(path)) + .await + .resolve(path) + .map(str::to_string) } #[cfg(test)] @@ -112,6 +130,7 @@ mod tests { use std::fs; use tempfile::tempdir; + const PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH: &str = ".codex-plugin/plugin.json"; const ALTERNATE_PLUGIN_MANIFEST_RELATIVE_PATH: &str = ".claude-plugin/plugin.json"; #[tokio::test] @@ -123,7 +142,7 @@ mod tests { fs::create_dir_all(skill_path.parent().expect("parent")).expect("mkdir"); fs::create_dir_all(plugin_root.join(".codex-plugin")).expect("mkdir manifest"); fs::write( - plugin_root.join(".codex-plugin/plugin.json"), + plugin_root.join(PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH), r#"{"name":"sample"}"#, ) .expect("write manifest"); @@ -154,4 +173,32 @@ mod tests { ); assert_eq!(find_plugin_manifest_path(&plugin_root), Some(manifest_path)); } + + #[tokio::test] + async fn invalid_primary_manifest_does_not_fall_back_to_alternate() { + let tmp = tempdir().expect("tempdir"); + let plugin_root = tmp.path().join("plugins/sample"); + let skill_path = plugin_root.join("skills/search/SKILL.md"); + let primary_manifest_path = plugin_root.join(PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH); + let alternate_manifest_path = plugin_root.join(ALTERNATE_PLUGIN_MANIFEST_RELATIVE_PATH); + + fs::create_dir_all(skill_path.parent().expect("parent")).expect("mkdir"); + fs::create_dir_all(primary_manifest_path.parent().expect("manifest parent")) + .expect("mkdir primary manifest"); + fs::create_dir_all(alternate_manifest_path.parent().expect("manifest parent")) + .expect("mkdir alternate manifest"); + fs::write(&primary_manifest_path, "invalid json").expect("write primary manifest"); + fs::write(&alternate_manifest_path, r#"{"name":"sample"}"#) + .expect("write alternate manifest"); + fs::write(&skill_path, "---\ndescription: search\n---\n").expect("write skill"); + + assert_eq!( + plugin_namespace_for_skill_path(LOCAL_FS.as_ref(), &skill_path.abs()).await, + None + ); + assert_eq!( + find_plugin_manifest_path(&plugin_root), + Some(primary_manifest_path) + ); + } } From baa2809950383c47694f87e33fb10a45fed5ab58 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 19 Jun 2026 18:30:54 +0200 Subject: [PATCH 3/7] Use generic RPC batch requests --- codex-rs/Cargo.lock | 1 + codex-rs/exec-server/src/client.rs | 25 ++ codex-rs/exec-server/src/connection.rs | 48 ++- codex-rs/exec-server/src/lib.rs | 5 +- .../src/noise_relay/executor_stream.rs | 2 +- .../src/noise_relay/executor_stream_tests.rs | 2 +- .../exec-server/src/noise_relay/harness.rs | 2 +- .../src/noise_relay/message_framing.rs | 9 +- .../src/noise_relay/message_framing_tests.rs | 5 +- codex-rs/exec-server/src/relay.rs | 18 +- .../exec-server/src/remote_file_system.rs | 48 +++ codex-rs/exec-server/src/rpc.rs | 154 ++++++-- codex-rs/exec-server/src/server/processor.rs | 333 ++++++++++-------- codex-rs/exec-server/src/server/registry.rs | 8 +- .../exec-server/tests/file_system/shared.rs | 51 --- codex-rs/file-system/Cargo.toml | 1 + codex-rs/file-system/src/lib.rs | 91 +---- 17 files changed, 458 insertions(+), 345 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 55d02cd8b155..1036f4815f75 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3049,6 +3049,7 @@ dependencies = [ "codex-utils-path-uri", "futures", "serde", + "serde_json", ] [[package]] diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 88c0c8034288..66a5708f4406 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -97,6 +97,7 @@ use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; +use crate::rpc::RpcBatchCall; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; @@ -761,6 +762,29 @@ impl ExecServerClient { self.call_rpc(&rpc_client, method, params).await } + pub(crate) async fn call_batch( + &self, + calls: Vec, + ) -> Result>, ExecServerError> { + let rpc_client = self.inner.rpc_client().await?; + match rpc_client.call_batch(calls).await { + Ok(results) => Ok(results + .into_iter() + .map(|result| result.map_err(ExecServerError::from)) + .collect()), + Err(error) => { + let error = ExecServerError::from(error); + if is_transport_closed_error(&error) { + Err(ExecServerError::Disconnected(disconnected_message( + /*reason*/ None, + ))) + } else { + Err(error) + } + } + } + } + async fn call_rpc( &self, rpc_client: &Arc, @@ -813,6 +837,7 @@ impl From for ExecServerError { code: error.code, message: error.message, }, + RpcCallError::InvalidBatch(message) => Self::Protocol(message), } } } diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 2d2e5afe0641..9a3da73e80d8 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -12,6 +12,8 @@ use futures::Sink; use futures::SinkExt; use futures::Stream; use futures::StreamExt; +use serde::Deserialize; +use serde::Serialize; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::process::Child; @@ -29,6 +31,7 @@ use tokio::io::BufReader; use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; +pub(crate) const MAX_RPC_BATCH_REQUESTS: usize = 512; const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2); #[cfg(test)] pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25); @@ -38,10 +41,33 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30 #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { Message(JSONRPCMessage), + Batch(Vec), MalformedMessage { reason: String }, Disconnected { reason: Option }, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum JsonRpcWireMessage { + Single(JSONRPCMessage), + Batch(Vec), +} + +impl JsonRpcWireMessage { + pub(crate) fn into_connection_event(self) -> JsonRpcConnectionEvent { + match self { + Self::Single(message) => JsonRpcConnectionEvent::Message(message), + Self::Batch(messages) => JsonRpcConnectionEvent::Batch(messages), + } + } +} + +impl From for JsonRpcWireMessage { + fn from(message: JSONRPCMessage) -> Self { + Self::Single(message) + } +} + #[derive(Clone)] pub(crate) enum JsonRpcTransport { // Plain means no child process; transport bytes may still be encrypted. @@ -221,7 +247,7 @@ fn log_stdio_child_wait_result(result: std::io::Result } pub(crate) struct JsonRpcConnection { - pub(crate) outgoing_tx: mpsc::Sender, + pub(crate) outgoing_tx: mpsc::Sender, pub(crate) incoming_rx: mpsc::Receiver, pub(crate) disconnected_rx: watch::Receiver, pub(crate) task_handles: Vec>, @@ -249,10 +275,10 @@ impl JsonRpcConnection { if line.trim().is_empty() { continue; } - match serde_json::from_str::(&line) { + match serde_json::from_str::(&line) { Ok(message) => { if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -395,7 +421,7 @@ impl JsonRpcConnection { Some(Ok(message)) => match message.parse_jsonrpc_frame() { Ok(JsonRpcWebSocketFrame::Message(message)) => { if incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -464,7 +490,7 @@ impl JsonRpcConnection { } enum JsonRpcWebSocketFrame { - Message(JSONRPCMessage), + Message(JsonRpcWireMessage), Close, Ignore, } @@ -549,7 +575,7 @@ async fn send_malformed_message( async fn write_jsonrpc_line_message( writer: &mut BufWriter, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> std::io::Result<()> where W: AsyncWrite + Unpin, @@ -564,7 +590,7 @@ where async fn send_websocket_jsonrpc_message( websocket_writer: &mut W, connection_label: &str, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> Result<(), String> where W: Sink + Unpin, @@ -584,7 +610,7 @@ where } } -fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { +fn serialize_jsonrpc_message(message: &JsonRpcWireMessage) -> Result { serde_json::to_string(message) } @@ -694,7 +720,7 @@ mod tests { ); let message = test_jsonrpc_message(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -706,7 +732,9 @@ mod tests { control.set_write_ready(); assert!(matches!( timeout(Duration::from_secs(1), outbound_rx.next()).await?, - Some(Message::Text(text)) if serde_json::from_str::(&text)? == message + Some(Message::Text(text)) + if serde_json::from_str::(&text)? + == JsonRpcWireMessage::Single(message) )); drop(connection); Ok(()) diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index d66d2f9b7e68..edcd7f43e1cc 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,11 +42,10 @@ pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; pub use codex_file_system::ExecutorFileSystem; pub use codex_file_system::ExecutorFileSystemFuture; +pub use codex_file_system::ExecutorRpcBatchCall; +pub use codex_file_system::ExecutorRpcBatchResult; pub use codex_file_system::FILE_READ_CHUNK_SIZE; pub use codex_file_system::FileMetadata; -pub use codex_file_system::FileSystemOperation; -pub use codex_file_system::FileSystemOperationOutput; -pub use codex_file_system::FileSystemOperationResult; pub use codex_file_system::FileSystemReadStream; pub use codex_file_system::FileSystemResult; pub use codex_file_system::FileSystemSandboxContext; diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 57a05579dcf5..04f25dc4462b 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -75,7 +75,7 @@ impl NoiseVirtualStream { }; for message in self.inbound_decoder.push(&plaintext)? { self.incoming_tx - .try_send(JsonRpcConnectionEvent::Message(message)) + .try_send(message.into_connection_event()) .map_err(|_| { ExecServerError::Protocol( "Noise virtual stream inbound queue is full or closed".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs index 9119236cfd28..d757a12393ef 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs @@ -51,7 +51,7 @@ async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { id: RequestId::Integer(1), result: serde_json::Value::Null, }); - let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message)?)?; + let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message.clone().into())?)?; stream.receive_data(RelayData { seq: 0, segment_index: 0, diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 1989525224f6..62e82e608c1a 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -420,7 +420,7 @@ async fn receive_data( // messages; emit only complete, successfully parsed messages. for message in decoder.push(&plaintext)? { incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .map_err(|_| ExecServerError::Closed)?; } diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index 24ee48173e2e..b84253695c07 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -1,6 +1,5 @@ -use codex_app_server_protocol::JSONRPCMessage; - use crate::ExecServerError; +use crate::connection::JsonRpcWireMessage; const LENGTH_PREFIX_BYTES: usize = size_of::(); const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024; @@ -12,7 +11,9 @@ pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024; /// exec-server responses can be much larger. A four-byte authenticated length /// prefix lets the caller split this byte stream into bounded Noise records and /// lets the receiver reconstruct exact JSON-RPC message boundaries. -pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn frame_jsonrpc_message( + message: &JsonRpcWireMessage, +) -> Result, ExecServerError> { let mut framed = vec![0; LENGTH_PREFIX_BYTES]; serde_json::to_writer(&mut framed, message)?; let message_len = framed.len() - LENGTH_PREFIX_BYTES; @@ -39,7 +40,7 @@ impl JsonRpcMessageDecoder { pub(crate) fn push( &mut self, plaintext_record: &[u8], - ) -> Result, ExecServerError> { + ) -> Result, ExecServerError> { if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN { return Err(ExecServerError::Protocol( "Noise relay plaintext record exceeds maximum length".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs index c3df14bf5622..5cc3ed10f1c1 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs @@ -7,6 +7,7 @@ use super::MAX_NOISE_JSONRPC_MESSAGE_LEN; use super::NOISE_RECORD_PLAINTEXT_LEN; use super::frame_jsonrpc_message; use crate::ExecServerError; +use crate::connection::JsonRpcWireMessage; #[test] fn fragments_and_reassembles_large_jsonrpc_message() { @@ -16,7 +17,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { "data": "x".repeat(128 * 1024), })), }); - let framed = frame_jsonrpc_message(&message).unwrap(); + let framed = frame_jsonrpc_message(&message.clone().into()).unwrap(); assert!(framed.len() > 128 * 1024); let mut decoder = JsonRpcMessageDecoder::default(); @@ -25,7 +26,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { decoded.extend(decoder.push(record).unwrap()); } - assert_eq!(decoded, vec![message]); + assert_eq!(decoded, vec![JsonRpcWireMessage::Single(message)]); } #[test] diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 033da4b2a746..a00707c67779 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use std::time::Duration; -use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; use futures::SinkExt; use futures::Stream; @@ -25,6 +24,7 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::JsonRpcWireMessage; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; @@ -170,7 +170,7 @@ impl RelayMessageFrame { } } - fn into_jsonrpc_message(self) -> Result { + fn into_jsonrpc_message(self) -> Result { let payload = self.into_data()?.payload; serde_json::from_slice(&payload).map_err(ExecServerError::Json) } @@ -211,7 +211,7 @@ pub(crate) fn decode_relay_message_frame( .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) } -pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn jsonrpc_payload(message: &JsonRpcWireMessage) -> Result, ExecServerError> { serde_json::to_vec(message).map_err(ExecServerError::Json) } @@ -346,7 +346,7 @@ where &mut websocket, &mut keepalive, &incoming_tx, - JsonRpcConnectionEvent::Message(message), + message.into_connection_event(), ) .await { @@ -847,6 +847,7 @@ mod tests { use std::task::Poll; use std::time::Duration; + use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; use futures::Sink; @@ -880,7 +881,7 @@ mod tests { encode_relay_message_frame(&RelayMessageFrame::data( stream_id, /*seq*/ 0, - jsonrpc_payload(&message)?, + jsonrpc_payload(&message.clone().into())?, )) .into(), )) @@ -1029,7 +1030,7 @@ mod tests { let message = test_jsonrpc_message(); control.set_write_blocked(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -1047,7 +1048,10 @@ mod tests { }; let frame = decode_relay_message_frame(data_payload.as_ref())?; assert_eq!(frame.stream_id, stream_id); - assert_eq!(frame.into_jsonrpc_message()?, message); + assert_eq!( + frame.into_jsonrpc_message()?, + JsonRpcWireMessage::Single(message) + ); drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index f9136e146108..5950e919f125 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -9,6 +9,8 @@ use crate::CreateDirectoryOptions; use crate::ExecServerError; use crate::ExecutorFileSystem; use crate::ExecutorFileSystemFuture; +use crate::ExecutorRpcBatchCall; +use crate::ExecutorRpcBatchResult; use crate::FileMetadata; use crate::FileSystemReadStream; use crate::FileSystemResult; @@ -16,6 +18,7 @@ use crate::FileSystemSandboxContext; use crate::ReadDirectoryEntry; use crate::RemoveOptions; use crate::client::LazyRemoteExecServerClient; +use crate::connection::MAX_RPC_BATCH_REQUESTS; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; @@ -24,6 +27,7 @@ use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadFileParams; use crate::protocol::FsRemoveParams; use crate::protocol::FsWriteFileParams; +use crate::rpc::RpcBatchCall; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const NOT_FOUND_ERROR_CODE: i64 = -32004; @@ -223,6 +227,43 @@ impl RemoteFileSystem { .map_err(map_remote_error)?; Ok(()) } + + async fn execute_rpc_batch( + &self, + calls: Vec, + ) -> FileSystemResult> { + if calls.is_empty() { + return Ok(Vec::new()); + } + + let client = self.client.get().await.map_err(map_remote_error)?; + let mut decoded = Vec::with_capacity(calls.len()); + let mut calls = calls.into_iter(); + loop { + let call_chunk = calls + .by_ref() + .take(MAX_RPC_BATCH_REQUESTS) + .map(|call| RpcBatchCall { + method: call.method, + params: call.params, + }) + .collect::>(); + if call_chunk.is_empty() { + break; + } + + let results = client + .call_batch(call_chunk) + .await + .map_err(map_remote_error)?; + decoded.extend( + results + .into_iter() + .map(|result| result.map_err(map_remote_error)), + ); + } + Ok(decoded) + } } impl ExecutorFileSystem for RemoteFileSystem { @@ -310,6 +351,13 @@ impl ExecutorFileSystem for RemoteFileSystem { sandbox, )) } + + fn execute_rpc_batch<'a>( + &'a self, + calls: Vec, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(RemoteFileSystem::execute_rpc_batch(self, calls)) + } } fn remote_sandbox_context( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index c2a5124541d3..b16a3215110e 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::collections::HashSet; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -26,6 +25,8 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::JsonRpcWireMessage; +use crate::connection::MAX_RPC_BATCH_REQUESTS; pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; @@ -37,6 +38,13 @@ pub(crate) enum RpcCallError { Json(serde_json::Error), /// The executor returned a JSON-RPC error response for this call. Server(JSONRPCErrorError), + /// The caller attempted to construct a batch that cannot be represented safely. + InvalidBatch(String), +} + +pub(crate) struct RpcBatchCall { + pub(crate) method: String, + pub(crate) params: Value, } type PendingRequest = oneshot::Sender>; @@ -64,6 +72,7 @@ pub(crate) enum RpcServerOutboundMessage { error: JSONRPCErrorError, }, Notification(JSONRPCNotification), + Batch(Vec), } #[derive(Clone)] @@ -107,7 +116,6 @@ impl RpcNotificationSender { pub(crate) struct RpcRouter { request_routes: HashMap<&'static str, RequestRoute>, - concurrent_request_methods: HashSet<&'static str>, notification_routes: HashMap<&'static str, NotificationRoute>, } @@ -115,7 +123,6 @@ impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), - concurrent_request_methods: HashSet::new(), notification_routes: HashMap::new(), } } @@ -136,7 +143,6 @@ where F: Fn(Arc, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -166,24 +172,12 @@ where ); } - pub(crate) fn concurrent_request(&mut self, method: &'static str, handler: F) - where - P: DeserializeOwned + Send + 'static, - R: Serialize + Send + 'static, - F: Fn(Arc, P) -> Fut + Send + Sync + 'static, - Fut: Future> + Send + 'static, - { - self.request(method, handler); - self.concurrent_request_methods.insert(method); - } - pub(crate) fn request_with_id(&mut self, method: &'static str, handler: F) where P: DeserializeOwned + Send + 'static, F: Fn(Arc, RequestId, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -232,17 +226,13 @@ where self.request_routes.get(method) } - pub(crate) fn is_request_concurrent(&self, method: &str) -> bool { - self.concurrent_request_methods.contains(method) - } - pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { self.notification_routes.get(method) } } pub(crate) struct RpcClient { - write_tx: mpsc::Sender, + write_tx: mpsc::Sender, pending: Arc>>, // Shared transport state from `JsonRpcConnection`. Calls use this to fail // immediately when the socket closes, even if no JSON-RPC error response @@ -272,7 +262,7 @@ impl RpcClient { let closed_for_reader = Arc::clone(&closed); let transport_for_reader = transport.clone(); let reader_task = tokio::spawn(async move { - let disconnect_reason = loop { + let disconnect_reason = 'reader: loop { let Some(event) = incoming_rx.recv().await else { break None; }; @@ -285,6 +275,16 @@ impl RpcClient { break None; } } + JsonRpcConnectionEvent::Batch(messages) => { + for message in messages { + if let Err(err) = + handle_server_message(&pending_for_reader, &event_tx, message).await + { + let _ = err; + break 'reader None; + } + } + } JsonRpcConnectionEvent::MalformedMessage { reason } => { let _ = reason; break None; @@ -330,10 +330,13 @@ impl RpcClient { return Err(RpcCallError::Closed); } self.write_tx - .send(JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - })) + .send( + JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }) + .into(), + ) .await .map_err(|_| RpcCallError::Closed) } @@ -378,12 +381,15 @@ impl RpcClient { }; if self .write_tx - .send(JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), - method: method.to_string(), - params: Some(params), - trace: None, - })) + .send( + JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + }) + .into(), + ) .await .is_err() { @@ -406,6 +412,70 @@ impl RpcClient { serde_json::from_value(response).map_err(RpcCallError::Json) } + /// Sends independent RPC requests as one JSON-RPC batch. + /// + /// The server may execute batch entries concurrently and does not guarantee execution order. + /// Callers must only batch requests that are safe to run in any order, and must send dependent + /// requests separately. Results are returned in the same order as `calls` by matching request + /// ids in the responses. + pub(crate) async fn call_batch( + &self, + calls: Vec, + ) -> Result>, RpcCallError> { + if calls.is_empty() { + return Ok(Vec::new()); + } + if calls.len() > MAX_RPC_BATCH_REQUESTS { + return Err(RpcCallError::InvalidBatch(format!( + "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", + calls.len() + ))); + } + + let mut request_ids = Vec::with_capacity(calls.len()); + let mut requests = Vec::with_capacity(calls.len()); + let mut response_receivers = Vec::with_capacity(calls.len()); + { + let mut pending = self.pending.lock().await; + if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { + return Err(RpcCallError::Closed); + } + for RpcBatchCall { method, params } in calls { + let request_id = + RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); + let (response_tx, response_rx) = oneshot::channel(); + pending.insert(request_id.clone(), response_tx); + request_ids.push(request_id.clone()); + requests.push(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id, + method, + params: Some(params), + trace: None, + })); + response_receivers.push(response_rx); + } + } + + if self + .write_tx + .send(JsonRpcWireMessage::Batch(requests)) + .await + .is_err() + { + let mut pending = self.pending.lock().await; + for request_id in request_ids { + pending.remove(&request_id); + } + return Err(RpcCallError::Closed); + } + + let mut results = Vec::with_capacity(response_receivers.len()); + for response_rx in response_receivers { + results.push(response_rx.await.map_err(|_| RpcCallError::Closed)?); + } + Ok(results) + } + #[cfg(test)] pub(crate) async fn pending_request_count(&self) -> usize { self.pending.lock().await.len() @@ -424,22 +494,34 @@ impl Drop for RpcClient { pub(crate) fn encode_server_message( message: RpcServerOutboundMessage, -) -> Result { +) -> Result { match message { RpcServerOutboundMessage::Response { request_id, result } => { Ok(JSONRPCMessage::Response(JSONRPCResponse { id: request_id, result, - })) + }) + .into()) } RpcServerOutboundMessage::Error { request_id, error } => { Ok(JSONRPCMessage::Error(JSONRPCError { id: request_id, error, - })) + }) + .into()) } RpcServerOutboundMessage::Notification(notification) => { - Ok(JSONRPCMessage::Notification(notification)) + Ok(JSONRPCMessage::Notification(notification).into()) + } + RpcServerOutboundMessage::Batch(messages) => { + let mut encoded = Vec::with_capacity(messages.len()); + for message in messages { + match encode_server_message(message)? { + JsonRpcWireMessage::Single(message) => encoded.push(message), + JsonRpcWireMessage::Batch(messages) => encoded.extend(messages), + } + } + Ok(JsonRpcWireMessage::Batch(encoded)) } } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 7d203dd49519..6837832c406a 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,7 +1,7 @@ use std::sync::Arc; +use futures::StreamExt; use tokio::sync::mpsc; -use tokio::task::JoinSet; use tracing::debug; use tracing::warn; @@ -9,7 +9,7 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::protocol::INITIALIZED_METHOD; +use crate::connection::MAX_RPC_BATCH_REQUESTS; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; @@ -20,7 +20,7 @@ use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; -const MAX_CONCURRENT_READ_ONLY_REQUESTS: usize = 32; +const MAX_RPC_BATCH_CONCURRENCY: usize = 32; #[derive(Clone)] pub(crate) struct ConnectionProcessor { @@ -67,8 +67,6 @@ async fn run_connection( notifications, runtime_paths, )); - let mut read_only_tasks = JoinSet::new(); - let mut read_only_requests_enabled = false; let outbound_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { @@ -85,33 +83,13 @@ async fn run_connection( } }); - // Process inbound events sequentially to preserve initialize/initialized ordering. - loop { + while let Some(event) = incoming_rx.recv().await { if !handler.is_session_attached() { debug!("exec-server connection evicted after session resume"); break; } - let event = tokio::select! { - result = read_only_tasks.join_next(), if !read_only_tasks.is_empty() => { - if !send_finished_request(result, &outgoing_tx).await { - break; - } - continue; - } - event = incoming_rx.recv() => { - let Some(event) = event else { - break; - }; - event - } - }; match event { JsonRpcConnectionEvent::MalformedMessage { reason } => { - if !drain_read_only_tasks(&mut read_only_tasks, &outgoing_tx, &mut disconnected_rx) - .await - { - break; - } warn!("ignoring malformed exec-server message: {reason}"); if outgoing_tx .send(RpcServerOutboundMessage::Error { @@ -126,36 +104,6 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { - if read_only_requests_enabled - && router.is_request_concurrent(request.method.as_str()) - { - if !limit_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } - spawn_read_only_request( - &mut read_only_tasks, - Arc::clone(&router), - Arc::clone(&handler), - request, - ); - continue; - } - - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } let message = tokio::select! { message = dispatch_request( Arc::clone(&router), @@ -174,15 +122,6 @@ async fn run_connection( } } codex_app_server_protocol::JSONRPCMessage::Notification(notification) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } let Some(route) = router.notification_route(notification.method.as_str()) else { warn!( @@ -191,7 +130,6 @@ async fn run_connection( ); break; }; - let enables_read_only_requests = notification.method == INITIALIZED_METHOD; let result = tokio::select! { result = route(Arc::clone(&handler), notification) => result, _ = disconnected_rx.changed() => { @@ -205,20 +143,8 @@ async fn run_connection( warn!("closing exec-server connection after protocol error: {err}"); break; } - if enables_read_only_requests { - read_only_requests_enabled = true; - } } codex_app_server_protocol::JSONRPCMessage::Response(response) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } warn!( "closing exec-server connection after unexpected client response: {:?}", response.id @@ -226,15 +152,6 @@ async fn run_connection( break; } codex_app_server_protocol::JSONRPCMessage::Error(error) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } warn!( "closing exec-server connection after unexpected client error: {:?}", error.id @@ -242,6 +159,88 @@ async fn run_connection( break; } }, + JsonRpcConnectionEvent::Batch(messages) => { + if messages.is_empty() || messages.len() > MAX_RPC_BATCH_REQUESTS { + let message = if messages.is_empty() { + "JSON-RPC batch must not be empty".to_string() + } else { + format!( + "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", + messages.len() + ) + }; + if outgoing_tx + .send(RpcServerOutboundMessage::Error { + request_id: codex_app_server_protocol::RequestId::Integer(-1), + error: invalid_request(message), + }) + .await + .is_err() + { + break; + } + continue; + } + + let batch = futures::stream::iter(messages) + .map(|message| { + let router = Arc::clone(&router); + let handler = Arc::clone(&handler); + async move { + match message { + codex_app_server_protocol::JSONRPCMessage::Request(request) => { + dispatch_request(router, handler, request).await + } + codex_app_server_protocol::JSONRPCMessage::Notification(_) => { + Some(RpcServerOutboundMessage::Error { + request_id: codex_app_server_protocol::RequestId::Integer( + -1, + ), + error: invalid_request( + "notifications cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + codex_app_server_protocol::JSONRPCMessage::Response(response) => { + Some(RpcServerOutboundMessage::Error { + request_id: response.id, + error: invalid_request( + "client responses cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + codex_app_server_protocol::JSONRPCMessage::Error(error) => { + Some(RpcServerOutboundMessage::Error { + request_id: error.id, + error: invalid_request( + "client errors cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + } + } + }) + .buffered(MAX_RPC_BATCH_CONCURRENCY) + .filter_map(std::future::ready) + .collect::>(); + let messages = tokio::select! { + messages = batch => messages, + _ = disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling batch"); + break; + } + }; + if outgoing_tx + .send(RpcServerOutboundMessage::Batch(messages)) + .await + .is_err() + { + break; + } + } JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); @@ -251,7 +250,6 @@ async fn run_connection( } } - read_only_tasks.abort_all(); handler.shutdown().await; drop(handler); drop(outgoing_tx); @@ -262,70 +260,6 @@ async fn run_connection( let _ = outbound_task.await; } -async fn limit_read_only_tasks( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - while tasks.len() >= MAX_CONCURRENT_READ_ONLY_REQUESTS { - if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { - return false; - } - } - true -} - -async fn drain_read_only_tasks( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - while !tasks.is_empty() { - if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { - return false; - } - } - true -} - -async fn drain_one_read_only_task( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - let result = tokio::select! { - result = tasks.join_next() => result, - _ = disconnected_rx.changed() => { - debug!("exec-server transport disconnected while handling read-only request"); - return false; - } - }; - send_finished_request(result, outgoing_tx).await -} - -async fn send_finished_request( - result: Option, tokio::task::JoinError>>, - outgoing_tx: &mpsc::Sender, -) -> bool { - match result { - Some(Ok(Some(message))) => outgoing_tx.send(message).await.is_ok(), - Some(Ok(None)) | None => true, - Some(Err(err)) => { - warn!("read-only exec-server request task failed: {err}"); - true - } - } -} - -fn spawn_read_only_request( - tasks: &mut JoinSet>, - router: Arc>, - handler: Arc, - request: codex_app_server_protocol::JSONRPCRequest, -) { - tasks.spawn(dispatch_request(router, handler, request)); -} - async fn dispatch_request( router: Arc>, handler: Arc, @@ -370,11 +304,16 @@ mod tests { use crate::ExecServerRuntimePaths; use crate::ProcessId; use crate::connection::JsonRpcConnection; + use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_READ_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; + use crate::protocol::EnvironmentInfo; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; + use crate::protocol::FS_READ_DIRECTORY_METHOD; + use crate::protocol::FsReadDirectoryParams; + use crate::protocol::FsReadDirectoryResponse; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; @@ -384,6 +323,81 @@ mod tests { use crate::protocol::TerminateResponse; use crate::server::session_registry::SessionRegistry; + #[tokio::test] + async fn json_rpc_batch_dispatches_generic_requests() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = spawn_test_connection(registry, "batch"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + let cwd = PathUri::from_path(std::env::current_dir().expect("cwd")).expect("cwd URI"); + write_batch( + &mut writer, + &[ + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(2), + method: ENVIRONMENT_INFO_METHOD.to_string(), + params: Some(serde_json::to_value(()).expect("serialize params")), + trace: None, + }), + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(3), + method: FS_READ_DIRECTORY_METHOD.to_string(), + params: Some( + serde_json::to_value(FsReadDirectoryParams { + path: cwd, + sandbox: None, + }) + .expect("serialize fs/readDirectory params"), + ), + trace: None, + }), + ], + ) + .await; + + let messages = read_batch(&mut lines).await; + let mut responses = HashMap::new(); + for message in messages { + match message { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { + responses.insert(id, result); + } + JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"), + other => panic!("expected JSON-RPC response, got {other:?}"), + } + } + let environment_info = responses + .remove(&RequestId::Integer(2)) + .expect("environment/info response"); + let _: EnvironmentInfo = + serde_json::from_value(environment_info).expect("decode environment/info response"); + let read_directory = responses + .remove(&RequestId::Integer(3)) + .expect("fs/readDirectory response"); + let _: FsReadDirectoryResponse = + serde_json::from_value(read_directory).expect("decode fs/readDirectory response"); + assert!(responses.is_empty()); + + drop(writer); + drop(lines); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + #[tokio::test] async fn transport_disconnect_detaches_session_during_in_flight_read() { let registry = SessionRegistry::new(); @@ -529,6 +543,12 @@ mod tests { writer.write_all(b"\n").await.expect("write newline"); } + async fn write_batch(writer: &mut DuplexStream, messages: &[JSONRPCMessage]) { + let encoded = serde_json::to_vec(messages).expect("serialize JSON-RPC batch"); + writer.write_all(&encoded).await.expect("write batch"); + writer.write_all(b"\n").await.expect("write newline"); + } + async fn read_response( lines: &mut Lines>, expected_id: i64, @@ -548,6 +568,15 @@ mod tests { } } + async fn read_batch(lines: &mut Lines>) -> Vec { + let line = lines + .next_line() + .await + .expect("read batch") + .expect("batch line"); + serde_json::from_str(&line).expect("decode JSON-RPC batch") + } + fn exec_params(process_id: ProcessId) -> ExecParams { let mut env = HashMap::new(); if let Some(path) = std::env::var_os("PATH") { diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 891a0885de79..8f48aeaf99b7 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -93,7 +93,7 @@ pub(crate) fn build_router() -> RpcRouter { handler.terminate(params).await }, ); - router.concurrent_request( + router.request( FS_READ_FILE_METHOD, |handler: Arc, params: FsReadFileParams| async move { handler.fs_read_file(params).await @@ -129,19 +129,19 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_create_directory(params).await }, ); - router.concurrent_request( + router.request( FS_GET_METADATA_METHOD, |handler: Arc, params: FsGetMetadataParams| async move { handler.fs_get_metadata(params).await }, ); - router.concurrent_request( + router.request( FS_CANONICALIZE_METHOD, |handler: Arc, params: FsCanonicalizeParams| async move { handler.fs_canonicalize(params).await }, ); - router.concurrent_request( + router.request( FS_READ_DIRECTORY_METHOD, |handler: Arc, params: FsReadDirectoryParams| async move { handler.fs_read_directory(params).await diff --git a/codex-rs/exec-server/tests/file_system/shared.rs b/codex-rs/exec-server/tests/file_system/shared.rs index 2d7cc1c668d9..08ba2efdd81b 100644 --- a/codex-rs/exec-server/tests/file_system/shared.rs +++ b/codex-rs/exec-server/tests/file_system/shared.rs @@ -4,8 +4,6 @@ use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FILE_READ_CHUNK_SIZE; use codex_exec_server::FileMetadata; -use codex_exec_server::FileSystemOperation; -use codex_exec_server::FileSystemOperationOutput; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; use codex_protocol::models::AdditionalPermissionProfile; @@ -198,55 +196,6 @@ async fn file_system_read_file_returns_bytes( Ok(()) } -#[test_case(FileSystemImplementation::Local ; "local")] -#[test_case(FileSystemImplementation::Remote ; "remote")] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn file_system_batch_preserves_operation_order( - implementation: FileSystemImplementation, -) -> Result<()> { - let context = create_file_system_context(implementation).await?; - let file_system = context.file_system; - - let tmp = TempDir::new()?; - let first_path = tmp.path().join("first.txt"); - let second_path = tmp.path().join("second.txt"); - std::fs::write(&first_path, "first")?; - std::fs::write(&second_path, "second")?; - let first_uri = PathUri::from_path(&first_path)?; - let second_uri = PathUri::from_path(&second_path)?; - - let outputs = file_system - .execute_batch( - vec![ - FileSystemOperation::ReadFile { - path: first_uri.clone(), - }, - FileSystemOperation::Canonicalize { - path: second_uri.clone(), - }, - FileSystemOperation::ReadFile { path: second_uri }, - ], - /*sandbox*/ None, - ) - .await - .with_context(|| format!("mode={implementation}"))? - .into_iter() - .collect::>>()?; - - assert_eq!( - outputs, - vec![ - FileSystemOperationOutput::ReadFile(b"first".to_vec()), - FileSystemOperationOutput::Canonicalize(PathUri::from_path(std::fs::canonicalize( - second_path - )?)?), - FileSystemOperationOutput::ReadFile(b"second".to_vec()), - ] - ); - - Ok(()) -} - #[test_case(FileSystemImplementation::Local ; "local")] #[test_case(FileSystemImplementation::Remote ; "remote")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/file-system/Cargo.toml b/codex-rs/file-system/Cargo.toml index 382eacb43d33..68a63620c886 100644 --- a/codex-rs/file-system/Cargo.toml +++ b/codex-rs/file-system/Cargo.toml @@ -14,6 +14,7 @@ codex-utils-absolute-path = { workspace = true } codex-utils-path-uri = { workspace = true } futures = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } [lib] test = false diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index 1fd1c3ebf378..18ecfc0cca26 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -12,7 +12,6 @@ use codex_protocol::protocol::SandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::Stream; -use futures::StreamExt; use std::future::Future; use std::io; use std::path::Path; @@ -22,7 +21,6 @@ use std::task::Poll; /// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`]. pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024; -const FILE_SYSTEM_BATCH_CONCURRENCY: usize = 32; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct CreateDirectoryOptions { @@ -58,25 +56,13 @@ pub struct ReadDirectoryEntry { pub is_file: bool, } -/// One independent, read-only filesystem operation eligible for batching. -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum FileSystemOperation { - Canonicalize { path: PathUri }, - ReadFile { path: PathUri }, - GetMetadata { path: PathUri }, - ReadDirectory { path: PathUri }, -} - -/// Typed output for a [`FileSystemOperation`]. -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum FileSystemOperationOutput { - Canonicalize(PathUri), - ReadFile(Vec), - GetMetadata(FileMetadata), - ReadDirectory(Vec), +#[derive(Clone, Debug, PartialEq)] +pub struct ExecutorRpcBatchCall { + pub method: String, + pub params: serde_json::Value, } -pub type FileSystemOperationResult = FileSystemResult; +pub type ExecutorRpcBatchResult = FileSystemResult; #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] @@ -286,61 +272,20 @@ pub trait ExecutorFileSystem: Send + Sync { sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, ()>; - /// Executes independent read-only filesystem operations concurrently. + /// Sends independent executor RPC requests as one batch. /// - /// Results retain request order. Operations have no ordering or transactional relationship; - /// callers must put dependent operations in separate batches. - fn execute_batch<'a>( + /// The executor may execute entries concurrently and does not guarantee execution order. + /// Callers must only batch requests that are safe to run in any order, and must send dependent + /// requests separately. Results are matched back to the input order by request id. + fn execute_rpc_batch<'a>( &'a self, - operations: Vec, - sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, Vec> { - Box::pin(execute_batch_with_scalar_operations( - self, operations, sandbox, - )) - } -} - -/// Default batch implementation used by local filesystems and by remote fallbacks. -pub async fn execute_batch_with_scalar_operations( - file_system: &F, - operations: Vec, - sandbox: Option<&FileSystemSandboxContext>, -) -> FileSystemResult> -where - F: ExecutorFileSystem + ?Sized, -{ - Ok(futures::stream::iter(operations) - .map(|operation| execute_scalar_operation(file_system, operation, sandbox)) - .buffered(FILE_SYSTEM_BATCH_CONCURRENCY) - .collect() - .await) -} - -async fn execute_scalar_operation( - file_system: &F, - operation: FileSystemOperation, - sandbox: Option<&FileSystemSandboxContext>, -) -> FileSystemOperationResult -where - F: ExecutorFileSystem + ?Sized, -{ - match operation { - FileSystemOperation::Canonicalize { path } => file_system - .canonicalize(&path, sandbox) - .await - .map(FileSystemOperationOutput::Canonicalize), - FileSystemOperation::ReadFile { path } => file_system - .read_file(&path, sandbox) - .await - .map(FileSystemOperationOutput::ReadFile), - FileSystemOperation::GetMetadata { path } => file_system - .get_metadata(&path, sandbox) - .await - .map(FileSystemOperationOutput::GetMetadata), - FileSystemOperation::ReadDirectory { path } => file_system - .read_directory(&path, sandbox) - .await - .map(FileSystemOperationOutput::ReadDirectory), + _calls: Vec, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(async move { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "filesystem does not support RPC batching", + )) + }) } } From 2cd053e8a723e92ec7408b6b8a5c52baea9a43b1 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 19 Jun 2026 18:30:54 +0200 Subject: [PATCH 4/7] Use generic RPC batch requests --- codex-rs/Cargo.lock | 1 + codex-rs/exec-server/src/client.rs | 25 ++ codex-rs/exec-server/src/connection.rs | 48 ++- codex-rs/exec-server/src/lib.rs | 9 +- .../src/noise_relay/executor_stream.rs | 2 +- .../src/noise_relay/executor_stream_tests.rs | 2 +- .../exec-server/src/noise_relay/harness.rs | 2 +- .../src/noise_relay/message_framing.rs | 9 +- .../src/noise_relay/message_framing_tests.rs | 5 +- codex-rs/exec-server/src/relay.rs | 18 +- .../exec-server/src/remote_file_system.rs | 48 +++ codex-rs/exec-server/src/rpc.rs | 154 ++++++-- codex-rs/exec-server/src/server/processor.rs | 333 ++++++++++-------- codex-rs/exec-server/src/server/registry.rs | 8 +- .../exec-server/tests/file_system/shared.rs | 51 --- codex-rs/file-system/Cargo.toml | 1 + codex-rs/file-system/src/lib.rs | 91 +---- 17 files changed, 462 insertions(+), 345 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 55d02cd8b155..1036f4815f75 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3049,6 +3049,7 @@ dependencies = [ "codex-utils-path-uri", "futures", "serde", + "serde_json", ] [[package]] diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 88c0c8034288..66a5708f4406 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -97,6 +97,7 @@ use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; +use crate::rpc::RpcBatchCall; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; @@ -761,6 +762,29 @@ impl ExecServerClient { self.call_rpc(&rpc_client, method, params).await } + pub(crate) async fn call_batch( + &self, + calls: Vec, + ) -> Result>, ExecServerError> { + let rpc_client = self.inner.rpc_client().await?; + match rpc_client.call_batch(calls).await { + Ok(results) => Ok(results + .into_iter() + .map(|result| result.map_err(ExecServerError::from)) + .collect()), + Err(error) => { + let error = ExecServerError::from(error); + if is_transport_closed_error(&error) { + Err(ExecServerError::Disconnected(disconnected_message( + /*reason*/ None, + ))) + } else { + Err(error) + } + } + } + } + async fn call_rpc( &self, rpc_client: &Arc, @@ -813,6 +837,7 @@ impl From for ExecServerError { code: error.code, message: error.message, }, + RpcCallError::InvalidBatch(message) => Self::Protocol(message), } } } diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 2d2e5afe0641..9a3da73e80d8 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -12,6 +12,8 @@ use futures::Sink; use futures::SinkExt; use futures::Stream; use futures::StreamExt; +use serde::Deserialize; +use serde::Serialize; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::process::Child; @@ -29,6 +31,7 @@ use tokio::io::BufReader; use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; +pub(crate) const MAX_RPC_BATCH_REQUESTS: usize = 512; const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2); #[cfg(test)] pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25); @@ -38,10 +41,33 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30 #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { Message(JSONRPCMessage), + Batch(Vec), MalformedMessage { reason: String }, Disconnected { reason: Option }, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum JsonRpcWireMessage { + Single(JSONRPCMessage), + Batch(Vec), +} + +impl JsonRpcWireMessage { + pub(crate) fn into_connection_event(self) -> JsonRpcConnectionEvent { + match self { + Self::Single(message) => JsonRpcConnectionEvent::Message(message), + Self::Batch(messages) => JsonRpcConnectionEvent::Batch(messages), + } + } +} + +impl From for JsonRpcWireMessage { + fn from(message: JSONRPCMessage) -> Self { + Self::Single(message) + } +} + #[derive(Clone)] pub(crate) enum JsonRpcTransport { // Plain means no child process; transport bytes may still be encrypted. @@ -221,7 +247,7 @@ fn log_stdio_child_wait_result(result: std::io::Result } pub(crate) struct JsonRpcConnection { - pub(crate) outgoing_tx: mpsc::Sender, + pub(crate) outgoing_tx: mpsc::Sender, pub(crate) incoming_rx: mpsc::Receiver, pub(crate) disconnected_rx: watch::Receiver, pub(crate) task_handles: Vec>, @@ -249,10 +275,10 @@ impl JsonRpcConnection { if line.trim().is_empty() { continue; } - match serde_json::from_str::(&line) { + match serde_json::from_str::(&line) { Ok(message) => { if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -395,7 +421,7 @@ impl JsonRpcConnection { Some(Ok(message)) => match message.parse_jsonrpc_frame() { Ok(JsonRpcWebSocketFrame::Message(message)) => { if incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .is_err() { @@ -464,7 +490,7 @@ impl JsonRpcConnection { } enum JsonRpcWebSocketFrame { - Message(JSONRPCMessage), + Message(JsonRpcWireMessage), Close, Ignore, } @@ -549,7 +575,7 @@ async fn send_malformed_message( async fn write_jsonrpc_line_message( writer: &mut BufWriter, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> std::io::Result<()> where W: AsyncWrite + Unpin, @@ -564,7 +590,7 @@ where async fn send_websocket_jsonrpc_message( websocket_writer: &mut W, connection_label: &str, - message: &JSONRPCMessage, + message: &JsonRpcWireMessage, ) -> Result<(), String> where W: Sink + Unpin, @@ -584,7 +610,7 @@ where } } -fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { +fn serialize_jsonrpc_message(message: &JsonRpcWireMessage) -> Result { serde_json::to_string(message) } @@ -694,7 +720,7 @@ mod tests { ); let message = test_jsonrpc_message(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -706,7 +732,9 @@ mod tests { control.set_write_ready(); assert!(matches!( timeout(Duration::from_secs(1), outbound_rx.next()).await?, - Some(Message::Text(text)) if serde_json::from_str::(&text)? == message + Some(Message::Text(text)) + if serde_json::from_str::(&text)? + == JsonRpcWireMessage::Single(message) )); drop(connection); Ok(()) diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index d66d2f9b7e68..c1779148d8b1 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,11 +42,10 @@ pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; pub use codex_file_system::ExecutorFileSystem; pub use codex_file_system::ExecutorFileSystemFuture; +pub use codex_file_system::ExecutorRpcBatchCall; +pub use codex_file_system::ExecutorRpcBatchResult; pub use codex_file_system::FILE_READ_CHUNK_SIZE; pub use codex_file_system::FileMetadata; -pub use codex_file_system::FileSystemOperation; -pub use codex_file_system::FileSystemOperationOutput; -pub use codex_file_system::FileSystemOperationResult; pub use codex_file_system::FileSystemReadStream; pub use codex_file_system::FileSystemResult; pub use codex_file_system::FileSystemSandboxContext; @@ -90,6 +89,10 @@ pub use protocol::ExecOutputDeltaNotification; pub use protocol::ExecOutputStream; pub use protocol::ExecParams; pub use protocol::ExecResponse; +pub use protocol::FS_CANONICALIZE_METHOD; +pub use protocol::FS_GET_METADATA_METHOD; +pub use protocol::FS_READ_DIRECTORY_METHOD; +pub use protocol::FS_READ_FILE_METHOD; pub use protocol::FsCanonicalizeParams; pub use protocol::FsCanonicalizeResponse; pub use protocol::FsCloseParams; diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 57a05579dcf5..04f25dc4462b 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -75,7 +75,7 @@ impl NoiseVirtualStream { }; for message in self.inbound_decoder.push(&plaintext)? { self.incoming_tx - .try_send(JsonRpcConnectionEvent::Message(message)) + .try_send(message.into_connection_event()) .map_err(|_| { ExecServerError::Protocol( "Noise virtual stream inbound queue is full or closed".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs index 9119236cfd28..d757a12393ef 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs @@ -51,7 +51,7 @@ async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { id: RequestId::Integer(1), result: serde_json::Value::Null, }); - let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message)?)?; + let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message.clone().into())?)?; stream.receive_data(RelayData { seq: 0, segment_index: 0, diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 1989525224f6..62e82e608c1a 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -420,7 +420,7 @@ async fn receive_data( // messages; emit only complete, successfully parsed messages. for message in decoder.push(&plaintext)? { incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) + .send(message.into_connection_event()) .await .map_err(|_| ExecServerError::Closed)?; } diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index 24ee48173e2e..b84253695c07 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -1,6 +1,5 @@ -use codex_app_server_protocol::JSONRPCMessage; - use crate::ExecServerError; +use crate::connection::JsonRpcWireMessage; const LENGTH_PREFIX_BYTES: usize = size_of::(); const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024; @@ -12,7 +11,9 @@ pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024; /// exec-server responses can be much larger. A four-byte authenticated length /// prefix lets the caller split this byte stream into bounded Noise records and /// lets the receiver reconstruct exact JSON-RPC message boundaries. -pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn frame_jsonrpc_message( + message: &JsonRpcWireMessage, +) -> Result, ExecServerError> { let mut framed = vec![0; LENGTH_PREFIX_BYTES]; serde_json::to_writer(&mut framed, message)?; let message_len = framed.len() - LENGTH_PREFIX_BYTES; @@ -39,7 +40,7 @@ impl JsonRpcMessageDecoder { pub(crate) fn push( &mut self, plaintext_record: &[u8], - ) -> Result, ExecServerError> { + ) -> Result, ExecServerError> { if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN { return Err(ExecServerError::Protocol( "Noise relay plaintext record exceeds maximum length".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs index c3df14bf5622..5cc3ed10f1c1 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs @@ -7,6 +7,7 @@ use super::MAX_NOISE_JSONRPC_MESSAGE_LEN; use super::NOISE_RECORD_PLAINTEXT_LEN; use super::frame_jsonrpc_message; use crate::ExecServerError; +use crate::connection::JsonRpcWireMessage; #[test] fn fragments_and_reassembles_large_jsonrpc_message() { @@ -16,7 +17,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { "data": "x".repeat(128 * 1024), })), }); - let framed = frame_jsonrpc_message(&message).unwrap(); + let framed = frame_jsonrpc_message(&message.clone().into()).unwrap(); assert!(framed.len() > 128 * 1024); let mut decoder = JsonRpcMessageDecoder::default(); @@ -25,7 +26,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { decoded.extend(decoder.push(record).unwrap()); } - assert_eq!(decoded, vec![message]); + assert_eq!(decoded, vec![JsonRpcWireMessage::Single(message)]); } #[test] diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 033da4b2a746..a00707c67779 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use std::time::Duration; -use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; use futures::SinkExt; use futures::Stream; @@ -25,6 +24,7 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::JsonRpcWireMessage; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; @@ -170,7 +170,7 @@ impl RelayMessageFrame { } } - fn into_jsonrpc_message(self) -> Result { + fn into_jsonrpc_message(self) -> Result { let payload = self.into_data()?.payload; serde_json::from_slice(&payload).map_err(ExecServerError::Json) } @@ -211,7 +211,7 @@ pub(crate) fn decode_relay_message_frame( .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) } -pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn jsonrpc_payload(message: &JsonRpcWireMessage) -> Result, ExecServerError> { serde_json::to_vec(message).map_err(ExecServerError::Json) } @@ -346,7 +346,7 @@ where &mut websocket, &mut keepalive, &incoming_tx, - JsonRpcConnectionEvent::Message(message), + message.into_connection_event(), ) .await { @@ -847,6 +847,7 @@ mod tests { use std::task::Poll; use std::time::Duration; + use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; use futures::Sink; @@ -880,7 +881,7 @@ mod tests { encode_relay_message_frame(&RelayMessageFrame::data( stream_id, /*seq*/ 0, - jsonrpc_payload(&message)?, + jsonrpc_payload(&message.clone().into())?, )) .into(), )) @@ -1029,7 +1030,7 @@ mod tests { let message = test_jsonrpc_message(); control.set_write_blocked(); - connection.outgoing_tx.send(message.clone()).await?; + connection.outgoing_tx.send(message.clone().into()).await?; control.wait_for_blocked_write().await?; control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( @@ -1047,7 +1048,10 @@ mod tests { }; let frame = decode_relay_message_frame(data_payload.as_ref())?; assert_eq!(frame.stream_id, stream_id); - assert_eq!(frame.into_jsonrpc_message()?, message); + assert_eq!( + frame.into_jsonrpc_message()?, + JsonRpcWireMessage::Single(message) + ); drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index f9136e146108..5950e919f125 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -9,6 +9,8 @@ use crate::CreateDirectoryOptions; use crate::ExecServerError; use crate::ExecutorFileSystem; use crate::ExecutorFileSystemFuture; +use crate::ExecutorRpcBatchCall; +use crate::ExecutorRpcBatchResult; use crate::FileMetadata; use crate::FileSystemReadStream; use crate::FileSystemResult; @@ -16,6 +18,7 @@ use crate::FileSystemSandboxContext; use crate::ReadDirectoryEntry; use crate::RemoveOptions; use crate::client::LazyRemoteExecServerClient; +use crate::connection::MAX_RPC_BATCH_REQUESTS; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; @@ -24,6 +27,7 @@ use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadFileParams; use crate::protocol::FsRemoveParams; use crate::protocol::FsWriteFileParams; +use crate::rpc::RpcBatchCall; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const NOT_FOUND_ERROR_CODE: i64 = -32004; @@ -223,6 +227,43 @@ impl RemoteFileSystem { .map_err(map_remote_error)?; Ok(()) } + + async fn execute_rpc_batch( + &self, + calls: Vec, + ) -> FileSystemResult> { + if calls.is_empty() { + return Ok(Vec::new()); + } + + let client = self.client.get().await.map_err(map_remote_error)?; + let mut decoded = Vec::with_capacity(calls.len()); + let mut calls = calls.into_iter(); + loop { + let call_chunk = calls + .by_ref() + .take(MAX_RPC_BATCH_REQUESTS) + .map(|call| RpcBatchCall { + method: call.method, + params: call.params, + }) + .collect::>(); + if call_chunk.is_empty() { + break; + } + + let results = client + .call_batch(call_chunk) + .await + .map_err(map_remote_error)?; + decoded.extend( + results + .into_iter() + .map(|result| result.map_err(map_remote_error)), + ); + } + Ok(decoded) + } } impl ExecutorFileSystem for RemoteFileSystem { @@ -310,6 +351,13 @@ impl ExecutorFileSystem for RemoteFileSystem { sandbox, )) } + + fn execute_rpc_batch<'a>( + &'a self, + calls: Vec, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(RemoteFileSystem::execute_rpc_batch(self, calls)) + } } fn remote_sandbox_context( diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index c2a5124541d3..b16a3215110e 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::collections::HashSet; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -26,6 +25,8 @@ use tokio::task::JoinHandle; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::JsonRpcWireMessage; +use crate::connection::MAX_RPC_BATCH_REQUESTS; pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; @@ -37,6 +38,13 @@ pub(crate) enum RpcCallError { Json(serde_json::Error), /// The executor returned a JSON-RPC error response for this call. Server(JSONRPCErrorError), + /// The caller attempted to construct a batch that cannot be represented safely. + InvalidBatch(String), +} + +pub(crate) struct RpcBatchCall { + pub(crate) method: String, + pub(crate) params: Value, } type PendingRequest = oneshot::Sender>; @@ -64,6 +72,7 @@ pub(crate) enum RpcServerOutboundMessage { error: JSONRPCErrorError, }, Notification(JSONRPCNotification), + Batch(Vec), } #[derive(Clone)] @@ -107,7 +116,6 @@ impl RpcNotificationSender { pub(crate) struct RpcRouter { request_routes: HashMap<&'static str, RequestRoute>, - concurrent_request_methods: HashSet<&'static str>, notification_routes: HashMap<&'static str, NotificationRoute>, } @@ -115,7 +123,6 @@ impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), - concurrent_request_methods: HashSet::new(), notification_routes: HashMap::new(), } } @@ -136,7 +143,6 @@ where F: Fn(Arc, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -166,24 +172,12 @@ where ); } - pub(crate) fn concurrent_request(&mut self, method: &'static str, handler: F) - where - P: DeserializeOwned + Send + 'static, - R: Serialize + Send + 'static, - F: Fn(Arc, P) -> Fut + Send + Sync + 'static, - Fut: Future> + Send + 'static, - { - self.request(method, handler); - self.concurrent_request_methods.insert(method); - } - pub(crate) fn request_with_id(&mut self, method: &'static str, handler: F) where P: DeserializeOwned + Send + 'static, F: Fn(Arc, RequestId, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { - self.concurrent_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -232,17 +226,13 @@ where self.request_routes.get(method) } - pub(crate) fn is_request_concurrent(&self, method: &str) -> bool { - self.concurrent_request_methods.contains(method) - } - pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { self.notification_routes.get(method) } } pub(crate) struct RpcClient { - write_tx: mpsc::Sender, + write_tx: mpsc::Sender, pending: Arc>>, // Shared transport state from `JsonRpcConnection`. Calls use this to fail // immediately when the socket closes, even if no JSON-RPC error response @@ -272,7 +262,7 @@ impl RpcClient { let closed_for_reader = Arc::clone(&closed); let transport_for_reader = transport.clone(); let reader_task = tokio::spawn(async move { - let disconnect_reason = loop { + let disconnect_reason = 'reader: loop { let Some(event) = incoming_rx.recv().await else { break None; }; @@ -285,6 +275,16 @@ impl RpcClient { break None; } } + JsonRpcConnectionEvent::Batch(messages) => { + for message in messages { + if let Err(err) = + handle_server_message(&pending_for_reader, &event_tx, message).await + { + let _ = err; + break 'reader None; + } + } + } JsonRpcConnectionEvent::MalformedMessage { reason } => { let _ = reason; break None; @@ -330,10 +330,13 @@ impl RpcClient { return Err(RpcCallError::Closed); } self.write_tx - .send(JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - })) + .send( + JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }) + .into(), + ) .await .map_err(|_| RpcCallError::Closed) } @@ -378,12 +381,15 @@ impl RpcClient { }; if self .write_tx - .send(JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), - method: method.to_string(), - params: Some(params), - trace: None, - })) + .send( + JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + }) + .into(), + ) .await .is_err() { @@ -406,6 +412,70 @@ impl RpcClient { serde_json::from_value(response).map_err(RpcCallError::Json) } + /// Sends independent RPC requests as one JSON-RPC batch. + /// + /// The server may execute batch entries concurrently and does not guarantee execution order. + /// Callers must only batch requests that are safe to run in any order, and must send dependent + /// requests separately. Results are returned in the same order as `calls` by matching request + /// ids in the responses. + pub(crate) async fn call_batch( + &self, + calls: Vec, + ) -> Result>, RpcCallError> { + if calls.is_empty() { + return Ok(Vec::new()); + } + if calls.len() > MAX_RPC_BATCH_REQUESTS { + return Err(RpcCallError::InvalidBatch(format!( + "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", + calls.len() + ))); + } + + let mut request_ids = Vec::with_capacity(calls.len()); + let mut requests = Vec::with_capacity(calls.len()); + let mut response_receivers = Vec::with_capacity(calls.len()); + { + let mut pending = self.pending.lock().await; + if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { + return Err(RpcCallError::Closed); + } + for RpcBatchCall { method, params } in calls { + let request_id = + RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); + let (response_tx, response_rx) = oneshot::channel(); + pending.insert(request_id.clone(), response_tx); + request_ids.push(request_id.clone()); + requests.push(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id, + method, + params: Some(params), + trace: None, + })); + response_receivers.push(response_rx); + } + } + + if self + .write_tx + .send(JsonRpcWireMessage::Batch(requests)) + .await + .is_err() + { + let mut pending = self.pending.lock().await; + for request_id in request_ids { + pending.remove(&request_id); + } + return Err(RpcCallError::Closed); + } + + let mut results = Vec::with_capacity(response_receivers.len()); + for response_rx in response_receivers { + results.push(response_rx.await.map_err(|_| RpcCallError::Closed)?); + } + Ok(results) + } + #[cfg(test)] pub(crate) async fn pending_request_count(&self) -> usize { self.pending.lock().await.len() @@ -424,22 +494,34 @@ impl Drop for RpcClient { pub(crate) fn encode_server_message( message: RpcServerOutboundMessage, -) -> Result { +) -> Result { match message { RpcServerOutboundMessage::Response { request_id, result } => { Ok(JSONRPCMessage::Response(JSONRPCResponse { id: request_id, result, - })) + }) + .into()) } RpcServerOutboundMessage::Error { request_id, error } => { Ok(JSONRPCMessage::Error(JSONRPCError { id: request_id, error, - })) + }) + .into()) } RpcServerOutboundMessage::Notification(notification) => { - Ok(JSONRPCMessage::Notification(notification)) + Ok(JSONRPCMessage::Notification(notification).into()) + } + RpcServerOutboundMessage::Batch(messages) => { + let mut encoded = Vec::with_capacity(messages.len()); + for message in messages { + match encode_server_message(message)? { + JsonRpcWireMessage::Single(message) => encoded.push(message), + JsonRpcWireMessage::Batch(messages) => encoded.extend(messages), + } + } + Ok(JsonRpcWireMessage::Batch(encoded)) } } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 7d203dd49519..6837832c406a 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,7 +1,7 @@ use std::sync::Arc; +use futures::StreamExt; use tokio::sync::mpsc; -use tokio::task::JoinSet; use tracing::debug; use tracing::warn; @@ -9,7 +9,7 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::protocol::INITIALIZED_METHOD; +use crate::connection::MAX_RPC_BATCH_REQUESTS; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; @@ -20,7 +20,7 @@ use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; -const MAX_CONCURRENT_READ_ONLY_REQUESTS: usize = 32; +const MAX_RPC_BATCH_CONCURRENCY: usize = 32; #[derive(Clone)] pub(crate) struct ConnectionProcessor { @@ -67,8 +67,6 @@ async fn run_connection( notifications, runtime_paths, )); - let mut read_only_tasks = JoinSet::new(); - let mut read_only_requests_enabled = false; let outbound_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { @@ -85,33 +83,13 @@ async fn run_connection( } }); - // Process inbound events sequentially to preserve initialize/initialized ordering. - loop { + while let Some(event) = incoming_rx.recv().await { if !handler.is_session_attached() { debug!("exec-server connection evicted after session resume"); break; } - let event = tokio::select! { - result = read_only_tasks.join_next(), if !read_only_tasks.is_empty() => { - if !send_finished_request(result, &outgoing_tx).await { - break; - } - continue; - } - event = incoming_rx.recv() => { - let Some(event) = event else { - break; - }; - event - } - }; match event { JsonRpcConnectionEvent::MalformedMessage { reason } => { - if !drain_read_only_tasks(&mut read_only_tasks, &outgoing_tx, &mut disconnected_rx) - .await - { - break; - } warn!("ignoring malformed exec-server message: {reason}"); if outgoing_tx .send(RpcServerOutboundMessage::Error { @@ -126,36 +104,6 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_app_server_protocol::JSONRPCMessage::Request(request) => { - if read_only_requests_enabled - && router.is_request_concurrent(request.method.as_str()) - { - if !limit_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } - spawn_read_only_request( - &mut read_only_tasks, - Arc::clone(&router), - Arc::clone(&handler), - request, - ); - continue; - } - - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } let message = tokio::select! { message = dispatch_request( Arc::clone(&router), @@ -174,15 +122,6 @@ async fn run_connection( } } codex_app_server_protocol::JSONRPCMessage::Notification(notification) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } let Some(route) = router.notification_route(notification.method.as_str()) else { warn!( @@ -191,7 +130,6 @@ async fn run_connection( ); break; }; - let enables_read_only_requests = notification.method == INITIALIZED_METHOD; let result = tokio::select! { result = route(Arc::clone(&handler), notification) => result, _ = disconnected_rx.changed() => { @@ -205,20 +143,8 @@ async fn run_connection( warn!("closing exec-server connection after protocol error: {err}"); break; } - if enables_read_only_requests { - read_only_requests_enabled = true; - } } codex_app_server_protocol::JSONRPCMessage::Response(response) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } warn!( "closing exec-server connection after unexpected client response: {:?}", response.id @@ -226,15 +152,6 @@ async fn run_connection( break; } codex_app_server_protocol::JSONRPCMessage::Error(error) => { - if !drain_read_only_tasks( - &mut read_only_tasks, - &outgoing_tx, - &mut disconnected_rx, - ) - .await - { - break; - } warn!( "closing exec-server connection after unexpected client error: {:?}", error.id @@ -242,6 +159,88 @@ async fn run_connection( break; } }, + JsonRpcConnectionEvent::Batch(messages) => { + if messages.is_empty() || messages.len() > MAX_RPC_BATCH_REQUESTS { + let message = if messages.is_empty() { + "JSON-RPC batch must not be empty".to_string() + } else { + format!( + "JSON-RPC batch contains {} requests; maximum is {MAX_RPC_BATCH_REQUESTS}", + messages.len() + ) + }; + if outgoing_tx + .send(RpcServerOutboundMessage::Error { + request_id: codex_app_server_protocol::RequestId::Integer(-1), + error: invalid_request(message), + }) + .await + .is_err() + { + break; + } + continue; + } + + let batch = futures::stream::iter(messages) + .map(|message| { + let router = Arc::clone(&router); + let handler = Arc::clone(&handler); + async move { + match message { + codex_app_server_protocol::JSONRPCMessage::Request(request) => { + dispatch_request(router, handler, request).await + } + codex_app_server_protocol::JSONRPCMessage::Notification(_) => { + Some(RpcServerOutboundMessage::Error { + request_id: codex_app_server_protocol::RequestId::Integer( + -1, + ), + error: invalid_request( + "notifications cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + codex_app_server_protocol::JSONRPCMessage::Response(response) => { + Some(RpcServerOutboundMessage::Error { + request_id: response.id, + error: invalid_request( + "client responses cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + codex_app_server_protocol::JSONRPCMessage::Error(error) => { + Some(RpcServerOutboundMessage::Error { + request_id: error.id, + error: invalid_request( + "client errors cannot be sent in a JSON-RPC batch" + .to_string(), + ), + }) + } + } + } + }) + .buffered(MAX_RPC_BATCH_CONCURRENCY) + .filter_map(std::future::ready) + .collect::>(); + let messages = tokio::select! { + messages = batch => messages, + _ = disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling batch"); + break; + } + }; + if outgoing_tx + .send(RpcServerOutboundMessage::Batch(messages)) + .await + .is_err() + { + break; + } + } JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); @@ -251,7 +250,6 @@ async fn run_connection( } } - read_only_tasks.abort_all(); handler.shutdown().await; drop(handler); drop(outgoing_tx); @@ -262,70 +260,6 @@ async fn run_connection( let _ = outbound_task.await; } -async fn limit_read_only_tasks( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - while tasks.len() >= MAX_CONCURRENT_READ_ONLY_REQUESTS { - if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { - return false; - } - } - true -} - -async fn drain_read_only_tasks( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - while !tasks.is_empty() { - if !drain_one_read_only_task(tasks, outgoing_tx, disconnected_rx).await { - return false; - } - } - true -} - -async fn drain_one_read_only_task( - tasks: &mut JoinSet>, - outgoing_tx: &mpsc::Sender, - disconnected_rx: &mut tokio::sync::watch::Receiver, -) -> bool { - let result = tokio::select! { - result = tasks.join_next() => result, - _ = disconnected_rx.changed() => { - debug!("exec-server transport disconnected while handling read-only request"); - return false; - } - }; - send_finished_request(result, outgoing_tx).await -} - -async fn send_finished_request( - result: Option, tokio::task::JoinError>>, - outgoing_tx: &mpsc::Sender, -) -> bool { - match result { - Some(Ok(Some(message))) => outgoing_tx.send(message).await.is_ok(), - Some(Ok(None)) | None => true, - Some(Err(err)) => { - warn!("read-only exec-server request task failed: {err}"); - true - } - } -} - -fn spawn_read_only_request( - tasks: &mut JoinSet>, - router: Arc>, - handler: Arc, - request: codex_app_server_protocol::JSONRPCRequest, -) { - tasks.spawn(dispatch_request(router, handler, request)); -} - async fn dispatch_request( router: Arc>, handler: Arc, @@ -370,11 +304,16 @@ mod tests { use crate::ExecServerRuntimePaths; use crate::ProcessId; use crate::connection::JsonRpcConnection; + use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_READ_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; + use crate::protocol::EnvironmentInfo; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; + use crate::protocol::FS_READ_DIRECTORY_METHOD; + use crate::protocol::FsReadDirectoryParams; + use crate::protocol::FsReadDirectoryResponse; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; @@ -384,6 +323,81 @@ mod tests { use crate::protocol::TerminateResponse; use crate::server::session_registry::SessionRegistry; + #[tokio::test] + async fn json_rpc_batch_dispatches_generic_requests() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = spawn_test_connection(registry, "batch"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + let cwd = PathUri::from_path(std::env::current_dir().expect("cwd")).expect("cwd URI"); + write_batch( + &mut writer, + &[ + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(2), + method: ENVIRONMENT_INFO_METHOD.to_string(), + params: Some(serde_json::to_value(()).expect("serialize params")), + trace: None, + }), + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(3), + method: FS_READ_DIRECTORY_METHOD.to_string(), + params: Some( + serde_json::to_value(FsReadDirectoryParams { + path: cwd, + sandbox: None, + }) + .expect("serialize fs/readDirectory params"), + ), + trace: None, + }), + ], + ) + .await; + + let messages = read_batch(&mut lines).await; + let mut responses = HashMap::new(); + for message in messages { + match message { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { + responses.insert(id, result); + } + JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"), + other => panic!("expected JSON-RPC response, got {other:?}"), + } + } + let environment_info = responses + .remove(&RequestId::Integer(2)) + .expect("environment/info response"); + let _: EnvironmentInfo = + serde_json::from_value(environment_info).expect("decode environment/info response"); + let read_directory = responses + .remove(&RequestId::Integer(3)) + .expect("fs/readDirectory response"); + let _: FsReadDirectoryResponse = + serde_json::from_value(read_directory).expect("decode fs/readDirectory response"); + assert!(responses.is_empty()); + + drop(writer); + drop(lines); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + #[tokio::test] async fn transport_disconnect_detaches_session_during_in_flight_read() { let registry = SessionRegistry::new(); @@ -529,6 +543,12 @@ mod tests { writer.write_all(b"\n").await.expect("write newline"); } + async fn write_batch(writer: &mut DuplexStream, messages: &[JSONRPCMessage]) { + let encoded = serde_json::to_vec(messages).expect("serialize JSON-RPC batch"); + writer.write_all(&encoded).await.expect("write batch"); + writer.write_all(b"\n").await.expect("write newline"); + } + async fn read_response( lines: &mut Lines>, expected_id: i64, @@ -548,6 +568,15 @@ mod tests { } } + async fn read_batch(lines: &mut Lines>) -> Vec { + let line = lines + .next_line() + .await + .expect("read batch") + .expect("batch line"); + serde_json::from_str(&line).expect("decode JSON-RPC batch") + } + fn exec_params(process_id: ProcessId) -> ExecParams { let mut env = HashMap::new(); if let Some(path) = std::env::var_os("PATH") { diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 891a0885de79..8f48aeaf99b7 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -93,7 +93,7 @@ pub(crate) fn build_router() -> RpcRouter { handler.terminate(params).await }, ); - router.concurrent_request( + router.request( FS_READ_FILE_METHOD, |handler: Arc, params: FsReadFileParams| async move { handler.fs_read_file(params).await @@ -129,19 +129,19 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_create_directory(params).await }, ); - router.concurrent_request( + router.request( FS_GET_METADATA_METHOD, |handler: Arc, params: FsGetMetadataParams| async move { handler.fs_get_metadata(params).await }, ); - router.concurrent_request( + router.request( FS_CANONICALIZE_METHOD, |handler: Arc, params: FsCanonicalizeParams| async move { handler.fs_canonicalize(params).await }, ); - router.concurrent_request( + router.request( FS_READ_DIRECTORY_METHOD, |handler: Arc, params: FsReadDirectoryParams| async move { handler.fs_read_directory(params).await diff --git a/codex-rs/exec-server/tests/file_system/shared.rs b/codex-rs/exec-server/tests/file_system/shared.rs index 2d7cc1c668d9..08ba2efdd81b 100644 --- a/codex-rs/exec-server/tests/file_system/shared.rs +++ b/codex-rs/exec-server/tests/file_system/shared.rs @@ -4,8 +4,6 @@ use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FILE_READ_CHUNK_SIZE; use codex_exec_server::FileMetadata; -use codex_exec_server::FileSystemOperation; -use codex_exec_server::FileSystemOperationOutput; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; use codex_protocol::models::AdditionalPermissionProfile; @@ -198,55 +196,6 @@ async fn file_system_read_file_returns_bytes( Ok(()) } -#[test_case(FileSystemImplementation::Local ; "local")] -#[test_case(FileSystemImplementation::Remote ; "remote")] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn file_system_batch_preserves_operation_order( - implementation: FileSystemImplementation, -) -> Result<()> { - let context = create_file_system_context(implementation).await?; - let file_system = context.file_system; - - let tmp = TempDir::new()?; - let first_path = tmp.path().join("first.txt"); - let second_path = tmp.path().join("second.txt"); - std::fs::write(&first_path, "first")?; - std::fs::write(&second_path, "second")?; - let first_uri = PathUri::from_path(&first_path)?; - let second_uri = PathUri::from_path(&second_path)?; - - let outputs = file_system - .execute_batch( - vec![ - FileSystemOperation::ReadFile { - path: first_uri.clone(), - }, - FileSystemOperation::Canonicalize { - path: second_uri.clone(), - }, - FileSystemOperation::ReadFile { path: second_uri }, - ], - /*sandbox*/ None, - ) - .await - .with_context(|| format!("mode={implementation}"))? - .into_iter() - .collect::>>()?; - - assert_eq!( - outputs, - vec![ - FileSystemOperationOutput::ReadFile(b"first".to_vec()), - FileSystemOperationOutput::Canonicalize(PathUri::from_path(std::fs::canonicalize( - second_path - )?)?), - FileSystemOperationOutput::ReadFile(b"second".to_vec()), - ] - ); - - Ok(()) -} - #[test_case(FileSystemImplementation::Local ; "local")] #[test_case(FileSystemImplementation::Remote ; "remote")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/file-system/Cargo.toml b/codex-rs/file-system/Cargo.toml index 382eacb43d33..68a63620c886 100644 --- a/codex-rs/file-system/Cargo.toml +++ b/codex-rs/file-system/Cargo.toml @@ -14,6 +14,7 @@ codex-utils-absolute-path = { workspace = true } codex-utils-path-uri = { workspace = true } futures = { workspace = true } serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } [lib] test = false diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index 1fd1c3ebf378..18ecfc0cca26 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -12,7 +12,6 @@ use codex_protocol::protocol::SandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::Stream; -use futures::StreamExt; use std::future::Future; use std::io; use std::path::Path; @@ -22,7 +21,6 @@ use std::task::Poll; /// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`]. pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024; -const FILE_SYSTEM_BATCH_CONCURRENCY: usize = 32; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct CreateDirectoryOptions { @@ -58,25 +56,13 @@ pub struct ReadDirectoryEntry { pub is_file: bool, } -/// One independent, read-only filesystem operation eligible for batching. -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum FileSystemOperation { - Canonicalize { path: PathUri }, - ReadFile { path: PathUri }, - GetMetadata { path: PathUri }, - ReadDirectory { path: PathUri }, -} - -/// Typed output for a [`FileSystemOperation`]. -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum FileSystemOperationOutput { - Canonicalize(PathUri), - ReadFile(Vec), - GetMetadata(FileMetadata), - ReadDirectory(Vec), +#[derive(Clone, Debug, PartialEq)] +pub struct ExecutorRpcBatchCall { + pub method: String, + pub params: serde_json::Value, } -pub type FileSystemOperationResult = FileSystemResult; +pub type ExecutorRpcBatchResult = FileSystemResult; #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] @@ -286,61 +272,20 @@ pub trait ExecutorFileSystem: Send + Sync { sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, ()>; - /// Executes independent read-only filesystem operations concurrently. + /// Sends independent executor RPC requests as one batch. /// - /// Results retain request order. Operations have no ordering or transactional relationship; - /// callers must put dependent operations in separate batches. - fn execute_batch<'a>( + /// The executor may execute entries concurrently and does not guarantee execution order. + /// Callers must only batch requests that are safe to run in any order, and must send dependent + /// requests separately. Results are matched back to the input order by request id. + fn execute_rpc_batch<'a>( &'a self, - operations: Vec, - sandbox: Option<&'a FileSystemSandboxContext>, - ) -> ExecutorFileSystemFuture<'a, Vec> { - Box::pin(execute_batch_with_scalar_operations( - self, operations, sandbox, - )) - } -} - -/// Default batch implementation used by local filesystems and by remote fallbacks. -pub async fn execute_batch_with_scalar_operations( - file_system: &F, - operations: Vec, - sandbox: Option<&FileSystemSandboxContext>, -) -> FileSystemResult> -where - F: ExecutorFileSystem + ?Sized, -{ - Ok(futures::stream::iter(operations) - .map(|operation| execute_scalar_operation(file_system, operation, sandbox)) - .buffered(FILE_SYSTEM_BATCH_CONCURRENCY) - .collect() - .await) -} - -async fn execute_scalar_operation( - file_system: &F, - operation: FileSystemOperation, - sandbox: Option<&FileSystemSandboxContext>, -) -> FileSystemOperationResult -where - F: ExecutorFileSystem + ?Sized, -{ - match operation { - FileSystemOperation::Canonicalize { path } => file_system - .canonicalize(&path, sandbox) - .await - .map(FileSystemOperationOutput::Canonicalize), - FileSystemOperation::ReadFile { path } => file_system - .read_file(&path, sandbox) - .await - .map(FileSystemOperationOutput::ReadFile), - FileSystemOperation::GetMetadata { path } => file_system - .get_metadata(&path, sandbox) - .await - .map(FileSystemOperationOutput::GetMetadata), - FileSystemOperation::ReadDirectory { path } => file_system - .read_directory(&path, sandbox) - .await - .map(FileSystemOperationOutput::ReadDirectory), + _calls: Vec, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(async move { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "filesystem does not support RPC batching", + )) + }) } } From 10596f252bd5ec87e66eeb39f8022facb944ec2e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Sun, 21 Jun 2026 11:45:58 +0200 Subject: [PATCH 5/7] Use generic RPC batches for skill discovery --- codex-rs/core-skills/src/loader.rs | 250 +++++++++++------- .../utils/plugins/src/plugin_namespace.rs | 72 ++++- 2 files changed, 215 insertions(+), 107 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 58641d4bd882..1953eb605342 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -13,8 +13,21 @@ use codex_config::default_project_root_markers; use codex_config::merge_toml_values; use codex_config::project_root_markers_from_config; use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::FileSystemOperation; -use codex_exec_server::FileSystemOperationOutput; +use codex_exec_server::ExecutorRpcBatchCall; +use codex_exec_server::ExecutorRpcBatchResult; +use codex_exec_server::FS_CANONICALIZE_METHOD; +use codex_exec_server::FS_GET_METADATA_METHOD; +use codex_exec_server::FS_READ_DIRECTORY_METHOD; +use codex_exec_server::FS_READ_FILE_METHOD; +use codex_exec_server::FileMetadata; +use codex_exec_server::FsCanonicalizeParams; +use codex_exec_server::FsCanonicalizeResponse; +use codex_exec_server::FsGetMetadataParams; +use codex_exec_server::FsGetMetadataResponse; +use codex_exec_server::FsReadDirectoryParams; +use codex_exec_server::FsReadDirectoryResponse; +use codex_exec_server::FsReadFileParams; +use codex_exec_server::FsReadFileResponse; use codex_exec_server::LOCAL_FS; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; @@ -25,6 +38,8 @@ use codex_utils_plugins::PluginNamespaceResolver; use codex_utils_plugins::PluginSkillRoot; use dirs::home_dir; use serde::Deserialize; +use serde::Serialize; +use serde::de::DeserializeOwned; use std::collections::HashSet; use std::error::Error; use std::fmt; @@ -534,17 +549,26 @@ async fn discover_skills_under_root( while !queue.is_empty() { let dirs = std::mem::take(&mut queue); - let directory_results = match fs - .execute_batch( - dirs.iter() - .map(|(dir, _)| FileSystemOperation::ReadDirectory { + let directory_calls = match dirs + .iter() + .map(|(dir, _)| { + rpc_batch_call( + FS_READ_DIRECTORY_METHOD, + FsReadDirectoryParams { path: PathUri::from_abs_path(dir), - }) - .collect(), - /*sandbox*/ None, - ) - .await + sandbox: None, + }, + ) + }) + .collect::>>() { + Ok(calls) => calls, + Err(err) => { + error!("failed to build batched skills directory reads: {err:#}"); + break; + } + }; + let directory_results = match fs.execute_rpc_batch(directory_calls).await { Ok(results) => results, Err(err) => { error!("failed to batch-read skills directories: {err:#}"); @@ -554,18 +578,15 @@ async fn discover_skills_under_root( let mut candidates = Vec::new(); for ((dir, depth), result) in dirs.into_iter().zip(directory_results) { - let entries = match result { - Ok(FileSystemOperationOutput::ReadDirectory(entries)) => entries, - Ok(_) => { - error!("filesystem returned the wrong result for {}", dir.display()); - continue; - } - Err(err) => { - error!("failed to read skills dir {}: {err:#}", dir.display()); - continue; - } - }; - candidates.extend(entries.into_iter().filter_map(|entry| { + let response: FsReadDirectoryResponse = + match decode_rpc_batch_result(Some(result), "read skills dir") { + Ok(response) => response, + Err(err) => { + error!("failed to read skills dir {}: {err:#}", dir.display()); + continue; + } + }; + candidates.extend(response.entries.into_iter().filter_map(|entry| { let file_name = entry.file_name; if file_name.starts_with('.') { return None; @@ -574,18 +595,26 @@ async fn discover_skills_under_root( })); } - let metadata_results = match fs - .execute_batch( - candidates - .iter() - .map(|(path, _, _)| FileSystemOperation::GetMetadata { + let metadata_calls = match candidates + .iter() + .map(|(path, _, _)| { + rpc_batch_call( + FS_GET_METADATA_METHOD, + FsGetMetadataParams { path: PathUri::from_abs_path(path), - }) - .collect(), - /*sandbox*/ None, - ) - .await + sandbox: None, + }, + ) + }) + .collect::>>() { + Ok(calls) => calls, + Err(err) => { + error!("failed to build batched skills path stats: {err:#}"); + break; + } + }; + let metadata_results = match fs.execute_rpc_batch(metadata_calls).await { Ok(results) => results, Err(err) => { error!("failed to batch-stat skills paths: {err:#}"); @@ -595,19 +624,21 @@ async fn discover_skills_under_root( let mut directories = Vec::new(); for ((path, depth, file_name), result) in candidates.into_iter().zip(metadata_results) { - let metadata = match result { - Ok(FileSystemOperationOutput::GetMetadata(metadata)) => metadata, - Ok(_) => { - error!( - "filesystem returned the wrong result for {}", - path.display() - ); - continue; - } - Err(err) => { - error!("failed to stat skills path {}: {err:#}", path.display()); - continue; - } + let response: FsGetMetadataResponse = + match decode_rpc_batch_result(Some(result), "stat skills path") { + Ok(response) => response, + Err(err) => { + error!("failed to stat skills path {}: {err:#}", path.display()); + continue; + } + }; + let metadata = FileMetadata { + is_directory: response.is_directory, + is_file: response.is_file, + is_symlink: response.is_symlink, + size: response.size, + created_at_ms: response.created_at_ms, + modified_at_ms: response.modified_at_ms, }; if metadata.is_symlink { if follow_symlinks && metadata.is_directory { @@ -620,18 +651,26 @@ async fn discover_skills_under_root( } } - let canonical_results = match fs - .execute_batch( - directories - .iter() - .map(|(path, _)| FileSystemOperation::Canonicalize { + let canonical_calls = match directories + .iter() + .map(|(path, _)| { + rpc_batch_call( + FS_CANONICALIZE_METHOD, + FsCanonicalizeParams { path: PathUri::from_abs_path(path), - }) - .collect(), - /*sandbox*/ None, - ) - .await + sandbox: None, + }, + ) + }) + .collect::>>() { + Ok(calls) => calls, + Err(err) => { + error!("failed to build batched skills directory canonicalization: {err:#}"); + break; + } + }; + let canonical_results = match fs.execute_rpc_batch(canonical_calls).await { Ok(results) => results, Err(err) => { error!("failed to batch-canonicalize skills directories: {err:#}"); @@ -639,18 +678,18 @@ async fn discover_skills_under_root( } }; for ((path, depth), result) in directories.into_iter().zip(canonical_results) { - let resolved_dir = match result { - Ok(FileSystemOperationOutput::Canonicalize(resolved_path)) => { - resolved_path.to_abs_path().unwrap_or_else(|_| path.clone()) - } - Ok(_) => { + let resolved_dir = match decode_rpc_batch_result::( + Some(result), + "canonicalize skills dir", + ) { + Ok(response) => response.path.to_abs_path().unwrap_or_else(|_| path.clone()), + Err(err) => { error!( - "filesystem returned the wrong result for {}", + "failed to canonicalize skills dir {}: {err:#}", path.display() ); path } - Err(_) => path, }; enqueue_dir( &mut queue, @@ -710,7 +749,7 @@ async fn preload_skill_files( fs: &dyn ExecutorFileSystem, skill_paths: &[AbsolutePathBuf], ) -> io::Result> { - let operations = skill_paths + let calls = skill_paths .iter() .flat_map(|path| { let metadata_path = path @@ -722,34 +761,43 @@ async fn preload_skill_files( }) .unwrap_or_else(|| path.clone()); [ - FileSystemOperation::ReadFile { - path: PathUri::from_abs_path(path), - }, - FileSystemOperation::ReadFile { - path: PathUri::from_abs_path(&metadata_path), - }, - FileSystemOperation::Canonicalize { - path: PathUri::from_abs_path(path), - }, + rpc_batch_call( + FS_READ_FILE_METHOD, + FsReadFileParams { + path: PathUri::from_abs_path(path), + sandbox: None, + }, + ), + rpc_batch_call( + FS_READ_FILE_METHOD, + FsReadFileParams { + path: PathUri::from_abs_path(&metadata_path), + sandbox: None, + }, + ), + rpc_batch_call( + FS_CANONICALIZE_METHOD, + FsCanonicalizeParams { + path: PathUri::from_abs_path(path), + sandbox: None, + }, + ), ] }) - .collect(); - let mut results = fs - .execute_batch(operations, /*sandbox*/ None) - .await? - .into_iter(); + .collect::>>()?; + let mut results = fs.execute_rpc_batch(calls).await?.into_iter(); Ok(skill_paths .iter() .map(|path| { - let contents = decode_text_operation_result(results.next(), "read skill file"); - let metadata_contents = - decode_text_operation_result(results.next(), "read skill metadata"); - let resolved_path = match results.next() { - Some(Ok(FileSystemOperationOutput::Canonicalize(resolved_path))) => { - resolved_path.to_abs_path().unwrap_or_else(|_| path.clone()) - } - Some(Ok(_)) | Some(Err(_)) | None => path.clone(), + let contents = decode_text_rpc_result(results.next(), "read skill file"); + let metadata_contents = decode_text_rpc_result(results.next(), "read skill metadata"); + let resolved_path = match decode_rpc_batch_result::( + results.next(), + "canonicalize skill path", + ) { + Ok(response) => response.path.to_abs_path().unwrap_or_else(|_| path.clone()), + Err(_) => path.clone(), }; PreloadedSkillFile { contents, @@ -760,17 +808,20 @@ async fn preload_skill_files( .collect()) } -fn decode_text_operation_result( - result: Option>, +fn rpc_batch_call(method: &str, params: P) -> io::Result { + Ok(ExecutorRpcBatchCall { + method: method.to_string(), + params: serde_json::to_value(params).map_err(io::Error::other)?, + }) +} + +fn decode_rpc_batch_result( + result: Option, operation: &str, -) -> io::Result { +) -> io::Result { match result { - Some(Ok(FileSystemOperationOutput::ReadFile(contents))) => String::from_utf8(contents) + Some(Ok(value)) => serde_json::from_value(value) .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)), - Some(Ok(_)) => Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("filesystem returned the wrong result for {operation}"), - )), Some(Err(error)) => Err(error), None => Err(io::Error::new( io::ErrorKind::InvalidData, @@ -779,6 +830,17 @@ fn decode_text_operation_result( } } +fn decode_text_rpc_result( + result: Option, + operation: &str, +) -> io::Result { + let response: FsReadFileResponse = decode_rpc_batch_result(result, operation)?; + let contents = response + .into_bytes() + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + String::from_utf8(contents).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) +} + fn parse_skill_file( path: &AbsolutePathBuf, scope: SkillScope, diff --git a/codex-rs/utils/plugins/src/plugin_namespace.rs b/codex-rs/utils/plugins/src/plugin_namespace.rs index c829d4c26d37..4d3573e13dd6 100644 --- a/codex-rs/utils/plugins/src/plugin_namespace.rs +++ b/codex-rs/utils/plugins/src/plugin_namespace.rs @@ -1,12 +1,19 @@ //! Resolve plugin namespace from skill file paths by walking ancestors for `plugin.json`. use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::FileSystemOperation; -use codex_exec_server::FileSystemOperationOutput; +use codex_exec_server::ExecutorRpcBatchCall; +use codex_exec_server::ExecutorRpcBatchResult; +use codex_exec_server::FS_GET_METADATA_METHOD; +use codex_exec_server::FS_READ_FILE_METHOD; +use codex_exec_server::FsGetMetadataParams; +use codex_exec_server::FsGetMetadataResponse; +use codex_exec_server::FsReadFileParams; +use codex_exec_server::FsReadFileResponse; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use std::collections::HashMap; use std::collections::HashSet; +use std::io; use std::path::Path; use std::path::PathBuf; @@ -54,7 +61,7 @@ impl PluginNamespaceResolver { .flat_map(AbsolutePathBuf::ancestors) .filter(|ancestor| seen.insert(ancestor.clone())) .collect::>(); - let operations = roots + let calls = roots .iter() .flat_map(|root| { DISCOVERABLE_PLUGIN_MANIFEST_PATHS @@ -62,14 +69,26 @@ impl PluginNamespaceResolver { .flat_map(|relative_path| { let path = PathUri::from_abs_path(&root.join(relative_path)); [ - FileSystemOperation::GetMetadata { path: path.clone() }, - FileSystemOperation::ReadFile { path }, + rpc_batch_call( + FS_GET_METADATA_METHOD, + FsGetMetadataParams { + path: path.clone(), + sandbox: None, + }, + ), + rpc_batch_call( + FS_READ_FILE_METHOD, + FsReadFileParams { + path, + sandbox: None, + }, + ), ] }) }) - .collect(); + .collect::>>(); let mut results = fs - .execute_batch(operations, /*sandbox*/ None) + .execute_rpc_batch(calls.unwrap_or_default()) .await .unwrap_or_default() .into_iter(); @@ -81,15 +100,17 @@ impl PluginNamespaceResolver { for _ in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { let metadata_result = results.next(); let contents_result = results.next(); - let is_file = matches!( + let is_file = decode_rpc_batch_result::( metadata_result, - Some(Ok(FileSystemOperationOutput::GetMetadata(metadata))) - if metadata.is_file - ); + "stat plugin manifest", + ) + .is_ok_and(|metadata| metadata.is_file); if !selected && is_file { selected = true; - if let Some(Ok(FileSystemOperationOutput::ReadFile(contents))) = - contents_result + if let Ok(response) = decode_rpc_batch_result::( + contents_result, + "read plugin manifest", + ) && let Ok(contents) = response.into_bytes() { name = plugin_manifest_name(&root, &contents); } @@ -109,6 +130,31 @@ impl PluginNamespaceResolver { } } +fn rpc_batch_call( + method: &str, + params: P, +) -> io::Result { + Ok(ExecutorRpcBatchCall { + method: method.to_string(), + params: serde_json::to_value(params).map_err(io::Error::other)?, + }) +} + +fn decode_rpc_batch_result( + result: Option, + operation: &str, +) -> io::Result { + match result { + Some(Ok(value)) => serde_json::from_value(value) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)), + Some(Err(error)) => Err(error), + None => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("filesystem returned no result for {operation}"), + )), + } +} + /// Returns the plugin manifest `name` for the nearest ancestor of `path` that contains a valid /// plugin manifest (same `name` rules as full manifest loading in codex-core). pub async fn plugin_namespace_for_skill_path( From cb3948fc2fd0b98efa6f99d31aaffe3d173caa36 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Sun, 21 Jun 2026 12:21:07 +0200 Subject: [PATCH 6/7] Drop skill discovery batching --- codex-rs/core-skills/src/loader.rs | 442 ++++++------------ codex-rs/utils/plugins/src/lib.rs | 1 - .../utils/plugins/src/plugin_namespace.rs | 185 ++------ 3 files changed, 160 insertions(+), 468 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 1953eb605342..2a3a3189be9b 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -13,34 +13,18 @@ use codex_config::default_project_root_markers; use codex_config::merge_toml_values; use codex_config::project_root_markers_from_config; use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::ExecutorRpcBatchCall; -use codex_exec_server::ExecutorRpcBatchResult; -use codex_exec_server::FS_CANONICALIZE_METHOD; -use codex_exec_server::FS_GET_METADATA_METHOD; -use codex_exec_server::FS_READ_DIRECTORY_METHOD; -use codex_exec_server::FS_READ_FILE_METHOD; -use codex_exec_server::FileMetadata; -use codex_exec_server::FsCanonicalizeParams; -use codex_exec_server::FsCanonicalizeResponse; -use codex_exec_server::FsGetMetadataParams; -use codex_exec_server::FsGetMetadataResponse; -use codex_exec_server::FsReadDirectoryParams; -use codex_exec_server::FsReadDirectoryResponse; -use codex_exec_server::FsReadFileParams; -use codex_exec_server::FsReadFileResponse; use codex_exec_server::LOCAL_FS; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::AbsolutePathBufGuard; use codex_utils_path_uri::PathUri; -use codex_utils_plugins::PluginNamespaceResolver; use codex_utils_plugins::PluginSkillRoot; +use codex_utils_plugins::plugin_namespace_for_skill_path; use dirs::home_dir; use serde::Deserialize; -use serde::Serialize; -use serde::de::DeserializeOwned; use std::collections::HashSet; +use std::collections::VecDeque; use std::error::Error; use std::fmt; use std::io; @@ -84,12 +68,6 @@ struct LoadedSkillMetadata { policy: Option, } -struct PreloadedSkillFile { - contents: io::Result, - metadata_contents: io::Result, - resolved_path: AbsolutePathBuf, -} - #[derive(Debug, Default, Deserialize)] struct Interface { display_name: Option, @@ -516,7 +494,7 @@ async fn discover_skills_under_root( } fn enqueue_dir( - queue: &mut Vec<(AbsolutePathBuf, usize)>, + queue: &mut VecDeque<(AbsolutePathBuf, usize)>, visited_dirs: &mut HashSet, truncated_by_dir_limit: &mut bool, path: AbsolutePathBuf, @@ -530,7 +508,7 @@ async fn discover_skills_under_root( return; } if visited_dirs.insert(path.clone()) { - queue.push((path, depth)); + queue.push_back((path, depth)); } } @@ -543,197 +521,102 @@ async fn discover_skills_under_root( let mut visited_dirs: HashSet = HashSet::new(); visited_dirs.insert(root.clone()); - let mut queue = vec![(root.clone(), 0)]; + let mut queue: VecDeque<(AbsolutePathBuf, usize)> = VecDeque::from([(root.clone(), 0)]); let mut truncated_by_dir_limit = false; - let mut skill_paths = Vec::new(); - - while !queue.is_empty() { - let dirs = std::mem::take(&mut queue); - let directory_calls = match dirs - .iter() - .map(|(dir, _)| { - rpc_batch_call( - FS_READ_DIRECTORY_METHOD, - FsReadDirectoryParams { - path: PathUri::from_abs_path(dir), - sandbox: None, - }, - ) - }) - .collect::>>() - { - Ok(calls) => calls, - Err(err) => { - error!("failed to build batched skills directory reads: {err:#}"); - break; + + while let Some((dir, depth)) = queue.pop_front() { + let dir_uri = PathUri::from_abs_path(&dir); + let entries = match fs.read_directory(&dir_uri, /*sandbox*/ None).await { + Ok(entries) => entries, + Err(e) => { + error!("failed to read skills dir {}: {e:#}", dir.display()); + continue; } }; - let directory_results = match fs.execute_rpc_batch(directory_calls).await { - Ok(results) => results, - Err(err) => { - error!("failed to batch-read skills directories: {err:#}"); - break; + + for entry in entries { + let file_name = entry.file_name; + if file_name.starts_with('.') { + continue; } - }; - let mut candidates = Vec::new(); - for ((dir, depth), result) in dirs.into_iter().zip(directory_results) { - let response: FsReadDirectoryResponse = - match decode_rpc_batch_result(Some(result), "read skills dir") { - Ok(response) => response, + let path = dir.join(&file_name); + let path_uri = PathUri::from_abs_path(&path); + let metadata = match fs.get_metadata(&path_uri, /*sandbox*/ None).await { + Ok(metadata) => metadata, + Err(e) => { + error!("failed to stat skills path {}: {e:#}", path.display()); + continue; + } + }; + + if metadata.is_symlink { + if !follow_symlinks { + continue; + } + match fs.read_directory(&path_uri, /*sandbox*/ None).await { + Ok(_) => { + let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; + enqueue_dir( + &mut queue, + &mut visited_dirs, + &mut truncated_by_dir_limit, + resolved_dir, + depth + 1, + ); + } + Err(err) + if matches!( + err.kind(), + io::ErrorKind::NotADirectory | io::ErrorKind::NotFound + ) => {} Err(err) => { - error!("failed to read skills dir {}: {err:#}", dir.display()); - continue; + error!( + "failed to read skills symlink dir {}: {err:#}", + path.display() + ); } - }; - candidates.extend(response.entries.into_iter().filter_map(|entry| { - let file_name = entry.file_name; - if file_name.starts_with('.') { - return None; } - Some((dir.join(&file_name), depth, file_name)) - })); - } - - let metadata_calls = match candidates - .iter() - .map(|(path, _, _)| { - rpc_batch_call( - FS_GET_METADATA_METHOD, - FsGetMetadataParams { - path: PathUri::from_abs_path(path), - sandbox: None, - }, - ) - }) - .collect::>>() - { - Ok(calls) => calls, - Err(err) => { - error!("failed to build batched skills path stats: {err:#}"); - break; + continue; } - }; - let metadata_results = match fs.execute_rpc_batch(metadata_calls).await { - Ok(results) => results, - Err(err) => { - error!("failed to batch-stat skills paths: {err:#}"); - break; + + if metadata.is_directory { + let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; + enqueue_dir( + &mut queue, + &mut visited_dirs, + &mut truncated_by_dir_limit, + resolved_dir, + depth + 1, + ); + continue; } - }; - let mut directories = Vec::new(); - for ((path, depth, file_name), result) in candidates.into_iter().zip(metadata_results) { - let response: FsGetMetadataResponse = - match decode_rpc_batch_result(Some(result), "stat skills path") { - Ok(response) => response, + if metadata.is_file && file_name == SKILLS_FILENAME { + match parse_skill_file( + fs, + &path, + scope, + plugin_id, + plugin_namespace, + plugin_root.as_ref(), + ) + .await + { + Ok(skill) => { + outcome.skills.push(skill); + } Err(err) => { - error!("failed to stat skills path {}: {err:#}", path.display()); - continue; + if scope != SkillScope::System { + outcome.errors.push(SkillError { + path: path.clone(), + message: err.to_string(), + }); + } } - }; - let metadata = FileMetadata { - is_directory: response.is_directory, - is_file: response.is_file, - is_symlink: response.is_symlink, - size: response.size, - created_at_ms: response.created_at_ms, - modified_at_ms: response.modified_at_ms, - }; - if metadata.is_symlink { - if follow_symlinks && metadata.is_directory { - directories.push((path, depth + 1)); } - } else if metadata.is_directory { - directories.push((path, depth + 1)); - } else if metadata.is_file && file_name == SKILLS_FILENAME { - skill_paths.push(path); } } - - let canonical_calls = match directories - .iter() - .map(|(path, _)| { - rpc_batch_call( - FS_CANONICALIZE_METHOD, - FsCanonicalizeParams { - path: PathUri::from_abs_path(path), - sandbox: None, - }, - ) - }) - .collect::>>() - { - Ok(calls) => calls, - Err(err) => { - error!("failed to build batched skills directory canonicalization: {err:#}"); - break; - } - }; - let canonical_results = match fs.execute_rpc_batch(canonical_calls).await { - Ok(results) => results, - Err(err) => { - error!("failed to batch-canonicalize skills directories: {err:#}"); - break; - } - }; - for ((path, depth), result) in directories.into_iter().zip(canonical_results) { - let resolved_dir = match decode_rpc_batch_result::( - Some(result), - "canonicalize skills dir", - ) { - Ok(response) => response.path.to_abs_path().unwrap_or_else(|_| path.clone()), - Err(err) => { - error!( - "failed to canonicalize skills dir {}: {err:#}", - path.display() - ); - path - } - }; - enqueue_dir( - &mut queue, - &mut visited_dirs, - &mut truncated_by_dir_limit, - resolved_dir, - depth, - ); - } - } - - let plugin_namespace_resolver = if plugin_namespace.is_none() { - Some(PluginNamespaceResolver::load(fs, &skill_paths).await) - } else { - None - }; - let preloaded_skills = match preload_skill_files(fs, &skill_paths).await { - Ok(preloaded_skills) => preloaded_skills, - Err(err) => { - error!("failed to batch-read skill files: {err:#}"); - Vec::new() - } - }; - for (path, preloaded) in skill_paths.into_iter().zip(preloaded_skills) { - let plugin_namespace = plugin_namespace.or_else(|| { - plugin_namespace_resolver - .as_ref() - .and_then(|resolver| resolver.resolve(&path)) - }); - match parse_skill_file( - &path, - scope, - plugin_id, - plugin_namespace, - plugin_root.as_ref(), - preloaded, - ) { - Ok(skill) => outcome.skills.push(skill), - Err(err) if scope != SkillScope::System => outcome.errors.push(SkillError { - path, - message: err.to_string(), - }), - Err(_) => {} - } } if truncated_by_dir_limit { @@ -745,116 +628,19 @@ async fn discover_skills_under_root( } } -async fn preload_skill_files( +async fn parse_skill_file( fs: &dyn ExecutorFileSystem, - skill_paths: &[AbsolutePathBuf], -) -> io::Result> { - let calls = skill_paths - .iter() - .flat_map(|path| { - let metadata_path = path - .parent() - .map(|parent| { - parent - .join(SKILLS_METADATA_DIR) - .join(SKILLS_METADATA_FILENAME) - }) - .unwrap_or_else(|| path.clone()); - [ - rpc_batch_call( - FS_READ_FILE_METHOD, - FsReadFileParams { - path: PathUri::from_abs_path(path), - sandbox: None, - }, - ), - rpc_batch_call( - FS_READ_FILE_METHOD, - FsReadFileParams { - path: PathUri::from_abs_path(&metadata_path), - sandbox: None, - }, - ), - rpc_batch_call( - FS_CANONICALIZE_METHOD, - FsCanonicalizeParams { - path: PathUri::from_abs_path(path), - sandbox: None, - }, - ), - ] - }) - .collect::>>()?; - let mut results = fs.execute_rpc_batch(calls).await?.into_iter(); - - Ok(skill_paths - .iter() - .map(|path| { - let contents = decode_text_rpc_result(results.next(), "read skill file"); - let metadata_contents = decode_text_rpc_result(results.next(), "read skill metadata"); - let resolved_path = match decode_rpc_batch_result::( - results.next(), - "canonicalize skill path", - ) { - Ok(response) => response.path.to_abs_path().unwrap_or_else(|_| path.clone()), - Err(_) => path.clone(), - }; - PreloadedSkillFile { - contents, - metadata_contents, - resolved_path, - } - }) - .collect()) -} - -fn rpc_batch_call(method: &str, params: P) -> io::Result { - Ok(ExecutorRpcBatchCall { - method: method.to_string(), - params: serde_json::to_value(params).map_err(io::Error::other)?, - }) -} - -fn decode_rpc_batch_result( - result: Option, - operation: &str, -) -> io::Result { - match result { - Some(Ok(value)) => serde_json::from_value(value) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)), - Some(Err(error)) => Err(error), - None => Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("filesystem returned no result for {operation}"), - )), - } -} - -fn decode_text_rpc_result( - result: Option, - operation: &str, -) -> io::Result { - let response: FsReadFileResponse = decode_rpc_batch_result(result, operation)?; - let contents = response - .into_bytes() - .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; - String::from_utf8(contents).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) -} - -fn parse_skill_file( path: &AbsolutePathBuf, scope: SkillScope, plugin_id: Option<&str>, plugin_namespace: Option<&str>, plugin_root: Option<&AbsolutePathBuf>, - preloaded: PreloadedSkillFile, ) -> Result { - let PreloadedSkillFile { - contents, - metadata_contents, - resolved_path, - } = preloaded; - let contents = contents.map_err(SkillParseError::Read)?; + let path_uri = PathUri::from_abs_path(path); + let contents = fs + .read_file_text(&path_uri, /*sandbox*/ None) + .await + .map_err(SkillParseError::Read)?; let frontmatter = extract_frontmatter(&contents).ok_or(SkillParseError::MissingFrontmatter)?; @@ -878,9 +664,7 @@ fn parse_skill_file( .map(sanitize_single_line) .filter(|value| !value.is_empty()) .unwrap_or_else(|| default_skill_name(path)); - let name = plugin_namespace - .map(|namespace| format!("{namespace}:{base_name}")) - .unwrap_or_else(|| base_name.clone()); + let name = namespaced_skill_name(fs, path, &base_name, plugin_namespace).await; let description = parsed .description .as_deref() @@ -896,7 +680,7 @@ fn parse_skill_file( interface, dependencies, policy, - } = load_skill_metadata(path, plugin_root, metadata_contents); + } = load_skill_metadata(fs, path, plugin_root).await; validate_len(&base_name, MAX_NAME_LEN, "name")?; validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name")?; @@ -909,6 +693,8 @@ fn parse_skill_file( )?; } + let resolved_path = canonicalize_for_skill_identity(fs, path).await; + Ok(SkillMetadata { name, description, @@ -934,10 +720,25 @@ fn default_skill_name(path: &AbsolutePathBuf) -> String { .unwrap_or_else(|| "skill".to_string()) } -fn load_skill_metadata( +async fn namespaced_skill_name( + fs: &dyn ExecutorFileSystem, + path: &AbsolutePathBuf, + base_name: &str, + plugin_namespace: Option<&str>, +) -> String { + if let Some(plugin_namespace) = plugin_namespace { + return format!("{plugin_namespace}:{base_name}"); + } + plugin_namespace_for_skill_path(fs, path) + .await + .map(|namespace| format!("{namespace}:{base_name}")) + .unwrap_or_else(|| base_name.to_string()) +} + +async fn load_skill_metadata( + fs: &dyn ExecutorFileSystem, skill_path: &AbsolutePathBuf, plugin_root: Option<&AbsolutePathBuf>, - contents: io::Result, ) -> LoadedSkillMetadata { // Fail open: optional metadata should not block loading SKILL.md. let Some(skill_dir) = skill_path.parent() else { @@ -946,11 +747,28 @@ fn load_skill_metadata( let metadata_path = skill_dir .join(SKILLS_METADATA_DIR) .join(SKILLS_METADATA_FILENAME); - let contents = match contents { - Ok(contents) => contents, + let metadata_path_uri = PathUri::from_abs_path(&metadata_path); + match fs.get_metadata(&metadata_path_uri, /*sandbox*/ None).await { + Ok(metadata) if metadata.is_file => {} + Ok(_) => return LoadedSkillMetadata::default(), Err(error) if error.kind() == io::ErrorKind::NotFound => { return LoadedSkillMetadata::default(); } + Err(error) => { + tracing::warn!( + "ignoring {path}: failed to stat {label}: {error}", + path = metadata_path.display(), + label = SKILLS_METADATA_FILENAME + ); + return LoadedSkillMetadata::default(); + } + } + + let contents = match fs + .read_file_text(&metadata_path_uri, /*sandbox*/ None) + .await + { + Ok(contents) => contents, Err(error) => { tracing::warn!( "ignoring {path}: failed to read {label}: {error}", diff --git a/codex-rs/utils/plugins/src/lib.rs b/codex-rs/utils/plugins/src/lib.rs index 238d11f7cb41..203b1f22e41d 100644 --- a/codex-rs/utils/plugins/src/lib.rs +++ b/codex-rs/utils/plugins/src/lib.rs @@ -8,7 +8,6 @@ pub mod mention_syntax; pub mod plugin_namespace; pub use plugin_namespace::DISCOVERABLE_PLUGIN_MANIFEST_PATHS; -pub use plugin_namespace::PluginNamespaceResolver; pub use plugin_namespace::find_plugin_manifest_path; pub use plugin_namespace::plugin_namespace_for_skill_path; diff --git a/codex-rs/utils/plugins/src/plugin_namespace.rs b/codex-rs/utils/plugins/src/plugin_namespace.rs index 4d3573e13dd6..b29ff5c5d684 100644 --- a/codex-rs/utils/plugins/src/plugin_namespace.rs +++ b/codex-rs/utils/plugins/src/plugin_namespace.rs @@ -1,19 +1,8 @@ //! Resolve plugin namespace from skill file paths by walking ancestors for `plugin.json`. use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::ExecutorRpcBatchCall; -use codex_exec_server::ExecutorRpcBatchResult; -use codex_exec_server::FS_GET_METADATA_METHOD; -use codex_exec_server::FS_READ_FILE_METHOD; -use codex_exec_server::FsGetMetadataParams; -use codex_exec_server::FsGetMetadataResponse; -use codex_exec_server::FsReadFileParams; -use codex_exec_server::FsReadFileResponse; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; -use std::collections::HashMap; -use std::collections::HashSet; -use std::io; use std::path::Path; use std::path::PathBuf; @@ -35,8 +24,29 @@ struct RawPluginManifestName { name: String, } -fn plugin_manifest_name(plugin_root: &AbsolutePathBuf, contents: &[u8]) -> Option { - let RawPluginManifestName { name: raw_name } = serde_json::from_slice(contents).ok()?; +async fn plugin_manifest_name( + fs: &dyn ExecutorFileSystem, + plugin_root: &AbsolutePathBuf, +) -> Option { + let mut manifest_path = None; + for relative_path in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { + let candidate = plugin_root.join(relative_path); + let candidate_uri = PathUri::from_abs_path(&candidate); + match fs.get_metadata(&candidate_uri, /*sandbox*/ None).await { + Ok(metadata) if metadata.is_file => { + manifest_path = Some(candidate); + break; + } + Ok(_) | Err(_) => {} + } + } + let manifest_path = manifest_path?; + let manifest_path_uri = PathUri::from_abs_path(&manifest_path); + let contents = fs + .read_file_text(&manifest_path_uri, /*sandbox*/ None) + .await + .ok()?; + let RawPluginManifestName { name: raw_name } = serde_json::from_str(&contents).ok()?; Some( plugin_root .file_name() @@ -47,124 +57,18 @@ fn plugin_manifest_name(plugin_root: &AbsolutePathBuf, contents: &[u8]) -> Optio ) } -/// Resolves plugin namespaces while caching and batching ancestor manifest probes. -pub struct PluginNamespaceResolver { - manifests_by_root: HashMap, -} - -impl PluginNamespaceResolver { - /// Loads plugin manifest names for every ancestor of `paths`. - pub async fn load(fs: &dyn ExecutorFileSystem, paths: &[AbsolutePathBuf]) -> Self { - let mut seen = HashSet::new(); - let roots = paths - .iter() - .flat_map(AbsolutePathBuf::ancestors) - .filter(|ancestor| seen.insert(ancestor.clone())) - .collect::>(); - let calls = roots - .iter() - .flat_map(|root| { - DISCOVERABLE_PLUGIN_MANIFEST_PATHS - .iter() - .flat_map(|relative_path| { - let path = PathUri::from_abs_path(&root.join(relative_path)); - [ - rpc_batch_call( - FS_GET_METADATA_METHOD, - FsGetMetadataParams { - path: path.clone(), - sandbox: None, - }, - ), - rpc_batch_call( - FS_READ_FILE_METHOD, - FsReadFileParams { - path, - sandbox: None, - }, - ), - ] - }) - }) - .collect::>>(); - let mut results = fs - .execute_rpc_batch(calls.unwrap_or_default()) - .await - .unwrap_or_default() - .into_iter(); - let manifests_by_root = roots - .into_iter() - .filter_map(|root| { - let mut name = None; - let mut selected = false; - for _ in DISCOVERABLE_PLUGIN_MANIFEST_PATHS { - let metadata_result = results.next(); - let contents_result = results.next(); - let is_file = decode_rpc_batch_result::( - metadata_result, - "stat plugin manifest", - ) - .is_ok_and(|metadata| metadata.is_file); - if !selected && is_file { - selected = true; - if let Ok(response) = decode_rpc_batch_result::( - contents_result, - "read plugin manifest", - ) && let Ok(contents) = response.into_bytes() - { - name = plugin_manifest_name(&root, &contents); - } - } - } - name.map(|name| (root, name)) - }) - .collect(); - - Self { manifests_by_root } - } - - /// Returns the nearest preloaded plugin namespace for `path`. - pub fn resolve(&self, path: &AbsolutePathBuf) -> Option<&str> { - path.ancestors() - .find_map(|ancestor| self.manifests_by_root.get(&ancestor).map(String::as_str)) - } -} - -fn rpc_batch_call( - method: &str, - params: P, -) -> io::Result { - Ok(ExecutorRpcBatchCall { - method: method.to_string(), - params: serde_json::to_value(params).map_err(io::Error::other)?, - }) -} - -fn decode_rpc_batch_result( - result: Option, - operation: &str, -) -> io::Result { - match result { - Some(Ok(value)) => serde_json::from_value(value) - .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)), - Some(Err(error)) => Err(error), - None => Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("filesystem returned no result for {operation}"), - )), - } -} - /// Returns the plugin manifest `name` for the nearest ancestor of `path` that contains a valid /// plugin manifest (same `name` rules as full manifest loading in codex-core). pub async fn plugin_namespace_for_skill_path( fs: &dyn ExecutorFileSystem, path: &AbsolutePathBuf, ) -> Option { - PluginNamespaceResolver::load(fs, std::slice::from_ref(path)) - .await - .resolve(path) - .map(str::to_string) + for ancestor in path.ancestors() { + if let Some(name) = plugin_manifest_name(fs, &ancestor).await { + return Some(name); + } + } + None } #[cfg(test)] @@ -176,7 +80,6 @@ mod tests { use std::fs; use tempfile::tempdir; - const PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH: &str = ".codex-plugin/plugin.json"; const ALTERNATE_PLUGIN_MANIFEST_RELATIVE_PATH: &str = ".claude-plugin/plugin.json"; #[tokio::test] @@ -188,7 +91,7 @@ mod tests { fs::create_dir_all(skill_path.parent().expect("parent")).expect("mkdir"); fs::create_dir_all(plugin_root.join(".codex-plugin")).expect("mkdir manifest"); fs::write( - plugin_root.join(PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH), + plugin_root.join(".codex-plugin/plugin.json"), r#"{"name":"sample"}"#, ) .expect("write manifest"); @@ -219,32 +122,4 @@ mod tests { ); assert_eq!(find_plugin_manifest_path(&plugin_root), Some(manifest_path)); } - - #[tokio::test] - async fn invalid_primary_manifest_does_not_fall_back_to_alternate() { - let tmp = tempdir().expect("tempdir"); - let plugin_root = tmp.path().join("plugins/sample"); - let skill_path = plugin_root.join("skills/search/SKILL.md"); - let primary_manifest_path = plugin_root.join(PRIMARY_PLUGIN_MANIFEST_RELATIVE_PATH); - let alternate_manifest_path = plugin_root.join(ALTERNATE_PLUGIN_MANIFEST_RELATIVE_PATH); - - fs::create_dir_all(skill_path.parent().expect("parent")).expect("mkdir"); - fs::create_dir_all(primary_manifest_path.parent().expect("manifest parent")) - .expect("mkdir primary manifest"); - fs::create_dir_all(alternate_manifest_path.parent().expect("manifest parent")) - .expect("mkdir alternate manifest"); - fs::write(&primary_manifest_path, "invalid json").expect("write primary manifest"); - fs::write(&alternate_manifest_path, r#"{"name":"sample"}"#) - .expect("write alternate manifest"); - fs::write(&skill_path, "---\ndescription: search\n---\n").expect("write skill"); - - assert_eq!( - plugin_namespace_for_skill_path(LOCAL_FS.as_ref(), &skill_path.abs()).await, - None - ); - assert_eq!( - find_plugin_manifest_path(&plugin_root), - Some(primary_manifest_path) - ); - } } From 80c86785cec56d113d0312e998c0ddeaf30254f4 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Sun, 21 Jun 2026 12:38:48 +0200 Subject: [PATCH 7/7] Batch skill metadata discovery requests --- codex-rs/core-skills/src/loader.rs | 97 +++++++++++-- codex-rs/core-skills/src/loader_tests.rs | 176 +++++++++++++++++++++++ 2 files changed, 264 insertions(+), 9 deletions(-) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 2a3a3189be9b..da0e01300a4b 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -13,6 +13,10 @@ use codex_config::default_project_root_markers; use codex_config::merge_toml_values; use codex_config::project_root_markers_from_config; use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FS_GET_METADATA_METHOD; +use codex_exec_server::FileMetadata; +use codex_exec_server::FsGetMetadataParams; +use codex_exec_server::FsGetMetadataResponse; use codex_exec_server::LOCAL_FS; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; @@ -534,15 +538,26 @@ async fn discover_skills_under_root( } }; - for entry in entries { - let file_name = entry.file_name; - if file_name.starts_with('.') { - continue; - } - - let path = dir.join(&file_name); - let path_uri = PathUri::from_abs_path(&path); - let metadata = match fs.get_metadata(&path_uri, /*sandbox*/ None).await { + let paths = entries + .into_iter() + .filter_map(|entry| { + let file_name = entry.file_name; + if file_name.starts_with('.') { + return None; + } + let path = dir.join(&file_name); + let path_uri = PathUri::from_abs_path(&path); + Some((file_name, path, path_uri)) + }) + .collect::>(); + let mut metadata_results = batched_get_metadata(fs, &paths).await.map(Vec::into_iter); + + for (file_name, path, path_uri) in paths { + let metadata_result = match metadata_results.as_mut().and_then(Iterator::next) { + Some(result) => result, + None => fs.get_metadata(&path_uri, /*sandbox*/ None).await, + }; + let metadata = match metadata_result { Ok(metadata) => metadata, Err(e) => { error!("failed to stat skills path {}: {e:#}", path.display()); @@ -628,6 +643,70 @@ async fn discover_skills_under_root( } } +async fn batched_get_metadata( + fs: &dyn ExecutorFileSystem, + paths: &[(String, AbsolutePathBuf, PathUri)], +) -> Option>> { + if paths.is_empty() { + return Some(Vec::new()); + } + + let mut requests = Vec::with_capacity(paths.len()); + for (_, _, path) in paths { + let params = match serde_json::to_value(FsGetMetadataParams { + path: path.clone(), + sandbox: None, + }) { + Ok(params) => params, + Err(err) => { + error!("failed to encode batched skills path stat: {err:#}"); + return None; + } + }; + requests.push((FS_GET_METADATA_METHOD.to_string(), params)); + } + + let results = match fs.execute_rpc_batch(requests).await { + Ok(results) if results.len() == paths.len() => results, + Ok(results) => { + error!( + "batched skills path stat returned {} results for {} requests", + results.len(), + paths.len() + ); + return None; + } + Err(err) if err.kind() == io::ErrorKind::Unsupported => return None, + Err(err) => { + error!("failed to batch-stat skills paths: {err:#}"); + return None; + } + }; + + Some( + results + .into_iter() + .map(|result| { + let response: FsGetMetadataResponse = + serde_json::from_value(result?).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("fs/getMetadata returned invalid response: {err}"), + ) + })?; + Ok(FileMetadata { + is_directory: response.is_directory, + is_file: response.is_file, + is_symlink: response.is_symlink, + size: response.size, + created_at_ms: response.created_at_ms, + modified_at_ms: response.modified_at_ms, + }) + }) + .collect(), + ) +} + async fn parse_skill_file( fs: &dyn ExecutorFileSystem, path: &AbsolutePathBuf, diff --git a/codex-rs/core-skills/src/loader_tests.rs b/codex-rs/core-skills/src/loader_tests.rs index e22ed33f3df1..be7cbc2175cc 100644 --- a/codex-rs/core-skills/src/loader_tests.rs +++ b/codex-rs/core-skills/src/loader_tests.rs @@ -4,7 +4,14 @@ use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerStack; use codex_config::ConfigRequirements; use codex_config::ConfigRequirementsToml; +use codex_exec_server::CopyOptions; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecutorFileSystemFuture; +use codex_exec_server::FileSystemReadStream; +use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::LOCAL_FS; +use codex_exec_server::ReadDirectoryEntry; +use codex_exec_server::RemoveOptions; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; use codex_utils_absolute_path::AbsolutePathBuf; @@ -16,6 +23,7 @@ use std::fs; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use tempfile::TempDir; use toml::Value as TomlValue; @@ -2157,3 +2165,171 @@ async fn skill_roots_include_admin_with_lowest_priority() { expected.push(SkillScope::Admin); assert_eq!(scopes, expected); } + +#[tokio::test] +async fn batched_get_metadata_sends_one_request_per_path() { + let tmp = tempfile::tempdir().expect("tempdir"); + let first_path = tmp.path().join("first").abs(); + let second_path = tmp.path().join("second").abs(); + let first_uri = PathUri::from_abs_path(&first_path); + let second_uri = PathUri::from_abs_path(&second_path); + let fs = BatchMetadataFileSystem::default(); + + let results = super::batched_get_metadata( + &fs, + &[ + ("first".to_string(), first_path, first_uri.clone()), + ("second".to_string(), second_path, second_uri.clone()), + ], + ) + .await + .expect("batch should be supported") + .into_iter() + .map(|result| result.expect("metadata result")) + .collect::>(); + + assert_eq!( + results, + vec![ + FileMetadata { + is_directory: true, + is_file: false, + is_symlink: false, + size: 0, + created_at_ms: 1, + modified_at_ms: 2, + }, + FileMetadata { + is_directory: true, + is_file: false, + is_symlink: false, + size: 0, + created_at_ms: 1, + modified_at_ms: 2, + }, + ] + ); + + let calls = fs.calls.lock().expect("calls lock"); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].len(), 2); + assert_eq!(calls[0][0].0, FS_GET_METADATA_METHOD); + assert_eq!(calls[0][1].0, FS_GET_METADATA_METHOD); + let first_params: FsGetMetadataParams = + serde_json::from_value(calls[0][0].1.clone()).expect("first params"); + let second_params: FsGetMetadataParams = + serde_json::from_value(calls[0][1].1.clone()).expect("second params"); + assert_eq!(first_params.path, first_uri); + assert_eq!(second_params.path, second_uri); +} + +#[derive(Default)] +struct BatchMetadataFileSystem { + calls: Mutex>>, +} + +impl ExecutorFileSystem for BatchMetadataFileSystem { + fn canonicalize<'a>( + &'a self, + _path: &'a PathUri, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, PathUri> { + Box::pin(async { unsupported() }) + } + + fn read_file<'a>( + &'a self, + _path: &'a PathUri, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(async { unsupported() }) + } + + fn read_file_stream<'a>( + &'a self, + _path: &'a PathUri, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> { + Box::pin(async { unsupported() }) + } + + fn write_file<'a>( + &'a self, + _path: &'a PathUri, + _contents: Vec, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + Box::pin(async { unsupported() }) + } + + fn create_directory<'a>( + &'a self, + _path: &'a PathUri, + _create_directory_options: CreateDirectoryOptions, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + Box::pin(async { unsupported() }) + } + + fn get_metadata<'a>( + &'a self, + _path: &'a PathUri, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileMetadata> { + Box::pin(async { unsupported() }) + } + + fn read_directory<'a>( + &'a self, + _path: &'a PathUri, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + Box::pin(async { unsupported() }) + } + + fn remove<'a>( + &'a self, + _path: &'a PathUri, + _remove_options: RemoveOptions, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + Box::pin(async { unsupported() }) + } + + fn copy<'a>( + &'a self, + _source_path: &'a PathUri, + _destination_path: &'a PathUri, + _copy_options: CopyOptions, + _sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + Box::pin(async { unsupported() }) + } + + fn execute_rpc_batch<'a>( + &'a self, + requests: Vec<(String, serde_json::Value)>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + Box::pin(async move { + let len = requests.len(); + self.calls.lock().expect("calls lock").push(requests); + Ok((0..len) + .map(|_| { + serde_json::to_value(FsGetMetadataResponse { + is_directory: true, + is_file: false, + is_symlink: false, + size: 0, + created_at_ms: 1, + modified_at_ms: 2, + }) + .map_err(io::Error::other) + }) + .collect()) + }) + } +} + +fn unsupported() -> io::Result { + Err(io::Error::new(io::ErrorKind::Unsupported, "unsupported")) +}