From d58d0e5841e0de08e251673db2d5af8cf3a1ad51 Mon Sep 17 00:00:00 2001 From: pmccrary-oai Date: Mon, 31 Aug 2026 08:19:39 +0000 Subject: [PATCH] Allow models to enable token budgeting by default (#41803) ## What changed - Add model metadata flags for enabling token budgeting and its history-notes extension. - Apply those activation defaults once at thread startup when the user has not explicitly configured token budgeting, while respecting managed feature requirements. - Continue resolving model-owned prompts for the active model so model switches can use their own token-budget guidance. ## Testing - Update the token-budget integration test to verify that model metadata activates the feature without manual configuration. GitOrigin-RevId: 73a32623eaeabb9243e240c43da7a5e7e93e9a19 --- codex-rs/core/src/session/mod.rs | 4 ++ .../core/src/session/step_activation_tests.rs | 2 + codex-rs/core/src/session/token_budget.rs | 39 +++++++++++++++++++ codex-rs/core/tests/suite/review.rs | 2 + codex-rs/core/tests/suite/step_settings.rs | 2 + codex-rs/core/tests/suite/token_budget.rs | 17 +++----- .../models-manager/src/model_info_tests.rs | 2 + codex-rs/protocol/src/openai_models.rs | 4 ++ 8 files changed, 60 insertions(+), 12 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8f0eb2ad8d..ea0651cc02 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -668,6 +668,10 @@ impl Session { let model_info = models_manager .get_model_info(model.as_str(), &config.to_models_manager_config()) .await; + // Intentionally resolve `enabled` and `use_history_notes_extension` only at + // thread startup. Both activation flags stay fixed for this thread runtime, + // even if the selected model changes later. + token_budget::apply_model_defaults(Arc::make_mut(&mut config), &model_info); let configured_config = Arc::clone(&config); let multi_agent_version = config.multi_agent_version_override().or_else(|| { resolve_multi_agent_version(&conversation_history, inherited_multi_agent_version) diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index a584dd4987..c3ca3b572a 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -277,6 +277,8 @@ async fn submitted_sparse_updates_preserve_captured_steps_and_ordering() { .as_mut() .expect("model messages") .token_budget = Some(ModelTokenBudgetConfig { + enabled: false, + use_history_notes_extension: false, reminder_threshold_tokens: 2_000, reminder_message_template: "{n_remaining} tokens remain.".to_string(), guidance_message: format!("Guidance for {}.", model.slug), diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index c2fe280f97..179db47caf 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -68,6 +68,45 @@ pub(super) fn resolve_token_budget( Some(token_budget) } +/// Applies model activation defaults before thread extensions are initialized. +pub(super) fn apply_model_defaults(config: &mut Config, model_info: &ModelInfo) { + let Some(model_defaults) = model_info + .model_messages + .as_ref() + .and_then(|messages| messages.token_budget.as_ref()) + else { + return; + }; + if !model_defaults.enabled { + return; + } + + let has_explicit_config = config.token_budget.is_some() + || config + .config_layer_stack + .effective_config() + .get("features") + .and_then(|features| features.get("token_budget")) + .is_some(); + if has_explicit_config { + return; + } + + if config.features.enable(Feature::TokenBudget).is_err() { + return; + } + // Managed requirements can pin the feature off even when enable() succeeds. + if !config.features.enabled(Feature::TokenBudget) { + return; + } + + // Keep prompts unresolved so later turns and model switches use their own defaults. + config.token_budget = Some(TokenBudgetConfig { + use_history_notes_extension: model_defaults.use_history_notes_extension, + ..TokenBudgetConfig::default() + }); +} + pub(super) async fn maybe_record( sess: &Session, turn_context: &TurnContext, diff --git a/codex-rs/core/tests/suite/review.rs b/codex-rs/core/tests/suite/review.rs index b4c9064880..830fd1c512 100644 --- a/codex-rs/core/tests/suite/review.rs +++ b/codex-rs/core/tests/suite/review.rs @@ -582,6 +582,8 @@ async fn review_uses_updated_turn_permissions_and_approval_policy() { fn model_defaults(guidance_message: &str) -> ModelTokenBudgetConfig { ModelTokenBudgetConfig { + enabled: false, + use_history_notes_extension: false, reminder_threshold_tokens: 6_144, reminder_message_template: "Reminder: {n_remaining} tokens remain.".to_string(), guidance_message: guidance_message.to_string(), diff --git a/codex-rs/core/tests/suite/step_settings.rs b/codex-rs/core/tests/suite/step_settings.rs index a347b6c85f..bf5651d64c 100644 --- a/codex-rs/core/tests/suite/step_settings.rs +++ b/codex-rs/core/tests/suite/step_settings.rs @@ -430,6 +430,8 @@ async fn active_model_switch_resolves_token_budget_from_original_preferences( .as_mut() .expect("model messages") .token_budget = (initial_model || destination_has_guidance).then(|| ModelTokenBudgetConfig { + enabled: false, + use_history_notes_extension: false, reminder_threshold_tokens: if initial_model { 8_000 } else { 2_000 }, reminder_message_template: format!( "Reminder for {slug}: {{n_remaining}} tokens remain." diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index e05ea06138..e4d1624ada 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -63,6 +63,8 @@ const AUTO_COMPACT_FALLBACK_PROMPT: &str = "Save the important state before roll fn model_token_budget_config() -> ModelTokenBudgetConfig { ModelTokenBudgetConfig { + enabled: false, + use_history_notes_extension: false, reminder_threshold_tokens: 6_144, reminder_message_template: "Model reminder: {n_remaining} tokens remain.".to_string(), guidance_message: "Use the model-owned context-window guidance.".to_string(), @@ -295,16 +297,11 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> { let server = start_mock_server().await; let response = mount_sse_once(&server, sse_completed("resp-1")).await; - let model_defaults = model_token_budget_config(); + let mut model_defaults = model_token_budget_config(); + model_defaults.enabled = true; + model_defaults.use_history_notes_extension = true; let expected_guidance = model_defaults.guidance_message.clone(); let test = test_codex() - .with_pre_build_hook(|home| { - std::fs::write( - home.join("config.toml"), - "[features.token_budget]\nenabled = true\nuse_history_notes_extension = true\n", - ) - .expect("write token-budget configuration"); - }) .with_model_info_override("gpt-5.2", move |model_info| { model_info .model_messages @@ -314,10 +311,6 @@ async fn token_budget_uses_model_message_defaults() -> Result<()> { }) .with_config(|config| { config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); - config - .features - .enable(Feature::TokenBudget) - .expect("test config should allow token budget"); }) .build_with_auto_env(&server) .await?; diff --git a/codex-rs/models-manager/src/model_info_tests.rs b/codex-rs/models-manager/src/model_info_tests.rs index d6f3a73bf4..921dd58189 100644 --- a/codex-rs/models-manager/src/model_info_tests.rs +++ b/codex-rs/models-manager/src/model_info_tests.rs @@ -62,6 +62,8 @@ fn base_instruction_override_is_literal_and_preserves_catalog_messages() { }), }; let token_budget = ModelTokenBudgetConfig { + enabled: false, + use_history_notes_extension: false, reminder_threshold_tokens: 128, reminder_message_template: "budget reminder".to_string(), guidance_message: "budget guidance".to_string(), diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 6b6d481bda..2e3104e864 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -594,6 +594,10 @@ pub struct ToolMessage { /// Model-owned defaults for the context-window token-budget feature. #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] pub struct ModelTokenBudgetConfig { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub use_history_notes_extension: bool, pub reminder_threshold_tokens: i64, pub reminder_message_template: String, pub guidance_message: String,