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:
qing committed 2026-08-01 03:46:02 +08:00
1 parent 90761a723d
commit 2f3e63e499
3 files changed
+494 -94

No files matched your search

+14 -2
View File
@@ -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 正确顺序加载
+215
View File
@@ -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
View File
@@ -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)