klieo-ops 3.3.0

//! `Governor` trait + RAII `Permit`.

use crate::types::{BudgetScope, ProviderId};
use async_trait::async_trait;
use thiserror::Error;

/// RAII permit. Dropping it releases the slot in the underlying bucket.
pub struct Permit {
    on_drop: Option<Box<dyn FnOnce() + Send>>,
}

impl Permit {
    /// Construct a permit. Impls call this when granting.
    pub fn new<F>(on_drop: F) -> Self
    where
        F: FnOnce() + Send + 'static,
    {
        Self {
            on_drop: Some(Box::new(on_drop)),
        }
    }
}

impl Drop for Permit {
    fn drop(&mut self) {
        if let Some(cb) = self.on_drop.take() {
            cb();
        }
    }
}

/// Budget snapshot.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct BudgetSnapshot {
    /// Remaining budget in the current window.
    pub remaining: i64,
    /// Limit in the current window.
    pub limit: i64,
}

/// Budget limit (Phase A models RPS only; TPM/RPM/cost in Phase B).
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub enum BudgetLimit {
    /// Requests per second.
    Rps(u32),
}

/// Egress host identifier.
pub type Host = String;

/// Governor primitive.
#[async_trait]
pub trait Governor: Send + Sync {
    /// Acquire a permit for an LLM call against `provider`.
    async fn acquire_llm(
        &self,
        provider: ProviderId,
        est_tokens: u32,
    ) -> Result<Permit, GovernorError>;

    /// Acquire a permit for outbound HTTP to `host`.
    async fn acquire_egress(&self, host: Host) -> Result<Permit, GovernorError>;

    /// Read current budget snapshot for `scope`.
    async fn budget(&self, scope: BudgetScope) -> BudgetSnapshot;

    /// Fence `scope` (refuse further acquires for the configured duration).
    async fn fence(&self, scope: BudgetScope, limit: BudgetLimit) -> Result<(), GovernorError>;

    /// Acquire a permit for an LLM call against a `scope`.
    ///
    /// Phase B B4 inbound-MCP entry point. The default impl delegates
    /// to [`Self::acquire_llm`] so adopter `Governor` impls keep
    /// working with no source edits — only `TokenBucketGovernor`
    /// applies a finer per-tenant cap. Scopes lacking a provider
    /// (`Global`, `Tenant`, `Agent`) surface
    /// [`GovernorError::Unavailable`]: silently downgrading to an
    /// unbounded path would defeat the inbound budget gate.
    async fn acquire_llm_scoped(
        &self,
        scope: BudgetScope,
        est_tokens: u32,
    ) -> Result<Permit, GovernorError> {
        match scope {
            BudgetScope::Provider(p) => self.acquire_llm(p, est_tokens).await,
            BudgetScope::TenantProvider(_, p) => self.acquire_llm(p, est_tokens).await,
            other => Err(GovernorError::Unavailable(format!(
                "acquire_llm_scoped requires a provider-bearing scope, got {other}"
            ))),
        }
    }
}

/// Governor errors.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum GovernorError {
    /// Budget exhausted in the current window.
    #[error("saturated: {scope}")]
    Saturated {
        /// Scope display.
        scope: String,
    },
    /// Acquire timed out.
    #[error("acquire timed out after {millis}ms")]
    TimedOut {
        /// Wait duration before timeout fired.
        millis: u64,
    },
    /// Storage / coordination error.
    #[error("unavailable: {0}")]
    Unavailable(String),
}