magi-code 0.102.1

//! Session-owned jobs. Live handles never enter history or survive an owner reset.
use super::{
    ToolResult, ToolResultDisplay,
    args::{BASH_STDERR_MAX_BYTES, BASH_STDOUT_MAX_BYTES, BashJobAction, BashJobArgs},
    process::{self, CleanupOutcome},
};
use crate::cancellation::AgentCancellation;
use serde_json::{Value, json};
use std::{
    collections::VecDeque,
    path::Path,
    process::Child,
    sync::{
        Arc, Condvar, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const MAX_RUNNING_JOBS: usize = 4;
const MAX_RETAINED_JOBS: usize = 16;
pub(super) const STREAM_CAPACITY: usize = 256 * 1024;
const POLL_INTERVAL: Duration = Duration::from_millis(20);

pub(crate) struct BashJobsCleanupGuard(BashJobs);

impl Drop for BashJobsCleanupGuard {
    fn drop(&mut self) {
        self.0.shutdown();
    }
}

#[derive(Clone, Default)]
pub(crate) struct BashJobs(Arc<Mutex<Registry>>);

impl std::fmt::Debug for BashJobs {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BashJobs").finish_non_exhaustive()
    }
}

#[derive(Default)]
struct Registry {
    jobs: VecDeque<Job>,
    closed: bool,
    children: Vec<BashJobs>,
}

struct Job {
    id: String,
    pid: u32,
    command: String,
    state: Arc<JobState>,
    worker: Option<JoinHandle<()>>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct BackgroundJob {
    pub(crate) id: String,
    pub(crate) pid: u32,
    pub(crate) command: String,
    pub(crate) uptime_seconds: u64,
    pub(crate) stopping: bool,
}

pub(super) struct JobState {
    data: Mutex<JobData>,
    changed: Condvar,
    stop: AtomicBool,
}

struct JobData {
    stdout: StreamBuffer,
    stderr: StreamBuffer,
    started: Instant,
    finished: Option<Duration>,
    exit_code: Option<i32>,
    cleanup_warning: Option<String>,
}

#[derive(Default)]
struct StreamBuffer {
    bytes: VecDeque<u8>,
    dropped: u64,
}

impl StreamBuffer {
    fn push(&mut self, bytes: &[u8]) {
        let overflow = (self.bytes.len() + bytes.len()).saturating_sub(STREAM_CAPACITY);
        self.bytes.drain(..overflow.min(self.bytes.len()));
        self.dropped = self.dropped.saturating_add(overflow as u64);
        self.bytes.extend(
            bytes
                .iter()
                .skip(bytes.len().saturating_sub(STREAM_CAPACITY)),
        );
    }

    fn take(&mut self, limit: usize, tail: bool) -> String {
        // Lossy UTF-8 expands each invalid byte to at most three bytes.
        let count = self.bytes.len().min(limit / 3);
        if tail {
            let discarded = self.bytes.len() - count;
            self.bytes.drain(..discarded);
            self.dropped = self.dropped.saturating_add(discarded as u64);
        }
        let bytes: Vec<_> = self.bytes.drain(..count).collect();
        String::from_utf8_lossy(&bytes).into_owned()
    }
}

impl JobState {
    pub(super) fn append(&self, stderr: bool, bytes: &[u8]) {
        let mut data = self.data.lock().unwrap_or_else(|e| e.into_inner());
        if stderr {
            &mut data.stderr
        } else {
            &mut data.stdout
        }
        .push(bytes);
        self.changed.notify_all();
    }
}

impl Job {
    fn snapshot(&self, output: bool, tail: bool) -> Value {
        let mut data = self.state.data.lock().unwrap_or_else(|e| e.into_inner());
        let mut result = json!({
            "job_id": self.id, "pid": self.pid,
            "status": if data.finished.is_some() { "exited" } else { "running" },
            "exit_code": data.exit_code,
            "uptime_seconds": data.finished.unwrap_or_else(|| data.started.elapsed()).as_secs_f64(),
            "stdout_bytes_dropped": data.stdout.dropped,
            "stderr_bytes_dropped": data.stderr.dropped,
            "cleanup_warning": data.cleanup_warning,
        });
        #[cfg(unix)]
        {
            result["pgid"] = json!(self.pid);
        }
        if output {
            result["stdout"] = json!(data.stdout.take(BASH_STDOUT_MAX_BYTES, tail));
            result["stderr"] = json!(data.stderr.take(BASH_STDERR_MAX_BYTES, tail));
            result["output_truncated"] = json!(data.stdout.dropped > 0 || data.stderr.dropped > 0);
            result["has_more_output"] =
                json!(!data.stdout.bytes.is_empty() || !data.stderr.bytes.is_empty());
            result["stdout_bytes_dropped"] = json!(data.stdout.dropped);
            result["stderr_bytes_dropped"] = json!(data.stderr.dropped);
        }
        result
    }

    fn finished(&self) -> bool {
        self.state
            .data
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .finished
            .is_some()
    }
}

impl Drop for Registry {
    fn drop(&mut self) {
        self.stop_all();
    }
}

impl Registry {
    fn stop_all(&mut self) {
        for job in &self.jobs {
            job.state.stop.store(true, Ordering::Release);
        }
        for child in &self.children {
            child.request_shutdown();
        }
        for child in &self.children {
            child.shutdown();
        }
        for job in &mut self.jobs {
            if let Some(worker) = job.worker.take() {
                let _ = worker.join();
            }
        }
    }
}

impl BashJobs {
    // UI callers must never wait behind process cleanup or output readers.
    pub(crate) fn running_snapshot(&self) -> Option<Vec<BackgroundJob>> {
        let registry = self.0.try_lock().ok()?;
        let mut jobs = Vec::new();
        for job in &registry.jobs {
            let data = job.state.data.try_lock().ok()?;
            if data.finished.is_none() {
                jobs.push(BackgroundJob {
                    id: job.id.clone(),
                    pid: job.pid,
                    command: job.command.clone(),
                    uptime_seconds: data.started.elapsed().as_secs(),
                    stopping: job.state.stop.load(Ordering::Acquire),
                });
            }
        }
        Some(jobs)
    }

    pub(crate) fn request_stop(&self, id: &str) -> anyhow::Result<()> {
        let registry = self
            .0
            .try_lock()
            .map_err(|_| anyhow::anyhow!("background jobs busy; try again"))?;
        let job = registry
            .jobs
            .iter()
            .find(|job| job.id == id)
            .ok_or_else(|| anyhow::anyhow!("background job no longer available"))?;
        // The existing supervisor owns termination and reaping, not the UI thread.
        job.state.stop.store(true, Ordering::Release);
        Ok(())
    }

    pub(super) fn child(&self) -> anyhow::Result<Self> {
        let mut registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
        anyhow::ensure!(!registry.closed, "background job runtime ended");
        registry
            .children
            .retain(|child| !child.0.lock().unwrap_or_else(|e| e.into_inner()).closed);
        anyhow::ensure!(
            registry.children.len() < 64,
            "background child registry limit (64) reached"
        );
        let child = Self::default();
        registry.children.push(child.clone());
        Ok(child)
    }

    pub(crate) fn request_shutdown(&self) {
        let mut registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
        registry.closed = true;
        for job in &registry.jobs {
            job.state.stop.store(true, Ordering::Release);
        }
        for child in &registry.children {
            child.request_shutdown();
        }
    }

    pub(super) fn cleanup_guard(&self) -> BashJobsCleanupGuard {
        BashJobsCleanupGuard(self.clone())
    }

    pub(crate) fn shutdown(&self) {
        let mut registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
        registry.closed = true;
        registry.stop_all();
    }

    // A prepared switch can be rejected as stale. Reap jobs without retiring its owner yet.
    pub(crate) fn stop_all(&self) {
        self.0.lock().unwrap_or_else(|e| e.into_inner()).stop_all();
    }

    pub(super) fn start(
        &self,
        command: &str,
        cwd: &Path,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        cancellation.check()?;
        let mut registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
        anyhow::ensure!(!registry.closed, "background job runtime ended");
        let running: Vec<_> = registry
            .jobs
            .iter()
            .filter(|job| !job.finished())
            .map(|job| job.id.as_str())
            .collect();
        anyhow::ensure!(
            running.len() < MAX_RUNNING_JOBS,
            "background job limit ({MAX_RUNNING_JOBS}) reached; running jobs: {}",
            running.join(", ")
        );
        if registry.jobs.len() >= MAX_RETAINED_JOBS {
            let index = registry
                .jobs
                .iter()
                .position(Job::finished)
                .expect("running jobs bounded");
            if let Some(mut job) = registry.jobs.remove(index)
                && let Some(worker) = job.worker.take()
            {
                let _ = worker.join();
            }
        }
        let child = crate::shell::runtime::spawn_agent_shell(command, cwd)?;
        let pid = child.id();
        let state = Arc::new(JobState {
            data: Mutex::new(JobData {
                stdout: StreamBuffer::default(),
                stderr: StreamBuffer::default(),
                started: Instant::now(),
                finished: None,
                exit_code: None,
                cleanup_warning: None,
            }),
            changed: Condvar::new(),
            stop: AtomicBool::new(false),
        });
        // Moving a cleanup guard into the closure also covers failure to spawn its thread.
        let child = OwnedChild {
            child,
            cleaned: false,
        };
        let worker_state = Arc::clone(&state);
        let worker = thread::Builder::new()
            .name("bash-job".into())
            .spawn(move || run_job(child, worker_state))?;
        let id = uuid::Uuid::new_v4().to_string();
        registry.jobs.push_back(Job {
            id: id.clone(),
            pid,
            command: crate::output::sanitize_display_text(command)
                .split_whitespace()
                .collect::<Vec<_>>()
                .join(" ")
                .chars()
                .take(512)
                .collect(),
            state,
            worker: Some(worker),
        });
        drop(registry);
        // Startup is a bounded observation, never the lifetime of the command.
        self.wait(&id, Duration::from_millis(200), cancellation, false)?;
        self.result(&id, true, false, "bash")
    }

    pub(super) fn call(
        &self,
        args: BashJobArgs,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        args.validate()?;
        if args.action == BashJobAction::List {
            let registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
            return Ok(tool_result(
                "bash_job",
                json!({"jobs": registry.jobs.iter().map(|job| job.snapshot(false, false)).collect::<Vec<_>>()}),
            ));
        }
        let id = args.job_id.as_deref().expect("validated job id");
        match args.action {
            BashJobAction::Output => self.wait(
                id,
                Duration::from_secs(args.wait_seconds.unwrap_or(0)),
                cancellation,
                true,
            )?,
            BashJobAction::Stop => {
                self.state(id)?.stop.store(true, Ordering::Release);
                // Stop is not undone by cancellation; the worker always finishes cleanup.
                self.wait(id, Duration::from_secs(5), cancellation, false)?;
            }
            BashJobAction::Status => {}
            BashJobAction::List => unreachable!(),
        }
        self.result(
            id,
            args.action != BashJobAction::Status,
            args.action == BashJobAction::Stop,
            "bash_job",
        )
    }

    fn state(&self, id: &str) -> anyhow::Result<Arc<JobState>> {
        self.0
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .jobs
            .iter()
            .find(|job| job.id == id)
            .map(|job| Arc::clone(&job.state))
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "job not found (runtime ended, job expired, or different owner): {id}"
                )
            })
    }

    fn wait(
        &self,
        id: &str,
        timeout: Duration,
        cancellation: &AgentCancellation,
        on_output: bool,
    ) -> anyhow::Result<()> {
        let state = self.state(id)?;
        let deadline = Instant::now() + timeout;
        let mut data = state.data.lock().unwrap_or_else(|e| e.into_inner());
        while data.finished.is_none() && Instant::now() < deadline {
            cancellation.check()?;
            if on_output && (!data.stdout.bytes.is_empty() || !data.stderr.bytes.is_empty()) {
                break;
            }
            data = state
                .changed
                .wait_timeout(
                    data,
                    POLL_INTERVAL.min(deadline.saturating_duration_since(Instant::now())),
                )
                .unwrap_or_else(|e| e.into_inner())
                .0;
        }
        cancellation.check()
    }

    fn result(&self, id: &str, output: bool, tail: bool, name: &str) -> anyhow::Result<ToolResult> {
        let registry = self.0.lock().unwrap_or_else(|e| e.into_inner());
        let job = registry
            .jobs
            .iter()
            .find(|job| job.id == id)
            .ok_or_else(|| anyhow::anyhow!("job not found (runtime ended): {id}"))?;
        Ok(tool_result(name, job.snapshot(output, tail)))
    }
}

struct OwnedChild {
    child: Child,
    cleaned: bool,
}

impl OwnedChild {
    fn terminate(&mut self) -> anyhow::Result<CleanupOutcome> {
        let result = process::terminate_child_tree_and_wait(&mut self.child);
        self.cleaned = result.is_ok();
        result
    }
}

impl Drop for OwnedChild {
    fn drop(&mut self) {
        if !self.cleaned {
            let _ = self.terminate();
        }
    }
}

fn run_job(mut child: OwnedChild, state: Arc<JobState>) {
    let outcome = (|| -> anyhow::Result<CleanupOutcome> {
        let stdout = child
            .child
            .stdout
            .take()
            .ok_or_else(|| anyhow::anyhow!("missing stdout"))?;
        let stderr = child
            .child
            .stderr
            .take()
            .ok_or_else(|| anyhow::anyhow!("missing stderr"))?;
        let stdout = process::spawn_job_pipe_reader(stdout, Arc::clone(&state), false)?;
        let stderr = match process::spawn_job_pipe_reader(stderr, Arc::clone(&state), true) {
            Ok(reader) => reader,
            Err(error) => {
                let _ = child.terminate();
                let _ = process::recv_pipe_reader_with_timeout(stdout);
                return Err(error);
            }
        };
        let monitor_error = loop {
            match child.child.try_wait() {
                Ok(Some(_)) => break None,
                Ok(None) if !state.stop.load(Ordering::Acquire) => thread::sleep(POLL_INTERVAL),
                Ok(None) => break None,
                Err(error) => break Some(error.to_string()),
            }
        };
        let cleanup = child.terminate();
        let stdout = process::recv_pipe_reader_with_timeout(stdout);
        let stderr = process::recv_pipe_reader_with_timeout(stderr);
        let mut cleanup = cleanup?;
        if stdout.timed_out
            || stderr.timed_out
            || stdout.stopped_before_eof
            || stderr.stopped_before_eof
        {
            cleanup.cleanup_warning = process::combine_cleanup_warning(
                cleanup.cleanup_warning.as_deref(),
                &["job pipe output incomplete after cleanup"],
            );
        }
        if let Some(error) = monitor_error {
            cleanup.cleanup_warning = Some(format!(
                "{}; {error}",
                cleanup.cleanup_warning.unwrap_or_default()
            ));
        }
        Ok(cleanup)
    })();
    // Cleanup guard runs before publishing exit, including setup/monitor errors.
    drop(child);
    let mut data = state.data.lock().unwrap_or_else(|e| e.into_inner());
    match outcome {
        Ok(cleanup) => {
            data.exit_code = cleanup.status.and_then(|status| status.code());
            data.cleanup_warning = cleanup.cleanup_warning;
        }
        Err(error) => data.cleanup_warning = Some(error.to_string()),
    }
    data.finished = Some(data.started.elapsed());
    state.changed.notify_all();
}

fn tool_result(name: &str, mut metadata: Value) -> ToolResult {
    metadata["bash_file_tracking"] = json!({"status": "omitted", "source": "background_job"});
    ToolResult {
        tool_name: name.into(),
        success: metadata.get("cleanup_warning").is_none_or(Value::is_null),
        content: metadata.to_string(),
        metadata,
        display: ToolResultDisplay::default(),
    }
}