Support model-owned token budget defaults (#35608)

## What changed

- Add token-budget settings to model catalog messages and apply them when the
  feature is enabled without explicit token-budget configuration.
- Keep explicit user settings authoritative, reject invalid catalog defaults,
  and preserve resolved defaults in exported config locks.
- Manage context-window guidance through world state so it updates once when
  the active model changes while retaining prior conversation history.

## Testing

- Cover catalog defaults, explicit overrides, disabled features, invalid
  values, config-lock replay, and model switching.

GitOrigin-RevId: 54544fefaa14b09f5ef5bad9967a13c52b87b0fe

codex-rs/core/src/config/mod.rs

@@ -1120,6 +1120,78 @@ pub struct TokenBudgetConfig {
 }
 
 impl TokenBudgetConfig {
+    pub(crate) fn validate(&self) -> std::io::Result<()> {
+        if self
+            .reminder_threshold_tokens
+            .is_some_and(|tokens| tokens <= 0)
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "features.token_budget.reminder_threshold_tokens must be positive",
+            ));
+        }
+
+        if self.reminder_message_template.trim().is_empty() {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "features.token_budget.reminder_message_template must not be empty",
+            ));
+        }
+        if self.reminder_message_template.len() > TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "features.token_budget.reminder_message_template must not exceed {TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES} bytes"
+                ),
+            ));
+        }
+
+        if self
+            .guidance_message
+            .as_ref()
+            .is_some_and(|message| message.len() > TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES)
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "features.token_budget.guidance_message must not exceed {TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES} bytes"
+                ),
+            ));
+        }
+
+        if self
+            .auto_compact_fallback_prompt
+            .as_ref()
+            .is_some_and(|prompt| prompt.len() > AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES)
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                format!(
+                    "features.token_budget.auto_compact_fallback_prompt must not exceed {AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES} bytes"
+                ),
+            ));
+        }
+        if self.auto_compact_fallback_prompt.is_some()
+            && self.auto_compact_fallback_buffer_tokens.is_none()
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "features.token_budget.auto_compact_fallback_buffer_tokens is required when auto_compact_fallback_prompt is set",
+            ));
+        }
+        if self
+            .auto_compact_fallback_buffer_tokens
+            .is_some_and(|tokens| tokens <= 0)
+        {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "features.token_budget.auto_compact_fallback_buffer_tokens must be positive",
+            ));
+        }
+
+        Ok(())
+    }
+
     pub(crate) fn fallback_buffer_tokens(&self) -> i64 {
         if self.auto_compact_fallback_prompt.is_some() {
             self.auto_compact_fallback_buffer_tokens.unwrap_or(0)
@@ -2684,85 +2756,29 @@ fn resolve_token_budget_config(
     let token_budget_config = token_budget_toml_config(config_toml.features.as_ref());
     let reminder_threshold_tokens =
         token_budget_config.and_then(|config| config.reminder_threshold_tokens);
-    if reminder_threshold_tokens.is_some_and(|tokens| tokens <= 0) {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            "features.token_budget.reminder_threshold_tokens must be positive",
-        ));
-    }
-
     let reminder_message_template = token_budget_config
         .and_then(|config| config.reminder_message_template.clone())
         .unwrap_or_else(|| DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string());
-    if reminder_message_template.trim().is_empty() {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            "features.token_budget.reminder_message_template must not be empty",
-        ));
-    }
-    if reminder_message_template.len() > TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            format!(
-                "features.token_budget.reminder_message_template must not exceed {TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES} bytes"
-            ),
-        ));
-    }
-
     let guidance_message = token_budget_config
         .and_then(|config| config.guidance_message.clone())
         .filter(|message| !message.trim().is_empty());
-    if guidance_message
-        .as_ref()
-        .is_some_and(|message| message.len() > TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES)
-    {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            format!(
-                "features.token_budget.guidance_message must not exceed {TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES} bytes"
-            ),
-        ));
-    }
-
     let auto_compact_fallback_prompt = token_budget_config
         .and_then(|config| config.auto_compact_fallback_prompt.as_deref())
         .map(str::trim)
         .filter(|value| !value.is_empty())
         .map(str::to_string);
-    if auto_compact_fallback_prompt
-        .as_ref()
-        .is_some_and(|prompt| prompt.len() > AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES)
-    {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            format!(
-                "features.token_budget.auto_compact_fallback_prompt must not exceed {AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES} bytes"
-            ),
-        ));
-    }
-
     let auto_compact_fallback_buffer_tokens =
         token_budget_config.and_then(|config| config.auto_compact_fallback_buffer_tokens);
-    if auto_compact_fallback_prompt.is_some() && auto_compact_fallback_buffer_tokens.is_none() {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            "features.token_budget.auto_compact_fallback_buffer_tokens is required when auto_compact_fallback_prompt is set",
-        ));
-    }
-    if auto_compact_fallback_buffer_tokens.is_some_and(|tokens| tokens <= 0) {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::InvalidInput,
-            "features.token_budget.auto_compact_fallback_buffer_tokens must be positive",
-        ));
-    }
 
-    Ok(Some(TokenBudgetConfig {
+    let token_budget = TokenBudgetConfig {
         reminder_threshold_tokens,
         reminder_message_template,
         guidance_message,
         auto_compact_fallback_prompt,
         auto_compact_fallback_buffer_tokens,
-    }))
+    };
+    token_budget.validate()?;
+    Ok(Some(token_budget))
 }
 
 fn resolve_rollout_budget_config(

codex-rs/core/src/context/world_state/context_window_guidance.rs

@@ -0,0 +1,54 @@
+use super::PreviousSectionState;
+use super::WorldStateSection;
+use crate::context::ContextWindowGuidance;
+use crate::context::ContextualUserFragment;
+
+/// Model-visible guidance for managing the current context window.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) struct ContextWindowGuidanceState {
+    message: String,
+}
+
+impl ContextWindowGuidanceState {
+    pub(crate) fn new(message: &str) -> Self {
+        Self {
+            message: message.to_string(),
+        }
+    }
+}
+
+impl WorldStateSection for ContextWindowGuidanceState {
+    const ID: &'static str = "context_window_guidance";
+    type Snapshot = String;
+
+    fn snapshot(&self) -> Self::Snapshot {
+        self.message.clone()
+    }
+
+    fn matches_legacy_fragment(role: &str, text: &str) -> bool {
+        role == "developer" && ContextWindowGuidance::matches_text(text)
+    }
+
+    fn has_retained_fragment_matcher() -> bool {
+        true
+    }
+
+    fn matches_retained_fragment(role: &str, text: &str) -> bool {
+        Self::matches_legacy_fragment(role, text)
+    }
+
+    fn render_diff(
+        &self,
+        previous: PreviousSectionState<'_, Self::Snapshot>,
+    ) -> Option<Box<dyn ContextualUserFragment>> {
+        if matches!(previous, PreviousSectionState::Known(message) if message == &self.message) {
+            return None;
+        }
+
+        Some(Box::new(ContextWindowGuidance::new(&self.message)))
+    }
+}
+
+#[cfg(test)]
+#[path = "context_window_guidance_tests.rs"]
+mod tests;

codex-rs/core/src/context/world_state/context_window_guidance_tests.rs

@@ -0,0 +1,52 @@
+use super::ContextWindowGuidanceState;
+use crate::context::ContextWindowGuidance;
+use crate::context::ContextualUserFragment;
+use crate::context::world_state::WorldState;
+use codex_protocol::models::ResponseItem;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn changed_guidance_is_rendered_once() {
+    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());
+
+    let mut refreshed = WorldState::default();
+    refreshed.add_section(ContextWindowGuidanceState::new("refreshed guidance"));
+    let fragments = refreshed.render_diff(&snapshot);
+
+    assert_eq!(fragments.len(), 1);
+    assert_eq!(
+        fragments[0].render(),
+        ContextWindowGuidance::new("refreshed guidance").render()
+    );
+}
+
+#[test]
+fn legacy_guidance_is_reconciled_once() {
+    let mut state = WorldState::default();
+    state.add_section(ContextWindowGuidanceState::new("current guidance"));
+    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);
+    assert_eq!(
+        fragments[0].render(),
+        ContextWindowGuidance::new("current guidance").render()
+    );
+
+    let snapshot = state.snapshot();
+    let reconciled: ResponseItem =
+        ContextualUserFragment::into(ContextWindowGuidance::new("current guidance"));
+    assert!(
+        state
+            .render_history_diff(Some(&snapshot), &[retained, reconciled])
+            .is_empty()
+    );
+}

codex-rs/core/src/context/world_state/mod.rs

@@ -1,6 +1,7 @@
 mod agents_md;
 mod apps_instructions;
 mod collaboration_mode;
+mod context_window_guidance;
 mod environment;
 mod environments_instructions;
 mod model;
@@ -32,6 +33,7 @@ use std::fmt;
 pub(crate) use agents_md::AgentsMdState;
 pub(crate) use apps_instructions::AppsInstructionsState;
 pub(crate) use collaboration_mode::CollaborationModeState;
+pub(crate) use context_window_guidance::ContextWindowGuidanceState;
 pub(crate) use environment::EnvironmentsState;
 pub(crate) use environments_instructions::EnvironmentsInstructionsState;
 pub(crate) use model::ModelInstructionsState;

codex-rs/core/src/guardian/review_session.rs

@@ -1453,6 +1453,7 @@ mod tests {
                 policy_template: Some(catalog_template.to_string()),
             }),
             permissions: None,
+            token_budget: None,
         };
 
         let guardian_config = build_guardian_review_session_config(
@@ -1485,6 +1486,7 @@ mod tests {
                 policy_template: None,
             }),
             permissions: None,
+            token_budget: None,
         };
 
         let guardian_config = build_guardian_review_session_config(
@@ -1525,6 +1527,7 @@ mod tests {
                 policy_template: Some(String::new()),
             }),
             permissions: None,
+            token_budget: None,
         };
 
         let guardian_config = build_guardian_review_session_config(

codex-rs/core/src/session/config_lock.rs

@@ -157,7 +157,9 @@ fn save_config_resolved_fields(
         resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?;
     multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2));
     features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2));
-    if let Some(token_budget) = config.token_budget.as_ref() {
+    if let Some(token_budget) = config.token_budget.as_ref()
+        && super::token_budget::has_explicit_settings(config)
+    {
         let mut token_budget: TokenBudgetConfigToml =
             resolved_config_to_toml(token_budget, "features.token_budget")?;
         token_budget.enabled = Some(config.features.enabled(Feature::TokenBudget));
@@ -387,6 +389,29 @@ mod tests {
         assert_eq!(lockfile.version, crate::config_lock::CONFIG_LOCK_VERSION);
     }
 
+    #[tokio::test]
+    async fn lock_preserves_model_owned_token_budget_defaults() {
+        let mut sc = crate::session::tests::make_session_configuration_for_tests().await;
+        let mut config = (*sc.original_config_do_not_use).clone();
+        config.token_budget = Some(crate::config::TokenBudgetConfig::default());
+        config
+            .features
+            .enable(Feature::TokenBudget)
+            .expect("token_budget should be enableable in tests");
+        sc.original_config_do_not_use = Arc::new(config);
+
+        let lockfile = sc.to_config_lockfile_toml().expect("lock should serialize");
+
+        assert_eq!(
+            lockfile
+                .config
+                .features
+                .as_ref()
+                .and_then(|features| features.token_budget.as_ref()),
+            Some(&FeatureToml::Enabled(true))
+        );
+    }
+
     #[tokio::test]
     async fn lock_skips_session_values_when_model_catalog_fields_are_not_saved() {
         let mut sc = crate::session::tests::make_session_configuration_for_tests().await;

codex-rs/core/src/session/mod.rs

@@ -602,7 +602,7 @@ impl Session {
             )
         };
 
-        let config = Arc::new(config);
+        let mut config = Arc::new(config);
         let refresh_strategy = if session_source.is_non_root_agent() {
             codex_models_manager::manager::RefreshStrategy::Offline
         } else {
@@ -645,6 +645,11 @@ impl Session {
         let model_info = models_manager
             .get_model_info(model.as_str(), &config.to_models_manager_config())
             .await;
+        if config.config_lock_export_dir.is_some()
+            && config.config_lock_save_fields_resolved_from_model_catalog
+        {
+            self::token_budget::apply_model_defaults(Arc::make_mut(&mut config), &model_info);
+        }
         let multi_agent_version = config.multi_agent_version_override().or_else(|| {
             resolve_multi_agent_version(&conversation_history, inherited_multi_agent_version)
         });
@@ -3489,16 +3494,6 @@ impl Session {
                 )
                 .render(),
             );
-            if let Some(guidance_message) = turn_context
-                .config
-                .token_budget
-                .as_ref()
-                .and_then(|config| config.guidance_message.as_deref())
-                .filter(|message| !message.trim().is_empty())
-            {
-                developer_sections
-                    .push(crate::context::ContextWindowGuidance::new(guidance_message).render());
-            }
         }
         // Render the active mode after the usage hint so it can override that hint.
         let mut initial_multi_agent_mode = None;

codex-rs/core/src/session/token_budget.rs

@@ -1,7 +1,59 @@
 use super::session::Session;
 use super::turn_context::TurnContext;
+use crate::config::Config;
+use crate::config::TokenBudgetConfig;
 use crate::context::ContextualUserFragment;
 use codex_features::Feature;
+use codex_protocol::openai_models::ModelInfo;
+
+pub(super) fn has_explicit_settings(config: &Config) -> bool {
+    config
+        .config_layer_stack
+        .effective_config()
+        .get("features")
+        .and_then(|features| features.get("token_budget"))
+        .and_then(|token_budget| token_budget.as_table())
+        .is_some_and(|settings| settings.keys().any(|key| key != "enabled"))
+        || config
+            .token_budget
+            .as_ref()
+            .is_some_and(|token_budget| token_budget != &TokenBudgetConfig::default())
+}
+
+pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) {
+    if !config.features.enabled(Feature::TokenBudget) || has_explicit_settings(config) {
+        return;
+    }
+
+    let Some(model_defaults) = model_info
+        .model_messages
+        .as_ref()
+        .and_then(|messages| messages.token_budget.as_ref())
+    else {
+        return;
+    };
+
+    let token_budget = TokenBudgetConfig {
+        reminder_threshold_tokens: Some(model_defaults.reminder_threshold_tokens),
+        reminder_message_template: model_defaults.reminder_message_template.clone(),
+        guidance_message: Some(model_defaults.guidance_message.clone()),
+        auto_compact_fallback_prompt: Some(model_defaults.auto_compact_fallback_prompt.clone()),
+        auto_compact_fallback_buffer_tokens: Some(
+            model_defaults.auto_compact_fallback_buffer_tokens,
+        ),
+    };
+
+    if let Err(error) = token_budget.validate() {
+        tracing::warn!(
+            model = %model_info.slug,
+            %error,
+            "ignoring invalid model-owned token-budget defaults"
+        );
+        return;
+    }
+
+    config.token_budget = Some(token_budget);
+}
 
 pub(super) async fn maybe_record(
     sess: &Session,

codex-rs/core/src/session/turn_context.rs

@@ -536,6 +536,7 @@ impl Session {
         );
 
         let mut per_turn_config = per_turn_config;
+        super::token_budget::apply_model_defaults(&mut per_turn_config, &model_info);
         per_turn_config.service_tier = get_service_tier(
             per_turn_config.service_tier,
             per_turn_config.features.enabled(Feature::FastMode),

codex-rs/core/src/session/world_state.rs

@@ -7,6 +7,7 @@ use crate::context::ApprovalPromptContext;
 use crate::context::world_state::AgentsMdState;
 use crate::context::world_state::AppsInstructionsState;
 use crate::context::world_state::CollaborationModeState;
+use crate::context::world_state::ContextWindowGuidanceState;
 use crate::context::world_state::EnvironmentsInstructionsState;
 use crate::context::world_state::EnvironmentsState;
 use crate::context::world_state::ModelInstructionsState;
@@ -82,6 +83,17 @@ impl Session {
                 personality_is_baked,
             ));
         }
+        if turn_context.config.features.enabled(Feature::TokenBudget)
+            && turn_context.model_context_window().is_some()
+            && 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));
+        }
         world_state.add_section(RealtimeState::new(
             turn_context.realtime_active,
             turn_context