1pub mod graph_handler;
23pub mod journal;
24pub mod sleep_handler;
25
26pub use graph_handler::GraphHandler;
27pub use journal::{EffectJournal, EffectSite};
28pub use sleep_handler::SleepHandler;
29
30use crate::event_bus::EventBus;
31use somatize_core::effect::{Effect, EffectResult, Usage};
32use somatize_core::error::{Result, SomaError};
33use somatize_core::event::Event;
34use somatize_core::step::{Step, StepCtx, Transition};
35use somatize_core::value::Value;
36use std::sync::Arc;
37use std::time::Instant;
38
39fn suspension_effect(reason: &somatize_core::effect::SuspendReason) -> Effect {
45 Effect::Custom {
46 kind: "soma.suspend".into(),
47 payload: Value::json(serde_json::to_value(reason).unwrap_or(serde_json::Value::Null)),
48 }
49}
50
51pub use somatize_core::effect::EffectHandler;
52
53pub use somatize_core::node::NodeOutcome;
54
55#[derive(Clone)]
57pub struct EffectDriver {
58 handlers: Vec<Arc<dyn EffectHandler>>,
59 journal: EffectJournal,
60 event_bus: Option<Arc<EventBus>>,
61 catalog: Option<Arc<crate::node_catalog::NodeCatalog>>,
64}
65
66impl EffectDriver {
67 pub fn new(journal: EffectJournal) -> Self {
70 Self {
71 handlers: Vec::new(),
72 journal,
73 event_bus: None,
74 catalog: None,
75 }
76 }
77
78 pub fn with_catalog(mut self, catalog: Arc<crate::node_catalog::NodeCatalog>) -> Self {
80 self.catalog = Some(catalog);
81 self
82 }
83
84 pub fn with_handler(mut self, handler: Arc<dyn EffectHandler>) -> Self {
86 self.handlers.push(handler);
87 self
88 }
89
90 pub fn with_event_bus(mut self, bus: Arc<EventBus>) -> Self {
92 self.event_bus = Some(bus);
93 self
94 }
95
96 fn emit(&self, event: Event) {
97 if let Some(bus) = &self.event_bus {
98 bus.emit(event);
99 }
100 }
101
102 pub fn run(
108 &self,
109 step: &dyn Step,
110 run_id: &str,
111 node_id: &str,
112 input: &Value,
113 ) -> Result<NodeOutcome> {
114 let meta = step.meta();
115 let journal = self
117 .journal
118 .clone()
119 .with_enabled(self.journal.is_enabled() && meta.journal);
120
121 let started = Instant::now();
122 let mut history: Vec<Vec<EffectResult>> = Vec::new();
125 let mut usage = Usage::default();
126
127 for turn in 0..meta.max_turns {
128 self.emit(Event::AgentTurnStarted {
129 run_id: run_id.to_string(),
130 node_id: node_id.to_string(),
131 turn,
132 });
133
134 let ctx = StepCtx::new(node_id, run_id, input, turn).with_history(&history);
135 let transition = match step.poll(&ctx) {
139 Ok(transition) => transition,
140 Err(e) => {
141 self.finish(run_id, node_id, turn + 1, started, usage, true);
142 return Err(e);
143 }
144 };
145
146 match transition {
147 Transition::Await(effects) => {
148 if effects.is_empty() {
149 self.finish(run_id, node_id, turn + 1, started, usage, true);
150 return Err(SomaError::Execution {
151 node_id: node_id.to_string(),
152 message: format!(
153 "step awaited nothing on turn {turn}; it would spin. \
154 Return `Done` to finish, or ask for at least one effect"
155 ),
156 });
157 }
158 match self.perform_all(&journal, run_id, node_id, turn, &effects, &mut usage) {
159 Ok(results) => history.push(results),
160 Err(e) => {
161 self.finish(run_id, node_id, turn + 1, started, usage, true);
162 return Err(e);
163 }
164 }
165 }
166
167 Transition::Done(value) => {
168 self.finish(run_id, node_id, turn + 1, started, usage, false);
169 return Ok(NodeOutcome::Produced(value));
170 }
171
172 Transition::Goto { target, carry } => {
173 self.emit(Event::Handoff {
174 run_id: run_id.to_string(),
175 from: node_id.to_string(),
176 to: target.clone(),
177 });
178 self.finish(run_id, node_id, turn + 1, started, usage, false);
179 return Ok(NodeOutcome::HandOff { target, carry });
180 }
181
182 Transition::Suspend { reason } => {
188 let site = EffectSite {
189 run_id,
190 node_id,
191 turn,
192 index: 0,
193 };
194 let effect = suspension_effect(&reason);
195
196 if let Some(answered) = journal.lookup(site, &effect)? {
197 self.emit(Event::Resumed {
198 run_id: run_id.to_string(),
199 node_id: node_id.to_string(),
200 turn,
201 });
202 history.push(vec![answered]);
203 continue;
204 }
205
206 self.emit(Event::Suspended {
207 run_id: run_id.to_string(),
208 node_id: node_id.to_string(),
209 reason: reason.kind().to_string(),
210 turns: turn + 1,
211 duration: started.elapsed(),
212 input_tokens: usage.input_tokens,
213 output_tokens: usage.output_tokens,
214 });
215 return Ok(NodeOutcome::Paused { turn, reason });
216 }
217
218 Transition::Spawn { specs, join } => {
219 if specs.is_empty() {
220 self.finish(run_id, node_id, turn + 1, started, usage, true);
221 return Err(SomaError::Execution {
222 node_id: node_id.to_string(),
223 message: format!(
224 "step spawned nothing on turn {turn}; it would spin. \
225 Return `Done` when there is no work to fan out"
226 ),
227 });
228 }
229 match self.spawn_all(run_id, node_id, turn, &specs, join) {
230 Ok(results) => history.push(results),
231 Err(e) => {
232 self.finish(run_id, node_id, turn + 1, started, usage, true);
233 return Err(e);
234 }
235 }
236 }
237 }
238 }
239
240 self.finish(run_id, node_id, meta.max_turns, started, usage, true);
241 Err(SomaError::Execution {
242 node_id: node_id.to_string(),
243 message: format!(
244 "step did not finish within {} turns. Raise `StepMeta::max_turns` if the \
245 work genuinely needs more, or check whether it is looping",
246 meta.max_turns
247 ),
248 })
249 }
250
251 fn finish(
252 &self,
253 run_id: &str,
254 node_id: &str,
255 turns: usize,
256 started: Instant,
257 usage: Usage,
258 failed: bool,
259 ) {
260 self.emit(Event::AgentStepCompleted {
261 run_id: run_id.to_string(),
262 node_id: node_id.to_string(),
263 turns,
264 duration: started.elapsed(),
265 input_tokens: usage.input_tokens,
266 output_tokens: usage.output_tokens,
267 failed,
268 });
269 }
270
271 pub fn resume_with(
280 &self,
281 run_id: &str,
282 node_id: &str,
283 turn: usize,
284 reason: &somatize_core::effect::SuspendReason,
285 answer: Value,
286 ) -> Result<()> {
287 if !self.journal.is_enabled() {
288 return Err(SomaError::Execution {
289 node_id: node_id.to_string(),
290 message: "cannot resume a run whose journal is disabled: there is \
291 nothing to replay up to the suspension point"
292 .into(),
293 });
294 }
295 let site = EffectSite {
296 run_id,
297 node_id,
298 turn,
299 index: 0,
300 };
301 self.journal.record(
302 site,
303 &suspension_effect(reason),
304 &EffectResult::Node(answer),
305 0,
306 )
307 }
308
309 fn spawn_all(
317 &self,
318 run_id: &str,
319 node_id: &str,
320 turn: usize,
321 specs: &[somatize_core::effect::NodeSpec],
322 join: somatize_core::effect::JoinPolicy,
323 ) -> Result<Vec<EffectResult>> {
324 use somatize_core::effect::JoinPolicy;
325
326 let catalog = self.catalog.as_ref().ok_or_else(|| SomaError::Execution {
327 node_id: node_id.to_string(),
328 message: "step spawned work, but the driver has no step library; \
329 build it with `EffectDriver::with_catalog(...)`"
330 .into(),
331 })?;
332
333 let child_ids: Vec<String> = specs
336 .iter()
337 .enumerate()
338 .map(|(index, spec)| {
339 let label = spec
340 .label
341 .clone()
342 .unwrap_or_else(|| format!("{turn}.{index}"));
343 format!("{node_id}/{label}")
344 })
345 .collect();
346
347 self.emit(Event::AgentSpawned {
348 run_id: run_id.to_string(),
349 node_id: node_id.to_string(),
350 turn,
351 children: child_ids.clone(),
352 join: join.label().to_string(),
353 });
354
355 let outcomes: Vec<Result<EffectResult>> = std::thread::scope(|scope| {
356 let handles: Vec<_> = specs
357 .iter()
358 .zip(&child_ids)
359 .map(|(spec, child_id)| {
360 let child_id = child_id.clone();
361 scope.spawn(move || {
362 let step = catalog
363 .step(&spec.runs)
364 .ok_or_else(|| SomaError::NodeNotFound(spec.runs.clone()))?;
365 match self.run(step.as_ref(), run_id, &child_id, &spec.input)? {
366 NodeOutcome::Produced(value) => Ok(EffectResult::Node(value)),
367 NodeOutcome::HandOff { target, .. } => Err(SomaError::Execution {
368 node_id: child_id.clone(),
369 message: format!(
370 "a spawned step handed control to `{target}`; spawned \
371 work must finish with `Done`, since it has no place \
372 in the graph to hand control to"
373 ),
374 }),
375 NodeOutcome::Paused { .. } => Err(SomaError::Execution {
376 node_id: child_id.clone(),
377 message: "a spawned step suspended; suspension is only \
378 supported for nodes in the graph"
379 .into(),
380 }),
381 }
382 })
383 })
384 .collect();
385
386 handles
387 .into_iter()
388 .map(|h| {
389 h.join().unwrap_or_else(|_| {
390 Err(SomaError::Execution {
391 node_id: node_id.to_string(),
392 message: "a spawned step panicked".into(),
393 })
394 })
395 })
396 .collect()
397 });
398
399 match join {
400 JoinPolicy::All => outcomes.into_iter().collect(),
402
403 JoinPolicy::AllSettled => Ok(outcomes
406 .into_iter()
407 .map(|o| match o {
408 Ok(result) => result,
409 Err(e) => EffectResult::Failed {
410 message: e.to_string(),
411 },
412 })
413 .collect()),
414
415 JoinPolicy::First => {
418 let mut last_error = None;
419 for outcome in outcomes {
420 match outcome {
421 Ok(result) => return Ok(vec![result]),
422 Err(e) => last_error = Some(e),
423 }
424 }
425 Err(last_error.unwrap_or_else(|| SomaError::Execution {
426 node_id: node_id.to_string(),
427 message: "no spawned step succeeded".into(),
428 }))
429 }
430
431 _ => Err(SomaError::Execution {
432 node_id: node_id.to_string(),
433 message: format!("unsupported join policy {join:?}"),
434 }),
435 }
436 }
437
438 fn perform_all(
441 &self,
442 journal: &EffectJournal,
443 run_id: &str,
444 node_id: &str,
445 turn: usize,
446 effects: &[Effect],
447 usage: &mut Usage,
448 ) -> Result<Vec<EffectResult>> {
449 for effect in effects {
450 self.emit(Event::EffectRequested {
451 run_id: run_id.to_string(),
452 node_id: node_id.to_string(),
453 turn,
454 effect: effect.label(),
455 });
456 }
457
458 let outcomes: Vec<Result<(EffectResult, bool, std::time::Duration)>> =
459 std::thread::scope(|scope| {
460 let handles: Vec<_> = effects
461 .iter()
462 .enumerate()
463 .map(|(index, effect)| {
464 let site = EffectSite {
465 run_id,
466 node_id,
467 turn,
468 index,
469 };
470 scope.spawn(move || self.perform_one(journal, site, effect))
471 })
472 .collect();
473
474 handles
475 .into_iter()
476 .map(|h| {
477 h.join().unwrap_or_else(|_| {
478 Err(SomaError::Execution {
479 node_id: node_id.to_string(),
480 message: "effect handler panicked".into(),
481 })
482 })
483 })
484 .collect()
485 });
486
487 let mut results = Vec::with_capacity(effects.len());
488 for (effect, outcome) in effects.iter().zip(outcomes) {
489 let (result, replayed, elapsed) = outcome?;
490
491 if let EffectResult::Llm(response) = &result {
492 *usage += response.usage;
493 }
494 if let Effect::Tool { name, .. } = effect {
495 self.emit(Event::ToolCalled {
496 run_id: run_id.to_string(),
497 node_id: node_id.to_string(),
498 tool: name.clone(),
499 is_error: result.is_error(),
500 });
501 }
502
503 self.emit(Event::EffectCompleted {
504 run_id: run_id.to_string(),
505 node_id: node_id.to_string(),
506 turn,
507 effect: effect.label(),
508 duration: elapsed,
509 replayed,
510 is_error: result.is_error(),
511 });
512
513 results.push(result);
514 }
515 Ok(results)
516 }
517
518 fn perform_one(
520 &self,
521 journal: &EffectJournal,
522 site: EffectSite<'_>,
523 effect: &Effect,
524 ) -> Result<(EffectResult, bool, std::time::Duration)> {
525 let started = Instant::now();
526
527 if let Some(recorded) = journal.lookup(site, effect)? {
528 return Ok((recorded, true, started.elapsed()));
529 }
530
531 let handler = self
532 .handlers
533 .iter()
534 .find(|h| h.handles(effect))
535 .ok_or_else(|| SomaError::Execution {
536 node_id: site.node_id.to_string(),
537 message: format!(
538 "no handler for effect `{}`. Register one on the driver",
539 effect.label()
540 ),
541 })?;
542
543 let result = handler.perform(effect)?;
544 let elapsed = started.elapsed();
545 journal.record(site, effect, &result, elapsed.as_millis() as u64)?;
546 Ok((result, false, elapsed))
547 }
548}
549
550#[cfg(test)]
551mod tests {
552 use super::*;
553 use crate::cache::fs_store::FsActionStore;
554 use somatize_core::cache::CacheKey;
555 use somatize_core::effect::{LlmRequest, LlmResponse, StopReason};
556 use somatize_core::message::Message;
557 use somatize_core::step::StepMeta;
558 use std::sync::atomic::{AtomicUsize, Ordering};
559
560 struct CountingLlm {
562 calls: AtomicUsize,
563 reply: String,
564 }
565
566 impl CountingLlm {
567 fn new(reply: &str) -> Arc<Self> {
568 Arc::new(Self {
569 calls: AtomicUsize::new(0),
570 reply: reply.to_string(),
571 })
572 }
573 }
574
575 impl EffectHandler for CountingLlm {
576 fn handles(&self, effect: &Effect) -> bool {
577 matches!(effect, Effect::Llm(_))
578 }
579 fn perform(&self, _effect: &Effect) -> Result<EffectResult> {
580 self.calls.fetch_add(1, Ordering::SeqCst);
581 Ok(EffectResult::Llm(LlmResponse {
582 message: Message::assistant(&self.reply),
583 stop_reason: StopReason::EndTurn,
584 usage: Usage {
585 input_tokens: 10,
586 output_tokens: 3,
587 ..Default::default()
588 },
589 model: None,
590 }))
591 }
592 }
593
594 struct MultiTurn {
596 rounds: usize,
597 }
598
599 impl Step for MultiTurn {
600 fn config_hash(&self) -> CacheKey {
601 CacheKey::from_parts(&[b"MultiTurn"])
602 }
603 fn meta(&self) -> StepMeta {
604 StepMeta::new("MultiTurn")
605 }
606 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
607 if ctx.turn < self.rounds {
608 return Ok(Transition::Await(vec![Effect::Llm(LlmRequest::new(
609 "claude-opus-5",
610 vec![Message::user(format!("turn {}", ctx.turn))].into(),
611 ))]));
612 }
613 let text = match ctx.result() {
614 Some(EffectResult::Llm(r)) => r.message.text(),
615 _ => String::new(),
616 };
617 Ok(Transition::Done(Value::text(text)))
618 }
619 }
620
621 fn driver(handler: Arc<dyn EffectHandler>) -> (EffectDriver, tempfile::TempDir) {
622 let dir = tempfile::tempdir().unwrap();
623 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
624 let journal = EffectJournal::new(store.clone(), store);
625 (EffectDriver::new(journal).with_handler(handler), dir)
626 }
627
628 #[test]
629 fn runs_a_multi_turn_step() {
630 let llm = CountingLlm::new("hello");
631 let (d, _dir) = driver(llm.clone());
632
633 let out = d
634 .run(&MultiTurn { rounds: 3 }, "r1", "agent", &Value::Empty)
635 .unwrap();
636
637 match out {
638 NodeOutcome::Produced(v) => assert_eq!(v.as_text(), Some("hello")),
639 other => panic!("expected Done, got {other:?}"),
640 }
641 assert_eq!(llm.calls.load(Ordering::SeqCst), 3);
642 }
643
644 #[test]
647 fn replaying_a_run_performs_nothing() {
648 let llm = CountingLlm::new("recorded answer");
649 let dir = tempfile::tempdir().unwrap();
650 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
651 let journal = EffectJournal::new(store.clone(), store);
652
653 let d = EffectDriver::new(journal).with_handler(llm.clone());
654
655 let first = d
656 .run(&MultiTurn { rounds: 3 }, "run-A", "agent", &Value::Empty)
657 .unwrap();
658 assert_eq!(llm.calls.load(Ordering::SeqCst), 3);
659
660 let second = d
662 .run(&MultiTurn { rounds: 3 }, "run-A", "agent", &Value::Empty)
663 .unwrap();
664
665 assert_eq!(
666 llm.calls.load(Ordering::SeqCst),
667 3,
668 "a replay called the model again"
669 );
670 match (first, second) {
671 (NodeOutcome::Produced(a), NodeOutcome::Produced(b)) => assert_eq!(a, b),
672 other => panic!("expected two Done outcomes, got {other:?}"),
673 }
674 }
675
676 #[test]
678 fn a_fresh_run_calls_the_model() {
679 let llm = CountingLlm::new("x");
680 let (d, _dir) = driver(llm.clone());
681
682 d.run(&MultiTurn { rounds: 2 }, "run-A", "agent", &Value::Empty)
683 .unwrap();
684 d.run(&MultiTurn { rounds: 2 }, "run-B", "agent", &Value::Empty)
685 .unwrap();
686
687 assert_eq!(llm.calls.load(Ordering::SeqCst), 4);
688 }
689
690 #[test]
693 fn concurrent_effects_keep_request_order() {
694 struct Echo;
695 impl EffectHandler for Echo {
696 fn handles(&self, e: &Effect) -> bool {
697 matches!(e, Effect::Tool { .. })
698 }
699 fn perform(&self, e: &Effect) -> Result<EffectResult> {
700 let Effect::Tool { args, .. } = e else {
701 unreachable!()
702 };
703 let n = args.as_text().unwrap_or("0").parse::<u64>().unwrap_or(0);
706 std::thread::sleep(std::time::Duration::from_millis(30 - n * 10));
707 Ok(EffectResult::Tool {
708 output: args.clone(),
709 is_error: false,
710 })
711 }
712 }
713
714 struct FanOut;
715 impl Step for FanOut {
716 fn config_hash(&self) -> CacheKey {
717 CacheKey::from_parts(&[b"FanOut"])
718 }
719 fn meta(&self) -> StepMeta {
720 StepMeta::new("FanOut")
721 }
722 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
723 if ctx.turn == 0 {
724 return Ok(Transition::Await(
725 (0..3)
726 .map(|i| Effect::Tool {
727 name: "echo".into(),
728 args: Value::text(i.to_string()),
729 })
730 .collect(),
731 ));
732 }
733 let joined: Vec<String> = ctx
734 .results
735 .iter()
736 .filter_map(|r| r.value().and_then(|v| v.as_text()).map(String::from))
737 .collect();
738 Ok(Transition::Done(Value::text(joined.join(","))))
739 }
740 }
741
742 let (d, _dir) = driver(Arc::new(Echo));
743 match d.run(&FanOut, "r", "n", &Value::Empty).unwrap() {
744 NodeOutcome::Produced(v) => assert_eq!(v.as_text(), Some("0,1,2")),
745 other => panic!("{other:?}"),
746 }
747 }
748
749 #[test]
751 fn a_runaway_step_is_capped() {
752 struct Forever;
753 impl Step for Forever {
754 fn config_hash(&self) -> CacheKey {
755 CacheKey::from_parts(&[b"Forever"])
756 }
757 fn meta(&self) -> StepMeta {
758 StepMeta::new("Forever").with_max_turns(3)
759 }
760 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
761 Ok(Transition::Await(vec![Effect::Llm(LlmRequest::new(
762 "claude-opus-5",
763 vec![Message::user(format!("{}", ctx.turn))].into(),
764 ))]))
765 }
766 }
767
768 let llm = CountingLlm::new("x");
769 let (d, _dir) = driver(llm.clone());
770 let err = d.run(&Forever, "r", "n", &Value::Empty).unwrap_err();
771
772 assert!(err.to_string().contains("max_turns"), "{err}");
773 assert_eq!(llm.calls.load(Ordering::SeqCst), 3, "ran past the cap");
774 }
775
776 #[test]
780 fn a_capped_step_still_reports_its_cost() {
781 struct Forever;
782 impl Step for Forever {
783 fn config_hash(&self) -> CacheKey {
784 CacheKey::from_parts(&[b"Forever"])
785 }
786 fn meta(&self) -> StepMeta {
787 StepMeta::new("Forever").with_max_turns(3)
788 }
789 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
790 Ok(Transition::Await(vec![Effect::Llm(LlmRequest::new(
791 "claude-opus-5",
792 vec![Message::user(format!("{}", ctx.turn))].into(),
793 ))]))
794 }
795 }
796
797 let bus = Arc::new(EventBus::new(64));
798 let mut rx = bus.subscribe();
799 let (d, _dir) = driver(CountingLlm::new("x"));
800 let d = d.with_event_bus(bus);
801
802 d.run(&Forever, "r", "n", &Value::Empty).unwrap_err();
803
804 let mut completed = None;
805 while let Ok(event) = rx.try_recv() {
806 if let Event::AgentStepCompleted {
807 turns,
808 output_tokens,
809 failed,
810 ..
811 } = event
812 {
813 completed = Some((turns, output_tokens, failed));
814 }
815 }
816 let (turns, output_tokens, failed) =
817 completed.expect("no AgentStepCompleted for the capped step");
818 assert!(failed, "turn exhaustion is a failure, not a completion");
819 assert_eq!(turns, 3);
820 assert!(output_tokens > 0, "the tokens it burned went uncounted");
821 }
822
823 #[test]
825 fn an_unhandled_effect_says_so() {
826 let (d, _dir) = driver(CountingLlm::new("x"));
827 struct WantsTool;
828 impl Step for WantsTool {
829 fn config_hash(&self) -> CacheKey {
830 CacheKey::from_parts(&[b"WantsTool"])
831 }
832 fn meta(&self) -> StepMeta {
833 StepMeta::new("WantsTool")
834 }
835 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
836 Ok(Transition::Await(vec![Effect::Tool {
837 name: "search".into(),
838 args: Value::Empty,
839 }]))
840 }
841 }
842
843 let err = d.run(&WantsTool, "r", "n", &Value::Empty).unwrap_err();
844 assert!(err.to_string().contains("tool:search"), "{err}");
845 assert!(err.to_string().contains("no handler"), "{err}");
846 }
847
848 #[test]
850 fn awaiting_nothing_is_an_error() {
851 struct Empty;
852 impl Step for Empty {
853 fn config_hash(&self) -> CacheKey {
854 CacheKey::from_parts(&[b"Empty"])
855 }
856 fn meta(&self) -> StepMeta {
857 StepMeta::new("Empty")
858 }
859 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
860 Ok(Transition::Await(vec![]))
861 }
862 }
863 let (d, _dir) = driver(CountingLlm::new("x"));
864 let err = d.run(&Empty, "r", "n", &Value::Empty).unwrap_err();
865 assert!(err.to_string().contains("awaited nothing"), "{err}");
866 }
867
868 struct Worker;
872 impl Step for Worker {
873 fn config_hash(&self) -> CacheKey {
874 CacheKey::from_parts(&[b"Worker"])
875 }
876 fn meta(&self) -> StepMeta {
877 StepMeta::new("Worker")
878 }
879 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
880 Ok(Transition::Done(Value::text(
881 ctx.input.as_text().unwrap_or_default().to_uppercase(),
882 )))
883 }
884 }
885
886 struct Orchestrator {
890 join: somatize_core::effect::JoinPolicy,
891 }
892
893 impl Step for Orchestrator {
894 fn config_hash(&self) -> CacheKey {
895 CacheKey::from_parts(&[b"Orchestrator"])
896 }
897 fn meta(&self) -> StepMeta {
898 StepMeta::new("Orchestrator")
899 }
900 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
901 if ctx.turn == 0 {
902 let specs = ctx
903 .input
904 .as_text()
905 .unwrap_or_default()
906 .split(',')
907 .enumerate()
908 .map(|(i, part)| {
909 somatize_core::effect::NodeSpec::new("worker", Value::text(part))
910 .with_label(format!("w{i}"))
911 })
912 .collect();
913 return Ok(Transition::Spawn {
914 specs,
915 join: self.join,
916 });
917 }
918 let joined: Vec<String> = ctx
919 .results
920 .iter()
921 .map(|r| match r {
922 EffectResult::Node(v) => v.as_text().unwrap_or_default().to_string(),
923 EffectResult::Failed { message } => format!("<{message}>"),
924 other => format!("<unexpected {other:?}>"),
925 })
926 .collect();
927 Ok(Transition::Done(Value::text(joined.join("|"))))
928 }
929 }
930
931 fn spawning_driver(
932 join: somatize_core::effect::JoinPolicy,
933 ) -> (EffectDriver, tempfile::TempDir) {
934 let dir = tempfile::tempdir().unwrap();
935 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
936 let journal = EffectJournal::new(store.clone(), store);
937
938 let mut steps = crate::node_catalog::NodeCatalog::new();
939 steps.register_step("worker", Box::new(Worker));
940 steps.register_step("orchestrator", Box::new(Orchestrator { join }));
941
942 (
943 EffectDriver::new(journal).with_catalog(Arc::new(steps)),
944 dir,
945 )
946 }
947
948 #[test]
949 fn spawns_a_worker_per_item_and_joins_in_order() {
950 use somatize_core::effect::JoinPolicy;
951
952 let (d, _dir) = spawning_driver(JoinPolicy::All);
953 let out = d
954 .run(
955 &Orchestrator {
956 join: JoinPolicy::All,
957 },
958 "r",
959 "orch",
960 &Value::text("alpha,beta,gamma"),
961 )
962 .unwrap();
963
964 match out {
965 NodeOutcome::Produced(v) => assert_eq!(v.as_text(), Some("ALPHA|BETA|GAMMA")),
966 other => panic!("{other:?}"),
967 }
968 }
969
970 #[test]
973 fn spawning_emits_the_fan_out() {
974 use somatize_core::effect::JoinPolicy;
975
976 let bus = Arc::new(EventBus::new(64));
977 let mut rx = bus.subscribe();
978 let (d, _dir) = spawning_driver(JoinPolicy::All);
979 let d = d.with_event_bus(bus);
980
981 d.run(
982 &Orchestrator {
983 join: JoinPolicy::All,
984 },
985 "r",
986 "orch",
987 &Value::text("alpha,beta"),
988 )
989 .unwrap();
990
991 let mut spawned = None;
992 let mut child_completions = 0;
993 while let Ok(event) = rx.try_recv() {
994 match event {
995 Event::AgentSpawned { children, join, .. } => spawned = Some((children, join)),
996 Event::AgentStepCompleted { node_id, .. } if node_id.contains('/') => {
997 child_completions += 1;
998 }
999 _ => {}
1000 }
1001 }
1002 let (children, join) = spawned.expect("no AgentSpawned event");
1003 assert_eq!(children, vec!["orch/w0".to_string(), "orch/w1".to_string()]);
1004 assert_eq!(join, "all");
1005 assert_eq!(
1006 child_completions, 2,
1007 "each spawned child should report its own completion under its hierarchical id"
1008 );
1009 }
1010
1011 #[test]
1015 fn spawned_siblings_journal_separately() {
1016 use somatize_core::effect::JoinPolicy;
1017
1018 let (d, _dir) = spawning_driver(JoinPolicy::All);
1019 let orch = Orchestrator {
1020 join: JoinPolicy::All,
1021 };
1022
1023 let first = d.run(&orch, "r", "orch", &Value::text("a,b,c")).unwrap();
1024 let replay = d.run(&orch, "r", "orch", &Value::text("a,b,c")).unwrap();
1025
1026 match (first, replay) {
1027 (NodeOutcome::Produced(a), NodeOutcome::Produced(b)) => {
1028 assert_eq!(a.as_text(), Some("A|B|C"));
1029 assert_eq!(a, b, "replay of a fan-out diverged");
1030 }
1031 other => panic!("{other:?}"),
1032 }
1033 }
1034
1035 #[test]
1037 fn spawning_an_unknown_step_names_it() {
1038 struct BadOrchestrator;
1039 impl Step for BadOrchestrator {
1040 fn config_hash(&self) -> CacheKey {
1041 CacheKey::from_parts(&[b"Bad"])
1042 }
1043 fn meta(&self) -> StepMeta {
1044 StepMeta::new("Bad")
1045 }
1046 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
1047 Ok(Transition::Spawn {
1048 specs: vec![somatize_core::effect::NodeSpec::new(
1049 "nonexistent",
1050 Value::Empty,
1051 )],
1052 join: somatize_core::effect::JoinPolicy::All,
1053 })
1054 }
1055 }
1056
1057 let (d, _dir) = spawning_driver(somatize_core::effect::JoinPolicy::All);
1058 let err = d
1059 .run(&BadOrchestrator, "r", "orch", &Value::Empty)
1060 .unwrap_err();
1061 assert!(err.to_string().contains("nonexistent"), "{err}");
1062 }
1063
1064 #[test]
1066 fn spawning_without_a_library_explains_itself() {
1067 use somatize_core::effect::JoinPolicy;
1068
1069 let dir = tempfile::tempdir().unwrap();
1070 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
1071 let d = EffectDriver::new(EffectJournal::new(store.clone(), store));
1072
1073 let err = d
1074 .run(
1075 &Orchestrator {
1076 join: JoinPolicy::All,
1077 },
1078 "r",
1079 "orch",
1080 &Value::text("a"),
1081 )
1082 .unwrap_err();
1083 assert!(err.to_string().contains("with_catalog"), "{err}");
1084 }
1085
1086 #[test]
1088 fn spawning_nothing_is_an_error() {
1089 use somatize_core::effect::JoinPolicy;
1090
1091 let (d, _dir) = spawning_driver(JoinPolicy::All);
1092
1093 struct SpawnsNothing;
1094 impl Step for SpawnsNothing {
1095 fn config_hash(&self) -> CacheKey {
1096 CacheKey::from_parts(&[b"SpawnsNothing"])
1097 }
1098 fn meta(&self) -> StepMeta {
1099 StepMeta::new("SpawnsNothing")
1100 }
1101 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
1102 Ok(Transition::Spawn {
1103 specs: vec![],
1104 join: JoinPolicy::All,
1105 })
1106 }
1107 }
1108 let err = d
1109 .run(&SpawnsNothing, "r", "orch", &Value::Empty)
1110 .unwrap_err();
1111 assert!(err.to_string().contains("spawned nothing"), "{err}");
1112 }
1113
1114 struct FlakyWorker;
1119 impl Step for FlakyWorker {
1120 fn config_hash(&self) -> CacheKey {
1121 CacheKey::from_parts(&[b"FlakyWorker"])
1122 }
1123 fn meta(&self) -> StepMeta {
1124 StepMeta::new("FlakyWorker")
1125 }
1126 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
1127 let text = ctx.input.as_text().unwrap_or_default();
1128 if text == "bad" {
1129 return Err(SomaError::Execution {
1130 node_id: ctx.node_id.to_string(),
1131 message: "worker refused".into(),
1132 });
1133 }
1134 Ok(Transition::Done(Value::text(text.to_uppercase())))
1135 }
1136 }
1137
1138 fn flaky_driver(join: somatize_core::effect::JoinPolicy) -> (EffectDriver, tempfile::TempDir) {
1140 let dir = tempfile::tempdir().unwrap();
1141 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
1142 let journal = EffectJournal::new(store.clone(), store);
1143
1144 let mut steps = crate::node_catalog::NodeCatalog::new();
1145 steps.register_step("worker", Box::new(FlakyWorker));
1146 steps.register_step("orchestrator", Box::new(Orchestrator { join }));
1147
1148 (
1149 EffectDriver::new(journal).with_catalog(Arc::new(steps)),
1150 dir,
1151 )
1152 }
1153
1154 #[test]
1159 fn all_settled_keeps_what_succeeded() {
1160 use somatize_core::effect::JoinPolicy;
1161
1162 let (d, _dir) = flaky_driver(JoinPolicy::AllSettled);
1163 let out = d
1164 .run(
1165 &Orchestrator {
1166 join: JoinPolicy::AllSettled,
1167 },
1168 "r",
1169 "orch",
1170 &Value::text("ok,bad,fine"),
1171 )
1172 .expect("a failed sibling must not fail the join");
1173
1174 let NodeOutcome::Produced(v) = out else {
1175 panic!("expected Done, got {out:?}");
1176 };
1177 let text = v.as_text().unwrap();
1178 assert!(text.starts_with("OK|"), "first success lost: {text}");
1179 assert!(text.ends_with("|FINE"), "last success lost: {text}");
1180 assert!(
1181 text.contains("worker refused"),
1182 "the failure should be reported in place, not dropped: {text}"
1183 );
1184 }
1185
1186 #[test]
1190 fn first_returns_the_first_answer() {
1191 use somatize_core::effect::JoinPolicy;
1192
1193 let (d, _dir) = flaky_driver(JoinPolicy::First);
1194 let orch = Orchestrator {
1195 join: JoinPolicy::First,
1196 };
1197
1198 match d
1199 .run(&orch, "r1", "orch", &Value::text("alpha,beta"))
1200 .unwrap()
1201 {
1202 NodeOutcome::Produced(v) => assert_eq!(
1203 v.as_text(),
1204 Some("ALPHA"),
1205 "exactly the first answer, alone"
1206 ),
1207 other => panic!("{other:?}"),
1208 }
1209
1210 match d
1212 .run(&orch, "r2", "orch", &Value::text("bad,good"))
1213 .unwrap()
1214 {
1215 NodeOutcome::Produced(v) => assert_eq!(v.as_text(), Some("GOOD")),
1216 other => panic!("{other:?}"),
1217 }
1218 }
1219
1220 struct Defector;
1223 impl Step for Defector {
1224 fn config_hash(&self) -> CacheKey {
1225 CacheKey::from_parts(&[b"Defector"])
1226 }
1227 fn meta(&self) -> StepMeta {
1228 StepMeta::new("Defector")
1229 }
1230 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
1231 Ok(Transition::Goto {
1232 target: "elsewhere".into(),
1233 carry: Value::Empty,
1234 })
1235 }
1236 }
1237
1238 #[test]
1240 fn a_spawned_child_that_hands_off_is_an_error() {
1241 use somatize_core::effect::JoinPolicy;
1242
1243 let dir = tempfile::tempdir().unwrap();
1244 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
1245 let mut steps = crate::node_catalog::NodeCatalog::new();
1246 steps.register_step("worker", Box::new(Defector));
1247 steps.register_step(
1248 "orchestrator",
1249 Box::new(Orchestrator {
1250 join: JoinPolicy::All,
1251 }),
1252 );
1253 let d = EffectDriver::new(EffectJournal::new(store.clone(), store))
1254 .with_catalog(Arc::new(steps));
1255
1256 let err = d
1257 .run(
1258 &Orchestrator {
1259 join: JoinPolicy::All,
1260 },
1261 "r",
1262 "orch",
1263 &Value::text("x"),
1264 )
1265 .unwrap_err();
1266 let msg = err.to_string();
1267 assert!(msg.contains("elsewhere"), "should name the target: {msg}");
1268 assert!(
1269 msg.contains("must finish with `Done`"),
1270 "should state the contract: {msg}"
1271 );
1272 }
1273
1274 struct PanickingWorker;
1276 impl Step for PanickingWorker {
1277 fn config_hash(&self) -> CacheKey {
1278 CacheKey::from_parts(&[b"PanickingWorker"])
1279 }
1280 fn meta(&self) -> StepMeta {
1281 StepMeta::new("PanickingWorker")
1282 }
1283 fn poll(&self, _ctx: &StepCtx<'_>) -> Result<Transition> {
1284 panic!("the worker fell over");
1285 }
1286 }
1287
1288 #[test]
1291 fn a_spawned_child_that_panics_is_contained() {
1292 use somatize_core::effect::JoinPolicy;
1293
1294 let dir = tempfile::tempdir().unwrap();
1295 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
1296 let mut steps = crate::node_catalog::NodeCatalog::new();
1297 steps.register_step("worker", Box::new(PanickingWorker));
1298 let d = EffectDriver::new(EffectJournal::new(store.clone(), store))
1299 .with_catalog(Arc::new(steps));
1300
1301 let previous = std::panic::take_hook();
1302 std::panic::set_hook(Box::new(|_| {}));
1303 let result = d.run(
1304 &Orchestrator {
1305 join: JoinPolicy::All,
1306 },
1307 "r",
1308 "orch",
1309 &Value::text("x"),
1310 );
1311 std::panic::set_hook(previous);
1312
1313 let err = result.expect_err("a panicking child must surface as an error");
1314 assert!(err.to_string().contains("a spawned step panicked"), "{err}");
1315 }
1316
1317 struct NeedsApproval;
1321
1322 impl Step for NeedsApproval {
1323 fn config_hash(&self) -> CacheKey {
1324 CacheKey::from_parts(&[b"NeedsApproval"])
1325 }
1326 fn meta(&self) -> StepMeta {
1327 StepMeta::new("NeedsApproval")
1328 }
1329 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
1330 match ctx.result() {
1331 None => Ok(Transition::Suspend {
1332 reason: somatize_core::effect::SuspendReason::Human {
1333 prompt: "Approve deleting 3 files?".into(),
1334 schema: None,
1335 },
1336 }),
1337 Some(EffectResult::Node(answer)) => Ok(Transition::Done(Value::text(format!(
1338 "decision: {}",
1339 answer.as_text().unwrap_or("?")
1340 )))),
1341 Some(other) => Ok(Transition::Done(Value::text(format!("odd: {other:?}")))),
1342 }
1343 }
1344 }
1345
1346 fn reason() -> somatize_core::effect::SuspendReason {
1347 somatize_core::effect::SuspendReason::Human {
1348 prompt: "Approve deleting 3 files?".into(),
1349 schema: None,
1350 }
1351 }
1352
1353 #[test]
1357 fn suspends_then_resumes_with_the_answer() {
1358 let (d, _dir) = driver(CountingLlm::new("unused"));
1359
1360 let first = d
1361 .run(&NeedsApproval, "run-hitl", "approve", &Value::Empty)
1362 .unwrap();
1363 let turn = match first {
1364 NodeOutcome::Paused { turn, .. } => turn,
1365 other => panic!("expected a suspension, got {other:?}"),
1366 };
1367 assert_eq!(turn, 0);
1368
1369 d.resume_with("run-hitl", "approve", turn, &reason(), Value::text("yes"))
1370 .unwrap();
1371
1372 match d
1373 .run(&NeedsApproval, "run-hitl", "approve", &Value::Empty)
1374 .unwrap()
1375 {
1376 NodeOutcome::Produced(v) => assert_eq!(v.as_text(), Some("decision: yes")),
1377 other => panic!("expected Done after resuming, got {other:?}"),
1378 }
1379 }
1380
1381 #[test]
1383 fn re_running_without_an_answer_suspends_again() {
1384 let (d, _dir) = driver(CountingLlm::new("unused"));
1385
1386 for _ in 0..2 {
1387 match d
1388 .run(&NeedsApproval, "r", "approve", &Value::Empty)
1389 .unwrap()
1390 {
1391 NodeOutcome::Paused { .. } => {}
1392 other => panic!("expected a suspension, got {other:?}"),
1393 }
1394 }
1395 }
1396
1397 #[test]
1399 fn an_answer_does_not_carry_to_another_run() {
1400 let (d, _dir) = driver(CountingLlm::new("unused"));
1401
1402 d.run(&NeedsApproval, "run-A", "approve", &Value::Empty)
1403 .unwrap();
1404 d.resume_with("run-A", "approve", 0, &reason(), Value::text("yes"))
1405 .unwrap();
1406
1407 match d
1408 .run(&NeedsApproval, "run-B", "approve", &Value::Empty)
1409 .unwrap()
1410 {
1411 NodeOutcome::Paused { .. } => {}
1412 other => panic!("run B reused run A's approval: {other:?}"),
1413 }
1414 }
1415
1416 #[test]
1419 fn resuming_without_a_journal_is_refused() {
1420 let dir = tempfile::tempdir().unwrap();
1421 let store = Arc::new(FsActionStore::new(dir.path()).unwrap());
1422 let d = EffectDriver::new(EffectJournal::disabled(store.clone(), store));
1423
1424 let err = d
1425 .resume_with("r", "approve", 0, &reason(), Value::text("yes"))
1426 .unwrap_err();
1427 assert!(err.to_string().contains("journal is disabled"), "{err}");
1428 }
1429
1430 #[test]
1433 fn a_step_can_decline_journaling() {
1434 struct Private;
1435 impl Step for Private {
1436 fn config_hash(&self) -> CacheKey {
1437 CacheKey::from_parts(&[b"Private"])
1438 }
1439 fn meta(&self) -> StepMeta {
1440 StepMeta::new("Private").without_journal()
1441 }
1442 fn poll(&self, ctx: &StepCtx<'_>) -> Result<Transition> {
1443 if ctx.turn == 0 {
1444 return Ok(Transition::Await(vec![Effect::Llm(LlmRequest::new(
1445 "claude-opus-5",
1446 vec![Message::user("sensitive")].into(),
1447 ))]));
1448 }
1449 Ok(Transition::Done(Value::Empty))
1450 }
1451 }
1452
1453 let llm = CountingLlm::new("x");
1454 let (d, _dir) = driver(llm.clone());
1455 d.run(&Private, "r", "n", &Value::Empty).unwrap();
1456 d.run(&Private, "r", "n", &Value::Empty).unwrap();
1457
1458 assert_eq!(
1459 llm.calls.load(Ordering::SeqCst),
1460 2,
1461 "an un-journaled step was replayed from disk"
1462 );
1463 }
1464}