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