Use captured step model settings for extension context (#44202)

## Why

Model switches within a turn can leave extension context using stale model metadata. Skill catalogs, context windows, and metric attribution need to match the model captured for each sampling step.

## What changed

- Pass captured model metadata and step-specific telemetry to world-state contributors.
- Use that metadata for skill catalog budgets and usage instructions, preserving configured budget overrides.
- Supply the captured model's usable context window to turn-context contributors, including when rebuilding context.

## Testing

Add regression coverage for model switches during skill discovery, catalog budgets and metric attribution, and extension context windows after `new_context` and retained-step context rebuilds.

GitOrigin-RevId: c51b40b739a380b9767a54ea8c8ee6de45bf01c8
This commit is contained in:
Abhinav
2026-09-09 16:48:49 +00:00
committed by copyberry
parent b4507997e0
commit 8ff4aa8ee4
9 changed files with 447 additions and 99 deletions
+2 -2
View File
@@ -3964,7 +3964,7 @@ impl Session {
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store: turn_context.extension_data.as_ref(),
model_context_window: turn_context.model_context_window(),
model_context_window: step_context.settings.model_info.usable_context_window(),
})
.await
{
@@ -4066,7 +4066,7 @@ impl Session {
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store: turn_context.extension_data.as_ref(),
model_context_window: turn_context.model_context_window(),
model_context_window: step_context.settings.model_info.usable_context_window(),
})
.await
{
+21 -1
View File
@@ -10810,7 +10810,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_split_fi
#[tokio::test]
async fn build_initial_context_uses_retained_step_after_model_change() {
let (session, mut turn_context, _rx_event) =
let (mut session, mut turn_context, _rx_event) =
make_session_and_context_with_auth_and_config_and_rx(
CodexAuth::from_api_key("Test API Key"),
Vec::new(),
@@ -10821,6 +10821,17 @@ async fn build_initial_context_uses_retained_step_after_model_change() {
},
)
.await;
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.prompt_contributor(Arc::new(TurnContextExtensionTestContributor));
Arc::get_mut(&mut session)
.expect("unshared test session")
.services
.extensions = Arc::new(builder.build());
turn_context
.extension_data
.insert(TurnContextExtensionTestState {
expected_model_context_window: Some(64_000),
});
update_turn_settings_for_test(Arc::get_mut(&mut turn_context).unwrap(), |settings| {
let model_info = Arc::make_mut(&mut settings.model_info);
model_info.slug = "model-a".to_string();
@@ -10859,6 +10870,7 @@ async fn build_initial_context_uses_retained_step_after_model_change() {
let mut model_b = step_a.settings.model_info.as_ref().clone();
model_b.slug = "model-b".to_string();
model_b.context_window = Some(128_000);
model_b.effective_context_window_percent = 50;
model_b
.model_messages
.as_mut()
@@ -10879,6 +10891,7 @@ async fn build_initial_context_uses_retained_step_after_model_change() {
let initial_b = session
.build_initial_context_with_world_state(&step_b, &world_b)
.await;
let turn_contributions_b = session.build_turn_context_contribution_items(&step_b).await;
let (restored_a, restored_world) =
crate::compact::build_compaction_initial_context(&session, &retained).await;
@@ -10894,6 +10907,13 @@ async fn build_initial_context_uses_retained_step_after_model_change() {
assert!(!a_text.contains("<context_window>"));
assert!(b_text.contains("B instructions: pragmatic"));
assert!(!b_text.contains("A instructions:"));
assert!(!a_text.contains("turn context extension enabled"));
assert!(b_text.contains("turn context extension enabled"));
assert!(
developer_input_texts(&turn_contributions_b)
.join("\n")
.contains("turn context extension enabled")
);
assert!(
b_text.contains("<context_window>"),
"full-context metadata must use B's window even though the turn started without one"
+2 -1
View File
@@ -290,13 +290,14 @@ impl Session {
.map(|root| root.selected_root().clone())
.collect::<Vec<_>>();
let extension_metrics = super::extension_metrics::from_session_telemetry(
turn_context.session_telemetry.clone(),
step_context.session_telemetry.clone(),
);
for contributor in self.services.extensions.context_contributors() {
for section in contributor
.contribute_world_state(WorldStateContributionInput {
thread_id: self.thread_id(),
turn_id: turn_context.sub_id.as_str(),
model_info: &step_context.settings.model_info,
environments: &environments,
ready_selected_capability_roots: &ready_selected_capability_roots,
executor_capability_discovery: step_context
+289 -82
View File
@@ -1,5 +1,7 @@
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use anyhow::Result;
use codex_config::ConfigLayerEntry;
@@ -24,9 +26,16 @@ use codex_extension_api::SkillInvocationInput;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_mcp::CODEX_APPS_MCP_SERVER_NAME;
use codex_models_manager::bundled_models_response;
use codex_otel::OtelExporter;
use codex_otel::OtelHttpProtocol;
use codex_otel::OtelProvider;
use codex_otel::OtelSettings;
use codex_otel::THREAD_SKILLS_KEPT_TOTAL_METRIC;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::models::PermissionProfile;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::openai_models::TruncationPolicyConfig;
use codex_protocol::permissions::FileSystemAccessMode;
use codex_protocol::permissions::FileSystemPath;
@@ -36,7 +45,10 @@ use codex_protocol::permissions::FileSystemSpecialPath;
use codex_protocol::permissions::NetworkSandboxPolicy;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::protocol::TurnSettingsUpdate;
use codex_protocol::protocol::TurnSettingsUpdateOutcome;
use codex_protocol::user_input::UserInput;
use codex_skills_extension::ExecutorSkillProvider;
use codex_skills_extension::HostSkillProvider;
@@ -80,6 +92,8 @@ use core_test_support::test_codex::test_env;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use core_test_support::wait_for_mcp_server;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
@@ -108,6 +122,45 @@ struct CatalogSkillProvider {
catalog: SkillCatalog,
}
struct PausedCatalogSkillProvider {
inner: CatalogSkillProvider,
first_discovery: AtomicBool,
started: tokio::sync::mpsc::UnboundedSender<String>,
resume: tokio::sync::Semaphore,
}
impl SkillProvider for PausedCatalogSkillProvider {
fn list(&self, query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> {
Box::pin(async move {
if query.host_snapshot.is_none() {
return Ok(SkillCatalog::default());
}
if self.first_discovery.swap(/*val*/ false, Ordering::SeqCst) {
self.started
.send(query.turn_id.clone())
.expect("discovery receiver");
self.resume
.acquire()
.await
.expect("resume discovery")
.forget();
}
self.inner.list(query).await
})
}
fn read<'a>(
&'a self,
request: SkillReadRequest<'a>,
) -> SkillProviderFuture<'a, SkillReadResult> {
self.inner.read(request)
}
fn search(&self, request: SkillSearchRequest) -> SkillProviderFuture<'_, SkillSearchResult> {
self.inner.search(request)
}
}
#[derive(Debug)]
enum CapturedExtensionEvent {
Event(Box<Event>),
@@ -2488,113 +2541,266 @@ async fn production_turn_aliases_combined_skill_catalogs_under_shared_budget() -
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn production_turn_scales_extension_catalog_from_resolved_model_window() -> Result<()> {
assert_catalog_model_switch(/*max_context_tokens*/ None).await
}
async fn assert_catalog_model_switch(max_context_tokens: Option<usize>) -> Result<()> {
const MODEL_A: &str = "skills-model-a";
const MODEL_B: &str = "skills-model-b";
let server = responses::start_mock_server().await;
let response = responses::mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-a"),
responses::ev_function_call(
"plan",
"update_plan",
&json!({
"plan": [{"step": "Inspect skills", "status": "completed"}]
})
.to_string(),
),
ev_completed("resp-a"),
]),
sse(vec![ev_response_created("resp-b"), ev_completed("resp-b")]),
],
)
.await;
let codex_home = Arc::new(TempDir::new()?);
// Use the normal metrics sink to verify core's model attribution.
let telemetry = OtelProvider::try_new(&OtelSettings {
environment: "test".to_string(),
service_name: "skills-model-switch".to_string(),
service_version: env!("CARGO_PKG_VERSION").to_string(),
codex_home: codex_home.path().to_path_buf(),
exporter: OtelExporter::None,
trace_exporter: OtelExporter::None,
metrics_exporter: OtelExporter::OtlpHttp {
endpoint: format!("{}/metrics", server.uri()),
headers: Default::default(),
protocol: OtelHttpProtocol::Json,
tls: None,
},
runtime_metrics: true,
span_attributes: Default::default(),
tracestate: Default::default(),
})
.map_err(|error| anyhow::anyhow!("{error}"))?
.expect("metrics provider");
let skill_count = 800;
let mut included_counts = Vec::new();
for (context_window, max_context_window, expected_budget) in
[(Some(10_000), None, 200), (None, Some(400_000), 8_000)]
{
let server = responses::start_mock_server().await;
let response = mount_sse_once(
&server,
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
)
.await;
let source_kind = SkillSourceKind::Custom("test".to_string());
let catalog = SkillCatalog {
entries: (0..skill_count)
.map(|index| {
let name = format!("skill-{index:03}");
SkillCatalogEntry::new(
SkillPackageId(format!("test/{name}")),
SkillAuthority::new(source_kind.clone(), "test"),
name.clone(),
"A description long enough to keep the catalog under sustained budget pressure.",
SkillResourceId::new(format!("{name}/SKILL.md")),
)
.with_display_path(format!("skill://test/{name}/SKILL.md"))
let catalog = SkillCatalog {
entries: (0..skill_count)
.map(|index| {
let name = format!("skill-{index:03}");
SkillCatalogEntry::new(
SkillPackageId(format!("test/{name}")),
SkillAuthority::new(SkillSourceKind::Host, "test"),
name.clone(),
"A description long enough to keep the catalog under sustained budget pressure.",
SkillResourceId::new(format!("{name}/SKILL.md")),
)
.with_display_path(format!("skill://test/{name}/SKILL.md"))
})
.collect(),
warnings: Vec::new(),
};
let (started, mut discovery) = tokio::sync::mpsc::unbounded_channel();
let provider = Arc::new(PausedCatalogSkillProvider {
inner: CatalogSkillProvider { catalog },
first_discovery: AtomicBool::new(/*v*/ true),
started,
resume: tokio::sync::Semaphore::new(/*permits*/ 0),
});
let (event_tx, event_rx) = std::sync::mpsc::channel();
let mut extensions =
ExtensionRegistryBuilder::<Config>::with_event_sink(Arc::new(ChannelEventSink(event_tx)));
install_with_providers(
&mut extensions,
SkillProviders::new().with_host_provider(provider.clone()),
|config: &Config| SkillsExtensionConfig {
include_instructions: config.include_skill_instructions,
max_context_tokens: config.skill_max_context_tokens,
bundled_skills_enabled: false,
orchestrator_skills_enabled: false,
shadow_selection_enabled: false,
},
);
let mut builder = test_codex()
.with_home(codex_home)
.with_model(MODEL_A)
.with_extensions(Arc::new(extensions.build()))
.with_config(move |config| {
configure_catalog_test(config);
config
.features
.enable(Feature::StepModelSwitching)
.expect("model switching");
config.skill_max_context_tokens =
max_context_tokens.and_then(std::num::NonZeroUsize::new);
let template = bundled_models_response()
.expect("bundled models")
.models
.into_iter()
.find(|model| model.slug == "gpt-5.5")
.expect("gpt-5.5");
config.model_catalog = Some(ModelsResponse {
models: [
(MODEL_A, Some(10_000), None, false),
(MODEL_B, None, Some(400_000), true),
]
.into_iter()
.map(|(slug, window, max_window, usage)| {
let mut model = template.clone();
model.slug = slug.to_string();
model.context_window = window;
model.max_context_window = max_window;
// Distinguish the resolved window from the smaller usable limit.
model.effective_context_window_percent = 50;
model.include_skills_usage_instructions = usage;
model
})
.collect(),
warnings: Vec::new(),
};
let (event_tx, event_rx) = std::sync::mpsc::channel();
let mut extensions = ExtensionRegistryBuilder::<Config>::with_event_sink(Arc::new(
ChannelEventSink(event_tx),
));
install_with_providers(
&mut extensions,
SkillProviders::new().with_provider(SkillProviderSource::new(
source_kind,
"test",
Arc::new(StaticSkillProvider {
catalog,
main_prompt_contents: None,
}),
)),
|config: &Config| SkillsExtensionConfig {
include_instructions: config.include_skill_instructions,
max_context_tokens: config.skill_max_context_tokens,
bundled_skills_enabled: false,
orchestrator_skills_enabled: false,
shadow_selection_enabled: false,
},
);
let mut builder = test_codex()
.with_extensions(Arc::new(extensions.build()))
.with_model_info_override("gpt-5.5", move |model_info| {
model_info.context_window = context_window;
model_info.max_context_window = max_context_window;
})
.with_config(|config| {
config.include_skill_instructions = true;
});
let test = builder.build_with_auto_env(&server).await?;
});
let test = builder.build_with_auto_env(&server).await?;
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "Inspect the available skills.".to_string(),
text_elements: Vec::new(),
}]))
.await?;
let turn_id = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), discovery.recv())
.await?
.expect("A catalog discovery started");
let (reply, outcome) = tokio::sync::oneshot::channel();
test.codex
.submit(Op::TurnSettings {
turn_id,
update: TurnSettingsUpdate {
model: Some(MODEL_B.to_string()),
..Default::default()
},
reply,
})
.await?;
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 10), outcome).await??,
TurnSettingsUpdateOutcome::Applied
);
provider.resume.add_permits(/*n*/ 1);
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
test.submit_turn("Inspect the available skills.").await?;
let request = response.single_request();
let requests = response.requests();
assert_eq!(requests.len(), 2);
let mut included_counts = Vec::new();
for (request, (model, default_budget, include_usage)) in requests
.iter()
.zip([(MODEL_A, 200, false), (MODEL_B, 8_000, true)])
{
assert_eq!(request.body_json()["model"], model);
let developer_texts = request.message_input_texts("developer");
// B retains A in history and appends its own changed catalog.
let catalog_text = developer_texts
.iter()
.rev()
.find(|text| text.contains("skill://test/"))
.unwrap_or_else(|| {
panic!(
"production request should include the extension skill catalog, got {developer_texts:?}"
)
});
.expect("catalog in provider request");
assert_eq!(
catalog_text.contains("### How to use skills"),
include_usage
);
let metadata_lines = catalog_text
.lines()
.skip_while(|line| *line != "### Available skills")
.skip(1)
.skip(/*n*/ 1)
.take_while(|line| !line.starts_with("### "))
.filter(|line| line.starts_with("- "))
.collect::<Vec<_>>();
let metadata_cost = metadata_lines.iter().fold(0usize, |cost, line| {
cost.saturating_add(approx_token_count(&format!("{line}\n")))
});
let metadata_cost = metadata_lines
.iter()
.map(|line| approx_token_count(&format!("{line}\n")))
.sum::<usize>();
let budget = max_context_tokens.unwrap_or(default_budget);
assert!(
metadata_cost <= budget,
"catalog cost {metadata_cost} exceeds {budget}"
);
assert!(
metadata_cost > budget / 2,
"catalog incorrectly used the usable inference limit"
);
let included_count = metadata_lines
.iter()
.filter(|line| line.starts_with("- skill-"))
.count();
let warning = event_rx.try_recv()?.into_warning();
let omitted_count = skill_count - included_count;
assert!(catalog_text.contains("additional skills omitted"));
assert!(included_count > 0 && included_count < skill_count);
assert!(!catalog_text.contains(
"A description long enough to keep the catalog under sustained budget pressure."
));
assert!(metadata_cost <= expected_budget);
assert_eq!(
warning.message,
format!(
"Exceeded skills context budget. All skill descriptions were removed and {omitted_count} additional skills were not included in the model-visible skills list."
)
);
included_counts.push(included_count);
}
if max_context_tokens.is_some() {
assert_eq!(included_counts[0], included_counts[1]);
} else {
assert!(included_counts[0] < included_counts[1]);
}
let warnings = event_rx
.try_iter()
.map(|event| event.into_warning().message)
.collect::<Vec<_>>();
let mut expected_warnings = included_counts
.iter()
.map(|count| {
format!(
"Exceeded skills context budget. All skill descriptions were removed and {} additional skills were not included in the model-visible skills list.",
skill_count - count,
)
})
.collect::<Vec<_>>();
expected_warnings.dedup();
assert_eq!(warnings, expected_warnings);
assert!(included_counts[0] > 0);
assert!(included_counts[0] < included_counts[1]);
let snapshot = telemetry.metrics().expect("metrics client").snapshot()?;
let metric = snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
.find(|metric| metric.name() == THREAD_SKILLS_KEPT_TOTAL_METRIC)
.expect("catalog metrics");
let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else {
panic!("catalog metric should be a histogram");
};
let mut samples = histogram
.data_points()
.filter_map(|point| {
let attributes = point
.attributes()
.map(|attribute| (attribute.key.as_str(), attribute.value.as_str().to_string()))
.collect::<std::collections::BTreeMap<_, _>>();
(attributes.get("catalog_surface").map(String::as_str) == Some("host_world_state")
&& matches!(
attributes.get("model").map(String::as_str),
Some(MODEL_A | MODEL_B)
))
.then(|| (attributes["model"].clone(), point.count(), point.sum()))
})
.collect::<Vec<_>>();
samples.sort_by(|left, right| left.0.cmp(&right.0));
assert_eq!(
samples,
vec![
(MODEL_A.to_string(), 1, included_counts[0] as f64),
(MODEL_B.to_string(), 1, included_counts[1] as f64),
]
);
telemetry
.shutdown_with_timeout(Duration::from_secs(/*secs*/ 5))
.await?;
Ok(())
}
@@ -2642,6 +2848,7 @@ async fn production_turn_shares_catalog_budget_across_host_and_executor_sections
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn production_turn_uses_configured_skill_catalog_token_budget() -> Result<()> {
let server = responses::start_mock_server().await;
let response = mount_sse_once(
@@ -2684,7 +2891,7 @@ async fn production_turn_uses_configured_skill_catalog_token_budget() -> Result<
assert_full_descriptions(&executor_lines, &EXECUTOR_CATALOG);
assert!(metadata_cost(&combined_lines) <= 800);
Ok(())
assert_catalog_model_switch(Some(800)).await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -4,6 +4,12 @@ use codex_core::CodexThread;
use codex_core::TurnInputRequest;
use codex_core::config::Constrained;
use codex_core::config::TokenBudgetConfig;
use codex_extension_api::ContentItemKind;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PromptFragment;
use codex_extension_api::TurnContextContributionInput;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_models_manager::bundled_models_response;
@@ -63,6 +69,7 @@ use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
use test_case::test_case;
const MODEL_A: &str = "step-settings-a";
@@ -366,7 +373,97 @@ async fn settings_updates_preserve_turn_identity_and_target(target: SettingsTarg
.collect::<Vec<_>>(),
expected_switch_counts
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn new_context_after_model_switch_uses_captured_extension_window() -> Result<()> {
struct WindowContributor;
impl ContextContributor for WindowContributor {
fn contribute_turn_context<'a>(
&'a self,
input: TurnContextContributionInput<'a>,
) -> ExtensionFuture<'a, Vec<PromptFragment>> {
Box::pin(async move {
input
.model_context_window
.map(|window| {
PromptFragment::developer_policy(
format!("Extension window: {window} tokens."),
ContentItemKind("test.turn_context".to_string()),
)
})
.into_iter()
.collect()
})
}
}
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
vec![
paused_response("resp-a", "pause-before-new-context"),
sse(vec![
ev_response_created("resp-b"),
ev_function_call("new-window", "new_context", "{}"),
ev_completed("resp-b"),
]),
sse_completed("resp-new-window"),
],
)
.await;
let mut extensions = ExtensionRegistryBuilder::new();
extensions.prompt_contributor(Arc::new(WindowContributor));
let test = step_settings_test()
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
config.features.enable(Feature::TokenBudget).unwrap();
config.model_context_window = None;
for model in &mut config.model_catalog.as_mut().unwrap().models {
model.context_window = None;
model.max_context_window = Some(if model.slug == MODEL_A {
128_000
} else {
256_000
});
model.effective_context_window_percent =
if model.slug == MODEL_A { 75 } else { 50 };
}
})
.build_with_auto_env(&server)
.await?;
let paused = start_paused_turn(&test.codex).await?;
assert_eq!(
submit_turn_settings(
&test.codex,
&paused.turn_id,
TurnSettingsUpdate {
model: Some(MODEL_B.to_string()),
..Default::default()
},
)
.await?,
TurnSettingsUpdateOutcome::Applied
);
answer_paused_turn(&test.codex, &paused.turn_id).await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let requests = responses.requests();
assert_eq!(
requests
.iter()
.map(|request| request.body_json()["model"].clone())
.collect::<Vec<_>>(),
vec![json!(MODEL_A), json!(MODEL_B), json!(MODEL_B)]
);
assert!(requests[0].body_contains_text("Extension window: 96000 tokens."));
assert!(requests[2].body_contains_text("Extension window: 128000 tokens."));
assert!(!requests[2].body_contains_text("Extension window: 96000 tokens."));
Ok(())
}
@@ -15,6 +15,6 @@ pub struct TurnContextContributionInput<'a> {
pub thread_store: &'a ExtensionData,
/// Store scoped to this turn.
pub turn_store: &'a ExtensionData,
/// Effective model context window for this turn, when known.
/// Usable context window of the captured model for this context build, when known.
pub model_context_window: Option<i64>,
}
@@ -3,6 +3,7 @@ use std::sync::Arc;
use codex_exec_server_protocol::ExecutorCapabilityDiscoverySnapshot;
use codex_protocol::ThreadId;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::protocol::TurnEnvironmentSelection;
use serde_json::Value;
@@ -13,12 +14,14 @@ use crate::ExtensionMetrics;
pub struct WorldStateContributionInput<'a> {
pub thread_id: ThreadId,
pub turn_id: &'a str,
/// Resolved model metadata captured for this sampling step, retained across discovery.
pub model_info: &'a ModelInfo,
pub environments: &'a [TurnEnvironmentSelection],
/// Selected roots whose stable environments are ready in this sampling step.
pub ready_selected_capability_roots: &'a [SelectedCapabilityRoot],
/// Executor-materialized capability files shared by all consumers in this exact step.
pub executor_capability_discovery: Option<&'a ExecutorCapabilityDiscoverySnapshot>,
/// Metrics bound to the effective model for this turn.
/// Metrics bound to the captured model for this sampling step.
pub extension_metrics: Option<Arc<dyn ExtensionMetrics>>,
pub session_store: &'a ExtensionData,
pub thread_store: &'a ExtensionData,
@@ -6,7 +6,6 @@ use codex_extension_api::ExtensionWarning;
use codex_extension_api::SelectedPluginSnapshot;
use codex_extension_api::WorldStateContributionInput;
use codex_extension_api::WorldStateSectionContribution;
use codex_protocol::openai_models::ModelInfo;
use crate::HostSkillsSnapshot;
use crate::SkillsExtensionConfig;
@@ -102,13 +101,8 @@ impl<'a> CatalogContext<'a> {
) -> Option<Self> {
let thread_state = input.thread_store.get::<SkillsThreadState>()?;
let config = thread_state.config();
let model_info = input.thread_store.get::<ModelInfo>();
let include_usage = model_info
.as_deref()
.is_some_and(|model_info| model_info.include_skills_usage_instructions);
let context_window = model_info
.as_deref()
.and_then(ModelInfo::resolved_context_window);
let include_usage = input.model_info.include_skills_usage_instructions;
let context_window = input.model_info.resolved_context_window();
let metadata_budget = skill_metadata_budget(context_window, config.max_context_tokens);
let emitted_warnings = input
.turn_store
+29 -3
View File
@@ -40,6 +40,7 @@ use codex_otel::THREAD_SKILLS_KEPT_TOTAL_METRIC;
use codex_otel::THREAD_SKILLS_TRUNCATED_METRIC;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::protocol::EnvironmentConfigState;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG;
@@ -92,6 +93,15 @@ const SKILLS_INTRO_WITH_ABSOLUTE_PATHS: &str = "A skill is a set of instructions
const DEMO_SKILL_CONTENTS: &str =
"---\nname: demo\ndescription: Demo skill.\n---\n# Demo\n\nUse the demo skill.\n";
fn catalog_model_info() -> ModelInfo {
ModelInfo {
context_window: None,
max_context_window: None,
include_skills_usage_instructions: false,
..model_info_from_slug("test-model")
}
}
fn world_state_section<'a>(
sections: &'a [WorldStateSectionContribution],
id: &str,
@@ -118,6 +128,7 @@ async fn skill_world_state_fragments(
}];
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id,
environments: &[],
@@ -272,6 +283,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -298,6 +310,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -334,6 +347,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes
turn_store.insert(HostSkillsSnapshot::new(Arc::new(outcome)));
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -400,11 +414,11 @@ async fn persisted_host_snapshot_deduplicates_warning_after_reinitialization() -
.await;
let mut model_info = model_info_from_slug("test-model");
model_info.context_window = Some(50);
thread_store.insert(model_info);
let turn_store = ExtensionData::new("turn-1");
turn_store.insert(host_snapshot.clone());
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &model_info,
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -441,11 +455,11 @@ async fn persisted_host_snapshot_deduplicates_warning_after_reinitialization() -
.await;
let mut model_info = model_info_from_slug("test-model");
model_info.context_window = Some(50);
resumed_thread_store.insert(model_info);
let resumed_turn_store = ExtensionData::new("turn-2");
resumed_turn_store.insert(host_snapshot);
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &model_info,
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-2",
environments: &[],
@@ -534,6 +548,7 @@ async fn executor_orchestrator_and_host_share_catalog_world_state_flow() -> Test
let metrics = Arc::new(RecordingMetrics::default());
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -641,6 +656,7 @@ async fn nonempty_executor_empty_host_records_catalog_metrics() -> TestResult {
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -718,6 +734,7 @@ async fn host_world_state_uses_provider_catalog_with_core_compatible_rendering()
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -797,6 +814,7 @@ async fn shadow_selection_uses_host_catalog_when_instructions_are_disabled() ->
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -1054,6 +1072,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
};
let available_sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: std::slice::from_ref(&turn_environment),
@@ -1109,6 +1128,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
let unavailable_turn_store = ExtensionData::new("turn-2");
let unavailable_sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-2",
environments: &[],
@@ -1133,6 +1153,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
let restored_turn_store = ExtensionData::new("turn-3");
let restored_sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-3",
environments: &[turn_environment],
@@ -1176,6 +1197,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
] {
registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id,
environments: &[],
@@ -1201,6 +1223,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach
let listing_disabled_turn_store = ExtensionData::new("turn-4");
let listing_disabled_sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-4",
environments: &[],
@@ -1972,6 +1995,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te
let turn_store = ExtensionData::new("turn-1");
registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[TurnEnvironmentSelection {
@@ -2083,7 +2107,6 @@ async fn model_context_window_scales_executor_and_orchestrator_catalogs() -> Tes
.await;
let mut model_info = model_info_from_slug("test-model");
model_info.context_window = Some(10_000);
thread_store.insert(model_info);
let thread_fragments = registry.context_contributors()[0]
.contribute_thread_context(&session_store, &thread_store)
@@ -2100,6 +2123,7 @@ async fn model_context_window_scales_executor_and_orchestrator_catalogs() -> Tes
let turn_store = ExtensionData::new("turn-1");
let sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &model_info,
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -2114,6 +2138,7 @@ async fn model_context_window_scales_executor_and_orchestrator_catalogs() -> Tes
// Core rebuilds world state before each sampling step.
let _repeated_sections = registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &model_info,
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],
@@ -2213,6 +2238,7 @@ async fn executor_catalog_emits_at_most_four_warnings() -> TestResult {
registry.context_contributors()[0]
.contribute_world_state(WorldStateContributionInput {
model_info: &catalog_model_info(),
thread_id: codex_protocol::ThreadId::new(),
turn_id: "turn-1",
environments: &[],