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) session_messenger: Option<Arc<bamboo_engine::SessionMessenger>>,
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) project_store: Option<Arc<bamboo_projects::ProjectStore>>,
pub(crate) workspace_resolver: bamboo_agent_core::workspace_state::WorkspaceResolver,
pub(crate) parent_wait_slots: Arc<dashmap::DashMap<String, Arc<ParentWaitSlot>>>,
}
#[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 {
fn finish_child_save(
&self,
child: &Session,
saved: std::io::Result<()>,
) -> Result<(), ChildSessionError> {
saved.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(())
}
#[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>>>>,
session_messenger: Option<Arc<bamboo_engine::SessionMessenger>>,
subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
config: Arc<RwLock<Config>>,
) -> Self {
Self {
session_store,
storage,
persistence,
session_messenger,
scheduler,
sessions_cache,
agent_runners,
session_event_senders,
subagent_model_resolver,
config,
project_store: None,
workspace_resolver:
bamboo_agent_core::workspace_state::WorkspaceResolver::from_process_globals(),
parent_wait_slots: Arc::new(dashmap::DashMap::new()),
}
}
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()
}
pub async fn terminal_child_ids(
&self,
parent_session_id: &str,
candidates: &[String],
) -> Vec<(String, String)> {
let statuses = self
.storage
.list_child_run_statuses(parent_session_id)
.await
.unwrap_or_default();
candidates
.iter()
.filter_map(|candidate| {
statuses.iter().find_map(|(id, status)| {
let status = status.as_deref()?;
(id == candidate && is_terminal_child_status(status))
.then(|| (candidate.clone(), status.to_string()))
})
})
.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 persisted_parent_workspace = parent_session.workspace_path_meta();
let parent_workspace_is_project_default = parent_session
.metadata
.get(bamboo_engine::project_context::WORKSPACE_SOURCE_METADATA_KEY)
.map(String::as_str)
== Some(bamboo_engine::project_context::WorkspaceSource::ProjectDefault.as_str());
let workspace_source = if parent_workspace_is_project_default
|| (persisted_parent_workspace.is_none()
&& matches!(
bamboo_engine::project_context::ProjectContextResolver::session_project_identity(
parent_session
),
bamboo_engine::project_context::SessionProjectIdentity::Assigned(_)
))
{
bamboo_engine::project_context::WorkspaceSource::ProjectDefault
} else {
match parent_session
.metadata
.get(bamboo_engine::project_context::WORKSPACE_SOURCE_METADATA_KEY)
.map(String::as_str)
{
Some("project_default") => {
bamboo_engine::project_context::WorkspaceSource::ProjectDefault
}
_ => bamboo_engine::project_context::WorkspaceSource::Session,
}
};
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: if parent_workspace_is_project_default {
String::new()
} else {
persisted_parent_workspace.unwrap_or_default()
},
workspace_source,
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 {
fn publish_child_workspace(
&self,
session_id: &str,
workspace: std::path::PathBuf,
source: &str,
) -> std::path::PathBuf {
self.workspace_resolver
.publish_resolved_workspace(session_id, workspace, source)
}
async fn validate_child_workspace(
&self,
project_id: Option<&bamboo_domain::ProjectId>,
requested_workspace: &str,
) -> Result<String, ChildSessionError> {
let Some(store) = self.project_store.as_deref() else {
if requested_workspace.trim().is_empty() {
return Err(ChildSessionError::InvalidArguments(
"child workspace must be a non-empty path".to_string(),
));
}
let requested = std::path::PathBuf::from(requested_workspace);
if requested.exists() && !requested.is_dir() {
return Err(ChildSessionError::InvalidArguments(format!(
"child workspace is not a directory: {requested_workspace}"
)));
}
let canonical = requested.canonicalize().unwrap_or(requested);
let final_workspace =
bamboo_agent_core::workspace_state::resolve_workspace_path(canonical);
return Ok(bamboo_config::paths::path_to_display_string(
&final_workspace,
));
};
let final_workspace = crate::project_context::validate_workspace_assignment_with_resolver(
store,
project_id,
Some(requested_workspace),
&self.workspace_resolver,
)
.map_err(|error| ChildSessionError::InvalidArguments(error.to_string()))?
.ok_or_else(|| {
ChildSessionError::InvalidArguments(
"child workspace must be a non-empty path".to_string(),
)
})?;
Ok(bamboo_config::paths::path_to_display_string(
&final_workspace,
))
}
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> {
let saved = self.persistence.merge_save_runtime(child).await;
self.finish_child_save(child, saved)
}
async fn save_child_session_authoritative_flags(
&self,
child: &mut Session,
) -> Result<(), ChildSessionError> {
let saved = self
.persistence
.save_runtime_authoritative_flags(child)
.await;
self.finish_child_save(child, saved)
}
async fn save_resident_reuse_state(
&self,
child: &mut Session,
workspace: &str,
workspace_source: bamboo_engine::project_context::WorkspaceSource,
permission_audit: bamboo_domain::PermissionAuditSeed,
no_human_approver: bool,
) -> Result<(), ChildSessionError> {
let child_id = child.id.clone();
let workspace_value = workspace.to_string();
let source_value = workspace_source.as_str().to_string();
let saved = self
.persistence
.update_authoritative_permission_posture_and_publish(
&child_id,
&permission_audit,
|latest| {
latest.workspace = Some(workspace_value.clone());
latest.set_workspace_path_meta(&workspace_value);
latest.metadata.insert(
bamboo_engine::project_context::WORKSPACE_SOURCE_METADATA_KEY.to_string(),
source_value,
);
latest
.agent_runtime_state
.get_or_insert_with(bamboo_domain::AgentRuntimeState::default)
.no_human_approver = no_human_approver;
},
|latest| {
self.sessions_cache.insert(
latest.id.clone(),
Arc::new(parking_lot::RwLock::new(latest.clone())),
);
},
)
.await
.map_err(|error| {
ChildSessionError::Execution(format!(
"failed to atomically re-seed resident child: {error}"
))
})?
.ok_or_else(|| ChildSessionError::NotFound(child_id.clone()))?;
*child = saved;
self.publish_child_workspace(
&child.id,
std::path::PathBuf::from(workspace),
workspace_source.as_str(),
);
Ok(())
}
async fn send_session_message(
&self,
source_session_id: &str,
target_session_id: &str,
message: &str,
idempotency_key: Option<&str>,
) -> Result<
bamboo_engine::session_app::child_session::ChildSessionMessageDelivery,
ChildSessionError,
> {
let messenger = self.session_messenger.as_ref().ok_or_else(|| {
ChildSessionError::Execution(
"logical SessionMessenger is not configured for this runtime".to_string(),
)
})?;
let id = idempotency_key.map_or_else(bamboo_domain::SessionMessageId::new, |key| {
bamboo_domain::SessionMessageId::stable(
"subagent_send_message",
&serde_json::json!({
"source_session_id": source_session_id,
"target_session_id": target_session_id,
"tool_call_id": key,
}),
)
});
let envelope = bamboo_domain::SessionMessageEnvelope {
id,
source: bamboo_domain::SessionMessageSource::Session {
session_id: source_session_id.to_string(),
},
target_session_id: target_session_id.to_string(),
kind: bamboo_domain::SessionMessageKind::PeerMessage,
body: bamboo_domain::SessionMessageBody::Content(
bamboo_domain::SessionMessageContent::text(message),
),
created_at: chrono::Utc::now(),
thread_id: None,
in_reply_to: None,
attempt: None,
correlation_id: None,
};
match messenger.send(envelope).await {
Ok(receipt) => Ok(
bamboo_engine::session_app::child_session::ChildSessionMessageDelivery::Activated(
receipt,
),
),
Err(bamboo_engine::SessionMessengerError::Activation {
receipt, source, ..
}) => Ok(
bamboo_engine::session_app::child_session::ChildSessionMessageDelivery::ActivationPending {
delivery: receipt,
error: source.to_string(),
},
),
Err(error) => Err(ChildSessionError::Execution(error.to_string())),
}
}
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>>());
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 terminal_child_ids(
&self,
parent_session_id: &str,
candidates: &[String],
) -> Vec<(String, String)> {
ChildSessionAdapter::terminal_child_ids(self, parent_session_id, candidates).await
}
async fn ensure_child_indexed(&self, child_session_id: &str) {
let _ = self.session_store.get_index_entry(child_session_id).await;
}
}