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
This commit is contained in:
1 parent
90761a723d
commit
2f3e63e499
3 files changed
+494
-94
No files matched your search
+14
-2
@@ -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 正确顺序加载
|
||||
|
||||
@@ -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
|
||||
+265
-92
@@ -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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user