diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 26adfa7542..c0cc2eb5eb 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -1,6 +1,7 @@ mod discovery; mod environment; mod namespace; +mod text_prefix; pub use environment::EnvironmentSkillLoadOutcome; pub use environment::EnvironmentSkillMetadata; @@ -21,7 +22,6 @@ 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_BATCH_MAX_PATHS; use codex_exec_server::LOCAL_FS; use codex_protocol::protocol::Product; use codex_protocol::protocol::SkillScope; @@ -596,27 +596,50 @@ async fn load_skills_under_root( } } }; - let skill_results = futures::stream::iter(resolved_skills) + let contents = text_prefix::read_skill_frontmatter_texts(fs, &skill_paths); + let metadata_requests = resolved_skills + .iter() .map(|skill| { - let plugin_root = plugin_root.as_ref(); - async move { - let result = parse_skill_file( - fs, - &skill.skill, - &skill.path, - &skill.path_uri, - scope, - plugin_id, - plugin_root, - ) - .await - .map_err(|err| err.to_string()); - (skill.path, skill.path_uri, result) + let metadata_path = skill + .path_uri + .parent() + .and_then(|parent| parent.join(SKILLS_METADATA_DIR).ok()) + .and_then(|directory| directory.join(SKILLS_METADATA_FILENAME).ok()); + let metadata = match &skill.skill.metadata { + SkillMetadataDiscovery::Present(_) => { + metadata_path.map(SkillMetadataDiscovery::Present) + } + SkillMetadataDiscovery::Probe(_) => { + metadata_path.map(SkillMetadataDiscovery::Probe) + } + SkillMetadataDiscovery::Absent => None, } + .unwrap_or(SkillMetadataDiscovery::Absent); + (skill.path.clone(), metadata) + }) + .collect::>(); + let plugin_root = plugin_root.as_ref(); + let metadata = futures::stream::iter(metadata_requests) + .map(|(path, metadata)| async move { + load_skill_metadata(fs, &path, &metadata, plugin_root).await }) .buffered(MAX_CONCURRENT_SKILL_LOADS) - .collect::>() - .boxed(); + .collect::>(); + let skill_results = async move { + let (contents, metadata) = tokio::join!(contents, metadata); + resolved_skills + .into_iter() + .zip(contents) + .zip(metadata) + .map(|((skill, contents), loaded_metadata)| { + let result = + parse_skill_contents(&skill.path, contents, loaded_metadata, scope, plugin_id) + .map_err(|err| err.to_string()); + (skill.path, skill.path_uri, result) + }) + .collect::>() + } + .boxed(); let (namespace_resolver, skill_results) = tokio::join!(namespace_resolver, skill_results); for (path, path_uri, result) in skill_results { let result = result.and_then(|mut skill| { @@ -637,29 +660,13 @@ async fn load_skills_under_root( } } -async fn parse_skill_file( - fs: &dyn ExecutorFileSystem, - skill: &DiscoveredSkill, +fn parse_skill_contents( path: &AbsolutePathBuf, - path_uri: &PathUri, + contents: io::Result, + loaded_metadata: LoadedSkillMetadata, scope: SkillScope, plugin_id: Option<&str>, - plugin_root: Option<&AbsolutePathBuf>, ) -> Result { - let metadata_path = path_uri - .parent() - .and_then(|parent| parent.join(SKILLS_METADATA_DIR).ok()) - .and_then(|directory| directory.join(SKILLS_METADATA_FILENAME).ok()); - let metadata = match &skill.metadata { - SkillMetadataDiscovery::Present(_) => metadata_path.map(SkillMetadataDiscovery::Present), - SkillMetadataDiscovery::Probe(_) => metadata_path.map(SkillMetadataDiscovery::Probe), - SkillMetadataDiscovery::Absent => None, - } - .unwrap_or(SkillMetadataDiscovery::Absent); - let (contents, loaded_metadata) = tokio::join!( - fs.read_file_text(path_uri, /*sandbox*/ None), - load_skill_metadata(fs, path, &metadata, plugin_root), - ); let contents = contents.map_err(SkillParseError::Read)?; let ParsedSkillFrontmatter { name: base_name, diff --git a/codex-rs/core-skills/src/loader/text_prefix.rs b/codex-rs/core-skills/src/loader/text_prefix.rs new file mode 100644 index 0000000000..265a5fec35 --- /dev/null +++ b/codex-rs/core-skills/src/loader/text_prefix.rs @@ -0,0 +1,63 @@ +use std::io; + +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS; +use codex_utils_path_uri::PathUri; +use futures::StreamExt; + +use super::discovery::MAX_CONCURRENT_SKILL_LOADS; +use super::extract_frontmatter; + +const SKILL_FRONTMATTER_PREFIX_BYTES: usize = 2 * 1024; + +enum PrefixRead { + Ready(io::Result), + Full(PathUri), +} + +pub(super) async fn read_skill_frontmatter_texts( + fs: &dyn ExecutorFileSystem, + paths: &[PathUri], +) -> Vec> { + let mut pending = Vec::with_capacity(paths.len()); + + for chunk in paths.chunks(FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS) { + let results = fs + .read_text_prefixes_batch(chunk, SKILL_FRONTMATTER_PREFIX_BYTES, /*sandbox*/ None) + .await; + let results = match results { + Ok(results) if results.len() == chunk.len() => results, + Ok(_) | Err(_) => { + pending.extend(chunk.iter().cloned().map(PrefixRead::Full)); + continue; + } + }; + for (path, result) in chunk.iter().cloned().zip(results) { + match result { + Ok(prefix) if extract_frontmatter(&prefix.text).is_some() => { + pending.push(PrefixRead::Ready(Ok(prefix.text))); + } + Ok(prefix) if prefix.complete => { + pending.push(PrefixRead::Ready(Ok(prefix.text))); + } + Ok(_) => pending.push(PrefixRead::Full(path)), + Err(error) => pending.push(PrefixRead::Ready(Err(error))), + } + } + } + + futures::stream::iter(pending) + .map(|read| async move { + match read { + PrefixRead::Ready(result) => result, + PrefixRead::Full(path) => fs.read_file_text(&path, /*sandbox*/ None).await, + } + }) + .buffered(MAX_CONCURRENT_SKILL_LOADS) + .collect() + .await +} + +#[cfg(test)] +#[path = "text_prefix_tests.rs"] +mod tests; diff --git a/codex-rs/core-skills/src/loader/text_prefix_tests.rs b/codex-rs/core-skills/src/loader/text_prefix_tests.rs new file mode 100644 index 0000000000..22f4248e94 --- /dev/null +++ b/codex-rs/core-skills/src/loader/text_prefix_tests.rs @@ -0,0 +1,241 @@ +use std::fs; +use std::io; +use std::sync::Arc; +use std::sync::Mutex; + +use codex_exec_server::CopyOptions; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::ExecutorFileSystemFuture; +use codex_exec_server::FileMetadata; +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_exec_server::TextFilePrefix; +use codex_exec_server::WalkOptions; +use codex_exec_server::WalkOutcome; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; + +use super::SKILL_FRONTMATTER_PREFIX_BYTES; +use super::read_skill_frontmatter_texts; + +#[derive(Clone, Copy, Debug)] +enum BatchBehavior { + Delegate, + Error, + WrongCardinality, +} + +struct TestFileSystem { + inner: Arc, + batch_behavior: BatchBehavior, + full_reads: Mutex>, +} + +impl TestFileSystem { + fn new(batch_behavior: BatchBehavior) -> Self { + Self { + inner: Arc::clone(&LOCAL_FS), + batch_behavior, + full_reads: Mutex::new(Vec::new()), + } + } + + fn full_reads(&self) -> Vec { + self.full_reads + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } +} + +impl ExecutorFileSystem for TestFileSystem { + fn canonicalize<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, PathUri> { + self.inner.canonicalize(path, sandbox) + } + + fn read_file<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + self.inner.read_file(path, sandbox) + } + + fn read_file_stream<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> { + self.inner.read_file_stream(path, sandbox) + } + + fn read_file_text<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, String> { + self.full_reads + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(path.clone()); + self.inner.read_file_text(path, sandbox) + } + + fn read_text_prefixes_batch<'a>( + &'a self, + paths: &'a [PathUri], + max_bytes: usize, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + match self.batch_behavior { + BatchBehavior::Delegate => self + .inner + .read_text_prefixes_batch(paths, max_bytes, sandbox), + BatchBehavior::Error => { + Box::pin(async { Err(io::Error::other("synthetic batch failure")) }) + } + BatchBehavior::WrongCardinality => Box::pin(async { Ok(Vec::new()) }), + } + } + + fn write_file<'a>( + &'a self, + path: &'a PathUri, + contents: Vec, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.write_file(path, contents, sandbox) + } + + fn create_directory<'a>( + &'a self, + path: &'a PathUri, + options: CreateDirectoryOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.create_directory(path, options, sandbox) + } + + fn get_metadata<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileMetadata> { + self.inner.get_metadata(path, sandbox) + } + + fn read_directory<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + self.inner.read_directory(path, sandbox) + } + + fn walk<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { + self.inner.walk(path, options, sandbox) + } + + fn remove<'a>( + &'a self, + path: &'a PathUri, + options: RemoveOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.remove(path, options, sandbox) + } + + fn copy<'a>( + &'a self, + source_path: &'a PathUri, + destination_path: &'a PathUri, + options: CopyOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner + .copy(source_path, destination_path, options, sandbox) + } +} + +fn write_file(root: &std::path::Path, name: &str, contents: &str) -> PathUri { + let path = root.join(name); + fs::write(&path, contents).expect("write test file"); + PathUri::from_abs_path(&AbsolutePathBuf::try_from(path).expect("absolute test path")) +} + +fn unwrap_texts(results: Vec>) -> Vec { + results + .into_iter() + .map(|result| result.expect("read text")) + .collect() +} + +#[tokio::test] +async fn prefix_results_preserve_input_order() { + let root = tempfile::tempdir().expect("tempdir"); + let first = format!( + "---\nname: first\ndescription: first\n---\n{}", + "a".repeat(4_000) + ); + let second = "---\nname: second\ndescription: second\n---\nbody"; + let paths = vec![ + write_file(root.path(), "first.md", &first), + write_file(root.path(), "second.md", second), + ]; + let file_system = TestFileSystem::new(BatchBehavior::Delegate); + + let texts = unwrap_texts(read_skill_frontmatter_texts(&file_system, &paths).await); + + assert_eq!(texts[0].len(), SKILL_FRONTMATTER_PREFIX_BYTES); + assert!(texts[0].starts_with("---\nname: first\n")); + assert_eq!(texts[1], second); + assert_eq!(file_system.full_reads(), Vec::::new()); +} + +#[tokio::test] +async fn incomplete_frontmatter_falls_back_to_full_read() { + let root = tempfile::tempdir().expect("tempdir"); + let contents = format!( + "---\nname: long\ndescription: {}\n---\nbody", + "x".repeat(4_000) + ); + let path = write_file(root.path(), "long.md", &contents); + let file_system = TestFileSystem::new(BatchBehavior::Delegate); + + let texts = + unwrap_texts(read_skill_frontmatter_texts(&file_system, std::slice::from_ref(&path)).await); + + assert_eq!(texts, vec![contents]); + assert_eq!(file_system.full_reads(), vec![path]); +} + +#[tokio::test] +async fn invalid_batch_results_fall_back_for_the_whole_chunk() { + for behavior in [BatchBehavior::Error, BatchBehavior::WrongCardinality] { + let root = tempfile::tempdir().expect("tempdir"); + let paths = vec![ + write_file(root.path(), "first.md", "first"), + write_file(root.path(), "second.md", "second"), + ]; + let file_system = TestFileSystem::new(behavior); + + let texts = unwrap_texts(read_skill_frontmatter_texts(&file_system, &paths).await); + + assert_eq!(texts, vec!["first", "second"], "behavior: {behavior:?}"); + assert_eq!(file_system.full_reads(), paths, "behavior: {behavior:?}"); + } +}