agent_engine/runtime/

budget.rs

//! Task 23 — per-turn budgets (spec §8.1).
//!
//! Every stream session runs under one [`TurnBudget`]: provider rounds,
//! tool calls, elapsed wall-clock, accumulated tool-result bytes, and the
//! optional context-token / USD-cost dimensions. Budgets differ by role
//! (foreground, autonomous/watcher, worker) but share ONE enforcement
//! mechanism — the [`TurnBudgetMeter`] consulted by the stream loop.
//!
//! Composition note: the chat frontend's reactor auto-turn cap
//! (`events.auto_turn_cap`, default `engine::reactor::AUTO_TURN_CAP`) limits consecutive auto-triggered
//! turns ACROSS turns; this budget bounds work WITHIN one turn. They
//! compose — neither replaces nor duplicates the other.

use std::time::{Duration, Instant};

use agent_core::BudgetDimension;

/// Start of every synthetic tool_result for a call the engine did not run
/// because a turn budget was used up. Hosts that count executed calls (the
/// watcher agent) skip results that start with it.
pub const TOOL_CALL_NOT_EXECUTED: &str = "Tool call not executed";

/// The synthetic tool_result for a call refused by the turn's tool-call
/// allowance.
pub const TOOL_CALL_BUDGET_REFUSAL: &str =
    "Tool call not executed: turn tool-call budget exhausted";

/// Per-turn limits (spec §8.1).
#[derive(Debug, Clone, PartialEq)]
pub struct TurnBudget {
    /// Provider rounds per turn (per renewal). `None` means no limit — the
    /// foreground default. Autonomous (24) and worker (64, x3 with renewals)
    /// turns are bounded by default. Set with
    /// `turn_budget.<role>.max_provider_rounds`.
    pub max_provider_rounds: Option<u32>,
    /// Tool calls per turn. `None` means no limit — the foreground and worker
    /// default; autonomous turns default to 96. Set with
    /// `turn_budget.<role>.max_tool_calls`.
    pub max_tool_calls: Option<u32>,
    /// Wall-clock limit for one turn. `None` (the default for every role)
    /// means no limit: a turn that is still making progress is never cut off
    /// for taking long. Opt in with `turn_budget.<role>.max_elapsed_secs` or
    /// `/budget time`.
    pub max_elapsed: Option<Duration>,
    /// Accumulated tool-result bytes per turn. `None` means no limit — the
    /// foreground and worker default; autonomous turns default to 8 MiB. Set
    /// with `turn_budget.<role>.max_accumulated_tool_result_bytes`.
    pub max_accumulated_tool_result_bytes: Option<usize>,
    pub max_context_tokens: Option<u64>,
    pub max_cost_usd: Option<f64>,
    /// Uncached input + output tokens summed over the turn's requests (cache
    /// reads and writes are priced into `max_cost_usd`, not counted here).
    /// No role sets it by default; the watcher agent uses it to carry its
    /// remaining `max_session_tokens` into the turn.
    pub max_total_tokens: Option<u64>,
    /// Bounded auto-renewals of the provider-round allowance (spec §8.1).
    /// On exhaustion the stream loop may reset the round counter this many
    /// times so a long, legitimate task continues instead of hard-failing.
    /// Round renewal never resets wall-clock (durable context successors do); a
    /// value of `0` means no graceful continuation (immediate hard stop). Only
    /// matters when `max_provider_rounds` is configured.
    pub max_round_renewals: u32,
}

/// Execution role a turn runs under.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TurnRole {
    /// Interactive foreground session.
    Foreground,
    /// Watcher agent sessions (`synaps agent`; no human in the loop).
    Autonomous,
    /// Turns a session starts by itself because events arrived (inbox
    /// messages, subagent completions). Nobody typed them.
    Event,
    /// Delegated worker/subagent sessions.
    Worker,
}

impl TurnBudget {
    /// Compiled per-role defaults. No role has a wall-clock limit unless
    /// configured. Foreground (a human is watching and can interrupt) has no
    /// round, tool-call or tool-result-byte limit either. Unattended roles
    /// stay bounded: autonomous (watcher) and event-started turns hard-stop
    /// at 24 rounds / 96 tool calls / 8 MiB (nobody is watching), and a
    /// worker stops after 64 rounds x (1 + 2 renewals) = 192 rounds per run
    /// (a `timeout: 0` worker has no other bound).
    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,
                // A human is watching and can interrupt; let long agentic
                // tasks continue through several checkpoints.
                max_round_renewals: 8,
            },
            // Same numbers as the watcher; configured separately.
            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,
                // Nobody is watching: no graceful extension — the round cap
                // is a hard stop.
                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,
                // Bounded delegation under a supervising turn: 64 x (1 + 2)
                // = 192 rounds.
                max_round_renewals: 2,
            },
        }
    }

    /// The stricter of two budgets, dimension by dimension (`None` = no
    /// limit loses to any limit). Used to narrow a role budget to what a
    /// session has left, or to combine a session's budget with the event
    /// role's.
    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),
        }
    }

    /// The budget for a turn the session starts by itself because events
    /// arrived: this (session) budget narrowed by the event role's
    /// (`turn_budget.event.*`), so a `/budget` the user set still applies.
    pub fn for_event_turn(&self, config: &agent_core::config::TurnBudgetsConfig) -> TurnBudget {
        self.tighter(&TurnBudget::from_config(TurnRole::Event, config))
    }

    /// Whether a wall-clock limit can end (or checkpoint) a turn. A zero
    /// limit counts as exhausted at once and never as a time checkpoint.
    pub fn has_time_limit(&self) -> bool {
        self.max_elapsed.is_some_and(|d| !d.is_zero())
    }

    /// The wall-clock limit in seconds, for diagnostics (`None` = no limit).
    pub fn max_elapsed_secs(&self) -> Option<u64> {
        self.max_elapsed.map(|d| d.as_secs())
    }

    /// Role defaults overlaid with the typed config. Unset fields keep the
    /// compiled defaults; `unlimited` clears a round, tool-call or
    /// tool-result-byte limit (`0` means none allowed, deprecated); `0`
    /// seconds means no wall-clock limit.
    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
    }
}

/// Moves an event-started turn onto the session's budget once the owner's
/// own message is delivered into it (a steer): from then on someone is
/// watching. Shared by the session actor (which sends the owner's steer) and
/// the turn's stream (which drains steering). Sending an owner steer and
/// draining steering hold the same lock, so the stream takes the budget in
/// the very drain that delivers the owner's message, never before it and
/// never for an event or note steered into the same channel.
#[derive(Debug, Default)]
pub struct OwnerSteerBudget {
    pending: std::sync::Mutex<Option<TurnBudget>>,
}

impl OwnerSteerBudget {
    /// Session side: `send` the owner's message into the turn; if it went in,
    /// the turn takes `budget` when the stream delivers it.
    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
    }

    /// Stream side: run `drain` (the steering drain) and return the budget
    /// an owner message delivered by it, or by an earlier drain, carries.
    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())
    }
}

/// Shared provider-usage counters, updated at the single authoritative
/// Usage emission sites in the transport and read by the stream loop for
/// the optional token/cost dimensions. Metadata only.
#[derive(Debug, Default)]
pub struct UsageCounters {
    /// Sum of individually priced requests, rounded UP to nanodollars.
    /// Saturation is fail-closed for any finite configured budget.
    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 request's context footprint (input + cache read + cache
    /// write) — the context-token dimension compares against this.
    latest_context_tokens: std::sync::atomic::AtomicU64,
    /// Number of requests booked; lets a caller tell whether the latest
    /// footprint belongs to the request it just sent.
    records: std::sync::atomic::AtomicU64,
}

impl UsageCounters {
    /// Book exactly one request, using its model and optional cache TTL split.
    /// Unknown models retain the pricing table's nonzero Sonnet fallback.
    #[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)
    }

    /// Accumulated (input, output, cache_read, cache_write) totals.
    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),
        )
    }
}

/// One turn's running budget meter. NOT thread-shared: owned by the
/// stream loop; the only cross-task input is the [`UsageCounters`] cell.
#[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
    }

    /// Switch the turn to `budget` mid-turn. Usage so far (rounds, tool
    /// calls, tool-result bytes, renewals, elapsed time) keeps counting.
    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)
    }

    /// A successfully committed durable context successor starts a new time
    /// segment. Never call on an attempted/failed save or a model checkpoint.
    /// Cumulative resource/cost limits and round renewals remain unchanged.
    pub(crate) fn start_context_segment(&mut self) {
        self.started = Instant::now();
    }

    /// Charge one provider round. Checked BEFORE the provider call:
    /// wall-clock first (a stale turn must not spend another request),
    /// then the exact round cap (when one is configured).
    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(())
    }

    /// Attempt a bounded graceful extension after the provider-round cap is
    /// hit. Resets the round counter and returns the number of renewals still
    /// remaining, or `None` when the renewal budget is exhausted. Wall-clock
    /// is deliberately untouched — [`Self::begin_round`] re-checks it, so a
    /// stale turn cannot extend past `max_elapsed`. Only meaningful right
    /// after `begin_round` returned [`BudgetDimension::ProviderRounds`].
    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)
    }

    /// Number of graceful round renewals already consumed this turn.
    pub fn round_renewals_used(&self) -> u32 {
        self.round_renewals
    }

    /// Wall-clock consumed by this turn so far. Diagnostics only — the
    /// enforcement path uses [`Self::wall_clock_exceeded`].
    pub fn elapsed(&self) -> Duration {
        self.started.elapsed()
    }

    /// Metadata-only diagnostic; preserve the typed outcome for every frontend
    /// and autonomous driver. Exhaustion is a local stop, not provider fallback.
    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
    }

    /// Provider rounds charged since the last renewal (reset by
    /// [`Self::try_renew_rounds`]). Diagnostics only.
    pub fn rounds_used(&self) -> u32 {
        self.rounds
    }

    /// Tool calls charged this turn. Diagnostics only.
    pub fn tool_calls_used(&self) -> u32 {
        self.tool_calls
    }

    /// Accumulated tool-result bytes charged this turn. Diagnostics only.
    pub fn tool_result_bytes_used(&self) -> usize {
        self.tool_result_bytes
    }

    /// Remaining tool-call allowance (for exact-cap batch splitting);
    /// `u32::MAX` when no tool-call limit is configured.
    pub fn remaining_tool_calls(&self) -> u32 {
        self.budget
            .max_tool_calls
            .map_or(u32::MAX, |max| max.saturating_sub(self.tool_calls))
    }

    /// Charge `n` executed tool calls (the caller split the batch against
    /// [`Self::remaining_tool_calls`], so this never over-charges).
    pub fn charge_tool_calls(&mut self, n: u32) {
        self.tool_calls = self.tool_calls.saturating_add(n);
    }

    /// Charge the bytes of one round's tool results; exceeding a configured
    /// accumulated cap exhausts the ToolResultBytes dimension.
    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(())
    }

    /// Optional dimensions, consulted after each provider round from the
    /// shared usage counters. Costs are summed per request, never repriced
    /// from token totals; both checks are no-ops when unconfigured.
    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();
        // Nothing sent: the drain carries no budget.
        assert_eq!(lift.drain(|| ()).1, None);
        // A send that fails (turn over) arms nothing.
        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);
        // Interactive turns have no round, tool-call or tool-result-byte
        // limit unless one is configured (#219).
        assert_eq!(fg.max_provider_rounds, None);
        assert_eq!(fg.max_tool_calls, None);
        assert_eq!(fg.max_accumulated_tool_result_bytes, None);
        // No role has a wall-clock limit unless one is configured.
        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());
        // Graceful continuation is most generous for foreground (human
        // watching) and disabled for autonomous (nobody watching).
        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() {
        // Nobody watches a watcher turn and its session limits are checked
        // between turns, so one turn must stay bounded by default.
        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() {
        // Turns started by events have nobody typing; they get the watcher's
        // per-turn numbers by default, configurable under turn_budget.event.
        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)
        );
        // The watcher role is configured separately.
        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();
        // Cache traffic does not count toward the token limit.
        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() {
        // 64 rounds per allowance x (1 + 2 renewals) = the pre-#219 cap.
        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)
            );
            // Renewals cannot turn a zero round allowance into rounds.
            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"));
        // Other dimensions keep the existing diagnostic and exact typed label.
        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)
        });
        // Wall clock outranks the round check.
        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)
        });
        // First round of the base allowance.
        assert!(meter.begin_round().is_ok());
        // Exhausted → renew #1, then a fresh round is available.
        assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
        assert_eq!(meter.try_renew_rounds(), Some(1));
        assert!(meter.begin_round().is_ok());
        // Exhausted again → renew #2 (last one).
        assert_eq!(meter.begin_round(), Err(BudgetDimension::ProviderRounds));
        assert_eq!(meter.try_renew_rounds(), Some(0));
        assert!(meter.begin_round().is_ok());
        // Exhausted a third time → no renewals left: hard stop.
        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)
        });
        // Renewal resets the round counter, but begin_round re-checks the
        // wall clock — a stale turn cannot extend past max_elapsed.
        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));
        // Nobody watching: the round cap is a hard stop.
        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());
    }
}