pub mod agent_context;
pub mod input_inject;
pub mod lua_layer;
pub mod project_name_alias;
pub mod resolver;
pub mod sink;
pub mod task_input;
pub mod worker_binding;
use crate::blueprint::compiler::CompiledAgentTable;
use crate::blueprint::{AuditDef, AuditMode};
use crate::core::ctx::{Ctx, OperatorKind};
use crate::core::engine::Engine;
use crate::core::state::{DispatchOutcome, Event, TaskSpec};
use crate::types::{CapToken, StepId};
use crate::worker::adapter::{SpawnError, SpawnerAdapter};
use crate::worker::output::{ContentRef, OutputEvent};
use crate::worker::{wrap_join, MiddlewareWorker, Worker, WorkerJoinHandler};
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::broadcast;
async fn pull_final_value_ok(
engine: &Engine,
task_id: &StepId,
attempt: u32,
) -> Option<(Value, bool)> {
let tail = engine.output_tail(task_id, attempt).await;
tail.iter().rev().find_map(|ev| match ev {
OutputEvent::Final {
content: ContentRef::Inline { value },
ok,
} => Some((value.clone(), *ok)),
OutputEvent::Final {
content: ContentRef::FileRef { path, .. },
ok,
} => Some((serde_json::json!({"file_ref": path.to_string_lossy()}), *ok)),
_ => None,
})
}
pub trait SpawnerLayer: Send + Sync + 'static {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter>;
}
pub struct SpawnerStack {
inner: Arc<dyn SpawnerAdapter>,
}
impl SpawnerStack {
pub fn new(base: Arc<dyn SpawnerAdapter>) -> Self {
Self { inner: base }
}
pub fn layer<L: SpawnerLayer>(mut self, layer: L) -> Self {
self.inner = layer.wrap(self.inner);
self
}
pub fn layer_dyn(mut self, layer: Arc<dyn SpawnerLayer>) -> Self {
self.inner = layer.wrap(self.inner);
self
}
pub fn build(self) -> Arc<dyn SpawnerAdapter> {
self.inner
}
}
pub type LayerFactory =
Arc<dyn Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static>;
#[derive(Default, Clone)]
pub struct LayerRegistry {
base: Vec<LayerFactory>,
hints: std::collections::HashMap<String, LayerFactory>,
}
impl LayerRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn with_base<F>(mut self, factory: F) -> Self
where
F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
{
self.base.push(Arc::new(factory));
self
}
pub fn with_hint<F>(mut self, key: impl Into<String>, factory: F) -> Self
where
F: Fn(&crate::core::engine::Engine) -> Arc<dyn SpawnerLayer> + Send + Sync + 'static,
{
self.hints.insert(key.into(), Arc::new(factory));
self
}
pub fn base_factories(&self) -> &[LayerFactory] {
&self.base
}
pub fn lookup_hint(&self, key: &str) -> Option<&LayerFactory> {
self.hints.get(key)
}
}
pub struct AuditMiddleware {
pub event_tx: broadcast::Sender<Event>,
}
impl AuditMiddleware {
pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
Self { event_tx }
}
}
impl SpawnerLayer for AuditMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(AuditWrapped {
inner,
event_tx: self.event_tx.clone(),
})
}
}
struct AuditWrapped {
inner: Arc<dyn SpawnerAdapter>,
event_tx: broadcast::Sender<Event>,
}
#[async_trait]
impl SpawnerAdapter for AuditWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let _ = self.event_tx.send(Event::TaskAttemptStarted {
task_id: task_id.clone(),
attempt,
});
self.inner.spawn(engine, ctx, task_id, attempt, token).await
}
}
pub struct MainAIMiddleware;
impl MainAIMiddleware {
pub fn new() -> Self {
Self
}
}
impl Default for MainAIMiddleware {
fn default() -> Self {
Self::new()
}
}
impl SpawnerLayer for MainAIMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(MainAIWrapped { inner })
}
}
struct MainAIWrapped {
inner: Arc<dyn SpawnerAdapter>,
}
#[async_trait]
impl SpawnerAdapter for MainAIWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let mainai = matches!(
ctx.operator.kind,
OperatorKind::MainAi | OperatorKind::Composite
);
if mainai {
if let Some(hook) = &ctx.operator.spawn_hook {
hook.before(ctx)
.await
.map_err(SpawnError::RejectedByMiddleware)?;
}
}
let handle = self
.inner
.spawn(engine, ctx, task_id.clone(), attempt, token)
.await?;
if !mainai {
return Ok(handle);
}
let Some(hook) = ctx.operator.spawn_hook.clone() else {
return Ok(handle);
};
let ctx_clone = ctx.clone();
let engine_clone = engine.clone();
let task_id_clone = task_id.clone();
Ok(wrap_join(handle, move |signal| {
let hook = hook.clone();
let ctx_clone = ctx_clone.clone();
let engine_clone = engine_clone.clone();
let task_id_clone = task_id_clone.clone();
async move {
let v = match &signal {
Ok(()) => pull_final_value_ok(&engine_clone, &task_id_clone, attempt)
.await
.map(|(v, _)| v)
.unwrap_or(Value::Null),
Err(e) => Value::String(e.to_string()),
};
let _ = hook.after(&ctx_clone, &v).await;
signal
}
}))
}
}
pub struct SeniorEscalationMiddleware;
impl SeniorEscalationMiddleware {
pub fn new() -> Self {
Self
}
}
impl Default for SeniorEscalationMiddleware {
fn default() -> Self {
Self::new()
}
}
impl SpawnerLayer for SeniorEscalationMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(SeniorWrapped { inner })
}
}
struct SeniorWrapped {
inner: Arc<dyn SpawnerAdapter>,
}
#[async_trait]
impl SpawnerAdapter for SeniorWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let bridge = ctx.operator.senior_bridge.clone();
let task_id_for_hook = task_id.clone();
let engine_clone = engine.clone();
let token_clone = token.clone();
let handle = self
.inner
.spawn(engine, ctx, task_id, attempt, token)
.await?;
let Some(bridge) = bridge else {
return Ok(handle);
};
Ok(wrap_join(handle, move |signal| {
let bridge = bridge.clone();
let task_id = task_id_for_hook.clone();
let engine = engine_clone.clone();
let token = token_clone.clone();
async move {
signal?;
let last = pull_final_value_ok(&engine, &task_id, attempt).await;
if let Some((value, false)) = last {
let question = serde_json::json!({
"reason": "worker reported ok=false",
"value": value.clone(),
});
if let Ok(answer) = bridge.ask(&task_id, question).await {
let override_val = serde_json::json!({
"original": value,
"senior_answer": answer,
});
let _ = engine
.submit_output(
&token,
&task_id,
attempt,
OutputEvent::Final {
content: ContentRef::Inline {
value: override_val,
},
ok: true,
},
)
.await;
}
}
Ok(())
}
}))
}
}
pub struct OperatorDelegateMiddleware;
impl OperatorDelegateMiddleware {
pub fn new() -> Self {
Self
}
}
impl Default for OperatorDelegateMiddleware {
fn default() -> Self {
Self::new()
}
}
impl SpawnerLayer for OperatorDelegateMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(OperatorDelegateWrapped { inner })
}
}
struct OperatorDelegateWrapped {
inner: Arc<dyn SpawnerAdapter>,
}
#[async_trait]
impl SpawnerAdapter for OperatorDelegateWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let Some(operator) = ctx.operator.operator.clone() else {
return self.inner.spawn(engine, ctx, task_id, attempt, token).await;
};
let prompt = engine
.fetch_prompt(&token, &task_id)
.await
.map_err(|e| SpawnError::Internal(format!("fetch_prompt: {e}")))?;
let worker: Option<crate::operator::WorkerBinding> = match ctx
.meta
.runtime
.get(crate::middleware::worker_binding::WORKER_BINDING_KEY)
{
Some(v) => Some(serde_json::from_value(v.clone()).map_err(|e| {
SpawnError::Internal(format!(
"ctx.meta.runtime['{}'] for agent '{}' is malformed: {e}",
crate::middleware::worker_binding::WORKER_BINDING_KEY,
ctx.agent
))
})?),
None => None,
};
let engine_clone = engine.clone();
let token_clone = token.clone();
let token_for_op = token.clone();
let task_id_clone = task_id.clone();
let ctx_clone = ctx.clone();
let (tx, rx) = tokio::sync::oneshot::channel();
let cancel = tokio_util::sync::CancellationToken::new();
let cancel_inner = cancel.clone();
let worker_id = crate::types::WorkerId::new();
tracing::debug!(worker_id = %worker_id, step_id = %task_id, "worker spawned (delegate axis)");
tokio::spawn(async move {
let result: Result<
crate::worker::adapter::WorkerResult,
crate::worker::adapter::WorkerError,
> = tokio::select! {
r = operator.execute(&ctx_clone, None, prompt, worker, token_for_op) => r,
_ = cancel_inner.cancelled() => Err(crate::worker::adapter::WorkerError::Cancelled),
};
if let Ok(wr) = &result {
let tail = engine_clone.output_tail(&task_id_clone, attempt).await;
let has_final = tail
.iter()
.any(|ev| matches!(ev, crate::worker::output::OutputEvent::Final { .. }));
if !has_final {
let ev = crate::worker::output::OutputEvent::Final {
content: crate::worker::output::ContentRef::Inline {
value: wr.value.clone(),
},
ok: wr.ok,
};
let _ = engine_clone
.submit_output(&token_clone, &task_id_clone, attempt, ev)
.await;
}
}
let signal: Result<(), crate::worker::adapter::WorkerError> = result.map(|_| ());
let _ = tx.send(signal);
});
Ok(Box::new(MiddlewareWorker {
handler: WorkerJoinHandler {
worker_id,
cancel,
completion: rx,
},
}))
}
}
pub struct LongHoldMiddleware {
pub default_hold: Duration,
pub event_tx: broadcast::Sender<Event>,
}
impl LongHoldMiddleware {
pub fn new(default_hold: Duration, event_tx: broadcast::Sender<Event>) -> Self {
Self {
default_hold,
event_tx,
}
}
}
impl SpawnerLayer for LongHoldMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(LongHoldWrapped {
inner,
default_hold: self.default_hold,
event_tx: self.event_tx.clone(),
})
}
}
struct LongHoldWrapped {
inner: Arc<dyn SpawnerAdapter>,
default_hold: Duration,
event_tx: broadcast::Sender<Event>,
}
#[async_trait]
impl SpawnerAdapter for LongHoldWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let handle = self
.inner
.spawn(engine, ctx, task_id.clone(), attempt, token)
.await?;
let started = Instant::now();
let default_hold = self.default_hold;
let event_tx = self.event_tx.clone();
let task_id_inner = task_id.clone();
Ok(wrap_join(handle, move |signal| {
let elapsed = started.elapsed();
let default_hold = default_hold;
let event_tx = event_tx.clone();
let task_id_inner = task_id_inner.clone();
async move {
if elapsed > default_hold {
let _ = event_tx.send(Event::TaskAttemptCompleted {
task_id: task_id_inner,
attempt,
result: serde_json::json!({
"long_hold_warn": true,
"elapsed_ms": elapsed.as_millis() as u64,
"default_hold_ms": default_hold.as_millis() as u64,
}),
});
}
signal
}
}))
}
}
const AUDIT_INSTRUCTION: &str = "Inspect this step's transcript/output for degradations, tool \
failures, or silent fallbacks, and emit your findings as a structured JSON object in your \
final output.";
pub struct AfterRunAuditMiddleware {
defs: Vec<AuditDef>,
router: Arc<CompiledAgentTable>,
}
impl AfterRunAuditMiddleware {
pub fn new(defs: Vec<AuditDef>, router: Arc<CompiledAgentTable>) -> Self {
Self { defs, router }
}
}
impl SpawnerLayer for AfterRunAuditMiddleware {
fn wrap(&self, inner: Arc<dyn SpawnerAdapter>) -> Arc<dyn SpawnerAdapter> {
Arc::new(AfterRunAuditWrapped {
inner,
defs: self.defs.clone(),
router: self.router.clone(),
})
}
}
struct AfterRunAuditWrapped {
inner: Arc<dyn SpawnerAdapter>,
defs: Vec<AuditDef>,
router: Arc<CompiledAgentTable>,
}
fn audit_def_matches_step(def: &AuditDef, step_ref: &str) -> bool {
match &def.steps {
None => true,
Some(list) => list.iter().any(|s| s == "*" || s == step_ref),
}
}
#[allow(clippy::too_many_arguments)]
async fn run_one_audit(
engine: &Engine,
router: &Arc<CompiledAgentTable>,
token: &CapToken,
audited_task_id: &StepId,
attempt: u32,
step_ref: &str,
audit_agent: &str,
directive: Value,
) {
let spec = TaskSpec {
agent: audit_agent.to_string(),
initial_directive: directive,
step_ctx: None,
};
let audit_task_id = match engine.start_task(token, spec).await {
Ok(tid) => tid,
Err(e) => {
tracing::warn!(
audited_task_id = %audited_task_id,
step_ref,
audit_agent,
error = %e,
"AfterRunAuditMiddleware: start_task failed for audit agent; \
audited step's outcome is unaffected"
);
return;
}
};
let spawner: Arc<dyn SpawnerAdapter> = router.clone();
let findings = match engine
.dispatch_attempt_with(token, &audit_task_id, &spawner, None)
.await
{
Ok(DispatchOutcome::Pass(v)) | Ok(DispatchOutcome::Blocked(v)) => v,
Ok(other) => {
tracing::warn!(
audited_task_id = %audited_task_id,
step_ref,
audit_agent,
outcome = ?other,
"AfterRunAuditMiddleware: audit agent did not settle (Pass/Blocked); \
audited step's outcome is unaffected"
);
return;
}
Err(e) => {
tracing::warn!(
audited_task_id = %audited_task_id,
step_ref,
audit_agent,
error = %e,
"AfterRunAuditMiddleware: dispatch_attempt_with failed for audit agent; \
audited step's outcome is unaffected"
);
return;
}
};
if let Err(e) = engine
.submit_output(
token,
audited_task_id,
attempt,
OutputEvent::Artifact {
name: format!("audit:{step_ref}"),
content: ContentRef::Inline { value: findings },
},
)
.await
{
tracing::warn!(
audited_task_id = %audited_task_id,
step_ref,
audit_agent,
error = %e,
"AfterRunAuditMiddleware: submit_output failed for audit findings; \
audited step's outcome is unaffected"
);
}
}
#[async_trait]
impl SpawnerAdapter for AfterRunAuditWrapped {
async fn spawn(
&self,
engine: &Engine,
ctx: &Ctx,
task_id: StepId,
attempt: u32,
token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
let step_ref = ctx.agent.clone();
let handle = self
.inner
.spawn(engine, ctx, task_id.clone(), attempt, token.clone())
.await?;
let is_auditor = self.defs.iter().any(|d| d.agent == step_ref);
let matched: Vec<AuditDef> = if is_auditor {
Vec::new()
} else {
self.defs
.iter()
.filter(|d| audit_def_matches_step(d, &step_ref))
.cloned()
.collect()
};
if matched.is_empty() {
return Ok(handle);
}
let engine = engine.clone();
let router = self.router.clone();
Ok(wrap_join(handle, move |signal| async move {
signal?;
let (final_value, ok) = pull_final_value_ok(&engine, &task_id, attempt)
.await
.unwrap_or((Value::Null, true));
for def in matched {
let directive = serde_json::json!({
"kind": "after_run_audit",
"task_id": task_id.to_string(),
"step_ref": step_ref.clone(),
"attempt": attempt,
"ok": ok,
"final_value": final_value.clone(),
"instruction": AUDIT_INSTRUCTION,
});
match def.mode {
AuditMode::Sync => {
run_one_audit(
&engine, &router, &token, &task_id, attempt, &step_ref, &def.agent,
directive,
)
.await;
}
AuditMode::Async => {
let engine = engine.clone();
let router = router.clone();
let token = token.clone();
let task_id = task_id.clone();
let step_ref = step_ref.clone();
let agent = def.agent.clone();
tokio::spawn(async move {
run_one_audit(
&engine, &router, &token, &task_id, attempt, &step_ref, &agent,
directive,
)
.await;
});
}
}
}
Ok(())
}))
}
}
#[cfg(test)]
mod operator_delegate_worker_binding_tests {
use super::*;
use crate::core::config::EngineCfg;
use crate::core::state::TaskSpec;
use crate::operator::WorkerBinding;
use crate::types::Role;
use crate::worker::adapter::{WorkerError, WorkerResult};
use std::sync::Mutex;
struct RecordingOperator {
seen: Arc<Mutex<Option<Option<WorkerBinding>>>>,
}
#[async_trait]
impl crate::operator::Operator for RecordingOperator {
async fn execute(
&self,
_ctx: &Ctx,
_system: Option<String>,
_prompt: Value,
worker: Option<WorkerBinding>,
_worker_token: CapToken,
) -> Result<WorkerResult, WorkerError> {
*self.seen.lock().unwrap() = Some(worker);
Ok(WorkerResult {
value: Value::Null,
ok: true,
})
}
}
struct MustNotSpawn;
#[async_trait]
impl SpawnerAdapter for MustNotSpawn {
async fn spawn(
&self,
_engine: &Engine,
_ctx: &Ctx,
_task_id: StepId,
_attempt: u32,
_token: CapToken,
) -> Result<Box<dyn Worker>, SpawnError> {
panic!("delegate axis must bypass inner.spawn when an operator is attached");
}
}
async fn seeded_engine() -> (Engine, CapToken, StepId) {
let engine = Engine::new(EngineCfg::default());
let op_token = engine
.attach("ut-op", Role::Operator, Duration::from_secs(30))
.await
.expect("attach");
let task_id = engine
.start_task(
&op_token,
TaskSpec {
agent: "planner".to_string(),
initial_directive: "do the thing".into(),
step_ctx: None,
},
)
.await
.expect("start_task");
let worker_token = engine.signer().session(
format!("worker-of-{task_id}"),
Role::Worker,
vec!["*".into()],
Duration::from_secs(600),
);
let fp = worker_token.fingerprint();
let record = crate::core::state::CapTokenRecord::from_worker_token(
worker_token.clone(),
task_id.clone(),
);
engine
.with_state("test.mint_worker", move |s| {
s.tokens.insert(fp, record);
})
.await
.expect("mint worker token");
(engine, worker_token, task_id)
}
fn delegate_stack() -> Arc<dyn SpawnerAdapter> {
OperatorDelegateMiddleware::new().wrap(Arc::new(MustNotSpawn))
}
async fn recorded_worker(
seen: &Arc<Mutex<Option<Option<WorkerBinding>>>>,
) -> Option<WorkerBinding> {
for _ in 0..100 {
if let Some(w) = seen.lock().unwrap().clone() {
return w;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("operator.execute was never called within 1s");
}
#[tokio::test]
async fn forwards_ctx_injected_binding_to_operator_execute() {
let (engine, token, task_id) = seeded_engine().await;
let seen = Arc::new(Mutex::new(None));
let op = Arc::new(RecordingOperator { seen: seen.clone() });
let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
ctx.operator.operator = Some(op);
ctx.meta.runtime.insert(
crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
serde_json::to_value(WorkerBinding {
variant: "mse-worker-coder".to_string(),
tools: vec!["Edit".to_string()],
})
.unwrap(),
);
let _worker = delegate_stack()
.spawn(&engine, &ctx, task_id, 1, token)
.await
.expect("delegate spawn ok");
let got = recorded_worker(&seen).await.expect("binding forwarded");
assert_eq!(got.variant, "mse-worker-coder");
assert_eq!(got.tools, vec!["Edit".to_string()]);
}
#[tokio::test]
async fn absent_binding_stays_none_no_silent_default() {
let (engine, token, task_id) = seeded_engine().await;
let seen = Arc::new(Mutex::new(None));
let op = Arc::new(RecordingOperator { seen: seen.clone() });
let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
ctx.operator.operator = Some(op);
let _worker = delegate_stack()
.spawn(&engine, &ctx, task_id, 1, token)
.await
.expect("delegate spawn ok");
assert!(
recorded_worker(&seen).await.is_none(),
"no binding declared must reach the operator as None (fail-loud stays downstream)"
);
}
#[tokio::test]
async fn malformed_binding_fails_loud_before_execute() {
let (engine, token, task_id) = seeded_engine().await;
let seen = Arc::new(Mutex::new(None));
let op = Arc::new(RecordingOperator { seen: seen.clone() });
let mut ctx = Ctx::new(task_id.clone(), 1, "planner");
ctx.operator.operator = Some(op);
ctx.meta.runtime.insert(
crate::middleware::worker_binding::WORKER_BINDING_KEY.to_string(),
serde_json::json!({ "not_a_binding": true }),
);
let err = match delegate_stack()
.spawn(&engine, &ctx, task_id, 1, token)
.await
{
Ok(_) => panic!("malformed binding must fail the spawn"),
Err(e) => e,
};
let msg = format!("{err:?}");
assert!(
msg.contains("worker_binding") && msg.contains("malformed"),
"error must name the malformed key: {msg}"
);
assert!(
seen.lock().unwrap().is_none(),
"operator.execute must not run on malformed binding"
);
}
}
#[cfg(test)]
mod after_run_audit_tests {
use super::*;
use crate::blueprint::compiler::{Compiler, RustFnInProcessSpawnerFactory, SpawnerRegistry};
use crate::blueprint::{
current_schema_version, AgentDef, AgentKind, Blueprint, BlueprintMetadata, CompilerHints,
CompilerStrategy,
};
use crate::core::config::EngineCfg;
use crate::types::Role;
use crate::worker::adapter::{WorkerError as StubWorkerError, WorkerResult};
use mlua_flow_ir::Node as FlowNode;
fn rustfn_agent(name: &str, fn_id: &str) -> AgentDef {
AgentDef {
name: name.to_string(),
kind: AgentKind::RustFn,
spec: serde_json::json!({ "fn_id": fn_id }),
profile: None,
meta: None,
}
}
fn minimal_bp(agents: Vec<AgentDef>, audits: Vec<AuditDef>) -> Blueprint {
crate::blueprint::Blueprint {
schema_version: current_schema_version(),
id: "afterrun-audit-ut".into(),
flow: FlowNode::Seq { children: vec![] },
agents,
operators: vec![],
metas: vec![],
hints: CompilerHints::default(),
strategy: CompilerStrategy::default(),
metadata: BlueprintMetadata::default(),
spawner_hints: Default::default(),
default_agent_kind: AgentKind::Operator,
default_operator_kind: None,
default_init_ctx: None,
default_agent_ctx: None,
default_context_policy: None,
projection_placement: None,
audits,
degradation_policy: None,
}
}
fn test_registry() -> SpawnerRegistry {
let factory = RustFnInProcessSpawnerFactory::new()
.register_fn("worker", |_inv| async move {
Ok(WorkerResult {
value: serde_json::json!({ "result": "done" }),
ok: true,
})
})
.register_fn("auditor", |_inv| async move {
Ok(WorkerResult {
value: serde_json::json!({ "finding": "clean" }),
ok: true,
})
})
.register_fn("bad-auditor", |_inv| async move {
Err(StubWorkerError::Failed("boom".to_string()))
});
let mut reg = SpawnerRegistry::new();
reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
reg
}
async fn run_step(
engine: &Engine,
op_token: &CapToken,
agent_name: &str,
spawner: &Arc<dyn SpawnerAdapter>,
) -> (
StepId,
Result<DispatchOutcome, crate::core::errors::EngineError>,
) {
let task_id = engine
.start_task(
op_token,
TaskSpec {
agent: agent_name.to_string(),
initial_directive: serde_json::json!("go"),
step_ctx: None,
},
)
.await
.expect("start_task");
let outcome = engine
.dispatch_attempt_with(op_token, &task_id, spawner, None)
.await;
(task_id, outcome)
}
async fn seeded_op_token(engine: &Engine) -> CapToken {
engine
.attach("ut-op", Role::Operator, Duration::from_secs(30))
.await
.expect("attach")
}
fn find_artifact(tail: &[OutputEvent], name: &str) -> Option<Value> {
tail.iter().find_map(|ev| match ev {
OutputEvent::Artifact {
name: n,
content: ContentRef::Inline { value },
} if n == name => Some(value.clone()),
_ => None,
})
}
#[tokio::test]
async fn audit_fires_after_step_and_appends_artifact() {
let agents = vec![
rustfn_agent("worker", "worker"),
rustfn_agent("auditor", "auditor"),
];
let audits = vec![AuditDef {
agent: "auditor".to_string(),
steps: None,
mode: AuditMode::Sync,
}];
let bp = minimal_bp(agents, audits.clone());
let compiled = Compiler::new(test_registry())
.compile(&bp)
.expect("compile");
let spawner: Arc<dyn SpawnerAdapter> =
AfterRunAuditMiddleware::new(audits, compiled.router.clone())
.wrap(compiled.router.clone());
let engine = Engine::new(EngineCfg::default());
let op_token = seeded_op_token(&engine).await;
let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
match outcome.expect("dispatch ok") {
DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
other => panic!("expected Pass (the worker's own outcome), got {other:?}"),
}
let tail = engine.output_tail(&task_id, 1).await;
let findings =
find_artifact(&tail, "audit:worker").expect("audit:worker artifact must be appended");
assert_eq!(findings, serde_json::json!({ "finding": "clean" }));
}
#[tokio::test]
async fn audit_failure_never_alters_outcome() {
let agents = vec![
rustfn_agent("worker", "worker"),
rustfn_agent("bad-auditor", "bad-auditor"),
];
let audits = vec![AuditDef {
agent: "bad-auditor".to_string(),
steps: None,
mode: AuditMode::Sync,
}];
let bp = minimal_bp(agents, audits.clone());
let compiled = Compiler::new(test_registry())
.compile(&bp)
.expect("compile");
let spawner: Arc<dyn SpawnerAdapter> =
AfterRunAuditMiddleware::new(audits, compiled.router.clone())
.wrap(compiled.router.clone());
let engine = Engine::new(EngineCfg::default());
let op_token = seeded_op_token(&engine).await;
let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
match outcome.expect("audited step's dispatch must still succeed despite auditor failure") {
DispatchOutcome::Pass(v) => assert_eq!(v, serde_json::json!({ "result": "done" })),
other => panic!("expected Pass identical to a no-audit run, got {other:?}"),
}
let tail = engine.output_tail(&task_id, 1).await;
assert!(
find_artifact(&tail, "audit:worker").is_none(),
"auditor failure must not append an audit artifact"
);
}
#[tokio::test]
async fn no_audit_defs_appends_no_artifact() {
let agents = vec![rustfn_agent("worker", "worker")];
let bp = minimal_bp(agents, vec![]);
let compiled = Compiler::new(test_registry())
.compile(&bp)
.expect("compile");
let spawner: Arc<dyn SpawnerAdapter> = compiled.router.clone();
let engine = Engine::new(EngineCfg::default());
let op_token = seeded_op_token(&engine).await;
let (task_id, outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
assert!(matches!(
outcome.expect("dispatch ok"),
DispatchOutcome::Pass(_)
));
let tail = engine.output_tail(&task_id, 1).await;
assert!(
!tail
.iter()
.any(|ev| matches!(ev, OutputEvent::Artifact { .. })),
"no audits declared must never append any audit artifact"
);
}
#[tokio::test]
async fn steps_filter_respected() {
let agents = vec![
rustfn_agent("a", "worker"),
rustfn_agent("b", "worker"),
rustfn_agent("auditor", "auditor"),
];
let audits = vec![AuditDef {
agent: "auditor".to_string(),
steps: Some(vec!["b".to_string()]),
mode: AuditMode::Sync,
}];
let bp = minimal_bp(agents, audits.clone());
let compiled = Compiler::new(test_registry())
.compile(&bp)
.expect("compile");
let spawner: Arc<dyn SpawnerAdapter> =
AfterRunAuditMiddleware::new(audits, compiled.router.clone())
.wrap(compiled.router.clone());
let engine = Engine::new(EngineCfg::default());
let op_token = seeded_op_token(&engine).await;
let (task_a, outcome_a) = run_step(&engine, &op_token, "a", &spawner).await;
outcome_a.expect("dispatch a ok");
let (task_b, outcome_b) = run_step(&engine, &op_token, "b", &spawner).await;
outcome_b.expect("dispatch b ok");
let tail_a = engine.output_tail(&task_a, 1).await;
assert!(
find_artifact(&tail_a, "audit:a").is_none(),
"step 'a' is not listed in AuditDef.steps and must not be audited"
);
let tail_b = engine.output_tail(&task_b, 1).await;
assert!(
find_artifact(&tail_b, "audit:b").is_some(),
"step 'b' is listed in AuditDef.steps and must be audited"
);
}
#[tokio::test]
async fn auditor_not_audited() {
let agents = vec![
rustfn_agent("worker", "worker"),
rustfn_agent("auditor", "auditor"),
];
let audits = vec![AuditDef {
agent: "auditor".to_string(),
steps: None,
mode: AuditMode::Sync,
}];
let bp = minimal_bp(agents, audits.clone());
let compiled = Compiler::new(test_registry())
.compile(&bp)
.expect("compile");
let spawner: Arc<dyn SpawnerAdapter> =
AfterRunAuditMiddleware::new(audits, compiled.router.clone())
.wrap(compiled.router.clone());
let engine = Engine::new(EngineCfg::default());
let op_token = seeded_op_token(&engine).await;
let (worker_task, worker_outcome) = run_step(&engine, &op_token, "worker", &spawner).await;
worker_outcome.expect("dispatch worker ok");
let worker_tail = engine.output_tail(&worker_task, 1).await;
assert!(find_artifact(&worker_tail, "audit:worker").is_some());
let (auditor_task, auditor_outcome) =
run_step(&engine, &op_token, "auditor", &spawner).await;
auditor_outcome.expect("dispatch auditor ok");
let auditor_tail = engine.output_tail(&auditor_task, 1).await;
assert!(
find_artifact(&auditor_tail, "audit:auditor").is_none(),
"an agent declared as an auditor must never audit itself"
);
}
}