Store turn skill state in extension data (#36740)

## What changed

- Store each turn's host skills snapshot in `ExtensionData` and expose it through
  `TurnContext::skills_snapshot`.
- Lazily keep implicit skill invocation deduplication state in the same turn-scoped
  store.
- Carry the skills snapshot into review turns without a separate
  `TurnSkillsContext` field.

GitOrigin-RevId: d71579b5d230bd0b5f3c2ad062861f74d614d88c
This commit is contained in:
jif
2026-08-03 11:00:40 +00:00
committed by copyberry
parent d6407d7359
commit 79479cdf09
6 changed files with 26 additions and 38 deletions
-1
View File
@@ -244,7 +244,6 @@ use self::turn::agent_message_text;
use self::turn::collect_explicit_app_ids_from_skill_items;
use self::turn::realtime_text_for_event;
use self::turn_context::TurnContext;
use self::turn_context::TurnSkillsContext;
#[cfg(test)]
mod rollout_reconstruction_tests;
+1 -2
View File
@@ -106,7 +106,7 @@ pub(super) async fn spawn_review_thread(
let extension_data = Arc::new(codex_extension_api::ExtensionData::new(
review_turn_id.clone(),
));
extension_data.insert(parent_turn_context.turn_skills.snapshot.clone());
extension_data.insert(parent_turn_context.skills_snapshot().as_ref().clone());
let review_turn_context = TurnContext {
sub_id: review_turn_id.clone(),
@@ -147,7 +147,6 @@ pub(super) async fn spawn_review_thread(
dynamic_tools: parent_turn_context.dynamic_tools.clone(),
turn_metadata_state,
extension_data,
turn_skills: TurnSkillsContext::new(parent_turn_context.turn_skills.snapshot.clone()),
turn_timing_state: Arc::new(TurnTimingState::default()),
terminal_error: Arc::new(Mutex::new(None)),
server_model_warning_emitted: AtomicBool::new(false),
+3 -8
View File
@@ -5236,20 +5236,15 @@ enabled = false
let child_turn = session
.new_default_turn_with_sub_id("role-skill-turn".to_string())
.await;
let child_skill = child_turn
.turn_skills
.snapshot
let skills_snapshot = child_turn.skills_snapshot();
let child_skill = skills_snapshot
.outcome()
.skills
.iter()
.find(|skill| skill.name == "demo-skill")
.expect("demo skill should be discovered");
assert_eq!(
child_turn
.turn_skills
.snapshot
.outcome()
.is_skill_enabled(child_skill),
skills_snapshot.outcome().is_skill_enabled(child_skill),
false
);
}
+4 -2
View File
@@ -685,7 +685,8 @@ async fn required_mcp_servers_for_input(
} else {
HashMap::new()
};
let skills_outcome = turn_context.turn_skills.snapshot.outcome();
let skills_snapshot = turn_context.skills_snapshot();
let skills_outcome = skills_snapshot.outcome();
let mentioned_skills = collect_explicit_skill_mentions(
user_input,
&skills_outcome.skills,
@@ -757,7 +758,8 @@ async fn build_skills_and_plugins(
} else {
Vec::new()
};
let skills_outcome = turn_context.turn_skills.snapshot.outcome();
let skills_snapshot = turn_context.skills_snapshot();
let skills_outcome = skills_snapshot.outcome();
let connector_slug_counts = build_connector_slug_counts(&available_connectors);
let extension_injection_items =
build_extension_turn_input_items(sess, step_context, user_input, cancellation_token)
+8 -19
View File
@@ -25,21 +25,6 @@ use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tracing::instrument;
#[derive(Clone, Debug)]
pub(crate) struct TurnSkillsContext {
pub(crate) snapshot: HostSkillsSnapshot,
pub(crate) implicit_invocation_seen_skills: Arc<Mutex<HashSet<String>>>,
}
impl TurnSkillsContext {
pub(crate) fn new(snapshot: HostSkillsSnapshot) -> Self {
Self {
snapshot,
implicit_invocation_seen_skills: Arc::new(Mutex::new(HashSet::new())),
}
}
}
pub(crate) type ShellSnapshotTask = Shared<BoxFuture<'static, Option<Arc<ShellSnapshotFile>>>>;
#[derive(Clone)]
@@ -151,7 +136,6 @@ pub struct TurnContext {
pub(crate) dynamic_tools: Vec<DynamicToolSpec>,
pub(crate) turn_metadata_state: Arc<TurnMetadataState>,
pub(crate) extension_data: Arc<codex_extension_api::ExtensionData>,
pub(crate) turn_skills: TurnSkillsContext,
pub(crate) turn_timing_state: Arc<TurnTimingState>,
pub(crate) terminal_error: Arc<Mutex<Option<ErrorEvent>>>,
pub(crate) server_model_warning_emitted: AtomicBool,
@@ -164,6 +148,13 @@ enum TurnMultiAgentRuntime {
}
impl TurnContext {
pub(crate) fn skills_snapshot(&self) -> Arc<HostSkillsSnapshot> {
let Some(snapshot) = self.extension_data.get::<HostSkillsSnapshot>() else {
unreachable!("every turn has a host skills snapshot");
};
snapshot
}
pub(crate) fn collaboration_mode(&self) -> CollaborationMode {
CollaborationMode {
mode: self.mode,
@@ -319,7 +310,6 @@ impl TurnContext {
dynamic_tools: self.dynamic_tools.clone(),
turn_metadata_state: self.turn_metadata_state.clone(),
extension_data: Arc::clone(&self.extension_data),
turn_skills: self.turn_skills.clone(),
turn_timing_state: Arc::clone(&self.turn_timing_state),
terminal_error: Arc::clone(&self.terminal_error),
server_model_warning_emitted: AtomicBool::new(
@@ -548,7 +538,7 @@ impl Session {
));
let (current_date, timezone) = local_time_context();
let extension_data = Arc::new(codex_extension_api::ExtensionData::new(sub_id.clone()));
extension_data.insert(skills_snapshot.clone());
extension_data.insert(skills_snapshot);
TurnContext {
sub_id,
trace_id: current_span_trace_id(),
@@ -589,7 +579,6 @@ impl Session {
dynamic_tools: session_configuration.dynamic_tools.clone(),
turn_metadata_state,
extension_data,
turn_skills: TurnSkillsContext::new(skills_snapshot),
turn_timing_state: Arc::new(TurnTimingState::default()),
terminal_error: Arc::new(Mutex::new(None)),
server_model_warning_emitted: AtomicBool::new(false),
+10 -6
View File
@@ -10,6 +10,8 @@ use codex_otel::sanitize_metric_tag_value;
use codex_protocol::protocol::SkillScope;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_plugins::PluginSkillRoot;
use std::collections::HashSet;
use tokio::sync::Mutex;
pub use codex_core_skills::SkillError;
pub use codex_core_skills::SkillLoadOutcome;
@@ -31,6 +33,9 @@ pub use codex_core_skills::system;
pub use codex_skills::SkillMetadata;
pub use codex_skills::SkillPolicy;
#[derive(Debug, Default)]
struct ImplicitSkillInvocations(Mutex<HashSet<String>>);
pub(crate) fn skills_load_input_from_config(
config: &Config,
effective_skill_roots: Vec<PluginSkillRoot>,
@@ -50,7 +55,7 @@ pub(crate) async fn maybe_emit_implicit_skill_invocation(
workdir: &AbsolutePathBuf,
) {
let Some(candidate) = detect_implicit_skill_invocation_for_command(
turn_context.turn_skills.snapshot.outcome(),
turn_context.skills_snapshot().outcome(),
command,
workdir,
) else {
@@ -74,11 +79,10 @@ pub(crate) async fn maybe_emit_implicit_skill_invocation(
let skill_name = invocation.skill_name.clone();
let seen_key = format!("{skill_scope}:{skill_path}:{skill_name}");
let inserted = {
let mut seen_skills = turn_context
.turn_skills
.implicit_invocation_seen_skills
.lock()
.await;
let skill_invocations = turn_context
.extension_data
.get_or_init(ImplicitSkillInvocations::default);
let mut seen_skills = skill_invocations.0.lock().await;
seen_skills.insert(seen_key)
};
if !inserted {