use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, PoisonError};
use tokio::runtime::Handle;
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::{broadcast, oneshot};
use super::spawner::{EmbedSpawner, StagedBlueprints, mint_run_id};
use super::{BasicToolService, EmbedError, EventStream};
use crate::components::AgentStatus;
use crate::host::{ControlOp, SpawnArgs, WorldEvent, WorldHost};
use crate::inference_pool::InferencePoolConfig;
use crate::interaction_hub::InteractionHub;
use crate::pipeline::{ModelDefaults, ToolService};
use crate::provider_creds::{ProviderCreds, build_provider_registry};
use crate::providers::ProviderRegistry;
use crate::world::PipelineWorld;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RunId(String);
impl std::fmt::Display for RunId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl AsRef<str> for RunId {
fn as_ref(&self) -> &str {
&self.0
}
}
pub enum BlueprintSource {
Path(PathBuf),
Toml(String),
Inline(Box<leviath_core::Blueprint>),
}
#[non_exhaustive]
pub struct SpawnSpec {
pub blueprint: BlueprintSource,
pub task: String,
pub workdir: PathBuf,
pub model: Option<String>,
pub regions: HashMap<String, String>,
pub metadata: HashMap<String, String>,
}
impl SpawnSpec {
pub fn new(
blueprint: BlueprintSource,
task: impl Into<String>,
workdir: impl Into<PathBuf>,
) -> Self {
Self {
blueprint,
task: task.into(),
workdir: workdir.into(),
model: None,
regions: HashMap::new(),
metadata: HashMap::new(),
}
}
}
pub struct AgentWorldBuilder {
creds: Vec<ProviderCreds>,
custom_providers: Vec<(String, Arc<dyn leviath_providers::Provider>)>,
tool_service: Option<Arc<dyn ToolService>>,
pool_config: InferencePoolConfig,
tool_concurrency: usize,
state_dir: Option<PathBuf>,
defaults: ModelDefaults,
hints: leviath_core::config::PromptHints,
runtime: Option<Handle>,
}
impl AgentWorldBuilder {
fn new() -> Self {
Self {
creds: Vec::new(),
custom_providers: Vec::new(),
tool_service: None,
pool_config: InferencePoolConfig::new(),
tool_concurrency: 4,
state_dir: None,
defaults: ModelDefaults::default(),
hints: leviath_core::config::PromptHints {
batch_tool: false,
shell: false,
},
runtime: None,
}
}
pub fn provider(mut self, creds: ProviderCreds) -> Self {
self.creds.push(creds);
self
}
pub fn register_provider(
mut self,
name: impl Into<String>,
provider: Arc<dyn leviath_providers::Provider>,
) -> Self {
self.custom_providers.push((name.into(), provider));
self
}
pub fn default_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
self.defaults.provider = provider.into();
self.defaults.model = Some(model.into());
self
}
pub fn fallback_model(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
self.defaults
.fallback_order
.push(leviath_core::blueprint::ModelEntry::new(
provider.into(),
model.into(),
));
self
}
pub fn tool_service(mut self, service: Arc<dyn ToolService>) -> Self {
self.tool_service = Some(service);
self
}
pub fn state_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.state_dir = Some(dir.into());
self
}
pub fn inference_pool(mut self, config: InferencePoolConfig) -> Self {
self.pool_config = config;
self
}
pub fn tool_concurrency(mut self, n: usize) -> Self {
self.tool_concurrency = n;
self
}
pub fn prompt_hints(mut self, hints: leviath_core::config::PromptHints) -> Self {
self.hints = hints;
self
}
pub fn runtime(mut self, handle: Handle) -> Self {
self.runtime = Some(handle);
self
}
pub fn build(self) -> Result<AgentWorld, EmbedError> {
if self.creds.is_empty() && self.custom_providers.is_empty() {
return Err(EmbedError::NoProviders);
}
let handle = match self.runtime {
Some(handle) => handle,
None => Handle::try_current().map_err(|_| EmbedError::NoRuntime)?,
};
let mut registry: ProviderRegistry = build_provider_registry(&self.creds);
for (name, provider) in self.custom_providers {
registry.register(name, provider);
}
let hub = InteractionHub::new();
let (service, basic_tools): (Arc<dyn ToolService>, Option<Arc<BasicToolService>>) =
match self.tool_service {
Some(service) => (service, None),
None => {
let basic = Arc::new(BasicToolService::new(hub.clone()));
(basic.clone(), Some(basic))
}
};
let mut world = PipelineWorld::new(
registry,
service,
self.pool_config,
self.tool_concurrency,
self.state_dir.map(|d| d.join("runs")),
handle.clone(),
);
world.insert_interaction_hub(hub.clone());
let mut host = WorldHost::with_interactions(world, hub.clone());
let staged: StagedBlueprints = Arc::new(Mutex::new(HashMap::new()));
let spawner = EmbedSpawner {
basic_tools: basic_tools.clone(),
defaults: self.defaults,
hints: self.hints,
staged: staged.clone(),
};
host.set_spawner(Box::new(move |world, args| spawner.spawn(world, args)));
if let Some(tools) = basic_tools {
host.set_reaper(Box::new(move |_world, entity| tools.unregister(entity)));
}
let events = host.event_sender();
let (control, control_rx) = tokio::sync::mpsc::unbounded_channel();
let serve_task = handle.spawn(async move {
host.serve(control_rx).await;
host
});
Ok(AgentWorld {
control,
events,
hub,
staged,
serve_task,
})
}
}
pub struct AgentWorld {
control: UnboundedSender<ControlOp>,
events: broadcast::Sender<WorldEvent>,
hub: InteractionHub,
staged: StagedBlueprints,
serve_task: tokio::task::JoinHandle<WorldHost>,
}
impl AgentWorld {
pub fn builder() -> AgentWorldBuilder {
AgentWorldBuilder::new()
}
async fn ask<T>(
&self,
build: impl FnOnce(oneshot::Sender<T>) -> ControlOp,
) -> Result<T, EmbedError> {
let (reply, rx) = oneshot::channel();
self.control
.send(build(reply))
.map_err(|_| EmbedError::ChannelClosed)?;
rx.await.map_err(|_| EmbedError::ChannelClosed)
}
pub async fn spawn(&self, spec: SpawnSpec) -> Result<RunId, EmbedError> {
enum Resolved {
Path(PathBuf),
Inline(Box<leviath_core::Blueprint>),
}
let resolved = match spec.blueprint {
BlueprintSource::Path(path) => Resolved::Path(path),
BlueprintSource::Toml(toml) => Resolved::Inline(Box::new(
leviath_core::manifest::parse_manifest(&toml)
.map_err(|e| EmbedError::Blueprint(format!("parse manifest: {e}")))?,
)),
BlueprintSource::Inline(blueprint) => Resolved::Inline(blueprint),
};
if let Resolved::Inline(blueprint) = &resolved {
blueprint
.validate()
.map_err(|e| EmbedError::Blueprint(format!("invalid blueprint: {e}")))?;
}
let stem = match &resolved {
Resolved::Path(path) => path
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_default(),
Resolved::Inline(blueprint) => blueprint.name.clone(),
};
let run_id = mint_run_id(&stem);
let blueprint_path = match resolved {
Resolved::Path(path) => path.to_string_lossy().into_owned(),
Resolved::Inline(blueprint) => {
self.staged
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(run_id.clone(), *blueprint);
format!("inline:{run_id}")
}
};
let args = SpawnArgs {
run_id,
blueprint_path,
task: spec.task,
regions: spec.regions,
model: spec.model,
workdir: spec.workdir.to_string_lossy().into_owned(),
metadata: spec.metadata,
..Default::default()
};
let run_id = self
.ask(|reply| ControlOp::Spawn {
args: Box::new(args),
reply,
})
.await?
.map_err(EmbedError::Spawn)?;
Ok(RunId(run_id))
}
pub fn events(&self) -> EventStream {
EventStream::new(self.events.subscribe())
}
pub async fn status(&self, id: &RunId) -> Option<AgentStatus> {
self.ask(|reply| ControlOp::Status {
run_id: id.0.clone(),
reply,
})
.await
.ok()
.flatten()
}
pub async fn send_message(&self, id: &RunId, content: &str) -> bool {
self.ask(|reply| ControlOp::Message {
agent_id: id.0.clone(),
content: content.to_string(),
target_region: None,
reply,
})
.await
.unwrap_or(false)
}
pub async fn pause(&self, id: &RunId) -> bool {
self.ask(|reply| ControlOp::Pause {
run_id: id.0.clone(),
reply,
})
.await
.unwrap_or(false)
}
pub async fn resume(&self, id: &RunId) -> bool {
self.ask(|reply| ControlOp::Resume {
run_id: id.0.clone(),
reply,
})
.await
.unwrap_or(false)
}
pub async fn cancel(&self, id: &RunId) -> bool {
self.ask(|reply| ControlOp::Cancel {
run_id: id.0.clone(),
reply,
})
.await
.unwrap_or(false)
}
pub fn pending_inputs(&self) -> Vec<(RunId, leviath_core::interaction::InteractionRequest)> {
self.hub
.pending()
.into_iter()
.map(|(agent_id, request)| (RunId(agent_id), request))
.collect()
}
pub fn answer(&self, response: leviath_core::interaction::InteractionResponse) -> bool {
self.hub.answer(response)
}
pub async fn shutdown(self) {
let _ = self.ask(|reply| ControlOp::Shutdown { reply }).await;
drop(self.serve_task.await);
}
}
#[cfg(test)]
mod tests {
use super::*;
use leviath_providers::{
FinishReason, InferenceRequest, InferenceResponse, ModelCapabilities, Provider,
ProviderError, TokenUsage, ToolCall,
};
use std::collections::VecDeque;
struct Mock {
responses: Mutex<VecDeque<InferenceResponse>>,
}
#[async_trait::async_trait]
impl Provider for Mock {
async fn infer(&self, _r: InferenceRequest) -> Result<InferenceResponse, ProviderError> {
self.responses
.lock()
.unwrap_or_else(PoisonError::into_inner)
.pop_front()
.ok_or_else(|| ProviderError::Other("script exhausted".to_string()))
}
async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
1
}
fn max_context_tokens(&self, _m: &str) -> usize {
100_000
}
fn name(&self) -> &str {
"mock"
}
fn capabilities(&self, _m: &str) -> ModelCapabilities {
ModelCapabilities::default()
}
}
fn text(content: &str) -> InferenceResponse {
InferenceResponse {
content: content.to_string(),
tool_calls: vec![],
tokens_used: TokenUsage {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
cached_tokens: 0,
cache_write_tokens: 0,
},
finish_reason: FinishReason::Complete,
}
}
fn with_tool(id: &str, name: &str, args: serde_json::Value) -> InferenceResponse {
let mut r = text("");
r.tool_calls.push(ToolCall {
id: id.to_string(),
name: name.to_string(),
arguments: args,
thought_signature: None,
});
r
}
struct Recorder {
seen: Arc<Mutex<Vec<Vec<String>>>>,
}
#[async_trait::async_trait]
impl Provider for Recorder {
async fn infer(&self, r: InferenceRequest) -> Result<InferenceResponse, ProviderError> {
self.seen
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(r.system.iter().map(|b| b.text.clone()).collect());
Ok(text("done"))
}
async fn count_tokens(&self, _t: &str, _m: &str) -> usize {
1
}
fn max_context_tokens(&self, _m: &str) -> usize {
100_000
}
fn name(&self) -> &str {
"mock"
}
fn capabilities(&self, _m: &str) -> ModelCapabilities {
ModelCapabilities::default()
}
}
fn mock_world(responses: Vec<InferenceResponse>) -> AgentWorld {
AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(responses.into_iter().collect()),
}),
)
.build()
.expect("world builds inside the test runtime")
}
const TWO_STAGE: &str = r#"[agent]
name = "embedded"
version = "0.0.0"
description = "Two stage embedded test agent."
entry_stage = "work"
[stages.work]
mode = "autonomous"
model = { provider = "mock", model = "m" }
description = "Do the work"
available_tools = ["read_file"]
system_prompt = "Work."
[stages.work.transitions.wrap]
transform = "direct"
[stages.wrap]
mode = "autonomous"
model = { provider = "mock", model = "m" }
description = "Wrap up"
allow_complete = true
system_prompt = "Wrap."
[context.regions]
conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
"#;
const ASKER: &str = r#"[agent]
name = "asker"
version = "0.0.0"
description = "Asks one question then finishes."
entry_stage = "chat"
[stages.chat]
mode = "autonomous"
model = { provider = "mock", model = "m" }
description = "Chat"
available_tools = ["ask_user_text"]
allow_complete = true
system_prompt = "Ask."
[context.regions]
conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
"#;
async fn events_until(
stream: &mut EventStream,
pred: impl Fn(&WorldEvent) -> bool,
) -> Vec<WorldEvent> {
let mut seen = Vec::new();
while let Some(event) = stream.next().await {
let done = pred(&event);
seen.push(event);
if done {
break;
}
}
seen
}
#[tokio::test]
async fn build_without_providers_is_refused() {
let err = AgentWorld::builder().build().map(|_| ()).unwrap_err();
assert_eq!(err, EmbedError::NoProviders);
}
#[test]
fn build_outside_a_tokio_runtime_is_refused() {
let err = AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(VecDeque::new()),
}),
)
.build()
.map(|_| ())
.unwrap_err();
assert_eq!(err, EmbedError::NoRuntime);
}
#[test]
fn build_accepts_an_explicit_runtime_handle() {
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let world = AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(VecDeque::new()),
}),
)
.runtime(rt.handle().clone())
.build()
.expect("explicit handle suffices");
rt.block_on(world.shutdown());
}
#[tokio::test]
async fn agent_runs_to_completion_with_stage_and_tool_events() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("notes.txt"), "the notes").unwrap();
let world = mock_world(vec![
with_tool("c1", "read_file", serde_json::json!({"path": "notes.txt"})),
text("moving on"),
text("done"),
text("done"),
text("done"),
text("done"),
]);
let mut events = world.events();
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(TWO_STAGE.to_string()),
"summarize the notes",
dir.path(),
))
.await
.expect("spawns");
assert!(run_id.as_ref().starts_with("embedded-"));
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completed before timeout");
let spawned = seen
.iter()
.any(|e| matches!(e, WorldEvent::Spawned { run_id: r, .. } if r == run_id.as_ref()));
assert!(spawned, "saw Spawned: {seen:?}");
let transitioned = seen.iter().any(|e| {
matches!(e, WorldEvent::StageTransition { from, to, .. }
if from == "work" && to == "wrap")
});
assert!(transitioned, "saw StageTransition: {seen:?}");
let started = seen
.iter()
.any(|e| matches!(e, WorldEvent::ToolCallStarted { tool, .. } if tool == "read_file"));
assert!(started, "saw ToolCallStarted: {seen:?}");
let finished = seen.iter().any(|e| {
matches!(e, WorldEvent::ToolCallFinished { tool, ok, summary, .. }
if tool == "read_file" && *ok && summary.contains("the notes"))
});
assert!(finished, "saw ToolCallFinished: {seen:?}");
let completed = seen
.iter()
.any(|e| matches!(e, WorldEvent::Completed { status, .. } if status == "complete"));
assert!(completed, "saw Completed: {seen:?}");
world.shutdown().await;
}
#[tokio::test]
async fn ask_user_surfaces_as_interaction_and_resumes_on_answer() {
let dir = tempfile::tempdir().unwrap();
let world = mock_world(vec![
with_tool(
"c1",
"ask_user_text",
serde_json::json!({"prompt": "Which database?"}),
),
text("done"),
]);
let mut events = world.events();
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(ASKER.to_string()),
"pick a database",
dir.path(),
))
.await
.expect("spawns");
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
)
.await
.expect("interaction before timeout");
let request = seen
.iter()
.find_map(|e| match e {
WorldEvent::Interaction {
run_id: r, request, ..
} if r == run_id.as_ref() => Some(request.clone()),
_ => None,
})
.expect("interaction event carries the request");
assert!(request.prompt.contains("Which database?"));
let pending = world.pending_inputs();
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].0, run_id);
assert!(!world.pause(&run_id).await);
assert!(world.send_message(&run_id, "prefer something boring").await);
assert!(
world.answer(leviath_core::interaction::InteractionResponse::text(
request.id.clone(),
"postgres"
))
);
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completed before timeout");
assert!(
seen.iter()
.any(|e| matches!(e, WorldEvent::Completed { .. }))
);
world.shutdown().await;
}
#[tokio::test]
async fn spawn_reports_blueprint_and_workdir_errors() {
let dir = tempfile::tempdir().unwrap();
let world = mock_world(vec![]);
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml("not = [valid".to_string()),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().starts_with("blueprint error"));
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Path(dir.path().join("missing.leviath")),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().starts_with("spawn error"));
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(TWO_STAGE.to_string()),
"t",
dir.path().join("nope"),
))
.await
.unwrap_err();
assert!(err.to_string().starts_with("spawn error"));
world.shutdown().await;
}
#[tokio::test]
async fn spawn_from_a_manifest_file_works() {
let dir = tempfile::tempdir().unwrap();
let manifest = dir.path().join("embedded.leviath");
std::fs::write(&manifest, TWO_STAGE).unwrap();
let world = mock_world(vec![
text("moving on"),
text("moving on"),
text("moving on"),
text("moving on"),
text("done"),
text("done"),
text("done"),
text("done"),
]);
let mut events = world.events();
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Path(manifest),
"just finish",
dir.path(),
))
.await
.expect("spawns from the file");
assert!(run_id.as_ref().starts_with("embedded-"));
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completed before timeout");
assert!(
seen.iter()
.any(|e| matches!(e, WorldEvent::Completed { .. }))
);
world.shutdown().await;
}
#[tokio::test]
async fn unknown_runs_answer_negatively() {
let world = mock_world(vec![]);
let ghost = RunId("no-such-run".to_string());
assert_eq!(world.status(&ghost).await, None);
assert!(!world.pause(&ghost).await);
assert!(!world.resume(&ghost).await);
assert!(!world.cancel(&ghost).await);
assert!(world.pending_inputs().is_empty());
world.shutdown().await;
}
#[tokio::test]
async fn shutdown_ends_the_event_stream_and_further_requests_fail() {
let world = mock_world(vec![]);
let mut events = world.events();
let control = world.control.clone();
world.shutdown().await;
assert_eq!(
tokio::time::timeout(std::time::Duration::from_secs(5), events.next())
.await
.expect("stream ends"),
None
);
assert!(control.is_closed());
let (reply, _rx) = oneshot::channel();
assert!(control.send(ControlOp::List { reply }).is_err());
}
#[tokio::test]
async fn cancel_stops_a_parked_run() {
let dir = tempfile::tempdir().unwrap();
let world = mock_world(vec![with_tool(
"c1",
"ask_user_text",
serde_json::json!({"prompt": "?"}),
)]);
let mut events = world.events();
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(ASKER.to_string()),
"ask",
dir.path(),
))
.await
.expect("spawns");
tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Interaction { .. })),
)
.await
.expect("parked on the question");
assert!(world.cancel(&run_id).await);
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("terminal event after cancel");
assert!(seen.iter().any(|e| {
matches!(e, WorldEvent::Completed { status, .. } if status == "cancelled")
}));
world.shutdown().await;
}
struct CannedService;
impl crate::pipeline::ToolService for CannedService {
fn exec_for(
&self,
_entity: bevy_ecs::entity::Entity,
calls: Vec<leviath_providers::ToolCall>,
_progress: crate::pipeline::ToolProgress,
) -> crate::tool_bridge::BoxedToolExec {
Box::new(move || {
Box::pin(async move {
calls
.into_iter()
.map(|c| (c.id, "canned".to_string()))
.collect()
})
})
}
}
#[test]
fn fallback_models_accumulate_in_order_beside_the_default() {
let builder = AgentWorldBuilder::new()
.fallback_model("anthropic", "sonnet")
.fallback_model("openai", "gpt")
.default_model("openrouter", "deepseek");
assert_eq!(builder.defaults.provider, "openrouter");
assert_eq!(builder.defaults.model.as_deref(), Some("deepseek"));
assert_eq!(
builder
.defaults
.fallback_order
.iter()
.map(|e| (e.provider.as_str(), e.model.as_str()))
.collect::<Vec<_>>(),
vec![("anthropic", "sonnet"), ("openai", "gpt")]
);
}
#[tokio::test]
async fn every_builder_option_composes_and_state_dir_persists_runs() {
let dir = tempfile::tempdir().unwrap();
let state = tempfile::tempdir().unwrap();
let world = AgentWorld::builder()
.provider(ProviderCreds::simple("ollama"))
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(
vec![
with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
text("done"),
text("done"),
]
.into_iter()
.collect(),
),
}),
)
.default_model("mock", "m")
.fallback_model("ollama", "llama")
.fallback_model("mock", "spare")
.state_dir(state.path())
.inference_pool(InferencePoolConfig::new())
.tool_concurrency(2)
.build()
.expect("all options compose");
let mut events = world.events();
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(TWO_STAGE.to_string()),
"persist me",
dir.path(),
))
.await
.expect("spawns");
tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completes");
world.shutdown().await;
let run_dir = state.path().join("runs").join(run_id.as_ref());
assert!(run_dir.join("meta.json").exists());
assert!(state.path().join("machine-id").exists());
}
const SHELL_STAGE: &str = r#"[agent]
name = "shelly"
version = "0.0.0"
description = "One stage that can run commands."
entry_stage = "work"
[stages.work]
mode = "autonomous"
model = { provider = "mock", model = "m" }
description = "Do the work"
available_tools = ["shell"]
allow_complete = true
system_prompt = "Work."
[context.regions]
instructions = { kind = "pinned", max_tokens = 2000 }
conversation = { kind = "sliding_window", max_items = 40, max_tokens = 20000 }
"#;
async fn system_blocks_with(hints: leviath_core::config::PromptHints) -> Vec<String> {
let dir = tempfile::tempdir().unwrap();
let seen = Arc::new(Mutex::new(Vec::new()));
let world = AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Recorder {
seen: Arc::clone(&seen),
}),
)
.prompt_hints(hints)
.build()
.expect("world builds inside the test runtime");
let mut events = world.events();
world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(SHELL_STAGE.to_string()),
"go",
dir.path(),
))
.await
.expect("spawns");
tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completes");
world.shutdown().await;
let seen = seen.lock().unwrap_or_else(PoisonError::into_inner);
seen.first().cloned().expect("one inference happened")
}
#[tokio::test]
async fn prompt_hints_reach_the_request_and_are_off_by_default() {
let default_blocks = system_blocks_with(leviath_core::config::PromptHints {
batch_tool: false,
shell: false,
})
.await;
assert!(
default_blocks
.iter()
.all(|b| b != crate::pipeline::BATCH_TOOL_HINT),
);
assert!(default_blocks.iter().any(|b| b.contains("Work.")));
let hinted = system_blocks_with(leviath_core::config::PromptHints {
batch_tool: true,
shell: true,
})
.await;
assert_eq!(
hinted.first().map(String::as_str),
Some(crate::pipeline::BATCH_TOOL_HINT)
);
let shell_hint = crate::pipeline::shell_guidance_for(std::env::consts::OS);
assert_eq!(
hinted.iter().any(|b| Some(b.as_str()) == shell_hint),
shell_hint.is_some(),
);
}
#[tokio::test]
async fn a_custom_tool_service_replaces_the_builtin_one() {
let dir = tempfile::tempdir().unwrap();
let world = AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(
vec![
with_tool("c1", "read_file", serde_json::json!({"path": "x"})),
text("done"),
text("done"),
]
.into_iter()
.collect(),
),
}),
)
.tool_service(Arc::new(CannedService))
.build()
.expect("builds with a custom service");
let mut events = world.events();
world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(TWO_STAGE.to_string()),
"use the canned tools",
dir.path(),
))
.await
.expect("spawns");
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completes");
assert!(seen.iter().any(|e| {
matches!(e, WorldEvent::ToolCallFinished { summary, .. } if summary == "canned")
}));
world.shutdown().await;
}
#[tokio::test]
async fn manifest_files_that_do_not_parse_or_validate_fail_the_spawn() {
let dir = tempfile::tempdir().unwrap();
let world = mock_world(vec![]);
let garbled = dir.path().join("garbled.leviath");
std::fs::write(&garbled, "not = [valid").unwrap();
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Path(garbled),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().contains("parse manifest"));
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Path(PathBuf::from("")),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().starts_with("spawn error"));
let demanding = format!(
"{TWO_STAGE}\nspec = {{ kind = \"pinned\", max_tokens = 2000, seed = \"input\", required = true }}\n"
);
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(demanding),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().contains("required region"));
world.shutdown().await;
}
#[tokio::test]
async fn requests_after_the_world_closes_fail_closed() {
let world = mock_world(vec![]);
let (reply, _rx) = oneshot::channel();
world
.control
.send(ControlOp::Shutdown { reply })
.expect("world is up");
while !world.control.is_closed() {
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
let ghost = RunId("ghost".to_string());
assert_eq!(world.status(&ghost).await, None);
assert!(!world.pause(&ghost).await);
assert!(!world.send_message(&ghost, "hello").await);
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(TWO_STAGE.to_string()),
"t",
std::env::temp_dir(),
))
.await
.unwrap_err();
assert_eq!(err, EmbedError::ChannelClosed);
}
#[tokio::test]
async fn inline_blueprints_spawn_and_invalid_ones_are_refused() {
let dir = tempfile::tempdir().unwrap();
let world = mock_world(vec![text("done"), text("done"), text("done"), text("done")]);
let mut events = world.events();
let mut invalid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
invalid.entry_stage = Some("ghost".to_string());
let err = world
.spawn(SpawnSpec::new(
BlueprintSource::Inline(Box::new(invalid)),
"t",
dir.path(),
))
.await
.unwrap_err();
assert!(err.to_string().contains("invalid blueprint"));
let mut valid = leviath_core::manifest::parse_manifest(TWO_STAGE).unwrap();
valid.stages.truncate(1);
valid.stages[0].transitions = None;
valid.entry_stage = Some(valid.stages[0].name.clone());
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Inline(Box::new(valid)),
"just answer",
dir.path(),
))
.await
.expect("inline blueprint spawns");
assert!(run_id.as_ref().starts_with("embedded-"));
let seen = tokio::time::timeout(
std::time::Duration::from_secs(20),
events_until(&mut events, |e| matches!(e, WorldEvent::Completed { .. })),
)
.await
.expect("completes");
assert!(
seen.iter()
.any(|e| matches!(e, WorldEvent::Completed { .. }))
);
world.shutdown().await;
}
#[tokio::test]
async fn shutdown_survives_an_aborted_serve_loop() {
let world = mock_world(vec![]);
world.serve_task.abort();
world.shutdown().await;
}
#[tokio::test]
async fn the_mock_provider_is_a_minimal_stub() {
let mock = Mock {
responses: Mutex::new(VecDeque::new()),
};
assert_eq!(mock.count_tokens("x", "m").await, 1);
assert_eq!(mock.max_context_tokens("m"), 100_000);
assert_eq!(mock.name(), "mock");
let _ = mock.capabilities("m");
let recorder = Recorder {
seen: Arc::new(Mutex::new(Vec::new())),
};
assert_eq!(recorder.count_tokens("x", "m").await, 1);
assert_eq!(recorder.max_context_tokens("m"), 100_000);
assert_eq!(recorder.name(), "mock");
let _ = recorder.capabilities("m");
assert!(
mock.infer(
serde_json::from_value(serde_json::json!({
"messages": [],
"model": "m",
"max_tokens": 1,
"temperature": 0.0,
"tools": [],
"extra": null,
}))
.unwrap()
)
.await
.is_err()
);
}
#[tokio::test]
async fn pause_and_resume_round_trip_on_an_active_run() {
let dir = tempfile::tempdir().unwrap();
let mut pool = InferencePoolConfig::new();
pool.set_limit("m", 0);
let world = AgentWorld::builder()
.register_provider(
"mock",
Arc::new(Mock {
responses: Mutex::new(VecDeque::new()),
}),
)
.inference_pool(pool)
.build()
.expect("builds");
let run_id = world
.spawn(SpawnSpec::new(
BlueprintSource::Toml(ASKER.to_string()),
"wait around",
dir.path(),
))
.await
.expect("spawns");
assert_eq!(world.status(&run_id).await, Some(AgentStatus::Active));
assert!(world.pause(&run_id).await);
assert!(world.resume(&run_id).await);
assert!(world.cancel(&run_id).await);
world.shutdown().await;
}
#[test]
fn run_id_displays_as_its_string() {
let id = RunId("coder-1-2".to_string());
assert_eq!(id.to_string(), "coder-1-2");
assert_eq!(id.as_ref(), "coder-1-2");
}
}