diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index fad074e9c..8547269ce 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -44,6 +44,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::TokenBudget { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::RolloutBudget { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 828ec69f0..2655f2e74 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -669,7 +669,7 @@ "type": "boolean" }, "token_budget": { - "type": "boolean" + "$ref": "#/definitions/FeatureToml_for_TokenBudgetConfigToml" }, "tool_call_mcp_elicitation": { "type": "boolean" @@ -967,6 +967,16 @@ } ] }, + "FeatureToml_for_TokenBudgetConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/TokenBudgetConfigToml" + } + ] + }, "FeedbackConfigToml": { "additionalProperties": false, "properties": { @@ -2836,6 +2846,27 @@ } ] }, + "TokenBudgetConfigToml": { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "reminder_message_template": { + "description": "Reminder template. `{n_remaining}` is replaced with the tokens remaining before auto-compaction.", + "maxLength": 1000, + "minLength": 1, + "type": "string" + }, + "reminder_threshold_tokens": { + "description": "Number of tokens remaining before auto-compaction when the wrap-up reminder is emitted.", + "format": "int64", + "minimum": 1.0, + "type": "integer" + } + }, + "type": "object" + }, "ToolSuggestConfig": { "additionalProperties": false, "properties": { @@ -4932,7 +4963,7 @@ "type": "boolean" }, "token_budget": { - "type": "boolean" + "$ref": "#/definitions/FeatureToml_for_TokenBudgetConfigToml" }, "tool_call_mcp_elicitation": { "type": "boolean" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 904018ebc..ce9374016 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -460,6 +460,90 @@ direct_only_tool_namespaces = ["mcp__history", "mcp__notes"] Ok(()) } +#[tokio::test] +async fn load_config_resolves_token_budget_config() -> std::io::Result<()> { + for (config_toml, expected) in [ + ( + "[features]\ntoken_budget = true\n", + TokenBudgetConfig::default(), + ), + ( + r#" +[features.token_budget] +enabled = true +reminder_threshold_tokens = 16000 +reminder_message_template = "Custom reminder: {n_remaining} tokens." +"#, + TokenBudgetConfig { + reminder_threshold_tokens: Some(16_000), + reminder_message_template: "Custom reminder: {n_remaining} tokens.".to_string(), + }, + ), + ] { + let codex_home = tempdir()?; + let config_toml = toml::from_str(config_toml).expect("TOML should deserialize"); + let config = Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await?; + + assert!(config.features.enabled(Feature::TokenBudget)); + assert_eq!(config.token_budget, Some(expected)); + } + Ok(()) +} + +#[tokio::test] +async fn load_config_rejects_invalid_token_budget_reminder_template() -> std::io::Result<()> { + for reminder_message_template in [ + String::new(), + "x".repeat(TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES + 1), + ] { + let codex_home = tempdir()?; + let config_toml = toml::from_str(&format!( + "[features.token_budget]\nenabled = true\nreminder_message_template = {reminder_message_template:?}\n" + )) + .expect("TOML should deserialize"); + let error = Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await + .expect_err("invalid reminder template should be rejected"); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + } + Ok(()) +} + +#[tokio::test] +async fn load_config_rejects_non_positive_token_budget_reminder_threshold() -> std::io::Result<()> { + for reminder_threshold_tokens in [-1, 0] { + let codex_home = tempdir()?; + let config_toml = toml::from_str(&format!( + "[features.token_budget]\nenabled = true\nreminder_threshold_tokens = {reminder_threshold_tokens}\n" + )) + .expect("TOML should deserialize"); + let error = Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await + .expect_err("non-positive reminder threshold should be rejected"); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + error.to_string(), + "features.token_budget.reminder_threshold_tokens must be positive" + ); + } + Ok(()) +} + #[tokio::test] async fn load_config_resolves_rollout_budget() -> std::io::Result<()> { let codex_home = tempdir()?; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 8e9ccc6c2..d5986501a 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -68,6 +68,7 @@ use codex_features::Features; use codex_features::FeaturesToml; use codex_features::MultiAgentV2ConfigToml; use codex_features::NetworkProxyConfigToml; +use codex_features::TokenBudgetConfigToml; use codex_git_utils::resolve_root_git_project_for_trust; use codex_install_context::InstallContext; use codex_login::AuthManagerConfig; @@ -1026,6 +1027,8 @@ pub struct Config { /// Settings specific to the task-path-based multi-agent tool surface. pub multi_agent_v2: MultiAgentV2Config, + /// Context-window token budget configuration, when enabled. + pub token_budget: Option, /// Shared token budget for the root thread and its sub-agents. pub rollout_budget: Option, /// Current-time reminder configuration, when enabled. @@ -1075,6 +1078,27 @@ pub struct CodeModeConfig { pub direct_only_tool_namespaces: Vec, } +pub(crate) const DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE: &str = concat!( + "Your context window is nearly exhausted (only {n_remaining} tokens remaining) and will be automatically reset for you soon. ", + "Once reset, message items in current context window will be cleared in the new window, but notes and history items will be persistent across windows." +); +const TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES: usize = 1000; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct TokenBudgetConfig { + pub reminder_threshold_tokens: Option, + pub reminder_message_template: String, +} + +impl Default for TokenBudgetConfig { + fn default() -> Self { + Self { + reminder_threshold_tokens: None, + reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(), + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Serialize)] pub struct RolloutBudgetConfig { pub limit_tokens: i64, @@ -2509,6 +2533,48 @@ fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config } } +fn resolve_token_budget_config( + config_toml: &ConfigToml, + features: &ManagedFeatures, +) -> std::io::Result> { + if !features.enabled(Feature::TokenBudget) { + return Ok(None); + } + + 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" + ), + )); + } + + Ok(Some(TokenBudgetConfig { + reminder_threshold_tokens, + reminder_message_template, + })) +} + fn resolve_rollout_budget_config( config_toml: &ConfigToml, features: &ManagedFeatures, @@ -2635,6 +2701,13 @@ fn multi_agent_v2_toml_config(features: Option<&FeaturesToml>) -> Option<&MultiA } } +fn token_budget_toml_config(features: Option<&FeaturesToml>) -> Option<&TokenBudgetConfigToml> { + match features?.token_budget.as_ref()? { + FeatureToml::Enabled(_) => None, + FeatureToml::Config(config) => Some(config), + } +} + fn current_time_reminder_toml_config( features: Option<&FeaturesToml>, ) -> Option<&CurrentTimeReminderConfigToml> { @@ -3287,6 +3360,7 @@ impl Config { resolve_experimental_request_user_input_enabled(&cfg); let code_mode = resolve_code_mode_config(&cfg); let multi_agent_v2 = resolve_multi_agent_v2_config(&cfg); + let token_budget = resolve_token_budget_config(&cfg, &features)?; let rollout_budget = resolve_rollout_budget_config(&cfg, &features)?; let current_time_reminder = resolve_current_time_reminder_config(&cfg, &features)?; let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg); @@ -3830,6 +3904,7 @@ impl Config { background_terminal_max_timeout, ghost_snapshot, multi_agent_v2, + token_budget, rollout_budget, current_time_reminder, features, diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 3347358bb..c9319ebef 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -73,6 +73,7 @@ pub(crate) use rollout_budget::RolloutBudgetContext; pub(crate) use subagent_notification::SubagentNotification; pub(crate) use token_budget_context::TokenBudgetContext; pub(crate) use token_budget_context::TokenBudgetRemainingContext; +pub(crate) use token_budget_context::TokenBudgetReminder; pub(crate) use turn_aborted::TurnAborted; pub(crate) use user_instructions::UserInstructions; pub(crate) use user_shell_command::UserShellCommand; diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 24e1869e8..93d5e1e6a 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -99,3 +99,34 @@ impl ContextualUserFragment for TokenBudgetRemainingContext { } } } + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TokenBudgetReminder { + message: String, +} + +impl TokenBudgetReminder { + pub(crate) fn new(message_template: &str, n_remaining: i64) -> Self { + Self { + message: message_template.replace("{n_remaining}", &n_remaining.to_string()), + } + } +} + +impl ContextualUserFragment for TokenBudgetReminder { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("\n", "\n") + } + + fn body(&self) -> String { + self.message.clone() + } +} diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index 79afd1151..2553f7150 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -10,6 +10,7 @@ use codex_features::FeatureToml; use codex_features::FeaturesToml; use codex_features::MultiAgentV2ConfigToml; use codex_features::RolloutBudgetConfigToml; +use codex_features::TokenBudgetConfigToml; use codex_protocol::ThreadId; use crate::config::Config; @@ -152,6 +153,12 @@ 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() { + let mut token_budget: TokenBudgetConfigToml = + resolved_config_to_toml(token_budget, "features.token_budget")?; + token_budget.enabled = Some(config.features.enabled(Feature::TokenBudget)); + features.token_budget = Some(FeatureToml::Config(token_budget)); + } if let Some(rollout_budget) = config.rollout_budget.as_ref() { let mut rollout_budget: RolloutBudgetConfigToml = resolved_config_to_toml(rollout_budget, "features.rollout_budget")?; @@ -238,6 +245,14 @@ mod tests { async fn lock_contains_prompts_and_materializes_features() { 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 { + reminder_threshold_tokens: Some(16_000), + reminder_message_template: "Locked reminder: {n_remaining} tokens.".to_string(), + }); + config + .features + .enable(Feature::TokenBudget) + .expect("token_budget should be enableable in tests"); config.rollout_budget = Some(crate::config::RolloutBudgetConfig { limit_tokens: 100_000, reminder_interval_tokens: 10_000, @@ -318,6 +333,17 @@ mod tests { }) )); + assert_eq!( + features.token_budget, + Some(FeatureToml::Config(TokenBudgetConfigToml { + enabled: Some(true), + reminder_threshold_tokens: Some(16_000), + reminder_message_template: Some( + "Locked reminder: {n_remaining} tokens.".to_string() + ), + })) + ); + assert_eq!( features.rollout_budget, Some(FeatureToml::Config(RolloutBudgetConfigToml { diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index 4548ed319..017b7d542 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -5,40 +5,61 @@ use codex_features::Feature; const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75]; -pub(super) async fn maybe_record_token_budget_remaining_context( +pub(super) async fn maybe_record( sess: &Session, turn_context: &TurnContext, tokens_before_sampling: i64, tokens_after_sampling: i64, + tokens_until_compaction: i64, ) { if !turn_context.config.features.enabled(Feature::TokenBudget) { return; } - let Some(model_context_window) = turn_context.model_context_window() else { - return; - }; - if model_context_window <= 0 || tokens_after_sampling <= tokens_before_sampling { - return; + + let mut response_items = Vec::new(); + if let Some(model_context_window) = turn_context.model_context_window() + && model_context_window > 0 + && tokens_after_sampling > tokens_before_sampling + { + let tokens_before_sampling = tokens_before_sampling.max(0); + let tokens_after_sampling = tokens_after_sampling.max(0); + let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| { + tokens_before_sampling.saturating_mul(100) + < model_context_window.saturating_mul(*threshold) + && tokens_after_sampling.saturating_mul(100) + >= model_context_window.saturating_mul(*threshold) + }); + if crossed_threshold { + let tokens_left = model_context_window + .saturating_sub(tokens_after_sampling) + .max(0); + response_items.push(ContextualUserFragment::into( + crate::context::TokenBudgetRemainingContext::new(tokens_left), + )); + } } - let tokens_before_sampling = tokens_before_sampling.max(0); - let tokens_after_sampling = tokens_after_sampling.max(0); - let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| { - tokens_before_sampling.saturating_mul(100) < model_context_window.saturating_mul(*threshold) - && tokens_after_sampling.saturating_mul(100) - >= model_context_window.saturating_mul(*threshold) - }); - if !crossed_threshold { - return; + if let Some(config) = turn_context.config.token_budget.as_ref().filter(|config| { + config + .reminder_threshold_tokens + .is_some_and(|threshold| tokens_until_compaction <= threshold) + }) { + let reminder_due = { + let mut state = sess.state.lock().await; + state.claim_token_budget_reminder() + }; + if reminder_due { + response_items.push(ContextualUserFragment::into( + crate::context::TokenBudgetReminder::new( + &config.reminder_message_template, + tokens_until_compaction, + ), + )); + } } - let tokens_left = model_context_window - .saturating_sub(tokens_after_sampling) - .max(0); - - let response_item = ContextualUserFragment::into( - crate::context::TokenBudgetRemainingContext::new(tokens_left), - ); - sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) - .await; + if !response_items.is_empty() { + sess.record_conversation_items(turn_context, &response_items) + .await; + } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 28c0f6561..df74841b1 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -313,11 +313,22 @@ pub(crate) async fn run_turn( ); let tokens_after_sampling = token_status.active_context_tokens; - super::token_budget::maybe_record_token_budget_remaining_context( + let full_context_remaining = token_status + .full_context_window_limit + .map_or(i64::MAX, |limit| { + limit.saturating_sub(tokens_after_sampling) + }); + let tokens_until_compaction = token_status + .auto_compact_scope_limit + .saturating_sub(token_status.auto_compact_scope_tokens) + .min(full_context_remaining) + .max(0); + super::token_budget::maybe_record( sess.as_ref(), turn_context.as_ref(), tokens_before_sampling, tokens_after_sampling, + tokens_until_compaction, ) .await; diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index ff5d7a2a2..419e85be0 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -30,6 +30,7 @@ pub(super) struct AutoCompactWindow { /// not the growth itself; server-observed usage replaces estimated /// resume/recompute baselines when available. prefill_input_tokens: Option, + token_budget_reminder_delivered: bool, } impl AutoCompactWindow { @@ -44,6 +45,7 @@ impl AutoCompactWindow { }, new_context_window_requested: false, prefill_input_tokens: None, + token_budget_reminder_delivered: false, } } @@ -69,9 +71,14 @@ impl AutoCompactWindow { self.ids.previous_window_id = Some(self.ids.window_id); self.ids.window_id = Uuid::now_v7(); self.new_context_window_requested = false; + self.token_budget_reminder_delivered = false; (self.window_number, self.ids) } + pub(super) fn claim_token_budget_reminder(&mut self) -> bool { + !std::mem::replace(&mut self.token_budget_reminder_delivered, true) + } + pub(super) fn request_new_context_window(&mut self) { self.new_context_window_requested = true; } @@ -154,6 +161,8 @@ mod tests { ); assert_eq!(window.window_number(), 3); assert_eq!(window.ids().window_id, restored_window_id); + assert!(window.claim_token_budget_reminder()); + assert!(!window.claim_token_budget_reminder()); window.request_new_context_window(); assert!(window.take_new_context_window_request()); assert!(!window.take_new_context_window_request()); @@ -167,6 +176,7 @@ mod tests { assert_eq!(ids.window_id.get_version_num(), 7); assert_ne!(ids.window_id, restored_window_id); assert!(!window.take_new_context_window_request()); + assert!(window.claim_token_budget_reminder()); assert_eq!( window.snapshot(), diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 824f5f636..6444dc8cf 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -148,6 +148,10 @@ impl SessionState { self.auto_compact_window.snapshot() } + pub(crate) fn claim_token_budget_reminder(&mut self) -> bool { + self.auto_compact_window.claim_token_budget_reminder() + } + pub(crate) fn auto_compact_window_number(&self) -> u64 { self.auto_compact_window.window_number() } diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 80f852a33..a1f8fb5de 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -1,6 +1,7 @@ use anyhow::Result; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; +use codex_core::config::TokenBudgetConfig; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; use codex_protocol::protocol::EventMsg; @@ -298,6 +299,57 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_reminder_emits_after_crossing_compaction_threshold() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 8_000), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_context_window = Some(10_000); + config.token_budget = Some(TokenBudgetConfig { + reminder_threshold_tokens: Some(2_000), + ..TokenBudgetConfig::default() + }); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("cross threshold").await?; + test.submit_turn("observe reminder").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + let initial_context = token_budget_texts(&requests[0]); + assert_eq!(initial_context.len(), 1); + let remaining_context = + "\nYou have 1500 tokens left in this context window.\n" + .to_string(); + let reminder = "\nYour context window is nearly exhausted (only 1000 tokens remaining) and will be automatically reset for you soon. Once reset, message items in current context window will be cleared in the new window, but notes and history items will be persistent across windows.\n" + .to_string(); + assert_eq!( + token_budget_texts(&requests[1]), + vec![initial_context[0].clone(), remaining_context, reminder] + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index dc78d8cd4..15a37fe36 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -73,6 +73,32 @@ impl FeatureConfig for MultiAgentV2ConfigToml { } } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct TokenBudgetConfigToml { + #[serde(skip_serializing_if = "Option::is_none")] + pub enabled: Option, + /// Number of tokens remaining before auto-compaction when the wrap-up reminder is emitted. + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1))] + pub reminder_threshold_tokens: Option, + /// Reminder template. `{n_remaining}` is replaced with the tokens remaining before + /// auto-compaction. + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(length(min = 1, max = 1000))] + pub reminder_message_template: Option, +} + +impl FeatureConfig for TokenBudgetConfigToml { + fn enabled(&self) -> Option { + self.enabled + } + + fn set_enabled(&mut self, enabled: bool) { + self.enabled = Some(enabled); + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct RolloutBudgetConfigToml { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 3bba6faf9..f2a30b07b 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -26,6 +26,7 @@ pub use feature_configs::NetworkProxyModeToml; pub use feature_configs::NetworkProxyUnixSocketPermissionToml; use feature_configs::RemovedAppsMcpPathOverrideConfigToml; pub use feature_configs::RolloutBudgetConfigToml; +pub use feature_configs::TokenBudgetConfigToml; use legacy::LegacyFeatureToggles; pub use legacy::legacy_feature_keys; @@ -632,6 +633,8 @@ pub struct FeaturesToml { #[serde(default, skip_serializing_if = "Option::is_none")] pub multi_agent_v2: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_budget: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] pub rollout_budget: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub current_time_reminder: Option>, @@ -667,6 +670,9 @@ impl FeaturesToml { if let Some(enabled) = self.multi_agent_v2.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::MultiAgentV2.key().to_string(), enabled); } + if let Some(enabled) = self.token_budget.as_ref().and_then(FeatureToml::enabled) { + entries.insert(Feature::TokenBudget.key().to_string(), enabled); + } if let Some(enabled) = self.rollout_budget.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::RolloutBudget.key().to_string(), enabled); } @@ -688,6 +694,7 @@ impl FeaturesToml { let Self { code_mode, multi_agent_v2, + token_budget, rollout_budget, current_time_reminder, removed_apps_mcp_path_override: _, @@ -703,6 +710,8 @@ impl FeaturesToml { materialize_resolved_feature_enabled(code_mode, enabled); } else if spec.id == Feature::MultiAgentV2 { materialize_resolved_feature_enabled(multi_agent_v2, enabled); + } else if spec.id == Feature::TokenBudget { + materialize_resolved_feature_enabled(token_budget, enabled); } else if spec.id == Feature::RolloutBudget { materialize_resolved_feature_enabled(rollout_budget, enabled); } else if spec.id == Feature::CurrentTimeReminder { diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index b9796a55c..3d8561a17 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -280,6 +280,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R background_terminal_max_timeout: 300_000, ghost_snapshot: GhostSnapshotConfig::default(), multi_agent_v2: MultiAgentV2Config::default(), + token_budget: None, rollout_budget: None, current_time_reminder: None, features: Default::default(),