1use crate::core::engine::Engine;
28use crate::core::projection_placement::ProjectionPlacement;
29use crate::core::state::{DispatchOutcome, TaskSpec};
30use crate::core::step_naming::StepNaming;
31use crate::store::run::{RunContext, StepEntry};
32use crate::types::{now_unix, CapToken};
33use crate::worker::adapter::SpawnerAdapter;
34use async_trait::async_trait;
35pub mod compiler;
36pub mod loader;
37pub mod store;
38
39use mlua_flow_ir::{AsyncDispatcher, EvalError};
40use serde_json::{Map, Value};
41use std::collections::HashMap;
42use std::sync::Arc;
43
44pub use mlua_swarm_schema::OperatorKind as SchemaOperatorKind;
51pub use mlua_swarm_schema::{
52 current_schema_version, default_global_agent_kind, AgentDef, AgentKind, AgentMeta,
53 AgentProfile, AuditDef, AuditMode, Blueprint, BlueprintMetadata, BlueprintOrigin,
54 CompilerHints, CompilerStrategy, MetaDef, OperatorDef, ProjectionPlacementSpec, SpawnerHints,
55 CURRENT_SCHEMA_VERSION,
56};
57
58pub struct EngineDispatcher {
97 engine: Engine,
98 op_token: CapToken,
99 spawner: Arc<dyn SpawnerAdapter>,
100 run_ctx: Option<RunContext>,
101 step_metas: HashMap<String, Value>,
102 step_naming: Option<Arc<StepNaming>>,
103 projection_placement: Option<Arc<ProjectionPlacement>>,
104}
105
106impl EngineDispatcher {
107 pub fn with_spawner(
113 engine: Engine,
114 op_token: CapToken,
115 spawner: Arc<dyn SpawnerAdapter>,
116 ) -> Self {
117 Self {
118 engine,
119 op_token,
120 spawner,
121 run_ctx: None,
122 step_metas: HashMap::new(),
123 step_naming: None,
124 projection_placement: None,
125 }
126 }
127
128 pub fn with_run(mut self, run_ctx: RunContext) -> Self {
132 self.run_ctx = Some(run_ctx);
133 self
134 }
135
136 pub fn with_step_metas(mut self, step_metas: HashMap<String, Value>) -> Self {
143 self.step_metas = step_metas;
144 self
145 }
146
147 pub fn with_step_naming(mut self, step_naming: Arc<StepNaming>) -> Self {
155 self.step_naming = Some(step_naming);
156 self
157 }
158
159 pub fn with_projection_placement(
167 mut self,
168 projection_placement: Arc<ProjectionPlacement>,
169 ) -> Self {
170 self.projection_placement = Some(projection_placement);
171 self
172 }
173}
174
175fn resolve_step_envelope(
211 step_metas: &HashMap<String, Value>,
212 ref_: &str,
213 input: Value,
214) -> Result<(Value, Option<Value>), EvalError> {
215 let mut obj = match input {
216 Value::Object(obj) => obj,
217 other => return Ok((other, None)),
218 };
219 let Some(envelope) = obj.remove("$step_meta") else {
220 return Ok((Value::Object(obj), None));
221 };
222 let envelope = match envelope {
223 Value::Object(map) => map,
224 other => {
225 return Err(EvalError::DispatcherError {
226 ref_: ref_.to_string(),
227 msg: format!(
228 "malformed $step_meta envelope for step '{ref_}': expected an object, got {other}"
229 ),
230 });
231 }
232 };
233
234 let ref_ctx: Option<Map<String, Value>> = match envelope.get("ref") {
235 None | Some(Value::Null) => None,
236 Some(Value::String(name)) => {
237 let resolved = step_metas.get(name).cloned().ok_or_else(|| {
238 EvalError::DispatcherError {
239 ref_: ref_.to_string(),
240 msg: format!(
241 "$step_meta.ref '{name}' (step '{ref_}') is not a defined Blueprint.metas entry (defined: {:?})",
242 step_metas.keys().collect::<Vec<_>>()
243 ),
244 }
245 })?;
246 match resolved {
247 Value::Object(map) => Some(map),
248 other => {
249 return Err(EvalError::DispatcherError {
250 ref_: ref_.to_string(),
251 msg: format!(
252 "malformed $step_meta: MetaDef '{name}'.ctx must be an object, got {other}"
253 ),
254 });
255 }
256 }
257 }
258 Some(other) => {
259 return Err(EvalError::DispatcherError {
260 ref_: ref_.to_string(),
261 msg: format!(
262 "malformed $step_meta.ref (step '{ref_}'): expected a string, got {other}"
263 ),
264 });
265 }
266 };
267
268 let inline: Option<Map<String, Value>> = match envelope.get("inline") {
269 None | Some(Value::Null) => None,
270 Some(Value::Object(map)) => Some(map.clone()),
271 Some(other) => {
272 return Err(EvalError::DispatcherError {
273 ref_: ref_.to_string(),
274 msg: format!(
275 "malformed $step_meta.inline (step '{ref_}'): expected an object, got {other}"
276 ),
277 });
278 }
279 };
280
281 let step_ctx = match (ref_ctx, inline) {
282 (None, None) => None,
283 (Some(base), None) => Some(Value::Object(base)),
284 (None, Some(inline)) => Some(Value::Object(inline)),
285 (Some(mut base), Some(inline)) => {
286 for (k, v) in inline {
287 base.insert(k, v);
288 }
289 Some(Value::Object(base))
290 }
291 };
292
293 let initial_directive = if let Some(in_value) = obj.remove("$in") {
296 in_value
297 } else if obj.is_empty() {
298 Value::String(String::new())
299 } else {
300 Value::Object(obj)
301 };
302
303 Ok((initial_directive, step_ctx))
304}
305
306#[async_trait]
307impl AsyncDispatcher for EngineDispatcher {
308 async fn dispatch(&self, ref_: &str, input: Value) -> Result<Value, EvalError> {
309 let (initial_directive, step_ctx) = resolve_step_envelope(&self.step_metas, ref_, input)?;
323 let tid = self
324 .engine
325 .start_task(
326 &self.op_token,
327 TaskSpec {
328 agent: ref_.to_string(),
329 initial_directive,
330 step_ctx,
331 },
332 )
333 .await
334 .map_err(|e| EvalError::DispatcherError {
335 ref_: ref_.to_string(),
336 msg: format!("start_task: {e}"),
337 })?;
338
339 if let Some(step_naming) = self.step_naming.clone() {
348 let tid_for_naming = tid.clone();
349 if let Err(e) = self
350 .engine
351 .with_state("EngineDispatcher::dispatch.step_naming", move |s| {
352 s.step_namings.insert(tid_for_naming, step_naming);
353 })
354 .await
355 {
356 tracing::warn!(
357 task_id = %tid,
358 error = %e,
359 "EngineDispatcher::dispatch: failed to snapshot StepNaming into EngineState"
360 );
361 }
362 }
363
364 if let Some(projection_placement) = self.projection_placement.clone() {
372 let tid_for_placement = tid.clone();
373 if let Err(e) = self
374 .engine
375 .with_state(
376 "EngineDispatcher::dispatch.projection_placement",
377 move |s| {
378 s.projection_placements
379 .insert(tid_for_placement, projection_placement);
380 },
381 )
382 .await
383 {
384 tracing::warn!(
385 task_id = %tid,
386 error = %e,
387 "EngineDispatcher::dispatch: failed to snapshot ProjectionPlacement into EngineState"
388 );
389 }
390 }
391
392 let run_id_for_ctx = self.run_ctx.as_ref().map(|rc| rc.run_id.clone());
393 let outcome = self
394 .engine
395 .dispatch_attempt_with(&self.op_token, &tid, &self.spawner, run_id_for_ctx.as_ref())
396 .await;
397
398 if let Some(rc) = &self.run_ctx {
406 let status = match &outcome {
407 Ok(DispatchOutcome::Pass(_)) => "passed",
408 Ok(DispatchOutcome::Blocked(_)) => "blocked",
409 Ok(DispatchOutcome::Suspended(_)) => "suspended",
410 Ok(DispatchOutcome::Cancelled) => "cancelled",
411 Ok(DispatchOutcome::Timeout) => "timeout",
412 Err(_) => "failed",
413 };
414 let entry = StepEntry {
415 step_id: tid.clone(),
416 step_ref: Some(ref_.to_string()),
417 status: Some(status.to_string()),
418 at: now_unix(),
419 };
420 if let Err(e) = rc.run_store.append_step_entry(&rc.run_id, entry).await {
421 tracing::warn!(
422 run_id = %rc.run_id,
423 step_id = %tid,
424 error = %e,
425 "EngineDispatcher::dispatch: append_step_entry failed"
426 );
427 }
428 }
429
430 match outcome {
431 Ok(DispatchOutcome::Pass(v)) => Ok(v),
432 Ok(DispatchOutcome::Blocked(v)) => Err(EvalError::DispatcherError {
433 ref_: ref_.to_string(),
434 msg: format!("blocked: {v}"),
435 }),
436 Ok(other) => Err(EvalError::DispatcherError {
437 ref_: ref_.to_string(),
438 msg: format!("non-terminal outcome: {:?}", other),
439 }),
440 Err(e) => Err(EvalError::DispatcherError {
441 ref_: ref_.to_string(),
442 msg: format!("dispatch_attempt: {e}"),
443 }),
444 }
445 }
446}
447
448#[cfg(test)]
454mod tests {
455 use super::*;
456 use serde_json::json;
457
458 fn metas(pairs: &[(&str, Value)]) -> HashMap<String, Value> {
459 pairs
460 .iter()
461 .map(|(k, v)| (k.to_string(), v.clone()))
462 .collect()
463 }
464
465 #[test]
466 fn no_envelope_string_input_passes_through_unchanged() {
467 let (directive, step_ctx) =
468 resolve_step_envelope(&HashMap::new(), "scout", json!("plain string")).unwrap();
469 assert_eq!(directive, json!("plain string"));
470 assert_eq!(step_ctx, None);
471 }
472
473 #[test]
474 fn no_envelope_plain_object_input_passes_through_unchanged() {
475 let input = json!({ "foo": "bar" });
476 let (directive, step_ctx) =
477 resolve_step_envelope(&HashMap::new(), "scout", input.clone()).unwrap();
478 assert_eq!(directive, input);
479 assert_eq!(step_ctx, None);
480 }
481
482 #[test]
483 fn envelope_with_only_ref_resolves_that_metadef_ctx() {
484 let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
485 let input = json!({ "$step_meta": { "ref": "heavy-scan" }, "$in": "go" });
486 let (directive, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
487 assert_eq!(directive, json!("go"));
488 assert_eq!(step_ctx, Some(json!({ "work_dir": "/x" })));
489 }
490
491 #[test]
492 fn envelope_with_only_inline_uses_inline_verbatim() {
493 let input = json!({
494 "$step_meta": { "inline": { "work_dir": "/inline-only" } },
495 "$in": "go"
496 });
497 let (directive, step_ctx) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
498 assert_eq!(directive, json!("go"));
499 assert_eq!(step_ctx, Some(json!({ "work_dir": "/inline-only" })));
500 }
501
502 #[test]
503 fn inline_wins_over_ref_on_key_collision() {
504 let step_metas = metas(&[(
505 "heavy-scan",
506 json!({ "work_dir": "/ref", "extra": "from-ref" }),
507 )]);
508 let input = json!({
509 "$step_meta": {
510 "ref": "heavy-scan",
511 "inline": { "work_dir": "/inline-wins" }
512 },
513 "$in": "go"
514 });
515 let (_, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
516 assert_eq!(
517 step_ctx,
518 Some(json!({ "work_dir": "/inline-wins", "extra": "from-ref" })),
519 "inline must win the collided key while ref-only keys survive the merge"
520 );
521 }
522
523 #[test]
524 fn dollar_in_rule_extracts_directive_and_ignores_other_sibling_keys() {
525 let input = json!({
526 "$step_meta": { "inline": { "k": "v" } },
527 "$in": "the real directive",
528 "unrelated_sibling": "ignored"
529 });
530 let (directive, step_ctx) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
531 assert_eq!(directive, json!("the real directive"));
532 assert_eq!(step_ctx, Some(json!({ "k": "v" })));
533 }
534
535 #[test]
536 fn no_dollar_in_remainder_becomes_the_directive() {
537 let input = json!({
538 "$step_meta": { "inline": { "k": "v" } },
539 "other_key": "other_value"
540 });
541 let (directive, _) = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap();
542 assert_eq!(directive, json!({ "other_key": "other_value" }));
543 }
544
545 #[test]
546 fn empty_remainder_becomes_empty_string_directive() {
547 let input = json!({ "$step_meta": { "ref": "heavy-scan" } });
548 let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
549 let (directive, step_ctx) = resolve_step_envelope(&step_metas, "scout", input).unwrap();
550 assert_eq!(directive, Value::String(String::new()));
551 assert_eq!(step_ctx, Some(json!({ "work_dir": "/x" })));
552 }
553
554 #[test]
555 fn unresolved_ref_is_a_loud_dispatcher_error_naming_ref_and_defined() {
556 let step_metas = metas(&[("known", json!({}))]);
557 let input = json!({ "$step_meta": { "ref": "unknown" }, "$in": "go" });
558 let err = resolve_step_envelope(&step_metas, "scout", input).unwrap_err();
559 match err {
560 EvalError::DispatcherError { ref_, msg } => {
561 assert_eq!(ref_, "scout");
562 assert!(
563 msg.contains("unknown"),
564 "message must name the unresolved ref: {msg}"
565 );
566 assert!(
567 msg.contains("known"),
568 "message must list defined names: {msg}"
569 );
570 }
571 other => panic!("expected DispatcherError, got {other:?}"),
572 }
573 }
574
575 #[test]
576 fn malformed_step_meta_not_an_object_is_a_loud_error() {
577 let input = json!({ "$step_meta": "not-an-object" });
578 let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
579 assert!(matches!(err, EvalError::DispatcherError { .. }));
580 }
581
582 #[test]
583 fn malformed_ref_non_string_is_a_loud_error() {
584 let input = json!({ "$step_meta": { "ref": 42 } });
585 let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
586 assert!(matches!(err, EvalError::DispatcherError { .. }));
587 }
588
589 #[test]
590 fn malformed_inline_non_object_is_a_loud_error() {
591 let input = json!({ "$step_meta": { "inline": "not-an-object" } });
592 let err = resolve_step_envelope(&HashMap::new(), "scout", input).unwrap_err();
593 assert!(matches!(err, EvalError::DispatcherError { .. }));
594 }
595
596 #[test]
597 fn ref_resolved_metadef_ctx_non_object_is_a_loud_error() {
598 let step_metas = metas(&[("bad", json!("not-an-object"))]);
599 let input = json!({ "$step_meta": { "ref": "bad" } });
600 let err = resolve_step_envelope(&step_metas, "scout", input).unwrap_err();
601 assert!(matches!(err, EvalError::DispatcherError { .. }));
602 }
603
604 #[tokio::test]
608 async fn dispatch_step_meta_envelope_never_leaks_into_stored_prompt() {
609 use crate::blueprint::compiler::{RustFnInProcessSpawnerFactory, SpawnerFactory};
610 use crate::core::config::EngineCfg;
611 use crate::types::{Role, StepId};
612 use crate::worker::adapter::WorkerResult;
613 use std::sync::Mutex as StdMutex;
614 use std::time::Duration;
615
616 let captured_tid: Arc<StdMutex<Option<StepId>>> = Arc::new(StdMutex::new(None));
617 let captured_tid_for_fn = captured_tid.clone();
618 let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", move |inv| {
619 let captured_tid = captured_tid_for_fn.clone();
620 async move {
621 *captured_tid.lock().unwrap() = Some(inv.task_id.clone());
622 Ok(WorkerResult {
623 value: json!({ "ok": true }),
624 ok: true,
625 })
626 }
627 });
628 let def = AgentDef {
629 name: "scout".into(),
630 kind: AgentKind::RustFn,
631 spec: json!({ "fn_id": "echo" }),
632 profile: None,
633 meta: None,
634 };
635 let spawner = factory.build(&def, None).expect("build");
636
637 let engine = Engine::new(EngineCfg::default());
638 let token = engine
639 .attach("ut-op", Role::Operator, Duration::from_secs(30))
640 .await
641 .expect("attach");
642 let step_metas = metas(&[("heavy-scan", json!({ "work_dir": "/x" }))]);
643 let dispatcher = EngineDispatcher::with_spawner(engine.clone(), token, spawner)
644 .with_step_metas(step_metas);
645
646 let input = json!({
647 "$step_meta": { "ref": "heavy-scan" },
648 "$in": "do the thing"
649 });
650 let out = dispatcher
651 .dispatch("scout", input)
652 .await
653 .expect("dispatch ok");
654 assert_eq!(out, json!({ "ok": true }));
655
656 let tid = captured_tid
657 .lock()
658 .unwrap()
659 .clone()
660 .expect("task_id captured");
661 let stored_prompt = engine
662 .with_state("test.read_prompt", move |s| {
663 s.prompts.get(&(tid, 1)).cloned()
664 })
665 .await
666 .expect("with_state")
667 .expect("prompt recorded for attempt 1");
668 assert_eq!(
669 stored_prompt,
670 json!("do the thing"),
671 "the stored prompt must be the post-envelope directive, with no $step_meta leakage"
672 );
673 }
674}