openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add an LRU baseline to skill shadow selection (#38197)

Browse Source 
## What changed

- Track the 50 most recently invoked skills for each thread and evaluate them
  as the `lru_v1` shadow-selection method.
- Filter stale and duplicate entries while preserving recency order.
- Increase the shadow result limit to 50 and add a `21_50` rank bucket.

## Testing

- Add unit coverage for recency refresh, eviction, filtering, limits, and rank
  buckets.
- Add an extension test showing that `lru_v1` recovers a skill invoked on an
  earlier turn.

GitOrigin-RevId: 7aa56514db57cf180cc3f57b4d493dee3304e766
This commit is contained in:
jif
2026-08-12 14:04:47 +00:00
committed by copyberry
parent eb752e43d9
commit 9dd22890f5
8 changed files with 343 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/ext/skills/src/dynamic_skill_selector.rs
View File

@@ -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;

61
codex-rs/ext/skills/src/dynamic_skill_selector/lru.rs Normal file
View File

@@ -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<usize>,
}
impl LruSkillSelector {
pub(crate) fn new(recent_skill_ids: Vec<usize>) -> 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::<HashSet<_>>();
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;

60
codex-rs/ext/skills/src/dynamic_skill_selector/lru_tests.rs Normal file
View File

@@ -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);
}

1
codex-rs/ext/skills/src/extension.rs
View File

@@ -399,6 +399,7 @@ where
&shadow_catalog,
&shadow_selected_entries,
host_snapshot.as_deref(),
Arc::clone(&thread_state.recent_skill_invocations),
))
} else {
None

74
codex-rs/ext/skills/src/shadow_selection_experiment.rs
View File

@@ -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<RecentSkillInvocations>,
) -> 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::<HashSet<_>>();
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::<HashMap<_, _>>();
let eligible_skill_resources = eligible_skill_ids_by_resource
.keys()
.cloned()
.collect::<HashSet<_>>();
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<String>,
seen_skill_resources: Mutex<HashSet<String>>,
recent_skill_invocations: Arc<RecentSkillInvocations>,
}
#[derive(Default)]
pub(crate) struct RecentSkillInvocations {
skill_resources: Mutex<VecDeque<String>>,
}
impl RecentSkillInvocations {
fn snapshot(&self) -> Vec<String> {
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<usize>) -> &'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;

27
codex-rs/ext/skills/src/shadow_selection_experiment_tests.rs Normal file
View File

@@ -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)));
}

3
codex-rs/ext/skills/src/state.rs
View File

@@ -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<Option<CachedExecutorDiscoveryCatalog>>,
orchestrator_cache: Mutex<Option<Arc<OrchestratorGenerationCache>>>,
shadow_selection_turn: Mutex<Option<ShadowSelectionTurn>>,
pub(crate) recent_skill_invocations: Arc<RecentSkillInvocations>,
}
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()),
}
}

122
codex-rs/ext/skills/tests/skills_extension.rs
View File

@@ -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::<Vec<_>>(),
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()));