From bd5bd953fb2a5d610a30d112eafe20b644924085 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Mon, 22 Jun 2026 13:25:48 -0700 Subject: [PATCH] [codex] configure rollout budget reminder thresholds (#29423) ## Summary Instead of: reminder_interval_tokens = 65_536 allow users to configure explicit remaining-token reminder thresholds: reminder_at_remaining_tokens = [65_536, 32_768, 16_384, 8_192, 4_096, 2_048, 1_024, 512] ## Validation - CARGO_INCREMENTAL=0 just test -p codex-core rollout_budget: 9 passed - just fix -p codex-core - just fmt --- codex-rs/core/config.schema.json | 11 ++++--- codex-rs/core/src/config/config_tests.rs | 4 +-- codex-rs/core/src/config/mod.rs | 26 ++++++++++----- codex-rs/core/src/rollout_budget.rs | 14 +++++--- codex-rs/core/src/session/config_lock.rs | 4 +-- codex-rs/core/src/thread_manager.rs | 2 +- codex-rs/core/tests/suite/rollout_budget.rs | 36 +++++++++++---------- codex-rs/features/src/feature_configs.rs | 4 +-- 8 files changed, 60 insertions(+), 41 deletions(-) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index b6736c940..4615f4e85 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2674,10 +2674,13 @@ "minimum": 0.0, "type": "number" }, - "reminder_interval_tokens": { - "format": "int64", - "minimum": 1.0, - "type": "integer" + "reminder_at_remaining_tokens": { + "description": "Remaining weighted-token values that trigger reminders when crossed.", + "items": { + "format": "int64", + "type": "integer" + }, + "type": "array" }, "sampling_token_weight": { "format": "double", diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index be332cd60..e0dc04a48 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -552,7 +552,7 @@ async fn load_config_resolves_rollout_budget() -> std::io::Result<()> { [features.rollout_budget] enabled = true limit_tokens = 100000 -reminder_interval_tokens = 10000 +reminder_at_remaining_tokens = [50000, 25000, 10000] sampling_token_weight = 1.0 prefill_token_weight = 0.1 "#, @@ -571,7 +571,7 @@ prefill_token_weight = 0.1 config.rollout_budget, Some(RolloutBudgetConfig { limit_tokens: 100_000, - reminder_interval_tokens: 10_000, + reminder_at_remaining_tokens: vec![50_000, 25_000, 10_000], sampling_token_weight: 1.0, prefill_token_weight: 0.1, }) diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 19d4bb44c..371257d18 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1098,10 +1098,10 @@ impl Default for TokenBudgetConfig { } } -#[derive(Debug, Clone, Copy, PartialEq, Serialize)] +#[derive(Debug, Clone, PartialEq, Serialize)] pub struct RolloutBudgetConfig { pub limit_tokens: i64, - pub reminder_interval_tokens: i64, + pub reminder_at_remaining_tokens: Vec, pub sampling_token_weight: f64, pub prefill_token_weight: f64, } @@ -2606,13 +2606,23 @@ fn resolve_rollout_budget_config( "features.rollout_budget.limit_tokens must be positive", )); } - let reminder_interval_tokens = config - .reminder_interval_tokens - .unwrap_or_else(|| (limit_tokens / 10).max(1)); - if reminder_interval_tokens <= 0 { + let reminder_at_remaining_tokens = + config + .reminder_at_remaining_tokens + .clone() + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "features.rollout_budget.reminder_at_remaining_tokens is required when rollout_budget is enabled", + ) + })?; + if reminder_at_remaining_tokens + .iter() + .any(|&tokens| tokens <= 0 || tokens >= limit_tokens) + { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, - "features.rollout_budget.reminder_interval_tokens must be positive", + "features.rollout_budget.reminder_at_remaining_tokens must contain only positive values below limit_tokens", )); } let sampling_token_weight = config.sampling_token_weight.unwrap_or(1.0); @@ -2630,7 +2640,7 @@ fn resolve_rollout_budget_config( } Ok(Some(RolloutBudgetConfig { limit_tokens, - reminder_interval_tokens, + reminder_at_remaining_tokens, sampling_token_weight, prefill_token_weight, })) diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs index b8901a9e3..334e9e73a 100644 --- a/codex-rs/core/src/rollout_budget.rs +++ b/codex-rs/core/src/rollout_budget.rs @@ -57,18 +57,22 @@ impl RolloutBudget { window_id: &str, ) -> Option { let state = self.lock()?; - let reminder_index = (state.weighted_tokens_used - / state.config.reminder_interval_tokens as f64) + let remaining_tokens = (state.config.limit_tokens as f64 - state.weighted_tokens_used) + .max(0.0) .floor() as i64; + let reminder_index = state + .config + .reminder_at_remaining_tokens + .iter() + .filter(|&&threshold| remaining_tokens <= threshold) + .count() as i64; if state.deliveries.get(&thread_id).is_some_and(|delivery| { delivery.window_id.as_str() == window_id && delivery.reminder_index >= reminder_index }) { return None; } Some(RolloutBudgetReminder { - remaining_tokens: (state.config.limit_tokens as f64 - state.weighted_tokens_used) - .max(0.0) - .floor() as i64, + remaining_tokens, reminder_index, }) } diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index e91c4f489..035f7bce4 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -255,7 +255,7 @@ mod tests { .expect("token_budget should be enableable in tests"); config.rollout_budget = Some(crate::config::RolloutBudgetConfig { limit_tokens: 100_000, - reminder_interval_tokens: 10_000, + reminder_at_remaining_tokens: vec![50_000, 25_000, 10_000], sampling_token_weight: 1.0, prefill_token_weight: 0.25, }); @@ -348,7 +348,7 @@ mod tests { Some(FeatureToml::Config(RolloutBudgetConfigToml { enabled: Some(true), limit_tokens: Some(100_000), - reminder_interval_tokens: Some(10_000), + reminder_at_remaining_tokens: Some(vec![50_000, 25_000, 10_000]), sampling_token_weight: Some(1.0), prefill_token_weight: Some(0.25), })) diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 2cce3bb8a..deba2864b 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1024,7 +1024,7 @@ impl ThreadManager { } fn agent_control_for_config(&self, config: &Config) -> AgentControl { - AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget) + AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget.clone()) } #[cfg(test)] diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index 61c6dfe42..feafdc51e 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -25,12 +25,14 @@ use std::time::Duration; use test_case::test_case; use tokio::time::timeout; -const ROLLOUT_BUDGET: RolloutBudgetConfig = RolloutBudgetConfig { - limit_tokens: 100, - reminder_interval_tokens: 25, - sampling_token_weight: 1.0, - prefill_token_weight: 1.0, -}; +fn rollout_budget() -> RolloutBudgetConfig { + RolloutBudgetConfig { + limit_tokens: 100, + reminder_at_remaining_tokens: vec![75, 50, 25], + sampling_token_weight: 1.0, + prefill_token_weight: 1.0, + } +} fn rollout_budget_texts(request: &ResponsesRequest) -> Vec { request @@ -51,7 +53,7 @@ fn wire_request_contains(request: &wiremock::Request, text: &str) -> bool { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adds_weighted_initial_and_periodic_reminders() -> Result<()> { +async fn adds_weighted_initial_and_threshold_reminders() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -83,7 +85,7 @@ async fn adds_weighted_initial_and_periodic_reminders() -> Result<()> { config.rollout_budget = Some(RolloutBudgetConfig { sampling_token_weight: 2.0, prefill_token_weight: 0.5, - ..ROLLOUT_BUDGET + ..rollout_budget() }); }) .build(&server) @@ -171,7 +173,7 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { .features .enable(Feature::MultiAgentV2) .expect("test config should allow multi-agent v2"); - config.rollout_budget = Some(ROLLOUT_BUDGET); + config.rollout_budget = Some(rollout_budget()); }) .build(&server) .await?; @@ -218,8 +220,8 @@ async fn exhausted_budget_aborts_current_and_later_turns() -> Result<()> { .with_config(|config| { config.rollout_budget = Some(RolloutBudgetConfig { limit_tokens: 30, - reminder_interval_tokens: 10, - ..ROLLOUT_BUDGET + reminder_at_remaining_tokens: vec![20, 10], + ..rollout_budget() }); }) .build(&server) @@ -293,8 +295,8 @@ async fn compaction_budget_exhaustion_aborts_without_error_or_retry(remote_v2: b config.model_provider = model_provider; config.rollout_budget = Some(RolloutBudgetConfig { limit_tokens: 10, - reminder_interval_tokens: 5, - ..ROLLOUT_BUDGET + reminder_at_remaining_tokens: vec![5], + ..rollout_budget() }); if remote_v2 { config @@ -354,8 +356,8 @@ async fn restates_the_current_remainder_after_compaction() -> Result<()> { .with_config(move |config| { config.model_provider = model_provider; config.rollout_budget = Some(RolloutBudgetConfig { - reminder_interval_tokens: 50, - ..ROLLOUT_BUDGET + reminder_at_remaining_tokens: vec![50], + ..rollout_budget() }); }) .build(&server) @@ -409,8 +411,8 @@ async fn restates_the_current_remainder_after_rollback() -> Result<()> { let test = test_codex() .with_config(|config| { config.rollout_budget = Some(RolloutBudgetConfig { - reminder_interval_tokens: 50, - ..ROLLOUT_BUDGET + reminder_at_remaining_tokens: vec![50], + ..rollout_budget() }); }) .build(&server) diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index a76aaaa3c..35b669b0a 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -109,8 +109,8 @@ pub struct RolloutBudgetConfigToml { #[schemars(range(min = 1))] pub limit_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] - #[schemars(range(min = 1))] - pub reminder_interval_tokens: Option, + /// Remaining weighted-token values that trigger reminders when crossed. + pub reminder_at_remaining_tokens: Option>, #[serde(skip_serializing_if = "Option::is_none")] #[schemars(range(min = 0.0))] pub sampling_token_weight: Option,