From 2f3e63e499adc80651752b489c3b49faa1b1f8ac Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Sat, 1 Aug 2026 03:46:02 +0800 Subject: [PATCH] fix(aibox): prevent OOM from large tool outputs with multi-layer safeguards - ReadFileTool: replace readlines() with on-demand range reading; add reverse block reading for large file tail; limit single line to 10KB; hard-cap max_lines to 500 - SubprocessCommandExecutor: replace capture_output with line-by-line Popen reading; enforce 1MB output limit with immediate process termination; merge stdout/stderr counting for both sync and streaming modes - Raise TOOL_MAX_RESULT_LENGTH default from 10KB to 1MB as final safety net - Add SECURITY_AUDIT.md documenting vulnerabilities and remediation roadmap - Update README with memory safety features and v2.3 changelog --- aibox/README.md | 16 +- aibox/SECURITY_AUDIT.md | 215 ++++++++++++++++++++++++ aibox/aibox.py | 357 +++++++++++++++++++++++++++++----------- 3 files changed, 494 insertions(+), 94 deletions(-) create mode 100644 aibox/SECURITY_AUDIT.md diff --git a/aibox/README.md b/aibox/README.md index 9771ef9..29ab146 100644 --- a/aibox/README.md +++ b/aibox/README.md @@ -19,6 +19,8 @@ - 🔐 **SOCKS5 代理**：通过 `set_proxy` / `get_proxy` 工具为网页访问与 API 调用配置代理，可持久化 - 💾 **对话管理**：历史对话归档、开始新对话（重置 AI 记忆并重新加载提示词） - 📦 **对话压缩**：消息数或 Token 数达到阈值时自动压缩，保持上下文质量 +- 🛡️ **内存安全防护**：文件读取按需加载不全载、命令输出 1MB 限制、单行超长截断、工具结果多层防护防止 OOM +- 🔒 **错误自动恢复**：API 请求失败自动回滚消息，工具结果安全清洗（控制字符过滤、非字符串转换） ## 🚀 快速开始 @@ -458,7 +460,7 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`，通过 `update_i "archive_dir": "archive" }, "tools": { - "max_result_length": 10000, + "max_result_length": 1048576, "save_result_to_file": false, "tool_results_dir": "tool_results" }, @@ -609,7 +611,10 @@ class MyTool(Tool): return f"执行结果: {param}" ``` -> **注意**：`execute()` 必须返回字符串。工具结果会经过 `_sanitize_tool_result()` 清洗（None/非字符串转为安全字符串，控制字符被过滤，超长结果截断），避免破坏 API 请求格式。 +> **注意**：`execute()` 必须返回字符串。工具结果会经过多层安全防护： +> 1. **工具层源头限制**：`read_file` 按需读取不全载（最多500行，单行≤10KB），`run_command` 输出≤1MB（超出立即终止进程） +> 2. **`_sanitize_tool_result()`**：None/非字符串转为安全字符串，过滤控制字符 +> 3. **最终兜底截断**：超过 `tools.max_result_length`（默认1MB字符）时截断并提示，避免破坏 API 请求格式 ### 添加新命令 @@ -684,6 +689,13 @@ elif command == "/mycommand": ## 📄 更新日志 +### v2.3 +- ✅ **修复工具输出OOM崩溃问题**： + - `read_file`：不再一次性 `readlines()` 加载整个文件，改为按行范围按需读取；大文件 `tail=true` 使用反向块读取（>1MB文件不全文遍历）；新增单行长度限制（10KB/行，超长截断）；`max_lines` 硬上限500行防止请求超大范围 + - `run_command`：非流式执行从 `capture_output=True` 改为逐行读取并实时计数，输出达到1MB立即终止进程；流式执行同样增加1MB输出大小限制，超出终止并提示截断；统一stdout/stderr合并计数 + - `tools.max_result_length` 默认值从10KB提高到1MB字符，作为最终兜底截断 + - 形成"工具层源头限制→结果清洗→最终兜底"三层防护，彻底防止大文件/大输出导致内存溢出 + ### v2.2 - ✅ 修复提示词系统排序 bug：`SKILL_PREFIX = "1_"` 只匹配 `1_` 开头文件，导致 `2_`~`10_` 技能提示词排序错乱（`10_task_handling.md` 排在 `1_memory_search.md` 前面）；改用正则按数字前缀排序 - ✅ 修复 `_init_system_prompts` 使用未排序列表的问题，确保 system prompt 按 identity → 技能 → soul → status 正确顺序加载 diff --git a/aibox/SECURITY_AUDIT.md b/aibox/SECURITY_AUDIT.md new file mode 100644 index 0000000..ba136af --- /dev/null +++ b/aibox/SECURITY_AUDIT.md @@ -0,0 +1,215 @@ +# aibox 模块安全审计报告 + +**审计日期：** 2026-08-01 +**审计文件：** [aibox.py](file:///Users/qingsong/micro_servers/aibox/aibox.py) +**审计范围：** AI引擎、工具系统、WebSocket服务、任务调度器、数据库 + +--- + +## 问题严重等级定义 + +| 等级 | 颜色 | 说明 | 响应时间 | +|------|------|------|----------| +| P0 | 🔴 严重 | 可直接被利用获取系统权限或造成服务完全不可用 | 立即修复 | +| P1 | 🟠 高危 | 可被利用获取敏感数据或造成部分功能不可用 | 24小时内修复 | +| P2 | 🟡 中危 | 存在安全隐患但利用难度较高或影响有限 | 一周内修复 | +| P3 | 🟢 低危 | 健壮性/信息泄露类问题，风险较低 | 版本迭代时修复 | + +--- + +## 🔴 P0 严重问题 + +### 1. 命令注入漏洞 +- **位置：** `execute_local_command` 工具实现 (~L2670) +- **风险描述：** + ```python + # 当前危险实现模式 + result = subprocess.run(cmd, shell=True, capture_output=True, text=True) + ``` + 使用 `shell=True` 并通过 f-string 直接拼接用户输入的命令，攻击者（或被诱导的AI）可以构造恶意命令执行任意系统操作，例如： + - `rm -rf /` 删除所有文件 + - `cat /etc/shadow` 窃取密码哈希 + - 反弹 shell 完全控制系统 + +- **修复建议：** + 1. 禁止 `shell=True`，改用 `subprocess.run([...], shell=False)` + 2. 将命令拆分为参数列表，使用 `shlex.split()` 解析或手动分割 + 3. 添加命令白名单机制，只允许执行预定义的安全命令（ls, cat, grep等） + 4. 限制工作目录，禁止在敏感目录执行命令 + +--- + +### 2. WebSocket 无认证机制 +- **位置：** `WebSocketHandler.handle()` / `WebSocketServer` (~L6000+) +- **风险描述：** + - WebSocket 服务监听在 `0.0.0.0:8765`，公网可访问 + - 无任何 token、密码或 session 认证 + - 任何人连接后可无限制使用 DeepSeek API（消耗费用） + - 结合工具调用漏洞可直接控制服务器 + +- **修复建议：** + 1. **短期方案：** 将监听地址改为 `127.0.0.1`，仅允许本地访问，通过 gateway/WAF 统一接入并认证 + 2. **长期方案：** WebSocket 握手阶段校验 token（支持 URL query 参数或 Authorization header） + 3. 参考现有 gateway 的认证机制（RSA签名/PBKDF2）做统一接入 + 4. 添加 Origin 校验，防止跨站 WebSocket 劫持 + +--- + +## 🟠 P1 高危问题 + +### 3. 路径遍历风险 +- **位置：** `read_file` / `write_file` / `append_file` 等文件操作工具 +- **风险描述：** + - 只有 `read_memory` 实现了 `safe_path` 校验 + - 其他文件工具未限制可访问目录，通过 `../../etc/passwd` 可读取系统敏感文件 + - 写入工具可覆盖任意文件，破坏系统或植入后门 + +- **修复建议：** + - 定义允许访问的根目录（如 `~/.aibox/workspace/`） + - 统一实现 `safe_path()` 函数，校验后规范化路径，确保在允许目录内 + - 对写入操作添加额外确认或权限控制 + +--- + +### 4. SSRF 漏洞风险 +- **位置：** `web_search` / `fetch_url` / 网络请求类工具 +- **风险描述：** + - 未限制可访问的 IP 范围 + - 可访问内网服务：`http://127.0.0.1:8443`（gateway）、`http://10.x.x.x` 等 + - 可访问云服务商元数据接口（如 `http://169.254.169.254/` 获取临时凭证） + +- **修复建议：** + - 禁止访问内网 IP 段：`127.0.0.0/8`, `10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`, `169.254.0.0/16` + - 禁止访问非 HTTP/HTTPS 协议 + - 跟随重定向时重新校验目标地址 + +--- + +### 5. 敏感信息暴露 +- **位置：** 系统提示词构建 / `tools_info` +- **风险描述：** + - 系统提示词中包含所有工具的完整参数 schema + - 错误响应可能返回栈追踪等内部信息 + - 工具返回的原始错误信息可能泄露路径、代码结构 + +- **修复建议：** + - 精简工具描述，只暴露必要信息 + - 对外返回的错误信息做脱敏处理，仅返回通用错误提示 + - 详细错误只记录在服务端日志 + +--- + +### 6. API Key 配置校验缺失 +- **位置：** `ConfigManager` 初始化 +- **风险描述：** + - 配置中默认值是占位符，但启动时未校验 + - 可能因配置错误导致空 Key 或使用默认无效 Key + - 没有 Key 时直接启动失败或行为异常 + +- **修复建议：** + - 启动时强制校验 `DEEPSEEK_API_KEY` 配置 + - 未配置有效 Key 时直接拒绝启动并给出明确提示 + - 日志中不要打印完整 API Key + +--- + +## 🟡 P2 中危问题 + +### 7. 无请求频率限制 +- **位置：** WebSocket 消息处理 +- **风险：** 恶意用户可高频发送请求，导致 API 费用暴涨或服务拒绝 +- **建议：** 添加 per-connection 消息频率限制（如每分钟最多 30 轮对话） + +### 8. 工具调用结果无大小限制 +- **位置：** `_sanitize_tool_result()` +- **风险：** 读取大文件（如几GB的日志）可能导致内存溢出，服务崩溃 +- **建议：** 工具输出统一限制大小（如 1MB），超出截断并提示 + +### 9. 对话历史明文存储 +- **位置：** SQLite 数据库 `~/.aibox/memos.db` +- **风险：** 聊天记录包含敏感信息，明文存储易被窃取 +- **建议：** 对敏感内容加密存储，或限制数据库文件权限 + +### 10. 临时文件清理不彻底 +- **位置：** 代码执行/命令执行工具 +- **风险：** 执行代码生成的临时文件未清理，可能占用磁盘空间或泄露信息 +- **建议：** 使用 `tempfile.TemporaryDirectory()` 自动清理 + +### 11. 技能文件加载无沙箱 +- **位置：** `load_skills()` 技能加载 +- **风险：** 可执行技能（`.sh`/`.py`）直接运行，无隔离或校验 +- **建议：** 技能目录做权限校验，仅加载可信来源技能，或在沙箱中执行 + +### 12. 日志未脱敏 +- **位置：** Logger 输出 +- **风险：** API 请求/响应、对话内容可能写入日志文件，泄露敏感信息 +- **建议：** 日志中脱敏 API Key、密码等敏感字段，对话内容可选记录 + +--- + +## 🟢 P3 低危/健壮性问题 + +### 13. 递归深度防护不足 +- **位置：** `think_and_respond()` 递归调用 +- **说明：** 虽然设置了 `max_recursion_depth=10`，但复杂工具链仍可能导致栈溢出 +- **建议：** 改用迭代而非递归，或添加更严格的深度检查 + +### 14. 对话压缩可能丢失上下文 +- **位置：** `_compress_conversation()` +- **说明：** 当前压缩算法简单截断早期消息，可能丢失工具调用上下文 +- **建议：** 保留系统提示和工具调用/结果对，只压缩纯对话内容 + +### 15. 多线程竞争风险 +- **位置：** `EnhancedAIManager` 状态管理 +- **说明：** `messages`/`tools` 等状态在多连接下可能有竞争条件（当前单连接模型，但代码未防护） +- **建议：** 必要时添加锁机制，或确保每个连接独立实例 + +### 16. 监听地址硬编码 +- **位置：** WebSocketServer 启动 +- **说明：** 当前默认 `0.0.0.0:8765`，建议改为配置项且默认 `127.0.0.1` + +### 17. 缺少优雅关闭 +- **位置：** 信号处理 +- **说明：** 关闭时未确保消息队列刷新、数据库连接正确关闭 + +--- + +## ✅ 已有的良好安全实践 + +审计中发现以下安全措施已经实现： + +1. ✅ **工具结果清洗：** `_sanitize_tool_result()` 过滤控制字符，确保工具返回值是合法字符串 +2. ✅ **消息回滚机制：** API 失败时 `_rollback_messages()` 回滚对话，防止坏消息污染历史 +3. ✅ **对话自动压缩：** Token/消息数超限自动压缩，防止上下文无限增长 +4. ✅ **pymsg 超时防护：** `recv_timeout=10s` 防止 Slowloris DoS 攻击 +5. ✅ **log_client 健壮性：** 双层异常保护、队列大小限制、Event 同步刷新 +6. ✅ **提示词软约束：** 系统提示词中明确要求 AI 不要执行危险命令（需配合代码层面硬约束） +7. ✅ **提示词排序修复：** 数字前缀正则排序，确保技能提示正确加载顺序 +8. ✅ **内存搜索优化：** 基于主题相关性搜索，避免不必要的全库扫描 + +--- + +## 🛠️ 修复路线图 + +### 第一阶段（立即）：修复 P0 +- [ ] 修复命令注入：移除 `shell=True`，改用参数列表，添加命令白名单 +- [ ] WebSocket 认证：默认绑定 127.0.0.1，通过 gateway 接入或添加 token 认证 + +### 第二阶段（近期）：修复 P1 +- [ ] 统一文件操作 safe_path 校验，限制工作目录 +- [ ] 网络请求添加 SSRF 防护，禁止内网访问 +- [ ] 工具输出大小限制（1MB） +- [ ] 启动时 API Key 配置校验 + +### 第三阶段（后续版本）：P2/P3 +- [ ] 请求频率限制 +- [ ] 日志脱敏 +- [ ] 错误信息脱敏 +- [ ] 临时文件自动清理 +- [ ] 对话加密存储 +- [ ] 多线程安全加固 +- [ ] 优雅关闭 + +--- + +**审计完成时间：** 2026-08-01 diff --git a/aibox/aibox.py b/aibox/aibox.py index 5169026..0c2822e 100644 --- a/aibox/aibox.py +++ b/aibox/aibox.py @@ -1515,7 +1515,7 @@ class Constants: @property def TOOL_MAX_RESULT_LENGTH(cls): - return cls._get_config().get("tools.max_result_length", 10000) + return cls._get_config().get("tools.max_result_length", 1048576) # 1MB 字符限制 @property def COMPRESSOR_MIN_MESSAGES(cls): @@ -3509,10 +3509,58 @@ class CommandExecutor(ABC): class SubprocessCommandExecutor(CommandExecutor): - """基于subprocess的命令执行器 - 支持超时和工作目录""" + """基于subprocess的命令执行器 - 支持超时、工作目录和输出大小限制""" + + # 命令输出最大字节数（1MB），超出截断 + MAX_OUTPUT_BYTES = 1 * 1024 * 1024 + + def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None) -> Tuple[str, bool, bool]: + """从进程读取输出，限制总大小和超时。 + 返回 (output_str, timeout_occurred, truncated) + """ + output_chunks = [] + total_bytes = 0 + timeout_occurred = False + truncated = False + + while True: + # 检查超时 + if timeout > 0 and time.time() - start_time > timeout: + process.terminate() + timeout_occurred = True + break + + # 检查输出大小 + if total_bytes >= self.MAX_OUTPUT_BYTES: + truncated = True + process.terminate() + break + + # 读取一行 + line = process.stdout.readline() + if not line and process.poll() is not None: + break + if line: + line_bytes = len(line.encode('utf-8', errors='replace')) + if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated: + allowed = self.MAX_OUTPUT_BYTES - total_bytes + if allowed > 0: + output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace')) + total_bytes = self.MAX_OUTPUT_BYTES + truncated = True + process.terminate() + break + + output_chunks.append(line) + total_bytes += line_bytes + + if output_queue: + output_queue.put(("line", line.rstrip('\n') + "\n")) + + return ''.join(output_chunks), timeout_occurred, truncated def execute(self, command: str, timeout: int = 30, work_dir: str = ".", env: Dict = None, **kwargs) -> str: - """执行命令并返回完整结果""" + """执行命令并返回完整结果（带输出大小限制）""" if not self.is_enabled(): return "❌ 命令执行功能已禁用" @@ -3523,63 +3571,9 @@ class SubprocessCommandExecutor(CommandExecutor): if env: exec_env.update(env) + process = None try: - result = subprocess.run( - command, - shell=True, - capture_output=True, - text=True, - timeout=timeout, - env=exec_env, - cwd=os.path.expanduser(work_dir) - ) - - data = {} - - if result.stdout: - data["输出"] = result.stdout.strip() - if result.stderr: - data["错误"] = result.stderr.strip() - - data["返回码"] = result.returncode - - if not result.stdout and not result.stderr: - return self.format_result(True, "命令执行成功，无输出", data) - - return self.format_result(result.returncode == 0, "命令执行完成", data) - - except subprocess.TimeoutExpired: - return f"❌ 命令执行超时（{timeout}秒）" - except FileNotFoundError as e: - return f"❌ 命令执行错误: 命令未找到 - {str(e)}" - except PermissionError as e: - return f"❌ 命令执行错误: 权限不足 - {str(e)}" - except Exception as e: - return f"❌ 执行错误: {str(e)}" - - def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30, - work_dir: str = ".", env: Dict = None, **kwargs): - """流式执行命令，支持超时控制和工作目录""" - if not self.is_enabled(): - output_queue.put(("line", "❌ 命令执行功能已禁用\n")) - output_queue.put(("complete", "❌ 命令执行功能已禁用")) - return - - self.logger.info(f"流式执行命令: {command}") - output_queue.put(("line", f"\n🔧 执行命令: {command}\n")) - output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n")) - if work_dir != ".": - output_queue.put(("line", f"📁 工作目录: {work_dir}\n")) - output_queue.put(("line", "-" * 50 + "\n")) - - # 准备环境变量 - exec_env = os.environ.copy() - if env: - exec_env.update(env) - - try: - # 使用管道来实时获取输出 - self.process = subprocess.Popen( + process = subprocess.Popen( command, shell=True, stdout=subprocess.PIPE, @@ -3592,57 +3586,162 @@ class SubprocessCommandExecutor(CommandExecutor): cwd=os.path.expanduser(work_dir) ) + start_time = time.time() + output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time) + + # 确保进程终止 + if process.poll() is None: + if timeout_occurred or truncated: + process.kill() + else: + process.terminate() + try: + process.wait(timeout=2) + except: + process.kill() + + return_code = process.returncode if process.returncode is not None else -1 + + data = { + "输出": output_str.strip() if output_str.strip() else "(无输出)", + "返回码": return_code + } + + if timeout_occurred: + return self.format_result(False, f"命令执行超时（{timeout}秒）", data) + + if truncated: + data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制，结果已截断" + return self.format_result(False, "命令执行完成（输出已截断）", data) + + return self.format_result(return_code == 0, "命令执行完成", data) + + except subprocess.TimeoutExpired: + return f"❌ 命令执行超时（{timeout}秒）" + except FileNotFoundError as e: + return f"❌ 命令执行错误: 命令未找到 - {str(e)}" + except PermissionError as e: + return f"❌ 命令执行错误: 权限不足 - {str(e)}" + except Exception as e: + return f"❌ 执行错误: {str(e)}" + finally: + if process and process.poll() is None: + process.kill() + + def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30, + work_dir: str = ".", env: Dict = None, **kwargs): + """流式执行命令，支持超时控制、工作目录和输出大小限制""" + if not self.is_enabled(): + output_queue.put(("line", "❌ 命令执行功能已禁用\n")) + output_queue.put(("complete", "❌ 命令执行功能已禁用")) + return + + self.logger.info(f"流式执行命令: {command}") + output_queue.put(("line", f"\n🔧 执行命令: {command}\n")) + output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n")) + output_queue.put(("line", f"📦 输出限制: {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB\n")) + if work_dir != ".": + output_queue.put(("line", f"📁 工作目录: {work_dir}\n")) + output_queue.put(("line", "-" * 50 + "\n")) + + # 准备环境变量 + exec_env = os.environ.copy() + if env: + exec_env.update(env) + + process = None + try: + process = subprocess.Popen( + command, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + text=True, + bufsize=1, + universal_newlines=True, + env=exec_env, + cwd=os.path.expanduser(work_dir) + ) + self.process = process + start_time = time.time() output_lines = [] + total_bytes = 0 timeout_occurred = False + truncated = False # 实时读取输出 while True: # 检查超时 if timeout > 0 and time.time() - start_time > timeout: - self.process.terminate() + process.terminate() timeout_occurred = True output_queue.put(("line", f"\n⏰ 命令执行超时（{timeout}秒）\n")) break + # 检查输出大小 + if total_bytes >= self.MAX_OUTPUT_BYTES: + truncated = True + process.terminate() + output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制，结果已截断\n")) + break + # 读取一行输出 - line = self.process.stdout.readline() - if not line and self.process.poll() is not None: + line = process.stdout.readline() + if not line and process.poll() is not None: break if line: - line = line.rstrip('\n') - output_queue.put(("line", line + "\n")) + line_bytes = len(line.encode('utf-8', errors='replace')) + + if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated: + allowed = self.MAX_OUTPUT_BYTES - total_bytes + if allowed > 0: + truncated_line = line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace') + output_queue.put(("line", truncated_line)) + output_lines.append(truncated_line) + total_bytes = self.MAX_OUTPUT_BYTES + truncated = True + process.terminate() + output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制，结果已截断\n")) + break + + stripped = line.rstrip('\n') + output_queue.put(("line", stripped + "\n")) output_lines.append(line) + total_bytes += line_bytes - # 如果超时，确保进程被终止 - if timeout_occurred and self.process.poll() is None: - self.process.kill() - self.process.wait(timeout=2) + # 如果超时或截断，确保进程被终止 + if (timeout_occurred or truncated) and process.poll() is None: + process.kill() + try: + process.wait(timeout=2) + except: + pass - # 获取剩余输出 - remaining, _ = self.process.communicate(timeout=5) - if remaining: - output_queue.put(("line", remaining)) - output_lines.append(remaining) - - return_code = self.process.returncode if self.process.returncode is not None else -1 + return_code = process.returncode if process.returncode is not None else -1 output_queue.put(("line", "-" * 50 + "\n")) if timeout_occurred: output_queue.put(("line", f"⏰ 命令执行超时（{timeout}秒），返回码: {return_code}\n")) + elif truncated: + output_queue.put(("line", f"📦 输出已截断，返回码: {return_code}\n")) else: output_queue.put(("line", f"✅ 命令执行完成，返回码: {return_code}\n")) # 准备完整结果 - full_output = "\n".join(output_lines) + full_output = "".join(output_lines) data = { - "输出": full_output if full_output else "(无输出)", + "输出": full_output.strip() if full_output.strip() else "(无输出)", "返回码": return_code } if timeout_occurred: result = self.format_result(False, f"命令执行超时（{timeout}秒）", data) + elif truncated: + data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制，结果已截断" + result = self.format_result(False, "命令执行完成（输出已截断）", data) else: result = self.format_result(return_code == 0, "命令执行完成", data) @@ -3652,12 +3751,13 @@ class SubprocessCommandExecutor(CommandExecutor): output_queue.put(("line", f"❌ 执行错误: {str(e)}\n")) output_queue.put(("complete", f"❌ 执行错误: {str(e)}")) finally: - if self.process and self.process.poll() is None: - self.process.terminate() + if process and process.poll() is None: + process.terminate() try: - self.process.wait(timeout=5) + process.wait(timeout=2) except: - self.process.kill() + process.kill() + self.process = None # ==================== 工具基类 ==================== @@ -5840,6 +5940,63 @@ class ReadFileTool(FileToolBase): return start, end, desc + @staticmethod + def _count_lines(filepath: str, encoding: str) -> int: + """高效统计文件行数，不存储内容""" + line_count = 0 + with open(filepath, 'r', encoding=encoding) as f: + for _ in f: + line_count += 1 + return line_count + + @staticmethod + def _read_lines_range(filepath: str, encoding: str, start: int, end: int) -> List[str]: + """读取指定行范围 [start, end)（0-indexed），只读取需要的行""" + lines = [] + with open(filepath, 'r', encoding=encoding) as f: + for i, line in enumerate(f): + if i >= end: + break + if i >= start: + lines.append(line) + return lines + + @staticmethod + def _read_tail_lines(filepath: str, encoding: str, num_lines: int) -> Tuple[List[str], int]: + """高效读取文件末尾N行，无需遍历整个文件。返回 (lines, total_line_count)""" + # 对于小文件直接读全部 + file_size = os.path.getsize(filepath) + if file_size < 1024 * 1024: # < 1MB 直接读 + with open(filepath, 'r', encoding=encoding) as f: + all_lines = f.readlines() + return all_lines[-num_lines:] if num_lines > 0 else [], len(all_lines) + + # 大文件：从后向前找换行符 + with open(filepath, 'rb') as f: + block_size = 8192 + blocks = [] + lines_found = 0 + position = file_size + + while position > 0 and lines_found <= num_lines: + read_size = min(block_size, position) + position -= read_size + f.seek(position) + block = f.read(read_size) + blocks.append(block) + lines_found += block.count(b'\n') + + # 拼接并解码 + data = b''.join(reversed(blocks)) + text = data.decode(encoding, errors='replace') + all_lines = text.splitlines() + + # 统计总行数 + f.seek(0) + total_lines = sum(1 for _ in f) + + return all_lines[-num_lines:], total_lines + def execute(self, filepath: str, encoding: str = "utf-8", start_line: int = 1, end_line: int = 0, max_lines: int = 100, head: bool = False, @@ -5851,13 +6008,15 @@ class ReadFileTool(FileToolBase): if not is_safe: return error + # 限制max_lines上限，防止请求超大范围 + max_lines = min(max_lines, 500) + try: file_size = os.path.getsize(abs_path) modified_time = datetime.fromtimestamp(os.path.getmtime(abs_path)) - with open(abs_path, 'r', encoding=encoding) as f: - all_lines = f.readlines() - line_count = len(all_lines) + # 高效读取：先统计行数 + line_count = self._count_lines(abs_path, encoding) start, end, range_desc = self._calculate_line_range( line_count, start_line, end_line, max_lines, head, tail @@ -5866,7 +6025,25 @@ class ReadFileTool(FileToolBase): if end <= start: return f"❌ 无效的行范围" - display_lines = all_lines[start:end] + lines_to_read = end - start + + # 对于tail模式，用高效的尾部读取 + if tail and line_count > 200: + display_lines_raw, _ = self._read_tail_lines(abs_path, encoding, lines_to_read) + else: + display_lines_raw = self._read_lines_range(abs_path, encoding, start, end) + + # 限制单行大小（防止单行超大导致内存问题） + MAX_LINE_LENGTH = 10000 + content_lines = [] + truncated_lines = 0 + for i, line in enumerate(display_lines_raw): + line_num = start + i + 1 + content = line.rstrip('\n\r') + if len(content) > MAX_LINE_LENGTH: + content = content[:MAX_LINE_LENGTH] + f"... (行超长，已截断，原始长度 {len(content)})" + truncated_lines += 1 + content_lines.append(f" {line_num:4d} | {content}") data = { "路径": abs_path, @@ -5876,17 +6053,13 @@ class ReadFileTool(FileToolBase): "显示范围": f"{range_desc} (行 {start+1}-{end})" } - if display_lines: - content_lines = [] - for i, line in enumerate(display_lines): - line_num = start + i + 1 - content = line.rstrip('\n\r') - content_lines.append(f" {line_num:4d} | {content}") - + if content_lines: data["内容"] = "\n" + "\n".join(content_lines) if end < line_count: data["提示"] = f"还有 {line_count - end} 行未显示" + if truncated_lines > 0: + data["提示"] += f"，{truncated_lines} 行因超长被截断" return self.format_result(True, "文件读取成功", data)