Skip to main content

atman_runtime/tools/
bash_bg.rs

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