From ee4676a591f0db511fbc14d3a8d2b31ca4eca90c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 15 Aug 2026 02:14:36 +0800 Subject: [PATCH] Add files via upload --- internal/audit/conversation_create.go | 55 ++ internal/audit/meta.go | 9 + internal/audit/record.go | 29 + internal/audit/resource_availability.go | 86 ++ internal/audit/retention.go | 27 + internal/audit/sanitize.go | 58 ++ internal/audit/service.go | 177 ++++ internal/audit/throttle.go | 55 ++ internal/audit/types.go | 16 + internal/authctx/principal.go | 72 ++ internal/hitl/retention.go | 71 ++ internal/hitl/retention_test.go | 50 + internal/knowledge/chunk_eino.go | 67 ++ internal/knowledge/eino_meta.go | 129 +++ internal/knowledge/eino_meta_test.go | 14 + internal/knowledge/eino_pipeline_retriever.go | 96 ++ internal/knowledge/eino_retrieve_chain.go | 24 + .../knowledge/eino_retrieve_chain_test.go | 23 + internal/knowledge/eino_retriever_adapter.go | 173 ++++ internal/knowledge/eino_sqlite_indexer.go | 142 +++ internal/knowledge/embedder.go | 251 +++++ internal/knowledge/index_pipeline.go | 91 ++ internal/knowledge/index_pipeline_test.go | 21 + internal/knowledge/indexer.go | 435 +++++++++ .../knowledge/indexer_rebuild_state_test.go | 20 + internal/knowledge/manager.go | 885 ++++++++++++++++++ internal/knowledge/rerank_http.go | 226 +++++ internal/knowledge/rerank_http_test.go | 97 ++ internal/knowledge/retrieval_postprocess.go | 216 +++++ .../knowledge/retrieval_postprocess_test.go | 62 ++ internal/knowledge/retriever.go | 334 +++++++ internal/knowledge/schema_migrate.go | 51 + internal/knowledge/tool.go | 323 +++++++ internal/knowledge/types.go | 123 +++ internal/knowledge/wire_retriever.go | 76 ++ internal/projectprompt/blackboard.go | 132 +++ internal/projectprompt/shell_tools.go | 11 + internal/reasoning/eino.go | 428 +++++++++ internal/reasoning/eino_test.go | 424 +++++++++ internal/workflow/agent_subgraph.go | 48 + internal/workflow/bindings.go | 153 +++ internal/workflow/checkpoint_store.go | 69 ++ internal/workflow/draft_generator.go | 782 ++++++++++++++++ internal/workflow/draft_generator_test.go | 213 +++++ internal/workflow/dry_run.go | 173 ++++ internal/workflow/eino_branch.go | 107 +++ internal/workflow/eino_callbacks.go | 22 + internal/workflow/eino_compile.go | 243 +++++ internal/workflow/eino_compile_test.go | 412 ++++++++ internal/workflow/eino_runtime.go | 64 ++ internal/workflow/engine.go | 239 +++++ internal/workflow/errors.go | 24 + internal/workflow/expression.go | 186 ++++ internal/workflow/expression_join_test.go | 107 +++ internal/workflow/graph_types.go | 153 +++ internal/workflow/hitl_wait.go | 119 +++ internal/workflow/join.go | 117 +++ internal/workflow/jsonpath.go | 115 +++ internal/workflow/metrics.go | 57 ++ internal/workflow/node_exec.go | 153 +++ internal/workflow/nodes.go | 333 +++++++ internal/workflow/package/exporter.go | 98 ++ internal/workflow/package/exporter_test.go | 58 ++ internal/workflow/package/inspector.go | 225 +++++ internal/workflow/package/inspector_test.go | 154 +++ internal/workflow/package/manifest.go | 78 ++ internal/workflow/runner.go | 223 +++++ internal/workflow/state.go | 214 +++++ internal/workflow/structured_outputs.go | 150 +++ internal/workflow/tool_output_budget_test.go | 23 + internal/workflow/types.go | 74 ++ internal/workflow/validation.go | 366 ++++++++ 72 files changed, 11131 insertions(+) create mode 100644 internal/audit/conversation_create.go create mode 100644 internal/audit/meta.go create mode 100644 internal/audit/record.go create mode 100644 internal/audit/resource_availability.go create mode 100644 internal/audit/retention.go create mode 100644 internal/audit/sanitize.go create mode 100644 internal/audit/service.go create mode 100644 internal/audit/throttle.go create mode 100644 internal/audit/types.go create mode 100644 internal/authctx/principal.go create mode 100644 internal/hitl/retention.go create mode 100644 internal/hitl/retention_test.go create mode 100644 internal/knowledge/chunk_eino.go create mode 100644 internal/knowledge/eino_meta.go create mode 100644 internal/knowledge/eino_meta_test.go create mode 100644 internal/knowledge/eino_pipeline_retriever.go create mode 100644 internal/knowledge/eino_retrieve_chain.go create mode 100644 internal/knowledge/eino_retrieve_chain_test.go create mode 100644 internal/knowledge/eino_retriever_adapter.go create mode 100644 internal/knowledge/eino_sqlite_indexer.go create mode 100644 internal/knowledge/embedder.go create mode 100644 internal/knowledge/index_pipeline.go create mode 100644 internal/knowledge/index_pipeline_test.go create mode 100644 internal/knowledge/indexer.go create mode 100644 internal/knowledge/indexer_rebuild_state_test.go create mode 100644 internal/knowledge/manager.go create mode 100644 internal/knowledge/rerank_http.go create mode 100644 internal/knowledge/rerank_http_test.go create mode 100644 internal/knowledge/retrieval_postprocess.go create mode 100644 internal/knowledge/retrieval_postprocess_test.go create mode 100644 internal/knowledge/retriever.go create mode 100644 internal/knowledge/schema_migrate.go create mode 100644 internal/knowledge/tool.go create mode 100644 internal/knowledge/types.go create mode 100644 internal/knowledge/wire_retriever.go create mode 100644 internal/projectprompt/blackboard.go create mode 100644 internal/projectprompt/shell_tools.go create mode 100644 internal/reasoning/eino.go create mode 100644 internal/reasoning/eino_test.go create mode 100644 internal/workflow/agent_subgraph.go create mode 100644 internal/workflow/bindings.go create mode 100644 internal/workflow/checkpoint_store.go create mode 100644 internal/workflow/draft_generator.go create mode 100644 internal/workflow/draft_generator_test.go create mode 100644 internal/workflow/dry_run.go create mode 100644 internal/workflow/eino_branch.go create mode 100644 internal/workflow/eino_callbacks.go create mode 100644 internal/workflow/eino_compile.go create mode 100644 internal/workflow/eino_compile_test.go create mode 100644 internal/workflow/eino_runtime.go create mode 100644 internal/workflow/engine.go create mode 100644 internal/workflow/errors.go create mode 100644 internal/workflow/expression.go create mode 100644 internal/workflow/expression_join_test.go create mode 100644 internal/workflow/graph_types.go create mode 100644 internal/workflow/hitl_wait.go create mode 100644 internal/workflow/join.go create mode 100644 internal/workflow/jsonpath.go create mode 100644 internal/workflow/metrics.go create mode 100644 internal/workflow/node_exec.go create mode 100644 internal/workflow/nodes.go create mode 100644 internal/workflow/package/exporter.go create mode 100644 internal/workflow/package/exporter_test.go create mode 100644 internal/workflow/package/inspector.go create mode 100644 internal/workflow/package/inspector_test.go create mode 100644 internal/workflow/package/manifest.go create mode 100644 internal/workflow/runner.go create mode 100644 internal/workflow/state.go create mode 100644 internal/workflow/structured_outputs.go create mode 100644 internal/workflow/tool_output_budget_test.go create mode 100644 internal/workflow/types.go create mode 100644 internal/workflow/validation.go diff --git a/internal/audit/conversation_create.go b/internal/audit/conversation_create.go new file mode 100644 index 00000000..82e19b54 --- /dev/null +++ b/internal/audit/conversation_create.go @@ -0,0 +1,55 @@ +package audit + +import ( + "strings" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" +) + +// RegisterConversationCreateHook records platform audit rows for every new conversation. +func RegisterConversationCreateHook(s *Service) { + if s == nil { + return + } + database.SetConversationCreateHook(func(conv *database.Conversation, meta database.ConversationCreateMeta) { + detail := map[string]interface{}{ + "title": conv.Title, + "source": meta.Source, + } + if meta.WebShellConnectionID != "" { + detail["webshell_connection_id"] = meta.WebShellConnectionID + } + s.Record(nil, Entry{ + Category: "conversation", + Action: "create", + Result: "success", + Message: "创建对话", + ResourceType: "conversation", + ResourceID: conv.ID, + Detail: detail, + ClientIP: meta.ClientIP, + SessionHint: meta.SessionHint, + }) + }) +} + +// ConversationCreateMeta builds audit metadata for conversation creation. +func ConversationCreateMeta(source string) database.ConversationCreateMeta { + return database.ConversationCreateMeta{Source: strings.TrimSpace(source)} +} + +// ConversationCreateMetaFromGin includes client IP and session hint when available. +func ConversationCreateMetaFromGin(c *gin.Context, source string) database.ConversationCreateMeta { + m := ConversationCreateMeta(source) + if c == nil { + return m + } + m.ClientIP = c.ClientIP() + if token := c.GetString(security.ContextAuthTokenKey); token != "" { + m.SessionHint = sessionHint(token) + } + return m +} diff --git a/internal/audit/meta.go b/internal/audit/meta.go new file mode 100644 index 00000000..33649e0c --- /dev/null +++ b/internal/audit/meta.go @@ -0,0 +1,9 @@ +package audit + +// RetentionDays returns configured retention; 0 means keep forever. +func (s *Service) RetentionDays() int { + if s == nil || s.cfg == nil { + return 0 + } + return s.cfg.Audit.RetentionDaysEffective() +} diff --git a/internal/audit/record.go b/internal/audit/record.go new file mode 100644 index 00000000..b1c1ad40 --- /dev/null +++ b/internal/audit/record.go @@ -0,0 +1,29 @@ +package audit + +import "github.com/gin-gonic/gin" + +// RecordAction writes a platform audit row with common defaults. +func (s *Service) RecordAction(c *gin.Context, category, action, result, message, resourceType, resourceID string, detail map[string]interface{}) { + if s == nil { + return + } + s.Record(c, Entry{ + Category: category, + Action: action, + Result: result, + Message: message, + ResourceType: resourceType, + ResourceID: resourceID, + Detail: detail, + }) +} + +// RecordOK is a shorthand for successful operations. +func (s *Service) RecordOK(c *gin.Context, category, action, message, resourceType, resourceID string, detail map[string]interface{}) { + s.RecordAction(c, category, action, "success", message, resourceType, resourceID, detail) +} + +// RecordFail is a shorthand for failed operations. +func (s *Service) RecordFail(c *gin.Context, category, action, message string, detail map[string]interface{}) { + s.RecordAction(c, category, action, "failure", message, "", "", detail) +} diff --git a/internal/audit/resource_availability.go b/internal/audit/resource_availability.go new file mode 100644 index 00000000..3b22871f --- /dev/null +++ b/internal/audit/resource_availability.go @@ -0,0 +1,86 @@ +package audit + +import ( + "strings" + + "cyberstrike-ai/internal/database" +) + +var auditActionsResourceRemoved = map[string]bool{ + "delete": true, + "item_delete": true, + "connection_delete": true, + "listener_delete": true, + "session_delete": true, + "task_delete": true, + "execution_delete": true, + "execution_delete_batch": true, + "delete_queue": true, + "delete_batch_task": true, + "markdown_delete": true, +} + +// ApplyResourceAvailability sets log.ResourceAvailable when the linked resource can be checked. +func ApplyResourceAvailability(db *database.DB, log *database.AuditLog) { + if log == nil || strings.TrimSpace(log.ResourceID) == "" { + return + } + if auditActionsResourceRemoved[log.Action] { + f := false + log.ResourceAvailable = &f + return + } + if db == nil { + return + } + available, known := resourceStillExists(db, log.ResourceType, log.ResourceID) + if known { + log.ResourceAvailable = &available + } +} + +func resourceStillExists(db *database.DB, resourceType, resourceID string) (bool, bool) { + resourceID = strings.TrimSpace(resourceID) + if resourceID == "" { + return false, false + } + t := strings.TrimSpace(resourceType) + if t == "" { + if len(resourceID) > 8 && !strings.HasPrefix(resourceID, "c2_") { + t = "conversation" + } else { + return false, false + } + } + switch t { + case "conversation": + ok, err := db.ConversationExists(resourceID) + return ok, err == nil + case "vulnerability": + _, err := db.GetVulnerability(resourceID) + if err != nil { + return false, strings.Contains(err.Error(), "不存在") + } + return true, true + case "batch_queue": + _, err := db.GetBatchQueue(resourceID) + return err == nil, true + case "c2_listener": + _, err := db.GetC2Listener(resourceID) + return err == nil, true + case "c2_session": + _, err := db.GetC2Session(resourceID) + return err == nil, true + case "c2_task": + _, err := db.GetC2Task(resourceID) + return err == nil, true + case "webshell_connection": + c, err := db.GetWebshellConnection(resourceID) + return err == nil && c != nil, true + case "tool_execution": + _, err := db.GetToolExecution(resourceID) + return err == nil, true + default: + return false, false + } +} diff --git a/internal/audit/retention.go b/internal/audit/retention.go new file mode 100644 index 00000000..f882595c --- /dev/null +++ b/internal/audit/retention.go @@ -0,0 +1,27 @@ +package audit + +import ( + "time" + + "go.uber.org/zap" +) + +// auditRetentionPurgeInterval is how often PurgeExpired runs while the process is up (startup also purges once). +const auditRetentionPurgeInterval = time.Hour + +// StartRetentionLoop periodically purges expired audit rows. +func StartRetentionLoop(s *Service, logger *zap.Logger) { + if s == nil { + return + } + go func() { + ticker := time.NewTicker(auditRetentionPurgeInterval) + defer ticker.Stop() + for range ticker.C { + s.PurgeExpired() + if logger != nil { + logger.Debug("audit retention tick completed") + } + } + }() +} diff --git a/internal/audit/sanitize.go b/internal/audit/sanitize.go new file mode 100644 index 00000000..34f2b439 --- /dev/null +++ b/internal/audit/sanitize.go @@ -0,0 +1,58 @@ +package audit + +import ( + "encoding/json" + "strings" +) + +var sensitiveKeySubstrings = []string{ + "password", "api_key", "apikey", "secret", "token", "authorization", + "credential", "private_key", "access_key", +} + +// SanitizeDetail redacts sensitive keys and truncates serialized size. +func SanitizeDetail(detail map[string]interface{}, maxBytes int) map[string]interface{} { + if detail == nil { + return nil + } + if maxBytes <= 0 { + maxBytes = 8192 + } + out := sanitizeValue("", detail) + if m, ok := out.(map[string]interface{}); ok { + b, _ := json.Marshal(m) + if len(b) > maxBytes { + return map[string]interface{}{ + "_truncated": true, + "_preview": string(b[:maxBytes]), + } + } + return m + } + return map[string]interface{}{"value": out} +} + +func sanitizeValue(key string, v interface{}) interface{} { + kl := strings.ToLower(key) + for _, sub := range sensitiveKeySubstrings { + if strings.Contains(kl, sub) { + return "***" + } + } + switch t := v.(type) { + case map[string]interface{}: + m := make(map[string]interface{}, len(t)) + for k, val := range t { + m[k] = sanitizeValue(k, val) + } + return m + case []interface{}: + arr := make([]interface{}, len(t)) + for i, val := range t { + arr[i] = sanitizeValue(key, val) + } + return arr + default: + return v + } +} diff --git a/internal/audit/service.go b/internal/audit/service.go new file mode 100644 index 00000000..60ea11b3 --- /dev/null +++ b/internal/audit/service.go @@ -0,0 +1,177 @@ +package audit + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Service persists platform audit logs. +type Service struct { + db *database.DB + cfg *config.Config + logger *zap.Logger + failThrottle *failureThrottle +} + +// NewService creates an audit service. +func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service { + return &Service{ + db: db, + cfg: cfg, + logger: logger, + failThrottle: newFailureThrottle(), + } +} + +// Enabled reports whether audit persistence is on. +func (s *Service) Enabled() bool { + if s == nil || s.cfg == nil { + return false + } + return s.cfg.Audit.EnabledEffective() +} + +// Record writes one audit row from a Gin request context. +func (s *Service) Record(c *gin.Context, e Entry) { + if s == nil || !s.Enabled() || s.db == nil { + return + } + if strings.TrimSpace(e.Category) == "" || strings.TrimSpace(e.Action) == "" { + return + } + if e.Result == "failure" && !s.allowFailureAudit(c, e) { + return + } + if strings.TrimSpace(e.Result) == "" { + e.Result = "success" + } + if strings.TrimSpace(e.Level) == "" { + if e.Result == "failure" { + e.Level = "warn" + } else { + e.Level = "info" + } + } + if strings.TrimSpace(e.Actor) == "" { + if c != nil { + e.Actor = strings.TrimSpace(c.GetString(security.ContextUsernameKey)) + } + if e.Actor == "" { + e.Actor = "admin" + } + } + maxDetail := s.cfg.Audit.MaxDetailBytesEffective() + detail := SanitizeDetail(e.Detail, maxDetail) + + sessionHintVal := e.SessionHint + if sessionHintVal == "" && c != nil { + if token := c.GetString(security.ContextAuthTokenKey); token != "" { + sessionHintVal = sessionHint(token) + } + } + clientIPVal := e.ClientIP + if clientIPVal == "" { + clientIPVal = clientIP(c) + } + + row := &database.AuditLog{ + ID: "audit_" + strings.ReplaceAll(uuid.New().String(), "-", ""), + CreatedAt: time.Now(), + Level: e.Level, + Category: e.Category, + Action: e.Action, + Result: e.Result, + Actor: e.Actor, + SessionHint: sessionHintVal, + ClientIP: clientIPVal, + UserAgent: userAgent(c), + ResourceType: e.ResourceType, + ResourceID: e.ResourceID, + Message: e.Message, + Detail: detail, + } + if err := s.db.AppendAuditLog(row); err != nil && s.logger != nil { + s.logger.Warn("写入审计日志失败", + zap.String("action", e.Action), + zap.Error(err), + ) + } +} + +// RecordSystem writes an audit row without HTTP context (e.g. retention cleanup). +func (s *Service) RecordSystem(e Entry) { + s.Record(nil, e) +} + +// PurgeExpired deletes rows older than retention_days when configured. +func (s *Service) PurgeExpired() { + if s == nil || s.db == nil || s.cfg == nil { + return + } + days := s.cfg.Audit.RetentionDaysEffective() + if days <= 0 { + return + } + cutoff := time.Now().AddDate(0, 0, -days) + n, err := s.db.DeleteAuditLogsBefore(cutoff) + if err != nil { + if s.logger != nil { + s.logger.Warn("清理过期审计日志失败", zap.Error(err)) + } + return + } + if n > 0 && s.logger != nil { + s.logger.Info("已清理过期审计日志", zap.Int64("deleted", n)) + } +} + +// HintFromToken returns a short stable hash prefix for a session token. +func HintFromToken(token string) string { + return sessionHint(token) +} + +func sessionHint(token string) string { + token = strings.TrimSpace(token) + if token == "" { + return "" + } + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:4]) +} + +func (s *Service) allowFailureAudit(c *gin.Context, e Entry) bool { + if !isAuthFailureThrottled(e.Category, e.Action) { + return true + } + cooldown := time.Duration(s.cfg.Audit.AuthFailureCooldownEffective()) * time.Second + key := authFailureThrottleKey(e.Category, e.Action, clientIP(c)) + return s.failThrottle.allow(key, cooldown) +} + +func clientIP(c *gin.Context) string { + if c == nil { + return "" + } + return c.ClientIP() +} + +func userAgent(c *gin.Context) string { + if c == nil { + return "" + } + ua := c.GetHeader("User-Agent") + if len(ua) > 512 { + return ua[:512] + } + return ua +} diff --git a/internal/audit/throttle.go b/internal/audit/throttle.go new file mode 100644 index 00000000..7364e07d --- /dev/null +++ b/internal/audit/throttle.go @@ -0,0 +1,55 @@ +package audit + +import ( + "sync" + "time" +) + +// failureThrottle deduplicates high-frequency failure audit rows (e.g. wrong password). +type failureThrottle struct { + mu sync.Mutex + last map[string]time.Time +} + +func newFailureThrottle() *failureThrottle { + return &failureThrottle{last: make(map[string]time.Time)} +} + +// allow reports whether a row with the given key may be written now. +func (t *failureThrottle) allow(key string, cooldown time.Duration) bool { + if t == nil || cooldown <= 0 || key == "" { + return true + } + now := time.Now() + t.mu.Lock() + defer t.mu.Unlock() + if prev, ok := t.last[key]; ok && now.Sub(prev) < cooldown { + return false + } + t.last[key] = now + if len(t.last) > 4096 { + for k, ts := range t.last { + if now.Sub(ts) > cooldown*2 { + delete(t.last, k) + } + } + } + return true +} + +// authFailureThrottleKey builds a per-IP key for auth failure deduplication. +func authFailureThrottleKey(category, action, clientIP string) string { + return category + ":" + action + ":" + clientIP +} + +func isAuthFailureThrottled(category, action string) bool { + if category != "auth" { + return false + } + switch action { + case "login", "change_password": + return true + default: + return false + } +} diff --git a/internal/audit/types.go b/internal/audit/types.go new file mode 100644 index 00000000..ff83ea58 --- /dev/null +++ b/internal/audit/types.go @@ -0,0 +1,16 @@ +package audit + +// Entry describes one platform audit record (not chat/tool execution bodies). +type Entry struct { + Level string + Category string + Action string + Result string // success | failure + Actor string + SessionHint string + ResourceType string + ResourceID string + Message string + Detail map[string]interface{} + ClientIP string // optional when c is nil (robot, batch, DB hook) +} diff --git a/internal/authctx/principal.go b/internal/authctx/principal.go new file mode 100644 index 00000000..f74199ac --- /dev/null +++ b/internal/authctx/principal.go @@ -0,0 +1,72 @@ +package authctx + +import ( + "context" + "strings" +) + +// Principal is the immutable authorization identity propagated beyond the +// transport layer into Agent, MCP and background task contexts. +type Principal struct { + UserID string + Username string + Permissions map[string]bool + PermissionScopes map[string]string + Scope string +} + +type principalContextKey struct{} + +func NewPrincipal(userID, username, scope string, permissions map[string]bool) Principal { + return NewPrincipalWithScopes(userID, username, scope, permissions, nil) +} + +func NewPrincipalWithScopes(userID, username, scope string, permissions map[string]bool, permissionScopes map[string]string) Principal { + permissionCopy := make(map[string]bool, len(permissions)) + scopeCopy := make(map[string]string, len(permissionScopes)) + for permission, allowed := range permissions { + if allowed { + permissionCopy[permission] = true + if permissionScope := strings.TrimSpace(permissionScopes[permission]); permissionScope != "" { + scopeCopy[permission] = permissionScope + } + } + } + return Principal{ + UserID: strings.TrimSpace(userID), Username: strings.TrimSpace(username), + Scope: strings.TrimSpace(scope), Permissions: permissionCopy, PermissionScopes: scopeCopy, + } +} + +func WithPrincipal(ctx context.Context, principal Principal) context.Context { + if ctx == nil { + ctx = context.Background() + } + if strings.TrimSpace(principal.UserID) == "" { + return ctx + } + return context.WithValue(ctx, principalContextKey{}, principal) +} + +func PrincipalFromContext(ctx context.Context) (Principal, bool) { + if ctx == nil { + return Principal{}, false + } + principal, ok := ctx.Value(principalContextKey{}).(Principal) + return principal, ok && strings.TrimSpace(principal.UserID) != "" +} + +func (p Principal) HasPermission(permission string) bool { + return p.Permissions[strings.TrimSpace(permission)] +} + +// ScopeFor returns the scope attached to the permission that authorizes the +// current action. Falling back to Scope keeps explicit service principals and +// legacy callers compatible without reintroducing cross-role scope widening. +func (p Principal) ScopeFor(permission string) string { + permission = strings.TrimSpace(permission) + if scope := strings.TrimSpace(p.PermissionScopes[permission]); scope != "" { + return scope + } + return strings.TrimSpace(p.Scope) +} diff --git a/internal/hitl/retention.go b/internal/hitl/retention.go new file mode 100644 index 00000000..2746d0f0 --- /dev/null +++ b/internal/hitl/retention.go @@ -0,0 +1,71 @@ +package hitl + +import ( + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +const retentionPurgeInterval = time.Hour + +// Service manages HITL audit log retention (decided hitl_interrupts rows). +type Service struct { + db *database.DB + cfg *config.Config + logger *zap.Logger +} + +// NewService creates a HITL audit log retention service. +func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service { + return &Service{db: db, cfg: cfg, logger: logger} +} + +// RetentionDays returns configured retention; 0 means keep forever. +func (s *Service) RetentionDays() int { + if s == nil || s.cfg == nil { + return config.HitlConfig{}.RetentionDaysEffective() + } + return s.cfg.Hitl.RetentionDaysEffective() +} + +// PurgeExpired deletes decided HITL log rows older than retention_days when configured. +func (s *Service) PurgeExpired() { + if s == nil || s.db == nil || s.cfg == nil { + return + } + days := s.cfg.Hitl.RetentionDaysEffective() + if days <= 0 { + return + } + cutoff := time.Now().AddDate(0, 0, -days) + n, err := s.db.PurgeHitlInterruptLogsBefore(cutoff) + if err != nil { + if s.logger != nil { + s.logger.Warn("清理过期人机协同审计日志失败", zap.Error(err)) + } + return + } + if n > 0 && s.logger != nil { + s.logger.Info("已清理过期人机协同审计日志", zap.Int64("deleted", n), zap.Int("retention_days", days)) + } +} + +// StartRetentionLoop periodically purges expired HITL audit log rows. +func StartRetentionLoop(s *Service, logger *zap.Logger) { + if s == nil { + return + } + go func() { + ticker := time.NewTicker(retentionPurgeInterval) + defer ticker.Stop() + for range ticker.C { + s.PurgeExpired() + if logger != nil { + logger.Debug("hitl audit log retention tick completed") + } + } + }() +} diff --git a/internal/hitl/retention_test.go b/internal/hitl/retention_test.go new file mode 100644 index 00000000..f2db086d --- /dev/null +++ b/internal/hitl/retention_test.go @@ -0,0 +1,50 @@ +package hitl + +import ( + "path/filepath" + "testing" + "time" + + appconfig "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestServicePurgeExpired_respectsZeroRetention(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "hitl.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS hitl_interrupts ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + mode TEXT NOT NULL, + tool_name TEXT NOT NULL, + status TEXT NOT NULL, + decision TEXT, + created_at DATETIME NOT NULL, + decided_at DATETIME + )`); err != nil { + t.Fatalf("create table: %v", err) + } + + old := time.Now().AddDate(0, 0, -100).UTC().Format(time.RFC3339) + if _, err := db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, mode, tool_name, status, decision, created_at, decided_at) + VALUES ('old-1', 'c1', 'approval', 'exec', 'decided', 'approve', ?, ?)`, old, old); err != nil { + t.Fatalf("insert: %v", err) + } + + zero := 0 + svc := NewService(db, &appconfig.Config{ + Hitl: appconfig.HitlConfig{RetentionDays: &zero}, + }, zap.NewNop()) + svc.PurgeExpired() + + if err := db.QueryRow(`SELECT id FROM hitl_interrupts WHERE id = 'old-1'`).Scan(new(string)); err != nil { + t.Fatalf("record should remain when retention_days=0: %v", err) + } +} diff --git a/internal/knowledge/chunk_eino.go b/internal/knowledge/chunk_eino.go new file mode 100644 index 00000000..6592f350 --- /dev/null +++ b/internal/knowledge/chunk_eino.go @@ -0,0 +1,67 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino-ext/components/document/transformer/splitter/markdown" + "github.com/cloudwego/eino-ext/components/document/transformer/splitter/recursive" + "github.com/cloudwego/eino/components/document" + "github.com/pkoukk/tiktoken-go" +) + +func tokenizerLenFunc(embeddingModel string) func(string) int { + fallback := func(s string) int { + r := []rune(s) + if len(r) == 0 { + return 0 + } + return (len(r) + 3) / 4 + } + m := strings.TrimSpace(embeddingModel) + if m == "" { + return fallback + } + tok, err := tiktoken.EncodingForModel(m) + if err != nil { + return fallback + } + return func(s string) int { + return len(tok.Encode(s, nil, nil)) + } +} + +// newKnowledgeSplitter builds an Eino recursive text splitter. LenFunc uses tiktoken for +// embeddingModel when available, else rune/4 approximation. +func newKnowledgeSplitter(chunkSize, overlap int, embeddingModel string) (document.Transformer, error) { + if chunkSize <= 0 { + return nil, fmt.Errorf("chunk size must be positive") + } + if overlap < 0 { + overlap = 0 + } + return recursive.NewSplitter(context.Background(), &recursive.Config{ + ChunkSize: chunkSize, + OverlapSize: overlap, + LenFunc: tokenizerLenFunc(embeddingModel), + Separators: []string{ + "\n\n", "\n## ", "\n### ", "\n#### ", "\n", + "。", "！", "？", ". ", "? ", "! ", + " ", + }, + }) +} + +// newMarkdownHeaderSplitter Eino-ext Markdown 按标题切分（#～####），适合技术/Markdown 知识库。 +func newMarkdownHeaderSplitter(ctx context.Context) (document.Transformer, error) { + return markdown.NewHeaderSplitter(ctx, &markdown.HeaderConfig{ + Headers: map[string]string{ + "#": "h1", + "##": "h2", + "###": "h3", + "####": "h4", + }, + TrimHeaders: false, + }) +} diff --git a/internal/knowledge/eino_meta.go b/internal/knowledge/eino_meta.go new file mode 100644 index 00000000..0ee7c41b --- /dev/null +++ b/internal/knowledge/eino_meta.go @@ -0,0 +1,129 @@ +package knowledge + +import ( + "fmt" + "strings" +) + +// Document metadata keys for Eino schema.Document flowing through the RAG pipeline. +const ( + metaKBCategory = "kb_category" + metaKBTitle = "kb_title" + metaKBItemID = "kb_item_id" + metaKBChunkIndex = "kb_chunk_index" + metaSimilarity = "similarity" +) + +// DSL keys for [VectorEinoRetriever.Retrieve] via [retriever.WithDSLInfo]. +const ( + DSLRiskType = "risk_type" + DSLSimilarityThreshold = "similarity_threshold" + DSLSubIndexFilter = "sub_index_filter" +) + +// FormatEmbeddingInput matches the historical indexing format so existing embeddings +// stay comparable if users skip reindex; new indexes use the same string shape. +func FormatEmbeddingInput(category, title, chunkText string) string { + return fmt.Sprintf("[风险类型：%s] [标题：%s]\n%s", category, title, chunkText) +} + +// FormatQueryEmbeddingText builds the string embedded at query time so it matches +// [FormatEmbeddingInput] for the same risk category (title left empty for queries). +func FormatQueryEmbeddingText(riskType, query string) string { + q := strings.TrimSpace(query) + rt := strings.TrimSpace(riskType) + if rt != "" { + return FormatEmbeddingInput(rt, "", q) + } + return q +} + +// MetaLookupString returns metadata string value or "" if absent. +func MetaLookupString(md map[string]any, key string) string { + if md == nil { + return "" + } + v, ok := md[key] + if !ok || v == nil { + return "" + } + switch t := v.(type) { + case string: + return t + default: + return strings.TrimSpace(fmt.Sprint(t)) + } +} + +// MetaStringOK returns trimmed non-empty string and true if present and non-empty. +func MetaStringOK(md map[string]any, key string) (string, bool) { + s := strings.TrimSpace(MetaLookupString(md, key)) + if s == "" { + return "", false + } + return s, true +} + +// RequireMetaString requires a non-empty string metadata field. +func RequireMetaString(md map[string]any, key string) (string, error) { + s, ok := MetaStringOK(md, key) + if !ok { + return "", fmt.Errorf("missing or empty metadata %q", key) + } + return s, nil +} + +// RequireMetaInt requires an integer metadata field. +func RequireMetaInt(md map[string]any, key string) (int, error) { + if md == nil { + return 0, fmt.Errorf("missing metadata key %q", key) + } + v, ok := md[key] + if !ok { + return 0, fmt.Errorf("missing metadata key %q", key) + } + switch t := v.(type) { + case int: + return t, nil + case int32: + return int(t), nil + case int64: + return int(t), nil + case float64: + return int(t), nil + default: + return 0, fmt.Errorf("metadata %q: unsupported type %T", key, v) + } +} + +// DSLNumeric coerces DSL map values (e.g. from JSON) to float64. +func DSLNumeric(v any) (float64, bool) { + switch t := v.(type) { + case float64: + return t, true + case float32: + return float64(t), true + case int: + return float64(t), true + case int64: + return float64(t), true + case uint32: + return float64(t), true + case uint64: + return float64(t), true + default: + return 0, false + } +} + +// MetaFloat64OK reads a float metadata value. +func MetaFloat64OK(md map[string]any, key string) (float64, bool) { + if md == nil { + return 0, false + } + v, ok := md[key] + if !ok { + return 0, false + } + return DSLNumeric(v) +} diff --git a/internal/knowledge/eino_meta_test.go b/internal/knowledge/eino_meta_test.go new file mode 100644 index 00000000..ba3f60da --- /dev/null +++ b/internal/knowledge/eino_meta_test.go @@ -0,0 +1,14 @@ +package knowledge + +import "testing" + +func TestFormatQueryEmbeddingText_AlignsWithIndexPrefix(t *testing.T) { + q := FormatQueryEmbeddingText("XSS", "payload") + want := FormatEmbeddingInput("XSS", "", "payload") + if q != want { + t.Fatalf("query embed text mismatch:\n got: %q\nwant: %q", q, want) + } + if FormatQueryEmbeddingText("", "hello") != "hello" { + t.Fatalf("expected bare query without risk type") + } +} diff --git a/internal/knowledge/eino_pipeline_retriever.go b/internal/knowledge/eino_pipeline_retriever.go new file mode 100644 index 00000000..487b439b --- /dev/null +++ b/internal/knowledge/eino_pipeline_retriever.go @@ -0,0 +1,96 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// knowledgePipelineRetriever: MultiQuery → vector candidates → rerank → post-process. +type knowledgePipelineRetriever struct { + inner retriever.Retriever + base *Retriever +} + +func newKnowledgePipelineRetriever(inner retriever.Retriever, base *Retriever) *knowledgePipelineRetriever { + if inner == nil || base == nil { + return nil + } + return &knowledgePipelineRetriever{inner: inner, base: base} +} + +func (p *knowledgePipelineRetriever) GetType() string { + return "KnowledgeRAGPipeline" +} + +func (p *knowledgePipelineRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (out []*schema.Document, err error) { + if p == nil || p.inner == nil || p.base == nil { + return nil, fmt.Errorf("knowledge pipeline retriever: nil") + } + q := strings.TrimSpace(query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + + ro := retriever.GetCommonOptions(nil, opts...) + finalTopK := p.base.config.TopK + if finalTopK <= 0 { + finalTopK = 5 + } + if ro.TopK != nil && *ro.TopK > 0 { + finalTopK = *ro.TopK + } + + ctx = callbacks.EnsureRunInfo(ctx, p.GetType(), components.ComponentOfRetriever) + ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{Query: q, TopK: finalTopK, Extra: ro.DSLInfo}) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: out}) + }() + + out, err = p.inner.Retrieve(ctx, q, opts...) + if err != nil { + return nil, err + } + if len(out) == 0 { + return out, nil + } + + if rr := p.base.documentReranker(); rr != nil && len(out) > 1 { + reranked, rerr := rr.Rerank(ctx, q, out) + if rerr != nil { + if p.base.logger != nil { + p.base.logger.Warn("知识检索重排失败，已使用融合序", zap.Error(rerr)) + } + } else if len(reranked) > 0 { + out = reranked + } + } + + tokenModel := "" + if p.base.embedder != nil { + tokenModel = p.base.embedder.EmbeddingModelName() + } + var postPO *config.PostRetrieveConfig + if p.base.config != nil { + postPO = &p.base.config.PostRetrieve + } + out, err = ApplyPostRetrieve(out, postPO, tokenModel, finalTopK) + if err != nil { + return nil, err + } + return out, nil +} + +var _ retriever.Retriever = (*knowledgePipelineRetriever)(nil) diff --git a/internal/knowledge/eino_retrieve_chain.go b/internal/knowledge/eino_retrieve_chain.go new file mode 100644 index 00000000..81fa4159 --- /dev/null +++ b/internal/knowledge/eino_retrieve_chain.go @@ -0,0 +1,24 @@ +package knowledge + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// BuildKnowledgeRetrieveChain 编译「查询字符串 → 文档列表」的 Eino Chain（MultiQuery → 向量 → 重排 → 后处理）。 +func BuildKnowledgeRetrieveChain(ctx context.Context, r *Retriever) (compose.Runnable[string, []*schema.Document], error) { + if r == nil { + return nil, fmt.Errorf("retriever is nil") + } + ch := compose.NewChain[string, []*schema.Document]() + ch.AppendRetriever(r.AsEinoRetriever()) + return ch.Compile(ctx) +} + +// CompileRetrieveChain 等价于 [BuildKnowledgeRetrieveChain](ctx, r)。 +func (r *Retriever) CompileRetrieveChain(ctx context.Context) (compose.Runnable[string, []*schema.Document], error) { + return BuildKnowledgeRetrieveChain(ctx, r) +} diff --git a/internal/knowledge/eino_retrieve_chain_test.go b/internal/knowledge/eino_retrieve_chain_test.go new file mode 100644 index 00000000..c74a6900 --- /dev/null +++ b/internal/knowledge/eino_retrieve_chain_test.go @@ -0,0 +1,23 @@ +package knowledge + +import ( + "context" + "testing" + + "go.uber.org/zap" +) + +func TestBuildKnowledgeRetrieveChain_Compile(t *testing.T) { + r := NewRetriever(nil, nil, &RetrievalConfig{TopK: 3, SimilarityThreshold: 0.5}, zap.NewNop()) + _, err := BuildKnowledgeRetrieveChain(context.Background(), r) + if err != nil { + t.Fatal(err) + } +} + +func TestBuildKnowledgeRetrieveChain_NilRetriever(t *testing.T) { + _, err := BuildKnowledgeRetrieveChain(context.Background(), nil) + if err == nil { + t.Fatal("expected error for nil retriever") + } +} diff --git a/internal/knowledge/eino_retriever_adapter.go b/internal/knowledge/eino_retriever_adapter.go new file mode 100644 index 00000000..712b4734 --- /dev/null +++ b/internal/knowledge/eino_retriever_adapter.go @@ -0,0 +1,173 @@ +package knowledge + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// VectorEinoRetriever implements [retriever.Retriever] on top of SQLite-stored embeddings + cosine similarity. +// It returns prefetch-sized vector candidates only; rerank and post-process run in [knowledgePipelineRetriever]. +type VectorEinoRetriever struct { + inner *Retriever +} + +// NewVectorEinoRetriever wraps r for Eino compose / tooling. +func NewVectorEinoRetriever(r *Retriever) *VectorEinoRetriever { + if r == nil { + return nil + } + return &VectorEinoRetriever{inner: r} +} + +// GetType identifies this retriever for Eino callbacks. +func (h *VectorEinoRetriever) GetType() string { + return "SQLiteVectorKnowledgeRetriever" +} + +// Retrieve runs vector search and returns [schema.Document] rows. +func (h *VectorEinoRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) (out []*schema.Document, err error) { + if h == nil || h.inner == nil { + return nil, fmt.Errorf("VectorEinoRetriever: nil retriever") + } + q := strings.TrimSpace(query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + + ro := retriever.GetCommonOptions(nil, opts...) + cfg := h.inner.config + + req := &SearchRequest{Query: q} + + if ro.TopK != nil && *ro.TopK > 0 { + req.TopK = *ro.TopK + } else if cfg != nil && cfg.TopK > 0 { + req.TopK = cfg.TopK + } else { + req.TopK = 5 + } + + req.Threshold = 0 + if ro.DSLInfo != nil { + if rt, ok := ro.DSLInfo[DSLRiskType].(string); ok { + req.RiskType = strings.TrimSpace(rt) + } + if v, ok := ro.DSLInfo[DSLSimilarityThreshold]; ok { + if f, ok2 := DSLNumeric(v); ok2 && f > 0 { + req.Threshold = f + } + } + if sf, ok := ro.DSLInfo[DSLSubIndexFilter].(string); ok { + req.SubIndexFilter = strings.TrimSpace(sf) + } + } + if req.SubIndexFilter == "" && cfg != nil && strings.TrimSpace(cfg.SubIndexFilter) != "" { + req.SubIndexFilter = strings.TrimSpace(cfg.SubIndexFilter) + } + if req.Threshold <= 0 && cfg != nil && cfg.SimilarityThreshold > 0 { + req.Threshold = cfg.SimilarityThreshold + } + if req.Threshold <= 0 { + req.Threshold = 0.7 + } + + finalTopK := req.TopK + var postPO *config.PostRetrieveConfig + if cfg != nil { + postPO = &cfg.PostRetrieve + } + fetchK := EffectivePrefetchTopK(finalTopK, postPO) + searchReq := *req + searchReq.TopK = fetchK + + ctx = callbacks.EnsureRunInfo(ctx, h.GetType(), components.ComponentOfRetriever) + th := req.Threshold + st := &th + ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{ + Query: q, + TopK: finalTopK, + ScoreThreshold: st, + Extra: ro.DSLInfo, + }) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: out}) + }() + + results, err := h.inner.vectorSearch(ctx, &searchReq) + if err != nil { + return nil, err + } + out = retrievalResultsToDocuments(results) + return out, nil +} + +func retrievalResultsToDocuments(results []*RetrievalResult) []*schema.Document { + out := make([]*schema.Document, 0, len(results)) + for _, res := range results { + if res == nil || res.Chunk == nil || res.Item == nil { + continue + } + d := &schema.Document{ + ID: res.Chunk.ID, + Content: res.Chunk.ChunkText, + MetaData: map[string]any{ + metaKBItemID: res.Item.ID, + metaKBCategory: res.Item.Category, + metaKBTitle: res.Item.Title, + metaKBChunkIndex: res.Chunk.ChunkIndex, + metaSimilarity: res.Similarity, + }, + } + d.WithScore(res.Score) + out = append(out, d) + } + return out +} + +func documentsToRetrievalResults(docs []*schema.Document) ([]*RetrievalResult, error) { + out := make([]*RetrievalResult, 0, len(docs)) + for i, d := range docs { + if d == nil { + continue + } + itemID, err := RequireMetaString(d.MetaData, metaKBItemID) + if err != nil { + return nil, fmt.Errorf("document %d: %w", i, err) + } + cat := MetaLookupString(d.MetaData, metaKBCategory) + title := MetaLookupString(d.MetaData, metaKBTitle) + chunkIdx, err := RequireMetaInt(d.MetaData, metaKBChunkIndex) + if err != nil { + return nil, fmt.Errorf("document %d: %w", i, err) + } + sim, _ := MetaFloat64OK(d.MetaData, metaSimilarity) + item := &KnowledgeItem{ID: itemID, Category: cat, Title: title} + chunk := &KnowledgeChunk{ + ID: d.ID, + ItemID: itemID, + ChunkIndex: chunkIdx, + ChunkText: d.Content, + } + out = append(out, &RetrievalResult{ + Chunk: chunk, + Item: item, + Similarity: sim, + Score: d.Score(), + }) + } + return out, nil +} + +var _ retriever.Retriever = (*VectorEinoRetriever)(nil) diff --git a/internal/knowledge/eino_sqlite_indexer.go b/internal/knowledge/eino_sqlite_indexer.go new file mode 100644 index 00000000..a0bbdcdc --- /dev/null +++ b/internal/knowledge/eino_sqlite_indexer.go @@ -0,0 +1,142 @@ +package knowledge + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" + "github.com/google/uuid" +) + +// SQLiteIndexer implements [indexer.Indexer] against knowledge_embeddings + existing schema. +type SQLiteIndexer struct { + db *sql.DB + batchSize int + embeddingModel string +} + +// NewSQLiteIndexer returns an indexer that writes chunk rows for one knowledge item per Store call. +// batchSize is the embedding batch size; if <= 0, default 64 is used. +// embeddingModel is persisted per row for retrieval-time consistency checks (may be empty). +func NewSQLiteIndexer(db *sql.DB, batchSize int, embeddingModel string) *SQLiteIndexer { + return &SQLiteIndexer{db: db, batchSize: batchSize, embeddingModel: strings.TrimSpace(embeddingModel)} +} + +// GetType implements eino callback run info. +func (s *SQLiteIndexer) GetType() string { + return "SQLiteKnowledgeIndexer" +} + +// Store embeds documents and inserts rows. Each doc must carry MetaData: +// kb_item_id, kb_category, kb_title, kb_chunk_index (int). Content is chunk text only. +func (s *SQLiteIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) (ids []string, err error) { + options := indexer.GetCommonOptions(nil, opts...) + if options.Embedding == nil { + return nil, fmt.Errorf("sqlite indexer: embedding is required") + } + if len(docs) == 0 { + return nil, nil + } + + ctx = callbacks.EnsureRunInfo(ctx, s.GetType(), components.ComponentOfIndexer) + ctx = callbacks.OnStart(ctx, &indexer.CallbackInput{Docs: docs}) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + return + } + _ = callbacks.OnEnd(ctx, &indexer.CallbackOutput{IDs: ids}) + }() + + subIdxStr := strings.Join(options.SubIndexes, ",") + + texts := make([]string, len(docs)) + for i, d := range docs { + if d == nil { + return nil, fmt.Errorf("sqlite indexer: nil document at %d", i) + } + cat := MetaLookupString(d.MetaData, metaKBCategory) + title := MetaLookupString(d.MetaData, metaKBTitle) + texts[i] = FormatEmbeddingInput(cat, title, d.Content) + } + + bs := s.batchSize + if bs <= 0 { + bs = 64 + } + + var allVecs [][]float64 + for start := 0; start < len(texts); start += bs { + end := start + bs + if end > len(texts) { + end = len(texts) + } + batch := texts[start:end] + vecs, embedErr := options.Embedding.EmbedStrings(ctx, batch) + if embedErr != nil { + return nil, fmt.Errorf("sqlite indexer: embed batch %d-%d: %w", start, end, embedErr) + } + if len(vecs) != len(batch) { + return nil, fmt.Errorf("sqlite indexer: embed count mismatch: got %d want %d", len(vecs), len(batch)) + } + allVecs = append(allVecs, vecs...) + } + + embedDim := 0 + if len(allVecs) > 0 { + embedDim = len(allVecs[0]) + } + + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("sqlite indexer: begin tx: %w", err) + } + defer tx.Rollback() + + ids = make([]string, 0, len(docs)) + for i, d := range docs { + chunkID := uuid.New().String() + itemID, metaErr := RequireMetaString(d.MetaData, metaKBItemID) + if metaErr != nil { + return nil, fmt.Errorf("sqlite indexer: doc %d: %w", i, metaErr) + } + chunkIdx, metaErr := RequireMetaInt(d.MetaData, metaKBChunkIndex) + if metaErr != nil { + return nil, fmt.Errorf("sqlite indexer: doc %d: %w", i, metaErr) + } + vec := allVecs[i] + if embedDim > 0 && len(vec) != embedDim { + return nil, fmt.Errorf("sqlite indexer: inconsistent embedding dim at doc %d: got %d want %d", i, len(vec), embedDim) + } + vec32 := make([]float32, len(vec)) + for j, v := range vec { + vec32[j] = float32(v) + } + embeddingJSON, jsonErr := json.Marshal(vec32) + if jsonErr != nil { + return nil, fmt.Errorf("sqlite indexer: marshal embedding: %w", jsonErr) + } + _, err = tx.ExecContext(ctx, + `INSERT INTO knowledge_embeddings (id, item_id, chunk_index, chunk_text, embedding, sub_indexes, embedding_model, embedding_dim, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))`, + chunkID, itemID, chunkIdx, d.Content, string(embeddingJSON), subIdxStr, s.embeddingModel, embedDim, + ) + if err != nil { + return nil, fmt.Errorf("sqlite indexer: insert chunk %d: %w", i, err) + } + ids = append(ids, chunkID) + } + + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("sqlite indexer: commit: %w", err) + } + return ids, nil +} + +var _ indexer.Indexer = (*SQLiteIndexer)(nil) diff --git a/internal/knowledge/embedder.go b/internal/knowledge/embedder.go new file mode 100644 index 00000000..d9ce8afa --- /dev/null +++ b/internal/knowledge/embedder.go @@ -0,0 +1,251 @@ +package knowledge + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + einoembedopenai "github.com/cloudwego/eino-ext/components/embedding/openai" + "github.com/cloudwego/eino/components/embedding" + "go.uber.org/zap" + "golang.org/x/time/rate" +) + +// Embedder 使用 CloudWeGo Eino 的 OpenAI Embedding 组件，并保留速率限制与重试。 +type Embedder struct { + eino embedding.Embedder + config *config.KnowledgeConfig + logger *zap.Logger + + rateLimiter *rate.Limiter + rateLimitDelay time.Duration + maxRetries int + retryDelay time.Duration + mu sync.Mutex +} + +// NewEmbedder 基于 Eino eino-ext OpenAI Embedder；openAIConfig 用于在知识库未单独配置 key 时回退 API Key。 +func NewEmbedder(ctx context.Context, cfg *config.KnowledgeConfig, openAIConfig *config.OpenAIConfig, logger *zap.Logger) (*Embedder, error) { + if cfg == nil { + return nil, fmt.Errorf("knowledge config is nil") + } + + var rateLimiter *rate.Limiter + var rateLimitDelay time.Duration + if cfg.Indexing.MaxRPM > 0 { + rpm := cfg.Indexing.MaxRPM + rateLimiter = rate.NewLimiter(rate.Every(time.Minute/time.Duration(rpm)), rpm) + if logger != nil { + logger.Info("知识库索引速率限制已启用", zap.Int("maxRPM", rpm)) + } + } else if cfg.Indexing.RateLimitDelayMs > 0 { + rateLimitDelay = time.Duration(cfg.Indexing.RateLimitDelayMs) * time.Millisecond + if logger != nil { + logger.Info("知识库索引固定延迟已启用", zap.Duration("delay", rateLimitDelay)) + } + } + + maxRetries := 3 + retryDelay := 1000 * time.Millisecond + if cfg.Indexing.MaxRetries > 0 { + maxRetries = cfg.Indexing.MaxRetries + } + if cfg.Indexing.RetryDelayMs > 0 { + retryDelay = time.Duration(cfg.Indexing.RetryDelayMs) * time.Millisecond + } + + model := strings.TrimSpace(cfg.Embedding.Model) + if model == "" { + model = "text-embedding-3-small" + } + + baseURL := strings.TrimSpace(cfg.Embedding.BaseURL) + baseURL = strings.TrimSuffix(baseURL, "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + apiKey := strings.TrimSpace(cfg.Embedding.APIKey) + if apiKey == "" && openAIConfig != nil { + apiKey = strings.TrimSpace(openAIConfig.APIKey) + } + if apiKey == "" { + return nil, fmt.Errorf("embedding API key 未配置") + } + + timeout := 120 * time.Second + if cfg.Indexing.RequestTimeoutSeconds > 0 { + timeout = time.Duration(cfg.Indexing.RequestTimeoutSeconds) * time.Second + } + httpClient := &http.Client{Timeout: timeout} + + inner, err := einoembedopenai.NewEmbedder(ctx, &einoembedopenai.EmbeddingConfig{ + APIKey: apiKey, + BaseURL: baseURL, + ByAzure: false, + Model: model, + HTTPClient: httpClient, + }) + if err != nil { + return nil, fmt.Errorf("eino OpenAI embedder: %w", err) + } + + return &Embedder{ + eino: inner, + config: cfg, + logger: logger, + rateLimiter: rateLimiter, + rateLimitDelay: rateLimitDelay, + maxRetries: maxRetries, + retryDelay: retryDelay, + }, nil +} + +// EmbeddingModelName 返回配置的嵌入模型名（用于 tiktoken 分块与向量行元数据）。 +func (e *Embedder) EmbeddingModelName() string { + if e == nil || e.config == nil { + return "" + } + s := strings.TrimSpace(e.config.Embedding.Model) + if s != "" { + return s + } + return "text-embedding-3-small" +} + +func (e *Embedder) waitRateLimiter() { + e.mu.Lock() + defer e.mu.Unlock() + + if e.rateLimiter != nil { + ctx := context.Background() + if err := e.rateLimiter.Wait(ctx); err != nil && e.logger != nil { + e.logger.Warn("速率限制器等待失败", zap.Error(err)) + } + } + if e.rateLimitDelay > 0 { + time.Sleep(e.rateLimitDelay) + } +} + +// EmbedText 单条嵌入（float32，与历史存储格式一致）。 +func (e *Embedder) EmbedText(ctx context.Context, text string) ([]float32, error) { + vecs, err := e.EmbedStrings(ctx, []string{text}) + if err != nil { + return nil, err + } + if len(vecs) != 1 { + return nil, fmt.Errorf("unexpected embedding count: %d", len(vecs)) + } + return vecs[0], nil +} + +// EmbedStrings 批量嵌入，带重试；实现 [embedding.Embedder]，可供 Eino Indexer 使用。 +func (e *Embedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float32, error) { + if e == nil || e.eino == nil { + return nil, fmt.Errorf("embedder not initialized") + } + if len(texts) == 0 { + return nil, nil + } + + var lastErr error + for attempt := 0; attempt < e.maxRetries; attempt++ { + if attempt > 0 { + wait := e.retryDelay * time.Duration(attempt) + if e.logger != nil { + e.logger.Debug("嵌入重试前等待", zap.Int("attempt", attempt+1), zap.Duration("wait", wait)) + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(wait): + } + } else { + e.waitRateLimiter() + } + + raw, err := e.eino.EmbedStrings(ctx, texts, opts...) + if err == nil { + out := make([][]float32, len(raw)) + for i, row := range raw { + out[i] = make([]float32, len(row)) + for j, v := range row { + out[i][j] = float32(v) + } + } + return out, nil + } + lastErr = err + if !e.isRetryableError(err) { + return nil, err + } + if e.logger != nil { + e.logger.Debug("嵌入失败，将重试", zap.Int("attempt", attempt+1), zap.Error(err)) + } + } + return nil, fmt.Errorf("达到最大重试次数 (%d): %v", e.maxRetries, lastErr) +} + +// EmbedTexts 批量 float32 嵌入（兼容旧调用；单次请求批量以减小延迟）。 +func (e *Embedder) EmbedTexts(ctx context.Context, texts []string) ([][]float32, error) { + return e.EmbedStrings(ctx, texts) +} + +func (e *Embedder) isRetryableError(err error) bool { + if err == nil { + return false + } + errStr := err.Error() + if strings.Contains(errStr, "429") || strings.Contains(errStr, "rate limit") { + return true + } + if strings.Contains(errStr, "500") || strings.Contains(errStr, "502") || + strings.Contains(errStr, "503") || strings.Contains(errStr, "504") { + return true + } + if strings.Contains(errStr, "timeout") || strings.Contains(errStr, "connection") || + strings.Contains(errStr, "network") || strings.Contains(errStr, "EOF") { + return true + } + return false +} + +// einoFloatEmbedder adapts [][]float32 embedder to Eino's [][]float64 [embedding.Embedder] for Indexer.Store. +type einoFloatEmbedder struct { + inner *Embedder +} + +func (w *einoFloatEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) { + vec32, err := w.inner.EmbedStrings(ctx, texts, opts...) + if err != nil { + return nil, err + } + out := make([][]float64, len(vec32)) + for i, row := range vec32 { + out[i] = make([]float64, len(row)) + for j, v := range row { + out[i][j] = float64(v) + } + } + return out, nil +} + +func (w *einoFloatEmbedder) GetType() string { + return "CyberStrikeKnowledgeEmbedder" +} + +func (w *einoFloatEmbedder) IsCallbacksEnabled() bool { + return false +} + +// EinoEmbeddingComponent returns an [embedding.Embedder] that uses the same retry/rate-limit path +// and produces float64 vectors expected by generic Eino indexer helpers. +func (e *Embedder) EinoEmbeddingComponent() embedding.Embedder { + return &einoFloatEmbedder{inner: e} +} diff --git a/internal/knowledge/index_pipeline.go b/internal/knowledge/index_pipeline.go new file mode 100644 index 00000000..a9b9a4c4 --- /dev/null +++ b/internal/knowledge/index_pipeline.go @@ -0,0 +1,91 @@ +package knowledge + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// normalizeChunkStrategy returns "recursive" or "markdown_then_recursive". +func normalizeChunkStrategy(s string) string { + v := strings.TrimSpace(strings.ToLower(s)) + switch v { + case "recursive": + return "recursive" + case "markdown_then_recursive", "markdown_recursive", "markdown": + return "markdown_then_recursive" + case "": + return "markdown_then_recursive" + default: + return "markdown_then_recursive" + } +} + +func buildKnowledgeIndexChain( + ctx context.Context, + indexingCfg *config.IndexingConfig, + db *sql.DB, + recursive document.Transformer, + embeddingModel string, +) (compose.Runnable[[]*schema.Document, []string], error) { + if recursive == nil { + return nil, fmt.Errorf("recursive transformer is nil") + } + if db == nil { + return nil, fmt.Errorf("db is nil") + } + strategy := normalizeChunkStrategy("markdown_then_recursive") + batch := 64 + maxChunks := 0 + if indexingCfg != nil { + strategy = normalizeChunkStrategy(indexingCfg.ChunkStrategy) + if indexingCfg.BatchSize > 0 { + batch = indexingCfg.BatchSize + } + maxChunks = indexingCfg.MaxChunksPerItem + } + + si := NewSQLiteIndexer(db, batch, embeddingModel) + ch := compose.NewChain[[]*schema.Document, []string]() + if strategy != "recursive" { + md, err := newMarkdownHeaderSplitter(ctx) + if err != nil { + return nil, fmt.Errorf("markdown splitter: %w", err) + } + ch.AppendDocumentTransformer(md) + } + ch.AppendDocumentTransformer(recursive) + ch.AppendLambda(newChunkEnrichLambda(maxChunks)) + ch.AppendIndexer(si) + return ch.Compile(ctx) +} + +func newChunkEnrichLambda(maxChunks int) *compose.Lambda { + return compose.InvokableLambda(func(ctx context.Context, docs []*schema.Document) ([]*schema.Document, error) { + _ = ctx + out := make([]*schema.Document, 0, len(docs)) + for _, d := range docs { + if d == nil || strings.TrimSpace(d.Content) == "" { + continue + } + out = append(out, d) + } + if maxChunks > 0 && len(out) > maxChunks { + out = out[:maxChunks] + } + for i, d := range out { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + d.MetaData[metaKBChunkIndex] = i + } + return out, nil + }) +} diff --git a/internal/knowledge/index_pipeline_test.go b/internal/knowledge/index_pipeline_test.go new file mode 100644 index 00000000..9e4b03fa --- /dev/null +++ b/internal/knowledge/index_pipeline_test.go @@ -0,0 +1,21 @@ +package knowledge + +import "testing" + +func TestNormalizeChunkStrategy(t *testing.T) { + cases := []struct { + in, want string + }{ + {"", "markdown_then_recursive"}, + {"recursive", "recursive"}, + {"RECURSIVE", "recursive"}, + {"markdown_then_recursive", "markdown_then_recursive"}, + {"markdown", "markdown_then_recursive"}, + {"unknown", "markdown_then_recursive"}, + } + for _, tc := range cases { + if got := normalizeChunkStrategy(tc.in); got != tc.want { + t.Errorf("normalizeChunkStrategy(%q) = %q, want %q", tc.in, got, tc.want) + } + } +} diff --git a/internal/knowledge/indexer.go b/internal/knowledge/indexer.go new file mode 100644 index 00000000..1c1e7cbe --- /dev/null +++ b/internal/knowledge/indexer.go @@ -0,0 +1,435 @@ +package knowledge + +import ( + "context" + "database/sql" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + fileloader "github.com/cloudwego/eino-ext/components/document/loader/file" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Indexer 使用 Eino Compose 索引链（Markdown/递归分块、Lambda enrich、SQLite 索引）与嵌入写入。 +type Indexer struct { + db *sql.DB + embedder *Embedder + logger *zap.Logger + chunkSize int + overlap int + indexingCfg *config.IndexingConfig + + indexChain compose.Runnable[[]*schema.Document, []string] + fileLoader *fileloader.FileLoader + + mu sync.RWMutex + lastError string + lastErrorTime time.Time + errorCount int + + rebuildMu sync.RWMutex + isRebuilding bool + rebuildTotalItems int + rebuildCurrent int + rebuildFailed int + rebuildStartTime time.Time + rebuildLastItemID string + rebuildLastChunks int +} + +// NewIndexer 创建索引器并编译 Eino 索引链；kcfg 为完整知识库配置（含 indexing 与路径相关行为）。 +func NewIndexer(ctx context.Context, db *sql.DB, embedder *Embedder, logger *zap.Logger, kcfg *config.KnowledgeConfig) (*Indexer, error) { + if db == nil { + return nil, fmt.Errorf("db is nil") + } + if embedder == nil { + return nil, fmt.Errorf("embedder is nil") + } + if err := EnsureKnowledgeEmbeddingsSchema(db); err != nil { + return nil, fmt.Errorf("knowledge_embeddings 结构迁移: %w", err) + } + if kcfg == nil { + kcfg = &config.KnowledgeConfig{} + } + indexingCfg := &kcfg.Indexing + + chunkSize := 512 + overlap := 50 + if indexingCfg.ChunkSize > 0 { + chunkSize = indexingCfg.ChunkSize + } + if indexingCfg.ChunkOverlap >= 0 { + overlap = indexingCfg.ChunkOverlap + } + + embedModel := embedder.EmbeddingModelName() + splitter, err := newKnowledgeSplitter(chunkSize, overlap, embedModel) + if err != nil { + return nil, fmt.Errorf("eino recursive splitter: %w", err) + } + + chain, err := buildKnowledgeIndexChain(ctx, indexingCfg, db, splitter, embedModel) + if err != nil { + return nil, fmt.Errorf("knowledge index chain: %w", err) + } + + var fl *fileloader.FileLoader + fl, err = fileloader.NewFileLoader(ctx, nil) + if err != nil { + if logger != nil { + logger.Warn("Eino FileLoader 初始化失败，prefer_source_file 将回退数据库正文", zap.Error(err)) + } + fl = nil + err = nil + } + + return &Indexer{ + db: db, + embedder: embedder, + logger: logger, + chunkSize: chunkSize, + overlap: overlap, + indexingCfg: indexingCfg, + indexChain: chain, + fileLoader: fl, + }, nil +} + +// RecompileIndexChain 在配置或嵌入模型变更后重建 Eino 索引链（无需重启进程）。 +func (idx *Indexer) RecompileIndexChain(ctx context.Context) error { + if idx == nil || idx.db == nil || idx.embedder == nil { + return fmt.Errorf("indexer 未初始化") + } + if err := EnsureKnowledgeEmbeddingsSchema(idx.db); err != nil { + return err + } + embedModel := idx.embedder.EmbeddingModelName() + splitter, err := newKnowledgeSplitter(idx.chunkSize, idx.overlap, embedModel) + if err != nil { + return fmt.Errorf("eino recursive splitter: %w", err) + } + chain, err := buildKnowledgeIndexChain(ctx, idx.indexingCfg, idx.db, splitter, embedModel) + if err != nil { + return fmt.Errorf("knowledge index chain: %w", err) + } + idx.indexChain = chain + return nil +} + +// IndexItem 索引单个知识项：先清空旧向量，再走 Compose 链（分块、嵌入、写入）。 +func (idx *Indexer) IndexItem(ctx context.Context, itemID string) error { + if idx.indexChain == nil { + return fmt.Errorf("索引链未初始化") + } + if idx.embedder == nil { + return fmt.Errorf("嵌入器未初始化") + } + + var content, category, title, filePath string + err := idx.db.QueryRow("SELECT content, category, title, file_path FROM knowledge_base_items WHERE id = ?", itemID).Scan(&content, &category, &title, &filePath) + if err != nil { + return fmt.Errorf("获取知识项失败：%w", err) + } + + if _, err := idx.db.Exec("DELETE FROM knowledge_embeddings WHERE item_id = ?", itemID); err != nil { + return fmt.Errorf("删除旧向量失败：%w", err) + } + + body := strings.TrimSpace(content) + if idx.indexingCfg != nil && idx.indexingCfg.PreferSourceFile && strings.TrimSpace(filePath) != "" && idx.fileLoader != nil { + docs, lerr := idx.fileLoader.Load(ctx, document.Source{URI: strings.TrimSpace(filePath)}) + if lerr == nil && len(docs) > 0 { + var b strings.Builder + for i, d := range docs { + if d == nil { + continue + } + if i > 0 { + b.WriteString("\n\n") + } + b.WriteString(d.Content) + } + if s := strings.TrimSpace(b.String()); s != "" { + body = s + } + } else if idx.logger != nil { + idx.logger.Warn("优先源文件读取失败，使用数据库正文", + zap.String("itemId", itemID), + zap.String("path", filePath), + zap.Error(lerr)) + } + } + + root := &schema.Document{ + ID: itemID, + Content: body, + MetaData: map[string]any{ + metaKBCategory: category, + metaKBTitle: title, + metaKBItemID: itemID, + }, + } + + idxOpts := []indexer.Option{indexer.WithEmbedding(idx.embedder.EinoEmbeddingComponent())} + if idx.indexingCfg != nil && len(idx.indexingCfg.SubIndexes) > 0 { + idxOpts = append(idxOpts, indexer.WithSubIndexes(idx.indexingCfg.SubIndexes)) + } + + ids, err := idx.indexChain.Invoke(ctx, []*schema.Document{root}, compose.WithIndexerOption(idxOpts...)) + if err != nil { + msg := fmt.Sprintf("索引写入失败 (知识项：%s): %v", itemID, err) + idx.mu.Lock() + idx.lastError = msg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + return err + } + + if idx.logger != nil { + idx.logger.Info("知识项索引完成", zap.String("itemId", itemID), zap.Int("chunks", len(ids))) + } + idx.rebuildMu.Lock() + idx.rebuildLastItemID = itemID + idx.rebuildLastChunks = len(ids) + idx.rebuildMu.Unlock() + return nil +} + +// HasIndex 检查是否存在索引 +func (idx *Indexer) HasIndex() (bool, error) { + var count int + err := idx.db.QueryRow("SELECT COUNT(*) FROM knowledge_embeddings").Scan(&count) + if err != nil { + return false, fmt.Errorf("检查索引失败：%w", err) + } + return count > 0, nil +} + +func (idx *Indexer) beginIndexRun() error { + idx.rebuildMu.Lock() + defer idx.rebuildMu.Unlock() + + if idx.isRebuilding { + return fmt.Errorf("索引任务已在进行中") + } + idx.isRebuilding = true + idx.rebuildTotalItems = 0 + idx.rebuildCurrent = 0 + idx.rebuildFailed = 0 + idx.rebuildStartTime = time.Now() + idx.rebuildLastItemID = "" + idx.rebuildLastChunks = 0 + return nil +} + +// TryBeginIndexRun 同步占用索引任务槽位；调用方必须在后台任务结束时调用 FinishIndexRun。 +func (idx *Indexer) TryBeginIndexRun() error { + return idx.beginIndexRun() +} + +func (idx *Indexer) FinishIndexRun() { + idx.rebuildMu.Lock() + idx.isRebuilding = false + idx.rebuildMu.Unlock() +} + +func (idx *Indexer) resetLastError() { + idx.mu.Lock() + idx.lastError = "" + idx.lastErrorTime = time.Time{} + idx.errorCount = 0 + idx.mu.Unlock() +} + +func (idx *Indexer) setIndexRunTotal(total int) { + idx.rebuildMu.Lock() + idx.rebuildTotalItems = total + idx.rebuildMu.Unlock() +} + +// IndexMissing 为尚无向量的知识项构建索引（默认推荐路径，适合冷启动与中断续跑）。 +func (idx *Indexer) IndexMissing(ctx context.Context) error { + if err := idx.beginIndexRun(); err != nil { + return err + } + defer idx.FinishIndexRun() + return idx.runIndexMissing(ctx) +} + +// RebuildIndex 全量重建所有知识项索引（显式 opt-in，成本更高）。 +func (idx *Indexer) RebuildIndex(ctx context.Context) error { + if err := idx.beginIndexRun(); err != nil { + return err + } + defer idx.FinishIndexRun() + return idx.runRebuildIndex(ctx) +} + +// RunRebuildIndex 在已占用索引任务槽位后执行全量重建（供 HTTP handler 后台任务使用）。 +func (idx *Indexer) RunRebuildIndex(ctx context.Context) error { + return idx.runRebuildIndex(ctx) +} + +// RunIndexMissing 在已占用索引任务槽位后执行缺失索引补齐（供 HTTP handler 后台任务使用）。 +func (idx *Indexer) RunIndexMissing(ctx context.Context) error { + return idx.runIndexMissing(ctx) +} + +func (idx *Indexer) runRebuildIndex(ctx context.Context) error { + idx.resetLastError() + + rows, err := idx.db.QueryContext(ctx, "SELECT id FROM knowledge_base_items ORDER BY updated_at ASC, id ASC") + if err != nil { + return fmt.Errorf("查询知识项失败：%w", err) + } + defer rows.Close() + + itemIDs, err := scanKnowledgeItemIDs(rows) + if err != nil { + return err + } + + idx.setIndexRunTotal(len(itemIDs)) + idx.logger.Info("开始重建索引", zap.Int("totalItems", len(itemIDs))) + + return idx.indexItemIDs(ctx, itemIDs, "索引重建完成") +} + +func (idx *Indexer) runIndexMissing(ctx context.Context) error { + idx.resetLastError() + + rows, err := idx.db.QueryContext(ctx, ` + SELECT i.id + FROM knowledge_base_items i + LEFT JOIN knowledge_embeddings e ON e.item_id = i.id + WHERE e.item_id IS NULL + ORDER BY i.updated_at ASC, i.id ASC + `) + if err != nil { + return fmt.Errorf("查询未索引知识项失败：%w", err) + } + defer rows.Close() + + itemIDs, err := scanKnowledgeItemIDs(rows) + if err != nil { + return fmt.Errorf("扫描未索引知识项 ID 失败：%w", err) + } + + idx.setIndexRunTotal(len(itemIDs)) + idx.logger.Info("开始补齐缺失索引", zap.Int("totalItems", len(itemIDs))) + + return idx.indexItemIDs(ctx, itemIDs, "索引构建完成") +} + +func scanKnowledgeItemIDs(rows *sql.Rows) ([]string, error) { + var itemIDs []string + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err) + } + itemIDs = append(itemIDs, id) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("扫描知识项 ID 失败：%w", err) + } + return itemIDs, nil +} + +func (idx *Indexer) indexItemIDs(ctx context.Context, itemIDs []string, doneMessage string) error { + failedCount := 0 + consecutiveFailures := 0 + maxConsecutiveFailures := 5 + firstFailureItemID := "" + var firstFailureError error + + for i, itemID := range itemIDs { + if err := idx.IndexItem(ctx, itemID); err != nil { + failedCount++ + consecutiveFailures++ + + if consecutiveFailures == 1 { + firstFailureItemID = itemID + firstFailureError = err + idx.logger.Warn("索引知识项失败", + zap.String("itemId", itemID), + zap.Int("totalItems", len(itemIDs)), + zap.Error(err), + ) + } + + if consecutiveFailures >= maxConsecutiveFailures { + errorMsg := fmt.Sprintf("连续 %d 个知识项索引失败，可能存在配置问题（如嵌入模型配置错误、API 密钥无效、余额不足等）。第一个失败项：%s, 错误：%v", consecutiveFailures, firstFailureItemID, firstFailureError) + idx.mu.Lock() + idx.lastError = errorMsg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + + idx.logger.Error("连续索引失败次数过多，立即停止索引", + zap.Int("consecutiveFailures", consecutiveFailures), + zap.Int("totalItems", len(itemIDs)), + zap.Int("processedItems", i+1), + zap.String("firstFailureItemId", firstFailureItemID), + zap.Error(firstFailureError), + ) + return fmt.Errorf("连续索引失败次数过多：%v", firstFailureError) + } + + if failedCount > len(itemIDs)*3/10 && failedCount == len(itemIDs)*3/10+1 { + errorMsg := fmt.Sprintf("索引失败的知识项过多 (%d/%d)，可能存在配置问题。第一个失败项：%s, 错误：%v", failedCount, len(itemIDs), firstFailureItemID, firstFailureError) + idx.mu.Lock() + idx.lastError = errorMsg + idx.lastErrorTime = time.Now() + idx.mu.Unlock() + + idx.logger.Error("索引失败的知识项过多，可能存在配置问题", + zap.Int("failedCount", failedCount), + zap.Int("totalItems", len(itemIDs)), + zap.String("firstFailureItemId", firstFailureItemID), + zap.Error(firstFailureError), + ) + } + continue + } + + if consecutiveFailures > 0 { + consecutiveFailures = 0 + firstFailureItemID = "" + firstFailureError = nil + } + + idx.rebuildMu.Lock() + idx.rebuildCurrent = i + 1 + idx.rebuildFailed = failedCount + idx.rebuildMu.Unlock() + + if (i+1)%10 == 0 || (len(itemIDs) > 0 && (i+1)*100/len(itemIDs)%10 == 0 && (i+1)*100/len(itemIDs) > 0) { + idx.logger.Info("索引进度", zap.Int("current", i+1), zap.Int("total", len(itemIDs)), zap.Int("failed", failedCount)) + } + } + + idx.logger.Info(doneMessage, zap.Int("totalItems", len(itemIDs)), zap.Int("failedCount", failedCount)) + return nil +} + +// GetLastError 获取最近一次错误信息 +func (idx *Indexer) GetLastError() (string, time.Time) { + idx.mu.RLock() + defer idx.mu.RUnlock() + return idx.lastError, idx.lastErrorTime +} + +// GetRebuildStatus 获取重建索引状态 +func (idx *Indexer) GetRebuildStatus() (isRebuilding bool, totalItems int, current int, failed int, lastItemID string, lastChunks int, startTime time.Time) { + idx.rebuildMu.RLock() + defer idx.rebuildMu.RUnlock() + return idx.isRebuilding, idx.rebuildTotalItems, idx.rebuildCurrent, idx.rebuildFailed, idx.rebuildLastItemID, idx.rebuildLastChunks, idx.rebuildStartTime +} diff --git a/internal/knowledge/indexer_rebuild_state_test.go b/internal/knowledge/indexer_rebuild_state_test.go new file mode 100644 index 00000000..1e04480e --- /dev/null +++ b/internal/knowledge/indexer_rebuild_state_test.go @@ -0,0 +1,20 @@ +package knowledge + +import "testing" + +func TestIndexerRejectsConcurrentIndexRuns(t *testing.T) { + idx := &Indexer{} + + if err := idx.beginIndexRun(); err != nil { + t.Fatalf("first index run should start: %v", err) + } + if err := idx.beginIndexRun(); err == nil { + t.Fatal("second index run should be rejected while one is active") + } + + idx.FinishIndexRun() + if err := idx.beginIndexRun(); err != nil { + t.Fatalf("index run should start again after finish: %v", err) + } + idx.FinishIndexRun() +} diff --git a/internal/knowledge/manager.go b/internal/knowledge/manager.go new file mode 100644 index 00000000..7309cc2a --- /dev/null +++ b/internal/knowledge/manager.go @@ -0,0 +1,885 @@ +package knowledge + +import ( + "database/sql" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "strings" + "time" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// Manager 知识库管理器 +type Manager struct { + db *sql.DB + basePath string + logger *zap.Logger +} + +// NewManager 创建新的知识库管理器 +func NewManager(db *sql.DB, basePath string, logger *zap.Logger) *Manager { + return &Manager{ + db: db, + basePath: basePath, + logger: logger, + } +} + +// ScanKnowledgeBase 扫描知识库目录，更新数据库 +// 返回需要索引的知识项ID列表（新添加的或更新的） +func (m *Manager) ScanKnowledgeBase() ([]string, error) { + if m.basePath == "" { + return nil, fmt.Errorf("知识库路径未配置") + } + + // 确保目录存在 + if err := os.MkdirAll(m.basePath, 0755); err != nil { + return nil, fmt.Errorf("创建知识库目录失败: %w", err) + } + + var itemsToIndex []string + + // 遍历知识库目录 + err := filepath.WalkDir(m.basePath, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + + // 跳过目录和非markdown文件 + if d.IsDir() || !strings.HasSuffix(strings.ToLower(path), ".md") { + return nil + } + + // 计算相对路径和分类 + relPath, err := filepath.Rel(m.basePath, path) + if err != nil { + return err + } + + // 第一个目录名作为分类（风险类型） + parts := strings.Split(relPath, string(filepath.Separator)) + category := "未分类" + if len(parts) > 1 { + category = parts[0] + } + + // 文件名为标题 + title := strings.TrimSuffix(filepath.Base(path), ".md") + + // 读取文件内容 + content, err := os.ReadFile(path) + if err != nil { + m.logger.Warn("读取知识库文件失败", zap.String("path", path), zap.Error(err)) + return nil // 继续处理其他文件 + } + + // 检查是否已存在 + var existingID string + var existingContent string + var existingUpdatedAt time.Time + err = m.db.QueryRow( + "SELECT id, content, updated_at FROM knowledge_base_items WHERE file_path = ?", + path, + ).Scan(&existingID, &existingContent, &existingUpdatedAt) + + if err == sql.ErrNoRows { + // 创建新项 + id := uuid.New().String() + now := time.Now() + _, err = m.db.Exec( + "INSERT INTO knowledge_base_items (id, category, title, file_path, content, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, category, title, path, string(content), now, now, + ) + if err != nil { + return fmt.Errorf("插入知识项失败: %w", err) + } + m.logger.Info("添加知识项", zap.String("id", id), zap.String("title", title), zap.String("category", category)) + // 新添加的项需要索引 + itemsToIndex = append(itemsToIndex, id) + } else if err == nil { + // 检查内容是否有变化 + contentChanged := existingContent != string(content) + if contentChanged { + // 更新现有项 + _, err = m.db.Exec( + "UPDATE knowledge_base_items SET category = ?, title = ?, content = ?, updated_at = ? WHERE id = ?", + category, title, string(content), time.Now(), existingID, + ) + if err != nil { + return fmt.Errorf("更新知识项失败: %w", err) + } + m.logger.Info("更新知识项", zap.String("id", existingID), zap.String("title", title)) + // 内容已更新的项需要重新索引 + itemsToIndex = append(itemsToIndex, existingID) + } else { + m.logger.Debug("知识项未变化，跳过", zap.String("id", existingID), zap.String("title", title)) + } + } else { + return fmt.Errorf("查询知识项失败: %w", err) + } + + return nil + }) + + if err != nil { + return nil, err + } + + return itemsToIndex, nil +} + +// GetCategories 获取所有分类（风险类型） +func (m *Manager) GetCategories() ([]string, error) { + rows, err := m.db.Query("SELECT DISTINCT category FROM knowledge_base_items ORDER BY category") + if err != nil { + return nil, fmt.Errorf("查询分类失败: %w", err) + } + defer rows.Close() + + var categories []string + for rows.Next() { + var category string + if err := rows.Scan(&category); err != nil { + return nil, fmt.Errorf("扫描分类失败: %w", err) + } + categories = append(categories, category) + } + + return categories, nil +} + +// GetStats 获取知识库统计信息 +func (m *Manager) GetStats() (int, int, error) { + // 获取分类总数 + categories, err := m.GetCategories() + if err != nil { + return 0, 0, fmt.Errorf("获取分类失败: %w", err) + } + totalCategories := len(categories) + + // 获取知识项总数 + var totalItems int + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&totalItems) + if err != nil { + return totalCategories, 0, fmt.Errorf("获取知识项总数失败: %w", err) + } + + return totalCategories, totalItems, nil +} + +// GetCategoriesWithItems 按分类分页获取知识项（每个分类包含其下的所有知识项） +// limit: 每页分类数量（0表示不限制） +// offset: 偏移量（按分类偏移） +func (m *Manager) GetCategoriesWithItems(limit, offset int) ([]*CategoryWithItems, int, error) { + // 首先获取所有分类（带数量统计） + rows, err := m.db.Query(` + SELECT category, COUNT(*) as item_count + FROM knowledge_base_items + GROUP BY category + ORDER BY category + `) + if err != nil { + return nil, 0, fmt.Errorf("查询分类失败: %w", err) + } + defer rows.Close() + + // 收集所有分类信息 + type categoryInfo struct { + name string + itemCount int + } + var allCategories []categoryInfo + for rows.Next() { + var info categoryInfo + if err := rows.Scan(&info.name, &info.itemCount); err != nil { + return nil, 0, fmt.Errorf("扫描分类失败: %w", err) + } + allCategories = append(allCategories, info) + } + + totalCategories := len(allCategories) + + // 应用分页（按分类分页） + var paginatedCategories []categoryInfo + if limit > 0 { + start := offset + end := offset + limit + if start >= totalCategories { + paginatedCategories = []categoryInfo{} + } else { + if end > totalCategories { + end = totalCategories + } + paginatedCategories = allCategories[start:end] + } + } else { + paginatedCategories = allCategories + } + + // 为每个分类获取其下的知识项（只返回摘要，不包含完整内容） + result := make([]*CategoryWithItems, 0, len(paginatedCategories)) + for _, catInfo := range paginatedCategories { + // 获取该分类下的所有知识项 + items, _, err := m.GetItemsSummary(catInfo.name, 0, 0) + if err != nil { + return nil, 0, fmt.Errorf("获取分类 %s 的知识项失败: %w", catInfo.name, err) + } + + result = append(result, &CategoryWithItems{ + Category: catInfo.name, + ItemCount: catInfo.itemCount, + Items: items, + }) + } + + return result, totalCategories, nil +} + +// GetItems 获取知识项列表（完整内容，用于向后兼容） +func (m *Manager) GetItems(category string) ([]*KnowledgeItem, error) { + return m.GetItemsWithOptions(category, 0, 0, true) +} + +// GetItemsWithOptions 获取知识项列表（支持分页和可选内容） +// category: 分类筛选（空字符串表示所有分类） +// limit: 每页数量（0表示不限制） +// offset: 偏移量 +// includeContent: 是否包含完整内容（false时只返回摘要） +func (m *Manager) GetItemsWithOptions(category string, limit, offset int, includeContent bool) ([]*KnowledgeItem, error) { + var rows *sql.Rows + var err error + + // 构建SQL查询 + var query string + var args []interface{} + + if includeContent { + query = "SELECT id, category, title, file_path, content, created_at, updated_at FROM knowledge_base_items" + } else { + query = "SELECT id, category, title, file_path, created_at, updated_at FROM knowledge_base_items" + } + + if category != "" { + query += " WHERE category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + if limit > 0 { + query += " LIMIT ?" + args = append(args, limit) + if offset > 0 { + query += " OFFSET ?" + args = append(args, offset) + } + } + + rows, err = m.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("查询知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItem + for rows.Next() { + item := &KnowledgeItem{} + var createdAt, updatedAt string + + if includeContent { + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &item.Content, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + } else { + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + // 不包含内容时，Content为空字符串 + item.Content = "" + } + + // 解析时间 - 支持多种格式 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + // 解析创建时间 + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + // 解析更新时间 + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + // 如果更新时间为空，使用创建时间 + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, nil +} + +// GetItemsCount 获取知识项总数 +func (m *Manager) GetItemsCount(category string) (int, error) { + var count int + var err error + + if category != "" { + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items WHERE category = ?", category).Scan(&count) + } else { + err = m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&count) + } + + if err != nil { + return 0, fmt.Errorf("查询知识项总数失败: %w", err) + } + + return count, nil +} + +// SearchItemsByKeyword 按关键字搜索知识项（在所有数据中搜索，支持标题、分类、路径、内容匹配） +func (m *Manager) SearchItemsByKeyword(keyword string, category string) ([]*KnowledgeItemSummary, error) { + if keyword == "" { + return nil, fmt.Errorf("搜索关键字不能为空") + } + + // 构建SQL查询，使用LIKE进行关键字匹配（不区分大小写） + var query string + var args []interface{} + + // SQLite的LIKE不区分大小写，使用COLLATE NOCASE或LOWER()函数 + // 使用%keyword%进行模糊匹配 + searchPattern := "%" + keyword + "%" + + query = ` + SELECT id, category, title, file_path, created_at, updated_at + FROM knowledge_base_items + WHERE (LOWER(title) LIKE LOWER(?) OR LOWER(category) LIKE LOWER(?) OR LOWER(file_path) LIKE LOWER(?) OR LOWER(content) LIKE LOWER(?)) + ` + args = append(args, searchPattern, searchPattern, searchPattern, searchPattern) + + // 如果指定了分类，添加分类过滤 + if category != "" { + query += " AND category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + rows, err := m.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("搜索知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItemSummary + for rows.Next() { + item := &KnowledgeItemSummary{} + var createdAt, updatedAt string + + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, fmt.Errorf("扫描知识项失败: %w", err) + } + + // 解析时间 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, nil +} + +// GetItemsSummary 获取知识项摘要列表（不包含完整内容，支持分页） +func (m *Manager) GetItemsSummary(category string, limit, offset int) ([]*KnowledgeItemSummary, int, error) { + // 获取总数 + total, err := m.GetItemsCount(category) + if err != nil { + return nil, 0, err + } + + // 获取列表数据（不包含内容） + var rows *sql.Rows + var query string + var args []interface{} + + query = "SELECT id, category, title, file_path, created_at, updated_at FROM knowledge_base_items" + + if category != "" { + query += " WHERE category = ?" + args = append(args, category) + } + + query += " ORDER BY category, title" + + if limit > 0 { + query += " LIMIT ?" + args = append(args, limit) + if offset > 0 { + query += " OFFSET ?" + args = append(args, offset) + } + } + + rows, err = m.db.Query(query, args...) + if err != nil { + return nil, 0, fmt.Errorf("查询知识项失败: %w", err) + } + defer rows.Close() + + var items []*KnowledgeItemSummary + for rows.Next() { + item := &KnowledgeItemSummary{} + var createdAt, updatedAt string + + if err := rows.Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &createdAt, &updatedAt); err != nil { + return nil, 0, fmt.Errorf("扫描知识项失败: %w", err) + } + + // 解析时间 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + items = append(items, item) + } + + return items, total, nil +} + +// GetItem 获取单个知识项 +func (m *Manager) GetItem(id string) (*KnowledgeItem, error) { + item := &KnowledgeItem{} + var createdAt, updatedAt string + err := m.db.QueryRow( + "SELECT id, category, title, file_path, content, created_at, updated_at FROM knowledge_base_items WHERE id = ?", + id, + ).Scan(&item.ID, &item.Category, &item.Title, &item.FilePath, &item.Content, &createdAt, &updatedAt) + + if err == sql.ErrNoRows { + return nil, fmt.Errorf("知识项不存在") + } + if err != nil { + return nil, fmt.Errorf("查询知识项失败: %w", err) + } + + // 解析时间 - 支持多种格式 + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + // 解析创建时间 + if createdAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, createdAt) + if err == nil && !parsed.IsZero() { + item.CreatedAt = parsed + break + } + } + } + + // 解析更新时间 + if updatedAt != "" { + for _, format := range timeFormats { + parsed, err := time.Parse(format, updatedAt) + if err == nil && !parsed.IsZero() { + item.UpdatedAt = parsed + break + } + } + } + + // 如果更新时间为空，使用创建时间 + if item.UpdatedAt.IsZero() && !item.CreatedAt.IsZero() { + item.UpdatedAt = item.CreatedAt + } + + return item, nil +} + +// CreateItem 创建知识项 +func (m *Manager) CreateItem(category, title, content string) (*KnowledgeItem, error) { + id := uuid.New().String() + now := time.Now() + + // 构建文件路径 + filePath := filepath.Join(m.basePath, category, title+".md") + + // 确保目录存在 + if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil { + return nil, fmt.Errorf("创建目录失败: %w", err) + } + + // 写入文件 + if err := os.WriteFile(filePath, []byte(content), 0644); err != nil { + return nil, fmt.Errorf("写入文件失败: %w", err) + } + + // 插入数据库 + _, err := m.db.Exec( + "INSERT INTO knowledge_base_items (id, category, title, file_path, content, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, category, title, filePath, content, now, now, + ) + if err != nil { + return nil, fmt.Errorf("插入知识项失败: %w", err) + } + + return &KnowledgeItem{ + ID: id, + Category: category, + Title: title, + FilePath: filePath, + Content: content, + CreatedAt: now, + UpdatedAt: now, + }, nil +} + +// UpdateItem 更新知识项 +func (m *Manager) UpdateItem(id, category, title, content string) (*KnowledgeItem, error) { + // 获取现有项 + item, err := m.GetItem(id) + if err != nil { + return nil, err + } + + // 构建新文件路径 + newFilePath := filepath.Join(m.basePath, category, title+".md") + + // 如果路径改变，需要移动文件 + if item.FilePath != newFilePath { + // 确保新目录存在 + if err := os.MkdirAll(filepath.Dir(newFilePath), 0755); err != nil { + return nil, fmt.Errorf("创建目录失败: %w", err) + } + + // 移动文件 + if err := os.Rename(item.FilePath, newFilePath); err != nil { + return nil, fmt.Errorf("移动文件失败: %w", err) + } + + // 删除旧目录（如果为空） + oldDir := filepath.Dir(item.FilePath) + if isEmpty, _ := isEmptyDir(oldDir); isEmpty { + // 只有当目录不是知识库根目录时才删除（避免删除根目录） + if oldDir != m.basePath { + if err := os.Remove(oldDir); err != nil { + m.logger.Warn("删除空目录失败", zap.String("dir", oldDir), zap.Error(err)) + } + } + } + } + + // 写入文件 + if err := os.WriteFile(newFilePath, []byte(content), 0644); err != nil { + return nil, fmt.Errorf("写入文件失败: %w", err) + } + + // 更新数据库 + _, err = m.db.Exec( + "UPDATE knowledge_base_items SET category = ?, title = ?, file_path = ?, content = ?, updated_at = ? WHERE id = ?", + category, title, newFilePath, content, time.Now(), id, + ) + if err != nil { + return nil, fmt.Errorf("更新知识项失败: %w", err) + } + + // 删除旧的向量嵌入（需要重新索引） + _, err = m.db.Exec("DELETE FROM knowledge_embeddings WHERE item_id = ?", id) + if err != nil { + m.logger.Warn("删除旧向量嵌入失败", zap.Error(err)) + } + + return m.GetItem(id) +} + +// DeleteItem 删除知识项 +func (m *Manager) DeleteItem(id string) error { + // 获取文件路径 + var filePath string + err := m.db.QueryRow("SELECT file_path FROM knowledge_base_items WHERE id = ?", id).Scan(&filePath) + if err != nil { + return fmt.Errorf("查询知识项失败: %w", err) + } + + // 删除文件 + if err := os.Remove(filePath); err != nil && !os.IsNotExist(err) { + m.logger.Warn("删除文件失败", zap.String("path", filePath), zap.Error(err)) + } + + // 删除数据库记录（级联删除向量） + _, err = m.db.Exec("DELETE FROM knowledge_base_items WHERE id = ?", id) + if err != nil { + return fmt.Errorf("删除知识项失败: %w", err) + } + + // 删除空目录（如果为空） + dir := filepath.Dir(filePath) + if isEmpty, _ := isEmptyDir(dir); isEmpty { + // 只有当目录不是知识库根目录时才删除（避免删除根目录） + if dir != m.basePath { + if err := os.Remove(dir); err != nil { + m.logger.Warn("删除空目录失败", zap.String("dir", dir), zap.Error(err)) + } + } + } + + return nil +} + +// isEmptyDir 检查目录是否为空（忽略隐藏文件和 . 开头的文件） +func isEmptyDir(dir string) (bool, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return false, err + } + for _, entry := range entries { + // 忽略隐藏文件（以 . 开头） + if !strings.HasPrefix(entry.Name(), ".") { + return false, nil + } + } + return true, nil +} + +// LogRetrieval 记录检索日志 +func (m *Manager) LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error { + id := uuid.New().String() + itemsJSON, _ := json.Marshal(retrievedItems) + + _, err := m.db.Exec( + "INSERT INTO knowledge_retrieval_logs (id, conversation_id, message_id, query, risk_type, retrieved_items, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)", + id, conversationID, messageID, query, riskType, string(itemsJSON), time.Now(), + ) + return err +} + +// GetIndexStatus 获取索引状态 +func (m *Manager) GetIndexStatus() (map[string]interface{}, error) { + // 获取总知识项数 + var totalItems int + err := m.db.QueryRow("SELECT COUNT(*) FROM knowledge_base_items").Scan(&totalItems) + if err != nil { + return nil, fmt.Errorf("查询总知识项数失败: %w", err) + } + + // 获取已索引的知识项数（有向量嵌入的） + var indexedItems int + err = m.db.QueryRow(` + SELECT COUNT(DISTINCT item_id) + FROM knowledge_embeddings + `).Scan(&indexedItems) + if err != nil { + return nil, fmt.Errorf("查询已索引项数失败: %w", err) + } + + // 计算进度百分比 + var progressPercent float64 + if totalItems > 0 { + progressPercent = float64(indexedItems) / float64(totalItems) * 100 + } else { + progressPercent = 100.0 + } + + // 判断是否完成 + isComplete := indexedItems >= totalItems && totalItems > 0 + + return map[string]interface{}{ + "total_items": totalItems, + "indexed_items": indexedItems, + "progress_percent": progressPercent, + "is_complete": isComplete, + }, nil +} + +// GetRetrievalLogs 获取检索日志 +func (m *Manager) GetRetrievalLogs(conversationID, messageID string, limit int) ([]*RetrievalLog, error) { + var rows *sql.Rows + var err error + + if messageID != "" { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs WHERE message_id = ? ORDER BY created_at DESC LIMIT ?", + messageID, limit, + ) + } else if conversationID != "" { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs WHERE conversation_id = ? ORDER BY created_at DESC LIMIT ?", + conversationID, limit, + ) + } else { + rows, err = m.db.Query( + "SELECT id, conversation_id, message_id, query, risk_type, retrieved_items, created_at FROM knowledge_retrieval_logs ORDER BY created_at DESC LIMIT ?", + limit, + ) + } + + if err != nil { + return nil, fmt.Errorf("查询检索日志失败: %w", err) + } + defer rows.Close() + + var logs []*RetrievalLog + for rows.Next() { + log := &RetrievalLog{} + var createdAt string + var itemsJSON sql.NullString + if err := rows.Scan(&log.ID, &log.ConversationID, &log.MessageID, &log.Query, &log.RiskType, &itemsJSON, &createdAt); err != nil { + return nil, fmt.Errorf("扫描检索日志失败: %w", err) + } + + // 解析时间 - 支持多种格式 + var err error + timeFormats := []string{ + "2006-01-02 15:04:05.999999999-07:00", + "2006-01-02 15:04:05.999999999", + "2006-01-02T15:04:05.999999999Z07:00", + "2006-01-02T15:04:05Z", + "2006-01-02 15:04:05", + time.RFC3339, + time.RFC3339Nano, + } + + for _, format := range timeFormats { + log.CreatedAt, err = time.Parse(format, createdAt) + if err == nil && !log.CreatedAt.IsZero() { + break + } + } + + // 如果所有格式都失败，记录警告但继续处理 + if log.CreatedAt.IsZero() { + m.logger.Warn("解析检索日志时间失败", + zap.String("timeStr", createdAt), + zap.Error(err), + ) + // 使用当前时间作为fallback + log.CreatedAt = time.Now() + } + + // 解析检索项 + if itemsJSON.Valid { + json.Unmarshal([]byte(itemsJSON.String), &log.RetrievedItems) + } + + logs = append(logs, log) + } + + return logs, nil +} + +// DeleteRetrievalLog 删除检索日志 +func (m *Manager) DeleteRetrievalLog(id string) error { + result, err := m.db.Exec("DELETE FROM knowledge_retrieval_logs WHERE id = ?", id) + if err != nil { + return fmt.Errorf("删除检索日志失败: %w", err) + } + + rowsAffected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("获取删除行数失败: %w", err) + } + + if rowsAffected == 0 { + return fmt.Errorf("检索日志不存在") + } + + return nil +} diff --git a/internal/knowledge/rerank_http.go b/internal/knowledge/rerank_http.go new file mode 100644 index 00000000..61e173ed --- /dev/null +++ b/internal/knowledge/rerank_http.go @@ -0,0 +1,226 @@ +package knowledge + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// HTTPReranker calls a hosted rerank API (DashScope or Cohere-compatible). +type HTTPReranker struct { + provider string + model string + baseURL string + apiKey string + client *http.Client + logger *zap.Logger +} + +// NewHTTPReranker builds a rerank client from knowledge retrieval config; openAI supplies fallback credentials. +func NewHTTPReranker(rc *config.RerankConfig, openAI *config.OpenAIConfig, logger *zap.Logger) (*HTTPReranker, error) { + if rc == nil { + return nil, fmt.Errorf("rerank config is nil") + } + baseURL := strings.TrimSpace(rc.BaseURL) + apiKey := strings.TrimSpace(rc.APIKey) + if openAI != nil { + if baseURL == "" { + baseURL = strings.TrimSpace(openAI.BaseURL) + } + if apiKey == "" { + apiKey = strings.TrimSpace(openAI.APIKey) + } + } + if apiKey == "" { + return nil, fmt.Errorf("rerank api_key is required") + } + provider := rc.ProviderEffective(baseURL) + model := rc.ModelEffective(provider) + return &HTTPReranker{ + provider: provider, + model: model, + baseURL: strings.TrimSuffix(baseURL, "/"), + apiKey: apiKey, + client: &http.Client{Timeout: 60 * time.Second}, + logger: logger, + }, nil +} + +func (r *HTTPReranker) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) { + if r == nil { + return docs, nil + } + q := strings.TrimSpace(query) + if q == "" || len(docs) == 0 { + return docs, nil + } + if len(docs) == 1 { + return docs, nil + } + texts := make([]string, 0, len(docs)) + for _, d := range docs { + if d == nil { + texts = append(texts, "") + continue + } + texts = append(texts, d.Content) + } + var order []int + var err error + switch r.provider { + case "dashscope": + order, err = r.rerankDashScope(ctx, q, texts, len(docs)) + default: + order, err = r.rerankCohere(ctx, q, texts, len(docs)) + } + if err != nil { + return nil, err + } + out := make([]*schema.Document, 0, len(order)) + for _, idx := range order { + if idx < 0 || idx >= len(docs) || docs[idx] == nil { + continue + } + out = append(out, docs[idx]) + } + if len(out) == 0 { + return docs, nil + } + return out, nil +} + +func (r *HTTPReranker) rerankCohere(ctx context.Context, query string, documents []string, topN int) ([]int, error) { + url := r.cohereRerankURL() + body := map[string]any{ + "model": r.model, + "query": query, + "documents": documents, + "top_n": topN, + } + raw, err := json.Marshal(body) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+r.apiKey) + resp, err := r.client.Do(req) + if err != nil { + return nil, fmt.Errorf("rerank request: %w", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody))) + } + var parsed struct { + Results []struct { + Index int `json:"index"` + } `json:"results"` + } + if err := json.Unmarshal(respBody, &parsed); err != nil { + return nil, fmt.Errorf("rerank decode: %w", err) + } + order := make([]int, 0, len(parsed.Results)) + for _, row := range parsed.Results { + order = append(order, row.Index) + } + return order, nil +} + +func (r *HTTPReranker) rerankDashScope(ctx context.Context, query string, documents []string, topN int) ([]int, error) { + url := r.dashscopeRerankURL() + body := map[string]any{ + "model": r.model, + "input": map[string]any{ + "query": query, + "documents": documents, + }, + "parameters": map[string]any{ + "return_documents": false, + "top_n": topN, + }, + } + raw, err := json.Marshal(body) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+r.apiKey) + resp, err := r.client.Do(req) + if err != nil { + return nil, fmt.Errorf("dashscope rerank: %w", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("dashscope rerank http %d: %s", resp.StatusCode, truncateForRerankLog(string(respBody))) + } + var parsed struct { + Output struct { + Results []struct { + Index int `json:"index"` + } `json:"results"` + } `json:"output"` + } + if err := json.Unmarshal(respBody, &parsed); err != nil { + return nil, fmt.Errorf("dashscope rerank decode: %w", err) + } + order := make([]int, 0, len(parsed.Output.Results)) + for _, row := range parsed.Output.Results { + order = append(order, row.Index) + } + return order, nil +} + +func (r *HTTPReranker) cohereRerankURL() string { + base := r.baseURL + if base == "" { + base = "https://api.cohere.com" + } + if strings.HasSuffix(base, "/v1") { + return base + "/rerank" + } + return base + "/v1/rerank" +} + +func (r *HTTPReranker) dashscopeRerankURL() string { + base := strings.TrimSpace(r.baseURL) + if base == "" { + return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank" + } + if strings.Contains(base, "/api/v1/services/rerank") { + return base + } + if strings.Contains(base, "dashscope.aliyuncs.com") || strings.Contains(base, "compatible-mode") { + return "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank" + } + return strings.TrimSuffix(base, "/") +} + +func truncateForRerankLog(s string) string { + s = strings.TrimSpace(s) + if len(s) > 512 { + return s[:512] + "..." + } + return s +} + +var _ DocumentReranker = (*HTTPReranker)(nil) diff --git a/internal/knowledge/rerank_http_test.go b/internal/knowledge/rerank_http_test.go new file mode 100644 index 00000000..013ad7d4 --- /dev/null +++ b/internal/knowledge/rerank_http_test.go @@ -0,0 +1,97 @@ +package knowledge + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" +) + +func TestHTTPReranker_CohereOrder(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/rerank" { + t.Fatalf("path %s", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "results": []map[string]any{ + {"index": 2, "relevance_score": 0.9}, + {"index": 0, "relevance_score": 0.5}, + }, + }) + })) + defer srv.Close() + + rr, err := NewHTTPReranker(&config.RerankConfig{ + Provider: "cohere", + Model: "rerank-multilingual-v3.0", + BaseURL: srv.URL, + APIKey: "test-key", + }, nil, nil) + if err != nil { + t.Fatal(err) + } + docs := []*schema.Document{ + {ID: "a", Content: "alpha"}, + {ID: "b", Content: "beta"}, + {ID: "c", Content: "gamma"}, + } + out, err := rr.Rerank(context.Background(), "query", docs) + if err != nil { + t.Fatal(err) + } + if len(out) != 2 || out[0].ID != "c" || out[1].ID != "a" { + t.Fatalf("order wrong: %#v", out) + } +} + +func TestHTTPReranker_DashScopeOrder(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "output": map[string]any{ + "results": []map[string]any{ + {"index": 1, "relevance_score": 0.88}, + }, + }, + }) + })) + defer srv.Close() + + rr, err := NewHTTPReranker(&config.RerankConfig{ + Provider: "dashscope", + Model: "gte-rerank", + BaseURL: srv.URL, + APIKey: "test-key", + }, nil, nil) + if err != nil { + t.Fatal(err) + } + docs := []*schema.Document{{ID: "a", Content: "a"}, {ID: "b", Content: "b"}} + out, err := rr.Rerank(context.Background(), "q", docs) + if err != nil { + t.Fatal(err) + } + if len(out) != 1 || out[0].ID != "b" { + t.Fatalf("got %#v", out) + } +} + +func TestRerankConfigDefaults(t *testing.T) { + t.Parallel() + rc := config.RerankConfig{} + if rc.ProviderEffective("https://dashscope.aliyuncs.com/x") != "dashscope" { + t.Fatal("dashscope detect") + } + if rc.ModelEffective("dashscope") != "gte-rerank" { + t.Fatal("dashscope model") + } + if rc.ModelEffective("cohere") != "rerank-multilingual-v3.0" { + t.Fatal("cohere model") + } +} diff --git a/internal/knowledge/retrieval_postprocess.go b/internal/knowledge/retrieval_postprocess.go new file mode 100644 index 00000000..20e07110 --- /dev/null +++ b/internal/knowledge/retrieval_postprocess.go @@ -0,0 +1,216 @@ +package knowledge + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" + "github.com/pkoukk/tiktoken-go" +) + +// postRetrieveMaxPrefetchCap 限制单次向量候选上限，避免误配置导致全表扫压力过大。 +const postRetrieveMaxPrefetchCap = 200 + +// DocumentReranker 精排（HTTP dashscope / Cohere 兼容 API），由 [WireRetrieverPipeline] 注入。 +type DocumentReranker interface { + Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) +} + +// NopDocumentReranker 占位实现，便于测试或未启用重排时显式注入。 +type NopDocumentReranker struct{} + +// Rerank implements [DocumentReranker] as no-op. +func (NopDocumentReranker) Rerank(_ context.Context, _ string, docs []*schema.Document) ([]*schema.Document, error) { + return docs, nil +} + +var tiktokenEncMu sync.Mutex +var tiktokenEncCache = map[string]*tiktoken.Tiktoken{} + +func encodingForTokenizerModel(model string) (*tiktoken.Tiktoken, error) { + m := strings.TrimSpace(model) + if m == "" { + m = "gpt-4" + } + tiktokenEncMu.Lock() + defer tiktokenEncMu.Unlock() + if enc, ok := tiktokenEncCache[m]; ok { + return enc, nil + } + enc, err := tiktoken.EncodingForModel(m) + if err != nil { + enc, err = tiktoken.GetEncoding("cl100k_base") + if err != nil { + return nil, err + } + } + tiktokenEncCache[m] = enc + return enc, nil +} + +func countDocTokens(text, model string) (int, error) { + enc, err := encodingForTokenizerModel(model) + if err != nil { + return 0, err + } + toks := enc.Encode(text, nil, nil) + return len(toks), nil +} + +// normalizeContentFingerprintKey 去重键：trim + 空白折叠（不改动大小写，避免合并仅大小写不同的代码片段）。 +func normalizeContentFingerprintKey(s string) string { + s = strings.TrimSpace(s) + var b strings.Builder + b.Grow(len(s)) + prevSpace := false + for _, r := range s { + if unicode.IsSpace(r) { + if !prevSpace { + b.WriteByte(' ') + prevSpace = true + } + continue + } + prevSpace = false + b.WriteRune(r) + } + return b.String() +} + +func contentNormKey(d *schema.Document) string { + if d == nil { + return "" + } + n := normalizeContentFingerprintKey(d.Content) + if n == "" { + return "" + } + sum := sha256.Sum256([]byte(n)) + return hex.EncodeToString(sum[:]) +} + +// dedupeByNormalizedContent 按规范化正文去重，保留向量检索顺序中首次出现的文档（同正文仅保留一条）。 +func dedupeByNormalizedContent(docs []*schema.Document) []*schema.Document { + if len(docs) < 2 { + return docs + } + seen := make(map[string]struct{}, len(docs)) + out := make([]*schema.Document, 0, len(docs)) + for _, d := range docs { + if d == nil { + continue + } + k := contentNormKey(d) + if k == "" { + out = append(out, d) + continue + } + if _, ok := seen[k]; ok { + continue + } + seen[k] = struct{}{} + out = append(out, d) + } + return out +} + +// truncateDocumentsByBudget 按检索顺序整段保留文档，直至字符数或 token 数（任一启用）超限则停止。 +func truncateDocumentsByBudget(docs []*schema.Document, maxRunes, maxTokens int, tokenModel string) ([]*schema.Document, error) { + if len(docs) == 0 { + return docs, nil + } + unlimitedChars := maxRunes <= 0 + unlimitedTok := maxTokens <= 0 + if unlimitedChars && unlimitedTok { + return docs, nil + } + + remRunes := maxRunes + remTok := maxTokens + out := make([]*schema.Document, 0, len(docs)) + + for _, d := range docs { + if d == nil || strings.TrimSpace(d.Content) == "" { + continue + } + runes := utf8.RuneCountInString(d.Content) + if !unlimitedChars && runes > remRunes { + break + } + var tok int + var err error + if !unlimitedTok { + tok, err = countDocTokens(d.Content, tokenModel) + if err != nil { + return nil, fmt.Errorf("token count: %w", err) + } + if tok > remTok { + break + } + } + out = append(out, d) + if !unlimitedChars { + remRunes -= runes + } + if !unlimitedTok { + remTok -= tok + } + } + return out, nil +} + +// EffectivePrefetchTopK 计算每条 MultiQuery 变体在向量阶段的候选条数（供融合 / 重排 / 后处理）。 +func EffectivePrefetchTopK(topK int, po *config.PostRetrieveConfig) int { + if topK < 1 { + topK = 5 + } + fetch := topK * 4 + if fetch < 20 { + fetch = 20 + } + if po != nil && po.PrefetchTopK > 0 { + fetch = po.PrefetchTopK + } + if fetch > postRetrieveMaxPrefetchCap { + fetch = postRetrieveMaxPrefetchCap + } + return fetch +} + +// ApplyPostRetrieve 检索后处理：规范化正文去重 → 预算截断 → 最终 TopK（精排已在流水线中完成）。 +func ApplyPostRetrieve(docs []*schema.Document, po *config.PostRetrieveConfig, tokenModel string, finalTopK int) ([]*schema.Document, error) { + if finalTopK < 1 { + finalTopK = 5 + } + if len(docs) == 0 { + return docs, nil + } + + maxChars := 0 + maxTok := 0 + if po != nil { + maxChars = po.MaxContextChars + maxTok = po.MaxContextTokens + } + + out := dedupeByNormalizedContent(docs) + + var err error + out, err = truncateDocumentsByBudget(out, maxChars, maxTok, tokenModel) + if err != nil { + return nil, err + } + + if len(out) > finalTopK { + out = out[:finalTopK] + } + return out, nil +} diff --git a/internal/knowledge/retrieval_postprocess_test.go b/internal/knowledge/retrieval_postprocess_test.go new file mode 100644 index 00000000..889d5c62 --- /dev/null +++ b/internal/knowledge/retrieval_postprocess_test.go @@ -0,0 +1,62 @@ +package knowledge + +import ( + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/schema" +) + +func doc(id, content string, score float64) *schema.Document { + d := &schema.Document{ID: id, Content: content, MetaData: map[string]any{metaKBItemID: "it1"}} + d.WithScore(score) + return d +} + +func TestDedupeByNormalizedContent(t *testing.T) { + a := doc("1", "hello world", 0.9) + b := doc("2", "hello world", 0.8) + c := doc("3", "other", 0.7) + out := dedupeByNormalizedContent([]*schema.Document{a, b, c}) + if len(out) != 2 { + t.Fatalf("len=%d want 2", len(out)) + } + if out[0].ID != "1" || out[1].ID != "3" { + t.Fatalf("order/ids wrong: %#v", out) + } +} + +func TestEffectivePrefetchTopK(t *testing.T) { + if g := EffectivePrefetchTopK(5, nil); g != 20 { + t.Fatalf("default prefetch got %d want 20", g) + } + if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 50}); g != 50 { + t.Fatalf("got %d", g) + } + if g := EffectivePrefetchTopK(5, &config.PostRetrieveConfig{PrefetchTopK: 9999}); g != postRetrieveMaxPrefetchCap { + t.Fatalf("cap: got %d", g) + } +} + +func TestApplyPostRetrieveTruncateAndTopK(t *testing.T) { + d1 := doc("1", "ab", 0.9) + d2 := doc("2", "cd", 0.8) + d3 := doc("3", "ef", 0.7) + po := &config.PostRetrieveConfig{MaxContextChars: 3} + out, err := ApplyPostRetrieve([]*schema.Document{d1, d2, d3}, po, "gpt-4", 5) + if err != nil { + t.Fatal(err) + } + if len(out) != 1 || out[0].ID != "1" { + t.Fatalf("got %#v", out) + } + + out2, err := ApplyPostRetrieve([]*schema.Document{d1, d2, d3}, nil, "gpt-4", 2) + if err != nil { + t.Fatal(err) + } + if len(out2) != 2 { + t.Fatalf("topk: len=%d", len(out2)) + } +} diff --git a/internal/knowledge/retriever.go b/internal/knowledge/retriever.go new file mode 100644 index 00000000..c75e8a13 --- /dev/null +++ b/internal/knowledge/retriever.go @@ -0,0 +1,334 @@ +package knowledge + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "math" + "sort" + "strings" + "sync" + + "cyberstrike-ai/internal/config" + + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" + "go.uber.org/zap" +) + +// Retriever 检索器：SQLite 存向量 + Eino 嵌入，**纯向量检索**（余弦相似度、TopK、阈值）， +// 实现语义与 [retriever.Retriever] 适配层 [VectorEinoRetriever] 一致。 +type Retriever struct { + db *sql.DB + embedder *Embedder + config *RetrievalConfig + logger *zap.Logger + + rerankMu sync.RWMutex + reranker DocumentReranker + + pipeline retriever.Retriever + wireOpenAI *config.OpenAIConfig +} + +// RetrievalConfig 检索配置 +type RetrievalConfig struct { + TopK int + SimilarityThreshold float64 + SubIndexFilter string + MultiQuery config.MultiQueryConfig + Rerank config.RerankConfig + PostRetrieve config.PostRetrieveConfig +} + +// NewRetriever 创建新的检索器 +func NewRetriever(db *sql.DB, embedder *Embedder, config *RetrievalConfig, logger *zap.Logger) *Retriever { + return &Retriever{ + db: db, + embedder: embedder, + config: config, + logger: logger, + } +} + +// UpdateConfig 更新检索配置并重建 Eino MultiQuery + 重排流水线。 +func (r *Retriever) UpdateConfig(cfg *RetrievalConfig) { + if cfg != nil { + r.config = cfg + if r.logger != nil { + r.logger.Info("检索器配置已更新", + zap.Int("top_k", cfg.TopK), + zap.Float64("similarity_threshold", cfg.SimilarityThreshold), + zap.String("sub_index_filter", cfg.SubIndexFilter), + zap.Int("multi_query_max", cfg.MultiQuery.MaxQueriesEffective()), + zap.Int("post_retrieve_prefetch_top_k", cfg.PostRetrieve.PrefetchTopK), + zap.Int("post_retrieve_max_context_chars", cfg.PostRetrieve.MaxContextChars), + zap.Int("post_retrieve_max_context_tokens", cfg.PostRetrieve.MaxContextTokens), + ) + } + } + if r.wireOpenAI != nil { + if err := WireRetrieverPipeline(context.Background(), r, r.wireOpenAI); err != nil && r.logger != nil { + r.logger.Warn("检索流水线重建失败", zap.Error(err)) + } + } +} + +// SetDocumentReranker 注入可选重排器（并发安全）；nil 表示禁用。 +func (r *Retriever) SetDocumentReranker(rr DocumentReranker) { + if r == nil { + return + } + r.rerankMu.Lock() + defer r.rerankMu.Unlock() + r.reranker = rr +} + +func (r *Retriever) documentReranker() DocumentReranker { + if r == nil { + return nil + } + r.rerankMu.RLock() + defer r.rerankMu.RUnlock() + return r.reranker +} + +func cosineSimilarity(a, b []float32) float64 { + if len(a) != len(b) { + return 0.0 + } + + var dotProduct, normA, normB float64 + for i := range a { + dotProduct += float64(a[i] * b[i]) + normA += float64(a[i] * a[i]) + normB += float64(b[i] * b[i]) + } + + if normA == 0 || normB == 0 { + return 0.0 + } + + return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)) +} + +// Search 搜索知识库（Eino MultiQuery → 向量检索 → 重排 → 后处理）。 +func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) { + if req == nil { + return nil, fmt.Errorf("请求不能为空") + } + q := strings.TrimSpace(req.Query) + if q == "" { + return nil, fmt.Errorf("查询不能为空") + } + opts := r.einoRetrieverOptions(req) + docs, err := r.activeEinoRetriever().Retrieve(ctx, q, opts...) + if err != nil { + return nil, err + } + return documentsToRetrievalResults(docs) +} + +func (r *Retriever) einoRetrieverOptions(req *SearchRequest) []retriever.Option { + var opts []retriever.Option + if req.TopK > 0 { + opts = append(opts, retriever.WithTopK(req.TopK)) + } + dsl := map[string]any{} + if strings.TrimSpace(req.RiskType) != "" { + dsl[DSLRiskType] = strings.TrimSpace(req.RiskType) + } + if req.Threshold > 0 { + dsl[DSLSimilarityThreshold] = req.Threshold + } + if strings.TrimSpace(req.SubIndexFilter) != "" { + dsl[DSLSubIndexFilter] = strings.TrimSpace(req.SubIndexFilter) + } + if len(dsl) > 0 { + opts = append(opts, retriever.WithDSLInfo(dsl)) + } + return opts +} + +// EinoRetrieve 直接返回 [schema.Document]，供 Eino Graph / Chain 使用。 +func (r *Retriever) EinoRetrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + return r.activeEinoRetriever().Retrieve(ctx, query, opts...) +} + +func (r *Retriever) activeEinoRetriever() retriever.Retriever { + if r != nil && r.pipeline != nil { + return r.pipeline + } + return NewVectorEinoRetriever(r) +} + +// AsEinoRetriever 将知识库检索流水线暴露为 Eino [retriever.Retriever]。 +func (r *Retriever) AsEinoRetriever() retriever.Retriever { + return r.activeEinoRetriever() +} + +func (r *Retriever) knowledgeEmbeddingSelectSQL(riskType, subIndexFilter string) (string, []interface{}) { + q := `SELECT e.id, e.item_id, e.chunk_index, e.chunk_text, e.embedding, e.embedding_model, e.embedding_dim, i.category, i.title +FROM knowledge_embeddings e +JOIN knowledge_base_items i ON e.item_id = i.id +WHERE 1=1` + var args []interface{} + if strings.TrimSpace(riskType) != "" { + q += ` AND TRIM(i.category) = TRIM(?) COLLATE NOCASE` + args = append(args, riskType) + } + if tag := strings.TrimSpace(subIndexFilter); tag != "" { + tag = strings.ToLower(strings.ReplaceAll(tag, " ", "")) + q += ` AND (TRIM(COALESCE(e.sub_indexes,'')) = '' OR INSTR(',' || LOWER(REPLACE(e.sub_indexes,' ','')) || ',', ',' || ? || ',') > 0)` + args = append(args, tag) + } + return q, args +} + +// vectorSearch 纯向量检索：余弦相似度排序，按相似度阈值与 TopK 截断（无 BM25、无混合分、无邻块扩展）。 +func (r *Retriever) vectorSearch(ctx context.Context, req *SearchRequest) ([]*RetrievalResult, error) { + if req.Query == "" { + return nil, fmt.Errorf("查询不能为空") + } + + topK := req.TopK + if topK <= 0 && r.config != nil { + topK = r.config.TopK + } + if topK <= 0 { + topK = 5 + } + + threshold := req.Threshold + if threshold <= 0 && r.config != nil { + threshold = r.config.SimilarityThreshold + } + if threshold <= 0 { + threshold = 0.7 + } + + subIdxFilter := strings.TrimSpace(req.SubIndexFilter) + if subIdxFilter == "" && r.config != nil { + subIdxFilter = strings.TrimSpace(r.config.SubIndexFilter) + } + + queryText := FormatQueryEmbeddingText(req.RiskType, req.Query) + queryEmbedding, err := r.embedder.EmbedText(ctx, queryText) + if err != nil { + return nil, fmt.Errorf("向量化查询失败: %w", err) + } + queryDim := len(queryEmbedding) + expectedModel := "" + if r.embedder != nil { + expectedModel = r.embedder.EmbeddingModelName() + } + + sqlStr, sqlArgs := r.knowledgeEmbeddingSelectSQL(strings.TrimSpace(req.RiskType), subIdxFilter) + rows, err := r.db.QueryContext(ctx, sqlStr, sqlArgs...) + if err != nil { + return nil, fmt.Errorf("查询向量失败: %w", err) + } + defer rows.Close() + + type candidate struct { + chunk *KnowledgeChunk + item *KnowledgeItem + similarity float64 + } + + candidates := make([]candidate, 0) + rowNum := 0 + for rows.Next() { + rowNum++ + if rowNum%48 == 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + } + + var chunkID, itemID, chunkText, embeddingJSON, category, title, rowModel string + var chunkIndex, rowDim int + + if err := rows.Scan(&chunkID, &itemID, &chunkIndex, &chunkText, &embeddingJSON, &rowModel, &rowDim, &category, &title); err != nil { + r.logger.Warn("扫描向量失败", zap.Error(err)) + continue + } + + var embedding []float32 + if err := json.Unmarshal([]byte(embeddingJSON), &embedding); err != nil { + r.logger.Warn("解析向量失败", zap.Error(err)) + continue + } + + if rowDim > 0 && len(embedding) != rowDim { + r.logger.Debug("跳过维度不一致的向量行", zap.String("chunkId", chunkID), zap.Int("rowDim", rowDim), zap.Int("got", len(embedding))) + continue + } + if queryDim > 0 && len(embedding) != queryDim { + r.logger.Debug("跳过与查询维度不一致的向量", zap.String("chunkId", chunkID), zap.Int("queryDim", queryDim), zap.Int("got", len(embedding))) + continue + } + if expectedModel != "" && strings.TrimSpace(rowModel) != "" && strings.TrimSpace(rowModel) != expectedModel { + r.logger.Debug("跳过嵌入模型不一致的行", zap.String("chunkId", chunkID), zap.String("rowModel", rowModel), zap.String("expected", expectedModel)) + continue + } + + similarity := cosineSimilarity(queryEmbedding, embedding) + candidates = append(candidates, candidate{ + chunk: &KnowledgeChunk{ + ID: chunkID, + ItemID: itemID, + ChunkIndex: chunkIndex, + ChunkText: chunkText, + Embedding: embedding, + }, + item: &KnowledgeItem{ + ID: itemID, + Category: category, + Title: title, + }, + similarity: similarity, + }) + } + + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].similarity > candidates[j].similarity + }) + + filtered := make([]candidate, 0, len(candidates)) + for _, c := range candidates { + if c.similarity >= threshold { + filtered = append(filtered, c) + } + } + + if len(filtered) > topK { + filtered = filtered[:topK] + } + + results := make([]*RetrievalResult, len(filtered)) + for i, c := range filtered { + results[i] = &RetrievalResult{ + Chunk: c.chunk, + Item: c.item, + Similarity: c.similarity, + Score: c.similarity, + } + } + return results, nil +} + +// RetrievalConfigFromYAML maps API/YAML retrieval settings into the knowledge package. +func RetrievalConfigFromYAML(r config.RetrievalConfig) *RetrievalConfig { + return &RetrievalConfig{ + TopK: r.TopK, + SimilarityThreshold: r.SimilarityThreshold, + SubIndexFilter: r.SubIndexFilter, + MultiQuery: r.MultiQuery, + Rerank: r.Rerank, + PostRetrieve: r.PostRetrieve, + } +} diff --git a/internal/knowledge/schema_migrate.go b/internal/knowledge/schema_migrate.go new file mode 100644 index 00000000..85fd26e2 --- /dev/null +++ b/internal/knowledge/schema_migrate.go @@ -0,0 +1,51 @@ +package knowledge + +import ( + "database/sql" + "fmt" +) + +// EnsureKnowledgeEmbeddingsSchema migrates knowledge_embeddings for sub_indexes + embedding metadata. +func EnsureKnowledgeEmbeddingsSchema(db *sql.DB) error { + if db == nil { + return fmt.Errorf("db is nil") + } + var n int + if err := db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='knowledge_embeddings'`).Scan(&n); err != nil { + return err + } + if n == 0 { + return nil + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "sub_indexes", + `ALTER TABLE knowledge_embeddings ADD COLUMN sub_indexes TEXT NOT NULL DEFAULT ''`); err != nil { + return err + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_model", + `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_model TEXT NOT NULL DEFAULT ''`); err != nil { + return err + } + if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_dim", + `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_dim INTEGER NOT NULL DEFAULT 0`); err != nil { + return err + } + return nil +} + +func addKnowledgeEmbeddingsColumnIfMissing(db *sql.DB, column, alterSQL string) error { + var colCount int + q := `SELECT COUNT(*) FROM pragma_table_info('knowledge_embeddings') WHERE name = ?` + if err := db.QueryRow(q, column).Scan(&colCount); err != nil { + return err + } + if colCount > 0 { + return nil + } + _, err := db.Exec(alterSQL) + return err +} + +// ensureKnowledgeEmbeddingsSubIndexesColumn 向后兼容；请使用 [EnsureKnowledgeEmbeddingsSchema]。 +func ensureKnowledgeEmbeddingsSubIndexesColumn(db *sql.DB) error { + return EnsureKnowledgeEmbeddingsSchema(db) +} diff --git a/internal/knowledge/tool.go b/internal/knowledge/tool.go new file mode 100644 index 00000000..ffceb7e2 --- /dev/null +++ b/internal/knowledge/tool.go @@ -0,0 +1,323 @@ +package knowledge + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +// RegisterKnowledgeTool 注册知识检索工具到MCP服务器 +func RegisterKnowledgeTool( + mcpServer *mcp.Server, + retriever *Retriever, + manager *Manager, + logger *zap.Logger, +) { + // 注册第一个工具：获取所有可用的风险类型列表 + listRiskTypesTool := mcp.Tool{ + Name: builtin.ToolListKnowledgeRiskTypes, + Description: "获取知识库中所有可用的风险类型（risk_type）列表。在搜索知识库之前，可以先调用此工具获取可用的风险类型，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间并提高检索准确性。", + ShortDescription: "获取知识库中所有可用的风险类型列表", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + "required": []string{}, + }, + } + + listRiskTypesHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + categories, err := manager.GetCategories() + if err != nil { + logger.Error("获取风险类型列表失败", zap.Error(err)) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("获取风险类型列表失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + if len(categories) == 0 { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "知识库中暂无风险类型。", + }, + }, + }, nil + } + + var resultText strings.Builder + resultText.WriteString(fmt.Sprintf("知识库中共有 %d 个风险类型：\n\n", len(categories))) + for i, category := range categories { + resultText.WriteString(fmt.Sprintf("%d. %s\n", i+1, category)) + } + resultText.WriteString("\n提示：在调用 " + builtin.ToolSearchKnowledgeBase + " 工具时，可以使用上述风险类型之一作为 risk_type 参数，以缩小搜索范围并提高检索效率。") + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: resultText.String(), + }, + }, + }, nil + } + + mcpServer.RegisterTool(listRiskTypesTool, listRiskTypesHandler) + logger.Debug("风险类型列表工具已注册", zap.String("toolName", listRiskTypesTool.Name)) + + // 注册第二个工具：搜索知识库（保持原有功能） + searchTool := mcp.Tool{ + Name: builtin.ToolSearchKnowledgeBase, + Description: "在知识库中搜索相关的安全知识。当你需要了解特定漏洞类型、攻击技术、检测方法等安全知识时，可以使用此工具进行检索。工具基于向量嵌入与余弦相似度检索（与 Eino retriever 语义一致）。建议：在搜索前可以先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型，然后使用正确的 risk_type 参数进行精确搜索，这样可以大幅减少检索时间。", + ShortDescription: "搜索知识库中的安全知识（向量语义检索）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "搜索查询内容，描述你想要了解的安全知识主题", + }, + "risk_type": map[string]interface{}{ + "type": "string", + "description": "可选：指定风险类型（如：SQL注入、XSS、文件上传等）。建议先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型列表，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间。如果不指定则搜索所有类型。", + }, + }, + "required": []string{"query"}, + }, + } + + searchHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + query, ok := args["query"].(string) + if !ok || query == "" { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: 查询参数不能为空", + }, + }, + IsError: true, + }, nil + } + + riskType := "" + if rt, ok := args["risk_type"].(string); ok && rt != "" { + riskType = rt + } + + logger.Info("执行知识库检索", + zap.String("query", query), + zap.String("riskType", riskType), + ) + + // 检索统一走 Retriever.Search → VectorEinoRetriever（Eino retriever 语义）。 + searchReq := &SearchRequest{ + Query: query, + RiskType: riskType, + TopK: 5, + } + + results, err := retriever.Search(ctx, searchReq) + if err != nil { + logger.Error("知识库检索失败", zap.Error(err)) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("检索失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + if len(results) == 0 { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("未找到与查询 '%s' 相关的知识。建议：\n1. 尝试使用不同的关键词\n2. 检查风险类型是否正确\n3. 确认知识库中是否包含相关内容", query), + }, + }, + }, nil + } + + // 格式化结果 + var resultText strings.Builder + + // 按余弦相似度（Score）降序 + sort.Slice(results, func(i, j int) bool { + return results[i].Score > results[j].Score + }) + + // 按文档分组结果，以便更好地展示上下文 + type itemGroup struct { + itemID string + results []*RetrievalResult + maxScore float64 // 该文档块的最高相似度 + } + itemGroups := make([]*itemGroup, 0) + itemMap := make(map[string]*itemGroup) + + for _, result := range results { + itemID := result.Item.ID + group, exists := itemMap[itemID] + if !exists { + group = &itemGroup{ + itemID: itemID, + results: make([]*RetrievalResult, 0), + maxScore: result.Score, + } + itemMap[itemID] = group + itemGroups = append(itemGroups, group) + } + group.results = append(group.results, result) + if result.Score > group.maxScore { + group.maxScore = result.Score + } + } + + // 按文档内最高相似度排序 + sort.Slice(itemGroups, func(i, j int) bool { + return itemGroups[i].maxScore > itemGroups[j].maxScore + }) + + // 收集检索到的知识项ID（用于日志） + retrievedItemIDs := make([]string, 0, len(itemGroups)) + + resultText.WriteString(fmt.Sprintf("找到 %d 条相关知识片段：\n\n", len(results))) + + resultIndex := 1 + for _, group := range itemGroups { + itemResults := group.results + mainResult := itemResults[0] + maxScore := mainResult.Score + for _, result := range itemResults { + if result.Score > maxScore { + maxScore = result.Score + mainResult = result + } + } + + // 按chunk_index排序，保证阅读的逻辑顺序（文档的原始顺序） + sort.Slice(itemResults, func(i, j int) bool { + return itemResults[i].Chunk.ChunkIndex < itemResults[j].Chunk.ChunkIndex + }) + + resultText.WriteString(fmt.Sprintf("--- 结果 %d (相似度: %.2f%%) ---\n", + resultIndex, mainResult.Similarity*100)) + resultText.WriteString(fmt.Sprintf("来源: [%s] %s (ID: %s)\n", mainResult.Item.Category, mainResult.Item.Title, mainResult.Item.ID)) + + // 按逻辑顺序显示所有chunk（包括主结果和扩展的chunk） + if len(itemResults) == 1 { + // 只有一个chunk，直接显示 + resultText.WriteString(fmt.Sprintf("内容片段:\n%s\n", mainResult.Chunk.ChunkText)) + } else { + // 多个chunk，按逻辑顺序显示 + resultText.WriteString("内容片段（按文档顺序）:\n") + for i, result := range itemResults { + // 标记主结果 + marker := "" + if result.Chunk.ID == mainResult.Chunk.ID { + marker = " [主匹配]" + } + resultText.WriteString(fmt.Sprintf(" [片段 %d%s]\n%s\n", i+1, marker, result.Chunk.ChunkText)) + } + } + resultText.WriteString("\n") + + if !contains(retrievedItemIDs, group.itemID) { + retrievedItemIDs = append(retrievedItemIDs, group.itemID) + } + resultIndex++ + } + + // 在结果末尾添加元数据（JSON格式，用于提取知识项ID） + // 使用特殊标记，避免影响AI阅读结果 + if len(retrievedItemIDs) > 0 { + metadataJSON, _ := json.Marshal(map[string]interface{}{ + "_metadata": map[string]interface{}{ + "retrievedItemIDs": retrievedItemIDs, + }, + }) + resultText.WriteString(fmt.Sprintf("\n", string(metadataJSON))) + } + + // 记录检索日志（异步，不阻塞） + // 注意：这里没有conversationID和messageID，需要在Agent层面记录 + // 实际的日志记录应该在Agent的progressCallback中完成 + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: resultText.String(), + }, + }, + }, nil + } + + mcpServer.RegisterTool(searchTool, searchHandler) + logger.Debug("知识检索工具已注册", zap.String("toolName", searchTool.Name)) +} + +// contains 检查切片是否包含元素 +func contains(slice []string, item string) bool { + for _, s := range slice { + if s == item { + return true + } + } + return false +} + +// GetRetrievalMetadata 从工具调用中提取检索元数据（用于日志记录） +func GetRetrievalMetadata(args map[string]interface{}) (query string, riskType string) { + if q, ok := args["query"].(string); ok { + query = q + } + if rt, ok := args["risk_type"].(string); ok { + riskType = rt + } + return +} + +// FormatRetrievalResults 格式化检索结果为字符串（用于日志） +func FormatRetrievalResults(results []*RetrievalResult) string { + if len(results) == 0 { + return "未找到相关结果" + } + + var builder strings.Builder + builder.WriteString(fmt.Sprintf("检索到 %d 条结果:\n", len(results))) + + itemIDs := make(map[string]bool) + for i, result := range results { + builder.WriteString(fmt.Sprintf("%d. [%s] %s (相似度: %.2f%%)\n", + i+1, result.Item.Category, result.Item.Title, result.Similarity*100)) + itemIDs[result.Item.ID] = true + } + + // 返回知识项ID列表（JSON格式） + ids := make([]string, 0, len(itemIDs)) + for id := range itemIDs { + ids = append(ids, id) + } + idsJSON, _ := json.Marshal(ids) + builder.WriteString(fmt.Sprintf("\n检索到的知识项ID: %s", string(idsJSON))) + + return builder.String() +} diff --git a/internal/knowledge/types.go b/internal/knowledge/types.go new file mode 100644 index 00000000..42e35e76 --- /dev/null +++ b/internal/knowledge/types.go @@ -0,0 +1,123 @@ +package knowledge + +import ( + "encoding/json" + "time" +) + +// formatTime 格式化时间为 RFC3339 格式，零时间返回空字符串 +func formatTime(t time.Time) string { + if t.IsZero() { + return "" + } + return t.Format(time.RFC3339) +} + +// KnowledgeItem 知识库项 +type KnowledgeItem struct { + ID string `json:"id"` + Category string `json:"category"` // 风险类型（文件夹名） + Title string `json:"title"` // 标题（文件名） + FilePath string `json:"filePath"` // 文件路径 + Content string `json:"content"` // 文件内容 + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// KnowledgeItemSummary 知识库项摘要（用于列表，不包含完整内容） +type KnowledgeItemSummary struct { + ID string `json:"id"` + Category string `json:"category"` + Title string `json:"title"` + FilePath string `json:"filePath"` + Content string `json:"content,omitempty"` // 可选：内容预览（如果提供，通常只包含前 150 字符） + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (k *KnowledgeItemSummary) MarshalJSON() ([]byte, error) { + type Alias KnowledgeItemSummary + aux := &struct { + *Alias + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + }{ + Alias: (*Alias)(k), + } + aux.CreatedAt = formatTime(k.CreatedAt) + aux.UpdatedAt = formatTime(k.UpdatedAt) + return json.Marshal(aux) +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (k *KnowledgeItem) MarshalJSON() ([]byte, error) { + type Alias KnowledgeItem + aux := &struct { + *Alias + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + }{ + Alias: (*Alias)(k), + } + aux.CreatedAt = formatTime(k.CreatedAt) + aux.UpdatedAt = formatTime(k.UpdatedAt) + return json.Marshal(aux) +} + +// KnowledgeChunk 知识块（用于向量化） +type KnowledgeChunk struct { + ID string `json:"id"` + ItemID string `json:"itemId"` + ChunkIndex int `json:"chunkIndex"` + ChunkText string `json:"chunkText"` + Embedding []float32 `json:"-"` // 向量嵌入，不序列化到 JSON + CreatedAt time.Time `json:"createdAt"` +} + +// RetrievalResult 检索结果 +type RetrievalResult struct { + Chunk *KnowledgeChunk `json:"chunk"` + Item *KnowledgeItem `json:"item"` + Similarity float64 `json:"similarity"` // 相似度分数 + Score float64 `json:"score"` // 与 Similarity 相同：余弦相似度 +} + +// RetrievalLog 检索日志 +type RetrievalLog struct { + ID string `json:"id"` + ConversationID string `json:"conversationId,omitempty"` + MessageID string `json:"messageId,omitempty"` + Query string `json:"query"` + RiskType string `json:"riskType,omitempty"` + RetrievedItems []string `json:"retrievedItems"` // 检索到的知识项 ID 列表 + CreatedAt time.Time `json:"createdAt"` +} + +// MarshalJSON 自定义 JSON 序列化，确保时间格式正确 +func (r *RetrievalLog) MarshalJSON() ([]byte, error) { + type Alias RetrievalLog + return json.Marshal(&struct { + *Alias + CreatedAt string `json:"createdAt"` + }{ + Alias: (*Alias)(r), + CreatedAt: formatTime(r.CreatedAt), + }) +} + +// CategoryWithItems 分类及其下的知识项（用于按分类分页） +type CategoryWithItems struct { + Category string `json:"category"` // 分类名称 + ItemCount int `json:"itemCount"` // 该分类下的知识项总数 + Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表 +} + +// SearchRequest 搜索请求 +type SearchRequest struct { + Query string `json:"query"` + RiskType string `json:"riskType,omitempty"` // 可选：指定风险类型 + SubIndexFilter string `json:"subIndexFilter,omitempty"` // 可选：仅保留 sub_indexes 含该标签的行（含未打标旧数据） + TopK int `json:"topK,omitempty"` // 返回 Top-K 结果，默认 5 + Threshold float64 `json:"threshold,omitempty"` // 相似度阈值，默认 0.7 +} diff --git a/internal/knowledge/wire_retriever.go b/internal/knowledge/wire_retriever.go new file mode 100644 index 00000000..142ee427 --- /dev/null +++ b/internal/knowledge/wire_retriever.go @@ -0,0 +1,76 @@ +package knowledge + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/openai" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/flow/retriever/multiquery" + "go.uber.org/zap" +) + +// WireRetrieverPipeline builds Eino MultiQuery + HTTP rerank + post-process pipeline on r. +// Call once after NewRetriever; UpdateConfig re-invokes when wireOpenAI is set. +func WireRetrieverPipeline(ctx context.Context, r *Retriever, openAI *config.OpenAIConfig) error { + if r == nil { + return fmt.Errorf("retriever is nil") + } + if openAI == nil { + return fmt.Errorf("openai config is nil") + } + if r.config == nil { + return fmt.Errorf("retrieval config is nil") + } + r.wireOpenAI = openAI + + httpClient := openai.NewEinoHTTPClient(openAI, &http.Client{Timeout: 120 * time.Second}) + maxCompletionTokens := openAI.MaxCompletionTokensEffective() + chatCfg := &einoopenai.ChatModelConfig{ + APIKey: strings.TrimSpace(openAI.APIKey), + BaseURL: strings.TrimSuffix(strings.TrimSpace(openAI.BaseURL), "/"), + Model: strings.TrimSpace(openAI.Model), + HTTPClient: httpClient, + MaxCompletionTokens: &maxCompletionTokens, + } + if chatCfg.Model == "" { + chatCfg.Model = "gpt-4o" + } + rewriteLLM, err := einoopenai.NewChatModel(ctx, chatCfg) + if err != nil { + return fmt.Errorf("multi_query rewrite model: %w", err) + } + + reranker, err := NewHTTPReranker(&r.config.Rerank, openAI, r.logger) + if err != nil { + return fmt.Errorf("reranker: %w", err) + } + r.SetDocumentReranker(reranker) + + vec := NewVectorEinoRetriever(r) + mq, err := multiquery.NewRetriever(ctx, &multiquery.Config{ + RewriteLLM: rewriteLLM, + MaxQueriesNum: r.config.MultiQuery.MaxQueriesEffective(), + OrigRetriever: vec, + }) + if err != nil { + return fmt.Errorf("multi_query: %w", err) + } + + r.pipeline = newKnowledgePipelineRetriever(mq, r) + if r.logger != nil { + provider := r.config.Rerank.ProviderEffective(strings.TrimSpace(openAI.BaseURL)) + r.logger.Info("知识库检索流水线已启用", + zap.String("pipeline", "MultiQuery→Vector→Rerank→PostRetrieve"), + zap.Int("multi_query_max", r.config.MultiQuery.MaxQueriesEffective()), + zap.String("rerank_provider", provider), + zap.String("rerank_model", r.config.Rerank.ModelEffective(provider)), + ) + } + return nil +} diff --git a/internal/projectprompt/blackboard.go b/internal/projectprompt/blackboard.go new file mode 100644 index 00000000..d3e3ae76 --- /dev/null +++ b/internal/projectprompt/blackboard.go @@ -0,0 +1,132 @@ +// Package projectprompt 提供项目黑板相关的系统提示文本（纯字符串，无 database 依赖）。 +// 供 agent / multiagent 等包引用，避免 agent → project 导入环导致 gopls 元数据失败。 +package projectprompt + +import ( + "strings" + + "cyberstrike-ai/internal/mcp/builtin" +) + +const ( + factRhythmCore = "勿等会话结束或收尾再批量写入。每**确认**一条新认知（开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化）后，**立即**调用 `upsert_project_fact`（同 fact_key 覆盖更新）。每**验证**出一条可复现漏洞（含 POC/影响）后，**立即**调用 `record_vulnerability`；与事实可各记一次。继续下一步工作前优先落库，避免上下文压缩后细节丢失。未绑项目时说明无法写黑板，仍在本轮保留证据摘要。" + factRhythmCoordinatorSuffix = "委派/子任务返回新认知或漏洞时，由协调者及时写入，勿假定子代理已记。" + factRhythmSubAgentSuffix = "若工具集中无上述工具，须在交付物末尾给出「待落库」结构化条目（fact_key 建议、summary、body/POC 要点），供协调者**立即**写入。" +) + +// FactRecordingIncrementalRhythmMarkdown 返回边渗透边记录节奏（Markdown，供 agents/*.md 与文档对齐）。 +func FactRecordingIncrementalRhythmMarkdown(coordinator, subAgent bool) string { + var b strings.Builder + b.WriteString("- **边渗透边记录（强制节奏）**：") + b.WriteString(factRhythmCore) + if coordinator { + b.WriteString(factRhythmCoordinatorSuffix) + } + if subAgent { + b.WriteString(factRhythmSubAgentSuffix) + } + return b.String() +} + +func factRecordingIncrementalRhythmBuiltin(coordinator, subAgent bool) string { + var b strings.Builder + b.WriteString("- **边渗透边记录（强制节奏）**：勿等会话结束或收尾再批量写入。每**确认**一条新认知（开放端口/服务版本、入口路径、认证态或凭据特征、可利用点或攻击面变化）后，**立即**调用 ") + b.WriteString(builtin.ToolUpsertProjectFact) + b.WriteString("（同 fact_key 覆盖更新）。每**验证**出一条可复现漏洞（含 POC/影响）后，**立即**调用 ") + b.WriteString(builtin.ToolRecordVulnerability) + b.WriteString("；与事实可各记一次。继续下一步工作前优先落库，避免上下文压缩后细节丢失。未绑项目时说明无法写黑板，仍在本轮保留证据摘要。") + if coordinator { + b.WriteString(factRhythmCoordinatorSuffix) + } + if subAgent { + b.WriteString(factRhythmSubAgentSuffix) + } + return b.String() +} + +func factEdgeRecordingGuidance() string { + return `### 事实关系边（links） + +- 写入 **finding / chain / exploit / poc** 时，**必须**在 ` + "`upsert_project_fact`" + ` 中提供 ` + "`links`" + `（**推荐 ` + "`from`" + `**：来源 fact 指向当前 fact，即 ` + "`from`" + ` → 当前 ` + "`fact_key`" + `）。 +- **最少要求**：finding 类至少 1 条 from=target/* + type=discovered_on（即 target → finding）；在 finding 上记录 exploit 用 from=exploit/* + type=exploits（即 exploit → finding）。 +- **常用 type**：` + "`discovered_on`" + `（发现在哪）、` + "`depends_on`" + `（复现前置）、` + "`leads_to`" + `（认知推进）、` + "`enables`" + `（扩大攻击面）、` + "`exploits`" + `（利用关系）、` + "`contains`" + `（资产包含）、` + "`part_of`" + `（属于链/组）、` + "`supports`" + `（证据支撑）。 +- 更新时：**省略 links 保留已有边**；传入 links 则**替换**全部关系边（from → 当前 fact）。 +- body 中「依赖事实」段落可与 links 并存（人读）；结构化关系以 links 为准。` +} + +func factRecordingGuidanceBlock() string { + return `### 事实写入规范（审计复现 / 知识沉淀） + +- **summary**：索引用一行，须含「什么 + 在哪 + 如何触发/验证」要点，禁止只写结论（如仅写「存在 SQLi」）。 +- **body**：完整可复现上下文，写入 ` + "`upsert_project_fact`" + ` 的 body 字段；索引不含 body，后续会话须靠 ` + "`get_project_fact`" + ` 取回。 +- **category / fact_key 建议**： + - 环境认知：` + "`target/`" + `、` + "`auth/`" + `、` + "`infra/`" + `、` + "`business/`" + `（body 用环境模板即可） + - 发现与利用：` + "`finding/`" + `、` + "`chain/`" + `、` + "`exploit/`" + `、` + "`poc/`" + `（**必须**用攻击链模板填满 body：入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID） +- **与漏洞记录分工**：` + "`record_vulnerability`" + ` 记可交付 findings；事实记**复现所需的全部上下文**（含失败尝试、绕过、依赖会话），二者可各记一次。 +- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入，勿散落多个 key 导致上下文丢失。` +} + +// FactRecordingBlackboardSection 项目黑板与漏洞记录的完整系统提示块（单/多 Agent 主代理共用）。 +func FactRecordingBlackboardSection(coordinatorDelegate bool) string { + var b strings.Builder + b.WriteString("## 项目黑板（事实）与漏洞记录（分离）\n\n") + b.WriteString("当前对话若已绑定项目，系统会自动注入「项目黑板索引」（仅 fact_key + 摘要）。**摘要不足时必须调用 ") + b.WriteString(builtin.ToolGetProjectFact) + b.WriteString("(fact_key) 获取 body，禁止凭摘要臆造细节。**\n\n") + b.WriteString(factRecordingIncrementalRhythmBuiltin(coordinatorDelegate, false)) + b.WriteString("\n\n") + b.WriteString("- **环境/目标/认证等认知**（非正式漏洞条目）：使用 ") + b.WriteString(builtin.ToolUpsertProjectFact) + b.WriteString("，fact_key 建议 `category/slug`（如 target/primary_domain），同 key 覆盖更新；body 记端口/版本/凭据特征与证据来源。\n") + b.WriteString("- **发现与利用上下文**（审计复现）：fact_key 建议 finding/、chain/、exploit/、poc/ 前缀；**body 必填**完整攻击链（入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 related_vulnerability_id），**禁止仅写结论**；summary 写「什么 + 在哪 + 如何验证」一行要点。\n") + b.WriteString("- **可交付漏洞**：使用 ") + b.WriteString(builtin.ToolRecordVulnerability) + b.WriteString("，含标题、严重程度、类型、目标、证明（POC）、影响、修复建议。记前可先 ") + b.WriteString(builtin.ToolListVulnerabilities) + b.WriteString(" 查重，详情用 ") + b.WriteString(builtin.ToolGetVulnerability) + b.WriteString("(id)（默认仅当前项目/会话）。\n") + b.WriteString("- 同一发现可能需**各记一次**（事实记**完整攻击链与 exploit 细节**供复现，漏洞记正式 findings）。误报用 ") + b.WriteString(builtin.ToolDeprecateProjectFact) + b.WriteString(" 或漏洞状态 false_positive。\n") + b.WriteString("- 事实多时用 ") + b.WriteString(builtin.ToolListProjectFacts) + b.WriteString(" / ") + b.WriteString(builtin.ToolSearchProjectFacts) + b.WriteString(" 检索。\n\n") + b.WriteString(factEdgeRecordingGuidance()) + b.WriteString("\n\n") + b.WriteString(factRecordingGuidanceBlock()) + b.WriteString("\n\n严重程度：critical / high / medium / low / info。证明须含足够证据（请求响应、截图、命令输出等）。") + return b.String() +} + +// FactRecordingSubAgentSection 子代理边渗透边记录（无工具时输出待落库条目）。 +func FactRecordingSubAgentSection() string { + return "## 边渗透边记录\n\n" + factRecordingIncrementalRhythmBuiltin(false, true) + "\n" +} + +// FactRecordingBlackboardSectionMarkdown 与 FactRecordingBlackboardSection 等价的 Markdown（工具名为字面量，供 agents/*.md）。 +func FactRecordingBlackboardSectionMarkdown(coordinatorDelegate bool) string { + var b strings.Builder + b.WriteString("## 项目黑板（事实）与漏洞记录（分离）\n\n") + b.WriteString("当前对话若已绑定项目，系统会自动注入「项目黑板索引」（仅 `fact_key` + 摘要）。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body，禁止凭摘要臆造细节。**\n\n") + b.WriteString(FactRecordingIncrementalRhythmMarkdown(coordinatorDelegate, false)) + b.WriteString("\n\n") + b.WriteString("- **环境/目标/认证等认知**（非正式漏洞）：使用 **`upsert_project_fact`**，`fact_key` 建议 `category/slug`（如 `target/primary_domain`），同 key 覆盖更新；body 记端口/版本/凭据特征与证据来源。\n") + b.WriteString("- **发现与利用上下文**（审计复现）：`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀；**body 必填**完整攻击链（入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`），**禁止仅写结论**；summary 写「什么 + 在哪 + 如何验证」一行要点。\n") + b.WriteString("- **可交付漏洞**：使用 **`record_vulnerability`**（标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议）。严重程度 critical / high / medium / low / info。\n") + b.WriteString("- 同一发现可能需**各记一次**（事实记可复现攻击链，漏洞记正式 findings）。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。\n") + b.WriteString("- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。\n\n") + b.WriteString(factEdgeRecordingGuidance()) + b.WriteString("\n\n") + b.WriteString(factRecordingGuidanceBlock()) + b.WriteString("\n\n严重程度：critical / high / medium / low / info。证明须含足够证据（请求响应、截图、命令输出等）。") + return b.String() +} + +// FactEdgeRecordingGuidance 写入边时的 Agent 规范（供 project 包复用）。 +func FactEdgeRecordingGuidance() string { return factEdgeRecordingGuidance() } + +// FactRecordingGuidanceBlock 事实写入规范块（供 project 包复用）。 +func FactRecordingGuidanceBlock() string { return factRecordingGuidanceBlock() } diff --git a/internal/projectprompt/shell_tools.go b/internal/projectprompt/shell_tools.go new file mode 100644 index 00000000..6314a181 --- /dev/null +++ b/internal/projectprompt/shell_tools.go @@ -0,0 +1,11 @@ +package projectprompt + +// ShellExecExecuteGuidanceSection 供单代理/多代理系统提示追加：exec 与 execute 分工（尽量短）。 +func ShellExecExecuteGuidanceSection() string { + return `Shell（exec/execute）：有专用 MCP 工具时优先专用工具；系统命令（管道、workdir、后台 &）用 exec；skills/ 内脚本（配合 read_file、skill）用 execute；多步扫描分拆调用，禁止一条 shell 串多个扫描器。长脚本、请求体或 Payload 必须先用 write_file 写入会话工作目录，再用 exec/execute 执行短命令；禁止把长内容嵌入 command。下载/临时文件须写入系统提示中的「会话工作目录」，禁止用 /tmp。` +} + +// ShellExecExecuteGuidanceReconSuffix 侦察子代理可选追加（一行）。 +func ShellExecExecuteGuidanceReconSuffix() string { + return `枚举优先 subfinder、amass 等专用 MCP，勿 exec/execute 拼长链。` +} diff --git a/internal/reasoning/eino.go b/internal/reasoning/eino.go new file mode 100644 index 00000000..f88741bd --- /dev/null +++ b/internal/reasoning/eino.go @@ -0,0 +1,428 @@ +// Package reasoning maps user/config intent to CloudWeGo Eino OpenAI ChatModel fields +// (ReasoningEffort, ExtraFields such as thinking / reasoning_effort / output_config). +package reasoning + +import ( + "strings" + + "cyberstrike-ai/internal/config" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" +) + +// ClientIntent is optional per-request override from ChatRequest.reasoning. +type ClientIntent struct { + Mode string + Effort string +} + +type wireProfile int + +const ( + wireNone wireProfile = iota + wireClaude + wireDeepseek + wireOpenAI + wireOutputConfig +) + +// ApplyPlanExecutePlannerModelConfig configures the plan_execute planner/replanner +// ChatModel. Those Eino agents call WithToolChoice(Forced); several gateways reject +// thinking / reasoning fields on the same request (tool_choice required/object). +// Executor should keep the normal ApplyToEinoChatModelConfig path. +func ApplyPlanExecutePlannerModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig) { + if cfg == nil || oa == nil { + return + } + mergeExtraRequestFields(cfg, oa.Reasoning.ExtraRequestFields) + clearReasoningFromChatModelConfig(cfg) + if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + // DeepSeek enables thinking by default, so omission would not actually + // disable it for the planner's forced tool-choice requests. + applyThinkingDisabled(cfg) + } +} + +func clearReasoningFromChatModelConfig(cfg *einoopenai.ChatModelConfig) { + if cfg == nil { + return + } + cfg.ReasoningEffort = "" + if cfg.ExtraFields != nil { + for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} { + delete(cfg.ExtraFields, key) + } + if len(cfg.ExtraFields) == 0 { + cfg.ExtraFields = nil + } + } +} + +func mergeExtraRequestFields(cfg *einoopenai.ChatModelConfig, fields map[string]interface{}) { + if cfg == nil || len(fields) == 0 { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any, len(fields)) + } + for k, v := range fields { + cfg.ExtraFields[k] = v + } +} + +// ApplyToEinoChatModelConfig merges reasoning-related options into cfg. +// Precondition: cfg already has APIKey, BaseURL, Model, HTTPClient set. +func ApplyToEinoChatModelConfig(cfg *einoopenai.ChatModelConfig, oa *config.OpenAIConfig, client *ClientIntent) { + if cfg == nil || oa == nil { + return + } + sr := &oa.Reasoning + allowClient := sr.AllowClientReasoningEffective() + mode := effectiveMode(sr, client, allowClient) + + // Admin-defined root fields are independent of the selected reasoning wire + // profile. Merge them first so mode=off can remove only reasoning controls + // while preserving unrelated gateway options. + mergeExtraRequestFields(cfg, sr.ExtraRequestFields) + if mode == "off" { + clearReasoningFromChatModelConfig(cfg) + // Strict OpenAI endpoints reject unknown `thinking` fields, whereas the + // DeepSeek API enables thinking by default and requires an explicit + // thinking.type=disabled switch. Detect the actual DeepSeek target even + // when the configured reasoning profile was left as openai_compat. + if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + applyThinkingDisabled(cfg) + } + return + } + + // Claude (Anthropic): merge admin extras first; optional extended thinking maps to top-level `thinking` + // (see internal/openai convertOpenAIToClaude). DeepSeek/OpenAI-style fields are not sent. + if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") || + strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") { + applyClaudeExtendedThinking(cfg, mode, effectiveEffort(sr, client, allowClient), oa.Model) + return + } + + effort := effectiveEffort(sr, client, allowClient) + prof := resolveWireProfile(oa, sr) + + switch prof { + case wireClaude, wireNone: + return + case wireDeepseek: + applyDeepseek(cfg, mode, effort) + case wireOutputConfig: + applyOutputConfigEffort(cfg, mode, effort) + default: // wireOpenAI + applyOpenAICompat(cfg, mode, effort) + } +} + +// AgenticOpenAIExtraFields returns reasoning-related request fields for +// agenticopenai.ChatConfig. The agentic chat backend currently exposes provider +// extensions through ExtraFields instead of typed ReasoningEffort fields. +func AgenticOpenAIExtraFields(oa *config.OpenAIConfig, client *ClientIntent) map[string]any { + if oa == nil { + return nil + } + sr := &oa.Reasoning + allowClient := sr.AllowClientReasoningEffective() + mode := effectiveMode(sr, client, allowClient) + fields := cloneExtraRequestFields(sr.ExtraRequestFields) + if mode == "off" { + clearReasoningExtraFields(fields) + if resolveWireProfile(oa, sr) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "disabled"} + } + return fields + } + if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") || + strings.EqualFold(strings.TrimSpace(oa.Provider), "anthropic") { + return fields + } + effort := effectiveEffort(sr, client, allowClient) + switch resolveWireProfile(oa, sr) { + case wireDeepseek: + if mode == "auto" || mode == "on" { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "enabled"} + } + if effort != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["reasoning_effort"] = effortStringForAPI(effort) + } + case wireOutputConfig: + e := effort + if mode == "on" && e == "" { + e = "high" + } + if e != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["output_config"] = map[string]any{"effort": effortStringForAPI(e)} + } + default: + e := effort + if mode == "on" && e == "" { + e = "medium" + } + if e != "" { + if fields == nil { + fields = make(map[string]any) + } + fields["reasoning_effort"] = effortStringForAPI(e) + } + } + return fields +} + +// AgenticOpenAIPlannerExtraFields mirrors ApplyPlanExecutePlannerModelConfig for +// agenticopenai.ChatConfig: keep admin extras, strip reasoning controls, and +// explicitly disable DeepSeek thinking where omission would still think. +func AgenticOpenAIPlannerExtraFields(oa *config.OpenAIConfig) map[string]any { + if oa == nil { + return nil + } + fields := cloneExtraRequestFields(oa.Reasoning.ExtraRequestFields) + clearReasoningExtraFields(fields) + if resolveWireProfile(oa, &oa.Reasoning) == wireDeepseek || oa.IsDeepSeekEndpointOrModel() { + if fields == nil { + fields = make(map[string]any) + } + fields["thinking"] = map[string]any{"type": "disabled"} + } + return fields +} + +func cloneExtraRequestFields(fields map[string]interface{}) map[string]any { + if len(fields) == 0 { + return nil + } + out := make(map[string]any, len(fields)) + for k, v := range fields { + out[k] = v + } + return out +} + +func clearReasoningExtraFields(fields map[string]any) { + for _, key := range []string{"thinking", "reasoning_effort", "output_config", "reasoning"} { + delete(fields, key) + } +} + +// applyClaudeExtendedThinking sets Anthropic Messages API fields per official guidance: +// - Adaptive models (4.6+): thinking.type=adaptive; output_config.effort only when user sets effort (API default is high). +// - Sonnet 3.7: thinking.type=enabled + budget_tokens=10000 (doc example); effort is not mapped — use extra_request_fields for custom budget. +func applyClaudeExtendedThinking(cfg *einoopenai.ChatModelConfig, mode, effort, model string) { + if cfg == nil || mode == "off" { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + m := strings.ToLower(strings.TrimSpace(model)) + sonnet37 := isClaudeSonnet37(m) + + if _, exists := cfg.ExtraFields["thinking"]; !exists { + cfg.ExtraFields["thinking"] = claudeThinkingForModel(m, sonnet37) + } + + applyClaudeOutputConfigEffort(cfg, effort, sonnet37) +} + +// claudeSonnet37DefaultBudgetTokens matches Anthropic extended-thinking documentation examples (budget_tokens with max_tokens 16000). +const claudeSonnet37DefaultBudgetTokens = 10000 + +func isClaudeSonnet37(m string) bool { + return strings.Contains(m, "claude-3-7-sonnet") || + strings.Contains(m, "3-7-sonnet") || + strings.Contains(m, "sonnet-3.7") +} + +func claudeThinkingForModel(m string, sonnet37 bool) map[string]any { + if sonnet37 { + return map[string]any{ + "type": "enabled", + "budget_tokens": claudeSonnet37DefaultBudgetTokens, + "display": "summarized", + } + } + // Opus 4.7+: manual enabled+budget rejected — adaptive only. + if strings.Contains(m, "opus-4-7") || strings.Contains(m, "opus-4.7") { + return map[string]any{ + "type": "adaptive", + "display": "summarized", + } + } + return map[string]any{ + "type": "adaptive", + "display": "summarized", + } +} + +// applyClaudeOutputConfigEffort sets top-level output_config.effort only when effort is explicitly configured. +// Omitted effort uses the API default (high); do not inject effort on mode:on alone. +func applyClaudeOutputConfigEffort(cfg *einoopenai.ChatModelConfig, effort string, sonnet37 bool) { + if cfg == nil || sonnet37 { + return + } + if _, exists := cfg.ExtraFields["output_config"]; exists { + return + } + e := effortStringForAPI(effort) + if e == "" { + return + } + cfg.ExtraFields["output_config"] = map[string]any{"effort": e} +} + +func effectiveMode(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string { + server := strings.ToLower(strings.TrimSpace(sr.ModeEffective())) + if server == "" || server == "default" { + server = "auto" + } + if !allowClient || client == nil { + return server + } + cm := strings.ToLower(strings.TrimSpace(client.Mode)) + if cm == "" || cm == "default" { + return server + } + return cm +} + +func effectiveEffort(sr *config.OpenAIReasoningConfig, client *ClientIntent, allowClient bool) string { + se := normalizeEffort(sr.Effort) + if !allowClient || client == nil { + return se + } + ce := normalizeEffort(client.Effort) + if ce != "" { + return ce + } + return se +} + +func normalizeEffort(s string) string { + e := strings.ToLower(strings.TrimSpace(s)) + switch e { + case "low", "medium", "high", "max", "xhigh": + return e + default: + return "" + } +} + +// usesExtraFieldsReasoningEffort 为 Eino 无枚举的最高档 effort，经 ExtraFields 原样下发（max / xhigh 由网关自行识别，不做互转）。 +func usesExtraFieldsReasoningEffort(e string) bool { + return e == "max" || e == "xhigh" +} + +func resolveWireProfile(oa *config.OpenAIConfig, sr *config.OpenAIReasoningConfig) wireProfile { + provider := strings.TrimSpace(oa.Provider) + if strings.EqualFold(provider, "claude") || strings.EqualFold(provider, "anthropic") { + return wireClaude + } + p := strings.ToLower(strings.TrimSpace(sr.ProfileEffective())) + switch p { + case "output_config", "output_config_effort": + return wireOutputConfig + case "openai", "openai_compat": + return wireOpenAI + case "deepseek", "deepseek_compat": + return wireDeepseek + case "auto", "": + if oa.IsDeepSeekEndpointOrModel() { + return wireDeepseek + } + return wireOpenAI + default: + return wireOpenAI + } +} + +func applyThinkingDisabled(cfg *einoopenai.ChatModelConfig) { + if cfg == nil { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["thinking"] = map[string]any{"type": "disabled"} +} + +func applyDeepseek(cfg *einoopenai.ChatModelConfig, mode, effort string) { + // auto: enable thinking for DeepSeek line; on: same; auto without effort still opens thinking. + if mode == "auto" || mode == "on" { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["thinking"] = map[string]any{"type": "enabled"} + } + if effort != "" { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(effort) + } +} + +func applyOpenAICompat(cfg *einoopenai.ChatModelConfig, mode, effort string) { + if mode == "auto" && effort == "" { + return + } + e := effort + if mode == "on" && e == "" { + e = "medium" + } + if e == "" { + return + } + if usesExtraFieldsReasoningEffort(e) { + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(e) + return + } + switch e { + case "low": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelLow + case "medium": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelMedium + case "high": + cfg.ReasoningEffort = einoopenai.ReasoningEffortLevelHigh + } +} + +func applyOutputConfigEffort(cfg *einoopenai.ChatModelConfig, mode, effort string) { + if mode == "auto" && effort == "" { + return + } + e := effort + if mode == "on" && e == "" { + e = "high" + } + if e == "" { + return + } + if cfg.ExtraFields == nil { + cfg.ExtraFields = make(map[string]any) + } + cfg.ExtraFields["output_config"] = map[string]any{"effort": effortStringForAPI(e)} +} + +func effortStringForAPI(e string) string { + // 原样透传：OpenAI 官方多为 xhigh，部分兼容网关为 max，由配置/对话 effort 选择。 + return strings.ToLower(strings.TrimSpace(e)) +} diff --git a/internal/reasoning/eino_test.go b/internal/reasoning/eino_test.go new file mode 100644 index 00000000..c9eadf8a --- /dev/null +++ b/internal/reasoning/eino_test.go @@ -0,0 +1,424 @@ +package reasoning + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/config" + + einoopenai "github.com/cloudwego/eino-ext/components/model/openai" + "github.com/cloudwego/eino/schema" +) + +var reasoningPayloadKeysForTest = []string{"thinking", "reasoning_effort", "output_config", "reasoning"} + +func assertNoReasoningFields(t *testing.T, cfg *einoopenai.ChatModelConfig) { + t.Helper() + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + for _, key := range reasoningPayloadKeysForTest { + if _, ok := cfg.ExtraFields[key]; ok { + t.Fatalf("expected %q omitted, got %#v", key, cfg.ExtraFields) + } + } +} + +func TestEffortStringForAPI_passthrough(t *testing.T) { + cases := map[string]string{ + "max": "max", + "xhigh": "xhigh", + "HIGH": "high", + "Medium": "medium", + } + for in, want := range cases { + if got := effortStringForAPI(in); got != want { + t.Fatalf("%q -> %q, want %q", in, got, want) + } + } +} + +func TestNormalizeEffort_maxAndXhigh(t *testing.T) { + if normalizeEffort("xhigh") != "xhigh" { + t.Fatal("xhigh not accepted") + } + if normalizeEffort("max") != "max" { + t.Fatal("max not accepted") + } +} + +func TestApplyOpenAICompat_xhighExtraField(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "xhigh", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if cfg.ExtraFields == nil { + t.Fatal("expected ExtraFields") + } + if got, _ := cfg.ExtraFields["reasoning_effort"].(string); got != "xhigh" { + t.Fatalf("reasoning_effort=%q", got) + } +} + +func TestAgenticOpenAIExtraFields_openAICompatReasoningEffort(t *testing.T) { + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIExtraFields(oa, nil) + if got["reasoning_effort"] != "high" { + t.Fatalf("reasoning_effort=%#v, want high in %#v", got["reasoning_effort"], got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIExtraFields_reasoningOffPreservesUnrelatedFields(t *testing.T) { + oa := &config.OpenAIConfig{ + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "off", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "thinking": map[string]any{"type": "enabled"}, + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIExtraFields(oa, nil) + for _, key := range reasoningPayloadKeysForTest { + if _, ok := got[key]; ok { + t.Fatalf("agentic fields unexpectedly contain %q: %#v", key, got) + } + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIPlannerExtraFields_deepseekDisablesThinking(t *testing.T) { + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com", + Model: "deepseek-chat", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "auto", + Mode: "on", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIPlannerExtraFields(oa) + if got["reasoning_effort"] != nil { + t.Fatalf("planner should strip reasoning_effort: %#v", got) + } + thinking, ok := got["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled, got %#v", got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestAgenticOpenAIPlannerExtraFields_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) { + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4-flash", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + ExtraRequestFields: map[string]interface{}{ + "reasoning_effort": "high", + "vendor_option": true, + }, + }, + } + got := AgenticOpenAIPlannerExtraFields(oa) + if _, ok := got["reasoning_effort"]; ok { + t.Fatalf("planner should strip reasoning_effort: %#v", got) + } + thinking, ok := got["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", got) + } + if got["vendor_option"] != true { + t.Fatalf("vendor option not preserved: %#v", got) + } +} + +func TestApplyPlanExecutePlannerModelConfig_stripsReasoningWhenGlobalOn(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://antchat.example.com/v1", + Model: "minimax-m3", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + }, + } + ApplyPlanExecutePlannerModelConfig(cfg, oa) + assertNoReasoningFields(t, cfg) + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyPlanExecutePlannerModelConfig_deepseekEndpointWinsOverOpenAIProfile(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com/v1", + Model: "deepseek-v4-flash", + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "high", + }, + } + ApplyPlanExecutePlannerModelConfig(cfg, oa) + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + if _, ok := cfg.ExtraFields["reasoning_effort"]; ok { + t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields) + } + thinking, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected deepseek thinking disabled despite openai_compat profile, got %#v", cfg.ExtraFields) + } + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyReasoningOff_omitsAllReasoningFields(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "thinking": map[string]any{"type": "enabled"}, + "output_config": map[string]any{"effort": "high"}, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", + Effort: "high", + Profile: "openai_compat", + ExtraRequestFields: map[string]interface{}{ + "thinking": map[string]any{"type": "disabled"}, + "reasoning": map[string]any{"effort": "high"}, + "vendor_option": true, + }, + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + assertNoReasoningFields(t, cfg) + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } +} + +func TestApplyReasoningOff_clientOverrideOmit(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", Effort: "high", Profile: "openai_compat", + }} + ApplyToEinoChatModelConfig(cfg, oa, &ClientIntent{Mode: "off", Effort: "high"}) + assertNoReasoningFields(t, cfg) +} + +func TestApplyReasoningOff_deepseekExplicitlyDisablesDefaultThinking(t *testing.T) { + for _, profile := range []string{"deepseek_compat", "auto", "openai_compat"} { + t.Run(profile, func(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{ExtraFields: map[string]any{ + "reasoning_effort": "high", + "vendor_option": true, + }} + oa := &config.OpenAIConfig{ + BaseURL: "https://api.deepseek.com", + Model: "deepseek-v4-pro", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", Effort: "high", Profile: profile, + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if cfg.ReasoningEffort != "" { + t.Fatalf("expected ReasoningEffort omitted, got %q", cfg.ReasoningEffort) + } + if _, ok := cfg.ExtraFields["reasoning_effort"]; ok { + t.Fatalf("expected reasoning_effort omitted, got %#v", cfg.ExtraFields) + } + thinking, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || thinking["type"] != "disabled" { + t.Fatalf("expected DeepSeek thinking disabled, got %#v", cfg.ExtraFields) + } + if cfg.ExtraFields["vendor_option"] != true { + t.Fatalf("expected unrelated extra field preserved, got %#v", cfg.ExtraFields) + } + }) + } +} + +func TestApplyReasoningOff_wirePayloadOmitsThinking(t *testing.T) { + var requestBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + if err := json.Unmarshal(body, &requestBody); err != nil { + t.Errorf("decode request body: %v; body=%s", err, body) + } + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"chatcmpl-test","object":"chat.completion","created":1,"model":"gpt-4o-mini","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`) + })) + defer srv.Close() + + cfg := &einoopenai.ChatModelConfig{ + APIKey: "test-key", + BaseURL: srv.URL, + Model: "gpt-4o-mini", + } + oa := &config.OpenAIConfig{ + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4o-mini", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "off", Effort: "high", Profile: "openai_compat", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + model, err := einoopenai.NewChatModel(context.Background(), cfg) + if err != nil { + t.Fatalf("new chat model: %v", err) + } + if _, err := model.Generate(context.Background(), []*schema.Message{schema.UserMessage("hello")}); err != nil { + t.Fatalf("generate: %v", err) + } + for _, key := range reasoningPayloadKeysForTest { + if _, ok := requestBody[key]; ok { + t.Fatalf("wire payload unexpectedly contains %q: %#v", key, requestBody) + } + } +} + +func TestApplyOpenAICompat_maxPassthrough(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Reasoning: config.OpenAIReasoningConfig{ + Profile: "openai_compat", + Mode: "on", + Effort: "max", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + got, _ := cfg.ExtraFields["reasoning_effort"].(string) + if got != "max" { + t.Fatalf("max effort wire=%q, want max", got) + } +} + +func TestApplyClaude_adaptiveOutputConfigEffort(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-opus-4-8", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + Effort: "xhigh", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + th, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || th["type"] != "adaptive" { + t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"]) + } + oc, ok := cfg.ExtraFields["output_config"].(map[string]any) + if !ok { + t.Fatal("expected output_config") + } + if oc["effort"] != "xhigh" { + t.Fatalf("effort=%v", oc["effort"]) + } +} + +func TestApplyClaude_sonnet37OfficialBudget(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-3-7-sonnet-latest", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + Effort: "low", // 3.7 has no output_config.effort; effort is not mapped to budget_tokens + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + th, ok := cfg.ExtraFields["thinking"].(map[string]any) + if !ok || th["type"] != "enabled" { + t.Fatalf("thinking=%#v", cfg.ExtraFields["thinking"]) + } + if th["budget_tokens"] != claudeSonnet37DefaultBudgetTokens { + t.Fatalf("budget_tokens=%v, want official example %d", th["budget_tokens"], claudeSonnet37DefaultBudgetTokens) + } + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("sonnet 3.7 should not set output_config") + } +} + +func TestApplyClaude_onWithoutEffortOmitsOutputConfig(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-sonnet-4-6", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "on", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("on without explicit effort should omit output_config (API default high)") + } +} + +func TestApplyClaude_autoWithoutEffortSkipsOutputConfig(t *testing.T) { + cfg := &einoopenai.ChatModelConfig{} + oa := &config.OpenAIConfig{ + Provider: "claude", + Model: "claude-sonnet-4-6", + Reasoning: config.OpenAIReasoningConfig{ + Mode: "auto", + }, + } + ApplyToEinoChatModelConfig(cfg, oa, nil) + if _, hasOC := cfg.ExtraFields["output_config"]; hasOC { + t.Fatal("auto without effort should omit output_config") + } +} diff --git a/internal/workflow/agent_subgraph.go b/internal/workflow/agent_subgraph.go new file mode 100644 index 00000000..5e5d1286 --- /dev/null +++ b/internal/workflow/agent_subgraph.go @@ -0,0 +1,48 @@ +package workflow + +import ( + "context" + + "github.com/cloudwego/eino/compose" +) + +// compileAgentSubgraph wraps an Agent canvas node as an Eino subgraph (AddGraphNode best practice). +func compileAgentSubgraph(_ context.Context, node graphNode) (compose.AnyGraph, error) { + n := node + prepareID := n.ID + "__agent_prepare" + executeID := n.ID + "__agent_execute" + finalizeID := n.ID + "__agent_finalize" + g := compose.NewGraph[WorkflowNodeOutput, WorkflowNodeOutput]() + _ = g.AddLambdaNode(prepareID, compose.InvokableLambda(func(_ context.Context, input WorkflowNodeOutput) (WorkflowNodeOutput, error) { + if input == nil { + input = WorkflowNodeOutput{} + } + input["agent_subgraph_stage"] = "prepare" + input["agent_node_id"] = n.ID + return input, nil + })) + _ = g.AddLambdaNode(executeID, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowNodeOutput) (WorkflowNodeOutput, error) { + return runWorkflowNodeLambda(runCtx, n) + })) + _ = g.AddLambdaNode(finalizeID, compose.InvokableLambda(func(_ context.Context, output WorkflowNodeOutput) (WorkflowNodeOutput, error) { + if output == nil { + output = WorkflowNodeOutput{} + } + output["agent_subgraph_stage"] = "finalize" + output["agent_node_id"] = n.ID + return output, nil + })) + if err := g.AddEdge(compose.START, prepareID); err != nil { + return nil, err + } + if err := g.AddEdge(prepareID, executeID); err != nil { + return nil, err + } + if err := g.AddEdge(executeID, finalizeID); err != nil { + return nil, err + } + if err := g.AddEdge(finalizeID, compose.END); err != nil { + return nil, err + } + return g, nil +} diff --git a/internal/workflow/bindings.go b/internal/workflow/bindings.go new file mode 100644 index 00000000..f5003cc2 --- /dev/null +++ b/internal/workflow/bindings.go @@ -0,0 +1,153 @@ +package workflow + +import ( + "encoding/json" + "fmt" + "strings" +) + +// FieldBinding selects a value from workflow state (replaces {{...}} templates). +type FieldBinding struct { + From string `json:"from"` // inputs | previous | + Field string `json:"field"` // e.g. output, message +} + +func parseFieldBinding(cfg map[string]any, keys ...string) (FieldBinding, bool) { + for _, key := range keys { + if cfg == nil { + continue + } + raw, ok := cfg[key] + if !ok || raw == nil { + continue + } + switch v := raw.(type) { + case map[string]any: + return FieldBinding{ + From: strings.TrimSpace(fmt.Sprint(v["from"])), + Field: strings.TrimSpace(fmt.Sprint(v["field"])), + }, true + case string: + s := strings.TrimSpace(v) + if s == "" { + continue + } + var b FieldBinding + if err := json.Unmarshal([]byte(s), &b); err == nil && (b.From != "" || b.Field != "") { + return b, true + } + } + } + return FieldBinding{}, false +} + +func defaultBinding(from, field string) FieldBinding { + return FieldBinding{From: from, Field: field} +} + +func resolveBinding(b FieldBinding, state *WorkflowLocalState) any { + from := strings.TrimSpace(b.From) + field := strings.TrimSpace(b.Field) + if field == "" { + field = "output" + } + if from == "" || from == "previous" || from == "prev" { + if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") { + return evalJSONPathValue(state.LastOutput, field) + } + if field == "output" && state.LastOutput != nil { + return state.LastOutput["output"] + } + return valueFromPath("previous."+field, state) + } + if from == "inputs" || from == "input" { + if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") { + return evalJSONPathValue(state.Inputs, field) + } + if field == "" { + return state.Inputs + } + return valueFromPath("inputs."+field, state) + } + if from == "outputs" { + if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") { + return evalJSONPathValue(state.Outputs, field) + } + return valueFromPath("outputs."+field, state) + } + if strings.HasPrefix(field, "$") || strings.HasPrefix(field, ".") { + return evalJSONPathValue(valueFromPath(from, state), field) + } + return valueFromPath(from+"."+field, state) +} + +func resolveBindingString(b FieldBinding, state *WorkflowLocalState) string { + return strings.TrimSpace(fmt.Sprint(resolveBinding(b, state))) +} + +func resolveNodeInputBinding(cfg map[string]any, state *WorkflowLocalState) string { + if b, ok := parseFieldBinding(cfg, "input_binding"); ok { + return resolveBindingString(b, state) + } + // legacy template field removed — default previous.output + return resolveBindingString(defaultBinding("previous", "output"), state) +} + +func resolveOutputSourceBinding(cfg map[string]any, state *WorkflowLocalState) any { + if b, ok := parseFieldBinding(cfg, "source_binding"); ok { + return resolveBinding(b, state) + } + return resolveBinding(defaultBinding("previous", "output"), state) +} + +func resolveHITLPromptBinding(cfg map[string]any, state *WorkflowLocalState) string { + if b, ok := parseFieldBinding(cfg, "prompt_binding"); ok { + return resolveBindingString(b, state) + } + if s := cfgString(cfg, "prompt"); s != "" { + return s + } + return resolveBindingString(defaultBinding("previous", "output"), state) +} + +func toolArgumentBindings(cfg map[string]any) map[string]FieldBinding { + raw, ok := cfg["argument_bindings"].(map[string]any) + if !ok || len(raw) == 0 { + return nil + } + out := make(map[string]FieldBinding, len(raw)) + for argName, v := range raw { + m, ok := v.(map[string]any) + if !ok { + continue + } + out[argName] = FieldBinding{ + From: strings.TrimSpace(fmt.Sprint(m["from"])), + Field: strings.TrimSpace(fmt.Sprint(m["field"])), + } + } + return out +} + +func resolveToolArguments(cfg map[string]any, state *WorkflowLocalState) (map[string]interface{}, error) { + bindings := toolArgumentBindings(cfg) + if len(bindings) > 0 { + args := make(map[string]interface{}, len(bindings)) + for k, b := range bindings { + args[k] = resolveBinding(b, state) + } + return args, nil + } + raw := cfgString(cfg, "arguments") + if raw == "" { + return map[string]interface{}{}, nil + } + var args map[string]interface{} + if err := json.Unmarshal([]byte(raw), &args); err != nil { + return nil, err + } + if args == nil { + args = map[string]interface{}{} + } + return args, nil +} diff --git a/internal/workflow/checkpoint_store.go b/internal/workflow/checkpoint_store.go new file mode 100644 index 00000000..5d254ac0 --- /dev/null +++ b/internal/workflow/checkpoint_store.go @@ -0,0 +1,69 @@ +package workflow + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync" +) + +// fileCheckPointStore persists Eino workflow checkpoints on disk (per run id). +type fileCheckPointStore struct { + dir string + mu sync.RWMutex +} + +func newFileCheckPointStore(dir string) (*fileCheckPointStore, error) { + dir = strings.TrimSpace(dir) + if dir == "" { + dir = filepath.Join("data", "workflow-checkpoints") + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("create workflow checkpoint dir: %w", err) + } + return &fileCheckPointStore{dir: dir}, nil +} + +func (s *fileCheckPointStore) path(id string) (string, error) { + id = strings.TrimSpace(id) + if id == "" { + return "", fmt.Errorf("checkpoint id is empty") + } + if strings.Contains(id, "..") || strings.ContainsAny(id, `/\`) { + return "", fmt.Errorf("invalid checkpoint id") + } + return filepath.Join(s.dir, id+".ckpt"), nil +} + +func (s *fileCheckPointStore) Get(_ context.Context, checkPointID string) ([]byte, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + p, err := s.path(checkPointID) + if err != nil { + return nil, false, err + } + data, err := os.ReadFile(p) + if err != nil { + if os.IsNotExist(err) { + return nil, false, nil + } + return nil, false, err + } + return data, true, nil +} + +func (s *fileCheckPointStore) Set(_ context.Context, checkPointID string, checkPoint []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + p, err := s.path(checkPointID) + if err != nil { + return err + } + tmp := p + ".tmp" + if err := os.WriteFile(tmp, checkPoint, 0o600); err != nil { + return err + } + return os.Rename(tmp, p) +} diff --git a/internal/workflow/draft_generator.go b/internal/workflow/draft_generator.go new file mode 100644 index 00000000..bfb8b680 --- /dev/null +++ b/internal/workflow/draft_generator.go @@ -0,0 +1,782 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "hash/fnv" + "regexp" + "strings" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/openai" + + "go.uber.org/zap" +) + +type DraftTool struct { + Key string `json:"key"` + Name string `json:"name,omitempty"` + Enabled bool `json:"enabled"` +} + +type DraftOptions struct { + IncludeObjective bool `json:"include_objective"` + AllowSchedule bool `json:"allow_schedule"` + AllowHighRisk bool `json:"allow_high_risk"` +} + +type DraftRequest struct { + Prompt string `json:"prompt"` + Options DraftOptions `json:"options"` + AvailableTools []DraftTool `json:"available_tools,omitempty"` +} + +type DraftMeta struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + Enabled bool `json:"enabled"` +} + +type DraftCapability struct { + Label string `json:"label"` + ToolName string `json:"tool_name,omitempty"` + ToolCandidates []string `json:"tool_candidates,omitempty"` +} + +type DraftAudit struct { + Savable bool `json:"savable"` + Validation []string `json:"validation,omitempty"` + MissingFields []string `json:"missing_fields,omitempty"` + RiskWarnings []string `json:"risk_warnings,omitempty"` + Assumptions []string `json:"assumptions,omitempty"` + HighRisk bool `json:"high_risk"` + NeedsHITL bool `json:"needs_hitl"` +} + +type DraftResult struct { + Graph *graphDef `json:"graph"` + Meta DraftMeta `json:"meta"` + Generator string `json:"generator"` + Audit DraftAudit `json:"audit"` + Capabilities []DraftCapability `json:"capabilities,omitempty"` + Stats map[string]int `json:"stats"` +} + +type llmDraftEnvelope struct { + Graph graphDef `json:"graph"` + Meta DraftMeta `json:"meta"` + Capabilities []DraftCapability `json:"capabilities,omitempty"` + Audit DraftAudit `json:"audit,omitempty"` +} + +type draftToolHint struct { + Label string + Keywords []string + Tools []string +} + +var draftToolHints = []draftToolHint{ + {Label: "子域名发现", Keywords: []string{"子域名", "subdomain", "subfinder", "amass"}, Tools: []string{"subfinder", "amass"}}, + {Label: "端口扫描", Keywords: []string{"端口", "port", "nmap", "rustscan", "masscan"}, Tools: []string{"nmap", "rustscan", "masscan"}}, + {Label: "漏洞扫描", Keywords: []string{"漏洞", "vuln", "漏洞扫描", "nuclei", "nikto", "zap"}, Tools: []string{"nuclei", "nikto", "zap"}}, + {Label: "暴露面探测", Keywords: []string{"目录", "路径", "暴露页面", "dir", "ffuf", "gobuster", "feroxbuster"}, Tools: []string{"ffuf", "gobuster", "feroxbuster", "dirsearch"}}, + {Label: "证书与域名线索收集", Keywords: []string{"证书", "certificate", "crt"}, Tools: []string{"subfinder"}}, + {Label: "云配置审计", Keywords: []string{"云", "cloud", "配置审计", "prowler", "scout"}, Tools: []string{"prowler", "scout-suite"}}, + {Label: "容器安全检查", Keywords: []string{"容器", "镜像", "k8s", "kubernetes", "trivy", "kube"}, Tools: []string{"trivy", "kube-bench", "kube-hunter"}}, + {Label: "威胁情报收集", Keywords: []string{"情报", "威胁情报", "threat", "ioc", "virustotal", "shodan", "fofa"}, Tools: []string{"virustotal_search", "shodan_search", "fofa_search"}}, +} + +var highRiskDraftRE = regexp.MustCompile(`(?i)(隔离|封禁|加固|修复|执行|命令|脚本|删除|清理|阻断|封锁|攻击|利用|getshell|shell|payload|exploit|isolate|block|execute|script|delete|exploit|payload)`) + +func GenerateDraftFromNaturalLanguage(ctx context.Context, req DraftRequest) (*DraftResult, error) { + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + return nil, fmt.Errorf("工作流需求不能为空") + } + capabilities := detectDraftCapabilities(prompt, req.AvailableTools) + wantsApproval := containsAnyFold(prompt, "审批", "确认", "审核", "负责人", "人工", "review", "approve", "approval", "human") + wantsReport := containsAnyFold(prompt, "报告", "汇总", "输出", "通知", "任务", "工单", "report", "summary", "notify", "ticket") + wantsCondition := containsAnyFold(prompt, "如果", "发现", "存在", "高危", "新增", "失败", "通过", "否则", "if", "when", "high", "critical", "new", "fail") + highRisk := highRiskDraftRE.MatchString(prompt) + + builder := &draftGraphBuilder{x: 120, y: 150} + assumptions := make([]string, 0) + riskWarnings := make([]string, 0) + missingFields := make([]string, 0) + + start := builder.add("start", "开始", map[string]any{"input_keys": "message, conversationId, projectId, target"}, 0) + previous := start + for _, capability := range capabilities { + hasTool := strings.TrimSpace(capability.ToolName) != "" + var id string + if hasTool { + id = builder.add("tool", capability.Label, map[string]any{ + "tool_name": capability.ToolName, + "arguments": `{"target":"{{inputs.target}}","message":"{{inputs.message}}"}`, + "timeout_seconds": "120", + "join_strategy": "all_merge", + }, 0) + } else { + id = builder.add("agent", capability.Label, map[string]any{ + "agent_mode": "eino_single", + "input_binding": map[string]any{"from": "previous", "field": "output"}, + "instruction": capability.Label + "。根据用户需求执行安全流程步骤，并输出结构化结果：" + prompt, + "output_key": "agent_result", + "join_strategy": "all_merge", + "missing_tool_candidates": strings.Join(capability.ToolCandidates, ", "), + }, 0) + if len(capability.ToolCandidates) > 0 { + assumptions = append(assumptions, capability.Label+" 未匹配到已启用工具，已生成 Agent 草稿节点。") + missingFields = append(missingFields, capability.Label+": 选择或启用对应 MCP 工具") + } + } + builder.connect(previous, id, "", nil) + previous = id + } + + openConditionID := "" + if wantsCondition { + expr := `{{previous.output}} != ""` + label := "是否满足触发条件" + if highRisk { + expr = `{{previous.output}} contains "高危"` + label = "是否需要高风险处置" + } + condition := builder.add("condition", label, map[string]any{"expression": expr, "join_strategy": "all_merge"}, 0) + builder.connect(previous, condition, "", nil) + openConditionID = condition + report := builder.add("output", draftOutputLabel(wantsReport), map[string]any{ + "output_key": "result", + "source_binding": map[string]any{"from": "previous", "field": "output"}, + "static_value": "", + "join_strategy": "all_merge", + }, 130) + builder.connect(condition, report, "否", map[string]any{"condition": `{{previous.matched}} == "false"`, "branch": "false"}) + previous = condition + } + + insertedHITL := false + if highRisk { + if !req.Options.AllowHighRisk || wantsApproval { + approval := builder.add("hitl", "人工审批", map[string]any{ + "prompt": "请确认是否允许继续执行高风险处置：" + prompt, + "prompt_binding": map[string]any{"from": "previous", "field": "output"}, + "reviewer": "human", + "join_strategy": "all_merge", + "risk_level": "high", + }, 0) + builder.connect(previous, approval, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true)) + if previous == openConditionID { + openConditionID = "" + } + previous = approval + insertedHITL = true + } + action := builder.add("agent", "执行受控处置", map[string]any{ + "agent_mode": "eino_single", + "input_binding": map[string]any{"from": "previous", "field": "output"}, + "instruction": "仅在授权范围内生成处置步骤草稿；实际执行前必须由人工确认。用户需求：" + prompt, + "output_key": "remediation_plan", + "join_strategy": "all_merge", + "risk_level": "high", + "requires_human_confirmation": "true", + }, 0) + builder.connect(previous, action, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true)) + if previous == openConditionID { + openConditionID = "" + } + previous = action + if insertedHITL { + riskWarnings = append(riskWarnings, "检测到高风险动作，已加入人工审批与 requires_human_confirmation 标记。") + } else { + riskWarnings = append(riskWarnings, "检测到高风险动作，已保留为草稿并添加 requires_human_confirmation 标记。") + } + } else if wantsApproval { + approval := builder.add("hitl", "人工审批", map[string]any{ + "prompt": "请审核工作流阶段结果：" + prompt, + "prompt_binding": map[string]any{"from": "previous", "field": "output"}, + "reviewer": "human", + "join_strategy": "all_merge", + }, 0) + builder.connect(previous, approval, "", nil) + previous = approval + insertedHITL = true + } + + output := builder.add("output", draftOutputLabel(wantsReport), map[string]any{ + "output_key": "result", + "source_binding": map[string]any{"from": "previous", "field": "output"}, + "static_value": "", + "join_strategy": "all_merge", + }, 0) + builder.connect(previous, output, branchLabel(previous, openConditionID), branchConfig(previous, openConditionID, true)) + + graph := &graphDef{Nodes: builder.nodes, Edges: builder.edges, Config: map[string]any{ + "schema_version": 1, + "generated_by": "natural_language", + "source_prompt": prompt, + }} + if req.Options.IncludeObjective { + graph.Config["objective"] = prompt + } + if req.Options.AllowSchedule && containsAnyFold(prompt, "每天", "每周", "定时", "周期", "持续", "daily", "weekly", "schedule", "monitor") { + if containsAnyFold(prompt, "每天", "daily") { + graph.Config["trigger_suggestion"] = "daily" + } else { + graph.Config["trigger_suggestion"] = "scheduled" + } + assumptions = append(assumptions, "已记录定时触发建议；保存后仍需在触发器或角色绑定处配置。") + } + + raw, _ := json.Marshal(graph) + validation := make([]string, 0) + if err := ValidateGraphJSON(ctx, string(raw)); err != nil { + validation = append(validation, err.Error()) + } + return &DraftResult{ + Graph: graph, + Meta: DraftMeta{ID: draftSlug(prompt), Name: draftName(prompt), Description: prompt, Enabled: true}, + Generator: "deterministic", + Audit: DraftAudit{ + Savable: len(validation) == 0, + Validation: validation, + MissingFields: missingFields, + RiskWarnings: riskWarnings, + Assumptions: assumptions, + HighRisk: highRisk, + NeedsHITL: insertedHITL, + }, + Capabilities: capabilities, + Stats: map[string]int{"nodes": len(graph.Nodes), "edges": len(graph.Edges)}, + }, nil +} + +func GenerateDraftFromLLM(ctx context.Context, req DraftRequest, oa config.OpenAIConfig, logger *zap.Logger) (*DraftResult, error) { + prompt := strings.TrimSpace(req.Prompt) + if prompt == "" { + return nil, fmt.Errorf("工作流需求不能为空") + } + if strings.TrimSpace(oa.APIKey) == "" || strings.TrimSpace(oa.Model) == "" { + return nil, fmt.Errorf("AI 通道未配置 api_key 或 model") + } + if logger == nil { + logger = zap.NewNop() + } + callCtx, cancel := context.WithTimeout(ctx, 90*time.Second) + defer cancel() + toolJSON, _ := json.Marshal(req.AvailableTools) + systemPrompt := `你是 CyberStrikeAI 的工作流编排助手。你必须把用户的一句话需求转换为可保存的工作流草稿 JSON。 +只返回 JSON 对象，不要 Markdown，不要解释。JSON 必须符合： +{ + "meta": {"id":"kebab-case-id","name":"短名称","description":"用户需求","enabled":true}, + "graph": { + "nodes": [{"id":"start-1","type":"start","label":"显示名","position":{"x":120,"y":150},"config":{}}], + "edges": [{"id":"edge-1","source":"start-1","target":"node-2","label":"","config":{}}], + "config": {"schema_version":1,"generated_by":"llm","source_prompt":"用户原文"} + }, + "capabilities": [{"label":"能力名","tool_name":"已匹配工具名","tool_candidates":["候选工具"]}], + "audit": {"assumptions":[],"missing_fields":[],"risk_warnings":[]} +} +硬性规则： +- 只能输出一个合法 JSON object；不要输出 JSON Schema、注释、解释文字、Markdown 代码块或多余前后缀。 +- 不要在 JSON 字符串值中使用竖线枚举写法；type 字段一次只能填写一个节点类型字符串。 +- 至少 1 个 start 和 1 个 output；output/end 不能有出边。 +- 节点 type 只能从这些字符串中选择：start、tool、agent、condition、hitl、output、end。 +- 每个 agent、tool、output 节点都必须配置唯一的 output_key；output 节点默认使用 result。 +- agent 节点必须配置 instruction 或 input_binding；默认 input_binding 为 {"from":"previous","field":"output"}。 +- output 节点必须配置 source_binding 或 static_value；默认 source_binding 为 {"from":"previous","field":"output"}。 +- tool 节点必须配置 tool_name、arguments、timeout_seconds；arguments 必须是合法 JSON 字符串。 +- 所有非 start 且可能有多个上游的节点必须配置 join_strategy:"all_merge"。 +- condition 最多 2 条出边，必须用 branch true/false，并用 label 是/否。 +- tool 节点只有在 available_tools 中存在启用工具时才使用，否则用 agent 节点并在 audit.missing_fields 写明缺失工具。 +- 高风险动作（执行脚本、隔离、封禁、删除、利用、payload、命令执行等）必须加入 hitl 审批，或在高风险节点 config 中标记 requires_human_confirmation:"true"、risk_level:"high"。 +- 不要生成会真实执行攻击的参数；工具参数使用 {{inputs.target}}、{{inputs.message}} 占位。 +- 所有节点 config 加 generated_by:"llm" 和 needs_review:"true"。` + userPrompt := fmt.Sprintf("用户需求：%s\n\n选项：%+v\n\n可用工具 JSON：%s", prompt, req.Options, string(toolJSON)) + requestBody := map[string]interface{}{ + "model": strings.TrimSpace(oa.Model), + "messages": []map[string]interface{}{ + {"role": "system", "content": systemPrompt}, + {"role": "user", "content": userPrompt}, + }, + "temperature": 0, + "max_completion_tokens": 4096, + "response_format": map[string]interface{}{"type": "json_object"}, + "thinking": map[string]interface{}{"type": "disabled"}, + } + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + } `json:"message"` + } `json:"choices"` + } + client := openai.NewClient(&oa, nil, logger) + if err := client.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil { + return nil, fmt.Errorf("调用大模型失败: %w", err) + } + if len(apiResponse.Choices) == 0 { + return nil, fmt.Errorf("大模型未返回候选结果") + } + raw := strings.TrimSpace(apiResponse.Choices[0].Message.Content) + if raw == "" { + raw = strings.TrimSpace(apiResponse.Choices[0].Message.ReasoningContent) + } + env, err := parseLLMDraftEnvelope(raw) + if err != nil { + return nil, err + } + result := normalizeLLMDraft(prompt, req, env) + graphRaw, _ := json.Marshal(result.Graph) + validation := make([]string, 0) + if err := ValidateGraphJSON(ctx, string(graphRaw)); err != nil { + validation = append(validation, err.Error()) + } + result.Audit.Validation = validation + result.Audit.Savable = len(validation) == 0 + if !result.Audit.Savable { + return nil, fmt.Errorf("大模型生成的工作流未通过校验: %s", strings.Join(validation, "；")) + } + return result, nil +} + +type draftGraphBuilder struct { + nodes []graphNode + edges []graphEdge + x float64 + y float64 + nodeSeq int + edgeSeq int +} + +func (b *draftGraphBuilder) add(nodeType, label string, config map[string]any, yOffset float64) string { + b.nodeSeq++ + id := fmt.Sprintf("%s-%d", nodeType, b.nodeSeq) + if config == nil { + config = make(map[string]any) + } + config["generated_by"] = "natural_language" + config["needs_review"] = "true" + b.nodes = append(b.nodes, graphNode{ + ID: id, + Type: nodeType, + Label: label, + Position: graphPosition{X: b.x, Y: b.y + yOffset}, + Config: config, + }) + b.x += 210 + return id +} + +func (b *draftGraphBuilder) connect(source, target, label string, config map[string]any) { + b.edgeSeq++ + if config == nil { + config = make(map[string]any) + } + b.edges = append(b.edges, graphEdge{ID: fmt.Sprintf("edge-ai-%d", b.edgeSeq), Source: source, Target: target, Label: label, Config: config}) +} + +func parseLLMDraftEnvelope(raw string) (llmDraftEnvelope, error) { + var lastErr error + for _, candidate := range jsonObjectCandidates(raw) { + var env llmDraftEnvelope + if err := json.Unmarshal([]byte(candidate), &env); err == nil { + if len(env.Graph.Nodes) == 0 { + lastErr = fmt.Errorf("大模型 JSON 缺少 graph.nodes") + continue + } + return env, nil + } else { + lastErr = err + } + } + if lastErr == nil { + lastErr = fmt.Errorf("大模型响应为空") + } + return llmDraftEnvelope{}, fmt.Errorf("解析大模型工作流 JSON 失败: %w", lastErr) +} + +func jsonObjectCandidates(raw string) []string { + s := strings.TrimSpace(raw) + s = strings.TrimPrefix(s, "```json") + s = strings.TrimPrefix(s, "```") + s = strings.TrimSuffix(s, "```") + s = strings.TrimSpace(s) + candidates := []string{s} + if start := strings.Index(s, "{"); start >= 0 { + if end := strings.LastIndex(s, "}"); end > start { + candidates = append(candidates, s[start:end+1]) + } + } + return candidates +} + +func normalizeLLMDraft(prompt string, req DraftRequest, env llmDraftEnvelope) *DraftResult { + g := env.Graph + if g.Config == nil { + g.Config = make(map[string]any) + } + g.Config["schema_version"] = 1 + g.Config["generated_by"] = "llm" + g.Config["source_prompt"] = prompt + if req.Options.IncludeObjective { + g.Config["objective"] = prompt + } + enabledTools := enabledDraftToolNames(req.AvailableTools) + usedOutputKeys := make(map[string]bool) + nodeTypes := make(map[string]string, len(g.Nodes)) + for i := range g.Nodes { + if strings.TrimSpace(g.Nodes[i].ID) == "" { + g.Nodes[i].ID = fmt.Sprintf("%s-%d", firstNonEmpty(g.Nodes[i].Type, "node"), i+1) + } + if strings.TrimSpace(g.Nodes[i].Type) == "" { + g.Nodes[i].Type = "agent" + } + if strings.TrimSpace(g.Nodes[i].Label) == "" { + g.Nodes[i].Label = displayNodeType(g.Nodes[i].Type) + } + if g.Nodes[i].Position.X == 0 && g.Nodes[i].Position.Y == 0 { + g.Nodes[i].Position = graphPosition{X: 120 + float64(i)*210, Y: 150} + } + if g.Nodes[i].Config == nil { + g.Nodes[i].Config = make(map[string]any) + } + g.Nodes[i].Config["generated_by"] = "llm" + g.Nodes[i].Config["needs_review"] = "true" + normalizeLLMNodeConfig(prompt, &g.Nodes[i], enabledTools, usedOutputKeys) + nodeTypes[g.Nodes[i].ID] = strings.ToLower(strings.TrimSpace(g.Nodes[i].Type)) + } + conditionBranchCounts := make(map[string]int) + for i := range g.Edges { + if strings.TrimSpace(g.Edges[i].ID) == "" { + g.Edges[i].ID = fmt.Sprintf("edge-llm-%d", i+1) + } + if g.Edges[i].Config == nil { + g.Edges[i].Config = make(map[string]any) + } + normalizeLLMEdgeConfig(&g.Edges[i], nodeTypes, conditionBranchCounts) + } + audit := env.Audit + highRisk := highRiskDraftRE.MatchString(prompt) || graphHasHighRisk(g) + audit.HighRisk = highRisk + audit.NeedsHITL = graphHasNodeType(g, "hitl") + if highRisk && !audit.NeedsHITL && !graphHasConfirmation(g) { + audit.RiskWarnings = append(audit.RiskWarnings, "大模型生成包含高风险语义，请补充人工审批或确认标记后再运行。") + } + if len(audit.RiskWarnings) == 0 && highRisk { + audit.RiskWarnings = append(audit.RiskWarnings, "检测到高风险动作，已标记为需要重点审计。") + } + meta := env.Meta + if strings.TrimSpace(meta.Description) == "" { + meta.Description = prompt + } + if strings.TrimSpace(meta.Name) == "" { + meta.Name = draftName(prompt) + } + if strings.TrimSpace(meta.ID) == "" { + meta.ID = draftSlug(prompt) + } + meta.Enabled = true + return &DraftResult{ + Graph: &g, + Meta: meta, + Generator: "llm", + Audit: audit, + Capabilities: env.Capabilities, + Stats: map[string]int{"nodes": len(g.Nodes), "edges": len(g.Edges)}, + } +} + +func normalizeLLMEdgeConfig(edge *graphEdge, nodeTypes map[string]string, conditionBranchCounts map[string]int) { + if nodeTypes[strings.TrimSpace(edge.Source)] != "condition" { + return + } + if conditionBranchHint(*edge) != "" { + return + } + conditionBranchCounts[edge.Source]++ + branch := "true" + label := "是" + if conditionBranchCounts[edge.Source] > 1 { + branch = "false" + label = "否" + } + edge.Label = label + edge.Config["branch"] = branch +} + +func normalizeLLMNodeConfig(prompt string, node *graphNode, enabledTools map[string]bool, usedOutputKeys map[string]bool) { + nodeType := strings.ToLower(strings.TrimSpace(node.Type)) + switch nodeType { + case "start": + if cfgString(node.Config, "input_keys") == "" { + node.Config["input_keys"] = "message, conversationId, projectId, target" + } + case "tool": + toolName := cfgString(node.Config, "tool_name") + if toolName == "" || !enabledTools[strings.ToLower(toolName)] { + node.Type = "agent" + node.Config["missing_tool_name"] = toolName + normalizeAgentDraftConfig(prompt, node, usedOutputKeys) + return + } + if cfgString(node.Config, "arguments") == "" { + node.Config["arguments"] = `{"target":"{{inputs.target}}","message":"{{inputs.message}}"}` + } + if cfgString(node.Config, "timeout_seconds") == "" { + node.Config["timeout_seconds"] = "120" + } + ensureNodeOutputKey(node, usedOutputKeys, draftOutputKeyBase(node, "tool_result")) + ensureJoinStrategy(node) + case "agent": + normalizeAgentDraftConfig(prompt, node, usedOutputKeys) + case "condition": + if cfgString(node.Config, "expression") == "" { + node.Config["expression"] = `{{previous.output}} != ""` + } + ensureJoinStrategy(node) + case "hitl": + if cfgString(node.Config, "prompt") == "" { + node.Config["prompt"] = "请审核工作流阶段结果：" + prompt + } + if cfgString(node.Config, "reviewer") == "" { + node.Config["reviewer"] = "human" + } + ensureJoinStrategy(node) + case "output": + ensureNodeOutputKey(node, usedOutputKeys, "result") + if cfgString(node.Config, "static_value") == "" { + if _, ok := parseFieldBinding(node.Config, "source_binding"); !ok { + node.Config["source_binding"] = map[string]any{"from": "previous", "field": "output"} + } + } + ensureJoinStrategy(node) + case "end": + ensureJoinStrategy(node) + } +} + +func normalizeAgentDraftConfig(prompt string, node *graphNode, usedOutputKeys map[string]bool) { + if cfgString(node.Config, "agent_mode") == "" { + node.Config["agent_mode"] = "eino_single" + } + if cfgString(node.Config, "instruction") == "" { + node.Config["instruction"] = node.Label + "。根据用户需求执行安全流程步骤，并输出结构化结果：" + prompt + } + if _, ok := parseFieldBinding(node.Config, "input_binding"); !ok { + node.Config["input_binding"] = map[string]any{"from": "previous", "field": "output"} + } + ensureNodeOutputKey(node, usedOutputKeys, draftOutputKeyBase(node, "agent_result")) + ensureJoinStrategy(node) +} + +func ensureJoinStrategy(node *graphNode) { + if cfgString(node.Config, "join_strategy") == "" { + node.Config["join_strategy"] = "all_merge" + } +} + +func ensureNodeOutputKey(node *graphNode, used map[string]bool, fallback string) { + current := sanitizeOutputKey(cfgString(node.Config, "output_key")) + if current == "" { + current = sanitizeOutputKey(fallback) + } + if current == "" { + current = "result" + } + base := current + for i := 2; used[current]; i++ { + current = fmt.Sprintf("%s_%d", base, i) + } + node.Config["output_key"] = current + used[current] = true +} + +func draftOutputKeyBase(node *graphNode, fallback string) string { + if name := cfgString(node.Config, "tool_name"); name != "" { + return name + "_result" + } + if node.ID != "" { + return node.ID + "_result" + } + return fallback +} + +func sanitizeOutputKey(value string) string { + value = strings.ToLower(strings.TrimSpace(value)) + var b strings.Builder + lastUnderscore := false + for _, r := range value { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') { + b.WriteRune(r) + lastUnderscore = false + continue + } + if b.Len() > 0 && !lastUnderscore { + b.WriteByte('_') + lastUnderscore = true + } + } + return strings.Trim(b.String(), "_") +} + +func enabledDraftToolNames(tools []DraftTool) map[string]bool { + names := make(map[string]bool, len(tools)*2) + for _, tool := range tools { + if !tool.Enabled { + continue + } + if key := strings.ToLower(strings.TrimSpace(tool.Key)); key != "" { + names[key] = true + } + if name := strings.ToLower(strings.TrimSpace(tool.Name)); name != "" { + names[name] = true + } + } + return names +} + +func graphHasNodeType(g graphDef, nodeType string) bool { + for _, node := range g.Nodes { + if strings.EqualFold(node.Type, nodeType) { + return true + } + } + return false +} + +func graphHasConfirmation(g graphDef) bool { + for _, node := range g.Nodes { + if cfgString(node.Config, "requires_human_confirmation") == "true" { + return true + } + } + return false +} + +func graphHasHighRisk(g graphDef) bool { + for _, node := range g.Nodes { + if cfgString(node.Config, "risk_level") == "high" || cfgString(node.Config, "requires_human_confirmation") == "true" { + return true + } + if highRiskDraftRE.MatchString(node.Label) || highRiskDraftRE.MatchString(cfgString(node.Config, "instruction")) { + return true + } + } + return false +} + +func detectDraftCapabilities(prompt string, tools []DraftTool) []DraftCapability { + capabilities := make([]DraftCapability, 0) + for _, hint := range draftToolHints { + if containsAnyFold(prompt, hint.Keywords...) { + capabilities = append(capabilities, DraftCapability{ + Label: hint.Label, + ToolName: matchDraftTool(hint.Tools, tools), + ToolCandidates: append([]string(nil), hint.Tools...), + }) + } + } + if len(capabilities) == 0 { + capabilities = append(capabilities, DraftCapability{Label: "节点能力", ToolCandidates: nil}) + } + return capabilities +} + +func matchDraftTool(candidates []string, tools []DraftTool) string { + if len(candidates) == 0 || len(tools) == 0 { + return "" + } + for _, enabledOnly := range []bool{true, false} { + for _, candidate := range candidates { + candidate = strings.ToLower(strings.TrimSpace(candidate)) + for _, tool := range tools { + if enabledOnly && !tool.Enabled { + continue + } + key := strings.ToLower(strings.TrimSpace(firstNonEmpty(tool.Key, tool.Name))) + if key != "" && strings.Contains(key, candidate) { + return firstNonEmpty(tool.Key, tool.Name) + } + } + } + } + return "" +} + +func containsAnyFold(text string, needles ...string) bool { + lower := strings.ToLower(text) + for _, needle := range needles { + if strings.Contains(lower, strings.ToLower(needle)) { + return true + } + } + return false +} + +func draftOutputLabel(wantsReport bool) string { + if wantsReport { + return "输出报告" + } + return "输出" +} + +func branchLabel(source, conditionID string) string { + if source == conditionID && conditionID != "" { + return "是" + } + return "" +} + +func branchConfig(source, conditionID string, yes bool) map[string]any { + if source != conditionID || conditionID == "" { + return nil + } + if yes { + return map[string]any{"condition": `{{previous.matched}} == "true"`, "branch": "true"} + } + return map[string]any{"condition": `{{previous.matched}} == "false"`, "branch": "false"} +} + +func draftName(prompt string) string { + runes := []rune(strings.TrimSpace(prompt)) + if len(runes) > 22 { + return string(runes[:22]) + "..." + } + return string(runes) +} + +func draftSlug(prompt string) string { + lower := strings.ToLower(strings.TrimSpace(prompt)) + var b strings.Builder + lastDash := false + for _, r := range lower { + if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' { + b.WriteRune(r) + lastDash = false + continue + } + if !lastDash && b.Len() > 0 { + b.WriteByte('-') + lastDash = true + } + } + slug := strings.Trim(b.String(), "-") + if slug != "" { + if len(slug) > 48 { + return strings.Trim(slug[:48], "-") + } + return slug + } + h := fnv.New32a() + _, _ = h.Write([]byte(lower)) + if !utf8.ValidString(lower) || lower == "" { + lower = "workflow" + } + return fmt.Sprintf("ai-workflow-%x", h.Sum32()) +} diff --git a/internal/workflow/draft_generator_test.go b/internal/workflow/draft_generator_test.go new file mode 100644 index 00000000..e5ec61e7 --- /dev/null +++ b/internal/workflow/draft_generator_test.go @@ -0,0 +1,213 @@ +package workflow + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "cyberstrike-ai/internal/config" +) + +func TestGenerateDraftFromNaturalLanguageHighRiskAddsHITLAndValidGraph(t *testing.T) { + result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{ + Prompt: "对目标资产做端口扫描，如果发现高危端口就执行加固脚本，最后输出报告", + Options: DraftOptions{ + IncludeObjective: true, + AllowSchedule: false, + AllowHighRisk: false, + }, + AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}}, + }) + if err != nil { + t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err) + } + raw, _ := json.Marshal(result.Graph) + if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil { + t.Fatalf("generated graph should validate: %v\n%s", err, raw) + } + if !result.Audit.HighRisk || !result.Audit.NeedsHITL || len(result.Audit.RiskWarnings) == 0 { + t.Fatalf("audit did not flag high-risk HITL path: %#v", result.Audit) + } + var hasTool, hasHITL, hasConfirmation bool + for _, node := range result.Graph.Nodes { + if node.Type == "tool" && cfgString(node.Config, "tool_name") == "nmap" { + hasTool = true + } + if node.Type == "hitl" { + hasHITL = true + } + if cfgString(node.Config, "requires_human_confirmation") == "true" { + hasConfirmation = true + } + } + if !hasTool || !hasHITL || !hasConfirmation { + t.Fatalf("expected nmap tool, HITL, and confirmation marker; tool=%v hitl=%v confirmation=%v", hasTool, hasHITL, hasConfirmation) + } +} + +func TestGenerateDraftAllowHighRiskStillLabelsConditionBranch(t *testing.T) { + result, err := GenerateDraftFromNaturalLanguage(context.Background(), DraftRequest{ + Prompt: "如果漏洞扫描发现高危漏洞，允许生成执行修复脚本的草稿并输出报告", + Options: DraftOptions{ + AllowHighRisk: true, + }, + AvailableTools: []DraftTool{{Key: "nuclei", Name: "nuclei", Enabled: true}}, + }) + if err != nil { + t.Fatalf("GenerateDraftFromNaturalLanguage: %v", err) + } + raw, _ := json.Marshal(result.Graph) + if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil { + t.Fatalf("generated graph should validate: %v\n%s", err, raw) + } + branches := map[string]bool{} + for _, edge := range result.Graph.Edges { + if branch := cfgString(edge.Config, "branch"); branch != "" { + branches[branch] = true + } + } + if !branches["true"] || !branches["false"] { + t.Fatalf("condition branches = %#v, want true and false", branches) + } +} + +func TestGenerateDraftFromLLMUsesOpenAICompatibleEndpoint(t *testing.T) { + called := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + if r.URL.Path != "/chat/completions" { + t.Fatalf("path = %s, want /chat/completions", r.URL.Path) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Fatalf("authorization = %q", got) + } + var payload struct { + Temperature float64 `json:"temperature"` + ResponseFormat struct { + Type string `json:"type"` + } `json:"response_format"` + Messages []struct { + Role string `json:"role"` + Content string `json:"content"` + } `json:"messages"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode request: %v", err) + } + if payload.Temperature != 0 || payload.ResponseFormat.Type != "json_object" { + t.Fatalf("unexpected structured output controls: temperature=%v response_format=%#v", payload.Temperature, payload.ResponseFormat) + } + if len(payload.Messages) == 0 || strings.Contains(payload.Messages[0].Content, "start|tool|agent") { + t.Fatalf("system prompt still contains pipe enum: %q", payload.Messages[0].Content) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"llm-port-scan\",\"name\":\"端口扫描\",\"description\":\"端口扫描\",\"enabled\":true},\"graph\":{\"nodes\":[{\"id\":\"start-1\",\"type\":\"start\",\"label\":\"开始\",\"position\":{\"x\":120,\"y\":150},\"config\":{\"input_keys\":\"message, target\"}},{\"id\":\"tool-2\",\"type\":\"tool\",\"label\":\"端口扫描\",\"position\":{\"x\":330,\"y\":150},\"config\":{\"tool_name\":\"nmap\",\"arguments\":\"{\\\"target\\\":\\\"{{inputs.target}}\\\"}\",\"timeout_seconds\":\"120\",\"join_strategy\":\"all_merge\"}},{\"id\":\"output-3\",\"type\":\"output\",\"label\":\"输出报告\",\"position\":{\"x\":540,\"y\":150},\"config\":{\"source_binding\":{\"from\":\"previous\",\"field\":\"output\"},\"join_strategy\":\"all_merge\"}}],\"edges\":[{\"id\":\"edge-1\",\"source\":\"start-1\",\"target\":\"tool-2\"},{\"id\":\"edge-2\",\"source\":\"tool-2\",\"target\":\"output-3\"}],\"config\":{\"schema_version\":1}},\"capabilities\":[{\"label\":\"端口扫描\",\"tool_name\":\"nmap\",\"tool_candidates\":[\"nmap\"]}],\"audit\":{\"assumptions\":[]}}"}}]}`)) + })) + defer srv.Close() + + result, err := GenerateDraftFromLLM(context.Background(), DraftRequest{ + Prompt: "对目标做端口扫描并输出报告", + AvailableTools: []DraftTool{{Key: "nmap", Name: "nmap", Enabled: true}}, + }, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil) + if err != nil { + t.Fatalf("GenerateDraftFromLLM: %v", err) + } + if !called { + t.Fatal("expected LLM endpoint to be called") + } + if result.Generator != "llm" || !result.Audit.Savable || result.Meta.ID != "llm-port-scan" { + t.Fatalf("unexpected result: %#v", result) + } + for _, node := range result.Graph.Nodes { + if node.Type == "output" && cfgString(node.Config, "output_key") != "result" { + t.Fatalf("output_key = %q, want result", cfgString(node.Config, "output_key")) + } + } +} + +func TestGenerateDraftFromLLMReturnsErrorOnMalformedJSON(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"meta\":{\"id\":\"bad\"},|\"graph\":{\"nodes\":[]}}"}}]}`)) + })) + defer srv.Close() + + _, err := GenerateDraftFromLLM(context.Background(), DraftRequest{ + Prompt: "随便生成一个工作流，要求所有节点都用到输出变量", + Options: DraftOptions{ + IncludeObjective: true, + }, + }, config.OpenAIConfig{APIKey: "test-key", BaseURL: srv.URL, Model: "test-model"}, nil) + if err == nil { + t.Fatal("expected malformed JSON error") + } + if !strings.Contains(err.Error(), "解析大模型工作流 JSON 失败") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestNormalizeLLMDraftRepairsMissingRequiredConfig(t *testing.T) { + result := normalizeLLMDraft("随便生成一个工作流", DraftRequest{}, llmDraftEnvelope{ + Graph: graphDef{ + Nodes: []graphNode{ + {ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}}, + {ID: "agent-1", Type: "agent", Label: "分析", Config: map[string]any{}}, + {ID: "out-1", Type: "output", Label: "输出结果", Config: map[string]any{}}, + }, + Edges: []graphEdge{ + {ID: "e1", Source: "start-1", Target: "agent-1"}, + {ID: "e2", Source: "agent-1", Target: "out-1"}, + }, + }, + }) + raw, _ := json.Marshal(result.Graph) + if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil { + t.Fatalf("normalized graph should validate: %v\n%s", err, raw) + } + var agentKey, outputKey string + for _, node := range result.Graph.Nodes { + switch node.Type { + case "agent": + agentKey = cfgString(node.Config, "output_key") + case "output": + outputKey = cfgString(node.Config, "output_key") + } + } + if agentKey == "" || outputKey != "result" { + t.Fatalf("agentKey=%q outputKey=%q", agentKey, outputKey) + } +} + +func TestNormalizeLLMDraftRepairsConditionBranches(t *testing.T) { + result := normalizeLLMDraft("如果发现异常则输出详情，否则输出正常", DraftRequest{}, llmDraftEnvelope{ + Graph: graphDef{ + Nodes: []graphNode{ + {ID: "start-1", Type: "start", Label: "开始", Config: map[string]any{}}, + {ID: "cond-1", Type: "condition", Label: "判断", Config: map[string]any{"expression": `{{inputs.message}} != ""`}}, + {ID: "out-yes", Type: "output", Label: "异常", Config: map[string]any{}}, + {ID: "out-no", Type: "output", Label: "正常", Config: map[string]any{}}, + }, + Edges: []graphEdge{ + {ID: "e1", Source: "start-1", Target: "cond-1"}, + {ID: "e2", Source: "cond-1", Target: "out-yes"}, + {ID: "e3", Source: "cond-1", Target: "out-no"}, + }, + }, + }) + raw, _ := json.Marshal(result.Graph) + if err := ValidateGraphJSON(context.Background(), string(raw)); err != nil { + t.Fatalf("normalized graph should validate: %v\n%s", err, raw) + } + branches := map[string]bool{} + for _, edge := range result.Graph.Edges { + if edge.Source == "cond-1" { + branches[cfgString(edge.Config, "branch")] = true + } + } + if !branches["true"] || !branches["false"] { + t.Fatalf("branches = %#v, want true and false", branches) + } +} diff --git a/internal/workflow/dry_run.go b/internal/workflow/dry_run.go new file mode 100644 index 00000000..0a79f0bf --- /dev/null +++ b/internal/workflow/dry_run.go @@ -0,0 +1,173 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" +) + +type DryRunResult struct { + Outputs map[string]any `json:"outputs"` + NodeOutputs map[string]map[string]any `json:"nodeOutputs"` + Executed []string `json:"executed"` + Skipped []string `json:"skipped"` + Trace []map[string]any `json:"trace"` + Metrics map[string]any `json:"metrics"` + ReplayScript []map[string]any `json:"replayScript"` +} + +func DryRunGraphJSON(ctx context.Context, graphJSON string, inputs map[string]any) (*DryRunResult, error) { + g, err := parseGraph(graphJSON) + if err != nil { + return nil, err + } + idx := indexGraph(g) + if err := validateGraphDefinition(g, idx); err != nil { + return nil, err + } + in := make(map[string]interface{}, len(inputs)) + for k, v := range inputs { + in[k] = v + } + if _, ok := in["message"]; !ok { + in["message"] = "" + } + state := newWorkflowLocalState(in, "dry-run") + rt := &workflowRuntime{runID: "dry-run", idx: idx, state: state} + trace := []map[string]any{} + executedIDs := map[string]bool{} + queue := findStartNodeIDs(idx) + for len(queue) > 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + nodeID := queue[0] + queue = queue[1:] + if executedIDs[nodeID] { + continue + } + node := idx.nodes[nodeID] + if !dryRunPredecessorsReady(idx, nodeID, executedIDs) { + queue = append(queue, nodeID) + continue + } + if err := prepareNodeInputState(rt, node); err != nil { + return nil, err + } + started := time.Now() + out, proceed, status, errText := dryRunNode(node, state) + out["duration_ms"] = time.Since(started).Milliseconds() + out["status"] = status + state.NodeOutputs[node.ID] = out + state.LastOutput = out + executedIDs[nodeID] = true + if status == "skipped" { + state.Skipped = append(state.Skipped, firstNonEmpty(node.Label, node.ID)) + } else { + state.Executed = append(state.Executed, firstNonEmpty(node.Label, node.ID)) + } + trace = append(trace, map[string]any{ + "nodeId": node.ID, + "label": firstNonEmpty(node.Label, node.ID), + "type": node.Type, + "status": status, + "error": errText, + "output": out, + "previous": state.LastOutput, + }) + if !proceed { + continue + } + for edgeIdx, edge := range idx.outgoing[nodeID] { + if edgeAllowed(edge, node, edgeIdx, state) { + queue = append(queue, edge.Target) + } + } + } + for id, node := range idx.nodes { + if !executedIDs[id] { + state.Skipped = append(state.Skipped, firstNonEmpty(node.Label, id)) + } + } + return &DryRunResult{ + Outputs: state.Outputs, + NodeOutputs: state.NodeOutputs, + Executed: state.Executed, + Skipped: state.Skipped, + Trace: trace, + Metrics: state.Metrics, + ReplayScript: buildReplayScript(trace), + }, nil +} + +func dryRunPredecessorsReady(idx *graphIndex, nodeID string, executed map[string]bool) bool { + for _, edge := range idx.incoming[nodeID] { + if !executed[edge.Source] { + return false + } + } + return true +} + +func dryRunNode(node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { + switch strings.ToLower(strings.TrimSpace(node.Type)) { + case "start": + return startOutputMap(node, state.Inputs["message"], state.Inputs["conversationId"], state.Inputs["projectId"]), true, "completed", "" + case "condition": + expr := cfgString(node.Config, "expression") + matched := evalCondition(expr, state) + return conditionOutputMap(node, expr, matched), true, "completed", "" + case "output": + key := cfgString(node.Config, "output_key") + value := resolveOutputSourceBinding(node.Config, state) + if static := cfgString(node.Config, "static_value"); static != "" { + value = static + } + state.Outputs[key] = value + return outputNodeOutputMap(node, key, value), true, "completed", "" + case "end": + value := resolveOutputSourceBinding(node.Config, state) + if b, ok := parseFieldBinding(node.Config, "result_binding"); ok { + value = resolveBinding(b, state) + } + return endOutputMap(node, value), false, "completed", "" + case "tool": + args, err := resolveToolArguments(node.Config, state) + if err != nil { + errText := fmt.Sprintf("工具参数不是合法 JSON：%v", err) + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText + } + return toolOutputMap(node, "[dry-run] tool call skipped", cfgString(node.Config, "tool_name"), args, "dry-run", false), true, "simulated", "" + case "agent": + mode := firstNonEmpty(cfgString(node.Config, "agent_mode"), "eino_single") + response := "[dry-run] agent execution skipped" + if key := cfgString(node.Config, "output_key"); key != "" { + state.Outputs[key] = response + } + return agentOutputMap(node, response, mode, nil), true, "simulated", "" + case "hitl": + prompt := resolveHITLPromptBinding(node.Config, state) + return hitlOutputMap(node, "simulated", prompt, prompt, firstNonEmpty(cfgString(node.Config, "reviewer"), "human"), true), true, "simulated", "" + default: + return outputMap(envelope("unknown", node.ID, node.Type, "skipped", ""), map[string]any{"reason": "未知节点类型"}), true, "skipped", "未知节点类型" + } +} + +func buildReplayScript(trace []map[string]any) []map[string]any { + out := make([]map[string]any, 0, len(trace)) + for i, step := range trace { + raw, _ := json.Marshal(step["output"]) + out = append(out, map[string]any{ + "step": i + 1, + "nodeId": step["nodeId"], + "type": step["type"], + "status": step["status"], + "output": string(raw), + }) + } + return out +} diff --git a/internal/workflow/eino_branch.go b/internal/workflow/eino_branch.go new file mode 100644 index 00000000..1b3b9b13 --- /dev/null +++ b/internal/workflow/eino_branch.go @@ -0,0 +1,107 @@ +package workflow + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/compose" +) + +func hasConditionalOutgoingEdges(idx *graphIndex, nodeID string) bool { + for _, edge := range idx.outgoing[nodeID] { + cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression")) + if cond != "" { + return true + } + } + return false +} + +func wireConditionBranch( + wf *compose.Workflow[WorkflowInput, WorkflowOutput], + nodeRefs map[string]*compose.WorkflowNode, + idx *graphIndex, + condID string, + condNode graphNode, +) error { + edges := idx.outgoing[condID] + if len(edges) == 0 { + return nil + } + branchID := branchNodeID(condID) + wf.AddPassthroughNode(branchID).AddInput(condID) + + endNodes := map[string]bool{compose.END: true} + for _, edge := range edges { + endNodes[edge.Target] = true + } + + sortedEdges := append([]graphEdge(nil), edges...) + sortEdgesByCanvas(sortedEdges, idx.nodes) + + branch := compose.NewGraphBranch(func(runCtx context.Context, _ map[string]any) (string, error) { + rt := workflowRuntimeFrom(runCtx) + if rt == nil { + return compose.END, fmt.Errorf("workflow runtime missing in context") + } + emitConditionBranchProgress(rt.args, rt.runID, condNode, sortedEdges, idx.nodes, rt.state) + for edgeIdx, edge := range sortedEdges { + if conditionBranchAllowed(edge, edgeIdx, rt.state) { + return edge.Target, nil + } + } + return compose.END, nil + }, endNodes) + wf.AddBranch(branchID, branch) + + for _, edge := range edges { + if target, ok := nodeRefs[edge.Target]; ok { + target.AddInput(branchID) + } + } + return nil +} + +func wireEdgeConditionBranch( + wf *compose.Workflow[WorkflowInput, WorkflowOutput], + nodeRefs map[string]*compose.WorkflowNode, + idx *graphIndex, + sourceID string, + sourceNode graphNode, +) error { + edges := idx.outgoing[sourceID] + if len(edges) == 0 { + return nil + } + branchID := edgeBranchNodeID(sourceID) + wf.AddPassthroughNode(branchID).AddInput(sourceID) + + endNodes := map[string]bool{compose.END: true} + for _, edge := range edges { + endNodes[edge.Target] = true + } + + sortedEdges := append([]graphEdge(nil), edges...) + sortEdgesByCanvas(sortedEdges, idx.nodes) + + branch := compose.NewGraphBranch(func(runCtx context.Context, _ map[string]any) (string, error) { + rt := workflowRuntimeFrom(runCtx) + if rt == nil { + return compose.END, fmt.Errorf("workflow runtime missing in context") + } + for edgeIdx, edge := range sortedEdges { + if edgeAllowed(edge, sourceNode, edgeIdx, rt.state) { + return edge.Target, nil + } + } + return compose.END, nil + }, endNodes) + wf.AddBranch(branchID, branch) + + for _, edge := range edges { + if target, ok := nodeRefs[edge.Target]; ok { + target.AddInput(branchID) + } + } + return nil +} diff --git a/internal/workflow/eino_callbacks.go b/internal/workflow/eino_callbacks.go new file mode 100644 index 00000000..e4761c89 --- /dev/null +++ b/internal/workflow/eino_callbacks.go @@ -0,0 +1,22 @@ +package workflow + +import ( + "context" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/einoobserve" +) + +func attachWorkflowCallbacks(ctx context.Context, cfg *config.Config, args RunArgs, workflowName string) context.Context { + if cfg == nil { + return ctx + } + cbCfg := &cfg.MultiAgent.EinoCallbacks + return einoobserve.AttachAgentRunCallbacks(ctx, cbCfg, einoobserve.Params{ + Logger: args.Logger, + Progress: args.Progress, + ConversationID: args.ConversationID, + OrchMode: "workflow", + OrchestratorName: workflowName, + }) +} diff --git a/internal/workflow/eino_compile.go b/internal/workflow/eino_compile.go new file mode 100644 index 00000000..a2ca9881 --- /dev/null +++ b/internal/workflow/eino_compile.go @@ -0,0 +1,243 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/cloudwego/eino/compose" +) + +func executeEinoGraph(ctx context.Context, args RunArgs, runID string, workflowID string, version int, g *graphDef, state *WorkflowLocalState) error { + _, err := invokeEinoGraph(ctx, args, runID, workflowID, version, g, state, false) + return err +} + +func invokeEinoGraph(ctx context.Context, args RunArgs, runID string, workflowID string, version int, g *graphDef, state *WorkflowLocalState, resume bool) (bool, error) { + wfInput := workflowInputFromMap(state.Inputs) + if resume { + wfInput = WorkflowInput{} + } + rt := &workflowRuntime{ + args: args, + runID: runID, + idx: indexGraph(g), + state: state, + } + + art, err := defaultEngine.getOrCompile(ctx, workflowID, version, g) + if err != nil { + return false, fmt.Errorf("编译 Eino Workflow 失败: %w", err) + } + rt.idx = art.idx + + runCtx := withWorkflowRuntime(ctx, rt) + runCtx = attachWorkflowCallbacks(runCtx, args.AppCfg, args, workflowID) + + invokeOpts := []compose.Option{compose.WithCheckPointID(runID)} + for { + _, err = art.runnable.Invoke(runCtx, wfInput, invokeOpts...) + if err == nil { + return false, nil + } + if hitlErr := extractAwaitingHITL(err, art, runID, args, state); hitlErr != nil { + return true, hitlErr + } + return false, err + } +} + +func extractAwaitingHITL(err error, art *compiledArtifact, runID string, args RunArgs, state *WorkflowLocalState) error { + info, ok := compose.ExtractInterruptInfo(err) + if !ok || len(art.hitlIDs) == 0 { + return nil + } + nodeID := nextHITLNodeID(info, art.hitlIDs) + node := art.idx.nodes[nodeID] + if nodeID == "" { + return nil + } + prompt := resolveHITLPromptBinding(node.Config, state) + label := firstNonEmpty(node.Label, nodeID) + if args.DB != nil { + pending := map[string]any{ + "nodeId": nodeID, + "label": label, + "prompt": prompt, + "reviewer": cfgString(node.Config, "reviewer"), + "checkpointId": runID, + "interrupt": workflowInterruptMetadata(info), + "resumePayload": map[string]any{"approved": "bool", "comment": "string"}, + } + pendingJSON, _ := json.Marshal(pending) + _ = args.DB.SetWorkflowRunAwaitingHITL(runID, nodeID, string(pendingJSON)) + } + if args.Progress != nil { + args.Progress("workflow_hitl_waiting", fmt.Sprintf("等待人工确认：%s", label), map[string]any{ + "workflowRunId": runID, + "nodeId": nodeID, + "label": label, + "prompt": prompt, + "reviewer": cfgString(node.Config, "reviewer"), + "mode": "interactive", + "resumeApi": fmt.Sprintf("/api/workflows/runs/%s/resume", runID), + }) + } + return &AwaitingHITLError{ + RunID: runID, + NodeID: nodeID, + NodeLabel: label, + Prompt: prompt, + Reviewer: cfgString(node.Config, "reviewer"), + } +} + +func workflowInterruptMetadata(info *compose.InterruptInfo) map[string]any { + if info == nil { + return map[string]any{} + } + before := append([]string(nil), info.BeforeNodes...) + return map[string]any{ + "beforeNodes": before, + "resumeTarget": firstString(before), + "address": map[string]any{ + "kind": "compose_interrupt", + "beforeNodes": before, + "path": strings.Join(before, "/"), + }, + "raw": fmt.Sprintf("%+v", info), + } +} + +func firstString(values []string) string { + if len(values) == 0 { + return "" + } + return values[0] +} + +func nextHITLNodeID(info *compose.InterruptInfo, hitlIDs []string) string { + if info != nil && len(info.BeforeNodes) > 0 { + for _, id := range info.BeforeNodes { + for _, hitl := range hitlIDs { + if id == hitl { + return id + } + } + } + return info.BeforeNodes[0] + } + if len(hitlIDs) == 0 { + return "" + } + return hitlIDs[0] +} + +// ResumeWorkflowRun continues a run paused at HITL after human decision. +func ResumeWorkflowRun(ctx context.Context, args RunArgs, runID string, approved bool, comment string) (*RunResult, error) { + run, err := args.DB.GetWorkflowRun(runID) + if err != nil { + return nil, err + } + if run == nil { + return nil, fmt.Errorf("工作流运行不存在") + } + if run.Status != "awaiting_hitl" { + return nil, fmt.Errorf("工作流运行不在等待审批状态: %s", run.Status) + } + wf, err := args.DB.GetWorkflowDefinition(run.WorkflowID) + if err != nil || wf == nil { + return nil, fmt.Errorf("工作流定义不存在") + } + graph, err := parseGraph(wf.GraphJSON) + if err != nil { + return nil, err + } + + var input map[string]interface{} + _ = json.Unmarshal([]byte(run.InputJSON), &input) + state := newWorkflowLocalState(input, runID) + if state.Inputs == nil { + state.Inputs = map[string]any{} + } + state.Inputs["_hitl_approved"] = approved + state.Inputs["_hitl_comment"] = strings.TrimSpace(comment) + state.Inputs["_hitl_node_id"] = run.PendingHITLNodeID + + if !approved { + errText := strings.TrimSpace(comment) + if errText == "" { + errText = "人工审批拒绝" + } + _ = args.DB.FinishWorkflowRun(runID, "rejected", "", errText) + if args.Progress != nil { + args.Progress("workflow_hitl_rejected", fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", run.PendingHITLNodeID), map[string]interface{}{ + "workflowRunId": runID, + "nodeId": run.PendingHITLNodeID, + "comment": errText, + }) + } + return &RunResult{ + RunID: runID, + Response: fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", run.PendingHITLNodeID), + Status: "rejected", + }, nil + } + + if args.Progress != nil { + args.Progress("workflow_hitl_resumed", "人工审批已通过，继续执行", map[string]interface{}{ + "workflowRunId": runID, + "nodeId": run.PendingHITLNodeID, + "comment": strings.TrimSpace(comment), + }) + } + + _ = args.DB.SetWorkflowRunStatus(runID, "running") + resumeArgs := args + if strings.TrimSpace(resumeArgs.ConversationID) == "" { + resumeArgs.ConversationID = run.ConversationID + } + + awaiting, err := invokeEinoGraph(ctx, resumeArgs, runID, wf.ID, run.WorkflowVersion, graph, state, true) + if err != nil { + if IsAwaitingHITL(err) { + return &RunResult{ + RunID: runID, + Status: "awaiting_hitl", + Response: fmt.Sprintf("工作流在节点「%s」等待下一次人工确认。", err.(*AwaitingHITLError).NodeID), + AwaitingHITL: true, + }, nil + } + _ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error()) + return nil, err + } + _ = awaiting + + output := map[string]interface{}{ + "workflowId": wf.ID, + "workflowName": wf.Name, + "workflowVersion": wf.Version, + "workflowRunId": runID, + "status": "completed", + "outputs": state.Outputs, + "metrics": state.Metrics, + "executedNodes": state.Executed, + "skippedNodes": state.Skipped, + "engine": "eino_workflow", + } + outputJSON, _ := json.Marshal(output) + response := renderWorkflowResponse(args.Role.Name, wf.Name, wf.Version, runID, state) + _ = args.DB.FinishWorkflowRun(runID, "completed", string(outputJSON), "") + if args.Progress != nil { + args.Progress("workflow_done", fmt.Sprintf("流程「%s」运行完成", wf.Name), map[string]interface{}{ + "workflowRunId": runID, + "workflowId": wf.ID, + "outputs": state.Outputs, + "metrics": state.Metrics, + "response": response, + "engine": "eino_workflow", + }) + } + return &RunResult{Response: response, RunID: runID, Status: "completed"}, nil +} diff --git a/internal/workflow/eino_compile_test.go b/internal/workflow/eino_compile_test.go new file mode 100644 index 00000000..691bdbce --- /dev/null +++ b/internal/workflow/eino_compile_test.go @@ -0,0 +1,412 @@ +package workflow + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "github.com/cloudwego/eino/compose" + "go.uber.org/zap" +) + +func testWorkflowDB(t *testing.T) *database.DB { + t.Helper() + dir := t.TempDir() + db, err := database.NewDB(filepath.Join(dir, "workflow.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + return db +} + +func linearStartOutputGraph() string { + return `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 120}, "config": {"output_key": "result", "source_binding": {"from": "inputs", "field": "message"}}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "out-1"} + ], + "config": {"schema_version": 1} +}` +} + +func conditionBranchGraph() string { + return `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "cond-1", "type": "condition", "label": "判断", "position": {"x": 0, "y": 80}, "config": {"expression": "{{inputs.message}} == yes"}}, + {"id": "out-yes", "type": "output", "label": "是", "position": {"x": -80, "y": 160}, "config": {"output_key": "branch", "static_value": "yes"}}, + {"id": "out-no", "type": "output", "label": "否", "position": {"x": 80, "y": 160}, "config": {"output_key": "branch", "static_value": "no"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "cond-1"}, + {"id": "e2", "source": "cond-1", "target": "out-yes", "label": "是"}, + {"id": "e3", "source": "cond-1", "target": "out-no", "label": "否"} + ], + "config": {"schema_version": 1} +}` +} + +func TestValidateGraphJSON_linear(t *testing.T) { + if err := ValidateGraphJSON(context.Background(), linearStartOutputGraph()); err != nil { + t.Fatalf("validate: %v", err) + } +} + +func TestValidateGraphJSON_rejectsInvalidGraphs(t *testing.T) { + tests := []struct { + name string + graph string + wantErr string + }{ + { + name: "start with incoming edge", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "agent-1", "type": "agent", "label": "Agent", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop"}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "agent-1"}, + {"id": "e2", "source": "agent-1", "target": "start-1"}, + {"id": "e3", "source": "agent-1", "target": "out-1"} + ] +}`, + wantErr: "开始节点", + }, + { + name: "output with outgoing edge", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {"output_key": "result"}}, + {"id": "end-1", "type": "end", "label": "结束", "position": {"x": 0, "y": 160}, "config": {}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "out-1"}, + {"id": "e2", "source": "out-1", "target": "end-1"} + ] +}`, + wantErr: "不能有出边", + }, + { + name: "tool without name", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "tool-1", "type": "tool", "label": "工具", "position": {"x": 0, "y": 80}, "config": {}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "tool-1"}, + {"id": "e2", "source": "tool-1", "target": "out-1"} + ] +}`, + wantErr: "必须选择 MCP 工具", + }, + { + name: "condition with too many branches", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "cond-1", "type": "condition", "label": "判断", "position": {"x": 0, "y": 80}, "config": {"expression": "{{inputs.message}}"}}, + {"id": "out-1", "type": "output", "label": "输出1", "position": {"x": -80, "y": 160}, "config": {"output_key": "a"}}, + {"id": "out-2", "type": "output", "label": "输出2", "position": {"x": 0, "y": 160}, "config": {"output_key": "b"}}, + {"id": "out-3", "type": "output", "label": "输出3", "position": {"x": 80, "y": 160}, "config": {"output_key": "c"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "cond-1"}, + {"id": "e2", "source": "cond-1", "target": "out-1"}, + {"id": "e3", "source": "cond-1", "target": "out-2"}, + {"id": "e4", "source": "cond-1", "target": "out-3"} + ] +}`, + wantErr: "1 到 2 条出边", + }, + { + name: "orphan node", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {"output_key": "result"}}, + {"id": "agent-1", "type": "agent", "label": "孤岛", "position": {"x": 200, "y": 80}, "config": {"instruction": "noop"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "out-1"} + ] +}`, + wantErr: "不可达", + }, + { + name: "cycle", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "agent-1", "type": "agent", "label": "Agent1", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop", "output_key": "a1"}}, + {"id": "agent-2", "type": "agent", "label": "Agent2", "position": {"x": 0, "y": 160}, "config": {"instruction": "noop", "output_key": "a2"}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 240}, "config": {"output_key": "result"}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "agent-1"}, + {"id": "e2", "source": "agent-1", "target": "agent-2"}, + {"id": "e3", "source": "agent-2", "target": "agent-1"}, + {"id": "e4", "source": "agent-2", "target": "out-1"} + ] +}`, + wantErr: "环路", + }, + { + name: "output without key", + graph: `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 80}, "config": {}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "out-1"} + ] +}`, + wantErr: "输出变量名", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateGraphJSON(context.Background(), tt.graph) + if err == nil { + t.Fatal("expected validation error") + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %q, want substring %q", err.Error(), tt.wantErr) + } + }) + } +} + +func TestCompileEngine_linear(t *testing.T) { + ctx := context.Background() + SetCheckpointDir(t.TempDir()) + g, err := parseGraph(linearStartOutputGraph()) + if err != nil { + t.Fatal(err) + } + if _, err := defaultEngine.compile(ctx, g); err != nil { + t.Fatalf("compile: %v", err) + } +} + +func createTestWorkflowRun(t *testing.T, db *database.DB, runID string) { + t.Helper() + if err := db.CreateWorkflowRun(&database.WorkflowRun{ + ID: runID, + WorkflowID: "test-wf", + Status: "running", + }); err != nil { + t.Fatalf("CreateWorkflowRun: %v", err) + } +} + +func TestExecuteEinoGraph_linearStartOutput(t *testing.T) { + ctx := context.Background() + SetCheckpointDir(t.TempDir()) + db := testWorkflowDB(t) + createTestWorkflowRun(t, db, "run-linear") + g, err := parseGraph(linearStartOutputGraph()) + if err != nil { + t.Fatal(err) + } + state := newWorkflowLocalState(map[string]interface{}{"message": "ping"}, "run-linear") + args := RunArgs{DB: db} + if err := executeEinoGraph(ctx, args, "run-linear", "test-wf", 1, g, state); err != nil { + t.Fatalf("execute: %v", err) + } + if got := state.Outputs["result"]; got != "ping" { + t.Fatalf("outputs[result] = %v, want ping", got) + } + if len(state.Executed) != 2 { + t.Fatalf("executed nodes = %d, want 2", len(state.Executed)) + } +} + +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) { + ctx := context.Background() + SetCheckpointDir(t.TempDir()) + db := testWorkflowDB(t) + createTestWorkflowRun(t, db, "run-yes") + createTestWorkflowRun(t, db, "run-no") + g, err := parseGraph(conditionBranchGraph()) + if err != nil { + t.Fatal(err) + } + + stateYes := newWorkflowLocalState(map[string]interface{}{"message": "yes"}, "run-yes") + if err := executeEinoGraph(ctx, RunArgs{DB: db}, "run-yes", "test-wf-branch", 1, g, stateYes); err != nil { + t.Fatalf("execute yes: %v", err) + } + if got := stateYes.Outputs["branch"]; got != "yes" { + t.Fatalf("yes branch output = %v", got) + } + + stateNo := newWorkflowLocalState(map[string]interface{}{"message": "no"}, "run-no") + if err := executeEinoGraph(ctx, RunArgs{DB: db}, "run-no", "test-wf-branch", 1, g, stateNo); err != nil { + t.Fatalf("execute no: %v", err) + } + if got := stateNo.Outputs["branch"]; got != "no" { + t.Fatalf("no branch output = %v", got) + } +} + +func TestRunRoleBoundWorkflow_integration(t *testing.T) { + ctx := context.Background() + SetCheckpointDir(t.TempDir()) + db := testWorkflowDB(t) + graph := linearStartOutputGraph() + if err := db.UpsertWorkflowDefinition(&database.WorkflowDefinition{ + ID: "wf-linear", + Name: "线性流程", + Version: 1, + GraphJSON: graph, + Enabled: true, + }); err != nil { + t.Fatal(err) + } + role := config.RoleConfig{ + Name: "tester", + Enabled: true, + WorkflowID: "wf-linear", + WorkflowPolicy: "auto", + } + result, err := RunRoleBoundWorkflow(ctx, RunArgs{ + DB: db, + Logger: zap.NewNop(), + Role: role, + UserMessage: "from-role", + }) + if err != nil { + t.Fatalf("RunRoleBoundWorkflow: %v", err) + } + if result == nil || result.RunID == "" { + t.Fatal("expected run result") + } +} + +func TestCompiledCache_reuse(t *testing.T) { + ctx := context.Background() + SetCheckpointDir(t.TempDir()) + InvalidateCompiledCache("cache-wf") + g, err := parseGraph(linearStartOutputGraph()) + if err != nil { + t.Fatal(err) + } + a1, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g) + if err != nil { + t.Fatal(err) + } + a2, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g) + if err != nil { + t.Fatal(err) + } + if a1 != a2 { + t.Fatal("expected cached artifact pointer reuse") + } + InvalidateCompiledCache("cache-wf") + a3, err := defaultEngine.getOrCompile(ctx, "cache-wf", 1, g) + if err != nil { + t.Fatal(err) + } + if a1 == a3 { + t.Fatal("expected new artifact after invalidation") + } +} diff --git a/internal/workflow/eino_runtime.go b/internal/workflow/eino_runtime.go new file mode 100644 index 00000000..904b6c5a --- /dev/null +++ b/internal/workflow/eino_runtime.go @@ -0,0 +1,64 @@ +package workflow + +import ( + "context" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +type workflowRuntimeCtxKey struct{} + +// workflowRuntime carries per-run execution context into Eino Workflow local state. +type workflowRuntime struct { + args RunArgs + runID string + idx *graphIndex + state *WorkflowLocalState +} + +func withWorkflowRuntime(ctx context.Context, rt *workflowRuntime) context.Context { + return context.WithValue(ctx, workflowRuntimeCtxKey{}, rt) +} + +func workflowRuntimeFrom(ctx context.Context) *workflowRuntime { + rt, _ := ctx.Value(workflowRuntimeCtxKey{}).(*workflowRuntime) + return rt +} + +func newWorkflowRuntime(args RunArgs, runID string, idx *graphIndex, inputs map[string]interface{}) *workflowRuntime { + return &workflowRuntime{ + args: args, + runID: runID, + idx: idx, + state: newWorkflowLocalState(inputs, runID), + } +} + +// RunArgs is the execution context for a role-bound workflow run. +type RunArgs struct { + DB *database.DB + Logger *zap.Logger + Role config.RoleConfig + AppCfg *config.Config + Agent *agent.Agent + ConversationID string + ProjectID string + UserMessage string + History []agent.ChatMessage + RoleTools []string + AgentsMarkdownDir string + SystemPromptExtra string + AssistantMessageID string + Progress agent.ProgressCallback +} + +type RunResult struct { + Response string + RunID string + Status string + AwaitingHITL bool +} diff --git a/internal/workflow/engine.go b/internal/workflow/engine.go new file mode 100644 index 00000000..888ca16d --- /dev/null +++ b/internal/workflow/engine.go @@ -0,0 +1,239 @@ +package workflow + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/cloudwego/eino/compose" +) + +type compiledArtifact struct { + runnable compose.Runnable[WorkflowInput, WorkflowOutput] + idx *graphIndex + hitlIDs []string +} + +// Engine compiles and caches Eino Workflow artifacts. +type Engine struct { + mu sync.RWMutex + cache map[string]*compiledArtifact + cpStore compose.CheckPointStore + cpStoreMu sync.Once + cpStoreErr error + checkpointDir string +} + +var defaultEngine = &Engine{ + cache: make(map[string]*compiledArtifact), + checkpointDir: "data/workflow-checkpoints", +} + +// SetCheckpointDir overrides the workflow checkpoint root (mainly for tests). +func SetCheckpointDir(dir string) { + defaultEngine.mu.Lock() + defer defaultEngine.mu.Unlock() + defaultEngine.checkpointDir = strings.TrimSpace(dir) + defaultEngine.cpStore = nil + defaultEngine.cpStoreErr = nil + defaultEngine.cpStoreMu = sync.Once{} +} + +func (e *Engine) checkpointStore() (compose.CheckPointStore, error) { + e.cpStoreMu.Do(func() { + e.cpStore, e.cpStoreErr = newFileCheckPointStore(e.checkpointDir) + }) + return e.cpStore, e.cpStoreErr +} + +// InvalidateCompiledCache drops cached compilations for a workflow id. +func InvalidateCompiledCache(workflowID string) { + workflowID = strings.TrimSpace(workflowID) + if workflowID == "" { + return + } + defaultEngine.mu.Lock() + defer defaultEngine.mu.Unlock() + for key := range defaultEngine.cache { + if strings.HasPrefix(key, workflowID+":") { + delete(defaultEngine.cache, key) + } + } +} + +// ValidateGraphJSON parses and trial-compiles a canvas graph (save-time gate). +func ValidateGraphJSON(ctx context.Context, graphJSON string) error { + g, err := parseGraph(graphJSON) + if err != nil { + return err + } + idx := indexGraph(g) + if err := validateGraphDefinition(g, idx); err != nil { + return err + } + _, err = defaultEngine.compile(ctx, g) + return err +} + +func hasTerminalNode(idx *graphIndex) bool { + for id, node := range idx.nodes { + if len(idx.outgoing[id]) == 0 { + return true + } + if strings.EqualFold(node.Type, "end") || strings.EqualFold(node.Type, "output") { + return true + } + } + return false +} + +func (e *Engine) getOrCompile(ctx context.Context, workflowID string, version int, g *graphDef) (*compiledArtifact, error) { + key := cacheKey(workflowID, version) + e.mu.RLock() + if art, ok := e.cache[key]; ok { + e.mu.RUnlock() + return art, nil + } + e.mu.RUnlock() + + art, err := e.compile(ctx, g) + if err != nil { + return nil, err + } + e.mu.Lock() + if existing, ok := e.cache[key]; ok { + e.mu.Unlock() + return existing, nil + } + e.cache[key] = art + e.mu.Unlock() + return art, nil +} + +func (e *Engine) compile(ctx context.Context, g *graphDef) (*compiledArtifact, error) { + cpStore, err := e.checkpointStore() + if err != nil { + return nil, err + } + idx := indexGraph(g) + if err := validateGraphDefinition(g, idx); err != nil { + return nil, err + } + hitlIDs := collectHITLNodeIDs(idx) + compileOpts := []compose.GraphCompileOption{ + compose.WithGraphName("CyberStrikeWorkflow"), + compose.WithCheckPointStore(cpStore), + } + if len(hitlIDs) > 0 { + compileOpts = append(compileOpts, compose.WithInterruptBeforeNodes(hitlIDs)) + } + + wf := compose.NewWorkflow[WorkflowInput, WorkflowOutput]( + compose.WithGenLocalState(func(runCtx context.Context) *WorkflowLocalState { + if rt := workflowRuntimeFrom(runCtx); rt != nil && rt.state != nil { + return rt.state + } + return &WorkflowLocalState{ + Outputs: make(map[string]any), + NodeOutputs: make(map[string]map[string]any), + NodeProceed: make(map[string]bool), + } + }), + ) + + nodeRefs := make(map[string]*compose.WorkflowNode, len(idx.nodes)) + for id, node := range idx.nodes { + n := node + if strings.EqualFold(n.Type, "agent") { + sub, err := compileAgentSubgraph(ctx, n) + if err != nil { + return nil, fmt.Errorf("编译 Agent 子图 %s 失败: %w", id, err) + } + nodeRefs[id] = wf.AddGraphNode(id, sub) + continue + } + if strings.EqualFold(n.Type, "start") { + nodeRefs[id] = wf.AddLambdaNode(id, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowInput) (WorkflowNodeOutput, error) { + return runWorkflowNodeLambda(runCtx, n) + })) + continue + } + nodeRefs[id] = wf.AddLambdaNode(id, compose.InvokableLambda(func(runCtx context.Context, _ WorkflowNodeOutput) (WorkflowNodeOutput, error) { + return runWorkflowNodeLambda(runCtx, n) + })) + } + + for id, node := range idx.nodes { + if strings.EqualFold(node.Type, "condition") { + if err := wireConditionBranch(wf, nodeRefs, idx, id, node); err != nil { + return nil, err + } + continue + } + if hasConditionalOutgoingEdges(idx, id) { + if err := wireEdgeConditionBranch(wf, nodeRefs, idx, id, node); err != nil { + return nil, err + } + continue + } + for _, edge := range idx.outgoing[id] { + if target, ok := nodeRefs[edge.Target]; ok { + target.AddInput(id) + } + } + } + + for _, startID := range findStartNodeIDs(idx) { + if ref, ok := nodeRefs[startID]; ok { + ref.AddInput(compose.START) + } + } + + endNode := wf.End() + for id, node := range idx.nodes { + if len(idx.outgoing[id]) == 0 || strings.EqualFold(node.Type, "end") { + endNode.AddInput(id, compose.ToField(id)) + } + } + + runnable, err := wf.Compile(ctx, compileOpts...) + if err != nil { + return nil, err + } + return &compiledArtifact{runnable: runnable, idx: idx, hitlIDs: hitlIDs}, nil +} + +func collectHITLNodeIDs(idx *graphIndex) []string { + var ids []string + for id, node := range idx.nodes { + if strings.EqualFold(node.Type, "hitl") { + ids = append(ids, id) + } + } + return ids +} + +func runWorkflowNodeLambda(runCtx context.Context, n graphNode) (WorkflowNodeOutput, error) { + localRT := workflowRuntimeFrom(runCtx) + if localRT == nil { + return nil, fmt.Errorf("workflow runtime missing in context") + } + if err := prepareNodeInputState(localRT, n); err != nil { + return nil, err + } + result, proceed, err := executeNode(runCtx, localRT.args, localRT.runID, n, localRT.state) + if err != nil { + return nil, err + } + localRT.state.NodeOutputs[n.ID] = result + localRT.state.LastOutput = result + if !proceed && !strings.EqualFold(n.Type, "end") { + label := firstNonEmpty(n.Label, n.ID) + if errText := cfgString(result, "error"); errText != "" { + return result, fmt.Errorf("节点「%s」失败: %s", label, errText) + } + return result, fmt.Errorf("节点「%s」未继续执行", label) + } + return result, nil +} diff --git a/internal/workflow/errors.go b/internal/workflow/errors.go new file mode 100644 index 00000000..840b9b0e --- /dev/null +++ b/internal/workflow/errors.go @@ -0,0 +1,24 @@ +package workflow + +import "errors" + +// AwaitingHITLError indicates the workflow paused before a HITL node for human approval. +type AwaitingHITLError struct { + RunID string + NodeID string + NodeLabel string + Prompt string + Reviewer string +} + +func (e *AwaitingHITLError) Error() string { + if e == nil { + return "workflow awaiting human approval" + } + return "workflow awaiting human approval at node " + e.NodeID +} + +func IsAwaitingHITL(err error) bool { + var target *AwaitingHITLError + return errors.As(err, &target) +} diff --git a/internal/workflow/expression.go b/internal/workflow/expression.go new file mode 100644 index 00000000..4a3d2b6d --- /dev/null +++ b/internal/workflow/expression.go @@ -0,0 +1,186 @@ +package workflow + +import ( + "fmt" + "regexp" + "strconv" + "strings" +) + +var expressionOps = []string{">=", "<=", "==", "!=", " contains ", " matches ", ">", "<"} +var jsonFuncRe = regexp.MustCompile(`^(jsonpath|jq)\((.*),\s*(['"][^'"]+['"])\)$`) +var jsonFuncFindRe = regexp.MustCompile(`(jsonpath|jq)\([^)]*\)`) +var singleTemplateVarRe = regexp.MustCompile(`^\{\{\s*([a-zA-Z0-9_.-]+)\s*\}\}$`) + +func validateConditionExpression(expr string) error { + expr = strings.TrimSpace(expr) + if expr == "" { + return fmt.Errorf("条件表达式不能为空") + } + for _, part := range splitBoolExpr(expr, "||") { + for _, atom := range splitBoolExpr(part, "&&") { + if err := validateConditionAtom(atom); err != nil { + return err + } + } + } + return nil +} + +func validateConditionAtom(expr string) error { + expr = strings.TrimSpace(expr) + if expr == "" { + return fmt.Errorf("条件表达式存在空片段") + } + if strings.Count(expr, "{{") != strings.Count(expr, "}}") { + return fmt.Errorf("条件表达式模板括号不匹配: %s", expr) + } + if err := validateJSONFunctions(expr); err != nil { + return err + } + if left, right, ok := splitExpressionAtom(expr, " matches "); ok { + if strings.TrimSpace(left) == "" || strings.TrimSpace(right) == "" { + return fmt.Errorf("matches 表达式两侧不能为空: %s", expr) + } + pattern := cleanComparable(resolveStaticTemplate(right)) + if _, err := regexp.Compile(pattern); err != nil { + return fmt.Errorf("matches 正则非法: %w", err) + } + return nil + } + for _, op := range expressionOps { + if op == " matches " { + continue + } + if left, right, ok := splitExpressionAtom(expr, op); ok { + if strings.TrimSpace(left) == "" || strings.TrimSpace(right) == "" { + return fmt.Errorf("表达式 %q 两侧不能为空: %s", strings.TrimSpace(op), expr) + } + return nil + } + } + return nil +} + +func evalCondition(expr string, state *WorkflowLocalState) bool { + expr = strings.TrimSpace(expr) + if expr == "" { + return true + } + orParts := splitBoolExpr(expr, "||") + for _, orPart := range orParts { + andOK := true + for _, atom := range splitBoolExpr(orPart, "&&") { + if !evalConditionAtom(atom, state) { + andOK = false + break + } + } + if andOK { + return true + } + } + return false +} + +func evalConditionAtom(expr string, state *WorkflowLocalState) bool { + expr = strings.TrimSpace(expr) + for _, op := range expressionOps { + if left, right, ok := splitExpressionAtom(expr, op); ok { + left = strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(left, state))) + right = strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(right, state))) + switch strings.TrimSpace(op) { + case "==": + return cleanComparable(left) == cleanComparable(right) + case "!=": + return cleanComparable(left) != cleanComparable(right) + case ">": + return compareNumeric(left, right, func(a, b float64) bool { return a > b }) + case ">=": + return compareNumeric(left, right, func(a, b float64) bool { return a >= b }) + case "<": + return compareNumeric(left, right, func(a, b float64) bool { return a < b }) + case "<=": + return compareNumeric(left, right, func(a, b float64) bool { return a <= b }) + case "contains": + return strings.Contains(cleanComparable(left), cleanComparable(right)) + case "matches": + matched, _ := regexp.MatchString(cleanComparable(right), cleanComparable(left)) + return matched + } + } + } + resolved := strings.TrimSpace(fmt.Sprint(resolveExpressionOperand(expr, state))) + v := strings.ToLower(cleanComparable(resolved)) + return v != "" && v != "false" && v != "0" && v != "null" +} + +func splitBoolExpr(expr, sep string) []string { + parts := strings.Split(expr, sep) + out := make([]string, 0, len(parts)) + for _, part := range parts { + if s := strings.TrimSpace(part); s != "" { + out = append(out, s) + } + } + if len(out) == 0 { + return []string{strings.TrimSpace(expr)} + } + return out +} + +func splitExpressionAtom(expr, op string) (string, string, bool) { + if strings.TrimSpace(op) == "contains" || strings.TrimSpace(op) == "matches" { + idx := strings.Index(expr, op) + if idx < 0 { + return "", "", false + } + return expr[:idx], expr[idx+len(op):], true + } + idx := strings.Index(expr, op) + if idx < 0 { + return "", "", false + } + return expr[:idx], expr[idx+len(op):], true +} + +func compareNumeric(left, right string, cmp func(float64, float64) bool) bool { + a, errA := strconv.ParseFloat(cleanComparable(left), 64) + b, errB := strconv.ParseFloat(cleanComparable(right), 64) + if errA != nil || errB != nil { + return false + } + return cmp(a, b) +} + +func resolveStaticTemplate(s string) string { + return templateVarRe.ReplaceAllString(s, "value") +} + +func resolveExpressionOperand(raw string, state *WorkflowLocalState) any { + raw = strings.TrimSpace(raw) + if m := jsonFuncRe.FindStringSubmatch(raw); len(m) == 4 { + inputExpr := strings.TrimSpace(m[2]) + path := strings.Trim(m[3], `"'`) + input := resolveExpressionOperand(inputExpr, state) + return evalJSONPathValue(input, path) + } + if m := singleTemplateVarRe.FindStringSubmatch(raw); len(m) == 2 { + return valueFromPath(m[1], state) + } + return resolveTemplate(raw, state) +} + +func validateJSONFunctions(expr string) error { + for _, candidate := range jsonFuncFindRe.FindAllString(expr, -1) { + candidate = strings.TrimSpace(candidate) + m := jsonFuncRe.FindStringSubmatch(candidate) + if len(m) != 4 { + return fmt.Errorf("JSONPath/JQ 函数格式应为 jsonpath(value, \"$.path\") 或 jq(value, \".path\")") + } + if err := validateJSONPathSyntax(strings.Trim(m[3], `"'`)); err != nil { + return err + } + } + return nil +} diff --git a/internal/workflow/expression_join_test.go b/internal/workflow/expression_join_test.go new file mode 100644 index 00000000..325a9159 --- /dev/null +++ b/internal/workflow/expression_join_test.go @@ -0,0 +1,107 @@ +package workflow + +import ( + "context" + "testing" +) + +func TestEvalCondition_extendedOperators(t *testing.T) { + state := newWorkflowLocalState(map[string]interface{}{"score": 9, "message": "status: ok"}, "run-expr") + state.LastOutput = map[string]any{"output": "asset-123.example.com"} + + tests := []string{ + "{{inputs.score}} >= 9", + "{{inputs.message}} contains ok", + "{{previous.output}} matches ^asset-[0-9]+\\.example\\.com$", + "{{inputs.score}} > 5 && {{inputs.message}} contains status", + } + for _, expr := range tests { + if err := validateConditionExpression(expr); err != nil { + t.Fatalf("validate %q: %v", expr, err) + } + if !evalCondition(expr, state) { + t.Fatalf("evalCondition(%q) = false, want true", expr) + } + } +} + +func TestEvalCondition_jsonPathAndJQSafeSubset(t *testing.T) { + state := newWorkflowLocalState(map[string]interface{}{ + "payload": map[string]any{ + "risk": 9, + "items": []any{ + map[string]any{"name": "first"}, + }, + }, + }, "run-jsonpath") + state.LastOutput = map[string]any{"output": `{"status":"ok","score":7}`} + + tests := []string{ + `jsonpath({{inputs.payload}}, "$.risk") >= 8`, + `jq({{inputs.payload}}, ".items[0].name") == first`, + `jsonpath({{previous.output}}, "$.status") == ok`, + } + for _, expr := range tests { + if err := validateConditionExpression(expr); err != nil { + t.Fatalf("validate %q: %v", expr, err) + } + if !evalCondition(expr, state) { + t.Fatalf("evalCondition(%q) = false, want true", expr) + } + } +} + +func TestMergeUpstreamOutputs_allMerge(t *testing.T) { + got := mergeUpstreamOutputs(JoinAllMerge, []map[string]any{ + {"output": "a", "left": 1}, + {"output": "b", "right": 2}, + }) + if got["kind"] != "join" || got["strategy"] != JoinAllMerge { + t.Fatalf("join metadata = %#v", got) + } + values, ok := got["output"].([]any) + if !ok || len(values) != 2 || values[0] != "a" || values[1] != "b" { + t.Fatalf("merged output = %#v", got["output"]) + } + if got["left"] != 1 || got["right"] != 2 { + t.Fatalf("merged fields = %#v", got) + } +} + +func TestMergeUpstreamOutputs_firstNonEmpty(t *testing.T) { + got := mergeUpstreamOutputs(JoinFirstNonEmpty, []map[string]any{ + {"output": ""}, + {"output": "winner"}, + }) + if got["output"] != "winner" { + t.Fatalf("output = %#v, want winner", got["output"]) + } +} + +func TestDryRunGraphJSON_simulatesUnsafeNodes(t *testing.T) { + graph := `{ + "nodes": [ + {"id": "start-1", "type": "start", "label": "开始", "position": {"x": 0, "y": 0}, "config": {}}, + {"id": "agent-1", "type": "agent", "label": "Agent", "position": {"x": 0, "y": 80}, "config": {"instruction": "noop", "output_key": "agent_result"}}, + {"id": "out-1", "type": "output", "label": "输出", "position": {"x": 0, "y": 160}, "config": {"output_key": "result", "source_binding": {"from": "outputs", "field": "agent_result"}}} + ], + "edges": [ + {"id": "e1", "source": "start-1", "target": "agent-1"}, + {"id": "e2", "source": "agent-1", "target": "out-1"} + ] +}` + result, err := DryRunGraphJSON(nilContext(), graph, map[string]any{"message": "hello"}) + if err != nil { + t.Fatalf("DryRunGraphJSON: %v", err) + } + if got := result.Outputs["result"]; got != "[dry-run] agent execution skipped" { + t.Fatalf("result output = %#v", got) + } + if len(result.Trace) != 3 { + t.Fatalf("trace len = %d, want 3", len(result.Trace)) + } +} + +func nilContext() context.Context { + return context.Background() +} diff --git a/internal/workflow/graph_types.go b/internal/workflow/graph_types.go new file mode 100644 index 00000000..8e49be11 --- /dev/null +++ b/internal/workflow/graph_types.go @@ -0,0 +1,153 @@ +package workflow + +import ( + "encoding/json" + "fmt" + "sort" + "strings" +) + +type graphDef struct { + Nodes []graphNode `json:"nodes"` + Edges []graphEdge `json:"edges"` + Config map[string]any `json:"config"` +} + +type graphNode struct { + ID string `json:"id"` + Type string `json:"type"` + Label string `json:"label"` + Position graphPosition `json:"position"` + Config map[string]any `json:"config"` +} + +type graphEdge struct { + ID string `json:"id"` + Source string `json:"source"` + Target string `json:"target"` + Label string `json:"label"` + Config map[string]any `json:"config"` +} + +type graphPosition struct { + X float64 `json:"x"` + Y float64 `json:"y"` +} + +type graphIndex struct { + nodes map[string]graphNode + outgoing map[string][]graphEdge + incoming map[string][]graphEdge +} + +func parseGraph(raw string) (*graphDef, error) { + var g graphDef + if err := json.Unmarshal([]byte(strings.TrimSpace(raw)), &g); err != nil { + return nil, fmt.Errorf("解析工作流图失败: %w", err) + } + if len(g.Nodes) == 0 { + return nil, fmt.Errorf("工作流没有节点") + } + if g.Config == nil { + g.Config = make(map[string]any) + } + return &g, nil +} + +func indexGraph(g *graphDef) *graphIndex { + idx := &graphIndex{ + nodes: make(map[string]graphNode, len(g.Nodes)), + outgoing: make(map[string][]graphEdge), + incoming: make(map[string][]graphEdge), + } + for _, node := range g.Nodes { + node.ID = strings.TrimSpace(node.ID) + if node.ID == "" { + continue + } + if strings.TrimSpace(node.Type) == "" { + node.Type = "tool" + } + if node.Config == nil { + node.Config = make(map[string]any) + } + idx.nodes[node.ID] = node + } + for _, edge := range g.Edges { + if _, ok := idx.nodes[edge.Source]; !ok { + continue + } + if _, ok := idx.nodes[edge.Target]; !ok { + continue + } + idx.outgoing[edge.Source] = append(idx.outgoing[edge.Source], edge) + idx.incoming[edge.Target] = append(idx.incoming[edge.Target], edge) + } + for source := range idx.outgoing { + sortEdgesByCanvas(idx.outgoing[source], idx.nodes) + } + return idx +} + +func sortEdgesByCanvas(edges []graphEdge, nodes map[string]graphNode) { + sort.SliceStable(edges, func(i, j int) bool { + a := nodes[edges[i].Target] + b := nodes[edges[j].Target] + if a.Position.Y != b.Position.Y { + return a.Position.Y < b.Position.Y + } + if a.Position.X != b.Position.X { + return a.Position.X < b.Position.X + } + return edges[i].Target < edges[j].Target + }) +} + +func sortNodeIDsByCanvas(ids []string, nodes map[string]graphNode) { + sort.SliceStable(ids, func(i, j int) bool { + a := nodes[ids[i]] + b := nodes[ids[j]] + if a.Position.Y != b.Position.Y { + return a.Position.Y < b.Position.Y + } + if a.Position.X != b.Position.X { + return a.Position.X < b.Position.X + } + return ids[i] < ids[j] + }) +} + +func findStartNodeIDs(idx *graphIndex) []string { + var queue []string + for id, node := range idx.nodes { + if strings.EqualFold(node.Type, "start") { + queue = append(queue, id) + } + } + if len(queue) == 0 { + inDegree := make(map[string]int, len(idx.nodes)) + for id := range idx.nodes { + inDegree[id] = 0 + } + for _, edges := range idx.outgoing { + for _, edge := range edges { + inDegree[edge.Target]++ + } + } + for id, deg := range inDegree { + if deg == 0 { + queue = append(queue, id) + } + } + } + sortNodeIDsByCanvas(queue, idx.nodes) + return queue +} + +func branchNodeID(nodeID string) string { + return nodeID + "__eino_branch" +} + +func edgeBranchNodeID(nodeID string) string { + return nodeID + "__eino_edge_branch" +} diff --git a/internal/workflow/hitl_wait.go b/internal/workflow/hitl_wait.go new file mode 100644 index 00000000..624cdf5b --- /dev/null +++ b/internal/workflow/hitl_wait.go @@ -0,0 +1,119 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" +) + +// HITLDecision is a human decision on a workflow approval node. +type HITLDecision struct { + Approved bool + Comment string +} + +var hitlWaiters sync.Map // runID -> chan HITLDecision + +func registerHITLWaiter(runID string) chan HITLDecision { + ch := make(chan HITLDecision, 1) + hitlWaiters.Store(runID, ch) + return ch +} + +func unregisterHITLWaiter(runID string, ch chan HITLDecision) { + hitlWaiters.CompareAndDelete(runID, ch) +} + +// NotifyHITLDecision wakes a streaming workflow run waiting at a HITL node. +// Returns true when an active waiter was signaled. +func NotifyHITLDecision(runID string, decision HITLDecision) bool { + v, ok := hitlWaiters.Load(runID) + if !ok { + return false + } + ch, ok := v.(chan HITLDecision) + if !ok { + return false + } + select { + case ch <- decision: + return true + default: + return true + } +} + +func readHITLDecisionFromDB(db *database.DB, runID string) (HITLDecision, bool, error) { + if db == nil { + return HITLDecision{}, false, nil + } + run, err := db.GetWorkflowRun(runID) + if err != nil { + return HITLDecision{}, false, err + } + if run == nil || strings.TrimSpace(run.PendingHITLJSON) == "" { + return HITLDecision{}, false, nil + } + var pending map[string]interface{} + if err := json.Unmarshal([]byte(run.PendingHITLJSON), &pending); err != nil { + return HITLDecision{}, false, nil + } + raw, ok := pending["decision"] + if !ok { + return HITLDecision{}, false, nil + } + decision := strings.ToLower(strings.TrimSpace(fmt.Sprint(raw))) + switch decision { + case "approved", "approve": + comment := "" + if v, ok := pending["comment"]; ok { + comment = strings.TrimSpace(fmt.Sprint(v)) + } + return HITLDecision{Approved: true, Comment: comment}, true, nil + case "rejected", "reject": + comment := "" + if v, ok := pending["comment"]; ok { + comment = strings.TrimSpace(fmt.Sprint(v)) + } + return HITLDecision{Approved: false, Comment: comment}, true, nil + default: + return HITLDecision{}, false, nil + } +} + +func waitWorkflowHITLDecision(ctx context.Context, db *database.DB, runID string) (HITLDecision, error) { + ch := registerHITLWaiter(runID) + defer unregisterHITLWaiter(runID, ch) + return waitWorkflowHITLDecisionWithChannel(ctx, db, runID, ch) +} + +func waitWorkflowHITLDecisionWithChannel(ctx context.Context, db *database.DB, runID string, ch chan HITLDecision) (HITLDecision, error) { + if d, ok, err := readHITLDecisionFromDB(db, runID); err != nil { + return HITLDecision{}, err + } else if ok { + return d, nil + } + + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return HITLDecision{}, ctx.Err() + case d := <-ch: + return d, nil + case <-ticker.C: + if d, ok, err := readHITLDecisionFromDB(db, runID); err != nil { + return HITLDecision{}, err + } else if ok { + return d, nil + } + } + } +} diff --git a/internal/workflow/join.go b/internal/workflow/join.go new file mode 100644 index 00000000..2d720ef5 --- /dev/null +++ b/internal/workflow/join.go @@ -0,0 +1,117 @@ +package workflow + +import ( + "fmt" + "strings" +) + +const ( + JoinAllMerge = "all_merge" + JoinLastByCanvas = "last_by_canvas" + JoinFirstNonEmpty = "first_non_empty" + JoinFailFast = "fail_fast" +) + +var allowedJoinStrategies = map[string]bool{ + JoinAllMerge: true, + JoinLastByCanvas: true, + JoinFirstNonEmpty: true, + JoinFailFast: true, +} + +func joinStrategy(node graphNode) string { + strategy := strings.ToLower(cfgString(node.Config, "join_strategy")) + if strategy == "" { + return JoinAllMerge + } + return strategy +} + +func prepareNodeInputState(rt *workflowRuntime, node graphNode) error { + if rt == nil || rt.idx == nil || rt.state == nil { + return nil + } + incoming := rt.idx.incoming[node.ID] + if len(incoming) <= 1 { + return nil + } + strategy := joinStrategy(node) + if !allowedJoinStrategies[strategy] { + return fmt.Errorf("节点「%s」使用了未知汇聚策略: %s", firstNonEmpty(node.Label, node.ID), strategy) + } + upstreams := make([]map[string]any, 0, len(incoming)) + for _, edge := range incoming { + out := rt.state.NodeOutputs[edge.Source] + if out == nil { + continue + } + if isFailedNodeOutput(out) && strategy == JoinFailFast { + return fmt.Errorf("上游节点「%s」失败，汇聚策略 fail_fast 中止", edge.Source) + } + upstreams = append(upstreams, out) + } + if len(upstreams) == 0 { + return nil + } + rt.state.LastOutput = mergeUpstreamOutputs(strategy, upstreams) + return nil +} + +func mergeUpstreamOutputs(strategy string, upstreams []map[string]any) map[string]any { + switch strategy { + case JoinLastByCanvas: + return cloneNodeOutput(upstreams[len(upstreams)-1]) + case JoinFirstNonEmpty: + for _, out := range upstreams { + if !isEmptyOutputValue(out["output"]) { + return cloneNodeOutput(out) + } + } + return cloneNodeOutput(upstreams[0]) + default: + merged := map[string]any{ + "kind": "join", + "strategy": strategy, + "upstreams": upstreams, + } + values := make([]any, 0, len(upstreams)) + for _, out := range upstreams { + values = append(values, out["output"]) + for k, v := range out { + if _, exists := merged[k]; !exists { + merged[k] = v + } + } + } + merged["output"] = values + return merged + } +} + +func cloneNodeOutput(in map[string]any) map[string]any { + out := make(map[string]any, len(in)) + for k, v := range in { + out[k] = v + } + return out +} + +func isEmptyOutputValue(v any) bool { + if v == nil { + return true + } + return strings.TrimSpace(fmt.Sprint(v)) == "" +} + +func isFailedNodeOutput(out map[string]any) bool { + if out == nil { + return false + } + if v, ok := out["error"]; ok && strings.TrimSpace(fmt.Sprint(v)) != "" { + return true + } + if v, ok := out["is_error"]; ok { + return strings.EqualFold(fmt.Sprint(v), "true") + } + return false +} diff --git a/internal/workflow/jsonpath.go b/internal/workflow/jsonpath.go new file mode 100644 index 00000000..c379d299 --- /dev/null +++ b/internal/workflow/jsonpath.go @@ -0,0 +1,115 @@ +package workflow + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" +) + +func evalJSONPathValue(input any, path string) any { + path = strings.TrimSpace(path) + if path == "" || path == "$" || path == "." { + return input + } + if strings.HasPrefix(path, "$.") { + path = strings.TrimPrefix(path, "$.") + } else if strings.HasPrefix(path, ".") { + path = strings.TrimPrefix(path, ".") + } else if strings.HasPrefix(path, "$") { + path = strings.TrimPrefix(path, "$") + } + cur := normalizeJSONInput(input) + for _, token := range parseJSONPathTokens(path) { + if token == "" { + continue + } + switch v := cur.(type) { + case map[string]any: + cur = v[token] + case []any: + idx, err := strconv.Atoi(token) + if err != nil || idx < 0 || idx >= len(v) { + return "" + } + cur = v[idx] + default: + return "" + } + } + if cur == nil { + return "" + } + return cur +} + +func normalizeJSONInput(input any) any { + switch v := input.(type) { + case string: + var decoded any + if err := json.Unmarshal([]byte(v), &decoded); err == nil { + return decoded + } + return v + case []byte: + var decoded any + if err := json.Unmarshal(v, &decoded); err == nil { + return decoded + } + return string(v) + default: + return input + } +} + +func parseJSONPathTokens(path string) []string { + var tokens []string + var buf strings.Builder + for i := 0; i < len(path); i++ { + ch := path[i] + switch ch { + case '.': + if buf.Len() > 0 { + tokens = append(tokens, buf.String()) + buf.Reset() + } + case '[': + if buf.Len() > 0 { + tokens = append(tokens, buf.String()) + buf.Reset() + } + j := i + 1 + for j < len(path) && path[j] != ']' { + j++ + } + if j <= len(path) { + token := strings.Trim(path[i+1:j], `"' `) + tokens = append(tokens, token) + i = j + } + default: + buf.WriteByte(ch) + } + } + if buf.Len() > 0 { + tokens = append(tokens, buf.String()) + } + return tokens +} + +func validateJSONPathSyntax(path string) error { + path = strings.TrimSpace(path) + if path == "" { + return fmt.Errorf("JSONPath 不能为空") + } + if !strings.HasPrefix(path, "$") && !strings.HasPrefix(path, ".") { + return fmt.Errorf("JSONPath/JQ 路径必须以 $ 或 . 开头") + } + if strings.Contains(path, "..") || strings.ContainsAny(path, "*?()|") { + return fmt.Errorf("仅支持安全路径子集，不支持通配符、递归或表达式") + } + if strings.Count(path, "[") != strings.Count(path, "]") { + return fmt.Errorf("JSONPath 方括号不匹配") + } + return nil +} diff --git a/internal/workflow/metrics.go b/internal/workflow/metrics.go new file mode 100644 index 00000000..9eb8285a --- /dev/null +++ b/internal/workflow/metrics.go @@ -0,0 +1,57 @@ +package workflow + +import ( + "fmt" + "strconv" +) + +func accumulateWorkflowMetric(state *WorkflowLocalState, key string, delta any) { + if state == nil { + return + } + if state.Metrics == nil { + state.Metrics = make(map[string]any) + } + current := numericMetric(state.Metrics[key]) + state.Metrics[key] = current + numericMetric(delta) +} + +func numericMetric(v any) float64 { + switch n := v.(type) { + case int: + return float64(n) + case int32: + return float64(n) + case int64: + return float64(n) + case float32: + return float64(n) + case float64: + return n + case string: + f, _ := strconv.ParseFloat(n, 64) + return f + default: + f, _ := strconv.ParseFloat(fmt.Sprint(v), 64) + return f + } +} + +func collectAgentMetrics(state *WorkflowLocalState, data interface{}) { + m, ok := data.(map[string]interface{}) + if !ok || state == nil { + return + } + for _, key := range []string{"prompt_tokens", "completion_tokens", "total_tokens", "cost", "input_tokens", "output_tokens"} { + if v, ok := m[key]; ok { + accumulateWorkflowMetric(state, key, v) + } + } + if usage, ok := m["usage"].(map[string]interface{}); ok { + for _, key := range []string{"prompt_tokens", "completion_tokens", "total_tokens", "input_tokens", "output_tokens"} { + if v, ok := usage[key]; ok { + accumulateWorkflowMetric(state, key, v) + } + } + } +} diff --git a/internal/workflow/node_exec.go b/internal/workflow/node_exec.go new file mode 100644 index 00000000..1e3b0fda --- /dev/null +++ b/internal/workflow/node_exec.go @@ -0,0 +1,153 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/google/uuid" +) + +func executeNode(ctx context.Context, args RunArgs, runID string, node graphNode, state *WorkflowLocalState) (map[string]any, bool, error) { + label := node.Label + if strings.TrimSpace(label) == "" { + label = node.ID + } + nodeRunID := uuid.NewString() + startedAt := time.Now() + incomingCount := 0 + if rt := workflowRuntimeFrom(ctx); rt != nil && rt.idx != nil { + incomingCount = len(rt.idx.incoming[node.ID]) + } + input := map[string]any{ + "nodeId": node.ID, + "nodeType": node.Type, + "label": label, + "inputs": state.Inputs, + "previous": state.LastOutput, + "join": map[string]any{ + "strategy": joinStrategy(node), + "incoming": incomingCount, + }, + } + inputJSON, _ := json.Marshal(input) + if err := args.DB.CreateWorkflowNodeRun(&database.WorkflowNodeRun{ + ID: nodeRunID, + RunID: runID, + NodeID: node.ID, + Status: "running", + InputJSON: string(inputJSON), + StartedAt: startedAt, + }); err != nil { + return nil, false, err + } + if args.Progress != nil { + args.Progress("workflow_node_start", fmt.Sprintf("开始节点：%s", label), map[string]any{ + "workflowRunId": runID, + "nodeRunId": nodeRunID, + "nodeId": node.ID, + "nodeType": node.Type, + "label": label, + }) + } + + result, proceed, status, errText := runBuiltinNode(ctx, args, node, state) + duration := time.Since(startedAt) + if result == nil { + result = map[string]any{} + } + result["duration_ms"] = duration.Milliseconds() + result["finished_at"] = time.Now().Format(time.RFC3339Nano) + result["status"] = status + accumulateWorkflowMetric(state, "node_count", 1) + accumulateWorkflowMetric(state, "duration_ms", duration.Milliseconds()) + if strings.EqualFold(node.Type, "tool") { + accumulateWorkflowMetric(state, "tool_call_count", 1) + } + outputJSON, _ := json.Marshal(result) + if err := args.DB.FinishWorkflowNodeRun(nodeRunID, status, string(outputJSON), errText); err != nil { + return nil, false, err + } + if status == "skipped" { + state.Skipped = append(state.Skipped, label) + } else { + state.Executed = append(state.Executed, label) + } + if args.Progress != nil { + progressData := map[string]any{ + "workflowRunId": runID, + "nodeRunId": nodeRunID, + "nodeId": node.ID, + "nodeType": node.Type, + "label": label, + "status": status, + "durationMs": duration.Milliseconds(), + "output": result, + } + progressMsg := fmt.Sprintf("节点完成：%s（%s）", label, status) + if strings.EqualFold(node.Type, "condition") { + matched := false + if v, ok := result["matched"].(bool); ok { + matched = v + } + expr := cfgString(node.Config, "expression") + if matched { + progressMsg = fmt.Sprintf("条件判断：%s → 是", label) + } else { + progressMsg = fmt.Sprintf("条件判断：%s → 否", label) + } + progressData["expression"] = expr + progressData["matched"] = matched + } + args.Progress("workflow_node_result", progressMsg, progressData) + } + state.NodeProceed[node.ID] = proceed + return result, proceed, nil +} + +func emitConditionBranchProgress(args RunArgs, runID string, node graphNode, edges []graphEdge, nodes map[string]graphNode, state *WorkflowLocalState) { + if args.Progress == nil || len(edges) == 0 { + return + } + for edgeIdx, edge := range edges { + allowed := edgeAllowed(edge, node, edgeIdx, state) + target := nodes[edge.Target] + targetLabel := strings.TrimSpace(target.Label) + if targetLabel == "" { + targetLabel = edge.Target + } + branchLabel := strings.TrimSpace(edge.Label) + if branchLabel == "" { + switch edgeIdx { + case 0: + branchLabel = "是" + case 1: + branchLabel = "否" + default: + branchLabel = fmt.Sprintf("分支 %d", edgeIdx+1) + } + } + cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression")) + eventType := "workflow_branch_skipped" + msg := fmt.Sprintf("跳过分支「%s」→ %s", branchLabel, targetLabel) + if allowed { + eventType = "workflow_branch_taken" + msg = fmt.Sprintf("执行分支「%s」→ %s", branchLabel, targetLabel) + } + args.Progress(eventType, msg, map[string]any{ + "workflowRunId": runID, + "nodeId": node.ID, + "nodeType": node.Type, + "label": node.Label, + "branchLabel": branchLabel, + "targetId": edge.Target, + "targetLabel": targetLabel, + "edgeCondition": cond, + "matched": conditionMatched(state), + }) + } +} diff --git a/internal/workflow/nodes.go b/internal/workflow/nodes.go new file mode 100644 index 00000000..eb0342e2 --- /dev/null +++ b/internal/workflow/nodes.go @@ -0,0 +1,333 @@ +package workflow + +import ( + "context" + "fmt" + "strings" + "unicode/utf8" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/multiagent" +) + +func runBuiltinNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { + cfg := node.Config + switch strings.ToLower(strings.TrimSpace(node.Type)) { + case "start": + return startOutputMap(node, state.Inputs["message"], state.Inputs["conversationId"], state.Inputs["projectId"]), true, "completed", "" + case "condition": + expr := cfgString(cfg, "expression") + ok := evalCondition(expr, state) + return conditionOutputMap(node, expr, ok), true, "completed", "" + case "output": + key := cfgString(cfg, "output_key") + if key == "" { + key = "result" + } + var value any + if v := cfgString(cfg, "static_value"); v != "" { + value = v + } else { + value = resolveOutputSourceBinding(cfg, state) + } + state.Outputs[key] = value + return outputNodeOutputMap(node, key, value), true, "completed", "" + case "end": + value := resolveOutputSourceBinding(cfg, state) + if b, ok := parseFieldBinding(cfg, "result_binding"); ok { + value = resolveBinding(b, state) + } + return endOutputMap(node, value), false, "completed", "" + case "tool": + return runToolNode(ctx, args, node, state) + case "agent": + return runAgentNode(ctx, args, node, state) + case "hitl": + return runHITLNode(args, node, state) + default: + reason := "未知节点类型" + out := outputMap(envelope("unknown", node.ID, node.Type, "skipped", ""), map[string]any{"skipped": true, "reason": reason}) + return out, true, "skipped", reason + } +} + +func runToolNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { + toolName := cfgString(node.Config, "tool_name") + if toolName == "" { + errText := "工具节点未选择 MCP 工具" + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText + } + if args.Agent == nil { + errText := "工具节点执行失败：Agent 为空" + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "error": errText}), false, "failed", errText + } + toolArgs, err := resolveToolArguments(node.Config, state) + if err != nil { + errText := fmt.Sprintf("工具参数不是合法 JSON：%v", err) + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "error": errText}), false, "failed", errText + } + if args.Progress != nil { + args.Progress("workflow_tool_start", fmt.Sprintf("调用工具：%s", toolName), map[string]any{ + "nodeId": node.ID, + "tool": toolName, + "args": toolArgs, + }) + } + result, err := args.Agent.ExecuteMCPToolForConversation(ctx, args.ConversationID, toolName, toolArgs) + if err != nil { + errText := err.Error() + return outputMap(envelope("tool", node.ID, node.Type, "failed", ""), map[string]any{"tool_name": toolName, "arguments": toolArgs, "error": errText}), false, "failed", errText + } + output := "" + executionID := "" + isError := false + if result != nil { + output = result.Result + executionID = result.ExecutionID + isError = result.IsError + } + maxToolOutputBytes := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective() + if args.AppCfg != nil { + maxToolOutputBytes = args.AppCfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective() + } + output = truncateWorkflowToolOutput(output, maxToolOutputBytes, executionID) + out := toolOutputMap(node, output, toolName, toolArgs, executionID, isError) + if key := cfgString(node.Config, "output_key"); key != "" { + state.Outputs[key] = output + } + if isError { + errText := strings.TrimSpace(output) + if errText == "" { + errText = "工具返回错误" + } + return out, false, "failed", errText + } + return out, true, "completed", "" +} + +func truncateWorkflowToolOutput(output string, maxBytes int, executionID string) string { + if maxBytes <= 0 || len(output) <= maxBytes { + return output + } + marker := fmt.Sprintf("\n\n...[workflow tool output truncated; full result is stored in execution %s]...\n\n", strings.TrimSpace(executionID)) + if strings.TrimSpace(executionID) == "" { + marker = "\n\n...[workflow tool output truncated; full result remains in the tool execution record]...\n\n" + } + budget := maxBytes - len(marker) + if budget <= 0 { + return marker + } + head := budget / 2 + tail := budget - head + for head > 0 && !utf8.RuneStart(output[head]) { + head-- + } + tailStart := len(output) - tail + for tailStart < len(output) && !utf8.RuneStart(output[tailStart]) { + tailStart++ + } + return output[:head] + marker + output[tailStart:] +} + +func runAgentNode(ctx context.Context, args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { + if args.AppCfg == nil || args.Agent == nil { + errText := "Agent 节点执行失败：应用配置或 Agent 为空" + return outputMap(envelope("agent", node.ID, node.Type, "failed", ""), map[string]any{"error": errText}), false, "failed", errText + } + mode := strings.ToLower(cfgString(node.Config, "agent_mode")) + if mode == "" { + mode = "eino_single" + } + inputSource := resolveNodeInputBinding(node.Config, state) + message := buildAgentNodeMessage(node, state, inputSource) + var result *multiagent.RunResult + var err error + state.SegmentMaxIteration = 0 + agentProgress := workflowAgentProgress(args.Progress, state, node) + switch mode { + case "eino_single", "single", "chat": + result, err = multiagent.RunEinoSingleChatModelAgent( + ctx, + args.AppCfg, + &args.AppCfg.MultiAgent, + args.Agent, + args.DB, + args.Logger, + args.ConversationID, + args.ProjectID, + message, + args.History, + args.RoleTools, + agentProgress, + nil, + args.SystemPromptExtra, + ) + default: + result, err = multiagent.RunDeepAgent( + ctx, + args.AppCfg, + &args.AppCfg.MultiAgent, + args.Agent, + args.DB, + args.Logger, + args.ConversationID, + args.ProjectID, + message, + args.History, + args.RoleTools, + agentProgress, + args.AgentsMarkdownDir, + mode, + nil, + args.SystemPromptExtra, + ) + } + if err != nil { + errText := err.Error() + state.MainIterationOffset += state.SegmentMaxIteration + return outputMap(envelope("agent", node.ID, node.Type, "failed", ""), map[string]any{"mode": mode, "error": errText}), false, "failed", errText + } + state.MainIterationOffset += state.SegmentMaxIteration + response := "" + mcpIDs := []string{} + if result != nil { + response = result.Response + mcpIDs = result.MCPExecutionIDs + } + if args.Progress != nil { + args.Progress("workflow_agent_output", response, map[string]any{ + "nodeId": node.ID, + "label": firstNonEmpty(node.Label, node.ID), + "mode": mode, + "inputSource": inputSource, + "inputPreview": truncateWorkflowPreview(inputSource, 500), + "mcpExecutionIds": mcpIDs, + }) + } + if key := cfgString(node.Config, "output_key"); key != "" { + state.Outputs[key] = response + } + return agentOutputMap(node, response, mode, mcpIDs), true, "completed", "" +} + +func buildAgentNodeMessage(node graphNode, state *WorkflowLocalState, upstreamInput string) string { + instruction := strings.TrimSpace(cfgString(node.Config, "instruction")) + upstreamInput = strings.TrimSpace(upstreamInput) + if instruction == "" { + if upstreamInput != "" { + return fmt.Sprintf("请基于上游节点输出继续处理：\n%s", upstreamInput) + } + return fmt.Sprintf("请基于上游节点输出继续处理：\n%v", state.LastOutput["output"]) + } + if upstreamInput == "" { + return instruction + } + return strings.TrimSpace(fmt.Sprintf("上游输入：\n%s\n\n节点指令：\n%s", upstreamInput, instruction)) +} + +func workflowAgentProgress(progress agent.ProgressCallback, state *WorkflowLocalState, node graphNode) agent.ProgressCallback { + if progress == nil { + return nil + } + return func(eventType, message string, data interface{}) { + switch eventType { + case "response_start", "response_delta", "response", "done": + return + default: + enrichWorkflowAgentEventData(data, state, node) + collectAgentMetrics(state, data) + if eventType == "iteration" { + applyWorkflowMainIterationOffset(data, state) + } + progress(eventType, message, data) + } + } +} + +func enrichWorkflowAgentEventData(data interface{}, state *WorkflowLocalState, node graphNode) { + m, ok := data.(map[string]interface{}) + if !ok || m == nil { + return + } + if node.ID != "" { + m["workflowNodeId"] = node.ID + } + if state != nil && strings.TrimSpace(state.WorkflowRunID) != "" { + m["workflowRunId"] = state.WorkflowRunID + } +} + +func applyWorkflowMainIterationOffset(data interface{}, state *WorkflowLocalState) { + if state == nil { + return + } + m, ok := data.(map[string]interface{}) + if !ok || m == nil { + return + } + scope, _ := m["einoScope"].(string) + if strings.TrimSpace(scope) != "main" { + return + } + raw := iterationNumberFromProgressData(m) + if raw <= 0 { + return + } + if raw > state.SegmentMaxIteration { + state.SegmentMaxIteration = raw + } + m["iteration"] = raw + state.MainIterationOffset +} + +func iterationNumberFromProgressData(m map[string]interface{}) int { + switch v := m["iteration"].(type) { + case int: + return v + case int32: + return int(v) + case int64: + return int(v) + case float64: + return int(v) + case float32: + return int(v) + default: + return 0 + } +} + +func runHITLNode(args RunArgs, node graphNode, state *WorkflowLocalState) (map[string]any, bool, string, string) { + prompt := resolveHITLPromptBinding(node.Config, state) + reviewer := cfgString(node.Config, "reviewer") + if reviewer == "" { + reviewer = "human" + } + approved := true + if state != nil && state.Inputs != nil { + if v, ok := state.Inputs["_hitl_approved"]; ok { + approved = fmt.Sprint(v) == "true" + } + } + if !approved { + reason := "人工审批已拒绝" + if state != nil && state.Inputs != nil { + if v, ok := state.Inputs["_hitl_comment"]; ok { + if s := strings.TrimSpace(fmt.Sprint(v)); s != "" { + reason = s + } + } + } + return hitlOutputMap(node, "failed", "", prompt, reviewer, false), false, "failed", reason + } + if args.Progress != nil { + args.Progress("workflow_hitl_checkpoint", "人工确认节点已通过", map[string]any{ + "nodeId": node.ID, + "prompt": prompt, + "reviewer": reviewer, + "mode": "interactive", + "approved": true, + }) + } + return hitlOutputMap(node, "completed", prompt, prompt, reviewer, true), true, "completed", "" +} diff --git a/internal/workflow/package/exporter.go b/internal/workflow/package/exporter.go new file mode 100644 index 00000000..46494354 --- /dev/null +++ b/internal/workflow/package/exporter.go @@ -0,0 +1,98 @@ +package workflowpackage + +import ( + "archive/zip" + "bytes" + "encoding/json" + "fmt" + "path" + "strings" + "time" +) + +// Export builds a deterministic, human-readable single-workflow package. +func Export(source Document) ([]byte, ExportMetadata, error) { + source.ID = strings.TrimSpace(source.ID) + source.Name = strings.TrimSpace(source.Name) + if source.ID == "" || source.Name == "" || source.Version <= 0 || !safePackageWorkflowID(source.ID) { + return nil, ExportMetadata{}, fmt.Errorf("workflow id, name and version are required") + } + if strings.TrimSpace(source.GraphJSON) == "" { + return nil, ExportMetadata{}, fmt.Errorf("workflow graph_json is required") + } + contentHash, graphHash, payload, err := DocumentHashes(source) + if err != nil { + return nil, ExportMetadata{}, err + } + workflowPath := path.Join("workflows", source.ID+".json") + createdAt := source.UpdatedAt.UTC() + if createdAt.IsZero() { + createdAt = time.Unix(0, 0).UTC() + } + manifest := Manifest{ + PackageFormat: PackageFormat, + FormatVersion: FormatVersion, + PackageID: "pkg_" + strings.TrimPrefix(contentHash, "sha256:")[:16], + CreatedAt: createdAt.Format(time.RFC3339), + Items: []ManifestItem{{ + Type: "workflow", + Path: workflowPath, + SourceID: source.ID, + SourceRevision: source.Version, + ContentHash: contentHash, + GraphHash: graphHash, + }}, + } + manifestBytes, err := json.Marshal(manifest) + if err != nil { + return nil, ExportMetadata{}, fmt.Errorf("marshal manifest: %w", err) + } + checksums := fmt.Sprintf("%s manifest.json\n%s %s\n", strings.TrimPrefix(sha256Prefixed(manifestBytes), "sha256:"), strings.TrimPrefix(contentHash, "sha256:"), workflowPath) + + var out bytes.Buffer + zw := zip.NewWriter(&out) + for _, entry := range []struct { + name string + data []byte + }{ + {name: "checksums.sha256", data: []byte(checksums)}, + {name: "manifest.json", data: manifestBytes}, + {name: workflowPath, data: payload}, + } { + header := &zip.FileHeader{Name: entry.name, Method: zip.Store} + header.SetModTime(time.Unix(0, 0).UTC()) + writer, err := zw.CreateHeader(header) + if err != nil { + return nil, ExportMetadata{}, fmt.Errorf("write %s: %w", entry.name, err) + } + if _, err := writer.Write(entry.data); err != nil { + return nil, ExportMetadata{}, fmt.Errorf("write %s: %w", entry.name, err) + } + } + if err := zw.Close(); err != nil { + return nil, ExportMetadata{}, fmt.Errorf("close package: %w", err) + } + pkg := out.Bytes() + return pkg, ExportMetadata{ + PackageHash: sha256Prefixed(pkg), + ContentHash: contentHash, + GraphHash: graphHash, + SourceRevision: source.Version, + FileName: source.ID + ".csapkg.zip", + }, nil +} + +// DocumentHashes returns the canonical package item and graph hashes together +// with the canonical item bytes used by export and inspection persistence. +func DocumentHashes(source Document) (string, string, []byte, error) { + graph, err := canonicalJSON([]byte(source.GraphJSON)) + if err != nil { + return "", "", nil, fmt.Errorf("canonicalize graph_json: %w", err) + } + source.GraphJSON = string(graph) + payload, err := json.Marshal(source) + if err != nil { + return "", "", nil, fmt.Errorf("marshal workflow payload: %w", err) + } + return sha256Prefixed(payload), sha256Prefixed(graph), payload, nil +} diff --git a/internal/workflow/package/exporter_test.go b/internal/workflow/package/exporter_test.go new file mode 100644 index 00000000..4b13fa48 --- /dev/null +++ b/internal/workflow/package/exporter_test.go @@ -0,0 +1,58 @@ +package workflowpackage + +import ( + "archive/zip" + "bytes" + "strings" + "testing" + "time" +) + +func testDocument() Document { + return Document{ + ID: "web-src-hunting", + Name: "Web SRC 猎洞", + Description: "面向 SRC Web 资产的侦察与漏洞候选流程", + Version: 18, + Enabled: true, + GraphJSON: `{"nodes":[{"id":"start-1","type":"start","label":"开始","position":{"x":0,"y":0},"config":{}},{"id":"out-1","type":"output","label":"输出","position":{"x":0,"y":120},"config":{"output_key":"result","source_binding":{"from":"inputs","field":"message"}}}],"edges":[{"id":"e1","source":"start-1","target":"out-1"}],"config":{"schema_version":1}}`, + UpdatedAt: time.Date(2026, 7, 13, 10, 0, 0, 0, time.UTC), + } +} + +func TestExportIsDeterministicAndSelfDescribing(t *testing.T) { + first, firstMeta, err := Export(testDocument()) + if err != nil { + t.Fatalf("first export: %v", err) + } + second, secondMeta, err := Export(testDocument()) + if err != nil { + t.Fatalf("second export: %v", err) + } + if !bytes.Equal(first, second) { + t.Fatal("identical document must produce byte-identical package") + } + if firstMeta.PackageHash != secondMeta.PackageHash || !strings.HasPrefix(firstMeta.PackageHash, "sha256:") { + t.Fatalf("unexpected deterministic package hash: %#v / %#v", firstMeta, secondMeta) + } + + zr, err := zip.NewReader(bytes.NewReader(first), int64(len(first))) + if err != nil { + t.Fatalf("open package: %v", err) + } + if len(zr.File) != 3 { + t.Fatalf("zip entry count = %d, want 3", len(zr.File)) + } + wantNames := []string{"checksums.sha256", "manifest.json", "workflows/web-src-hunting.json"} + for i, f := range zr.File { + if f.Name != wantNames[i] { + t.Fatalf("entry %d = %q, want %q", i, f.Name, wantNames[i]) + } + } + if firstMeta.SourceRevision != 18 { + t.Fatalf("source revision = %d, want 18", firstMeta.SourceRevision) + } + if !strings.HasPrefix(firstMeta.ContentHash, "sha256:") || !strings.HasPrefix(firstMeta.GraphHash, "sha256:") { + t.Fatalf("content/graph hashes must be sha256: %#v", firstMeta) + } +} diff --git a/internal/workflow/package/inspector.go b/internal/workflow/package/inspector.go new file mode 100644 index 00000000..7d533563 --- /dev/null +++ b/internal/workflow/package/inspector.go @@ -0,0 +1,225 @@ +package workflowpackage + +import ( + "archive/zip" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path" + "strings" + "unicode" +) + +const ( + MaxArchiveBytes = 10 << 20 + MaxExtractedBytes = 20 << 20 +) + +// PackageError contains only a contract error code and safe, client-facing fields. +type PackageError struct { + Code string + Message string + Details map[string]any +} + +func (e *PackageError) Error() string { return e.Code + ": " + e.Message } + +func packageError(code, message string) error { + return &PackageError{Code: code, Message: message} +} + +// ErrorCode returns a package contract code without exposing internal errors. +func ErrorCode(err error) string { + var target *PackageError + if errors.As(err, &target) { + return target.Code + } + return "" +} + +type InspectionResult struct { + PackageHash string + Manifest Manifest + Document Document + ContentHash string + GraphHash string + NodeCount int + EdgeCount int +} + +// InspectArchive verifies an archive without executing any package content. +// validateGraph is injected by the application so this format package has no +// dependency on the workflow runtime or database driver. +func InspectArchive(ctx context.Context, archive []byte, validateGraph func(context.Context, string) error) (*InspectionResult, error) { + if len(archive) == 0 { + return nil, packageError("WFPKG_FILE_REQUIRED", "必须上传工作流包文件") + } + if len(archive) > MaxArchiveBytes { + return nil, packageError("WFPKG_FILE_TOO_LARGE", "工作流包文件超过大小限制") + } + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包不是有效 ZIP 文件") + } + entries := make(map[string][]byte, len(zr.File)) + var extracted int64 + for _, file := range zr.File { + if !safeArchivePath(file.Name) || file.FileInfo().IsDir() || file.FileInfo().Mode()&os.ModeSymlink != 0 { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含不安全文件路径") + } + if _, exists := entries[file.Name]; exists { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含重复文件") + } + if file.UncompressedSize64 > MaxExtractedBytes || extracted+int64(file.UncompressedSize64) > MaxExtractedBytes { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包解压后超过大小限制") + } + reader, err := file.Open() + if err != nil { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "无法读取工作流包文件") + } + data, readErr := io.ReadAll(io.LimitReader(reader, int64(MaxExtractedBytes)-extracted+1)) + closeErr := reader.Close() + if readErr != nil || closeErr != nil || len(data) > MaxExtractedBytes-int(extracted) { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包解压后超过大小限制") + } + extracted += int64(len(data)) + entries[file.Name] = data + } + + manifestRaw, hasManifest := entries["manifest.json"] + checksumsRaw, hasChecksums := entries["checksums.sha256"] + if !hasManifest || !hasChecksums { + return nil, packageError("WFPKG_UNSUPPORTED_FORMAT", "工作流包缺少必需文件") + } + manifest, err := parseManifest(manifestRaw) + if err != nil { + return nil, err + } + if len(manifest.Items) != 1 || manifest.Items[0].Type != "workflow" { + return nil, packageError("WFPKG_MULTIPLE_WORKFLOWS", "工作流包必须且只能包含一个工作流") + } + item := manifest.Items[0] + workflowRaw, exists := entries[item.Path] + if !exists || !safeWorkflowPath(item.Path) || len(entries) != 3 { + return nil, packageError("WFPKG_INVALID_ARCHIVE", "工作流包包含未声明文件") + } + checksums, err := parseChecksums(checksumsRaw) + if err != nil { + return nil, err + } + if len(checksums) != 2 || checksums["manifest.json"] != sha256Prefixed(manifestRaw) || checksums[item.Path] != sha256Prefixed(workflowRaw) { + return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和不匹配") + } + if item.ContentHash != sha256Prefixed(workflowRaw) || !validHash(item.ContentHash) || !validHash(item.GraphHash) { + return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包内容校验和不匹配") + } + doc, err := parseDocument(workflowRaw) + if err != nil { + return nil, err + } + if !safePackageWorkflowID(doc.ID) || doc.ID != item.SourceID || doc.Version != item.SourceRevision { + return nil, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单与工作流内容不一致") + } + graph, err := canonicalJSON([]byte(doc.GraphJSON)) + if err != nil || item.GraphHash != sha256Prefixed(graph) { + return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流图校验和不匹配") + } + if validateGraph == nil || validateGraph(ctx, string(graph)) != nil { + return nil, packageError("WFPKG_WORKFLOW_INVALID", "工作流图校验失败") + } + var graphShape struct { + Nodes []json.RawMessage `json:"nodes"` + Edges []json.RawMessage `json:"edges"` + } + if err := json.Unmarshal(graph, &graphShape); err != nil { + return nil, packageError("WFPKG_WORKFLOW_INVALID", "工作流图不是有效 JSON") + } + return &InspectionResult{ + PackageHash: sha256Prefixed(archive), + Manifest: manifest, + Document: doc, + ContentHash: item.ContentHash, + GraphHash: item.GraphHash, + NodeCount: len(graphShape.Nodes), + EdgeCount: len(graphShape.Edges), + }, nil +} + +func safeArchivePath(name string) bool { + return name != "" && !strings.Contains(name, `\`) && !strings.HasPrefix(name, "/") && path.Clean(name) == name && !strings.HasPrefix(name, "../") && name != ".." +} + +func safeWorkflowPath(name string) bool { + rest := strings.TrimPrefix(name, "workflows/") + return safeArchivePath(name) && strings.HasPrefix(name, "workflows/") && rest != "" && !strings.Contains(rest, "/") && strings.HasSuffix(rest, ".json") +} + +func safePackageWorkflowID(id string) bool { + if id == "" || strings.ContainsAny(id, `/\`) { + return false + } + for _, r := range id { + if unicode.IsControl(r) { + return false + } + } + return true +} + +func parseManifest(raw []byte) (Manifest, error) { + var manifest Manifest + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.DisallowUnknownFields() + if err := dec.Decode(&manifest); err != nil { + return Manifest{}, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单格式无效") + } + if err := consumeJSONEnd(dec); err != nil || manifest.PackageFormat != PackageFormat || manifest.FormatVersion != FormatVersion || strings.TrimSpace(manifest.PackageID) == "" || len(manifest.Items) == 0 { + return Manifest{}, packageError("WFPKG_INVALID_MANIFEST", "工作流包清单格式不受支持") + } + return manifest, nil +} + +func parseDocument(raw []byte) (Document, error) { + var doc Document + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.DisallowUnknownFields() + if err := dec.Decode(&doc); err != nil || consumeJSONEnd(dec) != nil { + return Document{}, packageError("WFPKG_WORKFLOW_INVALID", "工作流定义格式无效") + } + doc.ID = strings.TrimSpace(doc.ID) + doc.Name = strings.TrimSpace(doc.Name) + if doc.ID == "" || doc.Name == "" || doc.Version <= 0 || strings.TrimSpace(doc.GraphJSON) == "" { + return Document{}, packageError("WFPKG_WORKFLOW_INVALID", "工作流定义缺少必需字段") + } + return doc, nil +} + +func consumeJSONEnd(dec *json.Decoder) error { + var extra any + if err := dec.Decode(&extra); err != io.EOF { + if err == nil { + return fmt.Errorf("multiple JSON values") + } + return err + } + return nil +} + +func parseChecksums(raw []byte) (map[string]string, error) { + entries := make(map[string]string) + for _, line := range strings.Split(strings.TrimSpace(string(raw)), "\n") { + parts := strings.SplitN(strings.TrimSpace(line), " ", 2) + if len(parts) != 2 || !validHash("sha256:"+parts[0]) || !safeArchivePath(parts[1]) { + return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和格式无效") + } + if _, exists := entries[parts[1]]; exists { + return nil, packageError("WFPKG_CHECKSUM_MISMATCH", "工作流包校验和重复") + } + entries[parts[1]] = "sha256:" + parts[0] + } + return entries, nil +} diff --git a/internal/workflow/package/inspector_test.go b/internal/workflow/package/inspector_test.go new file mode 100644 index 00000000..4847eabd --- /dev/null +++ b/internal/workflow/package/inspector_test.go @@ -0,0 +1,154 @@ +package workflowpackage + +import ( + "archive/zip" + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "testing" +) + +func TestInspectArchiveAcceptsSingleVerifiedWorkflow(t *testing.T) { + pkg, meta, err := Export(testDocument()) + if err != nil { + t.Fatal(err) + } + result, err := InspectArchive(context.Background(), pkg, func(context.Context, string) error { return nil }) + if err != nil { + t.Fatalf("InspectArchive: %v", err) + } + if result.PackageHash != meta.PackageHash || result.Document.ID != "web-src-hunting" { + t.Fatalf("unexpected inspection: %#v", result) + } + if result.NodeCount != 2 || result.EdgeCount != 1 { + t.Fatalf("counts = %d/%d, want 2/1", result.NodeCount, result.EdgeCount) + } +} + +func TestInspectArchiveRejectsUnsafeArchiveShapes(t *testing.T) { + valid, _, err := Export(testDocument()) + if err != nil { + t.Fatal(err) + } + cases := []struct { + name string + archive []byte + }{ + {name: "duplicate entry", archive: appendZipEntry(t, valid, "manifest.json", []byte(`{}`), 0)}, + {name: "path traversal", archive: appendZipEntry(t, valid, "../payload.json", []byte(`{}`), 0)}, + {name: "symlink", archive: appendZipEntry(t, valid, "workflows/link.json", []byte("target"), 0o120777)}, + {name: "undeclared file", archive: appendZipEntry(t, valid, "notes.txt", []byte("not allowed"), 0)}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := InspectArchive(context.Background(), tc.archive, func(context.Context, string) error { return nil }) + if ErrorCode(err) != "WFPKG_INVALID_ARCHIVE" { + t.Fatalf("code = %q, err = %v", ErrorCode(err), err) + } + }) + } +} + +func TestInspectArchiveRejectsChecksumMismatchAndInvalidWorkflow(t *testing.T) { + pkg, _, err := Export(testDocument()) + if err != nil { + t.Fatal(err) + } + badChecksum := replaceZipEntry(t, pkg, "checksums.sha256", []byte("00 manifest.json\n"), 0) + if _, err := InspectArchive(context.Background(), badChecksum, func(context.Context, string) error { return nil }); ErrorCode(err) != "WFPKG_CHECKSUM_MISMATCH" { + t.Fatalf("checksum code = %q, err = %v", ErrorCode(err), err) + } + if _, err := InspectArchive(context.Background(), pkg, func(context.Context, string) error { return errors.New("invalid graph") }); ErrorCode(err) != "WFPKG_WORKFLOW_INVALID" { + t.Fatalf("graph code = %q, err = %v", ErrorCode(err), err) + } +} + +func appendZipEntry(t *testing.T, archive []byte, name string, data []byte, mode os.FileMode) []byte { + t.Helper() + return rewriteZip(t, archive, func(zw *zip.Writer) error { + h := &zip.FileHeader{Name: name, Method: zip.Store} + h.SetMode(mode) + w, err := zw.CreateHeader(h) + if err != nil { + return err + } + _, err = w.Write(data) + return err + }) +} + +func replaceZipEntry(t *testing.T, archive []byte, name string, data []byte, mode os.FileMode) []byte { + t.Helper() + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + t.Fatal(err) + } + var out bytes.Buffer + zw := zip.NewWriter(&out) + for _, f := range zr.File { + if f.Name == name { + h := &zip.FileHeader{Name: name, Method: zip.Store} + h.SetMode(mode) + w, err := zw.CreateHeader(h) + if err != nil { + t.Fatal(err) + } + if _, err := w.Write(data); err != nil { + t.Fatal(err) + } + continue + } + r, err := f.Open() + if err != nil { + t.Fatal(err) + } + h := &zip.FileHeader{Name: f.Name, Method: zip.Store} + w, err := zw.CreateHeader(h) + if err == nil { + _, err = io.Copy(w, r) + } + _ = r.Close() + if err != nil { + t.Fatal(err) + } + } + if err := zw.Close(); err != nil { + t.Fatal(err) + } + return out.Bytes() +} + +func rewriteZip(t *testing.T, archive []byte, appendEntry func(*zip.Writer) error) []byte { + t.Helper() + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + t.Fatal(err) + } + var out bytes.Buffer + zw := zip.NewWriter(&out) + for _, f := range zr.File { + r, err := f.Open() + if err != nil { + t.Fatal(err) + } + h := &zip.FileHeader{Name: f.Name, Method: zip.Store} + w, err := zw.CreateHeader(h) + if err == nil { + _, err = io.Copy(w, r) + } + _ = r.Close() + if err != nil { + t.Fatal(err) + } + } + if err := appendEntry(zw); err != nil { + t.Fatal(fmt.Errorf("append entry: %w", err)) + } + if err := zw.Close(); err != nil { + t.Fatal(err) + } + return out.Bytes() +} diff --git a/internal/workflow/package/manifest.go b/internal/workflow/package/manifest.go new file mode 100644 index 00000000..a84e263b --- /dev/null +++ b/internal/workflow/package/manifest.go @@ -0,0 +1,78 @@ +package workflowpackage + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "strings" + "time" +) + +const ( + PackageFormat = "cyberstrikeai.workflow-package" + FormatVersion = "1.0" +) + +// Document is the single non-executable workflow definition carried by a package. +// Version is the source instance revision and is never applied as a target version. +type Document struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Version int `json:"version"` + GraphJSON string `json:"graph_json"` + Enabled bool `json:"enabled"` + UpdatedAt time.Time `json:"-"` +} + +type Manifest struct { + PackageFormat string `json:"package_format"` + FormatVersion string `json:"format_version"` + PackageID string `json:"package_id"` + CreatedAt string `json:"created_at"` + Items []ManifestItem `json:"items"` +} + +type ManifestItem struct { + Type string `json:"type"` + Path string `json:"path"` + SourceID string `json:"source_id"` + SourceRevision int `json:"source_revision"` + ContentHash string `json:"content_hash"` + GraphHash string `json:"graph_hash"` +} + +type ExportMetadata struct { + PackageHash string + ContentHash string + GraphHash string + SourceRevision int + FileName string +} + +func canonicalJSON(raw []byte) ([]byte, error) { + dec := json.NewDecoder(strings.NewReader(string(raw))) + dec.UseNumber() + var value any + if err := dec.Decode(&value); err != nil { + return nil, err + } + if dec.More() { + return nil, fmt.Errorf("extra JSON values") + } + return json.Marshal(value) +} + +func sha256Prefixed(b []byte) string { + sum := sha256.Sum256(b) + return "sha256:" + hex.EncodeToString(sum[:]) +} + +func validHash(value string) bool { + if !strings.HasPrefix(value, "sha256:") || len(value) != len("sha256:")+64 { + return false + } + _, err := hex.DecodeString(strings.TrimPrefix(value, "sha256:")) + return err == nil && value == strings.ToLower(value) +} diff --git a/internal/workflow/runner.go b/internal/workflow/runner.go new file mode 100644 index 00000000..e5b3876d --- /dev/null +++ b/internal/workflow/runner.go @@ -0,0 +1,223 @@ +package workflow + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// ShouldAutoRunRoleWorkflow returns true when a role explicitly binds a workflow +// and does not turn it off. Empty policy defaults to auto to keep role UX simple. +func ShouldAutoRunRoleWorkflow(role config.RoleConfig) bool { + if strings.TrimSpace(role.WorkflowID) == "" { + return false + } + policy := strings.ToLower(strings.TrimSpace(role.WorkflowPolicy)) + return policy == "" || policy == "auto" +} + +// RunRoleBoundWorkflow executes the persisted role-bound workflow via cached Eino Workflow. +func RunRoleBoundWorkflow(ctx context.Context, args RunArgs) (*RunResult, error) { + if args.DB == nil { + return nil, fmt.Errorf("workflow db is nil") + } + workflowID := strings.TrimSpace(args.Role.WorkflowID) + if workflowID == "" { + return nil, fmt.Errorf("角色未绑定工作流") + } + wf, err := args.DB.GetWorkflowDefinition(workflowID) + if err != nil { + return nil, err + } + if wf == nil { + return nil, fmt.Errorf("角色绑定的工作流不存在: %s", workflowID) + } + if !wf.Enabled { + return nil, fmt.Errorf("角色绑定的工作流已禁用: %s", workflowID) + } + + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + runID := uuid.NewString() + input := map[string]interface{}{ + "message": args.UserMessage, + "conversationId": args.ConversationID, + "projectId": args.ProjectID, + "role": args.Role.Name, + "workflowId": wf.ID, + "workflowVersion": wf.Version, + } + inputJSON, _ := json.Marshal(input) + run := &database.WorkflowRun{ + ID: runID, + WorkflowID: wf.ID, + WorkflowVersion: wf.Version, + ConversationID: args.ConversationID, + ProjectID: args.ProjectID, + RoleID: args.Role.Name, + Status: "running", + InputJSON: string(inputJSON), + StartedAt: time.Now(), + } + if err := args.DB.CreateWorkflowRun(run); err != nil { + return nil, err + } + if args.Progress != nil { + args.Progress("workflow_start", fmt.Sprintf("开始运行流程「%s」", wf.Name), map[string]interface{}{ + "workflowId": wf.ID, + "workflowName": wf.Name, + "workflowVersion": wf.Version, + "workflowRunId": runID, + "conversationId": args.ConversationID, + "engine": "eino_workflow", + }) + } + + graph, err := parseGraph(wf.GraphJSON) + if err != nil { + _ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error()) + return nil, err + } + state := newWorkflowLocalState(input, runID) + streaming := args.Progress != nil + resuming := false + for { + _, err := invokeEinoGraph(ctx, args, runID, wf.ID, wf.Version, graph, state, resuming) + if err == nil { + break + } + if !IsAwaitingHITL(err) { + _ = args.DB.FinishWorkflowRun(runID, "failed", "", err.Error()) + return nil, err + } + hitl := err.(*AwaitingHITLError) + partial := map[string]interface{}{ + "workflowId": wf.ID, + "workflowName": wf.Name, + "workflowVersion": wf.Version, + "workflowRunId": runID, + "status": "awaiting_hitl", + "outputs": state.Outputs, + "executedNodes": state.Executed, + "skippedNodes": state.Skipped, + "pendingHitl": map[string]interface{}{ + "nodeId": hitl.NodeID, + "label": hitl.NodeLabel, + "prompt": hitl.Prompt, + }, + "engine": "eino_workflow", + } + partialJSON, _ := json.Marshal(partial) + _ = args.DB.SetWorkflowRunAwaitingHITL(runID, hitl.NodeID, string(partialJSON)) + response := fmt.Sprintf("工作流「%s」已在节点「%s」暂停，等待人工审批。\n运行 ID：%s", wf.Name, firstNonEmpty(hitl.NodeLabel, hitl.NodeID), runID) + if args.Progress != nil { + args.Progress("workflow_paused", response, map[string]interface{}{ + "workflowRunId": runID, + "status": "awaiting_hitl", + "nodeId": hitl.NodeID, + "resumeApi": fmt.Sprintf("/api/workflows/runs/%s/resume", runID), + }) + } + if !streaming { + return &RunResult{ + Response: response, + RunID: runID, + Status: "awaiting_hitl", + AwaitingHITL: true, + }, nil + } + ch := registerHITLWaiter(runID) + decision, waitErr := waitWorkflowHITLDecisionWithChannel(ctx, args.DB, runID, ch) + unregisterHITLWaiter(runID, ch) + if waitErr != nil { + _ = args.DB.FinishWorkflowRun(runID, "cancelled", "", waitErr.Error()) + return nil, waitErr + } + if !decision.Approved { + errText := strings.TrimSpace(decision.Comment) + if errText == "" { + errText = "人工审批拒绝" + } + _ = args.DB.FinishWorkflowRun(runID, "rejected", "", errText) + rejectResponse := fmt.Sprintf("工作流已在审批节点「%s」被拒绝。", firstNonEmpty(hitl.NodeLabel, hitl.NodeID)) + if args.Progress != nil { + args.Progress("workflow_hitl_rejected", rejectResponse, map[string]interface{}{ + "workflowRunId": runID, + "nodeId": hitl.NodeID, + "comment": errText, + }) + } + return &RunResult{ + Response: rejectResponse, + RunID: runID, + Status: "rejected", + }, nil + } + if args.Progress != nil { + args.Progress("workflow_hitl_resumed", "人工审批已通过，继续执行", map[string]interface{}{ + "workflowRunId": runID, + "nodeId": hitl.NodeID, + "comment": decision.Comment, + }) + } + if state.Inputs == nil { + state.Inputs = map[string]any{} + } + state.Inputs["_hitl_approved"] = true + state.Inputs["_hitl_comment"] = decision.Comment + state.Inputs["_hitl_node_id"] = hitl.NodeID + _ = args.DB.SetWorkflowRunStatus(runID, "running") + resuming = true + } + + output := map[string]interface{}{ + "workflowId": wf.ID, + "workflowName": wf.Name, + "workflowVersion": wf.Version, + "workflowRunId": runID, + "status": "completed", + "outputs": state.Outputs, + "metrics": state.Metrics, + "executedNodes": state.Executed, + "skippedNodes": state.Skipped, + "engine": "eino_workflow", + } + outputJSON, _ := json.Marshal(output) + + response := renderWorkflowResponse(args.Role.Name, wf.Name, wf.Version, runID, state) + if err := args.DB.FinishWorkflowRun(runID, "completed", string(outputJSON), ""); err != nil { + return nil, err + } + if args.Progress != nil { + args.Progress("workflow_done", fmt.Sprintf("流程「%s」运行完成", wf.Name), map[string]interface{}{ + "workflowRunId": runID, + "workflowId": wf.ID, + "outputs": state.Outputs, + "metrics": state.Metrics, + "response": response, + "engine": "eino_workflow", + }) + } + if args.Logger != nil { + args.Logger.Info("role-bound workflow completed", + zap.String("workflow_id", wf.ID), + zap.String("workflow_run_id", runID), + zap.String("conversation_id", args.ConversationID), + zap.String("role", args.Role.Name), + zap.String("engine", "eino_workflow"), + ) + } + return &RunResult{Response: response, RunID: runID, Status: "completed"}, nil +} diff --git a/internal/workflow/state.go b/internal/workflow/state.go new file mode 100644 index 00000000..8b3a0376 --- /dev/null +++ b/internal/workflow/state.go @@ -0,0 +1,214 @@ +package workflow + +import ( + "fmt" + "regexp" + "sort" + "strings" + + "github.com/cloudwego/eino/schema" +) + +func init() { + 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). +type WorkflowLocalState struct { + Inputs map[string]any `json:"inputs,omitempty"` + Outputs map[string]any `json:"outputs,omitempty"` + NodeOutputs map[string]map[string]any `json:"nodeOutputs,omitempty"` + NodeProceed map[string]bool `json:"nodeProceed,omitempty"` + LastOutput map[string]any `json:"lastOutput,omitempty"` + Metrics map[string]any `json:"metrics,omitempty"` + Executed []string `json:"executed,omitempty"` + Skipped []string `json:"skipped,omitempty"` + WorkflowRunID string `json:"workflowRunId,omitempty"` + MainIterationOffset int `json:"mainIterationOffset,omitempty"` + SegmentMaxIteration int `json:"segmentMaxIteration,omitempty"` +} + +func newWorkflowLocalState(inputs map[string]interface{}, runID string) *WorkflowLocalState { + in := make(map[string]any, len(inputs)) + for k, v := range inputs { + in[k] = v + } + return &WorkflowLocalState{ + Inputs: in, + Outputs: make(map[string]any), + NodeOutputs: make(map[string]map[string]any), + NodeProceed: make(map[string]bool), + Metrics: make(map[string]any), + WorkflowRunID: runID, + } +} + +var templateVarRe = regexp.MustCompile(`\{\{\s*([a-zA-Z0-9_.-]+)\s*\}\}`) + +func resolveTemplate(s string, state *WorkflowLocalState) string { + if strings.TrimSpace(s) == "" { + return fmt.Sprint(valueFromPath("previous.output", state)) + } + return templateVarRe.ReplaceAllStringFunc(s, func(match string) string { + m := templateVarRe.FindStringSubmatch(match) + if len(m) != 2 { + return match + } + return fmt.Sprint(valueFromPath(m[1], state)) + }) +} + +func valueFromPath(path string, state *WorkflowLocalState) any { + parts := strings.Split(path, ".") + if len(parts) == 0 { + return "" + } + var cur any + switch parts[0] { + case "inputs", "input": + cur = state.Inputs + case "previous", "prev": + cur = state.LastOutput + case "outputs": + cur = state.Outputs + default: + if v, ok := state.Inputs[parts[0]]; ok { + cur = v + } else if v, ok := state.NodeOutputs[parts[0]]; ok { + cur = v + } else { + return "" + } + } + for _, p := range parts[1:] { + m, ok := cur.(map[string]any) + if !ok { + return "" + } + cur = m[p] + } + if cur == nil { + return "" + } + return cur +} + +func cleanComparable(s string) string { + s = strings.TrimSpace(s) + s = strings.Trim(s, `"'`) + return s +} + +func edgeAllowed(edge graphEdge, sourceNode graphNode, edgeIndex int, state *WorkflowLocalState) bool { + cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression")) + if cond != "" { + return evalCondition(cond, state) + } + if strings.EqualFold(strings.TrimSpace(sourceNode.Type), "condition") { + return conditionBranchAllowed(edge, edgeIndex, state) + } + return true +} + +func conditionBranchAllowed(edge graphEdge, edgeIndex int, state *WorkflowLocalState) bool { + matched := conditionMatched(state) + if branch := conditionBranchHint(edge); branch != "" { + return (branch == "true" && matched) || (branch == "false" && !matched) + } + switch edgeIndex { + case 0: + return matched + case 1: + return !matched + default: + return false + } +} + +func conditionMatched(state *WorkflowLocalState) bool { + v := strings.ToLower(cleanComparable(fmt.Sprint(valueFromPath("previous.matched", state)))) + return v == "true" || v == "1" +} + +func conditionBranchHint(edge graphEdge) string { + if edge.Config != nil { + switch strings.ToLower(strings.TrimSpace(cfgString(edge.Config, "branch"))) { + case "true", "yes", "y", "是": + return "true" + case "false", "no", "n", "否": + return "false" + } + } + switch strings.ToLower(strings.TrimSpace(edge.Label)) { + case "true", "yes", "y", "是": + return "true" + case "false", "no", "n", "否": + return "false" + } + return "" +} + +func cfgString(cfg map[string]any, key string) string { + if cfg == nil { + return "" + } + if v, ok := cfg[key]; ok { + return strings.TrimSpace(fmt.Sprint(v)) + } + return "" +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if s := strings.TrimSpace(value); s != "" { + return s + } + } + return "" +} + +func truncateWorkflowPreview(s string, limit int) string { + s = strings.TrimSpace(s) + if limit <= 0 || len([]rune(s)) <= limit { + return s + } + runes := []rune(s) + return string(runes[:limit]) + "..." +} + +func renderWorkflowResponse(roleName, workflowName string, version int, runID string, state *WorkflowLocalState) string { + var sb strings.Builder + sb.WriteString(fmt.Sprintf("角色「%s」已完成工作流「%s」（版本 %d）。\n\n", roleName, workflowName, version)) + sb.WriteString(fmt.Sprintf("运行 ID：%s\n", runID)) + sb.WriteString(fmt.Sprintf("已执行节点：%d", len(state.Executed))) + if len(state.Skipped) > 0 { + sb.WriteString(fmt.Sprintf("，跳过节点：%d", len(state.Skipped))) + } + sb.WriteString("\n\n") + if len(state.Outputs) > 0 { + sb.WriteString("输出：\n") + keys := make([]string, 0, len(state.Outputs)) + for k := range state.Outputs { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + sb.WriteString(fmt.Sprintf("- %s：%v\n", k, state.Outputs[k])) + } + } else { + sb.WriteString("暂无输出。请检查是否配置了输出节点，或条件分支是否命中。\n") + } + if len(state.Skipped) > 0 { + sb.WriteString("\n未执行的节点类型仍会保留运行记录：") + sb.WriteString(strings.Join(state.Skipped, "、")) + sb.WriteString("。") + } + return strings.TrimSpace(sb.String()) +} diff --git a/internal/workflow/structured_outputs.go b/internal/workflow/structured_outputs.go new file mode 100644 index 00000000..1c3200a5 --- /dev/null +++ b/internal/workflow/structured_outputs.go @@ -0,0 +1,150 @@ +package workflow + +type NodeOutputEnvelope struct { + Kind string `json:"kind"` + NodeID string `json:"node_id"` + NodeType string `json:"node_type"` + Status string `json:"status"` + Output any `json:"output"` +} + +type StartOutput struct { + NodeOutputEnvelope + Message any `json:"message"` + ConversationID any `json:"conversationId"` + ProjectID any `json:"projectId"` +} + +type ConditionOutput struct { + NodeOutputEnvelope + Condition string `json:"condition"` + Matched bool `json:"matched"` +} + +type ToolOutput struct { + NodeOutputEnvelope + ToolName string `json:"tool_name"` + Arguments map[string]any `json:"arguments"` + ExecutionID string `json:"execution_id"` + IsError bool `json:"is_error"` +} + +type AgentOutput struct { + NodeOutputEnvelope + Mode string `json:"mode"` + MCPExecutionIDs []string `json:"mcp_execution_ids"` +} + +type HITLOutput struct { + NodeOutputEnvelope + Prompt string `json:"prompt"` + Reviewer string `json:"reviewer"` + Approved bool `json:"approved"` + Mode string `json:"mode"` +} + +type OutputNodeOutput struct { + NodeOutputEnvelope + OutputKey string `json:"output_key"` + Outputs map[string]any `json:"outputs"` +} + +func envelope(kind, nodeID, nodeType, status string, output any) NodeOutputEnvelope { + return NodeOutputEnvelope{Kind: kind, NodeID: nodeID, NodeType: nodeType, Status: status, Output: output} +} + +func outputMap(env NodeOutputEnvelope, extra map[string]any) map[string]any { + out := map[string]any{ + "kind": env.Kind, + "node_id": env.NodeID, + "node_type": env.NodeType, + "status": env.Status, + "output": env.Output, + "typed": env, + } + for k, v := range extra { + out[k] = v + } + return out +} + +func startOutputMap(node graphNode, message, conversationID, projectID any) map[string]any { + typed := StartOutput{ + NodeOutputEnvelope: envelope("start", node.ID, node.Type, "completed", message), + Message: message, + ConversationID: conversationID, + ProjectID: projectID, + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{ + "message": typed.Message, + "conversationId": typed.ConversationID, + "projectId": typed.ProjectID, + "typed": typed, + }) +} + +func conditionOutputMap(node graphNode, expr string, matched bool) map[string]any { + typed := ConditionOutput{ + NodeOutputEnvelope: envelope("condition", node.ID, node.Type, "completed", matched), + Condition: expr, + Matched: matched, + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{"condition": expr, "matched": matched, "typed": typed}) +} + +func outputNodeOutputMap(node graphNode, key string, value any) map[string]any { + typed := OutputNodeOutput{ + NodeOutputEnvelope: envelope("output", node.ID, node.Type, "completed", value), + OutputKey: key, + Outputs: map[string]any{key: value}, + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{"output_key": key, "outputs": typed.Outputs, "typed": typed}) +} + +func endOutputMap(node graphNode, value any) map[string]any { + typed := envelope("end", node.ID, node.Type, "completed", value) + return outputMap(typed, nil) +} + +func toolOutputMap(node graphNode, output string, toolName string, args map[string]any, executionID string, isError bool) map[string]any { + typed := ToolOutput{ + NodeOutputEnvelope: envelope("tool", node.ID, node.Type, "completed", output), + ToolName: toolName, + Arguments: args, + ExecutionID: executionID, + IsError: isError, + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{ + "tool_name": toolName, + "arguments": args, + "execution_id": executionID, + "is_error": isError, + "typed": typed, + }) +} + +func agentOutputMap(node graphNode, response, mode string, mcpIDs []string) map[string]any { + typed := AgentOutput{ + NodeOutputEnvelope: envelope("agent", node.ID, node.Type, "completed", response), + Mode: mode, + MCPExecutionIDs: mcpIDs, + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{"mode": mode, "mcp_execution_ids": mcpIDs, "typed": typed}) +} + +func hitlOutputMap(node graphNode, status string, output string, prompt string, reviewer string, approved bool) map[string]any { + typed := HITLOutput{ + NodeOutputEnvelope: envelope("hitl", node.ID, node.Type, status, output), + Prompt: prompt, + Reviewer: reviewer, + Approved: approved, + Mode: "interactive", + } + return outputMap(typed.NodeOutputEnvelope, map[string]any{ + "prompt": prompt, + "reviewer": reviewer, + "approved": approved, + "mode": "interactive", + "typed": typed, + }) +} diff --git a/internal/workflow/tool_output_budget_test.go b/internal/workflow/tool_output_budget_test.go new file mode 100644 index 00000000..d66d27eb --- /dev/null +++ b/internal/workflow/tool_output_budget_test.go @@ -0,0 +1,23 @@ +package workflow + +import ( + "strings" + "testing" +) + +func TestTruncateWorkflowToolOutputBoundsBytesAndKeepsExecutionReference(t *testing.T) { + out := truncateWorkflowToolOutput(strings.Repeat("响应正文", 1000), 256, "exec-123") + if len(out) > 256 { + t.Fatalf("workflow output bytes=%d, want <=256", len(out)) + } + if !strings.Contains(out, "exec-123") || !strings.Contains(out, "truncated") { + t.Fatalf("missing truncation reference: %q", out) + } +} + +func TestTruncateWorkflowToolOutputLeavesBoundedContentUntouched(t *testing.T) { + const want = "small-result" + if got := truncateWorkflowToolOutput(want, 256, "exec-123"); got != want { + t.Fatalf("got %q want %q", got, want) + } +} diff --git a/internal/workflow/types.go b/internal/workflow/types.go new file mode 100644 index 00000000..218c024b --- /dev/null +++ b/internal/workflow/types.go @@ -0,0 +1,74 @@ +package workflow + +import ( + "fmt" + "strconv" +) + +// WorkflowInput is the typed entry for Eino compose.Workflow[I,O]. +type WorkflowInput struct { + Message string `json:"message"` + ConversationID string `json:"conversationId"` + ProjectID string `json:"projectId"` + Role string `json:"role"` + WorkflowID string `json:"workflowId"` + WorkflowVersion int `json:"workflowVersion"` +} + +// WorkflowOutput aggregates terminal node payloads keyed by canvas node id. +type WorkflowOutput map[string]any + +// WorkflowNodeOutput is the per-node lambda payload (alias for Eino edge type alignment). +type WorkflowNodeOutput = map[string]interface{} + +func workflowInputFromMap(m map[string]interface{}) WorkflowInput { + in := WorkflowInput{} + if m == nil { + return in + } + if v, ok := m["message"].(string); ok { + in.Message = v + } else if m["message"] != nil { + in.Message = fmt.Sprint(m["message"]) + } + if v, ok := m["conversationId"].(string); ok { + in.ConversationID = v + } + if v, ok := m["projectId"].(string); ok { + in.ProjectID = v + } + if v, ok := m["role"].(string); ok { + in.Role = v + } + if v, ok := m["workflowId"].(string); ok { + in.WorkflowID = v + } + switch v := m["workflowVersion"].(type) { + case int: + in.WorkflowVersion = v + case int64: + in.WorkflowVersion = int(v) + case float64: + in.WorkflowVersion = int(v) + case string: + if n, err := strconv.Atoi(v); err == nil { + in.WorkflowVersion = n + } + } + return in +} + +func (in WorkflowInput) toStateInputs() map[string]any { + return map[string]any{ + "message": in.Message, + "conversationId": in.ConversationID, + "projectId": in.ProjectID, + "role": in.Role, + "workflowId": in.WorkflowID, + "workflowVersion": in.WorkflowVersion, + } +} + +func cacheKey(workflowID string, version int) string { + return workflowID + ":" + strconv.Itoa(version) +} diff --git a/internal/workflow/validation.go b/internal/workflow/validation.go new file mode 100644 index 00000000..70502778 --- /dev/null +++ b/internal/workflow/validation.go @@ -0,0 +1,366 @@ +package workflow + +import ( + "fmt" + "strconv" + "strings" +) + +var allowedWorkflowNodeTypes = map[string]bool{ + "start": true, + "tool": true, + "agent": true, + "condition": true, + "hitl": true, + "output": true, + "end": true, +} + +func validateGraphDefinition(g *graphDef, idx *graphIndex) error { + if g == nil || idx == nil { + return fmt.Errorf("工作流图为空") + } + if err := validateNodeIDsAndTypes(g); err != nil { + return err + } + if err := validateEdges(g, idx); err != nil { + return err + } + if err := validateNodeTopology(idx); err != nil { + return err + } + if err := validateNodeConfigs(idx); err != nil { + return err + } + if err := validateDAG(idx); err != nil { + return err + } + if err := validateReachability(idx); err != nil { + return err + } + return nil +} + +func validateNodeIDsAndTypes(g *graphDef) error { + seen := make(map[string]bool, len(g.Nodes)) + for _, node := range g.Nodes { + id := strings.TrimSpace(node.ID) + if id == "" { + return fmt.Errorf("工作流存在空节点 ID") + } + if seen[id] { + return fmt.Errorf("工作流存在重复节点 ID: %s", id) + } + seen[id] = true + nodeType := strings.ToLower(strings.TrimSpace(node.Type)) + if nodeType == "" { + return fmt.Errorf("节点「%s」缺少节点类型", id) + } + if !allowedWorkflowNodeTypes[nodeType] { + return fmt.Errorf("节点「%s」使用了未知节点类型: %s", id, node.Type) + } + } + return nil +} + +func validateEdges(g *graphDef, idx *graphIndex) error { + seen := make(map[string]bool, len(g.Edges)) + for _, edge := range g.Edges { + if id := strings.TrimSpace(edge.ID); id != "" { + if seen[id] { + return fmt.Errorf("工作流存在重复连线 ID: %s", id) + } + seen[id] = true + } + source := strings.TrimSpace(edge.Source) + target := strings.TrimSpace(edge.Target) + if source == "" || target == "" { + return fmt.Errorf("工作流存在源或目标为空的连线") + } + if source == target { + return fmt.Errorf("连线「%s」不能自环", firstNonEmpty(edge.ID, source)) + } + if _, ok := idx.nodes[source]; !ok { + return fmt.Errorf("连线「%s」引用了不存在的源节点: %s", firstNonEmpty(edge.ID, source), source) + } + if _, ok := idx.nodes[target]; !ok { + return fmt.Errorf("连线「%s」引用了不存在的目标节点: %s", firstNonEmpty(edge.ID, target), target) + } + } + return nil +} + +func validateNodeTopology(idx *graphIndex) error { + starts := explicitStartNodeIDs(idx) + if len(starts) == 0 { + return fmt.Errorf("工作流至少需要一个开始节点") + } + outputs := outputNodeIDs(idx) + if len(outputs) == 0 { + return fmt.Errorf("工作流至少需要一个输出节点") + } + for id, node := range idx.nodes { + inDegree := len(idx.incoming[id]) + outDegree := len(idx.outgoing[id]) + nodeType := strings.ToLower(strings.TrimSpace(node.Type)) + switch nodeType { + case "start": + if inDegree > 0 { + return fmt.Errorf("开始节点「%s」不能有入边", firstNonEmpty(node.Label, id)) + } + if outDegree == 0 { + return fmt.Errorf("开始节点「%s」至少需要一条出边", firstNonEmpty(node.Label, id)) + } + case "output", "end": + if outDegree > 0 { + return fmt.Errorf("%s 节点「%s」不能有出边", displayNodeType(nodeType), firstNonEmpty(node.Label, id)) + } + if inDegree == 0 { + return fmt.Errorf("%s 节点「%s」至少需要一条入边", displayNodeType(nodeType), firstNonEmpty(node.Label, id)) + } + default: + if inDegree == 0 { + return fmt.Errorf("节点「%s」不可达：非开始节点必须有入边", firstNonEmpty(node.Label, id)) + } + if outDegree == 0 { + return fmt.Errorf("节点「%s」没有出边；请连接到 output/end 节点", firstNonEmpty(node.Label, id)) + } + } + } + return nil +} + +func validateNodeConfigs(idx *graphIndex) error { + for id, node := range idx.nodes { + label := firstNonEmpty(node.Label, id) + switch strings.ToLower(strings.TrimSpace(node.Type)) { + case "tool": + if cfgString(node.Config, "tool_name") == "" { + return fmt.Errorf("工具节点「%s」必须选择 MCP 工具", label) + } + if err := validateToolConfig(node); err != nil { + return err + } + case "agent": + if cfgString(node.Config, "instruction") == "" { + if _, ok := parseFieldBinding(node.Config, "input_binding"); !ok { + return fmt.Errorf("Agent 节点「%s」必须填写节点指令或输入绑定", label) + } + } + if cfgString(node.Config, "output_key") == "" { + return fmt.Errorf("Agent 节点「%s」必须填写输出变量名", label) + } + case "condition": + if cfgString(node.Config, "expression") == "" { + return fmt.Errorf("条件节点「%s」必须填写表达式", label) + } + if err := validateConditionExpression(cfgString(node.Config, "expression")); err != nil { + return fmt.Errorf("条件节点「%s」表达式非法: %w", label, err) + } + if n := len(idx.outgoing[id]); n < 1 || n > 2 { + return fmt.Errorf("条件节点「%s」需要 1 到 2 条出边（是/否）", label) + } + if err := validateConditionBranchLabels(idx, id, node); err != nil { + return err + } + case "output": + if cfgString(node.Config, "output_key") == "" { + return fmt.Errorf("输出节点「%s」必须填写输出变量名", label) + } + } + if err := validateJoinConfig(idx, id, node); err != nil { + return err + } + if hasConditionalOutgoingEdges(idx, id) { + if err := validateConditionalOutgoingEdges(idx, id, node); err != nil { + return err + } + } + } + return nil +} + +func validateConditionalOutgoingEdges(idx *graphIndex, nodeID string, node graphNode) error { + unconditional := 0 + for _, edge := range idx.outgoing[nodeID] { + cond := firstNonEmpty(cfgString(edge.Config, "condition"), cfgString(edge.Config, "expression")) + if cond != "" { + if err := validateConditionExpression(cond); err != nil { + return fmt.Errorf("节点「%s」的连线条件非法: %w", firstNonEmpty(node.Label, nodeID), err) + } + } + if cond == "" { + unconditional++ + } + } + if unconditional > 1 { + return fmt.Errorf("节点「%s」的条件出边最多只能有一条默认分支", firstNonEmpty(node.Label, nodeID)) + } + return nil +} + +func validateToolConfig(node graphNode) error { + rawArgs := cfgString(node.Config, "arguments") + if rawArgs != "" { + if _, err := resolveToolArguments(node.Config, &WorkflowLocalState{}); err != nil { + return fmt.Errorf("工具节点「%s」参数 JSON 非法: %w", firstNonEmpty(node.Label, node.ID), err) + } + } + if timeout := cfgString(node.Config, "timeout_seconds"); timeout != "" { + if _, err := parsePositiveInt(timeout); err != nil { + return fmt.Errorf("工具节点「%s」超时时间必须是正整数", firstNonEmpty(node.Label, node.ID)) + } + } + return nil +} + +func validateJoinConfig(idx *graphIndex, nodeID string, node graphNode) error { + strategy := joinStrategy(node) + if !allowedJoinStrategies[strategy] { + return fmt.Errorf("节点「%s」使用了未知汇聚策略: %s", firstNonEmpty(node.Label, nodeID), strategy) + } + if len(idx.incoming[nodeID]) > 1 && strategy == "" { + return fmt.Errorf("节点「%s」有多个上游时必须声明汇聚策略", firstNonEmpty(node.Label, nodeID)) + } + return nil +} + +func validateConditionBranchLabels(idx *graphIndex, nodeID string, node graphNode) error { + seen := map[string]bool{} + for _, edge := range idx.outgoing[nodeID] { + hint := conditionBranchHint(edge) + if hint == "" { + return fmt.Errorf("条件节点「%s」的出边必须标记为是/否或 true/false", firstNonEmpty(node.Label, nodeID)) + } + if seen[hint] { + return fmt.Errorf("条件节点「%s」存在重复分支标签: %s", firstNonEmpty(node.Label, nodeID), hint) + } + seen[hint] = true + } + return nil +} + +func validateDAG(idx *graphIndex) error { + color := make(map[string]int, len(idx.nodes)) + var visit func(string) error + visit = func(id string) error { + switch color[id] { + case 1: + return fmt.Errorf("工作流存在环路，Workflow 编排必须是 DAG: %s", id) + case 2: + return nil + } + color[id] = 1 + for _, edge := range idx.outgoing[id] { + if err := visit(edge.Target); err != nil { + return err + } + } + color[id] = 2 + return nil + } + for id := range idx.nodes { + if err := visit(id); err != nil { + return err + } + } + return nil +} + +func validateReachability(idx *graphIndex) error { + starts := explicitStartNodeIDs(idx) + reached := make(map[string]bool, len(idx.nodes)) + queue := append([]string(nil), starts...) + for len(queue) > 0 { + id := queue[0] + queue = queue[1:] + if reached[id] { + continue + } + reached[id] = true + for _, edge := range idx.outgoing[id] { + queue = append(queue, edge.Target) + } + } + for id, node := range idx.nodes { + if !reached[id] { + return fmt.Errorf("节点「%s」不可达：没有从开始节点连通到该节点", firstNonEmpty(node.Label, id)) + } + } + + canReachTerminal := make(map[string]bool, len(idx.nodes)) + visiting := make(map[string]bool, len(idx.nodes)) + var reachesTerminal func(string) bool + reachesTerminal = func(id string) bool { + if canReachTerminal[id] { + return true + } + if visiting[id] { + return false + } + visiting[id] = true + node := idx.nodes[id] + nodeType := strings.ToLower(strings.TrimSpace(node.Type)) + if nodeType == "output" || nodeType == "end" { + canReachTerminal[id] = true + visiting[id] = false + return true + } + for _, edge := range idx.outgoing[id] { + if reachesTerminal(edge.Target) { + canReachTerminal[id] = true + visiting[id] = false + return true + } + } + visiting[id] = false + return false + } + for id, node := range idx.nodes { + if !reachesTerminal(id) { + return fmt.Errorf("节点「%s」无法到达 output/end 终点", firstNonEmpty(node.Label, id)) + } + } + return nil +} + +func explicitStartNodeIDs(idx *graphIndex) []string { + var ids []string + for id, node := range idx.nodes { + if strings.EqualFold(node.Type, "start") { + ids = append(ids, id) + } + } + sortNodeIDsByCanvas(ids, idx.nodes) + return ids +} + +func outputNodeIDs(idx *graphIndex) []string { + var ids []string + for id, node := range idx.nodes { + if strings.EqualFold(node.Type, "output") { + ids = append(ids, id) + } + } + sortNodeIDsByCanvas(ids, idx.nodes) + return ids +} + +func displayNodeType(nodeType string) string { + switch strings.ToLower(strings.TrimSpace(nodeType)) { + case "output": + return "输出" + case "end": + return "结束" + default: + return nodeType + } +} + +func parsePositiveInt(s string) (int, error) { + n, err := strconv.Atoi(strings.TrimSpace(s)) + if err != nil || n <= 0 { + return 0, fmt.Errorf("not positive integer") + } + return n, nil +}