perf(skills): batch frontmatter prefix reads

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::<Vec<_>>();
+    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::<Vec<_>>()
-        .boxed();
+        .collect::<Vec<_>>();
+    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::<Vec<_>>()
+    }
+    .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<String>,
+    loaded_metadata: LoadedSkillMetadata,
     scope: SkillScope,
     plugin_id: Option<&str>,
-    plugin_root: Option<&AbsolutePathBuf>,
 ) -> Result<SkillMetadata, SkillParseError> {
-    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,

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<String>),
+    Full(PathUri),
+}
+
+pub(super) async fn read_skill_frontmatter_texts(
+    fs: &dyn ExecutorFileSystem,
+    paths: &[PathUri],
+) -> Vec<io::Result<String>> {
+    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;

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<dyn ExecutorFileSystem>,
+    batch_behavior: BatchBehavior,
+    full_reads: Mutex<Vec<PathUri>>,
+}
+
+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<PathUri> {
+        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<u8>> {
+        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<io::Result<TextFilePrefix>>> {
+        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<u8>,
+        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<ReadDirectoryEntry>> {
+        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<io::Result<String>>) -> Vec<String> {
+    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::<PathUri>::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:?}");
+    }
+}