klieo-ops 3.3.0

//! In-memory token-bucket governor. Each scope owns a Tokio `Semaphore`
//! sized to its per-second quota; a refill task tops semaphores back to
//! their limit each second.

use super::trait_::{BudgetLimit, BudgetSnapshot, Governor, GovernorError, Host, Permit};
use crate::types::{BudgetScope, ProviderId, TenantId};
use async_trait::async_trait;
use dashmap::DashMap;
use klieo_core::KvStore;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Semaphore;

const DEFAULT_EGRESS_RPS: u32 = 10;
/// Default RPS cap applied to a `(tenant, provider)` pair the operator
/// did not pre-seed via `tenant_llm_limit`. Bounded by design — an
/// unconfigured tenant must NOT bypass the inbound budget gate.
const DEFAULT_TENANT_LLM_RPS: u32 = 5;
const DEFAULT_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(30);

/// Builder for the token-bucket governor.
pub struct TokenBucketGovernorBuilder {
    kv: Arc<dyn KvStore>,
    llm_limits: Vec<(ProviderId, u32)>,
    egress_default_rps: u32,
    tenant_llm_default_rps: u32,
    tenant_llm_limits: Vec<(TenantId, ProviderId, u32)>,
}

impl TokenBucketGovernorBuilder {
    /// Set the per-provider LLM RPS cap.
    #[must_use]
    pub fn llm_limit(mut self, provider: ProviderId, rps: u32) -> Self {
        self.llm_limits.push((provider, rps));
        self
    }

    /// Override the default egress RPS cap.
    #[must_use]
    pub fn egress_default_rps(mut self, rps: u32) -> Self {
        self.egress_default_rps = rps;
        self
    }

    /// Override the default per-tenant LLM RPS cap applied to any
    /// `(tenant, provider)` pair the builder did not pre-seed via
    /// [`Self::tenant_llm_limit`]. Defaults to
    /// [`DEFAULT_TENANT_LLM_RPS`]. Setting `0` would never grant a
    /// permit; the builder accepts it because some test profiles want
    /// to assert universal refusal.
    #[must_use]
    pub fn tenant_llm_default_rps(mut self, rps: u32) -> Self {
        self.tenant_llm_default_rps = rps;
        self
    }

    /// Pre-seed an explicit `(tenant, provider)` LLM RPS cap. The
    /// override replaces the default cap on the first acquire for that
    /// pair and persists for the governor's lifetime.
    #[must_use]
    pub fn tenant_llm_limit(mut self, tenant: TenantId, provider: ProviderId, rps: u32) -> Self {
        self.tenant_llm_limits.push((tenant, provider, rps));
        self
    }

    /// Finalise.
    #[must_use]
    pub fn build(self) -> TokenBucketGovernor {
        let llm: DashMap<ProviderId, Arc<Bucket>> = DashMap::new();
        for (p, rps) in &self.llm_limits {
            llm.insert(p.clone(), Arc::new(Bucket::new(*rps)));
        }
        let tenant_llm: DashMap<(TenantId, ProviderId), Arc<Bucket>> = DashMap::new();
        for (t, p, rps) in &self.tenant_llm_limits {
            tenant_llm.insert((t.clone(), p.clone()), Arc::new(Bucket::new(*rps)));
        }
        let g = TokenBucketGovernor {
            kv: self.kv,
            llm_buckets: llm,
            tenant_llm_buckets: tenant_llm,
            egress_buckets: DashMap::new(),
            egress_default_rps: self.egress_default_rps,
            tenant_llm_default_rps: self.tenant_llm_default_rps,
        };
        g.spawn_refiller();
        g
    }
}

struct Bucket {
    sem: Arc<Semaphore>,
    rps: u32,
    fenced: std::sync::atomic::AtomicBool,
}

impl Bucket {
    fn new(rps: u32) -> Self {
        Self {
            sem: Arc::new(Semaphore::new(usize::try_from(rps).unwrap_or(usize::MAX))),
            rps,
            fenced: std::sync::atomic::AtomicBool::new(false),
        }
    }

    fn snapshot(&self) -> BudgetSnapshot {
        BudgetSnapshot {
            remaining: self.sem.available_permits() as i64,
            limit: i64::from(self.rps),
        }
    }
}

/// Token-bucket governor with per-provider LLM caps + per-host egress caps.
pub struct TokenBucketGovernor {
    #[allow(dead_code)]
    kv: Arc<dyn KvStore>,
    llm_buckets: DashMap<ProviderId, Arc<Bucket>>,
    tenant_llm_buckets: DashMap<(TenantId, ProviderId), Arc<Bucket>>,
    egress_buckets: DashMap<Host, Arc<Bucket>>,
    egress_default_rps: u32,
    tenant_llm_default_rps: u32,
}

impl TokenBucketGovernor {
    /// Start a builder; `kv` is reserved for cross-process coordination in
    /// Phase B but is required up-front so the API stays stable.
    #[must_use]
    pub fn builder(kv: Arc<dyn KvStore>) -> TokenBucketGovernorBuilder {
        TokenBucketGovernorBuilder {
            kv,
            llm_limits: Vec::new(),
            egress_default_rps: DEFAULT_EGRESS_RPS,
            tenant_llm_default_rps: DEFAULT_TENANT_LLM_RPS,
            tenant_llm_limits: Vec::new(),
        }
    }

    fn spawn_refiller(&self) {
        let llm = self.llm_buckets.clone();
        let tenant_llm = self.tenant_llm_buckets.clone();
        let egress = self.egress_buckets.clone();
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(Duration::from_secs(1));
            loop {
                tick.tick().await;
                refill_map(&llm);
                refill_map(&tenant_llm);
                refill_map(&egress);
            }
        });
    }

    /// Look up (or lazily allocate) the bucket for `provider`. Returns
    /// `None` when the provider was never configured — the inbound
    /// gate treats this as fail-closed `Saturated`, never lazy-grant.
    fn provider_bucket(&self, provider: &ProviderId) -> Option<Arc<Bucket>> {
        self.llm_buckets.get(provider).map(|e| e.value().clone())
    }

    /// Look up (or lazily allocate) the bucket for a `(tenant,
    /// provider)` pair. Tenants not pre-seeded inherit the bounded
    /// `tenant_llm_default_rps` — never unlimited.
    fn tenant_bucket(&self, tenant: &TenantId, provider: &ProviderId) -> Arc<Bucket> {
        let key = (tenant.clone(), provider.clone());
        self.tenant_llm_buckets
            .entry(key)
            .or_insert_with(|| Arc::new(Bucket::new(self.tenant_llm_default_rps)))
            .value()
            .clone()
    }

    async fn acquire_from(
        bucket: Arc<Bucket>,
        scope_label: String,
    ) -> Result<Permit, GovernorError> {
        if bucket.fenced.load(std::sync::atomic::Ordering::Acquire) {
            return Err(GovernorError::Saturated { scope: scope_label });
        }
        let sem = Arc::clone(&bucket.sem);
        let timeout = tokio::time::timeout(DEFAULT_ACQUIRE_TIMEOUT, sem.acquire_owned()).await;
        match timeout {
            Ok(Ok(p)) => {
                p.forget();
                let sem_ref = Arc::clone(&bucket.sem);
                Ok(Permit::new(move || {
                    sem_ref.add_permits(1);
                }))
            }
            Ok(Err(_)) => Err(GovernorError::Unavailable(format!(
                "semaphore closed for {scope_label}"
            ))),
            Err(_) => Err(GovernorError::TimedOut {
                millis: DEFAULT_ACQUIRE_TIMEOUT.as_millis() as u64,
            }),
        }
    }

    /// Fail-fast acquire used by tenant-scoped LLM requests. An empty
    /// bucket returns [`GovernorError::Saturated`] immediately so the
    /// inbound MCP gate surfaces `LlmError::RateLimit` to the peer
    /// instead of holding the request for the 30 s
    /// [`DEFAULT_ACQUIRE_TIMEOUT`]. Distinct from
    /// [`Self::acquire_from`] (which back-pressures supervisor-side
    /// callers) so the provider path stays unchanged.
    fn try_acquire_from(bucket: Arc<Bucket>, scope_label: String) -> Result<Permit, GovernorError> {
        if bucket.fenced.load(std::sync::atomic::Ordering::Acquire) {
            return Err(GovernorError::Saturated { scope: scope_label });
        }
        let sem = Arc::clone(&bucket.sem);
        match sem.try_acquire_owned() {
            Ok(p) => {
                p.forget();
                let sem_ref = Arc::clone(&bucket.sem);
                Ok(Permit::new(move || {
                    sem_ref.add_permits(1);
                }))
            }
            Err(tokio::sync::TryAcquireError::NoPermits) => {
                Err(GovernorError::Saturated { scope: scope_label })
            }
            Err(tokio::sync::TryAcquireError::Closed) => Err(GovernorError::Unavailable(format!(
                "semaphore closed for {scope_label}"
            ))),
        }
    }
}

fn refill_map<K>(map: &DashMap<K, Arc<Bucket>>)
where
    K: std::hash::Hash + Eq,
{
    for entry in map.iter() {
        let b = entry.value();
        if b.fenced.load(std::sync::atomic::Ordering::Acquire) {
            continue;
        }
        let want = usize::try_from(b.rps).unwrap_or(usize::MAX);
        let have = b.sem.available_permits();
        if have < want {
            b.sem.add_permits(want - have);
        }
    }
}

#[async_trait]
impl Governor for TokenBucketGovernor {
    async fn acquire_llm(
        &self,
        provider: ProviderId,
        _est_tokens: u32,
    ) -> Result<Permit, GovernorError> {
        let bucket = self
            .provider_bucket(&provider)
            .ok_or_else(|| GovernorError::Saturated {
                scope: format!("provider:{provider}"),
            })?;
        Self::acquire_from(bucket, format!("llm:{provider}")).await
    }

    async fn acquire_llm_scoped(
        &self,
        scope: BudgetScope,
        _est_tokens: u32,
    ) -> Result<Permit, GovernorError> {
        match scope {
            BudgetScope::Provider(p) => {
                let bucket = self
                    .provider_bucket(&p)
                    .ok_or_else(|| GovernorError::Saturated {
                        scope: format!("provider:{p}"),
                    })?;
                Self::acquire_from(bucket, format!("llm:{p}")).await
            }
            BudgetScope::TenantProvider(t, p) => {
                let bucket = self.tenant_bucket(&t, &p);
                // Fail-fast: a saturated tenant bucket surfaces
                // `Saturated` immediately so the inbound MCP gate can
                // map it to `LlmError::RateLimit` without holding the
                // peer for the 30 s acquire budget.
                Self::try_acquire_from(bucket, format!("llm:tenant:{t}:provider:{p}"))
            }
            other => Err(GovernorError::Unavailable(format!(
                "TokenBucketGovernor::acquire_llm_scoped: unsupported scope {other}"
            ))),
        }
    }

    async fn acquire_egress(&self, host: Host) -> Result<Permit, GovernorError> {
        let bucket = self
            .egress_buckets
            .entry(host.clone())
            .or_insert_with(|| Arc::new(Bucket::new(self.egress_default_rps)))
            .value()
            .clone();
        Self::acquire_from(bucket, format!("egress:{host}")).await
    }

    async fn budget(&self, scope: BudgetScope) -> BudgetSnapshot {
        match scope {
            BudgetScope::Provider(p) => self
                .llm_buckets
                .get(&p)
                .map(|e| e.value().snapshot())
                .unwrap_or(BudgetSnapshot {
                    remaining: 0,
                    limit: 0,
                }),
            _ => BudgetSnapshot {
                remaining: 0,
                limit: 0,
            },
        }
    }

    async fn fence(&self, scope: BudgetScope, _limit: BudgetLimit) -> Result<(), GovernorError> {
        if let BudgetScope::Provider(p) = scope {
            if let Some(b) = self.llm_buckets.get(&p) {
                b.fenced.store(true, std::sync::atomic::Ordering::Release);
            }
        }
        Ok(())
    }
}