Skip to main content

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}