130 lines
4.0 KiB
Rust
130 lines
4.0 KiB
Rust
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
|
use schemars::JsonSchema;
|
|
use serde::Deserialize;
|
|
use workdir::{CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle};
|
|
|
|
const DEFAULT_TIMEOUT_SECS: u64 = 120;
|
|
const MAX_TIMEOUT_SECS: u64 = 600;
|
|
const INLINE_BYTE_BUDGET: usize = 12 * 1024;
|
|
|
|
#[derive(Debug, Deserialize, JsonSchema)]
|
|
struct BashParams {
|
|
command: String,
|
|
#[serde(default)]
|
|
timeout: Option<u64>,
|
|
}
|
|
|
|
pub(crate) struct BashTool {
|
|
session: WorkdirSessionHandle,
|
|
}
|
|
|
|
struct CommandGuard {
|
|
session: WorkdirSessionHandle,
|
|
handle: Option<CommandHandle>,
|
|
}
|
|
|
|
impl Drop for CommandGuard {
|
|
fn drop(&mut self) {
|
|
if let Some(handle) = self.handle.take() {
|
|
let workdir = self.session.clone();
|
|
tokio::spawn(async move {
|
|
let _ = workdir.cancel_command(handle).await;
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Tool for BashTool {
|
|
async fn execute(
|
|
&self,
|
|
input_json: &str,
|
|
ctx: llm_engine::tool::ToolExecutionContext,
|
|
) -> Result<ToolOutput, ToolError> {
|
|
let params: BashParams = serde_json::from_str(input_json)
|
|
.map_err(|error| ToolError::InvalidArgument(format!("invalid Bash input: {error}")))?;
|
|
let timeout_secs = params
|
|
.timeout
|
|
.unwrap_or(DEFAULT_TIMEOUT_SECS)
|
|
.clamp(1, MAX_TIMEOUT_SECS);
|
|
let cmd_summary = truncate_for_summary(¶ms.command);
|
|
let handle = self
|
|
.session
|
|
.start_command(CommandRequest {
|
|
command: params.command,
|
|
timeout_secs,
|
|
output_limit: INLINE_BYTE_BUDGET,
|
|
tool_call_id: Some(ctx.call_id),
|
|
})
|
|
.await
|
|
.map_err(crate::ToolsError::from)?;
|
|
let mut guard = CommandGuard {
|
|
session: self.session.clone(),
|
|
handle: Some(handle.clone()),
|
|
};
|
|
let output = self
|
|
.session
|
|
.command_output(CommandOutputRequest {
|
|
handle,
|
|
cursor: 0,
|
|
limit: INLINE_BYTE_BUDGET,
|
|
wait: true,
|
|
})
|
|
.await
|
|
.map_err(crate::ToolsError::from)?;
|
|
guard.handle = None;
|
|
|
|
let summary = if output.timed_out {
|
|
format!("$ {cmd_summary} (timed out after {timeout_secs}s)")
|
|
} else {
|
|
match output.exit_code {
|
|
Some(0) => format!("$ {cmd_summary}"),
|
|
Some(code) => format!("$ {cmd_summary} (exit {code})"),
|
|
None => format!("$ {cmd_summary} (terminated)"),
|
|
}
|
|
};
|
|
let content = if output.content.is_empty() {
|
|
None
|
|
} else if output.truncated {
|
|
Some(format!(
|
|
"[showing bounded WorkdirSession command output; additional output was truncated]\n{}",
|
|
output.content
|
|
))
|
|
} else {
|
|
Some(output.content)
|
|
};
|
|
Ok(ToolOutput {
|
|
summary,
|
|
content,
|
|
attachments: Vec::new(),
|
|
})
|
|
}
|
|
}
|
|
|
|
fn truncate_for_summary(command: &str) -> String {
|
|
const MAX: usize = 100;
|
|
if command.chars().count() <= MAX {
|
|
return command.to_owned();
|
|
}
|
|
let mut summary = command.chars().take(MAX - 1).collect::<String>();
|
|
summary.push('…');
|
|
summary
|
|
}
|
|
|
|
pub fn bash_tool(session: WorkdirSessionHandle, _output_dir: PathBuf) -> ToolDefinition {
|
|
Arc::new(move || {
|
|
let schema = schemars::schema_for!(BashParams);
|
|
let meta = ToolMeta::new("Bash")
|
|
.description("Execute a shell command in the bound Workdir. Process start, bounded output, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.")
|
|
.input_schema(serde_json::to_value(schema).expect("Bash schema serialization"));
|
|
let tool: Arc<dyn Tool> = Arc::new(BashTool {
|
|
session: session.clone(),
|
|
});
|
|
(meta, tool)
|
|
})
|
|
}
|