oo-ide 0.0.4

//! Async task executor with incremental output streaming.
//!
//! [`TaskExecutor`] bridges [`crate::task_registry::TaskRegistry`] and the OS
//! process layer.  When the registry marks a task as running, the caller
//! invokes [`TaskExecutor::spawn`] to start the actual process asynchronously.
//!
//! # Output fan-out
//!
//! Raw output bytes are broadcast on a [`tokio::sync::broadcast`] channel so
//! that multiple independent consumers (future: log storage, diagnostic
//! parser, UI render) can each subscribe without blocking one another.  Use
//! [`TaskExecutor::subscribe`] to obtain a receiver.
//!
//! # Cancellation
//!
//! The `CancellationToken` stored in each `crate::task_registry::Task` is monitored inside the
//! spawned async task. When it fires, the child process is killed and the
//! task is reported back as `TaskStatus::Cancelled`.
//!
//! # Integration
//!
//! The executor sends [`Operation::TaskFinished`] (or [`Operation::TaskStarted`])
//! back to the main event loop via the `op_tx` channel so that
//! `apply_operation` can call [`crate::task_registry::TaskRegistry::mark_finished`] on the main
//! thread and keep the registry consistent.

use std::collections::HashMap;
use std::sync::Arc;

use tokio::io::AsyncReadExt as _;
use tokio::process::Command;
use tokio::sync::broadcast;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

use crate::diagnostics_extractor::DiagnosticsExtractor;
use crate::log_matcher::{MatcherEngine, MatcherRegistry};
use crate::log_store::LogStore;
use crate::operation::Operation;
use crate::task_registry::{TaskId, TaskStatus};
use crate::vt_parser::{StyledLine, TerminalParser};

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Which stdio stream an output chunk originated from.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OutputStream {
    Stdout,
    Stderr,
}

/// A raw byte chunk read from a task's stdout or stderr.
///
/// Retained for potential future consumers that need raw bytes (e.g. a
/// protocol bridge).  The primary pipeline uses [`ParsedOutputEvent`] instead.
#[derive(Clone, Debug)]
pub struct TaskOutputEvent {
    pub task_id: TaskId,
    pub stream: OutputStream,
    /// Raw bytes — may contain partial UTF-8 sequences or ANSI escape codes.
    pub chunk: Vec<u8>,
}

/// A parsed, styled output event produced after running bytes through
/// [`TerminalParser`].
#[derive(Clone, Debug)]
pub struct ParsedOutputEvent {
    pub task_id: TaskId,
    pub stream: OutputStream,
    /// Completed styled lines decoded from the raw byte stream.
    pub lines: Vec<StyledLine>,
}

/// High-level event broadcast to all subscribers.
#[derive(Clone, Debug)]
pub enum TaskEvent {
    /// The executor started working on this task (process about to be spawned).
    Started(TaskId),
    /// One or more styled lines of output arrived from the running process.
    ParsedOutput(ParsedOutputEvent),
    /// The task reached a terminal state.
    Finished(TaskId, TaskStatus),
}

// ---------------------------------------------------------------------------
// TaskExecutor
// ---------------------------------------------------------------------------

/// Spawns tasks as OS processes and streams their output over a broadcast
/// channel.
///
/// Lives on `AppState` on the main thread; all async work is delegated to
/// Tokio task handles stored internally.
pub struct TaskExecutor {
    event_tx: broadcast::Sender<TaskEvent>,
    /// Active async handles — used to abort tasks on cancellation.
    handles: HashMap<TaskId, JoinHandle<()>>,
    /// Snapshot provider for matchers – owned Arc to the central registry.
    matcher_registry: std::sync::Arc<MatcherRegistry>,
}

impl TaskExecutor {
    /// Capacity of the broadcast ring buffer.  Lagging receivers will miss
    /// older output chunks but the executor itself never blocks.
    const BROADCAST_CAPACITY: usize = 256;

    pub fn new() -> Self {
        let (event_tx, _) = broadcast::channel(Self::BROADCAST_CAPACITY);
        Self {
            event_tx,
            handles: HashMap::new(),
            matcher_registry: std::sync::Arc::new(MatcherRegistry::new()),
        }
    }

    /// Create a TaskExecutor that uses the given MatcherRegistry snapshot
    /// provider.  This is used by AppState to ensure spawned tasks use the
    /// central ExtensionManager matcher set.
    pub fn with_matcher_registry(matcher_registry: std::sync::Arc<MatcherRegistry>) -> Self {
        let (event_tx, _) = broadcast::channel(Self::BROADCAST_CAPACITY);
        Self {
            event_tx,
            handles: HashMap::new(),
            matcher_registry,
        }
    }

    /// Subscribe to the live event stream.
    ///
    /// Receivers are independent; a slow receiver simply misses old messages
    /// once the ring buffer wraps (broadcast semantics).
    pub fn subscribe(&self) -> broadcast::Receiver<TaskEvent> {
        self.event_tx.subscribe()
    }

    /// Start executing `command` in a background Tokio task.
    ///
    /// * Broadcasts [`TaskEvent::Started`] immediately.
    /// * Streams stdout and stderr as [`TaskEvent::ParsedOutput`] events,
    ///   with raw bytes parsed through per-stream [`TerminalParser`] instances.
    /// * Runs a [`DiagnosticsExtractor`] (shared between both streams) and
    ///   sends [`Operation::AddIssue`] for each matched diagnostic line.
    /// * Runs a [`MatcherEngine`] (separate instance per stream) for each set
    ///   of extension-defined log matchers; emits [`Operation::AddIssue`] for
    ///   every completed block and on end-of-stream flush.
    /// * Watches `cancellation_token`; kills the process if it fires.
    /// * Sends [`Operation::TaskFinished`] (or Cancelled) back through `op_tx`
    ///   so the main loop can update the registry.
    ///
    /// `marker` is the issue registry marker used for all extracted issues
    /// (typically `"task:{queue}:{target}"`).  The caller is responsible for
    /// clearing stale issues with this marker before calling `spawn`.
    ///
    /// `matchers` is the set of compiled log matchers contributed by loaded
    /// extensions.  Pass an empty `Vec` when no extensions are active.
    ///
    /// If `log_store` is provided, a dedicated log-writing sub-task subscribes
    /// to the broadcast channel and writes each line to
    /// `.oo/cache/tasks/<task_id>.log`, then compresses the file on completion.
    #[allow(clippy::too_many_arguments)]
    pub fn spawn(
        &mut self,
        task_id: TaskId,
        command: String,
        cancellation_token: CancellationToken,
        marker: String,
        log_store: Option<LogStore>,
        op_tx: &tokio::sync::mpsc::UnboundedSender<Vec<Operation>>,
    ) {
        // Abort any existing handle for this task ID (defensive — should be
        // a no-op in normal operation because schedule_task cancelled it first).
        if let Some(old) = self.handles.remove(&task_id) {
            old.abort();
        }

        let event_tx = self.event_tx.clone();
        let op_tx = op_tx.clone();

        // Subscribe to the broadcast channel BEFORE spawning the async task so
        // that the log writer's receiver is set up before any events are sent.
        // Broadcast ring-buffers up to BROADCAST_CAPACITY messages for us.
        let log_rx = log_store.as_ref().map(|_| event_tx.subscribe());

        // Build a shared diagnostics extractor for this task execution.
        // The source tag is derived from the marker (best-effort).
        let source = marker
            .strip_prefix("task:")
            .and_then(|s| s.split(':').next())
            .unwrap_or("task")
            .to_string();
        let extractor = Arc::new(DiagnosticsExtractor::new(marker.clone(), source));

        // Obtain a snapshot of active matchers from the central registry and
        // build per-stream MatcherEngine instances.  Clone is O(1) per matcher
        // because regex fields are Arc-wrapped.
        let matchers = self.matcher_registry.iter_active();
        let stdout_engine = if matchers.is_empty() {
            None
        } else {
            Some(MatcherEngine::new(matchers.clone(), marker.clone()))
        };
        let stderr_engine = if matchers.is_empty() {
            None
        } else {
            Some(MatcherEngine::new(matchers, marker))
        };

        let handle = tokio::spawn(async move {
            // Spawn the log-writing task BEFORE broadcasting Started so it is
            // ready to receive every subsequent event.
            if let (Some(ls), Some(mut rx)) = (log_store, log_rx) {
                tokio::spawn(async move {
                    let mut writer = match ls.open_writer(task_id).await {
                        Ok(w) => w,
                        Err(e) => {
                            log::warn!(
                                "log_store: failed to open writer for task {}: {e}",
                                task_id.0
                            );
                            return;
                        }
                    };

                    let mut finished_normally = false;
                    loop {
                        match rx.recv().await {
                            Ok(TaskEvent::ParsedOutput(ev)) if ev.task_id == task_id => {
                                for line in &ev.lines {
                                    if let Err(e) = writer.append_line(&line.text).await {
                                        log::warn!("log_store: write error: {e}");
                                    }
                                }
                            }
                            Ok(TaskEvent::Finished(id, _)) if id == task_id => {
                                finished_normally = true;
                                break;
                            }
                            Err(broadcast::error::RecvError::Closed) => break,
                            Err(broadcast::error::RecvError::Lagged(n)) => {
                                log::warn!("log_store: receiver lagged by {n} messages");
                            }
                            _ => {}
                        }
                    }

                    if finished_normally {
                        if let Err(e) = writer.close_and_compress(&ls).await {
                            log::warn!(
                                "log_store: compress error for task {}: {e}",
                                task_id.0
                            );
                            // Attempt plain close as a fallback.
                            let _ = ls.log_path(task_id); // path still exists
                        }
                    } else {
                        // Cancelled or channel closed — keep partial log uncompressed.
                        if let Err(e) = writer.close().await {
                            log::warn!(
                                "log_store: close error for task {}: {e}",
                                task_id.0
                            );
                        }
                    }
                });
            }

            // Signal that we are starting.
            let _ = event_tx.send(TaskEvent::Started(task_id));
            let _ = op_tx.send(vec![Operation::TaskStarted(task_id)]);

            // Spawn the OS process.
            let child_result = build_command(&command)
                .stdout(std::process::Stdio::piped())
                .stderr(std::process::Stdio::piped())
                // Ensure the process is killed if this future is dropped or
                // aborted before `child.wait()` is reached.
                .kill_on_drop(true)
                .spawn();

            let mut child = match child_result {
                Ok(c) => c,
                Err(_) => {
                    let _ = event_tx.send(TaskEvent::Finished(task_id, TaskStatus::Error));
                    let _ = op_tx.send(vec![Operation::TaskFinished {
                        id: task_id,
                        status: TaskStatus::Error,
                    }]);
                    return;
                }
            };

            // Take the stdio handles before `child` is borrowed by `wait`.
            let stdout = child.stdout.take().expect("stdout was piped");
            let stderr = child.stderr.take().expect("stderr was piped");

            // Each stream gets its own stateful VT100 parser and MatcherEngine;
            // they run concurrently in separate sub-tasks and flush on stream close.
            let stdout_tx = event_tx.clone();
            let stdout_op_tx = op_tx.clone();
            let stdout_extractor = Arc::clone(&extractor);
            let stdout_task = tokio::spawn(async move {
                stream_output_parsed(
                    task_id,
                    OutputStream::Stdout,
                    stdout,
                    TerminalParser::new(),
                    &stdout_extractor,
                    stdout_engine,
                    &stdout_tx,
                    &stdout_op_tx,
                )
                .await;
            });

            let stderr_tx = event_tx.clone();
            let stderr_op_tx = op_tx.clone();
            let stderr_extractor = Arc::clone(&extractor);
            let stderr_task = tokio::spawn(async move {
                stream_output_parsed(
                    task_id,
                    OutputStream::Stderr,
                    stderr,
                    TerminalParser::new(),
                    &stderr_extractor,
                    stderr_engine,
                    &stderr_tx,
                    &stderr_op_tx,
                )
                .await;
            });

            // Wait for either the process to exit or cancellation.
            let final_status = tokio::select! {
                _ = cancellation_token.cancelled() => {
                    // Kill the process and abort the streaming tasks immediately.
                    // On Windows, killing the shell (`cmd.exe`) may leave grandchild
                    // processes alive that still hold the pipe handles open; aborting
                    // the streaming tasks avoids an indefinite pipe-drain hang.
                    let _ = child.kill().await;
                    stdout_task.abort();
                    stderr_task.abort();
                    TaskStatus::Cancelled
                }
                result = child.wait() => {
                    // Process exited naturally — drain all remaining output first.
                    let _ = tokio::join!(stdout_task, stderr_task);
                    match result {
                        Ok(exit) if exit.success() => TaskStatus::Success,
                        _ => TaskStatus::Error,
                    }
                }
            };

            let _ = event_tx.send(TaskEvent::Finished(task_id, final_status.clone()));
            let _ = op_tx.send(vec![Operation::TaskFinished {
                id: task_id,
                status: final_status,
            }]);
        });

        self.handles.insert(task_id, handle);
    }

    /// Abort the async task for `task_id`, if one is running.
    ///
    /// This does *not* update the registry — the caller is responsible for
    /// calling [`crate::task_registry::TaskRegistry::cancel`] before or after this.
    pub fn abort(&mut self, task_id: TaskId) {
        if let Some(handle) = self.handles.remove(&task_id) {
            handle.abort();
        }
    }
}

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

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Read bytes from `reader`, parse them through `parser`, and broadcast
/// completed [`StyledLine`]s as [`TaskEvent::ParsedOutput`] events until the
/// stream closes.
///
/// Each completed line is also run through `extractor` and (if provided)
/// `engine`; any resulting [`crate::issue_registry::NewIssue`]s are sent back
/// via `op_tx` as [`Operation::AddIssue`] operations so the registry is
/// updated on the main thread.
///
/// When the stream reaches EOF:
/// * The parser is flushed so that any partial line without a trailing `\n`
///   is also broadcast and extracted.
/// * The engine (if any) is flushed so that any pending multi-line block is
///   emitted as an issue.
///
/// On cancellation the parent task aborts this sub-task; no flush occurs,
/// which is intentional — partial output from a cancelled task is discarded.
#[allow(clippy::too_many_arguments)]
async fn stream_output_parsed(
    task_id: TaskId,
    stream: OutputStream,
    mut reader: impl tokio::io::AsyncRead + Unpin,
    mut parser: TerminalParser,
    extractor: &DiagnosticsExtractor,
    mut engine: Option<MatcherEngine>,
    tx: &broadcast::Sender<TaskEvent>,
    op_tx: &tokio::sync::mpsc::UnboundedSender<Vec<Operation>>,
) {
    let mut buf = vec![0u8; 4096];
    loop {
        match reader.read(&mut buf).await {
            Ok(0) | Err(_) => break,
            Ok(n) => {
                let lines = parser.push(&buf[..n]);
                broadcast_and_extract(task_id, &stream, &lines, extractor, engine.as_mut(), tx, op_tx);
            }
        }
    }
    // Flush any partial line that had no trailing newline.
    if let Some(line) = parser.flush() {
        broadcast_and_extract(task_id, &stream, &[line], extractor, engine.as_mut(), tx, op_tx);
    }
    // Flush any pending engine block (end-of-stream).
    if let Some(eng) = engine.as_mut() {
        let issues = eng.flush();
        send_issues(issues, op_tx);
    }
}

/// Broadcast `lines` as a `ParsedOutput` event, run the extractor and engine
/// on each line, sending any resulting `AddIssue` operations through `op_tx`.
fn broadcast_and_extract(
    task_id: TaskId,
    stream: &OutputStream,
    lines: &[StyledLine],
    extractor: &DiagnosticsExtractor,
    engine: Option<&mut MatcherEngine>,
    tx: &broadcast::Sender<TaskEvent>,
    op_tx: &tokio::sync::mpsc::UnboundedSender<Vec<Operation>>,
) {
    if lines.is_empty() {
        return;
    }

    // Broadcast the styled lines to all subscribers.
    let _ = tx.send(TaskEvent::ParsedOutput(ParsedOutputEvent {
        task_id,
        stream: stream.clone(),
        lines: lines.to_vec(),
    }));

    // Extract diagnostics and send AddIssue ops for each match.
    for line in lines {
        let issues = extractor.extract_from_line(line);
        if !issues.is_empty() {
            let ops: Vec<Operation> =
                issues.into_iter().map(|i| Operation::AddIssue { issue: i }).collect();
            let _ = op_tx.send(ops);
        }
    }

    // Run log matcher engine on each line.
    if let Some(eng) = engine {
        for line in lines {
            let issues = eng.process_line(&line.text);
            send_issues(issues, op_tx);
        }
    }
}

/// Send a batch of issues via `op_tx` (no-op if the vec is empty).
fn send_issues(
    issues: Vec<crate::issue_registry::NewIssue>,
    op_tx: &tokio::sync::mpsc::UnboundedSender<Vec<Operation>>,
) {
    if !issues.is_empty() {
        let ops: Vec<Operation> =
            issues.into_iter().map(|i| Operation::AddIssue { issue: i }).collect();
        let _ = op_tx.send(ops);
    }
}

/// Build a platform-appropriate [`Command`] that runs `command` through the
/// system shell.
fn build_command(command: &str) -> Command {
    #[cfg(unix)]
    {
        let mut cmd = Command::new("sh");
        cmd.arg("-c").arg(command);
        cmd
    }
    #[cfg(windows)]
    {
        let mut cmd = Command::new("cmd");
        cmd.arg("/C").arg(command);
        cmd
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::task_registry::{TaskKey, TaskQueueId, TaskTrigger, TaskRegistry};
    use tokio::sync::mpsc::unbounded_channel;
    use std::time::Duration;

    /// Safety timeout for all async tests — prevents CI hangs on Windows.
    const TEST_TIMEOUT: Duration = Duration::from_secs(15);

    fn dummy_key() -> TaskKey {
        TaskKey {
            queue: TaskQueueId("test".into()),
            target: "t".into(),
        }
    }

    fn key_n(n: u64) -> TaskKey {
        TaskKey {
            queue: TaskQueueId("test".into()),
            target: format!("t{n}"),
        }
    }

    fn schedule(reg: &mut TaskRegistry, key: TaskKey, command: &str) -> (TaskId, CancellationToken, String) {
        let cmd = command.to_string();
        let id = reg.schedule_task(key, TaskTrigger::Manual, cmd.clone());
        let token = reg.get(id).unwrap().cancellation_token.clone();
        (id, token, cmd)
    }

    // Platform-specific test commands.
    // On Windows we use `cmd /C <command>` via build_command.
    // echo exits quickly; sleep_long is a long-running placeholder that
    // we cancel before it ends.

    #[cfg(unix)]
    mod cmds {
        pub const ECHO_STDOUT: &str = "echo hello";
        pub const ECHO_STDERR: &str = "echo err >&2";
        pub const SLEEP_LONG: &str = "sleep 60";
        pub const NO_SUCH: &str = "/no/such/binary/xyz_oo_test";
    }
    #[cfg(windows)]
    mod cmds {
        // On Windows `cmd /C echo hello` writes to stdout.
        pub const ECHO_STDOUT: &str = "echo hello";
        // Redirect echo output to stderr via cmd.
        pub const ECHO_STDERR: &str = "echo err>&2";
        // `ping` with a very large count acts as a long-running sleep.
        pub const SLEEP_LONG: &str = "ping -n 120 127.0.0.1";
        // A command that cmd.exe cannot find → exits with non-zero.
        pub const NO_SUCH: &str = "no_such_binary_xyz_oo_test";
    }

    // 1. Task runs and emits stdout
    #[tokio::test]
    async fn task_emits_stdout() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let mut rx = exec.subscribe();
            let (op_tx, mut op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id, token, cmd) = schedule(&mut reg, dummy_key(), cmds::ECHO_STDOUT);
            exec.spawn(id, cmd, token, "task:test:t".into(), None, &op_tx);

            let mut got_output = false;
            loop {
                match rx.recv().await.unwrap() {
                    TaskEvent::ParsedOutput(ev)
                        if ev.task_id == id && ev.stream == OutputStream::Stdout =>
                    {
                        got_output = true;
                    }
                    TaskEvent::Finished(tid, status) if tid == id => {
                        assert_eq!(status, TaskStatus::Success);
                        break;
                    }
                    _ => {}
                }
            }
            assert!(got_output, "expected at least one stdout output event");

            // op_tx must deliver TaskStarted followed by TaskFinished(Success).
            let ops: Vec<Operation> = op_rx.try_recv().unwrap();
            assert!(ops.iter().any(|o| matches!(o, Operation::TaskStarted(i) if *i == id)));
            let ops2: Vec<Operation> = op_rx.recv().await.unwrap();
            assert!(ops2.iter().any(|o| matches!(
                o,
                Operation::TaskFinished { id: i, status: TaskStatus::Success } if *i == id
            )));
        })
        .await
        .expect("test timed out");
    }

    // 2. Task emits stderr
    #[tokio::test]
    async fn task_emits_stderr() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let mut rx = exec.subscribe();
            let (op_tx, _op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id, token, cmd) = schedule(&mut reg, dummy_key(), cmds::ECHO_STDERR);
            exec.spawn(id, cmd, token, "task:test:t".into(), None, &op_tx);

            let mut got_stderr = false;
            loop {
                match rx.recv().await.unwrap() {
                    TaskEvent::ParsedOutput(ev)
                        if ev.task_id == id && ev.stream == OutputStream::Stderr =>
                    {
                        got_stderr = true;
                    }
                    TaskEvent::Finished(tid, _) if tid == id => break,
                    _ => {}
                }
            }
            assert!(got_stderr, "expected at least one stderr output event");
        })
        .await
        .expect("test timed out");
    }

    // 3. Cancellation stops execution
    #[tokio::test]
    async fn cancellation_stops_task() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let mut rx = exec.subscribe();
            let (op_tx, mut op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id, token, cmd) = schedule(&mut reg, dummy_key(), cmds::SLEEP_LONG);
            exec.spawn(id, cmd, token.clone(), "task:test:t".into(), None, &op_tx);

            // Wait for Started, then cancel.
            loop {
                if let Ok(TaskEvent::Started(tid)) = rx.recv().await
                    && tid == id { break; }
            }
            token.cancel();

            // Must receive Finished(Cancelled) on the broadcast channel.
            loop {
                if let Ok(TaskEvent::Finished(tid, status)) = rx.recv().await
                    && tid == id {
                        assert_eq!(status, TaskStatus::Cancelled);
                        break;
                    }
            }

            // op_tx must also deliver the Cancelled finish.
            let mut found = false;
            while let Some(ops) = op_rx.recv().await {
                if ops.iter().any(|o| matches!(
                    o,
                    Operation::TaskFinished { id: i, status: TaskStatus::Cancelled }
                    if *i == id
                )) {
                    found = true;
                    break;
                }
            }
            assert!(found, "expected TaskFinished(Cancelled) in op_tx");
        })
        .await
        .expect("test timed out");
    }

    // 4. Multiple subscribers receive the same output
    #[tokio::test]
    async fn multiple_subscribers_receive_same_output() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let rx1 = exec.subscribe();
            let rx2 = exec.subscribe();
            let (op_tx, _op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id, token, cmd) = schedule(&mut reg, dummy_key(), cmds::ECHO_STDOUT);
            exec.spawn(id, cmd, token, "task:test:t".into(), None, &op_tx);

            async fn collect(mut rx: broadcast::Receiver<TaskEvent>, id: TaskId) -> usize {
                let mut chunks = 0usize;
                loop {
                    match rx.recv().await.unwrap() {
                        TaskEvent::ParsedOutput(_) => chunks += 1,
                        TaskEvent::Finished(tid, _) if tid == id => return chunks,
                        _ => {}
                    }
                }
            }

            let (c1, c2) = tokio::join!(collect(rx1, id), collect(rx2, id));
            assert_eq!(c1, c2, "both subscribers should see the same chunk count");
            assert!(c1 > 0, "expected at least one output chunk");
        })
        .await
        .expect("test timed out");
    }

    // 5. Process spawn/exit failure → Error status
    #[tokio::test]
    async fn spawn_failure_reports_error() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let mut rx = exec.subscribe();
            let (op_tx, _op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id, token, _) = schedule(&mut reg, dummy_key(), cmds::NO_SUCH);
            exec.spawn(id, cmds::NO_SUCH.to_string(), token, "task:test:t".into(), None, &op_tx);

            let final_status;
            loop {
                match rx.recv().await.unwrap() {
                    TaskEvent::Finished(tid, status) if tid == id => {
                        final_status = status;
                        break;
                    }
                    _ => {}
                }
            }
            // Whether spawn itself fails or the shell exits non-zero, both → Error.
            assert!(
                matches!(final_status, TaskStatus::Error),
                "expected Error, got {final_status:?}"
            );
        })
        .await
        .expect("test timed out");
    }

    // 6. Rapid cancel + restart: old process is killed, new one succeeds
    #[tokio::test]
    async fn rapid_cancel_restart() {
        tokio::time::timeout(TEST_TIMEOUT, async {
            let mut reg = TaskRegistry::new();
            let mut exec = TaskExecutor::new();
            let mut rx = exec.subscribe();
            let (op_tx, _op_rx) = unbounded_channel::<Vec<Operation>>();

            let (id1, token1, cmd1) = schedule(&mut reg, key_n(1), cmds::SLEEP_LONG);
            exec.spawn(id1, cmd1, token1.clone(), "task:test:t1".into(), None, &op_tx);

            // Wait for the first task to start, then cancel it.
            loop {
                if let Ok(TaskEvent::Started(tid)) = rx.recv().await
                    && tid == id1 { break; }
            }
            token1.cancel();

            // Schedule a second task (different key so no registry interaction needed).
            let (id2, token2, cmd2) = schedule(&mut reg, key_n(2), cmds::ECHO_STDOUT);
            exec.spawn(id2, cmd2, token2, "task:test:t2".into(), None, &op_tx);

            // Expect: id1 → Cancelled, id2 → Success (order may interleave).
            let mut saw_cancel = false;
            let mut saw_success = false;
            loop {
                match rx.recv().await.unwrap() {
                    TaskEvent::Finished(tid, TaskStatus::Cancelled) if tid == id1 => {
                        saw_cancel = true;
                    }
                    TaskEvent::Finished(tid, TaskStatus::Success) if tid == id2 => {
                        saw_success = true;
                    }
                    _ => {}
                }
                if saw_cancel && saw_success { break; }
            }
            assert!(saw_cancel, "first task should be Cancelled");
            assert!(saw_success, "second task should succeed");
        })
        .await
        .expect("test timed out");
    }
}