klieo_ops/runtime/
supervised.rs

1//! Generic agent wrapper that auto-emits Supervisor / Governor / Gate
2//! events into the agent's [`EpisodicMemory`] without requiring user code
3//! to call [`emit_ops_event`] manually.
4
5use crate::emit::emit_ops_event;
6use crate::ops_event::OpsEvent;
7use crate::run_context::{scope_run_context, RunContext};
8use crate::tenant::{current_tenant, scope_tenant, TenantResolver};
9use crate::types::AgentId;
10use async_trait::async_trait;
11use klieo_core::{Agent, AgentContext, ToolDef};
12use std::sync::Arc;
13#[cfg(feature = "gates")]
14use std::time::Duration;
15
16#[cfg(feature = "gates")]
17use crate::runtime::gated_invoker::GatedToolInvoker;
18
19#[cfg(feature = "governor")]
20use crate::runtime::governed_llm::GovernedLlmClient;
21
22/// Wrapper around `A: Agent` that:
23///
24/// - Resolves a tenant and binds it via [`scope_tenant`] for the duration
25///   of `run`.
26/// - Wraps `ctx.llm` with a `GovernedLlmClient` (private wrapper) when a
27///   [`crate::governor::Governor`] is configured (`governor` feature flag required).
28/// - Wraps `ctx.tools` with a [`crate::runtime::gated_invoker::GatedToolInvoker`]
29///   when any [`crate::gates::Gate`]s are configured (`gates` feature flag required).
30/// - Emits [`crate::OpsEvent::SupervisorStateChange`] around the inner agent call.
31pub struct SupervisedAgent<A: Agent> {
32    pub(crate) inner: A,
33    pub(crate) tenant_resolver: Arc<dyn TenantResolver>,
34    /// Stable agent identity recorded in [`OpsEvent::SupervisorStateChange`].
35    pub(crate) agent_id: AgentId,
36    #[cfg(feature = "governor")]
37    pub(crate) governor: Option<Arc<dyn crate::governor::Governor>>,
38    #[cfg(feature = "governor")]
39    pub(crate) governor_provider: crate::types::ProviderId,
40    #[cfg(feature = "gates")]
41    pub(crate) gates: Vec<Arc<dyn crate::gates::Gate>>,
42    #[cfg(feature = "gates")]
43    pub(crate) approval_timeout: Duration,
44    #[cfg(feature = "gates")]
45    pub(crate) approval_timeout_policy: crate::gates::ApprovalTimeoutPolicy,
46}
47
48#[async_trait]
49impl<A> Agent for SupervisedAgent<A>
50where
51    A: Agent + 'static,
52    A::Input: Send + 'static,
53    A::Output: Send + 'static,
54    A::Error: Send + Sync + 'static,
55{
56    type Input = A::Input;
57    type Output = A::Output;
58    type Error = A::Error;
59
60    fn name(&self) -> &str {
61        self.inner.name()
62    }
63
64    fn system_prompt(&self) -> &str {
65        self.inner.system_prompt()
66    }
67
68    fn tools(&self) -> &[ToolDef] {
69        self.inner.tools()
70    }
71
72    async fn run(
73        &self,
74        ctx: AgentContext,
75        input: Self::Input,
76    ) -> Result<Self::Output, Self::Error> {
77        let tenant = self.tenant_resolver.resolve().await;
78        let agent_id = self.agent_id.clone();
79        let episodic = ctx.episodic.clone();
80        let run_id = ctx.run_id;
81
82        let ctx = self.wrap_llm(ctx);
83        let ctx = self.wrap_tools(ctx);
84
85        let run_ctx = RunContext {
86            episodic: episodic.clone(),
87            run_id,
88        };
89
90        scope_tenant(
91            tenant,
92            scope_run_context(Some(run_ctx), async move {
93                if let Err(err) = emit_ops_event(
94                    &*episodic,
95                    run_id,
96                    OpsEvent::SupervisorStateChange {
97                        tenant: current_tenant(),
98                        agent: agent_id.clone(),
99                        state: "running".into(),
100                        reason: None,
101                    },
102                )
103                .await
104                {
105                    tracing::warn!(
106                        target: "klieo.ops.audit",
107                        error = %err,
108                        "audit emit failed; episode not recorded"
109                    );
110                }
111
112                let result = self.inner.run(ctx, input).await;
113
114                let state = if result.is_ok() {
115                    "completed"
116                } else {
117                    "failed"
118                };
119                if let Err(err) = emit_ops_event(
120                    &*episodic,
121                    run_id,
122                    OpsEvent::SupervisorStateChange {
123                        tenant: current_tenant(),
124                        agent: agent_id,
125                        state: state.into(),
126                        reason: None,
127                    },
128                )
129                .await
130                {
131                    tracing::warn!(
132                        target: "klieo.ops.audit",
133                        error = %err,
134                        "audit emit failed; episode not recorded"
135                    );
136                }
137
138                result
139            }),
140        )
141        .await
142    }
143}
144
145impl<A: Agent> SupervisedAgent<A> {
146    /// Wrap `ctx.llm` in a [`GovernedLlmClient`] when a governor is present.
147    ///
148    /// When the `governor` feature is disabled the context is returned
149    /// unchanged.
150    fn wrap_llm(&self, ctx: AgentContext) -> AgentContext {
151        #[cfg(feature = "governor")]
152        if let Some(g) = &self.governor {
153            let llm = ctx.llm.clone();
154            let episodic = ctx.episodic.clone();
155            let run_id = ctx.run_id;
156            return ctx.with_llm(Arc::new(GovernedLlmClient::new(
157                llm,
158                g.clone(),
159                episodic,
160                run_id,
161                self.governor_provider.clone(),
162                None,
163            )));
164        }
165        ctx
166    }
167
168    /// Wrap `ctx.tools` in a [`GatedToolInvoker`] when gates are registered.
169    ///
170    /// When the `gates` feature is disabled the context is returned unchanged.
171    fn wrap_tools(&self, ctx: AgentContext) -> AgentContext {
172        #[cfg(feature = "gates")]
173        if !self.gates.is_empty() {
174            let tools = ctx.tools.clone();
175            let episodic = ctx.episodic.clone();
176            let run_id = ctx.run_id;
177            return ctx.with_tools(Arc::new(
178                GatedToolInvoker::new(tools, self.gates.clone(), episodic, run_id)
179                    .with_approval_timeout(self.approval_timeout)
180                    .with_approval_timeout_policy(self.approval_timeout_policy),
181            ));
182        }
183        ctx
184    }
185}
klieo_ops/runtime/supervised.rs

klieo_ops/runtime/
supervised.rs
