Skip to main content

runifold_agent/
gateway.rs

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