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/// A live native owner's once-only Agent execution entry, never persisted or
190/// reconstructed from Actor/Session timestamps. It proves no heartbeat or progress.
191/// Engine owns this physical fact; Orchestration decides its meaning. Server
192/// adapters only join it to current Ticket/Host authority. Stop, shutdown and
193/// restart provide no live sample and cannot recover execution permission.
194#[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/// A heartbeat is a successful renewal attempted after this exact Task entered.
213/// It measures owner liveness independently of the Task's progress.
214#[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/// External child session runtimes emit events and respect the cancel token.
231#[async_trait::async_trait]
232pub trait ExternalChildRunner: Send + Sync {
233    /// Live Host evidence only. Queue/preparation, persisted Actor timestamps
234    /// and absent/non-native owners cannot supply a Task-entry witness.
235    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    /// Read the current native owner and its last acknowledged renewal.
245    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    /// Capability of the actually assembled backend, independent of hot config.
259    /// Only the in-process backend consumes canonical Ticket request ceilings.
260    fn supports_native_ticket_budget(&self) -> bool {
261        false
262    }
263
264    /// App catalog/read services for the owned local backend only.
265    /// Process and remote adapters deliberately retain the default no-op.
266    fn set_local_child_skill_factory(
267        &self,
268        _factory: Option<Arc<dyn crate::external_agents::local_task::LocalChildSkillFactory>>,
269    ) {
270    }
271
272    /// Bind the Host's in-process execution resources. Process/remote adapters
273    /// deliberately ignore this capability and retain their existing mode.
274    fn set_local_task_runtime(
275        &self,
276        _runtime: Option<crate::external_agents::local_task::LocalTaskRuntime>,
277    ) {
278    }
279
280    /// Acquire this backend's exact durable writer before the canonical
281    /// running checkpoint. A default writer is never substituted after claim.
282    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    /// Roll back a prepared execution whose initial canonical checkpoint failed.
292    /// Local implementations must close the launch gate and prove resource drain.
293    async fn abandon_child_execution(
294        &self,
295        _session: &Session,
296        _run_id: &str,
297    ) -> Result<(), String> {
298        Ok(())
299    }
300
301    /// Record a setup panic without inventing a successful resource stop.
302    fn quarantine_child_execution(&self, _session: &Session, _run_id: &str) {}
303
304    /// Keep a local coordinator handle owned until it really exits. Returning
305    /// the handle leaves supervision to an outer composite or legacy caller.
306    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    /// End the exact backend writer after terminal persistence and the history
316    /// barrier, before publishing runner replacement eligibility.
317    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    /// Stop admission, request cancellation, and await owned task/resource
327    /// completion. This is called by the Host's graceful shutdown path.
328    async fn shutdown(&self) -> Result<(), String> {
329        Ok(())
330    }
331
332    /// Opt-in, host-owned scope authority. It is never copied into the Worker.
333    fn set_ticket_service(&self, _service: Option<Arc<bamboo_tickets::TicketService>>) {}
334    /// A narrow, pre-persistence compatibility check on the actual registered
335    /// runner. Unknown/custom routes fail closed for required one-shot packets.
336    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    /// Returns true if this runner should handle the given child session.
343    async fn should_handle(&self, session: &Session) -> bool;
344
345    /// Execute the child session using an external runtime.
346    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    /// Persist a Host-only broker receipt before the final Child checkpoint.
355    /// `Ok(true)` selects the append-safe final checkpoint. An unprepared
356    /// broker terminal must never be ACKed after a merely ordinary save.
357    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    /// Commit and read back any Host-only proof required for this terminal
366    /// delivery after the final Child Session save. Failure (including timeout)
367    /// prevents a committed completion source and successful runner finalization.
368    /// When the source is pending, publish its terminal status and seal with an
369    /// exact-snapshot canonical write before returning; the SDK verifies readback.
370    /// Legacy/default runners need no separate proof and retain save-only behavior.
371    async fn commit_durable_child_delivery(
372        &self,
373        _session: &Session,
374        _activation_run_id: &str,
375    ) -> Result<(), String> {
376        Ok(())
377    }
378
379    /// ACK transport delivery only after the final Child save and required
380    /// Host proof commit both succeed. The exact Child birth and activation run
381    /// identify the pending receipt; failed commitment leaves it unacknowledged.
382    /// ACK delivery failure does not invalidate an already-proven completion.
383    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    /// Bind this runner's per-run escalation host bridge (#68). A nested worker's
393    /// `run()` installs its OWN host bridge here so the runner can hand it to each
394    /// grandchild's `drive()` AT SPAWN time (captured into the drive task, not read
395    /// later), letting the grandchild re-proxy a non-bypass approval request UP to
396    /// its parent run for its whole lifetime — even when it outlives the run that
397    /// spawned it. Default no-op for runners that don't escalate (e.g. A2A).
398    fn set_escalation_bridge(&self, _bridge: Option<bamboo_subagent::executor::HostBridge>) {}
399
400    /// Bind the owning runtime's canonical SessionInbox resources. Actor
401    /// runners use this to bridge active local/remote/warm workers without a
402    /// process-global live-session registry.
403    fn set_session_inbox_runtime(&self, _binding: Option<SessionInboxRuntimeBinding>) {}
404
405    /// Optional actual host Store for the local zero-tool named-profile route.
406    /// This does not enable owned Inbox claims or grant authority to custom runners.
407    fn set_actor_directory_store(&self, _store: Option<Arc<bamboo_storage::SessionStoreV2>>) {}
408
409    /// Bind the application-owned logical SubAgent tool after its canonical
410    /// scheduler/adapter exist. Actor runners use it only for fenced child RPCs.
411    fn set_canonical_subagent_tool(&self, _tool: Option<Arc<dyn bamboo_agent_core::tools::Tool>>) {}
412
413    /// Optional host-owned, redacted canonical Actor event observer.
414    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    /// Optional application-layer completion hook. The engine still emits
432    /// `SubAgentCompleted` to the parent stream itself; this hook lets the
433    /// server persist parent wait state and resume the parent runner without
434    /// introducing an engine -> AppState dependency.
435    pub completion_handler: Option<Arc<dyn ChildCompletionHandler>>,
436    /// Optional application observer setup shared by queued tool launches and
437    /// reserved idle SessionInbox activation.
438    pub child_run_launch_hook: Option<Arc<dyn ChildRunLaunchHook>>,
439    /// Optional inbox to the account-wide change feed. When present, durable
440    /// change events from child-session execution are mirrored onto the feed
441    /// for resumable multi-client sync.
442    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        // This is a route-only probe for the Child we have not created yet.
504        // Keep it out of execution/persistence, but select the same backend as
505        // that Child; a Root-shaped probe cannot match the native Child runner.
506        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        // The worker loop is a single point of failure for ALL child spawning:
524        // if it unwinds, queued jobs are dropped with no completion published
525        // and every later enqueue fails with "spawn scheduler is not running"
526        // for the rest of the process lifetime. Run each job on its own task
527        // and await the JoinHandle — a panicking job is isolated, keeps the
528        // worker alive, and still publishes a terminal error completion so the
529        // waiting parent is woken instead of stranded.
530        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                    // The producer publishes Start before returning to the
544                    // tool. If it is aborted after queue admission, the worker
545                    // publishes Start before any child execution/completion.
546                    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    /// Reserve queue capacity, publish the observable child Start, and only
631    /// then release the job to the worker. This preserves S→C ordering even
632    /// when child loading/execution fails immediately.
633    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    /// A child action supplies the generation it durably prepared before
655    /// queue admission. Replayed or cancelled generations cannot borrow a
656    /// newer explicit retry's pending status.
657    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    /// Preserve the caller's live Host deadline/revocation check across queue
669    /// backpressure. Existing callers retain the unchecked compatibility path.
670    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    /// Reserve backpressure capacity before asynchronous Host preflight. This
693    /// only reserves a Tokio queue slot: it creates no Task, lease or authority.
694    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                // From publication through the acknowledgement there is no
751                // await. A dropped producer after this point must not make
752                // the worker publish a duplicate Start.
753                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    /// Publish replayable parent-session state through the scheduler's shared
767    /// runner/cache/account/broadcast boundary.
768    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    /// Launch through the canonical child core using a runner slot already
797    /// reserved by SessionInbox activation. This bypasses only queue mechanics;
798    /// placement and execution still flow through `run_child_spawn`.
799    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
858/// The queue send is the admission commit. The cancellation fence and send run
859/// in one synchronous critical section after capacity is reserved. Observer
860/// setup and `SubAgentStarted` run only for a committed queue entry, before
861/// the child execution starts.
862async 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): the child-wait watchdog (#546) reads these limits to decide
945    // when a Running runner entry whose task died (frozen last_event_at) is
946    // stale beyond what the per-child liveness watchdog could still act on.
947    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            // Parent waits may be longer, but child execution owns its own
956            // liveness. A one hour total cap avoids indefinitely orphaned
957            // sub-session runners.
958            max_total_secs: 60 * 60,
959            // No child event for 15 minutes is considered stalled.
960            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        // Contain a panicking handler: this call frequently runs on the caller's
1006        // only liveness-critical task (the child's terminal block, or the spawn
1007        // scheduler worker for early failures). Unwinding here would kill that
1008        // task after the child already looks terminal everywhere — the classic
1009        // stranded-parent signature. The child-wait watchdog backstops the wake
1010        // that a panicked handler failed to deliver.
1011        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    // Skip the immediate tick.
1063    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
1127/// Drive a single queued spawn job through the canonical child-spawn path.
1128///
1129/// ANTI-FORK: this is a 1-line delegator to [`crate::sdk::spawn::run_child_spawn`],
1130/// which is the single implementation of the spawn/execute/finalize logic. The
1131/// `SpawnScheduler` queue mechanics (above) remain here; the body lives in the SDK
1132/// core so both the scheduler and the ergonomic `ChildRunner` funnel into it.
1133async 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        // The Host expires while capacity is unavailable. Freeing capacity
1317        // must not turn that old permission into a new queue admission.
1318        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        // Reproduce the precise post-publish, pre-ack producer loss: the
1535        // worker observes a closed announcement channel, but Start is already
1536        // in the replayable event stream.
1537        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}