Skip to main content

atman_runtime/tools/
bash_bg.rs

1use std::collections::HashMap;
2use std::process::Stdio;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6
7use command_group::{AsyncCommandGroup, AsyncGroupChild};
8use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
9use tokio::sync::mpsc;
10use tokio_util::sync::CancellationToken;
11
12use crate::error::RuntimeError;
13use crate::task_registry::{TaskKind, TaskRegistry, TaskStatus};
14use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
15use crate::value::Value;
16
17const DEFAULT_SPAWN_TIMEOUT_MS: u64 = 1_800_000;
18const MAX_SPAWN_TIMEOUT_MS: u64 = 86_400_000;
19const DEFAULT_MAX_OUTPUT_BYTES: u64 = 10_485_760;
20const RING_BUFFER_BYTES: usize = 65_536;
21const IO_DRAIN_TIMEOUT: Duration = Duration::from_secs(2);
22const DEFAULT_OUTPUT_LIMIT: usize = 32_000;
23
24#[derive(Debug, Clone, Hash, PartialEq, Eq)]
25pub struct BgHandle {
26    session_id: String,
27    local_id: u64,
28}
29
30impl BgHandle {
31    #[allow(clippy::inherent_to_string)]
32    pub fn to_string(&self) -> String {
33        format!("bg_{}_{}", self.session_id, self.local_id)
34    }
35
36    pub fn parse(s: &str) -> Option<Self> {
37        let rest = s.strip_prefix("bg_")?;
38        let idx = rest.rfind('_')?;
39        let session_id = rest[..idx].to_string();
40        let local_id = rest[idx + 1..].parse().ok()?;
41        Some(Self {
42            session_id,
43            local_id,
44        })
45    }
46}
47
48#[derive(Debug, Clone)]
49pub enum BgStatus {
50    Running {
51        pid: u32,
52        started_at: i64,
53    },
54    Exited {
55        exit_code: i32,
56        started_at: i64,
57        ended_at: i64,
58    },
59    TimedOut {
60        started_at: i64,
61        ended_at: i64,
62    },
63    Killed {
64        started_at: i64,
65        ended_at: i64,
66    },
67    Failed {
68        error: String,
69        started_at: i64,
70        ended_at: i64,
71    },
72}
73
74impl BgStatus {
75    fn kind(&self) -> &'static str {
76        match self {
77            Self::Running { .. } => "running",
78            Self::Exited { .. } => "exited",
79            Self::TimedOut { .. } => "timed_out",
80            Self::Killed { .. } => "killed",
81            Self::Failed { .. } => "failed",
82        }
83    }
84
85    fn exit_code(&self) -> Option<i32> {
86        match self {
87            Self::Exited { exit_code, .. } => Some(*exit_code),
88            _ => None,
89        }
90    }
91
92    fn started_at(&self) -> i64 {
93        match self {
94            Self::Running { started_at, .. }
95            | Self::Exited { started_at, .. }
96            | Self::TimedOut { started_at, .. }
97            | Self::Killed { started_at, .. }
98            | Self::Failed { started_at, .. } => *started_at,
99        }
100    }
101
102    fn ended_at(&self) -> Option<i64> {
103        match self {
104            Self::Exited { ended_at, .. }
105            | Self::TimedOut { ended_at, .. }
106            | Self::Killed { ended_at, .. }
107            | Self::Failed { ended_at, .. } => Some(*ended_at),
108            _ => None,
109        }
110    }
111
112    fn is_finished(&self) -> bool {
113        !matches!(self, Self::Running { .. })
114    }
115}
116
117#[derive(Debug, Default)]
118pub struct BgOutput {
119    pub combined: Vec<u8>,
120    pub total_bytes: u64,
121    pub truncated: bool,
122}
123
124impl BgOutput {
125    fn push(&mut self, kind: StreamKind, data: &[u8], max: u64) {
126        let prefix: &[u8] = match kind {
127            StreamKind::Stdout => b"[out] ",
128            StreamKind::Stderr => b"[err] ",
129        };
130        let mut new_total = self.total_bytes + data.len() as u64;
131        let mut to_write = data;
132        if new_total > max {
133            let allowed = max.saturating_sub(self.total_bytes) as usize;
134            to_write = &data[..allowed.min(data.len())];
135            new_total = max;
136            self.truncated = true;
137        }
138        if !to_write.is_empty() {
139            self.combined.extend_from_slice(prefix);
140            self.combined.extend_from_slice(to_write);
141            if !to_write.ends_with(b"\n") {
142                self.combined.push(b'\n');
143            }
144        }
145        self.total_bytes = new_total;
146        let max_ring = RING_BUFFER_BYTES;
147        if self.combined.len() > max_ring {
148            let drop = self.combined.len() - max_ring;
149            self.combined.drain(..drop);
150        }
151    }
152
153    fn read_from(&self, cursor: usize, limit: usize) -> (Vec<u8>, usize, bool) {
154        let data = &self.combined;
155        if cursor >= data.len() {
156            return (Vec::new(), data.len(), true);
157        }
158        let remaining = &data[cursor..];
159        let take = remaining.len().min(limit);
160        let chunk = remaining[..take].to_vec();
161        let next = cursor + take;
162        let eof = next >= data.len();
163        (chunk, next, eof)
164    }
165}
166
167#[derive(Clone, Copy)]
168enum StreamKind {
169    Stdout,
170    Stderr,
171}
172
173pub(crate) enum BgControl {
174    Kill,
175}
176
177pub struct BgEntry {
178    pub session_id: String,
179    pub(crate) control_tx: mpsc::Sender<BgControl>,
180    pub status: Arc<Mutex<BgStatus>>,
181    pub output: Arc<Mutex<BgOutput>>,
182    pub log_path: std::path::PathBuf,
183    pub task_id: Option<crate::task_registry::TaskId>,
184}
185
186#[derive(Default)]
187pub struct BgRegistry {
188    next_id: AtomicU64,
189    entries: Mutex<HashMap<String, Arc<BgEntry>>>,
190    task_registry: Option<TaskRegistry>,
191}
192
193impl BgRegistry {
194    pub fn new() -> Self {
195        Self::default()
196    }
197
198    pub fn with_task_registry(mut self, tr: TaskRegistry) -> Self {
199        self.task_registry = Some(tr);
200        self
201    }
202
203    pub fn kill_all(&self) {
204        let entries = self.entries.lock().unwrap();
205        for (_, entry) in entries.iter() {
206            let _ = entry.control_tx.try_send(BgControl::Kill);
207        }
208    }
209
210    pub fn spawn(
211        self: &Arc<Self>,
212        cmd: String,
213        timeout_ms: Option<u64>,
214        max_output_bytes: u64,
215        ctx: &ToolCtx,
216    ) -> Result<Value, RuntimeError> {
217        let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
218        let local_id = self.next_id.fetch_add(1, Ordering::Relaxed);
219        let handle = BgHandle {
220            session_id: session_id.clone(),
221            local_id,
222        };
223        let handle_str = handle.to_string();
224
225        let dir = ctx.session_dir.clone().ok_or_else(|| {
226            RuntimeError::ToolFailed("bash.spawn: session_dir not available".into())
227        })?;
228        std::fs::create_dir_all(&dir).map_err(|e| {
229            RuntimeError::ToolFailed(format!("bash.spawn: create session_dir: {e}"))
230        })?;
231        let log_path = dir.join(format!("bg_{}.log", handle_str));
232
233        let timeout = match timeout_ms {
234            Some(0) => None,
235            Some(ms) => Some(Duration::from_millis(ms.min(MAX_SPAWN_TIMEOUT_MS))),
236            None => Some(Duration::from_millis(DEFAULT_SPAWN_TIMEOUT_MS)),
237        };
238
239        let (control_tx, control_rx) = mpsc::channel::<BgControl>(8);
240        let status = Arc::new(Mutex::new(BgStatus::Running {
241            pid: 0,
242            started_at: now_ms(),
243        }));
244        let output = Arc::new(Mutex::new(BgOutput::default()));
245        let cancel = ctx.cancel.clone();
246        let task_cancel = cancel.child_token();
247
248        let task_id = self.task_registry.as_ref().map(|tr| {
249            tr.register(
250                TaskKind::Bash,
251                cmd.clone(),
252                handle_str.clone(),
253                session_id.clone(),
254                task_cancel.clone(),
255            )
256        });
257
258        let entry = Arc::new(BgEntry {
259            session_id: session_id.clone(),
260            control_tx,
261            status: status.clone(),
262            output: output.clone(),
263            log_path: log_path.clone(),
264            task_id: task_id.clone(),
265        });
266        {
267            let mut entries = self.entries.lock().unwrap();
268            entries.insert(handle_str.clone(), entry.clone());
269        }
270
271        let registry = Arc::clone(self);
272        let handle_str_for_task = handle_str.clone();
273        let status_for_task = status.clone();
274        let log_path_for_return = log_path.clone();
275        let stream_tx = ctx.stream_tx.clone();
276        let handle_for_task = handle_str.clone();
277        let task_registry = self.task_registry.clone();
278        let task_id_for_spawn = task_id.clone();
279        tokio::spawn(async move {
280            run_bg_process(
281                handle_str_for_task,
282                cmd,
283                timeout,
284                max_output_bytes,
285                log_path,
286                status_for_task,
287                output,
288                control_rx,
289                task_cancel,
290                registry,
291                stream_tx,
292                handle_for_task,
293                task_registry,
294                task_id_for_spawn,
295            )
296            .await;
297        });
298
299        let pid = {
300            let s = status.lock().unwrap();
301            if let BgStatus::Running { pid, .. } = &*s {
302                *pid
303            } else {
304                0
305            }
306        };
307
308        Ok(Value::Struct(vec![
309            ("handle".into(), Value::Str(handle_str)),
310            ("status".into(), Value::Str("running".into())),
311            ("pid".into(), Value::Int(pid as i64)),
312            (
313                "log_path".into(),
314                Value::Str(log_path_for_return.to_string_lossy().into_owned()),
315            ),
316        ]))
317    }
318
319    fn lookup(&self, handle_str: &str, session_id: &str) -> Result<Arc<BgEntry>, RuntimeError> {
320        let handle = BgHandle::parse(handle_str).ok_or_else(|| {
321            RuntimeError::ToolFailed(format!("bash: invalid handle `{handle_str}`"))
322        })?;
323        if handle.session_id != session_id {
324            return Err(RuntimeError::ToolFailed(format!(
325                "bash: handle `{handle_str}` does not belong to session `{session_id}`"
326            )));
327        }
328        let entries = self.entries.lock().unwrap();
329        entries.get(handle_str).cloned().ok_or_else(|| {
330            RuntimeError::ToolFailed(format!("bash: handle `{handle_str}` not found"))
331        })
332    }
333
334    pub fn status(&self, handle_str: &str, session_id: &str) -> Result<Value, RuntimeError> {
335        let entry = self.lookup(handle_str, session_id)?;
336        let st = entry.status.lock().unwrap().clone();
337        let out = entry.output.lock().unwrap();
338        let mut fields = vec![
339            ("handle".into(), Value::Str(handle_str.into())),
340            ("status".into(), Value::Str(st.kind().into())),
341            ("started_at".into(), Value::Int(st.started_at())),
342        ];
343        if let Some(ec) = st.exit_code() {
344            fields.push(("exit_code".into(), Value::Int(ec as i64)));
345        }
346        if let Some(ended) = st.ended_at() {
347            fields.push(("ended_at".into(), Value::Int(ended)));
348        }
349        fields.push(("bytes_total".into(), Value::Int(out.total_bytes as i64)));
350        fields.push(("output_truncated".into(), Value::Bool(out.truncated)));
351        Ok(Value::Struct(fields))
352    }
353
354    pub fn output(
355        &self,
356        handle_str: &str,
357        session_id: &str,
358        session_dir: Option<&std::path::Path>,
359        cursor: usize,
360        limit: usize,
361    ) -> Result<Value, RuntimeError> {
362        if let Ok(entry) = self.lookup(handle_str, session_id) {
363            let st = entry.status.lock().unwrap().clone();
364            let out = entry.output.lock().unwrap();
365            let (chunk, next, eof) = out.read_from(cursor, limit);
366            return Ok(Value::Struct(vec![
367                ("handle".into(), Value::Str(handle_str.into())),
368                ("status".into(), Value::Str(st.kind().into())),
369                (
370                    "chunk".into(),
371                    Value::Str(String::from_utf8_lossy(&chunk).into_owned()),
372                ),
373                ("cursor".into(), Value::Int(cursor as i64)),
374                ("next_cursor".into(), Value::Int(next as i64)),
375                ("eof".into(), Value::Bool(eof)),
376                ("truncated".into(), Value::Bool(out.truncated)),
377                ("live".into(), Value::Bool(true)),
378            ]));
379        }
380
381        let Some(dir) = session_dir else {
382            return Err(RuntimeError::ToolFailed(format!(
383                "bash: handle `{handle_str}` not found"
384            )));
385        };
386        let log_path = dir.join(format!("bg_{handle_str}.log"));
387        let data = std::fs::read(&log_path).map_err(|_| {
388            RuntimeError::ToolFailed(format!("bash: handle `{handle_str}` not found"))
389        })?;
390        if cursor >= data.len() {
391            return Ok(Value::Struct(vec![
392                ("handle".into(), Value::Str(handle_str.into())),
393                ("status".into(), Value::Str("exited".into())),
394                ("chunk".into(), Value::Str(String::new())),
395                ("cursor".into(), Value::Int(cursor as i64)),
396                ("next_cursor".into(), Value::Int(data.len() as i64)),
397                ("eof".into(), Value::Bool(true)),
398                ("truncated".into(), Value::Bool(false)),
399                ("live".into(), Value::Bool(false)),
400            ]));
401        }
402        let take = (data.len() - cursor).min(limit);
403        let chunk = data[cursor..cursor + take].to_vec();
404        let next = cursor + take;
405        let eof = next >= data.len();
406        Ok(Value::Struct(vec![
407            ("handle".into(), Value::Str(handle_str.into())),
408            ("status".into(), Value::Str("exited".into())),
409            (
410                "chunk".into(),
411                Value::Str(String::from_utf8_lossy(&chunk).into_owned()),
412            ),
413            ("cursor".into(), Value::Int(cursor as i64)),
414            ("next_cursor".into(), Value::Int(next as i64)),
415            ("eof".into(), Value::Bool(eof)),
416            ("truncated".into(), Value::Bool(false)),
417            ("live".into(), Value::Bool(false)),
418        ]))
419    }
420
421    pub fn kill(&self, handle_str: &str, session_id: &str) -> Result<Value, RuntimeError> {
422        let entry = self.lookup(handle_str, session_id)?;
423        let _ = entry.control_tx.try_send(BgControl::Kill);
424        let st = entry.status.lock().unwrap().clone();
425        Ok(Value::Struct(vec![
426            ("handle".into(), Value::Str(handle_str.into())),
427            ("status".into(), Value::Str(st.kind().into())),
428        ]))
429    }
430
431    fn remove(&self, handle_str: &str) {
432        self.entries.lock().unwrap().remove(handle_str);
433    }
434
435    #[doc(hidden)]
436    pub fn clear_for_test(&self) {
437        self.entries.lock().unwrap().clear();
438    }
439
440    pub fn list(
441        &self,
442        session_id: &str,
443        session_dir: Option<&std::path::Path>,
444        all: bool,
445    ) -> Value {
446        let entries = self.entries.lock().unwrap();
447        let mut live_handles: std::collections::HashSet<String> = std::collections::HashSet::new();
448        let mut items: Vec<Value> = entries
449            .iter()
450            .filter(|(_, e)| e.session_id == session_id)
451            .map(|(handle, entry)| {
452                live_handles.insert(handle.clone());
453                let st = entry.status.lock().unwrap().clone();
454                let out = entry.output.lock().unwrap();
455                let mut fields = vec![
456                    ("handle".into(), Value::Str(handle.clone())),
457                    ("status".into(), Value::Str(st.kind().into())),
458                    ("started_at".into(), Value::Int(st.started_at())),
459                    ("live".into(), Value::Bool(true)),
460                ];
461                if let Some(ec) = st.exit_code() {
462                    fields.push(("exit_code".into(), Value::Int(ec as i64)));
463                }
464                fields.push(("bytes_total".into(), Value::Int(out.total_bytes as i64)));
465                Value::Struct(fields)
466            })
467            .collect();
468
469        if all {
470            if let Some(dir) = session_dir {
471                if let Ok(rd) = std::fs::read_dir(dir) {
472                    for entry in rd.flatten() {
473                        let name = entry.file_name();
474                        let name = name.to_string_lossy();
475                        let Some(rest) = name
476                            .strip_prefix("bg_")
477                            .and_then(|s| s.strip_suffix(".log"))
478                        else {
479                            continue;
480                        };
481                        let handle: String = rest.to_string();
482                        if live_handles.contains(&handle) {
483                            continue;
484                        }
485                        let Ok(meta) = entry.metadata() else {
486                            continue;
487                        };
488                        let modified = meta
489                            .modified()
490                            .ok()
491                            .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
492                            .map(|d| d.as_millis() as i64)
493                            .unwrap_or(0);
494                        items.push(Value::Struct(vec![
495                            ("handle".into(), Value::Str(handle)),
496                            ("status".into(), Value::Str("exited".into())),
497                            ("started_at".into(), Value::Int(modified)),
498                            ("live".into(), Value::Bool(false)),
499                            ("bytes_total".into(), Value::Int(meta.len() as i64)),
500                        ]));
501                    }
502                }
503            }
504        }
505
506        Value::List(items)
507    }
508}
509
510impl Drop for BgRegistry {
511    fn drop(&mut self) {
512        self.kill_all();
513    }
514}
515
516#[allow(clippy::too_many_arguments)]
517async fn run_bg_process(
518    handle_str: String,
519    cmd: String,
520    timeout: Option<Duration>,
521    max_output_bytes: u64,
522    log_path: std::path::PathBuf,
523    status: Arc<Mutex<BgStatus>>,
524    output: Arc<Mutex<BgOutput>>,
525    mut control_rx: mpsc::Receiver<BgControl>,
526    cancel: CancellationToken,
527    registry: Arc<BgRegistry>,
528    stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
529    handle_for_stream: String,
530    task_registry: Option<TaskRegistry>,
531    task_id: Option<crate::task_registry::TaskId>,
532) {
533    let started_at = now_ms();
534    let mut command = tokio::process::Command::new("sh");
535    command
536        .arg("-c")
537        .arg(&cmd)
538        .stdin(Stdio::null())
539        .stdout(Stdio::piped())
540        .stderr(Stdio::piped());
541    let mut child: AsyncGroupChild = match command.group().kill_on_drop(true).spawn() {
542        Ok(c) => c,
543        Err(e) => {
544            *status.lock().unwrap() = BgStatus::Failed {
545                error: format!("spawn: {e}"),
546                started_at,
547                ended_at: now_ms(),
548            };
549            registry.remove(&handle_str);
550            return;
551        }
552    };
553    let pid = child.id();
554    *status.lock().unwrap() = BgStatus::Running {
555        pid: pid.unwrap_or(0),
556        started_at,
557    };
558
559    let stdout = child.inner().stdout.take();
560    let stderr = child.inner().stderr.take();
561
562    let stdout_reader = stdout.map(|s| {
563        let output = output.clone();
564        let log_path = log_path.clone();
565        let stx = stream_tx.clone();
566        let h = handle_for_stream.clone();
567        tokio::spawn(read_stream(
568            BufReader::new(s),
569            output,
570            log_path,
571            StreamKind::Stdout,
572            max_output_bytes,
573            stx,
574            h,
575        ))
576    });
577    let stderr_reader = stderr.map(|s| {
578        let output = output.clone();
579        let log_path = log_path.clone();
580        let stx = stream_tx.clone();
581        let h = handle_for_stream.clone();
582        tokio::spawn(read_stream(
583            BufReader::new(s),
584            output,
585            log_path,
586            StreamKind::Stderr,
587            max_output_bytes,
588            stx,
589            h,
590        ))
591    });
592
593    let exit_reason = tokio::select! {
594        biased;
595        _ = cancel.cancelled() => ExitReason::Cancelled,
596        ctrl = control_rx.recv() => {
597            match ctrl {
598                Some(BgControl::Kill) => ExitReason::Kill,
599                None => ExitReason::Natural,
600            }
601        }
602        _ = async {
603            if let Some(t) = timeout {
604                tokio::time::sleep(t).await;
605            } else {
606                std::future::pending::<()>().await;
607            }
608        } => ExitReason::Timeout,
609        s = child.wait() => ExitReason::Exited(s),
610    };
611
612    let ended_at = now_ms();
613    let final_status = match &exit_reason {
614        ExitReason::Exited(Ok(s)) => BgStatus::Exited {
615            exit_code: s.code().unwrap_or(-1),
616            started_at,
617            ended_at,
618        },
619        ExitReason::Timeout => {
620            let _ = child.start_kill();
621            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
622            BgStatus::TimedOut {
623                started_at,
624                ended_at,
625            }
626        }
627        ExitReason::Kill => {
628            let _ = child.start_kill();
629            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
630            BgStatus::Killed {
631                started_at,
632                ended_at,
633            }
634        }
635        ExitReason::Cancelled => {
636            let _ = child.start_kill();
637            let _ = tokio::time::timeout(Duration::from_millis(500), child.wait()).await;
638            BgStatus::Killed {
639                started_at,
640                ended_at,
641            }
642        }
643        ExitReason::Exited(Err(_)) => {
644            let _ = child.start_kill();
645            BgStatus::Failed {
646                error: "wait failed".into(),
647                started_at,
648                ended_at,
649            }
650        }
651        ExitReason::Natural => {
652            let s = child.wait().await;
653            BgStatus::Exited {
654                exit_code: s.ok().and_then(|s| s.code()).unwrap_or(-1),
655                started_at,
656                ended_at: now_ms(),
657            }
658        }
659    };
660
661    if let Some(r) = stdout_reader {
662        let _ = tokio::time::timeout(IO_DRAIN_TIMEOUT, r).await;
663    }
664    if let Some(r) = stderr_reader {
665        let _ = tokio::time::timeout(IO_DRAIN_TIMEOUT, r).await;
666    }
667
668    let exit_code = match &final_status {
669        BgStatus::Exited { exit_code, .. } => Some(*exit_code),
670        _ => None,
671    };
672    *status.lock().unwrap() = final_status.clone();
673
674    if let Some(tx) = &stream_tx {
675        let _ = tx.send(crate::stream::StreamFrame::BashExited {
676            handle: handle_for_stream,
677            exit_code,
678        });
679    }
680
681    if let (Some(tr), Some(tid)) = (task_registry, task_id) {
682        let ts = match &final_status {
683            BgStatus::Exited { exit_code, .. } if *exit_code == 0 => TaskStatus::Ok,
684            BgStatus::Killed { .. } | BgStatus::TimedOut { .. } => TaskStatus::Killed,
685            _ => TaskStatus::Err,
686        };
687        tr.finish(&tid, ts);
688    }
689}
690
691async fn read_stream<R: tokio::io::AsyncBufRead + Unpin>(
692    mut reader: R,
693    output: Arc<Mutex<BgOutput>>,
694    log_path: std::path::PathBuf,
695    kind: StreamKind,
696    max_output_bytes: u64,
697    stream_tx: Option<tokio::sync::broadcast::Sender<crate::stream::StreamFrame>>,
698    handle: String,
699) {
700    let prefix: &[u8] = match kind {
701        StreamKind::Stdout => b"[out] ",
702        StreamKind::Stderr => b"[err] ",
703    };
704    let kind_str = match kind {
705        StreamKind::Stdout => "stdout",
706        StreamKind::Stderr => "stderr",
707    };
708    let mut buf = String::new();
709    let mut log_file = tokio::fs::OpenOptions::new()
710        .append(true)
711        .create(true)
712        .open(&log_path)
713        .await
714        .ok();
715    loop {
716        buf.clear();
717        match reader.read_line(&mut buf).await {
718            Ok(0) => break,
719            Ok(_) => {
720                let data = buf.as_bytes();
721                {
722                    let mut out = output.lock().unwrap();
723                    out.push(kind, data, max_output_bytes);
724                }
725                if let Some(file) = log_file.as_mut() {
726                    let _ = file.write_all(prefix).await;
727                    let _ = file.write_all(data).await;
728                    if !data.ends_with(b"\n") {
729                        let _ = file.write_all(b"\n").await;
730                    }
731                }
732                if let Some(tx) = &stream_tx {
733                    let _ = tx.send(crate::stream::StreamFrame::BashChunk {
734                        handle: handle.clone(),
735                        kind: kind_str.to_string(),
736                        line: buf.clone(),
737                    });
738                }
739            }
740            Err(_) => break,
741        }
742    }
743}
744
745enum ExitReason {
746    Exited(std::io::Result<std::process::ExitStatus>),
747    Timeout,
748    Kill,
749    Cancelled,
750    Natural,
751}
752
753fn now_ms() -> i64 {
754    chrono::Utc::now().timestamp_millis()
755}
756
757pub struct BashSpawn;
758
759impl Tool for BashSpawn {
760    fn name(&self) -> &str {
761        "bash.spawn"
762    }
763
764    fn tier(&self) -> Tier {
765        Tier::Four
766    }
767
768    fn description(&self) -> Option<&str> {
769        Some(
770            "Run a shell command via `sh -c`.\n\n\
771block=false (default): command runs in background, returns immediately with a\n\
772handle. The command keeps running — use bash.output to read its output later,\n\
773bash.status to check if it finished, bash.kill to stop it. Use this for:\n\
774- long-running commands (servers, watchers)\n\
775- commands where you need to check output incrementally\n\
776- when you want to do other things while the command runs\n\n\
777block=true: waits for the command to finish, then returns stdout/stderr/exit_code.\n\
778Use block_timeout_ms to set a max wait (default 30s). Use this for:\n\
779- short commands where you need the result immediately (ls, git status, echo)\n\
780- commands that finish quickly\n\n\
781Do NOT use `sleep` in your command to wait — use block=true with block_timeout_ms\n\
782instead, or use the sleep tool to pause the workflow.",
783        )
784    }
785
786    fn input_schema(&self) -> serde_json::Value {
787        serde_json::json!({
788            "type": "object",
789            "properties": {
790                "cmd": {"type": "string", "description": "Shell command line."},
791                "block": {"type": "boolean", "default": false, "description": "If true, wait for process to exit before returning."},
792                "block_timeout_ms": {"type": "integer", "description": "Only with block=true. Max wait. 0 = no timeout. Default 30000."},
793                "timeout_ms": {"type": "integer", "description": "Process kill timeout in ms. Default 1800000 (30min). 0 = no timeout."},
794                "max_output_bytes": {"type": "integer", "description": "Max combined output bytes. Default 10485760 (10MB)."}
795            },
796            "required": ["cmd"]
797        })
798    }
799
800    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
801        Box::pin(async move {
802            let cmd = extract_string(&args, "cmd", 0)?;
803            let block = args
804                .named("block")
805                .and_then(|v| {
806                    if let Value::Bool(b) = v {
807                        Some(*b)
808                    } else {
809                        None
810                    }
811                })
812                .unwrap_or(false);
813            let block_timeout_ms = extract_optional_int(&args, "block_timeout_ms")
814                .map(|v| v as u64)
815                .unwrap_or(30_000);
816            let timeout_ms = extract_optional_int(&args, "timeout_ms").map(|v| v as u64);
817            let max_output = extract_optional_int(&args, "max_output_bytes")
818                .map(|v| v as u64)
819                .unwrap_or(DEFAULT_MAX_OUTPUT_BYTES);
820            let registry = ctx.bg_registry.clone().ok_or_else(|| {
821                RuntimeError::ToolFailed("bash.spawn: registry not available".into())
822            })?;
823            let handle_str = registry.spawn(cmd, timeout_ms, max_output, ctx)?;
824
825            if !block {
826                return Ok(handle_str);
827            }
828
829            let handle_s = handle_str
830                .field("handle")
831                .and_then(|v| {
832                    if let Value::Str(s) = v {
833                        Some(s.clone())
834                    } else {
835                        None
836                    }
837                })
838                .ok_or_else(|| {
839                    RuntimeError::ToolFailed("bash.spawn: missing handle field".into())
840                })?;
841            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".into());
842
843            let deadline = if block_timeout_ms == 0 {
844                None
845            } else {
846                Some(tokio::time::Instant::now() + Duration::from_millis(block_timeout_ms))
847            };
848            loop {
849                let entry = registry.lookup(&handle_s, &session_id)?;
850                let finished = {
851                    let st = entry.status.lock().unwrap();
852                    st.is_finished()
853                };
854                if finished {
855                    break;
856                }
857                if let Some(d) = deadline {
858                    if tokio::time::Instant::now() >= d {
859                        break;
860                    }
861                }
862                tokio::time::sleep(Duration::from_millis(50)).await;
863            }
864
865            let entry = registry.lookup(&handle_s, &session_id)?;
866            let st = entry.status.lock().unwrap().clone();
867            let out = entry.output.lock().unwrap();
868            let combined = String::from_utf8_lossy(&out.combined).into_owned();
869            Ok(Value::Struct(vec![
870                ("handle".into(), Value::Str(handle_s)),
871                ("status".into(), Value::Str(st.kind().into())),
872                (
873                    "exit_code".into(),
874                    st.exit_code()
875                        .map(|c| Value::Int(c as i64))
876                        .unwrap_or(Value::Unit),
877                ),
878                ("output".into(), Value::Str(combined)),
879                ("bytes_total".into(), Value::Int(out.total_bytes as i64)),
880            ]))
881        })
882    }
883}
884
885pub struct BashStatus;
886
887impl Tool for BashStatus {
888    fn name(&self) -> &str {
889        "bash.status"
890    }
891
892    fn tier(&self) -> Tier {
893        Tier::Four
894    }
895
896    fn description(&self) -> Option<&str> {
897        Some("Check the status of a background bash process.")
898    }
899
900    fn input_schema(&self) -> serde_json::Value {
901        serde_json::json!({
902            "type": "object",
903            "properties": {"handle": {"type": "string"}},
904            "required": ["handle"]
905        })
906    }
907
908    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
909        Box::pin(async move {
910            let handle = extract_string(&args, "handle", 0)?;
911            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
912            let registry = ctx.bg_registry.clone().ok_or_else(|| {
913                RuntimeError::ToolFailed("bash.status: registry not available".into())
914            })?;
915            registry.status(&handle, &session_id)
916        })
917    }
918}
919
920pub struct BashOutput;
921
922impl Tool for BashOutput {
923    fn name(&self) -> &str {
924        "bash.output"
925    }
926
927    fn tier(&self) -> Tier {
928        Tier::Four
929    }
930
931    fn description(&self) -> Option<&str> {
932        Some("Read output from a background bash process by byte cursor.")
933    }
934
935    fn input_schema(&self) -> serde_json::Value {
936        serde_json::json!({
937            "type": "object",
938            "properties": {
939                "handle": {"type": "string"},
940                "cursor": {"type": "integer", "description": "Byte offset to start reading. Default 0."},
941                "limit_bytes": {"type": "integer", "description": "Max bytes to return. Default 32000."}
942            },
943            "required": ["handle"]
944        })
945    }
946
947    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
948        Box::pin(async move {
949            let handle = extract_string(&args, "handle", 0)?;
950            let cursor = extract_optional_int(&args, "cursor").unwrap_or(0).max(0) as usize;
951            let limit = extract_optional_int(&args, "limit_bytes")
952                .unwrap_or(DEFAULT_OUTPUT_LIMIT as i64)
953                .max(1) as usize;
954            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
955            let registry = ctx.bg_registry.clone().ok_or_else(|| {
956                RuntimeError::ToolFailed("bash.output: registry not available".into())
957            })?;
958            registry.output(
959                &handle,
960                &session_id,
961                ctx.session_dir.as_deref(),
962                cursor,
963                limit,
964            )
965        })
966    }
967}
968
969pub struct BashKill;
970
971impl Tool for BashKill {
972    fn name(&self) -> &str {
973        "bash.kill"
974    }
975
976    fn tier(&self) -> Tier {
977        Tier::Four
978    }
979
980    fn description(&self) -> Option<&str> {
981        Some(
982            "Kill a background bash process. signal=term (default) sends SIGTERM, signal=kill sends SIGKILL.",
983        )
984    }
985
986    fn input_schema(&self) -> serde_json::Value {
987        serde_json::json!({
988            "type": "object",
989            "properties": {
990                "handle": {"type": "string"},
991                "signal": {"type": "string", "enum": ["term", "kill"], "description": "Default term."}
992            },
993            "required": ["handle"]
994        })
995    }
996
997    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
998        Box::pin(async move {
999            let handle = extract_string(&args, "handle", 0)?;
1000            let _ = extract_string(&args, "signal", 1);
1001            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
1002            let registry = ctx.bg_registry.clone().ok_or_else(|| {
1003                RuntimeError::ToolFailed("bash.kill: registry not available".into())
1004            })?;
1005            registry.kill(&handle, &session_id)
1006        })
1007    }
1008}
1009
1010pub struct BashList;
1011
1012impl Tool for BashList {
1013    fn name(&self) -> &str {
1014        "bash.list"
1015    }
1016
1017    fn tier(&self) -> Tier {
1018        Tier::Four
1019    }
1020
1021    fn description(&self) -> Option<&str> {
1022        Some(
1023            "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).",
1024        )
1025    }
1026
1027    fn input_schema(&self) -> serde_json::Value {
1028        serde_json::json!({
1029            "type": "object",
1030            "properties": {
1031                "all": {"type": "boolean", "description": "Include historical processes (default false)."}
1032            }
1033        })
1034    }
1035
1036    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1037        Box::pin(async move {
1038            let all = extract_optional_bool(&args, "all").unwrap_or(false);
1039            let session_id = ctx.session_id.clone().unwrap_or_else(|| "anon".to_string());
1040            let registry = ctx.bg_registry.clone().ok_or_else(|| {
1041                RuntimeError::ToolFailed("bash.list: registry not available".into())
1042            })?;
1043            Ok(registry.list(&session_id, ctx.session_dir.as_deref(), all))
1044        })
1045    }
1046}
1047
1048fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
1049    let value = match args.named(name) {
1050        Some(v) => v,
1051        None => args.positional(pos)?,
1052    };
1053    match value {
1054        Value::Str(s) => Ok(s.clone()),
1055        other => Err(RuntimeError::TypeMismatch {
1056            expected: "string".into(),
1057            actual: other.kind_name().into(),
1058        }),
1059    }
1060}
1061
1062fn extract_optional_int(args: &ToolArgs, name: &str) -> Option<i64> {
1063    match args.named(name)? {
1064        Value::Int(n) => Some(*n),
1065        _ => None,
1066    }
1067}
1068
1069fn extract_optional_bool(args: &ToolArgs, name: &str) -> Option<bool> {
1070    match args.named(name)? {
1071        Value::Bool(b) => Some(*b),
1072        _ => None,
1073    }
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078    use super::*;
1079    use crate::tool::{ToolArgs, ToolCtx};
1080    use crate::value::Value;
1081    use std::sync::Arc;
1082    use tempfile::TempDir;
1083
1084    fn ctx_with_registry(registry: Arc<BgRegistry>, dir: &std::path::Path) -> ToolCtx {
1085        let mut ctx = ToolCtx::new();
1086        ctx.bg_registry = Some(registry);
1087        ctx.session_dir = Some(dir.to_path_buf());
1088        ctx.session_id = Some("test-session".to_string());
1089        ctx
1090    }
1091
1092    #[test]
1093    fn handle_parse_roundtrip() {
1094        let h = BgHandle {
1095            session_id: "abc".into(),
1096            local_id: 42,
1097        };
1098        let s = h.to_string();
1099        assert_eq!(s, "bg_abc_42");
1100        let back = BgHandle::parse(&s).unwrap();
1101        assert_eq!(back, h);
1102    }
1103
1104    #[test]
1105    fn handle_parse_rejects_bad_format() {
1106        assert!(BgHandle::parse("not_bg").is_none());
1107        assert!(BgHandle::parse("bg_nosuffix").is_none());
1108        assert!(BgHandle::parse("bg_x_notnum").is_none());
1109    }
1110
1111    #[tokio::test]
1112    async fn spawn_returns_immediately_with_running_status() {
1113        let registry = Arc::new(BgRegistry::new());
1114        let dir = TempDir::new().unwrap();
1115        let ctx = ctx_with_registry(registry.clone(), dir.path());
1116        let args = ToolArgs {
1117            positional: vec![Value::Str("echo hello".into())],
1118            named: vec![],
1119        };
1120        let v = BashSpawn.call(args, &ctx).await.unwrap();
1121        let Value::Struct(fields) = v else {
1122            panic!("expected struct")
1123        };
1124        let handle = fields
1125            .iter()
1126            .find(|(k, _)| k == "handle")
1127            .and_then(|(_, v)| {
1128                if let Value::Str(s) = v {
1129                    Some(s.clone())
1130                } else {
1131                    None
1132                }
1133            })
1134            .unwrap();
1135        assert!(handle.starts_with("bg_"));
1136        let status_val = fields.iter().find(|(k, _)| k == "status").unwrap();
1137        assert!(matches!(&status_val.1, Value::Str(s) if s == "running"));
1138    }
1139
1140    #[tokio::test]
1141    async fn spawn_then_status_reaches_exited() {
1142        let registry = Arc::new(BgRegistry::new());
1143        let dir = TempDir::new().unwrap();
1144        let ctx = ctx_with_registry(registry.clone(), dir.path());
1145        let spawn_args = ToolArgs {
1146            positional: vec![Value::Str("echo hello".into())],
1147            named: vec![],
1148        };
1149        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
1150        let Value::Struct(fields) = v else { panic!() };
1151        let handle = fields
1152            .iter()
1153            .find(|(k, _)| k == "handle")
1154            .and_then(|(_, v)| {
1155                if let Value::Str(s) = v {
1156                    Some(s.clone())
1157                } else {
1158                    None
1159                }
1160            })
1161            .unwrap();
1162
1163        for _ in 0..50 {
1164            tokio::time::sleep(Duration::from_millis(50)).await;
1165            let status_args = ToolArgs {
1166                positional: vec![Value::Str(handle.clone())],
1167                named: vec![],
1168            };
1169            let s = BashStatus.call(status_args, &ctx).await.unwrap();
1170            if let Value::Struct(sf) = s {
1171                let kind = sf.iter().find(|(k, _)| k == "status").unwrap();
1172                if matches!(&kind.1, Value::Str(s) if s == "exited") {
1173                    let ec = sf.iter().find(|(k, _)| k == "exit_code").unwrap();
1174                    assert!(matches!(ec.1, Value::Int(0)));
1175                    return;
1176                }
1177            }
1178        }
1179        panic!("process did not exit in time");
1180    }
1181
1182    #[tokio::test]
1183    async fn spawn_output_captures_stdout() {
1184        let registry = Arc::new(BgRegistry::new());
1185        let dir = TempDir::new().unwrap();
1186        let ctx = ctx_with_registry(registry.clone(), dir.path());
1187        let spawn_args = ToolArgs {
1188            positional: vec![Value::Str("echo line1; echo line2".into())],
1189            named: vec![],
1190        };
1191        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
1192        let Value::Struct(fields) = v else { panic!() };
1193        let handle = fields
1194            .iter()
1195            .find(|(k, _)| k == "handle")
1196            .and_then(|(_, v)| {
1197                if let Value::Str(s) = v {
1198                    Some(s.clone())
1199                } else {
1200                    None
1201                }
1202            })
1203            .unwrap();
1204
1205        tokio::time::sleep(Duration::from_millis(300)).await;
1206
1207        let out_args = ToolArgs {
1208            positional: vec![Value::Str(handle.clone())],
1209            named: vec![],
1210        };
1211        let o = BashOutput.call(out_args, &ctx).await.unwrap();
1212        let Value::Struct(of) = o else { panic!() };
1213        let chunk = of.iter().find(|(k, _)| k == "chunk").unwrap();
1214        if let Value::Str(s) = &chunk.1 {
1215            assert!(s.contains("line1"), "chunk should contain line1: {s}");
1216            assert!(s.contains("line2"), "chunk should contain line2: {s}");
1217        } else {
1218            panic!("chunk not str");
1219        }
1220    }
1221
1222    #[tokio::test]
1223    async fn kill_terminates_long_running_process() {
1224        let registry = Arc::new(BgRegistry::new());
1225        let dir = TempDir::new().unwrap();
1226        let ctx = ctx_with_registry(registry.clone(), dir.path());
1227        let spawn_args = ToolArgs {
1228            positional: vec![Value::Str("sleep 100".into())],
1229            named: vec![],
1230        };
1231        let v = BashSpawn.call(spawn_args, &ctx).await.unwrap();
1232        let Value::Struct(fields) = v else { panic!() };
1233        let handle = fields
1234            .iter()
1235            .find(|(k, _)| k == "handle")
1236            .and_then(|(_, v)| {
1237                if let Value::Str(s) = v {
1238                    Some(s.clone())
1239                } else {
1240                    None
1241                }
1242            })
1243            .unwrap();
1244
1245        let kill_args = ToolArgs {
1246            positional: vec![Value::Str(handle.clone())],
1247            named: vec![],
1248        };
1249        BashKill.call(kill_args, &ctx).await.unwrap();
1250
1251        for _ in 0..50 {
1252            tokio::time::sleep(Duration::from_millis(50)).await;
1253            let status_args = ToolArgs {
1254                positional: vec![Value::Str(handle.clone())],
1255                named: vec![],
1256            };
1257            let s = BashStatus.call(status_args, &ctx).await.unwrap();
1258            if let Value::Struct(sf) = s {
1259                let kind = sf.iter().find(|(k, _)| k == "status").unwrap();
1260                if matches!(&kind.1, Value::Str(s) if s == "killed") {
1261                    return;
1262                }
1263            }
1264        }
1265        panic!("process not killed in time");
1266    }
1267
1268    #[tokio::test]
1269    async fn cross_session_access_rejected() {
1270        let registry = Arc::new(BgRegistry::new());
1271        let dir = TempDir::new().unwrap();
1272        let ctx_a = ctx_with_registry(registry.clone(), dir.path());
1273
1274        let spawn_args = ToolArgs {
1275            positional: vec![Value::Str("sleep 10".into())],
1276            named: vec![],
1277        };
1278        let v = BashSpawn.call(spawn_args, &ctx_a).await.unwrap();
1279        let Value::Struct(fields) = v else { panic!() };
1280        let handle = fields
1281            .iter()
1282            .find(|(k, _)| k == "handle")
1283            .and_then(|(_, v)| {
1284                if let Value::Str(s) = v {
1285                    Some(s.clone())
1286                } else {
1287                    None
1288                }
1289            })
1290            .unwrap();
1291
1292        let mut ctx_b = ToolCtx::new();
1293        ctx_b.bg_registry = Some(registry.clone());
1294        ctx_b.session_dir = Some(dir.path().to_path_buf());
1295        ctx_b.session_id = Some("other-session".to_string());
1296        let status_args = ToolArgs {
1297            positional: vec![Value::Str(handle)],
1298            named: vec![],
1299        };
1300        let err = BashStatus.call(status_args, &ctx_b).await.err().unwrap();
1301        assert!(format!("{err}").contains("does not belong to session"));
1302    }
1303}