Allow advisory rollout budgets with optional limit enforcement (#57187)

## What changed

Add `features.rollout_budget.enforce_limit`, defaulting to `true`. Setting it to `false` keeps shared usage accounting and budget reminders active while allowing turns to continue after the budget is exhausted.

## Testing

Extend coverage to check enforced and advisory budgets on the current and subsequent turns, including a zero-remaining-budget reminder in advisory mode. Verify that configuration defaults to enforcing the limit.

GitOrigin-RevId: deeac8dae5714cc7c10d030cade6f3bfbfc8d3a7

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"
-}
+}

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,

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<i64>,
     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,

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<bool> {
         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(

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,

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,

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,

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(())
 }
 

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<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub enabled: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]