klieo-ops 3.3.0

//! [`LlmClient`] wrapper that acquires a [`Governor`] permit before each
//! call and emits [`OpsEvent::GovernorDenial`] into [`EpisodicMemory`]
//! on saturation or timeout.

use crate::emit::emit_ops_event;
use crate::governor::{Governor, GovernorError};
use crate::ops_event::OpsEvent;
use crate::tenant::current_tenant;
use crate::types::{BudgetScope, ProviderId, TenantId};
use async_trait::async_trait;
use klieo_core::error::LlmError;
use klieo_core::ids::RunId;
use klieo_core::llm::{Capabilities, ChatRequest, ChatResponse, ChunkStream};
use klieo_core::memory::EpisodicMemory;
use klieo_core::LlmClient;
use std::sync::Arc;

/// Conservative per-call token estimate handed to the governor when
/// the wrapper has no provider-reported figure yet. Lives behind a
/// named const so token-budget tuning (Phase B TPM follow-up) is a
/// single-site change.
const ESTIMATED_TOKENS_PER_CALL: u32 = 4096;

/// Wraps an inner [`LlmClient`] with a [`Governor`]-acquired permit per call.
///
/// Every call to [`complete`] or [`stream`] acquires an LLM permit from the
/// governor before delegating. Metadata-only calls ([`name`], [`capabilities`],
/// [`embed`]) are forwarded unconditionally — they carry no inference cost.
///
/// On governor refusal, [`OpsEvent::GovernorDenial`] is written into the
/// provided [`EpisodicMemory`] and the corresponding [`LlmError`] is returned.
#[non_exhaustive]
pub struct GovernedLlmClient {
    inner: Arc<dyn LlmClient>,
    governor: Arc<dyn Governor>,
    episodic: Arc<dyn EpisodicMemory>,
    run_id: RunId,
    provider: ProviderId,
    /// `Some(tenant)` routes every acquire through
    /// `BudgetScope::TenantProvider`, applying the per-tenant cap on
    /// top of the provider cap. `None` preserves the legacy
    /// supervisor-side `BudgetScope::Provider` semantics.
    tenant: Option<TenantId>,
}

impl GovernedLlmClient {
    /// Construct a governed wrapper around `inner`.
    ///
    /// Pass `tenant = Some(_)` when the caller has a verified tenant
    /// label (e.g. inbound MCP after `with_tenant_label`); pass `None`
    /// for the legacy in-process supervisor path that already runs
    /// inside a single tenant ambient.
    #[must_use]
    pub fn new(
        inner: Arc<dyn LlmClient>,
        governor: Arc<dyn Governor>,
        episodic: Arc<dyn EpisodicMemory>,
        run_id: RunId,
        provider: ProviderId,
        tenant: Option<TenantId>,
    ) -> Self {
        Self {
            inner,
            governor,
            episodic,
            run_id,
            provider,
            tenant,
        }
    }

    async fn acquire_permit(&self) -> Result<crate::governor::Permit, LlmError> {
        let scope = match self.tenant.clone() {
            Some(tenant) => BudgetScope::TenantProvider(tenant, self.provider.clone()),
            None => BudgetScope::Provider(self.provider.clone()),
        };
        self.governor
            .acquire_llm_scoped(scope, ESTIMATED_TOKENS_PER_CALL)
            .await
            .map_err(|e| self.translate_governor_error(e))
    }

    fn translate_governor_error(&self, e: GovernorError) -> LlmError {
        match e {
            GovernorError::Saturated { .. } => LlmError::RateLimit {
                retry_after_secs: 1,
            },
            GovernorError::TimedOut { .. } => LlmError::Timeout,
            GovernorError::Unavailable(s) => LlmError::Server(s),
        }
    }

    async fn record_denial(&self, reason: String) {
        if let Err(err) = emit_ops_event(
            &*self.episodic,
            self.run_id,
            OpsEvent::GovernorDenial {
                tenant: current_tenant(),
                resource: "llm".into(),
                provider: Some(self.provider.clone()),
                host: None,
                reason,
            },
        )
        .await
        {
            tracing::warn!(
                target: "klieo.ops.audit",
                error = %err,
                "audit emit failed; episode not recorded"
            );
        }
    }
}

#[async_trait]
impl LlmClient for GovernedLlmClient {
    fn name(&self) -> &str {
        self.inner.name()
    }

    fn capabilities(&self) -> &Capabilities {
        self.inner.capabilities()
    }

    async fn complete(&self, req: ChatRequest) -> Result<ChatResponse, LlmError> {
        match self.acquire_permit().await {
            Ok(_permit) => self.inner.complete(req).await,
            Err(e) => {
                self.record_denial(format!("{e}")).await;
                Err(e)
            }
        }
    }

    async fn stream(&self, req: ChatRequest) -> Result<ChunkStream, LlmError> {
        match self.acquire_permit().await {
            Ok(_permit) => self.inner.stream(req).await,
            Err(e) => {
                self.record_denial(format!("{e}")).await;
                Err(e)
            }
        }
    }

    async fn embed(&self, texts: &[String]) -> Result<Vec<klieo_core::llm::Embedding>, LlmError> {
        // Embeddings carry inference cost; govern them the same as completions.
        match self.acquire_permit().await {
            Ok(_permit) => self.inner.embed(texts).await,
            Err(e) => {
                self.record_denial(format!("{e}")).await;
                Err(e)
            }
        }
    }
}