From dd3d8973e9ba0bc4220faea722bdd83ebbc990bf Mon Sep 17 00:00:00 2001 From: bryanashley Date: Tue, 23 Jun 2026 21:51:15 -0700 Subject: [PATCH] [skills] limit executor catalog cache to first turn --- codex-rs/ext/skills/src/extension.rs | 35 ++++++++++++++++--- codex-rs/ext/skills/src/state.rs | 30 ++++++++++++---- codex-rs/ext/skills/tests/skills_extension.rs | 8 ++--- 3 files changed, 56 insertions(+), 17 deletions(-) diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index d86cdf10c2..47c15c4a4e 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -50,6 +50,12 @@ struct SkillsExtension { config_from_host: Arc SkillsExtensionConfig + Send + Sync>, } +#[derive(Clone, Copy)] +enum ExecutorCatalogCacheAction { + PopulateForFirstTurn, + ConsumeForTurn, +} + impl ThreadLifecycleContributor for SkillsExtension where C: Send + Sync + 'static, @@ -128,6 +134,7 @@ where mcp_resources: session_store.get::(), }, &thread_state, + ExecutorCatalogCacheAction::PopulateForFirstTurn, ) .await; for warning in &catalog.warnings { @@ -194,7 +201,13 @@ where include_orchestrator_skills: thread_state.orchestrator_skills_enabled(), mcp_resources: session_store.get::(), }; - let catalog = self.list_skills(query, &thread_state).await; + let catalog = self + .list_skills( + query, + &thread_state, + ExecutorCatalogCacheAction::ConsumeForTurn, + ) + .await; for warning in &catalog.warnings { self.emit_warning(&input.turn_id, warning.clone()); } @@ -292,6 +305,7 @@ impl SkillsExtension { &self, mut query: SkillListQuery, thread_state: &SkillsThreadState, + executor_catalog_cache_action: ExecutorCatalogCacheAction, ) -> SkillCatalog { let executor_query = query.clone(); let include_orchestrator_skills = query.include_orchestrator_skills; @@ -299,9 +313,22 @@ impl SkillsExtension { let mcp_resources = orchestrator_query.mcp_resources.clone(); query.include_orchestrator_skills = false; - let executor_catalog = thread_state - .executor_catalog_snapshot(self.providers.list_executor_for_turn(executor_query)) - .await; + let executor_catalog = match executor_catalog_cache_action { + ExecutorCatalogCacheAction::PopulateForFirstTurn => { + thread_state + .executor_catalog_for_thread_context( + self.providers.list_executor_for_turn(executor_query), + ) + .await + } + ExecutorCatalogCacheAction::ConsumeForTurn => { + thread_state + .executor_catalog_for_turn( + self.providers.list_executor_for_turn(executor_query), + ) + .await + } + }; let mut catalog = self.providers.list_for_turn(query, executor_catalog).await; if include_orchestrator_skills { let orchestrator_catalog = thread_state diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 3cd1f69ef4..038b37ce32 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -6,6 +6,7 @@ use std::sync::Mutex; use codex_mcp::McpResourceClient; use codex_mcp::McpResourceClientCacheKey; use codex_protocol::capabilities::SelectedCapabilityRoot; +use tokio::sync::Mutex as AsyncMutex; use tokio::sync::OnceCell; use crate::SkillsExtensionConfig; @@ -27,7 +28,7 @@ const MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES: usize = 8 * 1024 * 1024; pub(crate) struct SkillsThreadState { config: Mutex, selected_roots: Vec, - executor_catalog: OnceCell, + first_turn_executor_catalog: AsyncMutex>, orchestrator_skills_available: bool, orchestrator_cache: Mutex>>, } @@ -41,7 +42,7 @@ impl SkillsThreadState { Self { config: Mutex::new(config), selected_roots, - executor_catalog: OnceCell::new(), + first_turn_executor_catalog: AsyncMutex::new(None), orchestrator_skills_available, orchestrator_cache: Mutex::new(None), } @@ -69,14 +70,29 @@ impl SkillsThreadState { self.orchestrator_skills_available && self.config().orchestrator_skills_enabled } - pub(crate) async fn executor_catalog_snapshot( + pub(crate) async fn executor_catalog_for_thread_context( &self, initialize: impl Future + Send, ) -> SkillCatalog { - self.executor_catalog - .get_or_init(|| initialize) - .await - .clone() + let mut cached_catalog = self.first_turn_executor_catalog.lock().await; + if let Some(catalog) = cached_catalog.as_ref() { + return catalog.clone(); + } + + let catalog = initialize.await; + *cached_catalog = Some(catalog.clone()); + catalog + } + + pub(crate) async fn executor_catalog_for_turn( + &self, + initialize: impl Future + Send, + ) -> SkillCatalog { + if let Some(catalog) = self.first_turn_executor_catalog.lock().await.take() { + return catalog; + } + + initialize.await } pub(crate) async fn orchestrator_catalog_snapshot( diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index d3b08bd5cb..95bb05d071 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -234,11 +234,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in )], read_request_keys(&read_requests) ); - let rebuilt_prompt_fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) - .await; - assert_eq!(1, rebuilt_prompt_fragments.len()); - assert!(rebuilt_prompt_fragments[0].text().contains("lint-fix")); + assert_eq!(1, list_calls.load(Ordering::Relaxed)); let next_turn_store = ExtensionData::new("turn-2"); let next_fragments = registry.turn_input_contributors()[0] @@ -258,7 +254,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in .await; assert!(next_fragments.is_empty()); - assert_eq!(1, list_calls.load(Ordering::Relaxed)); + assert_eq!(2, list_calls.load(Ordering::Relaxed)); Ok(()) }