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::blueprint::compiler::CompiledAgentTable;
63use crate::blueprint::{AuditDef, AuditMode};
64use crate::core::ctx::{Ctx, OperatorKind};
65use crate::core::engine::Engine;
66use crate::core::state::{DispatchOutcome, Event, TaskSpec};
67use crate::types::{CapToken, StepId};
68use crate::worker::adapter::{SpawnError, SpawnerAdapter};
69use crate::worker::output::{ContentRef, OutputEvent};
70use crate::worker::{wrap_join, MiddlewareWorker, Worker, WorkerJoinHandler};
71use async_trait::async_trait;
72use serde_json::Value;
73use std::sync::Arc;
74use std::time::{Duration, Instant};
75use tokio::sync::broadcast;
76
77/// Pull the terminal `Final` event's `(value, ok)` out of the tail (works
78/// for both `Inline` and `FileRef` content).
79async fn pull_final_value_ok(
80    engine: &Engine,
81    task_id: &StepId,
82    attempt: u32,
83) -> Option<(Value, bool)> {
84    let tail = engine.output_tail(task_id, attempt).await;
85    tail.iter().rev().find_map(|ev| match ev {
86        OutputEvent::Final {
87            content: ContentRef::Inline { value },
88            ok,
89        } => Some((value.clone(), *ok)),
90        OutputEvent::Final {
91            content: ContentRef::FileRef { path, .. },
92            ok,
93        } => Some((serde_json::json!({"file_ref": path.to_string_lossy()}), *ok)),
94        _ => None,
95    })
96}
97
98/// Layer trait — one middleware stage wrapping a `SpawnerAdapter`.
99pub trait SpawnerLayer: Send + Sync + 'static {
100    /// Wraps `inner` in this layer's behaviour, returning a new
101    /// `SpawnerAdapter` that delegates to `inner` (directly or via
102    /// `wrap_join`) while adding this layer's cross-cutting effect.
103    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter>;
104}
105
106/// Stack builder that layers `SpawnerLayer`s on top of a base adapter.
107///
108/// Each `.layer(...)` call wraps a new **outer** stage — same ergonomics as
109/// `tower::ServiceBuilder`.
110pub struct SpawnerStack {
111    inner: Arc<dyn SpawnerAdapter>,
112}
113
114impl SpawnerStack {
115    /// Starts a stack with `base` as the innermost adapter.
116    pub fn new(base: Arc<dyn SpawnerAdapter>) -> Self {
117        Self { inner: base }
118    }
119
120    /// Wraps the current stack with a statically-typed `SpawnerLayer`,
121    /// becoming the new outermost stage.
122    pub fn layer<L: SpawnerLayer>(mut self, layer: L) -> Self {
123        self.inner = layer.wrap(self.inner);
124        self
125    }
126
127    /// Dynamically-typed variant taking `Arc<dyn SpawnerLayer>`. Used via
128    /// the `LayerRegistry` resolution path (where a factory returns
129    /// `Arc<dyn ...>`).
130    pub fn layer_dyn(mut self, layer: Arc<dyn SpawnerLayer>) -> Self {
131        self.inner = layer.wrap(self.inner);
132        self
133    }
134
135    /// Finishes the stack, returning the fully-wrapped adapter.
136    pub fn build(self) -> Arc<dyn SpawnerAdapter> {
137        self.inner
138    }
139}
140
141// ─── SpawnerLayerFactory + LayerRegistry ─────────────────────────────────
142//
143// # Design rationale
144//
145// Wiring is assembled per-launch through `TaskLaunchService.launch`:
146//
147//   Compiler.compile(bp) ─┬─→ compiled.router (CompiledAgentTable: agent name → SpawnerAdapter dispatch)
148//                         │
149//                         │   service::linker::link(router, bp.spawner_hints.layers, &engine)
150//                         │     internal:
151//                         │       SpawnerStack::new(router)
152//                         │         .layer_dyn(base_factory_n(engine))   ← every LayerRegistry.base entry
153//                         │         .layer_dyn(hint_factory(engine))     ← resolves each bp.spawner_hints.layers key
154//                         │         .build()
155//                         ▼
156//                   EngineDispatcher::with_spawner(engine, op_token, stacked)
157//                         ▼
158//                   engine.dispatch_attempt_with(op_token, task_id, &stacked)
159//
160// # base vs hint — when to use each
161//
162// - **base layer**: wrapped around every Blueprint. Example: AuditMiddleware
163//   (a mandatory EventLog audit). The caller registers with
164//   `LayerRegistry::with_base(|e| Arc::new(AuditMiddleware::new(e.event_tx())))`.
165//
166// - **hint layer**: wrapped **only when the Blueprint declares the key** in
167//   `spawner_hints.layers`. Examples: MainAIMiddleware /
168//   SeniorEscalationMiddleware / OperatorDelegateMiddleware. The Blueprint
169//   only declares a capability key (e.g. `"main_ai"`) without knowing the
170//   implementation; the engine-side LayerRegistry resolves key → factory,
171//   keeping the pure Flow layer separate from implementation details.
172//
173// # Factory pattern (handles layers that need Engine context)
174//
175// We do not hold `Arc<dyn SpawnerLayer>` directly because some layers
176// depend on the engine instance — for example AuditMiddleware needs
177// `engine.event_tx()` and can only be built after the engine exists. A
178// factory closure defers construction: the Layer instance is created only
179// when the engine is handed in.
180
181/// Factory closure for a `SpawnerLayer`. The caller registers these at
182/// startup, and they are called with the engine context at bind time.
183/// Stateless layers can use `|_engine| Arc::new(MyLayer)`; layers that need
184/// something like `event_tx` should do `|engine| Arc::new(MyLayer::new(engine.event_tx()))`.
185pub type LayerFactory =
186    Arc<dyn Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static>;
187
188/// Registry of `LayerFactory`s, split into `base` (always applied) and
189/// `hints` (applied only when a Blueprint declares the matching key in
190/// `spawner_hints.layers`). See the module-level `# Factory pattern`
191/// notes above for why factories rather than pre-built layers.
192#[derive(Default, Clone)]
193pub struct LayerRegistry {
194    base: Vec<LayerFactory>,
195    hints: std::collections::HashMap<String, LayerFactory>,
196}
197
198impl LayerRegistry {
199    /// Empty registry (no base layers, no hint layers).
200    pub fn new() -> Self {
201        Self::default()
202    }
203
204    /// Register a base layer factory that is applied on every Blueprint bind
205    /// (for layers that must fire for every task — e.g. `AuditMiddleware`).
206    pub fn with_base<F>(mut self, factory: F) -> Self
207    where
208        F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
209    {
210        self.base.push(Arc::new(factory));
211        self
212    }
213
214    /// Register a layer factory addressable by hint key. If
215    /// `Blueprint.spawner_hints.layers` lists the same key, it is wrapped at
216    /// bind time; otherwise it is a no-op.
217    pub fn with_hint<F>(mut self, key: impl Into<String>, factory: F) -> Self
218    where
219        F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
220    {
221        self.hints.insert(key.into(), Arc::new(factory));
222        self
223    }
224
225    /// All registered base-layer factories, in registration order.
226    pub fn base_factories(&self) -> &[LayerFactory] {
227        &self.base
228    }
229
230    /// Looks up the hint-layer factory registered under `key`, if any.
231    pub fn lookup_hint(&self, key: &str) -> Option<&LayerFactory> {
232        self.hints.get(key)
233    }
234}
235
236// ─── AuditMiddleware (pushes into the EventLog broadcast path) ────────────
237
238/// Mandatory base layer that emits `Event::TaskAttemptStarted` on every
239/// spawn, before delegating. This is the audit trail's entry point into
240/// the EventLog broadcast channel.
241pub struct AuditMiddleware {
242    /// Broadcast sender the EventLog subscribes to.
243    pub event_tx: broadcast::Sender<Event>,
244}
245
246impl AuditMiddleware {
247    /// Wraps a broadcast sender to notify on every spawn.
248    pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
249        Self { event_tx }
250    }
251}
252
253impl SpawnerLayer for AuditMiddleware {
254    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
255        Arc::new(AuditWrapped {
256            inner,
257            event_tx: self.event_tx.clone(),
258        })
259    }
260}
261
262struct AuditWrapped {
263    inner: Arc<dyn SpawnerAdapter>,
264    event_tx: broadcast::Sender<Event>,
265}
266
267#[async_trait]
268impl SpawnerAdapter for AuditWrapped {
269    async fn spawn(
270        &self,
271        engine: &Engine,
272        ctx: &Ctx,
273        task_id: StepId,
274        attempt: u32,
275        token: CapToken,
276    ) -> Result<Box<dyn Worker>, SpawnError> {
277        let _ = self.event_tx.send(Event::TaskAttemptStarted {
278            task_id: task_id.clone(),
279            attempt,
280        });
281        self.inner.spawn(engine, ctx, task_id, attempt, token).await
282    }
283}
284
285// ─── MainAIMiddleware (fires SpawnHook before/after for MainAI/Composite) ─
286
287/// Hint layer that fires `ctx.operator.spawn_hook.before`/`after` around
288/// a spawn, but only for `MainAi` / `Composite` sessions. No-op for
289/// other kinds (still delegates, just skips the hook calls).
290pub struct MainAIMiddleware;
291
292impl MainAIMiddleware {
293    /// Stateless constructor.
294    pub fn new() -> Self {
295        Self
296    }
297}
298
299impl Default for MainAIMiddleware {
300    fn default() -> Self {
301        Self::new()
302    }
303}
304
305impl SpawnerLayer for MainAIMiddleware {
306    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
307        Arc::new(MainAIWrapped { inner })
308    }
309}
310
311struct MainAIWrapped {
312    inner: Arc<dyn SpawnerAdapter>,
313}
314
315#[async_trait]
316impl SpawnerAdapter for MainAIWrapped {
317    async fn spawn(
318        &self,
319        engine: &Engine,
320        ctx: &Ctx,
321        task_id: StepId,
322        attempt: u32,
323        token: CapToken,
324    ) -> Result<Box<dyn Worker>, SpawnError> {
325        let mainai = matches!(
326            ctx.operator.kind,
327            OperatorKind::MainAi | OperatorKind::Composite
328        );
329        if mainai {
330            if let Some(hook) = &ctx.operator.spawn_hook {
331                hook.before(ctx)
332                    .await
333                    .map_err(SpawnError::RejectedByMiddleware)?;
334            }
335        }
336
337        let handle = self
338            .inner
339            .spawn(engine, ctx, task_id.clone(), attempt, token)
340            .await?;
341
342        if !mainai {
343            return Ok(handle);
344        }
345        let Some(hook) = ctx.operator.spawn_hook.clone() else {
346            return Ok(handle);
347        };
348
349        // Wrap the completion signal and call hook.after on finish.
350        // Pull the last Final from engine.output_tail as the value.
351        let ctx_clone = ctx.clone();
352        let engine_clone = engine.clone();
353        let task_id_clone = task_id.clone();
354        Ok(wrap_join(handle, move |signal| {
355            let hook = hook.clone();
356            let ctx_clone = ctx_clone.clone();
357            let engine_clone = engine_clone.clone();
358            let task_id_clone = task_id_clone.clone();
359            async move {
360                let v = match &signal {
361                    Ok(()) => pull_final_value_ok(&engine_clone, &task_id_clone, attempt)
362                        .await
363                        .map(|(v, _)| v)
364                        .unwrap_or(Value::Null),
365                    Err(e) => Value::String(e.to_string()),
366                };
367                let _ = hook.after(&ctx_clone, &v).await;
368                signal
369            }
370        }))
371    }
372}
373
374// ─── SeniorEscalationMiddleware ───────────────────────────────────────────
375//
376// When a spawn's completion is `ok=false` and `ctx.operator.senior_bridge` is
377// Some, this auxiliary layer calls `SeniorBridge.ask`, merges the answer into
378// `WorkerResult.value` under `"senior_answer"`, and upgrades the result to
379// `ok=true`. Retry / re-dispatch is the engine (operator) side's job; this
380// layer only injects fresh material for that decision.
381
382/// Hint layer: on `ok=false` completion with `ctx.operator.senior_bridge`
383/// set, asks the bridge for guidance and pushes an override `Final`
384/// (`ok=true`) carrying `senior_answer`. See the module comment above
385/// this type for the full contract.
386pub struct SeniorEscalationMiddleware;
387
388impl SeniorEscalationMiddleware {
389    /// Stateless constructor.
390    pub fn new() -> Self {
391        Self
392    }
393}
394
395impl Default for SeniorEscalationMiddleware {
396    fn default() -> Self {
397        Self::new()
398    }
399}
400
401impl SpawnerLayer for SeniorEscalationMiddleware {
402    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
403        Arc::new(SeniorWrapped { inner })
404    }
405}
406
407struct SeniorWrapped {
408    inner: Arc<dyn SpawnerAdapter>,
409}
410
411#[async_trait]
412impl SpawnerAdapter for SeniorWrapped {
413    async fn spawn(
414        &self,
415        engine: &Engine,
416        ctx: &Ctx,
417        task_id: StepId,
418        attempt: u32,
419        token: CapToken,
420    ) -> Result<Box<dyn Worker>, SpawnError> {
421        let bridge = ctx.operator.senior_bridge.clone();
422        let task_id_for_hook = task_id.clone();
423        let engine_clone = engine.clone();
424        let token_clone = token.clone();
425        let handle = self
426            .inner
427            .spawn(engine, ctx, task_id, attempt, token)
428            .await?;
429        let Some(bridge) = bridge else {
430            return Ok(handle);
431        };
432        Ok(wrap_join(handle, move |signal| {
433            let bridge = bridge.clone();
434            let task_id = task_id_for_hook.clone();
435            let engine = engine_clone.clone();
436            let token = token_clone.clone();
437            async move {
438                signal?;
439                // Read the existing Final.
440                let last = pull_final_value_ok(&engine, &task_id, attempt).await;
441                if let Some((value, false)) = last {
442                    // ok=false: escalate to senior and push an override Final.
443                    let question = serde_json::json!({
444                        "reason": "worker reported ok=false",
445                        "value": value.clone(),
446                    });
447                    if let Ok(answer) = bridge.ask(&task_id, question).await {
448                        let override_val = serde_json::json!({
449                            "original": value,
450                            "senior_answer": answer,
451                        });
452                        let _ = engine
453                            .submit_output(
454                                &token,
455                                &task_id,
456                                attempt,
457                                OutputEvent::Final {
458                                    content: ContentRef::Inline {
459                                        value: override_val,
460                                    },
461                                    ok: true,
462                                },
463                            )
464                            .await;
465                    }
466                }
467                Ok(())
468            }
469        }))
470    }
471}
472
473// ─── OperatorDelegateMiddleware (delegates the whole spawn to an external Operator when one is attached) ──
474
475/// When `ctx.operator.operator.is_some()` (the session has an Operator
476/// backend), **bypass** `inner.spawn`, call `operator.execute(ctx, prompt)`,
477/// and box the result up as a `WorkerHandle`. In other words: the path that
478/// hands "this spawn" to whatever external Operator backend the engine has
479/// registered.
480///
481/// # Independent of `OperatorKind` (Operator is a generic abstraction)
482///
483/// An earlier implementation gated on `kind == MainAi | Composite`, which
484/// tied the `Operator` abstraction to an "AI driver" assumption — a design
485/// weakness. The `Operator` trait is a generic **external processing backend**
486/// (LLM, human, external resource, side-effectful operation — anything), and
487/// is orthogonal to the kind axis.
488///
489/// The current implementation decides solely on `operator.is_some()`:
490/// - Automate session + operator backend registered → delegate
491///   (pure external-execution delegation).
492/// - MainAi session + operator backend registered → delegate.
493/// - Any kind + `operator` `None` → pass through (normal `inner.spawn`).
494///
495/// `kind` still matters as a firing condition for `SpawnHook`s over in
496/// `MainAIMiddleware`, but this middleware ignores it.
497///
498/// # Split of responsibilities with `OperatorSpawner`
499///
500/// The two axes exist for different reasons:
501///
502/// - **This middleware — the Blueprint-global (session) axis.** Delegate every
503///   agent to the same Operator backend. The `operator_backend_id` is set
504///   at session-attach time; `ctx.agent` is ignored and every spawn in that
505///   session is routed through the operator (e.g. a MainAI-wide driver, or a
506///   human-wide console). The Blueprint doesn't have to talk about `kind` —
507///   it just declares the capability hint `"operator_delegate"` (keeping the
508///   Blueprint clean).
509///
510/// - **`OperatorSpawner` — the AgentSpec axis.** Each `AgentDef` bakes its
511///   own Operator backend. `kind = Operator` `AgentDef`s pick a backend via
512///   `spec.operator_ref`; the compiler bakes an `Arc<dyn Operator>` into
513///   `routes[agent_name]`. Agents loaded via the `agent.md` loader come in
514///   through this path (their default is `kind = Operator`).
515///
516/// # Exclusivity
517///
518/// When both are effective — this middleware's hint is declared, the session
519/// has an operator backend, **and** the Blueprint has a `kind = Operator`
520/// `AgentDef` — this middleware sits at the outer end of the stack and
521/// **completely bypasses** `inner.spawn`. The `OperatorSpawner` is never
522/// reached, so a double fire cannot occur by construction; the AgentSpec
523/// axis is inert. Consistent use means picking one axis per use case.
524pub struct OperatorDelegateMiddleware;
525
526impl OperatorDelegateMiddleware {
527    /// Stateless constructor.
528    pub fn new() -> Self {
529        Self
530    }
531}
532
533impl Default for OperatorDelegateMiddleware {
534    fn default() -> Self {
535        Self::new()
536    }
537}
538
539impl SpawnerLayer for OperatorDelegateMiddleware {
540    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
541        Arc::new(OperatorDelegateWrapped { inner })
542    }
543}
544
545struct OperatorDelegateWrapped {
546    inner: Arc<dyn SpawnerAdapter>,
547}
548
549#[async_trait]
550impl SpawnerAdapter for OperatorDelegateWrapped {
551    async fn spawn(
552        &self,
553        engine: &Engine,
554        ctx: &Ctx,
555        task_id: StepId,
556        attempt: u32,
557        token: CapToken,
558    ) -> Result<Box<dyn Worker>, SpawnError> {
559        // Kind-independent: we decide purely on whether an operator backend is
560        // registered on the session. `kind` matters for SpawnHook-style layers
561        // (MainAIMiddleware); this middleware does not consult it.
562        let Some(operator) = ctx.operator.operator.clone() else {
563            return self.inner.spawn(engine, ctx, task_id, attempt, token).await;
564        };
565
566        // Delegate: same shape as OperatorSpawner — fetch_prompt + operator.execute + Final emit.
567        let prompt = engine
568            .fetch_prompt(&token, &task_id)
569            .await
570            .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
571
572        // Resolve the Blueprint-baked worker binding injected into
573        // `ctx.meta.runtime` by `WorkerBindingMiddleware` (launch-time layer,
574        // built from `AgentDef.profile.worker_binding`). Absent key = agent
575        // declared no binding → hand `None` and let binding-requiring
576        // backends fail loud (`requires_worker_binding`). A present-but-
577        // malformed value is a wiring bug, not a degrade case — fail here.
578        let worker: Option<crate::operator::WorkerBinding> = match ctx
579            .meta
580            .runtime
581            .get(crate::middleware::worker_binding::WORKER_BINDING_KEY)
582        {
583            Some(v) => Some(serde_json::from_value(v.clone()).map_err(|e| {
584                SpawnError::Internal(format!(
585                    "ctx.meta.runtime['{}'] for agent '{}' is malformed: {e}",
586                    crate::middleware::worker_binding::WORKER_BINDING_KEY,
587                    ctx.agent
588                ))
589            })?),
590            None => None,
591        };
592
593        let engine_clone = engine.clone();
594        let token_clone = token.clone();
595        let token_for_op = token.clone();
596        let task_id_clone = task_id.clone();
597        let ctx_clone = ctx.clone();
598        let (tx, rx) = tokio::sync::oneshot::channel();
599        let cancel = tokio_util::sync::CancellationToken::new();
600        let cancel_inner = cancel.clone();
601        let worker_id = crate::types::WorkerId::new();
602        // issue #11: WorkerId was minted but never observable anywhere;
603        // surface it in the trace log, tied to the step it serves.
604        tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (delegate axis)");
605
606        tokio::spawn(async move {
607            let result: Result<
608                crate::worker::adapter::WorkerResult,
609                crate::worker::adapter::WorkerError,
610            > = tokio::select! {
611                // OperatorDelegateMiddleware = session-global Operator delegation.
612                // Baking per-AgentDef profile.system_prompt is OperatorSpawner's
613                // job; this path has no per-agent spawner, so system stays None.
614                // The worker binding, however, IS resolved on this axis now:
615                // `WorkerBindingMiddleware` (launch-time layer) injects the
616                // Blueprint-baked binding into ctx.meta.runtime and we forward
617                // it here — the delegate axis is a first-class variant-dispatch
618                // path, not a binding-less fallback (issue 45db42a7).
619                // We hand the capability token (Role::Worker, TTL from
620                // `EngineCfg::worker_token_ttl_secs` — default 1800s —
621                // minted by `Engine::dispatch_attempt_with`) to the
622                // operator as `worker_token` — thin-spawn operators (e.g. a
623                // WebSocket-backed operator session) forward it to the SubAgent
624                // via encode(), while Operator impls that call the LLM directly
625                // may ignore it.
626                r = operator.execute(&ctx_clone, None, prompt, worker, token_for_op) => r,
627                _ = cancel_inner.cancelled() => Err(crate::worker::adapter::WorkerError::Cancelled),
628            };
629            let result = result.map(|wr| wr.ensure_worker_kind("operator"));
630            if let Ok(wr) = &result {
631                // Stats sidecar (operator axis): the WS ack may carry the
632                // Operator's proxy report of the SubAgent's usage — forward
633                // it to the engine so the dispatcher's outcome fold lands it
634                // on the terminal StepEntry (same funnel as the InProc /
635                // subprocess fold sites). Even without an ack-attached
636                // stats blob, `ensure_worker_kind` above guarantees a
637                // `worker_kind: "operator"` label always rides.
638                if let Some(stats) = wr.stats.clone() {
639                    engine_clone
640                        .record_worker_stats(&task_id_clone, attempt, stats)
641                        .await;
642                }
643                // If the SubAgent has already pushed a Final through
644                // /v1/worker/result or /v1/worker/submit POST, skip a second
645                // emit here — the POST value is the canonical one (protocol
646                // design intent). Operator impls that never POST (e.g. tests
647                // and inline Operators) still get the fallback emit.
648                let tail = engine_clone.output_tail(&task_id_clone, attempt).await;
649                let has_final = tail
650                    .iter()
651                    .any(|ev| matches!(ev, crate::worker::output::OutputEvent::Final { .. }));
652                if !has_final {
653                    let ev = crate::worker::output::OutputEvent::Final {
654                        content: crate::worker::output::ContentRef::Inline {
655                            value: wr.value.clone(),
656                        },
657                        ok: wr.ok,
658                    };
659                    let _ = engine_clone
660                        .submit_output(&token_clone, &task_id_clone, attempt, ev)
661                        .await;
662                }
663            }
664            let signal: Result<(), crate::worker::adapter::WorkerError> = result.map(|_| ());
665            let _ = tx.send(signal);
666        });
667
668        Ok(Box::new(MiddlewareWorker {
669            handler: WorkerJoinHandler {
670                worker_id,
671                cancel,
672                completion: rx,
673            },
674        }))
675    }
676}
677
678// ─── LongHoldMiddleware (warns on the EventLog if completion time exceeds default_hold) ─
679
680/// Base layer that emits `Event::TaskAttemptCompleted` with a
681/// `long_hold_warn` marker when a spawn's completion takes longer than
682/// `default_hold`. Purely observational — it never alters the signal or
683/// blocks completion.
684pub struct LongHoldMiddleware {
685    /// Threshold above which a completion is flagged as long-held.
686    pub default_hold: Duration,
687    /// Broadcast sender the EventLog subscribes to.
688    pub event_tx: broadcast::Sender<Event>,
689}
690
691impl LongHoldMiddleware {
692    /// Sets the hold threshold and the event sender to warn through.
693    pub fn new(default_hold: Duration, event_tx: broadcast::Sender<Event>) -> Self {
694        Self {
695            default_hold,
696            event_tx,
697        }
698    }
699}
700
701impl SpawnerLayer for LongHoldMiddleware {
702    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
703        Arc::new(LongHoldWrapped {
704            inner,
705            default_hold: self.default_hold,
706            event_tx: self.event_tx.clone(),
707        })
708    }
709}
710
711struct LongHoldWrapped {
712    inner: Arc<dyn SpawnerAdapter>,
713    default_hold: Duration,
714    event_tx: broadcast::Sender<Event>,
715}
716
717#[async_trait]
718impl SpawnerAdapter for LongHoldWrapped {
719    async fn spawn(
720        &self,
721        engine: &Engine,
722        ctx: &Ctx,
723        task_id: StepId,
724        attempt: u32,
725        token: CapToken,
726    ) -> Result<Box<dyn Worker>, SpawnError> {
727        let handle = self
728            .inner
729            .spawn(engine, ctx, task_id.clone(), attempt, token)
730            .await?;
731        let started = Instant::now();
732        let default_hold = self.default_hold;
733        let event_tx = self.event_tx.clone();
734        let task_id_inner = task_id.clone();
735        let engine_for_trace = engine.clone();
736        Ok(wrap_join(handle, move |signal| {
737            let elapsed = started.elapsed();
738            let default_hold = default_hold;
739            let event_tx = event_tx.clone();
740            let task_id_inner = task_id_inner.clone();
741            let engine_for_trace = engine_for_trace.clone();
742            async move {
743                if elapsed > default_hold {
744                    let _ = event_tx.send(Event::TaskAttemptCompleted {
745                        task_id: task_id_inner.clone(),
746                        attempt,
747                        result: serde_json::json!({
748                            "long_hold_warn": true,
749                            "elapsed_ms": elapsed.as_millis() as u64,
750                            "default_hold_ms": default_hold.as_millis() as u64,
751                        }),
752                    });
753                    // RunTrace rail: mirror the warn onto the persisted
754                    // per-Run stream via the dispatcher-registered handle
755                    // (`Engine::trace_handle`) — the middleware
756                    // insertion-point exemplar. No handle (traceless
757                    // dispatch) = no-op; append itself is best-effort.
758                    if let Some(trace) = engine_for_trace.trace_handle(&task_id_inner).await {
759                        trace
760                            .append(
761                                crate::store::trace::kind::LONG_HOLD_WARN,
762                                None,
763                                Some(attempt),
764                                serde_json::json!({
765                                    "elapsed_ms": elapsed.as_millis() as u64,
766                                    "default_hold_ms": default_hold.as_millis() as u64,
767                                }),
768                            )
769                            .await;
770                    }
771                }
772                signal
773            }
774        }))
775    }
776}
777
778// ─── AfterRunAuditMiddleware (GH #34: Blueprint-declared after-run audit hooks) ──
779
780/// One-paragraph instruction handed to the audit agent alongside the
781/// structured `after_run_audit` envelope (see [`AfterRunAuditMiddleware`]
782/// for the full contract).
783const AUDIT_INSTRUCTION: &str = "Inspect this step's transcript/output for degradations, tool \
784    failures, or silent fallbacks, and emit your findings as a structured JSON object in your \
785    final output.";
786
787/// Blueprint-declared after-run audit hook layer (GH #34).
788///
789/// Wraps every spawn. After a matched step's inner signal SETTLES (`Ok`),
790/// dispatches the Blueprint-declared audit agent(s) for that step as an
791/// independent, synthetic sub-task — via `Engine::start_task` +
792/// `Engine::dispatch_attempt_with`, the same "recursive swarming" path a
793/// `Role::Worker` token is allow-listed for (`types::WORKER_SWARM_VERBS`) —
794/// reusing the AUDITED step's own worker token. Findings are persisted as
795/// an `OutputEvent::Artifact` named `"audit:<step_ref>"` on the AUDITED
796/// step's own output tail. Downstream steps read those findings via
797/// `WorkerPayload.context.steps["audit:<step_ref>"]` (fold-final drops
798/// them from the BP-chain value, but `Engine::submit_output` dual-writes
799/// every Artifact into `OutputStore` keyed by its own name — see
800/// `src/core/engine.rs`).
801///
802/// # Invariant (observational-only, binding — issue.md #1/#2/#3)
803///
804/// Every failure in the audit path (spawn/dispatch failure, audit worker
805/// failure, submit failure) is `tracing::warn!`-logged and swallowed. The
806/// audited step's own signal, returned to the caller, is ALWAYS the
807/// original inner signal, bit-for-bit — same `signal?; ...; Ok(())` shape
808/// as `SeniorEscalationMiddleware` above, so an inner `Err` short-circuits
809/// the audit entirely and propagates untouched, and an inner `Ok(())`
810/// always returns as `Ok(())` regardless of what happens inside the audit.
811///
812/// # Recursion guard
813///
814/// An agent name declared as an `AuditDef.agent` (an "auditor") is never
815/// itself audited — even if a real flow Step happens to be named after a
816/// declared auditor (e.g. a Blueprint audits every step via `steps: None`
817/// and also has a flow Step literally named after the auditor). The
818/// audit's OWN dispatch additionally never revisits this layer to begin
819/// with: it goes through `router` (the raw `CompiledAgentTable` —
820/// `Compiler::compile`'s name→adapter table), not the fully-layered stack
821/// this middleware itself sits inside, so there is no path back into
822/// `AfterRunAuditWrapped::spawn` from an audit dispatch. The name-set
823/// check in `audit_def_matches_step` (below) is a second, independent
824/// belt-and-suspenders guard for the real-flow-Step scenario.
825///
826/// Wired conditionally by `service::task_launch::TaskLaunchService::launch`
827/// (empty `Blueprint.audits` → no layer, invariant #4 — byte-identical
828/// behavior).
829pub struct AfterRunAuditMiddleware {
830    defs: Vec<AuditDef>,
831    router: Arc<CompiledAgentTable>,
832}
833
834impl AfterRunAuditMiddleware {
835    /// Holds the audit defs relevant to wiring, and the compiled
836    /// name→adapter table (`Compiler::compile`'s `CompiledBlueprint.router`)
837    /// used to dispatch each audit agent by name via
838    /// `Engine::start_task` + `Engine::dispatch_attempt_with` — the
839    /// narrowest handle that resolves an agent name to its
840    /// `SpawnerAdapter` without re-entering this same layer (see the
841    /// module comment's Recursion guard section).
842    pub fn new(defs: Vec<AuditDef>, router: Arc<CompiledAgentTable>) -> Self {
843        Self { defs, router }
844    }
845}
846
847impl SpawnerLayer for AfterRunAuditMiddleware {
848    fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
849        Arc::new(AfterRunAuditWrapped {
850            inner,
851            defs: self.defs.clone(),
852            router: self.router.clone(),
853        })
854    }
855}
856
857struct AfterRunAuditWrapped {
858    inner: Arc<dyn SpawnerAdapter>,
859    defs: Vec<AuditDef>,
860    router: Arc<CompiledAgentTable>,
861}
862
863/// Whether `def` applies to a step whose agent ref is `step_ref`. `None`,
864/// or a list containing the literal `"*"`, matches every step; otherwise
865/// only an exact name match. `Some(vec![])` (declared-but-empty) matches
866/// nothing.
867fn audit_def_matches_step(def: &AuditDef, step_ref: &str) -> bool {
868    match &def.steps {
869        None => true,
870        Some(list) => list.iter().any(|s| s == "*" || s == step_ref),
871    }
872}
873
874/// Dispatches one audit agent as an independent sub-task and — best
875/// effort — appends its findings as an `OutputEvent::Artifact` named
876/// `"audit:<step_ref>"` on the AUDITED task's own output tail. See the
877/// module comment above [`AfterRunAuditMiddleware`] for the full
878/// contract; every failure path here only `tracing::warn!`s and returns
879/// (invariant #1 — the audited step's outcome is unaffected regardless).
880#[allow(clippy::too_many_arguments)]
881async fn run_one_audit(
882    engine: &Engine,
883    router: &Arc<CompiledAgentTable>,
884    token: &CapToken,
885    audited_task_id: &StepId,
886    attempt: u32,
887    step_ref: &str,
888    audit_agent: &str,
889    directive: Value,
890) {
891    let spec = TaskSpec {
892        agent: audit_agent.to_string(),
893        initial_directive: directive,
894        step_ctx: None,
895        check_policy: None,
896    };
897    let audit_task_id = match engine.start_task(token, spec).await {
898        Ok(tid) => tid,
899        Err(e) => {
900            tracing::warn!(
901                audited_task_id = %audited_task_id,
902                step_ref,
903                audit_agent,
904                error = %e,
905                "AfterRunAuditMiddleware: start_task failed for audit agent; \
906                 audited step's outcome is unaffected"
907            );
908            return;
909        }
910    };
911    let spawner: Arc<dyn SpawnerAdapter> = router.clone();
912    let findings = match engine
913        .dispatch_attempt_with(token, &audit_task_id, &spawner, None)
914        .await
915    {
916        Ok(DispatchOutcome::Pass(v)) | Ok(DispatchOutcome::Blocked(v)) => v,
917        Ok(other) => {
918            tracing::warn!(
919                audited_task_id = %audited_task_id,
920                step_ref,
921                audit_agent,
922                outcome = ?other,
923                "AfterRunAuditMiddleware: audit agent did not settle (Pass/Blocked); \
924                 audited step's outcome is unaffected"
925            );
926            return;
927        }
928        Err(e) => {
929            tracing::warn!(
930                audited_task_id = %audited_task_id,
931                step_ref,
932                audit_agent,
933                error = %e,
934                "AfterRunAuditMiddleware: dispatch_attempt_with failed for audit agent; \
935                 audited step's outcome is unaffected"
936            );
937            return;
938        }
939    };
940    if let Err(e) = engine
941        .submit_output(
942            token,
943            audited_task_id,
944            attempt,
945            OutputEvent::Artifact {
946                name: format!("audit:{step_ref}"),
947                content: ContentRef::Inline { value: findings },
948            },
949        )
950        .await
951    {
952        tracing::warn!(
953            audited_task_id = %audited_task_id,
954            step_ref,
955            audit_agent,
956            error = %e,
957            "AfterRunAuditMiddleware: submit_output failed for audit findings; \
958             audited step's outcome is unaffected"
959        );
960    }
961}
962
963#[async_trait]
964impl SpawnerAdapter for AfterRunAuditWrapped {
965    async fn spawn(
966        &self,
967        engine: &Engine,
968        ctx: &Ctx,
969        task_id: StepId,
970        attempt: u32,
971        token: CapToken,
972    ) -> Result<Box<dyn Worker>, SpawnError> {
973        let step_ref = ctx.agent.clone();
974        let handle = self
975            .inner
976            .spawn(engine, ctx, task_id.clone(), attempt, token.clone())
977            .await?;
978
979        // Recursion guard (see the module comment's Recursion guard
980        // section): an auditor's own spawn is never itself audited.
981        let is_auditor = self.defs.iter().any(|d| d.agent == step_ref);
982        let matched: Vec<AuditDef> = if is_auditor {
983            Vec::new()
984        } else {
985            self.defs
986                .iter()
987                .filter(|d| audit_def_matches_step(d, &step_ref))
988                .cloned()
989                .collect()
990        };
991
992        if matched.is_empty() {
993            return Ok(handle);
994        }
995
996        let engine = engine.clone();
997        let router = self.router.clone();
998        Ok(wrap_join(handle, move |signal| async move {
999            // INVARIANT (issue.md #1): `signal?` propagates an inner
1000            // `Err` untouched (short-circuits the audit entirely); an
1001            // inner `Ok(())` falls through to the `Ok(())` at the bottom
1002            // of this block — byte-identical to what we matched on. The
1003            // returned signal is ALWAYS the original inner signal,
1004            // bit-for-bit.
1005            signal?;
1006
1007            let (final_value, ok) = pull_final_value_ok(&engine, &task_id, attempt)
1008                .await
1009                .unwrap_or((Value::Null, true));
1010
1011            for def in matched {
1012                let directive = serde_json::json!({
1013                    "kind": "after_run_audit",
1014                    "task_id": task_id.to_string(),
1015                    "step_ref": step_ref.clone(),
1016                    "attempt": attempt,
1017                    "ok": ok,
1018                    "final_value": final_value.clone(),
1019                    "instruction": AUDIT_INSTRUCTION,
1020                });
1021                match def.mode {
1022                    AuditMode::Sync => {
1023                        run_one_audit(
1024                            &engine, &router, &token, &task_id, attempt, &step_ref, &def.agent,
1025                            directive,
1026                        )
1027                        .await;
1028                    }
1029                    AuditMode::Async => {
1030                        let engine = engine.clone();
1031                        let router = router.clone();
1032                        let token = token.clone();
1033                        let task_id = task_id.clone();
1034                        let step_ref = step_ref.clone();
1035                        let agent = def.agent.clone();
1036                        tokio::spawn(async move {
1037                            run_one_audit(
1038                                &engine, &router, &token, &task_id, attempt, &step_ref, &agent,
1039                                directive,
1040                            )
1041                            .await;
1042                        });
1043                    }
1044                }
1045            }
1046            Ok(())
1047        }))
1048    }
1049}
1050
1051// Boundary regression spec for the delegate-axis worker-binding handoff
1052// (issue 45db42a7): OperatorDelegateMiddleware must forward the binding
1053// injected into ctx.meta.runtime by WorkerBindingMiddleware — both the
1054// hit path (Some(worker) reaches Operator::execute) and the absent path
1055// (None reaches it), plus fail-loud on a malformed value.
1056#[cfg(test)]
1057mod operator_delegate_worker_binding_tests {
1058    use super::*;
1059    use crate::core::config::EngineCfg;
1060    use crate::core::state::TaskSpec;
1061    use crate::operator::WorkerBinding;
1062    use crate::types::Role;
1063    use crate::worker::adapter::{WorkerError, WorkerResult};
1064    use std::sync::Mutex;
1065
1066    /// Operator stub recording the `worker` argument it was executed with.
1067    struct RecordingOperator {
1068        seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
1069    }
1070
1071    #[async_trait]
1072    impl crate::operator::Operator for RecordingOperator {
1073        async fn execute(
1074            &self,
1075            _ctx: &Ctx,
1076            _system: Option<String>,
1077            _prompt: Value,
1078            worker: Option<WorkerBinding>,
1079            _worker_token: CapToken,
1080        ) -> Result<WorkerResult, WorkerError> {
1081            *self.seen.lock().unwrap() = Some(worker);
1082            Ok(WorkerResult {
1083                value: Value::Null,
1084                ok: true,
1085                stats: None,
1086            })
1087        }
1088    }
1089
1090    /// Inner spawner that must never be reached when an operator is attached.
1091    struct MustNotSpawn;
1092
1093    #[async_trait]
1094    impl SpawnerAdapter for MustNotSpawn {
1095        async fn spawn(
1096            &self,
1097            _engine: &Engine,
1098            _ctx: &Ctx,
1099            _task_id: StepId,
1100            _attempt: u32,
1101            _token: CapToken,
1102        ) -> Result<Box<dyn Worker>, SpawnError> {
1103            panic!("delegate axis must bypass inner.spawn when an operator is attached");
1104        }
1105    }
1106
1107    async fn seeded_engine() -> (Engine, CapToken, StepId) {
1108        let engine = Engine::new(EngineCfg::default());
1109        let op_token = engine
1110            .attach("ut-op", Role::Operator, Duration::from_secs(30))
1111            .await
1112            .expect("attach");
1113        let task_id = engine
1114            .start_task(
1115                &op_token,
1116                TaskSpec {
1117                    agent: "planner".to_string(),
1118                    initial_directive: "do the thing".into(),
1119                    step_ctx: None,
1120                    check_policy: None,
1121                },
1122            )
1123            .await
1124            .expect("start_task");
1125        // Mint + register a worker token the same way
1126        // `dispatch_attempt_with` does — the spawner path runs with a
1127        // `Role::Worker` token (FetchPrompt is worker-verb-gated).
1128        let worker_token = engine.signer().session(
1129            format!("worker-of-{task_id}"),
1130            Role::Worker,
1131            vec!["*".into()],
1132            Duration::from_secs(600),
1133        );
1134        let fp = worker_token.fingerprint();
1135        let record = crate::core::state::CapTokenRecord::from_worker_token(
1136            worker_token.clone(),
1137            task_id.clone(),
1138        );
1139        engine
1140            .with_state("test.mint_worker", move |s| {
1141                s.tokens.insert(fp, record);
1142            })
1143            .await
1144            .expect("mint worker token");
1145        (engine, worker_token, task_id)
1146    }
1147
1148    fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
1149        OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
1150    }
1151
1152    async fn recorded_worker(
1153        seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
1154    ) -> Option<WorkerBinding> {
1155        for _ in 0..100 {
1156            if let Some(w) = seen.lock().unwrap().clone() {
1157                return w;
1158            }
1159            tokio::time::sleep(Duration::from_millis(10)).await;
1160        }
1161        panic!("operator.execute was never called within 1s");
1162    }
1163
1164    #[tokio::test]
1165    async fn forwards_ctx_injected_binding_to_operator_execute() {
1166        let (engine, token, task_id) = seeded_engine().await;
1167        let seen = Arc::new(Mutex::new(None));
1168        let op = Arc::new(RecordingOperator { seen: seen.clone() });
1169
1170        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1171        ctx.operator.operator = Some(op);
1172        ctx.meta.runtime.insert(
1173            crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1174            serde_json::to_value(WorkerBinding {
1175                variant: "code-worker".to_string(),
1176                tools: vec!["Edit".to_string()],
1177                request_digest: None,
1178                requested_model: None,
1179            })
1180            .unwrap(),
1181        );
1182
1183        let _worker = delegate_stack()
1184            .spawn(&engine, &ctx, task_id, 1, token)
1185            .await
1186            .expect("delegate spawn ok");
1187
1188        let got = recorded_worker(&seen).await.expect("binding forwarded");
1189        assert_eq!(got.variant, "code-worker");
1190        assert_eq!(got.tools, vec!["Edit".to_string()]);
1191    }
1192
1193    #[tokio::test]
1194    async fn absent_binding_stays_none_no_silent_default() {
1195        let (engine, token, task_id) = seeded_engine().await;
1196        let seen = Arc::new(Mutex::new(None));
1197        let op = Arc::new(RecordingOperator { seen: seen.clone() });
1198
1199        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1200        ctx.operator.operator = Some(op);
1201
1202        let _worker = delegate_stack()
1203            .spawn(&engine, &ctx, task_id, 1, token)
1204            .await
1205            .expect("delegate spawn ok");
1206
1207        assert!(
1208            recorded_worker(&seen).await.is_none(),
1209            "no binding declared must reach the operator as None (fail-loud stays downstream)"
1210        );
1211    }
1212
1213    #[tokio::test]
1214    async fn malformed_binding_fails_loud_before_execute() {
1215        let (engine, token, task_id) = seeded_engine().await;
1216        let seen = Arc::new(Mutex::new(None));
1217        let op = Arc::new(RecordingOperator { seen: seen.clone() });
1218
1219        let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1220        ctx.operator.operator = Some(op);
1221        ctx.meta.runtime.insert(
1222            crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1223            serde_json::json!({ "not_a_binding": true }),
1224        );
1225
1226        let err = match delegate_stack()
1227            .spawn(&engine, &ctx, task_id, 1, token)
1228            .await
1229        {
1230            Ok(_) => panic!("malformed binding must fail the spawn"),
1231            Err(e) => e,
1232        };
1233        let msg = format!("{err:?}");
1234        assert!(
1235            msg.contains("worker_binding") && msg.contains("malformed"),
1236            "error must name the malformed key: {msg}"
1237        );
1238        assert!(
1239            seen.lock().unwrap().is_none(),
1240            "operator.execute must not run on malformed binding"
1241        );
1242    }
1243}
1244
1245// ─── GH #34: `AfterRunAuditMiddleware` ─────────────────────────────────────
1246#[cfg(test)]
1247mod after_run_audit_tests {
1248    use super::*;
1249    use crate::blueprint::compiler::{Compiler, RustFnInProcessSpawnerFactory, SpawnerRegistry};
1250    use crate::blueprint::{
1251        current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
1252        CompilerStrategy,
1253    };
1254    use crate::core::config::EngineCfg;
1255    use crate::types::Role;
1256    use crate::worker::adapter::{WorkerError as StubWorkerError, WorkerResult};
1257    use mlua_flow_ir::Node as FlowNode;
1258
1259    fn rustfn_agent(name: &str, fn_id: &str) -> AgentDef {
1260        AgentDef {
1261            name: name.to_string(),
1262            kind: AgentKind::RustFn,
1263            spec: serde_json::json!({ "fn_id": fn_id }),
1264            profile: None,
1265            meta: None,
1266            runner: None,
1267            runner_ref: None,
1268            verdict: None,
1269            lints: None,
1270        }
1271    }
1272
1273    fn minimal_bp(agents: Vec<AgentDef>, audits: Vec<AuditDef>) -> Blueprint {
1274        crate::blueprint::Blueprint {
1275            schema_version: current_schema_version(),
1276            id: "afterrun-audit-ut".into(),
1277            // Unused directly by these tests — each dispatches one agent's
1278            // step at a time via `run_step` (start_task +
1279            // dispatch_attempt_with), the same shape
1280            // `EngineDispatcher::dispatch` uses per flow.ir Step. The
1281            // AfterRunAudit layer keys off `ctx.agent`/`AuditDef.steps`
1282            // only, so a real multi-step flow.ir Seq is not needed to
1283            // exercise it.
1284            flow: FlowNode::Seq { children: vec![] },
1285            agents,
1286            operators: vec![],
1287            metas: vec![],
1288            hints: CompilerHints::default(),
1289            strategy: CompilerStrategy::default(),
1290            metadata: BlueprintMetadata::default(),
1291            spawner_hints: Default::default(),
1292            default_agent_kind: AgentKind::Operator,
1293            default_operator_kind: None,
1294            default_init_ctx: None,
1295            default_agent_ctx: None,
1296            default_context_policy: None,
1297            projection_placement: None,
1298            audits,
1299            degradation_policy: None,
1300            runners: vec![],
1301            default_runner: None,
1302            subprocesses: vec![],
1303            check_policy: None,
1304            blueprint_ref_includes: Vec::new(),
1305        }
1306    }
1307
1308    /// Registers three stub `RustFn` workers shared across this module's
1309    /// tests: `"worker"` (ok, generic step body), `"auditor"` (ok, fixed
1310    /// findings), `"bad-auditor"` (always fails — GH #34 test 2).
1311    fn test_registry() -> SpawnerRegistry {
1312        let factory = RustFnInProcessSpawnerFactory::new()
1313            .register_fn("worker", |_inv| async move {
1314                Ok(WorkerResult {
1315                    value: serde_json::json!({ "result": "done" }),
1316                    ok: true,
1317                    stats: None,
1318                })
1319            })
1320            .register_fn("auditor", |_inv| async move {
1321                Ok(WorkerResult {
1322                    value: serde_json::json!({ "finding": "clean" }),
1323                    ok: true,
1324                    stats: None,
1325                })
1326            })
1327            .register_fn("bad-auditor", |_inv| async move {
1328                Err(StubWorkerError::Failed("boom".to_string()))
1329            });
1330        let mut reg = SpawnerRegistry::new();
1331        reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1332        reg
1333    }
1334
1335    /// Dispatches `agent_name` as its own independent single-step task
1336    /// through `spawner` (start_task + dispatch_attempt_with — the same
1337    /// shape `EngineDispatcher::dispatch` uses per flow.ir Step), reusing
1338    /// `op_token` (a `Role::Operator` token — `start_task` mints a fresh
1339    /// `Role::Worker` token per attempt internally, exactly as
1340    /// `dispatch_attempt_with` always does).
1341    async fn run_step(
1342        engine: &Engine,
1343        op_token: &CapToken,
1344        agent_name: &str,
1345        spawner: &Arc<dyn SpawnerAdapter>,
1346    ) -> (
1347        StepId,
1348        Result<DispatchOutcome, crate::core::errors::EngineError>,
1349    ) {
1350        let task_id = engine
1351            .start_task(
1352                op_token,
1353                TaskSpec {
1354                    agent: agent_name.to_string(),
1355                    initial_directive: serde_json::json!("go"),
1356                    step_ctx: None,
1357                    check_policy: None,
1358                },
1359            )
1360            .await
1361            .expect("start_task");
1362        let outcome = engine
1363            .dispatch_attempt_with(op_token, &task_id, spawner, None)
1364            .await;
1365        (task_id, outcome)
1366    }
1367
1368    async fn seeded_op_token(engine: &Engine) -> CapToken {
1369        engine
1370            .attach("ut-op", Role::Operator, Duration::from_secs(30))
1371            .await
1372            .expect("attach")
1373    }
1374
1375    fn find_artifact(tail: &[OutputEvent], name: &str) -> Option<Value> {
1376        tail.iter().find_map(|ev| match ev {
1377            OutputEvent::Artifact {
1378                name: n,
1379                content: ContentRef::Inline { value },
1380            } if n == name => Some(value.clone()),
1381            _ => None,
1382        })
1383    }
1384
1385    /// GH #34 test 1: a matched step's Sync-mode audit appends
1386    /// `audit:<step_ref>` to the AUDITED step's own output tail, and the
1387    /// audited step's own outcome is unaffected (the worker's own value).
1388    #[tokio::test]
1389    async fn audit_fires_after_step_and_appends_artifact() {
1390        let agents = vec![
1391            rustfn_agent("worker", "worker"),
1392            rustfn_agent("auditor", "auditor"),
1393        ];
1394        let audits = vec![AuditDef {
1395            agent: "auditor".to_string(),
1396            steps: None,
1397            mode: AuditMode::Sync,
1398        }];
1399        let bp = minimal_bp(agents, audits.clone());
1400        let compiled = Compiler::new(test_registry())
1401            .compile(&bp)
1402            .expect("compile");
1403        let spawner: Arc<dyn SpawnerAdapter> =
1404            AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1405                .wrap(compiled.router.clone());
1406
1407        let engine = Engine::new(EngineCfg::default());
1408        let op_token = seeded_op_token(&engine).await;
1409        let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1410        match outcome.expect("dispatch ok") {
1411            DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1412            other => panic!("expected Pass (the worker's own outcome), got {other:?}"),
1413        }
1414
1415        let tail = engine.output_tail(&task_id, 1).await;
1416        let findings =
1417            find_artifact(&tail, "audit:worker").expect("audit:worker artifact must be appended");
1418        assert_eq!(findings, serde_json::json!({ "finding": "clean" }));
1419    }
1420
1421    /// GH #34 test 2: an auditor that errors never alters the audited
1422    /// step's own outcome or status — the failure is swallowed (a warn is
1423    /// logged, not asserted here — this asserts outcome + artifact-absence
1424    /// only, per the subtask spec).
1425    #[tokio::test]
1426    async fn audit_failure_never_alters_outcome() {
1427        let agents = vec![
1428            rustfn_agent("worker", "worker"),
1429            rustfn_agent("bad-auditor", "bad-auditor"),
1430        ];
1431        let audits = vec![AuditDef {
1432            agent: "bad-auditor".to_string(),
1433            steps: None,
1434            mode: AuditMode::Sync,
1435        }];
1436        let bp = minimal_bp(agents, audits.clone());
1437        let compiled = Compiler::new(test_registry())
1438            .compile(&bp)
1439            .expect("compile");
1440        let spawner: Arc<dyn SpawnerAdapter> =
1441            AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1442                .wrap(compiled.router.clone());
1443
1444        let engine = Engine::new(EngineCfg::default());
1445        let op_token = seeded_op_token(&engine).await;
1446        let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1447        match outcome.expect("audited step's dispatch must still succeed despite auditor failure") {
1448            DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1449            other => panic!("expected Pass identical to a no-audit run, got {other:?}"),
1450        }
1451
1452        let tail = engine.output_tail(&task_id, 1).await;
1453        assert!(
1454            find_artifact(&tail, "audit:worker").is_none(),
1455            "auditor failure must not append an audit artifact"
1456        );
1457    }
1458
1459    /// GH #34 test 3 (mirrors `audits_absent_no_layer`, exercised more
1460    /// directly against `derive_audits` in
1461    /// `service::task_launch::tests`): with no `AuditDef` at all, the base
1462    /// (unwrapped) adapter chain behaves identically — no artifact is ever
1463    /// appended.
1464    #[tokio::test]
1465    async fn no_audit_defs_appends_no_artifact() {
1466        let agents = vec![rustfn_agent("worker", "worker")];
1467        let bp = minimal_bp(agents, vec![]);
1468        let compiled = Compiler::new(test_registry())
1469            .compile(&bp)
1470            .expect("compile");
1471        let spawner: Arc<dyn SpawnerAdapter> = compiled.router.clone();
1472
1473        let engine = Engine::new(EngineCfg::default());
1474        let op_token = seeded_op_token(&engine).await;
1475        let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1476        assert!(matches!(
1477            outcome.expect("dispatch ok"),
1478            DispatchOutcome::Pass(_)
1479        ));
1480
1481        let tail = engine.output_tail(&task_id, 1).await;
1482        assert!(
1483            !tail
1484                .iter()
1485                .any(|ev| matches!(ev, OutputEvent::Artifact { .. })),
1486            "no audits declared must never append any audit artifact"
1487        );
1488    }
1489
1490    /// GH #34 test 4: `AuditDef.steps` filters which step names an audit
1491    /// applies to — only the listed step gets an artifact.
1492    #[tokio::test]
1493    async fn steps_filter_respected() {
1494        let agents = vec![
1495            rustfn_agent("a", "worker"),
1496            rustfn_agent("b", "worker"),
1497            rustfn_agent("auditor", "auditor"),
1498        ];
1499        let audits = vec![AuditDef {
1500            agent: "auditor".to_string(),
1501            steps: Some(vec!["b".to_string()]),
1502            mode: AuditMode::Sync,
1503        }];
1504        let bp = minimal_bp(agents, audits.clone());
1505        let compiled = Compiler::new(test_registry())
1506            .compile(&bp)
1507            .expect("compile");
1508        let spawner: Arc<dyn SpawnerAdapter> =
1509            AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1510                .wrap(compiled.router.clone());
1511
1512        let engine = Engine::new(EngineCfg::default());
1513        let op_token = seeded_op_token(&engine).await;
1514
1515        let (task_a, outcome_a) = run_step(&engine, &op_token, "a", &spawner).await;
1516        outcome_a.expect("dispatch a ok");
1517        let (task_b, outcome_b) = run_step(&engine, &op_token, "b", &spawner).await;
1518        outcome_b.expect("dispatch b ok");
1519
1520        let tail_a = engine.output_tail(&task_a, 1).await;
1521        assert!(
1522            find_artifact(&tail_a, "audit:a").is_none(),
1523            "step 'a' is not listed in AuditDef.steps and must not be audited"
1524        );
1525        let tail_b = engine.output_tail(&task_b, 1).await;
1526        assert!(
1527            find_artifact(&tail_b, "audit:b").is_some(),
1528            "step 'b' is listed in AuditDef.steps and must be audited"
1529        );
1530    }
1531
1532    /// GH #34 test 5: an agent name declared as an auditor is never
1533    /// itself audited, even when a Blueprint audits every step
1534    /// (`steps: None`) and a real flow Step happens to dispatch that same
1535    /// agent name.
1536    #[tokio::test]
1537    async fn auditor_not_audited() {
1538        let agents = vec![
1539            rustfn_agent("worker", "worker"),
1540            rustfn_agent("auditor", "auditor"),
1541        ];
1542        let audits = vec![AuditDef {
1543            agent: "auditor".to_string(),
1544            steps: None,
1545            mode: AuditMode::Sync,
1546        }];
1547        let bp = minimal_bp(agents, audits.clone());
1548        let compiled = Compiler::new(test_registry())
1549            .compile(&bp)
1550            .expect("compile");
1551        let spawner: Arc<dyn SpawnerAdapter> =
1552            AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1553                .wrap(compiled.router.clone());
1554
1555        let engine = Engine::new(EngineCfg::default());
1556        let op_token = seeded_op_token(&engine).await;
1557
1558        // The worker step gets audited as usual.
1559        let (worker_task, worker_outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1560        worker_outcome.expect("dispatch worker ok");
1561        let worker_tail = engine.output_tail(&worker_task, 1).await;
1562        assert!(find_artifact(&worker_tail, "audit:worker").is_some());
1563
1564        // A real flow Step happening to dispatch the "auditor" agent name
1565        // must not recurse into auditing itself.
1566        let (auditor_task, auditor_outcome) =
1567            run_step(&engine, &op_token, "auditor", &spawner).await;
1568        auditor_outcome.expect("dispatch auditor ok");
1569        let auditor_tail = engine.output_tail(&auditor_task, 1).await;
1570        assert!(
1571            find_artifact(&auditor_tail, "audit:auditor").is_none(),
1572            "an agent declared as an auditor must never audit itself"
1573        );
1574    }
1575}