use bevy_ecs::prelude::*;
use leviath_core::RegionKind;
use leviath_core::run_meta::{
ContextSnapshot, RegionEntrySnapshot, RegionSnapshot, RunMeta, RunStatus, StageRunStatus,
};
use crate::components::{AgentState, AgentStatus, ContextWindow};
#[derive(Component, Clone)]
pub struct RunMetadata {
pub run_id: String,
pub agent_name: String,
pub agent_path: String,
pub task: String,
pub model: Option<String>,
pub workdir: String,
pub num_stages: usize,
pub started_at: i64,
pub parent_run_id: Option<String>,
pub metadata: std::collections::HashMap<String, String>,
pub callback_url: Option<String>,
pub callback_secret: Option<String>,
pub title: Option<String>,
pub unattended: bool,
pub read_paths: Option<leviath_core::run_meta::ReadPathGrantCounts>,
}
#[derive(Component, Clone, Copy, Default, Debug, PartialEq, Eq)]
pub struct TokenTotals {
pub prompt_tokens: usize,
pub completion_tokens: usize,
pub cached_tokens: usize,
pub cache_write_tokens: usize,
pub tool_calls: usize,
}
#[derive(Component, Clone, Default, Debug, PartialEq)]
pub struct RunOutcomeFlags(pub leviath_core::run_meta::RunFlags);
impl RunOutcomeFlags {
pub fn for_blueprint(bp: &leviath_core::Blueprint) -> Self {
Self(leviath_core::run_meta::RunFlags {
no_output_tools: !bp.stages.iter().any(stage_can_modify),
..Default::default()
})
}
}
fn stage_can_modify(stage: &leviath_core::Stage) -> bool {
stage.available_tools.iter().any(|t| {
let canonical = leviath_tools::canonical_tool_name(t);
leviath_core::blueprint::MODIFYING_TOOLS.contains(&canonical)
|| stage
.transitions
.iter()
.flat_map(|edges| edges.values())
.filter_map(|edge| edge.gate.as_ref())
.any(|gate| {
gate.tools
.iter()
.any(|extra| leviath_tools::canonical_tool_name(extra) == canonical)
})
})
}
impl TokenTotals {
pub fn add_usage(&mut self, usage: &leviath_providers::TokenUsage) {
self.prompt_tokens += usage.prompt_tokens;
self.completion_tokens += usage.completion_tokens;
self.cached_tokens += usage.cached_tokens;
self.cache_write_tokens += usage.cache_write_tokens;
}
}
pub fn is_empty_output(status: &AgentStatus, flags: &leviath_core::run_meta::RunFlags) -> bool {
matches!(
run_status_from(status),
RunStatus::Complete | RunStatus::Error | RunStatus::Cancelled
) && flags.modified_file_count == 0
&& !flags.no_output_tools
}
pub fn run_status_from(status: &AgentStatus) -> RunStatus {
match status {
AgentStatus::Idle | AgentStatus::Active => RunStatus::Running,
AgentStatus::Paused => RunStatus::Paused,
AgentStatus::Waiting => RunStatus::WaitingInput,
AgentStatus::Complete => RunStatus::Complete,
AgentStatus::Error { .. } => RunStatus::Error,
AgentStatus::Cancelled => RunStatus::Cancelled,
}
}
pub fn stage_status_from(status: &AgentStatus) -> StageRunStatus {
match status {
AgentStatus::Idle | AgentStatus::Active | AgentStatus::Paused => StageRunStatus::Active,
AgentStatus::Waiting => StageRunStatus::WaitingInput,
AgentStatus::Complete => StageRunStatus::Complete,
AgentStatus::Error { .. } | AgentStatus::Cancelled => StageRunStatus::Error,
}
}
fn region_kind_str(kind: &RegionKind) -> &'static str {
match kind {
RegionKind::Pinned => "pinned",
RegionKind::Temporary => "temporary",
RegionKind::Clearable => "clearable",
RegionKind::SlidingWindow { .. } => "sliding",
RegionKind::Compacting { .. } => "compacting",
RegionKind::CompactHistory { .. } => "history",
RegionKind::HashMap { .. } => "hashmap",
RegionKind::Custom { .. } => "custom",
}
}
pub fn build_context_snapshot(window: &ContextWindow, stage_name: &str) -> ContextSnapshot {
let regions = window
.regions
.iter()
.map(|r| RegionSnapshot {
name: r.name.clone(),
kind: region_kind_str(&r.kind).to_string(),
current_tokens: r.current_tokens,
max_tokens: r.max_tokens,
entries: r
.content
.iter()
.enumerate()
.map(|(i, e)| RegionEntrySnapshot {
content: e.content.clone(),
tokens: e.tokens,
kind: e.kind.clone(),
metadata: e.metadata.clone(),
key: e.key.clone(),
taint: r
.taint
.as_ref()
.and_then(|t| t.entry_taint(i))
.unwrap_or_default(),
})
.collect(),
})
.collect();
ContextSnapshot {
stage_name: stage_name.to_string(),
total_tokens: window.current_tokens,
max_tokens: window.max_tokens,
regions,
}
}
#[allow(clippy::too_many_arguments)]
pub fn build_run_meta(
md: &RunMetadata,
state: &AgentState,
totals: &TokenTotals,
flags: &RunOutcomeFlags,
stage_index: usize,
now_secs: i64,
last_progress_at: Option<i64>,
depth: usize,
max_child_depth: usize,
) -> RunMeta {
let status = run_status_from(&state.status);
let mut flags = flags.0.clone();
flags.empty_output = is_empty_output(&state.status, &flags);
RunMeta {
run_id: md.run_id.clone(),
agent_name: md.agent_name.clone(),
agent_path: md.agent_path.clone(),
task: md.task.clone(),
model: md.model.clone(),
pid: 0, status,
current_stage: state.current_stage.clone(),
stage_index,
num_stages: md.num_stages,
iteration: state.iteration,
prompt_tokens: totals.prompt_tokens,
completion_tokens: totals.completion_tokens,
cached_tokens: totals.cached_tokens,
cache_write_tokens: totals.cache_write_tokens,
tool_calls: totals.tool_calls,
workdir: md.workdir.clone(),
started_at: md.started_at,
updated_at: now_secs,
last_progress_at,
error: match &state.status {
AgentStatus::Error { message } => Some(message.clone()),
_ => None,
},
title: md.title.clone(),
metadata: md.metadata.clone(),
callback_url: md.callback_url.clone(),
callback_secret: md.callback_secret.clone(),
parent_run_id: md.parent_run_id.clone(),
children: state.spawned_children_ids.clone(),
depth,
max_child_depth,
flags,
yolo: md.unattended,
read_paths: md.read_paths,
}
}
#[cfg(test)]
mod tests {
use super::*;
use leviath_core::Region;
use leviath_providers::TokenUsage;
fn state(status: AgentStatus) -> AgentState {
AgentState {
agent_id: "a".to_string(),
current_stage: "plan".to_string(),
iteration: 4,
status,
spawned_children_ids: vec![],
pending_wait: None,
accepts_messages: true,
}
}
fn metadata() -> RunMetadata {
RunMetadata {
run_id: "run-1".to_string(),
agent_name: "coder".to_string(),
agent_path: "/agents/coder".to_string(),
task: "do it".to_string(),
model: Some("anthropic/claude".to_string()),
workdir: "/work".to_string(),
num_stages: 3,
started_at: 1000,
parent_run_id: Some("parent".to_string()),
metadata: std::collections::HashMap::from([("k".to_string(), "v".to_string())]),
callback_url: Some("http://cb".to_string()),
callback_secret: Some("sekret".to_string()),
title: Some("Do It".to_string()),
unattended: false,
read_paths: None,
}
}
fn stage_with(tools: &[&str], gate_tools: Option<&[&str]>) -> leviath_core::Stage {
let mut stage = leviath_core::Stage::new(
"s".to_string(),
leviath_core::blueprint::ModelConfig::new("p".to_string(), "m".to_string()),
);
stage.available_tools = tools.iter().map(|t| (*t).to_string()).collect();
stage.transitions = gate_tools.map(|extra| {
let gate = (!extra.is_empty()).then(|| leviath_core::blueprint::TransitionGate {
require_modifications: true,
tools: extra.iter().map(|t| (*t).to_string()).collect(),
..Default::default()
});
std::collections::HashMap::from([(
"next".to_string(),
leviath_core::blueprint::TransitionEdge {
target: "next".to_string(),
condition: leviath_core::blueprint::TransitionCondition::Always,
hint: None,
transform: leviath_core::blueprint::EdgeTransform::Direct,
gate,
stuck: None,
},
)])
});
stage
}
fn blueprint_of(stages: Vec<leviath_core::Stage>) -> leviath_core::Blueprint {
leviath_core::Blueprint::new(
"bp".to_string(),
"d".to_string(),
stages,
leviath_core::ContextLayout::new(vec![], 1000),
)
}
fn no_output_tools(stages: Vec<leviath_core::Stage>) -> bool {
RunOutcomeFlags::for_blueprint(&blueprint_of(stages))
.0
.no_output_tools
}
#[test]
fn for_blueprint_asks_whether_any_stage_could_have_written() {
assert!(no_output_tools(vec![]));
assert!(no_output_tools(vec![stage_with(
&["read_file", "spawn_agent", "context_write"],
None
)]));
assert!(no_output_tools(vec![stage_with(&["bash"], None)]));
assert!(!no_output_tools(vec![stage_with(&["write_file"], None)]));
assert!(!no_output_tools(vec![stage_with(&["edit_file"], None)]));
assert!(!no_output_tools(vec![
stage_with(&["read_file"], None),
stage_with(&["write_file"], None),
]));
}
#[test]
fn for_blueprint_honors_a_gate_declaring_its_own_write_tool() {
assert!(!no_output_tools(vec![stage_with(
&["mcp__fs__put"],
Some(&["mcp__fs__put"])
)]));
assert!(no_output_tools(vec![stage_with(
&["read_file"],
Some(&["mcp__fs__put"])
)]));
assert!(no_output_tools(vec![stage_with(&["read_file"], Some(&[]))]));
assert!(no_output_tools(vec![stage_with(
&["mcp__fs__put"],
Some(&["mcp__other__put"])
)]));
}
#[test]
fn is_empty_output_needs_a_stopped_run_that_could_have_written() {
let nothing = leviath_core::run_meta::RunFlags::default();
assert!(!is_empty_output(&AgentStatus::Active, ¬hing));
assert!(!is_empty_output(&AgentStatus::Idle, ¬hing));
assert!(!is_empty_output(&AgentStatus::Paused, ¬hing));
assert!(!is_empty_output(&AgentStatus::Waiting, ¬hing));
for status in [
AgentStatus::Complete,
AgentStatus::Cancelled,
AgentStatus::Error {
message: "x".to_string(),
},
] {
assert!(is_empty_output(&status, ¬hing));
}
let mut wrote = leviath_core::run_meta::RunFlags::default();
wrote.record_modification("src/a.rs");
assert!(!is_empty_output(&AgentStatus::Complete, &wrote));
let incapable = leviath_core::run_meta::RunFlags {
no_output_tools: true,
..Default::default()
};
assert!(!is_empty_output(&AgentStatus::Complete, &incapable));
}
#[test]
fn status_mapping_covers_all_variants() {
assert_eq!(run_status_from(&AgentStatus::Idle), RunStatus::Running);
assert_eq!(run_status_from(&AgentStatus::Active), RunStatus::Running);
assert_eq!(run_status_from(&AgentStatus::Paused), RunStatus::Paused);
assert_eq!(
run_status_from(&AgentStatus::Waiting),
RunStatus::WaitingInput
);
assert_eq!(run_status_from(&AgentStatus::Complete), RunStatus::Complete);
assert_eq!(
run_status_from(&AgentStatus::Error {
message: "x".to_string()
}),
RunStatus::Error
);
assert_eq!(
run_status_from(&AgentStatus::Cancelled),
RunStatus::Cancelled
);
}
#[test]
fn stage_status_mapping_covers_all_variants() {
use leviath_core::run_meta::StageRunStatus;
assert_eq!(
stage_status_from(&AgentStatus::Idle),
StageRunStatus::Active
);
assert_eq!(
stage_status_from(&AgentStatus::Active),
StageRunStatus::Active
);
assert_eq!(
stage_status_from(&AgentStatus::Paused),
StageRunStatus::Active
);
assert_eq!(
stage_status_from(&AgentStatus::Waiting),
StageRunStatus::WaitingInput
);
assert_eq!(
stage_status_from(&AgentStatus::Complete),
StageRunStatus::Complete
);
assert_eq!(
stage_status_from(&AgentStatus::Error {
message: "x".to_string()
}),
StageRunStatus::Error
);
assert_eq!(
stage_status_from(&AgentStatus::Cancelled),
StageRunStatus::Error
);
}
#[test]
fn token_totals_accumulate() {
let mut t = TokenTotals::default();
t.add_usage(&TokenUsage {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15,
cached_tokens: 2,
cache_write_tokens: 1,
});
t.add_usage(&TokenUsage {
prompt_tokens: 3,
completion_tokens: 4,
total_tokens: 7,
cached_tokens: 0,
cache_write_tokens: 0,
});
t.tool_calls = 6;
assert_eq!(t.prompt_tokens, 13);
assert_eq!(t.completion_tokens, 9);
assert_eq!(t.cached_tokens, 2);
assert_eq!(t.cache_write_tokens, 1);
}
#[test]
fn build_run_meta_fills_dynamic_and_static_fields() {
let md = metadata();
let totals = TokenTotals {
prompt_tokens: 100,
completion_tokens: 50,
cached_tokens: 10,
cache_write_tokens: 5,
tool_calls: 7,
};
let mut st = state(AgentStatus::Active);
st.spawned_children_ids = vec!["child-a".to_string(), "child-b".to_string()];
let meta = build_run_meta(
&md,
&st,
&totals,
&RunOutcomeFlags::default(),
1,
2000,
Some(1900),
1,
4,
);
assert_eq!(meta.run_id, "run-1");
assert_eq!(meta.status, RunStatus::Running);
assert_eq!(meta.current_stage, "plan");
assert_eq!(meta.stage_index, 1);
assert_eq!(meta.iteration, 4);
assert_eq!(meta.prompt_tokens, 100);
assert_eq!(meta.tool_calls, 7);
assert_eq!(meta.updated_at, 2000);
assert_eq!(meta.last_progress_at, Some(1900));
assert_eq!(meta.parent_run_id.as_deref(), Some("parent"));
assert_eq!(meta.callback_url.as_deref(), Some("http://cb"));
assert_eq!(meta.callback_secret.as_deref(), Some("sekret"));
assert!(meta.error.is_none());
assert_eq!(
meta.children,
vec!["child-a".to_string(), "child-b".to_string()]
);
assert_eq!(meta.depth, 1);
assert_eq!(meta.max_child_depth, 4);
assert!(!meta.yolo);
}
#[test]
fn build_run_meta_records_an_unattended_run() {
let mut md = metadata();
md.unattended = true;
let meta = build_run_meta(
&md,
&state(AgentStatus::Active),
&TokenTotals::default(),
&RunOutcomeFlags::default(),
1,
2000,
None,
1,
4,
);
assert!(meta.yolo);
}
#[test]
fn build_run_meta_flags_empty_output_only_once_the_run_has_stopped() {
let mut flags = RunOutcomeFlags::default();
flags.0.gates_forced = 2;
let running = build_run_meta(
&metadata(),
&state(AgentStatus::Active),
&TokenTotals::default(),
&flags,
0,
1000,
None,
0,
0,
);
assert!(!running.flags.empty_output);
assert_eq!(running.flags.gates_forced, 2);
for status in [
AgentStatus::Complete,
AgentStatus::Cancelled,
AgentStatus::Error {
message: "x".to_string(),
},
] {
let meta = build_run_meta(
&metadata(),
&state(status),
&TokenTotals::default(),
&flags,
0,
1000,
None,
0,
0,
);
assert!(meta.flags.empty_output);
}
let mut wrote = RunOutcomeFlags::default();
wrote.0.record_modification("src/a.rs");
let meta = build_run_meta(
&metadata(),
&state(AgentStatus::Complete),
&TokenTotals::default(),
&wrote,
0,
1000,
None,
0,
0,
);
assert!(!meta.flags.empty_output);
assert_eq!(meta.flags.modified_files, vec!["src/a.rs".to_string()]);
let mut incapable = RunOutcomeFlags::default();
incapable.0.no_output_tools = true;
let meta = build_run_meta(
&metadata(),
&state(AgentStatus::Complete),
&TokenTotals::default(),
&incapable,
0,
1000,
None,
0,
0,
);
assert!(!meta.flags.empty_output);
assert!(meta.flags.no_output_tools);
}
#[test]
fn build_run_meta_carries_error_message() {
let meta = build_run_meta(
&metadata(),
&state(AgentStatus::Error {
message: "boom".to_string(),
}),
&TokenTotals::default(),
&RunOutcomeFlags::default(),
2,
3000,
None,
0,
0,
);
assert_eq!(meta.status, RunStatus::Error);
assert_eq!(meta.error.as_deref(), Some("boom"));
}
#[test]
fn context_snapshot_captures_all_region_kinds() {
let mut w = ContextWindow::new(1000);
w.add_region(Region::new("pin".to_string(), RegionKind::Pinned, 100));
w.add_region(Region::new("tmp".to_string(), RegionKind::Temporary, 100));
w.add_region(Region::new("clr".to_string(), RegionKind::Clearable, 100));
w.add_region(Region::new(
"slide".to_string(),
RegionKind::SlidingWindow {
max_items: 5,
eviction_strategy: leviath_core::EvictionStrategy::PerItem,
},
100,
));
w.add_region(Region::new(
"comp".to_string(),
RegionKind::Compacting {
threshold_tokens: 5,
},
100,
));
w.add_region(Region::new(
"hist".to_string(),
RegionKind::CompactHistory {
source_region: "comp".to_string(),
},
100,
));
w.add_region(Region::new(
"map".to_string(),
RegionKind::HashMap { max_entries: None },
100,
));
w.add_region(Region::new(
"brain".to_string(),
RegionKind::Custom {
script: "b.rhai".to_string(),
persistent: false,
},
100,
));
let _ = w.add_to_region("pin", "hello".to_string(), 3);
w.current_tokens = w.calculate_tokens();
let snap = build_context_snapshot(&w, "plan");
assert_eq!(snap.stage_name, "plan");
let kinds: Vec<&str> = snap.regions.iter().map(|r| r.kind.as_str()).collect();
assert_eq!(
kinds,
vec![
"pinned",
"temporary",
"clearable",
"sliding",
"compacting",
"history",
"hashmap",
"custom"
]
);
let pin = snap.regions.iter().find(|r| r.name == "pin").unwrap();
assert_eq!(pin.entries.len(), 1);
assert_eq!(pin.entries[0].content, "hello");
}
}