use std::time::{Duration, Instant};
use agent_core::BudgetDimension;
pub const TOOL_CALL_NOT_EXECUTED: &str = "Tool call not executed";
pub const TOOL_CALL_BUDGET_REFUSAL: &str =
"Tool call not executed: turn tool-call budget exhausted";
#[derive(Debug, Clone, PartialEq)]
pub struct TurnBudget {
pub max_provider_rounds: Option<u32>,
pub max_tool_calls: Option<u32>,
pub max_elapsed: Option<Duration>,
pub max_accumulated_tool_result_bytes: Option<usize>,
pub max_context_tokens: Option<u64>,
pub max_cost_usd: Option<f64>,
pub max_total_tokens: Option<u64>,
pub max_round_renewals: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TurnRole {
Foreground,
Autonomous,
Event,
Worker,
}
impl TurnBudget {
pub fn for_role(role: TurnRole) -> Self {
match role {
TurnRole::Foreground => Self {
max_provider_rounds: None,
max_tool_calls: None,
max_elapsed: None,
max_accumulated_tool_result_bytes: None,
max_context_tokens: None,
max_cost_usd: None,
max_total_tokens: None,
max_round_renewals: 8,
},
TurnRole::Autonomous | TurnRole::Event => Self {
max_provider_rounds: Some(24),
max_tool_calls: Some(96),
max_elapsed: None,
max_accumulated_tool_result_bytes: Some(8 * 1024 * 1024),
max_context_tokens: None,
max_cost_usd: None,
max_total_tokens: None,
max_round_renewals: 0,
},
TurnRole::Worker => Self {
max_provider_rounds: Some(64),
max_tool_calls: None,
max_elapsed: None,
max_accumulated_tool_result_bytes: None,
max_context_tokens: None,
max_cost_usd: None,
max_total_tokens: None,
max_round_renewals: 2,
},
}
}
pub fn tighter(&self, other: &TurnBudget) -> TurnBudget {
fn min<T: PartialOrd + Copy>(a: Option<T>, b: Option<T>) -> Option<T> {
match (a, b) {
(Some(a), Some(b)) => Some(if b < a { b } else { a }),
(a, None) => a,
(None, b) => b,
}
}
TurnBudget {
max_provider_rounds: min(self.max_provider_rounds, other.max_provider_rounds),
max_tool_calls: min(self.max_tool_calls, other.max_tool_calls),
max_elapsed: min(self.max_elapsed, other.max_elapsed),
max_accumulated_tool_result_bytes: min(
self.max_accumulated_tool_result_bytes,
other.max_accumulated_tool_result_bytes,
),
max_context_tokens: min(self.max_context_tokens, other.max_context_tokens),
max_cost_usd: min(self.max_cost_usd, other.max_cost_usd),
max_total_tokens: min(self.max_total_tokens, other.max_total_tokens),
max_round_renewals: self.max_round_renewals.min(other.max_round_renewals),
}
}
pub fn for_event_turn(&self, config: &agent_core::config::TurnBudgetsConfig) -> TurnBudget {
self.tighter(&TurnBudget::from_config(TurnRole::Event, config))
}
pub fn has_time_limit(&self) -> bool {
self.max_elapsed.is_some_and(|d| !d.is_zero())
}
pub fn max_elapsed_secs(&self) -> Option<u64> {
self.max_elapsed.map(|d| d.as_secs())
}
pub fn from_config(role: TurnRole, config: &agent_core::config::TurnBudgetsConfig) -> Self {
let overrides = match role {
TurnRole::Foreground => &config.foreground,
TurnRole::Autonomous => &config.autonomous,
TurnRole::Event => &config.event,
TurnRole::Worker => &config.worker,
};
let mut budget = Self::for_role(role);
if let Some(v) = overrides.max_provider_rounds {
budget.max_provider_rounds = v.limit();
}
if let Some(v) = overrides.max_tool_calls {
budget.max_tool_calls = v.limit();
}
if let Some(v) = overrides.max_elapsed_secs {
budget.max_elapsed = (v > 0).then(|| Duration::from_secs(v));
}
if let Some(v) = overrides.max_accumulated_tool_result_bytes {
budget.max_accumulated_tool_result_bytes = v.limit();
}
if overrides.max_context_tokens.is_some() {
budget.max_context_tokens = overrides.max_context_tokens;
}
if overrides.max_cost_usd.is_some() {
budget.max_cost_usd = overrides.max_cost_usd;
}
if let Some(v) = overrides.max_round_renewals {
budget.max_round_renewals = v;
}
budget
}
}
#[derive(Debug, Default)]
pub struct OwnerSteerBudget {
pending: std::sync::Mutex<Option<TurnBudget>>,
}
impl OwnerSteerBudget {
pub fn send_owner_steer(&self, budget: TurnBudget, send: impl FnOnce() -> bool) -> bool {
let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner());
let sent = send();
if sent {
*pending = Some(budget);
}
sent
}
pub fn drain<R>(&self, drain: impl FnOnce() -> R) -> (R, Option<TurnBudget>) {
let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner());
let drained = drain();
(drained, pending.take())
}
}
#[derive(Debug, Default)]
pub struct UsageCounters {
cost_nanodollars: std::sync::atomic::AtomicU64,
input_tokens: std::sync::atomic::AtomicU64,
output_tokens: std::sync::atomic::AtomicU64,
cache_read_tokens: std::sync::atomic::AtomicU64,
cache_write_tokens: std::sync::atomic::AtomicU64,
latest_context_tokens: std::sync::atomic::AtomicU64,
records: std::sync::atomic::AtomicU64,
}
impl UsageCounters {
#[allow(clippy::too_many_arguments)]
pub fn record_priced(
&self,
model: &str,
input: u64,
output: u64,
cache_read: u64,
cache_write: u64,
cache_write_5m: Option<u64>,
cache_write_1h: Option<u64>,
) {
use std::sync::atomic::Ordering;
let cost = agent_core::pricing::calculate_cost_optional_split(
model,
input,
output,
cache_read,
cache_write,
cache_write_5m,
cache_write_1h,
);
let units = cost * 1_000_000_000.0;
let units = if !units.is_finite() || units < 0.0 {
u64::MAX
} else {
units.ceil() as u64
};
Self::saturating_add(&self.cost_nanodollars, units);
Self::saturating_add(&self.input_tokens, input);
Self::saturating_add(&self.output_tokens, output);
Self::saturating_add(&self.cache_read_tokens, cache_read);
Self::saturating_add(&self.cache_write_tokens, cache_write);
self.latest_context_tokens.store(
input.saturating_add(cache_read).saturating_add(cache_write),
Ordering::Relaxed,
);
Self::saturating_add(&self.records, 1);
}
pub fn records(&self) -> u64 {
self.records.load(std::sync::atomic::Ordering::Relaxed)
}
fn saturating_add(counter: &std::sync::atomic::AtomicU64, value: u64) {
use std::sync::atomic::Ordering;
let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
Some(old.saturating_add(value))
});
}
pub fn cost_usd(&self) -> f64 {
let units = self
.cost_nanodollars
.load(std::sync::atomic::Ordering::Relaxed);
if units == u64::MAX {
f64::MAX
} else {
units as f64 / 1_000_000_000.0
}
}
pub fn latest_context_tokens(&self) -> u64 {
self.latest_context_tokens
.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn totals(&self) -> (u64, u64, u64, u64) {
use std::sync::atomic::Ordering;
(
self.input_tokens.load(Ordering::Relaxed),
self.output_tokens.load(Ordering::Relaxed),
self.cache_read_tokens.load(Ordering::Relaxed),
self.cache_write_tokens.load(Ordering::Relaxed),
)
}
}
#[derive(Debug)]
pub struct TurnBudgetMeter {
budget: TurnBudget,
started: Instant,
rounds: u32,
tool_calls: u32,
tool_result_bytes: usize,
round_renewals: u32,
}
impl TurnBudgetMeter {
pub fn new(budget: TurnBudget) -> Self {
Self {
budget,
started: Instant::now(),
rounds: 0,
tool_calls: 0,
tool_result_bytes: 0,
round_renewals: 0,
}
}
pub fn budget(&self) -> &TurnBudget {
&self.budget
}
pub fn switch_budget(&mut self, budget: TurnBudget) {
self.budget = budget;
}
pub(crate) fn wall_clock_exceeded(&self) -> bool {
self.budget
.max_elapsed
.is_some_and(|max| self.started.elapsed() >= max)
}
pub(crate) fn start_context_segment(&mut self) {
self.started = Instant::now();
}
pub fn begin_round(&mut self) -> Result<(), BudgetDimension> {
if self.wall_clock_exceeded() {
return Err(BudgetDimension::WallClock);
}
if self
.budget
.max_provider_rounds
.is_some_and(|max| self.rounds >= max)
{
return Err(BudgetDimension::ProviderRounds);
}
self.rounds = self.rounds.saturating_add(1);
Ok(())
}
pub fn try_renew_rounds(&mut self) -> Option<u32> {
if self.round_renewals >= self.budget.max_round_renewals {
return None;
}
self.round_renewals += 1;
self.rounds = 0;
Some(self.budget.max_round_renewals - self.round_renewals)
}
pub fn round_renewals_used(&self) -> u32 {
self.round_renewals
}
pub fn elapsed(&self) -> Duration {
self.started.elapsed()
}
pub fn exhaustion_error(&self, dimension: BudgetDimension) -> agent_core::TurnError {
let mut error = agent_core::TurnError::budget(dimension);
if dimension == BudgetDimension::WallClock {
error.message = format!(
"turn stopped at its {}s time limit (ran {}s). History is kept; send a prompt to continue. \
The limit comes from turn_budget.<role>.max_elapsed_secs or /budget time; /budget time off removes it.",
self.budget.max_elapsed_secs().unwrap_or(0),
self.elapsed().as_secs(),
);
}
error
}
pub fn rounds_used(&self) -> u32 {
self.rounds
}
pub fn tool_calls_used(&self) -> u32 {
self.tool_calls
}
pub fn tool_result_bytes_used(&self) -> usize {
self.tool_result_bytes
}
pub fn remaining_tool_calls(&self) -> u32 {
self.budget
.max_tool_calls
.map_or(u32::MAX, |max| max.saturating_sub(self.tool_calls))
}
pub fn charge_tool_calls(&mut self, n: u32) {
self.tool_calls = self.tool_calls.saturating_add(n);
}
pub fn charge_tool_result_bytes(&mut self, bytes: usize) -> Result<(), BudgetDimension> {
self.tool_result_bytes = self.tool_result_bytes.saturating_add(bytes);
if self
.budget
.max_accumulated_tool_result_bytes
.is_some_and(|max| self.tool_result_bytes > max)
{
return Err(BudgetDimension::ToolResultBytes);
}
Ok(())
}
pub fn check_usage(&self, usage: &UsageCounters) -> Result<(), BudgetDimension> {
if let Some(max_context) = self.budget.max_context_tokens {
if usage.latest_context_tokens() > max_context {
return Err(BudgetDimension::InputTokens);
}
}
if let Some(max_cost) = self.budget.max_cost_usd {
if usage.cost_usd() > max_cost {
return Err(BudgetDimension::CostUsd);
}
}
if let Some(max_tokens) = self.budget.max_total_tokens {
let (input, output, _, _) = usage.totals();
if input.saturating_add(output) > max_tokens {
return Err(BudgetDimension::TotalTokens);
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use agent_core::config::TurnLimit;
#[test]
fn an_owner_steer_switches_the_budget_and_keeps_the_usage() {
let event = TurnBudget::from_config(TurnRole::Event, &Default::default());
let session = TurnBudget::for_role(TurnRole::Foreground);
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_provider_rounds: Some(2),
..event
});
meter.begin_round().unwrap();
meter.begin_round().unwrap();
meter.charge_tool_calls(3);
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
let lift = OwnerSteerBudget::default();
assert_eq!(lift.drain(|| ()).1, None);
assert!(!lift.send_owner_steer(session.clone(), || false));
assert_eq!(lift.drain(|| ()).1, None);
assert!(lift.send_owner_steer(session.clone(), || true));
let (_, budget) = lift.drain(|| ());
meter.switch_budget(budget.expect("the owner's steer carries the session budget"));
assert_eq!(lift.drain(|| ()).1, None, "taken once");
assert_eq!(meter.budget(), &session);
assert_eq!(meter.rounds_used(), 2, "usage keeps counting");
assert_eq!(meter.tool_calls_used(), 3);
meter.begin_round().unwrap();
assert_eq!(meter.rounds_used(), 3);
}
#[test]
fn partial_ttl_metadata_cannot_drop_cost_or_lower_tier() {
for (input, aggregate, c5, c1, expected) in [
(0, 100_001, Some(0), None, 0.100001),
(0, 100_001, None, Some(0), 0.100001),
(60_000, 40_000, Some(10_000), None, 0.01325),
(60_000, 40_001, Some(10_000), None, 0.066251),
(60_000, 40_001, None, Some(10_000), 0.070001),
(60_000, 40_001, Some(10_000), Some(10_000), 0.066251),
(60_000, 1, Some(20_000), Some(20_001), 0.062501),
(60_000, 1, Some(40_001), None, 0.055000625),
(60_000, 40_001, Some(40_001), None, 0.055000625),
(60_000, 40_001, None, Some(40_001), 0.070001),
(60_000, 40_001, Some(10_000), Some(30_001), 0.066251),
(60_000, 40_001, None, None, 0.055000625),
] {
let usage = UsageCounters::default();
usage.record_priced("claude-haiku-5.5", input, 0, 0, aggregate, c5, c1);
assert!(
(usage.cost_usd() - expected).abs() < 2e-9,
"input={input}, aggregate={aggregate}, split={c5:?}/{c1:?}"
);
let mut budget = TurnBudget::for_role(TurnRole::Foreground);
budget.max_cost_usd = Some(expected / 2.0);
assert_eq!(
TurnBudgetMeter::new(budget).check_usage(&usage),
Err(BudgetDimension::CostUsd)
);
}
for (aggregate, c5, c1) in [
(u64::MAX, Some(1), None),
(1, Some(u64::MAX), Some(u64::MAX)),
] {
let usage = UsageCounters::default();
usage.record_priced("claude-haiku-5.5", 0, 0, 0, aggregate, c5, c1);
assert_eq!(usage.cost_usd(), f64::MAX);
}
}
#[test]
fn costs_are_per_request_not_cumulative_tiers_or_current_model() {
let usage = UsageCounters::default();
for _ in 0..2 {
usage.record_priced("claude-haiku-5.5", 60_000, 1_000, 0, 0, None, None);
}
assert!((usage.cost_usd() - 0.013).abs() < 3e-9);
let mut budget = TurnBudget::for_role(TurnRole::Foreground);
budget.max_cost_usd = Some(0.03);
let meter = TurnBudgetMeter::new(budget);
assert_eq!(meter.check_usage(&usage), Ok(()));
usage.record_priced("claude-haiku-5-5", 100_001, 1_000, 0, 0, None, None);
assert!((usage.cost_usd() - 0.0655005).abs() < 4e-9);
assert_eq!(meter.check_usage(&usage), Err(BudgetDimension::CostUsd));
usage.record_priced("claude-opus-4.1", 1_000_000, 1_000_000, 0, 0, None, None);
usage.record_priced("claude-opus-4.8", 1_000_000, 1_000_000, 0, 0, None, None);
assert!((usage.cost_usd() - 120.0655005).abs() < 6e-9);
assert_eq!(usage.latest_context_tokens(), 1_000_000);
}
#[test]
fn request_tier_includes_reads_and_both_write_ttls() {
for extra in [0, 1] {
let usage = UsageCounters::default();
usage.record_priced(
"claude-haiku-5-5",
60_000,
1_000,
20_000 + extra,
20_000,
Some(5_000),
Some(15_000),
);
let expected = (0.006 + 0.0002 + 0.000625 + 0.003 + 0.0005 + extra as f64 * 0.00000001)
* if extra == 0 { 1.0 } else { 5.0 };
assert!((usage.cost_usd() - expected).abs() < 2e-9);
assert_eq!(usage.latest_context_tokens(), 100_000 + extra);
}
}
#[test]
fn unknown_models_and_saturation_cannot_silently_zero_cost() {
let usage = UsageCounters::default();
usage.record_priced("unknown-model", 1_000_000, 0, 0, 0, None, None);
assert!((usage.cost_usd() - 3.0).abs() < 2e-9);
for _ in 0..2 {
usage.record_priced(
"claude-opus-4.1",
u64::MAX,
u64::MAX,
u64::MAX,
u64::MAX,
None,
None,
);
}
assert_eq!(usage.totals(), (u64::MAX, u64::MAX, u64::MAX, u64::MAX));
assert_eq!(usage.cost_usd(), f64::MAX);
assert_eq!(usage.latest_context_tokens(), u64::MAX);
}
#[test]
fn role_defaults_are_typed_and_only_foreground_is_unlimited() {
let fg = TurnBudget::for_role(TurnRole::Foreground);
let auto = TurnBudget::for_role(TurnRole::Autonomous);
let worker = TurnBudget::for_role(TurnRole::Worker);
assert_eq!(fg.max_provider_rounds, None);
assert_eq!(fg.max_tool_calls, None);
assert_eq!(fg.max_accumulated_tool_result_bytes, None);
for budget in [&fg, &auto, &worker] {
assert_eq!(budget.max_elapsed, None);
assert!(!budget.has_time_limit());
}
assert!(fg.max_context_tokens.is_none());
assert!(fg.max_cost_usd.is_none());
assert_eq!(auto.max_round_renewals, 0);
assert!(worker.max_round_renewals > 0);
assert!(fg.max_round_renewals >= worker.max_round_renewals);
}
#[test]
fn autonomous_role_keeps_finite_per_turn_defaults() {
let auto = TurnBudget::for_role(TurnRole::Autonomous);
assert_eq!(auto.max_provider_rounds, Some(24));
assert_eq!(auto.max_tool_calls, Some(96));
assert_eq!(
auto.max_accumulated_tool_result_bytes,
Some(8 * 1024 * 1024)
);
assert_eq!(auto.max_round_renewals, 0);
let mut meter = TurnBudgetMeter::new(auto);
for _ in 0..24 {
assert_eq!(meter.begin_round(), Ok(()));
}
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), None);
assert_eq!(meter.remaining_tool_calls(), 96);
}
#[test]
fn event_role_uses_the_watcher_numbers_and_its_own_config() {
let event = TurnBudget::for_role(TurnRole::Event);
assert_eq!(event, TurnBudget::for_role(TurnRole::Autonomous));
let cfg = agent_core::config::load_config_from_str(
"turn_budget.event.max_provider_rounds = unlimited\n\
turn_budget.event.max_tool_calls = 300\n\
turn_budget.autonomous.max_tool_calls = 5\n",
)
.turn_budgets;
let event = TurnBudget::from_config(TurnRole::Event, &cfg);
assert_eq!(event.max_provider_rounds, None);
assert_eq!(event.max_tool_calls, Some(300));
assert_eq!(
event.max_accumulated_tool_result_bytes,
Some(8 * 1024 * 1024)
);
assert_eq!(
TurnBudget::from_config(TurnRole::Autonomous, &cfg).max_tool_calls,
Some(5)
);
}
#[test]
fn tighter_takes_the_smaller_limit_per_dimension() {
let session = TurnBudget {
max_elapsed: Some(Duration::from_secs(600)),
max_cost_usd: Some(2.0),
max_tool_calls: Some(500),
..TurnBudget::for_role(TurnRole::Foreground)
};
let event = TurnBudget {
max_cost_usd: Some(5.0),
max_total_tokens: Some(1_000),
..TurnBudget::for_role(TurnRole::Event)
};
let combined = session.tighter(&event);
assert_eq!(combined.max_provider_rounds, Some(24));
assert_eq!(combined.max_tool_calls, Some(96));
assert_eq!(
combined.max_accumulated_tool_result_bytes,
Some(8 * 1024 * 1024)
);
assert_eq!(combined.max_elapsed, Some(Duration::from_secs(600)));
assert_eq!(combined.max_cost_usd, Some(2.0));
assert_eq!(combined.max_total_tokens, Some(1_000));
assert_eq!(combined.max_context_tokens, None);
assert_eq!(combined.max_round_renewals, 0);
assert_eq!(combined, event.tighter(&session));
}
#[test]
fn total_tokens_limit_counts_uncached_input_and_output() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_total_tokens: Some(100),
..TurnBudget::for_role(TurnRole::Autonomous)
});
let usage = UsageCounters::default();
usage.record_priced("claude-sonnet-4-5", 40, 60, 50_000, 50_000, None, None);
assert_eq!(meter.check_usage(&usage), Ok(()));
usage.record_priced("claude-sonnet-4-5", 1, 0, 0, 0, None, None);
assert_eq!(meter.check_usage(&usage), Err(BudgetDimension::TotalTokens));
meter.charge_tool_calls(0);
}
#[test]
fn worker_role_keeps_a_finite_round_cap_of_192() {
let worker = TurnBudget::for_role(TurnRole::Worker);
assert_eq!(worker.max_provider_rounds, Some(64));
assert_eq!(worker.max_round_renewals, 2);
let mut meter = TurnBudgetMeter::new(worker);
let mut rounds = 0u32;
loop {
match meter.begin_round() {
Ok(()) => rounds += 1,
Err(BudgetDimension::ProviderRounds) => {
if meter.try_renew_rounds().is_none() {
break;
}
}
Err(other) => panic!("unexpected dimension {other:?}"),
}
assert!(rounds <= 10_000, "worker round cap must be finite");
}
assert_eq!(rounds, 192);
}
#[test]
fn config_overlays_only_set_fields() {
let mut cfg = agent_core::config::TurnBudgetsConfig::default();
cfg.autonomous.max_provider_rounds = Some(TurnLimit::Max(2));
cfg.autonomous.max_elapsed_secs = Some(9);
cfg.autonomous.max_round_renewals = Some(3);
let budget = TurnBudget::from_config(TurnRole::Autonomous, &cfg);
assert_eq!(budget.max_provider_rounds, Some(2));
assert_eq!(budget.max_elapsed, Some(Duration::from_secs(9)));
assert_eq!(budget.max_round_renewals, 3);
assert_eq!(
budget.max_tool_calls,
TurnBudget::for_role(TurnRole::Autonomous).max_tool_calls
);
}
#[test]
fn zero_seconds_in_config_means_no_time_limit() {
let mut cfg = agent_core::config::TurnBudgetsConfig::default();
cfg.foreground.max_elapsed_secs = Some(0);
let budget = TurnBudget::from_config(TurnRole::Foreground, &cfg);
assert_eq!(budget.max_elapsed, None);
}
#[test]
fn config_sets_round_tool_call_and_byte_limits_and_unlimited_clears_them() {
for role in [
TurnRole::Foreground,
TurnRole::Autonomous,
TurnRole::Event,
TurnRole::Worker,
] {
let mut cfg = agent_core::config::TurnBudgetsConfig::default();
let overrides = match role {
TurnRole::Foreground => &mut cfg.foreground,
TurnRole::Autonomous => &mut cfg.autonomous,
TurnRole::Event => &mut cfg.event,
TurnRole::Worker => &mut cfg.worker,
};
overrides.max_provider_rounds = Some(TurnLimit::Max(5));
overrides.max_tool_calls = Some(TurnLimit::Max(7));
overrides.max_accumulated_tool_result_bytes = Some(TurnLimit::Max(11));
let budget = TurnBudget::from_config(role, &cfg);
assert_eq!(budget.max_provider_rounds, Some(5));
assert_eq!(budget.max_tool_calls, Some(7));
assert_eq!(budget.max_accumulated_tool_result_bytes, Some(11));
let overrides = match role {
TurnRole::Foreground => &mut cfg.foreground,
TurnRole::Autonomous => &mut cfg.autonomous,
TurnRole::Event => &mut cfg.event,
TurnRole::Worker => &mut cfg.worker,
};
overrides.max_provider_rounds = Some(TurnLimit::Unlimited);
overrides.max_tool_calls = Some(TurnLimit::Unlimited);
overrides.max_accumulated_tool_result_bytes = Some(TurnLimit::Unlimited);
let budget = TurnBudget::from_config(role, &cfg);
assert_eq!(budget.max_provider_rounds, None);
assert_eq!(budget.max_tool_calls, None);
assert_eq!(budget.max_accumulated_tool_result_bytes, None);
}
}
#[test]
fn configured_limits_are_enforced_exactly() {
let mut cfg = agent_core::config::TurnBudgetsConfig::default();
cfg.worker.max_provider_rounds = Some(TurnLimit::Max(2));
cfg.worker.max_round_renewals = Some(1);
cfg.worker.max_tool_calls = Some(TurnLimit::Max(3));
cfg.worker.max_accumulated_tool_result_bytes = Some(TurnLimit::Max(100));
let mut meter = TurnBudgetMeter::new(TurnBudget::from_config(TurnRole::Worker, &cfg));
for _ in 0..2 {
assert_eq!(meter.begin_round(), Ok(()));
}
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), Some(0));
for _ in 0..2 {
assert_eq!(meter.begin_round(), Ok(()));
}
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), None);
assert_eq!(meter.remaining_tool_calls(), 3);
meter.charge_tool_calls(3);
assert_eq!(meter.remaining_tool_calls(), 0);
assert_eq!(meter.charge_tool_result_bytes(100), Ok(()));
assert_eq!(
meter.charge_tool_result_bytes(1),
Err(BudgetDimension::ToolResultBytes)
);
}
#[test]
fn zero_in_config_means_none_allowed_for_every_role() {
let cfg = agent_core::config::load_config_from_str(
"turn_budget.foreground.max_provider_rounds = 0\n\
turn_budget.foreground.max_tool_calls = 0\n\
turn_budget.foreground.max_accumulated_tool_result_bytes = 0\n\
turn_budget.autonomous.max_provider_rounds = 0\n\
turn_budget.autonomous.max_tool_calls = 0\n\
turn_budget.autonomous.max_accumulated_tool_result_bytes = 0\n\
turn_budget.worker.max_provider_rounds = 0\n\
turn_budget.worker.max_tool_calls = 0\n\
turn_budget.worker.max_accumulated_tool_result_bytes = 0\n\
turn_budget.event.max_provider_rounds = 0\n\
turn_budget.event.max_tool_calls = 0\n\
turn_budget.event.max_accumulated_tool_result_bytes = 0\n",
)
.turn_budgets;
for role in [
TurnRole::Foreground,
TurnRole::Autonomous,
TurnRole::Event,
TurnRole::Worker,
] {
let budget = TurnBudget::from_config(role, &cfg);
assert_eq!(budget.max_provider_rounds, Some(0), "{role:?}");
assert_eq!(budget.max_tool_calls, Some(0), "{role:?}");
assert_eq!(
budget.max_accumulated_tool_result_bytes,
Some(0),
"{role:?}"
);
let mut meter = TurnBudgetMeter::new(budget);
assert_eq!(meter.remaining_tool_calls(), 0);
assert_eq!(
meter.charge_tool_result_bytes(1),
Err(BudgetDimension::ToolResultBytes)
);
loop {
match meter.begin_round() {
Ok(()) => panic!("{role:?}: a zero round limit allows no rounds"),
Err(BudgetDimension::ProviderRounds) => {
if meter.try_renew_rounds().is_none() {
break;
}
}
Err(other) => panic!("unexpected {other:?}"),
}
}
}
}
#[test]
fn configured_unlimited_never_exhausts() {
let mut cfg = agent_core::config::TurnBudgetsConfig::default();
cfg.autonomous.max_provider_rounds = Some(TurnLimit::Unlimited);
cfg.autonomous.max_tool_calls = Some(TurnLimit::Unlimited);
cfg.autonomous.max_accumulated_tool_result_bytes = Some(TurnLimit::Unlimited);
let mut meter = TurnBudgetMeter::new(TurnBudget::from_config(TurnRole::Autonomous, &cfg));
for _ in 0..1_000 {
assert_eq!(meter.begin_round(), Ok(()));
}
assert_eq!(meter.remaining_tool_calls(), u32::MAX);
meter.charge_tool_calls(1_000);
assert_eq!(meter.remaining_tool_calls(), u32::MAX);
assert_eq!(meter.charge_tool_result_bytes(1 << 20), Ok(()));
}
#[test]
fn an_unlimited_meter_never_exhausts_but_still_counts() {
let mut unlimited = agent_core::config::TurnBudgetsConfig::default();
unlimited.autonomous.max_provider_rounds = Some(TurnLimit::Unlimited);
unlimited.autonomous.max_tool_calls = Some(TurnLimit::Unlimited);
unlimited.autonomous.max_accumulated_tool_result_bytes = Some(TurnLimit::Unlimited);
unlimited.worker.max_provider_rounds = Some(TurnLimit::Unlimited);
for budget in [
TurnBudget::for_role(TurnRole::Foreground),
TurnBudget::from_config(TurnRole::Autonomous, &unlimited),
TurnBudget::from_config(TurnRole::Worker, &unlimited),
] {
let mut meter = TurnBudgetMeter::new(budget);
for _ in 0..10_000 {
assert_eq!(meter.begin_round(), Ok(()));
assert_eq!(meter.remaining_tool_calls(), u32::MAX);
meter.charge_tool_calls(1);
assert_eq!(meter.charge_tool_result_bytes(64 * 1024), Ok(()));
}
assert_eq!(meter.rounds_used(), 10_000);
assert_eq!(meter.tool_calls_used(), 10_000);
assert_eq!(meter.tool_result_bytes_used(), 10_000 * 64 * 1024);
assert_eq!(meter.round_renewals_used(), 0);
assert_eq!(meter.check_usage(&UsageCounters::default()), Ok(()));
}
}
#[test]
fn a_turn_without_a_time_limit_never_runs_out_of_time() {
let mut meter = TurnBudgetMeter::new(TurnBudget::for_role(TurnRole::Foreground));
meter.started = Instant::now() - Duration::from_secs(30 * 24 * 60 * 60);
assert!(!meter.wall_clock_exceeded());
assert_eq!(meter.begin_round(), Ok(()));
}
#[test]
fn committed_context_segment_resets_only_time() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_elapsed: Some(Duration::from_secs(7200)),
..TurnBudget::for_role(TurnRole::Foreground)
});
meter.started = Instant::now() - Duration::from_secs(7201);
meter.charge_tool_calls(3);
meter.charge_tool_result_bytes(1234).unwrap();
meter.try_renew_rounds();
let original = meter.budget().clone();
assert_eq!(meter.begin_round(), Err(BudgetDimension::WallClock));
meter.start_context_segment();
assert!(meter.begin_round().is_ok());
assert_eq!(meter.budget(), &original);
assert_eq!(meter.tool_calls_used(), 3);
assert_eq!(meter.tool_result_bytes_used(), 1234);
assert_eq!(meter.round_renewals_used(), 1);
}
#[test]
fn exhaustion_diagnostic_keeps_typed_outcome_and_reports_metadata() {
let meter = TurnBudgetMeter::new(TurnBudget {
max_elapsed: Some(Duration::from_secs(7200)),
..TurnBudget::for_role(TurnRole::Foreground)
});
let error = meter.exhaustion_error(BudgetDimension::WallClock);
assert_eq!(
error.outcome,
agent_core::TurnOutcome::BudgetExceeded {
dimension: BudgetDimension::WallClock
}
);
assert!(error.message.contains("its 7200s time limit (ran "));
assert!(error.message.contains("History is kept"));
assert!(error.message.contains("/budget time off"));
for dimension in [
BudgetDimension::ToolCalls,
BudgetDimension::ProviderRounds,
BudgetDimension::ToolResultBytes,
BudgetDimension::InputTokens,
BudgetDimension::OutputTokens,
BudgetDimension::CostUsd,
] {
assert_eq!(
meter.exhaustion_error(dimension),
agent_core::TurnError::budget(dimension)
);
}
}
#[test]
fn meter_enforces_exact_rounds_and_wall_clock_order() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_provider_rounds: Some(2),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert!(meter.begin_round().is_ok());
assert!(meter.begin_round().is_ok());
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
let mut stale = TurnBudgetMeter::new(TurnBudget {
max_elapsed: Some(Duration::ZERO),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert_eq!(stale.begin_round(), Err(BudgetDimension::WallClock));
}
#[test]
fn round_renewals_extend_a_bounded_number_of_times() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_provider_rounds: Some(1),
max_round_renewals: 2,
..TurnBudget::for_role(TurnRole::Foreground)
});
assert!(meter.begin_round().is_ok());
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), Some(1));
assert!(meter.begin_round().is_ok());
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), Some(0));
assert!(meter.begin_round().is_ok());
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), None);
assert_eq!(meter.round_renewals_used(), 2);
}
#[test]
fn wall_clock_still_bounds_a_renewed_turn() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_provider_rounds: Some(1),
max_round_renewals: 5,
max_elapsed: Some(Duration::ZERO),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert_eq!(meter.try_renew_rounds(), Some(4));
assert_eq!(meter.begin_round(), Err(BudgetDimension::WallClock));
}
#[test]
fn autonomous_role_has_no_graceful_renewals() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_provider_rounds: Some(1),
..TurnBudget::for_role(TurnRole::Autonomous)
});
assert!(meter.begin_round().is_ok());
assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
assert_eq!(meter.try_renew_rounds(), None);
}
#[test]
fn meter_tool_calls_and_bytes_are_exact() {
let mut meter = TurnBudgetMeter::new(TurnBudget {
max_tool_calls: Some(3),
max_accumulated_tool_result_bytes: Some(10),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert_eq!(meter.remaining_tool_calls(), 3);
meter.charge_tool_calls(2);
assert_eq!(meter.remaining_tool_calls(), 1);
meter.charge_tool_calls(1);
assert_eq!(meter.remaining_tool_calls(), 0);
assert!(meter.charge_tool_result_bytes(10).is_ok());
assert_eq!(
meter.charge_tool_result_bytes(1),
Err(BudgetDimension::ToolResultBytes)
);
}
#[test]
fn usage_counters_drive_optional_dimensions() {
let usage = UsageCounters::default();
usage.record_priced("claude-sonnet-4-5", 100, 50, 10, 5, None, None);
assert_eq!(usage.latest_context_tokens(), 115);
let meter = TurnBudgetMeter::new(TurnBudget {
max_context_tokens: Some(100),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert_eq!(meter.check_usage(&usage), Err(BudgetDimension::InputTokens));
let cost_meter = TurnBudgetMeter::new(TurnBudget {
max_cost_usd: Some(0.0),
..TurnBudget::for_role(TurnRole::Foreground)
});
assert_eq!(
cost_meter.check_usage(&usage),
Err(BudgetDimension::CostUsd)
);
let unlimited = TurnBudgetMeter::new(TurnBudget::for_role(TurnRole::Foreground));
assert!(unlimited.check_usage(&usage).is_ok());
}
}