Skip to main content

zeph_common/
task_supervisor.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Supervised lifecycle task manager for long-running named services.
5//!
6//! [`TaskSupervisor`] manages named, long-lived background tasks (config watcher,
7//! scheduler loop, gateway, MCP connections, etc.) with restart policies, health
8//! snapshots, and graceful shutdown. Unlike `BackgroundSupervisor`
9//! (which is `&mut self`-only, lossy, and turn-scoped), `TaskSupervisor` is
10//! `Clone + Send + Sync` and designed for the full agent session lifetime.
11//!
12//! # Design rationale
13//!
14//! - **Shared handle**: `Arc<Inner>` interior allows passing the supervisor to bootstrap
15//!   code, TUI status display, and shutdown orchestration without lifetime coupling.
16//! - **Event-driven reap**: An internal mpsc channel delivers completion events to a
17//!   reap driver task; no polling interval required.
18//! - **No `JoinSet`**: Individual `JoinHandle`s per task enable per-name abort, status
19//!   tracking, and restart policies — `JoinSet` is better for homogeneous work.
20//! - **Mutex held briefly**: `parking_lot::Mutex` guards only bookkeeping operations
21//!   (insert/remove from `HashMap`). The lock is **never held across `.await`**.
22//!
23//! # Examples
24//!
25//! ```rust,no_run
26//! use std::time::Duration;
27//! use tokio_util::sync::CancellationToken;
28//! use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskSupervisor};
29//!
30//! # #[tokio::main]
31//! # async fn main() {
32//! let cancel = CancellationToken::new();
33//! let supervisor = TaskSupervisor::new(cancel.clone());
34//!
35//! supervisor.spawn(TaskDescriptor {
36//!     name: "my-service",
37//!     restart: RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) },
38//!     factory: || async { /* service loop */ },
39//! });
40//!
41//! // Graceful shutdown — waits up to 5 s for all tasks to stop.
42//! supervisor.shutdown_all(Duration::from_secs(5)).await;
43//! # }
44//! ```
45
46use std::collections::HashMap;
47use std::future::Future;
48use std::pin::Pin;
49use std::sync::Arc;
50use std::time::{Duration, Instant};
51
52use tokio::sync::{mpsc, oneshot};
53use tokio::task::AbortHandle;
54use tokio_util::sync::CancellationToken;
55use tracing::Instrument as _;
56
57use crate::BlockingSpawner;
58
59// ── Public types ─────────────────────────────────────────────────────────────
60
61/// Policy governing what happens when a supervised task completes or panics.
62///
63/// Used in [`TaskDescriptor`] to configure restart behaviour for a task.
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65#[non_exhaustive]
66pub enum RestartPolicy {
67    /// Task runs once; normal completion removes it from the registry.
68    RunOnce,
69    /// Task is restarted **only on panic**, up to `max` times.
70    ///
71    /// Normal completion (the future returns `()`) does **not** trigger a restart.
72    /// The task is removed from the registry on normal exit.
73    ///
74    /// A `max` of `0` means the task is monitored but **never** restarted —
75    /// a panic leaves the entry as `Failed` in the registry for observability.
76    /// Use `RunOnce` when you want the entry removed on completion.
77    ///
78    /// Restart delays follow **exponential backoff**: the delay before attempt `n`
79    /// is `base_delay * 2^(n-1)`, capped at [`MAX_RESTART_DELAY`].
80    ///
81    /// # Examples
82    ///
83    /// ```
84    /// use std::time::Duration;
85    /// use zeph_common::task_supervisor::RestartPolicy;
86    ///
87    /// // Restart up to 3 times with exponential backoff starting at 1 s.
88    /// let policy = RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) };
89    /// ```
90    Restart { max: u32, base_delay: Duration },
91}
92
93/// Maximum delay between restart attempts (caps exponential backoff).
94pub const MAX_RESTART_DELAY: Duration = Duration::from_mins(1);
95
96/// Configuration passed to [`TaskSupervisor::spawn`] to describe a supervised task.
97///
98/// `F` must be `Fn` (not `FnOnce`) to support restarts: the factory is called once on
99/// initial spawn and once per restart attempt.
100pub struct TaskDescriptor<F> {
101    /// Unique name for this task (e.g., `"config-watcher"`, `"scheduler-loop"`).
102    ///
103    /// Names must be `'static` — they are typically compile-time string literals.
104    /// Spawning a task with a name that already exists aborts the prior instance.
105    pub name: &'static str,
106    /// Restart policy applied when the task exits unexpectedly.
107    pub restart: RestartPolicy,
108    /// Factory called to produce a new future. Must be `Fn` for restart support.
109    pub factory: F,
110}
111
112/// Opaque handle to a single supervised task.
113///
114/// Can be used to abort the task by name independently of the supervisor.
115#[derive(Debug, Clone)]
116pub struct TaskHandle {
117    name: &'static str,
118    abort: AbortHandle,
119}
120
121impl TaskHandle {
122    /// Abort the task immediately.
123    pub fn abort(&self) {
124        tracing::debug!(task.name = self.name, "task aborted via handle");
125        self.abort.abort();
126    }
127
128    /// Return the task's name.
129    #[must_use]
130    pub const fn name(&self) -> &'static str {
131        self.name
132    }
133}
134
135/// Error returned by [`BlockingHandle::join`].
136#[derive(Debug, PartialEq, Eq)]
137#[non_exhaustive]
138pub enum BlockingError {
139    /// The task panicked before producing a result.
140    Panicked,
141    /// The supervisor (or the task's abort handle) was dropped before the task completed.
142    SupervisorDropped,
143}
144
145impl std::fmt::Display for BlockingError {
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        match self {
148            Self::Panicked => write!(f, "supervised blocking task panicked"),
149            Self::SupervisorDropped => write!(f, "supervisor dropped before task completed"),
150        }
151    }
152}
153
154impl std::error::Error for BlockingError {}
155
156/// Handle returned by [`TaskSupervisor::spawn_blocking`].
157///
158/// Awaiting [`BlockingHandle::join`] blocks until the OS-thread task produces a
159/// value. Dropping the handle without joining does **not** cancel the task — it
160/// continues to run on the blocking thread pool but the result is discarded.
161///
162/// A panic inside the closure is captured and returned as
163/// [`BlockingError::Panicked`] rather than propagating to the caller.
164pub struct BlockingHandle<R> {
165    rx: oneshot::Receiver<Result<R, BlockingError>>,
166    abort: AbortHandle,
167}
168
169impl<R> BlockingHandle<R> {
170    /// Await the task result.
171    ///
172    /// # Errors
173    ///
174    /// - [`BlockingError::Panicked`] — the task closure panicked.
175    /// - [`BlockingError::SupervisorDropped`] — the task was aborted or the
176    ///   supervisor was dropped before a value was produced.
177    pub async fn join(self) -> Result<R, BlockingError> {
178        self.rx
179            .await
180            .unwrap_or(Err(BlockingError::SupervisorDropped))
181    }
182
183    /// Non-blocking poll: return the result if the task has already finished, or `None`
184    /// if it is still running.
185    ///
186    /// This is the `BlockingHandle` equivalent of `FutureExt::now_or_never` on a
187    /// [`tokio::task::JoinHandle`]. Call this inside a synchronous context (e.g., between
188    /// agent turns) to apply a completed background result without blocking.
189    ///
190    /// The handle is consumed on success. If the task is not yet done, the handle
191    /// is returned as `Err(self)` so the caller can re-store it.
192    ///
193    /// # Examples
194    ///
195    /// ```rust,no_run
196    /// # use zeph_common::task_supervisor::{BlockingHandle, BlockingError};
197    /// async fn example(mut handle: BlockingHandle<u32>) {
198    ///     // Try to get the result without blocking.
199    ///     match handle.try_join() {
200    ///         Ok(result) => println!("done: {result:?}"),
201    ///         Err(handle) => {
202    ///             // Task still running — `handle` is returned for re-storage.
203    ///             drop(handle);
204    ///         }
205    ///     }
206    /// }
207    /// ```
208    ///
209    /// # Errors
210    ///
211    /// Returns `Err(self)` when the task has not yet produced a result (still running).
212    /// The inner `Ok(Err(BlockingError::...))` variants are returned when the task
213    /// panicked or the supervisor was dropped before the task completed.
214    pub fn try_join(mut self) -> Result<Result<R, BlockingError>, Self> {
215        match self.rx.try_recv() {
216            Ok(result) => Ok(result),
217            Err(tokio::sync::oneshot::error::TryRecvError::Empty) => Err(self),
218            Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
219                Ok(Err(BlockingError::SupervisorDropped))
220            }
221        }
222    }
223
224    /// Abort the underlying task immediately.
225    pub fn abort(&self) {
226        self.abort.abort();
227    }
228}
229
230/// Point-in-time state of a supervised task.
231#[derive(Debug, Clone, PartialEq, Eq)]
232#[non_exhaustive]
233pub enum TaskStatus {
234    /// Task is actively running.
235    Running,
236    /// Task is waiting for the restart delay before the next attempt.
237    Restarting { attempt: u32, max: u32 },
238    /// Task completed normally.
239    Completed,
240    /// Task was force-aborted during shutdown.
241    Aborted,
242    /// Task exhausted all restart attempts and is permanently failed.
243    Failed { reason: String },
244}
245
246/// Point-in-time snapshot of a supervised task, returned by [`TaskSupervisor::snapshot`].
247#[derive(Debug, Clone)]
248/// Observability surface per field:
249///
250/// | Field | tokio-console | Jaeger / OTLP | TUI | `metrics` histogram |
251/// |-------|--------------|--------------|-----|---------------------|
252/// | `name` | span name | span name | task list | label `"task"` |
253/// | `task.wall_time_ms` | — | span field | — | `zeph.task.wall_time_ms` |
254/// | `task.cpu_time_ms` | — | span field | — | `zeph.task.cpu_time_ms` |
255/// | `status` | — | — | task list | — |
256/// | `restart_count` | — | — | task list | — |
257pub struct TaskSnapshot {
258    /// Task name.
259    pub name: Arc<str>,
260    /// Current status.
261    pub status: TaskStatus,
262    /// Instant the task was first spawned.
263    pub started_at: Instant,
264    /// Number of times the task has been restarted.
265    pub restart_count: u32,
266}
267
268// ── Internal types ───────────────────────────────────────────────────────────
269
270type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
271type BoxFactory = Box<dyn Fn() -> BoxFuture + Send + Sync>;
272
273struct TaskEntry {
274    name: Arc<str>,
275    status: TaskStatus,
276    started_at: Instant,
277    restart_count: u32,
278    restart_policy: RestartPolicy,
279    abort_handle: AbortHandle,
280    /// `Some` only for `Restart` policy tasks.
281    factory: Option<BoxFactory>,
282}
283
284/// How a supervised task ended.
285#[derive(Debug, Clone, Copy, PartialEq, Eq)]
286enum CompletionKind {
287    /// Future returned normally.
288    Normal,
289    /// Future returned a value that the caller-supplied classifier (see
290    /// [`TaskSupervisor::spawn_oneshot_classified`]) identified as an inner failure —
291    /// distinct from `Normal` even though the outer `JoinHandle` itself resolved
292    /// successfully (no panic, no cancellation).
293    Failed,
294    /// Future panicked.
295    Panicked,
296    /// Future was cancelled via the cancellation token or abort handle.
297    Cancelled,
298}
299
300struct Completion {
301    name: Arc<str>,
302    kind: CompletionKind,
303}
304
305struct SupervisorState {
306    tasks: HashMap<Arc<str>, TaskEntry>,
307}
308
309struct Inner {
310    state: parking_lot::Mutex<SupervisorState>,
311    /// Completion events from spawned tasks → reap driver.
312    /// Lives in `Inner` (not `SupervisorState`) to avoid double mutex acquisition
313    /// — callers clone it once during spawn without re-locking state.
314    completion_tx: mpsc::UnboundedSender<Completion>,
315    cancel: CancellationToken,
316    /// Limits the number of concurrently running `spawn_blocking` tasks to prevent
317    /// runaway thread-pool growth under burst load.
318    blocking_semaphore: Arc<tokio::sync::Semaphore>,
319    /// Notified when `active_count()` reaches zero so `shutdown_all` wakes immediately.
320    shutdown_notify: Arc<tokio::sync::Notify>,
321}
322
323// ── Main type ────────────────────────────────────────────────────────────────
324
325/// Shared, cloneable handle to the supervised lifecycle task registry.
326///
327/// `TaskSupervisor` manages named, long-lived background tasks with restart
328/// policies, health snapshots, and graceful shutdown. It is `Clone + Send + Sync`
329/// so it can be distributed to bootstrap code, TUI, and shutdown orchestration
330/// without any additional synchronisation.
331///
332/// # Thread safety
333///
334/// Interior state is guarded by a `parking_lot::Mutex`. The lock is **never**
335/// held across `.await` points.
336///
337/// # Examples
338///
339/// ```rust,no_run
340/// use std::time::Duration;
341/// use tokio_util::sync::CancellationToken;
342/// use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskSupervisor};
343///
344/// # #[tokio::main]
345/// # async fn main() {
346/// let cancel = CancellationToken::new();
347/// let sup = TaskSupervisor::new(cancel.clone());
348///
349/// let _handle = sup.spawn(TaskDescriptor {
350///     name: "watcher",
351///     restart: RestartPolicy::RunOnce,
352///     factory: || async { tokio::time::sleep(std::time::Duration::from_secs(1)).await },
353/// });
354///
355/// sup.shutdown_all(Duration::from_secs(5)).await;
356/// # }
357/// ```
358#[derive(Clone)]
359pub struct TaskSupervisor {
360    inner: Arc<Inner>,
361}
362
363impl std::fmt::Debug for TaskSupervisor {
364    /// Prints a stable placeholder rather than task internals (factory closures
365    /// inside `TaskEntry` are not `Debug`); callers needing task details should
366    /// use [`TaskSupervisor::snapshot`] instead.
367    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
368        f.debug_struct("TaskSupervisor").finish_non_exhaustive()
369    }
370}
371
372impl TaskSupervisor {
373    /// Create a new supervisor and start its reap driver.
374    ///
375    /// The `cancel` token is propagated into every spawned task via `tokio::select!`.
376    /// When the token is cancelled, all tasks exit cooperatively on their next
377    /// cancellation check. Call [`shutdown_all`][Self::shutdown_all] to wait for
378    /// them to finish.
379    ///
380    /// When called outside a Tokio runtime context (e.g. in synchronous unit tests),
381    /// the reap driver is skipped. The supervisor still accepts task registrations but
382    /// completion callbacks are not processed — safe because no tasks can actually be
383    /// spawned without a runtime.
384    #[must_use]
385    pub fn new(cancel: CancellationToken) -> Self {
386        // NOTE: unbounded channel is acceptable here because supervised tasks are
387        // O(10–20) lifecycle services, not high-throughput work. Backpressure would
388        // complicate the spawn path without practical benefit.
389        let (completion_tx, completion_rx) = mpsc::unbounded_channel();
390        let inner = Arc::new(Inner {
391            state: parking_lot::Mutex::new(SupervisorState {
392                tasks: HashMap::new(),
393            }),
394            completion_tx,
395            cancel: cancel.clone(),
396            blocking_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
397            shutdown_notify: Arc::new(tokio::sync::Notify::new()),
398        });
399
400        // Only start the reap driver when a Tokio runtime is available. In synchronous
401        // unit tests that construct Agent/LifecycleState directly there is no reactor,
402        // so we skip the spawn. Without a runtime no tasks can be spawned either, so
403        // the driver is not needed.
404        if tokio::runtime::Handle::try_current().is_ok() {
405            Self::start_reap_driver(Arc::clone(&inner), completion_rx, cancel);
406        }
407
408        Self { inner }
409    }
410
411    /// Spawn a named, supervised async task.
412    ///
413    /// If a task with the same `name` already exists, it is aborted before the
414    /// new one is started.
415    ///
416    /// # Examples
417    ///
418    /// ```rust,no_run
419    /// use std::time::Duration;
420    /// use tokio_util::sync::CancellationToken;
421    /// use zeph_common::task_supervisor::{RestartPolicy, TaskDescriptor, TaskHandle, TaskSupervisor};
422    ///
423    /// # #[tokio::main]
424    /// # async fn main() {
425    /// let cancel = CancellationToken::new();
426    /// let sup = TaskSupervisor::new(cancel.clone());
427    ///
428    /// let handle: TaskHandle = sup.spawn(TaskDescriptor {
429    ///     name: "config-watcher",
430    ///     restart: RestartPolicy::Restart { max: 3, base_delay: Duration::from_secs(1) },
431    ///     factory: || async { /* watch loop */ },
432    /// });
433    /// # }
434    /// ```
435    pub fn spawn<F, Fut>(&self, desc: TaskDescriptor<F>) -> TaskHandle
436    where
437        F: Fn() -> Fut + Send + Sync + 'static,
438        Fut: Future<Output = ()> + Send + 'static,
439    {
440        let factory: BoxFactory = Box::new(move || Box::pin((desc.factory)()));
441        let cancel = self.inner.cancel.clone();
442        let completion_tx = self.inner.completion_tx.clone();
443        let name: Arc<str> = Arc::from(desc.name);
444
445        let (abort_handle, jh) = Self::do_spawn(desc.name, &factory, cancel);
446        Self::wire_completion_reporter(Arc::clone(&name), jh, completion_tx);
447
448        let entry = TaskEntry {
449            name: Arc::clone(&name),
450            status: TaskStatus::Running,
451            started_at: Instant::now(),
452            restart_count: 0,
453            restart_policy: desc.restart,
454            abort_handle: abort_handle.clone(),
455            factory: match desc.restart {
456                RestartPolicy::RunOnce => None,
457                RestartPolicy::Restart { .. } => Some(factory),
458            },
459        };
460
461        {
462            let mut state = self.inner.state.lock();
463            if let Some(old) = state.tasks.remove(&name) {
464                old.abort_handle.abort();
465            }
466            state.tasks.insert(Arc::clone(&name), entry);
467        }
468
469        TaskHandle {
470            name: desc.name,
471            abort: abort_handle,
472        }
473    }
474
475    /// Spawn a CPU-bound closure on the OS blocking thread pool.
476    ///
477    /// The closure runs via [`tokio::task::spawn_blocking`] — it is never polled
478    /// on tokio worker threads and cannot block async I/O. The task is registered
479    /// in the supervisor registry and is visible to [`snapshot`][Self::snapshot]
480    /// and [`shutdown_all`][Self::shutdown_all].
481    ///
482    /// Dropping the returned [`BlockingHandle`] without calling `.join()` does
483    /// **not** cancel the task; it runs to completion but the result is discarded.
484    ///
485    /// A panic inside `f` is captured and returned as [`BlockingError::Panicked`]
486    /// rather than propagating to the caller.
487    ///
488    /// # Examples
489    ///
490    /// ```rust,no_run
491    /// use std::sync::Arc;
492    /// use tokio_util::sync::CancellationToken;
493    /// use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
494    ///
495    /// # #[tokio::main]
496    /// # async fn main() {
497    /// let cancel = CancellationToken::new();
498    /// let sup = TaskSupervisor::new(cancel);
499    ///
500    /// let handle: BlockingHandle<u32> = sup.spawn_blocking(Arc::from("compute"), || {
501    ///     // CPU-bound work — safe to block here
502    ///     42_u32
503    /// });
504    /// let result = handle.join().await.unwrap();
505    /// assert_eq!(result, 42);
506    /// # }
507    /// ```
508    ///
509    /// # Capacity limit
510    ///
511    /// At most 8 `spawn_blocking` tasks run concurrently. Additional tasks wait for a
512    /// semaphore permit, bounding thread-pool growth under burst load.
513    ///
514    /// # Panics
515    ///
516    /// Panics inside `f` are captured and returned as [`BlockingError::Panicked`] — they
517    /// do not propagate to the caller.
518    #[allow(clippy::needless_pass_by_value)] // `name` is cloned into async task and registry
519    pub fn spawn_blocking<F, R>(&self, name: Arc<str>, f: F) -> BlockingHandle<R>
520    where
521        F: FnOnce() -> R + Send + 'static,
522        R: Send + 'static,
523    {
524        let (tx, rx) = oneshot::channel::<Result<R, BlockingError>>();
525        let span = tracing::info_span!(
526            "supervised_blocking_task",
527            task.name = %name,
528            task.wall_time_ms = tracing::field::Empty,
529            task.cpu_time_ms = tracing::field::Empty,
530        );
531
532        let semaphore = Arc::clone(&self.inner.blocking_semaphore);
533        let inner = Arc::clone(&self.inner);
534        let name_clone = Arc::clone(&name);
535        let completion_tx = self.inner.completion_tx.clone();
536
537        // Wrap the blocking spawn in an async task that first acquires a semaphore
538        // permit, bounding the number of concurrently running blocking tasks to 8.
539        let outer = tokio::spawn(async move {
540            let _permit = semaphore
541                .acquire_owned()
542                .await
543                .expect("blocking semaphore closed");
544
545            let name_for_measure = Arc::clone(&name_clone);
546            let join_handle = tokio::task::spawn_blocking(move || {
547                let _enter = span.enter();
548                measure_blocking(&name_for_measure, f)
549            });
550            let abort = join_handle.abort_handle();
551
552            // Update registry with the real abort handle now that spawn_blocking is live.
553            {
554                let mut state = inner.state.lock();
555                if let Some(entry) = state.tasks.get_mut(&name_clone) {
556                    entry.abort_handle = abort;
557                }
558            }
559
560            let kind = match join_handle.await {
561                Ok(val) => {
562                    let _ = tx.send(Ok(val));
563                    CompletionKind::Normal
564                }
565                Err(e) if e.is_panic() => {
566                    let _ = tx.send(Err(BlockingError::Panicked));
567                    CompletionKind::Panicked
568                }
569                Err(_) => {
570                    // Aborted — drop tx so rx returns SupervisorDropped.
571                    CompletionKind::Cancelled
572                }
573            };
574            // _permit released here, freeing the semaphore slot.
575            let _ = completion_tx.send(Completion {
576                name: name_clone,
577                kind,
578            });
579        });
580        let abort = outer.abort_handle();
581
582        // Register in registry so snapshot/shutdown sees the task.
583        {
584            let mut state = self.inner.state.lock();
585            if let Some(old) = state.tasks.remove(&name) {
586                old.abort_handle.abort();
587            }
588            state.tasks.insert(
589                Arc::clone(&name),
590                TaskEntry {
591                    name: Arc::clone(&name),
592                    status: TaskStatus::Running,
593                    started_at: Instant::now(),
594                    restart_count: 0,
595                    restart_policy: RestartPolicy::RunOnce,
596                    abort_handle: abort.clone(),
597                    factory: None,
598                },
599            );
600        }
601
602        BlockingHandle { rx, abort }
603    }
604
605    /// Spawn an async task that produces a typed result value (runs on tokio worker thread).
606    ///
607    /// Unlike [`spawn`][Self::spawn], no restart policy is supported — the task
608    /// runs once. The task is registered in the supervisor registry under the
609    /// provided `name` and is visible to [`snapshot`][Self::snapshot] and
610    /// [`shutdown_all`][Self::shutdown_all].
611    ///
612    /// For CPU-bound work that must not block tokio workers, use
613    /// [`spawn_blocking`][Self::spawn_blocking] instead.
614    ///
615    /// # Examples
616    ///
617    /// ```rust,no_run
618    /// use std::sync::Arc;
619    /// use tokio_util::sync::CancellationToken;
620    /// use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
621    ///
622    /// # #[tokio::main]
623    /// # async fn main() {
624    /// let cancel = CancellationToken::new();
625    /// let sup = TaskSupervisor::new(cancel.clone());
626    ///
627    /// let handle: BlockingHandle<u32> = sup.spawn_oneshot(Arc::from("compute"), || async { 42_u32 });
628    /// let result = handle.join().await.unwrap();
629    /// assert_eq!(result, 42);
630    /// # }
631    /// ```
632    pub fn spawn_oneshot<F, Fut, R>(&self, name: Arc<str>, factory: F) -> BlockingHandle<R>
633    where
634        F: FnOnce() -> Fut + Send + 'static,
635        Fut: Future<Output = R> + Send + 'static,
636        R: Send + 'static,
637    {
638        self.spawn_oneshot_classified(name, factory, |_| true)
639    }
640
641    /// Spawn an async task like [`spawn_oneshot`][Self::spawn_oneshot], but additionally
642    /// classify the produced value's own success/failure via `is_success` instead of only
643    /// the outer `JoinHandle` outcome (normal exit vs. panic vs. cancellation).
644    ///
645    /// This matters whenever `Fut::Output` itself encodes failure (typically
646    /// `Result<T, E>`): without a classifier, a task that runs to completion but produces
647    /// `Err(..)` is classified and logged identically to one that produced `Ok(..)` —
648    /// `spawn_oneshot` alone only observes whether the *outer* future resolved, never what
649    /// value it resolved to. `is_success` is called once, by reference, before the value is
650    /// sent to the returned [`BlockingHandle`]'s receiver — it never consumes or blocks on
651    /// the value, and the value itself is delivered to the caller unchanged either way.
652    ///
653    /// # Examples
654    ///
655    /// ```rust,no_run
656    /// use std::sync::Arc;
657    /// use tokio_util::sync::CancellationToken;
658    /// use zeph_common::task_supervisor::TaskSupervisor;
659    ///
660    /// # #[tokio::main]
661    /// # async fn main() {
662    /// let sup = TaskSupervisor::new(CancellationToken::new());
663    /// let handle = sup.spawn_oneshot_classified(
664    ///     Arc::from("fallible-task"),
665    ///     || async { Result::<u32, String>::Err("boom".to_string()) },
666    ///     Result::is_ok,
667    /// );
668    /// let result = handle.join().await.unwrap();
669    /// assert!(result.is_err());
670    /// # }
671    /// ```
672    pub fn spawn_oneshot_classified<F, Fut, R>(
673        &self,
674        name: Arc<str>,
675        factory: F,
676        is_success: impl Fn(&R) -> bool + Send + 'static,
677    ) -> BlockingHandle<R>
678    where
679        F: FnOnce() -> Fut + Send + 'static,
680        Fut: Future<Output = R> + Send + 'static,
681        R: Send + 'static,
682    {
683        let (tx, rx) = oneshot::channel::<Result<R, BlockingError>>();
684        let cancel = self.inner.cancel.clone();
685        let span = tracing::info_span!("supervised_task", task.name = %name);
686        let join_handle: tokio::task::JoinHandle<Option<R>> = tokio::spawn(
687            async move {
688                let fut = factory();
689                tokio::select! {
690                    result = fut => Some(result),
691                    () = cancel.cancelled() => None,
692                }
693            }
694            .instrument(span),
695        );
696        let abort = join_handle.abort_handle();
697
698        {
699            let mut state = self.inner.state.lock();
700            if let Some(old) = state.tasks.remove(&name) {
701                old.abort_handle.abort();
702            }
703            state.tasks.insert(
704                Arc::clone(&name),
705                TaskEntry {
706                    name: Arc::clone(&name),
707                    status: TaskStatus::Running,
708                    started_at: Instant::now(),
709                    restart_count: 0,
710                    restart_policy: RestartPolicy::RunOnce,
711                    abort_handle: abort.clone(),
712                    factory: None,
713                },
714            );
715        }
716
717        let completion_tx = self.inner.completion_tx.clone();
718        tokio::spawn(async move {
719            let kind = match join_handle.await {
720                Ok(Some(val)) => {
721                    let kind = if is_success(&val) {
722                        CompletionKind::Normal
723                    } else {
724                        CompletionKind::Failed
725                    };
726                    let _ = tx.send(Ok(val));
727                    kind
728                }
729                Err(e) if e.is_panic() => {
730                    let _ = tx.send(Err(BlockingError::Panicked));
731                    CompletionKind::Panicked
732                }
733                _ => CompletionKind::Cancelled,
734            };
735            let _ = completion_tx.send(Completion { name, kind });
736        });
737        BlockingHandle { rx, abort }
738    }
739
740    /// Abort a task by name. No-op if no task with that name is registered.
741    pub fn abort(&self, name: &'static str) {
742        let state = self.inner.state.lock();
743        let key: Arc<str> = Arc::from(name);
744        if let Some(entry) = state.tasks.get(&key) {
745            entry.abort_handle.abort();
746            tracing::debug!(task.name = name, "task aborted via supervisor");
747        }
748    }
749
750    /// Gracefully shut down all supervised tasks.
751    ///
752    /// Cancels the supervisor's [`CancellationToken`] and waits up to `timeout`
753    /// for all tasks to exit. Tasks that do not exit within the timeout are
754    /// aborted forcefully and their registry entries updated to [`TaskStatus::Aborted`].
755    ///
756    /// # Note
757    ///
758    /// This cancels the token passed to [`TaskSupervisor::new`]. If you share
759    /// that token with other subsystems, they will be cancelled too. Use a child
760    /// token (`cancel.child_token()`) when the supervisor should not affect
761    /// unrelated components.
762    pub async fn shutdown_all(&self, timeout: Duration) {
763        self.inner.cancel.cancel();
764        let sleep = tokio::time::sleep(timeout);
765        tokio::pin!(sleep);
766        loop {
767            let active = self.active_count();
768            if active == 0 {
769                break;
770            }
771            // Subscribe before re-checking so we cannot miss a notification that
772            // fires between the active_count() call above and the select below.
773            let notified = self.inner.shutdown_notify.notified();
774            tokio::select! {
775                biased;
776                () = notified => {
777                    // reap driver decremented active count — re-check at top of loop
778                }
779                () = &mut sleep => {
780                    let mut remaining_names: Vec<Arc<str>> = Vec::new();
781                    {
782                        let mut state = self.inner.state.lock();
783                        for entry in state.tasks.values_mut() {
784                            if matches!(
785                                entry.status,
786                                TaskStatus::Running | TaskStatus::Restarting { .. }
787                            ) {
788                                remaining_names.push(Arc::clone(&entry.name));
789                                entry.abort_handle.abort();
790                                entry.status = TaskStatus::Aborted;
791                            }
792                        }
793                    }
794                    tracing::warn!(
795                        remaining = active,
796                        tasks = ?remaining_names,
797                        "shutdown timeout — aborting remaining tasks"
798                    );
799                    break;
800                }
801            }
802        }
803    }
804
805    /// Return a point-in-time snapshot of all registered tasks.
806    ///
807    /// Suitable for TUI status panels and structured logging. The returned
808    /// list is sorted by `started_at` ascending.
809    #[must_use]
810    pub fn snapshot(&self) -> Vec<TaskSnapshot> {
811        let state = self.inner.state.lock();
812        let mut snaps: Vec<TaskSnapshot> = state
813            .tasks
814            .values()
815            .map(|e| TaskSnapshot {
816                name: Arc::clone(&e.name),
817                status: e.status.clone(),
818                started_at: e.started_at,
819                restart_count: e.restart_count,
820            })
821            .collect();
822        snaps.sort_by_key(|s| s.started_at);
823        snaps
824    }
825
826    /// Return the number of tasks currently in `Running` or `Restarting` state.
827    #[must_use]
828    pub fn active_count(&self) -> usize {
829        let state = self.inner.state.lock();
830        state
831            .tasks
832            .values()
833            .filter(|e| {
834                matches!(
835                    e.status,
836                    TaskStatus::Running | TaskStatus::Restarting { .. }
837                )
838            })
839            .count()
840    }
841
842    /// Return a clone of the supervisor's [`CancellationToken`].
843    ///
844    /// Callers can use this to check whether shutdown has been initiated.
845    #[must_use]
846    pub fn cancellation_token(&self) -> CancellationToken {
847        self.inner.cancel.clone()
848    }
849
850    // ── Internal helpers ──────────────────────────────────────────────────────
851
852    /// Spawn the actual tokio task. Returns `(AbortHandle, JoinHandle)`.
853    fn do_spawn(
854        name: &'static str,
855        factory: &BoxFactory,
856        cancel: CancellationToken,
857    ) -> (AbortHandle, tokio::task::JoinHandle<()>) {
858        let fut = factory();
859        let span = tracing::info_span!("supervised_task", task.name = name);
860        let jh = tokio::spawn(
861            async move {
862                tokio::select! {
863                    () = fut => {},
864                    () = cancel.cancelled() => {},
865                }
866            }
867            .instrument(span),
868        );
869        let abort = jh.abort_handle();
870        (abort, jh)
871    }
872
873    /// Wire a completion reporter: drives `jh` and sends the result to `completion_tx`.
874    fn wire_completion_reporter(
875        name: Arc<str>,
876        jh: tokio::task::JoinHandle<()>,
877        completion_tx: mpsc::UnboundedSender<Completion>,
878    ) {
879        tokio::spawn(async move {
880            let kind = match jh.await {
881                Ok(()) => CompletionKind::Normal,
882                Err(e) if e.is_panic() => CompletionKind::Panicked,
883                Err(_) => CompletionKind::Cancelled,
884            };
885            let _ = completion_tx.send(Completion { name, kind });
886        });
887    }
888
889    /// Spawn the reap driver. The driver processes completion events from the mpsc channel.
890    ///
891    /// After the cancellation token fires, the driver continues draining the channel
892    /// until the registry reports no active tasks (or the channel closes) — this
893    /// ensures that tasks which complete after cancel, however long that takes,
894    /// still have their registry entries updated, allowing `shutdown_all` to observe
895    /// `active_count() == 0` correctly and wake early. The driver enforces no
896    /// deadline of its own: [`TaskSupervisor::shutdown_all`]'s own `sleep(timeout)`
897    /// is the sole authority for giving up on a stuck task, regardless of how long
898    /// the gap is between the token being cancelled (typically out-of-band, by a
899    /// shutdown bridge or signal handler — see `src/runner.rs`, `src/serve/mod.rs`)
900    /// and `shutdown_all` actually being invoked. See the Phase 2 comment below for
901    /// why this is correct and terminates.
902    fn start_reap_driver(
903        inner: Arc<Inner>,
904        mut completion_rx: mpsc::UnboundedReceiver<Completion>,
905        cancel: CancellationToken,
906    ) {
907        tokio::spawn(async move {
908            // Phase 1: normal operation — process completions until cancel fires.
909            loop {
910                tokio::select! {
911                    biased;
912                    Some(completion) = completion_rx.recv() => {
913                        Self::handle_completion(&inner, completion).await;
914                    }
915                    () = cancel.cancelled() => break,
916                }
917            }
918
919            // Phase 2: post-cancel drain — keep receiving completions until the
920            // registry reports no active tasks, or the channel closes. This prevents
921            // losing completions that arrive after tasks observe cancellation
922            // (#3161), for however long that takes.
923            //
924            // Deliberately no independent deadline here (#5926): an earlier version
925            // of this drain phase enforced its own short fallback timeout, which
926            // could — and at every real call site, reliably did — expire before
927            // `shutdown_all` was ever invoked (the shared `CancellationToken` is
928            // cancelled out-of-band by a shutdown bridge/signal handler well before
929            // `shutdown_all` runs; see `src/runner.rs`, `src/serve/mod.rs`), causing
930            // the driver to exit and silently drop any later completion. There is
931            // exactly one deadline authority for the whole shutdown sequence:
932            // `shutdown_all`'s own `sleep(timeout)` (below), which force-aborts any
933            // task still `Running`/`Restarting` under the registry lock directly —
934            // that abort itself generates a `Cancelled` completion for `spawn`/
935            // `spawn_oneshot` tasks via `wire_completion_reporter`, waking this loop
936            // promptly. The one exception is a `spawn_blocking` task force-aborted
937            // while its outer wrapper is still parked on the concurrency semaphore
938            // (see `spawn_blocking`, below) — no completion is sent in that case, but
939            // `shutdown_all` has already marked the entry `Aborted` directly under the
940            // state lock and returned correctly; only this now-idle reap-driver task
941            // may linger parked until process/test teardown, which is not a hang. This
942            // task holds its own `Arc<Inner>` (which also owns a `completion_tx`
943            // clone) for as long as it runs, so `completion_rx.recv()` can only
944            // resolve via a real completion or the channel genuinely closing — never a
945            // stale timeout.
946            let active = Self::has_active_tasks(&inner);
947            tracing::debug!(active, "reap driver entered post-cancel drain phase");
948            loop {
949                if !Self::has_active_tasks(&inner) {
950                    break;
951                }
952                match completion_rx.recv().await {
953                    Some(completion) => Self::handle_completion(&inner, completion).await,
954                    None => break, // channel closed — unreachable in practice
955                }
956            }
957            tracing::debug!(
958                active = Self::has_active_tasks(&inner),
959                "reap driver drain phase complete"
960            );
961        });
962    }
963
964    /// Returns `true` if any task is in `Running` or `Restarting` state.
965    fn has_active_tasks(inner: &Arc<Inner>) -> bool {
966        let state = inner.state.lock();
967        state.tasks.values().any(|e| {
968            matches!(
969                e.status,
970                TaskStatus::Running | TaskStatus::Restarting { .. }
971            )
972        })
973    }
974
975    /// Process a single task completion event.
976    ///
977    /// Lock is never held across `.await`. Phase 1 classifies the completion
978    /// under lock; Phase 2 sleeps with exponential backoff without a lock;
979    /// Phase 3 spawns the next instance and updates the registry.
980    async fn handle_completion(inner: &Arc<Inner>, completion: Completion) {
981        // Short-circuit: once cancellation has fired, never schedule restarts.
982        // Without this, Restart-policy tasks re-register as Running, causing
983        // has_active_tasks() to stay true and the drain loop to spin until timeout.
984        if inner.cancel.is_cancelled() {
985            {
986                let mut state = inner.state.lock();
987                state.tasks.remove(&completion.name);
988            }
989            inner.shutdown_notify.notify_waiters();
990            return;
991        }
992
993        let Some((attempt, max, delay)) = Self::classify_completion(inner, &completion) else {
994            // Task removed from registry (RunOnce completed or Restart exhausted) —
995            // wake shutdown_all so it can re-check active_count immediately.
996            inner.shutdown_notify.notify_waiters();
997            return;
998        };
999
1000        tracing::warn!(
1001            task.name = %completion.name,
1002            attempt,
1003            max,
1004            delay_ms = delay.as_millis(),
1005            "restarting supervised task"
1006        );
1007
1008        if !delay.is_zero() {
1009            tokio::time::sleep(delay).await;
1010        }
1011
1012        Self::do_restart(inner, &completion.name, attempt);
1013    }
1014
1015    /// Phase 1: classify the completion under lock and return restart parameters if needed.
1016    ///
1017    /// Returns `Some((attempt, max, backoff_delay))` when a restart should be scheduled.
1018    fn classify_completion(
1019        inner: &Arc<Inner>,
1020        completion: &Completion,
1021    ) -> Option<(u32, u32, Duration)> {
1022        let mut state = inner.state.lock();
1023        let entry = state.tasks.get_mut(&completion.name)?;
1024
1025        match completion.kind {
1026            CompletionKind::Panicked => {
1027                tracing::warn!(task.name = %completion.name, "supervised task panicked");
1028            }
1029            CompletionKind::Failed => {
1030                tracing::warn!(
1031                    task.name = %completion.name,
1032                    "supervised task completed with an inner failure"
1033                );
1034            }
1035            CompletionKind::Normal => {
1036                tracing::info!(task.name = %completion.name, "supervised task completed");
1037            }
1038            CompletionKind::Cancelled => {
1039                tracing::debug!(task.name = %completion.name, "supervised task cancelled");
1040            }
1041        }
1042
1043        match entry.restart_policy {
1044            RestartPolicy::RunOnce => {
1045                entry.status = if completion.kind == CompletionKind::Failed {
1046                    TaskStatus::Failed {
1047                        reason: "task completed with an inner failure".to_string(),
1048                    }
1049                } else {
1050                    TaskStatus::Completed
1051                };
1052                state.tasks.remove(&completion.name);
1053                None
1054            }
1055            RestartPolicy::Restart { max, base_delay } => {
1056                // Only restart on panic — normal exit and cancellation are not errors.
1057                if completion.kind != CompletionKind::Panicked {
1058                    entry.status = TaskStatus::Completed;
1059                    state.tasks.remove(&completion.name);
1060                    return None;
1061                }
1062                if entry.restart_count >= max {
1063                    let reason = format!("panicked after {max} restart(s)");
1064                    tracing::error!(
1065                        task.name = %completion.name,
1066                        attempts = max,
1067                        "task failed permanently"
1068                    );
1069                    entry.status = TaskStatus::Failed { reason };
1070                    None
1071                } else {
1072                    let attempt = entry.restart_count + 1;
1073                    entry.status = TaskStatus::Restarting { attempt, max };
1074                    // Exponential backoff: base_delay * 2^(attempt-1), capped at MAX_RESTART_DELAY.
1075                    let multiplier = 1_u32
1076                        .checked_shl(attempt.saturating_sub(1))
1077                        .unwrap_or(u32::MAX);
1078                    let delay = base_delay.saturating_mul(multiplier).min(MAX_RESTART_DELAY);
1079                    Some((attempt, max, delay))
1080                }
1081            }
1082        }
1083        // lock released here
1084    }
1085
1086    /// Phase 3: TOCTOU check, collect spawn params under lock, then spawn outside.
1087    fn do_restart(inner: &Arc<Inner>, name: &Arc<str>, attempt: u32) {
1088        let spawn_params = {
1089            let mut state = inner.state.lock();
1090            let Some(entry) = state.tasks.get_mut(name.as_ref()) else {
1091                tracing::debug!(
1092                    task.name = %name,
1093                    "task removed during restart delay — skipping"
1094                );
1095                return;
1096            };
1097            if !matches!(entry.status, TaskStatus::Restarting { .. }) {
1098                return;
1099            }
1100            let Some(factory) = &entry.factory else {
1101                return;
1102            };
1103            // Wrap factory() in catch_unwind to prevent a factory panic from crashing
1104            // the reap driver and orphaning the registry.
1105            match std::panic::catch_unwind(std::panic::AssertUnwindSafe(factory)) {
1106                Err(_) => {
1107                    let reason = format!("factory panicked on restart attempt {attempt}");
1108                    tracing::error!(task.name = %name, attempt, "factory panicked during restart");
1109                    entry.status = TaskStatus::Failed { reason };
1110                    None
1111                }
1112                Ok(fut) => Some((
1113                    fut,
1114                    inner.cancel.clone(),
1115                    inner.completion_tx.clone(),
1116                    name.clone(),
1117                )),
1118            }
1119            // lock released here
1120        };
1121
1122        let Some((fut, cancel, completion_tx, name)) = spawn_params else {
1123            return;
1124        };
1125
1126        let span = tracing::info_span!("supervised_task", task.name = %name);
1127        let jh = tokio::spawn(
1128            async move {
1129                tokio::select! {
1130                    () = fut => {},
1131                    () = cancel.cancelled() => {},
1132                }
1133            }
1134            .instrument(span),
1135        );
1136        let new_abort = jh.abort_handle();
1137
1138        {
1139            let mut state = inner.state.lock();
1140            if let Some(entry) = state.tasks.get_mut(name.as_ref()) {
1141                entry.restart_count = attempt;
1142                entry.status = TaskStatus::Running;
1143                entry.abort_handle = new_abort;
1144            }
1145        }
1146
1147        Self::wire_completion_reporter(name, jh, completion_tx);
1148    }
1149}
1150
1151// ── Task metrics helpers ──────────────────────────────────────────────────────
1152
1153/// Run `f` and record wall-time and CPU-time metrics via `metrics` crate.
1154#[inline]
1155fn measure_blocking<F, R>(name: &str, f: F) -> R
1156where
1157    F: FnOnce() -> R,
1158{
1159    use cpu_time::ThreadTime;
1160    let wall_start = std::time::Instant::now();
1161    let cpu_start = ThreadTime::now();
1162    let result = f();
1163    let wall_ms = wall_start.elapsed().as_secs_f64() * 1000.0;
1164    let cpu_ms = cpu_start.elapsed().as_secs_f64() * 1000.0;
1165    metrics::histogram!("zeph.task.wall_time_ms", "task" => name.to_owned()).record(wall_ms);
1166    metrics::histogram!("zeph.task.cpu_time_ms", "task" => name.to_owned()).record(cpu_ms);
1167    tracing::Span::current().record("task.wall_time_ms", wall_ms);
1168    tracing::Span::current().record("task.cpu_time_ms", cpu_ms);
1169    result
1170}
1171
1172// ── BlockingSpawner impl ──────────────────────────────────────────────────────
1173
1174impl BlockingSpawner for TaskSupervisor {
1175    /// Spawn a named blocking closure through the supervisor.
1176    ///
1177    /// The task is registered in the supervisor registry (visible in
1178    /// [`snapshot`][Self::snapshot] and subject to graceful shutdown) before
1179    /// the closure begins executing.
1180    fn spawn_blocking_named(
1181        &self,
1182        name: Arc<str>,
1183        f: Box<dyn FnOnce() + Send + 'static>,
1184    ) -> tokio::task::JoinHandle<()> {
1185        let handle = self.spawn_blocking(Arc::clone(&name), f);
1186        tokio::spawn(async move {
1187            if let Err(e) = handle.join().await {
1188                tracing::error!(task.name = %name, error = %e, "supervised blocking task failed");
1189            }
1190        })
1191    }
1192}
1193
1194// ── Unit tests ────────────────────────────────────────────────────────────────
1195
1196#[cfg(test)]
1197mod tests {
1198    use std::sync::Arc;
1199    use std::sync::atomic::{AtomicU32, Ordering};
1200    use std::time::Duration;
1201
1202    use tokio_util::sync::CancellationToken;
1203
1204    use super::*;
1205
1206    fn make_supervisor() -> (TaskSupervisor, CancellationToken) {
1207        let cancel = CancellationToken::new();
1208        let sup = TaskSupervisor::new(cancel.clone());
1209        (sup, cancel)
1210    }
1211
1212    #[tokio::test]
1213    async fn test_spawn_and_complete() {
1214        let (sup, _cancel) = make_supervisor();
1215
1216        let done = Arc::new(tokio::sync::Notify::new());
1217        let done2 = Arc::clone(&done);
1218
1219        sup.spawn(TaskDescriptor {
1220            name: "simple",
1221            restart: RestartPolicy::RunOnce,
1222            factory: move || {
1223                let d = Arc::clone(&done2);
1224                async move {
1225                    d.notify_one();
1226                }
1227            },
1228        });
1229
1230        tokio::time::timeout(Duration::from_secs(2), done.notified())
1231            .await
1232            .expect("task should complete");
1233
1234        tokio::time::sleep(Duration::from_millis(50)).await;
1235        assert_eq!(
1236            sup.active_count(),
1237            0,
1238            "RunOnce task should be removed after completion"
1239        );
1240    }
1241
1242    #[tokio::test]
1243    async fn test_panic_capture() {
1244        let (sup, _cancel) = make_supervisor();
1245
1246        sup.spawn(TaskDescriptor {
1247            name: "panicking",
1248            restart: RestartPolicy::RunOnce,
1249            factory: || async { panic!("intentional test panic") },
1250        });
1251
1252        tokio::time::sleep(Duration::from_millis(200)).await;
1253
1254        let snaps = sup.snapshot();
1255        assert!(
1256            snaps.iter().all(|s| s.name.as_ref() != "panicking"),
1257            "entry should be reaped"
1258        );
1259        assert_eq!(
1260            sup.active_count(),
1261            0,
1262            "active count must be 0 after RunOnce panic"
1263        );
1264    }
1265
1266    /// Regression test for S2: Restart-policy tasks must only restart on panic,
1267    /// not on normal completion.
1268    #[tokio::test]
1269    async fn test_restart_only_on_panic() {
1270        let (sup, _cancel) = make_supervisor();
1271
1272        // Part 1: normal completion — must NOT restart.
1273        let normal_counter = Arc::new(AtomicU32::new(0));
1274        let nc = Arc::clone(&normal_counter);
1275        sup.spawn(TaskDescriptor {
1276            name: "normal-exit",
1277            restart: RestartPolicy::Restart {
1278                max: 3,
1279                base_delay: Duration::from_millis(10),
1280            },
1281            factory: move || {
1282                let c = Arc::clone(&nc);
1283                async move {
1284                    c.fetch_add(1, Ordering::SeqCst);
1285                    // Returns normally — no panic.
1286                }
1287            },
1288        });
1289
1290        tokio::time::sleep(Duration::from_millis(300)).await;
1291        assert_eq!(
1292            normal_counter.load(Ordering::SeqCst),
1293            1,
1294            "normal exit must not restart"
1295        );
1296        assert!(
1297            sup.snapshot()
1298                .iter()
1299                .all(|s| s.name.as_ref() != "normal-exit"),
1300            "entry removed after normal exit"
1301        );
1302
1303        // Part 2: panic — MUST restart up to max times.
1304        let panic_counter = Arc::new(AtomicU32::new(0));
1305        let pc = Arc::clone(&panic_counter);
1306        sup.spawn(TaskDescriptor {
1307            name: "panic-exit",
1308            restart: RestartPolicy::Restart {
1309                max: 2,
1310                base_delay: Duration::from_millis(10),
1311            },
1312            factory: move || {
1313                let c = Arc::clone(&pc);
1314                async move {
1315                    c.fetch_add(1, Ordering::SeqCst);
1316                    panic!("test panic");
1317                }
1318            },
1319        });
1320
1321        // initial + 2 restarts = 3 total
1322        tokio::time::sleep(Duration::from_millis(500)).await;
1323        assert!(
1324            panic_counter.load(Ordering::SeqCst) >= 3,
1325            "panicking task must restart max times"
1326        );
1327        let snap = sup
1328            .snapshot()
1329            .into_iter()
1330            .find(|s| s.name.as_ref() == "panic-exit");
1331        assert!(
1332            matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1333            "task must be Failed after exhausting restarts"
1334        );
1335    }
1336
1337    #[tokio::test]
1338    async fn test_restart_policy() {
1339        let (sup, _cancel) = make_supervisor();
1340
1341        let counter = Arc::new(AtomicU32::new(0));
1342        let counter2 = Arc::clone(&counter);
1343
1344        sup.spawn(TaskDescriptor {
1345            name: "restartable",
1346            restart: RestartPolicy::Restart {
1347                max: 2,
1348                base_delay: Duration::from_millis(10),
1349            },
1350            factory: move || {
1351                let c = Arc::clone(&counter2);
1352                async move {
1353                    c.fetch_add(1, Ordering::SeqCst);
1354                    panic!("always panic");
1355                }
1356            },
1357        });
1358
1359        tokio::time::sleep(Duration::from_millis(500)).await;
1360
1361        let runs = counter.load(Ordering::SeqCst);
1362        assert!(
1363            runs >= 3,
1364            "expected at least 3 invocations (initial + 2 restarts), got {runs}"
1365        );
1366
1367        let snaps = sup.snapshot();
1368        let snap = snaps.iter().find(|s| s.name.as_ref() == "restartable");
1369        assert!(snap.is_some(), "failed task should remain in registry");
1370        assert!(
1371            matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1372            "task should be Failed after exhausting retries"
1373        );
1374    }
1375
1376    /// Verify exponential backoff: delay doubles on each restart attempt.
1377    #[tokio::test]
1378    async fn test_exponential_backoff() {
1379        let (sup, _cancel) = make_supervisor();
1380
1381        let timestamps = Arc::new(parking_lot::Mutex::new(Vec::<std::time::Instant>::new()));
1382        let ts = Arc::clone(&timestamps);
1383
1384        sup.spawn(TaskDescriptor {
1385            name: "backoff-task",
1386            restart: RestartPolicy::Restart {
1387                max: 3,
1388                base_delay: Duration::from_millis(50),
1389            },
1390            factory: move || {
1391                let t = Arc::clone(&ts);
1392                async move {
1393                    t.lock().push(std::time::Instant::now());
1394                    panic!("always panic");
1395                }
1396            },
1397        });
1398
1399        // Wait long enough for all restarts: 50 + 100 + 200 ms = 350 ms + overhead
1400        tokio::time::sleep(Duration::from_millis(800)).await;
1401
1402        let ts = timestamps.lock();
1403        assert!(
1404            ts.len() >= 3,
1405            "expected at least 3 invocations, got {}",
1406            ts.len()
1407        );
1408
1409        // Verify delays are roughly doubling (within 2x tolerance for CI jitter).
1410        if ts.len() >= 3 {
1411            let d1 = ts[1].duration_since(ts[0]);
1412            let d2 = ts[2].duration_since(ts[1]);
1413            // d2 should be at least 1.5x d1 (allowing for jitter).
1414            assert!(
1415                d2 >= d1.mul_f64(1.5),
1416                "expected exponential backoff: d1={d1:?} d2={d2:?}"
1417            );
1418        }
1419    }
1420
1421    #[tokio::test]
1422    async fn test_graceful_shutdown() {
1423        let (sup, _cancel) = make_supervisor();
1424
1425        for name in ["svc-a", "svc-b", "svc-c"] {
1426            sup.spawn(TaskDescriptor {
1427                name,
1428                restart: RestartPolicy::RunOnce,
1429                factory: || async {
1430                    tokio::time::sleep(Duration::from_mins(1)).await;
1431                },
1432            });
1433        }
1434
1435        assert_eq!(sup.active_count(), 3);
1436
1437        tokio::time::timeout(
1438            Duration::from_secs(2),
1439            sup.shutdown_all(Duration::from_secs(1)),
1440        )
1441        .await
1442        .expect("shutdown should complete within timeout");
1443    }
1444
1445    /// Verify that force-aborted tasks get `TaskStatus::Aborted` in the registry (A2 fix).
1446    #[tokio::test]
1447    async fn test_force_abort_marks_aborted() {
1448        let cancel = CancellationToken::new();
1449        let sup = TaskSupervisor::new(cancel.clone());
1450
1451        sup.spawn(TaskDescriptor {
1452            name: "stubborn-for-abort",
1453            restart: RestartPolicy::RunOnce,
1454            factory: || async {
1455                // Does not cooperate with cancellation.
1456                std::future::pending::<()>().await;
1457            },
1458        });
1459
1460        // Use a very short timeout to trigger force-abort.
1461        sup.shutdown_all(Duration::from_millis(1)).await;
1462
1463        // Entry should be Aborted, not Running.
1464        let snaps = sup.snapshot();
1465        if let Some(snap) = snaps
1466            .iter()
1467            .find(|s| s.name.as_ref() == "stubborn-for-abort")
1468        {
1469            assert_eq!(
1470                snap.status,
1471                TaskStatus::Aborted,
1472                "force-aborted task must have Aborted status"
1473            );
1474        }
1475        // If entry was already reaped (cooperative cancel won), that's also acceptable.
1476    }
1477
1478    #[tokio::test]
1479    async fn test_registry_snapshot() {
1480        let (sup, _cancel) = make_supervisor();
1481
1482        for name in ["alpha", "beta"] {
1483            sup.spawn(TaskDescriptor {
1484                name,
1485                restart: RestartPolicy::RunOnce,
1486                factory: || async {
1487                    tokio::time::sleep(Duration::from_secs(10)).await;
1488                },
1489            });
1490        }
1491
1492        let snaps = sup.snapshot();
1493        assert_eq!(snaps.len(), 2);
1494        let names: Vec<&str> = snaps.iter().map(|s| s.name.as_ref()).collect();
1495        assert!(names.contains(&"alpha"));
1496        assert!(names.contains(&"beta"));
1497        assert!(snaps.iter().all(|s| s.status == TaskStatus::Running));
1498    }
1499
1500    #[tokio::test]
1501    async fn test_blocking_returns_value() {
1502        let (sup, cancel) = make_supervisor();
1503
1504        let handle: BlockingHandle<u32> = sup.spawn_blocking(Arc::from("compute"), || 42_u32);
1505        let result = handle.join().await.expect("should return value");
1506        assert_eq!(result, 42);
1507        cancel.cancel();
1508    }
1509
1510    #[tokio::test]
1511    async fn test_blocking_panic() {
1512        let (sup, _cancel) = make_supervisor();
1513
1514        let handle: BlockingHandle<u32> =
1515            sup.spawn_blocking(Arc::from("panicking-compute"), || panic!("intentional"));
1516        let err = handle
1517            .join()
1518            .await
1519            .expect_err("should return error on panic");
1520        assert_eq!(err, BlockingError::Panicked);
1521    }
1522
1523    /// Verify `spawn_blocking` tasks appear in registry (M3 fix).
1524    #[tokio::test]
1525    async fn test_blocking_registered_in_registry() {
1526        let (sup, cancel) = make_supervisor();
1527
1528        let (tx, rx) = std::sync::mpsc::channel::<()>();
1529        let _handle: BlockingHandle<()> =
1530            sup.spawn_blocking(Arc::from("blocking-task"), move || {
1531                // Block until signalled.
1532                let _ = rx.recv();
1533            });
1534
1535        tokio::time::sleep(Duration::from_millis(10)).await;
1536        assert_eq!(
1537            sup.active_count(),
1538            1,
1539            "blocking task must appear in active_count"
1540        );
1541
1542        let _ = tx.send(());
1543        tokio::time::sleep(Duration::from_millis(100)).await;
1544        assert_eq!(
1545            sup.active_count(),
1546            0,
1547            "blocking task must be removed after completion"
1548        );
1549
1550        cancel.cancel();
1551    }
1552
1553    /// Verify `spawn_oneshot` tasks appear in registry (M3 fix).
1554    #[tokio::test]
1555    async fn test_oneshot_registered_in_registry() {
1556        let (sup, cancel) = make_supervisor();
1557
1558        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
1559        let _handle: BlockingHandle<()> =
1560            sup.spawn_oneshot(Arc::from("oneshot-task"), move || async move {
1561                let _ = rx.await;
1562            });
1563
1564        tokio::time::sleep(Duration::from_millis(10)).await;
1565        assert_eq!(
1566            sup.active_count(),
1567            1,
1568            "oneshot task must appear in active_count"
1569        );
1570
1571        let _ = tx.send(());
1572        tokio::time::sleep(Duration::from_millis(50)).await;
1573        assert_eq!(
1574            sup.active_count(),
1575            0,
1576            "oneshot task must be removed after completion"
1577        );
1578
1579        cancel.cancel();
1580    }
1581
1582    /// `spawn_oneshot_classified` must deliver the factory's actual return value to the
1583    /// caller unchanged, regardless of how `is_success` classifies it — the classifier only
1584    /// affects internal registry/log bookkeeping, never the value received via `join()`.
1585    #[tokio::test]
1586    async fn test_oneshot_classified_delivers_actual_value_on_success_and_failure() {
1587        let (sup, _cancel) = make_supervisor();
1588
1589        let ok_handle = sup.spawn_oneshot_classified(
1590            Arc::from("classified-ok"),
1591            || async { Result::<u32, String>::Ok(7) },
1592            Result::is_ok,
1593        );
1594        assert_eq!(ok_handle.join().await.unwrap(), Ok(7));
1595
1596        let err_handle = sup.spawn_oneshot_classified(
1597            Arc::from("classified-err"),
1598            || async { Result::<u32, String>::Err("boom".to_string()) },
1599            Result::is_ok,
1600        );
1601        assert_eq!(
1602            err_handle.join().await.unwrap(),
1603            Err("boom".to_string()),
1604            "the real Err value must reach the caller unchanged even though it is \
1605             classified as CompletionKind::Failed internally"
1606        );
1607    }
1608
1609    /// Regression test for #6257: an `Err`-returning factory must be classified as
1610    /// `CompletionKind::Failed` (logged at `warn` as "completed with an inner failure"),
1611    /// not `CompletionKind::Normal` (logged at `info` as a plain completion). This is the
1612    /// distinction `spawn_agent_task` (zeph-subagent) relies on so a subagent task that
1613    /// returns `Err` after a genuine completion (e.g. a worktree-setup failure) is reported
1614    /// as a supervisor-level failure instead of a normal completion. The classification
1615    /// only surfaces via the reap driver's log lines — for `RunOnce` the registry entry's
1616    /// `Failed`/`Completed` status assignment and its removal happen in the same locked
1617    /// critical section (see `classify_completion`), so there is no window in which
1618    /// `snapshot()` could observe the transient status; the log line is the only externally
1619    /// observable signal of which branch was taken.
1620    #[tokio::test]
1621    #[tracing_test::traced_test]
1622    async fn test_oneshot_classified_logs_failed_for_inner_err() {
1623        let (sup, _cancel) = make_supervisor();
1624
1625        let handle = sup.spawn_oneshot_classified(
1626            Arc::from("classified-inner-err"),
1627            || async { Result::<u32, String>::Err("boom".to_string()) },
1628            Result::is_ok,
1629        );
1630        handle.join().await.unwrap().unwrap_err();
1631
1632        // Give the reap driver a moment to process the completion event.
1633        tokio::time::sleep(Duration::from_millis(50)).await;
1634
1635        assert!(
1636            logs_contain("completed with an inner failure"),
1637            "an Err value must be classified as CompletionKind::Failed, not Normal"
1638        );
1639    }
1640
1641    /// Counterpart to the above: an `Ok`-returning factory must be classified as
1642    /// `CompletionKind::Normal` (logged at `info`), never as an inner failure — confirms
1643    /// the classifier doesn't over-fire on the success path.
1644    #[tokio::test]
1645    #[tracing_test::traced_test]
1646    async fn test_oneshot_classified_logs_normal_for_inner_ok() {
1647        let (sup, _cancel) = make_supervisor();
1648
1649        let handle = sup.spawn_oneshot_classified(
1650            Arc::from("classified-inner-ok"),
1651            || async { Result::<u32, String>::Ok(1) },
1652            Result::is_ok,
1653        );
1654        handle.join().await.unwrap().unwrap();
1655
1656        tokio::time::sleep(Duration::from_millis(50)).await;
1657
1658        assert!(
1659            !logs_contain("completed with an inner failure"),
1660            "an Ok value must not be classified as CompletionKind::Failed"
1661        );
1662    }
1663
1664    /// Regression guard for #6257: `spawn_oneshot`'s ~20 existing call sites across the
1665    /// workspace must remain unaffected by the addition of `spawn_oneshot_classified` —
1666    /// `spawn_oneshot` delegates with an always-true classifier (`|_| true`), so a plain
1667    /// (non-`Result`) factory value must still always be classified `Normal`/logged as a
1668    /// plain completion, never as an inner failure.
1669    #[tokio::test]
1670    #[tracing_test::traced_test]
1671    async fn test_oneshot_backward_compat_never_logs_inner_failure() {
1672        let (sup, _cancel) = make_supervisor();
1673
1674        let handle: BlockingHandle<u32> =
1675            sup.spawn_oneshot(Arc::from("plain-oneshot"), || async { 42_u32 });
1676        assert_eq!(handle.join().await.unwrap(), 42);
1677
1678        tokio::time::sleep(Duration::from_millis(50)).await;
1679
1680        assert!(
1681            !logs_contain("completed with an inner failure"),
1682            "spawn_oneshot's default always-true classifier must never report an inner failure"
1683        );
1684    }
1685
1686    #[tokio::test]
1687    async fn test_restart_max_zero() {
1688        let (sup, _cancel) = make_supervisor();
1689
1690        let counter = Arc::new(AtomicU32::new(0));
1691        let counter2 = Arc::clone(&counter);
1692
1693        sup.spawn(TaskDescriptor {
1694            name: "zero-max",
1695            restart: RestartPolicy::Restart {
1696                max: 0,
1697                base_delay: Duration::from_millis(10),
1698            },
1699            factory: move || {
1700                let c = Arc::clone(&counter2);
1701                async move {
1702                    c.fetch_add(1, Ordering::SeqCst);
1703                    panic!("always panic");
1704                }
1705            },
1706        });
1707
1708        tokio::time::sleep(Duration::from_millis(200)).await;
1709
1710        assert_eq!(
1711            counter.load(Ordering::SeqCst),
1712            1,
1713            "max=0 should not restart"
1714        );
1715
1716        let snaps = sup.snapshot();
1717        let snap = snaps.iter().find(|s| s.name.as_ref() == "zero-max");
1718        assert!(snap.is_some(), "entry should remain as Failed");
1719        assert!(
1720            matches!(snap.unwrap().status, TaskStatus::Failed { .. }),
1721            "status should be Failed"
1722        );
1723    }
1724
1725    /// Stress test: spawn 50 tasks concurrently, all must complete and registry must be accurate.
1726    #[tokio::test]
1727    async fn test_concurrent_spawns() {
1728        // All task names must be 'static — pre-defined before any let statements.
1729        static NAMES: [&str; 50] = [
1730            "t00", "t01", "t02", "t03", "t04", "t05", "t06", "t07", "t08", "t09", "t10", "t11",
1731            "t12", "t13", "t14", "t15", "t16", "t17", "t18", "t19", "t20", "t21", "t22", "t23",
1732            "t24", "t25", "t26", "t27", "t28", "t29", "t30", "t31", "t32", "t33", "t34", "t35",
1733            "t36", "t37", "t38", "t39", "t40", "t41", "t42", "t43", "t44", "t45", "t46", "t47",
1734            "t48", "t49",
1735        ];
1736        let (sup, cancel) = make_supervisor();
1737
1738        let completed = Arc::new(AtomicU32::new(0));
1739        for name in &NAMES {
1740            let c = Arc::clone(&completed);
1741            sup.spawn(TaskDescriptor {
1742                name,
1743                restart: RestartPolicy::RunOnce,
1744                factory: move || {
1745                    let c = Arc::clone(&c);
1746                    async move {
1747                        c.fetch_add(1, Ordering::SeqCst);
1748                    }
1749                },
1750            });
1751        }
1752
1753        // Wait for all tasks to complete.
1754        tokio::time::timeout(Duration::from_secs(5), async {
1755            loop {
1756                if completed.load(Ordering::SeqCst) == 50 {
1757                    break;
1758                }
1759                tokio::time::sleep(Duration::from_millis(10)).await;
1760            }
1761        })
1762        .await
1763        .expect("all 50 tasks should complete");
1764
1765        // Give reap driver time to process all completions.
1766        tokio::time::sleep(Duration::from_millis(100)).await;
1767        assert_eq!(sup.active_count(), 0, "all tasks must be reaped");
1768
1769        cancel.cancel();
1770    }
1771
1772    #[tokio::test]
1773    async fn test_shutdown_timeout_expiry() {
1774        let cancel = CancellationToken::new();
1775        let sup = TaskSupervisor::new(cancel.clone());
1776
1777        sup.spawn(TaskDescriptor {
1778            name: "stubborn",
1779            restart: RestartPolicy::RunOnce,
1780            factory: || async {
1781                tokio::time::sleep(Duration::from_mins(1)).await;
1782            },
1783        });
1784
1785        assert_eq!(sup.active_count(), 1);
1786
1787        tokio::time::timeout(
1788            Duration::from_secs(2),
1789            sup.shutdown_all(Duration::from_millis(50)),
1790        )
1791        .await
1792        .expect("shutdown_all should return even on timeout expiry");
1793
1794        assert!(
1795            cancel.is_cancelled(),
1796            "cancel token must be cancelled after shutdown"
1797        );
1798    }
1799
1800    #[tokio::test]
1801    async fn test_cancellation_token() {
1802        let cancel = CancellationToken::new();
1803        let sup = TaskSupervisor::new(cancel.clone());
1804
1805        assert!(!sup.cancellation_token().is_cancelled());
1806
1807        sup.shutdown_all(Duration::from_millis(100)).await;
1808
1809        assert!(
1810            sup.cancellation_token().is_cancelled(),
1811            "token must be cancelled after shutdown"
1812        );
1813    }
1814
1815    /// Regression test for #3161: after `shutdown_all`, all tasks must be reaped
1816    /// even when they complete *after* the cancel signal.
1817    ///
1818    /// The yield loop forces the reap driver to observe cancel and exit phase-1
1819    /// before the tasks send their completions — reliably reproducing the race.
1820    #[tokio::test]
1821    async fn test_shutdown_drains_post_cancel_completions() {
1822        let cancel = CancellationToken::new();
1823        let sup = TaskSupervisor::new(cancel.clone());
1824
1825        for name in [
1826            "loop-1", "loop-2", "loop-3", "loop-4", "loop-5", "loop-6", "loop-7",
1827        ] {
1828            let cancel_inner = cancel.clone();
1829            sup.spawn(TaskDescriptor {
1830                name,
1831                restart: RestartPolicy::RunOnce,
1832                factory: move || {
1833                    let c = cancel_inner.clone();
1834                    async move {
1835                        c.cancelled().await;
1836                        // Yield multiple times so the reap driver observes cancel first.
1837                        for _ in 0..64 {
1838                            tokio::task::yield_now().await;
1839                        }
1840                    }
1841                },
1842            });
1843        }
1844        assert_eq!(sup.active_count(), 7);
1845
1846        sup.shutdown_all(Duration::from_secs(2)).await;
1847
1848        assert_eq!(
1849            sup.active_count(),
1850            0,
1851            "all tasks must be reaped after shutdown (#3161)"
1852        );
1853    }
1854
1855    /// Regression test for #5926: `shutdown_all` must wait for a task that is
1856    /// still running when the caller-supplied `timeout` has not yet elapsed,
1857    /// rather than giving up on any independent, shorter deadline.
1858    ///
1859    /// `spawn_blocking` tasks don't observe the cancellation token — unlike
1860    /// `spawn`-based tasks, whose futures are dropped immediately by `do_spawn`'s
1861    /// `tokio::select!` against `cancel.cancelled()` — so a `spawn_blocking` task is
1862    /// the only way to have a task genuinely keep running past cancel and exercise
1863    /// the drain phase.
1864    #[tokio::test]
1865    async fn test_shutdown_all_reaps_blocking_task_finishing_after_cancel() {
1866        let (sup, _cancel) = make_supervisor();
1867
1868        let handle = sup.spawn_blocking(Arc::from("slow-blocking"), || {
1869            std::thread::sleep(Duration::from_millis(300));
1870        });
1871
1872        // Let the blocking task actually start before triggering shutdown.
1873        tokio::time::sleep(Duration::from_millis(20)).await;
1874
1875        let start = tokio::time::Instant::now();
1876        sup.shutdown_all(Duration::from_secs(2)).await;
1877        let elapsed = start.elapsed();
1878
1879        assert!(
1880            elapsed < Duration::from_secs(1),
1881            "shutdown_all must wake as soon as the blocking task actually finishes \
1882             (~300ms), not wait out the full 2s caller timeout; elapsed={elapsed:?}"
1883        );
1884
1885        // Successfully reaped tasks are removed from the registry entirely (see
1886        // `handle_completion`'s post-cancel short-circuit); only force-aborted tasks
1887        // remain with `TaskStatus::Aborted`. Absence here means the completion was
1888        // processed before shutdown_all's own timeout forced an abort.
1889        let snapshot = sup.snapshot();
1890        assert!(
1891            snapshot.iter().all(|t| t.name.as_ref() != "slow-blocking"),
1892            "task that completed cleanly after cancel must be reaped from the \
1893             registry, not left behind and force-aborted: {snapshot:?}"
1894        );
1895
1896        // `join` returns `Ok` only when the closure ran to completion and the
1897        // outer task wasn't aborted — the `BlockingHandle` equivalent of
1898        // `TaskStatus::Completed` vs `TaskStatus::Aborted`.
1899        handle
1900            .join()
1901            .await
1902            .expect("blocking task must complete, not be reported as aborted (#5926)");
1903    }
1904
1905    /// Regression test for #5926, critical topology (out-of-band cancel):
1906    /// production call sites (`src/runner.rs`, `src/serve/mod.rs`) always cancel
1907    /// the shared `CancellationToken` *before* calling `shutdown_all` — a
1908    /// separate shutdown bridge/signal handler owns the cancel, and `shutdown_all`
1909    /// is only invoked afterward as a waiter. A fix that only works when
1910    /// `shutdown_all` is the sole canceller would pass a naive test but stay
1911    /// broken in production. This test replicates the real topology explicitly:
1912    /// cancel first, let the reap driver start draining, *then* call
1913    /// `shutdown_all`, and confirm the caller's real timeout still governs the
1914    /// wait.
1915    #[tokio::test]
1916    async fn test_shutdown_all_reaps_task_after_out_of_band_cancel() {
1917        let cancel = CancellationToken::new();
1918        let sup = TaskSupervisor::new(cancel.clone());
1919
1920        let handle = sup.spawn_blocking(Arc::from("slow-blocking-oob"), || {
1921            std::thread::sleep(Duration::from_millis(300));
1922        });
1923        tokio::time::sleep(Duration::from_millis(20)).await;
1924
1925        // Out-of-band cancel — nothing to do with shutdown_all, mirrors a shutdown
1926        // bridge/signal handler cancelling the shared token directly.
1927        cancel.cancel();
1928        // Give the reap driver time to observe cancellation and enter Phase 2
1929        // before shutdown_all ever runs.
1930        tokio::time::sleep(Duration::from_millis(50)).await;
1931
1932        let start = tokio::time::Instant::now();
1933        sup.shutdown_all(Duration::from_secs(2)).await;
1934        let elapsed = start.elapsed();
1935
1936        assert!(
1937            elapsed < Duration::from_secs(1),
1938            "shutdown_all must wake once the blocking task actually finishes even \
1939             though the token was cancelled out-of-band before shutdown_all was \
1940             called; elapsed={elapsed:?}"
1941        );
1942
1943        handle.join().await.expect(
1944            "blocking task must complete, not be reported as aborted, even under \
1945             the out-of-band-cancel topology (#5926)",
1946        );
1947    }
1948
1949    /// Regression test for #5926 residual (critic finding S1): the reap driver's
1950    /// drain phase must never give up on a still-active, non-cooperative task on
1951    /// its own — only `shutdown_all`'s own `sleep(timeout)` may do that. An
1952    /// earlier version of the fix gave the drain phase its own short fallback
1953    /// deadline; if the gap between the out-of-band cancel and `shutdown_all`
1954    /// being invoked exceeded that fallback, the reap driver exited *before*
1955    /// `shutdown_all` ever ran, silently dropping the eventual completion and
1956    /// misreporting the task `Aborted`. This is realistic in production: e.g.
1957    /// `serve/mod.rs`'s `axum::serve(...).with_graceful_shutdown(...)` drains
1958    /// in-flight connections with no timeout of its own before `shutdown_all(30s)`
1959    /// is ever reached, and a long-running agent-turn request can hold that gap
1960    /// open well past any short fallback.
1961    ///
1962    /// Uses paused virtual time (`tokio::time::advance`) to jump the cancel-to-
1963    /// `shutdown_all` gap arbitrarily far forward without a real multi-second
1964    /// test. A `spawn_blocking` task runs on a real OS thread (unaffected by
1965    /// paused tokio time) and is held open via a channel until explicitly
1966    /// released, so its completion stays under exact manual control relative to
1967    /// the virtual-time advance and to when `shutdown_all` is actually called.
1968    #[tokio::test(start_paused = true)]
1969    async fn test_shutdown_all_reaps_task_after_long_pre_invocation_gap() {
1970        let cancel = CancellationToken::new();
1971        let sup = TaskSupervisor::new(cancel.clone());
1972
1973        let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
1974        let handle: BlockingHandle<()> = sup.spawn_blocking(Arc::from("held-task"), move || {
1975            let _ = release_rx.recv();
1976        });
1977
1978        // Out-of-band cancel — before shutdown_all is ever called.
1979        cancel.cancel();
1980        for _ in 0..8 {
1981            tokio::task::yield_now().await;
1982        }
1983
1984        // Advance virtual time by a large margin — deliberately larger than any
1985        // fallback window a naive fix might reintroduce — with `shutdown_all`
1986        // still not called and the task still held open. A drain phase with any
1987        // independent deadline of its own would have already exited by now.
1988        tokio::time::advance(Duration::from_mins(2)).await;
1989
1990        // *Now* the real caller shows up.
1991        let sup2 = sup.clone();
1992        let shutdown_jh =
1993            tokio::spawn(async move { sup2.shutdown_all(Duration::from_secs(30)).await });
1994        for _ in 0..8 {
1995            tokio::task::yield_now().await;
1996        }
1997
1998        // Let the real OS thread actually finish.
1999        let _ = release_tx.send(());
2000
2001        // Await directly (no wrapping virtual-time timeout): under `start_paused`,
2002        // an artificial timeout here would itself be subject to auto-advance and
2003        // cannot reliably distinguish outcomes. The real signal is the registry
2004        // state checked below, not how long this await took.
2005        shutdown_jh.await.expect("shutdown_all task must not panic");
2006
2007        // NOTE: `handle.join()` reflects `spawn_blocking`'s own internal oneshot
2008        // completion channel, which is fed directly by the real OS thread finishing
2009        // — entirely independent of the reap driver / registry path this test is
2010        // targeting. It succeeds either way and cannot distinguish the bug, so the
2011        // real assertion is the supervisor's registry state, not `handle.join()`.
2012        let snapshot = sup.snapshot();
2013        let entry = snapshot.iter().find(|t| t.name.as_ref() == "held-task");
2014        assert!(
2015            entry.is_none(),
2016            "a task finishing after shutdown_all was (belatedly) invoked must be \
2017             reaped from the registry (the post-cancel path removes completed \
2018             entries rather than marking them), not force-aborted, no matter how \
2019             long the cancel-to-shutdown_all gap was: {snapshot:?} (#5926 S1)"
2020        );
2021
2022        handle
2023            .join()
2024            .await
2025            .expect("blocking task must complete without panicking");
2026    }
2027
2028    #[tokio::test]
2029    async fn test_blocking_spawner_task_appears_in_snapshot() {
2030        // Verify that tasks spawned via BlockingSpawner appear in supervisor.snapshot().
2031        use crate::BlockingSpawner;
2032
2033        let cancel = CancellationToken::new();
2034        let sup = TaskSupervisor::new(cancel);
2035
2036        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
2037        let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
2038
2039        let handle = sup.spawn_blocking_named(
2040            Arc::from("chunk_file"),
2041            Box::new(move || {
2042                // Signal that the task has started.
2043                let _ = ready_tx.send(());
2044                // Block until test signals release.
2045                let _ = release_rx.blocking_recv();
2046            }),
2047        );
2048
2049        // Wait until the blocking task has actually started.
2050        ready_rx.await.expect("task should start");
2051
2052        let snapshot = sup.snapshot();
2053        assert!(
2054            snapshot.iter().any(|t| t.name.as_ref() == "chunk_file"),
2055            "chunk_file task must appear in supervisor snapshot"
2056        );
2057
2058        // Release the blocking task and await completion.
2059        let _ = release_tx.send(());
2060        handle.await.expect("task should complete");
2061    }
2062
2063    /// Verify that `measure_blocking` emits wall-time and CPU-time histograms.
2064    ///
2065    /// `measure_blocking` calls `metrics::histogram!` on the current thread.
2066    /// We test it directly using a `DebuggingRecorder` installed as the thread-local
2067    /// recorder via `metrics::with_local_recorder`.
2068    #[test]
2069    fn test_measure_blocking_emits_metrics() {
2070        use metrics_util::debugging::DebuggingRecorder;
2071
2072        let recorder = DebuggingRecorder::new();
2073        let snapshotter = recorder.snapshotter();
2074
2075        // Call measure_blocking inside the local recorder scope so histogram! calls
2076        // are captured. The closure runs synchronously on this thread.
2077        metrics::with_local_recorder(&recorder, || {
2078            measure_blocking("test_task", || std::hint::black_box(42_u64));
2079        });
2080
2081        let snapshot = snapshotter.snapshot();
2082        let metric_names: Vec<String> = snapshot
2083            .into_vec()
2084            .into_iter()
2085            .map(|(k, _, _, _)| k.key().name().to_owned())
2086            .collect();
2087
2088        assert!(
2089            metric_names.iter().any(|n| n == "zeph.task.wall_time_ms"),
2090            "expected zeph.task.wall_time_ms histogram; got: {metric_names:?}"
2091        );
2092        assert!(
2093            metric_names.iter().any(|n| n == "zeph.task.cpu_time_ms"),
2094            "expected zeph.task.cpu_time_ms histogram; got: {metric_names:?}"
2095        );
2096    }
2097
2098    /// Verify that `spawn_blocking` semaphore limits concurrent OS-thread tasks to 8.
2099    ///
2100    /// Spawns 16 tasks. Each holds a barrier until 8 are waiting; then releases in order.
2101    /// If more than 8 run concurrently the test would either deadlock (waiting for 9+ to reach
2102    /// the barrier) or the counter would exceed 8 — both are caught.
2103    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2104    async fn test_spawn_blocking_semaphore_cap() {
2105        let (sup, _cancel) = make_supervisor();
2106        let concurrent = Arc::new(AtomicU32::new(0));
2107        let max_concurrent = Arc::new(AtomicU32::new(0));
2108        let barrier = Arc::new(std::sync::Barrier::new(1)); // just a sync point
2109
2110        let mut handles = Vec::new();
2111        for i in 0u32..16 {
2112            let c = Arc::clone(&concurrent);
2113            let m = Arc::clone(&max_concurrent);
2114            let name: Arc<str> = Arc::from(format!("blocking-{i}").as_str());
2115            let h = sup.spawn_blocking(name, move || {
2116                let prev = c.fetch_add(1, Ordering::SeqCst);
2117                // Update observed maximum.
2118                let mut cur_max = m.load(Ordering::SeqCst);
2119                while prev + 1 > cur_max {
2120                    match m.compare_exchange(cur_max, prev + 1, Ordering::SeqCst, Ordering::SeqCst)
2121                    {
2122                        Ok(_) => break,
2123                        Err(x) => cur_max = x,
2124                    }
2125                }
2126                // Simulate work.
2127                std::thread::sleep(std::time::Duration::from_millis(20));
2128                c.fetch_sub(1, Ordering::SeqCst);
2129            });
2130            handles.push(h);
2131        }
2132
2133        for h in handles {
2134            h.join().await.expect("blocking task should succeed");
2135        }
2136        drop(barrier);
2137
2138        let observed = max_concurrent.load(Ordering::SeqCst);
2139        assert!(
2140            observed <= 8,
2141            "observed {observed} concurrent blocking tasks; expected ≤ 8 (semaphore cap)"
2142        );
2143    }
2144}