mirror of
https://github.com/openai/codex.git
synced 2026-11-14 21:40:39 +00:00
rhan-oai
c725e472e5
## 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
76 lines
2.5 KiB
Rust
76 lines
2.5 KiB
Rust
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;
|
|
|
|
pub(super) async fn maybe_record_reminder(
|
|
sess: &Session,
|
|
turn_context: &TurnContext,
|
|
window_id: &str,
|
|
) {
|
|
// Host runtimes own budget policy, including whether accounting is enabled.
|
|
if matches!(&sess.agent_runtime, AgentRuntimeServices::Host(_)) {
|
|
return;
|
|
}
|
|
let agent_control = &sess.services.agent_control;
|
|
let Some(reminder) = agent_control
|
|
.pending_budget_reminder(sess.thread_id(), window_id)
|
|
.await
|
|
else {
|
|
return;
|
|
};
|
|
let response_item = ContextualUserFragment::into(crate::context::RolloutBudgetContext {
|
|
remaining_tokens: reminder.remaining_tokens,
|
|
});
|
|
sess.record_conversation_items(
|
|
turn_context,
|
|
turn_context.model_info(),
|
|
std::slice::from_ref(&response_item),
|
|
)
|
|
.await;
|
|
agent_control
|
|
.mark_budget_reminder_delivered(sess.thread_id(), window_id, reminder)
|
|
.await;
|
|
}
|
|
|
|
impl Session {
|
|
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),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|