From 52f858bfb4cec46a40302b7486372a8a9d2d9672 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Mon, 17 Aug 2026 09:55:03 +0800 Subject: [PATCH] =?UTF-8?q?perf(tools):=20bash=20=E8=BE=93=E5=87=BA?= =?UTF-8?q?=E5=A2=9E=E9=87=8F=E5=8C=B9=E9=85=8D=E6=B6=88=E9=99=A4=20O(n^2)?= =?UTF-8?q?=20+=20=E7=BC=93=E5=86=B2=E4=B8=8A=E9=99=90=EF=BC=8Chttp=5Frequ?= =?UTF-8?q?est=20=E5=A4=8D=E7=94=A8=20Client=20+=20=E6=B5=81=E5=BC=8F?= =?UTF-8?q?=E9=99=90=E9=95=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - bash 增量窗口扫描 pending 短语，不再全量重扫；输出缓冲设上限（头尾保留） - 修复 wait 分支 stdout_buf 重复锁导致的 tokio Mutex 自死锁（此前正常命令挂到超时） - shell_session 沿用 cap_output_buffer 上限并修复截断后偏移兜底 - http_request 复用长生命周期 Client；响应体改为流式限长读取，防超大响应耗尽内存 --- src/tools/bash.rs | 218 +++++++++++++++++++++++-------------- src/tools/http_request.rs | 77 +++++++++---- src/tools/shell_session.rs | 38 +++++-- 3 files changed, 218 insertions(+), 115 deletions(-) diff --git a/src/tools/bash.rs b/src/tools/bash.rs index 7b07ff2..a665391 100644 --- a/src/tools/bash.rs +++ b/src/tools/bash.rs @@ -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, - deny_patterns: Vec, + /// 危险命令拦截正则：构造时预编译（模式串在构造后不变）， + /// 避免每次执行命令都重新编译全部正则。 + deny_patterns: Vec, shell: ShellKind, session_manager: Arc, } @@ -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 { 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 { 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); diff --git a/src/tools/http_request.rs b/src/tools/http_request.rs index e0cc1cc..d2ef6e2 100644 --- a/src/tools/http_request.rs +++ b/src/tools/http_request.rs @@ -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, 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() } } + + /// 下载字节上限：字符上限 × 4（UTF-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 { + let mut stream = response.bytes_stream(); + let mut body: Vec = 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) -> Vec { @@ -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{}", diff --git a/src/tools/shell_session.rs b/src/tools/shell_session.rs index 98419a3..50c8d0c 100644 --- a/src/tools/shell_session.rs +++ b/src/tools/shell_session.rs @@ -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() {