use crate::core::engine::Engine;
use crate::core::projection_placement::ProjectionPlacement;
use crate::core::state::{DispatchOutcome, TaskSpec};
use crate::core::step_naming::StepNaming;
use crate::store::run::{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, AgentDef, AgentKind, AgentMeta,
AgentProfile, AuditDef, AuditMode, Blueprint, BlueprintMetadata, BlueprintOrigin,
CompilerHints, CompilerStrategy, MetaDef, OperatorDef, ProjectionPlacementSpec, SpawnerHints,
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>>,
}
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,
}
}
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_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
}
}
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,
},
)
.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 run_id_for_ctx = self.run_ctx.as_ref().map(|rc| rc.run_id.clone());
let outcome = self
.engine
.dispatch_attempt_with(&self.op_token, &tid, &self.spawner, run_id_for_ctx.as_ref())
.await;
if let Some(rc) = &self.run_ctx {
let status = match &outcome {
Ok(DispatchOutcome::Pass(_)) => "passed",
Ok(DispatchOutcome::Blocked(_)) => "blocked",
Ok(DispatchOutcome::Suspended(_)) => "suspended",
Ok(DispatchOutcome::Cancelled) => "cancelled",
Ok(DispatchOutcome::Timeout) => "timeout",
Err(_) => "failed",
};
let entry = StepEntry {
step_id: tid.clone(),
step_ref: Some(ref_.to_string()),
status: Some(status.to_string()),
at: now_unix(),
};
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"
);
}
}
match outcome {
Ok(DispatchOutcome::Pass(v)) => Ok(v),
Ok(DispatchOutcome::Blocked(v)) => 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,
})
}
});
let def = AgentDef {
name: "scout".into(),
kind: AgentKind::RustFn,
spec: json!({ "fn_id": "echo" }),
profile: None,
meta: 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"
);
}
}