1use std::collections::HashMap;
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use chrono::Utc;
11use tokio::sync::{broadcast, RwLock};
12use tokio::time::{sleep, Duration, Instant};
13
14use crate::app_state::session_events::get_or_create_event_sender;
15use crate::app_state::{AgentRunner, AgentStatus};
16use bamboo_agent_core::storage::Storage;
17use bamboo_agent_core::{AgentEvent, Session, SessionKind};
18use bamboo_domain::session::runtime_state::{
19 AgentRuntimeState, ChildWaitPolicy, WaitingForChildrenState,
20};
21use bamboo_engine::execution::spawn::{SpawnJob, SpawnScheduler};
22use bamboo_engine::session_app::child_session::{
23 ChildRunnerInfo, ChildSessionEntry, ChildSessionError, ChildSessionPort, DeleteChildResult,
24 SubagentResolutionPort,
25};
26use bamboo_llm::Config;
27use bamboo_storage::{LockedSessionStore, SessionIndexEntry, SessionStoreV2};
28
29pub struct ChildSessionAdapter {
34 pub(crate) session_store: Arc<SessionStoreV2>,
35 pub(crate) storage: Arc<dyn Storage>,
36 pub(crate) persistence: Arc<LockedSessionStore>,
37 pub(crate) scheduler: Arc<SpawnScheduler>,
38 pub(crate) sessions_cache: bamboo_engine::SessionCache,
39 pub(crate) agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
40 pub(crate) session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
41 pub(crate) subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
43 pub(crate) config: Arc<RwLock<Config>>,
45 pub(crate) parent_wait_slots: Arc<dashmap::DashMap<String, Arc<ParentWaitSlot>>>,
50 pub(crate) notification_relay: Option<crate::app_state::session_events::NotificationRelayDeps>,
64}
65
66#[derive(Default)]
75pub(crate) struct ParentWaitSlot {
76 flush_lock: tokio::sync::Mutex<()>,
77 pending: parking_lot::Mutex<Vec<(String, Option<String>)>>,
78}
79
80const AGENT_RUNTIME_STATE_METADATA_KEY: &str = "agent.runtime.state";
81
82fn is_terminal_child_status(status: &str) -> bool {
85 matches!(
86 status,
87 "completed" | "error" | "timeout" | "cancelled" | "skipped"
88 )
89}
90
91fn read_runtime_state(session: &Session) -> AgentRuntimeState {
92 session
93 .agent_runtime_state
94 .clone()
95 .or_else(|| {
96 session
97 .metadata
98 .get(AGENT_RUNTIME_STATE_METADATA_KEY)
99 .and_then(|raw| serde_json::from_str::<AgentRuntimeState>(raw).ok())
100 })
101 .unwrap_or_else(|| AgentRuntimeState::new(format!("{}-wait", session.id)))
102}
103
104fn write_runtime_state(session: &mut Session, runtime_state: &AgentRuntimeState) {
105 session.agent_runtime_state = Some(runtime_state.clone());
106 if let Ok(serialized) = serde_json::to_string(runtime_state) {
107 session
108 .metadata
109 .insert(AGENT_RUNTIME_STATE_METADATA_KEY.to_string(), serialized);
110 }
111}
112
113impl ChildSessionAdapter {
114 #[allow(clippy::too_many_arguments)]
119 pub fn new(
120 session_store: Arc<SessionStoreV2>,
121 storage: Arc<dyn Storage>,
122 persistence: Arc<LockedSessionStore>,
123 scheduler: Arc<SpawnScheduler>,
124 sessions_cache: bamboo_engine::SessionCache,
125 agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
126 session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
127 subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
128 config: Arc<RwLock<Config>>,
129 ) -> Self {
130 Self {
131 session_store,
132 storage,
133 persistence,
134 scheduler,
135 sessions_cache,
136 agent_runners,
137 session_event_senders,
138 subagent_model_resolver,
139 config,
140 parent_wait_slots: Arc::new(dashmap::DashMap::new()),
143 notification_relay: None,
147 }
148 }
149
150 pub async fn resolve_subagent_model(
152 &self,
153 subagent_type: &str,
154 ) -> Option<bamboo_domain::ProviderModelRef> {
155 match &self.subagent_model_resolver {
156 Some(resolver) => resolver(subagent_type.to_string()).await,
157 None => None,
158 }
159 }
160
161 pub async fn resolve_runtime_metadata(&self, subagent_type: &str) -> HashMap<String, String> {
163 let config = self.config.read().await;
164 bamboo_engine::external_agents::config::resolve_runtime_metadata(&config, subagent_type)
165 }
166
167 pub async fn register_parent_wait_for_child(
181 &self,
182 parent_session_id: &str,
183 child_session_id: &str,
184 tool_call_id: Option<&str>,
185 ) -> Result<(), ChildSessionError> {
186 let slot = self
187 .parent_wait_slots
188 .entry(parent_session_id.to_string())
189 .or_default()
190 .clone();
191
192 slot.pending.lock().push((
194 child_session_id.to_string(),
195 tool_call_id.map(str::to_string),
196 ));
197
198 let _flush_guard = slot.flush_lock.lock().await;
200
201 let batch: Vec<(String, Option<String>)> = {
204 let mut pending = slot.pending.lock();
205 pending.drain(..).collect()
206 };
207 if batch.is_empty() {
208 return Ok(());
211 }
212
213 if let Err(error) = self
215 .flush_parent_waits(parent_session_id, &batch, ChildWaitPolicy::All)
216 .await
217 {
218 let mut pending = slot.pending.lock();
220 for item in batch {
221 pending.push(item);
222 }
223 return Err(error);
224 }
225
226 self.parent_wait_slots
240 .remove_if(parent_session_id, |_, slot| slot.pending.lock().is_empty());
241
242 Ok(())
243 }
244
245 pub async fn register_parent_wait_for_children(
252 &self,
253 parent_session_id: &str,
254 child_session_ids: &[String],
255 policy: ChildWaitPolicy,
256 ) -> Result<usize, ChildSessionError> {
257 if child_session_ids.is_empty() {
258 return Ok(0);
259 }
260 let batch: Vec<(String, Option<String>)> = child_session_ids
261 .iter()
262 .map(|id| (id.clone(), None))
263 .collect();
264 self.flush_parent_waits(parent_session_id, &batch, policy)
265 .await?;
266 Ok(batch.len())
267 }
268
269 pub async fn active_child_ids(&self, parent_session_id: &str) -> Vec<String> {
272 self.storage
273 .list_child_run_statuses(parent_session_id)
274 .await
275 .unwrap_or_default()
276 .into_iter()
277 .filter(|(_, status)| !status.as_deref().is_some_and(is_terminal_child_status))
278 .map(|(id, _)| id)
279 .collect()
280 }
281
282 async fn flush_parent_waits(
284 &self,
285 parent_session_id: &str,
286 batch: &[(String, Option<String>)],
287 policy: ChildWaitPolicy,
288 ) -> Result<(), ChildSessionError> {
289 let Some(mut parent) =
290 self.storage
291 .load_session(parent_session_id)
292 .await
293 .map_err(|error| {
294 ChildSessionError::Execution(format!(
295 "failed to load parent session {parent_session_id}: {error}"
296 ))
297 })?
298 else {
299 return Err(ChildSessionError::NotFound(parent_session_id.to_string()));
300 };
301
302 let mut runtime_state = read_runtime_state(&parent);
306
307 let now = Utc::now();
308 let mut wait = runtime_state
309 .waiting_for_children
310 .take()
311 .unwrap_or_else(|| WaitingForChildrenState::for_children(Vec::new(), policy, now));
312 wait.wait_for = policy;
314 for (child_session_id, tool_call_id) in batch {
315 if !wait
316 .child_session_ids
317 .iter()
318 .any(|id| id == child_session_id)
319 {
320 wait.child_session_ids.push(child_session_id.clone());
321 }
322 if wait.registered_by_tool_call_id.is_none() {
323 wait.registered_by_tool_call_id = tool_call_id.clone();
324 }
325 }
326 wait.child_session_ids.sort();
327 wait.child_session_ids.dedup();
328 runtime_state.waiting_for_children = Some(wait);
329
330 write_runtime_state(&mut parent, &runtime_state);
331 parent.metadata.insert(
332 "runtime.suspend_reason".to_string(),
333 "waiting_for_children".to_string(),
334 );
335 parent.updated_at = Utc::now();
336
337 self.persistence
342 .save_runtime_only(&mut parent)
343 .await
344 .map_err(|error| {
345 ChildSessionError::Execution(format!("failed to save parent wait state: {error}"))
346 })?;
347 self.sessions_cache.insert(
348 parent.id.clone(),
349 Arc::new(parking_lot::RwLock::new(parent)),
350 );
351
352 Ok(())
353 }
354}
355
356fn map_index_entry_to_child_entry(entry: &SessionIndexEntry) -> ChildSessionEntry {
357 ChildSessionEntry {
358 child_session_id: entry.id.clone(),
359 title: entry.title.clone(),
360 pinned: entry.pinned,
361 message_count: entry.message_count,
362 updated_at: entry.updated_at.to_rfc3339(),
363 last_run_status: entry.last_run_status.clone(),
364 last_run_error: entry.last_run_error.clone(),
365 }
366}
367
368#[async_trait]
369impl SubagentResolutionPort for ChildSessionAdapter {
370 async fn resolve_subagent_model(
371 &self,
372 subagent_type: &str,
373 ) -> Option<bamboo_domain::ProviderModelRef> {
374 ChildSessionAdapter::resolve_subagent_model(self, subagent_type).await
375 }
376
377 async fn resolve_runtime_metadata(
378 &self,
379 subagent_type: &str,
380 ) -> std::collections::HashMap<String, String> {
381 ChildSessionAdapter::resolve_runtime_metadata(self, subagent_type).await
382 }
383}
384
385#[async_trait]
393impl bamboo_engine::GuardianSpawner for ChildSessionAdapter {
394 async fn spawn_guardian_review(
395 &self,
396 parent_session: &Session,
397 review_prompt: String,
398 model: String,
399 disabled_tools: Option<std::collections::BTreeSet<String>>,
400 ) -> Result<String, String> {
401 let input = bamboo_engine::session_app::child_session::CreateChildInput {
402 parent_session: parent_session.clone(),
403 child_id: format!("guardian-{}", uuid::Uuid::new_v4()),
404 title: "Guardian review".to_string(),
405 responsibility: "Adversarially verify the parent agent's completed work.".to_string(),
406 assignment_prompt: review_prompt,
407 subagent_type: "guardian".to_string(),
410 workspace: parent_session.workspace_path_meta().unwrap_or_default(),
411 model_override: Some(model),
412 model_ref_override: None,
413 runtime_metadata: HashMap::new(),
414 auto_run: true,
415 reasoning_effort: None,
416 lifecycle: None,
417 resident_name: None,
418 resident_context: None,
419 disabled_tools,
420 context_fork: None,
421 };
422 bamboo_engine::session_app::child_session::create_child_action(self, input)
423 .await
424 .map(|result| result.child_session_id)
425 .map_err(|error| error.to_string())
426 }
427}
428
429#[async_trait]
430impl ChildSessionPort for ChildSessionAdapter {
431 async fn load_root_session(&self, root_session_id: &str) -> Result<Session, ChildSessionError> {
432 let Some(session) = self
433 .storage
434 .load_session(root_session_id)
435 .await
436 .map_err(|error| {
437 ChildSessionError::Execution(format!(
438 "failed to load session {root_session_id}: {error}"
439 ))
440 })?
441 else {
442 return Err(ChildSessionError::NotFound(root_session_id.to_string()));
443 };
444
445 if session.kind != SessionKind::Root {
446 return Err(ChildSessionError::NotRootSession(
447 root_session_id.to_string(),
448 ));
449 }
450
451 Ok(session)
452 }
453
454 async fn load_child_for_parent(
455 &self,
456 parent_session_id: &str,
457 child_session_id: &str,
458 ) -> Result<Session, ChildSessionError> {
459 let Some(child) = self
460 .storage
461 .load_session(child_session_id)
462 .await
463 .map_err(|error| {
464 ChildSessionError::Execution(format!(
465 "failed to load child session {child_session_id}: {error}"
466 ))
467 })?
468 else {
469 return Err(ChildSessionError::NotFound(child_session_id.to_string()));
470 };
471
472 if child.kind != SessionKind::Child {
473 return Err(ChildSessionError::NotChildSession(
474 child_session_id.to_string(),
475 ));
476 }
477
478 if child.parent_session_id.as_deref() != Some(parent_session_id) {
479 return Err(ChildSessionError::NotChildOfParent {
480 child_id: child_session_id.to_string(),
481 parent_id: parent_session_id.to_string(),
482 });
483 }
484
485 Ok(child)
486 }
487
488 async fn save_child_session(&self, child: &mut Session) -> Result<(), ChildSessionError> {
489 self.persistence
490 .merge_save_runtime(child)
491 .await
492 .map_err(|error| {
493 ChildSessionError::Execution(format!("failed to save child session: {error}"))
494 })?;
495
496 self.sessions_cache.insert(
497 child.id.clone(),
498 Arc::new(parking_lot::RwLock::new(child.clone())),
499 );
500
501 Ok(())
502 }
503
504 async fn is_child_running(&self, child_session_id: &str) -> bool {
505 let runners = self.agent_runners.read().await;
506 runners
507 .get(child_session_id)
508 .is_some_and(|runner| matches!(runner.status, AgentStatus::Running))
509 }
510
511 async fn list_children(&self, parent_session_id: &str) -> Vec<ChildSessionEntry> {
512 self.session_store
513 .list_index_entries()
514 .await
515 .into_iter()
516 .filter(|entry| {
517 entry.kind == SessionKind::Child
518 && entry.parent_session_id.as_deref() == Some(parent_session_id)
519 })
520 .map(|entry| map_index_entry_to_child_entry(&entry))
521 .collect()
522 }
523
524 async fn find_resident_child(
525 &self,
526 root_session_id: &str,
527 resident_name: &str,
528 ) -> Option<String> {
529 let name = resident_name.trim();
530 if name.is_empty() {
531 return None;
532 }
533 let mut best: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
537 for entry in self.session_store.list_index_entries().await {
538 if entry.kind == SessionKind::Child
539 && entry.root_session_id == root_session_id
540 && entry.resident_name.as_deref() == Some(name)
541 {
542 match &best {
543 Some((_, ts)) if *ts >= entry.updated_at => {}
544 _ => best = Some((entry.id.clone(), entry.updated_at)),
545 }
546 }
547 }
548 best.map(|(id, _)| id)
549 }
550
551 async fn enqueue_child_run(
552 &self,
553 parent: &Session,
554 child: &Session,
555 ) -> Result<(), ChildSessionError> {
556 let model = if child.model.trim().is_empty() {
557 parent.model.clone()
558 } else {
559 child.model.clone()
560 };
561 if model.trim().is_empty() {
562 return Err(ChildSessionError::Execution(
563 "child model is empty and parent model is unavailable".to_string(),
564 ));
565 }
566
567 let disabled_tools = child
573 .metadata
574 .get("disabled_tools")
575 .and_then(|raw| serde_json::from_str::<std::collections::BTreeSet<String>>(raw).ok())
576 .filter(|set| !set.is_empty())
577 .map(|set| set.into_iter().collect::<Vec<String>>());
578
579 if let Some(relay) = &self.notification_relay {
587 let child_tx = get_or_create_event_sender(&self.session_event_senders, &child.id).await;
588 crate::app_state::session_events::ensure_notification_relay(relay, &child.id, child_tx);
589 }
590
591 self.scheduler
597 .enqueue(SpawnJob {
598 parent_session_id: parent.id.clone(),
599 child_session_id: child.id.clone(),
600 model,
601 disabled_tools,
602 })
603 .await
604 .map_err(ChildSessionError::Execution)?;
605
606 let parent_tx = get_or_create_event_sender(&self.session_event_senders, &parent.id).await;
607 let _ = parent_tx.send(AgentEvent::SubAgentStarted {
608 parent_session_id: parent.id.clone(),
609 child_session_id: child.id.clone(),
610 title: Some(child.title.clone()),
611 });
612
613 Ok(())
614 }
615
616 async fn cancel_child_run_and_wait(
617 &self,
618 child_session_id: &str,
619 ) -> Result<(), ChildSessionError> {
620 let cancelled = {
621 let mut runners = self.agent_runners.write().await;
622 if let Some(runner) = runners.get_mut(child_session_id) {
623 if matches!(runner.status, AgentStatus::Running) {
624 runner.cancel_token.cancel();
625 true
626 } else {
627 false
628 }
629 } else {
630 false
631 }
632 };
633
634 if !cancelled {
635 return Ok(());
636 }
637
638 let deadline = Instant::now() + Duration::from_secs(10);
639 loop {
640 let still_running = {
641 let runners = self.agent_runners.read().await;
642 runners
643 .get(child_session_id)
644 .is_some_and(|runner| matches!(runner.status, AgentStatus::Running))
645 };
646 if !still_running {
647 return Ok(());
648 }
649 if Instant::now() >= deadline {
650 return Err(ChildSessionError::Execution(format!(
651 "timed out waiting for child session {child_session_id} to stop after cancellation"
652 )));
653 }
654 sleep(Duration::from_millis(50)).await;
655 }
656 }
657
658 async fn delete_child_session(
659 &self,
660 parent_session_id: &str,
661 child_id: &str,
662 ) -> Result<DeleteChildResult, ChildSessionError> {
663 let cancelled_running_child = {
664 let mut runners = self.agent_runners.write().await;
665 if let Some(runner) = runners.remove(child_id) {
666 runner.cancel_token.cancel();
667 true
668 } else {
669 false
670 }
671 };
672
673 let deleted = self
674 .storage
675 .delete_session(child_id)
676 .await
677 .map_err(|error| {
678 ChildSessionError::Execution(format!("failed to delete child session: {error}"))
679 })?;
680
681 self.sessions_cache.remove(child_id);
682 {
683 let mut senders = self.session_event_senders.write().await;
684 senders.remove(child_id);
685 if cancelled_running_child {
686 if let Some(parent_tx) = senders.get(parent_session_id) {
687 let _ = parent_tx.send(AgentEvent::SubAgentCompleted {
688 parent_session_id: parent_session_id.to_string(),
689 child_session_id: child_id.to_string(),
690 status: "cancelled".to_string(),
691 error: Some("Child session deleted while running".to_string()),
692 });
693 }
694 }
695 }
696
697 Ok(DeleteChildResult {
698 deleted,
699 cancelled_running_child,
700 })
701 }
702
703 async fn get_child_runner_info(&self, child_id: &str) -> Option<ChildRunnerInfo> {
704 let runners = self.agent_runners.read().await;
705 runners.get(child_id).map(|runner| ChildRunnerInfo {
706 started_at: Some(runner.started_at),
707 completed_at: runner.completed_at,
708 last_tool_name: runner.last_tool_name.clone(),
709 last_tool_phase: runner.last_tool_phase.clone(),
710 last_event_at: runner.last_event_at,
711 round_count: runner.round_count,
712 })
713 }
714
715 async fn register_parent_wait_for_child(
716 &self,
717 parent_session_id: &str,
718 child_session_id: &str,
719 tool_call_id: Option<&str>,
720 ) -> Result<(), ChildSessionError> {
721 ChildSessionAdapter::register_parent_wait_for_child(
722 self,
723 parent_session_id,
724 child_session_id,
725 tool_call_id,
726 )
727 .await
728 }
729
730 async fn register_parent_wait_for_children(
731 &self,
732 parent_session_id: &str,
733 child_session_ids: &[String],
734 policy: ChildWaitPolicy,
735 ) -> Result<usize, ChildSessionError> {
736 ChildSessionAdapter::register_parent_wait_for_children(
737 self,
738 parent_session_id,
739 child_session_ids,
740 policy,
741 )
742 .await
743 }
744
745 async fn active_child_ids(&self, parent_session_id: &str) -> Vec<String> {
746 ChildSessionAdapter::active_child_ids(self, parent_session_id).await
747 }
748
749 async fn ensure_child_indexed(&self, child_session_id: &str) {
750 let _ = self.session_store.get_index_entry(child_session_id).await;
751 }
752}