1use std::collections::VecDeque;
20use std::sync::Arc;
21
22use bevy_ecs::prelude::*;
23use leviath_core::blueprint::{FanOutConfig, StageMode, WorkerFailurePolicy};
24
25use crate::components::{
26 AgentState, AgentStatus, ContextWindow, InferenceResult, ParentRef, SubAgentChildren,
27};
28use crate::pipeline::{AgentBlueprint, ProcessResponse, ResolveTransition, StageCursor};
29
30const DEFAULT_FANOUT_DEPTH: usize = 3;
32
33#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Default)]
35pub struct WorkItem {
36 #[serde(default)]
38 pub id: String,
39 #[serde(default)]
41 pub context: serde_json::Value,
42}
43
44pub fn parse_work_items(content: &str) -> Result<Vec<WorkItem>, String> {
47 let trimmed = content.trim();
48 let slice = match (trimmed.find('['), trimmed.rfind(']')) {
51 (Some(s), Some(e)) if e > s => trimmed.get(s..=e),
52 _ => None,
53 }
54 .ok_or_else(|| "split output is not a JSON array".to_string())?;
55 serde_json::from_str(slice)
56 .map_err(|e| format!("split output is not a valid JSON array of work items: {e}"))
57}
58
59pub trait FanOutSpawner: Send + Sync {
65 fn spawn_worker(
68 &self,
69 world: &mut World,
70 parent: Entity,
71 config: &FanOutConfig,
72 item_id: &str,
73 item_context: &serde_json::Value,
74 ) -> Result<Entity, String>;
75}
76
77#[derive(Resource, Clone)]
80pub struct FanOutSpawnerRes(pub Arc<dyn FanOutSpawner>);
81
82struct ActiveWorker {
86 item_id: String,
87 entity: Entity,
88 run_id: String,
89}
90
91#[derive(Component)]
94pub struct FanOutWaiting {
95 config: FanOutConfig,
96 max_workers: usize,
97 pending: VecDeque<WorkItem>,
98 active: Vec<ActiveWorker>,
99 summaries: Vec<(String, String)>,
100 failures: Vec<(String, String)>,
101}
102
103#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
108pub struct FanOutState {
109 pub config: FanOutConfig,
111 pub max_workers: usize,
113 pub pending: Vec<WorkItem>,
115 pub active: Vec<(String, String)>,
117 pub summaries: Vec<(String, String)>,
119 pub failures: Vec<(String, String)>,
121}
122
123impl FanOutWaiting {
124 pub fn outstanding(&self) -> usize {
129 self.active.len() + self.pending.len()
130 }
131
132 pub(crate) fn to_state(&self) -> FanOutState {
134 FanOutState {
135 config: self.config.clone(),
136 max_workers: self.max_workers,
137 pending: self.pending.iter().cloned().collect(),
138 active: self
139 .active
140 .iter()
141 .map(|w| (w.item_id.clone(), w.run_id.clone()))
142 .collect(),
143 summaries: self.summaries.clone(),
144 failures: self.failures.clone(),
145 }
146 }
147}
148
149pub fn restore_fan_out_waiting(
155 world: &mut World,
156 parent: Entity,
157 state: FanOutState,
158 resolve: &dyn Fn(&str) -> Option<Entity>,
159) {
160 let mut active = Vec::new();
161 let mut failures = state.failures;
162 for (item_id, run_id) in state.active {
163 match resolve(&run_id) {
164 Some(entity) => active.push(ActiveWorker {
165 item_id,
166 entity,
167 run_id,
168 }),
169 None => failures.push((item_id, "worker did not reload after restart".to_string())),
170 }
171 }
172 world.entity_mut(parent).insert(FanOutWaiting {
173 config: state.config,
174 max_workers: state.max_workers,
175 pending: state.pending.into_iter().collect(),
176 active,
177 summaries: state.summaries,
178 failures,
179 });
180}
181
182pub fn fan_out_split(world: &mut World) {
188 crate::tick_scope::clear();
189 let mut candidates: Vec<(Entity, String, FanOutConfig)> = Vec::new();
190 {
191 let mut q = world.query_filtered::<(
192 Entity,
193 &AgentState,
194 &AgentBlueprint,
195 &StageCursor,
196 &InferenceResult,
197 ), With<ProcessResponse>>();
198 for (entity, state, bp, cursor, infer) in q.iter(world) {
199 if state.status != AgentStatus::Active {
200 continue;
201 }
202 if let StageMode::FanOut { config } = &bp.0.stages[cursor.index].mode {
203 candidates.push((entity, infer.response.clone(), config.clone()));
204 }
205 }
206 }
207
208 for (parent, response, config) in candidates {
209 crate::tick_scope::enter(parent);
210 world
211 .entity_mut(parent)
212 .remove::<ProcessResponse>()
213 .remove::<InferenceResult>();
214 match parse_work_items(&response) {
215 Ok(items) => {
216 let max_workers = config.max_workers.max(1);
217 let items = match config.max_items {
223 Some(cap) if items.len() > cap => {
224 tracing::warn!(
225 produced = items.len(),
226 cap,
227 "fan_out split produced more items than max_items; keeping the first"
228 );
229 items.into_iter().take(cap).collect::<Vec<_>>()
230 }
231 _ => items,
232 };
233 world.entity_mut(parent).insert(FanOutWaiting {
234 config,
235 max_workers,
236 pending: items.into_iter().collect(),
237 active: Vec::new(),
238 summaries: Vec::new(),
239 failures: Vec::new(),
240 });
241 set_status(world, parent, AgentStatus::Waiting);
242 }
243 Err(message) => {
244 set_status(
245 world,
246 parent,
247 AgentStatus::Error {
248 message: format!("fan_out split failed: {message}"),
249 },
250 );
251 }
252 }
253 }
254}
255
256pub fn fan_out_collect(world: &mut World) {
261 crate::tick_scope::clear();
262 let parents: Vec<Entity> = {
263 let mut q = world.query_filtered::<Entity, With<FanOutWaiting>>();
264 q.iter(world).collect()
265 };
266
267 for parent in parents {
268 crate::tick_scope::enter(parent);
269 if !matches!(agent_status(world, parent), Some(AgentStatus::Waiting)) {
272 world.entity_mut(parent).remove::<FanOutWaiting>();
273 continue;
274 }
275 let mut w = world
278 .entity_mut(parent)
279 .take::<FanOutWaiting>()
280 .expect("a Waiting fan-out parent still holds FanOutWaiting");
281
282 let mut still_active = Vec::with_capacity(w.active.len());
291 for aw in std::mem::take(&mut w.active) {
292 match worker_terminal_result(world, aw.entity) {
293 Some(result) => {
294 match result {
295 Ok(content) => w.summaries.push((aw.item_id, content)),
296 Err(message) => w.failures.push((aw.item_id, message)),
297 }
298 world.entity_mut(aw.entity).insert(MergedWorker);
299 }
300 None => still_active.push(aw),
301 }
302 }
303 w.active = still_active;
304
305 while w.active.len() < w.max_workers {
307 let Some(item) = w.pending.pop_front() else {
308 break;
309 };
310 match start_worker(world, parent, &w.config, &item) {
311 Ok(child) => {
312 let run_id = world
314 .get::<crate::persistence::RunMetadata>(child)
315 .map(|m| m.run_id.clone())
316 .unwrap_or_default();
317 w.active.push(ActiveWorker {
318 item_id: item.id,
319 entity: child,
320 run_id,
321 });
322 }
323 Err(message) => w.failures.push((item.id, message)),
324 }
325 }
326
327 if w.active.is_empty() && w.pending.is_empty() {
329 finish_fan_out(world, parent, w);
330 } else {
331 world.entity_mut(parent).insert(w);
332 }
333 }
334}
335
336#[derive(Component)]
339pub struct MergedWorker;
340
341pub fn slim_merged_workers(
351 workers: Query<(Entity, &crate::pipeline::PersistWatermark), With<MergedWorker>>,
352 mut commands: Commands,
353) {
354 crate::tick_scope::clear();
355 for (entity, watermark) in workers.iter() {
356 crate::tick_scope::enter(entity);
357 let terminal_persisted = matches!(
358 watermark.persisted_status(),
359 Some(
360 leviath_core::run_meta::RunStatus::Complete
361 | leviath_core::run_meta::RunStatus::Error
362 | leviath_core::run_meta::RunStatus::Cancelled
363 )
364 );
365 if !terminal_persisted {
366 continue; }
368 commands.entity(entity).remove::<(
369 ContextWindow,
370 InferenceResult,
371 crate::pipeline::StageInferences,
372 crate::pipeline::StageSetups,
373 AgentBlueprint,
374 MergedWorker,
375 )>();
376 }
377}
378
379fn finish_fan_out(world: &mut World, parent: Entity, w: FanOutWaiting) {
381 if !w.failures.is_empty() && w.config.on_worker_failure == WorkerFailurePolicy::FailAll {
382 set_status(
383 world,
384 parent,
385 AgentStatus::Error {
386 message: format!(
387 "fan_out: {} worker(s) failed (on_worker_failure = fail_all)",
388 w.failures.len()
389 ),
390 },
391 );
392 return;
393 }
394
395 let region = w
400 .config
401 .results_region
402 .clone()
403 .unwrap_or_else(|| "conversation".to_string());
404 let budget = world
405 .get::<ContextWindow>(parent)
406 .and_then(|window| window.get_region(®ion).map(|r| r.max_tokens));
407 let report = build_report(&w.summaries, &w.failures, budget);
408 inject_results(world, parent, ®ion, &report);
409
410 set_status(world, parent, AgentStatus::Active);
413 match w.config.merge_stage.as_deref().and_then(|name| {
414 world
415 .get::<AgentBlueprint>(parent)
416 .and_then(|bp| bp.0.stages.iter().position(|s| s.name == name))
417 }) {
418 Some(idx) => crate::pipeline::force_transition(
419 world,
420 crate::world::AgentId::in_world(world, parent),
421 idx,
422 ),
423 None => {
424 world.entity_mut(parent).insert(ResolveTransition);
425 }
426 }
427}
428
429fn start_worker(
432 world: &mut World,
433 parent: Entity,
434 config: &FanOutConfig,
435 item: &WorkItem,
436) -> Result<Entity, String> {
437 let max_depth = world
438 .get::<SubAgentChildren>(parent)
439 .map(|k| k.max_child_depth)
440 .or_else(|| {
441 world
442 .get::<AgentBlueprint>(parent)
443 .and_then(|bp| bp.0.max_child_depth)
444 })
445 .unwrap_or(DEFAULT_FANOUT_DEPTH);
446 let parent_depth = world.get::<ParentRef>(parent).map_or(0, |p| p.depth);
447 let child_depth = parent_depth + 1;
448 if child_depth > max_depth {
449 return Err(format!(
450 "fan-out worker depth limit ({max_depth}) reached; not spawning"
451 ));
452 }
453
454 let spawner = world
455 .get_resource::<FanOutSpawnerRes>()
456 .map(|r| r.0.clone())
457 .ok_or_else(|| "no fan-out spawner installed".to_string())?;
458 let child = spawner.spawn_worker(world, parent, config, &item.id, &item.context)?;
459
460 let parent_agent_id = world
461 .get::<AgentState>(parent)
462 .map(|s| s.agent_id.clone())
463 .unwrap_or_default();
464 world.entity_mut(child).insert(ParentRef {
465 parent_entity: parent,
466 parent_agent_id,
467 depth: child_depth,
468 });
469 match world.get_mut::<SubAgentChildren>(parent) {
470 Some(mut kids) => kids.children.push(child),
471 None => {
472 world.entity_mut(parent).insert(SubAgentChildren {
473 children: vec![child],
474 max_child_depth: max_depth,
475 });
476 }
477 }
478 let worker_id = world
482 .get::<crate::persistence::RunMetadata>(child)
483 .expect("a fan-out worker always has run metadata")
484 .run_id
485 .clone();
486 world
487 .get_mut::<AgentState>(parent)
488 .expect("a fan-out parent always has AgentState")
489 .spawned_children_ids
490 .push(worker_id);
491 crate::context_transform::apply_context_transforms(
494 world,
495 crate::world::AgentId::in_world(world, parent),
496 crate::world::AgentId::in_world(world, child),
497 );
498 Ok(child)
499}
500
501fn worker_terminal_result(world: &World, worker: Entity) -> Option<Result<String, String>> {
525 match agent_status(world, worker) {
526 None => Some(Err("worker vanished".to_string())),
527 Some(AgentStatus::Complete) => {
528 match world
529 .get::<crate::persistence::FinalOutput>(worker)
530 .map(|o| o.0.content.clone())
531 {
532 Some(content) => Some(Ok(content)),
533 None if worker_requires_output(world, worker) => Some(Err(
534 "worker finished without the final output its stage requires".to_string(),
535 )),
536 None => Some(Ok(world
537 .get::<InferenceResult>(worker)
538 .map(|r| r.response.clone())
539 .unwrap_or_default())),
540 }
541 }
542 Some(AgentStatus::Error { message }) => Some(Err(message)),
543 Some(AgentStatus::Cancelled) => Some(Err("worker cancelled".to_string())),
544 Some(_) => None,
545 }
546}
547
548fn worker_requires_output(world: &World, worker: Entity) -> bool {
550 let Some(bp) = world.get::<AgentBlueprint>(worker) else {
551 return false;
552 };
553 let Some(cursor) = world.get::<StageCursor>(worker) else {
554 return false;
555 };
556 bp.0.stages
557 .get(cursor.index)
558 .is_some_and(|s| s.require_output)
559}
560
561const MIN_REPORT_BYTES_PER_WORKER: usize = 200;
567
568const DEFAULT_REPORT_BYTES_PER_WORKER: usize = 4_000;
570
571const REPORT_TRUNCATION_MARKER: &str =
573 "\n[...truncated; read this worker's own run for the full answer]";
574
575fn bytes_per_worker(region_budget_tokens: Option<usize>, workers: usize) -> usize {
584 let Some(tokens) = region_budget_tokens.filter(|t| *t > 0) else {
585 return DEFAULT_REPORT_BYTES_PER_WORKER;
586 };
587 let usable = tokens.saturating_mul(4).saturating_mul(9) / 10;
590 (usable / workers.max(1)).max(MIN_REPORT_BYTES_PER_WORKER)
591}
592
593fn fit_worker_section(content: &str, budget: usize) -> String {
595 if content.len() <= budget {
596 return content.to_string();
597 }
598 let room = budget.saturating_sub(REPORT_TRUNCATION_MARKER.len());
599 format!(
600 "{}{REPORT_TRUNCATION_MARKER}",
601 leviath_core::truncate_at_boundary(content, room)
602 )
603}
604
605fn build_report(
610 summaries: &[(String, String)],
611 failures: &[(String, String)],
612 region_budget_tokens: Option<usize>,
613) -> String {
614 let sections = summaries.len().max(1);
615 let budget = bytes_per_worker(region_budget_tokens, sections);
616 let mut report = format!(
617 "[fan_out results: {} succeeded, {} failed]\n",
618 summaries.len(),
619 failures.len()
620 );
621 if summaries.iter().any(|(_, c)| c.len() > budget) {
624 report.push_str(&format!(
625 "[each worker's answer is shown up to {budget} characters; \
626 read a worker's own run for the whole thing]\n"
627 ));
628 }
629 for (id, content) in summaries {
630 report.push_str(&format!(
631 "\n## worker {id}\n{}\n",
632 fit_worker_section(content, budget)
633 ));
634 }
635 for (id, err) in failures {
636 report.push_str(&format!("\n## worker {id} FAILED\n{err}\n"));
637 }
638 report
639}
640
641fn inject_results(world: &mut World, parent: Entity, region: &str, text: &str) {
649 let Some(mut window) = world.get_mut::<ContextWindow>(parent) else {
650 return;
651 };
652 let region = match window.get_region(region).is_some() {
656 true => region,
657 false => {
658 tracing::warn!(
659 region = %region,
660 "fan-out results region is not in this agent's layout; using conversation"
661 );
662 "conversation"
663 }
664 };
665 let budget = window
666 .get_region(region)
667 .map(|r| r.max_tokens.saturating_sub(r.current_tokens))
668 .unwrap_or(0);
669 let allowed = budget.saturating_mul(4);
670 let fitted = match text.len() <= allowed {
671 true => text.to_string(),
672 false => {
673 let room = allowed.saturating_sub(REPORT_TRUNCATION_MARKER.len());
674 format!(
675 "{}{REPORT_TRUNCATION_MARKER}",
676 leviath_core::truncate_at_boundary(text, room)
677 )
678 }
679 };
680 let tokens = leviath_core::estimate_tokens(&fitted);
681 let _ = window.add_typed_entry(region, leviath_core::EntryKind::UserMessage, fitted, tokens);
682}
683
684fn agent_status(world: &World, entity: Entity) -> Option<AgentStatus> {
686 world.get::<AgentState>(entity).map(|s| s.status.clone())
687}
688
689fn set_status(world: &mut World, entity: Entity, status: AgentStatus) {
691 if let Some(mut state) = world.get_mut::<AgentState>(entity) {
692 state.status = status;
693 }
694}
695
696#[cfg(test)]
697mod tests {
698 use super::*;
699 use crate::components::{InferenceConfig, ToolResultRoutingComponent};
700 use crate::pipeline::{
701 ReadyToInfer, StageInference, StageInferences, StageProgress, StageSetup, StageSetups,
702 VisitCounts,
703 };
704 use leviath_core::blueprint::{ModelConfig, Stage};
705 use leviath_core::layout::{ContextLayout, RegionDefinition};
706 use leviath_core::{Blueprint, Region, RegionKind};
707 use std::collections::HashSet;
708
709 struct TestSpawner {
712 fail: HashSet<String>,
713 }
714
715 impl TestSpawner {
716 fn ok() -> Arc<dyn FanOutSpawner> {
717 Arc::new(TestSpawner {
718 fail: HashSet::new(),
719 })
720 }
721 fn refusing(ids: &[&str]) -> Arc<dyn FanOutSpawner> {
722 Arc::new(TestSpawner {
723 fail: ids.iter().map(|s| s.to_string()).collect(),
724 })
725 }
726 }
727
728 impl FanOutSpawner for TestSpawner {
729 fn spawn_worker(
730 &self,
731 world: &mut World,
732 _parent: Entity,
733 _config: &FanOutConfig,
734 item_id: &str,
735 _item_context: &serde_json::Value,
736 ) -> Result<Entity, String> {
737 if self.fail.contains(item_id) {
738 return Err(format!("spawn refused for '{item_id}'"));
739 }
740 Ok(world
741 .spawn((
742 AgentState {
743 agent_id: format!("worker-{item_id}"),
744 current_stage: "w".to_string(),
745 iteration: 0,
746 status: AgentStatus::Active,
747 spawned_children_ids: vec![],
748 pending_wait: None,
749 accepts_messages: true,
750 },
751 crate::persistence::RunMetadata {
754 run_id: format!("run-{item_id}"),
755 agent_name: "worker".to_string(),
756 agent_path: String::new(),
757 task: String::new(),
758 model: None,
759 workdir: String::new(),
760 num_stages: 1,
761 started_at: 0,
762 parent_run_id: None,
763 metadata: std::collections::HashMap::new(),
764 callback_url: None,
765 callback_secret: None,
766 title: None,
767 unattended: false,
768 read_paths: None,
769 output_request: None,
770 },
771 ))
772 .id())
773 }
774 }
775
776 fn cfg(merge: Option<&str>, max_workers: usize, policy: WorkerFailurePolicy) -> FanOutConfig {
777 FanOutConfig {
778 worker_agent: None,
779 worker_stage: Some("w".to_string()),
780 worker_query: None,
781 merge_stage: merge.map(String::from),
782 max_workers,
783 on_worker_failure: policy,
784 split_prompt: "split".to_string(),
785 results_region: None,
786 max_items: None,
787 }
788 }
789
790 fn window() -> ContextWindow {
791 let mut w = ContextWindow::new(12_000);
792 w.add_region(Region::new(
793 "conversation".to_string(),
794 RegionKind::Clearable,
795 10_000,
796 ));
797 w
798 }
799
800 fn stage_inf() -> StageInference {
801 StageInference {
802 provider_name: "script".to_string(),
803 model: "m".to_string(),
804 tools: vec![],
805 tool_filter: None,
806 fallbacks: Vec::new(),
807 output: None,
808 }
809 }
810
811 fn setup() -> StageSetup {
812 StageSetup {
813 inference_config: InferenceConfig {
814 temperature: None,
815 max_output_tokens: None,
816 extra_params: Default::default(),
817 batch_tool_hint: false,
818 shell_hint: false,
819 request_timeout_secs: None,
820 },
821 routing: None,
822 accepts_messages: true,
823 context_layout: None,
824 system_prompt: None,
825 output: None,
826 }
827 }
828
829 fn fanout_blueprint(config: FanOutConfig) -> Blueprint {
831 let layout = ContextLayout::new(
832 vec![RegionDefinition::new(
833 "conversation".to_string(),
834 RegionKind::Clearable,
835 10_000,
836 )],
837 12_000,
838 );
839 let mut s0 = Stage::new(
840 "fan".to_string(),
841 ModelConfig::new("script".to_string(), "m".to_string()),
842 );
843 s0.mode = StageMode::FanOut { config };
844 let s1 = Stage::new(
845 "merge".to_string(),
846 ModelConfig::new("script".to_string(), "m".to_string()),
847 );
848 Blueprint::new("t".to_string(), "d".to_string(), vec![s0, s1], layout)
849 }
850
851 fn parent_state() -> AgentState {
852 AgentState {
853 agent_id: "parent".to_string(),
854 current_stage: "fan".to_string(),
855 iteration: 0,
856 status: AgentStatus::Active,
857 spawned_children_ids: vec![],
858 pending_wait: None,
859 accepts_messages: true,
860 }
861 }
862
863 fn spawn_parent(world: &mut World, bp: Blueprint, response: &str) -> Entity {
866 world
867 .spawn((
868 AgentBlueprint(bp),
869 StageCursor { index: 0 },
870 parent_state(),
871 StageProgress::default(),
872 StageInferences(vec![stage_inf(), stage_inf()]),
873 StageSetups(vec![setup(), setup()]),
874 VisitCounts::default(),
875 window(),
876 InferenceResult {
877 response: response.to_string(),
878 tool_calls: vec![],
879 tokens_used: 0,
880 timestamp: 0,
881 },
882 ProcessResponse,
883 ))
884 .id()
885 }
886
887 fn install(world: &mut World, spawner: Arc<dyn FanOutSpawner>) {
888 world.insert_resource(FanOutSpawnerRes(spawner));
889 }
890
891 fn status_of(world: &World, e: Entity) -> AgentStatus {
892 world.get::<AgentState>(e).unwrap().status.clone()
893 }
894
895 fn assert_errored(world: &World, e: Entity) {
898 assert_eq!(
899 std::mem::discriminant(&status_of(world, e)),
900 std::mem::discriminant(&AgentStatus::Error {
901 message: String::new()
902 })
903 );
904 }
905
906 fn complete_worker(world: &mut World, worker: Entity, content: &str) {
907 set_status(world, worker, AgentStatus::Complete);
908 world.entity_mut(worker).insert(InferenceResult {
909 response: content.to_string(),
910 tool_calls: vec![],
911 tokens_used: 0,
912 timestamp: 0,
913 });
914 }
915
916 #[test]
919 fn parse_work_items_handles_array_prose_and_errors() {
920 let ok = parse_work_items(r#"[{"id":"a"},{"id":"b","context":{"k":1}}]"#).unwrap();
921 assert_eq!(ok.len(), 2);
922 assert_eq!(ok[0].id, "a");
923 assert_eq!(ok[1].context["k"], 1);
924 assert_eq!(parse_work_items("[{}]").unwrap()[0].id, "");
926 assert_eq!(
928 parse_work_items("Here you go:\n```json\n[{\"id\":\"x\"}]\n```")
929 .unwrap()
930 .len(),
931 1
932 );
933 assert!(parse_work_items("no array here").is_err());
935 assert!(parse_work_items("]nope[").is_err());
937 assert!(parse_work_items("[not json]").is_err());
939 }
940
941 #[test]
944 fn split_parks_a_fanout_stage_and_consumes_the_response() {
945 let mut world = World::new();
946 let e = spawn_parent(
947 &mut world,
948 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
949 r#"[{"id":"a"},{"id":"b"}]"#,
950 );
951 fan_out_split(&mut world);
952 assert!(world.get::<FanOutWaiting>(e).is_some());
953 assert_eq!(status_of(&world, e), AgentStatus::Waiting);
954 assert!(world.get::<ProcessResponse>(e).is_none());
956 assert!(world.get::<InferenceResult>(e).is_none());
957 let w = world.get::<FanOutWaiting>(e).unwrap();
958 assert_eq!(w.pending.len(), 2);
959 }
960
961 #[test]
967 fn split_keeps_only_the_first_max_items() {
968 let mut world = World::new();
969 let mut config = cfg(Some("merge"), 2, WorkerFailurePolicy::Continue);
970 config.max_items = Some(3);
971 let items: Vec<String> = (0..10).map(|i| format!(r#"{{"id":"w{i}"}}"#)).collect();
972 let e = spawn_parent(
973 &mut world,
974 fanout_blueprint(config),
975 &format!("[{}]", items.join(",")),
976 );
977
978 fan_out_split(&mut world);
979
980 let w = world.get::<FanOutWaiting>(e).expect("parked");
981 assert_eq!(w.pending.len(), 3, "kept the cap, not the ten produced");
982 let kept: Vec<&str> = w.pending.iter().map(|i| i.id.as_str()).collect();
983 assert_eq!(kept, ["w0", "w1", "w2"], "and kept the first of them");
984 }
985
986 #[test]
989 fn split_keeps_everything_under_the_cap() {
990 let mut world = World::new();
991 let mut config = cfg(Some("merge"), 2, WorkerFailurePolicy::Continue);
992 config.max_items = Some(9);
993 let e = spawn_parent(
994 &mut world,
995 fanout_blueprint(config),
996 r#"[{"id":"a"},{"id":"b"}]"#,
997 );
998
999 fan_out_split(&mut world);
1000
1001 assert_eq!(
1002 world.get::<FanOutWaiting>(e).expect("parked").pending.len(),
1003 2
1004 );
1005 }
1006
1007 #[test]
1008 fn split_errors_on_non_array_output() {
1009 let mut world = World::new();
1010 let e = spawn_parent(
1011 &mut world,
1012 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1013 "definitely not a json array",
1014 );
1015 fan_out_split(&mut world);
1016 assert!(world.get::<FanOutWaiting>(e).is_none());
1017 assert_errored(&world, e);
1018 }
1019
1020 #[test]
1021 fn split_skips_non_active_and_non_fanout_agents() {
1022 let mut world = World::new();
1024 let e = spawn_parent(
1025 &mut world,
1026 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1027 "[]",
1028 );
1029 set_status(&mut world, e, AgentStatus::Idle);
1030 fan_out_split(&mut world);
1031 assert!(world.get::<ProcessResponse>(e).is_some());
1032 assert!(world.get::<FanOutWaiting>(e).is_none());
1033
1034 let layout = ContextLayout::new(
1036 vec![RegionDefinition::new(
1037 "conversation".to_string(),
1038 RegionKind::Clearable,
1039 10_000,
1040 )],
1041 12_000,
1042 );
1043 let s = Stage::new(
1044 "plain".to_string(),
1045 ModelConfig::new("script".to_string(), "m".to_string()),
1046 );
1047 let bp = Blueprint::new("t".to_string(), "d".to_string(), vec![s], layout);
1048 let e2 = spawn_parent(&mut world, bp, "[]");
1049 fan_out_split(&mut world);
1050 assert!(world.get::<ProcessResponse>(e2).is_some());
1051 }
1052
1053 #[test]
1056 fn collect_starts_workers_then_merges_on_completion() {
1057 let mut world = World::new();
1058 install(&mut world, TestSpawner::ok());
1059 let e = spawn_parent(
1060 &mut world,
1061 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1062 r#"[{"id":"a"},{"id":"b"}]"#,
1063 );
1064 fan_out_split(&mut world);
1065 fan_out_collect(&mut world);
1066 let kids = world.get::<SubAgentChildren>(e).unwrap().children.clone();
1068 assert_eq!(kids.len(), 2);
1069 assert!(world.get::<FanOutWaiting>(e).is_some());
1070 for k in &kids {
1072 assert_eq!(world.get::<ParentRef>(*k).unwrap().depth, 1);
1073 }
1074
1075 for k in &kids {
1077 complete_worker(&mut world, *k, "fixed it");
1078 }
1079 fan_out_collect(&mut world);
1080 assert!(world.get::<FanOutWaiting>(e).is_none());
1081 assert_eq!(status_of(&world, e), AgentStatus::Active);
1082 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1083 assert!(world.get::<ReadyToInfer>(e).is_some());
1084 assert!(
1086 world
1087 .get::<ContextWindow>(e)
1088 .unwrap()
1089 .get_region("conversation")
1090 .unwrap()
1091 .current_tokens
1092 > 0
1093 );
1094 }
1095
1096 fn run_slim(world: &mut World) {
1098 let mut schedule = bevy_ecs::schedule::Schedule::default();
1099 schedule.add_systems(slim_merged_workers);
1100 schedule.run(world);
1101 }
1102
1103 #[test]
1108 fn merged_workers_are_slimmed_once_their_terminal_state_is_persisted() {
1109 let mut world = World::new();
1110 install(&mut world, TestSpawner::ok());
1111 let e = spawn_parent(
1112 &mut world,
1113 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1114 r#"[{"id":"a"}]"#,
1115 );
1116 fan_out_split(&mut world);
1117 fan_out_collect(&mut world);
1118 let worker = world.get::<SubAgentChildren>(e).unwrap().children[0];
1119 world
1121 .entity_mut(worker)
1122 .insert((window(), crate::pipeline::PersistWatermark::default()));
1123 complete_worker(&mut world, worker, "done");
1124 fan_out_collect(&mut world);
1125
1126 assert!(world.get::<MergedWorker>(worker).is_some());
1129 run_slim(&mut world);
1130 assert!(
1131 world.get::<ContextWindow>(worker).is_some(),
1132 "unpersisted terminal state stays resident"
1133 );
1134
1135 let mut wm = crate::pipeline::PersistWatermark::default();
1137 wm.stamp_status(leviath_core::run_meta::RunStatus::Complete);
1138 world.entity_mut(worker).insert(wm);
1139 run_slim(&mut world);
1140 assert!(world.get::<ContextWindow>(worker).is_none());
1141 assert!(world.get::<MergedWorker>(worker).is_none());
1142 assert!(world.get::<AgentState>(worker).is_some());
1144 }
1145
1146 #[test]
1147 fn collect_respects_max_workers_and_stages_pending() {
1148 let mut world = World::new();
1149 install(&mut world, TestSpawner::ok());
1150 let e = spawn_parent(
1151 &mut world,
1152 fanout_blueprint(cfg(Some("merge"), 1, WorkerFailurePolicy::Continue)),
1153 r#"[{"id":"a"},{"id":"b"}]"#,
1154 );
1155 fan_out_split(&mut world);
1156 fan_out_collect(&mut world);
1157 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 1);
1159 let first = world.get::<SubAgentChildren>(e).unwrap().children[0];
1160 fan_out_collect(&mut world);
1163 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 1);
1164 assert!(world.get::<FanOutWaiting>(e).is_some());
1165 complete_worker(&mut world, first, "one");
1166 fan_out_collect(&mut world);
1167 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 2);
1169 let second = world.get::<SubAgentChildren>(e).unwrap().children[1];
1170 complete_worker(&mut world, second, "two");
1171 fan_out_collect(&mut world);
1172 assert!(world.get::<FanOutWaiting>(e).is_none());
1173 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1174 }
1175
1176 #[test]
1177 fn fan_out_state_roundtrips_and_unresolved_workers_become_failures() {
1178 let mut world = World::new();
1179 install(&mut world, TestSpawner::ok());
1180 let e = spawn_parent(
1181 &mut world,
1182 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1183 r#"[{"id":"a"},{"id":"b"}]"#,
1184 );
1185 fan_out_split(&mut world);
1186 fan_out_collect(&mut world); let state = world.get::<FanOutWaiting>(e).unwrap().to_state();
1190 assert_eq!(state.active.len(), 2);
1191 assert!(state.active.iter().all(|(_id, run_id)| !run_id.is_empty()));
1192
1193 let by_run: std::collections::HashMap<String, Entity> = world
1195 .get::<SubAgentChildren>(e)
1196 .unwrap()
1197 .children
1198 .iter()
1199 .filter_map(|&c| {
1200 world
1201 .get::<crate::persistence::RunMetadata>(c)
1202 .map(|m| (m.run_id.clone(), c))
1203 })
1204 .collect();
1205 let fresh = world.spawn_empty().id();
1206 restore_fan_out_waiting(&mut world, fresh, state.clone(), &|rid| {
1207 by_run.get(rid).copied()
1208 });
1209 assert_eq!(
1210 world
1211 .get::<FanOutWaiting>(fresh)
1212 .unwrap()
1213 .to_state()
1214 .active
1215 .len(),
1216 2
1217 );
1218
1219 let orphaned = world.spawn_empty().id();
1222 restore_fan_out_waiting(&mut world, orphaned, state, &|_| None);
1223 let s = world.get::<FanOutWaiting>(orphaned).unwrap().to_state();
1224 assert!(s.active.is_empty());
1225 assert_eq!(s.failures.len(), 2);
1226 }
1227
1228 #[test]
1229 fn collect_fail_all_marks_parent_error() {
1230 let mut world = World::new();
1231 install(&mut world, TestSpawner::ok());
1232 let e = spawn_parent(
1233 &mut world,
1234 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::FailAll)),
1235 r#"[{"id":"a"}]"#,
1236 );
1237 fan_out_split(&mut world);
1238 fan_out_collect(&mut world);
1239 let worker = world.get::<SubAgentChildren>(e).unwrap().children[0];
1240 set_status(
1241 &mut world,
1242 worker,
1243 AgentStatus::Error {
1244 message: "boom".to_string(),
1245 },
1246 );
1247 fan_out_collect(&mut world);
1248 assert_errored(&world, e);
1249 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0); }
1251
1252 #[test]
1253 fn collect_continue_reports_failures_and_proceeds_without_merge() {
1254 let mut world = World::new();
1255 install(&mut world, TestSpawner::ok());
1256 let e = spawn_parent(
1258 &mut world,
1259 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1260 r#"[{"id":"a"},{"id":"b"}]"#,
1261 );
1262 fan_out_split(&mut world);
1263 fan_out_collect(&mut world);
1264 let kids = world.get::<SubAgentChildren>(e).unwrap().children.clone();
1265 set_status(
1266 &mut world,
1267 kids[0],
1268 AgentStatus::Error {
1269 message: "worker a died".to_string(),
1270 },
1271 );
1272 complete_worker(&mut world, kids[1], "b ok");
1273 fan_out_collect(&mut world);
1274 assert!(world.get::<FanOutWaiting>(e).is_none());
1275 assert!(world.get::<crate::pipeline::ResolveTransition>(e).is_some());
1276 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0);
1277 }
1278
1279 #[test]
1280 fn collect_finishes_immediately_when_there_are_no_work_items() {
1281 let mut world = World::new();
1282 install(&mut world, TestSpawner::ok());
1283 let e = spawn_parent(
1284 &mut world,
1285 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1286 "[]",
1287 );
1288 fan_out_split(&mut world);
1289 fan_out_collect(&mut world);
1290 assert!(world.get::<SubAgentChildren>(e).is_none());
1292 assert!(world.get::<FanOutWaiting>(e).is_none());
1293 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1294 }
1295
1296 #[test]
1297 fn collect_merge_stage_not_found_falls_through_to_transition() {
1298 let mut world = World::new();
1299 install(&mut world, TestSpawner::ok());
1300 let e = spawn_parent(
1301 &mut world,
1302 fanout_blueprint(cfg(Some("ghost"), 2, WorkerFailurePolicy::Continue)),
1303 "[]",
1304 );
1305 fan_out_split(&mut world);
1306 fan_out_collect(&mut world);
1307 assert!(world.get::<crate::pipeline::ResolveTransition>(e).is_some());
1309 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0);
1310 }
1311
1312 #[test]
1313 fn collect_abandons_a_cancelled_parent() {
1314 let mut world = World::new();
1315 install(&mut world, TestSpawner::ok());
1316 let e = spawn_parent(
1317 &mut world,
1318 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1319 r#"[{"id":"a"}]"#,
1320 );
1321 fan_out_split(&mut world);
1322 set_status(&mut world, e, AgentStatus::Cancelled);
1323 fan_out_collect(&mut world);
1324 assert!(world.get::<FanOutWaiting>(e).is_none());
1325 assert_eq!(status_of(&world, e), AgentStatus::Cancelled);
1326 }
1327
1328 #[test]
1329 fn collect_without_a_spawner_records_failures() {
1330 let mut world = World::new();
1332 let e = spawn_parent(
1333 &mut world,
1334 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1335 r#"[{"id":"a"}]"#,
1336 );
1337 fan_out_split(&mut world);
1338 fan_out_collect(&mut world);
1339 assert!(world.get::<FanOutWaiting>(e).is_none());
1341 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1342 }
1343
1344 #[test]
1345 fn collect_spawner_error_becomes_a_failure() {
1346 let mut world = World::new();
1347 install(&mut world, TestSpawner::refusing(&["a"]));
1348 let e = spawn_parent(
1349 &mut world,
1350 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::FailAll)),
1351 r#"[{"id":"a"}]"#,
1352 );
1353 fan_out_split(&mut world);
1354 fan_out_collect(&mut world);
1355 assert_errored(&world, e);
1357 }
1358
1359 #[test]
1362 fn start_worker_enforces_depth_cap() {
1363 let mut world = World::new();
1364 install(&mut world, TestSpawner::ok());
1365 let mut bp = fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue));
1366 bp.max_child_depth = Some(3);
1367 let e = spawn_parent(&mut world, bp, r#"[{"id":"deep"}]"#);
1368 world.entity_mut(e).insert(ParentRef {
1370 parent_entity: Entity::from_raw_u32(999)
1371 .expect("a small literal index is always a valid entity id"),
1372 parent_agent_id: "root".to_string(),
1373 depth: 3,
1374 });
1375 fan_out_split(&mut world);
1376 fan_out_collect(&mut world);
1377 assert!(world.get::<SubAgentChildren>(e).is_none());
1379 assert!(world.get::<FanOutWaiting>(e).is_none());
1380 }
1381
1382 #[test]
1383 fn start_worker_uses_existing_subagentchildren_cap_and_appends() {
1384 let mut world = World::new();
1385 install(&mut world, TestSpawner::ok());
1386 let e = spawn_parent(
1387 &mut world,
1388 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1389 r#"[{"id":"a"}]"#,
1390 );
1391 world.entity_mut(e).insert(SubAgentChildren {
1393 children: vec![
1394 Entity::from_raw_u32(1000)
1395 .expect("a small literal index is always a valid entity id"),
1396 ],
1397 max_child_depth: 9,
1398 });
1399 fan_out_split(&mut world);
1400 fan_out_collect(&mut world);
1401 let kids = world.get::<SubAgentChildren>(e).unwrap();
1402 assert_eq!(kids.max_child_depth, 9);
1403 assert_eq!(kids.children.len(), 2); }
1405
1406 #[test]
1414 fn a_submitted_answer_beats_the_last_assistant_text() {
1415 let mut world = World::new();
1416 let worker = world
1417 .spawn((
1418 parent_state(),
1419 InferenceResult {
1420 response: "Let me run the tests one more time.".to_string(),
1423 tool_calls: vec![],
1424 tokens_used: 0,
1425 timestamp: 0,
1426 },
1427 crate::persistence::FinalOutput(leviath_core::output::FinalOutput::new(
1428 "changed src/lib.rs; the failing test now passes",
1429 None,
1430 "fix_worker".to_string(),
1431 0,
1432 )),
1433 ))
1434 .id();
1435 set_status(&mut world, worker, AgentStatus::Complete);
1436 assert_eq!(
1437 worker_terminal_result(&world, worker),
1438 Some(Ok(
1439 "changed src/lib.rs; the failing test now passes".to_string()
1440 ))
1441 );
1442 }
1443
1444 #[test]
1447 fn a_worker_that_submitted_nothing_still_falls_back_to_its_text() {
1448 let mut world = World::new();
1449 let worker = world
1450 .spawn((
1451 parent_state(),
1452 InferenceResult {
1453 response: "the old behaviour".to_string(),
1454 tool_calls: vec![],
1455 tokens_used: 0,
1456 timestamp: 0,
1457 },
1458 ))
1459 .id();
1460 set_status(&mut world, worker, AgentStatus::Complete);
1461 assert_eq!(
1462 worker_terminal_result(&world, worker),
1463 Some(Ok("the old behaviour".to_string()))
1464 );
1465 }
1466
1467 fn spawn_required_output_worker(world: &mut World) -> Entity {
1469 let mut stage = Stage::new(
1470 "w".to_string(),
1471 ModelConfig::new("script".to_string(), "m".to_string()),
1472 );
1473 stage.require_output = true;
1474 let layout = ContextLayout::new(
1475 vec![RegionDefinition::new(
1476 "conversation".to_string(),
1477 RegionKind::Clearable,
1478 10_000,
1479 )],
1480 12_000,
1481 );
1482 let bp = Blueprint::new("w".to_string(), "d".to_string(), vec![stage], layout);
1483 let worker = world
1484 .spawn((parent_state(), AgentBlueprint(bp), StageCursor { index: 0 }))
1485 .id();
1486 set_status(world, worker, AgentStatus::Complete);
1487 worker
1488 }
1489
1490 #[test]
1499 fn a_worker_that_owes_an_output_and_has_none_is_a_failure() {
1500 let mut world = World::new();
1501 let worker = spawn_required_output_worker(&mut world);
1502
1503 assert_eq!(
1504 worker_terminal_result(&world, worker),
1505 Some(Err(
1506 "worker finished without the final output its stage requires".to_string()
1507 )),
1508 "the merge has to be told a worker failed, and why"
1509 );
1510 }
1511
1512 #[test]
1515 fn a_worker_that_owes_an_output_and_has_one_contributes_it() {
1516 let mut world = World::new();
1517 let worker = spawn_required_output_worker(&mut world);
1518 world
1519 .entity_mut(worker)
1520 .insert(crate::persistence::FinalOutput(
1521 leviath_core::output::FinalOutput {
1522 content: "the rows".to_string(),
1523 format: Some("csv".to_string()),
1524 stage: "w".to_string(),
1525 submitted_at: 0,
1526 truncated: false,
1527 artifacts: vec![],
1528 },
1529 ));
1530
1531 assert_eq!(
1532 worker_terminal_result(&world, worker),
1533 Some(Ok("the rows".to_string()))
1534 );
1535 }
1536
1537 #[test]
1541 fn a_worker_with_no_stage_to_read_owes_nothing() {
1542 let mut world = World::new();
1543 let bp = fanout_blueprint(cfg(None, 1, WorkerFailurePolicy::Continue));
1544
1545 let bare = world.spawn(parent_state()).id();
1547 assert!(!worker_requires_output(&world, bare));
1548
1549 let no_cursor = world.spawn((parent_state(), AgentBlueprint(bp))).id();
1551 assert!(!worker_requires_output(&world, no_cursor));
1552
1553 let past_end = world
1555 .spawn((
1556 parent_state(),
1557 AgentBlueprint(fanout_blueprint(cfg(
1558 None,
1559 1,
1560 WorkerFailurePolicy::Continue,
1561 ))),
1562 StageCursor { index: 99 },
1563 ))
1564 .id();
1565 assert!(!worker_requires_output(&world, past_end));
1566 }
1567
1568 #[test]
1572 fn a_worker_that_owes_nothing_keeps_the_last_turn_fallback() {
1573 let mut world = World::new();
1574 let worker = world.spawn(parent_state()).id();
1575 set_status(&mut world, worker, AgentStatus::Complete);
1576
1577 assert_eq!(
1578 worker_terminal_result(&world, worker),
1579 Some(Ok(String::new()))
1580 );
1581 }
1582
1583 #[test]
1584 fn worker_terminal_result_covers_every_status() {
1585 let mut world = World::new();
1586 let complete = world
1587 .spawn((
1588 parent_state(),
1589 InferenceResult {
1590 response: "done text".to_string(),
1591 tool_calls: vec![],
1592 tokens_used: 0,
1593 timestamp: 0,
1594 },
1595 ))
1596 .id();
1597 set_status(&mut world, complete, AgentStatus::Complete);
1598 assert_eq!(
1599 worker_terminal_result(&world, complete),
1600 Some(Ok("done text".to_string()))
1601 );
1602
1603 let complete_no_infer = world.spawn(parent_state()).id();
1604 set_status(&mut world, complete_no_infer, AgentStatus::Complete);
1605 assert_eq!(
1606 worker_terminal_result(&world, complete_no_infer),
1607 Some(Ok(String::new()))
1608 );
1609
1610 let errored = world.spawn(parent_state()).id();
1611 set_status(
1612 &mut world,
1613 errored,
1614 AgentStatus::Error {
1615 message: "x".to_string(),
1616 },
1617 );
1618 assert_eq!(
1619 worker_terminal_result(&world, errored),
1620 Some(Err("x".to_string()))
1621 );
1622
1623 let cancelled = world.spawn(parent_state()).id();
1624 set_status(&mut world, cancelled, AgentStatus::Cancelled);
1625 assert!(worker_terminal_result(&world, cancelled).is_some_and(|r| r.is_err()));
1626
1627 let running = world.spawn(parent_state()).id(); assert_eq!(worker_terminal_result(&world, running), None);
1629
1630 assert!(
1631 worker_terminal_result(
1632 &world,
1633 Entity::from_raw_u32(4242)
1634 .expect("a small literal index is always a valid entity id")
1635 )
1636 .is_some_and(|r| r.is_err())
1637 );
1638 }
1639
1640 #[test]
1645 fn a_huge_fan_out_still_reaches_the_merge_stage() {
1646 let mut world = World::new();
1647 let mut window = ContextWindow::new(100_000);
1648 window.add_region(leviath_core::Region::new(
1649 "conversation".to_string(),
1650 leviath_core::RegionKind::Clearable,
1651 10_000,
1652 ));
1653 let parent = world.spawn((parent_state(), window)).id();
1654
1655 let huge = "x".repeat(leviath_core::output::MAX_FINAL_OUTPUT_BYTES);
1657 let summaries: Vec<(String, String)> =
1658 (0..100).map(|i| (format!("w{i}"), huge.clone())).collect();
1659 let report = build_report(&summaries, &[], Some(10_000));
1660 inject_results(&mut world, parent, "conversation", &report);
1661
1662 let region = world
1663 .get::<ContextWindow>(parent)
1664 .expect("window")
1665 .get_region("conversation")
1666 .expect("region");
1667 assert!(
1668 !region.content.is_empty(),
1669 "the merge stage must receive something rather than nothing"
1670 );
1671 let landed = ®ion.content[0].content;
1672 assert!(landed.contains("100 succeeded"), "header survives");
1675 assert!(landed.contains("truncated"), "and says it was cut");
1676 assert!(region.current_tokens <= region.max_tokens, "within budget");
1677 }
1678
1679 #[test]
1686 fn a_report_larger_than_what_is_left_of_the_region_is_trimmed_not_dropped() {
1687 const REGION_TOKENS: usize = 2_000;
1688 let mut world = World::new();
1689 let mut window = ContextWindow::new(100_000);
1690 window.add_region(leviath_core::Region::new(
1691 "worker_results".to_string(),
1692 leviath_core::RegionKind::Clearable,
1693 REGION_TOKENS,
1694 ));
1695 let filler = "f".repeat(REGION_TOKENS * 4 * 8 / 10);
1697 let filler_tokens = leviath_core::estimate_tokens(&filler);
1698 window
1699 .add_typed_entry(
1700 "worker_results",
1701 leviath_core::EntryKind::UserMessage,
1702 filler,
1703 filler_tokens,
1704 )
1705 .expect("the filler fits");
1706 let parent = world.spawn((parent_state(), window)).id();
1707
1708 let long = "x".repeat(5_000);
1710 let summaries: Vec<(String, String)> =
1711 (0..8).map(|i| (format!("w{i}"), long.clone())).collect();
1712 let report = build_report(&summaries, &[], Some(REGION_TOKENS));
1713 assert!(report.len() > REGION_TOKENS * 4 / 5, "the report is big");
1714 inject_results(&mut world, parent, "worker_results", &report);
1715
1716 let region = world
1717 .get::<ContextWindow>(parent)
1718 .expect("window")
1719 .get_region("worker_results")
1720 .expect("region")
1721 .clone();
1722 assert_eq!(
1723 region.content.len(),
1724 2,
1725 "the report landed beside the filler"
1726 );
1727 let landed = ®ion.content[1].content;
1728 assert!(
1729 landed.contains("8 succeeded"),
1730 "the header survives the cut"
1731 );
1732 assert!(
1733 landed.contains(REPORT_TRUNCATION_MARKER.trim()),
1734 "and it says it was cut"
1735 );
1736 assert!(region.current_tokens <= region.max_tokens, "within budget");
1737 }
1738
1739 #[test]
1745 fn every_worker_appears_in_a_large_fan_out() {
1746 const REGION_TOKENS: usize = 40_000;
1751 let mut world = World::new();
1752 let mut window = ContextWindow::new(400_000);
1753 window.add_region(leviath_core::Region::new(
1754 "worker_results".to_string(),
1755 leviath_core::RegionKind::Clearable,
1756 REGION_TOKENS,
1757 ));
1758 let parent = world.spawn((parent_state(), window)).id();
1759
1760 let long = "x".repeat(50_000);
1761 let summaries: Vec<(String, String)> =
1762 (0..100).map(|i| (format!("w{i}"), long.clone())).collect();
1763 let report = build_report(&summaries, &[], Some(REGION_TOKENS));
1764 inject_results(&mut world, parent, "worker_results", &report);
1765
1766 let landed = world
1767 .get::<ContextWindow>(parent)
1768 .expect("window")
1769 .get_region("worker_results")
1770 .expect("region")
1771 .content[0]
1772 .content
1773 .clone();
1774 for i in 0..100 {
1775 assert!(
1776 landed.contains(&format!("## worker w{i}\n")),
1777 "worker w{i} never reached the merge stage"
1778 );
1779 }
1780 assert!(landed.contains("read a worker's own run"));
1783 }
1784
1785 #[test]
1787 fn the_share_shrinks_as_the_worker_count_grows() {
1788 assert!(bytes_per_worker(Some(40_000), 4) > bytes_per_worker(Some(40_000), 100));
1789 assert!(bytes_per_worker(Some(80_000), 10) > bytes_per_worker(Some(40_000), 10));
1791 assert_eq!(
1793 bytes_per_worker(Some(10), 10_000),
1794 MIN_REPORT_BYTES_PER_WORKER
1795 );
1796 assert_eq!(bytes_per_worker(None, 4), DEFAULT_REPORT_BYTES_PER_WORKER);
1798 }
1799
1800 #[test]
1803 fn results_go_to_the_named_region() {
1804 let mut world = World::new();
1805 let mut window = ContextWindow::new(100_000);
1806 window.add_region(leviath_core::Region::new(
1807 "conversation".to_string(),
1808 leviath_core::RegionKind::Clearable,
1809 10_000,
1810 ));
1811 window.add_region(leviath_core::Region::new(
1812 "worker_results".to_string(),
1813 leviath_core::RegionKind::Clearable,
1814 20_000,
1815 ));
1816 let parent = world.spawn((parent_state(), window)).id();
1817 inject_results(&mut world, parent, "worker_results", "the report");
1818
1819 let w = world.get::<ContextWindow>(parent).expect("window");
1820 assert_eq!(
1821 w.get_region("worker_results")
1822 .expect("region")
1823 .content
1824 .len(),
1825 1
1826 );
1827 assert!(
1828 w.get_region("conversation")
1829 .expect("region")
1830 .content
1831 .is_empty(),
1832 "the default region is left alone"
1833 );
1834 }
1835
1836 #[test]
1839 fn an_unknown_results_region_falls_back_to_the_conversation() {
1840 let mut world = World::new();
1841 let mut window = ContextWindow::new(100_000);
1842 window.add_region(leviath_core::Region::new(
1843 "conversation".to_string(),
1844 leviath_core::RegionKind::Clearable,
1845 10_000,
1846 ));
1847 let parent = world.spawn((parent_state(), window)).id();
1848 inject_results(&mut world, parent, "typo_region", "the report");
1849
1850 assert_eq!(
1851 world
1852 .get::<ContextWindow>(parent)
1853 .expect("window")
1854 .get_region("conversation")
1855 .expect("region")
1856 .content
1857 .len(),
1858 1
1859 );
1860 }
1861
1862 #[test]
1865 fn a_small_fan_out_report_is_not_trimmed() {
1866 let mut world = World::new();
1867 let mut window = ContextWindow::new(100_000);
1868 window.add_region(leviath_core::Region::new(
1869 "conversation".to_string(),
1870 leviath_core::RegionKind::Clearable,
1871 10_000,
1872 ));
1873 let parent = world.spawn((parent_state(), window)).id();
1874 let report = build_report(
1875 &[("a".to_string(), "did the thing".to_string())],
1876 &[],
1877 Some(10_000),
1878 );
1879 inject_results(&mut world, parent, "conversation", &report);
1880 let landed = world
1881 .get::<ContextWindow>(parent)
1882 .expect("window")
1883 .get_region("conversation")
1884 .expect("region")
1885 .content[0]
1886 .content
1887 .clone();
1888 assert_eq!(landed, report);
1889 }
1890
1891 #[test]
1892 fn build_report_lists_successes_and_failures() {
1893 let report = build_report(
1894 &[("a".to_string(), "ok-a".to_string())],
1895 &[("b".to_string(), "boom".to_string())],
1896 None,
1897 );
1898 assert!(report.contains("1 succeeded, 1 failed"));
1899 assert!(report.contains("## worker a\nok-a"));
1900 assert!(report.contains("## worker b FAILED\nboom"));
1901 }
1902
1903 #[test]
1904 fn inject_conversation_is_a_noop_without_a_window() {
1905 let mut world = World::new();
1906 let has_window = world.spawn(window()).id();
1907 inject_results(&mut world, has_window, "conversation", "hello");
1908 assert!(
1909 world
1910 .get::<ContextWindow>(has_window)
1911 .unwrap()
1912 .get_region("conversation")
1913 .unwrap()
1914 .current_tokens
1915 > 0
1916 );
1917 let no_window = world.spawn(parent_state()).id();
1919 inject_results(&mut world, no_window, "conversation", "hello");
1920 }
1921
1922 #[test]
1923 fn set_status_is_a_noop_for_a_missing_agent() {
1924 let mut world = World::new();
1925 set_status(
1926 &mut world,
1927 Entity::from_raw_u32(77).expect("a small literal index is always a valid entity id"),
1928 AgentStatus::Complete,
1929 );
1930 assert_eq!(
1931 agent_status(
1932 &world,
1933 Entity::from_raw_u32(77)
1934 .expect("a small literal index is always a valid entity id")
1935 ),
1936 None
1937 );
1938 }
1939
1940 #[test]
1943 fn force_transition_applies_routing_and_handles_despawn_and_overflow() {
1944 use crate::pipeline::force_transition;
1945 let mut world = World::new();
1947 let mut setups = vec![setup(), setup()];
1948 setups[1].routing = Some(leviath_core::ToolResultRouting::default());
1949 let e = world
1950 .spawn((
1951 AgentBlueprint(fanout_blueprint(cfg(
1952 Some("merge"),
1953 2,
1954 WorkerFailurePolicy::Continue,
1955 ))),
1956 StageCursor { index: 0 },
1957 parent_state(),
1958 StageProgress::default(),
1959 StageInferences(vec![stage_inf(), stage_inf()]),
1960 StageSetups(setups),
1961 VisitCounts::default(),
1962 window(),
1963 ))
1964 .id();
1965 let agent = crate::world::AgentId::in_world(&world, e);
1966 force_transition(&mut world, agent, 1);
1967 assert!(world.get::<ToolResultRoutingComponent>(e).is_some());
1968 assert!(world.get::<ReadyToInfer>(e).is_some());
1969
1970 let gone = crate::world::AgentId::in_world(
1972 &world,
1973 Entity::from_raw_u32(9191).expect("a small literal index is always a valid entity id"),
1974 );
1975 force_transition(&mut world, gone, 1);
1976 }
1977
1978 #[test]
1979 fn force_transition_marks_error_on_prompt_overflow() {
1980 use crate::pipeline::force_transition;
1981 let layout = ContextLayout::new(
1983 vec![RegionDefinition::new(
1984 "task".to_string(),
1985 RegionKind::Pinned,
1986 20,
1987 )],
1988 1000,
1989 );
1990 let mut s0 = Stage::new(
1991 "fan".to_string(),
1992 ModelConfig::new("script".to_string(), "m".to_string()),
1993 );
1994 s0.mode = StageMode::FanOut {
1995 config: cfg(Some("merge"), 2, WorkerFailurePolicy::Continue),
1996 };
1997 let mut s1 = Stage::new(
1998 "merge".to_string(),
1999 ModelConfig::new("script".to_string(), "m".to_string()),
2000 );
2001 s1.config.insert(
2002 "system_prompt".to_string(),
2003 serde_json::Value::String("x".repeat(10_000)),
2004 );
2005 let bp = Blueprint::new("t".to_string(), "d".to_string(), vec![s0, s1], layout);
2006
2007 let mut setups = vec![setup(), setup()];
2008 setups[1].system_prompt = Some("x".repeat(10_000));
2009 let mut w = ContextWindow::new(1000);
2010 w.add_region(Region::new("task".to_string(), RegionKind::Pinned, 20));
2011 let (mut world, e) = world_with(bp, setups, w);
2012 let agent = crate::world::AgentId::in_world(&world, e);
2013 force_transition(&mut world, agent, 1);
2014 assert_errored(&world, e);
2015 }
2016
2017 fn world_with(bp: Blueprint, setups: Vec<StageSetup>, w: ContextWindow) -> (World, Entity) {
2019 let mut world = World::new();
2020 let e = world
2021 .spawn((
2022 AgentBlueprint(bp),
2023 StageCursor { index: 0 },
2024 parent_state(),
2025 StageProgress::default(),
2026 StageInferences(vec![stage_inf(), stage_inf()]),
2027 StageSetups(setups),
2028 VisitCounts::default(),
2029 w,
2030 ))
2031 .id();
2032 (world, e)
2033 }
2034}