use crate::core::config::CheckPolicy;
use crate::core::engine::Engine;
use crate::core::projection_placement::ProjectionPlacement;
use crate::core::state::{wrap_skip_marker, DispatchOutcome, TaskSpec};
use crate::core::step_naming::StepNaming;
use crate::store::run::{LastFailure, RunContext, StepEntry};
use crate::types::{now_unix, CapToken};
use crate::worker::adapter::SpawnerAdapter;
use async_trait::async_trait;
pub mod compiler;
pub mod loader;
pub mod store;
use mlua_flow_ir::{AsyncDispatcher, EvalError};
use serde_json::{Map, Value};
use std::collections::HashMap;
use std::sync::Arc;
pub use mlua_swarm_schema::OperatorKind as SchemaOperatorKind;
pub use mlua_swarm_schema::{
current_schema_version, default_global_agent_kind, resolve_bound_agents,
resolve_bound_agents_strict, resolve_runner, AgentDef, AgentKind, AgentMeta, AgentProfile,
AgentProviderCapability, AgentProviderManifest, AuditDef, AuditMode, BindOutcome, BindReceipt,
BindRequest, BindingAttestation, BindingBackend, BindingDigest, BindingDigestParseError,
Blueprint, BlueprintMetadata, BlueprintOrigin, BoundAgent, BoundAgentResolveError,
CompilerHints, CompilerStrategy, MetaDef, OperatorDef, ProjectionPlacementSpec, Runner,
RunnerDef, RunnerResolutionSource, RunnerResolveError, SpawnerHints, SubprocessDef,
SubprocessOutput, SubprocessOverrides, WorkerModel, CURRENT_SCHEMA_VERSION,
};
pub struct EngineDispatcher {
engine: Engine,
op_token: CapToken,
spawner: Arc<dyn SpawnerAdapter>,
run_ctx: Option<RunContext>,
step_metas: HashMap<String, Value>,
step_naming: Option<Arc<StepNaming>>,
projection_placement: Option<Arc<ProjectionPlacement>>,
binding_digests: HashMap<String, BindingDigest>,
check_policy: Option<CheckPolicy>,
}
impl EngineDispatcher {
pub fn with_spawner(
engine: Engine,
op_token: CapToken,
spawner: Arc<dyn SpawnerAdapter>,
) -> Self {
Self {
engine,
op_token,
spawner,
run_ctx: None,
step_metas: HashMap::new(),
step_naming: None,
projection_placement: None,
binding_digests: HashMap::new(),
check_policy: None,
}
}
pub fn with_run(mut self, run_ctx: RunContext) -> Self {
self.run_ctx = Some(run_ctx);
self
}
pub fn with_step_metas(mut self, step_metas: HashMap<String, Value>) -> Self {
self.step_metas = step_metas;
self
}
pub fn with_binding_digests(mut self, binding_digests: HashMap<String, BindingDigest>) -> Self {
self.binding_digests = binding_digests;
self
}
pub fn with_step_naming(mut self, step_naming: Arc<StepNaming>) -> Self {
self.step_naming = Some(step_naming);
self
}
pub fn with_projection_placement(
mut self,
projection_placement: Arc<ProjectionPlacement>,
) -> Self {
self.projection_placement = Some(projection_placement);
self
}
pub fn with_check_policy(mut self, check_policy: Option<CheckPolicy>) -> Self {
self.check_policy = check_policy;
self
}
}
fn resolve_step_envelope(
step_metas: &HashMap<String, Value>,
ref_: &str,
input: Value,
) -> Result<(Value, Option<Value>), EvalError> {
let mut obj = match input {
Value::Object(obj) => obj,
other => return Ok((other, None)),
};
let Some(envelope) = obj.remove("$step_meta") else {
return Ok((Value::Object(obj), None));
};
let envelope = match envelope {
Value::Object(map) => map,
other => {
return Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!(
"malformed $step_meta envelope for step '{ref_}': expected an object, got {other}"
),
});
}
};
let ref_ctx: Option<Map<String, Value>> = match envelope.get("ref") {
None | Some(Value::Null) => None,
Some(Value::String(name)) => {
let resolved = step_metas.get(name).cloned().ok_or_else(|| {
EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!(
"$step_meta.ref '{name}' (step '{ref_}') is not a defined Blueprint.metas entry (defined: {:?})",
step_metas.keys().collect::<Vec<_>>()
),
}
})?;
match resolved {
Value::Object(map) => Some(map),
other => {
return Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!(
"malformed $step_meta: MetaDef '{name}'.ctx must be an object, got {other}"
),
});
}
}
}
Some(other) => {
return Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!(
"malformed $step_meta.ref (step '{ref_}'): expected a string, got {other}"
),
});
}
};
let inline: Option<Map<String, Value>> = match envelope.get("inline") {
None | Some(Value::Null) => None,
Some(Value::Object(map)) => Some(map.clone()),
Some(other) => {
return Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!(
"malformed $step_meta.inline (step '{ref_}'): expected an object, got {other}"
),
});
}
};
let step_ctx = match (ref_ctx, inline) {
(None, None) => None,
(Some(base), None) => Some(Value::Object(base)),
(None, Some(inline)) => Some(Value::Object(inline)),
(Some(mut base), Some(inline)) => {
for (k, v) in inline {
base.insert(k, v);
}
Some(Value::Object(base))
}
};
let initial_directive = if let Some(in_value) = obj.remove("$in") {
in_value
} else if obj.is_empty() {
Value::String(String::new())
} else {
Value::Object(obj)
};
Ok((initial_directive, step_ctx))
}
#[async_trait]
impl AsyncDispatcher for EngineDispatcher {
async fn dispatch(&self, ref_: &str, input: Value) -> Result<Value, EvalError> {
let (initial_directive, step_ctx) = resolve_step_envelope(&self.step_metas, ref_, input)?;
let tid = self
.engine
.start_task(
&self.op_token,
TaskSpec {
agent: ref_.to_string(),
initial_directive,
step_ctx,
check_policy: self.check_policy,
},
)
.await
.map_err(|e| EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("start_task: {e}"),
})?;
if let Some(step_naming) = self.step_naming.clone() {
let tid_for_naming = tid.clone();
if let Err(e) = self
.engine
.with_state("EngineDispatcher::dispatch.step_naming", move |s| {
s.step_namings.insert(tid_for_naming, step_naming);
})
.await
{
tracing::warn!(
task_id = %tid,
error = %e,
"EngineDispatcher::dispatch: failed to snapshot StepNaming into EngineState"
);
}
}
if let Some(projection_placement) = self.projection_placement.clone() {
let tid_for_placement = tid.clone();
if let Err(e) = self
.engine
.with_state(
"EngineDispatcher::dispatch.projection_placement",
move |s| {
s.projection_placements
.insert(tid_for_placement, projection_placement);
},
)
.await
{
tracing::warn!(
task_id = %tid,
error = %e,
"EngineDispatcher::dispatch: failed to snapshot ProjectionPlacement into EngineState"
);
}
}
let trace = self.run_ctx.as_ref().and_then(|rc| rc.trace.clone());
let started_at_ms = crate::store::trace::now_unix_ms();
let started = std::time::Instant::now();
if let Some(handle) = &trace {
self.engine
.set_trace_handle(&tid, Some(handle.clone()))
.await;
handle
.append(
crate::store::trace::kind::STEP_DISPATCHED,
Some(ref_),
None,
serde_json::json!({ "step_id": tid.to_string() }),
)
.await;
}
let outcome = self
.engine
.dispatch_attempt_with_run_ctx(
&self.op_token,
&tid,
&self.spawner,
self.run_ctx.as_ref(),
)
.await;
let completed_at_ms = crate::store::trace::now_unix_ms();
let duration_ms = started.elapsed().as_millis() as u64;
let worker_stats = self.engine.take_worker_stats(&tid).await;
if trace.is_some() {
self.engine.set_trace_handle(&tid, None).await;
}
if let Some(rc) = &self.run_ctx {
let status = match &outcome {
Ok(DispatchOutcome::Pass(_)) => "passed",
Ok(DispatchOutcome::Blocked(_)) => "blocked",
Ok(DispatchOutcome::Skip(_)) => "skipped",
Ok(DispatchOutcome::Suspended(_)) => "suspended",
Ok(DispatchOutcome::Cancelled) => "cancelled",
Ok(DispatchOutcome::Timeout) => "timeout",
Err(_) => "failed",
};
let mut entry = StepEntry::basic(
tid.clone(),
Some(ref_.to_string()),
Some(status.to_string()),
self.binding_digests.get(ref_).cloned(),
now_unix(),
);
entry.started_at_ms = Some(started_at_ms);
entry.completed_at_ms = Some(completed_at_ms);
entry.duration_ms = Some(duration_ms);
let attempt = if let Some((a, _)) = &worker_stats {
Some(*a)
} else {
self.engine.task_attempt(&tid).await.ok()
};
entry.attempt = attempt;
if let Some((_, stats)) = worker_stats {
entry = entry.with_worker_stats(stats);
}
if let Err(e) = rc.run_store.append_step_entry(&rc.run_id, entry).await {
tracing::warn!(
run_id = %rc.run_id,
step_id = %tid,
error = %e,
"EngineDispatcher::dispatch: append_step_entry failed"
);
}
if let Some(handle) = &trace {
handle
.append(
crate::store::trace::kind::STEP_COMPLETED,
Some(ref_),
attempt,
serde_json::json!({
"status": status,
"duration_ms": duration_ms,
}),
)
.await;
}
}
match outcome {
Ok(DispatchOutcome::Pass(v)) => Ok(v),
Ok(DispatchOutcome::Skip(v)) => Ok(wrap_skip_marker(v)),
Ok(DispatchOutcome::Blocked(v)) => {
if let Some(rc) = &self.run_ctx {
rc.set_last_failure(LastFailure {
step_id: tid.clone(),
step_ref: Some(ref_.to_string()),
verdict_value: v.clone(),
});
}
Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("blocked: {v}"),
})
}
Ok(other) => Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("non-terminal outcome: {:?}", other),
}),
Err(e) => Err(EvalError::DispatcherError {
ref_: ref_.to_string(),
msg: format!("dispatch_attempt: {e}"),
}),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn metas(pairs: &[(&str, Value)]) -> HashMap<String, Value> {
pairs
.iter()
.map(|(k, v)| (k.to_string(), v.clone()))
.collect()
}
#[test]
fn no_envelope_string_input_passes_through_unchanged() {
let (directive, step_ctx) =
resolve_step_envelope(&HashMap::new(), "scout", json!("plain string")).unwrap();
assert_eq!(directive, json!("plain string"));
assert_eq!(step_ctx, None);
}
#[test]
fn no_envelope_plain_object_input_passes_through_unchanged() {
let input = json!({ "foo": "bar" });
let (directive, step_ctx) =
resolve_step_envelope(&HashMap::new(), "scout", input.clone()).unwrap();
assert_eq!(directive, input);
assert_eq!(step_ctx, None);
}
#[test]
fn envelope_with_only_ref_resolves_that_metadef_ctx() {
let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
let input = json!({ "$step_meta": { "ref": "heavy-scan" }, "$in": "go" });
let (directive, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
assert_eq!(directive, json!("go"));
assert_eq!(step_ctx, Some(json!({ "work_dir": "/x" })));
}
#[test]
fn envelope_with_only_inline_uses_inline_verbatim() {
let input = json!({
"$step_meta": { "inline": { "work_dir": "/inline-only" } },
"$in": "go"
});
let (directive, step_ctx) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
assert_eq!(directive, json!("go"));
assert_eq!(step_ctx, Some(json!({ "work_dir": "/inline-only" })));
}
#[test]
fn inline_wins_over_ref_on_key_collision() {
let step_metas = metas(&[(
"heavy-scan",
json!({ "work_dir": "/ref", "extra": "from-ref" }),
)]);
let input = json!({
"$step_meta": {
"ref": "heavy-scan",
"inline": { "work_dir": "/inline-wins" }
},
"$in": "go"
});
let (_, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
assert_eq!(
step_ctx,
Some(json!({ "work_dir": "/inline-wins", "extra": "from-ref" })),
"inline must win the collided key while ref-only keys survive the merge"
);
}
#[test]
fn dollar_in_rule_extracts_directive_and_ignores_other_sibling_keys() {
let input = json!({
"$step_meta": { "inline": { "k": "v" } },
"$in": "the real directive",
"unrelated_sibling": "ignored"
});
let (directive, step_ctx) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
assert_eq!(directive, json!("the real directive"));
assert_eq!(step_ctx, Some(json!({ "k": "v" })));
}
#[test]
fn no_dollar_in_remainder_becomes_the_directive() {
let input = json!({
"$step_meta": { "inline": { "k": "v" } },
"other_key": "other_value"
});
let (directive, _) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
assert_eq!(directive, json!({ "other_key": "other_value" }));
}
#[test]
fn empty_remainder_becomes_empty_string_directive() {
let input = json!({ "$step_meta": { "ref": "heavy-scan" } });
let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
let (directive, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
assert_eq!(directive, Value::String(String::new()));
assert_eq!(step_ctx, Some(json!({ "work_dir": "/x" })));
}
#[test]
fn unresolved_ref_is_a_loud_dispatcher_error_naming_ref_and_defined() {
let step_metas = metas(&[("known", json!({}))]);
let input = json!({ "$step_meta": { "ref": "unknown" }, "$in": "go" });
let err = resolve_step_envelope(&step_metas, "scout", input).unwrap_err();
match err {
EvalError::DispatcherError { ref_, msg } => {
assert_eq!(ref_, "scout");
assert!(
msg.contains("unknown"),
"message must name the unresolved ref: {msg}"
);
assert!(
msg.contains("known"),
"message must list defined names: {msg}"
);
}
other => panic!("expected DispatcherError, got {other:?}"),
}
}
#[test]
fn malformed_step_meta_not_an_object_is_a_loud_error() {
let input = json!({ "$step_meta": "not-an-object" });
let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
assert!(matches!(err, EvalError::DispatcherError { .. }));
}
#[test]
fn malformed_ref_non_string_is_a_loud_error() {
let input = json!({ "$step_meta": { "ref": 42 } });
let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
assert!(matches!(err, EvalError::DispatcherError { .. }));
}
#[test]
fn malformed_inline_non_object_is_a_loud_error() {
let input = json!({ "$step_meta": { "inline": "not-an-object" } });
let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
assert!(matches!(err, EvalError::DispatcherError { .. }));
}
#[test]
fn ref_resolved_metadef_ctx_non_object_is_a_loud_error() {
let step_metas = metas(&[("bad", json!("not-an-object"))]);
let input = json!({ "$step_meta": { "ref": "bad" } });
let err = resolve_step_envelope(&step_metas, "scout", input).unwrap_err();
assert!(matches!(err, EvalError::DispatcherError { .. }));
}
#[tokio::test]
async fn dispatch_step_meta_envelope_never_leaks_into_stored_prompt() {
use crate::blueprint::compiler::{RustFnInProcessSpawnerFactory, SpawnerFactory};
use crate::core::config::EngineCfg;
use crate::types::{Role, StepId};
use crate::worker::adapter::WorkerResult;
use std::sync::Mutex as StdMutex;
use std::time::Duration;
let captured_tid: Arc<StdMutex<Option<StepId>>> = Arc::new(StdMutex::new(None));
let captured_tid_for_fn = captured_tid.clone();
let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", move |inv| {
let captured_tid = captured_tid_for_fn.clone();
async move {
*captured_tid.lock().unwrap() = Some(inv.task_id.clone());
Ok(WorkerResult {
value: json!({ "ok": true }),
ok: true,
stats: None,
})
}
});
let def = AgentDef {
name: "scout".into(),
kind: AgentKind::RustFn,
spec: json!({ "fn_id": "echo" }),
profile: None,
meta: None,
runner: None,
runner_ref: None,
verdict: None,
};
let spawner = factory.build(&def, None).expect("build");
let engine = Engine::new(EngineCfg::default());
let token = engine
.attach("ut-op", Role::Operator, Duration::from_secs(30))
.await
.expect("attach");
let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
let dispatcher = EngineDispatcher::with_spawner(engine.clone(), token, spawner)
.with_step_metas(step_metas);
let input = json!({
"$step_meta": { "ref": "heavy-scan" },
"$in": "do the thing"
});
let out = dispatcher
.dispatch("scout", input)
.await
.expect("dispatch ok");
assert_eq!(out, json!({ "ok": true }));
let tid = captured_tid
.lock()
.unwrap()
.clone()
.expect("task_id captured");
let stored_prompt = engine
.with_state("test.read_prompt", move |s| {
s.prompts.get(&(tid, 1)).cloned()
})
.await
.expect("with_state")
.expect("prompt recorded for attempt 1");
assert_eq!(
stored_prompt,
json!("do the thing"),
"the stored prompt must be the post-envelope directive, with no $step_meta leakage"
);
}
#[tokio::test]
async fn dispatcher_blocked_records_last_failure_breadcrumb() {
use crate::blueprint::compiler::{RustFnInProcessSpawnerFactory, SpawnerFactory};
use crate::core::config::EngineCfg;
use crate::store::run::{InMemoryRunStore, RunContext, RunRecord, RunStatus, RunStore};
use crate::types::{Role, RunId, TaskId};
use crate::worker::adapter::WorkerResult;
use std::time::Duration;
let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", |_inv| async move {
Ok(WorkerResult {
value: json!({ "verdict": "BLOCKED", "reason": "not-applicable" }),
ok: false,
stats: None,
})
});
let def = AgentDef {
name: "gate".into(),
kind: AgentKind::RustFn,
spec: json!({ "fn_id": "echo" }),
profile: None,
meta: None,
runner: None,
runner_ref: None,
verdict: None,
};
let spawner = factory.build(&def, None).expect("build");
let engine = Engine::new(EngineCfg::default());
let token = engine
.attach("ut-op", Role::Operator, Duration::from_secs(30))
.await
.expect("attach");
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let run_id = RunId::new();
run_store
.create(RunRecord {
id: run_id.clone(),
task_id: TaskId::new(),
status: RunStatus::Running,
step_entries: Vec::new(),
degradations: Vec::new(),
operator_sid: None,
result_ref: None,
input_json: Some("{}".to_string()),
created_at: 0,
updated_at: 0,
})
.await
.expect("seed RunRecord");
let run_ctx = RunContext::new(run_id, run_store);
let dispatcher =
EngineDispatcher::with_spawner(engine, token, spawner).with_run(run_ctx.clone());
let err = dispatcher
.dispatch("gate", json!("go"))
.await
.expect_err("expected DispatcherError for Blocked outcome");
assert!(
err.to_string().contains("blocked"),
"expected EvalError to mention blocked, got: {err}"
);
let breadcrumb = run_ctx
.last_failure
.lock()
.expect("last_failure mutex not poisoned")
.clone()
.expect("Blocked arm must have written LastFailure");
assert_eq!(
breadcrumb.step_ref,
Some("gate".to_string()),
"step_ref must be the Blueprint ref this dispatch was routed to"
);
assert_eq!(
breadcrumb.verdict_value,
json!({ "verdict": "BLOCKED", "reason": "not-applicable" }),
"verdict_value must be the exact value the worker returned"
);
assert!(!breadcrumb.step_id.to_string().is_empty());
}
#[tokio::test]
async fn run_context_snapshot_partial_ctx_reconstructs_step_entry_log() {
use crate::store::run::{
InMemoryRunStore, RunContext, RunRecord, RunStatus, RunStore, StepEntry,
};
use crate::types::{now_unix, RunId, StepId, TaskId};
let run_store: Arc<dyn RunStore> = Arc::new(InMemoryRunStore::new());
let run_id = RunId::new();
run_store
.create(RunRecord {
id: run_id.clone(),
task_id: TaskId::new(),
status: RunStatus::Running,
step_entries: Vec::new(),
degradations: Vec::new(),
operator_sid: None,
result_ref: None,
input_json: Some("{}".to_string()),
created_at: 0,
updated_at: 0,
})
.await
.expect("seed RunRecord");
let sid1 = StepId::new();
let sid2 = StepId::new();
run_store
.append_step_entry(
&run_id,
StepEntry::basic(
sid1.clone(),
Some("stage-1".to_string()),
Some("passed".to_string()),
None,
now_unix(),
),
)
.await
.expect("append 1");
run_store
.append_step_entry(
&run_id,
StepEntry::basic(
sid2.clone(),
Some("stage-2".to_string()),
Some("blocked".to_string()),
None,
now_unix(),
),
)
.await
.expect("append 2");
let run_ctx = RunContext::new(run_id, run_store);
let snap = run_ctx.snapshot_partial_ctx().await;
let steps = snap
.get("steps")
.and_then(|v| v.as_object())
.expect("steps object");
assert_eq!(steps.len(), 2);
assert_eq!(steps[&sid1.to_string()]["step_ref"], json!("stage-1"));
assert_eq!(steps[&sid1.to_string()]["status"], json!("passed"));
assert_eq!(steps[&sid2.to_string()]["step_ref"], json!("stage-2"));
assert_eq!(steps[&sid2.to_string()]["status"], json!("blocked"));
}
}