1use std::collections::HashMap;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10use std::time::Duration;
11
12use chrono::Utc;
13use dashmap::DashMap;
14use futures::future::BoxFuture;
15use tokio::sync::{broadcast, mpsc, oneshot, Mutex, OwnedMutexGuard, RwLock};
16use tokio_util::sync::CancellationToken;
17
18use bamboo_agent_core::storage::Storage;
19use bamboo_agent_core::tools::ToolExecutor;
20use bamboo_agent_core::{AgentEvent, Session};
21use bamboo_domain::{AdmissionCommit, AdmissionGate, RuntimeSessionPersistence, SessionInboxPort};
22use bamboo_llm::ProviderModelRouter;
23
24use crate::runtime::Agent;
25
26use super::agent_spawn::SessionExecutionReservation;
27use super::child_completion::{ChildCompletion, ChildCompletionHandler};
28use super::runner_state::AgentRunner;
29
30pub type QueueAdmissionCheck = dyn Fn() -> Result<(), String> + Send + Sync;
34
35pub struct QueueCapacityReservation<'a> {
39 scheduler: &'a SpawnScheduler,
40 gate: Option<&'a AdmissionGate>,
41 permit: mpsc::Permit<'a, QueuedSpawnJob>,
42}
43
44impl QueueCapacityReservation<'_> {
45 pub fn terminal_state(&self) -> Option<AdmissionCommit<()>> {
48 self.gate.and_then(|gate| {
49 if gate.is_committed() {
50 Some(AdmissionCommit::AlreadyCommitted)
51 } else if gate.is_cancelled() {
52 Some(AdmissionCommit::Cancelled)
53 } else {
54 None
55 }
56 })
57 }
58
59 pub async fn enqueue_announced_for_generation_revalidated(
63 self,
64 job: SpawnJob,
65 title: Option<String>,
66 launch_generation: Option<u64>,
67 admission_check: Option<&QueueAdmissionCheck>,
68 preflight: Result<(), String>,
69 ) -> Result<AdmissionCommit<()>, String> {
70 let Self {
71 scheduler,
72 gate,
73 permit,
74 } = self;
75 let preparation_job = job.clone();
76 let commit = send_reserved(
77 permit,
78 job,
79 title.clone(),
80 gate,
81 launch_generation,
82 admission_check,
83 preflight,
84 )?;
85 scheduler
86 .finish_admission(preparation_job, title, launch_generation, commit)
87 .await
88 }
89}
90
91#[derive(Debug, Clone)]
92pub struct SpawnJob {
93 pub parent_session_id: String,
94 pub child_session_id: String,
95 pub model: String,
96 pub disabled_tools: Option<Vec<String>>,
99}
100
101struct QueuedSpawnJob {
102 job: SpawnJob,
103 title: Option<String>,
104 launch_generation: Option<u64>,
105 announcement_ready: oneshot::Receiver<()>,
106 start_published: Arc<AtomicBool>,
107}
108
109#[derive(Clone, Default)]
112struct ChildLaunchLocks {
113 locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
114}
115
116pub struct ChildLaunchGuard {
117 guard: Option<OwnedMutexGuard<()>>,
118 locks: ChildLaunchLocks,
119 child_id: String,
120}
121
122impl Drop for ChildLaunchGuard {
123 fn drop(&mut self) {
124 self.guard.take();
125 self.locks
126 .locks
127 .remove_if(&self.child_id, |_, lock| Arc::strong_count(lock) == 1);
128 }
129}
130
131impl ChildLaunchLocks {
132 async fn acquire(&self, child_id: &str) -> ChildLaunchGuard {
133 let lock = self
134 .locks
135 .entry(child_id.to_string())
136 .or_insert_with(|| Arc::new(Mutex::new(())))
137 .clone();
138 let mut guard = ChildLaunchGuard {
139 guard: None,
140 locks: self.clone(),
141 child_id: child_id.to_string(),
142 };
143 guard.guard = Some(lock.lock_owned().await);
144 guard
145 }
146}
147
148pub trait ChildRunLaunchHook: Send + Sync {
156 fn before_child_launch(&self, job: &SpawnJob, child_events: broadcast::Sender<AgentEvent>);
157
158 fn before_queued_dequeue(&self, _job: &SpawnJob) -> BoxFuture<'static, ()> {
162 Box::pin(async {})
163 }
164
165 fn before_queued_reservation(&self, _job: &SpawnJob) -> BoxFuture<'static, ()> {
169 Box::pin(async {})
170 }
171}
172
173#[derive(Clone)]
179pub struct SessionInboxRuntimeBinding {
180 pub router: Arc<crate::SessionActivationRouter>,
181 pub inbox: Arc<dyn SessionInboxPort>,
182 pub storage: Arc<dyn Storage>,
183 pub persistence: Arc<dyn RuntimeSessionPersistence>,
184 pub parent_question_lock: Option<Arc<bamboo_storage::LockedSessionStore>>,
187}
188
189#[async_trait::async_trait]
194pub trait ExternalChildRunner: Send + Sync {
195 fn supports_native_ticket_budget(&self) -> bool {
198 false
199 }
200
201 fn set_local_child_skill_factory(
204 &self,
205 _factory: Option<Arc<dyn crate::external_agents::local_task::LocalChildSkillFactory>>,
206 ) {
207 }
208
209 fn set_local_task_runtime(
212 &self,
213 _runtime: Option<crate::external_agents::local_task::LocalTaskRuntime>,
214 ) {
215 }
216
217 async fn prepare_child_execution(
220 &self,
221 _session: &Session,
222 _run_id: &str,
223 _cancel: CancellationToken,
224 ) -> Result<Option<Arc<dyn RuntimeSessionPersistence>>, String> {
225 Ok(None)
226 }
227
228 async fn abandon_child_execution(
231 &self,
232 _session: &Session,
233 _run_id: &str,
234 ) -> Result<(), String> {
235 Ok(())
236 }
237
238 fn quarantine_child_execution(&self, _session: &Session, _run_id: &str) {}
240
241 fn track_child_task(
244 &self,
245 _session_id: &str,
246 _run_id: &str,
247 handle: tokio::task::JoinHandle<()>,
248 ) -> Option<tokio::task::JoinHandle<()>> {
249 Some(handle)
250 }
251
252 async fn finish_child_execution(
255 &self,
256 _session: &Session,
257 _run_id: &str,
258 _history_committed: bool,
259 ) -> Result<(), String> {
260 Ok(())
261 }
262
263 async fn shutdown(&self) -> Result<(), String> {
266 Ok(())
267 }
268
269 fn set_ticket_service(&self, _service: Option<Arc<bamboo_tickets::TicketService>>) {}
271 async fn validate_required_child_context_route(
274 &self,
275 _session: &Session,
276 ) -> Result<(), String> {
277 Err("required_child_context_unsupported: no supported registered worker route".into())
278 }
279 async fn should_handle(&self, session: &Session) -> bool;
281
282 async fn execute_external_child(
284 &self,
285 session: &mut Session,
286 job: &SpawnJob,
287 event_tx: tokio::sync::mpsc::Sender<AgentEvent>,
288 cancel_token: CancellationToken,
289 ) -> crate::runtime::runner::Result<()>;
290
291 async fn prepare_durable_child_delivery(
295 &self,
296 _session: &Session,
297 _activation_run_id: &str,
298 ) -> Result<bool, String> {
299 Ok(false)
300 }
301
302 async fn commit_durable_child_delivery(
309 &self,
310 _session: &Session,
311 _activation_run_id: &str,
312 ) -> Result<(), String> {
313 Ok(())
314 }
315
316 async fn confirm_durable_child_delivery(
321 &self,
322 _session: &Session,
323 _activation_run_id: &str,
324 _history_committed: bool,
325 ) -> Result<(), String> {
326 Ok(())
327 }
328
329 fn set_escalation_bridge(&self, _bridge: Option<bamboo_subagent::executor::HostBridge>) {}
336
337 fn set_session_inbox_runtime(&self, _binding: Option<SessionInboxRuntimeBinding>) {}
341
342 fn set_actor_directory_store(&self, _store: Option<Arc<bamboo_storage::SessionStoreV2>>) {}
345
346 fn set_canonical_subagent_tool(&self, _tool: Option<Arc<dyn bamboo_agent_core::tools::Tool>>) {}
349
350 fn set_actor_event_observer(
352 &self,
353 _observer: Option<Arc<dyn crate::external_agents::actor_event_stream::ActorEventObserver>>,
354 ) {
355 }
356}
357
358#[derive(Clone)]
359pub struct SpawnContext {
360 pub agent: Arc<Agent>,
361 pub tools: Arc<dyn ToolExecutor>,
362 pub sessions_cache: crate::SessionCache,
363 pub agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
364 pub session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
365 pub external_child_runner: Arc<dyn ExternalChildRunner>,
366 pub provider_router: Option<Arc<ProviderModelRouter>>,
367 pub app_data_dir: Option<std::path::PathBuf>,
368 pub completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
373 pub child_run_launch_hook: Option<Arc<dyn ChildRunLaunchHook>>,
376 pub account_feed_inbox: Option<super::event_forwarder::AccountFeedInbox>,
380}
381
382impl SpawnContext {
383 pub(crate) fn replayable_event_publisher(
384 &self,
385 ) -> super::session_events::ReplayableSessionEventPublisher {
386 super::session_events::ReplayableSessionEventPublisher::new(
387 self.agent_runners.clone(),
388 self.session_event_senders.clone(),
389 self.account_feed_inbox.clone(),
390 )
391 }
392}
393
394#[derive(Clone)]
395pub struct SpawnScheduler {
396 tx: mpsc::Sender<QueuedSpawnJob>,
397 ctx: SpawnContext,
398 launch_locks: ChildLaunchLocks,
399}
400
401impl SpawnScheduler {
402 pub fn supports_native_ticket_budget(&self) -> bool {
403 self.ctx
404 .external_child_runner
405 .supports_native_ticket_budget()
406 }
407
408 pub fn set_ticket_service(&self, service: Option<Arc<bamboo_tickets::TicketService>>) {
409 self.ctx.external_child_runner.set_ticket_service(service);
410 }
411
412 pub async fn validate_required_child_context_route(
413 &self,
414 metadata: &HashMap<String, String>,
415 role: &str,
416 ) -> Result<(), String> {
417 let mut candidate = Session::new("required-context-preflight", "");
418 candidate.kind = bamboo_domain::SessionKind::Child;
422 candidate.metadata = metadata.clone();
423 candidate
424 .metadata
425 .insert("subagent_type".into(), role.into());
426 self.ctx
427 .external_child_runner
428 .validate_required_child_context_route(&candidate)
429 .await
430 }
431
432 pub fn new(ctx: SpawnContext) -> Self {
433 let (tx, mut rx) = mpsc::channel::<QueuedSpawnJob>(128);
434 let worker_ctx = ctx.clone();
435 let launch_locks = ChildLaunchLocks::default();
436 let worker_launch_locks = launch_locks.clone();
437
438 tokio::spawn(async move {
446 while let Some(QueuedSpawnJob {
447 job,
448 title,
449 launch_generation,
450 announcement_ready,
451 start_published,
452 }) = rx.recv().await
453 {
454 let job_ctx = worker_ctx.clone();
455 let job_for_panic = job.clone();
456 let job_launch_locks = worker_launch_locks.clone();
457 let handle = tokio::spawn(async move {
458 let fallback_start =
462 needs_fallback_start(announcement_ready, &start_published).await;
463 if let Some(hook) = job_ctx.child_run_launch_hook.as_ref() {
464 hook.before_queued_dequeue(&job).await;
465 }
466 let launch_guard = if let Some(generation) = launch_generation {
467 let guard = job_launch_locks.acquire(&job.child_session_id).await;
468 if queued_launch_state(&job_ctx, &job, generation).await?
469 != QueuedLaunchState::Eligible
470 {
471 return Ok(());
472 }
473 Some(guard)
474 } else {
475 None
476 };
477 if let Some(hook) = job_ctx.child_run_launch_hook.as_ref() {
478 hook.before_queued_reservation(&job).await;
479 }
480 if fallback_start {
481 Self::prepare_child_launch(&job_ctx, &job).await;
482 job_ctx
483 .replayable_event_publisher()
484 .publish(
485 &job.parent_session_id,
486 AgentEvent::SubAgentStarted {
487 parent_session_id: job.parent_session_id.clone(),
488 child_session_id: job.child_session_id.clone(),
489 title,
490 },
491 )
492 .await;
493 }
494 if let Err(err) = run_spawn_job(job_ctx, job, launch_guard).await {
495 tracing::warn!("spawn job failed: {}", err);
496 }
497 Ok(())
498 });
499 let failure = match handle.await {
500 Ok(Ok(())) => continue,
501 Ok(Err(error)) => error,
502 Err(error) => format!("child spawn panicked: {error}"),
503 };
504 {
505 tracing::error!(
506 parent_session_id = %job_for_panic.parent_session_id,
507 child_session_id = %job_for_panic.child_session_id,
508 error = %failure,
509 "admitted spawn job failed; publishing terminal error completion"
510 );
511 let publisher = worker_ctx.replayable_event_publisher();
512 publish_child_completion_parts(
513 &publisher,
514 worker_ctx.completion_handler.clone(),
515 job_for_panic.parent_session_id.clone(),
516 job_for_panic.child_session_id.clone(),
517 "error".to_string(),
518 Some(failure),
519 )
520 .await;
521 }
522 }
523 });
524
525 Self {
526 tx,
527 ctx,
528 launch_locks,
529 }
530 }
531
532 async fn prepare_child_launch(ctx: &SpawnContext, job: &SpawnJob) {
533 let child_tx = super::session_events::get_or_create_event_sender(
534 &ctx.session_event_senders,
535 &job.child_session_id,
536 )
537 .await;
538 invoke_child_run_launch_hook(ctx.child_run_launch_hook.as_ref(), job, child_tx);
539 }
540
541 pub async fn enqueue(&self, job: SpawnJob) -> Result<(), String> {
542 self.enqueue_announced(job, None).await
543 }
544
545 pub async fn enqueue_announced(
549 &self,
550 job: SpawnJob,
551 title: Option<String>,
552 ) -> Result<(), String> {
553 match self.enqueue_announced_with_gate(job, title, None).await? {
554 AdmissionCommit::Committed(()) | AdmissionCommit::AlreadyCommitted => Ok(()),
555 AdmissionCommit::Cancelled => unreachable!("an ungated launch cannot be cancelled"),
556 }
557 }
558
559 pub async fn enqueue_announced_with_gate(
560 &self,
561 job: SpawnJob,
562 title: Option<String>,
563 gate: Option<&AdmissionGate>,
564 ) -> Result<AdmissionCommit<()>, String> {
565 self.enqueue_announced_for_generation(job, title, gate, None)
566 .await
567 }
568
569 pub async fn enqueue_announced_for_generation(
573 &self,
574 job: SpawnJob,
575 title: Option<String>,
576 gate: Option<&AdmissionGate>,
577 launch_generation: Option<u64>,
578 ) -> Result<AdmissionCommit<()>, String> {
579 self.enqueue_announced_for_generation_checked(job, title, gate, launch_generation, None)
580 .await
581 }
582
583 pub async fn enqueue_announced_for_generation_checked(
586 &self,
587 job: SpawnJob,
588 title: Option<String>,
589 gate: Option<&AdmissionGate>,
590 launch_generation: Option<u64>,
591 admission_check: Option<&QueueAdmissionCheck>,
592 ) -> Result<AdmissionCommit<()>, String> {
593 let preparation_job = job.clone();
594 let commit = reserve_and_send(
595 &self.tx,
596 job,
597 title.clone(),
598 gate,
599 launch_generation,
600 admission_check,
601 )
602 .await?;
603 self.finish_admission(preparation_job, title, launch_generation, commit)
604 .await
605 }
606
607 pub async fn reserve_queue_capacity<'a>(
610 &'a self,
611 gate: Option<&'a AdmissionGate>,
612 ) -> Result<AdmissionCommit<QueueCapacityReservation<'a>>, String> {
613 Ok(match reserve_queue_permit(&self.tx, gate).await? {
614 AdmissionCommit::Committed(permit) => {
615 AdmissionCommit::Committed(QueueCapacityReservation {
616 scheduler: self,
617 gate,
618 permit,
619 })
620 }
621 AdmissionCommit::AlreadyCommitted => AdmissionCommit::AlreadyCommitted,
622 AdmissionCommit::Cancelled => AdmissionCommit::Cancelled,
623 })
624 }
625
626 async fn finish_admission(
627 &self,
628 preparation_job: SpawnJob,
629 title: Option<String>,
630 launch_generation: Option<u64>,
631 commit: AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>,
632 ) -> Result<AdmissionCommit<()>, String> {
633 match commit {
634 AdmissionCommit::Committed((announce_ready, start_published)) => {
635 let launch_guard = if launch_generation.is_some() {
636 Some(
637 self.launch_locks
638 .acquire(&preparation_job.child_session_id)
639 .await,
640 )
641 } else {
642 None
643 };
644 let state = match launch_generation {
645 Some(generation) => {
646 queued_launch_state(&self.ctx, &preparation_job, generation).await?
647 }
648 None => QueuedLaunchState::Eligible,
649 };
650 if state == QueuedLaunchState::Eligible {
651 Self::prepare_child_launch(&self.ctx, &preparation_job).await;
652 self.ctx
653 .replayable_event_publisher()
654 .publish(
655 &preparation_job.parent_session_id,
656 AgentEvent::SubAgentStarted {
657 parent_session_id: preparation_job.parent_session_id.clone(),
658 child_session_id: preparation_job.child_session_id.clone(),
659 title,
660 },
661 )
662 .await;
663 }
664 drop(launch_guard);
665 start_published.store(true, Ordering::SeqCst);
669 let _ = announce_ready.send(());
670 match state {
671 QueuedLaunchState::Eligible => Ok(AdmissionCommit::Committed(())),
672 QueuedLaunchState::AlreadyStarted => Ok(AdmissionCommit::AlreadyCommitted),
673 QueuedLaunchState::Stale => Ok(AdmissionCommit::Cancelled),
674 }
675 }
676 AdmissionCommit::AlreadyCommitted => Ok(AdmissionCommit::AlreadyCommitted),
677 AdmissionCommit::Cancelled => Ok(AdmissionCommit::Cancelled),
678 }
679 }
680
681 pub async fn publish_parent_replayable_event(
684 &self,
685 parent_session_id: &str,
686 event: AgentEvent,
687 ) {
688 self.ctx
689 .replayable_event_publisher()
690 .publish(parent_session_id, event)
691 .await;
692 }
693
694 pub async fn lock_child_launch(&self, child_id: &str) -> ChildLaunchGuard {
695 self.launch_locks.acquire(child_id).await
696 }
697
698 pub async fn publish_queued_child_cancellation(&self, parent_id: &str, child_id: &str) {
699 let publisher = self.ctx.replayable_event_publisher();
700 publish_child_completion_parts(
701 &publisher,
702 self.ctx.completion_handler.clone(),
703 parent_id.to_string(),
704 child_id.to_string(),
705 "cancelled".to_string(),
706 Some("Cancelled by parent before activation".to_string()),
707 )
708 .await;
709 }
710
711 pub(crate) fn launch_reserved(
715 &self,
716 job: SpawnJob,
717 reservation: SessionExecutionReservation,
718 ) -> tokio::task::JoinHandle<()> {
719 let ctx = self.ctx.clone();
720 tokio::spawn(async move {
721 Self::prepare_child_launch(&ctx, &job).await;
722 ctx.replayable_event_publisher()
723 .publish(
724 &job.parent_session_id,
725 AgentEvent::SubAgentStarted {
726 parent_session_id: job.parent_session_id.clone(),
727 child_session_id: job.child_session_id.clone(),
728 title: None,
729 },
730 )
731 .await;
732 if let Err(error) =
733 crate::sdk::spawn::run_child_spawn_reserved(ctx, job.clone(), reservation).await
734 {
735 tracing::warn!(
736 parent_session_id = %job.parent_session_id,
737 child_session_id = %job.child_session_id,
738 %error,
739 "reserved child activation failed"
740 );
741 }
742 })
743 }
744}
745
746async fn needs_fallback_start(
747 announcement_ready: oneshot::Receiver<()>,
748 start_published: &AtomicBool,
749) -> bool {
750 announcement_ready.await.is_err() && !start_published.load(Ordering::SeqCst)
751}
752
753fn invoke_child_run_launch_hook(
754 hook: Option<&Arc<dyn ChildRunLaunchHook>>,
755 job: &SpawnJob,
756 child_tx: broadcast::Sender<AgentEvent>,
757) {
758 let Some(hook) = hook else {
759 return;
760 };
761 let invoked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
762 hook.before_child_launch(job, child_tx);
763 }));
764 if invoked.is_err() {
765 tracing::error!(
766 parent_session_id = %job.parent_session_id,
767 child_session_id = %job.child_session_id,
768 "child launch hook panicked; continuing with canonical execution"
769 );
770 }
771}
772
773async fn reserve_and_send(
778 tx: &mpsc::Sender<QueuedSpawnJob>,
779 job: SpawnJob,
780 title: Option<String>,
781 gate: Option<&AdmissionGate>,
782 launch_generation: Option<u64>,
783 admission_check: Option<&QueueAdmissionCheck>,
784) -> Result<AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>, String> {
785 match reserve_queue_permit(tx, gate).await? {
786 AdmissionCommit::Committed(permit) => send_reserved(
787 permit,
788 job,
789 title,
790 gate,
791 launch_generation,
792 admission_check,
793 Ok(()),
794 ),
795 AdmissionCommit::AlreadyCommitted => Ok(AdmissionCommit::AlreadyCommitted),
796 AdmissionCommit::Cancelled => Ok(AdmissionCommit::Cancelled),
797 }
798}
799
800async fn reserve_queue_permit<'a>(
801 tx: &'a mpsc::Sender<QueuedSpawnJob>,
802 gate: Option<&AdmissionGate>,
803) -> Result<AdmissionCommit<mpsc::Permit<'a, QueuedSpawnJob>>, String> {
804 if let Some(gate) = gate {
805 if gate.is_committed() {
806 return Ok(AdmissionCommit::AlreadyCommitted);
807 }
808 if gate.is_cancelled() {
809 return Ok(AdmissionCommit::Cancelled);
810 }
811 }
812 let permit = match gate {
813 Some(gate) => tokio::select! {
814 biased;
815 _ = gate.cancelled() => return Ok(AdmissionCommit::Cancelled),
816 permit = tx.reserve() => permit,
817 },
818 None => tx.reserve().await,
819 }
820 .map_err(|_| "spawn scheduler is not running".to_string())?;
821 Ok(AdmissionCommit::Committed(permit))
822}
823
824fn send_reserved(
825 permit: mpsc::Permit<'_, QueuedSpawnJob>,
826 job: SpawnJob,
827 title: Option<String>,
828 gate: Option<&AdmissionGate>,
829 launch_generation: Option<u64>,
830 admission_check: Option<&QueueAdmissionCheck>,
831 preflight: Result<(), String>,
832) -> Result<AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>, String> {
833 let (announce_ready, announcement_ready) = oneshot::channel();
834 let start_published = Arc::new(AtomicBool::new(false));
835 let queued = QueuedSpawnJob {
836 job,
837 title,
838 launch_generation,
839 announcement_ready,
840 start_published: start_published.clone(),
841 };
842 let commit = || {
843 preflight?;
844 if let Some(check) = admission_check {
845 check()?;
846 }
847 permit.send(queued);
848 Ok((announce_ready, start_published))
849 };
850 match gate {
851 Some(gate) => gate.commit(commit),
852 None => commit().map(AdmissionCommit::Committed),
853 }
854}
855
856#[derive(Debug, Clone, Copy)]
857pub(crate) struct ChildWatchdogPolicy {
858 check_interval_secs: i64,
859 pub(crate) max_total_secs: i64,
863 pub(crate) max_idle_secs: i64,
864}
865
866impl Default for ChildWatchdogPolicy {
867 fn default() -> Self {
868 Self {
869 check_interval_secs: 15,
870 max_total_secs: 60 * 60,
874 max_idle_secs: 15 * 60,
876 }
877 }
878}
879
880fn metadata_i64(session: &Session, key: &str) -> Option<i64> {
881 session
882 .metadata
883 .get(key)
884 .and_then(|value| value.trim().parse::<i64>().ok())
885 .filter(|value| *value > 0)
886}
887
888pub(crate) fn watchdog_policy_for_session(session: &Session) -> ChildWatchdogPolicy {
889 let mut policy = ChildWatchdogPolicy::default();
890 if let Some(value) = metadata_i64(session, "child_watchdog.max_total_secs") {
891 policy.max_total_secs = value;
892 }
893 if let Some(value) = metadata_i64(session, "child_watchdog.max_idle_secs") {
894 policy.max_idle_secs = value;
895 }
896 if let Some(value) = metadata_i64(session, "child_watchdog.check_interval_secs") {
897 policy.check_interval_secs = value;
898 }
899 policy
900}
901
902pub(crate) async fn publish_child_completion(
903 publisher: &super::session_events::ReplayableSessionEventPublisher,
904 completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
905 completion: ChildCompletion,
906) {
907 publisher
908 .publish(
909 &completion.parent_session_id,
910 AgentEvent::SubAgentCompleted {
911 parent_session_id: completion.parent_session_id.clone(),
912 child_session_id: completion.child_session_id.clone(),
913 status: completion.status.clone(),
914 error: completion.error.clone(),
915 },
916 )
917 .await;
918
919 if let Some(handler) = completion_handler {
920 use futures::FutureExt;
927 let parent_session_id = completion.parent_session_id.clone();
928 let child_session_id = completion.child_session_id.clone();
929 if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
930 .catch_unwind()
931 .await
932 .is_err()
933 {
934 tracing::error!(
935 %parent_session_id,
936 %child_session_id,
937 "child completion handler panicked; child-wait watchdog will backstop the parent wake"
938 );
939 }
940 }
941}
942
943pub(crate) async fn publish_child_completion_parts(
944 publisher: &super::session_events::ReplayableSessionEventPublisher,
945 completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
946 parent_session_id: String,
947 child_session_id: String,
948 status: String,
949 error: Option<String>,
950) {
951 publish_child_completion(
952 publisher,
953 completion_handler,
954 ChildCompletion {
955 parent_session_id,
956 child_session_id,
957 status,
958 error,
959 completed_at: Utc::now(),
960 source: None,
961 },
962 )
963 .await;
964}
965
966pub(crate) async fn watch_child_liveness(
967 parent_session_id: String,
968 child_session_id: String,
969 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
970 cancel_token: CancellationToken,
971 timeout_reason: Arc<RwLock<Option<String>>>,
972 done: CancellationToken,
973 policy: ChildWatchdogPolicy,
974) {
975 let mut ticker =
976 tokio::time::interval(Duration::from_secs(policy.check_interval_secs.max(1) as u64));
977 ticker.tick().await;
979
980 loop {
981 tokio::select! {
982 _ = done.cancelled() => return,
983 _ = ticker.tick() => {
984 if cancel_token.is_cancelled() {
985 return;
986 }
987
988 let snapshot = {
989 let guard = runners.read().await;
990 guard.get(&child_session_id).cloned()
991 };
992 let Some(runner) = snapshot else {
993 return;
994 };
995 if !matches!(runner.status, super::runner_state::AgentStatus::Running) {
996 return;
997 }
998
999 let now = Utc::now();
1000 let total_secs = now.signed_duration_since(runner.started_at).num_seconds();
1001 if total_secs >= policy.max_total_secs {
1002 let reason = format!(
1003 "Child session timed out after {} seconds (max_total_secs={})",
1004 total_secs, policy.max_total_secs
1005 );
1006 tracing::warn!(
1007 parent_session_id = %parent_session_id,
1008 child_session_id = %child_session_id,
1009 reason = %reason,
1010 "child session total timeout; cancelling child runner"
1011 );
1012 *timeout_reason.write().await = Some(reason);
1013 cancel_token.cancel();
1014 return;
1015 }
1016
1017 let last_activity_at = runner.last_activity_at().unwrap_or(runner.started_at);
1018 let idle_secs = now.signed_duration_since(last_activity_at).num_seconds();
1019 if idle_secs >= policy.max_idle_secs {
1020 let reason = format!(
1021 "Child session idle timeout after {} seconds without events (max_idle_secs={})",
1022 idle_secs, policy.max_idle_secs
1023 );
1024 tracing::warn!(
1025 parent_session_id = %parent_session_id,
1026 child_session_id = %child_session_id,
1027 reason = %reason,
1028 last_tool_name = ?runner.last_tool_name,
1029 last_tool_phase = ?runner.last_tool_phase,
1030 round_count = runner.round_count,
1031 "child session idle timeout; cancelling child runner"
1032 );
1033 *timeout_reason.write().await = Some(reason);
1034 cancel_token.cancel();
1035 return;
1036 }
1037 }
1038 }
1039 }
1040}
1041
1042async fn run_spawn_job(
1049 ctx: SpawnContext,
1050 job: SpawnJob,
1051 launch_guard: Option<ChildLaunchGuard>,
1052) -> Result<(), String> {
1053 match launch_guard {
1054 Some(guard) => crate::sdk::spawn::run_child_spawn_fenced(ctx, job, guard).await,
1055 None => crate::sdk::spawn::run_child_spawn(ctx, job).await,
1056 }
1057}
1058
1059#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1060enum QueuedLaunchState {
1061 Eligible,
1062 AlreadyStarted,
1063 Stale,
1064}
1065
1066async fn queued_launch_state(
1067 ctx: &SpawnContext,
1068 job: &SpawnJob,
1069 generation: u64,
1070) -> Result<QueuedLaunchState, String> {
1071 let child = ctx
1072 .agent
1073 .storage()
1074 .load_session(&job.child_session_id)
1075 .await
1076 .map_err(|error| format!("load queued child launch fence: {error}"))?;
1077 let Some(child) = child else {
1078 return Ok(QueuedLaunchState::Stale);
1079 };
1080 if child.kind != bamboo_domain::SessionKind::Child
1081 || child.parent_session_id.as_deref() != Some(job.parent_session_id.as_str())
1082 || child.child_launch_generation() != generation
1083 || child.is_child_launch_cancelled(generation)
1084 {
1085 return Ok(QueuedLaunchState::Stale);
1086 }
1087 match child.last_run_status().as_deref() {
1088 Some("pending") => Ok(QueuedLaunchState::Eligible),
1089 Some("running" | "completed" | "error" | "timeout" | "skipped" | "suspended") => {
1090 Ok(QueuedLaunchState::AlreadyStarted)
1091 }
1092 _ => Ok(QueuedLaunchState::Stale),
1093 }
1094}
1095
1096#[cfg(test)]
1097mod launch_hook_tests {
1098 use super::*;
1099 use std::sync::atomic::{AtomicUsize, Ordering};
1100
1101 fn job() -> SpawnJob {
1102 SpawnJob {
1103 parent_session_id: "parent".to_string(),
1104 child_session_id: "child".to_string(),
1105 model: "test".to_string(),
1106 disabled_tools: None,
1107 }
1108 }
1109
1110 struct PanickingHook {
1111 calls: AtomicUsize,
1112 }
1113
1114 impl ChildRunLaunchHook for PanickingHook {
1115 fn before_child_launch(
1116 &self,
1117 _job: &SpawnJob,
1118 _child_events: broadcast::Sender<AgentEvent>,
1119 ) {
1120 self.calls.fetch_add(1, Ordering::SeqCst);
1121 panic!("injected launch hook panic");
1122 }
1123 }
1124
1125 #[test]
1126 fn launch_hook_panic_is_contained() {
1127 let hook = Arc::new(PanickingHook {
1128 calls: AtomicUsize::new(0),
1129 });
1130 let hook_port: Arc<dyn ChildRunLaunchHook> = hook.clone();
1131 let (child_tx, _child_rx) = broadcast::channel(1);
1132
1133 invoke_child_run_launch_hook(Some(&hook_port), &job(), child_tx);
1134
1135 assert_eq!(hook.calls.load(Ordering::SeqCst), 1);
1136 }
1137
1138 #[tokio::test]
1139 async fn closed_scheduler_does_not_prepare_phantom_launch() {
1140 let (tx, rx) = mpsc::channel(1);
1141 drop(rx);
1142 let result = reserve_and_send(&tx, job(), None, None, None, None).await;
1143
1144 assert_eq!(result.unwrap_err(), "spawn scheduler is not running");
1145 }
1146
1147 #[tokio::test]
1148 async fn cancellation_while_waiting_for_queue_capacity_never_sends() {
1149 let (tx, mut rx) = mpsc::channel(1);
1150 let (_, announcement_ready) = oneshot::channel();
1151 tx.send(QueuedSpawnJob {
1152 job: job(),
1153 title: None,
1154 launch_generation: None,
1155 announcement_ready,
1156 start_published: Arc::new(AtomicBool::new(false)),
1157 })
1158 .await
1159 .unwrap();
1160 let gate = Arc::new(AdmissionGate::default());
1161 let waiting = {
1162 let tx = tx.clone();
1163 let gate = gate.clone();
1164 tokio::spawn(async move {
1165 reserve_and_send(&tx, job(), None, Some(&gate), None, None).await
1166 })
1167 };
1168 tokio::task::yield_now().await;
1169 assert!(!waiting.is_finished());
1170 gate.cancel_if_pending();
1171 assert!(matches!(
1172 tokio::time::timeout(Duration::from_millis(200), waiting)
1173 .await
1174 .expect("cancellation must wake without freeing queue capacity")
1175 .unwrap()
1176 .unwrap(),
1177 AdmissionCommit::Cancelled
1178 ));
1179 rx.recv().await.unwrap();
1180 assert!(matches!(
1181 rx.try_recv(),
1182 Err(mpsc::error::TryRecvError::Empty)
1183 ));
1184 }
1185
1186 #[tokio::test]
1187 async fn committed_queue_gate_does_not_enqueue_a_duplicate_job() {
1188 let (tx, mut rx) = mpsc::channel(1);
1189 let gate = AdmissionGate::default();
1190 assert!(matches!(
1191 reserve_and_send(&tx, job(), None, Some(&gate), None, None)
1192 .await
1193 .unwrap(),
1194 AdmissionCommit::Committed(_)
1195 ));
1196 assert!(matches!(
1197 reserve_and_send(&tx, job(), None, Some(&gate), None, None)
1198 .await
1199 .unwrap(),
1200 AdmissionCommit::AlreadyCommitted
1201 ));
1202 rx.recv().await.unwrap();
1203 assert!(matches!(
1204 rx.try_recv(),
1205 Err(mpsc::error::TryRecvError::Empty)
1206 ));
1207 }
1208
1209 async fn assert_expired_queue_check(gate: Option<&AdmissionGate>) {
1210 let (tx, mut rx) = mpsc::channel(1);
1211 let _first = reserve_and_send(&tx, job(), None, None, None, None)
1212 .await
1213 .unwrap();
1214 let expired = Arc::new(AtomicBool::new(false));
1215 let calls = Arc::new(AtomicUsize::new(0));
1216 let check = {
1217 let expired = expired.clone();
1218 let calls = calls.clone();
1219 move || {
1220 calls.fetch_add(1, Ordering::SeqCst);
1221 if expired.load(Ordering::SeqCst) {
1222 Err("Host admission expired".to_string())
1223 } else {
1224 Ok(())
1225 }
1226 }
1227 };
1228 let waiting = reserve_and_send(&tx, job(), None, gate, None, Some(&check));
1229 tokio::pin!(waiting);
1230 assert!(futures::poll!(waiting.as_mut()).is_pending());
1231 assert_eq!(calls.load(Ordering::SeqCst), 0);
1232
1233 expired.store(true, Ordering::SeqCst);
1236 rx.recv().await.unwrap();
1237 assert_eq!(waiting.await.unwrap_err(), "Host admission expired");
1238 assert_eq!(calls.load(Ordering::SeqCst), 1);
1239 assert!(matches!(
1240 rx.try_recv(),
1241 Err(mpsc::error::TryRecvError::Empty)
1242 ));
1243 }
1244
1245 #[tokio::test]
1246 async fn queue_admission_check_rejects_expiry_after_capacity_wait() {
1247 assert_expired_queue_check(None).await;
1248 }
1249
1250 #[tokio::test]
1251 async fn queue_admission_check_rejection_preserves_cancellation_gate() {
1252 let gate = AdmissionGate::default();
1253 assert_expired_queue_check(Some(&gate)).await;
1254 assert!(!gate.is_committed());
1255 gate.cancel_if_pending();
1256 assert!(gate.is_cancelled());
1257 }
1258
1259 #[tokio::test]
1260 async fn queue_admission_check_runs_once_for_successful_send() {
1261 let (tx, mut rx) = mpsc::channel(1);
1262 let gate = AdmissionGate::default();
1263 let calls = Arc::new(AtomicUsize::new(0));
1264 let check = {
1265 let calls = calls.clone();
1266 move || {
1267 calls.fetch_add(1, Ordering::SeqCst);
1268 Ok(())
1269 }
1270 };
1271 assert!(matches!(
1272 reserve_and_send(&tx, job(), None, Some(&gate), None, Some(&check))
1273 .await
1274 .unwrap(),
1275 AdmissionCommit::Committed(_)
1276 ));
1277 assert!(matches!(
1278 reserve_and_send(&tx, job(), None, Some(&gate), None, Some(&check))
1279 .await
1280 .unwrap(),
1281 AdmissionCommit::AlreadyCommitted
1282 ));
1283 assert_eq!(calls.load(Ordering::SeqCst), 1);
1284 rx.recv().await.unwrap();
1285 assert!(matches!(
1286 rx.try_recv(),
1287 Err(mpsc::error::TryRecvError::Empty)
1288 ));
1289 }
1290
1291 mod queue_capacity_revalidation {
1292 use super::*;
1293
1294 async fn reserve<'a>(
1295 tx: &'a mpsc::Sender<QueuedSpawnJob>,
1296 gate: &AdmissionGate,
1297 ) -> mpsc::Permit<'a, QueuedSpawnJob> {
1298 match reserve_queue_permit(tx, Some(gate)).await.unwrap() {
1299 AdmissionCommit::Committed(permit) => permit,
1300 _ => panic!("fresh gate must reserve without committing"),
1301 }
1302 }
1303
1304 #[tokio::test]
1305 async fn reservation_does_not_commit_the_original_gate() {
1306 let (tx, mut rx) = mpsc::channel(1);
1307 let gate = AdmissionGate::default();
1308 let permit = reserve(&tx, &gate).await;
1309 assert!(!gate.is_committed());
1310 assert!(!gate.is_cancelled());
1311 assert_eq!(tx.capacity(), 0);
1312 assert!(matches!(
1313 rx.try_recv(),
1314 Err(mpsc::error::TryRecvError::Empty)
1315 ));
1316 assert!(matches!(
1317 send_reserved(permit, job(), None, Some(&gate), Some(7), None, Ok(())).unwrap(),
1318 AdmissionCommit::Committed(_)
1319 ));
1320 assert!(gate.is_committed());
1321 assert_eq!(rx.recv().await.unwrap().launch_generation, Some(7));
1322 }
1323
1324 #[tokio::test]
1325 async fn preflight_error_keeps_pending_gate_and_releases_capacity() {
1326 let (tx, mut rx) = mpsc::channel(1);
1327 let gate = AdmissionGate::default();
1328 let permit = reserve(&tx, &gate).await;
1329 let error = send_reserved(
1330 permit,
1331 job(),
1332 None,
1333 Some(&gate),
1334 Some(7),
1335 None,
1336 Err("fresh Host resources rejected".into()),
1337 )
1338 .expect_err("resource failure cannot commit");
1339 assert_eq!(error, "fresh Host resources rejected");
1340 assert!(!gate.is_committed());
1341 assert!(!gate.is_cancelled());
1342 assert_eq!(tx.capacity(), 1);
1343 assert!(matches!(
1344 rx.try_recv(),
1345 Err(mpsc::error::TryRecvError::Empty)
1346 ));
1347 }
1348
1349 #[tokio::test]
1350 async fn concurrent_commit_wins_over_a_late_preflight_error() {
1351 let (tx, mut rx) = mpsc::channel(2);
1352 let gate = AdmissionGate::default();
1353 let late = reserve(&tx, &gate).await;
1354 let first = reserve(&tx, &gate).await;
1355 assert!(matches!(
1356 send_reserved(first, job(), None, Some(&gate), Some(7), None, Ok(())).unwrap(),
1357 AdmissionCommit::Committed(_)
1358 ));
1359 let calls = Arc::new(AtomicUsize::new(0));
1360 let check_calls = calls.clone();
1361 let check = move || {
1362 check_calls.fetch_add(1, Ordering::SeqCst);
1363 Ok(())
1364 };
1365 assert!(matches!(
1366 send_reserved(
1367 late,
1368 job(),
1369 None,
1370 Some(&gate),
1371 Some(7),
1372 Some(&check),
1373 Err("resource check finished after another commit".into())
1374 )
1375 .unwrap(),
1376 AdmissionCommit::AlreadyCommitted
1377 ));
1378 assert_eq!(calls.load(Ordering::SeqCst), 0);
1379 assert_eq!(rx.recv().await.unwrap().launch_generation, Some(7));
1380 assert!(matches!(
1381 rx.try_recv(),
1382 Err(mpsc::error::TryRecvError::Empty)
1383 ));
1384 assert_eq!(tx.capacity(), 2);
1385 }
1386
1387 #[tokio::test]
1388 async fn cancellation_wins_over_a_late_preflight_error() {
1389 let (tx, mut rx) = mpsc::channel(1);
1390 let gate = AdmissionGate::default();
1391 let permit = reserve(&tx, &gate).await;
1392 gate.cancel_if_pending();
1393 assert!(matches!(
1394 send_reserved(
1395 permit,
1396 job(),
1397 None,
1398 Some(&gate),
1399 Some(7),
1400 None,
1401 Err("preflight observed failure after cancellation".into())
1402 )
1403 .unwrap(),
1404 AdmissionCommit::Cancelled
1405 ));
1406 assert_eq!(tx.capacity(), 1);
1407 assert!(matches!(
1408 rx.try_recv(),
1409 Err(mpsc::error::TryRecvError::Empty)
1410 ));
1411 }
1412
1413 #[tokio::test]
1414 async fn successful_preflight_cannot_bypass_the_final_live_fence() {
1415 let (tx, mut rx) = mpsc::channel(1);
1416 let gate = AdmissionGate::default();
1417 let permit = reserve(&tx, &gate).await;
1418 let check = || Err("Ready policy expired during resource preflight".into());
1419 let error = send_reserved(
1420 permit,
1421 job(),
1422 None,
1423 Some(&gate),
1424 Some(7),
1425 Some(&check),
1426 Ok(()),
1427 )
1428 .expect_err("the final synchronous fence remains authoritative");
1429 assert_eq!(error, "Ready policy expired during resource preflight");
1430 assert!(!gate.is_committed());
1431 assert!(!gate.is_cancelled());
1432 assert_eq!(tx.capacity(), 1);
1433 assert!(matches!(
1434 rx.try_recv(),
1435 Err(mpsc::error::TryRecvError::Empty)
1436 ));
1437 }
1438 }
1439
1440 #[tokio::test]
1441 async fn producer_drop_after_start_publication_does_not_require_fallback_start() {
1442 let (announce_ready, announcement_ready) = oneshot::channel::<()>();
1443 let start_published = Arc::new(AtomicBool::new(false));
1444 let runners = Arc::new(RwLock::new(HashMap::new()));
1445 let senders = Arc::new(RwLock::new(HashMap::new()));
1446 let (sender, mut receiver) = broadcast::channel(4);
1447 senders.write().await.insert("parent".to_string(), sender);
1448 let publisher = super::super::session_events::ReplayableSessionEventPublisher::new(
1449 runners, senders, None,
1450 );
1451 publisher
1455 .publish(
1456 "parent",
1457 AgentEvent::SubAgentStarted {
1458 parent_session_id: "parent".to_string(),
1459 child_session_id: "child".to_string(),
1460 title: None,
1461 },
1462 )
1463 .await;
1464 start_published.store(true, Ordering::SeqCst);
1465 drop(announce_ready);
1466 assert!(!needs_fallback_start(announcement_ready, &start_published).await);
1467 assert!(matches!(
1468 receiver.try_recv(),
1469 Ok(AgentEvent::SubAgentStarted { .. })
1470 ));
1471 assert!(matches!(
1472 receiver.try_recv(),
1473 Err(broadcast::error::TryRecvError::Empty)
1474 ));
1475
1476 let (announce_ready, announcement_ready) = oneshot::channel::<()>();
1477 let start_published = AtomicBool::new(false);
1478 drop(announce_ready);
1479 assert!(needs_fallback_start(announcement_ready, &start_published).await);
1480 }
1481}