1use 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;
20use crate::providers::ProviderRegistry;
21use crate::world::PipelineWorld;
22
23#[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
39pub enum BlueprintSource {
41 Path(PathBuf),
43 Toml(String),
45 Inline(Box<leviath_core::Blueprint>),
48}
49
50#[non_exhaustive]
54pub struct SpawnSpec {
55 pub blueprint: BlueprintSource,
57 pub task: String,
59 pub workdir: PathBuf,
61 pub model: Option<String>,
63 pub regions: HashMap<String, String>,
65 pub metadata: HashMap<String, String>,
67 pub output: Option<leviath_core::output::OutputSpec>,
76}
77
78impl SpawnSpec {
79 pub fn new(
81 blueprint: BlueprintSource,
82 task: impl Into<String>,
83 workdir: impl Into<PathBuf>,
84 ) -> Self {
85 Self {
86 blueprint,
87 task: task.into(),
88 workdir: workdir.into(),
89 model: None,
90 regions: HashMap::new(),
91 metadata: HashMap::new(),
92 output: None,
93 }
94 }
95
96 pub fn output(mut self, format: impl Into<String>, instructions: Option<String>) -> Self {
101 self.output = Some(leviath_core::output::OutputSpec {
102 format: Some(format.into()),
103 instructions,
104 example: None,
105 schema: None,
106 validator: None,
107 });
108 self
109 }
110}
111
112pub struct AgentWorldBuilder {
120 creds: Vec<ProviderCreds>,
121 custom_providers: Vec<(String, Arc<dyn leviath_providers::Provider>)>,
122 tool_service: Option<Arc<dyn ToolService>>,
123 pool_config: InferencePoolConfig,
124 tool_concurrency: usize,
125 state_dir: Option<PathBuf>,
126 defaults: ModelDefaults,
127 hints: leviath_core::config::PromptHints,
128 runtime: Option<Handle>,
129}
130
131impl AgentWorldBuilder {
132 fn new() -> Self {
133 Self {
134 creds: Vec::new(),
135 custom_providers: Vec::new(),
136 tool_service: None,
137 pool_config: InferencePoolConfig::new(),
138 tool_concurrency: 4,
139 state_dir: None,
140 defaults: ModelDefaults::default(),
141 hints: leviath_core::config::PromptHints {
144 batch_tool: false,
145 shell: false,
146 },
147 runtime: None,
148 }
149 }
150
151 pub fn provider(mut self, creds: ProviderCreds) -> Self {
154 self.creds.push(creds);
155 self
156 }
157
158 pub fn register_provider(
162 mut self,
163 name: impl Into<String>,
164 provider: Arc<dyn leviath_providers::Provider>,
165 ) -> Self {
166 self.custom_providers.push((name.into(), provider));
167 self
168 }
169
170 pub fn default_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
173 self.defaults.provider = provider.into();
174 self.defaults.model = Some(model.into());
175 self
176 }
177
178 pub fn fallback_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
185 self.defaults
186 .fallback_order
187 .push(leviath_core::blueprint::ModelEntry::new(
188 provider.into(),
189 model.into(),
190 ));
191 self
192 }
193
194 pub fn tool_service(mut self, service: Arc<dyn ToolService>) -> Self {
198 self.tool_service = Some(service);
199 self
200 }
201
202 pub fn state_dir(mut self, dir: impl Into<PathBuf>) -> Self {
206 self.state_dir = Some(dir.into());
207 self
208 }
209
210 pub fn inference_pool(mut self, config: InferencePoolConfig) -> Self {
212 self.pool_config = config;
213 self
214 }
215
216 pub fn tool_concurrency(mut self, n: usize) -> Self {
218 self.tool_concurrency = n;
219 self
220 }
221
222 pub fn prompt_hints(mut self, hints: leviath_core::config::PromptHints) -> Self {
228 self.hints = hints;
229 self
230 }
231
232 pub fn runtime(mut self, handle: Handle) -> Self {
234 self.runtime = Some(handle);
235 self
236 }
237
238 pub fn build(self) -> Result<AgentWorld, EmbedError> {
240 self.build_with(&leviath_providers::provider::build_http_client)
241 }
242
243 pub fn build_with(
249 self,
250 build_client: leviath_providers::provider::HttpClientFactory<'_>,
251 ) -> Result<AgentWorld, EmbedError> {
252 if self.creds.is_empty() && self.custom_providers.is_empty() {
253 return Err(EmbedError::NoProviders);
254 }
255 let handle = match self.runtime {
256 Some(handle) => handle,
257 None => Handle::try_current().map_err(|_| EmbedError::NoRuntime)?,
258 };
259
260 let mut registry: ProviderRegistry =
261 crate::provider_creds::build_provider_registry_with(&self.creds, build_client)
262 .map_err(|e| EmbedError::ProviderClient(e.to_string()))?;
263 for (name, provider) in self.custom_providers {
264 registry.register(name, provider);
265 }
266
267 let hub = InteractionHub::new();
268 let (service, basic_tools): (Arc<dyn ToolService>, Option<Arc<BasicToolService>>) =
269 match self.tool_service {
270 Some(service) => (service, None),
271 None => {
272 let basic = Arc::new(BasicToolService::new(hub.clone()));
273 (basic.clone(), Some(basic))
274 }
275 };
276
277 let mut world = PipelineWorld::new(
278 registry,
279 service,
280 self.pool_config,
281 self.tool_concurrency,
282 self.state_dir.map(|d| d.join("runs")),
283 handle.clone(),
284 );
285 world.insert_interaction_hub(hub.clone());
286 let mut host = WorldHost::with_interactions(world, hub.clone());
287
288 let staged: StagedBlueprints = Arc::new(Mutex::new(HashMap::new()));
289 let spawner = EmbedSpawner {
290 basic_tools: basic_tools.clone(),
291 defaults: self.defaults,
292 hints: self.hints,
293 staged: staged.clone(),
294 };
295 host.set_spawner(Box::new(move |world, args| spawner.spawn(world, args)));
296 if let Some(tools) = basic_tools {
297 host.set_reaper(Box::new(move |_world, entity| tools.unregister(entity)));
298 }
299
300 let events = host.event_sender();
301 let (control, control_rx) = tokio::sync::mpsc::unbounded_channel();
302 let serve_task = handle.spawn(async move {
303 host.serve(control_rx).await;
304 host
305 });
306
307 Ok(AgentWorld {
308 control,
309 events,
310 hub,
311 staged,
312 serve_task,
313 })
314 }
315}
316
317pub struct AgentWorld {
322 control: UnboundedSender<ControlOp>,
323 events: broadcast::Sender<WorldEvent>,
324 hub: InteractionHub,
325 staged: StagedBlueprints,
326 serve_task: tokio::task::JoinHandle<WorldHost>,
327}
328
329impl AgentWorld {
330 pub fn builder() -> AgentWorldBuilder {
332 AgentWorldBuilder::new()
333 }
334
335 async fn ask<T>(
337 &self,
338 build: impl FnOnce(oneshot::Sender<T>) -> ControlOp,
339 ) -> Result<T, EmbedError> {
340 let (reply, rx) = oneshot::channel();
341 self.control
342 .send(build(reply))
343 .map_err(|_| EmbedError::ChannelClosed)?;
344 rx.await.map_err(|_| EmbedError::ChannelClosed)
345 }
346
347 pub async fn spawn(&self, spec: SpawnSpec) -> Result<RunId, EmbedError> {
350 enum Resolved {
354 Path(PathBuf),
355 Inline(Box<leviath_core::Blueprint>),
356 }
357 let resolved = match spec.blueprint {
358 BlueprintSource::Path(path) => Resolved::Path(path),
359 BlueprintSource::Toml(toml) => Resolved::Inline(Box::new(
360 leviath_core::manifest::parse_manifest(&toml)
361 .map_err(|e| EmbedError::Blueprint(format!("parse manifest: {e}")))?,
362 )),
363 BlueprintSource::Inline(blueprint) => Resolved::Inline(blueprint),
364 };
365 if let Resolved::Inline(blueprint) = &resolved {
366 blueprint
367 .validate()
368 .map_err(|e| EmbedError::Blueprint(format!("invalid blueprint: {e}")))?;
369 }
370 let stem = match &resolved {
371 Resolved::Path(path) => path
372 .file_stem()
373 .map(|s| s.to_string_lossy().into_owned())
374 .unwrap_or_default(),
375 Resolved::Inline(blueprint) => blueprint.name.clone(),
376 };
377 let run_id = mint_run_id(&stem);
378 let blueprint_path = match resolved {
379 Resolved::Path(path) => path.to_string_lossy().into_owned(),
380 Resolved::Inline(blueprint) => {
381 self.staged
382 .lock()
383 .unwrap_or_else(PoisonError::into_inner)
384 .insert(run_id.clone(), *blueprint);
385 format!("inline:{run_id}")
386 }
387 };
388 let args = SpawnArgs {
389 run_id,
390 blueprint_path,
391 task: spec.task,
392 regions: spec.regions,
393 model: spec.model,
394 workdir: spec.workdir.to_string_lossy().into_owned(),
395 metadata: spec.metadata,
396 output: spec.output,
397 ..Default::default()
398 };
399 let run_id = self
400 .ask(|reply| ControlOp::Spawn {
401 args: Box::new(args),
402 reply,
403 })
404 .await?
405 .map_err(EmbedError::Spawn)?;
406 Ok(RunId(run_id))
407 }
408
409 pub async fn result(&self, id: &RunId) -> Option<leviath_core::output::FinalOutput> {
417 self.ask(|reply| ControlOp::Result {
418 run_id: id.0.clone(),
419 reply,
420 })
421 .await
422 .ok()
423 .flatten()
424 }
425
426 pub fn events(&self) -> EventStream {
428 EventStream::new(self.events.subscribe())
429 }
430
431 pub async fn status(&self, id: &RunId) -> Option<AgentStatus> {
433 self.ask(|reply| ControlOp::Status {
434 run_id: id.0.clone(),
435 reply,
436 })
437 .await
438 .ok()
439 .flatten()
440 }
441
442 pub async fn send_message(&self, id: &RunId, content: &str) -> bool {
445 self.ask(|reply| ControlOp::Message {
446 agent_id: id.0.clone(),
447 content: content.to_string(),
448 target_region: None,
449 reply,
450 })
451 .await
452 .unwrap_or(false)
453 }
454
455 pub async fn pause(&self, id: &RunId) -> bool {
457 self.ask(|reply| ControlOp::Pause {
458 run_id: id.0.clone(),
459 reply,
460 })
461 .await
462 .unwrap_or(false)
463 }
464
465 pub async fn resume(&self, id: &RunId) -> bool {
467 self.ask(|reply| ControlOp::Resume {
468 run_id: id.0.clone(),
469 reply,
470 })
471 .await
472 .unwrap_or(false)
473 }
474
475 pub async fn cancel(&self, id: &RunId) -> bool {
477 self.ask(|reply| ControlOp::Cancel {
478 run_id: id.0.clone(),
479 reply,
480 })
481 .await
482 .unwrap_or(false)
483 }
484
485 pub fn pending_inputs(&self) -> Vec<(RunId, leviath_core::interaction::InteractionRequest)> {
488 self.hub
489 .pending()
490 .into_iter()
491 .map(|(agent_id, request)| (RunId(agent_id), request))
492 .collect()
493 }
494
495 pub fn answer(&self, response: leviath_core::interaction::InteractionResponse) -> bool {
498 self.hub.answer(response)
499 }
500
501 pub async fn shutdown(self) {
505 let _ = self.ask(|reply| ControlOp::Shutdown { reply }).await;
506 drop(self.serve_task.await);
511 }
512}
513
514#[cfg(test)]
515mod tests {
516 use super::*;
517 use leviath_providers::{
518 FinishReason, InferenceRequest, InferenceResponse, ModelCapabilities, Provider,
519 ProviderError, TokenUsage, ToolCall,
520 };
521 use std::collections::VecDeque;
522
523 struct Mock {
525 responses: Mutex<VecDeque<InferenceResponse>>,
526 }
527
528 #[async_trait::async_trait]
529 impl Provider for Mock {
530 async fn infer(&self, _r: &InferenceRequest) -> Result<InferenceResponse, ProviderError> {
531 self.responses
532 .lock()
533 .unwrap_or_else(PoisonError::into_inner)
534 .pop_front()
535 .ok_or_else(|| ProviderError::Other("script exhausted".to_string()))
536 }
537 async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
538 1
539 }
540 fn max_context_tokens(&self, _m: &str) -> usize {
541 100_000
542 }
543 fn name(&self) -> &str {
544 "mock"
545 }
546 fn capabilities(&self, _m: &str) -> ModelCapabilities {
547 ModelCapabilities::default()
548 }
549 }
550
551 fn text(content: &str) -> InferenceResponse {
552 InferenceResponse {
553 content: content.to_string(),
554 tool_calls: vec![],
555 tokens_used: TokenUsage {
556 prompt_tokens: 1,
557 completion_tokens: 1,
558 total_tokens: 2,
559 cached_tokens: 0,
560 cache_write_tokens: 0,
561 },
562 finish_reason: FinishReason::Complete,
563 }
564 }
565
566 fn with_tool(id: &str, name: &str, args: serde_json::Value) -> InferenceResponse {
567 let mut r = text("");
568 r.tool_calls.push(ToolCall {
569 id: id.to_string(),
570 name: name.to_string(),
571 arguments: args,
572 thought_signature: None,
573 });
574 r
575 }
576
577 struct Recorder {
580 seen: Arc<Mutex<Vec<Vec<String>>>>,
581 }
582
583 #[async_trait::async_trait]
584 impl Provider for Recorder {
585 async fn infer(&self, r: &InferenceRequest) -> Result<InferenceResponse, ProviderError> {
586 self.seen
587 .lock()
588 .unwrap_or_else(PoisonError::into_inner)
589 .push(r.system.iter().map(|b| b.text.clone()).collect());
590 Ok(text("done"))
591 }
592 async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
593 1
594 }
595 fn max_context_tokens(&self, _m: &str) -> usize {
596 100_000
597 }
598 fn name(&self) -> &str {
599 "mock"
600 }
601 fn capabilities(&self, _m: &str) -> ModelCapabilities {
602 ModelCapabilities::default()
603 }
604 }
605
606 fn mock_world(responses: Vec<InferenceResponse>) -> AgentWorld {
607 AgentWorld::builder()
608 .register_provider(
609 "mock",
610 Arc::new(Mock {
611 responses: Mutex::new(responses.into_iter().collect()),
612 }),
613 )
614 .build()
615 .expect("world builds inside the test runtime")
616 }
617
618 const TWO_STAGE: &str = r#"[agent]
619name = "embedded"
620version = "0.0.0"
621description = "Two stage embedded test agent."
622entry_stage = "work"
623
624[stages.work]
625mode = "autonomous"
626model = { provider = "mock", model = "m" }
627description = "Do the work"
628available_tools = ["read_file"]
629system_prompt = "Work."
630[stages.work.transitions.wrap]
631transform = "direct"
632
633[stages.wrap]
634mode = "autonomous"
635model = { provider = "mock", model = "m" }
636description = "Wrap up"
637allow_complete = true
638system_prompt = "Wrap."
639
640[context.regions]
641conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
642"#;
643
644 const ASKER: &str = r#"[agent]
645name = "asker"
646version = "0.0.0"
647description = "Asks one question then finishes."
648entry_stage = "chat"
649
650[stages.chat]
651mode = "autonomous"
652model = { provider = "mock", model = "m" }
653description = "Chat"
654available_tools = ["ask_user_text"]
655allow_complete = true
656system_prompt = "Ask."
657
658[context.regions]
659conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
660"#;
661
662 async fn events_until(
665 stream: &mut EventStream,
666 pred: impl Fn(&WorldEvent) -> bool,
667 ) -> Vec<WorldEvent> {
668 let mut seen = Vec::new();
669 while let Some(event) = stream.next().await {
670 let done = pred(&event);
671 seen.push(event);
672 if done {
673 break;
674 }
675 }
676 seen
677 }
678
679 #[tokio::test]
680 async fn build_without_providers_is_refused() {
681 let err = AgentWorld::builder().build().map(|_| ()).unwrap_err();
682 assert_eq!(err, EmbedError::NoProviders);
683 }
684
685 #[test]
689 fn a_spawn_spec_carries_the_requested_output_shape() {
690 let dir = tempfile::tempdir().expect("temp dir");
691 let spec = SpawnSpec::new(
692 BlueprintSource::Toml("[agent]\nname = \"x\"".to_string()),
693 "do the thing",
694 dir.path(),
695 )
696 .output("a2ui", Some("One card per finding.".to_string()));
697
698 let requested = spec.output.expect("the spec carries it");
699 assert_eq!(requested.format.as_deref(), Some("a2ui"));
700 assert_eq!(
701 requested.instructions.as_deref(),
702 Some("One card per finding.")
703 );
704 assert!(requested.schema.is_none());
707 assert!(requested.validator.is_none());
708 assert!(requested.example.is_none());
709 }
710
711 #[test]
714 fn a_spawn_spec_requests_no_shape_by_default() {
715 let dir = tempfile::tempdir().expect("temp dir");
716 let spec = SpawnSpec::new(
717 BlueprintSource::Toml("[agent]\nname = \"x\"".to_string()),
718 "do the thing",
719 dir.path(),
720 );
721 assert!(spec.output.is_none());
722 }
723
724 #[test]
725 fn build_outside_a_tokio_runtime_is_refused() {
726 let err = AgentWorld::builder()
727 .register_provider(
728 "mock",
729 Arc::new(Mock {
730 responses: Mutex::new(VecDeque::new()),
731 }),
732 )
733 .build()
734 .map(|_| ())
735 .unwrap_err();
736 assert_eq!(err, EmbedError::NoRuntime);
737 }
738
739 #[test]
740 fn build_accepts_an_explicit_runtime_handle() {
741 let rt = tokio::runtime::Builder::new_multi_thread()
744 .worker_threads(1)
745 .enable_all()
746 .build()
747 .unwrap();
748 let world = AgentWorld::builder()
749 .register_provider(
750 "mock",
751 Arc::new(Mock {
752 responses: Mutex::new(VecDeque::new()),
753 }),
754 )
755 .runtime(rt.handle().clone())
756 .build()
757 .expect("explicit handle suffices");
758 rt.block_on(world.shutdown());
759 }
760
761 #[tokio::test]
762 async fn agent_runs_to_completion_with_stage_and_tool_events() {
763 let dir = tempfile::tempdir().unwrap();
764 std::fs::write(dir.path().join("notes.txt"), "the notes").unwrap();
765 let world = mock_world(vec![
769 with_tool("c1", "read_file", serde_json::json!({"path": "notes.txt"})),
770 text("moving on"),
771 text("done"),
772 text("done"),
773 text("done"),
774 text("done"),
775 ]);
776 let mut events = world.events();
777
778 let run_id = world
779 .spawn(SpawnSpec::new(
780 BlueprintSource::Toml(TWO_STAGE.to_string()),
781 "summarize the notes",
782 dir.path(),
783 ))
784 .await
785 .expect("spawns");
786 assert!(run_id.as_ref().starts_with("embedded-"));
787
788 let seen = tokio::time::timeout(
789 std::time::Duration::from_secs(20),
790 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
791 )
792 .await
793 .expect("completed before timeout");
794
795 let spawned = seen
796 .iter()
797 .any(|e| matches!(e, WorldEvent::Spawned { run_id: r, .. } if r == run_id.as_ref()));
798 assert!(spawned, "saw Spawned: {seen:?}");
799 let transitioned = seen.iter().any(|e| {
800 matches!(e, WorldEvent::StageTransition { from, to, .. }
801 if from == "work" && to == "wrap")
802 });
803 assert!(transitioned, "saw StageTransition: {seen:?}");
804 let started = seen
805 .iter()
806 .any(|e| matches!(e, WorldEvent::ToolCallStarted { tool, .. } if tool == "read_file"));
807 assert!(started, "saw ToolCallStarted: {seen:?}");
808 let finished = seen.iter().any(|e| {
809 matches!(e, WorldEvent::ToolCallFinished { tool, ok, summary, .. }
810 if tool == "read_file" && *ok && summary.contains("the notes"))
811 });
812 assert!(finished, "saw ToolCallFinished: {seen:?}");
813 let completed = seen
814 .iter()
815 .any(|e| matches!(e, WorldEvent::Completed { status, .. } if status == "complete"));
816 assert!(completed, "saw Completed: {seen:?}");
817
818 world.shutdown().await;
819 }
820
821 #[tokio::test]
822 async fn ask_user_surfaces_as_interaction_and_resumes_on_answer() {
823 let dir = tempfile::tempdir().unwrap();
824 let world = mock_world(vec![
825 with_tool(
826 "c1",
827 "ask_user_text",
828 serde_json::json!({"prompt": "Which database?"}),
829 ),
830 text("done"),
831 ]);
832 let mut events = world.events();
833 let run_id = world
834 .spawn(SpawnSpec::new(
835 BlueprintSource::Toml(ASKER.to_string()),
836 "pick a database",
837 dir.path(),
838 ))
839 .await
840 .expect("spawns");
841
842 let seen = tokio::time::timeout(
844 std::time::Duration::from_secs(20),
845 events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
846 )
847 .await
848 .expect("interaction before timeout");
849 let request = seen
850 .iter()
851 .find_map(|e| match e {
852 WorldEvent::Interaction {
853 run_id: r, request, ..
854 } if r == run_id.as_ref() => Some(request.clone()),
855 _ => None,
856 })
857 .expect("interaction event carries the request");
858 assert!(request.prompt.contains("Which database?"));
859 let pending = world.pending_inputs();
860 assert_eq!(pending.len(), 1);
861 assert_eq!(pending[0].0, run_id);
862
863 assert!(!world.pause(&run_id).await);
867 assert!(world.send_message(&run_id, "prefer something boring").await);
868
869 assert!(
871 world.answer(leviath_core::interaction::InteractionResponse::text(
872 request.id.clone(),
873 "postgres"
874 ))
875 );
876 let seen = tokio::time::timeout(
877 std::time::Duration::from_secs(20),
878 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
879 )
880 .await
881 .expect("completed before timeout");
882 assert!(
883 seen.iter()
884 .any(|e| matches!(e, WorldEvent::Completed { .. }))
885 );
886
887 world.shutdown().await;
888 }
889
890 #[tokio::test]
891 async fn spawn_reports_blueprint_and_workdir_errors() {
892 let dir = tempfile::tempdir().unwrap();
893 let world = mock_world(vec![]);
894
895 let err = world
897 .spawn(SpawnSpec::new(
898 BlueprintSource::Toml("not = [valid".to_string()),
899 "t",
900 dir.path(),
901 ))
902 .await
903 .unwrap_err();
904 assert!(err.to_string().starts_with("blueprint error"));
905
906 let err = world
908 .spawn(SpawnSpec::new(
909 BlueprintSource::Path(dir.path().join("missing.leviath")),
910 "t",
911 dir.path(),
912 ))
913 .await
914 .unwrap_err();
915 assert!(err.to_string().starts_with("spawn error"));
916
917 let err = world
919 .spawn(SpawnSpec::new(
920 BlueprintSource::Toml(TWO_STAGE.to_string()),
921 "t",
922 dir.path().join("nope"),
923 ))
924 .await
925 .unwrap_err();
926 assert!(err.to_string().starts_with("spawn error"));
927
928 world.shutdown().await;
929 }
930
931 #[tokio::test]
932 async fn spawn_from_a_manifest_file_works() {
933 let dir = tempfile::tempdir().unwrap();
934 let manifest = dir.path().join("embedded.leviath");
935 std::fs::write(&manifest, TWO_STAGE).unwrap();
936 let world = mock_world(vec![
938 text("moving on"),
939 text("moving on"),
940 text("moving on"),
941 text("moving on"),
942 text("done"),
943 text("done"),
944 text("done"),
945 text("done"),
946 ]);
947 let mut events = world.events();
948
949 let run_id = world
950 .spawn(SpawnSpec::new(
951 BlueprintSource::Path(manifest),
952 "just finish",
953 dir.path(),
954 ))
955 .await
956 .expect("spawns from the file");
957 assert!(run_id.as_ref().starts_with("embedded-"));
958
959 let seen = tokio::time::timeout(
960 std::time::Duration::from_secs(20),
961 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
962 )
963 .await
964 .expect("completed before timeout");
965 assert!(
966 seen.iter()
967 .any(|e| matches!(e, WorldEvent::Completed { .. }))
968 );
969 world.shutdown().await;
970 }
971
972 #[tokio::test]
973 async fn unknown_runs_answer_negatively() {
974 let world = mock_world(vec![]);
975 let ghost = RunId("no-such-run".to_string());
976 assert_eq!(world.status(&ghost).await, None);
977 assert_eq!(world.result(&ghost).await, None);
978 assert!(!world.pause(&ghost).await);
979 assert!(!world.resume(&ghost).await);
980 assert!(!world.cancel(&ghost).await);
981 assert!(world.pending_inputs().is_empty());
982 world.shutdown().await;
983 }
984
985 #[tokio::test]
986 async fn shutdown_ends_the_event_stream_and_further_requests_fail() {
987 let world = mock_world(vec![]);
988 let mut events = world.events();
989 let control = world.control.clone();
990 world.shutdown().await;
991 assert_eq!(
992 tokio::time::timeout(std::time::Duration::from_secs(5), events.next())
993 .await
994 .expect("stream ends"),
995 None
996 );
997 assert!(control.is_closed());
1000 let (reply, _rx) = oneshot::channel();
1001 assert!(control.send(ControlOp::List { reply }).is_err());
1002 }
1003
1004 #[tokio::test]
1005 async fn cancel_stops_a_parked_run() {
1006 let dir = tempfile::tempdir().unwrap();
1007 let world = mock_world(vec![with_tool(
1008 "c1",
1009 "ask_user_text",
1010 serde_json::json!({"prompt": "?"}),
1011 )]);
1012 let mut events = world.events();
1013 let run_id = world
1014 .spawn(SpawnSpec::new(
1015 BlueprintSource::Toml(ASKER.to_string()),
1016 "ask",
1017 dir.path(),
1018 ))
1019 .await
1020 .expect("spawns");
1021 tokio::time::timeout(
1022 std::time::Duration::from_secs(20),
1023 events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
1024 )
1025 .await
1026 .expect("parked on the question");
1027
1028 assert!(world.cancel(&run_id).await);
1029 let seen = tokio::time::timeout(
1030 std::time::Duration::from_secs(20),
1031 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1032 )
1033 .await
1034 .expect("terminal event after cancel");
1035 assert!(seen.iter().any(|e| {
1036 matches!(e, WorldEvent::Completed { status, .. } if status == "cancelled")
1037 }));
1038 world.shutdown().await;
1039 }
1040
1041 struct CannedService;
1044 impl crate::pipeline::ToolService for CannedService {
1045 fn exec_for(
1046 &self,
1047 _entity: bevy_ecs::entity::Entity,
1048 calls: Vec<leviath_providers::ToolCall>,
1049 _progress: crate::pipeline::ToolProgress,
1050 ) -> crate::tool_bridge::BoxedToolExec {
1051 Box::new(move || {
1052 Box::pin(async move {
1053 calls
1054 .into_iter()
1055 .map(|c| (c.id, "canned".to_string()))
1056 .collect()
1057 })
1058 })
1059 }
1060 }
1061
1062 #[test]
1065 fn fallback_models_accumulate_in_order_beside_the_default() {
1066 let builder = AgentWorldBuilder::new()
1067 .fallback_model("anthropic", "sonnet")
1068 .fallback_model("openai", "gpt")
1069 .default_model("openrouter", "deepseek");
1070 assert_eq!(builder.defaults.provider, "openrouter");
1071 assert_eq!(builder.defaults.model.as_deref(), Some("deepseek"));
1072 assert_eq!(
1073 builder
1074 .defaults
1075 .fallback_order
1076 .iter()
1077 .map(|e| (e.provider.as_str(), e.model.as_str()))
1078 .collect::<Vec<_>>(),
1079 vec![("anthropic", "sonnet"), ("openai", "gpt")]
1080 );
1081 }
1082
1083 #[tokio::test]
1084 async fn every_builder_option_composes_and_state_dir_persists_runs() {
1085 let dir = tempfile::tempdir().unwrap();
1086 let state = tempfile::tempdir().unwrap();
1087 let world = AgentWorld::builder()
1088 .provider(ProviderCreds::simple("ollama"))
1089 .register_provider(
1090 "mock",
1091 Arc::new(Mock {
1092 responses: Mutex::new(
1093 vec![
1094 with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
1095 text("done"),
1096 text("done"),
1097 ]
1098 .into_iter()
1099 .collect(),
1100 ),
1101 }),
1102 )
1103 .default_model("mock", "m")
1104 .fallback_model("ollama", "llama")
1107 .fallback_model("mock", "spare")
1108 .state_dir(state.path())
1109 .inference_pool(InferencePoolConfig::new())
1110 .tool_concurrency(2)
1111 .build()
1112 .expect("all options compose");
1113 let mut events = world.events();
1114 let run_id = world
1115 .spawn(SpawnSpec::new(
1116 BlueprintSource::Toml(TWO_STAGE.to_string()),
1117 "persist me",
1118 dir.path(),
1119 ))
1120 .await
1121 .expect("spawns");
1122 tokio::time::timeout(
1123 std::time::Duration::from_secs(20),
1124 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1125 )
1126 .await
1127 .expect("completes");
1128 world.shutdown().await;
1129
1130 let run_dir = state.path().join("runs").join(run_id.as_ref());
1132 assert!(run_dir.join("meta.json").exists());
1133 assert!(state.path().join("machine-id").exists());
1134 }
1135
1136 const SHELL_STAGE: &str = r#"[agent]
1139name = "shelly"
1140version = "0.0.0"
1141description = "One stage that can run commands."
1142entry_stage = "work"
1143
1144[stages.work]
1145mode = "autonomous"
1146model = { provider = "mock", model = "m" }
1147description = "Do the work"
1148available_tools = ["shell"]
1149allow_complete = true
1150system_prompt = "Work."
1151
1152[context.regions]
1153instructions = { kind = "pinned", max_tokens = 2000 }
1154conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
1155"#;
1156
1157 async fn system_blocks_with(hints: leviath_core::config::PromptHints) -> Vec<String> {
1160 let dir = tempfile::tempdir().unwrap();
1161 let seen = Arc::new(Mutex::new(Vec::new()));
1162 let world = AgentWorld::builder()
1163 .register_provider(
1164 "mock",
1165 Arc::new(Recorder {
1166 seen: Arc::clone(&seen),
1167 }),
1168 )
1169 .prompt_hints(hints)
1170 .build()
1171 .expect("world builds inside the test runtime");
1172 let mut events = world.events();
1173 world
1174 .spawn(SpawnSpec::new(
1175 BlueprintSource::Toml(SHELL_STAGE.to_string()),
1176 "go",
1177 dir.path(),
1178 ))
1179 .await
1180 .expect("spawns");
1181 tokio::time::timeout(
1182 std::time::Duration::from_secs(20),
1183 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1184 )
1185 .await
1186 .expect("completes");
1187 world.shutdown().await;
1188 let seen = seen.lock().unwrap_or_else(PoisonError::into_inner);
1189 seen.first().cloned().expect("one inference happened")
1190 }
1191
1192 #[tokio::test]
1193 async fn prompt_hints_reach_the_request_and_are_off_by_default() {
1194 let default_blocks = system_blocks_with(leviath_core::config::PromptHints {
1197 batch_tool: false,
1198 shell: false,
1199 })
1200 .await;
1201 assert!(
1202 default_blocks
1203 .iter()
1204 .all(|b| b != crate::pipeline::BATCH_TOOL_HINT),
1205 );
1206 assert!(default_blocks.iter().any(|b| b.contains("Work.")));
1207
1208 let hinted = system_blocks_with(leviath_core::config::PromptHints {
1212 batch_tool: true,
1213 shell: true,
1214 })
1215 .await;
1216 assert_eq!(
1217 hinted.first().map(String::as_str),
1218 Some(crate::pipeline::BATCH_TOOL_HINT)
1219 );
1220 let shell_hint = crate::pipeline::shell_guidance_for(std::env::consts::OS);
1222 assert_eq!(
1223 hinted.iter().any(|b| Some(b.as_str()) == shell_hint),
1224 shell_hint.is_some(),
1225 );
1226 }
1227
1228 #[tokio::test]
1229 async fn a_custom_tool_service_replaces_the_builtin_one() {
1230 let dir = tempfile::tempdir().unwrap();
1231 let world = AgentWorld::builder()
1232 .register_provider(
1233 "mock",
1234 Arc::new(Mock {
1235 responses: Mutex::new(
1236 vec![
1237 with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
1238 text("done"),
1239 text("done"),
1240 ]
1241 .into_iter()
1242 .collect(),
1243 ),
1244 }),
1245 )
1246 .tool_service(Arc::new(CannedService))
1247 .build()
1248 .expect("builds with a custom service");
1249 let mut events = world.events();
1250 world
1251 .spawn(SpawnSpec::new(
1252 BlueprintSource::Toml(TWO_STAGE.to_string()),
1253 "use the canned tools",
1254 dir.path(),
1255 ))
1256 .await
1257 .expect("spawns");
1258 let seen = tokio::time::timeout(
1259 std::time::Duration::from_secs(20),
1260 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1261 )
1262 .await
1263 .expect("completes");
1264 assert!(seen.iter().any(|e| {
1266 matches!(e, WorldEvent::ToolCallFinished { summary, .. } if summary == "canned")
1267 }));
1268 world.shutdown().await;
1269 }
1270
1271 #[tokio::test]
1272 async fn manifest_files_that_do_not_parse_or_validate_fail_the_spawn() {
1273 let dir = tempfile::tempdir().unwrap();
1274 let world = mock_world(vec![]);
1275
1276 let garbled = dir.path().join("garbled.leviath");
1277 std::fs::write(&garbled, "not = [valid").unwrap();
1278 let err = world
1279 .spawn(SpawnSpec::new(
1280 BlueprintSource::Path(garbled),
1281 "t",
1282 dir.path(),
1283 ))
1284 .await
1285 .unwrap_err();
1286 assert!(err.to_string().contains("parse manifest"));
1287
1288 let err = world
1290 .spawn(SpawnSpec::new(
1291 BlueprintSource::Path(PathBuf::from("")),
1292 "t",
1293 dir.path(),
1294 ))
1295 .await
1296 .unwrap_err();
1297 assert!(err.to_string().starts_with("spawn error"));
1298
1299 let demanding = format!(
1302 "{TWO_STAGE}\nspec = {{ kind = \"pinned\", max_tokens = 2000, seed = \"input\", required = true }}\n"
1303 );
1304 let err = world
1305 .spawn(SpawnSpec::new(
1306 BlueprintSource::Toml(demanding),
1307 "t",
1308 dir.path(),
1309 ))
1310 .await
1311 .unwrap_err();
1312 assert!(err.to_string().contains("required region"));
1313
1314 world.shutdown().await;
1315 }
1316
1317 #[tokio::test]
1318 async fn requests_after_the_world_closes_fail_closed() {
1319 let world = mock_world(vec![]);
1320 let (reply, _rx) = oneshot::channel();
1323 world
1324 .control
1325 .send(ControlOp::Shutdown { reply })
1326 .expect("world is up");
1327 while !world.control.is_closed() {
1329 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1330 }
1331 let ghost = RunId("ghost".to_string());
1332 assert_eq!(world.status(&ghost).await, None);
1333 assert!(!world.pause(&ghost).await);
1334 assert!(!world.send_message(&ghost, "hello").await);
1335 let err = world
1336 .spawn(SpawnSpec::new(
1337 BlueprintSource::Toml(TWO_STAGE.to_string()),
1338 "t",
1339 std::env::temp_dir(),
1340 ))
1341 .await
1342 .unwrap_err();
1343 assert_eq!(err, EmbedError::ChannelClosed);
1344 }
1345
1346 #[tokio::test]
1347 async fn inline_blueprints_spawn_and_invalid_ones_are_refused() {
1348 let dir = tempfile::tempdir().unwrap();
1349 let world = mock_world(vec![text("done"), text("done"), text("done"), text("done")]);
1350 let mut events = world.events();
1351
1352 let mut invalid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
1355 invalid.entry_stage = Some("ghost".to_string());
1356 let err = world
1357 .spawn(SpawnSpec::new(
1358 BlueprintSource::Inline(Box::new(invalid)),
1359 "t",
1360 dir.path(),
1361 ))
1362 .await
1363 .unwrap_err();
1364 assert!(err.to_string().contains("invalid blueprint"));
1365
1366 let mut valid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
1368 valid.stages.truncate(1);
1369 valid.stages[0].transitions = None;
1370 valid.entry_stage = Some(valid.stages[0].name.clone());
1371 let run_id = world
1372 .spawn(SpawnSpec::new(
1373 BlueprintSource::Inline(Box::new(valid)),
1374 "just answer",
1375 dir.path(),
1376 ))
1377 .await
1378 .expect("inline blueprint spawns");
1379 assert!(run_id.as_ref().starts_with("embedded-"));
1380 let seen = tokio::time::timeout(
1381 std::time::Duration::from_secs(20),
1382 events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
1383 )
1384 .await
1385 .expect("completes");
1386 assert!(
1387 seen.iter()
1388 .any(|e| matches!(e, WorldEvent::Completed { .. }))
1389 );
1390 world.shutdown().await;
1391 }
1392
1393 #[tokio::test]
1394 async fn shutdown_survives_an_aborted_serve_loop() {
1395 let world = mock_world(vec![]);
1396 world.serve_task.abort();
1399 world.shutdown().await;
1400 }
1401
1402 #[tokio::test]
1403 async fn the_mock_provider_is_a_minimal_stub() {
1404 let mock = Mock {
1407 responses: Mutex::new(VecDeque::new()),
1408 };
1409 assert_eq!(mock.count_tokens("x", "m").await, 1);
1410 assert_eq!(mock.max_context_tokens("m"), 100_000);
1411 assert_eq!(mock.name(), "mock");
1412 let _ = mock.capabilities("m");
1413
1414 let recorder = Recorder {
1417 seen: Arc::new(Mutex::new(Vec::new())),
1418 };
1419 assert_eq!(recorder.count_tokens("x", "m").await, 1);
1420 assert_eq!(recorder.max_context_tokens("m"), 100_000);
1421 assert_eq!(recorder.name(), "mock");
1422 let _ = recorder.capabilities("m");
1423 assert!(
1424 mock.infer(
1425 &serde_json::from_value(serde_json::json!({
1426 "messages": [],
1427 "model": "m",
1428 "max_tokens": 1,
1429 "temperature": 0.0,
1430 "tools": [],
1431 "extra": null,
1432 }))
1433 .unwrap()
1434 )
1435 .await
1436 .is_err()
1437 );
1438 }
1439
1440 #[tokio::test]
1441 async fn pause_and_resume_round_trip_on_an_active_run() {
1442 let dir = tempfile::tempdir().unwrap();
1445 let mut pool = InferencePoolConfig::new();
1446 pool.set_limit("m", 0);
1447 let world = AgentWorld::builder()
1448 .register_provider(
1449 "mock",
1450 Arc::new(Mock {
1451 responses: Mutex::new(VecDeque::new()),
1452 }),
1453 )
1454 .inference_pool(pool)
1455 .build()
1456 .expect("builds");
1457 let run_id = world
1458 .spawn(SpawnSpec::new(
1459 BlueprintSource::Toml(ASKER.to_string()),
1460 "wait around",
1461 dir.path(),
1462 ))
1463 .await
1464 .expect("spawns");
1465
1466 assert_eq!(world.status(&run_id).await, Some(AgentStatus::Active));
1468 assert!(world.pause(&run_id).await);
1469 assert!(world.resume(&run_id).await);
1470 assert!(world.cancel(&run_id).await);
1471 world.shutdown().await;
1472 }
1473
1474 #[test]
1475 fn run_id_displays_as_its_string() {
1476 let id = RunId("coder-1-2".to_string());
1477 assert_eq!(id.to_string(), "coder-1-2");
1478 assert_eq!(id.as_ref(), "coder-1-2");
1479 }
1480
1481 #[tokio::test]
1482 async fn a_world_whose_provider_client_will_not_build_reports_it() {
1483 let failing = |_t: Option<u64>| Err(leviath_providers::provider::malformed_url_error());
1487 let mut cred = ProviderCreds::simple("anthropic");
1488 cred.api_key = Some("k".to_string());
1489 let err = AgentWorld::builder()
1490 .provider(cred)
1491 .build_with(&failing)
1492 .err()
1493 .expect("a failing client factory should fail the build");
1494 assert_eq!(
1499 std::mem::discriminant(&err),
1500 std::mem::discriminant(&EmbedError::ProviderClient(String::new()))
1501 );
1502 assert!(err.to_string().contains("provider HTTPS client error"));
1503 }
1504}