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 };
862 let audit_task_id = match engine.start_task(token, spec).await {
863 Ok(tid) => tid,
864 Err(e) => {
865 tracing::warn!(
866 audited_task_id = %audited_task_id,
867 step_ref,
868 audit_agent,
869 error = %e,
870 "AfterRunAuditMiddleware: start_task failed for audit agent; \
871 audited step's outcome is unaffected"
872 );
873 return;
874 }
875 };
876 let spawner: Arc<dyn SpawnerAdapter> = router.clone();
877 let findings = match engine
878 .dispatch_attempt_with(token, &audit_task_id, &spawner, None)
879 .await
880 {
881 Ok(DispatchOutcome::Pass(v)) | Ok(DispatchOutcome::Blocked(v)) => v,
882 Ok(other) => {
883 tracing::warn!(
884 audited_task_id = %audited_task_id,
885 step_ref,
886 audit_agent,
887 outcome = ?other,
888 "AfterRunAuditMiddleware: audit agent did not settle (Pass/Blocked); \
889 audited step's outcome is unaffected"
890 );
891 return;
892 }
893 Err(e) => {
894 tracing::warn!(
895 audited_task_id = %audited_task_id,
896 step_ref,
897 audit_agent,
898 error = %e,
899 "AfterRunAuditMiddleware: dispatch_attempt_with failed for audit agent; \
900 audited step's outcome is unaffected"
901 );
902 return;
903 }
904 };
905 if let Err(e) = engine
906 .submit_output(
907 token,
908 audited_task_id,
909 attempt,
910 OutputEvent::Artifact {
911 name: format!("audit:{step_ref}"),
912 content: ContentRef::Inline { value: findings },
913 },
914 )
915 .await
916 {
917 tracing::warn!(
918 audited_task_id = %audited_task_id,
919 step_ref,
920 audit_agent,
921 error = %e,
922 "AfterRunAuditMiddleware: submit_output failed for audit findings; \
923 audited step's outcome is unaffected"
924 );
925 }
926}
927
928#[async_trait]
929impl SpawnerAdapter for AfterRunAuditWrapped {
930 async fn spawn(
931 &self,
932 engine: &Engine,
933 ctx: &Ctx,
934 task_id: StepId,
935 attempt: u32,
936 token: CapToken,
937 ) -> Result<Box<dyn Worker>, SpawnError> {
938 let step_ref = ctx.agent.clone();
939 let handle = self
940 .inner
941 .spawn(engine, ctx, task_id.clone(), attempt, token.clone())
942 .await?;
943
944 let is_auditor = self.defs.iter().any(|d| d.agent == step_ref);
947 let matched: Vec<AuditDef> = if is_auditor {
948 Vec::new()
949 } else {
950 self.defs
951 .iter()
952 .filter(|d| audit_def_matches_step(d, &step_ref))
953 .cloned()
954 .collect()
955 };
956
957 if matched.is_empty() {
958 return Ok(handle);
959 }
960
961 let engine = engine.clone();
962 let router = self.router.clone();
963 Ok(wrap_join(handle, move |signal| async move {
964 signal?;
971
972 let (final_value, ok) = pull_final_value_ok(&engine, &task_id, attempt)
973 .await
974 .unwrap_or((Value::Null, true));
975
976 for def in matched {
977 let directive = serde_json::json!({
978 "kind": "after_run_audit",
979 "task_id": task_id.to_string(),
980 "step_ref": step_ref.clone(),
981 "attempt": attempt,
982 "ok": ok,
983 "final_value": final_value.clone(),
984 "instruction": AUDIT_INSTRUCTION,
985 });
986 match def.mode {
987 AuditMode::Sync => {
988 run_one_audit(
989 &engine, &router, &token, &task_id, attempt, &step_ref, &def.agent,
990 directive,
991 )
992 .await;
993 }
994 AuditMode::Async => {
995 let engine = engine.clone();
996 let router = router.clone();
997 let token = token.clone();
998 let task_id = task_id.clone();
999 let step_ref = step_ref.clone();
1000 let agent = def.agent.clone();
1001 tokio::spawn(async move {
1002 run_one_audit(
1003 &engine, &router, &token, &task_id, attempt, &step_ref, &agent,
1004 directive,
1005 )
1006 .await;
1007 });
1008 }
1009 }
1010 }
1011 Ok(())
1012 }))
1013 }
1014}
1015
1016#[cfg(test)]
1022mod operator_delegate_worker_binding_tests {
1023 use super::*;
1024 use crate::core::config::EngineCfg;
1025 use crate::core::state::TaskSpec;
1026 use crate::operator::WorkerBinding;
1027 use crate::types::Role;
1028 use crate::worker::adapter::{WorkerError, WorkerResult};
1029 use std::sync::Mutex;
1030
1031 struct RecordingOperator {
1033 seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
1034 }
1035
1036 #[async_trait]
1037 impl crate::operator::Operator for RecordingOperator {
1038 async fn execute(
1039 &self,
1040 _ctx: &Ctx,
1041 _system: Option<String>,
1042 _prompt: Value,
1043 worker: Option<WorkerBinding>,
1044 _worker_token: CapToken,
1045 ) -> Result<WorkerResult, WorkerError> {
1046 *self.seen.lock().unwrap() = Some(worker);
1047 Ok(WorkerResult {
1048 value: Value::Null,
1049 ok: true,
1050 })
1051 }
1052 }
1053
1054 struct MustNotSpawn;
1056
1057 #[async_trait]
1058 impl SpawnerAdapter for MustNotSpawn {
1059 async fn spawn(
1060 &self,
1061 _engine: &Engine,
1062 _ctx: &Ctx,
1063 _task_id: StepId,
1064 _attempt: u32,
1065 _token: CapToken,
1066 ) -> Result<Box<dyn Worker>, SpawnError> {
1067 panic!("delegate axis must bypass inner.spawn when an operator is attached");
1068 }
1069 }
1070
1071 async fn seeded_engine() -> (Engine, CapToken, StepId) {
1072 let engine = Engine::new(EngineCfg::default());
1073 let op_token = engine
1074 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1075 .await
1076 .expect("attach");
1077 let task_id = engine
1078 .start_task(
1079 &op_token,
1080 TaskSpec {
1081 agent: "planner".to_string(),
1082 initial_directive: "do the thing".into(),
1083 step_ctx: None,
1084 },
1085 )
1086 .await
1087 .expect("start_task");
1088 let worker_token = engine.signer().session(
1092 format!("worker-of-{task_id}"),
1093 Role::Worker,
1094 vec!["*".into()],
1095 Duration::from_secs(600),
1096 );
1097 let fp = worker_token.fingerprint();
1098 let record = crate::core::state::CapTokenRecord::from_worker_token(
1099 worker_token.clone(),
1100 task_id.clone(),
1101 );
1102 engine
1103 .with_state("test.mint_worker", move |s| {
1104 s.tokens.insert(fp, record);
1105 })
1106 .await
1107 .expect("mint worker token");
1108 (engine, worker_token, task_id)
1109 }
1110
1111 fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
1112 OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
1113 }
1114
1115 async fn recorded_worker(
1116 seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
1117 ) -> Option<WorkerBinding> {
1118 for _ in 0..100 {
1119 if let Some(w) = seen.lock().unwrap().clone() {
1120 return w;
1121 }
1122 tokio::time::sleep(Duration::from_millis(10)).await;
1123 }
1124 panic!("operator.execute was never called within 1s");
1125 }
1126
1127 #[tokio::test]
1128 async fn forwards_ctx_injected_binding_to_operator_execute() {
1129 let (engine, token, task_id) = seeded_engine().await;
1130 let seen = Arc::new(Mutex::new(None));
1131 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1132
1133 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1134 ctx.operator.operator = Some(op);
1135 ctx.meta.runtime.insert(
1136 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1137 serde_json::to_value(WorkerBinding {
1138 variant: "mse-worker-coder".to_string(),
1139 tools: vec!["Edit".to_string()],
1140 })
1141 .unwrap(),
1142 );
1143
1144 let _worker = delegate_stack()
1145 .spawn(&engine, &ctx, task_id, 1, token)
1146 .await
1147 .expect("delegate spawn ok");
1148
1149 let got = recorded_worker(&seen).await.expect("binding forwarded");
1150 assert_eq!(got.variant, "mse-worker-coder");
1151 assert_eq!(got.tools, vec!["Edit".to_string()]);
1152 }
1153
1154 #[tokio::test]
1155 async fn absent_binding_stays_none_no_silent_default() {
1156 let (engine, token, task_id) = seeded_engine().await;
1157 let seen = Arc::new(Mutex::new(None));
1158 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1159
1160 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1161 ctx.operator.operator = Some(op);
1162
1163 let _worker = delegate_stack()
1164 .spawn(&engine, &ctx, task_id, 1, token)
1165 .await
1166 .expect("delegate spawn ok");
1167
1168 assert!(
1169 recorded_worker(&seen).await.is_none(),
1170 "no binding declared must reach the operator as None (fail-loud stays downstream)"
1171 );
1172 }
1173
1174 #[tokio::test]
1175 async fn malformed_binding_fails_loud_before_execute() {
1176 let (engine, token, task_id) = seeded_engine().await;
1177 let seen = Arc::new(Mutex::new(None));
1178 let op = Arc::new(RecordingOperator { seen: seen.clone() });
1179
1180 let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
1181 ctx.operator.operator = Some(op);
1182 ctx.meta.runtime.insert(
1183 crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
1184 serde_json::json!({ "not_a_binding": true }),
1185 );
1186
1187 let err = match delegate_stack()
1188 .spawn(&engine, &ctx, task_id, 1, token)
1189 .await
1190 {
1191 Ok(_) => panic!("malformed binding must fail the spawn"),
1192 Err(e) => e,
1193 };
1194 let msg = format!("{err:?}");
1195 assert!(
1196 msg.contains("worker_binding") && msg.contains("malformed"),
1197 "error must name the malformed key: {msg}"
1198 );
1199 assert!(
1200 seen.lock().unwrap().is_none(),
1201 "operator.execute must not run on malformed binding"
1202 );
1203 }
1204}
1205
1206#[cfg(test)]
1208mod after_run_audit_tests {
1209 use super::*;
1210 use crate::blueprint::compiler::{Compiler, RustFnInProcessSpawnerFactory, SpawnerRegistry};
1211 use crate::blueprint::{
1212 current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
1213 CompilerStrategy,
1214 };
1215 use crate::core::config::EngineCfg;
1216 use crate::types::Role;
1217 use crate::worker::adapter::{WorkerError as StubWorkerError, WorkerResult};
1218 use mlua_flow_ir::Node as FlowNode;
1219
1220 fn rustfn_agent(name: &str, fn_id: &str) -> AgentDef {
1221 AgentDef {
1222 name: name.to_string(),
1223 kind: AgentKind::RustFn,
1224 spec: serde_json::json!({ "fn_id": fn_id }),
1225 profile: None,
1226 meta: None,
1227 }
1228 }
1229
1230 fn minimal_bp(agents: Vec<AgentDef>, audits: Vec<AuditDef>) -> Blueprint {
1231 crate::blueprint::Blueprint {
1232 schema_version: current_schema_version(),
1233 id: "afterrun-audit-ut".into(),
1234 flow: FlowNode::Seq { children: vec![] },
1242 agents,
1243 operators: vec![],
1244 metas: vec![],
1245 hints: CompilerHints::default(),
1246 strategy: CompilerStrategy::default(),
1247 metadata: BlueprintMetadata::default(),
1248 spawner_hints: Default::default(),
1249 default_agent_kind: AgentKind::Operator,
1250 default_operator_kind: None,
1251 default_init_ctx: None,
1252 default_agent_ctx: None,
1253 default_context_policy: None,
1254 projection_placement: None,
1255 audits,
1256 degradation_policy: None,
1257 }
1258 }
1259
1260 fn test_registry() -> SpawnerRegistry {
1264 let factory = RustFnInProcessSpawnerFactory::new()
1265 .register_fn("worker", |_inv| async move {
1266 Ok(WorkerResult {
1267 value: serde_json::json!({ "result": "done" }),
1268 ok: true,
1269 })
1270 })
1271 .register_fn("auditor", |_inv| async move {
1272 Ok(WorkerResult {
1273 value: serde_json::json!({ "finding": "clean" }),
1274 ok: true,
1275 })
1276 })
1277 .register_fn("bad-auditor", |_inv| async move {
1278 Err(StubWorkerError::Failed("boom".to_string()))
1279 });
1280 let mut reg = SpawnerRegistry::new();
1281 reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1282 reg
1283 }
1284
1285 async fn run_step(
1292 engine: &Engine,
1293 op_token: &CapToken,
1294 agent_name: &str,
1295 spawner: &Arc<dyn SpawnerAdapter>,
1296 ) -> (
1297 StepId,
1298 Result<DispatchOutcome, crate::core::errors::EngineError>,
1299 ) {
1300 let task_id = engine
1301 .start_task(
1302 op_token,
1303 TaskSpec {
1304 agent: agent_name.to_string(),
1305 initial_directive: serde_json::json!("go"),
1306 step_ctx: None,
1307 },
1308 )
1309 .await
1310 .expect("start_task");
1311 let outcome = engine
1312 .dispatch_attempt_with(op_token, &task_id, spawner, None)
1313 .await;
1314 (task_id, outcome)
1315 }
1316
1317 async fn seeded_op_token(engine: &Engine) -> CapToken {
1318 engine
1319 .attach("ut-op", Role::Operator, Duration::from_secs(30))
1320 .await
1321 .expect("attach")
1322 }
1323
1324 fn find_artifact(tail: &[OutputEvent], name: &str) -> Option<Value> {
1325 tail.iter().find_map(|ev| match ev {
1326 OutputEvent::Artifact {
1327 name: n,
1328 content: ContentRef::Inline { value },
1329 } if n == name => Some(value.clone()),
1330 _ => None,
1331 })
1332 }
1333
1334 #[tokio::test]
1338 async fn audit_fires_after_step_and_appends_artifact() {
1339 let agents = vec![
1340 rustfn_agent("worker", "worker"),
1341 rustfn_agent("auditor", "auditor"),
1342 ];
1343 let audits = vec![AuditDef {
1344 agent: "auditor".to_string(),
1345 steps: None,
1346 mode: AuditMode::Sync,
1347 }];
1348 let bp = minimal_bp(agents, audits.clone());
1349 let compiled = Compiler::new(test_registry())
1350 .compile(&bp)
1351 .expect("compile");
1352 let spawner: Arc<dyn SpawnerAdapter> =
1353 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1354 .wrap(compiled.router.clone());
1355
1356 let engine = Engine::new(EngineCfg::default());
1357 let op_token = seeded_op_token(&engine).await;
1358 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1359 match outcome.expect("dispatch ok") {
1360 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1361 other => panic!("expected Pass (the worker's own outcome), got {other:?}"),
1362 }
1363
1364 let tail = engine.output_tail(&task_id, 1).await;
1365 let findings =
1366 find_artifact(&tail, "audit:worker").expect("audit:worker artifact must be appended");
1367 assert_eq!(findings, serde_json::json!({ "finding": "clean" }));
1368 }
1369
1370 #[tokio::test]
1375 async fn audit_failure_never_alters_outcome() {
1376 let agents = vec![
1377 rustfn_agent("worker", "worker"),
1378 rustfn_agent("bad-auditor", "bad-auditor"),
1379 ];
1380 let audits = vec![AuditDef {
1381 agent: "bad-auditor".to_string(),
1382 steps: None,
1383 mode: AuditMode::Sync,
1384 }];
1385 let bp = minimal_bp(agents, audits.clone());
1386 let compiled = Compiler::new(test_registry())
1387 .compile(&bp)
1388 .expect("compile");
1389 let spawner: Arc<dyn SpawnerAdapter> =
1390 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1391 .wrap(compiled.router.clone());
1392
1393 let engine = Engine::new(EngineCfg::default());
1394 let op_token = seeded_op_token(&engine).await;
1395 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1396 match outcome.expect("audited step's dispatch must still succeed despite auditor failure") {
1397 DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
1398 other => panic!("expected Pass identical to a no-audit run, got {other:?}"),
1399 }
1400
1401 let tail = engine.output_tail(&task_id, 1).await;
1402 assert!(
1403 find_artifact(&tail, "audit:worker").is_none(),
1404 "auditor failure must not append an audit artifact"
1405 );
1406 }
1407
1408 #[tokio::test]
1414 async fn no_audit_defs_appends_no_artifact() {
1415 let agents = vec![rustfn_agent("worker", "worker")];
1416 let bp = minimal_bp(agents, vec![]);
1417 let compiled = Compiler::new(test_registry())
1418 .compile(&bp)
1419 .expect("compile");
1420 let spawner: Arc<dyn SpawnerAdapter> = compiled.router.clone();
1421
1422 let engine = Engine::new(EngineCfg::default());
1423 let op_token = seeded_op_token(&engine).await;
1424 let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1425 assert!(matches!(
1426 outcome.expect("dispatch ok"),
1427 DispatchOutcome::Pass(_)
1428 ));
1429
1430 let tail = engine.output_tail(&task_id, 1).await;
1431 assert!(
1432 !tail
1433 .iter()
1434 .any(|ev| matches!(ev, OutputEvent::Artifact { .. })),
1435 "no audits declared must never append any audit artifact"
1436 );
1437 }
1438
1439 #[tokio::test]
1442 async fn steps_filter_respected() {
1443 let agents = vec![
1444 rustfn_agent("a", "worker"),
1445 rustfn_agent("b", "worker"),
1446 rustfn_agent("auditor", "auditor"),
1447 ];
1448 let audits = vec![AuditDef {
1449 agent: "auditor".to_string(),
1450 steps: Some(vec!["b".to_string()]),
1451 mode: AuditMode::Sync,
1452 }];
1453 let bp = minimal_bp(agents, audits.clone());
1454 let compiled = Compiler::new(test_registry())
1455 .compile(&bp)
1456 .expect("compile");
1457 let spawner: Arc<dyn SpawnerAdapter> =
1458 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1459 .wrap(compiled.router.clone());
1460
1461 let engine = Engine::new(EngineCfg::default());
1462 let op_token = seeded_op_token(&engine).await;
1463
1464 let (task_a, outcome_a) = run_step(&engine, &op_token, "a", &spawner).await;
1465 outcome_a.expect("dispatch a ok");
1466 let (task_b, outcome_b) = run_step(&engine, &op_token, "b", &spawner).await;
1467 outcome_b.expect("dispatch b ok");
1468
1469 let tail_a = engine.output_tail(&task_a, 1).await;
1470 assert!(
1471 find_artifact(&tail_a, "audit:a").is_none(),
1472 "step 'a' is not listed in AuditDef.steps and must not be audited"
1473 );
1474 let tail_b = engine.output_tail(&task_b, 1).await;
1475 assert!(
1476 find_artifact(&tail_b, "audit:b").is_some(),
1477 "step 'b' is listed in AuditDef.steps and must be audited"
1478 );
1479 }
1480
1481 #[tokio::test]
1486 async fn auditor_not_audited() {
1487 let agents = vec![
1488 rustfn_agent("worker", "worker"),
1489 rustfn_agent("auditor", "auditor"),
1490 ];
1491 let audits = vec![AuditDef {
1492 agent: "auditor".to_string(),
1493 steps: None,
1494 mode: AuditMode::Sync,
1495 }];
1496 let bp = minimal_bp(agents, audits.clone());
1497 let compiled = Compiler::new(test_registry())
1498 .compile(&bp)
1499 .expect("compile");
1500 let spawner: Arc<dyn SpawnerAdapter> =
1501 AfterRunAuditMiddleware::new(audits, compiled.router.clone())
1502 .wrap(compiled.router.clone());
1503
1504 let engine = Engine::new(EngineCfg::default());
1505 let op_token = seeded_op_token(&engine).await;
1506
1507 let (worker_task, worker_outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
1509 worker_outcome.expect("dispatch worker ok");
1510 let worker_tail = engine.output_tail(&worker_task, 1).await;
1511 assert!(find_artifact(&worker_tail, "audit:worker").is_some());
1512
1513 let (auditor_task, auditor_outcome) =
1516 run_step(&engine, &op_token, "auditor", &spawner).await;
1517 auditor_outcome.expect("dispatch auditor ok");
1518 let auditor_tail = engine.output_tail(&auditor_task, 1).await;
1519 assert!(
1520 find_artifact(&auditor_tail, "audit:auditor").is_none(),
1521 "an agent declared as an auditor must never audit itself"
1522 );
1523 }
1524}