From 8ff4aa8ee402a0ab855c203d0436665cc0cc5470 Mon Sep 17 00:00:00 2001 From: Abhinav Date: Wed, 9 Sep 2026 16:33:05 +0000 Subject: [PATCH] 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 --- codex-rs/core/src/session/mod.rs | 4 +- codex-rs/core/src/session/tests.rs | 22 +- codex-rs/core/src/session/world_state.rs | 3 +- codex-rs/core/tests/suite/skills_extension.rs | 371 ++++++++++++++---- codex-rs/core/tests/suite/step_settings.rs | 97 +++++ .../extension-api/src/contributors/context.rs | 2 +- .../src/contributors/world_state.rs | 5 +- .../ext/skills/src/world_state_catalogs.rs | 10 +- codex-rs/ext/skills/tests/skills_extension.rs | 32 +- 9 files changed, 447 insertions(+), 99 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ed0b6981a1..b2189a6e49 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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 { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ca6fa99501..b7d67a461a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -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("")); 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(""), "full-context metadata must use B's window even though the turn started without one" diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 3ed069258f..1c952c51ab 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -290,13 +290,14 @@ impl Session { .map(|root| root.selected_root().clone()) .collect::>(); 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 diff --git a/codex-rs/core/tests/suite/skills_extension.rs b/codex-rs/core/tests/suite/skills_extension.rs index 0d461cb13a..d6a819185c 100644 --- a/codex-rs/core/tests/suite/skills_extension.rs +++ b/codex-rs/core/tests/suite/skills_extension.rs @@ -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, + 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), @@ -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) -> 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::::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::::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::>(); - 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::(); + 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::>(); + 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::>(); + 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::>(); + (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::>(); + 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)] diff --git a/codex-rs/core/tests/suite/step_settings.rs b/codex-rs/core/tests/suite/step_settings.rs index 19ad6daed2..bbfb3353bc 100644 --- a/codex-rs/core/tests/suite/step_settings.rs +++ b/codex-rs/core/tests/suite/step_settings.rs @@ -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::>(), 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> { + 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![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(()) } diff --git a/codex-rs/ext/extension-api/src/contributors/context.rs b/codex-rs/ext/extension-api/src/contributors/context.rs index fb6239ec4b..78dd8ae7b5 100644 --- a/codex-rs/ext/extension-api/src/contributors/context.rs +++ b/codex-rs/ext/extension-api/src/contributors/context.rs @@ -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, } diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index 3aba6b9418..c916fac6d0 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -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>, pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, diff --git a/codex-rs/ext/skills/src/world_state_catalogs.rs b/codex-rs/ext/skills/src/world_state_catalogs.rs index c30266f5ae..c1fab579f3 100644 --- a/codex-rs/ext/skills/src/world_state_catalogs.rs +++ b/codex-rs/ext/skills/src/world_state_catalogs.rs @@ -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 { let thread_state = input.thread_store.get::()?; let config = thread_state.config(); - let model_info = input.thread_store.get::(); - 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 diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index ac9cc37781..b0ee66b8de 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -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: &[],