diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector.rs b/codex-rs/ext/skills/src/dynamic_skill_selector.rs index 30a20f2ec0..fa71e01f85 100644 --- a/codex-rs/ext/skills/src/dynamic_skill_selector.rs +++ b/codex-rs/ext/skills/src/dynamic_skill_selector.rs @@ -1,6 +1,7 @@ mod character_ngram; mod character_routing_card; mod fielded_bm25; +mod lru; mod multi_query_lexical; mod routing_card_lexical; mod rrf_lexical_char; @@ -9,6 +10,7 @@ pub(crate) use character_ngram::CharacterNgramSkillSelector; pub(crate) use character_routing_card::CharacterRoutingCardSkillSelector; use codex_skills::SkillDependencies; pub(crate) use fielded_bm25::FieldedBm25SkillSelector; +pub(crate) use lru::LruSkillSelector; pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector; pub(crate) use routing_card_lexical::RoutingCardLexicalSkillSelector; pub(crate) use rrf_lexical_char::RrfLexicalCharSkillSelector; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/lru.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/lru.rs new file mode 100644 index 0000000000..8000f4e73a --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/lru.rs @@ -0,0 +1,61 @@ +use std::collections::HashSet; + +use super::CheapSkillSelection; +use super::CheapSkillSelector; +use super::SkillSelectionDocument; + +const MAX_CANDIDATES: usize = 1_000; +const MAX_QUERY_TERMS: usize = 64; +const MAX_RESULTS: usize = 50; + +#[derive(Clone, Debug, Default)] +pub(crate) struct LruSkillSelector { + recent_skill_ids: Vec, +} + +impl LruSkillSelector { + pub(crate) fn new(recent_skill_ids: Vec) -> Self { + Self { recent_skill_ids } + } +} + +impl CheapSkillSelector for LruSkillSelector { + fn method(&self) -> &'static str { + "lru_v1" + } + + fn select( + &self, + query: &str, + documents: &[SkillSelectionDocument<'_>], + limit: usize, + ) -> CheapSkillSelection { + let mut query_terms = query.split_whitespace(); + let query_term_count = query_terms.by_ref().take(MAX_QUERY_TERMS).count(); + let query_truncated = query_terms.next().is_some(); + let candidate_set_truncated = documents.len() > MAX_CANDIDATES; + let eligible_ids = documents + .iter() + .take(MAX_CANDIDATES) + .map(|document| document.id) + .collect::>(); + let mut seen_ids = HashSet::new(); + + CheapSkillSelection { + candidate_ids: self + .recent_skill_ids + .iter() + .copied() + .filter(|id| eligible_ids.contains(id) && seen_ids.insert(*id)) + .take(limit.min(MAX_RESULTS)) + .collect(), + query_term_count, + query_truncated, + candidate_set_truncated, + } + } +} + +#[cfg(test)] +#[path = "lru_tests.rs"] +mod tests; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/lru_tests.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/lru_tests.rs new file mode 100644 index 0000000000..3dc0c049ac --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/lru_tests.rs @@ -0,0 +1,60 @@ +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn lru_selector_preserves_recent_invocation_order() { + let documents = [ + SkillSelectionDocument { + id: 10, + name: "ci", + short_description: None, + description: "Investigate failing checks.", + dependencies: None, + }, + SkillSelectionDocument { + id: 20, + name: "python-tools", + short_description: None, + description: "Manage Python environments.", + dependencies: None, + }, + SkillSelectionDocument { + id: 30, + name: "monorepo", + short_description: None, + description: "Follow repository conventions.", + dependencies: None, + }, + ]; + + let selection = + LruSkillSelector::new(vec![30, 10]).select("continue", &documents, /*limit*/ 50); + + assert_eq!(vec![30, 10], selection.candidate_ids); +} + +#[test] +fn lru_selector_filters_stale_or_duplicate_skills_and_respects_the_limit() { + let documents = [ + SkillSelectionDocument { + id: 10, + name: "ci", + short_description: None, + description: "Investigate failing checks.", + dependencies: None, + }, + SkillSelectionDocument { + id: 20, + name: "python-tools", + short_description: None, + description: "Manage Python environments.", + dependencies: None, + }, + ]; + + let selection = + LruSkillSelector::new(vec![99, 10, 10, 20]).select("yes", &documents, /*limit*/ 1); + + assert_eq!(vec![10], selection.candidate_ids); +} diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index f1458cc1d5..2227688a97 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -399,6 +399,7 @@ where &shadow_catalog, &shadow_selected_entries, host_snapshot.as_deref(), + Arc::clone(&thread_state.recent_skill_invocations), )) } else { None diff --git a/codex-rs/ext/skills/src/shadow_selection_experiment.rs b/codex-rs/ext/skills/src/shadow_selection_experiment.rs index b43d377f16..556ab58543 100644 --- a/codex-rs/ext/skills/src/shadow_selection_experiment.rs +++ b/codex-rs/ext/skills/src/shadow_selection_experiment.rs @@ -1,6 +1,9 @@ // This shadow-selection experiment is temporary and should be removed after evaluation. +use std::collections::HashMap; use std::collections::HashSet; +use std::collections::VecDeque; +use std::sync::Arc; use std::sync::Mutex; use std::sync::PoisonError; use std::time::Duration; @@ -18,6 +21,7 @@ use crate::dynamic_skill_selector::CharacterRoutingCardSkillSelector; use crate::dynamic_skill_selector::CheapSkillSelection; use crate::dynamic_skill_selector::CheapSkillSelector; use crate::dynamic_skill_selector::FieldedBm25SkillSelector; +use crate::dynamic_skill_selector::LruSkillSelector; use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector; use crate::dynamic_skill_selector::RoutingCardLexicalSkillSelector; use crate::dynamic_skill_selector::RrfLexicalCharSkillSelector; @@ -25,7 +29,7 @@ use crate::dynamic_skill_selector::SkillSelectionDocument; use crate::dynamic_skill_selector::WeightedLexicalSkillSelector; const MAX_SHADOW_QUERY_BYTES: usize = 16 * 1024; -const MAX_SHADOW_RESULTS: usize = 20; +const MAX_SHADOW_RESULTS: usize = 50; const RUN_METRIC: &str = "codex.skills.shadow_selection"; const DURATION_METRIC: &str = "codex.skills.shadow_selection.duration_ms"; @@ -61,6 +65,7 @@ impl ShadowSelectionExperiment { catalog: &SkillCatalog, explicitly_selected: &[SkillCatalogEntry], host_snapshot: Option<&HostSkillsSnapshot>, + recent_skill_invocations: Arc, ) -> ShadowSelectionTurnState { let query = build_shadow_query(inputs); let query_script = query_script_tag(&query.text); @@ -95,22 +100,36 @@ impl ShadowSelectionExperiment { .iter() .map(|document| document.id) .collect::>(); - let eligible_skill_resources = documents + let eligible_skill_ids_by_resource = documents .iter() .map(|document| { - normalize_skill_resource(catalog.entries[document.id].main_prompt.as_str()) + ( + normalize_skill_resource(catalog.entries[document.id].main_prompt.as_str()), + document.id, + ) }) + .collect::>(); + let eligible_skill_resources = eligible_skill_ids_by_resource + .keys() + .cloned() .collect::>(); + let recent_skill_ids = recent_skill_invocations + .snapshot() + .iter() + .filter_map(|resource| eligible_skill_ids_by_resource.get(resource).copied()) + .collect(); let routing_selector = CharacterRoutingCardSkillSelector::new(catalog, host_snapshot); - let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 1); + let lru_selector = LruSkillSelector::new(recent_skill_ids); + let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 2); for selector in self .selectors .iter() .map(std::convert::AsRef::as_ref) - .chain(std::iter::once( + .chain([ &routing_selector as &dyn CheapSkillSelector, - )) + &lru_selector as &dyn CheapSkillSelector, + ]) { let start = Instant::now(); let selection = @@ -150,6 +169,7 @@ impl ShadowSelectionExperiment { query_script, eligible_skill_resources, seen_skill_resources: Mutex::new(HashSet::new()), + recent_skill_invocations, } } @@ -166,6 +186,9 @@ impl ShadowSelectionExperiment { { return; } + state + .recent_skill_invocations + .record(skill_resource.clone()); let Some(metrics_client) = self.metrics_client.as_ref() else { return; }; @@ -239,6 +262,38 @@ pub(crate) struct ShadowSelectionTurnState { query_script: &'static str, eligible_skill_resources: HashSet, seen_skill_resources: Mutex>, + recent_skill_invocations: Arc, +} + +#[derive(Default)] +pub(crate) struct RecentSkillInvocations { + skill_resources: Mutex>, +} + +impl RecentSkillInvocations { + fn snapshot(&self) -> Vec { + self.skill_resources + .lock() + .unwrap_or_else(PoisonError::into_inner) + .iter() + .cloned() + .collect() + } + + fn record(&self, skill_resource: String) { + let mut skill_resources = self + .skill_resources + .lock() + .unwrap_or_else(PoisonError::into_inner); + if let Some(index) = skill_resources + .iter() + .position(|resource| resource == &skill_resource) + { + skill_resources.remove(index); + } + skill_resources.push_front(skill_resource); + skill_resources.truncate(MAX_SHADOW_RESULTS); + } } struct RankedSelection { @@ -310,7 +365,8 @@ fn rank_bucket(rank: Option) -> &'static str { Some(1) => "1", Some(2..=5) => "2_5", Some(6..=10) => "6_10", - Some(11..=MAX_SHADOW_RESULTS) => "11_20", + Some(11..=20) => "11_20", + Some(21..=MAX_SHADOW_RESULTS) => "21_50", Some(_) | None => "miss", } } @@ -406,3 +462,7 @@ fn push_bounded(destination: &mut String, value: &str) -> bool { destination.push_str(&value[..end]); false } + +#[cfg(test)] +#[path = "shadow_selection_experiment_tests.rs"] +mod tests; diff --git a/codex-rs/ext/skills/src/shadow_selection_experiment_tests.rs b/codex-rs/ext/skills/src/shadow_selection_experiment_tests.rs new file mode 100644 index 0000000000..4fe6f5c1dc --- /dev/null +++ b/codex-rs/ext/skills/src/shadow_selection_experiment_tests.rs @@ -0,0 +1,27 @@ +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn recent_invocations_refresh_recency_and_evict_old_skills() { + let history = RecentSkillInvocations::default(); + for index in 0..=MAX_SHADOW_RESULTS { + history.record(format!("skill-{index}")); + } + history.record("skill-1".to_string()); + + let recent = history.snapshot(); + + assert_eq!(MAX_SHADOW_RESULTS, recent.len()); + assert_eq!(Some("skill-1"), recent.first().map(String::as_str)); + assert_eq!(Some("skill-2"), recent.last().map(String::as_str)); + assert!(!recent.iter().any(|skill| skill == "skill-0")); +} + +#[test] +fn rank_buckets_distinguish_results_above_twenty() { + assert_eq!("11_20", rank_bucket(Some(20))); + assert_eq!("21_50", rank_bucket(Some(21))); + assert_eq!("21_50", rank_bucket(Some(50))); + assert_eq!("miss", rank_bucket(Some(51))); +} diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 88f851a524..c3f0589a26 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -21,6 +21,7 @@ use crate::catalog::SkillResourceId; use crate::catalog::SkillSourceKind; use crate::provider::SkillListQuery; use crate::provider::SkillReadRequest; +use crate::shadow_selection_experiment::RecentSkillInvocations; use crate::shadow_selection_experiment::ShadowSelectionTurnState; use crate::sources::SkillProviders; @@ -39,6 +40,7 @@ pub(crate) struct SkillsThreadState { executor_discovery_cache: Mutex>, orchestrator_cache: Mutex>>, shadow_selection_turn: Mutex>, + pub(crate) recent_skill_invocations: Arc, } impl SkillsThreadState { @@ -50,6 +52,7 @@ impl SkillsThreadState { executor_discovery_cache: Mutex::new(None), orchestrator_cache: Mutex::new(None), shadow_selection_turn: Mutex::new(None), + recent_skill_invocations: Arc::new(RecentSkillInvocations::default()), } } diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index e0bb93ab58..eb4bcf38cb 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -18,6 +18,8 @@ use codex_extension_api::ExtensionWarning; use codex_extension_api::NoopTurnItemEmitter; use codex_extension_api::PreviousWorldStateSection; use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::SkillInvocationInput; +use codex_extension_api::SkillInvocationKind; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolPayload; @@ -842,6 +844,126 @@ async fn shadow_selection_uses_host_catalog_when_instructions_are_disabled() -> Ok(()) } +#[tokio::test] +async fn shadow_lru_selector_recovers_a_skill_invoked_on_an_earlier_turn() -> TestResult { + let provider = Arc::new(StaticSkillProvider { + catalog: SkillCatalog { + entries: vec![test_entry( + SkillSourceKind::Host, + "host", + "host/lint-fix", + "lint-fix/SKILL.md", + )], + warnings: Vec::new(), + }, + read_requests: Arc::new(Mutex::new(Vec::new())), + list_calls: None, + fail_first_list: false, + }); + let metrics = MetricsClient::new( + MetricsConfig::in_memory( + "test", + "codex-skills-extension", + env!("CARGO_PKG_VERSION"), + InMemoryMetricExporter::default(), + ) + .with_runtime_reader(), + )?; + let mut builder = ExtensionRegistryBuilder::new(); + install_with_providers_and_metrics( + &mut builder, + SkillProviders::new().with_host_provider(provider), + Some(metrics.clone()), + skills_extension_config, + ); + let registry = builder.build(); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + let mut config = default_config(); + config.include_instructions = false; + config.shadow_selection_enabled = true; + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &SessionSource::Cli, + persistent_thread_state_available: true, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + + for (turn_id, text) in [("turn-1", "Fix lint errors."), ("turn-2", "continue")] { + let turn_store = ExtensionData::new(turn_id); + let fragments = registry.turn_input_contributors()[0] + .contribute( + TurnInputContext { + turn_id: turn_id.to_string(), + user_input: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + environments: Vec::new(), + }, + /*extension_metrics*/ None, + &session_store, + &thread_store, + &turn_store, + ) + .await; + assert!(fragments.is_empty()); + registry.skill_invocation_contributors()[0] + .on_skill_invocation(SkillInvocationInput { + session_store: &session_store, + thread_store: &thread_store, + turn_store: &turn_store, + turn_id, + skill_resource: "lint-fix/SKILL.md", + kind: SkillInvocationKind::Implicit, + }) + .await; + } + + let snapshot = metrics.snapshot()?; + let metric = snapshot + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == "codex.skills.shadow_selection.invocation") + .ok_or("shadow invocation metric should be recorded")?; + let mut lru_hits = match metric.data() { + AggregatedMetrics::U64(MetricData::Sum(sum)) => sum + .data_points() + .filter_map(|point| { + let method = point + .attributes() + .find(|attribute| attribute.key.as_str() == "method")? + .value + .as_str(); + if method != "lru_v1" { + return None; + } + let hit = point + .attributes() + .find(|attribute| attribute.key.as_str() == "hit")? + .value + .as_str() + .to_string(); + Some((hit, point.value())) + }) + .collect::>(), + data => panic!("unexpected shadow invocation metric data: {data:?}"), + }; + lru_hits.sort(); + + assert_eq!( + vec![("false".to_string(), 1), ("true".to_string(), 1)], + lru_hits + ); + Ok(()) +} + #[tokio::test] async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cache() -> TestResult { let read_requests = Arc::new(Mutex::new(Vec::new()));