Skip to main content

leviath_runtime/
fanout.rs

1//! Fan-out stage handling as ECS systems.
2//!
3//! A `fan_out` stage (see [`leviath_core::blueprint::StageMode::FanOut`]) runs
4//! its single inference as a **split** - its prompt (with the config's
5//! `split_prompt` folded in by [`crate::pipeline`]) asks the model for a JSON
6//! array of work items. [`fan_out_split`] intercepts that response (before the
7//! normal `process_response` routing), parses the items, and parks the parent in
8//! [`FanOutWaiting`]. [`fan_out_collect`] then starts one worker per item -
9//! bounded by `max_workers` concurrent workers - via the daemon-installed
10//! [`FanOutSpawner`], tracks them as the parent's `SubAgentChildren`, and once
11//! every worker is terminal applies the failure policy, injects a consolidated
12//! report into the parent's conversation, and transitions to the `merge_stage`
13//! (or falls through to the stage's normal transition).
14//!
15//! The runtime only **starts and tracks** workers; resolving *which* blueprint a
16//! worker runs (self-at-worker-stage, a named agent, or a capability query) is
17//! the CLI's job, encapsulated behind the [`FanOutSpawner`] it installs.
18
19use 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
30/// Depth cap for fan-out workers when the parent's blueprint doesn't set one.
31const DEFAULT_FANOUT_DEPTH: usize = 3;
32
33/// How many times a malformed split is sent back to the model before the run
34/// fails.
35///
36/// A split asks for one exact shape, and a model that answers with prose or an
37/// apology has not failed at the work, only at the format. Failing the run on
38/// the first such answer throws away everything the parent has done, which is
39/// what a deep-researcher run reported: one non-conforming response ended it.
40/// Two corrections is enough to clear a formatting slip without letting a model
41/// that cannot produce the shape loop for ever.
42const MAX_SPLIT_RETRIES: usize = 2;
43
44/// How many corrective attempts a parent's split has already had.
45///
46/// Absent until the first malformed split, so a split that parses first time
47/// costs nothing.
48#[derive(Component, Debug, Clone, Copy, Default, PartialEq, Eq)]
49pub struct SplitAttempts(pub usize);
50
51/// What the model is told after a split that could not be parsed.
52///
53/// It names the failure and restates the shape rather than repeating the
54/// original instruction, because the original instruction is what just did not
55/// work.
56fn 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
64/// The first `MAX_SPLIT_SNIPPET` characters of what the model actually said,
65/// for the failure message.
66///
67/// The old message named the rule that was broken but never what came back, so
68/// an operator reading `split output is not a JSON array` could not tell a
69/// refusal from an empty response from prose. Bounded because a split response
70/// can be long, and truncated on a character boundary because model output is
71/// arbitrary UTF-8.
72fn 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    // Taken as characters rather than bytes: a byte ceiling can land inside a
78    // character, and model output is arbitrary UTF-8.
79    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
86/// How much of a failed split response the error message carries, in characters.
87const MAX_SPLIT_SNIPPET: usize = 200;
88
89/// One unit of work produced by a fan-out split.
90#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, Default)]
91pub struct WorkItem {
92    /// Stable id (used to label the worker in the consolidated report).
93    #[serde(default)]
94    pub id: String,
95    /// Free-form context handed to the worker (seeded into its pinned context).
96    #[serde(default)]
97    pub context: serde_json::Value,
98}
99
100/// Parse a split response into work items. Tolerates markdown fences and prose by
101/// extracting the outermost `[ … ]`. (Ported from the deleted imperative engine.)
102pub fn parse_work_items(content: &str) -> Result<Vec<WorkItem>, String> {
103    let trimmed = content.trim();
104    // Every rejection folds into one error: this parses model output, so
105    // "malformed input yields `Err`" has to hold for every shape of malformed.
106    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
115/// Starts one worker for a fan-out work item. The implementor resolves the
116/// worker's blueprint (per `config`'s `worker_stage` / `worker_agent` /
117/// `worker_query`), spawns it into `world` seeded with the work item, and returns
118/// the child entity. Parent/child linking is done by [`fan_out_collect`], not the
119/// spawner.
120pub trait FanOutSpawner: Send + Sync {
121    /// Spawn one worker under `parent` for the given work item, or `Err` with a
122    /// human-readable reason (recorded as that item's failure).
123    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/// The installed [`FanOutSpawner`], as a world resource. Absent in a pure-runtime
134/// world (then every fan-out item fails with "no fan-out spawner installed").
135#[derive(Resource, Clone)]
136pub struct FanOutSpawnerRes(pub Arc<dyn FanOutSpawner>);
137
138/// A currently-running fan-out worker: its work-item id, its live entity, and
139/// its run-id (kept so the waiting state can be persisted/restored without a
140/// cross-entity lookup - see [`FanOutState`]).
141struct ActiveWorker {
142    item_id: String,
143    entity: Entity,
144    run_id: String,
145}
146
147/// A parent parked while its fan-out workers run. Holds the not-yet-started
148/// `pending` items, the currently-`active` workers, and the accumulated results.
149#[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/// The serializable form of [`FanOutWaiting`], written to `<run_dir>/fanout.json`
160/// so a parent interrupted mid-split resumes its merge after a restart. `active`
161/// carries worker **run-ids** (not entities); recovery maps them back to the
162/// reloaded worker entities.
163#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
164pub struct FanOutState {
165    /// The fan-out configuration.
166    pub config: FanOutConfig,
167    /// The concurrency cap.
168    pub max_workers: usize,
169    /// Work items not yet started.
170    pub pending: Vec<WorkItem>,
171    /// In-flight workers as `(item_id, run_id)`.
172    pub active: Vec<(String, String)>,
173    /// Completed worker results as `(item_id, summary)`.
174    pub summaries: Vec<(String, String)>,
175    /// Failed worker results as `(item_id, message)`.
176    pub failures: Vec<(String, String)>,
177}
178
179impl FanOutWaiting {
180    /// Workers this parent is still parked on: in-flight plus not-yet-started.
181    ///
182    /// Surfaced by `lev ps` so "waiting" on a fan-out parent reads as progress
183    /// against a known denominator rather than an unexplained stall.
184    pub fn outstanding(&self) -> usize {
185        self.active.len() + self.pending.len()
186    }
187
188    /// Project to the serializable [`FanOutState`] (workers by run-id).
189    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
205/// Rebuild a parent's [`FanOutWaiting`] from a persisted [`FanOutState`] and
206/// insert it, mapping each active worker's run-id back to its reloaded entity
207/// via `resolve`. Workers whose entity didn't reload are treated as failures so
208/// the merge still completes rather than waiting forever. Used by restart
209/// recovery to resume an interrupted fan-out.
210pub 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
238/// Fan-out split system (exclusive): for each `ProcessResponse` agent whose
239/// current stage is a fan-out stage, consume its response as the split output -
240/// parse the work items and park the agent in [`FanOutWaiting`] (or mark it
241/// `Error` if the split output isn't a JSON array). Removing `ProcessResponse`
242/// here keeps the normal `process_response` routing from touching these agents.
243pub 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                // A split decides its own item count, so without a cap a model
274                // that returns five hundred items spawns five hundred runs. The
275                // cap also fixes each worker's share of the results region: past
276                // some number of ways to divide it, every section is too small
277                // to say anything.
278                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                // A split that did not parse is a formatting miss, not a dead
301                // run: send the model its own answer plus a correction and ask
302                // again, the way every other stage handles a response it cannot
303                // use. Only once the corrections are spent does the run fail.
304                let attempts = world.get::<SplitAttempts>(parent).map_or(0, |a| a.0);
305                // Correcting needs somewhere to put the correction, so an agent
306                // with no context window falls through to the failure below
307                // rather than looping without ever being told anything.
308                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                                // The model sees what it said before the
316                                // correction, so "reply with only the array"
317                                // has something to correct.
318                                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                    // Name what came back as well as the rule it broke: the
349                    // rule alone cannot tell a refusal from an empty response
350                    // from prose.
351                    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
368/// Fan-out collect system (exclusive): drive each [`FanOutWaiting`] parent - reap
369/// finished workers, start pending ones up to `max_workers`, and once none remain
370/// running apply the failure policy, inject the consolidated report, and
371/// transition to the merge stage (or resolve the stage's own transition).
372pub 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        // A cancelled/errored parent abandons the fan-out; its workers are reaped
382        // by the host's cascade cancel (which walks SubAgentChildren).
383        if !matches!(agent_status(world, parent), Some(AgentStatus::Waiting)) {
384            world.entity_mut(parent).remove::<FanOutWaiting>();
385            continue;
386        }
387        // A `Waiting` parent from the query above still holds its `FanOutWaiting`
388        // (only this system removes it, and each entity appears once per pass).
389        let mut w = world
390            .entity_mut(parent)
391            .take::<FanOutWaiting>()
392            .expect("a Waiting fan-out parent still holds FanOutWaiting");
393
394        // 1. Reap workers that have reached a terminal state. A consumed
395        // worker's result now lives in `w.summaries`/`w.failures`, so its heavy
396        // components are dead weight - mark it for `slim_merged_workers`, which
397        // drops them once the terminal snapshot has reached the persistence
398        // lane. The entity itself stays (the host only despawns it when the
399        // parent goes terminal), but without its context window: previously
400        // every finished fan-out worker kept a full window resident for the
401        // whole remainder of the parent's run.
402        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        // 2. Start pending workers up to the concurrency cap.
418        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                    // Capture the worker's run-id so the waiting state persists.
425                    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        // 3. Finished when nothing is running or queued.
440        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/// A fan-out worker whose terminal result the parent has already consumed.
449/// Set by [`fan_out_collect`]; consumed by [`slim_merged_workers`].
450#[derive(Component)]
451pub struct MergedWorker;
452
453/// Drop a merged worker's heavy components once its terminal snapshot has
454/// reached the persistence lane.
455///
456/// Ordering makes this safe on both sides: the marker is only set after the
457/// parent consumed the worker's result (so the merge no longer reads the
458/// worker), and the watermark gate (`PersistWatermark::persisted_status`)
459/// holds the slim back until the terminal state is on its way to disk (so
460/// nothing readable is lost - the entity's remaining metadata still identifies
461/// the run, and its full final state is in the run dir).
462pub 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; // the terminal snapshot has not been dispatched yet
479        }
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
491/// Apply the failure policy, inject the consolidated report, and transition.
492fn 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    // Where the results land, and how much room they have there. A blueprint
508    // that names a region of its own gets that region's budget to divide; the
509    // default is the conversation region, which is also carrying the message
510    // history.
511    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(&region).map(|r| r.max_tokens));
519    let report = build_report(&w.summaries, &w.failures, budget);
520    inject_results(world, parent, &region, &report);
521
522    // Ready the parent to run again, then jump to the merge stage (if any) or let
523    // the fan-out stage's own transition resolve.
524    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
541/// Start one worker and link it to `parent` (`ParentRef` + `SubAgentChildren`),
542/// enforcing the parent blueprint's child-depth cap. Returns the child entity.
543fn 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    // Record the worker's run-id on the parent's serializable state so the tree
591    // (fan-out workers included) is persisted for a deterministic restart rebuild.
592    // A freshly spawned worker always has run metadata; its parent always has state.
593    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    // Seed the worker's context from the parent per any declared blueprint
604    // context transform (when a fan-out worker runs a different blueprint).
605    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
613/// A worker's terminal result: `Some(Ok(deliverable))` if complete,
614/// `Some(Err(reason))` if it errored/was cancelled/vanished, `None` if still
615/// running.
616///
617/// A worker that called `submit_output` contributes exactly what it submitted.
618/// Otherwise this falls back to the text of its last assistant message, which is
619/// what every worker used to contribute and is usually wrong: a worker whose
620/// final turn was a tool call has no trailing text, so the merge stage received
621/// an empty string, which is silently indistinguishable from a worker that had
622/// nothing to say.
623///
624/// The fallback stays because it costs nothing and an existing blueprint that
625/// happens to end on a text turn keeps working. A blueprint that wants the
626/// guarantee sets `require_output` on its worker stage.
627///
628/// A worker whose stage set `require_output` and that finished without one is
629/// reported as a **failure**, not as a success with empty content. It reached
630/// `Complete` either way - the enforcement loop proceeds rather than stranding
631/// the run, and a worker that burns its iterations against a validator it cannot
632/// satisfy ends the same way. Counting that as success is how a fan-out reports
633/// "10 succeeded, 0 failed" over ten empty sections, which is worse than an
634/// error: the merge stage cannot tell an empty answer from a missing one, so it
635/// writes a confident merge of nothing.
636fn 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
660/// Whether the stage this worker is sitting in demands a final output.
661fn 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
673/// Smallest per-worker share worth writing, in bytes.
674///
675/// Below this a section says nothing useful, and the honest move is to tell the
676/// merge stage that the results are too many to carry rather than hand it a
677/// hundred fragments. That is what `max_items` on the fan-out config is for.
678const MIN_REPORT_BYTES_PER_WORKER: usize = 200;
679
680/// Per-worker share when the results region's budget cannot be read.
681const DEFAULT_REPORT_BYTES_PER_WORKER: usize = 4_000;
682
683/// Marker appended to a worker's section that was cut to fit the report.
684const REPORT_TRUNCATION_MARKER: &str =
685    "\n[...truncated; read this worker's own run for the full answer]";
686
687/// How many bytes each worker's section may use, given the region's token
688/// budget and how many workers there are.
689///
690/// An equal share, so every worker appears. The first cut at this capped each
691/// worker at a fixed size and then trimmed the finished report to fit, which
692/// meant the early workers got their full allowance and the late ones were cut
693/// off entirely - a hundred-way fan-out where only the first twenty were
694/// readable, with nothing saying so.
695fn 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    // The workspace's bytes-over-four estimate, minus a margin for the header
700    // and the per-worker `## worker <id>` lines.
701    let usable = tokens.saturating_mul(4).saturating_mul(9) / 10;
702    (usable / workers.max(1)).max(MIN_REPORT_BYTES_PER_WORKER)
703}
704
705/// One worker's contribution, trimmed to `budget` bytes.
706fn 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
717/// Build the consolidated `[fan_out results: …]` report from worker outcomes.
718///
719/// `region_budget_tokens` is the results region's budget, which the workers'
720/// sections divide equally between them.
721fn 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    // Say the share out loud when it is tight, so the merge stage knows it is
734    // reading extracts and can go to a worker's own run for the rest.
735    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
753/// Add `text` to the parent's results region, trimming it to fit.
754///
755/// The write used to be best-effort in the worst sense: `add_entry` rejects an
756/// over-budget entry outright, and the error was discarded, so a report too big
757/// for the region left the merge stage with nothing and said nothing about it.
758/// Trimming first means the merge always receives *something*, and a report that
759/// had to be cut says so where the model will read it.
760fn 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    // A named region the layout does not declare would silently swallow the
765    // whole report, so fall back to the one every agent has. `lev validate`
766    // catches the typo before a run gets here.
767    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
796/// An agent's status, if it still exists.
797fn agent_status(world: &World, entity: Entity) -> Option<AgentStatus> {
798    world.get::<AgentState>(entity).map(|s| s.status.clone())
799}
800
801/// Set an agent's status (no-op if it despawned).
802fn 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    /// A spawner that spawns a trivial `Active` worker per item, refusing the ids
822    /// in `fail`.
823    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                    // A real worker carries run metadata (attached by build_agent);
864                    // mirror that so the parent can record the worker's run-id.
865                    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    /// A blueprint whose stage 0 is a fan-out stage and stage 1 is `merge`.
942    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    /// Spawn a parent sitting on `ProcessResponse` with `response` as its
976    /// (split) inference output.
977    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    /// Assert an agent is in an `Error` state (by discriminant, so no unmatched
1008    /// `matches!` arm is left uncovered).
1009    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    // ── parse_work_items ──────────────────────────────────────────────────────
1029
1030    #[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        // Missing fields default.
1037        assert_eq!(parse_work_items("[{}]").unwrap()[0].id, "");
1038        // Prose around the array is tolerated.
1039        assert_eq!(
1040            parse_work_items("Here you go:\n```json\n[{\"id\":\"x\"}]\n```")
1041                .unwrap()
1042                .len(),
1043            1
1044        );
1045        // No brackets at all.
1046        assert!(parse_work_items("no array here").is_err());
1047        // Closing before opening (e <= s).
1048        assert!(parse_work_items("]nope[").is_err());
1049        // Brackets but not valid JSON.
1050        assert!(parse_work_items("[not json]").is_err());
1051    }
1052
1053    // ── fan_out_split ─────────────────────────────────────────────────────────
1054
1055    #[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        // ProcessResponse + InferenceResult were consumed.
1067        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    /// `max_items` is a ceiling on slices, not just on concurrency. A split that
1074    /// returns five hundred items would otherwise spawn five hundred runs, and
1075    /// each worker's share of the results region is the region's budget divided
1076    /// by how many there are: past some count every section is too small to say
1077    /// anything.
1078    #[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    /// Under the cap nothing is dropped, so a fan-out that sets one does not pay
1099    /// for it on every ordinary split.
1100    #[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        // The corrections have to be spent before the run dies, so this drives
1122        // the split until they are. Failing on the first answer is the bug the
1123        // retry exists to fix.
1124        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    /// Put the parent back where a fresh inference would leave it, so the split
1139    /// can be driven a second and third time without a real provider.
1140    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    /// The fix for the deep-researcher report: a split that comes back as prose
1167    /// is a formatting miss, and the run gets to correct it instead of dying.
1168    #[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    /// A correction that works is the whole point: the second answer parses and
1198    /// the run carries on into its fan-out.
1199    #[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    /// Once the corrections are spent the run does fail, and the message names
1217    /// what came back rather than only the rule it broke.
1218    #[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        // Read through Debug rather than a pattern: the arm a passing run does
1231        // not take reads to llvm-cov as an uncovered region.
1232        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    /// No context window means nowhere to put a correction, so the run fails on
1239    /// the first malformed split rather than looping while being told nothing.
1240    #[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        // Three bytes per character on purpose: cutting model output at a byte
1285        // offset is how a panic gets shipped, so the ceiling counts characters
1286        // and this proves no character was split.
1287        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        // Non-Active fan-out agent: left untouched.
1301        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        // Non-fan-out stage: not a candidate at all.
1313        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    // ── fan_out_collect: worker lifecycle + merge ─────────────────────────────
1332
1333    #[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        // Two workers started and tracked.
1345        let kids = world.get::<SubAgentChildren>(e).unwrap().children.clone();
1346        assert_eq!(kids.len(), 2);
1347        assert!(world.get::<FanOutWaiting>(e).is_some());
1348        // Each worker got a ParentRef at depth 1.
1349        for k in &kids {
1350            assert_eq!(world.get::<ParentRef>(*k).unwrap().depth, 1);
1351        }
1352
1353        // Complete both workers, then collect merges to the merge stage.
1354        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        // The consolidated report landed in the parent's conversation.
1363        assert!(
1364            world
1365                .get::<ContextWindow>(e)
1366                .unwrap()
1367                .get_region("conversation")
1368                .unwrap()
1369                .current_tokens
1370                > 0
1371        );
1372    }
1373
1374    /// Run the slim system once over `world`.
1375    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    /// A merged worker keeps its heavy components until its terminal snapshot
1382    /// has been dispatched, then sheds them - previously every finished
1383    /// fan-out worker kept a full context window resident until the parent
1384    /// went terminal.
1385    #[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        // Give the worker a context window so there is something to shed.
1398        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        // Consumed by the merge and marked - but its terminal snapshot has not
1405        // been dispatched, so it keeps its state.
1406        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        // Stamp the watermark terminal, and the worker sheds its heavy parts.
1414        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        // The entity itself survives for the host's bookkeeping.
1421        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        // Only one worker at a time.
1436        assert_eq!(world.get::<SubAgentChildren>(e).unwrap().children.len(), 1);
1437        let first = world.get::<SubAgentChildren>(e).unwrap().children[0];
1438        // A collect pass while the worker is still running keeps it active and
1439        // starts nothing new (worker still counts against max_workers).
1440        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        // Second worker started after the first finished.
1446        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); // starts both workers → active
1465
1466        // Projecting to the serializable state captures each worker's run-id.
1467        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        // Restore onto a fresh parent, resolving run-ids back to entities.
1472        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        // A resolver that can't map the workers → they become failures, so the
1498        // merge still completes rather than waiting forever.
1499        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); // no merge
1528    }
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        // No merge stage ⇒ ResolveTransition (proceed) rather than force_transition.
1535        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        // No workers; straight to merge.
1569        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        // Unknown merge stage ⇒ ResolveTransition, no stage jump.
1586        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        // No FanOutSpawnerRes installed ⇒ every item fails to start.
1609        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        // Item failed to start, Continue policy ⇒ still transitions to merge.
1618        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        // Spawn refused + FailAll ⇒ parent errors.
1634        assert_errored(&world, e);
1635    }
1636
1637    // ── start_worker: depth cap + existing SubAgentChildren ───────────────────
1638
1639    #[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        // Parent is itself a depth-3 sub-agent ⇒ child would be depth 4 > 3.
1647        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        // No worker spawned (depth cap hit before any container is created).
1656        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        // Pre-existing children container with a generous cap.
1670        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); // appended to the existing one
1682    }
1683
1684    // ── worker_terminal_result / build_report / inject_conversation ───────────
1685
1686    /// The bug this feature exists to fix. A worker's contribution used to be
1687    /// the text of its last assistant message, so a worker whose final turn was
1688    /// a tool call contributed an empty string - and the shipped
1689    /// a worker told to report what it did writes that report into exactly that
1690    /// channel.
1691    #[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                    // What the old code would have handed the merge stage: the
1699                    // trailing aside, not the deliverable.
1700                    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    /// The fallback stays, so a blueprint that happens to end on a text turn
1723    /// keeps working without declaring anything.
1724    #[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    /// Spawn a worker sitting in a stage that demands a final output.
1746    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    /// The fan-out reported "10 succeeded, 0 failed" over ten empty sections,
1769    /// because a worker that reached `Complete` without its required output was
1770    /// read as a success with nothing to say. The merge stage cannot tell those
1771    /// apart, so it writes a confident merge of nothing.
1772    ///
1773    /// This is the ordinary way it happens, not an edge case: a worker that
1774    /// cannot satisfy its validator retries until its iterations run out and
1775    /// leaves on the max-iterations path, which ends at `Complete`.
1776    #[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    /// The same worker, having actually submitted: its answer is what it
1791    /// contributes, and the requirement is discharged.
1792    #[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    /// A worker with a blueprint but no cursor cannot be placed in a stage, so
1816    /// there is no stage to read a requirement off. It keeps the fallback rather
1817    /// than being called a failure for a question that was never asked.
1818    #[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        // No blueprint at all.
1824        let bare = world.spawn(parent_state()).id();
1825        assert!(!worker_requires_output(&world, bare));
1826
1827        // A blueprint, but no cursor saying which stage it is in.
1828        let no_cursor = world.spawn((parent_state(), AgentBlueprint(bp))).id();
1829        assert!(!worker_requires_output(&world, no_cursor));
1830
1831        // A cursor pointing past the end of the stage list.
1832        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    /// A blueprint that never opted in keeps the old fallback, empty text and
1847    /// all. Turning that into a failure would break every fan-out written before
1848    /// `require_output` existed.
1849    #[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(); // Active
1906        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    /// The failure this bound exists for. A hundred workers answering at the
1919    /// size limit build a 25 MB report; `add_entry` rejects an over-budget entry
1920    /// rather than truncating, and the error was discarded - so the merge stage
1921    /// received nothing at all, silently, in exactly the case fan-out is for.
1922    #[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        // A hundred workers, each answering at the per-submission cap.
1934        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 = &region.content[0].content;
1950        // Every worker is still accounted for in the header, and the text says
1951        // it was cut rather than pretending to be whole.
1952        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    /// Dividing the region between the workers makes the report fit an *empty*
1958    /// region, which is the easy case. A region already carrying something has
1959    /// less room than that, and the report-level trim is what keeps the write
1960    /// from being rejected outright: `add_entry` refuses an over-budget entry
1961    /// rather than shortening it, so without this the merge stage receives
1962    /// nothing at all.
1963    #[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        // Most of the region is already spoken for.
1974        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        // A report sized for the whole region, landing in what is left of it.
1987        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 = &region.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    /// The share is equal, so every worker appears. The first cut capped each
2018    /// worker at a fixed size and trimmed the finished report to fit, which gave
2019    /// the early workers their full allowance and cut the late ones off
2020    /// entirely - a hundred-way fan-out where only the first twenty were
2021    /// readable, with nothing saying so.
2022    #[test]
2023    fn every_worker_appears_in_a_large_fan_out() {
2024        // End to end: building the report and landing it in the region. The
2025        // unfairness was in the second half - a fixed per-worker size makes a
2026        // report far too big, and trimming *that* keeps the front and drops the
2027        // back.
2028        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        // And it says the sections are extracts, so the merge stage knows to go
2059        // to a worker's own run for the rest.
2060        assert!(landed.contains("read a worker's own run"));
2061    }
2062
2063    /// Each worker gets the same room, whatever the count.
2064    #[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        // A bigger region means a bigger share for the same workers.
2068        assert!(bytes_per_worker(Some(80_000), 10) > bytes_per_worker(Some(40_000), 10));
2069        // Never so small a section says nothing at all.
2070        assert_eq!(
2071            bytes_per_worker(Some(10), 10_000),
2072            MIN_REPORT_BYTES_PER_WORKER
2073        );
2074        // No readable budget falls back rather than dividing by nothing.
2075        assert_eq!(bytes_per_worker(None, 4), DEFAULT_REPORT_BYTES_PER_WORKER);
2076    }
2077
2078    /// A blueprint can send the results somewhere other than the conversation,
2079    /// which is otherwise carrying the message history alongside them.
2080    #[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    /// A named region the layout does not declare falls back rather than
2115    /// swallowing the whole report.
2116    #[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    /// A report that fits is passed through untouched, so the common case reads
2141    /// exactly as it did.
2142    #[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        // Entity without a ContextWindow: silently ignored.
2196        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    // ── force_transition (pipeline helper) edge cases via fan-out ─────────────
2219
2220    #[test]
2221    fn force_transition_applies_routing_and_handles_despawn_and_overflow() {
2222        use crate::pipeline::force_transition;
2223        // Routing present on the target stage ⇒ ToolResultRoutingComponent added.
2224        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        // Despawned entity: no panic, no effect.
2249        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        // A tiny pinned region + a huge stage system prompt ⇒ overflow on entry.
2260        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    /// Build a world with one agent carrying the given blueprint/setups/window.
2296    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}