From c725e472e594fbdbb9d31af16f81396229efded1 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Mon, 2 Nov 2026 20:17:28 +0000 Subject: [PATCH] Add usage accounting for host-owned agent runtimes (#58975) ## Why Host-owned runtimes skip controller rollout-budget accounting. Give hosts an optional callback to account for inference spending and signal budget exhaustion. ## What changed - Add `AgentUsageAccounting`, `AgentUsageReport`, and `UsageAccountingDecision`, captured at startup and supported by host child spawn reservations. - Await per-response usage reports with thread and turn attribution, including provider budget units and foreground and background compaction usage. - Map host exhaustion to `SessionBudgetExceeded` and accounting failures to fatal errors without retrying inference. Preserve the original error for incomplete responses. - Keep controller accounting unchanged and leave accounting disabled for host-owned runtimes without a provider. ## Testing Add coverage for usage attribution, child accounting, budget exhaustion, nonretryable accounting failures, incomplete responses, and compaction reporting, including background usage charged once to its originating turn. GitOrigin-RevId: 887703f50fe882497e4e8cb33259072935ec82e0 --- codex-rs/core/src/agent/host/mod.rs | 4 + codex-rs/core/src/agent/host/ownership.rs | 7 + codex-rs/core/src/agent/host/spawn.rs | 3 + codex-rs/core/src/agent/host/spawn_tests.rs | 1 + codex-rs/core/src/agent/host/types.rs | 9 +- codex-rs/core/src/agent/host/usage.rs | 40 +++ codex-rs/core/src/compact.rs | 1 + codex-rs/core/src/compact_remote_v2.rs | 114 +++++- .../src/compact_remote_v2_installation.rs | 3 +- codex-rs/core/src/lib.rs | 3 + codex-rs/core/src/session/mod.rs | 4 +- codex-rs/core/src/session/rollout_budget.rs | 42 ++- codex-rs/core/src/session/turn.rs | 3 + .../core/tests/suite/agent_control_tests.rs | 26 ++ codex-rs/core/tests/suite/rollout_budget.rs | 333 ++++++++++++++++-- codex-rs/core/tests/suite/scenarios.rs | 25 +- .../tests/suite/scenarios_budget_context.rs | 11 +- .../tests/suite/scenarios_reserved_spawn.rs | 37 +- 18 files changed, 611 insertions(+), 55 deletions(-) create mode 100644 codex-rs/core/src/agent/host/usage.rs diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index 4e7ba4c663..9a88c60c77 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -4,6 +4,7 @@ mod mcp_metadata; mod ownership; mod spawn; mod types; +mod usage; pub use admission::AgentTurnAdmission; pub use admission::AgentTurnAdmissionRequest; @@ -19,3 +20,6 @@ pub(crate) use ownership::AgentRuntimeServices; pub use types::HostThreadSpawnReservation; pub use types::PreparedChildContext; pub(crate) use types::PreparedChildInheritance; +pub use usage::AgentUsageAccounting; +pub use usage::AgentUsageReport; +pub use usage::UsageAccountingDecision; diff --git a/codex-rs/core/src/agent/host/ownership.rs b/codex-rs/core/src/agent/host/ownership.rs index dea465b675..33005741ce 100644 --- a/codex-rs/core/src/agent/host/ownership.rs +++ b/codex-rs/core/src/agent/host/ownership.rs @@ -1,6 +1,7 @@ //! Captures coordination providers once, after the thread's startup contributors finish. use super::AgentTurnAdmission; +use super::AgentUsageAccounting; use codex_extension_api::ExtensionData; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; @@ -14,6 +15,7 @@ pub enum AgentRuntimeOwnership { Controller, /// Requires an [`AgentTurnAdmission`] provider. Core skips controller rollout-budget /// accounting, budget reminders, and subagent roster context for this runtime. + /// Install an [`AgentUsageAccounting`] provider to receive usage reports. Host, } @@ -24,6 +26,7 @@ pub(crate) enum AgentRuntimeServices { pub(crate) struct HostAgentServices { pub(crate) turn_admission: Arc, + pub(crate) usage_accounting: Option>, } impl AgentRuntimeServices { @@ -38,6 +41,10 @@ impl AgentRuntimeServices { })?; Ok(Self::Host(HostAgentServices { turn_admission: Arc::clone(turn_admission.as_ref()), + usage_accounting: data + .get::>() + .as_deref() + .cloned(), })) } } diff --git a/codex-rs/core/src/agent/host/spawn.rs b/codex-rs/core/src/agent/host/spawn.rs index 975f259e17..b0cffc5a4d 100644 --- a/codex-rs/core/src/agent/host/spawn.rs +++ b/codex-rs/core/src/agent/host/spawn.rs @@ -135,6 +135,9 @@ pub(super) async fn create_spawn( if let Some(admission) = reservation.turn_admission { options.thread_extension_init.insert(admission); } + if let Some(accounting) = reservation.usage_accounting { + options.thread_extension_init.insert(accounting); + } let child = state .spawn_host_child_thread( options, diff --git a/codex-rs/core/src/agent/host/spawn_tests.rs b/codex-rs/core/src/agent/host/spawn_tests.rs index 4ae544a30b..31bb8cf03c 100644 --- a/codex-rs/core/src/agent/host/spawn_tests.rs +++ b/codex-rs/core/src/agent/host/spawn_tests.rs @@ -199,6 +199,7 @@ async fn same_manager_creation_retains_step_executor_and_shares_exec_policy() -> thread_id: child_id, control: Arc::clone(&session.services.agent_control), turn_admission: None, + usage_accounting: None, }, ) .await?; diff --git a/codex-rs/core/src/agent/host/types.rs b/codex-rs/core/src/agent/host/types.rs index 758e49300b..f96d94a7b1 100644 --- a/codex-rs/core/src/agent/host/types.rs +++ b/codex-rs/core/src/agent/host/types.rs @@ -1,6 +1,7 @@ //! Prepared inputs and host reservations for native child startup. use super::admission::AgentTurnAdmission; +use super::usage::AgentUsageAccounting; use crate::agent::api::AgentControl; use crate::agents_md_manager::SessionInstructions; use crate::exec_policy::ExecPolicyManager; @@ -19,13 +20,17 @@ use std::sync::Arc; /// A fresh child ID reserved in the host's tree. /// The host commits membership, serializes startup attempts and deduplicates retries. /// Core derives the parent, session and path from the invocation and child request. -/// The controller supplies runtime coordination by default. For host-owned admission, select -/// `AgentRuntimeOwnership::Host` before startup finishes. +/// The controller supplies runtime coordination by default. To use host-owned admission and +/// accounting, a startup contributor must explicitly select `AgentRuntimeOwnership::Host`. pub struct HostThreadSpawnReservation { pub thread_id: ThreadId, pub control: Arc, /// Admission provider for a host-owned child. A lifecycle contributor may also install it. pub turn_admission: Option>, + /// Optional accounting for a host-owned child. Without a provider from the reservation or + /// a contributor, host-owned children disable accounting. Controller-owned children keep + /// their controller's accounting. + pub usage_accounting: Option>, } /// Parent inputs captured from the tool invocation for a host starting or forking a child. diff --git a/codex-rs/core/src/agent/host/usage.rs b/codex-rs/core/src/agent/host/usage.rs new file mode 100644 index 0000000000..073e15bbaf --- /dev/null +++ b/codex-rs/core/src/agent/host/usage.rs @@ -0,0 +1,40 @@ +//! Hands observed inference spending to one host-owned accounting authority. + +use codex_protocol::ThreadId; +use codex_protocol::error::Result; +use codex_protocol::protocol::TokenUsage; +use futures::future::BoxFuture; + +/// Provider-reported spending for one response, including compaction. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentUsageReport { + pub thread_id: ThreadId, + pub turn_id: String, + /// Per-response usage, not cumulative totals. Remote adapters must explicitly + /// carry `codex_rollout_budget_units`, which `TokenUsage` skips when serialized. + pub usage: TokenUsage, +} + +/// Both decisions mean the host recorded the reported spending. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum UsageAccountingDecision { + Continue, + Exhausted, +} + +/// Replaces this runtime's `AgentControl::record_usage` and native rollout-budget reminders. +/// +/// Install an `Arc` in thread extension data before host-owned startup; +/// supply reminder context through context contributors. Core awaits each report. The host +/// owns budget policy, timeouts, report-only retries, reporting IDs, deduplication and recovery. +/// A turn may emit multiple reports; thread/turn IDs give attribution, not unique charges. +/// `Err` may leave acceptance unknown and causes a nonretryable failure without inference +/// retry or Local fallback. Incomplete responses retain their original error; background +/// compaction errors fail that job. This callback neither reserves future spending nor +/// interrupts other agents in the tree. +pub trait AgentUsageAccounting: Send + Sync { + fn report_usage( + &self, + report: AgentUsageReport, + ) -> BoxFuture<'_, Result>; +} diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index cabf39001d..46a5eb5233 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -358,6 +358,7 @@ async fn run_compact_task_inner_impl( e.details(), CodexErrorDetails::SessionBudgetExceeded | CodexErrorDetails::ResponseIncomplete(_) + | CodexErrorDetails::Fatal(_) ) => { return Err(e); diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 10d04d26d9..76fc4c18fe 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -402,7 +402,9 @@ async fn collect_compaction_output( CompactionRequestKind::Background(_) => { if let Some(usage) = incomplete.token_usage.as_ref() { // Match completed background requests without replacing the terminal error. - let _ = sess.record_rollout_budget_usage(usage).await; + let _ = sess + .record_rollout_budget_usage(&turn_context.sub_id, usage) + .await; } } CompactionRequestKind::Foreground => { @@ -435,7 +437,8 @@ async fn collect_compaction_output( match request_kind { CompactionRequestKind::Background(_) => { if let Some(usage) = token_usage.as_ref() { - sess.record_rollout_budget_usage(usage).await?; + sess.record_rollout_budget_usage(&turn_context.sub_id, usage) + .await?; } } CompactionRequestKind::Foreground => { @@ -773,14 +776,54 @@ fn truncate_message_text_to_token_budget( #[cfg(test)] mod tests { use super::*; + use crate::agent::host::AgentUsageAccounting; + use crate::agent::host::AgentUsageReport; + use crate::agent::host::UsageAccountingDecision; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::MessagePhase; use pretty_assertions::assert_eq; + use std::sync::Mutex; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; + struct UsageRecorder { + reports: Mutex>, + decision: UsageAccountingDecision, + } + + impl crate::AgentTurnAdmission for UsageRecorder { + fn admit_turn( + &self, + _: crate::AgentTurnAdmissionRequest, + ) -> futures::future::BoxFuture<'_, CodexResult> { + Box::pin(async { Ok(crate::AgentExecutionGuard::new(())) }) + } + } + + impl UsageRecorder { + fn install(self: &Arc, session: &mut Arc) { + let data = &session.services.thread_extension_data; + data.insert(crate::AgentRuntimeOwnership::Host); + data.insert(self.clone() as Arc); + data.insert(self.clone() as Arc); + let runtime = crate::agent::host::AgentRuntimeServices::from_extensions(data) + .expect("valid host providers"); + Arc::get_mut(session).expect("unique session").agent_runtime = runtime; + } + } + + impl AgentUsageAccounting for UsageRecorder { + fn report_usage( + &self, + report: AgentUsageReport, + ) -> futures::future::BoxFuture<'_, CodexResult> { + self.reports.lock().unwrap().push(report); + Box::pin(async { Ok(self.decision) }) + } + } + fn message(role: &str, text: &str, phase: Option) -> ResponseItem { ResponseItem::Message { status: None, @@ -1244,23 +1287,29 @@ mod tests { #[tokio::test] async fn collect_background_incomplete_does_not_publish_foreground_usage() { + let usage = TokenUsage { + input_tokens: 100, + output_tokens: 20, + total_tokens: 120, + ..Default::default() + }; let stream = response_stream(vec![Err(CodexErr::new( CodexErrorDetails::ResponseIncomplete(Box::new( codex_protocol::error::ResponseIncompleteError { partial_output: vec![], reason: "max_output_tokens".to_string(), response_id: Some("resp-background".to_string()), - token_usage: Some(TokenUsage { - input_tokens: 100, - output_tokens: 20, - total_tokens: 120, - ..Default::default() - }), + token_usage: Some(usage.clone()), }, )), ))]); - let (sess, turn_context, rx) = + let (mut sess, turn_context, rx) = crate::session::tests::make_session_and_context_with_rx().await; + let accounting = Arc::new(UsageRecorder { + reports: Mutex::new(Vec::new()), + decision: UsageAccountingDecision::Exhausted, + }); + accounting.install(&mut sess); let cancellation = CancellationToken::new(); cancellation.cancel(); @@ -1276,6 +1325,53 @@ mod tests { if matches!(err.details(), CodexErrorDetails::ResponseIncomplete(_)))); assert_eq!(sess.token_usage_info().await, None); assert!(rx.is_empty()); + assert_eq!( + *accounting.reports.lock().unwrap(), + vec![AgentUsageReport { + thread_id: sess.thread_id(), + turn_id: turn_context.sub_id.clone(), + usage, + }] + ); + } + + #[tokio::test] + async fn background_compaction_charges_usage_even_if_output_is_missing() { + let (mut sess, turn, _) = crate::session::tests::make_session_and_context_with_rx().await; + let accounting = Arc::new(UsageRecorder { + reports: Mutex::new(Vec::new()), + decision: UsageAccountingDecision::Continue, + }); + accounting.install(&mut sess); + let usage = TokenUsage { + output_tokens: 42, + ..Default::default() + }; + let stream = response_stream(vec![Ok(ResponseEvent::Completed { + response_id: "resp-compact".into(), + token_usage: Some(usage.clone()), + usage_metadata: None, + end_turn: Some(true), + })]); + let result = collect_compaction_output( + &sess, + &turn, + stream, + CompactionRequestKind::Background(&CancellationToken::new()), + ) + .await; + + assert!( + matches!(result, Err(error) if matches!(error.details(), CodexErrorDetails::Fatal(_))) + ); + assert_eq!( + *accounting.reports.lock().unwrap(), + [AgentUsageReport { + thread_id: sess.thread_id(), + turn_id: turn.sub_id.clone(), + usage, + }] + ); } #[tokio::test] diff --git a/codex-rs/core/src/compact_remote_v2_installation.rs b/codex-rs/core/src/compact_remote_v2_installation.rs index 8397c5d402..aa3726491f 100644 --- a/codex-rs/core/src/compact_remote_v2_installation.rs +++ b/codex-rs/core/src/compact_remote_v2_installation.rs @@ -186,7 +186,8 @@ async fn build_and_install_compaction( if prefix.is_none() && let Some(token_usage) = token_usage.as_ref() { - sess.record_rollout_budget_usage(token_usage).await?; + sess.record_rollout_budget_usage(&compaction_turn_context.sub_id, token_usage) + .await?; analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); } let (compacted_history, retained_images) = build_v2_compacted_history( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e8db8a87c0..7e8fb8763b 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -160,10 +160,13 @@ pub use agent::host::AgentMcpMetadataRecipient; pub use agent::host::AgentRuntimeOwnership; pub use agent::host::AgentTurnAdmission; pub use agent::host::AgentTurnAdmissionRequest; +pub use agent::host::AgentUsageAccounting; +pub use agent::host::AgentUsageReport; pub use agent::host::HostThreadSpawnReservation; pub use agent::host::HostToolInvocationLookup; pub use agent::host::PreparedChildContext; pub use agent::host::ToolInvocationContext; +pub use agent::host::UsageAccountingDecision; pub use thread_spawn::ThreadCreateRequest; pub use thread_spawn::ThreadSpawnError; pub use thread_spawn::ThreadSpawnResult; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f2ee393b24..ff0872efa3 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -5462,7 +5462,9 @@ impl Session { token_usage, ); } - let budget_result = self.record_rollout_budget_usage(token_usage).await; + let budget_result = self + .record_rollout_budget_usage(&turn_context.sub_id, token_usage) + .await; if let Some(token_info) = token_info.as_ref() { for contributor in self.services.extensions.token_usage_contributors() { contributor diff --git a/codex-rs/core/src/session/rollout_budget.rs b/codex-rs/core/src/session/rollout_budget.rs index 6fae368a19..bc05663dc1 100644 --- a/codex-rs/core/src/session/rollout_budget.rs +++ b/codex-rs/core/src/session/rollout_budget.rs @@ -1,7 +1,10 @@ use super::session::Session; use super::turn_context::TurnContext; use crate::agent::host::AgentRuntimeServices; +use crate::agent::host::AgentUsageReport; +use crate::agent::host::UsageAccountingDecision; use crate::context::ContextualUserFragment; +use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::TokenUsage; @@ -10,6 +13,7 @@ pub(super) async fn maybe_record_reminder( turn_context: &TurnContext, window_id: &str, ) { + // Host runtimes own budget policy, including whether accounting is enabled. if matches!(&sess.agent_runtime, AgentRuntimeServices::Host(_)) { return; } @@ -35,13 +39,37 @@ pub(super) async fn maybe_record_reminder( } impl Session { - pub(crate) async fn record_rollout_budget_usage(&self, usage: &TokenUsage) -> CodexResult<()> { - if matches!(&self.agent_runtime, AgentRuntimeServices::Host(_)) { - return Ok(()); + pub(crate) async fn record_rollout_budget_usage( + &self, + turn_id: &str, + usage: &TokenUsage, + ) -> CodexResult<()> { + match &self.agent_runtime { + AgentRuntimeServices::Controller => { + self.services + .agent_control + .record_usage(usage.clone()) + .await + } + AgentRuntimeServices::Host(services) => { + let Some(accounting) = &services.usage_accounting else { + return Ok(()); + }; + let decision = accounting + .report_usage(AgentUsageReport { + thread_id: self.thread_id(), + turn_id: turn_id.to_owned(), + usage: usage.clone(), + }) + .await + .map_err(|error| { + CodexErr::Fatal(format!("Host usage accounting failed: {error}")) + })?; + match decision { + UsageAccountingDecision::Continue => Ok(()), + UsageAccountingDecision::Exhausted => Err(CodexErr::SessionBudgetExceeded), + } + } } - self.services - .agent_control - .record_usage(usage.clone()) - .await } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index ace2f54c78..c2f2a9b08f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1905,6 +1905,9 @@ async fn run_sampling_request( } return Err(err); } + CodexErrorDetails::SessionBudgetExceeded | CodexErrorDetails::Fatal(_) => { + return Err(err); + } _ => err, }, }; diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index f223133639..c6da8fa8f0 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -57,6 +57,32 @@ use serde_json::json; use test_case::test_case; use wiremock::MockServer; +pub(super) struct HostRuntime; + +impl codex_core::AgentTurnAdmission for HostRuntime { + fn admit_turn( + &self, + _: codex_core::AgentTurnAdmissionRequest, + ) -> BoxFuture<'_, CodexResult> { + Box::pin(async { Ok(AgentExecutionGuard::new(())) }) + } +} + +pub(super) fn install_host_runtime(data: &codex_extension_api::ExtensionData) { + data.insert(codex_core::AgentRuntimeOwnership::Host); + data.get_or_init(|| Arc::new(HostRuntime) as Arc); +} + +impl codex_extension_api::ThreadLifecycleContributor for HostRuntime { + fn on_thread_start<'a>( + &'a self, + input: codex_extension_api::ThreadStartInput<'a, codex_core::config::Config>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + install_host_runtime(input.thread_store); + Box::pin(async {}) + } +} + pub(super) struct TestAgentControl { thread_id: ThreadId, service_tier: Mutex>, diff --git a/codex-rs/core/tests/suite/rollout_budget.rs b/codex-rs/core/tests/suite/rollout_budget.rs index c0c244ab45..082dd2fc42 100644 --- a/codex-rs/core/tests/suite/rollout_budget.rs +++ b/codex-rs/core/tests/suite/rollout_budget.rs @@ -1,11 +1,25 @@ +use super::agent_control::HostRuntime; +use super::agent_control::install_host_runtime; use anyhow::Result; +use codex_core::AgentRuntimeOwnership; +use codex_core::AgentUsageAccounting; +use codex_core::AgentUsageReport; use codex_core::TurnInputRequest; +use codex_core::TurnInputSubmission; +use codex_core::UsageAccountingDecision; +use codex_core::config::Config; use codex_core::config::RolloutBudgetConfig; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; +use codex_protocol::error::CodexErr; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; use core_test_support::responses::ResponsesRequest; use core_test_support::responses::ev_assistant_message; @@ -13,22 +27,85 @@ use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::received_responses_requests; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use core_test_support::wait_for_event_match; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use test_case::test_case; +use tokio::sync::mpsc; +use tokio::sync::oneshot; use tokio::time::timeout; const MULTI_AGENT_V2_NAMESPACE: &str = "collaboration"; -fn rollout_budget() -> RolloutBudgetConfig { +#[derive(Clone, Copy, Default)] +enum HostAccountingBehavior { + #[default] + Continue, + Exhausted, + Fail, +} + +#[derive(Clone, Default)] +pub(super) struct HostAccounting { + pub(super) reports: Arc>>, + behavior: HostAccountingBehavior, + controller_owned: bool, + pause_report: Option>>, +} + +impl AgentUsageAccounting for HostAccounting { + fn report_usage( + &self, + report: AgentUsageReport, + ) -> BoxFuture<'_, codex_protocol::error::Result> { + Box::pin(async move { + self.reports.lock().expect("usage reports").push(report); + if let Some(pause) = &self.pause_report { + let (resume, resumed) = oneshot::channel(); + let _ = pause.send(resume); + let _ = resumed.await; + } + match self.behavior { + HostAccountingBehavior::Continue => Ok(UsageAccountingDecision::Continue), + HostAccountingBehavior::Exhausted => Ok(UsageAccountingDecision::Exhausted), + HostAccountingBehavior::Fail => { + Err(CodexErr::Stream("accounting unavailable".into())) + } + } + }) + } +} + +impl ThreadLifecycleContributor for HostAccounting { + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if !self.controller_owned { + install_host_runtime(input.thread_store); + } + input + .thread_store + .insert(Arc::new(self.clone()) as Arc); + }) + } +} + +pub(super) fn rollout_budget() -> RolloutBudgetConfig { RolloutBudgetConfig { enforce_limit: true, limit_tokens: 100, @@ -174,7 +251,9 @@ async fn invalid_provider_rollout_budget_units_fail_without_retry() -> Result<() } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { +#[test_case(false ; "local accounting")] +#[test_case(true ; "host accounting")] +async fn subagent_usage_draws_from_the_shared_budget(host_accounting: bool) -> Result<()> { skip_if_no_network!(Ok(())); const ROOT_PROMPT: &str = "spawn a budget worker"; @@ -231,8 +310,15 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { ) .await; + let accounting = Arc::new(HostAccounting { + controller_owned: !host_accounting, + ..Default::default() + }); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(accounting.clone()); let test = test_codex() - .with_config(|config| { + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { config .features .enable(Feature::Collab) @@ -241,7 +327,11 @@ 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(RolloutBudgetConfig { + // Host success must not also charge the Local budget. + limit_tokens: if host_accounting { 1 } else { 100 }, + ..rollout_budget() + }); }) .build(&server) .await?; @@ -250,10 +340,30 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { test.submit_turn(ROOT_PROMPT).await?; let child_thread_id = timeout(Duration::from_secs(10), created_threads.recv()).await??; let child_thread = test.thread_manager.get_thread(child_thread_id).await?; - wait_for_event(child_thread.as_ref(), |event| { - matches!(event, EventMsg::TurnComplete(_)) + let child_turn_id = wait_for_event_match(child_thread.as_ref(), |event| match event { + EventMsg::TurnComplete(completed) => Some(completed.turn_id.clone()), + _ => None, }) .await; + if host_accounting { + let mut reports = accounting.reports.lock().expect("usage reports").clone(); + reports.sort_by_key(|report| report.usage.total_tokens); + let root_thread_id = test.session_configured.thread_id; + assert_eq!( + reports + .iter() + .map(|report| (report.thread_id, report.usage.total_tokens)) + .collect::>(), + [ + (root_thread_id, 10), + (root_thread_id, 10), + (child_thread_id, 30) + ] + ); + assert!(!reports[0].turn_id.is_empty()); + assert_eq!(reports[0].turn_id, reports[1].turn_id); + assert_eq!(reports[2].turn_id, child_turn_id); + } test.submit_turn(FOLLOW_UP_PROMPT).await?; let requests = follow_up @@ -269,18 +379,25 @@ async fn subagent_usage_draws_from_the_shared_budget() -> Result<()> { let [request] = requests.as_slice() else { anyhow::bail!("expected 1 follow-up request, got {}", requests.len()); }; - assert_eq!( - rollout_budget_texts(request).last(), - Some(&rollout_budget_message(/*remaining_tokens*/ 50)) - ); + if host_accounting { + assert!(rollout_budget_texts(request).is_empty()); + } else { + assert_eq!( + rollout_budget_texts(request).last(), + Some(&rollout_budget_message(/*remaining_tokens*/ 50)) + ); + assert!(accounting.reports.lock().expect("usage reports").is_empty()); + } Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[test_case(true ; "enforced")] -#[test_case(false ; "advisory")] +#[test_case(AgentRuntimeOwnership::Controller, true ; "enforced")] +#[test_case(AgentRuntimeOwnership::Controller, false ; "advisory")] +#[test_case(AgentRuntimeOwnership::Host, true ; "host without accounting")] async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( + owner: AgentRuntimeOwnership, enforce_limit: bool, ) -> Result<()> { skip_if_no_network!(Ok(())); @@ -300,7 +417,12 @@ async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( ], ) .await; + let mut extensions = ExtensionRegistryBuilder::new(); + if owner == AgentRuntimeOwnership::Host { + extensions.thread_lifecycle_contributor(Arc::new(HostRuntime)); + } let test = test_codex() + .with_extensions(Arc::new(extensions.build())) .with_config(move |config| { config.rollout_budget = Some(RolloutBudgetConfig { enforce_limit, @@ -309,8 +431,15 @@ async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( ..rollout_budget() }); }) - .build(&server) + .build_with_auto_env(&server) .await?; + let late = Arc::new(HostAccounting::default()); + if owner == AgentRuntimeOwnership::Host { + // Neither the owner nor the accounting provider can change after startup. + let data = test.codex.thread_extension_data(); + data.insert(AgentRuntimeOwnership::Controller); + data.insert(late.clone() as Arc); + } for prompt in ["exhaust the budget", "try another turn"] { test.codex @@ -324,8 +453,8 @@ async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( matches!(event, EventMsg::Error(_) | EventMsg::TurnComplete(_)) }) .await; - if !enforce_limit { - assert!(matches!(event, EventMsg::TurnComplete(_))); + if owner == AgentRuntimeOwnership::Host || !enforce_limit { + assert!(matches!(event, EventMsg::TurnComplete(done) if done.error.is_none())); continue; } assert!(matches!( @@ -339,7 +468,14 @@ async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( .await; } - if !enforce_limit { + if owner == AgentRuntimeOwnership::Host { + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + for request in &requests { + assert!(rollout_budget_texts(request).is_empty()); + } + assert!(late.reports.lock().expect("usage reports").is_empty()); + } else if !enforce_limit { assert_eq!( rollout_budget_texts(&responses.requests()[1]).last(), Some(&rollout_budget_message(/*remaining_tokens*/ 0)) @@ -349,13 +485,15 @@ async fn exhausted_budget_respects_enforcement_on_current_and_later_turns( } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -#[test_case(false, false ; "local token usage")] -#[test_case(false, true ; "local provider units")] -#[test_case(true, false ; "remote v2 token usage")] -#[test_case(true, true ; "remote v2 provider units")] -async fn compaction_budget_exhaustion_fails_without_retry( +#[test_case(false, false, false ; "local token usage")] +#[test_case(false, true, false ; "local provider units")] +#[test_case(true, false, false ; "remote v2 token usage")] +#[test_case(true, true, false ; "remote v2 provider units")] +#[test_case(false, false, true ; "local compaction host accounting failure")] +async fn compaction_accounting_errors_fail_without_retry( remote_v2: bool, provider_units: bool, + host_failure: bool, ) -> Result<()> { skip_if_no_network!(Ok(())); @@ -386,8 +524,18 @@ async fn compaction_budget_exhaustion_fails_without_retry( ]) }; let responses = mount_sse_sequence(&server, vec![compact_response]).await; + let accounting = Arc::new(HostAccounting { + behavior: HostAccountingBehavior::Fail, + ..Default::default() + }); + let mut extensions = ExtensionRegistryBuilder::new(); + if host_failure { + extensions.thread_lifecycle_contributor(accounting.clone()); + } let test = test_codex() + .with_extensions(Arc::new(extensions.build())) .with_config(move |config| { + config.model_provider.stream_max_retries = Some(2); config.rollout_budget = Some(RolloutBudgetConfig { limit_tokens: 10, reminder_at_remaining_tokens: vec![5], @@ -400,12 +548,21 @@ async fn compaction_budget_exhaustion_fails_without_retry( .build(&server) .await?; + let expected_error = if host_failure { + CodexErrorInfo::Other + } else { + CodexErrorInfo::SessionBudgetExceeded + }; test.codex.submit(Op::Compact).await?; wait_for_event(&test.codex, |event| { + assert!( + !matches!(event, EventMsg::StreamError(_)), + "compaction must not retry accounting failures" + ); matches!( event, EventMsg::Error(error) - if error.codex_error_info == Some(CodexErrorInfo::SessionBudgetExceeded) + if error.codex_error_info == Some(expected_error.clone()) ) }) .await; @@ -414,6 +571,12 @@ async fn compaction_budget_exhaustion_fails_without_retry( }) .await; assert_eq!(responses.requests().len(), 1, "compaction should not retry"); + if host_failure { + let reports = accounting.reports.lock().expect("usage reports"); + assert_eq!(reports.len(), 1); + assert_eq!(reports[0].thread_id, test.session_configured.thread_id); + assert_eq!(reports[0].usage.total_tokens, 10); + } Ok(()) } @@ -482,3 +645,129 @@ async fn restates_the_current_remainder_after_compaction() -> Result<()> { Ok(()) } + +#[tokio::test] +#[test_case(HostAccountingBehavior::Exhausted, false ; "host exhausts budget")] +#[test_case(HostAccountingBehavior::Fail, false ; "accounting failure never retries inference")] +#[test_case(HostAccountingBehavior::Fail, true ; "incomplete response keeps original error")] +async fn host_accounting_reports_raw_usage_and_preserves_error_semantics( + behavior: HostAccountingBehavior, + incomplete: bool, +) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let mut response = json!({ + "usage": { + "input_tokens": 60, + "input_tokens_details": {"cached_tokens": 40}, + "output_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 5}, + "total_tokens": 75, + "codex_rollout_budget_units": 40.5 + } + }); + let response_type = if incomplete { + response["incomplete_details"] = json!({"reason": "max_output_tokens"}); + "response.incomplete" + } else { + response["id"] = json!("host-accounted"); + "response.completed" + }; + mount_sse_once( + &server, + sse(vec![json!({"type": response_type, "response": response})]), + ) + .await; + let (pause_report, mut paused) = mpsc::unbounded_channel(); + let accounting = Arc::new(HostAccounting { + behavior, + pause_report: (!incomplete).then_some(pause_report), + ..Default::default() + }); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(accounting.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config.features.enable(Feature::InstantInterrupt).unwrap(); + config.model_provider.stream_max_retries = Some(2); + config.model_provider.request_max_retries = Some(0); + config.rollout_budget = None; + }) + .build_with_auto_env(&server) + .await?; + let user_input = |text: &str| { + TurnInputRequest::user_input(vec![UserInput::Text { + text: text.into(), + text_elements: Vec::new(), + }]) + }; + test.codex + .start_or_steer_turn(user_input("report spending")) + .await?; + if !incomplete { + let resume = timeout(Duration::from_secs(10), paused.recv()) + .await? + .expect("accounting report"); + let submission = test + .codex + .start_or_steer_turn(user_input("continue")) + .await?; + assert!(matches!(submission, TurnInputSubmission::Steered { .. })); + // Virtual time advances once the runnable input watcher has processed the steer. + tokio::time::pause(); + tokio::time::sleep(Duration::from_millis(1)).await; + tokio::time::resume(); + resume.send(()).expect("resume accounting"); + } + drop(paused); + let mut errors = Vec::new(); + let completed = wait_for_event(&test.codex, |event| { + if let EventMsg::Error(error) = event { + errors.push(error.clone()); + } + assert!(!matches!(event, EventMsg::StreamError(_)), "{event:?}"); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let EventMsg::TurnComplete(completed) = completed else { + unreachable!(); + }; + let [error] = errors.as_slice() else { + anyhow::bail!("expected one turn error, got {errors:?}"); + }; + if incomplete { + assert_eq!( + error.message, + "Incomplete response returned, reason: max_output_tokens" + ); + } else if matches!(behavior, HostAccountingBehavior::Exhausted) { + assert_eq!( + error.codex_error_info, + Some(CodexErrorInfo::SessionBudgetExceeded) + ); + } else { + assert!(error.message.contains("Host usage accounting failed")); + assert_eq!(error.codex_error_info, Some(CodexErrorInfo::Other)); + } + assert_eq!(received_responses_requests(&server).await.len(), 1); + assert_eq!( + *accounting.reports.lock().expect("usage reports"), + vec![AgentUsageReport { + thread_id: test.session_configured.thread_id, + turn_id: completed.turn_id, + usage: TokenUsage { + input_tokens: 60, + cached_input_tokens: 40, + output_tokens: 15, + reasoning_output_tokens: 5, + total_tokens: 75, + codex_rollout_budget_units: Some( + serde_json::Number::from_f64(40.5).expect("finite budget units"), + ), + ..Default::default() + }, + }] + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 68251d5494..dc62ed14e4 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -3672,7 +3672,9 @@ async fn async_compaction_reuses_pending_work( outcome: PendingCompactionOutcome, trigger: CompactionTrigger, ) -> Result<()> { - use core_test_support::responses::ev_completed_with_tokens; + use super::rollout_budget::HostAccounting; + + const BACKGROUND_TOKENS: i64 = 3_001; skip_if_no_network!(Ok(())); let over_budget = outcome == PendingCompactionOutcome::SuccessOverBudget; @@ -3716,7 +3718,7 @@ async fn async_compaction_reuses_pending_work( } else { ev_assistant_message("invalid", "missing compaction output") }, - /*tokens*/ 3_000, + BACKGROUND_TOKENS, ), )]); } @@ -3740,8 +3742,16 @@ async fn async_compaction_reuses_pending_work( /*tokens*/ if over_budget { 10 } else { 3_000 }, ))); let (streaming, mut completions) = start_streaming_sse_server(responses).await; + let accounting = + (outcome == PendingCompactionOutcome::Success).then(|| Arc::new(HostAccounting::default())); + let mut builder = test_codex(); + if let Some(accounting) = &accounting { + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(accounting.clone()); + builder = builder.with_extensions(Arc::new(extensions.build())); + } let (_config_server, test) = async_compaction_test( - test_codex(), + builder, streaming.uri(), if over_budget { vec![Feature::RetainClientToolOutputs] @@ -3875,6 +3885,15 @@ async fn async_compaction_reuses_pending_work( assert!(history.items.iter().all(|item| !matches!(item, RolloutItem::TokenUsageRecord(record) if record.response_id == "background" ))); + if let Some(accounting) = accounting { + let reports = accounting.reports.lock().expect("usage reports"); + let charged_turns: Vec<_> = reports + .iter() + .filter(|report| report.usage.total_tokens == BACKGROUND_TOKENS) + .map(|report| &report.turn_id) + .collect(); + assert_eq!(charged_turns, [&usage("seed").turn_id]); + } let continuation = usage("continuation"); assert_eq!( ( diff --git a/codex-rs/core/tests/suite/scenarios_budget_context.rs b/codex-rs/core/tests/suite/scenarios_budget_context.rs index 433eb006a6..59c262485f 100644 --- a/codex-rs/core/tests/suite/scenarios_budget_context.rs +++ b/codex-rs/core/tests/suite/scenarios_budget_context.rs @@ -1,5 +1,7 @@ //! Host budget reminders use generic World State and refresh after compaction spending. +use super::super::rollout_budget::HostAccounting; +use super::super::rollout_budget::rollout_budget; use super::*; use codex_extension_api::ContextContributor; use codex_extension_api::ExtensionFuture; @@ -107,10 +109,10 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - // This fixture supplies host-owned policy and controlled shared spending. It does - // not implement usage reporting or configure Local's native rollout budget. + // Host accounting and context must override the configured native budget. let remaining = Arc::new(AtomicI64::new(1_000)); let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.thread_lifecycle_contributor(Arc::new(HostAccounting::default())); extensions.prompt_contributor(Arc::new(HostBudgetContext { remaining: Arc::clone(&remaining), })); @@ -121,7 +123,7 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { .with_config(|config| { configure_scenario_catalog(config); config.workspace_roots = vec![config.cwd.clone()]; - config.rollout_budget = None; + config.rollout_budget = Some(rollout_budget()); }) .build_with_auto_env(&server) .await?; @@ -185,6 +187,9 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { let requests = received_responses_requests(&server).await; assert_eq!(requests.len(), 6); + for request in &requests { + assert!(!request.has_content_kinds(&["rollout_budget.remaining_tokens"])); + } // Compare the actual wire input without normalization: both unchanged and // threshold-crossing reminders must only append within a context window. for (previous, next) in [(0, 1), (1, 2), (4, 5)] { diff --git a/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs index b281340d64..5e2a1125ba 100644 --- a/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs +++ b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs @@ -2,6 +2,7 @@ use super::super::agent_control::TestAgentControl; use super::super::agents_md::RecordingThreadInstructionsProvider; +use super::super::rollout_budget::HostAccounting; use super::*; use anyhow::Context; use codex_core::AgentControl; @@ -181,7 +182,15 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( let provider = Arc::new(RecordingThreadInstructionsProvider::with_text(INITIAL_RULE).shared()); let (test, parent, controller) = host_spawn_fixture( &server, - test_codex().with_config(|config| config.allow_persistent_dynamic_tool_updates = true), + test_codex() + .with_extensions({ + let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(Arc::new( + super::super::agent_control::HostRuntime, + )); + Arc::new(extensions.build()) + }) + .with_config(|config| config.allow_persistent_dynamic_tool_updates = true), network_policy.clone(), Some(provider.clone()), ) @@ -226,10 +235,12 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( ); let child_id = ThreadId::new(); let request = || create_request("worker", fork_history); + let accounting = Arc::new(HostAccounting::default()); let reservation = || HostThreadSpawnReservation { thread_id: child_id, control: controller.clone(), turn_admission: None, + usage_accounting: Some(accounting.clone()), }; let mut mismatched = reservation(); mismatched.control = Arc::new(TestAgentControl::new(ThreadId::new())); @@ -264,18 +275,29 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( ); // Once creation returns, ending the spawning invocation leaves the idle child available. finish_spawn_tool(&server, &parent, &invocation).await?; - let child_request = - responses::mount_sse_once(&server, responses::sse_completed("worker-done")).await; + let child_request = responses::mount_sse_once( + &server, + sse(vec![ev_completed_with_tokens( + "worker-done", + /*total_tokens*/ 42, + )]), + ) + .await; let submission = worker .start_turn_if_idle(TurnInputRequest::user_input(vec![text( "Report your status.", )])) .await?; - assert!( - matches!(submission, StartIfIdleSubmission::Started { .. }), - "expected the new child to be idle, got {submission:?}" - ); + let StartIfIdleSubmission::Started { turn_id, .. } = submission else { + panic!("expected the new child to be idle, got {submission:?}"); + }; wait_for_successful_turn(&worker).await; + assert!(matches!( + accounting.reports.lock().expect("usage reports").as_slice(), + [report] if report.thread_id == child_id + && report.turn_id == turn_id + && report.usage.total_tokens == 42 + )); let child_request = child_request.single_request(); assert!(child_request.body_contains_text(UPDATED_RULE)); assert_eq!( @@ -353,6 +375,7 @@ async fn cancelled_creation_cleans_registered_child() -> Result<()> { thread_id: child_id, control: controller, turn_admission: None, + usage_accounting: None, }, ) .await