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, 1800s TTL —
620 // minted by `Engine::dispatch_attempt_with`) to the
621 // operator as `worker_token` — thin-spawn operators (e.g. a
622 // WebSocket-backed operator session) forward it to the SubAgent
623 // via encode(), while Operator impls that call the LLM directly
624 // may ignore it.
625 r = operator.execute(&ctx_clone, None, prompt, worker, token_for_op) => r,
626 _ = cancel_inner.cancelled() => Err(crate::worker::adapter::WorkerError::Cancelled),
627 };
628 let result = result.map(|wr| wr.ensure_worker_kind("operator"));
629 if let Ok(wr) = &result {
630 // Stats sidecar (operator axis): the WS ack may carry the
631 // Operator's proxy report of the SubAgent's usage — forward
632 // it to the engine so the dispatcher's outcome fold lands it
633 // on the terminal StepEntry (same funnel as the InProc /
634 // subprocess fold sites). Even without an ack-attached
635 // stats blob, `ensure_worker_kind` above guarantees a
636 // `worker_kind: "operator"` label always rides.
637 if let Some(stats) = wr.stats.clone() {
638 engine_clone
639 .record_worker_stats(&task_id_clone, attempt, stats)
640 .await;
641 }
642 // If the SubAgent has already pushed a Final through
643 // /v1/worker/result or /v1/worker/submit POST, skip a second
644 // emit here — the POST value is the canonical one (protocol
645 // design intent). Operator impls that never POST (e.g. tests
646 // and inline Operators) still get the fallback emit.
647 let tail = engine_clone.output_tail(&task_id_clone, attempt).await;
648 let has_final = tail
649 .iter()
650 .any(|ev| matches!(ev, crate::worker::output::OutputEvent::Final { .. }));
651 if !has_final {
652 let ev = crate::worker::output::OutputEvent::Final {
653 content: crate::worker::output::ContentRef::Inline {
654 value: wr.value.clone(),
655 },
656 ok: wr.ok,
657 };
658 let _ = engine_clone
659 .submit_output(&token_clone, &task_id_clone, attempt, ev)
660 .await;
661 }
662 }
663 let signal: Result<(), crate::worker::adapter::WorkerError> = result.map(|_| ());
664 let _ = tx.send(signal);
665 });
666
667 Ok(Box::new(MiddlewareWorker {
668 handler: WorkerJoinHandler {
669 worker_id,
670 cancel,
671 completion: rx,
672 },
673 }))
674 }
675}
676
677// ─── LongHoldMiddleware (warns on the EventLog if completion time exceeds default_hold) ─
678
679/// Base layer that emits `Event::TaskAttemptCompleted` with a
680/// `long_hold_warn` marker when a spawn's completion takes longer than
681/// `default_hold`. Purely observational — it never alters the signal or
682/// blocks completion.
683pub struct LongHoldMiddleware {
684 /// Threshold above which a completion is flagged as long-held.
685 pub default_hold: Duration,
686 /// Broadcast sender the EventLog subscribes to.
687 pub event_tx: broadcast::Sender<Event>,
688}
689
690impl LongHoldMiddleware {
691 /// Sets the hold threshold and the event sender to warn through.
692 pub fn new(default_hold: Duration, event_tx: broadcast::Sender<Event>) -> Self {
693 Self {
694 default_hold,
695 event_tx,
696 }
697 }
698}
699
700impl SpawnerLayer for LongHoldMiddleware {
701 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
702 Arc::new(LongHoldWrapped {
703 inner,
704 default_hold: self.default_hold,
705 event_tx: self.event_tx.clone(),
706 })
707 }
708}
709
710struct LongHoldWrapped {
711 inner: Arc<dyn SpawnerAdapter>,
712 default_hold: Duration,
713 event_tx: broadcast::Sender<Event>,
714}
715
716#[async_trait]
717impl SpawnerAdapter for LongHoldWrapped {
718 async fn spawn(
719 &self,
720 engine: &Engine,
721 ctx: &Ctx,
722 task_id: StepId,
723 attempt: u32,
724 token: CapToken,
725 ) -> Result<Box<dyn Worker>, SpawnError> {
726 let handle = self
727 .inner
728 .spawn(engine, ctx, task_id.clone(), attempt, token)
729 .await?;
730 let started = Instant::now();
731 let default_hold = self.default_hold;
732 let event_tx = self.event_tx.clone();
733 let task_id_inner = task_id.clone();
734 let engine_for_trace = engine.clone();
735 Ok(wrap_join(handle, move |signal| {
736 let elapsed = started.elapsed();
737 let default_hold = default_hold;
738 let event_tx = event_tx.clone();
739 let task_id_inner = task_id_inner.clone();
740 let engine_for_trace = engine_for_trace.clone();
741 async move {
742 if elapsed > default_hold {
743 let _ = event_tx.send(Event::TaskAttemptCompleted {
744 task_id: task_id_inner.clone(),
745 attempt,
746 result: serde_json::json!({
747 "long_hold_warn": true,
748 "elapsed_ms": elapsed.as_millis() as u64,
749 "default_hold_ms": default_hold.as_millis() as u64,
750 }),
751 });
752 // RunTrace rail: mirror the warn onto the persisted
753 // per-Run stream via the dispatcher-registered handle
754 // (`Engine::trace_handle`) — the middleware
755 // insertion-point exemplar. No handle (traceless
756 // dispatch) = no-op; append itself is best-effort.
757 if let Some(trace) = engine_for_trace.trace_handle(&task_id_inner).await {
758 trace
759 .append(
760 crate::store::trace::kind::LONG_HOLD_WARN,
761 None,
762 Some(attempt),
763 serde_json::json!({
764 "elapsed_ms": elapsed.as_millis() as u64,
765 "default_hold_ms": default_hold.as_millis() as u64,
766 }),
767 )
768 .await;
769 }
770 }
771 signal
772 }
773 }))
774 }
775}
776
777// ─── AfterRunAuditMiddleware (GH #34: Blueprint-declared after-run audit hooks) ──
778
779/// One-paragraph instruction handed to the audit agent alongside the
780/// structured `after_run_audit` envelope (see [`AfterRunAuditMiddleware`]
781/// for the full contract).
782const AUDIT_INSTRUCTION: &str = "Inspect this step's transcript/output for degradations, tool \
783 failures, or silent fallbacks, and emit your findings as a structured JSON object in your \
784 final output.";
785
786/// Blueprint-declared after-run audit hook layer (GH #34).
787///
788/// Wraps every spawn. After a matched step's inner signal SETTLES (`Ok`),
789/// dispatches the Blueprint-declared audit agent(s) for that step as an
790/// independent, synthetic sub-task — via `Engine::start_task` +
791/// `Engine::dispatch_attempt_with`, the same "recursive swarming" path a
792/// `Role::Worker` token is allow-listed for (`types::WORKER_SWARM_VERBS`) —
793/// reusing the AUDITED step's own worker token. Findings are persisted as
794/// an `OutputEvent::Artifact` named `"audit:<step_ref>"` on the AUDITED
795/// step's own output tail. Downstream steps read those findings via
796/// `WorkerPayload.context.steps["audit:<step_ref>"]` (fold-final drops
797/// them from the BP-chain value, but `Engine::submit_output` dual-writes
798/// every Artifact into `OutputStore` keyed by its own name — see
799/// `src/core/engine.rs`).
800///
801/// # Invariant (observational-only, binding — issue.md #1/#2/#3)
802///
803/// Every failure in the audit path (spawn/dispatch failure, audit worker
804/// failure, submit failure) is `tracing::warn!`-logged and swallowed. The
805/// audited step's own signal, returned to the caller, is ALWAYS the
806/// original inner signal, bit-for-bit — same `signal?; ...; Ok(())` shape
807/// as `SeniorEscalationMiddleware` above, so an inner `Err` short-circuits
808/// the audit entirely and propagates untouched, and an inner `Ok(())`
809/// always returns as `Ok(())` regardless of what happens inside the audit.
810///
811/// # Recursion guard
812///
813/// An agent name declared as an `AuditDef.agent` (an "auditor") is never
814/// itself audited — even if a real flow Step happens to be named after a
815/// declared auditor (e.g. a Blueprint audits every step via `steps: None`
816/// and also has a flow Step literally named after the auditor). The
817/// audit's OWN dispatch additionally never revisits this layer to begin
818/// with: it goes through `router` (the raw `CompiledAgentTable` —
819/// `Compiler::compile`'s name→adapter table), not the fully-layered stack
820/// this middleware itself sits inside, so there is no path back into
821/// `AfterRunAuditWrapped::spawn` from an audit dispatch. The name-set
822/// check in `audit_def_matches_step` (below) is a second, independent
823/// belt-and-suspenders guard for the real-flow-Step scenario.
824///
825/// Wired conditionally by `service::task_launch::TaskLaunchService::launch`
826/// (empty `Blueprint.audits` → no layer, invariant #4 — byte-identical
827/// behavior).
828pub struct AfterRunAuditMiddleware {
829 defs: Vec<AuditDef>,
830 router: Arc<CompiledAgentTable>,
831}
832
833impl AfterRunAuditMiddleware {
834 /// Holds the audit defs relevant to wiring, and the compiled
835 /// name→adapter table (`Compiler::compile`'s `CompiledBlueprint.router`)
836 /// used to dispatch each audit agent by name via
837 /// `Engine::start_task` + `Engine::dispatch_attempt_with` — the
838 /// narrowest handle that resolves an agent name to its
839 /// `SpawnerAdapter` without re-entering this same layer (see the
840 /// module comment's Recursion guard section).
841 pub fn new(defs: Vec<AuditDef>, router: Arc<CompiledAgentTable>) -> Self {
842 Self { defs, router }
843 }
844}
845
846impl SpawnerLayer for AfterRunAuditMiddleware {
847 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
848 Arc::new(AfterRunAuditWrapped {
849 inner,
850 defs: self.defs.clone(),
851 router: self.router.clone(),
852 })
853 }
854}
855
856struct AfterRunAuditWrapped {
857 inner: Arc<dyn SpawnerAdapter>,
858 defs: Vec<AuditDef>,
859 router: Arc<CompiledAgentTable>,
860}
861
862/// Whether `def` applies to a step whose agent ref is `step_ref`. `None`,
863/// or a list containing the literal `"*"`, matches every step; otherwise
864/// only an exact name match. `Some(vec![])` (declared-but-empty) matches
865/// nothing.
866fn audit_def_matches_step(def: &AuditDef, step_ref: &str) -> bool {
867 match &def.steps {
868 None => true,
869 Some(list) => list.iter().any(|s| s == "*" || s == step_ref),
870 }
871}
872
873/// Dispatches one audit agent as an independent sub-task and — best
874/// effort — appends its findings as an `OutputEvent::Artifact` named
875/// `"audit:<step_ref>"` on the AUDITED task's own output tail. See the
876/// module comment above [`AfterRunAuditMiddleware`] for the full
877/// contract; every failure path here only `tracing::warn!`s and returns
878/// (invariant #1 — the audited step's outcome is unaffected regardless).
879#[allow(clippy::too_many_arguments)]
880async fn run_one_audit(
881 engine: &Engine,
882 router: &Arc<CompiledAgentTable>,
883 token: &CapToken,
884 audited_task_id: &StepId,
885 attempt: u32,
886 step_ref: &str,
887 audit_agent: &str,
888 directive: Value,
889) {
890 let spec = TaskSpec {
891 agent: audit_agent.to_string(),
892 initial_directive: directive,
893 step_ctx: None,
894 check_policy: None,
895 };
896 let audit_task_id = match engine.start_task(token, spec).await {
897 Ok(tid) => tid,
898 Err(e) => {
899 tracing::warn!(
900 audited_task_id = %audited_task_id,
901 step_ref,
902 audit_agent,
903 error = %e,
904 "AfterRunAuditMiddleware: start_task failed for audit agent; \
905 audited step's outcome is unaffected"
906 );
907 return;
908 }
909 };
910 let spawner: Arc<dyn SpawnerAdapter> = router.clone();
911 let findings = match engine
912 .dispatch_attempt_with(token, &audit_task_id, &spawner, None)
913 .await
914 {
915 Ok(DispatchOutcome::Pass(v)) | Ok(DispatchOutcome::Blocked(v)) => v,
916 Ok(other) => {
917 tracing::warn!(
918 audited_task_id = %audited_task_id,
919 step_ref,
920 audit_agent,
921 outcome = ?other,
922 "AfterRunAuditMiddleware: audit agent did not settle (Pass/Blocked); \
923 audited step's outcome is unaffected"
924 );
925 return;
926 }
927 Err(e) => {
928 tracing::warn!(
929 audited_task_id = %audited_task_id,
930 step_ref,
931 audit_agent,
932 error = %e,
933 "AfterRunAuditMiddleware: dispatch_attempt_with failed for audit agent; \
934 audited step's outcome is unaffected"
935 );
936 return;
937 }
938 };
939 if let Err(e) = engine
940 .submit_output(
941 token,
942 audited_task_id,
943 attempt,
944 OutputEvent::Artifact {
945 name: format!("audit:{step_ref}"),
946 content: ContentRef::Inline { value: findings },
947 },
948 )
949 .await
950 {
951 tracing::warn!(
952 audited_task_id = %audited_task_id,
953 step_ref,
954 audit_agent,
955 error = %e,
956 "AfterRunAuditMiddleware: submit_output failed for audit findings; \
957 audited step's outcome is unaffected"
958 );
959 }
960}
961
962#[async_trait]
963impl SpawnerAdapter for AfterRunAuditWrapped {
964 async fn spawn(
965 &self,
966 engine: &Engine,
967 ctx: &Ctx,
968 task_id: StepId,
969 attempt: u32,
970 token: CapToken,
971 ) -> Result<Box<dyn Worker>, SpawnError> {
972 let step_ref = ctx.agent.clone();
973 let handle = self
974 .inner
975 .spawn(engine, ctx, task_id.clone(), attempt, token.clone())
976 .await?;
977
978 // Recursion guard (see the module comment's Recursion guard
979 // section): an auditor's own spawn is never itself audited.
980 let is_auditor = self.defs.iter().any(|d| d.agent == step_ref);
981 let matched: Vec<AuditDef> = if is_auditor {
982 Vec::new()
983 } else {
984 self.defs
985 .iter()
986 .filter(|d| audit_def_matches_step(d, &step_ref))
987 .cloned()
988 .collect()
989 };
990
991 if matched.is_empty() {
992 return Ok(handle);
993 }
994
995 let engine = engine.clone();
996 let router = self.router.clone();
997 Ok(wrap_join(handle, move |signal| async move {
998 // INVARIANT (issue.md #1): `signal?` propagates an inner
999 // `Err` untouched (short-circuits the audit entirely); an
1000 // inner `Ok(())` falls through to the `Ok(())` at the bottom
1001 // of this block — byte-identical to what we matched on. The
1002 // returned signal is ALWAYS the original inner signal,
1003 // bit-for-bit.
1004 signal?;
1005
1006 let (final_value, ok) = pull_final_value_ok(&engine, &task_id, attempt)
1007 .await
1008 .unwrap_or((Value::Null, true));
1009
1010 for def in matched {
1011 let directive = serde_json::json!({
1012 "kind": "after_run_audit",
1013 "task_id": task_id.to_string(),
1014 "step_ref": step_ref.clone(),
1015 "attempt": attempt,
1016 "ok": ok,
1017 "final_value": final_value.clone(),
1018 "instruction": AUDIT_INSTRUCTION,
1019 });
1020 match def.mode {
1021 AuditMode::Sync => {
1022 run_one_audit(
1023 &engine, &router, &token, &task_id, attempt, &step_ref, &def.agent,
1024 directive,
1025 )
1026 .await;
1027 }
1028 AuditMode::Async => {
1029 let engine = engine.clone();
1030 let router = router.clone();
1031 let token = token.clone();
1032 let task_id = task_id.clone();
1033 let step_ref = step_ref.clone();
1034 let agent = def.agent.clone();
1035 tokio::spawn(async move {
1036 run_one_audit(
1037 &engine, &router, &token, &task_id, attempt, &step_ref, &agent,
1038 directive,
1039 )
1040 .await;
1041 });
1042 }
1043 }
1044 }
1045 Ok(())
1046 }))
1047 }
1048}
1049
1050// Boundary regression spec for the delegate-axis worker-binding handoff
1051// (issue 45db42a7): OperatorDelegateMiddleware must forward the binding
1052// injected into ctx.meta.runtime by WorkerBindingMiddleware — both the
1053// hit path (Some(worker) reaches Operator::execute) and the absent path
1054// (None reaches it), plus fail-loud on a malformed value.
1055#[cfg(test)]
1056mod operator_delegate_worker_binding_tests {
1057 use super::*;
1058 use crate::core::config::EngineCfg;
1059 use crate::core::state::TaskSpec;
1060 use crate::operator::WorkerBinding;
1061 use crate::types::Role;
1062 use crate::worker::adapter::{WorkerError, WorkerResult};
1063 use std::sync::Mutex;
1064
1065 /// Operator stub recording the `worker` argument it was executed with.
1066 struct RecordingOperator {
1067 seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
1068 }
1069
1070 #[async_trait]
1071 impl crate::operator::Operator for RecordingOperator {
1072 async fn execute(
1073 &self,
1074 _ctx: &Ctx,
1075 _system: Option<String>,
1076 _prompt: Value,
1077 worker: Option<WorkerBinding>,
1078 _worker_token: CapToken,
1079 ) -> Result<WorkerResult, WorkerError> {
1080 *self.seen.lock().unwrap() = Some(worker);
1081 Ok(WorkerResult {
1082 value: Value::Null,
1083 ok: true,
1084 stats: None,
1085 })
1086 }
1087 }
1088
1089 /// Inner spawner that must never be reached when an operator is attached.
1090 struct MustNotSpawn;
1091
1092 #[async_trait]
1093 impl SpawnerAdapter for MustNotSpawn {
1094 async fn spawn(
1095 &self,
1096 _engine: &Engine,
1097 _ctx: &Ctx,
1098 _task_id: StepId,
1099 _attempt: u32,
1100 _token: CapToken,
1101 ) -> Result<Box<dyn Worker>, SpawnError> {
1102 panic!("delegate axis must bypass inner.spawn when an operator is attached");
1103 }
1104 }
1105
1106 async fn seeded_engine() -> (Engine, CapToken, StepId) {
1107 let engine = Engine::new(EngineCfg::default());
1108 let op_token = engine
1109 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1110 .await
1111 .expect("attach");
1112 let task_id = engine
1113 .start_task(
1114 &op_token,
1115 TaskSpec {
1116 agent: "planner".to_string(),
1117 initial_directive: "do the thing".into(),
1118 step_ctx: None,
1119 check_policy: None,
1120 },
1121 )
1122 .await
1123 .expect("start_task");
1124 // Mint + register a worker token the same way
1125 // `dispatch_attempt_with` does — the spawner path runs with a
1126 // `Role::Worker` token (FetchPrompt is worker-verb-gated).
1127 let worker_token = engine.signer().session(
1128 format!("worker-of-{task_id}"),
1129 Role::Worker,
1130 vec!["*".into()],
1131 Duration::from_secs(600),
1132 );
1133 let fp = worker_token.fingerprint();
1134 let record = crate::core::state::CapTokenRecord::from_worker_token(
1135 worker_token.clone(),
1136 task_id.clone(),
1137 );
1138 engine
1139 .with_state("test.mint_worker", move |s| {
1140 s.tokens.insert(fp, record);
1141 })
1142 .await
1143 .expect("mint worker token");
1144 (engine, worker_token, task_id)
1145 }
1146
1147 fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
1148 OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
1149 }
1150
1151 async fn recorded_worker(
1152 seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
1153 ) -> Option<WorkerBinding> {
1154 for _ in 0..100 {
1155 if let Some(w) = seen.lock().unwrap().clone() {
1156 return w;
1157 }
1158 tokio::time::sleep(Duration::from_millis(10)).await;
1159 }
1160 panic!("operator.execute was never called within 1s");
1161 }
1162
1163 #[tokio::test]
1164 async fn forwards_ctx_injected_binding_to_operator_execute() {
1165 let (engine, token, task_id) = seeded_engine().await;
1166 let seen = Arc::new(Mutex::new(None));
1167 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1168
1169 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1170 ctx.operator.operator = Some(op);
1171 ctx.meta.runtime.insert(
1172 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1173 serde_json::to_value(WorkerBinding {
1174 variant: "code-worker".to_string(),
1175 tools: vec!["Edit".to_string()],
1176 request_digest: None,
1177 requested_model: None,
1178 })
1179 .unwrap(),
1180 );
1181
1182 let _worker = delegate_stack()
1183 .spawn(&engine, &ctx, task_id, 1, token)
1184 .await
1185 .expect("delegate spawn ok");
1186
1187 let got = recorded_worker(&seen).await.expect("binding forwarded");
1188 assert_eq!(got.variant, "code-worker");
1189 assert_eq!(got.tools, vec!["Edit".to_string()]);
1190 }
1191
1192 #[tokio::test]
1193 async fn absent_binding_stays_none_no_silent_default() {
1194 let (engine, token, task_id) = seeded_engine().await;
1195 let seen = Arc::new(Mutex::new(None));
1196 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1197
1198 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1199 ctx.operator.operator = Some(op);
1200
1201 let _worker = delegate_stack()
1202 .spawn(&engine, &ctx, task_id, 1, token)
1203 .await
1204 .expect("delegate spawn ok");
1205
1206 assert!(
1207 recorded_worker(&seen).await.is_none(),
1208 "no binding declared must reach the operator as None (fail-loud stays downstream)"
1209 );
1210 }
1211
1212 #[tokio::test]
1213 async fn malformed_binding_fails_loud_before_execute() {
1214 let (engine, token, task_id) = seeded_engine().await;
1215 let seen = Arc::new(Mutex::new(None));
1216 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1217
1218 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1219 ctx.operator.operator = Some(op);
1220 ctx.meta.runtime.insert(
1221 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1222 serde_json::json!({ "not_a_binding": true }),
1223 );
1224
1225 let err = match delegate_stack()
1226 .spawn(&engine, &ctx, task_id, 1, token)
1227 .await
1228 {
1229 Ok(_) => panic!("malformed binding must fail the spawn"),
1230 Err(e) => e,
1231 };
1232 let msg = format!("{err:?}");
1233 assert!(
1234 msg.contains("worker_binding") && msg.contains("malformed"),
1235 "error must name the malformed key: {msg}"
1236 );
1237 assert!(
1238 seen.lock().unwrap().is_none(),
1239 "operator.execute must not run on malformed binding"
1240 );
1241 }
1242}
1243
1244// ─── GH #34: `AfterRunAuditMiddleware` ─────────────────────────────────────
1245#[cfg(test)]
1246mod after_run_audit_tests {
1247 use super::*;
1248 use crate::blueprint::compiler::{Compiler, RustFnInProcessSpawnerFactory, SpawnerRegistry};
1249 use crate::blueprint::{
1250 current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
1251 CompilerStrategy,
1252 };
1253 use crate::core::config::EngineCfg;
1254 use crate::types::Role;
1255 use crate::worker::adapter::{WorkerError as StubWorkerError, WorkerResult};
1256 use mlua_flow_ir::Node as FlowNode;
1257
1258 fn rustfn_agent(name: &str, fn_id: &str) -> AgentDef {
1259 AgentDef {
1260 name: name.to_string(),
1261 kind: AgentKind::RustFn,
1262 spec: serde_json::json!({ "fn_id": fn_id }),
1263 profile: None,
1264 meta: None,
1265 runner: None,
1266 runner_ref: None,
1267 verdict: None,
1268 lints: None,
1269 }
1270 }
1271
1272 fn minimal_bp(agents: Vec<AgentDef>, audits: Vec<AuditDef>) -> Blueprint {
1273 crate::blueprint::Blueprint {
1274 schema_version: current_schema_version(),
1275 id: "afterrun-audit-ut".into(),
1276 // Unused directly by these tests — each dispatches one agent's
1277 // step at a time via `run_step` (start_task +
1278 // dispatch_attempt_with), the same shape
1279 // `EngineDispatcher::dispatch` uses per flow.ir Step. The
1280 // AfterRunAudit layer keys off `ctx.agent`/`AuditDef.steps`
1281 // only, so a real multi-step flow.ir Seq is not needed to
1282 // exercise it.
1283 flow: FlowNode::Seq { children: vec![] },
1284 agents,
1285 operators: vec![],
1286 metas: vec![],
1287 hints: CompilerHints::default(),
1288 strategy: CompilerStrategy::default(),
1289 metadata: BlueprintMetadata::default(),
1290 spawner_hints: Default::default(),
1291 default_agent_kind: AgentKind::Operator,
1292 default_operator_kind: None,
1293 default_init_ctx: None,
1294 default_agent_ctx: None,
1295 default_context_policy: None,
1296 projection_placement: None,
1297 audits,
1298 degradation_policy: None,
1299 runners: vec![],
1300 default_runner: None,
1301 subprocesses: vec![],
1302 check_policy: None,
1303 blueprint_ref_includes: Vec::new(),
1304 }
1305 }
1306
1307 /// Registers three stub `RustFn` workers shared across this module's
1308 /// tests: `"worker"` (ok, generic step body), `"auditor"` (ok, fixed
1309 /// findings), `"bad-auditor"` (always fails — GH #34 test 2).
1310 fn test_registry() -> SpawnerRegistry {
1311 let factory = RustFnInProcessSpawnerFactory::new()
1312 .register_fn("worker", |_inv| async move {
1313 Ok(WorkerResult {
1314 value: serde_json::json!({ "result": "done" }),
1315 ok: true,
1316 stats: None,
1317 })
1318 })
1319 .register_fn("auditor", |_inv| async move {
1320 Ok(WorkerResult {
1321 value: serde_json::json!({ "finding": "clean" }),
1322 ok: true,
1323 stats: None,
1324 })
1325 })
1326 .register_fn("bad-auditor", |_inv| async move {
1327 Err(StubWorkerError::Failed("boom".to_string()))
1328 });
1329 let mut reg = SpawnerRegistry::new();
1330 reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1331 reg
1332 }
1333
1334 /// Dispatches `agent_name` as its own independent single-step task
1335 /// through `spawner` (start_task + dispatch_attempt_with — the same
1336 /// shape `EngineDispatcher::dispatch` uses per flow.ir Step), reusing
1337 /// `op_token` (a `Role::Operator` token — `start_task` mints a fresh
1338 /// `Role::Worker` token per attempt internally, exactly as
1339 /// `dispatch_attempt_with` always does).
1340 async fn run_step(
1341 engine: &Engine,
1342 op_token: &CapToken,
1343 agent_name: &str,
1344 spawner: &Arc<dyn SpawnerAdapter>,
1345 ) -> (
1346 StepId,
1347 Result<DispatchOutcome, crate::core::errors::EngineError>,
1348 ) {
1349 let task_id = engine
1350 .start_task(
1351 op_token,
1352 TaskSpec {
1353 agent: agent_name.to_string(),
1354 initial_directive: serde_json::json!("go"),
1355 step_ctx: None,
1356 check_policy: None,
1357 },
1358 )
1359 .await
1360 .expect("start_task");
1361 let outcome = engine
1362 .dispatch_attempt_with(op_token, &task_id, spawner, None)
1363 .await;
1364 (task_id, outcome)
1365 }
1366
1367 async fn seeded_op_token(engine: &Engine) -> CapToken {
1368 engine
1369 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1370 .await
1371 .expect("attach")
1372 }
1373
1374 fn find_artifact(tail: &[OutputEvent], name: &str) -> Option<Value> {
1375 tail.iter().find_map(|ev| match ev {
1376 OutputEvent::Artifact {
1377 name: n,
1378 content: ContentRef::Inline { value },
1379 } if n == name => Some(value.clone()),
1380 _ => None,
1381 })
1382 }
1383
1384 /// GH #34 test 1: a matched step's Sync-mode audit appends
1385 /// `audit:<step_ref>` to the AUDITED step's own output tail, and the
1386 /// audited step's own outcome is unaffected (the worker's own value).
1387 #[tokio::test]
1388 async fn audit_fires_after_step_and_appends_artifact() {
1389 let agents = vec![
1390 rustfn_agent("worker", "worker"),
1391 rustfn_agent("auditor", "auditor"),
1392 ];
1393 let audits = vec![AuditDef {
1394 agent: "auditor".to_string(),
1395 steps: None,
1396 mode: AuditMode::Sync,
1397 }];
1398 let bp = minimal_bp(agents, audits.clone());
1399 let compiled = Compiler::new(test_registry())
1400 .compile(&bp)
1401 .expect("compile");
1402 let spawner: Arc<dyn SpawnerAdapter> =
1403 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1404 .wrap(compiled.router.clone());
1405
1406 let engine = Engine::new(EngineCfg::default());
1407 let op_token = seeded_op_token(&engine).await;
1408 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1409 match outcome.expect("dispatch ok") {
1410 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1411 other => panic!("expected Pass (the worker's own outcome), got {other:?}"),
1412 }
1413
1414 let tail = engine.output_tail(&task_id, 1).await;
1415 let findings =
1416 find_artifact(&tail, "audit:worker").expect("audit:worker artifact must be appended");
1417 assert_eq!(findings, serde_json::json!({ "finding": "clean" }));
1418 }
1419
1420 /// GH #34 test 2: an auditor that errors never alters the audited
1421 /// step's own outcome or status — the failure is swallowed (a warn is
1422 /// logged, not asserted here — this asserts outcome + artifact-absence
1423 /// only, per the subtask spec).
1424 #[tokio::test]
1425 async fn audit_failure_never_alters_outcome() {
1426 let agents = vec![
1427 rustfn_agent("worker", "worker"),
1428 rustfn_agent("bad-auditor", "bad-auditor"),
1429 ];
1430 let audits = vec![AuditDef {
1431 agent: "bad-auditor".to_string(),
1432 steps: None,
1433 mode: AuditMode::Sync,
1434 }];
1435 let bp = minimal_bp(agents, audits.clone());
1436 let compiled = Compiler::new(test_registry())
1437 .compile(&bp)
1438 .expect("compile");
1439 let spawner: Arc<dyn SpawnerAdapter> =
1440 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1441 .wrap(compiled.router.clone());
1442
1443 let engine = Engine::new(EngineCfg::default());
1444 let op_token = seeded_op_token(&engine).await;
1445 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1446 match outcome.expect("audited step's dispatch must still succeed despite auditor failure") {
1447 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1448 other => panic!("expected Pass identical to a no-audit run, got {other:?}"),
1449 }
1450
1451 let tail = engine.output_tail(&task_id, 1).await;
1452 assert!(
1453 find_artifact(&tail, "audit:worker").is_none(),
1454 "auditor failure must not append an audit artifact"
1455 );
1456 }
1457
1458 /// GH #34 test 3 (mirrors `audits_absent_no_layer`, exercised more
1459 /// directly against `derive_audits` in
1460 /// `service::task_launch::tests`): with no `AuditDef` at all, the base
1461 /// (unwrapped) adapter chain behaves identically — no artifact is ever
1462 /// appended.
1463 #[tokio::test]
1464 async fn no_audit_defs_appends_no_artifact() {
1465 let agents = vec![rustfn_agent("worker", "worker")];
1466 let bp = minimal_bp(agents, vec![]);
1467 let compiled = Compiler::new(test_registry())
1468 .compile(&bp)
1469 .expect("compile");
1470 let spawner: Arc<dyn SpawnerAdapter> = compiled.router.clone();
1471
1472 let engine = Engine::new(EngineCfg::default());
1473 let op_token = seeded_op_token(&engine).await;
1474 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1475 assert!(matches!(
1476 outcome.expect("dispatch ok"),
1477 DispatchOutcome::Pass(_)
1478 ));
1479
1480 let tail = engine.output_tail(&task_id, 1).await;
1481 assert!(
1482 !tail
1483 .iter()
1484 .any(|ev| matches!(ev, OutputEvent::Artifact { .. })),
1485 "no audits declared must never append any audit artifact"
1486 );
1487 }
1488
1489 /// GH #34 test 4: `AuditDef.steps` filters which step names an audit
1490 /// applies to — only the listed step gets an artifact.
1491 #[tokio::test]
1492 async fn steps_filter_respected() {
1493 let agents = vec![
1494 rustfn_agent("a", "worker"),
1495 rustfn_agent("b", "worker"),
1496 rustfn_agent("auditor", "auditor"),
1497 ];
1498 let audits = vec![AuditDef {
1499 agent: "auditor".to_string(),
1500 steps: Some(vec!["b".to_string()]),
1501 mode: AuditMode::Sync,
1502 }];
1503 let bp = minimal_bp(agents, audits.clone());
1504 let compiled = Compiler::new(test_registry())
1505 .compile(&bp)
1506 .expect("compile");
1507 let spawner: Arc<dyn SpawnerAdapter> =
1508 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1509 .wrap(compiled.router.clone());
1510
1511 let engine = Engine::new(EngineCfg::default());
1512 let op_token = seeded_op_token(&engine).await;
1513
1514 let (task_a, outcome_a) = run_step(&engine, &op_token, "a", &spawner).await;
1515 outcome_a.expect("dispatch a ok");
1516 let (task_b, outcome_b) = run_step(&engine, &op_token, "b", &spawner).await;
1517 outcome_b.expect("dispatch b ok");
1518
1519 let tail_a = engine.output_tail(&task_a, 1).await;
1520 assert!(
1521 find_artifact(&tail_a, "audit:a").is_none(),
1522 "step 'a' is not listed in AuditDef.steps and must not be audited"
1523 );
1524 let tail_b = engine.output_tail(&task_b, 1).await;
1525 assert!(
1526 find_artifact(&tail_b, "audit:b").is_some(),
1527 "step 'b' is listed in AuditDef.steps and must be audited"
1528 );
1529 }
1530
1531 /// GH #34 test 5: an agent name declared as an auditor is never
1532 /// itself audited, even when a Blueprint audits every step
1533 /// (`steps: None`) and a real flow Step happens to dispatch that same
1534 /// agent name.
1535 #[tokio::test]
1536 async fn auditor_not_audited() {
1537 let agents = vec![
1538 rustfn_agent("worker", "worker"),
1539 rustfn_agent("auditor", "auditor"),
1540 ];
1541 let audits = vec![AuditDef {
1542 agent: "auditor".to_string(),
1543 steps: None,
1544 mode: AuditMode::Sync,
1545 }];
1546 let bp = minimal_bp(agents, audits.clone());
1547 let compiled = Compiler::new(test_registry())
1548 .compile(&bp)
1549 .expect("compile");
1550 let spawner: Arc<dyn SpawnerAdapter> =
1551 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1552 .wrap(compiled.router.clone());
1553
1554 let engine = Engine::new(EngineCfg::default());
1555 let op_token = seeded_op_token(&engine).await;
1556
1557 // The worker step gets audited as usual.
1558 let (worker_task, worker_outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1559 worker_outcome.expect("dispatch worker ok");
1560 let worker_tail = engine.output_tail(&worker_task, 1).await;
1561 assert!(find_artifact(&worker_tail, "audit:worker").is_some());
1562
1563 // A real flow Step happening to dispatch the "auditor" agent name
1564 // must not recurse into auditing itself.
1565 let (auditor_task, auditor_outcome) =
1566 run_step(&engine, &op_token, "auditor", &spawner).await;
1567 auditor_outcome.expect("dispatch auditor ok");
1568 let auditor_tail = engine.output_tail(&auditor_task, 1).await;
1569 assert!(
1570 find_artifact(&auditor_tail, "audit:auditor").is_none(),
1571 "an agent declared as an auditor must never audit itself"
1572 );
1573 }
1574}