github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 19:08:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support model-owned token budget defaults (#35608)

Browse Source 
## 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
This commit is contained in:
rka-oai
2026-07-27 12:52:37 +00:00
committed by copyberry
parent 16d7daad7c
commit fbe65995bb
17 changed files with 742 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/config/mod.rs
+76 -60
View File
@@ -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
+54
View File
@@ -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
+52
View File
@@ -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
+2
View File
@@ -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
+3
View File
@@ -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
+26 -1
View File
@@ -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
+6 -11
View File
@@ -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
+52
View File
@@ -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
+1
View File
@@ -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
+12
View File
@@ -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
codex-rs/core/tests/suite/catalog_permission_messages.rs
+1
View File
@@ -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,
codex-rs/core/tests/suite/permissions_messages.rs
+3
View File
@@ -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)
codex-rs/core/tests/suite/personality.rs
+2
View File
@@ -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,
codex-rs/core/tests/suite/token_budget.rs
+425
View File
@@ -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<String> {
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("<model_switch>")),
"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")
codex-rs/models-manager/src/model_info.rs
+2
View File
@@ -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<ModelMessages> {
approvals: None,
auto_review: None,
permissions: None,
token_budget: None,
}),
_ => None,
}
codex-rs/models-manager/src/model_info_tests.rs
+9
View File
@@ -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);
codex-rs/protocol/src/openai_models.rs
+16
View File
@@ -509,6 +509,18 @@ pub struct ModelMessages {
pub approvals: Option<ApprovalMessages>,
pub auto_review: Option<AutoReviewMessages>,
pub permissions: Option<PermissionMessages>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub token_budget: Option<ModelTokenBudgetConfig>,
}
/// 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));