klieo-mcp-server 2.2.0

//! Inbound per-tenant LLM budget governor.
//!
//! Holds the `(Governor, ProviderId)` pair the builder accumulates via
//! [`crate::McpServerBuilder::with_governor`] and exposes a helper that
//! every workflow / governed agent invoker uses to wrap `ctx.llm` once
//! the per-request tenant label is known.
//!
//! Module gated on `feature = "governor"` so adopters that skip the
//! per-tenant budget gate keep the lean default dep graph (no
//! `klieo-ops` dependency).

use klieo_core::agent::AgentContext;
use klieo_ops::governor::{GovernedLlmClient, Governor};
use klieo_ops::types::{ProviderId, TenantId};
use std::sync::Arc;

/// The `(Governor, ProviderId)` pair every per-tenant wrap needs.
/// Held behind an [`Arc`] on the builder so adding a workflow or
/// governed agent does not clone the underlying `Governor` impl.
#[derive(Clone)]
pub(crate) struct GovernorBundle {
    pub(crate) governor: Arc<dyn Governor>,
    pub(crate) provider: ProviderId,
}

/// Wrap `ctx.llm` in a [`GovernedLlmClient`] keyed off the tenant
/// label already attached to `ctx`. A `ctx.tenant_label = None`
/// degrades to provider-scoped governance — the inbound gate's
/// builder-level hard gate guarantees this only happens on the
/// `add_agent` path (workflows refuse to build without a governor
/// AND a verified principal stamps the label before this fires).
pub(crate) fn wrap_ctx_with_governor(ctx: AgentContext, bundle: &GovernorBundle) -> AgentContext {
    let llm = ctx.llm.clone();
    let episodic = ctx.episodic.clone();
    let run_id = ctx.run_id;
    let tenant = ctx.tenant_label().map(|label| TenantId(label.to_string()));
    ctx.with_llm(Arc::new(GovernedLlmClient::new(
        llm,
        bundle.governor.clone(),
        episodic,
        run_id,
        bundle.provider.clone(),
        tenant,
    )))
}