diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 27cdb6f1d1..32e40a80bb 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/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, diff --git a/codex-rs/core/src/agent/control/budget.rs b/codex-rs/core/src/agent/control/budget.rs new file mode 100644 index 0000000000..855072b9c4 --- /dev/null +++ b/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 { + 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); + } +} diff --git a/codex-rs/core/src/guardian/input_budget.rs b/codex-rs/core/src/guardian/input_budget.rs index 5a775632f4..9514eb31e8 100644 --- a/codex-rs/core/src/guardian/input_budget.rs +++ b/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, diff --git a/codex-rs/core/src/session/rollout_budget.rs b/codex-rs/core/src/session/rollout_budget.rs index c5c640927c..55e4ff9ff9 100644 --- a/codex-rs/core/src/session/rollout_budget.rs +++ b/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) } } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 67e6c3a484..8c2df4e169 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/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");