Encapsulate rollout budget accounting in LocalAgentControl (#46559)

## What changed

Replace direct access to the shared rollout budget with controller methods for
recording token usage, retrieving pending reminders, and marking reminders as
delivered. Move the `SessionBudgetExceeded` check into the controller and update
session and guardian callers, preserving reminder acknowledgment after insertion
into conversation history.

## Testing

Replace the shared-budget pointer assertion with a behavioral check: recording
25 output tokens in a reviewer session leaves 75 tokens in the parent session's
budget reminder when the initial budget is 100 tokens.

GitOrigin-RevId: 003bc2b56e526787c1864c09729bbad1006e243d

codex-rs/core/src/agent/control.rs

@@ -70,6 +70,7 @@ pub(crate) use self::interrupt::AgentInterruptError;
 pub(crate) use self::interrupt::AgentInterruptOutcome;
 use self::residency::V2Residency;
 
+mod budget;
 mod completion;
 mod delivery;
 mod execution;
@@ -156,10 +157,6 @@ impl LocalAgentControl {
         (self.thread_id_generator)()
     }
 
-    pub(crate) fn rollout_budget(&self) -> &RolloutBudget {
-        self.rollout_budget.as_ref()
-    }
-
     /// Send rich user input items to an existing agent thread.
     pub(crate) async fn send_input(
         &self,

codex-rs/core/src/agent/control/budget.rs

@@ -0,0 +1,36 @@
+//! Keeps shared rollout-budget accounting and reminder state behind the controller.
+//! Callers acknowledge a reminder only after it has been added to conversation history.
+
+use super::LocalAgentControl;
+use crate::rollout_budget::RolloutBudgetReminder;
+use codex_protocol::ThreadId;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result as CodexResult;
+use codex_protocol::protocol::TokenUsage;
+
+impl LocalAgentControl {
+    pub(crate) fn record_rollout_budget_usage(&self, usage: &TokenUsage) -> CodexResult<()> {
+        if self.rollout_budget.record_usage(usage)? {
+            return Err(CodexErr::SessionBudgetExceeded);
+        }
+        Ok(())
+    }
+
+    pub(crate) fn pending_budget_reminder(
+        &self,
+        thread_id: ThreadId,
+        window_id: &str,
+    ) -> Option<RolloutBudgetReminder> {
+        self.rollout_budget.pending_reminder(thread_id, window_id)
+    }
+
+    pub(crate) fn mark_budget_reminder_delivered(
+        &self,
+        thread_id: ThreadId,
+        window_id: &str,
+        reminder: RolloutBudgetReminder,
+    ) {
+        self.rollout_budget
+            .mark_reminder_delivered(thread_id, window_id, reminder);
+    }
+}

codex-rs/core/src/guardian/input_budget.rs

@@ -123,8 +123,7 @@ pub(crate) async fn finalize(
     if let Some(reminder) = session
         .services
         .agent_control
-        .rollout_budget()
-        .pending_reminder(session.thread_id(), &session.current_window_id().await)
+        .pending_budget_reminder(session.thread_id(), &session.current_window_id().await)
     {
         let reminder = ContextualUserFragment::into(crate::context::RolloutBudgetContext {
             remaining_tokens: reminder.remaining_tokens,

codex-rs/core/src/session/rollout_budget.rs

@@ -1,7 +1,6 @@
 use super::session::Session;
 use super::turn_context::TurnContext;
 use crate::context::ContextualUserFragment;
-use codex_protocol::error::CodexErr;
 use codex_protocol::error::Result as CodexResult;
 use codex_protocol::protocol::TokenUsage;
 
@@ -10,8 +9,8 @@ pub(super) async fn maybe_record_reminder(
     turn_context: &TurnContext,
     window_id: &str,
 ) {
-    let budget = sess.services.agent_control.rollout_budget();
-    let Some(reminder) = budget.pending_reminder(sess.thread_id(), window_id) else {
+    let agent_control = &sess.services.agent_control;
+    let Some(reminder) = agent_control.pending_budget_reminder(sess.thread_id(), window_id) else {
         return;
     };
     let response_item = ContextualUserFragment::into(crate::context::RolloutBudgetContext {
@@ -23,19 +22,13 @@ pub(super) async fn maybe_record_reminder(
         std::slice::from_ref(&response_item),
     )
     .await;
-    budget.mark_reminder_delivered(sess.thread_id(), window_id, reminder);
+    agent_control.mark_budget_reminder_delivered(sess.thread_id(), window_id, reminder);
 }
 
 impl Session {
     pub(crate) fn record_rollout_budget_usage(&self, usage: &TokenUsage) -> CodexResult<()> {
-        if self
-            .services
+        self.services
             .agent_control
-            .rollout_budget()
-            .record_usage(usage)?
-        {
-            return Err(CodexErr::SessionBudgetExceeded);
-        }
-        Ok(())
+            .record_rollout_budget_usage(usage)
     }
 }

codex-rs/core/src/thread_manager_tests.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::agent::types::SpawnAgentOptions;
+use crate::config::RolloutBudgetConfig;
 use crate::config::test_config;
 use crate::init_state_db;
 use crate::installation_id::INSTALLATION_ID_FILENAME;
@@ -38,6 +39,7 @@ use codex_protocol::protocol::SessionMetaLine;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::SubAgentSource;
 use codex_protocol::protocol::ThreadSource;
+use codex_protocol::protocol::TokenUsage;
 use codex_protocol::protocol::TurnStartedEvent;
 use codex_protocol::protocol::UserMessageEvent;
 use codex_protocol::user_input::UserInput;
@@ -1249,6 +1251,12 @@ 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 {
+        limit_tokens: 100,
+        reminder_at_remaining_tokens: vec![75, 50, 25],
+        sampling_token_weight: 1.0,
+        prefill_token_weight: 1.0,
+    });
     config.codex_home = temp_dir.path().join("codex-home").abs();
     config.cwd = config.codex_home.abs();
     std::fs::create_dir_all(&config.codex_home).expect("create codex home");
@@ -1378,20 +1386,24 @@ async fn spawn_internal_session_preserves_parent_lineage_without_forking_history
         reviewer.session_configured.session_id,
         parent.session_configured.session_id
     );
-    assert!(std::ptr::eq(
-        reviewer
-            .thread
-            .session
-            .services
-            .agent_control
-            .rollout_budget(),
-        parent
-            .thread
-            .session
-            .services
-            .agent_control
-            .rollout_budget(),
-    ));
+    reviewer
+        .thread
+        .session
+        .services
+        .agent_control
+        .record_rollout_budget_usage(&TokenUsage {
+            output_tokens: 25,
+            ..Default::default()
+        })
+        .expect("record reviewer usage");
+    let reminder = parent
+        .thread
+        .session
+        .services
+        .agent_control
+        .pending_budget_reminder(parent.thread_id, "window")
+        .expect("parent budget reminder");
+    assert_eq!(reminder.remaining_tokens, 75);
     assert_eq!(reviewer_config.parent_thread_id, Some(parent.thread_id));
     assert_eq!(reviewer_config.forked_from_thread_id, None);
     assert_eq!(reviewer_config.originator, "codex_work_desktop");