codewhale-tui 0.9.3

Terminal UI for open-source and open-weight coding models
//! Hard per-turn tool-call admission budget (#4415).
//!
//! One counter per turn, built from the task's structured `max_tool_calls`
//! constraint. Every proposed tool call is admitted through the same gate in
//! proposal order, so a batch larger than the remaining budget is truncated
//! to exactly the calls that still fit and the excess are rejected — never
//! executed — with a typed reason carrying the remaining-call count.

use codewhale_tools::ToolError;

/// Countdown of tool calls one turn may still admit.
///
/// This is the turn's admission state: created when the turn starts and
/// decremented at the admission gate. It deliberately does not live in the
/// tool catalog or surface policy, which only carry the declared limit.
/// `None` means unlimited — the default when a task declares no budget —
/// and leaves the gate inert.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ToolCallBudget {
    max: Option<u32>,
    remaining: Option<u32>,
}

/// Typed rejection for a call that exceeds the remaining budget.
///
/// Rendered into the same `PermissionDenied` error the neighboring admission
/// gates (deny-list, allow-list) produce, so the denial is visible in the
/// transcript and to the model with the remaining-call count spelled out.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ToolCallBudgetExceeded {
    max: u32,
}

impl ToolCallBudgetExceeded {
    pub(crate) fn into_tool_error(self, tool_name: &str) -> ToolError {
        ToolError::permission_denied(format!(
            "Tool '{tool_name}' rejected: per-turn tool-call budget of {} exhausted (remaining=0). \
             The call was not executed.",
            self.max
        ))
    }
}

impl ToolCallBudget {
    pub(crate) fn new(max_tool_calls: Option<u32>) -> Self {
        Self {
            max: max_tool_calls,
            remaining: max_tool_calls,
        }
    }

    /// Admit one proposed tool call. Every proposed call counts: while budget
    /// remains, the call is admitted and the remaining count decrements; once
    /// exhausted, the call is rejected with [`ToolCallBudgetExceeded`].
    pub(crate) fn admit(&mut self) -> Result<(), ToolCallBudgetExceeded> {
        let (Some(max), Some(remaining)) = (self.max, self.remaining.as_mut()) else {
            return Ok(());
        };
        if *remaining == 0 {
            return Err(ToolCallBudgetExceeded { max });
        }
        *remaining -= 1;
        Ok(())
    }
}