mlua_swarm/middleware/resolver.rs
1//! Routing resolver — dynamic agent resolution at spawn time. Slotted into
2//! the `SpawnerStack` as a `SpawnerLayer`. Treats `Ctx.agent` as a hint,
3//! rewrites it to the actual dispatch target, and hands the updated `Ctx`
4//! down to the inner spawner.
5//!
6//! Examples:
7//! - `Ctx.agent = "best-coder"` — the resolver picks `claude` or `gpt-4` or `llama`.
8//! - `Ctx.agent = "router:by-prompt"` — branch on prompt contents.
9//! - `Ctx.agent = "ensemble"` — fan out to multiple agents (delegated to a
10//! different layer or a fanout stage).
11//!
12//! One of the axes that a plugin can drive (see handoff §Plugin).
13
14use crate::core::ctx::Ctx;
15use crate::core::engine::Engine;
16use crate::middleware::SpawnerLayer;
17use crate::types::{CapToken, TaskId};
18use crate::worker::adapter::{SpawnError, SpawnerAdapter};
19use crate::worker::Worker;
20use async_trait::async_trait;
21use std::sync::Arc;
22
23/// Routing resolver trait. Takes the `Ctx.agent` hint and returns a real
24/// agent name the inner spawner can resolve. Sync is enough — this is meant
25/// for light lookups. Push heavy resolvers into a separate spawner layer or
26/// a Lua plugin.
27///
28/// The `directive` argument was removed in the current design: prompts now travel
29/// through engine state and no longer appear in spawner arguments. If you
30/// need prompt-content-driven routing, either have the resolver call
31/// `engine.fetch_prompt(token, task_id)` from a separate layer, or implement
32/// a dedicated prompt-based routing layer (carry).
33pub trait AgentResolver: Send + Sync + 'static {
34 /// `agent_hint` is the raw `Ctx.agent` value. The returned string is
35 /// installed as the new `Ctx.agent` before the inner spawner is called.
36 fn resolve(&self, agent_hint: &str, ctx: &Ctx) -> String;
37}
38
39/// Wrapper that lets a closure act as an `AgentResolver` via a blanket impl.
40pub struct FnResolver<F>(
41 /// The closure implementing `Fn(&str, &Ctx) -> String`.
42 pub F,
43);
44
45impl<F> AgentResolver for FnResolver<F>
46where
47 F: Fn(&str, &Ctx) -> String + Send + Sync + 'static,
48{
49 fn resolve(&self, hint: &str, ctx: &Ctx) -> String {
50 (self.0)(hint, ctx)
51 }
52}
53
54/// `SpawnerLayer` implementation — inject into a `SpawnerStack` and use.
55pub struct ResolverMiddleware {
56 resolver: Arc<dyn AgentResolver>,
57}
58
59impl ResolverMiddleware {
60 /// Wraps an existing `AgentResolver` implementation.
61 pub fn new(resolver: Arc<dyn AgentResolver>) -> Self {
62 Self { resolver }
63 }
64
65 /// Convenience constructor: wraps a plain closure as the resolver via
66 /// `FnResolver`.
67 pub fn from_fn<F>(f: F) -> Self
68 where
69 F: Fn(&str, &Ctx) -> String + Send + Sync + 'static,
70 {
71 Self {
72 resolver: Arc::new(FnResolver(f)),
73 }
74 }
75}
76
77impl SpawnerLayer for ResolverMiddleware {
78 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
79 Arc::new(ResolverWrapped {
80 inner,
81 resolver: self.resolver.clone(),
82 })
83 }
84}
85
86struct ResolverWrapped {
87 inner: Arc<dyn SpawnerAdapter>,
88 resolver: Arc<dyn AgentResolver>,
89}
90
91#[async_trait]
92impl SpawnerAdapter for ResolverWrapped {
93 async fn spawn(
94 &self,
95 engine: &Engine,
96 ctx: &Ctx,
97 task_id: TaskId,
98 attempt: u32,
99 token: CapToken,
100 ) -> Result<Box<dyn Worker>, SpawnError> {
101 let resolved = self.resolver.resolve(&ctx.agent, ctx);
102 if resolved == ctx.agent {
103 // no-op: hint and resolved are the same
104 self.inner.spawn(engine, ctx, task_id, attempt, token).await
105 } else {
106 // Clone ctx and overwrite `agent`, then hand it to `inner`.
107 let mut new_ctx = ctx.clone();
108 new_ctx.agent = resolved;
109 self.inner
110 .spawn(engine, &new_ctx, task_id, attempt, token)
111 .await
112 }
113 }
114}