use std::sync::Arc;
use std::time::Duration;
use chrono::Utc;
use tokio::sync::broadcast;
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
use bamboo_agent_core::{AgentEvent, Role, SessionKind};
use bamboo_domain::{
AgentRuntimeState, AgentStatusState, SuspensionState, WaitingForBashState,
WaitingForChildrenState,
};
use crate::runtime::execution::event_forwarder::create_event_forwarder;
use crate::runtime::execution::runner_lifecycle::{finalize_runner, try_reserve_runner};
use crate::runtime::execution::session_events::get_or_create_event_sender;
use crate::runtime::execution::spawn::{
publish_child_completion_parts, watch_child_liveness, watchdog_policy_for_session,
SpawnContext, SpawnJob,
};
use crate::runtime::execution::SessionExecutionReservation;
type HostWaitSnapshot = (Option<WaitingForChildrenState>, Option<WaitingForBashState>);
fn host_wait_snapshot(session: &bamboo_agent_core::Session) -> HostWaitSnapshot {
crate::runtime::runner::state_bridge::read_runtime_state(session)
.map(|state| (state.waiting_for_children, state.waiting_for_bash))
.unwrap_or((None, None))
}
fn reconcile_actor_host_wait(
session: &mut bamboo_agent_core::Session,
durable: Option<HostWaitSnapshot>,
run_id: &str,
) {
if session
.metadata
.get("runtime.suspend_reason")
.is_some_and(|reason| !reason.trim().is_empty())
{
return;
}
let durable_available = durable.is_some();
let mut runtime_state = crate::runtime::runner::state_bridge::read_runtime_state(session)
.unwrap_or_else(|| AgentRuntimeState::new(run_id));
if let Some((waiting_for_children, waiting_for_bash)) = durable {
runtime_state.waiting_for_children = waiting_for_children;
runtime_state.waiting_for_bash = waiting_for_bash;
}
let reason = if runtime_state.waiting_for_children.is_some() {
Some(("waiting_for_children", "ChildCompletion"))
} else if runtime_state.waiting_for_bash.is_some() {
Some(("waiting_for_bash", "BashCompletion"))
} else {
None
};
if let Some((reason, hook_point)) = reason {
runtime_state.status = AgentStatusState::Suspended;
runtime_state.suspension = Some(SuspensionState {
reason: reason.to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some(hook_point.to_string()),
});
session
.metadata
.insert("runtime.suspend_reason".to_string(), reason.to_string());
} else if durable_available
&& runtime_state.suspension.as_ref().is_some_and(|suspension| {
matches!(
suspension.reason.as_str(),
"waiting_for_children" | "waiting_for_bash"
)
})
{
runtime_state.suspension = None;
if runtime_state.status == AgentStatusState::Suspended {
runtime_state.status = AgentStatusState::Idle;
}
}
crate::runtime::runner::state_bridge::write_runtime_state(session, &runtime_state);
if let Ok(serialized) = serde_json::to_string(&runtime_state) {
session
.metadata
.insert("agent.runtime.state".to_string(), serialized);
}
}
pub async fn run_child_spawn(ctx: SpawnContext, job: SpawnJob) -> Result<(), String> {
run_child_spawn_inner(ctx, job, None).await
}
pub(crate) async fn run_child_spawn_reserved(
ctx: SpawnContext,
job: SpawnJob,
reservation: SessionExecutionReservation,
) -> Result<(), String> {
run_child_spawn_inner(ctx, job, Some(reservation)).await
}
async fn run_child_spawn_inner(
ctx: SpawnContext,
job: SpawnJob,
reserved: Option<SessionExecutionReservation>,
) -> Result<(), String> {
let parent_tx =
get_or_create_event_sender(&ctx.session_event_senders, &job.parent_session_id).await;
let child_tx =
get_or_create_event_sender(&ctx.session_event_senders, &job.child_session_id).await;
let mut session = match ctx
.agent
.storage()
.load_session(&job.child_session_id)
.await
{
Ok(Some(s)) => s,
Ok(None) => {
let error = "child session not found".to_string();
publish_child_completion_parts(
&parent_tx,
ctx.completion_handler.clone(),
job.parent_session_id.clone(),
job.child_session_id.clone(),
"error".to_string(),
Some(error.clone()),
)
.await;
return Err(error);
}
Err(e) => {
let error = format!("failed to load child session: {e}");
publish_child_completion_parts(
&parent_tx,
ctx.completion_handler.clone(),
job.parent_session_id.clone(),
job.child_session_id.clone(),
"error".to_string(),
Some(error.clone()),
)
.await;
return Err(error);
}
};
let reserved_identity = reserved.as_ref().map(|reservation| {
(
reservation.run_id().to_string(),
reservation.matches_execution_target(
&job.child_session_id,
&session.id,
&ctx.agent_runners,
),
)
});
let reserved_activation = if let Some((reserved_run_id, exact_target)) = reserved_identity {
let exact_runner = ctx
.agent_runners
.read()
.await
.get(&job.child_session_id)
.is_some_and(|runner| {
runner.run_id == reserved_run_id
&& matches!(
runner.status,
crate::runtime::execution::AgentStatus::Running
)
});
let exact_owner = match ctx.agent.activation_router() {
Some(router) => {
router
.owns_run(&job.child_session_id, &reserved_run_id)
.await
}
None => false,
};
let authorized_prefix = match ctx.agent.session_inbox() {
Some(inbox) => inbox
.inspect(&job.child_session_id)
.await
.map_err(|error| {
format!(
"inspect reserved child SessionInbox {}: {error}",
job.child_session_id
)
})?
.activation_pending(),
None => false,
};
if !exact_target || !exact_runner || !exact_owner || !authorized_prefix {
return Err(format!(
"reserved child activation {} failed ownership/prefix validation for {}",
reserved_run_id, job.child_session_id
));
}
true
} else {
false
};
if session.kind != SessionKind::Child {
let error = "spawn job child session is not kind=child".to_string();
if !reserved_activation {
session.set_last_run_status("error");
session.set_last_run_error(error.clone());
let _ = ctx
.agent
.persistence()
.save_runtime_session(&mut session)
.await;
}
publish_child_completion_parts(
&parent_tx,
ctx.completion_handler.clone(),
job.parent_session_id.clone(),
job.child_session_id.clone(),
"error".to_string(),
Some(error.clone()),
)
.await;
return Err(error);
}
session.metadata.remove("runtime.suspend_reason");
session.clear_pending_question();
let last_is_user = session
.messages
.last()
.map(|m| matches!(m.role, Role::User))
.unwrap_or(false);
if !last_is_user && !reserved_activation {
session.set_last_run_status("skipped");
session.set_last_run_error("No pending message to execute");
let _ = ctx
.agent
.persistence()
.save_runtime_session(&mut session)
.await;
ctx.sessions_cache.insert(
job.child_session_id.clone(),
Arc::new(parking_lot::RwLock::new(session)),
);
publish_child_completion_parts(
&parent_tx,
ctx.completion_handler.clone(),
job.parent_session_id.clone(),
job.child_session_id.clone(),
"skipped".to_string(),
Some("No pending message to execute".to_string()),
)
.await;
return Ok(());
}
let mut execution_reservation = if let Some(reservation) = reserved {
reservation
} else {
let Some(reservation) = try_reserve_runner(
&ctx.agent_runners,
&ctx.session_event_senders,
&job.child_session_id,
&child_tx,
)
.await
else {
tracing::warn!(
parent_session_id = %job.parent_session_id,
child_session_id = %job.child_session_id,
"child spawn skipped: runner already registered for this child"
);
return Ok(());
};
SessionExecutionReservation::from_pending_registration(
job.child_session_id.clone(),
reservation,
ctx.agent.activation_router().cloned(),
ctx.agent_runners.clone(),
)
};
let run_id = execution_reservation.run_id().to_string();
if let Err(error) = execution_reservation.ensure_registered().await {
let same_run_duplicate = error.existing_run_id() == run_id;
tracing::warn!(
parent_session_id = %job.parent_session_id,
child_session_id = %job.child_session_id,
run_id = %run_id,
%error,
"child runner registration collided with an existing logical-session owner"
);
if same_run_duplicate {
return Ok(());
}
let setup_error = error.to_string();
publish_child_completion_parts(
&parent_tx,
ctx.completion_handler.clone(),
job.parent_session_id.clone(),
job.child_session_id.clone(),
"error".to_string(),
Some(setup_error.clone()),
)
.await;
return Err(setup_error);
}
let (cancel_token, mut activation_registration) = execution_reservation.disarm_for_execution();
if let Some(ref ws) = session.workspace {
let stored = bamboo_agent_core::workspace_state::set_workspace(
&session.id,
std::path::PathBuf::from(ws),
);
session.workspace = Some(stored.to_string_lossy().to_string());
}
session.set_last_run_status("running");
session.clear_last_run_error();
let _ = ctx
.agent
.persistence()
.save_runtime_session(&mut session)
.await;
let forwarder_done = CancellationToken::new();
{
let mut rx = child_tx.subscribe();
let parent_tx = parent_tx.clone();
let job_clone = job.clone();
let done = forwarder_done.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = done.cancelled() => break,
evt = rx.recv() => {
match evt {
Ok(event) => {
let _ = parent_tx.send(AgentEvent::SubAgentEvent {
parent_session_id: job_clone.parent_session_id.clone(),
child_session_id: job_clone.child_session_id.clone(),
event: Box::new(event),
});
}
Err(broadcast::error::RecvError::Lagged(_)) => {
continue;
}
Err(_) => break,
}
}
}
}
});
}
{
let parent_tx = parent_tx.clone();
let job_clone = job.clone();
let done = forwarder_done.clone();
tokio::spawn(async move {
let mut ticker = tokio::time::interval(Duration::from_secs(5));
loop {
tokio::select! {
_ = done.cancelled() => break,
_ = ticker.tick() => {
let _ = parent_tx.send(AgentEvent::SubAgentHeartbeat {
parent_session_id: job_clone.parent_session_id.clone(),
child_session_id: job_clone.child_session_id.clone(),
timestamp: Utc::now(),
});
}
}
}
});
}
let (mpsc_tx, _forwarder_handle) = create_event_forwarder(
job.child_session_id.clone(),
child_tx.clone(),
ctx.agent_runners.clone(),
ctx.account_feed_inbox.clone(),
);
let timeout_reason = Arc::new(RwLock::new(None::<String>));
let watchdog_policy = watchdog_policy_for_session(&session);
tokio::spawn(watch_child_liveness(
job.parent_session_id.clone(),
job.child_session_id.clone(),
ctx.agent_runners.clone(),
cancel_token.clone(),
timeout_reason.clone(),
forwarder_done.clone(),
watchdog_policy,
));
let model = job.model.clone();
let session_id_clone = job.child_session_id.clone();
let agent_runners_for_status = ctx.agent_runners.clone();
let sessions_cache = ctx.sessions_cache.clone();
let agent = ctx.agent.clone();
let external_runner = ctx.external_child_runner.clone();
let done = forwarder_done.clone();
let parent_tx_for_done = parent_tx.clone();
let parent_id_for_done = job.parent_session_id.clone();
let child_id_for_done = job.child_session_id.clone();
let session_event_senders = ctx.session_event_senders.clone();
let completion_handler = ctx.completion_handler.clone();
let activation_run_id = run_id;
tokio::spawn(async move {
crate::session_app::execution_prep::prepare_session_for_execution(
&mut session,
None,
Some(&model),
);
let result: crate::runtime::runner::Result<()> =
if external_runner.should_handle(&session).await {
use futures::FutureExt;
match std::panic::AssertUnwindSafe(external_runner.execute_external_child(
&mut session,
&job,
mpsc_tx,
cancel_token.clone(),
))
.catch_unwind()
.await
{
Ok(result) => result,
Err(panic) => {
let message = panic
.downcast_ref::<&str>()
.map(|s| (*s).to_string())
.or_else(|| panic.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "non-string panic payload".to_string());
tracing::error!(
parent_session_id = %job.parent_session_id,
child_session_id = %job.child_session_id,
panic = %message,
"child execution panicked; publishing terminal error completion"
);
Err(bamboo_agent_core::AgentError::LLM(format!(
"child execution panicked: {message}"
)))
}
}
} else {
Err(bamboo_agent_core::AgentError::LLM(format!(
"No child runner matched session runtime metadata: agent_id={:?}, protocol={:?}",
session.metadata.get("external.agent_id"),
session.metadata.get("external.protocol"),
)))
};
let durable_wait = match agent.storage().load_session(&session_id_clone).await {
Ok(Some(persisted)) => Some(host_wait_snapshot(&persisted)),
Ok(None) => None,
Err(error) => {
tracing::warn!(
session_id = %session_id_clone,
%error,
"could not reconcile host-persisted wait ownership after actor execution; preserving in-memory wait"
);
None
}
};
reconcile_actor_host_wait(&mut session, durable_wait, &activation_run_id);
let timeout_error = timeout_reason.read().await.clone();
let suspended_non_terminal = result.is_ok()
&& session
.metadata
.get("runtime.suspend_reason")
.is_some_and(|reason| !reason.trim().is_empty());
let (status, error) = if let Some(reason) = timeout_error {
("timeout".to_string(), Some(reason))
} else if suspended_non_terminal {
("suspended".to_string(), None)
} else {
match &result {
Ok(_) => ("completed".to_string(), None),
Err(e @ bamboo_agent_core::AgentError::Cancelled) => {
("cancelled".to_string(), Some(e.to_string()))
}
Err(e) => ("error".to_string(), Some(e.to_string())),
}
};
let executed_admitted_generation = session
.session_inbox_admission()
.map_or(0, |state| state.last_admitted_sequence);
let legacy_migration = crate::runtime::runner::state_bridge::migrate_legacy_pending_only(
&mut session,
Some(agent.storage()),
Some(agent.persistence()),
agent.session_inbox(),
)
.await;
let pending_boundary_generation = session
.session_inbox_admission()
.and_then(|state| state.pending_activation_generation());
let pending_generation = match (
pending_boundary_generation,
legacy_migration.highest_generation,
) {
(Some(left), Some(right)) => Some(left.max(right)),
(left, right) => left.or(right),
};
if let (Some(generation), Some(router)) = (pending_generation, agent.activation_router()) {
let activation_ready = if let Some(inbox) = agent.session_inbox() {
match inbox
.mark_activation_eligible(
&session_id_clone,
generation,
bamboo_domain::SessionActivationPolicy::InterruptSpecificWait,
)
.await
{
Ok(()) => true,
Err(error) => {
tracing::error!(
session_id = %session_id_clone,
%error,
"failed to persist unadmitted SessionInbox activation watermark"
);
false
}
}
} else {
false
};
if activation_ready {
if let Err(error) = bamboo_domain::SessionActivationPort::request_activation(
router.as_ref(),
&session_id_clone,
generation,
)
.await
{
tracing::error!(
session_id = %session_id_clone,
%error,
"failed to hand unadmitted SessionInbox generation to activation router"
);
}
}
}
session.set_last_run_status(status.clone());
if let Some(err) = &error {
session.set_last_run_error(err.clone());
} else {
session.clear_last_run_error();
}
if let Some(registration) = activation_registration.as_mut() {
registration.begin_finalization().await;
} else if let Some(router) = agent.activation_router() {
router
.begin_finalization(&session_id_clone, &activation_run_id)
.await;
}
let _ = agent.persistence().save_runtime_session(&mut session).await;
finalize_runner(&agent_runners_for_status, &session_id_clone, &result).await;
let finalization = if let Some(registration) = activation_registration.take() {
registration.finish(executed_admitted_generation).await
} else if let Some(router) = agent.activation_router() {
router
.finish_finalization(
&session_id_clone,
&activation_run_id,
executed_admitted_generation,
)
.await
} else {
Ok(None)
};
if let Err(error) = finalization {
tracing::error!(
session_id = %session_id_clone,
%error,
"failed to activate child successor for finalization-racing SessionInbox delivery"
);
}
sessions_cache.insert(
session_id_clone.clone(),
Arc::new(parking_lot::RwLock::new(session)),
);
done.cancel();
publish_child_completion_parts(
&parent_tx_for_done,
completion_handler,
parent_id_for_done,
child_id_for_done,
status,
error,
)
.await;
drop(session_event_senders);
});
Ok(())
}
#[cfg(test)]
mod actor_host_wait_tests {
use super::*;
use bamboo_domain::ChildWaitPolicy;
fn prepared_waiting_session() -> bamboo_agent_core::Session {
let mut session =
bamboo_agent_core::Session::new_child("child", "parent", "model", "Child");
let mut runtime = AgentRuntimeState::new("prepared-run");
runtime.status = AgentStatusState::Idle;
runtime.waiting_for_children = Some(WaitingForChildrenState::for_children(
vec!["grandchild".to_string()],
ChildWaitPolicy::All,
Utc::now(),
));
session.agent_runtime_state = Some(runtime);
session
}
#[test]
fn actor_host_read_failure_preserves_wait_and_non_terminal_status() {
let mut session = prepared_waiting_session();
reconcile_actor_host_wait(&mut session, None, "actor-run");
let runtime = session.agent_runtime_state.as_ref().unwrap();
assert!(runtime.waiting_for_children.is_some());
assert_eq!(runtime.status, AgentStatusState::Suspended);
assert_eq!(
session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str),
Some("waiting_for_children")
);
let legacy: AgentRuntimeState =
serde_json::from_str(session.metadata.get("agent.runtime.state").unwrap()).unwrap();
assert_eq!(legacy.status, AgentStatusState::Suspended);
assert!(legacy.waiting_for_children.is_some());
let suspended_non_terminal = session
.metadata
.get("runtime.suspend_reason")
.is_some_and(|reason| !reason.trim().is_empty());
assert!(suspended_non_terminal);
}
#[test]
fn actor_host_successful_durable_clear_wins_over_prepared_wait() {
let mut session = prepared_waiting_session();
session.agent_runtime_state.as_mut().unwrap().status = AgentStatusState::Suspended;
session.agent_runtime_state.as_mut().unwrap().suspension = Some(SuspensionState {
reason: "waiting_for_children".to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some("ChildCompletion".to_string()),
});
reconcile_actor_host_wait(&mut session, Some((None, None)), "actor-run");
let runtime = session.agent_runtime_state.as_ref().unwrap();
assert!(runtime.waiting_for_children.is_none());
assert!(runtime.waiting_for_bash.is_none());
assert_eq!(runtime.status, AgentStatusState::Idle);
assert!(runtime.suspension.is_none());
assert!(!session.metadata.contains_key("runtime.suspend_reason"));
let legacy: AgentRuntimeState =
serde_json::from_str(session.metadata.get("agent.runtime.state").unwrap()).unwrap();
assert_eq!(legacy, *runtime);
}
#[test]
fn actor_host_latest_legacy_only_wait_is_authoritative() {
let mut persisted =
bamboo_agent_core::Session::new_child("child", "parent", "model", "Child");
let mut legacy_runtime = AgentRuntimeState::new("legacy-run");
legacy_runtime.waiting_for_bash = Some(WaitingForBashState::for_bash(
vec!["shell-1".to_string()],
Utc::now(),
));
persisted.metadata.insert(
"agent.runtime.state".to_string(),
serde_json::to_string(&legacy_runtime).unwrap(),
);
assert!(persisted.agent_runtime_state.is_none());
let mut live = bamboo_agent_core::Session::new_child("child", "parent", "model", "Child");
reconcile_actor_host_wait(&mut live, Some(host_wait_snapshot(&persisted)), "actor-run");
let runtime = live.agent_runtime_state.as_ref().unwrap();
assert!(runtime.waiting_for_bash.is_some());
assert_eq!(runtime.status, AgentStatusState::Suspended);
assert_eq!(
live.metadata
.get("runtime.suspend_reason")
.map(String::as_str),
Some("waiting_for_bash")
);
}
}