diff --git a/.agents/design/code-sandbox/python-isolated-runner.md b/.agents/design/code-sandbox/python-isolated-runner.md new file mode 100644 index 000000000000..3df0c6fa0bb5 --- /dev/null +++ b/.agents/design/code-sandbox/python-isolated-runner.md @@ -0,0 +1,291 @@ +# Python Code Sandbox 隔离执行方案 + +## 背景 + +改造前，`projects/code-sandbox` 的 Python 执行路径使用长驻 `PythonProcessPool` +和 `worker.py`。同一个 Python 进程会多次执行用户代码，主要依赖 +`__import__` 白名单、受限 `open`、AST 检查、替换 builtin 等进程内限制。 + +GHSA-5jmh-5f2m-89jg 证明该模型存在结构性缺陷：Python 反射链可以绕过 +进程内 denylist，并最终获得容器内 OS 命令执行能力。继续补 AST 或字符串 +拦截只能覆盖已知 payload，不能作为多租户安全边界。 + +当前方案已经落地为 Python isolated one-shot warm pool： + +- 旧 `worker.py` 和 `PythonProcessPool` 已删除； +- `/sandbox/python` 统一使用 `PythonIsolatedRunner`； +- Python 进程最多执行一次用户代码，执行后销毁； +- Linux/Docker 环境固定启用 `chroot`、`no_new_privs`、`seccomp`、`setgid/setuid`； +- 网络请求统一走父进程 HTTP 代理，不允许 Python 子进程直接网络 syscall。 + +## 目标 + +- 保持 `/sandbox/python` API 兼容：成功返回 `{ success, data: { codeReturn, log } }`，失败返回 `{ success:false, message }`。 +- 兼容历史 Python Code 写法：`main()`、`main(variables)`、`main(a,b)`、全局变量注入、`print` 收集、内置 helper。 +- 多租户安全边界不依赖用户代码所在 Python 进程内的软限制。 +- 保留已有安全能力：模块白名单、文件限制、SSRF 防护、请求次数/大小限制、超时、输出限制、RSS 监控。 +- 降低冷启动延迟，但不复用执行过用户代码的 Python 解释器。 + +## 非目标 + +- 不移除 JS `ProcessPool`。 +- 不提供 Python pool 回滚模式。 +- 不把 AST 检查作为主要安全边界；它只作为纵深防御和兼容性辅助。 +- 不声称完整容器逃逸防护；OS 级隔离只约束当前 Python 执行进程的 syscall、根目录和权限。 + +## 总体架构 + +```text +POST /sandbox/python + -> queueIdLimiter.run(queueId, ...) + -> PythonIsolatedRunner.execute({ code, variables }) + -> 获取干净的 one-shot 预热 Python 进程 + -> 无空闲进程时按需 spawn + -> 预热进程已完成 chroot + no_new_privs + seccomp + setgid/setuid + -> bootstrap 注入 helper、variables、受限 builtins + -> exec 用户代码并调用 main + -> stdout 输出 JSON line result/http_request + -> 父进程解析结果、处理 HTTP 代理、监控超时/RSS/输出大小 + -> 该 Python 进程销毁，不归还池中 + -> 异步补充新的干净预热进程 +``` + +`/sandbox/modules` 仍返回 `env.SANDBOX_PYTHON_ALLOWED_MODULES`，表示 Python +用户可直接 import 的白名单模块，不暴露底层 runner 类型。 + +## 核心模块 + +### `src/isolated/python-isolated-runner.ts` + +`PythonIsolatedRunner` 负责父进程侧调度和资源控制： + +- 维护 one-shot warm pool，默认预热 `SANDBOX_POOL_SIZE` 个空闲 Python 进程； +- 预热进程只执行 `init` 协议，不执行用户代码； +- `execute()` 优先使用空闲预热进程，没有空闲进程时按需创建； +- 每个 Python 进程最多接受一条用户任务，任务结束后销毁； +- 执行完成后异步补充新的干净预热进程； +- 使用 `Semaphore` 控制 Python 任务最大并发，复用 `SANDBOX_POOL_SIZE`； +- 监控超时、输出大小和进程树 RSS； +- 清理整个进程树，避免子进程残留； +- 处理 Python 发起的 `http_request` IPC，并由父进程统一执行网络请求。 + +### `src/isolated/python-bootstrap.py` + +`python-bootstrap.py` 是 Python 子进程入口： + +- 支持两种协议： + - 兼容模式：直接读取一条 task JSON 并执行； + - 预热模式：先读取 `type:"init"`，完成 native 隔离后输出 `type:"ready"`，再等待一条 task JSON； +- 初始化 Python helper：`SystemHelper`、`system_helper`、`http_request`、`count_token`、`str_to_base64`、`create_hmac`、`delay`； +- 注入 `variables` 和历史全局变量写法； +- 安装受限 builtins、import 白名单、受限 `open`、危险属性拦截、audit hook； +- 捕获 `print` 到内存 log，避免污染 stdout JSON line 协议； +- 执行 `main()` / `main(variables)` / `main(a,b)`； +- 通过 stdout 输出 `result` 或 `http_request` JSON line。 + +### `native/python-sandbox` + +Go shared library `fastgpt_python_sandbox.so` 负责 Linux native 隔离： + +- `chroot` 到固定 Python sandbox root； +- `chdir("/")`； +- `setgroups([])`； +- `setgid(65537)` / `setuid(65537)`； +- `PR_SET_NO_NEW_PRIVS`； +- 安装 seccomp filter； +- 危险 syscall 默认拒绝，`execve/execveat`、`ptrace`、`mount` 等不能落地。 + +## 配置 + +Python 隔离相关配置已经收敛为内部安全默认，不提供运行时环境变量关闭或改弱： + +| 配置 | 当前值 | 说明 | +| --- | --- | --- | +| Python 最大并发 | `SANDBOX_POOL_SIZE` | 复用现有进程池大小配置 | +| 预热空闲进程数 | `SANDBOX_POOL_SIZE` | 与 Python 最大并发保持一致 | +| 直接网络 syscall | `false` | 统一走父进程 `http_request` 代理 | +| chroot 根目录 | `/tmp/fastgpt-python-sandbox` | Docker 构建阶段准备 | +| 用户代码 uid/gid | `65537:65537` | native 初始化后降权 | +| native seccomp/chroot/setuid | Linux 固定开启 | 缺少 native 库或 chroot root 时 fail-closed | + +保留的 Python 业务配置： + +| 变量 | 说明 | +| --- | --- | +| `SANDBOX_PYTHON_ALLOWED_MODULES` | 用户代码可直接 import 的 Python 模块白名单 | +| `SANDBOX_MAX_TIMEOUT` | 单次执行最大超时 | +| `SANDBOX_MAX_MEMORY_MB` | 子进程树 RSS 软限制 | +| `SANDBOX_MAX_OUTPUT_MB` | stdout/log 输出上限 | +| `SANDBOX_REQUEST_*` | 父进程 HTTP 代理的次数、超时、请求体和响应体限制 | + +## 网络代理 + +Python 子进程不允许直接使用网络 syscall。用户代码如需请求外部网络，必须调用： + +```python +http_request(url, method='GET', headers=None, body=None, timeout=None) +# 或 +system_helper.http_request(...) +SystemHelper.httpRequest(...) +``` + +Python bootstrap 向父进程写出： + +```json +{ + "type": "http_request", + "id": "http-1", + "payload": { + "url": "https://example.com", + "method": "GET", + "headers": {}, + "body": null, + "timeout": null + } +} +``` + +父进程执行： + +- URL 协议限制； +- SSRF / 内网地址检查； +- DNS pinning； +- 请求次数限制； +- 请求体大小限制； +- 响应体大小限制； +- 超时控制。 + +父进程再通过 stdin 返回： + +```json +{ + "type": "http_response", + "id": "http-1", + "success": true, + "payload": {} +} +``` + +每个 Python 进程只执行一次任务，因此请求次数计数天然按单次执行归零。 + +## 安全边界 + +### 语言层边界 + +语言层限制用于减少误用和拦截已知危险能力，但不是主安全边界： + +- `__import__` 白名单； +- `open` 受限； +- `eval` / `exec` / `compile` / `globals` / `locals` / `vars` / `dir` 等禁用； +- `__class__`、`__base__`、`__subclasses__`、`__globals__` 等危险属性拦截； +- `object` builtin 替换； +- audit hook 拦截 `os.system`、`subprocess`、`socket`、`ctypes` 等事件； +- AST 检查作为纵深防御。 + +### OS 层边界 + +OS 层是多租户核心安全边界： + +- 每个 Python 子进程最多执行一次用户代码； +- 执行用户代码前已经 chroot、降权、安装 seccomp； +- 执行完成后整个进程树销毁； +- 不复用执行过用户代码的解释器； +- seccomp 不允许命令执行和高危系统调用； +- chroot 只包含 Python 运行所需 stdlib、site-packages、动态库、证书、DNS 配置和 sandbox runtime 文件。 + +### 资源边界 + +- 父进程定时采样子进程树 RSS，超过 `SANDBOX_MAX_MEMORY_MB + RUNTIME_MEMORY_OVERHEAD_MB` 后杀进程树； +- 父进程控制总超时； +- stdout JSON line 和收集到的 log 受 `SANDBOX_MAX_OUTPUT_MB` 限制； +- 进程树清理由 `killProcessTree()` 处理，优先杀 descendant 和 process group。 + +## Docker 和构建 + +Docker 镜像使用 Debian bookworm/glibc。Alpine/musl 下 Go c-shared `.so` +存在兼容风险，不能作为当前 Python native isolation 运行基线。 + +构建流程： + +- `SANDBOX_BUILD_NATIVE_PYTHON=true pnpm build` 构建 Go shared library； +- `build.sh` 复制 `python-bootstrap.py` 和 `fastgpt_python_sandbox.so` 到 `dist`； +- Docker runner 阶段安装 Python、numpy、pandas、matplotlib 和 native 依赖； +- Docker 构建阶段准备 `/tmp/fastgpt-python-sandbox` chroot root； +- code-sandbox 主进程保留 root，以便 Python 子进程在 native 初始化阶段执行 chroot/setuid；用户代码进程会降权到 sandbox 用户。 + +## 测试覆盖 + +### 兼容性 + +- `main()`； +- `main(variables)`； +- `main(a,b)`； +- 全局变量注入； +- `print` log； +- 历史 helper； +- 旧 Python Code 节点写法。 + +### 安全 + +- GHSA-5jmh-5f2m-89jg 相关 `__base__` / `__subclasses__` 逃逸； +- 动态 `getattr`； +- import `os` / `sys` / `subprocess`； +- 文件系统访问； +- `os.system` / `subprocess` / `socket` / `ctypes`； +- 直接网络能力； +- 父进程 HTTP 代理 SSRF 防护； +- 请求大小、响应大小、请求次数、超时。 + +### 生命周期和资源 + +- init 后存在干净预热进程； +- 预热进程执行一次后销毁，不归还池中； +- 执行后自动补充新的干净预热进程； +- 并发超过上限时排队； +- 超时后可恢复； +- 内存超限后可恢复； +- shutdown 清理 running / idle / warming 子进程。 + +### Docker/Linux + +- native `.so` 加载； +- setuid/setgid 降权； +- chroot 生效； +- seccomp 阻断命令执行； +- numpy/pandas/matplotlib 在 chroot/seccomp 下可用； +- Docker 包可用性测试覆盖 Python 和 JS 白名单包。 + +## 性能和资源 + +one-shot warm pool 的收益主要在低并发和空闲命中场景： + +- 空闲预热进程命中时，省去 Python spawn/bootstrap/native init 的一部分延迟； +- 高并发短任务下，预热池可以覆盖 `SANDBOX_POOL_SIZE` 以内的首批请求，但每个进程执行一次后仍需销毁并补充，因此稳态吞吐仍不会接近旧长驻 pool； +- 预热进程不提前 import pandas/numpy，避免 idle 内存过高； +- 当前 Docker/seccomp 下，单个 idle Python bootstrap 进程 RSS 约 18.7MB； +- `SANDBOX_POOL_SIZE=20` 时，Python idle 进程 RSS 粗略约 374MB；实际容器 RSS 还包含 JS worker、Node 主进程、共享页和系统库统计口径。 + +安全优先级高于短任务吞吐。如果未来需要继续优化，应优先评估： + +- 是否为重包场景做专门的 package page-cache 预热，而不是让 idle 进程提前 import； +- 是否引入 clean forkserver。forkserver 父进程必须永不执行用户代码，且子进程在执行前完成 fd 清理、chroot、setuid 和 seccomp。 + +## 验收标准 + +- 旧 Python pool/worker 代码删除； +- `/sandbox/python` 只使用 `PythonIsolatedRunner`； +- Linux 缺少 native `.so` 或 chroot root 时 fail-closed； +- Python 子进程直接网络 syscall 固定关闭； +- 执行过用户代码的 Python 进程不会复用； +- 原有 Python 兼容用例通过； +- 原有安全边界测试迁移并通过； +- Docker/seccomp 下 Python 包可用性和 OS 隔离测试通过。 + +## 当前验证命令 + +```bash +pnpm --filter @fastgpt/code-sandbox exec tsc --noEmit +SANDBOX_MAX_MEMORY_MB=256 pnpm --filter @fastgpt/code-sandbox exec vitest run --coverage.enabled=false +pnpm --filter @fastgpt/code-sandbox build +docker build --build-arg proxy=1 -f projects/code-sandbox/Dockerfile -t fastgpt-code-sandbox-warm-pool . +``` diff --git a/.agents/issue/agent-context-tool-compression-analysis.md b/.agents/issue/agent-context-tool-compression-analysis.md new file mode 100644 index 000000000000..cd819d317cad --- /dev/null +++ b/.agents/issue/agent-context-tool-compression-analysis.md @@ -0,0 +1,164 @@ +# Agent 上下文和工具压缩逻辑分析 + +日期：2026-06-23 + +## 结论摘要 + +当前 Agent 上下文链路已经从旧的 `compressed_messages: ChatCompletionMessageParam[]` 转为 checkpoint 压缩模式：历史消息在超过阈值后被压成一条隐藏的 user message，并通过 `contextCheckpoint` 写入 AI history。工具结果压缩仍是单次 tool response 级别的压缩，执行后作为 `tool` message 回灌到同一条 agent loop 消息链。 + +整体设计方向是正确的：避免历史 `assistant.tool_calls` / `tool` message 被 LLM 改坏配对关系，同时保留 ask resume、plan、tool result 的连续上下文。但当前实现里有一个需要优先确认的风险：`compressRequestMessages` 的结构化工具 checkpoint 分支不产生 `usage`，而 `onCompressContext` 只有存在 `result.usage` 才返回压缩结果，导致这条无 LLM 压缩路径在 agent loop 中可能被忽略。 + +## 相关模块地图 + +| 模块 | 职责 | +| --- | --- | +| `packages/service/core/workflow/dispatch/ai/agent/index.ts` | Workflow Agent 节点入口，准备历史、用户上下文、工具、sandbox，并调用 unified loop。 | +| `packages/service/core/ai/llm/agentLoop/loop/unified.ts` | 单主 Agent Loop 适配层，注入 `ask_agent`、`update_plan` 和 runtime tools，处理 stop gate。 | +| `packages/service/core/ai/llm/agentLoop/loop/base.ts` | 底层循环：每轮请求前压缩上下文，请求 LLM，执行工具，压缩工具结果，回灌 tool message。 | +| `packages/service/core/ai/llm/compress/index.ts` | 压缩实现：历史 checkpoint、通用长文本压缩、JSON 工具结果结构摘要、tool response 压缩。 | +| `packages/service/core/workflow/dispatch/ai/agent/adapter/eventMapper.ts` | 将 loop 事件写入 `assistantResponses` 和 SSE；`after_message_compress` 在这里写 checkpoint。 | +| `packages/global/core/chat/adapt.ts` | history -> GPT messages 适配；识别最新 checkpoint，丢弃 checkpoint 前普通历史。 | +| `packages/service/core/workflow/dispatch/utils/index.ts` | 按节点 history 配置裁剪历史；存在 checkpoint 时优先从 checkpoint 开始保留。 | + +## 上下文构造链路 + +1. `dispatchRunAgent` 通过 `useUserContext` 拿到 `chatHistories`、改写后的历史和当前用户消息。 +2. `chats2GPTMessages({ reserveTool: true })` 将 FastGPT history 转为 LLM messages，并保留 agent/tool 结构。 +3. `runUnifiedAgentLoop` 注入 Main Agent system prompt，过滤历史里的 system message，组成初始 messages。 +4. `runAgentLoop` 每轮请求前调用 `onCompressContext`，由 `compressRequestMessages` 判断是否压缩。 +5. LLM 如果调用 runtime tool，工具结果会变成 `tool` message 追加回 `requestMessages`，下一轮继续沿同一条消息链请求。 +6. 如果触发 `ask_agent`，`pendingMainContext.messages` 会保存当时 messages；用户回答后作为对应 ask tool response 接回原链路。 + +关键代码： + +- `runAgentLoop` 每轮请求前压缩 request messages：`packages/service/core/ai/llm/agentLoop/loop/base.ts:286` +- LLM 请求使用压缩后的 `requestMessages`：`packages/service/core/ai/llm/agentLoop/loop/base.ts:331` +- ask resume 从 `pendingMainContext.messages` 接回 tool response：`packages/service/core/ai/llm/agentLoop/loop/unified.ts` + +## 历史 checkpoint 压缩 + +触发逻辑在 `compressRequestMessages`： + +1. 先拆出 `system/developer` 与其它消息。系统类消息不参与摘要，但最终保留在最前面。 +2. 使用完整 messages 计算 token，超过 `model.maxContext * 0.8` 才触发历史压缩。 +3. 优先尝试结构化工具 checkpoint：从历史 tool_calls / tool result 中确定性生成 checkpoint。 +4. 如果不能使用结构化路径，则调用 LLM 压缩为 `...`。 +5. 压缩结果作为 `{ role: user, hideInUI: true }` message 返回。 +6. 若 LLM 输出仍超阈值，尝试确定性 head-tail checkpoint 兜底；仍超限则返回原始 messages。 + +关键代码： + +- 拆分 system/developer 和其它消息：`packages/service/core/ai/llm/compress/index.ts:720` +- 80% 阈值判断：`packages/service/core/ai/llm/compress/index.ts:742` +- 结构化工具 checkpoint 分支：`packages/service/core/ai/llm/compress/index.ts:755` +- LLM checkpoint 压缩：`packages/service/core/ai/llm/compress/index.ts:791` +- 返回 checkpoint：`packages/service/core/ai/llm/compress/index.ts:914` + +## checkpoint 持久化和恢复 + +checkpoint 不在 `dispatchRunAgent` 末尾显式追加，而是通过 loop 事件写入： + +1. `runAgentLoop` 压缩成功后触发 `onAfterCompressContext`。 +2. `runUnifiedAgentLoop` 转发为 `after_message_compress` 事件。 +3. `eventMapper` 收到事件后向 `assistantResponses` push `{ contextCheckpoint, hideInUI: true }`。 +4. 本轮 chat 保存时该 value 随 AI history 落库。 +5. 下一轮 `getHistories` 发现 AI history 中有 checkpoint 时，从最新 checkpoint 所在 history 开始保留，避免先按最近 N 轮裁掉 checkpoint。 +6. `chats2GPTMessages` 再次从最新 checkpoint value 精确切片，把 checkpoint 转为隐藏 user message，并跳过同一 value 的其它字段。 + +关键代码： + +- 事件写入 checkpoint value：`packages/service/core/workflow/dispatch/ai/agent/adapter/eventMapper.ts:374` +- history 裁剪保留 checkpoint：`packages/service/core/workflow/dispatch/utils/index.ts:387` +- adapter 查找最新 checkpoint：`packages/global/core/chat/adapt.ts:73` +- checkpoint 转 hidden user message：`packages/global/core/chat/adapt.ts:479` + +## 工具压缩链路 + +工具执行和压缩在 `runAgentLoop` 内完成： + +1. LLM 产出 tool_calls 后，先把 assistant tool_calls message 追加到 `requestMessages`。 +2. 对每个 tool 执行 `onRunTool`，runtime 内部工具如 `ask_agent` / `update_plan` 会设置 `skipResponseCompress`。 +3. 对普通 runtime tool，调用 `compressToolResponse` 压缩结果。 +4. 压缩后的内容写成 `tool` message，追加到 `requestMessages` 和 `assistantMessages`。 +5. `onAfterToolCall` 把压缩后的 response 和压缩详情传给 workflow adapter，用于工具卡和运行详情。 + +关键代码： + +- 工具执行入口：`packages/service/core/ai/llm/agentLoop/loop/base.ts:425` +- 跳过内部工具压缩：`packages/service/core/ai/llm/agentLoop/loop/base.ts:461` +- 调用 `compressToolResponse`：`packages/service/core/ai/llm/agentLoop/loop/base.ts:469` +- tool message 回灌：`packages/service/core/ai/llm/agentLoop/loop/base.ts:506` + +`compressToolResponse` 的预算策略： + +1. 固定上限：`model.maxContext * 0.5`。 +2. 动态上限：`(model.maxContext - currentMessagesTokens) / toolLength`，避免并行工具结果整体打爆上下文。 +3. 调用方自定义上限。 +4. 三者取最小值。 +5. JSON 工具结果优先走本地结构摘要；否则走通用 `compressLargeContent`。 + +关键代码： + +- 工具压缩预算计算：`packages/service/core/ai/llm/compress/index.ts:1299` +- JSON 本地摘要优先：`packages/service/core/ai/llm/compress/index.ts:1316` +- 通用长文本压缩：`packages/service/core/ai/llm/compress/index.ts:1326` + +## 当前风险点 + +### R1：结构化工具 checkpoint 在 agent loop 中可能不生效 + +`compressRequestMessages` 的结构化工具 checkpoint 分支返回： + +```ts +return { + messages: finalStructuredMessages, + contextCheckpoint: structuredToolCheckpoint +}; +``` + +该返回没有 `usage`。但 `onCompressContext` 只有 `if (result.usage)` 才返回压缩结果。结果是：结构化 checkpoint 虽然在 `compressRequestMessages` 内生成了，但 `runAgentLoop` 不会替换 `requestMessages`，也不会向外传播 `contextCheckpoint`。 + +影响： + +- 工具调用历史很长时，本地确定性压缩路径可能被静默跳过。 +- 只能依赖后续 LLM checkpoint 分支；但当前代码在结构化分支成功后直接 return，不会落到 LLM 分支。 + +建议： + +- `onCompressContext` 应在 `result.messages !== requestMessages` 或 `result.contextCheckpoint` 存在时也返回压缩结果。 +- usage 可选；调用处 `usagePush` 和 `onAfterCompressContext` 需要允许无 usage 的压缩事件，或为本地压缩生成 0 usage 记录。 +- 增加 base loop 级测试，覆盖 `compressRequestMessages` 返回无 usage 但有 `contextCheckpoint` 的情况。 + +### R2：工具压缩动态预算可能为 0 + +`availableCompressedTokenLimit = max(0, floor((maxContext - currentMessagesTokens) / toolLength))`。当当前 messages 已接近或超过 maxContext 时，工具结果压缩目标可能为 0。后续 `compressLargeContent` 是否能稳定处理 0 token 预算，需要专项测试。 + +建议： + +- 设置最小压缩预算，例如 256 或 512 token；若连最小预算都无空间，应先触发 request message checkpoint，再执行/回灌工具结果。 +- 增加工具结果压缩预算为 0 的单测。 + +### R3：结构化 checkpoint 分支缺少运行详情事件 + +结构化 checkpoint 是本地压缩，不会产生 requestId 和 usage。即使修复 R1，也需要决定是否在运行详情里显示“本地上下文压缩”。否则用户只能看到上下文突然变短，缺少可观测性。 + +建议： + +- 若前端需要可观测性，可扩展 `after_message_compress` 事件，允许 `compressionMode: 'structured_tool_checkpoint'` 和 0 usage。 + +### R4：checkpoint 保存顺序依赖事件时机 + +当前 checkpoint 通过 `after_message_compress` 事件即时 push 到 `assistantResponses`。这能保证压缩发生在本轮中间时，checkpoint 排在后续 plan/tool/text value 之前。但如果未来某些压缩路径不触发事件，只在 `result.contextCheckpoint` 返回，正常完成路径不会兜底保存。 + +建议： + +- 明确约定：所有可持久化 checkpoint 必须通过 `after_message_compress` 写入。 +- 或在 `dispatchRunAgent` done/ask 分支增加去重兜底，避免事件丢失导致 checkpoint 不落库。 + +## 建议 TODO + +- [ ] 修复 `onCompressContext` 对无 usage checkpoint 的忽略问题。 +- [ ] 增加 base loop 测试：无 usage 的 structured checkpoint 应替换 request messages 并返回 `contextCheckpoint`。 +- [ ] 增加 workflow dispatch 集成测试：checkpoint value 写入顺序为 `checkpoint -> 后续 plan/tool/text`。 +- [ ] 增加工具压缩动态预算为 0 或极小值时的测试。 +- [ ] 明确本地结构化压缩是否需要运行详情展示。 diff --git a/document/content/self-host/upgrading/4-15/41505.mdx b/document/content/self-host/upgrading/4-15/41505.mdx index 2e8bf44988ce..6e32980a152f 100644 --- a/document/content/self-host/upgrading/4-15/41505.mdx +++ b/document/content/self-host/upgrading/4-15/41505.mdx @@ -57,6 +57,8 @@ curl --location --request POST 'https://{{host}}/api/admin/initSandboxArchive' \ 4. 知识库训练出现错误时的提示，同时支持一键全部重试。 5. 过滤掉无效的知识库引用角标。 6. 工具运行空响应时候，自动补充 "none"，避免部分模型报错。 +7. 系统工具运行前，再次进行二次权限校验。 +8. 优化重定向后 SSRF 校验。 ## 🐛 修复 diff --git a/document/data/doc-last-modified.json b/document/data/doc-last-modified.json index 7e2a6bf7353f..758fef7ff451 100644 --- a/document/data/doc-last-modified.json +++ b/document/data/doc-last-modified.json @@ -27,8 +27,8 @@ "content/guide/build/publish/mcp_server.mdx": "2026-06-22T11:01:59+08:00", "content/guide/build/publish/official_account.en.mdx": "2026-05-07T15:06:40+08:00", "content/guide/build/publish/official_account.mdx": "2026-05-07T15:06:40+08:00", - "content/guide/build/publish/openapi.en.mdx": "2026-06-23T11:27:33+08:00", - "content/guide/build/publish/openapi.mdx": "2026-06-23T11:27:33+08:00", + "content/guide/build/publish/openapi.en.mdx": "2026-06-23T13:54:06+08:00", + "content/guide/build/publish/openapi.mdx": "2026-06-23T13:54:06+08:00", "content/guide/build/publish/wechat.en.mdx": "2026-05-07T15:06:40+08:00", "content/guide/build/publish/wechat.mdx": "2026-05-07T15:06:40+08:00", "content/guide/build/publish/wecom.en.mdx": "2026-05-07T15:06:40+08:00", @@ -125,8 +125,8 @@ "content/guide/version/cloud/privacy.mdx": "2026-05-07T15:06:40+08:00", "content/guide/version/cloud/terms.en.mdx": "2026-05-07T15:06:40+08:00", "content/guide/version/cloud/terms.mdx": "2026-05-28T13:54:38+08:00", - "content/guide/version/commercial.en.mdx": "2026-05-07T15:06:40+08:00", - "content/guide/version/commercial.mdx": "2026-05-07T15:06:40+08:00", + "content/guide/version/commercial.en.mdx": "2026-06-23T15:18:25+08:00", + "content/guide/version/commercial.mdx": "2026-06-23T15:18:25+08:00", "content/guide/version/opensource/intro.en.mdx": "2026-05-07T15:06:40+08:00", "content/guide/version/opensource/intro.mdx": "2026-05-07T15:06:40+08:00", "content/guide/version/opensource/license.en.mdx": "2026-05-07T15:06:40+08:00", @@ -139,14 +139,14 @@ "content/guide/workspace/team/team_roles_permissions.mdx": "2026-05-07T15:06:40+08:00", "content/openapi/app.en.mdx": "2026-05-29T19:31:16+08:00", "content/openapi/app.mdx": "2026-05-29T19:31:16+08:00", - "content/openapi/chat.en.mdx": "2026-06-23T11:27:33+08:00", - "content/openapi/chat.mdx": "2026-06-23T12:10:31+08:00", + "content/openapi/chat.en.mdx": "2026-06-23T13:54:06+08:00", + "content/openapi/chat.mdx": "2026-06-23T13:54:06+08:00", "content/openapi/dataset.en.mdx": "2026-05-29T19:31:16+08:00", "content/openapi/dataset.mdx": "2026-05-29T19:31:16+08:00", "content/openapi/index.en.mdx": "2026-04-26T21:08:47+08:00", "content/openapi/index.mdx": "2026-04-26T21:08:47+08:00", - "content/openapi/intro.en.mdx": "2026-06-23T11:27:33+08:00", - "content/openapi/intro.mdx": "2026-06-23T12:10:31+08:00", + "content/openapi/intro.en.mdx": "2026-06-23T13:54:06+08:00", + "content/openapi/intro.mdx": "2026-06-23T13:54:06+08:00", "content/plugin/index.en.mdx": "2026-06-04T16:10:15+08:00", "content/plugin/index.mdx": "2026-06-04T16:10:15+08:00", "content/plugin/intro.en.mdx": "2026-06-09T16:03:58+08:00", @@ -155,8 +155,8 @@ "content/plugin/model-presets.mdx": "2026-06-04T16:10:15+08:00", "content/plugin/system-tool-development.en.mdx": "2026-06-09T16:03:58+08:00", "content/plugin/system-tool-development.mdx": "2026-06-09T16:03:58+08:00", - "content/self-host/config/env.en.mdx": "2026-06-22T11:01:59+08:00", - "content/self-host/config/env.mdx": "2026-06-22T11:01:59+08:00", + "content/self-host/config/env.en.mdx": "2026-06-23T13:54:06+08:00", + "content/self-host/config/env.mdx": "2026-06-23T13:54:06+08:00", "content/self-host/config/json.en.mdx": "2026-06-22T11:01:59+08:00", "content/self-host/config/json.mdx": "2026-06-22T11:01:59+08:00", "content/self-host/config/model/intro.en.mdx": "2026-06-04T16:10:15+08:00", @@ -277,17 +277,17 @@ "content/self-host/upgrading/4-14/41481.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/4-14/4149.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/4-14/4149.mdx": "2026-04-26T21:08:47+08:00", - "content/self-host/upgrading/4-15/41500.en.mdx": "2026-06-23T11:27:33+08:00", - "content/self-host/upgrading/4-15/41500.mdx": "2026-06-23T11:27:33+08:00", + "content/self-host/upgrading/4-15/41500.en.mdx": "2026-06-23T13:54:06+08:00", + "content/self-host/upgrading/4-15/41500.mdx": "2026-06-23T13:54:06+08:00", "content/self-host/upgrading/4-15/41501.mdx": "2026-06-22T11:01:59+08:00", "content/self-host/upgrading/4-15/41502.en.mdx": "2026-05-25T11:21:30+08:00", - "content/self-host/upgrading/4-15/41502.mdx": "2026-06-22T22:05:59+08:00", + "content/self-host/upgrading/4-15/41502.mdx": "2026-06-23T13:54:06+08:00", "content/self-host/upgrading/4-15/41503.en.mdx": "2026-05-28T16:21:09+08:00", "content/self-host/upgrading/4-15/41503.mdx": "2026-05-28T16:21:09+08:00", "content/self-host/upgrading/4-15/41504.en.mdx": "2026-06-10T19:02:59+08:00", "content/self-host/upgrading/4-15/41504.mdx": "2026-06-15T23:34:43+08:00", - "content/self-host/upgrading/4-15/41505.en.mdx": "2026-06-12T20:47:04+08:00", - "content/self-host/upgrading/4-15/41505.mdx": "2026-06-23T11:27:33+08:00", + "content/self-host/upgrading/4-15/41505.en.mdx": "2026-06-23T13:54:06+08:00", + "content/self-host/upgrading/4-15/41505.mdx": "2026-06-23T13:54:06+08:00", "content/self-host/upgrading/outdated/40.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/40.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/41.en.mdx": "2026-04-26T21:08:47+08:00", @@ -355,7 +355,7 @@ "content/self-host/upgrading/outdated/4810.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/4810.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/4811.en.mdx": "2026-04-26T21:08:47+08:00", - "content/self-host/upgrading/outdated/4811.mdx": "2026-06-22T22:05:59+08:00", + "content/self-host/upgrading/outdated/4811.mdx": "2026-06-23T13:54:06+08:00", "content/self-host/upgrading/outdated/4812.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/4812.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/4813.en.mdx": "2026-05-07T15:06:40+08:00", diff --git a/packages/global/core/ai/sandbox/constants.ts b/packages/global/core/ai/sandbox/constants.ts index 5d37c797482d..0cbb05db27a2 100644 --- a/packages/global/core/ai/sandbox/constants.ts +++ b/packages/global/core/ai/sandbox/constants.ts @@ -16,7 +16,7 @@ export const SandboxStatusEnum = { export type SandboxStatusType = (typeof SandboxStatusEnum)[keyof typeof SandboxStatusEnum]; // ---- 暂停阈值（分钟） ---- -export const SANDBOX_SUSPEND_MINUTES = 5; +export const SANDBOX_SUSPEND_MINUTES = 10; // ---- sandboxId 生成 ---- export const generateSandboxId = (appId: string, userId: string, chatId: string): string => { diff --git a/packages/global/core/workflow/template/system/agent/index.ts b/packages/global/core/workflow/template/system/agent/index.ts index 567daa3fbfd9..ac45c2a93f95 100644 --- a/packages/global/core/workflow/template/system/agent/index.ts +++ b/packages/global/core/workflow/template/system/agent/index.ts @@ -28,7 +28,7 @@ export const AgentNode: FlowNodeTemplateType = { templateType: FlowNodeTemplateTypeEnum.ai, showSourceHandle: true, showTargetHandle: true, - avatar: 'core/workflow/template/agent', + avatar: 'core/app/type/agentFill', avatarLinear: 'core/workflow/template/agentLinear', colorSchema: 'emerald', name: i18nT('workflow:template.agent_module'), diff --git a/packages/service/common/api/axios.ts b/packages/service/common/api/axios.ts index 1bbc7b835431..e52433c558d1 100644 --- a/packages/service/common/api/axios.ts +++ b/packages/service/common/api/axios.ts @@ -1,26 +1,84 @@ -import _, { type AxiosInstance, type AxiosRequestConfig } from 'axios'; +import _, { + type AxiosInstance, + type AxiosRequestConfig, + type AxiosResponse, + type InternalAxiosRequestConfig +} from 'axios'; import { ProxyAgent, type ProxyAgentOptions } from 'proxy-agent'; import { isDevEnv } from '@fastgpt/global/common/system/constants'; import { isInternalAddress, PRIVATE_URL_TEXT } from '../system/utils'; import { isAbsoluteUrl } from '../security/network'; import { SERVICE_LOCAL_HOST } from '../system/tools'; +/** + * 给 shared axios 实例添加 SSRF 防护。 + * + * 这里同时接管 axios 的重定向逻辑: axios/follow-redirects 的自动跳转发生在 + * request interceptor 之后,如果不手动处理,302 Location 指向内网时不会再次进入 + * isInternalAddress()。因此 safe axios 会强制关闭底层自动跳转,再在 response + * interceptor 中逐跳解析 Location、复用同一套 SSRF 策略校验后再发起下一跳请求。 + */ const addSSRFInterceptor = (instance: AxiosInstance) => { - instance.interceptors.request.use(async (config) => { - const requestUrl = (() => { - try { - return new URL(config.url || '', config.baseURL).toString(); - } catch { - return; - } - })(); + instance.interceptors.request.use(async (config): Promise => { + const safeConfig = config as SafeRedirectInternalConfig; + const requestUrl = buildRequestUrl(safeConfig); if (!requestUrl) return config; if (await isInternalAddress(requestUrl)) { return Promise.reject(new Error(PRIVATE_URL_TEXT)); } - return config; + const maxRedirects = + safeConfig.__safeRedirect?.maxRedirects ?? + (typeof safeConfig.maxRedirects === 'number' + ? safeConfig.maxRedirects + : SAFE_AXIOS_MAX_REDIRECTS); + + const nextConfig: SafeRedirectInternalConfig = { + ...safeConfig, + // 禁用底层自动跳转,保留调用方 maxRedirects 语义给手动跳转状态使用。 + maxRedirects: 0, + validateStatus: getRedirectValidateStatus(safeConfig.validateStatus, maxRedirects), + __safeRedirect: safeConfig.__safeRedirect ?? { + count: 0, + maxRedirects, + validateStatus: safeConfig.validateStatus + } + }; + + return nextConfig; + }); + + instance.interceptors.response.use(async (response) => { + if (!shouldRedirect(response)) return response; + + const config = response.config as SafeRedirectConfig; + const redirectState = config.__safeRedirect; + const currentUrl = buildRequestUrl(config); + // 理论上 request interceptor 会注入状态；缺失时按普通响应返回,避免误处理其它实例响应。 + if (!redirectState || !currentUrl) return response; + + if (redirectState.count >= redirectState.maxRedirects) { + return Promise.reject(new Error(`Maximum redirects exceeded: ${redirectState.maxRedirects}`)); + } + + const redirectUrl = resolveRedirectUrl(response.headers.location, currentUrl); + if (await isInternalAddress(redirectUrl)) { + return Promise.reject(new Error(PRIVATE_URL_TEXT)); + } + + const redirectConfig = getRedirectConfig(config, response, currentUrl, redirectUrl); + + const nextConfig: SafeRedirectConfig = { + ...redirectConfig, + validateStatus: redirectState.validateStatus, + __safeRedirect: { + ...redirectState, + count: redirectState.count + 1 + } + }; + + return instance.request(nextConfig); }); return instance; @@ -28,6 +86,166 @@ const addSSRFInterceptor = (instance: AxiosInstance) => { const createProxyAgent = (options?: ProxyAgentOptions) => new ProxyAgent(options); +const SAFE_AXIOS_MAX_REDIRECTS = 5; +const REDIRECT_STATUS_CODES = new Set([301, 302, 303, 307, 308]); + +/** + * 手动重定向状态。 + * + * count/maxRedirects 用于替代 axios/follow-redirects 的跳转次数控制。 + * validateStatus 保存调用方原始成功状态判定,因为 request interceptor 会临时允许 + * 3xx 进入 response interceptor,最终非重定向响应仍需按调用方语义判断成功/失败。 + */ +type SafeRedirectState = { + count: number; + maxRedirects: number; + validateStatus?: AxiosRequestConfig['validateStatus']; +}; + +type SafeRedirectConfig = AxiosRequestConfig & { + __safeRedirect?: SafeRedirectState; +}; + +type SafeRedirectInternalConfig = InternalAxiosRequestConfig & { + __safeRedirect?: { + count: number; + maxRedirects: number; + validateStatus?: AxiosRequestConfig['validateStatus']; + }; +}; + +const shouldRedirect = (response: AxiosResponse): boolean => + REDIRECT_STATUS_CODES.has(response.status) && typeof response.headers.location === 'string'; + +/** + * 按 axios baseURL + url 规则合成实际请求 URL。 + * + * 失败时返回 undefined,保持旧拦截器对非法/非标准 URL 的宽容行为； + * 真正的请求错误仍交给 axios 自身处理。 + */ +const buildRequestUrl = (config: AxiosRequestConfig): string | undefined => { + try { + return new URL(config.url || '', config.baseURL).toString(); + } catch { + return; + } +}; + +/** + * 解析 Location 头为下一跳绝对 URL,并限制协议。 + * + * file://、gopher:// 等非 HTTP 协议不能进入后续请求流程,否则 SSRF 防护的 + * 地址策略和 axios 出站语义都会变得不明确。 + */ +const resolveRedirectUrl = (location: string, currentUrl: string): string => { + const redirectUrl = new URL(location, currentUrl); + if (redirectUrl.protocol !== 'http:' && redirectUrl.protocol !== 'https:') { + throw new Error(`Unsupported redirect protocol: ${redirectUrl.protocol}`); + } + return redirectUrl.toString(); +}; + +/** + * 生成重定向请求头。 + * + * 行为对齐 follow-redirects 的安全取向: + * - Host 必须丢弃,由下一跳真实目标重新生成 + * - 301/302 POST 和 303 非 GET/HEAD 转 GET 时,不能继续携带 content-* 请求体头 + * - 跨 host/protocol 跳转时移除凭证类 header,避免用户配置的密钥被带到新域名 + */ +const filterRedirectHeaders = ({ + headers, + currentUrl, + redirectUrl, + shouldSwitchToGet +}: { + headers: AxiosRequestConfig['headers']; + currentUrl: string; + redirectUrl: string; + shouldSwitchToGet: boolean; +}): AxiosRequestConfig['headers'] => { + const nextHeaders = { ...(headers as Record) }; + const current = new URL(currentUrl); + const redirect = new URL(redirectUrl); + const shouldDropSensitiveHeaders = + current.protocol !== redirect.protocol || current.host !== redirect.host; + + for (const key of Object.keys(nextHeaders)) { + const normalizedKey = key.toLowerCase(); + if (normalizedKey === 'host') { + delete nextHeaders[key]; + continue; + } + if (shouldSwitchToGet && normalizedKey.startsWith('content-')) { + delete nextHeaders[key]; + continue; + } + if ( + shouldDropSensitiveHeaders && + ['authorization', 'cookie', 'proxy-authorization'].includes(normalizedKey) + ) { + delete nextHeaders[key]; + } + } + + return nextHeaders; +}; + +/** + * 根据 HTTP 重定向响应构造下一跳请求配置。 + * + * 仅实现服务端常见 301/302/303/307/308 语义: + * - 301/302 且原方法为 POST 时转 GET + * - 303 且原方法不是 GET/HEAD 时转 GET + * - 307/308 保留原方法和请求体 + */ +const getRedirectConfig = ( + config: AxiosRequestConfig, + response: AxiosResponse, + currentUrl: string, + redirectUrl: string +): AxiosRequestConfig => { + const currentMethod = (config.method || 'get').toUpperCase(); + const shouldSwitchToGet = + ((response.status === 301 || response.status === 302) && currentMethod === 'POST') || + (response.status === 303 && currentMethod !== 'GET' && currentMethod !== 'HEAD'); + + return { + ...config, + baseURL: undefined, + url: redirectUrl, + maxRedirects: 0, + method: shouldSwitchToGet ? 'GET' : config.method, + data: shouldSwitchToGet ? undefined : config.data, + headers: filterRedirectHeaders({ + headers: config.headers, + currentUrl, + redirectUrl, + shouldSwitchToGet + }) + }; +}; + +/** + * 临时放行重定向状态码,让 response interceptor 能拿到 3xx 响应并校验 Location。 + * + * maxRedirects=0 表示调用方明确禁止跟随重定向,此时完全保留调用方 validateStatus, + * 不把 3xx 转成“成功响应”。 + */ +const getRedirectValidateStatus = ( + validateStatus: AxiosRequestConfig['validateStatus'], + maxRedirects: number +): AxiosRequestConfig['validateStatus'] => { + if (maxRedirects === 0) { + return validateStatus; + } + + return (status) => { + if (REDIRECT_STATUS_CODES.has(status)) return true; + return validateStatus ? validateStatus(status) : status >= 200 && status < 300; + }; +}; + /** * 工作流 HTTP 节点跳过 HTTPS 证书校验专用 agent。 * 仍复用 ProxyAgent,只调整目标站 TLS 校验策略,避免改变部署环境的代理语义。 diff --git a/packages/service/core/ai/llm/agentLoop/constants.ts b/packages/service/core/ai/llm/agentLoop/constants.ts index c7411bfd3440..7f5119c5d2b2 100644 --- a/packages/service/core/ai/llm/agentLoop/constants.ts +++ b/packages/service/core/ai/llm/agentLoop/constants.ts @@ -9,7 +9,7 @@ export const AgentUsageModuleName = { export const AgentNodeResponseDisplay = { master: { moduleName: i18nT('chat:master_agent_call'), - moduleLogo: 'core/workflow/template/agent' + moduleLogo: 'core/app/type/agentFill' }, piMaster: { moduleName: i18nT('chat:master_agent_call'), diff --git a/packages/service/core/ai/llm/agentLoop/loop/base.ts b/packages/service/core/ai/llm/agentLoop/loop/base.ts index 1f58a59044c9..e89cefb096d5 100644 --- a/packages/service/core/ai/llm/agentLoop/loop/base.ts +++ b/packages/service/core/ai/llm/agentLoop/loop/base.ts @@ -42,7 +42,7 @@ type RunAgentCallProps = { childrenInteractiveParams?: AgentLoopChildrenInteractiveParams; // LLM 压缩后回调 onAfterCompressContext?: (e: { - usage: ChatNodeUsageType; + usage?: ChatNodeUsageType; requestIds: string[]; seconds: number; contextCheckpoint?: ContextCheckpointValueType; @@ -124,34 +124,50 @@ type RunAgentResponse = { */ export const onCompressContext = async ({ isAborted, + messageTokens, requestMessages, modelData, reasoningEffort, + tools, userKey }: { isAborted: RunAgentCallProps['isAborted']; + messageTokens?: number; requestMessages: ChatCompletionMessageParam[]; modelData: LLMModelItemType; reasoningEffort?: CreateLLMResponseProps['body']['reasoning_effort']; + tools?: ChatCompletionTool[]; userKey: RunAgentCallProps['userKey']; }) => { const compressStartTime = Date.now(); const result = await compressRequestMessages({ checkIsStopping: isAborted, + messageTokens, messages: requestMessages, model: modelData, reasoningEffort, + tools, userKey }); - if (result.usage) { + if (result.usage || result.contextCheckpoint || result.messages !== requestMessages) { return { messages: result.messages, + messageTokens: result.messageTokens, + hasCompressedContext: true, usage: result.usage, requestIds: result.requestIds ?? [], seconds: +((Date.now() - compressStartTime) / 1000).toFixed(2), contextCheckpoint: result.contextCheckpoint }; } + + return { + messages: result.messages, + messageTokens: result.messageTokens, + hasCompressedContext: false, + requestIds: result.requestIds ?? [], + seconds: +((Date.now() - compressStartTime) / 1000).toFixed(2) + }; }; /** @@ -210,6 +226,20 @@ export const runAgentLoop = async ({ } return item; }); + let requestMessagesTokenCount: number | undefined; + + /** + * requestMessages 在未压缩时主要是 append-only。 + * 已知基线 token 后，只统计新增 message，避免每轮压缩判断都重复统计完整历史。 + */ + const appendRequestMessages = async (...newMessages: ChatCompletionMessageParam[]) => { + requestMessages.push(...newMessages); + if (requestMessagesTokenCount !== undefined && newMessages.length > 0) { + requestMessagesTokenCount += await countGptMessagesTokens({ + messages: newMessages + }); + } + }; let inputTokens: number = 0; let outputTokens: number = 0; @@ -287,21 +317,28 @@ export const runAgentLoop = async ({ { const compressResult = await onCompressContext({ isAborted, + messageTokens: requestMessagesTokenCount, requestMessages, modelData, reasoningEffort: body.reasoning_effort, + tools: body.tools, userKey }); if (compressResult) { - requestMessages = compressResult.messages; - contextCheckpoint = compressResult.contextCheckpoint ?? contextCheckpoint; - usagePush?.([compressResult.usage]); - onAfterCompressContext?.({ - usage: compressResult.usage, - requestIds: compressResult.requestIds, - seconds: compressResult.seconds, - contextCheckpoint: compressResult.contextCheckpoint - }); + requestMessagesTokenCount = compressResult.messageTokens ?? requestMessagesTokenCount; + if (compressResult.hasCompressedContext) { + requestMessages = compressResult.messages; + contextCheckpoint = compressResult.contextCheckpoint ?? contextCheckpoint; + if (compressResult.usage) { + usagePush?.([compressResult.usage]); + } + onAfterCompressContext?.({ + usage: compressResult.usage, + requestIds: compressResult.requestIds, + seconds: compressResult.seconds, + contextCheckpoint: compressResult.contextCheckpoint + }); + } } } @@ -408,7 +445,7 @@ export const runAgentLoop = async ({ // 推送 AI 生成后的 assistantMessages if (llmAssistantMessage) { assistantMessages.push(llmAssistantMessage); - requestMessages.push(llmAssistantMessage); + await appendRequestMessages(llmAssistantMessage); } } @@ -423,11 +460,9 @@ export const runAgentLoop = async ({ }; const runTool = async ({ - tool, - currentMessagesTokens + tool }: { tool: ChatCompletionMessageToolCall; - currentMessagesTokens: number; }): Promise => { const toolStartTime = Date.now(); let toolErrorMessage: string | undefined; @@ -470,8 +505,6 @@ export const runAgentLoop = async ({ const compressionResult = await compressToolResponse({ response, model: modelData, - currentMessagesTokens, - toolLength: toolCalls.length, reasoningEffort: body.reasoning_effort, userKey }); @@ -517,31 +550,35 @@ export const runAgentLoop = async ({ }; // 按 toolCalls 原始顺序写回工具结果，保证后续 LLM 上下文稳定。 - const appendToolRunResults = (toolRunResults: ToolRunResult[]) => { - toolRunResults.forEach( - ({ tool, interactive, stopLoop, toolMessage, toolAssistantMessages }) => { - if (interactive) { - interactiveResponse = { - type: 'toolChildrenInteractive', - params: { - childrenResponse: interactive, - toolParams: { - memoryRequestMessages: [], - toolCallId: tool.id - } + const appendToolRunResults = async (toolRunResults: ToolRunResult[]) => { + for (const { + tool, + interactive, + stopLoop, + toolMessage, + toolAssistantMessages + } of toolRunResults) { + if (interactive) { + interactiveResponse = { + type: 'toolChildrenInteractive', + params: { + childrenResponse: interactive, + toolParams: { + memoryRequestMessages: [], + toolCallId: tool.id } - }; - } - if (stopLoop) { - stopAgentLoop = true; - } - - assistantMessages.push(toolMessage); - requestMessages.push(toolMessage); - // toolAssistantMessages 也需要记录成 AI 响应，所以这里需要推送。 - assistantMessages.push(...toolAssistantMessages); + } + }; } - ); + if (stopLoop) { + stopAgentLoop = true; + } + + assistantMessages.push(toolMessage); + await appendRequestMessages(toolMessage); + // toolAssistantMessages 也需要记录成 AI 响应，所以这里需要推送。 + assistantMessages.push(...toolAssistantMessages); + } }; const safeBatchToolSize = Math.max(1, batchToolSize); @@ -549,17 +586,13 @@ export const runAgentLoop = async ({ while (toolIndex < toolCalls.length) { const currentTool = toolCalls[toolIndex]; - const currentMessagesTokens = await countGptMessagesTokens({ - messages: requestMessages - }); // 只能串行的工具 if (!canBatchTool(currentTool)) { const result = await runTool({ - tool: currentTool, - currentMessagesTokens + tool: currentTool }); - appendToolRunResults([result]); + await appendToolRunResults([result]); toolIndex++; continue; } @@ -579,12 +612,11 @@ export const runAgentLoop = async ({ batchTools, async (tool) => runTool({ - tool, - currentMessagesTokens + tool }), safeBatchToolSize ); - appendToolRunResults(toolRunResults); + await appendToolRunResults(toolRunResults); } } @@ -618,7 +650,7 @@ export const runAgentLoop = async ({ ) { assistantMessages.pop(); } - requestMessages.push(stopResult.feedbackMessage); + await appendRequestMessages(stopResult.feedbackMessage); continue; } } diff --git a/packages/service/core/ai/llm/compress/constants.ts b/packages/service/core/ai/llm/compress/constants.ts index 21f8727c539f..ae0bbb4187e5 100644 --- a/packages/service/core/ai/llm/compress/constants.ts +++ b/packages/service/core/ai/llm/compress/constants.ts @@ -1,86 +1,136 @@ /** - * Agent 上下文压缩配置常量 + * request messages checkpoint 的起始标签。 + * 用于 compress/index.ts 中识别、规范化和生成隐藏历史 checkpoint。 + */ +export const CONTEXT_CHECKPOINT_START_TAG = ''; + +/** + * request messages checkpoint 的结束标签。 + * 用于 compress/index.ts 中识别、规范化和生成隐藏历史 checkpoint。 + */ +export const CONTEXT_CHECKPOINT_END_TAG = ''; + +/** + * token 预算转字符预算时使用的近似换算比例。 + * 用于 compress/index.ts 的本地 head-tail 截断、二分缩短和内容分块。 + */ +export const APPROX_CHARS_PER_TOKEN = 3; + +/** + * LLM 分块压缩合并结果仍超预算时，最多再次压缩合并结果的轮数。 + * 用于 compressLargeContent 的 merge compression 阶段。 + */ +export const MERGED_COMPRESSION_MAX_ROUNDS = 2; + +/** + * 本地 head-tail 截断时，字符预算分配给开头内容的比例。 + * 用于 compress/index.ts 的 truncateContentByHeadTail。 + */ +export const FINAL_HEAD_RATIO = 0.6; + +/** + * request messages 调用 LLM 生成 checkpoint 时的软目标比例。 + * 用于 getCompressRequestMessagesUserPrompt 的 output_budget。 + */ +export const CHECKPOINT_OUTPUT_TARGET_RATIO = 0.2; + +/** + * 压缩相关比例阈值的最小 token 数。 + * 用于避免小上下文模型按比例计算出的压缩目标过小；实际值不会超过模型 maxContext。 + */ +export const COMPRESSION_MIN_TOKEN_LIMIT = 4096; + +/** + * request messages 调用 LLM 生成 checkpoint 时的最小软目标 token 数。 + * 保留旧常量名，避免调用方理解成本；实际值与通用压缩最小值一致。 + */ +export const CHECKPOINT_OUTPUT_MIN_TOKENS = COMPRESSION_MIN_TOKEN_LIMIT; + +/** + * request checkpoint LLM 输出 token 可接受比例。 + * 用于判断 completion token 是否明显超过软目标；超过时记录警告但仍以最终 messages token 校验为准。 + */ +export const REQUEST_CHECKPOINT_COMPLETION_ACCEPT_CONTEXT_RATIO = 0.5; + +/** + * 大内容压缩结果过短时，若压缩结果已占目标预算的该比例以上，则不再追加原文摘录。 + * 用于 appendSourceExcerptForUnderfilledCompression。 + */ +export const SOURCE_ANCHOR_APPEND_SKIP_RATIO = 0.8; + +/** + * 大内容压缩后最多追加的原文结构锚点数量。 + * 用于 appendSourceAnchorsWithinBudget，避免锚点列表挤占摘要主体。 + */ +export const SOURCE_ANCHOR_APPEND_MAX_COUNT = 12; + +/** + * tool response 原文直接返回阈值比例。 + * 用于 compressToolResponse：响应不超过 20% context 时不做任何处理。 + */ +export const TOOL_RESPONSE_DIRECT_RETURN_CONTEXT_RATIO = 0.2; + +/** + * tool response 进入 LLM 压缩的阈值比例。 + * 用于 compressToolResponse：本地轻量处理后仍超过 50% context 才调用 LLM 压缩。 + */ +export const TOOL_RESPONSE_LIGHT_PROCESS_CONTEXT_RATIO = 0.5; + +/** + * 本地截断时插入中间省略内容的标记。 + * 用于 head-tail 截断，明确告知后续模型中间内容被省略。 + */ +export const TRUNCATED_MARKER = + '\n\n... [content truncated: middle omitted to fit token budget] ...\n\n'; + +/** + * Depends on（系统提示词中的步骤历史）压缩触发比例。 + * + * 拼接依赖步骤完整 response 后，超过模型上下文 15% 时触发压缩。 + * 用于 calculateCompressionThresholds().dependsOn。 + */ +export const DEPENDS_ON_THRESHOLD_RATIO = 0.15; + +/** + * 对话历史压缩触发比例。 * - * ## 设计原则 + * messages + tools schema 超过模型上下文 80% 时，尝试压缩成 checkpoint。 + * 用于 compressRequestMessages 的触发阈值。 + */ +export const MESSAGE_THRESHOLD_RATIO = 0.8; + +/** + * 文件读取结果压缩触发比例。 * - * 1. **压缩触发水位** - * - Depends on：超过上下文 15% 后压缩 - * - Agent 对话历史：超过上下文 80% 后压缩 - * - 单个 tool response / 文件读取结果：超过上下文 50% 后压缩 - * - 知识库检索结果：超过上下文 20% 后触发相关性筛选 + * 文件解析工具返回的大型文档内容超过模型上下文 50% 时触发压缩。 + * 用于 calculateCompressionThresholds().fileReadResponse。 + */ +export const FILE_READ_RESPONSE_MAX_RATIO = 0.5; + +/** + * 大文本分块压缩的单块比例。 * - * 2. **压缩策略** - * - 触发阈值：接近空间上限时触发 - * - 压缩目标：保留可续跑上下文，具体摘要粒度交给模型判断 - * - 约束机制：最终结果用真实 token 校验，LLM 输出长度只通过 prompt 软约束 + * 分块阶段单块不超过模型上下文 50%，避免单次压缩请求过大。 + * 用于 compressLargeContent 的 splitIntoChunks 分块预算。 + */ +export const CHUNK_SIZE_RATIO = 0.5; + +/** + * 知识库检索结果相关性筛选触发比例。 * - * 3. **协调关系** - * - Depends on 使用完整 response，先在较小水位触发 - * - Agent 历史包含多轮 user/assistant/tool 消息，接近上下文上限才整体 checkpoint - * - 单个 tool/file response 不能过大，避免挤占后续对话和模型输出空间 - */ - -export const COMPRESSION_CONFIG = { - /** - * === Depends on（系统提示词中的步骤历史）=== - * - * 触发场景：拼接依赖步骤的完整 response 后，token 数超过阈值 - * 内容特点：包含多个步骤的完整执行结果（使用 response 而非 summary） - * - * 示例（maxContext=100k）： - * - 依赖 3 个步骤，每个 4k → 12k (12%) ✅ 不触发 - * - 依赖 5 个步骤，每个 4k → 20k (20%) ⚠️ 触发压缩 - */ - DEPENDS_ON_THRESHOLD: 0.15, // 15% 触发压缩 - - /** - * === 对话历史 === - * - * 触发场景：对话历史（含所有 user/assistant/tool 消息）超过阈值 - * 内容特点：动态累积，触发后会整体压成 checkpoint string - * - * 示例（maxContext=100k）： - * - 初始 20k + 6 轮对话(34k) = 54k (54%) ✅ 不触发 - * - 再 1 轮 = 60k (60%) ⚠️ 触发压缩 → checkpoint - * - checkpoint 保存长程上下文、当前任务和未完成工具上下文 - */ - MESSAGE_THRESHOLD: 0.8, - - /** - * === 单个 tool response === - * - * 触发场景：单个 tool 返回的内容超过绝对大小限制 - * 内容特点：单次 tool 调用的响应（如搜索结果、文件内容等） - * - * 示例（maxContext=100k）： - * - tool response = 8k (8%) ✅ 不触发 - * - tool response = 15k (15%) ⚠️ 触发压缩 - */ - SINGLE_TOOL_MAX: 0.5, - - /** - * === 文件读取结果压缩 === - * - * 触发场景：文件解析工具返回的文件内容超过限制 - * 内容特点：通常是大型文档、PDF 等文件的完整文本内容 - * - * 示例（maxContext=100k）： - * - 文件内容 = 40k (40%) ✅ 不触发 - * - 文件内容 = 60k (60%) ⚠️ 触发压缩 - */ - FILE_READ_RESPONSE_MAX: 0.5, // 50% 触发压缩 - - /** - * === 分块压缩 === - */ - CHUNK_SIZE_RATIO: 0.5, // 单块不超过 maxContext 的 50% - - /** - * === 知识库检索工具的压缩阈值 === - * 策略：使用 LLM 根据查询相关性自动选择最相关的一半分块 - */ - DATASET_SEARCH_SELECTION_RATIO: 0.2 -} as const; + * 检索片段总 token 超过模型上下文 20% 时，调用 LLM 选择最相关片段。 + * 用于 dispatchAgentDatasetSearch 的 chunk selection 触发阈值。 + */ +export const DATASET_SEARCH_SELECTION_RATIO = 0.2; + +/** + * 计算压缩场景中的比例 token 阈值。 + * + * 压缩阈值按比例计算后至少保留 4096 token，避免小上下文模型过早压缩或 output_budget 过小； + * 但阈值不会超过模型 maxContext，避免压缩目标反向大于模型可承载上下文。 + */ +export const getCompressionTokenLimit = (maxContext: number, ratio: number) => + Math.min(maxContext, Math.max(COMPRESSION_MIN_TOKEN_LIMIT, Math.floor(maxContext * ratio))); /** * 计算各场景的压缩阈值 @@ -91,29 +141,27 @@ export const calculateCompressionThresholds = (maxContext: number) => { return { // Depends on 压缩阈值 dependsOn: { - threshold: Math.floor(maxContext * COMPRESSION_CONFIG.DEPENDS_ON_THRESHOLD) + threshold: Math.floor(maxContext * DEPENDS_ON_THRESHOLD_RATIO) }, // 对话历史压缩阈值 messages: { - threshold: Math.floor(maxContext * COMPRESSION_CONFIG.MESSAGE_THRESHOLD) + threshold: getCompressionTokenLimit(maxContext, MESSAGE_THRESHOLD_RATIO) }, - // 单个 tool response 压缩阈值 + // 单个 tool response 兼容阈值；新链路直接使用 0.2/0.5 分层常量。 singleTool: { - threshold: Math.floor(maxContext * COMPRESSION_CONFIG.SINGLE_TOOL_MAX) + threshold: getCompressionTokenLimit(maxContext, TOOL_RESPONSE_LIGHT_PROCESS_CONTEXT_RATIO) }, // 文件读取结果压缩阈值 fileReadResponse: { - threshold: Math.floor(maxContext * COMPRESSION_CONFIG.FILE_READ_RESPONSE_MAX) + threshold: Math.floor(maxContext * FILE_READ_RESPONSE_MAX_RATIO) }, // 分块压缩中每个分块的分割大小，用来划分原始大块的内容。 - chunkSize: Math.floor(maxContext * COMPRESSION_CONFIG.CHUNK_SIZE_RATIO), + chunkSize: Math.floor(maxContext * CHUNK_SIZE_RATIO), // 知识库检索工具阈值，到达阈值会触发选择最相关的一半分块内容（筛选） - datasetSearchSelection: Math.floor( - maxContext * COMPRESSION_CONFIG.DATASET_SEARCH_SELECTION_RATIO - ) + datasetSearchSelection: Math.floor(maxContext * DATASET_SEARCH_SELECTION_RATIO) }; }; diff --git a/packages/service/core/ai/llm/compress/index.ts b/packages/service/core/ai/llm/compress/index.ts index 827069f9026f..c122ce9db231 100644 --- a/packages/service/core/ai/llm/compress/index.ts +++ b/packages/service/core/ai/llm/compress/index.ts @@ -1,10 +1,28 @@ import type { LLMModelItemType } from '@fastgpt/global/core/ai/model.schema'; import { countGptMessagesTokens, countPromptTokens } from '../../../../common/string/tiktoken'; -import { calculateCompressionThresholds } from './constants'; +import { + APPROX_CHARS_PER_TOKEN, + CHECKPOINT_OUTPUT_TARGET_RATIO, + CONTEXT_CHECKPOINT_END_TAG, + CONTEXT_CHECKPOINT_START_TAG, + FINAL_HEAD_RATIO, + MERGED_COMPRESSION_MAX_ROUNDS, + REQUEST_CHECKPOINT_COMPLETION_ACCEPT_CONTEXT_RATIO, + SOURCE_ANCHOR_APPEND_MAX_COUNT, + SOURCE_ANCHOR_APPEND_SKIP_RATIO, + TOOL_RESPONSE_DIRECT_RETURN_CONTEXT_RATIO, + TOOL_RESPONSE_LIGHT_PROCESS_CONTEXT_RATIO, + TRUNCATED_MARKER, + calculateCompressionThresholds, + getCompressionTokenLimit +} from './constants'; import type { CreateLLMResponseProps } from '../request'; import { createLLMResponse } from '../request'; import { ChatCompletionRequestMessageRoleEnum } from '@fastgpt/global/core/ai/constants'; -import type { ChatCompletionMessageParam } from '@fastgpt/global/core/ai/llm/type'; +import type { + ChatCompletionMessageParam, + ChatCompletionTool +} from '@fastgpt/global/core/ai/llm/type'; import { extractExactAnchors, getCompressLargeContentPrompt, @@ -22,17 +40,6 @@ import type { OpenaiAccountType } from '@fastgpt/global/support/user/team/type'; const logger = getLogger(LogCategories.MODULE.AI.LLM_COMPRESS); -// Checkpoint 最终会作为纯字符串写入 chat history；固定标签用于后续 adapter 识别压缩边界。 -const CONTEXT_CHECKPOINT_START_TAG = ''; -const CONTEXT_CHECKPOINT_END_TAG = ''; -const APPROX_CHARS_PER_TOKEN = 3; -const MERGED_COMPRESSION_MAX_ROUNDS = 2; -const FINAL_HEAD_RATIO = 0.6; -const CHECKPOINT_OUTPUT_TARGET_RATIO = 0.15; -const SOURCE_ANCHOR_APPEND_SKIP_RATIO = 0.8; -const SOURCE_ANCHOR_APPEND_MAX_COUNT = 12; -const TRUNCATED_MARKER = '\n\n... [content truncated: middle omitted to fit token budget] ...\n\n'; - // LLM 可能会输出 markdown 代码块，或忘记补外层标签；入库前统一规整为一个可识别的 tagged string。 const normalizeContextCheckpointContent = (content: string) => { const trimmed = content.trim(); @@ -53,195 +60,6 @@ const normalizeContextCheckpointContent = (content: string) => { return `${CONTEXT_CHECKPOINT_START_TAG}\n${withoutFence}\n${CONTEXT_CHECKPOINT_END_TAG}`; }; -/** - * 将 OpenAI message content 统一转成可压缩的纯文本。 - * - * 压缩链路只消费文本；多模态消息里只有 text 部分对上下文摘要有稳定价值，其它结构交给原消息协议处理。 - */ -const getMessageContentText = (content: ChatCompletionMessageParam['content']) => { - if (typeof content === 'string') return content; - if (Array.isArray(content)) { - return content - .map((item) => { - if (typeof item === 'string') return item; - if (item && typeof item === 'object' && 'text' in item && typeof item.text === 'string') { - return item.text; - } - return ''; - }) - .filter(Boolean) - .join('\n'); - } - return ''; -}; - -/** - * 为 LLM checkpoint prompt 构造一段工具调用索引。 - * - * 这里不是最终压缩结果，而是把“用户意图 -> 函数名 -> 参数 -> 工具结果”提前整理出来， - * 避免模型从原始 message JSON 中自行配对 tool_call_id 时漏掉关键参数或结果。 - */ -const buildToolCallMemoryBlock = ({ messages }: { messages: ChatCompletionMessageParam[] }) => { - const toolResultByCallId = new Map(); - for (const message of messages) { - if (message.role !== ChatCompletionRequestMessageRoleEnum.Tool) continue; - - const toolCallId = message.tool_call_id; - if (!toolCallId) continue; - - toolResultByCallId.set(toolCallId, getMessageContentText(message.content)); - } - - const maxResultChars = 900; - const lines: string[] = []; - let latestUserIntent = ''; - - for (const message of messages) { - const content = getMessageContentText(message.content); - if (message.role === ChatCompletionRequestMessageRoleEnum.User && content) { - latestUserIntent = truncateByChars(content.replace(/\s+/g, ' ').trim(), 240); - continue; - } - - const toolCalls = - message.role === ChatCompletionRequestMessageRoleEnum.Assistant - ? message.tool_calls - : undefined; - if (!toolCalls || toolCalls.length === 0) continue; - - for (const toolCall of toolCalls) { - const functionCall = toolCall.function; - if (!functionCall?.name) continue; - const lineParts = [ - `- fn=${functionCall.name}`, - `args=${functionCall.arguments || '{}'}`, - latestUserIntent ? `user=${latestUserIntent}` : '', - toolResultByCallId.has(toolCall.id) - ? `result=${truncateByChars(toolResultByCallId.get(toolCall.id) || '', maxResultChars)}` - : '' - ].filter(Boolean); - - lines.push(lineParts.join('; ')); - } - } - - if (lines.length === 0) return; - - return ` -${lines.join('\n')} -`; -}; - -/** - * 将结构化工具调用历史压成确定性 checkpoint。 - * - * 这类历史的核心信息不是自然语言摘要，而是“用户意图 -> 已选择的工具 -> 参数”。这里只读取通用 - * message/tool_calls 结构，并使用生产压缩阈值作为安全上限，不接收 benchmark expected 这类评测目标。 - */ -const buildStructuredToolCallCheckpoint = ({ - maxCheckpointTokens, - messages -}: { - maxCheckpointTokens: number; - messages: ChatCompletionMessageParam[]; -}) => { - const hasToolCalls = messages.some( - (message) => - message.role === ChatCompletionRequestMessageRoleEnum.Assistant && - message.tool_calls && - message.tool_calls.length > 0 - ); - if (!hasToolCalls) return; - - const toolResultByCallId = new Map(); - for (const message of messages) { - if (message.role !== ChatCompletionRequestMessageRoleEnum.Tool || !message.tool_call_id) { - continue; - } - toolResultByCallId.set(message.tool_call_id, getMessageContentText(message.content)); - } - - const maxChars = Math.max(900, Math.floor(getApproxCharBudget(maxCheckpointTokens) * 0.75)); - const lines: string[] = [ - CONTEXT_CHECKPOINT_START_TAG, - '# Context Checkpoint', - '', - '## Structured Tool Calls' - ]; - let usedChars = lines.join('\n').length; - const pendingUserLines: string[] = []; - let turnIndex = 0; - - const pushLine = (line: string) => { - const normalized = line.replace(/\s+/g, ' ').trim(); - if (!normalized) return false; - const nextChars = usedChars + normalized.length + 1; - if (nextChars > maxChars) return false; - lines.push(normalized); - usedChars = nextChars; - return true; - }; - - const compactUserText = (content: string) => { - const normalized = content.replace(/\s+/g, ' ').trim(); - const toolNames = Array.from(normalized.matchAll(/"name"\s*:\s*"([^"]{1,120})"/g)) - .map((match) => match[1]) - .filter((name): name is string => Boolean(name)); - if (toolNames.length > 0) { - const firstToolDefinitionIndex = normalized.search(/\[\s*\{/); - const prefix = - firstToolDefinitionIndex >= 0 - ? normalized.slice(0, firstToolDefinitionIndex).trim() - : normalized.slice(0, 160).trim(); - return truncateByChars( - [prefix, Array.from(new Set(toolNames)).join(', ')].filter(Boolean).join(' '), - 260 - ); - } - - return truncateByChars(normalized, 220); - }; - - for (const message of messages) { - if (message.role === ChatCompletionRequestMessageRoleEnum.User) { - const content = getMessageContentText(message.content); - if (content) pendingUserLines.push(compactUserText(content)); - continue; - } - - const toolCalls = - message.role === ChatCompletionRequestMessageRoleEnum.Assistant - ? message.tool_calls - : undefined; - if (!toolCalls || toolCalls.length === 0) continue; - - turnIndex += 1; - if (!pushLine(`- turn ${turnIndex}`)) break; - const contextLines = pendingUserLines.splice(0); - contextLines.forEach((line) => pushLine(` source: ${line}`)); - - const assistantText = getMessageContentText(message.content); - if (assistantText) { - pushLine(` assistant: ${truncateByChars(assistantText, 220)}`); - } - - for (const toolCall of toolCalls) { - const functionCall = toolCall.function; - if (!functionCall?.name) continue; - const args = functionCall.arguments || '{}'; - if (!pushLine(` call: ${functionCall.name} args=${args}`)) break; - - const result = toolResultByCallId.get(toolCall.id); - if (result) { - pushLine(` result: ${truncateByChars(result, Math.max(180, maxChars * 0.08))}`); - } - } - } - - lines.push(CONTEXT_CHECKPOINT_END_TAG); - return lines.join('\n'); -}; - // 只用于切分和兜底截断前的粗估；最终是否超限仍以真实 token 统计为准。 const getApproxCharBudget = (tokenLimit: number) => Math.max(0, Math.floor(tokenLimit * APPROX_CHARS_PER_TOKEN)); @@ -457,68 +275,6 @@ const appendSourceAnchorsWithinBudget = async ({ .trim(); }; -/** - * 为非工具调用的长历史构造一个确定性 checkpoint 候选。 - * - * LLM checkpoint 在长会议/文档类历史上可能同义改写或超预算；这个候选只保留原始 history 的 - * head-tail 片段，并用真实 token 计数确认预算，不读取任何评测期望。 - */ -const buildDeterministicHistoryCheckpoint = async ({ - maxCheckpointTokens, - messages -}: { - maxCheckpointTokens: number; - messages: ChatCompletionMessageParam[]; -}) => { - const hasToolCalls = messages.some( - (message) => - message.role === ChatCompletionRequestMessageRoleEnum.Assistant && - message.tool_calls && - message.tool_calls.length > 0 - ); - if (hasToolCalls) return; - - const historyText = messages - .map((message) => { - const content = getMessageContentText(message.content).trim(); - return content ? `${message.role}: ${content}` : ''; - }) - .filter(Boolean) - .join('\n\n'); - if (!historyText) return; - - const wrapCheckpoint = (text: string) => - [ - CONTEXT_CHECKPOINT_START_TAG, - '# Context Checkpoint', - '', - '## Source History Excerpts', - text.trim(), - CONTEXT_CHECKPOINT_END_TAG - ].join('\n'); - - let charBudget = Math.min( - historyText.length, - Math.max(360, Math.floor(getApproxCharBudget(maxCheckpointTokens) * 0.85)) - ); - - for (let round = 0; round < 8 && charBudget > 0; round++) { - const candidate = wrapCheckpoint(truncateByChars(historyText, charBudget)); - const tokens = await countGptMessagesTokens({ - messages: [{ role: 'user', content: candidate }] - }); - - if (tokens <= maxCheckpointTokens) { - return { - checkpoint: candidate, - tokens - }; - } - - charBudget = Math.floor(charBudget * 0.7); - } -}; - /** * 将大型 JSON 工具返回压成通用结构摘要。 * @@ -631,82 +387,145 @@ const summarizeJsonStructure = ({ return true; } - const scalar = stringifyScalar(current); - if (!scalar) return true; - return pushLine(`${path}: ${truncateByChars(scalar, 60)}`); + return true; }; collectImportantScalars(value); visit(value, '', 0); - const scalarValues = importantScalars.map((scalar) => { - const separatorIndex = scalar.indexOf('='); - return separatorIndex >= 0 ? scalar.slice(separatorIndex + 1) : scalar; - }); - return JSON.stringify({ summaryType: 'JSON structural summary', - importantScalarSummary: `important scalar values: ${scalarValues.join('; ')}`, importantScalarValues: importantScalars, structure: structureLines }); }; /** - * 优先用本地确定性方式压缩 JSON 工具返回。 + * 对中等大小工具响应做本地轻量精简。 * - * JSON 的空白、结构 key、数组规模和代表性标量值可以由代码稳定保留；只有这条路径兜不住时， - * 调用方才会退回通用大文本压缩，避免为结构化数据无谓调用 LLM。 + * 这个阶段不调用 LLM，只做结构化 JSON 精简和低语义噪声清理；目标是减少上下文占用，同时避免 + * 对 20%~50% context 的工具结果做高成本压缩或过度摘要。 */ -const tryMinifyToolResponseJson = async ({ - compressedTokenLimit, - response +const lightProcessToolResponse = async ({ + response, + targetTokenLimit }: { - compressedTokenLimit: number; response: string; + targetTokenLimit: number; }) => { - let parsed: unknown; try { - parsed = JSON.parse(response); + const parsed = JSON.parse(response); + if (parsed) { + const compressed = JSON.stringify(parsed); + const compressedTokens = await countPromptTokens(compressed); + if (compressedTokens <= targetTokenLimit) return compressed; + + const structuralSummary = summarizeJsonStructure({ + compressedTokenLimit: targetTokenLimit, + value: parsed + }); + const structuralSummaryTokens = await countPromptTokens(structuralSummary); + if ( + structuralSummaryTokens <= targetTokenLimit && + structuralSummaryTokens < compressedTokens + ) { + return structuralSummary; + } + + return compressed; + } } catch { - return; + // Non-JSON tool responses continue through text cleanup. } - if (!parsed) return; - const compressed = JSON.stringify(parsed); - const tokens = await countPromptTokens(compressed); - if (tokens <= Math.min(200, compressedTokenLimit * 0.2)) return compressed; + return response + .replace(/!\[([^\]]*)\]\([^)]+\)/g, '[$1]') + .replace(/https?:\/\/[^\s"'(),}\]]+/gi, '') + .replace(/\b[a-zA-Z0-9+\/]{100,}={0,2}\b/g, '[BASE64_DATA]') + .replace(/[\/\w\-_]+\/[\w\-_]+\.\w+/g, (match) => { + const parts = match.split('/'); + return parts[parts.length - 1]; + }) + .replace(/\b[a-f0-9]{32,}\b/gi, '') + .replace(/\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b/gi, '') + .replace(/\b\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})?\b/g, '') + .replace(/[ \t]+/g, ' ') + .replace(/\n{3,}/g, '\n\n') + .trim(); +}; - const structuralSummary = summarizeJsonStructure({ - compressedTokenLimit, - value: parsed +/** + * 统计一次真实请求上下文 token。 + * + * request messages 压缩只会改写 messages，但模型请求仍然会携带 tools schema；因此所有触发判断、 + * 压缩后校验和缓存基线都必须把 tools 一起计入。 + */ +const countRequestMessagesTokens = ({ + messages, + tools +}: { + messages: ChatCompletionMessageParam[]; + tools?: ChatCompletionTool[]; +}) => + countGptMessagesTokens({ + messages, + ...(tools?.length ? { tools } : {}) }); - const structuralSummaryTokens = await countPromptTokens(structuralSummary); - if (structuralSummaryTokens <= compressedTokenLimit && structuralSummaryTokens < tokens) { - return structuralSummary; - } - if (tokens <= compressedTokenLimit) return compressed; +const createContextCheckpointMessage = ( + content: ContextCheckpointValueType +): ChatCompletionMessageParam => ({ + role: ChatCompletionRequestMessageRoleEnum.User, + content, + // checkpoint 是给下一轮模型看的历史上下文注入，不作为普通消息展示。 + hideInUI: true +}); + +const getRequestCheckpointOutputTargetTokens = (maxContext: number) => + getCompressionTokenLimit(maxContext, CHECKPOINT_OUTPUT_TARGET_RATIO); + +const getToolResponseCompressionLimits = ({ maxContext }: { maxContext: number }) => { + const directReturnTokenLimit = getCompressionTokenLimit( + maxContext, + TOOL_RESPONSE_DIRECT_RETURN_CONTEXT_RATIO + ); + + return { + // 原始结果不超过 20% context 时不处理；LLM 压缩目标也默认回到这个预算。 + directReturnTokenLimit, + lightProcessTokenLimit: getCompressionTokenLimit( + maxContext, + TOOL_RESPONSE_LIGHT_PROCESS_CONTEXT_RATIO + ), + llmCompressedTokenLimit: directReturnTokenLimit + }; }; /** - * 压缩 对话历史 - * 当 messages 的 token 长度超过阈值时，调用 LLM 进行压缩 + * 压缩对话历史。 + * + * 触发判断基于完整请求上下文（messages + tools schema）；真正压缩时只折叠非 system/developer + * 历史，system/developer 原样保留在最终请求前部。 */ export const compressRequestMessages = async ({ checkIsStopping, + messageTokens: cachedMessageTokens, messages, model, reasoningEffort, + tools, userKey }: { checkIsStopping?: CreateLLMResponseProps['isAborted']; + messageTokens?: number; messages: ChatCompletionMessageParam[]; model: LLMModelItemType; reasoningEffort?: CreateLLMResponseProps['body']['reasoning_effort']; + tools?: ChatCompletionTool[]; userKey?: OpenaiAccountType; }): Promise<{ messages: ChatCompletionMessageParam[]; + messageTokens?: number; usage?: ChatNodeUsageType; requestIds?: string[]; contextCheckpoint?: ContextCheckpointValueType; @@ -733,57 +552,39 @@ export const compressRequestMessages = async ({ } }); + // 触发阈值按完整请求上下文判断；压缩内容仍只包含非 system/developer 历史。 + // system/developer 和 tools schema 虽然不参与 checkpoint 压缩，但会真实占用模型上下文。 + const messageTokens = + cachedMessageTokens ?? + (await countRequestMessagesTokens({ + messages, + tools + })); + if (otherMessages.length === 0) { return { - messages + messages, + messageTokens }; } - // 触发阈值按完整请求上下文判断；压缩内容仍只包含非 system/developer 历史。 - // system/developer 虽然不参与 checkpoint 压缩，但会真实占用模型上下文。 - const messageTokens = await countGptMessagesTokens({ - messages - }); const thresholds = calculateCompressionThresholds(model.maxContext).messages; if (messageTokens <= thresholds.threshold) { return { - messages + messages, + messageTokens }; } - const structuredToolCheckpoint = buildStructuredToolCallCheckpoint({ - messages: otherMessages, - maxCheckpointTokens: thresholds.threshold - }); - if (structuredToolCheckpoint) { - const checkpointMessage: ChatCompletionMessageParam = { - role: ChatCompletionRequestMessageRoleEnum.User, - content: structuredToolCheckpoint, - hideInUI: true - }; - const finalStructuredMessages = [...systemMessages, checkpointMessage]; - const structuredTokens = await countGptMessagesTokens({ - messages: finalStructuredMessages - }); - - if ( - structuredTokens <= thresholds.threshold && - structuredTokens < Math.floor(messageTokens * 0.85) - ) { - return { - messages: finalStructuredMessages, - contextCheckpoint: structuredToolCheckpoint - }; - } - } - - const toolCallMemory = buildToolCallMemoryBlock({ - messages: otherMessages - }); - const checkpointTargetTokenLimit = Math.max( - 512, - Math.floor(model.maxContext * CHECKPOINT_OUTPUT_TARGET_RATIO) + const checkpointTargetTokenLimit = getRequestCheckpointOutputTargetTokens(model.maxContext); + const checkpointCompletionTokenLimit = getCompressionTokenLimit( + model.maxContext, + REQUEST_CHECKPOINT_COMPLETION_ACCEPT_CONTEXT_RATIO + ); + const checkpointWarnTokenLimit = getCompressionTokenLimit( + model.maxContext, + CHECKPOINT_OUTPUT_TARGET_RATIO ); logger.info('Message compression started'); @@ -793,8 +594,7 @@ export const compressRequestMessages = async ({ const compressPrompt = await getCompressRequestMessagesPrompt(); const userPrompt = await getCompressRequestMessagesUserPrompt({ messages: otherMessages, - outputTokenLimit: checkpointTargetTokenLimit, - toolCallMemory + outputTokenLimit: checkpointTargetTokenLimit }); const { answerText, usage, requestId, finish_reason } = await createLLMResponse({ throwError: false, @@ -834,76 +634,65 @@ export const compressRequestMessages = async ({ }; if (!answerText) { - logger.warn('Message compression failed: empty response'); - return { messages, usage: compressedUsage, requestIds: [requestId] }; + logger.warn('Message compression failed', { + reason: 'empty_response', + originalTokens: messageTokens, + threshold: thresholds.threshold, + requestId, + finishReason: finish_reason + }); + return { messages, messageTokens, usage: compressedUsage, requestIds: [requestId] }; } if (finish_reason === 'close') { - logger.info('Compression messages aborted: return original messages'); - return { messages, usage: compressedUsage, requestIds: [requestId] }; + logger.info('Message compression skipped', { + reason: 'request_closed', + originalTokens: messageTokens, + threshold: thresholds.threshold, + requestId + }); + return { messages, messageTokens, usage: compressedUsage, requestIds: [requestId] }; } - let checkpointContent = normalizeContextCheckpointContent(answerText); + const checkpointContent = normalizeContextCheckpointContent(answerText); if (!checkpointContent) { - logger.warn('Message compression failed: invalid checkpoint content'); - return { messages, usage: compressedUsage, requestIds: [requestId] }; + logger.warn('Message compression failed', { + reason: 'invalid_checkpoint_content', + originalTokens: messageTokens, + threshold: thresholds.threshold, + requestId, + answerTextLength: answerText.length + }); + return { messages, messageTokens, usage: compressedUsage, requestIds: [requestId] }; } - const checkpointMessage: ChatCompletionMessageParam = { - role: ChatCompletionRequestMessageRoleEnum.User, - content: checkpointContent, - // checkpoint 是给下一轮模型看的历史上下文注入，不作为普通消息展示。 - hideInUI: true - }; - let finalMessages = [...systemMessages, checkpointMessage]; - let compressedTokens = await countGptMessagesTokens({ - messages: finalMessages + const checkpointMessage = createContextCheckpointMessage(checkpointContent); + const finalMessages = [...systemMessages, checkpointMessage]; + const compressedTokens = await countRequestMessagesTokens({ + messages: finalMessages, + tools }); - // outputTokenLimit 只作为 prompt 软目标；只有最终消息仍超过生产安全阈值时，才退到确定性 head-tail 兜底。 - if (compressedTokens > thresholds.threshold) { - const systemTokens = await countGptMessagesTokens({ - messages: systemMessages + if (usage.outputTokens >= checkpointCompletionTokenLimit) { + logger.warn('Message compression completion exceeds soft limit, keeping LLM checkpoint', { + outputTokens: usage.outputTokens, + completionTokenLimit: checkpointCompletionTokenLimit, + originalTokens: messageTokens, + compressedTokens }); - const availableCheckpointTokens = thresholds.threshold - systemTokens; - - if (availableCheckpointTokens > 0) { - const deterministicCheckpoint = await buildDeterministicHistoryCheckpoint({ - maxCheckpointTokens: availableCheckpointTokens, - messages: otherMessages - }); - - if (deterministicCheckpoint) { - const deterministicMessage: ChatCompletionMessageParam = { - role: ChatCompletionRequestMessageRoleEnum.User, - content: deterministicCheckpoint.checkpoint, - hideInUI: true - }; - const deterministicMessages = [...systemMessages, deterministicMessage]; - const deterministicTokens = await countGptMessagesTokens({ - messages: deterministicMessages - }); - - if ( - deterministicTokens <= thresholds.threshold && - deterministicTokens < compressedTokens - ) { - checkpointContent = deterministicCheckpoint.checkpoint; - finalMessages = deterministicMessages; - compressedTokens = deterministicTokens; - } - } - } + } - if (compressedTokens > thresholds.threshold) { - logger.warn('Message compression failed: compressed checkpoint still exceeds threshold', { - originalTokens: messageTokens, - compressedTokens, - threshold: thresholds.threshold - }); - return { messages, usage: compressedUsage, requestIds: [requestId] }; - } + if (compressedTokens > checkpointWarnTokenLimit) { + logger.warn('Message compression result still exceeds target', { + reason: 'compressed_messages_over_threshold', + originalTokens: messageTokens, + compressedTokens, + compressedTokenLimit: checkpointWarnTokenLimit, + maxContext: model.maxContext, + requestId, + outputTokens: usage.outputTokens + }); } logger.info('Message compression succeeded', { @@ -913,13 +702,19 @@ export const compressRequestMessages = async ({ return { messages: finalMessages, + messageTokens: compressedTokens, usage: compressedUsage, requestIds: [requestId], contextCheckpoint: checkpointContent }; } catch (error) { - logger.error('Message compression failed', { error }); - return { messages }; + logger.error('Message compression failed', { + reason: 'exception', + originalTokens: messageTokens, + threshold: thresholds.threshold, + error + }); + return { messages, messageTokens }; } }; @@ -1040,8 +835,12 @@ export const compressLargeContent = async ({ }; if (!answerText) { - logger.warn('Chunk compression failed: empty response from LLM', { - chunkIndex + logger.warn('Chunk compression failed', { + reason: 'empty_response', + chunkIndex, + chunkLength: chunk.length, + chunkTokenLimit, + requestId }); return { compressed: chunk, @@ -1115,6 +914,7 @@ export const compressLargeContent = async ({ if (finalTokens > compressedTokenLimit) { logger.warn('LLM chunk compression exceeded limit, running merge compression', { + reason: 'chunk_merge_over_budget', finalTokens, compressedTokenLimit, exceedRatio: (finalTokens / compressedTokenLimit).toFixed(2) @@ -1149,6 +949,9 @@ export const compressLargeContent = async ({ if (needsDeterministicTruncate || finalTokens > compressedTokenLimit) { logger.warn('LLM merge compression still exceeded limit, applying head-tail truncate', { + reason: needsDeterministicTruncate + ? 'merge_compression_low_gain' + : 'merge_compression_over_budget', finalTokens, compressedTokenLimit, exceedRatio: (finalTokens / compressedTokenLimit).toFixed(2) @@ -1262,7 +1065,12 @@ export const compressLargeContent = async ({ requestIds: result.usage.requestIds }; } catch (error) { - logger.error('Chunk compression failed, fallback to binary truncate', { error }); + logger.error('Chunk compression failed, returning cleaned content without LLM compression', { + reason: 'exception', + error, + compressedTokenLimit: effectiveCompressedTokenLimit, + cleanedContentTokens: currentTokens + }); return { compressed: content.trim() }; @@ -1272,17 +1080,11 @@ export const compressLargeContent = async ({ export const compressToolResponse = async ({ response, model, - compressedTokenLimit: customCompressedTokenLimit, - currentMessagesTokens = 0, - toolLength = 1, reasoningEffort, userKey }: { response: string; model: LLMModelItemType; - compressedTokenLimit?: number; - currentMessagesTokens?: number; - toolLength?: number; reasoningEffort?: CreateLLMResponseProps['body']['reasoning_effort']; userKey?: OpenaiAccountType; }): Promise<{ @@ -1296,40 +1098,54 @@ export const compressToolResponse = async ({ }; } - // 单个 tool response 既受固定结果上限限制，也受当前请求剩余上下文窗口限制。 - const staticCompressedTokenLimit = calculateCompressionThresholds(model.maxContext).singleTool - .threshold; - - // 计算每个 tool response 的动态结果预算，预防多个 tool 同时返回的数据打爆上下文。 - const availableCompressedTokenLimit = Math.max( - 0, - Math.floor((model.maxContext - currentMessagesTokens) / toolLength) - ); + const responseTokens = await countPromptTokens(response); + const { directReturnTokenLimit, lightProcessTokenLimit, llmCompressedTokenLimit } = + getToolResponseCompressionLimits({ + maxContext: model.maxContext + }); - // 取静态结果上限、动态结果预算和调用方显式目标预算的较小值。 - const compressedTokenLimit = Math.min( - staticCompressedTokenLimit, - availableCompressedTokenLimit, - customCompressedTokenLimit ?? Number.POSITIVE_INFINITY - ); + // 不超过 20% context 的工具结果直接发给模型，保持原始结构和可读性。 + if (responseTokens <= directReturnTokenLimit) { + return { + compressed: response + }; + } - const jsonCompressed = await tryMinifyToolResponseJson({ + const lightProcessedResponse = await lightProcessToolResponse({ response, - compressedTokenLimit + targetTokenLimit: directReturnTokenLimit }); - if (jsonCompressed) { + const lightProcessedTokens = await countPromptTokens(lightProcessedResponse); + + // 大于 20% context 的工具结果先做本地精简；精简后不超过 50% context 则不进入 LLM 压缩。 + if (lightProcessedTokens <= lightProcessTokenLimit) { return { - compressed: jsonCompressed + compressed: lightProcessedResponse }; } // 调用通用压缩函数 - return compressLargeContent({ - content: response, + const result = await compressLargeContent({ + content: lightProcessedResponse, model, - compressedTokenLimit, + compressedTokenLimit: llmCompressedTokenLimit, moduleName: i18nT('account_usage:tool_response_compress'), reasoningEffort, userKey }); + + const compressedTokens = await countPromptTokens(result.compressed); + if (compressedTokens > llmCompressedTokenLimit) { + logger.warn('Tool response compression result still exceeds target', { + reason: 'compressed_tool_response_over_target', + originalTokens: responseTokens, + lightProcessedTokens, + compressedTokens, + compressedTokenLimit: llmCompressedTokenLimit, + maxContext: model.maxContext, + requestIds: result.requestIds + }); + } + + return result; }; diff --git a/packages/service/core/ai/llm/compress/prompt.ts b/packages/service/core/ai/llm/compress/prompt.ts index f369a42e6c21..10a28f860b18 100644 --- a/packages/service/core/ai/llm/compress/prompt.ts +++ b/packages/service/core/ai/llm/compress/prompt.ts @@ -70,20 +70,96 @@ ${anchors.map((anchor) => `- ${anchor}`).join('\n')} /** * 将历史消息格式化成 checkpoint 压缩模型可读的 JSON。 * - * 这里只暴露压缩需要的协议字段，避免把无关运行时字段塞进 prompt 增加 token 和噪声。 + * 这里只保留可继续工作的 user/assistant 文本历史。assistant 发起的 tool 调用会和对应 tool + * 结果合并成一条 assistant content，避免压缩模型看到旧工具协议字段，同时保留函数名、参数和结果绑定。 */ -const formatMessagesForCheckpoint = (messages: ChatCompletionMessageParam[]) => - JSON.stringify( - messages.map((message) => ({ - role: message.role, - content: message.content, - reasoning_content: message.reasoning_content, - tool_calls: message.role === 'assistant' ? message.tool_calls : undefined, - tool_call_id: message.role === 'tool' ? message.tool_call_id : undefined - })), - null, - 2 - ); +const getMessageContentText = (content: ChatCompletionMessageParam['content']) => { + if (typeof content === 'string') return content; + if (Array.isArray(content)) { + return content + .map((item) => { + if (typeof item === 'string') return item; + if (item && typeof item === 'object' && 'text' in item && typeof item.text === 'string') { + return item.text; + } + return ''; + }) + .filter(Boolean) + .join('\n'); + } + return ''; +}; + +const escapeXmlText = (text: string) => + text.replace(/&/g, '&').replace(//g, '>'); + +const escapeXmlAttribute = (text: string) => + escapeXmlText(text).replace(/"/g, '"').replace(/'/g, '''); + +const formatMessagesForCheckpoint = (messages: ChatCompletionMessageParam[]) => { + const toolResultByCallId = new Map(); + for (const message of messages) { + if (message.role !== 'tool' || !message.tool_call_id) continue; + + const content = getMessageContentText(message.content).trim(); + if (content) { + toolResultByCallId.set(message.tool_call_id, content); + } + } + + const consumedToolCallIds = new Set(); + + const formattedMessages = messages.flatMap((message) => { + const content = getMessageContentText(message.content).trim(); + + if (message.role === 'user') { + if (!content) return []; + return [{ role: 'user', content }]; + } + + if (message.role === 'assistant') { + const toolCalls = message.tool_calls ?? []; + const toolCallBlocks = toolCalls.flatMap((toolCall) => { + const functionCall = toolCall.function; + if (!functionCall?.name) return []; + + return [ + [ + ``, + `${escapeXmlText(functionCall.arguments || '{}')}`, + toolResultByCallId.has(toolCall.id) + ? `${escapeXmlText(toolResultByCallId.get(toolCall.id) || '')}` + : '', + '' + ].join('\n') + ]; + }); + toolCalls.forEach((toolCall) => { + if (toolResultByCallId.has(toolCall.id)) { + consumedToolCallIds.add(toolCall.id); + } + }); + + const toolsContent = + toolCallBlocks.length > 0 ? ['', ...toolCallBlocks, ''].join('\n') : ''; + const mergedContent = [content, toolsContent].filter(Boolean).join('\n\n').trim(); + if (!mergedContent) return []; + + return [{ role: 'assistant', content: mergedContent }]; + } + + if (message.role === 'tool') { + // 已经合并到对应 assistant tool call 中；孤立 tool 结果才作为 assistant 上下文保留。 + if (message.tool_call_id && consumedToolCallIds.has(message.tool_call_id)) return []; + if (!content) return []; + return [{ role: 'assistant', content }]; + } + + return []; + }); + + return JSON.stringify(formattedMessages, null, 2); +}; export const getCompressRequestMessagesPrompt = async () => { return `你是 Agent 历史上下文 checkpoint 压缩专家。你的任务是把用户提供的对话历史压缩成一段可继续工作的高保真上下文摘要 string。 @@ -120,7 +196,8 @@ export const getCompressRequestMessagesPrompt = async () => { ## 结构锚点规则 -- 对 structural_anchor_candidates 和 tool_call_memory 里的字段名、函数名、参数名、ID、路径、URL、错误码、数字、日期，优先原样保留。 +- 对 histories 和 structural_anchor_candidates 里的字段名、函数名、参数名、ID、路径、URL、错误码、数字、日期，优先原样保留。 +- histories 中 assistant content 里的 表示已经执行过的工具调用； 是工具名， 是调用参数， 是对应结果。压缩时必须保留仍影响后续任务的工具名、参数和关键结果。 - 不要为了保留锚点而堆无关关键词；锚点必须服务于后续执行、定位资源、复现工具结果或理解约束。 - 不要把具体工具名、机构名、文件名、接口名压成“相关工具/某机构/该文件/接口”。 @@ -165,12 +242,10 @@ export const getCompressRequestMessagesPrompt = async () => { export const getCompressRequestMessagesUserPrompt = async ({ messages, - outputTokenLimit, - toolCallMemory + outputTokenLimit }: { messages: ChatCompletionMessageParam[]; outputTokenLimit?: number; - toolCallMemory?: string; }) => { const histories = formatMessagesForCheckpoint(messages); @@ -178,7 +253,7 @@ export const getCompressRequestMessagesUserPrompt = async ({ ${histories} -${outputTokenLimit ? `\nTarget maximum output tokens: ${outputTokenLimit}. Use compact bullets and omit nonessential prose.\n\n\n` : ''}${toolCallMemory ? `${toolCallMemory}\n\n` : ''}${renderExactAnchors(histories, 100)} +${outputTokenLimit ? `\nTarget maximum output tokens: ${outputTokenLimit}. Use compact bullets and omit nonessential prose.\n\n\n` : ''}${renderExactAnchors(histories, 100)} 请执行历史上下文 checkpoint 压缩，只输出 ...。 `; }; diff --git a/packages/service/core/ai/sandbox/instance/repository.ts b/packages/service/core/ai/sandbox/instance/repository.ts index e553d41bc9ab..4263c836cc58 100644 --- a/packages/service/core/ai/sandbox/instance/repository.ts +++ b/packages/service/core/ai/sandbox/instance/repository.ts @@ -150,7 +150,7 @@ export async function markSandboxResourceStopped(resource: SandboxResourceRef) { } /** - * 查询可被 5 分钟 cron 暂停的运行中 sandbox。 + * 查询可被 sandbox stop cron 暂停的运行中 sandbox。 * * 正在归档或已归档的实例由归档流程管理生命周期，stop cron 不能插入中间状态。 */ @@ -249,7 +249,7 @@ export async function findSandboxInstanceArchiveState(params: { /** * 流式读取待归档 sandbox，避免迁移脚本或 cron 一次性把历史全量记录拉进 Node.js 内存。 * - * 一周后的实例应先由 5 分钟 cron 标记为 stopped，再由归档任务统一处理。 + * 一周后的实例应先由 sandbox stop cron 标记为 stopped，再由归档任务统一处理。 * 归档不直接抢 running 实例，避免和 stop cron 操作同一个远端资源。 */ export function createSandboxResourcesToArchiveCursor(params: { diff --git a/packages/service/core/ai/sandbox/service/cron.ts b/packages/service/core/ai/sandbox/service/cron.ts index f412b88fd72f..ab2a485e5fb8 100644 --- a/packages/service/core/ai/sandbox/service/cron.ts +++ b/packages/service/core/ai/sandbox/service/cron.ts @@ -23,7 +23,9 @@ export const cronJob = async () => { ); if (!instances.length) return; - logger.info('Found running sandboxes inactive > 5 min', { count: instances.length }); + logger.info(`Found running sandboxes inactive > ${SANDBOX_SUSPEND_MINUTES} min`, { + count: instances.length + }); await stopSandboxResources(instances); }); diff --git a/packages/service/core/app/mcp.ts b/packages/service/core/app/mcp.ts index 1d18e91d410c..f896e3770417 100644 --- a/packages/service/core/app/mcp.ts +++ b/packages/service/core/app/mcp.ts @@ -17,12 +17,184 @@ import { isInternalAddress, PRIVATE_URL_TEXT } from '../../common/system/utils'; const logger = getLogger(LogCategories.MODULE.APP.MCP_TOOLS); +const MCP_SAFE_FETCH_MAX_REDIRECTS = 5; +const MCP_REDIRECT_STATUS_CODES = new Set([301, 302, 303, 307, 308]); +const MCP_SENSITIVE_REDIRECT_HEADERS = new Set(['authorization', 'cookie', 'proxy-authorization']); + +type McpFetch = (url: string | URL, init?: RequestInit) => Promise; + export const assertMCPUrlNotInternal = async (url: string) => { if (await isInternalAddress(url)) { return Promise.reject(PRIVATE_URL_TEXT); } }; +const headersInitToRecord = (headers?: HeadersInit): Record => { + const record: Record = {}; + + if (!headers) return record; + + if (headers instanceof Headers) { + headers.forEach((value, key) => { + record[key] = value; + }); + return record; + } + + if (Array.isArray(headers)) { + headers.forEach(([key, value]) => { + record[key] = value; + }); + return record; + } + + Object.entries(headers).forEach(([key, value]) => { + record[key] = String(value); + }); + return record; +}; + +const isMcpRedirectResponse = (response: Response) => { + return MCP_REDIRECT_STATUS_CODES.has(response.status) && !!response.headers.get('location'); +}; + +const resolveMcpRedirectUrl = (location: string, currentUrl: string) => { + const redirectUrl = new URL(location, currentUrl); + + if (redirectUrl.protocol !== 'http:' && redirectUrl.protocol !== 'https:') { + throw new Error('MCP redirect target only supports http/https protocol'); + } + + return redirectUrl.toString(); +}; + +const getMcpRedirectHeaders = ({ + headers, + currentUrl, + redirectUrl, + shouldSwitchToGet +}: { + headers?: HeadersInit; + currentUrl: string; + redirectUrl: string; + shouldSwitchToGet: boolean; +}) => { + const current = new URL(currentUrl); + const redirect = new URL(redirectUrl); + const shouldDropSensitiveHeaders = + current.protocol !== redirect.protocol || current.host !== redirect.host; + + return Object.entries(headersInitToRecord(headers)).reduce>( + (acc, [key, value]) => { + const lowerKey = key.toLowerCase(); + + // 301/302 POST 与 303 会转成 GET，继续携带 content-* 容易让目标端误判请求体。 + if (shouldSwitchToGet && lowerKey.startsWith('content-')) { + return acc; + } + + // MCP header 中常带有鉴权密钥，跨 host/protocol 重定向时不能泄露给新目标。 + if (shouldDropSensitiveHeaders && MCP_SENSITIVE_REDIRECT_HEADERS.has(lowerKey)) { + return acc; + } + + if (lowerKey === 'host') { + return acc; + } + + acc[key] = value; + return acc; + }, + {} + ); +}; + +const getMcpRedirectRequestInit = ({ + init, + response, + currentUrl, + redirectUrl +}: { + init?: RequestInit; + response: Response; + currentUrl: string; + redirectUrl: string; +}): RequestInit => { + const method = (init?.method || 'GET').toUpperCase(); + const shouldSwitchToGet = + ((response.status === 301 || response.status === 302) && method === 'POST') || + (response.status === 303 && method !== 'GET' && method !== 'HEAD'); + + return { + ...init, + // Node fetch 默认会自动跟随重定向；这里必须保持 manual，才能逐跳做 SSRF 校验。 + redirect: 'manual', + method: shouldSwitchToGet ? 'GET' : init?.method, + body: shouldSwitchToGet ? undefined : init?.body, + headers: getMcpRedirectHeaders({ + headers: init?.headers, + currentUrl, + redirectUrl, + shouldSwitchToGet + }) + }; +}; + +/** + * 为 MCP SDK transport 注入安全 fetch。 + * + * MCP 连接本身会先校验初始 URL，但 SDK 内部默认使用 fetch 自动跟随重定向。 + * 这会让“初始 URL 合法，Location 跳到内网地址”的场景绕过 SSRF 防护。 + * 该 fetch 通过 `redirect: manual` 接管重定向流程，并对每一跳目标重新执行 + * 内网地址校验；跨 host/protocol 跳转时还会移除鉴权类 header，避免 MCP 密钥泄露。 + */ +export const createMcpSafeFetch = ({ + maxRedirects = MCP_SAFE_FETCH_MAX_REDIRECTS, + fetchImpl = fetch as McpFetch +}: { + maxRedirects?: number; + fetchImpl?: McpFetch; +} = {}): McpFetch => { + const redirectLimit = Math.max(0, maxRedirects); + + return async (url, init) => { + let currentUrl = new URL(url.toString()).toString(); + let currentInit: RequestInit = { + ...init, + redirect: 'manual' + }; + + for (let redirectCount = 0; redirectCount <= redirectLimit; redirectCount++) { + await assertMCPUrlNotInternal(currentUrl); + + const response = await fetchImpl(currentUrl, currentInit); + + if (!isMcpRedirectResponse(response)) { + return response; + } + + if (redirectCount === redirectLimit) { + throw new Error(`Maximum MCP redirects exceeded: ${redirectLimit}`); + } + + const redirectUrl = resolveMcpRedirectUrl(response.headers.get('location')!, currentUrl); + await assertMCPUrlNotInternal(redirectUrl); + + currentInit = getMcpRedirectRequestInit({ + init: currentInit, + response, + currentUrl, + redirectUrl + }); + currentUrl = redirectUrl; + + await response.body?.cancel().catch(() => undefined); + } + + throw new Error(`Maximum MCP redirects exceeded: ${redirectLimit}`); + }; +}; + const shouldFallbackToSSE = (error: unknown): boolean => { return ( error instanceof StreamableHTTPError && @@ -71,6 +243,7 @@ export class MCPClient { private async doConnect(): Promise { await assertMCPUrlNotInternal(this.url); + const safeFetch = createMcpSafeFetch(); // 避免连接重复，强制关闭一次 await this.client.close().catch(() => {}); @@ -78,6 +251,7 @@ export class MCPClient { logger.debug('Start connect mcp client', { url: this.url }); try { const transport = new StreamableHTTPClientTransport(new URL(this.url), { + fetch: safeFetch, requestInit: { headers: this.headers } @@ -92,26 +266,19 @@ export class MCPClient { try { await this.client.connect( new SSEClientTransport(new URL(this.url), { + fetch: safeFetch, requestInit: { headers: this.headers }, eventSourceInit: { fetch: (url, init) => { - const mergedHeaders: Record = { + const mergedHeaders = { ...this.headers }; - if (init?.headers) { - if (init.headers instanceof Headers) { - init.headers.forEach((value, key) => { - mergedHeaders[key] = value; - }); - } else if (typeof init.headers === 'object') { - Object.assign(mergedHeaders, init.headers); - } - } + Object.assign(mergedHeaders, headersInitToRecord(init?.headers)); - return fetch(url, { + return safeFetch(url, { ...init, headers: mergedHeaders }); diff --git a/packages/service/core/workflow/dispatch/ai/agent/sub/tool/index.ts b/packages/service/core/workflow/dispatch/ai/agent/sub/tool/index.ts index 8a3b7e20cf2c..66ebde1a087e 100644 --- a/packages/service/core/workflow/dispatch/ai/agent/sub/tool/index.ts +++ b/packages/service/core/workflow/dispatch/ai/agent/sub/tool/index.ts @@ -23,6 +23,8 @@ import { pluginClient } from '../../../../../../../thirdProvider/fastgptPlugin'; import { SystemToolRepo } from '../../../../../../app/tool/systemTool/systemTool.repo'; import { InvokeProcessor } from '../../../../../../../support/invoke/invoke'; import { getLogger, LogCategories } from '../../../../../../../common/logger'; +import { authAppByTmbId } from '../../../../../../../support/permission/app/auth'; +import { ReadPermissionVal } from '@fastgpt/global/support/permission/constant'; type SystemInputConfigType = { type: SystemToolSecretInputTypeEnum; @@ -85,6 +87,18 @@ export const dispatchTool = async ({ const logger = getLogger(LogCategories.MODULE.APP.TOOL); try { + /** + * Agent 工具调用也会按持久化 toolId 解析 HTTP/MCP 父工具集。 + * 这里必须使用当前运行工作流的 tmbId 做运行时授权，防止绕过保存阶段的脏引用被模型调用执行。 + */ + const authRuntimeToolset = async (parentId: string) => { + await authAppByTmbId({ + tmbId: runningAppInfo.tmbId, + appId: parentId, + per: ReadPermissionVal + }); + }; + if (toolConfig?.systemTool?.toolId) { const systemToolRepo = SystemToolRepo.getInstance(); const tool = await systemToolRepo.getSystemToolRuntime({ @@ -195,6 +209,7 @@ export const dispatchTool = async ({ if (!parentId || !toolName) { return Promise.reject(`Invalid MCP tool id: ${toolConfig.mcpTool.toolId}`); } + await authRuntimeToolset(parentId); const tool = await getAppVersionById({ appId: parentId, @@ -228,6 +243,7 @@ export const dispatchTool = async ({ if (!parentId || !toolName) { return Promise.reject(`Invalid HTTP tool id: ${toolConfig.httpTool.toolId}`); } + await authRuntimeToolset(parentId); const toolset = await getAppVersionById({ appId: parentId, diff --git a/packages/service/core/workflow/dispatch/ai/toolcall/toolCall.ts b/packages/service/core/workflow/dispatch/ai/toolcall/toolCall.ts index 89cc03d93499..21a0796b3cf2 100644 --- a/packages/service/core/workflow/dispatch/ai/toolcall/toolCall.ts +++ b/packages/service/core/workflow/dispatch/ai/toolcall/toolCall.ts @@ -156,6 +156,8 @@ export const runToolCall = async (props: DispatchToolModuleProps): Promise false, onAfterCompressContext({ usage, requestIds, seconds }) { + if (!usage) return; + appendContextCompressNodeResponse({ usage, requestIds, diff --git a/packages/service/core/workflow/dispatch/child/runTool.ts b/packages/service/core/workflow/dispatch/child/runTool.ts index 917b6eb735a4..97f93b87b983 100644 --- a/packages/service/core/workflow/dispatch/child/runTool.ts +++ b/packages/service/core/workflow/dispatch/child/runTool.ts @@ -24,6 +24,8 @@ import { pluginClient } from '../../../../thirdProvider/fastgptPlugin'; import { SystemToolRepo } from '../../../app/tool/systemTool/systemTool.repo'; import { InvokeProcessor } from '../../../../support/invoke/invoke'; import { getLogger, LogCategories } from '../../../../common/logger'; +import { authAppByTmbId } from '../../../../support/permission/app/auth'; +import { ReadPermissionVal } from '@fastgpt/global/support/permission/constant'; type SystemInputConfigType = { type: SystemToolSecretInputTypeEnum; @@ -66,6 +68,18 @@ export const dispatchRunTool = async (props: RunToolProps): Promise = {}; try { + /** + * HTTP/MCP 子工具的 toolId 可由工作流 JSON 持久化，运行时必须用当前工作流执行身份 + * 重新校验父工具集权限，避免脏数据或绕过保存接口的跨用户工具集引用被执行。 + */ + const authRuntimeToolset = async (parentId: string) => { + await authAppByTmbId({ + tmbId: runningAppInfo.tmbId, + appId: parentId, + per: ReadPermissionVal + }); + }; + // run system tool if (toolConfig?.systemTool?.toolId) { const systemToolRepo = SystemToolRepo.getInstance(); @@ -214,6 +228,7 @@ export const dispatchRunTool = async (props: RunToolProps): Promise { + const mutableServiceEnv = serviceEnv as { CHECK_INTERNAL_IP: boolean }; + const originalCheckInternalIp = serviceEnv.CHECK_INTERNAL_IP; + + afterEach(() => { + mutableServiceEnv.CHECK_INTERNAL_IP = originalCheckInternalIp; + }); + describe('createProxyAxios', () => { beforeEach(() => { vi.clearAllMocks(); @@ -154,7 +164,7 @@ describe('axios.ts', () => { describe('axios 导出实例', () => { it('应该导出一个默认的 axios 实例', async () => { - const { axios } = await import('@fastgpt/service/common/api/axios'); + const { axios } = await import('../../../common/api/axios'); expect(axios).toBeDefined(); expect(axios.defaults).toBeDefined(); @@ -164,6 +174,138 @@ describe('axios.ts', () => { }); }); + describe('safe axios redirect protection', () => { + const listen = (handler: http.RequestListener, host = '127.0.0.1') => + new Promise((resolve) => { + const server = http.createServer(handler); + server.listen(0, host, () => resolve(server)); + }); + + const closeServer = (server: http.Server) => + new Promise((resolve, reject) => { + server.close((err) => (err ? reject(err) : resolve())); + }); + + const getServerPort = (server: http.Server): number => { + const address = server.address(); + if (!address || typeof address === 'string') { + throw new Error('Invalid test server address'); + } + return address.port; + }; + + /** + * 构造一个非 loopback 的本机访问地址,用于模拟“初始 URL 通过 SSRF 校验”。 + * CHECK_INTERNAL_IP=false 时私网地址会放行,但 loopback/metadata 仍然恒拦截。 + */ + const getReachablePrivateHost = () => { + const interfaces = os.networkInterfaces(); + for (const items of Object.values(interfaces)) { + for (const item of items || []) { + if (item.family === 'IPv4' && !item.internal) { + return item.address; + } + } + } + return undefined; + }; + + it('应该逐跳跟随公网重定向', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + const targetServer = await listen((req, res) => { + res.end(JSON.stringify({ ok: true, url: req.url })); + }, '0.0.0.0'); + const targetPort = getServerPort(targetServer); + + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://${carrierHost}:${targetPort}/target`); + res.end('redirect'); + }, '0.0.0.0'); + const redirectPort = getServerPort(redirectServer); + + try { + const instance = createProxyAxios(); + const response = await instance.get<{ ok: boolean; url: string }>( + `http://${carrierHost}:${redirectPort}/fetch`, + { + httpAgent: new http.Agent() + } + ); + + expect(response.data).toEqual({ ok: true, url: '/target' }); + } finally { + await closeServer(redirectServer); + await closeServer(targetServer); + } + }); + + it('应该阻止重定向到 loopback 地址', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + const protectedServer = await listen((req, res) => { + res.end('INTERNAL-ONLY-RESPONSE'); + }); + const protectedPort = getServerPort(protectedServer); + + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://127.0.0.1:${protectedPort}/latest/meta-data/`); + res.end('redirect'); + }, '0.0.0.0'); + const redirectPort = getServerPort(redirectServer); + + try { + const instance = createProxyAxios(); + await expect( + instance.get(`http://${carrierHost}:${redirectPort}/fetch`, { + httpAgent: new http.Agent() + }) + ).rejects.toThrow(PRIVATE_URL_TEXT); + } finally { + await closeServer(redirectServer); + await closeServer(protectedServer); + } + }); + + it('应该限制最大重定向次数', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + let redirectPort = 0; + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://${carrierHost}:${redirectPort}/loop`); + res.end('redirect'); + }, '0.0.0.0'); + redirectPort = getServerPort(redirectServer); + + try { + const instance = createProxyAxios(); + await expect( + instance.get(`http://${carrierHost}:${redirectPort}/loop`, { + httpAgent: new http.Agent(), + maxRedirects: 1 + }) + ).rejects.toThrow('Maximum redirects exceeded'); + } finally { + await closeServer(redirectServer); + } + }); + }); + describe('pickOutboundAxios', () => { it.each([ 'http://example.com', @@ -171,14 +313,14 @@ describe('axios.ts', () => { 'http://169.254.169.254/latest/meta-data/', '//attacker.example/probe' // protocol-relative 也按绝对处理 ])('绝对 URL %j 返回 safe axios 实例', async (url) => { - const { axios, pickOutboundAxios } = await import('@fastgpt/service/common/api/axios'); + const { axios, pickOutboundAxios } = await import('../../../common/api/axios'); expect(pickOutboundAxios(url)).toBe(axios); }); it.each(['/api/foo', 'api/foo', '/support/outLink/feishu/abc'])( '相对路径 %j 返回内部 axios(baseURL 固定到本机)', async (url) => { - const { axios, pickOutboundAxios } = await import('@fastgpt/service/common/api/axios'); + const { axios, pickOutboundAxios } = await import('../../../common/api/axios'); const client = pickOutboundAxios(url); expect(client).not.toBe(axios); expect(client.defaults.baseURL).toMatch(/^http:\/\//); @@ -186,7 +328,7 @@ describe('axios.ts', () => { ); it('多次调用同一类型的 URL,内部 client 应被复用(避免每次新建实例)', async () => { - const { pickOutboundAxios } = await import('@fastgpt/service/common/api/axios'); + const { pickOutboundAxios } = await import('../../../common/api/axios'); const a = pickOutboundAxios('/api/a'); const b = pickOutboundAxios('/api/b'); expect(a).toBe(b); diff --git a/packages/service/test/core/ai/llm/agentLoop/baseLoop.test.ts b/packages/service/test/core/ai/llm/agentLoop/baseLoop.test.ts index 5b76afc45af7..6d6e6d2ad920 100644 --- a/packages/service/test/core/ai/llm/agentLoop/baseLoop.test.ts +++ b/packages/service/test/core/ai/llm/agentLoop/baseLoop.test.ts @@ -7,13 +7,17 @@ import type { ChatCompletionTool } from '@fastgpt/global/core/ai/llm/type'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { mockCreateLLMResponseQueue, text, toolCall } from './_mocks/llmQueue'; -const { createLLMResponseMock, compressRequestMessagesMock, compressToolResponseMock } = vi.hoisted( - () => ({ - createLLMResponseMock: vi.fn(), - compressRequestMessagesMock: vi.fn(), - compressToolResponseMock: vi.fn() - }) -); +const { + createLLMResponseMock, + compressRequestMessagesMock, + compressToolResponseMock, + countGptMessagesTokensMock +} = vi.hoisted(() => ({ + createLLMResponseMock: vi.fn(), + compressRequestMessagesMock: vi.fn(), + compressToolResponseMock: vi.fn(), + countGptMessagesTokensMock: vi.fn(async () => 100) +})); vi.mock('@fastgpt/service/core/ai/llm/request', () => ({ createLLMResponse: createLLMResponseMock @@ -46,7 +50,7 @@ vi.mock('@fastgpt/service/core/ai/llm/utils', () => ({ })); vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({ - countGptMessagesTokens: vi.fn(async () => 100) + countGptMessagesTokens: countGptMessagesTokensMock })); vi.mock('@fastgpt/service/support/wallet/usage/utils', () => ({ @@ -77,6 +81,7 @@ const searchTool: ChatCompletionTool = { describe('runAgentLoop with mocked createLLMResponse', () => { beforeEach(() => { vi.clearAllMocks(); + countGptMessagesTokensMock.mockResolvedValue(100); compressRequestMessagesMock.mockImplementation(async ({ messages }) => ({ messages })); @@ -195,6 +200,125 @@ describe('runAgentLoop with mocked createLLMResponse', () => { expect(usagePush).toHaveBeenCalledWith([compressedUsage]); }); + it('applies local context checkpoint compression even without usage', async () => { + const contextCheckpoint = 'local structured history'; + const usagePush = vi.fn(); + const onAfterCompressContext = vi.fn(); + + compressRequestMessagesMock.mockImplementation(async () => ({ + messages: [ + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: contextCheckpoint, + hideInUI: true + } + ], + requestIds: [], + contextCheckpoint + })); + mockCreateLLMResponseQueue(createLLMResponseMock, [ + text({ requestId: 'req_direct', content: 'direct answer' }) + ]); + + const result = await runAgentLoop({ + maxRunAgentTimes: 5, + body: { + model: 'gpt-4', + stream: true, + messages: [ + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: 'hello' + } + ], + tools: [] + }, + usagePush, + isAborted: () => false, + onRunTool: vi.fn(), + onRunInteractiveTool: vi.fn(), + onAfterCompressContext + }); + + expect(result.contextCheckpoint).toEqual(contextCheckpoint); + expect(createLLMResponseMock.mock.calls[0][0].body.messages[0]).toEqual({ + role: ChatCompletionRequestMessageRoleEnum.User, + content: contextCheckpoint, + hideInUI: true + }); + expect(onAfterCompressContext).toHaveBeenCalledWith( + expect.objectContaining({ + usage: undefined, + requestIds: [], + contextCheckpoint + }) + ); + expect(usagePush).not.toHaveBeenCalledWith([undefined]); + }); + + it('reuses request message tokens and counts only appended messages between turns', async () => { + countGptMessagesTokensMock + .mockResolvedValueOnce(11) + .mockResolvedValueOnce(17) + .mockResolvedValue(100); + compressRequestMessagesMock.mockImplementation(async ({ messages, messageTokens }) => ({ + messages, + messageTokens: messageTokens ?? 1000 + })); + const onRunTool = vi.fn(async () => ({ + response: 'search result', + assistantMessages: [], + usages: [] + })); + + mockCreateLLMResponseQueue(createLLMResponseMock, [ + toolCall({ + id: 'call_search', + name: 'search', + args: { + q: 'FastGPT' + } + }), + text({ requestId: 'req_final', content: 'final answer' }) + ]); + + await runAgentLoop({ + maxRunAgentTimes: 5, + body: { + model: 'gpt-4', + stream: true, + messages: [ + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: 'search FastGPT' + } + ], + tools: [searchTool] + }, + usagePush: vi.fn(), + isAborted: () => false, + onRunTool, + onRunInteractiveTool: vi.fn() + }); + + expect(compressRequestMessagesMock).toHaveBeenCalledTimes(2); + expect(compressRequestMessagesMock.mock.calls[0][0].messageTokens).toBeUndefined(); + expect(compressRequestMessagesMock.mock.calls[0][0].tools).toEqual([searchTool]); + expect(compressRequestMessagesMock.mock.calls[1][0].messageTokens).toBe(1028); + expect(compressRequestMessagesMock.mock.calls[1][0].tools).toEqual([searchTool]); + expect(countGptMessagesTokensMock.mock.calls[0][0].messages).toEqual([ + expect.objectContaining({ + role: ChatCompletionRequestMessageRoleEnum.Assistant + }) + ]); + expect(countGptMessagesTokensMock.mock.calls[1][0].messages).toEqual([ + expect.objectContaining({ + role: ChatCompletionRequestMessageRoleEnum.Tool, + content: 'search result' + }) + ]); + }); + it('uses request control tool choice while streaming immediately', async () => { const streamed: string[] = []; const order: string[] = []; diff --git a/packages/service/test/core/ai/llm/compress/index.test.ts b/packages/service/test/core/ai/llm/compress/index.test.ts index f8bba5affc8c..866530130410 100644 --- a/packages/service/test/core/ai/llm/compress/index.test.ts +++ b/packages/service/test/core/ai/llm/compress/index.test.ts @@ -2,7 +2,10 @@ import { ChatCompletionRequestMessageRoleEnum, ModelTypeEnum } from '@fastgpt/global/core/ai/constants'; -import type { ChatCompletionMessageParam } from '@fastgpt/global/core/ai/llm/type'; +import type { + ChatCompletionMessageParam, + ChatCompletionTool +} from '@fastgpt/global/core/ai/llm/type'; import type { LLMModelItemType } from '@fastgpt/global/core/ai/model.schema'; import { beforeEach, describe, expect, it, vi } from 'vitest'; @@ -10,12 +13,14 @@ const { createLLMResponseMock, countGptMessagesTokensMock, countPromptTokensMock, - formatModelChars2PointsMock + formatModelChars2PointsMock, + loggerWarnMock } = vi.hoisted(() => ({ createLLMResponseMock: vi.fn(), countGptMessagesTokensMock: vi.fn(), countPromptTokensMock: vi.fn(), - formatModelChars2PointsMock: vi.fn() + formatModelChars2PointsMock: vi.fn(), + loggerWarnMock: vi.fn() })); vi.mock('@fastgpt/service/core/ai/llm/request', () => ({ @@ -36,6 +41,22 @@ vi.mock('@fastgpt/service/support/wallet/usage/utils', () => ({ formatModelChars2Points: formatModelChars2PointsMock })); +vi.mock('@fastgpt/service/common/logger', () => ({ + LogCategories: { + MODULE: { + AI: { + LLM_COMPRESS: 'ai:llm_compress' + } + } + }, + getLogger: () => ({ + debug: vi.fn(), + error: vi.fn(), + info: vi.fn(), + warn: loggerWarnMock + }) +})); + import { compressLargeContent, compressRequestMessages, @@ -56,6 +77,16 @@ const model: LLMModelItemType = { reasoning: false }; +const largeContextModel: LLMModelItemType = { + ...model, + maxContext: 32000 +}; + +const toolCompressionModel: LLMModelItemType = { + ...model, + maxContext: 12000 +}; + const createMessages = (): ChatCompletionMessageParam[] => [ { role: ChatCompletionRequestMessageRoleEnum.System, @@ -88,6 +119,23 @@ const createMessages = (): ChatCompletionMessageParam[] => [ } ]; +const searchTool: ChatCompletionTool = { + type: 'function', + function: { + name: 'search', + description: 'Search test data', + parameters: { + type: 'object', + properties: { + q: { + type: 'string' + } + }, + required: ['q'] + } + } +}; + const mockDefaultUsagePoints = () => { formatModelChars2PointsMock.mockReturnValue({ totalPoints: 3 @@ -173,6 +221,7 @@ describe('compressRequestMessages', () => { expect(userPrompt).toContain(''); expect(userPrompt).toContain('recent user 3'); expect(userPrompt).toContain(''); + expect(userPrompt).toContain('Target maximum output tokens: 4000'); expect(compressPrompt).not.toContain('最近消息预览'); expect(createLLMResponseMock.mock.calls[0][0].body.max_tokens).toBeUndefined(); }); @@ -244,7 +293,7 @@ describe('compressRequestMessages', () => { ); }); - it('should not replace an over-target LLM checkpoint when final messages stay within the production threshold', async () => { + it('should keep the LLM checkpoint when completion output exceeds half context', async () => { countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(2800); createLLMResponseMock.mockResolvedValue({ answerText: @@ -267,49 +316,102 @@ describe('compressRequestMessages', () => { expect(createLLMResponseMock.mock.calls[0][0].body.max_tokens).toBeUndefined(); }); - it('should use deterministic fallback only when the final compressed messages fit the threshold', async () => { - countGptMessagesTokensMock - .mockResolvedValueOnce(5000) - .mockResolvedValueOnce(4500) - .mockResolvedValueOnce(1200) - .mockResolvedValueOnce(3000) - .mockResolvedValueOnce(2600); + it('should keep the LLM checkpoint when final checkpoint tokens still exceed 20 percent context target', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(30000).mockResolvedValueOnce(7000); createLLMResponseMock.mockResolvedValue({ - answerText: '\n超长 LLM checkpoint\n', + answerText: '\nLLM checkpoint summary\n', usage: { inputTokens: 500, - outputTokens: 3000 + outputTokens: 1000 }, - requestId: 'req_deterministic_budget', + requestId: 'req_llm_checkpoint_only', finish_reason: 'stop' }); const messages = createMessages(); const result = await compressRequestMessages({ messages, - model + model: largeContextModel }); - expect(result.contextCheckpoint).toContain('## Source History Excerpts'); + expect(createLLMResponseMock).toHaveBeenCalled(); expect(result.messages).not.toBe(messages); - expect(countGptMessagesTokensMock).toHaveBeenNthCalledWith(3, { - messages: [messages[0]] + expect(result.messageTokens).toBe(7000); + expect(result.contextCheckpoint).toBe( + '\nLLM checkpoint summary\n' + ); + expect(result.contextCheckpoint).not.toContain('## Source History Excerpts'); + expect(countGptMessagesTokensMock).toHaveBeenCalledTimes(2); + expect(loggerWarnMock).toHaveBeenCalledWith( + 'Message compression result still exceeds target', + expect.objectContaining({ + reason: 'compressed_messages_over_threshold', + compressedTokens: 7000, + compressedTokenLimit: 6400, + maxContext: 32000, + requestId: 'req_llm_checkpoint_only' + }) + ); + }); + + it('should keep original messages when below compression threshold', async () => { + countGptMessagesTokensMock.mockResolvedValue(100); + + const messages = createMessages(); + const result = await compressRequestMessages({ + messages, + model }); + + expect(result).toEqual({ messages, messageTokens: 100 }); + expect(createLLMResponseMock).not.toHaveBeenCalled(); }); - it('should return original messages when compressed checkpoint still exceeds the production threshold', async () => { - countGptMessagesTokensMock - .mockResolvedValueOnce(5000) - .mockResolvedValueOnce(4500) - .mockResolvedValueOnce(4500) - .mockResolvedValueOnce(1200); + it('should include tools schema when counting request message tokens', async () => { + countGptMessagesTokensMock.mockResolvedValue(100); + const messages: ChatCompletionMessageParam[] = [ + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: 'hello' + } + ]; + + const result = await compressRequestMessages({ + messages, + model, + tools: [searchTool] + }); + + expect(result).toEqual({ messages, messageTokens: 100 }); + expect(countGptMessagesTokensMock).toHaveBeenCalledWith({ + messages, + tools: [searchTool] + }); + expect(createLLMResponseMock).not.toHaveBeenCalled(); + }); + + it('should reuse provided request message token count', async () => { + const messages = createMessages(); + const result = await compressRequestMessages({ + messageTokens: 100, + messages, + model + }); + + expect(result).toEqual({ messages, messageTokens: 100 }); + expect(countGptMessagesTokensMock).not.toHaveBeenCalled(); + expect(createLLMResponseMock).not.toHaveBeenCalled(); + }); + + it('should send formatted user and assistant content to LLM for over-threshold tool-call histories', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(1200); createLLMResponseMock.mockResolvedValue({ - answerText: '\n超长 tool checkpoint\n', + answerText: '\norders summary\n', usage: { - inputTokens: 500, - outputTokens: 3000 + inputTokens: 50, + outputTokens: 10 }, - requestId: 'req_over_budget_tool_history', + requestId: 'req_formatted_tool_history', finish_reason: 'stop' }); const messages: ChatCompletionMessageParam[] = [ @@ -319,7 +421,12 @@ describe('compressRequestMessages', () => { }, { role: ChatCompletionRequestMessageRoleEnum.User, - content: 'Run search_orders.' + content: + 'Case alpha. Available tools: [{"type":"function","function":{"name":"search_orders","parameters":{"properties":{"customerId":{"type":"string"}}}}}]' + }, + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: 'Find recent orders for customer c_123.' }, { role: ChatCompletionRequestMessageRoleEnum.Assistant, @@ -330,10 +437,15 @@ describe('compressRequestMessages', () => { type: 'function', function: { name: 'search_orders', - arguments: '{"customerId":"c_123"}' + arguments: '{"customerId":"c_123","limit":5}' } } ] + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_search_orders', + content: '{"orders":[{"id":"ord_1","status":"paid"}]}' } ]; @@ -342,26 +454,96 @@ describe('compressRequestMessages', () => { model }); - expect(result.messages).toBe(messages); - expect(result.contextCheckpoint).toBeUndefined(); - expect(result.requestIds).toEqual(['req_over_budget_tool_history']); + expect(createLLMResponseMock).toHaveBeenCalledTimes(1); + expect(result.contextCheckpoint).toBe( + '\norders summary\n' + ); + const userPrompt = createLLMResponseMock.mock.calls[0][0].body.messages[1].content; + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('"role": "user"'); + expect(userPrompt).toContain('"role": "assistant"'); + expect(userPrompt).toContain('Find recent orders for customer c_123'); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('{\\"customerId\\":\\"c_123\\",\\"limit\\":5}'); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('ord_1'); + expect(userPrompt).toContain('paid'); + expect(userPrompt).not.toContain('"role": "tool"'); + expect(userPrompt).not.toContain('tool_calls'); + expect(userPrompt).not.toContain('tool_call_id'); + expect(userPrompt).not.toContain('call_search_orders'); }); - it('should keep original messages when below compression threshold', async () => { - countGptMessagesTokensMock.mockResolvedValue(100); + it('should merge assistant content and tool calls into the same assistant checkpoint item', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(1200); + createLLMResponseMock.mockResolvedValue({ + answerText: '\nassistant content tool summary\n', + usage: { + inputTokens: 50, + outputTokens: 10 + }, + requestId: 'req_assistant_content_tool_history', + finish_reason: 'stop' + }); + const messages: ChatCompletionMessageParam[] = [ + { + role: ChatCompletionRequestMessageRoleEnum.System, + content: 'system prompt' + }, + { + role: ChatCompletionRequestMessageRoleEnum.User, + content: 'Check customer c_456 before answering.' + }, + { + role: ChatCompletionRequestMessageRoleEnum.Assistant, + content: 'I will inspect the latest customer profile before answering.', + tool_calls: [ + { + id: 'call_get_customer', + type: 'function', + function: { + name: 'get_customer_profile', + arguments: '{"customerId":"c_456"}' + } + } + ] + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_get_customer', + content: '{"customer":{"id":"c_456","tier":"enterprise","region":"NA"}}' + } + ]; - const messages = createMessages(); - const result = await compressRequestMessages({ + await compressRequestMessages({ messages, model }); - expect(result).toEqual({ messages }); - expect(createLLMResponseMock).not.toHaveBeenCalled(); + const userPrompt = createLLMResponseMock.mock.calls[0][0].body.messages[1].content; + expect(userPrompt).toContain('I will inspect the latest customer profile before answering.'); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('{\\"customerId\\":\\"c_456\\"}'); + expect(userPrompt).toContain('enterprise'); + expect(userPrompt).not.toContain('"role": "tool"'); + expect(userPrompt).not.toContain('tool_calls'); + expect(userPrompt).not.toContain('tool_call_id'); + expect(userPrompt).not.toContain('call_get_customer'); }); - it('should build a local structured checkpoint for over-threshold tool-call histories', async () => { + it('should match multiple tool responses by tool call id when tool messages are out of order', async () => { countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(1200); + createLLMResponseMock.mockResolvedValue({ + answerText: '\nmulti tool summary\n', + usage: { + inputTokens: 50, + outputTokens: 10 + }, + requestId: 'req_multi_tool_history', + finish_reason: 'stop' + }); const messages: ChatCompletionMessageParam[] = [ { role: ChatCompletionRequestMessageRoleEnum.System, @@ -369,63 +551,154 @@ describe('compressRequestMessages', () => { }, { role: ChatCompletionRequestMessageRoleEnum.User, - content: - 'Case alpha. Available tools: [{"type":"function","function":{"name":"search_orders","parameters":{"properties":{"customerId":{"type":"string"}}}}}]' - }, - { - role: ChatCompletionRequestMessageRoleEnum.User, - content: 'Find recent orders for customer c_123.' + content: 'Compare recent orders and contracts for Acme.' }, { role: ChatCompletionRequestMessageRoleEnum.Assistant, - content: null, + content: 'I need both order and contract data.', tool_calls: [ { - id: 'call_search_orders', + id: 'call_orders_multi', type: 'function', function: { name: 'search_orders', - arguments: '{"customerId":"c_123","limit":5}' + arguments: '{"company":"Acme","limit":2}' + } + }, + { + id: 'call_contracts_multi', + type: 'function', + function: { + name: 'search_contracts', + arguments: '{"company":"Acme","year":2025}' } } ] + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_contracts_multi', + content: '{"contracts":[{"id":"ctr_multi_1","status":"active"}]}' + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_orders_multi', + content: '{"orders":[{"id":"ord_multi_1","status":"paid"}]}' } ]; - const result = await compressRequestMessages({ + await compressRequestMessages({ messages, model }); - expect(createLLMResponseMock).not.toHaveBeenCalled(); - expect(result.usage).toBeUndefined(); - expect(result.messages).toEqual([ - messages[0], + const userPrompt = createLLMResponseMock.mock.calls[0][0].body.messages[1].content; + const ordersToolIndex = userPrompt.indexOf(''); + const ordersResultIndex = userPrompt.indexOf('ord_multi_1'); + const contractsToolIndex = userPrompt.indexOf(''); + const contractsResultIndex = userPrompt.indexOf('ctr_multi_1'); + + expect(ordersToolIndex).toBeGreaterThan(-1); + expect(contractsToolIndex).toBeGreaterThan(-1); + expect(ordersResultIndex).toBeGreaterThan(ordersToolIndex); + expect(contractsResultIndex).toBeGreaterThan(contractsToolIndex); + expect(ordersToolIndex).toBeLessThan(contractsToolIndex); + expect(userPrompt).toContain('{\\"company\\":\\"Acme\\",\\"limit\\":2}'); + expect(userPrompt).toContain('{\\"company\\":\\"Acme\\",\\"year\\":2025}'); + expect(userPrompt).not.toContain('call_orders_multi'); + expect(userPrompt).not.toContain('call_contracts_multi'); + expect(userPrompt).not.toContain('"role": "tool"'); + }); + + it('should keep consecutive assistant messages separate while matching each tool response by id', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(1200); + createLLMResponseMock.mockResolvedValue({ + answerText: '\nconsecutive assistant summary\n', + usage: { + inputTokens: 50, + outputTokens: 10 + }, + requestId: 'req_consecutive_assistant_tools', + finish_reason: 'stop' + }); + const messages: ChatCompletionMessageParam[] = [ + { + role: ChatCompletionRequestMessageRoleEnum.System, + content: 'system prompt' + }, { role: ChatCompletionRequestMessageRoleEnum.User, - content: result.contextCheckpoint, - hideInUI: true + content: 'Gather order and contract context.' + }, + { + role: ChatCompletionRequestMessageRoleEnum.Assistant, + content: 'First I will inspect orders.', + tool_calls: [ + { + id: 'call_consecutive_orders', + type: 'function', + function: { + name: 'search_orders', + arguments: '{"customerId":"c_789"}' + } + } + ] + }, + { + role: ChatCompletionRequestMessageRoleEnum.Assistant, + content: 'Then I will inspect contracts.', + tool_calls: [ + { + id: 'call_consecutive_contracts', + type: 'function', + function: { + name: 'search_contracts', + arguments: '{"customerId":"c_789"}' + } + } + ] + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_consecutive_contracts', + content: '{"contracts":[{"id":"ctr_consecutive_1"}]}' + }, + { + role: ChatCompletionRequestMessageRoleEnum.Tool, + tool_call_id: 'call_consecutive_orders', + content: '{"orders":[{"id":"ord_consecutive_1"}]}' } - ]); - expect(result.contextCheckpoint).toContain(''); - expect(result.contextCheckpoint).toContain('Case alpha. Available tools:'); - expect(result.contextCheckpoint).toContain('search_orders'); - expect(result.contextCheckpoint).toContain('"customerId":"c_123"'); - expect(result.contextCheckpoint).toContain('Find recent orders for customer c_123'); + ]; + + await compressRequestMessages({ + messages, + model + }); + + const userPrompt = createLLMResponseMock.mock.calls[0][0].body.messages[1].content; + const firstAssistantIndex = userPrompt.indexOf('First I will inspect orders.'); + const secondAssistantIndex = userPrompt.indexOf('Then I will inspect contracts.'); + + expect(firstAssistantIndex).toBeGreaterThan(-1); + expect(secondAssistantIndex).toBeGreaterThan(firstAssistantIndex); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('ord_consecutive_1'); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('ctr_consecutive_1'); + expect(userPrompt).not.toContain('call_consecutive_orders'); + expect(userPrompt).not.toContain('call_consecutive_contracts'); + expect(userPrompt).not.toContain('"role": "tool"'); }); - it('should count system messages before accepting a local structured checkpoint', async () => { - countGptMessagesTokensMock - .mockResolvedValueOnce(5000) - .mockResolvedValueOnce(4800) - .mockResolvedValueOnce(260); + it('should count tools schema but not send tools schema to checkpoint compression LLM', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(260); createLLMResponseMock.mockResolvedValue({ - answerText: '\nllm fallback summary\n', + answerText: '\ntool schema counted summary\n', usage: { inputTokens: 50, outputTokens: 10 }, - requestId: 'req_structured_with_system_budget', + requestId: 'req_without_tools_schema', finish_reason: 'stop' }); const messages: ChatCompletionMessageParam[] = [ @@ -456,18 +729,25 @@ describe('compressRequestMessages', () => { const result = await compressRequestMessages({ messages, - model + model, + tools: [searchTool] }); expect(createLLMResponseMock).toHaveBeenCalled(); - expect(result.contextCheckpoint).toContain('llm fallback summary'); + expect(result.contextCheckpoint).toContain('tool schema counted summary'); + expect(createLLMResponseMock.mock.calls[0][0].body.tools).toBeUndefined(); + expect(countGptMessagesTokensMock).toHaveBeenNthCalledWith(1, { + messages, + tools: [searchTool] + }); + expect(countGptMessagesTokensMock).toHaveBeenNthCalledWith(2, { + messages: [messages[0], expect.objectContaining({ hideInUI: true })], + tools: [searchTool] + }); }); - it('should include generic tool call memory in checkpoint compression prompt', async () => { - countGptMessagesTokensMock - .mockResolvedValueOnce(5000) - .mockResolvedValueOnce(5000) - .mockResolvedValueOnce(2000); + it('should not include tool call memory in checkpoint compression prompt', async () => { + countGptMessagesTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(2000); createLLMResponseMock.mockResolvedValue({ answerText: '\ntool summary\n', usage: { @@ -519,11 +799,16 @@ describe('compressRequestMessages', () => { '\ntool summary\n' ); const userPrompt = createLLMResponseMock.mock.calls[0][0].body.messages[1].content; - expect(userPrompt).toContain(''); - expect(userPrompt).toContain('fn=search_contracts'); - expect(userPrompt).toContain('args={"company":"Acme","year":2025}'); - expect(userPrompt).toContain('user=Search enterprise contracts'); - expect(userPrompt).toContain('"id":"ctr_2025_001"'); + expect(userPrompt).not.toContain(''); + expect(userPrompt).not.toContain('fn=search_contracts'); + expect(userPrompt).not.toContain('args={"company":"Acme","year":2025}'); + expect(userPrompt).toContain(''); + expect(userPrompt).toContain('{\\"company\\":\\"Acme\\",\\"year\\":2025}'); + expect(userPrompt).toContain('Search enterprise contracts signed by Acme in 2025.'); + expect(userPrompt).toContain('ctr_2025_001'); + expect(userPrompt).not.toContain('"role": "tool"'); + expect(userPrompt).not.toContain('tool_calls'); + expect(userPrompt).not.toContain('tool_call_id'); }); it('should use the full request context to decide checkpoint compression', async () => { @@ -552,7 +837,7 @@ describe('compressRequestMessages', () => { ]; countGptMessagesTokensMock.mockImplementation( async (input: { messages: ChatCompletionMessageParam[] }) => - input.messages === messages ? 4000 : 100 + input.messages === messages ? 4001 : 100 ); const result = await compressRequestMessages({ @@ -691,7 +976,7 @@ describe('compressRequestMessages', () => { model }); - expect(result).toEqual({ messages }); + expect(result).toEqual({ messages, messageTokens: 6000 }); }); }); @@ -1131,8 +1416,36 @@ describe('compressToolResponse', () => { expect(createLLMResponseMock).not.toHaveBeenCalled(); }); + it('should keep small tool responses unchanged without any processing', async () => { + countPromptTokensMock.mockResolvedValue(800); + const response = JSON.stringify( + { + rows: [ + { + id: 'keep_format_001', + content: 'Small JSON should not be minified or structurally summarized.' + } + ] + }, + null, + 2 + ); + + const result = await compressToolResponse({ + response, + model + }); + + expect(result).toEqual({ compressed: response }); + expect(countPromptTokensMock).toHaveBeenCalledTimes(1); + expect(createLLMResponseMock).not.toHaveBeenCalled(); + }); + it('should minify JSON tool responses without LLM when minified content fits the budget', async () => { - countPromptTokensMock.mockResolvedValue(100); + countPromptTokensMock + .mockResolvedValueOnce(5000) + .mockResolvedValueOnce(3000) + .mockResolvedValueOnce(3000); const response = JSON.stringify( { source: 'tool_call_log', @@ -1194,13 +1507,11 @@ describe('compressToolResponse', () => { const result = await compressToolResponse({ response, - model, - compressedTokenLimit: 1200, - currentMessagesTokens: 0, - toolLength: 1 + model: toolCompressionModel }); expect(createLLMResponseMock).not.toHaveBeenCalled(); + expect(countPromptTokensMock).toHaveBeenCalledTimes(3); const compressed = JSON.parse(result.compressed); expect(compressed.source).toBe('tool_call_log'); expect(result.compressed).toContain('rows'); @@ -1212,7 +1523,11 @@ describe('compressToolResponse', () => { }); it('should summarize larger JSON tool responses structurally without LLM', async () => { - countPromptTokensMock.mockResolvedValueOnce(900).mockResolvedValueOnce(180); + countPromptTokensMock + .mockResolvedValueOnce(5000) + .mockResolvedValueOnce(4500) + .mockResolvedValueOnce(180) + .mockResolvedValueOnce(180); const response = JSON.stringify({ source: 'tool_call_log', rows: [ @@ -1247,27 +1562,48 @@ describe('compressToolResponse', () => { const result = await compressToolResponse({ response, - model, - compressedTokenLimit: 1200, - currentMessagesTokens: 0, - toolLength: 1 + model: toolCompressionModel }); expect(createLLMResponseMock).not.toHaveBeenCalled(); expect(result.compressed).toContain('JSON structural summary'); expect(result.compressed).toContain('root keys: source, rows'); - expect(result.compressed).toContain('important scalar values: tool_call_log; multiple_001'); + expect(result.compressed).not.toContain('importantScalarSummary'); + expect(result.compressed).toContain('"importantScalarValues"'); + expect(result.compressed).toContain('source=tool_call_log'); + expect(result.compressed).toContain('id=multiple_001'); expect(result.compressed).toContain('rows: array(length=2)'); - expect(result.compressed).toContain('rows[0].id: multiple_001'); - expect(result.compressed).toContain('rows[0].tools[0].function.name: lawsuits_search'); + expect(result.compressed).toContain('rows[0] keys: id, tools, messages'); + expect(result.compressed).not.toContain('rows[0].id: multiple_001'); + expect(result.compressed).not.toContain('rows[0].tools[0].function.name: lawsuits_search'); expect(result.compressed).not.toContain('{"source"'); }); - it('should use dynamic available context as the tool compressed token limit', async () => { - mockPromptTokensForLlmCompression({ - cleanedTokens: 1600, - initialTokens: 1600 + it('should lightly process medium tool responses without LLM compression', async () => { + countPromptTokensMock.mockResolvedValueOnce(5000).mockResolvedValueOnce(4500); + + const result = await compressToolResponse({ + response: + 'tool response https://example.com/a/b/c with image ![chart](https://example.com/chart.png)\n\n\nend', + model: toolCompressionModel, + reasoningEffort: 'high' }); + + expect(createLLMResponseMock).not.toHaveBeenCalled(); + expect(result.compressed).not.toContain('https://example.com'); + expect(result.compressed).toContain('tool response'); + expect(result.compressed).toContain('[chart]'); + }); + + it('should compress large tool responses with 20 percent context as target', async () => { + countPromptTokensMock + .mockResolvedValueOnce(7000) + .mockResolvedValueOnce(6500) + .mockResolvedValueOnce(6500) + .mockResolvedValueOnce(6500) + .mockResolvedValueOnce(6500) + .mockResolvedValueOnce(50) + .mockResolvedValue(4200); createLLMResponseMock.mockResolvedValue({ answerText: 'compressed tool response', usage: { @@ -1279,43 +1615,29 @@ describe('compressToolResponse', () => { const result = await compressToolResponse({ response: 'tool response', - model, - currentMessagesTokens: 1000, - toolLength: 2, + model: toolCompressionModel, reasoningEffort: 'high' }); - expect(createLLMResponseMock).toHaveBeenCalledTimes(1); + expect(createLLMResponseMock).toHaveBeenCalled(); expect(result.compressed).toBe('compressed tool response'); expect(result.usage?.moduleName).toBe('account_usage:tool_response_compress'); expect(result.requestIds).toEqual(['req_tool']); expect(createLLMResponseMock.mock.calls[0][0].body.reasoning_effort).toBe('high'); - }); - - it('should respect caller provided compressed token limit for tool response', async () => { - mockPromptTokensForLlmCompression({ - cleanedTokens: 1200, - initialTokens: 1200 - }); - createLLMResponseMock.mockResolvedValue({ - answerText: 'budgeted tool response', - usage: { - inputTokens: 30, - outputTokens: 6 - }, - requestId: 'req_tool_budget' - }); - - const result = await compressToolResponse({ - response: 'tool response', - model, - compressedTokenLimit: 1000, - currentMessagesTokens: 0, - toolLength: 1 - }); - - expect(createLLMResponseMock).toHaveBeenCalledTimes(1); - expect(result.compressed).toBe('budgeted tool response'); - expect(createLLMResponseMock.mock.calls[0][0].body.max_tokens).toBeUndefined(); + expect(createLLMResponseMock.mock.calls[0][0].body.messages[1].content).toContain( + 'Target maximum output tokens: 2662' + ); + expect(loggerWarnMock).toHaveBeenCalledWith( + 'Tool response compression result still exceeds target', + expect.objectContaining({ + reason: 'compressed_tool_response_over_target', + originalTokens: 7000, + lightProcessedTokens: 6500, + compressedTokens: 4200, + compressedTokenLimit: 4096, + maxContext: 12000, + requestIds: ['req_tool'] + }) + ); }); }); diff --git a/packages/service/test/core/ai/sandbox/service/cron.test.ts b/packages/service/test/core/ai/sandbox/service/cron.test.ts index 786ebd7de7c9..5720d4875e22 100644 --- a/packages/service/test/core/ai/sandbox/service/cron.test.ts +++ b/packages/service/test/core/ai/sandbox/service/cron.test.ts @@ -1,4 +1,4 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; const cronMocks = vi.hoisted(() => ({ setCron: vi.fn(), @@ -39,10 +39,16 @@ import { cronJob } from '@fastgpt/service/core/ai/sandbox/service/cron'; describe('sandbox cron service', () => { beforeEach(() => { vi.clearAllMocks(); + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-06-24T01:00:00.000Z')); cronMocks.checkTimerLock.mockResolvedValue(true); cronMocks.archiveInactiveSandboxes.mockResolvedValue(undefined); }); + afterEach(() => { + vi.useRealTimers(); + }); + it('registers a cron task and skips when no inactive sandbox exists', async () => { cronMocks.findInactiveRunningSandboxResources.mockResolvedValueOnce([]); @@ -51,7 +57,9 @@ describe('sandbox cron service', () => { await callback(); expect(cronMocks.setCron).toHaveBeenCalledWith('*/5 * * * *', expect.any(Function)); - expect(cronMocks.findInactiveRunningSandboxResources).toHaveBeenCalledWith(expect.any(Date)); + expect(cronMocks.findInactiveRunningSandboxResources).toHaveBeenCalledWith( + new Date('2026-06-24T00:50:00.000Z') + ); expect(cronMocks.stopSandboxResources).not.toHaveBeenCalled(); }); diff --git a/packages/service/test/core/app/mcp.test.ts b/packages/service/test/core/app/mcp.test.ts index fc0808fbb2b7..b93be41209e4 100644 --- a/packages/service/test/core/app/mcp.test.ts +++ b/packages/service/test/core/app/mcp.test.ts @@ -1,4 +1,6 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { afterEach, describe, it, expect, vi, beforeEach } from 'vitest'; +import http from 'http'; +import os from 'os'; // --- Hoisted mocks --- const { mockDereference, mockMongoAppFind } = vi.hoisted(() => ({ @@ -12,15 +14,22 @@ vi.mock('@apidevtools/json-schema-ref-parser', () => ({ } })); -vi.mock('@fastgpt/service/core/app/schema', () => ({ +vi.mock('../../../core/app/schema', () => ({ MongoApp: { find: mockMongoAppFind } })); import { StreamableHTTPError } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; -import { MCPClient, assertMCPUrlNotInternal, getMCPChildren } from '@fastgpt/service/core/app/mcp'; +import { + MCPClient, + assertMCPUrlNotInternal, + createMcpSafeFetch, + getMCPChildren +} from '../../../core/app/mcp'; import type { AppSchemaType } from '@fastgpt/global/core/app/type'; +import { PRIVATE_URL_TEXT } from '../../../common/system/utils'; +import { serviceEnv } from '../../../env'; // Access private client via prototype for spying const getPrivateClient = (mcpClient: MCPClient) => @@ -36,6 +45,48 @@ beforeEach(() => { vi.restoreAllMocks(); }); +const mutableServiceEnv = serviceEnv as { CHECK_INTERNAL_IP: boolean }; +const originalCheckInternalIp = serviceEnv.CHECK_INTERNAL_IP; + +afterEach(() => { + mutableServiceEnv.CHECK_INTERNAL_IP = originalCheckInternalIp; +}); + +const listen = (handler: http.RequestListener, host = '127.0.0.1') => + new Promise((resolve) => { + const server = http.createServer(handler); + server.listen(0, host, () => resolve(server)); + }); + +const closeServer = (server: http.Server) => + new Promise((resolve, reject) => { + server.close((err) => (err ? reject(err) : resolve())); + }); + +const getServerPort = (server: http.Server): number => { + const address = server.address(); + if (!address || typeof address === 'string') { + throw new Error('Invalid test server address'); + } + return address.port; +}; + +/** + * 构造一个非 loopback 的本机访问地址,用于模拟“初始 MCP URL 通过 SSRF 校验”。 + * CHECK_INTERNAL_IP=false 时私网地址会放行,但 loopback/metadata 仍然恒拦截。 + */ +const getReachablePrivateHost = () => { + const interfaces = os.networkInterfaces(); + for (const items of Object.values(interfaces)) { + for (const item of items || []) { + if (item.family === 'IPv4' && !item.internal) { + return item.address; + } + } + } + return undefined; +}; + describe('MCPClient', () => { const config = { url: 'https://example.com/mcp', headers: { Authorization: 'Bearer test' } }; @@ -200,7 +251,9 @@ describe('MCPClient', () => { const tools = await mcpClient.getTools(); - const homeSchema = tools[0].inputSchema.properties!['home'] as any; + expect(tools[0]).toBeDefined(); + const inputSchema = tools[0]!.inputSchema!; + const homeSchema = inputSchema.properties!['home'] as any; expect(homeSchema).toEqual({ type: 'object', properties: { street: { type: 'string' }, city: { type: 'string' } } @@ -271,7 +324,9 @@ describe('MCPClient', () => { const tools = await mcpClient.getTools(); // Verify nested refs are fully resolved - const ownerProps = (tools[0].inputSchema.properties!['owner'] as any).properties; + expect(tools[0]).toBeDefined(); + const inputSchema = tools[0]!.inputSchema!; + const ownerProps = (inputSchema.properties!['owner'] as any).properties; expect(ownerProps.name.properties).toEqual({ first: { type: 'string' }, last: { type: 'string' } @@ -311,7 +366,9 @@ describe('MCPClient', () => { const tools = await mcpClient.getTools(); - const tagsSchema = tools[0].inputSchema.properties!['tags'] as any; + expect(tools[0]).toBeDefined(); + const inputSchema = tools[0]!.inputSchema!; + const tagsSchema = inputSchema.properties!['tags'] as any; expect(tagsSchema.items).toEqual({ type: 'object', properties: { label: { type: 'string' } } @@ -485,6 +542,130 @@ describe('MCPClient', () => { }); }); +describe('createMcpSafeFetch', () => { + it('should follow safe redirects hop by hop', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + const targetServer = await listen((req, res) => { + res.end(JSON.stringify({ ok: true, url: req.url })); + }, '0.0.0.0'); + const targetPort = getServerPort(targetServer); + + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://${carrierHost}:${targetPort}/mcp-target`); + res.end('redirect'); + }, '0.0.0.0'); + const redirectPort = getServerPort(redirectServer); + + try { + const response = await createMcpSafeFetch()(`http://${carrierHost}:${redirectPort}/mcp`); + + expect(await response.json()).toEqual({ ok: true, url: '/mcp-target' }); + } finally { + await closeServer(redirectServer); + await closeServer(targetServer); + } + }); + + it('should block redirects to loopback addresses', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + const protectedServer = await listen((req, res) => { + res.end('INTERNAL-ONLY-RESPONSE'); + }); + const protectedPort = getServerPort(protectedServer); + + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://127.0.0.1:${protectedPort}/mcp`); + res.end('redirect'); + }, '0.0.0.0'); + const redirectPort = getServerPort(redirectServer); + + try { + await expect(createMcpSafeFetch()(`http://${carrierHost}:${redirectPort}/mcp`)).rejects.toBe( + PRIVATE_URL_TEXT + ); + } finally { + await closeServer(redirectServer); + await closeServer(protectedServer); + } + }); + + it('should enforce max redirect count', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + let redirectPort = 0; + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://${carrierHost}:${redirectPort}/loop`); + res.end('redirect'); + }, '0.0.0.0'); + redirectPort = getServerPort(redirectServer); + + try { + await expect( + createMcpSafeFetch({ maxRedirects: 1 })(`http://${carrierHost}:${redirectPort}/loop`) + ).rejects.toThrow('Maximum MCP redirects exceeded'); + } finally { + await closeServer(redirectServer); + } + }); + + it('should drop sensitive headers when redirect target changes', async () => { + mutableServiceEnv.CHECK_INTERNAL_IP = false; + const carrierHost = getReachablePrivateHost(); + if (!carrierHost) { + return; + } + + let receivedAuthorization: string | undefined; + let receivedCookie: string | undefined; + const targetServer = await listen((req, res) => { + receivedAuthorization = req.headers.authorization; + receivedCookie = req.headers.cookie; + res.end('ok'); + }, '0.0.0.0'); + const targetPort = getServerPort(targetServer); + + const redirectServer = await listen((req, res) => { + res.statusCode = 302; + res.setHeader('Location', `http://${carrierHost}:${targetPort}/mcp-target`); + res.end('redirect'); + }, '0.0.0.0'); + const redirectPort = getServerPort(redirectServer); + + try { + const response = await createMcpSafeFetch()(`http://${carrierHost}:${redirectPort}/mcp`, { + headers: { + Authorization: 'Bearer secret', + Cookie: 'token=secret' + } + }); + + expect(await response.text()).toBe('ok'); + expect(receivedAuthorization).toBeUndefined(); + expect(receivedCookie).toBeUndefined(); + } finally { + await closeServer(redirectServer); + await closeServer(targetServer); + } + }); +}); + describe('getMCPChildren', () => { it('should return tool list from new MCP format', async () => { const app = { diff --git a/packages/service/test/core/workflow/dispatch/ai/agent/adapter/runtime.test.ts b/packages/service/test/core/workflow/dispatch/ai/agent/adapter/runtime.test.ts index df95b92b5964..8a46310c6d63 100644 --- a/packages/service/test/core/workflow/dispatch/ai/agent/adapter/runtime.test.ts +++ b/packages/service/test/core/workflow/dispatch/ai/agent/adapter/runtime.test.ts @@ -241,7 +241,7 @@ describe('createWorkflowAgentLoopRuntime', () => { nodeId: 'agent_node-main_agent-1', moduleName: 'chat:master_agent_call', moduleType: FlowNodeTypeEnum.agent, - moduleLogo: 'core/workflow/template/agent', + moduleLogo: 'core/app/type/agentFill', model: 'GPT-4', llmRequestIds: ['req_1'], inputTokens: 10, @@ -257,7 +257,7 @@ describe('createWorkflowAgentLoopRuntime', () => { nodeId: 'agent_node-main_agent-2', moduleName: 'chat:master_agent_call', moduleType: FlowNodeTypeEnum.agent, - moduleLogo: 'core/workflow/template/agent', + moduleLogo: 'core/app/type/agentFill', model: 'GPT-4', llmRequestIds: ['req_2'], inputTokens: 6, @@ -392,19 +392,19 @@ describe('createWorkflowAgentLoopRuntime', () => { expect.objectContaining({ id: 'agent_node-1-req_empty_start', moduleName: 'chat:master_agent_call', - moduleLogo: 'core/workflow/template/agent', + moduleLogo: 'core/app/type/agentFill', llmRequestIds: ['req_empty_start'] }), expect.objectContaining({ id: 'agent_node-2-req_tool_round', moduleName: 'chat:master_agent_call', - moduleLogo: 'core/workflow/template/agent', + moduleLogo: 'core/app/type/agentFill', llmRequestIds: ['req_tool_round'] }), expect.objectContaining({ id: 'agent_node-3-req_empty_end', moduleName: 'chat:master_agent_call', - moduleLogo: 'core/workflow/template/agent', + moduleLogo: 'core/app/type/agentFill', llmRequestIds: ['req_empty_end'] }) ]); diff --git a/packages/service/test/core/workflow/dispatch/ai/agent/sub/tool/index.test.ts b/packages/service/test/core/workflow/dispatch/ai/agent/sub/tool/index.test.ts new file mode 100644 index 000000000000..02059a9c975f --- /dev/null +++ b/packages/service/test/core/workflow/dispatch/ai/agent/sub/tool/index.test.ts @@ -0,0 +1,203 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { dispatchTool } from '@fastgpt/service/core/workflow/dispatch/ai/agent/sub/tool'; +import { ReadPermissionVal } from '@fastgpt/global/support/permission/constant'; + +const { authAppByTmbIdMock, getAppVersionByIdMock, runHTTPToolMock, mcpToolCallMock } = vi.hoisted( + () => ({ + authAppByTmbIdMock: vi.fn(), + getAppVersionByIdMock: vi.fn(), + runHTTPToolMock: vi.fn(), + mcpToolCallMock: vi.fn() + }) +); + +vi.mock('@fastgpt/service/support/permission/app/auth', () => ({ + authAppByTmbId: authAppByTmbIdMock +})); + +vi.mock('@fastgpt/service/core/app/version/controller', () => ({ + getAppVersionById: getAppVersionByIdMock +})); + +vi.mock('@fastgpt/service/core/app/http', () => ({ + runHTTPTool: runHTTPToolMock +})); + +vi.mock('@fastgpt/service/core/app/mcp', () => ({ + assertMCPUrlNotInternal: vi.fn(), + MCPClient: vi.fn().mockImplementation(() => ({ + toolCall: mcpToolCallMock + })) +})); + +vi.mock('@fastgpt/service/common/logger', () => ({ + LogCategories: { + MODULE: { + APP: { + TOOL: 'tool' + }, + AI: { + LLM: 'llm' + } + } + }, + getLogger: vi.fn(() => ({ + error: vi.fn() + })) +})); + +vi.mock('@fastgpt/service/common/middle/tracks/utils', () => ({ + pushTrack: { + runSystemTool: vi.fn() + } +})); + +vi.mock('@fastgpt/service/thirdProvider/fastgptPlugin', () => ({ + pluginClient: { + runToolStream: vi.fn() + } +})); + +vi.mock('@fastgpt/service/core/app/tool/systemTool/systemTool.repo', () => ({ + SystemToolRepo: { + getInstance: vi.fn(() => ({ + getSystemToolRuntime: vi.fn() + })) + } +})); + +const createDispatchToolProps = (toolConfig: Record) => + ({ + tool: { + name: 'Agent tool', + avatar: '', + toolConfig + }, + params: { + keyword: 'fastgpt' + }, + runningAppInfo: { + id: 'attacker-app', + teamId: 'attacker-team', + tmbId: 'attacker-tmb', + name: 'Attacker workflow' + }, + runningUserInfo: { + username: 'attacker', + teamName: 'Attacker team', + memberName: 'Attacker member', + contact: '', + teamId: 'attacker-team', + tmbId: 'attacker-tmb' + }, + chatId: 'chat', + uid: 'uid', + variableState: { + get: vi.fn() + } + }) as any; + +describe('dispatchTool runtime toolset auth', () => { + beforeEach(() => { + vi.clearAllMocks(); + authAppByTmbIdMock.mockResolvedValue({ + app: { + _id: 'victim-toolset' + } + }); + }); + + it('should reject HTTP agent tool execution when running app tmb has no parent toolset permission', async () => { + authAppByTmbIdMock.mockRejectedValueOnce(new Error('unAuthApp')); + + const result = await dispatchTool( + createDispatchToolProps({ + httpTool: { + toolId: 'http-victim-toolset/sandbox_echo' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).not.toHaveBeenCalled(); + expect(runHTTPToolMock).not.toHaveBeenCalled(); + expect(result.response).toBeTruthy(); + }); + + it('should authorize HTTP parent toolset before agent tool execution', async () => { + getAppVersionByIdMock.mockResolvedValueOnce({ + nodes: [ + { + toolConfig: { + httpToolSet: { + baseUrl: 'https://example.com', + toolList: [ + { + name: 'sandbox_echo', + path: '/echo', + method: 'post' + } + ] + } + } + } + ] + }); + runHTTPToolMock.mockResolvedValueOnce({ + data: { + ok: true + } + }); + + const result = await dispatchTool( + createDispatchToolProps({ + httpTool: { + toolId: 'http-victim-toolset/sandbox_echo' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).toHaveBeenCalledWith({ + appId: 'victim-toolset', + versionId: undefined + }); + expect(runHTTPToolMock).toHaveBeenCalledWith( + expect.objectContaining({ + baseUrl: 'https://example.com', + toolPath: '/echo', + method: 'post' + }) + ); + expect(result.response).toBe(JSON.stringify({ ok: true })); + }); + + it('should reject MCP agent tool execution when running app tmb has no parent toolset permission', async () => { + authAppByTmbIdMock.mockRejectedValueOnce(new Error('unAuthApp')); + + const result = await dispatchTool( + createDispatchToolProps({ + mcpTool: { + toolId: 'mcp-victim-toolset/search' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).not.toHaveBeenCalled(); + expect(mcpToolCallMock).not.toHaveBeenCalled(); + expect(result.response).toBeTruthy(); + }); +}); diff --git a/packages/service/test/core/workflow/dispatch/ai/toolcall/toolCall.test.ts b/packages/service/test/core/workflow/dispatch/ai/toolcall/toolCall.test.ts index a250f50e9009..299deb86eadf 100644 --- a/packages/service/test/core/workflow/dispatch/ai/toolcall/toolCall.test.ts +++ b/packages/service/test/core/workflow/dispatch/ai/toolcall/toolCall.test.ts @@ -333,6 +333,32 @@ describe('runToolCall compression node responses', () => { ); }); + it('ignores context compression callbacks without usage', async () => { + runAgentLoopMock.mockImplementation(async (options) => { + options.onAfterCompressContext({ + requestIds: [], + seconds: 0.1 + }); + + return createLoopResult(); + }); + + const nodeResponseWriter = createWriter(); + const result = await runToolCall(createProps({ nodeResponseWriter })); + await Promise.resolve(); + + expect(result.requestIds).toEqual(['req_main']); + expect(nodeResponseWriter.recordWithParent).not.toHaveBeenCalled(); + expect(result.toolTotalPoints).toBe(0); + expect(result.runtimeNodeResponseSummary).toEqual( + expect.objectContaining({ + responseIds: [], + finishedNodeIds: [], + runningTime: 0 + }) + ); + }); + it('records only the compression child after onAfterToolCall when the tool workflow wrote details', async () => { const toolResponseCompressUsage = { moduleName: 'account_usage:tool_response_compress', diff --git a/packages/service/test/core/workflow/dispatch/tools/runTool.test.ts b/packages/service/test/core/workflow/dispatch/tools/runTool.test.ts index 0052e55884a8..b2ea4233155b 100644 --- a/packages/service/test/core/workflow/dispatch/tools/runTool.test.ts +++ b/packages/service/test/core/workflow/dispatch/tools/runTool.test.ts @@ -1,5 +1,226 @@ -import { describe, it, expect } from 'vitest'; -import { parseToolId } from '@fastgpt/service/core/workflow/dispatch/child/runTool'; +import { beforeEach, describe, it, expect, vi } from 'vitest'; +import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant'; +import { + dispatchRunTool, + parseToolId +} from '@fastgpt/service/core/workflow/dispatch/child/runTool'; +import { ReadPermissionVal } from '@fastgpt/global/support/permission/constant'; +import { NodeOutputKeyEnum } from '@fastgpt/global/core/workflow/constants'; + +const { authAppByTmbIdMock, getAppVersionByIdMock, runHTTPToolMock, mcpToolCallMock } = vi.hoisted( + () => ({ + authAppByTmbIdMock: vi.fn(), + getAppVersionByIdMock: vi.fn(), + runHTTPToolMock: vi.fn(), + mcpToolCallMock: vi.fn() + }) +); + +vi.mock('@fastgpt/service/support/permission/app/auth', () => ({ + authAppByTmbId: authAppByTmbIdMock +})); + +vi.mock('@fastgpt/service/core/app/version/controller', () => ({ + getAppVersionById: getAppVersionByIdMock +})); + +vi.mock('@fastgpt/service/core/app/http', () => ({ + runHTTPTool: runHTTPToolMock +})); + +vi.mock('@fastgpt/service/core/app/mcp', () => ({ + assertMCPUrlNotInternal: vi.fn(), + MCPClient: vi.fn().mockImplementation(() => ({ + toolCall: mcpToolCallMock + })) +})); + +vi.mock('@fastgpt/service/common/logger', () => ({ + LogCategories: { + MODULE: { + APP: { + TOOL: 'tool' + }, + AI: { + LLM: 'llm' + } + } + }, + getLogger: vi.fn(() => ({ + error: vi.fn() + })) +})); + +vi.mock('@fastgpt/service/common/middle/tracks/utils', () => ({ + pushTrack: { + runSystemTool: vi.fn() + } +})); + +vi.mock('@fastgpt/service/thirdProvider/fastgptPlugin', () => ({ + pluginClient: { + runToolStream: vi.fn() + } +})); + +vi.mock('@fastgpt/service/core/app/tool/systemTool/systemTool.repo', () => ({ + SystemToolRepo: { + getInstance: vi.fn(() => ({ + getSystemToolRuntime: vi.fn() + })) + } +})); + +vi.mock('@fastgpt/service/core/workflow/utils/context', () => ({ + getWorkflowContext: vi.fn(() => ({})) +})); + +const createRunToolProps = (toolConfig: Record) => + ({ + params: { + keyword: 'fastgpt' + }, + runningAppInfo: { + id: 'attacker-app', + teamId: 'attacker-team', + tmbId: 'attacker-tmb', + name: 'Attacker workflow' + }, + runningUserInfo: { + username: 'attacker', + teamName: 'Attacker team', + memberName: 'Attacker member', + contact: '', + teamId: 'attacker-team', + tmbId: 'attacker-tmb' + }, + variableState: { + get: vi.fn() + }, + node: { + nodeId: 'tool-node', + flowNodeType: FlowNodeTypeEnum.tool, + name: 'Tool node', + avatar: '', + toolConfig, + inputs: [], + outputs: [] + }, + uid: 'uid', + chatId: 'chat', + responseChatItemId: 'response', + usagePush: vi.fn() + }) as any; + +describe('dispatchRunTool runtime toolset auth', () => { + beforeEach(() => { + vi.clearAllMocks(); + authAppByTmbIdMock.mockResolvedValue({ + app: { + _id: 'victim-toolset' + } + }); + }); + + it('should reject HTTP tool execution when running app tmb has no parent toolset permission', async () => { + authAppByTmbIdMock.mockRejectedValueOnce(new Error('unAuthApp')); + + const result = await dispatchRunTool( + createRunToolProps({ + httpTool: { + toolId: 'http-victim-toolset/sandbox_echo' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).not.toHaveBeenCalled(); + expect(runHTTPToolMock).not.toHaveBeenCalled(); + expect(result.error?.[NodeOutputKeyEnum.errorText]).toBeTruthy(); + }); + + it('should authorize HTTP parent toolset before loading version and running tool', async () => { + getAppVersionByIdMock.mockResolvedValueOnce({ + nodes: [ + { + toolConfig: { + httpToolSet: { + baseUrl: 'https://example.com', + toolList: [ + { + name: 'sandbox_echo', + path: '/echo', + method: 'post' + } + ] + } + } + } + ] + }); + runHTTPToolMock.mockResolvedValueOnce({ + data: { + ok: true + } + }); + + const result = await dispatchRunTool( + createRunToolProps({ + httpTool: { + toolId: 'http-victim-toolset/sandbox_echo' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).toHaveBeenCalledWith({ + appId: 'victim-toolset', + versionId: undefined + }); + expect(runHTTPToolMock).toHaveBeenCalledWith( + expect.objectContaining({ + baseUrl: 'https://example.com', + toolPath: '/echo', + method: 'post' + }) + ); + expect(result.data).toEqual({ + [NodeOutputKeyEnum.rawResponse]: { + ok: true + }, + ok: true + }); + }); + + it('should reject MCP tool execution when running app tmb has no parent toolset permission', async () => { + authAppByTmbIdMock.mockRejectedValueOnce(new Error('unAuthApp')); + + const result = await dispatchRunTool( + createRunToolProps({ + mcpTool: { + toolId: 'mcp-victim-toolset/search' + } + }) + ); + + expect(authAppByTmbIdMock).toHaveBeenCalledWith({ + tmbId: 'attacker-tmb', + appId: 'victim-toolset', + per: ReadPermissionVal + }); + expect(getAppVersionByIdMock).not.toHaveBeenCalled(); + expect(mcpToolCallMock).not.toHaveBeenCalled(); + expect(result.error?.[NodeOutputKeyEnum.errorText]).toBeTruthy(); + }); +}); describe('parseToolId', () => { describe('新版格式: source-appId/toolName', () => { diff --git a/packages/web/components/common/DeleteConfirmInput/index.tsx b/packages/web/components/common/DeleteConfirmInput/index.tsx index 9d8dc1405ce1..d8655d11fe9a 100644 --- a/packages/web/components/common/DeleteConfirmInput/index.tsx +++ b/packages/web/components/common/DeleteConfirmInput/index.tsx @@ -1,6 +1,7 @@ import React from 'react'; import { Box, Input, VStack, type StackProps } from '@chakra-ui/react'; import { Trans, useTranslation } from 'next-i18next'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; type Props = Omit & { value: string; @@ -21,7 +22,7 @@ const DeleteConfirmInput = ({ value, confirmText, onChange, placeholder, ...prop diff --git a/packages/web/components/core/app/FileTypeSelector/index.tsx b/packages/web/components/core/app/FileTypeSelector/index.tsx index 5480171c3845..d6113cac0ae7 100644 --- a/packages/web/components/core/app/FileTypeSelector/index.tsx +++ b/packages/web/components/core/app/FileTypeSelector/index.tsx @@ -6,6 +6,7 @@ import { type FileExtensionKeyType } from '@fastgpt/global/core/app/constants'; import MyIcon from '../../../common/Icon'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; type FileTypeSelectorValue = { canSelectFile?: boolean; @@ -17,11 +18,14 @@ type FileTypeSelectorValue = { }; const fileExtensionTypeTranslationMap = new Map([ - ['canSelectFile', 'app:upload_file_extension_type_canSelectFile'], - ['canSelectImg', 'app:upload_file_extension_type_canSelectImg'], - ['canSelectVideo', 'app:upload_file_extension_type_canSelectVideo'], - ['canSelectAudio', 'app:upload_file_extension_type_canSelectAudio'], - ['canSelectCustomFileExtension', 'app:upload_file_extension_type_canSelectCustomFileExtension'] + ['canSelectFile', i18nT('app:upload_file_extension_type_canSelectFile')], + ['canSelectImg', i18nT('app:upload_file_extension_type_canSelectImg')], + ['canSelectVideo', i18nT('app:upload_file_extension_type_canSelectVideo')], + ['canSelectAudio', i18nT('app:upload_file_extension_type_canSelectAudio')], + [ + 'canSelectCustomFileExtension', + i18nT('app:upload_file_extension_type_canSelectCustomFileExtension') + ] ]); export const FileTypeSelectorPanel = ({ diff --git a/packages/web/i18n/en/app.json b/packages/web/i18n/en/app.json index 22df93a1648f..a28fd13f3bc3 100644 --- a/packages/web/i18n/en/app.json +++ b/packages/web/i18n/en/app.json @@ -250,9 +250,9 @@ "logs_source_count": "Channel users", "logs_source_count_description": "Number of users across channels", "logs_timespan_day": "Day", - "logs_timespan_week": "Week", "logs_timespan_month": "Month", "logs_timespan_quarter": "Quarter", + "logs_timespan_week": "Week", "logs_title": "Title", "logs_total": "Grand total", "logs_total_avg_duration": "Avg. duration", @@ -285,6 +285,7 @@ "my_agents": "my agents", "new_input_guide_lexicon": "New Lexicon", "no_mcp_tools_list": "No data yet, the MCP address needs to be parsed first", + "node_not_intro": "Node lacks introduction", "not_json_file": "Please select a JSON file", "not_the_newest": "Not the latest", "open_auto_execute": "Enable automatic execution", @@ -497,8 +498,12 @@ "unnamed_app": "Unnamed", "upgrade": "Upgrade", "upload_file_exists_filtered": "Duplicate files have been automatically filtered.", + "upload_file_extension_type_canSelectAudio": "Audio", "upload_file_extension_type_canSelectCustomFileExtension": "Custom file extension type", "upload_file_extension_type_canSelectCustomFileExtension_placeholder": "file extension name", + "upload_file_extension_type_canSelectFile": "Documents", + "upload_file_extension_type_canSelectImg": "Images", + "upload_file_extension_type_canSelectVideo": "Videos", "upload_file_extension_types": "Supported file types", "upload_file_max_amount": "Maximum File Quantity", "upload_file_max_amount_tip": "Maximum number of files uploaded in a single round of conversation", diff --git a/packages/web/i18n/zh-CN/app.json b/packages/web/i18n/zh-CN/app.json index 2a9f474bc569..5888cbea663b 100644 --- a/packages/web/i18n/zh-CN/app.json +++ b/packages/web/i18n/zh-CN/app.json @@ -250,9 +250,9 @@ "logs_source_count": "渠道用户", "logs_source_count_description": "各渠道用户的数量", "logs_timespan_day": "按天", - "logs_timespan_week": "按周", "logs_timespan_month": "按月", "logs_timespan_quarter": "按季度", + "logs_timespan_week": "按周", "logs_title": "标题", "logs_total": "累计", "logs_total_avg_duration": "平均时长", @@ -285,6 +285,7 @@ "my_agents": "我的 agents", "new_input_guide_lexicon": "新词库", "no_mcp_tools_list": "暂无数据，需先解析 MCP 地址", + "node_not_intro": "节点缺少介绍", "not_json_file": "请选择JSON文件", "not_the_newest": "非最新版", "open_auto_execute": "启用自动执行", @@ -497,8 +498,12 @@ "unnamed_app": "未命名", "upgrade": "去升级", "upload_file_exists_filtered": "已自动过滤重复文件", + "upload_file_extension_type_canSelectAudio": "音频", "upload_file_extension_type_canSelectCustomFileExtension": "自定义文件扩展类型", "upload_file_extension_type_canSelectCustomFileExtension_placeholder": "文件扩展名", + "upload_file_extension_type_canSelectFile": "文档", + "upload_file_extension_type_canSelectImg": "图片", + "upload_file_extension_type_canSelectVideo": "视频", "upload_file_extension_types": "支持上传的类型", "upload_file_max_amount": "最大文件数量", "upload_file_max_amount_tip": "单轮对话中最大上传文件数量", diff --git a/packages/web/i18n/zh-Hant/app.json b/packages/web/i18n/zh-Hant/app.json index d25f465b8ab4..d4b5b05d832f 100644 --- a/packages/web/i18n/zh-Hant/app.json +++ b/packages/web/i18n/zh-Hant/app.json @@ -244,9 +244,9 @@ "logs_source": "來源", "logs_source_count_description": "各渠道用戶的數量", "logs_timespan_day": "按天", - "logs_timespan_week": "按週", "logs_timespan_month": "按月", "logs_timespan_quarter": "按季度", + "logs_timespan_week": "按週", "logs_title": "標題", "logs_total": "累計", "logs_total_avg_points": "平均消耗", @@ -277,6 +277,7 @@ "my_agents": "我的 agents", "new_input_guide_lexicon": "新增詞彙庫", "no_mcp_tools_list": "暫無數據，需先解析 MCP 地址", + "node_not_intro": "節點缺少介紹", "not_json_file": "請選擇 JSON 檔案", "not_the_newest": "非最新版", "open_auto_execute": "啟用自動執行", @@ -300,6 +301,7 @@ "publish.show_whole_response": "顯示完整回應", "publish.show_whole_response_tip": "開啟後，聊天頁面可查看每個模組的詳細回應流程", "publish_channel": "發布通道", + "publish_channel.wecom.empty": "發布到企業微信機器人，請先 綁定自定義域名，並且通過域名校驗。", "publish_success": "發布成功", "question_guide_tip": "對話結束後，會為你產生 3 個引導性問題。", "raw_params": "原始參數", @@ -483,8 +485,12 @@ "unnamed_app": "未命名", "upgrade": "去升級", "upload_file_exists_filtered": "已自動過濾重複檔案", + "upload_file_extension_type_canSelectAudio": "音頻", "upload_file_extension_type_canSelectCustomFileExtension": "自定義文件擴展類型", "upload_file_extension_type_canSelectCustomFileExtension_placeholder": "文件擴展名", + "upload_file_extension_type_canSelectFile": "文件", + "upload_file_extension_type_canSelectImg": "圖片", + "upload_file_extension_type_canSelectVideo": "視頻", "upload_file_extension_types": "支持上傳的類型", "upload_file_max_amount": "最大檔案數量", "upload_file_max_amount_tip": "單輪對話中最大上傳檔案數量", diff --git a/pro b/pro index 1f9f0765e340..0103ae5b5303 160000 --- a/pro +++ b/pro @@ -1 +1 @@ -Subproject commit 1f9f0765e34037205ca59b453658b8a3caf2558a +Subproject commit 0103ae5b530380c8452abb2a479b49a3e767bf8d diff --git a/projects/app/src/components/core/app/VariableEditModal/index.tsx b/projects/app/src/components/core/app/VariableEditModal/index.tsx index c100f242c819..30fd0f1454f1 100644 --- a/projects/app/src/components/core/app/VariableEditModal/index.tsx +++ b/projects/app/src/components/core/app/VariableEditModal/index.tsx @@ -69,6 +69,12 @@ const VariableEditModal = ({ ) { setValue('defaultValue', ''); } + if ( + (typeEnum === VariableInputEnum.select || typeEnum === VariableInputEnum.multipleSelect) && + !value.list?.length + ) { + setValue('list', [{ label: '', value: '' }]); + } if (typeEnum === VariableInputEnum.datasetSelect && !value.datasetOptions) { setValue('datasetOptions', []); } diff --git a/projects/app/src/components/core/chat/components/AIResponseBox/RenderProcessingPreview.tsx b/projects/app/src/components/core/chat/components/AIResponseBox/RenderProcessingPreview.tsx index fe5c4fc06eae..ce845483aeda 100644 --- a/projects/app/src/components/core/chat/components/AIResponseBox/RenderProcessingPreview.tsx +++ b/projects/app/src/components/core/chat/components/AIResponseBox/RenderProcessingPreview.tsx @@ -1,5 +1,6 @@ import Markdown from '@/components/Markdown'; import { Box } from '@chakra-ui/react'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; import type { AIChatItemValueItemType } from '@fastgpt/global/core/chat/type'; import { useSize } from 'ahooks'; import React, { useEffect, useMemo, useRef } from 'react'; @@ -106,11 +107,8 @@ const ProcessingPreviewBody = React.memo(function ProcessingPreviewBody({ export const getProcessingPreviewLabelKey = (value: AIChatItemValueItemType) => { const tool = value.tools?.[value.tools.length - 1] || value.tool; if (tool) return tool.toolName; - if ( - (value.reasoning?.content || value.agentPlanUpdate?.reasoningText) && - !value.hideReason - ) { - return 'chat:history_generating'; + if ((value.reasoning?.content || value.agentPlanUpdate?.reasoningText) && !value.hideReason) { + return i18nT('chat:history_generating'); } return ''; diff --git a/projects/app/src/components/support/wallet/QRCodePayModal.tsx b/projects/app/src/components/support/wallet/QRCodePayModal.tsx index d7136ca19948..e66674c7240d 100644 --- a/projects/app/src/components/support/wallet/QRCodePayModal.tsx +++ b/projects/app/src/components/support/wallet/QRCodePayModal.tsx @@ -1,7 +1,7 @@ import MyModal from '@fastgpt/web/components/common/MyModal'; import React, { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation, Trans } from 'next-i18next'; -import { Box, ModalBody, Flex, Button, Text, Link } from '@chakra-ui/react'; +import { Box, ModalBody, Flex, Button, Link } from '@chakra-ui/react'; import { checkBalancePayResult, putUpdatePayment } from '@/web/support/wallet/bill/api'; import LightTip from '@fastgpt/web/components/common/LightTip'; import QRCode from 'qrcode'; @@ -16,6 +16,7 @@ import Markdown from '@/components/Markdown'; import MyIcon from '@fastgpt/web/components/common/Icon'; import { useToast } from '@fastgpt/web/hooks/useToast'; import type { CreateBillResponseType } from '@fastgpt/global/openapi/support/wallet/bill/api'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; export type QRPayProps = CreateBillResponseType & { billId: string; @@ -284,7 +285,7 @@ const QRCodePayModal = ({ {feConfigs.payFormUrl && ( }} diff --git a/projects/app/src/pageComponents/account/customDomain/createModal.tsx b/projects/app/src/pageComponents/account/customDomain/createModal.tsx index 44dd0cff45cf..3127bb0dad8a 100644 --- a/projects/app/src/pageComponents/account/customDomain/createModal.tsx +++ b/projects/app/src/pageComponents/account/customDomain/createModal.tsx @@ -36,6 +36,7 @@ import { createCustomDomain } from '@/web/support/customDomain/api'; import { getDocPath } from '@/web/common/system/doc'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const ProviderItem = ({ icon, @@ -92,16 +93,21 @@ function CreateCustomDomainModal({ const { feConfigs } = useSystemStore(); const { copyData } = useCopyData(); - const [provider, setProvider] = useState('tencent'); - const [domain, setDomain] = useState(''); + const [provider, setProvider] = useState(() => + type === 'refresh' ? data?.provider || 'tencent' : 'tencent' + ); + const [domain, setDomain] = useState(() => + type === 'refresh' ? data?.domain || '' : '' + ); const [editDomain, setEditDomain] = useState(true); - - useEffect(() => { - if (type === 'refresh') { - setProvider(data?.provider || 'tencent'); - setDomain(data?.domain || ''); - } - }, [data, type]); + const updateProvider = (provider: ProviderEnum) => { + setProvider(provider); + setDnsResolved(false); + }; + const updateDomain = (domain: string) => { + setDomain(domain); + setDnsResolved(false); + }; const cnameDomain = useMemo(() => { if (type === 'refresh') { @@ -151,12 +157,6 @@ function CreateCustomDomainModal({ return () => clearInterval(intervalId); }, [DnsResolved, checkDNSResolve, cnameDomain, domain, editDomain, startDnsResolve]); - useEffect(() => { - if (domain && provider) { - setDnsResolved(false); - } - }, [domain, provider]); - const loading = loadingCreatingDomain; return ( @@ -175,25 +175,25 @@ function CreateCustomDomainModal({ setProvider('tencent')} + onClick={() => updateProvider('tencent')} isDisabled={!editDomain || type === 'refresh'} /> setProvider('aliyun')} + onClick={() => updateProvider('aliyun')} isDisabled={!editDomain || type === 'refresh'} /> setProvider('volcengine')} + onClick={() => updateProvider('volcengine')} isDisabled={!editDomain || type === 'refresh'} /> }} /> @@ -204,7 +204,7 @@ function CreateCustomDomainModal({ h="40px" placeholder="www.example.com" value={domain} - onChange={(e) => setDomain(e.target.value)} + onChange={(e) => updateDomain(e.target.value)} isDisabled={!editDomain || type === 'refresh'} /> @@ -258,7 +258,7 @@ function CreateCustomDomainModal({ }} /> diff --git a/projects/app/src/pageComponents/account/usage/UsageTable.tsx b/projects/app/src/pageComponents/account/usage/UsageTable.tsx index 515cd4a9e7db..96c0ed114971 100644 --- a/projects/app/src/pageComponents/account/usage/UsageTable.tsx +++ b/projects/app/src/pageComponents/account/usage/UsageTable.tsx @@ -28,6 +28,7 @@ import PopoverConfirm from '@fastgpt/web/components/common/MyPopover/PopoverConf import { useRequest } from '@fastgpt/web/hooks/useRequest'; import { downloadFetch } from '@/web/common/system/utils'; import { useSafeTranslation } from '@fastgpt/web/hooks/useSafeTranslation'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const UsageDetail = dynamic(() => import('./UsageDetail')); const RechargeModal = dynamic(() => @@ -92,14 +93,14 @@ const UsageTableList = ({ ...requestParams, appNameMap: { ['core.app.Question Guide']: t('common:core.app.Question Guide'), - ['common:support.wallet.usage.Audio Speech']: t( + [i18nT('common:support.wallet.usage.Audio Speech')]: t( 'common:support.wallet.usage.Audio Speech' ), ['support.wallet.usage.Whisper']: t('common:support.wallet.usage.Whisper'), - ['account_usage:embedding_index']: t('account_usage:embedding_index'), - ['account_usage:qa']: t('account_usage:qa'), + [i18nT('account_usage:embedding_index')]: t('account_usage:embedding_index'), + [i18nT('account_usage:qa')]: t('account_usage:qa'), ['core.dataset.training.Auto mode']: t('common:core.dataset.training.Auto mode'), - ['common:core.module.template.ai_chat']: t('common:core.module.template.ai_chat') + [i18nT('common:core.module.template.ai_chat')]: t('common:core.module.template.ai_chat') }, sourcesMap: Object.fromEntries( Object.entries(UsageSourceMap).map(([key, config]) => [ diff --git a/projects/app/src/pageComponents/app/detail/Publish/Wecom/index.tsx b/projects/app/src/pageComponents/app/detail/Publish/Wecom/index.tsx index ff3a4f41776f..b7557b2e12ed 100644 --- a/projects/app/src/pageComponents/app/detail/Publish/Wecom/index.tsx +++ b/projects/app/src/pageComponents/app/detail/Publish/Wecom/index.tsx @@ -29,6 +29,7 @@ import EmptyTip from '@fastgpt/web/components/common/EmptyTip'; import { useRequest } from '@fastgpt/web/hooks/useRequest'; import { getDocPath } from '@/web/common/system/doc'; import { listCustomDomain } from '@/web/support/customDomain/api'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const WecomEditModal = dynamic(() => import('./WecomEditModal')); const ShowShareLinkModal = dynamic(() => import('../components/showShareLinkModal')); @@ -252,7 +253,7 @@ const Wecom = ({ appId }: { appId: string }) => { : { text: ( ) => { const { t } = useTranslation(); @@ -39,7 +40,10 @@ const NodeLoopRun = ({ data, selected }: NodeProps) => { WorkflowBufferDataContext, (v) => v ); - const childNodeIds = childrenNodeIdListMap[nodeId] ?? []; + const childNodeIds = useMemo( + () => childrenNodeIdListMap[nodeId] ?? [], + [childrenNodeIdListMap, nodeId] + ); const getRawNodeById = useContextSelector(WorkflowInitContext, (v) => v.getRawNodeById); const mode = @@ -111,8 +115,8 @@ const NodeLoopRun = ({ data, selected }: NodeProps) => { value: { id: NodeOutputKeyEnum.currentIndex, key: NodeOutputKeyEnum.currentIndex, - label: 'workflow:current_index', - description: 'workflow:current_index_desc', + label: i18nT('workflow:current_index'), + description: i18nT('workflow:current_index_desc'), type: FlowNodeOutputTypeEnum.static, valueType: WorkflowIOValueTypeEnum.number } @@ -125,8 +129,8 @@ const NodeLoopRun = ({ data, selected }: NodeProps) => { value: { id: NodeOutputKeyEnum.currentItem, key: NodeOutputKeyEnum.currentItem, - label: 'workflow:current_item', - description: 'workflow:current_item_desc', + label: i18nT('workflow:current_item'), + description: i18nT('workflow:current_item_desc'), type: FlowNodeOutputTypeEnum.static, valueType: WorkflowIOValueTypeEnum.any } @@ -154,8 +158,8 @@ const NodeLoopRun = ({ data, selected }: NodeProps) => { value: { id: NodeOutputKeyEnum.currentIteration, key: NodeOutputKeyEnum.currentIteration, - label: 'workflow:current_iteration', - description: 'workflow:current_iteration_desc', + label: i18nT('workflow:current_iteration'), + description: i18nT('workflow:current_iteration_desc'), type: FlowNodeOutputTypeEnum.static, valueType: WorkflowIOValueTypeEnum.number } diff --git a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodeFormInput/InputFormEditModal.tsx b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodeFormInput/InputFormEditModal.tsx index beae8c8a2c4d..7094486614fc 100644 --- a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodeFormInput/InputFormEditModal.tsx +++ b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodeFormInput/InputFormEditModal.tsx @@ -123,6 +123,13 @@ const InputFormEditModal = ({ onTypeChange={(type) => { setValue('type', type as FlowNodeInputTypeEnum); setValue('defaultValue', ''); + if ( + (type === FlowNodeInputTypeEnum.select || + type === FlowNodeInputTypeEnum.multipleSelect) && + !form.getValues('list')?.length + ) { + setValue('list', [{ label: '', value: '' }]); + } }} /> diff --git a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputEditModal.tsx b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputEditModal.tsx index b0a9f3667304..9aa637a6801e 100644 --- a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputEditModal.tsx +++ b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputEditModal.tsx @@ -192,6 +192,13 @@ const FieldEditModal = ({ if (targetItem) { setValue('renderTypeList', targetItem.value); setValue('defaultValue', ''); + if ( + (type === FlowNodeInputTypeEnum.select || + type === FlowNodeInputTypeEnum.multipleSelect) && + !form.getValues('list')?.length + ) { + setValue('list', [{ label: '', value: '' }]); + } } }} /> diff --git a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputTypeConfig.tsx b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputTypeConfig.tsx index 6968d1b79cfd..4101e3f5afe0 100644 --- a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputTypeConfig.tsx +++ b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/NodePluginIO/InputTypeConfig.tsx @@ -159,9 +159,21 @@ const InputTypeConfig = ({ name: 'list' }); - const mergedSelectEnums = selectEnums.map((field, index) => ({ - ...field, - ...listValue[index] + const isSelectInput = + inputType === FlowNodeInputTypeEnum.select || inputType === VariableInputEnum.select; + const isMultipleSelectInput = + inputType === FlowNodeInputTypeEnum.multipleSelect || + inputType === VariableInputEnum.multipleSelect; + const isOptionInput = isSelectInput || isMultipleSelectInput; + const optionFields = (listValue.length ? listValue : selectEnums) as { + id?: string; + label?: string; + value?: string; + }[]; + const optionDragList = optionFields.map((item, index) => ({ + id: item.id || `${index}`, + label: item.label || '', + value: item.value || item.label || '' })); const handleRemoveEnum = useCallback( @@ -171,7 +183,7 @@ const InputTypeConfig = ({ if (!removedValue) return; - if (inputType === FlowNodeInputTypeEnum.multipleSelect) { + if (isMultipleSelectInput) { const cur = getValues('defaultValue'); if (Array.isArray(cur) && cur.includes(removedValue)) { setValue( @@ -179,16 +191,16 @@ const InputTypeConfig = ({ cur.filter((v: string) => v !== removedValue) ); } - } else if (inputType === FlowNodeInputTypeEnum.select) { + } else if (isSelectInput) { if (getValues('defaultValue') === removedValue) { setValue('defaultValue', ''); } } }, - [removeEnums, inputType, getValues, setValue] + [removeEnums, isMultipleSelectInput, isSelectInput, getValues, setValue] ); - const isLastEnumEmpty = !mergedSelectEnums[mergedSelectEnums.length - 1]?.label; + const isOptionLimitReached = optionFields.length >= 50; const valueTypeSelectList = Object.values(FlowValueTypeMap) .filter((item) => !item.abandon) @@ -255,8 +267,8 @@ const InputTypeConfig = ({ [VariableInputEnum.datasetSelect]: true }; - return map[inputType as keyof typeof map]; - }, [inputType]); + return isOptionInput || map[inputType as keyof typeof map]; + }, [inputType, isOptionInput]); const showIsToolInput = useMemo(() => { const list = [ @@ -287,6 +299,21 @@ const InputTypeConfig = ({ defaultValue: data.defaultValue }; + if (isOptionInput) { + const cleanList = (data.list ?? []).filter( + (item: { label?: string; value?: string }) => !!item?.label + ); + commonData.list = cleanList; + const validValues = new Set(cleanList.map((item: { value: string }) => item.value)); + if (isMultipleSelectInput) { + commonData.defaultValue = Array.isArray(commonData.defaultValue) + ? commonData.defaultValue.filter((v: string) => validValues.has(v)) + : commonData.defaultValue; + } else if (commonData.defaultValue && !validValues.has(commonData.defaultValue)) { + commonData.defaultValue = ''; + } + } + switch (inputType) { case FlowNodeInputTypeEnum.input: case FlowNodeInputTypeEnum.textarea: @@ -296,22 +323,6 @@ const InputTypeConfig = ({ commonData.max = data.max; commonData.min = data.min; break; - case FlowNodeInputTypeEnum.select: - case FlowNodeInputTypeEnum.multipleSelect: { - const cleanList = (data.list ?? []).filter( - (item: { label?: string; value?: string }) => !!item?.label - ); - commonData.list = cleanList; - const validValues = new Set(cleanList.map((item: { value: string }) => item.value)); - if (inputType === FlowNodeInputTypeEnum.multipleSelect) { - commonData.defaultValue = Array.isArray(commonData.defaultValue) - ? commonData.defaultValue.filter((v: string) => validValues.has(v)) - : commonData.defaultValue; - } else if (commonData.defaultValue && !validValues.has(commonData.defaultValue)) { - commonData.defaultValue = ''; - } - break; - } case FlowNodeInputTypeEnum.addInputParam: commonData.customInputConfig = data.customInputConfig; break; @@ -372,7 +383,7 @@ const InputTypeConfig = ({ return commonData; }, - [inputType] + [inputType, isMultipleSelectInput, isOptionInput] ); return ( @@ -615,7 +626,7 @@ const InputTypeConfig = ({ )} - {inputType === FlowNodeInputTypeEnum.select && ( + {isSelectInput && ( list={[defaultListValue, ...listValue] .filter((item) => item.label !== '') @@ -634,7 +645,7 @@ const InputTypeConfig = ({ w={'200px'} /> )} - {inputType === FlowNodeInputTypeEnum.multipleSelect && ( + {isMultipleSelectInput && ( flex={'1 0 0'} size={'md'} @@ -799,18 +810,13 @@ const InputTypeConfig = ({ )} - {(inputType === FlowNodeInputTypeEnum.select || - inputType == FlowNodeInputTypeEnum.multipleSelect) && ( + {isOptionInput && ( <> - + onDragEndCb={(list) => { - const newOrder = list.map((item) => item.id); - const newSelectEnums = newOrder - .map((id) => mergedSelectEnums.find((item) => item.id === id)) - .filter(Boolean) as { id: string; value: string }[]; removeEnums(); - newSelectEnums.forEach((item) => - appendEnums({ label: item.value, value: item.value }) + list.forEach((item) => + appendEnums({ label: item.label || item.value, value: item.value || item.label }) ); // 防止最后一个元素被focus @@ -820,7 +826,7 @@ const InputTypeConfig = ({ } }, 0); }} - dataList={mergedSelectEnums} + dataList={optionDragList} renderClone={(provided, snapshot, rubric) => { return ( - {mergedSelectEnums[rubric.source.index].value} + {optionDragList[rubric.source.index]?.value} ); }} @@ -846,7 +852,7 @@ const InputTypeConfig = ({ flexDirection={'column'} gap={4} > - {mergedSelectEnums.map((item, i) => ( + {optionFields.map((item, i) => ( {(provided, snapshot) => ( } onClick={() => { - if (isLastEnumEmpty) return; + if (isOptionLimitReached) return; appendEnums({ label: '', value: '' }); }} - isDisabled={isLastEnumEmpty} + isDisabled={isOptionLimitReached} fontWeight={'medium'} fontSize={'12px'} w={'24'} diff --git a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/render/NodeCard.tsx b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/render/NodeCard.tsx index 8912af522c43..763901463ac1 100644 --- a/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/render/NodeCard.tsx +++ b/projects/app/src/pageComponents/app/detail/WorkflowComponents/Flow/nodes/render/NodeCard.tsx @@ -560,6 +560,7 @@ const NodeIntro = React.memo(function NodeIntro({ nodeId: string; intro?: string; }) { + const { t } = useTranslation(); const onChangeNode = useContextSelector(WorkflowActionsContext, (v) => v.onChangeNode); const handleSave = useCallback( @@ -585,7 +586,7 @@ const NodeIntro = React.memo(function NodeIntro({ onSave={handleSave} type={'textarea'} maxLength={500} - placeholder={'app:node_not_intro'} + placeholder={t('app:node_not_intro')} fontSize={'sm'} lineHeight={'short'} color={'myGray.500'} diff --git a/projects/app/src/pageComponents/chat/SandboxEditor/components/FileTree.tsx b/projects/app/src/pageComponents/chat/SandboxEditor/components/FileTree.tsx index e4ca39bcbe0d..3e67606a84bb 100644 --- a/projects/app/src/pageComponents/chat/SandboxEditor/components/FileTree.tsx +++ b/projects/app/src/pageComponents/chat/SandboxEditor/components/FileTree.tsx @@ -15,6 +15,7 @@ import MyTooltip from '@fastgpt/web/components/common/MyTooltip'; import { Trans, useTranslation } from 'next-i18next'; import { useToast } from '@fastgpt/web/hooks/useToast'; import { getErrText } from '@fastgpt/global/common/error/utils'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; import { DndContext, useSensor, @@ -153,7 +154,6 @@ const DroppableRootBox = ({ }; const FileTree = ({ - width = 250, filteredTree, searchQuery, setSearchQuery, @@ -525,7 +525,7 @@ const FileTree = ({ title: t('chat:sandbox_confirm_delete_title'), customContent: ( diff --git a/projects/app/src/pageComponents/dashboard/agent/context.tsx b/projects/app/src/pageComponents/dashboard/agent/context.tsx index 5aa65bc90ebf..8f2b24362276 100644 --- a/projects/app/src/pageComponents/dashboard/agent/context.tsx +++ b/projects/app/src/pageComponents/dashboard/agent/context.tsx @@ -15,6 +15,7 @@ import dynamic from 'next/dynamic'; import { AppTypeEnum } from '@fastgpt/global/core/app/constants'; import { useSystemStore } from '@/web/common/system/useSystemStore'; import { useTranslation } from 'next-i18next'; +import { resolveDashboardAppListTypes } from './utils/appListTypes'; const MoveModal = dynamic(() => import('@/components/common/folder/MoveModal')); type AppListContextType = { @@ -72,35 +73,10 @@ const AppListContextProvider = ({ children }: { children: ReactNode }) => { loading: isFetchingApps } = useRequest( () => { - const formatType = (() => { - // chat page show all apps - if (router.pathname.includes('/chat')) { - return [ - AppTypeEnum.folder, - AppTypeEnum.toolFolder, - AppTypeEnum.simple, - AppTypeEnum.workflow, - AppTypeEnum.workflowTool - ]; - } - - // agent page - if (router.pathname.includes('/agent')) { - return !type || type === 'all' - ? [AppTypeEnum.folder, AppTypeEnum.simple, AppTypeEnum.workflow, AppTypeEnum.chatAgent] - : [AppTypeEnum.folder, type]; - } - - // tool page - return !type || type === 'all' - ? [ - AppTypeEnum.toolFolder, - AppTypeEnum.workflowTool, - AppTypeEnum.mcpToolSet, - AppTypeEnum.httpToolSet - ] - : [AppTypeEnum.toolFolder, type]; - })(); + const formatType = resolveDashboardAppListTypes({ + pathname: router.pathname, + type + }); return getMyApps({ parentId, type: formatType, searchKey }); }, diff --git a/projects/app/src/pageComponents/dashboard/agent/utils/appListTypes.ts b/projects/app/src/pageComponents/dashboard/agent/utils/appListTypes.ts new file mode 100644 index 000000000000..2970f82094ae --- /dev/null +++ b/projects/app/src/pageComponents/dashboard/agent/utils/appListTypes.ts @@ -0,0 +1,59 @@ +import { AppTypeEnum } from '@fastgpt/global/core/app/constants'; + +type ResolveDashboardAppListTypesParams = { + pathname: string; + type?: AppTypeEnum | 'all'; +}; + +const allAgentAppTypes = [ + AppTypeEnum.folder, + AppTypeEnum.simple, + AppTypeEnum.workflow, + AppTypeEnum.chatAgent +]; + +const allToolAppTypes = [ + AppTypeEnum.toolFolder, + AppTypeEnum.workflowTool, + AppTypeEnum.mcpToolSet, + AppTypeEnum.httpToolSet, + AppTypeEnum.httpPlugin +]; + +/** + * 根据 dashboard 当前页面和类型筛选，生成 `/core/app/list` 的类型过滤条件。 + * + * `httpPlugin` 是旧版 HTTP 工具类型，工具页的“全部”和新版 HTTP 工具筛选都要带上它， + * 否则历史团队升级后仍存在的旧版 HTTP 工具不会出现在 dashboard/tool 列表里。 + */ +export const resolveDashboardAppListTypes = ({ + pathname, + type +}: ResolveDashboardAppListTypesParams): AppTypeEnum[] => { + // 聊天页只展示可直接对话/运行的应用和工具。 + if (pathname.includes('/chat')) { + return [ + AppTypeEnum.folder, + AppTypeEnum.toolFolder, + AppTypeEnum.simple, + AppTypeEnum.workflow, + AppTypeEnum.workflowTool + ]; + } + + // Agent 页保留原有 Agent 类型筛选行为。 + if (pathname.includes('/agent')) { + return !type || type === 'all' ? allAgentAppTypes : [AppTypeEnum.folder, type]; + } + + // 工具页需要兼容旧版工具类型。 + if (!type || type === 'all') { + return allToolAppTypes; + } + + if (type === AppTypeEnum.httpToolSet) { + return [AppTypeEnum.toolFolder, AppTypeEnum.httpToolSet, AppTypeEnum.httpPlugin]; + } + + return [AppTypeEnum.toolFolder, type]; +}; diff --git a/projects/app/src/pageComponents/dashboard/skill/List.tsx b/projects/app/src/pageComponents/dashboard/skill/List.tsx index 15cbf3708248..19e5c829c5b2 100644 --- a/projects/app/src/pageComponents/dashboard/skill/List.tsx +++ b/projects/app/src/pageComponents/dashboard/skill/List.tsx @@ -48,6 +48,7 @@ import ListCreateCard from '@/pageComponents/dashboard/ListCreateCard'; import { useVirtualGridList } from '@fastgpt/web/hooks/useVirtualGridList'; import { getWebReqUrl } from '@fastgpt/web/common/system/utils'; import MyTooltip from '@fastgpt/web/components/common/MyTooltip'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const EditResourceModal = dynamic(() => import('@/components/common/Modal/EditResourceModal')); const MoveModal = dynamic(() => import('@/components/common/folder/MoveModal')); @@ -367,7 +368,7 @@ const List = ({ openConfirmDelete({ customContent: ( }} /> diff --git a/projects/app/src/pageComponents/dashboard/skill/detail/Header.tsx b/projects/app/src/pageComponents/dashboard/skill/detail/Header.tsx index a5175f0535f8..83bc6d6c5037 100644 --- a/projects/app/src/pageComponents/dashboard/skill/detail/Header.tsx +++ b/projects/app/src/pageComponents/dashboard/skill/detail/Header.tsx @@ -31,6 +31,7 @@ import dynamic from 'next/dynamic'; import type { EditResourceInfoFormType } from '@/components/common/Modal/EditResourceModal'; import { useConfirm } from '@fastgpt/web/hooks/useConfirm'; import MyTooltip from '@fastgpt/web/components/common/MyTooltip'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const EditResourceModal = dynamic(() => import('@/components/common/Modal/EditResourceModal')); const ConfigPerModal = dynamic(() => import('@/components/support/permission/ConfigPerModal')); @@ -225,7 +226,7 @@ export const LeftHeader = () => { openConfirmDelete({ customContent: ( }} /> diff --git a/projects/app/src/pageComponents/dataset/detail/CollectionCard/TemplateImportModal.tsx b/projects/app/src/pageComponents/dataset/detail/CollectionCard/TemplateImportModal.tsx index 358411fb7ef5..7e5d146cd6ce 100644 --- a/projects/app/src/pageComponents/dataset/detail/CollectionCard/TemplateImportModal.tsx +++ b/projects/app/src/pageComponents/dataset/detail/CollectionCard/TemplateImportModal.tsx @@ -11,6 +11,7 @@ import { DatasetPageContext } from '@/web/core/dataset/context/datasetPageContex import { useContextSelector } from 'use-context-selector'; import { getDocPath } from '@/web/common/system/doc'; import { Trans } from 'next-i18next'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const TemplateImportModal = ({ onFinish, @@ -104,7 +105,7 @@ const TemplateImportModal = ({ FileTypeNode={ { const { feConfigs } = useSystemStore(); @@ -50,7 +51,7 @@ const PolicyTip = () => { whiteSpace={'pre-wrap'} > : , termsLink: ( @@ -79,7 +80,7 @@ const PolicyTip = () => { aria-hidden > , termsLink: , diff --git a/projects/app/src/pages/account/info/index.tsx b/projects/app/src/pages/account/info/index.tsx index 76d5b1130e25..6c9f8459f5d8 100644 --- a/projects/app/src/pages/account/info/index.tsx +++ b/projects/app/src/pages/account/info/index.tsx @@ -46,6 +46,7 @@ import MyDivider from '@fastgpt/web/components/common/MyDivider'; import { useUploadAvatar } from '@fastgpt/web/common/file/hooks/useUploadAvatar'; import { getUploadAvatarPresignedUrl } from '@/web/common/file/api'; import { TeamErrEnum } from '@fastgpt/global/common/error/code/team'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; const RedeemCouponModal = dynamic(() => import('@/pageComponents/account/info/RedeemCouponModal'), { ssr: false @@ -384,7 +385,7 @@ const PlanUsage = () => { const planName = useMemo(() => { if (!teamPlanStatus?.standard?.currentSubLevel) return ''; if (isWecomTeam && teamPlanStatus.standard.currentSubLevel === StandardSubLevelEnum.free) - return 'common:support.wallet.subscription.standardSubLevel.trial'; + return i18nT('common:support.wallet.subscription.standardSubLevel.trial'); return ( subPlans?.standard?.[teamPlanStatus.standard.currentSubLevel]?.name || diff --git a/projects/app/src/pages/api/core/chat/record/getCollectionQuote.ts b/projects/app/src/pages/api/core/chat/record/getCollectionQuote.ts index a8cc4190dda0..b9edcc351d4e 100644 --- a/projects/app/src/pages/api/core/chat/record/getCollectionQuote.ts +++ b/projects/app/src/pages/api/core/chat/record/getCollectionQuote.ts @@ -108,6 +108,7 @@ async function handleInitialLoad({ }): Promise { const centerNode = await MongoDatasetData.findOne( { + ...baseMatch, _id: new Types.ObjectId(initialId) }, quoteDataFieldSelector diff --git a/projects/app/src/pages/dashboard/evaluation/create.tsx b/projects/app/src/pages/dashboard/evaluation/create.tsx index 7b6877424dfd..d70e006243f6 100644 --- a/projects/app/src/pages/dashboard/evaluation/create.tsx +++ b/projects/app/src/pages/dashboard/evaluation/create.tsx @@ -5,7 +5,7 @@ import { Box, Button, Flex, Input, VStack } from '@chakra-ui/react'; import { useRouter } from 'next/router'; import { serviceSideProps } from '@/web/common/i18n/utils'; import AIModelSelector from '@/components/Select/AIModelSelector'; -import { useForm } from 'react-hook-form'; +import { useForm, useWatch } from 'react-hook-form'; import { useSystemStore } from '@/web/common/system/useSystemStore'; import FormLabel from '@fastgpt/web/components/common/MyBox/FormLabel'; import AppSelect from '@/components/Select/AppSelect'; @@ -19,12 +19,13 @@ import { useToast } from '@fastgpt/web/hooks/useToast'; import QuestionTip from '@fastgpt/web/components/common/MyTooltip/QuestionTip'; import { fileDownload } from '@/web/common/file/utils'; import { postCreateEvaluation } from '@/web/core/app/api/evaluation'; -import { useMemo, useState } from 'react'; +import { useState } from 'react'; import Markdown from '@/components/Markdown'; import { getEvaluationFileHeader } from '@fastgpt/global/core/app/evaluation/utils'; import { evaluationFileErrors } from '@fastgpt/global/core/app/evaluation/constants'; import { TeamErrEnum } from '@fastgpt/global/common/error/code/team'; import { getErrText } from '@fastgpt/global/common/error/utils'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; type EvaluationFormType = { name: string; @@ -43,7 +44,7 @@ const EvaluationCreating = () => { const { llmModelList } = useSystemStore(); - const { register, setValue, watch, handleSubmit } = useForm({ + const { register, setValue, control, handleSubmit } = useForm({ defaultValues: { name: '', evalModel: llmModelList[0]?.model, @@ -52,10 +53,10 @@ const EvaluationCreating = () => { } }); - const name = watch('name'); - const evalModel = watch('evalModel'); - const appId = watch('appId'); - const evaluationFiles = watch('evaluationFiles'); + const name = useWatch({ control, name: 'name' }); + const evalModel = useWatch({ control, name: 'evalModel' }); + const appId = useWatch({ control, name: 'appId' }); + const evaluationFiles = useWatch({ control, name: 'evaluationFiles' }); const { runAsync: getAppDetail, loading: isLoadingAppDetail } = useRequest(() => { if (appId) return getAppDetailById(appId); @@ -254,7 +255,7 @@ const EvaluationCreating = () => { FileTypeNode={ { { id: 'userChatInput', key: 'userChatInput', - label: 'common:core.module.input.label.user question', + label: i18nT('common:core.module.input.label.user question'), type: FlowNodeOutputTypeEnum.static, valueType: WorkflowIOValueTypeEnum.string } @@ -656,7 +656,7 @@ export const parsePluginFromCurlString = ( label: '', description: '自定义请求头，请严格填入 JSON 字符串。\n1. 确保最后一个属性没有逗号\n2. 确保 key 包含双引号\n例如：{"Authorization":"Bearer xxx"}', - placeholder: 'common:core.module.input.description.Http Request Header', + placeholder: i18nT('common:core.module.input.description.Http Request Header'), required: false, valueDesc: '', debugLabel: '', diff --git a/projects/app/src/web/core/workflow/utils.ts b/projects/app/src/web/core/workflow/utils.ts index 408b63bd0897..539a65a38c30 100644 --- a/projects/app/src/web/core/workflow/utils.ts +++ b/projects/app/src/web/core/workflow/utils.ts @@ -3,7 +3,6 @@ import type { FlowNodeTemplateType } from '@fastgpt/global/core/workflow/type/no import type { Edge, Node, XYPosition } from 'reactflow'; import { moduleTemplatesFlat } from '@fastgpt/global/core/workflow/template/constants'; import { - AppNodeFlowNodeTypeMap, EDGE_TYPE, FlowNodeInputTypeEnum, FlowNodeOutputTypeEnum, @@ -11,6 +10,7 @@ import { } from '@fastgpt/global/core/workflow/node/constant'; import { EmptyNode } from '@fastgpt/global/core/workflow/template/system/emptyNode'; import { type StoreEdgeItemType } from '@fastgpt/global/core/workflow/type/edge'; +import { i18nT } from '@fastgpt/global/common/i18n/utils'; import { getNanoid } from '@fastgpt/global/common/string/tools'; import { getGlobalVariableNode } from './adapt'; import { VARIABLE_NODE_ID, WorkflowIOValueTypeEnum } from '@fastgpt/global/core/workflow/constants'; @@ -61,7 +61,7 @@ export const adaptStoreNodeInputs = (storeNode: StoreNodeItemType): FlowNodeInpu return { ...input, key: NodeInputKeyEnum.datasetSearchInput, - label: 'workflow:search_query', + label: i18nT('workflow:search_query'), value: isReferenceValue ? [input.value] : input.value, valueType: WorkflowIOValueTypeEnum.arrayString, selectedTypeIndex: isReferenceValue ? 0 : 1 @@ -879,8 +879,6 @@ export const checkWorkflowNodeAndConnection = ({ continue; } - const hasIncoming = (incoming.get(nodeId) || []).length > 0; - const hasOutgoing = (outgoing.get(nodeId) || []).length > 0; const isStartNode = [ FlowNodeTypeEnum.workflowStart, FlowNodeTypeEnum.pluginInput, diff --git a/projects/app/test/api/core/chat/record/getCollectionQuote.test.ts b/projects/app/test/api/core/chat/record/getCollectionQuote.test.ts new file mode 100644 index 000000000000..e31f363f6ddc --- /dev/null +++ b/projects/app/test/api/core/chat/record/getCollectionQuote.test.ts @@ -0,0 +1,183 @@ +import type { ApiRequestProps } from '@fastgpt/service/type/next'; +import { Types } from 'mongoose'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + authChatCrud: vi.fn(), + authCollectionInChat: vi.fn(), + getCollectionWithDataset: vi.fn(), + findChatItem: vi.fn(), + findDatasetDataByInitialId: vi.fn(), + findDatasetDataList: vi.fn() +})); + +vi.mock('@/service/middleware/entry', () => ({ + NextAPI: (handler: unknown) => handler +})); + +vi.mock('@/service/support/permission/auth/chat', () => ({ + authChatCrud: mocks.authChatCrud, + authCollectionInChat: mocks.authCollectionInChat +})); + +vi.mock('@fastgpt/service/core/dataset/controller', () => ({ + getCollectionWithDataset: mocks.getCollectionWithDataset +})); + +vi.mock('@fastgpt/service/core/chat/chatItemSchema', () => ({ + MongoChatItem: { + findOne: mocks.findChatItem + } +})); + +vi.mock('@fastgpt/service/core/dataset/data/schema', () => ({ + MongoDatasetData: { + findOne: mocks.findDatasetDataByInitialId, + find: mocks.findDatasetDataList + } +})); + +import handler from '@/pages/api/core/chat/record/getCollectionQuote'; + +const appId = '68ad85a7463006c963799a05'; +const chatId = 'chat_quote_auth_regression'; +const chatItemDataId = 'quote_ai_response'; +const collectionId = '68ad85a7463006c963799a06'; +const teamId = '68ad85a7463006c963799a07'; +const datasetId = '68ad85a7463006c963799a08'; +const initialId = '68ad85a7463006c963799a09'; +const chatTime = new Date('2026-06-06T10:14:40.000Z'); + +const callHandler = () => + handler({ + body: { + appId, + chatId, + chatItemDataId, + collectionId, + initialId, + anchor: 0, + pageSize: 5 + } + } as ApiRequestProps); + +const createDatasetData = ({ + id, + q, + dataTeamId = teamId, + dataDatasetId = datasetId, + dataCollectionId = collectionId, + chunkIndex = 0 +}: { + id: string; + q: string; + dataTeamId?: string; + dataDatasetId?: string; + dataCollectionId?: string; + chunkIndex?: number; +}) => ({ + _id: new Types.ObjectId(id), + teamId: dataTeamId, + datasetId: dataDatasetId, + collectionId: dataCollectionId, + q, + a: '', + history: [], + updateTime: new Date('2026-06-06T10:00:00.000Z'), + chunkIndex +}); + +const mockFindDatasetDataList = (list: unknown[]) => { + const query = { + sort: vi.fn().mockReturnThis(), + limit: vi.fn().mockReturnThis(), + lean: vi.fn().mockResolvedValue(list) + }; + + mocks.findDatasetDataList.mockReturnValue(query); + return query; +}; + +describe('getCollectionQuote handler', () => { + beforeEach(() => { + vi.clearAllMocks(); + + mocks.authChatCrud.mockResolvedValue({ chat: { _id: chatId }, showFullText: true }); + mocks.authCollectionInChat.mockResolvedValue(undefined); + mocks.getCollectionWithDataset.mockResolvedValue({ + _id: collectionId, + teamId, + datasetId + }); + mocks.findChatItem.mockReturnValue({ + lean: () => Promise.resolve({ time: chatTime }) + }); + mocks.findDatasetDataByInitialId.mockReturnValue({ + lean: () => Promise.resolve(null) + }); + mockFindDatasetDataList([]); + }); + + it('binds initialId center-node lookup to the authorized team, dataset and collection', async () => { + mocks.findDatasetDataByInitialId.mockReturnValue({ + lean: () => + Promise.resolve( + createDatasetData({ + id: initialId, + q: 'authorized quote text' + }) + ) + }); + + await expect(callHandler()).resolves.toMatchObject({ + list: [expect.objectContaining({ q: 'authorized quote text' })], + hasMorePrev: false, + hasMoreNext: false + }); + + expect(mocks.authCollectionInChat).toHaveBeenCalledWith({ + appId, + chatId, + collectionIds: [collectionId] + }); + expect(mocks.findDatasetDataByInitialId).toHaveBeenCalledWith( + expect.objectContaining({ + teamId, + datasetId, + collectionId, + _id: new Types.ObjectId(initialId) + }), + expect.any(String) + ); + }); + + it('does not return a foreign initialId when it is outside the authorized baseMatch', async () => { + const authorizedFallback = createDatasetData({ + id: '68ad85a7463006c963799a10', + q: 'tenant B authorized quote' + }); + mockFindDatasetDataList([authorizedFallback]); + + const result = await callHandler(); + + expect(result.list.map((item) => item.q)).toEqual(['tenant B authorized quote']); + expect(result.list.map((item) => item.q)).not.toContain('tenant A secret quote'); + expect(mocks.findDatasetDataByInitialId).toHaveBeenCalledWith( + expect.objectContaining({ + teamId, + datasetId, + collectionId, + _id: new Types.ObjectId(initialId) + }), + expect.any(String) + ); + expect(mocks.findDatasetDataList).toHaveBeenCalledWith( + expect.objectContaining({ + teamId, + datasetId, + collectionId + }), + expect.any(String) + ); + }); +}); diff --git a/projects/app/test/pageComponents/dashboard/agent/utils/appListTypes.test.ts b/projects/app/test/pageComponents/dashboard/agent/utils/appListTypes.test.ts new file mode 100644 index 000000000000..d829ce85f056 --- /dev/null +++ b/projects/app/test/pageComponents/dashboard/agent/utils/appListTypes.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest'; +import { AppTypeEnum } from '@fastgpt/global/core/app/constants'; +import { resolveDashboardAppListTypes } from '@/pageComponents/dashboard/agent/utils/appListTypes'; + +describe('resolveDashboardAppListTypes', () => { + it('should include legacy HTTP plugin apps in the dashboard tool all list', () => { + expect( + resolveDashboardAppListTypes({ + pathname: '/dashboard/tool', + type: 'all' + }) + ).toEqual([ + AppTypeEnum.toolFolder, + AppTypeEnum.workflowTool, + AppTypeEnum.mcpToolSet, + AppTypeEnum.httpToolSet, + AppTypeEnum.httpPlugin + ]); + }); + + it('should include legacy HTTP plugin apps when filtering dashboard tools by HTTP toolset', () => { + expect( + resolveDashboardAppListTypes({ + pathname: '/dashboard/tool', + type: AppTypeEnum.httpToolSet + }) + ).toEqual([AppTypeEnum.toolFolder, AppTypeEnum.httpToolSet, AppTypeEnum.httpPlugin]); + }); + + it('should keep agent page type filters unchanged', () => { + expect( + resolveDashboardAppListTypes({ + pathname: '/dashboard/agent', + type: AppTypeEnum.workflow + }) + ).toEqual([AppTypeEnum.folder, AppTypeEnum.workflow]); + }); +}); diff --git a/projects/code-sandbox/Dockerfile b/projects/code-sandbox/Dockerfile index e412efea87ca..28ca0b4e3b7c 100644 --- a/projects/code-sandbox/Dockerfile +++ b/projects/code-sandbox/Dockerfile @@ -1,9 +1,14 @@ # --------- Build Stage ----------- -FROM node:24-alpine AS builder +FROM golang:1.22-bookworm AS go-runtime + +FROM node:24-bookworm-slim AS builder WORKDIR /app ARG proxy +COPY --from=go-runtime /usr/local/go /usr/local/go +ENV PATH="/usr/local/go/bin:${PATH}" + # 安装 pnpm RUN npm install -g pnpm@10.33.2 @@ -14,8 +19,11 @@ COPY packages/service ./packages/service COPY sdk ./sdk COPY projects/code-sandbox/ ./projects/code-sandbox/ -RUN [ -z "$proxy" ] || sed -i 's/dl-cdn.alpinelinux.org/mirrors.ustc.edu.cn/g' /etc/apk/repositories -RUN apk add --no-cache curl ca-certificates && update-ca-certificates +RUN [ -z "$proxy" ] || sed -i 's/deb.debian.org/mirrors.ustc.edu.cn/g' /etc/apt/sources.list.d/debian.sources +RUN apt-get -o Acquire::Retries=5 -o Acquire::http::Timeout=30 update && \ + apt-get -o Acquire::Retries=5 -o Acquire::http::Timeout=30 install -y --no-install-recommends \ + curl ca-certificates libseccomp-dev pkg-config gcc libc6-dev && \ + rm -rf /var/lib/apt/lists/* # 安装所有依赖（包括 devDependencies 用于编译） RUN if [ -z "$proxy" ]; then \ @@ -27,16 +35,16 @@ RUN if [ -z "$proxy" ]; then \ # 先构建 SDK workspace 包，确保 dist 入口可被打包工具解析 RUN pnpm --filter @fastgpt-sdk/otel --filter @fastgpt-sdk/storage build -# 编译主入口文件 -RUN cd /app/projects/code-sandbox && pnpm build +# 编译主入口文件和 Python native 隔离库 +RUN cd /app/projects/code-sandbox && SANDBOX_BUILD_NATIVE_PYTHON=true pnpm build # ===== Runner Stage ===== -FROM node:24-alpine AS runner +FROM node:24-bookworm-slim AS runner WORKDIR /app/code-sandbox ARG proxy -RUN [ -z "$proxy" ] || sed -i 's/dl-cdn.alpinelinux.org/mirrors.ustc.edu.cn/g' /etc/apk/repositories +RUN [ -z "$proxy" ] || sed -i 's/deb.debian.org/mirrors.ustc.edu.cn/g' /etc/apt/sources.list.d/debian.sources # 安装 pnpm（用于 runner 阶段安装 prod 依赖） RUN npm install -g pnpm@10.33.2 @@ -58,18 +66,60 @@ RUN if [ -z "$proxy" ]; then \ pnpm store prune || true # 安装 Python、依赖包及工具 -RUN apk add --no-cache python3 py3-pip libffi util-linux && \ - apk add --no-cache --virtual .build-deps gcc g++ musl-dev python3-dev libffi-dev +RUN apt-get -o Acquire::Retries=5 -o Acquire::http::Timeout=30 update && \ + apt-get -o Acquire::Retries=5 -o Acquire::http::Timeout=30 install -y --no-install-recommends \ + python3 python3-pip libffi8 util-linux libseccomp2 \ + gcc g++ python3-dev libffi-dev patch git && \ + rm -rf /var/lib/apt/lists/* COPY projects/code-sandbox/requirements.txt /tmp/requirements.txt RUN pip3 install --no-cache-dir --break-system-packages -r /tmp/requirements.txt && \ - rm /tmp/requirements.txt && \ - apk del .build-deps - - -# 创建非 root 用户运行沙箱 -RUN addgroup -S sandbox && adduser -S sandbox -G sandbox && \ - chown -R sandbox:sandbox /app -USER sandbox + rm /tmp/requirements.txt + +RUN set -eux; \ + root=/tmp/fastgpt-python-sandbox; \ + mkdir -p \ + "$root/app/code-sandbox" \ + "$root/bin" \ + "$root/usr/bin" \ + "$root/usr/lib" \ + "$root/usr/local/lib" \ + "$root/lib" \ + "$root/lib64" \ + "$root/etc" \ + "$root/dev" \ + "$root/tmp"; \ + chmod 755 "$root/tmp"; \ + cp -a /app/code-sandbox/. "$root/app/code-sandbox/"; \ + cp -a /usr/bin/python3 /usr/bin/python3.* "$root/usr/bin/" 2>/dev/null || true; \ + cp -a /usr/lib/python3* "$root/usr/lib/"; \ + cp -a /usr/local/lib/python3* "$root/usr/local/lib/" 2>/dev/null || true; \ + cp -a /usr/lib/aarch64-linux-gnu "$root/usr/lib/" 2>/dev/null || true; \ + cp -a /lib/aarch64-linux-gnu "$root/lib/" 2>/dev/null || true; \ + cp -a /lib/ld-linux-aarch64.so.1 "$root/lib/" 2>/dev/null || true; \ + cp -a /usr/lib/x86_64-linux-gnu "$root/usr/lib/" 2>/dev/null || true; \ + cp -a /lib/x86_64-linux-gnu "$root/lib/" 2>/dev/null || true; \ + cp -a /lib64/ld-linux-x86-64.so.2 "$root/lib64/" 2>/dev/null || true; \ + cp -a /etc/ssl "$root/etc/" 2>/dev/null || true; \ + cp -a /etc/ca-certificates "$root/etc/" 2>/dev/null || true; \ + cp -a /etc/resolv.conf /etc/hosts /etc/nsswitch.conf "$root/etc/" 2>/dev/null || true; \ + mkdir -p "$root/dev"; \ + mknod -m 666 "$root/dev/null" c 1 3 || true; \ + mknod -m 666 "$root/dev/zero" c 1 5 || true; \ + mknod -m 666 "$root/dev/random" c 1 8 || true; \ + mknod -m 666 "$root/dev/urandom" c 1 9 || true; \ + chmod -R a+rX "$root"; \ + chmod 755 "$root/tmp" + + +# 创建沙箱用户。code-sandbox 主进程默认保留 root，以便 Python 子进程 +# 在 native 隔离初始化阶段执行 chroot/setuid；用户代码会降权到 sandbox。 +RUN groupadd -g 65537 sandbox && useradd -u 65537 -g 65537 -M -r -s /usr/sbin/nologin sandbox && \ + mkdir -p /tmp/fastgpt-python-sandbox && \ + chown -R sandbox:sandbox /app && \ + chown root:root /tmp/fastgpt-python-sandbox && \ + chmod 755 /tmp/fastgpt-python-sandbox && \ + chown root:root /tmp/fastgpt-python-sandbox/tmp && \ + chmod 755 /tmp/fastgpt-python-sandbox/tmp ENV NODE_ENV=production ENV SANDBOX_PORT=3000 diff --git a/projects/code-sandbox/README.md b/projects/code-sandbox/README.md index c3977bf783c3..d929180cc454 100644 --- a/projects/code-sandbox/README.md +++ b/projects/code-sandbox/README.md @@ -1,28 +1,24 @@ # FastGPT Code Sandbox -基于 Node + Hono 的代码执行沙盒，支持 JS 和 Python。采用进程池架构，预热长驻 worker 进程，通过 stdin/stdout JSON 协议通信，消除每次请求的进程启动开销。 +基于 Node + Hono 的代码执行沙盒，支持 JS 和 Python。JS 采用长驻 worker 进程池；Python 采用 one-shot 预热进程池，Linux/Docker 环境固定启用 chroot、seccomp、setuid/setgid 隔离。 ## 架构 ``` -HTTP Request → Hono Server → Process Pool → Worker (long-lived) → Result - ↓ - ┌──────────────┐ - │ JS Workers │ node worker.js (×N) - │ Py Workers │ python3 worker.py (×N) - └──────────────┘ - stdin: JSON task → stdout: JSON result +HTTP Request → Hono Server + ├─ JS Process Pool → node worker.js (long-lived) → Result + └─ Python One-shot Warm Pool → clean python3 bootstrap → one task → exit ``` -- **进程池**：启动时预热 N 个 worker 进程（默认 20），请求到达时直接分配空闲 worker，执行完归还池中 +- **JS 进程池**：启动时预热 N 个 worker 进程（默认 20），请求到达时直接分配空闲 worker，执行完归还池中 - **JS 执行**：Node worker 进程 + 安全 shim（冻结 Function 构造器、危险全局对象遮蔽、require 白名单） -- **Python 执行**：python3 worker 进程 + `__import__` 拦截 + resource 资源限制 +- **Python 执行**：预热 `SANDBOX_POOL_SIZE` 个干净 python3 进程，进程进入 native seccomp/chroot/降权后等待一条任务；执行用户代码后立即销毁并异步补充新的干净进程 - **网络请求**：统一通过 `SystemHelper.httpRequest()` / `system_helper.http_request()` 收口，内置 SSRF 防护 -- **并发控制**：请求数超过池大小时自动排队，worker 崩溃自动重启补充 +- **并发控制**：JS 请求超过池大小时自动排队；Python 同时运行的独立子进程数复用 `SANDBOX_POOL_SIZE` ## 性能 -进程池 vs 旧版 spawn-per-request 对比（SANDBOX_POOL_SIZE=20）： +JS 仍保留进程池收益。Python 为了多租户安全改为 one-shot 预热池：空闲进程只在执行用户代码前复用，执行用户代码后立即销毁。它能降低 Python 冷启动成本，但吞吐仍会低于旧长驻 worker。 | 场景 | 旧版 QPS / P50 | 进程池 QPS / P50 | 提升 | |------|----------------|------------------|------| @@ -30,12 +26,7 @@ HTTP Request → Hono Server → Process Pool → Worker (long-lived) → Result | JS IO 500ms (c50) | 22 / 2,107ms | 38 / 1,005ms | 1.7x | | JS 高 CPU (c10) | 9 / 1,079ms | 12 / 796ms | 1.3x | | JS 高内存 (c10) | — | 13 / 787ms | — | -| Python 简单函数 (c50) | 14.7 / 2,897ms | 4,247 / 4ms | **289x** | -| Python IO 500ms (c50) | 14.2 / 3,066ms | 38 / 1,003ms | 2.7x | -| Python 高 CPU (c10) | 3.1 / 2,845ms | 4 / 2,191ms | 1.3x | -| Python 高内存 (c10) | — | 11 / 893ms | — | - -资源占用（20+20 workers）：空闲 ~1.5GB RSS，压测峰值 ~2GB RSS。 +资源占用由 `SANDBOX_POOL_SIZE`、Python 预热空闲进程、Python 包加载情况和 `SANDBOX_MAX_MEMORY_MB` 共同决定。 ## 快速开始 @@ -96,14 +87,23 @@ docker run -p 3000:3000 \ ### `GET /health` -健康检查，返回进程池状态。 +健康检查，返回 JS 进程池和 Python isolated runner 状态。 ```json { "status": "ok", - "version": "5.0.0", - "jsPool": { "total": 20, "idle": 18, "busy": 2, "queued": 0 }, - "pythonPool": { "total": 20, "idle": 20, "busy": 0, "queued": 0 } + "pools": { + "js": { "total": 20, "idle": 18, "busy": 2, "queued": 0, "poolSize": 20 }, + "python": { + "total": 20, + "idle": 18, + "busy": 2, + "warming": 0, + "queued": 0, + "poolSize": 20, + "ready": true + } + } } ``` @@ -139,13 +139,17 @@ docker run -p 3000:3000 \ | `SANDBOX_PORT` | 服务端口 | `3000` | | `SANDBOX_TOKEN` | Bearer Token 认证密钥 | 空（不鉴权） | -### 进程池 +### 并发控制 | 变量 | 说明 | 默认值 | |------|------|--------| -| `SANDBOX_POOL_SIZE` | 每种语言的 worker 进程数 | `20` | +| `SANDBOX_POOL_SIZE` | JS worker 进程数；也是 Python 同时运行和空闲预热的进程数 | `20` | | `SANDBOX_QUEUE_ID_CONCURRENCY` | 同一 `queueId` 同时可进入执行流程的请求数，空值表示不按 `queueId` 排队 | 空 | +### Python 隔离 + +Python 隔离不再提供运行时关闭开关。Linux 环境固定启用 native seccomp/chroot/降权，chroot 根目录固定为 `/tmp/fastgpt-python-sandbox`，用户代码进程固定降权到 `65537:65537`。Python 子进程不允许直接网络 syscall，外部请求必须通过父进程代理的 `http_request` 能力，并受请求次数、超时、请求体和响应体大小限制。 + ### 资源限制 | 变量 | 说明 | 默认值 | @@ -153,6 +157,7 @@ docker run -p 3000:3000 \ | `SANDBOX_API_MAX_BODY_MB` | API JSON 请求体总大小上限（包含 variables） | `8` | | `SANDBOX_MAX_TIMEOUT` | 超时上限（ms），请求不可超过此值 | `60000` | | `SANDBOX_MAX_MEMORY_MB` | 内存上限（MB） | `256` | +| `SANDBOX_MAX_TMP_MB` | Python 单任务临时目录写入上限（MB） | `16` | | `SANDBOX_MAX_OUTPUT_MB` | 单次执行输出 JSON 大小上限（包含返回值和日志） | `10` | ### 网络请求限制 @@ -174,9 +179,10 @@ src/ ├── types.ts # 类型定义 ├── pool/ │ ├── process-pool.ts # JS 进程池管理 -│ ├── python-process-pool.ts # Python 进程池管理 -│ ├── worker.ts # JS worker（长驻进程，含安全 shim） -│ └── worker.py # Python worker（长驻进程，含安全沙箱） +│ └── worker.ts # JS worker（长驻进程，含安全 shim） +├── isolated/ +│ ├── python-isolated-runner.ts # Python 独立进程执行器 +│ └── python-bootstrap.py # Python 单次执行 bootstrap └── utils/ └── semaphore.ts # 信号量（通用并发控制） diff --git a/projects/code-sandbox/build.sh b/projects/code-sandbox/build.sh index b0c733e0fb62..742f6dd7fd0b 100755 --- a/projects/code-sandbox/build.sh +++ b/projects/code-sandbox/build.sh @@ -6,6 +6,11 @@ echo "Building sandbox..." # 清理旧的构建产物 rm -rf dist +if [ "${SANDBOX_BUILD_NATIVE_PYTHON:-}" = "true" ]; then + echo "Building native Python sandbox library..." + pnpm run build:native:python +fi + # 编译入口（配置见 tsdown.config.ts）： # - 同时打包 index 和 worker 两个独立 bundle # - 所有 npm 依赖均打入 bundle（noExternal），仅保留 Node 内置模块外部化 @@ -16,16 +21,24 @@ pnpm exec tsdown # package.json 已声明 type:module，直接改后缀为 .js mv dist/index.mjs dist/index.js mv dist/worker.mjs dist/worker.js +mv dist/python-isolated-runner.mjs dist/python-isolated-runner.js -# Python worker 不需要编译，直接复制 -echo "Copying Python worker..." -cp src/pool/worker.py dist/worker.py +# Python bootstrap 不需要编译，直接复制 +echo "Copying Python runtime files..." +cp src/isolated/python-bootstrap.py dist/python-bootstrap.py +if [ -f src/isolated/fastgpt_python_sandbox.so ]; then + cp src/isolated/fastgpt_python_sandbox.so dist/fastgpt_python_sandbox.so +fi echo "" echo "Build complete!" echo " - index.js: $(du -h dist/index.js | cut -f1)" echo " - worker.js: $(du -h dist/worker.js | cut -f1)" -echo " - worker.py: $(du -h dist/worker.py | cut -f1)" +echo " - python-isolated-runner.js: $(du -h dist/python-isolated-runner.js | cut -f1)" +echo " - python-bootstrap.py: $(du -h dist/python-bootstrap.py | cut -f1)" +if [ -f dist/fastgpt_python_sandbox.so ]; then + echo " - fastgpt_python_sandbox.so: $(du -h dist/fastgpt_python_sandbox.so | cut -f1)" +fi echo "" echo "ℹ️ worker 通过 safeRequire(name) 在运行时动态加载白名单模块" echo " （lodash/dayjs/moment/uuid/crypto-js/qs），这些依赖必须以 node_modules" diff --git a/projects/code-sandbox/native/python-sandbox/cmd/lib/main.go b/projects/code-sandbox/native/python-sandbox/cmd/lib/main.go new file mode 100644 index 000000000000..c221d25e963e --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/cmd/lib/main.go @@ -0,0 +1,59 @@ +package main + +// FastGPT Python native sandbox. +// +// This package is project-local FastGPT glue code, not copied from Dify or any +// other sandbox project. It exposes a tiny C ABI for python-bootstrap.py and +// relies on the open-source github.com/seccomp/libseccomp-golang and +// golang.org/x/sys packages for Linux seccomp/prctl/syscall bindings. + +/* +#include +*/ +import "C" + +import ( + "sync" + "unsafe" + + "fastgpt-python-sandbox/internal/sandbox" +) + +var ( + lastErrMu sync.Mutex + lastErr string +) + +func setLastErr(err error) { + lastErrMu.Lock() + defer lastErrMu.Unlock() + if err == nil { + lastErr = "" + return + } + lastErr = err.Error() +} + +//export FastGPTInitPythonSandbox +func FastGPTInitPythonSandbox(uid C.int, gid C.int, enableNetwork C.int) C.int { + err := sandbox.Init(int(uid), int(gid), enableNetwork != 0) + setLastErr(err) + if err != nil { + return 1 + } + return 0 +} + +//export FastGPTLastError +func FastGPTLastError() *C.char { + lastErrMu.Lock() + defer lastErrMu.Unlock() + return C.CString(lastErr) +} + +//export FastGPTFreeCString +func FastGPTFreeCString(value *C.char) { + C.free(unsafe.Pointer(value)) +} + +func main() {} diff --git a/projects/code-sandbox/native/python-sandbox/go.mod b/projects/code-sandbox/native/python-sandbox/go.mod new file mode 100644 index 000000000000..2a7481d6b968 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/go.mod @@ -0,0 +1,8 @@ +module fastgpt-python-sandbox + +go 1.22 + +require ( + github.com/seccomp/libseccomp-golang v0.11.0 + golang.org/x/sys v0.29.0 +) diff --git a/projects/code-sandbox/native/python-sandbox/go.sum b/projects/code-sandbox/native/python-sandbox/go.sum new file mode 100644 index 000000000000..be9a8305b0d6 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/go.sum @@ -0,0 +1,4 @@ +github.com/seccomp/libseccomp-golang v0.11.0 h1:SDkcBRqGLP+sezmMACkxO1EfgbghxIxnRKfd6mHUEis= +github.com/seccomp/libseccomp-golang v0.11.0/go.mod h1:5m1Lk8E9OwgZTTVz4bBOer7JuazaBa+xTkM895tDiWc= +golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU= +golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_linux.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_linux.go new file mode 100644 index 000000000000..10b0e953ec03 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_linux.go @@ -0,0 +1,41 @@ +//go:build linux + +package sandbox + +import ( + "fmt" + "syscall" +) + +// Init installs the OS-level restrictions for the current Python process. +// +// The caller must load this shared library, open all required task fds, and set +// cwd to the prepared sandbox root before calling Init. After chroot/seccomp and +// uid/gid drop, the process cannot recover its previous privileges. +func Init(uid, gid int, enableNetwork bool) error { + if uid <= 0 || gid <= 0 { + return fmt.Errorf("uid/gid must be positive") + } + if err := syscall.Chroot("."); err != nil { + return fmt.Errorf("chroot: %w", err) + } + if err := syscall.Chdir("/"); err != nil { + return fmt.Errorf("chdir: %w", err) + } + if err := setNoNewPrivs(); err != nil { + return err + } + if err := loadSeccomp(enableNetwork); err != nil { + return err + } + if err := syscall.Setgroups([]int{}); err != nil { + return fmt.Errorf("setgroups: %w", err) + } + if err := syscall.Setgid(gid); err != nil { + return fmt.Errorf("setgid: %w", err) + } + if err := syscall.Setuid(uid); err != nil { + return fmt.Errorf("setuid: %w", err) + } + return nil +} diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_unsupported.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_unsupported.go new file mode 100644 index 000000000000..241eea6897e6 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/init_unsupported.go @@ -0,0 +1,9 @@ +//go:build !linux + +package sandbox + +import "fmt" + +func Init(uid, gid int, enableNetwork bool) error { + return fmt.Errorf("fastgpt python sandbox native isolation only supports linux") +} diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/no_new_privs_linux.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/no_new_privs_linux.go new file mode 100644 index 000000000000..ed4ec543bf0c --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/no_new_privs_linux.go @@ -0,0 +1,16 @@ +//go:build linux + +package sandbox + +import ( + "fmt" + + "golang.org/x/sys/unix" +) + +func setNoNewPrivs() error { + if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil { + return fmt.Errorf("prctl(PR_SET_NO_NEW_PRIVS): %w", err) + } + return nil +} diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/seccomp_linux.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/seccomp_linux.go new file mode 100644 index 000000000000..47f57fe8ca9e --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/seccomp_linux.go @@ -0,0 +1,36 @@ +//go:build linux + +package sandbox + +import ( + "fmt" + "syscall" + + seccomp "github.com/seccomp/libseccomp-golang" +) + +func loadSeccomp(enableNetwork bool) error { + filter, err := seccomp.NewFilter(seccomp.ActErrno.SetReturnCode(int16(syscall.EPERM))) + if err != nil { + return fmt.Errorf("create seccomp filter: %w", err) + } + + for _, syscallID := range allowBaseSyscalls { + if err := filter.AddRule(seccomp.ScmpSyscall(syscallID), seccomp.ActAllow); err != nil { + return fmt.Errorf("allow syscall %d: %w", syscallID, err) + } + } + + if enableNetwork { + for _, syscallID := range allowNetworkSyscalls { + if err := filter.AddRule(seccomp.ScmpSyscall(syscallID), seccomp.ActAllow); err != nil { + return fmt.Errorf("allow network syscall %d: %w", syscallID, err) + } + } + } + + if err := filter.Load(); err != nil { + return fmt.Errorf("load seccomp filter: %w", err) + } + return nil +} diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_amd64.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_amd64.go new file mode 100644 index 000000000000..9d74d6c60970 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_amd64.go @@ -0,0 +1,44 @@ +//go:build linux && amd64 + +package sandbox + +import "syscall" + +const ( + sysGetrandom = 318 + sysRseq = 334 + sysClone3 = 435 + sysSendmmsg = 307 +) + +var allowBaseSyscalls = []int{ + syscall.SYS_READ, syscall.SYS_WRITE, syscall.SYS_CLOSE, + syscall.SYS_NEWFSTATAT, syscall.SYS_FSTAT, syscall.SYS_FCNTL, syscall.SYS_IOCTL, + syscall.SYS_OPENAT, syscall.SYS_FACCESSAT, syscall.SYS_PREAD64, syscall.SYS_LSEEK, syscall.SYS_GETDENTS64, + // Python tempfile/matplotlib 需要在每个 task 临时目录内建目录；路径边界由 Python guard 和 chroot 权限控制。 + syscall.SYS_MKDIR, syscall.SYS_MKDIRAT, + syscall.SYS_MMAP, syscall.SYS_MPROTECT, syscall.SYS_MUNMAP, syscall.SYS_BRK, + syscall.SYS_MREMAP, syscall.SYS_MADVISE, + syscall.SYS_RT_SIGACTION, syscall.SYS_RT_SIGPROCMASK, syscall.SYS_SIGALTSTACK, syscall.SYS_RT_SIGRETURN, + syscall.SYS_CLOCK_GETTIME, syscall.SYS_GETTIMEOFDAY, syscall.SYS_NANOSLEEP, syscall.SYS_CLOCK_NANOSLEEP, + syscall.SYS_GETPID, syscall.SYS_GETPPID, syscall.SYS_GETTID, + syscall.SYS_GETUID, syscall.SYS_GETEUID, syscall.SYS_GETGID, syscall.SYS_GETEGID, syscall.SYS_GETGROUPS, + syscall.SYS_GETCWD, syscall.SYS_UNAME, syscall.SYS_SETITIMER, + syscall.SYS_SETGROUPS, syscall.SYS_SETGID, syscall.SYS_SETUID, + syscall.SYS_FUTEX, syscall.SYS_SCHED_GETAFFINITY, syscall.SYS_SCHED_YIELD, + syscall.SYS_EXIT, syscall.SYS_EXIT_GROUP, syscall.SYS_TGKILL, + sysGetrandom, syscall.SYS_EVENTFD2, + syscall.SYS_EPOLL_CREATE1, syscall.SYS_EPOLL_CTL, syscall.SYS_PSELECT6, + syscall.SYS_SET_ROBUST_LIST, syscall.SYS_GET_ROBUST_LIST, sysRseq, +} + +var errnoSyscalls = []int{ + syscall.SYS_CLONE, sysClone3, syscall.SYS_FORK, syscall.SYS_VFORK, +} + +var allowNetworkSyscalls = []int{ + syscall.SYS_SOCKET, syscall.SYS_CONNECT, syscall.SYS_BIND, syscall.SYS_LISTEN, syscall.SYS_ACCEPT, + syscall.SYS_SENDTO, syscall.SYS_RECVFROM, syscall.SYS_SENDMSG, sysSendmmsg, syscall.SYS_RECVMSG, + syscall.SYS_GETSOCKNAME, syscall.SYS_GETPEERNAME, syscall.SYS_SETSOCKOPT, syscall.SYS_GETSOCKOPT, + syscall.SYS_POLL, syscall.SYS_PPOLL, syscall.SYS_EPOLL_PWAIT, +} diff --git a/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_arm64.go b/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_arm64.go new file mode 100644 index 000000000000..8eedf33d8f25 --- /dev/null +++ b/projects/code-sandbox/native/python-sandbox/internal/sandbox/syscalls_arm64.go @@ -0,0 +1,43 @@ +//go:build linux && arm64 + +package sandbox + +import "syscall" + +const ( + sysRseq = 293 + sysClone3 = 435 +) + +var allowBaseSyscalls = []int{ + syscall.SYS_READ, syscall.SYS_WRITE, syscall.SYS_CLOSE, + syscall.SYS_FSTATAT, syscall.SYS_FSTAT, syscall.SYS_FCNTL, syscall.SYS_IOCTL, + syscall.SYS_OPENAT, syscall.SYS_FACCESSAT, syscall.SYS_READLINKAT, + // Python tempfile/matplotlib 需要在每个 task 临时目录内建目录；路径边界由 Python guard 和 chroot 权限控制。 + syscall.SYS_MKDIRAT, + syscall.SYS_PREAD64, syscall.SYS_LSEEK, syscall.SYS_GETDENTS64, + syscall.SYS_MMAP, syscall.SYS_MPROTECT, syscall.SYS_MUNMAP, syscall.SYS_BRK, + syscall.SYS_MREMAP, syscall.SYS_MADVISE, + syscall.SYS_RT_SIGACTION, syscall.SYS_RT_SIGPROCMASK, syscall.SYS_SIGALTSTACK, syscall.SYS_RT_SIGRETURN, + syscall.SYS_CLOCK_GETTIME, syscall.SYS_GETTIMEOFDAY, syscall.SYS_NANOSLEEP, syscall.SYS_CLOCK_NANOSLEEP, + syscall.SYS_GETPID, syscall.SYS_GETPPID, syscall.SYS_GETTID, + syscall.SYS_GETUID, syscall.SYS_GETEUID, syscall.SYS_GETGID, syscall.SYS_GETEGID, syscall.SYS_GETGROUPS, + syscall.SYS_GETCWD, syscall.SYS_UNAME, syscall.SYS_SETITIMER, + syscall.SYS_SETGROUPS, syscall.SYS_SETGID, syscall.SYS_SETUID, + syscall.SYS_FUTEX, syscall.SYS_SCHED_GETAFFINITY, syscall.SYS_SCHED_YIELD, + syscall.SYS_EXIT, syscall.SYS_EXIT_GROUP, syscall.SYS_TGKILL, + syscall.SYS_GETRANDOM, syscall.SYS_EVENTFD2, + syscall.SYS_EPOLL_CREATE1, syscall.SYS_EPOLL_CTL, syscall.SYS_PSELECT6, + syscall.SYS_SET_ROBUST_LIST, syscall.SYS_GET_ROBUST_LIST, sysRseq, +} + +var errnoSyscalls = []int{ + syscall.SYS_CLONE, sysClone3, +} + +var allowNetworkSyscalls = []int{ + syscall.SYS_SOCKET, syscall.SYS_CONNECT, syscall.SYS_BIND, syscall.SYS_LISTEN, syscall.SYS_ACCEPT, + syscall.SYS_SENDTO, syscall.SYS_RECVFROM, syscall.SYS_SENDMSG, syscall.SYS_SENDMMSG, syscall.SYS_RECVMSG, + syscall.SYS_GETSOCKNAME, syscall.SYS_GETPEERNAME, syscall.SYS_SETSOCKOPT, syscall.SYS_GETSOCKOPT, + syscall.SYS_PPOLL, syscall.SYS_EPOLL_PWAIT, +} diff --git a/projects/code-sandbox/package.json b/projects/code-sandbox/package.json index e7c6bb836031..7906bdcccab4 100644 --- a/projects/code-sandbox/package.json +++ b/projects/code-sandbox/package.json @@ -9,6 +9,7 @@ "scripts": { "dev": "tsx watch src/index.ts", "start": "node dist/index.js", + "build:native:python": "cd native/python-sandbox && go build -buildmode=c-shared -o ../../src/isolated/fastgpt_python_sandbox.so ./cmd/lib", "build": "sh build.sh", "test": "vitest run", "test:watch": "vitest" diff --git a/projects/code-sandbox/src/env.ts b/projects/code-sandbox/src/env.ts index 6cad72576a7c..29abfb02ee03 100644 --- a/projects/code-sandbox/src/env.ts +++ b/projects/code-sandbox/src/env.ts @@ -54,6 +54,7 @@ export const env = createEnv({ SANDBOX_API_MAX_BODY_MB: IntSchema.min(1).max(100).default(8), SANDBOX_MAX_TIMEOUT: IntSchema.min(1000).max(600000).default(60000), SANDBOX_MAX_MEMORY_MB: IntSchema.min(32).max(4096).default(256), + SANDBOX_MAX_TMP_MB: IntSchema.min(1).max(1024).default(16), SANDBOX_MAX_OUTPUT_MB: IntSchema.min(1).max(100).default(10), // ===== 网络请求限制 ===== @@ -81,7 +82,7 @@ export const env = createEnv({ 'json,csv,base64,binascii,struct,' + 'hashlib,hmac,secrets,uuid,' + 'typing,abc,enum,dataclasses,contextlib,' + - 'pprint,' + + 'pprint,weakref,' + 'numpy,pandas,matplotlib' ) .transform(parseAllowedModules) diff --git a/projects/code-sandbox/src/index.ts b/projects/code-sandbox/src/index.ts index cd04f957af08..91019dab7778 100644 --- a/projects/code-sandbox/src/index.ts +++ b/projects/code-sandbox/src/index.ts @@ -4,7 +4,7 @@ import { bearerAuth } from 'hono/bearer-auth'; import { serve } from '@hono/node-server'; import { z } from 'zod'; import { ProcessPool } from './pool/process-pool'; -import { PythonProcessPool } from './pool/python-process-pool'; +import { PythonIsolatedRunner } from './isolated/python-isolated-runner'; import type { ExecuteOptions } from './types'; import { getErrText } from './utils'; import { configureLogger, getLogger, LogCategories } from './utils/logger'; @@ -91,7 +91,7 @@ const app = new Hono(); /** 进程池 */ const jsPool = new ProcessPool(env.SANDBOX_POOL_SIZE); -const pythonPool = new PythonProcessPool(env.SANDBOX_POOL_SIZE); +const pythonRunner = new PythonIsolatedRunner(); const queueIdLimiter = new QueueIdLimiter(env.SANDBOX_QUEUE_ID_CONCURRENCY); if (queueIdLimiter.enabled) { @@ -100,11 +100,9 @@ if (queueIdLimiter.enabled) { ); } -const poolReady = Promise.all([jsPool.init(), pythonPool.init()]) +const poolReady = Promise.all([jsPool.init(), pythonRunner.init()]) .then(() => { - serverLogger.info( - `Process pools ready: JS=${env.SANDBOX_POOL_SIZE}, Python=${env.SANDBOX_POOL_SIZE} workers` - ); + serverLogger.info(`Process pools ready: JS=${env.SANDBOX_POOL_SIZE}, Python=isolated`); }) .catch((err) => { serverLogger.error('Failed to init process pool:', err.message); @@ -114,9 +112,18 @@ const poolReady = Promise.all([jsPool.init(), pythonPool.init()]) /** 健康检查（不需要认证） */ app.get('/health', (c) => { const jsStats = jsPool.stats; - const pyStats = pythonPool.stats; - const isReady = jsStats.total > 0 && pyStats.total > 0; - return c.json({ status: isReady ? 'ok' : 'degraded' }, isReady ? 200 : 503); + const pythonStats = pythonRunner.stats; + const isReady = jsStats.total > 0 && pythonStats.ready; + return c.json( + { + status: isReady ? 'ok' : 'degraded', + pools: { + js: jsStats, + python: pythonStats + } + }, + isReady ? 200 : 503 + ); }); // 增加日志中间件，打印请求信息 @@ -219,7 +226,7 @@ app.post('/sandbox/python', async (c) => { ); } const result = await queueIdLimiter.run(parsed.data.queueId, () => - pythonPool.execute(parsed.data as ExecuteOptions) + pythonRunner.execute(parsed.data as ExecuteOptions) ); return c.json(result); } catch (err: any) { diff --git a/projects/code-sandbox/src/isolated/python-bootstrap.py b/projects/code-sandbox/src/isolated/python-bootstrap.py new file mode 100644 index 000000000000..9af5be3fdb68 --- /dev/null +++ b/projects/code-sandbox/src/isolated/python-bootstrap.py @@ -0,0 +1,881 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Single-shot Python bootstrap for FastGPT isolated runner.""" + +import ast as _ast +import base64 as _base64 +import builtins as _builtins +import ctypes as _ctypes +import copy as _copy +import hashlib as _hashlib +import hmac as _hmac +import inspect as _inspect_mod +import ipaddress as _ipaddress +import json +import math as _math +import os as _os +import signal +import sys +import sysconfig as _sysconfig +import time as _time +import traceback as _tb +import types as _types +import urllib.parse as _urllib_parse +import encodings.idna as _encodings_idna # noqa: F401 + +_STDLIB_MODULES = sys.stdlib_module_names if hasattr(sys, 'stdlib_module_names') else frozenset() +_DANGEROUS_STDLIB = frozenset({ + 'os', 'subprocess', 'shutil', 'pathlib', 'glob', 'tempfile', + 'multiprocessing', 'threading', 'concurrent', + 'ctypes', 'importlib', 'runpy', 'code', 'codeop', 'compileall', + 'socket', 'http', 'urllib', 'ftplib', 'smtplib', 'poplib', 'imaplib', + 'xmlrpc', 'socketserver', 'ssl', 'asyncio', 'selectors', 'select', + 'signal', 'resource', 'pty', 'termios', 'tty', 'fcntl', + 'mmap', 'dbm', 'sqlite3', 'shelve', + 'webbrowser', 'turtle', 'tkinter', 'idlelib', + 'venv', 'ensurepip', 'pip', 'site', + 'gc', 'sys', 'builtins', 'marshal', 'pickle', +}) + +_REQUEST_LIMITS = { + 'max_requests': 30, + 'timeout': 60, + 'max_response_size': 10 * 1024 * 1024, + 'max_request_body_size': 5 * 1024 * 1024, + 'max_output_size': 10 * 1024 * 1024, + 'allowed_protocols': ['http:', 'https:'] +} + +_request_count = 0 +_allowed_modules = set() +_original_import = _builtins.__import__ +_original_json_dumps = json.dumps +_builtins_proxy = None +_import_guard = False +_original_open = open +_original_os_functions = {} +_open_guard = False +_path_guard = False +_logs = [] +_log_size = 0 +_MAX_LOG_SIZE = 1024 * 1024 +_timeout_stage = 0 +_audit_hook_installed = False +_native_isolation_ready = False +_task_tmpdir = None +_matplotlib_tmpdir = None + +_FORBIDDEN_ATTRS = frozenset({ + '__class__', '__base__', '__bases__', '__mro__', '__subclasses__', + '__globals__', '__code__', '__closure__', '__func__', '__self__', + '__dict__', '__getattribute__', '__getattr__', '__setattr__', +}) + +_FORBIDDEN_BUILTINS = frozenset({ + 'eval', 'exec', 'compile', 'input', 'breakpoint', + 'globals', 'locals', 'vars', 'dir', 'super', +}) + +_PROTECTED_BUILTINS = _FORBIDDEN_BUILTINS | frozenset({ + '__import__', 'open', 'getattr', 'setattr', 'delattr', +}) + + +def _write_result(payload): + sys.stdout.write(_original_json_dumps({'type': 'result', **payload}, ensure_ascii=False, default=str) + '\n') + sys.stdout.flush() + + +def _init_request_limits(limits): + if not limits: + return + if 'maxRequests' in limits: + _REQUEST_LIMITS['max_requests'] = limits['maxRequests'] + if 'timeoutMs' in limits: + _REQUEST_LIMITS['timeout'] = max(1, limits['timeoutMs'] // 1000) + if 'maxResponseSize' in limits: + _REQUEST_LIMITS['max_response_size'] = limits['maxResponseSize'] + if 'maxRequestBodySize' in limits: + _REQUEST_LIMITS['max_request_body_size'] = limits['maxRequestBodySize'] + if 'maxOutputSize' in limits: + _REQUEST_LIMITS['max_output_size'] = limits['maxOutputSize'] + + +class _SystemHelper: + __slots__ = () + + @staticmethod + def http_request(url, method='GET', headers=None, body=None, timeout=None, timeout_ms=None): + if headers is None: + headers = {} + if body is not None: + body_text = body if isinstance(body, str) else _original_json_dumps(body, ensure_ascii=False) + if len(body_text.encode('utf-8')) > _REQUEST_LIMITS['max_request_body_size']: + raise RuntimeError("Request body too large") + else: + body_text = None + + if body_text is not None and len(body_text.encode('utf-8')) > _REQUEST_LIMITS['max_request_body_size']: + raise RuntimeError("Request body too large") + + return _call_parent_http_proxy({ + 'url': url, + 'method': method, + 'headers': headers, + 'body': body if body_text is None or isinstance(body, str) else body, + 'timeout': timeout, + 'timeoutMs': timeout_ms + }) + + httpRequest = http_request + + +system_helper = _SystemHelper() +SystemHelper = system_helper + + +def count_token(text): + if not isinstance(text, str): + text = str(text) + return _math.ceil(len(text) / 4) + + +def str_to_base64(text, prefix=''): + b64 = _base64.b64encode(text.encode('utf-8')).decode('utf-8') + return prefix + b64 + + +def create_hmac(algorithm, secret): + timestamp = str(int(_time.time() * 1000)) + string_to_sign = timestamp + '\n' + secret + h = _hmac.new(secret.encode('utf-8'), string_to_sign.encode('utf-8'), algorithm) + sign = _urllib_parse.quote(_base64.b64encode(h.digest()).decode('utf-8')) + return {"timestamp": timestamp, "sign": sign} + + +def delay(ms): + if ms > 10000: + raise ValueError("Delay must be <= 10000ms") + _time.sleep(ms / 1000) + return None + + +_rpc_seq = 0 + + +def _call_parent_http_proxy(payload): + global _rpc_seq + _rpc_seq += 1 + req_id = f'http-{_rpc_seq}' + sys.stdout.write(_original_json_dumps({ + 'type': 'http_request', + 'id': req_id, + 'payload': payload + }, ensure_ascii=False, default=str) + '\n') + sys.stdout.flush() + + while True: + line = sys.stdin.readline() + if not line: + raise RuntimeError('HTTP proxy response channel closed') + try: + msg = json.loads(line) + except Exception: + continue + if msg.get('type') != 'http_response' or msg.get('id') != req_id: + continue + if msg.get('success'): + return msg.get('payload') + raise RuntimeError(msg.get('message') or 'HTTP request failed') + + +def _init_native_isolation(isolation): + global _native_isolation_ready + if _native_isolation_ready: + return + if not isolation or not isolation.get('enableSeccomp'): + return + + lib_path = isolation.get('libraryPath') or './fastgpt_python_sandbox.so' + try: + lib = _ctypes.CDLL(lib_path) + except Exception as e: + raise RuntimeError(f"Failed to load native python sandbox library: {e}") + + lib.FastGPTInitPythonSandbox.argtypes = [_ctypes.c_int, _ctypes.c_int, _ctypes.c_int] + lib.FastGPTInitPythonSandbox.restype = _ctypes.c_int + lib.FastGPTLastError.argtypes = [] + lib.FastGPTLastError.restype = _ctypes.c_void_p + lib.FastGPTFreeCString.argtypes = [_ctypes.c_void_p] + lib.FastGPTFreeCString.restype = None + + ret = lib.FastGPTInitPythonSandbox( + int(isolation.get('uid') or 65537), + int(isolation.get('gid') or 65537), + 1 if isolation.get('enableNetwork') else 0 + ) + if ret == 0: + _native_isolation_ready = True + return + + err_ptr = lib.FastGPTLastError() + try: + err_text = _ctypes.cast(err_ptr, _ctypes.c_char_p).value.decode('utf-8') if err_ptr else '' + finally: + if err_ptr: + lib.FastGPTFreeCString(err_ptr) + raise RuntimeError(err_text or f"Native python sandbox init failed: {ret}") + + +def _is_stdlib_frame(filename: str): + if not filename: + return False + try: + stdlib = _sysconfig.get_paths().get('stdlib') or '' + platstdlib = _sysconfig.get_paths().get('platstdlib') or '' + return (stdlib and filename.startswith(stdlib)) or (platstdlib and filename.startswith(platstdlib)) + except Exception: + return False + + +def _is_site_packages_frame(filename: str): + return 'site-packages' in filename or 'dist-packages' in filename + + +class _BuiltinsProxy(_types.ModuleType): + def __init__(self, original): + super().__init__('builtins') + object.__setattr__(self, '_original', original) + + def __getattr__(self, name): + if name == 'open': + return _restricted_open + if name == 'getattr': + return _safe_getattr + if name == 'setattr': + return _safe_setattr + if name == 'delattr': + return _safe_delattr + if name in _FORBIDDEN_BUILTINS: + raise AttributeError(f"builtins.{name} is not available in sandbox") + return getattr(object.__getattribute__(self, '_original'), name) + + def __setattr__(self, name, value): + if name in _PROTECTED_BUILTINS: + raise AttributeError(f"builtins.{name} cannot be modified in sandbox") + return setattr(object.__getattribute__(self, '_original'), name, value) + + def __delattr__(self, name): + if name in _PROTECTED_BUILTINS: + raise AttributeError(f"builtins.{name} cannot be deleted in sandbox") + return delattr(object.__getattribute__(self, '_original'), name) + + +def _safe_getattr(obj, name, *args): + if isinstance(name, str) and name in _FORBIDDEN_ATTRS: + raise AttributeError(f"Access to {name} is not allowed in sandbox") + return getattr(obj, name, *args) + + +def _safe_setattr(obj, name, value): + raise AttributeError("Use of setattr is not allowed in sandbox") + + +def _safe_delattr(obj, name): + raise AttributeError("Use of delattr is not allowed in sandbox") + + +def _is_direct_user_import_call(): + global _import_guard + _import_guard = True + try: + stack = _tb.extract_stack() + finally: + _import_guard = False + + if len(stack) < 3: + return False + caller_fn = stack[-3].filename or '' + return caller_fn in ('', '', '') + + +def _safe_import(name, *args, **kwargs): + if _import_guard: + return _original_import(name, *args, **kwargs) + top_level = name.split('.')[0] + if top_level == 'builtins' and _builtins_proxy is not None: + return _builtins_proxy + if top_level in _STDLIB_MODULES and top_level not in _DANGEROUS_STDLIB: + return _original_import(name, *args, **kwargs) + if top_level in _allowed_modules: + return _original_import(name, *args, **kwargs) + if _is_direct_user_import_call(): + raise ImportError(f"Module '{name}' is not in the allowlist.") + return _original_import(name, *args, **kwargs) + + +def _validate_user_code(code: str): + try: + tree = _ast.parse(code) + except SyntaxError: + return + + forbidden_builtin_calls = { + 'eval', 'exec', 'compile', 'globals', 'locals', 'vars', 'dir', + 'breakpoint', 'input', 'super' + } + + for node in _ast.walk(tree): + if isinstance(node, _ast.Attribute) and node.attr in _FORBIDDEN_ATTRS: + raise RuntimeError(f"Access to {node.attr} is not allowed in sandbox") + if isinstance(node, _ast.Call) and isinstance(node.func, _ast.Name): + if node.func.id in forbidden_builtin_calls: + raise RuntimeError(f"Use of {node.func.id} is not allowed in sandbox") + if node.func.id in ('setattr', 'delattr'): + raise RuntimeError(f"Use of {node.func.id} is not allowed in sandbox") + if node.func.id == 'getattr': + if len(node.args) < 2 or not isinstance(node.args[1], _ast.Constant): + raise RuntimeError("Dynamic getattr is not allowed in sandbox") + if node.args[1].value in _FORBIDDEN_ATTRS: + raise RuntimeError(f"Access to {node.args[1].value} is not allowed in sandbox") + + +def _restricted_open(*args, **kwargs): + global _open_guard + path = args[0] if args else None + mode = kwargs.get('mode', args[1] if len(args) > 1 else 'r') + if _open_guard: + return _original_open(*args, **kwargs) + _open_guard = True + try: + stack = _tb.extract_stack() + finally: + _open_guard = False + if len(stack) >= 2: + caller_fn = stack[-2].filename or '' + if _is_user_code_filename(caller_fn) and not _is_path_under_task_tmp(path): + raise PermissionError("File system access is not allowed in sandbox") + if ( + not _is_stdlib_frame(caller_fn) + and not _is_site_packages_frame(caller_fn) + and caller_fn != __file__ + and not _is_path_under_task_tmp(path) + ): + raise PermissionError("File system access is not allowed in sandbox") + if _is_write_mode(mode) and not _is_path_under_task_tmp(path): + raise PermissionError("File writes are only allowed in the task temporary directory") + return _original_open(*args, **kwargs) + + +def _is_user_code_filename(filename): + return filename in ('', '', '') + + +def _is_write_mode(mode): + mode_text = str(mode or 'r') + return any(flag in mode_text for flag in ('w', 'a', 'x', '+')) + + +def _is_write_flags(flags): + try: + flags_int = int(flags) + except Exception: + return False + return bool( + flags_int + & ( + _os.O_WRONLY + | _os.O_RDWR + | _os.O_CREAT + | _os.O_APPEND + | _os.O_TRUNC + ) + ) + + +def _is_path_under_task_tmp(path): + global _path_guard + if not _task_tmpdir: + return False + try: + path_text = _os.fspath(path) + except Exception: + return False + if not isinstance(path_text, str): + return False + try: + _path_guard = True + root = _os.path.realpath(_task_tmpdir) + candidate = _os.path.realpath(path_text if _os.path.isabs(path_text) else _os.path.join(_os.getcwd(), path_text)) + return candidate == root or candidate.startswith(root.rstrip(_os.sep) + _os.sep) + except Exception: + return False + finally: + _path_guard = False + + +def _is_called_from_matplotlib(): + try: + frame = sys._getframe(1) + while frame: + filename = frame.f_code.co_filename or '' + normalized = filename.replace('\\', '/') + if '/matplotlib/' in normalized: + return True + frame = frame.f_back + except Exception: + return False + return False + + +def _first_external_caller_filename(): + try: + frame = sys._getframe(1) + while frame: + filename = frame.f_code.co_filename or '' + if filename != __file__: + return filename + frame = frame.f_back + except Exception: + return '' + return '' + + +def _is_direct_user_fs_access(): + return _is_user_code_filename(_first_external_caller_filename()) + + +def _guard_fs_read_path(path): + if _path_guard: + return + if _is_direct_user_fs_access() and not _is_path_under_task_tmp(path): + raise PermissionError("File system access is only allowed in the task temporary directory") + + +def _guard_fs_write_path(path): + if _path_guard: + return + if not _is_path_under_task_tmp(path): + raise PermissionError("File system writes are only allowed in the task temporary directory") + + +def _install_os_guards(): + if _original_os_functions: + return + + def wrap_read_path(name): + original = getattr(_os, name, None) + if original is None: + return + _original_os_functions[name] = original + + def guarded(path, *args, **kwargs): + _guard_fs_read_path(path) + return original(path, *args, **kwargs) + + setattr(_os, name, guarded) + + def wrap_write_path(name): + original = getattr(_os, name, None) + if original is None: + return + _original_os_functions[name] = original + + def guarded(path, *args, **kwargs): + _guard_fs_write_path(path) + return original(path, *args, **kwargs) + + setattr(_os, name, guarded) + + def wrap_write_pair(name): + original = getattr(_os, name, None) + if original is None: + return + _original_os_functions[name] = original + + def guarded(src, dst, *args, **kwargs): + _guard_fs_write_path(src) + _guard_fs_write_path(dst) + return original(src, dst, *args, **kwargs) + + setattr(_os, name, guarded) + + original_open = _os.open + _original_os_functions['open'] = original_open + + def guarded_open(path, flags, *args, **kwargs): + if _is_write_flags(flags): + _guard_fs_write_path(path) + else: + _guard_fs_read_path(path) + return original_open(path, flags, *args, **kwargs) + + _os.open = guarded_open + + for name in ('listdir', 'scandir', 'stat', 'lstat', 'access', 'chdir'): + wrap_read_path(name) + for name in ('mkdir', 'makedirs', 'remove', 'unlink', 'rmdir', 'truncate', 'chmod', 'chown', 'utime'): + wrap_write_path(name) + for name in ('rename', 'replace', 'symlink', 'link'): + wrap_write_pair(name) + + +def _install_matplotlib_tmpdir_patch(): + if not _matplotlib_tmpdir: + return + try: + import tempfile as _tempfile + except Exception: + return + + original_gettempdir = _tempfile.gettempdir + original_gettempdirb = getattr(_tempfile, 'gettempdirb', None) + original_candidate_tempdir_list = getattr(_tempfile, '_candidate_tempdir_list', None) + + def matplotlib_gettempdir(): + if _is_called_from_matplotlib(): + return _matplotlib_tmpdir + return original_gettempdir() + + _tempfile.gettempdir = matplotlib_gettempdir + + if original_gettempdirb is not None: + def matplotlib_gettempdirb(): + if _is_called_from_matplotlib(): + return _os.fsencode(_matplotlib_tmpdir) + return original_gettempdirb() + + _tempfile.gettempdirb = matplotlib_gettempdirb + + if original_candidate_tempdir_list is not None: + def matplotlib_candidate_tempdir_list(): + if _is_called_from_matplotlib(): + return [_matplotlib_tmpdir] + return original_candidate_tempdir_list() + + _tempfile._candidate_tempdir_list = matplotlib_candidate_tempdir_list + + +def _init_task_tmpdir(path): + global _task_tmpdir, _matplotlib_tmpdir + _task_tmpdir = path or _os.environ.get('FASTGPT_TASK_TMPDIR') or '/tmp' + try: + _os.makedirs(_task_tmpdir, mode=0o700, exist_ok=True) + _os.environ['HOME'] = _task_tmpdir + _os.environ['TMPDIR'] = _task_tmpdir + mpl_config_dir = _os.path.join(_task_tmpdir, 'matplotlib') + mpl_cache_dir = _os.path.join(mpl_config_dir, 'cache') + mpl_xdg_config_dir = _os.path.join(mpl_config_dir, 'config') + mpl_tmp_dir = _os.path.join(mpl_config_dir, 'tmp') + _os.makedirs(mpl_config_dir, mode=0o700, exist_ok=True) + _os.makedirs(mpl_cache_dir, mode=0o700, exist_ok=True) + _os.makedirs(mpl_xdg_config_dir, mode=0o700, exist_ok=True) + _os.makedirs(mpl_tmp_dir, mode=0o700, exist_ok=True) + _os.environ['MPLCONFIGDIR'] = mpl_config_dir + _os.environ['XDG_CACHE_HOME'] = mpl_cache_dir + _os.environ['XDG_CONFIG_HOME'] = mpl_xdg_config_dir + _os.environ['MATPLOTLIB_TMPDIR'] = mpl_tmp_dir + _matplotlib_tmpdir = mpl_tmp_dir + except Exception as e: + raise RuntimeError(f"Failed to initialize task temporary directory: {e}") + + +def _safe_print(*args, **kwargs): + global _log_size + line = ' '.join(str(a) for a in args) + if _log_size + len(line) <= _MAX_LOG_SIZE: + _logs.append(line) + _log_size += len(line) + + +def _timeout_handler(signum, frame): + global _timeout_stage + _timeout_stage += 1 + if _timeout_stage >= 2: + raise SystemExit("Script execution timed out (forced)") + signal.alarm(1) + raise TimeoutError("Script execution timed out") + + +def _install_audit_hook(): + global _audit_hook_installed + if _audit_hook_installed: + return + _audit_hook_installed = True + + def _audit(event, args): + if event == 'open': + path = args[0] if len(args) > 0 else None + mode = args[1] if len(args) > 1 else 'r' + flags = args[2] if len(args) > 2 else 0 + if not _is_path_under_task_tmp(path) and (_is_write_mode(mode) or _is_write_flags(flags)): + raise RuntimeError("File writes are only allowed in the task temporary directory") + if _is_direct_user_fs_access() and not _is_path_under_task_tmp(path): + raise RuntimeError("File system access is only allowed in the task temporary directory") + return + if event in ( + 'os.listdir', + 'os.scandir', + 'os.stat', + 'os.lstat', + 'os.access', + 'os.chdir', + ): + path = args[0] if len(args) > 0 else None + if _is_direct_user_fs_access() and not _is_path_under_task_tmp(path): + raise RuntimeError("File system access is only allowed in the task temporary directory") + return + if event in ( + 'os.mkdir', + 'os.makedirs', + 'os.remove', + 'os.unlink', + 'os.rmdir', + 'os.rename', + 'os.replace', + 'os.symlink', + 'os.link', + 'os.truncate', + 'os.chmod', + 'os.chown', + 'os.utime', + 'shutil.rmtree' + ): + path = args[0] if len(args) > 0 else None + target = args[1] if event in ('os.rename', 'os.replace', 'os.symlink', 'os.link') and len(args) > 1 else None + if not _is_path_under_task_tmp(path) or (target is not None and not _is_path_under_task_tmp(target)): + raise RuntimeError("File system writes are only allowed in the task temporary directory") + return + if ( + event == 'os.system' + or event.startswith('os.exec') + or event.startswith('os.spawn') + or event.startswith('os.posix_spawn') + or event.startswith('socket.') + or event.startswith('ctypes.') + ): + raise RuntimeError(f"Operation {event} is not allowed in sandbox") + if event.startswith('subprocess.'): + if _is_allowed_matplotlib_font_probe(event, args): + return + raise RuntimeError(f"Operation {event} is not allowed in sandbox") + + sys.addaudithook(_audit) + + +def _is_allowed_matplotlib_font_probe(event, args): + if event != 'subprocess.Popen': + return False + if not _is_called_from_matplotlib_font_manager(): + return False + + command = args[1] if len(args) > 1 else args[0] if args else None + if isinstance(command, (list, tuple)): + parts = [str(item) for item in command] + elif isinstance(command, str): + parts = command.split() + else: + return False + + if parts[:2] == ['fc-list', '--help']: + return True + if parts[:2] == ['fc-list', '--format=%{file}\\n']: + return True + if parts == ['system_profiler', '-xml', 'SPFontsDataType']: + return True + return False + + +def _is_called_from_matplotlib_font_manager(): + try: + for frame in _inspect_mod.stack()[2:]: + filename = frame.filename or '' + normalized = filename.replace('\\', '/') + if normalized.endswith('/matplotlib/font_manager.py'): + return True + except Exception: + return False + return False + + +_PROTECTED_MODULES = [json, _math, _time, _base64, _hashlib, _hmac, _copy] + + +def _snapshot_modules(): + snapshots = [] + for mod in _PROTECTED_MODULES: + attrs = {} + for name in dir(mod): + if not name.startswith('__'): + try: + attrs[name] = getattr(mod, name) + except Exception: + pass + snapshots.append((mod, attrs)) + return snapshots + + +def _restore_modules(snapshots): + for mod, attrs in snapshots: + current_names = set(n for n in dir(mod) if not n.startswith('__')) + original_names = set(attrs.keys()) + for name in current_names - original_names: + try: + delattr(mod, name) + except Exception: + pass + for name, val in attrs.items(): + try: + setattr(mod, name, val) + except Exception: + pass + + +def _call_main(user_main, variables): + sig = _inspect_mod.signature(user_main) + params = list(sig.parameters.keys()) + + if len(params) == 0: + return user_main() + if len(params) == 1: + p = params[0] + if p in variables: + return user_main(variables[p]) + return user_main(variables) + + for p in params: + if p not in variables and sig.parameters[p].default is _inspect_mod.Parameter.empty: + raise TypeError(f"Missing required argument: '{p}'") + call_kwargs = {p: variables[p] for p in params if p in variables} + return user_main(**call_kwargs) + + +def _run_task(msg): + global _allowed_modules, _builtins_proxy, _request_count, _logs, _log_size, _timeout_stage + _allowed_modules = set(msg.get('allowedModules', [])) + _init_request_limits(msg.get('requestLimits')) + _request_count = 0 + _logs = [] + _log_size = 0 + _timeout_stage = 0 + + code = msg.get('code', '') + variables = msg.get('variables', {}) + timeout_ms = msg.get('timeoutMs', 10000) + timeout_s = max(1, -(-timeout_ms // 1000)) + + _builtins.__import__ = _safe_import + _builtins.print = _safe_print + _builtins_proxy = _BuiltinsProxy(_builtins) + + snapshots = _snapshot_modules() + + try: + _init_native_isolation(msg.get('isolation') or {}) + _init_task_tmpdir(msg.get('taskTmpDir')) + _install_os_guards() + _install_matplotlib_tmpdir_patch() + + signal.signal(signal.SIGALRM, _timeout_handler) + signal.alarm(timeout_s) + + safe_builtins = {} + for name in dir(_builtins): + if name.startswith('_') and name not in ( + '__name__', '__doc__', '__import__', '__build_class__', + ): + continue + if name in _FORBIDDEN_BUILTINS: + continue + safe_builtins[name] = getattr(_builtins, name) + safe_builtins['__import__'] = _safe_import + safe_builtins['__build_class__'] = _builtins.__build_class__ + safe_builtins['open'] = _restricted_open + safe_builtins['getattr'] = _safe_getattr + safe_builtins['setattr'] = _safe_setattr + safe_builtins['delattr'] = _safe_delattr + + class _SafeObject(object): + __subclasses__ = None + safe_builtins['object'] = _SafeObject + + exec_globals = { + '__builtins__': safe_builtins, + 'variables': variables, + 'SystemHelper': system_helper, + 'system_helper': system_helper, + 'task_tmpdir': _task_tmpdir, + 'count_token': count_token, + 'str_to_base64': str_to_base64, + 'create_hmac': create_hmac, + 'delay': delay, + 'http_request': system_helper.http_request, + 'print': _safe_print, + 'json': json, + 'math': _math, + 'time': _time, + } + reserved_keys = frozenset(exec_globals.keys()) + for k, v in variables.items(): + if k not in reserved_keys: + exec_globals[k] = v + + _validate_user_code(code) + _install_audit_hook() + exec(code, exec_globals) + + user_main = exec_globals.get('main') + if user_main is None: + raise RuntimeError("No 'main' function defined") + + result = _call_main(user_main, variables) + signal.alarm(0) + _write_result({'success': True, 'data': {'codeReturn': result, 'log': '\n'.join(_logs)}}) + except (Exception, SystemExit) as e: + signal.alarm(0) + _write_result({'success': False, 'message': str(e)}) + finally: + _restore_modules(snapshots) + _builtins.__import__ = _original_import + _builtins.open = _original_open + + +def _run_warm_worker(init_msg): + try: + _init_native_isolation(init_msg.get('isolation') or {}) + sys.stdout.write(_original_json_dumps({'type': 'ready'}, ensure_ascii=False) + '\n') + sys.stdout.flush() + except (Exception, SystemExit) as e: + _write_result({'success': False, 'message': str(e)}) + return + + line = sys.stdin.readline() + if not line: + _write_result({'success': False, 'message': 'Missing task input'}) + return + try: + msg = json.loads(line) + except Exception as e: + _write_result({'success': False, 'message': f'Invalid JSON input: {e}'}) + return + _run_task(msg) + + +def main(): + line = sys.stdin.readline() + if not line: + _write_result({'success': False, 'message': 'Missing task input'}) + return + try: + msg = json.loads(line) + except Exception as e: + _write_result({'success': False, 'message': f'Invalid JSON input: {e}'}) + return + if msg.get('type') == 'init': + _run_warm_worker(msg) + return + _run_task(msg) + + +if __name__ == '__main__': + main() diff --git a/projects/code-sandbox/src/isolated/python-isolated-runner.ts b/projects/code-sandbox/src/isolated/python-isolated-runner.ts new file mode 100644 index 000000000000..2e498b7198b2 --- /dev/null +++ b/projects/code-sandbox/src/isolated/python-isolated-runner.ts @@ -0,0 +1,617 @@ +import { spawn, type ChildProcess } from 'child_process'; +import { createInterface } from 'readline'; +import { basename, dirname, join } from 'path'; +import { fileURLToPath } from 'url'; +import { + chmodSync, + chownSync, + existsSync, + mkdirSync, + mkdtempSync, + readdirSync, + rmSync, + statSync +} from 'fs'; +import { tmpdir } from 'os'; +import { env, RUNTIME_MEMORY_OVERHEAD_MB } from '../env'; +import type { ExecuteOptions, ExecuteResult } from '../types'; +import { Semaphore } from '../utils/semaphore'; +import { getErrText } from '../utils'; +import { getLogger, LogCategories } from '../utils/logger'; +import { + getProcessTreeRSSMB, + killProcessTree, + PROCESS_GROUP_SUPPORTED +} from '../utils/process-tree'; +import { + runSandboxHttpRequest, + type SandboxHttpRequestPayload, + type SandboxHttpState +} from '../utils/sandbox-http'; +import { + assertPythonNativeIsolationReady, + getBundledPythonNativeLibraryPath, + PYTHON_ENABLE_NETWORK_SYSCALLS, + PYTHON_SANDBOX_GID, + PYTHON_SANDBOX_ROOT, + PYTHON_SANDBOX_UID, + shouldEnablePythonNativeIsolation +} from './python-isolation-config'; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const BOOTSTRAP_SCRIPT = join(__dirname, 'python-bootstrap.py'); +const NATIVE_SANDBOX_LIBRARY = getBundledPythonNativeLibraryPath(__dirname); +const RSS_POLL_INTERVAL = 500; +const TMP_USAGE_POLL_INTERVAL = 500; +const PYTHON_TASK_MATPLOTLIB_DIR = 'matplotlib'; +const PYTHON_TASK_MATPLOTLIB_CACHE_DIR = join(PYTHON_TASK_MATPLOTLIB_DIR, 'cache'); +const PYTHON_TASK_MATPLOTLIB_CONFIG_DIR = join(PYTHON_TASK_MATPLOTLIB_DIR, 'config'); +const PYTHON_TASK_MATPLOTLIB_TMP_DIR = join(PYTHON_TASK_MATPLOTLIB_DIR, 'tmp'); +const serverLogger = getLogger(LogCategories.MODULE.SANDBOX.SERVER); + +type RunningChild = { + proc: ChildProcess; + stderrBuf: string[]; + stdoutRl: ReturnType; + stderrRl: ReturnType; + taskTmpDir?: { + hostPath: string; + sandboxPath: string; + }; + lineHandler?: (line: string) => void; + closeHandler?: (code: number | null, signal: NodeJS.Signals | null) => void; + errorHandler?: (err: Error) => void; +}; + +/** + * PythonIsolatedRunner 使用 one-shot 预热池执行 Python 代码。 + * + * 预热进程只完成 bootstrap 和 native seccomp/chroot/降权初始化，尚未执行任何用户 + * 代码。每个预热进程最多接收一条任务，任务结束后销毁并异步补充新的干净进程， + * 避免用户代码污染后续任务。 + */ +export class PythonIsolatedRunner { + private readonly semaphore: Semaphore; + private readonly running = new Set(); + private readonly idleChildren = new Set(); + private readonly warmingChildren = new Set(); + private readonly warmIdleTarget: number; + private ready = false; + + constructor(private readonly maxConcurrency = env.SANDBOX_POOL_SIZE) { + this.semaphore = new Semaphore(maxConcurrency); + this.warmIdleTarget = maxConcurrency; + } + + async init(): Promise { + assertPythonNativeIsolationReady(NATIVE_SANDBOX_LIBRARY); + this.ready = true; + + try { + await this.replenishWarmChildren(true); + if (this.idleChildren.size < this.warmIdleTarget) { + throw new Error( + `Python isolated runner warmup failed: ready=${this.idleChildren.size}/${this.warmIdleTarget}` + ); + } + } catch (err) { + await this.shutdown(); + throw err; + } + + serverLogger.info( + `PythonIsolatedRunner ready: maxConcurrency=${this.maxConcurrency}, ` + + `warmIdleTarget=${this.warmIdleTarget}, nativeIsolation=${shouldEnablePythonNativeIsolation()}` + ); + } + + async shutdown(): Promise { + this.ready = false; + for (const child of [...this.running, ...this.idleChildren, ...this.warmingChildren]) { + killProcessTree(child.proc.pid); + this.cleanupChild(child); + } + this.running.clear(); + this.idleChildren.clear(); + this.warmingChildren.clear(); + } + + get stats() { + const semaphoreStats = this.semaphore.stats; + return { + total: this.running.size + this.idleChildren.size + this.warmingChildren.size, + idle: this.idleChildren.size, + busy: this.running.size, + warming: this.warmingChildren.size, + queued: semaphoreStats.queued, + poolSize: semaphoreStats.max, + ready: + this.ready && this.idleChildren.size + this.running.size + this.warmingChildren.size > 0 + }; + } + + async execute(options: ExecuteOptions): Promise { + const { code, variables } = options; + + if (!code || typeof code !== 'string' || !code.trim()) { + return { success: false, message: 'Code cannot be empty' }; + } + + await this.semaphore.acquire(); + try { + if (!this.ready) { + return { success: false, message: 'Python isolated runner is not ready' }; + } + return await this.executeOneShot({ code, variables: variables || {} }); + } finally { + this.semaphore.release(); + void this.replenishWarmChildren(); + } + } + + private async executeOneShot(task: { + code: string; + variables: Record; + }): Promise { + const child = this.takeIdleChild() ?? this.createChild(); + this.running.add(child); + return this.executeWithChild(child, task); + } + + private takeIdleChild(): RunningChild | undefined { + const child = this.idleChildren.values().next().value; + if (!child) return undefined; + this.idleChildren.delete(child); + if (child.closeHandler) child.proc.off('close', child.closeHandler); + if (child.errorHandler) child.proc.off('error', child.errorHandler); + child.closeHandler = undefined; + child.errorHandler = undefined; + return child; + } + + private createChild(): RunningChild { + const taskTmpDir = this.createTaskTmpDir(); + const proc = spawn('python3', ['-u', BOOTSTRAP_SCRIPT], { + stdio: ['pipe', 'pipe', 'pipe'], + detached: PROCESS_GROUP_SUPPORTED, + cwd: shouldEnablePythonNativeIsolation() ? PYTHON_SANDBOX_ROOT : undefined, + env: { + PATH: process.env.PATH || '/usr/local/bin:/usr/bin:/bin', + CHECK_INTERNAL_IP: String(env.CHECK_INTERNAL_IP), + PYTHONISOLATED: '1', + HOME: taskTmpDir.sandboxPath, + TMPDIR: taskTmpDir.sandboxPath, + FASTGPT_TASK_TMPDIR: taskTmpDir.sandboxPath, + MPLCONFIGDIR: `${taskTmpDir.sandboxPath}/${PYTHON_TASK_MATPLOTLIB_DIR}`, + XDG_CACHE_HOME: `${taskTmpDir.sandboxPath}/${PYTHON_TASK_MATPLOTLIB_CACHE_DIR}`, + XDG_CONFIG_HOME: `${taskTmpDir.sandboxPath}/${PYTHON_TASK_MATPLOTLIB_CONFIG_DIR}`, + MATPLOTLIB_TMPDIR: `${taskTmpDir.sandboxPath}/${PYTHON_TASK_MATPLOTLIB_TMP_DIR}`, + PYTHONDONTWRITEBYTECODE: '1', + // numpy/OpenBLAS may create worker threads while importing native extensions. + // Keep it single-threaded so seccomp does not need to allow clone/fork. + OPENBLAS_NUM_THREADS: '1', + OMP_NUM_THREADS: '1', + MKL_NUM_THREADS: '1', + NUMEXPR_NUM_THREADS: '1' + } + }); + const stdoutRl = createInterface({ input: proc.stdout!, terminal: false }); + const stderrRl = createInterface({ input: proc.stderr!, terminal: false }); + const child: RunningChild = { proc, stderrBuf: [], stdoutRl, stderrRl, taskTmpDir }; + + stderrRl.on('line', (line: string) => { + child.stderrBuf.push(line); + if (child.stderrBuf.length > 20) child.stderrBuf.shift(); + }); + + return child; + } + + /** + * 为 one-shot Python 子进程创建独立临时目录。 + * + * native chroot 模式下目录位于 sandbox root 的 /tmp 内，并 chown 给降权后的 sandbox + * 用户；非 Linux 本地开发模式下使用系统临时目录。父进程在 child cleanup 时递归删除， + * 覆盖超时/内存超限等 Python finally 无法执行的路径。 + */ + private createTaskTmpDir() { + const nativeIsolation = shouldEnablePythonNativeIsolation(); + const hostTmpRoot = nativeIsolation ? join(PYTHON_SANDBOX_ROOT, 'tmp') : tmpdir(); + if (!existsSync(hostTmpRoot)) { + mkdirSync(hostTmpRoot, { recursive: true }); + } + + const hostPath = mkdtempSync(join(hostTmpRoot, 'task-')); + const taskWritableDirs = [ + hostPath, + join(hostPath, PYTHON_TASK_MATPLOTLIB_DIR), + join(hostPath, PYTHON_TASK_MATPLOTLIB_CACHE_DIR), + join(hostPath, PYTHON_TASK_MATPLOTLIB_CONFIG_DIR), + join(hostPath, PYTHON_TASK_MATPLOTLIB_TMP_DIR) + ]; + for (const dir of taskWritableDirs) { + mkdirSync(dir, { recursive: true }); + } + + if (nativeIsolation) { + for (const dir of taskWritableDirs) { + chownSync(dir, PYTHON_SANDBOX_UID, PYTHON_SANDBOX_GID); + } + } + for (const dir of taskWritableDirs) { + chmodSync(dir, 0o700); + } + + return { + hostPath, + sandboxPath: nativeIsolation ? `/tmp/${basename(hostPath)}` : hostPath + }; + } + + private cleanupChild(child: RunningChild) { + try { + child.proc.stdin?.end(); + } catch {} + child.stdoutRl.close(); + child.stderrRl.close(); + if (child.lineHandler) child.stdoutRl.off('line', child.lineHandler); + if (child.closeHandler) child.proc.off('close', child.closeHandler); + if (child.errorHandler) child.proc.off('error', child.errorHandler); + child.proc.removeAllListeners(); + if (child.taskTmpDir) { + try { + rmSync(child.taskTmpDir.hostPath, { recursive: true, force: true }); + } catch (err) { + serverLogger.warn(`Failed to remove python task tmp dir: ${getErrText(err)}`); + } + } + this.running.delete(child); + this.idleChildren.delete(child); + this.warmingChildren.delete(child); + } + + private markChildIdle(child: RunningChild) { + child.closeHandler = (code, signal) => { + const stderr = child.stderrBuf.length > 0 ? ` | stderr: ${child.stderrBuf.join('\n')}` : ''; + serverLogger.warn( + `Python warm idle child exited (exit code: ${code}, signal: ${signal})${stderr}` + ); + this.cleanupChild(child); + void this.replenishWarmChildren(); + }; + child.errorHandler = (err) => { + serverLogger.warn(`Python warm idle child error: ${getErrText(err)}`); + this.cleanupChild(child); + void this.replenishWarmChildren(); + }; + child.proc.once('close', child.closeHandler); + child.proc.once('error', child.errorHandler); + this.idleChildren.add(child); + } + + private buildIsolationPayload() { + return { + enableSeccomp: shouldEnablePythonNativeIsolation(), + enableNetwork: PYTHON_ENABLE_NETWORK_SYSCALLS, + libraryPath: NATIVE_SANDBOX_LIBRARY, + sandboxRoot: PYTHON_SANDBOX_ROOT, + uid: PYTHON_SANDBOX_UID, + gid: PYTHON_SANDBOX_GID + }; + } + + private async replenishWarmChildren(waitForReady = false): Promise { + if (!this.ready || this.warmIdleTarget <= 0) return; + + const promises: Promise[] = []; + while (this.idleChildren.size + this.warmingChildren.size < this.warmIdleTarget) { + const child = this.createChild(); + this.warmingChildren.add(child); + promises.push(this.prepareWarmChild(child)); + } + + if (waitForReady && promises.length > 0) { + await Promise.all(promises); + } + } + + private prepareWarmChild(child: RunningChild): Promise { + return new Promise((resolve) => { + let settled = false; + const settle = (ready: boolean) => { + if (settled) return; + settled = true; + if (lineHandler) child.stdoutRl.off('line', lineHandler); + if (closeHandler) child.proc.off('close', closeHandler); + if (errorHandler) child.proc.off('error', errorHandler); + this.warmingChildren.delete(child); + if (ready && this.ready) { + this.markChildIdle(child); + } else { + killProcessTree(child.proc.pid); + this.cleanupChild(child); + } + resolve(); + }; + + const lineHandler = (line: string) => { + try { + const msg = JSON.parse(line); + if (msg.type === 'ready') { + settle(true); + return; + } + if (msg.type === 'result') { + serverLogger.warn(`Python warm child failed: ${msg.message || line}`); + } else { + serverLogger.warn(`Unexpected python warm child message: ${line}`); + } + } catch { + serverLogger.warn(`Invalid python warm child response: ${line}`); + } + settle(false); + }; + + const closeHandler = (code: number | null, signal: NodeJS.Signals | null) => { + const stderr = child.stderrBuf.length > 0 ? ` | stderr: ${child.stderrBuf.join('\n')}` : ''; + serverLogger.warn( + `Python warm child exited before ready (exit code: ${code}, signal: ${signal})${stderr}` + ); + settle(false); + }; + + const errorHandler = (err: Error) => { + serverLogger.warn(`Python warm child spawn error: ${getErrText(err)}`); + settle(false); + }; + + child.stdoutRl.on('line', lineHandler); + child.proc.once('close', closeHandler); + child.proc.once('error', errorHandler); + + try { + child.proc.stdin!.write( + JSON.stringify({ + type: 'init', + isolation: this.buildIsolationPayload() + }) + '\n' + ); + } catch (err) { + serverLogger.warn(`Python warm child communication error: ${getErrText(err)}`); + settle(false); + } + }); + } + + private executeWithChild( + child: RunningChild, + task: { code: string; variables: Record } + ): Promise { + return new Promise((resolve) => { + const timeoutMs = env.SANDBOX_MAX_TIMEOUT; + const proc = child.proc; + const stdoutRl = child.stdoutRl; + + let settled = false; + let outputBytes = 0; + let rssTimer: ReturnType | undefined; + let tmpUsageTimer: ReturnType | undefined; + const httpState: SandboxHttpState = { requestCount: 0 }; + const httpLimits = { + maxRequests: env.SANDBOX_REQUEST_MAX_COUNT, + timeoutMs: env.SANDBOX_REQUEST_TIMEOUT, + maxResponseSize: env.SANDBOX_REQUEST_MAX_RESPONSE_MB * 1024 * 1024, + maxRequestBodySize: env.SANDBOX_REQUEST_MAX_BODY_MB * 1024 * 1024 + }; + + const cleanup = () => { + clearTimeout(timer); + if (rssTimer) clearInterval(rssTimer); + if (tmpUsageTimer) clearInterval(tmpUsageTimer); + this.cleanupChild(child); + }; + + const settle = (result: ExecuteResult, opts: { kill?: boolean } = {}) => { + if (settled) return; + settled = true; + if (opts.kill) { + killProcessTree(proc.pid); + } + cleanup(); + resolve(result); + }; + + child.lineHandler = (line: string) => { + outputBytes += Buffer.byteLength(line, 'utf8') + 1; + const maxOutputBytes = env.SANDBOX_MAX_OUTPUT_MB * 1024 * 1024; + if (outputBytes > maxOutputBytes) { + settle( + { success: false, message: `Output too large (limit: ${maxOutputBytes} bytes)` }, + { kill: true } + ); + return; + } + + try { + const msg = JSON.parse(line); + if (msg.type === 'http_request') { + this.handleHttpRequestMessage({ + proc, + id: msg.id, + payload: msg.payload, + httpState, + httpLimits + }).catch((err) => { + serverLogger.warn( + `PythonIsolatedRunner http_request handler failed: ${getErrText(err)}` + ); + }); + return; + } + if (msg.type === 'result') { + delete msg.type; + settle(msg as ExecuteResult); + return; + } + settle( + { success: false, message: `Unknown python runner message: ${msg.type || line}` }, + { kill: true } + ); + } catch { + settle( + { success: false, message: `Invalid python runner response: ${line}` }, + { kill: true } + ); + } + }; + stdoutRl.on('line', child.lineHandler); + + child.errorHandler = (err) => { + settle({ success: false, message: `Python runner spawn error: ${getErrText(err)}` }); + }; + proc.once('error', child.errorHandler); + + // `exit` may fire before stdout/stderr streams are fully drained. Fast + // Python tasks can therefore exit with code 0 while the JSON result line + // is still buffered in the parent process. Wait for `close`, which is + // emitted after stdio streams are closed, before declaring "no result". + child.closeHandler = (code, signal) => { + if (settled) return; + const stderr = child.stderrBuf.length > 0 ? ` | stderr: ${child.stderrBuf.join('\n')}` : ''; + settle({ + success: false, + message: `Python runner exited before result (exit code: ${code}, signal: ${signal})${stderr}` + }); + }; + proc.once('close', child.closeHandler); + + const timer = setTimeout(() => { + settle( + { success: false, message: `Script execution timed out after ${timeoutMs}ms` }, + { kill: true } + ); + }, timeoutMs + 2000); + + if (env.SANDBOX_MAX_MEMORY_MB > 0 && proc.pid) { + const limitMB = env.SANDBOX_MAX_MEMORY_MB + RUNTIME_MEMORY_OVERHEAD_MB; + rssTimer = setInterval(async () => { + if (settled || !proc.pid) return; + const rss = await getProcessTreeRSSMB(proc.pid); + if (rss !== null && rss > limitMB) { + settle( + { + success: false, + message: `Memory limit exceeded (RSS: ${Math.round(rss)}MB, limit: ${limitMB}MB)` + }, + { kill: true } + ); + } + }, RSS_POLL_INTERVAL); + } + + if (env.SANDBOX_MAX_TMP_MB > 0 && child.taskTmpDir) { + const limitBytes = env.SANDBOX_MAX_TMP_MB * 1024 * 1024; + tmpUsageTimer = setInterval(() => { + if (settled || !child.taskTmpDir) return; + const size = this.getDirectorySize(child.taskTmpDir.hostPath); + if (size !== null && size > limitBytes) { + settle( + { + success: false, + message: `Temporary file limit exceeded (size: ${Math.ceil( + size / 1024 / 1024 + )}MB, limit: ${env.SANDBOX_MAX_TMP_MB}MB)` + }, + { kill: true } + ); + } + }, TMP_USAGE_POLL_INTERVAL); + } + + const payload = { + code: task.code, + variables: task.variables, + timeoutMs, + allowedModules: env.SANDBOX_PYTHON_ALLOWED_MODULES, + requestLimits: { + maxRequests: env.SANDBOX_REQUEST_MAX_COUNT, + timeoutMs: env.SANDBOX_REQUEST_TIMEOUT, + maxResponseSize: env.SANDBOX_REQUEST_MAX_RESPONSE_MB * 1024 * 1024, + maxRequestBodySize: env.SANDBOX_REQUEST_MAX_BODY_MB * 1024 * 1024, + maxOutputSize: env.SANDBOX_MAX_OUTPUT_MB * 1024 * 1024 + }, + taskTmpDir: child.taskTmpDir?.sandboxPath, + isolation: { + ...this.buildIsolationPayload() + } + }; + + try { + proc.stdin!.write(JSON.stringify(payload) + '\n'); + } catch (err) { + settle( + { success: false, message: `Python runner communication error: ${getErrText(err)}` }, + { kill: true } + ); + } + }); + } + + private async handleHttpRequestMessage({ + proc, + id, + payload, + httpState, + httpLimits + }: { + proc: ChildProcess; + id: string; + payload: SandboxHttpRequestPayload; + httpState: SandboxHttpState; + httpLimits: { + maxRequests: number; + timeoutMs: number; + maxResponseSize: number; + maxRequestBodySize: number; + }; + }) { + const writeResponse = (response: Record) => { + if (!proc.stdin?.writable) return; + proc.stdin.write(JSON.stringify({ type: 'http_response', id, ...response }) + '\n'); + }; + + try { + const data = await runSandboxHttpRequest({ + payload, + limits: httpLimits, + state: httpState + }); + writeResponse({ success: true, payload: data }); + } catch (err) { + writeResponse({ success: false, message: getErrText(err, 'HTTP request failed') }); + } + } + + private getDirectorySize(root: string): number | null { + let total = 0; + const stack = [root]; + + try { + while (stack.length > 0) { + const current = stack.pop()!; + const stat = statSync(current); + if (stat.isDirectory()) { + for (const entry of readdirSync(current)) { + stack.push(join(current, entry)); + } + } else { + total += stat.size; + } + } + return total; + } catch (err) { + serverLogger.warn(`Failed to calculate python task tmp dir size: ${getErrText(err)}`); + return null; + } + } +} diff --git a/projects/code-sandbox/src/isolated/python-isolation-config.ts b/projects/code-sandbox/src/isolated/python-isolation-config.ts new file mode 100644 index 000000000000..7abd8cfeb9e2 --- /dev/null +++ b/projects/code-sandbox/src/isolated/python-isolation-config.ts @@ -0,0 +1,33 @@ +import { existsSync } from 'fs'; +import { platform } from 'os'; +import { join } from 'path'; + +export const PYTHON_SANDBOX_ROOT = '/tmp/fastgpt-python-sandbox'; +export const PYTHON_SANDBOX_UID = 65537; +export const PYTHON_SANDBOX_GID = 65537; +export const PYTHON_ENABLE_NETWORK_SYSCALLS = false; + +export function shouldEnablePythonNativeIsolation(): boolean { + return platform() === 'linux'; +} + +/** + * Python native 隔离是 Linux 多租户安全边界的一部分。 + * + * Linux 环境固定启用 seccomp/chroot/setuid，缺失 native 库或 chroot 根目录时 + * 直接 fail-closed；macOS/Windows 仅保留本地开发兼容路径，不声明具备 OS 隔离。 + */ +export function assertPythonNativeIsolationReady(libraryPath: string) { + if (!shouldEnablePythonNativeIsolation()) return; + + if (!existsSync(libraryPath)) { + throw new Error(`Python native sandbox library does not exist: ${libraryPath}`); + } + if (!existsSync(PYTHON_SANDBOX_ROOT)) { + throw new Error(`Python sandbox root does not exist: ${PYTHON_SANDBOX_ROOT}`); + } +} + +export function getBundledPythonNativeLibraryPath(dirname: string) { + return join(dirname, 'fastgpt_python_sandbox.so'); +} diff --git a/projects/code-sandbox/src/pool/python-process-pool.ts b/projects/code-sandbox/src/pool/python-process-pool.ts deleted file mode 100644 index 67e32db3d5d4..000000000000 --- a/projects/code-sandbox/src/pool/python-process-pool.ts +++ /dev/null @@ -1,34 +0,0 @@ -/** - * PythonProcessPool - Python 子进程池 - * - * 继承 BaseProcessPool，提供 Python worker 的 spawn 配置。 - */ -import { dirname, join } from 'path'; -import { fileURLToPath } from 'url'; -import { env } from '../env'; -import { BaseProcessPool } from './base-process-pool'; - -const __dirname = dirname(fileURLToPath(import.meta.url)); -const WORKER_SCRIPT = join(__dirname, 'worker.py'); -const RECYCLE_AFTER_TASK_MODULES = new Set([ - 'subprocess', - 'multiprocessing', - 'threading', - 'concurrent' -]); - -function shouldRecycleAfterTask(allowedModules: readonly string[]): boolean { - return allowedModules.some((moduleName) => RECYCLE_AFTER_TASK_MODULES.has(moduleName)); -} - -export class PythonProcessPool extends BaseProcessPool { - constructor(poolSize?: number) { - super(poolSize, { - name: 'Python', - workerScript: WORKER_SCRIPT, - spawnCommand: (script) => `exec python3 -u ${script}`, - allowedModules: env.SANDBOX_PYTHON_ALLOWED_MODULES, - recycleAfterTask: shouldRecycleAfterTask(env.SANDBOX_PYTHON_ALLOWED_MODULES) - }); - } -} diff --git a/projects/code-sandbox/src/pool/worker.py b/projects/code-sandbox/src/pool/worker.py deleted file mode 100644 index 8a6777a721af..000000000000 --- a/projects/code-sandbox/src/pool/worker.py +++ /dev/null @@ -1,673 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -Python Worker 长驻进程 - 循环接收任务执行 - -协议： - 第 1 行 stdin: {"type":"init","allowedModules":["math","json",...]} - 后续每行 stdin: {"code":"...","variables":{},"timeoutMs":10000} - 每行 stdout: {"success":true,"data":{...}} 或 {"success":false,"message":"..."} -""" -import json -import sys -import copy as _copy -import hashlib as _hashlib -import hmac as _hmac -import base64 as _base64 -import urllib.parse as _urllib_parse -import urllib.request as _urllib_request -import socket as _socket -import ipaddress as _ipaddress -import time as _time -import math as _math -import inspect as _inspect_mod -import ast as _ast -import signal -import traceback as _tb -import sysconfig as _sysconfig -import builtins as _builtins -import types as _types -import encodings.idna as _encodings_idna # 预加载，避免 codec lazy load 被沙盒拦截 - -# stdlib 模块名集合，用于 _safe_import 快速放行 -_STDLIB_MODULES = sys.stdlib_module_names if hasattr(sys, 'stdlib_module_names') else frozenset() - -# 危险的 stdlib 模块，即使是 stdlib 也不允许用户代码直接 import -_DANGEROUS_STDLIB = frozenset({ - 'os', 'subprocess', 'shutil', 'pathlib', 'glob', 'tempfile', - 'multiprocessing', 'threading', 'concurrent', - 'ctypes', 'importlib', 'runpy', 'code', 'codeop', 'compileall', - 'socket', 'http', 'urllib', 'ftplib', 'smtplib', 'poplib', 'imaplib', - 'xmlrpc', 'socketserver', 'ssl', 'asyncio', 'selectors', 'select', - 'signal', 'resource', 'pty', 'termios', 'tty', 'fcntl', - 'mmap', 'dbm', 'sqlite3', 'shelve', - 'webbrowser', 'turtle', 'tkinter', 'idlelib', - 'venv', 'ensurepip', 'pip', 'site', - 'gc', 'sys', 'builtins', 'marshal', 'pickle', -}) - -# ===== 网络安全 ===== -_BLOCKED_CIDRS = [ - _ipaddress.ip_network('10.0.0.0/8'), - _ipaddress.ip_network('172.16.0.0/12'), - _ipaddress.ip_network('192.168.0.0/16'), - _ipaddress.ip_network('169.254.0.0/16'), - _ipaddress.ip_network('127.0.0.0/8'), - _ipaddress.ip_network('0.0.0.0/8'), - _ipaddress.ip_network('::1/128'), - _ipaddress.ip_network('fc00::/7'), - _ipaddress.ip_network('fe80::/10'), -] - -_REQUEST_LIMITS = { - 'max_requests': 30, - 'timeout': 60, - 'max_response_size': 10 * 1024 * 1024, - 'max_request_body_size': 5 * 1024 * 1024, - 'max_output_size': 10 * 1024 * 1024, - 'allowed_protocols': ['http:', 'https:'] -} - - -def _init_request_limits(limits): - """从 init 消息更新请求限制""" - if not limits: - return - if 'maxRequests' in limits: - _REQUEST_LIMITS['max_requests'] = limits['maxRequests'] - if 'timeoutMs' in limits: - _REQUEST_LIMITS['timeout'] = max(1, limits['timeoutMs'] // 1000) - if 'maxResponseSize' in limits: - _REQUEST_LIMITS['max_response_size'] = limits['maxResponseSize'] - if 'maxRequestBodySize' in limits: - _REQUEST_LIMITS['max_request_body_size'] = limits['maxRequestBodySize'] - if 'maxOutputSize' in limits: - _REQUEST_LIMITS['max_output_size'] = limits['maxOutputSize'] - - -def _is_blocked_ip(ip_str): - try: - addr = _ipaddress.ip_address(ip_str) - for net in _BLOCKED_CIDRS: - if addr in net: - return True - except ValueError: - return True - return False - - -# ===== DNS-pinned HTTP opener（防止 DNS rebinding）===== -import http.client as _http_client - -class _PinnedHTTPConnection(_http_client.HTTPConnection): - """强制连接到预解析的 IP，防止 DNS rebinding TOCTOU""" - def __init__(self, *args, pinned_ip=None, pinned_port=None, **kwargs): - super().__init__(*args, **kwargs) - self._pinned_ip = pinned_ip - self._pinned_port = pinned_port - - def connect(self): - self.sock = _socket.create_connection( - (self._pinned_ip or self.host, self._pinned_port or self.port), - self.timeout - ) - -class _PinnedHTTPSConnection(_http_client.HTTPSConnection): - def __init__(self, *args, pinned_ip=None, pinned_port=None, original_hostname=None, **kwargs): - super().__init__(*args, **kwargs) - self._pinned_ip = pinned_ip - self._pinned_port = pinned_port - self._original_hostname = original_hostname or self.host - - def connect(self): - import ssl as _ssl - self.sock = _socket.create_connection( - (self._pinned_ip or self.host, self._pinned_port or self.port), - self.timeout - ) - ctx = _ssl.create_default_context() - self.sock = ctx.wrap_socket(self.sock, server_hostname=self._original_hostname) - -class _PinnedHTTPHandler(_urllib_request.HTTPHandler): - def __init__(self, pinned_ip, pinned_port): - super().__init__() - self._pinned_ip = pinned_ip - self._pinned_port = pinned_port - - def http_open(self, req): - return self.do_open( - lambda host, **kwargs: _PinnedHTTPConnection( - host, pinned_ip=self._pinned_ip, pinned_port=self._pinned_port, **kwargs - ), - req - ) - -class _PinnedHTTPSHandler(_urllib_request.HTTPSHandler): - def __init__(self, pinned_ip, pinned_port, original_hostname): - super().__init__() - self._pinned_ip = pinned_ip - self._pinned_port = pinned_port - self._original_hostname = original_hostname - - def https_open(self, req): - return self.do_open( - lambda host, **kwargs: _PinnedHTTPSConnection( - host, pinned_ip=self._pinned_ip, pinned_port=self._pinned_port, - original_hostname=self._original_hostname, **kwargs - ), - req - ) - -def _build_pinned_opener(resolved_ip, port, hostname): - return _urllib_request.build_opener( - _PinnedHTTPHandler(resolved_ip, port), - _PinnedHTTPSHandler(resolved_ip, port, hostname) - ) - - -class _SystemHelper: - """安全的系统辅助类 — 所有模块引用通过闭包捕获，外部无法访问""" - __slots__ = () - - @staticmethod - def http_request(url, method='GET', headers=None, body=None, timeout=None): - global _request_count - _request_count += 1 - if _request_count > _REQUEST_LIMITS['max_requests']: - raise RuntimeError(f"Request limit exceeded: max {_REQUEST_LIMITS['max_requests']}") - - parsed = _urllib_parse.urlparse(url) - if parsed.scheme + ':' not in _REQUEST_LIMITS['allowed_protocols']: - raise RuntimeError(f"Protocol {parsed.scheme}: not allowed") - - hostname = parsed.hostname - try: - infos = _socket.getaddrinfo(hostname, parsed.port or (443 if parsed.scheme == 'https' else 80)) - for info in infos: - ip = info[4][0] - if _is_blocked_ip(ip): - raise RuntimeError("Request to private/internal network not allowed") - resolved_ip = infos[0][4][0] if infos else None - except _socket.gaierror as e: - raise RuntimeError(f"DNS resolution failed: {e}") - - if timeout is None: - timeout = _REQUEST_LIMITS['timeout'] - else: - timeout = min(timeout, _REQUEST_LIMITS['timeout']) - - if headers is None: - headers = {} - - data = None - if body is not None: - if isinstance(body, dict): - data = json.dumps(body).encode('utf-8') - if 'Content-Type' not in headers and 'content-type' not in headers: - headers['Content-Type'] = 'application/json' - elif isinstance(body, str): - data = body.encode('utf-8') - else: - data = body - - if data is not None and len(data) > _REQUEST_LIMITS['max_request_body_size']: - raise RuntimeError("Request body too large") - - # 使用自定义 handler 强制连接到预解析的 IP，防止 DNS rebinding - port = parsed.port or (443 if parsed.scheme == 'https' else 80) - opener = _build_pinned_opener(resolved_ip, port, hostname) - req = _urllib_request.Request(url, data=data, headers=headers, method=method.upper()) - try: - resp = opener.open(req, timeout=timeout) - resp_data = resp.read(_REQUEST_LIMITS['max_response_size'] + 1) - if len(resp_data) > _REQUEST_LIMITS['max_response_size']: - raise RuntimeError("Response too large") - return { - 'status': resp.status, - 'headers': dict(resp.headers), - 'data': resp_data.decode('utf-8', errors='replace') - } - except _urllib_request.URLError as e: - raise RuntimeError(f"HTTP request failed: {e}") - - # 驼峰别名，与 JS 端 SystemHelper API 保持一致 - httpRequest = http_request - - -system_helper = _SystemHelper() -SystemHelper = system_helper - -# Legacy global functions (backward compatibility, standalone) -def count_token(text): - if not isinstance(text, str): - text = str(text) - return _math.ceil(len(text) / 4) - -def str_to_base64(text, prefix=''): - b64 = _base64.b64encode(text.encode('utf-8')).decode('utf-8') - return prefix + b64 - -def create_hmac(algorithm, secret): - timestamp = str(int(_time.time() * 1000)) - string_to_sign = timestamp + '\n' + secret - h = _hmac.new(secret.encode('utf-8'), string_to_sign.encode('utf-8'), algorithm) - sign = _urllib_parse.quote(_base64.b64encode(h.digest()).decode('utf-8')) - return {"timestamp": timestamp, "sign": sign} - -def delay(ms): - if ms > 10000: - raise ValueError("Delay must be <= 10000ms") - _time.sleep(ms / 1000) - -_request_count = 0 - -# ===== __import__ 拦截（init 后设置）===== -_original_import = _builtins.__import__ -_allowed_modules = set() - -_stdlib_paths = [] -_site_packages_paths = [] -for _key in ('stdlib', 'platstdlib'): - _p = _sysconfig.get_path(_key) - if _p: - _stdlib_paths.append(_p) -for _key in ('purelib', 'platlib'): - _p = _sysconfig.get_path(_key) - if _p: - _stdlib_paths.append(_p) - _site_packages_paths.append(_p) -_stdlib_paths.append('', '', '') - - -def _safe_import(name, *args, **kwargs): - # 重入保护：避免 extract_stack / _original_import 内部触发 import 时无限递归 - if _import_guard: - return _original_import(name, *args, **kwargs) - top_level = name.split('.')[0] - # 拦截 builtins 模块，返回代理 - if top_level == 'builtins' and _builtins_proxy is not None: - return _builtins_proxy - # 安全的 stdlib 模块直接放行（不含危险模块） - if top_level in _STDLIB_MODULES and top_level not in _DANGEROUS_STDLIB: - return _original_import(name, *args, **kwargs) - # 在白名单中的模块直接放行 - if top_level in _allowed_modules: - return _original_import(name, *args, **kwargs) - # 不在白名单中的模块（含危险 stdlib）：检查是否由用户代码直接触发 - # 只拦截直接调用者是用户代码的情况（//） - # stdlib 内部的间接 import（如 locale -> os）放行 - if _is_direct_user_import_call(): - raise ImportError(f"Module '{name}' is not in the allowlist.") - return _original_import(name, *args, **kwargs) - - -# ===== 文件系统限制 ===== -_original_open = open - -_open_guard = False - -def _validate_user_code(code: str): - """执行前的静态安全检查，阻断已知高危反射链。""" - try: - tree = _ast.parse(code) - except SyntaxError: - # 语法错误交由后续 exec 抛出原始错误信息 - return - - for node in _ast.walk(tree): - # 直接属性访问：obj.__subclasses__ - if isinstance(node, _ast.Attribute) and node.attr == '__subclasses__': - raise RuntimeError("Access to __subclasses__ is not allowed in sandbox") - - # 动态属性访问：getattr(obj, '__subclasses__') - if ( - isinstance(node, _ast.Call) - and isinstance(node.func, _ast.Name) - and node.func.id == 'getattr' - and len(node.args) >= 2 - and isinstance(node.args[1], _ast.Constant) - and node.args[1].value == '__subclasses__' - ): - raise RuntimeError("Access to __subclasses__ is not allowed in sandbox") - - - -def _restricted_open(*args, **kwargs): - """限制 open() — 只允许第三方库内部调用，禁止用户代码直接读写文件""" - global _open_guard - if _open_guard: - return _original_open(*args, **kwargs) - _open_guard = True - try: - stack = _tb.extract_stack() - finally: - _open_guard = False - if len(stack) >= 2: - caller_fn = stack[-2].filename or '' - # 用户代码（）不允许直接 open - if caller_fn in ('', '', ''): - raise PermissionError("File system access is not allowed in sandbox") - # 非 stdlib、非 site-packages、非 worker 自身的帧也不允许 - if not _is_stdlib_frame(caller_fn) and not _is_site_packages_frame(caller_fn) and caller_fn != __file__: - raise PermissionError("File system access is not allowed in sandbox") - return _original_open(*args, **kwargs) - - -# ===== 日志收集 ===== -_logs = [] -_log_size = 0 -_MAX_LOG_SIZE = 1024 * 1024 # 1MB -_orig_print = print - - -def _safe_print(*args, **kwargs): - global _log_size - line = ' '.join(str(a) for a in args) - if _log_size + len(line) <= _MAX_LOG_SIZE: - _logs.append(line) - _log_size += len(line) - - -# ===== 输出 ===== -def write_line(obj): - try: - payload = json.dumps(obj, ensure_ascii=False, default=str) - except Exception as e: - payload = json.dumps({ - "success": False, - "message": f"Failed to serialize output: {e}", - "workerRecycle": "output_serialize" - }, ensure_ascii=False) - - if len(payload.encode('utf-8')) > _REQUEST_LIMITS['max_output_size']: - payload = json.dumps({ - "success": False, - "message": f"Output too large (limit: {_REQUEST_LIMITS['max_output_size']} bytes)", - "workerRecycle": "output_limit" - }, ensure_ascii=False) - - sys.stdout.write(payload + '\n') - sys.stdout.flush() - - -# ===== 超时信号处理 ===== -_timeout_stage = 0 # 0=未触发, 1=第一次(可恢复), 2=第二次(强制退出) - -def _timeout_handler(signum, frame): - global _timeout_stage - _timeout_stage += 1 - if _timeout_stage >= 2: - # 第二次 alarm：用户代码 catch 了第一次 TimeoutError，强制写错误并退出当前执行 - # 通过 SystemExit 强制终止（不可被 except Exception 捕获） - raise SystemExit("Script execution timed out (forced)") - # 第一次 alarm：设置 1s 后的兜底 alarm - signal.alarm(1) - raise TimeoutError("Script execution timed out") - - -# ===== 模块状态保护 ===== -# 用户代码可能污染共享模块（如 json.dumps = lambda x: "hacked"）， -# 需要在每次执行前快照、执行后恢复。 -_PROTECTED_MODULES = [json, _math, _time, _base64, _hashlib, _hmac, _copy] - - -def _snapshot_modules(): - """保存受保护模块的属性快照""" - snapshots = [] - for mod in _PROTECTED_MODULES: - attrs = {} - for name in dir(mod): - if not name.startswith('__'): - try: - attrs[name] = getattr(mod, name) - except Exception: - pass - snapshots.append((mod, attrs)) - return snapshots - - -def _restore_modules(snapshots): - """恢复受保护模块的属性""" - for mod, attrs in snapshots: - # 删除用户可能添加的新属性 - current_names = set(n for n in dir(mod) if not n.startswith('__')) - original_names = set(attrs.keys()) - for name in current_names - original_names: - try: - delattr(mod, name) - except Exception: - pass - # 恢复原始属性 - for name, val in attrs.items(): - try: - setattr(mod, name, val) - except Exception: - pass - - -# ===== 主循环 ===== -def main_loop(): - global _allowed_modules, _request_count, _logs, _log_size, _timeout_stage - - initialized = False - - for line in sys.stdin: - line = line.strip() - if not line: - continue - - try: - msg = json.loads(line) - except json.JSONDecodeError: - write_line({"success": False, "message": "Invalid JSON input"}) - continue - - # 初始化 - if not initialized: - if msg.get('type') == 'init': - _allowed_modules = set(msg.get('allowedModules', [])) - _init_request_limits(msg.get('requestLimits')) - _builtins.__import__ = _safe_import - global _builtins_proxy - _builtins_proxy = _BuiltinsProxy(_builtins) - write_line({"type": "ready"}) - initialized = True - else: - write_line({"success": False, "message": "Expected init message"}) - continue - - # ping 健康检查：立即回复 pong - if msg.get('type') == 'ping': - write_line({"type": "pong"}) - continue - - # 执行任务 - code = msg.get('code', '') - variables = msg.get('variables', {}) - timeout_ms = msg.get('timeoutMs', 10000) - # 修复精度：向上取整而非截断，最小 1 秒 - timeout_s = max(1, -(-timeout_ms // 1000)) # ceil division - - _request_count = 0 - _logs = [] - _log_size = 0 - _timeout_stage = 0 - - # 替换 print - _builtins.print = _safe_print - - # 每次执行前强制恢复 __import__，防止上次用户代码篡改 - _builtins.__import__ = _safe_import - - # 保存模块状态快照 - _mod_snapshots = _snapshot_modules() - - try: - # 设置超时（双重 alarm：第一次抛 TimeoutError，第二次兜底防止用户 catch） - signal.signal(signal.SIGALRM, _timeout_handler) - signal.alarm(timeout_s) - - # 构建受限的 __builtins__ 字典，移除 _original_import 等内部引用 - _safe_builtins = {} - for _name in dir(_builtins): - if _name.startswith('_') and _name not in ( - '__name__', '__doc__', '__import__', '__build_class__', - ): - continue - _safe_builtins[_name] = getattr(_builtins, _name) - # 确保 __import__ 指向安全版本 - _safe_builtins['__import__'] = _safe_import - _safe_builtins['__build_class__'] = _builtins.__build_class__ - # 限制 open() — 禁止用户代码直接读写文件系统 - _safe_builtins['open'] = _restricted_open - - # H3: 屏蔽 object.__subclasses__，防止通过子类树找到已加载的危险模块 - class _SafeObject(object): - __subclasses__ = None - _safe_builtins['object'] = _SafeObject - - # 构建执行环境 - exec_globals = { - '__builtins__': _safe_builtins, - 'variables': variables, - 'SystemHelper': system_helper, - 'system_helper': system_helper, - 'count_token': count_token, - 'str_to_base64': str_to_base64, - 'create_hmac': create_hmac, - 'delay': delay, - 'http_request': system_helper.http_request, - 'print': _safe_print, - 'json': json, - 'math': _math, - 'time': _time, - } - # 展开 variables 到全局（过滤保留关键字，防止覆盖沙箱安全全局变量） - _reserved_keys = frozenset(exec_globals.keys()) - for k, v in variables.items(): - if k not in _reserved_keys: - exec_globals[k] = v - - # 执行前静态检查：阻断 __subclasses__ 反射链 - _validate_user_code(code) - - # 执行用户代码 - exec(code, exec_globals) - - # 取出 main 函数 - user_main = exec_globals.get('main') - if user_main is None: - raise RuntimeError("No 'main' function defined") - - # 调用 main - sig = _inspect_mod.signature(user_main) - params = list(sig.parameters.keys()) - - if len(params) == 0: - result = user_main() - elif len(params) == 1: - p = params[0] - # If param name matches a variable key, pass that value directly. - # Otherwise fall back to passing the entire variables dict. - if p in variables: - result = user_main(variables[p]) - else: - result = user_main(variables) - else: - # 用 kwargs 调用：缺席参数走函数默认值，不影响后面的参数 - for p in params: - if p not in variables and sig.parameters[p].default is _inspect_mod.Parameter.empty: - raise TypeError(f"Missing required argument: '{p}'") - call_kwargs = {p: variables[p] for p in params if p in variables} - result = user_main(**call_kwargs) - - signal.alarm(0) - write_line({ - "success": True, - "data": {"codeReturn": result, "log": '\n'.join(_logs)} - }) - - except (Exception, SystemExit) as e: - signal.alarm(0) - result = {"success": False, "message": str(e)} - if isinstance(e, TimeoutError) or ( - isinstance(e, SystemExit) and 'timed out' in str(e).lower() - ): - result["workerRecycle"] = "timeout" - write_line(result) - - finally: - signal.alarm(0) - _builtins.print = _orig_print - # 恢复模块状态，防止用户代码污染影响后续请求 - _restore_modules(_mod_snapshots) - - -if __name__ == '__main__': - main_loop() diff --git a/projects/code-sandbox/src/pool/worker.ts b/projects/code-sandbox/src/pool/worker.ts index 0ddfd0640a37..b7684ca96b36 100644 --- a/projects/code-sandbox/src/pool/worker.ts +++ b/projects/code-sandbox/src/pool/worker.ts @@ -369,8 +369,8 @@ const earlyDangerousMethods = [ 'initgroups' ]; -// 延迟删除：会被 https/dns/tsx 等内部使用，要等 hardenRuntime 预加载完白名单后再删 -const lateDangerousMethods = ['kill', 'exit', 'emitWarning', 'abort']; +// 延迟处理：会被 https/dns/tsx/url 等内部使用，要等 hardenRuntime 预加载完白名单后再收紧。 +const lateDangerousMethods = ['kill', 'exit', 'abort']; function deleteProcessMethods(methods: readonly string[]): void { for (const method of methods) { @@ -384,6 +384,18 @@ function deleteProcessMethods(methods: readonly string[]): void { } } +function stubProcessMethods(methods: readonly string[]): void { + for (const method of methods) { + try { + Object.defineProperty(process, method, { + value: () => undefined, + writable: false, + configurable: false + }); + } catch {} + } +} + if (typeof process !== 'undefined') { deleteProcessMethods(earlyDangerousMethods); @@ -466,15 +478,6 @@ const SystemHelper = { if (!REQUEST_LIMITS.allowedProtocols.includes(parsed.protocol)) { throw new Error('Protocol not allowed'); } - // 先检查 URL 是否指向内部地址 - if (await isInternalAddress(url)) { - throw new Error('Request to private network not allowed'); - } - const ips = await dnsResolve(parsed.hostname); - // 防 DNS rebinding TOCTOU：对真正用于建连的 IP 再次校验 - if (ips.length === 0 || ips.some((ip) => isInternalResolvedIP(ip))) { - throw new Error('Request to private network not allowed'); - } const method = (opts.method || 'GET').toUpperCase(); const headers = opts.headers || {}; const body = @@ -486,11 +489,26 @@ const SystemHelper = { if (body && body.length > REQUEST_LIMITS.maxRequestBodySize) { throw new Error('Request body too large'); } - const timeoutSeconds = - typeof opts.timeout === 'number' && Number.isFinite(opts.timeout) && opts.timeout > 0 - ? opts.timeout - : REQUEST_LIMITS.timeoutMs / 1000; - const timeout = Math.min(Math.ceil(timeoutSeconds * 1000), REQUEST_LIMITS.timeoutMs); + // 先完成协议和请求体大小校验，再解析网络目标，避免本地错误被外部网络状态掩盖。 + if (await isInternalAddress(url)) { + throw new Error('Request to private network not allowed'); + } + const ips = await dnsResolve(parsed.hostname); + // 防 DNS rebinding TOCTOU：对真正用于建连的 IP 再次校验 + if (ips.length === 0 || ips.some((ip) => isInternalResolvedIP(ip))) { + throw new Error('Request to private network not allowed'); + } + const timeout = + typeof opts.timeoutMs === 'number' && Number.isFinite(opts.timeoutMs) && opts.timeoutMs > 0 + ? Math.min(Math.ceil(opts.timeoutMs), REQUEST_LIMITS.timeoutMs) + : Math.min( + Math.ceil( + typeof opts.timeout === 'number' && Number.isFinite(opts.timeout) && opts.timeout > 0 + ? opts.timeout * 1000 + : REQUEST_LIMITS.timeoutMs + ), + REQUEST_LIMITS.timeoutMs + ); if (body && !headers['Content-Type'] && !headers['content-type']) { headers['Content-Type'] = 'application/json'; } @@ -559,9 +577,12 @@ function hardenRuntime(): void { } catch {} } - // 白名单模块已加载完毕，此时再删除 kill/exit/emitWarning/abort： + // 白名单模块已加载完毕，此时再删除 kill/exit/abort： // 这些方法仅在模块初始化时被 https/dns/tsx 等使用，预加载后不再需要。 deleteProcessMethods(lateDangerousMethods); + // Node 内置 url.parse 在运行时仍会调用 process.emitWarning。 + // 用户代码拿到的是下方 _sandboxProcess，这里只给真实 process 保留不可写 no-op 兼容内置模块。 + stubProcessMethods(['emitWarning']); for (const intrinsic of hardenedIntrinsics) { if (intrinsic) _ObjectFreeze(intrinsic); @@ -885,13 +906,10 @@ rl.on('line', async (line: string) => { })(); const timeoutPromise = new _OriginalPromise((_, reject) => { - timer = _workerSetTimeout( - () => { - timedOut = true; - reject(new _OriginalError(`Script execution timed out after ${timeoutMs}ms`)); - }, - timeoutMs || 10000 - ); + timer = _workerSetTimeout(() => { + timedOut = true; + reject(new _OriginalError(`Script execution timed out after ${timeoutMs}ms`)); + }, timeoutMs || 10000); }); const result = await _PromiseRace([resultPromise, timeoutPromise]); diff --git a/projects/code-sandbox/src/utils/ipCheck.util.ts b/projects/code-sandbox/src/utils/ipCheck.util.ts index 432bd5c85b7d..bceea01ee0f9 100644 --- a/projects/code-sandbox/src/utils/ipCheck.util.ts +++ b/projects/code-sandbox/src/utils/ipCheck.util.ts @@ -2,11 +2,15 @@ import { createInternalAddressChecker, PRIVATE_URL_TEXT } from '@fastgpt/global/common/system/network'; +import { env } from '../env'; const truthyEnvValues = new Set(['true', '1', 'yes', 'y']); const { isInternalAddress, isInternalResolvedIP } = createInternalAddressChecker({ - checkInternalIp: () => truthyEnvValues.has((process.env.CHECK_INTERNAL_IP || '').toLowerCase()) + checkInternalIp: () => { + const raw = process.env.CHECK_INTERNAL_IP; + return raw === undefined ? env.CHECK_INTERNAL_IP : truthyEnvValues.has(raw.toLowerCase()); + } }); export { isInternalAddress, isInternalResolvedIP, PRIVATE_URL_TEXT }; diff --git a/projects/code-sandbox/src/utils/process-tree.ts b/projects/code-sandbox/src/utils/process-tree.ts new file mode 100644 index 000000000000..4827fdf3524e --- /dev/null +++ b/projects/code-sandbox/src/utils/process-tree.ts @@ -0,0 +1,122 @@ +import { execFile, execFileSync } from 'child_process'; +import { platform } from 'os'; +import { readdirSync, readFileSync } from 'fs'; +import { promisify } from 'util'; + +const execFileAsync = promisify(execFile); +export const PROCESS_GROUP_SUPPORTED = platform() !== 'win32'; + +export function readLinuxRSSKB(pid: number): number | null { + try { + const status = readFileSync(`/proc/${pid}/status`, 'utf-8'); + const match = status.match(/VmRSS:\s+(\d+)\s+kB/); + return match ? parseInt(match[1], 10) : null; + } catch { + return null; + } +} + +export function getLinuxChildPids(pid: number): number[] { + const children = new Set(); + try { + const taskIds = readdirSync(`/proc/${pid}/task`); + for (const taskId of taskIds) { + try { + const content = readFileSync(`/proc/${pid}/task/${taskId}/children`, 'utf-8'); + for (const child of content.trim().split(/\s+/)) { + const childPid = Number(child); + if (Number.isInteger(childPid) && childPid > 0) { + children.add(childPid); + } + } + } catch {} + } + } catch {} + return [...children]; +} + +export function getChildPids(pid: number): number[] { + if (platform() === 'linux') { + return getLinuxChildPids(pid); + } + + try { + const stdout = execFileSync('pgrep', ['-P', String(pid)], { + encoding: 'utf-8', + timeout: 1000 + }); + return stdout + .split(/\s+/) + .map((value) => Number(value)) + .filter((value) => Number.isInteger(value) && value > 0); + } catch { + return []; + } +} + +export function getDescendantPids(rootPid: number): number[] { + const descendants: number[] = []; + const queue = [rootPid]; + const seen = new Set(queue); + + while (queue.length > 0) { + const parentPid = queue.shift()!; + for (const childPid of getChildPids(parentPid)) { + if (seen.has(childPid)) continue; + seen.add(childPid); + descendants.push(childPid); + queue.push(childPid); + } + } + + return descendants; +} + +export async function getProcessTreeRSSMB(pid: number): Promise { + const pids = [pid, ...getDescendantPids(pid)]; + if (pids.length === 0) return null; + + try { + if (platform() === 'linux') { + let totalKB = 0; + for (const currentPid of pids) { + const rssKB = readLinuxRSSKB(currentPid); + if (rssKB !== null) totalKB += rssKB; + } + return totalKB > 0 ? totalKB / 1024 : null; + } + + const { stdout } = await execFileAsync('ps', ['-o', 'rss=', '-p', pids.join(',')], { + timeout: 2000 + }); + const totalKB = stdout + .split('\n') + .map((line) => parseInt(line.trim(), 10)) + .filter((rssKB) => !isNaN(rssKB)) + .reduce((sum, rssKB) => sum + rssKB, 0); + return totalKB > 0 ? totalKB / 1024 : null; + } catch { + return null; + } +} + +export function killProcessTree(pid?: number): void { + if (!pid) return; + + const descendantPids = getDescendantPids(pid).reverse(); + for (const childPid of descendantPids) { + try { + process.kill(childPid, 'SIGKILL'); + } catch {} + } + + if (PROCESS_GROUP_SUPPORTED) { + try { + process.kill(-pid, 'SIGKILL'); + } catch {} + } + + try { + process.kill(pid, 'SIGKILL'); + } catch {} +} diff --git a/projects/code-sandbox/src/utils/sandbox-http.ts b/projects/code-sandbox/src/utils/sandbox-http.ts new file mode 100644 index 000000000000..c612644b9a29 --- /dev/null +++ b/projects/code-sandbox/src/utils/sandbox-http.ts @@ -0,0 +1,145 @@ +import http from 'http'; +import https from 'https'; +import { isIP } from 'net'; +import dns from 'dns/promises'; +import { isInternalAddress, isInternalResolvedIP } from './ipCheck.util'; + +export type SandboxHttpLimits = { + maxRequests: number; + timeoutMs: number; + maxResponseSize: number; + maxRequestBodySize: number; + allowedProtocols?: string[]; +}; + +export type SandboxHttpState = { + requestCount: number; +}; + +export type SandboxHttpRequestPayload = { + url: string; + method?: string; + headers?: Record; + body?: any; + timeout?: number; + timeoutMs?: number; +}; + +const dnsResolve = async (hostname: string) => { + const res = await dns.lookup(hostname, { all: true }); + return res.map((r) => r.address); +}; + +/** + * 执行受控 HTTP 请求。 + * + * 该函数用于 sandbox 代理出网，集中维护协议、SSRF、DNS pinning、请求次数、 + * 请求体大小、响应大小和超时限制。调用方应为每次代码执行创建独立 state。 + */ +export async function runSandboxHttpRequest({ + payload, + limits, + state +}: { + payload: SandboxHttpRequestPayload; + limits: SandboxHttpLimits; + state: SandboxHttpState; +}): Promise { + if (++state.requestCount > limits.maxRequests) { + throw new Error('Request limit exceeded'); + } + + const allowedProtocols = limits.allowedProtocols ?? ['http:', 'https:']; + const parsed = new URL(payload.url); + if (!allowedProtocols.includes(parsed.protocol)) { + throw new Error('Protocol not allowed'); + } + + const method = (payload.method || 'GET').toUpperCase(); + const headers = { ...(payload.headers || {}) }; + const body = + payload.body != null + ? typeof payload.body === 'string' + ? payload.body + : JSON.stringify(payload.body) + : null; + + if (body && Buffer.byteLength(body, 'utf8') > limits.maxRequestBodySize) { + throw new Error('Request body too large'); + } + + if (await isInternalAddress(payload.url)) { + throw new Error('Request to private network not allowed'); + } + + const ips = await dnsResolve(parsed.hostname); + if (ips.length === 0 || ips.some((ip) => isInternalResolvedIP(ip))) { + throw new Error('Request to private network not allowed'); + } + + const timeout = (() => { + if (typeof payload.timeoutMs === 'number' && Number.isFinite(payload.timeoutMs)) { + return Math.min(Math.ceil(payload.timeoutMs), limits.timeoutMs); + } + if ( + typeof payload.timeout === 'number' && + Number.isFinite(payload.timeout) && + payload.timeout > 0 + ) { + return Math.min(Math.ceil(payload.timeout * 1000), limits.timeoutMs); + } + return limits.timeoutMs; + })(); + + if (body && !headers['Content-Type'] && !headers['content-type']) { + headers['Content-Type'] = 'application/json'; + } + + const resolvedIP = ips[0]; + if (!headers['Host'] && !headers['host']) { + headers['Host'] = parsed.hostname + (parsed.port ? ':' + parsed.port : ''); + } + + const lib = parsed.protocol === 'https:' ? https : http; + + return new Promise((resolve, reject) => { + const req = lib.request( + { + method, + headers, + timeout, + hostname: resolvedIP, + port: parsed.port || (parsed.protocol === 'https:' ? 443 : 80), + path: parsed.pathname + parsed.search, + ...(isIP(parsed.hostname) ? {} : { servername: parsed.hostname }) + }, + (res: any) => { + const chunks: Buffer[] = []; + let size = 0; + res.on('data', (chunk: Buffer) => { + size += chunk.length; + if (size > limits.maxResponseSize) { + req.destroy(); + reject(new Error('Response too large')); + return; + } + chunks.push(chunk); + }); + res.on('end', () => { + const data = Buffer.concat(chunks).toString('utf-8'); + const h: Record = {}; + for (const [k, v] of Object.entries(res.headers)) h[k] = v; + resolve({ status: res.statusCode, statusText: res.statusMessage, headers: h, data }); + }); + res.on('error', reject); + } + ); + req.on('timeout', () => { + req.destroy(); + reject(new Error('Request timeout')); + }); + req.on('error', reject); + if (body) req.write(body); + req.end(); + }); +} diff --git a/projects/code-sandbox/test/compat/legacy-python.test.ts b/projects/code-sandbox/test/compat/legacy-python.test.ts index cf2a370c1e9f..317e4aadac64 100644 --- a/projects/code-sandbox/test/compat/legacy-python.test.ts +++ b/projects/code-sandbox/test/compat/legacy-python.test.ts @@ -1,266 +1,244 @@ /** * 旧版 Python 代码兼容性测试 * - * 验证新沙盒能正确执行旧版 FastGPT 生成的 Python 代码写法，包括： - * - main() 无参数（通过全局变量访问） - * - main(variables) 单参数字典 - * - main(a, b, c) 多参数展开 - * - print 被收集到 log（旧版直接移除） - * - 危险模块拦截 - * - 各种返回值类型 + * 旧 Python 长驻 worker 已被移除；这批旧代码用例必须继续跑过 + * PythonIsolatedRunner，证明新执行器可以接住历史 Python Code 节点。 */ -import { describe, it, expect, beforeAll, afterAll } from 'vitest'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; - -let pool: PythonProcessPool; - -beforeAll(async () => { - pool = new PythonProcessPool(1); - await pool.init(); -}); - -afterAll(async () => { - await pool.shutdown(); -}); - -describe('旧版 Python 兼容性', () => { - // ===== main 函数签名兼容 ===== - - it('main(variables) 单参数字典写法', async () => { - const result = await pool.execute({ - code: `def main(variables): +import { afterAll, beforeAll, describe, expect, it } from 'vitest'; +import type { ExecuteOptions, ExecuteResult } from '../../src/types'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; + +type PythonRunner = { + execute(options: ExecuteOptions): Promise; +}; + +function legacyPythonCompatibilitySuite(name: string, getRunner: () => PythonRunner) { + describe(name, () => { + it('main(variables) 单参数字典写法', async () => { + const result = await getRunner().execute({ + code: `def main(variables): return {"name": variables["name"], "age": variables["age"]}`, - variables: { name: 'FastGPT', age: 3 } + variables: { name: 'FastGPT', age: 3 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ name: 'FastGPT', age: 3 }); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual({ name: 'FastGPT', age: 3 }); - }); - it('main(a, b) 多参数展开写法', async () => { - const result = await pool.execute({ - code: `def main(a, b): + it('main(a, b) 多参数展开写法', async () => { + const result = await getRunner().execute({ + code: `def main(a, b): return {"sum": a + b}`, - variables: { a: 10, b: 20 } + variables: { a: 10, b: 20 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.sum).toBe(30); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.sum).toBe(30); - }); - it('main(a, b, c) 三参数展开', async () => { - const result = await pool.execute({ - code: `def main(a, b, c): + it('main(a, b, c) 三参数展开', async () => { + const result = await getRunner().execute({ + code: `def main(a, b, c): return {"result": a * b + c}`, - variables: { a: 3, b: 4, c: 5 } + variables: { a: 3, b: 4, c: 5 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.result).toBe(17); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.result).toBe(17); - }); - it('main() 无参数写法', async () => { - const result = await pool.execute({ - code: `def main(): + it('main() 无参数写法', async () => { + const result = await getRunner().execute({ + code: `def main(): return {"ok": True}`, - variables: {} + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.ok).toBe(true); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.ok).toBe(true); - }); - it('旧版写法：main 外部直接访问全局变量', async () => { - const result = await pool.execute({ - code: `prefix = name + "_suffix" + it('旧版写法：main 外部直接访问全局变量', async () => { + const result = await getRunner().execute({ + code: `prefix = name + "_suffix" def main(name, age): return {"result": prefix, "name": name, "age": age}`, - variables: { name: 'hello', age: 18 } + variables: { name: 'hello', age: 18 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.result).toBe('hello_suffix'); + expect(result.data?.codeReturn.name).toBe('hello'); + expect(result.data?.codeReturn.age).toBe(18); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.result).toBe('hello_suffix'); - expect(result.data?.codeReturn.name).toBe('hello'); - expect(result.data?.codeReturn.age).toBe(18); - }); - it('旧版写法：无参 main 通过全局变量访问', async () => { - const result = await pool.execute({ - code: `def main(): + it('旧版写法：无参 main 通过全局变量访问', async () => { + const result = await getRunner().execute({ + code: `def main(): return {"name": name, "age": age}`, - variables: { name: 'test', age: 25 } + variables: { name: 'test', age: 25 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ name: 'test', age: 25 }); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual({ name: 'test', age: 25 }); - }); - it('main 带默认参数', async () => { - const result = await pool.execute({ - code: `def main(name, greeting="Hello"): + it('main 带默认参数', async () => { + const result = await getRunner().execute({ + code: `def main(name, greeting="Hello"): return {"msg": f"{greeting}, {name}!"}`, - variables: { name: 'World' } + variables: { name: 'World' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.msg).toBe('Hello, World!'); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.msg).toBe('Hello, World!'); - }); - it('main 前置命名参数缺失但有默认值，后置参数仍按 kwargs 注入', async () => { - const result = await pool.execute({ - code: `def main(a=None, b=None): + it('main 前置命名参数缺失但有默认值，后置参数仍按 kwargs 注入', async () => { + const result = await getRunner().execute({ + code: `def main(a=None, b=None): return {"a": a if a else "无", "b": b if b else "无"}`, - variables: { b: 'd2' } + variables: { b: 'd2' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ a: '无', b: 'd2' }); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual({ a: '无', b: 'd2' }); - }); - // ===== 返回值类型兼容 ===== - - it('返回列表（旧版常见）', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('返回列表（旧版常见）', async () => { + const result = await getRunner().execute({ + code: `def main(variables): return [variables["a"], variables["b"], variables["a"] + variables["b"]]`, - variables: { a: 1, b: 2 } + variables: { a: 1, b: 2 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual([1, 2, 3]); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual([1, 2, 3]); - }); - it('返回嵌套字典', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('返回嵌套字典', async () => { + const result = await getRunner().execute({ + code: `def main(variables): return { "user": {"name": variables["name"], "tags": ["admin"]}, "count": 42 }`, - variables: { name: 'test' } + variables: { name: 'test' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.user.name).toBe('test'); + expect(result.data?.codeReturn.count).toBe(42); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.user.name).toBe('test'); - expect(result.data?.codeReturn.count).toBe(42); - }); - it('返回布尔值和 None 转换', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('返回布尔值和 None 转换', async () => { + const result = await getRunner().execute({ + code: `def main(variables): return {"active": True, "deleted": False, "extra": None}`, - variables: {} + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.active).toBe(true); + expect(result.data?.codeReturn.deleted).toBe(false); + expect(result.data?.codeReturn.extra).toBeNull(); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.active).toBe(true); - expect(result.data?.codeReturn.deleted).toBe(false); - expect(result.data?.codeReturn.extra).toBeNull(); - }); - - // ===== print 行为 ===== - it('print 输出收集到 log（不影响返回值）', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('print 输出收集到 log（不影响返回值）', async () => { + const result = await getRunner().execute({ + code: `def main(variables): print("debug step 1") print("processing", variables["name"]) return {"done": True}`, - variables: { name: 'test' } + variables: { name: 'test' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.done).toBe(true); + expect(result.data?.log).toContain('debug step 1'); + expect(result.data?.log).toContain('processing'); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.done).toBe(true); - expect(result.data?.log).toContain('debug step 1'); - expect(result.data?.log).toContain('processing'); - }); - - // ===== 危险模块拦截 ===== - it('import os 被拦截', async () => { - const result = await pool.execute({ - code: `import os + it('import os 被拦截', async () => { + const result = await getRunner().execute({ + code: `import os def main(): return {"cwd": os.getcwd()}`, - variables: {} + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toContain('os'); }); - expect(result.success).toBe(false); - expect(result.message).toContain('os'); - }); - it('import subprocess 被拦截', async () => { - const result = await pool.execute({ - code: `import subprocess + it('import subprocess 被拦截', async () => { + const result = await getRunner().execute({ + code: `import subprocess def main(): return {"out": subprocess.check_output(["ls"])}`, - variables: {} + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toContain('subprocess'); }); - expect(result.success).toBe(false); - expect(result.message).toContain('subprocess'); - }); - it('from sys import path 被拦截', async () => { - const result = await pool.execute({ - code: `from sys import path + it('from sys import path 被拦截', async () => { + const result = await getRunner().execute({ + code: `from sys import path def main(): return {"path": path}`, - variables: {} + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toContain('sys'); }); - expect(result.success).toBe(false); - expect(result.message).toContain('sys'); - }); - - // ===== 安全模块允许 ===== - it('import json 允许', async () => { - const result = await pool.execute({ - code: `import json + it('import json 允许', async () => { + const result = await getRunner().execute({ + code: `import json def main(variables): data = json.dumps(variables) parsed = json.loads(data) return {"parsed": parsed}`, - variables: { key: 'value' } + variables: { key: 'value' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.parsed).toEqual({ key: 'value' }); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.parsed).toEqual({ key: 'value' }); - }); - it('import math 允许', async () => { - const result = await pool.execute({ - code: `import math + it('import math 允许', async () => { + const result = await getRunner().execute({ + code: `import math def main(variables): return {"sqrt": math.sqrt(variables["n"]), "pi": round(math.pi, 4)}`, - variables: { n: 16 } + variables: { n: 16 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.sqrt).toBe(4); + expect(result.data?.codeReturn.pi).toBe(3.1416); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.sqrt).toBe(4); - expect(result.data?.codeReturn.pi).toBe(3.1416); - }); - it('import re 允许', async () => { - const result = await pool.execute({ - code: `import re + it('import re 允许', async () => { + const result = await getRunner().execute({ + code: `import re def main(variables): matches = re.findall(r'\\d+', variables["text"]) return {"numbers": matches}`, - variables: { text: 'abc123def456' } + variables: { text: 'abc123def456' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.numbers).toEqual(['123', '456']); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.numbers).toEqual(['123', '456']); - }); - // ===== 典型旧版代码模式 ===== - - it('旧版典型写法：数据过滤', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('旧版典型写法：数据过滤', async () => { + const result = await getRunner().execute({ + code: `def main(variables): items = variables["items"] filtered = [x for x in items if x["score"] >= 60] return {"passed": len(filtered), "total": len(items)}`, - variables: { - items: [ - { name: 'A', score: 80 }, - { name: 'B', score: 45 }, - { name: 'C', score: 90 }, - { name: 'D', score: 55 } - ] - } + variables: { + items: [ + { name: 'A', score: 80 }, + { name: 'B', score: 45 }, + { name: 'C', score: 90 }, + { name: 'D', score: 55 } + ] + } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ passed: 2, total: 4 }); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual({ passed: 2, total: 4 }); - }); - it('旧版典型写法：字符串处理', async () => { - const result = await pool.execute({ - code: `def main(variables): + it('旧版典型写法：字符串处理', async () => { + const result = await getRunner().execute({ + code: `def main(variables): text = variables["text"] words = text.split() return { @@ -268,40 +246,58 @@ def main(variables): "upper": text.upper(), "reversed": text[::-1] }`, - variables: { text: 'hello world' } + variables: { text: 'hello world' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.word_count).toBe(2); + expect(result.data?.codeReturn.upper).toBe('HELLO WORLD'); + expect(result.data?.codeReturn.reversed).toBe('dlrow olleh'); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.word_count).toBe(2); - expect(result.data?.codeReturn.upper).toBe('HELLO WORLD'); - expect(result.data?.codeReturn.reversed).toBe('dlrow olleh'); - }); - it('旧版典型写法：日期处理', async () => { - const result = await pool.execute({ - code: `from datetime import datetime, timedelta + it('旧版典型写法：日期处理', async () => { + const result = await getRunner().execute({ + code: `from datetime import datetime, timedelta def main(variables): dt = datetime.strptime(variables["date"], "%Y-%m-%d") next_day = dt + timedelta(days=1) return {"next": next_day.strftime("%Y-%m-%d"), "weekday": dt.strftime("%A")}`, - variables: { date: '2024-01-01' } + variables: { date: '2024-01-01' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.next).toBe('2024-01-02'); + expect(result.data?.codeReturn.weekday).toBe('Monday'); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.next).toBe('2024-01-02'); - expect(result.data?.codeReturn.weekday).toBe('Monday'); - }); - it('旧版典型写法：辅助函数 + main', async () => { - const result = await pool.execute({ - code: `def calculate_tax(amount, rate): + it('旧版典型写法：辅助函数 + main', async () => { + const result = await getRunner().execute({ + code: `def calculate_tax(amount, rate): return round(amount * rate, 2) def main(variables): amount = variables["amount"] tax = calculate_tax(amount, 0.13) return {"amount": amount, "tax": tax, "total": amount + tax}`, - variables: { amount: 100 } + variables: { amount: 100 } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ amount: 100, tax: 13, total: 113 }); + }); + }); +} + +describe('旧版 Python 兼容性', () => { + describe('PythonIsolatedRunner', () => { + let runner: PythonIsolatedRunner; + + beforeAll(async () => { + runner = new PythonIsolatedRunner(1); + await runner.init(); }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn).toEqual({ amount: 100, tax: 13, total: 113 }); + + afterAll(async () => { + await runner.shutdown(); + }); + + legacyPythonCompatibilitySuite('legacy cases', () => runner); }); }); diff --git a/projects/code-sandbox/test/helpers/custom-process-pool.ts b/projects/code-sandbox/test/helpers/custom-process-pool.ts index 3810f65436f0..8afd393c4bbb 100644 --- a/projects/code-sandbox/test/helpers/custom-process-pool.ts +++ b/projects/code-sandbox/test/helpers/custom-process-pool.ts @@ -1,19 +1,16 @@ import { join } from 'path'; import { BaseProcessPool } from '../../src/pool/base-process-pool'; -export type SandboxLanguage = 'JS' | 'Python'; - export class CustomModuleProcessPool extends BaseProcessPool { constructor( - language: SandboxLanguage, + language: 'JS', allowedModules: readonly string[], options: { recycleAfterTask?: boolean } = {} ) { super(1, { name: `Custom${language}`, - workerScript: join(process.cwd(), 'src/pool', language === 'JS' ? 'worker.ts' : 'worker.py'), - spawnCommand: (script) => - language === 'JS' ? `exec tsx ${script}` : `exec python3 -u ${script}`, + workerScript: join(process.cwd(), 'src/pool', 'worker.ts'), + spawnCommand: (script) => `exec tsx ${script}`, allowedModules, recycleAfterTask: options.recycleAfterTask }); diff --git a/projects/code-sandbox/test/integration/allowed-modules-availability.test.ts b/projects/code-sandbox/test/integration/allowed-modules-availability.test.ts new file mode 100644 index 000000000000..728134e6919e --- /dev/null +++ b/projects/code-sandbox/test/integration/allowed-modules-availability.test.ts @@ -0,0 +1,246 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest'; +import { env } from '../../src/env'; +import { ProcessPool } from '../../src/pool/process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; + +type ModuleCase = { + code: string; + timeoutMs?: number; +}; + +const jsModuleCases: Record = { + lodash: { + code: `async function main() { + const _ = require('lodash'); + return { ok: _.sum([1, 2, 3]) === 6 }; + }` + }, + dayjs: { + code: `async function main() { + const dayjs = require('dayjs'); + return { ok: dayjs('2024-01-01').add(1, 'day').format('YYYY-MM-DD') === '2024-01-02' }; + }` + }, + moment: { + code: `async function main() { + const moment = require('moment'); + return { ok: moment.utc('2024-01-01').add(1, 'day').format('YYYY-MM-DD') === '2024-01-02' }; + }` + }, + uuid: { + code: `async function main() { + const { v5, validate } = require('uuid'); + return { ok: validate(v5('fastgpt', v5.URL)) }; + }` + }, + 'crypto-js': { + code: `async function main() { + const CryptoJS = require('crypto-js'); + return { ok: CryptoJS.SHA256('fastgpt').toString().length === 64 }; + }` + }, + qs: { + code: `async function main() { + const qs = require('qs'); + return { ok: qs.parse('a%5Bb%5D=1').a.b === '1' }; + }` + }, + url: { + code: `async function main() { + const url = require('url'); + const parsed = url.parse('https://example.com/a?b=1', true); + return { ok: parsed.hostname === 'example.com' && parsed.query.b === '1' }; + }` + }, + querystring: { + code: `async function main() { + const querystring = require('querystring'); + return { ok: querystring.parse('a=1').a === '1' }; + }` + } +}; + +const pythonModuleCases: Record = { + math: { code: `import math\ndef main():\n return {'ok': math.sqrt(9) == 3}` }, + cmath: { code: `import cmath\ndef main():\n return {'ok': cmath.sqrt(-1) == 1j}` }, + decimal: { + code: `import decimal\ndef main():\n return {'ok': str(decimal.Decimal('0.1') + decimal.Decimal('0.2')) == '0.3'}` + }, + fractions: { + code: `import fractions\ndef main():\n return {'ok': str(fractions.Fraction(1, 3) + fractions.Fraction(1, 6)) == '1/2'}` + }, + random: { + code: `import random\ndef main():\n random.seed(1)\n return {'ok': random.randint(1, 10) == 3}` + }, + statistics: { + code: `import statistics\ndef main():\n return {'ok': statistics.mean([1, 2, 3]) == 2}` + }, + collections: { + code: `import collections\ndef main():\n return {'ok': collections.Counter(['a', 'a'])['a'] == 2}` + }, + array: { + code: `import array\ndef main():\n return {'ok': array.array('i', [1, 2]).tolist() == [1, 2]}` + }, + heapq: { + code: `import heapq\ndef main():\n h = [3, 1, 2]\n heapq.heapify(h)\n return {'ok': heapq.heappop(h) == 1}` + }, + bisect: { + code: `import bisect\ndef main():\n return {'ok': bisect.bisect_left([1, 3, 5], 3) == 1}` + }, + queue: { + code: `import queue\ndef main():\n q = queue.Queue()\n q.put('ok')\n return {'ok': q.get() == 'ok'}` + }, + copy: { + code: `import copy\ndef main():\n a = {'x': [1]}\n b = copy.deepcopy(a)\n b['x'].append(2)\n return {'ok': a['x'] == [1]}` + }, + itertools: { + code: `import itertools\ndef main():\n return {'ok': list(itertools.combinations([1, 2, 3], 2))[0] == (1, 2)}` + }, + functools: { + code: `import functools\ndef main():\n return {'ok': functools.reduce(lambda a, b: a + b, [1, 2, 3], 0) == 6}` + }, + operator: { + code: `import operator\ndef main():\n return {'ok': operator.itemgetter('a')({'a': 1}) == 1}` + }, + string: { + code: `import string\ndef main():\n return {'ok': string.ascii_lowercase[:3] == 'abc'}` + }, + re: { + code: `import re\ndef main():\n return {'ok': re.search(r'\\d+', 'a123').group(0) == '123'}` + }, + difflib: { + code: `import difflib\ndef main():\n return {'ok': list(difflib.ndiff(['a'], ['b']))[0][0] == '-'}` + }, + textwrap: { + code: `import textwrap\ndef main():\n return {'ok': textwrap.shorten('hello world', width=8, placeholder='...') == 'hello...'}` + }, + unicodedata: { + code: `import unicodedata\ndef main():\n return {'ok': unicodedata.name('A') == 'LATIN CAPITAL LETTER A'}` + }, + codecs: { + code: `import codecs\ndef main():\n return {'ok': codecs.decode(b'Zm9v', 'base64').decode() == 'foo'}` + }, + datetime: { + code: `import datetime\ndef main():\n return {'ok': datetime.datetime(2024, 1, 1).year == 2024}` + }, + time: { code: `import time\ndef main():\n return {'ok': isinstance(time.time(), float)}` }, + calendar: { + code: `import calendar\ndef main():\n return {'ok': calendar.monthrange(2024, 2)[1] == 29}` + }, + _strptime: { + code: `import _strptime\ndef main():\n return {'ok': _strptime._strptime_time('2024-01-02', '%Y-%m-%d').tm_year == 2024}` + }, + json: { code: `import json\ndef main():\n return {'ok': json.loads('{"a":1}')['a'] == 1}` }, + csv: { + code: `import csv, io\ndef main():\n s = io.StringIO()\n csv.writer(s).writerow(['a', 'b'])\n return {'ok': s.getvalue().strip() == 'a,b'}` + }, + base64: { + code: `import base64\ndef main():\n return {'ok': base64.b64encode(b'ok').decode() == 'b2s='}` + }, + binascii: { + code: `import binascii\ndef main():\n return {'ok': binascii.hexlify(b'ok').decode() == '6f6b'}` + }, + struct: { + code: `import struct\ndef main():\n return {'ok': struct.unpack('>I', bytes([0, 0, 0, 42]))[0] == 42}` + }, + hashlib: { + code: `import hashlib\ndef main():\n return {'ok': hashlib.sha256(b'fastgpt').hexdigest().startswith('046ca27')}` + }, + hmac: { + code: `import hmac, hashlib\ndef main():\n return {'ok': len(hmac.new(b'k', b'v', hashlib.sha256).hexdigest()) == 64}` + }, + secrets: { + code: `import secrets\ndef main():\n return {'ok': len(secrets.token_hex(4)) == 8}` + }, + uuid: { + code: `import uuid\ndef main():\n return {'ok': str(uuid.uuid5(uuid.NAMESPACE_DNS, 'fastgpt')) == '8df8855a-2990-5a63-831d-be4be1d105bb'}` + }, + typing: { + code: `import typing\ndef main():\n return {'ok': str(typing.List[int]) == 'typing.List[int]'}` + }, + abc: { + code: `import abc\nclass Base(metaclass=abc.ABCMeta): pass\ndef main():\n return {'ok': isinstance(Base, abc.ABCMeta)}` + }, + enum: { + code: `import enum\nclass Color(enum.Enum):\n RED = 1\ndef main():\n return {'ok': Color.RED.name == 'RED'}` + }, + dataclasses: { + code: `import dataclasses\n@dataclasses.dataclass\nclass Item:\n name: str\ndef main():\n return {'ok': dataclasses.asdict(Item('ok')) == {'name': 'ok'}}` + }, + contextlib: { + code: `import contextlib\ndef main():\n with contextlib.suppress(ValueError):\n int('x')\n return {'ok': True}` + }, + pprint: { + code: `import pprint\ndef main():\n return {'ok': "'a': 1" in pprint.pformat({'b': 2, 'a': 1})}` + }, + weakref: { + code: `import weakref\nclass Box: pass\ndef main():\n b = Box()\n r = weakref.ref(b)\n return {'ok': r() is b}` + }, + numpy: { + code: `import numpy as np\ndef main():\n a = np.array([[1, 2, 3], [4, 5, 6]])\n return {'ok': list(a.shape) == [2, 3] and float(a.mean()) == 3.5 and int(np.dot(np.array([1, 2, 3]), np.array([4, 5, 6]))) == 32}` + }, + pandas: { + code: `import pandas as pd\ndef main():\n df = pd.DataFrame([{'team': 'a', 'score': 1}, {'team': 'a', 'score': 3}, {'team': 'b', 'score': 2}])\n grouped = df.groupby('team')['score'].sum().to_dict()\n return {'ok': int(grouped['a']) == 4 and int(grouped['b']) == 2}` + }, + matplotlib: { + code: `import matplotlib\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt\ndef main():\n fig, ax = plt.subplots(figsize=(2, 1))\n ax.plot([1, 2, 3], [1, 4, 9])\n axes_count = len(fig.axes)\n plt.close(fig)\n return {'ok': bool(matplotlib.__version__) and matplotlib.get_backend().lower() == 'agg' and axes_count == 1}`, + timeoutMs: 30000 + } +}; + +function expectSameMembers(actual: readonly string[], expected: Record) { + expect([...actual].sort()).toEqual(Object.keys(expected).sort()); +} + +describe('Sandbox allowed modules availability', () => { + let jsPool: ProcessPool; + let pyRunner: PythonIsolatedRunner; + + beforeAll(async () => { + jsPool = new ProcessPool(1); + await jsPool.init(); + pyRunner = new PythonIsolatedRunner(1); + await pyRunner.init(); + }); + + afterAll(async () => { + await jsPool?.shutdown(); + await pyRunner?.shutdown(); + }); + + it('JS 白名单中的每个模块都有本地可用性用例', () => { + expectSameMembers(env.SANDBOX_JS_ALLOWED_MODULES, jsModuleCases); + }); + + it('Python 白名单中的每个模块都有本地可用性用例', () => { + expectSameMembers(env.SANDBOX_PYTHON_ALLOWED_MODULES, pythonModuleCases); + }); + + it.each(env.SANDBOX_JS_ALLOWED_MODULES)( + 'JS 白名单模块 %s 可 require 并执行', + async (moduleName) => { + const result = await jsPool.execute({ + code: jsModuleCases[moduleName].code, + variables: {}, + timeoutMs: jsModuleCases[moduleName].timeoutMs + }); + + expect(result.success, result.message).toBe(true); + expect(result.data?.codeReturn?.ok).toBe(true); + } + ); + + it.each(env.SANDBOX_PYTHON_ALLOWED_MODULES)( + 'Python 白名单模块 %s 可 import 并执行', + async (moduleName) => { + const result = await pyRunner.execute({ + code: pythonModuleCases[moduleName].code, + variables: {}, + timeoutMs: pythonModuleCases[moduleName].timeoutMs + }); + + expect(result.success, result.message).toBe(true); + expect(result.data?.codeReturn?.ok).toBe(true); + } + ); +}); diff --git a/projects/code-sandbox/test/integration/api.test.ts b/projects/code-sandbox/test/integration/api.test.ts index 1ae114cb0db1..e2c57961dc0a 100644 --- a/projects/code-sandbox/test/integration/api.test.ts +++ b/projects/code-sandbox/test/integration/api.test.ts @@ -80,6 +80,9 @@ describe('API Routes', () => { expect(res.status).toBe(200); const data = await res.json(); expect(data.status).toBe('ok'); + expect(data.pools.js.total).toBeGreaterThan(0); + expect(data.pools.python.ready).toBe(true); + expect(data.pools.python.total).toBeGreaterThan(0); }); // ===== JS ===== diff --git a/projects/code-sandbox/test/integration/docker-js-packages.test.ts b/projects/code-sandbox/test/integration/docker-js-packages.test.ts new file mode 100644 index 000000000000..749fbcb37d7f --- /dev/null +++ b/projects/code-sandbox/test/integration/docker-js-packages.test.ts @@ -0,0 +1,199 @@ +import { describe, expect, it } from 'vitest'; + +const baseUrl = process.env.CODE_SANDBOX_URL?.replace(/\/$/, ''); +const token = process.env.SANDBOX_TOKEN || ''; +const shouldRun = Boolean(baseUrl); + +type SandboxResponse = { + success: boolean; + data?: { + codeReturn?: any; + log?: string; + }; + message?: string; +}; + +async function runJs(code: string, variables: Record = {}) { + if (!baseUrl) throw new Error('CODE_SANDBOX_URL is required'); + + const headers: Record = { + 'content-type': 'application/json' + }; + if (token) { + headers.Authorization = `Bearer ${token}`; + } + + const res = await fetch(`${baseUrl}/sandbox/js`, { + method: 'POST', + headers, + body: JSON.stringify({ + code, + variables, + timeoutMs: 30000 + }) + }); + + const json = (await res.json()) as SandboxResponse; + if (!res.ok) { + throw new Error(`HTTP ${res.status}: ${JSON.stringify(json)}`); + } + if (!json.success) { + throw new Error(json.message || JSON.stringify(json)); + } + return json; +} + +describe.skipIf(!shouldRun)('Docker JS 预装包集成测试', () => { + it('health ready', async () => { + const res = await fetch(`${baseUrl}/health`); + expect(res.status).toBe(200); + const json = await res.json(); + expect(json.status).toBe('ok'); + }); + + it('lodash 可 require 并执行常用集合函数', async () => { + const result = await runJs(` +async function main() { + const _ = require('lodash'); + const grouped = _.groupBy([ + { team: 'a', score: 1 }, + { team: 'a', score: 3 }, + { team: 'b', score: 2 } + ], 'team'); + return { + chunks: _.chunk([1, 2, 3, 4, 5], 2), + sumA: _.sumBy(grouped.a, 'score'), + sumB: _.sumBy(grouped.b, 'score') + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.chunks).toEqual([[1, 2], [3, 4], [5]]); + expect(result.data?.codeReturn.sumA).toBe(4); + expect(result.data?.codeReturn.sumB).toBe(2); + }); + + it('moment 可 require 并处理日期', async () => { + const result = await runJs(` +async function main() { + const moment = require('moment'); + return { + value: moment.utc('2024-01-15T12:00:00Z').add(3, 'hours').format('YYYY-MM-DD HH:mm') + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.value).toBe('2024-01-15 15:00'); + }); + + it('dayjs 可 require 并处理日期', async () => { + const result = await runJs(` +async function main() { + const dayjs = require('dayjs'); + return { + value: dayjs('2024-01-15T12:00:00Z').add(2, 'day').format('YYYY-MM-DD') + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.value).toBe('2024-01-17'); + }); + + it('crypto-js 可 require 并计算哈希', async () => { + const result = await runJs(` +async function main() { + const CryptoJS = require('crypto-js'); + return { + sha256: CryptoJS.SHA256('fastgpt').toString() + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.sha256).toBe( + '046ca27ed8a95d7aaec3fd577ba8d9eabd7f7915de4e7c0e120d06d758bff75a' + ); + }); + + it('uuid 可 require 并生成 UUID', async () => { + const result = await runJs(` +async function main() { + const { v5, validate } = require('uuid'); + const id = v5('fastgpt', v5.URL); + return { + id, + valid: validate(id) + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.valid).toBe(true); + expect(result.data?.codeReturn.id).toBe('8d10dc2b-3d86-5ae4-b41b-a42c6000e3eb'); + }); + + it('qs 可 require 并处理查询串', async () => { + const result = await runJs(` +async function main() { + const qs = require('qs'); + return { + parsed: qs.parse('a%5Bb%5D=1&list%5B0%5D=2&list%5B1%5D=3'), + text: qs.stringify({ a: { b: 1 }, list: [2, 3] }) + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.parsed).toEqual({ a: { b: '1' }, list: ['2', '3'] }); + expect(result.data?.codeReturn.text).toContain('a%5Bb%5D=1'); + }); + + it('url 可 require 并解析 URL', async () => { + const result = await runJs(` +async function main() { + const url = require('url'); + const parsed = url.parse('https://example.com/a?b=1', true); + return { + hostname: parsed.hostname, + query: parsed.query.b + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + hostname: 'example.com', + query: '1' + }); + }); + + it('querystring 可 require 并处理查询串', async () => { + const result = await runJs(` +async function main() { + const querystring = require('querystring'); + return { + parsed: querystring.parse('a=1&b=2'), + text: querystring.stringify({ a: 1, b: 2 }) + }; +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.parsed).toEqual({ a: '1', b: '2' }); + expect(result.data?.codeReturn.text).toContain('a=1'); + }); + + it.each(['child_process', 'fs', 'net', 'http', 'https'])( + '危险模块 %s 不能 require', + async (moduleName) => { + const result = await runJs(` +async function main() { + try { + require(${JSON.stringify(moduleName)}); + return { blocked: false }; + } catch (e) { + return { blocked: true, message: String(e.message || e) }; + } +}`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toBe(true); + expect(result.data?.codeReturn.message).toMatch(/not allowed|denied|forbidden/i); + } + ); +}); diff --git a/projects/code-sandbox/test/integration/docker-python-packages.test.ts b/projects/code-sandbox/test/integration/docker-python-packages.test.ts new file mode 100644 index 000000000000..07e4ad4c293f --- /dev/null +++ b/projects/code-sandbox/test/integration/docker-python-packages.test.ts @@ -0,0 +1,414 @@ +import { describe, expect, it } from 'vitest'; + +const baseUrl = process.env.CODE_SANDBOX_URL?.replace(/\/$/, ''); +const token = process.env.SANDBOX_TOKEN || ''; +const shouldRun = Boolean(baseUrl); + +type SandboxResponse = { + success: boolean; + data?: { + codeReturn?: any; + log?: string; + }; + message?: string; +}; + +async function runPython(code: string, variables: Record = {}) { + if (!baseUrl) throw new Error('CODE_SANDBOX_URL is required'); + + const headers: Record = { + 'content-type': 'application/json' + }; + if (token) { + headers.Authorization = `Bearer ${token}`; + } + + const res = await fetch(`${baseUrl}/sandbox/python`, { + method: 'POST', + headers, + body: JSON.stringify({ + code, + variables, + timeoutMs: 30000 + }) + }); + + const json = (await res.json()) as SandboxResponse; + if (!res.ok) { + throw new Error(`HTTP ${res.status}: ${JSON.stringify(json)}`); + } + if (!json.success) { + throw new Error(json.message || JSON.stringify(json)); + } + return json; +} + +describe.skipIf(!shouldRun)('Docker Python 预装包集成测试', () => { + it('health ready', async () => { + const res = await fetch(`${baseUrl}/health`); + expect(res.status).toBe(200); + const json = await res.json(); + expect(json.status).toBe('ok'); + }); + + it('数学和数值计算标准库均可 import 并运行', async () => { + const result = await runPython(` +import math +import cmath +import decimal +import fractions +import random +import statistics + +def main(): + random.seed(1) + return { + "math": math.isclose(math.sqrt(9), 3), + "cmath": cmath.sqrt(-1) == 1j, + "decimal": str(decimal.Decimal("0.1") + decimal.Decimal("0.2")), + "fractions": str(fractions.Fraction(1, 3) + fractions.Fraction(1, 6)), + "random": random.randint(1, 10), + "statistics": statistics.mean([1, 2, 3, 4]) + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toMatchObject({ + math: true, + cmath: true, + decimal: '0.3', + fractions: '1/2', + random: 3, + statistics: 2.5 + }); + }); + + it('数据结构和算法标准库均可 import 并运行', async () => { + const result = await runPython(` +import collections +import array +import heapq +import bisect +import queue +import copy + +def main(): + counter = collections.Counter(["a", "b", "a"]) + arr = array.array("i", [1, 2, 3]) + heap = [3, 1, 2] + heapq.heapify(heap) + q = queue.Queue() + q.put("ok") + original = {"a": [1]} + cloned = copy.deepcopy(original) + cloned["a"].append(2) + return { + "collections": counter["a"], + "array": arr.tolist(), + "heapq": heapq.heappop(heap), + "bisect": bisect.bisect_left([1, 3, 5], 3), + "queue": q.get(), + "copy": original["a"] + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + collections: 2, + array: [1, 2, 3], + heapq: 1, + bisect: 1, + queue: 'ok', + copy: [1] + }); + }); + + it('函数式编程标准库均可 import 并运行', async () => { + const result = await runPython(` +import itertools +import functools +import operator + +def main(): + pairs = list(itertools.combinations([1, 2, 3], 2)) + total = functools.reduce(operator.add, [1, 2, 3], 0) + picked = operator.itemgetter("name")({"name": "FastGPT"}) + return {"itertools": pairs, "functools": total, "operator": picked} +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + itertools: [ + [1, 2], + [1, 3], + [2, 3] + ], + functools: 6, + operator: 'FastGPT' + }); + }); + + it('字符串和文本处理标准库均可 import 并运行', async () => { + const result = await runPython(` +import string +import re +import difflib +import textwrap +import unicodedata +import codecs + +def main(): + return { + "string": string.ascii_lowercase[:3], + "re": re.search(r"\\d+", "a123").group(0), + "difflib": list(difflib.ndiff(["a"], ["b"]))[0][0], + "textwrap": textwrap.shorten("hello world", width=8, placeholder="..."), + "unicodedata": unicodedata.name("A"), + "codecs": codecs.decode(b"Zm9v", "base64").decode() + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + string: 'abc', + re: '123', + difflib: '-', + textwrap: 'hello...', + unicodedata: 'LATIN CAPITAL LETTER A', + codecs: 'foo' + }); + }); + + it('日期和时间标准库均可 import 并运行', async () => { + const result = await runPython(` +import datetime +import time +import calendar + +def main(): + dt = datetime.datetime(2024, 1, 15, 12, 0, 0) + return { + "datetime": dt.isoformat(), + "time": isinstance(time.time(), float), + "calendar": calendar.monthrange(2024, 2)[1] + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + datetime: '2024-01-15T12:00:00', + time: true, + calendar: 29 + }); + }); + + it('数据序列化标准库均可 import 并运行', async () => { + const result = await runPython(` +import json +import csv +import base64 +import binascii +import struct +import io + +def main(): + out = io.StringIO() + writer = csv.writer(out) + writer.writerow(["a", "b"]) + return { + "json": json.loads('{"a": 1}')["a"], + "csv": out.getvalue().strip(), + "base64": base64.b64encode(b"ok").decode(), + "binascii": binascii.hexlify(b"ok").decode(), + "struct": struct.unpack(">I", bytes([0, 0, 0, 42]))[0] + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + json: 1, + csv: 'a,b', + base64: 'b2s=', + binascii: '6f6b', + struct: 42 + }); + }); + + it('加密和哈希标准库均可 import 并运行', async () => { + const result = await runPython(` +import hashlib +import hmac +import secrets +import uuid + +def main(): + token = secrets.token_hex(4) + return { + "hashlib": hashlib.sha256(b"fastgpt").hexdigest(), + "hmac": hmac.new(b"k", b"v", hashlib.sha256).hexdigest(), + "secrets_len": len(token), + "uuid": str(uuid.uuid5(uuid.NAMESPACE_DNS, "fastgpt")) + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.hashlib).toBe( + '046ca27ed8a95d7aaec3fd577ba8d9eabd7f7915de4e7c0e120d06d758bff75a' + ); + expect(result.data?.codeReturn.hmac).toBeTruthy(); + expect(result.data?.codeReturn.secrets_len).toBe(8); + expect(result.data?.codeReturn.uuid).toBe('8df8855a-2990-5a63-831d-be4be1d105bb'); + }); + + it('类型和抽象标准库均可 import 并运行', async () => { + const result = await runPython(` +import typing +import abc +import enum +import dataclasses +import contextlib + +class Color(enum.Enum): + RED = 1 + +@dataclasses.dataclass +class Item: + name: str + +class Base(metaclass=abc.ABCMeta): + pass + +def main(): + with contextlib.suppress(ValueError): + int("x") + hint = typing.List[int] + return { + "typing": str(hint), + "abc": isinstance(Base, abc.ABCMeta), + "enum": Color.RED.name, + "dataclasses": dataclasses.asdict(Item("ok")), + "contextlib": True + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.abc).toBe(true); + expect(result.data?.codeReturn.enum).toBe('RED'); + expect(result.data?.codeReturn.dataclasses).toEqual({ name: 'ok' }); + expect(result.data?.codeReturn.contextlib).toBe(true); + }); + + it('其他实用工具标准库均可 import 并运行', async () => { + const result = await runPython(` +import pprint +import weakref + +class Box: + pass + +def main(): + box = Box() + ref = weakref.ref(box) + return { + "pprint": pprint.pformat({"b": 2, "a": 1}), + "weakref": ref() is box + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.pprint).toContain("'a': 1"); + expect(result.data?.codeReturn.weakref).toBe(true); + }); + + it('numpy 可 import 并执行基础矩阵运算', async () => { + const result = await runPython(` +import numpy as np + +def main(): + arr = np.array([[1, 2, 3], [4, 5, 6]]) + return { + "version": np.__version__, + "shape": list(arr.shape), + "mean": float(arr.mean()), + "dot": int(np.dot(np.array([1, 2, 3]), np.array([4, 5, 6]))) + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.shape).toEqual([2, 3]); + expect(result.data?.codeReturn.mean).toBe(3.5); + expect(result.data?.codeReturn.dot).toBe(32); + expect(result.data?.codeReturn.version).toBeTruthy(); + }); + + it('pandas 可 import 并执行 DataFrame 基础操作', async () => { + const result = await runPython(` +import pandas as pd + +def main(): + df = pd.DataFrame([ + {"team": "a", "score": 1}, + {"team": "a", "score": 3}, + {"team": "b", "score": 2} + ]) + grouped = df.groupby("team")["score"].sum().to_dict() + return { + "version": pd.__version__, + "rows": int(len(df)), + "a": int(grouped["a"]), + "b": int(grouped["b"]) + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.rows).toBe(3); + expect(result.data?.codeReturn.a).toBe(4); + expect(result.data?.codeReturn.b).toBe(2); + expect(result.data?.codeReturn.version).toBeTruthy(); + }); + + it('matplotlib 可使用 Agg 后端生成 PNG', async () => { + const result = await runPython(` +import io +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +def main(): + fig, ax = plt.subplots(figsize=(2, 1)) + ax.plot([1, 2, 3], [2, 4, 6]) + ax.set_title("ok") + buf = io.BytesIO() + fig.savefig(buf, format="png") + plt.close(fig) + data = buf.getvalue() + return { + "backend": matplotlib.get_backend(), + "size": len(data), + "png": data[:8].hex() + } +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.backend.toLowerCase()).toContain('agg'); + expect(result.data?.codeReturn.size).toBeGreaterThan(1000); + expect(result.data?.codeReturn.png).toBe('89504e470d0a1a0a'); + }); + + it('预装包间接暴露 os 时仍不能执行系统命令', async () => { + const result = await runPython(` +import platform + +def main(): + os_ref = getattr(platform, "os") + try: + rc = os_ref.system("id") + return {"blocked": rc != 0, "rc": rc} + except Exception as e: + return {"blocked": True, "error": str(e)} +`); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toBe(true); + }); +}); diff --git a/projects/code-sandbox/test/integration/functional.test.ts b/projects/code-sandbox/test/integration/functional.test.ts index 4684fd5b9584..bbf2248df3d9 100644 --- a/projects/code-sandbox/test/integration/functional.test.ts +++ b/projects/code-sandbox/test/integration/functional.test.ts @@ -6,7 +6,7 @@ */ import { describe, it, expect, beforeAll, afterAll } from 'vitest'; import { ProcessPool } from '../../src/pool/process-pool'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; // ============================================================ // 测试用例矩阵类型 @@ -25,7 +25,7 @@ interface TestCase { }; } -function runMatrix(getPool: () => ProcessPool | PythonProcessPool, cases: TestCase[]) { +function runMatrix(getPool: () => ProcessPool | PythonIsolatedRunner, cases: TestCase[]) { for (const tc of cases) { it(tc.name, async () => { const result = await getPool().execute({ @@ -608,15 +608,15 @@ async function main() { return recurse(); }`, { name: 'httpRequest GET', code: `async function main() { - const res = await httpRequest('https://1.1.1.1/cdn-cgi/trace'); - return { status: res.status, hasData: res.data.length > 0 }; + const res = await httpRequest('http://1.1.1.1/'); + return { statusOk: res.status >= 200 && res.status < 400, hasData: res.data.length > 0 }; }`, - expect: { success: true, codeReturnMatch: { status: 200, hasData: true } } + expect: { success: true, codeReturnMatch: { statusOk: true, hasData: true } } }, { name: 'httpRequest POST JSON', code: `async function main() { - const res = await httpRequest('https://1.1.1.1/cdn-cgi/trace', { + const res = await httpRequest('http://1.1.1.1/', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: { message: 'hello' } @@ -634,9 +634,9 @@ async function main() { return recurse(); }`, // Python 功能测试矩阵 // ============================================================ describe('Python 功能测试', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; beforeAll(async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); }); afterAll(async () => { @@ -903,12 +903,12 @@ describe('Python 功能测试', () => { [ { name: 'http_request GET', - code: `import json\ndef main():\n res = http_request('https://1.1.1.1/cdn-cgi/trace')\n return {'status': res['status'], 'hasData': len(res['data']) > 0}`, - expect: { success: true, codeReturnMatch: { status: 200, hasData: true } } + code: `import json\ndef main():\n res = http_request('http://1.1.1.1/')\n return {'statusOk': res['status'] >= 200 and res['status'] < 400, 'hasData': len(res['data']) > 0}`, + expect: { success: true, codeReturnMatch: { statusOk: true, hasData: true } } }, { name: 'http_request POST JSON', - code: `import json\ndef main():\n res = http_request('https://1.1.1.1/cdn-cgi/trace', method='POST', body={'message': 'hello'})\n return {'hasStatus': type(res['status']) == int}`, + code: `import json\ndef main():\n res = http_request('http://1.1.1.1/', method='POST', body={'message': 'hello'})\n return {'hasStatus': type(res['status']) == int}`, expect: { success: true, codeReturnMatch: { hasStatus: true } } } ] diff --git a/projects/code-sandbox/test/unit/boundary.test.ts b/projects/code-sandbox/test/unit/boundary.test.ts index 20678f77d25b..39fcb2f5798f 100644 --- a/projects/code-sandbox/test/unit/boundary.test.ts +++ b/projects/code-sandbox/test/unit/boundary.test.ts @@ -1,14 +1,14 @@ import { describe, it, expect, beforeAll, afterAll } from 'vitest'; import { ProcessPool } from '../../src/pool/process-pool'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; let jsPool: ProcessPool; -let pyPool: PythonProcessPool; +let pyPool: PythonIsolatedRunner; beforeAll(async () => { jsPool = new ProcessPool(1); await jsPool.init(); - pyPool = new PythonProcessPool(1); + pyPool = new PythonIsolatedRunner(1); await pyPool.init(); }); diff --git a/projects/code-sandbox/test/unit/ipCheck.test.ts b/projects/code-sandbox/test/unit/ipCheck.test.ts index 614992648949..70ea2217ee61 100644 --- a/projects/code-sandbox/test/unit/ipCheck.test.ts +++ b/projects/code-sandbox/test/unit/ipCheck.test.ts @@ -98,6 +98,15 @@ describe('isInternalResolvedIP', () => { }); describe('私网段：CHECK_INTERNAL_IP 控制', () => { + it('未显式设置 CHECK_INTERNAL_IP 时使用 env 默认值阻止私网', async () => { + vi.unstubAllEnvs(); + delete process.env.CHECK_INTERNAL_IP; + vi.stubEnv('NODE_ENV', 'production'); + vi.resetModules(); + const mod = await import('../../src/utils/ipCheck.util'); + expect(mod.isInternalResolvedIP('10.0.0.1')).toBe(true); + }); + it('CHECK_INTERNAL_IP=false 时放行 10.0.0.1', async () => { vi.stubEnv('CHECK_INTERNAL_IP', 'false'); vi.resetModules(); diff --git a/projects/code-sandbox/test/unit/process-pool.test.ts b/projects/code-sandbox/test/unit/process-pool.test.ts index b5b392fdca6e..87575bfc8943 100644 --- a/projects/code-sandbox/test/unit/process-pool.test.ts +++ b/projects/code-sandbox/test/unit/process-pool.test.ts @@ -1,5 +1,5 @@ /** - * ProcessPool / PythonProcessPool 单元测试 + * ProcessPool / PythonIsolatedRunner 单元测试 * * 覆盖进程池核心逻辑： * - 生命周期（init / shutdown / stats） @@ -10,7 +10,7 @@ */ import { describe, it, expect, beforeEach, afterEach } from 'vitest'; import { ProcessPool } from '../../src/pool/process-pool'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; // ============================================================ // JS ProcessPool @@ -343,7 +343,7 @@ describe('ProcessPool 返回值序列化与参数校验', () => { // ============================================================ describe('JS + Python 混合并发', () => { let jsPool: ProcessPool; - let pyPool: PythonProcessPool; + let pyPool: PythonIsolatedRunner; afterEach(async () => { try { @@ -355,7 +355,7 @@ describe('JS + Python 混合并发', () => { it('JS 和 Python 混合并发执行', async () => { jsPool = new ProcessPool(1); await jsPool.init(); - pyPool = new PythonProcessPool(1); + pyPool = new PythonIsolatedRunner(1); await pyPool.init(); const jsPromise = jsPool.execute({ @@ -579,10 +579,10 @@ describe('ProcessPool 健康检查失败路径', () => { }); // ============================================================ -// Python PythonProcessPool - Worker Ping/Pong 健康检查 +// PythonIsolatedRunner - 生命周期、恢复、排队 // ============================================================ -describe('PythonProcessPool Worker 健康检查 (ping/pong)', () => { - let pool: PythonProcessPool; +describe('PythonIsolatedRunner 生命周期', () => { + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -590,57 +590,48 @@ describe('PythonProcessPool Worker 健康检查 (ping/pong)', () => { } catch {} }); - it('worker 正常响应 ping 后仍可执行任务', async () => { - pool = new PythonProcessPool(1); + it('init 后 stats 正确', async () => { + pool = new PythonIsolatedRunner(2); await pool.init(); + const s = pool.stats; + expect(s.total).toBe(2); + expect(s.idle).toBe(2); + expect(s.busy).toBe(0); + expect(s.queued).toBe(0); + expect(s.poolSize).toBe(2); + }); - const r1 = await pool.execute({ - code: `def main():\n return {'step': 1}`, - variables: {} - }); - expect(r1.success).toBe(true); - expect(r1.data?.codeReturn.step).toBe(1); - - // 触发 ping - (pool as any).pingWorker((pool as any).idleWorkers[0]); - await new Promise((r) => setTimeout(r, 500)); + it('shutdown 后不再接受新任务', async () => { + pool = new PythonIsolatedRunner(1); + await pool.init(); + await pool.shutdown(); - const r2 = await pool.execute({ - code: `def main():\n return {'step': 2}`, + const result = await pool.execute({ + code: `def main():\n return {'ok': True}`, variables: {} }); - expect(r2.success).toBe(true); - expect(r2.data?.codeReturn.step).toBe(2); - expect(pool.stats.total).toBe(1); + expect(result.success).toBe(false); + expect(result.message).toMatch(/not ready/i); }); - it('连续多次 ping 不影响 worker 状态', async () => { - pool = new PythonProcessPool(2); + it('execute 后销毁已执行进程并补充干净空闲进程', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); - - for (let i = 0; i < 3; i++) { - for (const w of [...(pool as any).idleWorkers]) { - (pool as any).pingWorker(w); - } - await new Promise((r) => setTimeout(r, 300)); - } - - expect(pool.stats.total).toBe(2); - expect(pool.stats.idle).toBe(2); - const result = await pool.execute({ - code: `def main():\n return {'alive': True}`, + code: `def main():\n return {'ok': True}`, variables: {} }); expect(result.success).toBe(true); + await new Promise((resolve) => setTimeout(resolve, 200)); + const s = pool.stats; + expect(s.total).toBe(1); + expect(s.idle).toBe(1); + expect(s.busy).toBe(0); }); }); -// ============================================================ -// Python PythonProcessPool - 健康检查失败路径 -// ============================================================ -describe('PythonProcessPool 健康检查失败路径', () => { - let pool: PythonProcessPool; +describe('PythonIsolatedRunner 恢复能力', () => { + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -648,274 +639,75 @@ describe('PythonProcessPool 健康检查失败路径', () => { } catch {} }); - it('ping timeout: worker 不响应 pong 时被替换', async () => { - pool = new PythonProcessPool(1); + it('超时子进程被清理后，后续请求正常', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); - expect(pool.stats.total).toBe(1); - - const worker = (pool as any).idleWorkers[0]; - // 拦截 stdin.write 使 ping 不到达 worker，触发真正的 timeout - const origWrite = worker.proc.stdin!.write.bind(worker.proc.stdin!); - let interceptPing = true; - worker.proc.stdin!.write = (...args: any[]) => { - if (interceptPing) { - interceptPing = false; - return true; - } - return origWrite(...args); - }; - - (pool as any).pingWorker(worker); - - await new Promise((r) => setTimeout(r, 8000)); - - expect(pool.stats.total).toBe(1); const result = await pool.execute({ - code: `def main():\n return {'ok': True}`, - variables: {} - }); - expect(result.success).toBe(true); - }, 15000); - - it('stdin not writable: worker stdin 关闭时被替换', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - expect(pool.stats.total).toBe(1); - - const worker = (pool as any).idleWorkers[0]; - worker.proc.stdin!.destroy(); - - (pool as any).pingWorker(worker); - - await new Promise((r) => setTimeout(r, 3000)); - - expect(pool.stats.total).toBe(1); - - const result = await pool.execute({ - code: `def main():\n return {'replaced': True}`, - variables: {} - }); - expect(result.success).toBe(true); - }, 10000); - - it('health check invalid response: worker 返回错误类型时被替换', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - expect(pool.stats.total).toBe(1); - - const worker = (pool as any).idleWorkers[0]; - const origWrite = worker.proc.stdin!.write.bind(worker.proc.stdin!); - let intercepted = false; - worker.proc.stdin!.write = (...args: any[]) => { - if (!intercepted) { - intercepted = true; - setTimeout(() => worker.rl.emit('line', JSON.stringify({ type: 'wrong' })), 50); - return true; - } - return origWrite(...args); - }; - - (pool as any).pingWorker(worker); - - await new Promise((r) => setTimeout(r, 3000)); - - expect(pool.stats.total).toBe(1); - - const result = await pool.execute({ - code: `def main():\n return {'invalidResp': True}`, + code: `def main():\n while True:\n pass`, variables: {} }); - expect(result.success).toBe(true); - }, 10000); - - it('returnToIdle with waiter: ping 期间有等待请求时直接分配', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - - const worker = (pool as any).idleWorkers[0]; - (pool as any).pingWorker(worker); + expect(result.success).toBe(false); + expect(result.message).toContain('timed out'); - const p1 = pool.execute({ - code: `def main():\n return {'fromWaiter': True}`, + const result2 = await pool.execute({ + code: `def main():\n return {'ok': True}`, variables: {} }); - - const result = await p1; - expect(result.success).toBe(true); - expect(result.data?.codeReturn.fromWaiter).toBe(true); + expect(result2.success).toBe(true); + expect(result2.data?.codeReturn.ok).toBe(true); }); - it('health check parse error: worker 返回非 JSON 时被替换', async () => { - pool = new PythonProcessPool(1); + it('运行时异常不会污染下一次执行', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); - expect(pool.stats.total).toBe(1); - - const worker = (pool as any).idleWorkers[0]; - const origWrite = worker.proc.stdin!.write.bind(worker.proc.stdin!); - let intercepted = false; - worker.proc.stdin!.write = (...args: any[]) => { - if (!intercepted) { - intercepted = true; - setTimeout(() => worker.rl.emit('line', 'not-json'), 50); - return true; - } - return origWrite(...args); - }; - - (pool as any).pingWorker(worker); - - await new Promise((r) => setTimeout(r, 3000)); - - expect(pool.stats.total).toBe(1); const result = await pool.execute({ - code: `def main():\n return {'parseError': True}`, + code: `def main():\n raise ValueError('boom')`, variables: {} }); - expect(result.success).toBe(true); - }, 10000); - - it('health check write error: stdin.write 抛异常时被替换', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - expect(pool.stats.total).toBe(1); - - const worker = (pool as any).idleWorkers[0]; - worker.proc.stdin!.write = () => { - throw new Error('mock write error'); - }; - - (pool as any).pingWorker(worker); - - await new Promise((r) => setTimeout(r, 3000)); - - expect(pool.stats.total).toBe(1); + expect(result.success).toBe(false); + expect(result.message).toContain('boom'); - const result = await pool.execute({ - code: `def main():\n return {'writeError': True}`, + const result2 = await pool.execute({ + code: `def main():\n return {'recovered': True}`, variables: {} }); - expect(result.success).toBe(true); - }, 10000); + expect(result2.success).toBe(true); + expect(result2.data?.codeReturn.recovered).toBe(true); + }); }); -// ============================================================ -// Python PythonProcessPool - shutdown reject waiters -// ============================================================ -describe('PythonProcessPool shutdown reject waiters', () => { - it('shutdown 后 waitQueue 中的请求被 reject', async () => { - const pool = new PythonProcessPool(1); +describe('PythonIsolatedRunner shutdown reject waiters', () => { + it('shutdown 后排队中的请求返回 not ready 或被 reject', async () => { + const pool = new PythonIsolatedRunner(1); await pool.init(); - // 发起一个长时间运行的任务占住唯一 worker const p1 = pool.execute({ code: `import time\ndef main():\n time.sleep(3)\n return {'done': True}`, variables: {} }); - // 等一下确保 p1 已经拿到 worker await new Promise((r) => setTimeout(r, 200)); - // 发起第二个请求，它会进入 waitQueue const p2 = pool.execute({ code: `def main():\n return {'queued': True}`, variables: {} }); - // 确认有排队请求 expect(pool.stats.queued).toBe(1); - - // shutdown 应该 reject waitQueue 中的请求 await pool.shutdown(); - // p2 应该被 reject - await expect(p2).rejects.toThrow('shutting down'); + const r2 = await p2; + expect(r2.success).toBe(false); + expect(r2.message).toMatch(/not ready/i); - // p1 可能成功也可能因 worker 被 kill 而失败，不关心 await p1.catch(() => {}); }); }); -// ============================================================ -// Python PythonProcessPool -// ============================================================ -describe('PythonProcessPool 生命周期', () => { - let pool: PythonProcessPool; - - afterEach(async () => { - try { - await pool?.shutdown(); - } catch {} - }); - - it('init 后 stats 正确', async () => { - pool = new PythonProcessPool(2); - await pool.init(); - const s = pool.stats; - expect(s.total).toBe(2); - expect(s.idle).toBe(2); - expect(s.busy).toBe(0); - expect(s.queued).toBe(0); - expect(s.poolSize).toBe(2); - }); - - it('shutdown 后 stats 归零', async () => { - pool = new PythonProcessPool(2); - await pool.init(); - await pool.shutdown(); - const s = pool.stats; - expect(s.total).toBe(0); - expect(s.idle).toBe(0); - expect(s.busy).toBe(0); - }); - - it('execute 后 worker 归还到 idle', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - await pool.execute({ - code: `def main():\n return {'ok': True}`, - variables: {} - }); - const s = pool.stats; - expect(s.idle).toBe(1); - expect(s.busy).toBe(0); - }); -}); - -describe('PythonProcessPool Worker 恢复', () => { - let pool: PythonProcessPool; - - afterEach(async () => { - try { - await pool?.shutdown(); - } catch {} - }); - - it('超时后 worker 被 kill 并 respawn', async () => { - pool = new PythonProcessPool(1); - await pool.init(); - - const result = await pool.execute({ - code: `def main():\n while True:\n pass`, - variables: {} - }); - expect(result.success).toBe(false); - expect(result.message).toContain('timed out'); - - // 等 respawn - await new Promise((r) => setTimeout(r, 2000)); - - const result2 = await pool.execute({ - code: `def main():\n return {'ok': True}`, - variables: {} - }); - expect(result2.success).toBe(true); - }); -}); - -describe('PythonProcessPool 并发与排队', () => { - let pool: PythonProcessPool; +describe('PythonIsolatedRunner 并发与排队', () => { + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -923,17 +715,18 @@ describe('PythonProcessPool 并发与排队', () => { } catch {} }); - it('pool size=2，3 个并发请求，1 个排队', async () => { - pool = new PythonProcessPool(2); + it('maxConcurrency=2，3 个并发请求，1 个排队', async () => { + pool = new PythonIsolatedRunner(2); await pool.init(); const promises = Array.from({ length: 3 }, (_, i) => pool.execute({ - code: `import time\ndef main(variables):\n time.sleep(0.2)\n return {'idx': variables['idx']}`, + code: `import time\ndef main(idx):\n time.sleep(0.2)\n return {'idx': idx}`, variables: { idx: i } }) ); + expect(pool.stats.queued).toBe(1); const results = await Promise.all(promises); for (let i = 0; i < 3; i++) { expect(results[i].success).toBe(true); @@ -941,13 +734,13 @@ describe('PythonProcessPool 并发与排队', () => { } }); - it('pool size=1，10 个并发请求全部正确完成（串行排队）', async () => { - pool = new PythonProcessPool(1); + it('maxConcurrency=1，10 个并发请求全部正确完成（串行排队）', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); const promises = Array.from({ length: 10 }, (_, i) => pool.execute({ - code: `def main(variables):\n return {'n': variables['n'] * 2}`, + code: `def main(n):\n return {'n': n * 2}`, variables: { n: i } }) ); diff --git a/projects/code-sandbox/test/unit/python-isolated-runner.test.ts b/projects/code-sandbox/test/unit/python-isolated-runner.test.ts new file mode 100644 index 000000000000..f8ac34a7e8cb --- /dev/null +++ b/projects/code-sandbox/test/unit/python-isolated-runner.test.ts @@ -0,0 +1,361 @@ +import { afterEach, describe, expect, it } from 'vitest'; +import http from 'http'; +import { existsSync } from 'fs'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; +import { + PYTHON_SANDBOX_ROOT, + shouldEnablePythonNativeIsolation +} from '../../src/isolated/python-isolation-config'; + +describe('PythonIsolatedRunner 兼容性', () => { + let runner: PythonIsolatedRunner | undefined; + + afterEach(async () => { + await runner?.shutdown(); + runner = undefined; + }); + + async function createRunner(maxConcurrency = 2) { + runner = new PythonIsolatedRunner(maxConcurrency); + await runner.init(); + return runner; + } + + async function waitForIdlePid(r: PythonIsolatedRunner, previousPid?: number) { + const deadline = Date.now() + 3000; + while (Date.now() < deadline) { + const pid = (r as any).idleChildren.values().next().value?.proc?.pid; + if (pid && pid !== previousPid) return pid; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + return undefined; + } + + it('支持 main() 无参数、print log 和 JSON 返回', async () => { + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + print("debug") + return {"ok": True, "none": None}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ ok: true, none: null }); + expect(result.data?.log).toContain('debug'); + }); + + it('预热阶段没有 ready 子进程时 init fail closed', async () => { + const r = new PythonIsolatedRunner(1); + (r as any).replenishWarmChildren = async () => undefined; + runner = r; + + await expect(r.init()).rejects.toThrow(/warmup failed/); + expect(r.stats.ready).toBe(false); + expect(r.stats.total).toBe(0); + }); + + it('支持 main(variables) 和 main(a, b) 旧写法', async () => { + const r = await createRunner(); + + const byVariables = await r.execute({ + code: `def main(variables): + return {"name": variables["name"]}`, + variables: { name: 'FastGPT' } + }); + expect(byVariables.success).toBe(true); + expect(byVariables.data?.codeReturn.name).toBe('FastGPT'); + + const byArgs = await r.execute({ + code: `def main(a, b=1): + return {"sum": a + b}`, + variables: { a: 2 } + }); + expect(byArgs.success).toBe(true); + expect(byArgs.data?.codeReturn.sum).toBe(3); + }); + + it('保留 type() 正常判断能力', async () => { + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + value = 3 + return {"is_int": type(value) == int}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.is_int).toBe(true); + }); + + it('每次执行独立进程，不复用全局状态和模块污染', async () => { + const r = await createRunner(1); + + const first = await r.execute({ + code: `import json +json.dumps = lambda value: "polluted" +leaked = "yes" +def main(): + return {"polluted": json.dumps({})}`, + variables: {} + }); + expect(first.success).toBe(true); + expect(first.data?.codeReturn.polluted).toBe('polluted'); + + const second = await r.execute({ + code: `import json +def main(): + try: + leaked + has_leaked = True + except NameError: + has_leaked = False + return {"json": json.dumps({"a": 1}), "has_leaked": has_leaked}`, + variables: {} + }); + expect(second.success).toBe(true); + expect(second.data?.codeReturn.json).toBe('{"a": 1}'); + expect(second.data?.codeReturn.has_leaked).toBe(false); + }); + + it('每个任务使用独立临时目录，结束后由父进程清理', async () => { + const r = await createRunner(1); + + const result = await r.execute({ + code: `import pandas as pd +def main(): + path = task_tmpdir + '/allowed.csv' + pd.DataFrame({'a': [1]}).to_csv(path, index=False) + return {"tmp": task_tmpdir}`, + variables: {} + }); + + expect(result.success).toBe(true); + const taskTmp = result.data?.codeReturn.tmp; + expect(taskTmp).toMatch(/task-/); + const hostTaskTmp = shouldEnablePythonNativeIsolation() + ? `${PYTHON_SANDBOX_ROOT}${taskTmp}` + : taskTmp; + expect(existsSync(hostTaskTmp)).toBe(false); + }); + + it('预热进程执行一次后销毁，不归还给后续任务复用', async () => { + const r = await createRunner(1); + + const firstIdlePid = (r as any).idleChildren.values().next().value?.proc?.pid; + expect(firstIdlePid).toBeTruthy(); + + const first = await r.execute({ + code: `def main(): + return {"ok": True}`, + variables: {} + }); + expect(first.success).toBe(true); + expect(first.data?.codeReturn.ok).toBe(true); + + const secondIdlePid = await waitForIdlePid(r, firstIdlePid); + expect(secondIdlePid).toBeTruthy(); + expect(secondIdlePid).not.toBe(firstIdlePid); + + const second = await r.execute({ + code: `def main(): + return {"ok": True}`, + variables: {} + }); + expect(second.success).toBe(true); + expect(second.data?.codeReturn.ok).toBe(true); + }); + + it('并发超过上限时排队执行', async () => { + const r = await createRunner(1); + + const p1 = r.execute({ + code: `import time +def main(idx): + time.sleep(0.2) + return {"idx": idx}`, + variables: { idx: 1 } + }); + const p2 = r.execute({ + code: `def main(idx): + return {"idx": idx}`, + variables: { idx: 2 } + }); + + expect(r.stats.queued).toBe(1); + const results = await Promise.all([p1, p2]); + expect(results[0].success).toBe(true); + expect(results[1].success).toBe(true); + expect(results.map((item) => item.data?.codeReturn.idx)).toEqual([1, 2]); + }); + + it('高并发快速任务不会在 stdout drain 前被误判为无结果', async () => { + const r = await createRunner(20); + + const results = await Promise.all( + Array.from({ length: 80 }, (_, idx) => + r.execute({ + code: `def main(idx): + return {"idx": idx}`, + variables: { idx } + }) + ) + ); + + expect(results.every((item) => item.success)).toBe(true); + expect(results.map((item) => item.data?.codeReturn.idx).sort((a, b) => a - b)).toEqual( + Array.from({ length: 80 }, (_, idx) => idx) + ); + }); +}); + +describe('PythonIsolatedRunner 安全回归', () => { + let runner: PythonIsolatedRunner | undefined; + + afterEach(async () => { + await runner?.shutdown(); + runner = undefined; + }); + + async function createRunner() { + runner = new PythonIsolatedRunner(1); + await runner.init(); + return runner; + } + + it('阻断 GHSA-5jmh-5f2m-89jg 字符串拼接 __subclasses__ 绕过', async () => { + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + base = (1).__class__.__base__ + subs = getattr(base, "__subcl" + "asses__")() + for c in subs: + g = getattr(getattr(c, "__init__", None), "__globals__", None) + if g and "popen" in g: + return {"result": g["popen"]("id").read()} + return {"result": "os not found"}`, + variables: {} + }); + + expect(result.success).toBe(false); + expect(result.message).not.toMatch(/uid=/); + expect(result.message).toMatch(/__class__|Dynamic getattr|not allowed/i); + }); + + it('允许 type() 但阻断通过 type().__base__ 继续反射逃逸', async () => { + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + base = type(1).__base__ + return {"count": len(getattr(base, "__subclasses__")())}`, + variables: {} + }); + + expect(result.success).toBe(false); + expect(result.message).toMatch(/__base__|__subclasses__|not allowed/i); + }); + + it('阻断直接 import os/subprocess', async () => { + const r = await createRunner(); + + const osResult = await r.execute({ + code: `import os +def main(): + return {"cwd": os.getcwd()}`, + variables: {} + }); + expect(osResult.success).toBe(false); + expect(osResult.message).toContain('os'); + + const subprocessResult = await r.execute({ + code: `import subprocess +def main(): + return {"out": subprocess.check_output(["id"]).decode()}`, + variables: {} + }); + expect(subprocessResult.success).toBe(false); + expect(subprocessResult.message).toContain('subprocess'); + }); +}); + +describe('PythonIsolatedRunner HTTP 父进程代理', () => { + let runner: PythonIsolatedRunner | undefined; + let server: http.Server | undefined; + + afterEach(async () => { + await runner?.shutdown(); + runner = undefined; + await new Promise((resolve) => { + if (!server) return resolve(); + server.close(() => resolve()); + server = undefined; + }); + }); + + async function createRunner() { + runner = new PythonIsolatedRunner(1); + await runner.init(); + return runner; + } + + async function startPublicLocalServer() { + server = http.createServer((req, res) => { + let body = ''; + req.on('data', (chunk) => { + body += chunk.toString('utf8'); + }); + req.on('end', () => { + res.setHeader('content-type', 'application/json'); + res.end(JSON.stringify({ method: req.method, body })); + }); + }); + await new Promise((resolve) => server!.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (!address || typeof address === 'string') throw new Error('Failed to start test server'); + return address.port; + } + + it('http_request 通过 Node 代理层执行内网拦截', async () => { + const port = await startPublicLocalServer(); + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + try: + http_request('http://127.0.0.1:${port}/echo', method='POST', body={'hello': 'world'}) + return {'blocked': False} + except Exception as e: + return {'blocked': True, 'msg': str(e)}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toBe(true); + expect(result.data?.codeReturn.msg).toMatch(/private|internal|not allowed/i); + }); + + it('请求次数限制由父进程按单次执行计数', async () => { + const r = await createRunner(); + + const result = await r.execute({ + code: `def main(): + limit_error = None + for i in range(35): + try: + http_request('http://0.0.0.0:1') + except Exception as e: + if 'limit' in str(e).lower(): + limit_error = {'idx': i, 'msg': str(e)} + break + return {'limit_error': limit_error}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.limit_error).not.toBeNull(); + }); +}); diff --git a/projects/code-sandbox/test/unit/python-isolated-security.test.ts b/projects/code-sandbox/test/unit/python-isolated-security.test.ts new file mode 100644 index 000000000000..3f55a9c2b346 --- /dev/null +++ b/projects/code-sandbox/test/unit/python-isolated-security.test.ts @@ -0,0 +1,457 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest'; +import { statSync } from 'fs'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; +import { shouldEnablePythonNativeIsolation } from '../../src/isolated/python-isolation-config'; + +let runner: PythonIsolatedRunner; + +beforeAll(async () => { + runner = new PythonIsolatedRunner(1); + await runner.init(); +}); + +afterAll(async () => { + await runner.shutdown(); +}); + +describe('PythonIsolatedRunner 安全规则兼容旧 pool', () => { + it.each(['os', 'subprocess', 'sys', 'socket', 'threading', 'multiprocessing', 'signal'])( + '阻止 import %s', + async (moduleName) => { + const result = await runner.execute({ + code: `def main(): + try: + __import__(${JSON.stringify(moduleName)}) + return {'blocked': False} + except Exception as e: + return {'blocked': True, 'error': str(e)}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toBe(true); + expect(result.data?.codeReturn.error).toContain(moduleName); + } + ); + + it('builtins.__import__ 覆盖不能恢复危险 import', async () => { + const result = await runner.execute({ + code: `import builtins +def main(): + changed = False + try: + builtins.__import__ = lambda *a, **kw: None + changed = True + except Exception: + pass + try: + import os + escaped = True + except Exception: + escaped = False + return {'changed': changed, 'escaped': escaped}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.changed).toBe(false); + expect(result.data?.codeReturn.escaped).toBe(false); + }); + + it('运行时拼接 getattr 不能访问高危 dunder 属性', async () => { + const result = await runner.execute({ + code: `def main(): + name = "__" + "subclasses__" + return {'value': getattr(object, name)()}`, + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toMatch(/__subclasses__|not allowed/i); + }); + + it('setattr/delattr 即使通过变量名调用也不能改写 builtins 安全函数', async () => { + const result = await runner.execute({ + code: `import builtins +def main(): + setter = setattr + deleter = delattr + result = [] + for fn in (setter, deleter): + try: + if fn is setter: + fn(builtins, "__import__", lambda *a, **kw: None) + else: + fn(builtins, "__import__") + result.append(False) + except Exception: + result.append(True) + return {'blocked': result}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toEqual([true, true]); + }); + + it('即使通过允许的标准库间接拿到 os，命令执行被阻断且宿主 secret env 不泄露', async () => { + const oldToken = process.env.SANDBOX_TOKEN; + process.env.SANDBOX_TOKEN = 'test-host-secret-token'; + try { + const result = await runner.execute({ + code: `import platform +def main(): + os_ref = getattr(platform, 'os') + system_blocked = False + try: + os_ref.system('id') + except Exception: + system_blocked = True + env = dict(os_ref.environ) + return { + 'system_blocked': system_blocked, + 'has_sandbox_token': 'SANDBOX_TOKEN' in env, + 'has_test_secret': 'test-host-secret-token' in ''.join(env.values()) + }`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.system_blocked).toBe(true); + expect(result.data?.codeReturn.has_sandbox_token).toBe(false); + expect(result.data?.codeReturn.has_test_secret).toBe(false); + } finally { + if (oldToken === undefined) { + delete process.env.SANDBOX_TOKEN; + } else { + process.env.SANDBOX_TOKEN = oldToken; + } + } + }); + + it.each([ + ['exec', `exec("import subprocess")`], + ['eval', `eval("__import__('os')")`], + ['compile', `exec(compile("import subprocess", "", "exec"))`] + ])('%s 逃逸被拦截', async (_name, statement) => { + const result = await runner.execute({ + code: `def main(): + try: + ${statement} + return {'escaped': True} + except Exception: + return {'escaped': False}`, + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toMatch(/exec|eval|compile|not allowed/i); + }); + + it.each([ + ['object.__subclasses__', `object.__subclasses__()`], + ['__class__.__bases__ 链式访问', `().__class__.__bases__[0].__subclasses__()`], + ['getattr 常量 dunder', `getattr(object, "__subclasses__")()`], + ['getattr 动态 dunder', `getattr(object, "__sub" + "classes__")()`], + ['type().__base__ 链式访问', `type(1).__base__.__subclasses__()`] + ])('阻断 %s 反射逃逸', async (_name, expression) => { + const result = await runner.execute({ + code: `def main(): + return {'value': ${expression}}`, + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toMatch( + /__class__|__base__|__bases__|__subclasses__|Dynamic getattr|not allowed/i + ); + }); + + it('保留 type() 正常判断与动态创建类兼容', async () => { + const result = await runner.execute({ + code: `def main(): + MyClass = type('MyClass', (object,), {'x': 42}) + obj = MyClass() + return {'is_int': type(1) == int, 'x': obj.x}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ is_int: true, x: 42 }); + }); + + it.each(['/etc/passwd', '/proc/self/environ'])('阻止 open 读取 %s', async (path) => { + const result = await runner.execute({ + code: `def main(): + with open(${JSON.stringify(path)}, 'r') as f: + return {'data': f.read()}`, + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toContain('not allowed'); + }); + + it('阻止 open 写入文件', async () => { + const result = await runner.execute({ + code: `def main(): + with open('/tmp/evil.txt', 'w') as f: + f.write('hacked') + return {'ok': True}`, + variables: {} + }); + expect(result.success).toBe(false); + expect(result.message).toContain('not allowed'); + }); + + it('阻止标准库或三方库间接写入共享 /tmp，只允许当前任务临时目录', async () => { + const result = await runner.execute({ + code: `import pandas as pd +def main(): + pd.DataFrame({'a': [1]}).to_csv(task_tmpdir + '/allowed.csv', index=False) + try: + pd.DataFrame({'a': [2]}).to_csv('/tmp/shared.csv', index=False) + shared_tmp_blocked = False + except Exception: + shared_tmp_blocked = True + return {'tmp': task_tmpdir, 'shared_tmp_blocked': shared_tmp_blocked}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.tmp).toMatch(/task-/); + expect(result.data?.codeReturn.shared_tmp_blocked).toBe(true); + }); + + it('阻止通过允许标准库间接拿到 os 后读取宿主文件系统', async () => { + const result = await runner.execute({ + code: `import platform +def main(): + os_ref = platform.os + read_blocked = False + list_blocked = False + stat_blocked = False + try: + fd = os_ref.open('/etc/passwd', os_ref.O_RDONLY) + os_ref.close(fd) + except Exception: + read_blocked = True + try: + os_ref.listdir('/') + except Exception: + list_blocked = True + try: + os_ref.stat('/etc/passwd') + except Exception: + stat_blocked = True + return { + 'read_blocked': read_blocked, + 'list_blocked': list_blocked, + 'stat_blocked': stat_blocked + }`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn).toEqual({ + read_blocked: true, + list_blocked: true, + stat_blocked: true + }); + }); + + it('允许通过标准库 os 引用访问当前任务临时目录', async () => { + const result = await runner.execute({ + code: `import platform +def main(): + os_ref = platform.os + path = task_tmpdir + '/allowed.txt' + fd = os_ref.open(path, os_ref.O_CREAT | os_ref.O_WRONLY, 0o600) + os_ref.write(fd, b'ok') + os_ref.close(fd) + fd = os_ref.open(path, os_ref.O_RDONLY) + data = os_ref.read(fd, 16).decode() + os_ref.close(fd) + return {'data': data, 'items': os_ref.listdir(task_tmpdir)}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.data).toBe('ok'); + expect(result.data?.codeReturn.items).toContain('allowed.txt'); + }); + + it('允许三方库在当前任务临时目录内创建子目录，但不能写共享 /tmp', async () => { + const result = await runner.execute({ + code: `import platform +def main(): + os_ref = platform.os + nested = task_tmpdir + '/nested/cache' + os_ref.makedirs(nested, exist_ok=True) + try: + os_ref.makedirs('/tmp/shared-blocked', exist_ok=True) + outside_created = True + outside_error = '' + except Exception as e: + outside_created = False + outside_error = str(e) + return { + 'nested_exists': os_ref.path.isdir(nested), + 'outside_created': outside_created, + 'outside_error': outside_error + }`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.nested_exists).toBe(true); + expect(result.data?.codeReturn.outside_created).toBe(false); + expect(result.data?.codeReturn.outside_error).toMatch(/task temporary directory|not allowed/i); + }); + + it('matplotlib 可以使用任务临时目录初始化 config/cache', async () => { + const result = await runner.execute({ + code: `import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +def main(): + fig, ax = plt.subplots(figsize=(2, 1)) + ax.plot([1, 2, 3], [1, 4, 9]) + plt.close(fig) + return { + 'backend': matplotlib.get_backend(), + 'config': matplotlib.get_configdir(), + 'cache': matplotlib.get_cachedir(), + 'figure_axes': len(fig.axes) + }`, + variables: {} + }); + + expect(result.success, JSON.stringify(result)).toBe(true); + expect(result.data?.codeReturn.backend.toLowerCase()).toContain('agg'); + expect(result.data?.codeReturn.config).toContain('/matplotlib'); + expect(result.data?.codeReturn.cache).toContain('/matplotlib'); + expect(result.data?.codeReturn.figure_axes).toBe(1); + }); + + it('Linux native 隔离下 chroot /tmp 由 root 持有，仅 task 临时目录可写', async () => { + if (!shouldEnablePythonNativeIsolation()) { + return; + } + + const result = await runner.execute({ + code: `def main(): + return {'tmp': task_tmpdir}`, + variables: {} + }); + + expect(result.success).toBe(true); + const tmpStat = statSync('/tmp/fastgpt-python-sandbox/tmp'); + expect(tmpStat.uid).toBe(0); + expect(tmpStat.mode & 0o777).toBe(0o755); + }); + + it('变量值包含 Python 代码不会被执行', async () => { + const result = await runner.execute({ + code: `def main(v): + return {'val': v['code']}`, + variables: { code: '__import__("os").system("id")' } + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.val).toBe('__import__("os").system("id")'); + }); + + it('无法通过 os 模块读取环境变量', async () => { + const result = await runner.execute({ + code: `def main(): + try: + import os + return {'blocked': False, 'env': dict(os.environ)} + except Exception as e: + return {'blocked': True, 'error': str(e)}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.blocked).toBe(true); + }); + + it('上一次执行设置的全局变量，下一次读不到', async () => { + const r1 = await runner.execute({ + code: `def main(): + global secret_data + secret_data = 'leaked_password_123' + return {'written': True}`, + variables: {} + }); + expect(r1.success).toBe(true); + + const r2 = await runner.execute({ + code: `def main(): + try: + return {'leaked': True, 'val': secret_data} + except NameError: + return {'leaked': False}`, + variables: {} + }); + expect(r2.success).toBe(true); + expect(r2.data?.codeReturn.leaked).toBe(false); + }); + + it('上一次修改的模块状态不影响下一次', async () => { + const r1 = await runner.execute({ + code: `import json +def main(): + json._polluted = True + return {'polluted': hasattr(json, '_polluted')}`, + variables: {} + }); + expect(r1.success).toBe(true); + expect(r1.data?.codeReturn.polluted).toBe(true); + + const r2 = await runner.execute({ + code: `import json +def main(): + return { + 'has_pollution': hasattr(json, '_polluted'), + 'dumps_works': json.dumps({'test': 1}) == '{"test": 1}' + }`, + variables: {} + }); + expect(r2.success).toBe(true); + expect(r2.data?.codeReturn.has_pollution).toBe(false); + expect(r2.data?.codeReturn.dumps_works).toBe(true); + }); + + it('上一次的 print 输出不泄露到下一次', async () => { + await runner.execute({ + code: `def main(): + print('secret_token_abc123') + return {}`, + variables: {} + }); + + const result = await runner.execute({ + code: `def main(): + return {'ok': True}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.log || '').not.toContain('secret_token_abc123'); + }); + + it('上一次传入的 variables 不泄露到下一次', async () => { + await runner.execute({ + code: `def main(v): + return {'got': v['apiKey']}`, + variables: { apiKey: 'sk-secret-key-12345' } + }); + + const result = await runner.execute({ + code: `def main(v): + leaked = [] + if v and 'apiKey' in v: + leaked.append('apiKey from vars') + try: + _ = apiKey + leaked.append('apiKey from global') + except NameError: + pass + return {'clean': len(leaked) == 0, 'leaked': leaked}`, + variables: {} + }); + expect(result.success).toBe(true); + expect(result.data?.codeReturn.clean).toBe(true); + }); +}); diff --git a/projects/code-sandbox/test/unit/python-native-isolation.test.ts b/projects/code-sandbox/test/unit/python-native-isolation.test.ts new file mode 100644 index 000000000000..40e2b3fa3d6c --- /dev/null +++ b/projects/code-sandbox/test/unit/python-native-isolation.test.ts @@ -0,0 +1,62 @@ +import { existsSync, mkdtempSync, writeFileSync } from 'fs'; +import { tmpdir } from 'os'; +import { join } from 'path'; +import { spawnSync } from 'child_process'; +import { describe, expect, it } from 'vitest'; +import { + PYTHON_SANDBOX_GID, + PYTHON_SANDBOX_UID, + shouldEnablePythonNativeIsolation +} from '../../src/isolated/python-isolation-config'; + +const nativeLibraryPath = join(process.cwd(), 'dist', 'fastgpt_python_sandbox.so'); +const shouldRunNativeIsolation = + shouldEnablePythonNativeIsolation() && existsSync(nativeLibraryPath); + +/** + * 直接验证 Go native 隔离库，而不是通过 PythonIsolatedRunner。 + * + * Runner 的语言层会先拦截 os/subprocess/socket，这里用专用脚本在调用 + * FastGPTInitPythonSandbox 后直接尝试危险能力，确保 seccomp/chroot/setuid + * 本身也能形成边界。 + */ +describe.skipIf(!shouldRunNativeIsolation)('Python native seccomp/chroot isolation', () => { + it('降权到 sandbox uid/gid，并阻断 os.system 的 execve 落地', () => { + const sandboxRoot = mkdtempSync(join(tmpdir(), 'fastgpt-native-sandbox-')); + const probeScript = join(tmpdir(), `fastgpt-native-probe-${Date.now()}.py`); + + writeFileSync( + probeScript, + ` +import ctypes +import json +import os +import sys + +lib = ctypes.CDLL(${JSON.stringify(nativeLibraryPath)}) +lib.FastGPTInitPythonSandbox.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int] +lib.FastGPTInitPythonSandbox.restype = ctypes.c_int +ret = lib.FastGPTInitPythonSandbox(${PYTHON_SANDBOX_UID}, ${PYTHON_SANDBOX_GID}, 0) +if ret != 0: + print(json.dumps({"init": ret})) + sys.exit(1) + +print(json.dumps({"uid": os.getuid(), "gid": os.getgid()}), flush=True) +rc = os.system("id") +print(json.dumps({"system_rc": rc}), flush=True) +`, + 'utf8' + ); + + const result = spawnSync('python3', ['-u', probeScript], { + cwd: sandboxRoot, + encoding: 'utf8', + timeout: 5000 + }); + + expect(result.stdout).toContain(`"uid": ${PYTHON_SANDBOX_UID}`); + expect(result.stdout).toContain(`"gid": ${PYTHON_SANDBOX_GID}`); + expect(result.stdout).not.toContain('"system_rc": 0'); + expect(result.stdout + result.stderr).not.toMatch(/uid=\d+/); + }); +}); diff --git a/projects/code-sandbox/test/unit/resource-limits.test.ts b/projects/code-sandbox/test/unit/resource-limits.test.ts index f798c6f6392c..e22a8334f76e 100644 --- a/projects/code-sandbox/test/unit/resource-limits.test.ts +++ b/projects/code-sandbox/test/unit/resource-limits.test.ts @@ -5,13 +5,14 @@ * - 内存限制（RSS 轮询监控） * - CPU 密集型超时（JS / Python） * - 运行时长限制（wall-clock timeout 验证） + * - Python 任务临时目录大小限制 * - 网络请求限制（次数、请求体大小、响应大小） */ import { describe, it, expect, afterEach, beforeAll } from 'vitest'; import { ProcessPool } from '../../src/pool/process-pool'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; -import { BaseProcessPool } from '../../src/pool/base-process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; +import type { BaseProcessPool } from '../../src/pool/base-process-pool'; import { env, RUNTIME_MEMORY_OVERHEAD_MB } from '../../src/env'; import { CustomModuleProcessPool, @@ -100,7 +101,7 @@ describe('内存限制', () => { }); describe('Python 内存限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -108,10 +109,11 @@ describe('Python 内存限制', () => { } catch {} }); - it('Python 分配超大内存被 RSS 监控终止后自动 respawn', async () => { - pool = new PythonProcessPool(1); + it('Python 分配超大内存被 RSS 监控终止后 runner 可继续执行', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); expect(pool.stats.total).toBe(1); + expect(pool.stats.idle).toBe(1); const result = await pool.execute({ code: `import time\nimport random\ndef main():\n chunks = []\n for i in range(40):\n chunk = bytearray(10 * 1024 * 1024)\n chunk[0] = i % 256\n chunks.append(chunk)\n time.sleep(0.2)\n return {'size': len(chunks)}`, @@ -120,9 +122,6 @@ describe('Python 内存限制', () => { expect(result.success).toBe(false); expect(result.message).toMatch(/memory|Memory|crash|Worker|timed out/i); - // 等 respawn - await new Promise((r) => setTimeout(r, 2000)); - const result2 = await pool.execute({ code: `def main():\n return {'recovered': True}`, variables: {} @@ -132,9 +131,10 @@ describe('Python 内存限制', () => { }, 30000); it('Python 分配配置范围内的内存正常工作', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); expect(pool.stats.total).toBe(1); + expect(pool.stats.idle).toBe(1); const allocMB = 10; @@ -195,35 +195,25 @@ describe('显式放开后台执行模块时的 worker 回收', () => { expect(recovery.data?.codeReturn.ok).toBe(true); }, 20000); - it('Python 显式允许 subprocess 后，任务返回即回收 worker 并清理子进程', async () => { - pool = new CustomModuleProcessPool('Python', ['subprocess'], { recycleAfterTask: true }); - await pool.init(); - - const firstPid = getWorkerPid(pool); - expect(firstPid).toBeTypeOf('number'); + it('Python isolated runner 拒绝 subprocess 且后续任务可继续执行', async () => { + const runner = new PythonIsolatedRunner(1); + await runner.init(); - const result = await pool.execute({ - code: `import subprocess\ndef main():\n child = subprocess.Popen(['python3', '-c', 'import time; time.sleep(60)'])\n return {'childPid': child.pid}`, + const result = await runner.execute({ + code: `import subprocess\ndef main():\n return {'never': True}`, variables: {} }); + expect(result.success).toBe(false); + expect(result.message).toContain('subprocess'); - expect(result.success).toBe(true); - const childPid = result.data?.codeReturn.childPid; - expect(childPid).toBeTypeOf('number'); - - const workerRecycled = await waitForCondition(() => { - const currentPid = getWorkerPid(pool); - return typeof currentPid === 'number' && currentPid !== firstPid; - }); - expect(workerRecycled).toBe(true); - expect(await waitForPidExit(childPid)).toBe(true); - - const recovery = await pool.execute({ + const recovery = await runner.execute({ code: `def main():\n return {'ok': True}`, variables: {} }); expect(recovery.success).toBe(true); expect(recovery.data?.codeReturn.ok).toBe(true); + + await runner.shutdown(); }, 20000); }); @@ -273,7 +263,7 @@ describe('JS CPU 密集型超时', () => { expect(result.message).toMatch(/timed out|timeout/i); }); - it('CPU 超时后 worker 恢复正常', async () => { + it('CPU 超时后 runner 可继续执行新任务', async () => { pool = new ProcessPool(1); await pool.init(); @@ -293,7 +283,7 @@ describe('JS CPU 密集型超时', () => { }); describe('Python CPU 密集型超时', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -302,7 +292,7 @@ describe('Python CPU 密集型超时', () => { }); it('纯计算死循环被超时终止', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const start = Date.now(); @@ -318,7 +308,7 @@ describe('Python CPU 密集型超时', () => { }); it('CPU 超时后 worker 恢复正常', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); await pool.execute({ @@ -326,8 +316,6 @@ describe('Python CPU 密集型超时', () => { variables: {} }); - await new Promise((r) => setTimeout(r, 2000)); - const r2 = await pool.execute({ code: `def main():\n return {'ok': True}`, variables: {} @@ -420,7 +408,7 @@ describe('JS 运行时长限制', () => { }); describe('Python 运行时长限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -429,7 +417,7 @@ describe('Python 运行时长限制', () => { }); it('sleep 超过超时限制被终止', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const start = Date.now(); @@ -444,25 +432,28 @@ describe('Python 运行时长限制', () => { expect(elapsed).toBeLessThan(env.SANDBOX_MAX_TIMEOUT + 10000); }); - it('超时后回收 Python worker，避免信号和模块状态残留', async () => { - pool = new PythonProcessPool(1); + it('超时后清理 Python 子进程，避免状态残留到下一次执行', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); - const pidBefore = (pool as any).workers[0].proc.pid; const result = await pool.execute({ code: `import time\ndef main():\n time.sleep(${Math.ceil(env.SANDBOX_MAX_TIMEOUT / 1000) + 5})\n return {'done': True}`, variables: {} }); - await new Promise((r) => setTimeout(r, 1500)); - const pidAfter = (pool as any).workers[0].proc.pid; expect(result.success).toBe(false); expect(result.message).toMatch(/timed out|timeout/i); - expect(pidAfter).not.toBe(pidBefore); + + const recovery = await pool.execute({ + code: `def main():\n return {'ok': True}`, + variables: {} + }); + expect(recovery.success).toBe(true); + expect(recovery.data?.codeReturn.ok).toBe(true); }); it('在超时范围内完成的代码正常返回', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -474,7 +465,7 @@ describe('Python 运行时长限制', () => { }); it('delay() 超过 10s 上限被拒绝', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -528,7 +519,7 @@ describe('JS 输出大小限制', () => { }); describe('Python 输出大小限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -536,11 +527,10 @@ describe('Python 输出大小限制', () => { } catch {} }); - it('返回值超过 maxOutputSize 被拒绝且 worker 可恢复', async () => { - pool = new PythonProcessPool(1); + it('返回值超过 maxOutputSize 被拒绝且 runner 可恢复', async () => { + pool = new PythonIsolatedRunner(1); await pool.init(); - const pidBefore = getWorkerPid(pool); const result = await pool.execute({ code: `def main():\n return 'x' * (${env.SANDBOX_MAX_OUTPUT_MB} * 1024 * 1024 + 1)`, variables: {} @@ -549,23 +539,78 @@ describe('Python 输出大小限制', () => { expect(result.success).toBe(false); expect(result.message).toMatch(/output too large/i); - const workerRecycled = await waitForCondition(() => { - const currentPid = getWorkerPid(pool); - return typeof currentPid === 'number' && currentPid !== pidBefore; + const recovery = await pool.execute({ + code: `def main():\n return {'ok': True}`, + variables: {} + }); + expect(recovery.success).toBe(true); + expect(recovery.data?.codeReturn.ok).toBe(true); + }, 20000); +}); + +// ============================================================ +// 5. Python 临时目录大小限制 +// ============================================================ +describe('Python 临时目录大小限制', () => { + let pool: PythonIsolatedRunner; + + afterEach(async () => { + try { + await pool?.shutdown(); + } catch {} + }); + + it('写入 task_tmpdir 超过 maxTmpSize 被终止且 runner 可恢复', async () => { + pool = new PythonIsolatedRunner(1); + await pool.init(); + + const result = await pool.execute({ + code: `import time +def main(): + path = task_tmpdir + '/too-large.bin' + chunk = b'x' * (1024 * 1024) + with open(path, 'wb') as f: + for _ in range(${env.SANDBOX_MAX_TMP_MB + 1}): + f.write(chunk) + f.flush() + time.sleep(2) + return {'done': True}`, + variables: {} }); - expect(workerRecycled).toBe(true); + + expect(result.success).toBe(false); + expect(result.message).toMatch(/temporary file limit/i); const recovery = await pool.execute({ - code: `def main():\n return {'ok': True}`, + code: `def main(): + with open(task_tmpdir + '/small.txt', 'w') as f: + f.write('ok') + return {'ok': True}`, variables: {} }); expect(recovery.success).toBe(true); expect(recovery.data?.codeReturn.ok).toBe(true); }, 20000); + + it('写入 task_tmpdir 限制内文件正常', async () => { + pool = new PythonIsolatedRunner(1); + await pool.init(); + + const result = await pool.execute({ + code: `def main(): + with open(task_tmpdir + '/allowed.bin', 'wb') as f: + f.write(b'x' * 1024) + return {'ok': True}`, + variables: {} + }); + + expect(result.success).toBe(true); + expect(result.data?.codeReturn.ok).toBe(true); + }); }); // ============================================================ -// 5. 网络请求次数限制 +// 6. 网络请求次数限制 // ============================================================ describe('JS 网络请求次数限制', () => { let pool: ProcessPool; @@ -639,7 +684,7 @@ describe('JS 网络请求次数限制', () => { }); describe('Python 网络请求次数限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -648,7 +693,7 @@ describe('Python 网络请求次数限制', () => { }); it(`第 maxRequests+1 次请求被拒绝（计数器验证）`, async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -663,7 +708,7 @@ describe('Python 网络请求次数限制', () => { }); it('请求计数每次执行重置', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); await pool.execute({ @@ -681,7 +726,7 @@ describe('Python 网络请求次数限制', () => { }); // ============================================================ -// 6. 网络请求大小限制 +// 7. 网络请求大小限制 // ============================================================ describe('JS 请求体大小限制', () => { let pool: ProcessPool; @@ -702,7 +747,7 @@ describe('JS 请求体大小限制', () => { code: `async function main() { const bigBody = 'x'.repeat(${sizeMB} * 1024 * 1024 + 1); try { - await httpRequest('https://1.1.1.1/cdn-cgi/trace', { method: 'POST', body: bigBody }); + await httpRequest('https://example.com/', { method: 'POST', body: bigBody }); return { blocked: false }; } catch(e) { return { blocked: true, msg: e.message }; @@ -723,7 +768,7 @@ describe('JS 请求体大小限制', () => { code: `async function main() { const smallBody = JSON.stringify({ data: 'hello' }); try { - await httpRequest('https://1.1.1.1/cdn-cgi/trace', { method: 'POST', body: smallBody }); + await httpRequest('https://example.com/', { method: 'POST', body: smallBody }); return { sizeOk: true }; } catch(e) { // 网络错误可以接受，但不应该是 body too large @@ -738,7 +783,7 @@ describe('JS 请求体大小限制', () => { }); describe('Python 请求体大小限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -747,12 +792,12 @@ describe('Python 请求体大小限制', () => { }); it('请求体超过 maxRequestBodySize 被拒绝', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const sizeMB = env.SANDBOX_REQUEST_MAX_BODY_MB; const result = await pool.execute({ - code: `def main():\n big_body = 'x' * (${sizeMB} * 1024 * 1024 + 1)\n try:\n http_request('https://1.1.1.1/cdn-cgi/trace', method='POST', body=big_body)\n return {'blocked': False}\n except Exception as e:\n return {'blocked': True, 'msg': str(e)}`, + code: `def main():\n big_body = 'x' * (${sizeMB} * 1024 * 1024 + 1)\n try:\n http_request('https://example.com/', method='POST', body=big_body)\n return {'blocked': False}\n except Exception as e:\n return {'blocked': True, 'msg': str(e)}`, variables: {} }); expect(result.success).toBe(true); @@ -761,11 +806,11 @@ describe('Python 请求体大小限制', () => { }); it('请求体在限制内正常发送（不因大小被拒）', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ - code: `def main():\n try:\n http_request('https://1.1.1.1/cdn-cgi/trace', method='POST', body='hello')\n return {'size_ok': True}\n except Exception as e:\n return {'size_ok': 'too large' not in str(e).lower(), 'msg': str(e)}`, + code: `def main():\n try:\n http_request('https://example.com/', method='POST', body='hello')\n return {'size_ok': True}\n except Exception as e:\n return {'size_ok': 'too large' not in str(e).lower(), 'msg': str(e)}`, variables: {} }); expect(result.success).toBe(true); @@ -774,7 +819,7 @@ describe('Python 请求体大小限制', () => { }); // ============================================================ -// 7. 网络协议限制 +// 8. 网络协议限制 // ============================================================ describe('JS 网络协议限制', () => { let pool: ProcessPool; @@ -826,7 +871,7 @@ describe('JS 网络协议限制', () => { }); describe('Python 网络协议限制', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -835,7 +880,7 @@ describe('Python 网络协议限制', () => { }); it('ftp:// 协议被拒绝', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -848,7 +893,7 @@ describe('Python 网络协议限制', () => { }); it('file:// 协议被拒绝', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ diff --git a/projects/code-sandbox/test/unit/security.test.ts b/projects/code-sandbox/test/unit/security.test.ts index 080f47f575c1..1f5035c572a1 100644 --- a/projects/code-sandbox/test/unit/security.test.ts +++ b/projects/code-sandbox/test/unit/security.test.ts @@ -12,16 +12,16 @@ */ import { afterEach, describe, it, expect, beforeAll, afterAll } from 'vitest'; import { ProcessPool } from '../../src/pool/process-pool'; -import { PythonProcessPool } from '../../src/pool/python-process-pool'; +import { PythonIsolatedRunner } from '../../src/isolated/python-isolated-runner'; import { CustomModuleProcessPool } from '../helpers/custom-process-pool'; let jsPool: ProcessPool; -let pyPool: PythonProcessPool; +let pyPool: PythonIsolatedRunner; beforeAll(async () => { jsPool = new ProcessPool(1); await jsPool.init(); - pyPool = new PythonProcessPool(1); + pyPool = new PythonIsolatedRunner(1); await pyPool.init(); }); @@ -73,6 +73,7 @@ describe('模块拦截', () => { 'fs/promises', 'child_process', 'http', + 'url', 'lodash' ]); await pool.init(); @@ -82,7 +83,8 @@ describe('模块拦截', () => { `async function main() { const fs = require('node:fs'); return { ok: typeof fs.readFileSync === 'function' }; }`, `async function main() { const fs = require('fs/promises'); return { ok: typeof fs.readFile === 'function' }; }`, `async function main() { const cp = require('child_process'); return { ok: typeof cp.execFile === 'function' }; }`, - `async function main() { const http = require('http'); return { ok: typeof http.request === 'function' }; }` + `async function main() { const http = require('http'); return { ok: typeof http.request === 'function' }; }`, + `async function main() { const url = require('url'); const parsed = url.parse('https://example.com/a?b=1', true); return { ok: parsed.hostname === 'example.com' && parsed.query.b === '1' }; }` ]; for (const code of payloads) { @@ -278,33 +280,6 @@ describe('模块拦截', () => { expect(result.success).toBe(false); }); - it('显式加入白名单后允许危险 Python 标准库', async () => { - const pool = new CustomModuleProcessPool('Python', [ - 'os', - 'subprocess', - 'socket', - 'pathlib', - 'json' - ]); - await pool.init(); - try { - const payloads = [ - `import os\ndef main():\n return {'ok': hasattr(os, 'getcwd')}`, - `import subprocess\ndef main():\n return {'ok': hasattr(subprocess, 'Popen')}`, - `import socket\ndef main():\n return {'ok': hasattr(socket, 'socket')}`, - `from pathlib import Path\ndef main():\n return {'ok': Path('/tmp').name == 'tmp'}` - ]; - - for (const code of payloads) { - const result = await pool.execute({ code, variables: {} }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.ok).toBe(true); - } - } finally { - await pool.shutdown(); - } - }); - it('阻止 import requests（预检）', async () => { const result = await runner.execute({ code: `import requests\ndef main():\n return {}`, @@ -627,7 +602,7 @@ def main(): }); // --- exec/eval 逃逸 --- - it('exec 中导入危险模块被 __import__ hook 拦截', async () => { + it('exec 中导入危险模块在预检阶段被拒绝', async () => { const result = await runner.execute({ code: ` def main(): @@ -639,8 +614,8 @@ def main(): `, variables: {} }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.escaped).toBe(false); + expect(result.success).toBe(false); + expect(result.message).toMatch(/exec|not allowed/i); }); it('exec 字符串拼接绕过预检（运行时拦截兜底）', async () => { @@ -672,11 +647,11 @@ def main(): `, variables: {} }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.escaped).toBe(false); + expect(result.success).toBe(false); + expect(result.message).toMatch(/eval|not allowed/i); }); - it('compile + exec 导入危险模块被拦截', async () => { + it('compile + exec 导入危险模块在预检阶段被拒绝', async () => { const result = await runner.execute({ code: ` def main(): @@ -689,8 +664,8 @@ def main(): `, variables: {} }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.escaped).toBe(false); + expect(result.success).toBe(false); + expect(result.message).toMatch(/compile|exec|not allowed/i); }); // --- 内部变量隔离 --- @@ -742,7 +717,7 @@ def main(): expect(result.data?.codeReturn.escaped).toBe(false); }); - it('globals() 不泄露内部变量', async () => { + it('globals() 调用被拒绝，避免泄露内部变量', async () => { const result = await runner.execute({ code: `def main(v): g = globals() @@ -750,8 +725,8 @@ def main(): return {"has_original_import": has_orig}`, variables: {} }); - expect(result.success).toBe(true); - expect(result.data?.codeReturn.has_original_import).toBe(false); + expect(result.success).toBe(false); + expect(result.message).toMatch(/globals|not allowed/i); }); // --- __subclasses__ / type --- @@ -927,18 +902,19 @@ describe('网络请求安全', () => { it('httpRequest GET 公网地址正常', async () => { const result = await runner.execute({ - code: `async function main() { const res = await SystemHelper.httpRequest('https://1.1.1.1/cdn-cgi/trace'); return { status: res.status, hasData: res.data.length > 0 }; }`, + code: `async function main() { const res = await SystemHelper.httpRequest('http://1.1.1.1/'); return { status: res.status, hasData: res.data.length > 0 }; }`, variables: {} }); expect(result.success).toBe(true); - expect(result.data?.codeReturn.status).toBe(200); + expect(result.data?.codeReturn.status).toBeGreaterThanOrEqual(200); + expect(result.data?.codeReturn.status).toBeLessThan(400); expect(result.data?.codeReturn.hasData).toBe(true); }); it('httpRequest POST 带 body', async () => { const result = await runner.execute({ code: `async function main() { - const res = await SystemHelper.httpRequest('https://1.1.1.1/cdn-cgi/trace', { method: 'POST', body: { key: 'value' } }); + const res = await SystemHelper.httpRequest('http://1.1.1.1/', { method: 'POST', body: { key: 'value' } }); return { hasStatus: typeof res.status === 'number' }; }`, variables: {} @@ -949,11 +925,12 @@ describe('网络请求安全', () => { it('全局函数 httpRequest 可用', async () => { const result = await runner.execute({ - code: `async function main() { const res = await httpRequest('https://1.1.1.1/cdn-cgi/trace'); return { status: res.status }; }`, + code: `async function main() { const res = await httpRequest('http://1.1.1.1/'); return { status: res.status }; }`, variables: {} }); expect(result.success).toBe(true); - expect(result.data?.codeReturn.status).toBe(200); + expect(result.data?.codeReturn.status).toBeGreaterThanOrEqual(200); + expect(result.data?.codeReturn.status).toBeLessThan(400); }); }); @@ -994,17 +971,18 @@ describe('网络请求安全', () => { it('http_request GET 公网地址正常', async () => { const result = await runner.execute({ - code: `def main():\n res = system_helper.http_request('https://1.1.1.1/cdn-cgi/trace')\n return {'status': res['status'], 'hasData': len(res['data']) > 0}`, + code: `def main():\n res = system_helper.http_request('http://1.1.1.1/')\n return {'status': res['status'], 'hasData': len(res['data']) > 0}`, variables: {} }); expect(result.success).toBe(true); - expect(result.data?.codeReturn.status).toBe(200); + expect(result.data?.codeReturn.status).toBeGreaterThanOrEqual(200); + expect(result.data?.codeReturn.status).toBeLessThan(400); expect(result.data?.codeReturn.hasData).toBe(true); }); it('http_request POST 带 body', async () => { const result = await runner.execute({ - code: `import json\ndef main():\n res = system_helper.http_request('https://1.1.1.1/cdn-cgi/trace', method='POST', body={'key': 'value'})\n return {'hasStatus': type(res['status']) == int}`, + code: `import json\ndef main():\n res = system_helper.http_request('http://1.1.1.1/', method='POST', body={'key': 'value'})\n return {'hasStatus': type(res['status']) == int}`, variables: {} }); expect(result.success).toBe(true); @@ -1013,11 +991,12 @@ describe('网络请求安全', () => { it('全局函数 http_request 可用', async () => { const result = await runner.execute({ - code: `def main():\n res = http_request('https://1.1.1.1/cdn-cgi/trace')\n return {'status': res['status']}`, + code: `def main():\n res = http_request('http://1.1.1.1/')\n return {'status': res['status']}`, variables: {} }); expect(result.success).toBe(true); - expect(result.data?.codeReturn.status).toBe(200); + expect(result.data?.codeReturn.status).toBeGreaterThanOrEqual(200); + expect(result.data?.codeReturn.status).toBeLessThan(400); }); }); }); @@ -1302,19 +1281,26 @@ describe('沙盒环境加固，禁止用户代码篡改沙盒环境', () => { describe('Python', () => { const runner = { execute: (args: any) => pyPool.execute(args) }; - it('builtins.__import__ 覆盖被静默忽略', async () => { + it('builtins.__import__ 覆盖被拒绝', async () => { const result = await runner.execute({ code: `import builtins def main(): - builtins.__import__ = lambda *a, **kw: None + changed = False + try: + builtins.__import__ = lambda *a, **kw: None + changed = True + except Exception: + pass try: import os - return {'escaped': True} + escaped = True except (ImportError, Exception): - return {'escaped': False}`, + escaped = False + return {'changed': changed, 'escaped': escaped}`, variables: {} }); expect(result.success).toBe(true); + expect(result.data?.codeReturn.changed).toBe(false); expect(result.data?.codeReturn.escaped).toBe(false); }); @@ -1588,8 +1574,8 @@ describe('worker 状态隔离', () => { }); }); - describe('Python Worker 状态隔离', () => { - let pool: PythonProcessPool; + describe('Python isolated runner 状态隔离', () => { + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -1598,7 +1584,7 @@ describe('worker 状态隔离', () => { }); it('上一次执行设置的全局变量，下一次读不到', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); // 第一次：尝试写入全局 @@ -1617,7 +1603,7 @@ describe('worker 状态隔离', () => { }); it('上一次修改的模块状态不影响下一次（模块快照恢复）', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); // 第一次：给 json 模块添加自定义属性 @@ -1638,7 +1624,7 @@ describe('worker 状态隔离', () => { }); it('上一次的 print 输出不泄露到下一次', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); await pool.execute({ @@ -1656,7 +1642,7 @@ describe('worker 状态隔离', () => { }); it('上一次传入的 variables 不泄露到下一次', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); await pool.execute({ @@ -1742,7 +1728,7 @@ describe('环境变量隔离', () => { }); describe('Python 环境变量隔离', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -1751,7 +1737,7 @@ describe('环境变量隔离', () => { }); it('无法通过 os 模块读取环境变量', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1763,7 +1749,7 @@ describe('环境变量隔离', () => { }); it('无法通过 subprocess 执行 env 命令', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1775,7 +1761,7 @@ describe('环境变量隔离', () => { }); it('无法通过 open 读取 /etc/passwd', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1867,7 +1853,7 @@ describe('进程干扰', () => { }); describe('Python 进程干扰防护', () => { - let pool: PythonProcessPool; + let pool: PythonIsolatedRunner; afterEach(async () => { try { @@ -1876,7 +1862,7 @@ describe('进程干扰', () => { }); it('无法 import subprocess', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1888,7 +1874,7 @@ describe('进程干扰', () => { }); it('无法 import multiprocessing', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1900,7 +1886,7 @@ describe('进程干扰', () => { }); it('无法 import signal 发送信号', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ @@ -1912,7 +1898,7 @@ describe('进程干扰', () => { }); it('无法 import threading 创建线程', async () => { - pool = new PythonProcessPool(1); + pool = new PythonIsolatedRunner(1); await pool.init(); const result = await pool.execute({ diff --git a/projects/code-sandbox/tsdown.config.ts b/projects/code-sandbox/tsdown.config.ts index 66ace437a91c..c8bb87d49780 100644 --- a/projects/code-sandbox/tsdown.config.ts +++ b/projects/code-sandbox/tsdown.config.ts @@ -13,7 +13,8 @@ import { defineConfig } from 'tsdown'; export default defineConfig({ entry: { index: 'src/index.ts', - worker: 'src/pool/worker.ts' + worker: 'src/pool/worker.ts', + 'python-isolated-runner': 'src/isolated/python-isolated-runner.ts' }, format: 'esm', platform: 'node', diff --git a/projects/code-sandbox/vitest.config.ts b/projects/code-sandbox/vitest.config.ts index 8836fafdd85b..2646a1219c08 100644 --- a/projects/code-sandbox/vitest.config.ts +++ b/projects/code-sandbox/vitest.config.ts @@ -18,7 +18,8 @@ export default defineConfig({ env: { CHECK_INTERNAL_IP: 'true', SANDBOX_API_MAX_BODY_MB: '1', - SANDBOX_MAX_TIMEOUT: '5000', + SANDBOX_MAX_MEMORY_MB: '256', + SANDBOX_MAX_TIMEOUT: '30000', SANDBOX_QUEUE_ID_CONCURRENCY: '1', SANDBOX_TOKEN: 'test' }