1use crate::blueprint::{AgentDef, AgentKind, Blueprint, BlueprintMetadata};
31use crate::core::ctx::Ctx;
32use crate::core::engine::Engine;
33use crate::core::projection_placement::{ProjectionPlacement, ProjectionPlacementError};
34use crate::core::step_naming::{StepNaming, StepNamingError};
35use crate::operator::{Operator, OperatorSpawner, WorkerBinding};
36use crate::types::{CapToken, StepId};
37use crate::worker::adapter::{InProcSpawner, SpawnError, SpawnerAdapter, WorkerFn};
38use crate::worker::process_spawner::{ProcessSpawner, StreamMode};
39use crate::worker::Worker;
40use async_trait::async_trait;
41use mlua_flow_ir::{Expr, Node as FlowNode};
42use serde_json::Value;
43use std::collections::HashMap;
44use std::sync::Arc;
45use thiserror::Error;
46
47#[derive(Debug, Error)]
52pub enum CompileError {
53 #[error("unknown agent kind in SpawnerRegistry: {0:?}")]
56 UnknownKind(AgentKind),
57 #[error("agent '{name}' spec invalid: {msg}")]
60 InvalidSpec {
61 name: String,
63 msg: String,
65 },
66 #[error("flow references agent '{0}' but no AgentDef matches")]
69 UnresolvedRef(String),
70 #[error("duplicate AgentDef name: {0}")]
72 DuplicateAgent(String),
73 #[error("agent '{agent}' operator_ref '{op_ref}' does not match any OperatorDef.name in Blueprint.operators (defined: {defined:?})")]
76 UnresolvedOperatorRef {
77 agent: String,
79 op_ref: String,
81 defined: Vec<String>,
84 },
85 #[error("{where_} names an undefined MetaDef: '{meta_ref}' (defined: {defined:?})")]
89 UnresolvedMetaRef {
90 where_: String,
94 meta_ref: String,
96 defined: Vec<String>,
99 },
100 #[error("StepNaming collision: {0}")]
106 StepNamingCollision(#[from] StepNamingError),
107 #[error("invalid projection_placement: {0}")]
114 InvalidProjectionPlacement(#[from] ProjectionPlacementError),
115}
116
117pub trait SpawnerFactory: Send + Sync {
129 fn build(
132 &self,
133 agent_def: &AgentDef,
134 hint: Option<&Value>,
135 ) -> Result<Arc<dyn SpawnerAdapter>, CompileError>;
136}
137
138pub trait SpawnerFactoryKind: SpawnerFactory {
154 const KIND: AgentKind;
157 type Worker: crate::worker::Worker;
164}
165
166#[derive(Clone)]
169pub struct SpawnerRegistry {
170 factories: HashMap<AgentKind, Arc<dyn SpawnerFactory>>,
171}
172
173impl SpawnerRegistry {
174 pub fn new() -> Self {
176 Self {
177 factories: HashMap::new(),
178 }
179 }
180 pub fn register<F: SpawnerFactoryKind + 'static>(&mut self, factory: Arc<F>) -> &mut Self {
189 let f: Arc<dyn SpawnerFactory> = factory;
190 self.factories.insert(F::KIND, f);
191 self
192 }
193}
194
195impl Default for SpawnerRegistry {
196 fn default() -> Self {
197 Self::new()
198 }
199}
200
201pub struct Compiler {
208 registry: SpawnerRegistry,
209 default_spawner: Option<Arc<dyn SpawnerAdapter>>,
210}
211
212pub struct CompiledBlueprint {
216 pub router: Arc<CompiledAgentTable>,
218 pub flow: FlowNode,
220 pub metadata: BlueprintMetadata,
222 pub step_naming: Arc<StepNaming>,
227 pub projection_placement: Arc<ProjectionPlacement>,
233}
234
235impl Compiler {
236 pub fn new(registry: SpawnerRegistry) -> Self {
240 Self {
241 registry,
242 default_spawner: None,
243 }
244 }
245
246 pub fn with_default(mut self, sp: Arc<dyn SpawnerAdapter>) -> Self {
250 self.default_spawner = Some(sp);
251 self
252 }
253
254 pub fn compile(&self, bp: &Blueprint) -> Result<CompiledBlueprint, CompileError> {
259 let mut routes: HashMap<String, Arc<dyn SpawnerAdapter>> = HashMap::new();
260 let mut seen: HashMap<String, ()> = HashMap::new();
261
262 let defined: Vec<String> = bp.operators.iter().map(|o| o.name.clone()).collect();
268 for ad in &bp.agents {
269 if !matches!(ad.kind, AgentKind::Operator) {
270 continue;
271 }
272 let op_ref = ad.spec.get("operator_ref").and_then(|v| v.as_str());
273 if let Some(op_ref) = op_ref {
274 if !defined.iter().any(|n| n == op_ref) {
275 return Err(CompileError::UnresolvedOperatorRef {
276 agent: ad.name.clone(),
277 op_ref: op_ref.to_string(),
278 defined: defined.clone(),
279 });
280 }
281 }
282 }
284
285 let metas_defined: Vec<String> = bp.metas.iter().map(|m| m.name.clone()).collect();
289 for ad in &bp.agents {
290 let meta_ref = ad.meta.as_ref().and_then(|m| m.meta_ref.as_ref());
291 if let Some(meta_ref) = meta_ref {
292 if !metas_defined.iter().any(|n| n == meta_ref) {
293 return Err(CompileError::UnresolvedMetaRef {
294 where_: format!("AgentMeta.meta_ref of agent '{}'", ad.name),
295 meta_ref: meta_ref.clone(),
296 defined: metas_defined.clone(),
297 });
298 }
299 }
300 }
301 let mut static_step_meta_refs: Vec<(String, String)> = Vec::new();
307 collect_step_meta_refs(&bp.flow, &mut static_step_meta_refs);
308 for (where_, meta_ref) in static_step_meta_refs {
309 if !metas_defined.iter().any(|n| n == &meta_ref) {
310 return Err(CompileError::UnresolvedMetaRef {
311 where_,
312 meta_ref,
313 defined: metas_defined.clone(),
314 });
315 }
316 }
317
318 for ad in &bp.agents {
319 if seen.contains_key(&ad.name) {
320 return Err(CompileError::DuplicateAgent(ad.name.clone()));
321 }
322 seen.insert(ad.name.clone(), ());
323
324 let factory = match self.registry.factories.get(&ad.kind) {
325 Some(f) => f.clone(),
326 None => {
327 if bp.strategy.strict_kind {
328 return Err(CompileError::UnknownKind(ad.kind.clone()));
329 } else {
330 tracing::warn!(
331 agent = %ad.name,
332 kind = ?ad.kind,
333 "no spawner factory registered for agent kind; \
334 dropping agent from routing table (strict_kind=false)"
335 );
336 continue;
337 }
338 }
339 };
340 let hint = bp.hints.per_agent.get(&ad.name);
341 let spawner = factory.build(ad, hint)?;
342 routes.insert(ad.name.clone(), spawner);
343 }
344
345 if bp.strategy.strict_refs {
346 verify_refs(&bp.flow, &routes, self.default_spawner.is_some())?;
347 }
348
349 let (step_naming, step_naming_warnings) = StepNaming::from_blueprint(bp)?;
357 for warning in &step_naming_warnings {
358 tracing::warn!(
359 name = %warning.name,
360 first_step_ref = %warning.first_step_ref,
361 second_step_ref = %warning.second_step_ref,
362 "StepNaming: undeclared steps' canonical/alias names collide; \
363 the step whose own ref matches the name keeps it (data-plane priority)"
364 );
365 }
366
367 let projection_placement =
375 ProjectionPlacement::from_spec(bp.projection_placement.as_ref())?;
376
377 let router = Arc::new(CompiledAgentTable {
378 routes,
379 default: self.default_spawner.clone(),
380 });
381 Ok(CompiledBlueprint {
382 router,
383 flow: bp.flow.clone(),
384 metadata: bp.metadata.clone(),
385 step_naming: Arc::new(step_naming),
386 projection_placement: Arc::new(projection_placement),
387 })
388 }
389}
390
391fn verify_refs(
394 node: &FlowNode,
395 routes: &HashMap<String, Arc<dyn SpawnerAdapter>>,
396 has_default: bool,
397) -> Result<(), CompileError> {
398 let mut refs: Vec<String> = Vec::new();
399 collect_refs(node, &mut refs);
400 for r in refs {
401 if !routes.contains_key(&r) && !has_default {
402 return Err(CompileError::UnresolvedRef(r));
403 }
404 }
405 Ok(())
406}
407
408fn collect_refs(node: &FlowNode, out: &mut Vec<String>) {
409 match node {
410 FlowNode::Step { ref_, .. } => out.push(ref_.clone()),
411 FlowNode::Seq { children } => {
412 for c in children {
413 collect_refs(c, out);
414 }
415 }
416 FlowNode::Branch { then_, else_, .. } => {
417 collect_refs(then_, out);
418 collect_refs(else_, out);
419 }
420 FlowNode::Fanout { body, .. } => collect_refs(body, out),
421 FlowNode::Loop { body, .. } => collect_refs(body, out),
422 FlowNode::Try { body, catch, .. } => {
423 collect_refs(body, out);
424 collect_refs(catch, out);
425 }
426 FlowNode::Assign { .. } => {} }
428}
429
430fn collect_step_meta_refs(node: &FlowNode, out: &mut Vec<(String, String)>) {
438 match node {
439 FlowNode::Step { ref_, in_, .. } => {
440 if let Expr::Lit { value } = in_ {
441 if let Some(meta_ref) = static_step_meta_ref(value) {
442 out.push((format!("Step '{ref_}' $step_meta.ref"), meta_ref));
443 }
444 }
445 }
446 FlowNode::Seq { children } => {
447 for c in children {
448 collect_step_meta_refs(c, out);
449 }
450 }
451 FlowNode::Branch { then_, else_, .. } => {
452 collect_step_meta_refs(then_, out);
453 collect_step_meta_refs(else_, out);
454 }
455 FlowNode::Fanout { body, .. } => collect_step_meta_refs(body, out),
456 FlowNode::Loop { body, .. } => collect_step_meta_refs(body, out),
457 FlowNode::Try { body, catch, .. } => {
458 collect_step_meta_refs(body, out);
459 collect_step_meta_refs(catch, out);
460 }
461 FlowNode::Assign { .. } => {} }
463}
464
465fn static_step_meta_ref(value: &Value) -> Option<String> {
472 value
473 .as_object()?
474 .get("$step_meta")?
475 .as_object()?
476 .get("ref")?
477 .as_str()
478 .map(str::to_string)
479}
480
481pub struct CompiledAgentTable {
494 pub(crate) routes: HashMap<String, Arc<dyn SpawnerAdapter>>,
495 pub(crate) default: Option<Arc<dyn SpawnerAdapter>>,
496}
497
498impl CompiledAgentTable {
499 pub fn has_route(&self, agent: &str) -> bool {
502 self.routes.contains_key(agent)
503 }
504 pub fn routed_agents(&self) -> Vec<String> {
506 self.routes.keys().cloned().collect()
507 }
508}
509
510#[async_trait]
511impl SpawnerAdapter for CompiledAgentTable {
512 async fn spawn(
513 &self,
514 engine: &Engine,
515 ctx: &Ctx,
516 task_id: StepId,
517 attempt: u32,
518 token: CapToken,
519 ) -> Result<Box<dyn Worker>, SpawnError> {
520 let sp = self
521 .routes
522 .get(&ctx.agent)
523 .cloned()
524 .or_else(|| self.default.clone())
525 .ok_or_else(|| SpawnError::NotRegistered(ctx.agent.clone()))?;
526 sp.spawn(engine, ctx, task_id, attempt, token).await
527 }
528}
529
530pub struct SubprocessProcessSpawnerFactory;
548
549impl SpawnerFactoryKind for SubprocessProcessSpawnerFactory {
550 const KIND: AgentKind = AgentKind::Subprocess;
551 type Worker = crate::worker::process_spawner::ProcessWorker;
552}
553
554impl SpawnerFactory for SubprocessProcessSpawnerFactory {
555 fn build(
556 &self,
557 agent_def: &AgentDef,
558 _hint: Option<&Value>,
559 ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
560 let agent_name = &agent_def.name;
561 let spec = &agent_def.spec;
562 let invalid = |msg: String| CompileError::InvalidSpec {
563 name: agent_name.to_string(),
564 msg,
565 };
566 let program = spec
567 .get("program")
568 .and_then(|v| v.as_str())
569 .ok_or_else(|| invalid("shell spec: 'program' (string) required".into()))?
570 .to_string();
571 let args: Vec<String> = spec
572 .get("args")
573 .and_then(|v| v.as_array())
574 .map(|a| {
575 a.iter()
576 .filter_map(|x| x.as_str().map(|s| s.to_string()))
577 .collect()
578 })
579 .unwrap_or_default();
580 let use_stdin = spec
581 .get("use_stdin")
582 .and_then(|v| v.as_bool())
583 .unwrap_or(true);
584 let stream_mode = match spec.get("stream_mode").and_then(|v| v.as_str()) {
585 Some("ndjson_lines") => Some(StreamMode::NdjsonLines),
586 Some("sse_events") => Some(StreamMode::SseEvents),
587 Some("length_prefixed") => Some(StreamMode::LengthPrefixed),
588 Some(other) => return Err(invalid(format!("unknown stream_mode: {other}"))),
589 None => None,
590 };
591
592 let mut sp = ProcessSpawner {
593 program,
594 args,
595 use_stdin,
596 stream_mode,
597 };
598 if let Some(mode) = sp.stream_mode.clone() {
599 sp = sp.stream_mode(mode);
600 }
601 Ok(Arc::new(sp))
602 }
603}
604
605pub struct LuaInProcessSpawnerFactory {
636 registry: HashMap<String, WorkerFn>,
637 bridges: HashMap<String, HostBridge>,
638}
639
640#[derive(Clone)]
652pub struct HostBridge(
653 Arc<dyn Fn(serde_json::Value) -> Result<serde_json::Value, String> + Send + Sync>,
654);
655
656impl HostBridge {
657 pub fn new<F>(f: F) -> Self
659 where
660 F: Fn(serde_json::Value) -> Result<serde_json::Value, String> + Send + Sync + 'static,
661 {
662 Self(Arc::new(f))
663 }
664
665 pub fn call(&self, arg: serde_json::Value) -> Result<serde_json::Value, String> {
669 (self.0)(arg)
670 }
671}
672
673#[derive(Clone)]
680pub struct LuaScriptSource {
681 pub source: String,
683 pub label: String,
686}
687
688impl LuaScriptSource {
689 pub fn new(source: impl Into<String>, label: impl Into<String>) -> Self {
691 Self {
692 source: source.into(),
693 label: label.into(),
694 }
695 }
696}
697
698impl LuaInProcessSpawnerFactory {
699 pub fn new() -> Self {
701 Self {
702 registry: HashMap::new(),
703 bridges: HashMap::new(),
704 }
705 }
706
707 pub fn with_bridge(mut self, name: impl Into<String>, bridge: HostBridge) -> Self {
714 self.bridges.insert(name.into(), bridge);
715 self
716 }
717
718 pub fn register_lua(mut self, fn_id: impl Into<String>, source: LuaScriptSource) -> Self {
736 let source = Arc::new(source);
737 let bridges = Arc::new(self.bridges.clone());
738 let wrapped: WorkerFn = Arc::new(move |inv| {
739 let source = source.clone();
740 let bridges = bridges.clone();
741 Box::pin(run_lua_worker(source, bridges, inv))
742 });
743 self.registry.insert(fn_id.into(), wrapped);
744 self
745 }
746}
747
748async fn run_lua_worker(
750 source: Arc<LuaScriptSource>,
751 bridges: Arc<HashMap<String, HostBridge>>,
752 inv: crate::worker::adapter::WorkerInvocation,
753) -> Result<crate::worker::adapter::WorkerResult, crate::worker::adapter::WorkerError> {
754 use crate::worker::adapter::WorkerError;
755 use mlua::LuaSerdeExt;
756
757 let label = source.label.clone();
758 let outcome =
759 tokio::task::spawn_blocking(move || -> Result<(serde_json::Value, bool), String> {
760 let lua = mlua::Lua::new();
761 let g = lua.globals();
762
763 g.set("_PROMPT", inv.prompt.clone())
765 .map_err(|e| format!("set _PROMPT: {e}"))?;
766 g.set("_AGENT", inv.agent.clone())
767 .map_err(|e| format!("set _AGENT: {e}"))?;
768 g.set("_TASK_ID", inv.task_id.to_string())
769 .map_err(|e| format!("set _TASK_ID: {e}"))?;
770 g.set("_ATTEMPT", inv.attempt as i64)
771 .map_err(|e| format!("set _ATTEMPT: {e}"))?;
772
773 if let Ok(json_val) = serde_json::from_str::<serde_json::Value>(&inv.prompt) {
775 let lua_val = lua
776 .to_value(&json_val)
777 .map_err(|e| format!("_CTX to_value: {e}"))?;
778 g.set("_CTX", lua_val)
779 .map_err(|e| format!("set _CTX: {e}"))?;
780 }
781
782 if !bridges.is_empty() {
784 let host = lua
785 .create_table()
786 .map_err(|e| format!("create host table: {e}"))?;
787 for (name, bridge) in bridges.iter() {
788 let bridge = bridge.clone();
789 let bname = name.clone();
790 let f = lua
791 .create_function(move |lua, arg: mlua::Value| {
792 let json_arg: serde_json::Value = lua.from_value(arg).map_err(|e| {
793 mlua::Error::external(format!("bridge {bname} arg → json: {e}"))
794 })?;
795 let result_json =
796 bridge.call(json_arg).map_err(mlua::Error::external)?;
797 lua.to_value(&result_json).map_err(|e| {
798 mlua::Error::external(format!("bridge {bname} ret → lua: {e}"))
799 })
800 })
801 .map_err(|e| format!("create_function {name}: {e}"))?;
802 host.set(name.as_str(), f)
803 .map_err(|e| format!("host.{name} set: {e}"))?;
804 }
805 g.set("host", host).map_err(|e| format!("set host: {e}"))?;
806 }
807
808 let result: mlua::Value = lua
810 .load(&source.source)
811 .set_name(&source.label)
812 .eval()
813 .map_err(|e| format!("lua eval [{}]: {e}", source.label))?;
814
815 let json_result: serde_json::Value = lua
817 .from_value(result)
818 .map_err(|e| format!("lua → json [{}]: {e}", source.label))?;
819
820 let (value, ok) = match &json_result {
821 serde_json::Value::Object(map)
822 if map.contains_key("value") || map.contains_key("ok") =>
823 {
824 let ok = map.get("ok").and_then(|v| v.as_bool()).unwrap_or(true);
825 let value = map.get("value").cloned().unwrap_or(json_result.clone());
826 (value, ok)
827 }
828 _ => (json_result, true),
829 };
830 Ok((value, ok))
831 })
832 .await
833 .map_err(|e| WorkerError::Failed(format!("spawn_blocking join [{label}]: {e}")))?
834 .map_err(WorkerError::Failed)?;
835
836 Ok(crate::worker::adapter::WorkerResult {
837 value: outcome.0,
838 ok: outcome.1,
839 })
840}
841
842impl Default for LuaInProcessSpawnerFactory {
843 fn default() -> Self {
844 Self::new()
845 }
846}
847
848impl SpawnerFactoryKind for LuaInProcessSpawnerFactory {
849 const KIND: AgentKind = AgentKind::Lua;
850 type Worker = LuaWorker;
851}
852
853impl SpawnerFactory for LuaInProcessSpawnerFactory {
854 fn build(
855 &self,
856 agent_def: &AgentDef,
857 _hint: Option<&Value>,
858 ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
859 if let Some(source) = agent_def.spec.get("source").and_then(|v| v.as_str()) {
865 let label = agent_def
866 .spec
867 .get("label")
868 .and_then(|v| v.as_str())
869 .map(str::to_string)
870 .unwrap_or_else(|| format!("{}.lua", agent_def.name));
871 let script = Arc::new(LuaScriptSource::new(source.to_string(), label));
872 let bridges = Arc::new(self.bridges.clone());
873 let wrapped: WorkerFn = Arc::new(move |inv| {
874 let source = script.clone();
875 let bridges = bridges.clone();
876 Box::pin(run_lua_worker(source, bridges, inv))
877 });
878 let mut sp: InProcSpawner<LuaWorker> = InProcSpawner::<LuaWorker>::typed();
879 sp.registry.insert(agent_def.name.to_string(), wrapped);
880 return Ok(Arc::new(sp));
881 }
882 build_inproc_from_registry::<LuaWorker>(&self.registry, agent_def, "lua")
883 }
884}
885
886pub struct RustFnInProcessSpawnerFactory {
900 registry: HashMap<String, WorkerFn>,
901}
902
903impl RustFnInProcessSpawnerFactory {
904 pub fn new() -> Self {
906 Self {
907 registry: HashMap::new(),
908 }
909 }
910
911 pub fn register_fn<F, Fut>(mut self, fn_id: impl Into<String>, f: F) -> Self
914 where
915 F: Fn(crate::worker::adapter::WorkerInvocation) -> Fut + Send + Sync + 'static,
916 Fut: std::future::Future<
917 Output = Result<
918 crate::worker::adapter::WorkerResult,
919 crate::worker::adapter::WorkerError,
920 >,
921 > + Send
922 + 'static,
923 {
924 let f = Arc::new(f);
925 let wrapped: WorkerFn = Arc::new(move |inv| {
926 let f = f.clone();
927 Box::pin(f(inv))
928 });
929 self.registry.insert(fn_id.into(), wrapped);
930 self
931 }
932}
933
934impl Default for RustFnInProcessSpawnerFactory {
935 fn default() -> Self {
936 Self::new()
937 }
938}
939
940impl SpawnerFactoryKind for RustFnInProcessSpawnerFactory {
941 const KIND: AgentKind = AgentKind::RustFn;
942 type Worker = RustFnWorker;
943}
944
945impl SpawnerFactory for RustFnInProcessSpawnerFactory {
946 fn build(
947 &self,
948 agent_def: &AgentDef,
949 _hint: Option<&Value>,
950 ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
951 build_inproc_from_registry::<RustFnWorker>(&self.registry, agent_def, "rust_fn")
952 }
953}
954
955fn build_inproc_from_registry<W>(
961 registry: &HashMap<String, WorkerFn>,
962 agent_def: &AgentDef,
963 kind_label: &str,
964) -> Result<Arc<dyn SpawnerAdapter>, CompileError>
965where
966 W: crate::worker::Worker + From<crate::worker::WorkerJoinHandler> + Send + Sync + 'static,
967{
968 let agent_name = &agent_def.name;
969 let spec = &agent_def.spec;
970 let invalid = |msg: String| CompileError::InvalidSpec {
971 name: agent_name.to_string(),
972 msg,
973 };
974 let fn_id = spec
975 .get("fn_id")
976 .and_then(|v| v.as_str())
977 .ok_or_else(|| invalid(format!("{kind_label} spec: 'fn_id' (string) required")))?;
978 let f = registry
979 .get(fn_id)
980 .cloned()
981 .ok_or_else(|| invalid(format!("fn_id '{fn_id}' not registered in factory")))?;
982 let mut sp: InProcSpawner<W> = InProcSpawner::<W>::typed();
983 sp.registry.insert(agent_name.to_string(), f);
987 Ok(Arc::new(sp))
988}
989
990pub struct LuaWorker {
995 pub handler: crate::worker::WorkerJoinHandler,
997}
998
999impl From<crate::worker::WorkerJoinHandler> for LuaWorker {
1000 fn from(handler: crate::worker::WorkerJoinHandler) -> Self {
1001 Self { handler }
1002 }
1003}
1004
1005#[async_trait::async_trait]
1006impl crate::worker::Worker for LuaWorker {
1007 fn id(&self) -> &crate::types::WorkerId {
1008 &self.handler.worker_id
1009 }
1010 fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
1011 self.handler.cancel.clone()
1012 }
1013 async fn join(self: Box<Self>) -> Result<(), crate::worker::adapter::WorkerError> {
1014 self.handler.await_completion().await
1015 }
1016}
1017
1018pub struct RustFnWorker {
1023 pub handler: crate::worker::WorkerJoinHandler,
1025}
1026
1027impl From<crate::worker::WorkerJoinHandler> for RustFnWorker {
1028 fn from(handler: crate::worker::WorkerJoinHandler) -> Self {
1029 Self { handler }
1030 }
1031}
1032
1033#[async_trait::async_trait]
1034impl crate::worker::Worker for RustFnWorker {
1035 fn id(&self) -> &crate::types::WorkerId {
1036 &self.handler.worker_id
1037 }
1038 fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
1039 self.handler.cancel.clone()
1040 }
1041 async fn join(self: Box<Self>) -> Result<(), crate::worker::adapter::WorkerError> {
1042 self.handler.await_completion().await
1043 }
1044}
1045
1046pub struct OperatorSpawnerFactory {
1099 operators: Arc<std::sync::RwLock<HashMap<String, Arc<dyn Operator>>>>,
1100}
1101
1102impl OperatorSpawnerFactory {
1103 pub fn new() -> Self {
1105 Self {
1106 operators: Arc::new(std::sync::RwLock::new(HashMap::new())),
1107 }
1108 }
1109
1110 pub fn register_operator(&self, id: impl Into<String>, op: Arc<dyn Operator>) -> &Self {
1116 self.operators
1117 .write()
1118 .expect("OperatorSpawnerFactory.operators RwLock poisoned")
1119 .insert(id.into(), op);
1120 self
1121 }
1122
1123 pub fn unregister_operator(&self, id: &str) -> &Self {
1126 self.operators
1127 .write()
1128 .expect("OperatorSpawnerFactory.operators RwLock poisoned")
1129 .remove(id);
1130 self
1131 }
1132}
1133
1134impl Default for OperatorSpawnerFactory {
1135 fn default() -> Self {
1136 Self::new()
1137 }
1138}
1139
1140impl SpawnerFactoryKind for OperatorSpawnerFactory {
1141 const KIND: AgentKind = AgentKind::Operator;
1142 type Worker = crate::operator::OperatorWorker;
1143}
1144
1145impl SpawnerFactory for OperatorSpawnerFactory {
1146 fn build(
1147 &self,
1148 agent_def: &AgentDef,
1149 _hint: Option<&Value>,
1150 ) -> Result<Arc<dyn SpawnerAdapter>, CompileError> {
1151 let agent_name = &agent_def.name;
1152 let spec = &agent_def.spec;
1153 let system_prompt = agent_def.profile.as_ref().map(|p| p.system_prompt.clone());
1159 let invalid = |msg: String| CompileError::InvalidSpec {
1160 name: agent_name.to_string(),
1161 msg,
1162 };
1163 let op_ref = spec
1164 .get("operator_ref")
1165 .and_then(|v| v.as_str())
1166 .ok_or_else(|| invalid("operator spec: 'operator_ref' (string) required".into()))?;
1167 let operators = self
1168 .operators
1169 .read()
1170 .expect("OperatorSpawnerFactory.operators RwLock poisoned");
1171 let op = operators.get(op_ref).cloned().ok_or_else(|| {
1172 let mut names: Vec<String> = operators.keys().cloned().collect();
1173 names.sort();
1174 let names_list = if names.is_empty() {
1175 "<none>".to_string()
1176 } else {
1177 names.join(", ")
1178 };
1179 invalid(format!(
1180 "operator_ref '{op_ref}' not registered in factory. \
1181 Registered sids: [{names_list}]. \
1182 Hint: call mse_operator_join(roles=[...]) to mint the sid first."
1183 ))
1184 })?;
1185 drop(operators);
1186
1187 let worker_binding = agent_def
1194 .profile
1195 .as_ref()
1196 .and_then(|p| p.worker_binding.as_ref())
1197 .map(|variant| WorkerBinding {
1198 variant: variant.clone(),
1199 tools: agent_def
1200 .profile
1201 .as_ref()
1202 .map(|p| p.tools.clone())
1203 .unwrap_or_default(),
1204 });
1205 if op.requires_worker_binding() && worker_binding.is_none() {
1206 return Err(invalid(
1211 "profile.worker_binding is required for this operator backend. \
1212 Fix by either: \
1213 (a) if authoring the Blueprint JSON directly, add \
1214 `agents[N].profile.worker_binding: \"<subagent-type>\"` \
1215 to the JSON literal; or \
1216 (b) if using an $agent_md file ref, add \
1217 `worker_binding: <subagent-type>` to the agent .md frontmatter."
1218 .into(),
1219 ));
1220 }
1221 Ok(Arc::new(OperatorSpawner::new(
1222 op,
1223 system_prompt,
1224 worker_binding,
1225 )))
1226 }
1227}
1228
1229#[cfg(test)]
1230mod operator_spawner_factory_worker_binding_tests {
1231 use super::*;
1232 use crate::blueprint::AgentProfile;
1233 use crate::core::ctx::Ctx;
1234 use crate::types::CapToken;
1235 use crate::worker::adapter::{WorkerError, WorkerResult};
1236
1237 struct StubOperator {
1242 requires_binding: bool,
1243 }
1244
1245 #[async_trait]
1246 impl Operator for StubOperator {
1247 async fn execute(
1248 &self,
1249 _ctx: &Ctx,
1250 _system: Option<String>,
1251 _prompt: Value,
1252 _worker: Option<WorkerBinding>,
1253 _worker_token: CapToken,
1254 ) -> Result<WorkerResult, WorkerError> {
1255 Ok(WorkerResult {
1256 value: Value::Null,
1257 ok: true,
1258 })
1259 }
1260
1261 fn requires_worker_binding(&self) -> bool {
1262 self.requires_binding
1263 }
1264 }
1265
1266 fn agent_def_with(profile: Option<AgentProfile>) -> AgentDef {
1267 AgentDef {
1268 name: "test-agent".to_string(),
1269 kind: AgentKind::Operator,
1270 spec: serde_json::json!({ "operator_ref": "op1" }),
1271 profile,
1272 meta: None,
1273 }
1274 }
1275
1276 #[test]
1277 fn build_fails_loud_when_binding_required_but_absent() {
1278 let factory = OperatorSpawnerFactory::new();
1279 factory.register_operator(
1280 "op1",
1281 Arc::new(StubOperator {
1282 requires_binding: true,
1283 }) as Arc<dyn Operator>,
1284 );
1285 let def = agent_def_with(Some(AgentProfile::default()));
1286 match factory.build(&def, None) {
1287 Err(CompileError::InvalidSpec { name, msg }) => {
1288 assert_eq!(name, "test-agent");
1289 assert!(
1290 msg.contains("worker_binding is required"),
1291 "unexpected message: {msg}"
1292 );
1293 assert!(
1297 msg.contains("agents[N].profile.worker_binding"),
1298 "message missing JSON-direct hint (issue #9): {msg}"
1299 );
1300 assert!(
1301 msg.contains("agent .md frontmatter"),
1302 "message missing $agent_md hint: {msg}"
1303 );
1304 }
1305 Err(other) => panic!("expected InvalidSpec, got: {other:?}"),
1306 Ok(_) => panic!("expected compile-time failure, got Ok"),
1307 }
1308 }
1309
1310 #[test]
1311 fn build_succeeds_when_binding_required_and_present() {
1312 let factory = OperatorSpawnerFactory::new();
1313 factory.register_operator(
1314 "op1",
1315 Arc::new(StubOperator {
1316 requires_binding: true,
1317 }) as Arc<dyn Operator>,
1318 );
1319 let profile = AgentProfile {
1320 worker_binding: Some("mse-worker-coder".to_string()),
1321 tools: vec!["Read".to_string(), "Edit".to_string()],
1322 ..Default::default()
1323 };
1324 let def = agent_def_with(Some(profile));
1325 assert!(
1326 factory.build(&def, None).is_ok(),
1327 "expected Ok when worker_binding is declared"
1328 );
1329 }
1330
1331 #[test]
1332 fn build_succeeds_when_binding_not_required_and_absent() {
1333 let factory = OperatorSpawnerFactory::new();
1334 factory.register_operator(
1335 "op1",
1336 Arc::new(StubOperator {
1337 requires_binding: false,
1338 }) as Arc<dyn Operator>,
1339 );
1340 let def = agent_def_with(Some(AgentProfile::default()));
1341 assert!(
1342 factory.build(&def, None).is_ok(),
1343 "backends that don't require a binding must not be gated by its absence"
1344 );
1345 }
1346}
1347
1348#[cfg(test)]
1356mod lua_inline_source_tests {
1357 use super::*;
1358 use crate::types::{CapToken, Role, StepId};
1359
1360 fn agent(name: &str, spec: Value) -> AgentDef {
1361 AgentDef {
1362 name: name.to_string(),
1363 kind: AgentKind::Lua,
1364 spec,
1365 profile: None,
1366 meta: None,
1367 }
1368 }
1369
1370 fn test_invocation(prompt: &str) -> crate::worker::adapter::WorkerInvocation {
1371 crate::worker::adapter::WorkerInvocation {
1372 token: CapToken {
1373 agent_id: "a".into(),
1374 role: Role::Worker,
1375 scopes: vec!["*".into()],
1376 issued_at: 0,
1377 expire_at: u64::MAX / 2,
1378 max_uses: None,
1379 nonce: "test-nonce".into(),
1380 sig_hex: "".into(),
1381 },
1382 task_id: StepId::parse("ST-test").expect("StepId parse"),
1383 attempt: 1,
1384 agent: "g".into(),
1385 prompt: prompt.into(),
1386 sink: None,
1387 cancel_token: None,
1388 }
1389 }
1390
1391 #[test]
1392 fn build_accepts_inline_source_without_pre_registration() {
1393 let factory = LuaInProcessSpawnerFactory::new();
1394 let def = agent(
1395 "g",
1396 serde_json::json!({ "source": "return { value = 42, ok = true }" }),
1397 );
1398 assert!(
1399 factory.build(&def, None).is_ok(),
1400 "inline spec.source must build without a pre-registered fn_id"
1401 );
1402 }
1403
1404 #[test]
1405 fn build_rejects_when_neither_source_nor_fn_id_is_present() {
1406 let factory = LuaInProcessSpawnerFactory::new();
1407 let def = agent("g", serde_json::json!({}));
1408 match factory.build(&def, None) {
1409 Err(CompileError::InvalidSpec { msg, .. }) => {
1410 assert!(
1411 msg.contains("fn_id"),
1412 "empty spec must still surface the fn_id-required message: {msg}"
1413 );
1414 }
1415 Err(other) => panic!("expected InvalidSpec, got a different CompileError: {other}"),
1416 Ok(_) => panic!("expected InvalidSpec, got Ok(SpawnerAdapter)"),
1419 }
1420 }
1421
1422 #[tokio::test]
1426 async fn inline_source_evaluates_and_marshals_result() {
1427 let source =
1428 LuaScriptSource::new("return { value = _PROMPT .. '!', ok = true }", "smoke.lua");
1429 let out = run_lua_worker(
1430 std::sync::Arc::new(source),
1431 std::sync::Arc::new(HashMap::new()),
1432 test_invocation("hello"),
1433 )
1434 .await
1435 .expect("lua worker ok");
1436 assert_eq!(out.value, serde_json::json!("hello!"));
1437 assert!(out.ok);
1438 }
1439
1440 #[tokio::test]
1441 async fn inline_source_can_signal_agent_level_failure() {
1442 let source = LuaScriptSource::new("return { value = 'nope', ok = false }", "gate.lua");
1445 let out = run_lua_worker(
1446 std::sync::Arc::new(source),
1447 std::sync::Arc::new(HashMap::new()),
1448 test_invocation("input"),
1449 )
1450 .await
1451 .expect("lua worker ok");
1452 assert_eq!(out.value, serde_json::json!("nope"));
1453 assert!(!out.ok);
1454 }
1455}
1456
1457#[cfg(test)]
1460mod meta_ref_validation_tests {
1461 use super::*;
1462 use crate::blueprint::{AgentMeta, MetaDef};
1463 use crate::worker::adapter::WorkerResult;
1464
1465 fn registry_with_echo() -> SpawnerRegistry {
1466 let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", |inv| async move {
1467 Ok(WorkerResult {
1468 value: Value::String(inv.prompt),
1469 ok: true,
1470 })
1471 });
1472 let mut reg = SpawnerRegistry::new();
1473 reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1474 reg
1475 }
1476
1477 fn rustfn_agent(name: &str) -> AgentDef {
1478 AgentDef {
1479 name: name.to_string(),
1480 kind: AgentKind::RustFn,
1481 spec: serde_json::json!({ "fn_id": "echo" }),
1482 profile: None,
1483 meta: None,
1484 }
1485 }
1486
1487 fn simple_flow(agent_ref: &str, in_: Expr) -> FlowNode {
1488 FlowNode::Step {
1489 ref_: agent_ref.to_string(),
1490 in_,
1491 out: Expr::Path {
1492 at: "$.output".into(),
1493 },
1494 }
1495 }
1496
1497 fn minimal_bp(agents: Vec<AgentDef>, metas: Vec<MetaDef>, flow: FlowNode) -> Blueprint {
1498 Blueprint {
1499 schema_version: crate::blueprint::current_schema_version(),
1500 id: "meta-ref-ut".into(),
1501 flow,
1502 agents,
1503 operators: vec![],
1504 metas,
1505 hints: Default::default(),
1506 strategy: Default::default(),
1507 metadata: BlueprintMetadata::default(),
1508 spawner_hints: Default::default(),
1509 default_agent_kind: AgentKind::Operator,
1510 default_operator_kind: None,
1511 default_init_ctx: None,
1512 default_agent_ctx: None,
1513 default_context_policy: None,
1514 projection_placement: None,
1515 }
1516 }
1517
1518 #[test]
1519 fn valid_meta_ref_compiles() {
1520 let mut agent = rustfn_agent("worker");
1521 agent.meta = Some(AgentMeta {
1522 meta_ref: Some("shared".to_string()),
1523 ..Default::default()
1524 });
1525 let bp = minimal_bp(
1526 vec![agent],
1527 vec![MetaDef {
1528 name: "shared".into(),
1529 ctx: serde_json::json!({ "k": "v" }),
1530 }],
1531 simple_flow(
1532 "worker",
1533 Expr::Path {
1534 at: "$.input".into(),
1535 },
1536 ),
1537 );
1538 let compiler = Compiler::new(registry_with_echo());
1539 assert!(
1540 compiler.compile(&bp).is_ok(),
1541 "a resolvable AgentMeta.meta_ref must compile"
1542 );
1543 }
1544
1545 #[test]
1546 fn unknown_agent_meta_ref_is_unresolved_meta_ref() {
1547 let mut agent = rustfn_agent("worker");
1548 agent.meta = Some(AgentMeta {
1549 meta_ref: Some("missing".to_string()),
1550 ..Default::default()
1551 });
1552 let bp = minimal_bp(
1553 vec![agent],
1554 vec![],
1555 simple_flow(
1556 "worker",
1557 Expr::Path {
1558 at: "$.input".into(),
1559 },
1560 ),
1561 );
1562 let compiler = Compiler::new(registry_with_echo());
1563 match compiler.compile(&bp) {
1564 Err(CompileError::UnresolvedMetaRef {
1565 where_,
1566 meta_ref,
1567 defined,
1568 }) => {
1569 assert!(
1570 where_.contains("worker"),
1571 "where_ must name the agent: {where_}"
1572 );
1573 assert_eq!(meta_ref, "missing");
1574 assert!(defined.is_empty());
1575 }
1576 Err(other) => {
1577 panic!("expected UnresolvedMetaRef, got a different CompileError: {other}")
1578 }
1579 Ok(_) => panic!("expected compile-time failure, got Ok"),
1580 }
1581 }
1582
1583 #[test]
1584 fn unknown_static_step_meta_ref_in_lit_is_unresolved_meta_ref() {
1585 let agent = rustfn_agent("worker");
1586 let in_ = Expr::Lit {
1587 value: serde_json::json!({ "$step_meta": { "ref": "missing" }, "$in": "go" }),
1588 };
1589 let bp = minimal_bp(vec![agent], vec![], simple_flow("worker", in_));
1590 let compiler = Compiler::new(registry_with_echo());
1591 match compiler.compile(&bp) {
1592 Err(CompileError::UnresolvedMetaRef {
1593 where_, meta_ref, ..
1594 }) => {
1595 assert!(
1596 where_.contains("worker"),
1597 "where_ must name the offending step: {where_}"
1598 );
1599 assert_eq!(meta_ref, "missing");
1600 }
1601 Err(other) => {
1602 panic!("expected UnresolvedMetaRef, got a different CompileError: {other}")
1603 }
1604 Ok(_) => panic!("expected compile-time failure, got Ok"),
1605 }
1606 }
1607
1608 #[test]
1609 fn path_op_input_with_no_static_envelope_compiles_fine() {
1610 let agent = rustfn_agent("worker");
1611 let bp = minimal_bp(
1612 vec![agent],
1613 vec![],
1614 simple_flow(
1615 "worker",
1616 Expr::Path {
1617 at: "$.input".into(),
1618 },
1619 ),
1620 );
1621 let compiler = Compiler::new(registry_with_echo());
1622 assert!(
1623 compiler.compile(&bp).is_ok(),
1624 "a non-Lit Step.in must not trigger the best-effort static $step_meta check"
1625 );
1626 }
1627}
1628
1629#[cfg(test)]
1632mod projection_placement_compile_tests {
1633 use super::*;
1634 use crate::core::projection_placement::{ProjectionPlacement, RootPreference};
1635 use crate::worker::adapter::WorkerResult;
1636 use mlua_swarm_schema::ProjectionPlacementSpec;
1637
1638 fn registry_with_echo() -> SpawnerRegistry {
1639 let factory = RustFnInProcessSpawnerFactory::new().register_fn("echo", |inv| async move {
1640 Ok(WorkerResult {
1641 value: Value::String(inv.prompt),
1642 ok: true,
1643 })
1644 });
1645 let mut reg = SpawnerRegistry::new();
1646 reg.register::<RustFnInProcessSpawnerFactory>(Arc::new(factory));
1647 reg
1648 }
1649
1650 fn minimal_bp(projection_placement: Option<ProjectionPlacementSpec>) -> Blueprint {
1651 Blueprint {
1652 schema_version: crate::blueprint::current_schema_version(),
1653 id: "projection-placement-ut".into(),
1654 flow: FlowNode::Step {
1655 ref_: "worker".to_string(),
1656 in_: Expr::Path {
1657 at: "$.input".into(),
1658 },
1659 out: Expr::Path {
1660 at: "$.output".into(),
1661 },
1662 },
1663 agents: vec![AgentDef {
1664 name: "worker".to_string(),
1665 kind: AgentKind::RustFn,
1666 spec: serde_json::json!({ "fn_id": "echo" }),
1667 profile: None,
1668 meta: None,
1669 }],
1670 operators: vec![],
1671 metas: vec![],
1672 hints: Default::default(),
1673 strategy: Default::default(),
1674 metadata: BlueprintMetadata::default(),
1675 spawner_hints: Default::default(),
1676 default_agent_kind: AgentKind::Operator,
1677 default_operator_kind: None,
1678 default_init_ctx: None,
1679 default_agent_ctx: None,
1680 default_context_policy: None,
1681 projection_placement,
1682 }
1683 }
1684
1685 #[test]
1686 fn undeclared_projection_placement_compiles_to_byte_compat_default() {
1687 let bp = minimal_bp(None);
1688 let compiled = Compiler::new(registry_with_echo())
1689 .compile(&bp)
1690 .expect("undeclared projection_placement compiles");
1691 assert_eq!(
1692 *compiled.projection_placement,
1693 ProjectionPlacement::default()
1694 );
1695 }
1696
1697 #[test]
1698 fn declared_valid_projection_placement_compiles_to_matching_resolver() {
1699 let bp = minimal_bp(Some(ProjectionPlacementSpec {
1700 root: Some("project_root".to_string()),
1701 dir_template: Some("custom/{task_id}/out".to_string()),
1702 }));
1703 let compiled = Compiler::new(registry_with_echo())
1704 .compile(&bp)
1705 .expect("valid projection_placement compiles");
1706 assert_eq!(
1707 compiled.projection_placement.root_preference,
1708 RootPreference::ProjectRoot
1709 );
1710 assert_eq!(
1711 compiled.projection_placement.dir_template,
1712 "custom/{task_id}/out"
1713 );
1714 }
1715
1716 #[test]
1717 fn declared_invalid_dir_template_rejects_compile() {
1718 let bp = minimal_bp(Some(ProjectionPlacementSpec {
1719 root: None,
1720 dir_template: Some("workspace/tasks/ctx".to_string()), }));
1722 match Compiler::new(registry_with_echo()).compile(&bp) {
1723 Err(CompileError::InvalidProjectionPlacement(_)) => {}
1724 Err(other) => {
1725 panic!("expected InvalidProjectionPlacement, got a different CompileError: {other}")
1726 }
1727 Ok(_) => {
1728 panic!("expected compile-time rejection for a missing {{task_id}} placeholder")
1729 }
1730 }
1731 }
1732
1733 #[test]
1734 fn declared_invalid_root_literal_rejects_compile() {
1735 let bp = minimal_bp(Some(ProjectionPlacementSpec {
1736 root: Some("nope".to_string()),
1737 dir_template: None,
1738 }));
1739 match Compiler::new(registry_with_echo()).compile(&bp) {
1740 Err(CompileError::InvalidProjectionPlacement(_)) => {}
1741 Err(other) => {
1742 panic!("expected InvalidProjectionPlacement, got a different CompileError: {other}")
1743 }
1744 Ok(_) => panic!("expected compile-time rejection for an invalid root literal"),
1745 }
1746 }
1747}