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#[derive(Clone)]
19pub struct AgentRoute {
20 descriptor: AgentDescriptor,
21 agent: Arc<Agent>,
22 capabilities: CapabilitySet,
23}
24
25impl AgentRoute {
26 pub fn new(descriptor: AgentDescriptor, agent: Arc<Agent>) -> Self {
28 Self {
29 descriptor,
30 agent,
31 capabilities: CapabilitySet::new(),
32 }
33 }
34
35 #[must_use]
39 pub fn with_capabilities(mut self, capabilities: CapabilitySet) -> Self {
40 self.capabilities = capabilities;
41 self
42 }
43
44 pub const fn descriptor(&self) -> &AgentDescriptor {
46 &self.descriptor
47 }
48
49 pub fn agent(&self) -> &Arc<Agent> {
51 &self.agent
52 }
53
54 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#[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 pub fn new() -> Self {
92 Self::default()
93 }
94
95 #[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 #[must_use]
107 pub fn layer(mut self, middleware: Arc<dyn GatewayMiddleware>) -> Self {
108 self.middleware.push(middleware);
109 self
110 }
111
112 pub fn push_middleware(&mut self, middleware: Arc<dyn GatewayMiddleware>) {
114 self.middleware.push(middleware);
115 }
116
117 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 pub fn contains(&self, name: &str) -> bool {
137 self.routes.contains_key(name)
138 }
139
140 pub fn descriptor(&self, name: &str) -> Option<&AgentDescriptor> {
142 self.routes.get(name).map(|route| &route.descriptor)
143 }
144
145 pub fn model_specs(&self) -> Vec<ToolSpec> {
147 self.routes
148 .values()
149 .map(|route| route.descriptor.model_spec())
150 .collect()
151 }
152
153 pub fn len(&self) -> usize {
155 self.routes.len()
156 }
157
158 pub fn is_empty(&self) -> bool {
160 self.routes.is_empty()
161 }
162
163 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#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
339#[non_exhaustive]
340pub enum GatewayErrorKind {
341 NotFound,
343 InvalidInput,
345 CapabilityDenied,
347 AuthorityEscalation,
349 MaxDepth,
351 BudgetExceeded,
353 Cancelled,
355 DeadlineExceeded,
357 ChildFailed,
359 PolicyDenied,
361 ObservabilityFailed,
363}
364
365#[derive(Clone, Debug, Deserialize, Error, Eq, PartialEq, Serialize)]
367#[error("{kind:?}: {message}")]
368pub struct GatewayError {
369 pub kind: GatewayErrorKind,
371 pub message: String,
373}
374
375impl GatewayError {
376 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#[derive(Clone, Debug, Error, Eq, PartialEq)]
422#[non_exhaustive]
423pub enum AgentRegistrationError {
424 #[error("agent route name cannot be empty")]
426 EmptyName,
427 #[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}