mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-08 18:59:09 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6edc70f3fc | ||
|
|
377be6492a | ||
|
|
94cdf760af | ||
|
|
6ad9ea2d13 | ||
|
|
c70da22de7 | ||
|
|
fbbe984005 | ||
|
|
abc00c1fef | ||
|
|
d6d48c97c0 | ||
|
|
d7581a6373 | ||
|
|
4c011abb9d | ||
|
|
e4441f91ad | ||
|
|
21c6ad9bdf | ||
|
|
e0a2f01427 | ||
|
|
baff533196 | ||
|
|
474238cfc5 | ||
|
|
b47f8df3b0 | ||
|
|
a67761e843 | ||
|
|
b41596d51f | ||
|
|
d80e27e950 | ||
|
|
e218316c55 | ||
|
|
a34cab431a | ||
|
|
3bcf4458c5 | ||
|
|
bf761e9cd5 | ||
|
|
d640ef09c8 | ||
|
|
d88cfea761 | ||
|
|
bec2d2faf1 | ||
|
|
c7cc0bc9da | ||
|
|
24d06c5220 |
@@ -40,6 +40,7 @@ coverage.out
|
||||
coverage.html
|
||||
|
||||
# Logs and temporary files
|
||||
/log/
|
||||
*.log
|
||||
*.bak
|
||||
*~
|
||||
|
||||
@@ -126,13 +126,14 @@ CyberStrikeAI connects planning, execution, human oversight, evidence, and repla
|
||||
### Governance and audit
|
||||
|
||||
- 🧑⚖️ **Human in the loop** provides approval modes, tool allowlists, audit-agent review, and traceable decisions.
|
||||
- 🛡️ **Call blocking** under Security adds configurable regex checks before MCP execution, reminder templates, and dry runs, with government-domain protection enabled by default. See [Tool call blocking](docs/en-US/tool-call-guard.md).
|
||||
- 🔐 **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** 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.
|
||||
- 📁 **Conversation management** provides 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; 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.
|
||||
|
||||
@@ -125,6 +125,7 @@ CyberStrikeAI 将规划、执行、人工监督、证据与复盘连接在同一
|
||||
### 安全治理与审计
|
||||
|
||||
- 🧑⚖️ **人机协同**:支持审批模式、工具白名单、审计 Agent 复核和决策追踪。
|
||||
- 🛡️ **调用拦截**:「安全防护」下配置 MCP 执行前正则拦截、提醒模板和试匹配,默认启用政府域名保护。详见[调用拦截](docs/zh-CN/tool-call-guard.md)。
|
||||
- 🔐 **平台 RBAC**:支持多用户、系统及自定义角色、权限 Scope、资源归属和显式授权。
|
||||
- 🔒 **安全与审计**:提供登录保护、审计日志、SQLite 持久化和行动证据留存。
|
||||
- 📄 **结果治理**:数据库保存与 Agent 实际看到的同一份兜底后工具结果,恢复路径会再次防御历史超大输出,前端详情也有展示保护。详见[工具执行治理](docs/zh-CN/tool-execution-governance.md)。
|
||||
|
||||
+13
-1
@@ -5,6 +5,7 @@ import (
|
||||
"cyberstrike-ai/internal/logger"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
@@ -24,10 +25,21 @@ func main() {
|
||||
}
|
||||
|
||||
// 初始化日志(stdio 模式下使用 stderr 输出日志,避免干扰 JSON-RPC 通信)
|
||||
log := logger.New(cfg.Log.Level, "stderr")
|
||||
log := logger.New(cfg.Log.Level, "stderr", logger.DiagnosticOptions{
|
||||
Dir: cfg.Log.DiagnosticDir,
|
||||
Disabled: cfg.Log.DiagnosticDisabled,
|
||||
RetentionDays: cfg.Log.DiagnosticRetentionDays,
|
||||
})
|
||||
defer log.Sync()
|
||||
|
||||
// 创建MCP服务器
|
||||
mcpServer := mcp.NewServer(log.Logger)
|
||||
guard, err := toolguard.NewManager(cfg.EffectiveToolGuard())
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "初始化调用拦截失败: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
mcpServer.SetToolGuard(guard)
|
||||
|
||||
// 创建安全工具执行器
|
||||
executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger)
|
||||
|
||||
+6
-1
@@ -101,7 +101,12 @@ func main() {
|
||||
}
|
||||
|
||||
// 初始化日志
|
||||
log := logger.New(cfg.Log.Level, cfg.Log.Output)
|
||||
log := logger.New(cfg.Log.Level, cfg.Log.Output, logger.DiagnosticOptions{
|
||||
Dir: cfg.Log.DiagnosticDir,
|
||||
Disabled: cfg.Log.DiagnosticDisabled,
|
||||
RetentionDays: cfg.Log.DiagnosticRetentionDays,
|
||||
})
|
||||
defer log.Sync()
|
||||
|
||||
// 创建可取消的根 context,用于优雅关闭
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
+23
-3
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.7.15"
|
||||
version: "v1.7.18"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -36,6 +36,9 @@ auth:
|
||||
log:
|
||||
level: info # 日志级别: debug(调试), info(信息), warn(警告), error(错误)
|
||||
output: stdout # 日志输出位置: stdout(标准输出), stderr(标准错误), 或文件路径
|
||||
diagnostic_dir: log # 额外保存 warn 及以上诊断日志,按本地日期拆分;相对于进程工作目录
|
||||
diagnostic_retention_days: 14 # 保留天数(含当天);省略或 <= 0 使用 14 天
|
||||
diagnostic_disabled: false # true 关闭额外诊断日志;修改后需重启
|
||||
# 平台操作审计(系统设置 -> 日志审计;不记录对话正文与每次工具调用)
|
||||
audit:
|
||||
enabled: true
|
||||
@@ -130,13 +133,30 @@ agent:
|
||||
# system_prompt_path: prompts/single-agent.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
||||
|
||||
system_prompt_path: ""
|
||||
# 调用拦截:在「安全防护 → 调用拦截」编辑并保存后立即生效;与 HITL 审批白名单独立。
|
||||
# 旧配置省略 tool_guard 时也默认启用政府域名保护。手动修改 YAML 后需重启。
|
||||
tool_guard:
|
||||
enabled: true
|
||||
rules:
|
||||
- id: government-domains
|
||||
name: 政府网站保护
|
||||
enabled: true
|
||||
# Go/RE2 正则;命名捕获组 match 用于提醒中的 {match}。
|
||||
pattern: '(?i)(?:^|[^\p{L}\p{M}\p{N}_.-])(?P<match>(?:(?:[\p{L}\p{M}\p{N}_*-]+\.)+gov(?:\.[\p{L}\p{M}\p{N}_*-]+)*|gov(?:\.[\p{L}\p{M}\p{N}_*-]+)+|\.gov(?:\.[\p{L}\p{M}\p{N}_*-]+)*)\.?)(?:$|[^\p{L}\p{M}\p{N}_.-])'
|
||||
# 支持 {match}、{tool}、{rule};留空使用通用提醒。
|
||||
message: '识别到 {match},禁止攻击政府网站。请检查目标与授权范围,并更换为已获授权的非政府目标。'
|
||||
|
||||
# 人机协同(HITL)全局白名单:此处列出的工具始终免审批,与对话页「白名单工具(免审批,逗号分隔)」合并为并集;侧栏「应用」可合并写入本列表并立即生效。
|
||||
# 非白名单工具在审批方=审计 Agent 时,按会话 HITL 模式选用提示词:
|
||||
# approval → audit_agent_prompt
|
||||
# review_edit → audit_agent_prompt_review_edit(可改参后放行)
|
||||
hitl:
|
||||
# 全局默认审批方:human=人工审批,audit_agent=审计 Agent;未选会话时切换会写入本项,重启后仍生效
|
||||
# 全局默认人机协同模式:off=关闭,approval=审批模式,review_edit=审查编辑;新建会话无独立配置时沿用
|
||||
default_mode: off
|
||||
# 全局默认审批方:human=人工审批,audit_agent=审计 Agent;新建会话无独立配置时沿用
|
||||
default_reviewer: human
|
||||
# 全局默认审批等待时限(秒):300=5分钟,0=不限时;新建会话无独立配置时沿用
|
||||
default_timeout_seconds: 300
|
||||
# 审计 Agent 专用模型;字段留空则复用上方 openai 配置。建议 model 填小模型,用于降低审批成本。
|
||||
audit_model:
|
||||
provider: "" # openai / claude;留空跟随 openai.provider
|
||||
@@ -304,7 +324,7 @@ multi_agent:
|
||||
plan_execute_executed_steps_budget_ratio: 0.2 # plan_execute 中 executed_steps 预算比例
|
||||
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
|
||||
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
|
||||
checkpoint_dir: data/eino-checkpoints # P0:进程崩溃/OOM 后同会话自动 ADK Resume;正常结束会删 .ckpt;与「中断并继续」(last_react_*) 是两套机制
|
||||
checkpoint_dir: "" # 聊天链路不再使用 ADK checkpoint;跨轮模型态统一走 conversations.last_react_*,便于排查 stale context
|
||||
model_retry_max_retries: 0 # Eino 原生 ChatModel retry;408/409/425/429/5xx/网络抖动/空流式输出会重试;0=默认 4(永久性 4xx 不重试)
|
||||
model_retry_max_backoff_sec: 0 # Eino 原生 ChatModel retry 单次退避上限秒数;0=默认 30
|
||||
model_failover_channels: [] # Eino 原生 ChatModel failover;填写 ai.channels ID,例如 [qwen-plus];retry 耗尽后按顺序切换
|
||||
|
||||
@@ -143,3 +143,9 @@ After changing, validate the specific subsystem rather than trusting the save me
|
||||
- Config API and apply: `internal/handler/config.go`
|
||||
- Route registration: `internal/app/app.go`
|
||||
- C2 reconciliation: `internal/app/c2_lifecycle.go`
|
||||
|
||||
## Diagnostic logs
|
||||
|
||||
Alongside `log.output` (controlled by `log.level`), warnings and errors are saved as JSON Lines in `log/diagnostic-YYYY-MM-DD.log`, using the server’s local date. This independent warn-and-above output preserves existing context, caller information, and error stack traces; ordinary info/debug records are excluded and no extra request bodies or tool output are collected.
|
||||
|
||||
Configure `log.diagnostic_dir` (default `log`, relative to the working directory), `log.diagnostic_retention_days` (default 14, including today; nonpositive values use the default), or `log.diagnostic_disabled: true` to disable it. Restart after changing these settings. Files are created only when a diagnostic record is written; the first write each day removes expired files matching `diagnostic-YYYY-MM-DD.log`. Cleanup does not run while no diagnostic records are written. Write failures are reported to stderr without interrupting the primary log output.
|
||||
|
||||
@@ -89,7 +89,6 @@ Permissions use `module:action`. Common actions are `read`, `write`, `delete`, a
|
||||
| Attack chain | `attackchain:read`, `attackchain:write` |
|
||||
| Network-space search / Reconnaissance | `fofa:execute` |
|
||||
| OpenAPI | `openapi:read` |
|
||||
| Chat groups | `group:read`, `group:write`, `group:delete` |
|
||||
| Monitor | `monitor:read`, `monitor:write`, `monitor:delete` |
|
||||
|
||||
Important distinctions:
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# Tool call blocking
|
||||
|
||||
The **Security** sidebar groups the existing **Human in the loop** page with **Call blocking**. Call blocking checks internal MCP tools, external MCP tools, and HTTP MCP calls immediately before execution, independently of HITL approvals and allowlists.
|
||||
|
||||
The standalone `cmd/mcp-stdio` service also loads these rules. As a separate process, it requires a restart to pick up settings saved by the web application.
|
||||
|
||||
Government-domain protection is enabled by default, including when an older config omits `tool_guard`. The default rule matches government-domain forms such as `.gov`, `.gov.cn`, and wildcards, ignoring case. Add, edit, enable, or delete rules on the new page. Test an unsaved configuration with a tool name and JSON arguments before saving; testing never executes tools or updates the active policy.
|
||||
|
||||
**Add rule** opens a dialog with an independent draft and its own test inputs and results. **Add to list** checks the rule's RE2 syntax before adding it to the page draft; canceling leaves the list unchanged. Use the page's save button to apply the added rule.
|
||||
|
||||
Use **Test all rules** in the page header to check the configured order and enabled states. **Test this rule** inside an expanded rule opens a test area directly below that editor. Single-rule tests ignore the global and individual enable switches so disabled drafts can be checked; other rules cannot claim the match first. Each test area keeps separate inputs and displays its own matched text and rendered reminder.
|
||||
|
||||
Saving validates every rule, including disabled rules, writes only the `tool_guard` YAML section, and applies the policy immediately. Validation or write failure preserves the existing protection. Manual YAML edits require a restart; a non-null `tool_guard` section must explicitly provide `enabled` and `rules` (use `[]` for no rules). The first matching enabled rule supplies the reminder. Rules use Go/RE2 syntax; lookarounds and backreferences are unsupported. Up to 100 rules are allowed, with patterns and reminder templates capped at 4096 bytes each.
|
||||
|
||||
Reminder placeholders are `{match}` (matched text), `{tool}` (tool name), and `{rule}` (rule name). An optional named group `(?P<match>...)` selects the matched text. Empty templates use a default reminder.
|
||||
|
||||
Checks inspect the tool name, serialized JSON, nested strings and keys, and up to three rounds of common URL percent decoding. Blocked calls use a separate **Blocked** status in the UI and execution records, with the reason preserved. Monitoring counts blocks separately and excludes them from failed calls and the success-rate denominator. At startup after an upgrade, clearly identifiable legacy guard-block records are migrated to this status. MCP results retain `isError: true` alongside `blocked: true` so the agent knows the request did not execute. External MCP policy blocks do not count as provider failures for circuit breaking. Updated rules cannot cancel calls already dispatched.
|
||||
|
||||
Viewing and testing require `config:read`. Saving requires `config:write` with global scope, and configuration changes are audited. HITL approvals, edited arguments, and approval allowlists cannot bypass the execution check.
|
||||
|
||||
Text matching cannot establish target ownership from IP addresses, DNS aliases, redirects, file contents, or arbitrary obfuscation. Independent execution paths such as direct terminals and optional agent local tools are outside the MCP guard. Benign references to a protected domain in parameter text may also be blocked. Retain HITL and maintain rules against the actual authorized scope.
|
||||
|
||||
API endpoints:
|
||||
|
||||
- `GET /api/tool-guard`: active `{enabled, rules}` configuration.
|
||||
- `PUT /api/tool-guard`: save that structure, explicitly providing both fields. Rules contain `id`, `name`, `enabled`, `pattern`, and `message`.
|
||||
- `POST /api/tool-guard/test`: accepts `{config, toolName, arguments}` and returns `{blocked, match?}`. Match fields are `ruleId`, `ruleName`, `matchedText`, and `message`.
|
||||
@@ -27,7 +27,11 @@ log:
|
||||
- Chromium 浏览器插件的合法 `chrome-extension://<32位插件ID>` Origin 会被自动识别,无需配置。插件仍需按域授权,并使用密码登录与 Bearer Token 调用 API。
|
||||
- `server.cors_allowed_origins`:仅供其他可信 Web 集成使用的额外 Origin 精确白名单;不支持 `*`,修改后需重启服务。
|
||||
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
||||
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
||||
- `log.output`:可以是 `stdout`、`stderr` 或文件路径,由 `log.level` 控制级别。
|
||||
- 额外诊断日志默认开启,仅记录 `warn` 及以上(包括重试、连接异常和错误),独立于 `log.level`,不保存普通 `info` / `debug` 日志。保留原有结构化字段、时间、代码位置和 Error 及以上堆栈,不额外采集请求正文或工具输出。
|
||||
- `log.diagnostic_dir`:默认 `log`,相对于进程工作目录,文件名为 `diagnostic-YYYY-MM-DD.log`(JSON Lines,按服务器本地日期拆分)。只有出现诊断日志时才创建目录和文件;跨天后首次写入切换文件。
|
||||
- `log.diagnostic_retention_days`:默认 14 天(含当天);省略或小于等于 0 时使用默认值。每天首次写入时清理此目录内过期的 `diagnostic-日期.log`,不删除其他文件;没有新诊断日志时不执行清理。
|
||||
- `log.diagnostic_disabled: true`:关闭额外诊断落盘。以上日志配置修改后需重启;目录无法写入时保留原输出,并由 Zap 向 stderr 报告写入失败。
|
||||
|
||||
## AI 通道与模型配置
|
||||
|
||||
|
||||
@@ -96,7 +96,6 @@ AI 测试角色不是安全授权边界。即使选择了“渗透测试”角
|
||||
| 攻击链 | `attackchain:read`、`attackchain:write` |
|
||||
| 网络空间测绘 / 信息收集 | `fofa:execute` |
|
||||
| OpenAPI | `openapi:read` |
|
||||
| 对话分组 | `group:read`、`group:write`、`group:delete` |
|
||||
| 执行监控 | `monitor:read`、`monitor:write`、`monitor:delete` |
|
||||
|
||||
特殊权限说明:
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
# 调用拦截
|
||||
|
||||
侧边栏「安全防护」包含「人机协同」和「调用拦截」。人机协同保留原有审批、白名单、审计策略和日志功能;调用拦截对内部 MCP、外部 MCP 和 HTTP MCP 工具调用增加独立的执行前检查。
|
||||
|
||||
独立的 `cmd/mcp-stdio` 服务也会加载相同拦截规则;该服务是单独进程,网页保存后需重启它以加载更新。
|
||||
|
||||
## 使用
|
||||
|
||||
1. 打开「安全防护 → 调用拦截」。默认开启「政府网站保护」,匹配 `.gov`、`.gov.cn` 等政府域名及通配符写法,忽略大小写。旧配置没有 `tool_guard` 时也启用默认保护。
|
||||
2. 规则默认折叠,列表展示名称、提醒摘要和启停开关。点击规则展开编辑;「添加规则」打开独立弹窗,可以填写并验证尚未添加的规则。点击「添加到列表」通过正则校验后加入页面草稿,取消不会留下空规则。校验失败会定位到对应字段。可以单独启停规则,也可以关闭总开关。
|
||||
3. 点击页面顶部「全部规则验证」按当前顺序及启停状态检查全部规则。在已有规则编辑区点击「验证本条」,就在该规则下方输入工具名和 JSON 参数、查看命中文本与最终提醒;新增规则的单条验证直接在弹窗中完成。单条验证忽略总开关和该规则的启停状态,适合调试未启用规则。单条和全部验证分别保留输入与结果,均使用未保存的表单配置,不执行工具,也不改变运行中的规则。
|
||||
4. 点击保存。服务端校验全部规则后写入 `config.yaml` 的 `tool_guard`,立即生效,无需重启。校验或写入失败会保留原有规则。直接编辑 YAML 后需要重启服务;非空 `tool_guard` 必须明确填写 `enabled` 和 `rules`,清空规则使用 `[]`。
|
||||
|
||||
规则按列表顺序检查,首先命中的启用规则决定提醒。检查对象包括工具名称、参数的 JSON 表示、嵌套字符串和键名,以及最多三轮常见 URL 百分号解码后的文本。使用 Go/RE2 正则语法,例如 `(?i)` 表示忽略大小写;不支持回溯引用和环视。最多 100 条规则,正则和提醒各最多 4096 字节。禁用的规则也须通过校验。
|
||||
|
||||
提醒支持以下占位符,留空则使用通用提醒:
|
||||
|
||||
| 占位符 | 内容 |
|
||||
| --- | --- |
|
||||
| `{match}` | 匹配文本;正则含命名捕获组 `(?P<match>...)` 时使用该组 |
|
||||
| `{tool}` | 工具名称 |
|
||||
| `{rule}` | 规则名称 |
|
||||
|
||||
示例提醒:`识别到 {match},禁止攻击政府网站,请检查目标与授权范围。`
|
||||
|
||||
命中后,工具处理器或外部客户端不会执行该调用。界面和执行记录使用独立的「已拦截」状态,并保留拦截原因;监控单独统计拦截次数,不计入调用失败或成功率的分母。升级启动时,可明确识别的旧版安全规则拦截记录会自动归入此状态。返回给 Agent 的 MCP 结果仍保留 `isError: true`,同时携带 `blocked: true`,以明确表示请求未执行。外部 MCP 的规则拦截不会算作服务故障而触发熔断。
|
||||
|
||||
## 权限与边界
|
||||
|
||||
查看和试匹配需要 `config:read`;修改需要 `config:write` 和全局权限范围。规则配置变更写入系统审计日志。HITL 的关闭状态、免审批白名单和审批通过结果均不能覆盖调用拦截;审批后编辑的参数也会在实际执行入口检查。规则更新影响后续执行检查,不能撤销已经发出的调用。
|
||||
|
||||
这是文本规则防护,不能代替目标授权或网络隔离:它无法可靠识别仅以 IP 表示的政府目标、DNS 别名背后的机构、工具执行后的重定向、文件中才出现的目标或任意混淆编码。它只覆盖经过本应用 MCP 执行入口的调用;直接终端操作、可选的 Agent 本地执行工具等独立入口不在此范围内。参数中仅引用政府域名的说明文本也可能被保守拦截。请保留人机协同,并结合实际授权范围维护规则。
|
||||
|
||||
## API
|
||||
|
||||
- `GET /api/tool-guard`:返回生效配置 `{enabled, rules}`。
|
||||
- `PUT /api/tool-guard`:保存相同结构;每条规则含 `id`、`name`、`enabled`、`pattern`、`message`。必须明确提供总开关和规则数组。
|
||||
- `POST /api/tool-guard/test`:请求 `{config, toolName, arguments}`,响应 `{blocked, match?}`;`match` 含 `ruleId`、`ruleName`、`matchedText`、`message`。
|
||||
+34
-5
@@ -512,6 +512,7 @@ type ToolExecutionResult struct {
|
||||
Result string
|
||||
ExecutionID string
|
||||
IsError bool
|
||||
Blocked bool
|
||||
}
|
||||
|
||||
func buildToolFailureMessage(toolName, detail string, err error) string {
|
||||
@@ -612,6 +613,7 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
||||
Result: resultStr,
|
||||
ExecutionID: executionID,
|
||||
IsError: result != nil && result.IsError,
|
||||
Blocked: result != nil && result.Blocked,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -815,6 +817,10 @@ func (a *Agent) UpdateMCPExecutionDisplayResult(executionID, resultText string)
|
||||
tr := &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: text}},
|
||||
}
|
||||
if exec := a.mcpExecution(executionID); exec != nil && exec.Result != nil {
|
||||
tr.IsError = exec.Result.IsError
|
||||
tr.Blocked = exec.Result.Blocked
|
||||
}
|
||||
if a.mcpServer != nil {
|
||||
_ = a.mcpServer.UpdateToolExecutionResult(executionID, tr)
|
||||
}
|
||||
@@ -823,16 +829,39 @@ func (a *Agent) UpdateMCPExecutionDisplayResult(executionID, resultText string)
|
||||
// MCPExecutionResultText returns the monitor-facing result text after storage
|
||||
// guards such as large-output spilling have been applied.
|
||||
func (a *Agent) MCPExecutionResultText(executionID string) string {
|
||||
if a == nil || a.mcpServer == nil || strings.TrimSpace(executionID) == "" {
|
||||
return ""
|
||||
}
|
||||
exec, ok := a.mcpServer.GetExecution(executionID)
|
||||
if !ok || exec == nil || exec.Result == nil {
|
||||
exec := a.mcpExecution(executionID)
|
||||
if exec == nil || exec.Result == nil {
|
||||
return ""
|
||||
}
|
||||
return mcp.ToolResultPlainText(exec.Result)
|
||||
}
|
||||
|
||||
// MCPExecutionStatus returns the recorded outcome independently of model-facing
|
||||
// text reduction, which can remove the original refusal wording.
|
||||
func (a *Agent) MCPExecutionStatus(executionID string) string {
|
||||
if exec := a.mcpExecution(executionID); exec != nil {
|
||||
return exec.Status
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (a *Agent) mcpExecution(executionID string) *mcp.ToolExecution {
|
||||
if a == nil || strings.TrimSpace(executionID) == "" {
|
||||
return nil
|
||||
}
|
||||
if a.mcpServer != nil {
|
||||
if exec, ok := a.mcpServer.GetExecution(executionID); ok && exec != nil {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
if a.externalMCPMgr != nil {
|
||||
if exec, ok := a.externalMCPMgr.GetExecution(executionID); ok {
|
||||
return exec
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CancelMCPToolExecutionWithNote 取消一次进行中的 MCP 工具(先内部后外部),与监控页「终止工具」一致;note 非空时合并进返回给模型的文本。
|
||||
func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool {
|
||||
executionID = strings.TrimSpace(executionID)
|
||||
|
||||
@@ -65,12 +65,6 @@ func FromRunResult(db *database.DB, result *multiagent.RunResult, in Input) Deci
|
||||
if len(in.MCPExecutionIDs) == 0 {
|
||||
in.MCPExecutionIDs = result.MCPExecutionIDs
|
||||
}
|
||||
if strings.TrimSpace(in.Status) == "" {
|
||||
in.Status = result.Status
|
||||
}
|
||||
if strings.TrimSpace(in.CompletionReason) == "" {
|
||||
in.CompletionReason = result.CompletionReason
|
||||
}
|
||||
}
|
||||
d := Decide(db, in)
|
||||
if result != nil {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -130,3 +131,23 @@ func TestDecideAllowsInformationalAnswerWhenExecutionEvidenceIsNotRequired(t *te
|
||||
t.Fatalf("informational response should finalize when execution evidence is not required: %+v", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromRunResultDoesNotReusePreviousFinalizationStatusAsRunStatus(t *testing.T) {
|
||||
db := newDecisionTestDB(t)
|
||||
saveDecisionTestExecution(t, db, "run-slow", mcp.ToolExecutionStatusRunning)
|
||||
result := &multiagent.RunResult{
|
||||
Response: "工具已触发,按用户要求直接总结。",
|
||||
MCPExecutionIDs: []string{"run-slow"},
|
||||
}
|
||||
|
||||
first := FromRunResult(db, result, Input{})
|
||||
if first.Finalizable || first.CompletionReason != ReasonPendingTools || result.Status != StatusInProgress {
|
||||
t.Fatalf("first decision should mark pending and write metadata: decision=%+v result=%+v", first, result)
|
||||
}
|
||||
|
||||
saveDecisionTestExecution(t, db, "run-slow", mcp.ToolExecutionStatusCancelled)
|
||||
second := FromRunResult(db, result, Input{})
|
||||
if !second.Finalizable || !second.Finalized || second.Status != StatusCompleted {
|
||||
t.Fatalf("second decision should ignore previous result status after pending cleanup: decision=%+v result=%+v", second, result)
|
||||
}
|
||||
}
|
||||
|
||||
+16
-17
@@ -33,6 +33,7 @@ import (
|
||||
"cyberstrike-ai/internal/robot"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/skillpackage"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
@@ -76,6 +77,10 @@ type App struct {
|
||||
|
||||
// New 创建新应用
|
||||
func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) {
|
||||
toolGuard, err := toolguard.NewManager(cfg.EffectiveToolGuard())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("初始化调用拦截规则: %w", err)
|
||||
}
|
||||
if err := multiagent.InitADK(); err != nil {
|
||||
return nil, fmt.Errorf("初始化 Eino ADK: %w", err)
|
||||
}
|
||||
@@ -147,6 +152,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// 创建MCP服务器(带数据库持久化)
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
mcpServer.SetToolGuard(toolGuard)
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
|
||||
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
@@ -170,6 +176,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
||||
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
||||
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
|
||||
externalMCPMgr.SetToolGuard(toolGuard)
|
||||
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
|
||||
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||
externalMCPMgr.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||
@@ -391,7 +398,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
monitorHandler.SetTaskManager(agentHandler.TaskManager())
|
||||
monitorHandler.SetAgentHandler(agentHandler)
|
||||
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
|
||||
groupHandler := handler.NewGroupHandler(db, log.Logger)
|
||||
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
|
||||
authHandler.SetAudit(auditSvc)
|
||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||
@@ -413,6 +419,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
|
||||
configHandler.SetDB(db)
|
||||
configHandler.SetToolGuard(toolGuard)
|
||||
configHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetHitlToolWhitelistSaver(configHandler)
|
||||
agentHandler.SetHitlAuditStrategySaver(configHandler)
|
||||
@@ -567,7 +574,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
conversationHandler,
|
||||
robotHandler,
|
||||
wechatRobotHandler,
|
||||
groupHandler,
|
||||
configHandler,
|
||||
externalMCPHandler,
|
||||
attackChainHandler,
|
||||
@@ -870,7 +876,6 @@ func setupRoutes(
|
||||
conversationHandler *handler.ConversationHandler,
|
||||
robotHandler *handler.RobotHandler,
|
||||
wechatRobotHandler *handler.WechatRobotHandler,
|
||||
groupHandler *handler.GroupHandler,
|
||||
configHandler *handler.ConfigHandler,
|
||||
externalMCPHandler *handler.ExternalMCPHandler,
|
||||
attackChainHandler *handler.AttackChainHandler,
|
||||
@@ -972,6 +977,8 @@ func setupRoutes(
|
||||
protected.GET("/hitl/tool-whitelist", agentHandler.GetHITLGlobalToolWhitelist)
|
||||
protected.PUT("/hitl/tool-whitelist", agentHandler.SetHITLGlobalToolWhitelist)
|
||||
protected.POST("/hitl/tool-whitelist", agentHandler.MergeHITLGlobalToolWhitelist)
|
||||
protected.GET("/hitl/default-config", agentHandler.GetHITLDefaultConfig)
|
||||
protected.PUT("/hitl/default-config", agentHandler.UpdateHITLDefaultConfig)
|
||||
protected.GET("/hitl/default-reviewer", agentHandler.GetHITLDefaultReviewer)
|
||||
protected.PUT("/hitl/default-reviewer", agentHandler.UpdateHITLDefaultReviewer)
|
||||
protected.GET("/hitl/audit-strategy", agentHandler.GetHITLAuditStrategy)
|
||||
@@ -1027,9 +1034,11 @@ func setupRoutes(
|
||||
protected.DELETE("/batch-tasks/:queueId/tasks/:taskId", agentHandler.DeleteBatchTask)
|
||||
|
||||
// 对话历史
|
||||
protected.GET("/usage/tokens", conversationHandler.GetTokenUsageStats)
|
||||
protected.POST("/conversations", conversationHandler.CreateConversation)
|
||||
protected.GET("/conversations", conversationHandler.ListConversations)
|
||||
protected.GET("/conversations/:id", conversationHandler.GetConversation)
|
||||
protected.GET("/conversations/:id/token-usage", conversationHandler.GetConversationTokenUsageStats)
|
||||
protected.GET("/conversations/:id/plan-tasks", conversationHandler.GetConversationPlanTasks)
|
||||
protected.GET("/messages/:id/process-details", conversationHandler.GetMessageProcessDetails)
|
||||
protected.GET("/process-details/:id", conversationHandler.GetProcessDetail)
|
||||
@@ -1037,20 +1046,7 @@ func setupRoutes(
|
||||
protected.PUT("/conversations/:id/project", conversationHandler.SetConversationProject)
|
||||
protected.DELETE("/conversations/:id", conversationHandler.DeleteConversation)
|
||||
protected.POST("/conversations/:id/delete-turn", conversationHandler.DeleteConversationTurn)
|
||||
protected.PUT("/conversations/:id/pinned", groupHandler.UpdateConversationPinned)
|
||||
|
||||
// 对话分组
|
||||
protected.POST("/groups", groupHandler.CreateGroup)
|
||||
protected.GET("/groups", groupHandler.ListGroups)
|
||||
protected.GET("/groups/:id", groupHandler.GetGroup)
|
||||
protected.PUT("/groups/:id", groupHandler.UpdateGroup)
|
||||
protected.DELETE("/groups/:id", groupHandler.DeleteGroup)
|
||||
protected.PUT("/groups/:id/pinned", groupHandler.UpdateGroupPinned)
|
||||
protected.GET("/groups/:id/conversations", groupHandler.GetGroupConversations)
|
||||
protected.GET("/groups/mappings", groupHandler.GetAllMappings)
|
||||
protected.POST("/groups/conversations", groupHandler.AddConversationToGroup)
|
||||
protected.DELETE("/groups/:id/conversations/:conversationId", groupHandler.RemoveConversationFromGroup)
|
||||
protected.PUT("/groups/:id/conversations/:conversationId/pinned", groupHandler.UpdateConversationPinnedInGroup)
|
||||
protected.PUT("/conversations/:id/pinned", conversationHandler.UpdateConversationPinned)
|
||||
|
||||
// 监控
|
||||
protected.GET("/monitor", monitorHandler.Monitor)
|
||||
@@ -1066,6 +1062,9 @@ func setupRoutes(
|
||||
|
||||
// 配置管理
|
||||
protected.GET("/config", configHandler.GetConfig)
|
||||
protected.GET("/tool-guard", configHandler.GetToolGuard)
|
||||
protected.PUT("/tool-guard", configHandler.UpdateToolGuard)
|
||||
protected.POST("/tool-guard/test", configHandler.TestToolGuard)
|
||||
protected.GET("/config/tools", configHandler.GetTools)
|
||||
protected.GET("/config/tools/:name/schema", configHandler.GetToolSchema)
|
||||
protected.PUT("/config", configHandler.UpdateConfig)
|
||||
|
||||
@@ -6,12 +6,14 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/termout"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
@@ -29,6 +31,7 @@ type Config struct {
|
||||
Shodan SpaceSearchConfig `yaml:"shodan,omitempty" json:"shodan,omitempty"`
|
||||
Agent AgentConfig `yaml:"agent"`
|
||||
Hitl HitlConfig `yaml:"hitl,omitempty" json:"hitl,omitempty"`
|
||||
ToolGuard *toolguard.Config `yaml:"tool_guard,omitempty" json:"tool_guard,omitempty"`
|
||||
Security SecurityConfig `yaml:"security"`
|
||||
Database DatabaseConfig `yaml:"database"`
|
||||
Auth AuthConfig `yaml:"auth"`
|
||||
@@ -298,7 +301,8 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
PlanExecuteMaxStepResultRunes int `yaml:"plan_execute_max_step_result_runes,omitempty" json:"plan_execute_max_step_result_runes,omitempty"`
|
||||
// PlanExecuteKeepLastSteps keeps only the tail steps in prompt view (default 8).
|
||||
PlanExecuteKeepLastSteps int `yaml:"plan_execute_keep_last_steps,omitempty" json:"plan_execute_keep_last_steps,omitempty"`
|
||||
// CheckpointDir when non-empty enables adk.Runner CheckPointStore (file-backed) for interrupt/resume persistence.
|
||||
// CheckpointDir is retained for config compatibility. Chat agent runs do
|
||||
// not consume it; cross-turn recovery is centralized in conversations.last_react_*.
|
||||
CheckpointDir string `yaml:"checkpoint_dir,omitempty" json:"checkpoint_dir,omitempty"`
|
||||
// DeepOutputKey passed to deep.Config OutputKey (session final text); empty = off.
|
||||
DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"`
|
||||
@@ -803,8 +807,11 @@ type ServerConfig struct {
|
||||
}
|
||||
|
||||
type LogConfig struct {
|
||||
Level string `yaml:"level"`
|
||||
Output string `yaml:"output"`
|
||||
Level string `yaml:"level"`
|
||||
Output string `yaml:"output"`
|
||||
DiagnosticDir string `yaml:"diagnostic_dir"`
|
||||
DiagnosticDisabled bool `yaml:"diagnostic_disabled"`
|
||||
DiagnosticRetentionDays int `yaml:"diagnostic_retention_days"`
|
||||
}
|
||||
|
||||
type MCPConfig struct {
|
||||
@@ -945,10 +952,12 @@ func (c *Config) ApplyDefaultAIChannel() {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.NormalizeAIProviderProfiles()
|
||||
c.AI.EnsureDefaultFromOpenAI(c.OpenAI)
|
||||
if oa, _, ok := c.AI.ResolveChannel(c.AI.DefaultChannel); ok {
|
||||
c.OpenAI = oa
|
||||
}
|
||||
c.NormalizeAIProviderProfiles()
|
||||
}
|
||||
|
||||
func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
||||
@@ -959,13 +968,56 @@ func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
||||
}
|
||||
|
||||
// IsDeepSeekEndpointOrModel reports whether the channel targets DeepSeek's
|
||||
// official-compatible API or a DeepSeek model family. This is separate from the
|
||||
// reasoning profile: profile controls field mapping, while DeepSeek has provider
|
||||
// constraints such as default thinking mode and no tool_choice in thinking mode.
|
||||
// official-compatible API endpoint. The historical name is kept for compatibility;
|
||||
// model names alone are not enough to infer DeepSeek wire behavior behind
|
||||
// OpenAI-compatible gateways.
|
||||
func (c OpenAIConfig) IsDeepSeekEndpointOrModel() bool {
|
||||
baseURL := strings.ToLower(strings.TrimSpace(c.BaseURL))
|
||||
model := strings.ToLower(strings.TrimSpace(c.Model))
|
||||
return strings.Contains(baseURL, "deepseek") || strings.Contains(model, "deepseek")
|
||||
return strings.Contains(baseURL, "deepseek")
|
||||
}
|
||||
|
||||
func (c OpenAIConfig) IsDeepSeekOfficialEndpoint() bool {
|
||||
host := normalizedURLHost(c.BaseURL)
|
||||
return host == "api.deepseek.com"
|
||||
}
|
||||
|
||||
func normalizedURLHost(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil || parsed.Host == "" {
|
||||
parsed, err = url.Parse("https://" + strings.TrimLeft(raw, "/"))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return strings.ToLower(strings.TrimPrefix(parsed.Hostname(), "www."))
|
||||
}
|
||||
|
||||
func NormalizeOpenAIProviderProfile(oa *OpenAIConfig) {
|
||||
if oa == nil {
|
||||
return
|
||||
}
|
||||
if oa.IsDeepSeekOfficialEndpoint() {
|
||||
oa.Reasoning.Profile = "deepseek"
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) NormalizeAIProviderProfiles() {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
NormalizeOpenAIProviderProfile(&c.OpenAI)
|
||||
if c.AI.Channels != nil {
|
||||
for id, ch := range c.AI.Channels {
|
||||
oa := ch.ToOpenAIConfig()
|
||||
NormalizeOpenAIProviderProfile(&oa)
|
||||
ch.Reasoning = oa.Reasoning
|
||||
c.AI.Channels[id] = ch
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAIReasoningConfig 全局默认与网关 profile(对话页可通过 ChatRequest.reasoning 覆盖,受 AllowClientReasoning 约束)。
|
||||
@@ -1062,8 +1114,24 @@ type HitlConfig struct {
|
||||
AuditAgentPromptReviewEdit string `yaml:"audit_agent_prompt_review_edit,omitempty" json:"audit_agent_prompt_review_edit,omitempty"`
|
||||
// RetentionDays 已决策审计日志(hitl_interrupts 非 pending)保留天数;省略时默认 90;0 表示不自动清理。
|
||||
RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
|
||||
// DefaultReviewer 全局默认审批方(human | audit_agent);未选会话时切换会写入 config.yaml;新建会话无独立配置时沿用。
|
||||
// DefaultMode 全局默认人机协同模式(off | approval | review_edit);新建会话无独立配置时沿用。
|
||||
DefaultMode string `yaml:"default_mode,omitempty" json:"default_mode,omitempty"`
|
||||
// DefaultReviewer 全局默认审批方(human | audit_agent);新建会话无独立配置时沿用。
|
||||
DefaultReviewer string `yaml:"default_reviewer,omitempty" json:"default_reviewer,omitempty"`
|
||||
// DefaultTimeoutSeconds 全局默认审批等待秒数;nil 表示使用前端历史默认 300 秒,0 表示不限时。
|
||||
DefaultTimeoutSeconds *int `yaml:"default_timeout_seconds,omitempty" json:"default_timeout_seconds,omitempty"`
|
||||
}
|
||||
|
||||
// EffectiveDefaultMode returns off, approval, or review_edit; omitted or unknown values default to off.
|
||||
func (h HitlConfig) EffectiveDefaultMode() string {
|
||||
switch strings.ToLower(strings.TrimSpace(h.DefaultMode)) {
|
||||
case "feedback", "followup":
|
||||
return "approval"
|
||||
case "approval", "review_edit":
|
||||
return strings.ToLower(strings.TrimSpace(h.DefaultMode))
|
||||
default:
|
||||
return "off"
|
||||
}
|
||||
}
|
||||
|
||||
// EffectiveDefaultReviewer returns human or audit_agent; omitted or unknown values default to human.
|
||||
@@ -1076,6 +1144,17 @@ func (h HitlConfig) EffectiveDefaultReviewer() string {
|
||||
}
|
||||
}
|
||||
|
||||
// EffectiveDefaultTimeoutSeconds returns the default HITL approval timeout; nil defaults to 5 minutes.
|
||||
func (h HitlConfig) EffectiveDefaultTimeoutSeconds() int {
|
||||
if h.DefaultTimeoutSeconds == nil {
|
||||
return 300
|
||||
}
|
||||
if *h.DefaultTimeoutSeconds < 0 {
|
||||
return 0
|
||||
}
|
||||
return *h.DefaultTimeoutSeconds
|
||||
}
|
||||
|
||||
// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90.
|
||||
func (h HitlConfig) RetentionDaysEffective() int {
|
||||
if h.RetentionDays == nil {
|
||||
@@ -1365,6 +1444,14 @@ func Load(path string) (*Config, error) {
|
||||
if err := yaml.Unmarshal(data, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
if cfg.ToolGuard != nil {
|
||||
if err := validateToolGuardYAML(data); err != nil {
|
||||
return nil, fmt.Errorf("调用拦截配置无效: %w", err)
|
||||
}
|
||||
}
|
||||
if _, err := toolguard.Compile(cfg.EffectiveToolGuard()); err != nil {
|
||||
return nil, fmt.Errorf("调用拦截配置无效: %w", err)
|
||||
}
|
||||
|
||||
if cfg.Auth.SessionDurationHours <= 0 {
|
||||
cfg.Auth.SessionDurationHours = 12
|
||||
@@ -1372,6 +1459,7 @@ func Load(path string) (*Config, error) {
|
||||
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||
cfg.Audit.MaxDetailBytes = 8192
|
||||
}
|
||||
cfg.NormalizeAIProviderProfiles()
|
||||
cfg.ApplyDefaultAIChannel()
|
||||
if err := validateOpenAIOutputLimits(cfg.OpenAI); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -95,6 +95,29 @@ func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHitlDefaultConfigEffectiveValues(t *testing.T) {
|
||||
if got := (HitlConfig{}).EffectiveDefaultMode(); got != "off" {
|
||||
t.Fatalf("empty default mode = %q, want off", got)
|
||||
}
|
||||
if got := (HitlConfig{DefaultMode: "review-edit"}).EffectiveDefaultMode(); got != "off" {
|
||||
t.Fatalf("unknown default mode = %q, want off", got)
|
||||
}
|
||||
if got := (HitlConfig{DefaultMode: "review_edit"}).EffectiveDefaultMode(); got != "review_edit" {
|
||||
t.Fatalf("review_edit default mode = %q, want review_edit", got)
|
||||
}
|
||||
if got := (HitlConfig{}).EffectiveDefaultTimeoutSeconds(); got != 300 {
|
||||
t.Fatalf("empty default timeout = %d, want 300", got)
|
||||
}
|
||||
zero := 0
|
||||
if got := (HitlConfig{DefaultTimeoutSeconds: &zero}).EffectiveDefaultTimeoutSeconds(); got != 0 {
|
||||
t.Fatalf("zero default timeout = %d, want 0", got)
|
||||
}
|
||||
neg := -1
|
||||
if got := (HitlConfig{DefaultTimeoutSeconds: &neg}).EffectiveDefaultTimeoutSeconds(); got != 0 {
|
||||
t.Fatalf("negative default timeout = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
@@ -137,6 +160,84 @@ func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAIProviderProfilesForOfficialDeepSeekEndpoint(t *testing.T) {
|
||||
cfg := &Config{
|
||||
OpenAI: OpenAIConfig{
|
||||
BaseURL: "https://api.deepseek.com/v1",
|
||||
Model: "deepseek-chat",
|
||||
Reasoning: OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
},
|
||||
},
|
||||
AI: AIConfig{
|
||||
Channels: map[string]AIChannelConfig{
|
||||
"official": {
|
||||
BaseURL: "api.deepseek.com/v1",
|
||||
Model: "deepseek-chat",
|
||||
Reasoning: OpenAIReasoningConfig{
|
||||
Profile: "auto",
|
||||
},
|
||||
},
|
||||
"gateway": {
|
||||
BaseURL: "https://compatible.example.com/v1",
|
||||
Model: "deepseek-chat",
|
||||
Reasoning: OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cfg.NormalizeAIProviderProfiles()
|
||||
|
||||
if cfg.OpenAI.Reasoning.Profile != "deepseek" {
|
||||
t.Fatalf("openai profile = %q, want deepseek", cfg.OpenAI.Reasoning.Profile)
|
||||
}
|
||||
if got := cfg.AI.Channels["official"].Reasoning.Profile; got != "deepseek" {
|
||||
t.Fatalf("official channel profile = %q, want deepseek", got)
|
||||
}
|
||||
if got := cfg.AI.Channels["gateway"].Reasoning.Profile; got != "openai_compat" {
|
||||
t.Fatalf("gateway profile should be preserved, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadNormalizesDefaultChannelForOfficialDeepSeekEndpoint(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
initial := strings.Join([]string{
|
||||
"ai:",
|
||||
" default_channel: deepseek",
|
||||
" channels:",
|
||||
" deepseek:",
|
||||
" name: DeepSeek",
|
||||
" provider: openai_compatible",
|
||||
" base_url: https://api.deepseek.com/v1",
|
||||
" api_key: deepseek-key",
|
||||
" model: deepseek-chat",
|
||||
" reasoning:",
|
||||
" profile: openai_compat",
|
||||
"server:",
|
||||
" host: 127.0.0.1",
|
||||
" port: 8080",
|
||||
"",
|
||||
}, "\n")
|
||||
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
|
||||
t.Fatalf("write config: %v", err)
|
||||
}
|
||||
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.OpenAI.Reasoning.Profile != "deepseek" {
|
||||
t.Fatalf("runtime OpenAI profile = %q, want deepseek", cfg.OpenAI.Reasoning.Profile)
|
||||
}
|
||||
if got := cfg.AI.Channels["deepseek"].Reasoning.Profile; got != "deepseek" {
|
||||
t.Fatalf("channel profile = %q, want deepseek", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
||||
var zero MultiAgentEinoMiddlewareConfig
|
||||
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// EffectiveToolGuard enables the default government-domain protection for old
|
||||
// configurations as well as new installs. An explicit config may disable it.
|
||||
func (c *Config) EffectiveToolGuard() toolguard.Config {
|
||||
if c.ToolGuard == nil {
|
||||
return toolguard.DefaultConfig()
|
||||
}
|
||||
return *c.ToolGuard
|
||||
}
|
||||
|
||||
// validateToolGuardYAML requires an explicit decision for both protection and
|
||||
// its rules whenever a non-null section is supplied. Otherwise a typo or partial
|
||||
// section could silently turn the enabled-by-default protection off. Pointer
|
||||
// fields distinguish false/[] from omitted or null values, and the YAML decoder
|
||||
// continues to support aliases and merged configuration mappings.
|
||||
func validateToolGuardYAML(data []byte) error {
|
||||
var document struct {
|
||||
ToolGuard *struct {
|
||||
Enabled *bool `yaml:"enabled"`
|
||||
Rules *[]toolguard.Rule `yaml:"rules"`
|
||||
} `yaml:"tool_guard"`
|
||||
}
|
||||
if err := yaml.Unmarshal(data, &document); err != nil {
|
||||
return err
|
||||
}
|
||||
if section := document.ToolGuard; section != nil && (section.Enabled == nil || section.Rules == nil) {
|
||||
return fmt.Errorf("tool_guard 必须明确提供 enabled 和 rules;清空规则请提供空数组")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadToolGuardDefaultsAndValidation(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, yaml string
|
||||
enabled, wantErr bool
|
||||
}{
|
||||
{"legacy config", "server: {port: 8080}\n", true, false},
|
||||
{"null section", "tool_guard: null\n", true, false},
|
||||
{"implicit null section", "tool_guard:\n", true, false},
|
||||
{"explicit off", "tool_guard: {enabled: false, rules: []}\n", false, false},
|
||||
{"explicit empty", "tool_guard: {enabled: true, rules: []}\n", true, false},
|
||||
{"merged explicit config", "guard_defaults: &guard_defaults {enabled: false, rules: []}\ntool_guard: {<<: *guard_defaults}\n", false, false},
|
||||
{"empty section", "tool_guard: {}\n", false, true},
|
||||
{"missing enabled", "tool_guard: {rules: []}\n", false, true},
|
||||
{"null enabled", "tool_guard: {enabled: null, rules: []}\n", false, true},
|
||||
{"missing rules while off", "tool_guard: {enabled: false}\n", false, true},
|
||||
{"missing rules while on", "tool_guard: {enabled: true}\n", false, true},
|
||||
{"null rules", "tool_guard: {enabled: false, rules: null}\n", false, true},
|
||||
{"mistyped enabled field", "tool_guard: {enable: false, rules: []}\n", false, true},
|
||||
{"malformed rules while off", "tool_guard: {enabled: false, rules: disabled}\n", false, true},
|
||||
{"malformed rule while off", "tool_guard: {enabled: false, rules: [invalid]}\n", false, true},
|
||||
{"invalid pattern", "tool_guard:\n enabled: false\n rules:\n - {id: invalid, name: invalid, enabled: false, pattern: '['}\n", false, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.yaml")
|
||||
if err := os.WriteFile(path, []byte(tc.yaml), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := Load(path)
|
||||
if (err != nil) != tc.wantErr {
|
||||
t.Fatalf("load error: %v", err)
|
||||
}
|
||||
if err == nil && cfg.EffectiveToolGuard().Enabled != tc.enabled {
|
||||
t.Fatal("wrong effective enabled state")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestBlockedExecutionPersistenceStatsAndReconciliation(t *testing.T) {
|
||||
db, conversationID, _ := setupProcessDetailsSummaryTest(t)
|
||||
now := time.Now()
|
||||
for _, status := range []string{"completed", "failed", "blocked", "cancelled"} {
|
||||
result := &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "policy message"}}, IsError: status != "completed", Blocked: status == "blocked"}
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: status, ToolName: "test", Status: status, Result: result, StartTime: now.Add(-time.Minute), EndTime: &now, ConversationID: conversationID}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := db.UpdateToolStats("test", 4, 1, 1, &now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateToolExecutionResult("blocked", &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "reduced output"}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
reloaded, err := db.GetToolExecution("blocked")
|
||||
if err != nil || reloaded.Status != "blocked" || !reloaded.Result.Blocked || !reloaded.Result.IsError || reloaded.Result.Content[0].Text != "reduced output" {
|
||||
t.Fatalf("reduction/storage lost blocked classification: %#v err=%v", reloaded, err)
|
||||
}
|
||||
count, err := db.CancelOrphanedRunningToolExecutions(now, "restart")
|
||||
if err != nil || count != 0 {
|
||||
t.Fatalf("terminal blocks reclassified as orphaned: count=%d err=%v", count, err)
|
||||
}
|
||||
page, err := db.LoadToolExecutionListPage(0, 10, "blocked", "")
|
||||
if err != nil || len(page) != 1 || page[0].ID != "blocked" {
|
||||
t.Fatalf("blocked status filter failed: %#v err=%v", page, err)
|
||||
}
|
||||
summary, err := db.LoadToolStatsSummary(1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if summary.Summary.TotalCalls != 4 || summary.Summary.SuccessCalls != 1 || summary.Summary.FailedCalls != 1 || summary.Summary.BlockedCalls != 1 || summary.TopTools[0].BlockedCalls != 1 {
|
||||
t.Fatalf("incorrect summary: %#v top=%#v", summary.Summary, summary.TopTools)
|
||||
}
|
||||
stats, err := db.LoadToolStats()
|
||||
if err != nil || stats["test"].BlockedCalls != 1 || stats["test"].FailedCalls != 1 {
|
||||
t.Fatalf("incorrect legacy stats: %#v err=%v", stats, err)
|
||||
}
|
||||
for _, daily := range []bool{false, true} {
|
||||
buckets, err := db.LoadCallsTimeline(now.Add(-time.Hour), daily)
|
||||
if err != nil || len(buckets) != 1 || buckets[0].Total != 4 || buckets[0].Failed != 1 || buckets[0].Blocked != 1 {
|
||||
t.Fatalf("incorrect timeline daily=%v: %#v err=%v", daily, buckets, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyToolGuardBlockMigrationIsStrictAndIdempotent(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "legacy-guard.db")
|
||||
db, err := NewDB(path, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Now()
|
||||
refusal := "工具调用已被安全规则拦截:识别到 example.gov,禁止操作。\n规则: 政府网站保护 (government-domains)\n匹配内容: \"example.gov\""
|
||||
for i, reason := range []string{
|
||||
refusal,
|
||||
"upstream returned: " + refusal,
|
||||
"工具调用已被安全规则拦截:regular error without the envelope",
|
||||
"工具调用已被安全规则拦截:malformed match\n规则: Rule (id)\n匹配内容: unquoted",
|
||||
} {
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: fmt.Sprint(i), ToolName: "test", Status: "failed", Error: reason, StartTime: now, EndTime: &now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := db.UpdateToolStats("test", 4, 0, 4, &now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for run := 0; run < 2; run++ {
|
||||
db, err = NewDB(path, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
exec, err := db.GetToolExecution("0")
|
||||
if err != nil || exec.Status != "blocked" || !exec.Result.Blocked || !exec.Result.IsError || exec.Result.Content[0].Text != refusal {
|
||||
t.Fatalf("migration did not retain refusal: %#v err=%v", exec, err)
|
||||
}
|
||||
stats, err := db.LoadToolStats()
|
||||
if err != nil || stats["test"].TotalCalls != 4 || stats["test"].FailedCalls != 3 || stats["test"].BlockedCalls != 1 {
|
||||
t.Fatalf("migration run=%d stats=%#v err=%v", run, stats, err)
|
||||
}
|
||||
count, err := db.CountToolExecutions("failed", "")
|
||||
if err != nil || count != 3 {
|
||||
t.Fatalf("migration changed unrelated failures: count=%d err=%v", count, err)
|
||||
}
|
||||
if err := db.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolResultStatusFromPayloadDistinguishesBlocked(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
payload map[string]interface{}
|
||||
want string
|
||||
}{
|
||||
{map[string]interface{}{"blocked": true, "success": false, "isError": true}, "blocked"},
|
||||
{map[string]interface{}{"status": "blocked", "success": false}, "blocked"},
|
||||
{map[string]interface{}{"success": false, "isError": true, "result": "工具调用已被安全规则拦截"}, "failed"},
|
||||
{map[string]interface{}{"success": true}, "completed"},
|
||||
} {
|
||||
if got := toolResultStatusFromPayload(tc.payload, "tool_result"); got != tc.want {
|
||||
t.Fatalf("payload=%#v status=%s want=%s", tc.payload, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
+142
-131
@@ -665,81 +665,6 @@ func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
|
||||
return conversations, rows.Err()
|
||||
}
|
||||
|
||||
const ungroupedConversationsSQL = `
|
||||
FROM conversations c
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM conversation_group_mappings cgm WHERE cgm.conversation_id = c.id
|
||||
)`
|
||||
|
||||
// CountUngroupedConversations 统计不在任何分组中的对话数量。
|
||||
func (db *DB) CountUngroupedConversations(projectID string) (int, error) {
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
var count int
|
||||
if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil {
|
||||
return 0, fmt.Errorf("统计未分组对话失败: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (db *DB) CountUngroupedConversationsForAccess(projectID, userID, scope string) (int, error) {
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
var count int
|
||||
if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil {
|
||||
return 0, fmt.Errorf("统计未分组对话失败: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
|
||||
func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID string) ([]*Conversation, error) {
|
||||
orderClause := conversationOrderClause(sortBy, "c")
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
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, c.role_name, c.agent_mode `+
|
||||
where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
args...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
|
||||
return db.ListUngroupedConversations(limit, offset, sortBy, projectID)
|
||||
}
|
||||
orderClause := conversationOrderClause(sortBy, "c")
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
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, c.role_name, c.agent_mode `+
|
||||
where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
args...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
// GetConversationTitle 获取对话标题(轻量查询,不加载消息)
|
||||
func (db *DB) GetConversationTitle(id string) (string, error) {
|
||||
var title string
|
||||
@@ -766,6 +691,22 @@ func (db *DB) UpdateConversationTitle(id, title string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (db *DB) UpdateConversationPinned(id string, pinned bool) error {
|
||||
pinnedValue := 0
|
||||
if pinned {
|
||||
pinnedValue = 1
|
||||
}
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversations SET pinned = ?, updated_at = ? WHERE id = ?",
|
||||
pinnedValue, time.Now(), id,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新对话置顶状态失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateConversationTime 更新对话时间
|
||||
func (db *DB) UpdateConversationTime(id string) error {
|
||||
_, err := db.Exec(
|
||||
@@ -784,7 +725,6 @@ func (db *DB) UpdateConversationTime(id string) error {
|
||||
// - process_details(过程详情)
|
||||
// - attack_chain_nodes(攻击链节点)
|
||||
// - attack_chain_edges(攻击链边)
|
||||
// - conversation_group_mappings(分组映射)
|
||||
// 漏洞记录会保留:vulnerabilities.conversation_id 使用 ON DELETE SET NULL,仅解除与会话的关联。
|
||||
// 注意:knowledge_retrieval_logs 在删除前会被显式清理。
|
||||
func (db *DB) DeleteConversation(id string) error {
|
||||
@@ -1350,6 +1290,8 @@ func (db *DB) AddProcessDetailWithID(messageID, conversationID, eventType, messa
|
||||
return "", fmt.Errorf("添加过程详情失败: %w", err)
|
||||
}
|
||||
|
||||
db.maybeRecordModelTokenUsage(messageID, conversationID, id, eventType, data)
|
||||
|
||||
return id, nil
|
||||
}
|
||||
|
||||
@@ -1538,6 +1480,11 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
return nil, fmt.Errorf("统计工具调用详情失败: %w", err)
|
||||
}
|
||||
|
||||
pendingToolStatus := "result_missing"
|
||||
if summary.Status == "running" {
|
||||
pendingToolStatus = "running"
|
||||
}
|
||||
|
||||
execRows, err := db.Query(
|
||||
"SELECT id, event_type, data FROM process_details WHERE message_id = ? AND event_type IN ('tool_call', 'tool_result') ORDER BY created_at ASC, rowid ASC",
|
||||
messageID,
|
||||
@@ -1548,12 +1495,12 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
seenExecIDs := make(map[string]bool)
|
||||
// A provider may reuse a fallback toolCallId across streaming rounds. Keep a
|
||||
// FIFO per ID instead of a single index so every persisted call gets at most
|
||||
// one result. Results without a stable ID are kept separate instead of being
|
||||
// guessed by order; showing no link is safer than linking to the wrong tool.
|
||||
// one result. ID-less results still attach to an unmatched call with the same
|
||||
// tool name (parallel nmap 1/2, 2/2 often lose one ID); different tools stay
|
||||
// unlinked so a leftover preview cannot steal another call's slot.
|
||||
toolIndexesByCallID := make(map[string][]int)
|
||||
lastMatchedToolIndexByCallID := make(map[string]int)
|
||||
matchedToolIndexes := make([]bool, 0)
|
||||
nextUnmatchedToolIdx := 0
|
||||
for execRows.Next() {
|
||||
var detailID string
|
||||
var eventType string
|
||||
@@ -1569,33 +1516,19 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
if err := json.Unmarshal([]byte(dataJSON), &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
toolName, _ := payload["toolName"].(string)
|
||||
toolName = strings.TrimSpace(toolName)
|
||||
toolCallID, _ := payload["toolCallId"].(string)
|
||||
toolCallID = strings.TrimSpace(toolCallID)
|
||||
execID, _ := payload["executionId"].(string)
|
||||
execID = strings.TrimSpace(execID)
|
||||
status := ""
|
||||
if eventType == "tool_result" {
|
||||
if success, ok := payload["success"].(bool); ok {
|
||||
if success {
|
||||
status = "completed"
|
||||
} else {
|
||||
status = "failed"
|
||||
}
|
||||
} else if isErr, ok := payload["isError"].(bool); ok && isErr {
|
||||
status = "failed"
|
||||
}
|
||||
}
|
||||
toolName := processDetailString(payload, "toolName")
|
||||
toolCallID := processDetailString(payload, "toolCallId")
|
||||
execID := processDetailString(payload, "executionId")
|
||||
status := toolResultStatusFromPayload(payload, eventType)
|
||||
if eventType == "tool_call" {
|
||||
summary.ToolExecutions = append(summary.ToolExecutions, ProcessDetailsToolExecution{
|
||||
ProcessDetailID: strings.TrimSpace(detailID),
|
||||
ToolName: toolName,
|
||||
ToolCallID: toolCallID,
|
||||
// This summary is reconstructed from persisted history, not live
|
||||
// execution state. Until a matching result is found the honest state
|
||||
// is "result_missing", never "running".
|
||||
Status: "result_missing",
|
||||
// This summary is reconstructed from persisted history. For an
|
||||
// active assistant turn, a missing result means the call is still
|
||||
// pending; after the turn is terminal it is genuinely incomplete.
|
||||
Status: pendingToolStatus,
|
||||
})
|
||||
matchedToolIndexes = append(matchedToolIndexes, false)
|
||||
if toolCallID != "" {
|
||||
@@ -1603,36 +1536,14 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
}
|
||||
}
|
||||
if eventType == "tool_result" {
|
||||
idx := -1
|
||||
if toolCallID != "" {
|
||||
queue := toolIndexesByCallID[toolCallID]
|
||||
for len(queue) > 0 {
|
||||
candidate := queue[0]
|
||||
queue = queue[1:]
|
||||
if candidate >= 0 && candidate < len(matchedToolIndexes) && !matchedToolIndexes[candidate] {
|
||||
idx = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
toolIndexesByCallID[toolCallID] = queue
|
||||
if idx < 0 {
|
||||
// Multiple persisted result events for one call (for example an
|
||||
// agent-facing reduced result replacing an earlier preview) update
|
||||
// that call instead of consuming an unrelated FIFO entry.
|
||||
if previous, ok := lastMatchedToolIndexByCallID[toolCallID]; ok {
|
||||
idx = previous
|
||||
}
|
||||
}
|
||||
}
|
||||
if idx < 0 && toolCallID != "" {
|
||||
for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] {
|
||||
nextUnmatchedToolIdx++
|
||||
}
|
||||
if nextUnmatchedToolIdx < len(matchedToolIndexes) {
|
||||
idx = nextUnmatchedToolIdx
|
||||
nextUnmatchedToolIdx++
|
||||
}
|
||||
}
|
||||
idx := matchToolExecutionIndex(
|
||||
summary.ToolExecutions,
|
||||
matchedToolIndexes,
|
||||
toolCallID,
|
||||
toolName,
|
||||
toolIndexesByCallID,
|
||||
lastMatchedToolIndexByCallID,
|
||||
)
|
||||
if idx >= 0 && idx < len(summary.ToolExecutions) {
|
||||
matchedToolIndexes[idx] = true
|
||||
if toolCallID != "" {
|
||||
@@ -1648,6 +1559,8 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
summary.ToolExecutions[idx].ExecutionID = execID
|
||||
if status != "" {
|
||||
summary.ToolExecutions[idx].Status = status
|
||||
} else {
|
||||
summary.ToolExecutions[idx].Status = "completed"
|
||||
}
|
||||
} else {
|
||||
summary.ToolExecutions = append(summary.ToolExecutions, ProcessDetailsToolExecution{
|
||||
@@ -1670,6 +1583,7 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
return nil, fmt.Errorf("遍历工具执行摘要失败: %w", err)
|
||||
}
|
||||
execRows.Close()
|
||||
db.applyPersistedToolExecutionStatuses(summary.ToolExecutions)
|
||||
|
||||
rows, err := db.Query(
|
||||
"SELECT data FROM process_details WHERE message_id = ? AND event_type = 'iteration' ORDER BY created_at ASC, rowid ASC",
|
||||
@@ -1704,6 +1618,103 @@ LIMIT 1`, messageID).Scan(&terminalEvent, &terminalCreatedAt)
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
func processDetailString(payload map[string]interface{}, key string) string {
|
||||
if payload == nil {
|
||||
return ""
|
||||
}
|
||||
v, ok := payload[key]
|
||||
if !ok || v == nil {
|
||||
return ""
|
||||
}
|
||||
s := strings.TrimSpace(fmt.Sprint(v))
|
||||
if s == "" || s == "<nil>" {
|
||||
return ""
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func toolResultStatusFromPayload(payload map[string]interface{}, eventType string) string {
|
||||
if eventType != "tool_result" {
|
||||
return ""
|
||||
}
|
||||
if blocked, _ := payload["blocked"].(bool); blocked || strings.EqualFold(processDetailString(payload, "status"), "blocked") {
|
||||
return "blocked"
|
||||
}
|
||||
if status := processDetailString(payload, "status"); strings.EqualFold(status, "background_running") {
|
||||
return "background_running"
|
||||
}
|
||||
if success, ok := payload["success"].(bool); ok {
|
||||
if success {
|
||||
return "completed"
|
||||
}
|
||||
return "failed"
|
||||
}
|
||||
if isErr, ok := payload["isError"].(bool); ok && isErr {
|
||||
return "failed"
|
||||
}
|
||||
return "completed"
|
||||
}
|
||||
|
||||
func (db *DB) applyPersistedToolExecutionStatuses(executions []ProcessDetailsToolExecution) {
|
||||
for i := range executions {
|
||||
execID := strings.TrimSpace(executions[i].ExecutionID)
|
||||
if execID == "" {
|
||||
continue
|
||||
}
|
||||
var status string
|
||||
if err := db.QueryRow(`SELECT status FROM tool_executions WHERE id = ?`, execID).Scan(&status); err != nil {
|
||||
continue
|
||||
}
|
||||
status = strings.ToLower(strings.TrimSpace(status))
|
||||
if status == "" {
|
||||
continue
|
||||
}
|
||||
executions[i].Status = status
|
||||
}
|
||||
}
|
||||
|
||||
func matchToolExecutionIndex(
|
||||
executions []ProcessDetailsToolExecution,
|
||||
matched []bool,
|
||||
toolCallID, toolName string,
|
||||
toolIndexesByCallID map[string][]int,
|
||||
lastMatchedToolIndexByCallID map[string]int,
|
||||
) int {
|
||||
if toolCallID != "" {
|
||||
queue := toolIndexesByCallID[toolCallID]
|
||||
for len(queue) > 0 {
|
||||
candidate := queue[0]
|
||||
queue = queue[1:]
|
||||
if candidate >= 0 && candidate < len(matched) && !matched[candidate] {
|
||||
toolIndexesByCallID[toolCallID] = queue
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
toolIndexesByCallID[toolCallID] = queue
|
||||
if previous, ok := lastMatchedToolIndexByCallID[toolCallID]; ok {
|
||||
return previous
|
||||
}
|
||||
}
|
||||
if toolName != "" {
|
||||
for i := range matched {
|
||||
if matched[i] {
|
||||
continue
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(executions[i].ToolName), toolName) {
|
||||
return i
|
||||
}
|
||||
}
|
||||
}
|
||||
if toolCallID != "" {
|
||||
for i := range matched {
|
||||
if !matched[i] {
|
||||
return i
|
||||
}
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// GetProcessDetailsPage 分页获取消息的过程详情(按时间升序)。
|
||||
func (db *DB) GetProcessDetailsPage(messageID string, limit, offset int) ([]ProcessDetail, int, error) {
|
||||
var total int
|
||||
|
||||
@@ -155,6 +155,10 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("初始化表失败: %w", err)
|
||||
}
|
||||
if err := database.migrateLegacyToolGuardBlocks(); err != nil {
|
||||
_ = db.Close()
|
||||
return nil, fmt.Errorf("迁移历史安全拦截记录失败: %w", err)
|
||||
}
|
||||
database.startPassiveCheckpointLoop("conversations")
|
||||
|
||||
return database, nil
|
||||
@@ -216,6 +220,32 @@ func (db *DB) initTables() error {
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
// 创建模型 Token 用量表:process_details 负责时间线回放,本表负责结构化聚合统计。
|
||||
createModelTokenUsageTable := `
|
||||
CREATE TABLE IF NOT EXISTS model_token_usage (
|
||||
id TEXT PRIMARY KEY,
|
||||
process_detail_id TEXT NOT NULL UNIQUE,
|
||||
message_id TEXT NOT NULL,
|
||||
conversation_id TEXT NOT NULL,
|
||||
project_id TEXT,
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
orchestration TEXT NOT NULL DEFAULT '',
|
||||
reason TEXT NOT NULL DEFAULT '',
|
||||
model TEXT NOT NULL DEFAULT '',
|
||||
model_calls INTEGER NOT NULL DEFAULT 0,
|
||||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (process_detail_id) REFERENCES process_details(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
|
||||
);`
|
||||
|
||||
// 创建工具执行记录表
|
||||
createToolExecutionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS tool_executions (
|
||||
@@ -303,29 +333,6 @@ func (db *DB) initTables() error {
|
||||
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE SET NULL
|
||||
);`
|
||||
|
||||
// 创建对话分组表
|
||||
createConversationGroupsTable := `
|
||||
CREATE TABLE IF NOT EXISTS conversation_groups (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
icon TEXT,
|
||||
owner_user_id TEXT,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL
|
||||
);`
|
||||
|
||||
// 创建对话分组映射表
|
||||
createConversationGroupMappingsTable := `
|
||||
CREATE TABLE IF NOT EXISTS conversation_group_mappings (
|
||||
id TEXT PRIMARY KEY,
|
||||
conversation_id TEXT NOT NULL,
|
||||
group_id TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (group_id) REFERENCES conversation_groups(id) ON DELETE CASCADE,
|
||||
UNIQUE(conversation_id, group_id)
|
||||
);`
|
||||
|
||||
// 机器人会话绑定表(用于跨重启保持「平台+租户+用户」到 conversation 的映射)
|
||||
createRobotUserSessionsTable := `
|
||||
CREATE TABLE IF NOT EXISTS robot_user_sessions (
|
||||
@@ -719,6 +726,10 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_process_details_message_id ON process_details(message_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_process_details_conversation_id ON process_details(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_model_token_usage_created_at ON model_token_usage(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_model_token_usage_conversation ON model_token_usage(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_model_token_usage_project ON model_token_usage(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_model_token_usage_model ON model_token_usage(model);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_executions_tool_name ON tool_executions(tool_name);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_executions_start_time ON tool_executions(start_time);
|
||||
CREATE INDEX IF NOT EXISTS idx_tool_executions_status ON tool_executions(status);
|
||||
@@ -729,8 +740,6 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_conversation ON knowledge_retrieval_logs(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_message ON knowledge_retrieval_logs(message_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_conversation ON conversation_group_mappings(conversation_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversation_group_mappings_group ON conversation_group_mappings(group_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_robot_user_sessions_updated_at ON robot_user_sessions(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id);
|
||||
@@ -806,6 +815,10 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建process_details表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createModelTokenUsageTable); err != nil {
|
||||
return fmt.Errorf("创建model_token_usage表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createToolExecutionsTable); err != nil {
|
||||
return fmt.Errorf("创建tool_executions表失败: %w", err)
|
||||
}
|
||||
@@ -830,13 +843,6 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建knowledge_retrieval_logs表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createConversationGroupsTable); err != nil {
|
||||
return fmt.Errorf("创建conversation_groups表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createConversationGroupMappingsTable); err != nil {
|
||||
return fmt.Errorf("创建conversation_group_mappings表失败: %w", err)
|
||||
}
|
||||
if _, err := db.Exec(createRobotUserSessionsTable); err != nil {
|
||||
return fmt.Errorf("创建robot_user_sessions表失败: %w", err)
|
||||
}
|
||||
@@ -932,16 +938,6 @@ func (db *DB) initTables() error {
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateConversationGroupsTable(); err != nil {
|
||||
db.logger.Warn("迁移conversation_groups表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateConversationGroupMappingsTable(); err != nil {
|
||||
db.logger.Warn("迁移conversation_group_mappings表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateBatchTaskQueuesTable(); err != nil {
|
||||
db.logger.Warn("迁移batch_task_queues表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
@@ -981,6 +977,10 @@ func (db *DB) initTables() error {
|
||||
if _, err := db.Exec(createIndexes); err != nil {
|
||||
return fmt.Errorf("创建索引失败: %w", err)
|
||||
}
|
||||
|
||||
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
|
||||
return fmt.Errorf("回填模型Token用量失败: %w", err)
|
||||
}
|
||||
db.logger.Debug("数据库表初始化完成")
|
||||
return nil
|
||||
}
|
||||
@@ -1199,54 +1199,6 @@ func (db *DB) migrateConversationsTable() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateConversationGroupsTable 迁移conversation_groups表,添加新字段
|
||||
func (db *DB) migrateConversationGroupsTable() error {
|
||||
// 检查pinned字段是否存在
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversation_groups') WHERE name='pinned'").Scan(&count)
|
||||
if err != nil {
|
||||
// 如果查询失败,尝试添加字段
|
||||
if _, addErr := db.Exec("ALTER TABLE conversation_groups ADD COLUMN pinned INTEGER DEFAULT 0"); addErr != nil {
|
||||
// 如果字段已存在,忽略错误
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加pinned字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if count == 0 {
|
||||
// 字段不存在,添加它
|
||||
if _, err := db.Exec("ALTER TABLE conversation_groups ADD COLUMN pinned INTEGER DEFAULT 0"); err != nil {
|
||||
db.logger.Warn("添加pinned字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateConversationGroupMappingsTable 迁移conversation_group_mappings表,添加新字段
|
||||
func (db *DB) migrateConversationGroupMappingsTable() error {
|
||||
// 检查pinned字段是否存在
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversation_group_mappings') WHERE name='pinned'").Scan(&count)
|
||||
if err != nil {
|
||||
// 如果查询失败,尝试添加字段
|
||||
if _, addErr := db.Exec("ALTER TABLE conversation_group_mappings ADD COLUMN pinned INTEGER DEFAULT 0"); addErr != nil {
|
||||
// 如果字段已存在,忽略错误
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加pinned字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if count == 0 {
|
||||
// 字段不存在,添加它
|
||||
if _, err := db.Exec("ALTER TABLE conversation_group_mappings ADD COLUMN pinned INTEGER DEFAULT 0"); err != nil {
|
||||
db.logger.Warn("添加pinned字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateBatchTaskQueuesTable 迁移batch_task_queues表,补充新字段
|
||||
func (db *DB) migrateBatchTaskQueuesTable() error {
|
||||
// 检查title字段是否存在
|
||||
|
||||
@@ -1,486 +0,0 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ConversationGroup 对话分组
|
||||
type ConversationGroup struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
OwnerUserID string `json:"-"`
|
||||
}
|
||||
|
||||
// GroupExistsByName 检查分组名称是否已存在
|
||||
func (db *DB) GroupExistsByName(name string, excludeID string) (bool, error) {
|
||||
return db.groupExistsByNameForOwner(name, excludeID, "")
|
||||
}
|
||||
|
||||
func (db *DB) groupExistsByNameForOwner(name, excludeID, ownerUserID string) (bool, error) {
|
||||
var count int
|
||||
var err error
|
||||
if ownerUserID != "" && excludeID != "" {
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ? AND id != ?", name, ownerUserID, excludeID).Scan(&count)
|
||||
} else if ownerUserID != "" {
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ?", name, ownerUserID).Scan(&count)
|
||||
} else if excludeID != "" {
|
||||
err = db.QueryRow(
|
||||
"SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND id != ?",
|
||||
name, excludeID,
|
||||
).Scan(&count)
|
||||
} else {
|
||||
err = db.QueryRow(
|
||||
"SELECT COUNT(*) FROM conversation_groups WHERE name = ?",
|
||||
name,
|
||||
).Scan(&count)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("检查分组名称失败: %w", err)
|
||||
}
|
||||
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
// CreateGroup 创建分组
|
||||
func (db *DB) CreateGroup(name, icon string, owners ...string) (*ConversationGroup, error) {
|
||||
ownerUserID := ""
|
||||
if len(owners) > 0 {
|
||||
ownerUserID = owners[0]
|
||||
}
|
||||
// 检查名称是否已存在
|
||||
exists, err := db.groupExistsByNameForOwner(name, "", ownerUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if exists {
|
||||
return nil, fmt.Errorf("分组名称已存在")
|
||||
}
|
||||
|
||||
id := uuid.New().String()
|
||||
now := time.Now()
|
||||
|
||||
if icon == "" {
|
||||
icon = "📁"
|
||||
}
|
||||
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversation_groups (id, name, icon, pinned, owner_user_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
id, name, icon, 0, ownerUserID, now, now,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建分组失败: %w", err)
|
||||
}
|
||||
|
||||
return &ConversationGroup{
|
||||
ID: id,
|
||||
Name: name,
|
||||
Icon: icon,
|
||||
Pinned: false,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
OwnerUserID: ownerUserID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ListGroups 列出所有分组
|
||||
func (db *DB) ListGroups() ([]*ConversationGroup, error) {
|
||||
return db.ListGroupsForAccess("", RBACScopeAll)
|
||||
}
|
||||
|
||||
func (db *DB) ListGroupsForAccess(userID, scope string) ([]*ConversationGroup, error) {
|
||||
query := "SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups"
|
||||
args := []interface{}{}
|
||||
if scope != RBACScopeAll {
|
||||
query += " WHERE owner_user_id = ?"
|
||||
args = append(args, userID)
|
||||
}
|
||||
query += " ORDER BY COALESCE(pinned, 0) DESC, created_at ASC"
|
||||
rows, err := db.Query(
|
||||
query, args...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询分组列表失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var groups []*ConversationGroup
|
||||
for rows.Next() {
|
||||
var group ConversationGroup
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
if err := rows.Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt); err != nil {
|
||||
return nil, fmt.Errorf("扫描分组失败: %w", err)
|
||||
}
|
||||
|
||||
group.Pinned = pinned != 0
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if err1 != nil {
|
||||
group.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
|
||||
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||
if err2 != nil {
|
||||
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||
}
|
||||
if err2 != nil {
|
||||
group.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
|
||||
groups = append(groups, &group)
|
||||
}
|
||||
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
// GetGroup 获取分组
|
||||
func (db *DB) GetGroup(id string) (*ConversationGroup, error) {
|
||||
var group ConversationGroup
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
err := db.QueryRow(
|
||||
"SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups WHERE id = ?",
|
||||
id,
|
||||
).Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("分组不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("查询分组失败: %w", err)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
group.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if err1 != nil {
|
||||
group.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
|
||||
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||
if err2 != nil {
|
||||
group.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||
}
|
||||
if err2 != nil {
|
||||
group.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
|
||||
group.Pinned = pinned != 0
|
||||
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (db *DB) UserCanAccessGroup(userID, scope, groupID string) bool {
|
||||
if scope == RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
var count int
|
||||
err := db.QueryRow(`SELECT COUNT(*) FROM conversation_groups WHERE id = ? AND owner_user_id = ?`, groupID, userID).Scan(&count)
|
||||
return err == nil && count > 0
|
||||
}
|
||||
|
||||
// UpdateGroup 更新分组
|
||||
func (db *DB) UpdateGroup(id, name, icon string) error {
|
||||
existing, err := db.GetGroup(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 检查名称是否已存在(排除当前分组)
|
||||
exists, err := db.groupExistsByNameForOwner(name, id, existing.OwnerUserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return fmt.Errorf("分组名称已存在")
|
||||
}
|
||||
|
||||
_, err = db.Exec(
|
||||
"UPDATE conversation_groups SET name = ?, icon = ?, updated_at = ? WHERE id = ?",
|
||||
name, icon, time.Now(), id,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新分组失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteGroup 删除分组
|
||||
func (db *DB) DeleteGroup(id string) error {
|
||||
_, err := db.Exec("DELETE FROM conversation_groups WHERE id = ?", id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除分组失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddConversationToGroup 将对话添加到分组
|
||||
// 注意:一个对话只能属于一个分组,所以在添加新分组之前,会先删除该对话的所有旧分组关联
|
||||
func (db *DB) AddConversationToGroup(conversationID, groupID string) error {
|
||||
// 先删除该对话的所有旧分组关联,确保一个对话只属于一个分组
|
||||
_, err := db.Exec(
|
||||
"DELETE FROM conversation_group_mappings WHERE conversation_id = ?",
|
||||
conversationID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除对话旧分组关联失败: %w", err)
|
||||
}
|
||||
|
||||
// 然后插入新的分组关联
|
||||
id := uuid.New().String()
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversation_group_mappings (id, conversation_id, group_id, created_at) VALUES (?, ?, ?, ?)",
|
||||
id, conversationID, groupID, time.Now(),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("添加对话到分组失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RemoveConversationFromGroup 从分组中移除对话
|
||||
func (db *DB) RemoveConversationFromGroup(conversationID, groupID string) error {
|
||||
_, err := db.Exec(
|
||||
"DELETE FROM conversation_group_mappings WHERE conversation_id = ? AND group_id = ?",
|
||||
conversationID, groupID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("从分组中移除对话失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetConversationsByGroup 获取分组中的所有对话
|
||||
func (db *DB) GetConversationsByGroup(groupID string) ([]*Conversation, error) {
|
||||
rows, err := db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, COALESCE(cgm.pinned, 0) as group_pinned
|
||||
FROM conversations c
|
||||
INNER JOIN conversation_group_mappings cgm ON c.id = cgm.conversation_id
|
||||
WHERE cgm.group_id = ?
|
||||
ORDER BY COALESCE(cgm.pinned, 0) DESC, c.updated_at DESC`,
|
||||
groupID,
|
||||
)
|
||||
if err != nil {
|
||||
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 groupPinned int
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &groupPinned); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
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
|
||||
}
|
||||
|
||||
// SearchConversationsByGroup 搜索分组中的对话(按标题和消息内容模糊匹配)
|
||||
func (db *DB) SearchConversationsByGroup(groupID string, searchQuery string) ([]*Conversation, error) {
|
||||
// 构建SQL查询,支持按标题和消息内容搜索
|
||||
// 使用 DISTINCT 避免因为一个对话有多条匹配消息而重复
|
||||
query := `SELECT DISTINCT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, COALESCE(cgm.pinned, 0) as group_pinned
|
||||
FROM conversations c
|
||||
INNER JOIN conversation_group_mappings cgm ON c.id = cgm.conversation_id
|
||||
WHERE cgm.group_id = ?`
|
||||
|
||||
args := []interface{}{groupID}
|
||||
|
||||
// 如果有搜索关键词,添加标题和消息内容搜索条件
|
||||
if searchQuery != "" {
|
||||
searchPattern := "%" + searchQuery + "%"
|
||||
// 搜索标题或消息内容
|
||||
// 使用 LEFT JOIN 连接消息表,这样即使没有消息的对话也能被搜索到(通过标题)
|
||||
query += ` AND (
|
||||
LOWER(c.title) LIKE LOWER(?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM messages m
|
||||
WHERE m.conversation_id = c.id
|
||||
AND LOWER(m.content) LIKE LOWER(?)
|
||||
)
|
||||
)`
|
||||
args = append(args, searchPattern, searchPattern)
|
||||
}
|
||||
|
||||
query += " ORDER BY COALESCE(cgm.pinned, 0) DESC, c.updated_at DESC"
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
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 groupPinned int
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &groupPinned); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
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
|
||||
}
|
||||
|
||||
// GetGroupByConversation 获取对话所属的分组
|
||||
func (db *DB) GetGroupByConversation(conversationID string) (string, error) {
|
||||
var groupID string
|
||||
err := db.QueryRow(
|
||||
"SELECT group_id FROM conversation_group_mappings WHERE conversation_id = ? LIMIT 1",
|
||||
conversationID,
|
||||
).Scan(&groupID)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return "", nil // 没有分组
|
||||
}
|
||||
return "", fmt.Errorf("查询对话分组失败: %w", err)
|
||||
}
|
||||
return groupID, nil
|
||||
}
|
||||
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (db *DB) UpdateConversationPinned(id string, pinned bool) error {
|
||||
pinnedValue := 0
|
||||
if pinned {
|
||||
pinnedValue = 1
|
||||
}
|
||||
// 注意:不更新 updated_at,因为置顶操作不应该改变对话的更新时间
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversations SET pinned = ? WHERE id = ?",
|
||||
pinnedValue, id,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新对话置顶状态失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateGroupPinned 更新分组置顶状态
|
||||
func (db *DB) UpdateGroupPinned(id string, pinned bool) error {
|
||||
pinnedValue := 0
|
||||
if pinned {
|
||||
pinnedValue = 1
|
||||
}
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversation_groups SET pinned = ?, updated_at = ? WHERE id = ?",
|
||||
pinnedValue, time.Now(), id,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新分组置顶状态失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GroupMapping 分组映射关系
|
||||
type GroupMapping struct {
|
||||
ConversationID string `json:"conversationId"`
|
||||
GroupID string `json:"groupId"`
|
||||
}
|
||||
|
||||
// GetAllGroupMappings 批量获取所有分组映射(消除 N+1 查询)
|
||||
func (db *DB) GetAllGroupMappings() ([]GroupMapping, error) {
|
||||
rows, err := db.Query("SELECT conversation_id, group_id FROM conversation_group_mappings")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询分组映射失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var mappings []GroupMapping
|
||||
for rows.Next() {
|
||||
var m GroupMapping
|
||||
if err := rows.Scan(&m.ConversationID, &m.GroupID); err != nil {
|
||||
return nil, fmt.Errorf("扫描分组映射失败: %w", err)
|
||||
}
|
||||
mappings = append(mappings, m)
|
||||
}
|
||||
|
||||
if mappings == nil {
|
||||
mappings = []GroupMapping{}
|
||||
}
|
||||
return mappings, nil
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedInGroup 更新对话在分组中的置顶状态
|
||||
func (db *DB) UpdateConversationPinnedInGroup(conversationID, groupID string, pinned bool) error {
|
||||
pinnedValue := 0
|
||||
if pinned {
|
||||
pinnedValue = 1
|
||||
}
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversation_group_mappings SET pinned = ? WHERE conversation_id = ? AND group_id = ?",
|
||||
pinnedValue, conversationID, groupID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新分组对话置顶状态失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,485 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const modelTokenUsageEventType = "eino_usage_summary"
|
||||
|
||||
// ModelTokenUsage records one model-usage summary emitted by an Agent run.
|
||||
type ModelTokenUsage struct {
|
||||
ID string `json:"id"`
|
||||
ProcessDetailID string `json:"processDetailId"`
|
||||
MessageID string `json:"messageId"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
Source string `json:"source"`
|
||||
Orchestration string `json:"orchestration"`
|
||||
Reason string `json:"reason"`
|
||||
Model string `json:"model,omitempty"`
|
||||
ModelCalls int64 `json:"modelCalls"`
|
||||
PromptTokens int64 `json:"promptTokens"`
|
||||
CompletionTokens int64 `json:"completionTokens"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
CachedTokens int64 `json:"cachedTokens"`
|
||||
ReasoningTokens int64 `json:"reasoningTokens"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
// ModelTokenUsageSummary is the aggregate shape used by dashboard and APIs.
|
||||
type ModelTokenUsageSummary struct {
|
||||
Events int64 `json:"events"`
|
||||
ModelCalls int64 `json:"modelCalls"`
|
||||
PromptTokens int64 `json:"promptTokens"`
|
||||
CompletionTokens int64 `json:"completionTokens"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
CachedTokens int64 `json:"cachedTokens"`
|
||||
ReasoningTokens int64 `json:"reasoningTokens"`
|
||||
}
|
||||
|
||||
// ModelTokenUsageBreakdown is a grouped aggregate row.
|
||||
type ModelTokenUsageBreakdown struct {
|
||||
Key string `json:"key"`
|
||||
Label string `json:"label,omitempty"`
|
||||
Events int64 `json:"events"`
|
||||
ModelCalls int64 `json:"modelCalls"`
|
||||
PromptTokens int64 `json:"promptTokens"`
|
||||
CompletionTokens int64 `json:"completionTokens"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
CachedTokens int64 `json:"cachedTokens"`
|
||||
ReasoningTokens int64 `json:"reasoningTokens"`
|
||||
}
|
||||
|
||||
// ModelTokenUsageStats is a compact API response for usage dashboards.
|
||||
type ModelTokenUsageStats struct {
|
||||
Summary ModelTokenUsageSummary `json:"summary"`
|
||||
Today ModelTokenUsageSummary `json:"today"`
|
||||
ByDay []ModelTokenUsageBreakdown `json:"byDay"`
|
||||
ByModel []ModelTokenUsageBreakdown `json:"byModel"`
|
||||
ByOrchestration []ModelTokenUsageBreakdown `json:"byOrchestration"`
|
||||
Recent []ModelTokenUsage `json:"recent"`
|
||||
}
|
||||
|
||||
// ModelTokenUsageFilter scopes usage queries.
|
||||
type ModelTokenUsageFilter struct {
|
||||
ConversationID string
|
||||
ProjectID string
|
||||
Since time.Time
|
||||
Until time.Time
|
||||
Days int
|
||||
Access RBACListAccess
|
||||
Limit int
|
||||
}
|
||||
|
||||
func modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID string, data interface{}) (ModelTokenUsage, bool) {
|
||||
m := mapFromUsageData(data)
|
||||
if len(m) == 0 {
|
||||
return ModelTokenUsage{}, false
|
||||
}
|
||||
usage := ModelTokenUsage{
|
||||
ID: uuid.New().String(),
|
||||
ProcessDetailID: strings.TrimSpace(processDetailID),
|
||||
MessageID: strings.TrimSpace(messageID),
|
||||
ConversationID: strings.TrimSpace(conversationID),
|
||||
Source: strings.TrimSpace(fmt.Sprint(m["source"])),
|
||||
Orchestration: strings.TrimSpace(fmt.Sprint(m["orchestration"])),
|
||||
Reason: strings.TrimSpace(fmt.Sprint(m["reason"])),
|
||||
Model: strings.TrimSpace(fmt.Sprint(m["model"])),
|
||||
ModelCalls: usageInt64(m["modelCalls"]),
|
||||
PromptTokens: usageInt64(m["promptTokens"]),
|
||||
CompletionTokens: usageInt64(m["completionTokens"]),
|
||||
TotalTokens: usageInt64(m["totalTokens"]),
|
||||
CachedTokens: usageInt64(m["cachedTokens"]),
|
||||
ReasoningTokens: usageInt64(m["reasoningTokens"]),
|
||||
}
|
||||
if usage.TotalTokens == 0 && (usage.PromptTokens > 0 || usage.CompletionTokens > 0) {
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
}
|
||||
if usage.ProcessDetailID == "" || usage.MessageID == "" || usage.ConversationID == "" {
|
||||
return ModelTokenUsage{}, false
|
||||
}
|
||||
if usage.ModelCalls == 0 && usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 && usage.CachedTokens == 0 && usage.ReasoningTokens == 0 {
|
||||
return ModelTokenUsage{}, false
|
||||
}
|
||||
return usage, true
|
||||
}
|
||||
|
||||
func mapFromUsageData(data interface{}) map[string]interface{} {
|
||||
switch v := data.(type) {
|
||||
case nil:
|
||||
return nil
|
||||
case map[string]interface{}:
|
||||
return v
|
||||
case string:
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(v), &m); err == nil {
|
||||
return m
|
||||
}
|
||||
case []byte:
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(v, &m); err == nil {
|
||||
return m
|
||||
}
|
||||
default:
|
||||
raw, err := json.Marshal(v)
|
||||
if err == nil {
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(raw, &m); err == nil {
|
||||
return m
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func usageInt64(v interface{}) int64 {
|
||||
switch n := v.(type) {
|
||||
case int:
|
||||
return int64(n)
|
||||
case int8:
|
||||
return int64(n)
|
||||
case int16:
|
||||
return int64(n)
|
||||
case int32:
|
||||
return int64(n)
|
||||
case int64:
|
||||
return n
|
||||
case uint:
|
||||
return int64(n)
|
||||
case uint8:
|
||||
return int64(n)
|
||||
case uint16:
|
||||
return int64(n)
|
||||
case uint32:
|
||||
return int64(n)
|
||||
case uint64:
|
||||
if n > math.MaxInt64 {
|
||||
return math.MaxInt64
|
||||
}
|
||||
return int64(n)
|
||||
case float32:
|
||||
return int64(n)
|
||||
case float64:
|
||||
return int64(n)
|
||||
case json.Number:
|
||||
i, _ := n.Int64()
|
||||
return i
|
||||
case string:
|
||||
i, _ := strconv.ParseInt(strings.TrimSpace(n), 10, 64)
|
||||
return i
|
||||
default:
|
||||
i, _ := strconv.ParseInt(strings.TrimSpace(fmt.Sprint(v)), 10, 64)
|
||||
return i
|
||||
}
|
||||
}
|
||||
|
||||
func (db *DB) maybeRecordModelTokenUsage(messageID, conversationID, processDetailID, eventType string, data interface{}) {
|
||||
if db == nil || eventType != modelTokenUsageEventType {
|
||||
return
|
||||
}
|
||||
usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if err := db.UpsertModelTokenUsage(usage); err != nil && db.logger != nil {
|
||||
db.logger.Warn("保存模型Token用量失败",
|
||||
zap.String("processDetailId", processDetailID),
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// UpsertModelTokenUsage persists usage with process_detail_id idempotency.
|
||||
func (db *DB) UpsertModelTokenUsage(usage ModelTokenUsage) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("database is nil")
|
||||
}
|
||||
now := time.Now()
|
||||
createdAt := usage.CreatedAt
|
||||
if createdAt.IsZero() {
|
||||
createdAt = now
|
||||
}
|
||||
if usage.ID == "" {
|
||||
usage.ID = uuid.New().String()
|
||||
}
|
||||
var projectID sql.NullString
|
||||
if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, usage.ConversationID).Scan(&projectID); err != nil && err != sql.ErrNoRows {
|
||||
return fmt.Errorf("查询对话项目失败: %w", err)
|
||||
}
|
||||
projectValue := interface{}(nil)
|
||||
if projectID.Valid && strings.TrimSpace(projectID.String) != "" {
|
||||
projectValue = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO model_token_usage (
|
||||
id, process_detail_id, message_id, conversation_id, project_id,
|
||||
source, orchestration, reason, model, model_calls,
|
||||
prompt_tokens, completion_tokens, total_tokens, cached_tokens, reasoning_tokens,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(process_detail_id) DO UPDATE SET
|
||||
message_id = excluded.message_id,
|
||||
conversation_id = excluded.conversation_id,
|
||||
project_id = excluded.project_id,
|
||||
source = excluded.source,
|
||||
orchestration = excluded.orchestration,
|
||||
reason = excluded.reason,
|
||||
model = excluded.model,
|
||||
model_calls = excluded.model_calls,
|
||||
prompt_tokens = excluded.prompt_tokens,
|
||||
completion_tokens = excluded.completion_tokens,
|
||||
total_tokens = excluded.total_tokens,
|
||||
cached_tokens = excluded.cached_tokens,
|
||||
reasoning_tokens = excluded.reasoning_tokens,
|
||||
created_at = excluded.created_at,
|
||||
updated_at = excluded.updated_at`,
|
||||
usage.ID, usage.ProcessDetailID, usage.MessageID, usage.ConversationID, projectValue,
|
||||
usage.Source, usage.Orchestration, usage.Reason, usage.Model, usage.ModelCalls,
|
||||
usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens, usage.CachedTokens, usage.ReasoningTokens,
|
||||
createdAt, now,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("写入模型Token用量失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BackfillModelTokenUsageFromProcessDetails makes existing timeline usage events queryable.
|
||||
func (db *DB) BackfillModelTokenUsageFromProcessDetails() error {
|
||||
if db == nil {
|
||||
return nil
|
||||
}
|
||||
rows, err := db.Query(`
|
||||
SELECT pd.id, pd.message_id, pd.conversation_id, pd.data, pd.created_at
|
||||
FROM process_details pd
|
||||
LEFT JOIN model_token_usage mtu ON mtu.process_detail_id = pd.id
|
||||
WHERE pd.event_type = ?
|
||||
AND (mtu.id IS NULL OR mtu.created_at != pd.created_at)`, modelTokenUsageEventType)
|
||||
if err != nil {
|
||||
return fmt.Errorf("查询历史模型Token用量失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var processDetailID, messageID, conversationID string
|
||||
var data sql.NullString
|
||||
var createdAt string
|
||||
if err := rows.Scan(&processDetailID, &messageID, &conversationID, &data, &createdAt); err != nil {
|
||||
return fmt.Errorf("扫描历史模型Token用量失败: %w", err)
|
||||
}
|
||||
if !data.Valid {
|
||||
continue
|
||||
}
|
||||
usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data.String)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
usage.CreatedAt = parseModelTokenUsageTime(createdAt)
|
||||
if err := db.UpsertModelTokenUsage(usage); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return fmt.Errorf("遍历历史模型Token用量失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) GetModelTokenUsageStats(filter ModelTokenUsageFilter) (*ModelTokenUsageStats, error) {
|
||||
if db == nil {
|
||||
return nil, fmt.Errorf("database is nil")
|
||||
}
|
||||
if filter.Days <= 0 {
|
||||
filter.Days = 7
|
||||
}
|
||||
if filter.Limit <= 0 {
|
||||
filter.Limit = 10
|
||||
}
|
||||
where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
|
||||
summary, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
todayFilter := filter
|
||||
now := time.Now()
|
||||
todayFilter.Since = time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
|
||||
todayWhere, todayArgs := buildModelTokenUsageWhere(todayFilter, "mtu", "c")
|
||||
today, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+todayWhere, todayArgs...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byDay, err := db.queryModelTokenUsageBreakdown(
|
||||
"SELECT date(mtu.created_at) AS k, date(mtu.created_at) AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY date(mtu.created_at) ORDER BY k DESC LIMIT ?",
|
||||
append(args, filter.Days)...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byModel, err := db.queryModelTokenUsageBreakdown(
|
||||
"SELECT COALESCE(NULLIF(TRIM(mtu.model), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.model), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
|
||||
append(args, filter.Limit)...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byOrch, err := db.queryModelTokenUsageBreakdown(
|
||||
"SELECT COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
|
||||
append(args, filter.Limit)...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
recent, err := db.ListModelTokenUsage(filter)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ModelTokenUsageStats{
|
||||
Summary: summary,
|
||||
Today: today,
|
||||
ByDay: byDay,
|
||||
ByModel: byModel,
|
||||
ByOrchestration: byOrch,
|
||||
Recent: recent,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func modelTokenUsageSummarySelect(alias string) string {
|
||||
p := ""
|
||||
if alias != "" {
|
||||
p = alias + "."
|
||||
}
|
||||
return fmt.Sprintf(`COUNT(%sid),
|
||||
COALESCE(SUM(%smodel_calls), 0),
|
||||
COALESCE(SUM(%sprompt_tokens), 0),
|
||||
COALESCE(SUM(%scompletion_tokens), 0),
|
||||
COALESCE(SUM(%stotal_tokens), 0),
|
||||
COALESCE(SUM(%scached_tokens), 0),
|
||||
COALESCE(SUM(%sreasoning_tokens), 0)`, p, p, p, p, p, p, p)
|
||||
}
|
||||
|
||||
func buildModelTokenUsageWhere(filter ModelTokenUsageFilter, usageAlias, convAlias string) (string, []interface{}) {
|
||||
where := " WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
uPrefix := ""
|
||||
if usageAlias != "" {
|
||||
uPrefix = usageAlias + "."
|
||||
}
|
||||
if cid := strings.TrimSpace(filter.ConversationID); cid != "" {
|
||||
where += " AND " + uPrefix + "conversation_id = ?"
|
||||
args = append(args, cid)
|
||||
}
|
||||
where, args = appendConversationProjectFilter(where, args, filter.ProjectID, usageAlias)
|
||||
if !filter.Since.IsZero() {
|
||||
where += " AND " + uPrefix + "created_at >= ?"
|
||||
args = append(args, filter.Since)
|
||||
}
|
||||
if !filter.Until.IsZero() {
|
||||
where += " AND " + uPrefix + "created_at <= ?"
|
||||
args = append(args, filter.Until)
|
||||
}
|
||||
where, args = appendConversationAccessFilter(where, args, filter.Access.UserID, filter.Access.Scope, convAlias)
|
||||
return where, args
|
||||
}
|
||||
|
||||
func (db *DB) queryModelTokenUsageSummary(query string, args ...interface{}) (ModelTokenUsageSummary, error) {
|
||||
var s ModelTokenUsageSummary
|
||||
err := db.QueryRow(query, args...).Scan(
|
||||
&s.Events, &s.ModelCalls, &s.PromptTokens, &s.CompletionTokens,
|
||||
&s.TotalTokens, &s.CachedTokens, &s.ReasoningTokens,
|
||||
)
|
||||
if err != nil {
|
||||
return s, fmt.Errorf("查询模型Token用量汇总失败: %w", err)
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (db *DB) queryModelTokenUsageBreakdown(query string, args ...interface{}) ([]ModelTokenUsageBreakdown, error) {
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询模型Token用量分组失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []ModelTokenUsageBreakdown{}
|
||||
for rows.Next() {
|
||||
var row ModelTokenUsageBreakdown
|
||||
if err := rows.Scan(
|
||||
&row.Key, &row.Label, &row.Events, &row.ModelCalls, &row.PromptTokens,
|
||||
&row.CompletionTokens, &row.TotalTokens, &row.CachedTokens, &row.ReasoningTokens,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("扫描模型Token用量分组失败: %w", err)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("遍历模型Token用量分组失败: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (db *DB) ListModelTokenUsage(filter ModelTokenUsageFilter) ([]ModelTokenUsage, error) {
|
||||
if filter.Limit <= 0 {
|
||||
filter.Limit = 20
|
||||
}
|
||||
if filter.Limit > 500 {
|
||||
filter.Limit = 500
|
||||
}
|
||||
where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
|
||||
args = append(args, filter.Limit)
|
||||
rows, err := db.Query(`
|
||||
SELECT mtu.id, mtu.process_detail_id, mtu.message_id, mtu.conversation_id,
|
||||
COALESCE(mtu.project_id, ''), mtu.source, mtu.orchestration, mtu.reason, mtu.model,
|
||||
mtu.model_calls, mtu.prompt_tokens, mtu.completion_tokens, mtu.total_tokens,
|
||||
mtu.cached_tokens, mtu.reasoning_tokens, mtu.created_at, mtu.updated_at
|
||||
FROM model_token_usage mtu
|
||||
JOIN conversations c ON c.id = mtu.conversation_id`+where+`
|
||||
ORDER BY mtu.created_at DESC, mtu.rowid DESC
|
||||
LIMIT ?`, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询模型Token用量明细失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []ModelTokenUsage{}
|
||||
for rows.Next() {
|
||||
var u ModelTokenUsage
|
||||
var createdAt, updatedAt string
|
||||
if err := rows.Scan(
|
||||
&u.ID, &u.ProcessDetailID, &u.MessageID, &u.ConversationID, &u.ProjectID,
|
||||
&u.Source, &u.Orchestration, &u.Reason, &u.Model, &u.ModelCalls,
|
||||
&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens, &u.CachedTokens,
|
||||
&u.ReasoningTokens, &createdAt, &updatedAt,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("扫描模型Token用量明细失败: %w", err)
|
||||
}
|
||||
u.CreatedAt = parseModelTokenUsageTime(createdAt)
|
||||
u.UpdatedAt = parseModelTokenUsageTime(updatedAt)
|
||||
out = append(out, u)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("遍历模型Token用量明细失败: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func parseModelTokenUsageTime(s string) time.Time {
|
||||
for _, layout := range []string{
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05.999999-07:00",
|
||||
"2006-01-02 15:04:05",
|
||||
time.RFC3339Nano,
|
||||
time.RFC3339,
|
||||
} {
|
||||
if t, err := time.Parse(layout, strings.TrimSpace(s)); err == nil {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestModelTokenUsagePersistsFromUsageProcessDetail(t *testing.T) {
|
||||
db := newModelTokenUsageTestDB(t)
|
||||
conv, err := db.CreateConversation("usage", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
msg, err := db.AddMessage(conv.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(msg.ID, conv.ID, modelTokenUsageEventType, "usage", map[string]interface{}{
|
||||
"source": "eino",
|
||||
"orchestration": "deep",
|
||||
"reason": "final",
|
||||
"model": "gpt-test",
|
||||
"modelCalls": 2,
|
||||
"promptTokens": 10,
|
||||
"completionTokens": 3,
|
||||
"totalTokens": 13,
|
||||
"cachedTokens": 4,
|
||||
"reasoningTokens": 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail: %v", err)
|
||||
}
|
||||
|
||||
stats, err := db.GetModelTokenUsageStats(ModelTokenUsageFilter{})
|
||||
if err != nil {
|
||||
t.Fatalf("GetModelTokenUsageStats: %v", err)
|
||||
}
|
||||
if stats.Summary.Events != 1 || stats.Summary.ModelCalls != 2 || stats.Summary.TotalTokens != 13 || stats.Summary.CachedTokens != 4 || stats.Summary.ReasoningTokens != 1 {
|
||||
t.Fatalf("summary = %#v", stats.Summary)
|
||||
}
|
||||
if len(stats.ByModel) != 1 || stats.ByModel[0].Key != "gpt-test" || stats.ByModel[0].TotalTokens != 13 {
|
||||
t.Fatalf("by model = %#v", stats.ByModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelTokenUsageBackfillIsIdempotent(t *testing.T) {
|
||||
db := newModelTokenUsageTestDB(t)
|
||||
conv, err := db.CreateConversation("usage", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
msg, err := db.AddMessage(conv.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(msg.ID, conv.ID, modelTokenUsageEventType, "usage", map[string]interface{}{
|
||||
"source": "eino", "modelCalls": 1, "promptTokens": 7, "completionTokens": 5, "totalTokens": 12,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail: %v", err)
|
||||
}
|
||||
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
|
||||
t.Fatalf("Backfill 1: %v", err)
|
||||
}
|
||||
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
|
||||
t.Fatalf("Backfill 2: %v", err)
|
||||
}
|
||||
stats, err := db.GetModelTokenUsageStats(ModelTokenUsageFilter{})
|
||||
if err != nil {
|
||||
t.Fatalf("GetModelTokenUsageStats: %v", err)
|
||||
}
|
||||
if stats.Summary.Events != 1 || stats.Summary.TotalTokens != 12 {
|
||||
t.Fatalf("summary after backfill = %#v", stats.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func newModelTokenUsageTestDB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "usage.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db
|
||||
}
|
||||
@@ -91,6 +91,15 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
|
||||
if id == "" || result == nil {
|
||||
return nil
|
||||
}
|
||||
var status string
|
||||
if err := db.QueryRow(`SELECT status FROM tool_executions WHERE id = ?`, id).Scan(&status); err != nil && err != sql.ErrNoRows {
|
||||
return err
|
||||
}
|
||||
if status == mcp.ToolExecutionStatusBlocked {
|
||||
copy := *result
|
||||
copy.Blocked, copy.IsError = true, true
|
||||
result = ©
|
||||
}
|
||||
resultBytes, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -276,6 +285,7 @@ type ToolStatsSummary struct {
|
||||
TotalCalls int
|
||||
SuccessCalls int
|
||||
FailedCalls int
|
||||
BlockedCalls int
|
||||
LastCallTime *time.Time
|
||||
ToolCount int
|
||||
}
|
||||
@@ -304,6 +314,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
SELECT COUNT(*),
|
||||
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),
|
||||
COALESCE(SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), 0),
|
||||
MAX(start_time),
|
||||
COUNT(DISTINCT tool_name)
|
||||
FROM tool_executions
|
||||
@@ -313,6 +324,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
&result.Summary.TotalCalls,
|
||||
&result.Summary.SuccessCalls,
|
||||
&result.Summary.FailedCalls,
|
||||
&result.Summary.BlockedCalls,
|
||||
&lastCallRaw,
|
||||
&result.Summary.ToolCount,
|
||||
)
|
||||
@@ -334,6 +346,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
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,
|
||||
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked_calls,
|
||||
MAX(start_time) AS last_call_time
|
||||
FROM tool_executions
|
||||
GROUP BY tool_name
|
||||
@@ -354,6 +367,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
&stat.TotalCalls,
|
||||
&stat.SuccessCalls,
|
||||
&stat.FailedCalls,
|
||||
&stat.BlockedCalls,
|
||||
&lastCallTime,
|
||||
); err != nil {
|
||||
db.logger.Warn("加载 Top 工具统计失败", zap.Error(err))
|
||||
@@ -385,8 +399,9 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
||||
err := db.QueryRow(`SELECT COUNT(*),
|
||||
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),
|
||||
COALESCE(SUM(CASE WHEN status = 'blocked' 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,
|
||||
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls, &result.Summary.BlockedCalls,
|
||||
&lastCall, &result.Summary.ToolCount,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -398,7 +413,8 @@ 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', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END),
|
||||
SUM(CASE WHEN status = 'blocked' 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
|
||||
@@ -407,7 +423,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
||||
for rows.Next() {
|
||||
var stat mcp.ToolStats
|
||||
var last sql.NullString
|
||||
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &last); err != nil {
|
||||
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &stat.BlockedCalls, &last); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if last.Valid {
|
||||
@@ -916,8 +932,11 @@ func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
|
||||
// LoadToolStats 加载所有工具统计信息
|
||||
func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
|
||||
query := `
|
||||
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
|
||||
FROM tool_stats
|
||||
SELECT stats.tool_name, total_calls, success_calls, failed_calls, last_call_time,
|
||||
COALESCE(blocked.calls, 0)
|
||||
FROM tool_stats stats
|
||||
LEFT JOIN (SELECT tool_name, COUNT(*) AS calls FROM tool_executions WHERE status = 'blocked' GROUP BY tool_name) blocked
|
||||
ON blocked.tool_name = stats.tool_name
|
||||
`
|
||||
|
||||
rows, err := db.Query(query)
|
||||
@@ -937,6 +956,7 @@ func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
|
||||
&stat.SuccessCalls,
|
||||
&stat.FailedCalls,
|
||||
&lastCallTime,
|
||||
&stat.BlockedCalls,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Warn("加载统计信息失败", zap.Error(err))
|
||||
@@ -989,6 +1009,7 @@ type CallsTimelineBucket struct {
|
||||
BucketTime time.Time
|
||||
Total int
|
||||
Failed int
|
||||
Blocked int
|
||||
}
|
||||
|
||||
// truncateCallsTimelineBucket 将时间截断到趋势图桶边界(本地时区,与 handler 侧 truncateToBucket 一致)
|
||||
@@ -1008,7 +1029,8 @@ 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', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
|
||||
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
|
||||
FROM tool_executions
|
||||
WHERE start_time >= ?
|
||||
GROUP BY bucket
|
||||
@@ -1018,7 +1040,8 @@ 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', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
|
||||
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
|
||||
FROM tool_executions
|
||||
WHERE start_time >= ?
|
||||
GROUP BY bucket
|
||||
@@ -1035,8 +1058,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
||||
buckets := make([]CallsTimelineBucket, 0)
|
||||
for rows.Next() {
|
||||
var bucketStr string
|
||||
var total, failed int
|
||||
if err := rows.Scan(&bucketStr, &total, &failed); err != nil {
|
||||
var total, failed, blocked int
|
||||
if err := rows.Scan(&bucketStr, &total, &failed, &blocked); err != nil {
|
||||
db.logger.Warn("加载调用趋势失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
@@ -1049,6 +1072,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
||||
BucketTime: bucketTime,
|
||||
Total: total,
|
||||
Failed: failed,
|
||||
Blocked: blocked,
|
||||
})
|
||||
}
|
||||
return buckets, nil
|
||||
|
||||
@@ -8,7 +8,7 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) {
|
||||
func TestProcessDetailsSummaryDoesNotGuessIDLessResultsOntoDifferentTool(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
@@ -20,8 +20,8 @@ func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) {
|
||||
results := []map[string]interface{}{
|
||||
{"toolName": "http-framework-test", "toolCallId": "call-1", "success": true},
|
||||
{"toolName": "http-framework-test", "toolCallId": "call-2", "success": true},
|
||||
{"toolName": "http-framework-test", "success": true},
|
||||
{"toolName": "http-framework-test", "success": true},
|
||||
{"toolName": "other-tool", "success": true},
|
||||
{"toolName": "other-tool", "success": true},
|
||||
}
|
||||
var resultIDs []string
|
||||
for _, result := range results {
|
||||
@@ -53,12 +53,71 @@ func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) {
|
||||
}
|
||||
}
|
||||
for i, execution := range summary.ToolExecutions[4:] {
|
||||
if execution.Status != "completed" || execution.ToolCallID != "" {
|
||||
if execution.Status != "completed" || execution.ToolCallID != "" || execution.ToolName != "other-tool" {
|
||||
t.Fatalf("idless result %d = %#v, want separate completed result without toolCallId", i, execution)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryPairsIDLessResultsWithSameToolName(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
for i, id := range []string{"call-1", "call-2"} {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "nmap", "toolCallId": id, "index": i + 1, "total": 2,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
}
|
||||
var resultIDs []string
|
||||
for i := 0; i < 2; i++ {
|
||||
resultID, err := db.AddProcessDetailWithID(messageID, conversationID, "tool_result", "result", map[string]interface{}{
|
||||
"toolName": "nmap", "success": true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
resultIDs = append(resultIDs, resultID)
|
||||
}
|
||||
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 2 {
|
||||
t.Fatalf("tool executions = %d, want 2", len(summary.ToolExecutions))
|
||||
}
|
||||
for i, execution := range summary.ToolExecutions {
|
||||
if execution.Status != "completed" {
|
||||
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||
}
|
||||
if execution.ResultDetailID != resultIDs[i] {
|
||||
t.Fatalf("execution %d result detail id = %q, want %q", i, execution.ResultDetailID, resultIDs[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryPairedResultWithoutSuccessIsCompleted(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "nmap", "toolCallId": "call-1",
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", map[string]interface{}{
|
||||
"toolName": "nmap", "toolCallId": "call-1", "resultPreview": "open 22",
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "completed" {
|
||||
t.Fatalf("tool executions = %#v, want completed", summary.ToolExecutions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
for i := 0; i < 2; i++ {
|
||||
@@ -106,6 +165,32 @@ func TestProcessDetailsSummaryDoesNotReportPersistedOrphanAsRunning(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryReportsUnmatchedToolCallAsRunningForActiveTurn(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
if _, err := db.Exec(
|
||||
"UPDATE messages SET content = ?, updated_at = ? WHERE id = ?",
|
||||
"处理中...", "2026-08-10T08:00:00Z", messageID,
|
||||
); err != nil {
|
||||
t.Fatalf("update running message: %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "execute", "toolCallId": "pending",
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if summary.Status != "running" {
|
||||
t.Fatalf("summary status = %q, want running", summary.Status)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "running" {
|
||||
t.Fatalf("tool executions = %#v, want running", summary.ToolExecutions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryIncludesPersistedTurnTiming(t *testing.T) {
|
||||
db, _, messageID := setupProcessDetailsSummaryTest(t)
|
||||
startedAt := "2026-08-10T08:00:00Z"
|
||||
|
||||
@@ -201,7 +201,6 @@ func (db *DB) migrateRBACOwnershipColumns() error {
|
||||
{"webshell_connections", "owner_user_id", "ALTER TABLE webshell_connections ADD COLUMN owner_user_id TEXT"},
|
||||
{"batch_task_queues", "owner_user_id", "ALTER TABLE batch_task_queues ADD COLUMN owner_user_id TEXT"},
|
||||
{"c2_listeners", "owner_user_id", "ALTER TABLE c2_listeners ADD COLUMN owner_user_id TEXT"},
|
||||
{"conversation_groups", "owner_user_id", "ALTER TABLE conversation_groups ADD COLUMN owner_user_id TEXT"},
|
||||
{"tool_executions", "owner_user_id", "ALTER TABLE tool_executions ADD COLUMN owner_user_id TEXT"},
|
||||
{"tool_executions", "conversation_id", "ALTER TABLE tool_executions ADD COLUMN conversation_id TEXT"},
|
||||
} {
|
||||
|
||||
@@ -58,27 +58,8 @@ func TestRBACToolExecutionOwnershipAccess(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACGroupAndUploadOwnership(t *testing.T) {
|
||||
func TestRBACUploadOwnership(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
group1, err := db.CreateGroup("u1 group", "", "u1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
group2, err := db.CreateGroup("u2 group", "", "u2")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
groups, err := db.ListGroupsForAccess("u1", RBACScopeAssigned)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(groups) != 1 || groups[0].ID != group1.ID {
|
||||
t.Fatalf("groups = %#v, want only %s (not %s)", groups, group1.ID, group2.ID)
|
||||
}
|
||||
if db.UserCanAccessGroup("u1", RBACScopeAssigned, group2.ID) {
|
||||
t.Fatal("foreign group was accessible")
|
||||
}
|
||||
|
||||
conversation, err := db.CreateConversation("upload", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
)
|
||||
|
||||
const legacyToolGuardPrefix = "工具调用已被安全规则拦截"
|
||||
|
||||
// Only the exact envelope emitted by the old local guard is recognized here.
|
||||
// New executions use the structured marker and never infer policy from text.
|
||||
func isLegacyToolGuardRefusal(text string) bool {
|
||||
if !strings.HasPrefix(text, legacyToolGuardPrefix+":") && !strings.HasPrefix(text, legacyToolGuardPrefix+"\n规则: ") {
|
||||
return false
|
||||
}
|
||||
matchIndex := strings.LastIndex(text, "\n匹配内容: ")
|
||||
if matchIndex < 0 {
|
||||
return false
|
||||
}
|
||||
if _, err := strconv.Unquote(text[matchIndex+len("\n匹配内容: "):]); err != nil {
|
||||
return false
|
||||
}
|
||||
ruleIndex := strings.LastIndex(text[:matchIndex], "\n规则: ")
|
||||
if ruleIndex < 0 {
|
||||
return false
|
||||
}
|
||||
rule := text[ruleIndex+len("\n规则: ") : matchIndex]
|
||||
idIndex := strings.LastIndex(rule, " (")
|
||||
return idIndex > 0 && strings.HasSuffix(rule, ")") && len(rule[idIndex+2:len(rule)-1]) > 0 && !strings.Contains(rule, "\n")
|
||||
}
|
||||
|
||||
// migrateLegacyToolGuardBlocks is idempotent because only failed records qualify.
|
||||
// Keeping status and accumulated failure counts in one transaction makes monitor
|
||||
// filters, badges and statistics agree immediately after upgrading.
|
||||
func (db *DB) migrateLegacyToolGuardBlocks() error {
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
rows, err := tx.Query(`SELECT id, tool_name, error, COALESCE(result, '') FROM tool_executions WHERE status = 'failed' AND error LIKE ?`, legacyToolGuardPrefix+"%")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
type record struct{ id, tool, reason, result string }
|
||||
var records []record
|
||||
for rows.Next() {
|
||||
var r record
|
||||
if err := rows.Scan(&r.id, &r.tool, &r.reason, &r.result); err != nil {
|
||||
rows.Close()
|
||||
return err
|
||||
}
|
||||
if isLegacyToolGuardRefusal(r.reason) {
|
||||
records = append(records, r)
|
||||
}
|
||||
}
|
||||
err = rows.Err()
|
||||
rows.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, r := range records {
|
||||
var result mcp.ToolResult
|
||||
_ = json.Unmarshal([]byte(r.result), &result)
|
||||
if len(result.Content) == 0 {
|
||||
result.Content = []mcp.Content{{Type: "text", Text: r.reason}}
|
||||
}
|
||||
result.Blocked, result.IsError = true, true
|
||||
encoded, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE tool_executions SET status = 'blocked', result = ? WHERE id = ?`, string(encoded), r.id); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`UPDATE tool_stats SET failed_calls = MAX(0, failed_calls - 1) WHERE tool_name = ?`, r.tool); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
+55
-19
@@ -315,12 +315,13 @@ func (h *AgentHandler) SetHitlToolWhitelistSaver(s HitlToolWhitelistSaver) {
|
||||
h.hitlWhitelistSaver = s
|
||||
}
|
||||
|
||||
// HitlDefaultReviewerSaver 持久化全局默认审批方到 config.yaml。
|
||||
// HitlDefaultReviewerSaver 持久化全局默认人机协同配置到 config.yaml。
|
||||
type HitlDefaultReviewerSaver interface {
|
||||
UpdateHitlDefaultConfig(mode, reviewer string, timeoutSeconds int) error
|
||||
UpdateHitlDefaultReviewer(reviewer string) error
|
||||
}
|
||||
|
||||
// SetHitlDefaultReviewerSaver 设置 HITL 默认审批方落盘。
|
||||
// SetHitlDefaultReviewerSaver 设置 HITL 默认配置落盘。
|
||||
func (h *AgentHandler) SetHitlDefaultReviewerSaver(s HitlDefaultReviewerSaver) {
|
||||
h.hitlDefaultReviewerSaver = s
|
||||
}
|
||||
@@ -332,6 +333,35 @@ func (h *AgentHandler) hitlEffectiveDefaultReviewer() string {
|
||||
return "human"
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlEffectiveDefaultMode() string {
|
||||
if h != nil && h.config != nil {
|
||||
return normalizeHitlDefaultMode(h.config.Hitl.EffectiveDefaultMode())
|
||||
}
|
||||
return "off"
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlEffectiveDefaultTimeoutSeconds() int {
|
||||
if h != nil && h.config != nil {
|
||||
timeout := h.config.Hitl.EffectiveDefaultTimeoutSeconds()
|
||||
if timeout < 0 {
|
||||
return 0
|
||||
}
|
||||
return timeout
|
||||
}
|
||||
return 300
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlEffectiveDefaultRequest() *HITLRequest {
|
||||
mode := h.hitlEffectiveDefaultMode()
|
||||
return &HITLRequest{
|
||||
Enabled: mode != "off",
|
||||
Mode: mode,
|
||||
Reviewer: h.hitlEffectiveDefaultReviewer(),
|
||||
SensitiveTools: []string{},
|
||||
TimeoutSeconds: h.hitlEffectiveDefaultTimeoutSeconds(),
|
||||
}
|
||||
}
|
||||
|
||||
// HITLNeedsToolApproval 供 C2 危险任务门控:与会话侧人机协同及免审批白名单判定一致。
|
||||
func (h *AgentHandler) HITLNeedsToolApproval(conversationID, toolName string) bool {
|
||||
if h == nil || h.hitlManager == nil {
|
||||
@@ -698,25 +728,26 @@ func (h *AgentHandler) mergeAssistantMessagePartialOnCancel(messageID, partial s
|
||||
|
||||
// ChatResponse 聊天响应
|
||||
type ChatResponse struct {
|
||||
Response string `json:"response"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
||||
ConversationID string `json:"conversationId"` // 对话ID
|
||||
Time time.Time `json:"time"`
|
||||
Finalizable bool `json:"finalizable"`
|
||||
Finalized bool `json:"finalized"`
|
||||
Status string `json:"status,omitempty"`
|
||||
CompletionReason string `json:"completionReason,omitempty"`
|
||||
EvidenceVerified bool `json:"evidenceVerified"`
|
||||
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||
Response string `json:"response"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
||||
ConversationID string `json:"conversationId"` // 对话ID
|
||||
Time time.Time `json:"time"`
|
||||
Finalizable bool `json:"finalizable"`
|
||||
Finalized bool `json:"finalized"`
|
||||
Status string `json:"status,omitempty"`
|
||||
CompletionReason string `json:"completionReason,omitempty"`
|
||||
EvidenceVerified bool `json:"evidenceVerified"`
|
||||
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||
AutoCancelledPendingExecutionIDs []string `json:"autoCancelledPendingExecutionIds,omitempty"`
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(ctx) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
}
|
||||
errMsg := "执行失败: " + errMA.Error()
|
||||
errMsg := "执行失败: " + multiagent.EinoClientRunErrorMessage(errMA)
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
@@ -724,8 +755,13 @@ func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMes
|
||||
return "", conversationID, errMA
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput), true)
|
||||
func (h *AgentHandler) finalizeRobotAgentSuccess(taskCtx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||
reasoningContent := multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput)
|
||||
decision := h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, true)
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, true)
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, "robot", resultMA.MCPExecutionIDs, reasoningContent, decision)
|
||||
responseText := decision.FinalText
|
||||
if !decision.Finalizable {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
@@ -758,7 +794,7 @@ func (h *AgentHandler) runRobotEinoSingleWithRetry(
|
||||
*taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
return h.finalizeRobotAgentSuccess(taskCtx, assistantMessageID, conversationID, resultMA)
|
||||
}
|
||||
|
||||
func (h *AgentHandler) runRobotMultiAgentWithRetry(
|
||||
@@ -779,7 +815,7 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry(
|
||||
*taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
return h.finalizeRobotAgentSuccess(taskCtx, assistantMessageID, conversationID, resultMA)
|
||||
}
|
||||
|
||||
// ProcessMessageForRobot 供机器人(企业微信/钉钉/飞书)调用:Eino 单/多代理执行路径(含 progressCallback、过程详情),仅不发送 SSE,最后返回完整回复
|
||||
|
||||
@@ -281,7 +281,12 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
if useBatchMulti {
|
||||
agentMode = "batch_eino_" + batchOrch
|
||||
}
|
||||
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, reasoningContent, true)
|
||||
decision := h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, true)
|
||||
autoCancelledPendingExecutionIDs := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback)
|
||||
if len(autoCancelledPendingExecutionIDs) > 0 {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, true)
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, mcpIDs, reasoningContent, decision)
|
||||
resText := decision.FinalText
|
||||
if !decision.Finalizable {
|
||||
resText = finalizationBlockedMessage(decision)
|
||||
@@ -289,14 +294,15 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
sendEvent("finalization_check", resText, decision)
|
||||
}
|
||||
sendEvent("response", resText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": mcpIDs,
|
||||
"batchQueueId": queueID,
|
||||
"batchTaskId": task.ID,
|
||||
"batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable],
|
||||
"candidatePreview": safeTruncateString(resultMA.Response, 500),
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": mcpIDs,
|
||||
"batchQueueId": queueID,
|
||||
"batchTaskId": task.ID,
|
||||
"batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable],
|
||||
"candidatePreview": safeTruncateString(resultMA.Response, 500),
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
|
||||
if assistantMessageID == "" {
|
||||
@@ -385,7 +391,8 @@ func (h *AgentHandler) handleBatchSubTaskRunError(
|
||||
}
|
||||
|
||||
h.logger.Error("批量任务执行失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(runErr))
|
||||
errorMsg := "执行失败: " + runErr.Error()
|
||||
clientErr := multiagent.EinoClientRunErrorMessage(runErr)
|
||||
errorMsg := "执行失败: " + clientErr
|
||||
if assistantMessageID != "" {
|
||||
if _, updateErr := h.db.Exec(
|
||||
"UPDATE messages SET content = ?, updated_at = ? WHERE id = ?",
|
||||
@@ -398,5 +405,5 @@ func (h *AgentHandler) handleBatchSubTaskRunError(
|
||||
h.logger.Warn("保存错误详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", runErr.Error())
|
||||
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", clientErr)
|
||||
}
|
||||
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -95,6 +96,7 @@ type ConfigHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
toolGuard *toolguard.Manager
|
||||
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
|
||||
}
|
||||
|
||||
@@ -347,13 +349,13 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents))
|
||||
}
|
||||
multiPub := config.MultiAgentPublic{
|
||||
Enabled: h.config.MultiAgent.Enabled,
|
||||
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
|
||||
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
|
||||
Enabled: h.config.MultiAgent.Enabled,
|
||||
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
|
||||
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
|
||||
SummarizationUserIntentLedgerEntryMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective(),
|
||||
LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(),
|
||||
LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(),
|
||||
@@ -891,7 +893,14 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
if req.Hitl != nil {
|
||||
h.config.Hitl.AuditModel = req.Hitl.AuditModel
|
||||
h.config.Hitl.ToolWhitelist = mergeHitlToolWhitelistSlice(nil, req.Hitl.ToolWhitelist)
|
||||
if strings.TrimSpace(req.Hitl.DefaultMode) != "" {
|
||||
h.config.Hitl.DefaultMode = req.Hitl.EffectiveDefaultMode()
|
||||
}
|
||||
h.config.Hitl.DefaultReviewer = req.Hitl.EffectiveDefaultReviewer()
|
||||
if req.Hitl.DefaultTimeoutSeconds != nil {
|
||||
v := req.Hitl.EffectiveDefaultTimeoutSeconds()
|
||||
h.config.Hitl.DefaultTimeoutSeconds = &v
|
||||
}
|
||||
h.config.Hitl.AuditAgentPrompt = strings.TrimSpace(req.Hitl.AuditAgentPrompt)
|
||||
h.config.Hitl.AuditAgentPromptReviewEdit = strings.TrimSpace(req.Hitl.AuditAgentPromptReviewEdit)
|
||||
if req.Hitl.RetentionDays != nil {
|
||||
@@ -1162,6 +1171,8 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
h.config.NormalizeAIProviderProfiles()
|
||||
|
||||
// 保存配置到文件
|
||||
if err := h.saveConfig(); err != nil {
|
||||
h.logger.Error("保存配置失败", zap.Error(err))
|
||||
@@ -1737,6 +1748,10 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
|
||||
// saveConfig 保存配置到文件
|
||||
func (h *ConfigHandler) saveConfig() error {
|
||||
configFileMu.Lock()
|
||||
defer configFileMu.Unlock()
|
||||
h.config.NormalizeAIProviderProfiles()
|
||||
|
||||
// 读取现有配置文件并创建备份
|
||||
data, err := os.ReadFile(h.configPath)
|
||||
if err != nil {
|
||||
@@ -2141,12 +2156,35 @@ func updateHitlConfig(doc *yaml.Node, cfg config.HitlConfig) {
|
||||
setStringInMap(auditModelNode, "model", cfg.AuditModel.Model)
|
||||
// flow 样式 [a, b, c] 单行展示,工具多时比块序列省行数
|
||||
setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist)
|
||||
setStringInMap(hitlNode, "default_mode", cfg.EffectiveDefaultMode())
|
||||
setStringInMap(hitlNode, "default_reviewer", cfg.EffectiveDefaultReviewer())
|
||||
setIntInMap(hitlNode, "default_timeout_seconds", cfg.EffectiveDefaultTimeoutSeconds())
|
||||
setIntInMap(hitlNode, "retention_days", cfg.RetentionDaysEffective())
|
||||
setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt)
|
||||
setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit)
|
||||
}
|
||||
|
||||
// UpdateHitlDefaultConfig 更新全局默认人机协同配置并写入 config.yaml。
|
||||
func (h *ConfigHandler) UpdateHitlDefaultConfig(mode, reviewer string, timeoutSeconds int) error {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.config.Hitl.DefaultMode = config.HitlConfig{DefaultMode: mode}.EffectiveDefaultMode()
|
||||
h.config.Hitl.DefaultReviewer = config.HitlConfig{DefaultReviewer: reviewer}.EffectiveDefaultReviewer()
|
||||
if timeoutSeconds < 0 {
|
||||
timeoutSeconds = 0
|
||||
}
|
||||
h.config.Hitl.DefaultTimeoutSeconds = &timeoutSeconds
|
||||
if err := h.saveConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
h.logger.Info("HITL 全局默认配置已写入配置文件",
|
||||
zap.String("default_mode", h.config.Hitl.DefaultMode),
|
||||
zap.String("default_reviewer", h.config.Hitl.DefaultReviewer),
|
||||
zap.Int("default_timeout_seconds", timeoutSeconds),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateHitlDefaultReviewer 更新全局默认审批方并写入 config.yaml。
|
||||
func (h *ConfigHandler) UpdateHitlDefaultReviewer(reviewer string) error {
|
||||
h.mu.Lock()
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
package handler
|
||||
|
||||
import "sync"
|
||||
|
||||
// configFileMu serializes complete read-modify-write transactions across
|
||||
// handlers that share config.yaml. Per-handler locks cannot prevent lost
|
||||
// updates when another settings page saves a different YAML section.
|
||||
var configFileMu sync.Mutex
|
||||
@@ -160,24 +160,15 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
|
||||
limit = 1000
|
||||
}
|
||||
|
||||
excludeGrouped := strings.TrimSpace(search) == "" && projectID == "" &&
|
||||
(c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1")
|
||||
sortBy := strings.TrimSpace(c.Query("sort_by"))
|
||||
session, _ := security.CurrentSession(c)
|
||||
|
||||
var conversations []*database.Conversation
|
||||
var total int
|
||||
var err error
|
||||
if excludeGrouped {
|
||||
conversations, err = h.db.ListUngroupedConversationsForAccess(limit, offset, sortBy, projectID, session.UserID, session.Scope)
|
||||
if err == nil {
|
||||
total, err = h.db.CountUngroupedConversationsForAccess(projectID, session.UserID, session.Scope)
|
||||
}
|
||||
} else {
|
||||
conversations, err = h.db.ListConversationsForAccess(limit, offset, search, sortBy, projectID, session.UserID, session.Scope)
|
||||
if err == nil {
|
||||
total, err = h.db.CountConversationsForAccess(search, projectID, session.UserID, session.Scope)
|
||||
}
|
||||
conversations, err = h.db.ListConversationsForAccess(limit, offset, search, sortBy, projectID, session.UserID, session.Scope)
|
||||
if err == nil {
|
||||
total, err = h.db.CountConversationsForAccess(search, projectID, session.UserID, session.Scope)
|
||||
}
|
||||
if err != nil {
|
||||
h.logger.Error("获取对话列表失败", zap.Error(err))
|
||||
@@ -195,6 +186,35 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedRequest 更新对话置顶状态请求
|
||||
type UpdateConversationPinnedRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (h *ConversationHandler) UpdateConversationPinned(c *gin.Context) {
|
||||
conversationID := c.Param("id")
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateConversationPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationPinned(conversationID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新对话置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
// GetConversation 获取对话
|
||||
func (h *ConversationHandler) GetConversation(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
|
||||
@@ -76,6 +76,65 @@ func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testi
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsPageUsesPersistedExecutionStatusAfterBackgroundCancel(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-cancelled.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
conversation, err := db.CreateConversation("cancelled background", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
execID := "exec-cancelled-after-background"
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "exec", "toolCallId": "call-cancelled", "index": 1, "total": 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_result", "background", map[string]interface{}{
|
||||
"toolName": "exec", "toolCallId": "call-cancelled", "executionId": execID, "status": "background_running", "success": true,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||
ID: execID,
|
||||
ToolName: "exec",
|
||||
Status: mcp.ToolExecutionStatusCancelled,
|
||||
StartTime: now,
|
||||
EndTime: &now,
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveToolExecution: %v", err)
|
||||
}
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?limit=10&offset=0", nil)
|
||||
c.Params = gin.Params{{Key: "id", Value: message.ID}}
|
||||
NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("status = %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var response struct {
|
||||
ToolExecutions []database.ProcessDetailsToolExecution `json:"toolExecutions"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if len(response.ToolExecutions) != 1 {
|
||||
t.Fatalf("tool executions = %d, want 1", len(response.ToolExecutions))
|
||||
}
|
||||
if got := response.ToolExecutions[0].Status; got != mcp.ToolExecutionStatusCancelled {
|
||||
t.Fatalf("tool execution status = %q, want cancelled", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsFullBackfillsEmptyToolCallArgumentsFromExecution(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-args.db"), zap.NewNop())
|
||||
|
||||
@@ -192,6 +192,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
var emptyResponseContinueAttempt int
|
||||
var finalizationAutoContinueAttempt int
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
|
||||
for {
|
||||
segmentMainIterationMax := 0
|
||||
@@ -268,6 +269,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
@@ -366,15 +371,17 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
h.logger.Error("Eino ADK 单代理执行失败", zap.Error(runErr))
|
||||
taskStatus = "failed"
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
clientErr := multiagent.EinoClientRunErrorMessage(runErr)
|
||||
errMsg := "执行失败: " + clientErr
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
sendEvent("error", errMsg, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
errData := multiagent.EinoClientRunErrorFields(runErr)
|
||||
errData["conversationId"] = conversationID
|
||||
errData["messageId"] = assistantMessageID
|
||||
errData["error"] = errMsg
|
||||
sendEvent("error", errMsg, errData)
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
@@ -384,6 +391,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
|
||||
if decision.CompletionReason == "" {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, "eino_single", cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||
|
||||
@@ -401,10 +412,11 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
}
|
||||
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
}
|
||||
@@ -464,6 +476,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
var emptyResponseContinueAttempt int
|
||||
var finalizationAutoContinueAttempt int
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
for {
|
||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||
taskCtx,
|
||||
@@ -493,6 +506,10 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, prep.ConversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||
continue
|
||||
}
|
||||
@@ -509,18 +526,19 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"response": responseText,
|
||||
"conversationId": prep.ConversationID,
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"assistantMessageId": prep.AssistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"finalized": decision.Finalized,
|
||||
"finalizable": decision.Finalizable,
|
||||
"status": decision.Status,
|
||||
"completionReason": decision.CompletionReason,
|
||||
"evidenceVerified": decision.EvidenceVerified,
|
||||
"evidenceRefs": decision.EvidenceRefs,
|
||||
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||
"missingChecks": decision.MissingChecks,
|
||||
"response": responseText,
|
||||
"conversationId": prep.ConversationID,
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"assistantMessageId": prep.AssistantMessageID,
|
||||
"agentMode": "eino_single",
|
||||
"finalized": decision.Finalized,
|
||||
"finalizable": decision.Finalizable,
|
||||
"status": decision.Status,
|
||||
"completionReason": decision.CompletionReason,
|
||||
"evidenceVerified": decision.EvidenceVerified,
|
||||
"evidenceRefs": decision.EvidenceRefs,
|
||||
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||
"missingChecks": decision.MissingChecks,
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -379,6 +379,8 @@ func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool
|
||||
|
||||
// saveConfig 保存配置到文件
|
||||
func (h *ExternalMCPHandler) saveConfig() error {
|
||||
configFileMu.Lock()
|
||||
defer configFileMu.Unlock()
|
||||
data, err := os.ReadFile(h.configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取配置文件失败: %w", err)
|
||||
|
||||
@@ -2,16 +2,21 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/agentfinalizer"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const finalizationAutoContinueMaxAttempts = 2
|
||||
const finalizationPendingToolCancelWait = 2 * time.Second
|
||||
const finalizationPendingToolCancelPoll = 50 * time.Millisecond
|
||||
const finalizationPendingToolCancelNote = "Agent 迭代已结束,最终回复前自动终止未完成的工具执行"
|
||||
|
||||
func shouldAutoContinueAfterFinalization(d agentfinalizer.Decision, attempt int) bool {
|
||||
if d.Finalizable || d.Finalized {
|
||||
@@ -75,3 +80,105 @@ func finalizationAutoContinueBackoff(attempt int) time.Duration {
|
||||
}
|
||||
return time.Duration(attempt) * time.Second
|
||||
}
|
||||
|
||||
func (h *AgentHandler) cleanupPendingToolExecutionsAfterIteration(
|
||||
taskCtx context.Context,
|
||||
conversationID string,
|
||||
decision agentfinalizer.Decision,
|
||||
progressCallback func(eventType, message string, data interface{}),
|
||||
) []string {
|
||||
if h == nil || h.agent == nil || decision.CompletionReason != agentfinalizer.ReasonPendingTools {
|
||||
return nil
|
||||
}
|
||||
pending := uniqueNonEmptyStrings(decision.PendingExecutionIDs)
|
||||
if len(pending) == 0 {
|
||||
return nil
|
||||
}
|
||||
cancelled := make([]string, 0, len(pending))
|
||||
for _, executionID := range pending {
|
||||
if h.agent.CancelMCPToolExecutionWithNote(executionID, finalizationPendingToolCancelNote) {
|
||||
cancelled = append(cancelled, executionID)
|
||||
} else if h.logger != nil {
|
||||
h.logger.Warn("finalization pending tool cleanup could not cancel execution",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("executionId", executionID))
|
||||
}
|
||||
}
|
||||
if len(cancelled) == 0 {
|
||||
return nil
|
||||
}
|
||||
if progressCallback != nil {
|
||||
progressCallback("finalization_pending_tools_cancelled", "迭代结束,已自动终止仍在运行的工具执行。", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "finalizer",
|
||||
"autoCancelledPendingExecutionIds": cancelled,
|
||||
"pendingExecutionIds": pending,
|
||||
"reason": agentfinalizer.ReasonPendingTools,
|
||||
})
|
||||
}
|
||||
h.waitForToolExecutionsToLeavePending(taskCtx, cancelled, finalizationPendingToolCancelWait)
|
||||
return cancelled
|
||||
}
|
||||
|
||||
func (h *AgentHandler) waitForToolExecutionsToLeavePending(ctx context.Context, executionIDs []string, wait time.Duration) {
|
||||
if h == nil || h.db == nil || len(executionIDs) == 0 || wait <= 0 {
|
||||
return
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
defer timer.Stop()
|
||||
ticker := time.NewTicker(finalizationPendingToolCancelPoll)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
if !h.hasPendingToolExecutions(executionIDs) {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-contextDone(ctx):
|
||||
return
|
||||
case <-timer.C:
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hasPendingToolExecutions(executionIDs []string) bool {
|
||||
if h == nil || h.db == nil {
|
||||
return false
|
||||
}
|
||||
for _, executionID := range uniqueNonEmptyStrings(executionIDs) {
|
||||
exec, err := h.db.GetToolExecution(executionID)
|
||||
if err != nil || exec == nil {
|
||||
continue
|
||||
}
|
||||
switch strings.TrimSpace(exec.Status) {
|
||||
case mcp.ToolExecutionStatusQueued, mcp.ToolExecutionStatusRunning:
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func uniqueNonEmptyStrings(values []string) []string {
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
out := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[value]; ok {
|
||||
continue
|
||||
}
|
||||
seen[value] = struct{}{}
|
||||
out = append(out, value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func contextDone(ctx context.Context) <-chan struct{} {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
return ctx.Done()
|
||||
}
|
||||
|
||||
@@ -1,9 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
agentpkg "cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/agentfinalizer"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestShouldAutoContinueAfterFinalization(t *testing.T) {
|
||||
@@ -57,3 +66,66 @@ func TestRequestRequiresExecutionEvidenceUsesExplicitPolicyOnly(t *testing.T) {
|
||||
t.Fatal("explicit false policy should not require execution evidence")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanupPendingToolExecutionsAfterIterationAllowsFinalization(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "cleanup-finalization.db"), logger)
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
server := mcp.NewServerWithStorage(logger, db)
|
||||
server.ConfigureToolWaitTimeoutSeconds(1)
|
||||
server.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()
|
||||
})
|
||||
ag := agentpkg.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 10)
|
||||
h := &AgentHandler{agent: ag, db: db, logger: logger}
|
||||
|
||||
callCtx := mcp.WithMCPConversationID(context.Background(), "conv-cleanup")
|
||||
result, execID, err := server.CallTool(callCtx, "block", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool: %v", err)
|
||||
}
|
||||
if result == nil || !result.IsError || execID == "" {
|
||||
t.Fatalf("expected background wait result, result=%#v execID=%q", result, execID)
|
||||
}
|
||||
|
||||
decision := agentfinalizer.Decide(db, agentfinalizer.Input{
|
||||
Response: "基于已完成信息的阶段性总结。",
|
||||
MCPExecutionIDs: []string{execID},
|
||||
})
|
||||
if decision.CompletionReason != agentfinalizer.ReasonPendingTools {
|
||||
t.Fatalf("decision reason = %s, want pending tools: %+v", decision.CompletionReason, decision)
|
||||
}
|
||||
|
||||
var eventType string
|
||||
cancelled := h.cleanupPendingToolExecutionsAfterIteration(context.Background(), "conv-cleanup", decision, func(et, _ string, _ interface{}) {
|
||||
eventType = et
|
||||
})
|
||||
if len(cancelled) != 1 || cancelled[0] != execID {
|
||||
t.Fatalf("cancelled = %#v, want [%s]", cancelled, execID)
|
||||
}
|
||||
if eventType != "finalization_pending_tools_cancelled" {
|
||||
t.Fatalf("event type = %q", eventType)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
exec, err := db.GetToolExecution(execID)
|
||||
if err == nil && exec != nil && exec.Status == mcp.ToolExecutionStatusCancelled {
|
||||
after := agentfinalizer.Decide(db, agentfinalizer.Input{
|
||||
Response: "基于已完成信息的阶段性总结。",
|
||||
MCPExecutionIDs: []string{execID},
|
||||
})
|
||||
if !after.Finalizable || !after.Finalized {
|
||||
t.Fatalf("decision should finalize after cleanup: %+v", after)
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("execution did not become cancelled")
|
||||
}
|
||||
|
||||
@@ -1,438 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GroupHandler 分组处理器
|
||||
type GroupHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
const (
|
||||
maxGroupNameRunes = 64
|
||||
maxGroupIconRunes = 16
|
||||
)
|
||||
|
||||
// NewGroupHandler 创建新的分组处理器
|
||||
func NewGroupHandler(db *database.DB, logger *zap.Logger) *GroupHandler {
|
||||
return &GroupHandler{
|
||||
db: db,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
func validateGroupTextField(field, value string, maxRunes int, required bool) (string, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
if required {
|
||||
return "", errors.New(field + "不能为空")
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
if utf8.RuneCountInString(value) > maxRunes {
|
||||
return "", errors.New(field + "过长")
|
||||
}
|
||||
for _, r := range value {
|
||||
switch r {
|
||||
case '<', '>', '"', '\'', '`':
|
||||
return "", errors.New(field + "包含非法字符")
|
||||
}
|
||||
if r < 0x20 || r == 0x7f {
|
||||
return "", errors.New(field + "包含非法控制字符")
|
||||
}
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func validateGroupFields(name, icon string) (string, string, error) {
|
||||
validName, err := validateGroupTextField("分组名称", name, maxGroupNameRunes, true)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
validIcon, err := validateGroupTextField("分组图标", icon, maxGroupIconRunes, false)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return validName, validIcon, nil
|
||||
}
|
||||
|
||||
// CreateGroupRequest 创建分组请求
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
|
||||
// CreateGroup 创建分组
|
||||
func (h *GroupHandler) CreateGroup(c *gin.Context) {
|
||||
var req CreateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
name, icon, err := validateGroupFields(req.Name, req.Icon)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
session, _ := security.CurrentSession(c)
|
||||
group, err := h.db.CreateGroup(name, icon, session.UserID)
|
||||
if err != nil {
|
||||
h.logger.Error("创建分组失败", zap.Error(err))
|
||||
// 如果是名称重复错误,返回400状态码
|
||||
if err.Error() == "分组名称已存在" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// ListGroups 列出所有分组
|
||||
func (h *GroupHandler) ListGroups(c *gin.Context) {
|
||||
session, _ := security.CurrentSession(c)
|
||||
groups, err := h.db.ListGroupsForAccess(session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, groups)
|
||||
}
|
||||
|
||||
// GetGroup 获取分组
|
||||
func (h *GroupHandler) GetGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.GetGroup(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "分组不存在"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// UpdateGroupRequest 更新分组请求
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
}
|
||||
|
||||
// UpdateGroup 更新分组
|
||||
func (h *GroupHandler) UpdateGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
name, icon, err := validateGroupFields(req.Name, req.Icon)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateGroup(id, name, icon); err != nil {
|
||||
h.logger.Error("更新分组失败", zap.Error(err))
|
||||
// 如果是名称重复错误,返回400状态码
|
||||
if err.Error() == "分组名称已存在" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.GetGroup(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取更新后的分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, group)
|
||||
}
|
||||
|
||||
// DeleteGroup 删除分组
|
||||
func (h *GroupHandler) DeleteGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.DeleteGroup(id); err != nil {
|
||||
h.logger.Error("删除分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// AddConversationToGroupRequest 添加对话到分组请求
|
||||
type AddConversationToGroupRequest struct {
|
||||
ConversationID string `json:"conversationId"`
|
||||
GroupID string `json:"groupId"`
|
||||
}
|
||||
|
||||
// AddConversationToGroup 将对话添加到分组
|
||||
func (h *GroupHandler) AddConversationToGroup(c *gin.Context) {
|
||||
var req AddConversationToGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, req.ConversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if !h.groupAllowed(c, req.GroupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil {
|
||||
h.logger.Error("添加对话到分组失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "添加成功"})
|
||||
}
|
||||
|
||||
// RemoveConversationFromGroup 从分组中移除对话
|
||||
func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) {
|
||||
conversationID := c.Param("conversationId")
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.RemoveConversationFromGroup(conversationID, groupID); err != nil {
|
||||
h.logger.Error("从分组中移除对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "移除成功"})
|
||||
}
|
||||
|
||||
// GroupConversation 分组对话响应结构
|
||||
type GroupConversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
Pinned bool `json:"pinned"`
|
||||
GroupPinned bool `json:"groupPinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
// GetGroupConversations 获取分组中的所有对话
|
||||
func (h *GroupHandler) GetGroupConversations(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
searchQuery := c.Query("search") // 获取搜索参数
|
||||
|
||||
var conversations []*database.Conversation
|
||||
var err error
|
||||
|
||||
// 如果有搜索关键词,使用搜索方法;否则使用普通方法
|
||||
if searchQuery != "" {
|
||||
conversations, err = h.db.SearchConversationsByGroup(groupID, searchQuery)
|
||||
} else {
|
||||
conversations, err = h.db.GetConversationsByGroup(groupID)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 获取每个对话在分组中的置顶状态
|
||||
groupConvs := make([]GroupConversation, 0, len(conversations))
|
||||
for _, conv := range conversations {
|
||||
if conv == nil || !h.groupConversationAllowed(c, conv.ID) {
|
||||
continue
|
||||
}
|
||||
// 查询分组内置顶状态
|
||||
var groupPinned int
|
||||
err := h.db.QueryRow(
|
||||
"SELECT COALESCE(pinned, 0) FROM conversation_group_mappings WHERE conversation_id = ? AND group_id = ?",
|
||||
conv.ID, groupID,
|
||||
).Scan(&groupPinned)
|
||||
if err != nil {
|
||||
h.logger.Warn("查询分组内置顶状态失败", zap.String("conversationId", conv.ID), zap.Error(err))
|
||||
groupPinned = 0
|
||||
}
|
||||
|
||||
groupConvs = append(groupConvs, GroupConversation{
|
||||
ID: conv.ID,
|
||||
Title: conv.Title,
|
||||
Pinned: conv.Pinned,
|
||||
GroupPinned: groupPinned != 0,
|
||||
CreatedAt: conv.CreatedAt,
|
||||
UpdatedAt: conv.UpdatedAt,
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, groupConvs)
|
||||
}
|
||||
|
||||
// GetAllMappings 批量获取所有分组映射(消除前端 N+1 请求)
|
||||
func (h *GroupHandler) GetAllMappings(c *gin.Context) {
|
||||
mappings, err := h.db.GetAllGroupMappings()
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组映射失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
filtered := mappings[:0]
|
||||
for _, mapping := range mappings {
|
||||
if h.groupConversationAllowed(c, mapping.ConversationID) && h.groupAllowed(c, mapping.GroupID) {
|
||||
filtered = append(filtered, mapping)
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, filtered)
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedRequest 更新对话置顶状态请求
|
||||
type UpdateConversationPinnedRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (h *GroupHandler) UpdateConversationPinned(c *gin.Context) {
|
||||
conversationID := c.Param("id")
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateConversationPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationPinned(conversationID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新对话置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
// UpdateGroupPinnedRequest 更新分组置顶状态请求
|
||||
type UpdateGroupPinnedRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateGroupPinned 更新分组置顶状态
|
||||
func (h *GroupHandler) UpdateGroupPinned(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateGroupPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateGroupPinned(groupID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新分组置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedInGroupRequest 更新分组对话置顶状态请求
|
||||
type UpdateConversationPinnedInGroupRequest struct {
|
||||
Pinned bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedInGroup 更新对话在分组中的置顶状态
|
||||
func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
conversationID := c.Param("conversationId")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateConversationPinnedInGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.UpdateConversationPinnedInGroup(conversationID, groupID, req.Pinned); err != nil {
|
||||
h.logger.Error("更新分组对话置顶状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
func (h *GroupHandler) groupConversationAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
func (h *GroupHandler) groupAllowed(c *gin.Context, groupID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
return ok && h.db.UserCanAccessGroup(session.UserID, session.Scope, groupID)
|
||||
}
|
||||
@@ -1,41 +0,0 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestValidateGroupFieldsAllowsNormalNamesAndIcons(t *testing.T) {
|
||||
name, icon, err := validateGroupFields(" 日常安全巡检 ", " 📁 ")
|
||||
if err != nil {
|
||||
t.Fatalf("validateGroupFields returned error: %v", err)
|
||||
}
|
||||
if name != "日常安全巡检" {
|
||||
t.Fatalf("name = %q, want trimmed normal name", name)
|
||||
}
|
||||
if icon != "📁" {
|
||||
t.Fatalf("icon = %q, want trimmed icon", icon)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateGroupFieldsRejectsStoredXSSPayloads(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
icon string
|
||||
}{
|
||||
{name: `<img src=x onerror="alert(1)">`, icon: "📁"},
|
||||
{name: "日常安全巡检", icon: `<svg onload=alert(1)>`},
|
||||
{name: "日常安全巡检`onmouseover=alert(1)", icon: "📁"},
|
||||
{name: "日常安全巡检\x00", icon: "📁"},
|
||||
{name: strings.Repeat("分", maxGroupNameRunes+1), icon: "📁"},
|
||||
{name: "日常安全巡检", icon: strings.Repeat("📁", maxGroupIconRunes+1)},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name+"/"+tt.icon, func(t *testing.T) {
|
||||
if _, _, err := validateGroupFields(tt.name, tt.icon); err == nil {
|
||||
t.Fatal("validateGroupFields returned nil error for unsafe input")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -289,6 +289,18 @@ func normalizeHitlMode(mode string) string {
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeHitlDefaultMode(mode string) string {
|
||||
v := strings.ToLower(strings.TrimSpace(mode))
|
||||
switch v {
|
||||
case "feedback", "followup":
|
||||
return "approval"
|
||||
case "approval", "review_edit":
|
||||
return v
|
||||
default:
|
||||
return "off"
|
||||
}
|
||||
}
|
||||
|
||||
func (m *HITLManager) ActivateConversation(conversationID string, req *HITLRequest) {
|
||||
if req == nil || !req.Enabled {
|
||||
m.DeactivateConversation(conversationID)
|
||||
@@ -629,7 +641,7 @@ func (h *AgentHandler) loadHITLConversationConfig(conversationID string) (*HITLR
|
||||
return nil, err
|
||||
}
|
||||
if !has {
|
||||
cfg.Reviewer = h.hitlEffectiveDefaultReviewer()
|
||||
return h.hitlEffectiveDefaultRequest(), nil
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
@@ -994,7 +1006,9 @@ func (h *AgentHandler) GetHITLConversationConfig(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"conversationId": conversationID,
|
||||
"hitl": cfg,
|
||||
"defaultMode": h.hitlEffectiveDefaultMode(),
|
||||
"defaultReviewer": h.hitlEffectiveDefaultReviewer(),
|
||||
"defaultTimeoutSeconds": h.hitlEffectiveDefaultTimeoutSeconds(),
|
||||
"hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(),
|
||||
})
|
||||
}
|
||||
@@ -1051,11 +1065,64 @@ type setHitlDefaultReviewerReq struct {
|
||||
Reviewer string `json:"reviewer"`
|
||||
}
|
||||
|
||||
type setHitlDefaultConfigReq struct {
|
||||
Mode string `json:"mode"`
|
||||
Reviewer string `json:"reviewer"`
|
||||
TimeoutSeconds int `json:"timeoutSeconds"`
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlDefaultConfigResponse() gin.H {
|
||||
return gin.H{
|
||||
"defaultMode": h.hitlEffectiveDefaultMode(),
|
||||
"defaultReviewer": h.hitlEffectiveDefaultReviewer(),
|
||||
"defaultTimeoutSeconds": h.hitlEffectiveDefaultTimeoutSeconds(),
|
||||
"hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(),
|
||||
}
|
||||
}
|
||||
|
||||
// GetHITLDefaultConfig 返回 config.yaml 中的全局默认人机协同配置。
|
||||
func (h *AgentHandler) GetHITLDefaultConfig(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, h.hitlDefaultConfigResponse())
|
||||
}
|
||||
|
||||
// UpdateHITLDefaultConfig 将全局默认人机协同配置写入 config.yaml。
|
||||
func (h *AgentHandler) UpdateHITLDefaultConfig(c *gin.Context) {
|
||||
if h.hitlDefaultReviewerSaver == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 配置持久化不可用"})
|
||||
return
|
||||
}
|
||||
var req setHitlDefaultConfigReq
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
mode := normalizeHitlDefaultMode(req.Mode)
|
||||
reviewer := normalizeHitlReviewer(req.Reviewer)
|
||||
timeoutSeconds := req.TimeoutSeconds
|
||||
if timeoutSeconds < 0 {
|
||||
timeoutSeconds = 0
|
||||
}
|
||||
if err := h.hitlDefaultReviewerSaver.UpdateHitlDefaultConfig(mode, reviewer, timeoutSeconds); err != nil {
|
||||
h.logger.Warn("写入 HITL 默认配置到 config.yaml 失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.config != nil {
|
||||
h.config.Hitl.DefaultMode = mode
|
||||
h.config.Hitl.DefaultReviewer = reviewer
|
||||
h.config.Hitl.DefaultTimeoutSeconds = &timeoutSeconds
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "hitl", "default_config_update", "HITL 全局默认配置更新", "hitl_config", "default", nil)
|
||||
}
|
||||
out := h.hitlDefaultConfigResponse()
|
||||
out["ok"] = true
|
||||
c.JSON(http.StatusOK, out)
|
||||
}
|
||||
|
||||
// GetHITLDefaultReviewer 返回 config.yaml 中的全局默认审批方。
|
||||
func (h *AgentHandler) GetHITLDefaultReviewer(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"defaultReviewer": h.hitlEffectiveDefaultReviewer(),
|
||||
})
|
||||
c.JSON(http.StatusOK, h.hitlDefaultConfigResponse())
|
||||
}
|
||||
|
||||
// UpdateHITLDefaultReviewer 将全局默认审批方写入 config.yaml(未选会话时切换审批方)。
|
||||
@@ -1081,10 +1148,9 @@ func (h *AgentHandler) UpdateHITLDefaultReviewer(c *gin.Context) {
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "hitl", "default_reviewer_update", "HITL 全局默认审批方更新", "hitl_config", "default_reviewer", nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"defaultReviewer": reviewer,
|
||||
})
|
||||
out := h.hitlDefaultConfigResponse()
|
||||
out["ok"] = true
|
||||
c.JSON(http.StatusOK, out)
|
||||
}
|
||||
|
||||
// SetHITLGlobalToolWhitelist 整表替换 config.yaml 中的全局免审批工具白名单。
|
||||
|
||||
@@ -76,6 +76,7 @@ type MonitorStatsSummary struct {
|
||||
TotalCalls int `json:"totalCalls"`
|
||||
SuccessCalls int `json:"successCalls"`
|
||||
FailedCalls int `json:"failedCalls"`
|
||||
BlockedCalls int `json:"blockedCalls"`
|
||||
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
|
||||
ToolCount int `json:"toolCount"`
|
||||
}
|
||||
@@ -171,6 +172,8 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
||||
stat.FailedCalls++
|
||||
} else if exec.Status == "completed" {
|
||||
stat.SuccessCalls++
|
||||
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
|
||||
stat.BlockedCalls++
|
||||
}
|
||||
started := exec.StartTime
|
||||
if stat.LastCallTime == nil || started.After(*stat.LastCallTime) {
|
||||
@@ -448,6 +451,7 @@ func dbStatsSummaryToMonitor(result *database.ToolStatsSummaryResult) *MonitorSt
|
||||
TotalCalls: result.Summary.TotalCalls,
|
||||
SuccessCalls: result.Summary.SuccessCalls,
|
||||
FailedCalls: result.Summary.FailedCalls,
|
||||
BlockedCalls: result.Summary.BlockedCalls,
|
||||
ToolCount: result.Summary.ToolCount,
|
||||
}
|
||||
if result.Summary.LastCallTime != nil {
|
||||
@@ -472,6 +476,7 @@ func summarizeToolStats(stats map[string]*mcp.ToolStats, topN int) (*MonitorStat
|
||||
summary.TotalCalls += stat.TotalCalls
|
||||
summary.SuccessCalls += stat.SuccessCalls
|
||||
summary.FailedCalls += stat.FailedCalls
|
||||
summary.BlockedCalls += stat.BlockedCalls
|
||||
if stat.LastCallTime != nil && (summary.LastCallTime == nil || stat.LastCallTime.After(*summary.LastCallTime)) {
|
||||
t := *stat.LastCallTime
|
||||
summary.LastCallTime = &t
|
||||
@@ -528,6 +533,7 @@ func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats {
|
||||
existing.TotalCalls += v.TotalCalls
|
||||
existing.SuccessCalls += v.SuccessCalls
|
||||
existing.FailedCalls += v.FailedCalls
|
||||
existing.BlockedCalls += v.BlockedCalls
|
||||
// 使用最新的调用时间
|
||||
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
|
||||
existing.LastCallTime = v.LastCallTime
|
||||
@@ -734,9 +740,10 @@ func (h *MonitorHandler) GetStats(c *gin.Context) {
|
||||
|
||||
// CallsTimelinePoint 调用趋势数据点
|
||||
type CallsTimelinePoint struct {
|
||||
T time.Time `json:"t"`
|
||||
Total int `json:"total"`
|
||||
Failed int `json:"failed"`
|
||||
T time.Time `json:"t"`
|
||||
Total int `json:"total"`
|
||||
Failed int `json:"failed"`
|
||||
Blocked int `json:"blocked"`
|
||||
}
|
||||
|
||||
// CallsTimelineSummary 调用趋势汇总
|
||||
@@ -778,7 +785,7 @@ func truncateToBucket(t time.Time, bucketSize time.Duration, dailyBuckets bool)
|
||||
return t.Truncate(bucketSize)
|
||||
}
|
||||
|
||||
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed int }) []CallsTimelinePoint {
|
||||
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed, blocked int }) []CallsTimelinePoint {
|
||||
now := time.Now()
|
||||
start := truncateToBucket(now.Add(-cfg.duration), cfg.bucketSize, cfg.dailyBuckets)
|
||||
end := truncateToBucket(now, cfg.bucketSize, cfg.dailyBuckets)
|
||||
@@ -787,9 +794,10 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
|
||||
for current := start; !current.After(end); current = current.Add(cfg.bucketSize) {
|
||||
val := buckets[current]
|
||||
points = append(points, CallsTimelinePoint{
|
||||
T: current,
|
||||
Total: val.total,
|
||||
Failed: val.failed,
|
||||
T: current,
|
||||
Total: val.total,
|
||||
Failed: val.failed,
|
||||
Blocked: val.blocked,
|
||||
})
|
||||
}
|
||||
return points
|
||||
@@ -797,7 +805,7 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
|
||||
|
||||
func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimelinePoint {
|
||||
since := time.Now().Add(-cfg.duration)
|
||||
bucketMap := make(map[time.Time]struct{ total, failed int })
|
||||
bucketMap := make(map[time.Time]struct{ total, failed, blocked int })
|
||||
|
||||
if h.db != nil {
|
||||
dbBuckets, err := h.db.LoadCallsTimeline(since, cfg.dailyBuckets)
|
||||
@@ -809,6 +817,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
|
||||
entry := bucketMap[key]
|
||||
entry.total += b.Total
|
||||
entry.failed += b.Failed
|
||||
entry.blocked += b.Blocked
|
||||
bucketMap[key] = entry
|
||||
}
|
||||
return buildCallsTimelinePoints(cfg, bucketMap)
|
||||
@@ -824,6 +833,8 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
|
||||
entry.total++
|
||||
if monitorStatusCountsAsFailed(exec.Status) {
|
||||
entry.failed++
|
||||
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
|
||||
entry.blocked++
|
||||
}
|
||||
bucketMap[key] = entry
|
||||
}
|
||||
|
||||
@@ -205,6 +205,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
agentMode := "eino_" + effectiveOrch
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
|
||||
for {
|
||||
segmentMainIterationMax := 0
|
||||
@@ -282,6 +283,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
@@ -380,15 +385,17 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
taskStatus = "failed"
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
clientErr := multiagent.EinoClientRunErrorMessage(runErr)
|
||||
errMsg := "执行失败: " + clientErr
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
sendEvent("error", errMsg, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
errData := multiagent.EinoClientRunErrorFields(runErr)
|
||||
errData["conversationId"] = conversationID
|
||||
errData["messageId"] = assistantMessageID
|
||||
errData["error"] = errMsg
|
||||
sendEvent("error", errMsg, errData)
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
@@ -398,6 +405,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
if decision.CompletionReason == "" {
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, conversationID, decision, nil); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
}
|
||||
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||
|
||||
@@ -415,10 +426,11 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||
}
|
||||
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
"messageId": assistantMessageID,
|
||||
"agentMode": agentMode,
|
||||
"autoCancelledPendingExecutionIds": autoCancelledPendingExecutionIDs,
|
||||
}))
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
}
|
||||
@@ -478,6 +490,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
}
|
||||
agentMode := "eino_" + effectiveOrch
|
||||
var decision agentfinalizer.Decision
|
||||
var autoCancelledPendingExecutionIDs []string
|
||||
for {
|
||||
result, runErr = multiagent.RunDeepAgent(
|
||||
taskCtx,
|
||||
@@ -502,11 +515,14 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||
}
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
clientErr := multiagent.EinoClientRunErrorMessage(runErr)
|
||||
errMsg := "执行失败: " + clientErr
|
||||
if prep.AssistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), prep.AssistantMessageID)
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg})
|
||||
errData := multiagent.EinoClientRunErrorFields(runErr)
|
||||
errData["error"] = errMsg
|
||||
c.JSON(http.StatusInternalServerError, errData)
|
||||
return
|
||||
}
|
||||
mw := &h.config.MultiAgent.EinoMiddleware
|
||||
@@ -514,6 +530,10 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
if cancelled := h.cleanupPendingToolExecutionsAfterIteration(taskCtx, prep.ConversationID, decision, progressCallback); len(cancelled) > 0 {
|
||||
autoCancelledPendingExecutionIDs = mergeMCPExecutionIDLists(autoCancelledPendingExecutionIDs, cancelled)
|
||||
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||
}
|
||||
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||
continue
|
||||
}
|
||||
@@ -533,18 +553,19 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
responseText = finalizationBlockedMessage(decision)
|
||||
}
|
||||
c.JSON(http.StatusOK, ChatResponse{
|
||||
Response: responseText,
|
||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||
ConversationID: prep.ConversationID,
|
||||
Time: time.Now(),
|
||||
Finalizable: decision.Finalizable,
|
||||
Finalized: decision.Finalized,
|
||||
Status: decision.Status,
|
||||
CompletionReason: decision.CompletionReason,
|
||||
EvidenceVerified: decision.EvidenceVerified,
|
||||
EvidenceRefs: decision.EvidenceRefs,
|
||||
PendingExecutionIDs: decision.PendingExecutionIDs,
|
||||
MissingChecks: decision.MissingChecks,
|
||||
Response: responseText,
|
||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||
ConversationID: prep.ConversationID,
|
||||
Time: time.Now(),
|
||||
Finalizable: decision.Finalizable,
|
||||
Finalized: decision.Finalized,
|
||||
Status: decision.Status,
|
||||
CompletionReason: decision.CompletionReason,
|
||||
EvidenceVerified: decision.EvidenceVerified,
|
||||
EvidenceRefs: decision.EvidenceRefs,
|
||||
PendingExecutionIDs: decision.PendingExecutionIDs,
|
||||
MissingChecks: decision.MissingChecks,
|
||||
AutoCancelledPendingExecutionIDs: autoCancelledPendingExecutionIDs,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -456,75 +456,6 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"Group": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组ID",
|
||||
},
|
||||
"name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组名称",
|
||||
},
|
||||
"icon": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组图标",
|
||||
},
|
||||
"createdAt": map[string]interface{}{
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "创建时间",
|
||||
},
|
||||
"updatedAt": map[string]interface{}{
|
||||
"type": "string",
|
||||
"format": "date-time",
|
||||
"description": "更新时间",
|
||||
},
|
||||
},
|
||||
},
|
||||
"CreateGroupRequest": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组名称",
|
||||
},
|
||||
"icon": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组图标(可选)",
|
||||
},
|
||||
},
|
||||
},
|
||||
"UpdateGroupRequest": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组名称",
|
||||
},
|
||||
"icon": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组图标",
|
||||
},
|
||||
},
|
||||
},
|
||||
"AddConversationToGroupRequest": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"conversationId", "groupId"},
|
||||
"properties": map[string]interface{}{
|
||||
"conversationId": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "对话ID",
|
||||
},
|
||||
"groupId": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "分组ID",
|
||||
},
|
||||
},
|
||||
},
|
||||
"BatchTaskRequest": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"tasks"},
|
||||
@@ -1401,15 +1332,6 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "exclude_grouped",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "为 true 时排除已加入分组的对话(默认在未搜索且未按项目筛选时启用)",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "sort_by",
|
||||
"in": "query",
|
||||
@@ -2315,290 +2237,6 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups": map[string]interface{}{
|
||||
"post": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "创建分组",
|
||||
"description": "创建一个新的对话分组",
|
||||
"operationId": "createGroup",
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/CreateGroupRequest",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"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/Group",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"400": map[string]interface{}{
|
||||
"description": "请求参数错误或分组名称已存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "列出分组",
|
||||
"description": "获取所有对话分组",
|
||||
"operationId": "listGroups",
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "获取成功",
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "array",
|
||||
"items": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/Group",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/{id}": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "获取分组",
|
||||
"description": "获取指定分组的详细信息",
|
||||
"operationId": "getGroup",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"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/Group",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
"put": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "更新分组",
|
||||
"description": "更新分组信息",
|
||||
"operationId": "updateGroup",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/UpdateGroupRequest",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"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/Group",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"400": map[string]interface{}{
|
||||
"description": "请求参数错误或分组名称已存在",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
"delete": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "删除分组",
|
||||
"description": "删除指定分组",
|
||||
"operationId": "deleteGroup",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "删除成功",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/{id}/conversations": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "获取分组中的对话",
|
||||
"description": "获取指定分组中的所有对话",
|
||||
"operationId": "getGroupConversations",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "获取成功",
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "array",
|
||||
"items": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/Conversation",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/conversations": map[string]interface{}{
|
||||
"post": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "添加对话到分组",
|
||||
"description": "将对话添加到指定分组",
|
||||
"operationId": "addConversationToGroup",
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/AddConversationToGroupRequest",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "添加成功",
|
||||
},
|
||||
"400": map[string]interface{}{
|
||||
"description": "请求参数错误",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "对话或分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/{id}/conversations/{conversationId}": map[string]interface{}{
|
||||
"delete": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "从分组移除对话",
|
||||
"description": "从指定分组中移除对话",
|
||||
"operationId": "removeConversationFromGroup",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "conversationId",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "对话ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "移除成功",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "对话或分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/assets/import": map[string]interface{}{
|
||||
"post": map[string]interface{}{
|
||||
"tags": []string{"资产管理"},
|
||||
@@ -4266,109 +3904,6 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/{id}/pinned": map[string]interface{}{
|
||||
"put": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "设置分组置顶",
|
||||
"description": "设置或取消分组的置顶状态",
|
||||
"operationId": "updateGroupPinned",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"pinned"},
|
||||
"properties": map[string]interface{}{
|
||||
"pinned": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否置顶",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "更新成功",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/groups/{id}/conversations/{conversationId}/pinned": map[string]interface{}{
|
||||
"put": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "设置分组中对话的置顶",
|
||||
"description": "设置或取消分组中对话的置顶状态",
|
||||
"operationId": "updateConversationPinnedInGroup",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "分组ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "conversationId",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"description": "对话ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
},
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"required": []string{"pinned"},
|
||||
"properties": map[string]interface{}{
|
||||
"pinned": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否置顶",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "更新成功",
|
||||
},
|
||||
"404": map[string]interface{}{
|
||||
"description": "对话或分组不存在",
|
||||
},
|
||||
"401": map[string]interface{}{
|
||||
"description": "未授权",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/knowledge/categories": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"知识库"},
|
||||
@@ -5194,38 +4729,6 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
// ==================== 对话分组 - 缺失端点 ====================
|
||||
"/api/groups/mappings": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话分组"},
|
||||
"summary": "获取所有分组映射",
|
||||
"description": "获取所有对话与分组之间的映射关系列表。",
|
||||
"operationId": "getAllGroupMappings",
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
"description": "获取成功",
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "array",
|
||||
"items": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"conversation_id": map[string]interface{}{"type": "string", "description": "对话ID"},
|
||||
"group_id": map[string]interface{}{"type": "string", "description": "分组ID"},
|
||||
"pinned": map[string]interface{}{"type": "boolean", "description": "是否置顶"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"401": map[string]interface{}{"description": "未授权"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
// ==================== FOFA信息收集 ====================
|
||||
"/api/fofa/search": map[string]interface{}{
|
||||
"post": map[string]interface{}{
|
||||
|
||||
@@ -5,7 +5,7 @@ package handler
|
||||
|
||||
var apiDocI18nTagToKey = map[string]string{
|
||||
"认证": "auth", "对话管理": "conversationManagement", "对话交互": "conversationInteraction",
|
||||
"批量任务": "batchTasks", "对话分组": "conversationGroups", "漏洞管理": "vulnerabilityManagement",
|
||||
"批量任务": "batchTasks", "漏洞管理": "vulnerabilityManagement",
|
||||
"角色管理": "roleManagement", "Skills管理": "skillsManagement", "监控": "monitoring",
|
||||
"配置管理": "configManagement", "外部MCP管理": "externalMCPManagement", "攻击链": "attackChain",
|
||||
"知识库": "knowledgeBase", "MCP": "mcp",
|
||||
@@ -24,10 +24,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"删除批量任务队列": "deleteBatchQueue", "启动批量任务队列": "startBatchQueue", "暂停批量任务队列": "pauseBatchQueue",
|
||||
"添加任务到队列": "addTaskToQueue", "SQL注入扫描": "sqlInjectionScan", "端口扫描": "portScan",
|
||||
"更新批量任务": "updateBatchTask", "删除批量任务": "deleteBatchTask",
|
||||
"创建分组": "createGroup", "列出分组": "listGroups", "获取分组": "getGroup", "更新分组": "updateGroup",
|
||||
"删除分组": "deleteGroup", "获取分组中的对话": "getGroupConversations", "添加对话到分组": "addConversationToGroup",
|
||||
"从分组移除对话": "removeConversationFromGroup",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"获取漏洞": "getVulnerability", "更新漏洞": "updateVulnerability", "删除漏洞": "deleteVulnerability",
|
||||
"列出角色": "listRoles", "创建角色": "createRole", "获取角色": "getRole", "更新角色": "updateRole", "删除角色": "deleteRole",
|
||||
"获取可用Skills列表": "getAvailableSkills", "列出Skills": "listSkills", "创建Skill": "createSkill",
|
||||
@@ -40,8 +37,8 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"添加或更新外部MCP": "addOrUpdateExternalMCP", "stdio模式配置": "stdioModeConfig", "SSE模式配置": "sseModeConfig",
|
||||
"删除外部MCP": "deleteExternalMCP", "启动外部MCP": "startExternalMCP", "停止外部MCP": "stopExternalMCP",
|
||||
"获取攻击链": "getAttackChain", "重新生成攻击链": "regenerateAttackChain",
|
||||
"设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation",
|
||||
"获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem",
|
||||
"设置对话置顶": "pinConversation",
|
||||
"获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem",
|
||||
"获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem",
|
||||
"获取索引状态": "getIndexStatus", "构建索引": "startKnowledgeIndex", "扫描知识库": "scanKnowledgeBase",
|
||||
"搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType",
|
||||
@@ -52,8 +49,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"删除对话轮次": "deleteConversationTurn", "获取消息过程详情": "getMessageProcessDetails",
|
||||
"重跑批量任务队列": "rerunBatchQueue", "修改队列元数据": "updateBatchQueueMetadata",
|
||||
"修改队列调度配置": "updateBatchQueueSchedule", "开关Cron自动调度": "setBatchQueueScheduleEnabled",
|
||||
"获取所有分组映射": "getAllGroupMappings",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"测试OpenAI API连接": "testOpenAI",
|
||||
"执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS",
|
||||
"列出WebShell连接": "listWebshellConnections", "创建WebShell连接": "createWebshellConnection",
|
||||
@@ -84,7 +80,6 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"获取成功": "getSuccess", "未授权": "unauthorized", "未授权,需要有效的Token": "unauthorizedToken",
|
||||
"创建成功": "createSuccess", "请求参数错误": "badRequest", "对话不存在": "conversationNotFound",
|
||||
"对话不存在或结果不存在": "conversationOrResultNotFound", "请求参数错误(如task为空)": "badRequestTaskEmpty",
|
||||
"请求参数错误或分组名称已存在": "badRequestGroupNameExists", "分组不存在": "groupNotFound",
|
||||
"请求参数错误(如配置格式不正确、缺少必需字段等)": "badRequestConfig",
|
||||
"请求参数错误(如query为空)": "badRequestQueryEmpty", "方法不允许(仅支持POST请求)": "methodNotAllowed",
|
||||
"登录成功": "loginSuccess", "密码错误": "invalidPassword", "登出成功": "logoutSuccess",
|
||||
@@ -92,7 +87,7 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"对话创建成功": "conversationCreated", "服务器内部错误": "internalError", "更新成功": "updateSuccess",
|
||||
"删除成功": "deleteSuccess", "队列不存在": "queueNotFound", "启动成功": "startSuccess",
|
||||
"暂停成功": "pauseSuccess", "添加成功": "addSuccess",
|
||||
"任务不存在": "taskNotFound", "对话或分组不存在": "conversationOrGroupNotFound",
|
||||
"任务不存在": "taskNotFound",
|
||||
"取消请求已提交": "cancelSubmitted", "未找到正在执行的任务": "noRunningTask",
|
||||
"消息发送成功,返回AI回复": "messageSent", "流式响应(Server-Sent Events)": "streamResponse",
|
||||
// 新增缺失端点响应
|
||||
|
||||
@@ -3,6 +3,7 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -484,7 +485,10 @@ func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool,
|
||||
if runtimeCancel != nil {
|
||||
runtimeHandled = runtimeCancel(cause)
|
||||
}
|
||||
if cancel != nil && !runtimeHandled {
|
||||
// 「彻底停止」必须同时取消宿主 context:原生 Agent Cancel 即使已受理,
|
||||
// 也可能只在安全点返回或报告超时,不能据此让整条任务继续存活。
|
||||
// 中断并继续仍保留原语义:原生取消已处理时由运行时负责恢复。
|
||||
if cancel != nil && (!runtimeHandled || errors.Is(cause, ErrTaskCancelled)) {
|
||||
cancel(cause)
|
||||
}
|
||||
if toolCanceler != nil {
|
||||
@@ -591,6 +595,12 @@ func (m *AgentTaskManager) GetActiveTasks() []*AgentTask {
|
||||
Status: task.Status,
|
||||
})
|
||||
}
|
||||
sort.Slice(result, func(i, j int) bool {
|
||||
if result[i].StartedAt.Equal(result[j].StartedAt) {
|
||||
return result[i].ConversationID < result[j].ConversationID
|
||||
}
|
||||
return result[i].StartedAt.Before(result[j].StartedAt)
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestGetActiveTasksUsesStableCreationOrder(t *testing.T) {
|
||||
m := NewAgentTaskManager()
|
||||
started := time.Date(2026, 8, 19, 10, 0, 0, 0, time.UTC)
|
||||
m.mu.Lock()
|
||||
m.tasks = map[string]*AgentTask{
|
||||
"conversation-z": {ConversationID: "conversation-z", StartedAt: started, Status: "running"},
|
||||
"conversation-late": {ConversationID: "conversation-late", StartedAt: started.Add(time.Minute), Status: "running"},
|
||||
"conversation-a": {ConversationID: "conversation-a", StartedAt: started, Status: "running"},
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
want := []string{"conversation-a", "conversation-z", "conversation-late"}
|
||||
for attempt := 0; attempt < 20; attempt++ {
|
||||
gotTasks := m.GetActiveTasks()
|
||||
if len(gotTasks) != len(want) {
|
||||
t.Fatalf("GetActiveTasks() length = %d, want %d", len(gotTasks), len(want))
|
||||
}
|
||||
for i, task := range gotTasks {
|
||||
if task.ConversationID != want[i] {
|
||||
t.Fatalf("attempt %d order[%d] = %q, want %q", attempt, i, task.ConversationID, want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -32,7 +32,7 @@ func TestCancelTaskInvokesToolCancelerOnFullStop(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelTaskUsesAgentRuntimeCancelAsPrimaryPath(t *testing.T) {
|
||||
func TestCancelTaskFullStopCancelsRuntimeAndParentContext(t *testing.T) {
|
||||
tm := NewAgentTaskManager()
|
||||
var order []string
|
||||
tm.SetToolCanceler(func(conversationID string) {
|
||||
@@ -61,7 +61,7 @@ func TestCancelTaskUsesAgentRuntimeCancelAsPrimaryPath(t *testing.T) {
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("CancelTask: ok=%v err=%v", ok, err)
|
||||
}
|
||||
want := []string{"runtime", "tool"}
|
||||
want := []string{"runtime", "context", "tool"}
|
||||
if len(order) != len(want) {
|
||||
t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order)
|
||||
}
|
||||
@@ -72,6 +72,29 @@ func TestCancelTaskUsesAgentRuntimeCancelAsPrimaryPath(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelTaskInterruptContinueKeepsParentWhenRuntimeHandlesIt(t *testing.T) {
|
||||
tm := NewAgentTaskManager()
|
||||
ctx, cancel := context.WithCancelCause(context.Background())
|
||||
if _, err := tm.StartTask("conv-interrupt-native", "hello", cancel); err != nil {
|
||||
t.Fatalf("StartTask: %v", err)
|
||||
}
|
||||
unregister := tm.BindAgentRuntimeCancel("conv-interrupt-native", func(err error) bool {
|
||||
if !errors.Is(err, multiagent.ErrInterruptContinue) {
|
||||
t.Fatalf("runtime cancel got %v", err)
|
||||
}
|
||||
return true
|
||||
})
|
||||
defer unregister()
|
||||
|
||||
ok, err := tm.CancelTask("conv-interrupt-native", multiagent.ErrInterruptContinue)
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("CancelTask: ok=%v err=%v", ok, err)
|
||||
}
|
||||
if cause := context.Cause(ctx); cause != nil {
|
||||
t.Fatalf("interrupt-continue parent context cause = %v, want nil", cause)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelTaskFallsBackToContextWhenAgentRuntimeCancelMisses(t *testing.T) {
|
||||
tm := NewAgentTaskManager()
|
||||
var order []string
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GetTokenUsageStats returns model token usage aggregates for dashboard views.
|
||||
func (h *ConversationHandler) GetTokenUsageStats(c *gin.Context) {
|
||||
filter := tokenUsageFilterFromQuery(c)
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
filter.Access = database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
stats, err := h.db.GetModelTokenUsageStats(filter)
|
||||
if err != nil {
|
||||
h.logger.Error("获取Token用量统计失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
// GetConversationTokenUsageStats returns token usage scoped to one conversation.
|
||||
func (h *ConversationHandler) GetConversationTokenUsageStats(c *gin.Context) {
|
||||
filter := tokenUsageFilterFromQuery(c)
|
||||
filter.ConversationID = strings.TrimSpace(c.Param("id"))
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
filter.Access = database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
stats, err := h.db.GetModelTokenUsageStats(filter)
|
||||
if err != nil {
|
||||
h.logger.Error("获取对话Token用量统计失败", zap.Error(err), zap.String("conversationId", filter.ConversationID))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, stats)
|
||||
}
|
||||
|
||||
func tokenUsageFilterFromQuery(c *gin.Context) database.ModelTokenUsageFilter {
|
||||
days, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("days", "7")))
|
||||
if days <= 0 {
|
||||
days = 7
|
||||
}
|
||||
if days > 365 {
|
||||
days = 365
|
||||
}
|
||||
limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("limit", "10")))
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
filter := database.ModelTokenUsageFilter{
|
||||
ConversationID: strings.TrimSpace(c.Query("conversation_id")),
|
||||
ProjectID: strings.TrimSpace(c.Query("project_id")),
|
||||
Days: days,
|
||||
Limit: limit,
|
||||
}
|
||||
if since := parseTokenUsageQueryTime(c.Query("since")); !since.IsZero() {
|
||||
filter.Since = since
|
||||
} else if days > 0 {
|
||||
now := time.Now()
|
||||
start := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()).AddDate(0, 0, -(days - 1))
|
||||
filter.Since = start
|
||||
}
|
||||
if until := parseTokenUsageQueryTime(c.Query("until")); !until.IsZero() {
|
||||
filter.Until = until
|
||||
}
|
||||
return filter
|
||||
}
|
||||
|
||||
func parseTokenUsageQueryTime(raw string) time.Time {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return time.Time{}
|
||||
}
|
||||
for _, layout := range []string{time.RFC3339Nano, time.RFC3339, "2006-01-02"} {
|
||||
if t, err := time.Parse(layout, raw); err == nil {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
func (h *ConfigHandler) SetToolGuard(manager *toolguard.Manager) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.toolGuard = manager
|
||||
}
|
||||
|
||||
func (h *ConfigHandler) GetToolGuard(c *gin.Context) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
if h.toolGuard == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, h.toolGuard.Config())
|
||||
}
|
||||
|
||||
// decodeToolGuardRequest bounds both config and dry-run inputs, rejects unknown
|
||||
// fields and trailing JSON, and never invokes an actual tool.
|
||||
func decodeToolGuardRequest(c *gin.Context, dst interface{}) error {
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1<<20)
|
||||
decoder := json.NewDecoder(c.Request.Body)
|
||||
decoder.DisallowUnknownFields()
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(dst); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := decoder.Decode(new(interface{})); err != io.EOF {
|
||||
return fmt.Errorf("请求必须只包含一个 JSON 对象")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *ConfigHandler) UpdateToolGuard(c *gin.Context) {
|
||||
var req struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
Rules *[]toolguard.Rule `json:"rules"`
|
||||
}
|
||||
if err := decodeToolGuardRequest(c, &req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的调用拦截配置: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if req.Enabled == nil || req.Rules == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "必须明确提供 enabled 和 rules;清空规则请提供空数组"})
|
||||
return
|
||||
}
|
||||
cfg := toolguard.Config{Enabled: *req.Enabled, Rules: *req.Rules}
|
||||
if _, err := toolguard.Compile(cfg); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.toolGuard == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
|
||||
return
|
||||
}
|
||||
// Commit the file first; a validation/write failure must leave the current
|
||||
// effective policy and in-memory config intact.
|
||||
if err := h.saveToolGuardConfig(cfg); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存调用拦截配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.toolGuard.Update(cfg); err != nil {
|
||||
// The same immutable input was compiled above, so this cannot fail
|
||||
// unless validation gains an additional runtime dependency.
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "应用调用拦截配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
h.config.ToolGuard = &cfg
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "config", "tool_guard_update", "更新调用拦截规则", "config", "tool_guard", map[string]interface{}{
|
||||
"enabled": cfg.Enabled, "rule_count": len(cfg.Rules),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, h.toolGuard.Config())
|
||||
}
|
||||
|
||||
func (h *ConfigHandler) TestToolGuard(c *gin.Context) {
|
||||
var req struct {
|
||||
Config *toolguard.Config `json:"config"`
|
||||
ToolName string `json:"toolName"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
}
|
||||
if err := decodeToolGuardRequest(c, &req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的试匹配参数: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if req.Config == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请提供待测试的 config"})
|
||||
return
|
||||
}
|
||||
policy, err := toolguard.Compile(*req.Config)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if match := policy.Check(req.ToolName, req.Arguments); match != nil {
|
||||
c.JSON(http.StatusOK, gin.H{"blocked": true, "match": match})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"blocked": false})
|
||||
}
|
||||
|
||||
// saveToolGuardConfig changes only this YAML section, preserving unrelated
|
||||
// settings/comments and file permissions. Rename makes the write atomic.
|
||||
// h.mu protects the runtime configuration; configFileMu also covers independent
|
||||
// writers such as ExternalMCPHandler.
|
||||
func (h *ConfigHandler) saveToolGuardConfig(cfg toolguard.Config) error {
|
||||
configFileMu.Lock()
|
||||
defer configFileMu.Unlock()
|
||||
path, err := filepath.EvalSymlinks(h.configPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
doc, err := loadYAMLDocument(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var node yaml.Node
|
||||
if err := node.Encode(cfg); err != nil {
|
||||
return err
|
||||
}
|
||||
_, value := ensureKeyValue(doc.Content[0], "tool_guard")
|
||||
*value = node
|
||||
var buf bytes.Buffer
|
||||
encoder := yaml.NewEncoder(&buf)
|
||||
encoder.SetIndent(2)
|
||||
if err := encoder.Encode(doc); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := encoder.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp, err := os.CreateTemp(filepath.Dir(path), ".tool-guard-*.yaml")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer os.Remove(tmp.Name())
|
||||
defer tmp.Close()
|
||||
if err := tmp.Chmod(info.Mode().Perm()); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tmp.Write(buf.Bytes()); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tmp.Sync(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tmp.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp.Name(), path)
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func newToolGuardTestHandler(t *testing.T) *ConfigHandler {
|
||||
t.Helper()
|
||||
path := filepath.Join(t.TempDir(), "config.yaml")
|
||||
if err := os.WriteFile(path, []byte("# keep this comment\nserver:\n port: 8123\nhitl:\n tool_whitelist: [read_file]\n"), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
manager, err := toolguard.NewManager(toolguard.DefaultConfig())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return &ConfigHandler{configPath: path, config: &config.Config{}, toolGuard: manager}
|
||||
}
|
||||
|
||||
func toolGuardRequest(t *testing.T, handler gin.HandlerFunc, body interface{}) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
data, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPut, "/api/tool-guard", bytes.NewReader(data))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
handler(c)
|
||||
return w
|
||||
}
|
||||
|
||||
func TestToolGuardSavePersistsAndAppliesWithoutChangingHITL(t *testing.T) {
|
||||
h := newToolGuardTestHandler(t)
|
||||
cfg := toolguard.DefaultConfig()
|
||||
cfg.Rules[0].Message = "识别到 {match},禁止攻击政府网站,请检查目标。"
|
||||
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("save: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
loaded, err := config.Load(h.configPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) || !reflect.DeepEqual(h.toolGuard.Config(), cfg) {
|
||||
t.Fatal("saved and effective policies differ")
|
||||
}
|
||||
if loaded.Server.Port != 8123 || !reflect.DeepEqual(loaded.Hitl.ToolWhitelist, []string{"read_file"}) {
|
||||
t.Fatal("unrelated configuration was changed")
|
||||
}
|
||||
info, _ := os.Stat(h.configPath)
|
||||
data, _ := os.ReadFile(h.configPath)
|
||||
if info.Mode().Perm() != 0600 || !strings.Contains(string(data), "# keep this comment") {
|
||||
t.Fatal("file permissions or comments were lost")
|
||||
}
|
||||
match := h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov.cn"})
|
||||
if match == nil || !strings.Contains(match.Message, "agency.gov.cn") {
|
||||
t.Fatalf("updated message not applied: %+v", match)
|
||||
}
|
||||
cfg.Enabled = false
|
||||
w = toolGuardRequest(t, h.UpdateToolGuard, cfg)
|
||||
if w.Code != http.StatusOK || h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov"}) != nil {
|
||||
t.Fatal("explicitly disabling protection did not apply")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolGuardInvalidAndFailedSaveKeepEffectivePolicy(t *testing.T) {
|
||||
h := newToolGuardTestHandler(t)
|
||||
before, _ := os.ReadFile(h.configPath)
|
||||
cfg := toolguard.DefaultConfig()
|
||||
cfg.Enabled = false
|
||||
cfg.Rules[0].Pattern = "["
|
||||
for _, body := range []interface{}{cfg, map[string]interface{}{}, nil, map[string]interface{}{"enabled": false, "rules": nil}} {
|
||||
w := toolGuardRequest(t, h.UpdateToolGuard, body)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid update accepted: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
after, _ := os.ReadFile(h.configPath)
|
||||
if !bytes.Equal(before, after) || !h.toolGuard.Config().Enabled {
|
||||
t.Fatal("invalid input changed protection")
|
||||
}
|
||||
h.configPath = filepath.Join(t.TempDir(), "missing", "config.yaml")
|
||||
cfg = toolguard.DefaultConfig()
|
||||
cfg.Enabled = false
|
||||
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
|
||||
if w.Code != http.StatusInternalServerError || !h.toolGuard.Config().Enabled || h.config.ToolGuard != nil {
|
||||
t.Fatal("failed persistence changed live configuration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolGuardDryRunUsesUnsavedPolicyWithoutMutation(t *testing.T) {
|
||||
h := newToolGuardTestHandler(t)
|
||||
cfg := toolguard.DefaultConfig()
|
||||
cfg.Rules[0].Pattern = "example\\.org"
|
||||
w := toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{
|
||||
"config": cfg, "toolName": "scan", "arguments": map[string]interface{}{"target": "example.org"},
|
||||
})
|
||||
var got struct {
|
||||
Blocked bool `json:"blocked"`
|
||||
Match *toolguard.Match `json:"match"`
|
||||
}
|
||||
if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &got) != nil || !got.Blocked || got.Match == nil || got.Match.MatchedText != "example.org" {
|
||||
t.Fatalf("dry run failed: %s", w.Body.String())
|
||||
}
|
||||
if !reflect.DeepEqual(h.toolGuard.Config(), toolguard.DefaultConfig()) || h.config.ToolGuard != nil {
|
||||
t.Fatal("dry run changed live configuration")
|
||||
}
|
||||
w = toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{"config": cfg, "arguments": []string{"example.org"}})
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatal("non-object tool arguments accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolGuardRoutesEnforceConfigurationPermissions(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tc := range []struct {
|
||||
method, path, permission, scope string
|
||||
want int
|
||||
}{
|
||||
{"GET", "/api/tool-guard", "hitl:read", database.RBACScopeAll, 403},
|
||||
{"PUT", "/api/tool-guard", "hitl:write", database.RBACScopeAll, 403},
|
||||
{"GET", "/api/tool-guard", "config:read", database.RBACScopeAll, 200},
|
||||
{"POST", "/api/tool-guard/test", "config:read", database.RBACScopeAll, 200},
|
||||
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeAll, 200},
|
||||
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeOwn, 403},
|
||||
} {
|
||||
t.Run(tc.method+tc.permission+tc.scope, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
r.Use(func(c *gin.Context) {
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: "test", Permissions: map[string]bool{tc.permission: true}, Scope: tc.scope})
|
||||
})
|
||||
r.Use(security.RBACMiddleware(&database.DB{}))
|
||||
r.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(200) })
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil))
|
||||
if w.Code != tc.want {
|
||||
t.Fatalf("got %d, want %d: %s", w.Code, tc.want, w.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolGuardConcurrentOtherSettingsSavePreservesPolicy(t *testing.T) {
|
||||
h := newToolGuardTestHandler(t)
|
||||
external := &ExternalMCPHandler{configPath: h.configPath, config: h.config, logger: zap.NewNop()}
|
||||
cfg := toolguard.DefaultConfig()
|
||||
cfg.Rules[0].Message = "持久化策略 {match}"
|
||||
var wg sync.WaitGroup
|
||||
errors := make(chan error, 2)
|
||||
for _, save := range []func() error{func() error { return h.saveToolGuardConfig(cfg) }, external.saveConfig} {
|
||||
wg.Add(1)
|
||||
go func(save func() error) {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 20; i++ {
|
||||
if err := save(); err != nil {
|
||||
errors <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}(save)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
for err := range errors {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded, err := config.Load(h.configPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) {
|
||||
t.Fatal("another settings save overwrote the tool guard policy")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DiagnosticOptions controls the additional warn-and-above diagnostic output.
|
||||
type DiagnosticOptions struct {
|
||||
Dir string
|
||||
Disabled bool
|
||||
RetentionDays int // Values <= 0 use the default of 14 calendar days.
|
||||
}
|
||||
|
||||
// dailyWriter opens lazily: healthy runs create no diagnostic files. Opening
|
||||
// per write also avoids keeping descriptors open across rotation or shutdown.
|
||||
type dailyWriter struct {
|
||||
mu sync.Mutex
|
||||
dir string
|
||||
retentionDays int
|
||||
cleanedDay string
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
func newDailyWriter(options DiagnosticOptions) *dailyWriter {
|
||||
if options.Dir == "" {
|
||||
options.Dir = "log"
|
||||
}
|
||||
if options.RetentionDays <= 0 {
|
||||
options.RetentionDays = 14
|
||||
}
|
||||
return &dailyWriter{dir: options.Dir, retentionDays: options.RetentionDays, now: time.Now}
|
||||
}
|
||||
|
||||
func (w *dailyWriter) Write(p []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
now := w.now()
|
||||
day := now.Format(time.DateOnly)
|
||||
if err := os.MkdirAll(w.dir, 0700); err != nil {
|
||||
return 0, fmt.Errorf("create diagnostic log directory: %w", err)
|
||||
}
|
||||
f, err := os.OpenFile(filepath.Join(w.dir, "diagnostic-"+day+".log"), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("open diagnostic log: %w", err)
|
||||
}
|
||||
n, writeErr := f.Write(p)
|
||||
closeErr := f.Close()
|
||||
var cleanupErr error
|
||||
if w.cleanedDay != day {
|
||||
cleanupErr = w.cleanup(now)
|
||||
if cleanupErr == nil {
|
||||
w.cleanedDay = day
|
||||
}
|
||||
}
|
||||
return n, errors.Join(writeErr, closeErr, cleanupErr)
|
||||
}
|
||||
|
||||
// Writes are unbuffered and files are closed before Write returns.
|
||||
func (w *dailyWriter) Sync() error { return nil }
|
||||
|
||||
func (w *dailyWriter) cleanup(now time.Time) error {
|
||||
entries, err := os.ReadDir(w.dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cutoff := now.AddDate(0, 0, -(w.retentionDays - 1)).Format(time.DateOnly)
|
||||
var errs []error
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if !entry.Type().IsRegular() || !strings.HasPrefix(name, "diagnostic-") || !strings.HasSuffix(name, ".log") {
|
||||
continue
|
||||
}
|
||||
day := strings.TrimSuffix(strings.TrimPrefix(name, "diagnostic-"), ".log")
|
||||
if _, err := time.Parse(time.DateOnly, day); err != nil || day >= cutoff {
|
||||
continue
|
||||
}
|
||||
if err := os.Remove(filepath.Join(w.dir, name)); err != nil && !os.IsNotExist(err) {
|
||||
errs = append(errs, fmt.Errorf("remove expired diagnostic log %s: %w", name, err))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
@@ -11,7 +11,7 @@ type Logger struct {
|
||||
*zap.Logger
|
||||
}
|
||||
|
||||
func New(level, output string) *Logger {
|
||||
func New(level, output string, diagnostics ...DiagnosticOptions) *Logger {
|
||||
var zapLevel zapcore.Level
|
||||
switch level {
|
||||
case "debug":
|
||||
@@ -34,6 +34,8 @@ func New(level, output string) *Logger {
|
||||
var writeSyncer zapcore.WriteSyncer
|
||||
if output == "stdout" {
|
||||
writeSyncer = zapcore.AddSync(os.Stdout)
|
||||
} else if output == "stderr" {
|
||||
writeSyncer = zapcore.AddSync(os.Stderr)
|
||||
} else {
|
||||
file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
|
||||
if err != nil {
|
||||
@@ -49,6 +51,18 @@ func New(level, output string) *Logger {
|
||||
zapLevel,
|
||||
)
|
||||
|
||||
options := DiagnosticOptions{}
|
||||
if len(diagnostics) > 0 {
|
||||
options = diagnostics[0]
|
||||
}
|
||||
if !options.Disabled {
|
||||
// The diagnostic threshold is independent of the primary output level.
|
||||
core = zapcore.NewTee(core, zapcore.NewCore(
|
||||
zapcore.NewJSONEncoder(config.EncoderConfig),
|
||||
newDailyWriter(options), zapcore.WarnLevel,
|
||||
))
|
||||
}
|
||||
|
||||
logger := zap.New(core, zap.AddCaller(), zap.AddStacktrace(zapcore.ErrorLevel))
|
||||
|
||||
return &Logger{Logger: logger}
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestDiagnosticFiltering(t *testing.T) {
|
||||
for _, level := range []string{"debug", "error"} {
|
||||
t.Run(level, func(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
dir := filepath.Join(root, "log")
|
||||
log := New(level, filepath.Join(root, "primary.log"), DiagnosticOptions{Dir: dir})
|
||||
log.Debug("debug")
|
||||
log.Info("info")
|
||||
if _, err := os.Stat(dir); !os.IsNotExist(err) {
|
||||
t.Fatalf("ordinary logs created diagnostic directory: %v", err)
|
||||
}
|
||||
child := log.With(zap.String("conversation_id", "test-id"))
|
||||
child.Warn("retry", zap.Int("attempt", 2))
|
||||
child.Error("failed", zap.Error(fmt.Errorf("test failure")))
|
||||
files, _ := filepath.Glob(filepath.Join(dir, "*.log"))
|
||||
if len(files) != 1 {
|
||||
t.Fatalf("files: %v", files)
|
||||
}
|
||||
data, err := os.ReadFile(files[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
|
||||
if len(lines) != 2 {
|
||||
t.Fatalf("unexpected diagnostic records: %s", data)
|
||||
}
|
||||
for i, line := range lines {
|
||||
var record map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(line), &record); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if record["conversation_id"] != "test-id" || record["timestamp"] == nil || record["caller"] == nil {
|
||||
t.Fatalf("missing diagnostic context: %v", record)
|
||||
}
|
||||
if i == 1 && (record["stacktrace"] == nil || record["error"] != "test failure") {
|
||||
t.Fatalf("missing error details: %v", record)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiagnosticDisabled(t *testing.T) {
|
||||
dir := filepath.Join(t.TempDir(), "log")
|
||||
log := New("error", os.DevNull, DiagnosticOptions{Dir: dir, Disabled: true})
|
||||
log.Error("failure")
|
||||
if _, err := os.Stat(dir); !os.IsNotExist(err) {
|
||||
t.Fatalf("disabled diagnostics wrote files: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDailyRotationRetentionAndConcurrency(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
w := newDailyWriter(DiagnosticOptions{Dir: dir, RetentionDays: 2})
|
||||
now := time.Date(2026, 9, 8, 23, 59, 59, 0, time.Local)
|
||||
w.now = func() time.Time { return now }
|
||||
for _, name := range []string{"diagnostic-2026-09-06.log", "diagnostic-2026-09-07.log", "other.log", "diagnostic-invalid.log"} {
|
||||
if err := os.WriteFile(filepath.Join(dir, name), nil, 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if _, err := w.Write([]byte("before midnight\n")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "diagnostic-2026-09-06.log")); !os.IsNotExist(err) {
|
||||
t.Fatal("expired file remains")
|
||||
}
|
||||
now = now.Add(2 * time.Second)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 50; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := w.Write([]byte("after midnight\n")); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
for name, count := range map[string]int{"diagnostic-2026-09-08.log": 1, "diagnostic-2026-09-09.log": 50} {
|
||||
data, err := os.ReadFile(filepath.Join(dir, name))
|
||||
if err != nil || strings.Count(string(data), "\n") != count {
|
||||
t.Fatalf("%s: %q, %v", name, data, err)
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "diagnostic-2026-09-07.log")); !os.IsNotExist(err) {
|
||||
t.Fatal("rotation did not expire old file")
|
||||
}
|
||||
for _, name := range []string{"other.log", "diagnostic-invalid.log"} {
|
||||
if _, err := os.Stat(filepath.Join(dir, name)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiagnosticWriteFailureKeepsPrimaryOutput(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
primary := filepath.Join(root, "primary.log")
|
||||
log := New("info", primary, DiagnosticOptions{Dir: filepath.Join(primary, "invalid")})
|
||||
log.Error("still visible")
|
||||
data, err := os.ReadFile(primary)
|
||||
if err != nil || !strings.Contains(string(data), "still visible") {
|
||||
t.Fatalf("primary output lost: %s, %v", data, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
func TestBlockedExecutionIsTerminalAndNotFailed(t *testing.T) {
|
||||
for _, blocked := range []bool{true, false} {
|
||||
name := "error"
|
||||
want := ToolExecutionStatusFailed
|
||||
if blocked {
|
||||
name, want = "blocked", ToolExecutionStatusBlocked
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
service := NewExecutionService(nil, nil)
|
||||
handle, err := service.Submit(context.Background(), ExecutionRequest{
|
||||
ToolName: "test",
|
||||
Run: func(context.Context) (*ToolResult, error) {
|
||||
// Identical text must not turn ordinary failures into policy blocks.
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: toolGuardBlockedPrefix}}, IsError: true, Blocked: blocked}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
|
||||
if err != nil || snap.Execution.Status != want || snap.Execution.Result.Blocked != blocked || snap.Execution.Error == "" {
|
||||
t.Fatalf("incorrect classification: snapshot=%#v err=%v", snap, err)
|
||||
}
|
||||
if !isExecutionTerminal(want) || executionStatusCountsAsFailed(want) == blocked {
|
||||
t.Fatalf("incorrect terminal/failure classification for %s", want)
|
||||
}
|
||||
if service.Cancel(handle.ID, "cancel after completion") {
|
||||
t.Fatal("terminal execution must not be cancellable")
|
||||
}
|
||||
after, _ := service.Get(handle.ID)
|
||||
if after.Execution.Status != want {
|
||||
t.Fatalf("cancel reclassified terminal execution: %s", after.Execution.Status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBlockedMarkerSurvivesNormalizationAndMCPProtocol(t *testing.T) {
|
||||
original := &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("refused ", 2000)}}, IsError: true, Blocked: true}
|
||||
bounded := NormalizeToolResultForStorageWithSpill(original, 1000, ToolResultSpillConfig{RootDir: t.TempDir(), ExecutionID: "blocked"})
|
||||
if !bounded.Blocked || !bounded.IsError || ToolResultPlainText(bounded) == ToolResultPlainText(original) {
|
||||
t.Fatal("normalization must retain classification while bounding long output")
|
||||
}
|
||||
wire, err := json.Marshal(CallToolResponse{Content: bounded.Content, IsError: bounded.IsError, Blocked: bounded.Blocked, Meta: toolResultProtocolMeta(bounded)})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var decoded ToolResult
|
||||
if err := json.Unmarshal(wire, &decoded); err != nil || !decoded.Blocked || !decoded.IsError {
|
||||
t.Fatalf("application protocol lost block marker: %#v err=%v", decoded, err)
|
||||
}
|
||||
var sdkResult sdkmcp.CallToolResult
|
||||
if err := json.Unmarshal(wire, &sdkResult); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
converted := sdkCallToolResultToOurs(&sdkResult)
|
||||
if !converted.Blocked || !converted.IsError {
|
||||
t.Fatalf("SDK round trip lost block marker: %#v", converted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolStatsSeparateBlockedFromFailures(t *testing.T) {
|
||||
server := NewServer(nil)
|
||||
manager := NewExternalMCPManager(nil)
|
||||
for _, status := range []string{ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusBlocked, ToolExecutionStatusCancelled} {
|
||||
server.updateStats("test", status)
|
||||
manager.updateStats("test", status)
|
||||
}
|
||||
for name, stat := range map[string]*ToolStats{"internal": server.stats["test"], "external": manager.stats["test"]} {
|
||||
if stat.TotalCalls != 4 || stat.SuccessCalls != 1 || stat.FailedCalls != 1 || stat.BlockedCalls != 1 {
|
||||
t.Fatalf("%s stats = %#v", name, stat)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -308,9 +308,11 @@ func sdkCallToolResultToOurs(res *mcp.CallToolResult) *ToolResult {
|
||||
return &ToolResult{Content: []Content{}}
|
||||
}
|
||||
content := sdkContentToOurs(res.Content)
|
||||
blocked, _ := res.Meta[toolGuardBlockedMetaKey].(bool)
|
||||
return &ToolResult{
|
||||
Content: content,
|
||||
IsError: res.IsError,
|
||||
Blocked: blocked,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -190,6 +190,9 @@ func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) s
|
||||
if exec.Result != nil {
|
||||
payload["result"] = ToolResultPlainText(exec.Result)
|
||||
payload["is_error"] = exec.Result.IsError
|
||||
if exec.Result.Blocked {
|
||||
payload["blocked"] = true
|
||||
}
|
||||
}
|
||||
if opts.includePartialOutput && exec.PartialOutput != "" {
|
||||
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
|
||||
|
||||
@@ -18,6 +18,7 @@ const (
|
||||
ToolExecutionStatusQueued = "queued"
|
||||
ToolExecutionStatusRunning = "running"
|
||||
ToolExecutionStatusCompleted = "completed"
|
||||
ToolExecutionStatusBlocked = "blocked"
|
||||
ToolExecutionStatusFailed = "failed"
|
||||
ToolExecutionStatusCancelled = "cancelled"
|
||||
ToolExecutionStatusHardTimeout = "hard_timeout"
|
||||
@@ -224,6 +225,10 @@ func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
|
||||
|
||||
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
|
||||
id := entry.exec.ID
|
||||
var blockedErr *toolGuardBlockError
|
||||
if errors.As(err, &blockedErr) {
|
||||
result, err = blockedErr.result, nil
|
||||
}
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
|
||||
|
||||
now := time.Now()
|
||||
@@ -258,6 +263,10 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr
|
||||
entry.exec.Status = ToolExecutionStatusFailed
|
||||
entry.exec.Error = err.Error()
|
||||
}
|
||||
} else if result != nil && result.Blocked {
|
||||
entry.exec.Status = ToolExecutionStatusBlocked
|
||||
entry.exec.Error = firstToolResultText(result, "工具调用已被安全规则拦截")
|
||||
entry.exec.Result = result
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
entry.exec.Status = ToolExecutionStatusCancelled
|
||||
@@ -318,10 +327,11 @@ func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout
|
||||
if entry == nil {
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
if isExecutionTerminal(entry.exec.Status) {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
select {
|
||||
case <-entry.done:
|
||||
return s.snapshotEntry(entry), nil
|
||||
default:
|
||||
}
|
||||
|
||||
var timeoutCh <-chan time.Time
|
||||
var timer *time.Timer
|
||||
if timeout > 0 {
|
||||
@@ -332,18 +342,26 @@ func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout
|
||||
|
||||
select {
|
||||
case <-entry.done:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
return s.snapshotEntry(entry), nil
|
||||
case <-timeoutCh:
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
|
||||
return s.snapshotEntry(entry), ErrExecutionWaitTimeout
|
||||
case <-ctxDone(ctx):
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
|
||||
return s.snapshotEntry(entry), ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// snapshotEntry synchronizes snapshots with worker state and partial output
|
||||
// updates. Wait uses done to also observe persistence and completion callbacks.
|
||||
func (s *ExecutionService) snapshotEntry(entry *executionEntry) *ExecutionSnapshot {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}
|
||||
}
|
||||
|
||||
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
|
||||
entry := s.getEntry(executionID)
|
||||
if entry != nil {
|
||||
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
|
||||
return s.snapshotEntry(entry), nil
|
||||
}
|
||||
return s.getPersistedSnapshot(executionID)
|
||||
}
|
||||
@@ -464,6 +482,9 @@ func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID s
|
||||
}
|
||||
hasErr := err != nil && *err != nil
|
||||
hasRes := result != nil && *result != nil
|
||||
if hasRes && (*result).Blocked {
|
||||
return false
|
||||
}
|
||||
if !hasErr && !hasRes {
|
||||
return false
|
||||
}
|
||||
@@ -549,7 +570,16 @@ func isBackgroundWaitToolResult(result *ToolResult) bool {
|
||||
|
||||
func isExecutionTerminal(status string) bool {
|
||||
switch strings.TrimSpace(strings.ToLower(status)) {
|
||||
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
|
||||
case ToolExecutionStatusCompleted, ToolExecutionStatusBlocked, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func executionStatusCountsAsFailed(status string) bool {
|
||||
switch status {
|
||||
case ToolExecutionStatusFailed, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -74,6 +75,7 @@ type ExternalMCPManager struct {
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
toolGuard *toolguard.Manager
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
@@ -96,6 +98,23 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetToolGuard installs safety rules evaluated before dispatch to external MCPs.
|
||||
func (m *ExternalMCPManager) SetToolGuard(guard *toolguard.Manager) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.toolGuard = guard
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
func (m *ExternalMCPManager) checkToolGuard(toolName string, args map[string]interface{}) *ToolResult {
|
||||
m.mu.RLock()
|
||||
guard := m.toolGuard
|
||||
m.mu.RUnlock()
|
||||
return toolGuardBlockedResult(guard, toolName, args)
|
||||
}
|
||||
|
||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||
manager := &ExternalMCPManager{
|
||||
@@ -685,6 +704,7 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
}
|
||||
var mcpName, actualToolName string
|
||||
var client ExternalMCPClient
|
||||
var blockedByGuard bool
|
||||
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
@@ -702,6 +722,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
} else if authenticated {
|
||||
return nil, fmt.Errorf("external tool authorization policy is not configured")
|
||||
}
|
||||
if blocked := m.checkToolGuard(toolName, args); blocked != nil {
|
||||
blockedByGuard = true
|
||||
return nil, &toolGuardBlockError{result: blocked}
|
||||
}
|
||||
|
||||
// 解析工具名称:name::toolName
|
||||
if idx := findSubstring(toolName, "::"); idx > 0 {
|
||||
@@ -741,6 +765,11 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
return release, nil
|
||||
},
|
||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||
// Rules may have changed while this execution waited for a slot.
|
||||
if blocked := m.checkToolGuard(toolName, args); blocked != nil {
|
||||
blockedByGuard = true
|
||||
return blocked, nil
|
||||
}
|
||||
result, callErr := client.CallTool(runCtx, actualToolName, args)
|
||||
if callErr != nil {
|
||||
m.handleConnectionDead(mcpName, client, callErr)
|
||||
@@ -748,11 +777,13 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
return result, callErr
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||
if mcpName != "" {
|
||||
failed := exec != nil && executionStatusCountsAsFailed(exec.Status)
|
||||
if mcpName != "" && !blockedByGuard && (exec == nil || exec.Status != ToolExecutionStatusBlocked) {
|
||||
m.recordExternalMCPResult(mcpName, failed)
|
||||
}
|
||||
m.updateStats(toolName, failed)
|
||||
if exec != nil {
|
||||
m.updateStats(toolName, exec.Status)
|
||||
}
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
@@ -941,6 +972,9 @@ func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID
|
||||
}
|
||||
hasErr := err != nil && *err != nil
|
||||
hasRes := result != nil && *result != nil
|
||||
if hasRes && (*result).Blocked {
|
||||
return false
|
||||
}
|
||||
if !hasErr && !hasRes {
|
||||
return false
|
||||
}
|
||||
@@ -1098,15 +1132,15 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
|
||||
}
|
||||
|
||||
// updateStats 更新统计信息
|
||||
func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
|
||||
func (m *ExternalMCPManager) updateStats(toolName string, status string) {
|
||||
now := time.Now()
|
||||
if m.storage != nil {
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if failed {
|
||||
if executionStatusCountsAsFailed(status) {
|
||||
failedCalls = 1
|
||||
} else {
|
||||
} else if status == ToolExecutionStatusCompleted {
|
||||
successCalls = 1
|
||||
}
|
||||
if err := m.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
|
||||
@@ -1128,10 +1162,12 @@ func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
|
||||
stats.TotalCalls++
|
||||
stats.LastCallTime = &now
|
||||
|
||||
if failed {
|
||||
if executionStatusCountsAsFailed(status) {
|
||||
stats.FailedCalls++
|
||||
} else {
|
||||
} else if status == ToolExecutionStatusCompleted {
|
||||
stats.SuccessCalls++
|
||||
} else if status == ToolExecutionStatusBlocked {
|
||||
stats.BlockedCalls++
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -72,7 +72,9 @@ func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
|
||||
manager.ConfigureToolWaitTimeoutSeconds(1)
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("slow result ready")
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = client
|
||||
manager.mu.Unlock()
|
||||
|
||||
callCtx, callCancel := context.WithCancel(context.Background())
|
||||
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
|
||||
@@ -117,7 +119,9 @@ func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
client := newBlockingExternalMCPClient("control wait result")
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = client
|
||||
manager.mu.Unlock()
|
||||
|
||||
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
|
||||
if err != nil {
|
||||
@@ -157,7 +161,9 @@ func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
|
||||
CircuitCooldown: time.Second,
|
||||
})
|
||||
client := newBlockingExternalMCPClient("ok")
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = client
|
||||
manager.mu.Unlock()
|
||||
|
||||
done1 := make(chan struct{})
|
||||
go func() {
|
||||
@@ -217,7 +223,9 @@ func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
|
||||
CircuitFailureThreshold: 1,
|
||||
CircuitCooldown: time.Minute,
|
||||
})
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = &failingExternalMCPClient{}
|
||||
manager.mu.Unlock()
|
||||
|
||||
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "boom") {
|
||||
|
||||
+54
-16
@@ -16,6 +16,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
@@ -53,6 +54,7 @@ type Server struct {
|
||||
httpToolTimeoutMinutes *int
|
||||
httpToolTimeoutMu sync.RWMutex
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
toolGuard *toolguard.Manager
|
||||
executionService *ExecutionService
|
||||
toolWaitTimeout time.Duration
|
||||
toolResultMaxBytes int
|
||||
@@ -72,6 +74,23 @@ func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetToolGuard installs the runtime safety rules shared by HTTP and internal calls.
|
||||
func (s *Server) SetToolGuard(guard *toolguard.Manager) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.toolGuard = guard
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Server) checkToolGuard(toolName string, args map[string]interface{}) *ToolResult {
|
||||
s.mu.RLock()
|
||||
guard := s.toolGuard
|
||||
s.mu.RUnlock()
|
||||
return toolGuardBlockedResult(guard, toolName, args)
|
||||
}
|
||||
|
||||
type sseClient struct {
|
||||
id string
|
||||
send chan []byte
|
||||
@@ -566,7 +585,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
s.updateStats(req.Name, true)
|
||||
s.updateStats(req.Name, ToolExecutionStatusFailed)
|
||||
|
||||
return &Message{
|
||||
ID: msg.ID,
|
||||
@@ -590,10 +609,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
zap.Any("arguments", req.Arguments),
|
||||
)
|
||||
|
||||
result, err := handler(execCtx, req.Arguments)
|
||||
result := s.checkToolGuard(req.Name, req.Arguments)
|
||||
var err error
|
||||
if result == nil {
|
||||
result, err = handler(execCtx, req.Arguments)
|
||||
}
|
||||
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
|
||||
now := time.Now()
|
||||
var failed bool
|
||||
var finalResult *ToolResult
|
||||
|
||||
s.mu.Lock()
|
||||
@@ -604,13 +626,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
st, msg := executionStatusAndMessage(err)
|
||||
execution.Status = st
|
||||
execution.Error = msg
|
||||
failed = st != "cancelled"
|
||||
} else if result != nil && result.Blocked {
|
||||
execution.Status = ToolExecutionStatusBlocked
|
||||
execution.Error = firstToolResultText(result, toolGuardBlockedPrefix)
|
||||
execution.Result = result
|
||||
} else if result != nil && result.IsError {
|
||||
if cancelledWithUserNote {
|
||||
execution.Status = "cancelled"
|
||||
execution.Error = ""
|
||||
execution.Result = result
|
||||
failed = false
|
||||
} else {
|
||||
execution.Status = "failed"
|
||||
if len(result.Content) > 0 {
|
||||
@@ -619,7 +643,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
execution.Error = "工具执行返回错误结果"
|
||||
}
|
||||
execution.Result = result
|
||||
failed = true
|
||||
}
|
||||
} else {
|
||||
execution.Status = "completed"
|
||||
@@ -631,7 +654,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
}
|
||||
}
|
||||
execution.Result = result
|
||||
failed = false
|
||||
}
|
||||
|
||||
finalResult = execution.Result
|
||||
@@ -643,7 +665,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
}
|
||||
}
|
||||
|
||||
s.updateStats(req.Name, failed)
|
||||
s.updateStats(req.Name, execution.Status)
|
||||
|
||||
if s.storage != nil {
|
||||
s.mu.Lock()
|
||||
@@ -683,6 +705,8 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
errorResult, _ := json.Marshal(CallToolResponse{
|
||||
Content: finalResult.Content,
|
||||
IsError: true,
|
||||
Blocked: finalResult.Blocked,
|
||||
Meta: toolResultProtocolMeta(finalResult),
|
||||
})
|
||||
return &Message{
|
||||
ID: msg.ID,
|
||||
@@ -719,15 +743,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
||||
}
|
||||
|
||||
// updateStats 更新统计信息
|
||||
func (s *Server) updateStats(toolName string, failed bool) {
|
||||
func (s *Server) updateStats(toolName string, status string) {
|
||||
now := time.Now()
|
||||
if s.storage != nil {
|
||||
totalCalls := 1
|
||||
successCalls := 0
|
||||
failedCalls := 0
|
||||
if failed {
|
||||
if executionStatusCountsAsFailed(status) {
|
||||
failedCalls = 1
|
||||
} else {
|
||||
} else if status == ToolExecutionStatusCompleted {
|
||||
successCalls = 1
|
||||
}
|
||||
if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
|
||||
@@ -749,10 +773,12 @@ func (s *Server) updateStats(toolName string, failed bool) {
|
||||
stats.TotalCalls++
|
||||
stats.LastCallTime = &now
|
||||
|
||||
if failed {
|
||||
if executionStatusCountsAsFailed(status) {
|
||||
stats.FailedCalls++
|
||||
} else {
|
||||
} else if status == ToolExecutionStatusCompleted {
|
||||
stats.SuccessCalls++
|
||||
} else if status == ToolExecutionStatusBlocked {
|
||||
stats.BlockedCalls++
|
||||
}
|
||||
}
|
||||
|
||||
@@ -925,11 +951,15 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("工具 %s 未找到", toolName)
|
||||
}
|
||||
if blocked := s.checkToolGuard(toolName, args); blocked != nil {
|
||||
return blocked, nil
|
||||
}
|
||||
return handler(runCtx, args)
|
||||
},
|
||||
OnDone: func(exec *ToolExecution) {
|
||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||
s.updateStats(toolName, failed)
|
||||
if exec != nil {
|
||||
s.updateStats(toolName, exec.Status)
|
||||
}
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
@@ -1111,7 +1141,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
||||
}
|
||||
}
|
||||
|
||||
s.updateStats(exec.ToolName, failed)
|
||||
s.updateStats(exec.ToolName, exec.Status)
|
||||
|
||||
if s.storage != nil {
|
||||
s.mu.Lock()
|
||||
@@ -1155,6 +1185,11 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
|
||||
if executionID == "" || result == nil {
|
||||
return nil
|
||||
}
|
||||
if previous, ok := s.GetExecution(executionID); ok && previous != nil &&
|
||||
(previous.Status == ToolExecutionStatusBlocked || previous.Result != nil && previous.Result.Blocked) {
|
||||
result = cloneToolResult(result)
|
||||
result.Blocked, result.IsError = true, true
|
||||
}
|
||||
s.mu.Lock()
|
||||
spill := ToolResultSpillConfig{
|
||||
RootDir: s.spillRootDir,
|
||||
@@ -1270,6 +1305,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
|
||||
}
|
||||
hasErr := err != nil && *err != nil
|
||||
hasRes := result != nil && *result != nil
|
||||
if hasRes && (*result).Blocked {
|
||||
return false
|
||||
}
|
||||
if !hasErr && !hasRes {
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
)
|
||||
|
||||
const toolGuardBlockedPrefix = "工具调用已被安全规则拦截"
|
||||
const toolGuardBlockedMetaKey = "cyberstrike.ai/blocked"
|
||||
|
||||
// toolGuardBlockError carries structured policy results through pre-run hooks.
|
||||
type toolGuardBlockError struct{ result *ToolResult }
|
||||
|
||||
func (e *toolGuardBlockError) Error() string { return ToolResultPlainText(e.result) }
|
||||
|
||||
func toolResultProtocolMeta(result *ToolResult) map[string]interface{} {
|
||||
if result != nil && result.Blocked {
|
||||
return map[string]interface{}{toolGuardBlockedMetaKey: true}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// toolGuardBlockedResult uses the standard MCP error result so the refusal is
|
||||
// visible both to the model and in persisted execution monitoring records.
|
||||
func toolGuardBlockedResult(guard *toolguard.Manager, toolName string, args map[string]interface{}) *ToolResult {
|
||||
if guard == nil {
|
||||
return nil
|
||||
}
|
||||
match := guard.Check(toolName, args)
|
||||
if match == nil {
|
||||
return nil
|
||||
}
|
||||
message := toolGuardBlockedPrefix
|
||||
if custom := strings.TrimSpace(match.Message); custom != "" {
|
||||
message += ":" + custom
|
||||
}
|
||||
message += fmt.Sprintf("\n规则: %s (%s)\n匹配内容: %q", match.RuleName, match.RuleID, match.MatchedText)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: message}}, IsError: true, Blocked: true}
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func testToolGuard(t *testing.T, enabled bool) *toolguard.Manager {
|
||||
t.Helper()
|
||||
guard, err := toolguard.NewManager(toolguard.DefaultConfig())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := guard.Update(toolguard.Config{Enabled: enabled, Rules: []toolguard.Rule{{
|
||||
ID: "government", Name: "政府网站保护", Enabled: true,
|
||||
Pattern: `(?i)[a-z0-9.-]+\.gov(?:\.[a-z0-9.-]+)?`,
|
||||
Message: "识别到 {match},禁止攻击政府网站,请检查目标授权。",
|
||||
}}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return guard
|
||||
}
|
||||
|
||||
func assertGuardRefusal(t *testing.T, result *ToolResult, err error) {
|
||||
t.Helper()
|
||||
message := ToolResultPlainText(result)
|
||||
if err != nil {
|
||||
t.Fatalf("expected structured refusal, got error: %v", err)
|
||||
} else if result == nil || !result.IsError || !result.Blocked {
|
||||
t.Fatalf("expected tool error result, got %#v", result)
|
||||
}
|
||||
for _, text := range []string{toolGuardBlockedPrefix, "禁止攻击政府网站", "agency.gov.cn", "government"} {
|
||||
if !strings.Contains(message, text) {
|
||||
t.Errorf("refusal %q missing %q", message, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerToolGuardBlocksBeforeHandlerAndUpdatesLive(t *testing.T) {
|
||||
storage := newInMemoryMonitorStorage()
|
||||
server := NewServerWithStorage(zap.NewNop(), storage)
|
||||
guard := testToolGuard(t, true)
|
||||
server.SetToolGuard(guard)
|
||||
var calls, authorized atomic.Int32
|
||||
server.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error {
|
||||
authorized.Add(1)
|
||||
return nil
|
||||
})
|
||||
server.RegisterTool(Tool{Name: "scan"}, func(context.Context, map[string]interface{}) (*ToolResult, error) {
|
||||
calls.Add(1)
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
args := map[string]interface{}{"command": "scan https://agency.gov.cn"}
|
||||
result, executionID, err := server.CallTool(context.Background(), "scan", args)
|
||||
assertGuardRefusal(t, result, err)
|
||||
if calls.Load() != 0 || authorized.Load() != 1 {
|
||||
t.Fatalf("calls=%d authorized=%d, want 0 and 1", calls.Load(), authorized.Load())
|
||||
}
|
||||
execution, err := storage.GetToolExecution(executionID)
|
||||
if err != nil || execution == nil || execution.Status != ToolExecutionStatusBlocked || !strings.Contains(execution.Error, toolGuardBlockedPrefix) {
|
||||
t.Fatalf("expected persisted blocked execution, got %#v, err=%v", execution, err)
|
||||
}
|
||||
|
||||
result, _, err = server.CallTool(context.Background(), "scan", map[string]interface{}{"target": "example.org"})
|
||||
if err != nil || result.IsError || calls.Load() != 1 {
|
||||
t.Fatalf("allowed target did not execute: result=%#v calls=%d err=%v", result, calls.Load(), err)
|
||||
}
|
||||
cfg := guard.Config()
|
||||
cfg.Rules[0].Enabled = false
|
||||
if err := guard.Update(cfg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, _, err = server.CallTool(context.Background(), "scan", args)
|
||||
if err != nil || result.IsError || calls.Load() != 2 {
|
||||
t.Fatalf("disabled rule did not take effect: result=%#v calls=%d err=%v", result, calls.Load(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPToolGuardReturnsMCPErrorAndPersistsRefusal(t *testing.T) {
|
||||
storage := newInMemoryMonitorStorage()
|
||||
server := NewServerWithStorage(zap.NewNop(), storage)
|
||||
server.SetToolGuard(testToolGuard(t, true))
|
||||
var calls int
|
||||
server.RegisterTool(Tool{Name: "scan"}, func(context.Context, map[string]interface{}) (*ToolResult, error) {
|
||||
calls++
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
for _, tc := range []struct {
|
||||
target string
|
||||
blocked bool
|
||||
}{
|
||||
{target: "https://agency.gov.cn", blocked: true},
|
||||
{target: "https://example.org", blocked: false},
|
||||
} {
|
||||
body, err := json.Marshal(map[string]interface{}{
|
||||
"jsonrpc": "2.0", "id": 1, "method": "tools/call",
|
||||
"params": map[string]interface{}{"name": "scan", "arguments": map[string]interface{}{"target": tc.target}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
server.HandleHTTP(recorder, httptest.NewRequest(http.MethodPost, "/api/mcp", strings.NewReader(string(body))))
|
||||
var response Message
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if recorder.Code != http.StatusOK || response.Error != nil {
|
||||
t.Fatalf("expected MCP tool result, status=%d body=%s", recorder.Code, recorder.Body)
|
||||
}
|
||||
var result ToolResult
|
||||
if err := json.Unmarshal(response.Result, &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tc.blocked {
|
||||
assertGuardRefusal(t, &result, nil)
|
||||
if calls != 0 {
|
||||
t.Fatal("HTTP tool handler ran for a blocked target")
|
||||
}
|
||||
executions, err := storage.LoadToolExecutions()
|
||||
if err != nil || len(executions) != 1 || executions[0].Status != ToolExecutionStatusBlocked || !strings.Contains(executions[0].Error, toolGuardBlockedPrefix) {
|
||||
t.Fatalf("expected persisted HTTP refusal, got %#v err=%v", executions, err)
|
||||
}
|
||||
} else if result.IsError || calls != 1 {
|
||||
t.Fatalf("allowed HTTP target did not execute: result=%#v calls=%d", result, calls)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalToolGuardBlocksBeforeClientAndUpdatesLive(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
t.Cleanup(manager.StopAll)
|
||||
guard := testToolGuard(t, true)
|
||||
manager.SetToolGuard(guard)
|
||||
client := newBlockingExternalMCPClient("ok")
|
||||
close(client.release)
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = client
|
||||
manager.mu.Unlock()
|
||||
args := map[string]interface{}{"target": "https://agency.gov.cn"}
|
||||
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", args)
|
||||
assertGuardRefusal(t, result, err)
|
||||
if client.count.Load() != 0 {
|
||||
t.Fatal("external client ran for a blocked target")
|
||||
}
|
||||
execution, ok := manager.GetExecution(executionID)
|
||||
if !ok || execution.Status != ToolExecutionStatusBlocked || !strings.Contains(execution.Error, toolGuardBlockedPrefix) {
|
||||
t.Fatalf("expected blocked external execution, got %#v", execution)
|
||||
}
|
||||
cfg := guard.Config()
|
||||
cfg.Enabled = false
|
||||
if err := guard.Update(cfg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result, _, err = manager.CallTool(context.Background(), "lab::slow_tool", args)
|
||||
if err != nil || result.IsError || client.count.Load() != 1 {
|
||||
t.Fatalf("disabled guard did not take effect: result=%#v calls=%d err=%v", result, client.count.Load(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalToolGuardRechecksQueuedCallsWithoutTrippingCircuit(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
t.Cleanup(manager.StopAll)
|
||||
manager.toolWaitTimeout = 10 * time.Millisecond
|
||||
manager.ConfigureResilience(ExternalMCPResilienceConfig{
|
||||
MaxConcurrentPerServer: 1, MaxConcurrentTotal: 4,
|
||||
CircuitFailureThreshold: 1, CircuitCooldown: time.Minute,
|
||||
})
|
||||
guard := testToolGuard(t, false)
|
||||
manager.SetToolGuard(guard)
|
||||
client := newBlockingExternalMCPClient("ok")
|
||||
close(client.release)
|
||||
manager.mu.Lock()
|
||||
manager.clients["lab"] = client
|
||||
manager.mu.Unlock()
|
||||
|
||||
// Occupy the provider slot so the call passes its initial policy check and
|
||||
// remains queued until a live rule update is applied.
|
||||
release, err := manager.acquireExternalMCPCallSlot(context.Background(), "lab")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
released := false
|
||||
t.Cleanup(func() {
|
||||
if !released {
|
||||
release()
|
||||
}
|
||||
})
|
||||
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", map[string]interface{}{"target": "agency.gov.cn"})
|
||||
if err != nil || executionID == "" {
|
||||
t.Fatalf("failed to queue external call: id=%q err=%v", executionID, err)
|
||||
}
|
||||
deadline := time.After(time.Second)
|
||||
ticker := time.NewTicker(time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for len(manager.globalSemaphore) != 2 {
|
||||
select {
|
||||
case <-deadline:
|
||||
t.Fatal("execution did not reach the provider slot queue")
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
cfg := guard.Config()
|
||||
cfg.Enabled = true
|
||||
if err := guard.Update(cfg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
release()
|
||||
released = true
|
||||
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
|
||||
if err != nil || snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusBlocked {
|
||||
t.Fatalf("expected queued execution to be blocked on policy recheck, got %#v err=%v", snapshot, err)
|
||||
}
|
||||
assertGuardRefusal(t, snapshot.Execution.Result, nil)
|
||||
if client.count.Load() != 0 {
|
||||
t.Fatal("queued call bypassed the updated guard")
|
||||
}
|
||||
manager.mu.RLock()
|
||||
runtime := manager.serverRuntimes["lab"]
|
||||
failures, openUntil := runtime.consecutiveFailures, runtime.circuitOpenUntil
|
||||
manager.mu.RUnlock()
|
||||
if failures != 0 || !openUntil.IsZero() {
|
||||
t.Fatalf("local policy refusal affected provider circuit: failures=%d openUntil=%v", failures, openUntil)
|
||||
}
|
||||
result, _, err := manager.CallTool(context.Background(), "lab::slow_tool", map[string]interface{}{"target": "example.org"})
|
||||
if err != nil || result.IsError || client.count.Load() != 1 {
|
||||
t.Fatalf("allowed call failed after policy refusal: result=%#v calls=%d err=%v", result, client.count.Load(), err)
|
||||
}
|
||||
}
|
||||
@@ -116,6 +116,9 @@ type ToolCall struct {
|
||||
type ToolResult struct {
|
||||
Content []Content `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
// Blocked means policy stopped the call before execution. IsError remains
|
||||
// true for MCP/model handling, while monitoring uses a distinct status.
|
||||
Blocked bool `json:"blocked,omitempty"`
|
||||
}
|
||||
|
||||
// Content 表示内容
|
||||
@@ -184,8 +187,10 @@ type CallToolRequest struct {
|
||||
|
||||
// CallToolResponse 调用工具响应
|
||||
type CallToolResponse struct {
|
||||
Content []Content `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
Content []Content `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
Blocked bool `json:"blocked,omitempty"`
|
||||
Meta map[string]interface{} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// ToolExecution 工具执行记录
|
||||
@@ -193,7 +198,7 @@ type ToolExecution struct {
|
||||
ID string `json:"id"`
|
||||
ToolName string `json:"toolName"`
|
||||
Arguments map[string]interface{} `json:"arguments"`
|
||||
Status string `json:"status"` // pending, running, completed, failed, cancelled
|
||||
Status string `json:"status"` // queued, running, completed, blocked, failed, cancelled, hard_timeout, orphaned
|
||||
Result *ToolResult `json:"result,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
StartTime time.Time `json:"startTime"`
|
||||
@@ -216,6 +221,7 @@ type ToolStats struct {
|
||||
TotalCalls int `json:"totalCalls"`
|
||||
SuccessCalls int `json:"successCalls"`
|
||||
FailedCalls int `json:"failedCalls"`
|
||||
BlockedCalls int `json:"blockedCalls"`
|
||||
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// agenticOrphanToolPrunerMiddleware is the AgenticMessage equivalent of
|
||||
// orphanToolPrunerMiddleware. It removes user-role messages whose content
|
||||
// blocks are exclusively FunctionToolResult entries with CallIDs that do not
|
||||
// match any FunctionToolCall in the history.
|
||||
//
|
||||
// This is a defense-in-depth layer after agenticToolPairReconcilerMiddleware;
|
||||
// the reconciler handles the common case (assistant followed by its results)
|
||||
// while this pruner catches stray results that appear before their assistant
|
||||
// or in non-adjacent positions (e.g. after summarization rewriting).
|
||||
type agenticOrphanToolPrunerMiddleware struct {
|
||||
*adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]
|
||||
logger *zap.Logger
|
||||
phase string
|
||||
}
|
||||
|
||||
func newAgenticOrphanToolPrunerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] {
|
||||
return &agenticOrphanToolPrunerMiddleware{
|
||||
TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
|
||||
logger: logger,
|
||||
phase: phase,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *agenticOrphanToolPrunerMiddleware) BeforeModelRewriteState(
|
||||
ctx context.Context,
|
||||
state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
|
||||
mc *adk.TypedModelContext[*schema.AgenticMessage],
|
||||
) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) {
|
||||
_ = mc
|
||||
if m == nil || state == nil || len(state.Messages) == 0 {
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
// Pass 1: collect all provided CallIDs from assistant FunctionToolCall blocks.
|
||||
provided := make(map[string]struct{}, 8)
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
|
||||
continue
|
||||
}
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.FunctionToolCall != nil && block.FunctionToolCall.CallID != "" {
|
||||
provided[block.FunctionToolCall.CallID] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fast path: check if any orphan exists.
|
||||
hasOrphan := false
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil || !isPureAgenticToolResult(msg) {
|
||||
continue
|
||||
}
|
||||
for _, id := range agenticToolResultCallIDs(msg) {
|
||||
if _, ok := provided[id]; !ok {
|
||||
hasOrphan = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if hasOrphan {
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasOrphan {
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
// Pass 2: build pruned list.
|
||||
pruned := make([]*schema.AgenticMessage, 0, len(state.Messages))
|
||||
var droppedIDs []string
|
||||
var droppedNames []string
|
||||
for _, msg := range state.Messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if !isPureAgenticToolResult(msg) {
|
||||
pruned = append(pruned, msg)
|
||||
continue
|
||||
}
|
||||
// Check if ALL result call IDs are orphans. If any is matched, keep the
|
||||
// message (the reconciler already handled partial mismatches).
|
||||
allOrphan := true
|
||||
for _, id := range agenticToolResultCallIDs(msg) {
|
||||
if _, ok := provided[id]; ok {
|
||||
allOrphan = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if allOrphan {
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.FunctionToolResult != nil {
|
||||
droppedIDs = append(droppedIDs, block.FunctionToolResult.CallID)
|
||||
droppedNames = append(droppedNames, block.FunctionToolResult.Name)
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
pruned = append(pruned, msg)
|
||||
}
|
||||
|
||||
if len(droppedIDs) == 0 {
|
||||
return ctx, state, nil
|
||||
}
|
||||
if m.logger != nil {
|
||||
m.logger.Warn("agentic orphan tool messages pruned before model call",
|
||||
zap.String("phase", m.phase),
|
||||
zap.Int("dropped_count", len(droppedIDs)),
|
||||
zap.Strings("dropped_tool_call_ids", droppedIDs),
|
||||
zap.Strings("dropped_tool_names", droppedNames),
|
||||
zap.Int("messages_before", len(state.Messages)),
|
||||
zap.Int("messages_after", len(pruned)),
|
||||
)
|
||||
}
|
||||
ns := *state
|
||||
ns.Messages = pruned
|
||||
return ctx, &ns, nil
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// agenticToolPairReconcilerMiddleware is the AgenticMessage equivalent of
|
||||
// toolPairReconcilerMiddleware. It ensures every assistant FunctionToolCall
|
||||
// block is followed by a matching FunctionToolResult message, patching or
|
||||
// dropping as needed so the downstream model never receives an unpaired
|
||||
// tool-call history.
|
||||
//
|
||||
// In the AgenticMessage protocol:
|
||||
// - Assistant tool calls: Role=AgenticRoleTypeAssistant with FunctionToolCall content blocks.
|
||||
// - Tool results: Role=AgenticRoleTypeUser with FunctionToolResult content blocks.
|
||||
//
|
||||
// This middleware runs after summarization which may truncate history and
|
||||
// break pairings.
|
||||
type agenticToolPairReconcilerMiddleware struct {
|
||||
*adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]
|
||||
logger *zap.Logger
|
||||
phase string
|
||||
}
|
||||
|
||||
func newAgenticToolPairReconcilerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] {
|
||||
return &agenticToolPairReconcilerMiddleware{
|
||||
TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
|
||||
logger: logger,
|
||||
phase: phase,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *agenticToolPairReconcilerMiddleware) BeforeModelRewriteState(
|
||||
ctx context.Context,
|
||||
state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
|
||||
mc *adk.TypedModelContext[*schema.AgenticMessage],
|
||||
) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) {
|
||||
_ = mc
|
||||
if m == nil || state == nil || len(state.Messages) == 0 {
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
usedIDs := make(map[string]struct{}, 16)
|
||||
changed := false
|
||||
patched := 0
|
||||
dropped := 0
|
||||
out := make([]*schema.AgenticMessage, 0, len(state.Messages))
|
||||
|
||||
for i := 0; i < len(state.Messages); {
|
||||
msg := state.Messages[i]
|
||||
if msg == nil {
|
||||
changed = true
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
calls := agenticFunctionToolCalls(msg)
|
||||
|
||||
// Non-assistant or assistant without tool calls — but check for orphan
|
||||
// tool-result messages (user role with only FunctionToolResult blocks).
|
||||
if len(calls) == 0 {
|
||||
if isPureAgenticToolResult(msg) {
|
||||
// Orphan tool result not preceded by its assistant; drop it.
|
||||
changed = true
|
||||
dropped++
|
||||
i++
|
||||
continue
|
||||
}
|
||||
out = append(out, msg)
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
// Deduplicate / fix empty call IDs.
|
||||
idsChanged := false
|
||||
for ci := range calls {
|
||||
id := calls[ci].CallID
|
||||
_, duplicate := usedIDs[id]
|
||||
if id == "" || duplicate {
|
||||
base := fmt.Sprintf("patched_agentic_call_%d_%d", i, ci)
|
||||
id = base
|
||||
for suffix := 1; ; suffix++ {
|
||||
if _, exists := usedIDs[id]; !exists {
|
||||
break
|
||||
}
|
||||
id = fmt.Sprintf("%s_%d", base, suffix)
|
||||
}
|
||||
calls[ci].CallID = id
|
||||
idsChanged = true
|
||||
changed = true
|
||||
}
|
||||
usedIDs[id] = struct{}{}
|
||||
}
|
||||
|
||||
assistant := msg
|
||||
if idsChanged {
|
||||
assistant = cloneAgenticMessageWithCalls(msg, calls)
|
||||
}
|
||||
out = append(out, assistant)
|
||||
|
||||
// Build expected set.
|
||||
expected := make(map[string]*schema.FunctionToolCall, len(calls))
|
||||
for ci := range calls {
|
||||
expected[calls[ci].CallID] = calls[ci]
|
||||
}
|
||||
|
||||
// Consume following tool-result messages.
|
||||
results := make(map[string]*schema.AgenticMessage, len(calls))
|
||||
j := i + 1
|
||||
for j < len(state.Messages) {
|
||||
next := state.Messages[j]
|
||||
if next == nil {
|
||||
changed = true
|
||||
j++
|
||||
continue
|
||||
}
|
||||
if !isPureAgenticToolResult(next) {
|
||||
break
|
||||
}
|
||||
resultCallIDs := agenticToolResultCallIDs(next)
|
||||
consumed := false
|
||||
for _, rid := range resultCallIDs {
|
||||
if _, wanted := expected[rid]; !wanted {
|
||||
continue
|
||||
}
|
||||
if _, dup := results[rid]; dup {
|
||||
continue
|
||||
}
|
||||
results[rid] = next
|
||||
consumed = true
|
||||
}
|
||||
if !consumed {
|
||||
changed = true
|
||||
dropped++
|
||||
}
|
||||
j++
|
||||
}
|
||||
|
||||
// Emit results in call order, patching missing ones.
|
||||
for _, tc := range calls {
|
||||
if result, ok := results[tc.CallID]; ok {
|
||||
out = append(out, result)
|
||||
continue
|
||||
}
|
||||
out = append(out, makeAgenticPatchedToolResult(tc.CallID, tc.Name))
|
||||
changed = true
|
||||
patched++
|
||||
}
|
||||
i = j
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return ctx, state, nil
|
||||
}
|
||||
if m.logger != nil {
|
||||
m.logger.Warn("agentic tool-call/result pairs reconciled before model call",
|
||||
zap.String("phase", m.phase),
|
||||
zap.Int("patched_results", patched),
|
||||
zap.Int("dropped_results", dropped),
|
||||
zap.Int("messages_before", len(state.Messages)),
|
||||
zap.Int("messages_after", len(out)),
|
||||
)
|
||||
}
|
||||
ns := *state
|
||||
ns.Messages = out
|
||||
return ctx, &ns, nil
|
||||
}
|
||||
|
||||
// agenticFunctionToolCalls extracts FunctionToolCall pointers from an
|
||||
// assistant message's content blocks. Returns nil for non-assistant messages.
|
||||
func agenticFunctionToolCalls(msg *schema.AgenticMessage) []*schema.FunctionToolCall {
|
||||
if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
|
||||
return nil
|
||||
}
|
||||
var out []*schema.FunctionToolCall
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.FunctionToolCall != nil {
|
||||
out = append(out, block.FunctionToolCall)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// isPureAgenticToolResult returns true when the message is a user-role
|
||||
// message whose content blocks are exclusively FunctionToolResult entries.
|
||||
func isPureAgenticToolResult(msg *schema.AgenticMessage) bool {
|
||||
if msg == nil || msg.Role != schema.AgenticRoleTypeUser || len(msg.ContentBlocks) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block == nil {
|
||||
continue
|
||||
}
|
||||
if block.FunctionToolResult == nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// agenticToolResultCallIDs extracts all CallIDs from FunctionToolResult blocks.
|
||||
func agenticToolResultCallIDs(msg *schema.AgenticMessage) []string {
|
||||
if msg == nil {
|
||||
return nil
|
||||
}
|
||||
var ids []string
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.FunctionToolResult != nil && block.FunctionToolResult.CallID != "" {
|
||||
ids = append(ids, block.FunctionToolResult.CallID)
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func cloneAgenticMessageWithCalls(msg *schema.AgenticMessage, calls []*schema.FunctionToolCall) *schema.AgenticMessage {
|
||||
cloned := *msg
|
||||
cloned.ContentBlocks = make([]*schema.ContentBlock, 0, len(msg.ContentBlocks))
|
||||
callIdx := 0
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.FunctionToolCall != nil && callIdx < len(calls) {
|
||||
cloned.ContentBlocks = append(cloned.ContentBlocks, schema.NewContentBlock(calls[callIdx]))
|
||||
callIdx++
|
||||
} else {
|
||||
cloned.ContentBlocks = append(cloned.ContentBlocks, block)
|
||||
}
|
||||
}
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func makeAgenticPatchedToolResult(callID, name string) *schema.AgenticMessage {
|
||||
return &schema.AgenticMessage{
|
||||
Role: schema.AgenticRoleTypeUser,
|
||||
ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{
|
||||
CallID: callID,
|
||||
Name: name,
|
||||
Content: []*schema.FunctionToolResultContentBlock{{
|
||||
Type: schema.FunctionToolResultContentBlockTypeText,
|
||||
Text: &schema.UserInputText{Text: patchedMissingToolResult},
|
||||
}},
|
||||
})},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestAgenticToolPairReconcilerPatchesMissing(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
agenticAssistantToolCall("c1", "search", `{"q":"x"}`),
|
||||
agenticAssistantToolCall("c2", "execute", `{"cmd":"ls"}`),
|
||||
// c1 result present, c2 missing
|
||||
agenticToolResult("c1", "search", "found it"),
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Expected: assistant(c1) -> result(c1) -> assistant(c2) -> patched_result(c2)
|
||||
if len(out.Messages) != 4 {
|
||||
t.Fatalf("messages = %d, want 4", len(out.Messages))
|
||||
}
|
||||
// c1 assistant
|
||||
if calls := agenticFunctionToolCalls(out.Messages[0]); len(calls) != 1 || calls[0].CallID != "c1" {
|
||||
t.Fatal("msg[0] should be assistant(c1)")
|
||||
}
|
||||
// c1 result
|
||||
if ids := agenticToolResultCallIDs(out.Messages[1]); len(ids) != 1 || ids[0] != "c1" {
|
||||
t.Fatal("msg[1] should be result(c1)")
|
||||
}
|
||||
// c2 assistant
|
||||
if calls := agenticFunctionToolCalls(out.Messages[2]); len(calls) != 1 || calls[0].CallID != "c2" {
|
||||
t.Fatal("msg[2] should be assistant(c2)")
|
||||
}
|
||||
// c2 patched result
|
||||
if ids := agenticToolResultCallIDs(out.Messages[3]); len(ids) != 1 || ids[0] != "c2" {
|
||||
t.Fatal("msg[3] should be patched result(c2)")
|
||||
}
|
||||
resultText := out.Messages[3].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text
|
||||
if resultText != patchedMissingToolResult {
|
||||
t.Fatalf("patched text = %q", resultText)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgenticToolPairReconcilerDropsOrphan(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
// Orphan tool result with no preceding assistant
|
||||
agenticToolResult("orphan", "deleted_tool", "stale data"),
|
||||
{Role: schema.AgenticRoleTypeUser, ContentBlocks: []*schema.ContentBlock{
|
||||
schema.NewContentBlock(&schema.UserInputText{Text: "hello"}),
|
||||
}},
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Messages) != 1 {
|
||||
t.Fatalf("messages = %d, want 1 (orphan dropped)", len(out.Messages))
|
||||
}
|
||||
if out.Messages[0].ContentBlocks[0].UserInputText == nil {
|
||||
t.Fatal("remaining message should be the user text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgenticToolPairReconcilerNoopWhenPaired(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
agenticAssistantToolCall("c1", "search", `{}`),
|
||||
agenticToolResult("c1", "search", "ok"),
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Should return original state unchanged
|
||||
if &out.Messages[0] == &state.Messages[0] {
|
||||
// pointer equality on slice — state not cloned
|
||||
}
|
||||
if len(out.Messages) != 2 {
|
||||
t.Fatalf("messages = %d, want 2", len(out.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgenticToolPairReconcilerFixesEmptyCallID(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
{
|
||||
Role: schema.AgenticRoleTypeAssistant,
|
||||
ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{
|
||||
CallID: "", Name: "search", Arguments: `{}`,
|
||||
})},
|
||||
},
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
calls := agenticFunctionToolCalls(out.Messages[0])
|
||||
if len(calls) != 1 || calls[0].CallID == "" {
|
||||
t.Fatalf("empty call ID should be patched, got %q", calls[0].CallID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgenticOrphanToolPrunerRemovesOrphan(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticOrphanToolPrunerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
agenticAssistantToolCall("c1", "search", `{}`),
|
||||
agenticToolResult("c1", "search", "ok"),
|
||||
// Orphan: no assistant has call_id "c_orphan"
|
||||
agenticToolResult("c_orphan", "deleted", "stale"),
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Messages) != 2 {
|
||||
t.Fatalf("messages = %d, want 2 (orphan pruned)", len(out.Messages))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgenticOrphanToolPrunerNoopWhenClean(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticOrphanToolPrunerMiddleware(nil, "test")
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
agenticAssistantToolCall("c1", "search", `{}`),
|
||||
agenticToolResult("c1", "search", "ok"),
|
||||
},
|
||||
}
|
||||
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Messages) != 2 {
|
||||
t.Fatalf("messages = %d, want 2", len(out.Messages))
|
||||
}
|
||||
}
|
||||
@@ -443,22 +443,41 @@ func nextAgentEventWithContext(ctx context.Context, iter *adk.AsyncIterator[*adk
|
||||
|
||||
// recvSchemaMessageStream 消费 ADK Tool 流式结果;ctx 取消时立即返回,避免 amass 等无输出时永久阻塞。
|
||||
func recvSchemaMessageStream(ctx context.Context, stream *schema.StreamReader[*schema.Message]) (content, toolCallID, toolName string, recvErr error) {
|
||||
if stream == nil {
|
||||
return "", "", "", nil
|
||||
msgs, recvErr := recvSchemaToolResultMessages(ctx, stream)
|
||||
if len(msgs) == 0 {
|
||||
return "", "", "", recvErr
|
||||
}
|
||||
var buf strings.Builder
|
||||
recvErr = recvEinoSchemaMessageStreamWithContext(ctx, stream, 8, func(chunk *schema.Message) {
|
||||
if chunk.Content != "" {
|
||||
buf.WriteString(chunk.Content)
|
||||
parts := make([]string, 0, len(msgs))
|
||||
for _, msg := range msgs {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if tid := strings.TrimSpace(chunk.ToolCallID); tid != "" {
|
||||
toolCallID = tid
|
||||
parts = append(parts, msg.Content)
|
||||
if id := strings.TrimSpace(msg.ToolCallID); id != "" {
|
||||
toolCallID = id
|
||||
}
|
||||
if name := strings.TrimSpace(chunk.ToolName); name != "" {
|
||||
if name := strings.TrimSpace(msg.ToolName); name != "" {
|
||||
toolName = name
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, ""), toolCallID, toolName, recvErr
|
||||
}
|
||||
|
||||
// recvSchemaToolResultMessages 先收齐 Tool 流,再用 Eino ConcatMessages 合并。
|
||||
// EventSender 一 call 一条流时走 ConcatMessages;并行结果被摊平进同一条流时按 CallID 分列再合并。
|
||||
func recvSchemaToolResultMessages(ctx context.Context, stream *schema.StreamReader[*schema.Message]) (msgs []*schema.Message, recvErr error) {
|
||||
if stream == nil {
|
||||
return nil, nil
|
||||
}
|
||||
var chunks []*schema.Message
|
||||
recvErr = recvEinoSchemaMessageStreamWithContext(ctx, stream, 8, func(chunk *schema.Message) {
|
||||
chunks = append(chunks, chunk)
|
||||
})
|
||||
return buf.String(), toolCallID, toolName, recvErr
|
||||
msgs, concatErr := concatToolResultChunks(chunks)
|
||||
if concatErr != nil && recvErr == nil {
|
||||
return nil, concatErr
|
||||
}
|
||||
return msgs, recvErr
|
||||
}
|
||||
|
||||
func buildEinoCheckpointID(orchMode string) string {
|
||||
|
||||
@@ -30,6 +30,29 @@ func TestRecvSchemaMessageStream_EOF(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecvSchemaToolResultMessages_SplitsParallelIDs(t *testing.T) {
|
||||
sr, sw := schema.Pipe[*schema.Message](8)
|
||||
_ = sw.Send(schema.ToolMessage("one-", "tc-1", schema.WithToolName("nmap")), nil)
|
||||
_ = sw.Send(schema.ToolMessage("two-", "tc-2", schema.WithToolName("nmap")), nil)
|
||||
_ = sw.Send(schema.ToolMessage("a", "tc-1", schema.WithToolName("nmap")), nil)
|
||||
_ = sw.Send(schema.ToolMessage("b", "tc-2", schema.WithToolName("nmap")), nil)
|
||||
sw.Close()
|
||||
|
||||
msgs, err := recvSchemaToolResultMessages(context.Background(), sr)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected err: %v", err)
|
||||
}
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("msgs = %#v, want 2", msgs)
|
||||
}
|
||||
if msgs[0].ToolCallID != "tc-1" || msgs[0].Content != "one-a" {
|
||||
t.Fatalf("msg 0 = %#v", msgs[0])
|
||||
}
|
||||
if msgs[1].ToolCallID != "tc-2" || msgs[1].Content != "two-b" {
|
||||
t.Fatalf("msg 1 = %#v", msgs[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecvSchemaMessageStream_CapturesToolName(t *testing.T) {
|
||||
sr, sw := schema.Pipe[*schema.Message](4)
|
||||
_ = sw.Send(schema.ToolMessage("hello", "tc-1", schema.WithToolName("execute")), nil)
|
||||
|
||||
@@ -12,15 +12,24 @@ import (
|
||||
)
|
||||
|
||||
type capturingAgenticChatModel struct {
|
||||
mu sync.Mutex
|
||||
inputs [][]*schema.AgenticMessage
|
||||
output *schema.AgenticMessage
|
||||
mu sync.Mutex
|
||||
inputs [][]*schema.AgenticMessage
|
||||
output *schema.AgenticMessage
|
||||
outputs []*schema.AgenticMessage
|
||||
}
|
||||
|
||||
func (m *capturingAgenticChatModel) Generate(_ context.Context, input []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) {
|
||||
m.mu.Lock()
|
||||
m.inputs = append(m.inputs, input)
|
||||
callNo := len(m.inputs)
|
||||
m.mu.Unlock()
|
||||
if len(m.outputs) > 0 {
|
||||
idx := callNo - 1
|
||||
if idx >= len(m.outputs) {
|
||||
idx = len(m.outputs) - 1
|
||||
}
|
||||
return m.outputs[idx], nil
|
||||
}
|
||||
if m.output != nil {
|
||||
return m.output, nil
|
||||
}
|
||||
|
||||
@@ -20,8 +20,13 @@ func appendEinoAgenticChatModelTailMiddlewares(
|
||||
handlers = append(handlers, newAgenticSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
|
||||
handlers = append(handlers, newAgenticContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
|
||||
if cfg.agenticSummarization != nil {
|
||||
handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization"))
|
||||
handlers = append(handlers, cfg.agenticSummarization)
|
||||
}
|
||||
handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase))
|
||||
if !cfg.skipOrphanPruner {
|
||||
handlers = append(handlers, newAgenticOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
|
||||
}
|
||||
if !cfg.skipTrace && cfg.trace != nil {
|
||||
if capMw := newAgenticModelFacingTraceMiddleware(cfg.trace); capMw != nil {
|
||||
handlers = append(handlers, capMw)
|
||||
|
||||
@@ -106,7 +106,8 @@ func TestAppendEinoAgenticChatModelTailMiddlewares(t *testing.T) {
|
||||
phase: "agentic",
|
||||
trace: holder,
|
||||
})
|
||||
if len(handlers) != 3 {
|
||||
t.Fatalf("handlers = %d, want system + continuation + trace", len(handlers))
|
||||
// system + continuation + reconciler + orphan_pruner + trace
|
||||
if len(handlers) != 5 {
|
||||
t.Fatalf("handlers = %d, want system + continuation + reconciler + orphan_pruner + trace", len(handlers))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -31,6 +31,11 @@ func adaptAgenticEventToEinoEvents(ev *adk.TypedAgentEvent[*schema.AgenticMessag
|
||||
return []*adk.AgentEvent{base(&adk.AgentOutput{CustomizedOutput: customized})}
|
||||
}
|
||||
if mv.IsStreaming {
|
||||
// Tool 流保持 1 event ↔ 1 MessageStream,对齐 ADK EventSenderToolWrapper:
|
||||
// 每个 CallID 在工具包装层就已经是独立事件。这里不能再按 CallID 现场拆成
|
||||
// 多条 live pipe——drain 会阻塞读完当前流,交错的并行 chunk 会把另一列写满后死锁。
|
||||
// 若上游仍把 ToolsNode 的 MergeStreamReaders 摊成一条流,由
|
||||
// concatToolResultChunks 按列 ConcatMessages 恢复。
|
||||
return []*adk.AgentEvent{base(&adk.AgentOutput{
|
||||
MessageOutput: &adk.MessageVariant{
|
||||
IsStreaming: true,
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/middlewares/summarization"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
@@ -109,12 +108,10 @@ func newEinoAgenticSummarizationMiddleware(
|
||||
retryPolicy := einoTransientRunRetryPolicyFromMW(mwCfg)
|
||||
retryMax := retryPolicy.maxAttempts
|
||||
var summaryOverflowRetries int
|
||||
summaryModelOpts := []model.Option{
|
||||
einoopenai.WithMaxCompletionTokens(outputReserve),
|
||||
}
|
||||
summaryModelOpts := newEinoSummarizationModelOptions(outputReserve, modelName, "agentic", &appCfg.OpenAI, logger)
|
||||
|
||||
mw, err := summarization.NewTyped[*schema.AgenticMessage](ctx, &summarization.TypedConfig[*schema.AgenticMessage]{
|
||||
Model: summaryModel,
|
||||
Model: newNonEmptyAgenticSummaryModel(summaryModel),
|
||||
ModelOptions: summaryModelOpts,
|
||||
GenModelInput: func(ctx context.Context, sysInstruction, userInstruction *schema.AgenticMessage, originalMsgs []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) {
|
||||
classicOriginal := AgenticMessagesToEino(originalMsgs)
|
||||
|
||||
@@ -171,6 +171,59 @@ func TestEinoAgenticChatModelAgentCompactsContextBeforeBusinessModel(t *testing.
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoAgenticSummarizationMiddlewareRetriesWhenSummaryModelReturnsEmpty(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
emit := false
|
||||
summaryModel := &capturingAgenticChatModel{
|
||||
outputs: []*schema.AgenticMessage{
|
||||
agenticAssistantTextMessage(""),
|
||||
agenticAssistantTextMessage("<summary>有效摘要:继续验证 SQL 注入路径</summary>"),
|
||||
},
|
||||
}
|
||||
appCfg := &config.Config{}
|
||||
appCfg.OpenAI.Model = "gpt-4o"
|
||||
appCfg.OpenAI.MaxTotalTokens = 5000
|
||||
appCfg.Database.Path = filepath.Join(t.TempDir(), "cyberstrike.db")
|
||||
mwCfg := &config.MultiAgentEinoMiddlewareConfig{
|
||||
SummarizationEmitInternalEvents: &emit,
|
||||
SummarizationOutputReserveTokens: 1024,
|
||||
}
|
||||
|
||||
mw, err := newEinoAgenticSummarizationMiddleware(ctx, summaryModel, appCfg, mwCfg, "conv-agentic-empty-summary", nil, "", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("newEinoAgenticSummarizationMiddleware: %v", err)
|
||||
}
|
||||
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||
Messages: []*schema.AgenticMessage{
|
||||
schema.SystemAgenticMessage("system root"),
|
||||
schema.UserAgenticMessage("授权范围 example.com\n" + strings.Repeat("历史扫描输出 ", 12000)),
|
||||
agenticAssistantTextMessage("已记录范围"),
|
||||
schema.UserAgenticMessage("继续验证 SQL 注入路径"),
|
||||
},
|
||||
}
|
||||
|
||||
_, after, err := mw.BeforeModelRewriteState(ctx, state, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("BeforeModelRewriteState should retry instead of failing on empty summary: %v", err)
|
||||
}
|
||||
if after == nil {
|
||||
t.Fatal("after state is nil")
|
||||
}
|
||||
if inputs := summaryModel.snapshotInputs(); len(inputs) < 2 {
|
||||
t.Fatalf("summary model calls=%d, want retry after empty output", len(inputs))
|
||||
}
|
||||
joined := joinClassicMessageContent(AgenticMessagesToEino(after.Messages))
|
||||
for _, want := range []string{"有效摘要", "继续验证 SQL 注入路径", "原始用户输入与约束账本"} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Fatalf("retried compacted context missing %q:\n%s", want, joined)
|
||||
}
|
||||
}
|
||||
if strings.Contains(joined, "本地压缩摘要") {
|
||||
t.Fatalf("local fallback should not be used:\n%s", joined)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppendEinoAgenticChatModelTailMiddlewaresIncludesTypedSummarization(t *testing.T) {
|
||||
t.Parallel()
|
||||
mw := newAgenticSystemMessageNormalizerMiddleware(nil, "summary")
|
||||
|
||||
@@ -55,6 +55,70 @@ func TestEinoExtractFallbackAssistantFromMsgs_prefersToolOverEarlierAssistant(t
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoExtractFallbackAssistantFromMsgs_plainAssistant(t *testing.T) {
|
||||
msgs := []*schema.Message{
|
||||
schema.UserMessage("hi"),
|
||||
schema.AssistantMessage("plain answer", nil),
|
||||
}
|
||||
if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "plain answer" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoExtractFallbackAssistantFromMsgs_finalAssistantAfterToolResult(t *testing.T) {
|
||||
msgs := []*schema.Message{
|
||||
schema.UserMessage("hi"),
|
||||
schema.AssistantMessage("", []schema.ToolCall{{
|
||||
ID: "call-1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "execute",
|
||||
Arguments: `{"command":"pwd"}`,
|
||||
},
|
||||
}}),
|
||||
schema.ToolMessage("/tmp", "call-1", schema.WithToolName("execute")),
|
||||
schema.AssistantMessage("final after tool", nil),
|
||||
}
|
||||
if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "final after tool" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoExtractFallbackAssistantFromMsgs_doesNotUseAssistantBeforeUnfinishedToolResult(t *testing.T) {
|
||||
msgs := []*schema.Message{
|
||||
schema.UserMessage("hi"),
|
||||
schema.AssistantMessage("I will inspect that.", nil),
|
||||
schema.AssistantMessage("", []schema.ToolCall{{
|
||||
ID: "call-1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "execute",
|
||||
Arguments: `{"command":"pwd"}`,
|
||||
},
|
||||
}}),
|
||||
schema.ToolMessage("/tmp", "call-1", schema.WithToolName("execute")),
|
||||
}
|
||||
if got := einoExtractFallbackAssistantFromMsgs(msgs); got != "" {
|
||||
t.Fatalf("got %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunResultBuilderFinalFallsBackToPlainAssistantTrace(t *testing.T) {
|
||||
runMessages := newEinoRunMessageAccumulator(nil)
|
||||
runMessages.Append(schema.UserMessage("hi"))
|
||||
runMessages.Append(schema.AssistantMessage("plain answer", nil))
|
||||
|
||||
got := newEinoRunResultBuilder(einoRunResultBuilderConfig{
|
||||
OrchMode: "deep",
|
||||
EmptyHint: "empty",
|
||||
RunMessages: runMessages,
|
||||
}).BuildFinal()
|
||||
|
||||
if got.Response != "plain answer" {
|
||||
t.Fatalf("response = %q, want plain answer", got.Response)
|
||||
}
|
||||
}
|
||||
|
||||
func toolExitMsg(content, callID string) *schema.Message {
|
||||
m := schema.ToolMessage(content, callID)
|
||||
m.ToolName = "exit"
|
||||
|
||||
@@ -3,6 +3,8 @@ package multiagent
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
)
|
||||
@@ -82,12 +84,215 @@ func (h *einoRunErrorHandler) emitError(err error, kind string) {
|
||||
if h == nil || h.progress == nil || err == nil {
|
||||
return
|
||||
}
|
||||
userErr := einoUserFacingRunError(err)
|
||||
data := map[string]interface{}{
|
||||
"conversationId": h.conversationID,
|
||||
"source": "eino",
|
||||
"error": err.Error(),
|
||||
}
|
||||
if kind != "" {
|
||||
data["errorKind"] = kind
|
||||
} else if userErr.kind != "" {
|
||||
data["errorKind"] = userErr.kind
|
||||
}
|
||||
h.progress("error", err.Error(), data)
|
||||
if userErr.summary != "" {
|
||||
data["errorSummary"] = userErr.summary
|
||||
}
|
||||
if userErr.retryExhausted {
|
||||
data["retryExhausted"] = true
|
||||
if userErr.totalRetries > 0 {
|
||||
data["totalRetries"] = userErr.totalRetries
|
||||
}
|
||||
}
|
||||
if userErr.rawLastError != "" {
|
||||
data["lastError"] = userErr.rawLastError
|
||||
}
|
||||
if userErr.technicalError != "" {
|
||||
data["technicalError"] = userErr.technicalError
|
||||
}
|
||||
if userErr.hasModelOriginalError {
|
||||
data["modelOriginalError"] = userErr.rawLastError
|
||||
} else if userErr.retryExhausted {
|
||||
data["hasModelOriginalError"] = false
|
||||
}
|
||||
h.progress("error", EinoClientRunErrorMessage(err), data)
|
||||
}
|
||||
|
||||
type einoRunUserError struct {
|
||||
message string
|
||||
kind string
|
||||
summary string
|
||||
rawLastError string
|
||||
technicalError string
|
||||
retryExhausted bool
|
||||
totalRetries int
|
||||
hasModelOriginalError bool
|
||||
summarizationModelErr bool
|
||||
}
|
||||
|
||||
func einoUserFacingRunError(err error) einoRunUserError {
|
||||
var out einoRunUserError
|
||||
if err == nil {
|
||||
return out
|
||||
}
|
||||
var retryErr *adk.RetryExhaustedError
|
||||
if !errors.As(err, &retryErr) {
|
||||
return out
|
||||
}
|
||||
out.retryExhausted = true
|
||||
out.totalRetries = retryErr.TotalRetries
|
||||
lastErr := retryErr.LastErr
|
||||
if lastErr == nil {
|
||||
out.kind = "model_retry_exhausted"
|
||||
out.summary = "模型调用多次重试后仍未成功。"
|
||||
out.message = out.summary
|
||||
return out
|
||||
}
|
||||
out.rawLastError = strings.TrimSpace(lastErr.Error())
|
||||
if raw, ok := einoSummarizationModelRawErrorText(lastErr); ok {
|
||||
out.rawLastError = raw
|
||||
out.summarizationModelErr = true
|
||||
}
|
||||
if isEinoShouldRetryOutputRejected(lastErr) {
|
||||
out.kind = "model_output_rejected"
|
||||
out.summary = "模型未返回原始错误;输出被重试策略拒绝。"
|
||||
out.technicalError = out.rawLastError
|
||||
out.message = formatEinoRetryExhaustedMessage(out.summary, retryErr.TotalRetries)
|
||||
return out
|
||||
}
|
||||
kind, summary := einoTransientRunErrorUserDetail(lastErr)
|
||||
if strings.TrimSpace(summary) == "" {
|
||||
summary = einoTrimRetryErrorSummary(lastErr.Error())
|
||||
}
|
||||
if out.summarizationModelErr {
|
||||
summary = einoTrimRetryErrorSummary(out.rawLastError)
|
||||
}
|
||||
if kind == "" {
|
||||
kind = "model_retry_exhausted"
|
||||
}
|
||||
out.kind = kind
|
||||
out.summary = summary
|
||||
out.hasModelOriginalError = out.rawLastError != ""
|
||||
out.message = formatEinoRetryExhaustedMessage(summary, retryErr.TotalRetries)
|
||||
return out
|
||||
}
|
||||
|
||||
func isEinoShouldRetryOutputRejected(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
return strings.Contains(strings.ToLower(err.Error()), "model output rejected by shouldretry")
|
||||
}
|
||||
|
||||
func formatEinoRetryExhaustedMessage(summary string, totalRetries int) string {
|
||||
summary = strings.TrimSpace(summary)
|
||||
if summary == "" {
|
||||
summary = "模型调用多次重试后仍未成功。"
|
||||
}
|
||||
if totalRetries > 0 {
|
||||
return fmt.Sprintf("模型调用重试已耗尽(已重试 %d 次):%s", totalRetries, summary)
|
||||
}
|
||||
return "模型调用重试已耗尽:" + summary
|
||||
}
|
||||
|
||||
// EinoClientRunErrorMessage returns the error text that should be shown directly
|
||||
// to clients. When native retry hides the final provider failure behind a retry
|
||||
// wrapper, prefer the original last model error so summarization/model issues are
|
||||
// diagnosable from the frontend without opening server logs.
|
||||
func EinoClientRunErrorMessage(err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
userErr := einoUserFacingRunError(err)
|
||||
if userErr.retryExhausted {
|
||||
if userErr.hasModelOriginalError && userErr.rawLastError != "" {
|
||||
if userErr.summarizationModelErr {
|
||||
return formatEinoSummarizationRetryExhaustedRawModelMessage(userErr.rawLastError, userErr.totalRetries)
|
||||
}
|
||||
return formatEinoRetryExhaustedRawModelMessage(userErr.rawLastError, userErr.totalRetries)
|
||||
}
|
||||
if userErr.message != "" {
|
||||
return userErr.message
|
||||
}
|
||||
}
|
||||
if raw, ok := einoSummarizationModelRawErrorText(err); ok {
|
||||
return formatEinoSummarizationRawModelMessage(raw)
|
||||
}
|
||||
return err.Error()
|
||||
}
|
||||
|
||||
func formatEinoRetryExhaustedRawModelMessage(raw string, totalRetries int) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
prefix := "模型调用重试已耗尽"
|
||||
if totalRetries > 0 {
|
||||
prefix = fmt.Sprintf("模型调用重试已耗尽(已重试 %d 次)", totalRetries)
|
||||
}
|
||||
if raw == "" {
|
||||
return prefix
|
||||
}
|
||||
return prefix + ",最后一次模型原始错误:\n" + raw
|
||||
}
|
||||
|
||||
func formatEinoSummarizationRetryExhaustedRawModelMessage(raw string, totalRetries int) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
prefix := "摘要阶段大模型调用失败,模型调用重试已耗尽"
|
||||
if totalRetries > 0 {
|
||||
prefix = fmt.Sprintf("摘要阶段大模型调用失败,模型调用重试已耗尽(已重试 %d 次)", totalRetries)
|
||||
}
|
||||
if raw == "" {
|
||||
return prefix
|
||||
}
|
||||
return prefix + ",最后一次大模型报错原文:\n" + raw
|
||||
}
|
||||
|
||||
func formatEinoSummarizationRawModelMessage(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return "摘要阶段大模型调用失败。"
|
||||
}
|
||||
return "摘要阶段大模型调用失败,大模型报错原文:\n" + raw
|
||||
}
|
||||
|
||||
// EinoClientRunErrorFields returns structured diagnostic fields that handlers can
|
||||
// attach to their final error event in addition to the visible message.
|
||||
func EinoClientRunErrorFields(err error) map[string]interface{} {
|
||||
fields := make(map[string]interface{})
|
||||
if err == nil {
|
||||
return fields
|
||||
}
|
||||
userErr := einoUserFacingRunError(err)
|
||||
if userErr.kind != "" {
|
||||
fields["errorKind"] = userErr.kind
|
||||
}
|
||||
if userErr.summary != "" {
|
||||
fields["errorSummary"] = userErr.summary
|
||||
}
|
||||
if userErr.retryExhausted {
|
||||
fields["retryExhausted"] = true
|
||||
if userErr.totalRetries > 0 {
|
||||
fields["totalRetries"] = userErr.totalRetries
|
||||
}
|
||||
}
|
||||
if userErr.rawLastError != "" {
|
||||
fields["lastError"] = userErr.rawLastError
|
||||
}
|
||||
if userErr.technicalError != "" {
|
||||
fields["technicalError"] = userErr.technicalError
|
||||
}
|
||||
if userErr.hasModelOriginalError {
|
||||
fields["modelOriginalError"] = userErr.rawLastError
|
||||
} else if userErr.retryExhausted {
|
||||
fields["hasModelOriginalError"] = false
|
||||
}
|
||||
if userErr.summarizationModelErr {
|
||||
fields["errorPhase"] = "summarization"
|
||||
fields["summarizationModelError"] = true
|
||||
fields["modelOriginalError"] = userErr.rawLastError
|
||||
} else if raw, ok := einoSummarizationModelRawErrorText(err); ok {
|
||||
fields["errorPhase"] = "summarization"
|
||||
fields["summarizationModelError"] = true
|
||||
fields["modelOriginalError"] = raw
|
||||
fields["lastError"] = raw
|
||||
}
|
||||
return fields
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package multiagent
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
@@ -61,6 +62,177 @@ func TestEinoRunErrorHandlerTimeoutAndGeneralErrorProgress(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunErrorHandlerRetryExhaustedEmptyOutputProgress(t *testing.T) {
|
||||
err := &adk.RetryExhaustedError{
|
||||
LastErr: errors.New("model output rejected by ShouldRetry at attempt 5"),
|
||||
TotalRetries: 4,
|
||||
}
|
||||
var message string
|
||||
var data map[string]interface{}
|
||||
|
||||
got := newEinoRunErrorHandler(einoRunErrorHandlerConfig{
|
||||
ConversationID: "conv-1",
|
||||
Progress: func(eventType, msg string, raw interface{}) {
|
||||
if eventType == "error" {
|
||||
message = msg
|
||||
data, _ = raw.(map[string]interface{})
|
||||
}
|
||||
},
|
||||
}).Handle(err)
|
||||
|
||||
if !errors.Is(got, err) {
|
||||
t.Fatalf("err = %v", got)
|
||||
}
|
||||
if !strings.Contains(message, "模型调用重试已耗尽") ||
|
||||
!strings.Contains(message, "模型未返回原始错误;输出被重试策略拒绝。") ||
|
||||
strings.Contains(message, "model output rejected by ShouldRetry at attempt 5") {
|
||||
t.Fatalf("message = %q", message)
|
||||
}
|
||||
if data["errorKind"] != "model_output_rejected" {
|
||||
t.Fatalf("errorKind = %#v", data["errorKind"])
|
||||
}
|
||||
if data["errorSummary"] != "模型未返回原始错误;输出被重试策略拒绝。" {
|
||||
t.Fatalf("errorSummary = %#v", data["errorSummary"])
|
||||
}
|
||||
if data["hasModelOriginalError"] != false {
|
||||
t.Fatalf("hasModelOriginalError = %#v", data["hasModelOriginalError"])
|
||||
}
|
||||
if data["retryExhausted"] != true || data["totalRetries"] != 4 {
|
||||
t.Fatalf("retry metadata = %#v", data)
|
||||
}
|
||||
if data["lastError"] != "model output rejected by ShouldRetry at attempt 5" {
|
||||
t.Fatalf("lastError = %#v", data["lastError"])
|
||||
}
|
||||
if data["technicalError"] != "model output rejected by ShouldRetry at attempt 5" {
|
||||
t.Fatalf("technicalError = %#v", data["technicalError"])
|
||||
}
|
||||
if _, ok := data["modelOriginalError"]; ok {
|
||||
t.Fatalf("modelOriginalError should be absent for ShouldRetry rejection, got %#v", data["modelOriginalError"])
|
||||
}
|
||||
if data["error"] != err.Error() {
|
||||
t.Fatalf("raw error = %#v, want %#v", data["error"], err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunErrorHandlerRetryExhaustedOriginalErrorProgress(t *testing.T) {
|
||||
err := &adk.RetryExhaustedError{
|
||||
LastErr: errors.New("HTTP 429 Too Many Requests"),
|
||||
TotalRetries: 3,
|
||||
}
|
||||
var message string
|
||||
var data map[string]interface{}
|
||||
|
||||
got := newEinoRunErrorHandler(einoRunErrorHandlerConfig{
|
||||
ConversationID: "conv-1",
|
||||
Progress: func(eventType, msg string, raw interface{}) {
|
||||
if eventType == "error" {
|
||||
message = msg
|
||||
data, _ = raw.(map[string]interface{})
|
||||
}
|
||||
},
|
||||
}).Handle(err)
|
||||
|
||||
if !errors.Is(got, err) {
|
||||
t.Fatalf("err = %v", got)
|
||||
}
|
||||
if !strings.Contains(message, "HTTP 429 Too Many Requests") {
|
||||
t.Fatalf("message = %q", message)
|
||||
}
|
||||
if data["errorKind"] != "rate_limit" {
|
||||
t.Fatalf("errorKind = %#v", data["errorKind"])
|
||||
}
|
||||
if data["errorSummary"] != "HTTP 429 Too Many Requests" {
|
||||
t.Fatalf("errorSummary = %#v", data["errorSummary"])
|
||||
}
|
||||
if data["lastError"] != "HTTP 429 Too Many Requests" {
|
||||
t.Fatalf("lastError = %#v", data["lastError"])
|
||||
}
|
||||
if data["modelOriginalError"] != "HTTP 429 Too Many Requests" {
|
||||
t.Fatalf("modelOriginalError = %#v", data["modelOriginalError"])
|
||||
}
|
||||
if _, ok := data["hasModelOriginalError"]; ok {
|
||||
t.Fatalf("hasModelOriginalError should be absent when original error is present, got %#v", data["hasModelOriginalError"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoClientRunErrorMessageUsesRawRetryExhaustedModelError(t *testing.T) {
|
||||
raw := "summary content is empty: role=assistant content_runes=0 reasoning_runes=42\nprovider request id: req_123"
|
||||
err := &adk.RetryExhaustedError{
|
||||
LastErr: errors.New(raw),
|
||||
TotalRetries: 4,
|
||||
}
|
||||
|
||||
got := EinoClientRunErrorMessage(err)
|
||||
if !strings.Contains(got, "模型调用重试已耗尽(已重试 4 次)") {
|
||||
t.Fatalf("message missing retry prefix: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, raw) {
|
||||
t.Fatalf("message should include raw model error:\n%s", got)
|
||||
}
|
||||
|
||||
fields := EinoClientRunErrorFields(err)
|
||||
if fields["modelOriginalError"] != raw || fields["lastError"] != raw {
|
||||
t.Fatalf("raw fields = %#v", fields)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoClientRunErrorMessageMarksSummarizationModelRetryError(t *testing.T) {
|
||||
raw := `POST "https://api.deepseek.com/v1/chat/completions": 429 Too Many Requests: {"error":{"message":"Rate limit reached"}}`
|
||||
err := &adk.RetryExhaustedError{
|
||||
LastErr: newEinoSummarizationModelError(errors.New(raw)),
|
||||
TotalRetries: 4,
|
||||
}
|
||||
|
||||
got := EinoClientRunErrorMessage(err)
|
||||
for _, want := range []string{
|
||||
"摘要阶段大模型调用失败",
|
||||
"模型调用重试已耗尽(已重试 4 次)",
|
||||
"最后一次大模型报错原文",
|
||||
raw,
|
||||
} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("message missing %q:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
if strings.Contains(got, "summarization model error") {
|
||||
t.Fatalf("message leaked internal wrapper:\n%s", got)
|
||||
}
|
||||
|
||||
fields := EinoClientRunErrorFields(err)
|
||||
if fields["errorPhase"] != "summarization" || fields["summarizationModelError"] != true {
|
||||
t.Fatalf("phase fields = %#v", fields)
|
||||
}
|
||||
if fields["modelOriginalError"] != raw || fields["lastError"] != raw {
|
||||
t.Fatalf("raw fields = %#v", fields)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoClientRunErrorMessageMarksDirectSummarizationModelError(t *testing.T) {
|
||||
raw := `POST "https://api.deepseek.com/v1/chat/completions": 400 Bad Request: {"error":{"message":"invalid thinking parameter"}}`
|
||||
err := newEinoSummarizationModelError(errors.New(raw))
|
||||
|
||||
got := EinoClientRunErrorMessage(err)
|
||||
for _, want := range []string{
|
||||
"摘要阶段大模型调用失败",
|
||||
"大模型报错原文",
|
||||
raw,
|
||||
} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("message missing %q:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
if strings.Contains(got, "summarization model error") {
|
||||
t.Fatalf("message leaked internal wrapper:\n%s", got)
|
||||
}
|
||||
|
||||
fields := EinoClientRunErrorFields(err)
|
||||
if fields["errorPhase"] != "summarization" ||
|
||||
fields["modelOriginalError"] != raw ||
|
||||
fields["lastError"] != raw {
|
||||
t.Fatalf("fields = %#v", fields)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunErrorHandlerIterationLimitProgress(t *testing.T) {
|
||||
var events []string
|
||||
var errorKind interface{}
|
||||
|
||||
@@ -57,6 +57,16 @@ func (a *einoRunMessageAccumulator) Messages() []adk.Message {
|
||||
return a.msgs
|
||||
}
|
||||
|
||||
func (a *einoRunMessageAccumulator) NewMessages() []adk.Message {
|
||||
if a == nil {
|
||||
return nil
|
||||
}
|
||||
if a.baseCount < 0 || a.baseCount >= len(a.msgs) {
|
||||
return nil
|
||||
}
|
||||
return a.msgs[a.baseCount:]
|
||||
}
|
||||
|
||||
func (a *einoRunMessageAccumulator) BaseCount() int {
|
||||
if a == nil {
|
||||
return 0
|
||||
|
||||
@@ -27,6 +27,10 @@ func TestEinoRunMessageAccumulatorTracksBaseAndAppends(t *testing.T) {
|
||||
if len(msgs) != 2 || msgs[1].Role != schema.Assistant || msgs[1].Content != "hello" {
|
||||
t.Fatalf("messages = %#v", msgs)
|
||||
}
|
||||
newMsgs := acc.NewMessages()
|
||||
if len(newMsgs) != 1 || newMsgs[0].Content != "hello" {
|
||||
t.Fatalf("new messages = %#v", newMsgs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunMessageAccumulatorToolMessage(t *testing.T) {
|
||||
|
||||
@@ -98,6 +98,39 @@ func TestEinoRunProgressTrackerDedupesToolCalls(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunProgressTrackerDedupesSameToolCallIDsWithDifferentArgs(t *testing.T) {
|
||||
var toolCalls int
|
||||
progress := func(eventType, _ string, _ interface{}) {
|
||||
if eventType == "tool_call" {
|
||||
toolCalls++
|
||||
}
|
||||
}
|
||||
tracker := newEinoRunProgressTracker("deep", "lead", "conv-1", progress, nil, nil)
|
||||
first := &schema.Message{ToolCalls: []schema.ToolCall{{
|
||||
ID: "call-1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "nmap",
|
||||
Arguments: `{"host":"10.0.0.1"}`,
|
||||
},
|
||||
}}}
|
||||
second := &schema.Message{ToolCalls: []schema.ToolCall{{
|
||||
ID: "call-1",
|
||||
Type: "function",
|
||||
Function: schema.FunctionCall{
|
||||
Name: "nmap",
|
||||
Arguments: `{"host":"10.0.0.1","ports":"1-1024"}`,
|
||||
},
|
||||
}}}
|
||||
|
||||
tracker.EmitToolCalls(first, "lead", nil)
|
||||
tracker.EmitToolCalls(second, "lead", nil)
|
||||
|
||||
if toolCalls != 1 {
|
||||
t.Fatalf("tool call events = %d, want 1", toolCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunProgressTrackerHidesModelOutputRecoveryToolCalls(t *testing.T) {
|
||||
var eventTypes []string
|
||||
var marked []toolCallPendingInfo
|
||||
|
||||
@@ -45,7 +45,7 @@ func (b *einoRunResultBuilder) BuildFinal() *RunResult {
|
||||
func (b *einoRunResultBuilder) build(partial bool) *RunResult {
|
||||
var runMsgs []adk.Message
|
||||
if b.cfg.RunMessages != nil {
|
||||
runMsgs = b.cfg.RunMessages.Messages()
|
||||
runMsgs = b.cfg.RunMessages.NewMessages()
|
||||
}
|
||||
var lastAssistant string
|
||||
var lastPlanExecuteExecutor string
|
||||
@@ -107,6 +107,9 @@ func buildEinoRunResultFromAccumulated(
|
||||
if cleaned == "" {
|
||||
if fb := strings.TrimSpace(einoExtractFallbackAssistantFromMsgs(runAccumulatedMsgs)); fb != "" {
|
||||
cleaned = fb
|
||||
if orchMode == "plan_execute" {
|
||||
cleaned = UnwrapPlanExecuteUserText(cleaned)
|
||||
}
|
||||
}
|
||||
}
|
||||
cleaned = dedupeRepeatedParagraphs(cleaned, 80)
|
||||
@@ -146,32 +149,38 @@ func markModelFacingTraceForPersistence(msgs []adk.Message) []adk.Message {
|
||||
return out
|
||||
}
|
||||
|
||||
// einoExtractFallbackAssistantFromMsgs 在「主通道未产出助手正文」时,从 Eino ADK 轨迹中回填用户可见回复。
|
||||
// 典型场景:监督者仅调用 exit(final_result 落在 Tool 消息中),或工具结果已写入历史但 lastAssistant 未更新。
|
||||
// einoExtractFallbackAssistantFromMsgs 在「主通道未产出助手正文」时,从 Eino ADK
|
||||
// 原生消息轨迹中回填用户可见回复。这里保持克制:只采纳倒序最近的可交付终态,
|
||||
// 避免把工具调用前的过渡语或子任务过程误升为最终回复。
|
||||
//
|
||||
// 优先级:最后一次 exit 工具输出 → 最后一条含 exit 的助手 tool_calls 参数中的 final_result。
|
||||
// 可交付终态:
|
||||
// - exit 工具输出;
|
||||
// - assistant 调用 exit 时 arguments.final_result;
|
||||
// - 没有后续普通工具结果截断的纯 assistant 正文。
|
||||
func einoExtractFallbackAssistantFromMsgs(msgs []adk.Message) string {
|
||||
for i := len(msgs) - 1; i >= 0; i-- {
|
||||
m := msgs[i]
|
||||
if m == nil || m.Role != schema.Tool {
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(m.ToolName), adk.ToolInfoExit.Name) {
|
||||
continue
|
||||
}
|
||||
content := strings.TrimSpace(m.Content)
|
||||
if content == "" || strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||
continue
|
||||
}
|
||||
return content
|
||||
}
|
||||
for i := len(msgs) - 1; i >= 0; i-- {
|
||||
m := msgs[i]
|
||||
if m == nil || m.Role != schema.Assistant {
|
||||
continue
|
||||
}
|
||||
if s := einoExtractExitFinalFromAssistantToolCalls(m); s != "" {
|
||||
return s
|
||||
switch m.Role {
|
||||
case schema.Tool:
|
||||
if strings.EqualFold(strings.TrimSpace(m.ToolName), adk.ToolInfoExit.Name) {
|
||||
content := strings.TrimSpace(m.Content)
|
||||
if content != "" && !strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||
return content
|
||||
}
|
||||
}
|
||||
return ""
|
||||
case schema.Assistant:
|
||||
if s := einoExtractExitFinalFromAssistantToolCalls(m); s != "" {
|
||||
return s
|
||||
}
|
||||
if len(m.ToolCalls) == 0 {
|
||||
if content := strings.TrimSpace(m.Content); content != "" {
|
||||
return content
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
|
||||
@@ -55,6 +55,24 @@ func TestEinoRunResultBuilderFinalUsesSnapshots(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunResultBuilderFallbackIgnoresBaseHistory(t *testing.T) {
|
||||
runMessages := newEinoRunMessageAccumulator([]adk.Message{
|
||||
schema.UserMessage("previous request"),
|
||||
schema.AssistantMessage("previous answer", nil),
|
||||
schema.UserMessage("new request"),
|
||||
})
|
||||
|
||||
got := newEinoRunResultBuilder(einoRunResultBuilderConfig{
|
||||
OrchMode: "deep",
|
||||
EmptyHint: "empty",
|
||||
RunMessages: runMessages,
|
||||
}).BuildFinal()
|
||||
|
||||
if got.Response != "empty" {
|
||||
t.Fatalf("response = %q, want empty hint", got.Response)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunResultBuilderPlanExecutePrefersExecutorOutput(t *testing.T) {
|
||||
runMessages := newEinoRunMessageAccumulator(nil)
|
||||
runMessages.Append(schema.AssistantMessage(`{"response":"planner text"}`, nil))
|
||||
@@ -73,3 +91,18 @@ func TestEinoRunResultBuilderPlanExecutePrefersExecutorOutput(t *testing.T) {
|
||||
t.Fatalf("response = %q, want executor text", got.Response)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunResultBuilderPlanExecuteUnwrapsFallbackAssistant(t *testing.T) {
|
||||
runMessages := newEinoRunMessageAccumulator(nil)
|
||||
runMessages.Append(schema.AssistantMessage(`{"response":"fallback executor text"}`, nil))
|
||||
|
||||
got := newEinoRunResultBuilder(einoRunResultBuilderConfig{
|
||||
OrchMode: "plan_execute",
|
||||
EmptyHint: "empty",
|
||||
RunMessages: runMessages,
|
||||
}).BuildFinal()
|
||||
|
||||
if got.Response != "fallback executor text" {
|
||||
t.Fatalf("response = %q, want fallback executor text", got.Response)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -371,5 +371,9 @@ func (s *einoRunRuntimeSession) emitUsageSummary(reason string) bool {
|
||||
if s == nil || s.usage == nil {
|
||||
return false
|
||||
}
|
||||
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, s.progress, s.logger)
|
||||
modelName := ""
|
||||
if s.args != nil {
|
||||
modelName = s.args.ModelName
|
||||
}
|
||||
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, modelName, s.progress, s.logger)
|
||||
}
|
||||
|
||||
@@ -61,6 +61,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
|
||||
conversationID string,
|
||||
orchestration string,
|
||||
reason string,
|
||||
modelName string,
|
||||
progress func(eventType, message string, data interface{}),
|
||||
logger *zap.Logger,
|
||||
) bool {
|
||||
@@ -81,6 +82,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
|
||||
"source": "eino",
|
||||
"orchestration": orchestration,
|
||||
"reason": reason,
|
||||
"model": modelName,
|
||||
"modelCalls": s.ModelCalls,
|
||||
"promptTokens": s.PromptTokens,
|
||||
"completionTokens": s.CompletionTokens,
|
||||
@@ -96,6 +98,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("orchestration", orchestration),
|
||||
zap.String("reason", reason),
|
||||
zap.String("model", modelName),
|
||||
zap.Int("modelCalls", s.ModelCalls),
|
||||
zap.Int("promptTokens", s.PromptTokens),
|
||||
zap.Int("completionTokens", s.CompletionTokens),
|
||||
|
||||
@@ -49,16 +49,16 @@ func TestEinoRunUsageAccumulatorEmitOnce(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
if !acc.EmitOnce("conv-1", "deep", "final", progress, nil) {
|
||||
if !acc.EmitOnce("conv-1", "deep", "final", "gpt-test", progress, nil) {
|
||||
t.Fatal("first emit should return true")
|
||||
}
|
||||
if acc.EmitOnce("conv-1", "deep", "partial", progress, nil) {
|
||||
if acc.EmitOnce("conv-1", "deep", "partial", "gpt-test", progress, nil) {
|
||||
t.Fatal("second emit should return false")
|
||||
}
|
||||
if len(events) != 1 {
|
||||
t.Fatalf("events = %#v, want one usage summary", events)
|
||||
}
|
||||
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["totalTokens"] != 3 {
|
||||
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["model"] != "gpt-test" || events[0]["totalTokens"] != 3 {
|
||||
t.Fatalf("event = %#v", events[0])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,15 +203,18 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
|
||||
return runEinoADKAgentLoop(ctx, &einoADKRunLoopArgs{
|
||||
OrchMode: "eino_single",
|
||||
OrchestratorName: einoSingleAgentName,
|
||||
ConversationID: conversationID,
|
||||
Progress: progress,
|
||||
Logger: logger,
|
||||
SnapshotMCPIDs: snapshotMCPIDs,
|
||||
StreamsMainAssistant: streamsMainAssistant,
|
||||
EinoRoleTag: einoRoleTag,
|
||||
CheckpointDir: ma.EinoMiddleware.CheckpointDir,
|
||||
OrchMode: "eino_single",
|
||||
OrchestratorName: einoSingleAgentName,
|
||||
ConversationID: conversationID,
|
||||
Progress: progress,
|
||||
Logger: logger,
|
||||
SnapshotMCPIDs: snapshotMCPIDs,
|
||||
StreamsMainAssistant: streamsMainAssistant,
|
||||
EinoRoleTag: einoRoleTag,
|
||||
// Chat history recovery is intentionally centralized in last_react_*.
|
||||
// ADK checkpoints are a second persisted model-state channel and make
|
||||
// stale-context bugs hard to reason about across user turns.
|
||||
CheckpointDir: "",
|
||||
RunRetryMaxAttempts: RunRetryMaxAttemptsFromConfig(&ma.EinoMiddleware),
|
||||
RunRetryMaxBackoffSec: int(einoRunRetryMaxBackoffFromConfig(&ma.EinoMiddleware).Seconds()),
|
||||
McpIDsMu: &mcpIDsMu,
|
||||
|
||||
@@ -164,27 +164,10 @@ func newEinoSummarizationMiddleware(
|
||||
retryMax := retryPolicy.maxAttempts
|
||||
var summaryOverflowRetries int
|
||||
|
||||
// ModelOptions apply only to summarization Generate (same ChatModel instance as the agent).
|
||||
// Strip thinking/reasoning on this call path; mark requests for empty-choices diagnostics.
|
||||
summaryModelOpts := []model.Option{
|
||||
einoopenai.WithMaxCompletionTokens(outputReserve),
|
||||
einoopenai.WithExtraHeader(map[string]string{
|
||||
copenai.SummarizationRequestHeader: "1",
|
||||
}),
|
||||
einoopenai.WithRequestPayloadModifier(func(_ context.Context, in []*schema.Message, rawBody []byte) ([]byte, error) {
|
||||
if logger != nil {
|
||||
logger.Info("eino summarization generate request",
|
||||
zap.Int("input_messages", len(in)),
|
||||
zap.Int("payload_bytes", len(rawBody)),
|
||||
zap.String("model", modelName),
|
||||
)
|
||||
}
|
||||
return stripReasoningFromSummarizationPayload(rawBody)
|
||||
}),
|
||||
}
|
||||
summaryModelOpts := newEinoSummarizationModelOptions(outputReserve, modelName, "classic", &appCfg.OpenAI, logger)
|
||||
|
||||
mw, err := summarization.New(ctx, &summarization.Config{
|
||||
Model: summaryModel,
|
||||
Model: newNonEmptySummaryChatModel(summaryModel),
|
||||
ModelOptions: summaryModelOpts,
|
||||
GenModelInput: func(ctx context.Context, sysInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) {
|
||||
if transcriptPath != "" && len(originalMsgs) > 0 {
|
||||
@@ -308,6 +291,34 @@ func newEinoSummarizationMiddleware(
|
||||
return mw, nil
|
||||
}
|
||||
|
||||
// newEinoSummarizationModelOptions applies only to summarization Generate calls
|
||||
// on the shared main model. Summary generation should be plain-text and cheap:
|
||||
// strip provider reasoning/thinking controls so DeepSeek/OpenAI-compatible
|
||||
// endpoints do not spend the reserved output budget on invisible reasoning.
|
||||
func newEinoSummarizationModelOptions(outputReserve int, modelName, kind string, oa *config.OpenAIConfig, logger *zap.Logger) []model.Option {
|
||||
label := "eino summarization generate request"
|
||||
if strings.TrimSpace(kind) != "" && kind != "classic" {
|
||||
label = "eino " + kind + " summarization generate request"
|
||||
}
|
||||
return []model.Option{
|
||||
model.WithMaxTokens(outputReserve),
|
||||
einoopenai.WithMaxCompletionTokens(outputReserve),
|
||||
einoopenai.WithExtraHeader(map[string]string{
|
||||
copenai.SummarizationRequestHeader: "1",
|
||||
}),
|
||||
einoopenai.WithRequestPayloadModifier(func(_ context.Context, in []*schema.Message, rawBody []byte) ([]byte, error) {
|
||||
if logger != nil {
|
||||
logger.Info(label,
|
||||
zap.Int("input_messages", len(in)),
|
||||
zap.Int("payload_bytes", len(rawBody)),
|
||||
zap.String("model", modelName),
|
||||
)
|
||||
}
|
||||
return stripReasoningFromSummarizationPayload(rawBody, oa)
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// summarizationInputBudgetOpts controls spill/truncation behavior when a round alone exceeds budget.
|
||||
type summarizationInputBudgetOpts struct {
|
||||
toolMaxBytes int
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
type einoSummarizationModelError struct {
|
||||
err error
|
||||
}
|
||||
|
||||
func newEinoSummarizationModelError(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
var existing *einoSummarizationModelError
|
||||
if errors.As(err, &existing) {
|
||||
return err
|
||||
}
|
||||
return &einoSummarizationModelError{err: err}
|
||||
}
|
||||
|
||||
func (e *einoSummarizationModelError) Error() string {
|
||||
if e == nil || e.err == nil {
|
||||
return "summarization model error"
|
||||
}
|
||||
return "summarization model error: " + e.err.Error()
|
||||
}
|
||||
|
||||
func (e *einoSummarizationModelError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.err
|
||||
}
|
||||
|
||||
func einoSummarizationModelRawErrorText(err error) (string, bool) {
|
||||
var summaryErr *einoSummarizationModelError
|
||||
if !errors.As(err, &summaryErr) || summaryErr == nil || summaryErr.err == nil {
|
||||
return "", false
|
||||
}
|
||||
return strings.TrimSpace(summaryErr.err.Error()), true
|
||||
}
|
||||
|
||||
type nonEmptySummaryChatModel struct {
|
||||
base model.BaseChatModel
|
||||
}
|
||||
|
||||
func newNonEmptySummaryChatModel(base model.BaseChatModel) model.BaseChatModel {
|
||||
return &nonEmptySummaryChatModel{base: base}
|
||||
}
|
||||
|
||||
func (m *nonEmptySummaryChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
||||
out, err := m.base.Generate(ctx, input, opts...)
|
||||
if err != nil {
|
||||
return out, newEinoSummarizationModelError(err)
|
||||
}
|
||||
if strings.TrimSpace(classicAssistantTextContent(out)) == "" {
|
||||
return out, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", classicSummaryEmptyDiagnostics(out)))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *nonEmptySummaryChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
return m.base.Stream(ctx, input, opts...)
|
||||
}
|
||||
|
||||
type nonEmptyAgenticSummaryModel struct {
|
||||
base model.BaseModel[*schema.AgenticMessage]
|
||||
}
|
||||
|
||||
func newNonEmptyAgenticSummaryModel(base model.BaseModel[*schema.AgenticMessage]) model.BaseModel[*schema.AgenticMessage] {
|
||||
return &nonEmptyAgenticSummaryModel{base: base}
|
||||
}
|
||||
|
||||
func (m *nonEmptyAgenticSummaryModel) Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) {
|
||||
out, err := m.base.Generate(ctx, input, opts...)
|
||||
if err != nil {
|
||||
return out, newEinoSummarizationModelError(err)
|
||||
}
|
||||
if strings.TrimSpace(agenticAssistantTextContent(out)) == "" {
|
||||
return out, newEinoSummarizationModelError(fmt.Errorf("summary content is empty: %s", agenticSummaryEmptyDiagnostics(out)))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *nonEmptyAgenticSummaryModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) {
|
||||
return m.base.Stream(ctx, input, opts...)
|
||||
}
|
||||
|
||||
func classicAssistantTextContent(msg *schema.Message) string {
|
||||
if msg == nil || msg.Role != schema.Assistant {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, len(msg.AssistantGenMultiContent))
|
||||
for _, part := range msg.AssistantGenMultiContent {
|
||||
if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
|
||||
parts = append(parts, part.Text)
|
||||
}
|
||||
}
|
||||
if len(parts) > 0 {
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
return msg.Content
|
||||
}
|
||||
|
||||
func agenticAssistantTextContent(msg *schema.AgenticMessage) string {
|
||||
if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, len(msg.ContentBlocks))
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block != nil && block.AssistantGenText != nil {
|
||||
parts = append(parts, block.AssistantGenText.Text)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func classicSummaryEmptyDiagnostics(msg *schema.Message) string {
|
||||
if msg == nil {
|
||||
return "model returned nil message"
|
||||
}
|
||||
var textParts, reasoningParts, otherParts int
|
||||
var multiTextRunes, multiReasoningRunes int
|
||||
for _, part := range msg.AssistantGenMultiContent {
|
||||
switch part.Type {
|
||||
case schema.ChatMessagePartTypeText:
|
||||
textParts++
|
||||
multiTextRunes += len([]rune(strings.TrimSpace(part.Text)))
|
||||
case schema.ChatMessagePartTypeReasoning:
|
||||
reasoningParts++
|
||||
if part.Reasoning != nil {
|
||||
multiReasoningRunes += len([]rune(strings.TrimSpace(part.Reasoning.Text)))
|
||||
}
|
||||
default:
|
||||
otherParts++
|
||||
}
|
||||
}
|
||||
reasoningRunes := len([]rune(strings.TrimSpace(msg.ReasoningContent))) + multiReasoningRunes
|
||||
fields := []string{
|
||||
fmt.Sprintf("role=%s", msg.Role),
|
||||
fmt.Sprintf("content_runes=%d", len([]rune(strings.TrimSpace(msg.Content)))+multiTextRunes),
|
||||
fmt.Sprintf("reasoning_runes=%d", reasoningRunes),
|
||||
fmt.Sprintf("text_parts=%d", textParts),
|
||||
fmt.Sprintf("reasoning_parts=%d", reasoningParts),
|
||||
fmt.Sprintf("other_parts=%d", otherParts),
|
||||
fmt.Sprintf("tool_calls=%d", len(msg.ToolCalls)),
|
||||
}
|
||||
if msg.ResponseMeta != nil {
|
||||
fields = append(fields, fmt.Sprintf("finish_reason=%q", msg.ResponseMeta.FinishReason))
|
||||
if usage := msg.ResponseMeta.Usage; usage != nil {
|
||||
fields = append(fields,
|
||||
fmt.Sprintf("prompt_tokens=%d", usage.PromptTokens),
|
||||
fmt.Sprintf("completion_tokens=%d", usage.CompletionTokens),
|
||||
fmt.Sprintf("total_tokens=%d", usage.TotalTokens),
|
||||
fmt.Sprintf("reasoning_tokens=%d", usage.CompletionTokensDetails.ReasoningTokens),
|
||||
)
|
||||
}
|
||||
}
|
||||
if reasoningRunes > 0 {
|
||||
fields = append(fields, "hint=模型返回了 reasoning_content 但没有返回可作为摘要正文的 content;请检查 DeepSeek thinking 是否已在摘要请求中关闭")
|
||||
}
|
||||
return strings.Join(fields, " ")
|
||||
}
|
||||
|
||||
func agenticSummaryEmptyDiagnostics(msg *schema.AgenticMessage) string {
|
||||
if msg == nil {
|
||||
return "model returned nil agentic message"
|
||||
}
|
||||
var textBlocks, reasoningBlocks, otherBlocks int
|
||||
var textRunes, reasoningRunes int
|
||||
for _, block := range msg.ContentBlocks {
|
||||
if block == nil {
|
||||
continue
|
||||
}
|
||||
switch {
|
||||
case block.AssistantGenText != nil:
|
||||
textBlocks++
|
||||
textRunes += len([]rune(strings.TrimSpace(block.AssistantGenText.Text)))
|
||||
case block.Reasoning != nil:
|
||||
reasoningBlocks++
|
||||
reasoningRunes += len([]rune(strings.TrimSpace(block.Reasoning.Text)))
|
||||
default:
|
||||
otherBlocks++
|
||||
}
|
||||
}
|
||||
fields := []string{
|
||||
fmt.Sprintf("role=%s", msg.Role),
|
||||
fmt.Sprintf("content_runes=%d", textRunes),
|
||||
fmt.Sprintf("reasoning_runes=%d", reasoningRunes),
|
||||
fmt.Sprintf("text_blocks=%d", textBlocks),
|
||||
fmt.Sprintf("reasoning_blocks=%d", reasoningBlocks),
|
||||
fmt.Sprintf("other_blocks=%d", otherBlocks),
|
||||
}
|
||||
if msg.ResponseMeta != nil && msg.ResponseMeta.TokenUsage != nil {
|
||||
usage := msg.ResponseMeta.TokenUsage
|
||||
fields = append(fields,
|
||||
fmt.Sprintf("prompt_tokens=%d", usage.PromptTokens),
|
||||
fmt.Sprintf("completion_tokens=%d", usage.CompletionTokens),
|
||||
fmt.Sprintf("total_tokens=%d", usage.TotalTokens),
|
||||
fmt.Sprintf("reasoning_tokens=%d", usage.CompletionTokensDetails.ReasoningTokens),
|
||||
)
|
||||
}
|
||||
if reasoningRunes > 0 {
|
||||
fields = append(fields, "hint=模型返回了 reasoning block 但没有返回可作为摘要正文的 text block;请检查 DeepSeek thinking 是否已在摘要请求中关闭")
|
||||
}
|
||||
return strings.Join(fields, " ")
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
type guardClassicSummaryModel struct {
|
||||
out *schema.Message
|
||||
}
|
||||
|
||||
func (m *guardClassicSummaryModel) Generate(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) {
|
||||
return m.out, nil
|
||||
}
|
||||
|
||||
func (m *guardClassicSummaryModel) Stream(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
return schema.StreamReaderFromArray([]*schema.Message{m.out}), nil
|
||||
}
|
||||
|
||||
func TestNonEmptySummaryChatModelReportsEmptyContentDiagnostics(t *testing.T) {
|
||||
msg := schema.AssistantMessage("", nil)
|
||||
msg.ReasoningContent = "只返回了思考,没有最终摘要"
|
||||
msg.ResponseMeta = &schema.ResponseMeta{
|
||||
FinishReason: "stop",
|
||||
Usage: &schema.TokenUsage{
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 3,
|
||||
TotalTokens: 13,
|
||||
CompletionTokensDetails: schema.CompletionTokensDetails{
|
||||
ReasoningTokens: 3,
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := newNonEmptySummaryChatModel(&guardClassicSummaryModel{out: msg}).Generate(context.Background(), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected empty summary error")
|
||||
}
|
||||
text := err.Error()
|
||||
for _, want := range []string{
|
||||
"summary content is empty",
|
||||
"reasoning_runes=",
|
||||
`finish_reason="stop"`,
|
||||
"reasoning_tokens=3",
|
||||
"DeepSeek thinking",
|
||||
} {
|
||||
if !strings.Contains(text, want) {
|
||||
t.Fatalf("error missing %q:\n%s", want, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type guardAgenticSummaryModel struct {
|
||||
out *schema.AgenticMessage
|
||||
}
|
||||
|
||||
func (m *guardAgenticSummaryModel) Generate(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.AgenticMessage, error) {
|
||||
return m.out, nil
|
||||
}
|
||||
|
||||
func (m *guardAgenticSummaryModel) Stream(context.Context, []*schema.AgenticMessage, ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) {
|
||||
return schema.StreamReaderFromArray([]*schema.AgenticMessage{m.out}), nil
|
||||
}
|
||||
|
||||
func TestNonEmptyAgenticSummaryModelReportsEmptyContentDiagnostics(t *testing.T) {
|
||||
msg := &schema.AgenticMessage{
|
||||
Role: schema.AgenticRoleTypeAssistant,
|
||||
ContentBlocks: []*schema.ContentBlock{
|
||||
schema.NewContentBlock(&schema.Reasoning{Text: "只返回了思考,没有最终摘要"}),
|
||||
},
|
||||
ResponseMeta: &schema.AgenticResponseMeta{
|
||||
TokenUsage: &schema.TokenUsage{
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 3,
|
||||
TotalTokens: 13,
|
||||
CompletionTokensDetails: schema.CompletionTokensDetails{
|
||||
ReasoningTokens: 3,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
_, err := newNonEmptyAgenticSummaryModel(&guardAgenticSummaryModel{out: msg}).Generate(context.Background(), nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected empty summary error")
|
||||
}
|
||||
text := err.Error()
|
||||
for _, want := range []string{
|
||||
"summary content is empty",
|
||||
"reasoning_runes=",
|
||||
"reasoning_blocks=1",
|
||||
"reasoning_tokens=3",
|
||||
"DeepSeek thinking",
|
||||
} {
|
||||
if !strings.Contains(text, want) {
|
||||
t.Fatalf("error missing %q:\n%s", want, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,12 +1,36 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
copenai "cyberstrike-ai/internal/openai"
|
||||
)
|
||||
|
||||
// stripReasoningFromSummarizationPayload removes thinking / reasoning fields from a
|
||||
// chat-completions JSON body. Applied only to summarization Generate calls via
|
||||
// model.ModelOptions on the shared ChatModel — main-agent requests are unchanged.
|
||||
func stripReasoningFromSummarizationPayload(rawBody []byte) ([]byte, error) {
|
||||
func stripReasoningFromSummarizationPayload(rawBody []byte, oa *config.OpenAIConfig) ([]byte, error) {
|
||||
if shouldDisableDeepSeekThinkingForSummarization(oa) {
|
||||
return copenai.DisableThinkingForChatCompletionBody(rawBody)
|
||||
}
|
||||
return copenai.StripReasoningFromChatCompletionBody(rawBody)
|
||||
}
|
||||
|
||||
func shouldDisableDeepSeekThinkingForSummarization(oa *config.OpenAIConfig) bool {
|
||||
if oa == nil {
|
||||
return false
|
||||
}
|
||||
if oa.IsDeepSeekEndpointOrModel() {
|
||||
return true
|
||||
}
|
||||
profile := strings.ToLower(strings.TrimSpace(oa.Reasoning.ProfileEffective()))
|
||||
switch profile {
|
||||
case "deepseek", "deepseek_compat":
|
||||
return true
|
||||
case "", "auto":
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,11 +3,15 @@ package multiagent
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
)
|
||||
|
||||
func TestStripReasoningFromSummarizationPayload(t *testing.T) {
|
||||
in := []byte(`{"model":"deepseek-chat","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`)
|
||||
out, err := stripReasoningFromSummarizationPayload(in)
|
||||
out, err := stripReasoningFromSummarizationPayload(in, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -20,7 +24,7 @@ func TestStripReasoningFromSummarizationPayload(t *testing.T) {
|
||||
}
|
||||
|
||||
plain := []byte(`{"model":"gpt-4o","messages":[]}`)
|
||||
out2, err := stripReasoningFromSummarizationPayload(plain)
|
||||
out2, err := stripReasoningFromSummarizationPayload(plain, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -28,3 +32,75 @@ func TestStripReasoningFromSummarizationPayload(t *testing.T) {
|
||||
t.Fatalf("expected unchanged payload, got %s", out2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripReasoningFromSummarizationPayloadDisablesDeepSeekThinking(t *testing.T) {
|
||||
in := []byte(`{"model":"deepseek-v4-flash","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`)
|
||||
oa := &config.OpenAIConfig{
|
||||
BaseURL: "https://api.deepseek.com/v1",
|
||||
Model: "deepseek-v4-flash",
|
||||
}
|
||||
out, err := stripReasoningFromSummarizationPayload(in, oa)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := string(out)
|
||||
if strings.Contains(s, "reasoning_effort") {
|
||||
t.Fatalf("expected reasoning_effort stripped, got %s", s)
|
||||
}
|
||||
if !strings.Contains(s, `"thinking":{"type":"disabled"}`) {
|
||||
t.Fatalf("expected DeepSeek thinking disabled, got %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripReasoningFromSummarizationPayloadDisablesDeepSeekEndpointEvenWithOpenAICompatProfile(t *testing.T) {
|
||||
in := []byte(`{"model":"deepseek-v4-flash","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`)
|
||||
oa := &config.OpenAIConfig{
|
||||
BaseURL: "https://api.deepseek.com/v1",
|
||||
Model: "deepseek-v4-flash",
|
||||
Reasoning: config.OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
},
|
||||
}
|
||||
out, err := stripReasoningFromSummarizationPayload(in, oa)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := string(out)
|
||||
if strings.Contains(s, "reasoning_effort") {
|
||||
t.Fatalf("expected reasoning_effort stripped, got %s", s)
|
||||
}
|
||||
if !strings.Contains(s, `"thinking":{"type":"disabled"}`) {
|
||||
t.Fatalf("expected official DeepSeek endpoint thinking disabled, got %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripReasoningFromSummarizationPayloadHonorsOpenAICompatProfileForNonDeepSeekEndpoint(t *testing.T) {
|
||||
in := []byte(`{"model":"deepseek-v4-flash","messages":[],"thinking":{"type":"enabled"},"reasoning_effort":"high"}`)
|
||||
oa := &config.OpenAIConfig{
|
||||
BaseURL: "https://compatible.example.com/v1",
|
||||
Model: "deepseek-v4-flash",
|
||||
Reasoning: config.OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
},
|
||||
}
|
||||
out, err := stripReasoningFromSummarizationPayload(in, oa)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := string(out)
|
||||
if strings.Contains(s, "thinking") || strings.Contains(s, "reasoning_effort") {
|
||||
t.Fatalf("expected non-DeepSeek OpenAI-compatible endpoint to strip reasoning fields, got %s", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoSummarizationModelOptionsSetCommonMaxTokens(t *testing.T) {
|
||||
const outputReserve = 4096
|
||||
opts := newEinoSummarizationModelOptions(outputReserve, "minimax-m3", "agentic", nil, nil)
|
||||
common := model.GetCommonOptions(nil, opts...)
|
||||
if common == nil || common.MaxTokens == nil {
|
||||
t.Fatal("expected summarization options to set common max_tokens")
|
||||
}
|
||||
if *common.MaxTokens != outputReserve {
|
||||
t.Fatalf("max_tokens = %d, want %d", *common.MaxTokens, outputReserve)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/middlewares/summarization"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -385,6 +386,84 @@ func TestSummarizeFinalize_MergesSystemMessages(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoSummarizationMiddlewareRetriesWhenSummaryModelReturnsEmpty(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
emit := false
|
||||
summaryModel := &capturingClassicChatModel{outputs: []*schema.Message{
|
||||
schema.AssistantMessage("", nil),
|
||||
schema.AssistantMessage("<summary>有效摘要:继续验证 SQL 注入路径</summary>", nil),
|
||||
}}
|
||||
appCfg := &config.Config{}
|
||||
appCfg.OpenAI.Model = "gpt-4o"
|
||||
appCfg.OpenAI.MaxTotalTokens = 5000
|
||||
appCfg.Database.Path = filepath.Join(t.TempDir(), "cyberstrike.db")
|
||||
mwCfg := &config.MultiAgentEinoMiddlewareConfig{
|
||||
SummarizationEmitInternalEvents: &emit,
|
||||
SummarizationOutputReserveTokens: 1024,
|
||||
}
|
||||
|
||||
mw, err := newEinoSummarizationMiddleware(ctx, summaryModel, appCfg, mwCfg, "conv-empty-summary", nil, "", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("newEinoSummarizationMiddleware: %v", err)
|
||||
}
|
||||
state := &adk.ChatModelAgentState{Messages: []adk.Message{
|
||||
schema.SystemMessage("system root"),
|
||||
schema.UserMessage("授权范围 example.com\n" + strings.Repeat("历史扫描输出 ", 12000)),
|
||||
schema.AssistantMessage("已记录范围", nil),
|
||||
schema.UserMessage("继续验证 SQL 注入路径"),
|
||||
}}
|
||||
|
||||
_, after, err := mw.BeforeModelRewriteState(ctx, state, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("BeforeModelRewriteState should retry instead of failing on empty summary: %v", err)
|
||||
}
|
||||
if after == nil {
|
||||
t.Fatal("after state is nil")
|
||||
}
|
||||
if summaryModel.calls < 2 {
|
||||
t.Fatalf("summary model calls=%d, want retry after empty output", summaryModel.calls)
|
||||
}
|
||||
joined := joinClassicMessageContent(after.Messages)
|
||||
for _, want := range []string{"有效摘要", "继续验证 SQL 注入路径", "原始用户输入与约束账本"} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Fatalf("retried compacted context missing %q:\n%s", want, joined)
|
||||
}
|
||||
}
|
||||
if strings.Contains(joined, "本地压缩摘要") {
|
||||
t.Fatalf("local fallback should not be used:\n%s", joined)
|
||||
}
|
||||
}
|
||||
|
||||
type capturingClassicChatModel struct {
|
||||
output *schema.Message
|
||||
outputs []*schema.Message
|
||||
calls int
|
||||
}
|
||||
|
||||
func (m *capturingClassicChatModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
||||
m.calls++
|
||||
if len(m.outputs) > 0 {
|
||||
idx := m.calls - 1
|
||||
if idx >= len(m.outputs) {
|
||||
idx = len(m.outputs) - 1
|
||||
}
|
||||
return m.outputs[idx], nil
|
||||
}
|
||||
if m.output != nil {
|
||||
return m.output, nil
|
||||
}
|
||||
return schema.AssistantMessage("classic answer", nil), nil
|
||||
}
|
||||
|
||||
func (m *capturingClassicChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
||||
msg, err := m.Generate(ctx, input, opts...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return schema.StreamReaderFromArray([]*schema.Message{msg}), nil
|
||||
}
|
||||
|
||||
// assertNoOrphanTool 断言消息列表里的每个 role=tool 消息都能在更前面找到一个
|
||||
// assistant(tool_calls) 提供相同 ID,否则说明产生了孤儿(触发 LLM 400 的根因)。
|
||||
func assertNoOrphanTool(t *testing.T, msgs []adk.Message) {
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/toolguard"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestEinoToolResultPreservesBlockedOutcomeAfterTextReduction(t *testing.T) {
|
||||
for _, blocked := range []bool{true, false} {
|
||||
name := "execution_error"
|
||||
if blocked {
|
||||
name = "safety_block"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
logger := zap.NewNop()
|
||||
server := mcp.NewServer(logger)
|
||||
guard, err := toolguard.NewManager(toolguard.DefaultConfig())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
server.SetToolGuard(guard)
|
||||
calls := 0
|
||||
server.RegisterTool(mcp.Tool{Name: "inspect"}, func(context.Context, map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
calls++
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "execution failed"}}, IsError: true}, nil
|
||||
})
|
||||
ag := agent.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 1)
|
||||
target := "example.org"
|
||||
if blocked {
|
||||
target = "example.gov"
|
||||
}
|
||||
result, err := ag.ExecuteMCPToolForConversation(ctx, "conv-block", "inspect", map[string]interface{}{"target": target})
|
||||
if err != nil || result == nil || !result.IsError || result.Blocked != blocked {
|
||||
t.Fatalf("agent result = %#v, error = %v", result, err)
|
||||
}
|
||||
if blocked && calls != 0 || !blocked && calls != 1 {
|
||||
t.Fatalf("handler calls = %d, blocked = %v", calls, blocked)
|
||||
}
|
||||
binder := NewMCPExecutionBinder()
|
||||
binder.Bind("call-block", result.ExecutionID)
|
||||
var event map[string]interface{}
|
||||
emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{
|
||||
FilesystemMonitorAgent: ag,
|
||||
MCPExecutionBinder: binder,
|
||||
Progress: func(eventType, _ string, data interface{}) {
|
||||
if eventType == "tool_result" {
|
||||
event = data.(map[string]interface{})
|
||||
}
|
||||
},
|
||||
})
|
||||
// The reduced text deliberately contains no refusal wording. A blocked
|
||||
// outcome must survive even if reduction also loses the error prefix.
|
||||
const reduced = "The request did not run."
|
||||
if !emitter.Emit(ctx, "inspect", reduced, "call-block", !blocked, "worker") {
|
||||
t.Fatal("missing tool result event")
|
||||
}
|
||||
if event["success"] != false || event["isError"] != true || event["result"] != reduced {
|
||||
t.Fatalf("event = %#v", event)
|
||||
}
|
||||
if blocked && (event["blocked"] != true || event["status"] != "blocked" || event["executionId"] != result.ExecutionID) {
|
||||
t.Fatalf("blocked event lost its classification: %#v", event)
|
||||
}
|
||||
if !blocked && event["blocked"] != nil {
|
||||
t.Fatalf("ordinary failure was classified as blocked: %#v", event)
|
||||
}
|
||||
exec, ok := server.GetExecution(result.ExecutionID)
|
||||
if !ok || exec.Result == nil || !exec.Result.IsError || exec.Result.Blocked != blocked {
|
||||
t.Fatalf("display update lost result flags: %#v", exec)
|
||||
}
|
||||
wantStatus := "failed"
|
||||
if blocked {
|
||||
wantStatus = "blocked"
|
||||
}
|
||||
if ag.MCPExecutionStatus(result.ExecutionID) != wantStatus {
|
||||
t.Fatalf("execution status = %q, want %q", exec.Status, wantStatus)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
// concatToolResultChunks 按 Eino 原生语义合并工具结果流:
|
||||
// - 同一 CallID(EventSender 一 call 一 event):schema.ConcatMessages
|
||||
// - 并行工具被摊进同一条流(ToolsNode MergeStreamReaders 扁平化后):
|
||||
// 按 CallID 分列后再 ConcatMessages,等价于 schema.ConcatMessageArray
|
||||
func concatToolResultChunks(chunks []*schema.Message) ([]*schema.Message, error) {
|
||||
if len(chunks) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if toolResultChunksShareCallID(chunks) {
|
||||
merged, err := schema.ConcatMessages(chunks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []*schema.Message{merged}, nil
|
||||
}
|
||||
return concatToolResultChunksByCallID(chunks)
|
||||
}
|
||||
|
||||
func toolResultChunksShareCallID(chunks []*schema.Message) bool {
|
||||
id := ""
|
||||
for _, chunk := range chunks {
|
||||
if chunk == nil {
|
||||
continue
|
||||
}
|
||||
got := strings.TrimSpace(chunk.ToolCallID)
|
||||
if got == "" {
|
||||
continue
|
||||
}
|
||||
if id == "" {
|
||||
id = got
|
||||
continue
|
||||
}
|
||||
if got != id {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func concatToolResultChunksByCallID(chunks []*schema.Message) ([]*schema.Message, error) {
|
||||
type column struct {
|
||||
key string
|
||||
chunks []*schema.Message
|
||||
}
|
||||
var ordered []column
|
||||
index := make(map[string]int)
|
||||
lastKey := ""
|
||||
anon := 0
|
||||
for _, chunk := range chunks {
|
||||
if chunk == nil {
|
||||
continue
|
||||
}
|
||||
key := strings.TrimSpace(chunk.ToolCallID)
|
||||
if key == "" {
|
||||
if lastKey != "" {
|
||||
key = lastKey
|
||||
} else {
|
||||
key = fmt.Sprintf("\x00anon-%d", anon)
|
||||
anon++
|
||||
}
|
||||
}
|
||||
if idx, ok := index[key]; ok {
|
||||
ordered[idx].chunks = append(ordered[idx].chunks, chunk)
|
||||
} else {
|
||||
index[key] = len(ordered)
|
||||
ordered = append(ordered, column{key: key, chunks: []*schema.Message{chunk}})
|
||||
}
|
||||
lastKey = key
|
||||
}
|
||||
out := make([]*schema.Message, 0, len(ordered))
|
||||
for _, col := range ordered {
|
||||
merged, err := schema.ConcatMessages(col.chunks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.HasPrefix(col.key, "\x00anon-") {
|
||||
merged.ToolCallID = ""
|
||||
}
|
||||
out = append(out, merged)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestConcatToolResultChunksUsesEinoConcatForSingleCall(t *testing.T) {
|
||||
got, err := concatToolResultChunks([]*schema.Message{
|
||||
schema.ToolMessage("hel", "call-1", schema.WithToolName("execute")),
|
||||
schema.ToolMessage("lo", "call-1", schema.WithToolName("execute")),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("concat: %v", err)
|
||||
}
|
||||
if len(got) != 1 || got[0].ToolCallID != "call-1" || got[0].Content != "hello" || got[0].ToolName != "execute" {
|
||||
t.Fatalf("got = %#v, want one ConcatMessages result", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcatToolResultChunksSplitsParallelCalls(t *testing.T) {
|
||||
got, err := concatToolResultChunks([]*schema.Message{
|
||||
schema.ToolMessage("nmap 1/2 ", "call-1", schema.WithToolName("nmap")),
|
||||
schema.ToolMessage("nmap 2/2 ", "call-2", schema.WithToolName("nmap")),
|
||||
schema.ToolMessage("22/tcp", "call-1", schema.WithToolName("nmap")),
|
||||
schema.ToolMessage("80/tcp", "call-2", schema.WithToolName("nmap")),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("concat: %v", err)
|
||||
}
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("got = %#v, want two calls", got)
|
||||
}
|
||||
if got[0].ToolCallID != "call-1" || got[0].Content != "nmap 1/2 22/tcp" {
|
||||
t.Fatalf("call-1 = %#v", got[0])
|
||||
}
|
||||
if got[1].ToolCallID != "call-2" || got[1].Content != "nmap 2/2 80/tcp" {
|
||||
t.Fatalf("call-2 = %#v", got[1])
|
||||
}
|
||||
}
|
||||
@@ -42,27 +42,42 @@ func (h *einoToolResultEventHandler) HandleStreaming(mv *adk.MessageVariant, age
|
||||
if h == nil || mv == nil || !mv.IsStreaming || mv.MessageStream == nil || mv.Role != schema.Tool {
|
||||
return false
|
||||
}
|
||||
toolName := strings.TrimSpace(mv.ToolName)
|
||||
content, streamToolCallID, streamToolName, recvErr := recvSchemaMessageStream(h.ctx, mv.MessageStream)
|
||||
if toolName == "" {
|
||||
toolName = streamToolName
|
||||
defaultName := strings.TrimSpace(mv.ToolName)
|
||||
msgs, recvErr := recvSchemaToolResultMessages(h.ctx, mv.MessageStream)
|
||||
if isEinoVoluntaryCancelErr(recvErr) && len(msgs) == 0 {
|
||||
msgs = []*schema.Message{schema.ToolMessage("已中断并继续,当前工具调用已停止。", "", schema.WithToolName(defaultName))}
|
||||
}
|
||||
if isEinoVoluntaryCancelErr(recvErr) && strings.TrimSpace(content) == "" {
|
||||
content = "已中断并继续,当前工具调用已停止。"
|
||||
if len(msgs) == 0 {
|
||||
msgs = []*schema.Message{schema.ToolMessage("", "", schema.WithToolName(defaultName))}
|
||||
}
|
||||
isErr := einoToolResultIsError(toolName, content) || isEinoVoluntaryCancelErr(recvErr)
|
||||
content = einoToolResultBody(content)
|
||||
if streamToolCallID != "" && h.runMessages != nil {
|
||||
h.runMessages.AppendToolMessage(content, streamToolCallID, schema.WithToolName(toolName))
|
||||
}
|
||||
if h.emitter != nil {
|
||||
h.emitter.Emit(h.ctx, toolName, content, streamToolCallID, isErr, agentName)
|
||||
}
|
||||
if recvErr != nil && !isEinoVoluntaryCancelErr(recvErr) && h.logger != nil {
|
||||
h.logger.Warn("eino tool result stream recv error",
|
||||
zap.Error(recvErr),
|
||||
zap.String("agent", agentName),
|
||||
zap.String("tool", toolName))
|
||||
for _, msg := range msgs {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
toolName := strings.TrimSpace(msg.ToolName)
|
||||
if toolName == "" {
|
||||
toolName = defaultName
|
||||
}
|
||||
content := msg.Content
|
||||
if isEinoVoluntaryCancelErr(recvErr) && strings.TrimSpace(content) == "" {
|
||||
content = "已中断并继续,当前工具调用已停止。"
|
||||
}
|
||||
isErr := einoToolResultIsError(toolName, content) || isEinoVoluntaryCancelErr(recvErr)
|
||||
content = einoToolResultBody(content)
|
||||
toolCallID := strings.TrimSpace(msg.ToolCallID)
|
||||
if toolCallID != "" && h.runMessages != nil {
|
||||
h.runMessages.AppendToolMessage(content, toolCallID, schema.WithToolName(toolName))
|
||||
}
|
||||
if h.emitter != nil {
|
||||
h.emitter.Emit(h.ctx, toolName, content, toolCallID, isErr, agentName)
|
||||
}
|
||||
if recvErr != nil && !isEinoVoluntaryCancelErr(recvErr) && h.logger != nil {
|
||||
h.logger.Warn("eino tool result stream recv error",
|
||||
zap.Error(recvErr),
|
||||
zap.String("agent", agentName),
|
||||
zap.String("tool", toolName),
|
||||
zap.String("toolCallId", toolCallID))
|
||||
}
|
||||
}
|
||||
if recvErr == nil && h.confirmRecovery != nil {
|
||||
h.confirmRecovery()
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user