mirror of
https://github.com/openai/codex.git
synced 2026-09-13 11:47:17 +00:00
Add task context to shadow skill selection (#39008)
## Why Short continuation prompts such as `continue` do not contain enough context for skill selection on their own. ## What changed - Add the `task_context_fusion_v1` shadow selector, combining the current request with up to two prior substantive requests and recently relevant skills. - Record explicit skill intent and successful skill invocations for future turns, while excluding same-turn observations from predictions. - Bound retained requests, augmented queries, and skill history, including safe truncation at UTF-8 character boundaries. ## Testing Add unit and extension tests for continuation prompts, explicit intent, turn isolation, cold thread state, bounded history, and unchanged control selectors. GitOrigin-RevId: 72eca3f0d64620a0d24e95d5675f126ac982f8c7
This commit is contained in:
@@ -395,11 +395,12 @@ where
|
||||
let shadow_selected_entries =
|
||||
collect_explicit_skill_mentions(&input.user_input, &shadow_catalog);
|
||||
Some(self.shadow_selection.run(
|
||||
&input.user_input,
|
||||
&input,
|
||||
&shadow_catalog,
|
||||
&shadow_selected_entries,
|
||||
host_snapshot.as_deref(),
|
||||
Arc::clone(&thread_state.recent_skill_invocations),
|
||||
Arc::clone(&thread_state.shadow_task_context),
|
||||
))
|
||||
} else {
|
||||
None
|
||||
|
||||
@@ -1,5 +1,9 @@
|
||||
// This shadow-selection experiment is temporary and should be removed after evaluation.
|
||||
|
||||
mod task_context;
|
||||
|
||||
pub(crate) use task_context::ShadowTaskContext;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::collections::VecDeque;
|
||||
@@ -10,6 +14,7 @@ use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
use crate::HostSkillsSnapshot;
|
||||
use codex_extension_api::TurnInputContext;
|
||||
use codex_otel::MetricsClient;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
|
||||
@@ -64,14 +69,16 @@ impl ShadowSelectionExperiment {
|
||||
|
||||
pub(crate) fn run(
|
||||
&self,
|
||||
inputs: &[UserInput],
|
||||
input: &TurnInputContext,
|
||||
catalog: &SkillCatalog,
|
||||
explicitly_selected: &[SkillCatalogEntry],
|
||||
host_snapshot: Option<&HostSkillsSnapshot>,
|
||||
recent_skill_invocations: Arc<RecentSkillInvocations>,
|
||||
task_context: Arc<ShadowTaskContext>,
|
||||
) -> ShadowSelectionTurnState {
|
||||
let query = build_shadow_query(inputs);
|
||||
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 explicitly_selected_skill_resources = explicitly_selected
|
||||
.iter()
|
||||
.map(|entry| normalize_skill_resource(entry.main_prompt.as_str()))
|
||||
@@ -132,9 +139,19 @@ impl ShadowSelectionExperiment {
|
||||
lru_selector.clone(),
|
||||
routing_selector.clone(),
|
||||
);
|
||||
let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 5);
|
||||
let task_selector = LruPlusLexicalCharacterRoutingSkillSelector::new(
|
||||
LruSkillSelector::new(
|
||||
task_snapshot
|
||||
.recent_skills
|
||||
.iter()
|
||||
.filter_map(|resource| eligible_skill_ids_by_resource.get(resource).copied())
|
||||
.collect(),
|
||||
),
|
||||
routing_selector.clone(),
|
||||
);
|
||||
let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 6);
|
||||
|
||||
for selector in self
|
||||
for (method, selector, query) in self
|
||||
.selectors
|
||||
.iter()
|
||||
.map(std::convert::AsRef::as_ref)
|
||||
@@ -145,14 +162,21 @@ impl ShadowSelectionExperiment {
|
||||
&lru_plus_character_selector as &dyn CheapSkillSelector,
|
||||
&lru_plus_lexical_character_selector as &dyn CheapSkillSelector,
|
||||
])
|
||||
.map(|selector| (selector.method(), selector, &query))
|
||||
.chain([(
|
||||
"task_context_fusion_v1",
|
||||
&task_selector as &dyn CheapSkillSelector,
|
||||
&task_snapshot.query,
|
||||
)])
|
||||
{
|
||||
let query_script = query_script_tag(&query.text);
|
||||
let start = Instant::now();
|
||||
let selection =
|
||||
selector.select(&query.text, &documents, /*limit*/ MAX_SHADOW_RESULTS);
|
||||
let duration = start.elapsed();
|
||||
let selected_ids = sanitize_selected_ids(&selection, &eligible_ids);
|
||||
self.record_metrics(ShadowSelectionObservation {
|
||||
method: selector.method(),
|
||||
method,
|
||||
selection: &selection,
|
||||
query_truncated_before_selection: query.truncated,
|
||||
query_script,
|
||||
@@ -161,14 +185,14 @@ impl ShadowSelectionExperiment {
|
||||
duration,
|
||||
});
|
||||
ranked_selections.push(RankedSelection {
|
||||
method: selector.method(),
|
||||
method,
|
||||
skill_resources: selected_ids
|
||||
.iter()
|
||||
.map(|id| normalize_skill_resource(catalog.entries[*id].main_prompt.as_str()))
|
||||
.collect(),
|
||||
});
|
||||
tracing::debug!(
|
||||
method = selector.method(),
|
||||
method,
|
||||
catalog_entries = documents.len(),
|
||||
selected_entries = selected_ids.len(),
|
||||
query_terms = selection.query_term_count,
|
||||
@@ -179,12 +203,29 @@ 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::Orchestrator
|
||||
)
|
||||
}) {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -204,6 +245,9 @@ impl ShadowSelectionExperiment {
|
||||
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;
|
||||
};
|
||||
@@ -274,10 +318,12 @@ 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>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -435,6 +481,7 @@ fn is_cjk(character: char) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct ShadowQuery {
|
||||
text: String,
|
||||
truncated: bool,
|
||||
@@ -479,5 +526,5 @@ fn push_bounded(destination: &mut String, value: &str) -> bool {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "shadow_selection_experiment_tests.rs"]
|
||||
#[path = "experiment_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,155 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::PoisonError;
|
||||
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_utils_string::take_bytes_at_char_boundary;
|
||||
|
||||
use super::ShadowQuery;
|
||||
|
||||
const MAX_PRIOR_REQUESTS: usize = 2;
|
||||
const MAX_REQUEST_BYTES: usize = 2 * 1024;
|
||||
const MAX_QUERY_BYTES: usize = 4 * 1024;
|
||||
const MAX_RECENT_SKILLS: usize = 50;
|
||||
|
||||
/// Shadow-only relevance history. New and reconstructed thread runtimes start cold.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct ShadowTaskContext(Mutex<TaskContextState>);
|
||||
|
||||
#[derive(Default)]
|
||||
struct TaskContextState {
|
||||
prior_requests: VecDeque<ShadowQuery>,
|
||||
recent_skills: VecDeque<String>,
|
||||
pending: Option<PendingTurn>,
|
||||
}
|
||||
|
||||
struct PendingTurn {
|
||||
id: String,
|
||||
request: Option<ShadowQuery>,
|
||||
recent_skills: VecDeque<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(super) struct TaskContextSnapshot {
|
||||
pub(super) query: ShadowQuery,
|
||||
pub(super) recent_skills: Vec<String>,
|
||||
}
|
||||
|
||||
impl ShadowTaskContext {
|
||||
pub(super) fn begin_turn(
|
||||
&self,
|
||||
turn_id: &str,
|
||||
current: &ShadowQuery,
|
||||
inputs: &[UserInput],
|
||||
) -> TaskContextSnapshot {
|
||||
let mut state = self.0.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if state.pending.as_ref().is_none_or(|turn| turn.id != turn_id) {
|
||||
if let Some(previous) = state.pending.take() {
|
||||
if let Some(request) = previous.request {
|
||||
state
|
||||
.prior_requests
|
||||
.retain(|prior| prior.text != request.text);
|
||||
state.prior_requests.push_front(request);
|
||||
state.prior_requests.truncate(MAX_PRIOR_REQUESTS);
|
||||
}
|
||||
for resource in previous.recent_skills.into_iter().rev() {
|
||||
remember_skill(&mut state.recent_skills, resource);
|
||||
}
|
||||
}
|
||||
state.pending = Some(PendingTurn {
|
||||
id: turn_id.to_string(),
|
||||
request: None,
|
||||
recent_skills: VecDeque::new(),
|
||||
});
|
||||
}
|
||||
|
||||
let retained = take_bytes_at_char_boundary(¤t.text, MAX_REQUEST_BYTES);
|
||||
let text = take_bytes_at_char_boundary(¤t.text, MAX_QUERY_BYTES);
|
||||
let mut query = ShadowQuery {
|
||||
text: text.to_string(),
|
||||
truncated: current.truncated || text.len() < current.text.len(),
|
||||
};
|
||||
for prior in &state.prior_requests {
|
||||
if prior.text == retained {
|
||||
continue;
|
||||
}
|
||||
if !query.text.is_empty() && query.text.len() < MAX_QUERY_BYTES {
|
||||
query.text.push('\n');
|
||||
}
|
||||
let part = take_bytes_at_char_boundary(
|
||||
&prior.text,
|
||||
MAX_QUERY_BYTES.saturating_sub(query.text.len()),
|
||||
);
|
||||
query.text.push_str(part);
|
||||
query.truncated |= prior.truncated || part.len() < prior.text.len();
|
||||
if part.len() < prior.text.len() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let recent_skills = state.recent_skills.iter().cloned().collect();
|
||||
if is_substantive(¤t.text, inputs)
|
||||
&& let Some(pending) = state.pending.as_mut()
|
||||
{
|
||||
pending.request = Some(ShadowQuery {
|
||||
text: retained.to_string(),
|
||||
truncated: current.truncated || retained.len() < current.text.len(),
|
||||
});
|
||||
}
|
||||
TaskContextSnapshot {
|
||||
query,
|
||||
recent_skills,
|
||||
}
|
||||
}
|
||||
|
||||
/// Records relevance evidence for future turns, never the active turn's predictions.
|
||||
pub(super) fn record(&self, turn_id: &str, resource: String) {
|
||||
let mut state = self.0.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(pending) = state.pending.as_mut().filter(|turn| turn.id == turn_id) {
|
||||
remember_skill(&mut pending.recent_skills, resource);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn remember_skill(skills: &mut VecDeque<String>, resource: String) {
|
||||
skills.retain(|previous| previous != &resource);
|
||||
skills.push_front(resource);
|
||||
skills.truncate(MAX_RECENT_SKILLS);
|
||||
}
|
||||
|
||||
fn is_substantive(text: &str, inputs: &[UserInput]) -> bool {
|
||||
if inputs.iter().any(|input| {
|
||||
matches!(input, UserInput::Skill { name, .. } | UserInput::Mention { name, .. } if !name.trim().is_empty())
|
||||
}) {
|
||||
return true;
|
||||
}
|
||||
let normalized = text
|
||||
.split(|character: char| !character.is_alphanumeric())
|
||||
.filter(|part| !part.is_empty())
|
||||
.map(str::to_lowercase)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ");
|
||||
!matches!(
|
||||
normalized.as_str(),
|
||||
"" | "yes"
|
||||
| "yep"
|
||||
| "yeah"
|
||||
| "ok"
|
||||
| "okay"
|
||||
| "sure"
|
||||
| "go"
|
||||
| "go ahead"
|
||||
| "continue"
|
||||
| "please continue"
|
||||
| "proceed"
|
||||
| "do it"
|
||||
| "do that"
|
||||
| "try again"
|
||||
| "retry"
|
||||
| "thanks"
|
||||
| "thank you"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "task_context_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,112 @@
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn query(text: &str) -> ShadowQuery {
|
||||
ShadowQuery {
|
||||
text: text.to_string(),
|
||||
truncated: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn begin(context: &ShadowTaskContext, turn_id: &str, text: &str) -> TaskContextSnapshot {
|
||||
context.begin_turn(turn_id, &query(text), &[])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prior_requests_survive_continuations_and_keep_short_substantive_input() {
|
||||
let context = ShadowTaskContext::default();
|
||||
begin(&context, "1", "Fix lint errors.");
|
||||
assert_eq!(
|
||||
query("continue\nFix lint errors."),
|
||||
begin(&context, "2", "continue").query
|
||||
);
|
||||
begin(&context, "3", "OK!!!");
|
||||
begin(&context, "4", "deploy");
|
||||
begin(&context, "5", "修复测试");
|
||||
assert_eq!(
|
||||
query("continue\n修复测试\ndeploy"),
|
||||
begin(&context, "6", "continue").query
|
||||
);
|
||||
assert_eq!(
|
||||
query("deploy\n修复测试"),
|
||||
begin(&context, "7", "deploy").query
|
||||
);
|
||||
assert!(is_substantive(
|
||||
"go",
|
||||
&[UserInput::Mention {
|
||||
name: "go".to_string(),
|
||||
path: "skill://go".to_string(),
|
||||
}]
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_turn_observations_are_pending_and_late_observations_are_ignored() {
|
||||
let context = ShadowTaskContext::default();
|
||||
begin(&context, "1", "continue");
|
||||
context.record("1", "a".to_string());
|
||||
assert_eq!(
|
||||
Vec::<String>::new(),
|
||||
begin(&context, "1", "continue").recent_skills
|
||||
);
|
||||
assert_eq!(vec!["a"], begin(&context, "2", "continue").recent_skills);
|
||||
context.record("1", "late".to_string());
|
||||
context.record("2", "a".to_string());
|
||||
context.record("2", "b".to_string());
|
||||
assert_eq!(vec!["a"], begin(&context, "2", "continue").recent_skills);
|
||||
assert_eq!(
|
||||
TaskContextSnapshot {
|
||||
query: query("continue"),
|
||||
recent_skills: vec!["b".to_string(), "a".to_string()],
|
||||
},
|
||||
begin(&context, "3", "continue")
|
||||
);
|
||||
assert_eq!(
|
||||
TaskContextSnapshot {
|
||||
query: query("continue"),
|
||||
recent_skills: Vec::new(),
|
||||
},
|
||||
begin(&ShadowTaskContext::default(), "3", "continue")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn history_and_augmented_query_are_bounded_at_utf8_boundaries() {
|
||||
let context = ShadowTaskContext::default();
|
||||
let first = "界".repeat(MAX_REQUEST_BYTES);
|
||||
let initial = begin(&context, "1", &first);
|
||||
assert_eq!(
|
||||
ShadowQuery {
|
||||
text: take_bytes_at_char_boundary(&first, MAX_QUERY_BYTES).to_string(),
|
||||
truncated: true,
|
||||
},
|
||||
initial.query
|
||||
);
|
||||
for index in 0..=MAX_RECENT_SKILLS {
|
||||
context.record("1", format!("skill-{index}"));
|
||||
}
|
||||
begin(&context, "2", &"é".repeat(MAX_REQUEST_BYTES));
|
||||
context.record("2", "skill-1".to_string());
|
||||
let expected = format!(
|
||||
"continue\n{}\n{}",
|
||||
"é".repeat(MAX_REQUEST_BYTES / 2),
|
||||
take_bytes_at_char_boundary(&first, MAX_REQUEST_BYTES)
|
||||
);
|
||||
assert_eq!(
|
||||
TaskContextSnapshot {
|
||||
query: ShadowQuery {
|
||||
text: take_bytes_at_char_boundary(&expected, MAX_QUERY_BYTES).to_string(),
|
||||
truncated: true,
|
||||
},
|
||||
recent_skills: std::iter::once("skill-1".to_string())
|
||||
.chain(
|
||||
(2..=MAX_RECENT_SKILLS)
|
||||
.rev()
|
||||
.map(|index| format!("skill-{index}"))
|
||||
)
|
||||
.collect(),
|
||||
},
|
||||
begin(&context, "3", "continue")
|
||||
);
|
||||
}
|
||||
@@ -23,6 +23,7 @@ use crate::provider::SkillListQuery;
|
||||
use crate::provider::SkillReadRequest;
|
||||
use crate::shadow_selection_experiment::RecentSkillInvocations;
|
||||
use crate::shadow_selection_experiment::ShadowSelectionTurnState;
|
||||
use crate::shadow_selection_experiment::ShadowTaskContext;
|
||||
use crate::sources::SkillProviders;
|
||||
|
||||
const MAX_CACHED_ORCHESTRATOR_RESOURCES: usize = 100;
|
||||
@@ -41,6 +42,7 @@ pub(crate) struct SkillsThreadState {
|
||||
orchestrator_cache: Mutex<Option<Arc<OrchestratorGenerationCache>>>,
|
||||
shadow_selection_turn: Mutex<Option<ShadowSelectionTurn>>,
|
||||
pub(crate) recent_skill_invocations: Arc<RecentSkillInvocations>,
|
||||
pub(crate) shadow_task_context: Arc<ShadowTaskContext>,
|
||||
}
|
||||
|
||||
impl SkillsThreadState {
|
||||
@@ -53,6 +55,7 @@ impl SkillsThreadState {
|
||||
orchestrator_cache: Mutex::new(None),
|
||||
shadow_selection_turn: Mutex::new(None),
|
||||
recent_skill_invocations: Arc::new(RecentSkillInvocations::default()),
|
||||
shadow_task_context: Arc::new(ShadowTaskContext::default()),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -81,6 +81,9 @@ use pretty_assertions::assert_eq;
|
||||
|
||||
type TestResult = Result<(), Box<dyn std::error::Error>>;
|
||||
|
||||
#[path = "skills_extension/shadow_task_context_tests.rs"]
|
||||
mod shadow_task_context_tests;
|
||||
|
||||
static NEXT_CODEX_HOME_ID: AtomicUsize = AtomicUsize::new(0);
|
||||
const SKILLS_INTRO_WITH_ABSOLUTE_PATHS: &str = "A skill is a set of instructions provided through a `SKILL.md` source. Below is the list of skills that can be used. Each entry includes a name, description, and source locator. `file` locators are on the host filesystem, `executor package` locators are owned by their execution environment, `orchestrator package` locators are opaque package identifiers, and `custom resource` locators use their provider's access mechanism.";
|
||||
const DEMO_SKILL_CONTENTS: &str =
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct FailedReads(StaticSkillProvider);
|
||||
|
||||
impl SkillProvider for FailedReads {
|
||||
fn list(&self, query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> {
|
||||
self.0.list(query)
|
||||
}
|
||||
|
||||
fn read(&self, _request: SkillReadRequest) -> SkillProviderFuture<'_, SkillReadResult> {
|
||||
Box::pin(async { Err(SkillProviderError::new("read unavailable")) })
|
||||
}
|
||||
|
||||
fn search(&self, request: SkillSearchRequest) -> SkillProviderFuture<'_, SkillSearchResult> {
|
||||
self.0.search(request)
|
||||
}
|
||||
}
|
||||
|
||||
fn text_input(text: &str) -> Vec<UserInput> {
|
||||
vec![UserInput::Text {
|
||||
text: text.to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}]
|
||||
}
|
||||
|
||||
fn explicit_input() -> Vec<UserInput> {
|
||||
vec![UserInput::Skill {
|
||||
name: "x".to_string(),
|
||||
path: PathBuf::from("x/SKILL.md"),
|
||||
}]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_context_recovers_prior_requests_and_explicit_intent_without_changing_controls()
|
||||
-> TestResult {
|
||||
let mut opaque = test_entry(SkillSourceKind::Host, "host", "host/x", "x/SKILL.md");
|
||||
opaque.description = "zzzz".to_string();
|
||||
let provider = Arc::new(FailedReads(StaticSkillProvider {
|
||||
catalog: SkillCatalog {
|
||||
entries: vec![
|
||||
test_entry(
|
||||
SkillSourceKind::Host,
|
||||
"host",
|
||||
"host/lint-fix",
|
||||
"lint-fix/SKILL.md",
|
||||
),
|
||||
opaque,
|
||||
],
|
||||
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 mut config = default_config();
|
||||
config.include_instructions = false;
|
||||
config.shadow_selection_enabled = true;
|
||||
for (thread_id, turns, resource) in [
|
||||
(
|
||||
"prior-request",
|
||||
vec![
|
||||
("a1", text_input("Fix lint errors.")),
|
||||
("a2", text_input("continue")),
|
||||
],
|
||||
"lint-fix/SKILL.md",
|
||||
),
|
||||
(
|
||||
"explicit-intent",
|
||||
vec![("b1", explicit_input()), ("b2", text_input("continue"))],
|
||||
"x/SKILL.md",
|
||||
),
|
||||
(
|
||||
"same-turn",
|
||||
vec![("c1", explicit_input()), ("c1", text_input("continue"))],
|
||||
"x/SKILL.md",
|
||||
),
|
||||
(
|
||||
"cold-thread",
|
||||
vec![("d1", text_input("continue"))],
|
||||
"x/SKILL.md",
|
||||
),
|
||||
] {
|
||||
let thread_store = ExtensionData::new(thread_id);
|
||||
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;
|
||||
let observed_turn = turns.last().ok_or("test needs a turn")?.0;
|
||||
for (turn_id, user_input) in turns {
|
||||
let fragments = registry.turn_input_contributors()[0]
|
||||
.contribute(
|
||||
TurnInputContext {
|
||||
turn_id: turn_id.to_string(),
|
||||
user_input,
|
||||
environments: Vec::new(),
|
||||
},
|
||||
/*extension_metrics*/ None,
|
||||
&session_store,
|
||||
&thread_store,
|
||||
&ExtensionData::new(turn_id),
|
||||
)
|
||||
.await;
|
||||
assert!(fragments.is_empty());
|
||||
}
|
||||
registry.skill_invocation_contributors()[0]
|
||||
.on_skill_invocation(SkillInvocationInput {
|
||||
session_store: &session_store,
|
||||
thread_store: &thread_store,
|
||||
turn_store: &ExtensionData::new(observed_turn),
|
||||
turn_id: observed_turn,
|
||||
skill_resource: resource,
|
||||
kind: SkillInvocationKind::Implicit,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
let controls = [
|
||||
"lru_v1",
|
||||
"weighted_lexical_v1",
|
||||
"lru_plus_lexical_v1",
|
||||
"lru_plus_character_routing_v1",
|
||||
"lru_plus_lexical_character_routing_v1",
|
||||
];
|
||||
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 exist")?;
|
||||
let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
|
||||
panic!("unexpected shadow metric: {:?}", metric.data());
|
||||
};
|
||||
let actual = sum
|
||||
.data_points()
|
||||
.filter_map(|point| {
|
||||
let attribute = |key| {
|
||||
point
|
||||
.attributes()
|
||||
.find(|value| value.key.as_str() == key)
|
||||
.map(|value| value.value.as_str().to_string())
|
||||
};
|
||||
let method = attribute("method")?;
|
||||
(controls.contains(&method.as_str()) || method == "task_context_fusion_v1")
|
||||
.then(|| ((method, attribute("hit").expect("hit tag")), point.value()))
|
||||
})
|
||||
.fold(BTreeMap::new(), |mut totals, (key, count)| {
|
||||
*totals.entry(key).or_insert(0) += count;
|
||||
totals
|
||||
});
|
||||
let mut expected = controls
|
||||
.into_iter()
|
||||
.map(|method| ((method.to_string(), "false".to_string()), 4))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
expected.insert(
|
||||
("task_context_fusion_v1".to_string(), "false".to_string()),
|
||||
/*value*/ 2,
|
||||
);
|
||||
expected.insert(
|
||||
("task_context_fusion_v1".to_string(), "true".to_string()),
|
||||
/*value*/ 2,
|
||||
);
|
||||
assert_eq!(expected, actual);
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user