1pub 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
77async 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
98pub trait SpawnerLayer: Send + Sync + 'static {
100 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter>;
104}
105
106pub struct SpawnerStack {
111 inner: Arc<dyn SpawnerAdapter>,
112}
113
114impl SpawnerStack {
115 pub fn new(base: Arc<dyn SpawnerAdapter>) -> Self {
117 Self { inner: base }
118 }
119
120 pub fn layer<L: SpawnerLayer>(mut self, layer: L) -> Self {
123 self.inner = layer.wrap(self.inner);
124 self
125 }
126
127 pub fn layer_dyn(mut self, layer: Arc<dyn SpawnerLayer>) -> Self {
131 self.inner = layer.wrap(self.inner);
132 self
133 }
134
135 pub fn build(self) -> Arc<dyn SpawnerAdapter> {
137 self.inner
138 }
139}
140
141pub type LayerFactory =
186 Arc<dyn Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static>;
187
188#[derive(Default, Clone)]
193pub struct LayerRegistry {
194 base: Vec<LayerFactory>,
195 hints: std::collections::HashMap<String, LayerFactory>,
196}
197
198impl LayerRegistry {
199 pub fn new() -> Self {
201 Self::default()
202 }
203
204 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 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 pub fn base_factories(&self) -> &[LayerFactory] {
227 &self.base
228 }
229
230 pub fn lookup_hint(&self, key: &str) -> Option<&LayerFactory> {
232 self.hints.get(key)
233 }
234}
235
236pub struct AuditMiddleware {
242 pub event_tx: broadcast::Sender<Event>,
244}
245
246impl AuditMiddleware {
247 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
285pub struct MainAIMiddleware;
291
292impl MainAIMiddleware {
293 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 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
374pub struct SeniorEscalationMiddleware;
387
388impl SeniorEscalationMiddleware {
389 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 let last = pull_final_value_ok(&engine, &task_id, attempt).await;
441 if let Some((value, false)) = last {
442 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
473pub struct OperatorDelegateMiddleware;
525
526impl OperatorDelegateMiddleware {
527 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 let Some(operator) = ctx.operator.operator.clone() else {
563 return self.inner.spawn(engine, ctx, task_id, attempt, token).await;
564 };
565
566 let prompt = engine
568 .fetch_prompt(&token, &task_id)
569 .await
570 .map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
571
572 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 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 r = operator.execute(&ctx_clone, None, prompt, worker, token_for_op) => r,
626 _ = cancel_inner.cancelled() => Err(crate::worker::adapter::WorkerError::Cancelled),
627 };
628 if let Ok(wr) = &result {
629 let tail = engine_clone.output_tail(&task_id_clone, attempt).await;
635 let has_final = tail
636 .iter()
637 .any(|ev| matches!(ev, crate::worker::output::OutputEvent::Final { .. }));
638 if !has_final {
639 let ev = crate::worker::output::OutputEvent::Final {
640 content: crate::worker::output::ContentRef::Inline {
641 value: wr.value.clone(),
642 },
643 ok: wr.ok,
644 };
645 let _ = engine_clone
646 .submit_output(&token_clone, &task_id_clone, attempt, ev)
647 .await;
648 }
649 }
650 let signal: Result<(), crate::worker::adapter::WorkerError> = result.map(|_| ());
651 let _ = tx.send(signal);
652 });
653
654 Ok(Box::new(MiddlewareWorker {
655 handler: WorkerJoinHandler {
656 worker_id,
657 cancel,
658 completion: rx,
659 },
660 }))
661 }
662}
663
664pub struct LongHoldMiddleware {
671 pub default_hold: Duration,
673 pub event_tx: broadcast::Sender<Event>,
675}
676
677impl LongHoldMiddleware {
678 pub fn new(default_hold: Duration, event_tx: broadcast::Sender<Event>) -> Self {
680 Self {
681 default_hold,
682 event_tx,
683 }
684 }
685}
686
687impl SpawnerLayer for LongHoldMiddleware {
688 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
689 Arc::new(LongHoldWrapped {
690 inner,
691 default_hold: self.default_hold,
692 event_tx: self.event_tx.clone(),
693 })
694 }
695}
696
697struct LongHoldWrapped {
698 inner: Arc<dyn SpawnerAdapter>,
699 default_hold: Duration,
700 event_tx: broadcast::Sender<Event>,
701}
702
703#[async_trait]
704impl SpawnerAdapter for LongHoldWrapped {
705 async fn spawn(
706 &self,
707 engine: &Engine,
708 ctx: &Ctx,
709 task_id: StepId,
710 attempt: u32,
711 token: CapToken,
712 ) -> Result<Box<dyn Worker>, SpawnError> {
713 let handle = self
714 .inner
715 .spawn(engine, ctx, task_id.clone(), attempt, token)
716 .await?;
717 let started = Instant::now();
718 let default_hold = self.default_hold;
719 let event_tx = self.event_tx.clone();
720 let task_id_inner = task_id.clone();
721 Ok(wrap_join(handle, move |signal| {
722 let elapsed = started.elapsed();
723 let default_hold = default_hold;
724 let event_tx = event_tx.clone();
725 let task_id_inner = task_id_inner.clone();
726 async move {
727 if elapsed > default_hold {
728 let _ = event_tx.send(Event::TaskAttemptCompleted {
729 task_id: task_id_inner,
730 attempt,
731 result: serde_json::json!({
732 "long_hold_warn": true,
733 "elapsed_ms": elapsed.as_millis() as u64,
734 "default_hold_ms": default_hold.as_millis() as u64,
735 }),
736 });
737 }
738 signal
739 }
740 }))
741 }
742}
743
744const AUDIT_INSTRUCTION: &str = "Inspect this step's transcript/output for degradations, tool \
750 failures, or silent fallbacks, and emit your findings as a structured JSON object in your \
751 final output.";
752
753pub struct AfterRunAuditMiddleware {
796 defs: Vec<AuditDef>,
797 router: Arc<CompiledAgentTable>,
798}
799
800impl AfterRunAuditMiddleware {
801 pub fn new(defs: Vec<AuditDef>, router: Arc<CompiledAgentTable>) -> Self {
809 Self { defs, router }
810 }
811}
812
813impl SpawnerLayer for AfterRunAuditMiddleware {
814 fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
815 Arc::new(AfterRunAuditWrapped {
816 inner,
817 defs: self.defs.clone(),
818 router: self.router.clone(),
819 })
820 }
821}
822
823struct AfterRunAuditWrapped {
824 inner: Arc<dyn SpawnerAdapter>,
825 defs: Vec<AuditDef>,
826 router: Arc<CompiledAgentTable>,
827}
828
829fn audit_def_matches_step(def: &AuditDef, step_ref: &str) -> bool {
834 match &def.steps {
835 None => true,
836 Some(list) => list.iter().any(|s| s == "*" || s == step_ref),
837 }
838}
839
840#[allow(clippy::too_many_arguments)]
847async fn run_one_audit(
848 engine: &Engine,
849 router: &Arc<CompiledAgentTable>,
850 token: &CapToken,
851 audited_task_id: &StepId,
852 attempt: u32,
853 step_ref: &str,
854 audit_agent: &str,
855 directive: Value,
856) {
857 let spec = TaskSpec {
858 agent: audit_agent.to_string(),
859 initial_directive: directive,
860 step_ctx: None,
861 check_policy: None,
862 };
863 let audit_task_id = match engine.start_task(token, spec).await {
864 Ok(tid) => tid,
865 Err(e) => {
866 tracing::warn!(
867 audited_task_id = %audited_task_id,
868 step_ref,
869 audit_agent,
870 error = %e,
871 "AfterRunAuditMiddleware: start_task failed for audit agent; \
872 audited step's outcome is unaffected"
873 );
874 return;
875 }
876 };
877 let spawner: Arc<dyn SpawnerAdapter> = router.clone();
878 let findings = match engine
879 .dispatch_attempt_with(token, &audit_task_id, &spawner, None)
880 .await
881 {
882 Ok(DispatchOutcome::Pass(v)) | Ok(DispatchOutcome::Blocked(v)) => v,
883 Ok(other) => {
884 tracing::warn!(
885 audited_task_id = %audited_task_id,
886 step_ref,
887 audit_agent,
888 outcome = ?other,
889 "AfterRunAuditMiddleware: audit agent did not settle (Pass/Blocked); \
890 audited step's outcome is unaffected"
891 );
892 return;
893 }
894 Err(e) => {
895 tracing::warn!(
896 audited_task_id = %audited_task_id,
897 step_ref,
898 audit_agent,
899 error = %e,
900 "AfterRunAuditMiddleware: dispatch_attempt_with failed for audit agent; \
901 audited step's outcome is unaffected"
902 );
903 return;
904 }
905 };
906 if let Err(e) = engine
907 .submit_output(
908 token,
909 audited_task_id,
910 attempt,
911 OutputEvent::Artifact {
912 name: format!("audit:{step_ref}"),
913 content: ContentRef::Inline { value: findings },
914 },
915 )
916 .await
917 {
918 tracing::warn!(
919 audited_task_id = %audited_task_id,
920 step_ref,
921 audit_agent,
922 error = %e,
923 "AfterRunAuditMiddleware: submit_output failed for audit findings; \
924 audited step's outcome is unaffected"
925 );
926 }
927}
928
929#[async_trait]
930impl SpawnerAdapter for AfterRunAuditWrapped {
931 async fn spawn(
932 &self,
933 engine: &Engine,
934 ctx: &Ctx,
935 task_id: StepId,
936 attempt: u32,
937 token: CapToken,
938 ) -> Result<Box<dyn Worker>, SpawnError> {
939 let step_ref = ctx.agent.clone();
940 let handle = self
941 .inner
942 .spawn(engine, ctx, task_id.clone(), attempt, token.clone())
943 .await?;
944
945 let is_auditor = self.defs.iter().any(|d| d.agent == step_ref);
948 let matched: Vec<AuditDef> = if is_auditor {
949 Vec::new()
950 } else {
951 self.defs
952 .iter()
953 .filter(|d| audit_def_matches_step(d, &step_ref))
954 .cloned()
955 .collect()
956 };
957
958 if matched.is_empty() {
959 return Ok(handle);
960 }
961
962 let engine = engine.clone();
963 let router = self.router.clone();
964 Ok(wrap_join(handle, move |signal| async move {
965 signal?;
972
973 let (final_value, ok) = pull_final_value_ok(&engine, &task_id, attempt)
974 .await
975 .unwrap_or((Value::Null, true));
976
977 for def in matched {
978 let directive = serde_json::json!({
979 "kind": "after_run_audit",
980 "task_id": task_id.to_string(),
981 "step_ref": step_ref.clone(),
982 "attempt": attempt,
983 "ok": ok,
984 "final_value": final_value.clone(),
985 "instruction": AUDIT_INSTRUCTION,
986 });
987 match def.mode {
988 AuditMode::Sync => {
989 run_one_audit(
990 &engine, &router, &token, &task_id, attempt, &step_ref, &def.agent,
991 directive,
992 )
993 .await;
994 }
995 AuditMode::Async => {
996 let engine = engine.clone();
997 let router = router.clone();
998 let token = token.clone();
999 let task_id = task_id.clone();
1000 let step_ref = step_ref.clone();
1001 let agent = def.agent.clone();
1002 tokio::spawn(async move {
1003 run_one_audit(
1004 &engine, &router, &token, &task_id, attempt, &step_ref, &agent,
1005 directive,
1006 )
1007 .await;
1008 });
1009 }
1010 }
1011 }
1012 Ok(())
1013 }))
1014 }
1015}
1016
1017#[cfg(test)]
1023mod operator_delegate_worker_binding_tests {
1024 use super::*;
1025 use crate::core::config::EngineCfg;
1026 use crate::core::state::TaskSpec;
1027 use crate::operator::WorkerBinding;
1028 use crate::types::Role;
1029 use crate::worker::adapter::{WorkerError, WorkerResult};
1030 use std::sync::Mutex;
1031
1032 struct RecordingOperator {
1034 seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
1035 }
1036
1037 #[async_trait]
1038 impl crate::operator::Operator for RecordingOperator {
1039 async fn execute(
1040 &self,
1041 _ctx: &Ctx,
1042 _system: Option<String>,
1043 _prompt: Value,
1044 worker: Option<WorkerBinding>,
1045 _worker_token: CapToken,
1046 ) -> Result<WorkerResult, WorkerError> {
1047 *self.seen.lock().unwrap() = Some(worker);
1048 Ok(WorkerResult {
1049 value: Value::Null,
1050 ok: true,
1051 })
1052 }
1053 }
1054
1055 struct MustNotSpawn;
1057
1058 #[async_trait]
1059 impl SpawnerAdapter for MustNotSpawn {
1060 async fn spawn(
1061 &self,
1062 _engine: &Engine,
1063 _ctx: &Ctx,
1064 _task_id: StepId,
1065 _attempt: u32,
1066 _token: CapToken,
1067 ) -> Result<Box<dyn Worker>, SpawnError> {
1068 panic!("delegate axis must bypass inner.spawn when an operator is attached");
1069 }
1070 }
1071
1072 async fn seeded_engine() -> (Engine, CapToken, StepId) {
1073 let engine = Engine::new(EngineCfg::default());
1074 let op_token = engine
1075 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1076 .await
1077 .expect("attach");
1078 let task_id = engine
1079 .start_task(
1080 &op_token,
1081 TaskSpec {
1082 agent: "planner".to_string(),
1083 initial_directive: "do the thing".into(),
1084 step_ctx: None,
1085 check_policy: None,
1086 },
1087 )
1088 .await
1089 .expect("start_task");
1090 let worker_token = engine.signer().session(
1094 format!("worker-of-{task_id}"),
1095 Role::Worker,
1096 vec!["*".into()],
1097 Duration::from_secs(600),
1098 );
1099 let fp = worker_token.fingerprint();
1100 let record = crate::core::state::CapTokenRecord::from_worker_token(
1101 worker_token.clone(),
1102 task_id.clone(),
1103 );
1104 engine
1105 .with_state("test.mint_worker", move |s| {
1106 s.tokens.insert(fp, record);
1107 })
1108 .await
1109 .expect("mint worker token");
1110 (engine, worker_token, task_id)
1111 }
1112
1113 fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
1114 OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
1115 }
1116
1117 async fn recorded_worker(
1118 seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
1119 ) -> Option<WorkerBinding> {
1120 for _ in 0..100 {
1121 if let Some(w) = seen.lock().unwrap().clone() {
1122 return w;
1123 }
1124 tokio::time::sleep(Duration::from_millis(10)).await;
1125 }
1126 panic!("operator.execute was never called within 1s");
1127 }
1128
1129 #[tokio::test]
1130 async fn forwards_ctx_injected_binding_to_operator_execute() {
1131 let (engine, token, task_id) = seeded_engine().await;
1132 let seen = Arc::new(Mutex::new(None));
1133 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1134
1135 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1136 ctx.operator.operator = Some(op);
1137 ctx.meta.runtime.insert(
1138 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1139 serde_json::to_value(WorkerBinding {
1140 variant: "mse-worker-coder".to_string(),
1141 tools: vec!["Edit".to_string()],
1142 })
1143 .unwrap(),
1144 );
1145
1146 let _worker = delegate_stack()
1147 .spawn(&engine, &ctx, task_id, 1, token)
1148 .await
1149 .expect("delegate spawn ok");
1150
1151 let got = recorded_worker(&seen).await.expect("binding forwarded");
1152 assert_eq!(got.variant, "mse-worker-coder");
1153 assert_eq!(got.tools, vec!["Edit".to_string()]);
1154 }
1155
1156 #[tokio::test]
1157 async fn absent_binding_stays_none_no_silent_default() {
1158 let (engine, token, task_id) = seeded_engine().await;
1159 let seen = Arc::new(Mutex::new(None));
1160 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1161
1162 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1163 ctx.operator.operator = Some(op);
1164
1165 let _worker = delegate_stack()
1166 .spawn(&engine, &ctx, task_id, 1, token)
1167 .await
1168 .expect("delegate spawn ok");
1169
1170 assert!(
1171 recorded_worker(&seen).await.is_none(),
1172 "no binding declared must reach the operator as None (fail-loud stays downstream)"
1173 );
1174 }
1175
1176 #[tokio::test]
1177 async fn malformed_binding_fails_loud_before_execute() {
1178 let (engine, token, task_id) = seeded_engine().await;
1179 let seen = Arc::new(Mutex::new(None));
1180 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1181
1182 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1183 ctx.operator.operator = Some(op);
1184 ctx.meta.runtime.insert(
1185 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1186 serde_json::json!({ "not_a_binding": true }),
1187 );
1188
1189 let err = match delegate_stack()
1190 .spawn(&engine, &ctx, task_id, 1, token)
1191 .await
1192 {
1193 Ok(_) => panic!("malformed binding must fail the spawn"),
1194 Err(e) => e,
1195 };
1196 let msg = format!("{err:?}");
1197 assert!(
1198 msg.contains("worker_binding") && msg.contains("malformed"),
1199 "error must name the malformed key: {msg}"
1200 );
1201 assert!(
1202 seen.lock().unwrap().is_none(),
1203 "operator.execute must not run on malformed binding"
1204 );
1205 }
1206}
1207
1208#[cfg(test)]
1210mod after_run_audit_tests {
1211 use super::*;
1212 use crate::blueprint::compiler::{Compiler, RustFnInProcessSpawnerFactory, SpawnerRegistry};
1213 use crate::blueprint::{
1214 current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
1215 CompilerStrategy,
1216 };
1217 use crate::core::config::EngineCfg;
1218 use crate::types::Role;
1219 use crate::worker::adapter::{WorkerError as StubWorkerError, WorkerResult};
1220 use mlua_flow_ir::Node as FlowNode;
1221
1222 fn rustfn_agent(name: &str, fn_id: &str) -> AgentDef {
1223 AgentDef {
1224 name: name.to_string(),
1225 kind: AgentKind::RustFn,
1226 spec: serde_json::json!({ "fn_id": fn_id }),
1227 profile: None,
1228 meta: None,
1229 runner: None,
1230 runner_ref: None,
1231 verdict: None,
1232 }
1233 }
1234
1235 fn minimal_bp(agents: Vec<AgentDef>, audits: Vec<AuditDef>) -> Blueprint {
1236 crate::blueprint::Blueprint {
1237 schema_version: current_schema_version(),
1238 id: "afterrun-audit-ut".into(),
1239 flow: FlowNode::Seq { children: vec![] },
1247 agents,
1248 operators: vec![],
1249 metas: vec![],
1250 hints: CompilerHints::default(),
1251 strategy: CompilerStrategy::default(),
1252 metadata: BlueprintMetadata::default(),
1253 spawner_hints: Default::default(),
1254 default_agent_kind: AgentKind::Operator,
1255 default_operator_kind: None,
1256 default_init_ctx: None,
1257 default_agent_ctx: None,
1258 default_context_policy: None,
1259 projection_placement: None,
1260 audits,
1261 degradation_policy: None,
1262 runners: vec![],
1263 default_runner: None,
1264 check_policy: None,
1265 }
1266 }
1267
1268 fn test_registry() -> SpawnerRegistry {
1272 let factory = RustFnInProcessSpawnerFactory::new()
1273 .register_fn("worker", |_inv| async move {
1274 Ok(WorkerResult {
1275 value: serde_json::json!({ "result": "done" }),
1276 ok: true,
1277 })
1278 })
1279 .register_fn("auditor", |_inv| async move {
1280 Ok(WorkerResult {
1281 value: serde_json::json!({ "finding": "clean" }),
1282 ok: true,
1283 })
1284 })
1285 .register_fn("bad-auditor", |_inv| async move {
1286 Err(StubWorkerError::Failed("boom".to_string()))
1287 });
1288 let mut reg = SpawnerRegistry::new();
1289 reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1290 reg
1291 }
1292
1293 async fn run_step(
1300 engine: &Engine,
1301 op_token: &CapToken,
1302 agent_name: &str,
1303 spawner: &Arc<dyn SpawnerAdapter>,
1304 ) -> (
1305 StepId,
1306 Result<DispatchOutcome, crate::core::errors::EngineError>,
1307 ) {
1308 let task_id = engine
1309 .start_task(
1310 op_token,
1311 TaskSpec {
1312 agent: agent_name.to_string(),
1313 initial_directive: serde_json::json!("go"),
1314 step_ctx: None,
1315 check_policy: None,
1316 },
1317 )
1318 .await
1319 .expect("start_task");
1320 let outcome = engine
1321 .dispatch_attempt_with(op_token, &task_id, spawner, None)
1322 .await;
1323 (task_id, outcome)
1324 }
1325
1326 async fn seeded_op_token(engine: &Engine) -> CapToken {
1327 engine
1328 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1329 .await
1330 .expect("attach")
1331 }
1332
1333 fn find_artifact(tail: &[OutputEvent], name: &str) -> Option<Value> {
1334 tail.iter().find_map(|ev| match ev {
1335 OutputEvent::Artifact {
1336 name: n,
1337 content: ContentRef::Inline { value },
1338 } if n == name => Some(value.clone()),
1339 _ => None,
1340 })
1341 }
1342
1343 #[tokio::test]
1347 async fn audit_fires_after_step_and_appends_artifact() {
1348 let agents = vec![
1349 rustfn_agent("worker", "worker"),
1350 rustfn_agent("auditor", "auditor"),
1351 ];
1352 let audits = vec![AuditDef {
1353 agent: "auditor".to_string(),
1354 steps: None,
1355 mode: AuditMode::Sync,
1356 }];
1357 let bp = minimal_bp(agents, audits.clone());
1358 let compiled = Compiler::new(test_registry())
1359 .compile(&bp)
1360 .expect("compile");
1361 let spawner: Arc<dyn SpawnerAdapter> =
1362 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1363 .wrap(compiled.router.clone());
1364
1365 let engine = Engine::new(EngineCfg::default());
1366 let op_token = seeded_op_token(&engine).await;
1367 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1368 match outcome.expect("dispatch ok") {
1369 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1370 other => panic!("expected Pass (the worker's own outcome), got {other:?}"),
1371 }
1372
1373 let tail = engine.output_tail(&task_id, 1).await;
1374 let findings =
1375 find_artifact(&tail, "audit:worker").expect("audit:worker artifact must be appended");
1376 assert_eq!(findings, serde_json::json!({ "finding": "clean" }));
1377 }
1378
1379 #[tokio::test]
1384 async fn audit_failure_never_alters_outcome() {
1385 let agents = vec![
1386 rustfn_agent("worker", "worker"),
1387 rustfn_agent("bad-auditor", "bad-auditor"),
1388 ];
1389 let audits = vec![AuditDef {
1390 agent: "bad-auditor".to_string(),
1391 steps: None,
1392 mode: AuditMode::Sync,
1393 }];
1394 let bp = minimal_bp(agents, audits.clone());
1395 let compiled = Compiler::new(test_registry())
1396 .compile(&bp)
1397 .expect("compile");
1398 let spawner: Arc<dyn SpawnerAdapter> =
1399 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1400 .wrap(compiled.router.clone());
1401
1402 let engine = Engine::new(EngineCfg::default());
1403 let op_token = seeded_op_token(&engine).await;
1404 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1405 match outcome.expect("audited step's dispatch must still succeed despite auditor failure") {
1406 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1407 other => panic!("expected Pass identical to a no-audit run, got {other:?}"),
1408 }
1409
1410 let tail = engine.output_tail(&task_id, 1).await;
1411 assert!(
1412 find_artifact(&tail, "audit:worker").is_none(),
1413 "auditor failure must not append an audit artifact"
1414 );
1415 }
1416
1417 #[tokio::test]
1423 async fn no_audit_defs_appends_no_artifact() {
1424 let agents = vec![rustfn_agent("worker", "worker")];
1425 let bp = minimal_bp(agents, vec![]);
1426 let compiled = Compiler::new(test_registry())
1427 .compile(&bp)
1428 .expect("compile");
1429 let spawner: Arc<dyn SpawnerAdapter> = compiled.router.clone();
1430
1431 let engine = Engine::new(EngineCfg::default());
1432 let op_token = seeded_op_token(&engine).await;
1433 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1434 assert!(matches!(
1435 outcome.expect("dispatch ok"),
1436 DispatchOutcome::Pass(_)
1437 ));
1438
1439 let tail = engine.output_tail(&task_id, 1).await;
1440 assert!(
1441 !tail
1442 .iter()
1443 .any(|ev| matches!(ev, OutputEvent::Artifact { .. })),
1444 "no audits declared must never append any audit artifact"
1445 );
1446 }
1447
1448 #[tokio::test]
1451 async fn steps_filter_respected() {
1452 let agents = vec![
1453 rustfn_agent("a", "worker"),
1454 rustfn_agent("b", "worker"),
1455 rustfn_agent("auditor", "auditor"),
1456 ];
1457 let audits = vec![AuditDef {
1458 agent: "auditor".to_string(),
1459 steps: Some(vec!["b".to_string()]),
1460 mode: AuditMode::Sync,
1461 }];
1462 let bp = minimal_bp(agents, audits.clone());
1463 let compiled = Compiler::new(test_registry())
1464 .compile(&bp)
1465 .expect("compile");
1466 let spawner: Arc<dyn SpawnerAdapter> =
1467 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1468 .wrap(compiled.router.clone());
1469
1470 let engine = Engine::new(EngineCfg::default());
1471 let op_token = seeded_op_token(&engine).await;
1472
1473 let (task_a, outcome_a) = run_step(&engine, &op_token, "a", &spawner).await;
1474 outcome_a.expect("dispatch a ok");
1475 let (task_b, outcome_b) = run_step(&engine, &op_token, "b", &spawner).await;
1476 outcome_b.expect("dispatch b ok");
1477
1478 let tail_a = engine.output_tail(&task_a, 1).await;
1479 assert!(
1480 find_artifact(&tail_a, "audit:a").is_none(),
1481 "step 'a' is not listed in AuditDef.steps and must not be audited"
1482 );
1483 let tail_b = engine.output_tail(&task_b, 1).await;
1484 assert!(
1485 find_artifact(&tail_b, "audit:b").is_some(),
1486 "step 'b' is listed in AuditDef.steps and must be audited"
1487 );
1488 }
1489
1490 #[tokio::test]
1495 async fn auditor_not_audited() {
1496 let agents = vec![
1497 rustfn_agent("worker", "worker"),
1498 rustfn_agent("auditor", "auditor"),
1499 ];
1500 let audits = vec![AuditDef {
1501 agent: "auditor".to_string(),
1502 steps: None,
1503 mode: AuditMode::Sync,
1504 }];
1505 let bp = minimal_bp(agents, audits.clone());
1506 let compiled = Compiler::new(test_registry())
1507 .compile(&bp)
1508 .expect("compile");
1509 let spawner: Arc<dyn SpawnerAdapter> =
1510 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1511 .wrap(compiled.router.clone());
1512
1513 let engine = Engine::new(EngineCfg::default());
1514 let op_token = seeded_op_token(&engine).await;
1515
1516 let (worker_task, worker_outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1518 worker_outcome.expect("dispatch worker ok");
1519 let worker_tail = engine.output_tail(&worker_task, 1).await;
1520 assert!(find_artifact(&worker_tail, "audit:worker").is_some());
1521
1522 let (auditor_task, auditor_outcome) =
1525 run_step(&engine, &op_token, "auditor", &spawner).await;
1526 auditor_outcome.expect("dispatch auditor ok");
1527 let auditor_tail = engine.output_tail(&auditor_task, 1).await;
1528 assert!(
1529 find_artifact(&auditor_tail, "audit:auditor").is_none(),
1530 "an agent declared as an auditor must never audit itself"
1531 );
1532 }
1533}