IdeaSDK/crates/infrastructure/src/background_task/runner.rs

//! [`CommandBackgroundRunner`] — the concrete command-backed background task
//! runner (B8 sub-task 2).
//!
//! # Role
//!
//! This is the `impl` of the frozen [`BackgroundTaskRunner`] port that the
//! composition root wires under the [`BackgroundCompletionSink`]. It spawns a
//! command-backed task over a **composed** [`PtyPort`] (resolved by the project's
//! `RemoteHost`, so local/SSH/WSL are transparent — Liskov) and, for each task,
//! emits **exactly one** [`BackgroundTaskCompletion`] on the stream returned by
//! [`subscribe_completions`](BackgroundTaskRunner::subscribe_completions).
//!
//! It writes **neither** the store **nor** the inbox: the sink is the single
//! writer that persists the terminal state before signalling delivery
//! (persist-before-signal). The runner's only outward effect is the completion
//! event.
//!
//! # Lifecycle of one task
//!
//! `spawn` opens the PTY and starts a detached worker. The worker waits on the
//! PTY process status via [`PtyPort::wait`](domain::ports::PtyPort::wait), a
//! cancel notification, or the deadline, whichever comes first. It then snapshots
//! the bounded output tail from the PTY scrollback, cleans up via
//! [`PtyPort::kill`](domain::ports::PtyPort::kill), maps that to a terminal
//! [`BackgroundTaskResult`], and sends the completion.
//!
//! # Output tail & the live tee
//!
//! A PTY **merges** stdout and stderr onto one stream, so the completion carries
//! the merged output in `stdout_tail` and leaves `stderr_tail` empty. Completion
//! detection is deliberately decoupled from output subscriptions; the tail is read
//! from PTY scrollback, while process lifecycle comes from `wait`.

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex};

use async_trait::async_trait;

use application::BackgroundCommandArchive;
use domain::ports::{
    BackgroundCompletionStream, BackgroundTaskCompletion, BackgroundTaskHandle,
    BackgroundTaskPortError, BackgroundTaskRunner, BackgroundTaskSpec, Clock, PtyHandle, PtyPort,
    SpawnSpec,
};
use domain::terminal::PtySize;
use domain::{BackgroundTaskResult, TaskId};

use super::tail::{bounded_tail, tail_cap_bytes};

/// Default PTY geometry for a headless background command (no visible cell).
const DEFAULT_SIZE: PtySize = PtySize { rows: 24, cols: 80 };

/// Why the worker's terminal wait resolved.
enum WaitOutcome {
    /// The process exited on its own.
    Exited,
    /// A [`cancel`](BackgroundTaskRunner::cancel) was requested.
    Cancelled,
    /// The deadline elapsed before exit.
    Expired,
}

/// Live control handles for one running task (kept in the registry).
#[derive(Clone)]
struct TaskControl {
    pty_handle: PtyHandle,
    cancel: Arc<tokio::sync::Notify>,
    cancelled: Arc<AtomicBool>,
}

/// Command-backed [`BackgroundTaskRunner`] over a composed [`PtyPort`].
pub struct CommandBackgroundRunner {
    pty: Arc<dyn PtyPort>,
    clock: Arc<dyn Clock>,
    completion_tx: Sender<BackgroundTaskCompletion>,
    completion_rx: Mutex<Option<Receiver<BackgroundTaskCompletion>>>,
    running: Arc<Mutex<HashMap<TaskId, TaskControl>>>,
    /// Retained command invocations, so [`BackgroundCommandArchive::spec_for`]
    /// can recover a task's command for retry.
    specs: Arc<Mutex<HashMap<TaskId, SpawnSpec>>>,
}

impl CommandBackgroundRunner {
    /// Builds a runner over a composed PTY port and clock.
    #[must_use]
    pub fn new(pty: Arc<dyn PtyPort>, clock: Arc<dyn Clock>) -> Self {
        let (completion_tx, completion_rx) = mpsc::channel();
        Self {
            pty,
            clock,
            completion_tx,
            completion_rx: Mutex::new(Some(completion_rx)),
            running: Arc::new(Mutex::new(HashMap::new())),
            specs: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    fn now_ms(&self) -> u64 {
        u64::try_from(self.clock.now_millis().max(0)).unwrap_or(0)
    }

    /// Detached worker driving one command to its single completion.
    #[allow(clippy::too_many_arguments)]
    async fn run_to_completion(
        pty: Arc<dyn PtyPort>,
        clock: Arc<dyn Clock>,
        completion_tx: Sender<BackgroundTaskCompletion>,
        running: Arc<Mutex<HashMap<TaskId, TaskControl>>>,
        task_id: TaskId,
        control: TaskControl,
        deadline_ms: Option<u64>,
    ) {
        let outcome = Self::await_terminal(&pty, &control, deadline_ms, &clock).await;
        // A cancel that raced a natural EOF still wins: the user asked to cancel.
        let outcome = if control.cancelled.load(Ordering::SeqCst) {
            WaitOutcome::Cancelled
        } else {
            outcome
        };

        // Snapshot the merged output tail *before* kill tears the session down.
        let tail = pty
            .scrollback(&control.pty_handle)
            .ok()
            .map(|bytes| bounded_tail(&bytes, tail_cap_bytes()))
            .filter(|s| !s.is_empty());

        // Reap the real exit status (works for a natural exit and a forced kill).
        let exit_code = pty
            .kill(&control.pty_handle)
            .await
            .ok()
            .and_then(|status| status.code);

        let finished_at_ms = u64::try_from(clock.now_millis().max(0)).unwrap_or(0);
        let result = match outcome {
            WaitOutcome::Cancelled => BackgroundTaskResult::Cancelled {
                finished_at_ms,
                reason: "cancelled by request".to_owned(),
            },
            WaitOutcome::Expired => BackgroundTaskResult::Expired {
                finished_at_ms,
                reason: "deadline elapsed before completion".to_owned(),
            },
            WaitOutcome::Exited => {
                if exit_code == Some(0) {
                    BackgroundTaskResult::Success {
                        finished_at_ms,
                        exit_code,
                        summary: "command completed successfully".to_owned(),
                        stdout_tail: tail,
                        stderr_tail: None,
                    }
                } else {
                    BackgroundTaskResult::Failure {
                        finished_at_ms,
                        exit_code,
                        error: match exit_code {
                            Some(code) => format!("command exited with code {code}"),
                            None => "command terminated without an exit code".to_owned(),
                        },
                        stdout_tail: tail,
                        stderr_tail: None,
                    }
                }
            }
        };

        running
            .lock()
            .expect("runner registry poisoned")
            .remove(&task_id);
        let _ = completion_tx.send(BackgroundTaskCompletion { task_id, result });
    }

    /// Resolves when the process exits, is cancelled, or hits its deadline.
    async fn await_terminal(
        pty: &Arc<dyn PtyPort>,
        control: &TaskControl,
        deadline_ms: Option<u64>,
        clock: &Arc<dyn Clock>,
    ) -> WaitOutcome {
        let process_exit = pty.wait(&control.pty_handle);
        tokio::pin!(process_exit);

        let sleep_ms = deadline_ms.map(|deadline| {
            let now = u64::try_from(clock.now_millis().max(0)).unwrap_or(0);
            deadline.saturating_sub(now)
        });

        match sleep_ms {
            Some(ms) => tokio::select! {
                result = &mut process_exit => {
                    let _ = result;
                    WaitOutcome::Exited
                },
                () = control.cancel.notified() => WaitOutcome::Cancelled,
                () = tokio::time::sleep(std::time::Duration::from_millis(ms)) => WaitOutcome::Expired,
            },
            None => tokio::select! {
                result = &mut process_exit => {
                    let _ = result;
                    WaitOutcome::Exited
                },
                () = control.cancel.notified() => WaitOutcome::Cancelled,
            },
        }
    }
}

#[async_trait]
impl BackgroundTaskRunner for CommandBackgroundRunner {
    async fn spawn(
        &self,
        spec: BackgroundTaskSpec,
    ) -> Result<BackgroundTaskHandle, BackgroundTaskPortError> {
        let command = spec.command.clone().ok_or_else(|| {
            BackgroundTaskPortError::Runner(
                "command-backed runner requires a command spec".to_owned(),
            )
        })?;

        let pty_handle = self
            .pty
            .spawn(command.clone(), DEFAULT_SIZE)
            .await
            .map_err(|e| BackgroundTaskPortError::Runner(format!("pty spawn failed: {e}")))?;

        self.specs
            .lock()
            .expect("runner spec registry poisoned")
            .insert(spec.task_id, command);

        let control = TaskControl {
            pty_handle,
            cancel: Arc::new(tokio::sync::Notify::new()),
            cancelled: Arc::new(AtomicBool::new(false)),
        };
        self.running
            .lock()
            .expect("runner registry poisoned")
            .insert(spec.task_id, control.clone());

        tokio::spawn(Self::run_to_completion(
            Arc::clone(&self.pty),
            Arc::clone(&self.clock),
            self.completion_tx.clone(),
            Arc::clone(&self.running),
            spec.task_id,
            control,
            spec.deadline_ms,
        ));

        // Touch `now_ms` only to keep the clock wired for future scheduling; the
        // authoritative timestamps come from the worker at completion.
        let _ = self.now_ms();
        Ok(BackgroundTaskHandle {
            task_id: spec.task_id,
        })
    }

    async fn cancel(&self, task_id: TaskId) -> Result<(), BackgroundTaskPortError> {
        let control = self
            .running
            .lock()
            .expect("runner registry poisoned")
            .get(&task_id)
            .cloned();
        // Unknown task ⇒ already terminal or never ours: cancellation is idempotent.
        if let Some(control) = control {
            control.cancelled.store(true, Ordering::SeqCst);
            control.cancel.notify_one();
        }
        Ok(())
    }

    fn subscribe_completions(&self) -> BackgroundCompletionStream {
        let rx = self
            .completion_rx
            .lock()
            .expect("runner completion receiver poisoned")
            .take()
            .expect("subscribe_completions must be called exactly once");
        Box::new(rx.into_iter())
    }
}

impl BackgroundCommandArchive for CommandBackgroundRunner {
    fn spec_for(&self, task_id: TaskId) -> Option<SpawnSpec> {
        self.specs
            .lock()
            .expect("runner spec registry poisoned")
            .get(&task_id)
            .cloned()
    }
}