Skip to main content

bamboo_server/tools/
child_session_adapter.rs

1//! Shared adapter implementing `ChildSessionPort` for server-side child session tools.
2//!
3//! The unified `SubAgentTool` delegates to this adapter instead of
4//! duplicating `ChildSessionPort` implementations.
5
6use 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
29/// Server-side adapter that bridges domain `ChildSessionPort` to infrastructure.
30///
31/// Holds all shared state needed by `SubAgentTool`.
32/// Implements the full `ChildSessionPort` trait with real methods (no stubs).
33pub 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    /// Optional subagent model resolver: maps subagent_type → provider+model ref.
42    pub(crate) subagent_model_resolver: crate::tools::OptionalSubagentModelResolver,
43    /// Application config for resolving subagent routing and external agent profiles.
44    pub(crate) config: Arc<RwLock<Config>>,
45    /// Coalesces concurrent parent-wait registrations for the same parent that
46    /// arrive in one spawn round (the LLM emitting several `SubAgent.create`
47    /// calls at once → `join_all`) into a single parent persist. See
48    /// [`ChildSessionAdapter::register_parent_wait_for_child`].
49    pub(crate) parent_wait_slots: Arc<dashmap::DashMap<String, Arc<ParentWaitSlot>>>,
50    /// Deps to start the always-on notification relay for a NEWLY ENQUEUED
51    /// child session (see [`ChildSessionAdapter::enqueue_child_run`]).
52    ///
53    /// Without this, a child's own events (e.g. a `run_in_background` Bash
54    /// command finishing, or the child hitting critical context pressure)
55    /// only ever get classified if a client happens to be subscribed to that
56    /// specific child session's SSE/WS stream — the child's completion
57    /// (`SubAgentCompleted`) still reaches the owner via the PARENT's own
58    /// relay (already running for any in-process parent execution), but
59    /// events that occur only on the child's own stream would otherwise be
60    /// silently dropped for a headless/unwatched child. `None` for the
61    /// out-of-process worker binary (`ChildSessionAdapter::new`), which has
62    /// no local desktop/ntfy/bark config surface to deliver through.
63    pub(crate) notification_relay: Option<crate::app_state::session_events::NotificationRelayDeps>,
64}
65
66/// Per-parent coalescing slot for batched wait registration.
67///
68/// `flush_lock` is a barrier distinct from the persistence per-session lock
69/// (using the latter here would deadlock, since the flush itself takes it). The
70/// first registration to win the barrier drains `pending` and persists the whole
71/// batch once; concurrent registrations that find `pending` already drained were
72/// persisted by that holder before it released the barrier, so they return
73/// without an extra write.
74#[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
82/// Terminal child run statuses, as mirrored into the session index. A child not
83/// in one of these states is considered active (still pending/running).
84fn 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    /// Construct an adapter. Public so a self-orchestrating WORKER (Phase 6:
115    /// direct nested execution) can build its OWN child-session machinery
116    /// against its own store/scheduler — the struct fields are `pub(crate)`, so
117    /// out-of-crate callers (the worker binary) go through this constructor.
118    #[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            // Fresh per-adapter wait-coalescing map (the type is private to this
141            // crate, so out-of-crate callers can't supply it).
142            parent_wait_slots: Arc::new(dashmap::DashMap::new()),
143            // The worker binary runs on a different machine than the desktop
144            // it would notify; it reports events back to the orchestrating
145            // server over the wire instead. See the field doc.
146            notification_relay: None,
147        }
148    }
149
150    /// Resolve the provider+model ref for a given subagent_type using the configured resolver.
151    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    /// Resolve runtime metadata (e.g. external agent routing) for a subagent_type.
162    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    /// Register a durable parent wait for an enqueued child session.
168    ///
169    /// This is intentionally idempotent: repeated registrations for the same
170    /// child merge into the existing wait set. The child runner owns timeout
171    /// and liveness; the parent wait timeout is a long lease for observability.
172    ///
173    /// Registrations are **coalesced** per parent: when several children are
174    /// spawned in one round (the LLM issuing multiple `SubAgent.create` calls
175    /// that `join_all` runs concurrently), the first call to win the per-parent
176    /// barrier drains all currently-pending registrations and persists the parent
177    /// once, instead of each child triggering its own load+write. Callers whose
178    /// child was drained-and-persisted by that holder return without an extra
179    /// write — and only after the holder's write committed, so durability holds.
180    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        // 1. Enqueue this registration.
193        slot.pending.lock().push((
194            child_session_id.to_string(),
195            tool_call_id.map(str::to_string),
196        ));
197
198        // 2. Barrier: serialize flushers for this parent.
199        let _flush_guard = slot.flush_lock.lock().await;
200
201        // 3. Drain everything pending for this parent (siblings that enqueued
202        //    while we waited for the barrier are picked up here too).
203        let batch: Vec<(String, Option<String>)> = {
204            let mut pending = slot.pending.lock();
205            pending.drain(..).collect()
206        };
207        if batch.is_empty() {
208            // A prior barrier holder already persisted our child before releasing
209            // the barrier we just acquired — nothing left to write.
210            return Ok(());
211        }
212
213        // 4. Persist the whole batch in a single parent write.
214        if let Err(error) = self
215            .flush_parent_waits(parent_session_id, &batch, ChildWaitPolicy::All)
216            .await
217        {
218            // Re-queue so nothing is silently lost; a retry or sibling picks it up.
219            let mut pending = slot.pending.lock();
220            for item in batch {
221                pending.push(item);
222            }
223            return Err(error);
224        }
225
226        // 5. Self-clean: the slot exists only to coalesce a burst of sibling
227        //    registrations for THIS parent. Now that the batch is durably
228        //    persisted and nothing new is pending, drop the map entry so
229        //    `parent_wait_slots` does not retain one entry per parent-that-ever-
230        //    -spawned forever (issue #346). Still inside the flush barrier.
231        //
232        //    Race-freedom: `remove_if` re-checks `pending.is_empty()` under the
233        //    DashMap shard lock. A sibling that enqueued after our drain made
234        //    `pending` non-empty, so the predicate is false and we keep the slot;
235        //    that sibling (blocked on the barrier we still hold) will flush it and
236        //    run this same removal. A sibling that clones the slot Arc but has not
237        //    yet pushed keeps a live handle, so removing the map entry never loses
238        //    its child: whoever holds the barrier drains ALL pending on its Arc.
239        self.parent_wait_slots
240            .remove_if(parent_session_id, |_, slot| slot.pending.lock().is_empty());
241
242        Ok(())
243    }
244
245    /// Explicitly register a parent wait for an arbitrary set of children with a
246    /// chosen policy. Used by the `SubAgent.wait` action (wait on all active
247    /// children) and the end-of-turn safety net. A single parent write.
248    ///
249    /// Returns the number of children the wait now covers (0 means there was
250    /// nothing to wait on and no wait was registered).
251    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    /// The parent's currently-active (non-terminal) children, derived from the
270    /// session index (single source of truth).
271    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    /// Persist a batch of parent-wait registrations in one runtime-only save.
283    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        // The active/completed child sets are derived from the session index
303        // (single source of truth), so we no longer maintain a denormalized copy
304        // here. Only the durable wait state below is parent-owned.
305        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        // An explicit wait re-asserts the policy on any pre-existing wait state.
313        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        // Runtime-only save: registering a parent's wait mutates the
338        // control-plane (runtime_state + suspend metadata) but NEVER the message
339        // history. Writing just the sidecar keeps spawn O(1) in conversation
340        // length instead of rewriting the parent's full session.json per child.
341        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/// Lets a [`ChildSessionAdapter`] act as the engine's guardian-review spawner.
386///
387/// `Arc<ChildSessionAdapter>` therefore doubles as `Arc<dyn GuardianSpawner>`
388/// (wired onto `AppState`), so the terminal gate spawns the read-only reviewer
389/// through the same child-session machinery the `SubAgent` tool uses — no second
390/// spawn path. The reviewer is a real sub-agent: it fetches the diff and runs
391/// tests itself via its (read-only) toolset.
392#[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            // The coordinator branches on this subagent_type to recognize a
408            // guardian completion and parse its verdict.
409            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        // Scan the index for a child in this root tree tagged with the resident
534        // name. Prefer the most recently updated if (defensively) more than one
535        // exists. Index-backed: no session.json loads.
536        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        // Per-child tool denylist: persisted onto the child session by
568        // `create_child_action` (JSON in metadata). Most sub-agents are full
569        // agents and carry none; a read-only Guardian reviewer carries a
570        // denylist here so the worker trims its toolset. `SpawnJob` wants a
571        // `Vec<String>`, so collect the set.
572        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        // Start the always-on notification relay for the CHILD's own session
580        // before enqueueing the job — mirrors the execute handler starting it
581        // at execution entry (`spawn_event_forwarder`) rather than waiting
582        // for a client to subscribe. Idempotent (`try_begin_relay`) and races
583        // harmlessly with the engine's own `get_or_create_event_sender` call
584        // for the same id in `run_child_spawn`, since both resolve to the
585        // same map entry. See the `notification_relay` field doc.
586        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        // NOTE: enqueue only *runs* the child in the background. Registering the
592        // parent's wait (which suspends the parent) is now an explicit, separate
593        // step so the model can spawn several children without each one
594        // suspending it — see `register_parent_wait_for_child` /
595        // `register_parent_wait_for_children` and the `SubAgent.wait` action.
596        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}