Compare commits

...

105 Commits

Author SHA1 Message Date
公明 dbcf9b8418 Update config.yaml 2026-07-02 18:05:23 +08:00
公明 b3767b2deb Add files via upload 2026-07-02 18:03:35 +08:00
公明 7e764df0e8 Add files via upload 2026-07-02 18:02:45 +08:00
公明 a1ffb20d6e Add files via upload 2026-07-02 17:58:06 +08:00
公明 125685f08f Add files via upload 2026-07-02 17:50:09 +08:00
公明 b804635fa8 Add files via upload 2026-07-02 12:11:18 +08:00
公明 c9fb5d11d3 Add files via upload 2026-07-02 12:08:52 +08:00
公明 926491b746 Add files via upload 2026-07-02 12:08:14 +08:00
公明 4e17691717 Add files via upload 2026-07-02 12:06:49 +08:00
公明 2e2a6dedd4 Add files via upload 2026-07-02 12:02:37 +08:00
公明 b1323896c8 Add files via upload 2026-07-02 11:55:23 +08:00
公明 595074b7b0 Add files via upload 2026-07-02 11:52:32 +08:00
公明 2e063dd857 Add files via upload 2026-07-02 11:51:27 +08:00
公明 a110d233e1 Add files via upload 2026-07-02 11:49:03 +08:00
公明 2f58d0a457 Add files via upload 2026-07-01 16:06:15 +08:00
公明 5b7f157802 Add files via upload 2026-07-01 15:56:51 +08:00
公明 09890db635 Add files via upload 2026-07-01 14:37:36 +08:00
公明 c0171ef60a Add files via upload 2026-07-01 14:34:50 +08:00
公明 4eb73fb638 Add files via upload 2026-07-01 14:32:50 +08:00
公明 d1b49cb20d Add files via upload 2026-07-01 14:30:58 +08:00
公明 930eb47013 Add files via upload 2026-07-01 14:29:58 +08:00
公明 9964e13197 Add files via upload 2026-07-01 14:27:05 +08:00
公明 4f7b21cb7e Update config.yaml 2026-07-01 10:49:31 +08:00
公明 9fae9db906 Delete internal/project/user_verbatim_anchor_test.go 2026-07-01 10:48:29 +08:00
公明 7ecd8c61e8 Delete internal/project/user_verbatim_anchor.go 2026-07-01 10:48:09 +08:00
公明 bdb0326e47 Add files via upload 2026-07-01 10:46:53 +08:00
公明 8dccc6aa06 Add files via upload 2026-07-01 10:44:27 +08:00
公明 fd4bbe8d76 Update config.yaml 2026-06-30 20:22:19 +08:00
公明 d80651e4d8 Add files via upload 2026-06-30 20:16:43 +08:00
公明 f920ff0a5d Update config.yaml 2026-06-30 20:15:26 +08:00
公明 ce8b57501d Add files via upload 2026-06-30 20:14:28 +08:00
公明 ecb38a3959 Add files via upload 2026-06-30 20:13:31 +08:00
公明 e69fdb71ca Add files via upload 2026-06-30 20:11:54 +08:00
公明 6aa1631748 Add files via upload 2026-06-30 20:10:36 +08:00
公明 52de3b0f41 Add files via upload 2026-06-30 20:09:18 +08:00
公明 e537e55198 Add files via upload 2026-06-30 20:07:28 +08:00
公明 dc20b4804e Update config.yaml 2026-06-30 19:55:00 +08:00
公明 6245d69364 Add files via upload 2026-06-30 19:53:44 +08:00
公明 ede32951bf Add files via upload 2026-06-30 19:52:30 +08:00
公明 866a8ebccf Add files via upload 2026-06-30 19:10:46 +08:00
公明 276b3f7ef5 Add files via upload 2026-06-30 18:39:26 +08:00
公明 81e461db54 Update config.yaml 2026-06-30 18:38:27 +08:00
公明 02cd488a3d Add files via upload 2026-06-30 18:06:15 +08:00
公明 b4b2f55665 Add files via upload 2026-06-30 18:04:16 +08:00
公明 7aa0ebea6d Add files via upload 2026-06-30 18:02:08 +08:00
公明 63ef4399f8 Add files via upload 2026-06-30 18:00:00 +08:00
公明 553d0ed6bf Add files via upload 2026-06-30 17:59:02 +08:00
公明 d92bbbea07 Add files via upload 2026-06-30 17:56:40 +08:00
公明 f89ad1b42d Add files via upload 2026-06-30 16:00:00 +08:00
公明 bbe14c1861 Add files via upload 2026-06-30 15:00:50 +08:00
公明 2fc37fefd1 Add files via upload 2026-06-30 14:38:49 +08:00
公明 ded8ac5a3f Add files via upload 2026-06-30 13:03:40 +08:00
公明 bf44cf58d3 Add files via upload 2026-06-30 11:55:32 +08:00
公明 6d390e80d5 Add files via upload 2026-06-30 11:34:38 +08:00
公明 cfc49ba16f Add files via upload 2026-06-30 11:06:29 +08:00
公明 d03f2fcf2b Add files via upload 2026-06-30 10:50:29 +08:00
公明 6e67684bba Add files via upload 2026-06-30 00:16:31 +08:00
公明 8f9d2f381a Add files via upload 2026-06-29 16:57:32 +08:00
公明 89c275269f Update config.yaml 2026-06-29 16:52:45 +08:00
公明 cb4900c61d Add files via upload 2026-06-29 16:51:54 +08:00
公明 5c192cd308 Add files via upload 2026-06-29 16:46:26 +08:00
公明 8571e41138 Add files via upload 2026-06-29 16:24:43 +08:00
公明 e1a74b29b1 Add files via upload 2026-06-29 16:16:59 +08:00
公明 39f1c72755 Add files via upload 2026-06-29 14:35:52 +08:00
公明 dd3621e89d Add files via upload 2026-06-29 14:18:08 +08:00
公明 0bcb16e021 Add files via upload 2026-06-29 10:41:42 +08:00
公明 ed64803a51 Update config.yaml 2026-06-28 01:15:40 +08:00
公明 25e03dee84 Add files via upload 2026-06-28 01:15:10 +08:00
公明 58dcafd15f Add files via upload 2026-06-28 00:56:22 +08:00
公明 997c4e7262 Add files via upload 2026-06-27 01:44:08 +08:00
公明 ac370b0ada Add files via upload 2026-06-27 01:42:44 +08:00
公明 017db2b9a8 Add files via upload 2026-06-27 01:41:36 +08:00
公明 86b4803683 Add files via upload 2026-06-27 01:40:12 +08:00
公明 4d98264fc3 Add files via upload 2026-06-27 01:38:02 +08:00
公明 fd1de4ea94 Add files via upload 2026-06-27 01:36:09 +08:00
公明 41ba3baca9 Add files via upload 2026-06-27 01:35:46 +08:00
公明 2e908daebb Add files via upload 2026-06-27 00:34:19 +08:00
公明 c1763e1b9a Add files via upload 2026-06-27 00:03:16 +08:00
公明 70e5d28619 Add files via upload 2026-06-26 23:54:29 +08:00
公明 49990ecb4f Add files via upload 2026-06-26 23:50:13 +08:00
公明 c91806c0c4 Add files via upload 2026-06-26 23:11:52 +08:00
公明 e537236bf3 Add files via upload 2026-06-26 23:10:11 +08:00
公明 7eeffb1933 Add files via upload 2026-06-26 18:16:30 +08:00
公明 0556b29d40 Add files via upload 2026-06-26 14:34:45 +08:00
公明 be3c0cfa64 Add files via upload 2026-06-26 14:31:47 +08:00
公明 8e5f40d226 Add files via upload 2026-06-26 14:30:00 +08:00
公明 4b6719a6f3 Add files via upload 2026-06-26 14:27:32 +08:00
公明 7c8f3228f8 Add files via upload 2026-06-26 14:25:14 +08:00
公明 537843b6b8 Add files via upload 2026-06-26 14:24:01 +08:00
公明 4a57574cf9 Add files via upload 2026-06-26 14:21:51 +08:00
公明 0168530084 Add files via upload 2026-06-26 10:57:59 +08:00
公明 4184a7b6f0 Add files via upload 2026-06-26 10:54:59 +08:00
公明 fb3b4dd6e5 Add files via upload 2026-06-26 01:22:30 +08:00
公明 7e4a8db7af Add files via upload 2026-06-26 01:01:49 +08:00
公明 6a72c95b9f Add files via upload 2026-06-26 00:58:29 +08:00
公明 447be050cd Add files via upload 2026-06-25 21:28:46 +08:00
公明 9b75c43f7b Add files via upload 2026-06-25 15:15:01 +08:00
公明 a443454753 Add files via upload 2026-06-25 14:56:56 +08:00
公明 08822ba5df Update config.yaml 2026-06-25 14:56:31 +08:00
公明 eda75fb98f Add files via upload 2026-06-25 14:55:10 +08:00
公明 e6978a7994 Add files via upload 2026-06-25 14:52:39 +08:00
公明 1db0f4740f Add files via upload 2026-06-25 14:50:28 +08:00
公明 6e4ff96dcd Add files via upload 2026-06-25 14:48:25 +08:00
公明 95470fefbc Add files via upload 2026-06-25 14:47:16 +08:00
公明 5e075bb198 Add files via upload 2026-06-25 14:45:43 +08:00
121 changed files with 13949 additions and 1132 deletions
+39 -5
View File
@@ -35,7 +35,18 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
### System Dashboard Overview ### System Dashboard Overview
<img src="./images/dashboard.png" alt="System Dashboard" width="100%"> <table>
<tr>
<td width="50%" align="center">
<strong>Light Mode</strong><br/>
<img src="./images/dashboard.png" alt="System Dashboard (Light)" width="100%">
</td>
<td width="50%" align="center">
<strong>Dark Mode</strong><br/>
<img src="./images/dark.png" alt="System Dashboard (Dark)" width="100%">
</td>
</tr>
</table>
*The dashboard provides a comprehensive overview of system runtime status, security vulnerabilities, tool usage, and knowledge base, helping users quickly understand the platform's core features and current state.* *The dashboard provides a comprehensive overview of system runtime status, security vulnerabilities, tool usage, and knowledge base, helping users quickly understand the platform's core features and current state.*
@@ -110,7 +121,7 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
- 📄 Large-result pagination, compression, and searchable archives - 📄 Large-result pagination, compression, and searchable archives
- 🔗 Attack-chain graph, risk scoring, and step-by-step replay - 🔗 Attack-chain graph, risk scoring, and step-by-step replay
- 🔒 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): **Eino MultiQuery** query rewrite + multi-path vector retrieval + **HTTP rerank** (DashScope `gte-rerank` / Cohere-compatible) + post-processing (dedupe, budget); **Eino Compose** indexing pipeline
- 📁 Conversation grouping with pinning, rename, and batch management - 📁 Conversation grouping with pinning, rename, and batch management
- 📂 **Project management**: shared facts (blackboard) across sessions, `upsert_project_fact` + `links` to chain paths; attack-chain and project fact graph views - 📂 **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
@@ -455,9 +466,10 @@ A test SSE MCP server is available at `cmd/test-sse-mcp-server/` for validation
### Knowledge Base ### Knowledge Base
- **Vector search** AI agent can automatically search the knowledge base for relevant security knowledge during conversations using the `search_knowledge_base` tool. - **Vector search** AI agent can automatically search the knowledge base for relevant security knowledge during conversations using the `search_knowledge_base` tool.
- **Vector retrieval** cosine similarity over stored embeddings, aligned with Eino `retriever.Retriever` usage. - **RAG pipeline (always on)** **MultiQuery** (LLM query rewrite) → vector prefetch & fusion → **HTTP rerank** (DashScope `gte-rerank` or Cohere-compatible `/v1/rerank`) → post-processing (normalized dedupe, char/token budget, final top_k). Rerank failures degrade to fusion order without breaking search.
- **Auto-indexing** scans the `knowledge_base/` directory for Markdown files and automatically indexes them with embeddings. - **Vector retrieval** cosine similarity over stored embeddings with configurable threshold, aligned with Eino `retriever.Retriever` usage.
- **Web management** create, update, delete knowledge items through the web UI, with category-based organization. - **Auto-indexing** scans the `knowledge_base/` directory for Markdown files and automatically indexes them with embeddings (Markdown header split + recursive chunking via Eino).
- **Web management** create, update, delete knowledge items through the web UI, with category-based organization; settings page exposes MultiQuery / rerank / prefetch options.
- **Retrieval logs** tracks all knowledge retrieval operations for audit and debugging. - **Retrieval logs** tracks all knowledge retrieval operations for audit and debugging.
**Quick Start (Using Pre-built Knowledge Base):** **Quick Start (Using Pre-built Knowledge Base):**
@@ -479,6 +491,17 @@ A test SSE MCP server is available at `cmd/test-sse-mcp-server/` for validation
retrieval: retrieval:
top_k: 5 top_k: 5
similarity_threshold: 0.7 similarity_threshold: 0.7
multi_query:
max_queries: 4 # LLM rewrite variants (always on)
rerank: # always on; empty fields inherit openai/embedding credentials
provider: "" # auto: dashscope | cohere from base_url
model: "" # empty: gte-rerank (DashScope) or rerank-multilingual-v3.0 (Cohere)
base_url: ""
api_key: ""
post_retrieve:
prefetch_top_k: 20 # vector candidates per MultiQuery variant; 0 = max(top_k×4, 20)
max_context_chars: 0
max_context_tokens: 0
``` ```
2. **Add knowledge files** place Markdown files in `knowledge_base/` directory, organized by category (e.g., `knowledge_base/SQL Injection/README.md`). 2. **Add knowledge files** place Markdown files in `knowledge_base/` directory, organized by category (e.g., `knowledge_base/SQL Injection/README.md`).
3. **Scan and index** use the web UI to scan the knowledge base directory, which will automatically import files and build vector embeddings. 3. **Scan and index** use the web UI to scan the knowledge base directory, which will automatically import files and build vector embeddings.
@@ -539,6 +562,17 @@ knowledge:
retrieval: retrieval:
top_k: 5 # Number of top results to return top_k: 5 # Number of top results to return
similarity_threshold: 0.7 # Minimum cosine similarity (0-1) similarity_threshold: 0.7 # Minimum cosine similarity (0-1)
multi_query:
max_queries: 4 # MultiQuery rewrite variants (always on)
rerank: # HTTP rerank (always on); empty fields inherit openai/embedding credentials
provider: ""
model: ""
base_url: ""
api_key: ""
post_retrieve:
prefetch_top_k: 20 # per MultiQuery variant; 0 = max(top_k×4, 20)
max_context_chars: 0
max_context_tokens: 0
roles_dir: "roles" # Role configuration directory (relative to config file) roles_dir: "roles" # Role configuration directory (relative to config file)
skills_dir: "skills" # Skills directory (relative to config file) skills_dir: "skills" # Skills directory (relative to config file)
agents_dir: "agents" # Multi-agent Markdown definitions (orchestrator + sub-agents) agents_dir: "agents" # Multi-agent Markdown definitions (orchestrator + sub-agents)
+39 -5
View File
@@ -34,7 +34,18 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
### 系统仪表盘概览 ### 系统仪表盘概览
<img src="./images/dashboard.png" alt="系统仪表盘" width="100%"> <table>
<tr>
<td width="50%" align="center">
<strong>浅色模式</strong><br/>
<img src="./images/dashboard.png" alt="系统仪表盘(浅色)" width="100%">
</td>
<td width="50%" align="center">
<strong>深色模式</strong><br/>
<img src="./images/dark.png" alt="系统仪表盘(深色)" width="100%">
</td>
</tr>
</table>
*仪表盘提供系统运行状态、安全漏洞、工具使用情况和知识库的全面概览,帮助用户快速了解平台核心功能和当前状态。* *仪表盘提供系统运行状态、安全漏洞、工具使用情况和知识库的全面概览,帮助用户快速了解平台核心功能和当前状态。*
@@ -109,7 +120,7 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
- 📄 大结果分页、压缩与全文检索 - 📄 大结果分页、压缩与全文检索
- 🔗 攻击链可视化、风险打分与步骤回放 - 🔗 攻击链可视化、风险打分与步骤回放
- 🔒 Web 登录保护、审计日志、SQLite 持久化 - 🔒 Web 登录保护、审计日志、SQLite 持久化
- 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项) - 📚 知识库(RAG):**Eino MultiQuery** 查询改写 + 多路向量检索 + **HTTP 精排**DashScope `gte-rerank` / Cohere 兼容)+ 后处理(去重、预算);索引侧为 **Eino Compose** 流水线
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作 - 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
- 📂 **项目管理**:共享事实(黑板)跨会话沉淀认知,`upsert_project_fact` + `links` 串联攻击路径;聊天攻击链与项目事实图可视化 - 📂 **项目管理**:共享事实(黑板)跨会话沉淀认知,`upsert_project_fact` + `links` 串联攻击路径;聊天攻击链与项目事实图可视化
- 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板 - 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板
@@ -453,9 +464,10 @@ CyberStrikeAI 支持通过三种传输模式连接外部 MCP 服务器:
### 知识库功能 ### 知识库功能
- **向量检索**:AI 智能体在对话过程中可自动调用 `search_knowledge_base` 工具搜索知识库中的安全知识。 - **向量检索**:AI 智能体在对话过程中可自动调用 `search_knowledge_base` 工具搜索知识库中的安全知识。
- **向量检索**:基于嵌入余弦相似度与相似度阈值过滤(与 Eino `retriever.Retriever` 语义一致) - **RAG 管线(始终启用)****MultiQuery**LLM 查询改写)→ 向量预取与融合 → **HTTP 精排**DashScope `gte-rerank` 或 Cohere 兼容 `/v1/rerank`)→ 后处理(规范化去重、字符/token 预算、最终 top_k)。精排失败时自动降级为融合排序,检索仍可用
- **自动索引**:扫描 `knowledge_base/` 目录下的 Markdown 文件,自动构建向量嵌入索引 - **向量相似度**:基于嵌入余弦相似度与相似度阈值过滤(与 Eino `retriever.Retriever` 语义一致)
- **Web 管理**:通过 Web 界面创建、更新、删除知识项,支持分类管理 - **自动索引**:扫描 `knowledge_base/` 目录下的 Markdown 文件,自动构建向量嵌入索引(Eino Markdown 标题切分 + 递归分块)
- **Web 管理**:通过 Web 界面创建、更新、删除知识项,支持分类管理;设置页可配置 MultiQuery / 精排 / 预取候选数。
- **检索日志**:记录所有知识检索操作,便于审计与调试。 - **检索日志**:记录所有知识检索操作,便于审计与调试。
**快速开始(使用预构建知识库):** **快速开始(使用预构建知识库):**
@@ -477,6 +489,17 @@ CyberStrikeAI 支持通过三种传输模式连接外部 MCP 服务器:
retrieval: retrieval:
top_k: 5 top_k: 5
similarity_threshold: 0.7 similarity_threshold: 0.7
multi_query:
max_queries: 4 # LLM 改写变体上限(始终启用)
rerank: # 精排始终启用;留空则继承 openai/embedding 凭据
provider: "" # 空=按 base_url 推断 dashscope | cohere
model: "" # 空=DashScope→gte-rerankCohere→rerank-multilingual-v3.0
base_url: ""
api_key: ""
post_retrieve:
prefetch_top_k: 20 # 每条 MultiQuery 变体的向量候选数;0=max(top_k×4, 20)
max_context_chars: 0
max_context_tokens: 0
``` ```
2. **添加知识文件**:将 Markdown 文件放入 `knowledge_base/` 目录,按分类组织(如 `knowledge_base/SQL注入/README.md`)。 2. **添加知识文件**:将 Markdown 文件放入 `knowledge_base/` 目录,按分类组织(如 `knowledge_base/SQL注入/README.md`)。
3. **扫描索引**:在 Web 界面中点击"扫描知识库",系统会自动导入文件并构建向量索引。 3. **扫描索引**:在 Web 界面中点击"扫描知识库",系统会自动导入文件并构建向量索引。
@@ -537,6 +560,17 @@ knowledge:
retrieval: retrieval:
top_k: 5 # 检索返回的 Top-K 结果数量 top_k: 5 # 检索返回的 Top-K 结果数量
similarity_threshold: 0.7 # 余弦相似度阈值(0-1),低于此值的结果将被过滤 similarity_threshold: 0.7 # 余弦相似度阈值(0-1),低于此值的结果将被过滤
multi_query:
max_queries: 4 # MultiQuery 改写变体上限(始终启用)
rerank: # HTTP 精排(始终启用);留空则继承 openai/embedding 凭据
provider: ""
model: ""
base_url: ""
api_key: ""
post_retrieve:
prefetch_top_k: 20 # 每条 MultiQuery 变体;0=max(top_k×4, 20)
max_context_chars: 0
max_context_tokens: 0
roles_dir: "roles" # 角色配置文件目录(相对于配置文件所在目录) roles_dir: "roles" # 角色配置文件目录(相对于配置文件所在目录)
skills_dir: "skills" # Skills 目录(相对于配置文件所在目录) skills_dir: "skills" # Skills 目录(相对于配置文件所在目录)
agents_dir: "agents" # 多代理 Markdown(主代理 orchestrator.md + 子代理 *.md agents_dir: "agents" # 多代理 Markdown(主代理 orchestrator.md + 子代理 *.md
+76 -6
View File
@@ -10,7 +10,7 @@
# ============================================ # ============================================
# 前端显示的版本号(可选,不填则显示默认版本) # 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.6.46" version: "v1.6.49"
# 服务器配置 # 服务器配置
server: server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口 host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -102,9 +102,69 @@ agent:
system_prompt_path: "" system_prompt_path: ""
# 人机协同(HITL)全局白名单:此处列出的工具始终免审批,与对话页「白名单工具(免审批,逗号分隔)」合并为并集;侧栏「应用」可合并写入本列表并立即生效。 # 人机协同(HITL)全局白名单:此处列出的工具始终免审批,与对话页「白名单工具(免审批,逗号分隔)」合并为并集;侧栏「应用」可合并写入本列表并立即生效。
# 非白名单工具在审批方=审计 Agent 时,按会话 HITL 模式选用提示词:
# approval → audit_agent_prompt
# review_edit → audit_agent_prompt_review_edit(可改参后放行)
hitl: hitl:
# 已决策审计日志保留天数(与 MCP 监控一致;省略默认 90;0 表示不自动清理)
retention_days: 90
# 按你环境里的真实工具名增删(与侧栏一致、小写不敏感);不需要全局免审批可改为 [] # 按你环境里的真实工具名增删(与侧栏一致、小写不敏感);不需要全局免审批可改为 []
tool_whitelist: [read_file, list_dir, glob, grep] tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
# audit_agent_prompt: | # 审批模式;留空使用内置默认,可在「人机协同」页编辑
# audit_agent_prompt_review_edit: | # 审查编辑模式;留空使用内置默认
audit_agent_prompt: |-
你是 CyberStrikeAI 人机协同审计 Agent。审查 Agent 即将执行的工具调用是否会对系统造成实质性损害。
你会收到 JSON,包含 hitlMode、toolName、arguments/argumentsObj、userMessage、thinking、reasoningChain、planning 等字段。
裁决基调(默认放行):
- 常规、低风险的渗透测试操作 → approve(如信息收集、端口/服务扫描、目录枚举、只读查询、无害探测命令)
- 与用户授权、当前任务目标一致,且未见明确高危迹象 → approve
- 仅在「可能对系统造成实质影响」时 → reject
必须 reject 的高危情形(示例,非穷举):
- 删库、清表、批量删除数据、格式化磁盘、不可逆破坏
- 修改/重置密码、创建或篡改管理员账号、持久化后门、开机自启
- 向生产环境写入恶意载荷、勒索加密、停止关键服务、修改系统核心配置
- 明显越权:与任务/授权目标无关的破坏性操作
不应单独作为 reject 理由的情形:
- 常规 nmap/curl/grep/读文件/枚举类命令本身
- 参数略显宽泛但无明确破坏意图
- 仅因「信息不足」——若无上述高危迹象,应 approve 并可在 comment 中提示注意点
仅输出一行 JSON,不要 markdown 代码块:
{"decision":"approve"|"reject","comment":"简要理由"}
audit_agent_prompt_review_edit: |-
你是 CyberStrikeAI 人机协同审计 Agent。审查 Agent 即将执行的工具调用是否会对系统造成实质性损害。
你会收到 JSON,包含 hitlMode、toolName、arguments/argumentsObj、userMessage、thinking、reasoningChain、planning 等字段。
裁决基调(默认放行):
- 常规、低风险的渗透测试操作 → approve(如信息收集、端口/服务扫描、目录枚举、只读查询、无害探测命令)
- 与用户授权、当前任务目标一致,且未见明确高危迹象 → approve
- 仅在「可能对系统造成实质影响」时 → reject;参数可安全收窄时优先 approve + editedArguments
必须 reject 的高危情形(示例,非穷举):
- 删库、清表、批量删除数据、格式化磁盘、不可逆破坏
- 修改/重置密码、创建或篡改管理员账号、持久化后门、开机自启
- 向生产环境写入恶意载荷、勒索加密、停止关键服务、修改系统核心配置
- 明显越权:与任务/授权目标无关的破坏性操作
不应单独作为 reject 理由的情形:
- 常规 nmap/curl/grep/读文件/枚举类命令本身
- 参数略显宽泛但无明确破坏意图(应收窄参数后 approve)
- 仅因「信息不足」——若无上述高危迹象,应 approve 并可在 comment 中提示注意点
仅输出一行 JSON,不要 markdown 代码块:
{"decision":"approve"|"reject","comment":"简要理由","editedArguments":{...}}
editedArguments 规则(仅 approve 且需要改参时填写,否则省略该字段):
- 提供完整替换后的工具参数对象,键名与 argumentsObj 一致
- 只做最小必要修改以收窄范围、消除风险(如限制 path、去掉危险 flag)
- 禁止扩大攻击面:不得扩大目标范围、提升权限或引入破坏性参数
- 无法安全改参且存在上述高危情形时应 reject,不要勉强 approve
# 多代理与 Eino 单代理(CloudWeGo Eino ADK;单代理入口 /api/eino-agent*,多代理 /api/multi-agent* # 多代理与 Eino 单代理(CloudWeGo Eino ADK;单代理入口 /api/eino-agent*,多代理 /api/multi-agent*
# 依赖在 go.mod 中拉取;若下载失败可设置: go env -w GOPROXY=https://goproxy.cn,direct # 依赖在 go.mod 中拉取;若下载失败可设置: go env -w GOPROXY=https://goproxy.cn,direct
# Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体 orchestration 中指定;机器人按 robot_default_agent_mode # Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体 orchestration 中指定;机器人按 robot_default_agent_mode
@@ -114,7 +174,7 @@ multi_agent:
batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高) batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高)
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。主/子代理 ReAct 轮次见 agent.max_iterations。 # plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。主/子代理 ReAct 轮次见 agent.max_iterations。
plan_execute_loop_max_iterations: 0 plan_execute_loop_max_iterations: 0
sub_agent_user_context_max_runes: 0 # 子代理 task 描述中自动注入用户原始请求的字符上限;0=默认2000,负数=禁用 sub_agent_user_context_max_runes: 0 # 子代理 task 描述中注入用户原文;0=不截断(默认),>0=总字符上限,负数=禁用
without_general_sub_agent: false # false 时保留 Deep 内置 general-purpose 子代理 without_general_sub_agent: false # false 时保留 Deep 内置 general-purpose 子代理
without_write_todos: false without_write_todos: false
orchestrator_instruction: "" # Deep 主代理:agents/orchestrator.md(或 kind: orchestrator 的单个 .md)正文优先;正文为空时用此处;皆空则 Eino 默认 orchestrator_instruction: "" # Deep 主代理:agents/orchestrator.md(或 kind: orchestrator 的单个 .md)正文优先;正文为空时用此处;皆空则 Eino 默认
@@ -125,7 +185,8 @@ multi_agent:
disable: false # true:不注册 skill 渐进式披露中间件,也不挂本机 FS/Shell 工具;false:按下方开关加载 disable: false # true:不注册 skill 渐进式披露中间件,也不挂本机 FS/Shell 工具;false:按下方开关加载
filesystem_tools: true # true:注册 read_file/glob/grep/write/edit/execute(授权环境慎用);false:仅 skill,不暴露本机读写与 Shell filesystem_tools: true # true:注册 read_file/glob/grep/write/edit/execute(授权环境慎用);false:仅 skill,不暴露本机读写与 Shell
skill_tool_name: skill # 模型侧可调用的「加载技能」工具名,一般保持 skill;与技能包文档中的调用名一致即可 skill_tool_name: skill # 模型侧可调用的「加载技能」工具名,一般保持 skill;与技能包文档中的调用名一致即可
# Eino ADK 中间件与 Deep/Supervisor 调参(结构体见 internal/config/config.go → MultiAgentEinoMiddlewareConfig # Eino ADK 中间件与 Deep/Supervisor/plan_execute Executor 调参(结构体见 internal/config/config.go → MultiAgentEinoMiddlewareConfig
# plan_execute:下列 patch/reduction/tool_search/plantask 等同样作用于 Executor(经 ExecPreMiddlewares);Planner/Replanner 不挂 MCP 前置中间件。
eino_middleware: eino_middleware:
patch_tool_calls: true # true:修补历史中无 tool_result 的悬空 tool_call(流式中断/重试后更稳);false:关闭;字段省略时默认等同 true patch_tool_calls: true # true:修补历史中无 tool_result 的悬空 tool_call(流式中断/重试后更稳);false:关闭;字段省略时默认等同 true
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文 tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
@@ -149,6 +210,7 @@ multi_agent:
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 loop + summarization 共用 isEinoTransientRunError);0=默认 10 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
empty_response_continue_max_attempts: 0 # Run 成功但未捕获助手正文(含流式中断)时 Handler 退避续跑次数;0=默认 5
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: 0 # 已废弃,请用 run_retry_max_attempts;保留字段仅为兼容旧配置 deep_model_retry_max_retries: 0 # 已废弃,请用 run_retry_max_attempts;保留字段仅为兼容旧配置
task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑 task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑
@@ -221,9 +283,17 @@ knowledge:
retrieval: retrieval:
top_k: 5 # 检索返回的Top-K结果数量 top_k: 5 # 检索返回的Top-K结果数量
similarity_threshold: 0.4 # 余弦相似度阈值(0-1),低于此值的结果将被过滤 similarity_threshold: 0.4 # 余弦相似度阈值(0-1),低于此值的结果将被过滤
# 检索后处理:固定正文规范化去重;上下文预算;可选代码注入 DocumentReranker 做重排 # Eino MultiQueryLLM 改写查询后多路向量检索再融合(始终启用)
multi_query:
max_queries: 4 # 改写变体上限(含语义覆盖);建议 3~4
# 精排(始终启用):dashscope 用 gte-rerank;其他 OpenAI 兼容端点走 /v1/rerank
rerank:
provider: "" # 空=按 base_url 推断:dashscope | cohere
model: "" # 空=dashscope→gte-rerankcohere→rerank-multilingual-v3.0
base_url: "" # 留空则用 embedding / openai 的 base_url
api_key: "" # 留空则用 embedding / openai 的 api_key
post_retrieve: post_retrieve:
prefetch_top_k: 0 # 0 与 top_k 相同;可设为 15~30 以便去重后仍填满 top_k prefetch_top_k: 20 # 每条 MultiQuery 变体的向量候选数;0=内置 max(top_k*4,20)
max_context_chars: 0 # 0 不限制;否则返回的正文总 Unicode 字符上限(整段 chunk max_context_chars: 0 # 0 不限制;否则返回的正文总 Unicode 字符上限(整段 chunk
max_context_tokens: 0 # 0 不限制;tiktoken 总 token 上限 max_context_tokens: 0 # 0 不限制;tiktoken 总 token 上限
sub_index_filter: "" sub_index_filter: ""
+4 -2
View File
@@ -26,7 +26,7 @@
| OpenAPI | 多代理路径说明已更新(流式未启用为 SSE 错误事件)。 | | OpenAPI | 多代理路径说明已更新(流式未启用为 SSE 错误事件)。 |
| 机器人 | `ProcessMessageForRobot``robot_default_agent_mode`(默认 `eino_single`)调用 `RunEinoSingleChatModelAgent``RunDeepAgent`。 | | 机器人 | `ProcessMessageForRobot``robot_default_agent_mode`(默认 `eino_single`)调用 `RunEinoSingleChatModelAgent``RunDeepAgent`。 |
| 预置编排 | 聊天 / WebShell`POST /api/multi-agent*` 请求体 `orchestration``deep` \| `plan_execute` \| `supervisor`(缺省 `deep`)。`plan_execute` 不构建 YAML/Markdown 子代理;`plan_execute_loop_max_iterations` 仍来自配置。`supervisor` 至少需一个子代理。 | | 预置编排 | 聊天 / WebShell`POST /api/multi-agent*` 请求体 `orchestration``deep` \| `plan_execute` \| `supervisor`(缺省 `deep`)。`plan_execute` 不构建 YAML/Markdown 子代理;`plan_execute_loop_max_iterations` 仍来自配置。`supervisor` 至少需一个子代理。 |
| Eino 中间件 | `multi_agent.eino_middleware`(可选):`patchtoolcalls`(默认开)、`toolsearch`(按阈值拆分 MCP 工具列表)、`plantask`(需 `eino_skills`)、`reduction`(大工具输出截断/落盘)、`checkpoint_dir`Runner 断点)、`deep_output_key` / `deep_model_retry_max_retries` / `task_tool_description_prefix`Deep 与 supervisor 主代理共享其中模型重试与 OutputKey)。`plan_execute` 的 Executor 无 Handlers:仅继承 **ToolsConfig** 侧效果(如 `tool_search` 列表拆分),不挂载 patch/plantask/reduction 中间件。 | | Eino 中间件 | `multi_agent.eino_middleware`(可选):`patchtoolcalls`(默认开)、`toolsearch`(按阈值拆分 MCP 工具列表)、`plantask`(需 `eino_skills`)、`reduction`(大工具输出截断/落盘)、`checkpoint_dir`Runner 断点)、`deep_output_key` / `deep_model_retry_max_retries` / `task_tool_description_prefix`Deep 与 supervisor 主代理共享其中模型重试与 OutputKey)。**`plan_execute`**`runner.go``prependEinoMiddlewares(einoMWMain)` 产物作为 `ExecPreMiddlewares` 挂到 **Executor**(与 Deep/Supervisor 主代理同序:patch → reduction → toolsearch → plantask → filesystem → skill → summarization tail);Planner/Replanner 仅 summarization tail + prompt 预算截断,不跑 MCP 工具链。 |
## 进行中 / 待办( backlog ## 进行中 / 待办( backlog
@@ -37,7 +37,8 @@
## 关键文件索引 ## 关键文件索引
- `internal/multiagent/runner.go` — DeepAgent 组装与事件循环 - `internal/multiagent/runner.go` — DeepAgent / plan_execute / supervisor 组装与事件循环
- `internal/multiagent/eino_orchestration.go` — PlanExecute 根节点与 Executor 中间件栈(`buildPlanExecuteExecutorHandlers`
- `internal/handler/multi_agent.go` — SSE 与(同步)HTTP - `internal/handler/multi_agent.go` — SSE 与(同步)HTTP
- `internal/handler/multi_agent_prepare.go` — 会话准备(含 WebShell - `internal/handler/multi_agent_prepare.go` — 会话准备(含 WebShell
- `internal/einomcp/` — MCP → Eino Tool - `internal/einomcp/` — MCP → Eino Tool
@@ -59,4 +60,5 @@
| 2026-03-22 | `orchestrator.md` / `kind: orchestrator` 主代理、列表主/子标记、与 `orchestrator_instruction` 优先级。 | | 2026-03-22 | `orchestrator.md` / `kind: orchestrator` 主代理、列表主/子标记、与 `orchestrator_instruction` 优先级。 |
| 2026-04-19 | 主聊天「对话模式」:原生 ReAct 与 Deep / Plan-Execute / Supervisor`POST /api/multi-agent*` 请求体 `orchestration` 与界面一致;`config.yaml` / 设置页不再维护预置编排字段(机器人/批量默认 `deep`)。 | | 2026-04-19 | 主聊天「对话模式」:原生 ReAct 与 Deep / Plan-Execute / Supervisor`POST /api/multi-agent*` 请求体 `orchestration` 与界面一致;`config.yaml` / 设置页不再维护预置编排字段(机器人/批量默认 `deep`)。 |
| 2026-04-21 | 移除角色 `skills``/api/roles/skills/list``bind_role` 仅继承 toolsSkills 仅通过 Eino `skill` 工具按需加载。 | | 2026-04-21 | 移除角色 `skills``/api/roles/skills/list``bind_role` 仅继承 toolsSkills 仅通过 Eino `skill` 工具按需加载。 |
| 2026-07-02 | **plan_execute Executor 中间件对齐**`ExecPreMiddlewares` 与 Deep 主代理同源;`buildPlanExecuteExecutorHandlers` + 回归测试;文档更正。 |
| 2026-06-02 | **移除原生 ReAct**:删除 `/api/agent-loop*` 执行入口与 `AgentLoopWithProgress`;统一 Eino ADK(单代理 `/api/eino-agent*`,多代理 `/api/multi-agent*`);任务 cancel/tasks API 保留。 | | 2026-06-02 | **移除原生 ReAct**:删除 `/api/agent-loop*` 执行入口与 `AgentLoopWithProgress`;统一 Eino ADK(单代理 `/api/eino-agent*`,多代理 `/api/multi-agent*`);任务 cancel/tasks API 保留。 |
BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.0 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 88 KiB

After

Width:  |  Height:  |  Size: 265 KiB

+32 -14
View File
@@ -21,11 +21,13 @@ import (
"cyberstrike-ai/internal/database" "cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/einoobserve" "cyberstrike-ai/internal/einoobserve"
"cyberstrike-ai/internal/handler" "cyberstrike-ai/internal/handler"
"cyberstrike-ai/internal/hitl"
"cyberstrike-ai/internal/knowledge" "cyberstrike-ai/internal/knowledge"
"cyberstrike-ai/internal/logger" "cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/monitor" "cyberstrike-ai/internal/monitor"
"cyberstrike-ai/internal/multiagent"
"cyberstrike-ai/internal/robot" "cyberstrike-ai/internal/robot"
"cyberstrike-ai/internal/security" "cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/skillpackage" "cyberstrike-ai/internal/skillpackage"
@@ -67,6 +69,10 @@ type App struct {
// New 创建新应用 // New 创建新应用
func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) { func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) {
if err := multiagent.InitADK(); err != nil {
return nil, fmt.Errorf("初始化 Eino ADK: %w", err)
}
gin.SetMode(gin.ReleaseMode) gin.SetMode(gin.ReleaseMode)
router := gin.Default() router := gin.Default()
@@ -104,6 +110,10 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
monitorRetention.PurgeExpired() monitorRetention.PurgeExpired()
monitor.StartRetentionLoop(monitorRetention, log.Logger) monitor.StartRetentionLoop(monitorRetention, log.Logger)
hitlRetention := hitl.NewService(db, cfg, log.Logger)
hitlRetention.PurgeExpired()
hitl.StartRetentionLoop(hitlRetention, log.Logger)
// 创建MCP服务器(带数据库持久化) // 创建MCP服务器(带数据库持久化)
mcpServer := mcp.NewServerWithStorage(log.Logger, db) mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes) mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
@@ -135,6 +145,10 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
externalMCPMgr.StartAllEnabled() externalMCPMgr.StartAllEnabled()
} }
execReconciler := monitor.NewExecutionReconciler(db, mcpServer, externalMCPMgr, log.Logger)
execReconciler.ReconcileOnStartup()
monitor.StartStaleRunningReconcileLoop(execReconciler, log.Logger)
// 创建Agent // 创建Agent
maxIterations := cfg.Agent.MaxIterations maxIterations := cfg.Agent.MaxIterations
if maxIterations <= 0 { if maxIterations <= 0 {
@@ -193,14 +207,12 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
return nil, fmt.Errorf("初始化知识库嵌入器失败: %w", err) return nil, fmt.Errorf("初始化知识库嵌入器失败: %w", err)
} }
// 创建检索器 // 创建检索器Eino MultiQuery + 重排流水线)
retrievalConfig := &knowledge.RetrievalConfig{ retrievalConfig := knowledge.RetrievalConfigFromYAML(cfg.Knowledge.Retrieval)
TopK: cfg.Knowledge.Retrieval.TopK,
SimilarityThreshold: cfg.Knowledge.Retrieval.SimilarityThreshold,
SubIndexFilter: cfg.Knowledge.Retrieval.SubIndexFilter,
PostRetrieve: cfg.Knowledge.Retrieval.PostRetrieve,
}
knowledgeRetriever = knowledge.NewRetriever(knowledgeDB, embedder, retrievalConfig, log.Logger) knowledgeRetriever = knowledge.NewRetriever(knowledgeDB, embedder, retrievalConfig, log.Logger)
if err := knowledge.WireRetrieverPipeline(context.Background(), knowledgeRetriever, &cfg.OpenAI); err != nil {
return nil, fmt.Errorf("初始化知识库检索流水线失败: %w", err)
}
// 创建索引器(Eino Compose 链) // 创建索引器(Eino Compose 链)
knowledgeIndexer, err = knowledge.NewIndexer(context.Background(), knowledgeDB, embedder, log.Logger, &cfg.Knowledge) knowledgeIndexer, err = knowledge.NewIndexer(context.Background(), knowledgeDB, embedder, log.Logger, &cfg.Knowledge)
@@ -354,6 +366,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger) configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
configHandler.SetAudit(auditSvc) configHandler.SetAudit(auditSvc)
agentHandler.SetHitlToolWhitelistSaver(configHandler) agentHandler.SetHitlToolWhitelistSaver(configHandler)
agentHandler.SetHitlAuditStrategySaver(configHandler)
externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger) externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger)
externalMCPHandler.SetAudit(auditSvc) externalMCPHandler.SetAudit(auditSvc)
roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger) roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger)
@@ -803,11 +816,18 @@ func setupRoutes(
protected.POST("/eino-agent", agentHandler.EinoSingleAgentLoop) protected.POST("/eino-agent", agentHandler.EinoSingleAgentLoop)
protected.POST("/eino-agent/stream", agentHandler.EinoSingleAgentLoopStream) protected.POST("/eino-agent/stream", agentHandler.EinoSingleAgentLoopStream)
protected.GET("/hitl/pending", agentHandler.ListHITLPending) protected.GET("/hitl/pending", agentHandler.ListHITLPending)
protected.GET("/hitl/logs", agentHandler.ListHITLLogs)
protected.DELETE("/hitl/logs", agentHandler.DeleteHITLLogs)
protected.GET("/hitl/logs/:id", agentHandler.GetHITLLog)
protected.POST("/hitl/decision", agentHandler.DecideHITLInterrupt) protected.POST("/hitl/decision", agentHandler.DecideHITLInterrupt)
protected.POST("/hitl/dismiss", agentHandler.DismissHITLInterrupt) protected.POST("/hitl/dismiss", agentHandler.DismissHITLInterrupt)
protected.GET("/hitl/config/:conversationId", agentHandler.GetHITLConversationConfig) protected.GET("/hitl/config/:conversationId", agentHandler.GetHITLConversationConfig)
protected.PUT("/hitl/config", agentHandler.UpsertHITLConversationConfig) protected.PUT("/hitl/config", agentHandler.UpsertHITLConversationConfig)
protected.GET("/hitl/tool-whitelist", agentHandler.GetHITLGlobalToolWhitelist)
protected.PUT("/hitl/tool-whitelist", agentHandler.SetHITLGlobalToolWhitelist)
protected.POST("/hitl/tool-whitelist", agentHandler.MergeHITLGlobalToolWhitelist) protected.POST("/hitl/tool-whitelist", agentHandler.MergeHITLGlobalToolWhitelist)
protected.GET("/hitl/audit-strategy", agentHandler.GetHITLAuditStrategy)
protected.PUT("/hitl/audit-strategy", agentHandler.UpdateHITLAuditStrategy)
// Agent Loop 取消与任务列表 // Agent Loop 取消与任务列表
protected.POST("/agent-loop/cancel", agentHandler.CancelAgentLoop) protected.POST("/agent-loop/cancel", agentHandler.CancelAgentLoop)
protected.GET("/agent-loop/tasks", agentHandler.ListAgentTasks) protected.GET("/agent-loop/tasks", agentHandler.ListAgentTasks)
@@ -1778,14 +1798,12 @@ func initializeKnowledge(
return nil, fmt.Errorf("初始化知识库嵌入器失败: %w", err) return nil, fmt.Errorf("初始化知识库嵌入器失败: %w", err)
} }
// 创建检索器 // 创建检索器Eino MultiQuery + 重排流水线)
retrievalConfig := &knowledge.RetrievalConfig{ retrievalConfig := knowledge.RetrievalConfigFromYAML(cfg.Knowledge.Retrieval)
TopK: cfg.Knowledge.Retrieval.TopK,
SimilarityThreshold: cfg.Knowledge.Retrieval.SimilarityThreshold,
SubIndexFilter: cfg.Knowledge.Retrieval.SubIndexFilter,
PostRetrieve: cfg.Knowledge.Retrieval.PostRetrieve,
}
knowledgeRetriever := knowledge.NewRetriever(knowledgeDB, embedder, retrievalConfig, logger) knowledgeRetriever := knowledge.NewRetriever(knowledgeDB, embedder, retrievalConfig, logger)
if err := knowledge.WireRetrieverPipeline(context.Background(), knowledgeRetriever, &cfg.OpenAI); err != nil {
return nil, fmt.Errorf("初始化知识库检索流水线失败: %w", err)
}
// 创建索引器(Eino Compose 链) // 创建索引器(Eino Compose 链)
knowledgeIndexer, err := knowledge.NewIndexer(context.Background(), knowledgeDB, embedder, logger, &cfg.Knowledge) knowledgeIndexer, err := knowledge.NewIndexer(context.Background(), knowledgeDB, embedder, logger, &cfg.Knowledge)
+246 -31
View File
@@ -96,8 +96,8 @@ type MultiAgentConfig struct {
// OrchestratorInstructionSupervisor supervisor 主代理系统提示(transfer/exit 说明仍由运行追加);非空且 agents/orchestrator-supervisor.md 正文为空或未存在时生效。 // OrchestratorInstructionSupervisor supervisor 主代理系统提示(transfer/exit 说明仍由运行追加);非空且 agents/orchestrator-supervisor.md 正文为空或未存在时生效。
OrchestratorInstructionSupervisor string `yaml:"orchestrator_instruction_supervisor,omitempty" json:"orchestrator_instruction_supervisor,omitempty"` OrchestratorInstructionSupervisor string `yaml:"orchestrator_instruction_supervisor,omitempty" json:"orchestrator_instruction_supervisor,omitempty"`
SubAgents []MultiAgentSubConfig `yaml:"sub_agents" json:"sub_agents"` SubAgents []MultiAgentSubConfig `yaml:"sub_agents" json:"sub_agents"`
// SubAgentUserContextMaxRunes caps the user-context supplement appended to task descriptions for sub-agents. // SubAgentUserContextMaxRunes caps user-context supplement for sub-agent task descriptions.
// 0 (default) uses the built-in default of 2000 runes; negative value disables injection entirely. // 0 (default) preserves all user turns verbatim; >0 caps total runes; negative disables injection.
SubAgentUserContextMaxRunes int `yaml:"sub_agent_user_context_max_runes,omitempty" json:"sub_agent_user_context_max_runes,omitempty"` SubAgentUserContextMaxRunes int `yaml:"sub_agent_user_context_max_runes,omitempty" json:"sub_agent_user_context_max_runes,omitempty"`
// EinoSkills configures CloudWeGo Eino ADK skill middleware + optional local filesystem/execute on DeepAgent. // EinoSkills configures CloudWeGo Eino ADK skill middleware + optional local filesystem/execute on DeepAgent.
EinoSkills MultiAgentEinoSkillsConfig `yaml:"eino_skills,omitempty" json:"eino_skills,omitempty"` EinoSkills MultiAgentEinoSkillsConfig `yaml:"eino_skills,omitempty" json:"eino_skills,omitempty"`
@@ -107,6 +107,11 @@ type MultiAgentConfig struct {
EinoCallbacks MultiAgentEinoCallbacksConfig `yaml:"eino_callbacks,omitempty" json:"eino_callbacks,omitempty"` EinoCallbacks MultiAgentEinoCallbacksConfig `yaml:"eino_callbacks,omitempty" json:"eino_callbacks,omitempty"`
} }
// SubAgentUserContextMaxRunesEffective returns max runes for sub-agent task supplement; 0 = unlimited; negative = disabled.
func (c MultiAgentConfig) SubAgentUserContextMaxRunesEffective() int {
return c.SubAgentUserContextMaxRunes
}
// MultiAgentEinoCallbacksConfig enables Eino unified callbacks on each ADK agent run (deep / plan_execute / supervisor / eino_single). // MultiAgentEinoCallbacksConfig enables Eino unified callbacks on each ADK agent run (deep / plan_execute / supervisor / eino_single).
// Modes: log_only (zap + optional OTel; no SSE to browser), sse (adds client SSE eino_trace_* when sse_trace_to_client), full (sse rules + stream callback copies closed). // Modes: log_only (zap + optional OTel; no SSE to browser), sse (adds client SSE eino_trace_* when sse_trace_to_client), full (sse rules + stream callback copies closed).
type MultiAgentEinoCallbacksConfig struct { type MultiAgentEinoCallbacksConfig struct {
@@ -270,6 +275,8 @@ type MultiAgentEinoMiddlewareConfig struct {
RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"` RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"`
// RunRetryMaxBackoffSec 单次退避上限秒数;0=默认 30。 // RunRetryMaxBackoffSec 单次退避上限秒数;0=默认 30。
RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"` RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"`
// EmptyResponseContinueMaxAttempts Run 成功但未捕获助手正文时 Handler 层退避续跑次数;0=默认 5。
EmptyResponseContinueMaxAttempts int `yaml:"empty_response_continue_max_attempts,omitempty" json:"empty_response_continue_max_attempts,omitempty"`
// TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended). // TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended).
TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"` TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"`
} }
@@ -490,6 +497,17 @@ type RobotWecomConfig struct {
AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId
} }
// ValidateWecomConfig 校验企业微信机器人配置;启用时必须配置 token,否则回调无法防伪造。
func ValidateWecomConfig(w RobotWecomConfig) error {
if !w.Enabled {
return nil
}
if strings.TrimSpace(w.Token) == "" {
return fmt.Errorf("robots.wecom.enabled 为 true 时必须配置 robots.wecom.token")
}
return nil
}
// RobotDingtalkConfig 钉钉机器人配置 // RobotDingtalkConfig 钉钉机器人配置
type RobotDingtalkConfig struct { type RobotDingtalkConfig struct {
Enabled bool `yaml:"enabled" json:"enabled"` Enabled bool `yaml:"enabled" json:"enabled"`
@@ -614,10 +632,100 @@ type AgentConfig struct {
} }
// HitlConfig 人机协同全局选项;与会话侧栏/API 中的白名单合并为并集后参与判定。 // HitlConfig 人机协同全局选项;与会话侧栏/API 中的白名单合并为并集后参与判定。
// tool_whitelist 可在侧栏「应用」时合并写入 config.yaml 并立即生效;其他字段若仅改文件仍需重启 // tool_whitelist 可在侧栏「应用」时合并写入 config.yaml 并立即生效。
// audit_agent_prompt / audit_agent_prompt_review_edit 可在人机协同页编辑并立即生效;空则使用内置默认。
type HitlConfig struct { type HitlConfig struct {
// ToolWhitelist 全局免审批工具名(与每条会话配置的 sensitiveTools 语义相同:白名单内工具不触发 HITL)。 // ToolWhitelist 全局免审批工具名(与白名单内工具不触发 HITL 审批)。
ToolWhitelist []string `yaml:"tool_whitelist,omitempty" json:"tool_whitelist,omitempty"` ToolWhitelist []string `yaml:"tool_whitelist,omitempty" json:"tool_whitelist,omitempty"`
// AuditAgentPrompt 审批模式(approval)下审计 Agent 系统提示词。
AuditAgentPrompt string `yaml:"audit_agent_prompt,omitempty" json:"audit_agent_prompt,omitempty"`
// AuditAgentPromptReviewEdit 审查编辑模式(review_edit)下审计 Agent 系统提示词。
AuditAgentPromptReviewEdit string `yaml:"audit_agent_prompt_review_edit,omitempty" json:"audit_agent_prompt_review_edit,omitempty"`
// RetentionDays 已决策审计日志(hitl_interrupts 非 pending)保留天数;省略时默认 90;0 表示不自动清理。
RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
}
// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90.
func (h HitlConfig) RetentionDaysEffective() int {
if h.RetentionDays == nil {
return 90
}
if *h.RetentionDays < 0 {
return 0
}
return *h.RetentionDays
}
const hitlAuditAgentPromptBase = `你是 CyberStrikeAI 人机协同审计 Agent。审查 Agent 即将执行的工具调用是否会对系统造成实质性损害。
你会收到 JSON,包含 hitlMode、toolName、arguments/argumentsObj、userMessage、thinking、reasoningChain、planning 等字段。
裁决基调(默认放行):
- 常规、低风险的渗透测试操作 → approve(如信息收集、端口/服务扫描、目录枚举、只读查询、无害探测命令)
- 与用户授权、当前任务目标一致,且未见明确高危迹象 → approve
- 仅在「可能对系统造成实质影响」时 → reject
必须 reject 的高危情形(示例,非穷举):
- 删库、清表、批量删除数据、格式化磁盘、不可逆破坏
- 修改/重置密码、创建或篡改管理员账号、持久化后门、开机自启
- 向生产环境写入恶意载荷、勒索加密、停止关键服务、修改系统核心配置
- 明显越权:与任务/授权目标无关的破坏性操作
不应单独作为 reject 理由的情形:
- 常规 nmap/curl/grep/读文件/枚举类命令本身
- 参数略显宽泛但无明确破坏意图(审查编辑模式可收窄参数后 approve)
- 仅因「信息不足」——若无上述高危迹象,应 approve 并可在 comment 中提示注意点`
const hitlAuditAgentPromptApprovalOutput = `
仅输出一行 JSON,不要 markdown 代码块:
{"decision":"approve"|"reject","comment":"简要理由"}`
const hitlAuditAgentPromptReviewEditOutput = `
仅输出一行 JSON,不要 markdown 代码块:
{"decision":"approve"|"reject","comment":"简要理由","editedArguments":{...}}
editedArguments 规则(仅 approve 且需要改参时填写,否则省略该字段):
- 提供完整替换后的工具参数对象,键名与 argumentsObj 一致
- 只做最小必要修改以收窄范围、消除风险(如限制 path、去掉危险 flag)
- 禁止扩大攻击面:不得扩大目标范围、提升权限或引入破坏性参数
- 无法安全改参时应 reject,不要勉强 approve`
// DefaultHitlAuditAgentPrompt 内置审批模式审计 Agent 提示词。
func DefaultHitlAuditAgentPrompt() string {
return hitlAuditAgentPromptBase + hitlAuditAgentPromptApprovalOutput
}
// DefaultHitlAuditAgentPromptReviewEdit 内置审查编辑模式审计 Agent 提示词。
func DefaultHitlAuditAgentPromptReviewEdit() string {
return hitlAuditAgentPromptBase + hitlAuditAgentPromptReviewEditOutput
}
// EffectiveAuditAgentPrompt 返回审批模式生效的审计 Agent 提示词。
func (c HitlConfig) EffectiveAuditAgentPrompt() string {
return c.EffectiveAuditAgentPromptForMode("approval")
}
// EffectiveAuditAgentPromptForMode 按 HITL 模式返回生效的审计 Agent 提示词。
func (c HitlConfig) EffectiveAuditAgentPromptForMode(mode string) string {
if normalizeHitlModeForPrompt(mode) == "review_edit" {
if s := strings.TrimSpace(c.AuditAgentPromptReviewEdit); s != "" {
return s
}
return DefaultHitlAuditAgentPromptReviewEdit()
}
if s := strings.TrimSpace(c.AuditAgentPrompt); s != "" {
return s
}
return DefaultHitlAuditAgentPrompt()
}
func normalizeHitlModeForPrompt(mode string) string {
switch strings.ToLower(strings.TrimSpace(mode)) {
case "review_edit":
return "review_edit"
default:
return "approval"
}
} }
type AuthConfig struct { type AuthConfig struct {
@@ -804,33 +912,13 @@ func Load(path string) (*Config, error) {
// 如果配置了工具目录,从目录加载工具配置 // 如果配置了工具目录,从目录加载工具配置
if cfg.Security.ToolsDir != "" { if cfg.Security.ToolsDir != "" {
configDir := filepath.Dir(path) inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...)
toolsDir := cfg.Security.ToolsDir toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, path)
merged, err := MergeToolsFromDir(toolsDir, inlineTools)
// 如果是相对路径,相对于配置文件所在目录
if !filepath.IsAbs(toolsDir) {
toolsDir = filepath.Join(configDir, toolsDir)
}
tools, err := LoadToolsFromDir(toolsDir)
if err != nil { if err != nil {
return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err) return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err)
} }
cfg.Security.Tools = merged
// 合并工具配置:目录中的工具优先,主配置中的工具作为补充
existingTools := make(map[string]bool)
for _, tool := range tools {
existingTools[tool.Name] = true
}
// 添加主配置中不存在于目录中的工具(向后兼容)
for _, tool := range cfg.Security.Tools {
if !existingTools[tool.Name] {
tools = append(tools, tool)
}
}
cfg.Security.Tools = tools
} }
// 外部 MCP:迁移 + 环境变量展开 // 外部 MCP:迁移 + 环境变量展开
@@ -874,6 +962,10 @@ func Load(path string) (*Config, error) {
} }
} }
if err := ValidateWecomConfig(cfg.Robots.Wecom); err != nil {
return nil, err
}
return &cfg, nil return &cfg, nil
} }
@@ -1098,6 +1190,75 @@ func PrintMCPConfigJSON(mcp MCPConfig) {
fmt.Println("----------------------------------------------------------------") fmt.Println("----------------------------------------------------------------")
} }
// ResolveToolsDir 将 tools_dir 解析为绝对路径(相对路径相对于 configPath 所在目录)。
func ResolveToolsDir(toolsDir, configPath string) string {
toolsDir = strings.TrimSpace(toolsDir)
if toolsDir == "" {
return ""
}
if filepath.IsAbs(toolsDir) {
return toolsDir
}
return filepath.Join(filepath.Dir(configPath), toolsDir)
}
// MergeToolsFromDir 从目录加载工具并与 inline 列表合并:目录中的工具优先,主配置中的工具作为补充。
func MergeToolsFromDir(toolsDir string, inlineTools []ToolConfig) ([]ToolConfig, error) {
dirTools, err := LoadToolsFromDir(toolsDir)
if err != nil {
return nil, err
}
existing := make(map[string]bool, len(dirTools))
for _, tool := range dirTools {
existing[tool.Name] = true
}
merged := append([]ToolConfig(nil), dirTools...)
for _, tool := range inlineTools {
if !existing[tool.Name] {
merged = append(merged, tool)
}
}
return merged, nil
}
// loadInlineSecurityToolsFromYAML 读取 config.yaml 中 security.tools(不含 tools_dir 扫描结果)。
func loadInlineSecurityToolsFromYAML(configPath string) ([]ToolConfig, error) {
data, err := os.ReadFile(configPath)
if err != nil {
return nil, fmt.Errorf("读取配置文件失败: %w", err)
}
var partial struct {
Security struct {
Tools []ToolConfig `yaml:"tools"`
} `yaml:"security"`
}
if err := yaml.Unmarshal(data, &partial); err != nil {
return nil, fmt.Errorf("解析配置文件失败: %w", err)
}
if partial.Security.Tools == nil {
return []ToolConfig{}, nil
}
return partial.Security.Tools, nil
}
// ReloadSecurityToolsFromDir 从 tools_dir 重新加载工具并更新 cfg.Security.ToolsApplyConfig 热重载用)。
func ReloadSecurityToolsFromDir(cfg *Config, configPath string) error {
if cfg == nil || strings.TrimSpace(cfg.Security.ToolsDir) == "" {
return nil
}
inlineTools, err := loadInlineSecurityToolsFromYAML(configPath)
if err != nil {
return err
}
toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, configPath)
merged, err := MergeToolsFromDir(toolsDir, inlineTools)
if err != nil {
return fmt.Errorf("从工具目录加载工具配置失败: %w", err)
}
cfg.Security.Tools = merged
return nil
}
// LoadToolsFromDir 从目录加载所有工具配置文件 // LoadToolsFromDir 从目录加载所有工具配置文件
func LoadToolsFromDir(dir string) ([]ToolConfig, error) { func LoadToolsFromDir(dir string) ([]ToolConfig, error) {
var tools []ToolConfig var tools []ToolConfig
@@ -1316,7 +1477,12 @@ func Default() *Config {
}, },
Retrieval: RetrievalConfig{ Retrieval: RetrievalConfig{
TopK: 5, TopK: 5,
SimilarityThreshold: 0.65, // 降低阈值到 0.65,减少漏检 SimilarityThreshold: 0.65,
MultiQuery: MultiQueryConfig{MaxQueries: 4},
Rerank: RerankConfig{},
PostRetrieve: PostRetrieveConfig{
PrefetchTopK: 20,
},
}, },
Indexing: IndexingConfig{ Indexing: IndexingConfig{
ChunkStrategy: "markdown_then_recursive", ChunkStrategy: "markdown_then_recursive",
@@ -1412,7 +1578,7 @@ type EmbeddingConfig struct {
// PostRetrieveConfig 检索后处理:固定对正文做规范化去重(最佳实践)、上下文预算截断;PrefetchTopK 用于多取候选再收敛到 top_k。 // PostRetrieveConfig 检索后处理:固定对正文做规范化去重(最佳实践)、上下文预算截断;PrefetchTopK 用于多取候选再收敛到 top_k。
type PostRetrieveConfig struct { type PostRetrieveConfig struct {
// PrefetchTopK 向量检索阶段最多保留的候选数(余弦序),应 ≥ top_k,0 表示与 top_k 相同;上限见知识库包内常量 // PrefetchTopK 向量检索阶段每条 MultiQuery 变体最多保留的候选数;0 表示使用内置默认 max(top_k*4, 20)
PrefetchTopK int `yaml:"prefetch_top_k,omitempty" json:"prefetch_top_k,omitempty"` PrefetchTopK int `yaml:"prefetch_top_k,omitempty" json:"prefetch_top_k,omitempty"`
// MaxContextChars 返回文档内容总 Unicode 字符数上限(整段 chunk,不截断半段);0 表示不限制。 // MaxContextChars 返回文档内容总 Unicode 字符数上限(整段 chunk,不截断半段);0 表示不限制。
MaxContextChars int `yaml:"max_context_chars,omitempty" json:"max_context_chars,omitempty"` MaxContextChars int `yaml:"max_context_chars,omitempty" json:"max_context_chars,omitempty"`
@@ -1420,13 +1586,62 @@ type PostRetrieveConfig struct {
MaxContextTokens int `yaml:"max_context_tokens,omitempty" json:"max_context_tokens,omitempty"` MaxContextTokens int `yaml:"max_context_tokens,omitempty" json:"max_context_tokens,omitempty"`
} }
// MultiQueryConfig Eino MultiQuery 查询改写(始终启用,无关闭开关)。
type MultiQueryConfig struct {
// MaxQueries LLM 生成的检索变体上限(含原问语义覆盖);0 表示默认 4。
MaxQueries int `yaml:"max_queries,omitempty" json:"max_queries,omitempty"`
}
func (c MultiQueryConfig) MaxQueriesEffective() int {
if c.MaxQueries <= 0 {
return 4
}
if c.MaxQueries > 8 {
return 8
}
return c.MaxQueries
}
// RerankConfig 检索精排(始终启用);支持 dashscope 与 Cohere 兼容 HTTP API。
type RerankConfig struct {
// Provider: dashscope | cohere;空则按 base_url 自动推断。
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
Model string `yaml:"model,omitempty" json:"model,omitempty"`
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"`
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
}
func (c RerankConfig) ProviderEffective(baseURL string) string {
p := strings.TrimSpace(strings.ToLower(c.Provider))
if p != "" {
return p
}
u := strings.ToLower(baseURL)
if strings.Contains(u, "dashscope") {
return "dashscope"
}
return "cohere"
}
func (c RerankConfig) ModelEffective(provider string) string {
if m := strings.TrimSpace(c.Model); m != "" {
return m
}
if provider == "dashscope" {
return "gte-rerank"
}
return "rerank-multilingual-v3.0"
}
// RetrievalConfig 检索配置 // RetrievalConfig 检索配置
type RetrievalConfig struct { type RetrievalConfig struct {
TopK int `yaml:"top_k" json:"top_k"` // 检索Top-K TopK int `yaml:"top_k" json:"top_k"` // 检索Top-K
SimilarityThreshold float64 `yaml:"similarity_threshold" json:"similarity_threshold"` // 余弦相似度阈值 SimilarityThreshold float64 `yaml:"similarity_threshold" json:"similarity_threshold"` // 余弦相似度阈值
// SubIndexFilter 非空时仅保留 sub_indexes 含该标签(逗号分隔之一)的行;sub_indexes 为空的旧行仍返回。 // SubIndexFilter 非空时仅保留 sub_indexes 含该标签(逗号分隔之一)的行;sub_indexes 为空的旧行仍返回。
SubIndexFilter string `yaml:"sub_index_filter,omitempty" json:"sub_index_filter,omitempty"` SubIndexFilter string `yaml:"sub_index_filter,omitempty" json:"sub_index_filter,omitempty"`
// PostRetrieve 检索后处理(去重、预算截断);重排通过代码注入 [knowledge.DocumentReranker]。 MultiQuery MultiQueryConfig `yaml:"multi_query" json:"multi_query"`
Rerank RerankConfig `yaml:"rerank" json:"rerank"`
// PostRetrieve 检索后处理(去重、预算截断);精排在 MultiQuery 融合后执行。
PostRetrieve PostRetrieveConfig `yaml:"post_retrieve,omitempty" json:"post_retrieve,omitempty"` PostRetrieve PostRetrieveConfig `yaml:"post_retrieve,omitempty" json:"post_retrieve,omitempty"`
} }
+45
View File
@@ -0,0 +1,45 @@
package config
import "testing"
func TestValidateWecomConfig(t *testing.T) {
t.Parallel()
tests := []struct {
name string
cfg RobotWecomConfig
wantErr bool
}{
{
name: "disabled without token",
cfg: RobotWecomConfig{Enabled: false, Token: ""},
wantErr: false,
},
{
name: "enabled with token",
cfg: RobotWecomConfig{Enabled: true, Token: "secret"},
wantErr: false,
},
{
name: "enabled without token",
cfg: RobotWecomConfig{Enabled: true, Token: ""},
wantErr: true,
},
{
name: "enabled with whitespace token",
cfg: RobotWecomConfig{Enabled: true, Token: " "},
wantErr: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := ValidateWecomConfig(tt.cfg)
if (err != nil) != tt.wantErr {
t.Fatalf("ValidateWecomConfig() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
+111
View File
@@ -0,0 +1,111 @@
package config
import (
"os"
"path/filepath"
"testing"
)
func TestReloadSecurityToolsFromDir(t *testing.T) {
root := t.TempDir()
toolsDir := filepath.Join(root, "tools")
if err := os.MkdirAll(toolsDir, 0755); err != nil {
t.Fatal(err)
}
configPath := filepath.Join(root, "config.yaml")
if err := os.WriteFile(configPath, []byte(`security:
tools_dir: tools
tools:
- name: inline-only
command: inline-cmd
enabled: true
description: inline tool
`), 0644); err != nil {
t.Fatal(err)
}
writeTool := func(name, command string) {
t.Helper()
content := "name: " + name + "\ncommand: " + command + "\nenabled: true\ndescription: test\n"
if err := os.WriteFile(filepath.Join(toolsDir, name+".yaml"), []byte(content), 0644); err != nil {
t.Fatal(err)
}
}
writeTool("alpha", "alpha-cmd")
cfg := &Config{
Security: SecurityConfig{
ToolsDir: "tools",
Tools: []ToolConfig{
{Name: "stale", Command: "stale-cmd", Enabled: true, Description: "should be removed"},
},
},
}
if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil {
t.Fatalf("reload: %v", err)
}
if len(cfg.Security.Tools) != 2 {
t.Fatalf("expected 2 tools, got %d", len(cfg.Security.Tools))
}
names := map[string]string{}
for _, tool := range cfg.Security.Tools {
names[tool.Name] = tool.Command
}
if names["alpha"] != "alpha-cmd" {
t.Fatalf("alpha tool missing or wrong command: %#v", names)
}
if names["inline-only"] != "inline-cmd" {
t.Fatalf("inline-only tool missing: %#v", names)
}
if _, ok := names["stale"]; ok {
t.Fatal("stale in-memory tool should not survive reload")
}
writeTool("beta", "beta-cmd")
if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil {
t.Fatalf("second reload: %v", err)
}
if len(cfg.Security.Tools) != 3 {
t.Fatalf("expected 3 tools after add, got %d", len(cfg.Security.Tools))
}
foundBeta := false
for _, tool := range cfg.Security.Tools {
if tool.Name == "beta" {
foundBeta = true
break
}
}
if !foundBeta {
t.Fatal("beta tool not found after second reload")
}
}
func TestMergeToolsFromDir_DirOverridesInline(t *testing.T) {
root := t.TempDir()
toolsDir := filepath.Join(root, "tools")
if err := os.MkdirAll(toolsDir, 0755); err != nil {
t.Fatal(err)
}
content := "name: shared\ncommand: dir-cmd\nenabled: true\ndescription: from dir\n"
if err := os.WriteFile(filepath.Join(toolsDir, "shared.yaml"), []byte(content), 0644); err != nil {
t.Fatal(err)
}
inline := []ToolConfig{
{Name: "shared", Command: "inline-cmd", Enabled: true, Description: "from inline"},
}
merged, err := MergeToolsFromDir(toolsDir, inline)
if err != nil {
t.Fatal(err)
}
if len(merged) != 1 {
t.Fatalf("expected 1 tool, got %d", len(merged))
}
if merged[0].Command != "dir-cmd" {
t.Fatalf("dir tool should win, got command %q", merged[0].Command)
}
}
+149 -20
View File
@@ -3,6 +3,7 @@ package database
import ( import (
"database/sql" "database/sql"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
@@ -13,6 +14,9 @@ import (
"go.uber.org/zap" "go.uber.org/zap"
) )
// ProjectFilterUnbound 列表 API 中 project_id=__none__ 表示仅未绑定项目的对话。
const ProjectFilterUnbound = "__none__"
// Conversation 对话 // Conversation 对话
type Conversation struct { type Conversation struct {
ID string `json:"id"` ID string `json:"id"`
@@ -361,20 +365,44 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
return &conv, nil return &conv, nil
} }
func conversationProjectIDColumn(alias string) string {
if alias != "" {
return alias + ".project_id"
}
return "project_id"
}
func appendConversationProjectFilter(where string, args []interface{}, projectID, alias string) (string, []interface{}) {
pid := strings.TrimSpace(projectID)
if pid == "" {
return where, args
}
col := conversationProjectIDColumn(alias)
if pid == ProjectFilterUnbound {
return where + fmt.Sprintf(" AND (%s IS NULL OR TRIM(COALESCE(%s, '')) = '')", col, col), args
}
return where + fmt.Sprintf(" AND %s = ?", col), append(args, pid)
}
// CountConversations 统计对话数量。 // CountConversations 统计对话数量。
func (db *DB) CountConversations(search string) (int, error) { func (db *DB) CountConversations(search, projectID string) (int, error) {
var count int var count int
var err error var err error
if search != "" { if search != "" {
searchPattern := "%" + search + "%" searchPattern := "%" + search + "%"
err = db.QueryRow( where := ` WHERE (c.title LIKE ?
`SELECT COUNT(*) FROM conversations c OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
WHERE c.title LIKE ? args := []interface{}{searchPattern, searchPattern}
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)`, where, args = appendConversationProjectFilter(where, args, projectID, "c")
searchPattern, searchPattern, err = db.QueryRow(`SELECT COUNT(*) FROM conversations c`+where, args...).Scan(&count)
).Scan(&count)
} else { } else {
err = db.QueryRow(`SELECT COUNT(*) FROM conversations`).Scan(&count) where := ""
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "")
if where != "" {
where = " WHERE" + strings.TrimPrefix(where, " AND")
}
err = db.QueryRow(`SELECT COUNT(*) FROM conversations`+where, args...).Scan(&count)
} }
if err != nil { if err != nil {
return 0, fmt.Errorf("统计对话失败: %w", err) return 0, fmt.Errorf("统计对话失败: %w", err)
@@ -395,7 +423,7 @@ func conversationOrderClause(sortBy, tableAlias string) string {
} }
// ListConversations 列出所有对话 // ListConversations 列出所有对话
func (db *DB) ListConversations(limit, offset int, search, sortBy string) ([]*Conversation, error) { func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID string) ([]*Conversation, error) {
var rows *sql.Rows var rows *sql.Rows
var err error var err error
@@ -403,20 +431,30 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy string) ([]*Co
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积 // 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
searchPattern := "%" + search + "%" searchPattern := "%" + search + "%"
orderClause := conversationOrderClause(sortBy, "c") orderClause := conversationOrderClause(sortBy, "c")
where := ` WHERE (c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
args := []interface{}{searchPattern, searchPattern}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
args = append(args, limit, offset)
rows, err = db.Query( rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
FROM conversations c FROM conversations c`+where+`
WHERE c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
`+orderClause+` `+orderClause+`
LIMIT ? OFFSET ?`, LIMIT ? OFFSET ?`,
searchPattern, searchPattern, limit, offset, args...,
) )
} else { } else {
orderClause := conversationOrderClause(sortBy, "") orderClause := conversationOrderClause(sortBy, "")
where := ""
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "")
if where != "" {
where = " WHERE" + strings.TrimPrefix(where, " AND")
}
args = append(args, limit, offset)
rows, err = db.Query( rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations "+orderClause+" LIMIT ? OFFSET ?", "SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
limit, offset, args...,
) )
} }
@@ -472,23 +510,30 @@ const ungroupedConversationsSQL = `
)` )`
// CountUngroupedConversations 统计不在任何分组中的对话数量。 // CountUngroupedConversations 统计不在任何分组中的对话数量。
func (db *DB) CountUngroupedConversations() (int, error) { func (db *DB) CountUngroupedConversations(projectID string) (int, error) {
where := ungroupedConversationsSQL
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
var count int var count int
if err := db.QueryRow(`SELECT COUNT(*) ` + ungroupedConversationsSQL).Scan(&count); err != nil { if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil {
return 0, fmt.Errorf("统计未分组对话失败: %w", err) return 0, fmt.Errorf("统计未分组对话失败: %w", err)
} }
return count, nil return count, nil
} }
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。 // ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
func (db *DB) ListUngroupedConversations(limit, offset int, sortBy string) ([]*Conversation, error) { func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID string) ([]*Conversation, error) {
orderClause := conversationOrderClause(sortBy, "c") orderClause := conversationOrderClause(sortBy, "c")
where := ungroupedConversationsSQL
args := []interface{}{}
where, args = appendConversationProjectFilter(where, args, projectID, "c")
args = append(args, limit, offset)
rows, err := db.Query( rows, err := db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+ `SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
ungroupedConversationsSQL+` where+`
`+orderClause+` `+orderClause+`
LIMIT ? OFFSET ?`, LIMIT ? OFFSET ?`,
limit, offset, args...,
) )
if err != nil { if err != nil {
return nil, fmt.Errorf("查询未分组对话失败: %w", err) return nil, fmt.Errorf("查询未分组对话失败: %w", err)
@@ -533,6 +578,19 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy string) ([]*C
return conversations, rows.Err() return conversations, rows.Err()
} }
// GetConversationTitle 获取对话标题(轻量查询,不加载消息)
func (db *DB) GetConversationTitle(id string) (string, error) {
var title string
err := db.QueryRow("SELECT title FROM conversations WHERE id = ?", id).Scan(&title)
if err != nil {
if err == sql.ErrNoRows {
return "", fmt.Errorf("对话不存在")
}
return "", fmt.Errorf("查询对话标题失败: %w", err)
}
return title, nil
}
// UpdateConversationTitle 更新对话标题 // UpdateConversationTitle 更新对话标题
func (db *DB) UpdateConversationTitle(id, title string) error { func (db *DB) UpdateConversationTitle(id, title string) error {
// 注意:不更新 updated_at,因为重命名操作不应该改变对话的更新时间 // 注意:不更新 updated_at,因为重命名操作不应该改变对话的更新时间
@@ -1013,6 +1071,77 @@ type ProcessDetail struct {
CreatedAt time.Time `json:"createdAt"` CreatedAt time.Time `json:"createdAt"`
} }
// GetTurnUserMessage 返回锚点消息所在轮次中的用户原文(最近一条 user 消息,不含完整历史)。
func (db *DB) GetTurnUserMessage(conversationID, anchorMessageID string) (string, error) {
conversationID = strings.TrimSpace(conversationID)
anchorMessageID = strings.TrimSpace(anchorMessageID)
if conversationID == "" || anchorMessageID == "" {
return "", nil
}
var content string
err := db.QueryRow(`
SELECT m.content FROM messages m
WHERE m.conversation_id = ? AND m.role = 'user'
AND m.created_at <= COALESCE((SELECT created_at FROM messages WHERE id = ? AND conversation_id = ?), m.created_at)
ORDER BY m.created_at DESC, m.rowid DESC
LIMIT 1`, conversationID, anchorMessageID, conversationID).Scan(&content)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return "", nil
}
return "", fmt.Errorf("query turn user message: %w", err)
}
return content, nil
}
// AssistantCognitionTexts 单条助手消息上的思考/推理/规划文本。
type AssistantCognitionTexts struct {
Thinking string
ReasoningChain string
Planning string
}
// GetAssistantCognitionTexts 聚合助手消息在 process_details 中的 thinking / reasoning_chain / planning。
func (db *DB) GetAssistantCognitionTexts(assistantMessageID string) (AssistantCognitionTexts, error) {
assistantMessageID = strings.TrimSpace(assistantMessageID)
if assistantMessageID == "" {
return AssistantCognitionTexts{}, nil
}
rows, err := db.Query(`
SELECT event_type, message FROM process_details
WHERE message_id = ? AND event_type IN ('thinking', 'reasoning_chain', 'planning')
ORDER BY created_at ASC, rowid ASC`, assistantMessageID)
if err != nil {
return AssistantCognitionTexts{}, fmt.Errorf("query assistant cognition: %w", err)
}
defer rows.Close()
var thinkingParts, reasoningParts, planningParts []string
for rows.Next() {
var eventType, message string
if err := rows.Scan(&eventType, &message); err != nil {
continue
}
msg := strings.TrimSpace(message)
if msg == "" {
continue
}
switch eventType {
case "thinking":
thinkingParts = append(thinkingParts, msg)
case "reasoning_chain":
reasoningParts = append(reasoningParts, msg)
case "planning":
planningParts = append(planningParts, msg)
}
}
return AssistantCognitionTexts{
Thinking: strings.Join(thinkingParts, "\n\n"),
ReasoningChain: strings.Join(reasoningParts, "\n\n"),
Planning: strings.Join(planningParts, "\n\n"),
}, nil
}
// AddProcessDetail 添加过程详情事件 // AddProcessDetail 添加过程详情事件
func (db *DB) AddProcessDetail(messageID, conversationID, eventType, message string, data interface{}) error { func (db *DB) AddProcessDetail(messageID, conversationID, eventType, message string, data interface{}) error {
id := uuid.New().String() id := uuid.New().String()
@@ -0,0 +1,60 @@
package database
import (
"path/filepath"
"testing"
"go.uber.org/zap"
)
func TestConversationProjectFilter(t *testing.T) {
tmp := t.TempDir()
dbPath := filepath.Join(tmp, "conversations.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
p, err := db.CreateProject(&Project{Name: "target-a", Status: "active"})
if err != nil {
t.Fatalf("CreateProject: %v", err)
}
convNone, err := db.CreateConversation("unbound", ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation unbound: %v", err)
}
convBound, err := db.CreateConversation("bound", ConversationCreateMeta{ProjectID: p.ID})
if err != nil {
t.Fatalf("CreateConversation bound: %v", err)
}
totalAll, err := db.CountConversations("", "")
if err != nil || totalAll < 2 {
t.Fatalf("CountConversations all: total=%d err=%v", totalAll, err)
}
totalBound, err := db.CountConversations("", p.ID)
if err != nil || totalBound != 1 {
t.Fatalf("CountConversations project: total=%d err=%v", totalBound, err)
}
totalUnbound, err := db.CountConversations("", ProjectFilterUnbound)
if err != nil || totalUnbound != 1 {
t.Fatalf("CountConversations unbound: total=%d err=%v", totalUnbound, err)
}
listBound, err := db.ListConversations(10, 0, "", "", p.ID)
if err != nil || len(listBound) != 1 || listBound[0].ID != convBound.ID {
t.Fatalf("ListConversations project: %+v err=%v", listBound, err)
}
listUnbound, err := db.ListConversations(10, 0, "", "", ProjectFilterUnbound)
if err != nil || len(listUnbound) != 1 || listUnbound[0].ID != convNone.ID {
t.Fatalf("ListConversations unbound: %+v err=%v", listUnbound, err)
}
_ = convNone
_ = convBound
}
+75
View File
@@ -0,0 +1,75 @@
package database
import (
"fmt"
"strings"
"time"
"go.uber.org/zap"
)
// DeleteHitlInterruptLogsByIDs deletes decided HITL audit logs by id (pending rows are skipped).
func (db *DB) DeleteHitlInterruptLogsByIDs(ids []string) (int64, error) {
if db == nil {
return 0, fmt.Errorf("database is nil")
}
clean := make([]string, 0, len(ids))
for _, id := range ids {
id = strings.TrimSpace(id)
if id != "" {
clean = append(clean, id)
}
}
if len(clean) == 0 {
return 0, nil
}
placeholders := strings.TrimRight(strings.Repeat("?,", len(clean)), ",")
q := fmt.Sprintf(`DELETE FROM hitl_interrupts WHERE status != 'pending' AND id IN (%s)`, placeholders)
args := make([]interface{}, len(clean))
for i, id := range clean {
args[i] = id
}
res, err := db.Exec(q, args...)
if err != nil {
db.logger.Error("批量删除人机协同审计日志失败", zap.Error(err), zap.Int("count", len(clean)))
return 0, fmt.Errorf("批量删除人机协同审计日志失败: %w", err)
}
n, _ := res.RowsAffected()
return n, nil
}
// DeleteHitlInterruptLogsMatching deletes decided logs matching whereSQL (e.g. "WHERE 1=1 AND status != 'pending' ...").
func (db *DB) DeleteHitlInterruptLogsMatching(whereSQL string, args []interface{}) (int64, error) {
if db == nil {
return 0, fmt.Errorf("database is nil")
}
whereSQL = strings.TrimSpace(whereSQL)
if whereSQL == "" {
return 0, fmt.Errorf("where clause is required")
}
q := `DELETE FROM hitl_interrupts ` + whereSQL
res, err := db.Exec(q, args...)
if err != nil {
db.logger.Error("清空人机协同审计日志失败", zap.Error(err))
return 0, fmt.Errorf("清空人机协同审计日志失败: %w", err)
}
n, _ := res.RowsAffected()
return n, nil
}
// PurgeHitlInterruptLogsBefore deletes decided logs with decided/created time before cutoff.
func (db *DB) PurgeHitlInterruptLogsBefore(cutoff time.Time) (int64, error) {
if db == nil {
return 0, fmt.Errorf("database is nil")
}
res, err := db.Exec(
`DELETE FROM hitl_interrupts WHERE status != 'pending' AND datetime(COALESCE(decided_at, created_at)) < datetime(?)`,
cutoff.UTC().Format(time.RFC3339),
)
if err != nil {
db.logger.Error("清理过期人机协同审计日志失败", zap.Error(err))
return 0, fmt.Errorf("清理过期人机协同审计日志失败: %w", err)
}
n, _ := res.RowsAffected()
return n, nil
}
+106
View File
@@ -0,0 +1,106 @@
package database
import (
"path/filepath"
"testing"
"time"
"go.uber.org/zap"
)
func ensureHitlInterruptsTable(t *testing.T, db *DB) {
t.Helper()
if _, err := db.Exec(`
CREATE TABLE IF NOT EXISTS hitl_interrupts (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL,
message_id TEXT,
mode TEXT NOT NULL,
tool_name TEXT NOT NULL,
tool_call_id TEXT,
payload TEXT,
status TEXT NOT NULL,
decision TEXT,
decision_comment TEXT,
created_at DATETIME NOT NULL,
decided_at DATETIME
);`); err != nil {
t.Fatalf("create hitl_interrupts: %v", err)
}
}
func TestDeleteHitlInterruptLogsByIDs_skipsPending(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "hitl.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
ensureHitlInterruptsTable(t, db)
now := time.Now().UTC().Format(time.RFC3339)
if _, err := db.Exec(`INSERT INTO hitl_interrupts
(id, conversation_id, mode, tool_name, status, created_at)
VALUES ('pending-1', 'c1', 'approval', 'exec', 'pending', ?)`, now); err != nil {
t.Fatalf("insert pending: %v", err)
}
if _, err := db.Exec(`INSERT INTO hitl_interrupts
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
VALUES ('done-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, now, now); err != nil {
t.Fatalf("insert decided: %v", err)
}
deleted, err := db.DeleteHitlInterruptLogsByIDs([]string{"pending-1", "done-1"})
if err != nil {
t.Fatalf("DeleteHitlInterruptLogsByIDs: %v", err)
}
if deleted != 1 {
t.Fatalf("deleted = %d, want 1", deleted)
}
var status string
if err := db.QueryRow(`SELECT status FROM hitl_interrupts WHERE id = 'pending-1'`).Scan(&status); err != nil {
t.Fatalf("pending row missing: %v", err)
}
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'done-1'`).Scan(new(string)); err == nil {
t.Fatal("decided row should be deleted")
}
}
func TestPurgeHitlInterruptLogsBefore(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "hitl.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
ensureHitlInterruptsTable(t, db)
old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339)
recent := time.Now().AddDate(0, 0, -1).UTC().Format(time.RFC3339)
for _, row := range []struct{ id, decided string }{
{"old-1", old},
{"new-1", recent},
} {
if _, err := db.Exec(`INSERT INTO hitl_interrupts
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
VALUES (?, 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, row.id, row.decided, row.decided); err != nil {
t.Fatalf("insert %s: %v", row.id, err)
}
}
cutoff := time.Now().AddDate(0, 0, -90)
deleted, err := db.PurgeHitlInterruptLogsBefore(cutoff)
if err != nil {
t.Fatalf("PurgeHitlInterruptLogsBefore: %v", err)
}
if deleted != 1 {
t.Fatalf("deleted = %d, want 1", deleted)
}
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err == nil {
t.Fatal("old row should be purged")
}
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'new-1'`).Scan(new(string)); err != nil {
t.Fatalf("new row should remain: %v", err)
}
}
+288 -26
View File
@@ -3,7 +3,6 @@ package database
import ( import (
"database/sql" "database/sql"
"encoding/json" "encoding/json"
"sort"
"strings" "strings"
"time" "time"
@@ -227,6 +226,167 @@ func (db *DB) LoadToolExecutionsWithPagination(offset, limit int, status, toolNa
return executions, nil return executions, nil
} }
func toolExecutionsFilterSQL(status, toolName string) (string, []interface{}) {
args := []interface{}{}
conditions := []string{}
if status != "" {
conditions = append(conditions, "status = ?")
args = append(args, status)
}
if toolName != "" {
conditions = append(conditions, "LOWER(tool_name) LIKE ?")
args = append(args, "%"+strings.ToLower(toolName)+"%")
}
if len(conditions) == 0 {
return "", args
}
return ` WHERE ` + strings.Join(conditions, ` AND `), args
}
// ToolStatsSummary 工具调用汇总(全量聚合,不含逐工具明细)
type ToolStatsSummary struct {
TotalCalls int
SuccessCalls int
FailedCalls int
LastCallTime *time.Time
ToolCount int
}
// ToolStatsSummaryResult 汇总 + Top N 工具排行
type ToolStatsSummaryResult struct {
Summary ToolStatsSummary
TopTools []*mcp.ToolStats
}
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
if topN <= 0 {
topN = 6
}
if topN > 100 {
topN = 100
}
result := &ToolStatsSummaryResult{
TopTools: make([]*mcp.ToolStats, 0, topN),
}
summaryQuery := `
SELECT COUNT(*),
COALESCE(SUM(total_calls), 0),
COALESCE(SUM(success_calls), 0),
COALESCE(SUM(failed_calls), 0),
MAX(last_call_time)
FROM tool_stats
`
var lastCallRaw sql.NullString
err := db.QueryRow(summaryQuery).Scan(
&result.Summary.ToolCount,
&result.Summary.TotalCalls,
&result.Summary.SuccessCalls,
&result.Summary.FailedCalls,
&lastCallRaw,
)
if err != nil {
return nil, err
}
if lastCallRaw.Valid && strings.TrimSpace(lastCallRaw.String) != "" {
if t, parseErr := time.Parse(time.RFC3339Nano, lastCallRaw.String); parseErr == nil {
result.Summary.LastCallTime = &t
} else if t, parseErr := time.Parse("2006-01-02 15:04:05.999999999-07:00", lastCallRaw.String); parseErr == nil {
result.Summary.LastCallTime = &t
} else if t, parseErr := time.Parse("2006-01-02 15:04:05", lastCallRaw.String); parseErr == nil {
result.Summary.LastCallTime = &t
}
}
topQuery := `
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
FROM tool_stats
WHERE total_calls > 0
ORDER BY total_calls DESC, tool_name ASC
LIMIT ?
`
rows, err := db.Query(topQuery, topN)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var stat mcp.ToolStats
var lastCallTime sql.NullTime
if err := rows.Scan(
&stat.ToolName,
&stat.TotalCalls,
&stat.SuccessCalls,
&stat.FailedCalls,
&lastCallTime,
); err != nil {
db.logger.Warn("加载 Top 工具统计失败", zap.Error(err))
continue
}
if lastCallTime.Valid {
stat.LastCallTime = &lastCallTime.Time
}
result.TopTools = append(result.TopTools, &stat)
}
return result, nil
}
// LoadToolExecutionListPage 分页加载执行记录列表(不含 arguments/result,供监控列表使用)
func (db *DB) LoadToolExecutionListPage(offset, limit int, status, toolName string) ([]*mcp.ToolExecution, error) {
if limit <= 0 {
limit = 20
}
if limit > 100 {
limit = 100
}
query := `
SELECT id, tool_name, status, start_time, end_time, duration_ms
FROM tool_executions
`
whereSQL, args := toolExecutionsFilterSQL(status, toolName)
query += whereSQL + ` ORDER BY start_time DESC LIMIT ? OFFSET ?`
args = append(args, limit, offset)
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
executions := make([]*mcp.ToolExecution, 0, limit)
for rows.Next() {
var exec mcp.ToolExecution
var endTime sql.NullTime
var durationMs sql.NullInt64
if err := rows.Scan(
&exec.ID,
&exec.ToolName,
&exec.Status,
&exec.StartTime,
&endTime,
&durationMs,
); err != nil {
db.logger.Warn("加载执行记录列表失败", zap.Error(err))
continue
}
if endTime.Valid {
exec.EndTime = &endTime.Time
}
if durationMs.Valid {
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
}
executions = append(executions, &exec)
}
return executions, nil
}
// GetToolExecution 根据ID获取单条工具执行记录 // GetToolExecution 根据ID获取单条工具执行记录
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) { func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
query := ` query := `
@@ -288,6 +448,93 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
return &exec, nil return &exec, nil
} }
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled(如进程重启后无对应执行协程)。
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
if errMsg == "" {
errMsg = "执行已中断(服务重启或会话结束)"
}
query := `
UPDATE tool_executions
SET status = 'cancelled',
error = ?,
end_time = ?,
duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER))
WHERE status = 'running'
`
res, err := db.Exec(query, errMsg, endTime, endTime)
if err != nil {
return 0, err
}
return res.RowsAffected()
}
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。
// activeIDs 为当前进程内仍登记 cancel 的 executionId;不在集合内且已超时的视为孤儿记录。
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
if errMsg == "" {
errMsg = "执行已中断(会话已结束)"
}
if minAge < 0 {
minAge = 0
}
cutoff := endTime.Add(-minAge)
rows, err := db.Query(`
SELECT id, start_time FROM tool_executions
WHERE status = 'running' AND start_time <= ?
`, cutoff)
if err != nil {
return 0, err
}
defer rows.Close()
type staleRow struct {
id string
startTime time.Time
}
var stale []staleRow
for rows.Next() {
var row staleRow
if err := rows.Scan(&row.id, &row.startTime); err != nil {
db.logger.Warn("读取 stale running 执行记录失败", zap.Error(err))
continue
}
if activeIDs != nil {
if _, active := activeIDs[row.id]; active {
continue
}
}
stale = append(stale, row)
}
if err := rows.Err(); err != nil {
return 0, err
}
if len(stale) == 0 {
return 0, nil
}
var affected int64
for _, row := range stale {
durationMs := endTime.Sub(row.startTime).Milliseconds()
if durationMs < 0 {
durationMs = 0
}
res, err := db.Exec(`
UPDATE tool_executions
SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ?
WHERE id = ? AND status = 'running'
`, errMsg, endTime, durationMs, row.id)
if err != nil {
db.logger.Warn("更新 stale running 执行记录失败", zap.Error(err), zap.String("executionId", row.id))
continue
}
n, _ := res.RowsAffected()
affected += n
}
return affected, nil
}
// DeleteToolExecution 删除工具执行记录 // DeleteToolExecution 删除工具执行记录
func (db *DB) DeleteToolExecution(id string) error { func (db *DB) DeleteToolExecution(id string) error {
query := `DELETE FROM tool_executions WHERE id = ?` query := `DELETE FROM tool_executions WHERE id = ?`
@@ -600,13 +847,28 @@ func truncateCallsTimelineBucket(t time.Time, dailyBuckets bool) time.Time {
// LoadCallsTimeline 按时间范围加载调用趋势(since 起至今,含边界) // LoadCallsTimeline 按时间范围加载调用趋势(since 起至今,含边界)
func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTimelineBucket, error) { func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTimelineBucket, error) {
// 在 Go 侧按本地时区分桶,避免 SQLite strftime 对 UTC 存储时间分桶后再误当本地时间解析(差 8h 等问题) var query string
query := ` if dailyBuckets {
SELECT start_time, query = `
CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END AS failed SELECT date(start_time, 'localtime') AS bucket,
FROM tool_executions COUNT(*) AS total,
WHERE start_time >= ? SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
` FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
ORDER BY bucket
`
} else {
query = `
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
ORDER BY bucket
`
}
rows, err := db.Query(query, since) rows, err := db.Query(query, since)
if err != nil { if err != nil {
@@ -614,35 +876,35 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
} }
defer rows.Close() defer rows.Close()
bucketMap := make(map[time.Time]struct{ total, failed int }) buckets := make([]CallsTimelineBucket, 0)
for rows.Next() { for rows.Next() {
var startTime time.Time var bucketStr string
var failed int var total, failed int
if err := rows.Scan(&startTime, &failed); err != nil { if err := rows.Scan(&bucketStr, &total, &failed); err != nil {
db.logger.Warn("加载调用趋势失败", zap.Error(err)) db.logger.Warn("加载调用趋势失败", zap.Error(err))
continue continue
} }
key := truncateCallsTimelineBucket(startTime, dailyBuckets) bucketTime, err := parseCallsTimelineBucket(bucketStr, dailyBuckets)
entry := bucketMap[key] if err != nil {
entry.total++ db.logger.Warn("解析调用趋势时间桶失败", zap.Error(err), zap.String("bucket", bucketStr))
entry.failed += failed continue
bucketMap[key] = entry }
}
buckets := make([]CallsTimelineBucket, 0, len(bucketMap))
for bucketTime, counts := range bucketMap {
buckets = append(buckets, CallsTimelineBucket{ buckets = append(buckets, CallsTimelineBucket{
BucketTime: bucketTime, BucketTime: bucketTime,
Total: counts.total, Total: total,
Failed: counts.failed, Failed: failed,
}) })
} }
sort.Slice(buckets, func(i, j int) bool {
return buckets[i].BucketTime.Before(buckets[j].BucketTime)
})
return buckets, nil return buckets, nil
} }
func parseCallsTimelineBucket(bucketStr string, dailyBuckets bool) (time.Time, error) {
if dailyBuckets {
return time.ParseInLocation("2006-01-02", bucketStr, time.Local)
}
return time.ParseInLocation("2006-01-02 15:04:05", bucketStr, time.Local)
}
// DecreaseToolStats 减少工具统计信息(用于删除执行记录时) // DecreaseToolStats 减少工具统计信息(用于删除执行记录时)
// 如果统计信息变为0,则删除该统计记录 // 如果统计信息变为0,则删除该统计记录
func (db *DB) DecreaseToolStats(toolName string, totalCalls, successCalls, failedCalls int) error { func (db *DB) DecreaseToolStats(toolName string, totalCalls, successCalls, failedCalls int) error {
+102
View File
@@ -0,0 +1,102 @@
package database
import (
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "monitor.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
start := time.Now().Add(-2 * time.Hour)
exec := &mcp.ToolExecution{
ID: "orphan-hydra",
ToolName: "hydra",
Arguments: map[string]interface{}{"target": "127.0.0.1"},
Status: "running",
StartTime: start,
}
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
end := time.Now()
n, err := db.CancelOrphanedRunningToolExecutions(end, "执行已中断(服务重启)")
if err != nil {
t.Fatalf("CancelOrphanedRunningToolExecutions: %v", err)
}
if n != 1 {
t.Fatalf("expected 1 row updated, got %d", n)
}
got, err := db.GetToolExecution("orphan-hydra")
if err != nil {
t.Fatalf("GetToolExecution: %v", err)
}
if got.Status != "cancelled" {
t.Fatalf("expected cancelled, got %s", got.Status)
}
if got.EndTime == nil {
t.Fatal("expected end_time to be set")
}
if got.Duration <= 0 {
t.Fatalf("expected positive duration, got %v", got.Duration)
}
}
func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "monitor.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
now := time.Now()
oldStart := now.Add(-5 * time.Minute)
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "stale", ToolName: "hydra", Status: "running", StartTime: oldStart,
}); err != nil {
t.Fatalf("SaveToolExecution stale: %v", err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "active", ToolName: "hydra", Status: "running", StartTime: oldStart,
}); err != nil {
t.Fatalf("SaveToolExecution active: %v", err)
}
active := map[string]struct{}{"active": {}}
n, err := db.FinalizeStaleRunningToolExecutions(now, time.Minute, active, "执行已中断(会话已结束)")
if err != nil {
t.Fatalf("FinalizeStaleRunningToolExecutions: %v", err)
}
if n != 1 {
t.Fatalf("expected 1 stale row updated, got %d", n)
}
stale, err := db.GetToolExecution("stale")
if err != nil {
t.Fatalf("GetToolExecution stale: %v", err)
}
if stale.Status != "cancelled" {
t.Fatalf("stale expected cancelled, got %s", stale.Status)
}
activeExec, err := db.GetToolExecution("active")
if err != nil {
t.Fatalf("GetToolExecution active: %v", err)
}
if activeExec.Status != "running" {
t.Fatalf("active expected running, got %s", activeExec.Status)
}
}
+86
View File
@@ -0,0 +1,86 @@
package database
import (
"fmt"
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestLoadToolStatsSummaryAndListPage(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "monitor-summary.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
now := time.Now()
tools := []struct {
name string
calls int
ok int
fail int
result string
}{
{"alpha::run", 10, 9, 1, `{"content":[{"type":"text","text":"` + string(make([]byte, 64*1024)) + `"}]}`},
{"beta::scan", 5, 5, 0, `{"content":[{"type":"text","text":"ok"}]}`},
{"gamma::ping", 1, 1, 0, `{"content":[{"type":"text","text":"pong"}]}`},
}
for _, tool := range tools {
if err := db.UpdateToolStats(tool.name, tool.calls, tool.ok, tool.fail, &now); err != nil {
t.Fatalf("UpdateToolStats(%s): %v", tool.name, err)
}
for j := 0; j < tool.calls; j++ {
exec := &mcp.ToolExecution{
ID: fmt.Sprintf("%s-exec-%d", tool.name, j),
ToolName: tool.name,
Arguments: map[string]interface{}{"n": j},
Status: "completed",
StartTime: now.Add(-time.Duration(j) * time.Minute),
Result: &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: tool.result}}},
}
end := exec.StartTime.Add(time.Second)
exec.EndTime = &end
exec.Duration = time.Second
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
}
}
summary, err := db.LoadToolStatsSummary(2)
if err != nil {
t.Fatalf("LoadToolStatsSummary: %v", err)
}
if summary.Summary.ToolCount != 3 {
t.Fatalf("toolCount = %d, want 3", summary.Summary.ToolCount)
}
if summary.Summary.TotalCalls != 16 {
t.Fatalf("totalCalls = %d, want 16", summary.Summary.TotalCalls)
}
if len(summary.TopTools) != 2 {
t.Fatalf("top tools = %d, want 2", len(summary.TopTools))
}
if summary.TopTools[0].ToolName != "alpha::run" {
t.Fatalf("top tool = %q, want alpha::run", summary.TopTools[0].ToolName)
}
list, err := db.LoadToolExecutionListPage(0, 5, "", "")
if err != nil {
t.Fatalf("LoadToolExecutionListPage: %v", err)
}
if len(list) != 5 {
t.Fatalf("list len = %d, want 5", len(list))
}
for _, exec := range list {
if exec.Arguments != nil || exec.Result != nil || exec.Error != "" {
t.Fatalf("expected lite execution row, got args/result/error on %s", exec.ID)
}
}
}
+33 -17
View File
@@ -111,19 +111,43 @@ func (db *DB) GetProject(id string) (*Project, error) {
return &p, nil return &p, nil
} }
// CountProjects 统计项目数量。 func projectListSearchPattern(q string) string {
func (db *DB) CountProjects(status, search string) (int, error) { q = strings.TrimSpace(q)
query := `SELECT COUNT(*) FROM projects WHERE 1=1` if q == "" {
args := []interface{}{} return ""
}
var b strings.Builder
b.WriteByte('%')
for _, r := range q {
switch r {
case '%', '_', '\\':
b.WriteByte('\\')
b.WriteRune(r)
default:
b.WriteRune(r)
}
}
b.WriteByte('%')
return b.String()
}
func appendProjectListFilters(query string, args []interface{}, status, search string) (string, []interface{}) {
if s := strings.TrimSpace(status); s != "" { if s := strings.TrimSpace(status); s != "" {
query += " AND status = ?" query += " AND status = ?"
args = append(args, s) args = append(args, s)
} }
if q := strings.TrimSpace(search); q != "" { if pattern := projectListSearchPattern(search); pattern != "" {
pattern := "%" + q + "%" query += ` AND (LOWER(name) LIKE LOWER(?) ESCAPE '\' OR LOWER(COALESCE(description,'')) LIKE LOWER(?) ESCAPE '\' OR LOWER(id) LIKE LOWER(?) ESCAPE '\')`
query += " AND (name LIKE ? OR COALESCE(description,'') LIKE ?)" args = append(args, pattern, pattern, pattern)
args = append(args, pattern, pattern)
} }
return query, args
}
// CountProjects 统计项目数量。
func (db *DB) CountProjects(status, search string) (int, error) {
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
args := []interface{}{}
query, args = appendProjectListFilters(query, args, status, search)
var count int var count int
if err := db.QueryRow(query, args...).Scan(&count); err != nil { if err := db.QueryRow(query, args...).Scan(&count); err != nil {
return 0, fmt.Errorf("统计项目失败: %w", err) return 0, fmt.Errorf("统计项目失败: %w", err)
@@ -139,15 +163,7 @@ func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
FROM projects WHERE 1=1` FROM projects WHERE 1=1`
args := []interface{}{} args := []interface{}{}
if s := strings.TrimSpace(status); s != "" { query, args = appendProjectListFilters(query, args, status, search)
query += " AND status = ?"
args = append(args, s)
}
if q := strings.TrimSpace(search); q != "" {
pattern := "%" + q + "%"
query += " AND (name LIKE ? OR COALESCE(description,'') LIKE ?)"
args = append(args, pattern, pattern)
}
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?" query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset) args = append(args, limit, offset)
+82
View File
@@ -0,0 +1,82 @@
package database
import (
"path/filepath"
"testing"
"go.uber.org/zap"
)
func TestListProjectsSearchCaseInsensitive(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "projects-search.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p1, err := db.CreateProject(&Project{Name: "Alpha Security Review", Status: "active"})
if err != nil {
t.Fatal(err)
}
p2, err := db.CreateProject(&Project{Name: "beta-scan", Status: "active"})
if err != nil {
t.Fatal(err)
}
if _, err := db.CreateProject(&Project{Name: "Other", Status: "archived"}); err != nil {
t.Fatal(err)
}
cases := []struct {
name string
search string
status string
want []string
}{
{name: "case insensitive name", search: "alpha", status: "active", want: []string{p1.ID}},
{name: "upper query", search: "BETA", status: "active", want: []string{p2.ID}},
{name: "search by id substring", search: p1.ID[:8], status: "", want: []string{p1.ID}},
{name: "status filter", search: "alpha", status: "archived", want: nil},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
list, err := db.ListProjects(tc.status, tc.search, 50, 0)
if err != nil {
t.Fatal(err)
}
got := make([]string, 0, len(list))
for _, p := range list {
got = append(got, p.ID)
}
if len(got) != len(tc.want) {
t.Fatalf("got %v want %v", got, tc.want)
}
for i := range got {
if got[i] != tc.want[i] {
t.Fatalf("got %v want %v", got, tc.want)
}
}
})
}
}
func TestProjectListSearchPatternEscapesWildcards(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "projects-like.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&Project{Name: "100% coverage", Status: "active"})
if err != nil {
t.Fatal(err)
}
list, err := db.ListProjects("active", "100%", 50, 0)
if err != nil {
t.Fatal(err)
}
if len(list) != 1 || list[0].ID != p.ID {
t.Fatalf("expected exact match for literal %% query, got %#v", list)
}
}
+2 -2
View File
@@ -2,8 +2,8 @@ package einomcp
import "sync" import "sync"
// ToolInvokeNotifyHolder 由 Eino run loop 在迭代开始前 Set 回调;MCP/execute 桥在工具调用结束时 Fire, // ToolInvokeNotifyHolder 由 Eino run loop 与 MCP/execute 桥共享;Fire 在工具原始返回时触发。
// 用于清除 pending tool_calltool_result ADK schema.Tool 事件推送,含流式工具与 reduction 后正文)。 // UI 的 tool_result 须等 ADK schema.Tool 事件reduction 后正文),不在此 holder 的回调里推送
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)
+99 -12
View File
@@ -77,6 +77,13 @@ type responsePlanAgg struct {
b strings.Builder b strings.Builder
} }
// thinkingBuf aggregates thinking_stream_* / reasoning_chain_stream_* before flush to process_details.
type thinkingBuf struct {
b strings.Builder
meta map[string]interface{}
persistAs string // "thinking" | "reasoning_chain"
}
func normalizeProcessDetailText(s string) string { func normalizeProcessDetailText(s string) string {
s = strings.ReplaceAll(s, "\r\n", "\n") s = strings.ReplaceAll(s, "\r\n", "\n")
s = strings.ReplaceAll(s, "\r", "\n") s = strings.ReplaceAll(s, "\r", "\n")
@@ -179,6 +186,8 @@ type AgentHandler struct {
batchCronParser cron.Parser batchCronParser cron.Parser
// hitlWhitelistSaver 侧栏「应用」HITL 时将会话增量白名单合并写入 config.yaml(可选) // hitlWhitelistSaver 侧栏「应用」HITL 时将会话增量白名单合并写入 config.yaml(可选)
hitlWhitelistSaver HitlToolWhitelistSaver hitlWhitelistSaver HitlToolWhitelistSaver
hitlStrategySaver HitlAuditStrategySaver
auditLLM *openai.Client
audit *audit.Service audit *audit.Service
} }
@@ -200,9 +209,8 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
if h == nil || conversationID == "" || h.tasks == nil { if h == nil || conversationID == "" || h.tasks == nil {
return return
} }
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" { h.cancelActiveMCPToolForConversation(conversationID)
h.agent.CancelMCPToolExecutionWithNote(execID, "") h.tasks.AbortActiveEinoExecute(conversationID, "")
}
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok { if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID)) h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
} else if err != nil { } else if err != nil {
@@ -210,9 +218,19 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
} }
} }
// HitlToolWhitelistSaver 合并 HITL 免审批工具到全局配置并落盘 func (h *AgentHandler) cancelActiveMCPToolForConversation(conversationID string) {
if h == nil || h.tasks == nil || h.agent == nil {
return
}
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
h.agent.CancelMCPToolExecutionWithNote(execID, "")
}
}
// HitlToolWhitelistSaver 合并/设置 HITL 免审批工具到全局配置并落盘
type HitlToolWhitelistSaver interface { type HitlToolWhitelistSaver interface {
MergeHitlToolWhitelistIntoConfig(add []string) error MergeHitlToolWhitelistIntoConfig(add []string) error
SetHitlToolWhitelist(tools []string) error
} }
// NewAgentHandler 创建新的Agent处理器 // NewAgentHandler 创建新的Agent处理器
@@ -228,6 +246,11 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
bus := NewTaskEventBus() bus := NewTaskEventBus()
tm := NewAgentTaskManager() tm := NewAgentTaskManager()
tm.SetTaskEventBus(bus) tm.SetTaskEventBus(bus)
llmHTTP := &http.Client{Timeout: 2 * time.Minute}
var llmCfg *config.OpenAIConfig
if cfg != nil {
llmCfg = &cfg.OpenAI
}
handler := &AgentHandler{ handler := &AgentHandler{
agent: agent, agent: agent,
db: db, db: db,
@@ -238,7 +261,9 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
config: cfg, config: cfg,
hitlManager: NewHITLManager(db, logger), hitlManager: NewHITLManager(db, logger),
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor), batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
} }
tm.SetToolCanceler(handler.cancelActiveMCPToolForConversation)
if err := handler.hitlManager.EnsureSchema(); err != nil { if err := handler.hitlManager.EnsureSchema(); err != nil {
logger.Warn("初始化 HITL 表失败", zap.Error(err)) logger.Warn("初始化 HITL 表失败", zap.Error(err))
} }
@@ -311,6 +336,7 @@ func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientInten
type HITLRequest struct { type HITLRequest struct {
Enabled bool `json:"enabled"` Enabled bool `json:"enabled"`
Mode string `json:"mode,omitempty"` Mode string `json:"mode,omitempty"`
Reviewer string `json:"reviewer,omitempty"` // human | audit_agent
SensitiveTools []string `json:"sensitiveTools,omitempty"` SensitiveTools []string `json:"sensitiveTools,omitempty"`
TimeoutSeconds int `json:"timeoutSeconds,omitempty"` TimeoutSeconds int `json:"timeoutSeconds,omitempty"`
} }
@@ -840,11 +866,6 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
// thinking_stream_*ReAct 等助手正文流)与 reasoning_chain_stream_*Eino ReasoningContent): // thinking_stream_*ReAct 等助手正文流)与 reasoning_chain_stream_*Eino ReasoningContent):
// 不逐条落库,按 streamId 聚合,flush 时分别落 thinking / reasoning_chain。 // 不逐条落库,按 streamId 聚合,flush 时分别落 thinking / reasoning_chain。
type thinkingBuf struct {
b strings.Builder
meta map[string]interface{}
persistAs string // "thinking" | "reasoning_chain"
}
thinkingStreams := make(map[string]*thinkingBuf) // streamId -> buf thinkingStreams := make(map[string]*thinkingBuf) // streamId -> buf
flushedThinking := make(map[string]bool) // streamId -> flushed flushedThinking := make(map[string]bool) // streamId -> flushed
seenToolCallSigs := make(map[string]string) // toolCallId -> payload signature seenToolCallSigs := make(map[string]string) // toolCallId -> payload signature
@@ -857,6 +878,12 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
// response_start + response_delta:前端时间线显示为「📝 规划中」(monitor.js),不落逐条 delta // response_start + response_delta:前端时间线显示为「📝 规划中」(monitor.js),不落逐条 delta
// 聚合为一条 planning 写入 process_details,刷新后与线上一致。 // 聚合为一条 planning 写入 process_details,刷新后与线上一致。
var respPlan responsePlanAgg var respPlan responsePlanAgg
if assistantMessageID != "" {
h.tasks.SetHitlAssistantMessageID(conversationID, assistantMessageID)
}
syncHitlCognition := func() {
h.syncHitlCognitionFromProgress(conversationID, assistantMessageID, thinkingStreams, &respPlan)
}
flushResponsePlan := func() { flushResponsePlan := func() {
if assistantMessageID == "" { if assistantMessageID == "" {
return return
@@ -876,6 +903,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "planning", content, data); err != nil { if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "planning", content, data); err != nil {
h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", "planning")) h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", "planning"))
} }
syncHitlCognition()
respPlan.meta = nil respPlan.meta = nil
respPlan.b.Reset() respPlan.b.Reset()
} }
@@ -912,6 +940,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
} }
flushedThinking[sid] = true flushedThinking[sid] = true
} }
syncHitlCognition()
} }
return func(eventType, message string, data interface{}) { return func(eventType, message string, data interface{}) {
@@ -972,6 +1001,25 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
} }
} }
if eventType == "tool_result" {
if dataMap, ok := data.(map[string]interface{}); ok {
toolName, _ := dataMap["toolName"].(string)
toolCallID, _ := dataMap["toolCallId"].(string)
success := true
if v, ok := dataMap["success"].(bool); ok {
success = v
}
resultText := ""
if r, ok := dataMap["result"].(string); ok {
resultText = r
}
if strings.TrimSpace(resultText) == "" {
resultText = message
}
h.recordHitlToolExecutionResult(conversationID, toolCallID, toolName, success, resultText)
}
}
// 处理知识检索日志记录 // 处理知识检索日志记录
if eventType == "tool_result" && h.knowledgeManager != nil { if eventType == "tool_result" && h.knowledgeManager != nil {
if dataMap, ok := data.(map[string]interface{}); ok { if dataMap, ok := data.(map[string]interface{}); ok {
@@ -1179,6 +1227,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
respPlan.meta[k] = v respPlan.meta[k] = v
} }
} }
syncHitlCognition()
return return
} }
if eventType == "response" { if eventType == "response" {
@@ -1248,6 +1297,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
} }
} }
} }
syncHitlCognition()
return return
} }
@@ -1352,6 +1402,7 @@ func (h *AgentHandler) cancelToolContinueAfter(conversationID, preferredExecID,
func (h *AgentHandler) CancelAgentLoop(c *gin.Context) { func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
var req struct { var req struct {
ConversationID string `json:"conversationId" binding:"required"` ConversationID string `json:"conversationId" binding:"required"`
ExecutionID string `json:"executionId,omitempty"`
Reason string `json:"reason,omitempty"` Reason string `json:"reason,omitempty"`
ContinueAfter bool `json:"continueAfter,omitempty"` ContinueAfter bool `json:"continueAfter,omitempty"`
} }
@@ -1368,7 +1419,7 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
} }
note := strings.TrimSpace(req.Reason) note := strings.TrimSpace(req.Reason)
activeExec := strings.TrimSpace(h.tasks.ActiveMCPExecutionID(req.ConversationID)) activeExec := strings.TrimSpace(h.tasks.ActiveMCPExecutionID(req.ConversationID))
if ok, payload := h.cancelToolContinueAfter(req.ConversationID, "", note); ok { if ok, payload := h.cancelToolContinueAfter(req.ConversationID, strings.TrimSpace(req.ExecutionID), note); ok {
execID, _ := payload["executionId"].(string) execID, _ := payload["executionId"].(string)
h.logger.Info("对话页仅终止当前工具", h.logger.Info("对话页仅终止当前工具",
zap.String("conversationId", req.ConversationID), zap.String("conversationId", req.ConversationID),
@@ -1411,6 +1462,8 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
var cause error = ErrTaskCancelled var cause error = ErrTaskCancelled
msg := "已提交取消请求,任务将在当前步骤完成后停止。" msg := "已提交取消请求,任务将在当前步骤完成后停止。"
h.cancelActiveMCPToolForConversation(req.ConversationID)
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
ok, err := h.tasks.CancelTask(req.ConversationID, cause) ok, err := h.tasks.CancelTask(req.ConversationID, cause)
if err != nil { if err != nil {
h.logger.Error("取消任务失败", zap.Error(err)) h.logger.Error("取消任务失败", zap.Error(err))
@@ -1477,17 +1530,51 @@ func (h *AgentHandler) SubscribeAgentTaskEvents(c *gin.Context) {
} }
} }
// enrichAgentTasksWithConversationTitles 为任务列表附加当前会话标题(供顶栏/任务页展示,重命名后自动同步)
func (h *AgentHandler) enrichAgentTasksWithConversationTitles(tasks []*AgentTask) {
if h == nil || h.db == nil {
return
}
for _, task := range tasks {
if task == nil || strings.TrimSpace(task.ConversationID) == "" {
continue
}
if title, err := h.db.GetConversationTitle(task.ConversationID); err == nil {
task.Title = strings.TrimSpace(title)
}
}
}
// enrichCompletedTasksWithConversationTitles 为已完成任务附加当前会话标题
func (h *AgentHandler) enrichCompletedTasksWithConversationTitles(tasks []*CompletedTask) {
if h == nil || h.db == nil {
return
}
for _, task := range tasks {
if task == nil || strings.TrimSpace(task.ConversationID) == "" {
continue
}
if title, err := h.db.GetConversationTitle(task.ConversationID); err == nil {
task.Title = strings.TrimSpace(title)
}
}
}
// ListAgentTasks 列出所有运行中的任务 // ListAgentTasks 列出所有运行中的任务
func (h *AgentHandler) ListAgentTasks(c *gin.Context) { func (h *AgentHandler) ListAgentTasks(c *gin.Context) {
tasks := h.tasks.GetActiveTasks()
h.enrichAgentTasksWithConversationTitles(tasks)
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"tasks": h.tasks.GetActiveTasks(), "tasks": tasks,
}) })
} }
// ListCompletedTasks 列出最近完成的任务历史 // ListCompletedTasks 列出最近完成的任务历史
func (h *AgentHandler) ListCompletedTasks(c *gin.Context) { func (h *AgentHandler) ListCompletedTasks(c *gin.Context) {
tasks := h.tasks.GetCompletedTasks()
h.enrichCompletedTasksWithConversationTitles(tasks)
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"tasks": h.tasks.GetCompletedTasks(), "tasks": tasks,
}) })
} }
+52 -6
View File
@@ -798,6 +798,10 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
// 更新机器人配置 // 更新机器人配置
if req.Robots != nil { if req.Robots != nil {
if err := config.ValidateWecomConfig(req.Robots.Wecom); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.config.Robots = *req.Robots h.config.Robots = *req.Robots
h.logger.Info("更新机器人配置", h.logger.Info("更新机器人配置",
zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled), zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled),
@@ -1329,6 +1333,17 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
h.logger.Info("已更新嵌入模型配置记录") h.logger.Info("已更新嵌入模型配置记录")
} }
// 从 tools 目录重新加载工具配置(新增/修改/删除 yaml 后无需重启)
if err := config.ReloadSecurityToolsFromDir(h.config, h.configPath); err != nil {
h.logger.Error("重新加载工具配置失败", zap.Error(err))
if h.audit != nil {
h.audit.RecordFail(c, "config", "apply", "应用配置失败:重新加载工具", map[string]interface{}{"error": err.Error()})
}
c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()})
return
}
h.logger.Info("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools)))
// 重新注册工具(根据新的启用状态) // 重新注册工具(根据新的启用状态)
h.logger.Info("重新注册工具") h.logger.Info("重新注册工具")
@@ -1417,12 +1432,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
// 更新检索器配置(如果知识库启用) // 更新检索器配置(如果知识库启用)
if h.config.Knowledge.Enabled && h.retrieverUpdater != nil { if h.config.Knowledge.Enabled && h.retrieverUpdater != nil {
retrievalConfig := &knowledge.RetrievalConfig{ retrievalConfig := knowledge.RetrievalConfigFromYAML(h.config.Knowledge.Retrieval)
TopK: h.config.Knowledge.Retrieval.TopK,
SimilarityThreshold: h.config.Knowledge.Retrieval.SimilarityThreshold,
SubIndexFilter: h.config.Knowledge.Retrieval.SubIndexFilter,
PostRetrieve: h.config.Knowledge.Retrieval.PostRetrieve,
}
h.retrieverUpdater.UpdateConfig(retrievalConfig) h.retrieverUpdater.UpdateConfig(retrievalConfig)
h.logger.Info("检索器配置已更新", h.logger.Info("检索器配置已更新",
zap.Int("top_k", retrievalConfig.TopK), zap.Int("top_k", retrievalConfig.TopK),
@@ -1705,6 +1715,13 @@ func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) {
setIntInMap(retrievalNode, "top_k", cfg.Retrieval.TopK) setIntInMap(retrievalNode, "top_k", cfg.Retrieval.TopK)
setFloatInMap(retrievalNode, "similarity_threshold", cfg.Retrieval.SimilarityThreshold) setFloatInMap(retrievalNode, "similarity_threshold", cfg.Retrieval.SimilarityThreshold)
setStringInMap(retrievalNode, "sub_index_filter", cfg.Retrieval.SubIndexFilter) setStringInMap(retrievalNode, "sub_index_filter", cfg.Retrieval.SubIndexFilter)
mqNode := ensureMap(retrievalNode, "multi_query")
setIntInMap(mqNode, "max_queries", cfg.Retrieval.MultiQuery.MaxQueries)
rerankNode := ensureMap(retrievalNode, "rerank")
setStringInMap(rerankNode, "provider", cfg.Retrieval.Rerank.Provider)
setStringInMap(rerankNode, "model", cfg.Retrieval.Rerank.Model)
setStringInMap(rerankNode, "base_url", cfg.Retrieval.Rerank.BaseURL)
setStringInMap(rerankNode, "api_key", cfg.Retrieval.Rerank.APIKey)
postNode := ensureMap(retrievalNode, "post_retrieve") postNode := ensureMap(retrievalNode, "post_retrieve")
setIntInMap(postNode, "prefetch_top_k", cfg.Retrieval.PostRetrieve.PrefetchTopK) setIntInMap(postNode, "prefetch_top_k", cfg.Retrieval.PostRetrieve.PrefetchTopK)
setIntInMap(postNode, "max_context_chars", cfg.Retrieval.PostRetrieve.MaxContextChars) setIntInMap(postNode, "max_context_chars", cfg.Retrieval.PostRetrieve.MaxContextChars)
@@ -1751,6 +1768,20 @@ func mergeHitlToolWhitelistSlice(existing, add []string) []string {
return out return out
} }
// SetHitlToolWhitelist 将全局免审批工具白名单整表写入 config.yaml(替换,非合并)。
func (h *ConfigHandler) SetHitlToolWhitelist(tools []string) error {
h.mu.Lock()
defer h.mu.Unlock()
h.config.Hitl.ToolWhitelist = mergeHitlToolWhitelistSlice(nil, tools)
if err := h.saveConfig(); err != nil {
return err
}
h.logger.Info("HITL 全局工具白名单已写入配置文件",
zap.Int("count", len(h.config.Hitl.ToolWhitelist)),
)
return nil
}
// MergeHitlToolWhitelistIntoConfig 将会话侧栏提交的免审批工具名合并进内存配置并写入 config.yaml(与全局白名单去重规则一致:小写键、保留首次出现的原始大小写)。 // MergeHitlToolWhitelistIntoConfig 将会话侧栏提交的免审批工具名合并进内存配置并写入 config.yaml(与全局白名单去重规则一致:小写键、保留首次出现的原始大小写)。
func (h *ConfigHandler) MergeHitlToolWhitelistIntoConfig(add []string) error { func (h *ConfigHandler) MergeHitlToolWhitelistIntoConfig(add []string) error {
h.mu.Lock() h.mu.Lock()
@@ -1771,6 +1802,21 @@ func updateHitlConfig(doc *yaml.Node, cfg config.HitlConfig) {
hitlNode := ensureMap(root, "hitl") hitlNode := ensureMap(root, "hitl")
// flow 样式 [a, b, c] 单行展示,工具多时比块序列省行数 // flow 样式 [a, b, c] 单行展示,工具多时比块序列省行数
setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist) setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist)
setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt)
setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit)
}
// UpdateHitlAuditAgentStrategy 更新审批/审查编辑两套审计 Agent 提示词并写入 config.yaml。
func (h *ConfigHandler) UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt string) error {
h.mu.Lock()
defer h.mu.Unlock()
h.config.Hitl.AuditAgentPrompt = strings.TrimSpace(approvalPrompt)
h.config.Hitl.AuditAgentPromptReviewEdit = strings.TrimSpace(reviewEditPrompt)
if err := h.saveConfig(); err != nil {
return err
}
h.logger.Info("HITL 审计 Agent 提示词已写入配置文件")
return nil
} }
func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) { func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) {
+6 -5
View File
@@ -103,6 +103,7 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
limitStr := c.DefaultQuery("limit", "50") limitStr := c.DefaultQuery("limit", "50")
offsetStr := c.DefaultQuery("offset", "0") offsetStr := c.DefaultQuery("offset", "0")
search := c.Query("search") // 获取搜索参数 search := c.Query("search") // 获取搜索参数
projectID := strings.TrimSpace(c.Query("project_id"))
limit, _ := strconv.Atoi(limitStr) limit, _ := strconv.Atoi(limitStr)
offset, _ := strconv.Atoi(offsetStr) offset, _ := strconv.Atoi(offsetStr)
@@ -114,7 +115,7 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
limit = 1000 limit = 1000
} }
excludeGrouped := strings.TrimSpace(search) == "" && excludeGrouped := strings.TrimSpace(search) == "" && projectID == "" &&
(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")) sortBy := strings.TrimSpace(c.Query("sort_by"))
@@ -122,14 +123,14 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
var total int var total int
var err error var err error
if excludeGrouped { if excludeGrouped {
conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy) conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy, projectID)
if err == nil { if err == nil {
total, err = h.db.CountUngroupedConversations() total, err = h.db.CountUngroupedConversations(projectID)
} }
} else { } else {
conversations, err = h.db.ListConversations(limit, offset, search, sortBy) conversations, err = h.db.ListConversations(limit, offset, search, sortBy, projectID)
if err == nil { if err == nil {
total, err = h.db.CountConversations(search) total, err = h.db.CountConversations(search, projectID)
} }
} }
if err != nil { if err != nil {
@@ -0,0 +1,83 @@
package handler
import (
"context"
"fmt"
"time"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/multiagent"
"go.uber.org/zap"
)
// rebindEinoRunningTask 中断并继续 / 空正文续跑:重建 cancel 链与超时 ctx,保持任务 running。
func (h *AgentHandler) rebindEinoRunningTask(conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) {
if timeoutCancel != nil {
timeoutCancel()
}
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, newTimeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
return baseCtx, cancelWithCause, taskCtx, newTimeoutCancel
}
// tryContinueOnEinoEmptyResponse Run 成功但 Response 为 emptyHint 时退避续跑;true 表示已准备下一段 Run。
func (h *AgentHandler) tryContinueOnEinoEmptyResponse(
taskCtx context.Context,
mw *config.MultiAgentEinoMiddlewareConfig,
conversationID string,
result *multiagent.RunResult,
attempt *int,
curHistory *[]agent.ChatMessage,
curFinalMessage *string,
progressCallback func(eventType, message string, data interface{}),
) bool {
if result == nil || !multiagent.IsEinoEmptyResponseResult(result) || !multiagent.HasEinoResumeTrace(result) {
return false
}
maxAttempts := multiagent.EmptyResponseContinueMaxAttemptsFromConfig(mw)
if *attempt >= maxAttempts {
if h.logger != nil {
h.logger.Warn("eino empty response continue exhausted",
zap.String("conversationId", conversationID),
zap.Int("maxAttempts", maxAttempts))
}
return false
}
*attempt++
h.persistEinoAgentTraceForResume(conversationID, result)
backoff := multiagent.EmptyResponseContinueBackoff(*attempt-1, mw)
waitMsg := fmt.Sprintf("会话已结束但未捕获到助手正文,%d 秒后第 %d/%d 次自动续跑…",
int(backoff.Seconds()), *attempt, maxAttempts)
if progressCallback != nil {
progressCallback("eino_empty_response_continue", waitMsg, map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": *attempt,
"maxAttempts": maxAttempts,
"backoffSec": int(backoff.Seconds()),
})
}
select {
case <-taskCtx.Done():
return false
case <-time.After(backoff):
}
inject := multiagent.FormatEmptyResponseContinueUserMessage()
h.applyEinoTraceResumeSegment(conversationID, result, curHistory, curFinalMessage, inject)
if progressCallback != nil {
progressCallback("eino_empty_response_continue", "已恢复上下文,正在续跑…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": *attempt,
"maxAttempts": maxAttempts,
"contextSource": "empty_response_continue",
})
}
return true
}
+8
View File
@@ -178,6 +178,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
var cumulativeMCPExecutionIDs []string var cumulativeMCPExecutionIDs []string
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。 // 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int var mainIterationOffset int
var emptyResponseContinueAttempt int
for { for {
segmentMainIterationMax := 0 segmentMainIterationMax := 0
@@ -239,6 +240,13 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
} }
if runErr == nil { if runErr == nil {
mw := &h.config.MultiAgent.EinoMiddleware
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
mainIterationOffset += segmentMainIterationMax
timeoutCancel()
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel)
continue
}
timeoutCancel() timeoutCancel()
break break
} }
+136 -89
View File
@@ -23,6 +23,7 @@ import (
type hitlRuntimeConfig struct { type hitlRuntimeConfig struct {
Enabled bool Enabled bool
Mode string Mode string
Reviewer string
SensitiveTools map[string]struct{} SensitiveTools map[string]struct{}
Timeout time.Duration Timeout time.Duration
} }
@@ -49,6 +50,8 @@ type HITLManager struct {
mu sync.RWMutex mu sync.RWMutex
runtime map[string]hitlRuntimeConfig runtime map[string]hitlRuntimeConfig
pending map[string]*pendingInterrupt pending map[string]*pendingInterrupt
// approvedExec 审批通过、待回写 tool_result 的队列(按会话 FIFO
approvedExec map[string][]hitlApprovedExecTrack
} }
func NewHITLManager(db *database.DB, logger *zap.Logger) *HITLManager { func NewHITLManager(db *database.DB, logger *zap.Logger) *HITLManager {
@@ -90,6 +93,7 @@ CREATE TABLE IF NOT EXISTS hitl_conversation_configs (
if err != nil { if err != nil {
return err return err
} }
m.migrateHitlSchemaColumns()
// On startup, cancel all orphaned pending interrupts from previous process. // On startup, cancel all orphaned pending interrupts from previous process.
// Their in-memory channels are gone, so they can never be resolved. // Their in-memory channels are gone, so they can never be resolved.
@@ -141,6 +145,7 @@ func (m *HITLManager) ActivateConversation(conversationID string, req *HITLReque
m.runtime[conversationID] = hitlRuntimeConfig{ m.runtime[conversationID] = hitlRuntimeConfig{
Enabled: true, Enabled: true,
Mode: normalizeHitlMode(req.Mode), Mode: normalizeHitlMode(req.Mode),
Reviewer: normalizeHitlReviewer(req.Reviewer),
SensitiveTools: tools, SensitiveTools: tools,
Timeout: timeout, Timeout: timeout,
} }
@@ -153,17 +158,14 @@ func (m *HITLManager) DeactivateConversation(conversationID string) {
m.mu.Unlock() m.mu.Unlock()
} }
// hitlConfigGlobalToolWhitelist 来自 config.yaml hitl.tool_whitelist(去重、去空)。 // hitlConfigGlobalToolWhitelist 来自 config.yaml hitl.tool_whitelist(去重、去空),并合并内置元工具免审批项
func (h *AgentHandler) hitlConfigGlobalToolWhitelist() []string { func (h *AgentHandler) hitlConfigGlobalToolWhitelist() []string {
if h == nil || h.config == nil { if h == nil || h.config == nil {
return nil return multiagent.MergeHitlExemptMetaTools(nil)
} }
raw := h.config.Hitl.ToolWhitelist raw := h.config.Hitl.ToolWhitelist
if len(raw) == 0 {
return nil
}
seen := make(map[string]struct{}) seen := make(map[string]struct{})
out := make([]string, 0, len(raw)) out := make([]string, 0, len(raw)+len(multiagent.HitlExemptMetaTools))
for _, t := range raw { for _, t := range raw {
n := strings.ToLower(strings.TrimSpace(t)) n := strings.ToLower(strings.TrimSpace(t))
if n == "" { if n == "" {
@@ -175,44 +177,35 @@ func (h *AgentHandler) hitlConfigGlobalToolWhitelist() []string {
seen[n] = struct{}{} seen[n] = struct{}{}
out = append(out, strings.TrimSpace(t)) out = append(out, strings.TrimSpace(t))
} }
return out return multiagent.MergeHitlExemptMetaTools(out)
} }
// hitlRequestWithMergedConfigWhitelist 将会话/API 中的白名单与 config.yaml 全局白名单合并(并集),仅用于运行时 Activate;不写入数据库。 // hitlRequestWithMergedConfigWhitelist 将会话/API 中的白名单与 config.yaml 全局白名单及内置元工具免审批项合并(并集),仅用于运行时 Activate;不写入数据库。
func (h *AgentHandler) hitlRequestWithMergedConfigWhitelist(req *HITLRequest) *HITLRequest { func (h *AgentHandler) hitlRequestWithMergedConfigWhitelist(req *HITLRequest) *HITLRequest {
gw := h.hitlConfigGlobalToolWhitelist()
if len(gw) == 0 {
return req
}
if req == nil { if req == nil {
return nil return nil
} }
seen := make(map[string]struct{}) seen := make(map[string]struct{})
union := make([]string, 0, len(gw)+len(req.SensitiveTools)) union := make([]string, 0, len(req.SensitiveTools)+16)
for _, t := range gw { add := func(t string) {
n := strings.ToLower(strings.TrimSpace(t)) n := strings.ToLower(strings.TrimSpace(t))
if n == "" { if n == "" {
continue return
} }
if _, ok := seen[n]; ok { if _, ok := seen[n]; ok {
continue return
} }
seen[n] = struct{}{} seen[n] = struct{}{}
union = append(union, strings.TrimSpace(t)) union = append(union, strings.TrimSpace(t))
} }
for _, t := range h.hitlConfigGlobalToolWhitelist() {
add(t)
}
for _, t := range req.SensitiveTools { for _, t := range req.SensitiveTools {
n := strings.ToLower(strings.TrimSpace(t)) add(t)
if n == "" {
continue
}
if _, ok := seen[n]; ok {
continue
}
seen[n] = struct{}{}
union = append(union, strings.TrimSpace(t))
} }
out := *req out := *req
out.SensitiveTools = union out.SensitiveTools = multiagent.MergeHitlExemptMetaTools(union)
return &out return &out
} }
@@ -362,22 +355,22 @@ func (m *HITLManager) SaveConversationConfig(conversationID string, req *HITLReq
timeout = 0 timeout = 0
} }
_, err := m.db.Exec(`INSERT INTO hitl_conversation_configs _, err := m.db.Exec(`INSERT INTO hitl_conversation_configs
(conversation_id, enabled, mode, sensitive_tools, timeout_seconds, updated_at) (conversation_id, enabled, mode, reviewer, sensitive_tools, timeout_seconds, updated_at)
VALUES (?, ?, ?, ?, ?, ?) VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(conversation_id) DO UPDATE SET ON CONFLICT(conversation_id) DO UPDATE SET
enabled=excluded.enabled, mode=excluded.mode, sensitive_tools=excluded.sensitive_tools, timeout_seconds=excluded.timeout_seconds, updated_at=excluded.updated_at`, enabled=excluded.enabled, mode=excluded.mode, reviewer=excluded.reviewer, sensitive_tools=excluded.sensitive_tools, timeout_seconds=excluded.timeout_seconds, updated_at=excluded.updated_at`,
conversationID, boolToInt(req.Enabled), mode, string(tools), timeout, time.Now()) conversationID, boolToInt(req.Enabled), mode, normalizeHitlReviewer(req.Reviewer), string(tools), timeout, time.Now())
return err return err
} }
func (m *HITLManager) LoadConversationConfig(conversationID string) (*HITLRequest, error) { func (m *HITLManager) LoadConversationConfig(conversationID string) (*HITLRequest, error) {
var enabledInt int var enabledInt int
var mode, toolsJSON string var mode, reviewer, toolsJSON string
var timeout int var timeout int
err := m.db.QueryRow(`SELECT enabled, mode, sensitive_tools, timeout_seconds FROM hitl_conversation_configs WHERE conversation_id = ?`, conversationID). err := m.db.QueryRow(`SELECT enabled, mode, COALESCE(reviewer,'human'), sensitive_tools, timeout_seconds FROM hitl_conversation_configs WHERE conversation_id = ?`, conversationID).
Scan(&enabledInt, &mode, &toolsJSON, &timeout) Scan(&enabledInt, &mode, &reviewer, &toolsJSON, &timeout)
if errors.Is(err, sql.ErrNoRows) { if errors.Is(err, sql.ErrNoRows) {
return &HITLRequest{Enabled: false, Mode: "off", SensitiveTools: []string{}, TimeoutSeconds: 0}, nil return &HITLRequest{Enabled: false, Mode: "off", Reviewer: "human", SensitiveTools: []string{}, TimeoutSeconds: 0}, nil
} }
if err != nil { if err != nil {
return nil, err return nil, err
@@ -390,6 +383,7 @@ func (m *HITLManager) LoadConversationConfig(conversationID string) (*HITLReques
return &HITLRequest{ return &HITLRequest{
Enabled: enabledInt == 1, Enabled: enabledInt == 1,
Mode: mode, Mode: mode,
Reviewer: normalizeHitlReviewer(reviewer),
SensitiveTools: tools, SensitiveTools: tools,
TimeoutSeconds: timeout, TimeoutSeconds: timeout,
}, nil }, nil
@@ -413,15 +407,16 @@ func (m *HITLManager) waitDecision(ctx context.Context, p *pendingInterrupt, tim
if p.Mode != "review_edit" && len(d.EditedArguments) > 0 { if p.Mode != "review_edit" && len(d.EditedArguments) > 0 {
d.EditedArguments = nil d.EditedArguments = nil
} }
_, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=? WHERE id=?`, _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=?, decided_by='human' WHERE id=?`,
d.Decision, d.Comment, time.Now(), p.InterruptID) d.Decision, d.Comment, time.Now(), p.InterruptID)
return d, nil return d, nil
case <-timeoutCh: case <-timeoutCh:
_, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='approve', decision_comment='timeout auto approve', decided_at=? WHERE id=?`, comment := "HITL timeout auto-reject for safety"
time.Now(), p.InterruptID) _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='reject', decision_comment=?, decided_at=?, decided_by='system' WHERE id=?`,
return hitlDecision{Decision: "approve", Comment: "timeout auto approve"}, nil comment, time.Now(), p.InterruptID)
return hitlDecision{Decision: "reject", Comment: comment}, nil
case <-ctx.Done(): case <-ctx.Done():
_, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', decision_comment='task cancelled', decided_at=? WHERE id=?`, _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', decision_comment='task cancelled', decided_at=?, decided_by='system' WHERE id=?`,
time.Now(), p.InterruptID) time.Now(), p.InterruptID)
return hitlDecision{Decision: "reject", Comment: "task cancelled"}, ctx.Err() return hitlDecision{Decision: "reject", Comment: "task cancelled"}, ctx.Err()
} }
@@ -445,12 +440,57 @@ func (h *AgentHandler) waitHITLApproval(runCtx context.Context, cancelRun contex
if !need { if !need {
return nil, nil return nil, nil
} }
h.enrichHitlApprovalPayload(conversationID, assistantMessageID, payload)
payloadRaw, _ := json.Marshal(payload) payloadRaw, _ := json.Marshal(payload)
p, err := h.hitlManager.CreatePendingInterrupt(conversationID, assistantMessageID, cfg.Mode, toolName, toolCallID, string(payloadRaw)) p, err := h.hitlManager.CreatePendingInterrupt(conversationID, assistantMessageID, cfg.Mode, toolName, toolCallID, string(payloadRaw))
if err != nil { if err != nil {
h.logger.Warn("创建 HITL 中断失败", zap.Error(err)) h.logger.Warn("创建 HITL 中断失败", zap.Error(err))
return nil, err return nil, err
} }
if cfg.Reviewer == "audit_agent" {
ad := h.auditAgentReview(runCtx, cfg.Mode, toolName, payload)
now := time.Now()
_, _ = h.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=?, decided_by='audit_agent' WHERE id=?`,
ad.Decision, ad.Comment, now, p.InterruptID)
if sendEventFunc != nil {
sendEventFunc("hitl_audit_agent", "审计 Agent 已裁决", map[string]interface{}{
"conversationId": conversationID,
"interruptId": p.InterruptID,
"toolName": toolName,
"mode": cfg.Mode,
"decision": ad.Decision,
"comment": ad.Comment,
"editedArgs": ad.EditedArguments,
"decidedBy": "audit_agent",
})
}
if ad.Decision == "reject" {
if sendEventFunc != nil {
sendEventFunc("hitl_rejected", "审计 Agent 拒绝本次工具调用", map[string]interface{}{
"conversationId": conversationID,
"interruptId": p.InterruptID,
"toolName": toolName,
"comment": ad.Comment,
"decidedBy": "audit_agent",
})
}
return &ad, nil
}
if sendEventFunc != nil {
sendEventFunc("hitl_resumed", "审计 Agent 已通过,继续执行", map[string]interface{}{
"conversationId": conversationID,
"interruptId": p.InterruptID,
"toolName": toolName,
"comment": ad.Comment,
"editedArgs": ad.EditedArguments,
"decidedBy": "audit_agent",
})
}
h.hitlManager.TrackApprovedHitlExecution(p.InterruptID, conversationID, toolName, toolCallID)
return &ad, nil
}
if sendEventFunc != nil { if sendEventFunc != nil {
sendEventFunc("hitl_interrupt", "命中人机协同审批", map[string]interface{}{ sendEventFunc("hitl_interrupt", "命中人机协同审批", map[string]interface{}{
"conversationId": conversationID, "conversationId": conversationID,
@@ -479,8 +519,12 @@ func (h *AgentHandler) waitHITLApproval(runCtx context.Context, cancelRun contex
return nil, waitErr return nil, waitErr
} }
if d.Decision == "reject" { if d.Decision == "reject" {
rejectMsg := "人工拒绝本次工具调用,模型将基于反馈继续迭代"
if strings.Contains(strings.ToLower(strings.TrimSpace(d.Comment)), "timeout") {
rejectMsg = "审批超时,安全起见已自动拒绝,模型将基于反馈继续迭代"
}
if sendEventFunc != nil { if sendEventFunc != nil {
sendEventFunc("hitl_rejected", "人工拒绝本次工具调用,模型将基于反馈继续迭代", map[string]interface{}{ sendEventFunc("hitl_rejected", rejectMsg, map[string]interface{}{
"conversationId": conversationID, "conversationId": conversationID,
"interruptId": p.InterruptID, "interruptId": p.InterruptID,
"toolName": toolName, "toolName": toolName,
@@ -498,6 +542,7 @@ func (h *AgentHandler) waitHITLApproval(runCtx context.Context, cancelRun contex
"editedArgs": d.EditedArguments, "editedArgs": d.EditedArguments,
}) })
} }
h.hitlManager.TrackApprovedHitlExecution(p.InterruptID, conversationID, toolName, toolCallID)
return &d, nil return &d, nil
} }
@@ -527,11 +572,6 @@ func (h *AgentHandler) handleHITLToolCall(runCtx context.Context, cancelRun cont
} }
func (h *AgentHandler) ListHITLPending(c *gin.Context) { func (h *AgentHandler) ListHITLPending(c *gin.Context) {
conversationID := strings.TrimSpace(c.Query("conversationId"))
status := strings.TrimSpace(c.Query("status"))
if status == "" {
status = "pending"
}
page, _ := strconv.Atoi(c.DefaultQuery("page", "1")) page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
if page < 1 { if page < 1 {
page = 1 page = 1
@@ -539,15 +579,12 @@ func (h *AgentHandler) ListHITLPending(c *gin.Context) {
pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20")) pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20"))
pageSize = int(math.Max(1, math.Min(float64(pageSize), 200))) pageSize = int(math.Max(1, math.Min(float64(pageSize), 200)))
offset := (page - 1) * pageSize offset := (page - 1) * pageSize
q := `SELECT id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, decision, decision_comment, created_at, decided_at FROM hitl_interrupts WHERE 1=1` q, args := h.buildHitlListQuery(false)
args := []interface{}{} q, args = h.appendHitlListFilters(q, args, c)
if conversationID != "" { total, err := h.countHitlQuery(q, args)
q += " AND conversation_id = ?" if err != nil {
args = append(args, conversationID) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
} return
if status != "all" {
q += " AND status = ?"
args = append(args, status)
} }
q += " ORDER BY created_at DESC LIMIT ? OFFSET ?" q += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
args = append(args, pageSize, offset) args = append(args, pageSize, offset)
@@ -557,41 +594,12 @@ func (h *AgentHandler) ListHITLPending(c *gin.Context) {
return return
} }
defer rows.Close() defer rows.Close()
items := make([]map[string]interface{}, 0) items, err := h.scanHitlInterruptRows(rows)
for rows.Next() { if err != nil {
var id, cid, mode, toolName, toolCallID, payload, rowStatus string c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
var messageID sql.NullString return
var decision, comment sql.NullString
var createdAt time.Time
var decidedAt sql.NullTime
if err := rows.Scan(&id, &cid, &messageID, &mode, &toolName, &toolCallID, &payload, &rowStatus, &decision, &comment, &createdAt, &decidedAt); err != nil {
continue
}
msgID := ""
if messageID.Valid {
msgID = messageID.String
}
items = append(items, map[string]interface{}{
"id": id,
"conversationId": cid,
"messageId": msgID,
"mode": mode,
"toolName": toolName,
"toolCallId": toolCallID,
"payload": payload,
"status": rowStatus,
"decision": decision.String,
"comment": comment.String,
"createdAt": createdAt,
"decidedAt": func() interface{} {
if decidedAt.Valid {
return decidedAt.Time
}
return nil
}(),
})
} }
c.JSON(http.StatusOK, gin.H{"items": items, "page": page, "pageSize": pageSize}) c.JSON(http.StatusOK, gin.H{"items": items, "page": page, "pageSize": pageSize, "total": total})
} }
type hitlDecisionReq struct { type hitlDecisionReq struct {
@@ -636,7 +644,7 @@ func (h *AgentHandler) DismissHITLInterrupt(c *gin.Context) {
return return
} }
res, err := h.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', res, err := h.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject',
decision_comment='dismissed by user', decided_at=CURRENT_TIMESTAMP decision_comment='dismissed by user', decided_at=CURRENT_TIMESTAMP, decided_by='human'
WHERE id=? AND status='pending'`, req.InterruptID) WHERE id=? AND status='pending'`, req.InterruptID)
if err != nil { if err != nil {
c.JSON(500, gin.H{"error": err.Error()}) c.JSON(500, gin.H{"error": err.Error()})
@@ -732,6 +740,7 @@ func (h *AgentHandler) UpsertHITLConversationConfig(c *gin.Context) {
return return
} }
req.Mode = normalizeHitlMode(req.Mode) req.Mode = normalizeHitlMode(req.Mode)
req.Reviewer = normalizeHitlReviewer(req.Reviewer)
if err := h.hitlManager.SaveConversationConfig(req.ConversationID, &req.HITLRequest); err != nil { if err := h.hitlManager.SaveConversationConfig(req.ConversationID, &req.HITLRequest); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return return
@@ -753,6 +762,44 @@ type mergeHitlGlobalWhitelistReq struct {
SensitiveTools []string `json:"sensitiveTools"` SensitiveTools []string `json:"sensitiveTools"`
} }
type setHitlGlobalWhitelistReq struct {
ToolWhitelist []string `json:"toolWhitelist"`
}
// GetHITLGlobalToolWhitelist 返回 config.yaml 中的全局免审批工具白名单。
func (h *AgentHandler) GetHITLGlobalToolWhitelist(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{
"toolWhitelist": h.hitlConfigGlobalToolWhitelist(),
})
}
// SetHITLGlobalToolWhitelist 整表替换 config.yaml 中的全局免审批工具白名单。
func (h *AgentHandler) SetHITLGlobalToolWhitelist(c *gin.Context) {
if h.hitlWhitelistSaver == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 配置持久化不可用"})
return
}
var req setHitlGlobalWhitelistReq
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if err := h.hitlWhitelistSaver.SetHitlToolWhitelist(req.ToolWhitelist); err != nil {
h.logger.Warn("写入 HITL 工具白名单到 config.yaml 失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "hitl", "tool_whitelist_update", "HITL 全局白名单更新", "hitl_config", "tool_whitelist", nil)
}
c.JSON(http.StatusOK, gin.H{
"ok": true,
"toolWhitelist": h.hitlConfigGlobalToolWhitelist(),
"hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(),
"hitlGlobalWhitelistMerged": false,
})
}
// MergeHITLGlobalToolWhitelist 无会话 ID 时将侧栏提交的免审批工具合并进 config.yaml(与 PUT /hitl/config 中白名单落盘规则一致)。 // MergeHITLGlobalToolWhitelist 无会话 ID 时将侧栏提交的免审批工具合并进 config.yaml(与 PUT /hitl/config 中白名单落盘规则一致)。
func (h *AgentHandler) MergeHITLGlobalToolWhitelist(c *gin.Context) { func (h *AgentHandler) MergeHITLGlobalToolWhitelist(c *gin.Context) {
if h.hitlWhitelistSaver == nil { if h.hitlWhitelistSaver == nil {
+357
View File
@@ -0,0 +1,357 @@
package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
"cyberstrike-ai/internal/config"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// auditAgentReview 在 reviewer=audit_agent 时由 LLM 代行审批。
// 白名单工具在 shouldInterrupt 阶段已跳过,到达此处的一律需要裁决。
func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName string, payload map[string]interface{}) hitlDecision {
if h == nil {
return hitlDecision{Decision: "reject", Comment: "audit agent: handler unavailable"}
}
mode := normalizeHitlMode(hitlMode)
prompt := config.DefaultHitlAuditAgentPrompt()
if h.config != nil {
prompt = h.config.Hitl.EffectiveAuditAgentPromptForMode(mode)
}
if h.auditLLM == nil {
return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 未配置"}
}
if ctx == nil {
ctx = context.Background()
}
callCtx, cancel := context.WithTimeout(ctx, 90*time.Second)
defer cancel()
userContent := buildAuditAgentReviewInput(mode, toolName, payload)
requestBody := map[string]interface{}{
"model": h.auditLLMModel(),
"messages": []map[string]interface{}{
{"role": "system", "content": prompt},
{"role": "user", "content": userContent},
},
"temperature": 0.1,
"max_completion_tokens": 1024,
// 审计裁决需要结构化 JSON;关闭 thinking 避免 Qwen 等把正文放进 reasoning_content 导致解析失败。
"thinking": map[string]interface{}{"type": "disabled"},
}
var apiResponse struct {
Choices []struct {
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
} `json:"message"`
} `json:"choices"`
}
if err := h.auditLLM.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil {
h.logger.Warn("审计 Agent LLM 调用失败", zap.Error(err), zap.String("tool", toolName))
return hitlDecision{
Decision: "reject",
Comment: "audit agent: LLM 调用失败,保守拒绝",
}
}
if len(apiResponse.Choices) == 0 {
return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 无有效响应,保守拒绝"}
}
msg := apiResponse.Choices[0].Message
raw := strings.TrimSpace(msg.Content)
if raw == "" {
raw = strings.TrimSpace(msg.ReasoningContent)
}
dec, err := parseAuditAgentLLMContent(raw)
if err != nil {
snippet := raw
if len(snippet) > 240 {
snippet = snippet[:240] + "..."
}
h.logger.Warn("审计 Agent 响应解析失败",
zap.Error(err),
zap.String("tool", toolName),
zap.String("mode", mode),
zap.String("snippet", snippet),
)
return hitlDecision{Decision: "reject", Comment: "audit agent: 响应无法解析,保守拒绝"}
}
if mode != "review_edit" && len(dec.EditedArguments) > 0 {
h.logger.Warn("审计 Agent 在审批模式下返回 editedArguments,已忽略",
zap.String("tool", toolName),
)
dec.EditedArguments = nil
}
if dec.Comment == "" {
dec.Comment = "audit agent: " + dec.Decision
} else if !strings.HasPrefix(strings.ToLower(dec.Comment), "audit agent") {
dec.Comment = "audit agent: " + dec.Comment
}
return dec
}
func (h *AgentHandler) auditLLMModel() string {
if h.config != nil && strings.TrimSpace(h.config.OpenAI.Model) != "" {
return strings.TrimSpace(h.config.OpenAI.Model)
}
return ""
}
func buildAuditAgentReviewInput(hitlMode, toolName string, payload map[string]interface{}) string {
review := map[string]interface{}{
"hitlMode": normalizeHitlMode(hitlMode),
"toolName": strings.TrimSpace(toolName),
}
if payload != nil {
for _, k := range []string{"arguments", "argumentsObj", "command", hitlPayloadUserMessage, hitlPayloadThinking, hitlPayloadReasoningChain, hitlPayloadPlanning} {
if v, ok := payload[k]; ok && v != nil && fmt.Sprint(v) != "" {
review[k] = v
}
}
}
b, err := json.MarshalIndent(review, "", " ")
if err != nil {
return fmt.Sprintf(`{"hitlMode":%q,"toolName":%q}`, normalizeHitlMode(hitlMode), toolName)
}
return string(b)
}
func parseAuditAgentLLMContent(content string) (hitlDecision, error) {
s := strings.TrimSpace(content)
if s == "" {
return hitlDecision{}, errors.New("empty content")
}
for _, candidate := range auditAgentJSONCandidates(s) {
dec, comment, editedArgs, err := parseAuditAgentDecisionObject(candidate)
if err == nil {
return hitlDecision{
Decision: dec,
Comment: comment,
EditedArguments: editedArgs,
}, nil
}
}
return hitlDecision{}, fmt.Errorf("no valid decision json in response")
}
func auditAgentJSONCandidates(s string) []string {
out := make([]string, 0, 4)
seen := make(map[string]struct{})
add := func(c string) {
c = strings.TrimSpace(c)
if c == "" {
return
}
if _, ok := seen[c]; ok {
return
}
seen[c] = struct{}{}
out = append(out, c)
}
add(s)
add(stripMarkdownCodeFence(s))
if obj := extractFirstJSONObject(s); obj != "" {
add(obj)
}
if obj := extractFirstJSONObject(stripMarkdownCodeFence(s)); obj != "" {
add(obj)
}
return out
}
func stripMarkdownCodeFence(s string) string {
s = strings.TrimSpace(s)
for _, fence := range []string{"```json", "```JSON", "```"} {
if strings.HasPrefix(s, fence) {
s = strings.TrimPrefix(s, fence)
}
}
s = strings.TrimSuffix(s, "```")
return strings.TrimSpace(s)
}
func extractFirstJSONObject(s string) string {
start := strings.Index(s, "{")
if start < 0 {
return ""
}
depth := 0
inStr := false
esc := false
for i := start; i < len(s); i++ {
ch := s[i]
if inStr {
if esc {
esc = false
continue
}
if ch == '\\' {
esc = true
continue
}
if ch == '"' {
inStr = false
}
continue
}
switch ch {
case '"':
inStr = true
case '{':
depth++
case '}':
depth--
if depth == 0 {
return s[start : i+1]
}
}
}
return ""
}
func parseAuditAgentDecisionObject(jsonText string) (decision, comment string, editedArgs map[string]interface{}, err error) {
var parsed map[string]interface{}
if err := json.Unmarshal([]byte(jsonText), &parsed); err != nil {
return "", "", nil, err
}
rawDecision := auditAgentPickString(parsed, "decision", "Decision", "result", "action", "verdict", "决策", "决定")
decision = normalizeAuditAgentDecision(rawDecision)
if decision == "" {
return "", "", nil, fmt.Errorf("missing decision")
}
comment = auditAgentPickString(parsed, "comment", "Comment", "reason", "message", "rationale", "备注", "理由", "说明")
editedArgs = auditAgentPickObject(parsed, "editedArguments", "edited_arguments", "editedArgs")
return decision, strings.TrimSpace(comment), editedArgs, nil
}
func auditAgentPickString(m map[string]interface{}, keys ...string) string {
for _, k := range keys {
if v, ok := m[k]; ok && v != nil {
s := strings.TrimSpace(fmt.Sprint(v))
if s != "" {
return s
}
}
}
return ""
}
func auditAgentPickObject(m map[string]interface{}, keys ...string) map[string]interface{} {
for _, k := range keys {
v, ok := m[k]
if !ok || v == nil {
continue
}
switch t := v.(type) {
case map[string]interface{}:
if len(t) > 0 {
return t
}
case string:
s := strings.TrimSpace(t)
if s == "" || s == "{}" {
continue
}
var obj map[string]interface{}
if err := json.Unmarshal([]byte(s), &obj); err == nil && len(obj) > 0 {
return obj
}
}
}
return nil
}
func normalizeAuditAgentDecision(v string) string {
d := strings.ToLower(strings.TrimSpace(v))
switch d {
case "approve", "approved", "pass", "passed", "allow", "allowed", "yes", "ok", "accept", "accepted":
return "approve"
case "reject", "rejected", "deny", "denied", "no", "block", "blocked", "refuse", "refused":
return "reject"
}
switch strings.TrimSpace(v) {
case "通过", "批准", "允许", "同意", "放行":
return "approve"
case "拒绝", "驳回", "禁止", "否决":
return "reject"
}
return ""
}
type hitlAuditStrategyReq struct {
AuditAgentPrompt string `json:"auditAgentPrompt"`
AuditAgentPromptReviewEdit string `json:"auditAgentPromptReviewEdit"`
}
func (h *AgentHandler) GetHITLAuditStrategy(c *gin.Context) {
approvalPrompt := config.DefaultHitlAuditAgentPrompt()
reviewEditPrompt := config.DefaultHitlAuditAgentPromptReviewEdit()
approvalCustom := false
reviewEditCustom := false
if h.config != nil {
approvalPrompt = h.config.Hitl.EffectiveAuditAgentPromptForMode("approval")
reviewEditPrompt = h.config.Hitl.EffectiveAuditAgentPromptForMode("review_edit")
approvalCustom = strings.TrimSpace(h.config.Hitl.AuditAgentPrompt) != ""
reviewEditCustom = strings.TrimSpace(h.config.Hitl.AuditAgentPromptReviewEdit) != ""
}
c.JSON(http.StatusOK, gin.H{
"auditAgentPrompt": approvalPrompt,
"auditAgentPromptCustom": approvalCustom,
"auditAgentPromptReviewEdit": reviewEditPrompt,
"auditAgentPromptReviewEditCustom": reviewEditCustom,
"defaultAuditAgentPrompt": config.DefaultHitlAuditAgentPrompt(),
"defaultAuditAgentPromptReviewEdit": config.DefaultHitlAuditAgentPromptReviewEdit(),
})
}
func (h *AgentHandler) UpdateHITLAuditStrategy(c *gin.Context) {
if h.hitlStrategySaver == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 策略持久化不可用"})
return
}
var req hitlAuditStrategyReq
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
approvalPrompt := strings.TrimSpace(req.AuditAgentPrompt)
reviewEditPrompt := strings.TrimSpace(req.AuditAgentPromptReviewEdit)
if err := h.hitlStrategySaver.UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt); err != nil {
h.logger.Warn("保存审计 Agent 提示词失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "hitl", "audit_strategy_update", "HITL 审计策略更新", "hitl_config", "audit_agent_prompt", nil)
}
if h.config != nil {
h.config.Hitl.AuditAgentPrompt = approvalPrompt
h.config.Hitl.AuditAgentPromptReviewEdit = reviewEditPrompt
}
c.JSON(http.StatusOK, gin.H{
"ok": true,
"auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"),
"auditAgentPromptCustom": approvalPrompt != "",
"auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"),
"auditAgentPromptReviewEditCustom": reviewEditPrompt != "",
})
}
// HitlAuditStrategySaver 持久化审计 Agent 提示词到 config.yaml。
type HitlAuditStrategySaver interface {
UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt string) error
}
// SetHitlAuditStrategySaver 设置审计策略落盘。
func (h *AgentHandler) SetHitlAuditStrategySaver(s HitlAuditStrategySaver) {
h.hitlStrategySaver = s
}
+88
View File
@@ -0,0 +1,88 @@
package handler
import (
"strings"
"testing"
)
func TestParseAuditAgentLLMContentApprove(t *testing.T) {
d, err := parseAuditAgentLLMContent(`{"decision":"approve","comment":"与任务一致"}`)
if err != nil {
t.Fatal(err)
}
if d.Decision != "approve" || d.Comment != "与任务一致" {
t.Fatalf("unexpected %+v", d)
}
}
func TestParseAuditAgentLLMContentReject(t *testing.T) {
d, err := parseAuditAgentLLMContent("```json\n{\"decision\":\"reject\",\"comment\":\"风险过高\"}\n```")
if err != nil {
t.Fatal(err)
}
if d.Decision != "reject" {
t.Fatalf("expected reject, got %s", d.Decision)
}
}
func TestParseAuditAgentLLMContentInvalid(t *testing.T) {
_, err := parseAuditAgentLLMContent(`{"decision":"maybe"}`)
if err == nil {
t.Fatal("expected error for invalid decision")
}
}
func TestParseAuditAgentLLMContentProseWrapped(t *testing.T) {
d, err := parseAuditAgentLLMContent("好的,裁决如下:\n```json\n{\"decision\":\"approve\",\"comment\":\"只读 ls\"}\n```\n以上。")
if err != nil {
t.Fatal(err)
}
if d.Decision != "approve" {
t.Fatalf("expected approve, got %s", d.Decision)
}
}
func TestParseAuditAgentLLMContentChineseDecision(t *testing.T) {
d, err := parseAuditAgentLLMContent(`{"decision":"通过","comment":"风险低"}`)
if err != nil {
t.Fatal(err)
}
if d.Decision != "approve" {
t.Fatalf("expected approve, got %s", d.Decision)
}
}
func TestParseAuditAgentLLMContentWithEditedArguments(t *testing.T) {
d, err := parseAuditAgentLLMContent(`{"decision":"approve","comment":"收窄路径","editedArguments":{"path":"/safe"}}`)
if err != nil {
t.Fatal(err)
}
if d.Decision != "approve" {
t.Fatalf("expected approve, got %s", d.Decision)
}
if d.EditedArguments == nil || d.EditedArguments["path"] != "/safe" {
t.Fatalf("unexpected edited args: %+v", d.EditedArguments)
}
}
func TestBuildAuditAgentReviewInputIncludesMode(t *testing.T) {
s := buildAuditAgentReviewInput("review_edit", "execute", map[string]interface{}{
"arguments": `{"command":"pwd"}`,
})
if !strings.Contains(s, "review_edit") || !strings.Contains(s, "execute") {
t.Fatalf("unexpected input: %s", s)
}
}
func TestBuildAuditAgentReviewInput(t *testing.T) {
s := buildAuditAgentReviewInput("approval", "nmap", map[string]interface{}{
"arguments": `{"target":"10.0.0.1"}`,
"userMessage": "扫描内网",
})
if s == "" {
t.Fatal("expected non-empty input")
}
if !strings.Contains(s, "nmap") || !strings.Contains(s, "10.0.0.1") || !strings.Contains(s, "扫描内网") {
t.Fatalf("unexpected input: %s", s)
}
}
+97
View File
@@ -0,0 +1,97 @@
package handler
import (
"strings"
)
type hitlCognitionState struct {
AssistantMessageID string
UserMessage string
Thinking string
ReasoningChain string
Planning string
}
// GetHitlCognition 返回当前运行任务上缓存的本轮 HITL 上下文(不含会话历史)。
func (m *AgentTaskManager) GetHitlCognition(conversationID string) hitlCognitionFields {
conversationID = strings.TrimSpace(conversationID)
if m == nil || conversationID == "" {
return hitlCognitionFields{}
}
m.mu.RLock()
defer m.mu.RUnlock()
t, ok := m.tasks[conversationID]
if !ok || t == nil || t.hitlCognition == nil {
return hitlCognitionFields{}
}
c := t.hitlCognition
return hitlCognitionFields{
UserMessage: c.UserMessage,
Thinking: c.Thinking,
ReasoningChain: c.ReasoningChain,
Planning: c.Planning,
}
}
// ResetHitlCognition 新任务开始时重置本轮 HITL 上下文。
func (m *AgentTaskManager) ResetHitlCognition(conversationID, userMessage string) {
conversationID = strings.TrimSpace(conversationID)
if m == nil || conversationID == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
t, ok := m.tasks[conversationID]
if !ok || t == nil {
return
}
t.hitlCognition = &hitlCognitionState{UserMessage: strings.TrimSpace(userMessage)}
}
// SetHitlAssistantMessageID 记录当前助手消息 ID,供 HITL 与 DB 回退对齐。
func (m *AgentTaskManager) SetHitlAssistantMessageID(conversationID, assistantMessageID string) {
conversationID = strings.TrimSpace(conversationID)
assistantMessageID = strings.TrimSpace(assistantMessageID)
if m == nil || conversationID == "" || assistantMessageID == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
t, ok := m.tasks[conversationID]
if !ok || t == nil {
return
}
if t.hitlCognition == nil {
t.hitlCognition = &hitlCognitionState{}
}
t.hitlCognition.AssistantMessageID = assistantMessageID
}
// UpdateHitlCognitionSnapshot 从进行中的进度流快照更新 thinking / reasoning / planning。
func (m *AgentTaskManager) UpdateHitlCognitionSnapshot(conversationID, assistantMessageID, thinking, reasoningChain, planning string) {
conversationID = strings.TrimSpace(conversationID)
if m == nil || conversationID == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
t, ok := m.tasks[conversationID]
if !ok || t == nil {
return
}
if t.hitlCognition == nil {
t.hitlCognition = &hitlCognitionState{}
}
if id := strings.TrimSpace(assistantMessageID); id != "" {
t.hitlCognition.AssistantMessageID = id
}
if s := strings.TrimSpace(thinking); s != "" {
t.hitlCognition.Thinking = s
}
if s := strings.TrimSpace(reasoningChain); s != "" {
t.hitlCognition.ReasoningChain = s
}
if s := strings.TrimSpace(planning); s != "" {
t.hitlCognition.Planning = s
}
}
+102
View File
@@ -0,0 +1,102 @@
package handler
import (
"strings"
)
const (
hitlPayloadUserMessage = "userMessage"
hitlPayloadThinking = "thinking"
hitlPayloadReasoningChain = "reasoningChain"
hitlPayloadPlanning = "planning"
)
type hitlCognitionFields struct {
UserMessage string
Thinking string
ReasoningChain string
Planning string
}
func (h *AgentHandler) enrichHitlApprovalPayload(conversationID, assistantMessageID string, payload map[string]interface{}) {
if h == nil || payload == nil {
return
}
cog := h.collectHitlCognition(conversationID, assistantMessageID)
if s := strings.TrimSpace(cog.UserMessage); s != "" {
payload[hitlPayloadUserMessage] = s
}
if s := strings.TrimSpace(cog.Thinking); s != "" {
payload[hitlPayloadThinking] = s
}
if s := strings.TrimSpace(cog.ReasoningChain); s != "" {
payload[hitlPayloadReasoningChain] = s
}
if s := strings.TrimSpace(cog.Planning); s != "" {
payload[hitlPayloadPlanning] = s
}
}
func (h *AgentHandler) collectHitlCognition(conversationID, assistantMessageID string) hitlCognitionFields {
var out hitlCognitionFields
if h.tasks != nil {
out = h.tasks.GetHitlCognition(conversationID)
}
if strings.TrimSpace(out.UserMessage) == "" && h.db != nil {
if msg, err := h.db.GetTurnUserMessage(conversationID, assistantMessageID); err == nil {
out.UserMessage = msg
}
}
if h.db != nil && assistantMessageID != "" {
dbCog, err := h.db.GetAssistantCognitionTexts(assistantMessageID)
if err == nil {
if strings.TrimSpace(out.Thinking) == "" {
out.Thinking = dbCog.Thinking
}
if strings.TrimSpace(out.ReasoningChain) == "" {
out.ReasoningChain = dbCog.ReasoningChain
}
if strings.TrimSpace(out.Planning) == "" {
out.Planning = dbCog.Planning
}
}
}
return out
}
func snapshotHitlCognitionFromStreams(thinkingStreams map[string]*thinkingBuf, respPlan *responsePlanAgg) (thinking, reasoningChain, planning string) {
if len(thinkingStreams) > 0 {
var thinkingParts, reasoningParts []string
for _, tb := range thinkingStreams {
if tb == nil {
continue
}
content := strings.TrimSpace(tb.b.String())
if content == "" {
continue
}
if tb.persistAs == "reasoning_chain" {
reasoningParts = append(reasoningParts, content)
} else {
thinkingParts = append(thinkingParts, content)
}
}
thinking = strings.Join(thinkingParts, "\n\n")
reasoningChain = strings.Join(reasoningParts, "\n\n")
}
if respPlan != nil {
planning = strings.TrimSpace(respPlan.b.String())
}
return thinking, reasoningChain, planning
}
func (h *AgentHandler) syncHitlCognitionFromProgress(conversationID, assistantMessageID string, thinkingStreams map[string]*thinkingBuf, respPlan *responsePlanAgg) {
if h == nil || h.tasks == nil {
return
}
thinking, reasoning, planning := snapshotHitlCognitionFromStreams(thinkingStreams, respPlan)
if thinking == "" && reasoning == "" && planning == "" {
return
}
h.tasks.UpdateHitlCognitionSnapshot(conversationID, assistantMessageID, thinking, reasoning, planning)
}
+46
View File
@@ -0,0 +1,46 @@
package handler
import (
"os"
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestEnrichHitlApprovalPayload(t *testing.T) {
tmp := t.TempDir()
db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop())
if err != nil {
t.Fatalf("db: %v", err)
}
defer os.RemoveAll(tmp)
conv, err := db.CreateConversation("hitl ctx", database.ConversationCreateMeta{})
if err != nil {
t.Fatalf("conv: %v", err)
}
if _, err := db.AddMessage(conv.ID, "user", "scan 10.0.0.1 please", nil); err != nil {
t.Fatalf("user msg: %v", err)
}
asst, err := db.AddMessage(conv.ID, "assistant", "", nil)
if err != nil {
t.Fatalf("asst msg: %v", err)
}
if err := db.AddProcessDetail(asst.ID, conv.ID, "thinking", "need port scan first", nil); err != nil {
t.Fatalf("detail: %v", err)
}
h := &AgentHandler{db: db, tasks: NewAgentTaskManager()}
payload := map[string]interface{}{"toolName": "nmap", "arguments": "{}"}
h.enrichHitlApprovalPayload(conv.ID, asst.ID, payload)
if got := payload["userMessage"]; got != "scan 10.0.0.1 please" {
t.Fatalf("userMessage=%v", got)
}
if got := payload["thinking"]; got != "need port scan first" {
t.Fatalf("thinking=%v", got)
}
}
+132
View File
@@ -0,0 +1,132 @@
package handler
import (
"encoding/json"
"strings"
"time"
)
const hitlPayloadExecutionResult = "executionResult"
type hitlExecutionResult struct {
Success bool `json:"success"`
Result string `json:"result,omitempty"`
ToolName string `json:"toolName,omitempty"`
ToolCallID string `json:"toolCallId,omitempty"`
RecordedAt time.Time `json:"recordedAt"`
}
type hitlApprovedExecTrack struct {
InterruptID string
ConversationID string
ToolName string
ToolCallID string
}
// TrackApprovedHitlExecution 审批通过后登记,待 tool_result 回写执行结果。
func (m *HITLManager) TrackApprovedHitlExecution(interruptID, conversationID, toolName, toolCallID string) {
if m == nil {
return
}
interruptID = strings.TrimSpace(interruptID)
conversationID = strings.TrimSpace(conversationID)
if interruptID == "" || conversationID == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
if m.approvedExec == nil {
m.approvedExec = make(map[string][]hitlApprovedExecTrack)
}
m.approvedExec[conversationID] = append(m.approvedExec[conversationID], hitlApprovedExecTrack{
InterruptID: interruptID,
ConversationID: conversationID,
ToolName: strings.TrimSpace(toolName),
ToolCallID: strings.TrimSpace(toolCallID),
})
}
func (m *HITLManager) popApprovedInterruptForTool(conversationID, toolCallID, toolName string) string {
if m == nil {
return ""
}
conversationID = strings.TrimSpace(conversationID)
toolCallID = strings.TrimSpace(toolCallID)
toolName = strings.TrimSpace(toolName)
m.mu.Lock()
defer m.mu.Unlock()
queue := m.approvedExec[conversationID]
if len(queue) == 0 {
return ""
}
idx := -1
if toolCallID != "" {
for i, t := range queue {
if t.ToolCallID == toolCallID {
idx = i
break
}
}
}
if idx < 0 && toolName != "" {
for i, t := range queue {
if strings.EqualFold(t.ToolName, toolName) {
idx = i
break
}
}
}
if idx < 0 {
return ""
}
id := queue[idx].InterruptID
queue = append(queue[:idx], queue[idx+1:]...)
if len(queue) == 0 {
delete(m.approvedExec, conversationID)
} else {
m.approvedExec[conversationID] = queue
}
return id
}
func mergeHitlPayloadExecutionResult(payloadJSON string, exec hitlExecutionResult) (string, error) {
root := make(map[string]interface{})
if strings.TrimSpace(payloadJSON) != "" {
_ = json.Unmarshal([]byte(payloadJSON), &root)
}
if root == nil {
root = make(map[string]interface{})
}
root[hitlPayloadExecutionResult] = exec
out, err := json.Marshal(root)
if err != nil {
return payloadJSON, err
}
return string(out), nil
}
func (h *AgentHandler) recordHitlToolExecutionResult(conversationID, toolCallID, toolName string, success bool, result string) {
if h == nil || h.hitlManager == nil || h.db == nil {
return
}
interruptID := h.hitlManager.popApprovedInterruptForTool(conversationID, toolCallID, toolName)
if interruptID == "" {
return
}
var payloadJSON string
err := h.db.QueryRow(`SELECT payload FROM hitl_interrupts WHERE id = ?`, interruptID).Scan(&payloadJSON)
if err != nil {
return
}
merged, err := mergeHitlPayloadExecutionResult(payloadJSON, hitlExecutionResult{
Success: success,
Result: strings.TrimSpace(result),
ToolName: strings.TrimSpace(toolName),
ToolCallID: strings.TrimSpace(toolCallID),
RecordedAt: time.Now(),
})
if err != nil {
return
}
_, _ = h.db.Exec(`UPDATE hitl_interrupts SET payload = ? WHERE id = ?`, merged, interruptID)
}
+39
View File
@@ -0,0 +1,39 @@
package handler
import (
"encoding/json"
"testing"
)
func TestMergeHitlPayloadExecutionResult(t *testing.T) {
merged, err := mergeHitlPayloadExecutionResult(`{"userMessage":"hi","toolName":"nmap"}`, hitlExecutionResult{
Success: true,
Result: "open ports: 80",
})
if err != nil {
t.Fatal(err)
}
var root map[string]interface{}
if err := json.Unmarshal([]byte(merged), &root); err != nil {
t.Fatal(err)
}
if root["userMessage"] != "hi" {
t.Fatalf("userMessage lost: %v", root["userMessage"])
}
exec, ok := root["executionResult"].(map[string]interface{})
if !ok || exec["success"] != true {
t.Fatalf("executionResult missing: %v", root["executionResult"])
}
}
func TestPopApprovedInterruptForTool(t *testing.T) {
m := NewHITLManager(nil, nil)
m.TrackApprovedHitlExecution("hitl_a", "conv1", "nmap", "tc1")
m.TrackApprovedHitlExecution("hitl_b", "conv1", "exec", "")
if id := m.popApprovedInterruptForTool("conv1", "tc1", "nmap"); id != "hitl_a" {
t.Fatalf("tc1 match=%q", id)
}
if id := m.popApprovedInterruptForTool("conv1", "", "exec"); id != "hitl_b" {
t.Fatalf("tool name match=%q", id)
}
}
+263
View File
@@ -0,0 +1,263 @@
package handler
import (
"database/sql"
"errors"
"math"
"net/http"
"strconv"
"strings"
"time"
"cyberstrike-ai/internal/config"
"github.com/gin-gonic/gin"
)
func normalizeHitlReviewer(v string) string {
switch strings.ToLower(strings.TrimSpace(v)) {
case "audit_agent", "agent", "ai":
return "audit_agent"
default:
return "human"
}
}
func normalizeHitlDecidedBy(v string) string {
switch strings.ToLower(strings.TrimSpace(v)) {
case "audit_agent", "agent", "ai":
return "audit_agent"
case "system", "timeout":
return "system"
case "manual":
return "manual"
default:
return "human"
}
}
func (m *HITLManager) migrateHitlSchemaColumns() {
_, _ = m.db.Exec(`ALTER TABLE hitl_interrupts ADD COLUMN decided_by TEXT NOT NULL DEFAULT 'human'`)
_, _ = m.db.Exec(`ALTER TABLE hitl_conversation_configs ADD COLUMN reviewer TEXT NOT NULL DEFAULT 'human'`)
}
func hitlInterruptRowToMap(
id, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy string,
messageID sql.NullString,
decision, comment sql.NullString,
createdAt time.Time,
decidedAt sql.NullTime,
) map[string]interface{} {
msgID := ""
if messageID.Valid {
msgID = messageID.String
}
return map[string]interface{}{
"id": id,
"conversationId": cid,
"messageId": msgID,
"mode": mode,
"toolName": toolName,
"toolCallId": toolCallID,
"payload": payload,
"status": rowStatus,
"decision": decision.String,
"comment": comment.String,
"decidedBy": decidedBy,
"createdAt": createdAt,
"decidedAt": func() interface{} {
if decidedAt.Valid {
return decidedAt.Time
}
return nil
}(),
}
}
func (h *AgentHandler) buildHitlListQuery(logs bool) (string, []interface{}) {
where, args := h.buildHitlLogsWhere(logs)
q := `SELECT id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, decision, decision_comment, COALESCE(decided_by,'human'), created_at, decided_at FROM hitl_interrupts` + where
return q, args
}
func (h *AgentHandler) buildHitlLogsWhere(logs bool) (string, []interface{}) {
q := " WHERE 1=1"
args := []interface{}{}
if logs {
q += " AND status != 'pending'"
} else {
q += " AND status = 'pending'"
}
return q, args
}
func (h *AgentHandler) appendHitlListFilters(q string, args []interface{}, c *gin.Context) (string, []interface{}) {
conversationID := strings.TrimSpace(c.Query("conversationId"))
toolName := strings.TrimSpace(c.Query("toolName"))
decision := strings.TrimSpace(c.Query("decision"))
decidedBy := strings.TrimSpace(c.Query("decidedBy"))
status := strings.TrimSpace(c.Query("status"))
search := strings.TrimSpace(c.Query("q"))
if conversationID != "" {
q += " AND conversation_id = ?"
args = append(args, conversationID)
}
if toolName != "" {
q += " AND tool_name LIKE ?"
args = append(args, "%"+toolName+"%")
}
if decision != "" && decision != "all" {
q += " AND decision = ?"
args = append(args, decision)
}
if decidedBy != "" && decidedBy != "all" {
q += " AND COALESCE(decided_by,'human') = ?"
args = append(args, normalizeHitlDecidedBy(decidedBy))
}
if status != "" && status != "all" {
q += " AND status = ?"
args = append(args, status)
}
if search != "" {
like := "%" + search + "%"
q += " AND (id LIKE ? OR conversation_id LIKE ? OR tool_name LIKE ? OR payload LIKE ? OR COALESCE(decision_comment,'') LIKE ?)"
args = append(args, like, like, like, like, like)
}
return q, args
}
func (h *AgentHandler) scanHitlInterruptRows(rows *sql.Rows) ([]map[string]interface{}, error) {
items := make([]map[string]interface{}, 0)
for rows.Next() {
var id, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy string
var messageID sql.NullString
var decision, comment sql.NullString
var createdAt time.Time
var decidedAt sql.NullTime
if err := rows.Scan(&id, &cid, &messageID, &mode, &toolName, &toolCallID, &payload, &rowStatus, &decision, &comment, &decidedBy, &createdAt, &decidedAt); err != nil {
continue
}
items = append(items, hitlInterruptRowToMap(id, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy, messageID, decision, comment, createdAt, decidedAt))
}
return items, nil
}
func (h *AgentHandler) countHitlQuery(baseQ string, args []interface{}) (int, error) {
countQ := "SELECT COUNT(*) FROM (" + baseQ + ") AS hitl_cnt"
var total int
if err := h.db.QueryRow(countQ, args...).Scan(&total); err != nil {
return 0, err
}
return total, nil
}
func (h *AgentHandler) ListHITLLogs(c *gin.Context) {
page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
if page < 1 {
page = 1
}
pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20"))
pageSize = int(math.Max(1, math.Min(float64(pageSize), 200)))
offset := (page - 1) * pageSize
q, args := h.buildHitlListQuery(true)
q, args = h.appendHitlListFilters(q, args, c)
total, err := h.countHitlQuery(q, args)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
q += " ORDER BY COALESCE(decided_at, created_at) DESC LIMIT ? OFFSET ?"
args = append(args, pageSize, offset)
rows, err := h.db.Query(q, args...)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
defer rows.Close()
items, err := h.scanHitlInterruptRows(rows)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"items": items, "page": page, "pageSize": pageSize, "total": total, "retentionDays": h.hitlRetentionDays()})
}
func (h *AgentHandler) hitlRetentionDays() int {
if h.config != nil {
return h.config.Hitl.RetentionDaysEffective()
}
return config.HitlConfig{}.RetentionDaysEffective()
}
// DeleteHITLLogs 批量删除或按筛选清空已决策的人机协同审计日志(不删除 pending)。
func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) {
var request struct {
IDs []string `json:"ids"`
All bool `json:"all"`
}
if err := c.ShouldBindJSON(&request); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数无效: " + err.Error()})
return
}
var deleted int64
var err error
if request.All {
where, args := h.buildHitlLogsWhere(true)
where, args = h.appendHitlListFilters(where, args, c)
deleted, err = h.db.DeleteHitlInterruptLogsMatching(where, args)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "hitl", "logs_clear", "清空人机协同审计日志", "hitl_interrupt", "", map[string]interface{}{
"deleted": deleted,
})
}
} else {
if len(request.IDs) == 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "审计日志 ID 列表不能为空"})
return
}
deleted, err = h.db.DeleteHitlInterruptLogsByIDs(request.IDs)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "hitl", "logs_delete_batch", "批量删除人机协同审计日志", "hitl_interrupt", "", map[string]interface{}{
"count": len(request.IDs),
"deleted": deleted,
})
}
}
c.JSON(http.StatusOK, gin.H{"message": "删除成功", "deleted": deleted})
}
func (h *AgentHandler) GetHITLLog(c *gin.Context) {
id := strings.TrimSpace(c.Param("id"))
if id == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"})
return
}
q := `SELECT id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, decision, decision_comment, COALESCE(decided_by,'human'), created_at, decided_at FROM hitl_interrupts WHERE id = ?`
var rowID, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy string
var messageID sql.NullString
var decision, comment sql.NullString
var createdAt time.Time
var decidedAt sql.NullTime
err := h.db.QueryRow(q, id).Scan(&rowID, &cid, &messageID, &mode, &toolName, &toolCallID, &payload, &rowStatus, &decision, &comment, &decidedBy, &createdAt, &decidedAt)
if errors.Is(err, sql.ErrNoRows) {
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
return
}
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, hitlInterruptRowToMap(rowID, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy, messageID, decision, comment, createdAt, decidedAt))
}
+221 -15
View File
@@ -5,6 +5,7 @@ import (
"errors" "errors"
"io" "io"
"net/http" "net/http"
"sort"
"strconv" "strconv"
"strings" "strings"
"time" "time"
@@ -68,16 +69,34 @@ func (h *MonitorHandler) SetAgentHandler(ah *AgentHandler) {
h.agentHandler = ah h.agentHandler = ah
} }
const monitorPageTopTools = 6
// MonitorStatsSummary 工具调用汇总
type MonitorStatsSummary struct {
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
ToolCount int `json:"toolCount"`
}
// MonitorResponse 监控响应 // MonitorResponse 监控响应
type MonitorResponse struct { type MonitorResponse struct {
Executions []*mcp.ToolExecution `json:"executions"` Executions []*mcp.ToolExecution `json:"executions"`
Stats map[string]*mcp.ToolStats `json:"stats"` Summary *MonitorStatsSummary `json:"summary"`
Timestamp time.Time `json:"timestamp"` TopTools []*mcp.ToolStats `json:"topTools"`
Total int `json:"total,omitempty"` Timestamp time.Time `json:"timestamp"`
Page int `json:"page,omitempty"` Total int `json:"total"`
PageSize int `json:"page_size,omitempty"` Page int `json:"page"`
TotalPages int `json:"total_pages,omitempty"` PageSize int `json:"pageSize"`
RetentionDays int `json:"retention_days,omitempty"` TotalPages int `json:"totalPages"`
RetentionDays int `json:"retentionDays"`
}
// StatsResponse 统计信息响应(Dashboard 等)
type StatsResponse struct {
Summary *MonitorStatsSummary `json:"summary"`
TopTools []*mcp.ToolStats `json:"topTools"`
} }
// Monitor 获取监控信息 // Monitor 获取监控信息
@@ -101,9 +120,9 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
// 解析工具筛选参数(兼容 mcp__tool 与内部 mcp::tool // 解析工具筛选参数(兼容 mcp__tool 与内部 mcp::tool
toolName := normalizeToolNameFilter(c.Query("tool")) toolName := normalizeToolNameFilter(c.Query("tool"))
executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName) executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName)
h.enrichExecutionsConversationID(executions) h.enrichExecutionsConversationID(executions)
stats := h.loadStats() summary, topTools := h.loadStatsSummary(monitorPageTopTools)
totalPages := (total + pageSize - 1) / pageSize totalPages := (total + pageSize - 1) / pageSize
if totalPages == 0 { if totalPages == 0 {
@@ -112,7 +131,8 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
c.JSON(http.StatusOK, MonitorResponse{ c.JSON(http.StatusOK, MonitorResponse{
Executions: executions, Executions: executions,
Stats: stats, Summary: summary,
TopTools: topTools,
Timestamp: time.Now(), Timestamp: time.Now(),
Total: total, Total: total,
Page: page, Page: page,
@@ -134,6 +154,112 @@ func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
return executions return executions
} }
func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
if h.db == nil {
allExecutions := h.mcpServer.GetAllExecutions()
if status != "" || toolName != "" {
filtered := make([]*mcp.ToolExecution, 0)
for _, exec := range allExecutions {
matchStatus := status == "" || exec.Status == status
matchTool := toolNameFilterMatches(exec.ToolName, toolName)
if matchStatus && matchTool {
filtered = append(filtered, exec)
}
}
allExecutions = filtered
}
total := len(allExecutions)
offset := (page - 1) * pageSize
end := offset + pageSize
if end > total {
end = total
}
if offset >= total {
return []*mcp.ToolExecution{}, total
}
pageSlice := allExecutions[offset:end]
out := make([]*mcp.ToolExecution, 0, len(pageSlice))
for _, exec := range pageSlice {
if exec == nil {
continue
}
out = append(out, slimToolExecution(exec))
}
return out, total
}
offset := (page - 1) * pageSize
executions, err := h.db.LoadToolExecutionListPage(offset, pageSize, status, toolName)
if err != nil {
h.logger.Warn("从数据库加载执行记录列表失败,回退到内存数据", zap.Error(err))
return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName)
}
total, err := h.db.CountToolExecutions(status, toolName)
if err != nil {
h.logger.Warn("获取执行记录总数失败", zap.Error(err))
total = offset + len(executions)
if len(executions) == pageSize {
total = offset + len(executions) + 1
}
}
return executions, total
}
func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
allExecutions := h.mcpServer.GetAllExecutions()
if status != "" || toolName != "" {
filtered := make([]*mcp.ToolExecution, 0)
for _, exec := range allExecutions {
matchStatus := status == "" || exec.Status == status
matchTool := toolNameFilterMatches(exec.ToolName, toolName)
if matchStatus && matchTool {
filtered = append(filtered, exec)
}
}
allExecutions = filtered
}
total := len(allExecutions)
offset := (page - 1) * pageSize
end := offset + pageSize
if end > total {
end = total
}
if offset >= total {
return []*mcp.ToolExecution{}, total
}
pageSlice := allExecutions[offset:end]
out := make([]*mcp.ToolExecution, 0, len(pageSlice))
for _, exec := range pageSlice {
if exec == nil {
continue
}
out = append(out, slimToolExecution(exec))
}
return out, total
}
func slimToolExecution(exec *mcp.ToolExecution) *mcp.ToolExecution {
if exec == nil {
return nil
}
slim := &mcp.ToolExecution{
ID: exec.ID,
ToolName: exec.ToolName,
Status: exec.Status,
StartTime: exec.StartTime,
}
if exec.EndTime != nil {
end := *exec.EndTime
slim.EndTime = &end
}
if exec.Duration > 0 {
slim.Duration = exec.Duration
}
return slim
}
func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) { func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
if h.db == nil { if h.db == nil {
allExecutions := h.mcpServer.GetAllExecutions() allExecutions := h.mcpServer.GetAllExecutions()
@@ -206,7 +332,78 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status
return executions, total return executions, total
} }
func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats { func (h *MonitorHandler) loadStatsSummary(topN int) (*MonitorStatsSummary, []*mcp.ToolStats) {
if topN <= 0 {
topN = monitorPageTopTools
}
if h.db != nil {
result, err := h.db.LoadToolStatsSummary(topN)
if err == nil {
return dbStatsSummaryToMonitor(result), result.TopTools
}
h.logger.Warn("从数据库加载统计汇总失败,回退到内存数据", zap.Error(err))
}
stats := h.loadStatsMap()
return summarizeToolStats(stats, topN)
}
func dbStatsSummaryToMonitor(result *database.ToolStatsSummaryResult) *MonitorStatsSummary {
if result == nil {
return &MonitorStatsSummary{}
}
summary := &MonitorStatsSummary{
TotalCalls: result.Summary.TotalCalls,
SuccessCalls: result.Summary.SuccessCalls,
FailedCalls: result.Summary.FailedCalls,
ToolCount: result.Summary.ToolCount,
}
if result.Summary.LastCallTime != nil {
t := *result.Summary.LastCallTime
summary.LastCallTime = &t
}
return summary
}
func summarizeToolStats(stats map[string]*mcp.ToolStats, topN int) (*MonitorStatsSummary, []*mcp.ToolStats) {
summary := &MonitorStatsSummary{}
if len(stats) == 0 {
return summary, nil
}
all := make([]*mcp.ToolStats, 0, len(stats))
for _, stat := range stats {
if stat == nil {
continue
}
summary.ToolCount++
summary.TotalCalls += stat.TotalCalls
summary.SuccessCalls += stat.SuccessCalls
summary.FailedCalls += stat.FailedCalls
if stat.LastCallTime != nil && (summary.LastCallTime == nil || stat.LastCallTime.After(*summary.LastCallTime)) {
t := *stat.LastCallTime
summary.LastCallTime = &t
}
if stat.TotalCalls > 0 {
statCopy := *stat
all = append(all, &statCopy)
}
}
sort.Slice(all, func(i, j int) bool {
if all[i].TotalCalls == all[j].TotalCalls {
return all[i].ToolName < all[j].ToolName
}
return all[i].TotalCalls > all[j].TotalCalls
})
if len(all) > topN {
all = all[:topN]
}
return summary, all
}
func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats {
// 合并内部MCP服务器和外部MCP管理器的统计信息 // 合并内部MCP服务器和外部MCP管理器的统计信息
stats := make(map[string]*mcp.ToolStats) stats := make(map[string]*mcp.ToolStats)
@@ -334,7 +531,7 @@ func (h *MonitorHandler) CancelExecution(c *gin.Context) {
func (h *MonitorHandler) enrichExecutionsConversationID(executions []*mcp.ToolExecution) { func (h *MonitorHandler) enrichExecutionsConversationID(executions []*mcp.ToolExecution) {
for _, exec := range executions { for _, exec := range executions {
if exec == nil { if exec == nil || exec.Status != "running" {
continue continue
} }
exec.ConversationID = h.conversationIDForRunningExecution(exec.ID) exec.ConversationID = h.conversationIDForRunningExecution(exec.ID)
@@ -415,8 +612,17 @@ func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) {
// GetStats 获取统计信息 // GetStats 获取统计信息
func (h *MonitorHandler) GetStats(c *gin.Context) { func (h *MonitorHandler) GetStats(c *gin.Context) {
stats := h.loadStats() topN := 30
c.JSON(http.StatusOK, stats) if topStr := c.Query("top"); topStr != "" {
if t, err := strconv.Atoi(topStr); err == nil && t > 0 && t <= 100 {
topN = t
}
}
summary, topTools := h.loadStatsSummary(topN)
c.JSON(http.StatusOK, StatsResponse{
Summary: summary,
TopTools: topTools,
})
} }
// CallsTimelinePoint 调用趋势数据点 // CallsTimelinePoint 调用趋势数据点
+8
View File
@@ -188,6 +188,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
var cumulativeMCPExecutionIDs []string var cumulativeMCPExecutionIDs []string
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。 // 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int var mainIterationOffset int
var emptyResponseContinueAttempt int
for { for {
segmentMainIterationMax := 0 segmentMainIterationMax := 0
@@ -251,6 +252,13 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
} }
if runErr == nil { if runErr == nil {
mw := &h.config.MultiAgent.EinoMiddleware
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
mainIterationOffset += segmentMainIterationMax
timeoutCancel()
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel)
continue
}
timeoutCancel() timeoutCancel()
break break
} }
+45 -6
View File
@@ -740,14 +740,21 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"properties": map[string]interface{}{ "properties": map[string]interface{}{
"executions": map[string]interface{}{ "executions": map[string]interface{}{
"type": "array", "type": "array",
"description": "执行记录列表", "description": "执行记录列表(轻量字段,不含 arguments/result",
"items": map[string]interface{}{ "items": map[string]interface{}{
"$ref": "#/components/schemas/ToolExecution", "$ref": "#/components/schemas/ToolExecution",
}, },
}, },
"stats": map[string]interface{}{ "summary": map[string]interface{}{
"type": "object", "type": "object",
"description": "统计信息", "description": "工具调用汇总",
},
"topTools": map[string]interface{}{
"type": "array",
"description": "调用量 Top N 工具",
"items": map[string]interface{}{
"type": "object",
},
}, },
"timestamp": map[string]interface{}{ "timestamp": map[string]interface{}{
"type": "string", "type": "string",
@@ -756,20 +763,24 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
}, },
"total": map[string]interface{}{ "total": map[string]interface{}{
"type": "integer", "type": "integer",
"description": "总数", "description": "执行记录总数",
}, },
"page": map[string]interface{}{ "page": map[string]interface{}{
"type": "integer", "type": "integer",
"description": "当前页", "description": "当前页",
}, },
"page_size": map[string]interface{}{ "pageSize": map[string]interface{}{
"type": "integer", "type": "integer",
"description": "每页数量", "description": "每页数量",
}, },
"total_pages": map[string]interface{}{ "totalPages": map[string]interface{}{
"type": "integer", "type": "integer",
"description": "总页数", "description": "总页数",
}, },
"retentionDays": map[string]interface{}{
"type": "integer",
"description": "执行记录保留天数",
},
}, },
}, },
"ConfigResponse": map[string]interface{}{ "ConfigResponse": map[string]interface{}{
@@ -1232,6 +1243,34 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"type": "string", "type": "string",
}, },
}, },
{
"name": "project_id",
"in": "query",
"required": false,
"description": "按项目筛选;传 __none__ 表示仅未绑定项目的对话",
"schema": map[string]interface{}{
"type": "string",
},
},
{
"name": "exclude_grouped",
"in": "query",
"required": false,
"description": "为 true 时排除已加入分组的对话(默认在未搜索且未按项目筛选时启用)",
"schema": map[string]interface{}{
"type": "boolean",
},
},
{
"name": "sort_by",
"in": "query",
"required": false,
"description": "排序字段:updated_at(默认)或 created_at",
"schema": map[string]interface{}{
"type": "string",
"enum": []string{"updated_at", "created_at"},
},
},
}, },
"responses": map[string]interface{}{ "responses": map[string]interface{}{
"200": map[string]interface{}{ "200": map[string]interface{}{
+1
View File
@@ -8,6 +8,7 @@ import (
) )
// agentSessionContextBlock 注入会话工作目录与项目黑板(用于 system prompt 追加块)。 // agentSessionContextBlock 注入会话工作目录与项目黑板(用于 system prompt 追加块)。
// 用户输入由 message history 承载;压缩后由 summarization 摘要指令保留关键约束。
func (h *AgentHandler) agentSessionContextBlock(conversationID string) string { func (h *AgentHandler) agentSessionContextBlock(conversationID string) string {
var parts []string var parts []string
if ws := h.buildWorkspaceBlock(conversationID); ws != "" { if ws := h.buildWorkspaceBlock(conversationID); ws != "" {
+46 -23
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()
} }
@@ -711,12 +711,27 @@ type wecomReplyXML struct {
Content string `xml:"Content"` Content string `xml:"Content"`
} }
// wecomRequireToken 企业微信回调必须配置 Token;未配置时拒绝请求,防止未授权触发 Agent。
func (h *RobotHandler) wecomRequireToken(c *gin.Context) (string, bool) {
token := strings.TrimSpace(h.config.Robots.Wecom.Token)
if token == "" {
h.logger.Warn("企业微信已启用但未配置 token,已拒绝回调(请在配置中设置 robots.wecom.token")
c.String(http.StatusForbidden, "")
return "", false
}
return token, true
}
// HandleWecomGET 企业微信 URL 校验(GET // HandleWecomGET 企业微信 URL 校验(GET
func (h *RobotHandler) HandleWecomGET(c *gin.Context) { func (h *RobotHandler) HandleWecomGET(c *gin.Context) {
if !h.config.Robots.Wecom.Enabled { if !h.config.Robots.Wecom.Enabled {
c.String(http.StatusNotFound, "") c.String(http.StatusNotFound, "")
return return
} }
token, ok := h.wecomRequireToken(c)
if !ok {
return
}
// Gin 的 Query() 会自动 URL 解码,拿到的就是正确的 base64 字符串 // Gin 的 Query() 会自动 URL 解码,拿到的就是正确的 base64 字符串
echostr := c.Query("echostr") echostr := c.Query("echostr")
msgSignature := c.Query("msg_signature") msgSignature := c.Query("msg_signature")
@@ -724,7 +739,7 @@ func (h *RobotHandler) HandleWecomGET(c *gin.Context) {
nonce := c.Query("nonce") nonce := c.Query("nonce")
// 验证签名:将 token、timestamp、nonce、echostr 四个参数排序后拼接计算 SHA1 // 验证签名:将 token、timestamp、nonce、echostr 四个参数排序后拼接计算 SHA1
signature := h.signWecomRequest(h.config.Robots.Wecom.Token, timestamp, nonce, echostr) signature := h.signWecomRequest(token, timestamp, nonce, echostr)
if signature != msgSignature { if signature != msgSignature {
h.logger.Warn("企业微信 URL 验证签名失败", zap.String("expected", msgSignature), zap.String("got", signature)) h.logger.Warn("企业微信 URL 验证签名失败", zap.String("expected", msgSignature), zap.String("got", signature))
c.String(http.StatusBadRequest, "invalid signature") c.String(http.StatusBadRequest, "invalid signature")
@@ -865,27 +880,28 @@ func (h *RobotHandler) HandleWecomPOST(c *gin.Context) {
} }
h.logger.Debug("企业微信 POST 收到请求", zap.String("body", string(bodyRaw))) h.logger.Debug("企业微信 POST 收到请求", zap.String("body", string(bodyRaw)))
// 验证请求签名防止伪造。企业微信签名算法同 URL 验证,使用 token、timestamp、nonce、 Encrypt 四个字段 // 验证请求签名防止伪造。企业微信签名算法同 URL 验证,使用 token、timestamp、nonce、 Encrypt 四个字段
// 若配置了 Token 则必须校验签名,避免未授权请求触发 Agent(防止平台被接管) // 启用企业微信时必须配置 token 校验签名,避免未授权请求触发 Agent
token := h.config.Robots.Wecom.Token token, ok := h.wecomRequireToken(c)
if token != "" { if !ok {
if msgSignature == "" { return
h.logger.Warn("企业微信 POST 缺少签名,已拒绝(需配置 token 并确保回调携带 msg_signature") }
c.String(http.StatusOK, "") if msgSignature == "" {
return h.logger.Warn("企业微信 POST 缺少签名,已拒绝(需确保回调携带 msg_signature")
} c.String(http.StatusOK, "")
var tmp wecomXML return
if err := xml.Unmarshal(bodyRaw, &tmp); err != nil { }
h.logger.Warn("企业微信 POST 签名验证前解析 XML 失败", zap.Error(err)) var tmp wecomXML
c.String(http.StatusOK, "") if err := xml.Unmarshal(bodyRaw, &tmp); err != nil {
return h.logger.Warn("企业微信 POST 签名验证前解析 XML 失败", zap.Error(err))
} c.String(http.StatusOK, "")
expected := h.signWecomRequest(token, timestamp, nonce, tmp.Encrypt) return
if expected != msgSignature { }
h.logger.Warn("企业微信 POST 签名验证失败", zap.String("expected", expected), zap.String("got", msgSignature)) expected := h.signWecomRequest(token, timestamp, nonce, tmp.Encrypt)
c.String(http.StatusOK, "") if expected != msgSignature {
return h.logger.Warn("企业微信 POST 签名验证失败", zap.String("expected", expected), zap.String("got", msgSignature))
} c.String(http.StatusOK, "")
return
} }
var body wecomXML var body wecomXML
@@ -899,6 +915,13 @@ func (h *RobotHandler) HandleWecomPOST(c *gin.Context) {
// 保存企业 ID(用于明文模式回复) // 保存企业 ID(用于明文模式回复)
enterpriseID := body.ToUserName enterpriseID := body.ToUserName
// 配置了 EncodingAESKey 时必须走加密消息,拒绝明文 XML 绕过
if strings.TrimSpace(h.config.Robots.Wecom.EncodingAESKey) != "" && strings.TrimSpace(body.Encrypt) == "" {
h.logger.Warn("企业微信已配置加密模式但收到明文消息,已拒绝")
c.String(http.StatusOK, "")
return
}
// 加密模式:先解密再解析内层 XML // 加密模式:先解密再解析内层 XML
if body.Encrypt != "" && h.config.Robots.Wecom.EncodingAESKey != "" { if body.Encrypt != "" && h.config.Robots.Wecom.EncodingAESKey != "" {
h.logger.Debug("企业微信进入加密模式解密流程") h.logger.Debug("企业微信进入加密模式解密流程")
+78
View File
@@ -0,0 +1,78 @@
package handler
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func newWecomTestHandler(token string, aesKey string) *RobotHandler {
return &RobotHandler{
config: &config.Config{
Robots: config.RobotsConfig{
Wecom: config.RobotWecomConfig{
Enabled: true,
Token: token,
EncodingAESKey: aesKey,
},
},
},
logger: zap.NewNop(),
}
}
func TestHandleWecomPOST_rejectsWhenTokenEmpty(t *testing.T) {
gin.SetMode(gin.TestMode)
h := newWecomTestHandler("", "")
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
body := `<?xml version="1.0"?><xml><FromUserName>attacker</FromUserName><MsgType>text</MsgType><Content>hi</Content></xml>`
c.Request = httptest.NewRequest(http.MethodPost, "/api/robot/wecom", strings.NewReader(body))
h.HandleWecomPOST(c)
if w.Code != http.StatusForbidden {
t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden)
}
if w.Body.String() == "success" {
t.Fatal("expected rejection, got success")
}
}
func TestHandleWecomPOST_rejectsPlaintextWhenEncryptionConfigured(t *testing.T) {
gin.SetMode(gin.TestMode)
h := newWecomTestHandler("secret-token", "abcdefghijklmnopqrstuvwxyz0123456789ABCD")
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
body := `<?xml version="1.0"?><xml><FromUserName>attacker</FromUserName><MsgType>text</MsgType><Content>hi</Content></xml>`
c.Request = httptest.NewRequest(http.MethodPost, "/api/robot/wecom?timestamp=1&nonce=2&msg_signature=fake", strings.NewReader(body))
h.HandleWecomPOST(c)
if w.Body.String() == "success" {
t.Fatal("expected rejection for plaintext in encryption mode, got success")
}
}
func TestHandleWecomGET_rejectsWhenTokenEmpty(t *testing.T) {
gin.SetMode(gin.TestMode)
h := newWecomTestHandler("", "")
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodGet, "/api/robot/wecom?msg_signature=x&timestamp=1&nonce=2&echostr=abc", nil)
h.HandleWecomGET(c)
if w.Code != http.StatusForbidden {
t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden)
}
}
+24 -2
View File
@@ -26,6 +26,7 @@ func shouldPersistEinoAgentTraceAfterRunError(baseCtx context.Context) bool {
// AgentTask 描述正在运行的Agent任务 // AgentTask 描述正在运行的Agent任务
type AgentTask struct { type AgentTask struct {
ConversationID string `json:"conversationId"` ConversationID string `json:"conversationId"`
Title string `json:"title,omitempty"`
Message string `json:"message,omitempty"` Message string `json:"message,omitempty"`
StartedAt time.Time `json:"startedAt"` StartedAt time.Time `json:"startedAt"`
Status string `json:"status"` Status string `json:"status"`
@@ -42,6 +43,9 @@ type AgentTask struct {
// activeEinoExecuteAbortNote AbortActiveEinoExecute 写入的用户说明,由 execute 收尾时合并进工具结果 // activeEinoExecuteAbortNote AbortActiveEinoExecute 写入的用户说明,由 execute 收尾时合并进工具结果
activeEinoExecuteAbortNote string activeEinoExecuteAbortNote string
// hitlCognition 本轮运行中供 HITL/审计 Agent 读取的上下文(用户原话 + 思考,不含会话历史)
hitlCognition *hitlCognitionState
cancel func(error) cancel func(error)
} }
@@ -233,6 +237,7 @@ func (m *AgentTaskManager) ActiveMCPExecutionID(conversationID string) string {
// CompletedTask 已完成的任务(用于历史记录) // CompletedTask 已完成的任务(用于历史记录)
type CompletedTask struct { type CompletedTask struct {
ConversationID string `json:"conversationId"` ConversationID string `json:"conversationId"`
Title string `json:"title,omitempty"`
Message string `json:"message,omitempty"` Message string `json:"message,omitempty"`
StartedAt time.Time `json:"startedAt"` StartedAt time.Time `json:"startedAt"`
CompletedAt time.Time `json:"completedAt"` CompletedAt time.Time `json:"completedAt"`
@@ -247,6 +252,8 @@ type AgentTaskManager struct {
maxHistorySize int // 最大历史记录数 maxHistorySize int // 最大历史记录数
historyRetention time.Duration // 历史记录保留时间 historyRetention time.Duration // 历史记录保留时间
eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅 eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅
// toolCanceler 在用户整轮停止任务时终止当前 MCP 工具(非「中断并继续」)。
toolCanceler func(conversationID string)
} }
const ( const (
@@ -277,6 +284,13 @@ func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) {
m.eventBus = b m.eventBus = b
} }
// SetToolCanceler 设置整轮停止任务时终止当前 MCP 工具的回调(由 AgentHandler 注入)。
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
m.mu.Lock()
defer m.mu.Unlock()
m.toolCanceler = fn
}
// GetTask 返回运行中任务(无则 nil)。 // GetTask 返回运行中任务(无则 nil)。
func (m *AgentTaskManager) GetTask(conversationID string) *AgentTask { func (m *AgentTaskManager) GetTask(conversationID string) *AgentTask {
m.mu.RLock() m.mu.RLock()
@@ -343,6 +357,7 @@ func (m *AgentTaskManager) StartTask(conversationID, message string, cancel cont
} }
m.tasks[conversationID] = task m.tasks[conversationID] = task
task.hitlCognition = &hitlCognitionState{UserMessage: strings.TrimSpace(message)}
return task, nil return task, nil
} }
@@ -372,14 +387,21 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool,
task.InterruptContinueNote = "" task.InterruptContinueNote = ""
} }
cancel := task.cancel cancel := task.cancel
m.mu.Unlock()
if cause == nil { if cause == nil {
cause = ErrTaskCancelled cause = ErrTaskCancelled
} }
var toolCanceler func(string)
if errors.Is(cause, ErrTaskCancelled) {
toolCanceler = m.toolCanceler
}
m.mu.Unlock()
if cancel != nil { if cancel != nil {
cancel(cause) cancel(cause)
} }
if toolCanceler != nil {
toolCanceler(conversationID)
}
return true, nil return true, nil
} }
@@ -0,0 +1,80 @@
package handler
import (
"context"
"errors"
"testing"
"cyberstrike-ai/internal/multiagent"
)
func TestCancelTaskInvokesToolCancelerOnFullStop(t *testing.T) {
tm := NewAgentTaskManager()
called := false
tm.SetToolCanceler(func(conversationID string) {
if conversationID == "conv-1" {
called = true
}
})
_, cancel := context.WithCancelCause(context.Background())
_, err := tm.StartTask("conv-1", "hello", cancel)
if err != nil {
t.Fatalf("StartTask: %v", err)
}
ok, err := tm.CancelTask("conv-1", ErrTaskCancelled)
if err != nil || !ok {
t.Fatalf("CancelTask: ok=%v err=%v", ok, err)
}
if !called {
t.Fatal("expected tool canceler to be invoked on full task cancel")
}
}
func TestCancelTaskSkipsToolCancelerOnInterruptContinue(t *testing.T) {
tm := NewAgentTaskManager()
called := false
tm.SetToolCanceler(func(conversationID string) {
called = true
})
_, cancel := context.WithCancelCause(context.Background())
_, err := tm.StartTask("conv-1", "hello", cancel)
if err != nil {
t.Fatalf("StartTask: %v", err)
}
ok, err := tm.CancelTask("conv-1", multiagent.ErrInterruptContinue)
if err != nil || !ok {
t.Fatalf("CancelTask: ok=%v err=%v", ok, err)
}
if called {
t.Fatal("tool canceler must not run for interrupt-continue")
}
}
func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) {
tm := NewAgentTaskManager()
var gotCause error
tm.SetToolCanceler(func(conversationID string) {
if conversationID == "conv-2" {
gotCause = ErrTaskCancelled
}
})
ctx, cancel := context.WithCancelCause(context.Background())
if _, err := tm.StartTask("conv-2", "hello", cancel); err != nil {
t.Fatalf("StartTask: %v", err)
}
if _, err := tm.CancelTask("conv-2", nil); err != nil {
t.Fatalf("CancelTask: %v", err)
}
if !errors.Is(context.Cause(ctx), ErrTaskCancelled) {
t.Fatalf("expected ErrTaskCancelled cause, got %v", context.Cause(ctx))
}
if gotCause != ErrTaskCancelled {
t.Fatalf("expected tool canceler path for default cancel cause")
}
}
+16
View File
@@ -0,0 +1,16 @@
//go:build windows
package handler
import (
"net/http"
"github.com/gin-gonic/gin"
)
// RunCommandWS 交互式 PTY 终端依赖 Unix PTY(见 terminal_ws_unix.go);Windows 暂不支持。
func (h *TerminalHandler) RunCommandWS(c *gin.Context) {
c.JSON(http.StatusNotImplemented, gin.H{
"error": "Interactive WebSocket terminal is not supported on Windows; use POST /terminal/run or /terminal/run/stream instead.",
})
}
+71
View File
@@ -0,0 +1,71 @@
package hitl
import (
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
const retentionPurgeInterval = time.Hour
// Service manages HITL audit log retention (decided hitl_interrupts rows).
type Service struct {
db *database.DB
cfg *config.Config
logger *zap.Logger
}
// NewService creates a HITL audit log retention service.
func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service {
return &Service{db: db, cfg: cfg, logger: logger}
}
// RetentionDays returns configured retention; 0 means keep forever.
func (s *Service) RetentionDays() int {
if s == nil || s.cfg == nil {
return config.HitlConfig{}.RetentionDaysEffective()
}
return s.cfg.Hitl.RetentionDaysEffective()
}
// PurgeExpired deletes decided HITL log rows older than retention_days when configured.
func (s *Service) PurgeExpired() {
if s == nil || s.db == nil || s.cfg == nil {
return
}
days := s.cfg.Hitl.RetentionDaysEffective()
if days <= 0 {
return
}
cutoff := time.Now().AddDate(0, 0, -days)
n, err := s.db.PurgeHitlInterruptLogsBefore(cutoff)
if err != nil {
if s.logger != nil {
s.logger.Warn("清理过期人机协同审计日志失败", zap.Error(err))
}
return
}
if n > 0 && s.logger != nil {
s.logger.Info("已清理过期人机协同审计日志", zap.Int64("deleted", n), zap.Int("retention_days", days))
}
}
// StartRetentionLoop periodically purges expired HITL audit log rows.
func StartRetentionLoop(s *Service, logger *zap.Logger) {
if s == nil {
return
}
go func() {
ticker := time.NewTicker(retentionPurgeInterval)
defer ticker.Stop()
for range ticker.C {
s.PurgeExpired()
if logger != nil {
logger.Debug("hitl audit log retention tick completed")
}
}
}()
}
+50
View File
@@ -0,0 +1,50 @@
package hitl
import (
"path/filepath"
"testing"
"time"
appconfig "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestServicePurgeExpired_respectsZeroRetention(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "hitl.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS hitl_interrupts (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL,
mode TEXT NOT NULL,
tool_name TEXT NOT NULL,
status TEXT NOT NULL,
decision TEXT,
created_at DATETIME NOT NULL,
decided_at DATETIME
)`); err != nil {
t.Fatalf("create table: %v", err)
}
old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339)
if _, err := db.Exec(`INSERT INTO hitl_interrupts
(id, conversation_id, mode, tool_name, status, decision, created_at, decided_at)
VALUES ('old-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, old, old); err != nil {
t.Fatalf("insert: %v", err)
}
zero := 0
svc := NewService(db, &appconfig.Config{
Hitl: appconfig.HitlConfig{RetentionDays: &zero},
}, zap.NewNop())
svc.PurgeExpired()
if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err != nil {
t.Fatalf("record should remain when retention_days=0: %v", err)
}
}
@@ -0,0 +1,96 @@
package knowledge
import (
"context"
"fmt"
"strings"
"cyberstrike-ai/internal/config"
"github.com/cloudwego/eino/callbacks"
"github.com/cloudwego/eino/components"
"github.com/cloudwego/eino/components/retriever"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// knowledgePipelineRetriever: MultiQuery → vector candidates → rerank → post-process.
type knowledgePipelineRetriever struct {
inner retriever.Retriever
base *Retriever
}
func newKnowledgePipelineRetriever(inner retriever.Retriever, base *Retriever) *knowledgePipelineRetriever {
if inner == nil || base == nil {
return nil
}
return &knowledgePipelineRetriever{inner: inner, base: base}
}
func (p *knowledgePipelineRetriever) GetType() string {
return "KnowledgeRAGPipeline"
}
func (p *knowledgePipelineRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (out []*schema.Document, err error) {
if p == nil || p.inner == nil || p.base == nil {
return nil, fmt.Errorf("knowledge pipeline retriever: nil")
}
q := strings.TrimSpace(query)
if q == "" {
return nil, fmt.Errorf("查询不能为空")
}
ro := retriever.GetCommonOptions(nil, opts...)
finalTopK := p.base.config.TopK
if finalTopK <= 0 {
finalTopK = 5
}
if ro.TopK != nil && *ro.TopK > 0 {
finalTopK = *ro.TopK
}
ctx = callbacks.EnsureRunInfo(ctx, p.GetType(), components.ComponentOfRetriever)
ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{Query: q, TopK: finalTopK, Extra: ro.DSLInfo})
defer func() {
if err != nil {
_ = callbacks.OnError(ctx, err)
return
}
_ = callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: out})
}()
out, err = p.inner.Retrieve(ctx, q, opts...)
if err != nil {
return nil, err
}
if len(out) == 0 {
return out, nil
}
if rr := p.base.documentReranker(); rr != nil && len(out) > 1 {
reranked, rerr := rr.Rerank(ctx, q, out)
if rerr != nil {
if p.base.logger != nil {
p.base.logger.Warn("知识检索重排失败,已使用融合序", zap.Error(rerr))
}
} else if len(reranked) > 0 {
out = reranked
}
}
tokenModel := ""
if p.base.embedder != nil {
tokenModel = p.base.embedder.EmbeddingModelName()
}
var postPO *config.PostRetrieveConfig
if p.base.config != nil {
postPO = &p.base.config.PostRetrieve
}
out, err = ApplyPostRetrieve(out, postPO, tokenModel, finalTopK)
if err != nil {
return nil, err
}
return out, nil
}
var _ retriever.Retriever = (*knowledgePipelineRetriever)(nil)
+1 -2
View File
@@ -8,8 +8,7 @@ import (
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
) )
// BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino Chain,底层为 SQLite 向量检索([VectorEinoRetriever])。 // BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino ChainMultiQuery → 向量 → 重排 → 后处理)。
// 去重、上下文预算截断与最终 Top-K 均在 [VectorEinoRetriever.Retrieve] 内完成,与 HTTP/MCP 检索路径一致。
func BuildKnowledgeRetrieveChain(ctx context.Context, r *Retriever) (compose.Runnable[string, []*schema.Document], error) { func BuildKnowledgeRetrieveChain(ctx context.Context, r *Retriever) (compose.Runnable[string, []*schema.Document], error) {
if r == nil { if r == nil {
return nil, fmt.Errorf("retriever is nil") return nil, fmt.Errorf("retriever is nil")
+1 -30
View File
@@ -11,19 +11,10 @@ import (
"github.com/cloudwego/eino/components" "github.com/cloudwego/eino/components"
"github.com/cloudwego/eino/components/retriever" "github.com/cloudwego/eino/components/retriever"
"github.com/cloudwego/eino/schema" "github.com/cloudwego/eino/schema"
"go.uber.org/zap"
) )
// VectorEinoRetriever implements [retriever.Retriever] on top of SQLite-stored embeddings + cosine similarity. // VectorEinoRetriever implements [retriever.Retriever] on top of SQLite-stored embeddings + cosine similarity.
// // It returns prefetch-sized vector candidates only; rerank and post-process run in [knowledgePipelineRetriever].
// Options:
// - [retriever.WithTopK]
// - [retriever.WithDSLInfo] with [DSLRiskType] (string), [DSLSimilarityThreshold] (float, cosine 01), [DSLSubIndexFilter] (string)
//
// Document scores are cosine similarity; [retriever.WithScoreThreshold] is not mapped to a different metric.
//
// After vector search: optional [DocumentReranker] (see [Retriever.SetDocumentReranker]), then
// [ApplyPostRetrieve] (normalized-text dedupe, context budget, final Top-K) using [config.PostRetrieveConfig].
type VectorEinoRetriever struct { type VectorEinoRetriever struct {
inner *Retriever inner *Retriever
} }
@@ -119,26 +110,6 @@ func (h *VectorEinoRetriever) Retrieve(ctx context.Context, query string, opts .
return nil, err return nil, err
} }
out = retrievalResultsToDocuments(results) out = retrievalResultsToDocuments(results)
if rr := h.inner.documentReranker(); rr != nil && len(out) > 1 {
reranked, rerr := rr.Rerank(ctx, q, out)
if rerr != nil {
if h.inner.logger != nil {
h.inner.logger.Warn("知识检索重排失败,已使用向量序", zap.Error(rerr))
}
} else if len(reranked) > 0 {
out = reranked
}
}
tokenModel := ""
if h.inner.embedder != nil {
tokenModel = h.inner.embedder.EmbeddingModelName()
}
out, err = ApplyPostRetrieve(out, postPO, tokenModel, finalTopK)
if err != nil {
return nil, err
}
return out, nil return out, nil
} }
+226
View File
@@ -0,0 +1,226 @@
package knowledge
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"cyberstrike-ai/internal/config"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// HTTPReranker calls a hosted rerank API (DashScope or Cohere-compatible).
type HTTPReranker struct {
provider string
model string
baseURL string
apiKey string
client *http.Client
logger *zap.Logger
}
// NewHTTPReranker builds a rerank client from knowledge retrieval config; openAI supplies fallback credentials.
func NewHTTPReranker(rc *config.RerankConfig, openAI *config.OpenAIConfig, logger *zap.Logger) (*HTTPReranker, error) {
if rc == nil {
return nil, fmt.Errorf("rerank config is nil")
}
baseURL := strings.TrimSpace(rc.BaseURL)
apiKey := strings.TrimSpace(rc.APIKey)
if openAI != nil {
if baseURL == "" {
baseURL = strings.TrimSpace(openAI.BaseURL)
}
if apiKey == "" {
apiKey = strings.TrimSpace(openAI.APIKey)
}
}
if apiKey == "" {
return nil, fmt.Errorf("rerank api_key is required")
}
provider := rc.ProviderEffective(baseURL)
model := rc.ModelEffective(provider)
return &HTTPReranker{
provider: provider,
model: model,
baseURL: strings.TrimSuffix(baseURL, "/"),
apiKey: apiKey,
client: &http.Client{Timeout: 60 * time.Second},
logger: logger,
}, nil
}
func (r *HTTPReranker) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) {
if r == nil {
return docs, nil
}
q := strings.TrimSpace(query)
if q == "" || len(docs) == 0 {
return docs, nil
}
if len(docs) == 1 {
return docs, nil
}
texts := make([]string, 0, len(docs))
for _, d := range docs {
if d == nil {
texts = append(texts, "")
continue
}
texts = append(texts, d.Content)
}
var order []int
var err error
switch r.provider {
case "dashscope":
order, err = r.rerankDashScope(ctx, q, texts, len(docs))
default:
order, err = r.rerankCohere(ctx, q, texts, len(docs))
}
if err != nil {
return nil, err
}
out := make([]*schema.Document, 0, len(order))
for _, idx := range order {
if idx < 0 || idx >= len(docs) || docs[idx] == nil {
continue
}
out = append(out, docs[idx])
}
if len(out) == 0 {
return docs, nil
}
return out, nil
}
func (r *HTTPReranker) rerankCohere(ctx context.Context, query string, documents []string, topN int) ([]int, error) {
url := r.cohereRerankURL()
body := map[string]any{
"model": r.model,
"query": query,
"documents": documents,
"top_n": topN,
}
raw, err := json.Marshal(body)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+r.apiKey)
resp, err := r.client.Do(req)
if err != nil {
return nil, fmt.Errorf("rerank request: %w", err)
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody)))
}
var parsed struct {
Results []struct {
Index int `json:"index"`
} `json:"results"`
}
if err := json.Unmarshal(respBody, &parsed); err != nil {
return nil, fmt.Errorf("rerank decode: %w", err)
}
order := make([]int, 0, len(parsed.Results))
for _, row := range parsed.Results {
order = append(order, row.Index)
}
return order, nil
}
func (r *HTTPReranker) rerankDashScope(ctx context.Context, query string, documents []string, topN int) ([]int, error) {
url := r.dashscopeRerankURL()
body := map[string]any{
"model": r.model,
"input": map[string]any{
"query": query,
"documents": documents,
},
"parameters": map[string]any{
"return_documents": false,
"top_n": topN,
},
}
raw, err := json.Marshal(body)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+r.apiKey)
resp, err := r.client.Do(req)
if err != nil {
return nil, fmt.Errorf("dashscope rerank: %w", err)
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, fmt.Errorf("dashscope rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody)))
}
var parsed struct {
Output struct {
Results []struct {
Index int `json:"index"`
} `json:"results"`
} `json:"output"`
}
if err := json.Unmarshal(respBody, &parsed); err != nil {
return nil, fmt.Errorf("dashscope rerank decode: %w", err)
}
order := make([]int, 0, len(parsed.Output.Results))
for _, row := range parsed.Output.Results {
order = append(order, row.Index)
}
return order, nil
}
func (r *HTTPReranker) cohereRerankURL() string {
base := r.baseURL
if base == "" {
base = "https://api.cohere.com"
}
if strings.HasSuffix(base, "/v1") {
return base + "/rerank"
}
return base + "/v1/rerank"
}
func (r *HTTPReranker) dashscopeRerankURL() string {
base := strings.TrimSpace(r.baseURL)
if base == "" {
return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
}
if strings.Contains(base, "/api/v1/services/rerank") {
return base
}
if strings.Contains(base, "dashscope.aliyuncs.com") || strings.Contains(base, "compatible-mode") {
return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
}
return strings.TrimSuffix(base, "/")
}
func truncateForRerankLog(s string) string {
s = strings.TrimSpace(s)
if len(s) > 512 {
return s[:512] + "..."
}
return s
}
var _ DocumentReranker = (*HTTPReranker)(nil)
+97
View File
@@ -0,0 +1,97 @@
package knowledge
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"cyberstrike-ai/internal/config"
"github.com/cloudwego/eino/schema"
)
func TestHTTPReranker_CohereOrder(t *testing.T) {
t.Parallel()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/rerank" {
t.Fatalf("path %s", r.URL.Path)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"results": []map[string]any{
{"index": 2, "relevance_score": 0.9},
{"index": 0, "relevance_score": 0.5},
},
})
}))
defer srv.Close()
rr, err := NewHTTPReranker(&config.RerankConfig{
Provider: "cohere",
Model: "rerank-multilingual-v3.0",
BaseURL: srv.URL,
APIKey: "test-key",
}, nil, nil)
if err != nil {
t.Fatal(err)
}
docs := []*schema.Document{
{ID: "a", Content: "alpha"},
{ID: "b", Content: "beta"},
{ID: "c", Content: "gamma"},
}
out, err := rr.Rerank(context.Background(), "query", docs)
if err != nil {
t.Fatal(err)
}
if len(out) != 2 || out[0].ID != "c" || out[1].ID != "a" {
t.Fatalf("order wrong: %#v", out)
}
}
func TestHTTPReranker_DashScopeOrder(t *testing.T) {
t.Parallel()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_ = json.NewEncoder(w).Encode(map[string]any{
"output": map[string]any{
"results": []map[string]any{
{"index": 1, "relevance_score": 0.88},
},
},
})
}))
defer srv.Close()
rr, err := NewHTTPReranker(&config.RerankConfig{
Provider: "dashscope",
Model: "gte-rerank",
BaseURL: srv.URL,
APIKey: "test-key",
}, nil, nil)
if err != nil {
t.Fatal(err)
}
docs := []*schema.Document{{ID: "a", Content: "a"}, {ID: "b", Content: "b"}}
out, err := rr.Rerank(context.Background(), "q", docs)
if err != nil {
t.Fatal(err)
}
if len(out) != 1 || out[0].ID != "b" {
t.Fatalf("got %#v", out)
}
}
func TestRerankConfigDefaults(t *testing.T) {
t.Parallel()
rc := config.RerankConfig{}
if rc.ProviderEffective("https://dashscope.aliyuncs.com/x") != "dashscope" {
t.Fatal("dashscope detect")
}
if rc.ModelEffective("dashscope") != "gte-rerank" {
t.Fatal("dashscope model")
}
if rc.ModelEffective("cohere") != "rerank-multilingual-v3.0" {
t.Fatal("cohere model")
}
}
+8 -5
View File
@@ -19,7 +19,7 @@ import (
// postRetrieveMaxPrefetchCap 限制单次向量候选上限,避免误配置导致全表扫压力过大。 // postRetrieveMaxPrefetchCap 限制单次向量候选上限,避免误配置导致全表扫压力过大。
const postRetrieveMaxPrefetchCap = 200 const postRetrieveMaxPrefetchCap = 200
// DocumentReranker 可选重排(如交叉编码器 / 第三方 Rerank API),由 [Retriever.SetDocumentReranker] 注入;失败时在适配层降级为向量序 // DocumentReranker 精排(HTTP dashscope / Cohere 兼容 API),由 [WireRetrieverPipeline] 注入
type DocumentReranker interface { type DocumentReranker interface {
Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error)
} }
@@ -167,13 +167,16 @@ func truncateDocumentsByBudget(docs []*schema.Document, maxRunes, maxTokens int,
return out, nil return out, nil
} }
// EffectivePrefetchTopK 计算向量检索应拉取的候选条数(供粗排 / 重 / 重排)。 // EffectivePrefetchTopK 计算每条 MultiQuery 变体在向量阶段的候选条数(供融合 / 重 / 后处理)。
func EffectivePrefetchTopK(topK int, po *config.PostRetrieveConfig) int { func EffectivePrefetchTopK(topK int, po *config.PostRetrieveConfig) int {
if topK < 1 { if topK < 1 {
topK = 5 topK = 5
} }
fetch := topK fetch := topK * 4
if po != nil && po.PrefetchTopK > fetch { if fetch < 20 {
fetch = 20
}
if po != nil && po.PrefetchTopK > 0 {
fetch = po.PrefetchTopK fetch = po.PrefetchTopK
} }
if fetch > postRetrieveMaxPrefetchCap { if fetch > postRetrieveMaxPrefetchCap {
@@ -182,7 +185,7 @@ func EffectivePrefetchTopK(topK int, po *config.PostRetrieveConfig) int {
return fetch return fetch
} }
// ApplyPostRetrieve 检索后处理:规范化正文去重 → 预算截断 → 最终 TopK。重排在 [VectorEinoRetriever] 中单独调用以便失败时降级 // ApplyPostRetrieve 检索后处理:规范化正文去重 → 预算截断 → 最终 TopK(精排已在流水线中完成)
func ApplyPostRetrieve(docs []*schema.Document, po *config.PostRetrieveConfig, tokenModel string, finalTopK int) ([]*schema.Document, error) { func ApplyPostRetrieve(docs []*schema.Document, po *config.PostRetrieveConfig, tokenModel string, finalTopK int) ([]*schema.Document, error) {
if finalTopK < 1 { if finalTopK < 1 {
finalTopK = 5 finalTopK = 5
@@ -28,8 +28,8 @@ func TestDedupeByNormalizedContent(t *testing.T) {
} }
func TestEffectivePrefetchTopK(t *testing.T) { func TestEffectivePrefetchTopK(t *testing.T) {
if g := EffectivePrefetchTopK(5, nil); g != 5 { if g := EffectivePrefetchTopK(5, nil); g != 20 {
t.Fatalf("got %d", g) t.Fatalf("default prefetch got %d want 20", g)
} }
if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 50}); g != 50 { if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 50}); g != 50 {
t.Fatalf("got %d", g) t.Fatalf("got %d", g)
+39 -10
View File
@@ -27,15 +27,19 @@ type Retriever struct {
rerankMu sync.RWMutex rerankMu sync.RWMutex
reranker DocumentReranker reranker DocumentReranker
pipeline retriever.Retriever
wireOpenAI *config.OpenAIConfig
} }
// RetrievalConfig 检索配置 // RetrievalConfig 检索配置
type RetrievalConfig struct { type RetrievalConfig struct {
TopK int TopK int
SimilarityThreshold float64 SimilarityThreshold float64
// SubIndexFilter 非空时仅检索 sub_indexes 包含该标签(逗号分隔之一)的行;空 sub_indexes 的旧行仍保留以兼容。 SubIndexFilter string
SubIndexFilter string MultiQuery config.MultiQueryConfig
PostRetrieve config.PostRetrieveConfig Rerank config.RerankConfig
PostRetrieve config.PostRetrieveConfig
} }
// NewRetriever 创建新的检索器 // NewRetriever 创建新的检索器
@@ -48,7 +52,7 @@ func NewRetriever(db *sql.DB, embedder *Embedder, config *RetrievalConfig, logge
} }
} }
// UpdateConfig 更新检索配置 // UpdateConfig 更新检索配置并重建 Eino MultiQuery + 重排流水线。
func (r *Retriever) UpdateConfig(cfg *RetrievalConfig) { func (r *Retriever) UpdateConfig(cfg *RetrievalConfig) {
if cfg != nil { if cfg != nil {
r.config = cfg r.config = cfg
@@ -57,12 +61,18 @@ func (r *Retriever) UpdateConfig(cfg *RetrievalConfig) {
zap.Int("top_k", cfg.TopK), zap.Int("top_k", cfg.TopK),
zap.Float64("similarity_threshold", cfg.SimilarityThreshold), zap.Float64("similarity_threshold", cfg.SimilarityThreshold),
zap.String("sub_index_filter", cfg.SubIndexFilter), zap.String("sub_index_filter", cfg.SubIndexFilter),
zap.Int("multi_query_max", cfg.MultiQuery.MaxQueriesEffective()),
zap.Int("post_retrieve_prefetch_top_k", cfg.PostRetrieve.PrefetchTopK), zap.Int("post_retrieve_prefetch_top_k", cfg.PostRetrieve.PrefetchTopK),
zap.Int("post_retrieve_max_context_chars", cfg.PostRetrieve.MaxContextChars), zap.Int("post_retrieve_max_context_chars", cfg.PostRetrieve.MaxContextChars),
zap.Int("post_retrieve_max_context_tokens", cfg.PostRetrieve.MaxContextTokens), zap.Int("post_retrieve_max_context_tokens", cfg.PostRetrieve.MaxContextTokens),
) )
} }
} }
if r.wireOpenAI != nil {
if err := WireRetrieverPipeline(context.Background(), r, r.wireOpenAI); err != nil && r.logger != nil {
r.logger.Warn("检索流水线重建失败", zap.Error(err))
}
}
} }
// SetDocumentReranker 注入可选重排器(并发安全);nil 表示禁用。 // SetDocumentReranker 注入可选重排器(并发安全);nil 表示禁用。
@@ -103,7 +113,7 @@ func cosineSimilarity(a, b []float32) float64 {
return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)) return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
} }
// Search 搜索知识库。统一经 [VectorEinoRetriever]Eino retriever.Retriever 边界)。 // Search 搜索知识库Eino MultiQuery → 向量检索 → 重排 → 后处理)。
func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) { func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) {
if req == nil { if req == nil {
return nil, fmt.Errorf("请求不能为空") return nil, fmt.Errorf("请求不能为空")
@@ -113,7 +123,7 @@ func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*Retrieva
return nil, fmt.Errorf("查询不能为空") return nil, fmt.Errorf("查询不能为空")
} }
opts := r.einoRetrieverOptions(req) opts := r.einoRetrieverOptions(req)
docs, err := NewVectorEinoRetriever(r).Retrieve(ctx, q, opts...) docs, err := r.activeEinoRetriever().Retrieve(ctx, q, opts...)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -143,7 +153,19 @@ func (r *Retriever) einoRetrieverOptions(req *SearchRequest) []retriever.Option
// EinoRetrieve 直接返回 [schema.Document],供 Eino Graph / Chain 使用。 // EinoRetrieve 直接返回 [schema.Document],供 Eino Graph / Chain 使用。
func (r *Retriever) EinoRetrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { func (r *Retriever) EinoRetrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {
return NewVectorEinoRetriever(r).Retrieve(ctx, query, opts...) return r.activeEinoRetriever().Retrieve(ctx, query, opts...)
}
func (r *Retriever) activeEinoRetriever() retriever.Retriever {
if r != nil && r.pipeline != nil {
return r.pipeline
}
return NewVectorEinoRetriever(r)
}
// AsEinoRetriever 将知识库检索流水线暴露为 Eino [retriever.Retriever]。
func (r *Retriever) AsEinoRetriever() retriever.Retriever {
return r.activeEinoRetriever()
} }
func (r *Retriever) knowledgeEmbeddingSelectSQL(riskType, subIndexFilter string) (string, []interface{}) { func (r *Retriever) knowledgeEmbeddingSelectSQL(riskType, subIndexFilter string) (string, []interface{}) {
@@ -299,7 +321,14 @@ func (r *Retriever) vectorSearch(ctx context.Context, req *SearchRequest) ([]*Re
return results, nil return results, nil
} }
// AsEinoRetriever 将纯向量检索暴露为 Eino [retriever.Retriever]。 // RetrievalConfigFromYAML maps API/YAML retrieval settings into the knowledge package.
func (r *Retriever) AsEinoRetriever() retriever.Retriever { func RetrievalConfigFromYAML(r config.RetrievalConfig) *RetrievalConfig {
return NewVectorEinoRetriever(r) return &RetrievalConfig{
TopK: r.TopK,
SimilarityThreshold: r.SimilarityThreshold,
SubIndexFilter: r.SubIndexFilter,
MultiQuery: r.MultiQuery,
Rerank: r.Rerank,
PostRetrieve: r.PostRetrieve,
}
} }
+74
View File
@@ -0,0 +1,74 @@
package knowledge
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/openai"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/flow/retriever/multiquery"
"go.uber.org/zap"
)
// WireRetrieverPipeline builds Eino MultiQuery + HTTP rerank + post-process pipeline on r.
// Call once after NewRetriever; UpdateConfig re-invokes when wireOpenAI is set.
func WireRetrieverPipeline(ctx context.Context, r *Retriever, openAI *config.OpenAIConfig) error {
if r == nil {
return fmt.Errorf("retriever is nil")
}
if openAI == nil {
return fmt.Errorf("openai config is nil")
}
if r.config == nil {
return fmt.Errorf("retrieval config is nil")
}
r.wireOpenAI = openAI
httpClient := openai.NewEinoHTTPClient(openAI, &http.Client{Timeout: 120 * time.Second})
chatCfg := &einoopenai.ChatModelConfig{
APIKey: strings.TrimSpace(openAI.APIKey),
BaseURL: strings.TrimSuffix(strings.TrimSpace(openAI.BaseURL), "/"),
Model: strings.TrimSpace(openAI.Model),
HTTPClient: httpClient,
}
if chatCfg.Model == "" {
chatCfg.Model = "gpt-4o"
}
rewriteLLM, err := einoopenai.NewChatModel(ctx, chatCfg)
if err != nil {
return fmt.Errorf("multi_query rewrite model: %w", err)
}
reranker, err := NewHTTPReranker(&r.config.Rerank, openAI, r.logger)
if err != nil {
return fmt.Errorf("reranker: %w", err)
}
r.SetDocumentReranker(reranker)
vec := NewVectorEinoRetriever(r)
mq, err := multiquery.NewRetriever(ctx, &multiquery.Config{
RewriteLLM: rewriteLLM,
MaxQueriesNum: r.config.MultiQuery.MaxQueriesEffective(),
OrigRetriever: vec,
})
if err != nil {
return fmt.Errorf("multi_query: %w", err)
}
r.pipeline = newKnowledgePipelineRetriever(mq, r)
if r.logger != nil {
provider := r.config.Rerank.ProviderEffective(strings.TrimSpace(openAI.BaseURL))
r.logger.Info("知识库检索流水线已启用",
zap.String("pipeline", "MultiQuery→Vector→Rerank→PostRetrieve"),
zap.Int("multi_query_max", r.config.MultiQuery.MaxQueriesEffective()),
zap.String("rerank_provider", provider),
zap.String("rerank_model", r.config.Rerank.ModelEffective(provider)),
)
}
return nil
}
+17
View File
@@ -814,6 +814,23 @@ func (m *ExternalMCPManager) CancelToolExecution(id string) bool {
return m.CancelToolExecutionWithNote(id, "") return m.CancelToolExecutionWithNote(id, "")
} }
// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的外部 MCP executionId 快照。
func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
if m == nil {
return nil
}
m.mu.Lock()
defer m.mu.Unlock()
if len(m.runningCancels) == 0 {
return nil
}
out := make(map[string]struct{}, len(m.runningCancels))
for id := range m.runningCancels {
out[id] = struct{}{}
}
return out
}
// updateStats 更新统计信息 // updateStats 更新统计信息
func (m *ExternalMCPManager) updateStats(toolName string, failed bool) { func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
now := time.Now() now := time.Now()
+17
View File
@@ -1170,6 +1170,23 @@ func (s *Server) CancelToolExecution(id string) bool {
return s.CancelToolExecutionWithNote(id, "") return s.CancelToolExecutionWithNote(id, "")
} }
// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的 executionId 快照。
func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} {
if s == nil {
return nil
}
s.runningCancelsMu.Lock()
defer s.runningCancelsMu.Unlock()
if len(s.runningCancels) == 0 {
return nil
}
out := make(map[string]struct{}, len(s.runningCancels))
for id := range s.runningCancels {
out[id] = struct{}{}
}
return out
}
// initDefaultPrompts 初始化默认提示词模板 // initDefaultPrompts 初始化默认提示词模板
func (s *Server) initDefaultPrompts() { func (s *Server) initDefaultPrompts() {
s.mu.Lock() s.mu.Lock()
+101
View File
@@ -0,0 +1,101 @@
package monitor
import (
"time"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
const (
staleRunningMinAge = 45 * time.Second
staleRunningReconcileGap = 2 * time.Minute
)
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 cancelled。
type ExecutionReconciler struct {
db *database.DB
mcpServer *mcp.Server
externalMgr *mcp.ExternalMCPManager
logger *zap.Logger
}
// NewExecutionReconciler creates a reconciler for orphaned MCP tool executions.
func NewExecutionReconciler(db *database.DB, mcpServer *mcp.Server, externalMgr *mcp.ExternalMCPManager, logger *zap.Logger) *ExecutionReconciler {
return &ExecutionReconciler{
db: db,
mcpServer: mcpServer,
externalMgr: externalMgr,
logger: logger,
}
}
// ReconcileOnStartup marks every persisted running row as cancelled (safe right after process start).
func (r *ExecutionReconciler) ReconcileOnStartup() {
if r == nil || r.db == nil {
return
}
now := time.Now()
n, err := r.db.CancelOrphanedRunningToolExecutions(now, "执行已中断(服务重启)")
if err != nil {
if r.logger != nil {
r.logger.Warn("启动时清理孤儿 running 工具执行记录失败", zap.Error(err))
}
return
}
if n > 0 && r.logger != nil {
r.logger.Info("启动时已收尾孤儿 running 工具执行记录", zap.Int64("count", n))
}
}
func (r *ExecutionReconciler) activeExecutionIDs() map[string]struct{} {
ids := make(map[string]struct{})
if r.mcpServer != nil {
for id := range r.mcpServer.ActiveRunningExecutionIDs() {
ids[id] = struct{}{}
}
}
if r.externalMgr != nil {
for id := range r.externalMgr.ActiveRunningExecutionIDs() {
ids[id] = struct{}{}
}
}
return ids
}
// ReconcileStaleRunning finalizes running rows that are not tracked in-memory and older than staleRunningMinAge.
func (r *ExecutionReconciler) ReconcileStaleRunning() {
if r == nil || r.db == nil {
return
}
now := time.Now()
n, err := r.db.FinalizeStaleRunningToolExecutions(now, staleRunningMinAge, r.activeExecutionIDs(), "执行已中断(会话已结束)")
if err != nil {
if r.logger != nil {
r.logger.Warn("定期收尾 stale running 工具执行记录失败", zap.Error(err))
}
return
}
if n > 0 && r.logger != nil {
r.logger.Info("已收尾 stale running 工具执行记录", zap.Int64("count", n))
}
}
// StartStaleRunningReconcileLoop periodically reconciles orphaned running tool executions.
func StartStaleRunningReconcileLoop(r *ExecutionReconciler, logger *zap.Logger) {
if r == nil {
return
}
go func() {
ticker := time.NewTicker(staleRunningReconcileGap)
defer ticker.Stop()
for range ticker.C {
r.ReconcileStaleRunning()
if logger != nil {
logger.Debug("monitor stale running reconcile tick completed")
}
}
}()
}
+38
View File
@@ -0,0 +1,38 @@
package monitor
import (
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestExecutionReconciler_ReconcileOnStartup(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "monitor.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "run-1", ToolName: "hydra", Status: "running", StartTime: time.Now().Add(-time.Hour),
}); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
r := NewExecutionReconciler(db, mcp.NewServer(zap.NewNop()), nil, zap.NewNop())
r.ReconcileOnStartup()
got, err := db.GetToolExecution("run-1")
if err != nil {
t.Fatalf("GetToolExecution: %v", err)
}
if got.Status != "cancelled" {
t.Fatalf("expected cancelled after startup reconcile, got %s", got.Status)
}
}
+16
View File
@@ -0,0 +1,16 @@
package multiagent
import (
"fmt"
"github.com/cloudwego/eino/adk"
)
// InitADK configures global Eino ADK settings. Call once at process startup before
// any ADK middleware or agents are created.
func InitADK() error {
if err := adk.SetLanguage(adk.LanguageChinese); err != nil {
return fmt.Errorf("adk set language: %w", err)
}
return nil
}
+3 -19
View File
@@ -299,6 +299,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
var toolResultSent sync.Map // toolCallID -> struct{}ADK Tool 事件去重(权威正文来自 reduction 处理后的 agent 上下文) var toolResultSent sync.Map // toolCallID -> struct{}ADK Tool 事件去重(权威正文来自 reduction 处理后的 agent 上下文)
tryEmitToolResultProgress := func(toolName, content, toolCallID string, isErr bool, agentName string) { tryEmitToolResultProgress := func(toolName, content, toolCallID string, isErr bool, agentName string) {
// 仅由 ADK schema.Tool 事件调用;MCP/execute 桥在 reduction 前的 ToolInvokeNotify 不得推送 tool_result
// 否则全量输出会先占位并触发 toolResultSent 去重,导致 UI/监控展示与 agent 实际收到的截断正文不一致。
if progress == nil { if progress == nil {
return return
} }
@@ -316,6 +318,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"isError": isErr, "isError": isErr,
"result": content, "result": content,
"resultPreview": preview, "resultPreview": preview,
"agentFacing": true, // 与 reduction 后送入 ChatModel 的正文一致,供前端展示
"conversationId": conversationID, "conversationId": conversationID,
"einoAgent": agentName, "einoAgent": agentName,
"einoRole": einoRoleTag(agentName), "einoRole": einoRoleTag(agentName),
@@ -350,25 +353,6 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
} }
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data) progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
} }
if args.ToolInvokeNotify != nil {
args.ToolInvokeNotify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
// Eino execute / MCP 桥在工具返回时 Fire;若 ADK schema.Tool 事件迟迟不到,此处立即推送
// tool_result 解除 UI「执行中」。tryEmitToolResultProgress 经 toolResultSent 去重,ADK 晚到不重复。
isErr := !success || invokeErr != nil
body := einoToolResultBody(content)
if einoToolResultIsError(toolName, content) {
isErr = true
}
if tail := friendlyEinoExecuteInvokeTail(invokeErr); tail != "" {
if body == "" {
body = tail
} else if !strings.Contains(body, tail) {
body = strings.TrimSpace(body) + "\n\n" + tail
}
}
tryEmitToolResultProgress(toolName, body, toolCallID, isErr, einoAgent)
})
}
if args.EinoCallbacks != nil { if args.EinoCallbacks != nil {
ctx = einoobserve.AttachAgentRunCallbacks(ctx, args.EinoCallbacks, einoobserve.Params{ ctx = einoobserve.AttachAgentRunCallbacks(ctx, args.EinoCallbacks, einoobserve.Params{
@@ -0,0 +1,59 @@
package multiagent
import (
"strings"
"time"
"cyberstrike-ai/internal/config"
)
const defaultEmptyResponseContinueMaxAttempts = 5
// IsEinoEmptyResponseResult 判断 Run 是否以「未捕获助手正文」占位结束(非真实用户可见回复)。
func IsEinoEmptyResponseResult(result *RunResult) bool {
if result == nil {
return false
}
return isEinoEmptyResponseText(result.Response)
}
func isEinoEmptyResponseText(s string) bool {
s = strings.TrimSpace(s)
if s == "" {
return false
}
return strings.Contains(s, "no assistant text was captured") ||
strings.Contains(s, "未捕获到助手文本输出")
}
// HasEinoResumeTrace 轨迹非空,续跑才有上下文可恢复。
func HasEinoResumeTrace(result *RunResult) bool {
if result == nil {
return false
}
s := strings.TrimSpace(result.LastAgentTraceInput)
return s != "" && s != "[]" && s != "null"
}
// EmptyResponseContinueMaxAttemptsFromConfig 无助手正文时 Handler 层退避续跑上限;0=默认 5。
func EmptyResponseContinueMaxAttemptsFromConfig(mw *config.MultiAgentEinoMiddlewareConfig) int {
if mw != nil && mw.EmptyResponseContinueMaxAttempts > 0 {
return mw.EmptyResponseContinueMaxAttempts
}
return defaultEmptyResponseContinueMaxAttempts
}
// EmptyResponseContinueBackoff 与 run_retry 相同指数退避(2s, 4s, 8s… capped)。
func EmptyResponseContinueBackoff(attempt int, mw *config.MultiAgentEinoMiddlewareConfig) time.Duration {
maxBackoff := defaultEinoRunRetryMaxBackoff
if mw != nil && mw.RunRetryMaxBackoffSec > 0 {
maxBackoff = time.Duration(mw.RunRetryMaxBackoffSec) * time.Second
}
return einoTransientRetryBackoff(attempt, maxBackoff)
}
// FormatEmptyResponseContinueUserMessage 系统自动续跑时注入的 user 轮次(不写入 messages 表气泡)。
func FormatEmptyResponseContinueUserMessage() string {
return strings.TrimSpace(`系统自动续跑 / Auto resume
上一轮 Eino 会话未产出可见助手正文可能流式中断或仅完成工具调用请基于已有轨迹与工具结果继续推进并给出阶段性总结勿重复已完成步骤`)
}
@@ -0,0 +1,38 @@
package multiagent
import "testing"
func TestIsEinoEmptyResponseResult(t *testing.T) {
empty := &RunResult{
Response: "(Eino ADK single-agent session completed but no assistant text was captured. Check process details or logs.) " +
"Eino ADK 单代理会话已完成,但未捕获到助手文本输出。请查看过程详情或日志。)",
}
if !IsEinoEmptyResponseResult(empty) {
t.Fatal("expected empty placeholder response")
}
ok := &RunResult{Response: "扫描完成,发现 2 个开放端口。"}
if IsEinoEmptyResponseResult(ok) {
t.Fatalf("expected real response, got placeholder match")
}
if IsEinoEmptyResponseResult(nil) {
t.Fatal("nil result should be false")
}
}
func TestHasEinoResumeTrace(t *testing.T) {
if HasEinoResumeTrace(nil) {
t.Fatal("nil")
}
if HasEinoResumeTrace(&RunResult{LastAgentTraceInput: "[]"}) {
t.Fatal("enable resume on empty trace")
}
if !HasEinoResumeTrace(&RunResult{LastAgentTraceInput: `[{"role":"user","content":"hi"}]`}) {
t.Fatal("expected resume trace")
}
}
func TestEmptyResponseContinueMaxAttemptsFromConfig(t *testing.T) {
if got := EmptyResponseContinueMaxAttemptsFromConfig(nil); got != defaultEmptyResponseContinueMaxAttempts {
t.Fatalf("default: got %d want %d", got, defaultEmptyResponseContinueMaxAttempts)
}
}
@@ -84,7 +84,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround { if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround {
req.RunInBackendGround = true req.RunInBackendGround = true
} }
req.Command = security.PrepareNonInteractiveShellCommand(prependPythonUnbufferedEnv(req.Command)) req.Command = prependPythonUnbufferedEnv(req.Command)
convID := mcp.MCPConversationIDFromContext(ctx) convID := mcp.MCPConversationIDFromContext(ctx)
execReg := mcp.EinoExecuteRunRegistryFromContext(ctx) execReg := mcp.EinoExecuteRunRegistryFromContext(ctx)
@@ -61,12 +61,6 @@ func TestEinoStreamingShellWrap_PreparesNonInteractiveCommand(t *testing.T) {
t.Fatalf("recv: %v", rerr) t.Fatalf("recv: %v", rerr)
} }
} }
if !strings.Contains(inner.lastCommand, "exec </dev/null") {
t.Fatalf("missing stdin redirect in inner command: %q", inner.lastCommand)
}
if !strings.Contains(inner.lastCommand, "GIT_PAGER=cat") {
t.Fatalf("missing pager export in inner command: %q", inner.lastCommand)
}
if !strings.Contains(inner.lastCommand, "PYTHONUNBUFFERED=1") { if !strings.Contains(inner.lastCommand, "PYTHONUNBUFFERED=1") {
t.Fatalf("missing python unbuffer in inner command: %q", inner.lastCommand) t.Fatalf("missing python unbuffer in inner command: %q", inner.lastCommand)
} }
+36 -28
View File
@@ -80,34 +80,9 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
return nil, fmt.Errorf("plan_execute replanner: %w", err) return nil, fmt.Errorf("plan_execute replanner: %w", err)
} }
// 组装 executor handler 栈,顺序与 Deep/Supervisor 主代理一致(outermost first)。 execHandlers, err := buildPlanExecuteExecutorHandlers(ctx, a)
var execHandlers []adk.ChatModelAgentMiddleware if err != nil {
// 1. patchtoolcalls, reduction, toolsearch, plantask(来自 prependEinoMiddlewares return nil, err
if len(a.ExecPreMiddlewares) > 0 {
execHandlers = append(execHandlers, a.ExecPreMiddlewares...)
}
// 2. filesystem 中间件(可选)
if a.FilesystemMiddleware != nil {
execHandlers = append(execHandlers, a.FilesystemMiddleware)
}
// 3. skill 中间件(可选)
if a.SkillMiddleware != nil {
execHandlers = append(execHandlers, a.SkillMiddleware)
}
// 4. pre-summarization normalize + continuation dedup, then summarization (与 Deep/Supervisor 一致)
if a.AppCfg != nil {
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger)
if sumErr != nil {
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
}
execHandlers = appendEinoChatModelTailMiddlewares(execHandlers, einoChatModelTailConfig{
logger: a.Logger,
phase: "plan_execute_executor",
summarization: sumMw,
modelName: a.ModelName,
conversationID: a.ConversationID,
trace: a.ModelFacingTrace,
})
} }
executor, err := newPlanExecuteExecutor(ctx, &planexecute.ExecutorConfig{ executor, err := newPlanExecuteExecutor(ctx, &planexecute.ExecutorConfig{
Model: a.ExecModel, Model: a.ExecModel,
@@ -130,6 +105,39 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
}) })
} }
// buildPlanExecuteExecutorHandlers 组装 Executor 中间件栈(outermost first),与 Deep/Supervisor 主代理对齐:
// ExecPreMiddlewarespatch / reduction / toolsearch / plantask)→ filesystem → skill → summarization tail。
func buildPlanExecuteExecutorHandlers(ctx context.Context, a *PlanExecuteRootArgs) ([]adk.ChatModelAgentMiddleware, error) {
if a == nil {
return nil, fmt.Errorf("plan_execute: args 为空")
}
var execHandlers []adk.ChatModelAgentMiddleware
if len(a.ExecPreMiddlewares) > 0 {
execHandlers = append(execHandlers, a.ExecPreMiddlewares...)
}
if a.FilesystemMiddleware != nil {
execHandlers = append(execHandlers, a.FilesystemMiddleware)
}
if a.SkillMiddleware != nil {
execHandlers = append(execHandlers, a.SkillMiddleware)
}
if a.AppCfg != nil {
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger)
if sumErr != nil {
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
}
execHandlers = appendEinoChatModelTailMiddlewares(execHandlers, einoChatModelTailConfig{
logger: a.Logger,
phase: "plan_execute_executor",
summarization: sumMw,
modelName: a.ModelName,
conversationID: a.ConversationID,
trace: a.ModelFacingTrace,
})
}
return execHandlers, nil
}
// planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。 // planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。
// 返回 nil 时 Eino 使用内置默认 planner prompt。 // 返回 nil 时 Eino 使用内置默认 planner prompt。
func planExecutePlannerGenInput( func planExecutePlannerGenInput(
@@ -184,6 +184,7 @@ func RunEinoSingleChatModelAgent(
Name: einoSingleAgentName, Name: einoSingleAgentName,
Description: "Eino ADK ChatModelAgent with MCP tools for authorized security testing.", Description: "Eino ADK ChatModelAgent with MCP tools for authorized security testing.",
Instruction: ins, Instruction: ins,
GenModelInput: literalInstructionGenModelInput,
Model: mainModel, Model: mainModel,
ToolsConfig: mainToolsCfg, ToolsConfig: mainToolsCfg,
MaxIterations: maxIter, MaxIterations: maxIter,
+53 -7
View File
@@ -22,15 +22,60 @@ import (
"go.uber.org/zap" "go.uber.org/zap"
) )
// einoSummarizeUserInstruction:压缩历史时保留渗透测试关键信息。 // einoSummarizeUserInstruction:压缩历史时保留渗透测试与用户约束关键信息。
const einoSummarizeUserInstruction = `在保持所有关键安全测试信息完整的前提下压缩对话历史 // 结构对齐 Eino 最佳实践(禁止工具、<analysis>+<summary>、<all_user_messages>),章节为安全测试领域化。
const einoSummarizeUserInstruction = `关键仅以纯文本响应禁止调用任何工具read_fileexecgrepglobwriteedit
上述对话中已包含全部待压缩上下文不要要求用户粘贴历史不要输出请提供待压缩的对话历史等占位/meta 回复
工具调用将被拒绝并浪费唯一一次摘要机会
必须保留已确认漏洞与攻击路径工具输出中的核心发现凭证与认证细节架构与薄弱点当前进度失败尝试与死路策略决策 你的任务在保持所有关键安全测试信息完整的前提下压缩对话历史使后续代理能无缝继续同一授权测试任务
保留精确技术细节URL路径参数Payload版本号报错原文可摘要但要点不丢
将冗长扫描输出概括为结论重复发现合并表述
已枚举资产须保留**可继承的摘要**主域关键子域/主机短表或数量+代表样例高价值目标与已识别服务/端口要点避免后续子代理因看不见清单而重复全量枚举
输出须使后续代理能无缝继续同一授权测试任务` 压缩原则
- 必须保留已确认漏洞与攻击路径工具输出核心发现凭证与认证细节架构与薄弱点当前进度失败尝试与死路策略决策
- 保留精确技术细节URL路径参数Payload版本号报错原文可摘要但要点不丢
- 冗长扫描输出概括为结论重复发现合并表述
- 已枚举资产须保留可继承摘要主域关键子域/主机短表或数量+代表样例高价值目标已识别服务/端口要点
输出格式严格遵循仅一轮回复
1. 先输出 <analysis> 按时间顺序梳理对话检查是否涵盖下方各章节要点analysis 仅供自检保持简洁建议 400
2. 再输出 <summary> 按以下章节写入可继承的压缩报告无信息处写禁止留空模板占位符
<summary>
## 1. 授权范围与约束
- 目标/范围/禁止项域名路径IP环境
- 凭证/认证信息账号TokenCookie敏感值原文保留
- 用户指定的方法工具优先级与待办
- 否定约束不测什么不用什么手法
## 2. 资产与服务枚举摘要
- 主域/核心资产关键子域或主机短表或数量+代表样例
- 高价值目标已识别服务/端口要点
- 资产状态存活/可攻/已排除/待验证
## 3. 架构与已知薄弱点
- 技术栈/部署拓扑/信任边界
- 已识别薄弱点列表
## 4. 已确认漏洞与攻击路径
- 漏洞名/CVEURL/路径参数/PayloadPoC 要点影响等级
- 攻击链/利用路径步骤化
## 5. 工具核心发现与扫描结论
- 各工具结论概括核心输出非冗长日志
- 重复发现合并表述
## 6. 所有用户消息
<all_user_messages>
- [逐条列出非 tool 结果的用户消息要点敏感约束与原文措辞尽量保留]
</all_user_messages>
## 7. 当前进度策略决策与下一步
- 当前位置已完成/进行中/卡点
- 失败尝试与死路方法现象/报错摘要结论
- 策略决策与下一步具体操作须与最近用户请求及未完成任务一致
</summary>
提醒不要调用任何工具必须基于上文已有对话直接输出 <analysis> <summary>勿输出 analysis 以外的正文`
// newEinoSummarizationMiddleware 使用 Eino ADK Summarization 中间件(见 https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_summarization/)。 // newEinoSummarizationMiddleware 使用 Eino ADK Summarization 中间件(见 https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/eino_adk_chatmodelagentmiddleware/middleware_summarization/)。
// 触发阈值:估算 token 超过 openai.max_total_tokens * summarization_trigger_ratio(默认 0.8)时摘要。 // 触发阈值:估算 token 超过 openai.max_total_tokens * summarization_trigger_ratio(默认 0.8)时摘要。
@@ -144,6 +189,7 @@ 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) {
summary = stripAnalysisFromSummarizationMessage(summary)
out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax) out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax)
if ferr != nil { if ferr != nil {
return nil, ferr return nil, ferr
@@ -0,0 +1,73 @@
package multiagent
import (
"regexp"
"strings"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
var (
summarizationAnalysisBlockRegex = regexp.MustCompile(`(?is)<analysis>\s*.*?\s*</analysis>`)
summarizationSummaryBlockRegex = regexp.MustCompile(`(?is)<summary>\s*(.*?)\s*</summary>`)
)
// stripAnalysisFromSummarizationMessage removes the <analysis> block from a post-processed
// Eino summary user message. Analysis helps one-shot generation quality but should not
// occupy continuation context after compaction.
func stripAnalysisFromSummarizationMessage(msg adk.Message) adk.Message {
if msg == nil {
return msg
}
cloned := *msg
if cloned.Content != "" {
cloned.Content = stripAnalysisFromSummarizationText(cloned.Content)
}
if len(cloned.UserInputMultiContent) > 0 {
parts := make([]schema.MessageInputPart, len(cloned.UserInputMultiContent))
copy(parts, cloned.UserInputMultiContent)
// Only the first text part carries model output plus Eino preamble/transcript path.
for i := range parts {
if parts[i].Type != schema.ChatMessagePartTypeText || parts[i].Text == "" {
continue
}
if i == 0 {
parts[i].Text = stripAnalysisFromSummarizationText(parts[i].Text)
}
break
}
cloned.UserInputMultiContent = parts
}
return &cloned
}
func stripAnalysisFromSummarizationText(text string) string {
text = strings.TrimSpace(text)
if text == "" {
return text
}
stripped := strings.TrimSpace(summarizationAnalysisBlockRegex.ReplaceAllString(text, ""))
if stripped == "" {
return text
}
return stripped
}
// extractSummarizationSummaryBody returns the inner text of the last <summary> block when present.
// Used by tests and optional strict compaction paths.
func extractSummarizationSummaryBody(text string) (string, bool) {
text = strings.TrimSpace(text)
if text == "" {
return "", false
}
all := summarizationSummaryBlockRegex.FindAllStringSubmatch(text, -1)
if len(all) == 0 || len(all[len(all)-1]) < 2 {
return "", false
}
body := strings.TrimSpace(all[len(all)-1][1])
if body == "" {
return "", false
}
return body, true
}
@@ -0,0 +1,67 @@
package multiagent
import (
"strings"
"testing"
"github.com/cloudwego/eino/schema"
)
func TestStripAnalysisFromSummarizationText(t *testing.T) {
in := "<analysis>internal notes</analysis>\n\n<summary>\n## 1. 授权\n- example.com\n</summary>"
got := stripAnalysisFromSummarizationText(in)
if strings.Contains(got, "<analysis>") {
t.Fatalf("analysis block should be removed: %q", got)
}
if !strings.Contains(got, "## 1. 授权") {
t.Fatalf("summary body should remain: %q", got)
}
}
func TestStripAnalysisFromSummarizationMessage_UserInputMultiContent(t *testing.T) {
msg := &schema.Message{
Role: schema.User,
UserInputMultiContent: []schema.MessageInputPart{
{
Type: schema.ChatMessagePartTypeText,
Text: "此会话延续自此前一段因上下文耗尽而终止的对话。\n\n<analysis>draft</analysis>\n<summary>body</summary>\n\n完整记录位于:/tmp/transcript.txt",
},
{
Type: schema.ChatMessagePartTypeText,
Text: "请从我们中断的地方继续对话,无需向用户提出任何进一步的问题。",
},
},
}
out := stripAnalysisFromSummarizationMessage(msg)
if len(out.UserInputMultiContent) != 2 {
t.Fatalf("expected 2 parts, got %d", len(out.UserInputMultiContent))
}
if strings.Contains(out.UserInputMultiContent[0].Text, "<analysis>") {
t.Fatalf("part 0 should drop analysis: %q", out.UserInputMultiContent[0].Text)
}
if !strings.Contains(out.UserInputMultiContent[0].Text, "<summary>body</summary>") {
t.Fatalf("part 0 should keep summary: %q", out.UserInputMultiContent[0].Text)
}
if out.UserInputMultiContent[1].Text != "请从我们中断的地方继续对话,无需向用户提出任何进一步的问题。" {
t.Fatalf("continue instruction part should be unchanged: %q", out.UserInputMultiContent[1].Text)
}
}
func TestExtractSummarizationSummaryBody(t *testing.T) {
body, ok := extractSummarizationSummaryBody("<analysis>x</analysis><summary> kept </summary>")
if !ok || body != "kept" {
t.Fatalf("extract summary body: ok=%v body=%q", ok, body)
}
_, ok = extractSummarizationSummaryBody("plain text only")
if ok {
t.Fatal("expected false for plain text")
}
}
func TestStripAnalysisFromSummarizationText_NoAnalysisUnchanged(t *testing.T) {
in := "<summary>only summary</summary>"
got := stripAnalysisFromSummarizationText(in)
if got != in {
t.Fatalf("expected unchanged text, got %q", got)
}
}
+2 -25
View File
@@ -1,35 +1,12 @@
package multiagent package multiagent
import ( import (
"github.com/bytedance/sonic" copenai "cyberstrike-ai/internal/openai"
) )
// stripReasoningFromSummarizationPayload removes thinking / reasoning fields from a // stripReasoningFromSummarizationPayload removes thinking / reasoning fields from a
// chat-completions JSON body. Applied only to summarization Generate calls via // chat-completions JSON body. Applied only to summarization Generate calls via
// model.ModelOptions on the shared ChatModel — main-agent requests are unchanged. // model.ModelOptions on the shared ChatModel — main-agent requests are unchanged.
func stripReasoningFromSummarizationPayload(rawBody []byte) ([]byte, error) { func stripReasoningFromSummarizationPayload(rawBody []byte) ([]byte, error) {
var payload map[string]any return copenai.StripReasoningFromChatCompletionBody(rawBody)
if err := sonic.Unmarshal(rawBody, &payload); err != nil {
return rawBody, nil
}
changed := false
for _, key := range []string{
"thinking",
"reasoning_effort",
"output_config",
"reasoning",
} {
if _, ok := payload[key]; ok {
delete(payload, key)
changed = true
}
}
if !changed {
return rawBody, nil
}
out, err := sonic.Marshal(payload)
if err != nil {
return rawBody, err
}
return out, nil
} }
+5 -5
View File
@@ -409,9 +409,9 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
"需要写入请使用 upsert_project_fact。", "需要写入请使用 upsert_project_fact。",
project.FactIndexSectionEndMarker, project.FactIndexSectionEndMarker,
"", "",
"# Skills System", transcriptSkillsSystemMarker,
"**How to Use Skills**", "**如何使用 Skill(技能)(渐进式展示):**",
"Remember: Skills make you more capable", "记住:Skill 让你更加强大和稳定",
}, "\n") }, "\n")
out := sanitizeSystemContentForTranscript(system) out := sanitizeSystemContentForTranscript(system)
@@ -421,7 +421,7 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
if strings.Contains(out, "- nmap") || strings.Contains(out, "高强度扫描要求") { if strings.Contains(out, "- nmap") || strings.Contains(out, "高强度扫描要求") {
t.Fatalf("static persona should be stripped: %q", out) t.Fatalf("static persona should be stripped: %q", out)
} }
if strings.Contains(out, "# Skills System") || strings.Contains(out, "How to Use Skills") { if strings.Contains(out, transcriptSkillsSystemMarker) || strings.Contains(out, "如何使用 Skill") {
t.Fatalf("skills boilerplate should be stripped: %q", out) t.Fatalf("skills boilerplate should be stripped: %q", out)
} }
if !strings.Contains(out, transcriptStaticSystemOmitNote) { if !strings.Contains(out, transcriptStaticSystemOmitNote) {
@@ -435,7 +435,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.FactIndexSectionStartMarker + "\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + project.FactIndexSectionEndMarker + "\n# Skills System\nboiler"), schema.SystemMessage("以下是当前会话绑定的工具名称索引\n- nmap\n\n你是CyberStrikeAI\n" + project.FactIndexSectionStartMarker + "\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + project.FactIndexSectionEndMarker + "\n" + transcriptSkillsSystemMarker + "\nboiler"),
schema.UserMessage("hello"), schema.UserMessage("hello"),
schema.AssistantMessage("reply", nil), schema.AssistantMessage("reply", nil),
} }
@@ -20,7 +20,9 @@ const (
transcriptStaticSystemOmitNote = "[static system prompt omitted — unchanged in live context after compaction]" transcriptStaticSystemOmitNote = "[static system prompt omitted — unchanged in live context after compaction]"
transcriptToolIndexStartMarker = "以下是当前会话绑定的工具名称索引" transcriptToolIndexStartMarker = "以下是当前会话绑定的工具名称索引"
transcriptPersonaStartMarker = "你是CyberStrikeAI" transcriptPersonaStartMarker = "你是CyberStrikeAI"
transcriptSkillsSystemMarker = "# Skills System" // ADK LanguageChinese injects skill middleware prompt with this header (see eino adk/middlewares/skill/prompt.go).
transcriptSkillsSystemMarker = "# Skill 系统"
transcriptSkillsSystemMarkerEnglish = "# Skills System"
) )
type transcriptToolCall struct { type transcriptToolCall struct {
@@ -86,13 +88,23 @@ func stripToolNamesIndexFromSystem(s string) string {
} }
func stripSkillsSystemBoilerplate(s string) string { func stripSkillsSystemBoilerplate(s string) string {
idx := strings.Index(s, transcriptSkillsSystemMarker) idx := indexFirstSubstring(s, transcriptSkillsSystemMarker, transcriptSkillsSystemMarkerEnglish)
if idx < 0 { if idx < 0 {
return strings.TrimSpace(s) return strings.TrimSpace(s)
} }
return strings.TrimSpace(s[:idx]) return strings.TrimSpace(s[:idx])
} }
func indexFirstSubstring(s string, markers ...string) int {
first := -1
for _, m := range markers {
if i := strings.Index(s, m); i >= 0 && (first < 0 || i < first) {
first = i
}
}
return first
}
func extractProjectBlackboardSection(s string) string { func extractProjectBlackboardSection(s string) string {
start := strings.Index(s, project.FactIndexSectionStartMarker) start := strings.Index(s, project.FactIndexSectionStartMarker)
if start < 0 { if start < 0 {
+12 -15
View File
@@ -190,29 +190,26 @@ func einoMessagesForRunRestart(args *einoADKRunLoopArgs, baseMsgs, accumulated [
return append([]adk.Message(nil), baseMsgs...), einoRestartContextInitial return append([]adk.Message(nil), baseMsgs...), einoRestartContextInitial
} }
// adkMessagesHasUserContent 从尾部向前查找,是否已有与 want 相同的 user 消息(避免重复 append)。 // adkMessagesHasUserContent reports whether the conversation tail is already a user turn
// with the given content. Only the last message counts: matching text in an earlier round
// (e.g. user repeats the same prompt after an assistant reply) must not suppress appending
// the new user turn — Claude 4.6+ rejects requests whose final message is assistant.
func adkMessagesHasUserContent(msgs []adk.Message, want string) bool { func adkMessagesHasUserContent(msgs []adk.Message, want string) bool {
want = strings.TrimSpace(want) want = strings.TrimSpace(want)
if want == "" { if want == "" {
return true return true
} }
for i := len(msgs) - 1; i >= 0; i-- { if len(msgs) == 0 {
m := msgs[i] return false
if m == nil {
continue
}
if m.Role == schema.User {
return strings.TrimSpace(m.Content) == want
}
if m.Role == schema.Assistant || m.Role == schema.Tool {
continue
}
break
} }
return false last := msgs[len(msgs)-1]
if last == nil || last.Role != schema.User {
return false
}
return strings.TrimSpace(last.Content) == want
} }
// appendUserMessageIfNeeded 在 history 轨迹之后追加本轮 user 消息(仅当轨迹中尚未包含该句)。 // appendUserMessageIfNeeded 在 history 轨迹之后追加本轮 user 消息(仅当尾部已是相同 user 句)。
func appendUserMessageIfNeeded(msgs []adk.Message, userMessage string) []adk.Message { func appendUserMessageIfNeeded(msgs []adk.Message, userMessage string) []adk.Message {
if strings.TrimSpace(userMessage) == "" || adkMessagesHasUserContent(msgs, userMessage) { if strings.TrimSpace(userMessage) == "" || adkMessagesHasUserContent(msgs, userMessage) {
return msgs return msgs
@@ -143,3 +143,18 @@ func TestAppendUserMessageIfNeeded(t *testing.T) {
t.Fatalf("should not duplicate user message: len=%d", len(dup)) t.Fatalf("should not duplicate user message: len=%d", len(dup))
} }
} }
func TestAppendUserMessageIfNeeded_repeatPromptAfterAssistant(t *testing.T) {
t.Parallel()
msgs := []adk.Message{
schema.UserMessage("扫描 example.com"),
schema.AssistantMessage("开始扫描...", nil),
}
out := appendUserMessageIfNeeded(msgs, "扫描 example.com")
if len(out) != 3 {
t.Fatalf("should append new user turn after assistant reply: len=%d", len(out))
}
if out[2].Role != schema.User || out[2].Content != "扫描 example.com" {
t.Fatalf("tail should be repeated user prompt, got role=%s content=%q", out[2].Role, out[2].Content)
}
}
+23
View File
@@ -0,0 +1,23 @@
package multiagent
import (
"context"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
// literalInstructionGenModelInput passes Instruction through as a system message without
// FString template formatting. Eino defaultGenModelInput formats instruction whenever
// SessionValues exist; prompts with literal curly braces (project blackboard "{关系边: ...}",
// JSON examples, link syntax) then fail with "could not find key".
//
// Matches eino/adk/prebuilt/deep genModelInput — the supported fix per Eino docs.
func literalInstructionGenModelInput(ctx context.Context, instruction string, input *adk.AgentInput) ([]adk.Message, error) {
msgs := make([]adk.Message, 0, len(input.Messages)+1)
if instruction != "" {
msgs = append(msgs, schema.SystemMessage(instruction))
}
msgs = append(msgs, input.Messages...)
return msgs, nil
}
@@ -0,0 +1,33 @@
package multiagent
import (
"context"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestLiteralInstructionGenModelInput_PreservesLiteralCurlyBraces(t *testing.T) {
t.Parallel()
instruction := "- [finding/x] summary {关系边: discovered_on←target/dev}\n" +
"如 finding 上 {from:target/*, type:discovered_on}"
msgs, err := literalInstructionGenModelInput(context.Background(), instruction, &adk.AgentInput{
Messages: []adk.Message{schema.UserMessage("继续")},
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(msgs) != 2 {
t.Fatalf("expected 2 messages, got %d", len(msgs))
}
if msgs[0].Role != schema.System {
t.Fatalf("first message must be system, got %s", msgs[0].Role)
}
for _, want := range []string{"{关系边:", "{from:target/*, type:discovered_on}"} {
if !strings.Contains(msgs[0].Content, want) {
t.Fatalf("system content missing %q: %q", want, msgs[0].Content)
}
}
}
+3 -5
View File
@@ -3,7 +3,6 @@ package multiagent
import ( import (
"context" "context"
"errors" "errors"
"fmt"
"strings" "strings"
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
@@ -75,8 +74,8 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
if err != nil { if err != nil {
if IsHumanRejectError(err) { if IsHumanRejectError(err) {
// Human rejection should be a soft tool result so the model can continue iterating. // Human rejection should be a soft tool result so the model can continue iterating.
msg := fmt.Sprintf("[HITL Reject] Tool '%s' was rejected by human reviewer. Reason: %s\nPlease adjust parameters/plan and continue without this call.", // tool_search 须保持 JSON,否则 Eino toolsearch 中间件解析历史时会硬崩 ChatModel。
input.Name, strings.TrimSpace(err.Error())) msg := HitlRejectToolResult(input.Name, err.Error())
// transfer_to_agent 在 Eino 中标记为 returnDirectly:工具成功后 ReAct 子图会直接 END, // transfer_to_agent 在 Eino 中标记为 returnDirectly:工具成功后 ReAct 子图会直接 END,
// 并依赖真实工具内的 SendToolGenAction 触发移交。HITL 拒绝时不会执行真实工具, // 并依赖真实工具内的 SendToolGenAction 触发移交。HITL 拒绝时不会执行真实工具,
// 若仍走 returnDirectly 分支,监督者会在无 Transfer 动作的情况下结束,模型不再迭代。 // 若仍走 returnDirectly 分支,监督者会在无 Transfer 动作的情况下结束,模型不再迭代。
@@ -103,8 +102,7 @@ func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
edited, err := fn(ctx, input.Name, input.Arguments) edited, err := fn(ctx, input.Name, input.Arguments)
if err != nil { if err != nil {
if IsHumanRejectError(err) { if IsHumanRejectError(err) {
msg := fmt.Sprintf("[HITL Reject] Tool '%s' was rejected by human reviewer. Reason: %s\nPlease adjust parameters/plan and continue without this call.", msg := HitlRejectToolResult(input.Name, err.Error())
input.Name, strings.TrimSpace(err.Error()))
hitlClearReturnDirectlyIfTransfer(ctx, input.Name) hitlClearReturnDirectlyIfTransfer(ctx, input.Name)
return &compose.StreamToolOutput{ return &compose.StreamToolOutput{
Result: schema.StreamReaderFromArray([]string{msg}), Result: schema.StreamReaderFromArray([]string{msg}),
@@ -0,0 +1,85 @@
package multiagent
import (
"encoding/json"
"fmt"
"strings"
)
const toolSearchToolName = "tool_search"
// HitlExemptMetaTools 为编排/元工具:不直接执行攻击动作,但会阻塞 agent 控制流。
// tool_search 必须免审批,否则其 HITL 拒绝结果与 Eino toolsearch 中间件不兼容(会硬崩 ChatModel)。
var HitlExemptMetaTools = []string{
toolSearchToolName,
"skill",
"task",
"write_todos",
"transfer_to_agent",
"exit",
"TaskCreate",
"TaskGet",
"TaskUpdate",
"TaskList",
}
// IsToolSearchTool reports whether name is the Eino dynamictool tool_search meta-tool.
func IsToolSearchTool(name string) bool {
return strings.EqualFold(strings.TrimSpace(name), toolSearchToolName)
}
// MergeHitlExemptMetaTools unions configured whitelist with built-in meta-tool exemptions.
func MergeHitlExemptMetaTools(configured []string) []string {
merged := make([]string, 0, len(configured)+len(HitlExemptMetaTools))
seen := make(map[string]struct{}, len(configured)+len(HitlExemptMetaTools))
add := func(name string) {
n := strings.ToLower(strings.TrimSpace(name))
if n == "" {
return
}
if _, ok := seen[n]; ok {
return
}
seen[n] = struct{}{}
merged = append(merged, strings.TrimSpace(name))
}
for _, t := range configured {
add(t)
}
for _, t := range HitlExemptMetaTools {
add(t)
}
return merged
}
type toolSearchHitlRejectPayload struct {
SelectedTools []string `json:"selectedTools"`
HitlRejected bool `json:"_hitlRejected"`
Reason string `json:"reason"`
}
// HitlRejectToolResult returns a tool result body safe for downstream consumers.
// tool_search must stay JSON-shaped so toolsearch.extractSelectedTools does not terminate the graph.
func HitlRejectToolResult(toolName, reason string) string {
reason = strings.TrimSpace(reason)
if !IsToolSearchTool(toolName) {
if reason == "" {
reason = "rejected by reviewer"
}
return fmt.Sprintf("[HITL Reject] Tool '%s' was rejected by reviewer. Reason: %s\nPlease adjust parameters/plan and continue without this call.",
strings.TrimSpace(toolName), reason)
}
payload := toolSearchHitlRejectPayload{
SelectedTools: []string{},
HitlRejected: true,
Reason: reason,
}
if payload.Reason == "" {
payload.Reason = "tool_search rejected by reviewer; no dynamic tools unlocked"
}
out, err := json.Marshal(payload)
if err != nil {
return `{"selectedTools":[],"_hitlRejected":true,"reason":"tool_search rejected by reviewer"}`
}
return string(out)
}
@@ -0,0 +1,48 @@
package multiagent
import (
"encoding/json"
"strings"
"testing"
)
func TestHitlRejectToolResult_toolSearchIsJSON(t *testing.T) {
raw := HitlRejectToolResult("tool_search", "rejected by user: timeout")
var payload toolSearchHitlRejectPayload
if err := json.Unmarshal([]byte(raw), &payload); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if len(payload.SelectedTools) != 0 {
t.Fatalf("expected empty selectedTools, got %v", payload.SelectedTools)
}
if !payload.HitlRejected {
t.Fatal("expected _hitlRejected true")
}
if !strings.Contains(payload.Reason, "timeout") {
t.Fatalf("reason=%q", payload.Reason)
}
}
func TestHitlRejectToolResult_otherToolKeepsLegacyText(t *testing.T) {
raw := HitlRejectToolResult("nmap", "too risky")
if strings.HasPrefix(raw, "{") {
t.Fatalf("expected legacy text, got %q", raw)
}
if !strings.HasPrefix(raw, "[HITL Reject]") {
t.Fatalf("expected [HITL Reject] prefix, got %q", raw)
}
}
func TestMergeHitlExemptMetaTools_includesToolSearch(t *testing.T) {
merged := MergeHitlExemptMetaTools([]string{"read_file"})
found := false
for _, name := range merged {
if IsToolSearchTool(name) {
found = true
break
}
}
if !found {
t.Fatalf("tool_search missing from %v", merged)
}
}
@@ -0,0 +1,85 @@
package multiagent
import (
"context"
"fmt"
"testing"
"cyberstrike-ai/internal/config"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/components/tool"
)
type stubChatModelAgentMiddleware struct {
adk.BaseChatModelAgentMiddleware
tag string
}
func stubMW(tag string) adk.ChatModelAgentMiddleware {
return &stubChatModelAgentMiddleware{tag: tag}
}
func TestBuildPlanExecuteExecutorHandlers_IncludesExecPreMiddlewares(t *testing.T) {
t.Parallel()
pre := []adk.ChatModelAgentMiddleware{
stubMW("patch"),
stubMW("reduction"),
}
got, err := buildPlanExecuteExecutorHandlers(context.Background(), &PlanExecuteRootArgs{
ExecPreMiddlewares: pre,
FilesystemMiddleware: stubMW("filesystem"),
SkillMiddleware: stubMW("skill"),
})
if err != nil {
t.Fatalf("buildPlanExecuteExecutorHandlers: %v", err)
}
if len(got) != 4 {
t.Fatalf("expected 4 pre-tail handlers (2 pre + fs + skill), got %d", len(got))
}
for i, want := range []string{"patch", "reduction", "filesystem", "skill"} {
st, ok := got[i].(*stubChatModelAgentMiddleware)
if !ok || st.tag != want {
t.Fatalf("handler[%d]: got %#v want tag %q", i, got[i], want)
}
}
}
func stubTools(n int) []tool.BaseTool {
out := make([]tool.BaseTool, n)
for i := 0; i < n; i++ {
out[i] = stubTool{name: fmt.Sprintf("t%d", i)}
}
return out
}
func TestBuildPlanExecuteExecutorHandlers_NilArgs(t *testing.T) {
t.Parallel()
if _, err := buildPlanExecuteExecutorHandlers(context.Background(), nil); err == nil {
t.Fatal("expected error for nil args")
}
}
func TestPrependEinoMiddlewares_Main_IncludesPatch(t *testing.T) {
t.Parallel()
ctx := context.Background()
mw := configMultiAgentEinoMiddlewareForTest()
mw.ReductionEnable = false
mw.ToolSearchEnable = false
mw.PlantaskEnable = false
_, extra, _, err := prependEinoMiddlewares(ctx, mw, einoMWMain, stubTools(25), nil, "", "conv-test", "", nil)
if err != nil {
t.Fatalf("prependEinoMiddlewares: %v", err)
}
if len(extra) == 0 {
t.Fatal("expected patch middleware on einoMWMain when patch_tool_calls enabled")
}
}
func configMultiAgentEinoMiddlewareForTest() *config.MultiAgentEinoMiddlewareConfig {
patch := true
return &config.MultiAgentEinoMiddlewareConfig{
PatchToolCalls: &patch,
}
}
+33 -7
View File
@@ -254,10 +254,11 @@ func RunDeepAgent(
) )
} }
sa, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ sa, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
Name: id, Name: id,
Description: desc, Description: desc,
Instruction: subInstrFinal, Instruction: subInstrFinal,
Model: subModel, GenModelInput: literalInstructionGenModelInput,
Model: subModel,
ToolsConfig: adk.ToolsConfig{ ToolsConfig: adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{ ToolsNodeConfig: compose.ToolsNodeConfig{
Tools: subToolsForCfg, Tools: subToolsForCfg,
@@ -372,8 +373,15 @@ func RunDeepAgent(
// noNestedTaskMiddleware 必须在最外层(最先拦截),防止 skill 或其他中间件内部触发 task 调用绕过检测。 // noNestedTaskMiddleware 必须在最外层(最先拦截),防止 skill 或其他中间件内部触发 task 调用绕过检测。
deepHandlers := []adk.ChatModelAgentMiddleware{newNoNestedTaskMiddleware()} deepHandlers := []adk.ChatModelAgentMiddleware{newNoNestedTaskMiddleware()}
taskEnrichExtra := systemPromptExtra var taskBlackboardSupplement string
if mw := newTaskContextEnrichMiddleware(userMessage, history, ma.SubAgentUserContextMaxRunes, taskEnrichExtra); mw != nil { if appCfg.Project.Enabled && db != nil {
if pid := strings.TrimSpace(projectID); pid != "" {
if block, err := project.BuildFactIndexBlock(db, pid, appCfg.Project); err == nil {
taskBlackboardSupplement = strings.TrimSpace(block)
}
}
}
if mw := newTaskContextEnrichMiddleware(userMessage, history, ma.SubAgentUserContextMaxRunesEffective(), taskBlackboardSupplement); mw != nil {
deepHandlers = append(deepHandlers, mw) deepHandlers = append(deepHandlers, mw)
} }
if len(mainOrchestratorPre) > 0 { if len(mainOrchestratorPre) > 0 {
@@ -424,6 +432,22 @@ func RunDeepAgent(
var da adk.Agent var da adk.Agent
switch orchMode { switch orchMode {
case "plan_execute": case "plan_execute":
plannerModelCfg := &einoopenai.ChatModelConfig{
APIKey: appCfg.OpenAI.APIKey,
BaseURL: strings.TrimSuffix(appCfg.OpenAI.BaseURL, "/"),
Model: appCfg.OpenAI.Model,
HTTPClient: httpClient,
}
reasoning.ApplyPlanExecutePlannerModelConfig(plannerModelCfg, &appCfg.OpenAI)
peMainModel, perr := einoopenai.NewChatModel(ctx, plannerModelCfg)
if perr != nil {
return nil, fmt.Errorf("plan_execute 规划模型: %w", perr)
}
if logger != nil {
logger.Info("plan_execute: planner/replanner 使用无 reasoning 的独立 ChatModelToolChoiceForced 兼容)",
zap.String("model", appCfg.OpenAI.Model),
)
}
execModel, perr := einoopenai.NewChatModel(ctx, baseModelCfg) execModel, perr := einoopenai.NewChatModel(ctx, baseModelCfg)
if perr != nil { if perr != nil {
return nil, fmt.Errorf("plan_execute 执行器模型: %w", perr) return nil, fmt.Errorf("plan_execute 执行器模型: %w", perr)
@@ -437,7 +461,7 @@ func RunDeepAgent(
} }
} }
peRoot, perr := NewPlanExecuteRoot(ctx, &PlanExecuteRootArgs{ peRoot, perr := NewPlanExecuteRoot(ctx, &PlanExecuteRootArgs{
MainToolCallingModel: mainModel, MainToolCallingModel: peMainModel,
ExecModel: execModel, ExecModel: execModel,
OrchInstruction: orchInstruction, OrchInstruction: orchInstruction,
ToolsCfg: mainToolsCfg, ToolsCfg: mainToolsCfg,
@@ -450,6 +474,7 @@ func RunDeepAgent(
ProjectID: projectID, ProjectID: projectID,
Logger: logger, Logger: logger,
ModelName: appCfg.OpenAI.Model, ModelName: appCfg.OpenAI.Model,
// 与 Deep/Supervisor 主代理同源:patch / reduction / toolsearch / plantask(见 buildPlanExecuteExecutorHandlers)。
ExecPreMiddlewares: mainOrchestratorPre, ExecPreMiddlewares: mainOrchestratorPre,
SkillMiddleware: einoSkillMW, SkillMiddleware: einoSkillMW,
FilesystemMiddleware: peFsMw, FilesystemMiddleware: peFsMw,
@@ -472,6 +497,7 @@ func RunDeepAgent(
Name: orchestratorName, Name: orchestratorName,
Description: orchDescription, Description: orchDescription,
Instruction: supInstr, Instruction: supInstr,
GenModelInput: literalInstructionGenModelInput,
Model: mainModel, Model: mainModel,
ToolsConfig: mainToolsCfg, ToolsConfig: mainToolsCfg,
MaxIterations: deepMaxIter, MaxIterations: deepMaxIter,
+12 -9
View File
@@ -3,6 +3,7 @@ package multiagent
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"fmt"
"strings" "strings"
"cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/agent"
@@ -11,7 +12,7 @@ import (
"github.com/cloudwego/eino/components/tool" "github.com/cloudwego/eino/components/tool"
) )
const defaultSubAgentUserContextMaxRunes = 2000 const userContextSupplementHeader = "\n\n## 用户历史输入(原文,子代理必读)\n"
// taskContextEnrichMiddleware intercepts "task" tool calls on the orchestrator // taskContextEnrichMiddleware intercepts "task" tool calls on the orchestrator
// and appends the user's original conversation messages to the task description. // and appends the user's original conversation messages to the task description.
@@ -30,13 +31,14 @@ type taskContextEnrichMiddleware struct {
// newTaskContextEnrichMiddleware returns a middleware that enriches task // newTaskContextEnrichMiddleware returns a middleware that enriches task
// descriptions with user conversation context. Returns nil if disabled // descriptions with user conversation context. Returns nil if disabled
// (maxRunes < 0) or no user messages exist. // (maxRunes < 0) or no user messages exist.
// projectBlackboard 仅传项目黑板索引块(BuildFactIndexBlock);勿传完整 systemPromptExtra。
func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int, projectBlackboard string) adk.ChatModelAgentMiddleware { func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int, projectBlackboard string) adk.ChatModelAgentMiddleware {
supplement := buildUserContextSupplement(userMessage, history, maxRunes) supplement := buildUserContextSupplement(userMessage, history, maxRunes)
if bb := strings.TrimSpace(projectBlackboard); bb != "" { if bb := strings.TrimSpace(projectBlackboard); bb != "" {
if supplement != "" { if supplement != "" {
supplement += "\n\n## 项目黑板索引\n" + bb supplement += "\n\n" + bb
} else { } else {
supplement = "\n\n## 项目黑板索引\n" + bb supplement = "\n\n" + bb
} }
} }
if supplement == "" { if supplement == "" {
@@ -86,9 +88,6 @@ func buildUserContextSupplement(userMessage string, history []agent.ChatMessage,
if maxRunes < 0 { if maxRunes < 0 {
return "" return ""
} }
if maxRunes == 0 {
maxRunes = defaultSubAgentUserContextMaxRunes
}
var userMsgs []string var userMsgs []string
for _, h := range history { for _, h := range history {
@@ -107,12 +106,16 @@ func buildUserContextSupplement(userMessage string, history []agent.ChatMessage,
return "" return ""
} }
joined := strings.Join(userMsgs, "\n---\n") lines := make([]string, 0, len(userMsgs))
if len([]rune(joined)) > maxRunes { for i, msg := range userMsgs {
lines = append(lines, fmt.Sprintf("[第%d轮] %s", i+1, msg))
}
joined := strings.Join(lines, "\n")
if maxRunes > 0 && len([]rune(joined)) > maxRunes {
joined = truncateKeepFirstLast(userMsgs, maxRunes) joined = truncateKeepFirstLast(userMsgs, maxRunes)
} }
return "\n\n## 会话上下文(自动补充,确保你了解用户完整意图)\n" + joined return userContextSupplementHeader + joined
} }
// truncateKeepFirstLast keeps the first and last user messages, giving each // truncateKeepFirstLast keeps the first and last user messages, giving each
@@ -74,7 +74,7 @@ func TestBuildUserContextSupplement_DisabledByNegative(t *testing.T) {
func TestBuildUserContextSupplement_CustomMaxRunes(t *testing.T) { func TestBuildUserContextSupplement_CustomMaxRunes(t *testing.T) {
msg := strings.Repeat("A", 200) msg := strings.Repeat("A", 200)
result := buildUserContextSupplement(msg, nil, 50) result := buildUserContextSupplement(msg, nil, 50)
header := "\n\n## 会话上下文(自动补充,确保你了解用户完整意图)\n" header := userContextSupplementHeader
body := strings.TrimPrefix(result, header) body := strings.TrimPrefix(result, header)
if len([]rune(body)) > 50 { if len([]rune(body)) > 50 {
t.Errorf("body should be capped at 50 runes, got %d", len([]rune(body))) t.Errorf("body should be capped at 50 runes, got %d", len([]rune(body)))
@@ -89,7 +89,7 @@ func TestBuildUserContextSupplement_TruncateKeepsFirstAndLast(t *testing.T) {
history = append(history, agent.ChatMessage{Role: "user", Content: strings.Repeat("B", 500)}) history = append(history, agent.ChatMessage{Role: "user", Content: strings.Repeat("B", 500)})
} }
last := "最后一条指令" last := "最后一条指令"
result := buildUserContextSupplement(last, history, 0) result := buildUserContextSupplement(last, history, 800)
if !strings.Contains(result, "http://target.com") { if !strings.Contains(result, "http://target.com") {
t.Error("first message (target URL) should survive truncation") t.Error("first message (target URL) should survive truncation")
} }
+6 -3
View File
@@ -806,10 +806,12 @@ func isClaudeProvider(cfg *config.OpenAIConfig) bool {
// Eino HTTP Client Bridge // Eino HTTP Client Bridge
// ============================================================ // ============================================================
// NewEinoHTTPClient 为 einoopenai.ChatModelConfig 返回一个 http.Client,包含层 transport 包装: // NewEinoHTTPClient 为 einoopenai.ChatModelConfig 返回一个 http.Client,包含层 transport 包装:
// 1. 当 cfg.Provider 为 claude 时,最内层套 claudeRoundTripper,把 OpenAI /chat/completions 透明 // 1. 当 cfg.Provider 为 claude 时,套 claudeRoundTripper,把 OpenAI /chat/completions 透明
// 桥接为 Anthropic /v1/messages(并把 Claude SSE 翻译回 OpenAI SSE 格式)。 // 桥接为 Anthropic /v1/messages(并把 Claude SSE 翻译回 OpenAI SSE 格式)。
// 2. 最外层无条件套 einoSSESanitizingRoundTripper,吞掉中转站发的 SSE 心跳/注释/控制行 // 2. reasoningToolChoiceCompatRoundTrippertool_choice=required/object 时剥离 thinking 字段,避免
// plan_execute replanner 等强制工具调用与推理模式冲突(部分网关返回 400)。
// 3. 最外层无条件套 einoSSESanitizingRoundTripper,吞掉中转站发的 SSE 心跳/注释/控制行
// (": keepalive" / "event: ping" / "retry: 3000" 等),避免 Eino 用的 meguminnnnnnnnn/go-openai // (": keepalive" / "event: ping" / "retry: 3000" 等),避免 Eino 用的 meguminnnnnnnnn/go-openai
// SDK 在累计超过 300 个非 "data:" 行后抛 "stream has sent too many empty messages"。 // SDK 在累计超过 300 个非 "data:" 行后抛 "stream has sent too many empty messages"。
// //
@@ -825,6 +827,7 @@ func NewEinoHTTPClient(cfg *config.OpenAIConfig, base *http.Client) *http.Client
if transport == nil { if transport == nil {
transport = http.DefaultTransport transport = http.DefaultTransport
} }
transport = &reasoningToolChoiceCompatRoundTripper{base: transport}
if isClaudeProvider(cfg) { if isClaudeProvider(cfg) {
transport = &claudeRoundTripper{ transport = &claudeRoundTripper{
base: transport, base: transport,
+79
View File
@@ -0,0 +1,79 @@
package openai
import (
"github.com/bytedance/sonic"
)
// reasoningPayloadKeys are OpenAI-compatible root fields that enable "thinking" /
// extended-reasoning modes on gateways such as DashScope/Qwen and MiniMax.
var reasoningPayloadKeys = []string{
"thinking",
"reasoning_effort",
"output_config",
"reasoning",
}
// StripReasoningFromChatCompletionBody removes thinking / reasoning fields from a
// chat-completions JSON body.
func StripReasoningFromChatCompletionBody(rawBody []byte) ([]byte, error) {
var payload map[string]any
if err := sonic.Unmarshal(rawBody, &payload); err != nil {
return rawBody, nil
}
if !stripReasoningFields(payload) {
return rawBody, nil
}
out, err := sonic.Marshal(payload)
if err != nil {
return rawBody, err
}
return out, nil
}
// StripReasoningIfForcedToolChoice removes thinking / reasoning fields when the
// request sets tool_choice to "required" or an object. Several providers reject
// that combination (e.g. DashScope: "tool_choice does not support being set to
// required or object in thinking mode").
func StripReasoningIfForcedToolChoice(rawBody []byte) ([]byte, error) {
var payload map[string]any
if err := sonic.Unmarshal(rawBody, &payload); err != nil {
return rawBody, nil
}
if !forcedToolChoiceIncompatibleWithThinking(payload) {
return rawBody, nil
}
if !stripReasoningFields(payload) {
return rawBody, nil
}
out, err := sonic.Marshal(payload)
if err != nil {
return rawBody, err
}
return out, nil
}
func stripReasoningFields(payload map[string]any) bool {
changed := false
for _, key := range reasoningPayloadKeys {
if _, ok := payload[key]; ok {
delete(payload, key)
changed = true
}
}
return changed
}
func forcedToolChoiceIncompatibleWithThinking(payload map[string]any) bool {
tc, ok := payload["tool_choice"]
if !ok || tc == nil {
return false
}
switch v := tc.(type) {
case string:
return v == "required"
case map[string]any:
return true
default:
return false
}
}
+120
View File
@@ -0,0 +1,120 @@
package openai
import (
"io"
"net/http"
"strings"
"testing"
)
func TestStripReasoningFromChatCompletionBody(t *testing.T) {
in := []byte(`{"model":"deepseek-chat","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`)
out, err := StripReasoningFromChatCompletionBody(in)
if err != nil {
t.Fatal(err)
}
s := string(out)
if strings.Contains(s, "thinking") || strings.Contains(s, "reasoning_effort") {
t.Fatalf("expected reasoning fields stripped, got %s", s)
}
if !strings.Contains(s, `"model":"deepseek-chat"`) {
t.Fatalf("expected model preserved, got %s", s)
}
plain := []byte(`{"model":"gpt-4o","messages":[]}`)
out2, err := StripReasoningFromChatCompletionBody(plain)
if err != nil {
t.Fatal(err)
}
if string(out2) != string(plain) {
t.Fatalf("expected unchanged payload, got %s", out2)
}
}
func TestStripReasoningIfForcedToolChoice(t *testing.T) {
cases := []struct {
name string
in string
strip bool
contain string
}{
{
name: "required strips thinking",
in: `{"model":"minimax","messages":[],"thinking":{"type":"enabled"},"tool_choice":"required","tools":[]}`,
strip: true,
},
{
name: "object tool_choice strips thinking",
in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"},"tool_choice":{"type":"function","function":{"name":"respond"}}}`,
strip: true,
},
{
name: "auto keeps thinking",
in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"},"tool_choice":"auto"}`,
strip: false,
contain: "thinking",
},
{
name: "no tool_choice keeps thinking",
in: `{"model":"qwen","messages":[],"thinking":{"type":"enabled"}}`,
strip: false,
contain: "thinking",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
out, err := StripReasoningIfForcedToolChoice([]byte(tc.in))
if err != nil {
t.Fatal(err)
}
s := string(out)
hasThinking := strings.Contains(s, "thinking")
if tc.strip && hasThinking {
t.Fatalf("expected thinking stripped, got %s", s)
}
if !tc.strip && tc.contain != "" && !strings.Contains(s, tc.contain) {
t.Fatalf("expected %q in %s", tc.contain, s)
}
if !tc.strip && string(out) != tc.in {
t.Fatalf("expected unchanged payload, got %s", s)
}
})
}
}
func TestReasoningToolChoiceCompatRoundTripper(t *testing.T) {
var gotBody string
rt := &reasoningToolChoiceCompatRoundTripper{
base: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
b, _ := io.ReadAll(req.Body)
gotBody = string(b)
return &http.Response{
StatusCode: 200,
Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)),
Header: http.Header{"Content-Type": []string{"application/json"}},
}, nil
}),
}
req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/chat/completions", strings.NewReader(
`{"model":"m","thinking":{"type":"enabled"},"tool_choice":"required","messages":[]}`,
))
if err != nil {
t.Fatal(err)
}
_, err = rt.RoundTrip(req)
if err != nil {
t.Fatal(err)
}
if strings.Contains(gotBody, "thinking") {
t.Fatalf("expected thinking stripped in transit, got %s", gotBody)
}
if !strings.Contains(gotBody, `"tool_choice":"required"`) {
t.Fatalf("expected tool_choice preserved, got %s", gotBody)
}
}
type roundTripperFunc func(*http.Request) (*http.Response, error)
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
@@ -0,0 +1,43 @@
package openai
import (
"bytes"
"io"
"net/http"
"strconv"
"strings"
)
// reasoningToolChoiceCompatRoundTripper strips thinking/reasoning fields from
// chat/completions requests that force tool_choice, which some gateways reject
// when thinking mode is enabled on the same request.
type reasoningToolChoiceCompatRoundTripper struct {
base http.RoundTripper
}
func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if rt == nil || rt.base == nil || req == nil || req.Body == nil {
if rt != nil && rt.base != nil {
return rt.base.RoundTrip(req)
}
return http.DefaultTransport.RoundTrip(req)
}
if req.Method != http.MethodPost || !strings.HasSuffix(req.URL.Path, "/chat/completions") {
return rt.base.RoundTrip(req)
}
body, err := io.ReadAll(req.Body)
_ = req.Body.Close()
if err != nil {
return nil, err
}
patched, perr := StripReasoningIfForcedToolChoice(body)
if perr != nil {
patched = body
}
req.Body = io.NopCloser(bytes.NewReader(patched))
req.ContentLength = int64(len(patched))
req.Header.Set("Content-Length", strconv.Itoa(len(patched)))
return rt.base.RoundTrip(req)
}
+29
View File
@@ -26,6 +26,35 @@ const (
wireOutputConfig wireOutputConfig
) )
// ApplyPlanExecutePlannerModelConfig configures the plan_execute planner/replanner
// ChatModel. Those Eino agents call WithToolChoice(Forced); several gateways reject
// thinking / reasoning fields on the same request (tool_choice required/object).
// Executor should keep the normal ApplyToEinoChatModelConfig path.
func ApplyPlanExecutePlannerModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig) {
if cfg == nil || oa == nil {
return
}
offOA := *oa
offReasoning := oa.Reasoning
offReasoning.Mode = "off"
offOA.Reasoning = offReasoning
ApplyToEinoChatModelConfig(cfg, &offOA, nil)
clearReasoningFromChatModelConfig(cfg)
}
func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) {
if cfg == nil {
return
}
cfg.ReasoningEffort = ""
if cfg.ExtraFields != nil {
for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} {
delete(cfg.ExtraFields, key)
}
}
applyThinkingDisabled(cfg)
}
// ApplyToEinoChatModelConfig merges reasoning-related options into cfg. // ApplyToEinoChatModelConfig merges reasoning-related options into cfg.
// Precondition: cfg already has APIKey, BaseURL, Model, HTTPClient set. // Precondition: cfg already has APIKey, BaseURL, Model, HTTPClient set.
func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig, client *ClientIntent) { func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig, client *ClientIntent) {
+24
View File
@@ -49,6 +49,30 @@ func TestApplyOpenAICompat_xhighExtraField(t *testing.T) {
} }
} }
func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
BaseURL: "https://antchat.example.com/v1",
Model: "minimax-m3",
Reasoning: config.OpenAIReasoningConfig{
Profile: "openai_compat",
Mode: "on",
Effort: "high",
},
}
ApplyPlanExecutePlannerModelConfig(cfg, oa)
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort cleared, got %q", cfg.ReasoningEffort)
}
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "disabled" {
t.Fatalf("expected thinking disabled, got %#v", cfg.ExtraFields)
}
if _, ok := cfg.ExtraFields["reasoning_effort"]; ok {
t.Fatalf("expected reasoning_effort stripped, got %#v", cfg.ExtraFields)
}
}
func TestApplyReasoningOff_disablesThinking(t *testing.T) { func TestApplyReasoningOff_disablesThinking(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{} cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{ oa := &config.OpenAIConfig{
+83 -91
View File
@@ -162,9 +162,8 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
output, err = runCommandWithPTY(ctx, cmd2, cb) output, err = runCommandWithPTY(ctx, cmd2, cb)
} }
} else { } else {
outputBytes, err2 := cmd.CombinedOutput() // 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。
output = string(outputBytes) output, err = combinedOutputCancellable(ctx, cmd)
err = err2
if err != nil && shouldRetryWithPTY(output) { if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试", e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
zap.String("tool", toolName), zap.String("tool", toolName),
@@ -692,83 +691,21 @@ func (e *Executor) formatParamValue(param config.ParameterConfig, value interfac
// IsBackgroundShellCommand 检测命令是否为完全后台命令(末尾有独立 &,且不在引号内)。 // IsBackgroundShellCommand 检测命令是否为完全后台命令(末尾有独立 &,且不在引号内)。
// command1 & command2 不算完全后台(command2 仍在前台执行)。 // command1 & command2 不算完全后台(command2 仍在前台执行)。
func IsBackgroundShellCommand(command string) bool { func IsBackgroundShellCommand(command string) bool {
// 移除首尾空格
command = strings.TrimSpace(command) command = strings.TrimSpace(command)
if command == "" { if command == "" {
return false return false
} }
positions := findStandaloneAmpersandPositions(command)
// 检查命令中所有不在引号内的 & 符号 if len(positions) == 0 {
// 找到最后一个 & 符号,检查它是否在命令末尾
inSingleQuote := false
inDoubleQuote := false
escaped := false
lastAmpersandPos := -1
for i, r := range command {
if escaped {
escaped = false
continue
}
if r == '\\' {
escaped = true
continue
}
if r == '\'' && !inDoubleQuote {
inSingleQuote = !inSingleQuote
continue
}
if r == '"' && !inSingleQuote {
inDoubleQuote = !inDoubleQuote
continue
}
if r == '&' && !inSingleQuote && !inDoubleQuote {
// 检查 & 前后是否有空格或换行(确保是独立的 &,而不是变量名的一部分)
isStandalone := false
// 检查前面:空格、制表符、换行符,或者是命令开头
if i == 0 {
isStandalone = true
} else {
prev := command[i-1]
if prev == ' ' || prev == '\t' || prev == '\n' || prev == '\r' {
isStandalone = true
}
}
// 检查后面:空格、制表符、换行符,或者是命令末尾
if isStandalone {
if i == len(command)-1 {
// 在末尾,肯定是独立的 &
lastAmpersandPos = i
} else {
next := command[i+1]
if next == ' ' || next == '\t' || next == '\n' || next == '\r' {
// 后面有空格,是独立的 &
lastAmpersandPos = i
}
}
}
}
}
// 如果没有找到 & 符号,不是后台命令
if lastAmpersandPos == -1 {
return false return false
} }
last := positions[len(positions)-1]
// 检查最后一个 & 后面是否还有非空内容 afterAmpersand := strings.TrimSpace(command[last+1:])
afterAmpersand := strings.TrimSpace(command[lastAmpersandPos+1:]) if afterAmpersand != "" {
if afterAmpersand == "" { return false
// & 在末尾或后面只有空白字符,这是完全后台命令
// 检查 & 前面是否有内容
beforeAmpersand := strings.TrimSpace(command[:lastAmpersandPos])
return beforeAmpersand != ""
} }
beforeAmpersand := strings.TrimSpace(command[:last])
// 如果 & 后面还有非空内容,说明是 command1 & command2 的情况 return beforeAmpersand != ""
// 这种情况下,command2会在前台执行,所以不算完全后台命令
return false
} }
// executeSystemCommand 执行系统命令 // executeSystemCommand 执行系统命令
@@ -804,7 +741,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
zap.String("command", command), zap.String("command", command),
) )
command = PrepareNonInteractiveShellCommand(command) command = PrepareShellCommandForExecute(command)
// 获取shell类型(可选,默认为sh) // 获取shell类型(可选,默认为sh)
shell := "sh" shell := "sh"
@@ -845,10 +782,8 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&") commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&")
commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand) commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand)
// 构建新命令:将用户命令置于独立重定向的后台作业,再 echo $pid // 构建新命令:后台作业重定向标准流后 echo $pid(与 RedirectBackgroundJobStdio 一致)
// 若子进程与 echo 共享同一 stdout 管道,且长时间不向 stdout 写入换行, pidCommand := RedirectBackgroundJobStdio(commandWithoutAmpersand+" &") + " pid=$!; echo $pid"
// bufio.ReadString('\n') 会永久阻塞(例如 beacon 持续写二进制/单行日志)。
pidCommand := fmt.Sprintf("%s </dev/null >/dev/null 2>&1 & pid=$!; echo $pid", commandWithoutAmpersand)
// 创建新命令来获取PID // 创建新命令来获取PID
var pidCmd *exec.Cmd var pidCmd *exec.Cmd
@@ -981,9 +916,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
output, err = runCommandWithPTY(ctx, cmd2, cb) output, err = runCommandWithPTY(ctx, cmd2, cb)
} }
} else { } else {
outputBytes, err2 := cmd.CombinedOutput() output, err = combinedOutputCancellable(ctx, cmd)
output = string(outputBytes)
err = err2
if err != nil && shouldRetryWithPTY(output) { if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试") e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
cmd2 := exec.CommandContext(ctx, shell, "-c", command) cmd2 := exec.CommandContext(ctx, shell, "-c", command)
@@ -1027,12 +960,58 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
}, nil }, nil
} }
// combinedOutputCancellable 行为对齐 cmd.CombinedOutputstdout/stderr 写入内存缓冲),
// 但在 ctx 取消时 terminateCmdTree 终止整棵进程树。
// 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。
// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。
func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) {
var stdoutBuf, stderrBuf strings.Builder
cmd.Stdout = &stdoutBuf
cmd.Stderr = &stderrBuf
session, err := StartShellSession(cmd)
if err != nil {
return "", err
}
done := make(chan error, 1)
go func() {
done <- session.Wait()
}()
stopWatch := make(chan struct{})
go func() {
select {
case <-ctx.Done():
TerminateShellCmdSession(session)
case <-stopWatch:
}
}()
defer close(stopWatch)
var waitErr error
select {
case waitErr = <-done:
case <-ctx.Done():
waitErr = <-done
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), ctx.Err()
}
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), waitErr
}
func joinCommandOutput(stdout, stderr string) string {
if stderr == "" {
return stdout
}
if stdout == "" {
return stderr
}
return stdout + stderr
}
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。 // streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。 // 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int) (string, error) { func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int) (string, error) {
if err := prepareShellCmdSession(cmd); err != nil {
return "", err
}
stdoutPipe, err := cmd.StdoutPipe() stdoutPipe, err := cmd.StdoutPipe()
if err != nil { if err != nil {
return "", err return "", err
@@ -1042,7 +1021,8 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
_ = stdoutPipe.Close() _ = stdoutPipe.Close()
return "", err return "", err
} }
if err := cmd.Start(); err != nil { session, err := StartShellSession(cmd)
if err != nil {
_ = stdoutPipe.Close() _ = stdoutPipe.Close()
_ = stderrPipe.Close() _ = stderrPipe.Close()
return "", err return "", err
@@ -1052,7 +1032,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
go func() { go func() {
select { select {
case <-ctx.Done(): case <-ctx.Done():
terminateCmdTree(cmd) TerminateShellCmdSession(session)
case <-stopWatch: case <-stopWatch:
} }
}() }()
@@ -1091,7 +1071,9 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
if deltaBuilder.Len() == 0 { if deltaBuilder.Len() == 0 {
return return
} }
cb(deltaBuilder.String()) if cb != nil {
cb(deltaBuilder.String())
}
deltaBuilder.Reset() deltaBuilder.Reset()
lastFlush = time.Now() lastFlush = time.Now()
} }
@@ -1102,13 +1084,13 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
} }
fireInactivity := func() { fireInactivity := func() {
terminateCmdTree(cmd) TerminateShellCmdSession(session)
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec) msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
outBuilder.WriteString(msg) outBuilder.WriteString(msg)
if cb != nil { if cb != nil {
cb(msg) cb(msg)
} }
_ = cmd.Wait() _ = session.Wait()
} }
chunksLoop: chunksLoop:
@@ -1118,6 +1100,11 @@ chunksLoop:
idleCh = idleWatch.Expired idleCh = idleWatch.Expired
} }
select { select {
case <-ctx.Done():
TerminateShellCmdSession(session)
flush()
_ = session.Wait()
return outBuilder.String(), ctx.Err()
case <-idleCh: case <-idleCh:
fireInactivity() fireInactivity()
return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
@@ -1138,7 +1125,7 @@ chunksLoop:
flush() flush()
// 等待命令结束,返回最终退出状态 // 等待命令结束,返回最终退出状态
waitErr := cmd.Wait() waitErr := session.Wait()
return outBuilder.String(), waitErr return outBuilder.String(), waitErr
} }
@@ -1210,13 +1197,18 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
} }
defer func() { _ = ptmx.Close() }() defer func() { _ = ptmx.Close() }()
rootPID := 0
if cmd.Process != nil {
rootPID = cmd.Process.Pid
}
// ctx 取消时尽快终止子进程 // ctx 取消时尽快终止子进程
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
select { select {
case <-ctx.Done(): case <-ctx.Done():
_ = ptmx.Close() // 触发读退出 _ = ptmx.Close() // 触发读退出
terminateCmdTree(cmd) terminateProcessGroup(rootPID, cmd)
case <-done: case <-done:
} }
}() }()
+32
View File
@@ -2,6 +2,8 @@ package security
import ( import (
"context" "context"
"os/exec"
"runtime"
"strings" "strings"
"testing" "testing"
"time" "time"
@@ -147,3 +149,33 @@ func indexOf(slice []string, s string) int {
} }
return -1 return -1
} }
// TestCombinedOutputCancellable_ContextCancelKillsTree 验证 ctx 取消时能在数秒内结束(杀进程组,非挂死)。
func TestCombinedOutputCancellable_ContextCancelKillsTree(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
done := make(chan error, 1)
go func() {
_, err := combinedOutputCancellable(ctx, cmd)
done <- err
}()
time.Sleep(150 * time.Millisecond)
cancel()
select {
case err := <-done:
if err == nil {
t.Fatal("expected context cancel error")
}
case <-time.After(5 * time.Second):
t.Fatal("combinedOutputCancellable did not return within 5s after context cancel")
}
}
+15 -5
View File
@@ -19,13 +19,23 @@ func prepareShellCmdSession(cmd *exec.Cmd) error {
return nil return nil
} }
// terminateCmdTree 尽力终止 cmd 及其进程组(Unix 下 Setsid 后 PGID == 首进程 PID // terminateProcessGroup 对 rootPID 对应进程组发 SIGKILLrootPID 为 0 时回退到 cmd.Process.Pid
func terminateCmdTree(cmd *exec.Cmd) { func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
if cmd == nil || cmd.Process == nil { pid := rootPID
if pid <= 0 && cmd != nil && cmd.Process != nil {
pid = cmd.Process.Pid
}
if pid <= 0 {
return return
} }
pid := cmd.Process.Pid
if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil { if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil {
_ = cmd.Process.Kill() if cmd != nil && cmd.Process != nil {
_ = cmd.Process.Kill()
}
} }
} }
// terminateCmdTree 尽力终止 cmd 及其进程组(Unix 下 Setsid 后 PGID == 首进程 PID)。
func terminateCmdTree(cmd *exec.Cmd) {
terminateProcessGroup(0, cmd)
}
+31 -5
View File
@@ -2,16 +2,42 @@
package security package security
import "os/exec" import (
"os/exec"
"strconv"
"syscall"
)
func prepareShellCmdSession(cmd *exec.Cmd) error { func prepareShellCmdSession(cmd *exec.Cmd) error {
_ = cmd if cmd == nil {
return nil
}
// 独立进程组,便于 taskkill /T 终止整棵子进程树。
if cmd.SysProcAttr == nil {
cmd.SysProcAttr = &syscall.SysProcAttr{}
}
cmd.SysProcAttr.CreationFlags = syscall.CREATE_NEW_PROCESS_GROUP
return nil return nil
} }
func terminateCmdTree(cmd *exec.Cmd) { // terminateProcessGroup 使用 taskkill /F /T 终止进程及其子进程;rootPID 为 0 时回退到 cmd.Process.Pid。
if cmd == nil || cmd.Process == nil { func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
pid := rootPID
if pid <= 0 && cmd != nil && cmd.Process != nil {
pid = cmd.Process.Pid
}
if pid <= 0 {
return return
} }
_ = cmd.Process.Kill() tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid))
if err := tk.Run(); err != nil {
if cmd != nil && cmd.Process != nil {
_ = cmd.Process.Kill()
}
}
}
// terminateCmdTree 使用 taskkill /F /T 终止进程及其子进程(Windows 上 Process.Kill 无法保证杀掉 python 等孙进程)。
func terminateCmdTree(cmd *exec.Cmd) {
terminateProcessGroup(0, cmd)
} }
+111
View File
@@ -0,0 +1,111 @@
package security
import "strings"
const backgroundJobStdioRedirect = " </dev/null >/dev/null 2>&1"
// findStandaloneAmpersandPositions 返回不在引号内的独立 & 下标(排除 &&)。
func findStandaloneAmpersandPositions(command string) []int {
command = strings.TrimSpace(command)
if command == "" {
return nil
}
var positions []int
inSingleQuote := false
inDoubleQuote := false
escaped := false
for i := 0; i < len(command); i++ {
r := command[i]
if escaped {
escaped = false
continue
}
if r == '\\' {
escaped = true
continue
}
if r == '\'' && !inDoubleQuote {
inSingleQuote = !inSingleQuote
continue
}
if r == '"' && !inSingleQuote {
inDoubleQuote = !inDoubleQuote
continue
}
if r != '&' || inSingleQuote || inDoubleQuote {
continue
}
if i+1 < len(command) && command[i+1] == '&' {
continue
}
if i > 0 && command[i-1] == '&' {
continue
}
isStandalone := i == 0
if !isStandalone {
prev := command[i-1]
isStandalone = prev == ' ' || prev == '\t' || prev == '\n' || prev == '\r'
}
if !isStandalone {
continue
}
if i == len(command)-1 {
positions = append(positions, i)
continue
}
next := command[i+1]
if next == ' ' || next == '\t' || next == '\n' || next == '\r' {
positions = append(positions, i)
}
}
return positions
}
func segmentHasStdioRedirect(segment string) bool {
lower := strings.ToLower(strings.TrimSpace(segment))
if lower == "" {
return false
}
if strings.Contains(lower, ">/dev/null") || strings.Contains(lower, "2>/dev/null") {
return true
}
if strings.Contains(lower, "&>") || strings.Contains(lower, "&>>") {
return true
}
if strings.Contains(lower, "2>&1") && strings.Contains(lower, "/dev/null") {
return true
}
return false
}
// RedirectBackgroundJobStdio 为每个独立 & 前的后台段注入 </dev/null >/dev/null 2>&1
// 避免后台子进程占用 execute/exec 管道导致挂死。
func RedirectBackgroundJobStdio(command string) string {
positions := findStandaloneAmpersandPositions(command)
if len(positions) == 0 {
return command
}
out := command
for j := len(positions) - 1; j >= 0; j-- {
i := positions[j]
before := out[:i]
after := out[i:]
trimmed := strings.TrimRight(before, " \t\r\n")
if segmentHasStdioRedirect(trimmed) {
continue
}
trailing := before[len(trimmed):]
out = trimmed + backgroundJobStdioRedirect + trailing + after
}
return out
}
// PrepareShellCommandForExecute 组合 execute/exec 用的非交互包装与后台 IO 重定向。
// 须先注入 exec </dev/null,再改写 & 后台段,否则段内 </dev/null 会使 stdin 重定向被误判为已存在。
func PrepareShellCommandForExecute(shellCommand string) string {
return RedirectBackgroundJobStdio(PrepareNonInteractiveShellCommand(shellCommand))
}
@@ -0,0 +1,64 @@
package security
import (
"strings"
"testing"
)
func TestRedirectBackgroundJobStdio_mixedCommand(t *testing.T) {
in := "java -jar app.jar & JRMP_PID=$!; echo started"
out := RedirectBackgroundJobStdio(in)
if !strings.Contains(out, "java -jar app.jar </dev/null >/dev/null 2>&1 &") {
t.Fatalf("expected redirect before &: %q", out)
}
if !strings.Contains(out, "echo started") {
t.Fatalf("foreground tail preserved: %q", out)
}
}
func TestRedirectBackgroundJobStdio_trailingOnly(t *testing.T) {
in := "sleep 120 &"
out := RedirectBackgroundJobStdio(in)
want := "sleep 120 </dev/null >/dev/null 2>&1 &"
if strings.TrimSpace(out) != want {
t.Fatalf("got %q want %q", out, want)
}
}
func TestRedirectBackgroundJobStdio_skipsAlreadyRedirected(t *testing.T) {
in := "sleep 1 >/dev/null 2>&1 & echo ok"
out := RedirectBackgroundJobStdio(in)
if out != in {
t.Fatalf("should not double-redirect: %q", out)
}
}
func TestRedirectBackgroundJobStdio_skipsAndAnd(t *testing.T) {
in := "test -f /etc/passwd && echo ok"
out := RedirectBackgroundJobStdio(in)
if out != in {
t.Fatalf("&& must not be treated as background &: %q", out)
}
}
func TestPrepareShellCommandForExecute(t *testing.T) {
out := PrepareShellCommandForExecute("java -jar x & echo hi")
if !strings.Contains(out, "exec </dev/null") {
t.Fatalf("missing stdin redirect: %q", out)
}
if !strings.Contains(out, "GIT_PAGER=cat") {
t.Fatalf("missing pager export: %q", out)
}
if !strings.Contains(out, "java -jar x </dev/null >/dev/null 2>&1 &") {
t.Fatalf("missing background redirect: %q", out)
}
}
func TestIsBackgroundShellCommand_usesSharedParser(t *testing.T) {
if !IsBackgroundShellCommand("sleep 1 &") {
t.Fatal("trailing & should be background")
}
if IsBackgroundShellCommand("sleep 1 & echo hi") {
t.Fatal("mixed should not be fully background")
}
}
+20 -9
View File
@@ -27,6 +27,11 @@ func TerminateShellCmdTree(cmd *exec.Cmd) {
terminateCmdTree(cmd) terminateCmdTree(cmd)
} }
// TerminateShellCmdSession 使用 Start 时缓存的进程组 ID 终止(shell 已退出时仍有效)。
func TerminateShellCmdSession(session *ShellSession) {
TerminateShellSession(session)
}
// EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell,行为与 exec 对齐: // EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell,行为与 exec 对齐:
// 并发读取 stdout/stderr(定长块,非按行),避免官方 local.ExecuteStreaming 先排空 stdout // 并发读取 stdout/stderr(定长块,非按行),避免官方 local.ExecuteStreaming 先排空 stdout
// 导致 stderr 错误(如 sudo 密码提示)长时间不可见、UI 一直显示「执行中」。 // 导致 stderr 错误(如 sudo 密码提示)长时间不可见、UI 一直显示「执行中」。
@@ -55,8 +60,10 @@ func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy
func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close() defer w.Close()
command = PrepareShellCommandForExecute(command)
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
ConfigureShellCmdForAgentExecute(cmd) applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
stdout, err := cmd.StdoutPipe() stdout, err := cmd.StdoutPipe()
if err != nil { if err != nil {
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
@@ -68,7 +75,8 @@ func runShellInBackground(ctx context.Context, command string, w *schema.StreamW
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return return
} }
if err := cmd.Start(); err != nil { session, err := StartShellSession(cmd)
if err != nil {
_ = stdout.Close() _ = stdout.Close()
_ = stderr.Close() _ = stderr.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
@@ -78,14 +86,14 @@ func runShellInBackground(ctx context.Context, command string, w *schema.StreamW
done := make(chan struct{}) done := make(chan struct{})
go func() { go func() {
drainShellPipes(stdout, stderr) drainShellPipes(stdout, stderr)
_ = cmd.Wait() _ = session.Wait()
close(done) close(done)
}() }()
select { select {
case <-done: case <-done:
case <-ctx.Done(): case <-ctx.Done():
TerminateShellCmdTree(cmd) TerminateShellCmdSession(session)
} }
exitCode := 0 exitCode := 0
@@ -112,8 +120,10 @@ func drainShellPipes(stdout, stderr io.Reader) {
func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close() defer w.Close()
command = PrepareShellCommandForExecute(command)
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
ConfigureShellCmdForAgentExecute(cmd) applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
stdoutPipe, err := cmd.StdoutPipe() stdoutPipe, err := cmd.StdoutPipe()
if err != nil { if err != nil {
@@ -126,7 +136,8 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return return
} }
if err := cmd.Start(); err != nil { session, err := StartShellSession(cmd)
if err != nil {
_ = stdoutPipe.Close() _ = stdoutPipe.Close()
_ = stderrPipe.Close() _ = stderrPipe.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
@@ -137,7 +148,7 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
go func() { go func() {
select { select {
case <-ctx.Done(): case <-ctx.Done():
TerminateShellCmdTree(cmd) TerminateShellCmdSession(session)
case <-stopWatch: case <-stopWatch:
} }
}() }()
@@ -174,12 +185,12 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
} }
hadOutput = true hadOutput = true
if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) { if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) {
TerminateShellCmdTree(cmd) TerminateShellCmdSession(session)
return return
} }
} }
waitErr := cmd.Wait() waitErr := session.Wait()
if waitErr == nil { if waitErr == nil {
exitCode := 0 exitCode := 0
_ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil) _ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil)
@@ -115,3 +115,38 @@ func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) {
t.Fatalf("expected early stderr, got: %q", got.String()) t.Fatalf("expected early stderr, got: %q", got.String())
} }
} }
// TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe 模拟 cmd & 后继续前台逻辑:重定向后应快速结束。
func TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
shell := NewEinoStreamingShell()
cmd := `(sh -c 'printf x; sleep 120') & echo started; sleep 0`
sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
start := time.Now()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if time.Since(start) > 3*time.Second {
t.Fatalf("expected fast completion, took %v output=%q", time.Since(start), got.String())
}
if !strings.Contains(got.String(), "started") {
t.Fatalf("expected foreground echo, got: %q", got.String())
}
}
+47
View File
@@ -0,0 +1,47 @@
package security
import "os/exec"
// ShellSession 在 Start 时记录根 shell 的进程组 ID,取消/超时时可杀整组(即使 cmd.Process 已失效)。
type ShellSession struct {
Cmd *exec.Cmd
rootPID int
}
// StartShellSession 配置独立进程组并启动 shell,缓存 rootPIDUnix 下即 PGID)。
func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) {
if err := prepareShellCmdSession(cmd); err != nil {
return nil, err
}
if err := cmd.Start(); err != nil {
return nil, err
}
pid := 0
if cmd.Process != nil {
pid = cmd.Process.Pid
}
return &ShellSession{Cmd: cmd, rootPID: pid}, nil
}
// Wait 等待 shell 退出。
func (s *ShellSession) Wait() error {
if s == nil || s.Cmd == nil {
return nil
}
return s.Cmd.Wait()
}
// Terminate 终止 shell 及其进程组。
func (s *ShellSession) Terminate() {
if s == nil {
return
}
terminateProcessGroup(s.rootPID, s.Cmd)
}
// TerminateShellSession 终止由 StartShellSession 启动的会话。
func TerminateShellSession(session *ShellSession) {
if session != nil {
session.Terminate()
}
}
+65
View File
@@ -0,0 +1,65 @@
package security
import (
"context"
"os/exec"
"runtime"
"testing"
"time"
)
func TestShellSession_TerminateUsesCachedRootPID(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
cmd := exec.Command("sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
session, err := StartShellSession(cmd)
if err != nil {
t.Fatalf("StartShellSession: %v", err)
}
time.Sleep(100 * time.Millisecond)
session.Terminate()
done := make(chan error, 1)
go func() { done <- session.Wait() }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("session did not finish within 5s after Terminate")
}
}
func TestShellSession_TerminateAfterContextCancel(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
session, err := StartShellSession(cmd)
if err != nil {
t.Fatalf("StartShellSession: %v", err)
}
time.Sleep(100 * time.Millisecond)
cancel()
TerminateShellCmdSession(session)
done := make(chan error, 1)
go func() { done <- session.Wait() }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("session did not finish within 5s after cancel+terminate")
}
}
+63 -28
View File
@@ -1,60 +1,95 @@
name: "hydra" name: "hydra"
command: "hydra" command: "hydra"
args: ["-I"]
enabled: true enabled: true
short_description: "密码暴力破解工具,支持多种协议和服务" short_description: "密码暴力破解工具,支持多种协议和服务"
description: | description: |
Hydra是一个快速的网络登录破解工具,支持多种协议和服务的密码暴力破解 Hydra 是网络登录口令爆破工具,支持 SSH、FTP、HTTP、SMB 等多种协议
**主要功能** **调用约定(必读)**
- 支持多种协议(SSH, FTP, HTTP, SMB等) - 必须提供 **用户名**`username`-l)或 `username_file`-L)至少其一
- 快速并行破解 - 必须提供 **口令**`password`-p)、`password_file`-P)或 `-C`(经 `additional_args`)至少其一
- 支持用户名和密码字典 - **先用小字典试跑**(几十~几百条),确认目标可达再扩大;禁止默认使用 rockyou 等超大字典
- 可恢复的会话 - 默认已启用:找到即停(-f)、并行 4(-t)、忽略 restore-I);长任务请设 `output_file`
**使用场景:** **CLI 顺序:** `hydra [选项] <target> <service>`(本工具已按此顺序组参,勿把 target 写在选项前)
- 密码强度测试
- 渗透测试 **使用场景:** 授权环境下的弱口令检测、密码强度评估
- 安全评估
- 弱密码检测 **注意:** 仅用于已授权目标;对无响应目标请减小 `wait_time` 或缩小字典,避免长时间挂起。
parameters: parameters:
- name: "target"
type: "string"
description: "目标IP或主机名"
required: true
position: 0
format: "positional"
- name: "service"
type: "string"
description: "服务类型(ssh, ftp, http等)"
required: true
position: 1
format: "positional"
- name: "username" - name: "username"
type: "string" type: "string"
description: "单个用户名" description: "单个用户名-l);与 username_file 二选一至少填一个"
required: false required: false
flag: "-l" flag: "-l"
format: "flag" format: "flag"
- name: "username_file" - name: "username_file"
type: "string" type: "string"
description: "用户名字典文件" description: "用户名字典文件-L"
required: false required: false
flag: "-L" flag: "-L"
format: "flag" format: "flag"
- name: "password" - name: "password"
type: "string" type: "string"
description: "单个密码" description: "单个密码-p"
required: false required: false
flag: "-p" flag: "-p"
format: "flag" format: "flag"
- name: "password_file" - name: "password_file"
type: "string" type: "string"
description: "密码字典文件" description: "密码字典文件-P);优先使用小字典试跑"
required: false required: false
flag: "-P" flag: "-P"
format: "flag" format: "flag"
- name: "stop_on_first"
type: "bool"
description: "找到一对有效账密后立即退出(-f,默认 true)"
required: false
flag: "-f"
format: "flag"
default: true
- name: "tasks"
type: "int"
description: "每目标并行连接数(-t);SSH 等建议 4,默认 4"
required: false
flag: "-t"
format: "flag"
default: 4
- name: "wait_time"
type: "int"
description: "单次连接等待响应秒数(-w),默认 16(低于 Hydra 默认 32,减少挂起感)"
required: false
flag: "-w"
format: "flag"
default: 16
- name: "wait_between"
type: "int"
description: "每线程连接间隔秒数(-W),默认 1"
required: false
flag: "-W"
format: "flag"
default: 1
- name: "output_file"
type: "string"
description: "将结果写入文件(-o),长任务建议指定"
required: false
flag: "-o"
format: "flag"
- name: "target"
type: "string"
description: "目标 IP、主机名或 CIDR(须在选项之后)"
required: true
position: 1
format: "positional"
- name: "service"
type: "string"
description: "服务类型(ssh、ftp、http-get、http-post-form、smb 等,见 hydra -h"
required: true
position: 2
format: "positional"
- name: "additional_args" - name: "additional_args"
type: "string" type: "string"
description: "额外的Hydra参数" description: "额外参数(如 -s 端口、-S SSL、-m 模块选项、-C login:pass 文件),追加在命令末尾"
required: false required: false
format: "positional" format: "positional"

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