bamboo_tools/tools/
exec_command.rs

1//! Yielding shell execution whose results are pulled through ordinary tools.
2use async_trait::async_trait;
3use bamboo_agent_core::{Tool, ToolCtx, ToolError, ToolOutcome, ToolResult};
4use serde::{Deserialize, Serialize};
5use serde_json::json;
6use std::path::{Path, PathBuf};
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Arc, OnceLock};
9use tokio::io::{AsyncRead, AsyncReadExt};
10use tokio::sync::{watch, Mutex, Notify, OwnedSemaphorePermit, Semaphore};
11use tokio::time::{timeout, Duration, Instant};
12
13use super::{bash_runtime, workspace_state};
14
15const MAX_RETAINED_SESSIONS: usize = 64;
16const MAX_CAPTURE_BYTES: usize = 256 * 1024;
17const MAX_OUTPUT_TOKENS: usize = 8192;
18const DEFAULT_OUTPUT_TOKENS: usize = 2000;
19const MAX_COMMAND_BYTES: usize = 64 * 1024;
20// Advisory only: silence never changes process lifetime or terminal status.
21const QUIET_POLL_ADVISORY_AFTER: Duration = Duration::from_secs(90);
22const QUIET_POLL_ADVISORY_READOUTS: usize = 3;
23const QUIET_POLL_DIAGNOSIS_READOUTS: usize = 6;
24const QUIET_POLL_ADVISORY: &str = "No new output across repeated polls. Do other useful work or check a command-owned progress signal before polling again. Quiet output or an existing report does not prove completion. Continue waiting when appropriate, or report the concrete blocker; retain this handle for the final status.";
25
26#[derive(Debug)]
27struct CapturedOutput {
28    bytes: Vec<u8>,
29    omitted_bytes: usize,
30    completion: Option<(&'static str, Option<i32>)>,
31    error: Option<String>,
32    started_at: Instant,
33    last_output_at: Instant,
34    consecutive_empty_readouts: usize,
35}
36
37impl Default for CapturedOutput {
38    fn default() -> Self {
39        let now = Instant::now();
40        Self {
41            bytes: Vec::new(),
42            omitted_bytes: 0,
43            completion: None,
44            error: None,
45            started_at: now,
46            last_output_at: now,
47            consecutive_empty_readouts: 0,
48        }
49    }
50}
51
52/// Per-handle pull cursor and bounded capture. Process ownership remains in
53/// `ManagedBashChild`; this state can outlive a drained local resource scope.
54#[derive(Debug)]
55pub(crate) struct PolledShellState {
56    output: Mutex<CapturedOutput>,
57    completed: Notify,
58    // A reaped parent may still have descendants holding its output pipes.
59    // This observation does not publish a terminal capture result.
60    process_exited: AtomicBool,
61    #[cfg(test)]
62    fail_cleanup_after_reap: AtomicBool,
63    capture_stop: watch::Sender<bool>,
64    pub(crate) interaction: Mutex<()>,
65    _permit: OwnedSemaphorePermit,
66}
67
68impl PolledShellState {
69    fn new() -> Result<Arc<Self>, ToolError> {
70        static CAPACITY: OnceLock<Arc<Semaphore>> = OnceLock::new();
71        let permit = CAPACITY
72            .get_or_init(|| Arc::new(Semaphore::new(MAX_RETAINED_SESSIONS)))
73            .clone()
74            .try_acquire_owned()
75            .map_err(|_| ToolError::Execution("Exec session capacity exhausted; finish or stop an existing session before starting another command".into()))?;
76        Ok(Arc::new(Self {
77            output: Mutex::new(CapturedOutput::default()),
78            completed: Notify::new(),
79            process_exited: AtomicBool::new(false),
80            #[cfg(test)]
81            fail_cleanup_after_reap: AtomicBool::new(false),
82            capture_stop: watch::channel(false).0,
83            interaction: Mutex::new(()),
84            _permit: permit,
85        }))
86    }
87
88    async fn push(&self, bytes: &[u8]) {
89        if bytes.is_empty() {
90            return;
91        }
92        let mut output = self.output.lock().await;
93        output.last_output_at = Instant::now();
94        output.consecutive_empty_readouts = 0;
95        output.bytes.extend_from_slice(bytes);
96        let mut excess = output.bytes.len().saturating_sub(MAX_CAPTURE_BYTES);
97        if excess > 0 {
98            while excess < output.bytes.len() && output.bytes[excess] & 0b1100_0000 == 0b1000_0000 {
99                excess += 1;
100            }
101            output.bytes.drain(..excess);
102            output.omitted_bytes = output.omitted_bytes.saturating_add(excess);
103        }
104    }
105
106    pub(crate) async fn record_error(&self, error: String) {
107        self.output.lock().await.error.get_or_insert(error);
108    }
109
110    pub(crate) async fn record_cleanup_error(&self, error: String) {
111        // A real wait/group-drain failure takes priority over an earlier
112        // capture timeout so the retained continuation explains failed cleanup.
113        self.output.lock().await.error = Some(error);
114    }
115
116    #[cfg(test)]
117    pub(super) fn test_cleanup_after_reap(&self) -> std::io::Result<()> {
118        if self.fail_cleanup_after_reap.swap(false, Ordering::AcqRel) {
119            Err(std::io::Error::other("controlled post-reap cleanup error"))
120        } else {
121            Ok(())
122        }
123    }
124
125    pub(crate) fn stop_capture(&self) {
126        self.capture_stop.send_replace(true);
127    }
128
129    pub(crate) async fn finish(&self, status: &'static str, exit_code: Option<i32>) {
130        let mut output = self.output.lock().await;
131        let status = if output.error.is_some() {
132            "error"
133        } else {
134            status
135        };
136        output.completion = Some((status, exit_code));
137        drop(output);
138        self.completed.notify_waiters();
139    }
140
141    pub(super) fn mark_process_exited(&self) {
142        self.process_exited.store(true, Ordering::Release);
143    }
144
145    pub(super) fn process_exited(&self) -> bool {
146        self.process_exited.load(Ordering::Acquire)
147    }
148
149    pub(super) async fn wait(&self, yield_time_ms: u64) {
150        let notified = self.completed.notified();
151        tokio::pin!(notified);
152        // Register before observing completion: even an immediate child exit
153        // cannot slip between the state check and the notification wait.
154        notified.as_mut().enable();
155        let running = self.output.lock().await.completion.is_none();
156        if running {
157            let _ = timeout(Duration::from_millis(yield_time_ms), notified).await;
158        }
159    }
160
161    pub(crate) async fn is_complete(&self) -> bool {
162        self.output.lock().await.completion.is_some()
163    }
164
165    pub(super) async fn reset_empty_readouts(&self) {
166        self.output.lock().await.consecutive_empty_readouts = 0;
167    }
168
169    async fn take(&self, shell_id: &str, max_output_tokens: Option<usize>) -> ExecSessionResult {
170        let mut output = self.output.lock().await;
171        Self::take_locked(
172            &mut output,
173            shell_id,
174            max_output_tokens,
175            self.process_exited(),
176        )
177    }
178
179    fn take_locked(
180        output: &mut CapturedOutput,
181        shell_id: &str,
182        max_output_tokens: Option<usize>,
183        exit_observed: bool,
184    ) -> ExecSessionResult {
185        // Count captured bytes, including bytes omitted by either cap, rather
186        // than the rendered string. A tiny response cap can hide real output.
187        let had_output = !output.bytes.is_empty() || output.omitted_bytes > 0;
188        let polling = output.completion.is_none().then(|| {
189            output.consecutive_empty_readouts = if had_output {
190                0
191            } else {
192                output.consecutive_empty_readouts.saturating_add(1)
193            };
194            let now = Instant::now();
195            let idle = now.duration_since(output.last_output_at);
196            ExecSessionProgress {
197                terminal_result_ttl_ms: bash_runtime::COMPLETED_SESSION_TTL_SECS * 1000,
198                state: if idle >= QUIET_POLL_ADVISORY_AFTER
199                    && output.consecutive_empty_readouts >= QUIET_POLL_DIAGNOSIS_READOUTS
200                {
201                    "needs_diagnosis"
202                } else {
203                    "waiting"
204                },
205                elapsed_ms: duration_ms(now.duration_since(output.started_at)),
206                output_idle_ms: duration_ms(idle),
207                consecutive_empty_readouts: output.consecutive_empty_readouts,
208                exit_observed,
209                advisory: (idle >= QUIET_POLL_ADVISORY_AFTER
210                    && output.consecutive_empty_readouts >= QUIET_POLL_ADVISORY_READOUTS)
211                    .then_some(QUIET_POLL_ADVISORY),
212            }
213        });
214        let bytes = std::mem::take(&mut output.bytes);
215        let decoded = String::from_utf8_lossy(&bytes);
216        let budget = max_output_tokens
217            .unwrap_or(DEFAULT_OUTPUT_TOKENS)
218            .clamp(1, MAX_OUTPUT_TOKENS)
219            .saturating_mul(4);
220        let end = json_output_prefix_end(&decoded, budget);
221        let omitted = std::mem::take(&mut output.omitted_bytes)
222            .saturating_add(decoded.len().saturating_sub(end));
223        let (status, exit_code) = output.completion.unwrap_or(("running", None));
224        ExecSessionResult {
225            session_id: output.completion.is_none().then(|| shell_id.to_owned()),
226            status,
227            output: decoded[..end].to_owned(),
228            exit_code,
229            output_omitted_bytes: omitted,
230            error: output.error.clone(),
231            polling,
232        }
233    }
234}
235
236/// Bound the serialized `output` string, including its JSON quotes and escapes.
237/// Protocol fields such as status and session_id remain outside this body budget.
238fn json_output_prefix_end(output: &str, budget: usize) -> usize {
239    let mut encoded_bytes = 2; // The surrounding JSON string quotes.
240    let mut end = 0;
241    for (offset, character) in output.char_indices() {
242        // Match serde_json's string escaping without allocating an encoded copy.
243        let character_bytes = match character {
244            '"' | '\\' | '\n' | '\r' | '\t' | '\u{8}' | '\u{c}' => 2,
245            '\u{0}'..='\u{1f}' => 6,
246            _ => character.len_utf8(),
247        };
248        if encoded_bytes + character_bytes > budget {
249            break;
250        }
251        encoded_bytes += character_bytes;
252        end = offset + character.len_utf8();
253    }
254    end
255}
256
257/// Fixed-size reads capture prompts without a newline and keep an unterminated
258/// line from growing a BufReader allocation without a bound.
259pub(crate) async fn pump_polled_stream(
260    mut reader: impl AsyncRead + Unpin,
261    state: Arc<PolledShellState>,
262) {
263    let mut buffer = [0; 8192];
264    let mut pending = Vec::new();
265    let mut stop = state.capture_stop.subscribe();
266    loop {
267        let stopped = *stop.borrow();
268        if stopped {
269            if !pending.is_empty() {
270                state
271                    .push(String::from_utf8_lossy(&pending).as_bytes())
272                    .await;
273            }
274            break;
275        }
276        let read = tokio::select! {
277            biased;
278            _ = stop.changed() => continue,
279            result = reader.read(&mut buffer) => result,
280        };
281        match read {
282            Ok(0) => {
283                if !pending.is_empty() {
284                    state
285                        .push(String::from_utf8_lossy(&pending).as_bytes())
286                        .await;
287                }
288                break;
289            }
290            Ok(count) => {
291                pending.extend_from_slice(&buffer[..count]);
292                let mut consumed = 0;
293                let mut decoded = String::new();
294                while consumed < pending.len() {
295                    match std::str::from_utf8(&pending[consumed..]) {
296                        Ok(text) => {
297                            decoded.push_str(text);
298                            consumed = pending.len();
299                        }
300                        Err(error) => {
301                            let valid_end = consumed + error.valid_up_to();
302                            decoded.push_str(
303                                std::str::from_utf8(&pending[consumed..valid_end]).unwrap(),
304                            );
305                            consumed = valid_end;
306                            if let Some(invalid_len) = error.error_len() {
307                                decoded.push('\u{fffd}');
308                                consumed += invalid_len;
309                            } else {
310                                break;
311                            }
312                        }
313                    }
314                }
315                pending.drain(..consumed);
316                state.push(decoded.as_bytes()).await;
317            }
318            Err(error) => {
319                state
320                    .record_error(format!("Shell output read failed: {error}"))
321                    .await;
322                break;
323            }
324        }
325    }
326}
327
328#[derive(Debug, Serialize)]
329struct ExecSessionResult {
330    #[serde(skip_serializing_if = "Option::is_none")]
331    session_id: Option<String>,
332    status: &'static str,
333    output: String,
334    #[serde(skip_serializing_if = "Option::is_none")]
335    exit_code: Option<i32>,
336    #[serde(skip_serializing_if = "is_zero")]
337    output_omitted_bytes: usize,
338    #[serde(skip_serializing_if = "Option::is_none")]
339    error: Option<String>,
340    #[serde(skip_serializing_if = "Option::is_none")]
341    polling: Option<ExecSessionProgress>,
342}
343
344/// Capture observations, never a liveness probe or evidence of a hung process.
345#[derive(Debug, Serialize)]
346struct ExecSessionProgress {
347    terminal_result_ttl_ms: u64,
348    state: &'static str,
349    elapsed_ms: u64,
350    output_idle_ms: u64,
351    consecutive_empty_readouts: usize,
352    // False means no exit has been observed, not that liveness was verified.
353    exit_observed: bool,
354    #[serde(skip_serializing_if = "Option::is_none")]
355    advisory: Option<&'static str>,
356}
357
358fn duration_ms(duration: Duration) -> u64 {
359    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
360}
361
362fn is_zero(value: &usize) -> bool {
363    *value == 0
364}
365
366pub(crate) fn yield_duration(requested: Option<u64>) -> u64 {
367    requested.unwrap_or(10_000).clamp(250, 30_000)
368}
369
370pub(crate) fn owner_session(ctx: &ToolCtx) -> Result<&str, ToolError> {
371    ctx.session_id().filter(|id| !id.is_empty()).ok_or_else(|| {
372        ToolError::InvalidArguments("Exec sessions require an owning session".into())
373    })
374}
375
376pub(crate) async fn poll_result(
377    shell: &bash_runtime::ShellSession,
378    yield_time_ms: u64,
379    max_output_tokens: Option<usize>,
380) -> ToolOutcome {
381    let state = shell.polled.as_ref().expect("validated exec session");
382    state.wait(yield_time_ms).await;
383    let result = state.take(&shell.id, max_output_tokens).await;
384    if result.session_id.is_none() {
385        bash_runtime::remove_shell(&shell.id);
386    }
387    result_outcome(&result)
388}
389
390fn result_outcome(result: &ExecSessionResult) -> ToolOutcome {
391    ToolOutcome::Completed(ToolResult {
392        // A command's nonzero exit is an observed outcome, not tool transport
393        // failure. Models inspect exit_code/status to decide their next step.
394        success: true,
395        result: serde_json::to_string(result).expect("exec result is serializable"),
396        display_preference: Some("Collapsible".into()),
397        images: Vec::new(),
398    })
399}
400
401/// Exists only during the first invocation. The completion owner, not Drop,
402/// performs and certifies cleanup before discarding an unpublished entry.
403struct InitialDelivery<'a> {
404    shell: &'a bash_runtime::ShellSession,
405    delivered: bool,
406}
407
408impl Drop for InitialDelivery<'_> {
409    fn drop(&mut self) {
410        if !self.delivered {
411            self.shell.cancel_initial();
412        }
413    }
414}
415
416async fn poll_initial_result(
417    shell: &bash_runtime::ShellSession,
418    yield_time_ms: u64,
419    max_output_tokens: Option<usize>,
420) -> Result<ToolOutcome, ToolError> {
421    let mut delivery = InitialDelivery {
422        shell,
423        delivered: false,
424    };
425    let state = shell.polled.as_ref().expect("validated initial exec");
426    if let Some(owner) = &shell.initial_handoff {
427        tokio::select! {
428            _ = state.wait(yield_time_ms) => {},
429            _ = owner.capture_ready() => {},
430        }
431    } else {
432        state.wait(yield_time_ms).await;
433    }
434    let result = take_initial_result(shell, max_output_tokens, &mut delivery, || {}).await?;
435    if result.session_id.is_none() {
436        bash_runtime::remove_shell(&shell.id);
437    }
438    Ok(result_outcome(&result))
439}
440
441/// The last await acquires capture, then ownership and its one output snapshot
442/// commit synchronously under that same lock. A writer cannot finish draining
443/// between a stale snapshot and terminal delivery. The observer is a private
444/// scheduling seam; it does not supply ownership readiness or an exit status.
445async fn take_initial_result(
446    shell: &bash_runtime::ShellSession,
447    max_output_tokens: Option<usize>,
448    delivery: &mut InitialDelivery<'_>,
449    before_commit: impl FnOnce(),
450) -> Result<ExecSessionResult, ToolError> {
451    let state = shell.polled.as_ref().expect("validated initial exec");
452    let mut output = state.output.lock().await;
453    before_commit();
454    // Delivery errors must leave the output cursor untouched for cleanup.
455    let terminal = shell.deliver_initial().map_err(ToolError::Execution)?;
456    let mut result = PolledShellState::take_locked(
457        &mut output,
458        &shell.id,
459        max_output_tokens,
460        state.process_exited(),
461    );
462    if let Some(status) = terminal {
463        result.session_id = None;
464        result.polling = None;
465        result.exit_code = status.code();
466        result.status = if result.error.is_some() {
467            "error"
468        } else if status.code().is_none() {
469            "killed"
470        } else {
471            "completed"
472        };
473    } else if shell.initial_handoff.is_some() {
474        // A not-yet-drained initial owner still has a retrievable result handle.
475        result.session_id = Some(shell.id.clone());
476        result.status = "running";
477        result.exit_code = None;
478    }
479    delivery.delivered = true;
480    drop(output);
481    Ok(result)
482}
483
484#[derive(Deserialize)]
485#[serde(deny_unknown_fields)]
486struct ExecCommandArgs {
487    cmd: String,
488    workdir: Option<String>,
489    yield_time_ms: Option<u64>,
490    max_output_tokens: Option<usize>,
491    #[serde(default)]
492    interactive: bool,
493}
494
495fn resolve_cwd(owner: &str, workdir: Option<&str>) -> Result<PathBuf, ToolError> {
496    let workspace = workspace_state::workspace_or_process_cwd(Some(owner));
497    let workspace = workspace.canonicalize().map_err(|error| {
498        ToolError::InvalidArguments(format!(
499            "Invalid session workspace '{}': {error}",
500            workspace.display()
501        ))
502    })?;
503    let requested = match workdir {
504        Some(value) if value.trim().is_empty() => {
505            return Err(ToolError::InvalidArguments(
506                "'workdir' cannot be empty".into(),
507            ));
508        }
509        Some(value) => Path::new(value.trim()).to_path_buf(),
510        None => workspace.clone(),
511    };
512    let resolved = if requested.is_absolute() {
513        requested
514    } else {
515        workspace.join(requested)
516    };
517    let canonical = resolved.canonicalize().map_err(|error| {
518        ToolError::InvalidArguments(format!("Invalid workdir '{}': {error}", resolved.display()))
519    })?;
520    if !canonical.is_dir() {
521        return Err(ToolError::InvalidArguments(
522            "workdir must be a directory".into(),
523        ));
524    }
525    if !canonical.starts_with(&workspace) {
526        return Err(ToolError::InvalidArguments(
527            "workdir must stay within the current session workspace".into(),
528        ));
529    }
530    Ok(canonical)
531}
532
533#[derive(Default)]
534pub struct ExecCommandTool;
535
536impl ExecCommandTool {
537    pub fn new() -> Self {
538        Self
539    }
540}
541
542#[async_trait]
543impl Tool for ExecCommandTool {
544    fn name(&self) -> &str {
545        "exec_command"
546    }
547
548    fn description(&self) -> &str {
549        "Execute a shell command and wait up to yield_time_ms. A running command returns a session_id and polling progress; continue other work, then call write_stdin with that ID and empty chars to retrieve incremental output and its final exit code. Follow a quiet polling advisory before repeating empty polls; silence does not prove completion. Completion is delivered through these tool results. interactive=true keeps a stdin pipe open for later input; this is a pipe, not a terminal/PTY. Output is bounded and reports omitted bytes."
550    }
551
552    fn parameters_schema(&self) -> serde_json::Value {
553        json!({"type":"object","properties":{
554            "cmd":{"type":"string","description":"Shell command to execute"},
555            "workdir":{"type":"string","description":"Working directory within the current session workspace, relative or absolute; paths and symlinks outside the workspace are rejected"},
556            "yield_time_ms":{"type":"integer","description":"Wait before yielding a running session; default 10000, clamped to 250..30000 ms"},
557            "max_output_tokens":{"type":"integer","description":"Output string budget, default 2000 and capped at 8192; estimated as four JSON-encoded bytes per token, including string quotes and escapes. Status/session/exit metadata is additional"},
558            "interactive":{"type":"boolean","description":"Keep a piped stdin open for write_stdin; default false sends immediate EOF"}
559        },"required":["cmd"],"additionalProperties":false})
560    }
561
562    async fn invoke(
563        &self,
564        args: serde_json::Value,
565        ctx: ToolCtx,
566    ) -> Result<ToolOutcome, ToolError> {
567        let parsed: ExecCommandArgs = serde_json::from_value(args).map_err(|error| {
568            ToolError::InvalidArguments(format!("Invalid exec_command args: {error}"))
569        })?;
570        let owner = owner_session(&ctx)?;
571        if parsed.cmd.trim().is_empty() || parsed.cmd.len() > MAX_COMMAND_BYTES {
572            return Err(ToolError::InvalidArguments(
573                "cmd must be nonempty and at most 65536 bytes".into(),
574            ));
575        }
576        let cwd = resolve_cwd(owner, parsed.workdir.as_deref())?;
577        let state = PolledShellState::new()?;
578        let shell = bash_runtime::spawn_polled(
579            &parsed.cmd,
580            &cwd,
581            None,
582            owner.to_owned(),
583            parsed.interactive,
584            state,
585        )
586        .await
587        .map_err(ToolError::Execution)?;
588        poll_initial_result(
589            &shell,
590            yield_duration(parsed.yield_time_ms),
591            parsed.max_output_tokens,
592        )
593        .await
594    }
595}
596
597#[cfg(test)]
598#[path = "exec_sessions_tests.rs"]
599mod tests;
600
601#[cfg(all(test, unix))]
602#[path = "initial_exec_cleanup_tests.rs"]
603mod initial_cleanup_tests;
bamboo_tools/tools/exec_command.rs

bamboo_tools/tools/
exec_command.rs
