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