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