Compare commits

..
16 Commits
Author SHA1 Message Date
公明andGitHub 6529c13ffb Update config.example.yaml 2026-07-11 11:53:24 +08:00
公明andGitHub f2ba5093d7 Add files via upload 2026-07-11 11:47:26 +08:00
公明andGitHub 3c6cf633e1 Add files via upload 2026-07-11 11:44:25 +08:00
公明andGitHub 142977413e Add files via upload 2026-07-11 11:42:25 +08:00
公明andGitHub 62efc81993 Add files via upload 2026-07-11 11:41:00 +08:00
公明andGitHub 9ac5fd33ec Add files via upload 2026-07-11 11:40:09 +08:00
公明andGitHub d1d67b07d3 Add files via upload 2026-07-11 11:38:14 +08:00
公明andGitHub 211c36654a Add files via upload 2026-07-11 11:36:47 +08:00
公明andGitHub 3894ba6054 Add files via upload 2026-07-11 11:34:41 +08:00
公明andGitHub fa0dd6c721 Add files via upload 2026-07-11 11:32:44 +08:00
公明andGitHub 1cf10981ef Add files via upload 2026-07-11 11:30:50 +08:00
公明andGitHub 79d162ccb4 Add files via upload 2026-07-11 11:21:15 +08:00
公明andGitHub 722806797f Add files via upload 2026-07-11 10:47:31 +08:00
公明andGitHub 2fcda3a57d Add files via upload 2026-07-11 10:47:04 +08:00
公明andGitHub bec296ae3f Delete internal/multiagent directory 2026-07-11 10:45:21 +08:00
612015b6d7 修复工作流检查点序列化 (#193)
Co-authored-by: ruanmingchen <“ruanm@chenchen”>
2026-07-11 01:38:10 +08:00
55 changed files with 1309 additions and 637 deletions
+2 -3
View File
@@ -212,7 +212,7 @@ The `run.sh` script will automatically:
model: "gpt-4o" # or deepseek-chat, claude-3-opus, etc. model: "gpt-4o" # or deepseek-chat, claude-3-opus, etc.
``` ```
- Or edit `config.yaml` directly before launching - Or edit `config.yaml` directly before launching
2. **Login** - Use the auto-generated password shown in the console (or set `auth.password` in `config.yaml`) 2. **Login** - On first startup the console prints an auto-generated initial `admin` password; create accounts from **Platform permissions → User management**
3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples: 3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples:
**macOS (Homebrew):** **macOS (Homebrew):**
@@ -281,7 +281,7 @@ Requirements / tips:
### Built-in Safeguards ### Built-in Safeguards
- Required-field validation prevents accidental blank API credentials. - Required-field validation prevents accidental blank API credentials.
- Auto-generated strong passwords when `auth.password` is empty. - Auto-generated 24-character initial `admin` password on first startup when no RBAC users exist (stored in the database only, not in `config.yaml`).
- Unified auth middleware for every web/API call (Bearer token flow). - Unified auth middleware for every web/API call (Bearer token flow).
- Timeout and sandbox guards per tool, plus structured logging for triage. - Timeout and sandbox guards per tool, plus structured logging for triage.
@@ -537,7 +537,6 @@ A test SSE MCP server is available at `cmd/test-sse-mcp-server/` for validation
```yaml ```yaml
auth: auth:
password: "change-me"
session_duration_hours: 12 session_duration_hours: 12
server: server:
host: "0.0.0.0" host: "0.0.0.0"
+2 -3
View File
@@ -211,7 +211,7 @@ chmod +x run.sh && ./run.sh
model: "gpt-4o" # 或 deepseek-chat, claude-3-opus 等 model: "gpt-4o" # 或 deepseek-chat, claude-3-opus 等
``` ```
- 或启动前直接编辑 `config.yaml` 文件 - 或启动前直接编辑 `config.yaml` 文件
2. **登录系统** - 使用控制台显示自动生成密码(或在 `config.yaml` 中设置 `auth.password` 2. **登录系统** - 首次启动时控制台显示自动生成的 `admin` 初始密码;也可在「平台权限 → 用户管理」中创建账号
3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例: 3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例:
**macOSHomebrew):** **macOSHomebrew):**
@@ -279,7 +279,7 @@ go build -o cyberstrike-ai cmd/server/main.go
### 默认安全措施 ### 默认安全措施
- 设置面板内置必填校验,防止漏配 API Key/Base URL/模型。 - 设置面板内置必填校验,防止漏配 API Key/Base URL/模型。
- `auth.password` 为空时自动生成 24 位强口令并写回 `config.yaml`。 - 首次启动且无 RBAC 用户时,自动生成 24 位 `admin` 初始密码并在控制台输出(仅存于数据库,不再写入 `config.yaml`
- 所有 API(除登录外)都需携带 Bearer Token,统一鉴权中间件拦截。 - 所有 API(除登录外)都需携带 Bearer Token,统一鉴权中间件拦截。
- 每个工具执行都带有超时、日志和错误隔离。 - 每个工具执行都带有超时、日志和错误隔离。
@@ -535,7 +535,6 @@ CyberStrikeAI 支持通过三种传输模式连接外部 MCP 服务器:
```yaml ```yaml
auth: auth:
password: "change-me"
session_duration_hours: 12 session_duration_hours: 12
server: server:
host: "0.0.0.0" host: "0.0.0.0"
+8 -11
View File
@@ -5,6 +5,7 @@ import (
"cyberstrike-ai/internal/app" "cyberstrike-ai/internal/app"
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/logger" "cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/termout"
"flag" "flag"
"fmt" "fmt"
"os" "os"
@@ -47,8 +48,7 @@ func main() {
return return
} }
if localConfig.Created { if localConfig.Created {
cfg.Auth.GeneratedPassword = localConfig.GeneratedPassword termout.PrintConfigCreated()
cfg.Auth.GeneratedPasswordPersisted = true
} }
if *httpsBootstrap { if *httpsBootstrap {
@@ -63,15 +63,12 @@ func main() {
if config.MainWebUIUsesHTTPS(&cfg.Server) { if config.MainWebUIUsesHTTPS(&cfg.Server) {
scheme = "https" scheme = "https"
} }
fmt.Println() termout.PrintStartupWebUI(termout.StartupWebUIOptions{
fmt.Printf("→ Web 界面: %s://127.0.0.1:%d/\n", scheme, port) Scheme: scheme,
if scheme == "https" && cfg.Server.TLSAutoSelfSign { Port: port,
fmt.Println(" (内存自签证书:浏览器首次需确认「继续访问」)") SelfSigned: scheme == "https" && cfg.Server.TLSAutoSelfSign,
} HTTPRedirect: scheme == "https" && config.ServerHTTPRedirectEnabled(&cfg.Server),
if scheme == "https" && config.ServerHTTPRedirectEnabled(&cfg.Server) { })
fmt.Printf(" http://127.0.0.1:%d/ 将自动跳转到 HTTPS\n", port)
}
fmt.Println()
// MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置 // MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置
if err := config.EnsureMCPAuth(cp, cfg); err != nil { if err := config.EnsureMCPAuth(cp, cfg); err != nil {
+2 -2
View File
@@ -10,7 +10,7 @@
# ============================================ # ============================================
# 前端显示的版本号(可选,不填则显示默认版本) # 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.7.0" version: "v1.7.1"
# 服务器配置 # 服务器配置
server: server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口 host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -28,7 +28,6 @@ server:
tls_auto_self_sign: true tls_auto_self_sign: true
# 认证配置 # 认证配置
auth: auth:
password: # Web 登录密码,请修改为强密码
session_duration_hours: 12 # 登录有效期(小时),超时后需重新登录 session_duration_hours: 12 # 登录有效期(小时),超时后需重新登录
# 日志配置 # 日志配置
log: log:
@@ -210,6 +209,7 @@ multi_agent:
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写 reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
reduction_sub_agents: true # true:子代理也挂 reductionfalse:仅编排主代理使用 reduction reduction_sub_agents: true # true:子代理也挂 reductionfalse:仅编排主代理使用 reduction
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85 summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
summarization_output_reserve_tokens: 8192 # 摘要模型输出预留 token;摘要输入预算 = 触发阈值 - 该值
summarization_emit_internal_events: true # true:发出 summarization 内部事件(便于诊断) summarization_emit_internal_events: true # true:发出 summarization 内部事件(便于诊断)
summarization_user_intent_ledger_max_runes: 96000 # 压缩后注入模型上下文的「原始用户输入与约束账本」总字符上限;DB 原始消息不裁剪 summarization_user_intent_ledger_max_runes: 96000 # 压缩后注入模型上下文的「原始用户输入与约束账本」总字符上限;DB 原始消息不裁剪
summarization_user_intent_ledger_entry_max_runes: 16000 # 账本中单条用户消息的字符上限;超出仅裁剪模型可见账本,不影响 DB 原文 summarization_user_intent_ledger_entry_max_runes: 16000 # 账本中单条用户消息的字符上限;超出仅裁剪模型可见账本,不影响 DB 原文
+2 -3
View File
@@ -21,7 +21,7 @@ server:
tls_enabled: true tls_enabled: true
tls_auto_self_sign: true tls_auto_self_sign: true
auth: auth:
password: "dev-only-change-me" session_duration_hours: 12
audit: audit:
enabled: true enabled: true
retention_days: 7 retention_days: 7
@@ -45,7 +45,7 @@ server:
port: 8080 port: 8080
tls_enabled: false tls_enabled: false
auth: auth:
password: "<long-random-password>" session_duration_hours: 12
audit: audit:
enabled: true enabled: true
retention_days: 30 retention_days: 30
@@ -90,7 +90,6 @@ Goal: long-running production red-team or security platform.
```yaml ```yaml
auth: auth:
password: "<managed-secret>"
session_duration_hours: 8 session_duration_hours: 8
audit: audit:
enabled: true enabled: true
+1 -2
View File
@@ -12,7 +12,6 @@ server:
port: 8080 port: 8080
tls_enabled: true tls_enabled: true
auth: auth:
password: "change-me"
session_duration_hours: 12 session_duration_hours: 12
openai: openai:
provider: openai provider: openai
@@ -24,7 +23,7 @@ agent:
tool_timeout_minutes: 60 tool_timeout_minutes: 60
``` ```
Change the default password immediately. Use HTTPS or a trusted reverse proxy in any shared environment. Change the initial `admin` password from the Web UI after first login. Use HTTPS or a trusted reverse proxy in any shared environment.
## Hot-Apply Boundaries ## Hot-Apply Boundaries
+1 -1
View File
@@ -48,7 +48,7 @@ config.yaml
```yaml ```yaml
auth: auth:
password: "<long-random-password>" session_duration_hours: 12
server: server:
host: 127.0.0.1 host: 127.0.0.1
port: 8080 port: 8080
+1 -1
View File
@@ -6,7 +6,7 @@ This checklist covers pre-production and continuous hardening for CyberStrikeAI.
## Before Going Live ## Before Going Live
- Change `auth.password` to a long random secret. - Change the initial `admin` password from the Web UI after first login.
- Use HTTPS or a trusted reverse proxy. - Use HTTPS or a trusted reverse proxy.
- Restrict access by IP, VPN, or bastion. - Restrict access by IP, VPN, or bastion.
- Enable `audit.enabled`. - Enable `audit.enabled`.
+1 -1
View File
@@ -37,7 +37,7 @@ Page inaccessible:
Login fails: Login fails:
- wrong `auth.password`; - wrong RBAC user password;
- config not applied/restarted; - config not applied/restarted;
- stale cookie; - stale cookie;
- audit throttling repeated failures. - audit throttling repeated failures.
+2 -3
View File
@@ -21,7 +21,7 @@ server:
tls_enabled: true tls_enabled: true
tls_auto_self_sign: true tls_auto_self_sign: true
auth: auth:
password: "dev-only-change-me" session_duration_hours: 12
audit: audit:
enabled: true enabled: true
retention_days: 7 retention_days: 7
@@ -54,7 +54,7 @@ server:
port: 8080 port: 8080
tls_enabled: false tls_enabled: false
auth: auth:
password: "<long-random-password>" session_duration_hours: 12
audit: audit:
enabled: true enabled: true
retention_days: 30 retention_days: 30
@@ -107,7 +107,6 @@ multi_agent:
```yaml ```yaml
auth: auth:
password: "<managed-secret>"
session_duration_hours: 8 session_duration_hours: 8
audit: audit:
enabled: true enabled: true
+1 -2
View File
@@ -12,7 +12,6 @@ server:
tls_enabled: true tls_enabled: true
tls_auto_self_sign: true tls_auto_self_sign: true
auth: auth:
password: "change-me"
session_duration_hours: 12 session_duration_hours: 12
log: log:
level: info level: info
@@ -22,7 +21,7 @@ log:
- `version`:前端展示版本。 - `version`:前端展示版本。
- `server.host/port`Web 服务监听地址和端口。 - `server.host/port`Web 服务监听地址和端口。
- `server.tls_*`HTTPS 配置。生产环境建议使用 `tls_cert_path``tls_key_path` - `server.tls_*`HTTPS 配置。生产环境建议使用 `tls_cert_path``tls_key_path`
- `auth.password`Web 登录密码,必须改为强密码。 - `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
- `auth.session_duration_hours`:登录会话有效期。 - `auth.session_duration_hours`:登录会话有效期。
- `log.output`:可以是 `stdout``stderr` 或文件路径。 - `log.output`:可以是 `stdout``stderr` 或文件路径。
+1 -1
View File
@@ -48,7 +48,7 @@ config.yaml
```yaml ```yaml
auth: auth:
password: "<long-random-password>" session_duration_hours: 12
server: server:
host: 127.0.0.1 host: 127.0.0.1
port: 8080 port: 8080
+1 -1
View File
@@ -6,7 +6,7 @@
## 上线前必做 ## 上线前必做
- 修改 `auth.password` 为长随机密码 - 首次部署后立即修改 `admin` 初始密码(Web 界面或平台权限 → 用户管理)
- 使用 HTTPS,或放在可信反向代理之后。 - 使用 HTTPS,或放在可信反向代理之后。
- 限制来源 IP、VPN 或堡垒机访问。 - 限制来源 IP、VPN 或堡垒机访问。
- 开启 `audit.enabled` - 开启 `audit.enabled`
+2 -2
View File
@@ -16,9 +16,9 @@ CyberStrikeAI 面向授权安全测试场景,内置命令执行、MCP 工具
## 认证与会话 ## 认证与会话
`auth.password` 是 Web 登录密码。建议: Web 登录凭据由 RBAC 用户管理(默认内置 `admin` 账号)。建议:
- 首次部署立即修改默认密码 - 首次部署立即修改 `admin` 初始密码(控制台首次启动会输出)
- 使用长随机密码,并限制分享范围。 - 使用长随机密码,并限制分享范围。
- 将服务放在内网、VPN、堡垒机或反向代理认证后面。 - 将服务放在内网、VPN、堡垒机或反向代理认证后面。
- 生产环境开启 HTTPS,避免明文传输 Cookie。 - 生产环境开启 HTTPS,避免明文传输 Cookie。
+3 -3
View File
@@ -23,12 +23,12 @@ https://127.0.0.1:8080/
检查: 检查:
- `config.yaml` 中的 `auth.password` - RBAC 用户密码是否正确(默认 `admin`;首次启动密码见控制台输出)
- 是否修改后未重启或未应用配置 - 是否修改密码后旧会话已失效,需重新登录
- 浏览器 Cookie 是否异常,可尝试无痕窗口。 - 浏览器 Cookie 是否异常,可尝试无痕窗口。
- 审计日志中是否有登录失败节流。 - 审计日志中是否有登录失败节流。
生产环境忘记密码时,需在服务器上修改 `config.yaml` 并重启服务 生产环境忘记密码时,需在服务器上通过 RBAC 用户管理重置,或直接更新数据库中的用户密码哈希
## 模型无响应 ## 模型无响应
+1 -1
View File
@@ -661,7 +661,7 @@ func (a *Agent) UpdateToolDescriptionMode(mode string) {
mode = "short" mode = "short"
} }
a.toolDescriptionMode = mode a.toolDescriptionMode = mode
a.logger.Info("Agent工具描述模式已更新", zap.String("tool_description_mode", mode)) a.logger.Debug("Agent工具描述模式已更新", zap.String("tool_description_mode", mode))
} }
// RepairOrphanToolMessages 清理失去配对的tool消息和未完成的tool_calls,避免OpenAI报错 // RepairOrphanToolMessages 清理失去配对的tool消息和未完成的tool_calls,避免OpenAI报错
+17 -22
View File
@@ -101,13 +101,12 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
return nil, fmt.Errorf("初始化数据库失败: %w", err) return nil, fmt.Errorf("初始化数据库失败: %w", err)
} }
// 认证管理器(数据库初始化后挂载 RBAC,以兼容旧的单密码配置 // 认证管理器(数据库初始化后挂载 RBAC)
authManager, err := security.NewAuthManager(cfg.Auth.Password, cfg.Auth.SessionDurationHours) authManager := security.NewAuthManager(cfg.Auth.SessionDurationHours)
if err != nil { if generatedPassword, err := authManager.AttachRBACStore(db); err != nil {
return nil, fmt.Errorf("初始化认证失败: %w", err)
}
if err := authManager.AttachRBACStore(db); err != nil {
return nil, fmt.Errorf("初始化RBAC失败: %w", err) return nil, fmt.Errorf("初始化RBAC失败: %w", err)
} else if generatedPassword != "" {
config.PrintBootstrapAdminPassword(generatedPassword)
} }
for platform, userID := range cfg.Robots.ServiceAccountUserIDs() { for platform, userID := range cfg.Robots.ServiceAccountUserIDs() {
user, userErr := db.GetRBACUserByID(userID) user, userErr := db.GetRBACUserByID(userID)
@@ -125,6 +124,9 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
monitorRetention.PurgeExpired() monitorRetention.PurgeExpired()
monitor.StartRetentionLoop(monitorRetention, log.Logger) monitor.StartRetentionLoop(monitorRetention, log.Logger)
if err := handler.NewHITLManager(db, log.Logger).EnsureSchema(); err != nil {
log.Logger.Warn("初始化 HITL 表失败", zap.Error(err))
}
hitlRetention := hitl.NewService(db, cfg, log.Logger) hitlRetention := hitl.NewService(db, cfg, log.Logger)
hitlRetention.PurgeExpired() hitlRetention.PurgeExpired()
hitl.StartRetentionLoop(hitlRetention, log.Logger) hitl.StartRetentionLoop(hitlRetention, log.Logger)
@@ -146,13 +148,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
registerProjectFactTools(mcpServer, db, cfg, log.Logger) registerProjectFactTools(mcpServer, db, cfg, log.Logger)
registerVisionTools(mcpServer, cfg, log.Logger) registerVisionTools(mcpServer, cfg, log.Logger)
if cfg.Auth.GeneratedPassword != "" {
config.PrintGeneratedPasswordWarning(cfg.Auth.GeneratedPassword, cfg.Auth.GeneratedPasswordPersisted, cfg.Auth.GeneratedPasswordPersistErr)
cfg.Auth.GeneratedPassword = ""
cfg.Auth.GeneratedPasswordPersisted = false
cfg.Auth.GeneratedPasswordPersistErr = ""
}
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储) // 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db) externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer()) externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
@@ -181,7 +176,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
var knowledgeHandler *handler.KnowledgeHandler var knowledgeHandler *handler.KnowledgeHandler
var knowledgeDBConn *database.DB var knowledgeDBConn *database.DB
log.Logger.Info("检查知识库配置", zap.Bool("enabled", cfg.Knowledge.Enabled)) log.Logger.Debug("检查知识库配置", zap.Bool("enabled", cfg.Knowledge.Enabled))
if cfg.Knowledge.Enabled { if cfg.Knowledge.Enabled {
// 确定知识库数据库路径 // 确定知识库数据库路径
knowledgeDBPath := cfg.Database.KnowledgeDBPath knowledgeDBPath := cfg.Database.KnowledgeDBPath
@@ -324,7 +319,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
} }
skillsDir := skillpackage.SkillsRootFromConfig(cfg.SkillsDir, configPath) skillsDir := skillpackage.SkillsRootFromConfig(cfg.SkillsDir, configPath)
log.Logger.Info("Skills 目录(Eino ADK skill 中间件 + Web 管理 API", zap.String("skillsDir", skillsDir)) log.Logger.Debug("Skills 目录(Eino ADK skill 中间件 + Web 管理 API", zap.String("skillsDir", skillsDir))
configDir := filepath.Dir(configPath) configDir := filepath.Dir(configPath)
plantaskRel := strings.TrimSpace(cfg.MultiAgent.EinoMiddleware.PlantaskRelDir) plantaskRel := strings.TrimSpace(cfg.MultiAgent.EinoMiddleware.PlantaskRelDir)
if plantaskRel == "" { if plantaskRel == "" {
@@ -350,7 +345,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
} }
markdownAgentsHandler := handler.NewMarkdownAgentsHandler(agentsDir) markdownAgentsHandler := handler.NewMarkdownAgentsHandler(agentsDir)
markdownAgentsHandler.SetAudit(auditSvc) markdownAgentsHandler.SetAudit(auditSvc)
log.Logger.Info("多代理 Markdown 子 Agent 目录", zap.String("agentsDir", agentsDir)) log.Logger.Debug("多代理 Markdown 子 Agent 目录", zap.String("agentsDir", agentsDir))
// 创建处理器 // 创建处理器
agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger) agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger)
@@ -635,20 +630,20 @@ func (a *App) RunWithContext(ctx context.Context) error {
} }
switch tlsMode { switch tlsMode {
case mainTLSFromFiles: case mainTLSFromFiles:
a.logger.Info("启动 HTTPS 主服务(已启用 HTTP/2 协商)", a.logger.Debug("启动 HTTPS 主服务(已启用 HTTP/2 协商)",
zap.String("address", addr), zap.String("address", addr),
zap.String("cert", certFile), zap.String("cert", certFile),
) )
case mainTLSInMemorySelfSigned: case mainTLSInMemorySelfSigned:
a.logger.Info("启动 HTTPS 主服务(内存自签证书,仅测试;已启用 HTTP/2 协商)", a.logger.Debug("启动 HTTPS 主服务(内存自签证书,仅测试;已启用 HTTP/2 协商)",
zap.String("address", addr), zap.String("address", addr),
) )
} }
if httpRedirect { if httpRedirect {
a.logger.Info("已启用 HTTP→HTTPS 自动跳转(同端口嗅探分流)", zap.String("address", addr)) a.logger.Debug("已启用 HTTP→HTTPS 自动跳转(同端口嗅探分流)", zap.String("address", addr))
} }
} else { } else {
a.logger.Info("启动 HTTP 主服务", zap.String("address", addr)) a.logger.Debug("启动 HTTP 主服务", zap.String("address", addr))
} }
// 监听 context 取消,优雅关闭 HTTP 服务器 // 监听 context 取消,优雅关闭 HTTP 服务器
@@ -1508,7 +1503,7 @@ func registerWebshellTools(mcpServer *mcp.Server, db *database.DB, webshellHandl
} }
mcpServer.RegisterTool(writeTool, writeHandler) mcpServer.RegisterTool(writeTool, writeHandler)
logger.Info("WebShell 工具注册成功") logger.Debug("WebShell 工具注册成功")
} }
// registerWebshellManagementTools 注册 WebShell 连接管理 MCP 工具 // registerWebshellManagementTools 注册 WebShell 连接管理 MCP 工具
@@ -1879,7 +1874,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
} }
mcpServer.RegisterTool(testTool, testHandler) mcpServer.RegisterTool(testTool, testHandler)
logger.Info("WebShell 管理工具注册成功") logger.Debug("WebShell 管理工具注册成功")
} }
// initializeKnowledge 初始化知识库组件(用于动态初始化) // initializeKnowledge 初始化知识库组件(用于动态初始化)
+1 -1
View File
@@ -31,7 +31,7 @@ func registerC2Tools(mcpServer *mcp.Server, c2Manager *c2.Manager, logger *zap.L
registerC2EventTool(mcpServer, c2Manager, logger) registerC2EventTool(mcpServer, c2Manager, logger)
registerC2ProfileTool(mcpServer, c2Manager, logger) registerC2ProfileTool(mcpServer, c2Manager, logger)
registerC2FileTool(mcpServer, c2Manager, logger) registerC2FileTool(mcpServer, c2Manager, logger)
logger.Info("C2 MCP tools registered (8 unified tools)") logger.Debug("C2 MCP tools registered (8 unified tools)")
} }
func makeC2Result(data interface{}, err error) (*mcp.ToolResult, error) { func makeC2Result(data interface{}, err error) (*mcp.ToolResult, error) {
+6 -2
View File
@@ -21,11 +21,15 @@ func TestStandaloneMCPPrefersUserRBACAndDisablesGlobalTokenByDefault(t *testing.
t.Fatal(err) t.Fatal(err)
} }
t.Cleanup(func() { _ = db.Close() }) t.Cleanup(func() { _ = db.Close() })
auth, err := security.NewAuthManager("admin-secret", 12) auth := security.NewAuthManager(12)
if _, err := auth.AttachRBACStore(db); err != nil {
t.Fatal(err)
}
hash, err := security.HashPassword("admin-secret")
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if err := auth.AttachRBACStore(db); err != nil { if err := db.UpdateRBACAdminPassword(hash); err != nil {
t.Fatal(err) t.Fatal(err)
} }
token, _, err := auth.Authenticate("admin", "admin-secret") token, _, err := auth.Authenticate("admin", "admin-secret")
+2 -2
View File
@@ -90,7 +90,7 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
"description": "可选:关联的漏洞记录 ID", "description": "可选:关联的漏洞记录 ID",
}, },
"links": map[string]interface{}{ "links": map[string]interface{}{
"type": "array", "type": "array",
"description": "可选:关系边(from → 当前 fact)。finding 至少 1 条 {from:target/*, type:discovered_on}finding 上记录 exploit 用 {from:exploit/*, type:exploits}。省略保留已有边;传 [] 清空全部关系边。", "description": "可选:关系边(from → 当前 fact)。finding 至少 1 条 {from:target/*, type:discovered_on}finding 上记录 exploit 用 {from:exploit/*, type:exploits}。省略保留已有边;传 [] 清空全部关系边。",
"items": map[string]interface{}{ "items": map[string]interface{}{
"type": "object", "type": "object",
@@ -357,7 +357,7 @@ func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *confi
}) })
if logger != nil { if logger != nil {
logger.Info("项目黑板 MCP 工具注册成功") logger.Debug("项目黑板 MCP 工具注册成功")
} }
} }
+1 -1
View File
@@ -190,7 +190,7 @@ func registerVulnerabilityTools(mcpServer *mcp.Server, db *database.DB, logger *
registerListVulnerabilitiesTool(mcpServer, db, logger) registerListVulnerabilitiesTool(mcpServer, db, logger)
registerGetVulnerabilityTool(mcpServer, db, logger) registerGetVulnerabilityTool(mcpServer, db, logger)
if logger != nil { if logger != nil {
logger.Info("漏洞 MCP 工具注册成功", zap.Strings("tools", []string{ logger.Debug("漏洞 MCP 工具注册成功", zap.Strings("tools", []string{
builtin.ToolRecordVulnerability, builtin.ToolRecordVulnerability,
builtin.ToolListVulnerabilities, builtin.ToolListVulnerabilities,
builtin.ToolGetVulnerability, builtin.ToolGetVulnerability,
+20 -198
View File
@@ -2,7 +2,6 @@ package config
import ( import (
"crypto/rand" "crypto/rand"
"encoding/base64"
"encoding/hex" "encoding/hex"
"encoding/json" "encoding/json"
"fmt" "fmt"
@@ -12,6 +11,8 @@ import (
"strconv" "strconv"
"strings" "strings"
"cyberstrike-ai/internal/termout"
"gopkg.in/yaml.v3" "gopkg.in/yaml.v3"
) )
@@ -43,9 +44,8 @@ type Config struct {
} }
type EnsureLocalConfigResult struct { type EnsureLocalConfigResult struct {
Created bool Created bool
GeneratedPassword string ExamplePath string
ExamplePath string
} }
const ( const (
@@ -54,6 +54,7 @@ const (
DefaultLatestUserMessageMaxRunes = 48000 DefaultLatestUserMessageMaxRunes = 48000
DefaultLatestUserMessageHeadRunes = 24000 DefaultLatestUserMessageHeadRunes = 24000
DefaultLatestUserMessageTailRunes = 24000 DefaultLatestUserMessageTailRunes = 24000
DefaultSummarizationOutputReserveTokens = 8192
) )
// ProjectConfig 项目黑板(跨对话共享事实)配置。 // ProjectConfig 项目黑板(跨对话共享事实)配置。
@@ -268,6 +269,8 @@ type MultiAgentEinoMiddlewareConfig struct {
ReductionSubAgents bool `yaml:"reduction_sub_agents,omitempty" json:"reduction_sub_agents,omitempty"` // also attach to sub-agents ReductionSubAgents bool `yaml:"reduction_sub_agents,omitempty" json:"reduction_sub_agents,omitempty"` // also attach to sub-agents
// SummarizationTriggerRatio controls summarization trigger threshold as max_total_tokens * ratio (default 0.8). // SummarizationTriggerRatio controls summarization trigger threshold as max_total_tokens * ratio (default 0.8).
SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"` SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"`
// SummarizationOutputReserveTokens reserves completion headroom for the summarization model call (default 8192).
SummarizationOutputReserveTokens int `yaml:"summarization_output_reserve_tokens,omitempty" json:"summarization_output_reserve_tokens,omitempty"`
// SummarizationEmitInternalEvents controls middleware internal event emission (default true). // SummarizationEmitInternalEvents controls middleware internal event emission (default true).
SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"` SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"`
// SummarizationUserIntentLedgerMaxRunes caps the DB-backed immutable user input ledger injected into model context. // SummarizationUserIntentLedgerMaxRunes caps the DB-backed immutable user input ledger injected into model context.
@@ -320,6 +323,13 @@ func (c MultiAgentEinoMiddlewareConfig) SummarizationTriggerRatioEffective() flo
return v return v
} }
func (c MultiAgentEinoMiddlewareConfig) SummarizationOutputReserveTokensEffective() int {
if c.SummarizationOutputReserveTokens > 0 {
return c.SummarizationOutputReserveTokens
}
return DefaultSummarizationOutputReserveTokens
}
func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool { func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool {
if c.SummarizationEmitInternalEvents != nil { if c.SummarizationEmitInternalEvents != nil {
return *c.SummarizationEmitInternalEvents return *c.SummarizationEmitInternalEvents
@@ -995,11 +1005,7 @@ func normalizeHitlModeForPrompt(mode string) string {
} }
type AuthConfig struct { type AuthConfig struct {
Password string `yaml:"password" json:"password"` SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"`
SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"`
GeneratedPassword string `yaml:"-" json:"-"`
GeneratedPasswordPersisted bool `yaml:"-" json:"-"`
GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
} }
// MonitorConfig MCP 状态监控(tool_executions)保留策略。 // MonitorConfig MCP 状态监控(tool_executions)保留策略。
@@ -1159,23 +1165,6 @@ func Load(path string) (*Config, error) {
if cfg.Audit.MaxDetailBytes <= 0 { if cfg.Audit.MaxDetailBytes <= 0 {
cfg.Audit.MaxDetailBytes = 8192 cfg.Audit.MaxDetailBytes = 8192
} }
if strings.TrimSpace(cfg.Auth.Password) == "" {
password, err := generateStrongPassword(24)
if err != nil {
return nil, fmt.Errorf("生成默认密码失败: %w", err)
}
cfg.Auth.Password = password
cfg.Auth.GeneratedPassword = password
if err := PersistAuthPassword(path, password); err != nil {
cfg.Auth.GeneratedPasswordPersisted = false
cfg.Auth.GeneratedPasswordPersistErr = err.Error()
} else {
cfg.Auth.GeneratedPasswordPersisted = true
}
}
// 如果配置了工具目录,从目录加载工具配置 // 如果配置了工具目录,从目录加载工具配置
if cfg.Security.ToolsDir != "" { if cfg.Security.ToolsDir != "" {
inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...) inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...)
@@ -1246,7 +1235,7 @@ func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) {
if _, err := os.Stat(path); err == nil { if _, err := os.Stat(path); err == nil {
return EnsureLocalConfigResult{}, nil return EnsureLocalConfigResult{}, nil
} else if err != nil && !os.IsNotExist(err) { } else if !os.IsNotExist(err) {
return EnsureLocalConfigResult{}, fmt.Errorf("检查配置文件失败: %w", err) return EnsureLocalConfigResult{}, fmt.Errorf("检查配置文件失败: %w", err)
} }
@@ -1281,181 +1270,14 @@ func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) {
return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err) return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err)
} }
password, err := generateStrongPassword(24)
if err != nil {
return EnsureLocalConfigResult{}, fmt.Errorf("生成默认密码失败: %w", err)
}
if err := PersistAuthPassword(path, password); err != nil {
return EnsureLocalConfigResult{}, fmt.Errorf("写入默认密码失败: %w", err)
}
return EnsureLocalConfigResult{ return EnsureLocalConfigResult{
Created: true, Created: true,
GeneratedPassword: password, ExamplePath: examplePath,
ExamplePath: examplePath,
}, nil }, nil
} }
func generateStrongPassword(length int) (string, error) { func PrintBootstrapAdminPassword(password string) {
if length <= 0 { termout.PrintBootstrapAdminCredentials(password)
length = 24
}
bytesLen := length
randomBytes := make([]byte, bytesLen)
if _, err := rand.Read(randomBytes); err != nil {
return "", err
}
password := base64.RawURLEncoding.EncodeToString(randomBytes)
if len(password) > length {
password = password[:length]
}
return password, nil
}
func PersistAuthPassword(path, password string) error {
data, err := os.ReadFile(path)
if err != nil {
return err
}
lines := strings.Split(string(data), "\n")
inAuthBlock := false
authIndent := -1
for i, line := range lines {
trimmed := strings.TrimSpace(line)
if !inAuthBlock {
if strings.HasPrefix(trimmed, "auth:") {
inAuthBlock = true
authIndent = len(line) - len(strings.TrimLeft(line, " "))
}
continue
}
if trimmed == "" || strings.HasPrefix(trimmed, "#") {
continue
}
leadingSpaces := len(line) - len(strings.TrimLeft(line, " "))
if leadingSpaces <= authIndent {
// 离开 auth 块
inAuthBlock = false
authIndent = -1
// 继续寻找其它 auth 块(理论上没有)
if strings.HasPrefix(trimmed, "auth:") {
inAuthBlock = true
authIndent = leadingSpaces
}
continue
}
if strings.HasPrefix(strings.TrimSpace(line), "password:") {
prefix := line[:len(line)-len(strings.TrimLeft(line, " "))]
comment := ""
if idx := yamlLineCommentIndex(line); idx >= 0 {
comment = strings.TrimRight(line[idx:], " ")
}
newLine := fmt.Sprintf("%spassword: %s", prefix, quoteYAMLString(password))
if comment != "" {
if !strings.HasPrefix(comment, " ") {
newLine += " "
}
newLine += comment
}
lines[i] = newLine
break
}
}
return os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0644)
}
func quoteYAMLString(value string) string {
node := yaml.Node{
Kind: yaml.ScalarNode,
Tag: "!!str",
Style: yaml.DoubleQuotedStyle,
Value: value,
}
data, err := yaml.Marshal(&node)
if err != nil {
return strconv.Quote(value)
}
return strings.TrimSuffix(string(data), "\n")
}
func yamlLineCommentIndex(line string) int {
inSingleQuote := false
inDoubleQuote := false
escaped := false
for i, r := range line {
if inDoubleQuote {
if escaped {
escaped = false
continue
}
if r == '\\' {
escaped = true
continue
}
if r == '"' {
inDoubleQuote = false
}
continue
}
if inSingleQuote {
if r == '\'' {
inSingleQuote = false
}
continue
}
switch r {
case '"':
inDoubleQuote = true
case '\'':
inSingleQuote = true
case '#':
if i == 0 || isYAMLWhitespace(line[i-1]) {
return i
}
}
}
return -1
}
func isYAMLWhitespace(b byte) bool {
return b == ' ' || b == '\t'
}
func PrintGeneratedPasswordWarning(password string, persisted bool, persistErr string) {
if strings.TrimSpace(password) == "" {
return
}
if persisted {
fmt.Println("[CyberStrikeAI] ✅ 已为您自动生成并写入 Web 登录密码。")
} else {
if persistErr != "" {
fmt.Printf("[CyberStrikeAI] ⚠️ 无法自动写入配置文件中的密码: %s\n", persistErr)
} else {
fmt.Println("[CyberStrikeAI] ⚠️ 无法自动写入配置文件中的密码。")
}
fmt.Println("请手动将以下随机密码写入 config.yaml 的 auth.password")
}
fmt.Println("----------------------------------------------------------------")
fmt.Println("CyberStrikeAI Auto-Generated Web Password")
fmt.Printf("Password: %s\n", password)
fmt.Println("WARNING: Anyone with this password can fully control CyberStrikeAI.")
fmt.Println("Please store it securely and change it in config.yaml as soon as possible.")
fmt.Println("警告:持有此密码的人将拥有对 CyberStrikeAI 的完全控制权限。")
fmt.Println("请妥善保管,并尽快在 config.yaml 中修改 auth.password")
fmt.Println("----------------------------------------------------------------")
} }
// generateRandomToken 生成用于 MCP 鉴权的随机字符串(64 位十六进制) // generateRandomToken 生成用于 MCP 鉴权的随机字符串(64 位十六进制)
+39 -75
View File
@@ -7,78 +7,12 @@ import (
"testing" "testing"
) )
func TestPersistAuthPasswordQuotesYAMLSpecialCharacters(t *testing.T) { func TestEnsureLocalConfigCreatesFromExample(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
initial := strings.Join([]string{
"server:",
" host: 0.0.0.0",
"auth:",
" password: old-password # Web 登录密码",
" session_duration_hours: 12",
"log:",
" level: info",
"",
}, "\n")
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
t.Fatalf("write config: %v", err)
}
want := `@abc:def # still password`
if err := PersistAuthPassword(path, want); err != nil {
t.Fatalf("PersistAuthPassword: %v", err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read config: %v", err)
}
if !strings.Contains(string(data), `password: "@abc:def # still password" # Web 登录密码`) {
t.Fatalf("password was not safely quoted or comment was not preserved:\n%s", data)
}
cfg, err := Load(path)
if err != nil {
t.Fatalf("Load after PersistAuthPassword: %v", err)
}
if cfg.Auth.Password != want {
t.Fatalf("Auth.Password = %q, want %q", cfg.Auth.Password, want)
}
}
func TestPersistAuthPasswordDoesNotTreatQuotedHashAsComment(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
initial := strings.Join([]string{
"auth:",
` password: "old#password"`,
" session_duration_hours: 12",
"",
}, "\n")
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
t.Fatalf("write config: %v", err)
}
if err := PersistAuthPassword(path, "new-password"); err != nil {
t.Fatalf("PersistAuthPassword: %v", err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read config: %v", err)
}
if strings.Contains(string(data), "#password") {
t.Fatalf("old quoted password fragment was incorrectly preserved as a comment:\n%s", data)
}
}
func TestEnsureLocalConfigCreatesFromExampleWithGeneratedPassword(t *testing.T) {
dir := t.TempDir() dir := t.TempDir()
examplePath := filepath.Join(dir, "config.example.yaml") examplePath := filepath.Join(dir, "config.example.yaml")
configPath := filepath.Join(dir, "config.yaml") configPath := filepath.Join(dir, "config.yaml")
example := []byte(`auth: example := []byte(`auth:
password: "change-me-use-a-long-random-password"
session_duration_hours: 12 session_duration_hours: 12
server: server:
host: 127.0.0.1 host: 127.0.0.1
@@ -95,9 +29,6 @@ server:
if !result.Created { if !result.Created {
t.Fatal("Created = false, want true") t.Fatal("Created = false, want true")
} }
if result.GeneratedPassword == "" {
t.Fatal("GeneratedPassword is empty")
}
if result.ExamplePath != examplePath { if result.ExamplePath != examplePath {
t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath) t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath)
} }
@@ -106,11 +37,8 @@ server:
if err != nil { if err != nil {
t.Fatalf("Load generated config: %v", err) t.Fatalf("Load generated config: %v", err)
} }
if cfg.Auth.Password == "change-me-use-a-long-random-password" { if cfg.Auth.SessionDurationHours != 12 {
t.Fatal("auth.password still contains the template placeholder") t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours)
}
if cfg.Auth.Password != result.GeneratedPassword {
t.Fatalf("Auth.Password = %q, want generated password %q", cfg.Auth.Password, result.GeneratedPassword)
} }
second, err := EnsureLocalConfig(configPath) second, err := EnsureLocalConfig(configPath)
@@ -122,6 +50,31 @@ server:
} }
} }
func TestLoadIgnoresLegacyAuthPasswordField(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
initial := strings.Join([]string{
"auth:",
` password: "legacy-password"`,
" session_duration_hours: 12",
"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.Auth.SessionDurationHours != 12 {
t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours)
}
}
func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) { func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
main := OpenAIConfig{ main := OpenAIConfig{
Provider: "openai", Provider: "openai",
@@ -163,6 +116,17 @@ func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
} }
} }
func TestSummarizationOutputReserveTokensEffective(t *testing.T) {
var zero MultiAgentEinoMiddlewareConfig
if got := zero.SummarizationOutputReserveTokensEffective(); got != DefaultSummarizationOutputReserveTokens {
t.Fatalf("default output reserve = %d, want %d", got, DefaultSummarizationOutputReserveTokens)
}
custom := MultiAgentEinoMiddlewareConfig{SummarizationOutputReserveTokens: 4096}
if got := custom.SummarizationOutputReserveTokensEffective(); got != 4096 {
t.Fatalf("custom output reserve = %d", got)
}
}
func TestLatestUserMessageRunesEffective(t *testing.T) { func TestLatestUserMessageRunesEffective(t *testing.T) {
var zero MultiAgentEinoMiddlewareConfig var zero MultiAgentEinoMiddlewareConfig
if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes { if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes {
+3 -3
View File
@@ -113,10 +113,10 @@ func (db *DB) runPassiveCheckpoint(trigger string) {
return return
} }
if busy > 0 { if busy > 0 {
db.logger.Info("SQLite PASSIVE checkpoint 完成(部分推进)", fields...) db.logger.Debug("SQLite PASSIVE checkpoint 完成(部分推进)", fields...)
return return
} }
db.logger.Info("SQLite PASSIVE checkpoint 完成(成功)", fields...) db.logger.Debug("SQLite PASSIVE checkpoint 完成(成功)", fields...)
} }
// NewDB 创建数据库连接 // NewDB 创建数据库连接
@@ -869,7 +869,7 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建索引失败: %w", err) return fmt.Errorf("创建索引失败: %w", err)
} }
db.logger.Info("数据库表初始化完成") db.logger.Debug("数据库表初始化完成")
return nil return nil
} }
+28 -4
View File
@@ -227,11 +227,32 @@ func (db *DB) addColumnIfMissing(table, name, stmt string) error {
return nil return nil
} }
// RBACNeedsAdminPassword reports whether the built-in admin account still needs an initial password.
func (db *DB) RBACNeedsAdminPassword() (bool, error) {
var userCount int
if err := db.QueryRow(`SELECT COUNT(*) FROM rbac_users`).Scan(&userCount); err != nil {
return false, err
}
if userCount == 0 {
return true, nil
}
var hash sql.NullString
err := db.QueryRow(`
SELECT password_hash FROM rbac_users
WHERE username = 'admin' AND is_builtin = 1
LIMIT 1
`).Scan(&hash)
if err == sql.ErrNoRows {
return false, nil
}
if err != nil {
return false, err
}
return !hash.Valid || strings.TrimSpace(hash.String) == "", nil
}
// BootstrapRBAC seeds the local admin account and system roles. // BootstrapRBAC seeds the local admin account and system roles.
func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]string) error { func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]string) error {
if strings.TrimSpace(adminPasswordHash) == "" {
return errors.New("admin password hash is required")
}
now := time.Now() now := time.Now()
tx, err := db.Begin() tx, err := db.Begin()
if err != nil { if err != nil {
@@ -302,6 +323,9 @@ func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]str
return err return err
} }
if userCount == 0 { if userCount == 0 {
if strings.TrimSpace(adminPasswordHash) == "" {
return errors.New("admin password hash is required for initial bootstrap")
}
if _, err := tx.Exec(` if _, err := tx.Exec(`
INSERT INTO rbac_users (id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at) INSERT INTO rbac_users (id, username, display_name, password_hash, enabled, is_builtin, created_at, updated_at)
VALUES (?, 'admin', '管理员', ?, 1, 1, ?, ?) VALUES (?, 'admin', '管理员', ?, 1, 1, ?, ?)
@@ -311,7 +335,7 @@ func (db *DB) BootstrapRBAC(adminPasswordHash string, permissions map[string]str
if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_user_roles (user_id, role_id, created_at) VALUES ('admin', ?, ?)`, RBACSystemRoleAdmin, now); err != nil { if _, err := tx.Exec(`INSERT OR IGNORE INTO rbac_user_roles (user_id, role_id, created_at) VALUES ('admin', ?, ?)`, RBACSystemRoleAdmin, now); err != nil {
return err return err
} }
} else { } else if strings.TrimSpace(adminPasswordHash) != "" {
if _, err := tx.Exec(`UPDATE rbac_users SET password_hash = ?, updated_at = ? WHERE username = 'admin' AND is_builtin = 1 AND (password_hash = '' OR password_hash IS NULL)`, adminPasswordHash, now); err != nil { if _, err := tx.Exec(`UPDATE rbac_users SET password_hash = ?, updated_at = ? WHERE username = 'admin' AND is_builtin = 1 AND (password_hash = '' OR password_hash IS NULL)`, adminPasswordHash, now); err != nil {
return err return err
} }
+4 -22
View File
@@ -170,28 +170,10 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) {
return return
} }
if session.UserID == "" || session.UserID == "admin" { if session.UserID == "" {
if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil { session.UserID = "admin"
if h.logger != nil { }
h.logger.Error("保存新密码失败", zap.Error(err)) if err := h.manager.UpdateUserPassword(session.UserID, newPassword); err != nil {
}
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败,请重试"})
return
}
if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil {
if h.logger != nil {
h.logger.Error("更新认证配置失败", zap.Error(err))
}
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"})
return
}
h.config.Auth.Password = newPassword
h.config.Auth.GeneratedPassword = ""
h.config.Auth.GeneratedPasswordPersisted = false
h.config.Auth.GeneratedPasswordPersistErr = ""
} else if err := h.manager.UpdateUserPassword(session.UserID, newPassword); err != nil {
if h.logger != nil { if h.logger != nil {
h.logger.Error("更新用户密码失败", zap.Error(err)) h.logger.Error("更新用户密码失败", zap.Error(err))
} }
+1 -1
View File
@@ -664,7 +664,7 @@ schedule_mode 为 cron 时必须提供有效 cron_expr;为 manual 时会清除
return batchMCPJSONResult(queue) return batchMCPJSONResult(queue)
}) })
logger.Info("批量任务 MCP 工具已注册", zap.Int("count", 12)) logger.Debug("批量任务 MCP 工具已注册", zap.Int("count", 12))
} }
// --- batch_task_list 精简结构(避免把每条子任务的 result 等大段文本塞进列表上下文) --- // --- batch_task_list 精简结构(避免把每条子任务的 result 等大段文本塞进列表上下文) ---
+3 -3
View File
@@ -1362,7 +1362,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "初始化知识库失败: " + err.Error()}) c.JSON(http.StatusInternalServerError, gin.H{"error": "初始化知识库失败: " + err.Error()})
return return
} }
h.logger.Info("知识库动态初始化完成,工具已注册") h.logger.Debug("知识库动态初始化完成,工具已注册")
} }
// 检查嵌入模型配置是否变更(需要在锁外执行,避免阻塞) // 检查嵌入模型配置是否变更(需要在锁外执行,避免阻塞)
@@ -1441,10 +1441,10 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()}) c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()})
return return
} }
h.logger.Info("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools))) h.logger.Debug("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools)))
// 重新注册工具(根据新的启用状态) // 重新注册工具(根据新的启用状态)
h.logger.Info("重新注册工具") h.logger.Debug("重新注册工具")
// 清空MCP服务器中的工具 // 清空MCP服务器中的工具
h.mcpServer.ClearTools() h.mcpServer.ClearTools()
+2 -2
View File
@@ -76,7 +76,7 @@ func RegisterKnowledgeTool(
} }
mcpServer.RegisterTool(listRiskTypesTool, listRiskTypesHandler) mcpServer.RegisterTool(listRiskTypesTool, listRiskTypesHandler)
logger.Info("风险类型列表工具已注册", zap.String("toolName", listRiskTypesTool.Name)) logger.Debug("风险类型列表工具已注册", zap.String("toolName", listRiskTypesTool.Name))
// 注册第二个工具:搜索知识库(保持原有功能) // 注册第二个工具:搜索知识库(保持原有功能)
searchTool := mcp.Tool{ searchTool := mcp.Tool{
@@ -271,7 +271,7 @@ func RegisterKnowledgeTool(
} }
mcpServer.RegisterTool(searchTool, searchHandler) mcpServer.RegisterTool(searchTool, searchHandler)
logger.Info("知识检索工具已注册", zap.String("toolName", searchTool.Name)) logger.Debug("知识检索工具已注册", zap.String("toolName", searchTool.Name))
} }
// contains 检查切片是否包含元素 // contains 检查切片是否包含元素
+328
View File
@@ -0,0 +1,328 @@
package multiagent
import (
"context"
"fmt"
"strings"
"unicode/utf8"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
const (
toolOutputTruncationMarker = "\n\n...[tool output truncated; full text persisted in reduction cache or summarization transcript]...\n\n"
aggressiveToolTruncDivisor = 4
)
// isEinoContextOverflowError reports API-side context window rejections.
func isEinoContextOverflowError(err error) bool {
if err == nil {
return false
}
msg := strings.ToLower(strings.TrimSpace(err.Error()))
if msg == "" {
return false
}
markers := []string{
"context length",
"context_length",
"maximum context",
"max context",
"context window",
"context overflow",
"too many tokens",
"token limit",
"tokens exceed",
"exceeds the context",
"input is too long",
"prompt is too long",
"request too large",
}
for _, m := range markers {
if strings.Contains(msg, m) {
return true
}
}
return false
}
func truncateBytesWithMarker(content string, maxBytes int, marker string) string {
if maxBytes <= 0 || len(content) <= maxBytes {
return content
}
if marker == "" {
marker = toolOutputTruncationMarker
}
budget := maxBytes - len(marker)
if budget <= 0 {
if len(marker) > maxBytes {
return marker[:maxBytes]
}
return marker
}
head := budget / 2
tail := budget - head
for head > 0 && !utf8.RuneStart(content[head]) {
head--
}
tailStart := len(content) - tail
for tailStart < len(content) && !utf8.RuneStart(content[tailStart]) {
tailStart++
}
return content[:head] + marker + content[tailStart:]
}
func cloneMessage(msg adk.Message) adk.Message {
if msg == nil {
return nil
}
cloned := *msg
return &cloned
}
func truncateMessageToolContent(msg adk.Message, maxBytes int, spillRef string) adk.Message {
if msg == nil || maxBytes <= 0 {
return msg
}
out := cloneMessage(msg)
marker := toolOutputTruncationMarker
if spillRef != "" {
marker = fmt.Sprintf("\n\n...[tool output truncated; retrieve full text via: %s]...\n\n", spillRef)
}
switch out.Role {
case schema.Tool:
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
case schema.Assistant:
if out.ReasoningContent != "" {
out.ReasoningContent = truncateBytesWithMarker(out.ReasoningContent, maxBytes, marker)
}
if out.Content != "" {
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
}
case schema.User:
if out.Content != "" {
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
}
}
return out
}
func countMessagesTokens(
ctx context.Context,
msgs []adk.Message,
counter summarization.TokenCounterFunc,
tools []*schema.ToolInfo,
) (int, error) {
if counter == nil {
return 0, nil
}
n, err := counter(ctx, &summarization.TokenCounterInput{Messages: msgs, Tools: tools})
if err != nil {
return 0, err
}
return n, nil
}
func truncateRoundMessagesToTokenBudget(
ctx context.Context,
round messageRound,
tokenBudget int,
counter summarization.TokenCounterFunc,
toolMaxBytes int,
spillRef string,
) ([]adk.Message, error) {
if tokenBudget <= 0 || len(round.messages) == 0 {
return nil, nil
}
msgs := append([]adk.Message(nil), round.messages...)
if n, err := countMessagesTokens(ctx, msgs, counter, nil); err != nil {
return nil, err
} else if n <= tokenBudget {
return msgs, nil
}
if toolMaxBytes <= 0 {
toolMaxBytes = 12000
}
for pass := 0; pass < 8 && toolMaxBytes >= 32; pass++ {
out := make([]adk.Message, 0, len(msgs))
for _, msg := range msgs {
switch {
case msg != nil && msg.Role == schema.Tool:
out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef))
case msg != nil && msg.Role == schema.Assistant:
out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef))
default:
out = append(out, msg)
}
}
n, err := countMessagesTokens(ctx, out, counter, nil)
if err != nil {
return nil, err
}
if n <= tokenBudget {
return out, nil
}
msgs = out
toolMaxBytes /= 2
}
return msgs, nil
}
type compactMessagesOpts struct {
maxTokens int
counter summarization.TokenCounterFunc
toolMaxBytes int
spillRef string
aggressive bool
logger *zap.Logger
phase string
}
func compactMessagesByDroppingRounds(
ctx context.Context,
messages []adk.Message,
opts compactMessagesOpts,
) ([]adk.Message, bool) {
if opts.maxTokens <= 0 || len(messages) == 0 || opts.counter == nil {
return messages, false
}
before, err := countMessagesTokens(ctx, messages, opts.counter, nil)
if err != nil || before <= opts.maxTokens {
return messages, false
}
systems := make([]adk.Message, 0, 1)
contextMsgs := make([]adk.Message, 0, len(messages))
for _, msg := range messages {
if msg != nil && msg.Role == schema.System && len(contextMsgs) == 0 {
systems = append(systems, msg)
continue
}
if msg != nil {
contextMsgs = append(contextMsgs, msg)
}
}
rounds := splitMessagesIntoRounds(contextMsgs)
if len(rounds) == 0 {
return messages, false
}
startIdx := 0
if opts.aggressive {
startIdx = len(rounds) - 1
if startIdx < 0 {
startIdx = 0
}
}
dropped := 0
for len(rounds) > 1 || (opts.aggressive && len(rounds) == 1) {
if !opts.aggressive && len(rounds) <= 1 {
break
}
if opts.aggressive && len(rounds) == 1 {
// Fall through to latest-round truncation below.
break
}
rounds = rounds[1:]
dropped++
candidate := append([]adk.Message(nil), systems...)
for _, round := range rounds {
candidate = append(candidate, round.messages...)
}
after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil)
if countErr != nil {
break
}
if after <= opts.maxTokens {
if opts.logger != nil {
opts.logger.Warn("eino context compacted by dropping older rounds",
zap.String("phase", opts.phase),
zap.Int("tokens_before", before),
zap.Int("tokens_after", after),
zap.Int("max_tokens", opts.maxTokens),
zap.Int("dropped_rounds", dropped),
zap.Bool("aggressive", opts.aggressive),
)
}
return candidate, true
}
if opts.aggressive {
break
}
}
if len(rounds) == 0 {
return messages, false
}
latest := rounds[len(rounds)-1]
truncated, truncErr := truncateRoundMessagesToTokenBudget(
ctx, latest, opts.maxTokens, opts.counter, opts.toolMaxBytes, opts.spillRef,
)
if truncErr != nil || len(truncated) == 0 {
if opts.logger != nil {
opts.logger.Warn("eino context still above budget after round compaction; passing through without local error",
zap.String("phase", opts.phase),
zap.Int("tokens_before", before),
zap.Int("max_tokens", opts.maxTokens),
zap.Bool("aggressive", opts.aggressive),
)
}
return messages, false
}
candidate := append([]adk.Message(nil), systems...)
if dropped > 0 || startIdx > 0 {
for _, round := range rounds[:len(rounds)-1] {
candidate = append(candidate, round.messages...)
}
}
candidate = append(candidate, truncated...)
after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil)
if countErr != nil {
return messages, false
}
if opts.logger != nil {
opts.logger.Warn("eino context compacted by truncating latest round tool output",
zap.String("phase", opts.phase),
zap.Int("tokens_before", before),
zap.Int("tokens_after", after),
zap.Int("max_tokens", opts.maxTokens),
zap.Int("dropped_rounds", dropped),
zap.Bool("aggressive", opts.aggressive),
)
}
return candidate, true
}
func aggressiveCompactMessagesForOverflow(
ctx context.Context,
messages []adk.Message,
maxTotalTokens int,
modelName string,
toolMaxBytes int,
phase string,
logger *zap.Logger,
) []adk.Message {
if len(messages) == 0 || maxTotalTokens <= 0 {
return messages
}
budget := maxTotalTokens * 70 / 100
if budget < 4096 {
budget = 4096
}
aggressiveToolMax := toolMaxBytes / aggressiveToolTruncDivisor
if aggressiveToolMax < 2048 {
aggressiveToolMax = 2048
}
out, _ := compactMessagesByDroppingRounds(ctx, messages, compactMessagesOpts{
maxTokens: budget,
counter: einoSummarizationTokenCounter(modelName),
toolMaxBytes: aggressiveToolMax,
aggressive: true,
logger: logger,
phase: phase,
})
return out
}
+104
View File
@@ -0,0 +1,104 @@
package multiagent
import (
"context"
"errors"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestIsEinoContextOverflowError(t *testing.T) {
t.Parallel()
cases := []struct {
err error
want bool
}{
{nil, false},
{errors.New("context length exceeded"), true},
{errors.New("maximum context length"), true},
{errors.New("input is too long for model"), true},
{errors.New("HTTP 429 Too Many Requests"), false},
{errors.New("invalid api key"), false},
}
for _, tc := range cases {
if got := isEinoContextOverflowError(tc.err); got != tc.want {
t.Fatalf("isEinoContextOverflowError(%v) = %v, want %v", tc.err, got, tc.want)
}
}
}
func TestTruncateRoundMessagesToTokenBudget(t *testing.T) {
huge := strings.Repeat("x", 8000)
round := messageRound{messages: []adk.Message{
assistantToolCallsMsg("", "c1"),
schema.ToolMessage(huge, "c1"),
}}
out, err := truncateRoundMessagesToTokenBudget(
context.Background(), round, 256, einoSummarizationTokenCounter("gpt-4o"), 512, "",
)
if err != nil {
t.Fatal(err)
}
for _, msg := range out {
if msg != nil && msg.Role == schema.Tool && len(msg.Content) >= len(huge) {
t.Fatalf("expected truncated tool output, got len=%d", len(msg.Content))
}
}
}
func TestBuildBudgetedSummarizationModelInputTruncatesOversizedLatestRound(t *testing.T) {
huge := strings.Repeat("x", 8000)
msgs := []adk.Message{
assistantToolCallsMsg("", "call-latest"),
schema.ToolMessage(huge, "call-latest"),
}
counter := einoSummarizationTokenCounter("gpt-4o")
input, dropped, err := buildBudgetedSummarizationModelInput(
context.Background(),
schema.SystemMessage("sys"),
schema.UserMessage("instr"),
msgs,
counter,
512,
summarizationInputBudgetOpts{toolMaxBytes: 256},
)
if err != nil {
t.Fatal(err)
}
if dropped != 0 {
t.Fatalf("expected no dropped rounds, got %d", dropped)
}
toolContent := ""
for _, msg := range input {
if msg != nil && msg.Role == schema.Tool {
toolContent = msg.Content
}
}
if len(toolContent) >= len(huge) {
t.Fatalf("expected oversized tool output to be compacted, got len=%d", len(toolContent))
}
}
func TestModelInputSoftBudgetNeverErrors(t *testing.T) {
mw := &modelInputSoftBudgetMiddleware{
maxTokens: 4,
toolMaxBytes: 16,
counter: fixedTokenCounter(4),
phase: "test",
}
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.UserMessage("u"),
assistantToolCallsMsg("", "c1"),
schema.ToolMessage(strings.Repeat("t", 200), "c1"),
}}
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatalf("soft budget must not error: %v", err)
}
if out == nil || len(out.Messages) == 0 {
t.Fatal("expected compacted messages")
}
}
+31
View File
@@ -105,6 +105,11 @@ type einoADKRunLoopArgs struct {
// EinoCallbacks 可选:为 ADK Runner 注入 eino [callbacks] 全链路观测(见 internal/einoobserve)。 // EinoCallbacks 可选:为 ADK Runner 注入 eino [callbacks] 全链路观测(见 internal/einoobserve)。
EinoCallbacks *config.MultiAgentEinoCallbacksConfig EinoCallbacks *config.MultiAgentEinoCallbacksConfig
// MaxTotalTokens / ToolMaxBytes / ModelName 用于 context overflow 时的激进压缩续跑。
MaxTotalTokens int
ToolMaxBytes int
ModelName string
} }
func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs []adk.Message) (*RunResult, error) { func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs []adk.Message) (*RunResult, error) {
@@ -439,6 +444,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
iter = startRunnerIter(msgs) iter = startRunnerIter(msgs)
} }
transientRetrier := newEinoTransientRunRetrier(einoTransientRunRetryPolicyFromArgs(args)) transientRetrier := newEinoTransientRunRetrier(einoTransientRunRetryPolicyFromArgs(args))
var contextOverflowRetried bool
handleRunErr := func(runErr error) error { handleRunErr := func(runErr error) error {
if runErr == nil { if runErr == nil {
return nil return nil
@@ -495,6 +501,31 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
if runErr == nil { if runErr == nil {
return false, nil return false, nil
} }
if isEinoContextOverflowError(runErr) && !contextOverflowRetried {
contextOverflowRetried = true
restartMsgs, ctxSource := einoMessagesForRunRestart(args, baseMsgs, runAccumulatedMsgs, baseAccumulatedCount)
restartMsgs = aggressiveCompactMessagesForOverflow(
ctx, restartMsgs, args.MaxTotalTokens, args.ModelName, args.ToolMaxBytes, orchMode, logger,
)
if logger != nil {
logger.Warn("eino context overflow, retrying with aggressive compaction",
zap.Error(runErr),
zap.String("orchestration", orchMode),
zap.String("contextSource", string(ctxSource)),
)
}
if progress != nil {
progress("eino_context_overflow_retry", "上下文超限,正在激进压缩后重试…", map[string]interface{}{
"conversationId": conversationID,
"source": "eino",
"orchestration": orchMode,
"contextSource": string(ctxSource),
})
}
msgs = restartMsgs
iter = startRunnerIter(msgs)
return true, nil
}
if !isEinoTransientRunError(runErr) { if !isEinoTransientRunError(runErr) {
return false, handleRunErr(runErr) return false, handleRunErr(runErr)
} }
@@ -1,6 +1,8 @@
package multiagent package multiagent
import ( import (
"cyberstrike-ai/internal/config"
"github.com/cloudwego/eino/adk" "github.com/cloudwego/eino/adk"
"go.uber.org/zap" "go.uber.org/zap"
) )
@@ -13,7 +15,7 @@ import (
// 2. continuation user dedup — drop stale session-resume injections // 2. continuation user dedup — drop stale session-resume injections
// 3. pre-summarization tool-call/result reconciliation // 3. pre-summarization tool-call/result reconciliation
// 4. summarization // 4. summarization
// 5. total model-input hard budget // 5. soft model-input budget (warn/compact only, never fail locally)
// 6. final tool-call/result reconciliation // 6. final tool-call/result reconciliation
// 7. orphan tool prune (defense in depth) // 7. orphan tool prune (defense in depth)
// 8. telemetry // 8. telemetry
@@ -24,6 +26,7 @@ type einoChatModelTailConfig struct {
summarization adk.ChatModelAgentMiddleware summarization adk.ChatModelAgentMiddleware
modelName string modelName string
maxTotalTokens int maxTotalTokens int
toolMaxBytes int
conversationID string conversationID string
trace *modelFacingTraceHolder trace *modelFacingTraceHolder
skipOrphanPruner bool skipOrphanPruner bool
@@ -40,7 +43,7 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization")) handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization"))
handlers = append(handlers, cfg.summarization) handlers = append(handlers, cfg.summarization)
} }
handlers = append(handlers, newModelInputBudgetMiddleware(cfg.maxTotalTokens, cfg.modelName, cfg.logger, cfg.phase)) handlers = append(handlers, newModelInputSoftBudgetMiddleware(cfg.maxTotalTokens, cfg.toolMaxBytes, cfg.modelName, cfg.logger, cfg.phase))
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase)) handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase))
if !cfg.skipOrphanPruner { if !cfg.skipOrphanPruner {
handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase)) handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
@@ -57,3 +60,10 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
} }
return handlers return handlers
} }
func toolMaxBytesFromMW(mwCfg *config.MultiAgentEinoMiddlewareConfig) int {
if mwCfg != nil {
return mwCfg.ReductionMaxLengthForTruncEffective()
}
return config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective()
}
@@ -132,6 +132,7 @@ func buildPlanExecuteExecutorHandlers(ctx context.Context, a *PlanExecuteRootArg
summarization: sumMw, summarization: sumMw,
modelName: a.ModelName, modelName: a.ModelName,
maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens, maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(a.MwCfg),
conversationID: a.ConversationID, conversationID: a.ConversationID,
trace: a.ModelFacingTrace, trace: a.ModelFacingTrace,
}) })
@@ -151,6 +151,7 @@ func RunEinoSingleChatModelAgent(
summarization: mainSumMw, summarization: mainSumMw,
modelName: appCfg.OpenAI.Model, modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
conversationID: conversationID, conversationID: conversationID,
trace: modelFacingTrace, trace: modelFacingTrace,
}) })
@@ -235,6 +236,9 @@ func RunEinoSingleChatModelAgent(
DA: chatAgent, DA: chatAgent,
ModelFacingTrace: modelFacingTrace, ModelFacingTrace: modelFacingTrace,
EinoCallbacks: &ma.EinoCallbacks, EinoCallbacks: &ma.EinoCallbacks,
MaxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
ToolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
ModelName: appCfg.OpenAI.Model,
EmptyResponseMessage: "(Eino ADK single-agent session completed but no assistant text was captured. Check process details or logs.) " + EmptyResponseMessage: "(Eino ADK single-agent session completed but no assistant text was captured. Check process details or logs.) " +
"Eino ADK 单代理会话已完成,但未捕获到助手文本输出。请查看过程详情或日志。)", "Eino ADK 单代理会话已完成,但未捕获到助手文本输出。请查看过程详情或日志。)",
}, baseMsgs) }, baseMsgs)
+74 -9
View File
@@ -98,17 +98,21 @@ func newEinoSummarizationMiddleware(
} }
triggerRatio := 0.8 triggerRatio := 0.8
emitInternalEvents := true emitInternalEvents := true
outputReserve := config.DefaultSummarizationOutputReserveTokens
userLedgerMaxRunes := config.DefaultSummarizationUserIntentLedgerMaxRunes userLedgerMaxRunes := config.DefaultSummarizationUserIntentLedgerMaxRunes
userLedgerEntryMaxRunes := config.DefaultSummarizationUserIntentLedgerEntryMaxRunes userLedgerEntryMaxRunes := config.DefaultSummarizationUserIntentLedgerEntryMaxRunes
toolMaxBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective()
if mwCfg != nil { if mwCfg != nil {
triggerRatio = mwCfg.SummarizationTriggerRatioEffective() triggerRatio = mwCfg.SummarizationTriggerRatioEffective()
emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective() emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective()
outputReserve = mwCfg.SummarizationOutputReserveTokensEffective()
userLedgerMaxRunes = mwCfg.SummarizationUserIntentLedgerMaxRunesEffective() userLedgerMaxRunes = mwCfg.SummarizationUserIntentLedgerMaxRunesEffective()
userLedgerEntryMaxRunes = mwCfg.SummarizationUserIntentLedgerEntryMaxRunesEffective() userLedgerEntryMaxRunes = mwCfg.SummarizationUserIntentLedgerEntryMaxRunesEffective()
toolMaxBytes = mwCfg.ReductionMaxLengthForTruncEffective()
} }
// The ledger is merged into the leading system message and cannot be removed as // The ledger is merged into the leading system message and cannot be removed as
// an ordinary conversation round. Bound it relative to the configured window so // an ordinary conversation round. Bound it relative to the configured window so
// it cannot crowd out the summary/latest turn or trip the final fail-closed guard. // it cannot crowd out the summary/latest turn.
ledgerWindowCap := modelFacingRuneBudget(maxTotal, 0.20) ledgerWindowCap := modelFacingRuneBudget(maxTotal, 0.20)
userLedgerMaxRunes = minPositiveInt(userLedgerMaxRunes, ledgerWindowCap) userLedgerMaxRunes = minPositiveInt(userLedgerMaxRunes, ledgerWindowCap)
userLedgerEntryMaxRunes = minPositiveInt(userLedgerEntryMaxRunes, userLedgerMaxRunes) userLedgerEntryMaxRunes = minPositiveInt(userLedgerEntryMaxRunes, userLedgerMaxRunes)
@@ -137,11 +141,11 @@ func newEinoSummarizationMiddleware(
if recentTrailMax > trigger/2 { if recentTrailMax > trigger/2 {
recentTrailMax = trigger / 2 recentTrailMax = trigger / 2
} }
// The summarization request itself needs output headroom. A trigger is not a hard // Summarization input aligns with the trigger threshold, minus explicit output reserve.
// request limit: one large turn can jump far beyond it. Bound the actual summary summaryInputMax := trigger - outputReserve
// model input to 60% of the configured context window and keep complete recent if summaryInputMax < 4096 {
// rounds so tool_call/tool_result pairs are never split. summaryInputMax = trigger * 80 / 100
summaryInputMax := int(float64(maxTotal) * 0.6) }
if summaryInputMax < 4096 { if summaryInputMax < 4096 {
summaryInputMax = 4096 summaryInputMax = 4096
} }
@@ -153,6 +157,9 @@ func newEinoSummarizationMiddleware(
baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization") baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization")
} }
base := baseRoot base := baseRoot
if abs, err := filepath.Abs(base); err == nil {
base = abs
}
if mkErr := os.MkdirAll(base, 0o755); mkErr == nil { if mkErr := os.MkdirAll(base, 0o755); mkErr == nil {
transcriptPath = filepath.Join(base, "transcript.txt") transcriptPath = filepath.Join(base, "transcript.txt")
} }
@@ -160,6 +167,7 @@ func newEinoSummarizationMiddleware(
retryPolicy := einoTransientRunRetryPolicyFromMW(mwCfg) retryPolicy := einoTransientRunRetryPolicyFromMW(mwCfg)
retryMax := retryPolicy.maxAttempts retryMax := retryPolicy.maxAttempts
var summaryOverflowRetries int
// ModelOptions apply only to summarization Generate (same ChatModel instance as the agent). // ModelOptions apply only to summarization Generate (same ChatModel instance as the agent).
// Strip thinking/reasoning on this call path; mark requests for empty-choices diagnostics. // Strip thinking/reasoning on this call path; mark requests for empty-choices diagnostics.
@@ -189,13 +197,29 @@ func newEinoSummarizationMiddleware(
zap.String("path", transcriptPath), zap.Error(werr)) zap.String("path", transcriptPath), zap.Error(werr))
} }
} }
budget := summaryInputMax
aggressive := summaryOverflowRetries > 0
if aggressive {
budget = summaryInputMax * 70 / 100
if budget < 4096 {
budget = 4096
}
}
input, dropped, berr := buildBudgetedSummarizationModelInput( input, dropped, berr := buildBudgetedSummarizationModelInput(
ctx, sysInstruction, userInstruction, originalMsgs, tokenCounter, summaryInputMax, ctx, sysInstruction, userInstruction, originalMsgs, tokenCounter, budget,
summarizationInputBudgetOpts{
toolMaxBytes: toolMaxBytes,
spillRef: transcriptPath,
aggressive: aggressive,
},
) )
if logger != nil && (berr != nil || dropped > 0) { if logger != nil && (berr != nil || dropped > 0 || aggressive) {
fields := []zap.Field{ fields := []zap.Field{
zap.Int("max_input_tokens", summaryInputMax), zap.Int("max_input_tokens", budget),
zap.Int("trigger_context_tokens", trigger),
zap.Int("output_reserve_tokens", outputReserve),
zap.Int("dropped_rounds", dropped), zap.Int("dropped_rounds", dropped),
zap.Bool("aggressive", aggressive),
} }
if berr != nil { if berr != nil {
fields = append(fields, zap.Error(berr)) fields = append(fields, zap.Error(berr))
@@ -220,6 +244,15 @@ func newEinoSummarizationMiddleware(
Retry: &summarization.RetryConfig{ Retry: &summarization.RetryConfig{
MaxRetries: &retryMax, MaxRetries: &retryMax,
ShouldRetry: func(_ context.Context, _ adk.Message, err error) bool { ShouldRetry: func(_ context.Context, _ adk.Message, err error) bool {
if isEinoContextOverflowError(err) && summaryOverflowRetries < 1 {
summaryOverflowRetries++
if logger != nil {
logger.Warn("eino summarization context overflow, retrying with aggressive compaction",
zap.Error(err),
)
}
return true
}
retry := isEinoTransientRunError(err) retry := isEinoTransientRunError(err)
if retry && logger != nil { if retry && logger != nil {
logger.Warn("eino summarization generate transient error, will retry if attempts remain", logger.Warn("eino summarization generate transient error, will retry if attempts remain",
@@ -275,6 +308,13 @@ func newEinoSummarizationMiddleware(
return mw, nil return mw, nil
} }
// summarizationInputBudgetOpts controls spill/truncation behavior when a round alone exceeds budget.
type summarizationInputBudgetOpts struct {
toolMaxBytes int
spillRef string
aggressive bool
}
// buildBudgetedSummarizationModelInput builds the exact payload sent to the summary model. // buildBudgetedSummarizationModelInput builds the exact payload sent to the summary model.
// It retains the newest complete conversation rounds within budget and emits an explicit // It retains the newest complete conversation rounds within budget and emits an explicit
// marker when older rounds are omitted. The full pre-compaction transcript is persisted // marker when older rounds are omitted. The full pre-compaction transcript is persisted
@@ -286,6 +326,7 @@ func buildBudgetedSummarizationModelInput(
originalMsgs []adk.Message, originalMsgs []adk.Message,
tokenCounter summarization.TokenCounterFunc, tokenCounter summarization.TokenCounterFunc,
maxTokens int, maxTokens int,
opts summarizationInputBudgetOpts,
) ([]adk.Message, int, error) { ) ([]adk.Message, int, error) {
base := []adk.Message{sysInstruction, userInstruction} base := []adk.Message{sysInstruction, userInstruction}
baseTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: base}) baseTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: base})
@@ -312,12 +353,36 @@ func buildBudgetedSummarizationModelInput(
rounds := splitMessagesIntoRounds(contextMsgs) rounds := splitMessagesIntoRounds(contextMsgs)
selectedReverse := make([]messageRound, 0, len(rounds)) selectedReverse := make([]messageRound, 0, len(rounds))
used := 0 used := 0
toolMaxBytes := opts.toolMaxBytes
if toolMaxBytes <= 0 {
toolMaxBytes = 12000
}
if opts.aggressive {
toolMaxBytes /= aggressiveToolTruncDivisor
if toolMaxBytes < 2048 {
toolMaxBytes = 2048
}
}
for i := len(rounds) - 1; i >= 0; i-- { for i := len(rounds) - 1; i >= 0; i-- {
n, countErr := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: rounds[i].messages}) n, countErr := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: rounds[i].messages})
if countErr != nil { if countErr != nil {
return nil, 0, countErr return nil, 0, countErr
} }
if used+n > remaining { if used+n > remaining {
if len(selectedReverse) == 0 {
slot := remaining - used
if slot > 0 {
truncated, truncErr := truncateRoundMessagesToTokenBudget(
ctx, rounds[i], slot, tokenCounter, toolMaxBytes, opts.spillRef,
)
if truncErr != nil {
return nil, 0, truncErr
}
if len(truncated) > 0 {
selectedReverse = append(selectedReverse, messageRound{messages: truncated})
}
}
}
break break
} }
used += n used += n
+1 -1
View File
@@ -67,7 +67,7 @@ func TestBuildBudgetedSummarizationModelInputKeepsRecentCompleteRounds(t *testin
} }
input, dropped, err := buildBudgetedSummarizationModelInput( input, dropped, err := buildBudgetedSummarizationModelInput(
context.Background(), schema.SystemMessage("summary-system"), schema.UserMessage("summary-instruction"), context.Background(), schema.SystemMessage("summary-system"), schema.UserMessage("summary-instruction"),
msgs, fixedTokenCounter(2), 7, msgs, fixedTokenCounter(2), 7, summarizationInputBudgetOpts{},
) )
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@@ -35,9 +35,6 @@ func isEinoTransientRunError(err error) bool {
if msg == "" { if msg == "" {
return false return false
} }
if strings.Contains(msg, "model input exceeds configured hard budget") {
return false
}
transientMarkers := []string{ transientMarkers := []string{
"406", "406",
"429", "429",
@@ -33,7 +33,6 @@ func TestIsEinoTransientRunError(t *testing.T) {
{"iteration limit", errors.New("max iteration reached"), false}, {"iteration limit", errors.New("max iteration reached"), false},
{"canceled", context.Canceled, false}, {"canceled", context.Canceled, false},
{"deadline", context.DeadlineExceeded, false}, {"deadline", context.DeadlineExceeded, false},
{"model hard budget containing 500", errors.New("model input exceeds configured hard budget after preserving the latest round: tokens=20500 max=19500 phase=test"), false},
{"auth", errors.New("invalid api key"), false}, {"auth", errors.New("invalid api key"), false},
} }
for _, tc := range cases { for _, tc := range cases {
@@ -1,106 +0,0 @@
package multiagent
import (
"context"
"fmt"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// modelInputBudgetMiddleware is the final deterministic guard before a normal model call.
// Summarization remains the primary compaction strategy; this middleware only removes the
// oldest complete rounds if the finalized state still exceeds 65% of the configured
// window, leaving serialization/tokenizer headroom for the outbound HTTP guard.
type modelInputBudgetMiddleware struct {
adk.BaseChatModelAgentMiddleware
maxTokens int
counter summarization.TokenCounterFunc
logger *zap.Logger
phase string
}
func newModelInputBudgetMiddleware(maxTotalTokens int, modelName string, logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
if maxTotalTokens <= 0 {
maxTotalTokens = 120000
}
limit := int(float64(maxTotalTokens) * 0.65)
if limit < 4096 {
limit = 4096
}
return &modelInputBudgetMiddleware{
maxTokens: limit,
counter: einoSummarizationTokenCounter(modelName),
logger: logger,
phase: phase,
}
}
func (m *modelInputBudgetMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, error) {
if m == nil || state == nil || len(state.Messages) == 0 {
return ctx, state, nil
}
count := func(msgs []adk.Message) (int, error) {
input := &summarization.TokenCounterInput{Messages: msgs}
if mc != nil {
input.Tools = mc.Tools
}
return m.counter(ctx, input)
}
before, err := count(state.Messages)
if err != nil {
return ctx, state, err
}
if before <= m.maxTokens {
return ctx, state, nil
}
systems := make([]adk.Message, 0, 1)
contextMsgs := make([]adk.Message, 0, len(state.Messages))
for _, msg := range state.Messages {
if msg != nil && msg.Role == schema.System && len(contextMsgs) == 0 {
systems = append(systems, msg)
continue
}
if msg != nil {
contextMsgs = append(contextMsgs, msg)
}
}
rounds := splitMessagesIntoRounds(contextMsgs)
dropped := 0
var candidate []adk.Message
for len(rounds) > 1 {
rounds = rounds[1:]
dropped++
candidate = append(candidate[:0], systems...)
for _, round := range rounds {
candidate = append(candidate, round.messages...)
}
after, countErr := count(candidate)
if countErr != nil {
return ctx, state, countErr
}
if after <= m.maxTokens {
out := *state
out.Messages = append([]adk.Message(nil), candidate...)
if m.logger != nil {
m.logger.Warn("eino model input hard budget applied",
zap.String("phase", m.phase), zap.Int("tokens_before", before),
zap.Int("tokens_after", after), zap.Int("max_tokens", m.maxTokens),
zap.Int("dropped_rounds", dropped))
}
return ctx, &out, nil
}
}
return ctx, state, fmt.Errorf(
"model input exceeds configured hard budget after preserving the latest round: tokens=%d max=%d phase=%s",
before, m.maxTokens, m.phase,
)
}
@@ -1,50 +0,0 @@
package multiagent
import (
"context"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestModelInputBudgetDropsOldestCompleteRoundsAndPersistsLatest(t *testing.T) {
mw := &modelInputBudgetMiddleware{
maxTokens: 7,
counter: fixedTokenCounter(4),
phase: "test",
}
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("system"),
schema.UserMessage("old-user"),
schema.AssistantMessage("old-answer", nil),
schema.UserMessage("latest-user"),
assistantToolCallsMsg("", "latest-call"),
schema.ToolMessage("latest-result", "latest-call"),
}}
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatal(err)
}
joined := formatSummarizationTranscript(out.Messages)
if strings.Contains(joined, "old-user") || strings.Contains(joined, "old-answer") {
t.Fatalf("old rounds retained after hard budget: %s", joined)
}
if !strings.Contains(joined, "latest-user") || !strings.Contains(joined, "latest-result") {
t.Fatalf("latest rounds lost after hard budget: %s", joined)
}
}
func TestModelInputBudgetFailsLocallyWhenLatestRoundAloneCannotFit(t *testing.T) {
mw := &modelInputBudgetMiddleware{maxTokens: 2, counter: fixedTokenCounter(4), phase: "test"}
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("system"),
assistantToolCallsMsg("", "latest-call"),
schema.ToolMessage("latest-result", "latest-call"),
}}
_, _, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
if err == nil || !strings.Contains(err.Error(), "hard budget") {
t.Fatalf("expected local hard-budget error, got %v", err)
}
}
@@ -0,0 +1,66 @@
package multiagent
import (
"context"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization"
"go.uber.org/zap"
)
// modelInputSoftBudgetMiddleware is the final guard before a normal model call.
// It drops oldest complete rounds and truncates oversized tool output in the latest
// round, but never fails locally — API context limits are handled by overflow retry.
type modelInputSoftBudgetMiddleware struct {
adk.BaseChatModelAgentMiddleware
maxTokens int
toolMaxBytes int
counter summarization.TokenCounterFunc
logger *zap.Logger
phase string
}
func newModelInputSoftBudgetMiddleware(
maxTotalTokens int,
toolMaxBytes int,
modelName string,
logger *zap.Logger,
phase string,
) adk.ChatModelAgentMiddleware {
if maxTotalTokens <= 0 {
maxTotalTokens = 120000
}
if toolMaxBytes <= 0 {
toolMaxBytes = 12000
}
return &modelInputSoftBudgetMiddleware{
maxTokens: maxTotalTokens,
toolMaxBytes: toolMaxBytes,
counter: einoSummarizationTokenCounter(modelName),
logger: logger,
phase: phase,
}
}
func (m *modelInputSoftBudgetMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, error) {
if m == nil || state == nil || len(state.Messages) == 0 {
return ctx, state, nil
}
compacted, changed := compactMessagesByDroppingRounds(ctx, state.Messages, compactMessagesOpts{
maxTokens: m.maxTokens,
counter: m.counter,
toolMaxBytes: m.toolMaxBytes,
phase: m.phase,
logger: m.logger,
})
if !changed {
return ctx, state, nil
}
out := *state
out.Messages = compacted
return ctx, &out, nil
}
+7
View File
@@ -252,6 +252,7 @@ func RunDeepAgent(
summarization: subSumMw, summarization: subSumMw,
modelName: appCfg.OpenAI.Model, modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
conversationID: conversationID, conversationID: conversationID,
}) })
@@ -413,6 +414,7 @@ func RunDeepAgent(
summarization: mainSumMw, summarization: mainSumMw,
modelName: appCfg.OpenAI.Model, modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
conversationID: conversationID, conversationID: conversationID,
trace: modelFacingTrace, trace: modelFacingTrace,
}) })
@@ -430,6 +432,7 @@ func RunDeepAgent(
summarization: mainSumMw, summarization: mainSumMw,
modelName: appCfg.OpenAI.Model, modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
conversationID: conversationID, conversationID: conversationID,
trace: modelFacingTrace, trace: modelFacingTrace,
}) })
@@ -505,6 +508,7 @@ func RunDeepAgent(
summarization: mainSumMw, summarization: mainSumMw,
modelName: appCfg.OpenAI.Model, modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
toolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
conversationID: conversationID, conversationID: conversationID,
skipTrace: true, skipTrace: true,
}), }),
@@ -608,6 +612,9 @@ func RunDeepAgent(
DA: da, DA: da,
ModelFacingTrace: modelFacingTrace, ModelFacingTrace: modelFacingTrace,
EinoCallbacks: &ma.EinoCallbacks, EinoCallbacks: &ma.EinoCallbacks,
MaxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
ToolMaxBytes: toolMaxBytesFromMW(&ma.EinoMiddleware),
ModelName: appCfg.OpenAI.Model,
EmptyResponseMessage: "(Eino multi-agent orchestration completed but no assistant text was captured. Check process details or logs.) " + EmptyResponseMessage: "(Eino multi-agent orchestration completed but no assistant text was captured. Check process details or logs.) " +
"(Eino 多代理编排已完成,但未捕获到助手文本输出。请查看过程详情或日志。)", "(Eino 多代理编排已完成,但未捕获到助手文本输出。请查看过程详情或日志。)",
}, baseMsgs) }, baseMsgs)
+29 -70
View File
@@ -32,7 +32,6 @@ type Session struct {
// AuthManager manages password-based authentication and session lifecycle. // AuthManager manages password-based authentication and session lifecycle.
type AuthManager struct { type AuthManager struct {
password string
sessionDuration time.Duration sessionDuration time.Duration
db *database.DB db *database.DB
@@ -41,39 +40,49 @@ type AuthManager struct {
} }
// NewAuthManager creates a new AuthManager instance. // NewAuthManager creates a new AuthManager instance.
func NewAuthManager(password string, sessionDurationHours int) (*AuthManager, error) { func NewAuthManager(sessionDurationHours int) *AuthManager {
if strings.TrimSpace(password) == "" {
return nil, errors.New("auth password must be configured")
}
if sessionDurationHours <= 0 { if sessionDurationHours <= 0 {
sessionDurationHours = 12 sessionDurationHours = 12
} }
return &AuthManager{ return &AuthManager{
password: password,
sessionDuration: time.Duration(sessionDurationHours) * time.Hour, sessionDuration: time.Duration(sessionDurationHours) * time.Hour,
sessions: make(map[string]Session), sessions: make(map[string]Session),
}, nil }
} }
// AttachRBACStore enables multi-user RBAC authentication and bootstraps the // AttachRBACStore enables multi-user RBAC authentication. When no users exist yet,
// built-in admin account from the legacy auth.password value. // it bootstraps the built-in admin account and returns the generated initial password.
func (a *AuthManager) AttachRBACStore(db *database.DB) error { func (a *AuthManager) AttachRBACStore(db *database.DB) (generatedAdminPassword string, err error) {
if db == nil { if db == nil {
return nil return "", errors.New("database is required for authentication")
} }
hash, err := HashPassword(a.password)
needsAdminPassword, err := db.RBACNeedsAdminPassword()
if err != nil { if err != nil {
return err return "", err
} }
if err := db.BootstrapRBAC(hash, PermissionCatalog); err != nil {
return err adminPasswordHash := ""
if needsAdminPassword {
generatedAdminPassword, err = GenerateStrongPassword(24)
if err != nil {
return "", err
}
adminPasswordHash, err = HashPassword(generatedAdminPassword)
if err != nil {
return "", err
}
} }
if err := db.BootstrapRBAC(adminPasswordHash, PermissionCatalog); err != nil {
return "", err
}
a.mu.Lock() a.mu.Lock()
a.db = db a.db = db
a.mu.Unlock() a.mu.Unlock()
return nil return generatedAdminPassword, nil
} }
// Authenticate validates the password and creates a new session. // Authenticate validates the password and creates a new session.
@@ -94,23 +103,9 @@ func (a *AuthManager) authenticateSession(username, password string) (Session, e
a.mu.RLock() a.mu.RLock()
db := a.db db := a.db
legacyPassword := a.password
a.mu.RUnlock() a.mu.RUnlock()
if db == nil { if db == nil {
if password != legacyPassword { return Session{}, errors.New("authentication store is not configured")
return Session{}, ErrInvalidPassword
}
return Session{
Token: token,
ExpiresAt: expiresAt,
UserID: "admin",
Username: "admin",
DisplayName: "管理员",
Roles: []string{database.RBACSystemRoleAdmin},
Permissions: allPermissions(),
Scope: database.RBACScopeAll,
}, nil
} }
username = strings.TrimSpace(strings.ToLower(username)) username = strings.TrimSpace(strings.ToLower(username))
@@ -187,10 +182,9 @@ func (a *AuthManager) CheckPassword(password string) bool {
func (a *AuthManager) CheckUserPassword(username, password string) bool { func (a *AuthManager) CheckUserPassword(username, password string) bool {
a.mu.RLock() a.mu.RLock()
db := a.db db := a.db
legacyPassword := a.password
a.mu.RUnlock() a.mu.RUnlock()
if db == nil { if db == nil {
return password == legacyPassword return false
} }
user, err := db.GetRBACUserByUsername(username) user, err := db.GetRBACUserByUsername(username)
if err != nil { if err != nil {
@@ -212,7 +206,7 @@ func (a *AuthManager) UpdateUserPassword(userID, password string) error {
db := a.db db := a.db
a.mu.RUnlock() a.mu.RUnlock()
if db == nil { if db == nil {
return a.UpdateConfig(password, a.SessionDurationHours()) return errors.New("authentication store is not configured")
} }
if err := db.UpdateRBACUserPassword(userID, hash); err != nil { if err := db.UpdateRBACUserPassword(userID, hash); err != nil {
return err return err
@@ -263,41 +257,6 @@ func (a *AuthManager) SessionDurationHours() int {
return int(a.sessionDuration / time.Hour) return int(a.sessionDuration / time.Hour)
} }
// UpdateConfig updates the password and session duration, revoking existing sessions.
func (a *AuthManager) UpdateConfig(password string, sessionDurationHours int) error {
password = strings.TrimSpace(password)
if password == "" {
return errors.New("auth password must be configured")
}
if sessionDurationHours <= 0 {
sessionDurationHours = 12
}
hash := ""
if a.db != nil {
var err error
hash, err = HashPassword(password)
if err != nil {
return err
}
}
a.mu.Lock()
a.password = password
a.sessionDuration = time.Duration(sessionDurationHours) * time.Hour
a.sessions = make(map[string]Session)
db := a.db
a.mu.Unlock()
if db != nil {
if err := db.UpdateRBACAdminPassword(hash); err != nil {
return err
}
}
return nil
}
func allPermissions() map[string]bool { func allPermissions() map[string]bool {
out := make(map[string]bool, len(PermissionCatalog)) out := make(map[string]bool, len(PermissionCatalog))
for key := range PermissionCatalog { for key := range PermissionCatalog {
@@ -0,0 +1,38 @@
package security
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestAttachRBACStoreBootstrapsAdminPassword(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-bootstrap.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
manager := NewAuthManager(12)
generated, err := manager.AttachRBACStore(db)
if err != nil {
t.Fatalf("AttachRBACStore: %v", err)
}
if generated == "" {
t.Fatal("expected generated admin password on first bootstrap")
}
if !manager.CheckUserPassword("admin", generated) {
t.Fatal("generated password should authenticate admin")
}
second, err := manager.AttachRBACStore(db)
if err != nil {
t.Fatalf("AttachRBACStore second call: %v", err)
}
if second != "" {
t.Fatalf("expected no password on second bootstrap, got %q", second)
}
}
+2 -5
View File
@@ -20,11 +20,8 @@ func TestAuthManagerAuthenticatesCreatedRBACUser(t *testing.T) {
} }
t.Cleanup(func() { _ = db.Close() }) t.Cleanup(func() { _ = db.Close() })
manager, err := NewAuthManager("admin-secret", 12) manager := NewAuthManager(12)
if err != nil { if _, err := manager.AttachRBACStore(db); err != nil {
t.Fatalf("NewAuthManager: %v", err)
}
if err := manager.AttachRBACStore(db); err != nil {
t.Fatalf("AttachRBACStore: %v", err) t.Fatalf("AttachRBACStore: %v", err)
} }
hash, err := HashPassword("operator-secret") hash, err := HashPassword("operator-secret")
+13 -13
View File
@@ -65,7 +65,7 @@ func (e *Executor) buildToolIndex() {
e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i] e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i]
} }
} }
e.logger.Info("工具索引构建完成", e.logger.Debug("工具索引构建完成",
zap.Int("totalTools", len(e.config.Tools)), zap.Int("totalTools", len(e.config.Tools)),
zap.Int("enabledTools", len(e.toolIndex)), zap.Int("enabledTools", len(e.toolIndex)),
) )
@@ -73,14 +73,14 @@ func (e *Executor) buildToolIndex() {
// ExecuteTool 执行安全工具 // ExecuteTool 执行安全工具
func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) { func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) {
e.logger.Info("ExecuteTool被调用", e.logger.Debug("ExecuteTool被调用",
zap.String("toolName", toolName), zap.String("toolName", toolName),
zap.Any("args", args), zap.Any("args", args),
) )
// 特殊处理:exec工具直接执行系统命令 // 特殊处理:exec工具直接执行系统命令
if toolName == "exec" { if toolName == "exec" {
e.logger.Info("执行exec工具") e.logger.Debug("执行exec工具")
return e.executeSystemCommand(ctx, args) return e.executeSystemCommand(ctx, args)
} }
@@ -95,7 +95,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName) return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName)
} }
e.logger.Info("找到工具配置", e.logger.Debug("找到工具配置",
zap.String("toolName", toolName), zap.String("toolName", toolName),
zap.String("command", toolConfig.Command), zap.String("command", toolConfig.Command),
zap.Strings("args", toolConfig.Args), zap.Strings("args", toolConfig.Args),
@@ -103,7 +103,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
// 特殊处理:内部工具(command 以 "internal:" 开头) // 特殊处理:内部工具(command 以 "internal:" 开头)
if strings.HasPrefix(toolConfig.Command, "internal:") { if strings.HasPrefix(toolConfig.Command, "internal:") {
e.logger.Info("执行内部工具", e.logger.Debug("执行内部工具",
zap.String("toolName", toolName), zap.String("toolName", toolName),
zap.String("command", toolConfig.Command), zap.String("command", toolConfig.Command),
) )
@@ -113,7 +113,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
// 构建命令 - 根据工具类型使用不同的参数格式 // 构建命令 - 根据工具类型使用不同的参数格式
cmdArgs := e.buildCommandArgs(toolName, toolConfig, args) cmdArgs := e.buildCommandArgs(toolName, toolConfig, args)
e.logger.Info("构建命令参数完成", e.logger.Debug("构建命令参数完成",
zap.String("toolName", toolName), zap.String("toolName", toolName),
zap.Strings("cmdArgs", cmdArgs), zap.Strings("cmdArgs", cmdArgs),
zap.Int("argsCount", len(cmdArgs)), zap.Int("argsCount", len(cmdArgs)),
@@ -142,7 +142,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
attachNonInteractiveStdin(cmd) attachNonInteractiveStdin(cmd)
_ = prepareShellCmdSession(cmd) _ = prepareShellCmdSession(cmd)
e.logger.Info("执行安全工具", e.logger.Debug("执行安全工具",
zap.String("tool", toolName), zap.String("tool", toolName),
zap.Strings("args", cmdArgs), zap.Strings("args", cmdArgs),
) )
@@ -180,7 +180,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
if exitCode != nil && toolConfig.AllowedExitCodes != nil { if exitCode != nil && toolConfig.AllowedExitCodes != nil {
for _, allowedCode := range toolConfig.AllowedExitCodes { for _, allowedCode := range toolConfig.AllowedExitCodes {
if *exitCode == allowedCode { if *exitCode == allowedCode {
e.logger.Info("工具执行完成(退出码在允许列表中)", e.logger.Debug("工具执行完成(退出码在允许列表中)",
zap.String("tool", toolName), zap.String("tool", toolName),
zap.Int("exitCode", *exitCode), zap.Int("exitCode", *exitCode),
zap.String("output", string(output)), zap.String("output", string(output)),
@@ -215,7 +215,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
}, nil }, nil
} }
e.logger.Info("工具执行成功", e.logger.Debug("工具执行成功",
zap.String("tool", toolName), zap.String("tool", toolName),
zap.String("output", string(output)), zap.String("output", string(output)),
) )
@@ -233,7 +233,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
// RegisterTools 注册工具到MCP服务器 // RegisterTools 注册工具到MCP服务器
func (e *Executor) RegisterTools(mcpServer *mcp.Server) { func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
e.logger.Info("开始注册工具", e.logger.Debug("开始注册工具",
zap.Int("totalTools", len(e.config.Tools)), zap.Int("totalTools", len(e.config.Tools)),
zap.Int("enabledTools", len(e.toolIndex)), zap.Int("enabledTools", len(e.toolIndex)),
) )
@@ -281,7 +281,7 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
} }
handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
e.logger.Info("工具handler被调用", e.logger.Debug("工具handler被调用",
zap.String("toolName", toolName), zap.String("toolName", toolName),
zap.Any("args", args), zap.Any("args", args),
) )
@@ -289,14 +289,14 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
} }
mcpServer.RegisterTool(tool, handler) mcpServer.RegisterTool(tool, handler)
e.logger.Info("注册安全工具成功", e.logger.Debug("注册安全工具成功",
zap.String("tool", toolConfigCopy.Name), zap.String("tool", toolConfigCopy.Name),
zap.String("command", toolConfigCopy.Command), zap.String("command", toolConfigCopy.Command),
zap.Int("index", i), zap.Int("index", i),
) )
} }
e.logger.Info("工具注册完成", e.logger.Debug("工具注册完成",
zap.Int("registeredCount", len(e.config.Tools)), zap.Int("registeredCount", len(e.config.Tools)),
) )
} }
+24
View File
@@ -0,0 +1,24 @@
package security
import (
"crypto/rand"
"encoding/base64"
)
// GenerateStrongPassword returns a URL-safe random password of the given length.
func GenerateStrongPassword(length int) (string, error) {
if length <= 0 {
length = 24
}
randomBytes := make([]byte, length)
if _, err := rand.Read(randomBytes); err != nil {
return "", err
}
password := base64.RawURLEncoding.EncodeToString(randomBytes)
if len(password) > length {
password = password[:length]
}
return password, nil
}
+67
View File
@@ -0,0 +1,67 @@
package termout
import (
"fmt"
"os"
"strings"
)
// StartupWebUIOptions configures the startup Web UI banner.
type StartupWebUIOptions struct {
Scheme string
Port int
SelfSigned bool
HTTPRedirect bool
}
// PrintConfigCreated prints a short notice when config.yaml is bootstrapped.
func PrintConfigCreated() {
s := New(os.Stdout)
s.Println("")
s.Println(s.Green("✔ ") + s.Bold("已创建 config.yaml") + s.Dim("(来自 config.example.yaml"))
s.BlankLine()
}
// PrintStartupWebUI prints a colored startup banner for the Web UI.
func PrintStartupWebUI(opts StartupWebUIOptions) {
s := New(os.Stdout)
scheme := opts.Scheme
if scheme == "" {
scheme = "http"
}
port := opts.Port
if port <= 0 {
port = 8080
}
url := fmt.Sprintf("%s://127.0.0.1:%d/", scheme, port)
s.BlankLine()
s.Println(s.Bold(s.Cyan("CYBERSTRIKE AI")) + s.Dim(" / secure workspace"))
s.Println(s.Dim(strings.Repeat("─", 60)))
s.Println(s.Green("● ONLINE") + " " + s.Bold(s.White(url)))
if opts.SelfSigned {
s.Println(s.Dim(" TLS ") + s.Yellow("self-signed") + s.Dim(" · accept the browser warning once"))
}
if opts.HTTPRedirect {
s.Println(s.Dim(" Redirect ") + fmt.Sprintf("http://127.0.0.1:%d/ → HTTPS", port))
}
s.BlankLine()
}
// PrintBootstrapAdminCredentials prints the initial admin password banner.
func PrintBootstrapAdminCredentials(password string) {
password = strings.TrimSpace(password)
if password == "" {
return
}
s := New(os.Stdout)
s.Println(s.Bold(s.Yellow("ADMIN SETUP REQUIRED")))
s.Println(s.Dim(strings.Repeat("─", 60)))
s.Println(s.Dim(" Username ") + s.Bold(s.White("admin")))
s.Println(s.Dim(" Password ") + s.Bold(s.Yellow(password)))
s.BlankLine()
s.Println(s.Yellow(" ! ") + s.White("Store this password securely. It is shown only once."))
s.Println(s.Dim(" Change it in Settings immediately after signing in."))
s.BlankLine()
}
+76
View File
@@ -0,0 +1,76 @@
package termout
import (
"strings"
"testing"
)
func TestDisplayWidthEmoji(t *testing.T) {
if got := displayWidth("🚀"); got != 2 {
t.Fatalf("displayWidth(emoji) = %d, want 2", got)
}
if got := displayWidth("ab"); got != 2 {
t.Fatalf("displayWidth(ab) = %d, want 2", got)
}
}
func TestDisplayWidthIgnoresANSI(t *testing.T) {
s := New(nil)
colored := s.Bold("admin")
if got := displayWidth(colored); got != 5 {
t.Fatalf("displayWidth colored = %d, want 5", got)
}
}
func TestPadRightDisplay(t *testing.T) {
got := padRightDisplay("pwd", 10)
if displayWidth(got) != 10 {
t.Fatalf("padded width = %d, want 10", displayWidth(got))
}
}
func TestColorDisabledWithoutTTY(t *testing.T) {
s := New(nil)
if s.enabled {
t.Fatal("expected colors disabled for nil writer")
}
if got := s.Cyan("x"); got != "x" {
t.Fatalf("Cyan without TTY = %q, want plain text", got)
}
}
func TestPrintBootstrapAdminCredentialsEmpty(t *testing.T) {
PrintBootstrapAdminCredentials(" ")
}
func TestPrintStartupWebUIOptions(t *testing.T) {
PrintStartupWebUI(StartupWebUIOptions{
Scheme: "https",
Port: 8080,
SelfSigned: true,
HTTPRedirect: true,
})
}
func TestBoxRowAlignedWidth(t *testing.T) {
s := New(nil)
rows := []string{
s.Bold("CyberStrikeAI") + s.White(" is ready"),
s.Dim("Web UI ") + s.Bold("https://127.0.0.1:8080/"),
}
inner := maxDisplayWidth(rows...)
for _, row := range rows {
line := s.boxRow(inner, row)
if !strings.Contains(line, "│") {
t.Fatalf("box row missing border: %q", line)
}
}
}
func TestMaxDisplayWidth(t *testing.T) {
short := "abc"
long := "https://127.0.0.1:8080/"
if got := maxDisplayWidth(short, long); got != displayWidth(long) {
t.Fatalf("maxDisplayWidth = %d, want %d", got, displayWidth(long))
}
}
+108
View File
@@ -0,0 +1,108 @@
package termout
import (
"fmt"
"io"
"os"
"strings"
)
const (
codeReset = "\033[0m"
codeBold = "\033[1m"
codeDim = "\033[2m"
codeRed = "\033[31m"
codeGreen = "\033[32m"
codeYellow = "\033[33m"
codeBlue = "\033[34m"
codeCyan = "\033[36m"
codeWhite = "\033[97m"
)
// Style wraps ANSI styling with TTY / NO_COLOR awareness.
type Style struct {
out io.Writer
enabled bool
}
// New creates a Style writing to out (typically os.Stdout).
func New(out io.Writer) *Style {
return &Style{out: out, enabled: colorEnabled(out)}
}
func colorEnabled(w io.Writer) bool {
if strings.TrimSpace(os.Getenv("NO_COLOR")) != "" {
return false
}
force := strings.TrimSpace(os.Getenv("FORCE_COLOR"))
if force == "1" || strings.EqualFold(force, "true") || strings.EqualFold(force, "yes") {
return true
}
f, ok := w.(*os.File)
if !ok {
return false
}
stat, err := f.Stat()
if err != nil {
return false
}
return stat.Mode()&os.ModeCharDevice != 0
}
func (s *Style) paint(code, text string) string {
if !s.enabled || text == "" {
return text
}
return code + text + codeReset
}
func (s *Style) Bold(text string) string { return s.paint(codeBold, text) }
func (s *Style) Dim(text string) string { return s.paint(codeDim, text) }
func (s *Style) Red(text string) string { return s.paint(codeRed, text) }
func (s *Style) Green(text string) string { return s.paint(codeGreen, text) }
func (s *Style) Yellow(text string) string { return s.paint(codeYellow, text) }
func (s *Style) Blue(text string) string { return s.paint(codeBlue, text) }
func (s *Style) Cyan(text string) string { return s.paint(codeCyan, text) }
func (s *Style) White(text string) string { return s.paint(codeWhite, text) }
func (s *Style) Println(text string) {
_, _ = fmt.Fprintln(s.out, text)
}
func (s *Style) Printf(format string, args ...interface{}) {
_, _ = fmt.Fprintf(s.out, format, args...)
}
func (s *Style) BlankLine() {
s.Println("")
}
func (s *Style) boxTop(innerWidth int) string {
return s.Cyan("╭" + strings.Repeat("─", innerWidth+2) + "╮")
}
func (s *Style) boxBottom(innerWidth int) string {
return s.Cyan("╰" + strings.Repeat("─", innerWidth+2) + "╯")
}
func (s *Style) boxRow(innerWidth int, content string) string {
return s.Cyan("│ ") + padRightDisplay(content, innerWidth) + s.Cyan(" │")
}
func (s *Style) printBox(rows []string, minInner, maxInner int) {
inner := maxDisplayWidth(rows...)
if inner < minInner {
inner = minInner
}
if maxInner > 0 && inner > maxInner {
inner = maxInner
}
s.BlankLine()
s.Println(s.boxTop(inner))
for _, row := range rows {
s.Println(s.boxRow(inner, row))
}
s.Println(s.boxBottom(inner))
s.BlankLine()
}
+73
View File
@@ -0,0 +1,73 @@
package termout
import (
"regexp"
"strings"
"unicode/utf8"
"golang.org/x/text/width"
)
var ansiEscapeRe = regexp.MustCompile(`\x1b\[[0-9;]*m`)
// displayWidth returns the terminal display width of text, ignoring ANSI codes.
func displayWidth(text string) int {
plain := ansiEscapeRe.ReplaceAllString(text, "")
w := 0
for _, r := range plain {
w += runeDisplayWidth(r)
}
return w
}
func runeDisplayWidth(r rune) int {
if r == utf8.RuneError {
return 0
}
// Most emoji / symbols render as double-width in modern terminals.
if isEmojiLikeRune(r) {
return 2
}
switch width.LookupRune(r).Kind() {
case width.EastAsianWide, width.EastAsianFullwidth:
return 2
default:
return 1
}
}
func isEmojiLikeRune(r rune) bool {
switch {
case r >= 0x1F300 && r <= 0x1FAFF: // pictographs / emoji
return true
case r >= 0x2600 && r <= 0x27BF: // misc symbols
return true
case r >= 0x2300 && r <= 0x23FF: // misc technical (⌚ etc.)
return true
case r >= 0x2B50 && r <= 0x2B55:
return true
default:
return false
}
}
func padRightDisplay(text string, target int) string {
if target <= 0 {
return ""
}
gap := target - displayWidth(text)
if gap <= 0 {
return text
}
return text + strings.Repeat(" ", gap)
}
func maxDisplayWidth(rows ...string) int {
max := 0
for _, row := range rows {
if w := displayWidth(row); w > max {
max = w
}
}
return max
}
+1 -1
View File
@@ -89,7 +89,7 @@ func RegisterAnalyzeImageTool(mcpServer *mcp.Server, cfg *config.Config, logger
mcpServer.RegisterTool(tool, handler) mcpServer.RegisterTool(tool, handler)
if logger != nil { if logger != nil {
logger.Info("vision: analyze_image 工具已注册", zap.String("model", cfg.Vision.Model)) logger.Debug("vision: analyze_image 工具已注册", zap.String("model", cfg.Vision.Model))
} }
} }
+85
View File
@@ -2,6 +2,7 @@ package workflow
import ( import (
"context" "context"
"fmt"
"path/filepath" "path/filepath"
"strings" "strings"
"testing" "testing"
@@ -9,6 +10,7 @@ import (
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database" "cyberstrike-ai/internal/database"
"github.com/cloudwego/eino/compose"
"go.uber.org/zap" "go.uber.org/zap"
) )
@@ -235,6 +237,89 @@ func TestExecuteEinoGraph_linearStartOutput(t *testing.T) {
} }
} }
func TestExecuteEinoGraph_checkpointRestoresStartOutput(t *testing.T) {
ctx := context.Background()
checkpointStore, err := newFileCheckPointStore(t.TempDir())
if err != nil {
t.Fatalf("new checkpoint store: %v", err)
}
state := newWorkflowLocalState(map[string]interface{}{"message": "ping"}, "run-checkpoint")
node := graphNode{ID: "start-1", Type: "start"}
wf := compose.NewWorkflow[WorkflowInput, WorkflowOutput](
compose.WithGenLocalState(func(context.Context) *WorkflowLocalState { return state }),
)
start := wf.AddLambdaNode("start-1", compose.InvokableLambda(func(_ context.Context, input WorkflowInput) (WorkflowNodeOutput, error) {
result := startOutputMap(node, input.Message, input.ConversationID, input.ProjectID)
state.NodeOutputs[node.ID] = result
state.NodeOutputs["condition-1"] = conditionOutputMap(graphNode{ID: "condition-1", Type: "condition"}, "{{inputs.message}} == ping", true)
state.NodeOutputs["tool-1"] = toolOutputMap(graphNode{ID: "tool-1", Type: "tool"}, "tool result", "lookup", map[string]any{"id": "1"}, "exec-1", false)
state.NodeOutputs["agent-1"] = agentOutputMap(graphNode{ID: "agent-1", Type: "agent"}, "agent result", "chat", []string{"exec-1"})
state.NodeOutputs["hitl-1"] = hitlOutputMap(graphNode{ID: "hitl-1", Type: "hitl"}, "completed", "approved", "continue?", "reviewer", true)
state.NodeOutputs["output-1"] = outputNodeOutputMap(graphNode{ID: "output-1", Type: "output"}, "result", "ping")
state.NodeOutputs["end-1"] = endOutputMap(graphNode{ID: "end-1", Type: "end"}, "done")
state.LastOutput = result
state.Outputs["seed"] = "preserved"
return result, nil
}))
outputNode := wf.AddLambdaNode("out-1", compose.InvokableLambda(func(_ context.Context, input WorkflowNodeOutput) (WorkflowNodeOutput, error) {
return input, nil
}))
start.AddInput(compose.START)
outputNode.AddInput("start-1")
wf.End().AddInput("out-1", compose.ToField("out-1"))
runnable, err := wf.Compile(ctx,
compose.WithCheckPointStore(checkpointStore),
compose.WithInterruptAfterNodes([]string{"start-1"}),
)
if err != nil {
t.Fatalf("compile: %v", err)
}
_, err = runnable.Invoke(ctx, workflowInputFromMap(state.Inputs), compose.WithCheckPointID("run-checkpoint"))
info, ok := compose.ExtractInterruptInfo(err)
if !ok {
t.Fatalf("invoke error = %v, want checkpoint interrupt", err)
}
restored, ok := info.State.(*WorkflowLocalState)
if !ok {
t.Fatalf("checkpoint state = %T, want *WorkflowLocalState", info.State)
}
for nodeID, wantType := range map[string]string{
"start-1": "StartOutput",
"condition-1": "ConditionOutput",
"tool-1": "ToolOutput",
"agent-1": "AgentOutput",
"hitl-1": "HITLOutput",
"output-1": "OutputNodeOutput",
"end-1": "NodeOutputEnvelope",
} {
if got := fmt.Sprintf("%T", restored.NodeOutputs[nodeID]["typed"]); got != "workflow."+wantType {
t.Fatalf("restored %s typed output = %s, want workflow.%s", nodeID, got, wantType)
}
}
if got := valueFromPath("previous.message", restored); got != "ping" {
t.Fatalf("restored previous.message = %v, want ping", got)
}
if got := valueFromPath("inputs.message", restored); got != "ping" {
t.Fatalf("restored inputs.message = %v, want ping", got)
}
if got := valueFromPath("outputs.seed", restored); got != "preserved" {
t.Fatalf("restored outputs.seed = %v, want preserved", got)
}
result, err := runnable.Invoke(ctx, WorkflowInput{}, compose.WithCheckPointID("run-checkpoint"))
if err != nil {
t.Fatalf("resume checkpoint: %v", err)
}
output, ok := result["out-1"].(map[string]any)
if !ok {
t.Fatalf("resumed output type = %T, want map[string]any", result["out-1"])
}
if got := output["output"]; got != "ping" {
t.Fatalf("resumed output = %v, want ping", got)
}
}
func TestExecuteEinoGraph_conditionBranch(t *testing.T) { func TestExecuteEinoGraph_conditionBranch(t *testing.T) {
ctx := context.Background() ctx := context.Background()
SetCheckpointDir(t.TempDir()) SetCheckpointDir(t.TempDir())
+7
View File
@@ -11,6 +11,13 @@ import (
func init() { func init() {
schema.RegisterName[*WorkflowLocalState]("_cyberstrike_workflow_local_state") schema.RegisterName[*WorkflowLocalState]("_cyberstrike_workflow_local_state")
schema.RegisterName[NodeOutputEnvelope]("_cyberstrike_workflow_node_output_envelope")
schema.RegisterName[StartOutput]("_cyberstrike_workflow_start_output")
schema.RegisterName[ConditionOutput]("_cyberstrike_workflow_condition_output")
schema.RegisterName[ToolOutput]("_cyberstrike_workflow_tool_output")
schema.RegisterName[AgentOutput]("_cyberstrike_workflow_agent_output")
schema.RegisterName[HITLOutput]("_cyberstrike_workflow_hitl_output")
schema.RegisterName[OutputNodeOutput]("_cyberstrike_workflow_output_node_output")
} }
// WorkflowLocalState is the Eino WithGenLocalState payload (checkpoint-serializable). // WorkflowLocalState is the Eino WithGenLocalState payload (checkpoint-serializable).