bamboo_engine/runtime/execution/
spawn.rs

1//! Sub-session spawn scheduler.
2//!
3//! Provides a background queue for spawning child sessions. Spawn is async
4//! (tool returns immediately), but the UI can observe child progress via
5//! events forwarded to the parent session stream.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9use std::time::Duration;
10
11use chrono::Utc;
12use tokio::sync::{broadcast, mpsc, RwLock};
13use tokio_util::sync::CancellationToken;
14
15use bamboo_agent_core::tools::ToolExecutor;
16use bamboo_agent_core::{AgentEvent, Session};
17use bamboo_llm::ProviderModelRouter;
18
19use crate::runtime::Agent;
20
21use super::child_completion::{ChildCompletion, ChildCompletionHandler};
22use super::runner_state::AgentRunner;
23
24#[derive(Debug, Clone)]
25pub struct SpawnJob {
26    pub parent_session_id: String,
27    pub child_session_id: String,
28    pub model: String,
29    /// Tool names to hide from the LLM schema for this child session.
30    /// Computed from the child's `subagent_type` profile policy.
31    pub disabled_tools: Option<Vec<String>>,
32}
33
34/// Trait for external child session runtimes (e.g. A2A, CLI adapters).
35///
36/// Implementors are responsible for emitting AgentEvents via `event_tx`
37/// and respecting the `cancel_token`.
38#[async_trait::async_trait]
39pub trait ExternalChildRunner: Send + Sync {
40    /// Returns true if this runner should handle the given child session.
41    async fn should_handle(&self, session: &Session) -> bool;
42
43    /// Execute the child session using an external runtime.
44    async fn execute_external_child(
45        &self,
46        session: &mut Session,
47        job: &SpawnJob,
48        event_tx: tokio::sync::mpsc::Sender<AgentEvent>,
49        cancel_token: CancellationToken,
50    ) -> crate::runtime::runner::Result<()>;
51
52    /// Bind this runner's per-run escalation host bridge (#68). A nested worker's
53    /// `run()` installs its OWN host bridge here so the runner can hand it to each
54    /// grandchild's `drive()` AT SPAWN time (captured into the drive task, not read
55    /// later), letting the grandchild re-proxy a non-bypass approval request UP to
56    /// its parent run for its whole lifetime — even when it outlives the run that
57    /// spawned it. Default no-op for runners that don't escalate (e.g. A2A).
58    fn set_escalation_bridge(&self, _bridge: Option<bamboo_subagent::executor::HostBridge>) {}
59}
60
61#[derive(Clone)]
62pub struct SpawnContext {
63    pub agent: Arc<Agent>,
64    pub tools: Arc<dyn ToolExecutor>,
65    pub sessions_cache: crate::SessionCache,
66    pub agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
67    pub session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
68    pub external_child_runner: Arc<dyn ExternalChildRunner>,
69    pub provider_router: Option<Arc<ProviderModelRouter>>,
70    pub app_data_dir: Option<std::path::PathBuf>,
71    /// Optional application-layer completion hook. The engine still emits
72    /// `SubAgentCompleted` to the parent stream itself; this hook lets the
73    /// server persist parent wait state and resume the parent runner without
74    /// introducing an engine -> AppState dependency.
75    pub completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
76    /// Optional inbox to the account-wide change feed. When present, durable
77    /// change events from child-session execution are mirrored onto the feed
78    /// for resumable multi-client sync.
79    pub account_feed_inbox: Option<super::event_forwarder::AccountFeedInbox>,
80}
81
82#[derive(Clone)]
83pub struct SpawnScheduler {
84    tx: mpsc::Sender<SpawnJob>,
85}
86
87impl SpawnScheduler {
88    pub fn new(ctx: SpawnContext) -> Self {
89        let (tx, mut rx) = mpsc::channel::<SpawnJob>(128);
90
91        // The worker loop is a single point of failure for ALL child spawning:
92        // if it unwinds, queued jobs are dropped with no completion published
93        // and every later enqueue fails with "spawn scheduler is not running"
94        // for the rest of the process lifetime. Run each job on its own task
95        // and await the JoinHandle — a panicking job is isolated, keeps the
96        // worker alive, and still publishes a terminal error completion so the
97        // waiting parent is woken instead of stranded.
98        tokio::spawn(async move {
99            while let Some(job) = rx.recv().await {
100                let job_ctx = ctx.clone();
101                let job_for_panic = job.clone();
102                let handle = tokio::spawn(async move {
103                    if let Err(err) = run_spawn_job(job_ctx, job).await {
104                        tracing::warn!("spawn job failed: {}", err);
105                    }
106                });
107                if let Err(join_error) = handle.await {
108                    tracing::error!(
109                        parent_session_id = %job_for_panic.parent_session_id,
110                        child_session_id = %job_for_panic.child_session_id,
111                        error = %join_error,
112                        "spawn job panicked; publishing terminal error completion"
113                    );
114                    let parent_tx = super::session_events::get_or_create_event_sender(
115                        &ctx.session_event_senders,
116                        &job_for_panic.parent_session_id,
117                    )
118                    .await;
119                    publish_child_completion_parts(
120                        &parent_tx,
121                        ctx.completion_handler.clone(),
122                        job_for_panic.parent_session_id.clone(),
123                        job_for_panic.child_session_id.clone(),
124                        "error".to_string(),
125                        Some(format!("child spawn panicked: {join_error}")),
126                    )
127                    .await;
128                }
129            }
130        });
131
132        Self { tx }
133    }
134
135    pub async fn enqueue(&self, job: SpawnJob) -> Result<(), String> {
136        self.tx
137            .send(job)
138            .await
139            .map_err(|_| "spawn scheduler is not running".to_string())
140    }
141}
142
143#[derive(Debug, Clone, Copy)]
144pub(crate) struct ChildWatchdogPolicy {
145    check_interval_secs: i64,
146    // pub(crate): the child-wait watchdog (#546) reads these limits to decide
147    // when a Running runner entry whose task died (frozen last_event_at) is
148    // stale beyond what the per-child liveness watchdog could still act on.
149    pub(crate) max_total_secs: i64,
150    pub(crate) max_idle_secs: i64,
151}
152
153impl Default for ChildWatchdogPolicy {
154    fn default() -> Self {
155        Self {
156            check_interval_secs: 15,
157            // Parent waits may be longer, but child execution owns its own
158            // liveness. A one hour total cap avoids indefinitely orphaned
159            // sub-session runners.
160            max_total_secs: 60 * 60,
161            // No child event for 15 minutes is considered stalled.
162            max_idle_secs: 15 * 60,
163        }
164    }
165}
166
167fn metadata_i64(session: &Session, key: &str) -> Option<i64> {
168    session
169        .metadata
170        .get(key)
171        .and_then(|value| value.trim().parse::<i64>().ok())
172        .filter(|value| *value > 0)
173}
174
175pub(crate) fn watchdog_policy_for_session(session: &Session) -> ChildWatchdogPolicy {
176    let mut policy = ChildWatchdogPolicy::default();
177    if let Some(value) = metadata_i64(session, "child_watchdog.max_total_secs") {
178        policy.max_total_secs = value;
179    }
180    if let Some(value) = metadata_i64(session, "child_watchdog.max_idle_secs") {
181        policy.max_idle_secs = value;
182    }
183    if let Some(value) = metadata_i64(session, "child_watchdog.check_interval_secs") {
184        policy.check_interval_secs = value;
185    }
186    policy
187}
188
189async fn publish_child_completion(
190    parent_tx: &broadcast::Sender<AgentEvent>,
191    completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
192    completion: ChildCompletion,
193) {
194    let _ = parent_tx.send(AgentEvent::SubAgentCompleted {
195        parent_session_id: completion.parent_session_id.clone(),
196        child_session_id: completion.child_session_id.clone(),
197        status: completion.status.clone(),
198        error: completion.error.clone(),
199    });
200
201    if let Some(handler) = completion_handler {
202        // Contain a panicking handler: this call frequently runs on the caller's
203        // only liveness-critical task (the child's terminal block, or the spawn
204        // scheduler worker for early failures). Unwinding here would kill that
205        // task after the child already looks terminal everywhere — the classic
206        // stranded-parent signature. The child-wait watchdog backstops the wake
207        // that a panicked handler failed to deliver.
208        use futures::FutureExt;
209        let parent_session_id = completion.parent_session_id.clone();
210        let child_session_id = completion.child_session_id.clone();
211        if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
212            .catch_unwind()
213            .await
214            .is_err()
215        {
216            tracing::error!(
217                %parent_session_id,
218                %child_session_id,
219                "child completion handler panicked; child-wait watchdog will backstop the parent wake"
220            );
221        }
222    }
223}
224
225pub(crate) async fn publish_child_completion_parts(
226    parent_tx: &broadcast::Sender<AgentEvent>,
227    completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
228    parent_session_id: String,
229    child_session_id: String,
230    status: String,
231    error: Option<String>,
232) {
233    publish_child_completion(
234        parent_tx,
235        completion_handler,
236        ChildCompletion {
237            parent_session_id,
238            child_session_id,
239            status,
240            error,
241            completed_at: Utc::now(),
242        },
243    )
244    .await;
245}
246
247pub(crate) async fn watch_child_liveness(
248    parent_session_id: String,
249    child_session_id: String,
250    runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
251    cancel_token: CancellationToken,
252    timeout_reason: Arc<RwLock<Option<String>>>,
253    done: CancellationToken,
254    policy: ChildWatchdogPolicy,
255) {
256    let mut ticker =
257        tokio::time::interval(Duration::from_secs(policy.check_interval_secs.max(1) as u64));
258    // Skip the immediate tick.
259    ticker.tick().await;
260
261    loop {
262        tokio::select! {
263            _ = done.cancelled() => return,
264            _ = ticker.tick() => {
265                if cancel_token.is_cancelled() {
266                    return;
267                }
268
269                let snapshot = {
270                    let guard = runners.read().await;
271                    guard.get(&child_session_id).cloned()
272                };
273                let Some(runner) = snapshot else {
274                    return;
275                };
276                if !matches!(runner.status, super::runner_state::AgentStatus::Running) {
277                    return;
278                }
279
280                let now = Utc::now();
281                let total_secs = now.signed_duration_since(runner.started_at).num_seconds();
282                if total_secs >= policy.max_total_secs {
283                    let reason = format!(
284                        "Child session timed out after {} seconds (max_total_secs={})",
285                        total_secs, policy.max_total_secs
286                    );
287                    tracing::warn!(
288                        parent_session_id = %parent_session_id,
289                        child_session_id = %child_session_id,
290                        reason = %reason,
291                        "child session total timeout; cancelling child runner"
292                    );
293                    *timeout_reason.write().await = Some(reason);
294                    cancel_token.cancel();
295                    return;
296                }
297
298                let last_activity_at = runner.last_event_at.unwrap_or(runner.started_at);
299                let idle_secs = now.signed_duration_since(last_activity_at).num_seconds();
300                if idle_secs >= policy.max_idle_secs {
301                    let reason = format!(
302                        "Child session idle timeout after {} seconds without events (max_idle_secs={})",
303                        idle_secs, policy.max_idle_secs
304                    );
305                    tracing::warn!(
306                        parent_session_id = %parent_session_id,
307                        child_session_id = %child_session_id,
308                        reason = %reason,
309                        last_tool_name = ?runner.last_tool_name,
310                        last_tool_phase = ?runner.last_tool_phase,
311                        round_count = runner.round_count,
312                        "child session idle timeout; cancelling child runner"
313                    );
314                    *timeout_reason.write().await = Some(reason);
315                    cancel_token.cancel();
316                    return;
317                }
318            }
319        }
320    }
321}
322
323/// Drive a single queued spawn job through the canonical child-spawn path.
324///
325/// ANTI-FORK: this is a 1-line delegator to [`crate::sdk::spawn::run_child_spawn`],
326/// which is the single implementation of the spawn/execute/finalize logic. The
327/// `SpawnScheduler` queue mechanics (above) remain here; the body lives in the SDK
328/// core so both the scheduler and the ergonomic `ChildRunner` funnel into it.
329async fn run_spawn_job(ctx: SpawnContext, job: SpawnJob) -> Result<(), String> {
330    crate::sdk::spawn::run_child_spawn(ctx, job).await
331}
bamboo_engine/runtime/execution/spawn.rs

bamboo_engine/runtime/execution/
spawn.rs
