diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 264e19c8e0..bb6a528b0e 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -4495,6 +4495,10 @@ "enabled": { "type": "boolean" }, + "enforce_limit": { + "description": "Stop when the shared budget is exhausted. Defaults to true.", + "type": "boolean" + }, "limit_tokens": { "format": "int64", "minimum": 1.0, @@ -8514,4 +8518,4 @@ }, "title": "ConfigToml", "type": "object" -} +} \ No newline at end of file diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index b79f113fad..16c6fc8ef5 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -942,6 +942,7 @@ prefill_token_weight = 0.1 assert_eq!( config.rollout_budget, Some(RolloutBudgetConfig { + enforce_limit: true, limit_tokens: 100_000, reminder_at_remaining_tokens: vec![50_000, 25_000, 10_000], sampling_token_weight: 1.0, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 810df098ae..1f192747a4 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1355,6 +1355,7 @@ impl Default for TokenBudgetConfig { #[derive(Debug, Clone, PartialEq, Serialize)] pub struct RolloutBudgetConfig { + pub enforce_limit: bool, pub limit_tokens: i64, pub reminder_at_remaining_tokens: Vec, pub sampling_token_weight: f64, @@ -3046,6 +3047,7 @@ fn resolve_rollout_budget_config( } } Ok(Some(RolloutBudgetConfig { + enforce_limit: config.enforce_limit.unwrap_or(true), limit_tokens, reminder_at_remaining_tokens, sampling_token_weight, diff --git a/codex-rs/core/src/rollout_budget.rs b/codex-rs/core/src/rollout_budget.rs index 1b010110a8..332d15b17b 100644 --- a/codex-rs/core/src/rollout_budget.rs +++ b/codex-rs/core/src/rollout_budget.rs @@ -44,7 +44,7 @@ impl RolloutBudget { }); } - /// Returns true once the configured budget is exhausted, including on later calls. + /// Records usage and returns whether an enforced budget is exhausted. pub(crate) fn record_usage(&self, usage: &TokenUsage) -> CodexResult { let Some(mut state) = self.lock() else { return Ok(false); @@ -63,7 +63,8 @@ impl RolloutBudget { + usage.non_cached_input() as f64 * state.config.prefill_token_weight }; state.weighted_tokens_used += units; - Ok(state.weighted_tokens_used >= state.config.limit_tokens as f64) + Ok(state.config.enforce_limit + && state.weighted_tokens_used >= state.config.limit_tokens as f64) } pub(crate) fn pending_reminder( diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index a2afab3285..b46fe16e2e 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1288,6 +1288,7 @@ async fn spawn_internal_session_preserves_parent_lineage_without_forking_history let temp_dir = tempdir().expect("tempdir"); let mut config = test_config().await; config.rollout_budget = Some(RolloutBudgetConfig { + enforce_limit: true, limit_tokens: 100, reminder_at_remaining_tokens: vec![75, 50, 25], sampling_token_weight: 1.0, diff --git a/codex-rs/core/tests/suite/context_annotations.rs b/codex-rs/core/tests/suite/context_annotations.rs index d078c9bedf..30b6f81947 100644 --- a/codex-rs/core/tests/suite/context_annotations.rs +++ b/codex-rs/core/tests/suite/context_annotations.rs @@ -49,6 +49,7 @@ async fn first_request_item_types_roles_and_content_annotations() -> Result<()> ..TokenBudgetConfig::default() }); config.rollout_budget = Some(RolloutBudgetConfig { + enforce_limit: true, limit_tokens: 100, reminder_at_remaining_tokens: Vec::new(), sampling_token_weight: 1.0, diff --git a/codex-rs/core/tests/suite/guardian_context_budget.rs b/codex-rs/core/tests/suite/guardian_context_budget.rs index d966d37b92..ef336d3db3 100644 --- a/codex-rs/core/tests/suite/guardian_context_budget.rs +++ b/codex-rs/core/tests/suite/guardian_context_budget.rs @@ -277,6 +277,7 @@ async fn review_respects_complete_context_budget( .expect("test permits time reminders"); config.current_time_reminder = Some(CurrentTimeReminderConfig::default()); config.rollout_budget = Some(RolloutBudgetConfig { + enforce_limit: true, limit_tokens: 1_000_000, reminder_at_remaining_tokens: Vec::new(), sampling_token_weight: 1.0, diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index 8192a3f0c5..c0c244ab45 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -30,6 +30,7 @@ const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration"; fn rollout_budget() -> RolloutBudgetConfig { RolloutBudgetConfig { + enforce_limit: true, limit_tokens: 100, reminder_at_remaining_tokens: vec![75, 50, 25], sampling_token_weight: 1.0, @@ -277,11 +278,15 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn exhausted_budget_fails_current_and_later_turns() -> Result<()> { +#[test_case(true ; "enforced")] +#[test_case(false ; "advisory")] +async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( + enforce_limit: bool, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - mount_sse_sequence( + let responses = mount_sse_sequence( &server, vec![ sse(vec![ @@ -296,8 +301,9 @@ async fn exhausted_budget_fails_current_and_later_turns() -> Result<()> { ) .await; let test = test_codex() - .with_config(|config| { + .with_config(move |config| { config.rollout_budget = Some(RolloutBudgetConfig { + enforce_limit, limit_tokens: 30, reminder_at_remaining_tokens: vec![20, 10], ..rollout_budget() @@ -314,20 +320,31 @@ async fn exhausted_budget_fails_current_and_later_turns() -> Result<()> { }])) .await?; - wait_for_event(&test.codex, |event| { - matches!( - event, - EventMsg::Error(error) - if error.codex_error_info == Some(CodexErrorInfo::SessionBudgetExceeded) - ) + let event = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::Error(_) | EventMsg::TurnComplete(_)) }) .await; + if !enforce_limit { + assert!(matches!(event, EventMsg::TurnComplete(_))); + continue; + } + assert!(matches!( + event, + EventMsg::Error(error) + if error.codex_error_info == Some(CodexErrorInfo::SessionBudgetExceeded) + )); wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) .await; } + if !enforce_limit { + assert_eq!( + rollout_budget_texts(&responses.requests()[1]).last(), + Some(&rollout_budget_message(/*remaining_tokens*/ 0)) + ); + } Ok(()) } diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 0ca7e1a1bb..1995ea9945 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -421,6 +421,9 @@ impl FeatureConfig for TokenBudgetConfigToml { #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct RolloutBudgetConfigToml { + /// Stop when the shared budget is exhausted. Defaults to true. + #[serde(skip_serializing_if = "Option::is_none")] + pub enforce_limit: Option, #[serde(skip_serializing_if = "Option::is_none")] pub enabled: Option, #[serde(skip_serializing_if = "Option::is_none")]