Skip to main content

runifold_agent/
gateway.rs

1use std::{collections::BTreeMap, fmt, sync::Arc, time::Instant};
2
3use runifold_core::{BudgetEvent, CapabilitySet, ChildEvent, RunContext, RunEventKind, Usage};
4use runifold_model::ToolSpec;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use thiserror::Error;
8
9use crate::{
10    Agent, AgentDescriptor, AgentError, AgentOutcome, DelegationRequest, GatewayMiddleware,
11    GatewayNext,
12};
13
14const DELEGATION_DEPTH_KEY: &str = "runifold.agent.delegation_depth";
15const DELEGATED_AGENT_KEY: &str = "runifold.agent.delegated_agent";
16
17/// One explicitly configured route from a caller to a child agent.
18#[derive(Clone)]
19pub struct AgentRoute {
20    descriptor: AgentDescriptor,
21    agent: Arc<Agent>,
22    capabilities: CapabilitySet,
23}
24
25impl AgentRoute {
26    /// Creates a route whose child starts with no capabilities.
27    pub fn new(descriptor: AgentDescriptor, agent: Arc<Agent>) -> Self {
28        Self {
29            descriptor,
30            agent,
31            capabilities: CapabilitySet::new(),
32        }
33    }
34
35    /// Sets the exact capabilities requested for the child run.
36    ///
37    /// Invocation still rejects any grant not held by the parent.
38    #[must_use]
39    pub fn with_capabilities(mut self, capabilities: CapabilitySet) -> Self {
40        self.capabilities = capabilities;
41        self
42    }
43
44    /// Returns the model-facing and policy-facing route contract.
45    pub const fn descriptor(&self) -> &AgentDescriptor {
46        &self.descriptor
47    }
48
49    /// Returns the configured child agent.
50    pub fn agent(&self) -> &Arc<Agent> {
51        &self.agent
52    }
53
54    /// Returns the exact child capability grant.
55    pub const fn capabilities(&self) -> &CapabilitySet {
56        &self.capabilities
57    }
58}
59
60impl fmt::Debug for AgentRoute {
61    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
62        formatter
63            .debug_struct("AgentRoute")
64            .field("descriptor", &self.descriptor)
65            .field("agent", &self.agent.name())
66            .field("capabilities", &self.capabilities)
67            .finish()
68    }
69}
70
71/// Capability-gated router for parent-to-child agent delegation.
72#[derive(Clone)]
73pub struct AgentGateway {
74    routes: BTreeMap<String, AgentRoute>,
75    middleware: Vec<Arc<dyn GatewayMiddleware>>,
76    max_depth: u32,
77}
78
79impl Default for AgentGateway {
80    fn default() -> Self {
81        Self {
82            routes: BTreeMap::new(),
83            middleware: Vec::new(),
84            max_depth: 8,
85        }
86    }
87}
88
89impl AgentGateway {
90    /// Creates an empty gateway.
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    /// Sets the maximum delegation depth accepted by this gateway.
96    #[must_use]
97    pub fn with_max_depth(mut self, max_depth: u32) -> Self {
98        self.max_depth = max_depth;
99        self
100    }
101
102    /// Appends around-middleware to the gateway chain.
103    ///
104    /// Middleware runs in registration order before the child boundary and in
105    /// reverse order after `next` completes.
106    #[must_use]
107    pub fn layer(mut self, middleware: Arc<dyn GatewayMiddleware>) -> Self {
108        self.middleware.push(middleware);
109        self
110    }
111
112    /// Appends around-middleware without consuming the gateway.
113    pub fn push_middleware(&mut self, middleware: Arc<dyn GatewayMiddleware>) {
114        self.middleware.push(middleware);
115    }
116
117    /// Registers a route without replacing an existing model-facing name.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`AgentRegistrationError`] when the route name is blank or
122    /// already registered.
123    pub fn register(&mut self, route: AgentRoute) -> Result<(), AgentRegistrationError> {
124        let name = route.descriptor.name.trim();
125        if name.is_empty() {
126            return Err(AgentRegistrationError::EmptyName);
127        }
128        if self.routes.contains_key(name) {
129            return Err(AgentRegistrationError::DuplicateName(name.into()));
130        }
131        self.routes.insert(name.into(), route);
132        Ok(())
133    }
134
135    /// Returns whether a route is registered under `name`.
136    pub fn contains(&self, name: &str) -> bool {
137        self.routes.contains_key(name)
138    }
139
140    /// Returns the immutable route descriptor registered under `name`.
141    pub fn descriptor(&self, name: &str) -> Option<&AgentDescriptor> {
142        self.routes.get(name).map(|route| &route.descriptor)
143    }
144
145    /// Returns model-facing route specifications in deterministic name order.
146    pub fn model_specs(&self) -> Vec<ToolSpec> {
147        self.routes
148            .values()
149            .map(|route| route.descriptor.model_spec())
150            .collect()
151    }
152
153    /// Returns the number of registered routes.
154    pub fn len(&self) -> usize {
155        self.routes.len()
156    }
157
158    /// Returns whether no routes are registered.
159    pub fn is_empty(&self) -> bool {
160        self.routes.is_empty()
161    }
162
163    /// Invokes a child agent through an explicitly authorized route.
164    ///
165    /// # Errors
166    ///
167    /// Returns [`GatewayError`] when lifecycle, capability, authority, depth,
168    /// budget, or child execution checks fail.
169    pub async fn delegate(
170        &self,
171        name: &str,
172        input: impl Into<String>,
173        parent: &RunContext,
174    ) -> Result<AgentOutcome, GatewayError> {
175        let route = self.routes.get(name).ok_or_else(|| {
176            GatewayError::new(
177                GatewayErrorKind::NotFound,
178                format!("agent route `{name}` is not registered"),
179            )
180        })?;
181        let request =
182            DelegationRequest::new(route.descriptor.clone(), input.into(), parent.clone());
183        GatewayNext {
184            middleware: &self.middleware,
185            route,
186            max_depth: self.max_depth,
187            index: 0,
188        }
189        .run(request)
190        .await
191    }
192}
193
194impl fmt::Debug for AgentGateway {
195    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
196        formatter
197            .debug_struct("AgentGateway")
198            .field("routes", &self.routes)
199            .field("middleware_count", &self.middleware.len())
200            .field("max_depth", &self.max_depth)
201            .finish()
202    }
203}
204
205pub(crate) async fn execute_route(
206    route: &AgentRoute,
207    max_depth: u32,
208    request: DelegationRequest,
209) -> Result<AgentOutcome, GatewayError> {
210    let parent = request.parent();
211    let name = &request.descriptor().name;
212    let depth = validate_route(route, max_depth, parent, name)?;
213    let usage = parent
214        .budget()
215        .try_consume(Usage {
216            delegations: 1,
217            ..Usage::default()
218        })
219        .map_err(|error| GatewayError::new(GatewayErrorKind::BudgetExceeded, error.to_string()))?;
220    parent
221        .record(RunEventKind::Budget(BudgetEvent::Updated { usage }), None)
222        .map_err(observability_error)?;
223
224    let mut child = parent.child(route.capabilities.clone());
225    child
226        .metadata_mut()
227        .insert(DELEGATION_DEPTH_KEY.into(), Value::from(depth + 1));
228    child
229        .metadata_mut()
230        .insert(DELEGATED_AGENT_KEY.into(), Value::from(name.clone()));
231
232    let child_started = parent
233        .record(
234            RunEventKind::Child(ChildEvent::Started {
235                child_run_id: child.run_id(),
236            }),
237            None,
238        )
239        .map_err(observability_error)?
240        .map(|event| event.meta.event_id);
241    if let Some(event_id) = child_started {
242        child = child.with_cause(event_id);
243    }
244
245    let result = route
246        .agent
247        .run(request.input().to_owned(), &child)
248        .await
249        .map_err(|error| GatewayError::from_agent(&error));
250    record_child_terminal(parent, child.run_id(), child_started, &result)?;
251    result
252}
253
254fn validate_route(
255    route: &AgentRoute,
256    max_depth: u32,
257    parent: &RunContext,
258    name: &str,
259) -> Result<u64, GatewayError> {
260    if parent.cancellation().is_cancelled() {
261        return Err(GatewayError::new(
262            GatewayErrorKind::Cancelled,
263            "delegation was cancelled before the child run started",
264        ));
265    }
266    if parent
267        .deadline()
268        .is_some_and(|deadline| deadline <= Instant::now())
269    {
270        return Err(GatewayError::new(
271            GatewayErrorKind::DeadlineExceeded,
272            "delegation deadline elapsed before the child run started",
273        ));
274    }
275    if max_depth == 0 {
276        return Err(GatewayError::new(
277            GatewayErrorKind::MaxDepth,
278            "gateway max_depth must be greater than zero",
279        ));
280    }
281    if !parent.capabilities().contains(route.descriptor.id) {
282        return Err(GatewayError::new(
283            GatewayErrorKind::CapabilityDenied,
284            format!("run is not granted agent capability `{name}`"),
285        ));
286    }
287    if let Some(missing) = route.capabilities.first_missing_from(parent.capabilities()) {
288        return Err(GatewayError::new(
289            GatewayErrorKind::AuthorityEscalation,
290            format!(
291                "child agent `{name}` requested capability `{}` not held by its parent",
292                missing.name
293            ),
294        ));
295    }
296
297    let depth = delegation_depth(parent);
298    if depth >= u64::from(max_depth) {
299        return Err(GatewayError::new(
300            GatewayErrorKind::MaxDepth,
301            format!("delegation depth {depth} reached gateway maximum {max_depth}"),
302        ));
303    }
304    Ok(depth)
305}
306
307fn record_child_terminal(
308    parent: &RunContext,
309    child_run_id: runifold_core::RunId,
310    child_started: Option<runifold_core::EventId>,
311    result: &Result<AgentOutcome, GatewayError>,
312) -> Result<(), GatewayError> {
313    let child_event = match result {
314        Ok(_) => ChildEvent::Completed { child_run_id },
315        Err(error) if error.kind == GatewayErrorKind::Cancelled => {
316            ChildEvent::Cancelled { child_run_id }
317        }
318        Err(_) => ChildEvent::Failed { child_run_id },
319    };
320    parent
321        .record(RunEventKind::Child(child_event), child_started)
322        .map_err(observability_error)?;
323    Ok(())
324}
325
326fn observability_error(error: runifold_core::JournalError) -> GatewayError {
327    GatewayError::new(GatewayErrorKind::ObservabilityFailed, error.message)
328}
329
330fn delegation_depth(run: &RunContext) -> u64 {
331    run.metadata()
332        .get(DELEGATION_DEPTH_KEY)
333        .and_then(Value::as_u64)
334        .unwrap_or(0)
335}
336
337/// Normalized gateway failure category.
338#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
339#[non_exhaustive]
340pub enum GatewayErrorKind {
341    /// The requested route is not registered.
342    NotFound,
343    /// The delegation input did not match the canonical schema.
344    InvalidInput,
345    /// The parent was not granted the agent capability.
346    CapabilityDenied,
347    /// The configured child grant would amplify parent authority.
348    AuthorityEscalation,
349    /// The configured delegation-depth bound was reached.
350    MaxDepth,
351    /// The shared run-tree budget rejected the delegation.
352    BudgetExceeded,
353    /// Delegated work was cancelled.
354    Cancelled,
355    /// Delegated work exceeded its effective deadline.
356    DeadlineExceeded,
357    /// The child agent failed without violating a hard runtime invariant.
358    ChildFailed,
359    /// Gateway middleware or policy denied the delegation.
360    PolicyDenied,
361    /// The configured journal rejected a gateway event.
362    ObservabilityFailed,
363}
364
365/// Structured failure from the agent delegation boundary.
366#[derive(Clone, Debug, Deserialize, Error, Eq, PartialEq, Serialize)]
367#[error("{kind:?}: {message}")]
368pub struct GatewayError {
369    /// Normalized category.
370    pub kind: GatewayErrorKind,
371    /// Safe human-readable explanation.
372    pub message: String,
373}
374
375impl GatewayError {
376    /// Creates a gateway error.
377    pub fn new(kind: GatewayErrorKind, message: impl Into<String>) -> Self {
378        Self {
379            kind,
380            message: message.into(),
381        }
382    }
383
384    fn from_agent(error: &AgentError) -> Self {
385        let kind = match error {
386            AgentError::Model(error)
387                if matches!(error.kind, runifold_model::ModelErrorKind::Cancelled) =>
388            {
389                GatewayErrorKind::Cancelled
390            }
391            AgentError::Model(error)
392                if matches!(error.kind, runifold_model::ModelErrorKind::DeadlineExceeded) =>
393            {
394                GatewayErrorKind::DeadlineExceeded
395            }
396            AgentError::Tool(error)
397                if matches!(error.kind, runifold_tool::ToolErrorKind::CapabilityDenied) =>
398            {
399                GatewayErrorKind::CapabilityDenied
400            }
401            AgentError::Tool(error)
402                if matches!(error.kind, runifold_tool::ToolErrorKind::Cancelled) =>
403            {
404                GatewayErrorKind::Cancelled
405            }
406            AgentError::Tool(error)
407                if matches!(error.kind, runifold_tool::ToolErrorKind::DeadlineExceeded) =>
408            {
409                GatewayErrorKind::DeadlineExceeded
410            }
411            AgentError::Budget(_) => GatewayErrorKind::BudgetExceeded,
412            AgentError::Gateway(error) => error.kind.clone(),
413            AgentError::Journal(_) => GatewayErrorKind::ObservabilityFailed,
414            _ => GatewayErrorKind::ChildFailed,
415        };
416        Self::new(kind, error.to_string())
417    }
418}
419
420/// Failure to add an agent route to a gateway.
421#[derive(Clone, Debug, Error, Eq, PartialEq)]
422#[non_exhaustive]
423pub enum AgentRegistrationError {
424    /// Agent route names must not be blank.
425    #[error("agent route name cannot be empty")]
426    EmptyName,
427    /// Another route already owns the model-facing name.
428    #[error("agent route `{0}` is already registered")]
429    DuplicateName(String),
430}
431
432#[cfg(test)]
433mod tests {
434    use std::sync::Arc;
435
436    use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
437    use runifold_model::ModelRef;
438    use runifold_testkit::ScriptedModel;
439
440    use crate::{
441        Agent, AgentDescriptor, AgentGateway, AgentRoute, GatewayErrorKind,
442        gateway::DELEGATION_DEPTH_KEY,
443    };
444
445    fn gateway_and_run() -> (AgentGateway, RunContext, ScriptedModel) {
446        let model = ScriptedModel::new();
447        let child = Arc::new(Agent::new(
448            "child",
449            Arc::new(model.clone()),
450            ModelRef::new("test", "child"),
451        ));
452        let descriptor = AgentDescriptor::new("ask_child", "Delegate work");
453        let mut gateway = AgentGateway::new();
454        gateway
455            .register(AgentRoute::new(descriptor.clone(), child))
456            .unwrap();
457        let mut capabilities = CapabilitySet::new();
458        capabilities.grant(descriptor.capability());
459        let run = RunContext::root(BudgetTracker::new(Budget::default()), capabilities);
460        (gateway, run, model)
461    }
462
463    #[test]
464    fn preexisting_cancellation_stops_before_budget_or_child_execution() {
465        let (gateway, run, model) = gateway_and_run();
466        run.cancellation().cancel();
467
468        let error =
469            futures_executor::block_on(gateway.delegate("ask_child", "work", &run)).unwrap_err();
470
471        assert_eq!(error.kind, GatewayErrorKind::Cancelled);
472        assert_eq!(run.budget().usage().delegations, 0);
473        assert!(model.recorded_requests().is_empty());
474    }
475
476    #[test]
477    fn depth_limit_stops_before_budget_or_child_execution() {
478        let (gateway, mut run, model) = gateway_and_run();
479        let gateway = gateway.with_max_depth(1);
480        run.metadata_mut()
481            .insert(DELEGATION_DEPTH_KEY.into(), serde_json::json!(1));
482
483        let error =
484            futures_executor::block_on(gateway.delegate("ask_child", "work", &run)).unwrap_err();
485
486        assert_eq!(error.kind, GatewayErrorKind::MaxDepth);
487        assert_eq!(run.budget().usage().delegations, 0);
488        assert!(model.recorded_requests().is_empty());
489    }
490}