Skip to main content

leviath_runtime/embed/
world.rs

1//! [`AgentWorld`]: the embedder-facing runtime. Build one with plain values,
2//! spawn agents, watch the event stream, answer their questions, shut down.
3
4use std::collections::HashMap;
5use std::path::PathBuf;
6use std::sync::{Arc, Mutex, PoisonError};
7
8use tokio::runtime::Handle;
9use tokio::sync::mpsc::UnboundedSender;
10use tokio::sync::{broadcast, oneshot};
11
12use super::spawner::{EmbedSpawner, StagedBlueprints, mint_run_id};
13use super::{BasicToolService, EmbedError, EventStream};
14use crate::components::AgentStatus;
15use crate::host::{ControlOp, SpawnArgs, WorldEvent, WorldHost};
16use crate::inference_pool::InferencePoolConfig;
17use crate::interaction_hub::InteractionHub;
18use crate::pipeline::{ModelDefaults, ToolService};
19use crate::provider_creds::{ProviderCreds, build_provider_registry};
20use crate::providers::ProviderRegistry;
21use crate::world::PipelineWorld;
22
23/// An opaque run identifier, minted by [`AgentWorld::spawn`].
24#[derive(Debug, Clone, PartialEq, Eq, Hash)]
25pub struct RunId(String);
26
27impl std::fmt::Display for RunId {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        f.write_str(&self.0)
30    }
31}
32
33impl AsRef<str> for RunId {
34    fn as_ref(&self) -> &str {
35        &self.0
36    }
37}
38
39/// Where a spawn's blueprint comes from.
40pub enum BlueprintSource {
41    /// A `.leviath` manifest file on disk.
42    Path(PathBuf),
43    /// Manifest TOML held in memory (parsed and validated at spawn).
44    Toml(String),
45    /// An already-constructed blueprint value (boxed: a Blueprint is a
46    /// large value, and boxing keeps the enum small).
47    Inline(Box<leviath_core::Blueprint>),
48}
49
50/// One spawn request. Build with [`SpawnSpec::new`], then set the optional
51/// fields directly; the struct is non-exhaustive so new options stay
52/// additive.
53#[non_exhaustive]
54pub struct SpawnSpec {
55    /// The agent's blueprint.
56    pub blueprint: BlueprintSource,
57    /// The task prompt, seeded into the blueprint's `task` region.
58    pub task: String,
59    /// Working directory tools are confined to. Must exist.
60    pub workdir: PathBuf,
61    /// Optional model override (`provider/model` or a bare model name).
62    pub model: Option<String>,
63    /// Seed content for named caller-input regions.
64    pub regions: HashMap<String, String>,
65    /// Custom key/value metadata carried in the run's metadata.
66    pub metadata: HashMap<String, String>,
67}
68
69impl SpawnSpec {
70    /// A spec with the required fields; optional fields start empty.
71    pub fn new(
72        blueprint: BlueprintSource,
73        task: impl Into<String>,
74        workdir: impl Into<PathBuf>,
75    ) -> Self {
76        Self {
77            blueprint,
78            task: task.into(),
79            workdir: workdir.into(),
80            model: None,
81            regions: HashMap::new(),
82            metadata: HashMap::new(),
83        }
84    }
85}
86
87/// Builds an [`AgentWorld`] from plain values - no config file, no daemon.
88///
89/// ```ignore
90/// let world = AgentWorld::builder()
91///     .provider(ProviderCreds::anthropic(api_key))
92///     .build()?;
93/// ```
94pub struct AgentWorldBuilder {
95    creds: Vec<ProviderCreds>,
96    custom_providers: Vec<(String, Arc<dyn leviath_providers::Provider>)>,
97    tool_service: Option<Arc<dyn ToolService>>,
98    pool_config: InferencePoolConfig,
99    tool_concurrency: usize,
100    state_dir: Option<PathBuf>,
101    defaults: ModelDefaults,
102    hints: leviath_core::config::PromptHints,
103    runtime: Option<Handle>,
104}
105
106impl AgentWorldBuilder {
107    fn new() -> Self {
108        Self {
109            creds: Vec::new(),
110            custom_providers: Vec::new(),
111            tool_service: None,
112            pool_config: InferencePoolConfig::new(),
113            tool_concurrency: 4,
114            state_dir: None,
115            defaults: ModelDefaults::default(),
116            // Off unless asked for: an embedder owns its prompts, and the
117            // daemon's `config.toml` defaults do not reach this path.
118            hints: leviath_core::config::PromptHints {
119                batch_tool: false,
120                shell: false,
121            },
122            runtime: None,
123        }
124    }
125
126    /// Add a provider from credentials (repeatable). See [`ProviderCreds`]
127    /// for the supported providers.
128    pub fn provider(mut self, creds: ProviderCreds) -> Self {
129        self.creds.push(creds);
130        self
131    }
132
133    /// Register a custom [`Provider`](leviath_providers::Provider)
134    /// implementation under `name` (repeatable). Wins over a credentials
135    /// entry with the same name.
136    pub fn register_provider(
137        mut self,
138        name: impl Into<String>,
139        provider: Arc<dyn leviath_providers::Provider>,
140    ) -> Self {
141        self.custom_providers.push((name.into(), provider));
142        self
143    }
144
145    /// The user-default provider/model, the fallback when none of a stage's
146    /// listed models has a registered provider.
147    pub fn default_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
148        self.defaults.provider = provider.into();
149        self.defaults.model = Some(model.into());
150        self
151    }
152
153    /// Append a host-wide failover target, tried after a stage's own entries
154    /// and the default model when the provider in use stops answering.
155    ///
156    /// Call it once per target, best first. This is what keeps a blueprint
157    /// that names exactly one model running when that provider runs out of
158    /// credits (issue #201).
159    pub fn fallback_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
160        self.defaults
161            .fallback_order
162            .push(leviath_core::blueprint::ModelEntry::new(
163                provider.into(),
164                model.into(),
165            ));
166        self
167    }
168
169    /// Replace the default [`BasicToolService`] with a custom tool service.
170    /// The embed spawner then skips per-agent tool registration; the custom
171    /// service sees agents through its own `exec_for`.
172    pub fn tool_service(mut self, service: Arc<dyn ToolService>) -> Self {
173        self.tool_service = Some(service);
174        self
175    }
176
177    /// Persist run state on disk under `dir`, in the daemon's layout
178    /// (`<dir>/runs/<run_id>/`, machine id at `<dir>/machine-id`). Without
179    /// this the world runs entirely in memory and never touches disk.
180    pub fn state_dir(mut self, dir: impl Into<PathBuf>) -> Self {
181        self.state_dir = Some(dir.into());
182        self
183    }
184
185    /// Per-model inference concurrency limits.
186    pub fn inference_pool(mut self, config: InferencePoolConfig) -> Self {
187        self.pool_config = config;
188        self
189    }
190
191    /// How many tool batches may execute concurrently (default 4).
192    pub fn tool_concurrency(mut self, n: usize) -> Self {
193        self.tool_concurrency = n;
194        self
195    }
196
197    /// Opt in to the framework-authored system-prompt hints, off by default on
198    /// this path. `shell` is worth turning on for a blueprint that grants the
199    /// shell tool and may run on Windows: it tells the model that commands go
200    /// through `cmd.exe` rather than a POSIX shell. A blueprint's `[agent]` or
201    /// `[stages.<name>]` `batch_tool_hint` / `shell_hint` still overrides this.
202    pub fn prompt_hints(mut self, hints: leviath_core::config::PromptHints) -> Self {
203        self.hints = hints;
204        self
205    }
206
207    /// Run the world on `handle` instead of the ambient Tokio runtime.
208    pub fn runtime(mut self, handle: Handle) -> Self {
209        self.runtime = Some(handle);
210        self
211    }
212
213    /// Assemble the world and start its serve loop on the Tokio runtime.
214    pub fn build(self) -> Result<AgentWorld, EmbedError> {
215        if self.creds.is_empty() && self.custom_providers.is_empty() {
216            return Err(EmbedError::NoProviders);
217        }
218        let handle = match self.runtime {
219            Some(handle) => handle,
220            None => Handle::try_current().map_err(|_| EmbedError::NoRuntime)?,
221        };
222
223        let mut registry: ProviderRegistry = build_provider_registry(&self.creds);
224        for (name, provider) in self.custom_providers {
225            registry.register(name, provider);
226        }
227
228        let hub = InteractionHub::new();
229        let (service, basic_tools): (Arc<dyn ToolService>, Option<Arc<BasicToolService>>) =
230            match self.tool_service {
231                Some(service) => (service, None),
232                None => {
233                    let basic = Arc::new(BasicToolService::new(hub.clone()));
234                    (basic.clone(), Some(basic))
235                }
236            };
237
238        let mut world = PipelineWorld::new(
239            registry,
240            service,
241            self.pool_config,
242            self.tool_concurrency,
243            self.state_dir.map(|d| d.join("runs")),
244            handle.clone(),
245        );
246        world.insert_interaction_hub(hub.clone());
247        let mut host = WorldHost::with_interactions(world, hub.clone());
248
249        let staged: StagedBlueprints = Arc::new(Mutex::new(HashMap::new()));
250        let spawner = EmbedSpawner {
251            basic_tools: basic_tools.clone(),
252            defaults: self.defaults,
253            hints: self.hints,
254            staged: staged.clone(),
255        };
256        host.set_spawner(Box::new(move |world, args| spawner.spawn(world, args)));
257        if let Some(tools) = basic_tools {
258            host.set_reaper(Box::new(move |_world, entity| tools.unregister(entity)));
259        }
260
261        let events = host.event_sender();
262        let (control, control_rx) = tokio::sync::mpsc::unbounded_channel();
263        let serve_task = handle.spawn(async move {
264            host.serve(control_rx).await;
265            host
266        });
267
268        Ok(AgentWorld {
269            control,
270            events,
271            hub,
272            staged,
273            serve_task,
274        })
275    }
276}
277
278/// A running embedded world: agents spawn into it, events stream out of it.
279///
280/// Internally this is the same [`WorldHost`] the daemon serves - addressed
281/// in-process over a channel instead of over the control socket.
282pub struct AgentWorld {
283    control: UnboundedSender<ControlOp>,
284    events: broadcast::Sender<WorldEvent>,
285    hub: InteractionHub,
286    staged: StagedBlueprints,
287    serve_task: tokio::task::JoinHandle<WorldHost>,
288}
289
290impl AgentWorld {
291    /// Start building a world.
292    pub fn builder() -> AgentWorldBuilder {
293        AgentWorldBuilder::new()
294    }
295
296    /// Send one control op and await its reply.
297    async fn ask<T>(
298        &self,
299        build: impl FnOnce(oneshot::Sender<T>) -> ControlOp,
300    ) -> Result<T, EmbedError> {
301        let (reply, rx) = oneshot::channel();
302        self.control
303            .send(build(reply))
304            .map_err(|_| EmbedError::ChannelClosed)?;
305        rx.await.map_err(|_| EmbedError::ChannelClosed)
306    }
307
308    /// Spawn an agent. Returns its [`RunId`] once the agent is live in the
309    /// world (blueprint loaded, stages resolved, seeds applied).
310    pub async fn spawn(&self, spec: SpawnSpec) -> Result<RunId, EmbedError> {
311        // Resolve the blueprint source: a path passes through to the spawner;
312        // in-memory blueprints validate here and park in the staged map under
313        // the freshly minted run id.
314        enum Resolved {
315            Path(PathBuf),
316            Inline(Box<leviath_core::Blueprint>),
317        }
318        let resolved = match spec.blueprint {
319            BlueprintSource::Path(path) => Resolved::Path(path),
320            BlueprintSource::Toml(toml) => Resolved::Inline(Box::new(
321                leviath_core::manifest::parse_manifest(&toml)
322                    .map_err(|e| EmbedError::Blueprint(format!("parse manifest: {e}")))?,
323            )),
324            BlueprintSource::Inline(blueprint) => Resolved::Inline(blueprint),
325        };
326        if let Resolved::Inline(blueprint) = &resolved {
327            blueprint
328                .validate()
329                .map_err(|e| EmbedError::Blueprint(format!("invalid blueprint: {e}")))?;
330        }
331        let stem = match &resolved {
332            Resolved::Path(path) => path
333                .file_stem()
334                .map(|s| s.to_string_lossy().into_owned())
335                .unwrap_or_default(),
336            Resolved::Inline(blueprint) => blueprint.name.clone(),
337        };
338        let run_id = mint_run_id(&stem);
339        let blueprint_path = match resolved {
340            Resolved::Path(path) => path.to_string_lossy().into_owned(),
341            Resolved::Inline(blueprint) => {
342                self.staged
343                    .lock()
344                    .unwrap_or_else(PoisonError::into_inner)
345                    .insert(run_id.clone(), *blueprint);
346                format!("inline:{run_id}")
347            }
348        };
349        let args = SpawnArgs {
350            run_id,
351            blueprint_path,
352            task: spec.task,
353            regions: spec.regions,
354            model: spec.model,
355            workdir: spec.workdir.to_string_lossy().into_owned(),
356            metadata: spec.metadata,
357            ..Default::default()
358        };
359        let run_id = self
360            .ask(|reply| ControlOp::Spawn {
361                args: Box::new(args),
362                reply,
363            })
364            .await?
365            .map_err(EmbedError::Spawn)?;
366        Ok(RunId(run_id))
367    }
368
369    /// Subscribe to the world's events, from this moment on.
370    pub fn events(&self) -> EventStream {
371        EventStream::new(self.events.subscribe())
372    }
373
374    /// A run's current status, or `None` if the world doesn't know it.
375    pub async fn status(&self, id: &RunId) -> Option<AgentStatus> {
376        self.ask(|reply| ControlOp::Status {
377            run_id: id.0.clone(),
378            reply,
379        })
380        .await
381        .ok()
382        .flatten()
383    }
384
385    /// Deliver a message into a running agent's inbox. `false` when the
386    /// world can no longer accept messages (shut down or shutting down).
387    pub async fn send_message(&self, id: &RunId, content: &str) -> bool {
388        self.ask(|reply| ControlOp::Message {
389            agent_id: id.0.clone(),
390            content: content.to_string(),
391            target_region: None,
392            reply,
393        })
394        .await
395        .unwrap_or(false)
396    }
397
398    /// Pause a run. `false` if there is no such live run.
399    pub async fn pause(&self, id: &RunId) -> bool {
400        self.ask(|reply| ControlOp::Pause {
401            run_id: id.0.clone(),
402            reply,
403        })
404        .await
405        .unwrap_or(false)
406    }
407
408    /// Resume a paused run. `false` if there is no such live run.
409    pub async fn resume(&self, id: &RunId) -> bool {
410        self.ask(|reply| ControlOp::Resume {
411            run_id: id.0.clone(),
412            reply,
413        })
414        .await
415        .unwrap_or(false)
416    }
417
418    /// Cancel a run. `false` if there is no such live run.
419    pub async fn cancel(&self, id: &RunId) -> bool {
420        self.ask(|reply| ControlOp::Cancel {
421            run_id: id.0.clone(),
422            reply,
423        })
424        .await
425        .unwrap_or(false)
426    }
427
428    /// Every open question agents are waiting on, as `(run, request)`. Each
429    /// also arrived as an [`Interaction`](WorldEvent::Interaction) event.
430    pub fn pending_inputs(&self) -> Vec<(RunId, leviath_core::interaction::InteractionRequest)> {
431        self.hub
432            .pending()
433            .into_iter()
434            .map(|(agent_id, request)| (RunId(agent_id), request))
435            .collect()
436    }
437
438    /// Answer an open question (matched by the response's `request_id`).
439    /// `false` if no such request is open.
440    pub fn answer(&self, response: leviath_core::interaction::InteractionResponse) -> bool {
441        self.hub.answer(response)
442    }
443
444    /// Shut the world down and wait for it to finish. The serve loop drains
445    /// every queued persistence write before it returns (its own
446    /// flush-and-stop), so once this resolves nothing is left in flight.
447    pub async fn shutdown(self) {
448        let _ = self.ask(|reply| ControlOp::Shutdown { reply }).await;
449        // Joining is enough: `WorldHost::serve` flushes on its way out, and a
450        // second flush would tick a world whose persistence resource is
451        // already gone. `Err` here means the task was aborted; there is
452        // nothing left to wait for either way.
453        drop(self.serve_task.await);
454    }
455}
456
457#[cfg(test)]
458mod tests {
459    use super::*;
460    use leviath_providers::{
461        FinishReason, InferenceRequest, InferenceResponse, ModelCapabilities, Provider,
462        ProviderError, TokenUsage, ToolCall,
463    };
464    use std::collections::VecDeque;
465
466    /// A scripted provider: pops one canned response per inference call.
467    struct Mock {
468        responses: Mutex<VecDeque<InferenceResponse>>,
469    }
470
471    #[async_trait::async_trait]
472    impl Provider for Mock {
473        async fn infer(&self, _r: InferenceRequest) -> Result<InferenceResponse, ProviderError> {
474            self.responses
475                .lock()
476                .unwrap_or_else(PoisonError::into_inner)
477                .pop_front()
478                .ok_or_else(|| ProviderError::Other("script exhausted".to_string()))
479        }
480        async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
481            1
482        }
483        fn max_context_tokens(&self, _m: &str) -> usize {
484            100_000
485        }
486        fn name(&self) -> &str {
487            "mock"
488        }
489        fn capabilities(&self, _m: &str) -> ModelCapabilities {
490            ModelCapabilities::default()
491        }
492    }
493
494    fn text(content: &str) -> InferenceResponse {
495        InferenceResponse {
496            content: content.to_string(),
497            tool_calls: vec![],
498            tokens_used: TokenUsage {
499                prompt_tokens: 1,
500                completion_tokens: 1,
501                total_tokens: 2,
502                cached_tokens: 0,
503                cache_write_tokens: 0,
504            },
505            finish_reason: FinishReason::Complete,
506        }
507    }
508
509    fn with_tool(id: &str, name: &str, args: serde_json::Value) -> InferenceResponse {
510        let mut r = text("");
511        r.tool_calls.push(ToolCall {
512            id: id.to_string(),
513            name: name.to_string(),
514            arguments: args,
515            thought_signature: None,
516        });
517        r
518    }
519
520    /// A provider that records the system blocks of every request it is handed,
521    /// then answers "done". For asserting what the framework prepends.
522    struct Recorder {
523        seen: Arc<Mutex<Vec<Vec<String>>>>,
524    }
525
526    #[async_trait::async_trait]
527    impl Provider for Recorder {
528        async fn infer(&self, r: InferenceRequest) -> Result<InferenceResponse, ProviderError> {
529            self.seen
530                .lock()
531                .unwrap_or_else(PoisonError::into_inner)
532                .push(r.system.iter().map(|b| b.text.clone()).collect());
533            Ok(text("done"))
534        }
535        async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
536            1
537        }
538        fn max_context_tokens(&self, _m: &str) -> usize {
539            100_000
540        }
541        fn name(&self) -> &str {
542            "mock"
543        }
544        fn capabilities(&self, _m: &str) -> ModelCapabilities {
545            ModelCapabilities::default()
546        }
547    }
548
549    fn mock_world(responses: Vec<InferenceResponse>) -> AgentWorld {
550        AgentWorld::builder()
551            .register_provider(
552                "mock",
553                Arc::new(Mock {
554                    responses: Mutex::new(responses.into_iter().collect()),
555                }),
556            )
557            .build()
558            .expect("world builds inside the test runtime")
559    }
560
561    const TWO_STAGE: &str = r#"[agent]
562name = "embedded"
563version = "0.0.0"
564description = "Two stage embedded test agent."
565entry_stage = "work"
566
567[stages.work]
568mode = "autonomous"
569model = { provider = "mock", model = "m" }
570description = "Do the work"
571available_tools = ["read_file"]
572system_prompt = "Work."
573[stages.work.transitions.wrap]
574transform = "direct"
575
576[stages.wrap]
577mode = "autonomous"
578model = { provider = "mock", model = "m" }
579description = "Wrap up"
580allow_complete = true
581system_prompt = "Wrap."
582
583[context.regions]
584conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
585"#;
586
587    const ASKER: &str = r#"[agent]
588name = "asker"
589version = "0.0.0"
590description = "Asks one question then finishes."
591entry_stage = "chat"
592
593[stages.chat]
594mode = "autonomous"
595model = { provider = "mock", model = "m" }
596description = "Chat"
597available_tools = ["ask_user_text"]
598allow_complete = true
599system_prompt = "Ask."
600
601[context.regions]
602conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
603"#;
604
605    /// Drain events until `pred` matches (or the stream ends), collecting
606    /// everything seen. Bounded by the caller's `tokio::time::timeout`.
607    async fn events_until(
608        stream: &mut EventStream,
609        pred: impl Fn(&WorldEvent) -> bool,
610    ) -> Vec<WorldEvent> {
611        let mut seen = Vec::new();
612        while let Some(event) = stream.next().await {
613            let done = pred(&event);
614            seen.push(event);
615            if done {
616                break;
617            }
618        }
619        seen
620    }
621
622    #[tokio::test]
623    async fn build_without_providers_is_refused() {
624        let err = AgentWorld::builder().build().map(|_| ()).unwrap_err();
625        assert_eq!(err, EmbedError::NoProviders);
626    }
627
628    #[test]
629    fn build_outside_a_tokio_runtime_is_refused() {
630        let err = AgentWorld::builder()
631            .register_provider(
632                "mock",
633                Arc::new(Mock {
634                    responses: Mutex::new(VecDeque::new()),
635                }),
636            )
637            .build()
638            .map(|_| ())
639            .unwrap_err();
640        assert_eq!(err, EmbedError::NoRuntime);
641    }
642
643    #[test]
644    fn build_accepts_an_explicit_runtime_handle() {
645        // A plain test (no ambient runtime): the handle passed via
646        // `.runtime()` is what makes build succeed.
647        let rt = tokio::runtime::Builder::new_multi_thread()
648            .worker_threads(1)
649            .enable_all()
650            .build()
651            .unwrap();
652        let world = AgentWorld::builder()
653            .register_provider(
654                "mock",
655                Arc::new(Mock {
656                    responses: Mutex::new(VecDeque::new()),
657                }),
658            )
659            .runtime(rt.handle().clone())
660            .build()
661            .expect("explicit handle suffices");
662        rt.block_on(world.shutdown());
663    }
664
665    #[tokio::test]
666    async fn agent_runs_to_completion_with_stage_and_tool_events() {
667        let dir = tempfile::tempdir().unwrap();
668        std::fs::write(dir.path().join("notes.txt"), "the notes").unwrap();
669        // The wrap stage makes no tool calls, so its text-only responses get
670        // the "use your tools" nudge up to the cap before the last is
671        // accepted; script enough of them.
672        let world = mock_world(vec![
673            with_tool("c1", "read_file", serde_json::json!({"path": "notes.txt"})),
674            text("moving on"),
675            text("done"),
676            text("done"),
677            text("done"),
678            text("done"),
679        ]);
680        let mut events = world.events();
681
682        let run_id = world
683            .spawn(SpawnSpec::new(
684                BlueprintSource::Toml(TWO_STAGE.to_string()),
685                "summarize the notes",
686                dir.path(),
687            ))
688            .await
689            .expect("spawns");
690        assert!(run_id.as_ref().starts_with("embedded-"));
691
692        let seen = tokio::time::timeout(
693            std::time::Duration::from_secs(20),
694            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
695        )
696        .await
697        .expect("completed before timeout");
698
699        let spawned = seen
700            .iter()
701            .any(|e| matches!(e, WorldEvent::Spawned { run_id: r, .. } if r == run_id.as_ref()));
702        assert!(spawned, "saw Spawned: {seen:?}");
703        let transitioned = seen.iter().any(|e| {
704            matches!(e, WorldEvent::StageTransition { from, to, .. }
705                if from == "work" && to == "wrap")
706        });
707        assert!(transitioned, "saw StageTransition: {seen:?}");
708        let started = seen
709            .iter()
710            .any(|e| matches!(e, WorldEvent::ToolCallStarted { tool, .. } if tool == "read_file"));
711        assert!(started, "saw ToolCallStarted: {seen:?}");
712        let finished = seen.iter().any(|e| {
713            matches!(e, WorldEvent::ToolCallFinished { tool, ok, summary, .. }
714                if tool == "read_file" && *ok && summary.contains("the notes"))
715        });
716        assert!(finished, "saw ToolCallFinished: {seen:?}");
717        let completed = seen
718            .iter()
719            .any(|e| matches!(e, WorldEvent::Completed { status, .. } if status == "complete"));
720        assert!(completed, "saw Completed: {seen:?}");
721
722        world.shutdown().await;
723    }
724
725    #[tokio::test]
726    async fn ask_user_surfaces_as_interaction_and_resumes_on_answer() {
727        let dir = tempfile::tempdir().unwrap();
728        let world = mock_world(vec![
729            with_tool(
730                "c1",
731                "ask_user_text",
732                serde_json::json!({"prompt": "Which database?"}),
733            ),
734            text("done"),
735        ]);
736        let mut events = world.events();
737        let run_id = world
738            .spawn(SpawnSpec::new(
739                BlueprintSource::Toml(ASKER.to_string()),
740                "pick a database",
741                dir.path(),
742            ))
743            .await
744            .expect("spawns");
745
746        // The question arrives on the event stream and in pending_inputs.
747        let seen = tokio::time::timeout(
748            std::time::Duration::from_secs(20),
749            events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
750        )
751        .await
752        .expect("interaction before timeout");
753        let request = seen
754            .iter()
755            .find_map(|e| match e {
756                WorldEvent::Interaction {
757                    run_id: r, request, ..
758                } if r == run_id.as_ref() => Some(request.clone()),
759                _ => None,
760            })
761            .expect("interaction event carries the request");
762        assert!(request.prompt.contains("Which database?"));
763        let pending = world.pending_inputs();
764        assert_eq!(pending.len(), 1);
765        assert_eq!(pending[0].0, run_id);
766
767        // A live, parked agent still accepts messages. (Pause is refused
768        // while the agent waits on input - see the capacity test below for
769        // the pause/resume round-trip.)
770        assert!(!world.pause(&run_id).await);
771        assert!(world.send_message(&run_id, "prefer something boring").await);
772
773        // Answering resumes the run to completion.
774        assert!(
775            world.answer(leviath_core::interaction::InteractionResponse::text(
776                request.id.clone(),
777                "postgres"
778            ))
779        );
780        let seen = tokio::time::timeout(
781            std::time::Duration::from_secs(20),
782            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
783        )
784        .await
785        .expect("completed before timeout");
786        assert!(
787            seen.iter()
788                .any(|e| matches!(e, WorldEvent::Completed { .. }))
789        );
790
791        world.shutdown().await;
792    }
793
794    #[tokio::test]
795    async fn spawn_reports_blueprint_and_workdir_errors() {
796        let dir = tempfile::tempdir().unwrap();
797        let world = mock_world(vec![]);
798
799        // Unparseable TOML.
800        let err = world
801            .spawn(SpawnSpec::new(
802                BlueprintSource::Toml("not = [valid".to_string()),
803                "t",
804                dir.path(),
805            ))
806            .await
807            .unwrap_err();
808        assert!(err.to_string().starts_with("blueprint error"));
809
810        // A manifest path that does not exist fails in the spawner.
811        let err = world
812            .spawn(SpawnSpec::new(
813                BlueprintSource::Path(dir.path().join("missing.leviath")),
814                "t",
815                dir.path(),
816            ))
817            .await
818            .unwrap_err();
819        assert!(err.to_string().starts_with("spawn error"));
820
821        // A workdir that does not exist is refused before anything spawns.
822        let err = world
823            .spawn(SpawnSpec::new(
824                BlueprintSource::Toml(TWO_STAGE.to_string()),
825                "t",
826                dir.path().join("nope"),
827            ))
828            .await
829            .unwrap_err();
830        assert!(err.to_string().starts_with("spawn error"));
831
832        world.shutdown().await;
833    }
834
835    #[tokio::test]
836    async fn spawn_from_a_manifest_file_works() {
837        let dir = tempfile::tempdir().unwrap();
838        let manifest = dir.path().join("embedded.leviath");
839        std::fs::write(&manifest, TWO_STAGE).unwrap();
840        // Both stages are text-only here, so each needs its nudge budget.
841        let world = mock_world(vec![
842            text("moving on"),
843            text("moving on"),
844            text("moving on"),
845            text("moving on"),
846            text("done"),
847            text("done"),
848            text("done"),
849            text("done"),
850        ]);
851        let mut events = world.events();
852
853        let run_id = world
854            .spawn(SpawnSpec::new(
855                BlueprintSource::Path(manifest),
856                "just finish",
857                dir.path(),
858            ))
859            .await
860            .expect("spawns from the file");
861        assert!(run_id.as_ref().starts_with("embedded-"));
862
863        let seen = tokio::time::timeout(
864            std::time::Duration::from_secs(20),
865            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
866        )
867        .await
868        .expect("completed before timeout");
869        assert!(
870            seen.iter()
871                .any(|e| matches!(e, WorldEvent::Completed { .. }))
872        );
873        world.shutdown().await;
874    }
875
876    #[tokio::test]
877    async fn unknown_runs_answer_negatively() {
878        let world = mock_world(vec![]);
879        let ghost = RunId("no-such-run".to_string());
880        assert_eq!(world.status(&ghost).await, None);
881        assert!(!world.pause(&ghost).await);
882        assert!(!world.resume(&ghost).await);
883        assert!(!world.cancel(&ghost).await);
884        assert!(world.pending_inputs().is_empty());
885        world.shutdown().await;
886    }
887
888    #[tokio::test]
889    async fn shutdown_ends_the_event_stream_and_further_requests_fail() {
890        let world = mock_world(vec![]);
891        let mut events = world.events();
892        let control = world.control.clone();
893        world.shutdown().await;
894        assert_eq!(
895            tokio::time::timeout(std::time::Duration::from_secs(5), events.next())
896                .await
897                .expect("stream ends"),
898            None
899        );
900        // The serve loop is gone (shutdown consumed the world after joining
901        // it), so its receiver is dropped and a late op cannot be delivered.
902        assert!(control.is_closed());
903        let (reply, _rx) = oneshot::channel();
904        assert!(control.send(ControlOp::List { reply }).is_err());
905    }
906
907    #[tokio::test]
908    async fn cancel_stops_a_parked_run() {
909        let dir = tempfile::tempdir().unwrap();
910        let world = mock_world(vec![with_tool(
911            "c1",
912            "ask_user_text",
913            serde_json::json!({"prompt": "?"}),
914        )]);
915        let mut events = world.events();
916        let run_id = world
917            .spawn(SpawnSpec::new(
918                BlueprintSource::Toml(ASKER.to_string()),
919                "ask",
920                dir.path(),
921            ))
922            .await
923            .expect("spawns");
924        tokio::time::timeout(
925            std::time::Duration::from_secs(20),
926            events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
927        )
928        .await
929        .expect("parked on the question");
930
931        assert!(world.cancel(&run_id).await);
932        let seen = tokio::time::timeout(
933            std::time::Duration::from_secs(20),
934            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
935        )
936        .await
937        .expect("terminal event after cancel");
938        assert!(seen.iter().any(|e| {
939            matches!(e, WorldEvent::Completed { status, .. } if status == "cancelled")
940        }));
941        world.shutdown().await;
942    }
943
944    /// A tool service that answers every call with a canned string; used to
945    /// exercise the custom-service seam (no per-agent registration).
946    struct CannedService;
947    impl crate::pipeline::ToolService for CannedService {
948        fn exec_for(
949            &self,
950            _entity: bevy_ecs::entity::Entity,
951            calls: Vec<leviath_providers::ToolCall>,
952            _progress: crate::pipeline::ToolProgress,
953        ) -> crate::tool_bridge::BoxedToolExec {
954            Box::new(move || {
955                Box::pin(async move {
956                    calls
957                        .into_iter()
958                        .map(|c| (c.id, "canned".to_string()))
959                        .collect()
960                })
961            })
962        }
963    }
964
965    /// The failover chain is ordered and additive, and setting a default model
966    /// afterwards must not wipe it (issue #201).
967    #[test]
968    fn fallback_models_accumulate_in_order_beside_the_default() {
969        let builder = AgentWorldBuilder::new()
970            .fallback_model("anthropic", "sonnet")
971            .fallback_model("openai", "gpt")
972            .default_model("openrouter", "deepseek");
973        assert_eq!(builder.defaults.provider, "openrouter");
974        assert_eq!(builder.defaults.model.as_deref(), Some("deepseek"));
975        assert_eq!(
976            builder
977                .defaults
978                .fallback_order
979                .iter()
980                .map(|e| (e.provider.as_str(), e.model.as_str()))
981                .collect::<Vec<_>>(),
982            vec![("anthropic", "sonnet"), ("openai", "gpt")]
983        );
984    }
985
986    #[tokio::test]
987    async fn every_builder_option_composes_and_state_dir_persists_runs() {
988        let dir = tempfile::tempdir().unwrap();
989        let state = tempfile::tempdir().unwrap();
990        let world = AgentWorld::builder()
991            .provider(ProviderCreds::simple("ollama"))
992            .register_provider(
993                "mock",
994                Arc::new(Mock {
995                    responses: Mutex::new(
996                        vec![
997                            with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
998                            text("done"),
999                            text("done"),
1000                        ]
1001                        .into_iter()
1002                        .collect(),
1003                    ),
1004                }),
1005            )
1006            .default_model("mock", "m")
1007            // Repeated on purpose: the chain is ordered, so it must accumulate
1008            // rather than replace, and it must not disturb the default model.
1009            .fallback_model("ollama", "llama")
1010            .fallback_model("mock", "spare")
1011            .state_dir(state.path())
1012            .inference_pool(InferencePoolConfig::new())
1013            .tool_concurrency(2)
1014            .build()
1015            .expect("all options compose");
1016        let mut events = world.events();
1017        let run_id = world
1018            .spawn(SpawnSpec::new(
1019                BlueprintSource::Toml(TWO_STAGE.to_string()),
1020                "persist me",
1021                dir.path(),
1022            ))
1023            .await
1024            .expect("spawns");
1025        tokio::time::timeout(
1026            std::time::Duration::from_secs(20),
1027            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1028        )
1029        .await
1030        .expect("completes");
1031        world.shutdown().await;
1032
1033        // The daemon's on-disk layout appeared under the state dir.
1034        let run_dir = state.path().join("runs").join(run_id.as_ref());
1035        assert!(run_dir.join("meta.json").exists());
1036        assert!(state.path().join("machine-id").exists());
1037    }
1038
1039    /// The single-stage blueprint the hint tests drive, with a shell tool so the
1040    /// shell hint's tool guard is satisfied.
1041    const SHELL_STAGE: &str = r#"[agent]
1042name = "shelly"
1043version = "0.0.0"
1044description = "One stage that can run commands."
1045entry_stage = "work"
1046
1047[stages.work]
1048mode = "autonomous"
1049model = { provider = "mock", model = "m" }
1050description = "Do the work"
1051available_tools = ["shell"]
1052allow_complete = true
1053system_prompt = "Work."
1054
1055[context.regions]
1056instructions = { kind = "pinned", max_tokens = 2000 }
1057conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
1058"#;
1059
1060    /// Run `SHELL_STAGE` once against a [`Recorder`] and hand back the system
1061    /// blocks of the first request, with `hints` as the world's global toggles.
1062    async fn system_blocks_with(hints: leviath_core::config::PromptHints) -> Vec<String> {
1063        let dir = tempfile::tempdir().unwrap();
1064        let seen = Arc::new(Mutex::new(Vec::new()));
1065        let world = AgentWorld::builder()
1066            .register_provider(
1067                "mock",
1068                Arc::new(Recorder {
1069                    seen: Arc::clone(&seen),
1070                }),
1071            )
1072            .prompt_hints(hints)
1073            .build()
1074            .expect("world builds inside the test runtime");
1075        let mut events = world.events();
1076        world
1077            .spawn(SpawnSpec::new(
1078                BlueprintSource::Toml(SHELL_STAGE.to_string()),
1079                "go",
1080                dir.path(),
1081            ))
1082            .await
1083            .expect("spawns");
1084        tokio::time::timeout(
1085            std::time::Duration::from_secs(20),
1086            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1087        )
1088        .await
1089        .expect("completes");
1090        world.shutdown().await;
1091        let seen = seen.lock().unwrap_or_else(PoisonError::into_inner);
1092        seen.first().cloned().expect("one inference happened")
1093    }
1094
1095    #[tokio::test]
1096    async fn prompt_hints_reach_the_request_and_are_off_by_default() {
1097        // Off unless asked for: an embedder that never calls `prompt_hints`
1098        // gets exactly the blueprint's own prompt.
1099        let default_blocks = system_blocks_with(leviath_core::config::PromptHints {
1100            batch_tool: false,
1101            shell: false,
1102        })
1103        .await;
1104        assert!(
1105            default_blocks
1106                .iter()
1107                .all(|b| b != crate::pipeline::BATCH_TOOL_HINT),
1108        );
1109        assert!(default_blocks.iter().any(|b| b.contains("Work.")));
1110
1111        // Turned on, the hint leads the prefix. The shell hint rides the same
1112        // path but only says anything on Windows, so this asserts on the batch
1113        // hint, which is the platform-independent half of the plumbing.
1114        let hinted = system_blocks_with(leviath_core::config::PromptHints {
1115            batch_tool: true,
1116            shell: true,
1117        })
1118        .await;
1119        assert_eq!(
1120            hinted.first().map(String::as_str),
1121            Some(crate::pipeline::BATCH_TOOL_HINT)
1122        );
1123        // Whatever the host OS says about its shell is what the run carries.
1124        let shell_hint = crate::pipeline::shell_guidance_for(std::env::consts::OS);
1125        assert_eq!(
1126            hinted.iter().any(|b| Some(b.as_str()) == shell_hint),
1127            shell_hint.is_some(),
1128        );
1129    }
1130
1131    #[tokio::test]
1132    async fn a_custom_tool_service_replaces_the_builtin_one() {
1133        let dir = tempfile::tempdir().unwrap();
1134        let world = AgentWorld::builder()
1135            .register_provider(
1136                "mock",
1137                Arc::new(Mock {
1138                    responses: Mutex::new(
1139                        vec![
1140                            with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
1141                            text("done"),
1142                            text("done"),
1143                        ]
1144                        .into_iter()
1145                        .collect(),
1146                    ),
1147                }),
1148            )
1149            .tool_service(Arc::new(CannedService))
1150            .build()
1151            .expect("builds with a custom service");
1152        let mut events = world.events();
1153        world
1154            .spawn(SpawnSpec::new(
1155                BlueprintSource::Toml(TWO_STAGE.to_string()),
1156                "use the canned tools",
1157                dir.path(),
1158            ))
1159            .await
1160            .expect("spawns");
1161        let seen = tokio::time::timeout(
1162            std::time::Duration::from_secs(20),
1163            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1164        )
1165        .await
1166        .expect("completes");
1167        // The canned result (not a real file read) came back through the lane.
1168        assert!(seen.iter().any(|e| {
1169            matches!(e, WorldEvent::ToolCallFinished { summary, .. } if summary == "canned")
1170        }));
1171        world.shutdown().await;
1172    }
1173
1174    #[tokio::test]
1175    async fn manifest_files_that_do_not_parse_or_validate_fail_the_spawn() {
1176        let dir = tempfile::tempdir().unwrap();
1177        let world = mock_world(vec![]);
1178
1179        let garbled = dir.path().join("garbled.leviath");
1180        std::fs::write(&garbled, "not = [valid").unwrap();
1181        let err = world
1182            .spawn(SpawnSpec::new(
1183                BlueprintSource::Path(garbled),
1184                "t",
1185                dir.path(),
1186            ))
1187            .await
1188            .unwrap_err();
1189        assert!(err.to_string().contains("parse manifest"));
1190
1191        // A path with no file stem still spawns an attempt (and fails to read).
1192        let err = world
1193            .spawn(SpawnSpec::new(
1194                BlueprintSource::Path(PathBuf::from("")),
1195                "t",
1196                dir.path(),
1197            ))
1198            .await
1199            .unwrap_err();
1200        assert!(err.to_string().starts_with("spawn error"));
1201
1202        // A required caller-input region that was not provided fails before
1203        // any inference.
1204        let demanding = format!(
1205            "{TWO_STAGE}\nspec = {{ kind = \"pinned\", max_tokens = 2000, seed = \"input\", required = true }}\n"
1206        );
1207        let err = world
1208            .spawn(SpawnSpec::new(
1209                BlueprintSource::Toml(demanding),
1210                "t",
1211                dir.path(),
1212            ))
1213            .await
1214            .unwrap_err();
1215        assert!(err.to_string().contains("required region"));
1216
1217        world.shutdown().await;
1218    }
1219
1220    #[tokio::test]
1221    async fn requests_after_the_world_closes_fail_closed() {
1222        let world = mock_world(vec![]);
1223        // Stop the serve loop out from under the handle (without consuming
1224        // the AgentWorld, as shutdown() would).
1225        let (reply, _rx) = oneshot::channel();
1226        world
1227            .control
1228            .send(ControlOp::Shutdown { reply })
1229            .expect("world is up");
1230        // Wait until the serve loop is really gone (its rx dropped).
1231        while !world.control.is_closed() {
1232            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1233        }
1234        let ghost = RunId("ghost".to_string());
1235        assert_eq!(world.status(&ghost).await, None);
1236        assert!(!world.pause(&ghost).await);
1237        assert!(!world.send_message(&ghost, "hello").await);
1238        let err = world
1239            .spawn(SpawnSpec::new(
1240                BlueprintSource::Toml(TWO_STAGE.to_string()),
1241                "t",
1242                std::env::temp_dir(),
1243            ))
1244            .await
1245            .unwrap_err();
1246        assert_eq!(err, EmbedError::ChannelClosed);
1247    }
1248
1249    #[tokio::test]
1250    async fn inline_blueprints_spawn_and_invalid_ones_are_refused() {
1251        let dir = tempfile::tempdir().unwrap();
1252        let world = mock_world(vec![text("done"), text("done"), text("done"), text("done")]);
1253        let mut events = world.events();
1254
1255        // An invalid inline blueprint (entry stage names nothing) is refused
1256        // before it reaches the world.
1257        let mut invalid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
1258        invalid.entry_stage = Some("ghost".to_string());
1259        let err = world
1260            .spawn(SpawnSpec::new(
1261                BlueprintSource::Inline(Box::new(invalid)),
1262                "t",
1263                dir.path(),
1264            ))
1265            .await
1266            .unwrap_err();
1267        assert!(err.to_string().contains("invalid blueprint"));
1268
1269        // A valid one runs. Trim it to a single text-only stage.
1270        let mut valid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
1271        valid.stages.truncate(1);
1272        valid.stages[0].transitions = None;
1273        valid.entry_stage = Some(valid.stages[0].name.clone());
1274        let run_id = world
1275            .spawn(SpawnSpec::new(
1276                BlueprintSource::Inline(Box::new(valid)),
1277                "just answer",
1278                dir.path(),
1279            ))
1280            .await
1281            .expect("inline blueprint spawns");
1282        assert!(run_id.as_ref().starts_with("embedded-"));
1283        let seen = tokio::time::timeout(
1284            std::time::Duration::from_secs(20),
1285            events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1286        )
1287        .await
1288        .expect("completes");
1289        assert!(
1290            seen.iter()
1291                .any(|e| matches!(e, WorldEvent::Completed { .. }))
1292        );
1293        world.shutdown().await;
1294    }
1295
1296    #[tokio::test]
1297    async fn shutdown_survives_an_aborted_serve_loop() {
1298        let world = mock_world(vec![]);
1299        // Kill the serve task out from under the world: shutdown must not
1300        // hang or panic when the join fails.
1301        world.serve_task.abort();
1302        world.shutdown().await;
1303    }
1304
1305    #[tokio::test]
1306    async fn the_mock_provider_is_a_minimal_stub() {
1307        // Pins the fixture's inert answers so its impl stays measured (the
1308        // pipeline only calls infer when exact token counting is off).
1309        let mock = Mock {
1310            responses: Mutex::new(VecDeque::new()),
1311        };
1312        assert_eq!(mock.count_tokens("x", "m").await, 1);
1313        assert_eq!(mock.max_context_tokens("m"), 100_000);
1314        assert_eq!(mock.name(), "mock");
1315        let _ = mock.capabilities("m");
1316
1317        // Same for the recording fixture, which answers identically and is
1318        // registered under the same name.
1319        let recorder = Recorder {
1320            seen: Arc::new(Mutex::new(Vec::new())),
1321        };
1322        assert_eq!(recorder.count_tokens("x", "m").await, 1);
1323        assert_eq!(recorder.max_context_tokens("m"), 100_000);
1324        assert_eq!(recorder.name(), "mock");
1325        let _ = recorder.capabilities("m");
1326        assert!(
1327            mock.infer(
1328                serde_json::from_value(serde_json::json!({
1329                    "messages": [],
1330                    "model": "m",
1331                    "max_tokens": 1,
1332                    "temperature": 0.0,
1333                    "tools": [],
1334                    "extra": null,
1335                }))
1336                .unwrap()
1337            )
1338            .await
1339            .is_err()
1340        );
1341    }
1342
1343    #[tokio::test]
1344    async fn pause_and_resume_round_trip_on_an_active_run() {
1345        // Zero inference permits for the model: the agent stays Active,
1346        // parked on the pool, which is exactly when pause applies.
1347        let dir = tempfile::tempdir().unwrap();
1348        let mut pool = InferencePoolConfig::new();
1349        pool.set_limit("m", 0);
1350        let world = AgentWorld::builder()
1351            .register_provider(
1352                "mock",
1353                Arc::new(Mock {
1354                    responses: Mutex::new(VecDeque::new()),
1355                }),
1356            )
1357            .inference_pool(pool)
1358            .build()
1359            .expect("builds");
1360        let run_id = world
1361            .spawn(SpawnSpec::new(
1362                BlueprintSource::Toml(ASKER.to_string()),
1363                "wait around",
1364                dir.path(),
1365            ))
1366            .await
1367            .expect("spawns");
1368
1369        // Agents spawn Active, and with no permits nothing can change that.
1370        assert_eq!(world.status(&run_id).await, Some(AgentStatus::Active));
1371        assert!(world.pause(&run_id).await);
1372        assert!(world.resume(&run_id).await);
1373        assert!(world.cancel(&run_id).await);
1374        world.shutdown().await;
1375    }
1376
1377    #[test]
1378    fn run_id_displays_as_its_string() {
1379        let id = RunId("coder-1-2".to_string());
1380        assert_eq!(id.to_string(), "coder-1-2");
1381        assert_eq!(id.as_ref(), "coder-1-2");
1382    }
1383}