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