Compare commits

..
57 Commits
Author SHA1 Message Date
公明 6edc70f3fc Update config.example.yaml 2026-09-08 16:10:29 +08:00
Ed1s0nZ 377be6492a fix: refresh navigation permissions before login and on scoped renders 2026-09-08 15:15:18 +08:00
Ed1s0nZ 94cdf760af feat: persist diagnostic logs with daily rotation and retention 2026-09-08 14:41:46 +08:00
Ed1s0nZ 6ad9ea2d13 feat: add configurable tool call blocking and monitoring 2026-09-08 09:44:32 +08:00
temp c70da22de7 Remove conversation grouping feature 2026-08-26 15:00:09 +08:00
temp fbbe984005 fix: refresh ai channel reasoning selects 2026-08-25 18:51:19 +08:00
temp abc00c1fef fix: auto-detect deepseek profile in settings 2026-08-25 18:40:49 +08:00
temp d6d48c97c0 fix: normalize deepseek config profile 2026-08-25 18:37:08 +08:00
temp d7581a6373 fix: surface summarization model errors 2026-08-25 18:28:47 +08:00
temp 4c011abb9d fix: bound summarization max tokens for claude models 2026-08-25 16:37:45 +08:00
temp e4441f91ad fix: stabilize eino summarization for deepseek 2026-08-25 14:54:02 +08:00
temp 21c6ad9bdf feat: persist hitl default config 2026-08-24 19:45:07 +08:00
公明 e0a2f01427 Update version number to v1.7.17 2026-08-24 18:46:58 +08:00
temp baff533196 fix: disable chat checkpoint resume 2026-08-24 18:42:06 +08:00
temp 474238cfc5 fix: distinguish model original errors 2026-08-24 14:53:02 +08:00
temp b47f8df3b0 fix: avoid inferred retry failure reason 2026-08-24 14:48:37 +08:00
temp a67761e843 fix: surface original Eino retry errors 2026-08-24 14:43:48 +08:00
temp b41596d51f Fix Eino final output fallback capture 2026-08-24 14:32:27 +08:00
temp d80e27e950 Respect OpenAI reasoning profile for DeepSeek-named models 2026-08-24 13:55:07 +08:00
temp e218316c55 Add token usage tracking and UI refinements 2026-08-24 00:05:05 +08:00
Codex a34cab431a Improve conversation preview layout 2026-08-23 20:09:56 +08:00
temp 3bcf4458c5 Fix finalization cleanup for pending tool executions 2026-08-23 19:57:17 +08:00
公明 bf761e9cd5 Update config.example.yaml 2026-08-19 17:36:35 +08:00
d640ef09c8 fix: 为 Eino agentic 路径补充 tool_call/tool_result 配对防御中间件 (#265) (#266)
agentic 路径(单代理模式)缺少 toolPairReconciler 和 orphanToolPruner,
当 summarization 截断历史破坏配对后,序列化到 OpenAI API 触发 400
"insufficient tool messages following tool_calls message"。

新增 AgenticMessage 版本的 reconciler 和 pruner,与 classic 路径对齐。

Co-authored-by: temp <temp@tempdeMacBook-Pro.local>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-19 17:29:25 +08:00
tempandCursor d88cfea761 fix: 为 Eino agentic 路径补充 tool_call/tool_result 配对防御中间件
agentic 路径(单代理模式)缺少 toolPairReconciler 和 orphanToolPruner,
当 summarization 截断历史破坏配对后,序列化到 OpenAI API 触发 400
"insufficient tool messages following tool_calls message"。

新增 AgenticMessage 版本的 reconciler 和 pruner,与 classic 路径对齐。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-19 17:24:03 +08:00
RuoJi6 bec2d2faf1 修复多会话任务栏抖动、会话跳回与停止失效 (#264)
* fix(web): stabilize active task ordering

* fix(chat): preserve navigation during conversation startup

* fix(tasks): guarantee hard cancellation

* fix(chat): bind navigation and hard stop targets

* fix(chat): prevent replay from reclaiming navigation
2026-08-19 13:54:30 +08:00
RuoJi6 c7cc0bc9da 修复人工审批长内容遮挡及刷新审批人状态异常 (#263)
* docs(hitl): remove stale asm_list_resources references

* fix(hitl): constrain approval layout and preserve reviewer
2026-08-19 13:31:37 +08:00
公明andCursor 24d06c5220 Fix missing results for parallel Eino tool calls.
Merge streaming tool outputs by CallID with ConcatMessages, pair same-name historical results, and FIFO-match duplicate IDs so concurrent nmap 1/2 and 2/2 stay distinct.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-19 00:29:12 +08:00
公明 ac6e04a94c Update config.example.yaml 2026-08-19 00:14:46 +08:00
公明 cb58b90083 Update .gitignore 2026-08-18 20:53:37 +08:00
公明 5ccabdc3e4 Add files via upload 2026-08-18 20:52:02 +08:00
公明 286511d0a8 Add files via upload 2026-08-18 20:48:50 +08:00
公明 1f66b93866 Add files via upload 2026-08-18 20:41:19 +08:00
公明 c21ba6c419 Add files via upload 2026-08-18 20:38:40 +08:00
公明 5720de0d27 Add files via upload 2026-08-18 20:36:34 +08:00
公明 5aa2c382b0 Add files via upload 2026-08-18 20:34:53 +08:00
公明 98c380ed5f Add files via upload 2026-08-18 20:31:48 +08:00
公明 11537713ce Add files via upload 2026-08-18 20:30:18 +08:00
公明 acc2ebd7e2 Add files via upload 2026-08-18 20:29:17 +08:00
公明 5eaddd9e8c Add files via upload 2026-08-18 20:25:43 +08:00
公明 3421d8b047 Add files via upload 2026-08-18 20:23:28 +08:00
公明 5b000d1e3e Add files via upload 2026-08-18 20:22:40 +08:00
公明 cecbe5a086 Add files via upload 2026-08-18 20:20:39 +08:00
公明 c7603c2b66 Add files via upload 2026-08-18 20:17:38 +08:00
公明 d3bfb1c98e Add files via upload 2026-08-18 20:15:30 +08:00
公明 55ee5fbac9 Delete internal directory 2026-08-18 20:14:29 +08:00
RuoJi6 910d07ea0a docs(hitl): remove stale asm_list_resources references (#259) 2026-08-17 17:19:16 +08:00
公明 217f5570c2 Update config.example.yaml 2026-08-17 15:55:03 +08:00
RuoJi6 f30a9de01d feat(hitl): expand safe read-only tool allowlist (#258) 2026-08-17 15:50:32 +08:00
公明 893f1342ae Add files via upload 2026-08-15 11:11:40 +08:00
公明 6db0975223 Add files via upload 2026-08-15 11:09:28 +08:00
公明 b935e1f90f Add files via upload 2026-08-15 11:06:53 +08:00
公明 3b7fd8d582 Add files via upload 2026-08-15 11:06:08 +08:00
公明 053c75b2cc Add files via upload 2026-08-15 11:04:35 +08:00
公明 7823fb6b0c Add files via upload 2026-08-15 11:03:13 +08:00
公明 9c0621819f Add files via upload 2026-08-15 10:13:12 +08:00
公明 aa11b32ce5 Add files via upload 2026-08-15 09:45:37 +08:00
192 changed files with 13733 additions and 8186 deletions
+25 -1
View File
@@ -1,30 +1,49 @@
# Runtime data
data/
chat_uploads/
*.db
*.db-shm
*.db-wal
*.db-journal
*.sqlite
*.sqlite3
# Local configuration and secrets
config.yaml
config.local.yaml
config.yaml.backup
*.backup
.env
.env.*
!.env.example
*.pem
*.key
*.crt
# Upgrade / settings backups
.upgrade-backup/
# Eino runtime task board (under skills_dir)
skills/.eino/
# Build outputs
cyberstrike-ai
cyberstrike-ai.exe
*.test
bin/
dist/
build/
.build/
target/
.gradle/
coverage.out
coverage.html
# Logs and temporary files
/log/
*.log
*.bak
*~
tmp/
temp/
@@ -38,7 +57,12 @@ __pycache__/
# Go
vendor/
# macOS / editors
# Node / browser extension
node_modules/
# macOS / Windows / editors
.DS_Store
Thumbs.db
desktop.ini
.idea/
.vscode/
+2 -1
View File
@@ -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.
+1
View File
@@ -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
View File
@@ -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
View File
@@ -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())
+27 -7
View File
@@ -10,7 +10,7 @@
# ============================================
# 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.7.14"
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
@@ -57,7 +60,7 @@ monitor:
# - DeepSeek: https://api.deepseek.com/v1
# - 其他兼容 OpenAI 协议的 API
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
# provider: 可选值 openai_compatible(默认) | claude(自动桥接到 Anthropic Claude Messages API)
# provider: 可选值 openai_compatible(默认) | claude(Eino 原生 Anthropic Messages API)
ai:
default_channel: qwen-max
channels:
@@ -107,10 +110,10 @@ fofa:
base_url: https://fofa.info/api/v1/search/all # 可选,留空则使用默认
api_key: "" # FOFA API Key(可选,建议在系统设置中填写)
zoomeye:
base_url: https://api.zoomeye.org/v2/search # 可选,留空则使用默认
base_url: https://api.zoomeye.ai/v2/search # 可选,留空则使用默认
api_key: "" # ZoomEye API Key;也可通过 ZOOMEYE_API_KEY 环境变量配置
quake:
base_url: https://quake.360.cn/api/v3/search/quake_service # 可选,留空则使用默认
base_url: https://quake.360.net/api/v3/search/quake_service # 可选,留空则使用默认
api_key: "" # Quake API Token;也可通过 QUAKE_API_KEY 环境变量配置
shodan:
base_url: https://api.shodan.io # 可选,留空则使用默认
@@ -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
@@ -146,7 +166,7 @@ hitl:
# 已决策审计日志保留天数(与 MCP 监控一致;省略默认 90;0 表示不自动清理)
retention_days: 90
# 按你环境里的真实工具名增删(与侧栏一致、小写不敏感);不需要全局免审批可改为 []
tool_whitelist: [read_file, list_dir, glob, grep, tool_search, upsert_project_fact, get_project_fact]
tool_whitelist: [read_file, ls, list_dir, glob, grep, tool_search, upsert_project_fact, get_project_fact, list_project_facts, search_project_facts, list_vulnerabilities, get_vulnerability, get_asset, query_assets, list_knowledge_risk_types, get_tool_execution, wait_tool_execution, batch_task_list, batch_task_get, manage_webshell_list, c2_event, c2_file]
# audit_agent_prompt: | # 审批模式;留空使用内置默认,可在「人机协同」页编辑
# audit_agent_prompt_review_edit: | # 审查编辑模式;留空使用内置默认
@@ -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 retry408/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 耗尽后按顺序切换
+7 -1
View File
@@ -68,7 +68,7 @@ ai:
| `ai.default_channel` | Default channel ID for new conversations and requests without an explicit channel. |
| `ai.channels.<id>` | Channel config. IDs are normalized to lowercase letters, digits, and hyphens. |
| `name` | Display name in the Web UI; falls back to the ID. |
| `provider` | `openai_compatible` or `claude`. OpenAI-compatible channels map to runtime `openai`; Claude channels bridge to Anthropic Messages API. |
| `provider` | `openai_compatible` or `claude`. OpenAI-compatible channels map to runtime `openai`; Claude channels use Eino's native Anthropic Messages API component. |
| `base_url/api_key/model` | Required. Base URL usually includes a version path such as `/v1`. |
| `max_total_tokens` | Shared context budget for compression, attack-chain generation, multi-agent summaries, and similar paths. |
| `max_completion_tokens` | Per-response output cap; default is used when empty. |
@@ -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 servers 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.
+1 -1
View File
@@ -10,7 +10,7 @@ CyberStrikeAI can run as a local testing tool, an internal team service, or a pr
- Python for some MCP servers and tool scripts.
- SQLite files under `data/`; no external DB is required by default.
- Actual security tools installed in PATH. YAML files under `tools/` only describe commands.
- At least one `ai.channels` entry. Use `provider: openai_compatible` for OpenAI-compatible endpoints, or `provider: claude` for the Claude bridge.
- At least one `ai.channels` entry. Use `provider: openai_compatible` for OpenAI-compatible endpoints, or `provider: claude` for Eino's native Claude component.
Important persistent paths:
+10 -2
View File
@@ -25,7 +25,7 @@ hitl:
api_key: ""
model: "" # set a small model here; blank reuses the default AI channel model
retention_days: 90
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
tool_whitelist: [read_file, ls, glob, grep, tool_search, get_project_fact, list_project_facts, search_project_facts, list_vulnerabilities, get_vulnerability, get_asset, query_assets, list_knowledge_risk_types, get_tool_execution, wait_tool_execution, batch_task_list, batch_task_get, manage_webshell_list, c2_event, c2_file]
```
`audit_model` supports partial configuration. Empty fields inherit from the resolved default AI channel, so the common setup is to fill only `model` and run approvals on a cheaper small model.
@@ -91,10 +91,16 @@ In review-edit mode, you may narrow paths, targets, or command arguments before
Allowlisted tools skip approval, so keep the list stable and low-risk. Recommended examples:
- `read_file`
- `list_dir`
- `ls`
- `glob`
- `grep`
- `tool_search`
- Project and vulnerability reads: `get_project_fact`, `list_project_facts`, `search_project_facts`, `list_vulnerabilities`, `get_vulnerability`
- Asset and knowledge-metadata reads: `get_asset`, `query_assets`, `list_knowledge_risk_types`
- Execution and task-state reads: `get_tool_execution`, `wait_tool_execution`, `batch_task_list`, `batch_task_get`
- Local management-metadata reads: `manage_webshell_list`, `c2_event`, `c2_file`
These built-in MCP reads remain constrained by RBAC and resource scope; the allowlist only bypasses HITL approval and does not grant additional access. `list_dir` is effective only when that is the actual tool name; the Eino filesystem directory-listing tool is named `ls`.
Avoid globally allowlisting:
@@ -102,6 +108,8 @@ Avoid globally allowlisting:
- File write/delete tools
- C2 task tools
- WebShell command execution tools
- Nominally read-only tools that send requests to a target or external service, such as `webshell_file_read`, `webshell_file_list`, and `search_knowledge_base`
- Multiplexed tools whose actions include both reads and writes, such as `c2_session`, `c2_listener`, `c2_profile`, and `c2_task_manage`
## Mode Selection
-1
View File
@@ -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:
+27
View File
@@ -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`.
+6 -2
View File
@@ -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 通道与模型配置
@@ -65,7 +69,7 @@ ai:
| `ai.default_channel` | 默认通道 ID。新对话、机器人、批量任务和未显式选择通道的请求使用它。 |
| `ai.channels.<id>` | 通道配置。ID 会归一化为小写、数字和短横线,例如 `Qwen_Max` 会变成 `qwen-max`。 |
| `name` | Web UI 展示名。留空时使用通道 ID。 |
| `provider` | `openai_compatible``claude``openai_compatible` 会在运行时映射为 `openai``claude` 会桥接到 Anthropic Messages API。 |
| `provider` | `openai_compatible``claude``openai_compatible` 会在运行时映射为 `openai``claude` 使用 Eino 原生 Anthropic Messages API。 |
| `base_url/api_key/model` | 必填。Base URL 通常需要包含版本路径,如 OpenAI/兼容网关的 `/v1`。 |
| `max_total_tokens` | 上下文压缩、攻击链构建、多代理摘要等共用的总预算。 |
| `max_completion_tokens` | 单次模型输出上限;未填时使用默认值。 |
+1 -1
View File
@@ -10,7 +10,7 @@
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
- SQLite:默认使用文件型数据库,无需单独服务。
- 安全工具:`tools/` 中的 YAML 只是工具定义,实际命令如 `nmap``sqlmap``nuclei` 仍需安装到系统 PATH。
- 模型服务:至少配置一个 `ai.channels` 通道;`provider: openai_compatible` 适用于 OpenAI 兼容 API`provider: claude` 会走 Claude 桥接
- 模型服务:至少配置一个 `ai.channels` 通道;`provider: openai_compatible` 适用于 OpenAI 兼容 API`provider: claude` 使用 Eino 原生 Claude 组件
建议目录:
+10 -2
View File
@@ -25,7 +25,7 @@ hitl:
api_key: ""
model: "" # 可填小模型;留空复用默认 AI 通道的模型
retention_days: 90
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
tool_whitelist: [read_file, ls, glob, grep, tool_search, get_project_fact, list_project_facts, search_project_facts, list_vulnerabilities, get_vulnerability, get_asset, query_assets, list_knowledge_risk_types, get_tool_execution, wait_tool_execution, batch_task_list, batch_task_get, manage_webshell_list, c2_event, c2_file]
```
`audit_model` 的字段可以只填一部分。空字段会自动继承默认 AI 通道解析后的模型配置,因此常见做法是只填 `model`,让审计 Agent 使用更便宜的小模型。
@@ -91,10 +91,16 @@ hitl:
白名单工具会跳过审批,因此要保守维护。推荐放:
- `read_file`
- `list_dir`
- `ls`
- `glob`
- `grep`
- `tool_search`
- 项目与漏洞查询:`get_project_fact``list_project_facts``search_project_facts``list_vulnerabilities``get_vulnerability`
- 资产与知识元数据查询:`get_asset``query_assets``list_knowledge_risk_types`
- 执行与任务状态查询:`get_tool_execution``wait_tool_execution``batch_task_list``batch_task_get`
- 本地管理元数据查询:`manage_webshell_list``c2_event``c2_file`
上述内置 MCP 查询仍受 RBAC 和资源范围约束;白名单只跳过 HITL 审批,不扩大访问权限。`list_dir` 仅在实际工具名为该值时有效;Eino 文件系统的目录列表工具名是 `ls`
不建议直接全局白名单:
@@ -102,6 +108,8 @@ hitl:
- 文件写入/删除工具
- C2 任务工具
- WebShell 命令执行工具
- 会向目标或外部服务发起请求的“只读”工具,例如 `webshell_file_read``webshell_file_list``search_knowledge_base`
- 同一工具名同时包含读写 action 的复合工具,例如 `c2_session``c2_listener``c2_profile``c2_task_manage`
## 模式选择
-1
View File
@@ -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` |
特殊权限说明:
+38
View File
@@ -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`
+37 -3
View File
@@ -14,6 +14,7 @@ require (
github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown v0.0.0-20260427010451-749e3706378b
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260427010451-749e3706378b
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b
github.com/cloudwego/eino-ext/components/model/agenticclaude v0.1.3
github.com/cloudwego/eino-ext/components/model/agenticopenai v0.2.1
github.com/cloudwego/eino-ext/components/model/openai v0.1.13
github.com/creack/pty v1.1.24
@@ -24,7 +25,7 @@ require (
github.com/gorilla/websocket v1.5.3
github.com/larksuite/oapi-sdk-go/v3 v3.4.22
github.com/mattn/go-sqlite3 v1.14.18
github.com/modelcontextprotocol/go-sdk v1.2.0
github.com/modelcontextprotocol/go-sdk v1.3.1
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
github.com/pkoukk/tiktoken-go v0.1.8
github.com/robfig/cron/v3 v3.0.1
@@ -46,11 +47,29 @@ require (
)
require (
cloud.google.com/go/auth v0.7.2 // indirect
cloud.google.com/go/auth/oauth2adapt v0.2.3 // indirect
cloud.google.com/go/compute/metadata v0.9.0 // indirect
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0 // indirect
github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
github.com/anthropics/anthropic-sdk-go v1.56.0 // indirect
github.com/aws/aws-sdk-go-v2 v1.30.3 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 // indirect
github.com/aws/aws-sdk-go-v2/config v1.27.27 // indirect
github.com/aws/aws-sdk-go-v2/credentials v1.17.27 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 // indirect
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 // indirect
github.com/aws/smithy-go v1.20.3 // indirect
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/bmatcuk/doublestar/v4 v4.10.0 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
github.com/buger/jsonparser v1.1.2 // indirect
github.com/bytedance/gopkg v0.1.3 // indirect
github.com/bytedance/sonic/loader v0.5.0 // indirect
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
@@ -60,6 +79,7 @@ require (
github.com/dlclark/regexp2 v1.10.0 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/evanphx/json-patch v0.5.2 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
github.com/gabriel-vasile/mimetype v1.4.2 // indirect
github.com/gin-contrib/sse v0.1.0 // indirect
github.com/go-logr/logr v1.4.3 // indirect
@@ -71,9 +91,14 @@ require (
github.com/go-viper/mapstructure/v2 v2.5.0 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/google/jsonschema-go v0.3.0 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/protobuf v1.5.4 // indirect
github.com/google/jsonschema-go v0.4.2 // indirect
github.com/google/s2a-go v0.1.7 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
github.com/goph/emperror v0.17.2 // indirect
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1 // indirect
github.com/invopop/jsonschema v0.14.0 // indirect
github.com/jolestar/go-commons-pool/v2 v2.1.2 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.2.10 // indirect
@@ -86,10 +111,14 @@ require (
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/nikolalohinski/gonja v1.5.3 // indirect
github.com/openai/openai-go/v3 v3.35.0 // indirect
github.com/pb33f/ordered-map/v2 v2.3.1 // indirect
github.com/pelletier/go-toml/v2 v2.2.3 // indirect
github.com/pkg/errors v0.9.1 // indirect
github.com/segmentio/asm v1.1.3 // indirect
github.com/segmentio/encoding v0.5.4 // indirect
github.com/sirupsen/logrus v1.9.3 // indirect
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 // indirect
github.com/tetratelabs/wazero v1.11.0 // indirect
github.com/tidwall/gjson v1.18.0 // indirect
github.com/tidwall/match v1.1.1 // indirect
@@ -100,17 +129,22 @@ require (
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
github.com/yargevad/filepathx v1.0.0 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
go.opencensus.io v0.24.0 // indirect
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.34.0 // indirect
go.opentelemetry.io/otel/metric v1.39.0 // indirect
go.opentelemetry.io/proto/otlp v1.5.0 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.yaml.in/yaml/v4 v4.0.0-rc.2 // indirect
golang.org/x/arch v0.15.0 // indirect
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 // indirect
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 // indirect
golang.org/x/oauth2 v0.34.0 // indirect
golang.org/x/sync v0.20.0 // indirect
golang.org/x/sys v0.43.0 // indirect
google.golang.org/api v0.189.0 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
google.golang.org/grpc v1.79.3 // indirect
+144 -12
View File
@@ -1,39 +1,82 @@
cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
cloud.google.com/go/auth v0.7.2 h1:uiha352VrCDMXg+yoBtaD0tUF4Kv9vrtrWPYXwutnDE=
cloud.google.com/go/auth v0.7.2/go.mod h1:VEc4p5NNxycWQTMQEDQF0bd6aTMb6VgYDXEwiJJQAbs=
cloud.google.com/go/auth/oauth2adapt v0.2.3 h1:MlxF+Pd3OmSudg/b1yZ5lJwoXCEaeedAguodky1PcKI=
cloud.google.com/go/auth/oauth2adapt v0.2.3/go.mod h1:tMQXOfZzFuNuUxOypHlQEXgdfX5cuhwU+ffUuXRJE8I=
cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs=
cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0 h1:g0EZJwz7xkXQiZAI5xi9f3WWFYBlX1CPTrR+NDToRkQ=
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0/go.mod h1:XCW7KnZet0Opnr7HccfUw1PLc4CjHqpcaxW8DHklNkQ=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0 h1:tfLQ34V6F7tVSwoTf/4lH5sE0o6eCJuNDTmH09nDpbc=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0/go.mod h1:9kIvujWAA58nmPmWB1m23fyWic1kYZMxD9CxaWn4Qpg=
github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 h1:ywEEhmNahHBihViHepv3xPBn1663uRv2t2q/ESv9seY=
github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0/go.mod h1:iZDifYGJTIgIIkYRNWPENUnqx6bJ2xnSDFI2tjwZNuY=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 h1:XHOnouVk1mxXfQidrMEnLlPk9UMeRtyBTnEFtxkV0kU=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0=
github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o=
github.com/anthropics/anthropic-sdk-go v1.56.0 h1:idVU14wOZ06D0GBNEvuhn927xXmBVEquo0469iDwLsc=
github.com/anthropics/anthropic-sdk-go v1.56.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/aws/aws-sdk-go-v2 v1.30.3 h1:jUeBtG0Ih+ZIFH0F4UkmL9w3cSpaMv9tYYDbzILP8dY=
github.com/aws/aws-sdk-go-v2 v1.30.3/go.mod h1:nIQjQVp5sfpQcTc9mPSr1B0PaWK5ByX9MOoDadSN4lc=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 h1:tW1/Rkad38LA15X4UQtjXZXNKsCgkshC3EbmcUmghTg=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3/go.mod h1:UbnqO+zjqk3uIt9yCACHJ9IVNhyhOCnYk8yA19SAWrM=
github.com/aws/aws-sdk-go-v2/config v1.27.27 h1:HdqgGt1OAP0HkEDDShEl0oSYa9ZZBSOmKpdpsDMdO90=
github.com/aws/aws-sdk-go-v2/config v1.27.27/go.mod h1:MVYamCg76dFNINkZFu4n4RjDixhVr51HLj4ErWzrVwg=
github.com/aws/aws-sdk-go-v2/credentials v1.17.27 h1:2raNba6gr2IfA0eqqiP2XiQ0UVOpGPgDSi0I9iAP+UI=
github.com/aws/aws-sdk-go-v2/credentials v1.17.27/go.mod h1:gniiwbGahQByxan6YjQUMcW4Aov6bLC3m+evgcoN4r4=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 h1:KreluoV8FZDEtI6Co2xuNk/UqI9iwMrOx/87PBNIKqw=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11/go.mod h1:SeSUYBLsMYFoRvHE0Tjvn7kbxaUhl75CJi1sbfhMxkU=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 h1:SoNJ4RlFEQEbtDcCEt+QG56MY4fm4W8rYirAmq+/DdU=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15/go.mod h1:U9ke74k1n2bf+RIgoX1SXFed1HLs51OgUSs+Ph0KJP8=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 h1:C6WHdGnTDIYETAm5iErQUiVNsclNx9qbJVPIt03B6bI=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15/go.mod h1:ZQLZqhcu+JhSrA9/NXRm8SkDvsycE+JkV3WGY41e+IM=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 h1:dT3MqvGhSoaIhRseqw2I0yH81l7wiR2vjs57O51EAm8=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3/go.mod h1:GlAeCkHwugxdHaueRr4nhPuY+WW+gR8UjlcqzPr1SPI=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 h1:HGErhhrxZlQ044RiM+WdoZxp0p+EGM62y3L6pwA4olE=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17/go.mod h1:RkZEx4l0EHYDJpWppMJ3nD9wZJAa8/0lq9aVC+r2UII=
github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 h1:BXx0ZIxvrJdSgSvKTZ+yRBeSqqgPM89VPlulEcl37tM=
github.com/aws/aws-sdk-go-v2/service/sso v1.22.4/go.mod h1:ooyCOXjvJEsUw7x+ZDHeISPMhtwI3ZCB7ggFMcFfWLU=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 h1:yiwVzJW2ZxZTurVbYWA7QOrAaCYQR72t0wrSBfoesUE=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4/go.mod h1:0oxfLkpz3rQ/CHlx5hB7H69YUpFiI1tql6Q6Ne+1bCw=
github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 h1:ZsDKRLXGWHk8WdtyYMoGNO7bTudrvuKpDKgMVRlepGE=
github.com/aws/aws-sdk-go-v2/service/sts v1.30.3/go.mod h1:zwySh8fpFyXp9yOr/KVzxOl8SRqgf/IDw5aUt9UKFcQ=
github.com/aws/smithy-go v1.20.3 h1:ryHwveWzPV5BIof6fyDvor6V3iUL7nTfiTKXHiW05nE=
github.com/aws/smithy-go v1.20.3/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA=
github.com/bmatcuk/doublestar/v4 v4.10.0 h1:zU9WiOla1YA122oLM6i4EXvGW62DvKZVxIe6TYWexEs=
github.com/bmatcuk/doublestar/v4 v4.10.0/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc=
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk=
github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8=
github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE=
github.com/bwmarrin/discordgo v0.29.0 h1:FmWeXFaKUwrcL3Cx65c20bTRW+vOb6k8AnaP+EgjDno=
github.com/bwmarrin/discordgo v0.29.0/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=
github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM=
github.com/bytedance/mockey v1.3.0 h1:ONLRdvhqmCfr9rTasUB8ZKCfvbdD2tohOg4u+4Q/ed0=
github.com/bytedance/mockey v1.3.0/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY=
github.com/bytedance/mockey v1.4.6 h1:pPkAFB6yiaaybvgp7DP1Rj4Ztiew3nsaMizoNkzsvNA=
github.com/bytedance/mockey v1.4.6/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY=
github.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uSE=
github.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k=
github.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE=
github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo=
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4=
github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
github.com/cloudwego/eino v0.9.14 h1:suNVibjtkPMiW7csFBdBqN3FRG0nRlCqoGKE7t0UDwY=
@@ -48,14 +91,17 @@ github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive
github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive v0.0.0-20260427010451-749e3706378b/go.mod h1:9R0RQrQSpg1JaNnRtw7+RfRAAv0HgdE348YnrlZ6coo=
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b h1:pOqupZQyc46rw2Z0HeybtTmSMTwqfTrbRuGDuDsNf2A=
github.com/cloudwego/eino-ext/components/embedding/openai v0.0.0-20260427010451-749e3706378b/go.mod h1:zyPrZT2bO6LyRJgVksQowR18jVgyLSvqK93hnO53/Lc=
github.com/cloudwego/eino-ext/components/model/agenticclaude v0.1.3 h1:nfO9jUPhIKTHDPyL7rso6y7Mx8cAmrIgkeJC6KqTuXI=
github.com/cloudwego/eino-ext/components/model/agenticclaude v0.1.3/go.mod h1:/dX1P7HNC7D3/nYLDoXBvdKAFnGS4tBbtJ8NrkaXuUk=
github.com/cloudwego/eino-ext/components/model/agenticopenai v0.2.1 h1:bIMyUxEZdBEofWbXQCozfGxXCsbyNjpZcRMfW+Kt5u0=
github.com/cloudwego/eino-ext/components/model/agenticopenai v0.2.1/go.mod h1:84jUcVAcmlGDRVX0ExABVJEfrIaplaAeRbjFYW9mzkI=
github.com/cloudwego/eino-ext/components/model/openai v0.1.13 h1:5XHRTiTD5bt9KQrMHcfvuWNklEC3tpm3XHejdozt9vM=
github.com/cloudwego/eino-ext/components/model/openai v0.1.13/go.mod h1:mgIoqYYOc0eECCqvLbEYpOJrQNTNxkwXzSJzFU+v5sQ=
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.17 h1:EeVcR1TslRA2IdNW1h/2LaGbPlffwGhQm99jM3zWZiI=
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.17/go.mod h1:Zkcx6DPTR2NfWmtSXbhItswGw6hqUezNPhNcke0pOG8=
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.18-0.20260527084435-846f52bd97c6 h1:ES/xufN5eqJ3h+9tw/tq6F8kkgnAxBAHVUB6nqKsIDU=
github.com/cloudwego/eino-ext/libs/acl/openai v0.1.18-0.20260527084435-846f52bd97c6/go.mod h1:5Xj74dGrfHo1z7I07Fzp3SlTF7Bt4tss3A2FSt8SqQ4=
github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 h1:6xNmx7iTtyBRev0+D/Tv1FZd4SCg8axKApyNyRsAt/w=
github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5/go.mod h1:KdCmV+x/BuvyMxRnYBlmVaq4OLiKW6iRQfvC62cvdkI=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s=
github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE=
@@ -67,12 +113,25 @@ github.com/disintegration/imaging v1.6.2 h1:w1LecBlG2Lnp8B3jk5zSuNqd7b4DXhcjwek1
github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4=
github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/eino-contrib/jsonschema v1.0.3 h1:2Kfsm1xlMV0ssY2nuxshS4AwbLFuqmPmzIjLVJ1Fsp0=
github.com/eino-contrib/jsonschema v1.0.3/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4=
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA=
github.com/envoyproxy/go-control-plane/envoy v1.36.0 h1:yg/JjO5E7ubRyKX3m07GF3reDNEnfOboJ0QySbH736g=
github.com/envoyproxy/go-control-plane/envoy v1.36.0/go.mod h1:ty89S1YCCVruQAm9OtKeEkQLTb+Lkz0k8v9W0Oxsv98=
github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
github.com/envoyproxy/protoc-gen-validate v1.3.0 h1:TvGH1wof4H33rezVKWSpqKz5NXWg5VPuZ0uONDT6eb4=
github.com/envoyproxy/protoc-gen-validate v1.3.0/go.mod h1:HvYl7zwPa5mffgyeTUHA9zHIH36nmrm7oCbo4YKoSWA=
github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k=
github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ=
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw=
github.com/fortytw2/leaktest v1.3.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g=
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
@@ -117,33 +176,49 @@ github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8=
github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA=
github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs=
github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w=
github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0=
github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8=
github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
github.com/google/pprof v0.0.0-20260402051712-545e8a4df936 h1:EwtI+Al+DeppwYX2oXJCETMO23COyaKGP6fHVpkpWpg=
github.com/google/pprof v0.0.0-20260402051712-545e8a4df936/go.mod h1:MxpfABSjhmINe3F1It9d+8exIHFvUqtLIRCdOGNXqiI=
github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfFxPRy3Bf7vr3h0cechB90XaQs=
github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18=
github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic=
github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g=
@@ -155,6 +230,8 @@ github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/ad
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1 h1:VNqngBF40hVlDloBruUehVYC3ArSgIyScOAyMRqBxRg=
github.com/grpc-ecosystem/grpc-gateway/v2 v2.25.1/go.mod h1:RBRO7fro65R6tjKzYgLAFo0t1QEXY1Dp+i/bvpRiqiQ=
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg=
github.com/invopop/jsonschema v0.14.0/go.mod h1:ygm6C2EaVNMBDPpaPlnOA2pFAxBnxGjFlMZABxm9n2I=
github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
github.com/jolestar/go-commons-pool/v2 v2.1.2 h1:E+XGo58F23t7HtZiC/W6jzO2Ux2IccSH/yx4nD+J1CM=
github.com/jolestar/go-commons-pool/v2 v2.1.2/go.mod h1:r4NYccrkS5UqP1YQI1COyTZ9UjPJAAGTUxzcsK1kqhY=
@@ -179,6 +256,8 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
github.com/larksuite/oapi-sdk-go/v3 v3.4.22 h1:57daKuslQPX9X3hC2idc5bu8bl2krfsBGWGJ6b5FlD8=
github.com/larksuite/oapi-sdk-go/v3 v3.4.22/go.mod h1:ZEplY+kwuIrj/nqw5uSCINNATcH3KdxSN7y+UxYY5fI=
github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
@@ -195,8 +274,8 @@ github.com/meguminnnnnnnnn/go-openai v0.1.2 h1:iXombGGjqjBrmE9WaSidUhhi3YQhf42QT
github.com/meguminnnnnnnnn/go-openai v0.1.2/go.mod h1:qs96ysDmxhE4BZoU45I43zcyfnaYxU3X+aRzLko/htY=
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=
github.com/modelcontextprotocol/go-sdk v1.2.0 h1:Y23co09300CEk8iZ/tMxIX1dVmKZkzoSBZOpJwUnc/s=
github.com/modelcontextprotocol/go-sdk v1.2.0/go.mod h1:6fM3LCm3yV7pAs8isnKLn07oKtB0MP9LHd3DfAcKw10=
github.com/modelcontextprotocol/go-sdk v1.3.1 h1:TfqtNKOIWN4Z1oqmPAiWDC2Jq7K9OdJaooe0teoXASI=
github.com/modelcontextprotocol/go-sdk v1.3.1/go.mod h1:DgVX498dMD8UJlseK1S5i1T4tFz2fkBk4xogC3D15nw=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -221,16 +300,23 @@ github.com/onsi/gomega v1.40.0 h1:Vtol0e1MghCD2ZVIilPDIg44XSL9l2QAn8ZNaljWcJc=
github.com/onsi/gomega v1.40.0/go.mod h1:M/Uqpu/8qTjtzCLUA2zJHX9Iilrau25x1PdoSRbWh5A=
github.com/openai/openai-go/v3 v3.35.0 h1:109x3epXMSE423KW2euR506GGFezcEt0s87MoWejpH0=
github.com/openai/openai-go/v3 v3.35.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
github.com/pb33f/ordered-map/v2 v2.3.1 h1:5319HDO0aw4DA4gzi+zv4FXU9UlSs3xGZ40wcP1nBjY=
github.com/pb33f/ordered-map/v2 v2.3.1/go.mod h1:qxFQgd0PkVUtOMCkTapqotNgzRhMPL7VvaHKbd1HnmQ=
github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pkoukk/tiktoken-go v0.1.8 h1:85ENo+3FpWgAACBaEUVp+lctuTcYUO7BtmfhlN/QTRo=
github.com/pkoukk/tiktoken-go v0.1.8/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
@@ -238,6 +324,10 @@ github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/f
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ=
github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc=
github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg=
github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0=
github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0=
github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
@@ -251,6 +341,8 @@ github.com/smarty/assertions v1.16.0 h1:EvHNkdRA4QHMrn75NZSoUQ/mAUXAYWfatfB01yTC
github.com/smarty/assertions v1.16.0/go.mod h1:duaaFdCS0K9dnoM50iyek/eYINOZ64gbh1Xlf6LG7AI=
github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY=
github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 h1:uOfcYT+3QungH6tIGSVCR/Y3KJmgJiHcojJbMTPDZAI=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1/go.mod h1:L1MQhA6x4dn9r007T033lsaZMv9EmBAdXyU/+EF40fo=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
@@ -272,14 +364,12 @@ github.com/tencent-connect/botgo v0.2.1 h1:+BrTt9Zh+awL28GWC4g5Na3nQaGRWb0N5IctS
github.com/tencent-connect/botgo v0.2.1/go.mod h1:oO1sG9ybhXNickvt+CVym5khwQ+uKhTR+IhTqEfOVsI=
github.com/tetratelabs/wazero v1.11.0 h1:+gKemEuKCTevU4d7ZTzlsvgd1uaToIDtlQlmNbwqYhA=
github.com/tetratelabs/wazero v1.11.0/go.mod h1:eV28rsN8Q+xwjogd7f4/Pp4xFxO7uOGbLcD/LzB1wiU=
github.com/tidwall/gjson v1.9.3 h1:hqzS9wAHMO+KVBBkLxYdkEeeFHuqr95GfClRLKlgK0E=
github.com/tidwall/gjson v1.9.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY=
github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
@@ -302,8 +392,14 @@ github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT0
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=
go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 h1:4Pp6oUg3+e/6M4C0A/3kJ2VYa++dsWVTtGgLVj5xtHg=
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0/go.mod h1:Mjt1i1INqiaoZOMGR1RIUJN+i3ChKoFRqzrRQhlkbs0=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 h1:jq9TW8u3so/bN+JPT166wjOI6/vQPF6Xe7nMNIltagk=
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0/go.mod h1:p8pYQP+m5XfbZm9fxtSKAbM6oIllS7s2AfxrChvc7iw=
go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48=
go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8=
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.34.0 h1:OeNbIYk/2C15ckl7glBlOBp5+WlYsOElzTNmiPW/x60=
@@ -332,6 +428,8 @@ go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo=
go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so=
go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
go.yaml.in/yaml/v4 v4.0.0-rc.2 h1:/FrI8D64VSr4HtGIlUtlFMGsm7H7pWTbj6vOLVZcA6s=
go.yaml.in/yaml/v4 v4.0.0-rc.2/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=
golang.org/x/arch v0.15.0 h1:QtOrQd0bTUnhNVNndMpLHNWrDmYzZ2KDqSrEymqInZw=
golang.org/x/arch v0.15.0/go.mod h1:JmwW7aLIoRUKgaTzhkiEFxvcEiQGyOg9BMonBJUS7EE=
golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
@@ -343,22 +441,31 @@ golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5y
golang.org/x/crypto v0.16.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI=
golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 h1:nDVHiLt8aIbd/VzvPWN6kSOPE7+F/fNFDSXLVYkE/Iw=
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394/go.mod h1:sIifuuw/Yco/y6yb6+bDNfyeQ/MdPUy/hKEMYQV17cM=
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8 h1:hVwzHzIUGRjiF7EcUjqNxk3NCfkPxbDKRdnNE1Rpg0U=
golang.org/x/image v0.0.0-20191009234506-e7c1f5e7dbb8/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM=
golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU=
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
@@ -368,10 +475,12 @@ golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA=
golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs=
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
golang.org/x/oauth2 v0.23.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
golang.org/x/oauth2 v0.34.0 h1:hqK/t4AKgbqWkdkcAeI8XLmbK+4m4G5YeQRrmiotGlw=
golang.org/x/oauth2 v0.34.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -379,6 +488,7 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -421,6 +531,10 @@ golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
@@ -435,10 +549,22 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
google.golang.org/api v0.189.0 h1:equMo30LypAkdkLMBqfeIqtyAnlyig1JSZArl4XPwdI=
google.golang.org/api v0.189.0/go.mod h1:FLWGJKb0hb+pU2j+rJqwbnsF+ym+fQs73rbJ+KAUgy8=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc=
google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=
google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww=
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg=
google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY=
google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
google.golang.org/grpc v1.33.2/go.mod h1:JMHMWHQWaTccqQQlmk3MJZS+GWXOdAesneDmEnv2fbc=
google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
@@ -446,7 +572,10 @@ google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ
google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE=
google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo=
google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
@@ -462,7 +591,10 @@ gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
+42
View File
@@ -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,11 +817,51 @@ 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)
}
}
// 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 {
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)
-6
View File
@@ -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 {
+21
View File
@@ -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
View File
@@ -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)
+98 -10
View File
@@ -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 {
@@ -817,7 +824,7 @@ type MCPConfig struct {
}
type OpenAIConfig struct {
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` // API 提供商: "openai"(默认) 或 "claude"claude 时自动桥接为 Anthropic Messages API
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` // API 提供商: "openai"(默认) 或 "claude"claude 使用 Eino 原生 Anthropic Messages API
APIKey string `yaml:"api_key" json:"api_key"`
BaseURL string `yaml:"base_url" json:"base_url"`
Model string `yaml:"model" json:"model"`
@@ -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
+101
View File
@@ -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 {
+39
View File
@@ -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
}
+45
View File
@@ -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")
}
})
}
}
+120
View File
@@ -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
View File
@@ -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
+42 -90
View File
@@ -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字段是否存在
-486
View File
@@ -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
}
+485
View File
@@ -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
}
+33 -9
View File
@@ -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 = &copy
}
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"
-1
View File
@@ -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"},
} {
+1 -20
View File
@@ -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)
@@ -20,16 +20,20 @@ func (db *DB) FindNearestToolExecutionArguments(conversationID, toolName string,
if window <= 0 {
window = 5 * time.Second
}
names := []string{toolName}
if !strings.Contains(toolName, "::") {
names = append(names, "eino_fs::"+toolName)
}
start := at.Add(-window)
end := at.Add(window)
rows, err := db.Query(`
SELECT id, arguments
FROM tool_executions
WHERE conversation_id = ?
AND tool_name = ?
AND tool_name IN (?, ?)
AND julianday(start_time) BETWEEN julianday(?) AND julianday(?)
ORDER BY ABS(julianday(start_time) - julianday(?)) ASC, start_time ASC
LIMIT 1`, conversationID, toolName, start, end, at)
LIMIT 1`, conversationID, names[0], names[len(names)-1], start, end, at)
if err != nil {
return "", nil, err
}
+84
View File
@@ -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
View File
@@ -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,最后返回完整回复
+18 -11
View File
@@ -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)
}
+70 -8
View File
@@ -18,11 +18,14 @@ import (
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/knowledge"
"cyberstrike-ai/internal/llm"
"cyberstrike-ai/internal/mcp"
"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"
"go.uber.org/zap"
"gopkg.in/yaml.v3"
@@ -93,6 +96,7 @@ type ConfigHandler struct {
db *database.DB
logger *zap.Logger
mu sync.RWMutex
toolGuard *toolguard.Manager
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
}
@@ -345,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(),
@@ -889,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 {
@@ -1160,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))
@@ -1216,7 +1229,7 @@ func (h *ConfigHandler) TestOpenAI(c *gin.Context) {
"max_completion_tokens": 5,
}
// 使用内部 openai Client 进行测试,若 provider 为 claude 会自动走桥接层
// OpenAI-compatible 通道使用内部客户端;Claude 通道在下方直接使用 Eino agenticclaude。
tmpCfg := &config.OpenAIConfig{
Provider: req.Provider,
BaseURL: baseURL,
@@ -1229,6 +1242,28 @@ func (h *ConfigHandler) TestOpenAI(c *gin.Context) {
defer cancel()
start := time.Now()
if llm.IsClaudeProvider(req.Provider) {
nativeModel, err := llm.NewClaudeAgenticModel(ctx, *tmpCfg, nil, 5, nil)
if err == nil {
_, err = nativeModel.Generate(ctx, []*schema.AgenticMessage{
schema.UserAgenticMessage("Hi"),
})
}
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"error": "连接失败: " + err.Error(),
})
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"model": tmpCfg.Model,
"latency_ms": time.Since(start).Milliseconds(),
})
return
}
var chatResp struct {
ID string `json:"id"`
Object string `json:"object"`
@@ -1713,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 {
@@ -2117,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()
+8
View File
@@ -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
+33 -12
View File
@@ -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")
@@ -497,6 +517,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{}
}
allow := map[string]bool{
"toolName": true, "toolCallId": true, "index": true, "total": true,
"arguments": true, "argumentsObj": true,
"success": true, "isError": true, "executionId": true,
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
"agentFacing": true,
@@ -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())
@@ -139,3 +198,67 @@ func TestProcessDetailsFullBackfillsEmptyToolCallArgumentsFromExecution(t *testi
t.Fatalf("executionId = %#v, want exec-whoami", data["executionId"])
}
}
func TestProcessDetailsPageBackfillsEinoFilesystemArgumentsFromPrefixedExecution(t *testing.T) {
gin.SetMode(gin.TestMode)
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-eino-fs-args.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
conversation, err := db.CreateConversation("eino fs args", 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)
}
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "calling read_file", map[string]interface{}{
"toolName": "read_file", "toolCallId": "call-read", "arguments": "", "argumentsObj": nil,
}); err != nil {
t.Fatalf("AddProcessDetail(tool_call): %v", err)
}
if err := db.SaveToolExecution(&mcp.ToolExecution{
ID: "exec-read",
ToolName: "eino_fs::read_file",
Arguments: map[string]interface{}{"file_path": "/tmp/requirements.txt", "limit": float64(2000)},
Status: "completed",
StartTime: time.Now(),
ConversationID: conversation.ID,
}); 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=50&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 {
ProcessDetails []map[string]interface{} `json:"processDetails"`
}
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
t.Fatalf("decode response: %v", err)
}
if len(response.ProcessDetails) != 1 {
t.Fatalf("process details = %d, want 1", len(response.ProcessDetails))
}
data, ok := response.ProcessDetails[0]["data"].(map[string]interface{})
if !ok {
t.Fatalf("data = %#v", response.ProcessDetails[0]["data"])
}
args, ok := data["argumentsObj"].(map[string]interface{})
if !ok {
t.Fatalf("argumentsObj = %#v", data["argumentsObj"])
}
if args["file_path"] != "/tmp/requirements.txt" {
t.Fatalf("file_path = %#v, want /tmp/requirements.txt", args["file_path"])
}
if data["arguments"] == nil {
t.Fatalf("arguments should be preserved in summarized page data: %#v", data)
}
}
+40 -22
View File
@@ -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,
})
}
+2
View File
@@ -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")
}
+210 -54
View File
@@ -1,11 +1,13 @@
package handler
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
@@ -87,6 +89,24 @@ type fofaSearchResponse struct {
Results []map[string]interface{} `json:"results"`
}
type spaceSearchEnvelope struct {
Code interface{} `json:"code"`
Message string `json:"message"`
Error string `json:"error"`
Query string `json:"query"`
Total int `json:"total"`
TotalCount int `json:"total_count"`
Page int `json:"page"`
PageSize int `json:"pagesize"`
Data json.RawMessage `json:"data"`
Matches json.RawMessage `json:"matches"`
Meta struct {
Pagination struct {
Total int `json:"total"`
} `json:"pagination"`
} `json:"meta"`
}
func normalizeSpaceSearchProvider(provider string) string {
switch strings.ToLower(strings.TrimSpace(provider)) {
case "", "fofa":
@@ -156,19 +176,20 @@ func (h *FofaHandler) resolveAPIKey(provider string) string {
}
func (h *FofaHandler) resolveBaseURL(provider string) string {
provider = normalizeSpaceSearchProvider(provider)
if h.cfg != nil {
switch normalizeSpaceSearchProvider(provider) {
switch provider {
case "zoomeye":
if v := strings.TrimSpace(h.cfg.ZoomEye.BaseURL); v != "" {
return v
return canonicalizeSpaceSearchBaseURL(provider, v)
}
case "quake":
if v := strings.TrimSpace(h.cfg.Quake.BaseURL); v != "" {
return v
return canonicalizeSpaceSearchBaseURL(provider, v)
}
case "shodan":
if v := strings.TrimSpace(h.cfg.Shodan.BaseURL); v != "" {
return v
return canonicalizeSpaceSearchBaseURL(provider, v)
}
default:
if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" {
@@ -176,11 +197,11 @@ func (h *FofaHandler) resolveBaseURL(provider string) string {
}
}
}
switch normalizeSpaceSearchProvider(provider) {
switch provider {
case "zoomeye":
return "https://api.zoomeye.org/v2/search"
return "https://api.zoomeye.ai/v2/search"
case "quake":
return "https://quake.360.cn/api/v3/search/quake_service"
return "https://quake.360.net/api/v3/search/quake_service"
case "shodan":
return "https://api.shodan.io"
default:
@@ -188,6 +209,20 @@ func (h *FofaHandler) resolveBaseURL(provider string) string {
}
}
func canonicalizeSpaceSearchBaseURL(provider, raw string) string {
v := strings.TrimSpace(raw)
if v == "" {
return v
}
switch normalizeSpaceSearchProvider(provider) {
case "zoomeye":
v = strings.Replace(v, "://api.zoomeye.org", "://api.zoomeye.ai", 1)
case "quake":
v = strings.Replace(v, "://quake.360.cn", "://quake.360.net", 1)
}
return v
}
// ParseNaturalLanguage 将自然语言解析为 FOFA 查询语法(仅生成,不执行查询)
func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) {
var req fofaParseRequest
@@ -716,23 +751,17 @@ func (h *FofaHandler) searchZoomEye(c *gin.Context, req fofaSearchRequest, apiKe
if fields := strings.TrimSpace(req.Fields); fields != "" {
body["fields"] = fields
}
var apiResp struct {
Code int `json:"code"`
Message string `json:"message"`
Query string `json:"query"`
Total int `json:"total"`
Page int `json:"page"`
PageSize int `json:"pagesize"`
Data []map[string]interface{} `json:"data"`
}
var apiResp spaceSearchEnvelope
if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "API-KEY", body, &apiResp, "ZoomEye") {
return
}
if apiResp.Code != 60000 {
msg := strings.TrimSpace(apiResp.Message)
if msg == "" {
msg = "ZoomEye 返回错误"
}
rows, err := decodeSpaceSearchRows(apiResp.Data)
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "解析 ZoomEye 响应失败: " + err.Error()})
return
}
if zoomEyeRequestFailed(apiResp.Code, apiResp.Message) {
msg := firstNonEmptySpaceSearchValue(apiResp.Message, messageFromRawObject(apiResp.Data), "ZoomEye 返回错误")
c.JSON(http.StatusBadGateway, gin.H{"error": msg})
return
}
@@ -744,8 +773,8 @@ func (h *FofaHandler) searchZoomEye(c *gin.Context, req fofaSearchRequest, apiKe
Page: firstPositive(apiResp.Page, req.Page),
Total: apiResp.Total,
Fields: fields,
ResultsCount: len(apiResp.Data),
Results: projectRows(apiResp.Data, fields),
ResultsCount: len(rows),
Results: projectRows(rows, fields),
})
}
@@ -766,25 +795,17 @@ func (h *FofaHandler) searchQuake(c *gin.Context, req fofaSearchRequest, apiKey
if len(fields) > 0 {
body["include"] = fields
}
var apiResp struct {
Code interface{} `json:"code"`
Message string `json:"message"`
TotalCount int `json:"total_count"`
Data []map[string]interface{} `json:"data"`
Meta struct {
Pagination struct {
Total int `json:"total"`
} `json:"pagination"`
} `json:"meta"`
}
var apiResp spaceSearchEnvelope
if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "X-QuakeToken", body, &apiResp, "Quake") {
return
}
rows, err := decodeSpaceSearchRows(apiResp.Data)
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "解析 Quake 响应失败: " + err.Error()})
return
}
if !isZeroSpaceSearchCode(apiResp.Code) {
msg := strings.TrimSpace(apiResp.Message)
if msg == "" {
msg = "Quake 返回错误"
}
msg := firstNonEmptySpaceSearchValue(apiResp.Message, messageFromRawObject(apiResp.Data), "Quake 返回错误")
c.JSON(http.StatusBadGateway, gin.H{"error": msg})
return
}
@@ -796,8 +817,8 @@ func (h *FofaHandler) searchQuake(c *gin.Context, req fofaSearchRequest, apiKey
Page: req.Page,
Total: total,
Fields: fields,
ResultsCount: len(apiResp.Data),
Results: projectRows(apiResp.Data, fields),
ResultsCount: len(rows),
Results: projectRows(rows, fields),
})
}
@@ -811,6 +832,9 @@ func isZeroSpaceSearchCode(code interface{}) bool {
return v == 0
case float64:
return v == 0
case json.Number:
n, err := v.Int64()
return err == nil && n == 0
case string:
return strings.TrimSpace(v) == "0"
default:
@@ -818,6 +842,128 @@ func isZeroSpaceSearchCode(code interface{}) bool {
}
}
func isZoomEyeSuccessCode(code interface{}) bool {
switch v := code.(type) {
case int:
return v == 60000
case int64:
return v == 60000
case float64:
return v == 60000
case json.Number:
n, err := v.Int64()
return err == nil && n == 60000
case string:
return strings.TrimSpace(v) == "60000"
default:
return false
}
}
func zoomEyeRequestFailed(code interface{}, message string) bool {
if isZoomEyeSuccessCode(code) {
return false
}
msg := strings.ToLower(strings.TrimSpace(message))
if code == nil || isZeroSpaceSearchCode(code) {
return msg != "" && msg != "success" && msg != "ok" && msg != "successful."
}
return true
}
func decodeSpaceSearchRows(raw json.RawMessage) ([]map[string]interface{}, error) {
raw = bytes.TrimSpace(raw)
if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
return nil, nil
}
switch raw[0] {
case '[':
var rows []map[string]interface{}
if err := json.Unmarshal(raw, &rows); err != nil {
return nil, err
}
if rows == nil {
return []map[string]interface{}{}, nil
}
return rows, nil
case '{':
var obj map[string]interface{}
if err := json.Unmarshal(raw, &obj); err != nil {
return nil, err
}
if len(obj) == 0 {
return []map[string]interface{}{}, nil
}
for _, key := range []string{"data", "items", "matches", "results", "list", "records"} {
nested, ok := obj[key]
if !ok {
continue
}
switch rows := nested.(type) {
case []map[string]interface{}:
return rows, nil
case []interface{}:
return interfaceSliceToRowMaps(rows), nil
}
}
return []map[string]interface{}{}, nil
default:
return nil, fmt.Errorf("unexpected JSON value")
}
}
func interfaceSliceToRowMaps(items []interface{}) []map[string]interface{} {
out := make([]map[string]interface{}, 0, len(items))
for _, item := range items {
if row, ok := item.(map[string]interface{}); ok {
out = append(out, row)
}
}
return out
}
func messageFromRawObject(raw json.RawMessage) string {
raw = bytes.TrimSpace(raw)
if len(raw) == 0 || raw[0] != '{' {
return ""
}
var obj map[string]interface{}
if err := json.Unmarshal(raw, &obj); err != nil {
return ""
}
for _, key := range []string{"message", "error", "errmsg", "msg"} {
if s, ok := obj[key].(string); ok {
if msg := strings.TrimSpace(s); msg != "" {
return msg
}
}
}
return ""
}
func extractRemoteAPIError(body []byte, statusCode int, label string) string {
trimmed := bytes.TrimSpace(body)
if len(trimmed) > 0 && trimmed[0] == '{' {
var obj map[string]interface{}
if err := json.Unmarshal(trimmed, &obj); err == nil {
for _, key := range []string{"error", "message", "errmsg", "msg"} {
if s, ok := obj[key].(string); ok {
if msg := strings.TrimSpace(s); msg != "" {
return msg
}
}
}
}
}
if len(trimmed) > 0 && trimmed[0] == '<' {
return fmt.Sprintf("%s 返回了网页而不是 JSONHTTP %d),请检查 Base URL 或网络是否被拦截", label, statusCode)
}
if statusCode < 200 || statusCode >= 300 {
return fmt.Sprintf("%s 返回非 2xx: %d", label, statusCode)
}
return ""
}
func (h *FofaHandler) searchShodan(c *gin.Context, req fofaSearchRequest, apiKey string) {
baseURL := strings.TrimRight(h.resolveBaseURL("shodan"), "/") + "/shodan/host/search"
u, err := url.Parse(baseURL)
@@ -826,11 +972,7 @@ func (h *FofaHandler) searchShodan(c *gin.Context, req fofaSearchRequest, apiKey
return
}
var apiResp struct {
Total int `json:"total"`
Matches []map[string]interface{} `json:"matches"`
Error string `json:"error"`
}
var apiResp spaceSearchEnvelope
targetSize := req.Size
if targetSize <= 0 {
targetSize = 100
@@ -852,19 +994,23 @@ func (h *FofaHandler) searchShodan(c *gin.Context, req fofaSearchRequest, apiKey
params.Set("fields", fields)
}
pageURL.RawQuery = params.Encode()
apiResp.Matches = nil
apiResp.Error = ""
apiResp = spaceSearchEnvelope{}
if !h.doJSONRequest(c, http.MethodGet, pageURL.String(), "", "", nil, &apiResp, "Shodan") {
return
}
if strings.TrimSpace(apiResp.Error) != "" {
c.JSON(http.StatusBadGateway, gin.H{"error": apiResp.Error})
if errMsg := firstNonEmptySpaceSearchValue(apiResp.Error, apiResp.Message); errMsg != "" {
c.JSON(http.StatusBadGateway, gin.H{"error": errMsg})
return
}
if len(apiResp.Matches) == 0 {
pageMatches, err := decodeSpaceSearchRows(apiResp.Matches)
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "解析 Shodan 响应失败: " + err.Error()})
return
}
if len(pageMatches) == 0 {
break
}
matches = append(matches, apiResp.Matches...)
matches = append(matches, pageMatches...)
if len(matches) >= targetSize {
matches = matches[:targetSize]
break
@@ -947,11 +1093,21 @@ func (h *FofaHandler) doJSONRequest(c *gin.Context, method, endpoint, apiKey, he
return false
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("%s 返回非 2xx: %d", label, resp.StatusCode)})
respBody, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
if err != nil {
c.JSON(http.StatusBadGateway, gin.H{"error": "读取 " + label + " 响应失败: " + err.Error()})
return false
}
if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
msg := extractRemoteAPIError(respBody, resp.StatusCode, label)
c.JSON(http.StatusBadGateway, gin.H{"error": msg})
return false
}
if err := json.Unmarshal(respBody, out); err != nil {
if msg := extractRemoteAPIError(respBody, resp.StatusCode, label); strings.Contains(msg, "网页") {
c.JSON(http.StatusBadGateway, gin.H{"error": msg})
return false
}
c.JSON(http.StatusBadGateway, gin.H{"error": "解析 " + label + " 响应失败: " + err.Error()})
return false
}
+139 -1
View File
@@ -164,7 +164,7 @@ func TestQuakeSearchHandlesStringErrorCode(t *testing.T) {
t.Fatalf("Quake token = %q, want test-quake-key", got)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"code":"q5000","message":"查询语法错误"}`))
_, _ = w.Write([]byte(`{"code":"q5000","message":"查询语法错误","data":{}}`))
}))
defer quakeServer.Close()
@@ -195,6 +195,144 @@ func TestQuakeSearchHandlesStringErrorCode(t *testing.T) {
}
}
func TestQuakeSearchAcceptsArrayData(t *testing.T) {
gin.SetMode(gin.TestMode)
t.Setenv("QUAKE_API_KEY", "")
quakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"code":0,"message":"Successful.","data":[{"ip":"1.1.1.1","port":53}],"meta":{"pagination":{"total":1}}}`))
}))
defer quakeServer.Close()
h := NewFofaHandler(&config.Config{
Quake: config.SpaceSearchConfig{
BaseURL: quakeServer.URL,
APIKey: "test-quake-key",
},
}, zap.NewNop())
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
body := `{"provider":"quake","query":"ip:\"1.1.1.1\"","fields":"ip,port","size":10,"page":1}`
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
ctx.Request.Header.Set("Content-Type", "application/json")
h.Search(ctx)
if recorder.Code != http.StatusOK {
t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String())
}
var response fofaSearchResponse
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("decode response: %v", err)
}
if response.ResultsCount != 1 || response.Total != 1 {
t.Fatalf("results_count=%d total=%d, want 1/1", response.ResultsCount, response.Total)
}
}
func TestZoomEyeSearchHandlesObjectDataError(t *testing.T) {
gin.SetMode(gin.TestMode)
t.Setenv("ZOOMEYE_API_KEY", "")
zoomeyeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if got := r.Header.Get("API-KEY"); got != "test-zoomeye-key" {
t.Fatalf("ZoomEye API-KEY = %q, want test-zoomeye-key", got)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"code":40001,"message":"invalid api key","data":{}}`))
}))
defer zoomeyeServer.Close()
h := NewFofaHandler(&config.Config{
ZoomEye: config.SpaceSearchConfig{
BaseURL: zoomeyeServer.URL,
APIKey: "test-zoomeye-key",
},
}, zap.NewNop())
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
body := `{"provider":"zoomeye","query":"ip=\"1.1.1.1\"","fields":"ip,port","size":10,"page":1}`
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
ctx.Request.Header.Set("Content-Type", "application/json")
h.Search(ctx)
if recorder.Code != http.StatusBadGateway {
t.Fatalf("Search() status = %d, want %d, body = %s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
bodyText := recorder.Body.String()
if !strings.Contains(bodyText, "invalid api key") {
t.Fatalf("response should include ZoomEye error message, got %s", bodyText)
}
if strings.Contains(bodyText, "cannot unmarshal") {
t.Fatalf("response exposed JSON type decoding failure: %s", bodyText)
}
}
func TestShodanSearchSurfacesJSONErrorOnUnauthorized(t *testing.T) {
gin.SetMode(gin.TestMode)
t.Setenv("SHODAN_API_KEY", "")
shodanServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":"Invalid API key"}`))
}))
defer shodanServer.Close()
h := NewFofaHandler(&config.Config{
Shodan: config.SpaceSearchConfig{
BaseURL: shodanServer.URL,
APIKey: "test-shodan-key",
},
}, zap.NewNop())
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
body := `{"provider":"shodan","query":"product:nginx","fields":"ip_str,port","size":10,"page":1}`
ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
ctx.Request.Header.Set("Content-Type", "application/json")
h.Search(ctx)
if recorder.Code != http.StatusBadGateway {
t.Fatalf("Search() status = %d, want %d, body = %s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
bodyText := recorder.Body.String()
if !strings.Contains(bodyText, "Invalid API key") {
t.Fatalf("response should include Shodan error message, got %s", bodyText)
}
if strings.Contains(bodyText, "非 2xx") {
t.Fatalf("response should not hide Shodan error behind generic status, got %s", bodyText)
}
}
func TestCanonicalizeSpaceSearchBaseURLMigratesLegacyHosts(t *testing.T) {
t.Parallel()
cases := []struct {
provider string
in string
want string
}{
{provider: "zoomeye", in: "https://api.zoomeye.org/v2/search", want: "https://api.zoomeye.ai/v2/search"},
{provider: "quake", in: "https://quake.360.cn/api/v3/search/quake_service", want: "https://quake.360.net/api/v3/search/quake_service"},
{provider: "shodan", in: "https://api.shodan.io", want: "https://api.shodan.io"},
}
for _, tc := range cases {
tc := tc
t.Run(tc.provider, func(t *testing.T) {
t.Parallel()
got := canonicalizeSpaceSearchBaseURL(tc.provider, tc.in)
if got != tc.want {
t.Fatalf("canonicalizeSpaceSearchBaseURL() = %q, want %q", got, tc.want)
}
})
}
}
func TestExtractInfoCollectJSONObject(t *testing.T) {
t.Parallel()
cases := []struct {
-438
View File
@@ -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)
}
-41
View File
@@ -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")
}
})
}
}
+74 -8
View File
@@ -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 中的全局免审批工具白名单。
+19 -8
View File
@@ -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
}
+44 -23
View File
@@ -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,
})
}
-497
View File
@@ -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{}{
+6 -11
View File
@@ -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",
// 新增缺失端点响应
+11 -1
View File
@@ -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
+92
View File
@@ -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{}
}
+175
View File
@@ -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)
}
+193
View File
@@ -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")
}
}
+34 -15
View File
@@ -8,9 +8,11 @@ import (
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/llm"
"cyberstrike-ai/internal/openai"
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/flow/retriever/multiquery"
"go.uber.org/zap"
)
@@ -29,21 +31,38 @@ func WireRetrieverPipeline(ctx context.Context, r *Retriever, openAI *config.Ope
}
r.wireOpenAI = openAI
httpClient := openai.NewEinoHTTPClient(openAI, &http.Client{Timeout: 120 * time.Second})
maxCompletionTokens := openAI.MaxCompletionTokensEffective()
chatCfg := &einoopenai.ChatModelConfig{
APIKey: strings.TrimSpace(openAI.APIKey),
BaseURL: strings.TrimSuffix(strings.TrimSpace(openAI.BaseURL), "/"),
Model: strings.TrimSpace(openAI.Model),
HTTPClient: httpClient,
MaxCompletionTokens: &maxCompletionTokens,
}
if chatCfg.Model == "" {
chatCfg.Model = "gpt-4o"
}
rewriteLLM, err := einoopenai.NewChatModel(ctx, chatCfg)
if err != nil {
return fmt.Errorf("multi_query rewrite model: %w", err)
baseHTTPClient := &http.Client{Timeout: 120 * time.Second}
var rewriteLLM model.ChatModel
if llm.IsClaudeProvider(openAI.Provider) {
nativeModel, err := llm.NewClaudeAgenticModel(
ctx,
*openAI,
baseHTTPClient,
openAI.MaxCompletionTokensEffective(),
nil,
)
if err != nil {
return fmt.Errorf("multi_query native Claude rewrite model: %w", err)
}
rewriteLLM = llm.NewAgenticChatModelAdapter(nativeModel)
} else {
httpClient := openai.NewEinoHTTPClient(openAI, baseHTTPClient)
maxCompletionTokens := openAI.MaxCompletionTokensEffective()
chatCfg := &einoopenai.ChatModelConfig{
APIKey: strings.TrimSpace(openAI.APIKey),
BaseURL: strings.TrimSuffix(strings.TrimSpace(openAI.BaseURL), "/"),
Model: strings.TrimSpace(openAI.Model),
HTTPClient: httpClient,
MaxCompletionTokens: &maxCompletionTokens,
}
if chatCfg.Model == "" {
chatCfg.Model = "gpt-4o"
}
var err error
rewriteLLM, err = einoopenai.NewChatModel(ctx, chatCfg)
if err != nil {
return fmt.Errorf("multi_query rewrite model: %w", err)
}
}
reranker, err := NewHTTPReranker(&r.config.Rerank, openAI, r.logger)
+125
View File
@@ -0,0 +1,125 @@
package llm
import (
"context"
"fmt"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
// AgenticChatModelAdapter exposes a text-oriented AgenticModel as a classic
// BaseChatModel for Eino components that have not adopted AgenticMessage yet.
// It adapts only Eino's in-memory message shape; no HTTP protocol is translated.
type AgenticChatModelAdapter struct {
model model.AgenticModel
tools []*schema.ToolInfo
}
func NewAgenticChatModelAdapter(agenticModel model.AgenticModel) model.ChatModel {
return &AgenticChatModelAdapter{model: agenticModel}
}
func (a *AgenticChatModelAdapter) BindTools(tools []*schema.ToolInfo) error {
if a == nil || a.model == nil {
return fmt.Errorf("agentic chat adapter: model is nil")
}
a.tools = append([]*schema.ToolInfo(nil), tools...)
return nil
}
func (a *AgenticChatModelAdapter) Generate(
ctx context.Context,
input []*schema.Message,
opts ...model.Option,
) (*schema.Message, error) {
if a == nil || a.model == nil {
return nil, fmt.Errorf("agentic chat adapter: model is nil")
}
out, err := a.model.Generate(ctx, classicMessagesToAgentic(input), commonAgenticOptions(a.tools, opts...)...)
if err != nil {
return nil, err
}
return agenticMessageToClassic(out), nil
}
func (a *AgenticChatModelAdapter) Stream(
ctx context.Context,
input []*schema.Message,
opts ...model.Option,
) (*schema.StreamReader[*schema.Message], error) {
if a == nil || a.model == nil {
return nil, fmt.Errorf("agentic chat adapter: model is nil")
}
stream, err := a.model.Stream(ctx, classicMessagesToAgentic(input), commonAgenticOptions(a.tools, opts...)...)
if err != nil {
return nil, err
}
return schema.StreamReaderWithConvert(stream, func(msg *schema.AgenticMessage) (*schema.Message, error) {
return agenticMessageToClassic(msg), nil
}), nil
}
func classicMessagesToAgentic(input []*schema.Message) []*schema.AgenticMessage {
out := make([]*schema.AgenticMessage, 0, len(input))
for _, msg := range input {
if msg == nil {
continue
}
role := schema.AgenticRoleTypeUser
switch msg.Role {
case schema.System:
role = schema.AgenticRoleTypeSystem
case schema.Assistant:
role = schema.AgenticRoleTypeAssistant
}
agentic := &schema.AgenticMessage{Role: role}
if msg.Role == schema.Assistant {
if msg.Content != "" {
agentic.ContentBlocks = append(agentic.ContentBlocks, schema.NewContentBlock(&schema.AssistantGenText{Text: msg.Content}))
}
} else if msg.Content != "" {
agentic.ContentBlocks = append(agentic.ContentBlocks, schema.NewContentBlock(&schema.UserInputText{Text: msg.Content}))
}
out = append(out, agentic)
}
return out
}
func agenticMessageToClassic(msg *schema.AgenticMessage) *schema.Message {
if msg == nil {
return nil
}
content, reasoning := AgenticText(msg)
return &schema.Message{
Role: schema.Assistant,
Content: content,
ReasoningContent: reasoning,
}
}
func commonAgenticOptions(boundTools []*schema.ToolInfo, opts ...model.Option) []model.Option {
common := model.GetCommonOptions(&model.Options{
Tools: append([]*schema.ToolInfo(nil), boundTools...),
}, opts...)
out := make([]model.Option, 0, 6)
if common.Temperature != nil {
out = append(out, model.WithTemperature(*common.Temperature))
}
if common.Model != nil {
out = append(out, model.WithModel(*common.Model))
}
if common.TopP != nil {
out = append(out, model.WithTopP(*common.TopP))
}
if common.MaxTokens != nil {
out = append(out, model.WithMaxTokens(*common.MaxTokens))
}
if len(common.Stop) > 0 {
out = append(out, model.WithStop(common.Stop))
}
if common.Tools != nil {
out = append(out, model.WithTools(common.Tools))
}
return out
}
+60
View File
@@ -0,0 +1,60 @@
package llm
import (
"context"
"net/http"
"strings"
"cyberstrike-ai/internal/config"
agenticclaude "github.com/cloudwego/eino-ext/components/model/agenticclaude"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
func IsClaudeProvider(provider string) bool {
provider = strings.ToLower(strings.TrimSpace(provider))
return provider == "claude" || provider == "anthropic"
}
func NewClaudeAgenticModel(
ctx context.Context,
cfg config.OpenAIConfig,
httpClient *http.Client,
maxTokens int,
extraFields map[string]any,
) (model.AgenticModel, error) {
if maxTokens <= 0 {
maxTokens = cfg.MaxCompletionTokensEffective()
}
if cfg.IsDeepSeekEndpointOrModel() {
httpClient = newDeepSeekAnthropicCompatibleClient(httpClient)
}
return agenticclaude.New(ctx, &agenticclaude.Config{
APIKey: strings.TrimSpace(cfg.APIKey),
BaseURL: strings.TrimSuffix(strings.TrimSpace(cfg.BaseURL), "/"),
Model: strings.TrimSpace(cfg.Model),
MaxTokens: maxTokens,
HTTPClient: httpClient,
ExtraFields: extraFields,
})
}
func AgenticText(msg *schema.AgenticMessage) (content, reasoning string) {
if msg == nil {
return "", ""
}
var contentParts, reasoningParts []string
for _, block := range msg.ContentBlocks {
if block == nil {
continue
}
switch {
case block.AssistantGenText != nil:
contentParts = append(contentParts, block.AssistantGenText.Text)
case block.Reasoning != nil:
reasoningParts = append(reasoningParts, block.Reasoning.Text)
}
}
return strings.Join(contentParts, ""), strings.Join(reasoningParts, "")
}
+76
View File
@@ -0,0 +1,76 @@
package llm
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
)
// newDeepSeekAnthropicCompatibleClient compensates for DeepSeek's Anthropic
// endpoint lagging behind the current Anthropic SDK. The SDK emits
// {"type":"custom"} for function tools, while DeepSeek expects the older
// name/input_schema/description shape without that discriminator.
//
// This is a field-level compatibility fix; requests still originate from
// Eino's native agenticclaude model and remain Anthropic Messages API requests.
func newDeepSeekAnthropicCompatibleClient(base *http.Client) *http.Client {
if base == nil {
base = http.DefaultClient
}
cloned := *base
transport := base.Transport
if transport == nil {
transport = http.DefaultTransport
}
cloned.Transport = &deepSeekAnthropicCompatRoundTripper{base: transport}
return &cloned
}
type deepSeekAnthropicCompatRoundTripper struct {
base http.RoundTripper
}
func (rt *deepSeekAnthropicCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if req == nil || req.Body == nil || req.Method != http.MethodPost {
return rt.base.RoundTrip(req)
}
body, err := io.ReadAll(req.Body)
if err != nil {
return nil, fmt.Errorf("read DeepSeek Anthropic request: %w", err)
}
_ = req.Body.Close()
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
req.Body = io.NopCloser(bytes.NewReader(body))
return rt.base.RoundTrip(req)
}
tools, ok := payload["tools"].([]any)
if !ok {
req.Body = io.NopCloser(bytes.NewReader(body))
return rt.base.RoundTrip(req)
}
changed := false
for _, rawTool := range tools {
tool, ok := rawTool.(map[string]any)
if !ok || tool["type"] != "custom" {
continue
}
delete(tool, "type")
changed = true
}
if changed {
body, err = json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("marshal DeepSeek Anthropic request: %w", err)
}
}
req.Body = io.NopCloser(bytes.NewReader(body))
req.ContentLength = int64(len(body))
req.GetBody = func() (io.ReadCloser, error) {
return io.NopCloser(bytes.NewReader(body)), nil
}
return rt.base.RoundTrip(req)
}
@@ -0,0 +1,55 @@
package llm
import (
"bytes"
"io"
"net/http"
"strings"
"testing"
)
type captureRoundTripper struct {
body string
}
func (rt *captureRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
body, err := io.ReadAll(req.Body)
if err != nil {
return nil, err
}
rt.body = string(body)
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(bytes.NewReader(nil)),
Request: req,
}, nil
}
func TestDeepSeekAnthropicCompatStripsOnlyCustomToolType(t *testing.T) {
t.Parallel()
capture := &captureRoundTripper{}
client := newDeepSeekAnthropicCompatibleClient(&http.Client{Transport: capture})
req, err := http.NewRequest(
http.MethodPost,
"https://api.deepseek.com/anthropic/v1/messages",
strings.NewReader(`{"tools":[{"type":"custom","name":"mcp_tool","input_schema":{"type":"object"}},{"type":"web_search_20260209","name":"web_search"}]}`),
)
if err != nil {
t.Fatalf("NewRequest: %v", err)
}
resp, err := client.Do(req)
if err != nil {
t.Fatalf("Do: %v", err)
}
_ = resp.Body.Close()
if strings.Contains(capture.body, `"type":"custom"`) {
t.Fatalf("custom discriminator was not removed: %s", capture.body)
}
if !strings.Contains(capture.body, `"type":"web_search_20260209"`) {
t.Fatalf("server tool discriminator was removed: %s", capture.body)
}
if !strings.Contains(capture.body, `"name":"mcp_tool"`) {
t.Fatalf("custom tool definition was removed: %s", capture.body)
}
}
+88
View File
@@ -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...)
}
+15 -1
View File
@@ -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}
+120
View File
@@ -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)
}
}
+86
View File
@@ -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)
}
}
}
+2
View File
@@ -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,
}
}
+3
View File
@@ -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)
+38 -8
View File
@@ -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
+44 -8
View File
@@ -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
View File
@@ -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
}
+41
View File
@@ -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}
}
+239
View File
@@ -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)
}
}
+9 -3
View File
@@ -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))
}
}
+71 -10
View File
@@ -48,6 +48,48 @@ func isInterruptContinue(ctx context.Context) bool {
return errors.Is(context.Cause(ctx), ErrInterruptContinue)
}
func isEinoStreamCanceled(err error) bool {
if err == nil {
return false
}
if errors.Is(err, adk.ErrStreamCanceled) {
return true
}
var streamCanceled *adk.StreamCanceledError
return errors.As(err, &streamCanceled)
}
func isEinoCancelError(err error) bool {
if err == nil {
return false
}
var cancelErr *adk.CancelError
return errors.As(err, &cancelErr)
}
// isEinoVoluntaryCancelErr reports cancel signals produced by Agent Cancel /
// TurnLoop preempt (CancelError, ErrStreamCanceled, context.Canceled).
func isEinoVoluntaryCancelErr(err error) bool {
if err == nil {
return false
}
return isEinoCancelError(err) || isEinoStreamCanceled(err) || errors.Is(err, context.Canceled)
}
// isEinoTurnLoopPreemptErr is true when a cancel/stream-cancel leaked from the
// current agent turn while the host task context is still alive. TurnLoop
// interrupt-continue does not cancel the parent context; treating that leak as
// fatal would abort the whole run instead of starting the queued next turn.
func isEinoTurnLoopPreemptErr(ctx context.Context, err error) bool {
if err == nil || !isEinoVoluntaryCancelErr(err) {
return false
}
if ctx != nil && ctx.Err() != nil {
return false
}
return true
}
func isEinoIterationLimitError(err error) bool {
if err == nil {
return false
@@ -401,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,
+8 -2
View File
@@ -110,10 +110,16 @@ func AgenticMessageToEino(msg *schema.AgenticMessage) []*schema.Message {
base.Role = schema.Assistant
base.Content += block.AssistantGenText.Text
case block.FunctionToolCall != nil:
var index *int
if block.StreamingMeta != nil {
i := block.StreamingMeta.Index
index = &i
}
base.Role = schema.Assistant
base.ToolCalls = append(base.ToolCalls, schema.ToolCall{
ID: block.FunctionToolCall.CallID,
Type: "function",
Index: index,
ID: block.FunctionToolCall.CallID,
Type: "function",
Function: schema.FunctionCall{
Name: block.FunctionToolCall.Name,
Arguments: block.FunctionToolCall.Arguments,
@@ -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")
@@ -0,0 +1,118 @@
package multiagent
import (
"context"
"fmt"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
// agenticToolCallingChatModelAdapter lets Eino's classic plan_execute
// Planner/Replanner consume a native AgenticModel without translating the HTTP
// protocol. Only the in-memory Eino message and option shapes are adapted.
type agenticToolCallingChatModelAdapter struct {
model model.AgenticModel
tools []*schema.ToolInfo
}
func newAgenticToolCallingChatModelAdapter(agenticModel model.AgenticModel) model.ToolCallingChatModel {
return &agenticToolCallingChatModelAdapter{model: agenticModel}
}
func (m *agenticToolCallingChatModelAdapter) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
if m == nil || m.model == nil {
return nil, fmt.Errorf("agentic tool-calling adapter: model is nil")
}
clonedTools := append([]*schema.ToolInfo(nil), tools...)
return &agenticToolCallingChatModelAdapter{model: m.model, tools: clonedTools}, nil
}
func (m *agenticToolCallingChatModelAdapter) Generate(
ctx context.Context,
input []*schema.Message,
opts ...model.Option,
) (*schema.Message, error) {
if m == nil || m.model == nil {
return nil, fmt.Errorf("agentic tool-calling adapter: model is nil")
}
out, err := m.model.Generate(ctx, EinoMessagesToAgentic(input), m.agenticOptions(opts...)...)
if err != nil {
return nil, err
}
converted := AgenticMessageToEino(out)
if len(converted) == 0 {
return nil, fmt.Errorf("agentic tool-calling adapter: model returned no message")
}
return converted[0], nil
}
func (m *agenticToolCallingChatModelAdapter) Stream(
ctx context.Context,
input []*schema.Message,
opts ...model.Option,
) (*schema.StreamReader[*schema.Message], error) {
if m == nil || m.model == nil {
return nil, fmt.Errorf("agentic tool-calling adapter: model is nil")
}
stream, err := m.model.Stream(ctx, EinoMessagesToAgentic(input), m.agenticOptions(opts...)...)
if err != nil {
return nil, err
}
return agenticStreamToEinoStream(stream), nil
}
func (m *agenticToolCallingChatModelAdapter) agenticOptions(opts ...model.Option) []model.Option {
common := model.GetCommonOptions(&model.Options{
Tools: append([]*schema.ToolInfo(nil), m.tools...),
}, opts...)
out := make([]model.Option, 0, 8)
if common.Temperature != nil {
out = append(out, model.WithTemperature(*common.Temperature))
}
if common.Model != nil {
out = append(out, model.WithModel(*common.Model))
}
if common.TopP != nil {
out = append(out, model.WithTopP(*common.TopP))
}
if common.MaxTokens != nil {
out = append(out, model.WithMaxTokens(*common.MaxTokens))
}
if len(common.Stop) > 0 {
out = append(out, model.WithStop(common.Stop))
}
if common.Tools != nil {
out = append(out, model.WithTools(common.Tools))
}
if common.AgenticToolChoice != nil {
out = append(out, model.WithAgenticToolChoice(common.AgenticToolChoice))
} else if common.ToolChoice != nil {
out = append(out, model.WithAgenticToolChoice(classicToolChoiceToAgentic(
*common.ToolChoice,
common.AllowedToolNames,
)))
}
return out
}
func classicToolChoiceToAgentic(choice schema.ToolChoice, allowedNames []string) *schema.AgenticToolChoice {
allowed := make([]*schema.AllowedTool, 0, len(allowedNames))
for _, name := range allowedNames {
if name != "" {
allowed = append(allowed, &schema.AllowedTool{FunctionName: name})
}
}
out := &schema.AgenticToolChoice{Type: choice}
switch choice {
case schema.ToolChoiceAllowed:
if len(allowed) > 0 {
out.Allowed = &schema.AgenticAllowedToolChoice{Tools: allowed}
}
case schema.ToolChoiceForced:
if len(allowed) > 0 {
out.Forced = &schema.AgenticForcedToolChoice{Tools: allowed}
}
}
return out
}
@@ -0,0 +1,118 @@
package multiagent
import (
"context"
"testing"
"github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
)
type capturingAgenticToolCallingModel struct {
input []*schema.AgenticMessage
options *model.Options
}
func (m *capturingAgenticToolCallingModel) Generate(
_ context.Context,
input []*schema.AgenticMessage,
opts ...model.Option,
) (*schema.AgenticMessage, error) {
m.input = input
m.options = model.GetCommonOptions(nil, opts...)
return &schema.AgenticMessage{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlock(&schema.FunctionToolCall{
CallID: "call-1",
Name: "emit_plan",
Arguments: `{"steps":["inspect"]}`,
}),
},
}, nil
}
func (m *capturingAgenticToolCallingModel) Stream(
context.Context,
[]*schema.AgenticMessage,
...model.Option,
) (*schema.StreamReader[*schema.AgenticMessage], error) {
return schema.StreamReaderFromArray([]*schema.AgenticMessage{
{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlockChunk(&schema.FunctionToolCall{
CallID: "call-1",
Name: "emit_plan",
Arguments: `{"steps":[`,
}, &schema.StreamingMeta{Index: 0}),
},
},
{
Role: schema.AgenticRoleTypeAssistant,
ContentBlocks: []*schema.ContentBlock{
schema.NewContentBlockChunk(&schema.FunctionToolCall{
Arguments: `"inspect"]}`,
}, &schema.StreamingMeta{Index: 0}),
},
},
}), nil
}
func TestAgenticToolCallingAdapterConvertsForcedToolChoice(t *testing.T) {
t.Parallel()
native := &capturingAgenticToolCallingModel{}
adapter, err := newAgenticToolCallingChatModelAdapter(native).WithTools([]*schema.ToolInfo{{
Name: "emit_plan",
Desc: "emit a structured plan",
}})
if err != nil {
t.Fatalf("WithTools: %v", err)
}
out, err := adapter.Generate(
context.Background(),
[]*schema.Message{schema.UserMessage("plan this task")},
model.WithToolChoice(schema.ToolChoiceForced),
)
if err != nil {
t.Fatalf("Generate: %v", err)
}
if len(native.input) != 1 || native.input[0].Role != schema.AgenticRoleTypeUser {
t.Fatalf("native input = %#v", native.input)
}
if native.options == nil || len(native.options.Tools) != 1 || native.options.Tools[0].Name != "emit_plan" {
t.Fatalf("native tools = %#v", native.options)
}
if native.options.AgenticToolChoice == nil || native.options.AgenticToolChoice.Type != schema.ToolChoiceForced {
t.Fatalf("agentic tool choice = %#v", native.options.AgenticToolChoice)
}
if len(out.ToolCalls) != 1 || out.ToolCalls[0].Function.Name != "emit_plan" {
t.Fatalf("classic output = %#v", out)
}
}
func TestAgenticToolCallingAdapterPreservesStreamingToolCallIndex(t *testing.T) {
t.Parallel()
adapter := newAgenticToolCallingChatModelAdapter(&capturingAgenticToolCallingModel{})
stream, err := adapter.Stream(context.Background(), []*schema.Message{
schema.UserMessage("plan this task"),
})
if err != nil {
t.Fatalf("Stream: %v", err)
}
out, err := schema.ConcatMessageStream(stream)
if err != nil {
t.Fatalf("ConcatMessageStream: %v", err)
}
if len(out.ToolCalls) != 1 {
t.Fatalf("tool calls = %#v, want one merged call", out.ToolCalls)
}
call := out.ToolCalls[0]
if call.Index == nil || *call.Index != 0 {
t.Fatalf("tool call index = %#v", call.Index)
}
if call.Function.Name != "emit_plan" || call.Function.Arguments != `{"steps":["inspect"]}` {
t.Fatalf("merged tool call = %#v", call)
}
}
@@ -2,7 +2,6 @@ package multiagent
import (
"context"
"errors"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
@@ -139,7 +138,7 @@ func (h *einoAssistantStreamEventHandler) Handle(mv *adk.MessageVariant, agentNa
streamUsage = maxEinoTokenUsage(streamUsage, chunk.ResponseMeta.Usage)
}
})
if recvErr != nil && !errors.Is(recvErr, context.Canceled) && h.logger != nil {
if recvErr != nil && !isEinoVoluntaryCancelErr(recvErr) && h.logger != nil {
h.logger.Warn("eino stream recv error, flushing incomplete stream",
zap.Error(recvErr),
zap.String("agent", agentName),
@@ -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"
@@ -64,6 +64,17 @@ func toolCallArgsFromAccumulated(msgs []adk.Message, toolCallID, expectToolName
return map[string]interface{}{}
}
func mustMarshalToolArguments(args map[string]interface{}) string {
if len(args) == 0 {
return "{}"
}
raw, err := json.Marshal(args)
if err != nil {
return "{}"
}
return string(raw)
}
// beginEinoADKFilesystemToolMonitor 在 Eino ADK filesystem 工具开始调用时写入 running 状态。
func beginEinoADKFilesystemToolMonitor(
ctx context.Context,
@@ -71,6 +82,7 @@ func beginEinoADKFilesystemToolMonitor(
rec einomcp.ExecutionRecorder,
binder *MCPExecutionBinder,
toolCallID, toolName string,
args map[string]interface{},
) {
if ag == nil || rec == nil {
return
@@ -87,7 +99,7 @@ func beginEinoADKFilesystemToolMonitor(
return
}
storedName := "eino_fs::" + strings.ToLower(name)
id := ag.BeginLocalToolExecution(ctx, storedName, map[string]interface{}{})
id := ag.BeginLocalToolExecution(ctx, storedName, args)
if id == "" {
return
}
@@ -108,18 +120,21 @@ func recordEinoADKFilesystemToolMonitor(
msgs []adk.Message,
resultText string,
isErr bool,
) {
) string {
if ag == nil || rec == nil {
return
return ""
}
name := strings.TrimSpace(toolName)
if name == "" || strings.EqualFold(name, "execute") {
return
return ""
}
if !isBuiltinEinoADKFilesystemToolName(name) {
return
return ""
}
args := toolCallArgsFromAccumulated(msgs, toolCallID, name)
if len(args) == 0 && binder != nil {
args = binder.Arguments(toolCallID)
}
storedName := "eino_fs::" + strings.ToLower(name)
var invErr error
if isErr {
@@ -138,4 +153,5 @@ func recordEinoADKFilesystemToolMonitor(
if id != "" && execID == "" {
rec(id, toolCallID)
}
return id
}
@@ -2,6 +2,7 @@ package multiagent
import (
"context"
"strings"
"testing"
"cyberstrike-ai/internal/agent"
@@ -10,6 +11,7 @@ import (
"cyberstrike-ai/internal/mcp"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/components/tool"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
@@ -26,7 +28,7 @@ func TestEinoADKFilesystemToolMonitorBindsFinishesAndUpdatesDisplayResult(t *tes
recorded = append(recorded, executionID+"|"+toolCallID)
})
beginEinoADKFilesystemToolMonitor(ctx, ag, rec, binder, "call-read", "read_file")
beginEinoADKFilesystemToolMonitor(ctx, ag, rec, binder, "call-read", "read_file", map[string]interface{}{"path": "/tmp/secret.txt"})
execID := binder.ExecutionID("call-read")
if execID == "" {
t.Fatal("expected begin to bind execution id")
@@ -38,6 +40,9 @@ func TestEinoADKFilesystemToolMonitorBindsFinishesAndUpdatesDisplayResult(t *tes
if len(recorded) != 1 || recorded[0] != execID+"|call-read" {
t.Fatalf("recorded begin ids = %#v", recorded)
}
if got, _ := exec.Arguments["path"].(string); got != "/tmp/secret.txt" {
t.Fatalf("begin execution args = %#v", exec.Arguments)
}
runMessages := newEinoRunMessageAccumulator([]adk.Message{
&schema.Message{
@@ -80,3 +85,86 @@ func TestEinoADKFilesystemToolMonitorBindsFinishesAndUpdatesDisplayResult(t *tes
t.Fatalf("finish should reuse existing execution without recording a second id, got %#v", recorded)
}
}
func TestEinoADKFilesystemToolMonitorSpillsLargeReadFileResultForProgress(t *testing.T) {
t.Parallel()
ctx := context.Background()
logger := zap.NewNop()
server := mcp.NewServer(logger)
server.ConfigureToolResultMaxBytes(400)
server.ConfigureToolResultSpillRoot(t.TempDir())
ag := agent.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 1)
binder := NewMCPExecutionBinder()
rec := einomcp.ExecutionRecorder(func(executionID, toolCallID string) {})
var event map[string]interface{}
runMessages := newEinoRunMessageAccumulator([]adk.Message{
&schema.Message{
Role: schema.Assistant,
ToolCalls: []schema.ToolCall{{
ID: "call-read",
Type: "function",
Function: schema.FunctionCall{
Name: "read_file",
Arguments: `{"path":"/tmp/large.txt"}`,
},
}},
},
})
beginEinoADKFilesystemToolMonitor(ctx, ag, rec, binder, "call-read", "read_file", map[string]interface{}{"path": "/tmp/large.txt"})
emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{
ConversationID: "conv-1",
RunMessages: runMessages,
FilesystemMonitorAgent: ag,
FilesystemMonitorRecord: rec,
MCPExecutionBinder: binder,
Progress: func(eventType, _ string, data interface{}) {
if eventType == "tool_result" {
event, _ = data.(map[string]interface{})
}
},
})
if !emitter.Emit(ctx, "read_file", strings.Repeat("0123456789", 100), "call-read", false, "lead") {
t.Fatal("expected tool result emit")
}
result, _ := event["result"].(string)
if !strings.Contains(result, "<persisted-output>") || !strings.Contains(result, "Full output saved to:") {
t.Fatalf("large read_file result was not spilled in progress event: %q", result)
}
if len(result) > 400 {
t.Fatalf("progress result exceeded configured max: len=%d text=%q", len(result), result)
}
}
func TestEinoAgenticFilesystemWrapperCapturesArgumentsAndSpillsResult(t *testing.T) {
t.Parallel()
binder := NewMCPExecutionBinder()
mw := &einoAgenticFilesystemToolMiddleware{
TypedChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
conversationID: "conv-1",
toolMaxBytes: 400,
reductionRootDir: t.TempDir(),
binder: binder,
}
endpoint, err := mw.WrapInvokableToolCall(context.Background(), func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
return strings.Repeat("0123456789", 100), nil
}, &adk.ToolContext{Name: "read_file", CallID: "call-read"})
if err != nil {
t.Fatalf("WrapInvokableToolCall: %v", err)
}
result, err := endpoint(context.Background(), `{"file_path":"/tmp/requirements.txt","limit":2000}`)
if err != nil {
t.Fatalf("endpoint: %v", err)
}
args := binder.Arguments("call-read")
if args["file_path"] != "/tmp/requirements.txt" {
t.Fatalf("captured args = %#v", args)
}
if !strings.Contains(result, "<persisted-output>") || !strings.Contains(result, "Full output saved to:") {
t.Fatalf("expected persisted-output summary, got %q", result)
}
if len(result) > 400 {
t.Fatalf("summary exceeded max bytes: len=%d", len(result))
}
}
+44 -6
View File
@@ -11,6 +11,7 @@ import (
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/llm"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/reasoning"
@@ -49,7 +50,7 @@ func newEinoBaseHTTPClient() *http.Client {
}
}
func newEinoOpenAIChatModelFactory(
func newEinoToolCallingChatModelFactory(
baseHTTPClient *http.Client,
reasoningClient *reasoning.ClientIntent,
logger *zap.Logger,
@@ -58,6 +59,13 @@ func newEinoOpenAIChatModelFactory(
baseHTTPClient = newEinoBaseHTTPClient()
}
return func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.ToolCallingChatModel, error) {
if isEinoAgenticClaudeProvider(oa.Provider) {
nativeModel, err := newEinoClaudeAgenticChatModel(ctx, oa, mode, baseHTTPClient, reasoningClient)
if err != nil {
return nil, err
}
return newAgenticToolCallingChatModelAdapter(nativeModel), nil
}
httpClient := openai.NewEinoHTTPClient(&oa, baseHTTPClient)
openai.AttachSummarizationDiagTransport(httpClient, logger)
maxCompletionTokens := oa.MaxCompletionTokensEffective()
@@ -81,7 +89,7 @@ func newEinoOpenAIChatModelFactory(
}
}
func newEinoOpenAIAgenticChatModelFactory(
func newEinoAgenticChatModelFactory(
baseHTTPClient *http.Client,
reasoningClient *reasoning.ClientIntent,
logger *zap.Logger,
@@ -90,8 +98,11 @@ func newEinoOpenAIAgenticChatModelFactory(
baseHTTPClient = newEinoBaseHTTPClient()
}
return func(ctx context.Context, oa config.OpenAIConfig, mode einoModelMode) (model.AgenticModel, error) {
if !supportsEinoAgenticOpenAIBackend(oa) {
return nil, fmt.Errorf("eino agentic model: provider %q is not enabled for agenticopenai backend", strings.TrimSpace(oa.Provider))
if !supportsEinoAgenticBackend(oa) {
return nil, fmt.Errorf("eino agentic model: provider %q is not supported", strings.TrimSpace(oa.Provider))
}
if isEinoAgenticClaudeProvider(oa.Provider) {
return newEinoClaudeAgenticChatModel(ctx, oa, mode, baseHTTPClient, reasoningClient)
}
httpClient := openai.NewEinoHTTPClient(&oa, baseHTTPClient)
openai.AttachSummarizationDiagTransport(httpClient, logger)
@@ -111,9 +122,36 @@ func newEinoOpenAIAgenticChatModelFactory(
}
}
func supportsEinoAgenticOpenAIBackend(oa config.OpenAIConfig) bool {
func newEinoClaudeAgenticChatModel(
ctx context.Context,
oa config.OpenAIConfig,
mode einoModelMode,
httpClient *http.Client,
reasoningClient *reasoning.ClientIntent,
) (model.AgenticModel, error) {
extraFields := reasoning.AgenticOpenAIExtraFields(&oa, reasoningClient)
if mode == einoModelModePlanner {
extraFields = reasoning.AgenticOpenAIPlannerExtraFields(&oa)
}
return llm.NewClaudeAgenticModel(
ctx,
oa,
httpClient,
oa.MaxCompletionTokensEffective(),
extraFields,
)
}
func supportsEinoAgenticBackend(oa config.OpenAIConfig) bool {
provider := strings.ToLower(strings.TrimSpace(oa.Provider))
return provider == "" || provider == "openai" || provider == "openai_compatible"
return provider == "" ||
provider == "openai" ||
provider == "openai_compatible" ||
isEinoAgenticClaudeProvider(provider)
}
func isEinoAgenticClaudeProvider(provider string) bool {
return llm.IsClaudeProvider(provider)
}
func agenticModelGateFactory(factory einoAgenticModelConfigFactory, oa config.OpenAIConfig, mode einoModelMode) einoAgenticModelFactory {

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