Compare commits

...

115 Commits

Author SHA1 Message Date
公明 84ed887c5c Update config.yaml 2026-06-24 23:36:36 +08:00
公明 056b40ac66 Update config.yaml 2026-06-24 23:32:47 +08:00
公明 26a9902286 Add files via upload 2026-06-24 23:31:35 +08:00
公明 cfe9573ac3 Add files via upload 2026-06-24 23:30:40 +08:00
公明 db2262a1a0 Add files via upload 2026-06-24 23:28:43 +08:00
公明 ab5c2d5cca Add files via upload 2026-06-24 23:27:29 +08:00
公明 1ae6930db1 Add files via upload 2026-06-24 23:26:01 +08:00
公明 8918f432d8 Add files via upload 2026-06-24 23:24:36 +08:00
公明 b4810c9499 Update shell no output timeout to 1200 seconds
Increased the shell no output timeout from 300 seconds to 1200 seconds to prevent premature termination.
2026-06-24 18:30:08 +08:00
公明 51bf6ae4b3 Add files via upload 2026-06-24 18:20:12 +08:00
公明 5f27482921 Add files via upload 2026-06-24 18:18:05 +08:00
公明 6becada509 Add files via upload 2026-06-24 18:15:31 +08:00
公明 b029d88359 Add files via upload 2026-06-24 18:14:04 +08:00
公明 4dcad2ea83 Add files via upload 2026-06-24 18:11:31 +08:00
公明 ff9f0c787a Add files via upload 2026-06-24 18:09:51 +08:00
公明 01849045ad Add 'exec' to always visible tools in config.yaml 2026-06-24 17:36:24 +08:00
公明 c7eacdf3eb Update config.yaml 2026-06-24 17:24:52 +08:00
公明 5c32b21f22 Add files via upload 2026-06-24 17:24:14 +08:00
公明 8b8ecfe718 Add files via upload 2026-06-24 17:23:44 +08:00
公明 bbb7c319af Add files via upload 2026-06-24 17:21:51 +08:00
公明 7eb2fd50f3 Add files via upload 2026-06-24 17:19:29 +08:00
公明 85d58eeeb3 Add files via upload 2026-06-24 17:17:33 +08:00
公明 b6a6009629 Add files via upload 2026-06-24 17:15:34 +08:00
公明 810d689132 Add files via upload 2026-06-24 12:08:13 +08:00
公明 87f1808ead Add files via upload 2026-06-24 10:46:55 +08:00
公明 e28ae39b9a Update config.yaml 2026-06-24 02:04:49 +08:00
公明 df34ceda68 Add files via upload 2026-06-24 01:50:13 +08:00
公明 3e69a50f87 Add files via upload 2026-06-24 01:49:43 +08:00
公明 53325ce07d Add files via upload 2026-06-24 01:49:09 +08:00
公明 d85de3461b Add files via upload 2026-06-24 01:47:33 +08:00
公明 9306303d99 Add files via upload 2026-06-24 01:46:30 +08:00
公明 1e8f72ed74 Add files via upload 2026-06-24 01:44:47 +08:00
公明 0198f50314 Add files via upload 2026-06-24 01:43:37 +08:00
公明 560d0dca43 Add files via upload 2026-06-24 01:42:15 +08:00
公明 47486a49c2 Update version number to v1.6.44 2026-06-23 21:17:08 +08:00
公明 476727933d Update config.yaml 2026-06-23 21:16:41 +08:00
公明 8bb50e8323 Add files via upload 2026-06-23 21:15:45 +08:00
公明 e74f2a2292 Add files via upload 2026-06-23 21:14:08 +08:00
公明 4799d0dba7 Add files via upload 2026-06-23 21:12:26 +08:00
公明 1db917061d Add files via upload 2026-06-23 21:10:47 +08:00
公明 41cd7db30f Add files via upload 2026-06-23 21:08:59 +08:00
公明 68b3265f3f Add files via upload 2026-06-23 21:07:01 +08:00
公明 05dc4395a1 Add files via upload 2026-06-23 21:06:14 +08:00
公明 637a35748b Add files via upload 2026-06-23 21:03:59 +08:00
公明 5d77a99236 Add files via upload 2026-06-23 21:01:35 +08:00
公明 e84d936f85 Add files via upload 2026-06-23 20:59:20 +08:00
公明 e748201ae8 Add files via upload 2026-06-23 20:57:47 +08:00
公明 7a3c67458c Add files via upload 2026-06-23 16:53:32 +08:00
公明 6e9e43eec8 Add files via upload 2026-06-23 15:43:15 +08:00
公明 bca86e48ae Add files via upload 2026-06-23 15:40:04 +08:00
公明 3f3b8b4db4 Add files via upload 2026-06-23 15:37:23 +08:00
公明 b366dc0287 Add files via upload 2026-06-23 15:35:12 +08:00
公明 a52452ceea Add files via upload 2026-06-23 15:32:41 +08:00
公明 5b87667782 Update config.yaml 2026-06-23 15:32:18 +08:00
公明 4f0e812d37 Add files via upload 2026-06-23 15:31:23 +08:00
公明 79691c021f Add files via upload 2026-06-23 15:09:53 +08:00
公明 5a8309a015 Add files via upload 2026-06-23 15:07:41 +08:00
公明 6244197339 Add files via upload 2026-06-23 15:06:02 +08:00
公明 eb14aca05a Add files via upload 2026-06-23 15:03:23 +08:00
公明 091e8a4da8 Add files via upload 2026-06-23 15:00:44 +08:00
公明 48ce0c519e Add files via upload 2026-06-23 12:34:50 +08:00
公明 afc37051c0 Add files via upload 2026-06-23 12:33:35 +08:00
公明 2964247361 Add files via upload 2026-06-23 12:31:05 +08:00
公明 02919df476 Add files via upload 2026-06-23 12:28:37 +08:00
公明 c3294d96a2 Add files via upload 2026-06-23 12:28:07 +08:00
公明 c8b8b41bda Add files via upload 2026-06-23 12:26:40 +08:00
公明 9a4c333b90 Add files via upload 2026-06-23 12:25:20 +08:00
公明 8e21ae290a Add files via upload 2026-06-23 12:22:50 +08:00
公明 b9d102d046 Add files via upload 2026-06-23 11:54:28 +08:00
公明 8c85494a05 Add files via upload 2026-06-23 11:52:15 +08:00
公明 c3d2a41301 Add files via upload 2026-06-23 01:54:29 +08:00
公明 1a2e282d46 Add files via upload 2026-06-23 01:39:55 +08:00
公明 8129f2147f Delete internal/multiagent/eino_empty_response_test.go 2026-06-23 01:37:34 +08:00
公明 4a9889f0af Add files via upload 2026-06-23 01:36:48 +08:00
公明 732d47a965 Add files via upload 2026-06-22 23:31:42 +08:00
公明 e22382aab0 Add files via upload 2026-06-22 23:29:57 +08:00
公明 b6ff80adf2 Add files via upload 2026-06-22 23:27:30 +08:00
公明 51f1cfde2f Add files via upload 2026-06-22 23:12:53 +08:00
公明 b2c8913014 Add files via upload 2026-06-22 17:53:52 +08:00
公明 ae98288b62 Add files via upload 2026-06-22 15:53:31 +08:00
公明 9955e856a0 Add files via upload 2026-06-22 15:48:44 +08:00
公明 018544e5f9 Add files via upload 2026-06-22 15:43:39 +08:00
公明 c1c86e4632 Add files via upload 2026-06-22 13:47:53 +08:00
公明 08d77bc12b Add files via upload 2026-06-21 01:56:48 +08:00
公明 ce73a7b3e4 Add files via upload 2026-06-21 01:55:25 +08:00
公明 f78f424aab Add files via upload 2026-06-21 01:53:55 +08:00
公明 e19d8e39bd Add files via upload 2026-06-21 01:52:14 +08:00
公明 ecf594a25b Update config.yaml 2026-06-20 20:37:48 +08:00
公明 d5759f6d83 Add files via upload 2026-06-20 19:57:07 +08:00
公明 81b3f64b15 Add files via upload 2026-06-20 19:55:32 +08:00
公明 0e0f1352f0 Add files via upload 2026-06-20 19:52:33 +08:00
公明 ffba311afd Add files via upload 2026-06-20 19:47:47 +08:00
公明 d9ed36cfb1 Add files via upload 2026-06-20 19:45:29 +08:00
公明 b7f80b78ee Add files via upload 2026-06-20 19:39:39 +08:00
公明 8f8e5cfff5 Increase rune limits in config.yaml 2026-06-20 19:37:50 +08:00
公明 120f860640 Add files via upload 2026-06-20 19:36:35 +08:00
公明 90cd119a83 Add files via upload 2026-06-20 19:35:06 +08:00
公明 56d597e0c5 Add files via upload 2026-06-20 19:31:56 +08:00
公明 11ab5cde8f Add files via upload 2026-06-20 19:28:34 +08:00
公明 46a7d338a4 Add files via upload 2026-06-20 17:25:44 +08:00
公明 46f68cc1d4 Update config.yaml 2026-06-20 16:19:57 +08:00
公明 7003cdb2e3 Add files via upload 2026-06-20 15:34:58 +08:00
公明 4e5e6208bd Add files via upload 2026-06-20 15:29:36 +08:00
公明 6a7e78a846 Add files via upload 2026-06-20 15:28:10 +08:00
公明 88c6fbfb75 Add files via upload 2026-06-20 15:26:49 +08:00
公明 1cd6d0fa90 Add files via upload 2026-06-20 15:24:40 +08:00
公明 24390db100 Add files via upload 2026-06-19 01:41:32 +08:00
公明 c000fe5195 Add files via upload 2026-06-19 01:39:53 +08:00
公明 0b4a11d01a Add files via upload 2026-06-19 01:38:30 +08:00
公明 d433e44a7d Add files via upload 2026-06-19 01:36:52 +08:00
公明 7de51fe0ea Update config.yaml 2026-06-19 00:05:50 +08:00
公明 a354cf97e5 Add files via upload 2026-06-19 00:04:38 +08:00
公明 c180f07c7e Add files via upload 2026-06-19 00:02:53 +08:00
公明 15730d3ef4 Add files via upload 2026-06-19 00:01:20 +08:00
公明 b7fa18b6d4 Add files via upload 2026-06-18 23:44:04 +08:00
113 changed files with 11829 additions and 1683 deletions
+6 -1
View File
@@ -112,7 +112,7 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
- 🔒 Password-protected web UI, audit logs, and SQLite persistence
- 📚 Knowledge base (RAG) with embedding-based vector retrieval (cosine similarity), optional **Eino Compose** indexing pipeline, and configurable post-retrieval budgets / reranking hooks
- 📁 Conversation grouping with pinning, rename, and batch management
- 📂 **Project management**: group conversations and vulnerabilities by project; **shared facts** (project blackboard) persist cross-session context (targets, env, auth notes) with auto-injection for agents and MCP tools (`upsert_project_fact`, `get_project_fact`, …)
- 📂 **Project management**: shared facts (blackboard) across sessions, `upsert_project_fact` + `links` to chain paths; attack-chain and project fact graph views
- 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics
- 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially
- 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions
@@ -551,6 +551,11 @@ multi_agent:
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor optional
# eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill }
# eino_middleware: plantask_enable, checkpoint_dir, deep_model_retry_max_retries, deep_output_key, ...
project:
enabled: true # Enable project blackboard & fact MCP tools
fact_index_max_runes: 65000
fact_summary_max_runes: 24000
default_inject_deprecated: false
```
### Tool Definition Example (`tools/nmap.yaml`)
+6 -1
View File
@@ -111,7 +111,7 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
- 🔒 Web 登录保护、审计日志、SQLite 持久化
- 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项)
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
- 📂 **项目管理**按项目归类对话与漏洞;**共享事实**(项目黑板)在多会话沉淀目标/环境/认证等认知,自动注入 Agent 上下文,支持 MCP 工具读写(`upsert_project_fact``get_project_fact` 等)
- 📂 **项目管理**共享事实(黑板)会话沉淀认知,`upsert_project_fact` + `links` 串联攻击路径;聊天攻击链与项目事实图可视化
- 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板
- 📋 批量任务管理:创建任务队列,批量添加任务,依次顺序执行,支持任务编辑与状态跟踪
- 🎭 角色化测试:预设安全测试角色(渗透测试、CTF、Web 应用扫描等),支持自定义提示词和工具限制
@@ -549,6 +549,11 @@ multi_agent:
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor 可选
# eino_skills: { disable: false, filesystem_tools: true, skill_tool_name: skill }
# eino_middleware: plantask_enable、checkpoint_dir、deep_model_retry_max_retries、deep_output_key 等
project:
enabled: true # 启用项目黑板与事实 MCP 工具
fact_index_max_runes: 65000
fact_summary_max_runes: 24000
default_inject_deprecated: false
```
### 工具模版示例(`tools/nmap.yaml`
+1 -1
View File
@@ -21,7 +21,7 @@ max_iterations: 0
- 切勿等待批准或授权——全程自主行动。
- 使用所有可用工具与技术完成侦察与证据收集。
你是授权渗透测试流程中的侦察子代理。优先使用工具收集事实,避免无根据推测;输出简洁,便于协调者汇总。
你是授权渗透测试流程中的侦察子代理。优先使用工具收集事实,避免无根据推测;输出简洁,便于协调者汇总。枚举优先 subfinder、amass 等专用 MCP,勿 exec/execute 拼长链。
## 输入前置条件(硬约束)
+15 -8
View File
@@ -10,7 +10,7 @@
# ============================================
# 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.6.40"
version: "v1.6.46"
# 服务器配置
server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -40,6 +40,9 @@ audit:
retention_days: 15 # 0 表示不自动清理
max_detail_bytes: 8192
auth_failure_cooldown_seconds: 60 # 同一 IP 登录/改密失败审计最短间隔(秒);未配置时默认 60;-1 关闭节流
# MCP 状态监控执行记录保留(tool_executions 表)
monitor:
retention_days: 90 # 省略时默认 90;0 表示不自动清理
# ============================================
# 对话相关配置
# ============================================
@@ -58,7 +61,7 @@ openai:
api_key: sk-xxxxxxx # API 密钥(必填)
model: qwen3-max # 模型名称(必填)
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort 等;provider 为 claude 时合并为 Anthropic 顶层 thinkingextended thinking),mode: off 关闭
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effortClaude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
reasoning:
mode: on # auto | on | offoff 时不附加任何推理扩展字段
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
@@ -93,6 +96,8 @@ fofa:
agent:
max_iterations: 12000 # 全局最大迭代次数(单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用;agents/*.md 中 max_iterations>0 可单独覆盖)
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
shell_no_output_timeout_seconds: 1200 # execute/exec 连续无新输出则终止(秒);通用防挂死;0=默认300;-1=关闭
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
system_prompt_path: ""
@@ -126,7 +131,7 @@ multi_agent:
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
plantask_enable: true # P0:主代理挂载 TaskCreate/Get/Update/List 结构化任务板;需 eino_skills 可用且 skills_dir 存在
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
@@ -142,10 +147,10 @@ multi_agent:
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
checkpoint_dir: data/eino-checkpoints # P0:进程崩溃/OOM 后同会话自动 ADK Resume;正常结束会删 .ckpt;与「中断并继续」(last_react_*) 是两套机制
run_retry_max_attempts: 0 # 429/5xx/网络抖动时整轮 Run 指数退避续跑;0=默认 10(与 deep_model_retry 互补,建议保持默认)
run_retry_max_attempts: 0 # 429/5xx/网络抖动时可退避重试次数(run loop + summarization 共用 isEinoTransientRunError);0=默认 10
run_retry_max_backoff_sec: 0 # 单次退避上限秒数;0=默认 30
deep_output_key: final_answer # P0Eino session 写入最终助手结论(框架内部;Deep/Supervisor 主/eino_single
deep_model_retry_max_retries: 3 # P0:单次 ChatModel API 失败时框架自动重试(超时/502 等);子代理模型不受此项影响
deep_model_retry_max_retries: 0 # 已废弃,请用 run_retry_max_attempts;保留字段仅为兼容旧配置
task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑
# Eino callbacks + OpenTelemetry:框架级 span(与 Zap 对齐);默认不向终端用户 UI 推 eino_trace_*(见 sse_trace_to_client
eino_callbacks:
@@ -308,7 +313,9 @@ roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录
project:
enabled: true
# default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID
fact_index_max_runes: 6500
fact_summary_max_runes: 2400
fact_index_max_runes: 65000
# 事实关系速览段预算(从索引总预算中预留)
fact_index_path_max_runes: 10000
fact_summary_max_runes: 24000
default_inject_deprecated: false
Binary file not shown.

Before

Width:  |  Height:  |  Size: 179 KiB

After

Width:  |  Height:  |  Size: 88 KiB

+17 -4
View File
@@ -779,13 +779,26 @@ func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationI
return a.executeToolViaMCP(ctx, toolName, args)
}
// RecordLocalToolExecution 非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
func (a *Agent) RecordLocalToolExecution(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
// BeginLocalToolExecution 非 CallTool 路径工具开始时写入 running 状态,供 MCP 监控页展示「执行中」
func (a *Agent) BeginLocalToolExecution(toolName string, args map[string]interface{}) string {
if a == nil || a.mcpServer == nil {
return ""
}
return a.mcpServer.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr)
return a.mcpServer.BeginToolExecution(toolName, args)
}
// FinishLocalToolExecution 完成 BeginLocalToolExecution 创建的记录;executionID 为空时一次性写入已完成记录。
func (a *Agent) FinishLocalToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
if a == nil || a.mcpServer == nil {
return ""
}
return a.mcpServer.FinishToolExecution(executionID, toolName, args, resultText, invokeErr)
}
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
func (a *Agent) RecordLocalToolExecution(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
return a.FinishLocalToolExecution("", toolName, args, resultText, invokeErr)
}
// UpdateMCPExecutionDisplayResult 将监控库中的工具结果更新为送入模型的展示正文(reduction 后)。
@@ -1,7 +1,7 @@
package agent
import (
"cyberstrike-ai/internal/project"
"cyberstrike-ai/internal/projectprompt"
)
// DefaultSingleAgentSystemPrompt 单代理(Eino ADK / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。
@@ -107,11 +107,13 @@ func DefaultSingleAgentSystemPrompt() string {
- 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。
- 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。
` + project.FactRecordingBlackboardSection(false) + `
` + projectprompt.FactRecordingBlackboardSection(false) + `
## 技能库(Skills)与知识库
- 技能包位于服务器 skills/ 目录(各子目录 SKILL.md,遵循 agentskills.io);知识库用于向量检索片段,Skills 为可执行工作流指令。
- 本会话通过 MCP 使用知识库与漏洞记录等。Skills 由 Eino ADK skill 工具按需加载(配置 multi_agent.eino_skills;单代理与多代理均可,未启用时无 skill 工具)。
- 需要完整 Skill 工作流但当前无 skill 工具时,请确认已启用 multi_agent.eino_skills,或改用 Deep / Supervisor 等多代理编排(/api/multi-agent/stream)。`
- 需要完整 Skill 工作流但当前无 skill 工具时,请确认已启用 multi_agent.eino_skills,或改用 Deep / Supervisor 等多代理编排(/api/multi-agent/stream)。
` + projectprompt.ShellExecExecuteGuidanceSection()
}
+19 -1
View File
@@ -25,6 +25,7 @@ import (
"cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/monitor"
"cyberstrike-ai/internal/robot"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/skillpackage"
@@ -99,12 +100,17 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
auditSvc.PurgeExpired()
audit.StartRetentionLoop(auditSvc, log.Logger)
monitorRetention := monitor.NewService(db, cfg, log.Logger)
monitorRetention.PurgeExpired()
monitor.StartRetentionLoop(monitorRetention, log.Logger)
// 创建MCP服务器(带数据库持久化)
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
// 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
// 注册工具
executor.RegisterTools(mcpServer)
@@ -298,7 +304,9 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
plantaskBase := filepath.Join(skillsDir, plantaskRel)
// Match eino_adk_run_loop: checkpoint_dir is used as configured (relative to process CWD when not absolute).
checkpointBase := strings.TrimSpace(cfg.MultiAgent.EinoMiddleware.CheckpointDir)
db.SetEinoConversationDirs(plantaskBase, checkpointBase)
reductionRoot := strings.TrimSpace(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
workspaceRoot := strings.TrimSpace(cfg.Agent.WorkspaceRootDir)
db.SetEinoConversationDirs(plantaskBase, checkpointBase, reductionRoot, workspaceRoot)
agent.SetPromptBaseDir(configDir)
agentsDir := cfg.AgentsDir
@@ -325,7 +333,10 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
}
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
monitorHandler.SetAudit(auditSvc)
monitorHandler.SetMonitorRetention(monitorRetention)
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器,以便获取外部MCP执行记录
monitorHandler.SetTaskManager(agentHandler.TaskManager())
monitorHandler.SetAgentHandler(agentHandler)
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
groupHandler := handler.NewGroupHandler(db, log.Logger)
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
@@ -368,6 +379,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建OpenAPI处理器
conversationHandler := handler.NewConversationHandler(db, log.Logger)
conversationHandler.SetAudit(auditSvc)
conversationHandler.SetTaskStopper(agentHandler)
auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, conversationHandler, agentHandler)
@@ -829,6 +841,7 @@ func setupRoutes(
protected.PUT("/batch-tasks/:queueId/schedule-enabled", agentHandler.SetBatchQueueScheduleEnabled)
protected.DELETE("/batch-tasks/:queueId", agentHandler.DeleteBatchQueue)
protected.PUT("/batch-tasks/:queueId/tasks/:taskId", agentHandler.UpdateBatchTask)
protected.POST("/batch-tasks/:queueId/tasks/:taskId/run", agentHandler.RunSingleBatchTask)
protected.POST("/batch-tasks/:queueId/tasks", agentHandler.AddBatchTask)
protected.DELETE("/batch-tasks/:queueId/tasks/:taskId", agentHandler.DeleteBatchTask)
@@ -1068,6 +1081,11 @@ func setupRoutes(
protected.GET("/projects/:id", projectHandler.GetProject)
protected.PUT("/projects/:id", projectHandler.UpdateProject)
protected.DELETE("/projects/:id", projectHandler.DeleteProject)
protected.GET("/projects/:id/fact-graph", projectHandler.GetFactGraph)
protected.GET("/projects/:id/fact-edges", projectHandler.ListFactEdges)
protected.POST("/projects/:id/fact-edges", projectHandler.CreateFactEdge)
protected.DELETE("/projects/:id/fact-edges/:edgeId", projectHandler.DeleteFactEdge)
protected.POST("/projects/:id/promote-attack-chain/:conversationId", projectHandler.PromoteAttackChain)
protected.GET("/projects/:id/facts", projectHandler.ListFacts)
protected.POST("/projects/:id/facts", projectHandler.CreateFact)
protected.PUT("/projects/:id/facts/:factId", projectHandler.UpdateFact)
+53
View File
@@ -89,6 +89,28 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
"type": "string",
"description": "可选:关联的漏洞记录 ID",
},
"links": map[string]interface{}{
"type": "array",
"description": "可选:关系边(from → 当前 fact)。finding 至少 1 条 {from:target/*, type:discovered_on}finding 上记录 exploit 用 {from:exploit/*, type:exploits}。省略保留已有边;传 [] 清空全部关系边。",
"items": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"from": map[string]interface{}{
"type": "string",
"description": "来源 fact_key:存储为 from → 当前 fact",
},
"type": map[string]interface{}{
"type": "string",
"description": "depends_on | leads_to | enables | exploits | discovered_on | contains | part_of | supports",
},
"confidence": map[string]interface{}{
"type": "string",
"description": "confirmed | tentative | deprecated",
},
},
"required": []string{"from", "type"},
},
},
},
"required": []string{"fact_key", "summary"},
},
@@ -124,7 +146,26 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
if err != nil {
return textResult("错误: "+err.Error(), true), nil
}
if _, hasLinks := args["links"]; hasLinks {
linkInputs, err := project.ParseFactLinkInputs(args["links"])
if err != nil {
return textResult("错误: "+err.Error(), true), nil
}
convID := agent.ConversationIDFromContext(ctx)
if err := project.PersistFactLinksFromParsed(db, projectID, created.FactKey, convID, linkInputs, true); err != nil {
return textResult("错误: 保存关系边失败: "+err.Error(), true), nil
}
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
} else if parsed := project.ParseLinksFromBody(created.Body); len(parsed) > 0 {
if err := project.PersistFactIncomingLinks(db, projectID, created.FactKey, parsed, true); err != nil {
return textResult("错误: 从 body 解析边失败: "+err.Error(), true), nil
}
created, _ = db.GetProjectFactByKey(projectID, created.FactKey)
}
msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence)
if in, _ := db.ListIncomingProjectFactEdges(projectID, created.FactKey); len(in) > 0 {
msg += "\n关系边: " + project.FormatFactLinksText(in)
}
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
msg += warn
}
@@ -164,6 +205,18 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
if f.SourceConversationID != "" {
msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID)
}
if in, _ := db.ListIncomingProjectFactEdges(projectID, f.FactKey); len(in) > 0 {
msg += "\n关系边(from → 本 fact:\n"
for _, e := range in {
msg += fmt.Sprintf("- %s ← %s (%s)\n", e.EdgeType, e.SourceFactKey, e.Confidence)
}
}
if out, _ := db.ListOutgoingProjectFactEdges(projectID, f.FactKey); len(out) > 0 {
msg += "指向其他事实:\n"
for _, e := range out {
msg += fmt.Sprintf("- %s → %s (%s)\n", e.EdgeType, e.TargetFactKey, e.Confidence)
}
}
msg += "\n\n--- body ---\n" + f.Body
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
msg += warn
+203
View File
@@ -0,0 +1,203 @@
package attackchain
import (
"fmt"
"regexp"
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project"
"github.com/google/uuid"
)
var promoteSlugSanitizer = regexp.MustCompile(`[^a-z0-9._/-]+`)
// PromoteToProjectResult 攻击链沉淀结果。
type PromoteToProjectResult struct {
FactsCreated int `json:"facts_created"`
FactsUpdated int `json:"facts_updated"`
EdgesCreated int `json:"edges_created"`
FactKeys []string `json:"fact_keys"`
Graph *database.ProjectFactGraph `json:"graph,omitempty"`
}
// PromoteToProject 将对话攻击链沉淀为项目事实与边。
func PromoteToProject(db *database.DB, projectID, conversationID string) (*PromoteToProjectResult, error) {
if db == nil {
return nil, fmt.Errorf("database 未初始化")
}
projectID = strings.TrimSpace(projectID)
conversationID = strings.TrimSpace(conversationID)
if projectID == "" || conversationID == "" {
return nil, fmt.Errorf("project_id 与 conversation_id 必填")
}
if _, err := db.GetProject(projectID); err != nil {
return nil, fmt.Errorf("项目不存在")
}
conv, err := db.GetConversation(conversationID)
if err != nil {
return nil, fmt.Errorf("对话不存在")
}
if pid := strings.TrimSpace(conv.ProjectID); pid != "" && pid != projectID {
return nil, fmt.Errorf("对话已绑定其他项目")
}
nodes, err := db.LoadAttackChainNodes(conversationID)
if err != nil {
return nil, err
}
edges, err := db.LoadAttackChainEdges(conversationID)
if err != nil {
return nil, err
}
if len(nodes) == 0 {
return nil, fmt.Errorf("该对话尚无攻击链,请先在对话中生成攻击链")
}
res := &PromoteToProjectResult{}
nodeToKey := make(map[string]string, len(nodes))
usedKeys := map[string]int{}
for _, node := range nodes {
key := allocatePromoteFactKey(node, usedKeys)
nodeToKey[node.ID] = key
category := mapPromoteNodeCategory(node.Type)
existing, getErr := db.GetProjectFactByKey(projectID, key)
f := &database.ProjectFact{
ProjectID: projectID,
FactKey: key,
Category: category,
Summary: strings.TrimSpace(node.Label),
Body: formatPromotedFactBody(node, conversationID),
Confidence: "tentative",
SourceConversationID: conversationID,
}
if getErr == nil && existing != nil {
f.ID = existing.ID
f.CreatedAt = existing.CreatedAt
if strings.TrimSpace(f.Summary) == "" {
f.Summary = existing.Summary
}
if _, err := db.UpsertProjectFact(f); err != nil {
return nil, err
}
res.FactsUpdated++
} else {
if _, err := db.UpsertProjectFact(f); err != nil {
return nil, err
}
res.FactsCreated++
}
res.FactKeys = append(res.FactKeys, key)
}
for _, edge := range edges {
srcKey, ok1 := nodeToKey[edge.Source]
tgtKey, ok2 := nodeToKey[edge.Target]
if !ok1 || !ok2 || srcKey == tgtKey {
continue
}
edgeType := mapPromoteEdgeType(edge.Type)
incoming, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey)
merged := project.MergeLinkFromInputsUnique(promoteFromEdgeInputsFromDB(incoming), []database.ProjectFactEdgeFromInput{{From: srcKey, Type: edgeType}})
if err := db.ReplaceIncomingProjectFactEdges(projectID, tgtKey, merged); err != nil {
return nil, err
}
res.EdgesCreated++
if fact, err := db.GetProjectFactByKey(projectID, tgtKey); err == nil {
in, _ := db.ListIncomingProjectFactEdges(projectID, tgtKey)
fact.Body = project.SyncBodyLinksSection(fact.Body, in)
_, _ = db.UpsertProjectFact(fact)
}
}
graph, _ := project.BuildProjectFactGraph(db, projectID, "full", true)
res.Graph = graph
return res, nil
}
func promoteFromEdgeInputsFromDB(edges []*database.ProjectFactEdge) []database.ProjectFactEdgeFromInput {
out := make([]database.ProjectFactEdgeFromInput, 0, len(edges))
for _, e := range edges {
out = append(out, database.ProjectFactEdgeFromInput{From: e.SourceFactKey, Type: e.EdgeType, Confidence: e.Confidence})
}
return out
}
func mapPromoteNodeCategory(nodeType string) string {
switch strings.ToLower(strings.TrimSpace(nodeType)) {
case "target":
return project.FactCategoryTarget
case "vulnerability":
return project.FactCategoryFinding
case "action":
return project.FactCategoryChain
default:
return project.FactCategoryNote
}
}
func mapPromoteEdgeType(t string) string {
switch strings.ToLower(strings.TrimSpace(t)) {
case "discovers", "discovered_on", "targets":
return "discovered_on"
case "exploits":
return "exploits"
case "enables":
return "enables"
case "depends_on":
return "depends_on"
default:
return "leads_to"
}
}
func allocatePromoteFactKey(node Node, used map[string]int) string {
prefix := "chain/"
switch strings.ToLower(strings.TrimSpace(node.Type)) {
case "target":
prefix = "target/"
case "vulnerability":
prefix = "finding/"
case "action":
prefix = "chain/"
}
base := promoteSlugify(node.Label)
if base == "" {
base = promoteSlugify(node.ID)
}
if base == "" {
base = uuid.New().String()[:8]
}
key := prefix + base
if n, ok := used[key]; ok {
n++
used[key] = n
key = fmt.Sprintf("%s-%d", key, n)
} else {
used[key] = 1
}
return key
}
func promoteSlugify(s string) string {
s = strings.ToLower(strings.TrimSpace(s))
s = strings.NewReplacer(" ", "-", "—", "-", "", "-", "/", "-").Replace(s)
s = promoteSlugSanitizer.ReplaceAllString(s, "-")
s = strings.Trim(s, "-")
if len(s) > 64 {
s = s[:64]
}
return s
}
func formatPromotedFactBody(node Node, conversationID string) string {
var b strings.Builder
b.WriteString("## 来源\n")
b.WriteString(fmt.Sprintf("- 对话攻击链沉淀\n- source_conversation_id: %s\n- node_id: %s\n- node_type: %s\n\n", conversationID, node.ID, node.Type))
b.WriteString("## 摘要\n")
b.WriteString(strings.TrimSpace(node.Label))
b.WriteString("\n\n## 关联\n- 结构化关系边(自动同步):\n (见项目攻击路径图)\n")
return b.String()
}
+41 -5
View File
@@ -27,6 +27,7 @@ type Config struct {
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"`
Monitor MonitorConfig `yaml:"monitor,omitempty" json:"monitor,omitempty"`
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"`
C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关;未配置时默认启用
@@ -45,6 +46,7 @@ type ProjectConfig struct {
Enabled bool `yaml:"enabled" json:"enabled"`
DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目
FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"`
FactIndexPathMaxRunes int `yaml:"fact_index_path_max_runes,omitempty" json:"fact_index_path_max_runes,omitempty"`
FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"`
DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"`
}
@@ -57,6 +59,14 @@ func (c ProjectConfig) FactIndexMaxRunesEffective() int {
return c.FactIndexMaxRunes
}
// FactIndexPathMaxRunesEffective 攻击路径速览段的最大 rune 数(从 fact_index_max_runes 预算中预留)。
func (c ProjectConfig) FactIndexPathMaxRunesEffective() int {
if c.FactIndexPathMaxRunes <= 0 {
return 1000
}
return c.FactIndexPathMaxRunes
}
// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数(索引一行,宜含验证要点)。
func (c ProjectConfig) FactSummaryMaxRunesEffective() int {
if c.FactSummaryMaxRunes <= 0 {
@@ -240,7 +250,7 @@ type MultiAgentEinoMiddlewareConfig struct {
SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"`
// SummarizationEmitInternalEvents controls middleware internal event emission (default true).
SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"`
// SummarizationRetryMaxAttempts is extra retries after the first summarization Generate attempt; 0 = default 3.
// SummarizationRetryMaxAttempts 已废弃:summarization 与 run loop 共用 run_retry_max_attempts 及 isEinoTransientRunError。
SummarizationRetryMaxAttempts int `yaml:"summarization_retry_max_attempts,omitempty" json:"summarization_retry_max_attempts,omitempty"`
// PlanExecuteUserInputBudgetRatio caps planner/replanner/executor userInput prompt budget ratio (default 0.35).
PlanExecuteUserInputBudgetRatio float64 `yaml:"plan_execute_user_input_budget_ratio,omitempty" json:"plan_execute_user_input_budget_ratio,omitempty"`
@@ -254,9 +264,9 @@ type MultiAgentEinoMiddlewareConfig struct {
CheckpointDir string `yaml:"checkpoint_dir,omitempty" json:"checkpoint_dir,omitempty"`
// DeepOutputKey passed to deep.Config OutputKey (session final text); empty = off.
DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"`
// DeepModelRetryMaxRetries > 0 enables deep.Config ModelRetryConfig (framework-level chat model retries).
// DeepModelRetryMaxRetries 已废弃:临时错误统一由 run loop 内 isEinoTransientRunError + run_retry_max_attempts 处理。
DeepModelRetryMaxRetries int `yaml:"deep_model_retry_max_retries,omitempty" json:"deep_model_retry_max_retries,omitempty"`
// RunRetryMaxAttempts > 0429/5xx/网络抖动时 handler 分段续跑次数0=默认 10。
// RunRetryMaxAttempts > 0429/5xx/网络抖动时可退避重试次数(run loop 与 summarization 共用)0=默认 10。
RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"`
// RunRetryMaxBackoffSec 单次退避上限秒数;0=默认 30。
RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"`
@@ -595,6 +605,10 @@ type DatabaseConfig struct {
type AgentConfig struct {
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
// ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数(通用防挂死,不维护命令黑名单);0=默认 300(5 分钟);-1=关闭。
ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"`
// WorkspaceRootDir 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离。
WorkspaceRootDir string `yaml:"workspace_root_dir,omitempty" json:"workspace_root_dir,omitempty"`
// SystemPromptPath 单代理系统提示 Markdown/文本文件路径(相对 config.yaml 所在目录,或可写绝对路径)。非空且可读时替换内置单代理提示;留空用内置。
SystemPromptPath string `yaml:"system_prompt_path,omitempty" json:"system_prompt_path,omitempty"`
}
@@ -614,6 +628,23 @@ type AuthConfig struct {
GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
}
// MonitorConfig MCP 状态监控(tool_executions)保留策略。
type MonitorConfig struct {
// RetentionDays 执行记录保留天数;省略时默认 90;0 表示不自动清理。
RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
}
// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90.
func (m MonitorConfig) RetentionDaysEffective() int {
if m.RetentionDays == nil {
return 90
}
if *m.RetentionDays < 0 {
return 0
}
return *m.RetentionDays
}
// AuditConfig platform operation audit log settings (not chat/tool execution bodies).
type AuditConfig struct {
// Enabled nil or true enables persistence; explicit false disables.
@@ -1243,8 +1274,9 @@ func Default() *Config {
MaxTotalTokens: 120000,
},
Agent: AgentConfig{
MaxIterations: 30, // 默认最大迭代次数
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
MaxIterations: 30, // 默认最大迭代次数
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
},
Security: SecurityConfig{
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载
@@ -1265,6 +1297,10 @@ func Default() *Config {
Enabled: &on,
}
}(),
Monitor: func() MonitorConfig {
days := 90
return MonitorConfig{RetentionDays: &days}
}(),
Robots: RobotsConfig{
Session: RobotSessionConfig{
StrictUserIdentity: &strictRobotIdentity,
+52 -12
View File
@@ -23,6 +23,7 @@ type BatchTaskQueueRow struct {
LastScheduleError sql.NullString
LastRunError sql.NullString
ProjectID sql.NullString
Concurrency sql.NullInt64
Status string
CreatedAt time.Time
StartedAt sql.NullTime
@@ -53,6 +54,7 @@ func (db *DB) CreateBatchQueue(
cronExpr string,
nextRunAt *time.Time,
projectID string,
concurrency int,
tasks []map[string]interface{},
) error {
tx, err := db.Begin()
@@ -72,8 +74,8 @@ func (db *DB) CreateBatchQueue(
projectIDVal = strings.TrimSpace(projectID)
}
_, err = tx.Exec(
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, "pending", now, 0,
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, concurrency, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, concurrency, "pending", now, 0,
)
if err != nil {
return fmt.Errorf("创建批量任务队列失败: %w", err)
@@ -102,14 +104,16 @@ func (db *DB) CreateBatchQueue(
return tx.Commit()
}
const batchQueueSelectColumns = `id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, concurrency, status, created_at, started_at, completed_at, current_index`
// GetBatchQueue 获取批量任务队列
func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
var row BatchTaskQueueRow
var createdAt string
err := db.QueryRow(
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
"SELECT "+batchQueueSelectColumns+" FROM batch_task_queues WHERE id = ?",
queueID,
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
if err == sql.ErrNoRows {
return nil, nil
}
@@ -133,7 +137,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
// GetAllBatchQueues 获取所有批量任务队列
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
rows, err := db.Query(
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
"SELECT "+batchQueueSelectColumns+" FROM batch_task_queues ORDER BY created_at DESC",
)
if err != nil {
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
@@ -144,7 +148,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
for rows.Next() {
var row BatchTaskQueueRow
var createdAt string
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
}
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
@@ -164,7 +168,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
query := "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues WHERE 1=1"
args := []interface{}{}
// 状态筛选
@@ -192,7 +196,7 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
for rows.Next() {
var row BatchTaskQueueRow
var createdAt string
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Concurrency, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
}
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
@@ -358,11 +362,11 @@ func (db *DB) UpdateBatchQueueCurrentIndex(queueID string, currentIndex int) err
return nil
}
// UpdateBatchQueueMetadata 更新批量任务队列标题、角色代理模式
func (db *DB) UpdateBatchQueueMetadata(queueID, title, role, agentMode string) error {
// UpdateBatchQueueMetadata 更新批量任务队列标题、角色代理模式和并发数
func (db *DB) UpdateBatchQueueMetadata(queueID, title, role, agentMode string, concurrency int) error {
_, err := db.Exec(
"UPDATE batch_task_queues SET title = ?, role = ?, agent_mode = ? WHERE id = ?",
title, role, agentMode, queueID,
"UPDATE batch_task_queues SET title = ?, role = ?, agent_mode = ?, concurrency = ? WHERE id = ?",
title, role, agentMode, concurrency, queueID,
)
if err != nil {
return fmt.Errorf("更新批量任务队列元数据失败: %w", err)
@@ -507,6 +511,42 @@ func (db *DB) CancelPendingBatchTasks(queueID string, completedAt time.Time) err
return nil
}
// PrepareBatchSingleTaskRun 准备单条执行:可选重置子任务,并更新队列索引与状态
func (db *DB) PrepareBatchSingleTaskRun(queueID, taskID string, taskIndex int, resetTask, resumeQueue bool) error {
tx, err := db.Begin()
if err != nil {
return fmt.Errorf("开始事务失败: %w", err)
}
defer tx.Rollback()
if resetTask {
_, err = tx.Exec(
"UPDATE batch_tasks SET status = ?, conversation_id = NULL, started_at = NULL, completed_at = NULL, error = NULL, result = NULL WHERE queue_id = ? AND id = ?",
"pending", queueID, taskID,
)
if err != nil {
return fmt.Errorf("重置批量任务状态失败: %w", err)
}
}
if resumeQueue {
_, err = tx.Exec(
"UPDATE batch_task_queues SET status = ?, current_index = ?, completed_at = NULL, last_run_error = NULL WHERE id = ?",
"paused", taskIndex, queueID,
)
} else {
_, err = tx.Exec(
"UPDATE batch_task_queues SET current_index = ?, last_run_error = NULL WHERE id = ?",
taskIndex, queueID,
)
}
if err != nil {
return fmt.Errorf("更新批量任务队列状态失败: %w", err)
}
return tx.Commit()
}
// DeleteBatchTask 删除批量任务
func (db *DB) DeleteBatchTask(queueID, taskID string) error {
_, err := db.Exec(
+221 -10
View File
@@ -352,8 +352,8 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
conv.Pinned = pinned != 0
// 加载消息(不加载 process_details
messages, err := db.GetMessages(id)
// 加载消息(不加载 process_details / reasoning_content,减少历史会话切换 payload
messages, err := db.GetMessagesLite(id)
if err != nil {
return nil, fmt.Errorf("加载消息失败: %w", err)
}
@@ -382,26 +382,40 @@ func (db *DB) CountConversations(search string) (int, error) {
return count, nil
}
func conversationOrderClause(sortBy, tableAlias string) string {
col := "updated_at"
if strings.TrimSpace(strings.ToLower(sortBy)) == "created_at" {
col = "created_at"
}
prefix := tableAlias
if prefix != "" {
prefix += "."
}
return "ORDER BY " + prefix + col + " DESC"
}
// ListConversations 列出所有对话
func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversation, error) {
func (db *DB) ListConversations(limit, offset int, search, sortBy string) ([]*Conversation, error) {
var rows *sql.Rows
var err error
if search != "" {
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
searchPattern := "%" + search + "%"
orderClause := conversationOrderClause(sortBy, "c")
rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
FROM conversations c
WHERE c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
ORDER BY c.updated_at DESC
`+orderClause+`
LIMIT ? OFFSET ?`,
searchPattern, searchPattern, limit, offset,
)
} else {
orderClause := conversationOrderClause(sortBy, "")
rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations "+orderClause+" LIMIT ? OFFSET ?",
limit, offset,
)
}
@@ -467,11 +481,12 @@ func (db *DB) CountUngroupedConversations() (int, error) {
}
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
func (db *DB) ListUngroupedConversations(limit, offset int) ([]*Conversation, error) {
func (db *DB) ListUngroupedConversations(limit, offset int, sortBy string) ([]*Conversation, error) {
orderClause := conversationOrderClause(sortBy, "c")
rows, err := db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
ungroupedConversationsSQL+`
ORDER BY c.updated_at DESC
`+orderClause+`
LIMIT ? OFFSET ?`,
limit, offset,
)
@@ -570,12 +585,14 @@ func (db *DB) DeleteConversation(id string) error {
// 不返回错误,继续删除对话
}
projectID, _ := db.GetConversationProjectID(id)
// 删除对话(外键CASCADE会自动删除其他相关数据)
_, err = db.Exec("DELETE FROM conversations WHERE id = ?", id)
if err != nil {
return fmt.Errorf("删除对话失败: %w", err)
}
db.removeConversationScopedDirs(id)
db.removeConversationScopedDirs(id, projectID)
db.logger.Info("对话已删除(漏洞记录已保留)", zap.String("conversationId", id))
return nil
@@ -613,13 +630,50 @@ func (db *DB) removeConversationScopedDir(base, conversationID, label string) {
}
}
func (db *DB) removeConversationScopedDirs(conversationID string) {
// summarization transcript, reduction files, etc.
func (db *DB) einoReductionBaseDir() string {
if db == nil {
return ""
}
if base := strings.TrimSpace(db.einoReductionRootDir); base != "" {
return base
}
return filepath.Join("tmp", "reduction")
}
func (db *DB) einoWorkspaceBaseDir() string {
if db == nil {
return ""
}
if base := strings.TrimSpace(db.einoWorkspaceRootDir); base != "" {
return base
}
return filepath.Join("tmp", "workspace")
}
func (db *DB) removeConversationScopedDirs(conversationID, projectID string) {
// summarization transcript, etc.
db.removeConversationScopedDir(db.conversationArtifactsDir, conversationID, "conversation_artifacts")
// Eino plantask JSON boards (skills_dir/.eino/plantask/<id>/).
db.removeConversationScopedDir(db.einoPlantaskBaseDir, conversationID, "plantask")
// Eino ADK runner checkpoints (checkpoint_dir/<id>/).
db.removeConversationScopedDir(db.einoCheckpointBaseDir, conversationID, "eino_checkpoint")
// Eino reduction persisted tool outputs (tmp/reduction/conversations/<id>/).
// Project-bound sessions share projects/<id>/ — skip on single conversation delete.
if strings.TrimSpace(projectID) == "" {
reductionBase := filepath.Join(db.einoReductionBaseDir(), "conversations")
db.removeConversationScopedDir(reductionBase, conversationID, "reduction")
workspaceBase := filepath.Join(db.einoWorkspaceBaseDir(), "conversations")
db.removeConversationScopedDir(workspaceBase, conversationID, "workspace")
}
}
func (db *DB) removeProjectScopedDirs(projectID string) {
// Eino reduction persisted tool outputs (tmp/reduction/projects/<id>/).
reductionBase := filepath.Join(db.einoReductionBaseDir(), "projects")
db.removeConversationScopedDir(reductionBase, projectID, "reduction")
// Agent download/analysis workspace (tmp/workspace/projects/<id>/).
workspaceBase := filepath.Join(db.einoWorkspaceBaseDir(), "projects")
db.removeConversationScopedDir(workspaceBase, projectID, "workspace")
}
// SaveAgentTrace 保存最后一轮代理消息轨迹与助手输出摘要。
@@ -796,6 +850,62 @@ func (db *DB) GetMessages(conversationID string) ([]Message, error) {
return messages, nil
}
// GetMessagesLite 获取对话消息(不含 reasoning_content),用于历史会话快速切换。
func (db *DB) GetMessagesLite(conversationID string) ([]Message, error) {
rows, err := db.Query(
"SELECT id, conversation_id, role, content, mcp_execution_ids, created_at, updated_at FROM messages WHERE conversation_id = ? ORDER BY created_at ASC, rowid ASC",
conversationID,
)
if err != nil {
return nil, fmt.Errorf("查询消息失败: %w", err)
}
defer rows.Close()
var messages []Message
for rows.Next() {
var msg Message
var mcpIDsJSON sql.NullString
var createdAt string
var updatedAt sql.NullString
if err := rows.Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &mcpIDsJSON, &createdAt, &updatedAt); err != nil {
return nil, fmt.Errorf("扫描消息失败: %w", err)
}
var err error
msg.CreatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err != nil {
msg.CreatedAt, err = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err != nil {
msg.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
}
if updatedAt.Valid && strings.TrimSpace(updatedAt.String) != "" {
msg.UpdatedAt, err = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt.String)
if err != nil {
msg.UpdatedAt, err = time.Parse("2006-01-02 15:04:05", updatedAt.String)
}
if err != nil {
msg.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt.String)
}
}
if msg.UpdatedAt.IsZero() {
msg.UpdatedAt = msg.CreatedAt
}
if mcpIDsJSON.Valid && mcpIDsJSON.String != "" {
if err := json.Unmarshal([]byte(mcpIDsJSON.String), &msg.MCPExecutionIDs); err != nil {
db.logger.Warn("解析MCP执行ID失败", zap.Error(err))
}
}
messages = append(messages, msg)
}
return messages, nil
}
// turnSliceRange 根据任意一条消息 ID 定位「一轮对话」在 msgs 中的 [start, end) 下标区间(msgs 须已按时间升序,与 GetMessages 一致)。
// 一轮 = 从某条 user 消息起,至下一条 user 之前(含中间所有 assistant)。
func turnSliceRange(msgs []Message, anchorID string) (start, end int, err error) {
@@ -964,6 +1074,107 @@ func (db *DB) GetProcessDetails(messageID string) ([]ProcessDetail, error) {
return details, nil
}
// ProcessDetailsSummary 过程详情摘要(用于折叠态展示,避免全量加载)。
type ProcessDetailsSummary struct {
Total int `json:"total"`
IterationCount int `json:"iterationCount"`
MaxIteration int `json:"maxIteration"`
}
// GetProcessDetailsSummary 统计消息的过程详情数量与迭代轮次。
func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary, error) {
var total int
if err := db.QueryRow(
"SELECT COUNT(*) FROM process_details WHERE message_id = ?",
messageID,
).Scan(&total); err != nil {
return nil, fmt.Errorf("统计过程详情失败: %w", err)
}
summary := &ProcessDetailsSummary{Total: total}
if total == 0 {
return summary, nil
}
rows, err := db.Query(
"SELECT data FROM process_details WHERE message_id = ? AND event_type = 'iteration' ORDER BY created_at ASC, rowid ASC",
messageID,
)
if err != nil {
return nil, fmt.Errorf("查询迭代详情失败: %w", err)
}
defer rows.Close()
maxIter := 0
iterCount := 0
for rows.Next() {
var dataJSON string
if err := rows.Scan(&dataJSON); err != nil {
return nil, fmt.Errorf("扫描迭代详情失败: %w", err)
}
iterCount++
if dataJSON == "" {
continue
}
var payload map[string]interface{}
if err := json.Unmarshal([]byte(dataJSON), &payload); err != nil {
continue
}
if n, ok := payload["iteration"].(float64); ok && int(n) > maxIter {
maxIter = int(n)
}
}
summary.IterationCount = iterCount
summary.MaxIteration = maxIter
return summary, nil
}
// GetProcessDetailsPage 分页获取消息的过程详情(按时间升序)。
func (db *DB) GetProcessDetailsPage(messageID string, limit, offset int) ([]ProcessDetail, int, error) {
var total int
if err := db.QueryRow(
"SELECT COUNT(*) FROM process_details WHERE message_id = ?",
messageID,
).Scan(&total); err != nil {
return nil, 0, fmt.Errorf("统计过程详情失败: %w", err)
}
if total == 0 || offset >= total {
return nil, total, nil
}
rows, err := db.Query(
"SELECT id, message_id, conversation_id, event_type, message, data, created_at FROM process_details WHERE message_id = ? ORDER BY created_at ASC, rowid ASC LIMIT ? OFFSET ?",
messageID, limit, offset,
)
if err != nil {
return nil, 0, fmt.Errorf("查询过程详情失败: %w", err)
}
defer rows.Close()
var details []ProcessDetail
for rows.Next() {
var detail ProcessDetail
var createdAt string
if err := rows.Scan(&detail.ID, &detail.MessageID, &detail.ConversationID, &detail.EventType, &detail.Message, &detail.Data, &createdAt); err != nil {
return nil, 0, fmt.Errorf("扫描过程详情失败: %w", err)
}
var parseErr error
detail.CreatedAt, parseErr = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if parseErr != nil {
detail.CreatedAt, parseErr = time.Parse("2006-01-02 15:04:05", createdAt)
}
if parseErr != nil {
detail.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
}
details = append(details, detail)
}
return details, total, nil
}
// GetProcessDetailsByConversation 获取对话的所有过程详情(按消息分组)
func (db *DB) GetProcessDetailsByConversation(conversationID string) (map[string][]ProcessDetail, error) {
rows, err := db.Query(
+53 -2
View File
@@ -19,7 +19,9 @@ func TestDeleteConversationRemovesEinoScopedDirs(t *testing.T) {
plantaskBase := filepath.Join(tmp, "skills", ".eino", "plantask")
checkpointBase := filepath.Join(tmp, "eino-checkpoints")
db.SetEinoConversationDirs(plantaskBase, checkpointBase)
reductionBase := filepath.Join(tmp, "reduction")
workspaceBase := filepath.Join(tmp, "workspace")
db.SetEinoConversationDirs(plantaskBase, checkpointBase, reductionBase, workspaceBase)
conv, err := db.CreateConversation("cleanup test", ConversationCreateMeta{})
if err != nil {
@@ -34,6 +36,8 @@ func TestDeleteConversationRemovesEinoScopedDirs(t *testing.T) {
{db.conversationArtifactsDir, "transcript.txt"},
{plantaskBase, "task-1.json"},
{checkpointBase, "runner-deep.ckpt"},
{filepath.Join(reductionBase, "conversations"), "tool-output.txt"},
{filepath.Join(workspaceBase, "conversations"), "page.html"},
} {
dir := filepath.Join(base.root, seg)
if err := os.MkdirAll(dir, 0o755); err != nil {
@@ -48,10 +52,57 @@ func TestDeleteConversationRemovesEinoScopedDirs(t *testing.T) {
t.Fatalf("DeleteConversation: %v", err)
}
for _, base := range []string{db.conversationArtifactsDir, plantaskBase, checkpointBase} {
for _, base := range []string{db.conversationArtifactsDir, plantaskBase, checkpointBase, filepath.Join(reductionBase, "conversations"), filepath.Join(workspaceBase, "conversations")} {
dir := filepath.Join(base, seg)
if _, statErr := os.Stat(dir); !os.IsNotExist(statErr) {
t.Fatalf("expected removed dir %s, stat err=%v", dir, statErr)
}
}
}
func TestDeleteProjectRemovesReductionDir(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()
reductionBase := filepath.Join(tmp, "reduction")
workspaceBase := filepath.Join(tmp, "workspace")
db.SetEinoConversationDirs("", "", reductionBase, workspaceBase)
project, err := db.CreateProject(&Project{Name: "cleanup test"})
if err != nil {
t.Fatalf("CreateProject: %v", err)
}
seg := sanitizeConversationPathSegment(project.ID)
reductionDir := filepath.Join(reductionBase, "projects", seg, "clear")
if err := os.MkdirAll(reductionDir, 0o755); err != nil {
t.Fatalf("mkdir %s: %v", reductionDir, err)
}
if err := os.WriteFile(filepath.Join(reductionDir, "call-1.txt"), []byte("x"), 0o644); err != nil {
t.Fatalf("write: %v", err)
}
workspaceDir := filepath.Join(workspaceBase, "projects", seg, "downloads")
if err := os.MkdirAll(workspaceDir, 0o755); err != nil {
t.Fatalf("mkdir %s: %v", workspaceDir, err)
}
if err := os.WriteFile(filepath.Join(workspaceDir, "app.js"), []byte("x"), 0o644); err != nil {
t.Fatalf("write workspace: %v", err)
}
if err := db.DeleteProject(project.ID); err != nil {
t.Fatalf("DeleteProject: %v", err)
}
projectReductionDir := filepath.Join(reductionBase, "projects", seg)
if _, statErr := os.Stat(projectReductionDir); !os.IsNotExist(statErr) {
t.Fatalf("expected removed dir %s, stat err=%v", projectReductionDir, statErr)
}
projectWorkspaceDir := filepath.Join(workspaceBase, "projects", seg)
if _, statErr := os.Stat(projectWorkspaceDir); !os.IsNotExist(statErr) {
t.Fatalf("expected removed dir %s, stat err=%v", projectWorkspaceDir, statErr)
}
}
+47 -1
View File
@@ -51,6 +51,8 @@ type DB struct {
conversationArtifactsDir string
einoPlantaskBaseDir string // skills_dir + plantask_rel_dir (per-conversation subdirs)
einoCheckpointBaseDir string // checkpoint_dir root (per-conversation subdirs)
einoReductionRootDir string // reduction_root_dir or default tmp/reduction (conversations/<id> subdirs)
einoWorkspaceRootDir string // workspace_root_dir or default tmp/workspace (projects|conversations/<id> subdirs)
checkpointLoopName string
checkpointStop chan struct{}
checkpointDone chan struct{}
@@ -159,12 +161,16 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
// SetEinoConversationDirs configures best-effort filesystem cleanup on DeleteConversation.
// plantaskBase is skills_root/plantask_rel (no conversation id); checkpointBase is checkpoint_dir root.
func (db *DB) SetEinoConversationDirs(plantaskBase, checkpointBase string) {
// reductionRoot is reduction_root_dir from config; empty uses tmp/reduction (conversation-scoped subdirs only).
// workspaceRoot is agent.workspace_root_dir from config; empty uses tmp/workspace.
func (db *DB) SetEinoConversationDirs(plantaskBase, checkpointBase, reductionRoot, workspaceRoot string) {
if db == nil {
return
}
db.einoPlantaskBaseDir = strings.TrimSpace(plantaskBase)
db.einoCheckpointBaseDir = strings.TrimSpace(checkpointBase)
db.einoReductionRootDir = strings.TrimSpace(reductionRoot)
db.einoWorkspaceRootDir = strings.TrimSpace(workspaceRoot)
}
// initTables 初始化数据库表
@@ -353,6 +359,22 @@ func (db *DB) initTables() error {
UNIQUE(project_id, fact_key)
);`
// 项目事实关系边(黑板 DAG
createProjectFactEdgesTable := `
CREATE TABLE IF NOT EXISTS project_fact_edges (
id TEXT PRIMARY KEY,
project_id TEXT NOT NULL,
source_fact_key TEXT NOT NULL,
target_fact_key TEXT NOT NULL,
edge_type TEXT NOT NULL,
confidence TEXT NOT NULL DEFAULT 'tentative',
source_conversation_id TEXT,
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL,
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE,
UNIQUE(project_id, source_fact_key, target_fact_key, edge_type)
);`
// 创建漏洞表
createVulnerabilitiesTable := `
CREATE TABLE IF NOT EXISTS vulnerabilities (
@@ -389,6 +411,8 @@ func (db *DB) initTables() error {
last_schedule_trigger_at DATETIME,
last_schedule_error TEXT,
last_run_error TEXT,
project_id TEXT,
concurrency INTEGER NOT NULL DEFAULT 1,
status TEXT NOT NULL,
created_at DATETIME NOT NULL,
started_at DATETIME,
@@ -591,6 +615,9 @@ func (db *DB) initTables() error {
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence);
CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_project ON project_fact_edges(project_id);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_source ON project_fact_edges(project_id, source_fact_key);
CREATE INDEX IF NOT EXISTS idx_project_fact_edges_target ON project_fact_edges(project_id, target_fact_key);
CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id);
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id);
CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id);
@@ -672,6 +699,10 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建project_facts表失败: %w", err)
}
if _, err := db.Exec(createProjectFactEdgesTable); err != nil {
return fmt.Errorf("创建project_fact_edges表失败: %w", err)
}
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
}
@@ -1111,6 +1142,21 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
}
}
var concurrencyCount int
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='concurrency'").Scan(&concurrencyCount)
if err != nil {
if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN concurrency INTEGER NOT NULL DEFAULT 1"); addErr != nil {
errMsg := strings.ToLower(addErr.Error())
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
db.logger.Warn("添加batch_task_queues.concurrency字段失败", zap.Error(addErr))
}
}
} else if concurrencyCount == 0 {
if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN concurrency INTEGER NOT NULL DEFAULT 1"); err != nil {
db.logger.Warn("添加batch_task_queues.concurrency字段失败", zap.Error(err))
}
}
return nil
}
+70
View File
@@ -410,6 +410,76 @@ func (db *DB) GetToolExecutionsByIds(ids []string) ([]*mcp.ToolExecution, error)
return executions, nil
}
type toolExecutionStatDelta struct {
totalCalls int
successCalls int
failedCalls int
}
// PurgeToolExecutionsBefore deletes executions older than cutoff and adjusts tool_stats.
func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
query := `
SELECT tool_name, status, COUNT(*) AS cnt
FROM tool_executions
WHERE ` + sqliteEpochGE("start_time", "<") + `
GROUP BY tool_name, status
`
rows, err := db.Query(query, formatSQLiteUTC(cutoff))
if err != nil {
return 0, err
}
defer rows.Close()
deltas := make(map[string]*toolExecutionStatDelta)
for rows.Next() {
var toolName, status string
var count int
if err := rows.Scan(&toolName, &status, &count); err != nil {
db.logger.Warn("读取待清理执行记录统计失败", zap.Error(err))
continue
}
toolName = strings.TrimSpace(toolName)
if toolName == "" || count <= 0 {
continue
}
delta := deltas[toolName]
if delta == nil {
delta = &toolExecutionStatDelta{}
deltas[toolName] = delta
}
delta.totalCalls += count
switch status {
case "failed", "cancelled":
delta.failedCalls += count
case "completed":
delta.successCalls += count
}
}
if err := rows.Err(); err != nil {
return 0, err
}
res, err := db.Exec(`DELETE FROM tool_executions WHERE `+sqliteEpochGE("start_time", "<"), formatSQLiteUTC(cutoff))
if err != nil {
return 0, err
}
deleted, err := res.RowsAffected()
if err != nil {
return 0, err
}
for toolName, delta := range deltas {
if err := db.DecreaseToolStats(toolName, delta.totalCalls, delta.successCalls, delta.failedCalls); err != nil {
db.logger.Warn("清理过期执行记录后更新统计失败",
zap.Error(err),
zap.String("toolName", toolName),
)
}
}
return deleted, nil
}
// SaveToolStats 保存工具统计信息
func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
var lastCallTime sql.NullTime
+122
View File
@@ -0,0 +1,122 @@
package database
import (
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestPurgeToolExecutionsBefore(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()
oldStart := time.Now().AddDate(0, 0, -100)
newStart := time.Now().AddDate(0, 0, -1)
oldExec := &mcp.ToolExecution{
ID: "old-completed",
ToolName: "nmap::scan",
Arguments: map[string]interface{}{"target": "127.0.0.1"},
Status: "completed",
StartTime: oldStart,
}
oldFailed := &mcp.ToolExecution{
ID: "old-failed",
ToolName: "nmap::scan",
Arguments: map[string]interface{}{"target": "127.0.0.1"},
Status: "failed",
Error: "timeout",
StartTime: oldStart,
}
newExec := &mcp.ToolExecution{
ID: "new-completed",
ToolName: "nmap::scan",
Arguments: map[string]interface{}{"target": "127.0.0.1"},
Status: "completed",
StartTime: newStart,
}
for _, exec := range []*mcp.ToolExecution{oldExec, oldFailed, newExec} {
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution(%s): %v", exec.ID, err)
}
}
if err := db.UpdateToolStats("nmap::scan", 3, 2, 1, &newStart); err != nil {
t.Fatalf("UpdateToolStats: %v", err)
}
cutoff := time.Now().AddDate(0, 0, -90)
deleted, err := db.PurgeToolExecutionsBefore(cutoff)
if err != nil {
t.Fatalf("PurgeToolExecutionsBefore: %v", err)
}
if deleted != 2 {
t.Fatalf("deleted = %d, want 2", deleted)
}
if _, err := db.GetToolExecution("old-completed"); err == nil {
t.Fatal("old-completed should be deleted")
}
if _, err := db.GetToolExecution("old-failed"); err == nil {
t.Fatal("old-failed should be deleted")
}
if _, err := db.GetToolExecution("new-completed"); err != nil {
t.Fatalf("new-completed should remain: %v", err)
}
stats, err := db.LoadToolStats()
if err != nil {
t.Fatalf("LoadToolStats: %v", err)
}
stat := stats["nmap::scan"]
if stat == nil {
t.Fatal("expected stats for nmap::scan")
}
if stat.TotalCalls != 1 || stat.SuccessCalls != 1 || stat.FailedCalls != 0 {
t.Fatalf("stats after purge = %+v, want total=1 success=1 failed=0", stat)
}
total, err := db.CountToolExecutions("", "")
if err != nil {
t.Fatalf("CountToolExecutions: %v", err)
}
if total != 1 {
t.Fatalf("remaining executions = %d, want 1", total)
}
}
func TestPurgeToolExecutionsBefore_zeroRetentionSkipsViaService(t *testing.T) {
// RetentionDaysEffective: 0 means no purge at service layer; DB method still works when called directly.
dbPath := filepath.Join(t.TempDir(), "monitor.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
exec := &mcp.ToolExecution{
ID: "ancient",
ToolName: "curl::get",
Arguments: map[string]interface{}{},
Status: "completed",
StartTime: time.Now().AddDate(-1, 0, 0),
}
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
deleted, err := db.PurgeToolExecutionsBefore(time.Now())
if err != nil {
t.Fatalf("PurgeToolExecutionsBefore: %v", err)
}
if deleted != 1 {
t.Fatalf("deleted = %d, want 1", deleted)
}
}
+12 -4
View File
@@ -195,6 +195,7 @@ func (db *DB) DeleteProject(id string) error {
if err != nil {
return fmt.Errorf("删除项目失败: %w", err)
}
db.removeProjectScopedDirs(id)
return nil
}
@@ -389,7 +390,7 @@ func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
return f, nil
}
// DeprecateProjectFact 将事实标记为 deprecated。
// DeprecateProjectFact 将事实标记为 deprecated(关联边同步 deprecated
func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
res, err := db.Exec(
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
@@ -402,7 +403,7 @@ func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
if n == 0 {
return fmt.Errorf("事实不存在")
}
return nil
return db.DeprecateProjectFactEdgesForKey(projectID, factKey)
}
// RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。
@@ -430,9 +431,16 @@ func (db *DB) RestoreProjectFact(projectID, factKey, confidence string) error {
return err
}
// DeleteProjectFact 删除事实。
// DeleteProjectFact 删除事实(级联删除相关边)
func (db *DB) DeleteProjectFact(id string) error {
_, err := db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
f, err := db.GetProjectFact(id)
if err != nil {
return err
}
if err := db.DeleteProjectFactEdgesForKey(f.ProjectID, f.FactKey); err != nil {
return err
}
_, err = db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
return err
}
+410
View File
@@ -0,0 +1,410 @@
package database
import (
"database/sql"
"fmt"
"strings"
"time"
"github.com/google/uuid"
)
// ValidProjectFactEdgeTypes 项目事实图允许的边类型。
var ValidProjectFactEdgeTypes = map[string]struct{}{
"depends_on": {},
"leads_to": {},
"enables": {},
"exploits": {},
"discovered_on": {},
"contains": {},
"part_of": {},
"supports": {},
}
// ProjectFactEdge 项目事实关系边(source → target)。
type ProjectFactEdge struct {
ID string `json:"id"`
ProjectID string `json:"project_id"`
SourceFactKey string `json:"source_fact_key"`
TargetFactKey string `json:"target_fact_key"`
EdgeType string `json:"edge_type"`
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
SourceConversationID string `json:"source_conversation_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ProjectFactEdgeInput 写入边时的输入(出边:source → To)。
type ProjectFactEdgeInput struct {
To string `json:"to"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
// ProjectFactEdgeFromInput 写入入边时的输入(From → 当前事实)。
type ProjectFactEdgeFromInput struct {
From string `json:"from"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
// ProjectFactGraphNode 图 API 节点。
type ProjectFactGraphNode struct {
ID string `json:"id"`
FactKey string `json:"fact_key"`
Category string `json:"category"`
Label string `json:"label"` // 图节点短标签(截断)
Summary string `json:"summary"` // 完整摘要(侧栏等详情用)
Confidence string `json:"confidence"`
Type string `json:"type"`
Pinned bool `json:"pinned"`
}
// ProjectFactGraphEdge 图 API 边。
type ProjectFactGraphEdge struct {
ID string `json:"id"`
Source string `json:"source"`
Target string `json:"target"`
Type string `json:"type"`
Confidence string `json:"confidence"`
}
// ProjectFactGraph 项目事实图。
type ProjectFactGraph struct {
Nodes []ProjectFactGraphNode `json:"nodes"`
Edges []ProjectFactGraphEdge `json:"edges"`
}
// ValidateProjectFactEdgeType 校验边类型。
func ValidateProjectFactEdgeType(edgeType string) error {
edgeType = strings.TrimSpace(strings.ToLower(edgeType))
if edgeType == "" {
return fmt.Errorf("edge type 不能为空")
}
if _, ok := ValidProjectFactEdgeTypes[edgeType]; !ok {
return fmt.Errorf("无效的 edge type: %s", edgeType)
}
return nil
}
func normalizeEdgeConfidence(confidence string) string {
confidence = strings.TrimSpace(strings.ToLower(confidence))
switch confidence {
case "confirmed", "deprecated":
return confidence
default:
return "tentative"
}
}
// ListProjectFactEdgesByProject 列出项目全部边。
func (db *DB) ListProjectFactEdgesByProject(projectID string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ?
ORDER BY created_at ASC, rowid ASC`,
projectID,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ListOutgoingProjectFactEdges 列出某事实的全部出边。
func (db *DB) ListOutgoingProjectFactEdges(projectID, sourceFactKey string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND source_fact_key = ?
ORDER BY created_at ASC, rowid ASC`,
projectID, sourceFactKey,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ListIncomingProjectFactEdges 列出某事实的全部入边。
func (db *DB) ListIncomingProjectFactEdges(projectID, targetFactKey string) ([]*ProjectFactEdge, error) {
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND target_fact_key = ?
ORDER BY created_at ASC, rowid ASC`,
projectID, targetFactKey,
)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFactEdges(rows)
}
// ReplaceOutgoingProjectFactEdges 替换某事实的全部出边(links 省略时不调用)。
func (db *DB) ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID string, inputs []ProjectFactEdgeInput) error {
sourceFactKey = strings.TrimSpace(sourceFactKey)
if sourceFactKey == "" {
return fmt.Errorf("source_fact_key 不能为空")
}
if _, err := db.Exec(
`DELETE FROM project_fact_edges WHERE project_id = ? AND source_fact_key = ?`,
projectID, sourceFactKey,
); err != nil {
return fmt.Errorf("清除旧边失败: %w", err)
}
for _, in := range inputs {
target := strings.TrimSpace(in.To)
if target == "" {
continue
}
if err := ValidateFactKey(target); err != nil {
return fmt.Errorf("target fact_key 无效 (%s): %w", target, err)
}
if target == sourceFactKey {
return fmt.Errorf("边不能指向自身: %s", sourceFactKey)
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return err
}
edge := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: sourceFactKey,
TargetFactKey: target,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
if err := db.insertProjectFactEdge(edge); err != nil {
return err
}
}
return nil
}
// ReplaceIncomingProjectFactEdges 替换某事实的全部入边(From 为来源 fact_key)。
func (db *DB) ReplaceIncomingProjectFactEdges(projectID, targetFactKey string, inputs []ProjectFactEdgeFromInput) error {
targetFactKey = strings.TrimSpace(targetFactKey)
if targetFactKey == "" {
return fmt.Errorf("target_fact_key 不能为空")
}
if _, err := db.Exec(
`DELETE FROM project_fact_edges WHERE project_id = ? AND target_fact_key = ?`,
projectID, targetFactKey,
); err != nil {
return fmt.Errorf("清除旧入边失败: %w", err)
}
for _, in := range inputs {
source := strings.TrimSpace(in.From)
if source == "" {
continue
}
if err := ValidateFactKey(source); err != nil {
return fmt.Errorf("source fact_key 无效 (%s): %w", source, err)
}
if source == targetFactKey {
return fmt.Errorf("边不能指向自身: %s", targetFactKey)
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return err
}
sourceConversationID := ""
if srcFact, err := db.GetProjectFactByKey(projectID, source); err == nil && srcFact != nil {
sourceConversationID = srcFact.SourceConversationID
}
edge := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: source,
TargetFactKey: targetFactKey,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
if err := db.insertProjectFactEdge(edge); err != nil {
return err
}
}
return nil
}
// GetProjectFactEdge 按 ID 获取边。
func (db *DB) GetProjectFactEdge(edgeID string) (*ProjectFactEdge, error) {
var e ProjectFactEdge
var createdAt, updatedAt string
err := db.QueryRow(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges WHERE id = ?`, edgeID,
).Scan(&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
&e.SourceConversationID, &createdAt, &updatedAt)
if err != nil {
return nil, fmt.Errorf("边不存在")
}
e.CreatedAt = parseDBTime(createdAt)
e.UpdatedAt = parseDBTime(updatedAt)
return &e, nil
}
// AddProjectFactEdge 新增单条边(已存在则更新 confidence)。
func (db *DB) AddProjectFactEdge(projectID string, in ProjectFactEdgeInput, sourceFactKey, sourceConversationID string) (*ProjectFactEdge, error) {
sourceFactKey = strings.TrimSpace(sourceFactKey)
target := strings.TrimSpace(in.To)
if sourceFactKey == "" || target == "" {
return nil, fmt.Errorf("source 与 target 必填")
}
if sourceFactKey == target {
return nil, fmt.Errorf("边不能指向自身")
}
if err := ValidateProjectFactEdgeType(in.Type); err != nil {
return nil, err
}
if err := ValidateFactKey(target); err != nil {
return nil, err
}
now := time.Now()
e := &ProjectFactEdge{
ID: uuid.New().String(),
ProjectID: projectID,
SourceFactKey: sourceFactKey,
TargetFactKey: target,
EdgeType: strings.ToLower(strings.TrimSpace(in.Type)),
Confidence: normalizeEdgeConfidence(in.Confidence),
SourceConversationID: sourceConversationID,
CreatedAt: now,
UpdatedAt: now,
}
_, err := db.Exec(
`INSERT INTO project_fact_edges (
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
source_conversation_id, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(project_id, source_fact_key, target_fact_key, edge_type)
DO UPDATE SET confidence = excluded.confidence, updated_at = excluded.updated_at`,
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
)
if err != nil {
return nil, fmt.Errorf("添加边失败: %w", err)
}
// 返回最新
rows, err := db.Query(
`SELECT id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
COALESCE(source_conversation_id,''), created_at, updated_at
FROM project_fact_edges
WHERE project_id = ? AND source_fact_key = ? AND target_fact_key = ? AND edge_type = ?`,
projectID, sourceFactKey, target, e.EdgeType,
)
if err != nil {
return e, nil
}
defer rows.Close()
list, err := scanProjectFactEdges(rows)
if err != nil || len(list) == 0 {
return e, nil
}
return list[0], nil
}
// DeleteProjectFactEdge 删除单条边。
func (db *DB) DeleteProjectFactEdge(edgeID string) error {
res, err := db.Exec(`DELETE FROM project_fact_edges WHERE id = ?`, edgeID)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n == 0 {
return fmt.Errorf("边不存在")
}
return nil
}
func (db *DB) insertProjectFactEdge(e *ProjectFactEdge) error {
_, err := db.Exec(
`INSERT INTO project_fact_edges (
id, project_id, source_fact_key, target_fact_key, edge_type, confidence,
source_conversation_id, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
e.ID, e.ProjectID, e.SourceFactKey, e.TargetFactKey, e.EdgeType, e.Confidence,
nullIfEmpty(e.SourceConversationID), e.CreatedAt, e.UpdatedAt,
)
if err != nil {
return fmt.Errorf("写入边失败: %w", err)
}
return nil
}
// RenameProjectFactKeyEdges 事实 key 变更时同步边上的引用。
func (db *DB) RenameProjectFactKeyEdges(projectID, oldKey, newKey string) error {
oldKey = strings.TrimSpace(oldKey)
newKey = strings.TrimSpace(newKey)
if oldKey == "" || newKey == "" || oldKey == newKey {
return nil
}
now := time.Now()
if _, err := db.Exec(
`UPDATE project_fact_edges SET source_fact_key = ?, updated_at = ?
WHERE project_id = ? AND source_fact_key = ?`,
newKey, now, projectID, oldKey,
); err != nil {
return err
}
_, err := db.Exec(
`UPDATE project_fact_edges SET target_fact_key = ?, updated_at = ?
WHERE project_id = ? AND target_fact_key = ?`,
newKey, now, projectID, oldKey,
)
return err
}
// DeleteProjectFactEdgesForKey 删除与某 fact_key 相关的全部边。
func (db *DB) DeleteProjectFactEdgesForKey(projectID, factKey string) error {
_, err := db.Exec(
`DELETE FROM project_fact_edges
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)`,
projectID, factKey, factKey,
)
return err
}
// DeprecateProjectFactEdgesForKey 将关联边标记为 deprecated。
func (db *DB) DeprecateProjectFactEdgesForKey(projectID, factKey string) error {
now := time.Now()
_, err := db.Exec(
`UPDATE project_fact_edges SET confidence = 'deprecated', updated_at = ?
WHERE project_id = ? AND (source_fact_key = ? OR target_fact_key = ?)
AND confidence != 'deprecated'`,
now, projectID, factKey, factKey,
)
return err
}
func scanProjectFactEdges(rows *sql.Rows) ([]*ProjectFactEdge, error) {
var out []*ProjectFactEdge
for rows.Next() {
var e ProjectFactEdge
var createdAt, updatedAt string
if err := rows.Scan(
&e.ID, &e.ProjectID, &e.SourceFactKey, &e.TargetFactKey, &e.EdgeType, &e.Confidence,
&e.SourceConversationID, &createdAt, &updatedAt,
); err != nil {
return nil, err
}
e.CreatedAt = parseDBTime(createdAt)
e.UpdatedAt = parseDBTime(updatedAt)
out = append(out, &e)
}
return out, rows.Err()
}
+169 -428
View File
@@ -21,7 +21,6 @@ import (
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/reasoning"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/multiagent"
"cyberstrike-ai/internal/openai"
@@ -178,8 +177,6 @@ type AgentHandler struct {
}
agentsMarkdownDir string // 多代理:Markdown 子 Agent 目录(绝对路径,空则不从磁盘合并)
batchCronParser cron.Parser
batchRunnerMu sync.Mutex
batchRunning map[string]struct{}
// hitlWhitelistSaver 侧栏「应用」HITL 时将会话增量白名单合并写入 config.yaml(可选)
hitlWhitelistSaver HitlToolWhitelistSaver
audit *audit.Service
@@ -190,6 +187,29 @@ func (h *AgentHandler) SetAudit(s *audit.Service) {
h.audit = s
}
// TaskManager 返回 Agent 任务管理器(供 MCP 监控页终止 Eino execute 等)。
func (h *AgentHandler) TaskManager() *AgentTaskManager {
if h == nil {
return nil
}
return h.tasks
}
// CancelRunningTaskForConversation stops any in-flight agent work for the conversation (idempotent).
func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
if h == nil || conversationID == "" || h.tasks == nil {
return
}
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
h.agent.CancelMCPToolExecutionWithNote(execID, "")
}
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
} else if err != nil {
h.logger.Warn("取消会话运行中任务失败", zap.String("conversationId", conversationID), zap.Error(err))
}
}
// HitlToolWhitelistSaver 合并 HITL 免审批工具到全局配置并落盘
type HitlToolWhitelistSaver interface {
MergeHitlToolWhitelistIntoConfig(add []string) error
@@ -218,7 +238,6 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
config: cfg,
hitlManager: NewHITLManager(db, logger),
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
batchRunning: make(map[string]struct{}),
}
if err := handler.hitlManager.EnsureSchema(); err != nil {
logger.Warn("初始化 HITL 表失败", zap.Error(err))
@@ -631,40 +650,11 @@ func (h *AgentHandler) runRobotEinoSingleWithRetry(
assistantMessageID string,
taskStatus *string,
) (string, string, error) {
curHist := history
curMsg := finalMessage
segmentUserMessage := finalMessage
var resultMA *multiagent.RunResult
var errMA error
var transientRunAttempts int
var emptyResponseAttempts int
for {
resultMA, errMA = multiagent.RunEinoSingleChatModelAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger,
conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
taskCtx, conversationID, resultMA, errMA, &emptyResponseAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
)
if exhaustedEmpty {
errMA = nil
break
}
if handledEmpty {
continue
}
if errMA == nil {
transientRunAttempts = 0
emptyResponseAttempts = 0
break
}
if handled, _ := h.handleEinoTransientRetryContinue(
taskCtx, conversationID, resultMA, errMA, &transientRunAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
); handled {
continue
}
resultMA, errMA := multiagent.RunEinoSingleChatModelAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger,
conversationID, h.conversationProjectID(conversationID), finalMessage, history, roleTools, progressCallback, nil, h.agentSessionContextBlock(conversationID),
)
if errMA != nil {
*taskStatus = "failed"
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
}
@@ -680,41 +670,12 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry(
assistantMessageID string,
taskStatus *string,
) (string, string, error) {
curHist := history
curMsg := finalMessage
segmentUserMessage := finalMessage
var resultMA *multiagent.RunResult
var errMA error
var transientRunAttempts int
var emptyResponseAttempts int
for {
resultMA, errMA = multiagent.RunDeepAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger,
conversationID, h.conversationProjectID(conversationID), curMsg, curHist, roleTools, progressCallback,
h.agentsMarkdownDir, orchestration, nil, h.projectBlackboardBlock(conversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
taskCtx, conversationID, resultMA, errMA, &emptyResponseAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
)
if exhaustedEmpty {
errMA = nil
break
}
if handledEmpty {
continue
}
if errMA == nil {
transientRunAttempts = 0
emptyResponseAttempts = 0
break
}
if handled, _ := h.handleEinoTransientRetryContinue(
taskCtx, conversationID, resultMA, errMA, &transientRunAttempts,
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
); handled {
continue
}
resultMA, errMA := multiagent.RunDeepAgent(
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger,
conversationID, h.conversationProjectID(conversationID), finalMessage, history, roleTools, progressCallback,
h.agentsMarkdownDir, orchestration, nil, h.agentSessionContextBlock(conversationID),
)
if errMA != nil {
*taskStatus = "failed"
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
}
@@ -1309,7 +1270,10 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
// 保存过程详情到数据库(排除 response/doneresponse 正文已在 messages 表)
// response_start/response_delta 已聚合为 planning,不落逐条。
// [Eino] agent 心跳 progress 仅用于实时进度标题,不落库以免时间线刷屏。
skipEinoAgentHeartbeat := eventType == "progress" && strings.HasPrefix(strings.TrimSpace(message), "[Eino] ")
if assistantMessageID != "" &&
!skipEinoAgentHeartbeat &&
eventType != "response" &&
eventType != "done" &&
eventType != "response_start" &&
@@ -1335,6 +1299,55 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
}
}
// cancelToolContinueAfter 仅终止当前工具调用,不停止整条 Agent 任务(对话「中断并继续」与 MCP 监控终止共用)。
func (h *AgentHandler) cancelToolContinueAfter(conversationID, preferredExecID, note string) (bool, gin.H) {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" || h.tasks.GetTask(conversationID) == nil {
return false, nil
}
note = strings.TrimSpace(note)
execID := strings.TrimSpace(preferredExecID)
if execID == "" {
execID = h.tasks.ActiveMCPExecutionID(conversationID)
}
if execID != "" {
if h.agent.CancelMCPToolExecutionWithNote(execID, note) {
return true, gin.H{
"status": "tool_abort_requested",
"conversationId": conversationID,
"executionId": execID,
"message": "已请求终止当前工具调用;工具返回后本轮推理将继续(与 MCP 监控页终止一致)。",
"continueAfter": true,
"interruptWithNote": note != "",
"continueWithoutTool": false,
}
}
if h.tasks.AbortActiveEinoExecute(conversationID, note) {
return true, gin.H{
"status": "tool_abort_requested",
"conversationId": conversationID,
"executionId": execID,
"message": "已请求终止当前 execute 命令;命令返回后本轮推理将继续。",
"continueAfter": true,
"interruptWithNote": note != "",
"continueWithoutTool": false,
}
}
return false, nil
}
if h.tasks.AbortActiveEinoExecute(conversationID, note) {
return true, gin.H{
"status": "tool_abort_requested",
"conversationId": conversationID,
"message": "已请求终止当前 execute 命令;命令返回后本轮推理将继续。",
"continueAfter": true,
"interruptWithNote": note != "",
"continueWithoutTool": false,
}
}
return false, nil
}
// CancelAgentLoop 取消正在执行的任务
func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
var req struct {
@@ -1353,27 +1366,20 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": "未找到正在执行的任务"})
return
}
execID := h.tasks.ActiveMCPExecutionID(req.ConversationID)
note := strings.TrimSpace(req.Reason)
if execID != "" {
if !h.agent.CancelMCPToolExecutionWithNote(execID, note) {
c.JSON(http.StatusNotFound, gin.H{"error": "未找到进行中的工具执行或该调用已结束"})
return
}
h.logger.Info("对话页仅终止当前 MCP 工具",
activeExec := strings.TrimSpace(h.tasks.ActiveMCPExecutionID(req.ConversationID))
if ok, payload := h.cancelToolContinueAfter(req.ConversationID, "", note); ok {
execID, _ := payload["executionId"].(string)
h.logger.Info("对话页仅终止当前工具",
zap.String("conversationId", req.ConversationID),
zap.String("executionId", execID),
zap.Bool("hasNote", note != ""),
)
c.JSON(http.StatusOK, gin.H{
"status": "tool_abort_requested",
"conversationId": req.ConversationID,
"executionId": execID,
"message": "已请求终止当前工具调用;工具返回后本轮推理将继续(与 MCP 监控页终止一致)。",
"continueAfter": true,
"interruptWithNote": note != "",
"continueWithoutTool": false,
})
c.JSON(http.StatusOK, payload)
return
}
if activeExec != "" {
c.JSON(http.StatusNotFound, gin.H{"error": "未找到进行中的工具执行或该调用已结束"})
return
}
// 无进行中的 MCP 工具(模型纯推理/流式输出阶段):取消当前上下文并由 Eino 流式处理器合并用户补充后自动续跑。
@@ -1495,6 +1501,7 @@ type BatchTaskRequest struct {
CronExpr string `json:"cronExpr,omitempty"` // scheduleMode=cron 时必填
ExecuteNow bool `json:"executeNow,omitempty"` // 创建后是否立即执行(默认 false)
ProjectID string `json:"projectId,omitempty"` // 队列内子对话绑定的项目(可选)
Concurrency int `json:"concurrency,omitempty"` // 同时执行的子任务数,默认 1,最大 8
}
// batchQueueWantsEino 队列是否配置为走 Eino 多代理。
@@ -1554,7 +1561,7 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
nextRunAt = &next
}
queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, req.ProjectID, nextRunAt, validTasks)
queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, req.ProjectID, nextRunAt, req.Concurrency, validTasks)
if createErr != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()})
return
@@ -1678,6 +1685,7 @@ func (h *AgentHandler) ListBatchQueues(c *gin.Context) {
// StartBatchQueue 开始执行批量任务队列
func (h *AgentHandler) StartBatchQueue(c *gin.Context) {
queueID := c.Param("queueId")
h.batchTaskManager.ClearSingleRunTask(queueID)
ok, err := h.startBatchQueueExecution(queueID, false)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -1709,6 +1717,7 @@ func (h *AgentHandler) RerunBatchQueue(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "重置队列失败"})
return
}
h.batchTaskManager.ClearSingleRunTask(queueID)
ok, err := h.startBatchQueueExecution(queueID, false)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -1742,15 +1751,16 @@ func (h *AgentHandler) PauseBatchQueue(c *gin.Context) {
func (h *AgentHandler) UpdateBatchQueueMetadata(c *gin.Context) {
queueID := c.Param("queueId")
var req struct {
Title string `json:"title"`
Role string `json:"role"`
AgentMode string `json:"agentMode"`
Title string `json:"title"`
Role string `json:"role"`
AgentMode string `json:"agentMode"`
Concurrency *int `json:"concurrency"`
}
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if err := h.batchTaskManager.UpdateQueueMetadata(queueID, req.Title, req.Role, req.AgentMode); err != nil {
if err := h.batchTaskManager.UpdateQueueMetadata(queueID, req.Title, req.Role, req.AgentMode, req.Concurrency); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
@@ -1825,9 +1835,17 @@ func (h *AgentHandler) SetBatchQueueScheduleEnabled(c *gin.Context) {
// DeleteBatchQueue 删除批量任务队列
func (h *AgentHandler) DeleteBatchQueue(c *gin.Context) {
queueID := c.Param("queueId")
success := h.batchTaskManager.DeleteQueue(queueID)
if !success {
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
if err := h.batchTaskManager.DeleteQueue(queueID); err != nil {
switch {
case errors.Is(err, ErrBatchQueueNotFound):
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
case errors.Is(err, ErrBatchQueueExecutorActive):
c.JSON(http.StatusConflict, gin.H{"error": "队列执行器仍在运行,请稍后再删除"})
case errors.Is(err, ErrBatchQueueStillRunning):
c.JSON(http.StatusConflict, gin.H{"error": "队列正在运行中,无法删除"})
default:
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
}
return
}
if h.audit != nil {
@@ -1908,6 +1926,53 @@ func (h *AgentHandler) AddBatchTask(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"message": "任务已添加", "task": task, "queue": queue})
}
// RunSingleBatchTask 单条执行指定子任务(可覆盖已成功项),完成后暂停队列
func (h *AgentHandler) RunSingleBatchTask(c *gin.Context) {
queueID := c.Param("queueId")
taskID := c.Param("taskId")
if err := h.batchTaskManager.PrepareSingleTaskRun(queueID, taskID); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.batchTaskManager.SetSingleRunTask(queueID, taskID)
// 暂停态单条执行:旧批量协程可能仍占用执行槽,先回收以便重新启动
if queue, ok := h.batchTaskManager.GetBatchQueue(queueID); ok && queue.Status == BatchQueueStatusPaused {
h.batchTaskManager.ForceUnmarkQueueExecutor(queueID)
}
autoStarted := true
autoStartMsg := "已开始单条执行"
ok, startErr := h.startBatchQueueExecution(queueID, false)
if startErr != nil {
h.batchTaskManager.ClearSingleRunTask(queueID)
autoStarted = false
autoStartMsg = "任务已准备就绪,但自动启动失败: " + startErr.Error()
} else if !ok {
h.batchTaskManager.ClearSingleRunTask(queueID)
autoStarted = false
autoStartMsg = "任务已准备就绪,但队列不存在"
}
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
return
}
if h.audit != nil {
h.audit.RecordOK(c, "task", "run_single_batch_task", "单条执行批量子任务", "batch_task", taskID, map[string]interface{}{
"batch_queue_id": queueID,
"auto_started": autoStarted,
})
}
c.JSON(http.StatusOK, gin.H{
"message": autoStartMsg,
"queue": queue,
"autoStarted": autoStarted,
})
}
// DeleteBatchTask 删除批量任务
func (h *AgentHandler) DeleteBatchTask(c *gin.Context) {
queueID := c.Param("queueId")
@@ -1933,22 +1998,6 @@ func (h *AgentHandler) DeleteBatchTask(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"message": "任务已删除", "queue": queue})
}
func (h *AgentHandler) markBatchQueueRunning(queueID string) bool {
h.batchRunnerMu.Lock()
defer h.batchRunnerMu.Unlock()
if _, exists := h.batchRunning[queueID]; exists {
return false
}
h.batchRunning[queueID] = struct{}{}
return true
}
func (h *AgentHandler) unmarkBatchQueueRunning(queueID string) {
h.batchRunnerMu.Lock()
defer h.batchRunnerMu.Unlock()
delete(h.batchRunning, queueID)
}
func (h *AgentHandler) nextBatchQueueRunAt(cronExpr string, from time.Time) (*time.Time, error) {
expr := strings.TrimSpace(cronExpr)
if expr == "" {
@@ -1964,43 +2013,43 @@ func (h *AgentHandler) nextBatchQueueRunAt(cronExpr string, from time.Time) (*ti
func (h *AgentHandler) startBatchQueueExecution(queueID string, scheduled bool) (bool, error) {
// 先获取执行互斥门,再读取队列状态,避免基于过时快照做判断
if !h.markBatchQueueRunning(queueID) {
if !h.batchTaskManager.TryMarkQueueExecutor(queueID) {
return true, nil
}
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
return false, nil
}
if scheduled {
if queue.ScheduleMode != "cron" {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
err := fmt.Errorf("队列未启用 cron 调度")
h.batchTaskManager.SetLastScheduleError(queueID, err.Error())
return true, err
}
if queue.Status == "running" || queue.Status == "paused" || queue.Status == "cancelled" {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
err := fmt.Errorf("当前队列状态不允许被调度执行")
h.batchTaskManager.SetLastScheduleError(queueID, err.Error())
return true, err
}
if !h.batchTaskManager.ResetQueueForRerun(queueID) {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
err := fmt.Errorf("重置队列失败")
h.batchTaskManager.SetLastScheduleError(queueID, err.Error())
return true, err
}
queue, _ = h.batchTaskManager.GetBatchQueue(queueID)
} else if queue.Status != "pending" && queue.Status != "paused" {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
return true, fmt.Errorf("队列状态不允许启动")
}
if queue != nil && batchQueueWantsEino(queue.AgentMode) && (h.config == nil || !h.config.MultiAgent.Enabled) {
h.unmarkBatchQueueRunning(queueID)
h.batchTaskManager.UnmarkQueueExecutor(queueID)
err := fmt.Errorf("当前队列配置为 Eino 多代理,但系统未启用多代理")
if scheduled {
h.batchTaskManager.SetLastScheduleError(queueID, err.Error())
@@ -2052,314 +2101,6 @@ func (h *AgentHandler) batchQueueSchedulerLoop() {
}
}
// executeBatchQueue 执行批量任务队列
func (h *AgentHandler) executeBatchQueue(queueID string) {
defer h.unmarkBatchQueueRunning(queueID)
h.logger.Info("开始执行批量任务队列", zap.String("queueId", queueID))
for {
// 检查队列状态
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists || queue.Status == "cancelled" || queue.Status == "completed" || queue.Status == "paused" {
break
}
// 获取下一个任务
task, hasNext := h.batchTaskManager.GetNextTask(queueID)
if !hasNext {
// 所有任务完成:汇总子任务失败信息便于排障
q, ok := h.batchTaskManager.GetBatchQueue(queueID)
lastRunErr := ""
if ok {
for _, t := range q.Tasks {
if t.Status == "failed" && t.Error != "" {
lastRunErr = t.Error
}
}
}
h.batchTaskManager.SetLastRunError(queueID, lastRunErr)
h.batchTaskManager.UpdateQueueStatus(queueID, "completed")
h.logger.Info("批量任务队列执行完成", zap.String("queueId", queueID))
break
}
// 更新任务状态为运行中
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "running", "", "")
// 创建新对话
title := safeTruncateString(task.Message, 50)
batchMeta := audit.ConversationCreateMeta("batch_task")
batchMeta.ProjectID = effectiveProjectID(h.config, queue.ProjectID)
conv, err := h.db.CreateConversation(title, batchMeta)
var conversationID string
if err != nil {
h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "failed", "", "创建对话失败: "+err.Error())
h.batchTaskManager.MoveToNextTask(queueID)
continue
}
conversationID = conv.ID
// 保存conversationId到任务中(即使是运行中状态也要保存,以便查看对话)
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "running", "", "", conversationID)
// 应用角色用户提示词和工具配置
finalMessage := task.Message
var roleTools []string // 角色配置的工具列表
if queue.Role != "" && queue.Role != "默认" {
if h.config.Roles != nil {
if role, exists := h.config.Roles[queue.Role]; exists && role.Enabled {
// 应用用户提示词
if role.UserPrompt != "" {
finalMessage = role.UserPrompt + "\n\n" + task.Message
h.logger.Info("应用角色用户提示词", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role))
}
// 获取角色配置的工具列表(优先使用tools字段,向后兼容mcps字段)
if len(role.Tools) > 0 {
roleTools = role.Tools
h.logger.Info("使用角色配置的工具列表", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("toolCount", len(roleTools)))
}
}
}
}
// 保存用户消息(保存原始消息,不包含角色提示词)
_, err = h.db.AddMessage(conversationID, "user", task.Message, nil)
if err != nil {
h.logger.Error("保存用户消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
// 预先创建助手消息,以便关联过程详情
assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil)
if err != nil {
h.logger.Error("创建助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
// 如果创建失败,继续执行但不保存过程详情
assistantMsg = nil
}
// 创建进度回调函数,复用统一逻辑(批量任务不需要流式事件,所以传入nil)
var assistantMessageID string
if assistantMsg != nil {
assistantMessageID = assistantMsg.ID
}
// 注意:批量任务没有前端直连的 POST /stream,因此若要支持「刷新后补流」,
// 需要把进度事件镜像到 TaskEventBusGET /api/agent-loop/task-events 会订阅这里)。
// progressCallback 将在子任务的 IIFE 内创建,以便拿到 taskCtx/cancelWithCause 与 sendEvent。
var progressCallback func(eventType, message string, data interface{})
// 执行任务(使用包含角色提示词的finalMessage和角色工具列表)
h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID))
func() {
// 与对话流式接口一致:同 conversationId 仅允许一个运行中任务,并支持 /api/agent-loop/cancel 与会话锁对齐。
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
// 单个子任务超时:6 小时(与原先 WithTimeout(Background) 一致)
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour)
registered := false
finishStatus := "completed"
defer func() {
h.batchTaskManager.SetTaskCancel(queueID, nil)
timeoutCancel()
if registered {
// 与流式接口保持一致:结束前补一个 done,便于前端 task-events 侧及时收口 UI。
if h.taskEventBus != nil {
ev := StreamEvent{Type: "done", Message: "", Data: map[string]interface{}{"conversationId": conversationID}}
if b, err := json.Marshal(ev); err == nil {
h.taskEventBus.Publish(conversationID, append(append([]byte("data: "), b...), '\n', '\n'))
}
}
h.tasks.FinishTask(conversationID, finishStatus)
}
cancelWithCause(nil)
}()
// 事件镜像:只发布到 TaskEventBus,不直接写 HTTP Response(用于刷新后的补流)。
sendEvent := func(eventType, message string, data interface{}) {
if h.taskEventBus == nil {
return
}
ev := StreamEvent{Type: eventType, Message: message, Data: data}
b, err := json.Marshal(ev)
if err != nil {
b = []byte(`{"type":"error","message":"marshal failed"}`)
}
line := make([]byte, 0, len(b)+8)
line = append(line, []byte("data: ")...)
line = append(line, b...)
line = append(line, '\n', '\n')
h.taskEventBus.Publish(conversationID, line)
}
if _, err := h.tasks.StartTask(conversationID, task.Message, cancelWithCause); err != nil {
h.logger.Warn("批量队列子任务注册会话运行状态失败",
zap.String("queueId", queueID),
zap.String("taskId", task.ID),
zap.String("conversationId", conversationID),
zap.Error(err))
failMsg := err.Error()
if errors.Is(err, ErrTaskAlreadyRunning) {
failMsg = "会话已有任务正在执行,无法在该会话上并行启动批量子任务"
}
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "failed", "", failMsg)
return
}
registered = true
// 存储取消函数:暂停队列时取消子任务 context(与原先语义一致)
h.batchTaskManager.SetTaskCancel(queueID, timeoutCancel)
// 创建进度回调函数:写 DB + 镜像到 task-events,支持刷新后继续流式展示。
progressCallback = h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
taskCtx = mcp.WithMCPConversationID(taskCtx, conversationID)
taskCtx = mcp.WithToolRunRegistry(taskCtx, h.tasks)
// 使用队列配置的角色工具列表(如果为空,表示使用所有工具)
useBatchMulti := false
batchOrch := "deep"
am := strings.TrimSpace(strings.ToLower(queue.AgentMode))
if am == "multi" {
am = "deep"
}
if batchQueueWantsEino(queue.AgentMode) && h.config != nil && h.config.MultiAgent.Enabled {
useBatchMulti = true
batchOrch = config.NormalizeMultiAgentOrchestration(am)
} else if queue.AgentMode == "" && h.config != nil && h.config.MultiAgent.Enabled && h.config.MultiAgent.BatchUseMultiAgent {
// 兼容历史数据:未配置队列代理模式时,沿用旧的系统级开关
useBatchMulti = true
batchOrch = "deep"
}
var resultMA *multiagent.RunResult
var runErr error
switch {
case useBatchMulti:
resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID))
default:
if h.config == nil {
runErr = fmt.Errorf("服务器配置未加载")
} else {
resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID))
}
}
if runErr != nil {
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(conversationID, resultMA)
}
errStr := runErr.Error()
partialResp := ""
if resultMA != nil {
partialResp = resultMA.Response
}
isCancelled := errors.Is(context.Cause(baseCtx), ErrTaskCancelled) ||
errors.Is(runErr, context.Canceled) ||
strings.Contains(strings.ToLower(errStr), "context canceled") ||
strings.Contains(strings.ToLower(errStr), "context cancelled") ||
(partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断")))
isTimeout := errors.Is(runErr, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded)
if isTimeout {
finishStatus = "timeout"
} else if isCancelled {
finishStatus = "cancelled"
} else {
finishStatus = "failed"
}
if isCancelled {
h.logger.Info("批量任务被取消", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
cancelMsg := "任务已被用户取消,后续操作已停止。"
// 如果执行结果中有更具体的取消消息,使用它
if partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断")) {
cancelMsg = partialResp
}
// 更新助手消息内容
if assistantMessageID != "" {
if updateErr := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); updateErr != nil {
h.logger.Warn("更新取消后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
}
// 保存取消详情到数据库
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil); err != nil {
h.logger.Warn("保存取消详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
}
} else {
// 如果没有预先创建的助手消息,创建一个新的
_, errMsg := h.db.AddMessage(conversationID, "assistant", cancelMsg, nil)
if errMsg != nil {
h.logger.Warn("保存取消消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(errMsg))
}
}
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "cancelled", cancelMsg, "", conversationID)
} else {
h.logger.Error("批量任务执行失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(runErr))
errorMsg := "执行失败: " + runErr.Error()
// 更新助手消息内容
if assistantMessageID != "" {
if _, updateErr := h.db.Exec(
"UPDATE messages SET content = ?, updated_at = ? WHERE id = ?",
errorMsg,
time.Now(), assistantMessageID,
); updateErr != nil {
h.logger.Warn("更新失败后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
}
// 保存错误详情到数据库
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errorMsg, nil); err != nil {
h.logger.Warn("保存错误详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
}
}
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, "failed", "", runErr.Error())
}
} else {
h.logger.Info("批量任务执行成功", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
resText := resultMA.Response
mcpIDs := resultMA.MCPExecutionIDs
lastIn := resultMA.LastAgentTraceInput
lastOut := resultMA.LastAgentTraceOutput
// 更新助手消息内容
if assistantMessageID != "" {
if updateErr := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resText, mcpIDs, multiagent.AggregatedReasoningFromTraceJSON(lastIn)); updateErr != nil {
h.logger.Warn("更新助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
// 如果更新失败,尝试创建新消息
_, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs)
if err != nil {
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
}
} else {
// 如果没有预先创建的助手消息,创建一个新的
_, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs)
if err != nil {
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
}
// 保存代理轨迹
if lastIn != "" || lastOut != "" {
if err := h.db.SaveAgentTrace(conversationID, lastIn, lastOut); err != nil {
h.logger.Warn("保存代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
} else {
h.logger.Info("已保存代理轨迹", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
}
}
// 保存结果
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "completed", resText, "", conversationID)
}
}()
// 移动到下一个任务
h.batchTaskManager.MoveToNextTask(queueID)
// 检查是否被取消或暂停
queue, _ = h.batchTaskManager.GetBatchQueue(queueID)
if queue.Status == "cancelled" || queue.Status == "paused" {
break
}
}
}
// loadHistoryFromAgentTrace 从库中保存的代理消息轨迹恢复历史(列 last_react_*;含单代理与 Eino)。
// 逻辑与攻击链一致:优先用已保存的 JSON 消息带 + 最后一轮助手摘要,否则回退消息表。
func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent.ChatMessage, error) {
+352
View File
@@ -0,0 +1,352 @@
package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/audit"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/multiagent"
"go.uber.org/zap"
)
const batchQueueWorkerIdlePoll = 200 * time.Millisecond
// executeBatchQueue 使用并发 worker 池执行批量任务队列。
func (h *AgentHandler) executeBatchQueue(queueID string) {
defer h.batchTaskManager.UnmarkQueueExecutor(queueID)
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists {
return
}
concurrency := normalizeBatchQueueConcurrency(queue.Concurrency)
h.logger.Info("开始执行批量任务队列", zap.String("queueId", queueID), zap.Int("concurrency", concurrency))
var wg sync.WaitGroup
for i := 0; i < concurrency; i++ {
wg.Add(1)
go func() {
defer wg.Done()
h.runBatchQueueWorker(queueID)
}()
}
wg.Wait()
h.tryFinalizeBatchQueue(queueID)
}
func (h *AgentHandler) runBatchQueueWorker(queueID string) {
for {
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if batchQueueExecutionShouldStop(queue, exists) {
return
}
task, ok := h.batchTaskManager.ClaimNextPendingTask(queueID)
if !ok {
if !h.batchTaskManager.HasRunningTasks(queueID) {
return
}
time.Sleep(batchQueueWorkerIdlePoll)
continue
}
queue, _ = h.batchTaskManager.GetBatchQueue(queueID)
if queue == nil {
return
}
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusRunning, "", "")
h.executeOneBatchSubTask(queueID, queue, task)
if h.batchTaskManager.TakeSingleRunTaskIfMatch(queueID, task.ID) {
h.batchTaskManager.UpdateQueueStatus(queueID, BatchQueueStatusPaused)
h.logger.Info("单条执行完成,队列已暂停", zap.String("queueId", queueID), zap.String("taskId", task.ID))
return
}
queue, exists = h.batchTaskManager.GetBatchQueue(queueID)
if batchQueueExecutionShouldStop(queue, exists) {
if !exists {
h.logger.Warn("批量队列在执行收尾时已不存在,安全退出", zap.String("queueId", queueID))
}
return
}
}
}
func (h *AgentHandler) tryFinalizeBatchQueue(queueID string) {
queue, exists := h.batchTaskManager.GetBatchQueue(queueID)
if !exists || queue == nil {
return
}
if queue.Status != BatchQueueStatusRunning {
return
}
if h.batchTaskManager.HasPendingOrRunningTasks(queueID) {
return
}
lastRunErr := ""
for _, t := range queue.Tasks {
if t != nil && t.Status == BatchTaskStatusFailed && t.Error != "" {
lastRunErr = t.Error
}
}
h.batchTaskManager.SetLastRunError(queueID, lastRunErr)
h.batchTaskManager.UpdateQueueStatus(queueID, BatchQueueStatusCompleted)
h.logger.Info("批量任务队列执行完成", zap.String("queueId", queueID))
}
// executeOneBatchSubTask 执行单条批量子任务(各自独立会话)。
func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQueue, task *BatchTask) {
title := safeTruncateString(task.Message, 50)
batchMeta := audit.ConversationCreateMeta("batch_task")
batchMeta.ProjectID = effectiveProjectID(h.config, queue.ProjectID)
conv, err := h.db.CreateConversation(title, batchMeta)
if err != nil {
h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "创建对话失败: "+err.Error())
return
}
conversationID := conv.ID
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusRunning, "", "", conversationID)
finalMessage := task.Message
var roleTools []string
if queue.Role != "" && queue.Role != "默认" {
if h.config.Roles != nil {
if role, exists := h.config.Roles[queue.Role]; exists && role.Enabled {
if role.UserPrompt != "" {
finalMessage = role.UserPrompt + "\n\n" + task.Message
h.logger.Info("应用角色用户提示词", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role))
}
if len(role.Tools) > 0 {
roleTools = role.Tools
h.logger.Info("使用角色配置的工具列表", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("toolCount", len(roleTools)))
}
}
}
}
if _, err = h.db.AddMessage(conversationID, "user", task.Message, nil); err != nil {
h.logger.Error("保存用户消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil)
if err != nil {
h.logger.Error("创建助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
assistantMsg = nil
}
var assistantMessageID string
if assistantMsg != nil {
assistantMessageID = assistantMsg.ID
}
h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID))
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour)
registered := false
finishStatus := "completed"
defer func() {
h.batchTaskManager.SetTaskCancel(queueID, task.ID, nil)
timeoutCancel()
if registered {
if h.taskEventBus != nil {
ev := StreamEvent{Type: "done", Message: "", Data: map[string]interface{}{"conversationId": conversationID}}
if b, err := json.Marshal(ev); err == nil {
h.taskEventBus.Publish(conversationID, append(append([]byte("data: "), b...), '\n', '\n'))
}
}
h.tasks.FinishTask(conversationID, finishStatus)
}
cancelWithCause(nil)
}()
sendEvent := func(eventType, message string, data interface{}) {
if h.taskEventBus == nil {
return
}
ev := StreamEvent{Type: eventType, Message: message, Data: data}
b, err := json.Marshal(ev)
if err != nil {
b = []byte(`{"type":"error","message":"marshal failed"}`)
}
line := make([]byte, 0, len(b)+8)
line = append(line, []byte("data: ")...)
line = append(line, b...)
line = append(line, '\n', '\n')
h.taskEventBus.Publish(conversationID, line)
}
if _, err := h.tasks.StartTask(conversationID, task.Message, cancelWithCause); err != nil {
h.logger.Warn("批量队列子任务注册会话运行状态失败",
zap.String("queueId", queueID),
zap.String("taskId", task.ID),
zap.String("conversationId", conversationID),
zap.Error(err))
failMsg := err.Error()
if errors.Is(err, ErrTaskAlreadyRunning) {
failMsg = "会话已有任务正在执行,无法在该会话上并行启动批量子任务"
}
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", failMsg)
return
}
registered = true
h.batchTaskManager.SetTaskCancel(queueID, task.ID, timeoutCancel)
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
taskCtx = mcp.WithMCPConversationID(taskCtx, conversationID)
taskCtx = mcp.WithToolRunRegistry(taskCtx, h.tasks)
taskCtx = mcp.WithEinoExecuteRunRegistry(taskCtx, h.tasks)
useBatchMulti := false
batchOrch := "deep"
am := strings.TrimSpace(strings.ToLower(queue.AgentMode))
if am == "multi" {
am = "deep"
}
if batchQueueWantsEino(queue.AgentMode) && h.config != nil && h.config.MultiAgent.Enabled {
useBatchMulti = true
batchOrch = config.NormalizeMultiAgentOrchestration(am)
} else if queue.AgentMode == "" && h.config != nil && h.config.MultiAgent.Enabled && h.config.MultiAgent.BatchUseMultiAgent {
useBatchMulti = true
batchOrch = "deep"
}
var resultMA *multiagent.RunResult
var runErr error
switch {
case useBatchMulti:
resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.agentSessionContextBlock(conversationID))
default:
if h.config == nil {
runErr = fmt.Errorf("服务器配置未加载")
} else {
resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.agentSessionContextBlock(conversationID))
}
}
if runErr != nil {
h.handleBatchSubTaskRunError(queueID, task, conversationID, assistantMessageID, baseCtx, taskCtx, resultMA, runErr, &finishStatus)
return
}
if resultMA == nil {
h.logger.Error("批量任务执行成功但无结果对象",
zap.String("queueId", queueID),
zap.String("taskId", task.ID),
zap.String("conversationId", conversationID))
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "内部错误:无执行结果")
return
}
h.logger.Info("批量任务执行成功", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
resText := resultMA.Response
mcpIDs := resultMA.MCPExecutionIDs
lastIn := resultMA.LastAgentTraceInput
lastOut := resultMA.LastAgentTraceOutput
if assistantMessageID != "" {
if updateErr := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resText, mcpIDs, multiagent.AggregatedReasoningFromTraceJSON(lastIn)); updateErr != nil {
h.logger.Warn("更新助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
if _, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs); err != nil {
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
}
} else if _, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs); err != nil {
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
}
if lastIn != "" || lastOut != "" {
if err := h.db.SaveAgentTrace(conversationID, lastIn, lastOut); err != nil {
h.logger.Warn("保存代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
}
}
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCompleted, resText, "", conversationID)
}
func (h *AgentHandler) handleBatchSubTaskRunError(
queueID string,
task *BatchTask,
conversationID, assistantMessageID string,
baseCtx, taskCtx context.Context,
resultMA *multiagent.RunResult,
runErr error,
finishStatus *string,
) {
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(conversationID, resultMA)
}
errStr := runErr.Error()
partialResp := ""
if resultMA != nil {
partialResp = resultMA.Response
}
isCancelled := errors.Is(context.Cause(baseCtx), ErrTaskCancelled) ||
errors.Is(runErr, context.Canceled) ||
strings.Contains(strings.ToLower(errStr), "context canceled") ||
strings.Contains(strings.ToLower(errStr), "context cancelled") ||
(partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断")))
isTimeout := errors.Is(runErr, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded)
if isTimeout {
*finishStatus = "timeout"
} else if isCancelled {
*finishStatus = "cancelled"
} else {
*finishStatus = "failed"
}
if isCancelled {
h.logger.Info("批量任务被取消", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
cancelMsg := "任务已被用户取消,后续操作已停止。"
if partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断")) {
cancelMsg = partialResp
}
if assistantMessageID != "" {
if updateErr := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); updateErr != nil {
h.logger.Warn("更新取消后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
}
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil); err != nil {
h.logger.Warn("保存取消详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
}
} else if _, errMsg := h.db.AddMessage(conversationID, "assistant", cancelMsg, nil); errMsg != nil {
h.logger.Warn("保存取消消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(errMsg))
}
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCancelled, cancelMsg, "", conversationID)
return
}
h.logger.Error("批量任务执行失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(runErr))
errorMsg := "执行失败: " + runErr.Error()
if assistantMessageID != "" {
if _, updateErr := h.db.Exec(
"UPDATE messages SET content = ?, updated_at = ? WHERE id = ?",
errorMsg,
time.Now(), assistantMessageID,
); updateErr != nil {
h.logger.Warn("更新失败后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
}
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errorMsg, nil); err != nil {
h.logger.Warn("保存错误详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
}
}
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", runErr.Error())
}
+368 -42
View File
@@ -4,6 +4,7 @@ import (
"context"
"crypto/rand"
"encoding/hex"
"errors"
"fmt"
"sort"
"strings"
@@ -17,6 +18,15 @@ import (
"go.uber.org/zap"
)
var (
// ErrBatchQueueNotFound 队列不存在或已从内存卸载。
ErrBatchQueueNotFound = errors.New("batch queue not found")
// ErrBatchQueueExecutorActive executeBatchQueue 协程仍在收尾,禁止删除。
ErrBatchQueueExecutorActive = errors.New("batch queue executor is still active")
// ErrBatchQueueStillRunning 队列状态仍为 running(无活跃执行器时的兜底保护)。
ErrBatchQueueStillRunning = errors.New("batch queue is still running")
)
// 批量任务状态常量
const (
BatchQueueStatusPending = "pending"
@@ -39,6 +49,12 @@ const (
// MaxBatchQueueRoleLen 角色名最大长度
MaxBatchQueueRoleLen = 100
// DefaultBatchQueueConcurrency 批量队列默认并发数(串行)
DefaultBatchQueueConcurrency = 1
// MaxBatchQueueConcurrency 批量队列最大并发数
MaxBatchQueueConcurrency = 8
)
// BatchTask 批量任务项
@@ -67,6 +83,7 @@ type BatchTaskQueue struct {
LastScheduleError string `json:"lastScheduleError,omitempty"`
LastRunError string `json:"lastRunError,omitempty"`
ProjectID string `json:"projectId,omitempty"`
Concurrency int `json:"concurrency"` // 同时执行的子任务数,默认 1
Tasks []*BatchTask `json:"tasks"`
Status string `json:"status"` // pending, running, paused, completed, cancelled
CreatedAt time.Time `json:"createdAt"`
@@ -77,11 +94,13 @@ type BatchTaskQueue struct {
// BatchTaskManager 批量任务管理器
type BatchTaskManager struct {
db *database.DB
logger *zap.Logger
queues map[string]*BatchTaskQueue
taskCancels map[string]context.CancelFunc // 存储每个队列当前任务的取消函数
mu sync.RWMutex
db *database.DB
logger *zap.Logger
queues map[string]*BatchTaskQueue
taskCancels map[string]map[string]context.CancelFunc // queueID -> taskID -> 取消函数
singleRunTasks map[string]string // queueID -> taskID,单条执行完成后暂停队列
queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记(与队列 status 解耦)
mu sync.RWMutex
}
// NewBatchTaskManager 创建批量任务管理器
@@ -90,12 +109,58 @@ func NewBatchTaskManager(logger *zap.Logger) *BatchTaskManager {
logger = zap.NewNop()
}
return &BatchTaskManager{
logger: logger,
queues: make(map[string]*BatchTaskQueue),
taskCancels: make(map[string]context.CancelFunc),
logger: logger,
queues: make(map[string]*BatchTaskQueue),
taskCancels: make(map[string]map[string]context.CancelFunc),
singleRunTasks: make(map[string]string),
queueExecutors: make(map[string]struct{}),
}
}
// batchQueueExecutionShouldStop 判断 executeBatchQueue 主循环是否应退出。
func batchQueueExecutionShouldStop(queue *BatchTaskQueue, exists bool) bool {
if !exists || queue == nil {
return true
}
switch queue.Status {
case BatchQueueStatusCancelled, BatchQueueStatusCompleted, BatchQueueStatusPaused:
return true
default:
return false
}
}
// TryMarkQueueExecutor 标记队列执行协程已启动;若已有执行协程则返回 false。
func (m *BatchTaskManager) TryMarkQueueExecutor(queueID string) bool {
m.mu.Lock()
defer m.mu.Unlock()
if _, exists := m.queueExecutors[queueID]; exists {
return false
}
m.queueExecutors[queueID] = struct{}{}
return true
}
// UnmarkQueueExecutor 清除队列执行协程标记(executeBatchQueue defer 调用)。
func (m *BatchTaskManager) UnmarkQueueExecutor(queueID string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.queueExecutors, queueID)
}
// ForceUnmarkQueueExecutor 强制清除执行协程标记(暂停态单条重跑等场景回收陈旧槽位)。
func (m *BatchTaskManager) ForceUnmarkQueueExecutor(queueID string) {
m.UnmarkQueueExecutor(queueID)
}
// IsQueueExecutorActive 队列 executeBatchQueue 协程是否仍在运行。
func (m *BatchTaskManager) IsQueueExecutorActive(queueID string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
_, ok := m.queueExecutors[queueID]
return ok
}
// SetDB 设置数据库连接
func (m *BatchTaskManager) SetDB(db *database.DB) {
m.mu.Lock()
@@ -103,10 +168,22 @@ func (m *BatchTaskManager) SetDB(db *database.DB) {
m.db = db
}
// normalizeBatchQueueConcurrency 规范化队列并发数。
func normalizeBatchQueueConcurrency(n int) int {
if n < 1 {
return DefaultBatchQueueConcurrency
}
if n > MaxBatchQueueConcurrency {
return MaxBatchQueueConcurrency
}
return n
}
// CreateBatchQueue 创建批量任务队列
func (m *BatchTaskManager) CreateBatchQueue(
title, role, agentMode, scheduleMode, cronExpr, projectID string,
nextRunAt *time.Time,
concurrency int,
tasks []string,
) (*BatchTaskQueue, error) {
// 输入校验
@@ -134,6 +211,7 @@ func (m *BatchTaskManager) CreateBatchQueue(
CronExpr: strings.TrimSpace(cronExpr),
NextRunAt: nextRunAt,
ScheduleEnabled: true,
Concurrency: normalizeBatchQueueConcurrency(concurrency),
Tasks: make([]*BatchTask, 0, len(tasks)),
Status: BatchQueueStatusPending,
CreatedAt: time.Now(),
@@ -175,6 +253,7 @@ func (m *BatchTaskManager) CreateBatchQueue(
queue.CronExpr,
queue.NextRunAt,
queue.ProjectID,
queue.Concurrency,
dbTasks,
); err != nil {
m.logger.Warn("batch queue DB create failed", zap.String("queueId", queueID), zap.Error(err))
@@ -270,6 +349,7 @@ func (m *BatchTaskManager) loadQueueFromDB(queueID string) *BatchTaskQueue {
if queueRow.ProjectID.Valid {
queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String)
}
queue.Concurrency = batchQueueConcurrencyFromRow(queueRow)
if queueRow.StartedAt.Valid {
queue.StartedAt = &queueRow.StartedAt.Time
}
@@ -509,6 +589,7 @@ func (m *BatchTaskManager) LoadFromDB() error {
if queueRow.ProjectID.Valid {
queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String)
}
queue.Concurrency = batchQueueConcurrencyFromRow(queueRow)
if queueRow.StartedAt.Valid {
queue.StartedAt = &queueRow.StartedAt.Time
}
@@ -649,8 +730,16 @@ func (m *BatchTaskManager) UpdateQueueSchedule(queueID, scheduleMode, cronExpr s
}
}
// UpdateQueueMetadata 更新队列标题、角色和代理模式(非 running 时可用)
func (m *BatchTaskManager) UpdateQueueMetadata(queueID, title, role, agentMode string) error {
// batchQueueConcurrencyFromRow 从数据库行读取并发数(缺省为 1)。
func batchQueueConcurrencyFromRow(row *database.BatchTaskQueueRow) int {
if row == nil || !row.Concurrency.Valid {
return DefaultBatchQueueConcurrency
}
return normalizeBatchQueueConcurrency(int(row.Concurrency.Int64))
}
// UpdateQueueMetadata 更新队列标题、角色、代理模式和并发数(非 running 时可用)
func (m *BatchTaskManager) UpdateQueueMetadata(queueID, title, role, agentMode string, concurrency *int) error {
if utf8.RuneCountInString(title) > MaxBatchQueueTitleLen {
return fmt.Errorf("标题不能超过 %d 个字符", MaxBatchQueueTitleLen)
}
@@ -678,9 +767,12 @@ func (m *BatchTaskManager) UpdateQueueMetadata(queueID, title, role, agentMode s
queue.Title = title
queue.Role = role
queue.AgentMode = agentMode
if concurrency != nil {
queue.Concurrency = normalizeBatchQueueConcurrency(*concurrency)
}
if m.db != nil {
if err := m.db.UpdateBatchQueueMetadata(queueID, title, role, agentMode); err != nil {
if err := m.db.UpdateBatchQueueMetadata(queueID, title, role, agentMode, queue.Concurrency); err != nil {
m.logger.Warn("batch queue DB metadata update failed", zap.String("queueId", queueID), zap.Error(err))
}
}
@@ -864,6 +956,137 @@ func (m *BatchTaskManager) AddTaskToQueue(queueID, message string) (*BatchTask,
return task, nil
}
// PrepareSingleTaskRun 准备单条执行:重置目标任务(若已有结果)并定位队列索引
func (m *BatchTaskManager) PrepareSingleTaskRun(queueID, taskID string) error {
var siblingRunningIDs []string
m.mu.Lock()
queue, exists := m.queues[queueID]
if !exists {
m.mu.Unlock()
return fmt.Errorf("队列不存在")
}
var task *BatchTask
taskIndex := -1
for i, t := range queue.Tasks {
if t.ID == taskID {
taskIndex = i
task = t
break
}
}
if task == nil {
m.mu.Unlock()
return fmt.Errorf("任务不存在")
}
if !queueAllowsSingleTaskRunLocked(queue, task) {
m.mu.Unlock()
return fmt.Errorf("队列正在执行或未就绪,无法单条执行")
}
// 暂停态:中止在途子任务并收口仍标记 running 的其它子任务,以便单条执行非冲突项
var cancelFuncs []context.CancelFunc
if queue.Status == BatchQueueStatusPaused {
cancelFuncs = m.drainTaskCancelsLocked(queueID)
for _, t := range queue.Tasks {
if t != nil && t.ID != taskID && t.Status == BatchTaskStatusRunning {
siblingRunningIDs = append(siblingRunningIDs, t.ID)
}
}
}
needsReset := task.Status != BatchTaskStatusPending
resumeQueue := queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusCancelled
m.mu.Unlock()
for _, c := range cancelFuncs {
if c != nil {
c()
}
}
const staleRunMsg = "为单条执行其它任务,已中止"
for _, sid := range siblingRunningIDs {
m.UpdateTaskStatus(queueID, sid, BatchTaskStatusCancelled, "", staleRunMsg)
}
m.mu.Lock()
defer m.mu.Unlock()
queue, exists = m.queues[queueID]
if !exists {
return fmt.Errorf("队列不存在")
}
task = nil
taskIndex = -1
for i, t := range queue.Tasks {
if t.ID == taskID {
taskIndex = i
task = t
break
}
}
if task == nil {
return fmt.Errorf("任务不存在")
}
if m.db != nil {
if err := m.db.PrepareBatchSingleTaskRun(queueID, taskID, taskIndex, needsReset, resumeQueue); err != nil {
return fmt.Errorf("准备单条执行失败: %w", err)
}
}
if needsReset {
task.Status = BatchTaskStatusPending
task.ConversationID = ""
task.StartedAt = nil
task.CompletedAt = nil
task.Error = ""
task.Result = ""
}
queue.CurrentIndex = taskIndex
queue.LastRunError = ""
if resumeQueue {
queue.Status = BatchQueueStatusPaused
queue.CompletedAt = nil
}
return nil
}
// SetSingleRunTask 标记队列仅执行指定子任务,完成后自动暂停
func (m *BatchTaskManager) SetSingleRunTask(queueID, taskID string) {
m.mu.Lock()
defer m.mu.Unlock()
if m.singleRunTasks == nil {
m.singleRunTasks = make(map[string]string)
}
m.singleRunTasks[queueID] = taskID
}
// ClearSingleRunTask 清除单条执行标记
func (m *BatchTaskManager) ClearSingleRunTask(queueID string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.singleRunTasks, queueID)
}
// TakeSingleRunTaskIfMatch 若刚完成的子任务为单条执行目标,则清除标记并返回 true
func (m *BatchTaskManager) TakeSingleRunTaskIfMatch(queueID, taskID string) bool {
m.mu.Lock()
defer m.mu.Unlock()
if m.singleRunTasks == nil {
return false
}
if m.singleRunTasks[queueID] != taskID {
return false
}
delete(m.singleRunTasks, queueID)
return true
}
// DeleteTask 删除任务(队列空闲时可删;执行中任务不可删)
func (m *BatchTaskManager) DeleteTask(queueID, taskID string) error {
m.mu.Lock()
@@ -936,7 +1159,109 @@ func queueAllowsTaskListMutationLocked(queue *BatchTaskQueue) bool {
}
}
// GetNextTask 获取下一个待执行的任务
// queueAllowsSingleTaskRunLocked 是否允许对指定子任务发起单条执行(必须在持有 BatchTaskManager.mu 下调用)
func queueAllowsSingleTaskRunLocked(queue *BatchTaskQueue, task *BatchTask) bool {
if queue == nil || task == nil {
return false
}
if task.Status == BatchTaskStatusRunning {
return false
}
if queue.Status == BatchQueueStatusRunning {
return false
}
switch queue.Status {
case BatchQueueStatusPending, BatchQueueStatusPaused, BatchQueueStatusCompleted, BatchQueueStatusCancelled:
return true
default:
return false
}
}
// ClaimNextPendingTask 原子领取下一个待执行子任务(并发 worker 安全)。
func (m *BatchTaskManager) ClaimNextPendingTask(queueID string) (*BatchTask, bool) {
m.mu.Lock()
defer m.mu.Unlock()
queue, exists := m.queues[queueID]
if !exists || queue == nil {
return nil, false
}
if queue.Status == BatchQueueStatusCancelled || queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusPaused {
return nil, false
}
onlyTaskID := ""
if m.singleRunTasks != nil {
onlyTaskID = m.singleRunTasks[queueID]
}
for i, task := range queue.Tasks {
if task == nil || task.Status != BatchTaskStatusPending {
continue
}
if onlyTaskID != "" && task.ID != onlyTaskID {
continue
}
task.Status = BatchTaskStatusRunning
queue.CurrentIndex = i
return task, true
}
return nil, false
}
// HasRunningTasks 队列是否仍有 running 状态的子任务。
func (m *BatchTaskManager) HasRunningTasks(queueID string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
queue, exists := m.queues[queueID]
if !exists || queue == nil {
return false
}
for _, task := range queue.Tasks {
if task != nil && task.Status == BatchTaskStatusRunning {
return true
}
}
return false
}
// HasPendingOrRunningTasks 队列是否仍有未完成的子任务。
func (m *BatchTaskManager) HasPendingOrRunningTasks(queueID string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
queue, exists := m.queues[queueID]
if !exists || queue == nil {
return false
}
for _, task := range queue.Tasks {
if task == nil {
continue
}
if task.Status == BatchTaskStatusPending || task.Status == BatchTaskStatusRunning {
return true
}
}
return false
}
// drainTaskCancelsLocked 取出并清空队列下所有子任务取消函数(调用方须已持 m.mu)。
func (m *BatchTaskManager) drainTaskCancelsLocked(queueID string) []context.CancelFunc {
taskMap, ok := m.taskCancels[queueID]
if !ok || len(taskMap) == 0 {
return nil
}
cancels := make([]context.CancelFunc, 0, len(taskMap))
for _, c := range taskMap {
if c != nil {
cancels = append(cancels, c)
}
}
delete(m.taskCancels, queueID)
return cancels
}
// GetNextTask 获取下一个待执行的任务(串行兼容,优先使用 ClaimNextPendingTask
func (m *BatchTaskManager) GetNextTask(queueID string) (*BatchTask, bool) {
m.mu.Lock()
defer m.mu.Unlock()
@@ -977,20 +1302,28 @@ func (m *BatchTaskManager) MoveToNextTask(queueID string) {
}
}
// SetTaskCancel 设置当前任务的取消函数
func (m *BatchTaskManager) SetTaskCancel(queueID string, cancel context.CancelFunc) {
// SetTaskCancel 设置任务的取消函数
func (m *BatchTaskManager) SetTaskCancel(queueID, taskID string, cancel context.CancelFunc) {
m.mu.Lock()
defer m.mu.Unlock()
if cancel != nil {
m.taskCancels[queueID] = cancel
} else {
delete(m.taskCancels, queueID)
if cancel == nil {
if taskMap, ok := m.taskCancels[queueID]; ok {
delete(taskMap, taskID)
if len(taskMap) == 0 {
delete(m.taskCancels, queueID)
}
}
return
}
if m.taskCancels[queueID] == nil {
m.taskCancels[queueID] = make(map[string]context.CancelFunc)
}
m.taskCancels[queueID][taskID] = cancel
}
// PauseQueue 暂停队列
func (m *BatchTaskManager) PauseQueue(queueID string) bool {
var cancelFunc context.CancelFunc
var cancelFuncs []context.CancelFunc
m.mu.Lock()
queue, exists := m.queues[queueID]
@@ -1015,17 +1348,11 @@ func (m *BatchTaskManager) PauseQueue(queueID string) bool {
}
queue.Status = BatchQueueStatusPaused
// 取消当前正在执行的任务(通过取消context)
if cancel, ok := m.taskCancels[queueID]; ok {
cancelFunc = cancel
delete(m.taskCancels, queueID)
}
cancelFuncs = m.drainTaskCancelsLocked(queueID)
m.mu.Unlock()
// 释放锁后执行取消回调(cancel 可能阻塞,不应持锁)
if cancelFunc != nil {
cancelFunc()
for _, c := range cancelFuncs {
c()
}
return true
@@ -1034,7 +1361,7 @@ func (m *BatchTaskManager) PauseQueue(queueID string) bool {
// CancelQueue 取消队列(保留此方法以保持向后兼容,但建议使用PauseQueue)
func (m *BatchTaskManager) CancelQueue(queueID string) bool {
now := time.Now()
var cancelFunc context.CancelFunc
var cancelFuncs []context.CancelFunc
m.mu.Lock()
queue, exists := m.queues[queueID]
@@ -1075,34 +1402,33 @@ func (m *BatchTaskManager) CancelQueue(queueID string) bool {
}
}
// 取消当前正在执行的任务
if cancel, ok := m.taskCancels[queueID]; ok {
cancelFunc = cancel
delete(m.taskCancels, queueID)
}
cancelFuncs = m.drainTaskCancelsLocked(queueID)
m.mu.Unlock()
// 释放锁后执行取消回调(cancel 可能阻塞,不应持锁)
if cancelFunc != nil {
cancelFunc()
for _, c := range cancelFuncs {
c()
}
return true
}
// DeleteQueue 删除队列(运行中的队列不允许删除)
func (m *BatchTaskManager) DeleteQueue(queueID string) bool {
// DeleteQueue 删除队列。执行协程活跃或 status 为 running 时拒绝删除,避免 executeBatchQueue 空指针 panic。
func (m *BatchTaskManager) DeleteQueue(queueID string) error {
m.mu.Lock()
defer m.mu.Unlock()
queue, exists := m.queues[queueID]
if !exists {
return false
return ErrBatchQueueNotFound
}
if _, exec := m.queueExecutors[queueID]; exec {
return ErrBatchQueueExecutorActive
}
// 运行中的队列不允许删除,防止孤儿协程和数据丢失
if queue.Status == BatchQueueStatusRunning {
return false
return ErrBatchQueueStillRunning
}
// 清理取消函数
@@ -1116,7 +1442,7 @@ func (m *BatchTaskManager) DeleteQueue(queueID string) bool {
}
delete(m.queues, queueID)
return true
return nil
}
// generateShortID 生成短ID
+121
View File
@@ -0,0 +1,121 @@
package handler
import (
"errors"
"testing"
"go.uber.org/zap"
)
func TestNormalizeBatchQueueConcurrency(t *testing.T) {
if got := normalizeBatchQueueConcurrency(0); got != DefaultBatchQueueConcurrency {
t.Fatalf("expected default %d, got %d", DefaultBatchQueueConcurrency, got)
}
if got := normalizeBatchQueueConcurrency(99); got != MaxBatchQueueConcurrency {
t.Fatalf("expected max %d, got %d", MaxBatchQueueConcurrency, got)
}
}
func TestClaimNextPendingTaskParallel(t *testing.T) {
m := NewBatchTaskManager(zap.NewNop())
queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 3, []string{"a", "b", "c"})
if err != nil {
t.Fatalf("CreateBatchQueue: %v", err)
}
m.UpdateQueueStatus(queue.ID, BatchQueueStatusRunning)
t1, ok1 := m.ClaimNextPendingTask(queue.ID)
t2, ok2 := m.ClaimNextPendingTask(queue.ID)
if !ok1 || !ok2 || t1.ID == t2.ID {
t.Fatalf("expected two distinct claims, got ok1=%v ok2=%v t1=%v t2=%v", ok1, ok2, t1, t2)
}
if t1.Status != BatchTaskStatusRunning || t2.Status != BatchTaskStatusRunning {
t.Fatalf("claimed tasks should be running")
}
t3, ok3 := m.ClaimNextPendingTask(queue.ID)
if !ok3 {
t.Fatal("expected third claim")
}
_, ok4 := m.ClaimNextPendingTask(queue.ID)
if ok4 {
t.Fatal("expected no fourth pending task")
}
_ = t3
}
func TestBatchQueueExecutionShouldStop(t *testing.T) {
t.Parallel()
if !batchQueueExecutionShouldStop(nil, false) {
t.Fatal("expected stop when queue missing")
}
if !batchQueueExecutionShouldStop(nil, true) {
t.Fatal("expected stop when queue is nil but exists=true")
}
q := &BatchTaskQueue{Status: BatchQueueStatusRunning}
if batchQueueExecutionShouldStop(q, true) {
t.Fatal("expected continue when running")
}
q.Status = BatchQueueStatusCancelled
if !batchQueueExecutionShouldStop(q, true) {
t.Fatal("expected stop when cancelled")
}
}
func TestDeleteQueueBlockedWhileExecutorActive(t *testing.T) {
t.Parallel()
m := NewBatchTaskManager(zap.NewNop())
queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 1, []string{"hello"})
if err != nil {
t.Fatalf("CreateBatchQueue: %v", err)
}
if !m.TryMarkQueueExecutor(queue.ID) {
t.Fatal("expected to mark executor")
}
m.UpdateQueueStatus(queue.ID, BatchQueueStatusCancelled)
err = m.DeleteQueue(queue.ID)
if !errors.Is(err, ErrBatchQueueExecutorActive) {
t.Fatalf("expected ErrBatchQueueExecutorActive, got %v", err)
}
if _, ok := m.GetBatchQueue(queue.ID); !ok {
t.Fatal("queue should still exist while executor active")
}
m.UnmarkQueueExecutor(queue.ID)
if err := m.DeleteQueue(queue.ID); err != nil {
t.Fatalf("expected delete after executor unmarked, got %v", err)
}
if _, ok := m.GetBatchQueue(queue.ID); ok {
t.Fatal("queue should be deleted")
}
}
func TestDeleteQueueBlockedWhileRunning(t *testing.T) {
t.Parallel()
m := NewBatchTaskManager(zap.NewNop())
queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 1, []string{"hello"})
if err != nil {
t.Fatalf("CreateBatchQueue: %v", err)
}
m.UpdateQueueStatus(queue.ID, BatchQueueStatusRunning)
err = m.DeleteQueue(queue.ID)
if !errors.Is(err, ErrBatchQueueStillRunning) {
t.Fatalf("expected ErrBatchQueueStillRunning, got %v", err)
}
}
func TestTryMarkQueueExecutorDedupes(t *testing.T) {
t.Parallel()
m := NewBatchTaskManager(zap.NewNop())
if !m.TryMarkQueueExecutor("q-1") {
t.Fatal("first mark should succeed")
}
if m.TryMarkQueueExecutor("q-1") {
t.Fatal("second mark should fail")
}
m.UnmarkQueueExecutor("q-1")
if !m.TryMarkQueueExecutor("q-1") {
t.Fatal("mark after unmark should succeed")
}
}
+30 -4
View File
@@ -3,6 +3,7 @@ package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
@@ -181,6 +182,10 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
"type": "string",
"description": "队列内子对话绑定的项目 ID(可选,未指定时使用 config.project.default_project_id",
},
"concurrency": map[string]interface{}{
"type": "integer",
"description": "同时执行的子任务数,默认 1(串行),最大 8。含扫描类工具时建议 1-2。",
},
},
},
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
@@ -210,7 +215,8 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
executeNow = false
}
projectID := strings.TrimSpace(mcpArgString(args, "project_id"))
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, tasks)
concurrency := int(mcpArgFloat(args, "concurrency"))
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, concurrency, tasks)
if createErr != nil {
return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil
}
@@ -365,8 +371,17 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
if qid == "" {
return batchMCPTextResult("queue_id 不能为空", true), nil
}
if !h.batchTaskManager.DeleteQueue(qid) {
return batchMCPTextResult("删除失败:队列不存在", true), nil
if err := h.batchTaskManager.DeleteQueue(qid); err != nil {
switch {
case errors.Is(err, ErrBatchQueueNotFound):
return batchMCPTextResult("删除失败:队列不存在", true), nil
case errors.Is(err, ErrBatchQueueExecutorActive):
return batchMCPTextResult("删除失败:队列执行器仍在运行,请稍后再试", true), nil
case errors.Is(err, ErrBatchQueueStillRunning):
return batchMCPTextResult("删除失败:队列正在运行中", true), nil
default:
return batchMCPTextResult("删除失败:"+err.Error(), true), nil
}
}
logger.Info("MCP batch_task_delete", zap.String("queueId", qid))
return batchMCPTextResult("队列已删除。", false), nil
@@ -397,6 +412,10 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
"description": "代理模式:eino_single、deep、plan_execute、supervisor",
"enum": []string{"eino_single", "deep", "plan_execute", "supervisor"},
},
"concurrency": map[string]interface{}{
"type": "integer",
"description": "同时执行的子任务数,默认 1,最大 8",
},
},
"required": []string{"queue_id"},
},
@@ -408,7 +427,12 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
title := mcpArgString(args, "title")
role := mcpArgString(args, "role")
agentMode := mcpArgString(args, "agent_mode")
if err := h.batchTaskManager.UpdateQueueMetadata(qid, title, role, agentMode); err != nil {
var concurrency *int
if raw, ok := args["concurrency"]; ok && raw != nil {
v := int(mcpArgFloat(args, "concurrency"))
concurrency = &v
}
if err := h.batchTaskManager.UpdateQueueMetadata(qid, title, role, agentMode, concurrency); err != nil {
return batchMCPTextResult(err.Error(), true), nil
}
updated, _ := h.batchTaskManager.GetBatchQueue(qid)
@@ -652,6 +676,7 @@ type batchTaskQueueMCPListItem struct {
StartedAt *time.Time `json:"startedAt,omitempty"`
CompletedAt *time.Time `json:"completedAt,omitempty"`
CurrentIndex int `json:"currentIndex"`
Concurrency int `json:"concurrency"`
TaskTotal int `json:"task_total"`
TaskCounts map[string]int `json:"task_counts"`
Tasks []batchTaskMCPListSummary `json:"tasks"`
@@ -715,6 +740,7 @@ func toBatchTaskQueueMCPListItem(q *BatchTaskQueue) batchTaskQueueMCPListItem {
StartedAt: q.StartedAt,
CompletedAt: q.CompletedAt,
CurrentIndex: q.CurrentIndex,
Concurrency: q.Concurrency,
TaskTotal: len(tasks),
TaskCounts: counts,
Tasks: tasks,
+75 -9
View File
@@ -12,11 +12,17 @@ import (
"go.uber.org/zap"
)
// ConversationTaskStopper cancels in-flight agent work when a conversation is removed.
type ConversationTaskStopper interface {
CancelRunningTaskForConversation(conversationID string)
}
// ConversationHandler 对话处理器
type ConversationHandler struct {
db *database.DB
logger *zap.Logger
audit *audit.Service
db *database.DB
logger *zap.Logger
audit *audit.Service
taskStopper ConversationTaskStopper
}
// SetAudit wires platform audit logging.
@@ -24,6 +30,11 @@ func (h *ConversationHandler) SetAudit(s *audit.Service) {
h.audit = s
}
// SetTaskStopper wires cancellation of in-flight agent tasks on conversation delete.
func (h *ConversationHandler) SetTaskStopper(stopper ConversationTaskStopper) {
h.taskStopper = stopper
}
// NewConversationHandler 创建新的对话处理器
func NewConversationHandler(db *database.DB, logger *zap.Logger) *ConversationHandler {
return &ConversationHandler{
@@ -105,17 +116,18 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
excludeGrouped := strings.TrimSpace(search) == "" &&
(c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1")
sortBy := strings.TrimSpace(c.Query("sort_by"))
var conversations []*database.Conversation
var total int
var err error
if excludeGrouped {
conversations, err = h.db.ListUngroupedConversations(limit, offset)
conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy)
if err == nil {
total, err = h.db.CountUngroupedConversations()
}
} else {
conversations, err = h.db.ListConversations(limit, offset, search)
conversations, err = h.db.ListConversations(limit, offset, search, sortBy)
if err == nil {
total, err = h.db.CountConversations(search)
}
@@ -164,6 +176,9 @@ func (h *ConversationHandler) GetConversation(c *gin.Context) {
}
// GetMessageProcessDetails 获取指定消息的过程详情(按需加载)
// 查询参数:
// - summary=1:仅返回摘要(total / iterationCount / maxIteration
// - limit + offset:分页返回 processDetails(未指定 limit 时保持全量兼容)
func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
messageID := c.Param("id")
if messageID == "" {
@@ -171,6 +186,51 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
return
}
summaryStr := strings.TrimSpace(c.Query("summary"))
if summaryStr == "1" || strings.EqualFold(summaryStr, "true") || strings.EqualFold(summaryStr, "yes") {
summary, err := h.db.GetProcessDetailsSummary(messageID)
if err != nil {
h.logger.Error("获取过程详情摘要失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"summary": summary})
return
}
limitStr := strings.TrimSpace(c.Query("limit"))
if limitStr != "" {
limit, err := strconv.Atoi(limitStr)
if err != nil || limit <= 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"})
return
}
if limit > 500 {
limit = 500
}
offset, _ := strconv.Atoi(strings.TrimSpace(c.Query("offset")))
if offset < 0 {
offset = 0
}
details, total, err := h.db.GetProcessDetailsPage(messageID, limit, offset)
if err != nil {
h.logger.Error("分页获取过程详情失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
details = database.DedupeConsecutiveProcessDetails(details)
out := processDetailsToJSON(h.logger, details)
c.JSON(http.StatusOK, gin.H{
"processDetails": out,
"total": total,
"offset": offset,
"limit": limit,
"hasMore": offset+len(out) < total,
})
return
}
details, err := h.db.GetProcessDetails(messageID)
if err != nil {
h.logger.Error("获取过程详情失败", zap.Error(err))
@@ -179,14 +239,17 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
}
details = database.DedupeConsecutiveProcessDetails(details)
out := processDetailsToJSON(h.logger, details)
c.JSON(http.StatusOK, gin.H{"processDetails": out, "total": len(out)})
}
// 转换为前端期望的 JSON 结构(与 GetConversation 中 processDetails 结构一致)
func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail) []map[string]interface{} {
out := make([]map[string]interface{}, 0, len(details))
for _, d := range details {
var data interface{}
if d.Data != "" {
if err := json.Unmarshal([]byte(d.Data), &data); err != nil {
h.logger.Warn("解析过程详情数据失败", zap.Error(err))
logger.Warn("解析过程详情数据失败", zap.Error(err))
}
}
out = append(out, map[string]interface{}{
@@ -199,8 +262,7 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
"createdAt": d.CreatedAt,
})
}
c.JSON(http.StatusOK, gin.H{"processDetails": out})
return out
}
// UpdateConversationRequest 更新对话请求
@@ -244,6 +306,10 @@ func (h *ConversationHandler) UpdateConversation(c *gin.Context) {
func (h *ConversationHandler) DeleteConversation(c *gin.Context) {
id := c.Param("id")
if h.taskStopper != nil {
h.taskStopper.CancelRunningTaskForConversation(id)
}
if err := h.db.DeleteConversation(id); err != nil {
h.logger.Error("删除对话失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
@@ -0,0 +1,30 @@
package handler
import (
"context"
"testing"
"time"
"go.uber.org/zap"
)
func TestConversationHandlerDeleteConversationCancelsRunningTask(t *testing.T) {
tm := NewAgentTaskManager()
ctx, cancel := context.WithCancelCause(context.Background())
_, err := tm.StartTask("conv-1", "hello", cancel)
if err != nil {
t.Fatalf("StartTask: %v", err)
}
h := &AgentHandler{tasks: tm, logger: zap.NewNop()}
h.CancelRunningTaskForConversation("conv-1")
select {
case <-ctx.Done():
case <-time.After(2 * time.Second):
t.Fatal("task context was not cancelled")
}
if cause := context.Cause(ctx); cause != ErrTaskCancelled {
t.Fatalf("expected ErrTaskCancelled, got %v", cause)
}
}
-153
View File
@@ -2,31 +2,11 @@ package handler
import (
"context"
"errors"
"fmt"
"strings"
"time"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/multiagent"
"go.uber.org/zap"
)
func (h *AgentHandler) einoRunRetryMaxAttempts() int {
if h.config != nil {
return multiagent.RunRetryMaxAttemptsFromConfig(&h.config.MultiAgent.EinoMiddleware)
}
return multiagent.RunRetryMaxAttemptsFromConfig(nil)
}
func (h *AgentHandler) einoRunRetryMaxBackoffSec() int {
if h.config != nil && h.config.MultiAgent.EinoMiddleware.RunRetryMaxBackoffSec > 0 {
return h.config.MultiAgent.EinoMiddleware.RunRetryMaxBackoffSec
}
return 0
}
// applyEinoTraceResumeSegment 中断并继续:persist last_react_* → loadHistory,可选替换下一段 user 文案。
func (h *AgentHandler) applyEinoTraceResumeSegment(
conversationID string,
@@ -45,136 +25,3 @@ func (h *AgentHandler) applyEinoTraceResumeSegment(
*curFinalMessage = segmentUserMessage
}
}
// applyEinoTransientRetrySegment 临时错误重试:恢复轨迹并保留本请求原始 user 文案(不注入续跑说明)。
// segmentUserMessage 为本轮 HTTP 请求开始时用户发送的内容,避免因清空 finalMessage 而丢失「你好」等短句。
func (h *AgentHandler) applyEinoTransientRetrySegment(
conversationID string,
result *multiagent.RunResult,
curHistory *[]agent.ChatMessage,
curFinalMessage *string,
segmentUserMessage string,
) {
if shouldPersistEinoAgentTraceAfterRunError(context.Background()) {
h.persistEinoAgentTraceForResume(conversationID, result)
}
if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 {
*curHistory = hist
}
if s := strings.TrimSpace(segmentUserMessage); s != "" {
*curFinalMessage = segmentUserMessage
}
}
// handleEinoTransientRetryContinue 在 SSE 任务循环内处理临时错误重试;返回 true 表示外层 for 应 continue。
func (h *AgentHandler) handleEinoTransientRetryContinue(
baseCtx context.Context,
conversationID string,
result *multiagent.RunResult,
runErr error,
transientAttempts *int,
curHistory *[]agent.ChatMessage,
curFinalMessage *string,
segmentUserMessage string,
progressCallback func(eventType, message string, data interface{}),
sendProgress func(msg string, extra map[string]interface{}),
) (handled bool, fatal error) {
if !errors.Is(runErr, multiagent.ErrTransientRetryContinue) {
return false, nil
}
maxAttempts := h.einoRunRetryMaxAttempts()
*transientAttempts++
if *transientAttempts > maxAttempts {
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(conversationID, result)
}
return false, errors.New("transient retry exhausted: " + runErr.Error())
}
attemptNo := *transientAttempts
backoff := multiagent.TransientRetryBackoff(attemptNo-1, h.einoRunRetryMaxBackoffSec())
if progressCallback != nil {
progressCallback("eino_run_retry", fmt.Sprintf("遇到临时错误,%d 秒后第 %d/%d 次重试…", int(backoff.Seconds()), attemptNo, maxAttempts), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"backoffSec": int(backoff.Seconds()),
})
}
select {
case <-baseCtx.Done():
return false, context.Cause(baseCtx)
case <-time.After(backoff):
}
h.applyEinoTransientRetrySegment(conversationID, result, curHistory, curFinalMessage, segmentUserMessage)
if progressCallback != nil {
progressCallback("eino_run_retry", "已恢复上下文,正在重试…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": attemptNo,
})
}
if sendProgress != nil {
sendProgress("正在重试…", map[string]interface{}{
"conversationId": conversationID,
"source": "transient_retry",
})
}
return true, nil
}
// handleEinoEmptyResponseContinue 在 SSE 任务循环内处理「正常结束但无助手正文」;返回 exhausted=true 时由外层按成功结束(保留占位文案)。
// 与临时错误重试一致:仅恢复轨迹并保留本请求原始 user 文案,不向模型注入续跑说明。
func (h *AgentHandler) handleEinoEmptyResponseContinue(
baseCtx context.Context,
conversationID string,
result *multiagent.RunResult,
runErr error,
emptyResponseAttempts *int,
curHistory *[]agent.ChatMessage,
curFinalMessage *string,
segmentUserMessage string,
progressCallback func(eventType, message string, data interface{}),
sendProgress func(msg string, extra map[string]interface{}),
) (handled bool, exhausted bool) {
if !errors.Is(runErr, multiagent.ErrEmptyResponseContinue) {
return false, false
}
maxAttempts := h.einoRunRetryMaxAttempts()
*emptyResponseAttempts++
if *emptyResponseAttempts > maxAttempts {
if h.logger != nil {
h.logger.Warn("eino empty response auto resume exhausted",
zap.String("conversationId", conversationID),
zap.Int("maxAttempts", maxAttempts))
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(conversationID, result)
}
return false, true
}
attemptNo := *emptyResponseAttempts
if h.logger != nil {
h.logger.Info("eino empty response, auto resume from trace",
zap.String("conversationId", conversationID),
zap.Int("attempt", attemptNo),
zap.Int("maxAttempts", maxAttempts))
}
if progressCallback != nil {
progressCallback("eino_empty_response_continue", fmt.Sprintf("未捕获到助手正文,正在基于轨迹自动续跑(%d/%d)…", attemptNo, maxAttempts), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"resumeKind": "trace_segment",
})
}
h.applyEinoTransientRetrySegment(conversationID, result, curHistory, curFinalMessage, segmentUserMessage)
if sendProgress != nil {
sendProgress("已恢复上下文,正在继续推理…", map[string]interface{}{
"conversationId": conversationID,
"source": "empty_response_continue",
})
}
return true, false
}
+5 -71
View File
@@ -119,7 +119,6 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
var cancelWithCause context.CancelCauseFunc
curFinalMessage := prep.FinalMessage
segmentUserMessage := prep.FinalMessage // 本请求原始用户句,临时重试时不得丢失
curHistory := prep.History
roleTools := prep.RoleTools
@@ -177,8 +176,6 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
taskOwned = true
var cumulativeMCPExecutionIDs []string
var transientRunAttempts int
var emptyResponseAttempts int
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int
@@ -215,6 +212,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
}
taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID)
taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks)
taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks)
taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) {
return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments)
})
@@ -224,6 +222,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
h.config,
&h.config.MultiAgent,
h.agent,
h.db,
h.logger,
conversationID,
h.conversationProjectID(conversationID),
@@ -232,61 +231,18 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
roleTools,
progressCallback,
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(conversationID),
h.agentSessionContextBlock(conversationID),
)
if result != nil && len(result.MCPExecutionIDs) > 0 {
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
}
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, conversationID, result, runErr, &emptyResponseAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if exhaustedEmpty {
runErr = nil
transientRunAttempts = 0
timeoutCancel()
break
}
if handledEmpty {
mainIterationOffset += segmentMainIterationMax
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if runErr == nil {
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
transientRunAttempts = 0
emptyResponseAttempts = 0
timeoutCancel()
break
}
handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, conversationID, result, runErr, &transientRunAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if handled {
mainIterationOffset += segmentMainIterationMax
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if fatalErr != nil {
runErr = fatalErr
}
cause := context.Cause(baseCtx)
if errors.Is(cause, multiagent.ErrInterruptContinue) {
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
@@ -311,8 +267,6 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
"source": "interrupt_continue",
})
mainIterationOffset += segmentMainIterationMax
// 非临时错误分段续跑(用户中断并继续)时,清空 transient 计数,避免跨分段累加。
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
@@ -447,14 +401,13 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
curMsg := prep.FinalMessage
var result *multiagent.RunResult
var runErr error
var transientRunAttempts int
var emptyResponseAttempts int
for {
result, runErr = multiagent.RunEinoSingleChatModelAgent(
taskCtx,
h.config,
&h.config.MultiAgent,
h.agent,
h.db,
h.logger,
prep.ConversationID,
h.conversationProjectID(prep.ConversationID),
@@ -463,30 +416,11 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
prep.RoleTools,
progressCallback,
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(prep.ConversationID),
h.agentSessionContextBlock(prep.ConversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, prep.ConversationID, result, runErr, &emptyResponseAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
)
if exhaustedEmpty {
runErr = nil
break
}
if handledEmpty {
continue
}
if runErr == nil {
break
}
if handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, prep.ConversationID, result, runErr, &transientRunAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
); handled {
continue
} else if fatalErr != nil {
runErr = fatalErr
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
}
+112 -20
View File
@@ -10,8 +10,10 @@ import (
"time"
"cyberstrike-ai/internal/audit"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/monitor"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
@@ -19,12 +21,20 @@ import (
// MonitorHandler 监控处理器
type MonitorHandler struct {
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager
executor *security.Executor
db *database.DB
logger *zap.Logger
audit *audit.Service
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager
taskManager *AgentTaskManager
agentHandler *AgentHandler
executor *security.Executor
db *database.DB
logger *zap.Logger
audit *audit.Service
monitorRetention *monitor.Service
}
// SetMonitorRetention wires MCP execution retention settings.
func (h *MonitorHandler) SetMonitorRetention(s *monitor.Service) {
h.monitorRetention = s
}
// SetAudit wires platform audit logging.
@@ -48,15 +58,26 @@ func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) {
h.externalMCPMgr = mgr
}
// SetTaskManager 设置 Agent 任务管理器(用于 Eino execute 等按 executionId 终止)。
func (h *MonitorHandler) SetTaskManager(mgr *AgentTaskManager) {
h.taskManager = mgr
}
// SetAgentHandler 设置 Agent 处理器(MCP 监控终止与对话页「中断并继续」共用逻辑)。
func (h *MonitorHandler) SetAgentHandler(ah *AgentHandler) {
h.agentHandler = ah
}
// MonitorResponse 监控响应
type MonitorResponse struct {
Executions []*mcp.ToolExecution `json:"executions"`
Stats map[string]*mcp.ToolStats `json:"stats"`
Timestamp time.Time `json:"timestamp"`
Total int `json:"total,omitempty"`
Page int `json:"page,omitempty"`
PageSize int `json:"page_size,omitempty"`
TotalPages int `json:"total_pages,omitempty"`
Executions []*mcp.ToolExecution `json:"executions"`
Stats map[string]*mcp.ToolStats `json:"stats"`
Timestamp time.Time `json:"timestamp"`
Total int `json:"total,omitempty"`
Page int `json:"page,omitempty"`
PageSize int `json:"page_size,omitempty"`
TotalPages int `json:"total_pages,omitempty"`
RetentionDays int `json:"retention_days,omitempty"`
}
// Monitor 获取监控信息
@@ -81,6 +102,7 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
toolName := normalizeToolNameFilter(c.Query("tool"))
executions, total := h.loadExecutionsWithPagination(page, pageSize, status, toolName)
h.enrichExecutionsConversationID(executions)
stats := h.loadStats()
totalPages := (total + pageSize - 1) / pageSize
@@ -89,16 +111,24 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
}
c.JSON(http.StatusOK, MonitorResponse{
Executions: executions,
Stats: stats,
Timestamp: time.Now(),
Total: total,
Page: page,
PageSize: pageSize,
TotalPages: totalPages,
Executions: executions,
Stats: stats,
Timestamp: time.Now(),
Total: total,
Page: page,
PageSize: pageSize,
TotalPages: totalPages,
RetentionDays: h.monitorRetentionDays(),
})
}
func (h *MonitorHandler) monitorRetentionDays() int {
if h.monitorRetention != nil {
return h.monitorRetention.RetentionDays()
}
return config.MonitorConfig{}.RetentionDaysEffective()
}
func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
executions, _ := h.loadExecutionsWithPagination(1, 1000, "", "")
return executions
@@ -230,6 +260,7 @@ func (h *MonitorHandler) GetExecution(c *gin.Context) {
// 先从内部MCP服务器查找
exec, exists := h.mcpServer.GetExecution(id)
if exists {
h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec})
c.JSON(http.StatusOK, exec)
return
}
@@ -238,6 +269,7 @@ func (h *MonitorHandler) GetExecution(c *gin.Context) {
if h.externalMCPMgr != nil {
exec, exists = h.externalMCPMgr.GetExecution(id)
if exists {
h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec})
c.JSON(http.StatusOK, exec)
return
}
@@ -247,6 +279,7 @@ func (h *MonitorHandler) GetExecution(c *gin.Context) {
if h.db != nil {
exec, err := h.db.GetToolExecution(id)
if err == nil && exec != nil {
h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec})
c.JSON(http.StatusOK, exec)
return
}
@@ -273,6 +306,19 @@ func (h *MonitorHandler) CancelExecution(c *gin.Context) {
return
}
note = strings.TrimSpace(body.Note)
convID := h.conversationIDForRunningExecution(id)
if convID != "" && h.agentHandler != nil {
if ok, payload := h.agentHandler.cancelToolContinueAfter(convID, id, note); ok {
h.logger.Info("MCP 监控页终止工具(与对话中断并继续一致)",
zap.String("executionId", id),
zap.String("conversationId", convID),
zap.Bool("hasNote", note != ""),
)
c.JSON(http.StatusOK, payload)
return
}
}
if h.mcpServer.CancelToolExecutionWithNote(id, note) {
h.logger.Info("已请求取消 MCP 工具执行", zap.String("executionId", id), zap.String("source", "internal"), zap.Bool("hasNote", note != ""))
c.JSON(http.StatusOK, gin.H{"message": "已发送终止信号", "executionId": id})
@@ -286,6 +332,52 @@ func (h *MonitorHandler) CancelExecution(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": "未找到进行中的工具执行,或该任务已结束"})
}
func (h *MonitorHandler) enrichExecutionsConversationID(executions []*mcp.ToolExecution) {
for _, exec := range executions {
if exec == nil {
continue
}
exec.ConversationID = h.conversationIDForRunningExecution(exec.ID)
}
}
func (h *MonitorHandler) conversationIDForRunningExecution(executionID string) string {
executionID = strings.TrimSpace(executionID)
if executionID == "" || h.taskManager == nil {
return ""
}
if conv := h.taskManager.ConversationIDForActiveMCPExecution(executionID); conv != "" {
return conv
}
exec := h.lookupExecution(executionID)
if exec == nil || exec.Status != "running" {
return ""
}
if strings.TrimSpace(exec.ToolName) == "execute" {
if onlyConv, ok := h.taskManager.ConversationIDForActiveEinoExecute(); ok {
return onlyConv
}
}
return ""
}
func (h *MonitorHandler) lookupExecution(id string) *mcp.ToolExecution {
if exec, ok := h.mcpServer.GetExecution(id); ok {
return exec
}
if h.externalMCPMgr != nil {
if exec, ok := h.externalMCPMgr.GetExecution(id); ok {
return exec
}
}
if h.db != nil {
if exec, err := h.db.GetToolExecution(id); err == nil && exec != nil {
return exec
}
}
return nil
}
// BatchGetToolNames 批量获取工具执行的工具名称(消除前端 N+1 请求)
func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) {
var req struct {
+5 -71
View File
@@ -136,7 +136,6 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
var cancelWithCause context.CancelCauseFunc
curFinalMessage := prep.FinalMessage
segmentUserMessage := prep.FinalMessage // 本请求原始用户句,临时重试时不得丢失
curHistory := prep.History
roleTools := prep.RoleTools
orch := strings.TrimSpace(req.Orchestration)
@@ -187,8 +186,6 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
// 同一 HTTP 流内多段 Run(如中断并继续)合并 MCP execution id,供最终 response / 库表与工具芯片展示完整列表
var cumulativeMCPExecutionIDs []string
var transientRunAttempts int
var emptyResponseAttempts int
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
var mainIterationOffset int
@@ -225,6 +222,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
}
taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID)
taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks)
taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks)
taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) {
return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments)
})
@@ -234,6 +232,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
h.config,
&h.config.MultiAgent,
h.agent,
h.db,
h.logger,
conversationID,
h.conversationProjectID(conversationID),
@@ -244,61 +243,18 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
h.agentsMarkdownDir,
orch,
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(conversationID),
h.agentSessionContextBlock(conversationID),
)
if result != nil && len(result.MCPExecutionIDs) > 0 {
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
}
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, conversationID, result, runErr, &emptyResponseAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if exhaustedEmpty {
runErr = nil
transientRunAttempts = 0
timeoutCancel()
break
}
if handledEmpty {
mainIterationOffset += segmentMainIterationMax
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if runErr == nil {
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
transientRunAttempts = 0
emptyResponseAttempts = 0
timeoutCancel()
break
}
handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, conversationID, result, runErr, &transientRunAttempts,
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
)
if handled {
mainIterationOffset += segmentMainIterationMax
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
h.tasks.UpdateTaskStatus(conversationID, "running")
continue
}
if fatalErr != nil {
runErr = fatalErr
}
cause := context.Cause(baseCtx)
if errors.Is(cause, multiagent.ErrInterruptContinue) {
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
@@ -323,8 +279,6 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
"source": "interrupt_continue",
})
mainIterationOffset += segmentMainIterationMax
// 非临时错误分段续跑(用户中断并继续)时,清空 transient 计数,避免跨分段累加。
transientRunAttempts = 0
timeoutCancel()
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
@@ -459,14 +413,13 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
curMsg := prep.FinalMessage
var result *multiagent.RunResult
var runErr error
var transientRunAttempts int
var emptyResponseAttempts int
for {
result, runErr = multiagent.RunDeepAgent(
taskCtx,
h.config,
&h.config.MultiAgent,
h.agent,
h.db,
h.logger,
prep.ConversationID,
h.conversationProjectID(prep.ConversationID),
@@ -477,30 +430,11 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
h.agentsMarkdownDir,
strings.TrimSpace(req.Orchestration),
chatReasoningToClientIntent(req.Reasoning),
h.projectBlackboardBlock(prep.ConversationID),
h.agentSessionContextBlock(prep.ConversationID),
)
handledEmpty, exhaustedEmpty := h.handleEinoEmptyResponseContinue(
baseCtx, prep.ConversationID, result, runErr, &emptyResponseAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
)
if exhaustedEmpty {
runErr = nil
break
}
if handledEmpty {
continue
}
if runErr == nil {
break
}
if handled, fatalErr := h.handleEinoTransientRetryContinue(
baseCtx, prep.ConversationID, result, runErr, &transientRunAttempts,
&curHist, &curMsg, prep.FinalMessage, progressCallback, nil,
); handled {
continue
} else if fatalErr != nil {
runErr = fatalErr
}
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
}
+92 -1
View File
@@ -2464,17 +2464,108 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "fact_key", "in": "query", "schema": map[string]interface{}{"type": "string"}},
{"name": "include_links", "in": "query", "schema": map[string]interface{}{"type": "boolean"}},
{"name": "include_link_counts", "in": "query", "schema": map[string]interface{}{"type": "boolean"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条"}},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条(可含 link_counts / outgoing_links"}},
},
"post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "创建/更新事实", "operationId": "upsertProjectFactREST",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"fact_key": map[string]interface{}{"type": "string"},
"summary": map[string]interface{}{"type": "string"},
"links": map[string]interface{}{
"type": "array",
"items": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"to": map[string]interface{}{"type": "string"},
"type": map[string]interface{}{"type": "string"},
},
},
},
"links_text": map[string]interface{}{"type": "string", "description": "type: fact_key 每行一条"},
},
},
},
},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "成功"}},
},
},
"/api/projects/{id}/fact-graph": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "获取项目事实攻击路径图", "operationId": "getProjectFactGraph",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "view", "in": "query", "schema": map[string]interface{}{"type": "string", "enum": []string{"path", "full"}, "default": "path"}},
{"name": "exclude_deprecated", "in": "query", "schema": map[string]interface{}{"type": "boolean", "default": true}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "nodes + edges"}},
},
},
"/api/projects/{id}/fact-edges": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "列出项目全部事实边", "operationId": "listProjectFactEdges",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "边列表"}},
},
"post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "添加事实边", "operationId": "createProjectFactEdge",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{
"type": "object",
"required": []string{"source_fact_key", "target_fact_key", "edge_type"},
"properties": map[string]interface{}{
"source_fact_key": map[string]interface{}{"type": "string"},
"target_fact_key": map[string]interface{}{"type": "string"},
"edge_type": map[string]interface{}{"type": "string"},
"confidence": map[string]interface{}{"type": "string"},
},
},
},
},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "边已创建"}},
},
},
"/api/projects/{id}/fact-edges/{edgeId}": map[string]interface{}{
"delete": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "删除事实边", "operationId": "deleteProjectFactEdge",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "edgeId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "删除成功"}},
},
},
"/api/projects/{id}/promote-attack-chain/{conversationId}": map[string]interface{}{
"post": map[string]interface{}{
"tags": []string{"项目管理"}, "summary": "将对话攻击链沉淀到项目事实图", "operationId": "promoteAttackChainToProject",
"parameters": []map[string]interface{}{
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
{"name": "conversationId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "沉淀结果(facts/edges/graph"}},
},
},
"/api/vulnerabilities": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"漏洞管理"},
+255 -21
View File
@@ -1,10 +1,12 @@
package handler
import (
"fmt"
"net/http"
"strconv"
"strings"
"cyberstrike-ai/internal/attackchain"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project"
@@ -223,26 +225,102 @@ func (h *ProjectHandler) DeleteProject(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"success": true})
}
type factLinkRequest struct {
From string `json:"from"`
Type string `json:"type"`
Confidence string `json:"confidence,omitempty"`
}
type upsertFactRequest struct {
FactKey string `json:"fact_key" binding:"required"`
Category string `json:"category"`
Summary string `json:"summary" binding:"required"`
Body string `json:"body"`
Confidence string `json:"confidence"`
Pinned bool `json:"pinned"`
RelatedVulnerabilityID string `json:"related_vulnerability_id"`
FactKey string `json:"fact_key" binding:"required"`
Category string `json:"category"`
Summary string `json:"summary" binding:"required"`
Body string `json:"body"`
Confidence string `json:"confidence"`
Pinned bool `json:"pinned"`
RelatedVulnerabilityID string `json:"related_vulnerability_id"`
Links []factLinkRequest `json:"links"`
LinksText *string `json:"links_text"`
}
// updateFactRequest 部分更新事实;指针字段省略=不修改,body 传 "" 可清空(仍走 merge 逻辑见 Upsert)。
type updateFactRequest struct {
FactKey *string `json:"fact_key"`
Category *string `json:"category"`
Summary *string `json:"summary"`
Body *string `json:"body"`
Confidence *string `json:"confidence"`
Pinned *bool `json:"pinned"`
RelatedVulnerabilityID *string `json:"related_vulnerability_id"`
ClearBody bool `json:"clear_body"`
FactKey *string `json:"fact_key"`
Category *string `json:"category"`
Summary *string `json:"summary"`
Body *string `json:"body"`
Confidence *string `json:"confidence"`
Pinned *bool `json:"pinned"`
RelatedVulnerabilityID *string `json:"related_vulnerability_id"`
ClearBody bool `json:"clear_body"`
Links *[]factLinkRequest `json:"links"`
LinksText *string `json:"links_text"`
}
func factLinksFromRequest(links []factLinkRequest, linksText *string) (*project.ParsedFactLinks, error) {
if len(links) > 0 {
parsed := &project.ParsedFactLinks{}
for i, l := range links {
from := strings.TrimSpace(l.From)
edgeType := strings.TrimSpace(l.Type)
if from == "" {
return nil, fmt.Errorf("links[%d] 须含 from", i)
}
if edgeType == "" {
return nil, fmt.Errorf("links[%d] 须含 type", i)
}
parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{
From: from, Type: edgeType, Confidence: strings.TrimSpace(l.Confidence),
})
}
return parsed, nil
}
if linksText != nil {
in, err := project.ParseFactLinksText(*linksText)
if err != nil {
return nil, err
}
return &project.ParsedFactLinks{Incoming: in}, nil
}
return &project.ParsedFactLinks{Incoming: []database.ProjectFactEdgeFromInput{}}, nil
}
type factWithLinksResponse struct {
*database.ProjectFact
OutgoingLinks []*database.ProjectFactEdge `json:"outgoing_links,omitempty"`
IncomingLinks []*database.ProjectFactEdge `json:"incoming_links,omitempty"`
LinkCounts *project.LinkCounts `json:"link_counts,omitempty"`
}
func (h *ProjectHandler) applyFactLinksAfterUpsert(projectID string, fact *database.ProjectFact, links []factLinkRequest, linksText *string, explicitLinks, parseBody bool) error {
if explicitLinks {
parsed, err := factLinksFromRequest(links, linksText)
if err != nil {
return err
}
return project.PersistFactLinksFromParsed(h.db, projectID, fact.FactKey, fact.SourceConversationID, parsed, true)
}
if parseBody {
inputs := project.ParseLinksFromBody(fact.Body)
if inputs == nil {
return nil
}
return project.PersistFactIncomingLinks(h.db, projectID, fact.FactKey, inputs, true)
}
return nil
}
func (h *ProjectHandler) factResponseWithLinks(projectID string, f *database.ProjectFact, includeLinks bool) interface{} {
if !includeLinks || f == nil {
return f
}
out, _ := h.db.ListOutgoingProjectFactEdges(projectID, f.FactKey)
in, _ := h.db.ListIncomingProjectFactEdges(projectID, f.FactKey)
return &factWithLinksResponse{
ProjectFact: f,
OutgoingLinks: out,
IncomingLinks: in,
}
}
// ListFacts GET /api/projects/:id/facts fact_key 查询参数可获取单条详情)
@@ -254,7 +332,8 @@ func (h *ProjectHandler) ListFacts(c *gin.Context) {
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, f)
includeLinks := c.Query("include_links") == "1" || c.Query("include_links") == "true"
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, f, includeLinks))
return
}
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100"))
@@ -285,7 +364,52 @@ func (h *ProjectHandler) ListFacts(c *gin.Context) {
}
list = filtered
}
c.JSON(http.StatusOK, list)
includeLinkCounts := c.Query("include_link_counts") == "1" || c.Query("include_link_counts") == "true"
if !includeLinkCounts {
c.JSON(http.StatusOK, list)
return
}
counts, err := project.LoadProjectFactLinkCounts(h.db, projectID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
out := make([]factWithLinksResponse, 0, len(list))
for _, f := range list {
item := factWithLinksResponse{ProjectFact: f}
if c, ok := counts[f.FactKey]; ok {
cc := c
item.LinkCounts = &cc
}
out = append(out, item)
}
c.JSON(http.StatusOK, out)
}
// GetFactGraph GET /api/projects/:id/fact-graph?view=path|full
func (h *ProjectHandler) GetFactGraph(c *gin.Context) {
projectID := c.Param("id")
if _, err := h.db.GetProject(projectID); err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"})
return
}
view := c.DefaultQuery("view", "path")
excludeDeprecated := true
if v := c.Query("exclude_deprecated"); v == "0" || v == "false" {
excludeDeprecated = false
}
graph, err := project.BuildProjectFactGraph(h.db, projectID, view, excludeDeprecated)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if graph.Nodes == nil {
graph.Nodes = []database.ProjectFactGraphNode{}
}
if graph.Edges == nil {
graph.Edges = []database.ProjectFactGraphEdge{}
}
c.JSON(http.StatusOK, graph)
}
// CreateFact POST /api/projects/:id/facts
@@ -295,8 +419,9 @@ func (h *ProjectHandler) CreateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
projectID := c.Param("id")
f := &database.ProjectFact{
ProjectID: c.Param("id"),
ProjectID: projectID,
FactKey: req.FactKey,
Category: req.Category,
Summary: req.Summary,
@@ -310,16 +435,24 @@ func (h *ProjectHandler) CreateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, created)
explicitLinks := req.Links != nil || req.LinksText != nil
if err := h.applyFactLinksAfterUpsert(projectID, created, req.Links, req.LinksText, explicitLinks, !explicitLinks); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
created, _ = h.db.GetProjectFactByKey(projectID, created.FactKey)
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, created, true))
}
// UpdateFact PUT /api/projects/:id/facts/:factId
func (h *ProjectHandler) UpdateFact(c *gin.Context) {
projectID := c.Param("id")
existing, err := h.db.GetProjectFact(c.Param("factId"))
if err != nil || existing.ProjectID != c.Param("id") {
if err != nil || existing.ProjectID != projectID {
c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"})
return
}
oldFactKey := existing.FactKey
var req updateFactRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -355,7 +488,29 @@ func (h *ProjectHandler) UpdateFact(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, updated)
if oldFactKey != updated.FactKey {
if err := h.db.RenameProjectFactKeyEdges(projectID, oldFactKey, updated.FactKey); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
}
if req.Links != nil || req.LinksText != nil {
var links []factLinkRequest
if req.Links != nil {
links = *req.Links
}
if err := h.applyFactLinksAfterUpsert(projectID, updated, links, req.LinksText, true, false); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
} else if req.ClearBody || req.Body != nil {
if err := h.applyFactLinksAfterUpsert(projectID, updated, nil, nil, false, true); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
}
updated, _ = h.db.GetProjectFactByKey(projectID, updated.FactKey)
c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, updated, true))
}
// DeleteFact DELETE /api/projects/:id/facts/:factId
@@ -408,3 +563,82 @@ func (h *ProjectHandler) RestoreFact(c *gin.Context) {
}
c.JSON(http.StatusOK, gin.H{"success": true})
}
type createFactEdgeRequest struct {
SourceFactKey string `json:"source_fact_key" binding:"required"`
TargetFactKey string `json:"target_fact_key" binding:"required"`
EdgeType string `json:"edge_type" binding:"required"`
Confidence string `json:"confidence"`
}
// ListFactEdges GET /api/projects/:id/fact-edges
func (h *ProjectHandler) ListFactEdges(c *gin.Context) {
projectID := c.Param("id")
edges, err := h.db.ListProjectFactEdgesByProject(projectID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if edges == nil {
edges = []*database.ProjectFactEdge{}
}
c.JSON(http.StatusOK, edges)
}
// CreateFactEdge POST /api/projects/:id/fact-edges
func (h *ProjectHandler) CreateFactEdge(c *gin.Context) {
projectID := c.Param("id")
var req createFactEdgeRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
edge, err := h.db.AddProjectFactEdge(projectID, database.ProjectFactEdgeInput{
To: req.TargetFactKey,
Type: req.EdgeType,
Confidence: req.Confidence,
}, req.SourceFactKey, "")
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if f, err := h.db.GetProjectFactByKey(projectID, req.TargetFactKey); err == nil {
in, _ := h.db.ListIncomingProjectFactEdges(projectID, req.TargetFactKey)
f.Body = project.SyncBodyLinksSection(f.Body, in)
_, _ = h.db.UpsertProjectFact(f)
}
c.JSON(http.StatusOK, edge)
}
// DeleteFactEdge DELETE /api/projects/:id/fact-edges/:edgeId
func (h *ProjectHandler) DeleteFactEdge(c *gin.Context) {
projectID := c.Param("id")
edgeID := c.Param("edgeId")
edge, err := h.db.GetProjectFactEdge(edgeID)
if err != nil || edge.ProjectID != projectID {
c.JSON(http.StatusNotFound, gin.H{"error": "边不存在"})
return
}
if err := h.db.DeleteProjectFactEdge(edgeID); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
if f, err := h.db.GetProjectFactByKey(projectID, edge.TargetFactKey); err == nil {
in, _ := h.db.ListIncomingProjectFactEdges(projectID, edge.TargetFactKey)
f.Body = project.SyncBodyLinksSection(f.Body, in)
_, _ = h.db.UpsertProjectFact(f)
}
c.JSON(http.StatusOK, gin.H{"success": true})
}
// PromoteAttackChain POST /api/projects/:id/promote-attack-chain/:conversationId
func (h *ProjectHandler) PromoteAttackChain(c *gin.Context) {
projectID := c.Param("id")
conversationID := c.Param("conversationId")
result, err := attackchain.PromoteToProject(h.db, projectID, conversationID)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, result)
}
+36
View File
@@ -7,6 +7,42 @@ import (
"go.uber.org/zap"
)
// agentSessionContextBlock 注入会话工作目录与项目黑板(用于 system prompt 追加块)。
func (h *AgentHandler) agentSessionContextBlock(conversationID string) string {
var parts []string
if ws := h.buildWorkspaceBlock(conversationID); ws != "" {
parts = append(parts, ws)
}
if bb := h.projectBlackboardBlock(conversationID); bb != "" {
parts = append(parts, bb)
}
return strings.Join(parts, "\n\n")
}
func (h *AgentHandler) buildWorkspaceBlock(conversationID string) string {
if h == nil || h.config == nil {
return ""
}
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" {
return ""
}
projectID := h.conversationProjectID(conversationID)
rel := project.WorkspaceRootDir(h.config.Agent.WorkspaceRootDir, projectID, conversationID)
abs, err := project.EnsureWorkspace(rel)
if err != nil {
if h.logger != nil {
h.logger.Warn("创建会话工作目录失败",
zap.String("conversationId", conversationID),
zap.String("projectId", projectID),
zap.String("path", rel),
zap.Error(err))
}
return ""
}
return project.BuildWorkspaceBlock(abs)
}
// projectBlackboardBlock 根据对话 ID 构建项目事实索引块(用于注入 system prompt)。
func (h *AgentHandler) projectBlackboardBlock(conversationID string) string {
if h == nil || h.db == nil || h.config == nil {
+4 -1
View File
@@ -447,7 +447,7 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string {
}
func (h *RobotHandler) cmdList() string {
convs, err := h.db.ListConversations(50, 0, "")
convs, err := h.db.ListConversations(50, 0, "", "")
if err != nil {
return "获取对话列表失败: " + err.Error()
}
@@ -594,6 +594,9 @@ func (h *RobotHandler) cmdDelete(platform, userID, convID string) string {
h.mu.Unlock()
h.deleteSessionBinding(sk)
}
if h.agentHandler != nil {
h.agentHandler.CancelRunningTaskForConversation(convID)
}
if err := h.db.DeleteConversation(convID); err != nil {
return "删除失败: " + err.Error()
}
+102
View File
@@ -37,6 +37,11 @@ type AgentTask struct {
// InterruptContinueNote 无 MCP 时「中断并继续」由用户在弹窗中填写的补充说明(Cancel 前写入,续跑轮次读取后清空)
InterruptContinueNote string `json:"-"`
// activeEinoExecuteCancel 当前进行中的 Eino filesystem execute 取消函数(与 MCP 工具并行,供中断并继续)
activeEinoExecuteCancel context.CancelFunc
// activeEinoExecuteAbortNote AbortActiveEinoExecute 写入的用户说明,由 execute 收尾时合并进工具结果
activeEinoExecuteAbortNote string
cancel func(error)
}
@@ -70,6 +75,103 @@ func (m *AgentTaskManager) UnregisterRunningTool(conversationID, executionID str
}
}
// RegisterActiveEinoExecute 登记进行中的 Eino filesystem execute(每会话同时仅一条)。
func (m *AgentTaskManager) RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc) {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" || cancel == nil {
return
}
m.mu.Lock()
defer m.mu.Unlock()
if t, ok := m.tasks[conversationID]; ok && t != nil {
t.activeEinoExecuteCancel = cancel
t.activeEinoExecuteAbortNote = ""
}
}
// UnregisterActiveEinoExecute execute 正常结束或已取消后清除登记。
func (m *AgentTaskManager) UnregisterActiveEinoExecute(conversationID string) {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
if t, ok := m.tasks[conversationID]; ok && t != nil {
t.activeEinoExecuteCancel = nil
t.activeEinoExecuteAbortNote = ""
}
}
// ConversationIDForActiveMCPExecution 根据当前登记的工具 executionId 反查会话 ID(供 MCP 监控页按 executionId 终止)。
func (m *AgentTaskManager) ConversationIDForActiveMCPExecution(executionID string) string {
executionID = strings.TrimSpace(executionID)
if executionID == "" {
return ""
}
m.mu.Lock()
defer m.mu.Unlock()
for convID, t := range m.tasks {
if t != nil && t.ActiveMCPExecutionID == executionID {
return convID
}
}
return ""
}
// ConversationIDForActiveEinoExecute 返回当前唯一进行 Eino execute 的会话 ID;多会话并行时返回空。
func (m *AgentTaskManager) ConversationIDForActiveEinoExecute() (string, bool) {
m.mu.Lock()
defer m.mu.Unlock()
var found string
count := 0
for convID, t := range m.tasks {
if t != nil && t.activeEinoExecuteCancel != nil {
found = convID
count++
}
}
if count == 1 {
return found, true
}
return "", false
}
// AbortActiveEinoExecute 终止当前 Eino execute 并暂存用户说明(与 MCP 工具终止一致)。
func (m *AgentTaskManager) AbortActiveEinoExecute(conversationID, note string) bool {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" {
return false
}
m.mu.Lock()
t, ok := m.tasks[conversationID]
if !ok || t == nil || t.activeEinoExecuteCancel == nil {
m.mu.Unlock()
return false
}
t.activeEinoExecuteAbortNote = strings.TrimSpace(note)
cancel := t.activeEinoExecuteCancel
m.mu.Unlock()
cancel()
return true
}
// TakeEinoExecuteAbortNote 读取并清空 execute 终止说明(execute 收尾时调用一次)。
func (m *AgentTaskManager) TakeEinoExecuteAbortNote(conversationID string) string {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" {
return ""
}
m.mu.Lock()
defer m.mu.Unlock()
if t, ok := m.tasks[conversationID]; ok && t != nil {
n := t.activeEinoExecuteAbortNote
t.activeEinoExecuteAbortNote = ""
return n
}
return ""
}
// SetInterruptContinueNote 在发起 ErrInterruptContinue 取消前写入用户补充说明(仅内存)。
func (m *AgentTaskManager) SetInterruptContinueNote(conversationID, note string) {
conversationID = strings.TrimSpace(conversationID)
@@ -0,0 +1,56 @@
package handler
import (
"context"
"testing"
"time"
)
func TestAbortActiveEinoExecute(t *testing.T) {
m := NewAgentTaskManager()
conv := "conv-eino-exec-abort"
ctx, cancel := context.WithCancel(context.Background())
_, err := m.StartTask(conv, "test", func(error) {})
if err != nil {
t.Fatalf("StartTask: %v", err)
}
m.RegisterActiveEinoExecute(conv, cancel)
done := make(chan struct{})
go func() {
<-ctx.Done()
close(done)
}()
if !m.AbortActiveEinoExecute(conv, "跳过域名收集") {
t.Fatal("expected abort to succeed")
}
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("execute cancel did not propagate")
}
if got := m.TakeEinoExecuteAbortNote(conv); got != "跳过域名收集" {
t.Fatalf("abort note = %q, want 跳过域名收集", got)
}
m.UnregisterActiveEinoExecute(conv)
if m.AbortActiveEinoExecute(conv, "") {
t.Fatal("second abort should fail when no active execute")
}
}
func TestConversationIDForActiveMCPExecution(t *testing.T) {
m := NewAgentTaskManager()
conv := "conv-mcp-exec"
_, err := m.StartTask(conv, "test", func(error) {})
if err != nil {
t.Fatalf("StartTask: %v", err)
}
m.RegisterRunningTool(conv, "exec-123")
if got := m.ConversationIDForActiveMCPExecution("exec-123"); got != conv {
t.Fatalf("got %q, want %q", got, conv)
}
if got := m.ConversationIDForActiveMCPExecution("missing"); got != "" {
t.Fatalf("missing should be empty, got %q", got)
}
}
+26
View File
@@ -11,7 +11,16 @@ type ToolRunRegistry interface {
UnregisterRunningTool(conversationID, executionID string)
}
// EinoExecuteRunRegistry 登记进行中的 Eino filesystem execute,供「中断并继续」终止 amass 等长命令。
type EinoExecuteRunRegistry interface {
RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc)
UnregisterActiveEinoExecute(conversationID string)
AbortActiveEinoExecute(conversationID, note string) bool
TakeEinoExecuteAbortNote(conversationID string) string
}
type toolRunRegistryCtxKey struct{}
type einoExecuteRunRegistryCtxKey struct{}
type mcpConversationIDCtxKey struct{}
// WithToolRunRegistry 将登记器注入 ctxEino / 原生 Agent 任务 ctx)。
@@ -31,6 +40,23 @@ func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry {
return v
}
// WithEinoExecuteRunRegistry 将 Eino execute 取消登记器注入 ctx。
func WithEinoExecuteRunRegistry(ctx context.Context, reg EinoExecuteRunRegistry) context.Context {
if ctx == nil || reg == nil {
return ctx
}
return context.WithValue(ctx, einoExecuteRunRegistryCtxKey{}, reg)
}
// EinoExecuteRunRegistryFromContext 取出 Eino execute 登记器(无则 nil)。
func EinoExecuteRunRegistryFromContext(ctx context.Context) EinoExecuteRunRegistry {
if ctx == nil {
return nil
}
v, _ := ctx.Value(einoExecuteRunRegistryCtxKey{}).(EinoExecuteRunRegistry)
return v
}
// WithMCPConversationID 将对话 ID 注入 ctx,供 CallTool 内与 executionId 关联。
func WithMCPConversationID(ctx context.Context, conversationID string) context.Context {
if ctx == nil {
+83 -16
View File
@@ -921,9 +921,8 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
return finalResult, executionID, nil
}
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}) string {
if s == nil {
return ""
}
@@ -931,21 +930,73 @@ func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]
args = map[string]interface{}{}
}
executionID := uuid.New().String()
now := time.Now()
failed := invokeErr != nil
exec := &ToolExecution{
execution := &ToolExecution{
ID: executionID,
ToolName: toolName,
Arguments: args,
StartTime: now,
EndTime: &now,
Duration: 0,
Status: "running",
StartTime: time.Now(),
}
s.mu.Lock()
s.executions[executionID] = execution
s.cleanupOldExecutions()
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
return executionID
}
// FinishToolExecution 完成先前 BeginToolExecution 创建的记录;executionID 为空时等同 RecordCompletedToolInvocation。
func (s *Server) FinishToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
if s == nil {
return ""
}
if args == nil {
args = map[string]interface{}{}
}
id := strings.TrimSpace(executionID)
if id == "" {
return s.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr)
}
now := time.Now()
failed := invokeErr != nil
var finalResult *ToolResult
s.mu.Lock()
exec, inMem := s.executions[id]
if !inMem || exec == nil {
exec = &ToolExecution{
ID: id,
ToolName: toolName,
Arguments: args,
StartTime: now,
}
s.executions[id] = exec
} else if toolName != "" {
exec.ToolName = toolName
}
if len(args) > 0 {
exec.Arguments = args
}
exec.EndTime = &now
if exec.StartTime.IsZero() {
exec.StartTime = now
}
exec.Duration = now.Sub(exec.StartTime)
if failed {
exec.Status = "failed"
exec.Error = invokeErr.Error()
st, msg := executionStatusAndMessage(invokeErr)
exec.Status = st
exec.Error = msg
if strings.TrimSpace(resultText) != "" {
exec.Result = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
exec.Result = finalResult
}
} else {
exec.Status = "completed"
@@ -953,15 +1004,31 @@ func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]
if strings.TrimSpace(text) == "" {
text = "(无输出)"
}
exec.Result = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
exec.Result = finalResult
}
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(exec); err != nil {
s.logger.Warn("RecordCompletedToolInvocation 保存失败", zap.Error(err))
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
s.updateStats(toolName, failed)
return executionID
s.updateStats(exec.ToolName, failed)
if s.storage != nil {
s.mu.Lock()
delete(s.executions, id)
s.mu.Unlock()
}
return id
}
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
return s.FinishToolExecution("", toolName, args, resultText, invokeErr)
}
// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文(如 reduction 后的 persisted-output)。
+2
View File
@@ -199,6 +199,8 @@ type ToolExecution struct {
StartTime time.Time `json:"startTime"`
EndTime *time.Time `json:"endTime,omitempty"`
Duration time.Duration `json:"duration,omitempty"`
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
ConversationID string `json:"conversationId,omitempty"`
}
// ToolStats 工具统计信息
+71
View File
@@ -0,0 +1,71 @@
package monitor
import (
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
const retentionPurgeInterval = time.Hour
// Service manages MCP tool execution monitor retention.
type Service struct {
db *database.DB
cfg *config.Config
logger *zap.Logger
}
// NewService creates a monitor 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.MonitorConfig{}.RetentionDaysEffective()
}
return s.cfg.Monitor.RetentionDaysEffective()
}
// PurgeExpired deletes tool execution rows older than retention_days when configured.
func (s *Service) PurgeExpired() {
if s == nil || s.db == nil || s.cfg == nil {
return
}
days := s.cfg.Monitor.RetentionDaysEffective()
if days <= 0 {
return
}
cutoff := time.Now().AddDate(0, 0, -days)
n, err := s.db.PurgeToolExecutionsBefore(cutoff)
if err != nil {
if s.logger != nil {
s.logger.Warn("清理过期 MCP 执行记录失败", zap.Error(err))
}
return
}
if n > 0 && s.logger != nil {
s.logger.Info("已清理过期 MCP 执行记录", zap.Int64("deleted", n), zap.Int("retention_days", days))
}
}
// StartRetentionLoop periodically purges expired tool execution 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("monitor retention tick completed")
}
}
}()
}
+94
View File
@@ -0,0 +1,94 @@
package monitor
import (
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestServicePurgeExpired_respectsZeroRetention(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()
exec := &mcp.ToolExecution{
ID: "ancient",
ToolName: "curl::get",
Arguments: map[string]interface{}{},
Status: "completed",
StartTime: mustParseTime(t, "2020-01-01T00:00:00Z"),
}
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
zero := 0
svc := NewService(db, &config.Config{
Monitor: config.MonitorConfig{RetentionDays: &zero},
}, zap.NewNop())
svc.PurgeExpired()
if _, err := db.GetToolExecution("ancient"); err != nil {
t.Fatalf("record should remain when retention_days=0: %v", err)
}
}
func TestServicePurgeExpired_deletesOldRows(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()
exec := &mcp.ToolExecution{
ID: "ancient",
ToolName: "curl::get",
Arguments: map[string]interface{}{},
Status: "completed",
StartTime: mustParseTime(t, "2020-01-01T00:00:00Z"),
}
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution: %v", err)
}
days := 90
svc := NewService(db, &config.Config{
Monitor: config.MonitorConfig{RetentionDays: &days},
}, zap.NewNop())
svc.PurgeExpired()
if _, err := db.GetToolExecution("ancient"); err == nil {
t.Fatal("record should be purged when older than retention_days")
}
}
func TestRetentionDaysEffective_defaults(t *testing.T) {
got := config.MonitorConfig{}.RetentionDaysEffective()
if got != 90 {
t.Fatalf("default = %d, want 90", got)
}
zero := 0
cfg := config.MonitorConfig{RetentionDays: &zero}
if cfg.RetentionDaysEffective() != 0 {
t.Fatalf("zero = %d, want 0", cfg.RetentionDaysEffective())
}
}
func mustParseTime(t *testing.T, value string) time.Time {
t.Helper()
parsed, err := time.Parse(time.RFC3339, value)
if err != nil {
t.Fatalf("parse time: %v", err)
}
return parsed
}
@@ -0,0 +1,104 @@
package multiagent
import (
"context"
"strings"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// continuationSessionMarker matches Cursor / IDE session-resume user injections.
const continuationSessionMarker = "This session is being continued from a previous conversation"
// continuationUserDedupMiddleware keeps only the latest session-resume user message when
// multiple continuation injections were stacked (e.g. after repeated out-of-context resumes).
type continuationUserDedupMiddleware struct {
adk.BaseChatModelAgentMiddleware
logger *zap.Logger
phase string
}
func newContinuationUserDedupMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
return &continuationUserDedupMiddleware{logger: logger, phase: phase}
}
func (m *continuationUserDedupMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, error) {
_ = mc
if m == nil || state == nil || len(state.Messages) == 0 {
return ctx, state, nil
}
deduped, dropped := dedupContinuationUserMessages(state.Messages)
if dropped == 0 {
return ctx, state, nil
}
if m.logger != nil {
m.logger.Info("eino continuation user messages deduplicated",
zap.String("phase", m.phase),
zap.Int("dropped", dropped),
zap.Int("messages_before", len(state.Messages)),
zap.Int("messages_after", len(deduped)),
)
}
out := *state
out.Messages = deduped
return ctx, &out, nil
}
func adkUserMessageText(msg adk.Message) string {
if msg == nil {
return ""
}
var b strings.Builder
if s := strings.TrimSpace(msg.Content); s != "" {
b.WriteString(s)
}
for _, part := range msg.UserInputMultiContent {
if part.Type == schema.ChatMessagePartTypeText {
if s := strings.TrimSpace(part.Text); s != "" {
if b.Len() > 0 {
b.WriteByte('\n')
}
b.WriteString(s)
}
}
}
return b.String()
}
func isContinuationUserMessage(msg adk.Message) bool {
if msg == nil || msg.Role != schema.User {
return false
}
return strings.Contains(adkUserMessageText(msg), continuationSessionMarker)
}
func dedupContinuationUserMessages(msgs []adk.Message) ([]adk.Message, int) {
lastIdx := -1
contCount := 0
for i, msg := range msgs {
if !isContinuationUserMessage(msg) {
continue
}
contCount++
lastIdx = i
}
if contCount <= 1 {
return msgs, 0
}
out := make([]adk.Message, 0, len(msgs)-(contCount-1))
dropped := 0
for i, msg := range msgs {
if isContinuationUserMessage(msg) && i != lastIdx {
dropped++
continue
}
out = append(out, msg)
}
return out, dropped
}
@@ -0,0 +1,65 @@
package multiagent
import (
"context"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func continuationUser(text string) adk.Message {
return &schema.Message{
Role: schema.User,
UserInputMultiContent: []schema.MessageInputPart{
{Type: schema.ChatMessagePartTypeText, Text: continuationSessionMarker + "\n" + text},
{Type: schema.ChatMessagePartTypeText, Text: "Please continue the conversation from where we left it off."},
},
}
}
func TestDedupContinuationUserMessages_KeepsLatest(t *testing.T) {
msgs := []adk.Message{
continuationUser("summary old"),
schema.UserMessage("real task"),
continuationUser("summary new"),
}
out, dropped := dedupContinuationUserMessages(msgs)
if dropped != 1 {
t.Fatalf("dropped=%d want 1", dropped)
}
if len(out) != 2 {
t.Fatalf("len=%d want 2", len(out))
}
if out[0].Role != schema.User || adkUserMessageText(out[0]) != "real task" {
t.Fatalf("first should remain real task, got %q", adkUserMessageText(out[0]))
}
if !strings.Contains(adkUserMessageText(out[1]), "summary new") {
t.Fatalf("latest continuation not kept: %q", adkUserMessageText(out[1]))
}
}
func TestDedupContinuationUserMessages_NoOpSingle(t *testing.T) {
msgs := []adk.Message{continuationUser("only"), schema.UserMessage("task")}
out, dropped := dedupContinuationUserMessages(msgs)
if dropped != 0 || len(out) != 2 {
t.Fatalf("unexpected change dropped=%d len=%d", dropped, len(out))
}
}
func TestContinuationUserDedupMiddleware(t *testing.T) {
mw := newContinuationUserDedupMiddleware(nil, "test")
state := &adk.ChatModelAgentState{Messages: []adk.Message{
continuationUser("old"),
continuationUser("new"),
schema.UserMessage("task"),
}}
_, out, err := mw.(*continuationUserDedupMiddleware).BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatal(err)
}
if len(out.Messages) != 2 {
t.Fatalf("want 2 messages after dedup, got %d", len(out.Messages))
}
}
+226 -93
View File
@@ -18,6 +18,7 @@ import (
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/einoobserve"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/security"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
@@ -90,7 +91,7 @@ type einoADKRunLoopArgs struct {
FilesystemMonitorRecord einomcp.ExecutionRecorder
MCPExecutionBinder *MCPExecutionBinder
// ToolInvokeNotify 与 einomcp.ToolsFromDefinitions 共享:run loop 在迭代前 SetMCP 桥 Fire 以补全 tool_result。
// ToolInvokeNotify 与 einomcp.ToolsFromDefinitions 共享:run loop 在迭代前 Setexecute/MCP 桥 Fire 时立即推送 tool_resultADK 晚到经 toolResultSent 去重)
ToolInvokeNotify *einomcp.ToolInvokeNotifyHolder
DA adk.Agent
@@ -196,6 +197,16 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
pendingByID[tc.ToolCallID] = tc
pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
}
markPendingWithMonitor := func(tc toolCallPendingInfo) {
markPending(tc)
beginEinoADKFilesystemToolMonitor(
args.FilesystemMonitorAgent,
args.FilesystemMonitorRecord,
args.MCPExecutionBinder,
tc.ToolCallID,
tc.ToolName,
)
}
popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
pendingMu.Lock()
defer pendingMu.Unlock()
@@ -331,7 +342,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
toolCallID = tid
}
recordPendingExecuteStdoutDup(toolName, content, isErr)
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
@@ -341,8 +352,21 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
if args.ToolInvokeNotify != nil {
args.ToolInvokeNotify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
removePendingByID(strings.TrimSpace(toolCallID))
// tool_result 仅由下方 ADK schema.Tool 事件推送,正文与送入模型的上下文一致(含 reduction 截断)
// 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)
})
}
@@ -383,6 +407,12 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
}
runner := adk.NewRunner(ctx, runnerCfg)
startRunnerIter := func(runMsgs []adk.Message) *adk.AsyncIterator[*adk.AgentEvent] {
if checkPointID != "" {
return runner.Run(ctx, runMsgs, adk.WithCheckPointID(checkPointID))
}
return runner.Run(ctx, runMsgs)
}
var iter *adk.AsyncIterator[*adk.AgentEvent]
if cpStore != nil && checkPointID != "" {
if _, existed, getErr := cpStore.Get(ctx, checkPointID); getErr != nil {
@@ -422,12 +452,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
}
if iter == nil {
if checkPointID != "" {
iter = runner.Run(ctx, msgs, adk.WithCheckPointID(checkPointID))
} else {
iter = runner.Run(ctx, msgs)
}
iter = startRunnerIter(msgs)
}
transientRetrier := newEinoTransientRunRetrier(einoTransientRunRetryPolicyFromArgs(args))
handleRunErr := func(runErr error) error {
if runErr == nil {
return nil
@@ -480,26 +507,67 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
return runErr
}
// maybeRetryTransientRun:不在此层 runner.Run/Resume;由 handler 落库 + loadHistoryFromAgentTrace 分段续跑(同中断并继续)。
maybeRetryTransientRun := func(runErr error) (retry bool, fatal error) {
if runErr == nil || !isEinoTransientRunError(runErr) {
maybeRetryTransientRun := func(runErr error) (restarted bool, fatal error) {
if runErr == nil {
return false, nil
}
if !isEinoTransientRunError(runErr) {
return false, handleRunErr(runErr)
}
restarted, restartMsgs, ctxSource, backoff, retErr := transientRetrier.tryRetry(
ctx, runErr, args, baseMsgs, runAccumulatedMsgs, baseAccumulatedCount,
)
if retErr != nil {
flushAllPendingAsFailed(runErr)
if logger != nil {
logger.Warn("eino transient retry exhausted",
zap.Error(retErr),
zap.String("orchestration", orchMode),
zap.Int("maxAttempts", transientRetrier.maxAttempts()))
}
return false, retErr
}
if !restarted {
return false, nil
}
attemptNo := transientRetrier.attempt()
maxAttempts := transientRetrier.maxAttempts()
if logger != nil {
logger.Warn("eino transient error, ending run segment for handler resume",
logger.Warn("eino transient error, retrying after backoff",
zap.Error(runErr),
zap.String("orchestration", orchMode))
zap.String("orchestration", orchMode),
zap.Int("attempt", attemptNo),
zap.Int("maxAttempts", maxAttempts),
zap.Duration("backoff", backoff))
}
if progress != nil {
progress("eino_run_retry", "遇到临时错误(限流或网络波动),将保存上下文并重试…", map[string]interface{}{
progress("eino_run_retry", fmt.Sprintf("遇到临时错误(限流或网络波动),%d 秒后第 %d/%d 次重试…", int(backoff.Seconds()), attemptNo, maxAttempts), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"orchestration": orchMode,
"error": runErr.Error(),
"resumeKind": "trace_segment",
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"backoffSec": int(backoff.Seconds()),
})
progress("eino_run_retry", "已恢复上下文,正在重试…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"orchestration": orchMode,
"attempt": attemptNo,
"contextSource": string(ctxSource),
})
}
return false, ErrTransientRetryContinue
msgs = restartMsgs
iter = startRunnerIter(msgs)
return true, nil
}
// 仅在退避重试后真正收到数据/完成一步时清零,避免重启后首个无错 ADK 事件误把计数打回 0。
confirmTransientRetryRecovery := func() {
if transientRetrier.attempt() > 0 {
transientRetrier.reset()
}
}
takePartial := func(runErr error) (*RunResult, error) {
@@ -514,10 +582,10 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
for {
// 检测 context 取消(用户关闭浏览器、请求超时等),flush pending 工具状态避免 UI 卡在 "执行中"
select {
case <-ctx.Done():
flushAllPendingAsFailed(ctx.Err())
// iter.Next 可能长时间阻塞(工具执行、模型推理);须与 ctx 联动,否则取消/超时无法及时 flush pending
ev, ok, iterCtxErr := nextAgentEventWithContext(ctx, iter)
if iterCtxErr != nil {
flushAllPendingAsFailed(iterCtxErr)
if progress != nil {
if isInterruptContinue(ctx) {
progress("progress", "已暂停当前输出,正在合并用户补充并继续…", map[string]interface{}{
@@ -526,17 +594,14 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"kind": "interrupt_continue",
})
} else {
progress("error", "Request cancelled / 请求已取消", map[string]interface{}{
progress("error", iterCtxErr.Error(), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
})
}
}
return takePartial(ctx.Err())
default:
return takePartial(iterCtxErr)
}
ev, ok := iter.Next()
if !ok {
// iter 结束并不总是“正常完成”:
// 当取消/超时发生在 iter.Next() 阻塞期间时,可能直接返回 !ok。
@@ -583,9 +648,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
continue
}
if ev.Err != nil {
if _, retErr := maybeRetryTransientRun(ev.Err); retErr != nil {
restarted, retErr := maybeRetryTransientRun(ev.Err)
if retErr != nil {
return takePartial(retErr)
}
if restarted {
continue
}
}
if ev.AgentName != "" && progress != nil {
iterEinoAgent := orchestratorName
@@ -630,13 +699,16 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
}
}
// 仅在代理切换时更新进度标题;同一代理的每个 ADK 事件不再重复刷 progress。
if einoLastAgent != ev.AgentName {
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
"conversationId": conversationID,
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"orchestration": orchMode,
})
}
einoLastAgent = ev.AgentName
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
"conversationId": conversationID,
"einoAgent": ev.AgentName,
"einoRole": einoRoleTag(ev.AgentName),
"orchestration": orchMode,
})
}
if ev.Output == nil || ev.Output.MessageOutput == nil {
continue
@@ -645,34 +717,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
if mv.IsStreaming && mv.MessageStream != nil && mv.Role == schema.Tool {
toolName := strings.TrimSpace(mv.ToolName)
var toolBuf strings.Builder
streamToolCallID := ""
var toolStreamRecvErr error
for {
chunk, rerr := mv.MessageStream.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
toolStreamRecvErr = rerr
break
}
if chunk == nil {
continue
}
if chunk.Content != "" {
toolBuf.WriteString(chunk.Content)
}
if tid := strings.TrimSpace(chunk.ToolCallID); tid != "" {
streamToolCallID = tid
}
}
content := toolBuf.String()
isErr := false
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
isErr = true
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
}
content, streamToolCallID, toolStreamRecvErr := recvSchemaMessageStream(ctx, mv.MessageStream)
isErr := einoToolResultIsError(toolName, content)
content = einoToolResultBody(content)
if streamToolCallID != "" {
opts := []schema.ToolMessageOption{schema.WithToolName(toolName)}
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.ToolMessage(content, streamToolCallID, opts...))
@@ -684,6 +731,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
zap.String("agent", ev.AgentName),
zap.String("tool", toolName))
}
if toolStreamRecvErr == nil {
confirmTransientRetryRecovery()
}
continue
}
@@ -931,7 +981,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
lastToolChunk = mergeMessageToolCalls(&schema.Message{ToolCalls: merged})
}
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPendingWithMonitor)
// 流式路径此前只把 tool_calls 推给进度 UI,未写入 runAccumulatedMsgs;落库后 loadHistory→RepairOrphan 会删掉全部 tool 结果,表现为「续跑/下轮失忆」。
if lastToolChunk != nil && len(lastToolChunk.ToolCalls) > 0 {
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage("", lastToolChunk.ToolCalls))
@@ -948,9 +998,15 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"einoRole": einoRoleTag(ev.AgentName),
})
}
if _, retErr := maybeRetryTransientRun(streamRecvErr); retErr != nil {
restarted, retErr := maybeRetryTransientRun(streamRecvErr)
if retErr != nil {
return takePartial(retErr)
}
if restarted {
continue
}
} else {
confirmTransientRetryRecovery()
}
continue
}
@@ -960,7 +1016,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
continue
}
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPendingWithMonitor)
if mv.Role == schema.Assistant {
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
@@ -1035,15 +1091,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
content := msg.Content
isErr := false
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
isErr = true
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
}
isErr := einoToolResultIsError(toolName, content)
content = einoToolResultBody(content)
toolCallID := strings.TrimSpace(msg.ToolCallID)
tryEmitToolResultProgress(toolName, content, toolCallID, isErr, ev.AgentName)
}
confirmTransientRetryRecovery()
}
mcpIDsMu.Lock()
@@ -1054,32 +1108,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false,
)
if shouldEinoEmptyResponseContinue(out, emptyHint, len(runAccumulatedMsgs), baseAccumulatedCount) {
if logger != nil {
logger.Info("eino empty response, ending run segment for handler resume",
zap.String("conversationId", conversationID),
zap.String("orchestration", orchMode),
zap.Int("traceMessages", len(runAccumulatedMsgs)))
}
if progress != nil {
progress("eino_empty_response_continue", "会话已结束但未产生助手正文,正在基于轨迹自动续跑…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"resumeKind": "trace_segment",
})
}
return out, ErrEmptyResponseContinue
}
return out, nil
}
func shouldEinoEmptyResponseContinue(out *RunResult, emptyHint string, accumulatedLen, baseCount int) bool {
if out == nil || accumulatedLen <= baseCount {
return false
}
return strings.TrimSpace(out.Response) == strings.TrimSpace(emptyHint)
}
func persistTraceSource(args *einoADKRunLoopArgs, fallback []adk.Message) []adk.Message {
if args != nil && args.ModelFacingTrace != nil {
if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 {
@@ -1094,17 +1125,119 @@ func einoPartialRunLastOutputHint() string {
"[Run ended abnormally; continue from the trace above without repeating completed steps.]"
}
// friendlyEinoExecuteInvokeTail 将 Eino execute 等非 MCP 路径的结尾错误转成简短提示;其它情况保留原 error 文本
// friendlyEinoExecuteInvokeTail 将 Eino execute 超时/中断/流异常转为简短提示
// 命令非零退出(ExecuteExitError)已有 exec 对齐的正文,不再追加「执行未正常结束」。
func friendlyEinoExecuteInvokeTail(invokeErr error) string {
if invokeErr == nil {
return ""
}
var exitErr *ExecuteExitError
if errors.As(invokeErr, &exitErr) {
return ""
}
if errors.Is(invokeErr, context.DeadlineExceeded) {
return einoExecuteTimeoutUserHint()
}
if errors.Is(invokeErr, context.Canceled) {
return ""
}
if strings.Contains(invokeErr.Error(), "shell inactivity timeout") {
return ""
}
return "[执行未正常结束] " + invokeErr.Error()
}
// einoToolResultIsError 统一判断 Eino 工具结果是否应标记为错误(与 MCP exec 的 IsError 对齐)。
func einoToolResultIsError(toolName, content string) bool {
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
return true
}
if strings.TrimSpace(toolName) == "execute" && security.IsCommandFailureResult(content) {
return true
}
return false
}
// einoToolResultBody 去掉工具错误前缀,返回展示/持久化正文。
func einoToolResultBody(content string) string {
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
return strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
}
return content
}
// nextAgentEventWithContext 在 ctx 取消时不再无限阻塞于 iter.Next()(工具执行/模型推理期间常见)。
func nextAgentEventWithContext(ctx context.Context, iter *adk.AsyncIterator[*adk.AgentEvent]) (ev *adk.AgentEvent, ok bool, ctxErr error) {
if iter == nil {
return nil, false, nil
}
type nextRes struct {
ev *adk.AgentEvent
ok bool
}
ch := make(chan nextRes, 1)
go func() {
e, o := iter.Next()
ch <- nextRes{e, o}
}()
select {
case <-ctx.Done():
return nil, false, ctx.Err()
case res := <-ch:
return res.ev, res.ok, nil
}
}
// recvSchemaMessageStream 消费 ADK Tool 流式结果;ctx 取消时立即返回,避免 amass 等无输出时永久阻塞。
func recvSchemaMessageStream(ctx context.Context, stream *schema.StreamReader[*schema.Message]) (content, toolCallID string, recvErr error) {
if stream == nil {
return "", "", nil
}
type streamMsg struct {
chunk *schema.Message
err error
}
recvCh := make(chan streamMsg, 8)
go func() {
defer close(recvCh)
for {
ch, rerr := stream.Recv()
recvCh <- streamMsg{chunk: ch, err: rerr}
if rerr != nil {
return
}
}
}()
var buf strings.Builder
for {
select {
case <-ctx.Done():
return buf.String(), toolCallID, ctx.Err()
case sm, open := <-recvCh:
if !open {
return buf.String(), toolCallID, nil
}
rerr := sm.err
if errors.Is(rerr, io.EOF) {
return buf.String(), toolCallID, nil
}
if rerr != nil {
return buf.String(), toolCallID, rerr
}
chunk := sm.chunk
if chunk == nil {
continue
}
if chunk.Content != "" {
buf.WriteString(chunk.Content)
}
if tid := strings.TrimSpace(chunk.ToolCallID); tid != "" {
toolCallID = tid
}
}
}
}
func buildEinoRunResultFromAccumulated(
orchMode string,
runAccumulatedMsgs []adk.Message,
@@ -0,0 +1,74 @@
package multiagent
import (
"context"
"errors"
"io"
"testing"
"time"
"github.com/cloudwego/eino/schema"
)
func TestRecvSchemaMessageStream_EOF(t *testing.T) {
sr, sw := schema.Pipe[*schema.Message](4)
_ = sw.Send(schema.ToolMessage("hello", "tc-1"), nil)
sw.Close()
content, tid, err := recvSchemaMessageStream(context.Background(), sr)
if err != nil {
t.Fatalf("unexpected err: %v", err)
}
if content != "hello" {
t.Fatalf("content=%q want hello", content)
}
if tid != "tc-1" {
t.Fatalf("toolCallID=%q want tc-1", tid)
}
}
func TestRecvSchemaMessageStream_ContextCancel(t *testing.T) {
sr, sw := schema.Pipe[*schema.Message](4)
t.Cleanup(func() { sw.Close() })
ctx, cancel := context.WithCancel(context.Background())
go func() {
time.Sleep(30 * time.Millisecond)
cancel()
}()
content, _, err := recvSchemaMessageStream(ctx, sr)
if !errors.Is(err, context.Canceled) {
t.Fatalf("want context.Canceled, got %v content=%q", err, content)
}
}
func TestRecvSchemaMessageStream_RecvError(t *testing.T) {
sr, sw := schema.Pipe[*schema.Message](4)
want := errors.New("stream broken")
_ = sw.Send(nil, want)
sw.Close()
_, _, err := recvSchemaMessageStream(context.Background(), sr)
if !errors.Is(err, want) {
t.Fatalf("want %v, got %v", want, err)
}
}
func TestRecvSchemaMessageStream_NilStream(t *testing.T) {
content, tid, err := recvSchemaMessageStream(context.Background(), nil)
if err != nil || content != "" || tid != "" {
t.Fatalf("nil stream: content=%q tid=%q err=%v", content, tid, err)
}
}
func TestRecvSchemaMessageStream_EOFViaEmptyRead(t *testing.T) {
sr, sw := schema.Pipe[*schema.Message](4)
_ = sw.Send(nil, io.EOF)
sw.Close()
_, _, err := recvSchemaMessageStream(context.Background(), sr)
if err != nil {
t.Fatalf("EOF should not surface as error, got %v", err)
}
}
@@ -0,0 +1,50 @@
package multiagent
import (
"github.com/cloudwego/eino/adk"
"go.uber.org/zap"
)
// einoChatModelTailConfig configures middleware appended after reduction/skill/plantask
// and immediately before each ChatModel invocation pipeline completes.
//
// Order (best practice):
// 1. system merge — accurate token count for summarization
// 2. continuation user dedup — drop stale session-resume injections
// 3. summarization
// 4. orphan tool prune
// 5. telemetry
// 6. model-facing trace snapshot
type einoChatModelTailConfig struct {
logger *zap.Logger
phase string
summarization adk.ChatModelAgentMiddleware
modelName string
conversationID string
trace *modelFacingTraceHolder
skipOrphanPruner bool
skipTelemetry bool
skipTrace bool
}
func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware, cfg einoChatModelTailConfig) []adk.ChatModelAgentMiddleware {
handlers = append(handlers, newSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
handlers = append(handlers, newContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
if cfg.summarization != nil {
handlers = append(handlers, cfg.summarization)
}
if !cfg.skipOrphanPruner {
handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
}
if !cfg.skipTelemetry {
if teleMw := newEinoModelInputTelemetryMiddleware(cfg.logger, cfg.modelName, cfg.conversationID, cfg.phase); teleMw != nil {
handlers = append(handlers, teleMw)
}
}
if !cfg.skipTrace && cfg.trace != nil {
if capMw := newModelFacingTraceMiddleware(cfg.trace); capMw != nil {
handlers = append(handlers, capMw)
}
}
return handlers
}
@@ -1,21 +0,0 @@
package multiagent
import "testing"
func TestShouldEinoEmptyResponseContinue(t *testing.T) {
t.Parallel()
hint := "(empty hint)"
out := &RunResult{Response: hint}
if !shouldEinoEmptyResponseContinue(out, hint, 3, 1) {
t.Fatal("expected continue when response is empty hint and trace grew")
}
if shouldEinoEmptyResponseContinue(out, hint, 1, 1) {
t.Fatal("expected no continue when trace did not grow")
}
if shouldEinoEmptyResponseContinue(&RunResult{Response: "hello"}, hint, 3, 1) {
t.Fatal("expected no continue when response has content")
}
if shouldEinoEmptyResponseContinue(nil, hint, 3, 1) {
t.Fatal("expected no continue for nil result")
}
}
@@ -0,0 +1,114 @@
package multiagent
import (
"context"
"errors"
"io"
"strings"
"testing"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/security"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/schema"
)
type mockStreamingShellExitFail struct {
output string
code int
}
func (m *mockStreamingShellExitFail) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
go func() {
defer outW.Close()
if m.output != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: m.output}, nil)
}
code := m.code
_ = outW.Send(&filesystem.ExecuteResponse{ExitCode: &code}, nil)
}()
return outR, nil
}
func TestEinoStreamingShellWrap_CommandFailureFormat(t *testing.T) {
inner := &mockStreamingShellExitFail{
output: "sudo: a password is required\n",
code: 1,
}
notify := einomcp.NewToolInvokeNotifyHolder()
var firedBody string
var firedSuccess bool
var firedErr error
notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
firedBody = content
firedSuccess = success
firedErr = invokeErr
})
wrap := &einoStreamingShellWrap{inner: inner, invokeNotify: notify}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var stream 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 {
stream.WriteString(resp.Output)
}
}
if firedSuccess {
t.Fatal("expected success=false")
}
var exitErr *ExecuteExitError
if !errors.As(firedErr, &exitErr) || exitErr.Code != 1 {
t.Fatalf("expected ExecuteExitError code 1, got %v", firedErr)
}
if !strings.HasPrefix(firedBody, einomcp.ToolErrorPrefix) {
t.Fatalf("missing tool error prefix: %q", firedBody)
}
body := strings.TrimPrefix(firedBody, einomcp.ToolErrorPrefix)
if body != security.FormatCommandFailureResult(1, "sudo: a password is required\n") {
t.Fatalf("fire body = %q", body)
}
if !strings.Contains(stream.String(), "sudo:") {
t.Fatalf("stream missing sudo output: %q", stream.String())
}
if strings.Contains(stream.String(), "command exited with non-zero") {
t.Fatalf("stream has legacy noise: %q", stream.String())
}
if strings.Contains(stream.String(), "执行未正常结束") {
t.Fatalf("stream has abnormal tail: %q", stream.String())
}
if !security.IsCommandFailureResult(stream.String()) {
t.Fatalf("stream missing failure status line: %q", stream.String())
}
if tail := friendlyEinoExecuteInvokeTail(firedErr); tail != "" {
t.Fatalf("unexpected invoke tail: %q", tail)
}
if !einoToolResultIsError("execute", firedBody) {
t.Fatal("expected isError for execute failure")
}
}
func TestFriendlyEinoExecuteInvokeTail(t *testing.T) {
if friendlyEinoExecuteInvokeTail(&ExecuteExitError{Code: 1}) != "" {
t.Fatal("exit error should not get abnormal tail")
}
if !strings.Contains(friendlyEinoExecuteInvokeTail(context.DeadlineExceeded), "Timed out") {
t.Fatal("deadline should get timeout hint")
}
if friendlyEinoExecuteInvokeTail(errors.New("broken pipe")) == "" {
t.Fatal("unexpected error should get tail")
}
}
+22 -7
View File
@@ -7,11 +7,25 @@ import (
"cyberstrike-ai/internal/einomcp"
)
// newEinoExecuteMonitorCallback 在 Eino filesystem execute 结束时写入 MCP 监控库并 recorder(executionId)
// 与 CallTool 路径一致,供助手消息展示「渗透测试详情」芯片
func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRecorder) func(toolCallID, command, stdout string, success bool, invokeErr error) {
return func(toolCallID, command, stdout string, success bool, invokeErr error) {
if ag == nil || recorder == nil {
// newEinoExecuteMonitorCallbacks 在 Eino filesystem execute 开始/结束时写入 MCP 监控库并 recorder(executionId)
// 与 CallTool 路径一致,使监控页能展示「执行中」状态
func newEinoExecuteMonitorCallbacks(ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
begin func(toolCallID, command string) string,
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
) {
begin = func(toolCallID, command string) string {
if ag == nil {
return ""
}
args := map[string]interface{}{"command": command}
id := ag.BeginLocalToolExecution("execute", args)
if id != "" && recorder != nil {
recorder(id, toolCallID)
}
return id
}
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
if ag == nil {
return
}
var err error
@@ -23,9 +37,10 @@ func newEinoExecuteMonitorCallback(ag *agent.Agent, recorder einomcp.ExecutionRe
}
}
args := map[string]interface{}{"command": command}
id := ag.RecordLocalToolExecution("execute", args, stdout, err)
if id != "" {
id := ag.FinishLocalToolExecution(executionID, "execute", args, stdout, err)
if id != "" && recorder != nil && executionID == "" {
recorder(id, toolCallID)
}
}
return begin, finish
}
@@ -6,9 +6,11 @@ import (
"fmt"
"io"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/security"
"github.com/cloudwego/eino/adk/filesystem"
@@ -34,13 +36,22 @@ func einoExecuteTimeoutUserHint() string {
return "已超时终止 · Timed out"
}
// einoExecuteRecvErrIsToolTimeout 判断 Recv 错误是否由 agent.tool_timeout_minutes 触发。
// WithTimeout 到期后 local 侧常报 canceled / exit -1,但 execCtx.Err() 仍为 DeadlineExceeded。
func einoExecuteRecvErrIsToolTimeout(rerr error, tctx context.Context) bool {
if tctx != nil && errors.Is(tctx.Err(), context.DeadlineExceeded) {
return true
}
return errors.Is(rerr, context.DeadlineExceeded)
}
// einoStreamingShellWrap 包装 Eino filesystem 使用的 StreamingShellcloudwego eino-ext local.Local)。
// 官方 execute 工具默认走 ExecuteStreaming 且不设 RunInBackendGround;末尾带 & 时子进程仍与管道相连,
// streamStdout 按行读取会在无换行输出时长时间阻塞(与 MCP 工具 exec 的独立实现不同)。
// 对「完全后台」命令自动开启 RunInBackendGround,与 local.runCmdInBackground 行为对齐。
//
// 使用 Pipe 将内层流转发给调用方:在 inner EOF 后、关闭 Pipe 前同步调用 ToolInvokeNotify.Fire
// 保证 run loop 在模型开始下一轮输出前已记录 execute 结果(用于 UI 与「重复助手复述」去重)
// run loop 收到 Fire 后立即推送 tool_resulttoolResultSent 去重),避免 ADK Tool 事件迟到时 UI 卡在「执行中」
//
// 若 inner 在校验阶段直接返回 error(未建立 reader),不会进入下方 goroutine,也必须 Fire
// 否则 pending tool_call 要等整轮 run 结束才被 force-close,与已展示的助手/工具软错误文案不同步。
@@ -52,8 +63,11 @@ type einoStreamingShellWrap struct {
outputChunk func(toolName, toolCallID, chunk string)
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
toolTimeoutMinutes int
// recordMonitor 在 execute 流结束后写入 tool_executions 并 recorder(executionId),使「渗透测试详情」与常规 MCP 一致
recordMonitor func(toolCallID, command, stdout string, success bool, invokeErr error)
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭
shellNoOutputTimeoutSec int
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
beginMonitor func(toolCallID, command string) string
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
}
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
@@ -65,33 +79,65 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
}
req := *input
userCmd := strings.TrimSpace(req.Command)
tid := strings.TrimSpace(compose.GetToolCallID(ctx))
agentTag := strings.TrimSpace(w.einoAgentName)
if security.IsBackgroundShellCommand(req.Command) && !req.RunInBackendGround {
req.RunInBackendGround = true
}
req.Command = prependPythonUnbufferedEnv(req.Command)
tid := strings.TrimSpace(compose.GetToolCallID(ctx))
agentTag := strings.TrimSpace(w.einoAgentName)
req.Command = security.PrepareNonInteractiveShellCommand(prependPythonUnbufferedEnv(req.Command))
convID := mcp.MCPConversationIDFromContext(ctx)
execReg := mcp.EinoExecuteRunRegistryFromContext(ctx)
execCtx := ctx
var execCancel context.CancelFunc
var monitorExecID string
if w.beginMonitor != nil {
monitorExecID = w.beginMonitor(tid, userCmd)
}
if monitorExecID != "" && convID != "" {
if toolReg := mcp.ToolRunRegistryFromContext(ctx); toolReg != nil {
toolReg.RegisterRunningTool(convID, monitorExecID)
}
}
toolRunReg := mcp.ToolRunRegistryFromContext(ctx)
execCtx, execCancel := context.WithCancel(ctx)
var timeoutCancel context.CancelFunc
if w.toolTimeoutMinutes > 0 {
execCtx, execCancel = context.WithTimeout(ctx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
}
if execReg != nil && convID != "" {
execReg.RegisterActiveEinoExecute(convID, execCancel)
}
sr, err := w.inner.ExecuteStreaming(execCtx, &req)
if err != nil {
if timeoutCancel != nil {
timeoutCancel()
}
if execCancel != nil {
execCancel()
}
if w.recordMonitor != nil {
w.recordMonitor(tid, userCmd, "", false, err)
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
if w.finishMonitor != nil {
w.finishMonitor(monitorExecID, tid, userCmd, hint, false, context.DeadlineExceeded)
}
if w.invokeNotify != nil && tid != "" {
w.invokeNotify.Fire(tid, "execute", agentTag, false, hint, context.DeadlineExceeded)
}
return schema.StreamReaderFromArray([]*filesystem.ExecuteResponse{{Output: hint}}), nil
}
if w.finishMonitor != nil {
w.finishMonitor(monitorExecID, tid, userCmd, "", false, err)
}
if w.invokeNotify != nil && tid != "" {
w.invokeNotify.Fire(tid, "execute", agentTag, false, "", err)
}
return nil, err
}
if sr == nil || w.invokeNotify == nil || tid == "" {
if sr == nil {
if timeoutCancel != nil {
timeoutCancel()
}
if execCancel != nil {
execCancel()
}
@@ -100,11 +146,35 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, tctx context.Context) {
defer inner.Close()
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int) {
var innerCloseOnce sync.Once
closeInner := func() {
innerCloseOnce.Do(func() { inner.Close() })
}
defer closeInner()
if timeoutCleanup != nil {
defer timeoutCleanup()
}
if cancel != nil {
defer cancel()
}
if reg != nil && conversationID != "" {
defer reg.UnregisterActiveEinoExecute(conversationID)
}
if toolReg != nil && conversationID != "" && execID != "" {
defer toolReg.UnregisterRunningTool(conversationID, execID)
}
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
stopWatch := make(chan struct{})
go func() {
select {
case <-tctx.Done():
closeInner()
case <-stopWatch:
}
}()
defer close(stopWatch)
var sb strings.Builder
success := true
@@ -112,41 +182,103 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
exitCode := 0
hasExitCode := false
idleWatch := security.NewShellInactivityWatch(noOutputSec)
if idleWatch != nil {
defer idleWatch.Stop()
}
type execRecvMsg struct {
resp *filesystem.ExecuteResponse
err error
}
recvCh := make(chan execRecvMsg, 1)
go func() {
for {
resp, rerr := inner.Recv()
recvCh <- execRecvMsg{resp: resp, err: rerr}
if rerr != nil {
return
}
}
}()
fireInactivityTimeout := func() {
success = false
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
sb.WriteString(msg)
if w.outputChunk != nil && toolCallID != "" {
w.outputChunk("execute", toolCallID, msg)
}
if cancel != nil {
cancel()
}
closeInner()
}
recvLoop:
for {
resp, rerr := inner.Recv()
if errors.Is(rerr, io.EOF) {
break
var idleCh <-chan struct{}
if idleWatch != nil {
idleCh = idleWatch.Expired
}
if rerr != nil {
success = false
invokeErr = rerr
_ = outW.Send(nil, rerr)
break
}
if resp != nil {
if resp.ExitCode != nil {
hasExitCode = true
exitCode = *resp.ExitCode
select {
case <-idleCh:
fireInactivityTimeout()
break recvLoop
case msg := <-recvCh:
rerr := msg.err
resp := msg.resp
if errors.Is(rerr, io.EOF) {
break recvLoop
}
var appended string
if resp.Output != "" {
sb.WriteString(resp.Output)
appended = resp.Output
}
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
w.outputChunk("execute", tid, appended)
}
if outW.Send(resp, nil) {
if rerr != nil {
success = false
invokeErr = fmt.Errorf("execute stream closed by consumer")
break
invokeErr = rerr
if einoExecuteRecvErrIsToolTimeout(rerr, tctx) {
invokeErr = context.DeadlineExceeded
break recvLoop
}
if errors.Is(rerr, context.Canceled) || (tctx != nil && errors.Is(tctx.Err(), context.Canceled)) {
invokeErr = context.Canceled
break recvLoop
}
_ = outW.Send(nil, rerr)
break recvLoop
}
if resp != nil {
if resp.ExitCode != nil {
hasExitCode = true
exitCode = *resp.ExitCode
continue
}
var appended string
if resp.Output != "" {
if security.IsLegacyShellExitNoise(resp.Output) {
continue
}
if idleWatch != nil {
idleWatch.Bump()
}
sb.WriteString(resp.Output)
appended = resp.Output
}
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
w.outputChunk("execute", toolCallID, appended)
}
if outW.Send(resp, nil) {
success = false
invokeErr = fmt.Errorf("execute stream closed by consumer")
break recvLoop
}
}
}
}
if success && hasExitCode && exitCode != 0 {
success = false
invokeErr = fmt.Errorf("execute exited with code %d", exitCode)
invokeErr = &ExecuteExitError{Code: exitCode}
}
// WithTimeout 触发后,子进程常被信号结束,local 侧多报 exit -1 / canceled,错误链里不一定带 DeadlineExceeded。
// 用执行所用 ctx 归一化,便于 UI 展示「超时」而非含糊的 -1。
@@ -154,6 +286,21 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
success = false
invokeErr = context.DeadlineExceeded
}
// 用户「中断并继续」终止 execute:合并说明进工具结果(与 MCP CancelToolExecutionWithNote 一致)。
partialStreamed := sb.String()
var abortNote string
if reg != nil && conversationID != "" && (invokeErr != nil || errors.Is(tctx.Err(), context.Canceled)) {
if note := reg.TakeEinoExecuteAbortNote(conversationID); note != "" {
abortNote = note
merged := mcp.MergePartialToolOutputAndAbortNote(partialStreamed, note)
sb.Reset()
sb.WriteString(merged)
if invokeErr == nil {
success = false
invokeErr = context.Canceled
}
}
}
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
@@ -163,12 +310,32 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
}
sb.WriteString(hint)
}
if w.recordMonitor != nil {
w.recordMonitor(tid, command, sb.String(), success, invokeErr)
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
if partialStreamed != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
} else if text := strings.TrimSpace(sb.String()); text != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
}
}
rawOutput := sb.String()
fireBody := rawOutput
if !success && hasExitCode && exitCode != 0 {
statusLine := security.ExecuteFailureStatusLine(exitCode)
if !strings.Contains(rawOutput, "命令执行失败:") {
_ = outW.Send(&filesystem.ExecuteResponse{Output: statusLine}, nil)
sb.WriteString(statusLine)
}
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
}
if w.finishMonitor != nil {
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
}
if w.invokeNotify != nil {
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
}
w.invokeNotify.Fire(tid, "execute", agentTag, success, sb.String(), invokeErr)
outW.Close()
}(sr, userCmd, execCancel, execCtx)
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
return outR, nil
}
@@ -0,0 +1,362 @@
package multiagent
import (
"context"
"errors"
"io"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/mcp"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/schema"
)
type mockStreamingShell struct {
immediateErr error
recvErr error
output string
called bool
lastCommand string
}
func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
m.called = true
if input != nil {
m.lastCommand = input.Command
}
if m.immediateErr != nil {
return nil, m.immediateErr
}
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
go func() {
defer outW.Close()
if strings.TrimSpace(m.output) != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: m.output}, nil)
}
if m.recvErr != nil {
_ = outW.Send(nil, m.recvErr)
}
}()
return outR, nil
}
func TestEinoStreamingShellWrap_PreparesNonInteractiveCommand(t *testing.T) {
inner := &mockStreamingShell{output: "ok\n"}
wrap := &einoStreamingShellWrap{inner: inner}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "echo ok"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
for {
_, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
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") {
t.Fatalf("missing python unbuffer in inner command: %q", inner.lastCommand)
}
}
func TestEinoStreamingShellWrap_NoOutputTimeout(t *testing.T) {
inner := &mockStreamingShellHanging{}
notify := einomcp.NewToolInvokeNotifyHolder()
var fired string
notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
fired = content
})
wrap := &einoStreamingShellWrap{
inner: inner,
invokeNotify: notify,
shellNoOutputTimeoutSec: 1,
}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil {
got.WriteString(resp.Output)
}
}
if !inner.called {
t.Fatal("inner shell should run (no command blacklist)")
}
out := got.String()
if !strings.Contains(out, "没有新的输出") && !strings.Contains(out, "no new output") {
t.Fatalf("expected inactivity timeout message, got: %q notify=%q", out, fired)
}
}
type mockStreamingShellPartialThenHang struct {
called bool
}
func (m *mockStreamingShellPartialThenHang) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
m.called = true
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
go func() {
_ = outW.Send(&filesystem.ExecuteResponse{Output: "[sudo] password:\n"}, nil)
<-ctx.Done()
outW.Close()
}()
return outR, nil
}
func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) {
inner := &mockStreamingShellPartialThenHang{}
wrap := &einoStreamingShellWrap{
inner: inner,
shellNoOutputTimeoutSec: 1,
}
start := time.Now()
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "sudo whoami"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil {
got.WriteString(resp.Output)
}
}
if time.Since(start) > 5*time.Second {
t.Fatalf("expected inactivity timeout ~1s, took %v", time.Since(start))
}
if !strings.Contains(got.String(), "没有新的输出") && !strings.Contains(got.String(), "no new output") {
t.Fatalf("expected inactivity message, got: %q", got.String())
}
}
type mockStreamingShellHanging struct {
called bool
}
func (m *mockStreamingShellHanging) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
m.called = true
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](4)
go func() {
<-ctx.Done()
outW.Close()
}()
return outR, nil
}
func TestEinoExecuteRecvErrIsToolTimeout(t *testing.T) {
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
time.Sleep(2 * time.Millisecond)
<-tctx.Done()
if !einoExecuteRecvErrIsToolTimeout(context.Canceled, tctx) {
t.Fatal("expected canceled recv with deadline exec ctx to count as tool timeout")
}
if !einoExecuteRecvErrIsToolTimeout(context.DeadlineExceeded, nil) {
t.Fatal("expected DeadlineExceeded recv without tctx")
}
if einoExecuteRecvErrIsToolTimeout(errors.New("exit status 1"), context.Background()) {
t.Fatal("unexpected timeout for generic error")
}
}
func TestEinoStreamingShellWrap_ToolTimeoutImmediateErrIsSoft(t *testing.T) {
inner := &mockStreamingShell{immediateErr: context.DeadlineExceeded}
wrap := &einoStreamingShellWrap{
inner: inner,
toolTimeoutMinutes: 60,
}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"})
if err != nil {
t.Fatalf("immediate tool timeout must return soft stream, got err: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("outer stream must not hard-fail, got: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) {
t.Fatalf("expected timeout hint, got: %q", got.String())
}
}
func TestEinoStreamingShellWrap_ToolTimeoutRecvErrIsSoft(t *testing.T) {
inner := &mockStreamingShell{recvErr: context.DeadlineExceeded}
notify := einomcp.NewToolInvokeNotifyHolder()
wrap := &einoStreamingShellWrap{
inner: inner,
invokeNotify: notify,
toolTimeoutMinutes: 60,
}
// 生产路径由 Eino compose 注入 toolCallID;单测通过已过期 execCtx 识别 tool_timeout 软错误。
tctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
time.Sleep(2 * time.Millisecond)
<-tctx.Done()
sr, err := wrap.ExecuteStreaming(tctx, &filesystem.ExecuteRequest{Command: "sleep 999"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("outer stream must not hard-fail on tool timeout, got: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if !strings.Contains(got.String(), einoExecuteTimeoutUserHint()) {
t.Fatalf("expected timeout hint in stream, got: %q", got.String())
}
}
func TestEinoStreamingShellWrap_CapturesOutputWithToolTimeout(t *testing.T) {
inner := &mockStreamingShell{output: "100\n"}
notify := einomcp.NewToolInvokeNotifyHolder()
var firedContent string
notify.Set(func(toolCallID, toolName, einoAgent string, success bool, content string, invokeErr error) {
firedContent = content
})
wrap := &einoStreamingShellWrap{
inner: inner,
invokeNotify: notify,
toolTimeoutMinutes: 60,
}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "echo 100"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("unexpected stream error: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if !strings.Contains(got.String(), "100") {
t.Fatalf("stream output = %q, want contains 100", got.String())
}
if !strings.Contains(firedContent, "100") {
t.Fatalf("notify content = %q, want contains 100", firedContent)
}
}
func TestEinoStreamingShellWrap_AbortNoteDoesNotDuplicateStreamedOutput(t *testing.T) {
inner := &mockStreamingShell{output: "line1\nline2\n", recvErr: context.Canceled}
notify := einomcp.NewToolInvokeNotifyHolder()
wrap := &einoStreamingShellWrap{
inner: inner,
invokeNotify: notify,
}
reg := &abortNoteTestRegistry{note: "改成20次"}
ctx := mcp.WithEinoExecuteRunRegistry(
mcp.WithMCPConversationID(context.Background(), "conv-abort-dup"),
reg,
)
sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "ping -c 10 baidu.com"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("unexpected stream error: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
out := got.String()
if strings.Count(out, "line1") != 1 || strings.Count(out, "line2") != 1 {
t.Fatalf("stream duplicated stdout: %q", out)
}
if !strings.Contains(out, "改成20次") {
t.Fatalf("stream missing abort note: %q", out)
}
}
type abortNoteTestRegistry struct {
note string
}
func (r *abortNoteTestRegistry) RegisterActiveEinoExecute(string, context.CancelFunc) {}
func (r *abortNoteTestRegistry) UnregisterActiveEinoExecute(string) {}
func (r *abortNoteTestRegistry) AbortActiveEinoExecute(string, string) bool { return false }
func (r *abortNoteTestRegistry) TakeEinoExecuteAbortNote(string) string { return r.note }
func TestEinoStreamingShellWrap_NonTimeoutRecvErrStillHard(t *testing.T) {
inner := &mockStreamingShell{recvErr: errors.New("broken pipe")}
wrap := &einoStreamingShellWrap{inner: inner}
sr, err := wrap.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: "true"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
_, rerr := sr.Recv()
if rerr == nil || errors.Is(rerr, io.EOF) {
t.Fatal("expected hard stream error for non-timeout failure")
}
}
@@ -63,10 +63,43 @@ func toolCallArgsFromAccumulated(msgs []adk.Message, toolCallID, expectToolName
return map[string]interface{}{}
}
// beginEinoADKFilesystemToolMonitor 在 Eino ADK filesystem 工具开始调用时写入 running 状态。
func beginEinoADKFilesystemToolMonitor(
ag *agent.Agent,
rec einomcp.ExecutionRecorder,
binder *MCPExecutionBinder,
toolCallID, toolName string,
) {
if ag == nil || rec == nil {
return
}
name := strings.TrimSpace(toolName)
if name == "" || strings.EqualFold(name, "execute") {
return
}
if !isBuiltinEinoADKFilesystemToolName(name) {
return
}
tid := strings.TrimSpace(toolCallID)
if tid == "" {
return
}
storedName := "eino_fs::" + strings.ToLower(name)
id := ag.BeginLocalToolExecution(storedName, map[string]interface{}{})
if id == "" {
return
}
rec(id, tid)
if binder != nil {
binder.Bind(tid, id)
}
}
// recordEinoADKFilesystemToolMonitor 将 Eino ADK filesystem 中间件工具结果写入 MCP 监控(与 execute / MCP 桥芯片一致)。
func recordEinoADKFilesystemToolMonitor(
ag *agent.Agent,
rec einomcp.ExecutionRecorder,
binder *MCPExecutionBinder,
toolName string,
toolCallID string,
msgs []adk.Message,
@@ -94,8 +127,12 @@ func recordEinoADKFilesystemToolMonitor(
invErr = errors.New(t)
}
}
id := ag.RecordLocalToolExecution(storedName, args, resultText, invErr)
if id != "" {
execID := ""
if binder != nil {
execID = binder.ExecutionID(toolCallID)
}
id := ag.FinishLocalToolExecution(execID, storedName, args, resultText, invErr)
if id != "" && execID == "" {
rec(id, toolCallID)
}
}
+3 -6
View File
@@ -243,17 +243,14 @@ func prependEinoMiddlewares(
return outTools, extraHandlers, toolSearchActive, nil
}
func deepExtrasFromConfig(ma *config.MultiAgentConfig) (outputKey string, retry *adk.ModelRetryConfig, taskDesc func(context.Context, []adk.Agent) (string, error)) {
func deepExtrasFromConfig(ma *config.MultiAgentConfig) (outputKey string, taskDesc func(context.Context, []adk.Agent) (string, error)) {
if ma == nil {
return "", nil, nil
return "", nil
}
mw := ma.EinoMiddleware
if k := strings.TrimSpace(mw.DeepOutputKey); k != "" {
outputKey = k
}
if mw.DeepModelRetryMaxRetries > 0 {
retry = &adk.ModelRetryConfig{MaxRetries: mw.DeepModelRetryMaxRetries}
}
prefix := strings.TrimSpace(mw.TaskToolDescriptionPrefix)
if prefix != "" {
taskDesc = func(ctx context.Context, agents []adk.Agent) (string, error) {
@@ -274,5 +271,5 @@ func deepExtrasFromConfig(ma *config.MultiAgentConfig) (outputKey string, retry
return prefix + "\n可用子代理(按名称 transfer / task 调用):" + strings.Join(names, "、"), nil
}
}
return outputKey, retry, taskDesc
return outputKey, taskDesc
}
+14 -15
View File
@@ -7,6 +7,7 @@ import (
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/adk"
@@ -29,7 +30,9 @@ type PlanExecuteRootArgs struct {
MwCfg *config.MultiAgentEinoMiddlewareConfig
// ConversationID is used for transcript/isolation paths in middleware.
ConversationID string
Logger *zap.Logger
DB *database.DB
ProjectID string
Logger *zap.Logger
// ModelName is used for model input token estimation logs.
ModelName string
// ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件(patchtoolcalls, reduction, toolsearch, plantask),
@@ -91,24 +94,20 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
if a.SkillMiddleware != nil {
execHandlers = append(execHandlers, a.SkillMiddleware)
}
// 4. summarization(最后,与 Deep/Supervisor 一致
// 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.Logger)
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 = append(execHandlers, sumMw)
}
// 5. 孤儿 tool 消息兜底:必须挂在所有改写历史中间件(summarization/reduction/skill)之后、
// telemetry 之前,保证送入 ChatModel 的消息序列 tool_call ↔ tool_result 配对完整。
execHandlers = append(execHandlers, newOrphanToolPrunerMiddleware(a.Logger, "plan_execute_executor"))
if teleMw := newEinoModelInputTelemetryMiddleware(a.Logger, a.ModelName, a.ConversationID, "plan_execute_executor"); teleMw != nil {
execHandlers = append(execHandlers, teleMw)
}
if a.ModelFacingTrace != nil {
if capMw := newModelFacingTraceMiddleware(a.ModelFacingTrace); capMw != nil {
execHandlers = append(execHandlers, capMw)
}
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{
Model: a.ExecModel,
+14 -14
View File
@@ -11,6 +11,7 @@ import (
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project"
@@ -32,6 +33,7 @@ func RunEinoSingleChatModelAgent(
appCfg *config.Config,
ma *config.MultiAgentConfig,
ag *agent.Agent,
db *database.DB,
logger *zap.Logger,
conversationID string,
projectID string,
@@ -79,7 +81,7 @@ func RunEinoSingleChatModelAgent(
}
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder)
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ag, recorder)
mainDefs := ag.ToolsForRole(roleTools)
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
if err != nil {
@@ -121,7 +123,7 @@ func RunEinoSingleChatModelAgent(
return nil, fmt.Errorf("eino single 模型: %w", err)
}
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil {
return nil, fmt.Errorf("eino single summarization: %w", err)
}
@@ -134,7 +136,7 @@ func RunEinoSingleChatModelAgent(
}
if einoSkillMW != nil {
if einoFSTools && einoLoc != nil {
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if fsErr != nil {
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
}
@@ -142,13 +144,14 @@ func RunEinoSingleChatModelAgent(
}
handlers = append(handlers, einoSkillMW)
}
handlers = append(handlers, mainSumMw)
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "eino_single"); teleMw != nil {
handlers = append(handlers, teleMw)
}
if capMw := newModelFacingTraceMiddleware(modelFacingTrace); capMw != nil {
handlers = append(handlers, capMw)
}
handlers = appendEinoChatModelTailMiddlewares(handlers, einoChatModelTailConfig{
logger: logger,
phase: "eino_single",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
conversationID: conversationID,
trace: modelFacingTrace,
})
maxIter := agentMaxIterations(appCfg)
@@ -186,13 +189,10 @@ func RunEinoSingleChatModelAgent(
MaxIterations: maxIter,
Handlers: handlers,
}
outKey, modelRetry, _ := deepExtrasFromConfig(ma)
outKey, _ := deepExtrasFromConfig(ma)
if outKey != "" {
chatCfg.OutputKey = outKey
}
if modelRetry != nil {
chatCfg.ModelRetryConfig = modelRetry
}
chatAgent, err := adk.NewChatModelAgent(ctx, chatCfg)
if err != nil {
+27 -7
View File
@@ -9,6 +9,7 @@ import (
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/security"
localbk "github.com/cloudwego/eino-ext/adk/backend/local"
"github.com/cloudwego/eino/adk"
@@ -81,8 +82,10 @@ func subAgentFilesystemMiddleware(
loc *localbk.Local,
invokeNotify *einomcp.ToolInvokeNotifyHolder,
einoAgentName string,
recordMonitor func(toolCallID, command, stdout string, success bool, invokeErr error),
beginMonitor func(toolCallID, command string) string,
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
toolTimeoutMinutes int,
shellNoOutputTimeoutSec int,
outputChunk func(toolName, toolCallID, chunk string),
) (adk.ChatModelAgentMiddleware, error) {
if loc == nil {
@@ -91,12 +94,14 @@ func subAgentFilesystemMiddleware(
return filesystem.New(ctx, &filesystem.MiddlewareConfig{
Backend: loc,
StreamingShell: &einoStreamingShellWrap{
inner: loc,
invokeNotify: invokeNotify,
einoAgentName: strings.TrimSpace(einoAgentName),
outputChunk: outputChunk,
recordMonitor: recordMonitor,
toolTimeoutMinutes: toolTimeoutMinutes,
inner: security.NewEinoStreamingShell(),
invokeNotify: invokeNotify,
einoAgentName: strings.TrimSpace(einoAgentName),
outputChunk: outputChunk,
beginMonitor: beginMonitor,
finishMonitor: finishMonitor,
toolTimeoutMinutes: toolTimeoutMinutes,
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
},
})
}
@@ -108,3 +113,18 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
}
return cfg.Agent.ToolTimeoutMinutes
}
// agentShellNoOutputTimeoutSeconds0=默认 300s5 分钟);-1=关闭;>0=自定义秒数。
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
if cfg == nil {
return 300
}
v := cfg.Agent.ShellNoOutputTimeoutSeconds
if v < 0 {
return 0
}
if v == 0 {
return 300
}
return v
}
+70 -16
View File
@@ -9,7 +9,9 @@ import (
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
copenai "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project"
"github.com/bytedance/sonic"
"github.com/cloudwego/eino/adk"
@@ -20,8 +22,6 @@ import (
"go.uber.org/zap"
)
const defaultSummarizationRetryMax = 3
// einoSummarizeUserInstruction:压缩历史时保留渗透测试关键信息。
const einoSummarizeUserInstruction = `在保持所有关键安全测试信息完整的前提下压缩对话历史
@@ -40,6 +40,8 @@ func newEinoSummarizationMiddleware(
appCfg *config.Config,
mwCfg *config.MultiAgentEinoMiddlewareConfig,
conversationID string,
db *database.DB,
projectID string,
logger *zap.Logger,
) (adk.ChatModelAgentMiddleware, error) {
if summaryModel == nil || appCfg == nil {
@@ -93,10 +95,8 @@ func newEinoSummarizationMiddleware(
}
}
retryMax := defaultSummarizationRetryMax
if mwCfg != nil && mwCfg.SummarizationRetryMaxAttempts > 0 {
retryMax = mwCfg.SummarizationRetryMaxAttempts
}
retryPolicy := einoTransientRunRetryPolicyFromMW(mwCfg)
retryMax := retryPolicy.maxAttempts
// ModelOptions apply only to summarization Generate (same ChatModel instance as the agent).
// Strip thinking/reasoning on this call path; mark requests for empty-choices diagnostics.
@@ -133,17 +133,25 @@ func newEinoSummarizationMiddleware(
Retry: &summarization.RetryConfig{
MaxRetries: &retryMax,
ShouldRetry: func(_ context.Context, _ adk.Message, err error) bool {
if err != nil && logger != nil {
logger.Warn("eino summarization generate attempt failed, will retry if attempts remain",
retry := isEinoTransientRunError(err)
if retry && logger != nil {
logger.Warn("eino summarization generate transient error, will retry if attempts remain",
zap.Error(err),
zap.Int("max_retries", retryMax),
)
}
return err != nil
return retry
},
},
Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) {
return summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax)
out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax)
if ferr != nil {
return nil, ferr
}
if appCfg != nil {
out = refreshFactIndexInMessages(out, db, projectID, appCfg.Project, logger)
}
return out, nil
},
Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error {
if transcriptPath != "" && len(before.Messages) > 0 {
@@ -176,6 +184,50 @@ func newEinoSummarizationMiddleware(
return mw, nil
}
// refreshFactIndexInMessages 在 summarization 压缩后,用 DB 最新索引替换 system 中已有的项目黑板索引段。
func refreshFactIndexInMessages(msgs []adk.Message, db *database.DB, projectID string, cfg config.ProjectConfig, logger *zap.Logger) []adk.Message {
if db == nil || !cfg.Enabled {
return msgs
}
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return msgs
}
freshIndex, err := project.BuildFactIndexBlock(db, projectID, cfg)
if err != nil {
if logger != nil {
logger.Warn("summarization: 刷新项目黑板索引失败", zap.String("projectId", projectID), zap.Error(err))
}
return msgs
}
freshIndex = strings.TrimSpace(freshIndex)
if freshIndex == "" {
return msgs
}
changed := false
out := make([]adk.Message, len(msgs))
for i, msg := range msgs {
if msg == nil || msg.Role != schema.System {
out[i] = msg
continue
}
newContent, ok := project.ReplaceFactIndexSection(msg.Content, freshIndex)
if !ok {
out[i] = msg
continue
}
cloned := *msg
cloned.Content = newContent
out[i] = &cloned
changed = true
}
if changed && logger != nil {
logger.Info("summarization: 已刷新项目黑板索引", zap.String("projectId", projectID))
}
return out
}
// summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹,避免压缩后执行链断裂。
//
// 关键不变量:tool_call ↔ tool_result 的 pair 必须整体保留或整体丢弃。
@@ -205,17 +257,19 @@ func summarizeFinalizeWithRecentAssistantToolTrail(
nonSystem = append(nonSystem, msg)
}
mergedSystem := mergeCollectedSystemMessages(systemMsgs)
if recentTrailTokenBudget <= 0 || len(nonSystem) == 0 {
out := make([]adk.Message, 0, len(systemMsgs)+1)
out = append(out, systemMsgs...)
out := make([]adk.Message, 0, len(mergedSystem)+1)
out = append(out, mergedSystem...)
out = append(out, summary)
return out, nil
}
rounds := splitMessagesIntoRounds(nonSystem)
if len(rounds) == 0 {
out := make([]adk.Message, 0, len(systemMsgs)+1)
out = append(out, systemMsgs...)
out := make([]adk.Message, 0, len(mergedSystem)+1)
out = append(out, mergedSystem...)
out = append(out, summary)
return out, nil
}
@@ -267,8 +321,8 @@ func summarizeFinalizeWithRecentAssistantToolTrail(
selectedMsgs = append(selectedMsgs, selectedRoundsReverse[i].messages...)
}
out := make([]adk.Message, 0, len(systemMsgs)+1+len(selectedMsgs))
out = append(out, systemMsgs...)
out := make([]adk.Message, 0, len(mergedSystem)+1+len(selectedMsgs))
out = append(out, mergedSystem...)
out = append(out, summary)
out = append(out, selectedMsgs...)
return out, nil
+71 -7
View File
@@ -7,9 +7,14 @@ import (
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/project"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// fixedTokenCounter 让 tool 消息按 tokensPerToolMessage 计,其它消息按 1 计。
@@ -187,8 +192,8 @@ func TestSummarizeFinalize_KeepsToolRoundIntact(t *testing.T) {
if len(out) < 2 {
t.Fatalf("output too short: %d", len(out))
}
if out[0] != sys {
t.Fatalf("first message must be system")
if out[0].Role != schema.System || out[0].Content != "sys" {
t.Fatalf("first message must be system sys, got %s: %q", out[0].Role, out[0].Content)
}
if out[1] != summary {
t.Fatalf("second message must be summary")
@@ -288,12 +293,12 @@ func TestSummarizeFinalize_BudgetZeroFallsBackToSummaryOnly(t *testing.T) {
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(out) != 2 || out[0] != sys || out[1] != summary {
if len(out) != 2 || out[0].Role != schema.System || out[0].Content != "sys" || out[1] != summary {
t.Fatalf("budget=0 must yield [system, summary] only, got %+v", out)
}
}
func TestSummarizeFinalize_PreservesAllSystemMessages(t *testing.T) {
func TestSummarizeFinalize_MergesSystemMessages(t *testing.T) {
sys1 := schema.SystemMessage("sys1")
sys2 := schema.SystemMessage("sys2")
summary := schema.AssistantMessage("s", nil)
@@ -316,10 +321,13 @@ func TestSummarizeFinalize_PreservesAllSystemMessages(t *testing.T) {
for _, m := range out {
if m != nil && m.Role == schema.System {
systemCount++
if got := m.Content; got != "sys1\n\nsys2" {
t.Fatalf("unexpected merged system content: %q", got)
}
}
}
if systemCount != 2 {
t.Fatalf("want 2 system messages retained, got %d", systemCount)
if systemCount != 1 {
t.Fatalf("want 1 merged system message, got %d", systemCount)
}
}
@@ -373,6 +381,12 @@ func TestWriteSummarizationTranscript(t *testing.T) {
if !strings.Contains(text, "tool_calls:") || !strings.Contains(text, "nmap output") {
t.Fatalf("missing tool round: %q", text)
}
if !strings.Contains(text, `"name":"stub_tool"`) || !strings.Contains(text, `"arguments":"{}"`) {
t.Fatalf("missing tool name/arguments: %q", text)
}
if strings.Contains(text, "tool_call_id") || strings.Contains(text, `"id":"tc1"`) {
t.Fatalf("transcript should omit tool_call_id: %q", text)
}
}
func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
@@ -389,9 +403,11 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
"你是CyberStrikeAI,是一个专业的网络安全渗透测试专家。",
"高强度扫描要求:全力出击",
"",
project.FactIndexSectionStartMarker,
"## 项目黑板索引(project: 123, id: abc",
"(暂无事实)",
"需要写入请使用 upsert_project_fact。",
project.FactIndexSectionEndMarker,
"",
"# Skills System",
"**How to Use Skills**",
@@ -419,7 +435,7 @@ func TestSanitizeSystemContentForTranscript_BestPractice(t *testing.T) {
func TestFormatSummarizationTranscript_OmitsBloatedSystem(t *testing.T) {
t.Parallel()
msgs := []adk.Message{
schema.SystemMessage("以下是当前会话绑定的工具名称索引\n- nmap\n\n你是CyberStrikeAI\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n# Skills System\nboiler"),
schema.SystemMessage("以下是当前会话绑定的工具名称索引\n- nmap\n\n你是CyberStrikeAI\n" + project.FactIndexSectionStartMarker + "\n## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + project.FactIndexSectionEndMarker + "\n# Skills System\nboiler"),
schema.UserMessage("hello"),
schema.AssistantMessage("reply", nil),
}
@@ -434,3 +450,51 @@ func TestFormatSummarizationTranscript_OmitsBloatedSystem(t *testing.T) {
t.Fatalf("dynamic blackboard missing: %q", out)
}
}
func TestRefreshFactIndexInMessages(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "summarize-facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "summarize-proj"})
if err != nil {
t.Fatal(err)
}
cfg := config.ProjectConfig{Enabled: true}
oldIndex, err := project.BuildFactIndexBlock(db, proj.ID, cfg)
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "target/host",
Category: "target",
Summary: "fresh host fact",
})
if err != nil {
t.Fatal(err)
}
msgs := []adk.Message{
schema.SystemMessage("instruction\n\n" + oldIndex),
schema.UserMessage("hi"),
}
out := refreshFactIndexInMessages(msgs, db, proj.ID, cfg, nil)
sys := out[0].Content
if strings.Contains(sys, "(暂无事实)") {
t.Fatalf("expected refreshed index, got: %q", sys)
}
if !strings.Contains(sys, "fresh host fact") {
t.Fatalf("expected new fact in index: %q", sys)
}
if !strings.Contains(sys, "instruction") {
t.Fatalf("non-index system content should be preserved: %q", sys)
}
}
@@ -6,6 +6,8 @@ import (
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
"cyberstrike-ai/internal/project"
"github.com/bytedance/sonic"
)
@@ -19,9 +21,13 @@ const (
transcriptToolIndexStartMarker = "以下是当前会话绑定的工具名称索引"
transcriptPersonaStartMarker = "你是CyberStrikeAI"
transcriptSkillsSystemMarker = "# Skills System"
transcriptProjectBlackboardMarker = "## 项目黑板索引"
)
type transcriptToolCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
// formatSummarizationTranscript renders pre-compaction messages for transcript.txt.
// Best practice: keep full user/assistant/tool turns; slim system to dynamic blocks only.
func formatSummarizationTranscript(msgs []adk.Message) string {
@@ -88,11 +94,17 @@ func stripSkillsSystemBoilerplate(s string) string {
}
func extractProjectBlackboardSection(s string) string {
idx := strings.Index(s, transcriptProjectBlackboardMarker)
if idx < 0 {
start := strings.Index(s, project.FactIndexSectionStartMarker)
if start < 0 {
return ""
}
return strings.TrimSpace(s[idx:])
section := s[start:]
end := strings.Index(section, project.FactIndexSectionEndMarker)
if end < 0 {
return ""
}
section = section[:end+len(project.FactIndexSectionEndMarker)]
return strings.TrimSpace(section)
}
func appendTranscriptSection(sb *strings.Builder, role schema.RoleType, body string) {
@@ -131,15 +143,21 @@ func appendTranscriptMessage(sb *strings.Builder, msg adk.Message) {
}
}
if len(msg.ToolCalls) > 0 {
if b, err := sonic.Marshal(msg.ToolCalls); err == nil {
if b, err := sonic.Marshal(formatTranscriptToolCalls(msg.ToolCalls)); err == nil {
sb.WriteString("tool_calls: ")
sb.Write(b)
sb.WriteByte('\n')
}
}
if msg.ToolCallID != "" {
sb.WriteString("tool_call_id: ")
sb.WriteString(msg.ToolCallID)
sb.WriteByte('\n')
}
}
func formatTranscriptToolCalls(calls []schema.ToolCall) []transcriptToolCall {
out := make([]transcriptToolCall, 0, len(calls))
for _, tc := range calls {
out = append(out, transcriptToolCall{
Name: tc.Function.Name,
Arguments: tc.Function.Arguments,
})
}
return out
}
@@ -46,6 +46,10 @@ func injectToolNamesOnlyInstruction(ctx context.Context, instruction string, too
sb.WriteString("2) 调用具体工具前,请先确认该工具的参数要求(以当前请求中的工具定义为准);不确定时先澄清再调用。\n")
sb.WriteString("3) 不要臆造不存在的工具名。\n\n")
}
if s := strings.TrimSpace(injectShellToolGuidance("", names)); s != "" {
sb.WriteString(s)
sb.WriteString("\n\n")
}
if s := strings.TrimSpace(instruction); s != "" {
sb.WriteString(s)
}
+74 -14
View File
@@ -3,6 +3,7 @@ package multiagent
import (
"context"
"errors"
"fmt"
"strings"
"time"
@@ -17,8 +18,9 @@ const (
defaultEinoRunRetryMaxBackoff = 30 * time.Second
)
// isEinoTransientRunError 判断 ADK 运行期错误是否适合指数退避续跑(429、5xx、网络抖动等)
// 用户取消、超时、迭代上限等由 run loop 单独处理,不在此列
// isEinoTransientRunError 是 Eino 运行期「可退避重试 vs 直接失败」的唯一判据
// 429/5xx/网络抖动等返回 true;用户取消、超时、迭代上限、鉴权失败等返回 false
// 其它模块(run loop、summarization 等)只调用本函数,不在别处维护平行规则。
func isEinoTransientRunError(err error) bool {
if err == nil {
return false
@@ -60,6 +62,7 @@ func isEinoTransientRunError(err error) bool {
"dial tcp",
"tls handshake timeout",
"stream error",
"goaway", // http2: server sent GOAWAY and closed the connection
"unexpected eof",
`": eof`, // net/http: Post "url": EOF (often wraps io.EOF)
"unexpected end of json",
@@ -78,6 +81,71 @@ func isEinoTransientRunError(err error) bool {
return false
}
type einoTransientRunRetryPolicy struct {
maxAttempts int
maxBackoff time.Duration
}
func einoTransientRunRetryPolicyFromArgs(args *einoADKRunLoopArgs) einoTransientRunRetryPolicy {
return einoTransientRunRetryPolicy{
maxAttempts: einoRunRetryMaxAttempts(args),
maxBackoff: einoRunRetryMaxBackoff(args),
}
}
func einoTransientRunRetryPolicyFromMW(mw *config.MultiAgentEinoMiddlewareConfig) einoTransientRunRetryPolicy {
maxBackoff := defaultEinoRunRetryMaxBackoff
if mw != nil && mw.RunRetryMaxBackoffSec > 0 {
maxBackoff = time.Duration(mw.RunRetryMaxBackoffSec) * time.Second
}
return einoTransientRunRetryPolicy{
maxAttempts: RunRetryMaxAttemptsFromConfig(mw),
maxBackoff: maxBackoff,
}
}
// einoTransientRunRetrier 在 run loop 内对临时错误做指数退避并重启 Runner(唯一重试执行层)。
type einoTransientRunRetrier struct {
policy einoTransientRunRetryPolicy
attempts int
}
func newEinoTransientRunRetrier(policy einoTransientRunRetryPolicy) *einoTransientRunRetrier {
return &einoTransientRunRetrier{policy: policy}
}
// tryRetry 对临时错误退避后返回重启消息;次数用尽返回 exhausted 错误。
func (r *einoTransientRunRetrier) tryRetry(
ctx context.Context,
runErr error,
args *einoADKRunLoopArgs,
baseMsgs, accumulated []adk.Message,
baseCount int,
) (restarted bool, restartMsgs []adk.Message, ctxSource einoRunRestartContextSource, backoff time.Duration, fatal error) {
if runErr == nil || !isEinoTransientRunError(runErr) {
return false, nil, "", 0, runErr
}
r.attempts++
if r.attempts > r.policy.maxAttempts {
return false, nil, "", 0, fmt.Errorf("transient retry exhausted after %d attempts: %w", r.policy.maxAttempts, runErr)
}
backoff = einoTransientRetryBackoff(r.attempts-1, r.policy.maxBackoff)
select {
case <-ctx.Done():
return false, nil, "", 0, ctx.Err()
case <-time.After(backoff):
}
restartMsgs, ctxSource = einoMessagesForRunRestart(args, baseMsgs, accumulated, baseCount)
return true, restartMsgs, ctxSource, backoff, nil
}
func (r *einoTransientRunRetrier) attempt() int { return r.attempts }
func (r *einoTransientRunRetrier) maxAttempts() int { return r.policy.maxAttempts }
// reset 在退避重试后成功推进(流/消息完整接收)时清零计数,使后续临时错误从第 1 次退避重新开始。
func (r *einoTransientRunRetrier) reset() { r.attempts = 0 }
func einoRunRetryMaxAttempts(args *einoADKRunLoopArgs) int {
if args != nil && args.RunRetryMaxAttempts > 0 {
return args.RunRetryMaxAttempts
@@ -85,7 +153,7 @@ func einoRunRetryMaxAttempts(args *einoADKRunLoopArgs) int {
return defaultEinoRunRetryMaxAttempts
}
// RunRetryMaxAttemptsFromConfig 供 handler 分段续跑计数(与 eino_middleware.run_retry_max_attempts 一致
// RunRetryMaxAttemptsFromConfig 与 eino_middleware.run_retry_max_attempts 一致。
func RunRetryMaxAttemptsFromConfig(mw *config.MultiAgentEinoMiddlewareConfig) int {
if mw != nil && mw.RunRetryMaxAttempts > 0 {
return mw.RunRetryMaxAttempts
@@ -93,15 +161,6 @@ func RunRetryMaxAttemptsFromConfig(mw *config.MultiAgentEinoMiddlewareConfig) in
return defaultEinoRunRetryMaxAttempts
}
// TransientRetryBackoff 供 handler 在分段续跑前退避。
func TransientRetryBackoff(attempt int, maxBackoffSec int) time.Duration {
max := defaultEinoRunRetryMaxBackoff
if maxBackoffSec > 0 {
max = time.Duration(maxBackoffSec) * time.Second
}
return einoTransientRetryBackoff(attempt, max)
}
func einoRunRetryMaxBackoff(args *einoADKRunLoopArgs) time.Duration {
if args != nil && args.RunRetryMaxBackoffSec > 0 {
return time.Duration(args.RunRetryMaxBackoffSec) * time.Second
@@ -122,10 +181,11 @@ const (
// 1) ModelFacingTrace(与模型实际入参一致) 2) 事件流累积的 runAccumulatedMsgs 3) 初始 msgs。
func einoMessagesForRunRestart(args *einoADKRunLoopArgs, baseMsgs, accumulated []adk.Message, baseCount int) ([]adk.Message, einoRunRestartContextSource) {
if trace := persistTraceSource(args, nil); len(trace) > 0 {
return append([]adk.Message(nil), trace...), einoRestartContextModelTrace
// modelFacingTrace includes prior Instruction system message(s); genModelInput will prepend again.
return stripADKSystemMessages(trace), einoRestartContextModelTrace
}
if len(accumulated) > baseCount {
return append([]adk.Message(nil), accumulated...), einoRestartContextAccumulated
return stripADKSystemMessages(accumulated), einoRestartContextAccumulated
}
return append([]adk.Message(nil), baseMsgs...), einoRestartContextInitial
}
@@ -27,6 +27,7 @@ func TestIsEinoTransientRunError(t *testing.T) {
{"429", errors.New("HTTP 429 Too Many Requests"), true},
{"rate limit", errors.New(`{"error":"rate limit exceeded"}`), true},
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
{"http2 goaway", errors.New("failed to receive stream chunk: error, http2: server sent GOAWAY and closed the connection; LastStreamID=791, ErrCode=NO_ERROR"), true},
{"unexpected eof", errors.New("unexpected EOF"), true},
{"503", errors.New("upstream returned 503"), true},
{"iteration limit", errors.New("max iteration reached"), false},
@@ -90,6 +91,46 @@ func TestEinoRunRetryMaxAttemptsFromArgs(t *testing.T) {
}
}
func TestEinoTransientRunRetrierReset(t *testing.T) {
t.Parallel()
r := newEinoTransientRunRetrier(einoTransientRunRetryPolicy{maxAttempts: 10, maxBackoff: 30 * time.Second})
r.attempts = 3
r.reset()
if r.attempt() != 0 {
t.Fatalf("after reset: attempt=%d, want 0", r.attempt())
}
// 重置后下一次退避应从 2s 起算(attempt index 0)。
if got := einoTransientRetryBackoff(r.attempt(), r.policy.maxBackoff); got != 2*time.Second {
t.Fatalf("backoff after reset: got %v, want 2s", got)
}
}
func TestEinoTransientRunRetrierConsecutiveFailures(t *testing.T) {
t.Parallel()
r := newEinoTransientRunRetrier(einoTransientRunRetryPolicy{maxAttempts: 10, maxBackoff: 30 * time.Second})
ctx := context.Background()
runErr := errors.New("internal server error")
args := &einoADKRunLoopArgs{}
base := []adk.Message{schema.UserMessage("hi")}
for want := 1; want <= 3; want++ {
restarted, _, _, _, err := r.tryRetry(ctx, runErr, args, base, nil, len(base))
if err != nil {
t.Fatalf("tryRetry attempt %d: %v", want, err)
}
if !restarted {
t.Fatalf("tryRetry attempt %d: want restarted", want)
}
if got := r.attempt(); got != want {
t.Fatalf("after failure %d: attempt=%d, want %d", want, got, want)
}
}
r.reset()
if r.attempt() != 0 {
t.Fatalf("after successful recovery reset: attempt=%d, want 0", r.attempt())
}
}
func TestAppendUserMessageIfNeeded(t *testing.T) {
t.Parallel()
msgs := []adk.Message{schema.UserMessage("old task")}
@@ -102,10 +143,3 @@ func TestAppendUserMessageIfNeeded(t *testing.T) {
t.Fatalf("should not duplicate user message: len=%d", len(dup))
}
}
func TestErrTransientRetryContinue(t *testing.T) {
t.Parallel()
if !errors.Is(ErrTransientRetryContinue, ErrTransientRetryContinue) {
t.Fatal("sentinel should match")
}
}
+15
View File
@@ -0,0 +1,15 @@
package multiagent
import "fmt"
// ExecuteExitError 表示 execute 命令非零退出(预期失败,非超时/中断/流异常)。
type ExecuteExitError struct {
Code int
}
func (e *ExecuteExitError) Error() string {
if e == nil {
return "exit status unknown"
}
return fmt.Sprintf("exit status %d", e.Code)
}
-8
View File
@@ -5,11 +5,3 @@ import "errors"
// ErrInterruptContinue 作为 context.CancelCause 使用:用户选择「中断并继续」且当前无进行中的 MCP 工具时,
// 取消当前推理/流式输出,并在同一会话任务内携带用户补充说明自动续跑下一轮(类似 Hermes 式人机回合)。
var ErrInterruptContinue = errors.New("agent interrupt: continue with user-supplied context")
// ErrTransientRetryContinue 表示 Run 因 429/网络等临时错误结束,应由 handler 落库轨迹后
// loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue 同级的「分段续跑」语义)。
var ErrTransientRetryContinue = errors.New("agent transient: retry after persisting trace")
// ErrEmptyResponseContinue 表示 Eino ADK 会话正常结束但未捕获到助手正文,应由 handler 落库轨迹后
// loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue / ErrTransientRetryContinue 同级)。
var ErrEmptyResponseContinue = errors.New("agent empty response: continue after persisting trace")
@@ -6,6 +6,7 @@ import (
"cyberstrike-ai/internal/agents"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/project"
"cyberstrike-ai/internal/projectprompt"
)
// DefaultPlanExecuteOrchestratorInstruction 当未配置 plan_execute 专用 Markdown / YAML 时的内置主代理(规划/重规划侧)提示。
@@ -122,7 +123,9 @@ func DefaultPlanExecuteOrchestratorInstruction() string {
## 表达
在调用工具或给出计划变更前 25 句中文说明当前决策依据与期望证据形态最终对用户交付结构化结论发现摘要证据风险下一步`
在调用工具或给出计划变更前 25 句中文说明当前决策依据与期望证据形态最终对用户交付结构化结论发现摘要证据风险下一步
` + projectprompt.ShellExecExecuteGuidanceSection()
}
// DefaultSupervisorOrchestratorInstruction 当未配置 supervisor 专用 Markdown / YAML 时的内置监督者提示(transfer / exit 说明仍由运行时在末尾追加)。
@@ -27,7 +27,7 @@ import (
// 本中间件与之互补,专职兜底正向孤儿。
// - 仅剔除消息,不向历史里注入虚构 assistant(tc):虚构 tool_calls 反而会误导模型后续推理。
// 摘要已覆盖被裁剪段的语义,丢一条原始 tool 结果对对话连贯性影响最小。
// - 位置建议:挂在所有可能改写历史的中间件(summarization / reduction / skill / plantask /
// - 位置建议:挂在 summarization / reduction / skill / plantask / system 合并 / 续聊 dedup 之后,
// tool_search)之后,靠近 ChatModel 调用的那一端。
type orphanToolPrunerMiddleware struct {
adk.BaseChatModelAgentMiddleware
+50 -47
View File
@@ -15,10 +15,12 @@ import (
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/agents"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/project"
"cyberstrike-ai/internal/reasoning"
"cyberstrike-ai/internal/security"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/adk"
@@ -56,6 +58,7 @@ func RunDeepAgent(
appCfg *config.Config,
ma *config.MultiAgentConfig,
ag *agent.Agent,
db *database.DB,
logger *zap.Logger,
conversationID string,
projectID string,
@@ -118,7 +121,7 @@ func RunDeepAgent(
mcpIDs = append(mcpIDs, id)
mcpIDsMu.Unlock()
}
einoExecMonitor := newEinoExecuteMonitorCallback(ag, recorder)
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ag, recorder)
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
snapshotMCPIDs := func() []string {
@@ -210,7 +213,7 @@ func RunDeepAgent(
subMax := resolveMaxIterations(appCfg, sub.MaxIterations)
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil {
return nil, fmt.Errorf("子代理 %q summarization 中间件: %w", id, err)
}
@@ -221,7 +224,7 @@ func RunDeepAgent(
}
if einoSkillMW != nil {
if einoFSTools && einoLoc != nil {
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if fsErr != nil {
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
}
@@ -229,13 +232,13 @@ func RunDeepAgent(
}
subHandlers = append(subHandlers, einoSkillMW)
}
subHandlers = append(subHandlers, subSumMw)
// 孤儿 tool 消息兜底:放在 summarization 之后,telemetry 之前,
// 以便 telemetry 记录的 token 数与 LLM 实际入参一致。
subHandlers = append(subHandlers, newOrphanToolPrunerMiddleware(logger, "sub_agent:"+id))
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "sub_agent"); teleMw != nil {
subHandlers = append(subHandlers, teleMw)
}
subHandlers = appendEinoChatModelTailMiddlewares(subHandlers, einoChatModelTailConfig{
logger: logger,
phase: "sub_agent:" + id,
summarization: subSumMw,
modelName: appCfg.OpenAI.Model,
conversationID: conversationID,
})
subInstrFinal := project.AppendVisionImageAnalysisIfReady(instr, appCfg.Vision.Ready())
subInstrFinal = injectToolNamesOnlyInstruction(ctx, subInstrFinal, subTools, subToolSearchActive)
@@ -281,7 +284,7 @@ func RunDeepAgent(
return nil, fmt.Errorf("多代理主模型: %w", err)
}
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, db, projectID, logger)
if err != nil {
return nil, fmt.Errorf("多代理主 summarization 中间件: %w", err)
}
@@ -356,12 +359,14 @@ func RunDeepAgent(
if einoLoc != nil && einoFSTools {
deepBackend = einoLoc
deepShell = &einoStreamingShellWrap{
inner: einoLoc,
invokeNotify: toolInvokeNotify,
einoAgentName: orchestratorName,
outputChunk: nil,
recordMonitor: einoExecMonitor,
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
inner: security.NewEinoStreamingShell(),
invokeNotify: toolInvokeNotify,
einoAgentName: orchestratorName,
outputChunk: nil,
beginMonitor: einoExecBegin,
finishMonitor: einoExecFinish,
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
}
}
@@ -377,14 +382,14 @@ func RunDeepAgent(
if einoSkillMW != nil {
deepHandlers = append(deepHandlers, einoSkillMW)
}
deepHandlers = append(deepHandlers, mainSumMw)
deepHandlers = append(deepHandlers, newOrphanToolPrunerMiddleware(logger, "deep_orchestrator"))
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "deep_orchestrator"); teleMw != nil {
deepHandlers = append(deepHandlers, teleMw)
}
if capMw := newModelFacingTraceMiddleware(modelFacingTrace); capMw != nil {
deepHandlers = append(deepHandlers, capMw)
}
deepHandlers = appendEinoChatModelTailMiddlewares(deepHandlers, einoChatModelTailConfig{
logger: logger,
phase: "deep_orchestrator",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
conversationID: conversationID,
trace: modelFacingTrace,
})
supHandlers := []adk.ChatModelAgentMiddleware{}
if len(mainOrchestratorPre) > 0 {
@@ -393,14 +398,14 @@ func RunDeepAgent(
if einoSkillMW != nil {
supHandlers = append(supHandlers, einoSkillMW)
}
supHandlers = append(supHandlers, mainSumMw)
supHandlers = append(supHandlers, newOrphanToolPrunerMiddleware(logger, "supervisor_orchestrator"))
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "supervisor_orchestrator"); teleMw != nil {
supHandlers = append(supHandlers, teleMw)
}
if capMw := newModelFacingTraceMiddleware(modelFacingTrace); capMw != nil {
supHandlers = append(supHandlers, capMw)
}
supHandlers = appendEinoChatModelTailMiddlewares(supHandlers, einoChatModelTailConfig{
logger: logger,
phase: "supervisor_orchestrator",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
conversationID: conversationID,
trace: modelFacingTrace,
})
mainToolsCfg := adk.ToolsConfig{
ToolsNodeConfig: compose.ToolsNodeConfig{
@@ -414,7 +419,7 @@ func RunDeepAgent(
EmitInternalEvents: true,
}
deepOutKey, modelRetry, taskGen := deepExtrasFromConfig(ma)
deepOutKey, taskGen := deepExtrasFromConfig(ma)
var da adk.Agent
switch orchMode {
@@ -426,7 +431,7 @@ func RunDeepAgent(
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
var peFsMw adk.ChatModelAgentMiddleware
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecMonitor, agentToolTimeoutMinutes(appCfg), nil)
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if err != nil {
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
}
@@ -441,18 +446,22 @@ func RunDeepAgent(
AppCfg: appCfg,
MwCfg: &ma.EinoMiddleware,
ConversationID: conversationID,
DB: db,
ProjectID: projectID,
Logger: logger,
ModelName: appCfg.OpenAI.Model,
ExecPreMiddlewares: mainOrchestratorPre,
SkillMiddleware: einoSkillMW,
FilesystemMiddleware: peFsMw,
ModelFacingTrace: modelFacingTrace,
PlannerReplannerRewriteHandlers: []adk.ChatModelAgentMiddleware{
mainSumMw,
// 孤儿 tool 消息兜底:必须挂在 summarization 之后、telemetry 之前。
newOrphanToolPrunerMiddleware(logger, "plan_execute_planner_replanner"),
newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "plan_execute_planner_replanner_rewrite"),
},
PlannerReplannerRewriteHandlers: appendEinoChatModelTailMiddlewares(nil, einoChatModelTailConfig{
logger: logger,
phase: "plan_execute_planner_replanner",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
conversationID: conversationID,
skipTrace: true,
}),
})
if perr != nil {
return nil, perr
@@ -469,9 +478,6 @@ func RunDeepAgent(
Handlers: supHandlers,
Exit: &adk.ExitTool{},
}
if modelRetry != nil {
supCfg.ModelRetryConfig = modelRetry
}
if deepOutKey != "" {
supCfg.OutputKey = deepOutKey
}
@@ -505,9 +511,6 @@ func RunDeepAgent(
if deepOutKey != "" {
dcfg.OutputKey = deepOutKey
}
if modelRetry != nil {
dcfg.ModelRetryConfig = modelRetry
}
if taskGen != nil {
dcfg.TaskToolDescriptionGenerator = taskGen
}
@@ -0,0 +1,33 @@
package multiagent
import (
"strings"
"cyberstrike-ai/internal/projectprompt"
)
func shellToolsPresent(toolNames []string) bool {
for _, n := range toolNames {
switch strings.ToLower(strings.TrimSpace(n)) {
case "exec", "execute":
return true
}
}
return false
}
// injectShellToolGuidance 在系统提示末尾追加 exec/execute 分工(仅当工具列表含 exec 或 execute)。
func injectShellToolGuidance(instruction string, toolNames []string) string {
if !shellToolsPresent(toolNames) {
return instruction
}
block := strings.TrimSpace(projectprompt.ShellExecExecuteGuidanceSection())
if block == "" {
return instruction
}
s := strings.TrimSpace(instruction)
if s == "" {
return block
}
return s + "\n\n" + block
}
@@ -0,0 +1,17 @@
package multiagent
import (
"strings"
"testing"
)
func TestInjectShellToolGuidance(t *testing.T) {
got := injectShellToolGuidance("base", []string{"nmap"})
if got != "base" {
t.Fatalf("expected unchanged, got %q", got)
}
got = injectShellToolGuidance("base", []string{"exec", "nmap"})
if !strings.Contains(got, "exec/execute") || !strings.Contains(got, "base") {
t.Fatalf("expected shell guidance appended, got %q", got)
}
}
@@ -0,0 +1,86 @@
package multiagent
import (
"context"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// systemMessageNormalizerMiddleware merges duplicate role=system messages into a single
// leading system message before summarization and each ChatModel call.
type systemMessageNormalizerMiddleware struct {
adk.BaseChatModelAgentMiddleware
logger *zap.Logger
phase string
}
func newSystemMessageNormalizerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
return &systemMessageNormalizerMiddleware{logger: logger, phase: phase}
}
func (m *systemMessageNormalizerMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, error) {
_ = mc
if m == nil || state == nil || len(state.Messages) == 0 {
return ctx, state, nil
}
before := countADKSystemMessages(state.Messages)
if before <= 1 {
return ctx, state, nil
}
normalized := normalizeSingleLeadingSystemMessage(state.Messages, "")
if len(normalized) == len(state.Messages) && countADKSystemMessages(normalized) >= before {
return ctx, state, nil
}
if m.logger != nil {
m.logger.Info("eino system messages merged",
zap.String("phase", m.phase),
zap.Int("system_before", before),
zap.Int("system_after", countADKSystemMessages(normalized)),
zap.Int("messages_before", len(state.Messages)),
zap.Int("messages_after", len(normalized)),
)
}
out := *state
out.Messages = normalized
return ctx, &out, nil
}
func countADKSystemMessages(msgs []adk.Message) int {
n := 0
for _, msg := range msgs {
if msg != nil && msg.Role == schema.System {
n++
}
}
return n
}
// stripADKSystemMessages removes all system messages. Use before runner.Run restart when
// genModelInput will prepend a fresh Instruction.
func stripADKSystemMessages(msgs []adk.Message) []adk.Message {
if len(msgs) == 0 {
return msgs
}
out := make([]adk.Message, 0, len(msgs))
for _, msg := range msgs {
if msg == nil || msg.Role == schema.System {
continue
}
out = append(out, msg)
}
return out
}
// mergeCollectedSystemMessages collapses multiple system messages into one (or none).
func mergeCollectedSystemMessages(systemMsgs []adk.Message) []adk.Message {
if len(systemMsgs) == 0 {
return nil
}
return normalizeSingleLeadingSystemMessage(systemMsgs, "")
}
@@ -0,0 +1,75 @@
package multiagent
import (
"context"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestStripADKSystemMessages(t *testing.T) {
in := []adk.Message{
schema.SystemMessage("a"),
schema.UserMessage("u"),
schema.SystemMessage("b"),
schema.AssistantMessage("x", nil),
}
out := stripADKSystemMessages(in)
if len(out) != 2 {
t.Fatalf("got %d messages, want 2", len(out))
}
if out[0].Role != schema.User || out[1].Role != schema.Assistant {
t.Fatalf("unexpected roles: %s, %s", out[0].Role, out[1].Role)
}
}
func TestEinoMessagesForRunRestart_StripsSystemFromTrace(t *testing.T) {
holder := newModelFacingTraceHolder()
holder.storeFromState(&adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("sys-1"),
schema.SystemMessage("sys-2"),
schema.UserMessage("task"),
}})
msgs, src := einoMessagesForRunRestart(&einoADKRunLoopArgs{ModelFacingTrace: holder}, nil, nil, 0)
if src != einoRestartContextModelTrace {
t.Fatalf("source: got %q want model_trace", src)
}
if len(msgs) != 1 || msgs[0].Role != schema.User {
t.Fatalf("expected user-only restart msgs, got %+v", msgs)
}
}
func TestSystemMessageNormalizerMiddleware_MergesDuplicates(t *testing.T) {
mw := newSystemMessageNormalizerMiddleware(nil, "test")
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("a"),
schema.SystemMessage("b"),
schema.UserMessage("u"),
}}
_, out, err := mw.(*systemMessageNormalizerMiddleware).BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatal(err)
}
if countADKSystemMessages(out.Messages) != 1 {
t.Fatalf("want 1 system, got %d", countADKSystemMessages(out.Messages))
}
if out.Messages[0].Content != "a\n\nb" {
t.Fatalf("merged content: %q", out.Messages[0].Content)
}
}
func TestSystemMessageNormalizerMiddleware_NoOpSingleSystem(t *testing.T) {
mw := newSystemMessageNormalizerMiddleware(nil, "test")
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("only"),
schema.UserMessage("u"),
}}
_, out, err := mw.(*systemMessageNormalizerMiddleware).BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatal(err)
}
if out != state {
t.Fatalf("expected same state pointer for no-op")
}
}
+10 -4
View File
@@ -10,7 +10,7 @@ package openai
// Auth: Bearer → x-api-key
// Tools: OpenAI tools[] → Claude tools[] (input_schema)
//
// Extended thinking: 顶层 `thinking` 从 OpenAI 请求体透传;响应中 `thinking` block 映射为
// Extended thinking: 顶层 `thinking` / `output_config` 从 OpenAI 请求体透传;响应中 `thinking` block 映射为
// `reasoning_content`(可读前缀 + 内部 JSON 尾缀以保留 signature,供多轮工具续跑;UI 用 openai.DisplayReasoningContent 剥离)。
import (
@@ -40,8 +40,9 @@ type claudeRequest struct {
System string `json:"system,omitempty"`
Messages []claudeMessage `json:"messages"`
Tools []claudeTool `json:"tools,omitempty"`
Stream bool `json:"stream,omitempty"`
Thinking json.RawMessage `json:"thinking,omitempty"`
Stream bool `json:"stream,omitempty"`
Thinking json.RawMessage `json:"thinking,omitempty"`
OutputConfig json.RawMessage `json:"output_config,omitempty"`
}
type claudeMessage struct {
@@ -304,12 +305,17 @@ func convertOpenAIToClaude(payload interface{}) (*claudeRequest, error) {
}
}
// Extended thinking (Anthropic top-level); merged from Eino ExtraFields / admin extras.
// Extended thinking + effort (Anthropic top-level); merged from Eino ExtraFields / admin extras.
if th, ok := oai["thinking"]; ok && th != nil {
if raw, err := json.Marshal(th); err == nil && len(raw) > 0 && string(raw) != "null" {
req.Thinking = json.RawMessage(raw)
}
}
if oc, ok := oai["output_config"]; ok && oc != nil {
if raw, err := json.Marshal(oc); err == nil && len(raw) > 0 && string(raw) != "null" {
req.OutputConfig = json.RawMessage(raw)
}
}
return req, nil
}
@@ -73,6 +73,39 @@ func TestConvertOpenAIToClaude_AssistantReasoningReplay(t *testing.T) {
}
}
func TestConvertOpenAIToClaude_OutputConfigEffort(t *testing.T) {
payload := map[string]interface{}{
"model": "claude-opus-4-8",
"messages": []interface{}{
map[string]interface{}{"role": "user", "content": "hi"},
},
"thinking": map[string]interface{}{
"type": "adaptive",
"display": "summarized",
},
"output_config": map[string]interface{}{
"effort": "high",
},
}
req, err := convertOpenAIToClaude(payload)
if err != nil {
t.Fatal(err)
}
if len(req.Thinking) == 0 {
t.Fatal("expected thinking")
}
if len(req.OutputConfig) == 0 {
t.Fatal("expected output_config")
}
var oc map[string]interface{}
if err := json.Unmarshal(req.OutputConfig, &oc); err != nil {
t.Fatal(err)
}
if oc["effort"] != "high" {
t.Fatalf("effort=%v", oc["effort"])
}
}
func TestClaudeToOpenAIResponseJSON_Thinking(t *testing.T) {
claudeBody := []byte(`{
"id":"msg_1","type":"message","role":"assistant","model":"x","stop_reason":"end_turn",
+35 -14
View File
@@ -2,7 +2,6 @@ package project
import (
"fmt"
"sort"
"strings"
"cyberstrike-ai/internal/config"
@@ -22,7 +21,13 @@ func AppendSystemPromptBlock(base, block string) string {
return base + "\n\n" + block
}
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(仅 key + summary,不含 body)。
const (
factIndexFooterGetDetail = "需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。"
factIndexFooterWriteHint = "写入事实 links 时用 from(来源 fact_key → 当前 fact),如 finding 上 {from:target/*, type:discovered_on}body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。"
factIndexFooterEmpty = "需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。"
)
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(key + summary + 关系边 + 攻击路径,不含 body)。
func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
if db == nil || !cfg.Enabled {
return "", nil
@@ -41,27 +46,38 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo
if err != nil {
return "", err
}
allEdges, _ := db.ListProjectFactEdgesByProject(projectID)
_, incomingByTarget := indexEdgeGroupMaps(allEdges)
if len(facts) == 0 {
return fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n(暂无事实)\n需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。", proj.Name, proj.ID), nil
return wrapFactIndexBlock(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n(暂无事实)\n%s", proj.Name, proj.ID, factIndexFooterEmpty)), nil
}
sort.SliceStable(facts, func(i, j int) bool {
if facts[i].Pinned != facts[j].Pinned {
return facts[i].Pinned
}
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
})
sortFactsForIndex(facts)
maxRunes := cfg.FactIndexMaxRunesEffective()
pathMaxRunes := cfg.FactIndexPathMaxRunesEffective()
footer := factIndexFooterGetDetail + "\n" + factIndexFooterWriteHint
footerRunes := len([]rune(footer))
factsBudget := maxRunes - pathMaxRunes - footerRunes
if factsBudget < 800 {
factsBudget = maxRunes - footerRunes
pathMaxRunes = 0
}
indexedKeys := make(map[string]struct{}, len(facts))
var b strings.Builder
b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s\n", proj.Name, proj.ID))
used := len([]rune(b.String()))
omitted := 0
for _, f := range facts {
line := fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
indexedKeys[f.FactKey] = struct{}{}
line := fmt.Sprintf("- [%s] %s — %s (%s)", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
line += FormatFactIndexLinksHint(f.FactKey, incomingByTarget[f.FactKey])
line += "\n"
lineRunes := len([]rune(line))
if used+lineRunes > maxRunes {
if used+lineRunes > factsBudget {
omitted++
continue
}
@@ -72,7 +88,12 @@ func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectCo
if omitted > 0 {
b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted))
}
b.WriteString("需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。\n")
b.WriteString("写入事实时:summary 写「什么+在哪+如何验证」;body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。\n")
return b.String(), nil
if pathSection := BuildFactPathOverviewSection(allEdges, indexedKeys, pathMaxRunes); pathSection != "" {
b.WriteString("\n")
b.WriteString(pathSection)
}
b.WriteString(footer)
return wrapFactIndexBlock(b.String()), nil
}
+56
View File
@@ -0,0 +1,56 @@
package project
import "strings"
// FactIndexSectionHeading 黑板索引可读标题行前缀(块内保留,供 Agent 阅读)。
const FactIndexSectionHeading = "## 项目黑板索引"
// FactIndexSectionStartMarker / EndMarkerHTML 注释边界,供程序化替换;对模型无指令语义。
const (
FactIndexSectionStartMarker = "<!-- fact-index-start -->"
FactIndexSectionEndMarker = "<!-- fact-index-end -->"
)
// ReplaceFactIndexSection 用 freshIndex 替换 content 中已有的项目黑板索引段。
// freshIndex 须为 BuildFactIndexBlock 的完整输出。起止 HTML 注释缺失时返回 (_, false)。
func ReplaceFactIndexSection(content, freshIndex string) (string, bool) {
freshIndex = strings.TrimSpace(freshIndex)
if freshIndex == "" {
return content, false
}
start, ok := factIndexSectionStart(content)
if !ok {
return content, false
}
end, ok := factIndexSectionEnd(content, start)
if !ok || end <= start {
return content, false
}
return content[:start] + freshIndex + content[end:], true
}
// wrapFactIndexBlock 为 BuildFactIndexBlock 正文加上统一起止 HTML 注释边界。
func wrapFactIndexBlock(content string) string {
content = strings.TrimSpace(content)
return FactIndexSectionStartMarker + "\n" + content + "\n" + FactIndexSectionEndMarker + "\n"
}
func factIndexSectionStart(content string) (int, bool) {
idx := strings.Index(content, FactIndexSectionStartMarker)
if idx < 0 {
return 0, false
}
return idx, true
}
func factIndexSectionEnd(content string, start int) (int, bool) {
if start < 0 || start >= len(content) {
return 0, false
}
tail := content[start:]
idx := strings.LastIndex(tail, FactIndexSectionEndMarker)
if idx < 0 {
return 0, false
}
return start + idx + len(FactIndexSectionEndMarker), true
}
+154
View File
@@ -0,0 +1,154 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func sampleFactIndexWithFacts(projectLabel, summary string) string {
return wrapFactIndexBlock("## 项目黑板索引(project: " + projectLabel + ", id: x\n" +
"- [target/a] target — " + summary + " (tentative)\n" +
factIndexFooterGetDetail + "\n" +
factIndexFooterWriteHint)
}
func TestReplaceFactIndexSection(t *testing.T) {
t.Parallel()
oldIndex := sampleFactIndexWithFacts("p1", "old summary")
newIndex := sampleFactIndexWithFacts("p1", "new summary")
t.Run("replaces index before next section", func(t *testing.T) {
content := "你是助手\n\n" + oldIndex + "\n\n## 图片分析\n看截图"
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "old summary") {
t.Fatalf("old index should be gone: %q", out)
}
if !strings.Contains(out, "new summary") || !strings.Contains(out, "## 图片分析") {
t.Fatalf("expected new index and preserved vision section: %q", out)
}
if strings.Count(out, FactIndexSectionStartMarker) != 1 || strings.Count(out, FactIndexSectionEndMarker) != 1 {
t.Fatalf("expected exactly one start/end marker pair: %q", out)
}
})
t.Run("replaces index at end", func(t *testing.T) {
content := "## 项目测试范围\nscope\n\n" + oldIndex
out, ok := ReplaceFactIndexSection(content, newIndex)
if !ok {
t.Fatal("expected replacement")
}
if !strings.Contains(out, "## 项目测试范围") || !strings.Contains(out, "new summary") {
t.Fatalf("scope preserved, index updated: %q", out)
}
})
t.Run("summary with false markdown header does not truncate early", func(t *testing.T) {
summaryWithFakeHeader := "see\n\n## fake header in summary"
old := sampleFactIndexWithFacts("p1", summaryWithFakeHeader)
newIdx := sampleFactIndexWithFacts("p1", "new summary")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "fake header in summary") {
t.Fatalf("old index tail should be fully removed: %q", out)
}
})
t.Run("summary containing end marker text does not truncate early", func(t *testing.T) {
summary := "note " + FactIndexSectionEndMarker + " in summary"
old := sampleFactIndexWithFacts("p1", summary)
newIdx := sampleFactIndexWithFacts("p1", "clean")
content := old + "\n\n## 图片分析\nvision"
out, ok := ReplaceFactIndexSection(content, newIdx)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "in summary") {
t.Fatalf("old block should be fully removed: %q", out)
}
})
t.Run("missing html markers does not replace", func(t *testing.T) {
legacy := "## 项目黑板索引(project: p1, id: x\n- [a] note — old (tentative)\n"
newIdx := sampleFactIndexWithFacts("p1", "new")
out, ok := ReplaceFactIndexSection("prefix\n\n"+legacy, newIdx)
if ok {
t.Fatalf("expected no replacement without markers: %q", out)
}
})
t.Run("empty facts block", func(t *testing.T) {
oldEmpty := wrapFactIndexBlock("## 项目黑板索引(project: p1, id: x\n(暂无事实)\n" + factIndexFooterEmpty)
newEmpty := sampleFactIndexWithFacts("p1", "first fact")
out, ok := ReplaceFactIndexSection(oldEmpty, newEmpty)
if !ok {
t.Fatal("expected replacement")
}
if strings.Contains(out, "(暂无事实)") {
t.Fatalf("old empty block should be gone: %q", out)
}
})
t.Run("no marker", func(t *testing.T) {
_, ok := ReplaceFactIndexSection("no blackboard here", newIndex)
if ok {
t.Fatal("expected false when marker missing")
}
})
t.Run("empty fresh index", func(t *testing.T) {
_, ok := ReplaceFactIndexSection(oldIndex, " ")
if ok {
t.Fatal("expected false for empty fresh index")
}
})
}
func TestFactIndexSectionBounds_useHTMLMarkers(t *testing.T) {
t.Parallel()
body := sampleFactIndexWithFacts("p", "line with\n\n## not a real section") + "TAIL_SHOULD_DROP"
start, ok := factIndexSectionStart(body)
if !ok || !strings.HasPrefix(body[start:], FactIndexSectionStartMarker) {
t.Fatalf("start should be at html start marker, got %d", start)
}
end, ok := factIndexSectionEnd(body, start)
if !ok || body[end:] != "\nTAIL_SHOULD_DROP" {
t.Fatalf("end should be after end marker, got remainder %q", body[end:])
}
}
func TestBuildFactIndexBlock_includesHTMLMarkers(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "marker-proj"})
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true})
if err != nil {
t.Fatal(err)
}
if !strings.HasPrefix(strings.TrimSpace(block), FactIndexSectionStartMarker) {
t.Fatalf("block should start with start marker: %q", block)
}
if !strings.Contains(block, FactIndexSectionEndMarker) {
t.Fatalf("block should include end marker: %q", block)
}
}
+256
View File
@@ -0,0 +1,256 @@
package project
import (
"fmt"
"regexp"
"strings"
"cyberstrike-ai/internal/database"
)
var (
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[:]\s*([a-z0-9][a-z0-9._/-]*)`)
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[:]\s*([a-z0-9][a-z0-9._/-]*)`)
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
bodySyncLinksHead = "结构化关系边(自动同步)"
)
// ParseLinksFromBody 从 body「关联」段落解析 from 语义的关系边(无显式 links 时的兜底)。
func ParseLinksFromBody(body string) []database.ProjectFactEdgeFromInput {
body = strings.TrimSpace(body)
if body == "" {
return nil
}
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
add := func(key, edgeType string) {
key = strings.TrimSpace(key)
if key == "" {
return
}
if err := database.ValidateFactKey(key); err != nil {
return
}
sig := edgeType + "\x00" + key
if _, ok := seen[sig]; ok {
return
}
seen[sig] = struct{}{}
out = append(out, database.ProjectFactEdgeFromInput{From: key, Type: edgeType})
}
for _, m := range bodyDepFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "depends_on")
}
}
for _, m := range bodyRelFactLine.FindAllStringSubmatch(body, -1) {
if len(m) > 1 {
add(m[1], "supports")
}
}
// 自动同步块:type: key
syncBlock := extractBodySyncLinksBlock(body)
for _, line := range strings.Split(syncBlock, "\n") {
line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "-"))
if line == "" {
continue
}
edgeType, source, ok := strings.Cut(line, ":")
if !ok {
continue
}
edgeType = strings.TrimSpace(edgeType)
source = strings.TrimSpace(source)
if err := database.ValidateProjectFactEdgeType(edgeType); err != nil {
continue
}
add(source, edgeType)
}
if len(out) == 0 {
return nil
}
return out
}
func extractBodySyncLinksBlock(body string) string {
lines := strings.Split(body, "\n")
var b strings.Builder
inAssoc := false
inSync := false
for _, line := range lines {
trim := strings.TrimSpace(line)
if bodyAssocSection.MatchString(trim) {
inAssoc = true
inSync = false
continue
}
if inAssoc && strings.HasPrefix(trim, "## ") && !strings.HasPrefix(trim, "## 关联") {
break
}
if inAssoc && strings.Contains(trim, bodySyncLinksHead) {
inSync = true
continue
}
if inSync {
if trim == "" || strings.HasPrefix(trim, "-") || strings.Contains(trim, ":") {
if strings.HasPrefix(trim, "-") || (strings.Contains(trim, ":") && !strings.Contains(trim, "related_vulnerability")) {
b.WriteString(trim)
b.WriteByte('\n')
}
} else if strings.HasPrefix(trim, "##") {
break
}
}
}
return b.String()
}
// SyncBodyLinksSection 将入边镜像写入 body 的「关联」段(人读用;结构化以 links 为准)。
func SyncBodyLinksSection(body string, edges []*database.ProjectFactEdge) string {
body = strings.TrimSpace(body)
block := formatBodySyncLinksBlock(edges)
if block == "" {
return body
}
if body == "" {
return "## 关联\n" + block
}
lines := strings.Split(body, "\n")
var out []string
inAssoc := false
replaced := false
for i := 0; i < len(lines); i++ {
trim := strings.TrimSpace(lines[i])
if bodyAssocSection.MatchString(trim) {
inAssoc = true
out = append(out, lines[i])
// 跳过旧同步块
j := i + 1
for j < len(lines) {
t := strings.TrimSpace(lines[j])
if strings.HasPrefix(t, "## ") {
break
}
if strings.Contains(t, bodySyncLinksHead) {
for j < len(lines) {
t2 := strings.TrimSpace(lines[j])
if t2 != "" && !strings.HasPrefix(t2, "-") && !strings.Contains(t2, ":") && !strings.Contains(t2, bodySyncLinksHead) {
if strings.HasPrefix(t2, "##") {
break
}
}
j++
if j < len(lines) && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
if j >= len(lines) {
break
}
if j > i+1 && strings.TrimSpace(lines[j-1]) == "" && strings.HasPrefix(strings.TrimSpace(lines[j]), "## ") {
break
}
}
break
}
j++
}
out = append(out, block)
i = j - 1
replaced = true
continue
}
out = append(out, lines[i])
}
if !replaced {
if !inAssoc {
out = append(out, "", "## 关联", block)
} else {
out = append(out, block)
}
}
return strings.TrimSpace(strings.Join(out, "\n"))
}
func formatBodySyncLinksBlock(edges []*database.ProjectFactEdge) string {
if len(edges) == 0 {
return fmt.Sprintf("- %s:\n (暂无)", bodySyncLinksHead)
}
var b strings.Builder
b.WriteString("- ")
b.WriteString(bodySyncLinksHead)
b.WriteString(":\n")
for _, e := range edges {
b.WriteString(fmt.Sprintf(" - %s: %s\n", e.EdgeType, e.SourceFactKey))
}
return strings.TrimRight(b.String(), "\n")
}
// ResolveFactLinksForUpsert 合并显式 links、links_text 与 body 解析结果。
func ResolveFactLinksForUpsert(explicit []database.ProjectFactEdgeFromInput, linksText *string, body string, explicitSet bool) ([]database.ProjectFactEdgeFromInput, bool, error) {
if explicitSet {
if len(explicit) > 0 {
return explicit, true, nil
}
if linksText != nil {
parsed, err := ParseFactLinksText(*linksText)
if err != nil {
return nil, true, err
}
if parsed == nil {
return []database.ProjectFactEdgeFromInput{}, true, nil
}
return parsed, true, nil
}
return []database.ProjectFactEdgeFromInput{}, true, nil
}
if parsed := ParseLinksFromBody(body); len(parsed) > 0 {
return parsed, true, nil
}
return nil, false, nil
}
// MergeLinkFromInputsUnique 合并多组 from 入边输入并去重。
func MergeLinkFromInputsUnique(groups ...[]database.ProjectFactEdgeFromInput) []database.ProjectFactEdgeFromInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeFromInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.From
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.From); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
// MergeLinkInputsUnique 合并多组 link 输入并去重(内部出边写入用)。
func MergeLinkInputsUnique(groups ...[]database.ProjectFactEdgeInput) []database.ProjectFactEdgeInput {
seen := map[string]struct{}{}
var out []database.ProjectFactEdgeInput
for _, g := range groups {
for _, in := range g {
sig := in.Type + "\x00" + in.To
if _, ok := seen[sig]; ok {
continue
}
if err := database.ValidateProjectFactEdgeType(in.Type); err != nil {
continue
}
if err := database.ValidateFactKey(in.To); err != nil {
continue
}
seen[sig] = struct{}{}
out = append(out, in)
}
}
return out
}
+68
View File
@@ -0,0 +1,68 @@
package project
import (
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseLinksFromBodyDependsOn(t *testing.T) {
t.Parallel()
body := "## 关联\n- 依赖事实: target/api\n- 相关 fact_key: auth/session"
links := ParseLinksFromBody(body)
if len(links) != 2 {
t.Fatalf("want 2 links, got %d", len(links))
}
}
func TestSyncBodyLinksSection(t *testing.T) {
t.Parallel()
body := "## 结论\nx\n\n## 关联\n- 依赖事实: old/key"
edges := []*database.ProjectFactEdge{{EdgeType: "discovered_on", SourceFactKey: "target/a"}}
out := SyncBodyLinksSection(body, edges)
if !strings.Contains(out, "discovered_on: target/a") {
t.Fatalf("missing synced edge: %q", out)
}
}
func TestFactGraphIntegration(t *testing.T) {
dir := t.TempDir()
dbPath := filepath.Join(dir, "test.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "g"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat, summary string }{
{"target/root", "target", "root"},
{"finding/x", "finding", "finding x"},
} {
_, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.summary, Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/x", []database.ProjectFactEdgeFromInput{
{From: "target/root", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
if len(graph.Nodes) < 2 || len(graph.Edges) < 1 {
t.Fatalf("expected graph nodes/edges, got %d/%d", len(graph.Nodes), len(graph.Edges))
}
}
+407
View File
@@ -0,0 +1,407 @@
package project
import (
"fmt"
"strings"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/projectprompt"
)
// PathGraphCategories 攻击路径视图包含的事实分类。
var PathGraphCategories = map[string]struct{}{
FactCategoryTarget: {},
FactCategoryFinding: {},
FactCategoryChain: {},
FactCategoryExploit: {},
FactCategoryPOC: {},
"vuln": {},
}
// GraphNodeType 将 fact category 映射为图节点类型(供前端样式与 ELK 分层)。
// 优先使用 category;仅 synthetic 节点(vuln:)或无 category 时才回退到 fact_key 前缀。
func GraphNodeType(category, factKey string) string {
key := strings.ToLower(strings.TrimSpace(factKey))
if strings.HasPrefix(key, "vuln:") {
return "vulnerability"
}
c := strings.ToLower(strings.TrimSpace(category))
if c != "" {
switch c {
case FactCategoryTarget:
return "target"
case FactCategoryExploit:
return "exploit"
case FactCategoryPOC:
return "poc"
case FactCategoryChain:
return "chain"
case FactCategoryFinding:
return "finding"
case "vuln":
return "vulnerability"
case FactCategoryAuth:
return "auth"
case FactCategoryInfra, FactCategoryBusiness:
return "infra"
case FactCategoryNote:
return "note"
case "missing":
return "missing"
default:
return c
}
}
switch {
case strings.HasPrefix(key, "target/"):
return "target"
case strings.HasPrefix(key, "exploit/"), strings.HasPrefix(key, "evidence/"):
return "exploit"
case strings.HasPrefix(key, "poc/"):
return "poc"
case strings.HasPrefix(key, "chain/"):
return "chain"
case strings.HasPrefix(key, "finding/"):
return "finding"
case strings.HasPrefix(key, "auth/"):
return "auth"
case strings.HasPrefix(key, "infra/"), strings.HasPrefix(key, "business/"):
return "infra"
default:
return "note"
}
}
func truncateGraphLabel(summary string, maxRunes int) string {
summary = strings.TrimSpace(summary)
if summary == "" {
return "—"
}
r := []rune(summary)
if len(r) <= maxRunes {
return summary
}
return string(r[:maxRunes]) + "…"
}
// BuildProjectFactGraph 构建项目事实图(nodes + edges)。
func BuildProjectFactGraph(db *database.DB, projectID string, view string, excludeDeprecated bool) (*database.ProjectFactGraph, error) {
if db == nil {
return nil, fmt.Errorf("database 未初始化")
}
projectID = strings.TrimSpace(projectID)
if projectID == "" {
return nil, fmt.Errorf("project_id 不能为空")
}
view = strings.TrimSpace(strings.ToLower(view))
if view == "" {
view = "path"
}
filter := database.ProjectFactListFilter{}
if excludeDeprecated {
filter.ExcludeDeprecated = true
}
facts, err := db.ListProjectFacts(projectID, filter, 1000, 0)
if err != nil {
return nil, err
}
edges, err := db.ListProjectFactEdgesByProject(projectID)
if err != nil {
return nil, err
}
if excludeDeprecated {
edges = filterDeprecatedEdges(edges)
}
factByKey := make(map[string]*database.ProjectFact, len(facts))
for _, f := range facts {
factByKey[f.FactKey] = f
}
pathMode := view == "path"
nodeKeys := make(map[string]struct{})
if pathMode {
for _, f := range facts {
if isPathGraphFact(f.Category, f.FactKey) {
nodeKeys[f.FactKey] = struct{}{}
}
}
// 路径视图中保留作为依赖目标的 auth/infra 节点
for _, e := range edges {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if f, ok := factByKey[e.TargetFactKey]; ok && isDependencyGraphFact(f.Category, f.FactKey) {
nodeKeys[e.TargetFactKey] = struct{}{}
}
}
} else {
for _, f := range facts {
nodeKeys[f.FactKey] = struct{}{}
}
}
// 边上引用的 endpoint 纳入节点集
for _, e := range edges {
if pathMode {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; ok {
// already included
} else if f, ok := factByKey[e.TargetFactKey]; !ok {
nodeKeys[e.TargetFactKey] = struct{}{} // 占位节点
} else if isPathGraphFact(f.Category, f.FactKey) || isDependencyGraphFact(f.Category, f.FactKey) {
nodeKeys[e.TargetFactKey] = struct{}{}
} else {
continue
}
} else {
nodeKeys[e.SourceFactKey] = struct{}{}
nodeKeys[e.TargetFactKey] = struct{}{}
}
}
nodes := make([]database.ProjectFactGraphNode, 0, len(nodeKeys))
for key := range nodeKeys {
if f, ok := factByKey[key]; ok {
nodes = append(nodes, database.ProjectFactGraphNode{
ID: f.FactKey,
FactKey: f.FactKey,
Category: f.Category,
Label: truncateGraphLabel(f.Summary, 48),
Summary: strings.TrimSpace(f.Summary),
Confidence: f.Confidence,
Type: GraphNodeType(f.Category, f.FactKey),
Pinned: f.Pinned,
})
continue
}
nodes = append(nodes, database.ProjectFactGraphNode{
ID: key,
FactKey: key,
Category: "missing",
Label: key,
Confidence: "tentative",
Type: "missing",
Pinned: false,
})
}
graphEdges := make([]database.ProjectFactGraphEdge, 0, len(edges))
for _, e := range edges {
if pathMode {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
continue
}
} else {
if _, ok := nodeKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := nodeKeys[e.TargetFactKey]; !ok {
continue
}
}
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
ID: e.ID,
Source: e.SourceFactKey,
Target: e.TargetFactKey,
Type: e.EdgeType,
Confidence: e.Confidence,
})
}
// related_vulnerability_id 合成边(source=fact → target=vuln:<id>
for _, f := range facts {
if _, ok := nodeKeys[f.FactKey]; !ok {
continue
}
vid := strings.TrimSpace(f.RelatedVulnerabilityID)
if vid == "" {
continue
}
vulnNodeID := "vuln:" + vid
if _, exists := nodeKeys[vulnNodeID]; !exists {
nodeKeys[vulnNodeID] = struct{}{}
label := "漏洞"
if len(vid) >= 8 {
label += " " + vid[:8] + "…"
} else {
label += " " + vid
}
nodes = append(nodes, database.ProjectFactGraphNode{
ID: vulnNodeID,
FactKey: vulnNodeID,
Category: "vuln",
Label: label,
Confidence: f.Confidence,
Type: "vulnerability",
Pinned: false,
})
}
graphEdges = append(graphEdges, database.ProjectFactGraphEdge{
ID: "vuln-link:" + f.FactKey + ":" + vid,
Source: f.FactKey,
Target: vulnNodeID,
Type: "links_vuln",
Confidence: f.Confidence,
})
}
return &database.ProjectFactGraph{Nodes: nodes, Edges: graphEdges}, nil
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
func isPathGraphFact(category, factKey string) bool {
c := strings.ToLower(strings.TrimSpace(category))
if _, ok := PathGraphCategories[c]; ok {
return true
}
if c != "" {
return false
}
key := strings.ToLower(strings.TrimSpace(factKey))
for _, p := range []string{"target/", "finding/", "chain/", "exploit/", "poc/", "evidence/"} {
if strings.HasPrefix(key, p) {
return true
}
}
return false
}
func isDependencyGraphFact(category, factKey string) bool {
c := strings.ToLower(strings.TrimSpace(category))
if c == FactCategoryAuth || c == FactCategoryInfra || c == FactCategoryBusiness {
return true
}
if c != "" {
return false
}
key := strings.ToLower(strings.TrimSpace(factKey))
return strings.HasPrefix(key, "auth/") || strings.HasPrefix(key, "infra/") || strings.HasPrefix(key, "business/")
}
func filterDeprecatedEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
out := make([]*database.ProjectFactEdge, 0, len(edges))
for _, e := range edges {
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
continue
}
out = append(out, e)
}
return out
}
// ParsedFactLinks 解析 links 参数(from → 当前 fact)。
type ParsedFactLinks struct {
Incoming []database.ProjectFactEdgeFromInput
}
// ParseFactLinkInputs 从 MCP links 参数解析;空数组表示清空全部入边。
func ParseFactLinkInputs(raw interface{}) (*ParsedFactLinks, error) {
if raw == nil {
return nil, nil
}
items, ok := raw.([]interface{})
if !ok {
return nil, fmt.Errorf("links 须为数组")
}
if len(items) == 0 {
return &ParsedFactLinks{
Incoming: []database.ProjectFactEdgeFromInput{},
}, nil
}
parsed := &ParsedFactLinks{}
for i, item := range items {
m, ok := item.(map[string]interface{})
if !ok {
return nil, fmt.Errorf("links[%d] 格式无效", i)
}
from, _ := m["from"].(string)
edgeType, _ := m["type"].(string)
from = strings.TrimSpace(from)
edgeType = strings.TrimSpace(edgeType)
if from == "" {
return nil, fmt.Errorf("links[%d] 须含 from", i)
}
if edgeType == "" {
return nil, fmt.Errorf("links[%d] 须含 type", i)
}
conf, _ := m["confidence"].(string)
parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{
From: from, Type: edgeType, Confidence: strings.TrimSpace(conf),
})
}
return parsed, nil
}
// ParseFactLinksText 解析 UI 文本:`type: source_fact_key` 每行一条(from 语义)。
func ParseFactLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
return ParseFactIncomingLinksText(text)
}
// FormatFactLinksText 将入边格式化为 UI 文本。
func FormatFactLinksText(edges []*database.ProjectFactEdge) string {
return FormatFactIncomingLinksText(edges)
}
// ParseFactIncomingLinksText 解析 UI 入边文本:`type: source_fact_key` 每行一条。
func ParseFactIncomingLinksText(text string) ([]database.ProjectFactEdgeFromInput, error) {
text = strings.TrimSpace(text)
if text == "" {
return nil, nil
}
var out []database.ProjectFactEdgeFromInput
for i, line := range strings.Split(text, "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
edgeType, source, ok := strings.Cut(line, ":")
if !ok {
return nil, fmt.Errorf("第 %d 行格式无效,应为 type: fact_key", i+1)
}
edgeType = strings.TrimSpace(edgeType)
source = strings.TrimSpace(source)
if edgeType == "" || source == "" {
return nil, fmt.Errorf("第 %d 行 type 或 fact_key 为空", i+1)
}
out = append(out, database.ProjectFactEdgeFromInput{From: source, Type: edgeType})
}
return out, nil
}
// FormatFactIncomingLinksText 将入边格式化为 UI 文本。
func FormatFactIncomingLinksText(edges []*database.ProjectFactEdge) string {
if len(edges) == 0 {
return ""
}
var b strings.Builder
for i, e := range edges {
if i > 0 {
b.WriteByte('\n')
}
b.WriteString(e.EdgeType)
b.WriteString(": ")
b.WriteString(e.SourceFactKey)
}
return b.String()
}
// FactEdgeRecordingGuidance 写入边时的 Agent 规范。
func FactEdgeRecordingGuidance() string {
return projectprompt.FactEdgeRecordingGuidance()
}
+96
View File
@@ -0,0 +1,96 @@
package project
import (
"cyberstrike-ai/internal/database"
)
// ApplyFactOutgoingLinks 替换某事实的出边(links 为 nil 时不修改)。
func ApplyFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput) error {
if links == nil {
return nil
}
return db.ReplaceOutgoingProjectFactEdges(projectID, sourceFactKey, sourceConversationID, links)
}
// ResolveFactLinkInputs 合并 links 数组与 links_text 文本(数组优先)。
func ResolveFactLinkInputs(links []database.ProjectFactEdgeFromInput, linksText string) ([]database.ProjectFactEdgeFromInput, error) {
if len(links) > 0 {
return links, nil
}
return ParseFactLinksText(linksText)
}
// ApplyFactIncomingLinks 替换某事实的入边(links 为 nil 时不修改)。
func ApplyFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput) error {
if links == nil {
return nil
}
return db.ReplaceIncomingProjectFactEdges(projectID, targetFactKey, links)
}
// PersistFactIncomingLinks 写入入边并可选同步当前事实 body「关联」段。
func PersistFactIncomingLinks(db *database.DB, projectID, targetFactKey string, links []database.ProjectFactEdgeFromInput, syncBody bool) error {
if links == nil {
return nil
}
if err := ApplyFactIncomingLinks(db, projectID, targetFactKey, links); err != nil {
return err
}
if !syncBody {
return nil
}
f, err := db.GetProjectFactByKey(projectID, targetFactKey)
if err != nil {
return nil
}
in, err := db.ListIncomingProjectFactEdges(projectID, targetFactKey)
if err != nil {
return err
}
f.Body = SyncBodyLinksSection(f.Body, in)
_, err = db.UpsertProjectFact(f)
return err
}
// PersistFactLinksFromParsed 写入解析后的 linksparsed 为 nil 表示不修改)。
func PersistFactLinksFromParsed(db *database.DB, projectID, factKey, sourceConversationID string, parsed *ParsedFactLinks, syncBody bool) error {
if parsed == nil || parsed.Incoming == nil {
return nil
}
return PersistFactIncomingLinks(db, projectID, factKey, parsed.Incoming, syncBody)
}
// PersistFactOutgoingLinks 写入出边(图连线等低层 APIbody 同步请用 PersistFactIncomingLinks)。
func PersistFactOutgoingLinks(db *database.DB, projectID, sourceFactKey, sourceConversationID string, links []database.ProjectFactEdgeInput, syncBody bool) error {
if links == nil {
return nil
}
return ApplyFactOutgoingLinks(db, projectID, sourceFactKey, sourceConversationID, links)
}
// LinkCountMap 项目内各 fact 的入/出边计数。
type LinkCountMap map[string]LinkCounts
// LinkCounts 单 fact 的入/出边数。
type LinkCounts struct {
Outgoing int `json:"outgoing"`
Incoming int `json:"incoming"`
}
// LoadProjectFactLinkCounts 批量加载边计数。
func LoadProjectFactLinkCounts(db *database.DB, projectID string) (LinkCountMap, error) {
edges, err := db.ListProjectFactEdgesByProject(projectID)
if err != nil {
return nil, err
}
m := LinkCountMap{}
for _, e := range edges {
c := m[e.SourceFactKey]
c.Outgoing++
m[e.SourceFactKey] = c
c = m[e.TargetFactKey]
c.Incoming++
m[e.TargetFactKey] = c
}
return m, nil
}
+296
View File
@@ -0,0 +1,296 @@
package project
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestParseFactLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactLinksText("discovered_on: target/api\nleads_to: finding/swagger")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "discovered_on" || inputs[0].From != "target/api" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestParseFactIncomingLinksText(t *testing.T) {
t.Parallel()
inputs, err := ParseFactIncomingLinksText("leads_to: finding/swagger\ndepends_on: target/api")
if err != nil {
t.Fatal(err)
}
if len(inputs) != 2 {
t.Fatalf("want 2 links, got %d", len(inputs))
}
if inputs[0].Type != "leads_to" || inputs[0].From != "finding/swagger" {
t.Fatalf("unexpected first link: %+v", inputs[0])
}
}
func TestFormatFactIncomingLinksText(t *testing.T) {
t.Parallel()
text := FormatFactIncomingLinksText([]*database.ProjectFactEdge{
{EdgeType: "leads_to", SourceFactKey: "finding/a"},
{EdgeType: "depends_on", SourceFactKey: "target/b"},
})
want := "leads_to: finding/a\ndepends_on: target/b"
if text != want {
t.Fatalf("got %q want %q", text, want)
}
}
func TestParseFactLinkInputsEmptyClears(t *testing.T) {
t.Parallel()
parsed, err := ParseFactLinkInputs([]interface{}{})
if err != nil {
t.Fatal(err)
}
if parsed == nil || parsed.Incoming == nil || len(parsed.Incoming) != 0 {
t.Fatalf("empty array should clear incoming links, got %v", parsed)
}
}
func TestParseFactLinkInputsFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"from": "target/primary_domain",
"type": "discovered_on",
},
}
parsed, err := ParseFactLinkInputs(raw)
if err != nil {
t.Fatal(err)
}
if len(parsed.Incoming) != 1 || parsed.Incoming[0].From != "target/primary_domain" {
t.Fatalf("unexpected incoming: %+v", parsed.Incoming)
}
}
func TestParseFactLinkInputsRequiresFrom(t *testing.T) {
t.Parallel()
raw := []interface{}{
map[string]interface{}{
"to": "target/primary_domain",
"type": "discovered_on",
},
}
_, err := ParseFactLinkInputs(raw)
if err == nil {
t.Fatal("expected error when from is missing")
}
}
func TestGraphNodeType(t *testing.T) {
t.Parallel()
if GraphNodeType("chain", "chain/x") != "chain" {
t.Fatal("chain category")
}
if GraphNodeType("finding", "finding/x") != "finding" {
t.Fatal("finding category")
}
if GraphNodeType("exploit", "exploit/x") != "exploit" {
t.Fatal("exploit category")
}
if GraphNodeType("finding", "evidence/x") != "finding" {
t.Fatal("category should override evidence key prefix")
}
if GraphNodeType("note", "target/x") != "note" {
t.Fatal("category should override target key prefix")
}
if GraphNodeType("vuln", "finding/x") != "vulnerability" {
t.Fatal("vuln category maps to vulnerability node type")
}
if GraphNodeType("", "target/x") != "target" {
t.Fatal("empty category falls back to target key prefix")
}
}
func TestBuildProjectFactGraphPreservesStoredEdgeDirection(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "path-edges"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"chain/full_attack_path", "chain"},
{"finding/mysql_public", "finding"},
{"exploit/mysql_creds_extract", "exploit"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "finding/mysql_public", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
{From: "exploit/mysql_creds_extract", Type: "exploits"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "chain/full_attack_path", []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/mysql_creds_extract", []database.ProjectFactEdgeFromInput{
{From: "chain/full_attack_path", Type: "leads_to"},
}); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := map[string]struct{}{
"target/primary_domain|discovered_on|finding/mysql_public": {},
"exploit/mysql_creds_extract|exploits|finding/mysql_public": {},
"target/primary_domain|discovered_on|chain/full_attack_path": {},
"chain/full_attack_path|leads_to|exploit/mysql_creds_extract": {},
}
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
delete(want, key)
}
if len(want) > 0 {
t.Fatalf("missing expected stored-direction edges: %v", want)
}
countInOut := func(factKey string) (out, in int) {
for _, e := range graph.Edges {
if e.Source == factKey {
out++
}
if e.Target == factKey {
in++
}
}
return out, in
}
if out, in := countInOut("chain/full_attack_path"); out != 1 || in != 1 {
t.Fatalf("chain/full_attack_path want out=1 in=1 got out=%d in=%d", out, in)
}
if out, in := countInOut("exploit/mysql_creds_extract"); out != 1 || in != 1 {
t.Fatalf("exploit/mysql_creds_extract want out=1 in=1 got out=%d in=%d", out, in)
}
}
func TestPersistFactLinksFromUsesFromAsIncoming(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "from-links"})
if err != nil {
t.Fatal(err)
}
for _, spec := range []struct{ key, cat string }{
{"target/primary_domain", "target"},
{"finding/sqli", "finding"},
} {
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: spec.key, Category: spec.cat, Summary: spec.key, Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
}
parsed := &ParsedFactLinks{
Incoming: []database.ProjectFactEdgeFromInput{
{From: "target/primary_domain", Type: "discovered_on"},
},
}
if err := PersistFactLinksFromParsed(db, p.ID, "finding/sqli", "", parsed, false); err != nil {
t.Fatal(err)
}
graph, err := BuildProjectFactGraph(db, p.ID, "path", true)
if err != nil {
t.Fatal(err)
}
want := "target/primary_domain|discovered_on|finding/sqli"
for _, e := range graph.Edges {
key := e.Source + "|" + e.Type + "|" + e.Target
if key == want {
return
}
}
t.Fatalf("expected edge %s, got %+v", want, graph.Edges)
}
func TestFormatOutgoingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatOutgoingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", TargetFactKey: "target/a"},
})
if hint == "" || hint[0] != ' ' {
t.Fatalf("unexpected hint: %q", hint)
}
}
func TestReplaceIncomingAllowsNotYetCreatedSource(t *testing.T) {
dir := t.TempDir()
db, err := database.NewDB(filepath.Join(dir, "test.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
p, err := db.CreateProject(&database.Project{Name: "parallel-links"})
if err != nil {
t.Fatal(err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "exploit/sqli", Category: "exploit", Summary: "exploit", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
if err := db.ReplaceIncomingProjectFactEdges(p.ID, "exploit/sqli", []database.ProjectFactEdgeFromInput{
{From: "finding/sqli_endpoint", Type: "exploits"},
}); err != nil {
t.Fatalf("incoming edge should not require source fact to exist yet: %v", err)
}
if _, err := db.UpsertProjectFact(&database.ProjectFact{
ProjectID: p.ID, FactKey: "finding/sqli_endpoint", Category: "finding", Summary: "finding", Confidence: "confirmed",
}); err != nil {
t.Fatal(err)
}
in, err := db.ListIncomingProjectFactEdges(p.ID, "exploit/sqli")
if err != nil || len(in) != 1 || in[0].SourceFactKey != "finding/sqli_endpoint" {
t.Fatalf("expected persisted edge from finding, got %+v err=%v", in, err)
}
}
func TestValidateProjectFactEdgeType(t *testing.T) {
t.Parallel()
if err := database.ValidateProjectFactEdgeType("leads_to"); err != nil {
t.Fatal(err)
}
if err := database.ValidateProjectFactEdgeType("invalid"); err == nil {
t.Fatal("expected error")
}
}
+231
View File
@@ -0,0 +1,231 @@
package project
import (
"fmt"
"sort"
"strings"
"cyberstrike-ai/internal/database"
)
var factIndexEdgeTypeOrder = []string{
"discovered_on", "leads_to", "enables", "depends_on", "exploits", "contains", "part_of", "supports",
}
func filterIndexEdges(edges []*database.ProjectFactEdge) []*database.ProjectFactEdge {
if len(edges) == 0 {
return nil
}
out := make([]*database.ProjectFactEdge, 0, len(edges))
for _, e := range edges {
if e == nil {
continue
}
if strings.EqualFold(strings.TrimSpace(e.Confidence), "deprecated") {
continue
}
edgeType := strings.ToLower(strings.TrimSpace(e.EdgeType))
if _, ok := database.ValidProjectFactEdgeTypes[edgeType]; !ok {
continue
}
out = append(out, e)
}
return out
}
func edgeConfidenceSuffix(confidence string) string {
c := strings.ToLower(strings.TrimSpace(confidence))
if c == "" || c == "confirmed" {
return ""
}
return " (" + c + ")"
}
func formatRelationHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s←%s%s", e.EdgeType, e.SourceFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatOutgoingHintPart(e *database.ProjectFactEdge) string {
return fmt.Sprintf("%s→%s%s", e.EdgeType, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence))
}
func formatIncomingHintPart(e *database.ProjectFactEdge) string {
return formatRelationHintPart(e)
}
func joinEdgeHintParts(edges []*database.ProjectFactEdge, formatter func(*database.ProjectFactEdge) string) string {
parts := make([]string, 0, len(edges))
for _, e := range edges {
parts = append(parts, formatter(e))
}
return strings.Join(parts, ", ")
}
// FormatOutgoingLinksHint 黑板索引用出边摘要(全部有效边类型,不截断)。
func FormatOutgoingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {出边: " + joinEdgeHintParts(edges, formatOutgoingHintPart) + "}"
}
// FormatIncomingLinksHint 黑板索引用入边摘要(全部有效边类型,不截断)。
func FormatIncomingLinksHint(edges []*database.ProjectFactEdge) string {
edges = filterIndexEdges(edges)
if len(edges) == 0 {
return ""
}
return " {入边: " + joinEdgeHintParts(edges, formatIncomingHintPart) + "}"
}
// FormatFactIndexLinksHint 黑板索引行内关系边(from → 当前 fact,与 upsert links 一致)。
func FormatFactIndexLinksHint(_ string, incoming []*database.ProjectFactEdge) string {
in := filterIndexEdges(incoming)
if len(in) == 0 {
return ""
}
return " {关系边: " + joinEdgeHintParts(in, formatRelationHintPart) + "}"
}
func indexEdgeGroupMaps(edges []*database.ProjectFactEdge) (outgoing, incoming map[string][]*database.ProjectFactEdge) {
outgoing = map[string][]*database.ProjectFactEdge{}
incoming = map[string][]*database.ProjectFactEdge{}
for _, e := range filterIndexEdges(edges) {
outgoing[e.SourceFactKey] = append(outgoing[e.SourceFactKey], e)
incoming[e.TargetFactKey] = append(incoming[e.TargetFactKey], e)
}
return outgoing, incoming
}
func relationOverviewLine(e *database.ProjectFactEdge) string {
return fmt.Sprintf("- %s → %s%s · %s", e.SourceFactKey, e.TargetFactKey, edgeConfidenceSuffix(e.Confidence), e.EdgeType)
}
func indexEdgeSortKey(e *database.ProjectFactEdge) (int, int, string) {
confRank := 0
if strings.EqualFold(strings.TrimSpace(e.Confidence), "tentative") {
confRank = 1
}
typeRank := len(factIndexEdgeTypeOrder) + 1
for i, t := range factIndexEdgeTypeOrder {
if strings.EqualFold(e.EdgeType, t) {
typeRank = i
break
}
}
return confRank, typeRank, e.SourceFactKey + ">" + e.TargetFactKey + ">" + e.EdgeType
}
func sortIndexOverviewEdges(edges []*database.ProjectFactEdge) {
sort.SliceStable(edges, func(i, j int) bool {
ci, ti, ki := indexEdgeSortKey(edges[i])
cj, tj, kj := indexEdgeSortKey(edges[j])
if ci != cj {
return ci < cj
}
if ti != tj {
return ti < tj
}
return ki < kj
})
}
// BuildFactPathOverviewSection 生成事实关系速览(全部有效边类型,不含 body)。
func BuildFactPathOverviewSection(edges []*database.ProjectFactEdge, indexedKeys map[string]struct{}, maxRunes int) string {
if maxRunes <= 0 {
return ""
}
candidates := filterIndexEdges(edges)
if len(candidates) == 0 {
return ""
}
filtered := make([]*database.ProjectFactEdge, 0, len(candidates))
for _, e := range candidates {
if len(indexedKeys) > 0 {
if _, ok := indexedKeys[e.SourceFactKey]; !ok {
continue
}
if _, ok := indexedKeys[e.TargetFactKey]; !ok {
continue
}
}
filtered = append(filtered, e)
}
if len(filtered) == 0 {
return ""
}
sortIndexOverviewEdges(filtered)
header := "### 攻击路径(事实关系)\n"
header += "source → target · type(与攻击路径图/库中方向一致;写入时在目标 fact 的 links 用 from 声明来源)\n"
var b strings.Builder
b.WriteString(header)
used := len([]rune(header))
omitted := 0
for _, e := range filtered {
line := relationOverviewLine(e) + "\n"
lineRunes := len([]rune(line))
if used+lineRunes > maxRunes {
omitted++
continue
}
b.WriteString(line)
used += lineRunes
}
if omitted > 0 {
extra := fmt.Sprintf("(另有 %d 条关系边未列入,请 get_project_fact 查看完整关系。)\n", omitted)
if used+len([]rune(extra)) <= maxRunes {
b.WriteString(extra)
}
}
if used <= len([]rune(header)) {
return ""
}
return b.String()
}
func factIndexSortPriority(f *database.ProjectFact) int {
if f == nil {
return 0
}
score := 0
if f.Pinned {
score += 1000
}
c := strings.ToLower(strings.TrimSpace(f.Category))
switch c {
case FactCategoryTarget:
score += 400
case FactCategoryFinding, FactCategoryChain:
score += 300
case FactCategoryExploit, FactCategoryPOC:
score += 250
case "auth", "infra", "business":
score += 200
case "note":
score += 50
default:
key := strings.ToLower(strings.TrimSpace(f.FactKey))
if strings.HasPrefix(key, "target/") {
score += 400
} else if strings.HasPrefix(key, "finding/") || strings.HasPrefix(key, "chain/") {
score += 300
}
}
if strings.EqualFold(strings.TrimSpace(f.Confidence), "confirmed") {
score += 80
}
return score
}
func sortFactsForIndex(facts []*database.ProjectFact) {
sort.SliceStable(facts, func(i, j int) bool {
pi, pj := factIndexSortPriority(facts[i]), factIndexSortPriority(facts[j])
if pi != pj {
return pi > pj
}
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
})
}
+161
View File
@@ -0,0 +1,161 @@
package project
import (
"fmt"
"path/filepath"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestFormatIncomingLinksHint(t *testing.T) {
t.Parallel()
hint := FormatIncomingLinksHint([]*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "finding/x", Confidence: "tentative"},
})
if !strings.Contains(hint, "入边:") {
t.Fatalf("expected 入边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←finding/x") {
t.Fatalf("unexpected hint: %q", hint)
}
if !strings.Contains(hint, "tentative") {
t.Fatalf("expected tentative in hint: %q", hint)
}
}
func TestFormatIncomingLinksHint_allEdges(t *testing.T) {
t.Parallel()
edges := make([]*database.ProjectFactEdge, 0, 5)
for i := 1; i <= 5; i++ {
edges = append(edges, &database.ProjectFactEdge{
EdgeType: "discovered_on",
SourceFactKey: fmt.Sprintf("finding/f%d", i),
Confidence: "tentative",
})
}
hint := FormatIncomingLinksHint(edges)
if strings.Contains(hint, "+") {
t.Fatalf("should not truncate with +N: %q", hint)
}
for i := 1; i <= 5; i++ {
if !strings.Contains(hint, fmt.Sprintf("finding/f%d", i)) {
t.Fatalf("missing edge f%d in hint: %q", i, hint)
}
}
}
func TestFormatFactIndexLinksHint_incomingOnly(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", Confidence: "confirmed"},
}
hint := FormatFactIndexLinksHint("finding/sqli", in)
if !strings.Contains(hint, "关系边:") {
t.Fatalf("missing 关系边 label: %q", hint)
}
if !strings.Contains(hint, "discovered_on←target/dev") {
t.Fatalf("missing discovered_on: %q", hint)
}
if !strings.Contains(hint, "exploits←exploit/rce") {
t.Fatalf("missing exploits: %q", hint)
}
if strings.Contains(hint, "出边") || strings.Contains(hint, "入边") {
t.Fatalf("should not use legacy 出边/入边 labels: %q", hint)
}
}
func TestFormatFactIndexLinksHint_includesAuxiliaryEdgeTypes(t *testing.T) {
t.Parallel()
in := []*database.ProjectFactEdge{{EdgeType: "supports", SourceFactKey: "note/log"}}
hint := FormatFactIndexLinksHint("finding/x", in)
if !strings.Contains(hint, "supports←note/log") {
t.Fatalf("supports edge should be included: %q", hint)
}
}
func TestBuildFactPathOverviewSection(t *testing.T) {
t.Parallel()
edges := []*database.ProjectFactEdge{
{EdgeType: "discovered_on", SourceFactKey: "target/dev", TargetFactKey: "finding/sqli", Confidence: "tentative"},
{EdgeType: "exploits", SourceFactKey: "exploit/rce", TargetFactKey: "finding/sqli", Confidence: "confirmed"},
{EdgeType: "supports", SourceFactKey: "note/log", TargetFactKey: "finding/sqli"},
}
keys := map[string]struct{}{
"target/dev": {}, "finding/sqli": {}, "exploit/rce": {}, "note/log": {},
}
section := BuildFactPathOverviewSection(edges, keys, 800)
if !strings.Contains(section, "### 攻击路径(事实关系)") {
t.Fatalf("missing header: %q", section)
}
if !strings.Contains(section, "target/dev → finding/sqli") {
t.Fatalf("missing discovered_on line: %q", section)
}
if !strings.Contains(section, "exploit/rce → finding/sqli") {
t.Fatalf("missing exploits line: %q", section)
}
if !strings.Contains(section, "note/log → finding/sqli") {
t.Fatalf("supports edge should be included: %q", section)
}
}
func TestBuildFactIndexBlock_withLinksAndPathOverview(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "facts.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
proj, err := db.CreateProject(&database.Project{Name: "path-proj"})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "target/dev",
Category: "target",
Summary: "dev 子域",
Confidence: "confirmed",
})
if err != nil {
t.Fatal(err)
}
_, err = db.UpsertProjectFact(&database.ProjectFact{
ProjectID: proj.ID,
FactKey: "finding/sqli",
Category: "finding",
Summary: "时间盲注",
Confidence: "tentative",
})
if err != nil {
t.Fatal(err)
}
_, err = db.AddProjectFactEdge(proj.ID, database.ProjectFactEdgeInput{
To: "finding/sqli",
Type: "discovered_on",
}, "target/dev", "")
if err != nil {
t.Fatal(err)
}
block, err := BuildFactIndexBlock(db, proj.ID, config.ProjectConfig{Enabled: true, FactIndexMaxRunes: 6500, FactIndexPathMaxRunes: 1000})
if err != nil {
t.Fatal(err)
}
if !strings.Contains(block, "关系边: discovered_on←target/dev") {
t.Fatalf("finding line should include relation hint: %q", block)
}
if !strings.Contains(block, "### 攻击路径(事实关系)") {
t.Fatalf("missing relation overview: %q", block)
}
if !strings.Contains(block, "target/dev → finding/sqli") {
t.Fatalf("missing overview edge: %q", block)
}
}
+9 -86
View File
@@ -1,100 +1,23 @@
package project
import (
"strings"
import "cyberstrike-ai/internal/projectprompt"
"cyberstrike-ai/internal/mcp/builtin"
)
// 边渗透边记录:统一节奏文案(agents/*.md 须与 FactRecordingIncrementalRhythmMarkdown 保持一致)。
const (
factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。"
factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。"
factRhythmSubAgentSuffix = "若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。"
)
// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏(Markdown,供 agents/*.md 与文档对齐)。
// FactRecordingIncrementalRhythmMarkdown 见 projectprompt。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:")
b.WriteString(factRhythmCore)
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
return projectprompt.FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent)
}
func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(";与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。")
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块(单/多 Agent 主代理共用)。
// coordinatorDelegate 为 true 时追加「协调者代子代理落库」说明(Deep / plan_execute / supervisor)。
// FactRecordingBlackboardSection 见 projectprompt。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ")
b.WriteString(builtin.ToolGetProjectFact)
b.WriteString("(fact_key) 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞条目):使用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("fact_key 建议 `category/slug`(如 target/primary_domain),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):fact_key 建议 finding/、chain/、exploit/、poc/ 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(",含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ")
b.WriteString(builtin.ToolListVulnerabilities)
b.WriteString(" 查重,详情用 ")
b.WriteString(builtin.ToolGetVulnerability)
b.WriteString("(id)(默认仅当前项目/会话)。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ")
b.WriteString(builtin.ToolDeprecateProjectFact)
b.WriteString(" 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 ")
b.WriteString(builtin.ToolListProjectFacts)
b.WriteString(" / ")
b.WriteString(builtin.ToolSearchProjectFacts)
b.WriteString(" 检索。\n\n")
b.WriteString(FactRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
return projectprompt.FactRecordingBlackboardSection(coordinatorDelegate)
}
// FactRecordingSubAgentSection 子代理边渗透边记录(无工具时输出待落库条目)
// FactRecordingSubAgentSection 见 projectprompt
func FactRecordingSubAgentSection() string {
return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n"
return projectprompt.FactRecordingSubAgentSection()
}
// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown(工具名为字面量,供 agents/*.md
// FactRecordingBlackboardSectionMarkdown 见 projectprompt
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n")
b.WriteString(FactRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
return projectprompt.FactRecordingBlackboardSectionMarkdown(coordinatorDelegate)
}
+5 -10
View File
@@ -3,6 +3,8 @@ package project
import (
"fmt"
"strings"
"cyberstrike-ai/internal/projectprompt"
)
// 事实 category 常量(写入 upsert_project_fact 的 category 字段)。
@@ -90,7 +92,8 @@ const attackChainFactBodyTemplate = `## 结论(可验证,一句话)
## 关联
- related_vulnerability_id: <可选对应 record_vulnerability id>
- 依赖事实: <fact_key auth/session_cookie>
- linksupsert 参数: [{ "from": "<fact_key>", "type": "discovered_on|..." }]from 当前 fact
- 依赖事实body 可读镜像: <fact_key auth/session_cookie>
## 备注与不确定性
<待验证假设环境差异绕过尝试记录>`
@@ -109,15 +112,7 @@ const envFactBodyTemplate = `## 摘要
// FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。
func FactRecordingGuidanceBlock() string {
return `### 事实写入规范审计复现 / 知识沉淀
- **summary**索引用一行须含什么 + 在哪 + 如何触发/验证要点禁止只写结论如仅写存在 SQLi
- **body**完整可复现上下文写入 ` + "`upsert_project_fact`" + ` body 字段索引不含 body后续会话须靠 ` + "`get_project_fact`" + ` 取回
- **category / fact_key 建议**
- 环境认知` + "`target/`" + `` + "`auth/`" + `` + "`infra/`" + `` + "`business/`" + `body 用环境模板即可
- 发现与利用` + "`finding/`" + `` + "`chain/`" + `` + "`exploit/`" + `` + "`poc/`" + `**必须**用攻击链模板填满 body入口逐步攻击链原始请求/响应或命令证据关联漏洞 ID
- **与漏洞记录分工**` + "`record_vulnerability`" + ` 记可交付 findings事实记**复现所需的全部上下文**含失败尝试绕过依赖会话二者可各记一次
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入勿散落多个 key 导致上下文丢失`
return projectprompt.FactRecordingGuidanceBlock()
}
// SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。
+5 -1
View File
@@ -2,10 +2,14 @@ package project
import "strings"
// VisionImageSectionMarker 图片分析 section 标题(与 AppendVisionImageAnalysisIfReady 注入一致)。
const VisionImageSectionMarker = "## 图片分析"
// VisionImageAnalysisSection 单/多代理共用的图片分析提示(analyze_image;上下文仅保留文字摘要)。
func VisionImageAnalysisSection() string {
var b strings.Builder
b.WriteString("## 图片分析\n\n")
b.WriteString(VisionImageSectionMarker)
b.WriteString("\n\n")
b.WriteString("- 遇到图片文件(截图、验证码、登录页、报告配图)时,若存在工具 analyze_image,请传入服务器上的文件路径进行分析。\n")
b.WriteString("- 不要对二进制图片使用 read_file 指望理解内容;用户消息中「📎 xxx.png: /path」即为可传给 analyze_image 的路径。\n")
b.WriteString("- 验证码类:若已从页面或接口保存为本地图片(如 captcha.png),用 analyze_imagequestion 写明「只输出验证码字符」;识别失败则刷新验证码后重新保存再识;复杂滑块/行为验证码勿指望单次识图成功。\n")
+68
View File
@@ -0,0 +1,68 @@
package project
import (
"fmt"
"os"
"path/filepath"
"strings"
)
func sanitizeWorkspacePathSegment(s string) string {
s = strings.TrimSpace(s)
if s == "" {
return "default"
}
s = strings.ReplaceAll(s, string(filepath.Separator), "-")
s = strings.ReplaceAll(s, "/", "-")
s = strings.ReplaceAll(s, "\\", "-")
s = strings.ReplaceAll(s, "..", "__")
if len(s) > 180 {
s = s[:180]
}
return s
}
// WorkspaceRootDir returns the relative workspace root for downloads and local analysis.
// Project-bound sessions share projects/<id>/; otherwise conversations/<id>/.
func WorkspaceRootDir(configuredBase, projectID, conversationID string) string {
base := strings.TrimSpace(configuredBase)
if base == "" {
base = filepath.Join("tmp", "workspace")
}
if pid := strings.TrimSpace(projectID); pid != "" {
return filepath.Join(base, "projects", sanitizeWorkspacePathSegment(pid))
}
conv := strings.TrimSpace(conversationID)
if conv == "" {
conv = "default"
}
return filepath.Join(base, "conversations", sanitizeWorkspacePathSegment(conv))
}
// EnsureWorkspace creates the workspace directory and returns its absolute path.
func EnsureWorkspace(root string) (string, error) {
abs, err := filepath.Abs(strings.TrimSpace(root))
if err != nil {
return "", fmt.Errorf("workspace abs: %w", err)
}
if err := os.MkdirAll(abs, 0o755); err != nil {
return "", fmt.Errorf("workspace mkdir: %w", err)
}
return abs, nil
}
// BuildWorkspaceBlock instructs the agent to use the session workspace instead of /tmp.
func BuildWorkspaceBlock(absPath string) string {
absPath = strings.TrimSpace(absPath)
if absPath == "" {
return ""
}
return fmt.Sprintf(`## 会话工作目录下载与本地分析
**必须使用以下目录**保存 curl/wget 下载的文件临时 HTML/JS以及 read_file/glob/grep 的检索范围
`+"`%s`"+`
- **禁止**使用系统 `+"`/tmp`"+` 或其它全局临时目录多项目/多会话会互窜遗留文件
- 下载示例`+"`curl -o '%s/page.html' 'https://target/'`"+`exec 时可将 `+"`workdir`"+` 设为该目录。
- 读取前用 glob/grep/read_file **限定在该目录**下搜索勿在 `+"`/tmp`"+` 盲目检索`, absPath, absPath)
}
+52
View File
@@ -0,0 +1,52 @@
package project
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestWorkspaceRootDirProjectScoped(t *testing.T) {
got := WorkspaceRootDir("", "proj-1", "conv-1")
want := filepath.Join("tmp", "workspace", "projects", "proj-1")
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestWorkspaceRootDirConversationScoped(t *testing.T) {
got := WorkspaceRootDir("/data/ws", "", "conv-abc")
want := filepath.Join("/data/ws", "conversations", "conv-abc")
if got != want {
t.Fatalf("got %q want %q", got, want)
}
}
func TestEnsureWorkspaceCreatesDir(t *testing.T) {
root := filepath.Join(t.TempDir(), "nested", "workspace")
abs, err := EnsureWorkspace(root)
if err != nil {
t.Fatalf("EnsureWorkspace: %v", err)
}
st, err := os.Stat(abs)
if err != nil {
t.Fatalf("Stat: %v", err)
}
if !st.IsDir() {
t.Fatal("expected directory")
}
}
func TestBuildWorkspaceBlockMentionsPath(t *testing.T) {
block := BuildWorkspaceBlock("/opt/csai/tmp/workspace/projects/p1")
if block == "" {
t.Fatal("expected non-empty block")
}
if !strings.Contains(block, "/opt/csai/tmp/workspace/projects/p1") {
t.Fatalf("block missing path: %s", block)
}
if !strings.Contains(block, "/tmp") {
t.Fatalf("block should warn about /tmp: %s", block)
}
}
+132
View File
@@ -0,0 +1,132 @@
// Package projectprompt 提供项目黑板相关的系统提示文本(纯字符串,无 database 依赖)。
// 供 agent / multiagent 等包引用,避免 agent → project 导入环导致 gopls 元数据失败。
package projectprompt
import (
"strings"
"cyberstrike-ai/internal/mcp/builtin"
)
const (
factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 `upsert_project_fact`(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 `record_vulnerability`;与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。"
factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时,由协调者及时写入,勿假定子代理已记。"
factRhythmSubAgentSuffix = "若工具集中无上述工具,须在交付物末尾给出「待落库」结构化条目(fact_key 建议、summary、body/POC 要点),供协调者**立即**写入。"
)
// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏(Markdown,供 agents/*.md 与文档对齐)。
func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:")
b.WriteString(factRhythmCore)
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string {
var b strings.Builder
b.WriteString("- **边渗透边记录(强制节奏)**:勿等会话结束或收尾再批量写入。每**确认**一条新认知(开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化)后,**立即**调用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("(同 fact_key 覆盖更新)。每**验证**出一条可复现漏洞(含 POC/影响)后,**立即**调用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(";与事实可各记一次。继续下一步工作前优先落库,避免上下文压缩后细节丢失。未绑项目时说明无法写黑板,仍在本轮保留证据摘要。")
if coordinator {
b.WriteString(factRhythmCoordinatorSuffix)
}
if subAgent {
b.WriteString(factRhythmSubAgentSuffix)
}
return b.String()
}
func factEdgeRecordingGuidance() string {
return `### 事实关系边links
- 写入 **finding / chain / exploit / poc** **必须** ` + "`upsert_project_fact`" + ` 中提供 ` + "`links`" + `**推荐 ` + "`from`" + `**来源 fact 指向当前 fact ` + "`from`" + ` 当前 ` + "`fact_key`" + `
- **最少要求**finding 类至少 1 from=target/* + type=discovered_on target finding finding 上记录 exploit from=exploit/* + type=exploits exploit finding
- **常用 type**` + "`discovered_on`" + `发现在哪` + "`depends_on`" + `复现前置` + "`leads_to`" + `认知推进` + "`enables`" + `扩大攻击面` + "`exploits`" + `利用关系` + "`contains`" + `资产包含` + "`part_of`" + `属于链/` + "`supports`" + `证据支撑
- 更新时**省略 links 保留已有边**传入 links **替换**全部关系边from 当前 fact
- body 依赖事实段落可与 links 并存人读结构化关系以 links 为准`
}
func factRecordingGuidanceBlock() string {
return `### 事实写入规范审计复现 / 知识沉淀
- **summary**索引用一行须含什么 + 在哪 + 如何触发/验证要点禁止只写结论如仅写存在 SQLi
- **body**完整可复现上下文写入 ` + "`upsert_project_fact`" + ` body 字段索引不含 body后续会话须靠 ` + "`get_project_fact`" + ` 取回
- **category / fact_key 建议**
- 环境认知` + "`target/`" + `` + "`auth/`" + `` + "`infra/`" + `` + "`business/`" + `body 用环境模板即可
- 发现与利用` + "`finding/`" + `` + "`chain/`" + `` + "`exploit/`" + `` + "`poc/`" + `**必须**用攻击链模板填满 body入口逐步攻击链原始请求/响应或命令证据关联漏洞 ID
- **与漏洞记录分工**` + "`record_vulnerability`" + ` 记可交付 findings事实记**复现所需的全部上下文**含失败尝试绕过依赖会话二者可各记一次
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入勿散落多个 key 导致上下文丢失`
}
// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块(单/多 Agent 主代理共用)。
func FactRecordingBlackboardSection(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ")
b.WriteString(builtin.ToolGetProjectFact)
b.WriteString("(fact_key) 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞条目):使用 ")
b.WriteString(builtin.ToolUpsertProjectFact)
b.WriteString("fact_key 建议 `category/slug`(如 target/primary_domain),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):fact_key 建议 finding/、chain/、exploit/、poc/ 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 ")
b.WriteString(builtin.ToolRecordVulnerability)
b.WriteString(",含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ")
b.WriteString(builtin.ToolListVulnerabilities)
b.WriteString(" 查重,详情用 ")
b.WriteString(builtin.ToolGetVulnerability)
b.WriteString("(id)(默认仅当前项目/会话)。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ")
b.WriteString(builtin.ToolDeprecateProjectFact)
b.WriteString(" 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 ")
b.WriteString(builtin.ToolListProjectFacts)
b.WriteString(" / ")
b.WriteString(builtin.ToolSearchProjectFacts)
b.WriteString(" 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactRecordingSubAgentSection 子代理边渗透边记录(无工具时输出待落库条目)。
func FactRecordingSubAgentSection() string {
return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n"
}
// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown(工具名为字面量,供 agents/*.md)。
func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string {
var b strings.Builder
b.WriteString("## 项目黑板(事实)与漏洞记录(分离)\n\n")
b.WriteString("当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**\n\n")
b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false))
b.WriteString("\n\n")
b.WriteString("- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。\n")
b.WriteString("- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**summary 写「什么 + 在哪 + 如何验证」一行要点。\n")
b.WriteString("- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。\n")
b.WriteString("- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n")
b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n")
b.WriteString(factEdgeRecordingGuidance())
b.WriteString("\n\n")
b.WriteString(factRecordingGuidanceBlock())
b.WriteString("\n\n严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。")
return b.String()
}
// FactEdgeRecordingGuidance 写入边时的 Agent 规范(供 project 包复用)。
func FactEdgeRecordingGuidance() string { return factEdgeRecordingGuidance() }
// FactRecordingGuidanceBlock 事实写入规范块(供 project 包复用)。
func FactRecordingGuidanceBlock() string { return factRecordingGuidanceBlock() }
+11
View File
@@ -0,0 +1,11 @@
package projectprompt
// ShellExecExecuteGuidanceSection 供单代理/多代理系统提示追加:exec 与 execute 分工(尽量短)。
func ShellExecExecuteGuidanceSection() string {
return `Shellexec/execute):有专用 MCP 工具时优先专用工具;系统命令(管道、workdir、后台 &)用 execskills/ 内脚本(配合 read_file、skill)用 execute;多步扫描分拆调用,禁止一条 shell 串多个扫描器。下载/临时文件须写入系统提示中的「会话工作目录」,禁止用 /tmp。`
}
// ShellExecExecuteGuidanceReconSuffix 侦察子代理可选追加(一行)。
func ShellExecExecuteGuidanceReconSuffix() string {
return `枚举优先 subfinder、amass 等专用 MCP,勿 exec/execute 拼长链。`
}
+46 -16
View File
@@ -84,8 +84,9 @@ func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.Open
}
}
// applyClaudeExtendedThinking sets Anthropic Messages API `thinking` when absent from ExtraRequestFields.
// Uses adaptive + summarized display by default (per Anthropic guidance for Claude 4.x); Sonnet 3.7 uses enabled+budget.
// applyClaudeExtendedThinking sets Anthropic Messages API fields per official guidance:
// - Adaptive models (4.6+): thinking.type=adaptive; output_config.effort only when user sets effort (API default is high).
// - Sonnet 3.7: thinking.type=enabled + budget_tokens=10000 (doc example); effort is not mapped — use extra_request_fields for custom budget.
func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) {
if cfg == nil || mode == "off" {
return
@@ -93,31 +94,60 @@ func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort,
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
if _, exists := cfg.ExtraFields["thinking"]; exists {
return
}
m := strings.ToLower(strings.TrimSpace(model))
thinking := map[string]any{
"type": "adaptive",
"display": "summarized",
sonnet37 := isClaudeSonnet37(m)
if _, exists := cfg.ExtraFields["thinking"]; !exists {
cfg.ExtraFields["thinking"] = claudeThinkingForModel(m, sonnet37)
}
// Sonnet 3.7: manual extended thinking is the documented path.
if strings.Contains(m, "claude-3-7-sonnet") || strings.Contains(m, "3-7-sonnet") || strings.Contains(m, "sonnet-3.7") {
thinking = map[string]any{
applyClaudeOutputConfigEffort(cfg, effort, sonnet37)
}
// claudeSonnet37DefaultBudgetTokens matches Anthropic extended-thinking documentation examples (budget_tokens with max_tokens 16000).
const claudeSonnet37DefaultBudgetTokens = 10000
func isClaudeSonnet37(m string) bool {
return strings.Contains(m, "claude-3-7-sonnet") ||
strings.Contains(m, "3-7-sonnet") ||
strings.Contains(m, "sonnet-3.7")
}
func claudeThinkingForModel(m string, sonnet37 bool) map[string]any {
if sonnet37 {
return map[string]any{
"type": "enabled",
"budget_tokens": 10000,
"budget_tokens": claudeSonnet37DefaultBudgetTokens,
"display": "summarized",
}
}
// Opus 4.7+: manual enabled+budget rejected — keep adaptive only.
// Opus 4.7+: manual enabled+budget rejected — adaptive only.
if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") {
thinking = map[string]any{
return map[string]any{
"type": "adaptive",
"display": "summarized",
}
}
_ = effort // reserved: map to Anthropic effort / output_config when API stabilizes in one place
cfg.ExtraFields["thinking"] = thinking
return map[string]any{
"type": "adaptive",
"display": "summarized",
}
}
// applyClaudeOutputConfigEffort sets top-level output_config.effort only when effort is explicitly configured.
// Omitted effort uses the API default (high); do not inject effort on mode:on alone.
func applyClaudeOutputConfigEffort(cfg *einoopenai.ChatModelConfig, effort string, sonnet37 bool) {
if cfg == nil || sonnet37 {
return
}
if _, exists := cfg.ExtraFields["output_config"]; exists {
return
}
e := effortStringForAPI(effort)
if e == "" {
return
}
cfg.ExtraFields["output_config"] = map[string]any{"effort": e}
}
func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string {
+77
View File
@@ -80,3 +80,80 @@ func TestApplyOpenAICompat_maxPassthrough(t *testing.T) {
t.Fatalf("max effort wire=%q, want max", got)
}
}
func TestApplyClaude_adaptiveOutputConfigEffort(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-opus-4-8",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
Effort: "xhigh",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "adaptive" {
t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"])
}
oc, ok := cfg.ExtraFields["output_config"].(map[string]any)
if !ok {
t.Fatal("expected output_config")
}
if oc["effort"] != "xhigh" {
t.Fatalf("effort=%v", oc["effort"])
}
}
func TestApplyClaude_sonnet37OfficialBudget(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-3-7-sonnet-latest",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
Effort: "low", // 3.7 has no output_config.effort; effort is not mapped to budget_tokens
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "enabled" {
t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"])
}
if th["budget_tokens"] != claudeSonnet37DefaultBudgetTokens {
t.Fatalf("budget_tokens=%v, want official example %d", th["budget_tokens"], claudeSonnet37DefaultBudgetTokens)
}
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("sonnet 3.7 should not set output_config")
}
}
func TestApplyClaude_onWithoutEffortOmitsOutputConfig(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-sonnet-4-6",
Reasoning: config.OpenAIReasoningConfig{
Mode: "on",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("on without explicit effort should omit output_config (API default high)")
}
}
func TestApplyClaude_autoWithoutEffortSkipsOutputConfig(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{
Provider: "claude",
Model: "claude-sonnet-4-6",
Reasoning: config.OpenAIReasoningConfig{
Mode: "auto",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if _, hasOC := cfg.ExtraFields["output_config"]; hasOC {
t.Fatal("auto without effort should omit output_config")
}
}
@@ -0,0 +1,56 @@
package security
import (
"errors"
"fmt"
"os/exec"
"strings"
)
// FormatCommandFailureResult 与 exec 工具 ToolResult 文案一致(不含 ToolErrorPrefix)。
func FormatCommandFailureResult(exitCode int, output string) string {
output = strings.TrimSpace(output)
errMsg := fmt.Sprintf("exit status %d", exitCode)
if output == "" {
return fmt.Sprintf("命令执行失败: %s", errMsg)
}
if strings.HasPrefix(output, "命令执行失败:") {
return output
}
return fmt.Sprintf("命令执行失败: %s\n输出: %s", errMsg, output)
}
// FormatCommandFailureFromErr 根据 exec/execute 返回的 error 生成统一失败文案(IsError 正文)。
func FormatCommandFailureFromErr(err error, output string) string {
if err == nil {
return strings.TrimSpace(output)
}
var exitError *exec.ExitError
if errors.As(err, &exitError) {
return FormatCommandFailureResult(exitError.ExitCode(), output)
}
output = strings.TrimSpace(output)
if output == "" {
return fmt.Sprintf("命令执行失败: %v", err)
}
if strings.HasPrefix(output, "命令执行失败:") {
return output
}
return fmt.Sprintf("命令执行失败: %v\n输出: %s", err, output)
}
// ExecuteFailureStatusLine 流式 execute 结束时追加的单行状态(输出正文已在流中推送过)。
func ExecuteFailureStatusLine(exitCode int) string {
return fmt.Sprintf("\n命令执行失败: exit status %d", exitCode)
}
// IsCommandFailureResult 判断工具结果正文是否表示命令非零退出(用于 execute / exec 对齐 isError)。
func IsCommandFailureResult(content string) bool {
return strings.Contains(content, "命令执行失败:")
}
// IsLegacyShellExitNoise 过滤旧版 shell 流中冗余的 exit code 行。
func IsLegacyShellExitNoise(s string) bool {
trimmed := strings.TrimSpace(s)
return strings.HasPrefix(trimmed, "command exited with non-zero code ")
}
@@ -0,0 +1,54 @@
package security
import (
"errors"
"os/exec"
"strings"
"testing"
)
func TestFormatCommandFailureResult(t *testing.T) {
got := FormatCommandFailureResult(1, "sudo: password required")
want := "命令执行失败: exit status 1\n输出: sudo: password required"
if got != want {
t.Fatalf("got %q want %q", got, want)
}
if FormatCommandFailureResult(2, "") != "命令执行失败: exit status 2" {
t.Fatal("empty output format")
}
if FormatCommandFailureResult(1, "命令执行失败: exit status 1") != "命令执行失败: exit status 1" {
t.Fatal("should not double-wrap")
}
}
func TestIsCommandFailureResult(t *testing.T) {
if !IsCommandFailureResult("sudo: err\n命令执行失败: exit status 1") {
t.Fatal("expected true")
}
if IsCommandFailureResult("sudo: err only") {
t.Fatal("expected false")
}
}
func TestFormatCommandFailureFromErr(t *testing.T) {
cmd := exec.Command("sh", "-c", "exit 42")
err := cmd.Run()
got := FormatCommandFailureFromErr(err, "oops")
if got != "命令执行失败: exit status 42\n输出: oops" {
t.Fatalf("got %q", got)
}
timeoutErr := errors.New("shell inactivity timeout (300s)")
got2 := FormatCommandFailureFromErr(timeoutErr, "already timed out")
if !strings.Contains(got2, "shell inactivity timeout") || !strings.Contains(got2, "already timed out") {
t.Fatalf("got %q", got2)
}
}
func TestIsLegacyShellExitNoise(t *testing.T) {
if !IsLegacyShellExitNoise("command exited with non-zero code 1\n") {
t.Fatal("expected legacy noise")
}
if IsLegacyShellExitNoise("sudo: failed") {
t.Fatal("unexpected noise")
}
}
+60 -23
View File
@@ -32,10 +32,11 @@ var ToolOutputCallbackCtxKey = toolOutputCallbackCtxKey{}
// Executor 安全工具执行器
type Executor struct {
config *config.SecurityConfig
toolIndex map[string]*config.ToolConfig // 工具索引,用于 O(1) 查找
mcpServer *mcp.Server
logger *zap.Logger
config *config.SecurityConfig
toolIndex map[string]*config.ToolConfig // 工具索引,用于 O(1) 查找
mcpServer *mcp.Server
logger *zap.Logger
shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300-1=关闭(见 SetShellNoOutputTimeoutSeconds
}
// NewExecutor 创建新的执行器
@@ -51,6 +52,11 @@ func NewExecutor(cfg *config.SecurityConfig, mcpServer *mcp.Server, logger *zap.
return executor
}
// SetShellNoOutputTimeoutSeconds 配置 exec 工具无输出空闲终止(与 agent.shell_no_output_timeout_seconds 一致)。
func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) {
e.shellNoOutputTimeoutSec = sec
}
// buildToolIndex 构建工具索引,将 O(n) 查找优化为 O(1)
func (e *Executor) buildToolIndex() {
e.toolIndex = make(map[string]*config.ToolConfig)
@@ -133,6 +139,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
// 执行命令
cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
_ = prepareShellCmdSession(cmd)
e.logger.Info("执行安全工具",
@@ -144,7 +151,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
var err error
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(ctx, cmd, cb)
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
zap.String("tool", toolName),
@@ -797,6 +804,8 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
zap.String("command", command),
)
command = PrepareNonInteractiveShellCommand(command)
// 获取shell类型(可选,默认为sh)
shell := "sh"
if s, ok := args["shell"].(string); ok && s != "" {
@@ -820,8 +829,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
} else {
cmd = exec.CommandContext(ctx, shell, "-c", command)
}
applyDefaultTerminalEnv(cmd)
_ = prepareShellCmdSession(cmd)
ConfigureShellCmdForAgentExecute(cmd)
// 执行命令
e.logger.Info("执行系统命令",
@@ -850,8 +858,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
} else {
pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand)
}
applyDefaultTerminalEnv(pidCmd)
_ = prepareShellCmdSession(pidCmd)
ConfigureShellCmdForAgentExecute(pidCmd)
// 获取stdout管道
stdout, err := pidCmd.StdoutPipe()
@@ -963,15 +970,14 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
var err error
// 若上层提供工具输出增量回调,则边执行边流式读取。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(ctx, cmd, cb)
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
if workDir != "" {
cmd2.Dir = workDir
}
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
ConfigureShellCmdForAgentExecute(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, cb)
}
} else {
@@ -984,8 +990,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
if workDir != "" {
cmd2.Dir = workDir
}
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
ConfigureShellCmdForAgentExecute(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, nil)
}
}
@@ -999,7 +1004,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
Content: []mcp.Content{
{
Type: "text",
Text: fmt.Sprintf("命令执行失败: %v\n输出: %s", err, string(output)),
Text: FormatCommandFailureFromErr(err, output),
},
},
IsError: true,
@@ -1024,7 +1029,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int) (string, error) {
if err := prepareShellCmdSession(cmd); err != nil {
return "", err
}
@@ -1091,12 +1096,43 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
lastFlush = time.Now()
}
for chunk := range chunks {
outBuilder.WriteString(chunk)
deltaBuilder.WriteString(chunk)
// 简单节流:buffer 大于 2KB 或 200ms 就刷新一次
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
flush()
idleWatch := NewShellInactivityWatch(noOutputSec)
if idleWatch != nil {
defer idleWatch.Stop()
}
fireInactivity := func() {
terminateCmdTree(cmd)
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
outBuilder.WriteString(msg)
if cb != nil {
cb(msg)
}
_ = cmd.Wait()
}
chunksLoop:
for {
var idleCh <-chan struct{}
if idleWatch != nil {
idleCh = idleWatch.Expired
}
select {
case <-idleCh:
fireInactivity()
return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
case chunk, ok := <-chunks:
if !ok {
break chunksLoop
}
if chunk != "" && idleWatch != nil {
idleWatch.Bump()
}
outBuilder.WriteString(chunk)
deltaBuilder.WriteString(chunk)
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
flush()
}
}
}
flush()
@@ -1116,6 +1152,7 @@ func applyDefaultTerminalEnv(cmd *exec.Cmd) {
if cmd.Env == nil {
cmd.Env = os.Environ()
}
cmd.Env = ApplyNonInteractivePagerEnv(cmd.Env)
// 如果用户已设置 TERM/COLUMNS/LINES,则不覆盖
has := func(k string) bool {
prefix := k + "="
@@ -1159,7 +1196,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
if runtime.GOOS == "windows" {
// PTY 方案为类 UnixWindows 走原逻辑
if cb != nil {
return streamCommandOutput(ctx, cmd, cb)
return streamCommandOutput(ctx, cmd, cb, 0)
}
_ = prepareShellCmdSession(cmd)
out, err := cmd.CombinedOutput()
+21
View File
@@ -71,6 +71,27 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
}
}
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
executor, _ := setupTestExecutor(t)
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
"command": "echo fail-msg >&2; exit 7",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
if res == nil || !res.IsError {
t.Fatalf("expected IsError, got %+v", res)
}
text := res.Content[0].Text
if text != FormatCommandFailureResult(7, "fail-msg\n") && text != FormatCommandFailureResult(7, "fail-msg") {
t.Fatalf("unexpected failure text: %q", text)
}
if !strings.Contains(text, "exit status 7") || !strings.Contains(text, "fail-msg") {
t.Fatalf("unexpected failure text: %q", text)
}
}
func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
pos1 := 1
executor, _ := setupTestExecutor(t)
+200
View File
@@ -0,0 +1,200 @@
package security
import (
"context"
"errors"
"fmt"
"io"
"os/exec"
"sync"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/schema"
)
// ConfigureShellCmdForAgentExecute 与 exec 工具一致:非交互 stdin、pager/TERM 环境、独立进程组。
func ConfigureShellCmdForAgentExecute(cmd *exec.Cmd) {
if cmd == nil {
return
}
applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
_ = prepareShellCmdSession(cmd)
}
// TerminateShellCmdTree 尽力终止 shell 及其子进程组(与 exec/execute 超时取消一致)。
func TerminateShellCmdTree(cmd *exec.Cmd) {
terminateCmdTree(cmd)
}
// EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell,行为与 exec 对齐:
// 并发读取 stdout/stderr(定长块,非按行),避免官方 local.ExecuteStreaming 先排空 stdout
// 导致 stderr 错误(如 sudo 密码提示)长时间不可见、UI 一直显示「执行中」。
type EinoStreamingShell struct{}
// NewEinoStreamingShell 创建 execute 流式 shell 实现。
func NewEinoStreamingShell() *EinoStreamingShell {
return &EinoStreamingShell{}
}
// ExecuteStreaming 实现 filesystem.StreamingShell。
func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
if input == nil || input.Command == "" {
return nil, fmt.Errorf("command is required")
}
sr, w := schema.Pipe[*filesystem.ExecuteResponse](100)
if input.RunInBackendGround {
go runShellInBackground(ctx, input.Command, w)
return sr, nil
}
go streamShellForeground(ctx, input.Command, w)
return sr, nil
}
func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close()
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
ConfigureShellCmdForAgentExecute(cmd)
stdout, err := cmd.StdoutPipe()
if err != nil {
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
return
}
stderr, err := cmd.StderrPipe()
if err != nil {
_ = stdout.Close()
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return
}
if err := cmd.Start(); err != nil {
_ = stdout.Close()
_ = stderr.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
return
}
done := make(chan struct{})
go func() {
drainShellPipes(stdout, stderr)
_ = cmd.Wait()
close(done)
}()
select {
case <-done:
case <-ctx.Done():
TerminateShellCmdTree(cmd)
}
exitCode := 0
_ = w.Send(&filesystem.ExecuteResponse{
Output: "command started in background\n",
ExitCode: &exitCode,
}, nil)
}
func drainShellPipes(stdout, stderr io.Reader) {
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
_, _ = io.Copy(io.Discard, stdout)
}()
go func() {
defer wg.Done()
_, _ = io.Copy(io.Discard, stderr)
}()
wg.Wait()
}
func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close()
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
ConfigureShellCmdForAgentExecute(cmd)
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
return
}
stderrPipe, err := cmd.StderrPipe()
if err != nil {
_ = stdoutPipe.Close()
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return
}
if err := cmd.Start(); err != nil {
_ = stdoutPipe.Close()
_ = stderrPipe.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
return
}
stopWatch := make(chan struct{})
go func() {
select {
case <-ctx.Done():
TerminateShellCmdTree(cmd)
case <-stopWatch:
}
}()
defer close(stopWatch)
chunks := make(chan string, 64)
var wg sync.WaitGroup
readFn := func(r io.Reader) {
defer wg.Done()
buf := make([]byte, 8192)
for {
n, readErr := r.Read(buf)
if n > 0 {
chunks <- string(buf[:n])
}
if readErr != nil {
return
}
}
}
wg.Add(2)
go readFn(stdoutPipe)
go readFn(stderrPipe)
go func() {
wg.Wait()
close(chunks)
}()
hadOutput := false
for chunk := range chunks {
if chunk == "" {
continue
}
hadOutput = true
if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) {
TerminateShellCmdTree(cmd)
return
}
}
waitErr := cmd.Wait()
if waitErr == nil {
exitCode := 0
_ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil)
return
}
var exitError *exec.ExitError
if errors.As(waitErr, &exitError) {
exitCode := exitError.ExitCode()
resp := &filesystem.ExecuteResponse{ExitCode: &exitCode}
if !hadOutput {
resp.Output = FormatCommandFailureResult(exitCode, "")
}
_ = w.Send(resp, nil)
return
}
_ = w.Send(nil, fmt.Errorf("command failed: %w", waitErr))
}
@@ -0,0 +1,117 @@
package security
import (
"context"
"errors"
"io"
"strings"
"testing"
"time"
"github.com/cloudwego/eino/adk/filesystem"
)
func TestEinoStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) {
shell := NewEinoStreamingShell()
cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1")
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", time.Since(start))
}
if !strings.Contains(got.String(), "err-only") {
t.Fatalf("expected stderr in output, got: %q", got.String())
}
}
func TestEinoStreamingShell_SudoFailsFast(t *testing.T) {
shell := NewEinoStreamingShell()
cmd := PrepareNonInteractiveShellCommand("sudo whoami && sudo cat /etc/os-release")
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 {
continue
}
got.WriteString(resp.Output)
}
if time.Since(start) > 5*time.Second {
t.Fatalf("sudo should fail quickly, took %v output=%q", time.Since(start), got.String())
}
out := got.String()
if strings.Contains(out, "command exited with non-zero code") {
t.Fatalf("legacy exit line present: %q", out)
}
if !strings.Contains(out, "sudo") && !strings.Contains(out, "password") && !strings.Contains(out, "terminal") {
t.Fatalf("expected sudo error text, got: %q", out)
}
}
func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) {
shell := NewEinoStreamingShell()
// 模拟 sudostderr 先有输出,stdout 侧进程仍挂起;旧 eino local 在首包 stderr 前不会向流写任何内容。
cmd := PrepareNonInteractiveShellCommand(`echo "password prompt" >&2; sleep 30`)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
sr, err := shell.ExecuteStreaming(ctx, &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 {
break
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
if strings.Contains(got.String(), "password prompt") {
break
}
}
}
if time.Since(start) > 1500*time.Millisecond {
t.Fatalf("expected stderr promptly, took %v output=%q", time.Since(start), got.String())
}
if !strings.Contains(got.String(), "password prompt") {
t.Fatalf("expected early stderr, got: %q", got.String())
}
}
+163
View File
@@ -0,0 +1,163 @@
package security
import (
"fmt"
"os"
"os/exec"
"strings"
"sync"
"time"
)
// ShellNoOutputTimeoutMessage 长时间无新 stdout/stderr 时的提示(软失败,模型可见)。
func ShellNoOutputTimeoutMessage(idleSec int) string {
return fmt.Sprintf(`命令已终止超过 %d 秒没有新的输出疑似在等待交互输入或已挂起
长时静默任务请使用末尾 & 后台运行或增大 agent.shell_no_output_timeout_seconds-1=关闭此检测
Command terminated: no new output for %d seconds (possible interactive wait or hung process).`, idleSec, idleSec)
}
// ShellInactivityWatch 在 noOutputSec 内无任何新输出时向 expired 发送信号;每次 Bump 重置计时。
// 与「仅有首包输出就永久取消计时」不同,可兜住 sudo 打印 Password 提示后继续挂起等情况。
type ShellInactivityWatch struct {
Sec int
mu sync.Mutex
timer *time.Timer
Expired chan struct{}
}
func NewShellInactivityWatch(noOutputSec int) *ShellInactivityWatch {
sec := ResolveShellNoOutputTimeoutSeconds(noOutputSec)
if sec <= 0 {
return nil
}
w := &ShellInactivityWatch{
Sec: sec,
Expired: make(chan struct{}, 1),
}
w.Bump()
return w
}
func (w *ShellInactivityWatch) Bump() {
if w == nil || w.Sec <= 0 {
return
}
w.mu.Lock()
defer w.mu.Unlock()
if w.timer != nil {
w.timer.Stop()
}
w.timer = time.AfterFunc(time.Duration(w.Sec)*time.Second, func() {
select {
case w.Expired <- struct{}{}:
default:
}
})
}
func (w *ShellInactivityWatch) Stop() {
if w == nil {
return
}
w.mu.Lock()
defer w.mu.Unlock()
if w.timer != nil {
w.timer.Stop()
w.timer = nil
}
}
// ResolveShellNoOutputTimeoutSeconds0=默认 3005 分钟);-1=关闭;>0=自定义。
func ResolveShellNoOutputTimeoutSeconds(sec int) int {
if sec < 0 {
return 0
}
if sec == 0 {
return 300
}
return sec
}
// PrependNonInteractiveShellExports 为 sh -c 注入通用非交互环境(pager 等),不维护命令黑名单。
func PrependNonInteractiveShellExports(shellCommand string) string {
if strings.TrimSpace(shellCommand) == "" {
return shellCommand
}
upper := strings.ToUpper(shellCommand)
var pairs []string
add := func(key, val string) {
if strings.Contains(upper, strings.ToUpper(key)) {
return
}
pairs = append(pairs, key+"="+val)
}
add("GIT_PAGER", "cat")
add("PAGER", "cat")
add("SYSTEMD_PAGER", "cat")
add("DEBIAN_FRONTEND", "noninteractive")
if len(pairs) == 0 {
return shellCommand
}
return "export " + strings.Join(pairs, " ") + "\n" + shellCommand
}
// PrependNonInteractiveStdinRedirect 为 sh -c 关闭 stdin(与 attachNonInteractiveStdin 等价),
// 使 read/input()/sudo -S 等从 stdin 读取的程序快速失败而非挂起。已含 </dev/null 时不重复注入。
func PrependNonInteractiveStdinRedirect(shellCommand string) string {
if strings.TrimSpace(shellCommand) == "" {
return shellCommand
}
lower := strings.ToLower(shellCommand)
if strings.Contains(lower, "</dev/null") || strings.Contains(lower, "0</dev/null") {
return shellCommand
}
return "exec </dev/null\n" + shellCommand
}
// PrepareNonInteractiveShellCommand 组合非交互包装:stdin 关闭 + pager 等环境变量(零名单)。
func PrepareNonInteractiveShellCommand(shellCommand string) string {
return PrependNonInteractiveStdinRedirect(PrependNonInteractiveShellExports(shellCommand))
}
// ApplyNonInteractivePagerEnv 为 exec.Cmd 补齐与 PrependNonInteractiveShellExports 一致的环境变量。
func ApplyNonInteractivePagerEnv(cmdEnv []string) []string {
if cmdEnv == nil {
cmdEnv = []string{}
}
has := func(k string) bool {
prefix := k + "="
for _, e := range cmdEnv {
if strings.HasPrefix(e, prefix) {
return true
}
}
return false
}
if !has("GIT_PAGER") {
cmdEnv = append(cmdEnv, "GIT_PAGER=cat")
}
if !has("PAGER") {
cmdEnv = append(cmdEnv, "PAGER=cat")
}
if !has("SYSTEMD_PAGER") {
cmdEnv = append(cmdEnv, "SYSTEMD_PAGER=cat")
}
if !has("DEBIAN_FRONTEND") {
cmdEnv = append(cmdEnv, "DEBIAN_FRONTEND=noninteractive")
}
return cmdEnv
}
// attachNonInteractiveStdin 关闭交互式 stdin,使部分命令快速失败而非等待输入。
func attachNonInteractiveStdin(cmd *exec.Cmd) {
if cmd == nil || cmd.Stdin != nil {
return
}
f, err := os.Open(os.DevNull)
if err != nil {
return
}
cmd.Stdin = f
}
@@ -0,0 +1,128 @@
package security
import (
"context"
"os"
"os/exec"
"strings"
"testing"
"time"
)
func TestPrependNonInteractiveShellExports(t *testing.T) {
out := PrependNonInteractiveShellExports("echo hi")
if !strings.Contains(out, "GIT_PAGER=cat") || !strings.Contains(out, "PAGER=cat") {
t.Fatalf("missing pager exports: %q", out)
}
if !strings.HasSuffix(strings.TrimSpace(out), "echo hi") {
t.Fatalf("command suffix lost: %q", out)
}
skip := PrependNonInteractiveShellExports("GIT_PAGER=less echo hi")
if strings.Contains(skip, "export GIT_PAGER=cat") {
t.Fatalf("should not override existing GIT_PAGER: %q", skip)
}
}
func TestPrependNonInteractiveStdinRedirect(t *testing.T) {
out := PrependNonInteractiveStdinRedirect("echo hi")
if !strings.HasPrefix(out, "exec </dev/null") {
t.Fatalf("missing stdin redirect: %q", out)
}
if !strings.HasSuffix(strings.TrimSpace(out), "echo hi") {
t.Fatalf("command suffix lost: %q", out)
}
skip := PrependNonInteractiveStdinRedirect("cmd </dev/null")
if strings.HasPrefix(skip, "exec </dev/null") {
t.Fatalf("should not double redirect: %q", skip)
}
}
func TestPrepareNonInteractiveShellCommand(t *testing.T) {
out := PrepareNonInteractiveShellCommand("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)
}
}
func TestNewShellInactivityWatch(t *testing.T) {
w := NewShellInactivityWatch(1)
if w == nil {
t.Fatal("expected watch")
}
w.Bump()
select {
case <-w.Expired:
case <-time.After(3 * time.Second):
t.Fatal("expected inactivity fire within 3s")
}
}
func TestResolveShellNoOutputTimeoutSeconds(t *testing.T) {
if ResolveShellNoOutputTimeoutSeconds(0) != 300 {
t.Fatal("zero should default to 300")
}
if ResolveShellNoOutputTimeoutSeconds(-1) != 0 {
t.Fatal("-1 should disable")
}
if ResolveShellNoOutputTimeoutSeconds(30) != 30 {
t.Fatal("explicit value")
}
}
// TestNonInteractiveStdinReadExitsQuickly 验证 exec </dev/null + attachNonInteractiveStdin 时 read 立即 EOF,不挂起。
func TestNonInteractiveStdinReadExitsQuickly(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", PrepareNonInteractiveShellCommand(`read x; echo "x=<$x>"`))
attachNonInteractiveStdin(cmd)
start := time.Now()
out, err := cmd.CombinedOutput()
elapsed := time.Since(start)
if elapsed > 2*time.Second {
t.Fatalf("read with closed stdin took %v, want <2s", elapsed)
}
if err != nil {
t.Fatalf("unexpected error: %v output=%q", err, out)
}
if !strings.Contains(string(out), "x=<>") {
t.Fatalf("unexpected output: %q", out)
}
}
// TestNonInteractiveStdinReadBlocksWithoutRedirect 对照:stdin 为永不写入的管道时 read 会挂起。
func TestNonInteractiveStdinReadBlocksWithoutRedirect(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
r, w, err := os.Pipe()
if err != nil {
t.Fatal(err)
}
defer r.Close()
// 保持 w 打开且不写数据,模拟「等待用户输入」
cmd := exec.Command("sh", "-c", `read x; echo done`)
cmd.Stdin = r
done := make(chan error, 1)
go func() { done <- cmd.Run() }()
select {
case err := <-done:
t.Fatalf("expected hang, but command finished: %v", err)
case <-time.After(500 * time.Millisecond):
if cmd.Process != nil {
_ = cmd.Process.Kill()
}
_ = w.Close()
<-done // 等待 goroutine 退出
}
}

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