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

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,

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),

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<ModelInfo>) -> 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));

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<TurnContext>,
     /// One immutable settings version captured before request preparation.
     pub(crate) settings: Arc<ResolvedStepSettings>,
+    /// Frozen turn preferences resolved against this step's captured model.
+    pub(crate) token_budget: Option<TokenBudgetConfig>,
     /// Telemetry context tagged with this sampling request's model.
     pub(crate) session_telemetry: SessionTelemetry,
     pub(crate) environments: TurnEnvironmentSnapshot,

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),

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<TokenBudgetConfig> {
+    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(

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<Config>,
+    /// Preferences captured before token-budget defaults from the turn's initial model.
+    pub(crate) configured_token_budget: Option<TokenBudgetConfig>,
+    /// 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<Arc<AuthManager>>,
     /// 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<i64> {
-        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),

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,