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

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<dyn AgentTurnAdmission>,
+    pub(crate) usage_accounting: Option<Arc<dyn AgentUsageAccounting>>,
 }
 
 impl AgentRuntimeServices {
@@ -38,6 +41,10 @@ impl AgentRuntimeServices {
         })?;
         Ok(Self::Host(HostAgentServices {
             turn_admission: Arc::clone(turn_admission.as_ref()),
+            usage_accounting: data
+                .get::<Arc<dyn AgentUsageAccounting>>()
+                .as_deref()
+                .cloned(),
         }))
     }
 }

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,

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?;

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<dyn AgentControl>,
     /// Admission provider for a host-owned child. A lifecycle contributor may also install it.
     pub turn_admission: Option<Arc<dyn AgentTurnAdmission>>,
+    /// 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<Arc<dyn AgentUsageAccounting>>,
 }
 
 /// Parent inputs captured from the tool invocation for a host starting or forking a child.

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<dyn AgentUsageAccounting>` 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<UsageAccountingDecision>>;
+}

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);

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<Vec<AgentUsageReport>>,
+        decision: UsageAccountingDecision,
+    }
+
+    impl crate::AgentTurnAdmission for UsageRecorder {
+        fn admit_turn(
+            &self,
+            _: crate::AgentTurnAdmissionRequest,
+        ) -> futures::future::BoxFuture<'_, CodexResult<crate::AgentExecutionGuard>> {
+            Box::pin(async { Ok(crate::AgentExecutionGuard::new(())) })
+        }
+    }
+
+    impl UsageRecorder {
+        fn install(self: &Arc<Self>, session: &mut Arc<Session>) {
+            let data = &session.services.thread_extension_data;
+            data.insert(crate::AgentRuntimeOwnership::Host);
+            data.insert(self.clone() as Arc<dyn crate::AgentTurnAdmission>);
+            data.insert(self.clone() as Arc<dyn AgentUsageAccounting>);
+            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<UsageAccountingDecision>> {
+            self.reports.lock().unwrap().push(report);
+            Box::pin(async { Ok(self.decision) })
+        }
+    }
+
     fn message(role: &str, text: &str, phase: Option<MessagePhase>) -> 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]

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(

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;

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

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

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,
             },
         };

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<AgentExecutionGuard>> {
+        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<dyn codex_core::AgentTurnAdmission>);
+}
+
+impl codex_extension_api::ThreadLifecycleContributor<codex_core::config::Config> 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<Option<String>>,

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<Mutex<Vec<AgentUsageReport>>>,
+    behavior: HostAccountingBehavior,
+    controller_owned: bool,
+    pause_report: Option<mpsc::UnboundedSender<oneshot::Sender<()>>>,
+}
+
+impl AgentUsageAccounting for HostAccounting {
+    fn report_usage(
+        &self,
+        report: AgentUsageReport,
+    ) -> BoxFuture<'_, codex_protocol::error::Result<UsageAccountingDecision>> {
+        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<Config> 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<dyn AgentUsageAccounting>);
+        })
+    }
+}
+
+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::<Vec<_>>(),
+            [
+                (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<dyn AgentUsageAccounting>);
+    }
 
     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(())
+}

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!(
         (

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::<Config>::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)] {

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