agent_engine/runtime/

subagent.rs

use serde_json::Value;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tokio::sync::{mpsc, oneshot};

// ── SubagentResult ───────────────────────────────────────────────────────────────

#[derive(Debug)]
pub struct SubagentResult {
    pub text: String,
    pub model: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read: u64,
    pub cache_creation: u64,
    /// TTL split of `cache_creation` across the subagent's turns.
    /// `None` only if no turn ever reported a split; otherwise the sum.
    pub cache_creation_5m: Option<u64>,
    pub cache_creation_1h: Option<u64>,
    pub tool_count: u32,
    pub timed_out: bool,
}

#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct AuthorizationDiagnostic {
    pub selection_source: String,
    pub policy_digest: String,
    pub catalog_snapshot_id: String,
    pub catalog_digest: String,
    pub correlation_id: String,
    pub network_attempted: bool,
    pub cross_provider_grant_id: Option<String>,
}

// ── SubagentStatus ───────────────────────────────────────────────────────────────

#[derive(Debug, Clone, PartialEq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum TerminalCategory {
    Completed,
    Route,
    Credential,
    Runtime,
    ProviderClient,
    Transport,
    StartupFailed,
    ExecutionFailed,
    TimedOut,
    Cancelled,
}

#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct TerminalDiagnostic {
    pub category: TerminalCategory,
    pub code: String,
    pub stage: String,
    pub correlation_id: String,
    pub network_attempted: bool,
    pub safe_message: String,
}

pub fn safe_failure_message(category: &TerminalCategory) -> &'static str {
    match category {
        TerminalCategory::Route => "worker route could not be resolved",
        TerminalCategory::Credential => "worker credentials are unavailable",
        TerminalCategory::Runtime | TerminalCategory::StartupFailed => {
            "worker runtime could not be started"
        }
        TerminalCategory::ProviderClient => "provider client could not be initialized",
        TerminalCategory::Transport => "provider transport failed",
        TerminalCategory::TimedOut => "worker timed out",
        TerminalCategory::Cancelled => "worker was cancelled",
        TerminalCategory::ExecutionFailed => "worker execution failed",
        TerminalCategory::Completed => "worker completed",
    }
}

#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SubagentStatus {
    Running,
    Completed,
    /// User-aborted via TUI cancel. Never publishes a wake event.
    Cancelled,
    TimedOut,
    Failed(String),
}

// ── SubagentState ────────────────────────────────────────────────────────────────

/// All mutable state shared between the subagent thread and its handle.
/// Collapsed behind a single RwLock so a status poll takes exactly one lock.
#[derive(Debug)]
pub struct SubagentState {
    pub status: SubagentStatus,
    pub partial_text: String,
    pub tool_log: Vec<String>,
    pub conversation_state: Vec<Value>,
    /// Stamped once by finalize_subagent at thread exit.
    pub finished_at: Option<std::time::Instant>,
    /// Set by cancel() before the shutdown signal is sent.
    /// Read by finalize_subagent to label the terminal status correctly.
    pub cancel_requested: bool,
    pub terminal: Option<TerminalDiagnostic>,
}

impl SubagentState {
    pub fn new() -> Self {
        Self {
            status: SubagentStatus::Running,
            partial_text: String::new(),
            tool_log: Vec::new(),
            conversation_state: Vec::new(),
            finished_at: None,
            cancel_requested: false,
            terminal: None,
        }
    }

    /// Persist only typed, allowlisted failure data. Raw provider errors never enter state.
    pub fn safe_fail(&mut self, mut diagnostic: TerminalDiagnostic) {
        diagnostic.safe_message = safe_failure_message(&diagnostic.category).into();
        self.status = SubagentStatus::Failed(diagnostic.safe_message.clone());
        self.terminal = Some(diagnostic);
    }
}

impl Default for SubagentState {
    fn default() -> Self {
        Self::new()
    }
}

// ── SubagentDisplayRow ────────────────────────────────────────────────────────────

/// Snapshot row produced by SubagentRegistry::display_rows().
/// Used by the TUI reconcile path as the registry's liveness authority.
/// Serde: travels in `SessionEventWire::SubagentRows` (phase 3).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SubagentDisplayRow {
    pub subagent_id: u64,
    pub agent_name: String,
    pub status: SubagentStatus,
    pub cancel_requested: bool,
    pub elapsed_secs: f64,
    pub finished_elapsed: Option<std::time::Duration>,
}

// ── SubagentHandle ───────────────────────────────────────────────────────────────

pub struct SubagentHandle {
    pub id: String,
    /// Numeric ID — same value that was passed to SubagentStart/SubagentDone events
    /// so the TUI can correlate by id without parsing "sa_N" strings.
    pub subagent_id: u64,
    pub agent_name: String,
    pub task_preview: String,
    pub model: String,
    pub system_prompt: String,
    pub started_at: std::time::Instant,
    pub timeout_secs: u64,
    pub authorization: Option<AuthorizationDiagnostic>,

    // Shared state updated by the subagent thread — one lock for everything.
    state: Arc<RwLock<SubagentState>>,

    // Channels
    steer_tx: Option<mpsc::UnboundedSender<String>>,
    shutdown_tx: Option<oneshot::Sender<()>>,
    /// OS thread running the subagent. Stored for graceful shutdown (join).
    // OS thread handle for graceful shutdown
    pub(crate) thread_handle: Option<std::thread::JoinHandle<()>>,

    // Final result
    result_rx: Option<oneshot::Receiver<SubagentResult>>,

    /// Set to true once subagent_collect has read the terminal result.
    collected: bool,
}

impl std::fmt::Debug for SubagentHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SubagentHandle")
            .field("id", &self.id)
            .field("subagent_id", &self.subagent_id)
            .field("agent_name", &self.agent_name)
            .field("model", &self.model)
            .finish_non_exhaustive()
    }
}

impl SubagentHandle {
    /// Construct a new handle. The state Arc is shared with the spawned subagent thread.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        id: String,
        subagent_id: u64,
        agent_name: String,
        task_preview: String,
        model: String,
        system_prompt: String,
        timeout_secs: u64,
        state: Arc<RwLock<SubagentState>>,
        steer_tx: Option<mpsc::UnboundedSender<String>>,
        shutdown_tx: Option<oneshot::Sender<()>>,
        result_rx: Option<oneshot::Receiver<SubagentResult>>,
    ) -> Self {
        Self {
            id,
            subagent_id,
            agent_name,
            task_preview,
            model,
            system_prompt,
            started_at: std::time::Instant::now(),
            timeout_secs,
            authorization: None,
            state,
            steer_tx,
            shutdown_tx,
            thread_handle: None,
            result_rx,
            collected: false,
        }
    }

    pub fn with_authorization(
        mut self,
        decision: &crate::orchestration::AuthorizedWorkerModel,
    ) -> Self {
        self.authorization = Some(AuthorizationDiagnostic {
            selection_source: decision.selection_source.as_str().into(),
            policy_digest: decision.policy_digest.clone(),
            catalog_snapshot_id: decision.catalog_snapshot_id.clone(),
            catalog_digest: decision.catalog_digest.clone(),
            correlation_id: decision.correlation_id.clone(),
            network_attempted: decision.network_attempted,
            cross_provider_grant_id: decision.cross_provider_grant_id.clone(),
        });
        self
    }

    pub fn set_terminal_diagnostic(&self, diagnostic: TerminalDiagnostic) {
        let mut state = self.state.write().unwrap();
        state.safe_fail(diagnostic);
    }

    /// Current status snapshot.
    pub fn status(&self) -> SubagentStatus {
        self.state
            .read()
            .unwrap_or_else(|p| p.into_inner())
            .status
            .clone()
    }

    /// Partial output accumulated so far.
    pub fn partial_output(&self) -> String {
        self.state
            .read()
            .unwrap_or_else(|p| p.into_inner())
            .partial_text
            .clone()
    }

    /// Snapshot of the tool log.
    pub fn tool_log(&self) -> Vec<String> {
        self.state
            .read()
            .unwrap_or_else(|p| p.into_inner())
            .tool_log
            .clone()
    }

    pub fn terminal_diagnostic(&self) -> Option<TerminalDiagnostic> {
        if let Some(diagnostic) = self.state.read().unwrap().terminal.clone() {
            return Some(diagnostic);
        }
        match self.status() {
            SubagentStatus::Running => None,
            SubagentStatus::Completed => Some(TerminalDiagnostic {
                category: TerminalCategory::Completed,
                code: "completed".into(),
                stage: "inference".into(),
                correlation_id: self.id.clone(),
                network_attempted: true,
                safe_message: safe_failure_message(&TerminalCategory::Completed).into(),
            }),
            SubagentStatus::TimedOut => Some(TerminalDiagnostic {
                category: TerminalCategory::TimedOut,
                code: "worker_timeout".into(),
                stage: "inference".into(),
                correlation_id: self.id.clone(),
                network_attempted: true,
                safe_message: safe_failure_message(&TerminalCategory::TimedOut).into(),
            }),
            SubagentStatus::Cancelled => Some(TerminalDiagnostic {
                category: TerminalCategory::Cancelled,
                code: "worker_cancelled".into(),
                stage: "inference".into(),
                correlation_id: self.id.clone(),
                network_attempted: true,
                safe_message: safe_failure_message(&TerminalCategory::Cancelled).into(),
            }),
            SubagentStatus::Failed(_) => Some(TerminalDiagnostic {
                category: TerminalCategory::ExecutionFailed,
                code: "inference_failed".into(),
                stage: "inference".into(),
                correlation_id: self.id.clone(),
                network_attempted: true,
                safe_message: safe_failure_message(&TerminalCategory::ExecutionFailed).into(),
            }),
        }
    }

    /// Snapshot of conversation state (for resume).
    pub fn conversation_state(&self) -> Vec<Value> {
        self.state
            .read()
            .unwrap_or_else(|p| p.into_inner())
            .conversation_state
            .clone()
    }

    /// Seconds since this handle was created.
    pub fn elapsed_secs(&self) -> f64 {
        self.started_at.elapsed().as_secs_f64()
    }

    /// Release execution-only resources after a worker has expired while
    /// retaining bounded terminal identity/output for collection and reconciliation.
    fn make_tombstone(&mut self) {
        debug_assert!(self.is_finished());
        self.steer_tx.take();
        self.shutdown_tx.take();
        self.result_rx.take();
        if let Some(thread) = self.thread_handle.take() {
            let _ = thread.join();
        }
        let mut state = self.state.write().unwrap_or_else(|p| p.into_inner());
        if state.partial_text.len() > TOMBSTONE_OUTPUT_MAX_BYTES {
            let mut start = state.partial_text.len() - TOMBSTONE_OUTPUT_MAX_BYTES;
            while !state.partial_text.is_char_boundary(start) {
                start += 1;
            }
            state.partial_text.drain(..start);
        }
        state.tool_log.clear();
        state.conversation_state.clear();
    }

    /// Send a steering message into the running subagent.
    pub fn steer(&self, message: &str) -> Result<(), String> {
        match &self.steer_tx {
            Some(tx) => tx
                .send(message.to_string())
                .map_err(|e| format!("steer channel closed: {e}")),
            None => Err("no steer channel on this handle".to_string()),
        }
    }

    /// Signal the subagent to shut down.
    /// Store the OS thread handle for graceful shutdown.
    pub fn set_thread_handle(&mut self, handle: std::thread::JoinHandle<()>) {
        self.thread_handle = Some(handle);
    }

    pub fn cancel(&mut self) {
        // Flag FIRST so finalize_subagent can read it before the thread exits.
        // Use poison-safe write in case the thread panicked holding the lock.
        self.state
            .write()
            .unwrap_or_else(|p| p.into_inner())
            .cancel_requested = true;
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
    }

    /// True if the subagent is no longer running.
    pub fn is_finished(&self) -> bool {
        !matches!(self.status(), SubagentStatus::Running)
    }

    /// Mark this handle as collected by subagent_collect.
    pub fn mark_collected(&mut self) {
        self.collected = true;
    }

    /// Whether subagent_collect has already read the result.
    pub fn is_collected(&self) -> bool {
        self.collected
    }

    /// Whether execution transports and resumable context have been discarded.
    pub fn is_tombstone(&self) -> bool {
        self.is_finished()
            && self.steer_tx.is_none()
            && self.shutdown_tx.is_none()
            && self.result_rx.is_none()
            && self.thread_handle.is_none()
    }

    /// Time elapsed since the subagent reached a terminal state.
    /// Returns `None` if still running or if `finished_at` has not been stamped yet.
    pub fn finished_elapsed(&self) -> Option<std::time::Duration> {
        self.state
            .read()
            .unwrap_or_else(|p| p.into_inner())
            .finished_at
            .map(|t| t.elapsed())
    }

    /// Consume the handle and wait for the final result.
    pub async fn collect(mut self) -> Result<SubagentResult, String> {
        match self.result_rx.take() {
            Some(rx) => rx
                .await
                .map_err(|_| "subagent result channel dropped".to_string()),
            None => Err("no result receiver — already collected or never set".to_string()),
        }
    }
}

// ── SubagentRegistry ─────────────────────────────────────────────────────────────

/// Finished-but-uncollected handles are retained this long before GC.
pub const FINISHED_HANDLE_TTL: std::time::Duration = std::time::Duration::from_secs(900); // 15 min
/// Maximum terminal text retained by a transport-free tombstone.
pub const TOMBSTONE_OUTPUT_MAX_BYTES: usize = 32 * 1024;

#[derive(Debug)]
pub struct SubagentRegistry {
    pub(crate) handles: HashMap<String, SubagentHandle>,
    // A local session driver owns worker launches until explicit user takeover.
    // Keep the cancelled fence after revocation so racing late registrations
    // are cancelled before their thread can start inference.
    spawn_cancellation: Option<crate::CancellationToken>,
    spawn_epoch: u64,
}

impl SubagentRegistry {
    pub fn new() -> Self {
        Self {
            handles: HashMap::new(),
            spawn_cancellation: None,
            spawn_epoch: 0,
        }
    }

    /// Register a handle and return its id.
    pub fn register(&mut self, handle: SubagentHandle) -> String {
        self.register_with_cancellation(handle, None)
    }

    pub fn register_with_cancellation(
        &mut self,
        mut handle: SubagentHandle,
        origin: Option<&crate::CancellationToken>,
    ) -> String {
        if origin.is_some_and(|c| c.is_cancelled())
            || self
                .spawn_cancellation
                .as_ref()
                .is_some_and(|c| c.is_cancelled())
        {
            handle.cancel();
        }
        let id = handle.id.clone();
        self.handles.insert(id.clone(), handle);
        id
    }

    /// Frontend-owned launch fence. Model tools cannot clear or replace it.
    pub fn set_spawn_cancellation(&mut self, cancel: Option<crate::CancellationToken>) -> u64 {
        self.spawn_epoch = self
            .spawn_epoch
            .checked_add(1)
            .expect("worker launch epoch exhausted");
        self.spawn_cancellation = cancel;
        self.spawn_epoch
    }

    /// A stale deadline/Drop must not cancel workers belonging to newer work.
    pub fn cancel_spawn_epoch(&mut self, epoch: u64) {
        if self.spawn_epoch == epoch {
            for handle in self.iter_mut_handles() {
                handle.cancel();
            }
        }
    }

    pub fn get(&self, id: &str) -> Option<&SubagentHandle> {
        self.handles.get(id)
    }

    pub fn get_mut(&mut self, id: &str) -> Option<&mut SubagentHandle> {
        self.handles.get_mut(id)
    }

    pub fn remove(&mut self, id: &str) -> Option<SubagentHandle> {
        self.handles.remove(id)
    }

    /// Drop finished execution transports and resumable context while retaining
    /// bounded terminal diagnostics and output under normal session retention.
    pub fn release_finished_resources(&mut self, id: &str) {
        if let Some(handle) = self.handles.get_mut(id) {
            if handle.is_finished() {
                handle.make_tombstone();
            }
        }
    }

    /// Returns (id, agent_name, status) for every tracked handle.
    pub fn list_active(&self) -> Vec<(String, String, SubagentStatus)> {
        self.handles
            .values()
            .map(|h| (h.id.clone(), h.agent_name.clone(), h.status()))
            .collect()
    }

    /// Snapshot of every tracked handle for TUI reconcile.
    /// Reads all state under the lock POISON-SAFE; exposes cancel_requested
    /// which list_active() does not.
    pub fn display_rows(&self) -> Vec<SubagentDisplayRow> {
        self.handles
            .values()
            .map(|h| {
                let s = h.state.read().unwrap_or_else(|p| p.into_inner());
                SubagentDisplayRow {
                    subagent_id: h.subagent_id,
                    agent_name: h.agent_name.clone(),
                    status: s.status.clone(),
                    cancel_requested: s.cancel_requested,
                    elapsed_secs: h.started_at.elapsed().as_secs_f64(),
                    finished_elapsed: s.finished_at.map(|t| t.elapsed()),
                }
            })
            .collect()
    }

    /// Drop handles that are no longer running.
    /// Iterate over all handles mutably (for bulk operations like cancel-all).
    pub fn iter_mut_handles(&mut self) -> impl Iterator<Item = &mut SubagentHandle> {
        self.handles.values_mut()
    }

    /// Reap a finished handle iff:
    ///   (a) its result was collected via subagent_collect, OR
    ///   (b) it has been finished longer than `ttl` (abandoned).
    /// Finished-but-uncollected handles inside the TTL are RETAINED so the
    /// completion event can wake the parent and collect still succeeds.
    ///
    /// Additionally, handles whose OS thread is still running (e.g. still in
    /// the finalizer path) are deferred — joining a live thread blocks the TUI
    /// loop. They will be reaped on the next cleanup pass once the thread exits.
    /// Orchestration-required expired handles become bounded tombstones instead
    /// of being removed; a later collect/reconcile makes them normally reapable.
    pub fn cleanup_finished_with_ttl(&mut self, ttl: std::time::Duration) {
        self.cleanup_finished_with_ttl_and_retention(ttl, |_| false);
    }

    /// Reap normally, except that an expired uncollected handle still required by
    /// orchestration is downgraded to a bounded, collectible tombstone.
    pub fn cleanup_finished_with_ttl_and_retention(
        &mut self,
        ttl: std::time::Duration,
        retain_unreconciled: impl Fn(&str) -> bool,
    ) {
        let candidates: Vec<(String, bool)> = self
            .handles
            .iter()
            .filter(|(_, h)| {
                h.is_finished()
                    && h.thread_handle.as_ref().map_or(true, |t| t.is_finished())
                    && (h.is_collected() || h.finished_elapsed().is_some_and(|d| d >= ttl))
            })
            .map(|(id, h)| (id.clone(), !h.is_collected() && retain_unreconciled(id)))
            .collect();
        for (id, retain) in candidates {
            if retain {
                self.release_finished_resources(&id);
            } else if let Some(mut handle) = self.handles.remove(&id) {
                if let Some(th) = handle.thread_handle.take() {
                    let _ = th.join();
                }
            }
        }
    }

    /// Reap finished handles using the production TTL.
    /// Finished-but-uncollected handles within the TTL window are retained.
    pub fn cleanup_finished(&mut self) {
        self.cleanup_finished_with_ttl(FINISHED_HANDLE_TTL)
    }
}

/// Engine-owned housekeeping seam: reap finished subagent handles at end of
/// each turn.  Acquires the lock poison-safely — a panicked subagent thread
/// must not prevent the turn from completing.
///
/// Intended to be called from the `tokio::spawn` wrapper in `runtime/mod.rs`
/// after `run_stream_internal` returns and BEFORE sending `SessionEvent::Done`.
pub(crate) fn reap_finished_with_ttl(
    registry: &Arc<std::sync::Mutex<SubagentRegistry>>,
    orchestration: Option<&crate::orchestration::OrchestrationRuntime>,
    ttl: std::time::Duration,
) {
    let retained = orchestration
        .map(|runtime| runtime.unreconciled_runtime_handles())
        .unwrap_or_default();
    let cleanup = |guard: &mut SubagentRegistry| {
        let before = guard
            .handles
            .keys()
            .cloned()
            .collect::<std::collections::HashSet<_>>();
        guard.cleanup_finished_with_ttl_and_retention(ttl, |id| retained.contains(id));
        before
            .into_iter()
            .filter(|id| !guard.handles.contains_key(id))
            .collect::<std::collections::HashSet<_>>()
    };
    // Release the registry lock before entering orchestration (same lock order
    // as the retention snapshot above). Live workers are never retired.
    let removed = match registry.lock() {
        Ok(mut guard) => cleanup(&mut guard),
        Err(poisoned) => cleanup(&mut poisoned.into_inner()),
    };
    if let Some(orchestration) = orchestration {
        orchestration.retire_reconciled(&removed);
    }
}

pub fn reap_finished(
    registry: &Arc<std::sync::Mutex<SubagentRegistry>>,
    orchestration: Option<&crate::orchestration::OrchestrationRuntime>,
) {
    reap_finished_with_ttl(registry, orchestration, FINISHED_HANDLE_TTL);
}

impl Default for SubagentRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl SubagentStatus {
    pub fn as_str(&self) -> &str {
        match self {
            SubagentStatus::Running => "running",
            SubagentStatus::Completed => "completed",
            SubagentStatus::Cancelled => "cancelled",
            SubagentStatus::TimedOut => "timed_out",
            SubagentStatus::Failed(_) => "failed",
        }
    }

    /// Returns the failure reason string if this is a `Failed` variant.
    pub fn failure_reason(&self) -> Option<&str> {
        match self {
            SubagentStatus::Failed(r) => Some(r.as_str()),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::{mpsc, oneshot};

    // Keep receivers alive so channels don't close during tests
    struct TestHandle {
        handle: SubagentHandle,
        _steer_rx: mpsc::UnboundedReceiver<String>,
        _shutdown_rx: oneshot::Receiver<()>,
    }

    fn make_test_handle(id: &str) -> TestHandle {
        let state = Arc::new(RwLock::new(SubagentState::new()));
        let (steer_tx, steer_rx) = mpsc::unbounded_channel();
        let (shutdown_tx, shutdown_rx) = oneshot::channel();
        let (_result_tx, result_rx) = oneshot::channel();
        // Parse numeric id from "sa_N" or use 0 as fallback for tests
        let numeric_id: u64 = id
            .strip_prefix("sa_")
            .and_then(|n| n.parse().ok())
            .unwrap_or(0);
        TestHandle {
            handle: SubagentHandle::new(
                id.to_string(),
                numeric_id,
                "test-agent".to_string(),
                "test task".to_string(),
                "claude-sonnet-4-6".to_string(),
                "You are a test agent.".to_string(),
                300,
                state,
                Some(steer_tx),
                Some(shutdown_tx),
                Some(result_rx),
            ),
            _steer_rx: steer_rx,
            _shutdown_rx: shutdown_rx,
        }
    }

    fn make_handle(id: &str) -> SubagentHandle {
        make_test_handle(id).handle
    }

    fn make_finished_handle(id: &str) -> SubagentHandle {
        let h = make_handle(id);
        {
            let mut s = h.state.write().unwrap();
            s.status = SubagentStatus::Completed;
            s.finished_at = Some(std::time::Instant::now());
        }
        h
    }

    #[test]
    fn typed_failure_drops_raw_provider_secret_canary() {
        let mut state = SubagentState::new();
        let canary = "Bearer secret-canary-provider-body";
        state.safe_fail(TerminalDiagnostic {
            category: TerminalCategory::Credential,
            code: "credential_unavailable".into(),
            stage: "credential".into(),
            correlation_id: "sa_safe".into(),
            network_attempted: false,
            safe_message: canary.into(),
        });
        let rendered = format!("{:?}{:?}", state.status, state.terminal);
        assert!(!rendered.contains(canary));
        assert!(rendered.contains("credential_unavailable"));
    }

    #[test]
    fn late_registration_cannot_inherit_new_turn_and_stale_sweep_cannot_cancel_it() {
        let mut registry = SubagentRegistry::new();
        let old = crate::CancellationToken::new();
        let epoch = registry.set_spawn_cancellation(Some(old.clone()));
        old.cancel();
        registry.cancel_spawn_epoch(epoch);
        // User takeover/new grant happened while old start was preempted.
        registry.set_spawn_cancellation(None);
        registry.register(make_handle("new-user"));
        registry.register_with_cancellation(make_handle("old-late"), Some(&old));
        assert!(
            registry
                .get("old-late")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        assert!(
            !registry
                .get("new-user")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        registry.cancel_spawn_epoch(epoch);
        assert!(
            !registry
                .get("new-user")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        let new_grant = crate::CancellationToken::new();
        let next = registry.set_spawn_cancellation(Some(new_grant));
        registry.register(make_handle("new-grant"));
        registry.register_with_cancellation(make_handle("old-later"), Some(&old));
        registry.cancel_spawn_epoch(epoch);
        assert!(
            !registry
                .get("new-grant")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        assert!(
            registry
                .get("old-later")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        registry.cancel_spawn_epoch(next);
        assert!(
            registry
                .get("new-grant")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
    }

    #[test]
    fn revoked_launch_fence_cancels_late_registration_until_user_takeover() {
        let mut registry = SubagentRegistry::new();
        let cancel = crate::CancellationToken::new();
        registry.set_spawn_cancellation(Some(cancel.clone()));
        registry.register(make_handle("before"));
        assert!(!registry.display_rows()[0].cancel_requested);
        cancel.cancel();
        let late = make_test_handle("late");
        let mut shutdown = late._shutdown_rx;
        registry.register(late.handle);
        assert!(shutdown.try_recv().is_ok());
        assert!(
            registry
                .get("late")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
        registry.set_spawn_cancellation(None);
        registry.register(make_handle("user"));
        assert!(
            !registry
                .get("user")
                .unwrap()
                .state
                .read()
                .unwrap()
                .cancel_requested
        );
    }

    #[test]
    fn handle_initial_status_is_running() {
        let h = make_handle("sa_1");
        assert_eq!(h.status(), SubagentStatus::Running);
        assert!(!h.is_finished());
    }

    #[test]
    fn handle_partial_output_empty_initially() {
        let h = make_handle("sa_1");
        assert_eq!(h.partial_output(), "");
        assert!(h.tool_log().is_empty());
        assert!(h.conversation_state().is_empty());
    }

    #[test]
    fn handle_status_reflects_state_change() {
        let h = make_handle("sa_1");
        {
            let mut s = h.state.write().unwrap();
            s.status = SubagentStatus::Completed;
            s.partial_text = "done!".to_string();
        }
        assert_eq!(h.status(), SubagentStatus::Completed);
        assert!(h.is_finished());
        assert_eq!(h.partial_output(), "done!");
    }

    #[test]
    fn handle_steer_sends_message() {
        let th = make_test_handle("sa_1");
        assert!(th.handle.steer("redirect").is_ok());
    }

    #[test]
    fn handle_steer_fails_without_channel() {
        let state = Arc::new(RwLock::new(SubagentState::new()));
        let (_shutdown_tx, _) = oneshot::channel::<()>();
        let (_, result_rx) = oneshot::channel();
        let h = SubagentHandle::new(
            "sa_1".into(),
            1,
            "test".into(),
            "task".into(),
            "model".into(),
            "prompt".into(),
            300,
            state,
            None,
            None,
            Some(result_rx),
        );
        assert!(h.steer("msg").is_err());
    }

    #[test]
    fn handle_cancel_consumes_shutdown() {
        let mut h = make_handle("sa_1");
        h.cancel(); // first call sends
        h.cancel(); // second call is no-op (already taken)
    }

    #[test]
    fn handle_elapsed_increases() {
        let h = make_handle("sa_1");
        std::thread::sleep(std::time::Duration::from_millis(10));
        assert!(h.elapsed_secs() > 0.0);
    }

    #[test]
    fn registry_register_and_get() {
        let mut reg = SubagentRegistry::new();
        let h = make_handle("sa_1");
        reg.register(h);
        assert!(reg.get("sa_1").is_some());
        assert!(reg.get("sa_99").is_none());
    }

    #[test]
    fn registry_remove() {
        let mut reg = SubagentRegistry::new();
        reg.register(make_handle("sa_1"));
        assert!(reg.remove("sa_1").is_some());
        assert!(reg.get("sa_1").is_none());
    }

    #[test]
    fn registry_list_active() {
        let mut reg = SubagentRegistry::new();
        reg.register(make_handle("sa_1"));
        reg.register(make_handle("sa_2"));
        let active = reg.list_active();
        assert_eq!(active.len(), 2);
    }

    // Rewritten: finished-but-uncollected handles are RETAINED (new semantics).
    // The old test encoded the buggy behavior where any finished handle was removed.
    #[test]
    fn registry_cleanup_finished() {
        let mut reg = SubagentRegistry::new();
        let mut h = make_finished_handle("sa_1");
        h.mark_collected(); // mark as collected — should be reaped
        reg.register(h);
        reg.register(make_handle("sa_2")); // still running — must survive

        reg.cleanup_finished();

        assert!(
            reg.get("sa_1").is_none(),
            "collected finished handle must be reaped"
        );
        assert!(
            reg.get("sa_2").is_some(),
            "running handle must survive cleanup"
        );
    }

    // U1: finished, not collected, fresh → retained within TTL
    #[test]
    fn reaper_retains_finished_uncollected_within_ttl() {
        let mut reg = SubagentRegistry::new();
        let h = make_finished_handle("sa_1"); // not collected
        reg.register(h);

        reg.cleanup_finished_with_ttl(std::time::Duration::from_secs(900));

        assert!(
            reg.get("sa_1").is_some(),
            "finished-uncollected handle within TTL must be retained"
        );
    }

    // U2: finished + mark_collected → removed immediately
    #[test]
    fn reaper_reaps_finished_and_collected() {
        let mut reg = SubagentRegistry::new();
        let mut h = make_finished_handle("sa_1");
        h.mark_collected();
        reg.register(h);

        reg.cleanup_finished_with_ttl(std::time::Duration::from_secs(900));

        assert!(
            reg.get("sa_1").is_none(),
            "finished+collected handle must be reaped"
        );
    }

    // U3: finished, not collected, and TTL-expired, but still required by
    // orchestration → retained as a resource-free tombstone.
    #[test]
    fn reaper_tombstones_unreconciled_after_ttl() {
        let mut reg = SubagentRegistry::new();
        let h = make_finished_handle("sa_1"); // not collected
        h.state.write().unwrap().partial_text = "🦀".repeat(25_000);
        reg.register(h);

        reg.cleanup_finished_with_ttl_and_retention(std::time::Duration::ZERO, |_| true);

        let h = reg
            .get("sa_1")
            .expect("gate-relevant handle must remain collectible");
        assert!(h.is_tombstone(), "expired handle must release transports");
        assert!(h.partial_output().len() <= TOMBSTONE_OUTPUT_MAX_BYTES);
    }

    #[test]
    fn reaper_reaps_abandoned_after_ttl_without_orchestration() {
        let mut reg = SubagentRegistry::new();
        reg.register(make_finished_handle("sa_1"));

        reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);

        assert!(reg.get("sa_1").is_none());
    }

    // U5-reaper: finished-status handle with a still-sleeping OS thread is deferred;
    // after thread exits a second cleanup reaps it.
    #[test]
    fn reaper_defers_handle_with_live_thread() {
        let mut reg = SubagentRegistry::new();
        let mut h = make_finished_handle("sa_live");
        h.mark_collected(); // collected — would normally be reaped immediately

        // Attach a real thread that sleeps briefly (simulates thread still in finalizer)
        let barrier = Arc::new(std::sync::Barrier::new(2));
        let barrier_c = Arc::clone(&barrier);
        let thread = std::thread::spawn(move || {
            // Wait for the test to run its first cleanup
            barrier_c.wait();
            // Now sleep a bit more before truly exiting
            std::thread::sleep(std::time::Duration::from_millis(20));
        });
        h.set_thread_handle(thread);
        reg.register(h);

        // Signal thread (it's waiting on barrier)
        barrier.wait();

        // First cleanup: thread is still alive → handle must be deferred (retained)
        // The thread may or may not have exited by the time we get here, so we only
        // assert the positive case when we know the thread is still live.
        // Instead: sleep briefly to be sure the thread has NOT exited, then check.
        // (The thread sleeps 20ms after barrier; we check immediately after barrier.)
        reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);
        // Handle may still be present (thread live) or absent (thread exited) — both ok.
        // What must NOT happen: a blocking join. The test itself completes in < 1s.

        // Wait for thread to truly exit
        std::thread::sleep(std::time::Duration::from_millis(50));

        // Second cleanup: thread is finished → handle must be reaped
        reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);
        assert!(
            reg.get("sa_live").is_none(),
            "handle must be reaped once OS thread exits"
        );
    }

    // U4: running handles never touched by any reaper variant
    #[test]
    fn reaper_never_touches_running() {
        let mut reg = SubagentRegistry::new();
        reg.register(make_handle("sa_1")); // Running
        reg.register(make_handle("sa_2")); // Running

        reg.cleanup_finished_with_ttl(std::time::Duration::ZERO);
        assert!(reg.get("sa_1").is_some());
        assert!(reg.get("sa_2").is_some());

        reg.cleanup_finished(); // production TTL
        assert!(reg.get("sa_1").is_some());
        assert!(reg.get("sa_2").is_some());
    }

    #[test]
    fn subagent_state_new_defaults() {
        let s = SubagentState::new();
        assert_eq!(s.status, SubagentStatus::Running);
        assert!(s.partial_text.is_empty());
        assert!(s.tool_log.is_empty());
        assert!(s.conversation_state.is_empty());
        assert!(s.finished_at.is_none());
        assert!(
            !s.cancel_requested,
            "cancel_requested must default to false"
        );
    }

    #[test]
    fn subagent_status_as_str() {
        assert_eq!(SubagentStatus::Running.as_str(), "running");
        assert_eq!(SubagentStatus::Completed.as_str(), "completed");
        assert_eq!(SubagentStatus::Cancelled.as_str(), "cancelled");
        assert_eq!(SubagentStatus::TimedOut.as_str(), "timed_out");
        assert_eq!(SubagentStatus::Failed("oops".into()).as_str(), "failed");
    }

    // V1: cancel() sets cancel_requested on shared state before sending signal.
    #[test]
    fn cancel_sets_cancel_requested_flag() {
        let mut h = make_handle("sa_cancel");
        {
            let s = h.state.read().unwrap();
            assert!(!s.cancel_requested, "cancel_requested must start false");
        }
        h.cancel();
        {
            let s = h.state.read().unwrap();
            assert!(
                s.cancel_requested,
                "cancel_requested must be true after cancel()"
            );
        }
        // Second call is a no-op (shutdown_tx already taken)
        h.cancel();
    }

    // V2: cancelled status counts as finished, is_finished() returns true.
    #[test]
    fn cancelled_status_is_finished() {
        let h = make_handle("sa_c1");
        {
            let mut s = h.state.write().unwrap();
            s.status = SubagentStatus::Cancelled;
        }
        assert!(h.is_finished(), "Cancelled must count as finished");
    }

    // D1: display_rows returns a row for each registered handle
    #[test]
    fn display_rows_running_handle() {
        let mut reg = SubagentRegistry::new();
        reg.register(make_handle("sa_1"));
        reg.register(make_handle("sa_2"));
        let rows = reg.display_rows();
        assert_eq!(rows.len(), 2);
        for row in &rows {
            assert_eq!(row.status, SubagentStatus::Running);
            assert!(!row.cancel_requested);
            assert!(row.elapsed_secs >= 0.0);
            assert!(row.finished_elapsed.is_none());
        }
    }

    // D2: display_rows reflects finished status and cancel_requested flag
    #[test]
    fn display_rows_finished_and_cancel_requested() {
        let mut reg = SubagentRegistry::new();
        let h = make_finished_handle("sa_1");
        reg.register(h);

        // Make a handle with cancel_requested = true
        let th = make_test_handle("sa_2");
        {
            let mut s = th.handle.state.write().unwrap();
            s.cancel_requested = true;
        }
        reg.register(th.handle);

        let rows = reg.display_rows();
        assert_eq!(rows.len(), 2);

        let finished = rows.iter().find(|r| r.subagent_id == 1).unwrap();
        assert_eq!(finished.status, SubagentStatus::Completed);
        assert!(!finished.cancel_requested);
        assert!(finished.finished_elapsed.is_some());

        let cancelling = rows.iter().find(|r| r.subagent_id == 2).unwrap();
        assert!(cancelling.cancel_requested);
    }

    // D3: display_rows on empty registry returns empty vec
    #[test]
    fn display_rows_empty_registry() {
        let reg = SubagentRegistry::new();
        assert!(reg.display_rows().is_empty());
    }

    // ── reap_finished free-function tests ────────────────────────────────────

    /// R1: reap_finished with a collected+finished handle in an Arc<Mutex<>> reaps it.
    #[test]
    fn reap_finished_headless_reaps_collected() {
        let registry = Arc::new(std::sync::Mutex::new(SubagentRegistry::new()));
        {
            let mut reg = registry.lock().unwrap();
            let mut h = make_finished_handle("sa_r1");
            h.mark_collected();
            reg.register(h);
        }
        super::reap_finished(&registry, None);
        assert!(
            registry.lock().unwrap().get("sa_r1").is_none(),
            "reap_finished must reap collected+finished handle"
        );
    }

    /// R2: reap_finished retains running handles — must not touch live work.
    #[test]
    fn reap_finished_headless_retains_running() {
        let registry = Arc::new(std::sync::Mutex::new(SubagentRegistry::new()));
        {
            let mut reg = registry.lock().unwrap();
            reg.register(make_handle("sa_running"));
        }
        super::reap_finished(&registry, None);
        assert!(
            registry.lock().unwrap().get("sa_running").is_some(),
            "reap_finished must not touch running handles"
        );
    }

    /// R3: reap_finished recovers from a poisoned lock without panicking.
    #[test]
    fn reap_finished_poisoned_lock_recovery() {
        let registry = Arc::new(std::sync::Mutex::new(SubagentRegistry::new()));
        {
            let mut reg = registry.lock().unwrap();
            let mut h = make_finished_handle("sa_poison");
            h.mark_collected();
            reg.register(h);
        }

        // Poison the mutex by panicking inside a lock guard.
        let registry_c = Arc::clone(&registry);
        let _ = std::panic::catch_unwind(|| {
            let _guard = registry_c.lock().unwrap();
            panic!("deliberately poison the mutex");
        });

        // reap_finished must not panic despite poisoned lock.
        super::reap_finished(&registry, None);

        // The handle must be reaped (cleanup ran despite poison).
        let result = registry
            .lock()
            .unwrap_or_else(|p| p.into_inner())
            .get("sa_poison")
            .is_none();
        assert!(result, "reap_finished must run cleanup even after poison");
    }
}

// Separate module so finalize_subagent is accessible via its pub(crate) path.
#[cfg(test)]
mod cancelled_wake_tests {
    use super::*;
    use std::sync::{Arc, RwLock};

    // V3: finalize_subagent with cancel_requested=true → status Cancelled,
    // no event published, finished_at stamped.
    // Test name matches spec: cancelled_suppresses_wake.
    #[test]
    fn cancelled_suppresses_wake() {
        use crate::events::EventQueue;
        use crate::tools::finalize_subagent;

        let state = Arc::new(RwLock::new(SubagentState::new()));
        {
            let mut s = state.write().unwrap();
            s.status = SubagentStatus::Completed; // thread set Completed before cancel flag noticed
            s.cancel_requested = true; // cancel() was called
        }
        let queue = Arc::new(EventQueue::new(100));

        finalize_subagent(
            &state,
            Some(&queue),
            "sa_v3",
            42,
            "test-agent",
            std::time::Instant::now(),
            None,
        );

        // Queue must be empty — no wake event for cancelled subagents.
        assert!(
            queue.is_empty(),
            "cancelled subagent must not publish a wake event"
        );

        // Status must have been re-labelled to Cancelled.
        let s = state.read().unwrap();
        assert_eq!(
            s.status,
            SubagentStatus::Cancelled,
            "status must be Cancelled when cancel_requested"
        );

        // finished_at must be stamped.
        assert!(
            s.finished_at.is_some(),
            "finished_at must be stamped by finalizer"
        );
    }
}