use async_trait::async_trait;
use once_cell::sync::Lazy;
use serde_json::{Value, json};
use std::collections::{HashMap, VecDeque};
use std::process::Stdio;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use theway_core::{AgentTool, AgentToolError, AgentToolResult, AgentToolUpdate};
use theway_llm_provider::{Tool, UserContentBlock};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
const DEFAULT_FG_TIMEOUT_SECS: u64 = 60;
const MAX_OUTPUT_BYTES: usize = 200 * 1024;
const EXITED_KEEP_ALIVE: Duration = Duration::from_secs(60);
const EXIT_DRAIN_WINDOW: Duration = Duration::from_millis(250);
struct OutputBuffer {
chunks: VecDeque<String>,
bytes: usize,
dropped_chars: usize,
version: u64,
}
impl OutputBuffer {
fn new() -> Self {
Self {
chunks: VecDeque::new(),
bytes: 0,
dropped_chars: 0,
version: 0,
}
}
fn append(&mut self, mut chunk: String) {
self.version += 1;
let chunk = if chunk.len() > MAX_OUTPUT_BYTES {
let keep_start = tail_start(&chunk);
self.dropped_chars += chunk[..keep_start].chars().count();
chunk.split_off(keep_start)
} else {
chunk
};
self.bytes += chunk.len();
self.chunks.push_back(chunk);
while self.bytes > MAX_OUTPUT_BYTES {
if let Some(old) = self.chunks.pop_front() {
self.bytes -= old.len();
self.dropped_chars += old.chars().count();
}
}
}
fn snapshot(&self) -> (String, usize) {
let mut out = String::with_capacity(self.bytes.min(MAX_OUTPUT_BYTES));
for chunk in &self.chunks {
out.push_str(chunk);
}
(out, self.dropped_chars)
}
}
fn tail_start(chunk: &str) -> usize {
let mut start = 0;
for (i, _) in chunk.char_indices() {
if chunk.len() - i <= MAX_OUTPUT_BYTES {
start = i;
break;
}
}
start
}
struct ShellRegistry {
shells: Mutex<HashMap<String, Arc<ShellHandle>>>,
next_id: AtomicU64,
}
impl ShellRegistry {
fn insert(&self, id: String, handle: Arc<ShellHandle>) {
self.shells.lock().unwrap().insert(id, handle);
}
fn get(&self, id: &str) -> Option<Arc<ShellHandle>> {
self.shells.lock().unwrap().get(id).cloned()
}
fn remove(&self, id: &str) -> Option<Arc<ShellHandle>> {
self.shells.lock().unwrap().remove(id)
}
fn remove_if_exited(&self, id: &str) {
let mut guard = self.shells.lock().unwrap();
if let Some(handle) = guard.get(id) {
if handle.exited.load(Ordering::SeqCst) {
guard.remove(id);
}
}
}
fn ids(&self) -> Vec<String> {
self.shells.lock().unwrap().keys().cloned().collect()
}
}
static REGISTRY: OnceLock<Arc<ShellRegistry>> = OnceLock::new();
fn registry() -> &'static Arc<ShellRegistry> {
REGISTRY.get_or_init(|| {
Arc::new(ShellRegistry {
shells: Mutex::new(HashMap::new()),
next_id: AtomicU64::new(1),
})
})
}
fn next_shell_id() -> String {
format!(
"shell-{}",
registry().next_id.fetch_add(1, Ordering::SeqCst)
)
}
struct ShellHandle {
id: String,
pid: u32,
stdin: tokio::sync::Mutex<Option<tokio::process::ChildStdin>>,
stdout: Mutex<OutputBuffer>,
stderr: Mutex<OutputBuffer>,
notify: Notify,
exited: AtomicBool,
exit_code: Mutex<Option<i32>>,
killed: AtomicBool,
}
struct OutputSnapshot {
version: u64,
stdout: String,
stdout_dropped: usize,
stderr: String,
stderr_dropped: usize,
exited: bool,
exit_code: Option<i32>,
}
impl ShellHandle {
fn append_output(&self, stderr: bool, chunk: String) {
if stderr {
self.stderr.lock().unwrap().append(chunk);
} else {
self.stdout.lock().unwrap().append(chunk);
}
self.notify.notify_waiters();
}
fn mark_exited(&self, code: Option<i32>) {
self.exited.store(true, Ordering::SeqCst);
*self.exit_code.lock().unwrap() = code;
self.notify.notify_waiters();
let id = self.id.clone();
tokio::spawn(async move {
tokio::time::sleep(EXITED_KEEP_ALIVE).await;
registry().remove_if_exited(&id);
});
}
fn snapshot(&self) -> OutputSnapshot {
let (stdout, stdout_dropped) = self.stdout.lock().unwrap().snapshot();
let (stderr, stderr_dropped) = self.stderr.lock().unwrap().snapshot();
OutputSnapshot {
version: self.version(),
stdout,
stdout_dropped,
stderr,
stderr_dropped,
exited: self.exited.load(Ordering::SeqCst),
exit_code: *self.exit_code.lock().unwrap(),
}
}
fn version(&self) -> u64 {
self.stdout.lock().unwrap().version + self.stderr.lock().unwrap().version
}
async fn kill(&self) -> Result<(), AgentToolError> {
if self.killed.swap(true, Ordering::SeqCst) || self.exited.load(Ordering::SeqCst) {
return Ok(());
}
super::exec::process_group::kill(self.pid).await
}
}
pub struct BackgroundShell {
pub id: String,
pub pid: u32,
}
pub async fn run_in_background(command: &str) -> Result<BackgroundShell, AgentToolError> {
run_in_background_with_cwd(command, None).await
}
pub async fn run_in_background_with_cwd(
command: &str,
cwd: Option<&std::path::Path>,
) -> Result<BackgroundShell, AgentToolError> {
let mut cmd = tokio::process::Command::new(shell_program());
cmd.arg(shell_flag())
.arg(command)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
super::exec::process_group::prepare_command(&mut cmd);
let mut child = cmd
.spawn()
.map_err(|e| AgentToolError::from(format!("spawn: {e}")))?;
let pid = child
.id()
.ok_or_else(|| AgentToolError::from("spawned child has no pid"))?;
let id = next_shell_id();
let handle = Arc::new(ShellHandle {
id: id.clone(),
pid,
stdin: tokio::sync::Mutex::new(child.stdin.take()),
stdout: Mutex::new(OutputBuffer::new()),
stderr: Mutex::new(OutputBuffer::new()),
notify: Notify::new(),
exited: AtomicBool::new(false),
exit_code: Mutex::new(None),
killed: AtomicBool::new(false),
});
if let Some(pipe) = child.stdout.take() {
let h = handle.clone();
tokio::spawn(drain_pipe(pipe, h, false));
}
if let Some(pipe) = child.stderr.take() {
let h = handle.clone();
tokio::spawn(drain_pipe(pipe, h, true));
}
let h = handle.clone();
tokio::spawn(async move {
let status = child.wait().await;
h.mark_exited(status.ok().and_then(|s| s.code()));
});
registry().insert(id.clone(), handle);
Ok(BackgroundShell { id, pid })
}
async fn drain_pipe<R>(mut pipe: R, handle: Arc<ShellHandle>, stderr: bool)
where
R: AsyncRead + Unpin,
{
let mut buf = [0u8; 8192];
loop {
match pipe.read(&mut buf).await {
Ok(0) | Err(_) => break,
Ok(n) => handle.append_output(stderr, String::from_utf8_lossy(&buf[..n]).into_owned()),
}
}
}
fn shell_program() -> &'static str {
#[cfg(windows)]
{
"cmd"
}
#[cfg(not(windows))]
{
"sh"
}
}
fn shell_flag() -> &'static str {
#[cfg(windows)]
{
"/C"
}
#[cfg(not(windows))]
{
"-c"
}
}
async fn get_output_text(
handle: &ShellHandle,
timeout_secs: Option<u64>,
cancel: &CancellationToken,
) -> String {
let seen_version = handle.version();
loop {
let notified = handle.notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
let snap = handle.snapshot();
if snap.exited {
return render_snapshot(&handle.id, &drain_after_exit(handle, snap).await);
}
if snap.version != seen_version {
return render_snapshot(&handle.id, &snap);
}
let has_timeout = timeout_secs.is_some();
let timeout_future =
tokio::time::sleep(Duration::from_secs(timeout_secs.unwrap_or(u64::MAX / 2)));
tokio::pin!(timeout_future);
tokio::select! {
biased;
_ = cancel.cancelled() => return render_snapshot(&handle.id, &handle.snapshot()),
_ = &mut timeout_future, if has_timeout => return render_snapshot(&handle.id, &handle.snapshot()),
_ = &mut notified => {}
}
}
}
async fn drain_after_exit(handle: &ShellHandle, first: OutputSnapshot) -> OutputSnapshot {
let deadline = Instant::now() + EXIT_DRAIN_WINDOW;
let mut last = first;
loop {
tokio::time::sleep(Duration::from_millis(20)).await;
let snap = handle.snapshot();
if snap.version == last.version || Instant::now() >= deadline {
return snap;
}
last = snap;
}
}
fn render_snapshot(id: &str, snap: &OutputSnapshot) -> String {
let status = if snap.exited {
match snap.exit_code {
Some(code) => format!("exited (code {code})"),
None => "exited".to_string(),
}
} else {
"running".to_string()
};
let mut text = format!("[{id}] {status}\n\nstdout:\n{}", snap.stdout);
if snap.stdout_dropped > 0 {
text.push_str(&format!("\n…({} 字符, 截断)", snap.stdout_dropped));
}
text.push_str("\nstderr:\n");
text.push_str(&snap.stderr);
if snap.stderr_dropped > 0 {
text.push_str(&format!("\n…({} 字符, 截断)", snap.stderr_dropped));
}
text
}
fn decode_bytes_input(input: &str) -> String {
let mut out = String::with_capacity(input.len());
let mut rest = input;
while let Some(pos) = rest.find('<') {
out.push_str(&rest[..pos]);
rest = &rest[pos..];
if let Some(end) = rest.find('>') {
if let Some(byte) = decode_marker(&rest[1..end]) {
out.push(byte as char);
rest = &rest[end + 1..];
continue;
}
}
out.push('<');
rest = &rest[1..];
}
out.push_str(rest);
out
}
fn decode_marker(body: &str) -> Option<u8> {
match body {
"CR" => Some(b'\r'),
"LF" => Some(b'\n'),
"ESC" => Some(b'\x1b'),
"BS" => Some(b'\x7f'),
_ => {
let ctrl = body.strip_prefix("C-")?;
let mut chars = ctrl.chars();
let c = chars.next()?;
if chars.next().is_some() || !c.is_ascii_alphabetic() {
return None;
}
Some(c.to_ascii_uppercase() as u8 & 0x1F)
}
}
}
pub struct ExecTool;
pub struct GetOutputTool;
pub struct KillShellTool;
pub struct WriteToProcessTool;
#[async_trait]
impl AgentTool for ExecTool {
fn definition(&self) -> &Tool {
&EXEC_DEFINITION
}
fn label(&self) -> &str {
"exec"
}
async fn execute(
&self,
_tool_call_id: &str,
params: Value,
cancel: CancellationToken,
_on_update: Option<AgentToolUpdate>,
) -> Result<AgentToolResult, AgentToolError> {
let command = params
.get("command")
.and_then(|v| v.as_str())
.ok_or_else(|| AgentToolError::from("missing `command`"))?;
let background = params
.get("run_in_background")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let cwd = params.get("cwd").and_then(|v| v.as_str()).map(String::from);
if background {
let bg = run_in_background_with_cwd(command, cwd.as_deref().map(std::path::Path::new))
.await?;
let text = format!("background shell started: {} (pid {})", bg.id, bg.pid);
return Ok(AgentToolResult {
content: vec![UserContentBlock::text(text)],
details: json!({ "command": command, "shellId": bg.id, "pid": bg.pid }),
terminate: None,
});
}
let timeout_secs = Some(
params
.get("timeout")
.and_then(|v| v.as_u64())
.unwrap_or(DEFAULT_FG_TIMEOUT_SECS),
);
let outcome = super::exec::run_with_kill_on_timeout_or_cancel(
command,
timeout_secs.map(Duration::from_secs),
cwd.as_deref().map(std::path::Path::new),
None,
&cancel,
)
.await?;
let exit = outcome.rendered_exit();
let mut stderr_full = outcome.stderr;
if let Some(suffix) = &outcome.stderr_suffix {
if !stderr_full.is_empty() && !stderr_full.ends_with('\n') {
stderr_full.push('\n');
}
stderr_full.push_str(suffix);
}
let mut text = format!("$ {command}\n");
if !outcome.stdout.is_empty() {
text.push_str(&outcome.stdout);
if !outcome.stdout.ends_with('\n') {
text.push('\n');
}
}
if !stderr_full.is_empty() {
text.push_str(&stderr_full);
if !stderr_full.ends_with('\n') {
text.push('\n');
}
}
text.push_str(&format!("[exit {exit}]"));
Ok(AgentToolResult {
content: vec![UserContentBlock::text(text)],
details: json!({
"command": command,
"exitCode": exit,
"isError": exit != 0,
}),
terminate: None,
})
}
}
#[async_trait]
impl AgentTool for GetOutputTool {
fn definition(&self) -> &Tool {
&GET_OUTPUT_DEFINITION
}
fn label(&self) -> &str {
"Get Output"
}
async fn execute(
&self,
_id: &str,
params: Value,
cancel: CancellationToken,
_on_update: Option<AgentToolUpdate>,
) -> Result<AgentToolResult, AgentToolError> {
let shell_id = params
.get("shell_id")
.and_then(|v| v.as_str())
.ok_or_else(|| AgentToolError::from("missing `shell_id`"))?;
let timeout_secs = params.get("timeout").and_then(|v| v.as_u64());
let handle = registry().get(shell_id).ok_or_else(|| {
let available = registry().ids();
let available = if available.is_empty() {
"none".to_string()
} else {
available.join(", ")
};
AgentToolError::from(format!(
"Unknown shell_id: {shell_id}. Available: {available}"
))
})?;
let text = get_output_text(&handle, timeout_secs, &cancel).await;
Ok(AgentToolResult {
content: vec![UserContentBlock::text(text)],
details: json!({ "shellId": shell_id }),
terminate: None,
})
}
}
#[async_trait]
impl AgentTool for KillShellTool {
fn definition(&self) -> &Tool {
&KILL_SHELL_DEFINITION
}
fn label(&self) -> &str {
"Kill Shell"
}
async fn execute(
&self,
_id: &str,
params: Value,
_cancel: CancellationToken,
_on_update: Option<AgentToolUpdate>,
) -> Result<AgentToolResult, AgentToolError> {
let shell_id = params
.get("shell_id")
.and_then(|v| v.as_str())
.ok_or_else(|| AgentToolError::from("missing `shell_id`"))?;
let handle = registry()
.remove(shell_id)
.ok_or_else(|| AgentToolError::from(format!("Unknown shell_id: {shell_id}")))?;
handle.kill().await?;
Ok(AgentToolResult {
content: vec![UserContentBlock::text(format!("Killed {shell_id}"))],
details: json!({ "shellId": shell_id }),
terminate: None,
})
}
}
#[async_trait]
impl AgentTool for WriteToProcessTool {
fn definition(&self) -> &Tool {
&WRITE_TO_PROCESS_DEFINITION
}
fn label(&self) -> &str {
"Write To Process"
}
async fn execute(
&self,
_id: &str,
params: Value,
cancel: CancellationToken,
_on_update: Option<AgentToolUpdate>,
) -> Result<AgentToolResult, AgentToolError> {
let shell_id = params
.get("shell_id")
.and_then(|v| v.as_str())
.ok_or_else(|| AgentToolError::from("missing `shell_id`"))?;
let handle = registry()
.get(shell_id)
.ok_or_else(|| AgentToolError::from(format!("Unknown shell_id: {shell_id}")))?;
let input = match params.get("bytes_input").and_then(|v| v.as_str()) {
Some(bytes) => decode_bytes_input(bytes),
None => params
.get("text_input")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
};
let bytes_written = {
let mut guard = handle.stdin.lock().await;
match guard.as_mut() {
Some(stdin) => {
let write = stdin.write_all(input.as_bytes());
tokio::pin!(write);
tokio::select! {
biased;
_ = cancel.cancelled() => return Err(AgentToolError::from("write aborted")),
r = &mut write => r.map_err(|e| AgentToolError::from(format!("write error: {e}")))?,
}
input.len()
}
None => return Err(AgentToolError::from("shell stdin is closed")),
}
};
Ok(AgentToolResult {
content: vec![UserContentBlock::text(format!(
"Wrote {bytes_written} bytes to {shell_id}"
))],
details: json!({ "shellId": shell_id, "bytesWritten": bytes_written }),
terminate: None,
})
}
}
static EXEC_DEFINITION: Lazy<Tool> = Lazy::new(|| {
Tool {
name: "exec".into(),
description: "Execute a shell command. With `run_in_background: true` the command runs in a background shell and the tool immediately returns its shell_id — use get_output to read output, kill_shell to terminate, write_to_process to send input. Foreground mode behaves exactly like `bash` (captures stdout+stderr, optional `timeout` in seconds, kills the process tree on timeout/cancel)."
.into(),
parameters: json!({
"type": "object",
"properties": {
"command": { "type": "string", "description": "Shell command to execute" },
"run_in_background": { "type": "boolean", "description": "If true, run in background and return shell_id" },
"cwd": { "type": "string", "description": "Working directory to run the command in (absolute path). Optional; defaults to the session cwd" },
"timeout": { "type": "integer", "description": "Timeout in seconds (foreground only)" },
},
"required": ["command"],
}),
}
});
static GET_OUTPUT_DEFINITION: Lazy<Tool> = Lazy::new(|| {
Tool {
name: "get_output".into(),
description: "Read output from a background shell started with `run_in_background: true`. Blocks until new output arrives or the process exits; optional `timeout` in seconds caps the wait. Returns the accumulated stdout/stderr (tail-kept, ~200 KiB cap per stream with a truncation marker when bytes were dropped)."
.into(),
parameters: json!({
"type": "object",
"properties": {
"shell_id": { "type": "string", "description": "Background shell ID" },
"timeout": { "type": "integer", "description": "Max wait in seconds (optional; without it the tool waits until new output or exit)" },
},
"required": ["shell_id"],
}),
}
});
static KILL_SHELL_DEFINITION: Lazy<Tool> = Lazy::new(|| {
Tool {
name: "kill_shell".into(),
description: "Terminate a background shell (kills its whole process tree) and remove it from the registry."
.into(),
parameters: json!({
"type": "object",
"properties": {
"shell_id": { "type": "string", "description": "Background shell ID to kill" },
},
"required": ["shell_id"],
}),
}
});
static WRITE_TO_PROCESS_DEFINITION: Lazy<Tool> = Lazy::new(|| {
Tool {
name: "write_to_process".into(),
description: "Write input to a background shell's stdin (no newline is appended). `bytes_input` decodes markers: <CR>, <LF>, <ESC>, <BS>, <C-c> and other <C-x> control bytes; unknown markers are written literally."
.into(),
parameters: json!({
"type": "object",
"properties": {
"shell_id": { "type": "string", "description": "Background shell ID" },
"text_input": { "type": "string", "description": "Text to write" },
"bytes_input": { "type": "string", "description": "Special chars: <ESC>, <CR>, <C-c> etc." },
},
"required": ["shell_id"],
}),
}
});
#[cfg(test)]
tests_bridge_macro::tests_bridge!("tools/exec_shell");