atman_runtime/tools/

bash_bg.rs

use std::collections::HashMap;
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use command_group::{AsyncCommandGroup, AsyncGroupChild};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

use crate::error::RuntimeError;
use crate::task_registry::{TaskKind, TaskRegistry, TaskStatus};
use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
use crate::value::Value;

const DEFAULT_SPAWN_TIMEOUT_MS: u64 = 1_800_000;
const MAX_SPAWN_TIMEOUT_MS: u64 = 86_400_000;
const DEFAULT_MAX_OUTPUT_BYTES: u64 = 10_485_760;
const RING_BUFFER_BYTES: usize = 65_536;
const IO_DRAIN_TIMEOUT: Duration = Duration::from_secs(2);
const DEFAULT_OUTPUT_LIMIT: usize = 32_000;

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct BgHandle {
    session_id: String,
    local_id: u64,
}

impl BgHandle {
    #[allow(clippy::inherent_to_string)]
    pub fn to_string(&self) -> String {
        format!("bg_{}_{}", self.session_id, self.local_id)
    }

    pub fn parse(s: &str) -> Option<Self> {
        let rest = s.strip_prefix("bg_")?;
        let idx = rest.rfind('_')?;
        let session_id = rest[..idx].to_string();
        let local_id = rest[idx + 1..].parse().ok()?;
        Some(Self {
            session_id,
            local_id,
        })
    }
}

#[derive(Debug, Clone)]
pub enum BgStatus {
    Running {
        pid: u32,
        started_at: i64,
    },
    Exited {
        exit_code: i32,
        started_at: i64,
        ended_at: i64,
    },
    TimedOut {
        started_at: i64,
        ended_at: i64,
    },
    Killed {
        started_at: i64,
        ended_at: i64,
    },
    Failed {
        error: String,
        started_at: i64,
        ended_at: i64,
    },
}

impl BgStatus {
    fn kind(&self) -> &'static str {
        match self {
            Self::Running { .. } => "running",
            Self::Exited { .. } => "exited",
            Self::TimedOut { .. } => "timed_out",
            Self::Killed { .. } => "killed",
            Self::Failed { .. } => "failed",
        }
    }

    fn exit_code(&self) -> Option<i32> {
        match self {
            Self::Exited { exit_code, .. } => Some(*exit_code),
            _ => None,
        }
    }

    fn started_at(&self) -> i64 {
        match self {
            Self::Running { started_at, .. }
            | Self::Exited { started_at, .. }
            | Self::TimedOut { started_at, .. }
            | Self::Killed { started_at, .. }
            | Self::Failed { started_at, .. } => *started_at,
        }
    }

    fn ended_at(&self) -> Option<i64> {
        match self {
            Self::Exited { ended_at, .. }
            | Self::TimedOut { ended_at, .. }
            | Self::Killed { ended_at, .. }
            | Self::Failed { ended_at, .. } => Some(*ended_at),
            _ => None,
        }
    }

    fn is_finished(&self) -> bool {
        !matches!(self, Self::Running { .. })
    }
}

#[derive(Debug, Default)]
pub struct BgOutput {
    pub combined: Vec<u8>,
    pub total_bytes: u64,
    pub truncated: bool,
}

impl BgOutput {
    fn push(&mut self, kind: StreamKind, data: &[u8], max: u64) {
        let prefix: &[u8] = match kind {
            StreamKind::Stdout => b"[out] ",
            StreamKind::Stderr => b"[err] ",
        };
        let mut new_total = self.total_bytes + data.len() as u64;
        let mut to_write = data;
        if new_total > max {
            let allowed = max.saturating_sub(self.total_bytes) as usize;
            to_write = &data[..allowed.min(data.len())];
            new_total = max;
            self.truncated = true;
        }
        if !to_write.is_empty() {
            self.combined.extend_from_slice(prefix);
            self.combined.extend_from_slice(to_write);
            if !to_write.ends_with(b"\n") {
                self.combined.push(b'\n');
            }
        }
        self.total_bytes = new_total;
        let max_ring = RING_BUFFER_BYTES;
        if self.combined.len() > max_ring {
            let drop = self.combined.len() - max_ring;
            self.combined.drain(..drop);
        }
    }

    fn read_from(&self, cursor: usize, limit: usize) -> (Vec<u8>, usize, bool) {
        let data = &self.combined;
        if cursor >= data.len() {
            return (Vec::new(), data.len(), true);
        }
        let remaining = &data[cursor..];
        let take = remaining.len().min(limit);
        let chunk = remaining[..take].to_vec();
        let next = cursor + take;
        let eof = next >= data.len();
        (chunk, next, eof)
    }
}

#[derive(Clone, Copy)]
enum StreamKind {
    Stdout,
    Stderr,
}

pub(crate) enum BgControl {
    Kill,
}

pub struct BgEntry {
    pub session_id: String,
    pub(crate) control_tx: mpsc::Sender<BgControl>,
    pub status: Arc<Mutex<BgStatus>>,
    pub output: Arc<Mutex<BgOutput>>,
    pub log_path: std::path::PathBuf,
    pub task_id: Option<crate::task_registry::TaskId>,
}

impl crate::watch::Watchable for BgEntry {
    fn watch_output(
        self: std::sync::Arc<Self>,
        pattern: String,
        cancel: tokio_util::sync::CancellationToken,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = crate::watch::WatchResult> + Send>>
    {
        let output = self.output.clone();
        let status = self.status.clone();
        Box::pin(async move {
            loop {
                tokio::select! {
                    _ = cancel.cancelled() => return crate::watch::WatchResult::Cancelled,
                    _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
                        let text = {
                            let out = output.lock().unwrap();
                            String::from_utf8_lossy(&out.combined).into_owned()
                        };
                        if let Some(pos) = text.find(&pattern) {
                            return crate::watch::WatchResult::Matched {
                                row: None,
                                col: Some(pos as u16),
                                text: pattern,
                            };
                        }
                        let st = status.lock().unwrap().clone();
                        if !matches!(st, BgStatus::Running { .. }) {
                            return crate::watch::WatchResult::SourceExited;
                        }
                    }
                }
            }
        })
    }
}

#[derive(Default)]
pub struct BgRegistry {
    entries: Mutex<HashMap<String, Arc<BgEntry>>>,
    task_registry: Option<TaskRegistry>,
}

impl BgRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_task_registry(mut self, tr: TaskRegistry) -> Self {
        self.task_registry = Some(tr);
        self
    }

    pub fn kill_all(&self) {
        let entries = self.entries.lock().unwrap();
        for (_, entry) in entries.iter() {
            let _ = entry.control_tx.try_send(BgControl::Kill);
        }
    }

    pub fn spawn(
        self: &Arc<Self>,
        cmd: String,
        timeout_ms: Option<u64>,
        max_output_bytes: u64,
        ctx: &ToolCtx,
    ) -> Result<Value, RuntimeError> {
        let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
        let local_id = uuid::Uuid::now_v7().as_u64_pair().0;
        let handle = BgHandle {
            session_id: session_id.clone(),
            local_id,
        };
        let handle_str = handle.to_string();

        let dir = ctx.session_dir.clone().ok_or_else(|| {
            RuntimeError::ToolFailed("bash.spawn: session_dir not available".into())
        })?;
        std::fs::create_dir_all(&dir).map_err(|e| {
            RuntimeError::ToolFailed(format!("bash.spawn: create session_dir: {e}"))
        })?;
        let log_path = dir.join(format!("bg_{}.log", handle_str));

        let timeout = match timeout_ms {
            Some(0) => None,
            Some(ms) => Some(Duration::from_millis(ms.min(MAX_SPAWN_TIMEOUT_MS))),
            None => Some(Duration::from_millis(DEFAULT_SPAWN_TIMEOUT_MS)),
        };

        let (control_tx, control_rx) = mpsc::channel::<BgControl>(8);
        let status = Arc::new(Mutex::new(BgStatus::Running {
            pid: 0,
            started_at: now_ms(),
        }));
        let output = Arc::new(Mutex::new(BgOutput::default()));
        let cancel = ctx.cancel.clone();
        let task_cancel = cancel.child_token();

        let task_id = self.task_registry.as_ref().map(|tr| {
            tr.register(
                TaskKind::Bash,
                cmd.clone(),
                handle_str.clone(),
                session_id.clone(),
                task_cancel.clone(),
            )
        });

        let entry = Arc::new(BgEntry {
            session_id: session_id.clone(),
            control_tx,
            status: status.clone(),
            output: output.clone(),
            log_path: log_path.clone(),
            task_id: task_id.clone(),
        });
        {
            let mut entries = self.entries.lock().unwrap();
            entries.insert(handle_str.clone(), entry.clone());
        }

        let registry = Arc::clone(self);
        let handle_str_for_task = handle_str.clone();
        let status_for_task = status.clone();
        let log_path_for_return = log_path.clone();
        let stream_tx = ctx.stream_tx.clone();
        let handle_for_task = handle_str.clone();
        let task_registry = self.task_registry.clone();
        let task_id_for_spawn = task_id.clone();
        let flow_run_id = ctx.flow_run_id.as_ref().map(|r| r.0.to_string());
        tokio::spawn(async move {
            run_bg_process(
                handle_str_for_task,
                cmd,
                timeout,
                max_output_bytes,
                log_path,
                status_for_task,
                output,
                control_rx,
                task_cancel,
                registry,
                stream_tx,
                handle_for_task,
                task_registry,
                task_id_for_spawn,
                flow_run_id,
            )
            .await;
        });

        let pid = {
            let s = status.lock().unwrap();
            if let BgStatus::Running { pid, .. } = &*s {
                *pid
            } else {
                0
            }
        };

        Ok(Value::Struct(vec![
            ("handle".into(), Value::Str(handle_str)),
            ("status".into(), Value::Str("running".into())),
            ("pid".into(), Value::Int(pid as i64)),
            (
                "log_path".into(),
                Value::Str(log_path_for_return.to_string_lossy().into_owned()),
            ),
        ]))
    }

    pub fn lookup(&self, handle_str: &str, session_id: &str) -> Result<Arc<BgEntry>, RuntimeError> {
        let handle = BgHandle::parse(handle_str).ok_or_else(|| {
            RuntimeError::ToolFailed(format!("bash: invalid handle `{handle_str}`"))
        })?;
        if handle.session_id != session_id {
            return Err(RuntimeError::ToolFailed(format!(
                "bash: handle `{handle_str}` does not belong to session `{session_id}`"
            )));
        }
        let entries = self.entries.lock().unwrap();
        entries.get(handle_str).cloned().ok_or_else(|| {
            RuntimeError::ToolFailed(format!("bash: handle `{handle_str}` not found"))
        })
    }

    pub fn status(&self, handle_str: &str, session_id: &str) -> Result<Value, RuntimeError> {
        let entry = self.lookup(handle_str, session_id)?;
        let st = entry.status.lock().unwrap().clone();
        let out = entry.output.lock().unwrap();
        let mut fields = vec![
            ("handle".into(), Value::Str(handle_str.into())),
            ("status".into(), Value::Str(st.kind().into())),
            ("started_at".into(), Value::Int(st.started_at())),
            (
                "log_path".into(),
                Value::Str(entry.log_path.to_string_lossy().into_owned()),
            ),
        ];
        if let Some(ec) = st.exit_code() {
            fields.push(("exit_code".into(), Value::Int(ec as i64)));
        }
        if let Some(ended) = st.ended_at() {
            fields.push(("ended_at".into(), Value::Int(ended)));
        }
        fields.push(("bytes_total".into(), Value::Int(out.total_bytes as i64)));
        fields.push(("output_truncated".into(), Value::Bool(out.truncated)));
        Ok(Value::Struct(fields))
    }

    pub fn output(
        &self,
        handle_str: &str,
        session_id: &str,
        session_dir: Option<&std::path::Path>,
        cursor: usize,
        limit: usize,
    ) -> Result<Value, RuntimeError> {
        if let Ok(entry) = self.lookup(handle_str, session_id) {
            let st = entry.status.lock().unwrap().clone();
            let out = entry.output.lock().unwrap();
            let (chunk, next, eof) = out.read_from(cursor, limit);
            return Ok(Value::Struct(vec![
                ("handle".into(), Value::Str(handle_str.into())),
                ("status".into(), Value::Str(st.kind().into())),
                (
                    "chunk".into(),
                    Value::Str(String::from_utf8_lossy(&chunk).into_owned()),
                ),
                ("cursor".into(), Value::Int(cursor as i64)),
                ("next_cursor".into(), Value::Int(next as i64)),
                ("eof".into(), Value::Bool(eof)),
                ("truncated".into(), Value::Bool(out.truncated)),
                ("live".into(), Value::Bool(true)),
            ]));
        }

        let Some(dir) = session_dir else {
            return Err(RuntimeError::ToolFailed(format!(
                "bash: handle `{handle_str}` not found"
            )));
        };
        let log_path = dir.join(format!("bg_{handle_str}.log"));
        let data = std::fs::read(&log_path).map_err(|_| {
            RuntimeError::ToolFailed(format!("bash: handle `{handle_str}` not found"))
        })?;
        if cursor >= data.len() {
            return Ok(Value::Struct(vec![
                ("handle".into(), Value::Str(handle_str.into())),
                ("status".into(), Value::Str("exited".into())),
                ("chunk".into(), Value::Str(String::new())),
                ("cursor".into(), Value::Int(cursor as i64)),
                ("next_cursor".into(), Value::Int(data.len() as i64)),
                ("eof".into(), Value::Bool(true)),
                ("truncated".into(), Value::Bool(false)),
                ("live".into(), Value::Bool(false)),
            ]));
        }
        let take = (data.len() - cursor).min(limit);
        let chunk = data[cursor..cursor + take].to_vec();
        let next = cursor + take;
        let eof = next >= data.len();
        Ok(Value::Struct(vec![
            ("handle".into(), Value::Str(handle_str.into())),
            ("status".into(), Value::Str("exited".into())),
            (
                "chunk".into(),
                Value::Str(String::from_utf8_lossy(&chunk).into_owned()),
            ),
            ("cursor".into(), Value::Int(cursor as i64)),
            ("next_cursor".into(), Value::Int(next as i64)),
            ("eof".into(), Value::Bool(eof)),
            ("truncated".into(), Value::Bool(false)),
            ("live".into(), Value::Bool(false)),
        ]))
    }

    pub fn kill(&self, handle_str: &str, session_id: &str) -> Result<Value, RuntimeError> {
        let entry = self.lookup(handle_str, session_id)?;
        let _ = entry.control_tx.try_send(BgControl::Kill);
        let st = entry.status.lock().unwrap().clone();
        Ok(Value::Struct(vec![
            ("handle".into(), Value::Str(handle_str.into())),
            ("status".into(), Value::Str(st.kind().into())),
        ]))
    }

    fn remove(&self, handle_str: &str) {
        self.entries.lock().unwrap().remove(handle_str);
    }

    #[doc(hidden)]
    pub fn clear_for_test(&self) {
        self.entries.lock().unwrap().clear();
    }

    pub fn list(
        &self,
        session_id: &str,
        session_dir: Option<&std::path::Path>,
        all: bool,
    ) -> Value {
        let entries = self.entries.lock().unwrap();
        let mut live_handles: std::collections::HashSet<String> = std::collections::HashSet::new();
        let mut items: Vec<Value> = entries
            .iter()
            .filter(|(_, e)| e.session_id == session_id)
            .map(|(handle, entry)| {
                live_handles.insert(handle.clone());
                let st = entry.status.lock().unwrap().clone();
                let out = entry.output.lock().unwrap();
                let mut fields = vec![
                    ("handle".into(), Value::Str(handle.clone())),
                    ("status".into(), Value::Str(st.kind().into())),
                    ("started_at".into(), Value::Int(st.started_at())),
                    ("live".into(), Value::Bool(true)),
                ];
                if let Some(ec) = st.exit_code() {
                    fields.push(("exit_code".into(), Value::Int(ec as i64)));
                }
                fields.push(("bytes_total".into(), Value::Int(out.total_bytes as i64)));
                Value::Struct(fields)
            })
            .collect();

        if all {
            if let Some(dir) = session_dir {
                if let Ok(rd) = std::fs::read_dir(dir) {
                    for entry in rd.flatten() {
                        let name = entry.file_name();
                        let name = name.to_string_lossy();
                        let Some(rest) = name
                            .strip_prefix("bg_")
                            .and_then(|s| s.strip_suffix(".log"))
                        else {
                            continue;
                        };
                        let handle: String = rest.to_string();
                        if live_handles.contains(&handle) {
                            continue;
                        }
                        let Ok(meta) = entry.metadata() else {
                            continue;
                        };
                        let modified = meta
                            .modified()
                            .ok()
                            .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
                            .map(|d| d.as_millis() as i64)
                            .unwrap_or(0);
                        items.push(Value::Struct(vec![
                            ("handle".into(), Value::Str(handle)),
                            ("status".into(), Value::Str("exited".into())),
                            ("started_at".into(), Value::Int(modified)),
                            ("live".into(), Value::Bool(false)),
                            ("bytes_total".into(), Value::Int(meta.len() as i64)),
                        ]));
                    }
                }
            }
        }

        Value::List(items)
    }
}

impl Drop for BgRegistry {
    fn drop(&mut self) {
        self.kill_all();
    }
}

#[allow(clippy::too_many_arguments)]
async fn run_bg_process(
    handle_str: String,
    cmd: String,
    timeout: Option<Duration>,
    max_output_bytes: u64,
    log_path: std::path::PathBuf,
    status: Arc<Mutex<BgStatus>>,
    output: Arc<Mutex<BgOutput>>,
    mut control_rx: mpsc::Receiver<BgControl>,
    cancel: CancellationToken,
    registry: Arc<BgRegistry>,
    stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
    handle_for_stream: String,
    task_registry: Option<TaskRegistry>,
    task_id: Option<crate::task_registry::TaskId>,
    flow_run_id: Option<String>,
) {
    let started_at = now_ms();
    let mut command = tokio::process::Command::new("sh");
    command
        .arg("-c")
        .arg(&cmd)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut child: AsyncGroupChild = match command.group().kill_on_drop(true).spawn() {
        Ok(c) => c,
        Err(e) => {
            *status.lock().unwrap() = BgStatus::Failed {
                error: format!("spawn: {e}"),
                started_at,
                ended_at: now_ms(),
            };
            registry.remove(&handle_str);
            return;
        }
    };
    let pid = child.id();
    *status.lock().unwrap() = BgStatus::Running {
        pid: pid.unwrap_or(0),
        started_at,
    };

    let stdout = child.inner().stdout.take();
    let stderr = child.inner().stderr.take();

    let stdout_reader = stdout.map(|s| {
        let ctx = ReadStreamCtx {
            output: output.clone(),
            log_path: log_path.clone(),
            kind: StreamKind::Stdout,
            max_output_bytes,
            stream_tx: stream_tx.clone(),
            handle: handle_for_stream.clone(),
            flow_run_id: flow_run_id.clone(),
        };
        tokio::spawn(read_stream(BufReader::new(s), ctx))
    });
    let stderr_reader = stderr.map(|s| {
        let ctx = ReadStreamCtx {
            output: output.clone(),
            log_path: log_path.clone(),
            kind: StreamKind::Stderr,
            max_output_bytes,
            stream_tx: stream_tx.clone(),
            handle: handle_for_stream.clone(),
            flow_run_id: flow_run_id.clone(),
        };
        tokio::spawn(read_stream(BufReader::new(s), ctx))
    });

    let exit_reason = tokio::select! {
        biased;
        _ = cancel.cancelled() => ExitReason::Cancelled,
        ctrl = control_rx.recv() => {
            match ctrl {
                Some(BgControl::Kill) => ExitReason::Kill,
                None => ExitReason::Natural,
            }
        }
        _ = async {
            if let Some(t) = timeout {
                tokio::time::sleep(t).await;
            } else {
                std::future::pending::<()>().await;
            }
        } => ExitReason::Timeout,
        s = child.wait() => ExitReason::Exited(s),
    };

    let ended_at = now_ms();
    let final_status = match &exit_reason {
        ExitReason::Exited(Ok(s)) => BgStatus::Exited {
            exit_code: s.code().unwrap_or(-1),
            started_at,
            ended_at,
        },
        ExitReason::Timeout => {
            let _ = child.start_kill();
            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
            BgStatus::TimedOut {
                started_at,
                ended_at,
            }
        }
        ExitReason::Kill => {
            let _ = child.start_kill();
            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
            BgStatus::Killed {
                started_at,
                ended_at,
            }
        }
        ExitReason::Cancelled => {
            let _ = child.start_kill();
            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
            BgStatus::Killed {
                started_at,
                ended_at,
            }
        }
        ExitReason::Exited(Err(_)) => {
            let _ = child.start_kill();
            BgStatus::Failed {
                error: "wait failed".into(),
                started_at,
                ended_at,
            }
        }
        ExitReason::Natural => {
            let s = child.wait().await;
            BgStatus::Exited {
                exit_code: s.ok().and_then(|s| s.code()).unwrap_or(-1),
                started_at,
                ended_at: now_ms(),
            }
        }
    };

    if let Some(r) = stdout_reader {
        let _ = tokio::time::timeout(IO_DRAIN_TIMEOUT, r).await;
    }
    if let Some(r) = stderr_reader {
        let _ = tokio::time::timeout(IO_DRAIN_TIMEOUT, r).await;
    }

    let exit_code = match &final_status {
        BgStatus::Exited { exit_code, .. } => Some(*exit_code),
        _ => None,
    };
    *status.lock().unwrap() = final_status.clone();

    if let Some(tx) = &stream_tx {
        let _ = tx.send(crate::stream::StreamFrame::BashExited {
            handle: handle_for_stream,
            exit_code,
            run_id: flow_run_id,
        });
    }

    if let (Some(tr), Some(tid)) = (task_registry, task_id) {
        let ts = match &final_status {
            BgStatus::Exited { exit_code, .. } if *exit_code == 0 => TaskStatus::Ok,
            BgStatus::Killed { .. } | BgStatus::TimedOut { .. } => TaskStatus::Killed,
            _ => TaskStatus::Err,
        };
        tr.finish(&tid, ts);
    }
}

struct ReadStreamCtx {
    output: Arc<Mutex<BgOutput>>,
    log_path: std::path::PathBuf,
    kind: StreamKind,
    max_output_bytes: u64,
    stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
    handle: String,
    flow_run_id: Option<String>,
}

async fn read_stream<R: tokio::io::AsyncBufRead + Unpin>(mut reader: R, ctx: ReadStreamCtx) {
    let prefix: &[u8] = match ctx.kind {
        StreamKind::Stdout => b"[out] ",
        StreamKind::Stderr => b"[err] ",
    };
    let kind_str = match ctx.kind {
        StreamKind::Stdout => "stdout",
        StreamKind::Stderr => "stderr",
    };
    let mut buf = String::new();
    let mut log_file = tokio::fs::OpenOptions::new()
        .append(true)
        .create(true)
        .open(&ctx.log_path)
        .await
        .ok();
    loop {
        buf.clear();
        match reader.read_line(&mut buf).await {
            Ok(0) => break,
            Ok(_) => {
                let data = buf.as_bytes();
                {
                    let mut out = ctx.output.lock().unwrap();
                    out.push(ctx.kind, data, ctx.max_output_bytes);
                }
                if let Some(file) = log_file.as_mut() {
                    let _ = file.write_all(prefix).await;
                    let _ = file.write_all(data).await;
                    if !data.ends_with(b"\n") {
                        let _ = file.write_all(b"\n").await;
                    }
                }
                if let Some(tx) = &ctx.stream_tx {
                    let _ = tx.send(crate::stream::StreamFrame::BashChunk {
                        handle: ctx.handle.clone(),
                        kind: kind_str.to_string(),
                        line: buf.clone(),
                        run_id: ctx.flow_run_id.clone(),
                    });
                }
            }
            Err(_) => break,
        }
    }
}

enum ExitReason {
    Exited(std::io::Result<std::process::ExitStatus>),
    Timeout,
    Kill,
    Cancelled,
    Natural,
}

fn now_ms() -> i64 {
    chrono::Utc::now().timestamp_millis()
}

pub struct BashSpawn;

impl Tool for BashSpawn {
    fn name(&self) -> &str {
        "bash.spawn"
    }

    fn tier(&self) -> Tier {
        Tier::Four
    }

    fn description(&self) -> Option<&str> {
        Some(
            "Run a shell command via `sh -c`.\n\n\
block=false (default): command runs in background, returns immediately with a\n\
handle. The command keeps running — use bash.output to read its output later,\n\
bash.status to check if it finished, bash.kill to stop it. Use this for:\n\
- long-running commands (servers, watchers)\n\
- commands where you need to check output incrementally\n\
- when you want to do other things while the command runs\n\n\
block=true: waits for the command to finish, then returns stdout/stderr/exit_code.\n\
Use block_timeout_ms to set a max wait (default 30s). Use this for:\n\
- short commands where you need the result immediately (ls, git status, echo)\n\
- commands that finish quickly\n\n\
Do NOT use `sleep` in your command to wait — use block=true with block_timeout_ms\n\
instead, or use the sleep tool to pause the workflow.",
        )
    }

    fn input_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "cmd": {"type": "string", "description": "Shell command line."},
                "block": {"type": "boolean", "default": false, "description": "If true, wait for process to exit before returning."},
                "block_timeout_ms": {"type": "integer", "description": "Only with block=true. Max wait. 0 = no timeout. Default 30000."},
                "timeout_ms": {"type": "integer", "description": "Process kill timeout in ms. Default 1800000 (30min). 0 = no timeout."},
                "max_output_bytes": {"type": "integer", "description": "Max combined output bytes. Default 10485760 (10MB)."}
            },
            "required": ["cmd"]
        })
    }

    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
        Box::pin(async move {
            let cmd = extract_string(&args, "cmd", 0)?;
            let block = args
                .named("block")
                .and_then(|v| {
                    if let Value::Bool(b) = v {
                        Some(*b)
                    } else {
                        None
                    }
                })
                .unwrap_or(false);
            let block_timeout_ms = extract_optional_int(&args, "block_timeout_ms")
                .map(|v| v as u64)
                .unwrap_or(30_000);
            let timeout_ms = extract_optional_int(&args, "timeout_ms").map(|v| v as u64);
            let max_output = extract_optional_int(&args, "max_output_bytes")
                .map(|v| v as u64)
                .unwrap_or(DEFAULT_MAX_OUTPUT_BYTES);
            let registry = ctx.bg_registry.clone().ok_or_else(|| {
                RuntimeError::ToolFailed("bash.spawn: registry not available".into())
            })?;
            let handle_str = registry.spawn(cmd, timeout_ms, max_output, ctx)?;

            if !block {
                return Ok(handle_str);
            }

            let handle_s = handle_str
                .field("handle")
                .and_then(|v| {
                    if let Value::Str(s) = v {
                        Some(s.clone())
                    } else {
                        None
                    }
                })
                .ok_or_else(|| {
                    RuntimeError::ToolFailed("bash.spawn: missing handle field".into())
                })?;
            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".into());

            let deadline = if block_timeout_ms == 0 {
                None
            } else {
                Some(tokio::time::Instant::now() + Duration::from_millis(block_timeout_ms))
            };
            loop {
                let entry = registry.lookup(&handle_s, &session_id)?;
                let finished = {
                    let st = entry.status.lock().unwrap();
                    st.is_finished()
                };
                if finished {
                    break;
                }
                if let Some(d) = deadline {
                    if tokio::time::Instant::now() >= d {
                        break;
                    }
                }
                tokio::time::sleep(Duration::from_millis(50)).await;
            }

            let entry = registry.lookup(&handle_s, &session_id)?;
            let st = entry.status.lock().unwrap().clone();
            let out = entry.output.lock().unwrap();
            let combined = String::from_utf8_lossy(&out.combined).into_owned();
            let log_path = entry.log_path.to_string_lossy().into_owned();
            Ok(Value::Struct(vec![
                ("handle".into(), Value::Str(handle_s)),
                ("status".into(), Value::Str(st.kind().into())),
                (
                    "exit_code".into(),
                    st.exit_code()
                        .map(|c| Value::Int(c as i64))
                        .unwrap_or(Value::Unit),
                ),
                ("output".into(), Value::Str(combined)),
                ("bytes_total".into(), Value::Int(out.total_bytes as i64)),
                ("log_path".into(), Value::Str(log_path)),
            ]))
        })
    }
}

pub struct BashStatus;

impl Tool for BashStatus {
    fn name(&self) -> &str {
        "bash.status"
    }

    fn tier(&self) -> Tier {
        Tier::Four
    }

    fn description(&self) -> Option<&str> {
        Some("Check the status of a background bash process.")
    }

    fn input_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {"handle": {"type": "string"}},
            "required": ["handle"]
        })
    }

    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
        Box::pin(async move {
            let handle = extract_string(&args, "handle", 0)?;
            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
            let registry = ctx.bg_registry.clone().ok_or_else(|| {
                RuntimeError::ToolFailed("bash.status: registry not available".into())
            })?;
            registry.status(&handle, &session_id)
        })
    }
}

pub struct BashOutput;

impl Tool for BashOutput {
    fn name(&self) -> &str {
        "bash.output"
    }

    fn tier(&self) -> Tier {
        Tier::Four
    }

    fn description(&self) -> Option<&str> {
        Some("Read output from a background bash process by byte cursor.")
    }

    fn input_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "handle": {"type": "string"},
                "cursor": {"type": "integer", "description": "Byte offset to start reading. Default 0."},
                "limit_bytes": {"type": "integer", "description": "Max bytes to return. Default 32000."}
            },
            "required": ["handle"]
        })
    }

    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
        Box::pin(async move {
            let handle = extract_string(&args, "handle", 0)?;
            let cursor = extract_optional_int(&args, "cursor").unwrap_or(0).max(0) as usize;
            let limit = extract_optional_int(&args, "limit_bytes")
                .unwrap_or(DEFAULT_OUTPUT_LIMIT as i64)
                .max(1) as usize;
            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
            let registry = ctx.bg_registry.clone().ok_or_else(|| {
                RuntimeError::ToolFailed("bash.output: registry not available".into())
            })?;
            registry.output(
                &handle,
                &session_id,
                ctx.session_dir.as_deref(),
                cursor,
                limit,
            )
        })
    }
}

pub struct BashKill;

impl Tool for BashKill {
    fn name(&self) -> &str {
        "bash.kill"
    }

    fn tier(&self) -> Tier {
        Tier::Four
    }

    fn description(&self) -> Option<&str> {
        Some(
            "Kill a background bash process. signal=term (default) sends SIGTERM, signal=kill sends SIGKILL.",
        )
    }

    fn input_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "handle": {"type": "string"},
                "signal": {"type": "string", "enum": ["term", "kill"], "description": "Default term."}
            },
            "required": ["handle"]
        })
    }

    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
        Box::pin(async move {
            let handle = extract_string(&args, "handle", 0)?;
            let _ = extract_string(&args, "signal", 1);
            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
            let registry = ctx.bg_registry.clone().ok_or_else(|| {
                RuntimeError::ToolFailed("bash.kill: registry not available".into())
            })?;
            registry.kill(&handle, &session_id)
        })
    }
}

pub struct BashList;

impl Tool for BashList {
    fn name(&self) -> &str {
        "bash.list"
    }

    fn tier(&self) -> Tier {
        Tier::Four
    }

    fn description(&self) -> Option<&str> {
        Some(
            "List background bash processes for the current session. Default: only live processes. Pass all=true to include historical processes whose log files persist in session_dir (status=exited, live=false).",
        )
    }

    fn input_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "all": {"type": "boolean", "description": "Include historical processes (default false)."}
            }
        })
    }

    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
        Box::pin(async move {
            let all = extract_optional_bool(&args, "all").unwrap_or(false);
            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
            let registry = ctx.bg_registry.clone().ok_or_else(|| {
                RuntimeError::ToolFailed("bash.list: registry not available".into())
            })?;
            Ok(registry.list(&session_id, ctx.session_dir.as_deref(), all))
        })
    }
}

fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
    let value = match args.named(name) {
        Some(v) => v,
        None => args.positional(pos)?,
    };
    match value {
        Value::Str(s) => Ok(s.clone()),
        other => Err(RuntimeError::TypeMismatch {
            expected: "string".into(),
            actual: other.kind_name().into(),
        }),
    }
}

fn extract_optional_int(args: &ToolArgs, name: &str) -> Option<i64> {
    match args.named(name)? {
        Value::Int(n) => Some(*n),
        _ => None,
    }
}

fn extract_optional_bool(args: &ToolArgs, name: &str) -> Option<bool> {
    match args.named(name)? {
        Value::Bool(b) => Some(*b),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tool::{ToolArgs, ToolCtx};
    use crate::value::Value;
    use std::sync::Arc;
    use tempfile::TempDir;

    fn ctx_with_registry(registry: Arc<BgRegistry>, dir: &std::path::Path) -> ToolCtx {
        let mut ctx = ToolCtx::new();
        ctx.bg_registry = Some(registry);
        ctx.session_dir = Some(dir.to_path_buf());
        ctx.session_id = Some("test-session".to_string());
        ctx
    }

    #[test]
    fn handle_parse_roundtrip() {
        let h = BgHandle {
            session_id: "abc".into(),
            local_id: 42,
        };
        let s = h.to_string();
        assert_eq!(s, "bg_abc_42");
        let back = BgHandle::parse(&s).unwrap();
        assert_eq!(back, h);
    }

    #[test]
    fn handle_parse_rejects_bad_format() {
        assert!(BgHandle::parse("not_bg").is_none());
        assert!(BgHandle::parse("bg_nosuffix").is_none());
        assert!(BgHandle::parse("bg_x_notnum").is_none());
    }

    #[tokio::test]
    async fn spawn_returns_immediately_with_running_status() {
        let registry = Arc::new(BgRegistry::new());
        let dir = TempDir::new().unwrap();
        let ctx = ctx_with_registry(registry.clone(), dir.path());
        let args = ToolArgs {
            positional: vec![Value::Str("echo hello".into())],
            named: vec![],
        };
        let v = BashSpawn.call(args, &ctx).await.unwrap();
        let Value::Struct(fields) = v else {
            panic!("expected struct")
        };
        let handle = fields
            .iter()
            .find(|(k, _)| k == "handle")
            .and_then(|(_, v)| {
                if let Value::Str(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .unwrap();
        assert!(handle.starts_with("bg_"));
        let status_val = fields.iter().find(|(k, _)| k == "status").unwrap();
        assert!(matches!(&status_val.1, Value::Str(s) if s == "running"));
    }

    #[tokio::test]
    async fn spawn_then_status_reaches_exited() {
        let registry = Arc::new(BgRegistry::new());
        let dir = TempDir::new().unwrap();
        let ctx = ctx_with_registry(registry.clone(), dir.path());
        let spawn_args = ToolArgs {
            positional: vec![Value::Str("echo hello".into())],
            named: vec![],
        };
        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
        let Value::Struct(fields) = v else { panic!() };
        let handle = fields
            .iter()
            .find(|(k, _)| k == "handle")
            .and_then(|(_, v)| {
                if let Value::Str(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .unwrap();

        for _ in 0..50 {
            tokio::time::sleep(Duration::from_millis(50)).await;
            let status_args = ToolArgs {
                positional: vec![Value::Str(handle.clone())],
                named: vec![],
            };
            let s = BashStatus.call(status_args, &ctx).await.unwrap();
            if let Value::Struct(sf) = s {
                let kind = sf.iter().find(|(k, _)| k == "status").unwrap();
                if matches!(&kind.1, Value::Str(s) if s == "exited") {
                    let ec = sf.iter().find(|(k, _)| k == "exit_code").unwrap();
                    assert!(matches!(ec.1, Value::Int(0)));
                    return;
                }
            }
        }
        panic!("process did not exit in time");
    }

    #[tokio::test]
    async fn spawn_output_captures_stdout() {
        let registry = Arc::new(BgRegistry::new());
        let dir = TempDir::new().unwrap();
        let ctx = ctx_with_registry(registry.clone(), dir.path());
        let spawn_args = ToolArgs {
            positional: vec![Value::Str("echo line1; echo line2".into())],
            named: vec![],
        };
        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
        let Value::Struct(fields) = v else { panic!() };
        let handle = fields
            .iter()
            .find(|(k, _)| k == "handle")
            .and_then(|(_, v)| {
                if let Value::Str(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .unwrap();

        tokio::time::sleep(Duration::from_millis(300)).await;

        let out_args = ToolArgs {
            positional: vec![Value::Str(handle.clone())],
            named: vec![],
        };
        let o = BashOutput.call(out_args, &ctx).await.unwrap();
        let Value::Struct(of) = o else { panic!() };
        let chunk = of.iter().find(|(k, _)| k == "chunk").unwrap();
        if let Value::Str(s) = &chunk.1 {
            assert!(s.contains("line1"), "chunk should contain line1: {s}");
            assert!(s.contains("line2"), "chunk should contain line2: {s}");
        } else {
            panic!("chunk not str");
        }
    }

    #[tokio::test]
    async fn kill_terminates_long_running_process() {
        let registry = Arc::new(BgRegistry::new());
        let dir = TempDir::new().unwrap();
        let ctx = ctx_with_registry(registry.clone(), dir.path());
        let spawn_args = ToolArgs {
            positional: vec![Value::Str("sleep 100".into())],
            named: vec![],
        };
        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
        let Value::Struct(fields) = v else { panic!() };
        let handle = fields
            .iter()
            .find(|(k, _)| k == "handle")
            .and_then(|(_, v)| {
                if let Value::Str(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .unwrap();

        let kill_args = ToolArgs {
            positional: vec![Value::Str(handle.clone())],
            named: vec![],
        };
        BashKill.call(kill_args, &ctx).await.unwrap();

        for _ in 0..50 {
            tokio::time::sleep(Duration::from_millis(50)).await;
            let status_args = ToolArgs {
                positional: vec![Value::Str(handle.clone())],
                named: vec![],
            };
            let s = BashStatus.call(status_args, &ctx).await.unwrap();
            if let Value::Struct(sf) = s {
                let kind = sf.iter().find(|(k, _)| k == "status").unwrap();
                if matches!(&kind.1, Value::Str(s) if s == "killed") {
                    return;
                }
            }
        }
        panic!("process not killed in time");
    }

    #[tokio::test]
    async fn cross_session_access_rejected() {
        let registry = Arc::new(BgRegistry::new());
        let dir = TempDir::new().unwrap();
        let ctx_a = ctx_with_registry(registry.clone(), dir.path());

        let spawn_args = ToolArgs {
            positional: vec![Value::Str("sleep 10".into())],
            named: vec![],
        };
        let v = BashSpawn.call(spawn_args, &ctx_a).await.unwrap();
        let Value::Struct(fields) = v else { panic!() };
        let handle = fields
            .iter()
            .find(|(k, _)| k == "handle")
            .and_then(|(_, v)| {
                if let Value::Str(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .unwrap();

        let mut ctx_b = ToolCtx::new();
        ctx_b.bg_registry = Some(registry.clone());
        ctx_b.session_dir = Some(dir.path().to_path_buf());
        ctx_b.session_id = Some("other-session".to_string());
        let status_args = ToolArgs {
            positional: vec![Value::Str(handle)],
            named: vec![],
        };
        let err = BashStatus.call(status_args, &ctx_b).await.err().unwrap();
        assert!(format!("{err}").contains("does not belong to session"));
    }
}