use std::{
collections::{HashMap, HashSet},
path::PathBuf,
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use serde_json::json;
use tokio::sync::watch;
use crate::{
a2a::{A2aArtifact, Part},
child_launch::LaunchedChild,
delegation_plane::DelegationLaunch,
detached::DetachedRegistry,
};
pub(crate) const RESIDUE_ARTIFACT: &str = "residue";
pub(crate) const PHASE_QUEUED: &str = "queued";
pub(crate) const PHASE_TURN: &str = "turn";
pub(crate) const PHASE_INFERENCE: &str = "inference";
pub(crate) const PHASE_INPUT: &str = "input";
pub(crate) const PHASE_DELEGATION: &str = "delegation";
pub(crate) const PHASE_MEMBER_CALL: &str = "member_call";
pub(crate) const PHASE_HARNESS: &str = "harness";
pub(crate) const PHASE_PLAN: &str = "plan";
pub(crate) const PHASE_BACKGROUND_WAIT: &str = "background_wait";
pub(crate) const PHASE_SHELL: &str = "shell";
pub(crate) const PHASE_DONE_CHECK: &str = "done_check";
pub(crate) const CANCELED_STATUS_MESSAGE: &str = "task canceled";
pub(crate) const CANCELED_RESULT_TEXT: &str =
"canceled: the task was canceled before it completed, so it has no result";
#[derive(Debug, Clone)]
pub(crate) struct CancelSignal {
tx: Arc<watch::Sender<bool>>,
phase: Arc<Mutex<Option<&'static str>>>,
}
impl CancelSignal {
pub(crate) fn new() -> Self {
let (tx, _rx) = watch::channel(false);
Self {
tx: Arc::new(tx),
phase: Arc::new(Mutex::new(None)),
}
}
pub(crate) fn note_phase(&self, phase: &'static str) {
self.phase
.lock()
.expect("cancel phase mutex")
.get_or_insert(phase);
}
pub(crate) fn phase(&self) -> &'static str {
self.phase
.lock()
.expect("cancel phase mutex")
.unwrap_or(PHASE_TURN)
}
pub(crate) fn cancel(&self) {
self.tx.send_replace(true);
}
pub(crate) fn is_canceled(&self) -> bool {
*self.tx.borrow()
}
pub(crate) async fn canceled(&self) {
let mut rx = self.tx.subscribe();
loop {
if *rx.borrow_and_update() {
return;
}
if rx.changed().await.is_err() {
std::future::pending::<()>().await;
}
}
}
}
impl Default for CancelSignal {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub(crate) struct LiveDelegation {
pub(crate) task_id: String,
pub(crate) workdir: PathBuf,
pub(crate) capsule: String,
pub(crate) version: String,
pub(crate) child_session_id: String,
pub(crate) child_workdir: String,
pub(crate) child_records_dir: PathBuf,
pub(crate) started: Instant,
pub(crate) child: Option<LaunchedChild>,
pub(crate) arrived: bool,
}
impl LiveDelegation {
pub(crate) fn launched(task_id: String, notice: &DelegationLaunch) -> Self {
Self {
task_id,
workdir: PathBuf::from(¬ice.child_workdir),
capsule: notice.capsule.clone(),
version: notice.version.clone(),
child_session_id: notice.child_session_id.clone(),
child_workdir: notice.child_workdir.clone(),
child_records_dir: PathBuf::from(¬ice.child_records_dir),
started: Instant::now(),
child: None,
arrived: false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Arrival {
Delivered,
AlreadyDelivered,
NotOutstanding,
BeingEnded,
}
#[derive(Debug, Default)]
struct Delegations {
task_id: Option<String>,
entries: HashMap<String, LiveDelegation>,
delivered: HashSet<String>,
}
#[derive(Debug, Default)]
pub(crate) struct LiveDelegations {
inner: Mutex<Delegations>,
arrival: tokio::sync::Notify,
}
impl LiveDelegations {
pub(crate) fn new() -> Self {
Self::default()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Delegations> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
pub(crate) fn scope_task(self: &Arc<Self>, task_id: &str) -> DelegationScope {
let left = {
let mut delegations = self.lock();
delegations.task_id = Some(task_id.to_string());
std::mem::take(&mut delegations.entries)
};
end_together(left.into_values());
DelegationScope(Arc::clone(self))
}
pub(crate) fn task_id(&self) -> Option<String> {
self.lock().task_id.clone()
}
pub(crate) fn register(&self, delegation_id: String, delegation: LiveDelegation) -> bool {
let mut delegations = self.lock();
if delegations.task_id.as_deref() != Some(delegation.task_id.as_str()) {
return false;
}
delegations.entries.insert(delegation_id, delegation);
true
}
pub(crate) fn adopt(&self, delegation_id: &str, child: LaunchedChild) {
let orphan = {
let mut delegations = self.lock();
match delegations.entries.get_mut(delegation_id) {
Some(entry) => {
entry.child = Some(child);
None
}
None => Some(child),
}
};
drop(orphan);
}
pub(crate) fn discard(&self, delegation_id: &str) -> bool {
let entry = self.lock().entries.remove(delegation_id);
let removed = entry.is_some();
drop(entry);
removed
}
pub(crate) fn arrive(&self, delegation_id: &str) -> Arrival {
let mut delegations = self.lock();
if delegations.delivered.contains(delegation_id) {
return Arrival::AlreadyDelivered;
}
let Some(entry) = delegations.entries.get_mut(delegation_id) else {
return Arrival::NotOutstanding;
};
if entry.arrived {
return Arrival::AlreadyDelivered;
}
if entry
.child
.as_ref()
.is_some_and(LaunchedChild::is_being_ended)
{
return Arrival::BeingEnded;
}
entry.arrived = true;
drop(delegations);
self.arrival.notify_one();
Arrival::Delivered
}
pub(crate) fn counts(&self) -> (usize, usize) {
let delegations = self.lock();
let arrived = delegations
.entries
.values()
.filter(|entry| entry.arrived)
.count();
(delegations.entries.len() - arrived, arrived)
}
pub(crate) async fn wait_for_arrival(&self) {
self.arrival.notified().await;
}
pub(crate) fn take_arrived(&self) -> Vec<(String, LiveDelegation)> {
self.take_where(|entry| entry.arrived)
}
pub(crate) fn take_overdue(&self, bound: Duration) -> Vec<(String, LiveDelegation)> {
let now = Instant::now();
self.take_where(|entry| {
!entry.arrived
&& entry
.started
.checked_add(bound)
.is_some_and(|due| due <= now)
})
}
pub(crate) fn next_overdue(&self, bound: Duration) -> Option<Instant> {
self.lock()
.entries
.values()
.filter(|entry| !entry.arrived)
.filter_map(|entry| entry.started.checked_add(bound))
.min()
}
pub(crate) fn take_all(&self) -> Vec<(String, LiveDelegation)> {
self.take_where(|_| true)
}
fn take_where(&self, take: impl Fn(&LiveDelegation) -> bool) -> Vec<(String, LiveDelegation)> {
let mut delegations = self.lock();
let ids: Vec<String> = delegations
.entries
.iter()
.filter(|(_, entry)| take(entry))
.map(|(id, _)| id.clone())
.collect();
let mut taken: Vec<(String, LiveDelegation)> = ids
.into_iter()
.filter_map(|id| {
let entry = delegations.entries.remove(&id)?;
if entry.arrived {
delegations.delivered.insert(id.clone());
}
Some((id, entry))
})
.collect();
taken.sort_by(|a, b| a.0.cmp(&b.0));
taken
}
pub(crate) fn live(&self) -> Vec<ResidueItem> {
let mut items: Vec<ResidueItem> = self
.lock()
.entries
.iter()
.map(|(delegation_id, delegation)| ResidueItem::Delegation {
delegation_id: delegation_id.clone(),
capsule: delegation.capsule.clone(),
version: delegation.version.clone(),
child_session_id: delegation.child_session_id.clone(),
child_workdir: delegation.child_workdir.clone(),
})
.collect();
items.sort_by(|a, b| a.id().cmp(b.id()));
items
}
}
pub(crate) struct DelegationScope(Arc<LiveDelegations>);
impl Drop for DelegationScope {
fn drop(&mut self) {
let left = {
let mut delegations = self.0.lock();
delegations.task_id = None;
std::mem::take(&mut delegations.entries)
};
end_together(left.into_values());
}
}
pub(crate) fn begin_ending<'a>(delegations: impl IntoIterator<Item = &'a LiveDelegation>) {
for delegation in delegations {
if delegation.arrived {
continue;
}
if let Some(child) = &delegation.child {
child.begin_ending();
}
}
}
fn end_together(delegations: impl Iterator<Item = LiveDelegation>) {
let delegations: Vec<LiveDelegation> = delegations.collect();
begin_ending(&delegations);
drop(delegations);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ResidueItem {
DetachedShell {
work_id: String,
binary: String,
command: String,
started_at_ms: u64,
},
Delegation {
delegation_id: String,
capsule: String,
version: String,
child_session_id: String,
child_workdir: String,
},
}
impl ResidueItem {
pub(crate) fn id(&self) -> &str {
match self {
Self::DetachedShell { work_id, .. } => work_id,
Self::Delegation { delegation_id, .. } => delegation_id,
}
}
pub(crate) fn kind(&self) -> &'static str {
match self {
Self::DetachedShell { .. } => "detached_shell",
Self::Delegation { .. } => "delegation",
}
}
fn to_json(&self) -> serde_json::Value {
match self {
Self::DetachedShell {
work_id,
binary,
command,
started_at_ms,
} => json!({
"kind": self.kind(),
"work_id": work_id,
"binary": binary,
"command": command,
"started_at_ms": started_at_ms,
}),
Self::Delegation {
delegation_id,
capsule,
version,
child_session_id,
child_workdir,
} => json!({
"kind": self.kind(),
"delegation_id": delegation_id,
"capsule": capsule,
"version": version,
"child_session_id": child_session_id,
"child_workdir": child_workdir,
}),
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct Residue {
items: Vec<ResidueItem>,
}
impl Residue {
pub(crate) fn snapshot(
detached: Option<&Arc<DetachedRegistry>>,
delegations: &LiveDelegations,
) -> Self {
let mut items: Vec<ResidueItem> = detached
.map(|registry| {
let mut work: Vec<_> = registry.outstanding();
work.sort_by(|a, b| a.work_id.cmp(&b.work_id));
work.into_iter()
.map(|work| ResidueItem::DetachedShell {
work_id: work.work_id,
binary: work.binary,
command: work.command,
started_at_ms: work.started_at_ms,
})
.collect()
})
.unwrap_or_default();
items.extend(delegations.live());
Self { items }
}
pub(crate) fn detached_work_ids(&self) -> Vec<String> {
self.items
.iter()
.filter(|item| matches!(item, ResidueItem::DetachedShell { .. }))
.map(|item| item.id().to_string())
.collect()
}
pub(crate) fn delegation_ids(&self) -> Vec<String> {
self.items
.iter()
.filter(|item| matches!(item, ResidueItem::Delegation { .. }))
.map(|item| item.id().to_string())
.collect()
}
pub(crate) fn into_json_items(self) -> Vec<serde_json::Value> {
self.items.iter().map(ResidueItem::to_json).collect()
}
pub(crate) fn into_artifact(self) -> Option<A2aArtifact> {
if self.items.is_empty() {
return None;
}
Some(A2aArtifact::named(
RESIDUE_ARTIFACT,
self.items
.iter()
.map(|item| Part::data(item.to_json()))
.collect(),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn live_delegation(capsule: &str) -> LiveDelegation {
LiveDelegation {
task_id: "tsk_a".to_string(),
workdir: PathBuf::from("/tmp/child"),
capsule: capsule.to_string(),
version: "0.1.0".to_string(),
child_session_id: "ses_child".to_string(),
child_workdir: "/tmp/child".to_string(),
child_records_dir: std::path::PathBuf::from(
"/home/test/.murmur/records/work/ses_child",
),
started: Instant::now(),
child: None,
arrived: false,
}
}
fn scoped() -> (Arc<LiveDelegations>, DelegationScope) {
let live = Arc::new(LiveDelegations::new());
let scope = live.scope_task("tsk_a");
(live, scope)
}
#[test]
fn every_phase_has_its_own_word() {
let phases = [
PHASE_QUEUED,
PHASE_TURN,
PHASE_INFERENCE,
PHASE_INPUT,
PHASE_DELEGATION,
PHASE_HARNESS,
PHASE_MEMBER_CALL,
PHASE_PLAN,
PHASE_BACKGROUND_WAIT,
];
assert_eq!(phases[5], "harness");
assert_eq!(PHASE_PLAN, "plan");
assert_eq!(PHASE_BACKGROUND_WAIT, "background_wait");
assert_eq!(
phases
.iter()
.collect::<std::collections::HashSet<_>>()
.len(),
phases.len()
);
assert_eq!(CANCELED_STATUS_MESSAGE, "task canceled");
}
#[test]
fn a_registered_delegation_is_live_until_it_is_taken() {
let (live, _scope) = scoped();
assert!(live.live().is_empty());
assert!(live.register("dlg_one".to_string(), live_delegation("worker")));
let items = live.live();
assert_eq!(items.len(), 1);
assert_eq!(items[0].id(), "dlg_one");
assert_eq!(items[0].kind(), "delegation");
let taken = live.take_all();
assert_eq!(taken.len(), 1);
assert_eq!(taken[0].1.capsule, "worker");
assert!(live.live().is_empty());
assert!(live.take_all().is_empty());
}
#[test]
fn discarding_reports_whether_it_removed_the_row() {
let (live, _scope) = scoped();
live.register("dlg_one".to_string(), live_delegation("worker"));
assert!(live.discard("dlg_one"));
assert!(live.live().is_empty());
assert!(!live.discard("dlg_one"), "a second discard removes nothing");
assert!(!live.discard("dlg_missing"));
live.register("dlg_two".to_string(), live_delegation("worker"));
assert_eq!(live.take_all().len(), 1);
assert!(
!live.discard("dlg_two"),
"a row the task took to account for is not the caller's to close"
);
}
#[test]
fn live_delegations_are_listed_in_id_order() {
let (live, _scope) = scoped();
live.register("dlg_b".to_string(), live_delegation("second"));
live.register("dlg_a".to_string(), live_delegation("first"));
let listed = live.live();
let ids: Vec<&str> = listed.iter().map(ResidueItem::id).collect();
assert_eq!(ids, vec!["dlg_a", "dlg_b"]);
}
#[test]
fn a_delegation_is_kept_only_while_its_task_is_in_scope() {
let live = Arc::new(LiveDelegations::new());
assert!(live.task_id().is_none());
assert!(
!live.register("dlg_early".to_string(), live_delegation("worker")),
"no task is in scope"
);
let scope = live.scope_task("tsk_b");
assert_eq!(live.task_id().as_deref(), Some("tsk_b"));
assert!(
!live.register("dlg_other".to_string(), live_delegation("worker")),
"started for tsk_a, which is not the task in scope"
);
let mut mine = live_delegation("worker");
mine.task_id = "tsk_b".to_string();
assert!(live.register("dlg_mine".to_string(), mine));
assert_eq!(live.counts(), (1, 0));
drop(scope);
assert!(live.task_id().is_none());
assert!(live.live().is_empty(), "the scope's end takes what is left");
assert_eq!(live.arrive("dlg_mine"), Arrival::NotOutstanding);
}
#[test]
fn the_set_delivers_each_outcome_once_and_refuses_one_nobody_waits_for() {
let (live, _scope) = scoped();
live.register("dlg_a".to_string(), live_delegation("worker"));
live.register("dlg_b".to_string(), live_delegation("worker"));
assert_eq!(live.arrive("dlg_a"), Arrival::Delivered);
assert_eq!(live.arrive("dlg_a"), Arrival::AlreadyDelivered);
assert_eq!(live.arrive("dlg_nobody"), Arrival::NotOutstanding);
assert_eq!(live.counts(), (1, 1));
let arrived = live.take_arrived();
assert_eq!(arrived.len(), 1);
assert_eq!(arrived[0].0, "dlg_a");
assert_eq!(live.counts(), (1, 0));
assert_eq!(live.arrive("dlg_a"), Arrival::AlreadyDelivered);
let left = live.take_all();
assert_eq!(left.len(), 1);
assert_eq!(left[0].0, "dlg_b");
assert_eq!(live.arrive("dlg_b"), Arrival::NotOutstanding);
}
#[test]
fn the_backstop_takes_only_an_outstanding_delegation_past_its_bound() {
let (live, _scope) = scoped();
let mut old = live_delegation("worker");
old.started = Instant::now() - Duration::from_secs(120);
let old_started = old.started;
live.register("dlg_old".to_string(), old);
live.register("dlg_new".to_string(), live_delegation("worker"));
let mut arrived = live_delegation("worker");
arrived.started = Instant::now() - Duration::from_secs(600);
live.register("dlg_arrived".to_string(), arrived);
live.arrive("dlg_arrived");
let bound = Duration::from_secs(60);
assert_eq!(live.next_overdue(bound), Some(old_started + bound));
let overdue = live.take_overdue(bound);
let ids: Vec<&str> = overdue.iter().map(|(id, _)| id.as_str()).collect();
assert_eq!(ids, vec!["dlg_old"]);
assert_eq!(live.counts(), (1, 1));
}
#[tokio::test]
async fn each_arrival_wakes_the_wait_for_the_next_one() {
let (live, _scope) = scoped();
live.register("dlg_a".to_string(), live_delegation("worker"));
live.register("dlg_b".to_string(), live_delegation("worker"));
assert_eq!(live.arrive("dlg_a"), Arrival::Delivered);
tokio::time::timeout(Duration::from_secs(5), live.wait_for_arrival())
.await
.expect("an arrival nobody was waiting for resolves the next wait");
assert!(
tokio::time::timeout(Duration::from_millis(100), live.wait_for_arrival())
.await
.is_err(),
"an arrival already woken for does not resolve the wait again"
);
assert_eq!(
live.counts(),
(1, 1),
"the arrived delegation stays in the set"
);
let door = Arc::clone(&live);
let waiting = tokio::spawn(async move {
tokio::time::timeout(Duration::from_secs(5), live.wait_for_arrival())
.await
.expect("the arrival wakes the wait");
live
});
tokio::task::yield_now().await;
assert_eq!(door.arrive("dlg_b"), Arrival::Delivered);
let live = waiting.await.unwrap();
assert_eq!(live.counts(), (0, 2));
}
#[test]
fn an_empty_residue_produces_no_artifact() {
assert!(Residue::default().into_artifact().is_none());
assert!(Residue::snapshot(None, &LiveDelegations::new())
.into_artifact()
.is_none());
}
#[test]
fn a_residue_produces_one_data_part_per_item() {
let (live, _scope) = scoped();
live.register("dlg_one".to_string(), live_delegation("worker"));
let residue = Residue {
items: vec![
ResidueItem::DetachedShell {
work_id: "wrk_abc".to_string(),
binary: "sleep".to_string(),
command: "sleep 30".to_string(),
started_at_ms: 17,
},
live.live().remove(0),
],
};
assert_eq!(residue.detached_work_ids(), vec!["wrk_abc".to_string()]);
assert_eq!(residue.delegation_ids(), vec!["dlg_one".to_string()]);
let artifact = residue.into_artifact().expect("two items is not empty");
assert_eq!(artifact.artifact_id, "residue");
assert_eq!(artifact.name, "residue");
assert_eq!(artifact.parts.len(), 2);
for part in &artifact.parts {
assert_eq!(part.media_type.as_deref(), Some("application/json"));
}
let wire = serde_json::to_value(&artifact).unwrap();
crate::a2a_conformance::check_message(&wire, "lf.a2a.v1.Artifact")
.unwrap_or_else(|errors| panic!("{wire}: {errors:?}"));
let shell = artifact.parts[0].as_data().expect("a data part");
assert_eq!(shell["kind"], "detached_shell");
assert_eq!(shell["work_id"], "wrk_abc");
assert_eq!(shell["command"], "sleep 30");
assert_eq!(shell["binary"], "sleep");
assert_eq!(shell["started_at_ms"], 17);
let delegation = artifact.parts[1].as_data().expect("a data part");
assert_eq!(delegation["kind"], "delegation");
assert_eq!(delegation["delegation_id"], "dlg_one");
assert_eq!(delegation["capsule"], "worker");
assert_eq!(delegation["child_session_id"], "ses_child");
}
#[tokio::test]
async fn a_cancel_signal_resolves_before_and_after_it_is_raised() {
let signal = CancelSignal::new();
assert!(!signal.is_canceled());
let observer = signal.clone();
let waiting = tokio::spawn(async move { observer.canceled().await });
signal.cancel();
waiting.await.unwrap();
assert!(signal.is_canceled());
tokio::time::timeout(std::time::Duration::from_secs(5), signal.canceled())
.await
.expect("an already-cancelled signal does not wait");
}
}