Skip to main content

mlua_swarm/blueprint/
compiler.rs

1//! Blueprint `Compiler`, `CompiledAgentTable`, and the three default
2//! `SpawnerFactory` implementations.
3//!
4//! ## Pipeline
5//!
6//! ```text
7//! Blueprint (= flow + agents + hints + strategy + spawner_hints)
8//!     │
9//!     │ Compiler.compile(&bp)          ← this module (AgentDef → SpawnerAdapter table)
10//!     ▼
11//! CompiledBlueprint {
12//!     router: Arc<CompiledAgentTable>, // ctx.agent → SpawnerAdapter lookup
13//!     flow:   FlowNode,                // the flow.ir source (evaluated via EngineDispatcher)
14//!     metadata: BlueprintMetadata,
15//! }
16//!     │
17//!     │ service::linker::link(router, blueprint.spawner_hints.layers, &engine)
18//!     ▼                                   ↑ Layer wrapping is done separately (src/service/linker.rs)
19//! `Arc<dyn SpawnerAdapter>`            (already wrapped with base + hint SpawnerLayers)
20//!     │
21//!     ▼ EngineDispatcher::with_spawner → engine.dispatch_attempt_with
22//! ```
23//!
24//! `CompiledAgentTable` is a thin table: it looks up `routes[name]` by
25//! `ctx.agent` and hands the spawn off to the matching `SpawnerAdapter`.
26//! The `routes` map is built at compile time through `SpawnerFactory`
27//! implementations. Layer wrapping is not part of this module — it lives
28//! in `service::linker::link`.
29
30use crate::blueprint::{AgentDef, AgentKind, Blueprint, BlueprintMetadata};
31use crate::core::ctx::Ctx;
32use crate::core::engine::Engine;
33use crate::operator::{Operator, OperatorSpawner, WorkerBinding};
34use crate::types::{CapToken, StepId};
35use crate::worker::adapter::{InProcSpawner, SpawnError, SpawnerAdapter, WorkerFn};
36use crate::worker::process_spawner::{ProcessSpawner, StreamMode};
37use crate::worker::Worker;
38use async_trait::async_trait;
39use mlua_flow_ir::{Expr, Node as FlowNode};
40use serde_json::Value;
41use std::collections::HashMap;
42use std::sync::Arc;
43use thiserror::Error;
44
45// ─── error ───────────────────────────────────────────────────────────────
46
47/// Everything that can go wrong while `Compiler::compile` turns a
48/// `Blueprint` into a `CompiledBlueprint`.
49#[derive(Debug, Error)]
50pub enum CompileError {
51    /// An `AgentDef.kind` has no matching entry in the `SpawnerRegistry`
52    /// and `Blueprint.strategy.strict_kind` is set.
53    #[error("unknown agent kind in SpawnerRegistry: {0:?}")]
54    UnknownKind(AgentKind),
55    /// The `AgentDef.spec` shape did not match what the factory for its
56    /// kind requires (missing/mistyped field, etc.).
57    #[error("agent '{name}' spec invalid: {msg}")]
58    InvalidSpec {
59        /// The offending agent's name.
60        name: String,
61        /// Human-readable description of what was wrong with the spec.
62        msg: String,
63    },
64    /// The flow references an agent name that has no corresponding
65    /// `AgentDef` (and no default spawner is configured).
66    #[error("flow references agent '{0}' but no AgentDef matches")]
67    UnresolvedRef(String),
68    /// Two `AgentDef`s in the same `Blueprint` share a name.
69    #[error("duplicate AgentDef name: {0}")]
70    DuplicateAgent(String),
71    /// A `kind = Operator` agent's `spec.operator_ref` does not match
72    /// any `OperatorDef.name` declared in `Blueprint.operators`.
73    #[error("agent '{agent}' operator_ref '{op_ref}' does not match any OperatorDef.name in Blueprint.operators (defined: {defined:?})")]
74    UnresolvedOperatorRef {
75        /// The agent whose `operator_ref` didn't resolve.
76        agent: String,
77        /// The `operator_ref` value that was looked up.
78        op_ref: String,
79        /// The `OperatorDef.name`s that *are* declared, for the error
80        /// message.
81        defined: Vec<String>,
82    },
83    /// GH #21 Phase 2: an `AgentMeta.meta_ref` or a statically-visible
84    /// `$step_meta.ref` (inside a `Step.in` **Lit** expr) does not match
85    /// any `MetaDef.name` declared in `Blueprint.metas`.
86    #[error("{where_} names an undefined MetaDef: '{meta_ref}' (defined: {defined:?})")]
87    UnresolvedMetaRef {
88        /// Human-readable description of where the reference was found
89        /// (e.g. `"AgentMeta.meta_ref of agent 'planner'"` or `"Step
90        /// 'scout' $step_meta.ref"`).
91        where_: String,
92        /// The `meta_ref` value that was looked up.
93        meta_ref: String,
94        /// The `MetaDef.name`s that *are* declared, for the error
95        /// message.
96        defined: Vec<String>,
97    },
98}
99
100// ─── SpawnerFactory + Registry ───────────────────────────────────────────
101
102/// Factory trait that interprets an `AgentDef` and builds the concrete
103/// `SpawnerAdapter`. Register one per kind. Parsing the spec,
104/// validating it, and baking the profile are the implementation's job.
105///
106/// The signature was widened in v9 from `(name, spec, hint)` to
107/// `(&AgentDef, hint)` so the profile can be passed through. Most
108/// implementations still just pull `&agent_def.name` and
109/// `&agent_def.spec`, but Operator-backend factories consume
110/// `agent_def.profile` to bake the persona in.
111pub trait SpawnerFactory: Send + Sync {
112    /// Build the concrete `SpawnerAdapter` for one `AgentDef`. `hint` is
113    /// the matching entry (if any) from `Blueprint.hints.per_agent`.
114    fn build(
115        &self,
116        agent_def: &AgentDef,
117        hint: Option<&Value>,
118    ) -> Result<Arc<dyn SpawnerAdapter>, CompileError>;
119}
120
121/// Companion trait that carries the **type-side source of truth** for
122/// the Adapter ↔ `AgentKind` correspondence.
123///
124/// The base [`SpawnerFactory`] trait deliberately does not carry an
125/// associated const so it stays dyn-compatible — that is, so it can be
126/// stored and dispatched as `Arc<dyn SpawnerFactory>`. This companion
127/// trait splits `const KIND: AgentKind` out, and
128/// [`SpawnerRegistry::register`] uses `F::KIND` as the `HashMap` key.
129/// That physically removes the string-lookup failure mode at the type
130/// layer.
131///
132/// The three built-in factories (`Shell` / `InProc` / `Operator`)
133/// implement this. Extension backends (say, `AgentBlockSpawnerFactory`)
134/// follow the same explicit two-step recipe: add a new `AgentKind`
135/// variant and implement this trait.
136pub trait SpawnerFactoryKind: SpawnerFactory {
137    /// The `AgentKind` this factory handles — used as the `HashMap` key
138    /// by `SpawnerRegistry::register`.
139    const KIND: AgentKind;
140    /// The concrete Worker type produced by this `AgentKind` — this
141    /// binds the type chain all the way from `AgentKind` down to `Worker`.
142    /// Every factory declares it so the `AgentKind → Worker` mapping is
143    /// explicit across all four layers. It is the source of truth for
144    /// preserving the concrete type right up until `SpawnerAdapter::spawn`
145    /// erases it into `Box<dyn Worker>`.
146    type Worker: crate::worker::Worker;
147}
148
149/// `AgentKind → SpawnerFactory` mapping. The compiler looks entries up
150/// during `compile()`.
151#[derive(Clone)]
152pub struct SpawnerRegistry {
153    factories: HashMap<AgentKind, Arc<dyn SpawnerFactory>>,
154}
155
156impl SpawnerRegistry {
157    /// Start with an empty `AgentKind → SpawnerFactory` mapping.
158    pub fn new() -> Self {
159        Self {
160            factories: HashMap::new(),
161        }
162    }
163    /// **Type-driven registration** — takes `F::KIND` and uses it as the
164    /// `HashMap` key.
165    ///
166    /// Callers use the form
167    /// `reg.register::<SubprocessProcessSpawnerFactory>(Arc::new(...))`
168    /// and never have to pass an `AgentKind` literal. The Adapter ↔ Kind
169    /// correspondence is enforced at the type layer, physically removing
170    /// the string / enum-literal lookup failure mode.
171    pub fn register<F: SpawnerFactoryKind + 'static>(&mut self, factory: Arc<F>) -> &mut Self {
172        let f: Arc<dyn SpawnerFactory> = factory;
173        self.factories.insert(F::KIND, f);
174        self
175    }
176}
177
178impl Default for SpawnerRegistry {
179    fn default() -> Self {
180        Self::new()
181    }
182}
183
184// ─── Compiler ────────────────────────────────────────────────────────────
185
186/// Turns a `Blueprint` into a `CompiledBlueprint` by resolving every
187/// `AgentDef` against a `SpawnerRegistry`. One-shot: build a fresh
188/// `Compiler` per `compile()` call (or reuse it — it holds no
189/// per-compile state).
190pub struct Compiler {
191    registry: SpawnerRegistry,
192    default_spawner: Option<Arc<dyn SpawnerAdapter>>,
193}
194
195/// The result of `Compiler::compile` — a routing table plus the
196/// unmodified flow and metadata, ready to hand to
197/// `EngineDispatcher::with_spawner` / `mlua_flow_ir::eval_async`.
198pub struct CompiledBlueprint {
199    /// `ctx.agent → SpawnerAdapter` lookup table.
200    pub router: Arc<CompiledAgentTable>,
201    /// The flow.ir source, copied verbatim from `Blueprint.flow`.
202    pub flow: FlowNode,
203    /// Copied verbatim from `Blueprint.metadata`.
204    pub metadata: BlueprintMetadata,
205}
206
207impl Compiler {
208    /// Build a `Compiler` around the given `SpawnerRegistry`, with no
209    /// default spawner (unresolved flow refs are an error unless
210    /// `with_default` is chained on).
211    pub fn new(registry: SpawnerRegistry) -> Self {
212        Self {
213            registry,
214            default_spawner: None,
215        }
216    }
217
218    /// Set a default spawner — used for flow refs (and unregistered
219    /// `AgentKind`s under non-strict strategy) that don't resolve
220    /// against any `AgentDef`/`SpawnerRegistry` entry.
221    pub fn with_default(mut self, sp: Arc<dyn SpawnerAdapter>) -> Self {
222        self.default_spawner = Some(sp);
223        self
224    }
225
226    /// Resolve every `Blueprint.agents` entry through the registry,
227    /// validate `operator_ref`s and flow refs per `Blueprint.strategy`,
228    /// and return the routing table alongside the untouched flow and
229    /// metadata.
230    pub fn compile(&self, bp: &Blueprint) -> Result<CompiledBlueprint, CompileError> {
231        let mut routes: HashMap<String, Arc<dyn SpawnerAdapter>> = HashMap::new();
232        let mut seen: HashMap<String, ()> = HashMap::new();
233
234        // Design-time validation (OperatorDef as a first-class value):
235        // every `kind = Operator` agent's `spec.operator_ref` must point at
236        // one of `bp.operators[].name`. A Blueprint with any Operator agent
237        // must therefore declare its operators up front; the empty-operators
238        // backward-compat bypass is retired.
239        let defined: Vec<String> = bp.operators.iter().map(|o| o.name.clone()).collect();
240        for ad in &bp.agents {
241            if !matches!(ad.kind, AgentKind::Operator) {
242                continue;
243            }
244            let op_ref = ad.spec.get("operator_ref").and_then(|v| v.as_str());
245            if let Some(op_ref) = op_ref {
246                if !defined.iter().any(|n| n == op_ref) {
247                    return Err(CompileError::UnresolvedOperatorRef {
248                        agent: ad.name.clone(),
249                        op_ref: op_ref.to_string(),
250                        defined: defined.clone(),
251                    });
252                }
253            }
254            // A missing `op_ref` is reported through OperatorSpawnerFactory.build under a different error.
255        }
256
257        // GH #21 Phase 2: named `MetaDef` pool (`Blueprint.metas`) —
258        // validate every reference against it, mirroring the
259        // `operator_ref` validation above.
260        let metas_defined: Vec<String> = bp.metas.iter().map(|m| m.name.clone()).collect();
261        for ad in &bp.agents {
262            let meta_ref = ad.meta.as_ref().and_then(|m| m.meta_ref.as_ref());
263            if let Some(meta_ref) = meta_ref {
264                if !metas_defined.iter().any(|n| n == meta_ref) {
265                    return Err(CompileError::UnresolvedMetaRef {
266                        where_: format!("AgentMeta.meta_ref of agent '{}'", ad.name),
267                        meta_ref: meta_ref.clone(),
268                        defined: metas_defined.clone(),
269                    });
270                }
271            }
272        }
273        // Best-effort static walk of the flow for `$step_meta.ref`
274        // envelopes embedded in a Step's **Lit** `in` expr — this is a
275        // design-time hint only: a non-`Lit` `Step.in` (e.g. `Path`) is
276        // invisible here and skipped silently; `EngineDispatcher::dispatch`
277        // is the authoritative, loud validation line for those.
278        let mut static_step_meta_refs: Vec<(String, String)> = Vec::new();
279        collect_step_meta_refs(&bp.flow, &mut static_step_meta_refs);
280        for (where_, meta_ref) in static_step_meta_refs {
281            if !metas_defined.iter().any(|n| n == &meta_ref) {
282                return Err(CompileError::UnresolvedMetaRef {
283                    where_,
284                    meta_ref,
285                    defined: metas_defined.clone(),
286                });
287            }
288        }
289
290        for ad in &bp.agents {
291            if seen.contains_key(&ad.name) {
292                return Err(CompileError::DuplicateAgent(ad.name.clone()));
293            }
294            seen.insert(ad.name.clone(), ());
295
296            let factory = match self.registry.factories.get(&ad.kind) {
297                Some(f) => f.clone(),
298                None => {
299                    if bp.strategy.strict_kind {
300                        return Err(CompileError::UnknownKind(ad.kind.clone()));
301                    } else {
302                        tracing::warn!(
303                            agent = %ad.name,
304                            kind = ?ad.kind,
305                            "no spawner factory registered for agent kind; \
306                             dropping agent from routing table (strict_kind=false)"
307                        );
308                        continue;
309                    }
310                }
311            };
312            let hint = bp.hints.per_agent.get(&ad.name);
313            let spawner = factory.build(ad, hint)?;
314            routes.insert(ad.name.clone(), spawner);
315        }
316
317        if bp.strategy.strict_refs {
318            verify_refs(&bp.flow, &routes, self.default_spawner.is_some())?;
319        }
320
321        let router = Arc::new(CompiledAgentTable {
322            routes,
323            default: self.default_spawner.clone(),
324        });
325        Ok(CompiledBlueprint {
326            router,
327            flow: bp.flow.clone(),
328            metadata: bp.metadata.clone(),
329        })
330    }
331}
332
333/// Walk the flow `Node`, collect every `Step.ref`, and check that no ref
334/// is unresolved against `routes` (or the default, when one exists).
335fn verify_refs(
336    node: &FlowNode,
337    routes: &HashMap<String, Arc<dyn SpawnerAdapter>>,
338    has_default: bool,
339) -> Result<(), CompileError> {
340    let mut refs: Vec<String> = Vec::new();
341    collect_refs(node, &mut refs);
342    for r in refs {
343        if !routes.contains_key(&r) && !has_default {
344            return Err(CompileError::UnresolvedRef(r));
345        }
346    }
347    Ok(())
348}
349
350fn collect_refs(node: &FlowNode, out: &mut Vec<String>) {
351    match node {
352        FlowNode::Step { ref_, .. } => out.push(ref_.clone()),
353        FlowNode::Seq { children } => {
354            for c in children {
355                collect_refs(c, out);
356            }
357        }
358        FlowNode::Branch { then_, else_, .. } => {
359            collect_refs(then_, out);
360            collect_refs(else_, out);
361        }
362        FlowNode::Fanout { body, .. } => collect_refs(body, out),
363        FlowNode::Loop { body, .. } => collect_refs(body, out),
364        FlowNode::Try { body, catch, .. } => {
365            collect_refs(body, out);
366            collect_refs(catch, out);
367        }
368        FlowNode::Assign { .. } => {} // The Assign node carries no ref.
369    }
370}
371
372/// GH #21 Phase 2: walk the flow `Node` (same recursion shape as
373/// [`collect_refs`]) and collect every statically-visible `$step_meta.ref`
374/// found inside a Step's `in` **Lit** expr, as `(where_, meta_ref)` pairs
375/// for [`CompileError::UnresolvedMetaRef`] reporting. Non-`Lit` `in`
376/// exprs (e.g. `Expr::Path`) cannot be inspected statically and are
377/// silently skipped — `EngineDispatcher::dispatch` (the `mlua-swarm` core
378/// crate) is the authoritative, loud validation line for those.
379fn collect_step_meta_refs(node: &FlowNode, out: &mut Vec<(String, String)>) {
380    match node {
381        FlowNode::Step { ref_, in_, .. } => {
382            if let Expr::Lit { value } = in_ {
383                if let Some(meta_ref) = static_step_meta_ref(value) {
384                    out.push((format!("Step '{ref_}' $step_meta.ref"), meta_ref));
385                }
386            }
387        }
388        FlowNode::Seq { children } => {
389            for c in children {
390                collect_step_meta_refs(c, out);
391            }
392        }
393        FlowNode::Branch { then_, else_, .. } => {
394            collect_step_meta_refs(then_, out);
395            collect_step_meta_refs(else_, out);
396        }
397        FlowNode::Fanout { body, .. } => collect_step_meta_refs(body, out),
398        FlowNode::Loop { body, .. } => collect_step_meta_refs(body, out),
399        FlowNode::Try { body, catch, .. } => {
400            collect_step_meta_refs(body, out);
401            collect_step_meta_refs(catch, out);
402        }
403        FlowNode::Assign { .. } => {} // The Assign node carries no `in`.
404    }
405}
406
407/// Extract the `$step_meta.ref` string out of a literal `Step.in` value,
408/// if present and well-formed: `{"$step_meta": {"ref": "<name>", ...},
409/// ...}`. Any other shape (no `$step_meta` key, `ref` absent/null, `ref`
410/// not a string) yields `None` — this is a best-effort static hint only;
411/// a malformed envelope is caught loudly at dispatch time instead (see
412/// `EngineDispatcher::dispatch`'s doc in the `mlua-swarm` core crate).
413fn static_step_meta_ref(value: &Value) -> Option<String> {
414    value
415        .as_object()?
416        .get("$step_meta")?
417        .as_object()?
418        .get("ref")?
419        .as_str()
420        .map(str::to_string)
421}
422
423// ─── CompiledAgentTable ───────────────────────────────────────────────────────
424
425/// The compile result: an `agent name → SpawnerAdapter` lookup table.
426///
427/// Looks `routes` up by `ctx.agent` (the flow.ir `Step.ref`) and hands
428/// the spawn to the matching `SpawnerAdapter`. If the name is not
429/// registered and a `default` is configured, the default is used; if
430/// there is no default, `SpawnError::NotRegistered` is returned.
431///
432/// Layer wrapping (`AuditMiddleware` / `MainAIMiddleware` and friends) is
433/// not this type's concern — that is done separately in
434/// `service::linker::link`.
435pub struct CompiledAgentTable {
436    pub(crate) routes: HashMap<String, Arc<dyn SpawnerAdapter>>,
437    pub(crate) default: Option<Arc<dyn SpawnerAdapter>>,
438}
439
440impl CompiledAgentTable {
441    /// Whether the given agent name is registered in the table — i.e.,
442    /// whether its spawner has been resolved.
443    pub fn has_route(&self, agent: &str) -> bool {
444        self.routes.contains_key(agent)
445    }
446    /// List every resolved agent name.
447    pub fn routed_agents(&self) -> Vec<String> {
448        self.routes.keys().cloned().collect()
449    }
450}
451
452#[async_trait]
453impl SpawnerAdapter for CompiledAgentTable {
454    async fn spawn(
455        &self,
456        engine: &Engine,
457        ctx: &Ctx,
458        task_id: StepId,
459        attempt: u32,
460        token: CapToken,
461    ) -> Result<Box<dyn Worker>, SpawnError> {
462        let sp = self
463            .routes
464            .get(&ctx.agent)
465            .cloned()
466            .or_else(|| self.default.clone())
467            .ok_or_else(|| SpawnError::NotRegistered(ctx.agent.clone()))?;
468        sp.spawn(engine, ctx, task_id, attempt, token).await
469    }
470}
471
472// ─── default factories (three variants) ───────────────────────────────────
473
474/// Factory for `AgentKind::Subprocess`. Turns the spec into a
475/// [`ProcessSpawner`].
476///
477/// Naming convention: `<WorkerIMPL><AdapterType>SpawnerFactory`. Factory
478/// names carry both the worker implementation and the host adapter so
479/// they are not confused with each other; the old
480/// `ShellSpawnerFactory` was renamed to this.
481///
482/// Spec shape:
483/// ```jsonc
484/// { "program": "agent-block", "args": ["-s","s.lua"],
485///   "use_stdin": true,                       // optional, default = true
486///   "stream_mode": "ndjson_lines" | "sse_events" | "length_prefixed" | null  // optional, default = null (plain)
487/// }
488/// ```
489pub struct SubprocessProcessSpawnerFactory;
490
491impl SpawnerFactoryKind for SubprocessProcessSpawnerFactory {
492    const KIND: AgentKind = AgentKind::Subprocess;
493    type Worker = crate::worker::process_spawner::ProcessWorker;
494}
495
496impl SpawnerFactory for SubprocessProcessSpawnerFactory {
497    fn build(
498        &self,
499        agent_def: &AgentDef,
500        _hint: Option<&Value>,
501    ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
502        let agent_name = &agent_def.name;
503        let spec = &agent_def.spec;
504        let invalid = |msg: String| CompileError::InvalidSpec {
505            name: agent_name.to_string(),
506            msg,
507        };
508        let program = spec
509            .get("program")
510            .and_then(|v| v.as_str())
511            .ok_or_else(|| invalid("shell spec: 'program' (string) required".into()))?
512            .to_string();
513        let args: Vec<String> = spec
514            .get("args")
515            .and_then(|v| v.as_array())
516            .map(|a| {
517                a.iter()
518                    .filter_map(|x| x.as_str().map(|s| s.to_string()))
519                    .collect()
520            })
521            .unwrap_or_default();
522        let use_stdin = spec
523            .get("use_stdin")
524            .and_then(|v| v.as_bool())
525            .unwrap_or(true);
526        let stream_mode = match spec.get("stream_mode").and_then(|v| v.as_str()) {
527            Some("ndjson_lines") => Some(StreamMode::NdjsonLines),
528            Some("sse_events") => Some(StreamMode::SseEvents),
529            Some("length_prefixed") => Some(StreamMode::LengthPrefixed),
530            Some(other) => return Err(invalid(format!("unknown stream_mode: {other}"))),
531            None => None,
532        };
533
534        let mut sp = ProcessSpawner {
535            program,
536            args,
537            use_stdin,
538            stream_mode,
539        };
540        if let Some(mode) = sp.stream_mode.clone() {
541            sp = sp.stream_mode(mode);
542        }
543        Ok(Arc::new(sp))
544    }
545}
546
547/// Factory for `AgentKind::Lua`. At `build` time it inspects the
548/// `AgentDef.spec` and returns an [`InProcSpawner`] with the Lua-eval
549/// `WorkerFn` registered under `agent_name` — one `InProcSpawner`
550/// instance per agent.
551///
552/// Naming convention: `<WorkerIMPL><AdapterType>SpawnerFactory` (Lua
553/// worker on InProcess adapter). One half of the old
554/// `InProcSpawnerFactory`, split into Lua and RustFn variants.
555///
556/// Spec shape (choose one; `source` wins when both are present):
557///
558/// ```jsonc
559/// // (a) Registry lookup — Lua source id pre-registered with the
560/// //     factory via `register_lua` (used by the enhance flow's built-in
561/// //     workers). Requires the factory to know the id at construction
562/// //     time.
563/// { "fn_id": "patch-spawner" }
564///
565/// // (b) Inline source — a Lua chunk carried by the Blueprint itself,
566/// //     wrapped on the fly at `build` time. Combined with the loader's
567/// //     `$file` ref expansion (`"source": {"$file": "gates/foo.lua"}`)
568/// //     this lets a BP ship deterministic Lua gates without any
569/// //     pre-registration. `label` is optional and defaults to
570/// //     `"<agent_name>.lua"` for error messages.
571/// { "source": "return { value = 42, ok = true }",
572///   "label": "psim-gate.lua" }
573/// ```
574///
575/// Host bridges registered on the factory (see [`Self::with_bridge`])
576/// apply to both spec shapes.
577pub struct LuaInProcessSpawnerFactory {
578    registry: HashMap<String, WorkerFn>,
579    bridges: HashMap<String, HostBridge>,
580}
581
582/// Rust-side bridge function callable from Lua.
583///
584/// Inputs and outputs are both `serde_json::Value` (i.e. JSON). Lua
585/// invokes it as `host.<name>(arg_table)`. If the implementation needs
586/// to call async Rust, the caller does the sync-ification (typically
587/// `tokio::runtime::Handle::current().block_on(...)`).
588///
589/// Design intent: keep Lua scripts focused on flow control and `ctx`
590/// walking, while the heavy lifting (LLM calls, RFC 6902 apply,
591/// verifiers, and so on) stays on the Rust side. Going "pure Lua" —
592/// removing the bridge — is a carry.
593#[derive(Clone)]
594pub struct HostBridge(
595    Arc<dyn Fn(serde_json::Value) -> Result<serde_json::Value, String> + Send + Sync>,
596);
597
598impl HostBridge {
599    /// Wrap a Rust closure as a bridge callable from Lua.
600    pub fn new<F>(f: F) -> Self
601    where
602        F: Fn(serde_json::Value) -> Result<serde_json::Value, String> + Send + Sync + 'static,
603    {
604        Self(Arc::new(f))
605    }
606
607    /// Invoke the bridge directly — a thin trampoline over the inner
608    /// `Fn`. The production path goes through the Lua runtime, but this
609    /// stays `pub` so unit tests can exercise the primitive directly.
610    pub fn call(&self, arg: serde_json::Value) -> Result<serde_json::Value, String> {
611        (self.0)(arg)
612    }
613}
614
615/// Carrier type for Lua script sources. Paths are not required — a
616/// source string plus an identifying label is all it holds.
617///
618/// Callers bring in the source (via `include_str!` or similar) and
619/// register it with the factory through
620/// [`LuaInProcessSpawnerFactory::register_lua`].
621#[derive(Clone)]
622pub struct LuaScriptSource {
623    /// The Lua chunk source.
624    pub source: String,
625    /// Label used in error messages — typically the script's logical id
626    /// (for example `"patch_spawner.lua"`).
627    pub label: String,
628}
629
630impl LuaScriptSource {
631    /// Wrap a Lua chunk source and its error-message label.
632    pub fn new(source: impl Into<String>, label: impl Into<String>) -> Self {
633        Self {
634            source: source.into(),
635            label: label.into(),
636        }
637    }
638}
639
640impl LuaInProcessSpawnerFactory {
641    /// Start with no registered scripts and no host bridges.
642    pub fn new() -> Self {
643        Self {
644            registry: HashMap::new(),
645            bridges: HashMap::new(),
646        }
647    }
648
649    /// Register a host bridge. Subsequent `register_lua` calls snapshot
650    /// the current bridge set.
651    ///
652    /// Ordering rule: register bridges first, then call `register_lua`;
653    /// bridges added after `register_lua` will not be visible to that
654    /// script.
655    pub fn with_bridge(mut self, name: impl Into<String>, bridge: HostBridge) -> Self {
656        self.bridges.insert(name.into(), bridge);
657        self
658    }
659
660    /// Register a **Lua-eval Worker** under `fn_id`.
661    ///
662    /// Each dispatch spins up a fresh `mlua::Lua` VM, injects globals
663    /// (`_PROMPT` / `_AGENT` / `_TASK_ID` / `_ATTEMPT` / `_CTX` — the last
664    /// is `_PROMPT` parsed as JSON, or `nil` if that fails), evaluates
665    /// the script, and marshals the returned table into a `WorkerResult`.
666    ///
667    /// Marshalling rules for the return value:
668    /// - `{ value = ..., ok = bool }` → `WorkerResult.value` /
669    ///   `WorkerResult.ok` verbatim.
670    /// - Anything else → `value = <returned value>`, `ok = true`.
671    ///
672    /// Execution runs on `tokio::task::spawn_blocking` because `mlua::Lua`
673    /// is `!Send` and needs to stay away from the tokio async context.
674    /// Host bridges (the Lua-to-Rust callback path) previously registered
675    /// with [`Self::with_bridge`] are snapshotted at call time and
676    /// injected into every dispatch inside `run_lua_worker`.
677    pub fn register_lua(mut self, fn_id: impl Into<String>, source: LuaScriptSource) -> Self {
678        let source = Arc::new(source);
679        let bridges = Arc::new(self.bridges.clone());
680        let wrapped: WorkerFn = Arc::new(move |inv| {
681            let source = source.clone();
682            let bridges = bridges.clone();
683            Box::pin(run_lua_worker(source, bridges, inv))
684        });
685        self.registry.insert(fn_id.into(), wrapped);
686        self
687    }
688}
689
690/// Body of a single Lua-eval invocation (called from `register_lua`).
691async fn run_lua_worker(
692    source: Arc<LuaScriptSource>,
693    bridges: Arc<HashMap<String, HostBridge>>,
694    inv: crate::worker::adapter::WorkerInvocation,
695) -> Result<crate::worker::adapter::WorkerResult, crate::worker::adapter::WorkerError> {
696    use crate::worker::adapter::WorkerError;
697    use mlua::LuaSerdeExt;
698
699    let label = source.label.clone();
700    let outcome =
701        tokio::task::spawn_blocking(move || -> Result<(serde_json::Value, bool), String> {
702            let lua = mlua::Lua::new();
703            let g = lua.globals();
704
705            // 1. Base globals.
706            g.set("_PROMPT", inv.prompt.clone())
707                .map_err(|e| format!("set _PROMPT: {e}"))?;
708            g.set("_AGENT", inv.agent.clone())
709                .map_err(|e| format!("set _AGENT: {e}"))?;
710            g.set("_TASK_ID", inv.task_id.to_string())
711                .map_err(|e| format!("set _TASK_ID: {e}"))?;
712            g.set("_ATTEMPT", inv.attempt as i64)
713                .map_err(|e| format!("set _ATTEMPT: {e}"))?;
714
715            // 2. _CTX = JSON parse(_PROMPT); nil on parse failure (co-exists with the plain-string prompt path).
716            if let Ok(json_val) = serde_json::from_str::<serde_json::Value>(&inv.prompt) {
717                let lua_val = lua
718                    .to_value(&json_val)
719                    .map_err(|e| format!("_CTX to_value: {e}"))?;
720                g.set("_CTX", lua_val)
721                    .map_err(|e| format!("set _CTX: {e}"))?;
722            }
723
724            // 3. Inject the host bridge (Lua can call `host.<name>(arg)`).
725            if !bridges.is_empty() {
726                let host = lua
727                    .create_table()
728                    .map_err(|e| format!("create host table: {e}"))?;
729                for (name, bridge) in bridges.iter() {
730                    let bridge = bridge.clone();
731                    let bname = name.clone();
732                    let f = lua
733                        .create_function(move |lua, arg: mlua::Value| {
734                            let json_arg: serde_json::Value = lua.from_value(arg).map_err(|e| {
735                                mlua::Error::external(format!("bridge {bname} arg → json: {e}"))
736                            })?;
737                            let result_json =
738                                bridge.call(json_arg).map_err(mlua::Error::external)?;
739                            lua.to_value(&result_json).map_err(|e| {
740                                mlua::Error::external(format!("bridge {bname} ret → lua: {e}"))
741                            })
742                        })
743                        .map_err(|e| format!("create_function {name}: {e}"))?;
744                    host.set(name.as_str(), f)
745                        .map_err(|e| format!("host.{name} set: {e}"))?;
746                }
747                g.set("host", host).map_err(|e| format!("set host: {e}"))?;
748            }
749
750            // 4. eval
751            let result: mlua::Value = lua
752                .load(&source.source)
753                .set_name(&source.label)
754                .eval()
755                .map_err(|e| format!("lua eval [{}]: {e}", source.label))?;
756
757            // 5. Marshal: shape `{ value=..., ok=true }` or raw value.
758            let json_result: serde_json::Value = lua
759                .from_value(result)
760                .map_err(|e| format!("lua → json [{}]: {e}", source.label))?;
761
762            let (value, ok) = match &json_result {
763                serde_json::Value::Object(map)
764                    if map.contains_key("value") || map.contains_key("ok") =>
765                {
766                    let ok = map.get("ok").and_then(|v| v.as_bool()).unwrap_or(true);
767                    let value = map.get("value").cloned().unwrap_or(json_result.clone());
768                    (value, ok)
769                }
770                _ => (json_result, true),
771            };
772            Ok((value, ok))
773        })
774        .await
775        .map_err(|e| WorkerError::Failed(format!("spawn_blocking join [{label}]: {e}")))?
776        .map_err(WorkerError::Failed)?;
777
778    Ok(crate::worker::adapter::WorkerResult {
779        value: outcome.0,
780        ok: outcome.1,
781    })
782}
783
784impl Default for LuaInProcessSpawnerFactory {
785    fn default() -> Self {
786        Self::new()
787    }
788}
789
790impl SpawnerFactoryKind for LuaInProcessSpawnerFactory {
791    const KIND: AgentKind = AgentKind::Lua;
792    type Worker = LuaWorker;
793}
794
795impl SpawnerFactory for LuaInProcessSpawnerFactory {
796    fn build(
797        &self,
798        agent_def: &AgentDef,
799        _hint: Option<&Value>,
800    ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
801        // Inline `spec.source` (a Lua chunk carried by the BP itself) takes
802        // precedence over `spec.fn_id`. This is the path a BP author uses to
803        // ship a deterministic Lua gate without pre-registering it with the
804        // factory — the plumbing (`run_lua_worker` / `LuaScriptSource`) is
805        // the same, only the entry point differs.
806        if let Some(source) = agent_def.spec.get("source").and_then(|v| v.as_str()) {
807            let label = agent_def
808                .spec
809                .get("label")
810                .and_then(|v| v.as_str())
811                .map(str::to_string)
812                .unwrap_or_else(|| format!("{}.lua", agent_def.name));
813            let script = Arc::new(LuaScriptSource::new(source.to_string(), label));
814            let bridges = Arc::new(self.bridges.clone());
815            let wrapped: WorkerFn = Arc::new(move |inv| {
816                let source = script.clone();
817                let bridges = bridges.clone();
818                Box::pin(run_lua_worker(source, bridges, inv))
819            });
820            let mut sp: InProcSpawner<LuaWorker> = InProcSpawner::<LuaWorker>::typed();
821            sp.registry.insert(agent_def.name.to_string(), wrapped);
822            return Ok(Arc::new(sp));
823        }
824        build_inproc_from_registry::<LuaWorker>(&self.registry, agent_def, "lua")
825    }
826}
827
828/// Factory for `AgentKind::RustFn`. At `build` time it looks the `fn_id`
829/// up in its internal registry and returns an [`InProcSpawner`] with the
830/// Rust closure `WorkerFn` registered under `agent_name`.
831///
832/// Naming convention: `<WorkerIMPL><AdapterType>SpawnerFactory` (RustFn
833/// worker on InProcess adapter). Sibling to
834/// [`LuaInProcessSpawnerFactory`] — the Lua-worker half of the same
835/// split.
836///
837/// Spec shape:
838/// ```jsonc
839/// { "fn_id": "echo" }     // Rust closure id pre-registered with the factory
840/// ```
841pub struct RustFnInProcessSpawnerFactory {
842    registry: HashMap<String, WorkerFn>,
843}
844
845impl RustFnInProcessSpawnerFactory {
846    /// Start with no registered closures.
847    pub fn new() -> Self {
848        Self {
849            registry: HashMap::new(),
850        }
851    }
852
853    /// Register a Rust closure `WorkerFn` under `fn_id`, wrapping it so
854    /// it matches the `WorkerFn` signature (boxed, pinned future).
855    pub fn register_fn<F, Fut>(mut self, fn_id: impl Into<String>, f: F) -> Self
856    where
857        F: Fn(crate::worker::adapter::WorkerInvocation) -> Fut + Send + Sync + 'static,
858        Fut: std::future::Future<
859                Output = Result<
860                    crate::worker::adapter::WorkerResult,
861                    crate::worker::adapter::WorkerError,
862                >,
863            > + Send
864            + 'static,
865    {
866        let f = Arc::new(f);
867        let wrapped: WorkerFn = Arc::new(move |inv| {
868            let f = f.clone();
869            Box::pin(f(inv))
870        });
871        self.registry.insert(fn_id.into(), wrapped);
872        self
873    }
874}
875
876impl Default for RustFnInProcessSpawnerFactory {
877    fn default() -> Self {
878        Self::new()
879    }
880}
881
882impl SpawnerFactoryKind for RustFnInProcessSpawnerFactory {
883    const KIND: AgentKind = AgentKind::RustFn;
884    type Worker = RustFnWorker;
885}
886
887impl SpawnerFactory for RustFnInProcessSpawnerFactory {
888    fn build(
889        &self,
890        agent_def: &AgentDef,
891        _hint: Option<&Value>,
892    ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
893        build_inproc_from_registry::<RustFnWorker>(&self.registry, agent_def, "rust_fn")
894    }
895}
896
897/// Shared build helper used by both the Lua and the RustFn factories —
898/// look `spec.fn_id` up in the registry and return an `InProcSpawner`.
899/// The generic type parameter `W` fixes the per-kind Worker concrete
900/// type at the type level (the build-site half of the trait's
901/// associated-type binding across the four-layer cascade).
902fn build_inproc_from_registry<W>(
903    registry: &HashMap<String, WorkerFn>,
904    agent_def: &AgentDef,
905    kind_label: &str,
906) -> Result<Arc<dyn SpawnerAdapter>, CompileError>
907where
908    W: crate::worker::Worker + From<crate::worker::WorkerJoinHandler> + Send + Sync + 'static,
909{
910    let agent_name = &agent_def.name;
911    let spec = &agent_def.spec;
912    let invalid = |msg: String| CompileError::InvalidSpec {
913        name: agent_name.to_string(),
914        msg,
915    };
916    let fn_id = spec
917        .get("fn_id")
918        .and_then(|v| v.as_str())
919        .ok_or_else(|| invalid(format!("{kind_label} spec: 'fn_id' (string) required")))?;
920    let f = registry
921        .get(fn_id)
922        .cloned()
923        .ok_or_else(|| invalid(format!("fn_id '{fn_id}' not registered in factory")))?;
924    let mut sp: InProcSpawner<W> = InProcSpawner::<W>::typed();
925    // Register under `agent_name` (the flow's `Step.ref`). Both
926    // `CompiledAgentTable` and the `InProcSpawner` look the function up
927    // by name, so the same key is needed at both layers.
928    sp.registry.insert(agent_name.to_string(), f);
929    Ok(Arc::new(sp))
930}
931
932/// Concrete Worker type for the Lua kind — a handle to a Lua-eval task
933/// inside an mlua VM. Embeds a `WorkerJoinHandler`. Reserved as the home
934/// for future Lua-specific extensions (an mlua VM cancellation
935/// mechanism, Lua-side error type retention, and so on).
936pub struct LuaWorker {
937    /// The join handle / cancellation token for the underlying task.
938    pub handler: crate::worker::WorkerJoinHandler,
939}
940
941impl From<crate::worker::WorkerJoinHandler> for LuaWorker {
942    fn from(handler: crate::worker::WorkerJoinHandler) -> Self {
943        Self { handler }
944    }
945}
946
947#[async_trait::async_trait]
948impl crate::worker::Worker for LuaWorker {
949    fn id(&self) -> &crate::types::WorkerId {
950        &self.handler.worker_id
951    }
952    fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
953        self.handler.cancel.clone()
954    }
955    async fn join(self: Box<Self>) -> Result<(), crate::worker::adapter::WorkerError> {
956        self.handler.await_completion().await
957    }
958}
959
960/// Concrete Worker type for the RustFn kind — a handle to a task that
961/// directly calls a Rust closure. Embeds a `WorkerJoinHandler`. Being a
962/// pure function, there is minimal kind-specific extension surface here;
963/// the primary purpose is to nail down the type binding.
964pub struct RustFnWorker {
965    /// The join handle / cancellation token for the underlying task.
966    pub handler: crate::worker::WorkerJoinHandler,
967}
968
969impl From<crate::worker::WorkerJoinHandler> for RustFnWorker {
970    fn from(handler: crate::worker::WorkerJoinHandler) -> Self {
971        Self { handler }
972    }
973}
974
975#[async_trait::async_trait]
976impl crate::worker::Worker for RustFnWorker {
977    fn id(&self) -> &crate::types::WorkerId {
978        &self.handler.worker_id
979    }
980    fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
981        self.handler.cancel.clone()
982    }
983    async fn join(self: Box<Self>) -> Result<(), crate::worker::adapter::WorkerError> {
984        self.handler.await_completion().await
985    }
986}
987
988/// Factory for `AgentKind::Operator`. Looks up the `Arc<dyn Operator>`
989/// pre-registered under `spec.operator_ref` and wraps it in an
990/// `OperatorSpawner`. Also resolves `AgentDef.profile.worker_binding` into
991/// a `WorkerBinding` at compile time and fails loud (`CompileError::InvalidSpec`)
992/// when the resolved operator's `Operator::requires_worker_binding` is `true`
993/// and no binding was declared.
994///
995/// Spec shape:
996/// ```jsonc
997/// { "operator_ref": "main_ai" }     // Operator id pre-registered with the factory
998/// ```
999///
1000/// # Split of responsibilities with `OperatorDelegateMiddleware`
1001///
1002/// The two axes exist for different reasons:
1003///
1004/// - **This factory (`OperatorSpawnerFactory` → `OperatorSpawner`) — the
1005///   AgentSpec axis.** Bakes a separate Operator backend into each
1006///   `AgentDef`. A `kind = Operator` `AgentDef` names its backend through
1007///   `spec.operator_ref`; at `compile()` time the `Arc<dyn Operator>` is
1008///   baked into `routes[agent_name]`. Because the `agent.md` loader
1009///   (`agent_md_loader`) defaults `kind` to `Operator`, agents that flow
1010///   in through agent-profiles land here.
1011///
1012/// - **`OperatorDelegateMiddleware` — the Blueprint-global (session)
1013///   axis.** Delegates every agent to the same Operator backend. At
1014///   session-attach time you call `engine.register_operator(id, op)`
1015///   plus `attach_with_ids(.., operator_backend_id = Some(id))` to bind
1016///   it session-wide, and declare
1017///   `spawner_hints.layers = ["operator_delegate"]` to opt in. `ctx.agent`
1018///   is ignored; the operator handles every spawn in that session (a
1019///   MainAI-wide driver, a human-wide console, that sort of thing).
1020///
1021/// # Exclusivity (a double fire is structurally impossible)
1022///
1023/// When both are effective — the hint is declared, the session has an
1024/// operator backend, **and** the Blueprint has a `kind = Operator`
1025/// `AgentDef` — `OperatorDelegateMiddleware` sits at the outer end of
1026/// the stack and **completely bypasses** `inner.spawn`. The
1027/// `OperatorSpawner` is never reached, so under those conditions this
1028/// factory's routes entry is inert. This is not a double fire — the
1029/// session axis is overriding the agent axis. Consistent usage means
1030/// picking one axis per use case.
1031///
1032/// Interior mutability is provided by an `Arc<RwLock>`. Even after the
1033/// factory has been stored as `Arc<dyn SpawnerFactory>` in
1034/// `SpawnerRegistry`, a caller holding an `Arc` clone can still add
1035/// Operator backends dynamically via `register_operator(&self, id, op)`.
1036/// Typical uses: registering a `WSOperatorSession` under the session id
1037/// on WebSocket connect, binding agents that arrive via the `agent.md`
1038/// loader to arbitrary backends, and so on. `build()` performs a
1039/// `read()` lookup each time.
1040pub struct OperatorSpawnerFactory {
1041    operators: Arc<std::sync::RwLock<HashMap<String, Arc<dyn Operator>>>>,
1042}
1043
1044impl OperatorSpawnerFactory {
1045    /// Start with no registered Operator backends.
1046    pub fn new() -> Self {
1047        Self {
1048            operators: Arc::new(std::sync::RwLock::new(HashMap::new())),
1049        }
1050    }
1051
1052    /// Register an Operator backend dynamically through `&self`.
1053    /// Overwrites are allowed — later wins. Callers can still reach this
1054    /// after the factory has been stored as `Arc<dyn SpawnerFactory>` in
1055    /// `SpawnerRegistry`, as long as they hold an `Arc` clone; interior
1056    /// mutability is provided by the inner `RwLock`.
1057    pub fn register_operator(&self, id: impl Into<String>, op: Arc<dyn Operator>) -> &Self {
1058        self.operators
1059            .write()
1060            .expect("OperatorSpawnerFactory.operators RwLock poisoned")
1061            .insert(id.into(), op);
1062        self
1063    }
1064
1065    /// Dynamically unregister an id (used to clean up when a WebSocket
1066    /// disconnects, for example). A missing id is a no-op.
1067    pub fn unregister_operator(&self, id: &str) -> &Self {
1068        self.operators
1069            .write()
1070            .expect("OperatorSpawnerFactory.operators RwLock poisoned")
1071            .remove(id);
1072        self
1073    }
1074}
1075
1076impl Default for OperatorSpawnerFactory {
1077    fn default() -> Self {
1078        Self::new()
1079    }
1080}
1081
1082impl SpawnerFactoryKind for OperatorSpawnerFactory {
1083    const KIND: AgentKind = AgentKind::Operator;
1084    type Worker = crate::operator::OperatorWorker;
1085}
1086
1087impl SpawnerFactory for OperatorSpawnerFactory {
1088    fn build(
1089        &self,
1090        agent_def: &AgentDef,
1091        _hint: Option<&Value>,
1092    ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
1093        let agent_name = &agent_def.name;
1094        let spec = &agent_def.spec;
1095        // Bake AgentDef.profile.system_prompt into the OperatorSpawner at compile time.
1096        // `Some` → adopted first at spawn time; `None` → falls back to fetch_prompt (initial_directive).
1097        // Fallback path. Sibling: AgentBlockInProcessSpawnerFactory
1098        // (agent_block/runtime.rs) does the same compile-time bake by stuffing
1099        // the profile into BlockConfig.context.
1100        let system_prompt = agent_def.profile.as_ref().map(|p| p.system_prompt.clone());
1101        let invalid = |msg: String| CompileError::InvalidSpec {
1102            name: agent_name.to_string(),
1103            msg,
1104        };
1105        let op_ref = spec
1106            .get("operator_ref")
1107            .and_then(|v| v.as_str())
1108            .ok_or_else(|| invalid("operator spec: 'operator_ref' (string) required".into()))?;
1109        let operators = self
1110            .operators
1111            .read()
1112            .expect("OperatorSpawnerFactory.operators RwLock poisoned");
1113        let op = operators.get(op_ref).cloned().ok_or_else(|| {
1114            let mut names: Vec<String> = operators.keys().cloned().collect();
1115            names.sort();
1116            let names_list = if names.is_empty() {
1117                "<none>".to_string()
1118            } else {
1119                names.join(", ")
1120            };
1121            invalid(format!(
1122                "operator_ref '{op_ref}' not registered in factory. \
1123                 Registered sids: [{names_list}]. \
1124                 Hint: call mse_operator_join(roles=[...]) to mint the sid first."
1125            ))
1126        })?;
1127        drop(operators);
1128
1129        // Resolve the Blueprint-baked worker binding from
1130        // `AgentDef.profile.worker_binding` — the SoT for the
1131        // declaration↔executor binding (see `WorkerBinding` doc). Fail
1132        // loud at compile time when the operator backend requires one
1133        // and the Blueprint didn't declare it; this is a compile-time
1134        // gate, not a runtime guess.
1135        let worker_binding = agent_def
1136            .profile
1137            .as_ref()
1138            .and_then(|p| p.worker_binding.as_ref())
1139            .map(|variant| WorkerBinding {
1140                variant: variant.clone(),
1141                tools: agent_def
1142                    .profile
1143                    .as_ref()
1144                    .map(|p| p.tools.clone())
1145                    .unwrap_or_default(),
1146            });
1147        if op.requires_worker_binding() && worker_binding.is_none() {
1148            // Issue #9: the two Blueprint authoring paths (direct JSON
1149            // and `$agent_md` file ref) both land here. Old message
1150            // pointed only at the `.md` frontmatter, which was
1151            // confusing for authors on the JSON-direct path.
1152            return Err(invalid(
1153                "profile.worker_binding is required for this operator backend. \
1154                 Fix by either: \
1155                 (a) if authoring the Blueprint JSON directly, add \
1156                 `agents[N].profile.worker_binding: \"<subagent-type>\"` \
1157                 to the JSON literal; or \
1158                 (b) if using an $agent_md file ref, add \
1159                 `worker_binding: <subagent-type>` to the agent .md frontmatter."
1160                    .into(),
1161            ));
1162        }
1163        Ok(Arc::new(OperatorSpawner::new(
1164            op,
1165            system_prompt,
1166            worker_binding,
1167        )))
1168    }
1169}
1170
1171#[cfg(test)]
1172mod operator_spawner_factory_worker_binding_tests {
1173    use super::*;
1174    use crate::blueprint::AgentProfile;
1175    use crate::core::ctx::Ctx;
1176    use crate::types::CapToken;
1177    use crate::worker::adapter::{WorkerError, WorkerResult};
1178
1179    /// Minimal `Operator` stub whose `requires_worker_binding` is
1180    /// configurable — enough to exercise the compile-time fail-loud gate
1181    /// without standing up a real backend (e.g. `WSOperatorSession`,
1182    /// which lives in a downstream crate).
1183    struct StubOperator {
1184        requires_binding: bool,
1185    }
1186
1187    #[async_trait]
1188    impl Operator for StubOperator {
1189        async fn execute(
1190            &self,
1191            _ctx: &Ctx,
1192            _system: Option<String>,
1193            _prompt: Value,
1194            _worker: Option<WorkerBinding>,
1195            _worker_token: CapToken,
1196        ) -> Result<WorkerResult, WorkerError> {
1197            Ok(WorkerResult {
1198                value: Value::Null,
1199                ok: true,
1200            })
1201        }
1202
1203        fn requires_worker_binding(&self) -> bool {
1204            self.requires_binding
1205        }
1206    }
1207
1208    fn agent_def_with(profile: Option<AgentProfile>) -> AgentDef {
1209        AgentDef {
1210            name: "test-agent".to_string(),
1211            kind: AgentKind::Operator,
1212            spec: serde_json::json!({ "operator_ref": "op1" }),
1213            profile,
1214            meta: None,
1215        }
1216    }
1217
1218    #[test]
1219    fn build_fails_loud_when_binding_required_but_absent() {
1220        let factory = OperatorSpawnerFactory::new();
1221        factory.register_operator(
1222            "op1",
1223            Arc::new(StubOperator {
1224                requires_binding: true,
1225            }) as Arc<dyn Operator>,
1226        );
1227        let def = agent_def_with(Some(AgentProfile::default()));
1228        match factory.build(&def, None) {
1229            Err(CompileError::InvalidSpec { name, msg }) => {
1230                assert_eq!(name, "test-agent");
1231                assert!(
1232                    msg.contains("worker_binding is required"),
1233                    "unexpected message: {msg}"
1234                );
1235                // Issue #9: the message must be actionable for both
1236                // authoring paths — the JSON-direct hint and the
1237                // $agent_md hint both surface.
1238                assert!(
1239                    msg.contains("agents[N].profile.worker_binding"),
1240                    "message missing JSON-direct hint (issue #9): {msg}"
1241                );
1242                assert!(
1243                    msg.contains("agent .md frontmatter"),
1244                    "message missing $agent_md hint: {msg}"
1245                );
1246            }
1247            Err(other) => panic!("expected InvalidSpec, got: {other:?}"),
1248            Ok(_) => panic!("expected compile-time failure, got Ok"),
1249        }
1250    }
1251
1252    #[test]
1253    fn build_succeeds_when_binding_required_and_present() {
1254        let factory = OperatorSpawnerFactory::new();
1255        factory.register_operator(
1256            "op1",
1257            Arc::new(StubOperator {
1258                requires_binding: true,
1259            }) as Arc<dyn Operator>,
1260        );
1261        let profile = AgentProfile {
1262            worker_binding: Some("mse-worker-coder".to_string()),
1263            tools: vec!["Read".to_string(), "Edit".to_string()],
1264            ..Default::default()
1265        };
1266        let def = agent_def_with(Some(profile));
1267        assert!(
1268            factory.build(&def, None).is_ok(),
1269            "expected Ok when worker_binding is declared"
1270        );
1271    }
1272
1273    #[test]
1274    fn build_succeeds_when_binding_not_required_and_absent() {
1275        let factory = OperatorSpawnerFactory::new();
1276        factory.register_operator(
1277            "op1",
1278            Arc::new(StubOperator {
1279                requires_binding: false,
1280            }) as Arc<dyn Operator>,
1281        );
1282        let def = agent_def_with(Some(AgentProfile::default()));
1283        assert!(
1284            factory.build(&def, None).is_ok(),
1285            "backends that don't require a binding must not be gated by its absence"
1286        );
1287    }
1288}
1289
1290// ─── LuaInProcessSpawnerFactory: inline `spec.source` support ─────────────
1291//
1292// Issue `ab3d1145`: BPs served by `mse serve` couldn't declare `kind: lua`
1293// without pre-registering a `fn_id` on the factory. These tests cover the
1294// new inline path — `spec.source = "<lua chunk>"` (optionally with `label`)
1295// wraps a fresh `LuaScriptSource` at `build` time and runs it through the
1296// same `run_lua_worker` plumbing as the registry path.
1297#[cfg(test)]
1298mod lua_inline_source_tests {
1299    use super::*;
1300    use crate::types::{CapToken, Role, StepId};
1301
1302    fn agent(name: &str, spec: Value) -> AgentDef {
1303        AgentDef {
1304            name: name.to_string(),
1305            kind: AgentKind::Lua,
1306            spec,
1307            profile: None,
1308            meta: None,
1309        }
1310    }
1311
1312    fn test_invocation(prompt: &str) -> crate::worker::adapter::WorkerInvocation {
1313        crate::worker::adapter::WorkerInvocation {
1314            token: CapToken {
1315                agent_id: "a".into(),
1316                role: Role::Worker,
1317                scopes: vec!["*".into()],
1318                issued_at: 0,
1319                expire_at: u64::MAX / 2,
1320                max_uses: None,
1321                nonce: "test-nonce".into(),
1322                sig_hex: "".into(),
1323            },
1324            task_id: StepId::parse("ST-test").expect("StepId parse"),
1325            attempt: 1,
1326            agent: "g".into(),
1327            prompt: prompt.into(),
1328            sink: None,
1329            cancel_token: None,
1330        }
1331    }
1332
1333    #[test]
1334    fn build_accepts_inline_source_without_pre_registration() {
1335        let factory = LuaInProcessSpawnerFactory::new();
1336        let def = agent(
1337            "g",
1338            serde_json::json!({ "source": "return { value = 42, ok = true }" }),
1339        );
1340        assert!(
1341            factory.build(&def, None).is_ok(),
1342            "inline spec.source must build without a pre-registered fn_id"
1343        );
1344    }
1345
1346    #[test]
1347    fn build_rejects_when_neither_source_nor_fn_id_is_present() {
1348        let factory = LuaInProcessSpawnerFactory::new();
1349        let def = agent("g", serde_json::json!({}));
1350        match factory.build(&def, None) {
1351            Err(CompileError::InvalidSpec { msg, .. }) => {
1352                assert!(
1353                    msg.contains("fn_id"),
1354                    "empty spec must still surface the fn_id-required message: {msg}"
1355                );
1356            }
1357            Err(other) => panic!("expected InvalidSpec, got a different CompileError: {other}"),
1358            // `SpawnerAdapter` is not Debug, so we can't `unwrap_err()` /
1359            // pattern-print the Ok arm — describe the mismatch directly.
1360            Ok(_) => panic!("expected InvalidSpec, got Ok(SpawnerAdapter)"),
1361        }
1362    }
1363
1364    /// The inline path shares `run_lua_worker` with the registry path, so
1365    /// exercising the marshaller once through it is enough to prove the
1366    /// wrap is faithful.
1367    #[tokio::test]
1368    async fn inline_source_evaluates_and_marshals_result() {
1369        let source =
1370            LuaScriptSource::new("return { value = _PROMPT .. '!', ok = true }", "smoke.lua");
1371        let out = run_lua_worker(
1372            std::sync::Arc::new(source),
1373            std::sync::Arc::new(HashMap::new()),
1374            test_invocation("hello"),
1375        )
1376        .await
1377        .expect("lua worker ok");
1378        assert_eq!(out.value, serde_json::json!("hello!"));
1379        assert!(out.ok);
1380    }
1381
1382    #[tokio::test]
1383    async fn inline_source_can_signal_agent_level_failure() {
1384        // Deterministic gate pattern: return `ok = false` to flip the
1385        // dispatch outcome to `Blocked` (the flow.ir Try catch path).
1386        let source = LuaScriptSource::new("return { value = 'nope', ok = false }", "gate.lua");
1387        let out = run_lua_worker(
1388            std::sync::Arc::new(source),
1389            std::sync::Arc::new(HashMap::new()),
1390            test_invocation("input"),
1391        )
1392        .await
1393        .expect("lua worker ok");
1394        assert_eq!(out.value, serde_json::json!("nope"));
1395        assert!(!out.ok);
1396    }
1397}
1398
1399// ─── GH #21 Phase 2: `Blueprint.metas` / `AgentMeta.meta_ref` / static
1400// `$step_meta.ref` compile-time validation ─────────────────────────────────
1401#[cfg(test)]
1402mod meta_ref_validation_tests {
1403    use super::*;
1404    use crate::blueprint::{AgentMeta, MetaDef};
1405    use crate::worker::adapter::WorkerResult;
1406
1407    fn registry_with_echo() -> SpawnerRegistry {
1408        let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", |inv| async move {
1409            Ok(WorkerResult {
1410                value: Value::String(inv.prompt),
1411                ok: true,
1412            })
1413        });
1414        let mut reg = SpawnerRegistry::new();
1415        reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1416        reg
1417    }
1418
1419    fn rustfn_agent(name: &str) -> AgentDef {
1420        AgentDef {
1421            name: name.to_string(),
1422            kind: AgentKind::RustFn,
1423            spec: serde_json::json!({ "fn_id": "echo" }),
1424            profile: None,
1425            meta: None,
1426        }
1427    }
1428
1429    fn simple_flow(agent_ref: &str, in_: Expr) -> FlowNode {
1430        FlowNode::Step {
1431            ref_: agent_ref.to_string(),
1432            in_,
1433            out: Expr::Path {
1434                at: "$.output".into(),
1435            },
1436        }
1437    }
1438
1439    fn minimal_bp(agents: Vec<AgentDef>, metas: Vec<MetaDef>, flow: FlowNode) -> Blueprint {
1440        Blueprint {
1441            schema_version: crate::blueprint::current_schema_version(),
1442            id: "meta-ref-ut".into(),
1443            flow,
1444            agents,
1445            operators: vec![],
1446            metas,
1447            hints: Default::default(),
1448            strategy: Default::default(),
1449            metadata: BlueprintMetadata::default(),
1450            spawner_hints: Default::default(),
1451            default_agent_kind: AgentKind::Operator,
1452            default_operator_kind: None,
1453            default_init_ctx: None,
1454            default_agent_ctx: None,
1455            default_context_policy: None,
1456        }
1457    }
1458
1459    #[test]
1460    fn valid_meta_ref_compiles() {
1461        let mut agent = rustfn_agent("worker");
1462        agent.meta = Some(AgentMeta {
1463            meta_ref: Some("shared".to_string()),
1464            ..Default::default()
1465        });
1466        let bp = minimal_bp(
1467            vec![agent],
1468            vec![MetaDef {
1469                name: "shared".into(),
1470                ctx: serde_json::json!({ "k": "v" }),
1471            }],
1472            simple_flow(
1473                "worker",
1474                Expr::Path {
1475                    at: "$.input".into(),
1476                },
1477            ),
1478        );
1479        let compiler = Compiler::new(registry_with_echo());
1480        assert!(
1481            compiler.compile(&bp).is_ok(),
1482            "a resolvable AgentMeta.meta_ref must compile"
1483        );
1484    }
1485
1486    #[test]
1487    fn unknown_agent_meta_ref_is_unresolved_meta_ref() {
1488        let mut agent = rustfn_agent("worker");
1489        agent.meta = Some(AgentMeta {
1490            meta_ref: Some("missing".to_string()),
1491            ..Default::default()
1492        });
1493        let bp = minimal_bp(
1494            vec![agent],
1495            vec![],
1496            simple_flow(
1497                "worker",
1498                Expr::Path {
1499                    at: "$.input".into(),
1500                },
1501            ),
1502        );
1503        let compiler = Compiler::new(registry_with_echo());
1504        match compiler.compile(&bp) {
1505            Err(CompileError::UnresolvedMetaRef {
1506                where_,
1507                meta_ref,
1508                defined,
1509            }) => {
1510                assert!(
1511                    where_.contains("worker"),
1512                    "where_ must name the agent: {where_}"
1513                );
1514                assert_eq!(meta_ref, "missing");
1515                assert!(defined.is_empty());
1516            }
1517            Err(other) => {
1518                panic!("expected UnresolvedMetaRef, got a different CompileError: {other}")
1519            }
1520            Ok(_) => panic!("expected compile-time failure, got Ok"),
1521        }
1522    }
1523
1524    #[test]
1525    fn unknown_static_step_meta_ref_in_lit_is_unresolved_meta_ref() {
1526        let agent = rustfn_agent("worker");
1527        let in_ = Expr::Lit {
1528            value: serde_json::json!({ "$step_meta": { "ref": "missing" }, "$in": "go" }),
1529        };
1530        let bp = minimal_bp(vec![agent], vec![], simple_flow("worker", in_));
1531        let compiler = Compiler::new(registry_with_echo());
1532        match compiler.compile(&bp) {
1533            Err(CompileError::UnresolvedMetaRef {
1534                where_, meta_ref, ..
1535            }) => {
1536                assert!(
1537                    where_.contains("worker"),
1538                    "where_ must name the offending step: {where_}"
1539                );
1540                assert_eq!(meta_ref, "missing");
1541            }
1542            Err(other) => {
1543                panic!("expected UnresolvedMetaRef, got a different CompileError: {other}")
1544            }
1545            Ok(_) => panic!("expected compile-time failure, got Ok"),
1546        }
1547    }
1548
1549    #[test]
1550    fn path_op_input_with_no_static_envelope_compiles_fine() {
1551        let agent = rustfn_agent("worker");
1552        let bp = minimal_bp(
1553            vec![agent],
1554            vec![],
1555            simple_flow(
1556                "worker",
1557                Expr::Path {
1558                    at: "$.input".into(),
1559                },
1560            ),
1561        );
1562        let compiler = Compiler::new(registry_with_echo());
1563        assert!(
1564            compiler.compile(&bp).is_ok(),
1565            "a non-Lit Step.in must not trigger the best-effort static $step_meta check"
1566        );
1567    }
1568}