From 8228e9b867251f544a5e0c6c80bb5ebc9d5446a1 Mon Sep 17 00:00:00 2001 From: Abhinav Date: Thu, 27 Aug 2026 14:50:48 +0000 Subject: [PATCH] Resolve token budgets from each step's active model (#41162) ## Why Model settings can change between steps in the same turn. Token-budget defaults, context-window limits, and model-visible guidance need to stay aligned with the model captured for each step. ## What changed - Preserve the turn's original token-budget preferences and resolve model-owned defaults when each `StepContext` is created. - Use the step's model to determine context-window availability and token-budget guidance, while keeping explicit user settings unchanged. - Emit one-time replacement or removal notices when context-window guidance changes, including compatibility with previously persisted world-state snapshots. - Add `ModelInfo::usable_context_window()` to distinguish reserved-headroom capacity from the resolved context window and auto-compaction limit. ## Testing - Cover mid-turn model switches with model defaults, explicit settings, config reloads, missing context windows, and missing destination guidance. - Cover unchanged, replaced, removed, blank, and legacy guidance transitions. GitOrigin-RevId: 91100a9f3625bef8a79ecd158eb028051c309eee --- .../world_state/context_window_guidance.rs | 32 ++- .../context_window_guidance_tests.rs | 97 ++++++--- codex-rs/core/src/session/mod.rs | 6 + codex-rs/core/src/session/review.rs | 3 + .../core/src/session/step_activation_tests.rs | 62 +++++- codex-rs/core/src/session/step_context.rs | 3 + codex-rs/core/src/session/tests.rs | 5 + codex-rs/core/src/session/token_budget.rs | 22 ++- codex-rs/core/src/session/turn_context.rs | 28 ++- codex-rs/core/src/session/world_state.rs | 25 +-- codex-rs/core/tests/suite/step_settings.rs | 186 ++++++++++++++++++ codex-rs/core/tests/suite/token_budget.rs | 2 +- codex-rs/protocol/src/openai_models.rs | 28 +++ 13 files changed, 434 insertions(+), 65 deletions(-) diff --git a/codex-rs/core/src/context/world_state/context_window_guidance.rs b/codex-rs/core/src/context/world_state/context_window_guidance.rs index 6228a06b7f..05a289fdd4 100644 --- a/codex-rs/core/src/context/world_state/context_window_guidance.rs +++ b/codex-rs/core/src/context/world_state/context_window_guidance.rs @@ -3,16 +3,24 @@ use super::WorldStateSection; use crate::context::ContextWindowGuidance; use crate::context::ContextualUserFragment; +const REPLACEMENT_NOTICE: &str = + "This context-window guidance replaces all previously provided context-window guidance."; +const REMOVAL_NOTICE: &str = "The previously provided context-window guidance no longer applies."; + /// Model-visible guidance for managing the current context window. #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct ContextWindowGuidanceState { + // Empty string means no guidance. message: String, } impl ContextWindowGuidanceState { - pub(crate) fn new(message: &str) -> Self { + pub(crate) fn new(message: Option<&str>) -> Self { Self { - message: message.to_string(), + message: message + .filter(|message| !message.trim().is_empty()) + .unwrap_or_default() + .to_string(), } } } @@ -41,11 +49,25 @@ impl WorldStateSection for ContextWindowGuidanceState { &self, previous: PreviousSectionState<'_, Self::Snapshot>, ) -> Option> { - if matches!(previous, PreviousSectionState::Known(message) if message == &self.message) { + if matches!(previous, PreviousSectionState::Known(previous) if previous == &self.message) { return None; } - - Some(Box::new(ContextWindowGuidance::new(&self.message))) + let previous_may_contain_guidance = match previous { + PreviousSectionState::Known(previous) => !previous.is_empty(), + PreviousSectionState::Unknown => true, + PreviousSectionState::Absent => false, + }; + let message = if self.message.is_empty() { + if !previous_may_contain_guidance { + return None; + } + REMOVAL_NOTICE.to_string() + } else if previous_may_contain_guidance { + format!("{REPLACEMENT_NOTICE}\n\n{}", self.message) + } else { + self.message.clone() + }; + Some(Box::new(ContextWindowGuidance::new(&message))) } } diff --git a/codex-rs/core/src/context/world_state/context_window_guidance_tests.rs b/codex-rs/core/src/context/world_state/context_window_guidance_tests.rs index 67c960de79..71bef5264f 100644 --- a/codex-rs/core/src/context/world_state/context_window_guidance_tests.rs +++ b/codex-rs/core/src/context/world_state/context_window_guidance_tests.rs @@ -1,52 +1,95 @@ use super::ContextWindowGuidanceState; +use super::REMOVAL_NOTICE; use crate::context::ContextWindowGuidance; use crate::context::ContextualUserFragment; use crate::context::world_state::WorldState; +use crate::context::world_state::WorldStateSnapshot; use codex_protocol::models::ResponseItem; use pretty_assertions::assert_eq; +use serde_json::json; +use test_case::test_case; -#[test] -fn changed_guidance_is_rendered_once() { +#[test_case(None, None, None; "absent guidance stays silent")] +#[test_case(None, Some("guidance"), Some("guidance"); "initial guidance")] +#[test_case(Some("guidance"), Some("guidance"), None; "unchanged guidance")] +#[test_case( + Some("original guidance"), + Some("refreshed guidance"), + Some("This context-window guidance replaces all previously provided context-window guidance.\n\nrefreshed guidance"); + "changed guidance" +)] +#[test_case(Some("guidance"), None, Some(REMOVAL_NOTICE); "removed guidance")] +#[test_case(Some("guidance"), Some(" "), Some(REMOVAL_NOTICE); "blank guidance")] +fn guidance_transitions_render_once( + previous: Option<&str>, + current: Option<&str>, + expected: Option<&str>, +) { let mut original = WorldState::default(); - original.add_section(ContextWindowGuidanceState::new("original guidance")); - let snapshot = original.snapshot(); - - let mut unchanged = WorldState::default(); - unchanged.add_section(ContextWindowGuidanceState::new("original guidance")); - assert!(unchanged.render_diff(&snapshot).is_empty()); + original.add_section(ContextWindowGuidanceState::new(previous)); let mut refreshed = WorldState::default(); - refreshed.add_section(ContextWindowGuidanceState::new("refreshed guidance")); - let fragments = refreshed.render_diff(&snapshot); - - assert_eq!(fragments.len(), 1); + refreshed.add_section(ContextWindowGuidanceState::new(current)); + let fragments = refreshed.render_diff(&original.snapshot()); assert_eq!( - fragments[0].render(), - ContextWindowGuidance::new("refreshed guidance").render() + fragments + .iter() + .map(|fragment| fragment.render()) + .collect::>(), + expected + .map(|message| ContextWindowGuidance::new(message).render()) + .into_iter() + .collect::>() ); + + // Empty guidance must survive persistence as a known state, not a deleted section. + let snapshot: WorldStateSnapshot = + serde_json::from_value(serde_json::to_value(refreshed.snapshot()).unwrap()).unwrap(); + assert!(refreshed.render_diff(&snapshot).is_empty()); } -#[test] -fn legacy_guidance_is_reconciled_once() { +#[test_case( + false, + Some("current guidance"), + Some("This context-window guidance replaces all previously provided context-window guidance.\n\ncurrent guidance"); + "legacy history is replaced" +)] +#[test_case(false, None, Some(REMOVAL_NOTICE); "legacy history is cleared")] +#[test_case(true, None, Some(REMOVAL_NOTICE); "legacy string snapshot is cleared")] +#[test_case(true, Some("previous guidance"), None; "existing string snapshot is unchanged")] +fn legacy_guidance_is_reconciled_once( + has_snapshot: bool, + current: Option<&str>, + expected: Option<&str>, +) { let mut state = WorldState::default(); - state.add_section(ContextWindowGuidanceState::new("current guidance")); + state.add_section(ContextWindowGuidanceState::new(current)); + let snapshot: Option = has_snapshot.then(|| { + serde_json::from_value(json!({"context_window_guidance": "previous guidance"})).unwrap() + }); let retained: ResponseItem = ContextualUserFragment::into(ContextWindowGuidance::new("previous guidance")); - - let fragments = - state.render_history_diff(/*previous*/ None, std::slice::from_ref(&retained)); - assert_eq!(fragments.len(), 1); + let fragments = state.render_history_diff(snapshot.as_ref(), std::slice::from_ref(&retained)); assert_eq!( - fragments[0].render(), - ContextWindowGuidance::new("current guidance").render() + fragments + .iter() + .map(|fragment| fragment.render()) + .collect::>(), + expected + .map(|message| ContextWindowGuidance::new(message).render()) + .into_iter() + .collect::>() ); - let snapshot = state.snapshot(); - let reconciled: ResponseItem = - ContextualUserFragment::into(ContextWindowGuidance::new("current guidance")); + let mut history = vec![retained]; + if let Some(message) = expected { + history.push(ContextualUserFragment::into(ContextWindowGuidance::new( + message, + ))); + } assert!( state - .render_history_diff(Some(&snapshot), &[retained, reconciled]) + .render_history_diff(Some(&state.snapshot()), &history) .is_empty() ); } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1d7fe51c1d..7a87eb97b6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3319,6 +3319,11 @@ impl Session { // Capture once before asynchronous planning; all request consumers // retain this immutable settings version even if the turn is updated. let settings = turn_context.current_settings.load_full(); + let token_budget = token_budget::resolve_token_budget( + turn_context.configured_token_budget.as_ref(), + turn_context.use_model_token_budget_defaults, + settings.model_info.as_ref(), + ); let session_telemetry = settings.telemetry(&turn_context.session_telemetry); // Keep selections fixed for the turn while allowing their startup work to finish. let environments = turn_context.environments.refresh_readiness(); @@ -3402,6 +3407,7 @@ impl Session { .await??; Ok(Arc::new(StepContext { settings, + token_budget, session_telemetry, turn: turn_context, environments, diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 0cff774179..b2be830235 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -144,6 +144,9 @@ pub(super) async fn spawn_review_thread( trace_id: current_span_trace_id(), realtime_active: parent_turn_context.realtime_active, code_mode_available: parent_turn_context.code_mode_available, + configured_token_budget: per_turn_config.token_budget.clone(), + use_model_token_budget_defaults: per_turn_config.features.enabled(Feature::TokenBudget) + && !super::token_budget::has_explicit_settings(&per_turn_config), config: per_turn_config, auth_manager: auth_manager_for_context, initial_settings: Arc::clone(&step_settings), diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index 0e0bcd8c52..f0f8ed0d77 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -33,6 +33,7 @@ use codex_protocol::config_types::ServiceTier; use codex_protocol::openai_models::AutoReviewMessages; use codex_protocol::openai_models::GuardianV2ModelConfig; use codex_protocol::openai_models::GuardianV2TranscriptModelConfig; +use codex_protocol::openai_models::ModelTokenBudgetConfig; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; @@ -204,7 +205,11 @@ async fn activation_fixture(models: Vec) -> ActivationFixture { let mutable = Arc::get_mut(&mut session).expect("unshared test session"); let (models, lookup) = GatedModelsManager::new(models); mutable.services.models_manager = models; - for feature in [Feature::StepModelSwitching, Feature::FastMode] { + for feature in [ + Feature::StepModelSwitching, + Feature::FastMode, + Feature::TokenBudget, + ] { mutable .features .enable(feature) @@ -296,6 +301,27 @@ fn settings_submission( #[tokio::test] async fn submitted_sparse_updates_preserve_captured_steps_and_ordering() { let mut models = activation_models(); + for model in &mut models { + model + .model_messages + .as_mut() + .expect("model messages") + .token_budget = Some(ModelTokenBudgetConfig { + reminder_threshold_tokens: 2_000, + reminder_message_template: "{n_remaining} tokens remain.".to_string(), + guidance_message: format!("Guidance for {}.", model.slug), + auto_compact_fallback_prompt: "Save state before rollover.".to_string(), + auto_compact_fallback_buffer_tokens: 4_000, + }); + } + let initial_model = models + .iter_mut() + .find(|model| model.slug == MODEL_A) + .expect("initial model"); + initial_model.context_window = Some(272_000); + initial_model.max_context_window = Some(272_000); + initial_model.auto_compact_token_limit = None; + initial_model.effective_context_window_percent = 95; let destination = models .iter_mut() .find(|model| model.slug == MODEL_B) @@ -303,6 +329,7 @@ async fn submitted_sparse_updates_preserve_captured_steps_and_ordering() { destination.context_window = Some(190_000); destination.max_context_window = Some(190_000); destination.auto_compact_token_limit = Some(150_000); + destination.effective_context_window_percent = 80; destination.default_reasoning_summary = ReasoningSummary::Detailed; destination .model_messages @@ -435,9 +462,36 @@ async fn submitted_sparse_updates_preserve_captured_steps_and_ordering() { ); assert!(Arc::ptr_eq(&before.turn, &after.turn)); assert_eq!(after.settings.model_info.as_ref(), &expected_destination); - assert_ne!( - before.settings.model_info.context_window, - after.settings.model_info.context_window + let initial_budget = turn + .config + .token_budget + .clone() + .expect("initial model budget"); + let destination_budget = crate::config::TokenBudgetConfig { + guidance_message: Some(format!("Guidance for {MODEL_B}.")), + ..initial_budget.clone() + }; + assert_eq!( + [&before, &during, &after].map(|step| step.token_budget.clone()), + [ + Some(initial_budget.clone()), + Some(initial_budget), + Some(destination_budget), + ] + ); + assert_eq!( + [&before, &during, &after].map(|step| { + ( + step.settings.model_info.resolved_context_window(), + step.settings.model_info.usable_context_window(), + step.settings.model_info.auto_compact_token_limit(), + ) + }), + [ + (Some(272_000), Some(258_400), Some(244_800)), + (Some(272_000), Some(258_400), Some(244_800)), + (Some(190_000), Some(152_000), Some(150_000)), + ] ); assert_eq!(desired_step_settings(&session).await, desired); assert!(Arc::ptr_eq(&before.settings, &during.settings)); diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 41c65c5c31..ff91aad3c9 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; +use crate::config::TokenBudgetConfig; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::session::step_settings::ResolvedStepSettings; use crate::session::turn_context::TurnContext; @@ -15,6 +16,8 @@ pub(crate) struct StepContext { pub(crate) turn: Arc, /// One immutable settings version captured before request preparation. pub(crate) settings: Arc, + /// Frozen turn preferences resolved against this step's captured model. + pub(crate) token_budget: Option, /// Telemetry context tagged with this sampling request's model. pub(crate) session_telemetry: SessionTelemetry, pub(crate) environments: TurnEnvironmentSnapshot, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e9b3227aee..ecfe7b39e9 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -240,6 +240,11 @@ impl StepContext { }); settings.service_tier = turn.config.service_tier.clone(); Arc::new(Self { + token_budget: token_budget::resolve_token_budget( + turn.configured_token_budget.as_ref(), + turn.use_model_token_budget_defaults, + settings.model_info.as_ref(), + ), settings: Arc::new(settings), session_telemetry: turn.session_telemetry.clone(), turn: Arc::clone(&turn), diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 747d8cf5bc..c2fe280f97 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -6,6 +6,7 @@ use crate::context::ContextualUserFragment; use codex_features::Feature; use codex_protocol::openai_models::ModelInfo; +/// Detects explicit preferences before model defaults are applied to the turn config. pub(super) fn has_explicit_settings(config: &Config) -> bool { config .config_layer_stack @@ -25,9 +26,14 @@ pub(super) fn has_explicit_settings(config: &Config) -> bool { }) } -pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) { - if !config.features.enabled(Feature::TokenBudget) || has_explicit_settings(config) { - return; +/// Resolves user-configured token-budget preferences against the current model's defaults. +pub(super) fn resolve_token_budget( + configured_token_budget: Option<&TokenBudgetConfig>, + use_model_defaults: bool, + model_info: &ModelInfo, +) -> Option { + if !use_model_defaults { + return configured_token_budget.cloned(); } let Some(model_defaults) = model_info @@ -35,13 +41,11 @@ pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) .as_ref() .and_then(|messages| messages.token_budget.as_ref()) else { - return; + return configured_token_budget.cloned(); }; let token_budget = TokenBudgetConfig { - use_history_notes_extension: config - .token_budget - .as_ref() + use_history_notes_extension: configured_token_budget .is_some_and(|token_budget| token_budget.use_history_notes_extension), reminder_threshold_tokens: Some(model_defaults.reminder_threshold_tokens), reminder_message_template: model_defaults.reminder_message_template.clone(), @@ -58,10 +62,10 @@ pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) %error, "ignoring invalid model-owned token-budget defaults" ); - return; + return configured_token_budget.cloned(); } - config.token_budget = Some(token_budget); + Some(token_budget) } pub(super) async fn maybe_record( diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 04d04a7bcb..9ce1429f7d 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1,5 +1,8 @@ use super::step_settings::ResolvedStepSettings; +use super::token_budget::has_explicit_settings; +use super::token_budget::resolve_token_budget; use super::*; +use crate::config::TokenBudgetConfig; use crate::environment_selection::EnvironmentConfigOrigin; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::exec_policy::AllowPrefixRules; @@ -187,6 +190,10 @@ pub struct TurnContext { /// Turn-scoped configuration. Read step-specific settings such as service tier and /// approvals reviewer from the corresponding `StepContext` instead. pub config: Arc, + /// Preferences captured before token-budget defaults from the turn's initial model. + pub(crate) configured_token_budget: Option, + /// Captured once so later steps do not re-read config layers to detect user preferences. + pub(crate) use_model_token_budget_defaults: bool, pub(crate) auth_manager: Option>, /// Frozen settings used to construct this context. Legacy turn consumers /// keep this view even when later steps use different settings. @@ -425,12 +432,7 @@ impl TurnContext { /// Legacy: returns the frozen initial-turn model context window. /// Step-scoped consumers should use their captured `StepContext::settings`. pub(crate) fn model_context_window(&self) -> Option { - let effective_context_window_percent = self.model_info().effective_context_window_percent; - self.model_info() - .resolved_context_window() - .map(|context_window| { - context_window.saturating_mul(effective_context_window_percent) / 100 - }) + self.model_info().usable_context_window() } pub(crate) fn apps_enabled(&self) -> bool { @@ -503,6 +505,8 @@ impl TurnContext { realtime_active: self.realtime_active, code_mode_available: self.code_mode_available, config: Arc::new(config), + configured_token_budget: self.configured_token_budget.clone(), + use_model_token_budget_defaults: self.use_model_token_budget_defaults, auth_manager: self.auth_manager.clone(), initial_settings: Arc::clone(&step_settings), current_settings: ArcSwap::from(step_settings), @@ -725,7 +729,15 @@ impl Session { ); let mut per_turn_config = per_turn_config; - super::token_budget::apply_model_defaults(&mut per_turn_config, model_info); + let configured_token_budget = per_turn_config.token_budget.clone(); + let use_model_token_budget_defaults = + per_turn_config.features.enabled(Feature::TokenBudget) + && !has_explicit_settings(&per_turn_config); + per_turn_config.token_budget = resolve_token_budget( + configured_token_budget.as_ref(), + use_model_token_budget_defaults, + model_info, + ); if step_settings.reasoning_effort() == Some(&ReasoningEffort::Persistent) { super::time_reminder::apply_persistent_defaults(&mut per_turn_config); } @@ -764,6 +776,8 @@ impl Session { realtime_active: false, code_mode_available: true, config: per_turn_config, + configured_token_budget, + use_model_token_budget_defaults, auth_manager, initial_settings: Arc::clone(&step_settings), current_settings: ArcSwap::from(step_settings), diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index ec1a382a78..219a0c5345 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -105,9 +105,13 @@ impl Session { personality_is_baked, )); } - if turn_context.config.features.enabled(Feature::TokenBudget) - && turn_context.model_context_window().is_some() - { + let token_budget_enabled = turn_context.config.features.enabled(Feature::TokenBudget) + && step_context + .settings + .model_info + .resolved_context_window() + .is_some(); + if token_budget_enabled { let window_ids = self.state.lock().await.auto_compact_window_ids(); world_state.add_section(TokenBudgetContext::new( turn_context @@ -119,16 +123,13 @@ impl Session { window_ids.window_id, /*thread_hint*/ None, )); - if let Some(guidance) = turn_context - .config - .token_budget - .as_ref() - .and_then(|config| config.guidance_message.as_deref()) - .filter(|message| !message.trim().is_empty()) - { - world_state.add_section(ContextWindowGuidanceState::new(guidance)); - } } + let guidance = step_context + .token_budget + .as_ref() + .and_then(|config| config.guidance_message.as_deref()) + .filter(|_| token_budget_enabled); + world_state.add_section(ContextWindowGuidanceState::new(guidance)); let realtime_mode_instructions = self.conversation.mode_instructions().await; world_state.add_section(RealtimeState::new( turn_context.realtime_active, diff --git a/codex-rs/core/tests/suite/step_settings.rs b/codex-rs/core/tests/suite/step_settings.rs index 62fb0071b1..845812f5ab 100644 --- a/codex-rs/core/tests/suite/step_settings.rs +++ b/codex-rs/core/tests/suite/step_settings.rs @@ -3,6 +3,7 @@ use codex_config::McpServerConfig; use codex_core::CodexThread; use codex_core::TurnInputRequest; use codex_core::config::Constrained; +use codex_core::config::TokenBudgetConfig; use codex_features::Feature; use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; @@ -12,10 +13,13 @@ use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::ServiceTier; use codex_protocol::openai_models::ConfirmationPolicies; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelTokenBudgetConfig; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ReasoningEffortPreset; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; +use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::SafetyBufferingEvent; @@ -341,6 +345,188 @@ async fn settings_updates_preserve_turn_identity_and_target(target: SettingsTarg Ok(()) } +#[derive(Clone, Copy)] +enum TokenBudgetScenario { + ModelDefaults, + ExplicitDefaultTemplate, + ReloadPreferences, + DestinationWindowOnly, + InitialWindowOnly, + DestinationWithoutGuidance, +} + +#[test_case(TokenBudgetScenario::ModelDefaults)] +#[test_case(TokenBudgetScenario::ExplicitDefaultTemplate)] +#[test_case(TokenBudgetScenario::ReloadPreferences)] +#[test_case(TokenBudgetScenario::DestinationWindowOnly)] +#[test_case(TokenBudgetScenario::InitialWindowOnly)] +#[test_case(TokenBudgetScenario::DestinationWithoutGuidance)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn active_model_switch_resolves_token_budget_from_original_preferences( + scenario: TokenBudgetScenario, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let explicit_default_template = + matches!(scenario, TokenBudgetScenario::ExplicitDefaultTemplate); + let reload_user_config = matches!(scenario, TokenBudgetScenario::ReloadPreferences); + let context_window_model = match scenario { + TokenBudgetScenario::DestinationWindowOnly => Some(MODEL_B), + TokenBudgetScenario::InitialWindowOnly => Some(MODEL_A), + TokenBudgetScenario::ModelDefaults + | TokenBudgetScenario::ExplicitDefaultTemplate + | TokenBudgetScenario::ReloadPreferences + | TokenBudgetScenario::DestinationWithoutGuidance => None, + }; + let destination_has_guidance = + !matches!(scenario, TokenBudgetScenario::DestinationWithoutGuidance); + let server = start_mock_server().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + paused_response("resp-1", "pause-before-model-budget-switch"), + paused_response("resp-2", "pause-after-model-budget-switch"), + sse_completed("resp-3"), + ], + ) + .await; + let test = step_settings_test() + .with_pre_build_hook(move |home| { + let config = if explicit_default_template { + let default_template = TokenBudgetConfig::default().reminder_message_template; + format!( + "[features.token_budget]\nenabled = true\nreminder_message_template = {default_template:?}\n" + ) + } else { + "[features.token_budget]\nenabled = true\n".to_string() + }; + std::fs::write(home.join("config.toml"), config) + .expect("write token-budget preferences"); + }) + .with_config(move |config| { + config + .features + .enable(Feature::TokenBudget) + .expect("enable token-budget feature"); + if context_window_model.is_some() { + config.model_context_window = None; + } + for model in &mut config + .model_catalog + .as_mut() + .expect("controlled model catalog") + .models + { + let slug = model.slug.clone(); + let initial_model = slug == MODEL_A; + if let Some(context_window_model) = context_window_model { + model.context_window = (slug == context_window_model).then_some(128_000); + model.max_context_window = None; + } + model + .model_messages + .as_mut() + .expect("model messages") + .token_budget = (initial_model || destination_has_guidance).then(|| ModelTokenBudgetConfig { + reminder_threshold_tokens: if initial_model { 8_000 } else { 2_000 }, + reminder_message_template: format!( + "Reminder for {slug}: {{n_remaining}} tokens remain." + ), + guidance_message: format!("Use {slug} token-budget guidance."), + auto_compact_fallback_prompt: format!("Save {slug} state before rollover."), + auto_compact_fallback_buffer_tokens: if initial_model { + 16_000 + } else { + 4_000 + }, + }); + } + }) + .build_with_auto_env(&server) + .await?; + let request = start_paused_turn(&test.codex).await?; + + if reload_user_config { + std::fs::write( + test.codex_home_path().join("config.toml"), + "[features.token_budget]\nenabled = true\nreminder_message_template = \"Reloaded reminder\"\n", + )?; + test.codex.submit(Op::ReloadUserConfig).await?; + } + + assert_eq!( + submit_turn_settings( + &test.codex, + &request.turn_id, + TurnSettingsUpdate { + model: Some(MODEL_B.to_string()), + ..Default::default() + }, + ) + .await?, + TurnSettingsUpdateOutcome::Applied + ); + answer_paused_turn(&test.codex, &request.turn_id).await?; + let request = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RequestUserInput(request) => Some(request.clone()), + _ => None, + }) + .await; + answer_paused_turn(&test.codex, &request.turn_id).await?; + wait_for_event(&test.codex, |event| match event { + EventMsg::Error(error) => panic!("settings activation failed: {}", error.message), + EventMsg::TurnComplete(_) => true, + _ => false, + }) + .await; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let initial_guidance = format!("Use {MODEL_A} token-budget guidance."); + let initial_guidance_expected = + !explicit_default_template && context_window_model != Some(MODEL_B); + assert_eq!( + requests[0].body_contains_text(&initial_guidance), + initial_guidance_expected + ); + let mut expected_guidance = Vec::new(); + if initial_guidance_expected { + expected_guidance.push(format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{initial_guidance}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + )); + } + if !explicit_default_template + && destination_has_guidance + && context_window_model != Some(MODEL_A) + { + let replacement_notice = if initial_guidance_expected { + "This context-window guidance replaces all previously provided context-window guidance.\n\n" + } else { + "" + }; + expected_guidance.push(format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{replacement_notice}Use {MODEL_B} token-budget guidance.\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + )); + } else if initial_guidance_expected { + expected_guidance.push(format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\nThe previously provided context-window guidance no longer applies.\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + )); + } + for request in &requests[1..] { + let guidance = request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.starts_with(CONTEXT_WINDOW_GUIDANCE_OPEN_TAG)) + .collect::>(); + assert_eq!( + guidance, expected_guidance, + "preserve history and append the guidance transition only once" + ); + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_confirmation_policy_follows_step_model_changes() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 8f6d1fe421..2ca920bb44 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -461,7 +461,7 @@ async fn token_budget_defaults_follow_the_active_model() -> Result<()> { "switching models should preserve the existing model-switch message" ); let expected_guidance = format!( - "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\nUse second-model context-window guidance.\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\nThis context-window guidance replaces all previously provided context-window guidance.\n\nUse second-model context-window guidance.\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" ); assert_eq!( developer_texts diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 65d63723a5..69347fe7f7 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -486,6 +486,13 @@ impl ModelInfo { self.context_window.or(self.max_context_window) } + /// Context available to inference after reserving this model's configured headroom. + pub fn usable_context_window(&self) -> Option { + self.resolved_context_window().map(|context_window| { + context_window.saturating_mul(self.effective_context_window_percent) / 100 + }) + } + pub fn auto_compact_token_limit(&self) -> Option { let context_limit = self .resolved_context_window() @@ -1828,9 +1835,30 @@ mod tests { }; assert_eq!(model.resolved_context_window(), Some(400_000)); + assert_eq!(model.usable_context_window(), Some(380_000)); assert_eq!(model.auto_compact_token_limit(), Some(360_000)); } + #[test] + fn model_context_window_limits_preserve_their_distinct_meanings() { + let model = ModelInfo { + context_window: Some(272_000), + max_context_window: Some(400_000), + auto_compact_token_limit: Some(250_000), + effective_context_window_percent: 95, + ..test_model(/*spec*/ None) + }; + + assert_eq!( + ( + model.resolved_context_window(), + model.usable_context_window(), + model.auto_compact_token_limit(), + ), + (Some(272_000), Some(258_400), Some(244_800)) + ); + } + #[test] fn model_preset_preserves_availability_nux() { let preset = ModelPreset::from(ModelInfo {