opencrabs 0.3.79

//! Background task manager (#722).
//!
//! Runs a genuinely long command detached (so it doesn't churn the bash 600s
//! cap) and, on completion, enqueues a synthetic `QueuedUserMessage` into the
//! originating session via the surface enqueue callback. The tool loop drains
//! that at the next iteration boundary — injected mid-turn if the agent is still
//! working, or starting a fresh turn if it went idle.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Mutex;

use uuid::Uuid;

use super::types::{MessageEnqueueCallback, QueuedUserMessage};

/// Where a session's background-task completion must be delivered, keyed by
/// session rather than by whichever surface happened to run the command.
///
/// Every surface builds its own `BackgroundTaskManager` from its own enqueue
/// callback, so the completion used to follow the *executing* service. A
/// channel-bound session driven from the TUI therefore reported back to the
/// TUI, and the channel that started the work never heard the answer (#940).
/// A channel registers its session here when it binds one; the manager
/// consults this first and only falls back to its own callback when nothing
/// claims the session (a genuinely TUI-local or CLI-local session).
static SESSION_ROUTES: Mutex<Option<HashMap<Uuid, MessageEnqueueCallback>>> = Mutex::new(None);

/// Bind `session_id`'s background-task completions to `enqueue`.
///
/// Idempotent: re-binding the same session replaces the route, which is what
/// a reconnect or a bot restart needs.
pub fn register_session_route(session_id: Uuid, enqueue: MessageEnqueueCallback) {
    match SESSION_ROUTES.lock() {
        Ok(mut guard) => {
            guard
                .get_or_insert_with(HashMap::new)
                .insert(session_id, enqueue);
        }
        Err(e) => {
            // Worth saying out loud: without the route this session's next
            // background completion silently goes to the wrong surface.
            tracing::error!(
                target: "background_task",
                "Could not register resume route for session {session_id}: {e}"
            );
        }
    }
}

/// Who should receive `session_id`'s completion: the surface that claimed the
/// session, falling back to `executing` when nothing did.
///
/// The whole fix in one line — pick by session, never by who ran the command —
/// so it is a pure function and directly testable.
pub fn resolve_route(
    session_id: Uuid,
    executing: &MessageEnqueueCallback,
) -> MessageEnqueueCallback {
    session_route(session_id).unwrap_or_else(|| executing.clone())
}

/// The surface that owns `session_id`'s completions, if one claimed it.
fn session_route(session_id: Uuid) -> Option<MessageEnqueueCallback> {
    match SESSION_ROUTES.lock() {
        Ok(guard) => guard.as_ref()?.get(&session_id).cloned(),
        Err(e) => {
            tracing::error!(
                target: "background_task",
                "Could not read resume route for session {session_id}: {e}"
            );
            None
        }
    }
}

/// Result of a finished background command.
#[derive(Debug, Clone)]
pub struct CmdResult {
    pub success: bool,
    pub code: i32,
    pub output: String,
}

/// One in-flight background command.
#[derive(Debug, Clone)]
pub struct RunningTask {
    /// Short label for the command, e.g. `cargo test`.
    pub label: String,
    /// When it was spawned, for the elapsed time a surface displays.
    pub started: std::time::Instant,
}

/// Manages background commands and resumes their sessions on completion.
pub struct BackgroundTaskManager {
    enqueue: MessageEnqueueCallback,
    /// In-flight background tasks per session.
    ///
    /// Holds the label and start time, not just a count, because a surface has
    /// to be able to say WHAT is running and for how long. A detached task
    /// takes the turn idle, so without this the TUI has nothing at all to draw
    /// while a long build runs and the wait looks like a hang (#762).
    running: Mutex<HashMap<Uuid, Vec<RunningTask>>>,
}

impl BackgroundTaskManager {
    pub fn new(enqueue: MessageEnqueueCallback) -> Self {
        Self {
            enqueue,
            running: Mutex::new(HashMap::new()),
        }
    }

    /// How many background tasks are currently running for `session_id`.
    pub fn running_for(&self, session_id: Uuid) -> usize {
        self.running
            .lock()
            .map(|m| m.get(&session_id).map(Vec::len).unwrap_or(0))
            .unwrap_or(0)
    }

    /// What is running for `session_id`, oldest first, for surfaces that show
    /// progress. Returns owned data so the caller never holds the lock.
    pub fn running_tasks(&self, session_id: Uuid) -> Vec<RunningTask> {
        self.running
            .lock()
            .map(|m| m.get(&session_id).cloned().unwrap_or_default())
            .unwrap_or_default()
    }

    fn mark_started(&self, session_id: Uuid, label: &str) {
        if let Ok(mut m) = self.running.lock() {
            m.entry(session_id).or_default().push(RunningTask {
                label: label.to_string(),
                started: std::time::Instant::now(),
            });
        }
    }

    fn mark_finished(&self, session_id: Uuid, label: &str) {
        if let Ok(mut m) = self.running.lock()
            && let Some(tasks) = m.get_mut(&session_id)
        {
            // Remove the OLDEST entry with this label: two `cargo test` runs are
            // indistinguishable here, and dropping the oldest keeps the elapsed
            // time shown for the survivor honest.
            if let Some(pos) = tasks.iter().position(|t| t.label == label) {
                tasks.remove(pos);
            }
            if tasks.is_empty() {
                m.remove(&session_id);
            }
        }
    }

    /// Spawn `command` (via `sh -c`) in `cwd`, detached; on completion enqueue a
    /// system message into `session_id` summarizing the result. Returns
    /// immediately — the caller's turn is free to end.
    pub fn spawn_command(
        self: std::sync::Arc<Self>,
        session_id: Uuid,
        cwd: PathBuf,
        label: String,
        command: String,
    ) {
        self.mark_started(session_id, &label);
        let this = std::sync::Arc::clone(&self);
        let task_id = Uuid::new_v4();
        tokio::spawn(async move {
            // Persist BEFORE running: a restart mid-command must find a row to
            // report as interrupted, otherwise the session waits forever on a
            // resume that can no longer come (#763).
            if let Some(repo) = task_repo() {
                let cwd_str = cwd.to_string_lossy().to_string();
                if let Err(e) = repo
                    .record(task_id, session_id, &label, &command, &cwd_str)
                    .await
                {
                    // Not fatal: the command still runs and still resumes the
                    // session in this process. Only restart accounting is lost.
                    tracing::error!(
                        target: "background_task",
                        "Failed to persist background task '{label}': {e:#}"
                    );
                }
            }
            let result = run_detached(&command, &cwd).await;
            tracing::info!(
                target: "background_task",
                "Background task '{label}' for session {session_id} finished (success={})",
                result.success
            );
            let msg = completion_message(&label, &command, &result);
            if let Some(repo) = task_repo()
                && let Err(e) = repo.clear(task_id).await
            {
                // A stale row makes the NEXT startup report a phantom
                // interruption, so this must be visible even though the
                // command itself succeeded.
                tracing::error!(
                    target: "background_task",
                    "Failed to clear background task '{label}' after completion: {e:#}"
                );
            }
            // Clear the indicator BEFORE delivering, not after. The task is
            // over the moment the process exits, but mark_finished sat behind
            // the enqueue callback, so the "running" badge outlived the work by
            // however long delivery took — on a killed task the user saw the
            // agent confirm it had stopped while the input border still showed
            // it running.
            //
            // Only touches the in-memory map, so moving it earlier cannot
            // affect what gets delivered.
            this.mark_finished(session_id, &label);
            // Deliver to the surface that OWNS the session, not to whichever
            // one executed the command. A channel session driven from the TUI
            // runs on the TUI's service, so `this.enqueue` would answer into
            // the TUI and leave the channel that asked for the work waiting on
            // a reply that never comes (#940).
            resolve_route(session_id, &this.enqueue)(session_id, msg);
        });
    }
}

/// The background-task repository, when a pool exists.
///
/// Resolved per call through the global pool rather than threaded through the
/// manager, because `spawn_command` is reached from the bash tool which has no
/// pool in its context. `None` before the DB is initialized (early startup,
/// tests), which simply means restart accounting is skipped.
fn task_repo() -> Option<crate::db::BackgroundTaskRepository> {
    crate::db::global_pool().map(|p| crate::db::BackgroundTaskRepository::new(p.clone()))
}

/// Account for background tasks that were running when a previous process
/// died, then clear them.
///
/// Every surviving row belonged to a process that no longer exists, so its
/// child is gone too: there is nothing to reattach to and no result coming.
/// Each one is reported into its session as an interruption so the agent can
/// decide whether to re-run it, rather than waiting forever on a resume that
/// can never arrive (#763). Returns how many were reported.
pub async fn report_interrupted(enqueue: &MessageEnqueueCallback) -> usize {
    let Some(repo) = task_repo() else {
        return 0;
    };
    let rows = match repo.all().await {
        Ok(rows) => rows,
        Err(e) => {
            tracing::error!(target: "background_task", "Failed to read background tasks: {e:#}");
            return 0;
        }
    };
    if rows.is_empty() {
        return 0;
    }
    let count = rows.len();
    for row in rows {
        tracing::warn!(
            target: "background_task",
            "Background task '{}' for session {} was interrupted by a restart",
            row.label,
            row.session_id
        );
        enqueue(row.session_id, interrupted_message(&row));
    }
    if let Err(e) = repo.clear_all().await {
        // Leaving rows behind would re-report the same interruption on the
        // next start, so this is loud even though the reports already landed.
        tracing::error!(
            target: "background_task",
            "Failed to clear background tasks after reporting: {e:#}"
        );
    }
    count
}

/// What the agent is told about a command a restart killed. Deliberately
/// states that it did NOT finish and hands the decision back, rather than
/// re-running something expensive on the agent's behalf.
fn interrupted_message(row: &crate::db::BackgroundTaskRow) -> QueuedUserMessage {
    let context_text = format!(
        "[BACKGROUND TASK INTERRUPTED] `{}` was still running when OpenCrabs restarted, so it \
         was killed and produced no result. The command was:\n\n```\n{}\n```\n\nIt did NOT \
         complete. Decide whether to run it again based on what you were doing; do not assume \
         it passed or failed.",
        row.label, row.command
    );
    QueuedUserMessage {
        context_text,
        display_text: format!("⚠️ Background task interrupted by restart: {}", row.label),
    }
}

/// Run `command` through `sh -c` in `cwd`, capturing merged stdout+stderr.
async fn run_detached(command: &str, cwd: &std::path::Path) -> CmdResult {
    use tokio::process::Command;
    let output = Command::new("sh")
        .arg("-c")
        .arg(command)
        .current_dir(cwd)
        .output()
        .await;
    match output {
        Ok(out) => {
            let mut combined = String::from_utf8_lossy(&out.stdout).into_owned();
            let err = String::from_utf8_lossy(&out.stderr);
            if !err.trim().is_empty() {
                if !combined.is_empty() {
                    combined.push('\n');
                }
                combined.push_str(&err);
            }
            CmdResult {
                success: out.status.success(),
                code: out.status.code().unwrap_or(-1),
                output: combined,
            }
        }
        Err(e) => CmdResult {
            success: false,
            code: -1,
            output: format!("failed to launch: {e}"),
        },
    }
}

/// Command substrings that mark a genuinely long-running task (#722). When bash
/// sees one AND a background manager is wired, it runs the command detached and
/// resumes the session on completion instead of blocking toward the 600s cap.
/// Matched as a substring so `cd x && cargo test` still counts.
const KNOWN_LONG_MARKERS: &[&str] = &[
    "cargo test",
    "cargo build",
    "cargo run",
    "cargo clippy",
    "npm test",
    "npm run build",
    "pnpm test",
    "pnpm build",
    "yarn test",
    "yarn build",
    "npx remotion render",
    "remotion render",
    "gh run watch",
    "gh pr checks --watch",
];

/// Is `command` a known long-running task that should run in the background?
pub(crate) fn is_known_long(command: &str) -> bool {
    let lower = command.to_lowercase();
    KNOWN_LONG_MARKERS.iter().any(|m| lower.contains(m))
}

/// A short human label for a command (first meaningful token sequence), for the
/// "running in the background" acknowledgement and the completion tag.
pub(crate) fn short_label(command: &str) -> String {
    let after_cd = crate::utils::command_label::command_label(command);
    let label: String = after_cd.chars().take(60).collect();
    if after_cd.chars().count() > 60 {
        format!("{label}…")
    } else {
        label
    }
}

/// Keep only the last `n` lines of `text`.
pub(crate) fn tail_lines(text: &str, n: usize) -> String {
    let lines: Vec<&str> = text.lines().collect();
    let start = lines.len().saturating_sub(n);
    lines[start..].join("\n")
}

/// Build the resume message from a finished background command (#722). Pure so
/// the framing is unit-testable without spawning anything.
pub(crate) fn completion_message(
    label: &str,
    command: &str,
    result: &CmdResult,
) -> QueuedUserMessage {
    let status = if result.success {
        "exit 0 (success)".to_string()
    } else {
        format!("exit {} (failure)", result.code)
    };
    let tail = tail_lines(&result.output, 50);
    let context = format!(
        "[System: the background task you started has finished.\n\
         Task: {label}\n\
         Command: {command}\n\
         Status: {status}\n\
         Output (last 50 lines):\n{tail}\n\n\
         Report the result to the user and continue anything that was waiting on it. \
         Do not re-run the command — this IS its result.]"
    );
    let display = format!(
        "🔧 background task {}: {label}",
        if result.success { "finished" } else { "failed" }
    );
    QueuedUserMessage::system(context, display)
}