use super::*;
pub(crate) const PERSIST_HEARTBEAT_SECS: i64 = 30;
#[derive(Component, Default)]
pub struct PersistWatermark {
last: Option<(usize, usize, leviath_core::run_meta::RunStatus)>,
last_written_at: Option<i64>,
}
#[derive(Resource)]
pub struct PersistenceStage(pub UnboundedSender<PersistJob>);
#[allow(clippy::type_complexity)]
pub fn reflect_interaction_status(
hub: Option<Res<InteractionHub>>,
mut agents: Query<
(Entity, &mut AgentState, Option<&AwaitingInteraction>),
Without<FanOutWaiting>,
>,
mut commands: Commands,
) {
crate::tick_scope::clear();
let Some(hub) = hub else { return };
let pending: std::collections::HashSet<String> =
hub.pending().into_iter().map(|(id, _)| id).collect();
for (entity, mut state, marked) in agents.iter_mut() {
crate::tick_scope::enter(entity);
match (pending.contains(&state.agent_id), marked.is_some()) {
(true, false) => {
if state.status == AgentStatus::Active {
state.status = AgentStatus::Waiting;
commands.entity(entity).insert(AwaitingInteraction);
}
}
(false, true) => {
commands.entity(entity).remove::<AwaitingInteraction>();
if state.status == AgentStatus::Waiting {
state.status = AgentStatus::Active;
}
}
_ => {}
}
}
}
pub(crate) fn reconcile_stage_ledger(
ledger: &mut StageLedger,
cursor_index: usize,
status: &AgentStatus,
now: i64,
) {
use leviath_core::run_meta::StageRunStatus;
let active = crate::persistence::stage_status_from(status);
for rec in ledger.0.iter_mut() {
match rec.index.cmp(&cursor_index) {
std::cmp::Ordering::Less => {
if rec.started_at.is_none() {
rec.started_at = Some(now);
}
rec.status = StageRunStatus::Complete;
if rec.ended_at.is_none() {
rec.ended_at = Some(now);
}
}
std::cmp::Ordering::Equal => {
if rec.started_at.is_none() {
rec.started_at = Some(now);
}
if active == StageRunStatus::Complete && rec.ended_at.is_none() {
rec.ended_at = Some(now);
}
rec.status = active.clone();
}
std::cmp::Ordering::Greater => {
rec.status = StageRunStatus::Pending;
}
}
}
}
#[allow(clippy::type_complexity)]
pub fn dispatch_persistence(
mut agents: Query<(
Entity,
&RunMetadata,
&AgentState,
&ContextWindow,
&StageCursor,
&TokenTotals,
&mut PersistWatermark,
Option<&mut StageLedger>,
Option<&mut StageIoBuffer>,
Option<&crate::taint::TaintGate>,
Option<&crate::components::ParentRef>,
Option<&crate::components::SubAgentChildren>,
Option<&crate::fanout::FanOutWaiting>,
(
Option<&crate::interaction_points::AwaitingInteractionPoint>,
Option<&crate::interaction_points::InteractionPointCursor>,
Option<&crate::interaction_points::InteractionPointRounds>,
Option<&crate::persistence::RunOutcomeFlags>,
),
)>,
stage: Res<PersistenceStage>,
hub: Option<Res<InteractionHub>>,
sink: Option<Res<crate::host::WorldEventSink>>,
) {
crate::tick_scope::clear();
for (
entity,
md,
state,
window,
cursor,
totals,
mut watermark,
mut ledger,
buffer,
taint_gate,
parent_ref,
children,
fan_out_waiting,
(awaiting_point, ip_cursor, ip_rounds, outcome_flags),
) in agents.iter_mut()
{
crate::tick_scope::enter(entity);
let now = chrono::Utc::now().timestamp();
if let Some(ledger) = ledger.as_deref_mut() {
reconcile_stage_ledger(ledger, cursor.index, &state.status, now);
}
let (output_appends, log_appends) = match buffer {
Some(mut buf) => (
std::mem::take(&mut buf.output),
std::mem::take(&mut buf.logs),
),
None => (Vec::new(), Vec::new()),
};
let has_appends = !output_appends.is_empty() || !log_appends.is_empty();
let status = crate::persistence::run_status_from(&state.status);
let current = (state.iteration, cursor.index, status);
let watermark_changed = watermark.last.as_ref() != Some(¤t);
let due_for_heartbeat = watermark
.last_written_at
.is_none_or(|at| now.saturating_sub(at) >= PERSIST_HEARTBEAT_SECS);
if !watermark_changed && !has_appends && !due_for_heartbeat {
continue; }
if watermark_changed {
watermark.last = Some(current);
}
watermark.last_written_at = Some(now);
if let Some(sink) = &sink {
for (_idx, line) in output_appends.iter().chain(log_appends.iter()) {
let _ = sink.0.send(crate::host::WorldEvent::Log {
run_id: md.run_id.clone(),
agent_id: state.agent_id.clone(),
line: line.clone(),
});
}
}
let depth = parent_ref.map(|p| p.depth).unwrap_or(0);
let max_child_depth = children.map(|c| c.max_child_depth).unwrap_or(0);
let flags = outcome_flags.cloned().unwrap_or_default();
let meta = build_run_meta(
md,
state,
totals,
&flags,
cursor.index,
now,
depth,
max_child_depth,
);
let context = build_context_snapshot(window, &state.current_stage);
let stages = ledger.as_deref().map(|l| l.0.clone()).unwrap_or_default();
let taint_audit = taint_gate.filter(|g| !g.audit_log().is_empty()).map(|g| {
(
cursor.index,
serde_json::to_string_pretty(g.audit_log())
.expect("GateEvent slice always serializes"),
)
});
let fanout = fan_out_waiting
.map(|w| serde_json::to_string(&w.to_state()).expect("FanOutState always serializes"));
let interactions = awaiting_point.and_then(|_| {
let request = hub
.as_ref()?
.pending()
.into_iter()
.find(|(aid, req)| aid == &state.agent_id && req.id.contains("-point-"))?;
let ip_state = crate::interaction_points::InteractionPointState {
cursor: ip_cursor.map_or(0, |c| c.0),
round: ip_rounds.map_or(0, |r| r.0),
body: request.1.body.unwrap_or_default(),
};
Some(serde_json::to_string(&ip_state).expect("InteractionPointState always serializes"))
});
let _ = stage.0.send(PersistJob {
run_id: md.run_id.clone(),
meta,
context,
stages,
output_appends,
log_appends,
taint_audit,
fanout,
interactions,
});
}
}