use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
use tokio::sync::{broadcast, RwLock};
use tokio::time::{sleep, Duration, Instant};
use crate::app_state::session_events::get_or_create_event_sender;
use crate::app_state::{AgentRunner, AgentStatus};
use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::{AgentEvent, Session, SessionKind};
use bamboo_domain::session::runtime_state::{
AgentRuntimeState, ChildWaitPolicy, WaitingForChildrenState,
};
use bamboo_engine::execution::spawn::{SpawnJob, SpawnScheduler};
use bamboo_engine::session_app::child_session::{
ChildRunnerInfo, ChildSessionEntry, ChildSessionError, ChildSessionPort, DeleteChildResult,
SubagentResolutionPort,
};
use bamboo_llm::Config;
use bamboo_storage::{LockedSessionStore, SessionIndexEntry, SessionStoreV2};
pub struct ChildSessionAdapter {
pub(crate) session_store: Arc<SessionStoreV2>,
pub(crate) storage: Arc<dyn Storage>,
pub(crate) persistence: Arc<LockedSessionStore>,
pub(crate) scheduler: Arc<SpawnScheduler>,
pub(crate) sessions_cache: bamboo_engine::SessionCache,
pub(crate) agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
pub(crate) session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
pub(crate) subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
pub(crate) config: Arc<RwLock<Config>>,
pub(crate) parent_wait_slots: Arc<dashmap::DashMap<String, Arc<ParentWaitSlot>>>,
pub(crate) notification_relay: Option<crate::app_state::session_events::NotificationRelayDeps>,
}
#[derive(Default)]
pub(crate) struct ParentWaitSlot {
flush_lock: tokio::sync::Mutex<()>,
pending: parking_lot::Mutex<Vec<(String, Option<String>)>>,
}
const AGENT_RUNTIME_STATE_METADATA_KEY: &str = "agent.runtime.state";
fn is_terminal_child_status(status: &str) -> bool {
matches!(
status,
"completed" | "error" | "timeout" | "cancelled" | "skipped"
)
}
fn read_runtime_state(session: &Session) -> AgentRuntimeState {
session
.agent_runtime_state
.clone()
.or_else(|| {
session
.metadata
.get(AGENT_RUNTIME_STATE_METADATA_KEY)
.and_then(|raw| serde_json::from_str::<AgentRuntimeState>(raw).ok())
})
.unwrap_or_else(|| AgentRuntimeState::new(format!("{}-wait", session.id)))
}
fn write_runtime_state(session: &mut Session, runtime_state: &AgentRuntimeState) {
session.agent_runtime_state = Some(runtime_state.clone());
if let Ok(serialized) = serde_json::to_string(runtime_state) {
session
.metadata
.insert(AGENT_RUNTIME_STATE_METADATA_KEY.to_string(), serialized);
}
}
impl ChildSessionAdapter {
#[allow(clippy::too_many_arguments)]
pub fn new(
session_store: Arc<SessionStoreV2>,
storage: Arc<dyn Storage>,
persistence: Arc<LockedSessionStore>,
scheduler: Arc<SpawnScheduler>,
sessions_cache: bamboo_engine::SessionCache,
agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
config: Arc<RwLock<Config>>,
) -> Self {
Self {
session_store,
storage,
persistence,
scheduler,
sessions_cache,
agent_runners,
session_event_senders,
subagent_model_resolver,
config,
parent_wait_slots: Arc::new(dashmap::DashMap::new()),
notification_relay: None,
}
}
pub async fn resolve_subagent_model(
&self,
subagent_type: &str,
) -> Option<bamboo_domain::ProviderModelRef> {
match &self.subagent_model_resolver {
Some(resolver) => resolver(subagent_type.to_string()).await,
None => None,
}
}
pub async fn resolve_runtime_metadata(&self, subagent_type: &str) -> HashMap<String, String> {
let config = self.config.read().await;
bamboo_engine::external_agents::config::resolve_runtime_metadata(&config, subagent_type)
}
pub async fn register_parent_wait_for_child(
&self,
parent_session_id: &str,
child_session_id: &str,
tool_call_id: Option<&str>,
) -> Result<(), ChildSessionError> {
let slot = self
.parent_wait_slots
.entry(parent_session_id.to_string())
.or_default()
.clone();
slot.pending.lock().push((
child_session_id.to_string(),
tool_call_id.map(str::to_string),
));
let _flush_guard = slot.flush_lock.lock().await;
let batch: Vec<(String, Option<String>)> = {
let mut pending = slot.pending.lock();
pending.drain(..).collect()
};
if batch.is_empty() {
return Ok(());
}
if let Err(error) = self
.flush_parent_waits(parent_session_id, &batch, ChildWaitPolicy::All)
.await
{
let mut pending = slot.pending.lock();
for item in batch {
pending.push(item);
}
return Err(error);
}
self.parent_wait_slots
.remove_if(parent_session_id, |_, slot| slot.pending.lock().is_empty());
Ok(())
}
pub async fn register_parent_wait_for_children(
&self,
parent_session_id: &str,
child_session_ids: &[String],
policy: ChildWaitPolicy,
) -> Result<usize, ChildSessionError> {
if child_session_ids.is_empty() {
return Ok(0);
}
let batch: Vec<(String, Option<String>)> = child_session_ids
.iter()
.map(|id| (id.clone(), None))
.collect();
self.flush_parent_waits(parent_session_id, &batch, policy)
.await?;
Ok(batch.len())
}
pub async fn active_child_ids(&self, parent_session_id: &str) -> Vec<String> {
self.storage
.list_child_run_statuses(parent_session_id)
.await
.unwrap_or_default()
.into_iter()
.filter(|(_, status)| !status.as_deref().is_some_and(is_terminal_child_status))
.map(|(id, _)| id)
.collect()
}
async fn flush_parent_waits(
&self,
parent_session_id: &str,
batch: &[(String, Option<String>)],
policy: ChildWaitPolicy,
) -> Result<(), ChildSessionError> {
let Some(mut parent) =
self.storage
.load_session(parent_session_id)
.await
.map_err(|error| {
ChildSessionError::Execution(format!(
"failed to load parent session {parent_session_id}: {error}"
))
})?
else {
return Err(ChildSessionError::NotFound(parent_session_id.to_string()));
};
let mut runtime_state = read_runtime_state(&parent);
let now = Utc::now();
let mut wait = runtime_state
.waiting_for_children
.take()
.unwrap_or_else(|| WaitingForChildrenState::for_children(Vec::new(), policy, now));
wait.wait_for = policy;
for (child_session_id, tool_call_id) in batch {
if !wait
.child_session_ids
.iter()
.any(|id| id == child_session_id)
{
wait.child_session_ids.push(child_session_id.clone());
}
if wait.registered_by_tool_call_id.is_none() {
wait.registered_by_tool_call_id = tool_call_id.clone();
}
}
wait.child_session_ids.sort();
wait.child_session_ids.dedup();
runtime_state.waiting_for_children = Some(wait);
write_runtime_state(&mut parent, &runtime_state);
parent.metadata.insert(
"runtime.suspend_reason".to_string(),
"waiting_for_children".to_string(),
);
parent.updated_at = Utc::now();
self.persistence
.save_runtime_only(&mut parent)
.await
.map_err(|error| {
ChildSessionError::Execution(format!("failed to save parent wait state: {error}"))
})?;
self.sessions_cache.insert(
parent.id.clone(),
Arc::new(parking_lot::RwLock::new(parent)),
);
Ok(())
}
}
fn map_index_entry_to_child_entry(entry: &SessionIndexEntry) -> ChildSessionEntry {
ChildSessionEntry {
child_session_id: entry.id.clone(),
title: entry.title.clone(),
pinned: entry.pinned,
message_count: entry.message_count,
updated_at: entry.updated_at.to_rfc3339(),
last_run_status: entry.last_run_status.clone(),
last_run_error: entry.last_run_error.clone(),
}
}
#[async_trait]
impl SubagentResolutionPort for ChildSessionAdapter {
async fn resolve_subagent_model(
&self,
subagent_type: &str,
) -> Option<bamboo_domain::ProviderModelRef> {
ChildSessionAdapter::resolve_subagent_model(self, subagent_type).await
}
async fn resolve_runtime_metadata(
&self,
subagent_type: &str,
) -> std::collections::HashMap<String, String> {
ChildSessionAdapter::resolve_runtime_metadata(self, subagent_type).await
}
}
#[async_trait]
impl bamboo_engine::GuardianSpawner for ChildSessionAdapter {
async fn spawn_guardian_review(
&self,
parent_session: &Session,
review_prompt: String,
model: String,
disabled_tools: Option<std::collections::BTreeSet<String>>,
) -> Result<String, String> {
let input = bamboo_engine::session_app::child_session::CreateChildInput {
parent_session: parent_session.clone(),
child_id: format!("guardian-{}", uuid::Uuid::new_v4()),
title: "Guardian review".to_string(),
responsibility: "Adversarially verify the parent agent's completed work.".to_string(),
assignment_prompt: review_prompt,
subagent_type: "guardian".to_string(),
workspace: parent_session.workspace_path_meta().unwrap_or_default(),
model_override: Some(model),
model_ref_override: None,
runtime_metadata: HashMap::new(),
auto_run: true,
reasoning_effort: None,
lifecycle: None,
resident_name: None,
resident_context: None,
disabled_tools,
context_fork: None,
};
bamboo_engine::session_app::child_session::create_child_action(self, input)
.await
.map(|result| result.child_session_id)
.map_err(|error| error.to_string())
}
}
#[async_trait]
impl ChildSessionPort for ChildSessionAdapter {
async fn load_root_session(&self, root_session_id: &str) -> Result<Session, ChildSessionError> {
let Some(session) = self
.storage
.load_session(root_session_id)
.await
.map_err(|error| {
ChildSessionError::Execution(format!(
"failed to load session {root_session_id}: {error}"
))
})?
else {
return Err(ChildSessionError::NotFound(root_session_id.to_string()));
};
if session.kind != SessionKind::Root {
return Err(ChildSessionError::NotRootSession(
root_session_id.to_string(),
));
}
Ok(session)
}
async fn load_child_for_parent(
&self,
parent_session_id: &str,
child_session_id: &str,
) -> Result<Session, ChildSessionError> {
let Some(child) = self
.storage
.load_session(child_session_id)
.await
.map_err(|error| {
ChildSessionError::Execution(format!(
"failed to load child session {child_session_id}: {error}"
))
})?
else {
return Err(ChildSessionError::NotFound(child_session_id.to_string()));
};
if child.kind != SessionKind::Child {
return Err(ChildSessionError::NotChildSession(
child_session_id.to_string(),
));
}
if child.parent_session_id.as_deref() != Some(parent_session_id) {
return Err(ChildSessionError::NotChildOfParent {
child_id: child_session_id.to_string(),
parent_id: parent_session_id.to_string(),
});
}
Ok(child)
}
async fn save_child_session(&self, child: &mut Session) -> Result<(), ChildSessionError> {
self.persistence
.merge_save_runtime(child)
.await
.map_err(|error| {
ChildSessionError::Execution(format!("failed to save child session: {error}"))
})?;
self.sessions_cache.insert(
child.id.clone(),
Arc::new(parking_lot::RwLock::new(child.clone())),
);
Ok(())
}
async fn is_child_running(&self, child_session_id: &str) -> bool {
let runners = self.agent_runners.read().await;
runners
.get(child_session_id)
.is_some_and(|runner| matches!(runner.status, AgentStatus::Running))
}
async fn list_children(&self, parent_session_id: &str) -> Vec<ChildSessionEntry> {
self.session_store
.list_index_entries()
.await
.into_iter()
.filter(|entry| {
entry.kind == SessionKind::Child
&& entry.parent_session_id.as_deref() == Some(parent_session_id)
})
.map(|entry| map_index_entry_to_child_entry(&entry))
.collect()
}
async fn find_resident_child(
&self,
root_session_id: &str,
resident_name: &str,
) -> Option<String> {
let name = resident_name.trim();
if name.is_empty() {
return None;
}
let mut best: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
for entry in self.session_store.list_index_entries().await {
if entry.kind == SessionKind::Child
&& entry.root_session_id == root_session_id
&& entry.resident_name.as_deref() == Some(name)
{
match &best {
Some((_, ts)) if *ts >= entry.updated_at => {}
_ => best = Some((entry.id.clone(), entry.updated_at)),
}
}
}
best.map(|(id, _)| id)
}
async fn enqueue_child_run(
&self,
parent: &Session,
child: &Session,
) -> Result<(), ChildSessionError> {
let model = if child.model.trim().is_empty() {
parent.model.clone()
} else {
child.model.clone()
};
if model.trim().is_empty() {
return Err(ChildSessionError::Execution(
"child model is empty and parent model is unavailable".to_string(),
));
}
let disabled_tools = child
.metadata
.get("disabled_tools")
.and_then(|raw| serde_json::from_str::<std::collections::BTreeSet<String>>(raw).ok())
.filter(|set| !set.is_empty())
.map(|set| set.into_iter().collect::<Vec<String>>());
if let Some(relay) = &self.notification_relay {
let child_tx = get_or_create_event_sender(&self.session_event_senders, &child.id).await;
crate::app_state::session_events::ensure_notification_relay(relay, &child.id, child_tx);
}
self.scheduler
.enqueue(SpawnJob {
parent_session_id: parent.id.clone(),
child_session_id: child.id.clone(),
model,
disabled_tools,
})
.await
.map_err(ChildSessionError::Execution)?;
let parent_tx = get_or_create_event_sender(&self.session_event_senders, &parent.id).await;
let _ = parent_tx.send(AgentEvent::SubAgentStarted {
parent_session_id: parent.id.clone(),
child_session_id: child.id.clone(),
title: Some(child.title.clone()),
});
Ok(())
}
async fn cancel_child_run_and_wait(
&self,
child_session_id: &str,
) -> Result<(), ChildSessionError> {
let cancelled = {
let mut runners = self.agent_runners.write().await;
if let Some(runner) = runners.get_mut(child_session_id) {
if matches!(runner.status, AgentStatus::Running) {
runner.cancel_token.cancel();
true
} else {
false
}
} else {
false
}
};
if !cancelled {
return Ok(());
}
let deadline = Instant::now() + Duration::from_secs(10);
loop {
let still_running = {
let runners = self.agent_runners.read().await;
runners
.get(child_session_id)
.is_some_and(|runner| matches!(runner.status, AgentStatus::Running))
};
if !still_running {
return Ok(());
}
if Instant::now() >= deadline {
return Err(ChildSessionError::Execution(format!(
"timed out waiting for child session {child_session_id} to stop after cancellation"
)));
}
sleep(Duration::from_millis(50)).await;
}
}
async fn delete_child_session(
&self,
parent_session_id: &str,
child_id: &str,
) -> Result<DeleteChildResult, ChildSessionError> {
let cancelled_running_child = {
let mut runners = self.agent_runners.write().await;
if let Some(runner) = runners.remove(child_id) {
runner.cancel_token.cancel();
true
} else {
false
}
};
let deleted = self
.storage
.delete_session(child_id)
.await
.map_err(|error| {
ChildSessionError::Execution(format!("failed to delete child session: {error}"))
})?;
self.sessions_cache.remove(child_id);
{
let mut senders = self.session_event_senders.write().await;
senders.remove(child_id);
if cancelled_running_child {
if let Some(parent_tx) = senders.get(parent_session_id) {
let _ = parent_tx.send(AgentEvent::SubAgentCompleted {
parent_session_id: parent_session_id.to_string(),
child_session_id: child_id.to_string(),
status: "cancelled".to_string(),
error: Some("Child session deleted while running".to_string()),
});
}
}
}
Ok(DeleteChildResult {
deleted,
cancelled_running_child,
})
}
async fn get_child_runner_info(&self, child_id: &str) -> Option<ChildRunnerInfo> {
let runners = self.agent_runners.read().await;
runners.get(child_id).map(|runner| ChildRunnerInfo {
started_at: Some(runner.started_at),
completed_at: runner.completed_at,
last_tool_name: runner.last_tool_name.clone(),
last_tool_phase: runner.last_tool_phase.clone(),
last_event_at: runner.last_event_at,
round_count: runner.round_count,
})
}
async fn register_parent_wait_for_child(
&self,
parent_session_id: &str,
child_session_id: &str,
tool_call_id: Option<&str>,
) -> Result<(), ChildSessionError> {
ChildSessionAdapter::register_parent_wait_for_child(
self,
parent_session_id,
child_session_id,
tool_call_id,
)
.await
}
async fn register_parent_wait_for_children(
&self,
parent_session_id: &str,
child_session_ids: &[String],
policy: ChildWaitPolicy,
) -> Result<usize, ChildSessionError> {
ChildSessionAdapter::register_parent_wait_for_children(
self,
parent_session_id,
child_session_ids,
policy,
)
.await
}
async fn active_child_ids(&self, parent_session_id: &str) -> Vec<String> {
ChildSessionAdapter::active_child_ids(self, parent_session_id).await
}
async fn ensure_child_indexed(&self, child_session_id: &str) {
let _ = self.session_store.get_index_entry(child_session_id).await;
}
}