perf(tools): bash 输出增量匹配消除 O(n^2) + 缓冲上限,http_request 复用 Client + 流式限长

- bash 增量窗口扫描 pending 短语,不再全量重扫;输出缓冲设上限(头尾保留)
- 修复 wait 分支 stdout_buf 重复锁导致的 tokio Mutex 自死锁(此前正常命令挂到超时)
- shell_session 沿用 cap_output_buffer 上限并修复截断后偏移兜底
- http_request 复用长生命周期 Client;响应体改为流式限长读取,防超大响应耗尽内存
This commit is contained in:
oudecheng 2026-08-17 09:55:03 +08:00
parent 4517e4a724
commit 52f858bfb4
3 changed files with 218 additions and 115 deletions

View File

@ -13,10 +13,16 @@ use tokio::time::{Instant, sleep_until};
use crate::platform::{ShellInfo, dangerous_command_patterns};
use crate::tools::shell_session::ShellSessionManager;
use crate::tools::traits::{Tool, ToolResult};
use crate::tools::{check_null_args, extract_bool, extract_u64};
use crate::tools::{check_null_args, extract_u64};
const MAX_TIMEOUT_SECS: u64 = 600;
const MAX_OUTPUT_CHARS: usize = 50_000;
/// 运行时单流输出缓冲上限(字节):超出后保留头尾、丢弃中段,
/// 防止长输出命令(或交互式会话的 drain 任务)无限增长吃满内存。
const MAX_RUNTIME_BUFFER_BYTES: usize = 1024 * 1024;
/// pending 短语增量检测的尾部窗口(字节):交互提示只出现在输出尾部,
/// 只需"新 chunk + 尾部窗口"即可捕获(窗口 ≥ 最长短语长度,覆盖跨 chunk 边界)。
const PENDING_WINDOW_BYTES: usize = 2048;
/// 子进程退出后,等待 read_stream 把管道残余输出排空的最长时间。
///
/// 不能无界等待 EOF若子进程派生了继承 stdout 管道的守护进程
@ -29,6 +35,79 @@ const INTERACTIVE_HINT: &str =
const NON_INTERACTIVE_HINT: &str =
"该命令正在等待你完成外部操作。完成后请告诉我继续,或重新运行后续检查命令。";
/// "等待用户操作"检测短语(全部小写;检测前对输出做 to_lowercase
/// 新增短语时保持小写,并确保长度不超过 PENDING_WINDOW_BYTES。
const PENDING_USER_ACTION_PHRASES: &[&str] = &[
// 中文 — 原有
"等待用户授权",
"等待授权",
"等待你授权",
"在浏览器中打开以下链接进行认证",
// 中文 — 新增lark-cli 等工具的常见提示)
"请在浏览器中",
"请打开以下链接",
"打开以下链接",
"打开链接",
"访问以下",
"访问此链接",
"复制链接",
"输入验证码",
"输入授权码",
"完成认证",
"完成授权",
"请登录",
"正在等待",
"等待用户",
"手动授权",
// 英文 — 原有
"open the following link",
"waiting for authorization",
"waiting for user authorization",
"waiting for approval",
"device/verify",
"user_code=",
// 英文 — 新增
"visit the following url",
"visit this url",
"open the following url",
"browser to authenticate",
"browser to complete",
"enter the code",
"enter code",
"verification code",
"authorization code",
"one-time code",
"device code",
"oauth",
"go to the following",
"navigate to the following",
"paste the code",
];
/// 在小写化文本中检测 pending 短语(调用方需先 to_lowercase
fn contains_pending_phrase(lowercase_text: &str) -> bool {
PENDING_USER_ACTION_PHRASES
.iter()
.any(|phrase| lowercase_text.contains(phrase))
}
/// 缓冲超限时保留头尾(头 1/2 + 尾 1/4在字符边界处截断。
/// 头尾之外的中段对最终 truncate_output头+尾各 25K 字符)已无贡献。
pub(crate) fn cap_output_buffer(buf: &mut String) {
if buf.len() <= MAX_RUNTIME_BUFFER_BYTES {
return;
}
let head_len = MAX_RUNTIME_BUFFER_BYTES / 2;
let tail_len = MAX_RUNTIME_BUFFER_BYTES / 4;
let head_end = buf.floor_char_boundary(head_len);
let tail_start = buf.floor_char_boundary(buf.len().saturating_sub(tail_len));
let mut capped = String::with_capacity(head_end + tail_len + 64);
capped.push_str(&buf[..head_end]);
capped.push_str("\n[... output trimmed in memory (buffer limit reached) ...]\n");
capped.push_str(&buf[tail_start..]);
*buf = capped;
}
/// Shell 类型枚举,支持跨平台
///
/// 这是 ShellInfo 的兼容包装,提供更方便的 API。
@ -121,7 +200,9 @@ impl ShellKind {
pub struct BashTool {
timeout_secs: u64,
working_dir: Option<String>,
deny_patterns: Vec<String>,
/// 危险命令拦截正则:构造时预编译(模式串在构造后不变),
/// 避免每次执行命令都重新编译全部正则。
deny_patterns: Vec<regex::Regex>,
shell: ShellKind,
session_manager: Arc<ShellSessionManager>,
}
@ -131,7 +212,16 @@ impl BashTool {
Self {
timeout_secs: 60,
working_dir: None,
deny_patterns: dangerous_command_patterns(),
deny_patterns: dangerous_command_patterns()
.into_iter()
.filter_map(|p| match regex::Regex::new(&p) {
Ok(re) => Some(re),
Err(e) => {
tracing::warn!(pattern = %p, error = %e, "Invalid deny pattern skipped");
None
}
})
.collect(),
shell: ShellKind::detect(),
session_manager,
}
@ -154,15 +244,11 @@ impl BashTool {
fn guard_command(&self, command: &str) -> Option<String> {
let lower = command.to_lowercase();
for pattern in &self.deny_patterns {
if regex::Regex::new(pattern)
.ok()
.map(|re| re.is_match(&lower))
.unwrap_or(false)
{
for re in &self.deny_patterns {
if re.is_match(&lower) {
return Some(format!(
"Command blocked by safety guard (dangerous pattern: {})",
pattern
re.as_str()
));
}
}
@ -208,63 +294,6 @@ impl BashTool {
PENDING_USER_ACTION_MARKER, session_line, hint, output_section
)
}
fn should_return_pending(&self, _interactive: bool, output: &str) -> bool {
let normalized = output.to_lowercase();
let has_auth_phrase = [
// 中文 — 原有
"等待用户授权",
"等待授权",
"等待你授权",
"在浏览器中打开以下链接进行认证",
// 中文 — 新增lark-cli 等工具的常见提示)
"请在浏览器中",
"请打开以下链接",
"打开以下链接",
"打开链接",
"访问以下",
"访问此链接",
"复制链接",
"输入验证码",
"输入授权码",
"完成认证",
"完成授权",
"请登录",
"正在等待",
"等待用户",
"手动授权",
// 英文 — 原有
"open the following link",
"waiting for authorization",
"waiting for user authorization",
"waiting for approval",
"device/verify",
"user_code=",
// 英文 — 新增
"visit the following url",
"visit this url",
"open the following url",
"browser to authenticate",
"browser to complete",
"enter the code",
"enter code",
"verification code",
"authorization code",
"one-time code",
"device code",
"oauth",
"go to the following",
"navigate to the following",
"paste the code",
]
.iter()
.any(|pattern| normalized.contains(pattern));
// 仅 auth 短语命中才转 pending超时前早退不构成 deadline 绕过)。
// 此前的 `|| (interactive && !output.trim().is_empty())` 会在 interactive=true
// 时对任意输出转 pending绕过超时与严格硬超时冲突已移除。
has_auth_phrase
}
}
async fn drain_available_chunks(
@ -279,6 +308,10 @@ async fn drain_available_chunks(
stdout_buf.lock().await.push_str(&chunk);
}
}
let mut stdout_guard = stdout_buf.lock().await;
cap_output_buffer(&mut stdout_guard);
let mut stderr_guard = stderr_buf.lock().await;
cap_output_buffer(&mut stderr_guard);
}
impl Default for BashTool {
@ -380,7 +413,6 @@ impl Tool for BashTool {
let timeout_secs = extract_u64(&args, "timeout")
.unwrap_or(self.timeout_secs)
.clamp(1, MAX_TIMEOUT_SECS); // 下界 1 防止 timeout:0 误杀刚 spawn 的子进程;上界 600s与 schema minimum/maximum 对齐)
let interactive = extract_bool(&args, "interactive").unwrap_or(false);
let cwd = self
.working_dir
@ -388,10 +420,7 @@ impl Tool for BashTool {
.map(Path::new)
.unwrap_or_else(|| Path::new("."));
match self
.run_command(command, cwd, timeout_secs, interactive)
.await
{
match self.run_command(command, cwd, timeout_secs).await {
Ok(output) => Ok(ToolResult {
success: true,
output,
@ -410,15 +439,17 @@ impl BashTool {
/// 强制终止子进程并回收,避免 `wait()` 永久挂起导致超时未生效。
///
/// `start_kill` 在 Unix 发 SIGKILL、在 Windows 调 TerminateProcess均为强制终止
/// 用 `timeout` 包裹 `wait()` 防止 reap 在异常情况下永久挂起5s 内未回收则再
/// `kill().await`(重发信号并等待)+ 3s 兜底,保证工具调用必然返回。
/// 用 `timeout` 包裹 `wait()` 防止 reap 在异常情况下永久挂起5s 内未回收则
/// 重发终止信号(`start_kill` 非阻塞;不用 `kill().await`——其内部无超时地
/// await wait(),在进程被外部挂起/保护时会永久阻塞),再等 3s 兜底,
/// 保证工具调用必然返回。
async fn kill_and_reap(child: &mut tokio::process::Child) {
let _ = child.start_kill();
if tokio::time::timeout(Duration::from_secs(5), child.wait())
.await
.is_err()
{
let _ = child.kill().await;
let _ = child.start_kill();
let _ = tokio::time::timeout(Duration::from_secs(3), child.wait()).await;
}
}
@ -428,7 +459,6 @@ impl BashTool {
command: &str,
cwd: &Path,
timeout_secs: u64,
interactive: bool,
) -> Result<String, String> {
let mut cmd = Command::new(self.shell.executable());
cmd.args(self.shell.command_args(command))
@ -457,6 +487,9 @@ impl BashTool {
let stdout_buf = Arc::new(Mutex::new(String::new()));
let stderr_buf = Arc::new(Mutex::new(String::new()));
// pending 短语增量检测窗口:仅保留最近输出的小写化尾部,
// 每个 chunk 只扫描"窗口 + 新 chunk"O(chunk)不再全量重扫O(累计输出))。
let mut pending_window = String::new();
let deadline = Instant::now() + Duration::from_secs(timeout_secs);
loop {
@ -481,13 +514,19 @@ impl BashTool {
},
)
.await;
let mut stdout_guard = stdout_buf.lock().await;
cap_output_buffer(&mut stdout_guard);
let mut stderr_guard = stderr_buf.lock().await;
cap_output_buffer(&mut stderr_guard);
// 终止可能仍阻塞在 reader.read() 上的 read_stream 任务,避免泄漏。
for t in &read_tasks {
t.abort();
}
// 注意:直接复用上方已持有的 guard。tokio::sync::Mutex 不可重入,
// 此处若再次 lock().await 会自死锁(同一任务持锁等待自身释放)。
let output = format_command_output(
&stdout_buf.lock().await,
&stderr_buf.lock().await,
&stdout_guard,
&stderr_guard,
Some(status.code().unwrap_or(-1)),
);
return Ok(self.truncate_output(&output));
@ -498,14 +537,25 @@ impl BashTool {
None => std::future::pending().await,
}
} => {
if is_stderr {
stderr_buf.lock().await.push_str(&chunk);
} else {
stdout_buf.lock().await.push_str(&chunk);
{
let mut buf = if is_stderr {
stderr_buf.lock().await
} else {
stdout_buf.lock().await
};
buf.push_str(&chunk);
cap_output_buffer(&mut buf);
}
let combined = format_command_output(&stdout_buf.lock().await, &stderr_buf.lock().await, None);
if self.should_return_pending(interactive, &combined) {
// 增量 pending 检测:交互提示只出现在输出尾部,只扫描
// "尾部窗口 + 新 chunk",避免对全量输出做 O(n²) 重扫
pending_window.push_str(&chunk.to_lowercase());
if pending_window.len() > PENDING_WINDOW_BYTES * 2 {
let cut = pending_window
.floor_char_boundary(pending_window.len() - PENDING_WINDOW_BYTES);
pending_window.drain(..cut);
}
if contains_pending_phrase(&pending_window) {
let mut rx_val = rx.take().unwrap();
drain_available_chunks(&mut rx_val, &stdout_buf, &stderr_buf).await;
let combined = format_command_output(&stdout_buf.lock().await, &stderr_buf.lock().await, None);

View File

@ -1,17 +1,22 @@
use std::time::Duration;
use async_trait::async_trait;
use futures_util::StreamExt;
use reqwest::header::HeaderMap;
use serde_json::json;
use crate::text::take_prefix_chars;
use crate::tools::traits::{Tool, ToolResult};
/// 未配置响应大小限制时的硬性下载上限(防止无限响应打满内存)。
const HARD_DOWNLOAD_CAP_BYTES: usize = 32 * 1024 * 1024;
pub struct HttpRequestTool {
allowed_domains: Vec<String>,
max_response_size: usize,
timeout_secs: u64,
allow_private_hosts: bool,
/// 长生命周期 HTTP 客户端(连接池 + TLS 上下文 + 超时配置),构造一次全程复用。
client: reqwest::Client,
}
impl HttpRequestTool {
@ -21,11 +26,16 @@ impl HttpRequestTool {
timeout_secs: u64,
allow_private_hosts: bool,
) -> Self {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(timeout_secs))
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("valid HTTP client configuration");
Self {
allowed_domains: normalize_domains(allowed_domains),
max_response_size,
timeout_secs,
allow_private_hosts,
client,
}
}
@ -102,6 +112,41 @@ impl HttpRequestTool {
text.to_string()
}
}
/// 下载字节上限:字符上限 × 4UTF-8 单字符最多 4 字节)保证字符截断前必然读够;
/// 未配置字符上限时使用硬性上限,任何情况下下载量都有界。
fn download_byte_limit(&self) -> usize {
if self.max_response_size == 0 {
HARD_DOWNLOAD_CAP_BYTES
} else {
self.max_response_size
.saturating_mul(4)
.min(HARD_DOWNLOAD_CAP_BYTES)
}
}
}
/// 流式读取响应体,累计达到 `max_bytes` 即提前中止下载。
/// 限制在下载过程中生效(而非全量载入后截断),防止超大响应耗尽内存。
async fn read_body_limited(
response: reqwest::Response,
max_bytes: usize,
) -> Result<String, String> {
let mut stream = response.bytes_stream();
let mut body: Vec<u8> = Vec::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| format!("Failed to read response body: {}", e))?;
let remaining = max_bytes.saturating_sub(body.len());
if remaining == 0 {
break;
}
if chunk.len() > remaining {
body.extend_from_slice(&chunk[..remaining]);
break;
}
body.extend_from_slice(&chunk);
}
Ok(String::from_utf8_lossy(&body).into_owned())
}
fn normalize_domains(domains: Vec<String>) -> Vec<String> {
@ -308,22 +353,7 @@ impl Tool for HttpRequestTool {
let headers = self.parse_headers(&headers_val);
let client = match reqwest::Client::builder()
.timeout(Duration::from_secs(self.timeout_secs))
.redirect(reqwest::redirect::Policy::none())
.build()
{
Ok(c) => c,
Err(e) => {
return Ok(ToolResult {
success: false,
output: String::new(),
error: Some(format!("Failed to create HTTP client: {}", e)),
});
}
};
let mut request = client.request(method, &url).headers(headers);
let mut request = self.client.request(method, &url).headers(headers);
if let Some(body_str) = body {
request = request.body(body_str.to_string());
@ -334,11 +364,12 @@ impl Tool for HttpRequestTool {
let status = response.status();
let status_code = status.as_u16();
let response_text = response
.text()
.await
.map(|t| self.truncate_response(&t))
.unwrap_or_else(|_| "[Failed to read response body]".to_string());
// 流式限长读取:下载量在读取过程中即被约束,超限提前中止
let response_text =
match read_body_limited(response, self.download_byte_limit()).await {
Ok(text) => self.truncate_response(&text),
Err(_) => "[Failed to read response body]".to_string(),
};
let output = format!(
"Status: {} {}\n\nResponse Body:\n{}",

View File

@ -19,6 +19,8 @@ use tokio::time::Instant;
use uuid::Uuid;
use crate::tools::bash::cap_output_buffer;
const SESSION_TIMEOUT_SECS: u64 = 300; // 5 minutes
const OUTPUT_WAIT_MS: u64 = 2000;
@ -72,14 +74,20 @@ impl ShellSessionManager {
let stderr_buf = Arc::new(Mutex::new(initial_stderr));
// Spawn a background task that drains the channel into buffers.
// Buffers are capped in size (head+tail retained) so a long-lived
// chatty process cannot grow memory without bound during the session TTL.
let stdout_clone = stdout_buf.clone();
let stderr_clone = stderr_buf.clone();
let drain_task = tokio::spawn(async move {
while let Some((is_stderr, chunk)) = rx.recv().await {
if is_stderr {
stderr_clone.lock().await.push_str(&chunk);
let mut buf = stderr_clone.lock().await;
buf.push_str(&chunk);
cap_output_buffer(&mut buf);
} else {
stdout_clone.lock().await.push_str(&chunk);
let mut buf = stdout_clone.lock().await;
buf.push_str(&chunk);
cap_output_buffer(&mut buf);
}
}
});
@ -134,7 +142,9 @@ impl ShellSessionManager {
return Err("Session stdin is closed".to_string());
}
// Record output length before wait
// Record output length before wait (byte offsets — buffers only grow via
// push_str, so a previous end is always a valid char boundary, unless a
// buffer-cap trim happened in between, which is handled below).
let prev_stdout_len = session.stdout_buf.lock().await.len();
let prev_stderr_len = session.stderr_buf.lock().await.len();
@ -154,11 +164,23 @@ impl ShellSessionManager {
}
}
let stdout = session.stdout_buf.lock().await.clone();
let stderr = session.stderr_buf.lock().await.clone();
let new_stdout: String = stdout.chars().skip(prev_stdout_len).collect();
let new_stderr: String = stderr.chars().skip(prev_stderr_len).collect();
// 按字节偏移在锁内直接切片取新增输出:避免对整个缓冲做全量克隆,
// 也修复了旧实现"字节长度当字符数 skip"导致多字节输出丢失的问题。
let new_stdout = {
let buf = session.stdout_buf.lock().await;
match buf.get(prev_stdout_len..) {
Some(new_part) => new_part.to_string(),
// 偏移失效(缓冲被上限截断过):退化为返回截断后的全部内容
None => buf.clone(),
}
};
let new_stderr = {
let buf = session.stderr_buf.lock().await;
match buf.get(prev_stderr_len..) {
Some(new_part) => new_part.to_string(),
None => buf.clone(),
}
};
let mut result = String::new();
if !new_stdout.is_empty() {