klieo-ops 3.3.0

//! Generic agent wrapper that auto-emits Supervisor / Governor / Gate
//! events into the agent's [`EpisodicMemory`] without requiring user code
//! to call [`emit_ops_event`] manually.

use crate::emit::emit_ops_event;
use crate::ops_event::OpsEvent;
use crate::run_context::{scope_run_context, RunContext};
use crate::tenant::{current_tenant, scope_tenant, TenantResolver};
use crate::types::AgentId;
use async_trait::async_trait;
use klieo_core::{Agent, AgentContext, ToolDef};
use std::sync::Arc;
#[cfg(feature = "gates")]
use std::time::Duration;

#[cfg(feature = "gates")]
use crate::runtime::gated_invoker::GatedToolInvoker;

#[cfg(feature = "governor")]
use crate::runtime::governed_llm::GovernedLlmClient;

/// Wrapper around `A: Agent` that:
///
/// - Resolves a tenant and binds it via [`scope_tenant`] for the duration
///   of `run`.
/// - Wraps `ctx.llm` with a `GovernedLlmClient` (private wrapper) when a
///   [`crate::governor::Governor`] is configured (`governor` feature flag required).
/// - Wraps `ctx.tools` with a [`crate::runtime::gated_invoker::GatedToolInvoker`]
///   when any [`crate::gates::Gate`]s are configured (`gates` feature flag required).
/// - Emits [`crate::OpsEvent::SupervisorStateChange`] around the inner agent call.
pub struct SupervisedAgent<A: Agent> {
    pub(crate) inner: A,
    pub(crate) tenant_resolver: Arc<dyn TenantResolver>,
    /// Stable agent identity recorded in [`OpsEvent::SupervisorStateChange`].
    pub(crate) agent_id: AgentId,
    #[cfg(feature = "governor")]
    pub(crate) governor: Option<Arc<dyn crate::governor::Governor>>,
    #[cfg(feature = "governor")]
    pub(crate) governor_provider: crate::types::ProviderId,
    #[cfg(feature = "gates")]
    pub(crate) gates: Vec<Arc<dyn crate::gates::Gate>>,
    #[cfg(feature = "gates")]
    pub(crate) approval_timeout: Duration,
    #[cfg(feature = "gates")]
    pub(crate) approval_timeout_policy: crate::gates::ApprovalTimeoutPolicy,
}

#[async_trait]
impl<A> Agent for SupervisedAgent<A>
where
    A: Agent + 'static,
    A::Input: Send + 'static,
    A::Output: Send + 'static,
    A::Error: Send + Sync + 'static,
{
    type Input = A::Input;
    type Output = A::Output;
    type Error = A::Error;

    fn name(&self) -> &str {
        self.inner.name()
    }

    fn system_prompt(&self) -> &str {
        self.inner.system_prompt()
    }

    fn tools(&self) -> &[ToolDef] {
        self.inner.tools()
    }

    async fn run(
        &self,
        ctx: AgentContext,
        input: Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        let tenant = self.tenant_resolver.resolve().await;
        let agent_id = self.agent_id.clone();
        let episodic = ctx.episodic.clone();
        let run_id = ctx.run_id;

        let ctx = self.wrap_llm(ctx);
        let ctx = self.wrap_tools(ctx);

        let run_ctx = RunContext {
            episodic: episodic.clone(),
            run_id,
        };

        scope_tenant(
            tenant,
            scope_run_context(Some(run_ctx), async move {
                if let Err(err) = emit_ops_event(
                    &*episodic,
                    run_id,
                    OpsEvent::SupervisorStateChange {
                        tenant: current_tenant(),
                        agent: agent_id.clone(),
                        state: "running".into(),
                        reason: None,
                    },
                )
                .await
                {
                    tracing::warn!(
                        target: "klieo.ops.audit",
                        error = %err,
                        "audit emit failed; episode not recorded"
                    );
                }

                let result = self.inner.run(ctx, input).await;

                let state = if result.is_ok() {
                    "completed"
                } else {
                    "failed"
                };
                if let Err(err) = emit_ops_event(
                    &*episodic,
                    run_id,
                    OpsEvent::SupervisorStateChange {
                        tenant: current_tenant(),
                        agent: agent_id,
                        state: state.into(),
                        reason: None,
                    },
                )
                .await
                {
                    tracing::warn!(
                        target: "klieo.ops.audit",
                        error = %err,
                        "audit emit failed; episode not recorded"
                    );
                }

                result
            }),
        )
        .await
    }
}

impl<A: Agent> SupervisedAgent<A> {
    /// Wrap `ctx.llm` in a [`GovernedLlmClient`] when a governor is present.
    ///
    /// When the `governor` feature is disabled the context is returned
    /// unchanged.
    fn wrap_llm(&self, ctx: AgentContext) -> AgentContext {
        #[cfg(feature = "governor")]
        if let Some(g) = &self.governor {
            let llm = ctx.llm.clone();
            let episodic = ctx.episodic.clone();
            let run_id = ctx.run_id;
            return ctx.with_llm(Arc::new(GovernedLlmClient::new(
                llm,
                g.clone(),
                episodic,
                run_id,
                self.governor_provider.clone(),
                None,
            )));
        }
        ctx
    }

    /// Wrap `ctx.tools` in a [`GatedToolInvoker`] when gates are registered.
    ///
    /// When the `gates` feature is disabled the context is returned unchanged.
    fn wrap_tools(&self, ctx: AgentContext) -> AgentContext {
        #[cfg(feature = "gates")]
        if !self.gates.is_empty() {
            let tools = ctx.tools.clone();
            let episodic = ctx.episodic.clone();
            let run_id = ctx.run_id;
            return ctx.with_tools(Arc::new(
                GatedToolInvoker::new(tools, self.gates.clone(), episodic, run_id)
                    .with_approval_timeout(self.approval_timeout)
                    .with_approval_timeout_policy(self.approval_timeout_policy),
            ));
        }
        ctx
    }
}