Skip to main content

mlua_swarm/
middleware.rs

1//! Middleware overlay — cross-cutting concerns (Audit / MainAI / Senior /
2//! LongHold).
3//!
4//! Ships four `SpawnerLayer` implementations plus the `SpawnerStack` builder.
5//! Some layers key off `Ctx.operator.kind` and only fire for
6//! `MainAi` / `Composite` sessions; others (`Audit` / `LongHold`) apply
7//! uniformly across every kind.
8//!
9//! # Extension discipline — this layer is THE extension point (canonical)
10//!
11//! Background: an earlier iteration grew a verdict-specialised machinery
12//! (`judgment.rs` canonical type + 3-form parser + `state.agent_verdicts`
13//! map + dedicated accessor) that re-interpreted agent output *inside the
14//! engine core* and banned string-literal conds in favour of a Blueprint
15//! compile-layer translation. That whole complex was dismantled: the value
16//! it added over plain data was zero, while it created an IN-side dialect
17//! that every consumer had to learn. The design conclusion is a
18//! three-principle layering:
19//!
20//! 1. **IN is immutable, canonical form is JSON.** `Blueprint` /
21//!    `mlua_flow_ir::Node` are plain serde data. No compile pass, no schema
22//!    field that the engine expands, no Rust helper that builds `Expr`s.
23//!    Flow control is written literally in Flow.ir:
24//!    `Eq(Path("$.<step>.verdict"), Lit("blocked"))` — domain verdicts are
25//!    plain strings inside step output, consumed by plain conds.
26//! 2. **Generation (authoring sugar) lives OUT**, on the consumer side
27//!    (e.g. a vendored pure-Lua builder that prints Blueprint JSON). It
28//!    never leaks into engine / schema crates, whatever language it is
29//!    written in — the ban is on the *placement*, not the language.
30//! 3. **Runtime extension lives HERE, as a `SpawnerLayer`.** A middleware
31//!    (or any future extension mechanism) may interpret the *results* of a
32//!    Flow.ir run — `Ctx`, the `output_tail`, `Final { ok }` — in its own
33//!    way and transform them. What it must NOT do:
34//!    - introduce a new dialect on the IN side (schema fields / node
35//!      rewriting / cond translation) — extensions read and transform, the
36//!      wire format stays plain Flow.ir + JSON;
37//!    - hide its effect: overrides are *appended* to the output tail
38//!      (e.g. `SeniorEscalationMiddleware` pushes an override `Final`
39//!      rather than mutating the recorded one), so the trace stays
40//!      replayable and the flow stays observable;
41//!    - accumulate private engine state keyed by its own semantics (the
42//!      `agent_verdicts` anti-pattern) — state lives in ctx / output store
43//!      as plain data.
44//!
45//! `AgentResolver`, `ProjectNameAliasMiddleware`, `SinkMiddleware`,
46//! `InputInjectMiddleware`, `LuaMiddleware`, `SeniorEscalationMiddleware`,
47//! `TaskInputMiddleware` all follow this shape: edit `ctx` / wrap the
48//! worker, call the inner spawner, append observable output. Note
49//! `LuaMiddleware`'s scripts are host-constructed — embedding Lua source
50//! in a Blueprint is the IN-side dialect this discipline forbids, and
51//! would require its own guard design if ever revisited).
52
53pub mod agent_context;
54pub mod input_inject;
55pub mod lua_layer;
56pub mod project_name_alias;
57pub mod resolver;
58pub mod sink;
59pub mod task_input;
60pub mod worker_binding;
61
62use crate::core::ctx::{Ctx, OperatorKind};
63use crate::core::engine::Engine;
64use crate::core::state::Event;
65use crate::types::{CapToken, StepId};
66use crate::worker::adapter::{SpawnError, SpawnerAdapter};
67use crate::worker::output::{ContentRef, OutputEvent};
68use crate::worker::{wrap_join, MiddlewareWorker, Worker, WorkerJoinHandler};
69use async_trait::async_trait;
70use serde_json::Value;
71use std::sync::Arc;
72use std::time::{Duration, Instant};
73use tokio::sync::broadcast;
74
75/// Pull the terminal `Final` event's `(value, ok)` out of the tail (works
76/// for both `Inline` and `FileRef` content).
77async fn pull_final_value_ok(
78    engine: &Engine,
79    task_id: &StepId,
80    attempt: u32,
81) -> Option<(Value, bool)> {
82    let tail = engine.output_tail(task_id, attempt).await;
83    tail.iter().rev().find_map(|ev| match ev {
84        OutputEvent::Final {
85            content: ContentRef::Inline { value },
86            ok,
87        } => Some((value.clone(), *ok)),
88        OutputEvent::Final {
89            content: ContentRef::FileRef { path, .. },
90            ok,
91        } => Some((serde_json::json!({"file_ref": path.to_string_lossy()}), *ok)),
92        _ => None,
93    })
94}
95
96/// Layer trait — one middleware stage wrapping a `SpawnerAdapter`.
97pub trait SpawnerLayer: Send + Sync + 'static {
98    /// Wraps `inner` in this layer's behaviour, returning a new
99    /// `SpawnerAdapter` that delegates to `inner` (directly or via
100    /// `wrap_join`) while adding this layer's cross-cutting effect.
101    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter>;
102}
103
104/// Stack builder that layers `SpawnerLayer`s on top of a base adapter.
105///
106/// Each `.layer(...)` call wraps a new **outer** stage — same ergonomics as
107/// `tower::ServiceBuilder`.
108pub struct SpawnerStack {
109    inner: Arc<dyn SpawnerAdapter>,
110}
111
112impl SpawnerStack {
113    /// Starts a stack with `base` as the innermost adapter.
114    pub fn new(base: Arc<dyn SpawnerAdapter>) -> Self {
115        Self { inner: base }
116    }
117
118    /// Wraps the current stack with a statically-typed `SpawnerLayer`,
119    /// becoming the new outermost stage.
120    pub fn layer<L: SpawnerLayer>(mut self, layer: L) -> Self {
121        self.inner = layer.wrap(self.inner);
122        self
123    }
124
125    /// Dynamically-typed variant taking `Arc<dyn SpawnerLayer>`. Used via
126    /// the `LayerRegistry` resolution path (where a factory returns
127    /// `Arc<dyn ...>`).
128    pub fn layer_dyn(mut self, layer: Arc<dyn SpawnerLayer>) -> Self {
129        self.inner = layer.wrap(self.inner);
130        self
131    }
132
133    /// Finishes the stack, returning the fully-wrapped adapter.
134    pub fn build(self) -> Arc<dyn SpawnerAdapter> {
135        self.inner
136    }
137}
138
139// ─── SpawnerLayerFactory + LayerRegistry ─────────────────────────────────
140//
141// # Design rationale
142//
143// Wiring is assembled per-launch through `TaskLaunchService.launch`:
144//
145//   Compiler.compile(bp) ─┬─→ compiled.router (CompiledAgentTable: agent name → SpawnerAdapter dispatch)
146//                         │
147//                         │   service::linker::link(router, bp.spawner_hints.layers, &engine)
148//                         │     internal:
149//                         │       SpawnerStack::new(router)
150//                         │         .layer_dyn(base_factory_n(engine))   ← every LayerRegistry.base entry
151//                         │         .layer_dyn(hint_factory(engine))     ← resolves each bp.spawner_hints.layers key
152//                         │         .build()
153//                         ▼
154//                   EngineDispatcher::with_spawner(engine, op_token, stacked)
155//                         ▼
156//                   engine.dispatch_attempt_with(op_token, task_id, &stacked)
157//
158// # base vs hint — when to use each
159//
160// - **base layer**: wrapped around every Blueprint. Example: AuditMiddleware
161//   (a mandatory EventLog audit). The caller registers with
162//   `LayerRegistry::with_base(|e| Arc::new(AuditMiddleware::new(e.event_tx())))`.
163//
164// - **hint layer**: wrapped **only when the Blueprint declares the key** in
165//   `spawner_hints.layers`. Examples: MainAIMiddleware /
166//   SeniorEscalationMiddleware / OperatorDelegateMiddleware. The Blueprint
167//   only declares a capability key (e.g. `"main_ai"`) without knowing the
168//   implementation; the engine-side LayerRegistry resolves key → factory,
169//   keeping the pure Flow layer separate from implementation details.
170//
171// # Factory pattern (handles layers that need Engine context)
172//
173// We do not hold `Arc<dyn SpawnerLayer>` directly because some layers
174// depend on the engine instance — for example AuditMiddleware needs
175// `engine.event_tx()` and can only be built after the engine exists. A
176// factory closure defers construction: the Layer instance is created only
177// when the engine is handed in.
178
179/// Factory closure for a `SpawnerLayer`. The caller registers these at
180/// startup, and they are called with the engine context at bind time.
181/// Stateless layers can use `|_engine| Arc::new(MyLayer)`; layers that need
182/// something like `event_tx` should do `|engine| Arc::new(MyLayer::new(engine.event_tx()))`.
183pub type LayerFactory =
184    Arc<dyn Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static>;
185
186/// Registry of `LayerFactory`s, split into `base` (always applied) and
187/// `hints` (applied only when a Blueprint declares the matching key in
188/// `spawner_hints.layers`). See the module-level `# Factory pattern`
189/// notes above for why factories rather than pre-built layers.
190#[derive(Default, Clone)]
191pub struct LayerRegistry {
192    base: Vec<LayerFactory>,
193    hints: std::collections::HashMap<String, LayerFactory>,
194}
195
196impl LayerRegistry {
197    /// Empty registry (no base layers, no hint layers).
198    pub fn new() -> Self {
199        Self::default()
200    }
201
202    /// Register a base layer factory that is applied on every Blueprint bind
203    /// (for layers that must fire for every task — e.g. `AuditMiddleware`).
204    pub fn with_base<F>(mut self, factory: F) -> Self
205    where
206        F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
207    {
208        self.base.push(Arc::new(factory));
209        self
210    }
211
212    /// Register a layer factory addressable by hint key. If
213    /// `Blueprint.spawner_hints.layers` lists the same key, it is wrapped at
214    /// bind time; otherwise it is a no-op.
215    pub fn with_hint<F>(mut self, key: impl Into<String>, factory: F) -> Self
216    where
217        F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
218    {
219        self.hints.insert(key.into(), Arc::new(factory));
220        self
221    }
222
223    /// All registered base-layer factories, in registration order.
224    pub fn base_factories(&self) -> &[LayerFactory] {
225        &self.base
226    }
227
228    /// Looks up the hint-layer factory registered under `key`, if any.
229    pub fn lookup_hint(&self, key: &str) -> Option<&LayerFactory> {
230        self.hints.get(key)
231    }
232}
233
234// ─── AuditMiddleware (pushes into the EventLog broadcast path) ────────────
235
236/// Mandatory base layer that emits `Event::TaskAttemptStarted` on every
237/// spawn, before delegating. This is the audit trail's entry point into
238/// the EventLog broadcast channel.
239pub struct AuditMiddleware {
240    /// Broadcast sender the EventLog subscribes to.
241    pub event_tx: broadcast::Sender<Event>,
242}
243
244impl AuditMiddleware {
245    /// Wraps a broadcast sender to notify on every spawn.
246    pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
247        Self { event_tx }
248    }
249}
250
251impl SpawnerLayer for AuditMiddleware {
252    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
253        Arc::new(AuditWrapped {
254            inner,
255            event_tx: self.event_tx.clone(),
256        })
257    }
258}
259
260struct AuditWrapped {
261    inner: Arc<dyn SpawnerAdapter>,
262    event_tx: broadcast::Sender<Event>,
263}
264
265#[async_trait]
266impl SpawnerAdapter for AuditWrapped {
267    async fn spawn(
268        &self,
269        engine: &Engine,
270        ctx: &Ctx,
271        task_id: StepId,
272        attempt: u32,
273        token: CapToken,
274    ) -> Result<Box<dyn Worker>, SpawnError> {
275        let _ = self.event_tx.send(Event::TaskAttemptStarted {
276            task_id: task_id.clone(),
277            attempt,
278        });
279        self.inner.spawn(engine, ctx, task_id, attempt, token).await
280    }
281}
282
283// ─── MainAIMiddleware (fires SpawnHook before/after for MainAI/Composite) ─
284
285/// Hint layer that fires `ctx.operator.spawn_hook.before`/`after` around
286/// a spawn, but only for `MainAi` / `Composite` sessions. No-op for
287/// other kinds (still delegates, just skips the hook calls).
288pub struct MainAIMiddleware;
289
290impl MainAIMiddleware {
291    /// Stateless constructor.
292    pub fn new() -> Self {
293        Self
294    }
295}
296
297impl Default for MainAIMiddleware {
298    fn default() -> Self {
299        Self::new()
300    }
301}
302
303impl SpawnerLayer for MainAIMiddleware {
304    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
305        Arc::new(MainAIWrapped { inner })
306    }
307}
308
309struct MainAIWrapped {
310    inner: Arc<dyn SpawnerAdapter>,
311}
312
313#[async_trait]
314impl SpawnerAdapter for MainAIWrapped {
315    async fn spawn(
316        &self,
317        engine: &Engine,
318        ctx: &Ctx,
319        task_id: StepId,
320        attempt: u32,
321        token: CapToken,
322    ) -> Result<Box<dyn Worker>, SpawnError> {
323        let mainai = matches!(
324            ctx.operator.kind,
325            OperatorKind::MainAi | OperatorKind::Composite
326        );
327        if mainai {
328            if let Some(hook) = &ctx.operator.spawn_hook {
329                hook.before(ctx)
330                    .await
331                    .map_err(SpawnError::RejectedByMiddleware)?;
332            }
333        }
334
335        let handle = self
336            .inner
337            .spawn(engine, ctx, task_id.clone(), attempt, token)
338            .await?;
339
340        if !mainai {
341            return Ok(handle);
342        }
343        let Some(hook) = ctx.operator.spawn_hook.clone() else {
344            return Ok(handle);
345        };
346
347        // Wrap the completion signal and call hook.after on finish.
348        // Pull the last Final from engine.output_tail as the value.
349        let ctx_clone = ctx.clone();
350        let engine_clone = engine.clone();
351        let task_id_clone = task_id.clone();
352        Ok(wrap_join(handle, move |signal| {
353            let hook = hook.clone();
354            let ctx_clone = ctx_clone.clone();
355            let engine_clone = engine_clone.clone();
356            let task_id_clone = task_id_clone.clone();
357            async move {
358                let v = match &signal {
359                    Ok(()) => pull_final_value_ok(&engine_clone, &task_id_clone, attempt)
360                        .await
361                        .map(|(v, _)| v)
362                        .unwrap_or(Value::Null),
363                    Err(e) => Value::String(e.to_string()),
364                };
365                let _ = hook.after(&ctx_clone, &v).await;
366                signal
367            }
368        }))
369    }
370}
371
372// ─── SeniorEscalationMiddleware ───────────────────────────────────────────
373//
374// When a spawn's completion is `ok=false` and `ctx.operator.senior_bridge` is
375// Some, this auxiliary layer calls `SeniorBridge.ask`, merges the answer into
376// `WorkerResult.value` under `"senior_answer"`, and upgrades the result to
377// `ok=true`. Retry / re-dispatch is the engine (operator) side's job; this
378// layer only injects fresh material for that decision.
379
380/// Hint layer: on `ok=false` completion with `ctx.operator.senior_bridge`
381/// set, asks the bridge for guidance and pushes an override `Final`
382/// (`ok=true`) carrying `senior_answer`. See the module comment above
383/// this type for the full contract.
384pub struct SeniorEscalationMiddleware;
385
386impl SeniorEscalationMiddleware {
387    /// Stateless constructor.
388    pub fn new() -> Self {
389        Self
390    }
391}
392
393impl Default for SeniorEscalationMiddleware {
394    fn default() -> Self {
395        Self::new()
396    }
397}
398
399impl SpawnerLayer for SeniorEscalationMiddleware {
400    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
401        Arc::new(SeniorWrapped { inner })
402    }
403}
404
405struct SeniorWrapped {
406    inner: Arc<dyn SpawnerAdapter>,
407}
408
409#[async_trait]
410impl SpawnerAdapter for SeniorWrapped {
411    async fn spawn(
412        &self,
413        engine: &Engine,
414        ctx: &Ctx,
415        task_id: StepId,
416        attempt: u32,
417        token: CapToken,
418    ) -> Result<Box<dyn Worker>, SpawnError> {
419        let bridge = ctx.operator.senior_bridge.clone();
420        let task_id_for_hook = task_id.clone();
421        let engine_clone = engine.clone();
422        let token_clone = token.clone();
423        let handle = self
424            .inner
425            .spawn(engine, ctx, task_id, attempt, token)
426            .await?;
427        let Some(bridge) = bridge else {
428            return Ok(handle);
429        };
430        Ok(wrap_join(handle, move |signal| {
431            let bridge = bridge.clone();
432            let task_id = task_id_for_hook.clone();
433            let engine = engine_clone.clone();
434            let token = token_clone.clone();
435            async move {
436                signal?;
437                // Read the existing Final.
438                let last = pull_final_value_ok(&engine, &task_id, attempt).await;
439                if let Some((value, false)) = last {
440                    // ok=false: escalate to senior and push an override Final.
441                    let question = serde_json::json!({
442                        "reason": "worker reported ok=false",
443                        "value": value.clone(),
444                    });
445                    if let Ok(answer) = bridge.ask(&task_id, question).await {
446                        let override_val = serde_json::json!({
447                            "original": value,
448                            "senior_answer": answer,
449                        });
450                        let _ = engine
451                            .submit_output(
452                                &token,
453                                &task_id,
454                                attempt,
455                                OutputEvent::Final {
456                                    content: ContentRef::Inline {
457                                        value: override_val,
458                                    },
459                                    ok: true,
460                                },
461                            )
462                            .await;
463                    }
464                }
465                Ok(())
466            }
467        }))
468    }
469}
470
471// ─── OperatorDelegateMiddleware (delegates the whole spawn to an external Operator when one is attached) ──
472
473/// When `ctx.operator.operator.is_some()` (the session has an Operator
474/// backend), **bypass** `inner.spawn`, call `operator.execute(ctx, prompt)`,
475/// and box the result up as a `WorkerHandle`. In other words: the path that
476/// hands "this spawn" to whatever external Operator backend the engine has
477/// registered.
478///
479/// # Independent of `OperatorKind` (Operator is a generic abstraction)
480///
481/// An earlier implementation gated on `kind == MainAi | Composite`, which
482/// tied the `Operator` abstraction to an "AI driver" assumption — a design
483/// weakness. The `Operator` trait is a generic **external processing backend**
484/// (LLM, human, external resource, side-effectful operation — anything), and
485/// is orthogonal to the kind axis.
486///
487/// The current implementation decides solely on `operator.is_some()`:
488/// - Automate session + operator backend registered → delegate
489///   (pure external-execution delegation).
490/// - MainAi session + operator backend registered → delegate.
491/// - Any kind + `operator` `None` → pass through (normal `inner.spawn`).
492///
493/// `kind` still matters as a firing condition for `SpawnHook`s over in
494/// `MainAIMiddleware`, but this middleware ignores it.
495///
496/// # Split of responsibilities with `OperatorSpawner`
497///
498/// The two axes exist for different reasons:
499///
500/// - **This middleware — the Blueprint-global (session) axis.** Delegate every
501///   agent to the same Operator backend. The `operator_backend_id` is set
502///   at session-attach time; `ctx.agent` is ignored and every spawn in that
503///   session is routed through the operator (e.g. a MainAI-wide driver, or a
504///   human-wide console). The Blueprint doesn't have to talk about `kind` —
505///   it just declares the capability hint `"operator_delegate"` (keeping the
506///   Blueprint clean).
507///
508/// - **`OperatorSpawner` — the AgentSpec axis.** Each `AgentDef` bakes its
509///   own Operator backend. `kind = Operator` `AgentDef`s pick a backend via
510///   `spec.operator_ref`; the compiler bakes an `Arc<dyn Operator>` into
511///   `routes[agent_name]`. Agents loaded via the `agent.md` loader come in
512///   through this path (their default is `kind = Operator`).
513///
514/// # Exclusivity
515///
516/// When both are effective — this middleware's hint is declared, the session
517/// has an operator backend, **and** the Blueprint has a `kind = Operator`
518/// `AgentDef` — this middleware sits at the outer end of the stack and
519/// **completely bypasses** `inner.spawn`. The `OperatorSpawner` is never
520/// reached, so a double fire cannot occur by construction; the AgentSpec
521/// axis is inert. Consistent use means picking one axis per use case.
522pub struct OperatorDelegateMiddleware;
523
524impl OperatorDelegateMiddleware {
525    /// Stateless constructor.
526    pub fn new() -> Self {
527        Self
528    }
529}
530
531impl Default for OperatorDelegateMiddleware {
532    fn default() -> Self {
533        Self::new()
534    }
535}
536
537impl SpawnerLayer for OperatorDelegateMiddleware {
538    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
539        Arc::new(OperatorDelegateWrapped { inner })
540    }
541}
542
543struct OperatorDelegateWrapped {
544    inner: Arc<dyn SpawnerAdapter>,
545}
546
547#[async_trait]
548impl SpawnerAdapter for OperatorDelegateWrapped {
549    async fn spawn(
550        &self,
551        engine: &Engine,
552        ctx: &Ctx,
553        task_id: StepId,
554        attempt: u32,
555        token: CapToken,
556    ) -> Result<Box<dyn Worker>, SpawnError> {
557        // Kind-independent: we decide purely on whether an operator backend is
558        // registered on the session. `kind` matters for SpawnHook-style layers
559        // (MainAIMiddleware); this middleware does not consult it.
560        let Some(operator) = ctx.operator.operator.clone() else {
561            return self.inner.spawn(engine, ctx, task_id, attempt, token).await;
562        };
563
564        // Delegate: same shape as OperatorSpawner — fetch_prompt + operator.execute + Final emit.
565        let prompt = engine
566            .fetch_prompt(&token, &task_id)
567            .await
568            .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
569
570        // Resolve the Blueprint-baked worker binding injected into
571        // `ctx.meta.runtime` by `WorkerBindingMiddleware` (launch-time layer,
572        // built from `AgentDef.profile.worker_binding`). Absent key = agent
573        // declared no binding → hand `None` and let binding-requiring
574        // backends fail loud (`requires_worker_binding`). A present-but-
575        // malformed value is a wiring bug, not a degrade case — fail here.
576        let worker: Option<crate::operator::WorkerBinding> = match ctx
577            .meta
578            .runtime
579            .get(crate::middleware::worker_binding::WORKER_BINDING_KEY)
580        {
581            Some(v) => Some(serde_json::from_value(v.clone()).map_err(|e| {
582                SpawnError::Internal(format!(
583                    "ctx.meta.runtime['{}'] for agent '{}' is malformed: {e}",
584                    crate::middleware::worker_binding::WORKER_BINDING_KEY,
585                    ctx.agent
586                ))
587            })?),
588            None => None,
589        };
590
591        let engine_clone = engine.clone();
592        let token_clone = token.clone();
593        let token_for_op = token.clone();
594        let task_id_clone = task_id.clone();
595        let ctx_clone = ctx.clone();
596        let (tx, rx) = tokio::sync::oneshot::channel();
597        let cancel = tokio_util::sync::CancellationToken::new();
598        let cancel_inner = cancel.clone();
599        let worker_id = crate::types::WorkerId::new();
600        // issue #11: WorkerId was minted but never observable anywhere;
601        // surface it in the trace log, tied to the step it serves.
602        tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (delegate axis)");
603
604        tokio::spawn(async move {
605            let result: Result<
606                crate::worker::adapter::WorkerResult,
607                crate::worker::adapter::WorkerError,
608            > = tokio::select! {
609                // OperatorDelegateMiddleware = session-global Operator delegation.
610                // Baking per-AgentDef profile.system_prompt is OperatorSpawner's
611                // job; this path has no per-agent spawner, so system stays None.
612                // The worker binding, however, IS resolved on this axis now:
613                // `WorkerBindingMiddleware` (launch-time layer) injects the
614                // Blueprint-baked binding into ctx.meta.runtime and we forward
615                // it here — the delegate axis is a first-class variant-dispatch
616                // path, not a binding-less fallback (issue 45db42a7).
617                // We hand the capability token (Role::Worker, 1800s TTL —
618                // minted by `Engine::dispatch_attempt_with`) to the
619                // operator as `worker_token` — thin-spawn operators (e.g. a
620                // WebSocket-backed operator session) forward it to the SubAgent
621                // via encode(), while Operator impls that call the LLM directly
622                // may ignore it.
623                r = operator.execute(&ctx_clone, None, prompt, worker, token_for_op) => r,
624                _ = cancel_inner.cancelled() => Err(crate::worker::adapter::WorkerError::Cancelled),
625            };
626            if let Ok(wr) = &result {
627                // If the SubAgent has already pushed a Final through
628                // /v1/worker/result or /v1/worker/submit POST, skip a second
629                // emit here — the POST value is the canonical one (protocol
630                // design intent). Operator impls that never POST (e.g. tests
631                // and inline Operators) still get the fallback emit.
632                let tail = engine_clone.output_tail(&task_id_clone, attempt).await;
633                let has_final = tail
634                    .iter()
635                    .any(|ev| matches!(ev, crate::worker::output::OutputEvent::Final { .. }));
636                if !has_final {
637                    let ev = crate::worker::output::OutputEvent::Final {
638                        content: crate::worker::output::ContentRef::Inline {
639                            value: wr.value.clone(),
640                        },
641                        ok: wr.ok,
642                    };
643                    let _ = engine_clone
644                        .submit_output(&token_clone, &task_id_clone, attempt, ev)
645                        .await;
646                }
647            }
648            let signal: Result<(), crate::worker::adapter::WorkerError> = result.map(|_| ());
649            let _ = tx.send(signal);
650        });
651
652        Ok(Box::new(MiddlewareWorker {
653            handler: WorkerJoinHandler {
654                worker_id,
655                cancel,
656                completion: rx,
657            },
658        }))
659    }
660}
661
662// ─── LongHoldMiddleware (warns on the EventLog if completion time exceeds default_hold) ─
663
664/// Base layer that emits `Event::TaskAttemptCompleted` with a
665/// `long_hold_warn` marker when a spawn's completion takes longer than
666/// `default_hold`. Purely observational — it never alters the signal or
667/// blocks completion.
668pub struct LongHoldMiddleware {
669    /// Threshold above which a completion is flagged as long-held.
670    pub default_hold: Duration,
671    /// Broadcast sender the EventLog subscribes to.
672    pub event_tx: broadcast::Sender<Event>,
673}
674
675impl LongHoldMiddleware {
676    /// Sets the hold threshold and the event sender to warn through.
677    pub fn new(default_hold: Duration, event_tx: broadcast::Sender<Event>) -> Self {
678        Self {
679            default_hold,
680            event_tx,
681        }
682    }
683}
684
685impl SpawnerLayer for LongHoldMiddleware {
686    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
687        Arc::new(LongHoldWrapped {
688            inner,
689            default_hold: self.default_hold,
690            event_tx: self.event_tx.clone(),
691        })
692    }
693}
694
695struct LongHoldWrapped {
696    inner: Arc<dyn SpawnerAdapter>,
697    default_hold: Duration,
698    event_tx: broadcast::Sender<Event>,
699}
700
701#[async_trait]
702impl SpawnerAdapter for LongHoldWrapped {
703    async fn spawn(
704        &self,
705        engine: &Engine,
706        ctx: &Ctx,
707        task_id: StepId,
708        attempt: u32,
709        token: CapToken,
710    ) -> Result<Box<dyn Worker>, SpawnError> {
711        let handle = self
712            .inner
713            .spawn(engine, ctx, task_id.clone(), attempt, token)
714            .await?;
715        let started = Instant::now();
716        let default_hold = self.default_hold;
717        let event_tx = self.event_tx.clone();
718        let task_id_inner = task_id.clone();
719        Ok(wrap_join(handle, move |signal| {
720            let elapsed = started.elapsed();
721            let default_hold = default_hold;
722            let event_tx = event_tx.clone();
723            let task_id_inner = task_id_inner.clone();
724            async move {
725                if elapsed > default_hold {
726                    let _ = event_tx.send(Event::TaskAttemptCompleted {
727                        task_id: task_id_inner,
728                        attempt,
729                        result: serde_json::json!({
730                            "long_hold_warn": true,
731                            "elapsed_ms": elapsed.as_millis() as u64,
732                            "default_hold_ms": default_hold.as_millis() as u64,
733                        }),
734                    });
735                }
736                signal
737            }
738        }))
739    }
740}
741
742// Boundary regression spec for the delegate-axis worker-binding handoff
743// (issue 45db42a7): OperatorDelegateMiddleware must forward the binding
744// injected into ctx.meta.runtime by WorkerBindingMiddleware — both the
745// hit path (Some(worker) reaches Operator::execute) and the absent path
746// (None reaches it), plus fail-loud on a malformed value.
747#[cfg(test)]
748mod operator_delegate_worker_binding_tests {
749    use super::*;
750    use crate::core::config::EngineCfg;
751    use crate::core::state::TaskSpec;
752    use crate::operator::WorkerBinding;
753    use crate::types::Role;
754    use crate::worker::adapter::{WorkerError, WorkerResult};
755    use std::sync::Mutex;
756
757    /// Operator stub recording the `worker` argument it was executed with.
758    struct RecordingOperator {
759        seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
760    }
761
762    #[async_trait]
763    impl crate::operator::Operator for RecordingOperator {
764        async fn execute(
765            &self,
766            _ctx: &Ctx,
767            _system: Option<String>,
768            _prompt: Value,
769            worker: Option<WorkerBinding>,
770            _worker_token: CapToken,
771        ) -> Result<WorkerResult, WorkerError> {
772            *self.seen.lock().unwrap() = Some(worker);
773            Ok(WorkerResult {
774                value: Value::Null,
775                ok: true,
776            })
777        }
778    }
779
780    /// Inner spawner that must never be reached when an operator is attached.
781    struct MustNotSpawn;
782
783    #[async_trait]
784    impl SpawnerAdapter for MustNotSpawn {
785        async fn spawn(
786            &self,
787            _engine: &Engine,
788            _ctx: &Ctx,
789            _task_id: StepId,
790            _attempt: u32,
791            _token: CapToken,
792        ) -> Result<Box<dyn Worker>, SpawnError> {
793            panic!("delegate axis must bypass inner.spawn when an operator is attached");
794        }
795    }
796
797    async fn seeded_engine() -> (Engine, CapToken, StepId) {
798        let engine = Engine::new(EngineCfg::default());
799        let op_token = engine
800            .attach("ut-op", Role::Operator, Duration::from_secs(30))
801            .await
802            .expect("attach");
803        let task_id = engine
804            .start_task(
805                &op_token,
806                TaskSpec {
807                    agent: "planner".to_string(),
808                    initial_directive: "do the thing".into(),
809                    step_ctx: None,
810                },
811            )
812            .await
813            .expect("start_task");
814        // Mint + register a worker token the same way
815        // `dispatch_attempt_with` does — the spawner path runs with a
816        // `Role::Worker` token (FetchPrompt is worker-verb-gated).
817        let worker_token = engine.signer().session(
818            format!("worker-of-{task_id}"),
819            Role::Worker,
820            vec!["*".into()],
821            Duration::from_secs(600),
822        );
823        let fp = worker_token.fingerprint();
824        let record = crate::core::state::CapTokenRecord::from_worker_token(
825            worker_token.clone(),
826            task_id.clone(),
827        );
828        engine
829            .with_state("test.mint_worker", move |s| {
830                s.tokens.insert(fp, record);
831            })
832            .await
833            .expect("mint worker token");
834        (engine, worker_token, task_id)
835    }
836
837    fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
838        OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
839    }
840
841    async fn recorded_worker(
842        seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
843    ) -> Option<WorkerBinding> {
844        for _ in 0..100 {
845            if let Some(w) = seen.lock().unwrap().clone() {
846                return w;
847            }
848            tokio::time::sleep(Duration::from_millis(10)).await;
849        }
850        panic!("operator.execute was never called within 1s");
851    }
852
853    #[tokio::test]
854    async fn forwards_ctx_injected_binding_to_operator_execute() {
855        let (engine, token, task_id) = seeded_engine().await;
856        let seen = Arc::new(Mutex::new(None));
857        let op = Arc::new(RecordingOperator { seen: seen.clone() });
858
859        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
860        ctx.operator.operator = Some(op);
861        ctx.meta.runtime.insert(
862            crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
863            serde_json::to_value(WorkerBinding {
864                variant: "mse-worker-coder".to_string(),
865                tools: vec!["Edit".to_string()],
866            })
867            .unwrap(),
868        );
869
870        let _worker = delegate_stack()
871            .spawn(&engine, &ctx, task_id, 1, token)
872            .await
873            .expect("delegate spawn ok");
874
875        let got = recorded_worker(&seen).await.expect("binding forwarded");
876        assert_eq!(got.variant, "mse-worker-coder");
877        assert_eq!(got.tools, vec!["Edit".to_string()]);
878    }
879
880    #[tokio::test]
881    async fn absent_binding_stays_none_no_silent_default() {
882        let (engine, token, task_id) = seeded_engine().await;
883        let seen = Arc::new(Mutex::new(None));
884        let op = Arc::new(RecordingOperator { seen: seen.clone() });
885
886        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
887        ctx.operator.operator = Some(op);
888
889        let _worker = delegate_stack()
890            .spawn(&engine, &ctx, task_id, 1, token)
891            .await
892            .expect("delegate spawn ok");
893
894        assert!(
895            recorded_worker(&seen).await.is_none(),
896            "no binding declared must reach the operator as None (fail-loud stays downstream)"
897        );
898    }
899
900    #[tokio::test]
901    async fn malformed_binding_fails_loud_before_execute() {
902        let (engine, token, task_id) = seeded_engine().await;
903        let seen = Arc::new(Mutex::new(None));
904        let op = Arc::new(RecordingOperator { seen: seen.clone() });
905
906        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
907        ctx.operator.operator = Some(op);
908        ctx.meta.runtime.insert(
909            crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
910            serde_json::json!({ "not_a_binding": true }),
911        );
912
913        let err = match delegate_stack()
914            .spawn(&engine, &ctx, task_id, 1, token)
915            .await
916        {
917            Ok(_) => panic!("malformed binding must fail the spawn"),
918            Err(e) => e,
919        };
920        let msg = format!("{err:?}");
921        assert!(
922            msg.contains("worker_binding") && msg.contains("malformed"),
923            "error must name the malformed key: {msg}"
924        );
925        assert!(
926            seen.lock().unwrap().is_none(),
927            "operator.execute must not run on malformed binding"
928        );
929    }
930}