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:
jif
2026-08-17 13:23:26 +00:00
committed by copyberry
parent def7ed5572
commit 21cfd369ef
8 changed files with 524 additions and 9 deletions

View File

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

View File

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

View File

@@ -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(&current.text, MAX_REQUEST_BYTES);
let text = take_bytes_at_char_boundary(&current.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(&current.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;

View File

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

View File

@@ -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()),
}
}

View File

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

View File

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