mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
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:
@@ -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
|
||||
{
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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: &[],
|
||||
|
||||
Reference in New Issue
Block a user