diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 216af994bc..0289b65bde 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/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( 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 new file mode 100644 index 0000000000..6228a06b7f --- /dev/null +++ b/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> { + 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; 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 new file mode 100644 index 0000000000..67c960de79 --- /dev/null +++ b/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() + ); +} diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 7fddf2da87..6f2520b2e7 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/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; diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 0340043660..403b599082 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/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( diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index 00c75df317..f01cab93f5 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/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; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 450a1a3765..39f7a0627e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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; diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 35914119ac..80388076cd 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/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, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 49896774a0..70fe47fb5c 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/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), diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index c5c4556dbd..7663beb7c9 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/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 diff --git a/codex-rs/core/tests/suite/catalog_permission_messages.rs b/codex-rs/core/tests/suite/catalog_permission_messages.rs index 3856bae8bb..23796f0dc2 100644 --- a/codex-rs/core/tests/suite/catalog_permission_messages.rs +++ b/codex-rs/core/tests/suite/catalog_permission_messages.rs @@ -40,6 +40,7 @@ async fn catalog_permission_message_loaded_from_remote_models_is_sent() -> Resul workspace_write: None, read_only: Some("remote catalog permissions: {{ network_access }}".to_string()), }), + token_budget: None, }); let models_mock = mount_models_once( &server, diff --git a/codex-rs/core/tests/suite/permissions_messages.rs b/codex-rs/core/tests/suite/permissions_messages.rs index e9fa9f40dd..49982f6908 100644 --- a/codex-rs/core/tests/suite/permissions_messages.rs +++ b/codex-rs/core/tests/suite/permissions_messages.rs @@ -56,6 +56,7 @@ fn model_with_approval_messages( }), auto_review: None, permissions: None, + token_budget: None, }); model } @@ -71,6 +72,7 @@ fn model_with_permission_messages( approvals: None, auto_review: None, permissions: Some(permissions), + token_budget: None, }); model } @@ -223,6 +225,7 @@ async fn catalog_non_on_request_approval_messages_are_sent_in_initial_permission approvals: Some(approvals), auto_review: None, permissions: None, + token_budget: None, }); let mut builder = test_codex() .with_model(model_slug) diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index aa71da4e80..dc19784339 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -642,6 +642,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow approvals: None, auto_review: None, permissions: None, + token_budget: None, }), include_skills_usage_instructions: false, supports_reasoning_summary_parameter: true, @@ -761,6 +762,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - approvals: None, auto_review: None, permissions: None, + token_budget: None, }), include_skills_usage_instructions: false, supports_reasoning_summary_parameter: true, diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 38c7c47e53..62f7a699da 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -1,11 +1,15 @@ use anyhow::Result; +use codex_config::config_toml::ConfigLockfileToml; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; use codex_core::config::TokenBudgetConfig; use codex_features::Feature; +use codex_features::FeatureToml; +use codex_features::TokenBudgetConfigToml; use codex_model_provider_info::built_in_model_providers; use codex_protocol::config_types::AutoCompactTokenLimitScope; use codex_protocol::items::TurnItem; +use codex_protocol::openai_models::ModelTokenBudgetConfig; use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG; use codex_protocol::protocol::CONTEXT_WINDOW_GUIDANCE_OPEN_TAG; @@ -16,6 +20,8 @@ use codex_protocol::protocol::HookRunStatus; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ItemStartedEvent; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; use core_test_support::PathBufExt; use core_test_support::assert_regex_match; use core_test_support::context_snapshot; @@ -29,8 +35,10 @@ use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; use core_test_support::responses::ev_shell_command_call; use core_test_support::responses::mount_compact_json_once; +use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +use core_test_support::responses::sse_completed; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; @@ -44,11 +52,23 @@ use serde_json::Value; use serde_json::json; use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; use std::time::Duration; +use tempfile::TempDir; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; const AUTO_COMPACT_FALLBACK_PROMPT: &str = "Save the important state before rollover."; +fn model_token_budget_config() -> ModelTokenBudgetConfig { + ModelTokenBudgetConfig { + 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(), + auto_compact_fallback_prompt: AUTO_COMPACT_FALLBACK_PROMPT.to_string(), + auto_compact_fallback_buffer_tokens: 16_384, + } +} + fn token_budget_contexts(request: &ResponsesRequest) -> Vec { let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: "); request @@ -261,6 +281,385 @@ async fn token_budget_guidance_follows_context_window() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_uses_model_message_defaults() -> Result<()> { + skip_if_no_network!(Ok(())); + + 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 expected_guidance = model_defaults.guidance_message.clone(); + let test = test_codex() + .with_model_info_override("gpt-5.2", move |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_defaults); + }) + .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?; + + test.submit_turn("inspect model-owned context guidance") + .await?; + + let developer_texts = response.single_request().message_input_texts("developer"); + assert!(developer_texts.iter().any(|text| { + text == &format!( + "{CONTEXT_WINDOW_GUIDANCE_OPEN_TAG}\n{expected_guidance}\n{CONTEXT_WINDOW_GUIDANCE_CLOSE_TAG}" + ) + })); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_explicit_default_template_overrides_model_defaults() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response = mount_sse_once(&server, sse_completed("resp-1")).await; + let test = test_codex() + .with_pre_build_hook(|home| { + let default_template = TokenBudgetConfig::default().reminder_message_template; + std::fs::write( + home.join("config.toml"), + format!( + "[features.token_budget]\nenabled = true\nreminder_message_template = {default_template:?}\n" + ), + ) + .expect("write explicit default token-budget configuration"); + }) + .with_model_info_override("gpt-5.2", |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_token_budget_config()); + }) + .with_config(|config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + }) + .build_with_auto_env(&server) + .await?; + + test.submit_turn("inspect explicitly configured default template") + .await?; + + assert!( + !response + .single_request() + .message_input_texts("developer") + .iter() + .any(|text| text.contains("Use the model-owned context-window guidance.")), + "an explicitly configured default template must not inherit model-owned settings" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_model_defaults_survive_config_lock_replay() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![sse_completed("resp-1"), sse_completed("resp-2")], + ) + .await; + let home = Arc::new(TempDir::new()?); + let initial = test_codex() + .with_home(Arc::clone(&home)) + .with_model_info_override("gpt-5.2", |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_token_budget_config()); + }) + .with_config(|config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config.config_lock_export_dir = Some(config.codex_home.join("memento-config-locks")); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build_with_auto_env(&server) + .await?; + initial + .submit_turn("inspect guidance before lock replay") + .await?; + + let lock_path = initial + .codex_home_path() + .join("memento-config-locks") + .join(format!( + "{}.config.lock.toml", + initial.session_configured.thread_id + )); + let lock: ConfigLockfileToml = toml::from_str(&std::fs::read_to_string(&lock_path)?)?; + assert_eq!( + lock.config + .features + .as_ref() + .and_then(|features| features.token_budget.as_ref()), + Some(&FeatureToml::Config(TokenBudgetConfigToml { + enabled: Some(true), + reminder_threshold_tokens: Some(6_144), + reminder_message_template: Some( + "Model reminder: {n_remaining} tokens remain.".to_string() + ), + guidance_message: Some("Use the model-owned context-window guidance.".to_string()), + auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()), + auto_compact_fallback_buffer_tokens: Some(16_384), + })) + ); + let replay = test_codex() + .with_home(home) + .with_pre_build_hook(move |home| { + std::fs::write( + home.join("config.toml"), + format!("[debug.config_lockfile]\nload_path = {lock_path:?}\n"), + ) + .expect("write exported config lock replay settings"); + }) + .with_model_info_override("gpt-5.2", |model_info| { + let mut updated_defaults = model_token_budget_config(); + updated_defaults.reminder_threshold_tokens = 4_096; + updated_defaults.reminder_message_template = + "Updated model reminder: {n_remaining} tokens remain.".to_string(); + updated_defaults.guidance_message = + "Use the updated model-owned context-window guidance.".to_string(); + updated_defaults.auto_compact_fallback_prompt = + "Use the updated fallback before rollover.".to_string(); + updated_defaults.auto_compact_fallback_buffer_tokens = 8_192; + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(updated_defaults); + }) + .with_config(|config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + }) + .build_with_auto_env(&server) + .await?; + replay + .submit_turn("inspect guidance after lock replay") + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + for request in requests { + assert!( + request.body_contains_text("Use the model-owned context-window guidance."), + "exporting and replaying a config lock must preserve model-owned guidance" + ); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_defaults_follow_the_active_model() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![sse_completed("resp-1"), sse_completed("resp-2")], + ) + .await; + let test = test_codex() + .with_model_info_override("gpt-5.2", |model_info| { + let mut defaults = model_token_budget_config(); + defaults.guidance_message = "Use first-model context-window guidance.".to_string(); + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(defaults); + }) + .with_model_info_override("gpt-5.4", |model_info| { + let mut defaults = model_token_budget_config(); + defaults.guidance_message = "Use second-model context-window guidance.".to_string(); + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(defaults); + }) + .with_model("gpt-5.2") + .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?; + + test.submit_turn("inspect first-model guidance").await?; + core_test_support::submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + model: Some("gpt-5.4".to_string()), + ..Default::default() + }, + ) + .await?; + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "inspect second-model guidance".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: ThreadSettingsOverrides::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + assert!(requests[0].body_contains_text("Use first-model context-window guidance.")); + assert!( + requests[1].body_contains_text("Use first-model context-window guidance."), + "switching models must preserve the existing conversation history" + ); + let developer_texts = requests[1].message_input_texts("developer"); + assert!( + developer_texts + .iter() + .any(|text| text.contains("")), + "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}" + ); + assert_eq!( + developer_texts + .iter() + .filter(|text| **text == expected_guidance) + .count(), + 1, + "the active model's guidance should be emitted once by world state" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_ignores_invalid_model_message_defaults() -> Result<()> { + skip_if_no_network!(Ok(())); + + let mut zero_reminder_threshold = model_token_budget_config(); + zero_reminder_threshold.reminder_threshold_tokens = 0; + let mut negative_fallback_buffer = model_token_budget_config(); + negative_fallback_buffer.auto_compact_fallback_buffer_tokens = -1; + let mut oversized_reminder = model_token_budget_config(); + oversized_reminder.reminder_message_template = "x".repeat(2_001); + let mut oversized_guidance = model_token_budget_config(); + oversized_guidance.guidance_message = "x".repeat(2_001); + let mut oversized_fallback = model_token_budget_config(); + oversized_fallback.auto_compact_fallback_prompt = "x".repeat(2_001); + + for model_defaults in [ + zero_reminder_threshold, + negative_fallback_buffer, + oversized_reminder, + oversized_guidance, + oversized_fallback, + ] { + let server = start_mock_server().await; + let response = mount_sse_once(&server, sse_completed("resp-1")).await; + let test = test_codex() + .with_model_info_override("gpt-5.2", move |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_defaults); + }) + .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?; + + test.submit_turn("inspect invalid model-owned context guidance") + .await?; + + assert!( + !response + .single_request() + .message_input_texts("developer") + .iter() + .any(|text| text.starts_with(CONTEXT_WINDOW_GUIDANCE_OPEN_TAG)), + "invalid model-owned token-budget settings must not enter the model context" + ); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn model_token_budget_defaults_do_not_enable_disabled_feature() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response = mount_sse_once(&server, sse_completed("resp-1")).await; + let test = test_codex() + .with_model_info_override("gpt-5.2", |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_token_budget_config()); + }) + .with_config(|config| { + config + .features + .disable(Feature::TokenBudget) + .expect("test config should allow token budget to be disabled"); + }) + .build_with_auto_env(&server) + .await?; + + test.submit_turn("inspect disabled context guidance") + .await?; + + let request = response.single_request(); + assert!(token_budget_contexts(&request).is_empty()); + assert!( + !request + .message_input_texts("developer") + .iter() + .any(|text| text.contains("Use the model-owned context-window guidance.")) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { skip_if_no_network!(Ok(())); @@ -807,9 +1206,22 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() model_provider.base_url = Some(format!("{}/v1", server.uri())); model_provider.supports_websockets = false; let test = test_codex() + .with_model_info_override("gpt-5.2", |model_info| { + model_info + .model_messages + .as_mut() + .expect("bundled model should have model messages") + .token_budget = Some(model_token_budget_config()); + }) .with_config(move |config| { config.model_provider = model_provider; config.model_context_window = Some(10_000); + config.token_budget = Some(TokenBudgetConfig { + reminder_threshold_tokens: Some(16_384), + reminder_message_template: "Bridge reminder: {n_remaining} tokens remain." + .to_string(), + ..TokenBudgetConfig::default() + }); config .features .enable(Feature::TokenBudget) @@ -826,6 +1238,19 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() 2, "token-budget auto-compaction should reset locally before the continuation" ); + assert!( + !requests + .iter() + .any(|request| request + .body_contains_text("Use the model-owned context-window guidance.")), + "an explicit token-budget config must not inherit model-owned guidance" + ); + assert!( + !requests + .iter() + .any(|request| request.body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)), + "an explicit token-budget config must not enable model-owned fallback" + ); assert!( !requests[1].input().iter().any(|item| { item.get("type").and_then(Value::as_str) == Some("function_call_output") diff --git a/codex-rs/models-manager/src/model_info.rs b/codex-rs/models-manager/src/model_info.rs index f8c2ee3ecc..c06f74bc65 100644 --- a/codex-rs/models-manager/src/model_info.rs +++ b/codex-rs/models-manager/src/model_info.rs @@ -115,6 +115,7 @@ fn clear_instruction_messages(model: &mut ModelInfo) { if model_messages.approvals.is_none() && model_messages.auto_review.is_none() && model_messages.permissions.is_none() + && model_messages.token_budget.is_none() { model.model_messages = None; } @@ -181,6 +182,7 @@ fn local_personality_messages_for_slug(slug: &str) -> Option { approvals: None, auto_review: None, permissions: None, + token_budget: None, }), _ => None, } diff --git a/codex-rs/models-manager/src/model_info_tests.rs b/codex-rs/models-manager/src/model_info_tests.rs index 9cc3f3cecb..cd85915f2a 100644 --- a/codex-rs/models-manager/src/model_info_tests.rs +++ b/codex-rs/models-manager/src/model_info_tests.rs @@ -33,6 +33,7 @@ fn base_instruction_override_preserves_catalog_approval_messages() { approvals: Some(approvals.clone()), auto_review: None, permissions: None, + token_budget: None, }); let config = ModelsManagerConfig { base_instructions: Some("override".to_string()), @@ -49,6 +50,7 @@ fn base_instruction_override_preserves_catalog_approval_messages() { approvals: Some(approvals), auto_review: None, permissions: None, + token_budget: None, }) ); } @@ -68,6 +70,7 @@ fn disabled_personality_preserves_catalog_approval_messages() { approvals: Some(approvals.clone()), auto_review: None, permissions: None, + token_budget: None, }); let config = ModelsManagerConfig { personality_enabled: false, @@ -84,6 +87,7 @@ fn disabled_personality_preserves_catalog_approval_messages() { approvals: Some(approvals), auto_review: None, permissions: None, + token_budget: None, }) ); } @@ -101,6 +105,7 @@ fn base_instruction_override_preserves_catalog_auto_review_messages() { approvals: None, auto_review: Some(auto_review.clone()), permissions: None, + token_budget: None, }); let config = ModelsManagerConfig { base_instructions: Some("override".to_string()), @@ -117,6 +122,7 @@ fn base_instruction_override_preserves_catalog_auto_review_messages() { approvals: None, auto_review: Some(auto_review), permissions: None, + token_budget: None, }) ); } @@ -135,6 +141,7 @@ fn base_instruction_override_preserves_catalog_permission_messages() { approvals: None, auto_review: None, permissions: Some(permissions.clone()), + token_budget: None, }); let config = ModelsManagerConfig { base_instructions: Some("override".to_string()), @@ -151,6 +158,7 @@ fn base_instruction_override_preserves_catalog_permission_messages() { approvals: None, auto_review: None, permissions: Some(permissions), + token_budget: None, }) ); } @@ -187,6 +195,7 @@ fn personality_none_strips_catalog_instruction_sources_through_the_next_h1() { approvals: None, auto_review: None, permissions: None, + token_budget: None, }); let updated = with_config_overrides(model, &config); diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 5f53bab037..e397a45cb3 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -509,6 +509,18 @@ pub struct ModelMessages { pub approvals: Option, pub auto_review: Option, pub permissions: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_budget: Option, +} + +/// Model-owned defaults for the context-window token-budget feature. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] +pub struct ModelTokenBudgetConfig { + pub reminder_threshold_tokens: i64, + pub reminder_message_template: String, + pub guidance_message: String, + pub auto_compact_fallback_prompt: String, + pub auto_compact_fallback_buffer_tokens: i64, } #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] @@ -922,6 +934,7 @@ mod tests { approvals: None, auto_review: None, permissions: None, + token_budget: None, })); let instructions = model.get_model_instructions(Some(Personality::Friendly)); @@ -941,6 +954,7 @@ mod tests { approvals: None, auto_review: None, permissions: None, + token_budget: None, })); assert_eq!( model.get_model_instructions(Some(Personality::Friendly)), @@ -969,6 +983,7 @@ mod tests { approvals: None, auto_review: None, permissions: None, + token_budget: None, })); assert_eq!( model_no_personality.get_model_instructions(Some(Personality::Friendly)), @@ -1000,6 +1015,7 @@ mod tests { approvals: None, auto_review: None, permissions: None, + token_budget: None, })); let instructions = model.get_model_instructions(Some(Personality::Friendly));