Skip to main content

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::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
30/// A process-local Host admission check, evaluated at the synchronous queue
31/// commit after capacity waits. It must not block or perform I/O and cannot
32/// reconstruct authority from a persisted job.
33pub type QueueAdmissionCheck = dyn Fn() -> Result<(), String> + Send + Sync;
34
35/// One in-memory queue slot, not an execution grant. Dropping it releases
36/// capacity without publishing a job or Start. The original cancellation gate
37/// stays bound to the reservation and cannot be replaced before commit.
38pub struct QueueCapacityReservation<'a> {
39    scheduler: &'a SpawnScheduler,
40    gate: Option<&'a AdmissionGate>,
41    permit: mpsc::Permit<'a, QueuedSpawnJob>,
42}
43
44impl QueueCapacityReservation<'_> {
45    /// A terminal original gate needs no additional Host I/O. Commit repeats
46    /// this arbitration atomically, including when preflight returned an error.
47    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    /// The caller may perform asynchronous Host validation while holding this
60    /// slot. This method retains the final synchronous, non-I/O admission fence
61    /// and publishes Start only after that commit succeeds.
62    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    /// Tool names to hide from the LLM schema for this child session.
97    /// Computed from the child's `subagent_type` profile policy.
98    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/// Serializes a queued launch's durable eligibility check and runner
110/// reservation against parent cancellation of that same logical child.
111#[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
148/// Optional application-layer preparation for a child run launched through
149/// the canonical scheduler.
150///
151/// The engine owns runner reservation and execution. Applications may use
152/// this synchronous, no-fail hook to attach observers to the already-created
153/// child event sender (for example, the server's always-on notification
154/// relay) without introducing an engine dependency on application services.
155pub trait ChildRunLaunchHook: Send + Sync {
156    fn before_child_launch(&self, job: &SpawnJob, child_events: broadcast::Sender<AgentEvent>);
157
158    /// Optional queue boundary observer. The production notification hook is
159    /// a no-op; deterministic scheduler tests can hold dequeue before the
160    /// cancellation fence is acquired.
161    fn before_queued_dequeue(&self, _job: &SpawnJob) -> BoxFuture<'static, ()> {
162        Box::pin(async {})
163    }
164
165    /// Optional observer inside the cancellation fence, immediately before
166    /// runner reservation. It must never make a model decision or perform
167    /// unbounded work while this guard is held.
168    fn before_queued_reservation(&self, _job: &SpawnJob) -> BoxFuture<'static, ()> {
169        Box::pin(async {})
170    }
171}
172
173/// Runtime-scoped durable inbox resources used by external actor drivers.
174///
175/// The host store remains canonical. A worker confirmation is only permission
176/// for the driver to checkpoint that canonical logical Session and then ack the
177/// exact claim; it never turns transport state into authority.
178#[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    /// The production Host's exact SessionRepository write coordinator.
185    /// Canonical Worker checkpoints fail closed when this is absent.
186    pub parent_question_lock: Option<Arc<bamboo_storage::LockedSessionStore>>,
187}
188
189/// Trait for external child session runtimes (e.g. A2A, CLI adapters).
190///
191/// Implementors are responsible for emitting AgentEvents via `event_tx`
192/// and respecting the `cancel_token`.
193#[async_trait::async_trait]
194pub trait ExternalChildRunner: Send + Sync {
195    /// Capability of the actually assembled backend, independent of hot config.
196    /// Only the in-process backend consumes canonical Ticket request ceilings.
197    fn supports_native_ticket_budget(&self) -> bool {
198        false
199    }
200
201    /// App catalog/read services for the owned local backend only.
202    /// Process and remote adapters deliberately retain the default no-op.
203    fn set_local_child_skill_factory(
204        &self,
205        _factory: Option<Arc<dyn crate::external_agents::local_task::LocalChildSkillFactory>>,
206    ) {
207    }
208
209    /// Bind the Host's in-process execution resources. Process/remote adapters
210    /// deliberately ignore this capability and retain their existing mode.
211    fn set_local_task_runtime(
212        &self,
213        _runtime: Option<crate::external_agents::local_task::LocalTaskRuntime>,
214    ) {
215    }
216
217    /// Acquire this backend's exact durable writer before the canonical
218    /// running checkpoint. A default writer is never substituted after claim.
219    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    /// Roll back a prepared execution whose initial canonical checkpoint failed.
229    /// Local implementations must close the launch gate and prove resource drain.
230    async fn abandon_child_execution(
231        &self,
232        _session: &Session,
233        _run_id: &str,
234    ) -> Result<(), String> {
235        Ok(())
236    }
237
238    /// Record a setup panic without inventing a successful resource stop.
239    fn quarantine_child_execution(&self, _session: &Session, _run_id: &str) {}
240
241    /// Keep a local coordinator handle owned until it really exits. Returning
242    /// the handle leaves supervision to an outer composite or legacy caller.
243    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    /// End the exact backend writer after terminal persistence and the history
253    /// barrier, before publishing runner replacement eligibility.
254    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    /// Stop admission, request cancellation, and await owned task/resource
264    /// completion. This is called by the Host's graceful shutdown path.
265    async fn shutdown(&self) -> Result<(), String> {
266        Ok(())
267    }
268
269    /// Opt-in, host-owned scope authority. It is never copied into the Worker.
270    fn set_ticket_service(&self, _service: Option<Arc<bamboo_tickets::TicketService>>) {}
271    /// A narrow, pre-persistence compatibility check on the actual registered
272    /// runner. Unknown/custom routes fail closed for required one-shot packets.
273    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    /// Returns true if this runner should handle the given child session.
280    async fn should_handle(&self, session: &Session) -> bool;
281
282    /// Execute the child session using an external runtime.
283    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    /// Persist a Host-only broker receipt before the final Child checkpoint.
292    /// `Ok(true)` selects the append-safe final checkpoint. An unprepared
293    /// broker terminal must never be ACKed after a merely ordinary save.
294    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    /// Commit and read back any Host-only proof required for this terminal
303    /// delivery after the final Child Session save. Failure (including timeout)
304    /// prevents a committed completion source and successful runner finalization.
305    /// When the source is pending, publish its terminal status and seal with an
306    /// exact-snapshot canonical write before returning; the SDK verifies readback.
307    /// Legacy/default runners need no separate proof and retain save-only behavior.
308    async fn commit_durable_child_delivery(
309        &self,
310        _session: &Session,
311        _activation_run_id: &str,
312    ) -> Result<(), String> {
313        Ok(())
314    }
315
316    /// ACK transport delivery only after the final Child save and required
317    /// Host proof commit both succeed. The exact Child birth and activation run
318    /// identify the pending receipt; failed commitment leaves it unacknowledged.
319    /// ACK delivery failure does not invalidate an already-proven completion.
320    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    /// Bind this runner's per-run escalation host bridge (#68). A nested worker's
330    /// `run()` installs its OWN host bridge here so the runner can hand it to each
331    /// grandchild's `drive()` AT SPAWN time (captured into the drive task, not read
332    /// later), letting the grandchild re-proxy a non-bypass approval request UP to
333    /// its parent run for its whole lifetime — even when it outlives the run that
334    /// spawned it. Default no-op for runners that don't escalate (e.g. A2A).
335    fn set_escalation_bridge(&self, _bridge: Option<bamboo_subagent::executor::HostBridge>) {}
336
337    /// Bind the owning runtime's canonical SessionInbox resources. Actor
338    /// runners use this to bridge active local/remote/warm workers without a
339    /// process-global live-session registry.
340    fn set_session_inbox_runtime(&self, _binding: Option<SessionInboxRuntimeBinding>) {}
341
342    /// Optional actual host Store for the local zero-tool named-profile route.
343    /// This does not enable owned Inbox claims or grant authority to custom runners.
344    fn set_actor_directory_store(&self, _store: Option<Arc<bamboo_storage::SessionStoreV2>>) {}
345
346    /// Bind the application-owned logical SubAgent tool after its canonical
347    /// scheduler/adapter exist. Actor runners use it only for fenced child RPCs.
348    fn set_canonical_subagent_tool(&self, _tool: Option<Arc<dyn bamboo_agent_core::tools::Tool>>) {}
349
350    /// Optional host-owned, redacted canonical Actor event observer.
351    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    /// Optional application-layer completion hook. The engine still emits
369    /// `SubAgentCompleted` to the parent stream itself; this hook lets the
370    /// server persist parent wait state and resume the parent runner without
371    /// introducing an engine -> AppState dependency.
372    pub completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
373    /// Optional application observer setup shared by queued tool launches and
374    /// reserved idle SessionInbox activation.
375    pub child_run_launch_hook: Option<Arc<dyn ChildRunLaunchHook>>,
376    /// Optional inbox to the account-wide change feed. When present, durable
377    /// change events from child-session execution are mirrored onto the feed
378    /// for resumable multi-client sync.
379    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        // This is a route-only probe for the Child we have not created yet.
419        // Keep it out of execution/persistence, but select the same backend as
420        // that Child; a Root-shaped probe cannot match the native Child runner.
421        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        // The worker loop is a single point of failure for ALL child spawning:
439        // if it unwinds, queued jobs are dropped with no completion published
440        // and every later enqueue fails with "spawn scheduler is not running"
441        // for the rest of the process lifetime. Run each job on its own task
442        // and await the JoinHandle — a panicking job is isolated, keeps the
443        // worker alive, and still publishes a terminal error completion so the
444        // waiting parent is woken instead of stranded.
445        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                    // The producer publishes Start before returning to the
459                    // tool. If it is aborted after queue admission, the worker
460                    // publishes Start before any child execution/completion.
461                    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    /// Reserve queue capacity, publish the observable child Start, and only
546    /// then release the job to the worker. This preserves S→C ordering even
547    /// when child loading/execution fails immediately.
548    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    /// A child action supplies the generation it durably prepared before
570    /// queue admission. Replayed or cancelled generations cannot borrow a
571    /// newer explicit retry's pending status.
572    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    /// Preserve the caller's live Host deadline/revocation check across queue
584    /// backpressure. Existing callers retain the unchecked compatibility path.
585    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    /// Reserve backpressure capacity before asynchronous Host preflight. This
608    /// only reserves a Tokio queue slot: it creates no Task, lease or authority.
609    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                // From publication through the acknowledgement there is no
666                // await. A dropped producer after this point must not make
667                // the worker publish a duplicate Start.
668                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    /// Publish replayable parent-session state through the scheduler's shared
682    /// runner/cache/account/broadcast boundary.
683    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    /// Launch through the canonical child core using a runner slot already
712    /// reserved by SessionInbox activation. This bypasses only queue mechanics;
713    /// placement and execution still flow through `run_child_spawn`.
714    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
773/// The queue send is the admission commit. The cancellation fence and send run
774/// in one synchronous critical section after capacity is reserved. Observer
775/// setup and `SubAgentStarted` run only for a committed queue entry, before
776/// the child execution starts.
777async 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): the child-wait watchdog (#546) reads these limits to decide
860    // when a Running runner entry whose task died (frozen last_event_at) is
861    // stale beyond what the per-child liveness watchdog could still act on.
862    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            // Parent waits may be longer, but child execution owns its own
871            // liveness. A one hour total cap avoids indefinitely orphaned
872            // sub-session runners.
873            max_total_secs: 60 * 60,
874            // No child event for 15 minutes is considered stalled.
875            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        // Contain a panicking handler: this call frequently runs on the caller's
921        // only liveness-critical task (the child's terminal block, or the spawn
922        // scheduler worker for early failures). Unwinding here would kill that
923        // task after the child already looks terminal everywhere — the classic
924        // stranded-parent signature. The child-wait watchdog backstops the wake
925        // that a panicked handler failed to deliver.
926        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    // Skip the immediate tick.
978    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
1042/// Drive a single queued spawn job through the canonical child-spawn path.
1043///
1044/// ANTI-FORK: this is a 1-line delegator to [`crate::sdk::spawn::run_child_spawn`],
1045/// which is the single implementation of the spawn/execute/finalize logic. The
1046/// `SpawnScheduler` queue mechanics (above) remain here; the body lives in the SDK
1047/// core so both the scheduler and the ergonomic `ChildRunner` funnel into it.
1048async 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        // The Host expires while capacity is unavailable. Freeing capacity
1234        // must not turn that old permission into a new queue admission.
1235        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        // Reproduce the precise post-publish, pre-ack producer loss: the
1452        // worker observes a closed announcement channel, but Start is already
1453        // in the replayable event stream.
1454        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}