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#[derive(Clone, Debug, PartialEq, Eq)]
195pub struct NativeTaskStart {
196 pub session_id: String,
197 pub child_birth: chrono::DateTime<chrono::Utc>,
198 pub run_id: String,
199 pub instance_id: uuid::Uuid,
200 pub entered_at: chrono::DateTime<chrono::Utc>,
201}
202
203impl NativeTaskStart {
204 pub fn started_at_ms(&self) -> Option<u64> {
205 if self.entered_at < self.child_birth || self.entered_at > chrono::Utc::now() {
206 return None;
207 }
208 self.entered_at.timestamp_millis().try_into().ok()
209 }
210}
211
212#[derive(Clone, Debug)]
215pub struct NativeTaskLiveness {
216 pub start: NativeTaskStart,
217 pub renewal_started_at: Option<chrono::DateTime<chrono::Utc>>,
218}
219
220impl NativeTaskLiveness {
221 pub fn heartbeat_at_ms(&self) -> Option<u64> {
222 let at = self.renewal_started_at?;
223 if at < self.start.entered_at || at > chrono::Utc::now() {
224 return None;
225 }
226 at.timestamp_millis().try_into().ok()
227 }
228}
229
230#[async_trait::async_trait]
232pub trait ExternalChildRunner: Send + Sync {
233 fn native_task_start(
236 &self,
237 _session_id: &str,
238 _run_id: &str,
239 _child_birth: chrono::DateTime<chrono::Utc>,
240 ) -> Option<NativeTaskStart> {
241 None
242 }
243
244 fn native_task_liveness(
246 &self,
247 session_id: &str,
248 run_id: &str,
249 child_birth: chrono::DateTime<chrono::Utc>,
250 ) -> Option<NativeTaskLiveness> {
251 self.native_task_start(session_id, run_id, child_birth)
252 .map(|start| NativeTaskLiveness {
253 start,
254 renewal_started_at: None,
255 })
256 }
257
258 fn supports_native_ticket_budget(&self) -> bool {
261 false
262 }
263
264 fn set_local_child_skill_factory(
267 &self,
268 _factory: Option<Arc<dyn crate::external_agents::local_task::LocalChildSkillFactory>>,
269 ) {
270 }
271
272 fn set_local_task_runtime(
275 &self,
276 _runtime: Option<crate::external_agents::local_task::LocalTaskRuntime>,
277 ) {
278 }
279
280 async fn prepare_child_execution(
283 &self,
284 _session: &Session,
285 _run_id: &str,
286 _cancel: CancellationToken,
287 ) -> Result<Option<Arc<dyn RuntimeSessionPersistence>>, String> {
288 Ok(None)
289 }
290
291 async fn abandon_child_execution(
294 &self,
295 _session: &Session,
296 _run_id: &str,
297 ) -> Result<(), String> {
298 Ok(())
299 }
300
301 fn quarantine_child_execution(&self, _session: &Session, _run_id: &str) {}
303
304 fn track_child_task(
307 &self,
308 _session_id: &str,
309 _run_id: &str,
310 handle: tokio::task::JoinHandle<()>,
311 ) -> Option<tokio::task::JoinHandle<()>> {
312 Some(handle)
313 }
314
315 async fn finish_child_execution(
318 &self,
319 _session: &Session,
320 _run_id: &str,
321 _history_committed: bool,
322 ) -> Result<(), String> {
323 Ok(())
324 }
325
326 async fn shutdown(&self) -> Result<(), String> {
329 Ok(())
330 }
331
332 fn set_ticket_service(&self, _service: Option<Arc<bamboo_tickets::TicketService>>) {}
334 async fn validate_required_child_context_route(
337 &self,
338 _session: &Session,
339 ) -> Result<(), String> {
340 Err("required_child_context_unsupported: no supported registered worker route".into())
341 }
342 async fn should_handle(&self, session: &Session) -> bool;
344
345 async fn execute_external_child(
347 &self,
348 session: &mut Session,
349 job: &SpawnJob,
350 event_tx: tokio::sync::mpsc::Sender<AgentEvent>,
351 cancel_token: CancellationToken,
352 ) -> crate::runtime::runner::Result<()>;
353
354 async fn prepare_durable_child_delivery(
358 &self,
359 _session: &Session,
360 _activation_run_id: &str,
361 ) -> Result<bool, String> {
362 Ok(false)
363 }
364
365 async fn commit_durable_child_delivery(
372 &self,
373 _session: &Session,
374 _activation_run_id: &str,
375 ) -> Result<(), String> {
376 Ok(())
377 }
378
379 async fn confirm_durable_child_delivery(
384 &self,
385 _session: &Session,
386 _activation_run_id: &str,
387 _history_committed: bool,
388 ) -> Result<(), String> {
389 Ok(())
390 }
391
392 fn set_escalation_bridge(&self, _bridge: Option<bamboo_subagent::executor::HostBridge>) {}
399
400 fn set_session_inbox_runtime(&self, _binding: Option<SessionInboxRuntimeBinding>) {}
404
405 fn set_actor_directory_store(&self, _store: Option<Arc<bamboo_storage::SessionStoreV2>>) {}
408
409 fn set_canonical_subagent_tool(&self, _tool: Option<Arc<dyn bamboo_agent_core::tools::Tool>>) {}
412
413 fn set_actor_event_observer(
415 &self,
416 _observer: Option<Arc<dyn crate::external_agents::actor_event_stream::ActorEventObserver>>,
417 ) {
418 }
419}
420
421#[derive(Clone)]
422pub struct SpawnContext {
423 pub agent: Arc<Agent>,
424 pub tools: Arc<dyn ToolExecutor>,
425 pub sessions_cache: crate::SessionCache,
426 pub agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
427 pub session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
428 pub external_child_runner: Arc<dyn ExternalChildRunner>,
429 pub provider_router: Option<Arc<ProviderModelRouter>>,
430 pub app_data_dir: Option<std::path::PathBuf>,
431 pub completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
436 pub child_run_launch_hook: Option<Arc<dyn ChildRunLaunchHook>>,
439 pub account_feed_inbox: Option<super::event_forwarder::AccountFeedInbox>,
443}
444
445impl SpawnContext {
446 pub(crate) fn replayable_event_publisher(
447 &self,
448 ) -> super::session_events::ReplayableSessionEventPublisher {
449 super::session_events::ReplayableSessionEventPublisher::new(
450 self.agent_runners.clone(),
451 self.session_event_senders.clone(),
452 self.account_feed_inbox.clone(),
453 )
454 }
455}
456
457#[derive(Clone)]
458pub struct SpawnScheduler {
459 tx: mpsc::Sender<QueuedSpawnJob>,
460 ctx: SpawnContext,
461 launch_locks: ChildLaunchLocks,
462}
463
464impl SpawnScheduler {
465 pub fn native_task_start(
466 &self,
467 session_id: &str,
468 run_id: &str,
469 child_birth: chrono::DateTime<chrono::Utc>,
470 ) -> Option<NativeTaskStart> {
471 self.ctx
472 .external_child_runner
473 .native_task_start(session_id, run_id, child_birth)
474 }
475
476 pub fn native_task_liveness(
477 &self,
478 session_id: &str,
479 run_id: &str,
480 child_birth: chrono::DateTime<chrono::Utc>,
481 ) -> Option<NativeTaskLiveness> {
482 self.ctx
483 .external_child_runner
484 .native_task_liveness(session_id, run_id, child_birth)
485 }
486
487 pub fn supports_native_ticket_budget(&self) -> bool {
488 self.ctx
489 .external_child_runner
490 .supports_native_ticket_budget()
491 }
492
493 pub fn set_ticket_service(&self, service: Option<Arc<bamboo_tickets::TicketService>>) {
494 self.ctx.external_child_runner.set_ticket_service(service);
495 }
496
497 pub async fn validate_required_child_context_route(
498 &self,
499 metadata: &HashMap<String, String>,
500 role: &str,
501 ) -> Result<(), String> {
502 let mut candidate = Session::new("required-context-preflight", "");
503 candidate.kind = bamboo_domain::SessionKind::Child;
507 candidate.metadata = metadata.clone();
508 candidate
509 .metadata
510 .insert("subagent_type".into(), role.into());
511 self.ctx
512 .external_child_runner
513 .validate_required_child_context_route(&candidate)
514 .await
515 }
516
517 pub fn new(ctx: SpawnContext) -> Self {
518 let (tx, mut rx) = mpsc::channel::<QueuedSpawnJob>(128);
519 let worker_ctx = ctx.clone();
520 let launch_locks = ChildLaunchLocks::default();
521 let worker_launch_locks = launch_locks.clone();
522
523 tokio::spawn(async move {
531 while let Some(QueuedSpawnJob {
532 job,
533 title,
534 launch_generation,
535 announcement_ready,
536 start_published,
537 }) = rx.recv().await
538 {
539 let job_ctx = worker_ctx.clone();
540 let job_for_panic = job.clone();
541 let job_launch_locks = worker_launch_locks.clone();
542 let handle = tokio::spawn(async move {
543 let fallback_start =
547 needs_fallback_start(announcement_ready, &start_published).await;
548 if let Some(hook) = job_ctx.child_run_launch_hook.as_ref() {
549 hook.before_queued_dequeue(&job).await;
550 }
551 let launch_guard = if let Some(generation) = launch_generation {
552 let guard = job_launch_locks.acquire(&job.child_session_id).await;
553 if queued_launch_state(&job_ctx, &job, generation).await?
554 != QueuedLaunchState::Eligible
555 {
556 return Ok(());
557 }
558 Some(guard)
559 } else {
560 None
561 };
562 if let Some(hook) = job_ctx.child_run_launch_hook.as_ref() {
563 hook.before_queued_reservation(&job).await;
564 }
565 if fallback_start {
566 Self::prepare_child_launch(&job_ctx, &job).await;
567 job_ctx
568 .replayable_event_publisher()
569 .publish(
570 &job.parent_session_id,
571 AgentEvent::SubAgentStarted {
572 parent_session_id: job.parent_session_id.clone(),
573 child_session_id: job.child_session_id.clone(),
574 title,
575 },
576 )
577 .await;
578 }
579 if let Err(err) = run_spawn_job(job_ctx, job, launch_guard).await {
580 tracing::warn!("spawn job failed: {}", err);
581 }
582 Ok(())
583 });
584 let failure = match handle.await {
585 Ok(Ok(())) => continue,
586 Ok(Err(error)) => error,
587 Err(error) => format!("child spawn panicked: {error}"),
588 };
589 {
590 tracing::error!(
591 parent_session_id = %job_for_panic.parent_session_id,
592 child_session_id = %job_for_panic.child_session_id,
593 error = %failure,
594 "admitted spawn job failed; publishing terminal error completion"
595 );
596 let publisher = worker_ctx.replayable_event_publisher();
597 publish_child_completion_parts(
598 &publisher,
599 worker_ctx.completion_handler.clone(),
600 job_for_panic.parent_session_id.clone(),
601 job_for_panic.child_session_id.clone(),
602 "error".to_string(),
603 Some(failure),
604 )
605 .await;
606 }
607 }
608 });
609
610 Self {
611 tx,
612 ctx,
613 launch_locks,
614 }
615 }
616
617 async fn prepare_child_launch(ctx: &SpawnContext, job: &SpawnJob) {
618 let child_tx = super::session_events::get_or_create_event_sender(
619 &ctx.session_event_senders,
620 &job.child_session_id,
621 )
622 .await;
623 invoke_child_run_launch_hook(ctx.child_run_launch_hook.as_ref(), job, child_tx);
624 }
625
626 pub async fn enqueue(&self, job: SpawnJob) -> Result<(), String> {
627 self.enqueue_announced(job, None).await
628 }
629
630 pub async fn enqueue_announced(
634 &self,
635 job: SpawnJob,
636 title: Option<String>,
637 ) -> Result<(), String> {
638 match self.enqueue_announced_with_gate(job, title, None).await? {
639 AdmissionCommit::Committed(()) | AdmissionCommit::AlreadyCommitted => Ok(()),
640 AdmissionCommit::Cancelled => unreachable!("an ungated launch cannot be cancelled"),
641 }
642 }
643
644 pub async fn enqueue_announced_with_gate(
645 &self,
646 job: SpawnJob,
647 title: Option<String>,
648 gate: Option<&AdmissionGate>,
649 ) -> Result<AdmissionCommit<()>, String> {
650 self.enqueue_announced_for_generation(job, title, gate, None)
651 .await
652 }
653
654 pub async fn enqueue_announced_for_generation(
658 &self,
659 job: SpawnJob,
660 title: Option<String>,
661 gate: Option<&AdmissionGate>,
662 launch_generation: Option<u64>,
663 ) -> Result<AdmissionCommit<()>, String> {
664 self.enqueue_announced_for_generation_checked(job, title, gate, launch_generation, None)
665 .await
666 }
667
668 pub async fn enqueue_announced_for_generation_checked(
671 &self,
672 job: SpawnJob,
673 title: Option<String>,
674 gate: Option<&AdmissionGate>,
675 launch_generation: Option<u64>,
676 admission_check: Option<&QueueAdmissionCheck>,
677 ) -> Result<AdmissionCommit<()>, String> {
678 let preparation_job = job.clone();
679 let commit = reserve_and_send(
680 &self.tx,
681 job,
682 title.clone(),
683 gate,
684 launch_generation,
685 admission_check,
686 )
687 .await?;
688 self.finish_admission(preparation_job, title, launch_generation, commit)
689 .await
690 }
691
692 pub async fn reserve_queue_capacity<'a>(
695 &'a self,
696 gate: Option<&'a AdmissionGate>,
697 ) -> Result<AdmissionCommit<QueueCapacityReservation<'a>>, String> {
698 Ok(match reserve_queue_permit(&self.tx, gate).await? {
699 AdmissionCommit::Committed(permit) => {
700 AdmissionCommit::Committed(QueueCapacityReservation {
701 scheduler: self,
702 gate,
703 permit,
704 })
705 }
706 AdmissionCommit::AlreadyCommitted => AdmissionCommit::AlreadyCommitted,
707 AdmissionCommit::Cancelled => AdmissionCommit::Cancelled,
708 })
709 }
710
711 async fn finish_admission(
712 &self,
713 preparation_job: SpawnJob,
714 title: Option<String>,
715 launch_generation: Option<u64>,
716 commit: AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>,
717 ) -> Result<AdmissionCommit<()>, String> {
718 match commit {
719 AdmissionCommit::Committed((announce_ready, start_published)) => {
720 let launch_guard = if launch_generation.is_some() {
721 Some(
722 self.launch_locks
723 .acquire(&preparation_job.child_session_id)
724 .await,
725 )
726 } else {
727 None
728 };
729 let state = match launch_generation {
730 Some(generation) => {
731 queued_launch_state(&self.ctx, &preparation_job, generation).await?
732 }
733 None => QueuedLaunchState::Eligible,
734 };
735 if state == QueuedLaunchState::Eligible {
736 Self::prepare_child_launch(&self.ctx, &preparation_job).await;
737 self.ctx
738 .replayable_event_publisher()
739 .publish(
740 &preparation_job.parent_session_id,
741 AgentEvent::SubAgentStarted {
742 parent_session_id: preparation_job.parent_session_id.clone(),
743 child_session_id: preparation_job.child_session_id.clone(),
744 title,
745 },
746 )
747 .await;
748 }
749 drop(launch_guard);
750 start_published.store(true, Ordering::SeqCst);
754 let _ = announce_ready.send(());
755 match state {
756 QueuedLaunchState::Eligible => Ok(AdmissionCommit::Committed(())),
757 QueuedLaunchState::AlreadyStarted => Ok(AdmissionCommit::AlreadyCommitted),
758 QueuedLaunchState::Stale => Ok(AdmissionCommit::Cancelled),
759 }
760 }
761 AdmissionCommit::AlreadyCommitted => Ok(AdmissionCommit::AlreadyCommitted),
762 AdmissionCommit::Cancelled => Ok(AdmissionCommit::Cancelled),
763 }
764 }
765
766 pub async fn publish_parent_replayable_event(
769 &self,
770 parent_session_id: &str,
771 event: AgentEvent,
772 ) {
773 self.ctx
774 .replayable_event_publisher()
775 .publish(parent_session_id, event)
776 .await;
777 }
778
779 pub async fn lock_child_launch(&self, child_id: &str) -> ChildLaunchGuard {
780 self.launch_locks.acquire(child_id).await
781 }
782
783 pub async fn publish_queued_child_cancellation(&self, parent_id: &str, child_id: &str) {
784 let publisher = self.ctx.replayable_event_publisher();
785 publish_child_completion_parts(
786 &publisher,
787 self.ctx.completion_handler.clone(),
788 parent_id.to_string(),
789 child_id.to_string(),
790 "cancelled".to_string(),
791 Some("Cancelled by parent before activation".to_string()),
792 )
793 .await;
794 }
795
796 pub(crate) fn launch_reserved(
800 &self,
801 job: SpawnJob,
802 reservation: SessionExecutionReservation,
803 ) -> tokio::task::JoinHandle<()> {
804 let ctx = self.ctx.clone();
805 tokio::spawn(async move {
806 Self::prepare_child_launch(&ctx, &job).await;
807 ctx.replayable_event_publisher()
808 .publish(
809 &job.parent_session_id,
810 AgentEvent::SubAgentStarted {
811 parent_session_id: job.parent_session_id.clone(),
812 child_session_id: job.child_session_id.clone(),
813 title: None,
814 },
815 )
816 .await;
817 if let Err(error) =
818 crate::sdk::spawn::run_child_spawn_reserved(ctx, job.clone(), reservation).await
819 {
820 tracing::warn!(
821 parent_session_id = %job.parent_session_id,
822 child_session_id = %job.child_session_id,
823 %error,
824 "reserved child activation failed"
825 );
826 }
827 })
828 }
829}
830
831async fn needs_fallback_start(
832 announcement_ready: oneshot::Receiver<()>,
833 start_published: &AtomicBool,
834) -> bool {
835 announcement_ready.await.is_err() && !start_published.load(Ordering::SeqCst)
836}
837
838fn invoke_child_run_launch_hook(
839 hook: Option<&Arc<dyn ChildRunLaunchHook>>,
840 job: &SpawnJob,
841 child_tx: broadcast::Sender<AgentEvent>,
842) {
843 let Some(hook) = hook else {
844 return;
845 };
846 let invoked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
847 hook.before_child_launch(job, child_tx);
848 }));
849 if invoked.is_err() {
850 tracing::error!(
851 parent_session_id = %job.parent_session_id,
852 child_session_id = %job.child_session_id,
853 "child launch hook panicked; continuing with canonical execution"
854 );
855 }
856}
857
858async fn reserve_and_send(
863 tx: &mpsc::Sender<QueuedSpawnJob>,
864 job: SpawnJob,
865 title: Option<String>,
866 gate: Option<&AdmissionGate>,
867 launch_generation: Option<u64>,
868 admission_check: Option<&QueueAdmissionCheck>,
869) -> Result<AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>, String> {
870 match reserve_queue_permit(tx, gate).await? {
871 AdmissionCommit::Committed(permit) => send_reserved(
872 permit,
873 job,
874 title,
875 gate,
876 launch_generation,
877 admission_check,
878 Ok(()),
879 ),
880 AdmissionCommit::AlreadyCommitted => Ok(AdmissionCommit::AlreadyCommitted),
881 AdmissionCommit::Cancelled => Ok(AdmissionCommit::Cancelled),
882 }
883}
884
885async fn reserve_queue_permit<'a>(
886 tx: &'a mpsc::Sender<QueuedSpawnJob>,
887 gate: Option<&AdmissionGate>,
888) -> Result<AdmissionCommit<mpsc::Permit<'a, QueuedSpawnJob>>, String> {
889 if let Some(gate) = gate {
890 if gate.is_committed() {
891 return Ok(AdmissionCommit::AlreadyCommitted);
892 }
893 if gate.is_cancelled() {
894 return Ok(AdmissionCommit::Cancelled);
895 }
896 }
897 let permit = match gate {
898 Some(gate) => tokio::select! {
899 biased;
900 _ = gate.cancelled() => return Ok(AdmissionCommit::Cancelled),
901 permit = tx.reserve() => permit,
902 },
903 None => tx.reserve().await,
904 }
905 .map_err(|_| "spawn scheduler is not running".to_string())?;
906 Ok(AdmissionCommit::Committed(permit))
907}
908
909fn send_reserved(
910 permit: mpsc::Permit<'_, QueuedSpawnJob>,
911 job: SpawnJob,
912 title: Option<String>,
913 gate: Option<&AdmissionGate>,
914 launch_generation: Option<u64>,
915 admission_check: Option<&QueueAdmissionCheck>,
916 preflight: Result<(), String>,
917) -> Result<AdmissionCommit<(oneshot::Sender<()>, Arc<AtomicBool>)>, String> {
918 let (announce_ready, announcement_ready) = oneshot::channel();
919 let start_published = Arc::new(AtomicBool::new(false));
920 let queued = QueuedSpawnJob {
921 job,
922 title,
923 launch_generation,
924 announcement_ready,
925 start_published: start_published.clone(),
926 };
927 let commit = || {
928 preflight?;
929 if let Some(check) = admission_check {
930 check()?;
931 }
932 permit.send(queued);
933 Ok((announce_ready, start_published))
934 };
935 match gate {
936 Some(gate) => gate.commit(commit),
937 None => commit().map(AdmissionCommit::Committed),
938 }
939}
940
941#[derive(Debug, Clone, Copy)]
942pub(crate) struct ChildWatchdogPolicy {
943 check_interval_secs: i64,
944 pub(crate) max_total_secs: i64,
948 pub(crate) max_idle_secs: i64,
949}
950
951impl Default for ChildWatchdogPolicy {
952 fn default() -> Self {
953 Self {
954 check_interval_secs: 15,
955 max_total_secs: 60 * 60,
959 max_idle_secs: 15 * 60,
961 }
962 }
963}
964
965fn metadata_i64(session: &Session, key: &str) -> Option<i64> {
966 session
967 .metadata
968 .get(key)
969 .and_then(|value| value.trim().parse::<i64>().ok())
970 .filter(|value| *value > 0)
971}
972
973pub(crate) fn watchdog_policy_for_session(session: &Session) -> ChildWatchdogPolicy {
974 let mut policy = ChildWatchdogPolicy::default();
975 if let Some(value) = metadata_i64(session, "child_watchdog.max_total_secs") {
976 policy.max_total_secs = value;
977 }
978 if let Some(value) = metadata_i64(session, "child_watchdog.max_idle_secs") {
979 policy.max_idle_secs = value;
980 }
981 if let Some(value) = metadata_i64(session, "child_watchdog.check_interval_secs") {
982 policy.check_interval_secs = value;
983 }
984 policy
985}
986
987pub(crate) async fn publish_child_completion(
988 publisher: &super::session_events::ReplayableSessionEventPublisher,
989 completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
990 completion: ChildCompletion,
991) {
992 publisher
993 .publish(
994 &completion.parent_session_id,
995 AgentEvent::SubAgentCompleted {
996 parent_session_id: completion.parent_session_id.clone(),
997 child_session_id: completion.child_session_id.clone(),
998 status: completion.status.clone(),
999 error: completion.error.clone(),
1000 },
1001 )
1002 .await;
1003
1004 if let Some(handler) = completion_handler {
1005 use futures::FutureExt;
1012 let parent_session_id = completion.parent_session_id.clone();
1013 let child_session_id = completion.child_session_id.clone();
1014 if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
1015 .catch_unwind()
1016 .await
1017 .is_err()
1018 {
1019 tracing::error!(
1020 %parent_session_id,
1021 %child_session_id,
1022 "child completion handler panicked; child-wait watchdog will backstop the parent wake"
1023 );
1024 }
1025 }
1026}
1027
1028pub(crate) async fn publish_child_completion_parts(
1029 publisher: &super::session_events::ReplayableSessionEventPublisher,
1030 completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
1031 parent_session_id: String,
1032 child_session_id: String,
1033 status: String,
1034 error: Option<String>,
1035) {
1036 publish_child_completion(
1037 publisher,
1038 completion_handler,
1039 ChildCompletion {
1040 parent_session_id,
1041 child_session_id,
1042 status,
1043 error,
1044 completed_at: Utc::now(),
1045 source: None,
1046 },
1047 )
1048 .await;
1049}
1050
1051pub(crate) async fn watch_child_liveness(
1052 parent_session_id: String,
1053 child_session_id: String,
1054 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
1055 cancel_token: CancellationToken,
1056 timeout_reason: Arc<RwLock<Option<String>>>,
1057 done: CancellationToken,
1058 policy: ChildWatchdogPolicy,
1059) {
1060 let mut ticker =
1061 tokio::time::interval(Duration::from_secs(policy.check_interval_secs.max(1) as u64));
1062 ticker.tick().await;
1064
1065 loop {
1066 tokio::select! {
1067 _ = done.cancelled() => return,
1068 _ = ticker.tick() => {
1069 if cancel_token.is_cancelled() {
1070 return;
1071 }
1072
1073 let snapshot = {
1074 let guard = runners.read().await;
1075 guard.get(&child_session_id).cloned()
1076 };
1077 let Some(runner) = snapshot else {
1078 return;
1079 };
1080 if !matches!(runner.status, super::runner_state::AgentStatus::Running) {
1081 return;
1082 }
1083
1084 let now = Utc::now();
1085 let total_secs = now.signed_duration_since(runner.started_at).num_seconds();
1086 if total_secs >= policy.max_total_secs {
1087 let reason = format!(
1088 "Child session timed out after {} seconds (max_total_secs={})",
1089 total_secs, policy.max_total_secs
1090 );
1091 tracing::warn!(
1092 parent_session_id = %parent_session_id,
1093 child_session_id = %child_session_id,
1094 reason = %reason,
1095 "child session total timeout; cancelling child runner"
1096 );
1097 *timeout_reason.write().await = Some(reason);
1098 cancel_token.cancel();
1099 return;
1100 }
1101
1102 let last_activity_at = runner.last_activity_at().unwrap_or(runner.started_at);
1103 let idle_secs = now.signed_duration_since(last_activity_at).num_seconds();
1104 if idle_secs >= policy.max_idle_secs {
1105 let reason = format!(
1106 "Child session idle timeout after {} seconds without events (max_idle_secs={})",
1107 idle_secs, policy.max_idle_secs
1108 );
1109 tracing::warn!(
1110 parent_session_id = %parent_session_id,
1111 child_session_id = %child_session_id,
1112 reason = %reason,
1113 last_tool_name = ?runner.last_tool_name,
1114 last_tool_phase = ?runner.last_tool_phase,
1115 round_count = runner.round_count,
1116 "child session idle timeout; cancelling child runner"
1117 );
1118 *timeout_reason.write().await = Some(reason);
1119 cancel_token.cancel();
1120 return;
1121 }
1122 }
1123 }
1124 }
1125}
1126
1127async fn run_spawn_job(
1134 ctx: SpawnContext,
1135 job: SpawnJob,
1136 launch_guard: Option<ChildLaunchGuard>,
1137) -> Result<(), String> {
1138 match launch_guard {
1139 Some(guard) => crate::sdk::spawn::run_child_spawn_fenced(ctx, job, guard).await,
1140 None => crate::sdk::spawn::run_child_spawn(ctx, job).await,
1141 }
1142}
1143
1144#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1145enum QueuedLaunchState {
1146 Eligible,
1147 AlreadyStarted,
1148 Stale,
1149}
1150
1151async fn queued_launch_state(
1152 ctx: &SpawnContext,
1153 job: &SpawnJob,
1154 generation: u64,
1155) -> Result<QueuedLaunchState, String> {
1156 let child = ctx
1157 .agent
1158 .storage()
1159 .load_session(&job.child_session_id)
1160 .await
1161 .map_err(|error| format!("load queued child launch fence: {error}"))?;
1162 let Some(child) = child else {
1163 return Ok(QueuedLaunchState::Stale);
1164 };
1165 if child.kind != bamboo_domain::SessionKind::Child
1166 || child.parent_session_id.as_deref() != Some(job.parent_session_id.as_str())
1167 || child.child_launch_generation() != generation
1168 || child.is_child_launch_cancelled(generation)
1169 {
1170 return Ok(QueuedLaunchState::Stale);
1171 }
1172 match child.last_run_status().as_deref() {
1173 Some("pending") => Ok(QueuedLaunchState::Eligible),
1174 Some("running" | "suspended") => Ok(QueuedLaunchState::AlreadyStarted),
1175 _ => Ok(QueuedLaunchState::Stale),
1176 }
1177}
1178
1179#[cfg(test)]
1180mod launch_hook_tests {
1181 use super::*;
1182 use std::sync::atomic::{AtomicUsize, Ordering};
1183
1184 fn job() -> SpawnJob {
1185 SpawnJob {
1186 parent_session_id: "parent".to_string(),
1187 child_session_id: "child".to_string(),
1188 model: "test".to_string(),
1189 disabled_tools: None,
1190 }
1191 }
1192
1193 struct PanickingHook {
1194 calls: AtomicUsize,
1195 }
1196
1197 impl ChildRunLaunchHook for PanickingHook {
1198 fn before_child_launch(
1199 &self,
1200 _job: &SpawnJob,
1201 _child_events: broadcast::Sender<AgentEvent>,
1202 ) {
1203 self.calls.fetch_add(1, Ordering::SeqCst);
1204 panic!("injected launch hook panic");
1205 }
1206 }
1207
1208 #[test]
1209 fn launch_hook_panic_is_contained() {
1210 let hook = Arc::new(PanickingHook {
1211 calls: AtomicUsize::new(0),
1212 });
1213 let hook_port: Arc<dyn ChildRunLaunchHook> = hook.clone();
1214 let (child_tx, _child_rx) = broadcast::channel(1);
1215
1216 invoke_child_run_launch_hook(Some(&hook_port), &job(), child_tx);
1217
1218 assert_eq!(hook.calls.load(Ordering::SeqCst), 1);
1219 }
1220
1221 #[tokio::test]
1222 async fn closed_scheduler_does_not_prepare_phantom_launch() {
1223 let (tx, rx) = mpsc::channel(1);
1224 drop(rx);
1225 let result = reserve_and_send(&tx, job(), None, None, None, None).await;
1226
1227 assert_eq!(result.unwrap_err(), "spawn scheduler is not running");
1228 }
1229
1230 #[tokio::test]
1231 async fn cancellation_while_waiting_for_queue_capacity_never_sends() {
1232 let (tx, mut rx) = mpsc::channel(1);
1233 let (_, announcement_ready) = oneshot::channel();
1234 tx.send(QueuedSpawnJob {
1235 job: job(),
1236 title: None,
1237 launch_generation: None,
1238 announcement_ready,
1239 start_published: Arc::new(AtomicBool::new(false)),
1240 })
1241 .await
1242 .unwrap();
1243 let gate = Arc::new(AdmissionGate::default());
1244 let waiting = {
1245 let tx = tx.clone();
1246 let gate = gate.clone();
1247 tokio::spawn(async move {
1248 reserve_and_send(&tx, job(), None, Some(&gate), None, None).await
1249 })
1250 };
1251 tokio::task::yield_now().await;
1252 assert!(!waiting.is_finished());
1253 gate.cancel_if_pending();
1254 assert!(matches!(
1255 tokio::time::timeout(Duration::from_millis(200), waiting)
1256 .await
1257 .expect("cancellation must wake without freeing queue capacity")
1258 .unwrap()
1259 .unwrap(),
1260 AdmissionCommit::Cancelled
1261 ));
1262 rx.recv().await.unwrap();
1263 assert!(matches!(
1264 rx.try_recv(),
1265 Err(mpsc::error::TryRecvError::Empty)
1266 ));
1267 }
1268
1269 #[tokio::test]
1270 async fn committed_queue_gate_does_not_enqueue_a_duplicate_job() {
1271 let (tx, mut rx) = mpsc::channel(1);
1272 let gate = AdmissionGate::default();
1273 assert!(matches!(
1274 reserve_and_send(&tx, job(), None, Some(&gate), None, None)
1275 .await
1276 .unwrap(),
1277 AdmissionCommit::Committed(_)
1278 ));
1279 assert!(matches!(
1280 reserve_and_send(&tx, job(), None, Some(&gate), None, None)
1281 .await
1282 .unwrap(),
1283 AdmissionCommit::AlreadyCommitted
1284 ));
1285 rx.recv().await.unwrap();
1286 assert!(matches!(
1287 rx.try_recv(),
1288 Err(mpsc::error::TryRecvError::Empty)
1289 ));
1290 }
1291
1292 async fn assert_expired_queue_check(gate: Option<&AdmissionGate>) {
1293 let (tx, mut rx) = mpsc::channel(1);
1294 let _first = reserve_and_send(&tx, job(), None, None, None, None)
1295 .await
1296 .unwrap();
1297 let expired = Arc::new(AtomicBool::new(false));
1298 let calls = Arc::new(AtomicUsize::new(0));
1299 let check = {
1300 let expired = expired.clone();
1301 let calls = calls.clone();
1302 move || {
1303 calls.fetch_add(1, Ordering::SeqCst);
1304 if expired.load(Ordering::SeqCst) {
1305 Err("Host admission expired".to_string())
1306 } else {
1307 Ok(())
1308 }
1309 }
1310 };
1311 let waiting = reserve_and_send(&tx, job(), None, gate, None, Some(&check));
1312 tokio::pin!(waiting);
1313 assert!(futures::poll!(waiting.as_mut()).is_pending());
1314 assert_eq!(calls.load(Ordering::SeqCst), 0);
1315
1316 expired.store(true, Ordering::SeqCst);
1319 rx.recv().await.unwrap();
1320 assert_eq!(waiting.await.unwrap_err(), "Host admission expired");
1321 assert_eq!(calls.load(Ordering::SeqCst), 1);
1322 assert!(matches!(
1323 rx.try_recv(),
1324 Err(mpsc::error::TryRecvError::Empty)
1325 ));
1326 }
1327
1328 #[tokio::test]
1329 async fn queue_admission_check_rejects_expiry_after_capacity_wait() {
1330 assert_expired_queue_check(None).await;
1331 }
1332
1333 #[tokio::test]
1334 async fn queue_admission_check_rejection_preserves_cancellation_gate() {
1335 let gate = AdmissionGate::default();
1336 assert_expired_queue_check(Some(&gate)).await;
1337 assert!(!gate.is_committed());
1338 gate.cancel_if_pending();
1339 assert!(gate.is_cancelled());
1340 }
1341
1342 #[tokio::test]
1343 async fn queue_admission_check_runs_once_for_successful_send() {
1344 let (tx, mut rx) = mpsc::channel(1);
1345 let gate = AdmissionGate::default();
1346 let calls = Arc::new(AtomicUsize::new(0));
1347 let check = {
1348 let calls = calls.clone();
1349 move || {
1350 calls.fetch_add(1, Ordering::SeqCst);
1351 Ok(())
1352 }
1353 };
1354 assert!(matches!(
1355 reserve_and_send(&tx, job(), None, Some(&gate), None, Some(&check))
1356 .await
1357 .unwrap(),
1358 AdmissionCommit::Committed(_)
1359 ));
1360 assert!(matches!(
1361 reserve_and_send(&tx, job(), None, Some(&gate), None, Some(&check))
1362 .await
1363 .unwrap(),
1364 AdmissionCommit::AlreadyCommitted
1365 ));
1366 assert_eq!(calls.load(Ordering::SeqCst), 1);
1367 rx.recv().await.unwrap();
1368 assert!(matches!(
1369 rx.try_recv(),
1370 Err(mpsc::error::TryRecvError::Empty)
1371 ));
1372 }
1373
1374 mod queue_capacity_revalidation {
1375 use super::*;
1376
1377 async fn reserve<'a>(
1378 tx: &'a mpsc::Sender<QueuedSpawnJob>,
1379 gate: &AdmissionGate,
1380 ) -> mpsc::Permit<'a, QueuedSpawnJob> {
1381 match reserve_queue_permit(tx, Some(gate)).await.unwrap() {
1382 AdmissionCommit::Committed(permit) => permit,
1383 _ => panic!("fresh gate must reserve without committing"),
1384 }
1385 }
1386
1387 #[tokio::test]
1388 async fn reservation_does_not_commit_the_original_gate() {
1389 let (tx, mut rx) = mpsc::channel(1);
1390 let gate = AdmissionGate::default();
1391 let permit = reserve(&tx, &gate).await;
1392 assert!(!gate.is_committed());
1393 assert!(!gate.is_cancelled());
1394 assert_eq!(tx.capacity(), 0);
1395 assert!(matches!(
1396 rx.try_recv(),
1397 Err(mpsc::error::TryRecvError::Empty)
1398 ));
1399 assert!(matches!(
1400 send_reserved(permit, job(), None, Some(&gate), Some(7), None, Ok(())).unwrap(),
1401 AdmissionCommit::Committed(_)
1402 ));
1403 assert!(gate.is_committed());
1404 assert_eq!(rx.recv().await.unwrap().launch_generation, Some(7));
1405 }
1406
1407 #[tokio::test]
1408 async fn preflight_error_keeps_pending_gate_and_releases_capacity() {
1409 let (tx, mut rx) = mpsc::channel(1);
1410 let gate = AdmissionGate::default();
1411 let permit = reserve(&tx, &gate).await;
1412 let error = send_reserved(
1413 permit,
1414 job(),
1415 None,
1416 Some(&gate),
1417 Some(7),
1418 None,
1419 Err("fresh Host resources rejected".into()),
1420 )
1421 .expect_err("resource failure cannot commit");
1422 assert_eq!(error, "fresh Host resources rejected");
1423 assert!(!gate.is_committed());
1424 assert!(!gate.is_cancelled());
1425 assert_eq!(tx.capacity(), 1);
1426 assert!(matches!(
1427 rx.try_recv(),
1428 Err(mpsc::error::TryRecvError::Empty)
1429 ));
1430 }
1431
1432 #[tokio::test]
1433 async fn concurrent_commit_wins_over_a_late_preflight_error() {
1434 let (tx, mut rx) = mpsc::channel(2);
1435 let gate = AdmissionGate::default();
1436 let late = reserve(&tx, &gate).await;
1437 let first = reserve(&tx, &gate).await;
1438 assert!(matches!(
1439 send_reserved(first, job(), None, Some(&gate), Some(7), None, Ok(())).unwrap(),
1440 AdmissionCommit::Committed(_)
1441 ));
1442 let calls = Arc::new(AtomicUsize::new(0));
1443 let check_calls = calls.clone();
1444 let check = move || {
1445 check_calls.fetch_add(1, Ordering::SeqCst);
1446 Ok(())
1447 };
1448 assert!(matches!(
1449 send_reserved(
1450 late,
1451 job(),
1452 None,
1453 Some(&gate),
1454 Some(7),
1455 Some(&check),
1456 Err("resource check finished after another commit".into())
1457 )
1458 .unwrap(),
1459 AdmissionCommit::AlreadyCommitted
1460 ));
1461 assert_eq!(calls.load(Ordering::SeqCst), 0);
1462 assert_eq!(rx.recv().await.unwrap().launch_generation, Some(7));
1463 assert!(matches!(
1464 rx.try_recv(),
1465 Err(mpsc::error::TryRecvError::Empty)
1466 ));
1467 assert_eq!(tx.capacity(), 2);
1468 }
1469
1470 #[tokio::test]
1471 async fn cancellation_wins_over_a_late_preflight_error() {
1472 let (tx, mut rx) = mpsc::channel(1);
1473 let gate = AdmissionGate::default();
1474 let permit = reserve(&tx, &gate).await;
1475 gate.cancel_if_pending();
1476 assert!(matches!(
1477 send_reserved(
1478 permit,
1479 job(),
1480 None,
1481 Some(&gate),
1482 Some(7),
1483 None,
1484 Err("preflight observed failure after cancellation".into())
1485 )
1486 .unwrap(),
1487 AdmissionCommit::Cancelled
1488 ));
1489 assert_eq!(tx.capacity(), 1);
1490 assert!(matches!(
1491 rx.try_recv(),
1492 Err(mpsc::error::TryRecvError::Empty)
1493 ));
1494 }
1495
1496 #[tokio::test]
1497 async fn successful_preflight_cannot_bypass_the_final_live_fence() {
1498 let (tx, mut rx) = mpsc::channel(1);
1499 let gate = AdmissionGate::default();
1500 let permit = reserve(&tx, &gate).await;
1501 let check = || Err("Ready policy expired during resource preflight".into());
1502 let error = send_reserved(
1503 permit,
1504 job(),
1505 None,
1506 Some(&gate),
1507 Some(7),
1508 Some(&check),
1509 Ok(()),
1510 )
1511 .expect_err("the final synchronous fence remains authoritative");
1512 assert_eq!(error, "Ready policy expired during resource preflight");
1513 assert!(!gate.is_committed());
1514 assert!(!gate.is_cancelled());
1515 assert_eq!(tx.capacity(), 1);
1516 assert!(matches!(
1517 rx.try_recv(),
1518 Err(mpsc::error::TryRecvError::Empty)
1519 ));
1520 }
1521 }
1522
1523 #[tokio::test]
1524 async fn producer_drop_after_start_publication_does_not_require_fallback_start() {
1525 let (announce_ready, announcement_ready) = oneshot::channel::<()>();
1526 let start_published = Arc::new(AtomicBool::new(false));
1527 let runners = Arc::new(RwLock::new(HashMap::new()));
1528 let senders = Arc::new(RwLock::new(HashMap::new()));
1529 let (sender, mut receiver) = broadcast::channel(4);
1530 senders.write().await.insert("parent".to_string(), sender);
1531 let publisher = super::super::session_events::ReplayableSessionEventPublisher::new(
1532 runners, senders, None,
1533 );
1534 publisher
1538 .publish(
1539 "parent",
1540 AgentEvent::SubAgentStarted {
1541 parent_session_id: "parent".to_string(),
1542 child_session_id: "child".to_string(),
1543 title: None,
1544 },
1545 )
1546 .await;
1547 start_published.store(true, Ordering::SeqCst);
1548 drop(announce_ready);
1549 assert!(!needs_fallback_start(announcement_ready, &start_published).await);
1550 assert!(matches!(
1551 receiver.try_recv(),
1552 Ok(AgentEvent::SubAgentStarted { .. })
1553 ));
1554 assert!(matches!(
1555 receiver.try_recv(),
1556 Err(broadcast::error::TryRecvError::Empty)
1557 ));
1558
1559 let (announce_ready, announcement_ready) = oneshot::channel::<()>();
1560 let start_published = AtomicBool::new(false);
1561 drop(announce_ready);
1562 assert!(needs_fallback_start(announcement_ready, &start_published).await);
1563 }
1564}