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:
parent
4517e4a724
commit
52f858bfb4
@ -13,10 +13,16 @@ use tokio::time::{Instant, sleep_until};
|
|||||||
use crate::platform::{ShellInfo, dangerous_command_patterns};
|
use crate::platform::{ShellInfo, dangerous_command_patterns};
|
||||||
use crate::tools::shell_session::ShellSessionManager;
|
use crate::tools::shell_session::ShellSessionManager;
|
||||||
use crate::tools::traits::{Tool, ToolResult};
|
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_TIMEOUT_SECS: u64 = 600;
|
||||||
const MAX_OUTPUT_CHARS: usize = 50_000;
|
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 把管道残余输出排空的最长时间。
|
/// 子进程退出后,等待 read_stream 把管道残余输出排空的最长时间。
|
||||||
///
|
///
|
||||||
/// 不能无界等待 EOF:若子进程派生了继承 stdout 管道的守护进程
|
/// 不能无界等待 EOF:若子进程派生了继承 stdout 管道的守护进程
|
||||||
@ -29,6 +35,79 @@ const INTERACTIVE_HINT: &str =
|
|||||||
const NON_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 类型枚举,支持跨平台
|
/// Shell 类型枚举,支持跨平台
|
||||||
///
|
///
|
||||||
/// 这是 ShellInfo 的兼容包装,提供更方便的 API。
|
/// 这是 ShellInfo 的兼容包装,提供更方便的 API。
|
||||||
@ -121,7 +200,9 @@ impl ShellKind {
|
|||||||
pub struct BashTool {
|
pub struct BashTool {
|
||||||
timeout_secs: u64,
|
timeout_secs: u64,
|
||||||
working_dir: Option<String>,
|
working_dir: Option<String>,
|
||||||
deny_patterns: Vec<String>,
|
/// 危险命令拦截正则:构造时预编译(模式串在构造后不变),
|
||||||
|
/// 避免每次执行命令都重新编译全部正则。
|
||||||
|
deny_patterns: Vec<regex::Regex>,
|
||||||
shell: ShellKind,
|
shell: ShellKind,
|
||||||
session_manager: Arc<ShellSessionManager>,
|
session_manager: Arc<ShellSessionManager>,
|
||||||
}
|
}
|
||||||
@ -131,7 +212,16 @@ impl BashTool {
|
|||||||
Self {
|
Self {
|
||||||
timeout_secs: 60,
|
timeout_secs: 60,
|
||||||
working_dir: None,
|
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(),
|
shell: ShellKind::detect(),
|
||||||
session_manager,
|
session_manager,
|
||||||
}
|
}
|
||||||
@ -154,15 +244,11 @@ impl BashTool {
|
|||||||
|
|
||||||
fn guard_command(&self, command: &str) -> Option<String> {
|
fn guard_command(&self, command: &str) -> Option<String> {
|
||||||
let lower = command.to_lowercase();
|
let lower = command.to_lowercase();
|
||||||
for pattern in &self.deny_patterns {
|
for re in &self.deny_patterns {
|
||||||
if regex::Regex::new(pattern)
|
if re.is_match(&lower) {
|
||||||
.ok()
|
|
||||||
.map(|re| re.is_match(&lower))
|
|
||||||
.unwrap_or(false)
|
|
||||||
{
|
|
||||||
return Some(format!(
|
return Some(format!(
|
||||||
"Command blocked by safety guard (dangerous pattern: {})",
|
"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
|
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(
|
async fn drain_available_chunks(
|
||||||
@ -279,6 +308,10 @@ async fn drain_available_chunks(
|
|||||||
stdout_buf.lock().await.push_str(&chunk);
|
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 {
|
impl Default for BashTool {
|
||||||
@ -380,7 +413,6 @@ impl Tool for BashTool {
|
|||||||
let timeout_secs = extract_u64(&args, "timeout")
|
let timeout_secs = extract_u64(&args, "timeout")
|
||||||
.unwrap_or(self.timeout_secs)
|
.unwrap_or(self.timeout_secs)
|
||||||
.clamp(1, MAX_TIMEOUT_SECS); // 下界 1 防止 timeout:0 误杀刚 spawn 的子进程;上界 600s(与 schema minimum/maximum 对齐)
|
.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
|
let cwd = self
|
||||||
.working_dir
|
.working_dir
|
||||||
@ -388,10 +420,7 @@ impl Tool for BashTool {
|
|||||||
.map(Path::new)
|
.map(Path::new)
|
||||||
.unwrap_or_else(|| Path::new("."));
|
.unwrap_or_else(|| Path::new("."));
|
||||||
|
|
||||||
match self
|
match self.run_command(command, cwd, timeout_secs).await {
|
||||||
.run_command(command, cwd, timeout_secs, interactive)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(output) => Ok(ToolResult {
|
Ok(output) => Ok(ToolResult {
|
||||||
success: true,
|
success: true,
|
||||||
output,
|
output,
|
||||||
@ -410,15 +439,17 @@ impl BashTool {
|
|||||||
/// 强制终止子进程并回收,避免 `wait()` 永久挂起导致超时未生效。
|
/// 强制终止子进程并回收,避免 `wait()` 永久挂起导致超时未生效。
|
||||||
///
|
///
|
||||||
/// `start_kill` 在 Unix 发 SIGKILL、在 Windows 调 TerminateProcess(均为强制终止)。
|
/// `start_kill` 在 Unix 发 SIGKILL、在 Windows 调 TerminateProcess(均为强制终止)。
|
||||||
/// 用 `timeout` 包裹 `wait()` 防止 reap 在异常情况下永久挂起;5s 内未回收则再
|
/// 用 `timeout` 包裹 `wait()` 防止 reap 在异常情况下永久挂起;5s 内未回收则
|
||||||
/// `kill().await`(重发信号并等待)+ 3s 兜底,保证工具调用必然返回。
|
/// 重发终止信号(`start_kill` 非阻塞;不用 `kill().await`——其内部无超时地
|
||||||
|
/// await wait(),在进程被外部挂起/保护时会永久阻塞),再等 3s 兜底,
|
||||||
|
/// 保证工具调用必然返回。
|
||||||
async fn kill_and_reap(child: &mut tokio::process::Child) {
|
async fn kill_and_reap(child: &mut tokio::process::Child) {
|
||||||
let _ = child.start_kill();
|
let _ = child.start_kill();
|
||||||
if tokio::time::timeout(Duration::from_secs(5), child.wait())
|
if tokio::time::timeout(Duration::from_secs(5), child.wait())
|
||||||
.await
|
.await
|
||||||
.is_err()
|
.is_err()
|
||||||
{
|
{
|
||||||
let _ = child.kill().await;
|
let _ = child.start_kill();
|
||||||
let _ = tokio::time::timeout(Duration::from_secs(3), child.wait()).await;
|
let _ = tokio::time::timeout(Duration::from_secs(3), child.wait()).await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -428,7 +459,6 @@ impl BashTool {
|
|||||||
command: &str,
|
command: &str,
|
||||||
cwd: &Path,
|
cwd: &Path,
|
||||||
timeout_secs: u64,
|
timeout_secs: u64,
|
||||||
interactive: bool,
|
|
||||||
) -> Result<String, String> {
|
) -> Result<String, String> {
|
||||||
let mut cmd = Command::new(self.shell.executable());
|
let mut cmd = Command::new(self.shell.executable());
|
||||||
cmd.args(self.shell.command_args(command))
|
cmd.args(self.shell.command_args(command))
|
||||||
@ -457,6 +487,9 @@ impl BashTool {
|
|||||||
|
|
||||||
let stdout_buf = Arc::new(Mutex::new(String::new()));
|
let stdout_buf = Arc::new(Mutex::new(String::new()));
|
||||||
let stderr_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);
|
let deadline = Instant::now() + Duration::from_secs(timeout_secs);
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
@ -481,13 +514,19 @@ impl BashTool {
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.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 任务,避免泄漏。
|
// 终止可能仍阻塞在 reader.read() 上的 read_stream 任务,避免泄漏。
|
||||||
for t in &read_tasks {
|
for t in &read_tasks {
|
||||||
t.abort();
|
t.abort();
|
||||||
}
|
}
|
||||||
|
// 注意:直接复用上方已持有的 guard。tokio::sync::Mutex 不可重入,
|
||||||
|
// 此处若再次 lock().await 会自死锁(同一任务持锁等待自身释放)。
|
||||||
let output = format_command_output(
|
let output = format_command_output(
|
||||||
&stdout_buf.lock().await,
|
&stdout_guard,
|
||||||
&stderr_buf.lock().await,
|
&stderr_guard,
|
||||||
Some(status.code().unwrap_or(-1)),
|
Some(status.code().unwrap_or(-1)),
|
||||||
);
|
);
|
||||||
return Ok(self.truncate_output(&output));
|
return Ok(self.truncate_output(&output));
|
||||||
@ -498,14 +537,25 @@ impl BashTool {
|
|||||||
None => std::future::pending().await,
|
None => std::future::pending().await,
|
||||||
}
|
}
|
||||||
} => {
|
} => {
|
||||||
if is_stderr {
|
{
|
||||||
stderr_buf.lock().await.push_str(&chunk);
|
let mut buf = if is_stderr {
|
||||||
|
stderr_buf.lock().await
|
||||||
} else {
|
} else {
|
||||||
stdout_buf.lock().await.push_str(&chunk);
|
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);
|
// 增量 pending 检测:交互提示只出现在输出尾部,只扫描
|
||||||
if self.should_return_pending(interactive, &combined) {
|
// "尾部窗口 + 新 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();
|
let mut rx_val = rx.take().unwrap();
|
||||||
drain_available_chunks(&mut rx_val, &stdout_buf, &stderr_buf).await;
|
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);
|
let combined = format_command_output(&stdout_buf.lock().await, &stderr_buf.lock().await, None);
|
||||||
|
|||||||
@ -1,17 +1,22 @@
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use futures_util::StreamExt;
|
||||||
use reqwest::header::HeaderMap;
|
use reqwest::header::HeaderMap;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
use crate::text::take_prefix_chars;
|
use crate::text::take_prefix_chars;
|
||||||
use crate::tools::traits::{Tool, ToolResult};
|
use crate::tools::traits::{Tool, ToolResult};
|
||||||
|
|
||||||
|
/// 未配置响应大小限制时的硬性下载上限(防止无限响应打满内存)。
|
||||||
|
const HARD_DOWNLOAD_CAP_BYTES: usize = 32 * 1024 * 1024;
|
||||||
|
|
||||||
pub struct HttpRequestTool {
|
pub struct HttpRequestTool {
|
||||||
allowed_domains: Vec<String>,
|
allowed_domains: Vec<String>,
|
||||||
max_response_size: usize,
|
max_response_size: usize,
|
||||||
timeout_secs: u64,
|
|
||||||
allow_private_hosts: bool,
|
allow_private_hosts: bool,
|
||||||
|
/// 长生命周期 HTTP 客户端(连接池 + TLS 上下文 + 超时配置),构造一次全程复用。
|
||||||
|
client: reqwest::Client,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl HttpRequestTool {
|
impl HttpRequestTool {
|
||||||
@ -21,11 +26,16 @@ impl HttpRequestTool {
|
|||||||
timeout_secs: u64,
|
timeout_secs: u64,
|
||||||
allow_private_hosts: bool,
|
allow_private_hosts: bool,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(Duration::from_secs(timeout_secs))
|
||||||
|
.redirect(reqwest::redirect::Policy::none())
|
||||||
|
.build()
|
||||||
|
.expect("valid HTTP client configuration");
|
||||||
Self {
|
Self {
|
||||||
allowed_domains: normalize_domains(allowed_domains),
|
allowed_domains: normalize_domains(allowed_domains),
|
||||||
max_response_size,
|
max_response_size,
|
||||||
timeout_secs,
|
|
||||||
allow_private_hosts,
|
allow_private_hosts,
|
||||||
|
client,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -102,6 +112,41 @@ impl HttpRequestTool {
|
|||||||
text.to_string()
|
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<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> {
|
fn normalize_domains(domains: Vec<String>) -> Vec<String> {
|
||||||
@ -308,22 +353,7 @@ impl Tool for HttpRequestTool {
|
|||||||
|
|
||||||
let headers = self.parse_headers(&headers_val);
|
let headers = self.parse_headers(&headers_val);
|
||||||
|
|
||||||
let client = match reqwest::Client::builder()
|
let mut request = self.client.request(method, &url).headers(headers);
|
||||||
.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);
|
|
||||||
|
|
||||||
if let Some(body_str) = body {
|
if let Some(body_str) = body {
|
||||||
request = request.body(body_str.to_string());
|
request = request.body(body_str.to_string());
|
||||||
@ -334,11 +364,12 @@ impl Tool for HttpRequestTool {
|
|||||||
let status = response.status();
|
let status = response.status();
|
||||||
let status_code = status.as_u16();
|
let status_code = status.as_u16();
|
||||||
|
|
||||||
let response_text = response
|
// 流式限长读取:下载量在读取过程中即被约束,超限提前中止
|
||||||
.text()
|
let response_text =
|
||||||
.await
|
match read_body_limited(response, self.download_byte_limit()).await {
|
||||||
.map(|t| self.truncate_response(&t))
|
Ok(text) => self.truncate_response(&text),
|
||||||
.unwrap_or_else(|_| "[Failed to read response body]".to_string());
|
Err(_) => "[Failed to read response body]".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
let output = format!(
|
let output = format!(
|
||||||
"Status: {} {}\n\nResponse Body:\n{}",
|
"Status: {} {}\n\nResponse Body:\n{}",
|
||||||
|
|||||||
@ -19,6 +19,8 @@ use tokio::time::Instant;
|
|||||||
|
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::tools::bash::cap_output_buffer;
|
||||||
|
|
||||||
const SESSION_TIMEOUT_SECS: u64 = 300; // 5 minutes
|
const SESSION_TIMEOUT_SECS: u64 = 300; // 5 minutes
|
||||||
const OUTPUT_WAIT_MS: u64 = 2000;
|
const OUTPUT_WAIT_MS: u64 = 2000;
|
||||||
|
|
||||||
@ -72,14 +74,20 @@ impl ShellSessionManager {
|
|||||||
let stderr_buf = Arc::new(Mutex::new(initial_stderr));
|
let stderr_buf = Arc::new(Mutex::new(initial_stderr));
|
||||||
|
|
||||||
// Spawn a background task that drains the channel into buffers.
|
// 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 stdout_clone = stdout_buf.clone();
|
||||||
let stderr_clone = stderr_buf.clone();
|
let stderr_clone = stderr_buf.clone();
|
||||||
let drain_task = tokio::spawn(async move {
|
let drain_task = tokio::spawn(async move {
|
||||||
while let Some((is_stderr, chunk)) = rx.recv().await {
|
while let Some((is_stderr, chunk)) = rx.recv().await {
|
||||||
if is_stderr {
|
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 {
|
} 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());
|
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_stdout_len = session.stdout_buf.lock().await.len();
|
||||||
let prev_stderr_len = session.stderr_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();
|
// 也修复了旧实现"字节长度当字符数 skip"导致多字节输出丢失的问题。
|
||||||
|
let new_stdout = {
|
||||||
let new_stdout: String = stdout.chars().skip(prev_stdout_len).collect();
|
let buf = session.stdout_buf.lock().await;
|
||||||
let new_stderr: String = stderr.chars().skip(prev_stderr_len).collect();
|
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();
|
let mut result = String::new();
|
||||||
if !new_stdout.is_empty() {
|
if !new_stdout.is_empty() {
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user