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