mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Move shadow skill ranking off the turn preparation path (#49812)
## Why Shadow skill selection runs synchronously during turn input construction, making turn preparation wait for experimental rankings. ## What changed Run ranking on background blocking workers with a process-wide limit of two concurrent rankings. Wait for capacity in the background and hold each slot until its blocking worker finishes, even if the async evaluation is cancelled. Freeze catalog, query, and relevance history at turn start. Record eligible invocations immediately for future turns, then score queued observations against the frozen rankings once they are ready, including observations queued before the turn state is dropped. ## Testing Add tests for turn preparation and history updates while ranking is blocked, deferred invocation scoring against each turn's history, and shared ranking capacity that remains occupied until cancelled workers finish. Update existing metric assertions to collect asynchronous observations across snapshots. GitOrigin-RevId: f7776e9aaeb45d3a792de167abbe35f7ec5a513e
This commit is contained in:
@@ -33,7 +33,7 @@ schemars = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml = { workspace = true }
|
||||
tokio = { workspace = true, features = ["sync", "time"] }
|
||||
tokio = { workspace = true, features = ["rt", "sync", "time"] }
|
||||
toml = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
|
||||
@@ -362,8 +362,7 @@ where
|
||||
.get::<SkillsThreadState>()
|
||||
.and_then(|state| state.shadow_selection_turn(input.turn_id))
|
||||
{
|
||||
self.shadow_selection
|
||||
.record_invocation(&state, input.skill_resource);
|
||||
state.record_invocation(input.skill_resource);
|
||||
}
|
||||
}
|
||||
SkillInvocationKind::Explicit => {}
|
||||
@@ -425,11 +424,11 @@ where
|
||||
}
|
||||
let shadow_selected_entries =
|
||||
collect_explicit_skill_mentions(&input.user_input, &shadow_catalog);
|
||||
Some(self.shadow_selection.run(
|
||||
Some(self.shadow_selection.start(
|
||||
&input,
|
||||
&shadow_catalog,
|
||||
shadow_catalog,
|
||||
&shadow_selected_entries,
|
||||
host_snapshot.as_deref(),
|
||||
host_snapshot.clone(),
|
||||
Arc::clone(&thread_state.recent_skill_invocations),
|
||||
Arc::clone(&thread_state.shadow_task_context),
|
||||
))
|
||||
@@ -591,7 +590,6 @@ impl<C> SkillsExtension<C> {
|
||||
executor_query,
|
||||
selected_plugins,
|
||||
sandbox_contexts,
|
||||
Arc::clone(&self.shadow_selection),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -25,3 +25,326 @@ fn rank_buckets_distinguish_results_above_twenty() {
|
||||
assert_eq!("21_50", rank_bucket(Some(50)));
|
||||
assert_eq!("miss", rank_bucket(Some(51)));
|
||||
}
|
||||
|
||||
// These tests deliberately occupy the process-wide budget and must drain their
|
||||
// workers before the next test starts.
|
||||
static RANKING_TESTS: Semaphore = Semaphore::const_new(1);
|
||||
|
||||
#[tokio::test]
|
||||
async fn turn_input_and_history_advance_while_ranking_is_blocked()
|
||||
-> Result<(), Box<dyn std::error::Error>> {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use codex_extension_api::ExtensionData;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::SkillInvocationInput;
|
||||
use codex_extension_api::SkillInvocationKind;
|
||||
use codex_extension_api::ThreadStartInput;
|
||||
use codex_otel::MetricsConfig;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
|
||||
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
|
||||
use opentelemetry_sdk::metrics::data::MetricData;
|
||||
use opentelemetry_sdk::metrics::data::ScopeMetrics;
|
||||
|
||||
use crate::SkillProviders;
|
||||
use crate::SkillsExtensionConfig;
|
||||
use crate::SkillsThreadState;
|
||||
use crate::catalog::SkillAuthority;
|
||||
use crate::catalog::SkillPackageId;
|
||||
use crate::catalog::SkillResourceId;
|
||||
use crate::install_with_providers_and_metrics;
|
||||
use crate::state::HostSkillsStepState;
|
||||
|
||||
let _serial = RANKING_TESTS.acquire().await.unwrap();
|
||||
let slots = RANKING_SLOTS.acquire_many(2).await?;
|
||||
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(),
|
||||
Some(metrics.clone()),
|
||||
SkillsExtensionConfig::clone,
|
||||
);
|
||||
let registry = builder.build();
|
||||
let config = SkillsExtensionConfig {
|
||||
include_instructions: false,
|
||||
max_context_tokens: None,
|
||||
bundled_skills_enabled: false,
|
||||
cloud_skill_enabled: false,
|
||||
shadow_selection_enabled: true,
|
||||
};
|
||||
let session_store = ExtensionData::new("session");
|
||||
let thread_store = ExtensionData::new("thread");
|
||||
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 in ["first", "second"] {
|
||||
let turn_store = ExtensionData::new(turn_id);
|
||||
turn_store.insert(HostSkillsStepState(SkillCatalog {
|
||||
entries: ["x", "hidden", "selected"]
|
||||
.into_iter()
|
||||
.map(|name| {
|
||||
let mut entry = SkillCatalogEntry::new(
|
||||
SkillPackageId(name.to_string()),
|
||||
SkillAuthority::new(SkillSourceKind::Host, "host"),
|
||||
name,
|
||||
"zzzz",
|
||||
SkillResourceId::new(format!("{name}/SKILL.md")),
|
||||
);
|
||||
entry.prompt_visible = name != "hidden";
|
||||
entry
|
||||
})
|
||||
.collect(),
|
||||
warnings: Vec::new(),
|
||||
}));
|
||||
let fragments = tokio::time::timeout(
|
||||
Duration::from_secs(10),
|
||||
registry.turn_input_contributors()[0].contribute(
|
||||
TurnInputContext {
|
||||
turn_id: turn_id.to_string(),
|
||||
user_input: vec![UserInput::Skill {
|
||||
name: "selected".to_string(),
|
||||
path: "selected/SKILL.md".into(),
|
||||
}],
|
||||
environments: Vec::new(),
|
||||
},
|
||||
/*extension_metrics*/ None,
|
||||
&session_store,
|
||||
&thread_store,
|
||||
&turn_store,
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
assert!(fragments.is_empty());
|
||||
tokio::task::yield_now().await;
|
||||
// Synchronous ranking would emit run metrics despite the occupied budget;
|
||||
// awaiting admission in contribute would instead time out above.
|
||||
assert!(
|
||||
metrics
|
||||
.snapshot()?
|
||||
.scope_metrics()
|
||||
.all(|scope| { scope.metrics().all(|metric| metric.name() != RUN_METRIC) })
|
||||
);
|
||||
for resource in [
|
||||
"x/SKILL.md",
|
||||
"x\\SKILL.md",
|
||||
"hidden/SKILL.md",
|
||||
"selected/SKILL.md",
|
||||
] {
|
||||
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: resource,
|
||||
kind: SkillInvocationKind::Implicit,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
assert_eq!(
|
||||
vec!["x/SKILL.md"],
|
||||
thread_store
|
||||
.get::<SkillsThreadState>()
|
||||
.unwrap()
|
||||
.recent_skill_invocations
|
||||
.snapshot()
|
||||
);
|
||||
}
|
||||
// Both observers disappear before either ranking starts. Their queued unique
|
||||
// invocations must still be scored against their own turn-start histories.
|
||||
drop(thread_store);
|
||||
drop(slots);
|
||||
|
||||
let mut runs = BTreeMap::new();
|
||||
let mut hits = BTreeMap::new();
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
loop {
|
||||
// snapshot() exports deltas, so retain observations across polls.
|
||||
let snapshot = metrics.snapshot().unwrap();
|
||||
for metric in snapshot.scope_metrics().flat_map(ScopeMetrics::metrics) {
|
||||
let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
|
||||
continue;
|
||||
};
|
||||
for point in sum.data_points() {
|
||||
let attribute = |key: &str| {
|
||||
point
|
||||
.attributes()
|
||||
.find(|attribute| attribute.key.as_str() == key)
|
||||
.unwrap()
|
||||
.value
|
||||
.as_str()
|
||||
.to_string()
|
||||
};
|
||||
match metric.name() {
|
||||
RUN_METRIC => {
|
||||
*runs.entry(attribute("method")).or_insert(0) += point.value()
|
||||
}
|
||||
"codex.skills.shadow_selection.invocation" => {
|
||||
assert_eq!("ascii_latin", attribute("query_script"));
|
||||
*hits
|
||||
.entry((attribute("method"), attribute("hit"), attribute("rank")))
|
||||
.or_insert(0) += point.value();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
if hits.values().sum::<u64>() >= 24 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
let mut expected_runs = BTreeMap::new();
|
||||
let mut expected_hits = BTreeMap::new();
|
||||
for (method, uses_history) in [
|
||||
("weighted_lexical_v1", false),
|
||||
("fielded_bm25_v1", false),
|
||||
("character_ngram_v1", false),
|
||||
("multi_query_lexical_v1", false),
|
||||
("rrf_lexical_char_v1", false),
|
||||
("routing_card_exact_v1", false),
|
||||
("character_routing_card_v1", false),
|
||||
("lru_v1", true),
|
||||
("lru_plus_lexical_v1", true),
|
||||
("lru_plus_character_routing_v1", true),
|
||||
("lru_plus_lexical_character_routing_v1", true),
|
||||
("task_context_fusion_v1", true),
|
||||
] {
|
||||
expected_runs.insert(method.to_string(), 2);
|
||||
expected_hits.insert(
|
||||
(method.to_string(), "false".to_string(), "miss".to_string()),
|
||||
if uses_history { 1 } else { 2 },
|
||||
);
|
||||
if uses_history {
|
||||
expected_hits.insert((method.to_string(), "true".to_string(), "1".to_string()), 1);
|
||||
}
|
||||
}
|
||||
assert_eq!(expected_runs, runs);
|
||||
assert_eq!(expected_hits, hits);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Hold an actual blocking worker until the test releases it. The timeout keeps a
|
||||
/// failing assertion from leaving runtime shutdown stuck on the test worker.
|
||||
struct BlockedSelector {
|
||||
started: tokio::sync::mpsc::UnboundedSender<()>,
|
||||
release: Mutex<std::sync::mpsc::Receiver<()>>,
|
||||
}
|
||||
|
||||
impl CheapSkillSelector for BlockedSelector {
|
||||
fn method(&self) -> &'static str {
|
||||
"blocked_test"
|
||||
}
|
||||
|
||||
fn select(
|
||||
&self,
|
||||
_query: &str,
|
||||
_documents: &[SkillSelectionDocument<'_>],
|
||||
_limit: usize,
|
||||
) -> CheapSkillSelection {
|
||||
self.started.send(()).unwrap();
|
||||
self.release
|
||||
.lock()
|
||||
.unwrap()
|
||||
.recv_timeout(Duration::from_secs(10))
|
||||
.unwrap();
|
||||
CheapSkillSelection::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ranking_budget_is_shared_and_held_until_blocking_work_finishes() {
|
||||
let _serial = RANKING_TESTS.acquire().await.unwrap();
|
||||
let mut workers = Vec::new();
|
||||
for index in 0..3 {
|
||||
let (started, mut starts) = tokio::sync::mpsc::unbounded_channel();
|
||||
let (release, releases) = std::sync::mpsc::channel();
|
||||
// Separate experiment instances must share the same process-wide budget.
|
||||
let experiment = Arc::new(ShadowSelectionExperiment {
|
||||
selectors: vec![Box::new(BlockedSelector {
|
||||
started,
|
||||
release: Mutex::new(releases),
|
||||
})],
|
||||
metrics_client: None,
|
||||
});
|
||||
let query = ShadowQuery {
|
||||
text: "test".to_string(),
|
||||
truncated: false,
|
||||
};
|
||||
let input = ShadowSelectionInput {
|
||||
catalog: SkillCatalog::default(),
|
||||
host_snapshot: None,
|
||||
task_snapshot: TaskContextSnapshot {
|
||||
query: query.clone(),
|
||||
recent_skills: Vec::new(),
|
||||
},
|
||||
query,
|
||||
eligible_ids: HashSet::new(),
|
||||
recent_skill_resources: Vec::new(),
|
||||
};
|
||||
let (invocations, pending) = tokio::sync::mpsc::unbounded_channel();
|
||||
drop(invocations);
|
||||
let mut worker = Box::pin(experiment.evaluate(input, pending, tracing::Span::none()));
|
||||
// Polling to the first await deterministically attempts admission.
|
||||
assert!(futures::poll!(worker.as_mut()).is_pending());
|
||||
if index < 2 {
|
||||
assert_eq!(
|
||||
Some(()),
|
||||
tokio::time::timeout(Duration::from_secs(10), starts.recv())
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert_eq!(1 - index, RANKING_SLOTS.available_permits());
|
||||
} else {
|
||||
assert!(starts.try_recv().is_err());
|
||||
}
|
||||
workers.push((tokio::spawn(worker), release, starts));
|
||||
}
|
||||
let (first, release_first, _) = workers.remove(0);
|
||||
release_first.send(()).unwrap();
|
||||
first.await.unwrap();
|
||||
let (second, release_second, _) = workers.remove(0);
|
||||
let (third, release_third, mut third_starts) = workers.remove(0);
|
||||
assert_eq!(
|
||||
Some(()),
|
||||
tokio::time::timeout(Duration::from_secs(10), third_starts.recv())
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
assert_eq!(0, RANKING_SLOTS.available_permits());
|
||||
|
||||
// Cancelling evaluation cannot free a slot still occupied by blocking work.
|
||||
third.abort();
|
||||
assert!(third.await.unwrap_err().is_cancelled());
|
||||
assert_eq!(0, RANKING_SLOTS.available_permits());
|
||||
release_second.send(()).unwrap();
|
||||
release_third.send(()).unwrap();
|
||||
second.await.unwrap();
|
||||
let _all_returned =
|
||||
tokio::time::timeout(Duration::from_secs(10), RANKING_SLOTS.acquire_many(2))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
// This shadow-selection experiment is temporary and should be removed after evaluation.
|
||||
//! Shadow-only skill ranking. Capture turn inputs before recording new relevance signals,
|
||||
//! then evaluate on a blocking worker without delaying model input construction.
|
||||
//! This experiment is temporary and should be removed after evaluation.
|
||||
|
||||
mod task_context;
|
||||
use task_context::TaskContextSnapshot;
|
||||
|
||||
pub(crate) use task_context::ShadowTaskContext;
|
||||
|
||||
@@ -17,6 +20,7 @@ use crate::HostSkillsSnapshot;
|
||||
use codex_extension_api::TurnInputContext;
|
||||
use codex_otel::MetricsClient;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::catalog::SkillCatalog;
|
||||
use crate::catalog::SkillCatalogEntry;
|
||||
@@ -39,13 +43,16 @@ use crate::dynamic_skill_selector::WeightedLexicalSkillSelector;
|
||||
const MAX_SHADOW_QUERY_BYTES: usize = 16 * 1024;
|
||||
const MAX_SHADOW_RESULTS: usize = 50;
|
||||
|
||||
// Share the CPU budget across threads and extension registries. Admission happens
|
||||
// only in the background; turn preparation never waits for a ranking slot.
|
||||
static RANKING_SLOTS: Semaphore = Semaphore::const_new(2);
|
||||
|
||||
const RUN_METRIC: &str = "codex.skills.shadow_selection";
|
||||
const DURATION_METRIC: &str = "codex.skills.shadow_selection.duration_ms";
|
||||
const CATALOG_ENTRY_COUNT_METRIC: &str = "codex.skills.shadow_selection.catalog_entries";
|
||||
const SELECTED_ENTRY_COUNT_METRIC: &str = "codex.skills.shadow_selection.selected_entries";
|
||||
const QUERY_TERM_COUNT_METRIC: &str = "codex.skills.shadow_selection.query_terms";
|
||||
const REDUCTION_BPS_METRIC: &str = "codex.skills.shadow_selection.reduction_bps";
|
||||
const INVOCATION_METRIC: &str = "codex.skills.shadow_selection.invocation";
|
||||
|
||||
pub(crate) struct ShadowSelectionExperiment {
|
||||
selectors: Vec<Box<dyn CheapSkillSelector>>,
|
||||
@@ -67,37 +74,156 @@ impl ShadowSelectionExperiment {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn run(
|
||||
&self,
|
||||
/// Freeze prediction inputs before this turn can contribute new invocation history.
|
||||
/// The returned state observes invocations even while its rankings are still pending.
|
||||
pub(crate) fn start(
|
||||
self: &Arc<Self>,
|
||||
input: &TurnInputContext<'_>,
|
||||
catalog: &SkillCatalog,
|
||||
catalog: SkillCatalog,
|
||||
explicitly_selected: &[SkillCatalogEntry],
|
||||
host_snapshot: Option<&HostSkillsSnapshot>,
|
||||
host_snapshot: Option<Arc<HostSkillsSnapshot>>,
|
||||
recent_skill_invocations: Arc<RecentSkillInvocations>,
|
||||
task_context: Arc<ShadowTaskContext>,
|
||||
) -> ShadowSelectionTurnState {
|
||||
let query = build_shadow_query(&input.user_input);
|
||||
let query_script = query_script_tag(&query.text);
|
||||
let task_snapshot = task_context.begin_turn(&input.turn_id, &query, &input.user_input);
|
||||
let recent_skill_resources = recent_skill_invocations.snapshot();
|
||||
let explicitly_selected_skill_resources = explicitly_selected
|
||||
.iter()
|
||||
.map(|entry| normalize_skill_resource(entry.main_prompt.as_str()))
|
||||
.collect::<HashSet<_>>();
|
||||
let mut eligible_ids = HashSet::new();
|
||||
let mut eligible_skill_resources = HashSet::new();
|
||||
for (id, entry) in catalog.entries.iter().enumerate() {
|
||||
// Invocation observation exists only for host shell use and cloud reads.
|
||||
if !entry.is_model_visible()
|
||||
|| !matches!(
|
||||
&entry.authority.kind,
|
||||
SkillSourceKind::Host | SkillSourceKind::Cloud
|
||||
)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let resource = normalize_skill_resource(entry.main_prompt.as_str());
|
||||
if !explicitly_selected_skill_resources.contains(&resource) {
|
||||
eligible_ids.insert(id);
|
||||
eligible_skill_resources.insert(resource);
|
||||
}
|
||||
}
|
||||
let (invocations, pending_invocations) = tokio::sync::mpsc::unbounded_channel();
|
||||
let state = ShadowSelectionTurnState {
|
||||
turn_id: input.turn_id.clone(),
|
||||
eligible_skill_resources,
|
||||
seen_skill_resources: Mutex::new(HashSet::new()),
|
||||
recent_skill_invocations,
|
||||
task_context: Arc::clone(&task_context),
|
||||
invocations,
|
||||
};
|
||||
|
||||
// Explicit intent is a relevance signal even if the subsequent prompt read fails.
|
||||
// It can update future turns now that this turn's prediction inputs are frozen.
|
||||
for entry in explicitly_selected.iter().filter(|entry| {
|
||||
entry.is_model_visible()
|
||||
&& matches!(
|
||||
&entry.authority.kind,
|
||||
SkillSourceKind::Host | SkillSourceKind::Cloud
|
||||
)
|
||||
}) {
|
||||
task_context.record(
|
||||
&input.turn_id,
|
||||
normalize_skill_resource(entry.main_prompt.as_str()),
|
||||
);
|
||||
}
|
||||
|
||||
let selection_input = ShadowSelectionInput {
|
||||
catalog,
|
||||
host_snapshot,
|
||||
query,
|
||||
eligible_ids,
|
||||
recent_skill_resources,
|
||||
task_snapshot,
|
||||
};
|
||||
tokio::spawn(Arc::clone(self).evaluate(
|
||||
selection_input,
|
||||
pending_invocations,
|
||||
tracing::Span::current(),
|
||||
));
|
||||
state
|
||||
}
|
||||
|
||||
/// Own the rankings and drain observations for this turn, including those queued
|
||||
/// during ranking. Dropping the turn's observer closes the channel after its last read.
|
||||
async fn evaluate(
|
||||
self: Arc<Self>,
|
||||
input: ShadowSelectionInput,
|
||||
mut invocations: tokio::sync::mpsc::UnboundedReceiver<String>,
|
||||
parent: tracing::Span,
|
||||
) {
|
||||
let query_script = query_script_tag(&input.query.text);
|
||||
let experiment = Arc::clone(&self);
|
||||
let Ok(permit) = RANKING_SLOTS.acquire().await else {
|
||||
return;
|
||||
};
|
||||
let ranked_selections = match tokio::task::spawn_blocking(move || {
|
||||
// A cancelled async task must not release capacity while its blocking
|
||||
// worker is still running. Release it before waiting for invocations.
|
||||
let _permit = permit;
|
||||
experiment.rank(input, parent)
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(rankings) => rankings,
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "shadow skill selection failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(metrics_client) = self.metrics_client.as_ref() else {
|
||||
return;
|
||||
};
|
||||
while let Some(skill_resource) = invocations.recv().await {
|
||||
for selection in &ranked_selections {
|
||||
let rank = selection
|
||||
.skill_resources
|
||||
.iter()
|
||||
.position(|candidate| candidate == &skill_resource)
|
||||
.map(|index| index + 1);
|
||||
let tags = [
|
||||
("method", selection.method),
|
||||
("hit", bool_tag(rank.is_some())),
|
||||
("rank", rank_bucket(rank)),
|
||||
("query_script", query_script),
|
||||
];
|
||||
let _ = metrics_client.counter(
|
||||
"codex.skills.shadow_selection.invocation",
|
||||
/*inc*/ 1,
|
||||
&tags,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "skills.shadow_selection.rank",
|
||||
level = "info",
|
||||
parent = &parent,
|
||||
skip_all,
|
||||
fields(catalog_entries = input.eligible_ids.len())
|
||||
)]
|
||||
fn rank(&self, input: ShadowSelectionInput, parent: tracing::Span) -> Vec<RankedSelection> {
|
||||
let ShadowSelectionInput {
|
||||
catalog,
|
||||
host_snapshot,
|
||||
query,
|
||||
eligible_ids,
|
||||
recent_skill_resources,
|
||||
task_snapshot,
|
||||
} = input;
|
||||
let documents = catalog
|
||||
.entries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, entry)| {
|
||||
entry.is_model_visible()
|
||||
// Invocation observation currently exists only for host shell use and
|
||||
// cloud reads. Keep the candidate set aligned with that universe.
|
||||
&& matches!(
|
||||
&entry.authority.kind,
|
||||
SkillSourceKind::Host | SkillSourceKind::Cloud
|
||||
)
|
||||
&& !explicitly_selected_skill_resources
|
||||
.contains(&normalize_skill_resource(entry.main_prompt.as_str()))
|
||||
})
|
||||
.filter(|(id, _)| eligible_ids.contains(id))
|
||||
.map(|(id, entry)| SkillSelectionDocument {
|
||||
id,
|
||||
name: entry.name.as_str(),
|
||||
@@ -106,10 +232,6 @@ impl ShadowSelectionExperiment {
|
||||
dependencies: entry.dependencies.as_ref(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let eligible_ids = documents
|
||||
.iter()
|
||||
.map(|document| document.id)
|
||||
.collect::<HashSet<_>>();
|
||||
let eligible_skill_ids_by_resource = documents
|
||||
.iter()
|
||||
.map(|document| {
|
||||
@@ -119,16 +241,12 @@ impl ShadowSelectionExperiment {
|
||||
)
|
||||
})
|
||||
.collect::<HashMap<_, _>>();
|
||||
let eligible_skill_resources = eligible_skill_ids_by_resource
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<HashSet<_>>();
|
||||
let recent_skill_ids = recent_skill_invocations
|
||||
.snapshot()
|
||||
let recent_skill_ids = recent_skill_resources
|
||||
.iter()
|
||||
.filter_map(|resource| eligible_skill_ids_by_resource.get(resource).copied())
|
||||
.collect();
|
||||
let routing_selector = CharacterRoutingCardSkillSelector::new(catalog, host_snapshot);
|
||||
let routing_selector =
|
||||
CharacterRoutingCardSkillSelector::new(&catalog, host_snapshot.as_deref());
|
||||
let lru_selector = LruSkillSelector::new(recent_skill_ids);
|
||||
let lru_plus_lexical_selector = LruPlusLexicalSkillSelector::new(lru_selector.clone());
|
||||
let lru_plus_character_selector = LruPlusCharacterRoutingSkillSelector::new(
|
||||
@@ -203,68 +321,7 @@ impl ShadowSelectionExperiment {
|
||||
);
|
||||
}
|
||||
|
||||
// Explicit intent is a relevance signal even if the subsequent prompt read fails.
|
||||
// Keep it out of the implicit-only controls and freeze predictions before recording it.
|
||||
for entry in explicitly_selected.iter().filter(|entry| {
|
||||
entry.is_model_visible()
|
||||
&& matches!(
|
||||
&entry.authority.kind,
|
||||
SkillSourceKind::Host | SkillSourceKind::Cloud
|
||||
)
|
||||
}) {
|
||||
task_context.record(
|
||||
&input.turn_id,
|
||||
normalize_skill_resource(entry.main_prompt.as_str()),
|
||||
);
|
||||
}
|
||||
|
||||
ShadowSelectionTurnState {
|
||||
ranked_selections,
|
||||
turn_id: input.turn_id.clone(),
|
||||
query_script,
|
||||
eligible_skill_resources,
|
||||
seen_skill_resources: Mutex::new(HashSet::new()),
|
||||
recent_skill_invocations,
|
||||
task_context,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn record_invocation(&self, state: &ShadowSelectionTurnState, skill_resource: &str) {
|
||||
let skill_resource = normalize_skill_resource(skill_resource);
|
||||
if !state.eligible_skill_resources.contains(&skill_resource) {
|
||||
return;
|
||||
}
|
||||
if !state
|
||||
.seen_skill_resources
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.insert(skill_resource.clone())
|
||||
{
|
||||
return;
|
||||
}
|
||||
state
|
||||
.recent_skill_invocations
|
||||
.record(skill_resource.clone());
|
||||
state
|
||||
.task_context
|
||||
.record(&state.turn_id, skill_resource.clone());
|
||||
let Some(metrics_client) = self.metrics_client.as_ref() else {
|
||||
return;
|
||||
};
|
||||
for selection in &state.ranked_selections {
|
||||
let rank = selection
|
||||
.skill_resources
|
||||
.iter()
|
||||
.position(|candidate| candidate == &skill_resource)
|
||||
.map(|index| index + 1);
|
||||
let tags = [
|
||||
("method", selection.method),
|
||||
("hit", bool_tag(rank.is_some())),
|
||||
("rank", rank_bucket(rank)),
|
||||
("query_script", state.query_script),
|
||||
];
|
||||
let _ = metrics_client.counter(INVOCATION_METRIC, /*inc*/ 1, &tags);
|
||||
}
|
||||
ranked_selections
|
||||
}
|
||||
|
||||
fn record_metrics(&self, observation: ShadowSelectionObservation<'_>) {
|
||||
@@ -317,13 +374,42 @@ impl ShadowSelectionExperiment {
|
||||
}
|
||||
|
||||
pub(crate) struct ShadowSelectionTurnState {
|
||||
ranked_selections: Vec<RankedSelection>,
|
||||
turn_id: String,
|
||||
query_script: &'static str,
|
||||
eligible_skill_resources: HashSet<String>,
|
||||
seen_skill_resources: Mutex<HashSet<String>>,
|
||||
recent_skill_invocations: Arc<RecentSkillInvocations>,
|
||||
task_context: Arc<ShadowTaskContext>,
|
||||
invocations: tokio::sync::mpsc::UnboundedSender<String>,
|
||||
}
|
||||
|
||||
impl ShadowSelectionTurnState {
|
||||
/// Advance relevance history immediately; scoring can wait for the frozen rankings.
|
||||
pub(crate) fn record_invocation(&self, skill_resource: &str) {
|
||||
let skill_resource = normalize_skill_resource(skill_resource);
|
||||
if !self.eligible_skill_resources.contains(&skill_resource)
|
||||
|| !self
|
||||
.seen_skill_resources
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.insert(skill_resource.clone())
|
||||
{
|
||||
return;
|
||||
}
|
||||
self.recent_skill_invocations.record(skill_resource.clone());
|
||||
self.task_context
|
||||
.record(&self.turn_id, skill_resource.clone());
|
||||
let _ = self.invocations.send(skill_resource);
|
||||
}
|
||||
}
|
||||
|
||||
/// Owned turn-start inputs; workers never read the evolving relevance histories.
|
||||
struct ShadowSelectionInput {
|
||||
catalog: SkillCatalog,
|
||||
host_snapshot: Option<Arc<HostSkillsSnapshot>>,
|
||||
query: ShadowQuery,
|
||||
eligible_ids: HashSet<usize>,
|
||||
recent_skill_resources: Vec<String>,
|
||||
task_snapshot: TaskContextSnapshot,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
|
||||
@@ -44,7 +44,6 @@ use crate::catalog::SkillCatalogEntry;
|
||||
use crate::catalog::SkillSourceKind;
|
||||
use crate::provider::SkillListQuery;
|
||||
use crate::provider::attribute_executor_plugins;
|
||||
use crate::shadow_selection_experiment::ShadowSelectionExperiment;
|
||||
use crate::sources::SkillProviders;
|
||||
use crate::state::SkillsSessionState;
|
||||
use crate::state::SkillsThreadState;
|
||||
@@ -66,7 +65,6 @@ pub(crate) fn skill_tools(
|
||||
executor_query: Option<SkillListQuery>,
|
||||
selected_plugins: Option<Arc<SelectedPluginSnapshot>>,
|
||||
sandbox_contexts: Option<Arc<HashMap<String, FileSystemSandboxContext>>>,
|
||||
shadow_selection: Arc<ShadowSelectionExperiment>,
|
||||
) -> Vec<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
|
||||
let Some(thread_state) = thread_store.get::<SkillsThreadState>() else {
|
||||
return Vec::new();
|
||||
@@ -89,7 +87,6 @@ pub(crate) fn skill_tools(
|
||||
selected_plugins,
|
||||
sandbox_contexts,
|
||||
executor_catalog: Arc::new(OnceCell::new()),
|
||||
shadow_selection,
|
||||
};
|
||||
vec![
|
||||
Arc::new(list::ListTool {
|
||||
@@ -225,7 +222,6 @@ struct SkillToolContext {
|
||||
selected_plugins: Option<Arc<SelectedPluginSnapshot>>,
|
||||
sandbox_contexts: Option<Arc<HashMap<String, FileSystemSandboxContext>>>,
|
||||
executor_catalog: Arc<OnceCell<SkillCatalog>>,
|
||||
shadow_selection: Arc<ShadowSelectionExperiment>,
|
||||
}
|
||||
|
||||
impl SkillToolContext {
|
||||
|
||||
@@ -221,9 +221,7 @@ impl<'call> ToolExecutor<ToolCall<'call>> for ReadTool {
|
||||
.thread_state
|
||||
.shadow_selection_turn(&call.turn_id)
|
||||
{
|
||||
self.context
|
||||
.shadow_selection
|
||||
.record_invocation(&state, main_prompt.as_str());
|
||||
state.record_invocation(main_prompt.as_str());
|
||||
}
|
||||
let start = parse_pagination_cursor(
|
||||
args.cursor.as_deref(),
|
||||
|
||||
@@ -893,19 +893,24 @@ async fn shadow_selection_uses_host_catalog_when_instructions_are_disabled() ->
|
||||
.is_none()
|
||||
);
|
||||
assert!(fragments.is_empty());
|
||||
let snapshot = metrics.snapshot()?;
|
||||
let catalog_entry_counts = snapshot
|
||||
.scope_metrics()
|
||||
let snapshots = shadow_task_context_tests::collect_shadow_observations(
|
||||
&metrics,
|
||||
"codex.skills.shadow_selection.catalog_entries",
|
||||
/*expected*/ 12,
|
||||
)
|
||||
.await?;
|
||||
let catalog_entry_counts = snapshots
|
||||
.iter()
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ResourceMetrics::scope_metrics)
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
|
||||
.find(|metric| metric.name() == "codex.skills.shadow_selection.catalog_entries")
|
||||
.map(|metric| match metric.data() {
|
||||
.filter(|metric| metric.name() == "codex.skills.shadow_selection.catalog_entries")
|
||||
.flat_map(|metric| match metric.data() {
|
||||
AggregatedMetrics::F64(MetricData::Histogram(histogram)) => histogram
|
||||
.data_points()
|
||||
.map(opentelemetry_sdk::metrics::data::HistogramDataPoint::sum)
|
||||
.collect::<Vec<_>>(),
|
||||
.map(opentelemetry_sdk::metrics::data::HistogramDataPoint::sum),
|
||||
data => panic!("unexpected shadow catalog metric data: {data:?}"),
|
||||
})
|
||||
.ok_or("shadow catalog metric should be recorded")?;
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert!(
|
||||
catalog_entry_counts.iter().all(|count| *count == 1.0),
|
||||
@@ -997,42 +1002,54 @@ async fn shadow_lru_selector_recovers_a_skill_invoked_on_an_earlier_turn() -> Te
|
||||
.await;
|
||||
}
|
||||
|
||||
let snapshot = metrics.snapshot()?;
|
||||
let metric = snapshot
|
||||
.scope_metrics()
|
||||
let snapshots = shadow_task_context_tests::collect_shadow_observations(
|
||||
&metrics,
|
||||
"codex.skills.shadow_selection.invocation",
|
||||
/*expected*/ 24,
|
||||
)
|
||||
.await?;
|
||||
let selector_hits = snapshots
|
||||
.iter()
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ResourceMetrics::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 selector_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 !matches!(
|
||||
method.as_ref(),
|
||||
"lru_v1"
|
||||
| "lru_plus_lexical_v1"
|
||||
| "lru_plus_character_routing_v1"
|
||||
| "lru_plus_lexical_character_routing_v1"
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
let hit = point
|
||||
.attributes()
|
||||
.find(|attribute| attribute.key.as_str() == "hit")?
|
||||
.value
|
||||
.as_str()
|
||||
.to_string();
|
||||
Some((method.to_string(), hit, point.value()))
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
data => panic!("unexpected shadow invocation metric data: {data:?}"),
|
||||
};
|
||||
selector_hits.sort();
|
||||
.filter(|metric| metric.name() == "codex.skills.shadow_selection.invocation")
|
||||
.flat_map(|metric| match metric.data() {
|
||||
AggregatedMetrics::U64(MetricData::Sum(sum)) => sum.data_points(),
|
||||
data => panic!("unexpected shadow invocation metric data: {data:?}"),
|
||||
})
|
||||
.filter_map(|point| {
|
||||
let method = point
|
||||
.attributes()
|
||||
.find(|attribute| attribute.key.as_str() == "method")?
|
||||
.value
|
||||
.as_str();
|
||||
if !matches!(
|
||||
method.as_ref(),
|
||||
"lru_v1"
|
||||
| "lru_plus_lexical_v1"
|
||||
| "lru_plus_character_routing_v1"
|
||||
| "lru_plus_lexical_character_routing_v1"
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
let hit = point
|
||||
.attributes()
|
||||
.find(|attribute| attribute.key.as_str() == "hit")?
|
||||
.value
|
||||
.as_str()
|
||||
.to_string();
|
||||
Some((method.to_string(), hit, point.value()))
|
||||
})
|
||||
.fold(
|
||||
std::collections::BTreeMap::new(),
|
||||
|mut totals, (method, hit, count)| {
|
||||
*totals.entry((method, hit)).or_insert(0) += count;
|
||||
totals
|
||||
},
|
||||
)
|
||||
.into_iter()
|
||||
.map(|((method, hit), count)| (method, hit, count))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(
|
||||
vec![
|
||||
|
||||
@@ -1,9 +1,53 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use opentelemetry_sdk::metrics::data::ResourceMetrics;
|
||||
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
// The runtime reader emits deltas. Retain every snapshot while waiting for the
|
||||
// background evaluations so assertions include observations from earlier polls.
|
||||
pub(super) async fn collect_shadow_observations(
|
||||
metrics: &MetricsClient,
|
||||
name: &str,
|
||||
expected: u64,
|
||||
) -> Result<Vec<ResourceMetrics>, Box<dyn std::error::Error>> {
|
||||
tokio::time::timeout(std::time::Duration::from_secs(10), async {
|
||||
let mut snapshots = Vec::new();
|
||||
let mut total = 0;
|
||||
loop {
|
||||
let snapshot = metrics.snapshot()?;
|
||||
total += snapshot
|
||||
.scope_metrics()
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
|
||||
.find(|metric| metric.name() == name)
|
||||
.map(|metric| match metric.data() {
|
||||
AggregatedMetrics::U64(MetricData::Sum(sum)) => sum
|
||||
.data_points()
|
||||
.map(opentelemetry_sdk::metrics::data::SumDataPoint::value)
|
||||
.sum::<u64>(),
|
||||
AggregatedMetrics::F64(MetricData::Histogram(histogram)) => histogram
|
||||
.data_points()
|
||||
.map(opentelemetry_sdk::metrics::data::HistogramDataPoint::count)
|
||||
.sum::<u64>(),
|
||||
data => panic!("unexpected shadow metric data: {data:?}"),
|
||||
})
|
||||
.unwrap_or(0);
|
||||
snapshots.push(snapshot);
|
||||
if total >= expected {
|
||||
assert_eq!(
|
||||
expected, total,
|
||||
"shadow observations must not be duplicated"
|
||||
);
|
||||
return Ok(snapshots);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
struct FailedReads(StaticSkillProvider);
|
||||
|
||||
impl SkillProvider for FailedReads {
|
||||
@@ -154,17 +198,21 @@ async fn task_context_recovers_prior_requests_and_explicit_intent_without_changi
|
||||
"lru_plus_character_routing_v1",
|
||||
"lru_plus_lexical_character_routing_v1",
|
||||
];
|
||||
let snapshot = metrics.snapshot()?;
|
||||
let metric = snapshot
|
||||
.scope_metrics()
|
||||
let snapshots = collect_shadow_observations(
|
||||
&metrics,
|
||||
"codex.skills.shadow_selection.invocation",
|
||||
/*expected*/ 48,
|
||||
)
|
||||
.await?;
|
||||
let actual = snapshots
|
||||
.iter()
|
||||
.flat_map(ResourceMetrics::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 exist")?;
|
||||
let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
|
||||
panic!("unexpected shadow metric: {:?}", metric.data());
|
||||
};
|
||||
let actual = sum
|
||||
.data_points()
|
||||
.filter(|metric| metric.name() == "codex.skills.shadow_selection.invocation")
|
||||
.flat_map(|metric| match metric.data() {
|
||||
AggregatedMetrics::U64(MetricData::Sum(sum)) => sum.data_points(),
|
||||
data => panic!("unexpected shadow metric: {data:?}"),
|
||||
})
|
||||
.filter_map(|point| {
|
||||
let attribute = |key| {
|
||||
point
|
||||
|
||||
Reference in New Issue
Block a user