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#[derive(Clone)]
21pub struct AgentRoute {
22 descriptor: AgentDescriptor,
23 agent: Arc<Agent>,
24 capabilities: CapabilitySet,
25}
26
27impl AgentRoute {
28 pub fn new(descriptor: AgentDescriptor, agent: Arc<Agent>) -> Self {
30 Self {
31 descriptor,
32 agent,
33 capabilities: CapabilitySet::new(),
34 }
35 }
36
37 #[must_use]
41 pub fn with_capabilities(mut self, capabilities: CapabilitySet) -> Self {
42 self.capabilities = capabilities;
43 self
44 }
45
46 pub const fn descriptor(&self) -> &AgentDescriptor {
48 &self.descriptor
49 }
50
51 pub fn agent(&self) -> &Arc<Agent> {
53 &self.agent
54 }
55
56 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#[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 pub fn new() -> Self {
94 Self::default()
95 }
96
97 #[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 #[must_use]
109 pub fn layer(mut self, middleware: Arc<dyn GatewayMiddleware>) -> Self {
110 self.middleware.push(middleware);
111 self
112 }
113
114 pub fn push_middleware(&mut self, middleware: Arc<dyn GatewayMiddleware>) {
116 self.middleware.push(middleware);
117 }
118
119 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 pub fn contains(&self, name: &str) -> bool {
139 self.routes.contains_key(name)
140 }
141
142 pub fn descriptor(&self, name: &str) -> Option<&AgentDescriptor> {
144 self.routes.get(name).map(|route| &route.descriptor)
145 }
146
147 pub fn model_specs(&self) -> Vec<ToolSpec> {
149 self.routes
150 .values()
151 .map(|route| route.descriptor.model_spec())
152 .collect()
153 }
154
155 pub fn len(&self) -> usize {
157 self.routes.len()
158 }
159
160 pub fn is_empty(&self) -> bool {
162 self.routes.is_empty()
163 }
164
165 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#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
343#[non_exhaustive]
344pub enum GatewayErrorKind {
345 NotFound,
347 InvalidInput,
349 CapabilityDenied,
351 AuthorityEscalation,
353 MaxDepth,
355 BudgetExceeded,
357 Cancelled,
359 DeadlineExceeded,
361 ChildFailed,
363 PolicyDenied,
365 ObservabilityFailed,
367}
368
369#[derive(Clone, Debug, Deserialize, Error, Eq, PartialEq, Serialize)]
371#[error("{kind:?}: {message}")]
372pub struct GatewayError {
373 pub kind: GatewayErrorKind,
375 pub message: String,
377}
378
379impl GatewayError {
380 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#[derive(Clone, Debug, Error, Eq, PartialEq)]
426#[non_exhaustive]
427pub enum AgentRegistrationError {
428 #[error("agent route name cannot be empty")]
430 EmptyName,
431 #[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}