use std::collections::{HashMap, HashSet};
use bevy_ecs::entity::Entity;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
use tokio::sync::{broadcast, oneshot};
use crate::components::{
AgentMessage, AgentState, AgentStatus, ContextWindow, ParentRef, SubAgentChildren,
};
use crate::interaction_hub::InteractionHub;
use crate::persistence::{RunMetadata, TokenTotals};
use crate::world::PipelineWorld;
use leviath_core::interaction::{InteractionRequest, InteractionResponse};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct SpawnArgs {
pub run_id: String,
pub blueprint_path: String,
pub task: String,
#[serde(default)]
pub regions: HashMap<String, String>,
#[serde(default)]
pub model: Option<String>,
pub workdir: String,
#[serde(default)]
pub metadata: HashMap<String, String>,
#[serde(default)]
pub callback_url: Option<String>,
#[serde(default)]
pub callback_secret: Option<String>,
#[serde(default)]
pub yolo: bool,
#[serde(default)]
pub no_seed_commands: bool,
#[serde(default)]
pub allow: Vec<String>,
#[serde(default)]
pub max_depth: Option<usize>,
#[serde(default)]
pub parent_run_id: Option<String>,
}
pub type Spawner = Box<dyn FnMut(&mut PipelineWorld, &SpawnArgs) -> Result<Entity, String> + Send>;
pub type Reloader = Box<dyn FnMut(&mut PipelineWorld, &str) -> Option<Entity> + Send>;
pub type ForceTerminator = Box<dyn FnMut(&str) -> bool + Send>;
pub type Reaper = Box<dyn FnMut(&mut PipelineWorld, Entity) + Send>;
pub type SpawnPreprocessor = Box<
dyn Fn(&SpawnArgs) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> + Send,
>;
pub enum SubAgentOp {
Spawn {
args: Box<SpawnArgs>,
parent_run_id: String,
max_depth: usize,
reply: oneshot::Sender<Result<String, String>>,
},
Check {
run_id: String,
reply: oneshot::Sender<Option<AgentStatus>>,
},
Send {
run_id: String,
caller_run_id: String,
content: String,
reply: oneshot::Sender<bool>,
},
Kill {
run_id: String,
caller_run_id: String,
reply: oneshot::Sender<bool>,
},
}
pub enum ControlOp {
Spawn {
args: Box<SpawnArgs>,
reply: oneshot::Sender<Result<String, String>>,
},
Status {
run_id: String,
reply: oneshot::Sender<Option<AgentStatus>>,
},
Pause {
run_id: String,
reply: oneshot::Sender<bool>,
},
Resume {
run_id: String,
reply: oneshot::Sender<bool>,
},
Cancel {
run_id: String,
reply: oneshot::Sender<bool>,
},
List {
reply: oneshot::Sender<Vec<(String, AgentStatus)>>,
},
Message {
agent_id: String,
content: String,
target_region: Option<String>,
reply: oneshot::Sender<bool>,
},
ListInteractions {
reply: oneshot::Sender<Vec<(String, InteractionRequest)>>,
},
AnswerInteraction {
response: InteractionResponse,
reply: oneshot::Sender<bool>,
},
CancelInteraction {
request_id: String,
reply: oneshot::Sender<bool>,
},
Shutdown {
reply: oneshot::Sender<bool>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum WorldEvent {
Spawned {
run_id: String,
agent_id: String,
blueprint: String,
},
Status {
run_id: String,
agent_id: String,
status: String,
stage: String,
iteration: usize,
tool_calls: usize,
accepts_messages: bool,
},
Tokens {
run_id: String,
agent_id: String,
prompt_tokens: usize,
completion_tokens: usize,
cached_tokens: usize,
cache_write_tokens: usize,
},
Context {
run_id: String,
agent_id: String,
total_tokens: usize,
max_tokens: usize,
},
Interaction {
run_id: String,
agent_id: String,
request: InteractionRequest,
},
Completed {
run_id: String,
agent_id: String,
status: String,
},
Log {
run_id: String,
agent_id: String,
line: String,
},
}
#[derive(bevy_ecs::resource::Resource, Clone)]
pub struct WorldEventSink(pub broadcast::Sender<WorldEvent>);
fn status_str(status: &AgentStatus) -> &'static str {
match status {
AgentStatus::Idle => "idle",
AgentStatus::Active => "active",
AgentStatus::Waiting => "waiting",
AgentStatus::Complete => "complete",
AgentStatus::Error { .. } => "error",
AgentStatus::Cancelled => "cancelled",
}
}
#[derive(Clone)]
struct Emitted {
status: &'static str,
stage: String,
iteration: usize,
tool_calls: usize,
accepts_messages: bool,
prompt_tokens: usize,
completion_tokens: usize,
cached_tokens: usize,
cache_write_tokens: usize,
context_tokens: usize,
terminal: bool,
}
pub struct WorldHost {
world: PipelineWorld,
by_run_id: HashMap<String, Entity>,
interactions: InteractionHub,
spawner: Option<Spawner>,
spawn_preprocessor: Option<SpawnPreprocessor>,
reloader: Option<Reloader>,
force_terminator: Option<ForceTerminator>,
reaper: Option<Reaper>,
events: broadcast::Sender<WorldEvent>,
emitted: HashMap<String, Emitted>,
emitted_interactions: HashSet<String>,
subagent_tx: UnboundedSender<SubAgentOp>,
subagent_rx: UnboundedReceiver<SubAgentOp>,
}
impl WorldHost {
pub fn new(world: PipelineWorld) -> Self {
Self::with_interactions(world, InteractionHub::new())
}
pub fn with_interactions(mut world: PipelineWorld, interactions: InteractionHub) -> Self {
let (events, _) = broadcast::channel(1024);
world
.world_mut()
.insert_resource(WorldEventSink(events.clone()));
let (subagent_tx, subagent_rx) = tokio::sync::mpsc::unbounded_channel();
Self {
world,
by_run_id: HashMap::new(),
interactions,
spawner: None,
spawn_preprocessor: None,
reloader: None,
force_terminator: None,
reaper: None,
events,
emitted: HashMap::new(),
emitted_interactions: HashSet::new(),
subagent_tx,
subagent_rx,
}
}
pub fn subagent_sender(&self) -> UnboundedSender<SubAgentOp> {
self.subagent_tx.clone()
}
pub fn subscribe(&self) -> broadcast::Receiver<WorldEvent> {
self.events.subscribe()
}
pub fn event_sender(&self) -> broadcast::Sender<WorldEvent> {
self.events.clone()
}
fn emit_events(&mut self) {
self.adopt_unregistered_runs();
let pairs: Vec<(String, Entity)> = self
.by_run_id
.iter()
.map(|(k, &v)| (k.clone(), v))
.collect();
let mut to_reap: Vec<(String, Entity)> = Vec::new();
for (run_id, entity) in pairs {
let Some(state) = self.world.world().get::<AgentState>(entity) else {
continue; };
let agent_id = state.agent_id.clone();
let status = status_str(&state.status);
let terminal = matches!(
state.status,
AgentStatus::Complete | AgentStatus::Error { .. } | AgentStatus::Cancelled
);
let cur = {
let totals = self
.world
.world()
.get::<TokenTotals>(entity)
.copied()
.unwrap_or_default();
let (context_tokens, _) = self
.world
.world()
.get::<ContextWindow>(entity)
.map(|w| (w.current_tokens, w.max_tokens))
.unwrap_or((0, 0));
Emitted {
status,
stage: state.current_stage.clone(),
iteration: state.iteration,
tool_calls: totals.tool_calls,
accepts_messages: state.accepts_messages,
prompt_tokens: totals.prompt_tokens,
completion_tokens: totals.completion_tokens,
cached_tokens: totals.cached_tokens,
cache_write_tokens: totals.cache_write_tokens,
context_tokens,
terminal,
}
};
let max_tokens = self
.world
.world()
.get::<ContextWindow>(entity)
.map(|w| w.max_tokens)
.unwrap_or(0);
let prev = self.emitted.get(&run_id).cloned();
if prev.is_none() {
let blueprint = self
.world
.world()
.get::<RunMetadata>(entity)
.map(|m| m.agent_name.clone())
.unwrap_or_default();
let _ = self.events.send(WorldEvent::Spawned {
run_id: run_id.clone(),
agent_id: agent_id.clone(),
blueprint,
});
}
let status_key = |e: &Emitted| {
(
e.status,
e.stage.clone(),
e.iteration,
e.tool_calls,
e.accepts_messages,
)
};
if prev.as_ref().map(status_key) != Some(status_key(&cur)) {
let _ = self.events.send(WorldEvent::Status {
run_id: run_id.clone(),
agent_id: agent_id.clone(),
status: status.to_string(),
stage: cur.stage.clone(),
iteration: cur.iteration,
tool_calls: cur.tool_calls,
accepts_messages: cur.accepts_messages,
});
}
let token_key = |e: &Emitted| {
(
e.prompt_tokens,
e.completion_tokens,
e.cached_tokens,
e.cache_write_tokens,
)
};
if prev.as_ref().map(token_key) != Some(token_key(&cur)) {
let _ = self.events.send(WorldEvent::Tokens {
run_id: run_id.clone(),
agent_id: agent_id.clone(),
prompt_tokens: cur.prompt_tokens,
completion_tokens: cur.completion_tokens,
cached_tokens: cur.cached_tokens,
cache_write_tokens: cur.cache_write_tokens,
});
}
if prev.as_ref().map(|e| e.context_tokens) != Some(cur.context_tokens) {
let _ = self.events.send(WorldEvent::Context {
run_id: run_id.clone(),
agent_id: agent_id.clone(),
total_tokens: cur.context_tokens,
max_tokens,
});
}
let was_terminal = prev.as_ref().map(|e| e.terminal) == Some(true);
if cur.terminal && !was_terminal {
let _ = self.events.send(WorldEvent::Completed {
run_id: run_id.clone(),
agent_id: agent_id.clone(),
status: status.to_string(),
});
}
if cur.terminal && was_terminal && self.no_live_parent(entity) {
to_reap.push((run_id.clone(), entity));
}
self.emitted.insert(run_id, cur);
}
let mut reaper = self.reaper.take();
for (run_id, entity) in to_reap {
if let Some(reaper) = reaper.as_mut() {
reaper(&mut self.world, entity);
}
self.world.world_mut().despawn(entity);
self.by_run_id.remove(&run_id);
self.emitted.remove(&run_id);
}
self.reaper = reaper;
for (agent_id, request) in self.interactions.pending() {
if self.emitted_interactions.insert(request.id.clone()) {
let _ = self.events.send(WorldEvent::Interaction {
run_id: agent_id.clone(),
agent_id,
request,
});
}
}
}
fn adopt_unregistered_runs(&mut self) {
let live: Vec<(String, Entity)> = self
.world
.world_mut()
.query::<(Entity, &RunMetadata)>()
.iter(self.world.world())
.map(|(entity, md)| (md.run_id.clone(), entity))
.collect();
for (run_id, entity) in live {
if self.live_entity(&run_id) != Some(entity) {
self.by_run_id.insert(run_id, entity);
}
}
}
fn no_live_parent(&self, entity: Entity) -> bool {
let world = self.world.world();
match world.get::<crate::components::ParentRef>(entity) {
None => true,
Some(parent_ref) => match world.get::<AgentState>(parent_ref.parent_entity) {
None => true,
Some(state) => matches!(
state.status,
AgentStatus::Complete | AgentStatus::Error { .. } | AgentStatus::Cancelled
),
},
}
}
pub fn set_spawner(&mut self, spawner: Spawner) {
self.spawner = Some(spawner);
}
pub fn set_spawn_preprocessor(&mut self, pp: SpawnPreprocessor) {
self.spawn_preprocessor = Some(pp);
}
pub fn set_reloader(&mut self, reloader: Reloader) {
self.reloader = Some(reloader);
}
pub fn set_force_terminator(&mut self, force_terminator: ForceTerminator) {
self.force_terminator = Some(force_terminator);
}
pub fn set_reaper(&mut self, reaper: Reaper) {
self.reaper = Some(reaper);
}
fn resolve_or_reload(&mut self, run_id: &str) -> Option<Entity> {
if let Some(entity) = self.live_entity(run_id) {
return Some(entity);
}
let entity = (self.reloader.as_mut()?)(&mut self.world, run_id)?;
self.by_run_id.insert(run_id.to_string(), entity);
Some(entity)
}
pub fn interactions(&self) -> InteractionHub {
self.interactions.clone()
}
pub fn world_mut(&mut self) -> &mut PipelineWorld {
&mut self.world
}
pub fn register(&mut self, run_id: impl Into<String>, entity: Entity) {
self.by_run_id.insert(run_id.into(), entity);
}
fn live_entity(&self, run_id: &str) -> Option<Entity> {
let entity = *self.by_run_id.get(run_id)?;
self.world.world().get::<AgentState>(entity).map(|_| entity)
}
fn handle_subagent(&mut self, op: SubAgentOp) {
match op {
SubAgentOp::Spawn {
args,
parent_run_id,
max_depth,
reply,
} => {
let _ = reply.send(self.spawn_child(*args, &parent_run_id, max_depth));
}
SubAgentOp::Check { run_id, reply } => {
let status = self
.live_entity(&run_id)
.and_then(|e| self.world.agent_status(e));
let _ = reply.send(status);
}
SubAgentOp::Send {
run_id,
caller_run_id,
content,
reply,
} => {
if !self.is_within_tree(&run_id, &caller_run_id) {
let _ = reply.send(false);
return;
}
self.resolve_or_reload(&run_id);
let ok = self
.world
.send_message(AgentMessage {
agent_id: run_id,
content,
target_region: None,
priority: 0,
})
.is_ok();
let _ = reply.send(ok);
}
SubAgentOp::Kill {
run_id,
caller_run_id,
reply,
} => {
let within = self.is_within_tree(&run_id, &caller_run_id);
let _ = reply.send(within && self.cancel_tree(&run_id));
}
}
}
fn spawn_child(
&mut self,
mut args: SpawnArgs,
parent_run_id: &str,
max_depth: usize,
) -> Result<String, String> {
args.parent_run_id = Some(parent_run_id.to_string());
let parent = self
.live_entity(parent_run_id)
.ok_or_else(|| format!("parent run '{parent_run_id}' is not live"))?;
let parent_depth = self
.world
.world()
.get::<ParentRef>(parent)
.map_or(0, |p| p.depth);
let child_depth = parent_depth + 1;
if child_depth > max_depth {
return Err(format!(
"sub-agent depth limit ({max_depth}) reached; not spawning deeper"
));
}
let run_id = args.run_id.clone();
let child = match self.spawner.as_mut() {
Some(spawner) => spawner(&mut self.world, &args)?,
None => return Err("this daemon cannot spawn agents".to_string()),
};
let world = self.world.world_mut();
world.entity_mut(child).insert(ParentRef {
parent_entity: parent,
parent_agent_id: parent_run_id.to_string(),
depth: child_depth,
});
match world.get_mut::<SubAgentChildren>(parent) {
Some(mut kids) => kids.children.push(child),
None => {
world.entity_mut(parent).insert(SubAgentChildren {
children: vec![child],
max_child_depth: max_depth,
});
}
}
world
.get_mut::<crate::components::AgentState>(parent)
.expect("a spawning parent always has AgentState")
.spawned_children_ids
.push(run_id.clone());
crate::context_transform::apply_context_transforms(world, parent, child);
self.by_run_id.insert(run_id.clone(), child);
Ok(run_id)
}
fn is_within_tree(&mut self, run_id: &str, ancestor: &str) -> bool {
if run_id == ancestor {
return true;
}
let (Some(target), Some(root)) = (
self.resolve_or_reload(run_id),
self.resolve_or_reload(ancestor),
) else {
return false;
};
let mut stack = vec![root];
while let Some(e) = stack.pop() {
if e == target {
return true;
}
if let Some(kids) = self.world.world().get::<SubAgentChildren>(e) {
stack.extend(kids.children.iter().copied());
}
}
false
}
fn cancel_tree(&mut self, run_id: &str) -> bool {
let Some(root) = self.resolve_or_reload(run_id) else {
return false;
};
let mut subtree = Vec::new();
let mut stack = vec![root];
while let Some(e) = stack.pop() {
subtree.push(e);
if let Some(kids) = self.world.world().get::<SubAgentChildren>(e) {
stack.extend(kids.children.iter().copied());
}
}
let mut cancelled = false;
for e in subtree {
let agent_id = self
.world
.world()
.get::<AgentState>(e)
.map(|s| s.agent_id.clone());
cancelled |= self.world.cancel(e);
if let Some(agent_id) = agent_id {
self.interactions.cancel_for_agent(&agent_id);
let still_open: HashSet<String> = self
.interactions
.pending()
.into_iter()
.map(|(_, req)| req.id)
.collect();
self.emitted_interactions
.retain(|id| still_open.contains(id));
}
}
cancelled
}
fn list(&self) -> Vec<(String, AgentStatus)> {
self.by_run_id
.iter()
.filter_map(|(run_id, &entity)| {
self.world
.world()
.get::<AgentState>(entity)
.map(|s| (run_id.clone(), s.status.clone()))
})
.collect()
}
pub fn handle(&mut self, op: ControlOp) {
match op {
ControlOp::Spawn { args, reply } => {
let result = match self.spawner.as_mut() {
Some(spawner) => {
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
spawner(&mut self.world, &args)
})) {
Ok(Ok(entity)) => {
self.by_run_id.insert(args.run_id.clone(), entity);
Ok(args.run_id.clone())
}
Ok(Err(e)) => Err(e),
Err(_) => Err("agent spawn panicked".to_string()),
}
}
None => Err("this daemon cannot spawn agents".to_string()),
};
if let Err(error) = &result {
tracing::error!(
run_id = %args.run_id,
blueprint = %args.blueprint_path,
workdir = %args.workdir,
error = %error,
"agent spawn failed"
);
}
let _ = reply.send(result);
}
ControlOp::Status { run_id, reply } => {
let status = self
.live_entity(&run_id)
.and_then(|e| self.world.agent_status(e));
let _ = reply.send(status);
}
ControlOp::Pause { run_id, reply } => {
let ok = self
.resolve_or_reload(&run_id)
.is_some_and(|e| self.world.pause(e));
let _ = reply.send(ok);
}
ControlOp::Resume { run_id, reply } => {
let ok = self
.resolve_or_reload(&run_id)
.is_some_and(|e| self.world.resume(e));
let _ = reply.send(ok);
}
ControlOp::Cancel { run_id, reply } => {
let ok = self.cancel_tree(&run_id)
|| self
.force_terminator
.as_mut()
.is_some_and(|terminate| terminate(&run_id));
let _ = reply.send(ok);
}
ControlOp::List { reply } => {
let _ = reply.send(self.list());
}
ControlOp::Message {
agent_id,
content,
target_region,
reply,
} => {
self.resolve_or_reload(&agent_id);
let ok = self
.world
.send_message(AgentMessage {
agent_id,
content,
target_region,
priority: 0,
})
.is_ok();
let _ = reply.send(ok);
}
ControlOp::ListInteractions { reply } => {
let _ = reply.send(self.interactions.pending());
}
ControlOp::AnswerInteraction { response, reply } => {
let _ = reply.send(self.interactions.answer(response));
}
ControlOp::CancelInteraction { request_id, reply } => {
let _ = reply.send(self.interactions.cancel(&request_id));
}
ControlOp::Shutdown { reply } => {
let _ = reply.send(true);
self.world.shutdown();
}
}
}
pub async fn flush_and_stop(&mut self) {
self.world.flush_and_stop().await;
}
pub async fn serve(&mut self, mut control_rx: UnboundedReceiver<ControlOp>) {
let wake = self.world.wake_handle();
let shutdown = self.world.shutdown_handle();
'serve: loop {
self.world.run_to_fixed_point();
self.emit_events();
tokio::select! {
_ = wake.notified() => {}
_ = shutdown.notified() => break 'serve,
op = control_rx.recv() => {
match op {
Some(op) => {
let pre = match &op {
ControlOp::Spawn { args, .. } => {
self.spawn_preprocessor.as_ref().map(|pp| pp(args))
}
_ => None,
};
if let Some(fut) = pre {
fut.await;
}
self.handle(op);
}
None => break 'serve, }
}
Some(sub) = self.subagent_rx.recv() => {
let pre = match &sub {
SubAgentOp::Spawn { args, .. } => {
self.spawn_preprocessor.as_ref().map(|pp| pp(args))
}
_ => None,
};
if let Some(fut) = pre {
fut.await;
}
self.handle_subagent(sub);
}
}
}
self.flush_and_stop().await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dynamic_interaction::InteractionBackend;
use crate::inference_pool::InferencePoolConfig;
use crate::pipeline::{
AgentBlueprint, ReadyToInfer, StageCursor, StageInference, StageInferences, StageProgress,
StageSetup, StageSetups, ToolService, VisitCounts, WaitingForChildren,
};
use crate::tool_bridge::BoxedToolExec;
use leviath_core::{Region, RegionKind};
use leviath_providers::{
FinishReason, InferenceRequest, InferenceResponse, ModelCapabilities, Provider,
ProviderError, TokenUsage,
};
use std::sync::Arc;
use std::sync::Mutex;
use tokio::runtime::Handle;
use tokio::sync::mpsc;
struct Script {
responses: Mutex<std::collections::VecDeque<InferenceResponse>>,
}
#[async_trait::async_trait]
impl Provider for Script {
async fn infer(
&self,
_req: InferenceRequest,
) -> leviath_providers::Result<InferenceResponse> {
self.responses
.lock()
.unwrap()
.pop_front()
.ok_or_else(|| ProviderError::Other("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 {
"script"
}
fn capabilities(&self, _m: &str) -> ModelCapabilities {
ModelCapabilities::default()
}
}
struct NoTools;
impl ToolService for NoTools {
fn exec_for(&self, _e: Entity, calls: Vec<leviath_providers::ToolCall>) -> BoxedToolExec {
Box::new(move || {
Box::pin(async move { calls.into_iter().map(|c| (c.id, String::new())).collect() })
})
}
}
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 host_with(responses: Vec<InferenceResponse>) -> WorldHost {
let mut registry = crate::providers::ProviderRegistry::new();
registry.register(
"script".to_string(),
Arc::new(Script {
responses: Mutex::new(responses.into_iter().collect()),
}),
);
let world = PipelineWorld::new(
registry,
Arc::new(NoTools),
InferencePoolConfig::new(),
1,
std::env::temp_dir(),
Handle::current(),
);
WorldHost::new(world)
}
fn blueprint() -> leviath_core::Blueprint {
let layout = leviath_core::layout::ContextLayout::new(
vec![leviath_core::layout::RegionDefinition::new(
"conversation".to_string(),
RegionKind::Clearable,
10_000,
)],
12_000,
);
let s = leviath_core::Stage::new(
"s".to_string(),
leviath_core::blueprint::ModelConfig::new("script".to_string(), "m".to_string()),
);
leviath_core::Blueprint::new("t".to_string(), "d".to_string(), vec![s], layout)
}
fn window() -> crate::components::ContextWindow {
let mut w = crate::components::ContextWindow::new(10_000);
w.add_region(Region::new(
"conversation".to_string(),
RegionKind::Clearable,
10_000,
));
w
}
fn agent_state(agent_id: &str) -> AgentState {
AgentState {
agent_id: agent_id.to_string(),
current_stage: "s".to_string(),
iteration: 0,
status: AgentStatus::Active,
spawned_children_ids: vec![],
pending_wait: None,
accepts_messages: true,
}
}
fn si() -> StageInference {
StageInference {
provider_name: "script".to_string(),
model: "m".to_string(),
tools: vec![],
tool_filter: None,
}
}
fn setup() -> StageSetup {
StageSetup {
inference_config: crate::components::InferenceConfig {
temperature: None,
max_output_tokens: None,
extra_params: Default::default(),
batch_tool_hint: false,
request_timeout_secs: None,
},
routing: None,
accepts_messages: true,
context_layout: None,
system_prompt: None,
}
}
fn spawn(host: &mut WorldHost, run_id: &str, agent_id: &str) -> Entity {
let e = host.world_mut().spawn_agent((
AgentBlueprint(blueprint()),
StageCursor { index: 0 },
agent_state(agent_id),
crate::components::MessageInbox::default(),
StageProgress::default(),
StageInferences(vec![si()]),
StageSetups(vec![setup()]),
VisitCounts::default(),
window(),
si(),
setup().inference_config,
ReadyToInfer,
));
host.register(run_id, e);
e
}
fn recording_terminator(seen: Arc<Mutex<Vec<String>>>) -> ForceTerminator {
Box::new(move |run_id| {
seen.lock().unwrap().push(run_id.to_string());
run_id != "never-existed"
})
}
fn paging_reloader() -> Reloader {
Box::new(|world, run_id| Some(world.spawn_agent((agent_state(run_id),))))
}
async fn ask<T>(host: &mut WorldHost, make: impl FnOnce(oneshot::Sender<T>) -> ControlOp) -> T {
let (tx, rx) = oneshot::channel();
host.handle(make(tx));
rx.await.unwrap()
}
#[tokio::test]
async fn status_and_list_reflect_registered_runs() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "agent-a");
let status = ask(&mut host, |reply| ControlOp::Status {
run_id: "run-a".to_string(),
reply,
})
.await;
assert_eq!(status, Some(AgentStatus::Active));
let list = ask(&mut host, |reply| ControlOp::List { reply }).await;
assert_eq!(list, vec![("run-a".to_string(), AgentStatus::Active)]);
let none = ask(&mut host, |reply| ControlOp::Status {
run_id: "ghost".to_string(),
reply,
})
.await;
assert_eq!(none, None);
}
#[tokio::test]
async fn pause_resume_cancel_by_run_id() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "agent-a");
assert!(
ask(&mut host, |reply| ControlOp::Pause {
run_id: "run-a".to_string(),
reply
})
.await
);
assert_eq!(
host.world.agent_status(host.by_run_id["run-a"]),
Some(AgentStatus::Idle)
);
assert!(
ask(&mut host, |reply| ControlOp::Resume {
run_id: "run-a".to_string(),
reply
})
.await
);
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "run-a".to_string(),
reply
})
.await
);
assert_eq!(
host.world.agent_status(host.by_run_id["run-a"]),
Some(AgentStatus::Cancelled)
);
assert!(
!ask(&mut host, |reply| ControlOp::Pause {
run_id: "ghost".to_string(),
reply
})
.await
);
assert!(
!ask(&mut host, |reply| ControlOp::Resume {
run_id: "ghost".to_string(),
reply
})
.await
);
assert!(
!ask(&mut host, |reply| ControlOp::Cancel {
run_id: "ghost".to_string(),
reply
})
.await
);
}
#[tokio::test]
async fn spawn_op_uses_installed_spawner_and_registers() {
let mut host = host_with(vec![]);
host.set_spawner(Box::new(|world, args| {
Ok(world.spawn_agent((agent_state(&args.run_id),)))
}));
let result = ask(&mut host, |reply| ControlOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "r1".to_string(),
..Default::default()
}),
reply,
})
.await;
assert_eq!(result, Ok("r1".to_string()));
let status = ask(&mut host, |reply| ControlOp::Status {
run_id: "r1".to_string(),
reply,
})
.await;
assert_eq!(status, Some(AgentStatus::Active));
}
#[tokio::test]
async fn spawn_op_propagates_spawner_error() {
let mut host = host_with(vec![]);
host.set_spawner(Box::new(|_world, _args| Err("bad blueprint".to_string())));
let result = ask(&mut host, |reply| ControlOp::Spawn {
args: Box::new(SpawnArgs::default()),
reply,
})
.await;
assert_eq!(result, Err("bad blueprint".to_string()));
}
#[tokio::test]
async fn spawn_op_contains_a_panicking_spawner() {
let mut host = host_with(vec![]);
host.set_spawner(Box::new(|_world, _args| panic!("simulated spawn panic")));
let (tx, rx) = oneshot::channel();
crate::test_support::with_silenced_panics(|| {
host.handle(ControlOp::Spawn {
args: Box::new(SpawnArgs::default()),
reply: tx,
});
});
assert_eq!(rx.await.unwrap(), Err("agent spawn panicked".to_string()));
let status = ask(&mut host, |reply| ControlOp::Status {
run_id: SpawnArgs::default().run_id,
reply,
})
.await;
assert!(status.is_none());
}
#[tokio::test]
async fn spawn_op_errors_without_a_spawner() {
let mut host = host_with(vec![]);
let result = ask(&mut host, |reply| ControlOp::Spawn {
args: Box::new(SpawnArgs::default()),
reply,
})
.await;
assert!(result.unwrap_err().contains("cannot spawn"));
}
async fn ask_sub<T>(
host: &mut WorldHost,
make: impl FnOnce(oneshot::Sender<T>) -> SubAgentOp,
) -> T {
let (tx, rx) = oneshot::channel();
host.handle_subagent(make(tx));
rx.await.unwrap()
}
fn child_spawner() -> Spawner {
Box::new(|world, args| Ok(world.spawn_agent((agent_state(&args.run_id),))))
}
#[tokio::test]
async fn subagent_spawn_links_child_and_registers() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
let parent = spawn(&mut host, "parent", "parent");
let result = ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "child".to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await;
assert_eq!(result, Ok("child".to_string()));
let child = host.by_run_id["child"];
let pref = host.world.world().get::<ParentRef>(child).unwrap();
assert_eq!(pref.parent_entity, parent);
assert_eq!(pref.depth, 1);
let kids = host.world.world().get::<SubAgentChildren>(parent).unwrap();
assert_eq!(kids.children, vec![child]);
}
#[tokio::test]
async fn subagent_spawn_appends_to_existing_children() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
spawn(&mut host, "parent", "parent");
for id in ["c1", "c2"] {
let r = ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: id.to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await;
assert!(r.is_ok());
}
let parent = host.by_run_id["parent"];
let kids = host.world.world().get::<SubAgentChildren>(parent).unwrap();
assert_eq!(kids.children.len(), 2);
}
#[tokio::test]
async fn subagent_spawn_rejects_beyond_max_depth() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
spawn(&mut host, "parent", "parent");
let result = ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "child".to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 0, reply,
})
.await;
assert!(result.unwrap_err().contains("depth limit"));
assert!(!host.by_run_id.contains_key("child"));
}
#[tokio::test]
async fn subagent_spawn_unknown_parent_and_no_spawner_and_spawner_error() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
let r = ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs::default()),
parent_run_id: "ghost".to_string(),
max_depth: 3,
reply,
})
.await;
assert!(r.unwrap_err().contains("not live"));
let mut host2 = host_with(vec![]);
spawn(&mut host2, "parent", "parent");
let r = ask_sub(&mut host2, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs::default()),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await;
assert!(r.unwrap_err().contains("cannot spawn"));
let mut host3 = host_with(vec![]);
host3.set_spawner(Box::new(|_w, _a| Err("bad blueprint".to_string())));
spawn(&mut host3, "parent", "parent");
let r = ask_sub(&mut host3, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs::default()),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await;
assert_eq!(r, Err("bad blueprint".to_string()));
}
#[tokio::test]
async fn subagent_check_reports_status_or_none() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "run-a");
let status = ask_sub(&mut host, |reply| SubAgentOp::Check {
run_id: "run-a".to_string(),
reply,
})
.await;
assert_eq!(status, Some(AgentStatus::Active));
let none = ask_sub(&mut host, |reply| SubAgentOp::Check {
run_id: "ghost".to_string(),
reply,
})
.await;
assert_eq!(none, None);
}
#[tokio::test]
async fn subagent_ops_reach_a_run_the_caller_spawned() {
let mut host = host_with(vec![]);
let parent = spawn(&mut host, "parent", "parent");
let child = spawn(&mut host, "child", "child");
host.world_mut()
.world_mut()
.entity_mut(parent)
.insert(SubAgentChildren {
children: vec![child],
max_child_depth: 3,
});
let delivered = ask_sub(&mut host, |reply| SubAgentOp::Send {
run_id: "child".to_string(),
caller_run_id: "parent".to_string(),
content: "carry on".to_string(),
reply,
})
.await;
assert!(delivered, "a run we spawned is ours to message");
}
#[tokio::test]
async fn subagent_ops_refuse_a_run_outside_the_callers_tree() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "run-a");
spawn(&mut host, "outsider", "outsider");
let delivered = ask_sub(&mut host, |reply| SubAgentOp::Send {
run_id: "outsider".to_string(),
caller_run_id: "run-a".to_string(),
content: "take this".to_string(),
reply,
})
.await;
assert!(!delivered, "a run we did not spawn is not ours to message");
let killed = ask_sub(&mut host, |reply| SubAgentOp::Kill {
run_id: "outsider".to_string(),
caller_run_id: "run-a".to_string(),
reply,
})
.await;
assert!(!killed, "nor ours to cancel");
let phantom = ask_sub(&mut host, |reply| SubAgentOp::Send {
run_id: "no-such-run".to_string(),
caller_run_id: "run-a".to_string(),
content: "hello?".to_string(),
reply,
})
.await;
assert!(!phantom, "an unknown run id is in nobody's tree");
}
#[tokio::test]
async fn subagent_send_delivers_to_inbox() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "run-a");
let ok = ask_sub(&mut host, |reply| SubAgentOp::Send {
run_id: "run-a".to_string(),
caller_run_id: "run-a".to_string(),
content: "hello child".to_string(),
reply,
})
.await;
assert!(ok);
}
#[tokio::test]
async fn subagent_kill_cancels_the_whole_tree() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
spawn(&mut host, "parent", "parent");
ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "child".to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await
.unwrap();
let ok = ask_sub(&mut host, |reply| SubAgentOp::Kill {
run_id: "parent".to_string(),
caller_run_id: "parent".to_string(),
reply,
})
.await;
assert!(ok);
assert_eq!(
host.world.agent_status(host.by_run_id["parent"]),
Some(AgentStatus::Cancelled)
);
assert_eq!(
host.world.agent_status(host.by_run_id["child"]),
Some(AgentStatus::Cancelled)
);
let miss = ask_sub(&mut host, |reply| SubAgentOp::Kill {
run_id: "ghost".to_string(),
caller_run_id: "ghost".to_string(),
reply,
})
.await;
assert!(!miss);
}
#[tokio::test]
async fn cancel_cascades_to_the_whole_tree() {
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
spawn(&mut host, "parent", "parent");
ask_sub(&mut host, |reply| SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "child".to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply,
})
.await
.unwrap();
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "parent".to_string(),
reply
})
.await
);
assert_eq!(
host.world.agent_status(host.by_run_id["child"]),
Some(AgentStatus::Cancelled),
"cancelling the parent cancels its children"
);
}
#[tokio::test]
async fn cancel_tolerates_a_child_that_has_already_been_reaped() {
let mut host = host_with(vec![]);
let parent = spawn(&mut host, "parent", "parent");
let ghost = host.world_mut().spawn_agent((agent_state("ghost"),));
host.world_mut()
.world_mut()
.entity_mut(parent)
.insert(SubAgentChildren {
children: vec![ghost],
max_child_depth: 3,
});
host.world_mut().world_mut().despawn(ghost);
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "parent".to_string(),
reply
})
.await,
"the parent is still cancelled"
);
assert_eq!(
host.world.agent_status(parent),
Some(AgentStatus::Cancelled)
);
}
#[tokio::test]
async fn cancel_closes_the_runs_open_interactions() {
let mut host = host_with(vec![]);
let hub = host.interactions();
spawn(&mut host, "run-a", "agent-a");
let backend = hub.backend_for("agent-a");
let asking = tokio::spawn(async move {
backend
.ask(InteractionRequest::free_text("q", "ask", "stage", true))
.await
});
while hub.pending().is_empty() {
tokio::task::yield_now().await;
}
host.emit_events();
assert!(
!host.emitted_interactions.is_empty(),
"the open request was emitted"
);
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "run-a".to_string(),
reply,
})
.await;
tokio::time::timeout(std::time::Duration::from_secs(5), asking)
.await
.expect("cancelling the run releases its blocked ask")
.expect("the ask task did not panic");
assert!(hub.pending().is_empty(), "no orphaned prompt is left open");
assert!(
host.emitted_interactions.is_empty(),
"and it is pruned from the emitted set, not re-announced forever"
);
}
#[tokio::test]
async fn cancel_falls_back_to_the_force_terminator_when_the_world_cannot_hold_the_run() {
let mut host = host_with(vec![]);
host.set_reloader(Box::new(|_world, _run_id| None));
let terminated = Arc::new(Mutex::new(Vec::new()));
host.set_force_terminator(recording_terminator(terminated.clone()));
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "unreloadable".to_string(),
reply
})
.await,
"a run that can't be reloaded is still terminated"
);
assert!(
!ask(&mut host, |reply| ControlOp::Cancel {
run_id: "never-existed".to_string(),
reply
})
.await,
"`false` is reserved for a run that exists nowhere"
);
assert_eq!(
*terminated.lock().unwrap(),
vec!["unreloadable".to_string(), "never-existed".to_string()]
);
}
#[tokio::test]
async fn cancel_does_not_force_terminate_a_run_it_could_cancel() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "agent-a");
let terminated = Arc::new(Mutex::new(Vec::new()));
host.set_force_terminator(recording_terminator(terminated.clone()));
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "run-a".to_string(),
reply
})
.await
);
assert_eq!(
host.world.agent_status(host.by_run_id["run-a"]),
Some(AgentStatus::Cancelled)
);
assert!(
terminated.lock().unwrap().is_empty(),
"the disk fallback stayed unused"
);
}
#[tokio::test]
async fn unregistered_world_agents_are_adopted_and_become_cancellable() {
let mut host = host_with(vec![]);
let entity = host.world_mut().spawn_agent((
agent_state("worker"),
RunMetadata {
run_id: "worker-run".to_string(),
agent_name: "w".to_string(),
agent_path: String::new(),
task: String::new(),
model: None,
workdir: String::new(),
num_stages: 1,
started_at: 0,
parent_run_id: None,
metadata: Default::default(),
callback_url: None,
callback_secret: None,
title: None,
},
));
assert!(
!host.by_run_id.contains_key("worker-run"),
"not registered by the spawn itself"
);
host.emit_events();
assert_eq!(host.live_entity("worker-run"), Some(entity), "adopted");
host.set_reloader(paging_reloader());
assert!(
ask(&mut host, |reply| ControlOp::Cancel {
run_id: "worker-run".to_string(),
reply
})
.await
);
assert_eq!(
host.world.agent_status(entity),
Some(AgentStatus::Cancelled),
"the original entity is cancelled, not a reloaded copy"
);
}
#[tokio::test]
async fn interaction_ops_list_answer_and_cancel() {
let mut host = host_with(vec![]);
let hub = host.interactions();
let backend = hub.backend_for("agent-a");
let asking = tokio::spawn(async move {
backend
.ask(leviath_core::interaction::InteractionRequest::free_text(
"q1", "prompt?", "stage", true,
))
.await
});
for _ in 0..8 {
tokio::task::yield_now().await;
}
let list = ask(&mut host, |reply| ControlOp::ListInteractions { reply }).await;
assert_eq!(list.len(), 1);
assert_eq!(list[0].0, "agent-a");
let ok = ask(&mut host, |reply| ControlOp::AnswerInteraction {
response: leviath_core::interaction::InteractionResponse::text("q1", "hi"),
reply,
})
.await;
assert!(ok);
assert_eq!(asking.await.unwrap().value.as_deref(), Some("hi"));
let cancelled = ask(&mut host, |reply| ControlOp::CancelInteraction {
request_id: "gone".to_string(),
reply,
})
.await;
assert!(!cancelled);
}
#[tokio::test]
async fn cancel_interaction_op_wakes_asker() {
let mut host = host_with(vec![]);
let backend = host.interactions().backend_for("agent-a");
let asking = tokio::spawn(async move {
backend
.ask(leviath_core::interaction::InteractionRequest::free_text(
"q2", "p", "s", true,
))
.await
});
for _ in 0..8 {
tokio::task::yield_now().await;
}
let ok = ask(&mut host, |reply| ControlOp::CancelInteraction {
request_id: "q2".to_string(),
reply,
})
.await;
assert!(ok);
assert_eq!(asking.await.unwrap().request_id, "q2");
}
#[tokio::test]
async fn message_op_is_delivered() {
let mut host = host_with(vec![]);
let e = spawn(&mut host, "run-a", "agent-a");
let ok = ask(&mut host, |reply| ControlOp::Message {
agent_id: "agent-a".to_string(),
content: "hi".to_string(),
target_region: Some("conversation".to_string()),
reply,
})
.await;
assert!(ok);
host.world_mut().tick();
assert!(
host.world
.world()
.get::<crate::components::ContextWindow>(e)
.unwrap()
.get_region("conversation")
.unwrap()
.current_tokens
> 0
);
}
#[tokio::test]
async fn serve_drives_agents_and_handles_ops_until_shutdown() {
let mut host = host_with(vec![text("t1"), text("t2"), text("t3"), text("t4")]);
let e = spawn(&mut host, "run-a", "agent-a");
let shutdown = host.world_mut().shutdown_handle();
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move {
host.serve(op_rx).await;
host
});
let (tx, rx) = oneshot::channel();
op_tx
.send(ControlOp::Status {
run_id: "run-a".to_string(),
reply: tx,
})
.unwrap();
let _ = rx.await.unwrap();
shutdown.notify_one();
let host = handle.await.unwrap();
assert_eq!(host.world.agent_status(e), Some(AgentStatus::Complete));
}
#[tokio::test]
async fn serve_awaits_spawn_preprocessor_before_spawning() {
use std::sync::atomic::{AtomicBool, Ordering};
let mut host = host_with(vec![]);
let ran = Arc::new(AtomicBool::new(false));
let ran_pp = ran.clone();
host.set_spawn_preprocessor(Box::new(move |_args| {
let ran = ran_pp.clone();
Box::pin(async move {
ran.store(true, Ordering::SeqCst);
})
}));
let ran_spawn = ran.clone();
host.set_spawner(Box::new(move |world, args| {
assert!(ran_spawn.load(Ordering::SeqCst));
Ok(world.spawn_agent((agent_state(&args.run_id),)))
}));
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move {
host.serve(op_rx).await;
});
let (tx, rx) = oneshot::channel();
op_tx
.send(ControlOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "rp".to_string(),
..Default::default()
}),
reply: tx,
})
.unwrap();
let result = rx.await.unwrap();
drop(op_tx); handle.await.unwrap();
assert_eq!(result, Ok("rp".to_string()));
assert!(ran.load(Ordering::SeqCst), "preprocessor ran");
}
#[tokio::test]
async fn serve_awaits_preprocessor_for_subagent_spawn() {
use std::sync::atomic::{AtomicUsize, Ordering};
let mut host = host_with(vec![]);
host.set_spawner(child_spawner());
let _parent = spawn(&mut host, "parent", "parent");
let calls = Arc::new(AtomicUsize::new(0));
let calls_pp = calls.clone();
host.set_spawn_preprocessor(Box::new(move |_args| {
let calls = calls_pp.clone();
Box::pin(async move {
calls.fetch_add(1, Ordering::SeqCst);
})
}));
let sub_tx = host.subagent_sender();
let shutdown = host.world_mut().shutdown_handle();
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move {
host.serve(op_rx).await;
});
let (ctx, crx) = oneshot::channel();
sub_tx
.send(SubAgentOp::Check {
run_id: "parent".to_string(),
reply: ctx,
})
.unwrap();
let _ = crx.await.unwrap();
let (stx, srx) = oneshot::channel();
sub_tx
.send(SubAgentOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "child".to_string(),
..Default::default()
}),
parent_run_id: "parent".to_string(),
max_depth: 3,
reply: stx,
})
.unwrap();
assert_eq!(srx.await.unwrap(), Ok("child".to_string()));
shutdown.notify_one();
drop(op_tx);
handle.await.unwrap();
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"only the Spawn preprocessed"
);
}
#[tokio::test]
async fn serve_spawns_without_a_preprocessor() {
let mut host = host_with(vec![]);
host.set_spawner(Box::new(|world, args| {
Ok(world.spawn_agent((agent_state(&args.run_id),)))
}));
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move {
host.serve(op_rx).await;
});
let (tx, rx) = oneshot::channel();
op_tx
.send(ControlOp::Spawn {
args: Box::new(SpawnArgs {
run_id: "np".to_string(),
..Default::default()
}),
reply: tx,
})
.unwrap();
let result = rx.await.unwrap();
drop(op_tx);
handle.await.unwrap();
assert_eq!(result, Ok("np".to_string()));
}
#[tokio::test]
async fn shutdown_op_stops_the_serve_loop() {
let mut host = host_with(vec![]);
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move { host.serve(op_rx).await });
let (tx, rx) = oneshot::channel();
op_tx.send(ControlOp::Shutdown { reply: tx }).unwrap();
assert!(rx.await.unwrap());
handle.await.unwrap();
}
#[tokio::test]
async fn flush_and_stop_delegates_to_the_world() {
let mut host = host_with(vec![]);
host.flush_and_stop().await;
host.flush_and_stop().await; }
#[tokio::test]
async fn serve_loop_services_subagent_ops_via_the_sender() {
let mut host = host_with(vec![]);
spawn(&mut host, "run-a", "run-a");
let sub_tx = host.subagent_sender();
let (op_tx, op_rx) = mpsc::unbounded_channel();
let handle = tokio::spawn(async move { host.serve(op_rx).await });
let (tx, rx) = oneshot::channel();
sub_tx
.send(SubAgentOp::Check {
run_id: "run-a".to_string(),
reply: tx,
})
.unwrap();
assert!(rx.await.unwrap().is_some());
let (stx, srx) = oneshot::channel();
op_tx.send(ControlOp::Shutdown { reply: stx }).unwrap();
assert!(srx.await.unwrap());
handle.await.unwrap();
}
#[test]
fn status_str_covers_all_variants() {
assert_eq!(status_str(&AgentStatus::Idle), "idle");
assert_eq!(status_str(&AgentStatus::Active), "active");
assert_eq!(status_str(&AgentStatus::Waiting), "waiting");
assert_eq!(status_str(&AgentStatus::Complete), "complete");
assert_eq!(
status_str(&AgentStatus::Error {
message: "x".to_string()
}),
"error"
);
assert_eq!(status_str(&AgentStatus::Cancelled), "cancelled");
}
#[tokio::test]
async fn emit_events_broadcasts_agent_changes() {
let mut host = host_with(vec![text("done")]);
let mut rx = host.subscribe();
let entity = spawn(&mut host, "run-a", "agent-a");
host.world_mut()
.world_mut()
.entity_mut(entity)
.insert(RunMetadata {
run_id: "run-a".to_string(),
agent_name: "coder".to_string(),
agent_path: "/a".to_string(),
task: "t".to_string(),
model: None,
workdir: "/w".to_string(),
num_stages: 1,
started_at: 0,
parent_run_id: None,
metadata: std::collections::HashMap::new(),
callback_url: None,
callback_secret: None,
title: None,
});
host.emit_events();
let first: Vec<WorldEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
first
.iter()
.any(|e| matches!(e, WorldEvent::Spawned { .. }))
);
assert!(first.iter().any(|e| matches!(e, WorldEvent::Status { .. })));
assert!(first.iter().any(|e| matches!(e, WorldEvent::Tokens { .. })));
assert!(
first
.iter()
.any(|e| matches!(e, WorldEvent::Context { .. }))
);
host.emit_events();
assert!(rx.try_recv().is_err());
host.world_mut().run_until_idle(20).await;
host.emit_events();
let done: Vec<WorldEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
done.iter()
.any(|e| matches!(e, WorldEvent::Completed { .. }))
);
host.emit_events();
assert!(
std::iter::from_fn(|| rx.try_recv().ok())
.collect::<Vec<_>>()
.is_empty()
);
}
#[tokio::test]
async fn emit_events_unloads_terminal_agents_when_safe() {
let mut host = host_with(vec![]);
let root = {
let mut s = agent_state("root");
s.status = AgentStatus::Complete;
host.world.world_mut().spawn(s).id()
};
host.register("root", root);
host.emit_events();
assert!(
host.live_entity("root").is_some(),
"not reaped on the first terminal pass (event must go out first)"
);
host.emit_events();
assert!(host.live_entity("root").is_none(), "reaped after emit");
assert!(
host.world.world().get::<AgentState>(root).is_none(),
"entity despawned"
);
let parent = host.world.world_mut().spawn(agent_state("parent")).id();
host.register("parent", parent);
let child = {
let mut s = agent_state("child");
s.status = AgentStatus::Complete;
host.world
.world_mut()
.spawn((
s,
ParentRef {
parent_entity: parent,
parent_agent_id: "parent".to_string(),
depth: 1,
},
))
.id()
};
host.register("child", child);
host.emit_events();
host.emit_events();
assert!(
host.live_entity("child").is_some(),
"not reaped while its parent is live"
);
host.world
.world_mut()
.get_mut::<AgentState>(parent)
.unwrap()
.status = AgentStatus::Complete;
host.emit_events();
host.emit_events();
assert!(
host.live_entity("child").is_none(),
"reaped once its parent is terminal"
);
let ghost = host.world.world_mut().spawn_empty().id();
host.world.world_mut().despawn(ghost);
let orphan = {
let mut s = agent_state("orphan");
s.status = AgentStatus::Complete;
host.world
.world_mut()
.spawn((
s,
ParentRef {
parent_entity: ghost,
parent_agent_id: "gone".to_string(),
depth: 1,
},
))
.id()
};
host.register("orphan", orphan);
host.emit_events();
host.emit_events();
assert!(
host.live_entity("orphan").is_none(),
"reaped: parent entity despawned"
);
}
#[tokio::test]
async fn emit_events_does_not_reap_non_terminal_agents() {
let mut host = host_with(vec![]);
let active = host.world.world_mut().spawn(agent_state("active")).id();
host.register("active", active);
host.emit_events();
host.emit_events();
assert!(host.live_entity("active").is_some());
}
#[tokio::test]
async fn reaper_runs_once_per_agent_before_despawn() {
use std::sync::atomic::{AtomicUsize, Ordering};
let mut host = host_with(vec![]);
static SEEN_LIVE: AtomicUsize = AtomicUsize::new(0);
SEEN_LIVE.store(0, Ordering::SeqCst);
host.set_reaper(Box::new(|world, entity| {
let live = world.world().get::<AgentState>(entity).is_some();
SEEN_LIVE.fetch_add(live as usize, Ordering::SeqCst);
}));
let root = {
let mut s = agent_state("root");
s.status = AgentStatus::Complete;
host.world.world_mut().spawn(s).id()
};
host.register("root", root);
host.emit_events(); assert_eq!(SEEN_LIVE.load(Ordering::SeqCst), 0);
host.emit_events(); assert!(host.live_entity("root").is_none(), "reaped after emit");
assert_eq!(
SEEN_LIVE.load(Ordering::SeqCst),
1,
"reaper ran exactly once, while the entity was still live"
);
}
fn register_waiting(host: &mut WorldHost, run_id: &str) -> Entity {
let mut s = agent_state(run_id);
s.status = AgentStatus::Waiting;
let e = host.world.world_mut().spawn(s).id();
host.register(run_id, e);
e
}
#[tokio::test]
async fn emit_events_never_unloads_waiting_agents() {
use crate::components::AwaitingInteraction;
let mut host = host_with(vec![]);
let asking = register_waiting(&mut host, "asking");
host.world
.world_mut()
.entity_mut(asking)
.insert(AwaitingInteraction);
let gated = register_waiting(&mut host, "gated");
host.world
.world_mut()
.entity_mut(gated)
.insert(WaitingForChildren);
register_waiting(&mut host, "parked");
for _ in 0..5 {
host.emit_events();
}
for run_id in ["asking", "gated", "parked"] {
assert!(
host.live_entity(run_id).is_some(),
"a Waiting agent was unloaded and can no longer be resumed"
);
}
}
#[tokio::test]
async fn resolve_or_reload_pages_in_and_registers() {
let mut host = host_with(vec![]);
assert!(host.resolve_or_reload("ghost").is_none());
host.set_reloader(Box::new(|_world, _run_id| None));
assert!(host.resolve_or_reload("gone").is_none());
assert!(
host.live_entity("gone").is_none(),
"a declined reload registers nothing"
);
host.set_reloader(Box::new(|world, run_id| {
Some(world.spawn_agent((agent_state(run_id),)))
}));
let paged = host.resolve_or_reload("paged").expect("reloaded");
assert_eq!(
host.live_entity("paged"),
Some(paged),
"registered after reload"
);
assert_eq!(host.resolve_or_reload("paged"), Some(paged));
}
#[tokio::test]
async fn cancel_pages_in_an_unloaded_run() {
let mut host = host_with(vec![]);
host.set_reloader(paging_reloader());
let cancelled = ask(&mut host, |reply| ControlOp::Cancel {
run_id: "unloaded".to_string(),
reply,
})
.await;
assert!(cancelled, "reloaded then cancelled");
assert_eq!(
host.world
.agent_status(host.live_entity("unloaded").unwrap()),
Some(AgentStatus::Cancelled)
);
}
#[tokio::test]
async fn emit_events_broadcasts_new_interactions_once() {
let mut host = host_with(vec![]);
let mut rx = host.subscribe();
let backend = host.interactions().backend_for("agent-a");
let asking = tokio::spawn(async move {
backend
.ask(leviath_core::interaction::InteractionRequest::free_text(
"q1", "p", "s", true,
))
.await
});
for _ in 0..8 {
tokio::task::yield_now().await;
}
host.emit_events();
let evs: Vec<WorldEvent> = std::iter::from_fn(|| rx.try_recv().ok()).collect();
assert!(
evs.iter()
.any(|e| matches!(e, WorldEvent::Interaction { .. }))
);
host.emit_events();
assert!(rx.try_recv().is_err());
assert!(
host.interactions()
.answer(leviath_core::interaction::InteractionResponse::text(
"q1", "ok"
))
);
let _ = asking.await;
}
#[tokio::test]
async fn event_sender_feeds_subscribers() {
let host = host_with(vec![]);
let mut rx = host.subscribe();
let event = WorldEvent::Completed {
run_id: "r".to_string(),
agent_id: "a".to_string(),
status: "complete".to_string(),
};
host.event_sender().send(event.clone()).unwrap();
assert_eq!(rx.try_recv().unwrap(), event);
}
#[tokio::test]
async fn emit_events_skips_despawned_agents() {
let mut host = host_with(vec![]);
let e = spawn(&mut host, "run-a", "agent-a");
host.world_mut().world_mut().despawn(e);
host.emit_events();
}
#[tokio::test]
async fn serve_returns_when_control_channel_closes() {
let mut host = host_with(vec![text("done")]);
let (op_tx, op_rx) = mpsc::unbounded_channel();
drop(op_tx); host.serve(op_rx).await; }
#[tokio::test]
async fn mock_helpers_are_exercised() {
let p = Script {
responses: Mutex::new(std::collections::VecDeque::new()),
};
assert_eq!(p.name(), "script");
assert_eq!(p.count_tokens("t", "m").await, 1);
assert_eq!(p.max_context_tokens("m"), 100_000);
let _ = p.capabilities("m");
let req = InferenceRequest {
system: vec![],
messages: vec![],
model: "m".to_string(),
max_tokens: 1,
temperature: 0.0,
tools: vec![],
extra: serde_json::Value::Null,
request_timeout_secs: None,
};
assert!(p.infer(req).await.is_err());
let exec = NoTools.exec_for(
Entity::from_raw_u32(1).expect("a small literal index is always a valid entity id"),
vec![leviath_providers::ToolCall {
id: "c".to_string(),
name: "n".to_string(),
arguments: serde_json::Value::Null,
thought_signature: None,
}],
);
assert_eq!(exec().await, vec![("c".to_string(), String::new())]);
}
#[tokio::test]
async fn list_skips_despawned_entity() {
let mut host = host_with(vec![]);
let e = spawn(&mut host, "run-a", "agent-a");
host.world_mut().world_mut().despawn(e);
let list = ask(&mut host, |reply| ControlOp::List { reply }).await;
assert!(list.is_empty()); let status = ask(&mut host, |reply| ControlOp::Status {
run_id: "run-a".to_string(),
reply,
})
.await;
assert_eq!(status, None);
}
}