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
33const MAX_SPLIT_RETRIES: usize = 2;
43
44#[derive(Component, Debug, Clone, Copy, Default, PartialEq, Eq)]
49pub struct SplitAttempts(pub usize);
50
51fn split_correction(reason: &str) -> String {
57 format!(
58 "Your previous response could not be used: {reason}. Reply with the JSON \
59 array of work items and nothing else - no prose before or after it, no \
60 markdown fences, no explanation. It must start with `[` and end with `]`."
61 )
62}
63
64fn response_snippet(response: &str) -> String {
73 let trimmed = response.trim();
74 if trimmed.is_empty() {
75 return "the response was empty".to_string();
76 }
77 let kept: String = trimmed.chars().take(MAX_SPLIT_SNIPPET).collect();
80 match kept.len() < trimmed.len() {
81 true => format!("the response began: {kept}…"),
82 false => format!("the response was: {kept}"),
83 }
84}
85
86const MAX_SPLIT_SNIPPET: usize = 200;
88
89#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Default)]
91pub struct WorkItem {
92 #[serde(default)]
94 pub id: String,
95 #[serde(default)]
97 pub context: serde_json::Value,
98}
99
100pub fn parse_work_items(content: &str) -> Result<Vec<WorkItem>, String> {
103 let trimmed = content.trim();
104 let slice = match (trimmed.find('['), trimmed.rfind(']')) {
107 (Some(s), Some(e)) if e > s => trimmed.get(s..=e),
108 _ => None,
109 }
110 .ok_or_else(|| "split output is not a JSON array".to_string())?;
111 serde_json::from_str(slice)
112 .map_err(|e| format!("split output is not a valid JSON array of work items: {e}"))
113}
114
115pub trait FanOutSpawner: Send + Sync {
121 fn spawn_worker(
124 &self,
125 world: &mut World,
126 parent: Entity,
127 config: &FanOutConfig,
128 item_id: &str,
129 item_context: &serde_json::Value,
130 ) -> Result<Entity, String>;
131}
132
133#[derive(Resource, Clone)]
136pub struct FanOutSpawnerRes(pub Arc<dyn FanOutSpawner>);
137
138struct ActiveWorker {
142 item_id: String,
143 entity: Entity,
144 run_id: String,
145}
146
147#[derive(Component)]
150pub struct FanOutWaiting {
151 config: FanOutConfig,
152 max_workers: usize,
153 pending: VecDeque<WorkItem>,
154 active: Vec<ActiveWorker>,
155 summaries: Vec<(String, String)>,
156 failures: Vec<(String, String)>,
157}
158
159#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
164pub struct FanOutState {
165 pub config: FanOutConfig,
167 pub max_workers: usize,
169 pub pending: Vec<WorkItem>,
171 pub active: Vec<(String, String)>,
173 pub summaries: Vec<(String, String)>,
175 pub failures: Vec<(String, String)>,
177}
178
179impl FanOutWaiting {
180 pub fn outstanding(&self) -> usize {
185 self.active.len() + self.pending.len()
186 }
187
188 pub(crate) fn to_state(&self) -> FanOutState {
190 FanOutState {
191 config: self.config.clone(),
192 max_workers: self.max_workers,
193 pending: self.pending.iter().cloned().collect(),
194 active: self
195 .active
196 .iter()
197 .map(|w| (w.item_id.clone(), w.run_id.clone()))
198 .collect(),
199 summaries: self.summaries.clone(),
200 failures: self.failures.clone(),
201 }
202 }
203}
204
205pub fn restore_fan_out_waiting(
211 world: &mut World,
212 parent: Entity,
213 state: FanOutState,
214 resolve: &dyn Fn(&str) -> Option<Entity>,
215) {
216 let mut active = Vec::new();
217 let mut failures = state.failures;
218 for (item_id, run_id) in state.active {
219 match resolve(&run_id) {
220 Some(entity) => active.push(ActiveWorker {
221 item_id,
222 entity,
223 run_id,
224 }),
225 None => failures.push((item_id, "worker did not reload after restart".to_string())),
226 }
227 }
228 world.entity_mut(parent).insert(FanOutWaiting {
229 config: state.config,
230 max_workers: state.max_workers,
231 pending: state.pending.into_iter().collect(),
232 active,
233 summaries: state.summaries,
234 failures,
235 });
236}
237
238pub fn fan_out_split(world: &mut World) {
244 crate::tick_scope::clear();
245 let mut candidates: Vec<(Entity, String, FanOutConfig)> = Vec::new();
246 {
247 let mut q = world.query_filtered::<(
248 Entity,
249 &AgentState,
250 &AgentBlueprint,
251 &StageCursor,
252 &InferenceResult,
253 ), With<ProcessResponse>>();
254 for (entity, state, bp, cursor, infer) in q.iter(world) {
255 if state.status != AgentStatus::Active {
256 continue;
257 }
258 if let StageMode::FanOut { config } = &bp.0.stages[cursor.index].mode {
259 candidates.push((entity, infer.response.clone(), config.clone()));
260 }
261 }
262 }
263
264 for (parent, response, config) in candidates {
265 crate::tick_scope::enter(parent);
266 world
267 .entity_mut(parent)
268 .remove::<ProcessResponse>()
269 .remove::<InferenceResult>();
270 match parse_work_items(&response) {
271 Ok(items) => {
272 let max_workers = config.max_workers.max(1);
273 let items = match config.max_items {
279 Some(cap) if items.len() > cap => {
280 tracing::warn!(
281 produced = items.len(),
282 cap,
283 "fan_out split produced more items than max_items; keeping the first"
284 );
285 items.into_iter().take(cap).collect::<Vec<_>>()
286 }
287 _ => items,
288 };
289 world.entity_mut(parent).insert(FanOutWaiting {
290 config,
291 max_workers,
292 pending: items.into_iter().collect(),
293 active: Vec::new(),
294 summaries: Vec::new(),
295 failures: Vec::new(),
296 });
297 set_status(world, parent, AgentStatus::Waiting);
298 }
299 Err(message) => {
300 let attempts = world.get::<SplitAttempts>(parent).map_or(0, |a| a.0);
305 let corrected = match attempts < MAX_SPLIT_RETRIES {
309 false => false,
310 true => {
311 let mut entity = world.entity_mut(parent);
312 match entity.get_mut::<ContextWindow>() {
313 None => false,
314 Some(mut window) => {
315 let tokens = leviath_core::estimate_tokens(&response);
319 let _ = window.add_typed_entry(
320 "conversation",
321 leviath_core::EntryKind::AssistantTurn {
322 tool_calls: Vec::new(),
323 },
324 response.clone(),
325 tokens,
326 );
327 crate::pipeline::inject_system_nudge(
328 &mut window,
329 &split_correction(&message),
330 );
331 true
332 }
333 }
334 }
335 };
336 if corrected {
337 tracing::warn!(
338 attempt = attempts + 1,
339 max = MAX_SPLIT_RETRIES,
340 error = %message,
341 "fan_out split did not parse; asking the model again"
342 );
343 world
344 .entity_mut(parent)
345 .insert(SplitAttempts(attempts + 1))
346 .insert(crate::pipeline::ReadyToInfer);
347 } else {
348 set_status(
352 world,
353 parent,
354 AgentStatus::Error {
355 message: format!(
356 "fan_out split failed after {attempts} correction(s): \
357 {message} ({})",
358 response_snippet(&response)
359 ),
360 },
361 );
362 }
363 }
364 }
365 }
366}
367
368pub fn fan_out_collect(world: &mut World) {
373 crate::tick_scope::clear();
374 let parents: Vec<Entity> = {
375 let mut q = world.query_filtered::<Entity, With<FanOutWaiting>>();
376 q.iter(world).collect()
377 };
378
379 for parent in parents {
380 crate::tick_scope::enter(parent);
381 if !matches!(agent_status(world, parent), Some(AgentStatus::Waiting)) {
384 world.entity_mut(parent).remove::<FanOutWaiting>();
385 continue;
386 }
387 let mut w = world
390 .entity_mut(parent)
391 .take::<FanOutWaiting>()
392 .expect("a Waiting fan-out parent still holds FanOutWaiting");
393
394 let mut still_active = Vec::with_capacity(w.active.len());
403 for aw in std::mem::take(&mut w.active) {
404 match worker_terminal_result(world, aw.entity) {
405 Some(result) => {
406 match result {
407 Ok(content) => w.summaries.push((aw.item_id, content)),
408 Err(message) => w.failures.push((aw.item_id, message)),
409 }
410 world.entity_mut(aw.entity).insert(MergedWorker);
411 }
412 None => still_active.push(aw),
413 }
414 }
415 w.active = still_active;
416
417 while w.active.len() < w.max_workers {
419 let Some(item) = w.pending.pop_front() else {
420 break;
421 };
422 match start_worker(world, parent, &w.config, &item) {
423 Ok(child) => {
424 let run_id = world
426 .get::<crate::persistence::RunMetadata>(child)
427 .map(|m| m.run_id.clone())
428 .unwrap_or_default();
429 w.active.push(ActiveWorker {
430 item_id: item.id,
431 entity: child,
432 run_id,
433 });
434 }
435 Err(message) => w.failures.push((item.id, message)),
436 }
437 }
438
439 if w.active.is_empty() && w.pending.is_empty() {
441 finish_fan_out(world, parent, w);
442 } else {
443 world.entity_mut(parent).insert(w);
444 }
445 }
446}
447
448#[derive(Component)]
451pub struct MergedWorker;
452
453pub fn slim_merged_workers(
463 workers: Query<(Entity, &crate::pipeline::PersistWatermark), With<MergedWorker>>,
464 mut commands: Commands,
465) {
466 crate::tick_scope::clear();
467 for (entity, watermark) in workers.iter() {
468 crate::tick_scope::enter(entity);
469 let terminal_persisted = matches!(
470 watermark.persisted_status(),
471 Some(
472 leviath_core::run_meta::RunStatus::Complete
473 | leviath_core::run_meta::RunStatus::Error
474 | leviath_core::run_meta::RunStatus::Cancelled
475 )
476 );
477 if !terminal_persisted {
478 continue; }
480 commands.entity(entity).remove::<(
481 ContextWindow,
482 InferenceResult,
483 crate::pipeline::StageInferences,
484 crate::pipeline::StageSetups,
485 AgentBlueprint,
486 MergedWorker,
487 )>();
488 }
489}
490
491fn finish_fan_out(world: &mut World, parent: Entity, w: FanOutWaiting) {
493 if !w.failures.is_empty() && w.config.on_worker_failure == WorkerFailurePolicy::FailAll {
494 set_status(
495 world,
496 parent,
497 AgentStatus::Error {
498 message: format!(
499 "fan_out: {} worker(s) failed (on_worker_failure = fail_all)",
500 w.failures.len()
501 ),
502 },
503 );
504 return;
505 }
506
507 let region = w
512 .config
513 .results_region
514 .clone()
515 .unwrap_or_else(|| "conversation".to_string());
516 let budget = world
517 .get::<ContextWindow>(parent)
518 .and_then(|window| window.get_region(®ion).map(|r| r.max_tokens));
519 let report = build_report(&w.summaries, &w.failures, budget);
520 inject_results(world, parent, ®ion, &report);
521
522 set_status(world, parent, AgentStatus::Active);
525 match w.config.merge_stage.as_deref().and_then(|name| {
526 world
527 .get::<AgentBlueprint>(parent)
528 .and_then(|bp| bp.0.stages.iter().position(|s| s.name == name))
529 }) {
530 Some(idx) => crate::pipeline::force_transition(
531 world,
532 crate::world::AgentId::in_world(world, parent),
533 idx,
534 ),
535 None => {
536 world.entity_mut(parent).insert(ResolveTransition);
537 }
538 }
539}
540
541fn start_worker(
544 world: &mut World,
545 parent: Entity,
546 config: &FanOutConfig,
547 item: &WorkItem,
548) -> Result<Entity, String> {
549 let max_depth = world
550 .get::<SubAgentChildren>(parent)
551 .map(|k| k.max_child_depth)
552 .or_else(|| {
553 world
554 .get::<AgentBlueprint>(parent)
555 .and_then(|bp| bp.0.max_child_depth)
556 })
557 .unwrap_or(DEFAULT_FANOUT_DEPTH);
558 let parent_depth = world.get::<ParentRef>(parent).map_or(0, |p| p.depth);
559 let child_depth = parent_depth + 1;
560 if child_depth > max_depth {
561 return Err(format!(
562 "fan-out worker depth limit ({max_depth}) reached; not spawning"
563 ));
564 }
565
566 let spawner = world
567 .get_resource::<FanOutSpawnerRes>()
568 .map(|r| r.0.clone())
569 .ok_or_else(|| "no fan-out spawner installed".to_string())?;
570 let child = spawner.spawn_worker(world, parent, config, &item.id, &item.context)?;
571
572 let parent_agent_id = world
573 .get::<AgentState>(parent)
574 .map(|s| s.agent_id.clone())
575 .unwrap_or_default();
576 world.entity_mut(child).insert(ParentRef {
577 parent_entity: parent,
578 parent_agent_id,
579 depth: child_depth,
580 });
581 match world.get_mut::<SubAgentChildren>(parent) {
582 Some(mut kids) => kids.children.push(child),
583 None => {
584 world.entity_mut(parent).insert(SubAgentChildren {
585 children: vec![child],
586 max_child_depth: max_depth,
587 });
588 }
589 }
590 let worker_id = world
594 .get::<crate::persistence::RunMetadata>(child)
595 .expect("a fan-out worker always has run metadata")
596 .run_id
597 .clone();
598 world
599 .get_mut::<AgentState>(parent)
600 .expect("a fan-out parent always has AgentState")
601 .spawned_children_ids
602 .push(worker_id);
603 crate::context_transform::apply_context_transforms(
606 world,
607 crate::world::AgentId::in_world(world, parent),
608 crate::world::AgentId::in_world(world, child),
609 );
610 Ok(child)
611}
612
613fn worker_terminal_result(world: &World, worker: Entity) -> Option<Result<String, String>> {
637 match agent_status(world, worker) {
638 None => Some(Err("worker vanished".to_string())),
639 Some(AgentStatus::Complete) => {
640 match world
641 .get::<crate::persistence::FinalOutput>(worker)
642 .map(|o| o.0.content.clone())
643 {
644 Some(content) => Some(Ok(content)),
645 None if worker_requires_output(world, worker) => Some(Err(
646 "worker finished without the final output its stage requires".to_string(),
647 )),
648 None => Some(Ok(world
649 .get::<InferenceResult>(worker)
650 .map(|r| r.response.clone())
651 .unwrap_or_default())),
652 }
653 }
654 Some(AgentStatus::Error { message }) => Some(Err(message)),
655 Some(AgentStatus::Cancelled) => Some(Err("worker cancelled".to_string())),
656 Some(_) => None,
657 }
658}
659
660fn worker_requires_output(world: &World, worker: Entity) -> bool {
662 let Some(bp) = world.get::<AgentBlueprint>(worker) else {
663 return false;
664 };
665 let Some(cursor) = world.get::<StageCursor>(worker) else {
666 return false;
667 };
668 bp.0.stages
669 .get(cursor.index)
670 .is_some_and(|s| s.require_output)
671}
672
673const MIN_REPORT_BYTES_PER_WORKER: usize = 200;
679
680const DEFAULT_REPORT_BYTES_PER_WORKER: usize = 4_000;
682
683const REPORT_TRUNCATION_MARKER: &str =
685 "\n[...truncated; read this worker's own run for the full answer]";
686
687fn bytes_per_worker(region_budget_tokens: Option<usize>, workers: usize) -> usize {
696 let Some(tokens) = region_budget_tokens.filter(|t| *t > 0) else {
697 return DEFAULT_REPORT_BYTES_PER_WORKER;
698 };
699 let usable = tokens.saturating_mul(4).saturating_mul(9) / 10;
702 (usable / workers.max(1)).max(MIN_REPORT_BYTES_PER_WORKER)
703}
704
705fn fit_worker_section(content: &str, budget: usize) -> String {
707 if content.len() <= budget {
708 return content.to_string();
709 }
710 let room = budget.saturating_sub(REPORT_TRUNCATION_MARKER.len());
711 format!(
712 "{}{REPORT_TRUNCATION_MARKER}",
713 leviath_core::truncate_at_boundary(content, room)
714 )
715}
716
717fn build_report(
722 summaries: &[(String, String)],
723 failures: &[(String, String)],
724 region_budget_tokens: Option<usize>,
725) -> String {
726 let sections = summaries.len().max(1);
727 let budget = bytes_per_worker(region_budget_tokens, sections);
728 let mut report = format!(
729 "[fan_out results: {} succeeded, {} failed]\n",
730 summaries.len(),
731 failures.len()
732 );
733 if summaries.iter().any(|(_, c)| c.len() > budget) {
736 report.push_str(&format!(
737 "[each worker's answer is shown up to {budget} characters; \
738 read a worker's own run for the whole thing]\n"
739 ));
740 }
741 for (id, content) in summaries {
742 report.push_str(&format!(
743 "\n## worker {id}\n{}\n",
744 fit_worker_section(content, budget)
745 ));
746 }
747 for (id, err) in failures {
748 report.push_str(&format!("\n## worker {id} FAILED\n{err}\n"));
749 }
750 report
751}
752
753fn inject_results(world: &mut World, parent: Entity, region: &str, text: &str) {
761 let Some(mut window) = world.get_mut::<ContextWindow>(parent) else {
762 return;
763 };
764 let region = match window.get_region(region).is_some() {
768 true => region,
769 false => {
770 tracing::warn!(
771 region = %region,
772 "fan-out results region is not in this agent's layout; using conversation"
773 );
774 "conversation"
775 }
776 };
777 let budget = window
778 .get_region(region)
779 .map(|r| r.max_tokens.saturating_sub(r.current_tokens))
780 .unwrap_or(0);
781 let allowed = budget.saturating_mul(4);
782 let fitted = match text.len() <= allowed {
783 true => text.to_string(),
784 false => {
785 let room = allowed.saturating_sub(REPORT_TRUNCATION_MARKER.len());
786 format!(
787 "{}{REPORT_TRUNCATION_MARKER}",
788 leviath_core::truncate_at_boundary(text, room)
789 )
790 }
791 };
792 let tokens = leviath_core::estimate_tokens(&fitted);
793 let _ = window.add_typed_entry(region, leviath_core::EntryKind::UserMessage, fitted, tokens);
794}
795
796fn agent_status(world: &World, entity: Entity) -> Option<AgentStatus> {
798 world.get::<AgentState>(entity).map(|s| s.status.clone())
799}
800
801fn set_status(world: &mut World, entity: Entity, status: AgentStatus) {
803 if let Some(mut state) = world.get_mut::<AgentState>(entity) {
804 state.status = status;
805 }
806}
807
808#[cfg(test)]
809mod tests {
810 use super::*;
811 use crate::components::{InferenceConfig, ToolResultRoutingComponent};
812 use crate::pipeline::{
813 ReadyToInfer, StageInference, StageInferences, StageProgress, StageSetup, StageSetups,
814 VisitCounts,
815 };
816 use leviath_core::blueprint::{ModelConfig, Stage};
817 use leviath_core::layout::{ContextLayout, RegionDefinition};
818 use leviath_core::{Blueprint, Region, RegionKind};
819 use std::collections::HashSet;
820
821 struct TestSpawner {
824 fail: HashSet<String>,
825 }
826
827 impl TestSpawner {
828 fn ok() -> Arc<dyn FanOutSpawner> {
829 Arc::new(TestSpawner {
830 fail: HashSet::new(),
831 })
832 }
833 fn refusing(ids: &[&str]) -> Arc<dyn FanOutSpawner> {
834 Arc::new(TestSpawner {
835 fail: ids.iter().map(|s| s.to_string()).collect(),
836 })
837 }
838 }
839
840 impl FanOutSpawner for TestSpawner {
841 fn spawn_worker(
842 &self,
843 world: &mut World,
844 _parent: Entity,
845 _config: &FanOutConfig,
846 item_id: &str,
847 _item_context: &serde_json::Value,
848 ) -> Result<Entity, String> {
849 if self.fail.contains(item_id) {
850 return Err(format!("spawn refused for '{item_id}'"));
851 }
852 Ok(world
853 .spawn((
854 AgentState {
855 agent_id: format!("worker-{item_id}"),
856 current_stage: "w".to_string(),
857 iteration: 0,
858 status: AgentStatus::Active,
859 spawned_children_ids: vec![],
860 pending_wait: None,
861 accepts_messages: true,
862 },
863 crate::persistence::RunMetadata {
866 run_id: format!("run-{item_id}"),
867 agent_name: "worker".to_string(),
868 agent_path: String::new(),
869 task: String::new(),
870 model: None,
871 workdir: String::new(),
872 num_stages: 1,
873 started_at: 0,
874 parent_run_id: None,
875 metadata: std::collections::HashMap::new(),
876 callback_url: None,
877 callback_secret: None,
878 title: None,
879 unattended: false,
880 read_paths: None,
881 output_request: None,
882 },
883 ))
884 .id())
885 }
886 }
887
888 fn cfg(merge: Option<&str>, max_workers: usize, policy: WorkerFailurePolicy) -> FanOutConfig {
889 FanOutConfig {
890 worker_agent: None,
891 worker_stage: Some("w".to_string()),
892 worker_query: None,
893 merge_stage: merge.map(String::from),
894 max_workers,
895 on_worker_failure: policy,
896 split_prompt: "split".to_string(),
897 results_region: None,
898 max_items: None,
899 }
900 }
901
902 fn window() -> ContextWindow {
903 let mut w = ContextWindow::new(12_000);
904 w.add_region(Region::new(
905 "conversation".to_string(),
906 RegionKind::Clearable,
907 10_000,
908 ));
909 w
910 }
911
912 fn stage_inf() -> StageInference {
913 StageInference {
914 provider_name: "script".to_string(),
915 model: "m".to_string(),
916 tools: vec![],
917 tool_filter: None,
918 fallbacks: Vec::new(),
919 output: None,
920 }
921 }
922
923 fn setup() -> StageSetup {
924 StageSetup {
925 inference_config: InferenceConfig {
926 temperature: None,
927 max_output_tokens: None,
928 extra_params: Default::default(),
929 batch_tool_hint: false,
930 shell_hint: false,
931 request_timeout_secs: None,
932 },
933 routing: None,
934 accepts_messages: true,
935 context_layout: None,
936 system_prompt: None,
937 output: None,
938 }
939 }
940
941 fn fanout_blueprint(config: FanOutConfig) -> Blueprint {
943 let layout = ContextLayout::new(
944 vec![RegionDefinition::new(
945 "conversation".to_string(),
946 RegionKind::Clearable,
947 10_000,
948 )],
949 12_000,
950 );
951 let mut s0 = Stage::new(
952 "fan".to_string(),
953 ModelConfig::new("script".to_string(), "m".to_string()),
954 );
955 s0.mode = StageMode::FanOut { config };
956 let s1 = Stage::new(
957 "merge".to_string(),
958 ModelConfig::new("script".to_string(), "m".to_string()),
959 );
960 Blueprint::new("t".to_string(), "d".to_string(), vec![s0, s1], layout)
961 }
962
963 fn parent_state() -> AgentState {
964 AgentState {
965 agent_id: "parent".to_string(),
966 current_stage: "fan".to_string(),
967 iteration: 0,
968 status: AgentStatus::Active,
969 spawned_children_ids: vec![],
970 pending_wait: None,
971 accepts_messages: true,
972 }
973 }
974
975 fn spawn_parent(world: &mut World, bp: Blueprint, response: &str) -> Entity {
978 world
979 .spawn((
980 AgentBlueprint(bp),
981 StageCursor { index: 0 },
982 parent_state(),
983 StageProgress::default(),
984 StageInferences(vec![stage_inf(), stage_inf()]),
985 StageSetups(vec![setup(), setup()]),
986 VisitCounts::default(),
987 window(),
988 InferenceResult {
989 response: response.to_string(),
990 tool_calls: vec![],
991 tokens_used: 0,
992 timestamp: 0,
993 },
994 ProcessResponse,
995 ))
996 .id()
997 }
998
999 fn install(world: &mut World, spawner: Arc<dyn FanOutSpawner>) {
1000 world.insert_resource(FanOutSpawnerRes(spawner));
1001 }
1002
1003 fn status_of(world: &World, e: Entity) -> AgentStatus {
1004 world.get::<AgentState>(e).unwrap().status.clone()
1005 }
1006
1007 fn assert_errored(world: &World, e: Entity) {
1010 assert_eq!(
1011 std::mem::discriminant(&status_of(world, e)),
1012 std::mem::discriminant(&AgentStatus::Error {
1013 message: String::new()
1014 })
1015 );
1016 }
1017
1018 fn complete_worker(world: &mut World, worker: Entity, content: &str) {
1019 set_status(world, worker, AgentStatus::Complete);
1020 world.entity_mut(worker).insert(InferenceResult {
1021 response: content.to_string(),
1022 tool_calls: vec![],
1023 tokens_used: 0,
1024 timestamp: 0,
1025 });
1026 }
1027
1028 #[test]
1031 fn parse_work_items_handles_array_prose_and_errors() {
1032 let ok = parse_work_items(r#"[{"id":"a"},{"id":"b","context":{"k":1}}]"#).unwrap();
1033 assert_eq!(ok.len(), 2);
1034 assert_eq!(ok[0].id, "a");
1035 assert_eq!(ok[1].context["k"], 1);
1036 assert_eq!(parse_work_items("[{}]").unwrap()[0].id, "");
1038 assert_eq!(
1040 parse_work_items("Here you go:\n```json\n[{\"id\":\"x\"}]\n```")
1041 .unwrap()
1042 .len(),
1043 1
1044 );
1045 assert!(parse_work_items("no array here").is_err());
1047 assert!(parse_work_items("]nope[").is_err());
1049 assert!(parse_work_items("[not json]").is_err());
1051 }
1052
1053 #[test]
1056 fn split_parks_a_fanout_stage_and_consumes_the_response() {
1057 let mut world = World::new();
1058 let e = spawn_parent(
1059 &mut world,
1060 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1061 r#"[{"id":"a"},{"id":"b"}]"#,
1062 );
1063 fan_out_split(&mut world);
1064 assert!(world.get::<FanOutWaiting>(e).is_some());
1065 assert_eq!(status_of(&world, e), AgentStatus::Waiting);
1066 assert!(world.get::<ProcessResponse>(e).is_none());
1068 assert!(world.get::<InferenceResult>(e).is_none());
1069 let w = world.get::<FanOutWaiting>(e).unwrap();
1070 assert_eq!(w.pending.len(), 2);
1071 }
1072
1073 #[test]
1079 fn split_keeps_only_the_first_max_items() {
1080 let mut world = World::new();
1081 let mut config = cfg(Some("merge"), 2, WorkerFailurePolicy::Continue);
1082 config.max_items = Some(3);
1083 let items: Vec<String> = (0..10).map(|i| format!(r#"{{"id":"w{i}"}}"#)).collect();
1084 let e = spawn_parent(
1085 &mut world,
1086 fanout_blueprint(config),
1087 &format!("[{}]", items.join(",")),
1088 );
1089
1090 fan_out_split(&mut world);
1091
1092 let w = world.get::<FanOutWaiting>(e).expect("parked");
1093 assert_eq!(w.pending.len(), 3, "kept the cap, not the ten produced");
1094 let kept: Vec<&str> = w.pending.iter().map(|i| i.id.as_str()).collect();
1095 assert_eq!(kept, ["w0", "w1", "w2"], "and kept the first of them");
1096 }
1097
1098 #[test]
1101 fn split_keeps_everything_under_the_cap() {
1102 let mut world = World::new();
1103 let mut config = cfg(Some("merge"), 2, WorkerFailurePolicy::Continue);
1104 config.max_items = Some(9);
1105 let e = spawn_parent(
1106 &mut world,
1107 fanout_blueprint(config),
1108 r#"[{"id":"a"},{"id":"b"}]"#,
1109 );
1110
1111 fan_out_split(&mut world);
1112
1113 assert_eq!(
1114 world.get::<FanOutWaiting>(e).expect("parked").pending.len(),
1115 2
1116 );
1117 }
1118
1119 #[test]
1120 fn split_errors_on_non_array_output() {
1121 let mut world = World::new();
1125 let e = spawn_parent(
1126 &mut world,
1127 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1128 "definitely not a json array",
1129 );
1130 for _ in 0..=MAX_SPLIT_RETRIES {
1131 fan_out_split(&mut world);
1132 redrive_split(&mut world, e, "definitely not a json array");
1133 }
1134 assert!(world.get::<FanOutWaiting>(e).is_none());
1135 assert_errored(&world, e);
1136 }
1137
1138 fn redrive_split(world: &mut World, e: Entity, response: &str) {
1141 world
1142 .entity_mut(e)
1143 .remove::<crate::pipeline::ReadyToInfer>()
1144 .insert(InferenceResult {
1145 response: response.to_string(),
1146 tool_calls: vec![],
1147 tokens_used: 0,
1148 timestamp: 0,
1149 })
1150 .insert(ProcessResponse);
1151 }
1152
1153 fn conversation_text(world: &World, e: Entity) -> String {
1154 world
1155 .get::<ContextWindow>(e)
1156 .unwrap()
1157 .get_region("conversation")
1158 .unwrap()
1159 .content
1160 .iter()
1161 .map(|entry| entry.content.clone())
1162 .collect::<Vec<_>>()
1163 .join("\n")
1164 }
1165
1166 #[test]
1169 fn a_split_that_is_not_an_array_is_corrected_rather_than_fatal() {
1170 let mut world = World::new();
1171 let e = spawn_parent(
1172 &mut world,
1173 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1174 "Sure! I will research these topics for you.",
1175 );
1176
1177 fan_out_split(&mut world);
1178
1179 assert_eq!(status_of(&world, e), AgentStatus::Active, "still running");
1180 assert_eq!(world.get::<SplitAttempts>(e), Some(&SplitAttempts(1)));
1181 assert!(
1182 world.get::<crate::pipeline::ReadyToInfer>(e).is_some(),
1183 "the parent is queued for another attempt"
1184 );
1185 assert!(world.get::<FanOutWaiting>(e).is_none());
1186 let convo = conversation_text(&world, e);
1187 assert!(
1188 convo.contains("Sure! I will research"),
1189 "the model sees its own answer: {convo}"
1190 );
1191 assert!(
1192 convo.contains("[System]") && convo.contains("start with `[`"),
1193 "and the correction: {convo}"
1194 );
1195 }
1196
1197 #[test]
1200 fn a_corrected_split_proceeds_to_the_fan_out() {
1201 let mut world = World::new();
1202 let e = spawn_parent(
1203 &mut world,
1204 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1205 "no array here",
1206 );
1207 fan_out_split(&mut world);
1208 redrive_split(&mut world, e, r#"[{"id":"a","context":{}}]"#);
1209
1210 fan_out_split(&mut world);
1211
1212 assert!(world.get::<FanOutWaiting>(e).is_some(), "the split took");
1213 assert_eq!(status_of(&world, e), AgentStatus::Waiting);
1214 }
1215
1216 #[test]
1219 fn the_failure_message_quotes_what_the_model_actually_said() {
1220 let mut world = World::new();
1221 let e = spawn_parent(
1222 &mut world,
1223 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1224 "I cannot help with that request.",
1225 );
1226 for _ in 0..=MAX_SPLIT_RETRIES {
1227 fan_out_split(&mut world);
1228 redrive_split(&mut world, e, "I cannot help with that request.");
1229 }
1230 let status = format!("{:?}", status_of(&world, e));
1233 assert!(status.contains("Error"), "{status}");
1234 assert!(status.contains("I cannot help with that"), "{status}");
1235 assert!(status.contains("correction(s)"), "{status}");
1236 }
1237
1238 #[test]
1241 fn a_split_with_nowhere_to_put_a_correction_fails_at_once() {
1242 let mut world = World::new();
1243 let e = world
1244 .spawn((
1245 AgentBlueprint(fanout_blueprint(cfg(
1246 None,
1247 2,
1248 WorkerFailurePolicy::Continue,
1249 ))),
1250 StageCursor { index: 0 },
1251 parent_state(),
1252 StageProgress::default(),
1253 StageInferences(vec![stage_inf(), stage_inf()]),
1254 StageSetups(vec![setup(), setup()]),
1255 VisitCounts::default(),
1256 InferenceResult {
1257 response: "not an array".to_string(),
1258 tool_calls: vec![],
1259 tokens_used: 0,
1260 timestamp: 0,
1261 },
1262 ProcessResponse,
1263 ))
1264 .id();
1265
1266 fan_out_split(&mut world);
1267
1268 assert_errored(&world, e);
1269 assert!(world.get::<SplitAttempts>(e).is_none());
1270 }
1271
1272 #[test]
1273 fn the_snippet_reports_an_empty_response_as_empty() {
1274 assert_eq!(response_snippet(" \n "), "the response was empty");
1275 }
1276
1277 #[test]
1278 fn the_snippet_quotes_a_short_response_whole() {
1279 assert_eq!(response_snippet(" nope "), "the response was: nope");
1280 }
1281
1282 #[test]
1283 fn the_snippet_truncates_a_long_response_on_a_character_boundary() {
1284 let long = "€".repeat(MAX_SPLIT_SNIPPET + 10);
1288 let snippet = response_snippet(&long);
1289 assert!(snippet.starts_with("the response began: "), "{snippet}");
1290 assert!(snippet.ends_with('…'), "{snippet}");
1291 let kept = snippet
1292 .trim_start_matches("the response began: ")
1293 .trim_end_matches('…');
1294 assert_eq!(kept.chars().count(), MAX_SPLIT_SNIPPET);
1295 assert!(kept.chars().all(|c| c == '€'), "no character was split");
1296 }
1297
1298 #[test]
1299 fn split_skips_non_active_and_non_fanout_agents() {
1300 let mut world = World::new();
1302 let e = spawn_parent(
1303 &mut world,
1304 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1305 "[]",
1306 );
1307 set_status(&mut world, e, AgentStatus::Idle);
1308 fan_out_split(&mut world);
1309 assert!(world.get::<ProcessResponse>(e).is_some());
1310 assert!(world.get::<FanOutWaiting>(e).is_none());
1311
1312 let layout = ContextLayout::new(
1314 vec![RegionDefinition::new(
1315 "conversation".to_string(),
1316 RegionKind::Clearable,
1317 10_000,
1318 )],
1319 12_000,
1320 );
1321 let s = Stage::new(
1322 "plain".to_string(),
1323 ModelConfig::new("script".to_string(), "m".to_string()),
1324 );
1325 let bp = Blueprint::new("t".to_string(), "d".to_string(), vec![s], layout);
1326 let e2 = spawn_parent(&mut world, bp, "[]");
1327 fan_out_split(&mut world);
1328 assert!(world.get::<ProcessResponse>(e2).is_some());
1329 }
1330
1331 #[test]
1334 fn collect_starts_workers_then_merges_on_completion() {
1335 let mut world = World::new();
1336 install(&mut world, TestSpawner::ok());
1337 let e = spawn_parent(
1338 &mut world,
1339 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1340 r#"[{"id":"a"},{"id":"b"}]"#,
1341 );
1342 fan_out_split(&mut world);
1343 fan_out_collect(&mut world);
1344 let kids = world.get::<SubAgentChildren>(e).unwrap().children.clone();
1346 assert_eq!(kids.len(), 2);
1347 assert!(world.get::<FanOutWaiting>(e).is_some());
1348 for k in &kids {
1350 assert_eq!(world.get::<ParentRef>(*k).unwrap().depth, 1);
1351 }
1352
1353 for k in &kids {
1355 complete_worker(&mut world, *k, "fixed it");
1356 }
1357 fan_out_collect(&mut world);
1358 assert!(world.get::<FanOutWaiting>(e).is_none());
1359 assert_eq!(status_of(&world, e), AgentStatus::Active);
1360 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1361 assert!(world.get::<ReadyToInfer>(e).is_some());
1362 assert!(
1364 world
1365 .get::<ContextWindow>(e)
1366 .unwrap()
1367 .get_region("conversation")
1368 .unwrap()
1369 .current_tokens
1370 > 0
1371 );
1372 }
1373
1374 fn run_slim(world: &mut World) {
1376 let mut schedule = bevy_ecs::schedule::Schedule::default();
1377 schedule.add_systems(slim_merged_workers);
1378 schedule.run(world);
1379 }
1380
1381 #[test]
1386 fn merged_workers_are_slimmed_once_their_terminal_state_is_persisted() {
1387 let mut world = World::new();
1388 install(&mut world, TestSpawner::ok());
1389 let e = spawn_parent(
1390 &mut world,
1391 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1392 r#"[{"id":"a"}]"#,
1393 );
1394 fan_out_split(&mut world);
1395 fan_out_collect(&mut world);
1396 let worker = world.get::<SubAgentChildren>(e).unwrap().children[0];
1397 world
1399 .entity_mut(worker)
1400 .insert((window(), crate::pipeline::PersistWatermark::default()));
1401 complete_worker(&mut world, worker, "done");
1402 fan_out_collect(&mut world);
1403
1404 assert!(world.get::<MergedWorker>(worker).is_some());
1407 run_slim(&mut world);
1408 assert!(
1409 world.get::<ContextWindow>(worker).is_some(),
1410 "unpersisted terminal state stays resident"
1411 );
1412
1413 let mut wm = crate::pipeline::PersistWatermark::default();
1415 wm.stamp_status(leviath_core::run_meta::RunStatus::Complete);
1416 world.entity_mut(worker).insert(wm);
1417 run_slim(&mut world);
1418 assert!(world.get::<ContextWindow>(worker).is_none());
1419 assert!(world.get::<MergedWorker>(worker).is_none());
1420 assert!(world.get::<AgentState>(worker).is_some());
1422 }
1423
1424 #[test]
1425 fn collect_respects_max_workers_and_stages_pending() {
1426 let mut world = World::new();
1427 install(&mut world, TestSpawner::ok());
1428 let e = spawn_parent(
1429 &mut world,
1430 fanout_blueprint(cfg(Some("merge"), 1, WorkerFailurePolicy::Continue)),
1431 r#"[{"id":"a"},{"id":"b"}]"#,
1432 );
1433 fan_out_split(&mut world);
1434 fan_out_collect(&mut world);
1435 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 1);
1437 let first = world.get::<SubAgentChildren>(e).unwrap().children[0];
1438 fan_out_collect(&mut world);
1441 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 1);
1442 assert!(world.get::<FanOutWaiting>(e).is_some());
1443 complete_worker(&mut world, first, "one");
1444 fan_out_collect(&mut world);
1445 assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 2);
1447 let second = world.get::<SubAgentChildren>(e).unwrap().children[1];
1448 complete_worker(&mut world, second, "two");
1449 fan_out_collect(&mut world);
1450 assert!(world.get::<FanOutWaiting>(e).is_none());
1451 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1452 }
1453
1454 #[test]
1455 fn fan_out_state_roundtrips_and_unresolved_workers_become_failures() {
1456 let mut world = World::new();
1457 install(&mut world, TestSpawner::ok());
1458 let e = spawn_parent(
1459 &mut world,
1460 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1461 r#"[{"id":"a"},{"id":"b"}]"#,
1462 );
1463 fan_out_split(&mut world);
1464 fan_out_collect(&mut world); let state = world.get::<FanOutWaiting>(e).unwrap().to_state();
1468 assert_eq!(state.active.len(), 2);
1469 assert!(state.active.iter().all(|(_id, run_id)| !run_id.is_empty()));
1470
1471 let by_run: std::collections::HashMap<String, Entity> = world
1473 .get::<SubAgentChildren>(e)
1474 .unwrap()
1475 .children
1476 .iter()
1477 .filter_map(|&c| {
1478 world
1479 .get::<crate::persistence::RunMetadata>(c)
1480 .map(|m| (m.run_id.clone(), c))
1481 })
1482 .collect();
1483 let fresh = world.spawn_empty().id();
1484 restore_fan_out_waiting(&mut world, fresh, state.clone(), &|rid| {
1485 by_run.get(rid).copied()
1486 });
1487 assert_eq!(
1488 world
1489 .get::<FanOutWaiting>(fresh)
1490 .unwrap()
1491 .to_state()
1492 .active
1493 .len(),
1494 2
1495 );
1496
1497 let orphaned = world.spawn_empty().id();
1500 restore_fan_out_waiting(&mut world, orphaned, state, &|_| None);
1501 let s = world.get::<FanOutWaiting>(orphaned).unwrap().to_state();
1502 assert!(s.active.is_empty());
1503 assert_eq!(s.failures.len(), 2);
1504 }
1505
1506 #[test]
1507 fn collect_fail_all_marks_parent_error() {
1508 let mut world = World::new();
1509 install(&mut world, TestSpawner::ok());
1510 let e = spawn_parent(
1511 &mut world,
1512 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::FailAll)),
1513 r#"[{"id":"a"}]"#,
1514 );
1515 fan_out_split(&mut world);
1516 fan_out_collect(&mut world);
1517 let worker = world.get::<SubAgentChildren>(e).unwrap().children[0];
1518 set_status(
1519 &mut world,
1520 worker,
1521 AgentStatus::Error {
1522 message: "boom".to_string(),
1523 },
1524 );
1525 fan_out_collect(&mut world);
1526 assert_errored(&world, e);
1527 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0); }
1529
1530 #[test]
1531 fn collect_continue_reports_failures_and_proceeds_without_merge() {
1532 let mut world = World::new();
1533 install(&mut world, TestSpawner::ok());
1534 let e = spawn_parent(
1536 &mut world,
1537 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue)),
1538 r#"[{"id":"a"},{"id":"b"}]"#,
1539 );
1540 fan_out_split(&mut world);
1541 fan_out_collect(&mut world);
1542 let kids = world.get::<SubAgentChildren>(e).unwrap().children.clone();
1543 set_status(
1544 &mut world,
1545 kids[0],
1546 AgentStatus::Error {
1547 message: "worker a died".to_string(),
1548 },
1549 );
1550 complete_worker(&mut world, kids[1], "b ok");
1551 fan_out_collect(&mut world);
1552 assert!(world.get::<FanOutWaiting>(e).is_none());
1553 assert!(world.get::<crate::pipeline::ResolveTransition>(e).is_some());
1554 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0);
1555 }
1556
1557 #[test]
1558 fn collect_finishes_immediately_when_there_are_no_work_items() {
1559 let mut world = World::new();
1560 install(&mut world, TestSpawner::ok());
1561 let e = spawn_parent(
1562 &mut world,
1563 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1564 "[]",
1565 );
1566 fan_out_split(&mut world);
1567 fan_out_collect(&mut world);
1568 assert!(world.get::<SubAgentChildren>(e).is_none());
1570 assert!(world.get::<FanOutWaiting>(e).is_none());
1571 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1572 }
1573
1574 #[test]
1575 fn collect_merge_stage_not_found_falls_through_to_transition() {
1576 let mut world = World::new();
1577 install(&mut world, TestSpawner::ok());
1578 let e = spawn_parent(
1579 &mut world,
1580 fanout_blueprint(cfg(Some("ghost"), 2, WorkerFailurePolicy::Continue)),
1581 "[]",
1582 );
1583 fan_out_split(&mut world);
1584 fan_out_collect(&mut world);
1585 assert!(world.get::<crate::pipeline::ResolveTransition>(e).is_some());
1587 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 0);
1588 }
1589
1590 #[test]
1591 fn collect_abandons_a_cancelled_parent() {
1592 let mut world = World::new();
1593 install(&mut world, TestSpawner::ok());
1594 let e = spawn_parent(
1595 &mut world,
1596 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1597 r#"[{"id":"a"}]"#,
1598 );
1599 fan_out_split(&mut world);
1600 set_status(&mut world, e, AgentStatus::Cancelled);
1601 fan_out_collect(&mut world);
1602 assert!(world.get::<FanOutWaiting>(e).is_none());
1603 assert_eq!(status_of(&world, e), AgentStatus::Cancelled);
1604 }
1605
1606 #[test]
1607 fn collect_without_a_spawner_records_failures() {
1608 let mut world = World::new();
1610 let e = spawn_parent(
1611 &mut world,
1612 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1613 r#"[{"id":"a"}]"#,
1614 );
1615 fan_out_split(&mut world);
1616 fan_out_collect(&mut world);
1617 assert!(world.get::<FanOutWaiting>(e).is_none());
1619 assert_eq!(world.get::<StageCursor>(e).unwrap().index, 1);
1620 }
1621
1622 #[test]
1623 fn collect_spawner_error_becomes_a_failure() {
1624 let mut world = World::new();
1625 install(&mut world, TestSpawner::refusing(&["a"]));
1626 let e = spawn_parent(
1627 &mut world,
1628 fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::FailAll)),
1629 r#"[{"id":"a"}]"#,
1630 );
1631 fan_out_split(&mut world);
1632 fan_out_collect(&mut world);
1633 assert_errored(&world, e);
1635 }
1636
1637 #[test]
1640 fn start_worker_enforces_depth_cap() {
1641 let mut world = World::new();
1642 install(&mut world, TestSpawner::ok());
1643 let mut bp = fanout_blueprint(cfg(None, 2, WorkerFailurePolicy::Continue));
1644 bp.max_child_depth = Some(3);
1645 let e = spawn_parent(&mut world, bp, r#"[{"id":"deep"}]"#);
1646 world.entity_mut(e).insert(ParentRef {
1648 parent_entity: Entity::from_raw_u32(999)
1649 .expect("a small literal index is always a valid entity id"),
1650 parent_agent_id: "root".to_string(),
1651 depth: 3,
1652 });
1653 fan_out_split(&mut world);
1654 fan_out_collect(&mut world);
1655 assert!(world.get::<SubAgentChildren>(e).is_none());
1657 assert!(world.get::<FanOutWaiting>(e).is_none());
1658 }
1659
1660 #[test]
1661 fn start_worker_uses_existing_subagentchildren_cap_and_appends() {
1662 let mut world = World::new();
1663 install(&mut world, TestSpawner::ok());
1664 let e = spawn_parent(
1665 &mut world,
1666 fanout_blueprint(cfg(Some("merge"), 2, WorkerFailurePolicy::Continue)),
1667 r#"[{"id":"a"}]"#,
1668 );
1669 world.entity_mut(e).insert(SubAgentChildren {
1671 children: vec![
1672 Entity::from_raw_u32(1000)
1673 .expect("a small literal index is always a valid entity id"),
1674 ],
1675 max_child_depth: 9,
1676 });
1677 fan_out_split(&mut world);
1678 fan_out_collect(&mut world);
1679 let kids = world.get::<SubAgentChildren>(e).unwrap();
1680 assert_eq!(kids.max_child_depth, 9);
1681 assert_eq!(kids.children.len(), 2); }
1683
1684 #[test]
1692 fn a_submitted_answer_beats_the_last_assistant_text() {
1693 let mut world = World::new();
1694 let worker = world
1695 .spawn((
1696 parent_state(),
1697 InferenceResult {
1698 response: "Let me run the tests one more time.".to_string(),
1701 tool_calls: vec![],
1702 tokens_used: 0,
1703 timestamp: 0,
1704 },
1705 crate::persistence::FinalOutput(leviath_core::output::FinalOutput::new(
1706 "changed src/lib.rs; the failing test now passes",
1707 None,
1708 "fix_worker".to_string(),
1709 0,
1710 )),
1711 ))
1712 .id();
1713 set_status(&mut world, worker, AgentStatus::Complete);
1714 assert_eq!(
1715 worker_terminal_result(&world, worker),
1716 Some(Ok(
1717 "changed src/lib.rs; the failing test now passes".to_string()
1718 ))
1719 );
1720 }
1721
1722 #[test]
1725 fn a_worker_that_submitted_nothing_still_falls_back_to_its_text() {
1726 let mut world = World::new();
1727 let worker = world
1728 .spawn((
1729 parent_state(),
1730 InferenceResult {
1731 response: "the old behaviour".to_string(),
1732 tool_calls: vec![],
1733 tokens_used: 0,
1734 timestamp: 0,
1735 },
1736 ))
1737 .id();
1738 set_status(&mut world, worker, AgentStatus::Complete);
1739 assert_eq!(
1740 worker_terminal_result(&world, worker),
1741 Some(Ok("the old behaviour".to_string()))
1742 );
1743 }
1744
1745 fn spawn_required_output_worker(world: &mut World) -> Entity {
1747 let mut stage = Stage::new(
1748 "w".to_string(),
1749 ModelConfig::new("script".to_string(), "m".to_string()),
1750 );
1751 stage.require_output = true;
1752 let layout = ContextLayout::new(
1753 vec![RegionDefinition::new(
1754 "conversation".to_string(),
1755 RegionKind::Clearable,
1756 10_000,
1757 )],
1758 12_000,
1759 );
1760 let bp = Blueprint::new("w".to_string(), "d".to_string(), vec![stage], layout);
1761 let worker = world
1762 .spawn((parent_state(), AgentBlueprint(bp), StageCursor { index: 0 }))
1763 .id();
1764 set_status(world, worker, AgentStatus::Complete);
1765 worker
1766 }
1767
1768 #[test]
1777 fn a_worker_that_owes_an_output_and_has_none_is_a_failure() {
1778 let mut world = World::new();
1779 let worker = spawn_required_output_worker(&mut world);
1780
1781 assert_eq!(
1782 worker_terminal_result(&world, worker),
1783 Some(Err(
1784 "worker finished without the final output its stage requires".to_string()
1785 )),
1786 "the merge has to be told a worker failed, and why"
1787 );
1788 }
1789
1790 #[test]
1793 fn a_worker_that_owes_an_output_and_has_one_contributes_it() {
1794 let mut world = World::new();
1795 let worker = spawn_required_output_worker(&mut world);
1796 world
1797 .entity_mut(worker)
1798 .insert(crate::persistence::FinalOutput(
1799 leviath_core::output::FinalOutput {
1800 content: "the rows".to_string(),
1801 format: Some("csv".to_string()),
1802 stage: "w".to_string(),
1803 submitted_at: 0,
1804 truncated: false,
1805 artifacts: vec![],
1806 },
1807 ));
1808
1809 assert_eq!(
1810 worker_terminal_result(&world, worker),
1811 Some(Ok("the rows".to_string()))
1812 );
1813 }
1814
1815 #[test]
1819 fn a_worker_with_no_stage_to_read_owes_nothing() {
1820 let mut world = World::new();
1821 let bp = fanout_blueprint(cfg(None, 1, WorkerFailurePolicy::Continue));
1822
1823 let bare = world.spawn(parent_state()).id();
1825 assert!(!worker_requires_output(&world, bare));
1826
1827 let no_cursor = world.spawn((parent_state(), AgentBlueprint(bp))).id();
1829 assert!(!worker_requires_output(&world, no_cursor));
1830
1831 let past_end = world
1833 .spawn((
1834 parent_state(),
1835 AgentBlueprint(fanout_blueprint(cfg(
1836 None,
1837 1,
1838 WorkerFailurePolicy::Continue,
1839 ))),
1840 StageCursor { index: 99 },
1841 ))
1842 .id();
1843 assert!(!worker_requires_output(&world, past_end));
1844 }
1845
1846 #[test]
1850 fn a_worker_that_owes_nothing_keeps_the_last_turn_fallback() {
1851 let mut world = World::new();
1852 let worker = world.spawn(parent_state()).id();
1853 set_status(&mut world, worker, AgentStatus::Complete);
1854
1855 assert_eq!(
1856 worker_terminal_result(&world, worker),
1857 Some(Ok(String::new()))
1858 );
1859 }
1860
1861 #[test]
1862 fn worker_terminal_result_covers_every_status() {
1863 let mut world = World::new();
1864 let complete = world
1865 .spawn((
1866 parent_state(),
1867 InferenceResult {
1868 response: "done text".to_string(),
1869 tool_calls: vec![],
1870 tokens_used: 0,
1871 timestamp: 0,
1872 },
1873 ))
1874 .id();
1875 set_status(&mut world, complete, AgentStatus::Complete);
1876 assert_eq!(
1877 worker_terminal_result(&world, complete),
1878 Some(Ok("done text".to_string()))
1879 );
1880
1881 let complete_no_infer = world.spawn(parent_state()).id();
1882 set_status(&mut world, complete_no_infer, AgentStatus::Complete);
1883 assert_eq!(
1884 worker_terminal_result(&world, complete_no_infer),
1885 Some(Ok(String::new()))
1886 );
1887
1888 let errored = world.spawn(parent_state()).id();
1889 set_status(
1890 &mut world,
1891 errored,
1892 AgentStatus::Error {
1893 message: "x".to_string(),
1894 },
1895 );
1896 assert_eq!(
1897 worker_terminal_result(&world, errored),
1898 Some(Err("x".to_string()))
1899 );
1900
1901 let cancelled = world.spawn(parent_state()).id();
1902 set_status(&mut world, cancelled, AgentStatus::Cancelled);
1903 assert!(worker_terminal_result(&world, cancelled).is_some_and(|r| r.is_err()));
1904
1905 let running = world.spawn(parent_state()).id(); assert_eq!(worker_terminal_result(&world, running), None);
1907
1908 assert!(
1909 worker_terminal_result(
1910 &world,
1911 Entity::from_raw_u32(4242)
1912 .expect("a small literal index is always a valid entity id")
1913 )
1914 .is_some_and(|r| r.is_err())
1915 );
1916 }
1917
1918 #[test]
1923 fn a_huge_fan_out_still_reaches_the_merge_stage() {
1924 let mut world = World::new();
1925 let mut window = ContextWindow::new(100_000);
1926 window.add_region(leviath_core::Region::new(
1927 "conversation".to_string(),
1928 leviath_core::RegionKind::Clearable,
1929 10_000,
1930 ));
1931 let parent = world.spawn((parent_state(), window)).id();
1932
1933 let huge = "x".repeat(leviath_core::output::MAX_FINAL_OUTPUT_BYTES);
1935 let summaries: Vec<(String, String)> =
1936 (0..100).map(|i| (format!("w{i}"), huge.clone())).collect();
1937 let report = build_report(&summaries, &[], Some(10_000));
1938 inject_results(&mut world, parent, "conversation", &report);
1939
1940 let region = world
1941 .get::<ContextWindow>(parent)
1942 .expect("window")
1943 .get_region("conversation")
1944 .expect("region");
1945 assert!(
1946 !region.content.is_empty(),
1947 "the merge stage must receive something rather than nothing"
1948 );
1949 let landed = ®ion.content[0].content;
1950 assert!(landed.contains("100 succeeded"), "header survives");
1953 assert!(landed.contains("truncated"), "and says it was cut");
1954 assert!(region.current_tokens <= region.max_tokens, "within budget");
1955 }
1956
1957 #[test]
1964 fn a_report_larger_than_what_is_left_of_the_region_is_trimmed_not_dropped() {
1965 const REGION_TOKENS: usize = 2_000;
1966 let mut world = World::new();
1967 let mut window = ContextWindow::new(100_000);
1968 window.add_region(leviath_core::Region::new(
1969 "worker_results".to_string(),
1970 leviath_core::RegionKind::Clearable,
1971 REGION_TOKENS,
1972 ));
1973 let filler = "f".repeat(REGION_TOKENS * 4 * 8 / 10);
1975 let filler_tokens = leviath_core::estimate_tokens(&filler);
1976 window
1977 .add_typed_entry(
1978 "worker_results",
1979 leviath_core::EntryKind::UserMessage,
1980 filler,
1981 filler_tokens,
1982 )
1983 .expect("the filler fits");
1984 let parent = world.spawn((parent_state(), window)).id();
1985
1986 let long = "x".repeat(5_000);
1988 let summaries: Vec<(String, String)> =
1989 (0..8).map(|i| (format!("w{i}"), long.clone())).collect();
1990 let report = build_report(&summaries, &[], Some(REGION_TOKENS));
1991 assert!(report.len() > REGION_TOKENS * 4 / 5, "the report is big");
1992 inject_results(&mut world, parent, "worker_results", &report);
1993
1994 let region = world
1995 .get::<ContextWindow>(parent)
1996 .expect("window")
1997 .get_region("worker_results")
1998 .expect("region")
1999 .clone();
2000 assert_eq!(
2001 region.content.len(),
2002 2,
2003 "the report landed beside the filler"
2004 );
2005 let landed = ®ion.content[1].content;
2006 assert!(
2007 landed.contains("8 succeeded"),
2008 "the header survives the cut"
2009 );
2010 assert!(
2011 landed.contains(REPORT_TRUNCATION_MARKER.trim()),
2012 "and it says it was cut"
2013 );
2014 assert!(region.current_tokens <= region.max_tokens, "within budget");
2015 }
2016
2017 #[test]
2023 fn every_worker_appears_in_a_large_fan_out() {
2024 const REGION_TOKENS: usize = 40_000;
2029 let mut world = World::new();
2030 let mut window = ContextWindow::new(400_000);
2031 window.add_region(leviath_core::Region::new(
2032 "worker_results".to_string(),
2033 leviath_core::RegionKind::Clearable,
2034 REGION_TOKENS,
2035 ));
2036 let parent = world.spawn((parent_state(), window)).id();
2037
2038 let long = "x".repeat(50_000);
2039 let summaries: Vec<(String, String)> =
2040 (0..100).map(|i| (format!("w{i}"), long.clone())).collect();
2041 let report = build_report(&summaries, &[], Some(REGION_TOKENS));
2042 inject_results(&mut world, parent, "worker_results", &report);
2043
2044 let landed = world
2045 .get::<ContextWindow>(parent)
2046 .expect("window")
2047 .get_region("worker_results")
2048 .expect("region")
2049 .content[0]
2050 .content
2051 .clone();
2052 for i in 0..100 {
2053 assert!(
2054 landed.contains(&format!("## worker w{i}\n")),
2055 "worker w{i} never reached the merge stage"
2056 );
2057 }
2058 assert!(landed.contains("read a worker's own run"));
2061 }
2062
2063 #[test]
2065 fn the_share_shrinks_as_the_worker_count_grows() {
2066 assert!(bytes_per_worker(Some(40_000), 4) > bytes_per_worker(Some(40_000), 100));
2067 assert!(bytes_per_worker(Some(80_000), 10) > bytes_per_worker(Some(40_000), 10));
2069 assert_eq!(
2071 bytes_per_worker(Some(10), 10_000),
2072 MIN_REPORT_BYTES_PER_WORKER
2073 );
2074 assert_eq!(bytes_per_worker(None, 4), DEFAULT_REPORT_BYTES_PER_WORKER);
2076 }
2077
2078 #[test]
2081 fn results_go_to_the_named_region() {
2082 let mut world = World::new();
2083 let mut window = ContextWindow::new(100_000);
2084 window.add_region(leviath_core::Region::new(
2085 "conversation".to_string(),
2086 leviath_core::RegionKind::Clearable,
2087 10_000,
2088 ));
2089 window.add_region(leviath_core::Region::new(
2090 "worker_results".to_string(),
2091 leviath_core::RegionKind::Clearable,
2092 20_000,
2093 ));
2094 let parent = world.spawn((parent_state(), window)).id();
2095 inject_results(&mut world, parent, "worker_results", "the report");
2096
2097 let w = world.get::<ContextWindow>(parent).expect("window");
2098 assert_eq!(
2099 w.get_region("worker_results")
2100 .expect("region")
2101 .content
2102 .len(),
2103 1
2104 );
2105 assert!(
2106 w.get_region("conversation")
2107 .expect("region")
2108 .content
2109 .is_empty(),
2110 "the default region is left alone"
2111 );
2112 }
2113
2114 #[test]
2117 fn an_unknown_results_region_falls_back_to_the_conversation() {
2118 let mut world = World::new();
2119 let mut window = ContextWindow::new(100_000);
2120 window.add_region(leviath_core::Region::new(
2121 "conversation".to_string(),
2122 leviath_core::RegionKind::Clearable,
2123 10_000,
2124 ));
2125 let parent = world.spawn((parent_state(), window)).id();
2126 inject_results(&mut world, parent, "typo_region", "the report");
2127
2128 assert_eq!(
2129 world
2130 .get::<ContextWindow>(parent)
2131 .expect("window")
2132 .get_region("conversation")
2133 .expect("region")
2134 .content
2135 .len(),
2136 1
2137 );
2138 }
2139
2140 #[test]
2143 fn a_small_fan_out_report_is_not_trimmed() {
2144 let mut world = World::new();
2145 let mut window = ContextWindow::new(100_000);
2146 window.add_region(leviath_core::Region::new(
2147 "conversation".to_string(),
2148 leviath_core::RegionKind::Clearable,
2149 10_000,
2150 ));
2151 let parent = world.spawn((parent_state(), window)).id();
2152 let report = build_report(
2153 &[("a".to_string(), "did the thing".to_string())],
2154 &[],
2155 Some(10_000),
2156 );
2157 inject_results(&mut world, parent, "conversation", &report);
2158 let landed = world
2159 .get::<ContextWindow>(parent)
2160 .expect("window")
2161 .get_region("conversation")
2162 .expect("region")
2163 .content[0]
2164 .content
2165 .clone();
2166 assert_eq!(landed, report);
2167 }
2168
2169 #[test]
2170 fn build_report_lists_successes_and_failures() {
2171 let report = build_report(
2172 &[("a".to_string(), "ok-a".to_string())],
2173 &[("b".to_string(), "boom".to_string())],
2174 None,
2175 );
2176 assert!(report.contains("1 succeeded, 1 failed"));
2177 assert!(report.contains("## worker a\nok-a"));
2178 assert!(report.contains("## worker b FAILED\nboom"));
2179 }
2180
2181 #[test]
2182 fn inject_conversation_is_a_noop_without_a_window() {
2183 let mut world = World::new();
2184 let has_window = world.spawn(window()).id();
2185 inject_results(&mut world, has_window, "conversation", "hello");
2186 assert!(
2187 world
2188 .get::<ContextWindow>(has_window)
2189 .unwrap()
2190 .get_region("conversation")
2191 .unwrap()
2192 .current_tokens
2193 > 0
2194 );
2195 let no_window = world.spawn(parent_state()).id();
2197 inject_results(&mut world, no_window, "conversation", "hello");
2198 }
2199
2200 #[test]
2201 fn set_status_is_a_noop_for_a_missing_agent() {
2202 let mut world = World::new();
2203 set_status(
2204 &mut world,
2205 Entity::from_raw_u32(77).expect("a small literal index is always a valid entity id"),
2206 AgentStatus::Complete,
2207 );
2208 assert_eq!(
2209 agent_status(
2210 &world,
2211 Entity::from_raw_u32(77)
2212 .expect("a small literal index is always a valid entity id")
2213 ),
2214 None
2215 );
2216 }
2217
2218 #[test]
2221 fn force_transition_applies_routing_and_handles_despawn_and_overflow() {
2222 use crate::pipeline::force_transition;
2223 let mut world = World::new();
2225 let mut setups = vec![setup(), setup()];
2226 setups[1].routing = Some(leviath_core::ToolResultRouting::default());
2227 let e = world
2228 .spawn((
2229 AgentBlueprint(fanout_blueprint(cfg(
2230 Some("merge"),
2231 2,
2232 WorkerFailurePolicy::Continue,
2233 ))),
2234 StageCursor { index: 0 },
2235 parent_state(),
2236 StageProgress::default(),
2237 StageInferences(vec![stage_inf(), stage_inf()]),
2238 StageSetups(setups),
2239 VisitCounts::default(),
2240 window(),
2241 ))
2242 .id();
2243 let agent = crate::world::AgentId::in_world(&world, e);
2244 force_transition(&mut world, agent, 1);
2245 assert!(world.get::<ToolResultRoutingComponent>(e).is_some());
2246 assert!(world.get::<ReadyToInfer>(e).is_some());
2247
2248 let gone = crate::world::AgentId::in_world(
2250 &world,
2251 Entity::from_raw_u32(9191).expect("a small literal index is always a valid entity id"),
2252 );
2253 force_transition(&mut world, gone, 1);
2254 }
2255
2256 #[test]
2257 fn force_transition_marks_error_on_prompt_overflow() {
2258 use crate::pipeline::force_transition;
2259 let layout = ContextLayout::new(
2261 vec![RegionDefinition::new(
2262 "task".to_string(),
2263 RegionKind::Pinned,
2264 20,
2265 )],
2266 1000,
2267 );
2268 let mut s0 = Stage::new(
2269 "fan".to_string(),
2270 ModelConfig::new("script".to_string(), "m".to_string()),
2271 );
2272 s0.mode = StageMode::FanOut {
2273 config: cfg(Some("merge"), 2, WorkerFailurePolicy::Continue),
2274 };
2275 let mut s1 = Stage::new(
2276 "merge".to_string(),
2277 ModelConfig::new("script".to_string(), "m".to_string()),
2278 );
2279 s1.config.insert(
2280 "system_prompt".to_string(),
2281 serde_json::Value::String("x".repeat(10_000)),
2282 );
2283 let bp = Blueprint::new("t".to_string(), "d".to_string(), vec![s0, s1], layout);
2284
2285 let mut setups = vec![setup(), setup()];
2286 setups[1].system_prompt = Some("x".repeat(10_000));
2287 let mut w = ContextWindow::new(1000);
2288 w.add_region(Region::new("task".to_string(), RegionKind::Pinned, 20));
2289 let (mut world, e) = world_with(bp, setups, w);
2290 let agent = crate::world::AgentId::in_world(&world, e);
2291 force_transition(&mut world, agent, 1);
2292 assert_errored(&world, e);
2293 }
2294
2295 fn world_with(bp: Blueprint, setups: Vec<StageSetup>, w: ContextWindow) -> (World, Entity) {
2297 let mut world = World::new();
2298 let e = world
2299 .spawn((
2300 AgentBlueprint(bp),
2301 StageCursor { index: 0 },
2302 parent_state(),
2303 StageProgress::default(),
2304 StageInferences(vec![stage_inf(), stage_inf()]),
2305 StageSetups(setups),
2306 VisitCounts::default(),
2307 w,
2308 ))
2309 .id();
2310 (world, e)
2311 }
2312}