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/monitor/reconcile.go b/internal/monitor/reconcile.go new file mode 100644 index 00000000..526c0817 --- /dev/null +++ b/internal/monitor/reconcile.go @@ -0,0 +1,101 @@ +package monitor + +import ( + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +const ( + staleRunningMinAge = 45 * time.Second + staleRunningReconcileGap = 2 * time.Minute +) + +// ExecutionReconciler 在启动或运行期将无对应协程的 running 执行记录收尾为 orphaned。 +type ExecutionReconciler struct { + db *database.DB + mcpServer *mcp.Server + externalMgr *mcp.ExternalMCPManager + logger *zap.Logger +} + +// NewExecutionReconciler creates a reconciler for orphaned MCP tool executions. +func NewExecutionReconciler(db *database.DB, mcpServer *mcp.Server, externalMgr *mcp.ExternalMCPManager, logger *zap.Logger) *ExecutionReconciler { + return &ExecutionReconciler{ + db: db, + mcpServer: mcpServer, + externalMgr: externalMgr, + logger: logger, + } +} + +// ReconcileOnStartup marks every persisted running row as orphaned (safe right after process start). +func (r *ExecutionReconciler) ReconcileOnStartup() { + if r == nil || r.db == nil { + return + } + now := time.Now() + n, err := r.db.CancelOrphanedRunningToolExecutions(now, "执行已中断（服务重启）") + if err != nil { + if r.logger != nil { + r.logger.Warn("启动时清理孤儿 running 工具执行记录失败", zap.Error(err)) + } + return + } + if n > 0 && r.logger != nil { + r.logger.Info("启动时已收尾孤儿 running 工具执行记录", zap.Int64("count", n)) + } +} + +func (r *ExecutionReconciler) activeExecutionIDs() map[string]struct{} { + ids := make(map[string]struct{}) + if r.mcpServer != nil { + for id := range r.mcpServer.ActiveRunningExecutionIDs() { + ids[id] = struct{}{} + } + } + if r.externalMgr != nil { + for id := range r.externalMgr.ActiveRunningExecutionIDs() { + ids[id] = struct{}{} + } + } + return ids +} + +// ReconcileStaleRunning finalizes running rows that are not tracked in-memory and older than staleRunningMinAge. +func (r *ExecutionReconciler) ReconcileStaleRunning() { + if r == nil || r.db == nil { + return + } + now := time.Now() + n, err := r.db.FinalizeStaleRunningToolExecutions(now, staleRunningMinAge, r.activeExecutionIDs(), "执行已中断（会话已结束）") + if err != nil { + if r.logger != nil { + r.logger.Warn("定期收尾 stale running 工具执行记录失败", zap.Error(err)) + } + return + } + if n > 0 && r.logger != nil { + r.logger.Info("已收尾 stale running 工具执行记录", zap.Int64("count", n)) + } +} + +// StartStaleRunningReconcileLoop periodically reconciles orphaned running tool executions. +func StartStaleRunningReconcileLoop(r *ExecutionReconciler, logger *zap.Logger) { + if r == nil { + return + } + go func() { + ticker := time.NewTicker(staleRunningReconcileGap) + defer ticker.Stop() + for range ticker.C { + r.ReconcileStaleRunning() + if logger != nil { + logger.Debug("monitor stale running reconcile tick completed") + } + } + }() +} diff --git a/internal/monitor/reconcile_test.go b/internal/monitor/reconcile_test.go new file mode 100644 index 00000000..0dfad0e9 --- /dev/null +++ b/internal/monitor/reconcile_test.go @@ -0,0 +1,38 @@ +package monitor + +import ( + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +func TestExecutionReconciler_ReconcileOnStartup(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "monitor.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: "run-1", ToolName: "hydra", Status: "running", StartTime: time.Now().Add(-time.Hour), + }); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + r := NewExecutionReconciler(db, mcp.NewServer(zap.NewNop()), nil, zap.NewNop()) + r.ReconcileOnStartup() + + got, err := db.GetToolExecution("run-1") + if err != nil { + t.Fatalf("GetToolExecution: %v", err) + } + if got.Status != "orphaned" { + t.Fatalf("expected orphaned after startup reconcile, got %s", got.Status) + } +} diff --git a/internal/monitor/retention.go b/internal/monitor/retention.go new file mode 100644 index 00000000..d1ffb295 --- /dev/null +++ b/internal/monitor/retention.go @@ -0,0 +1,71 @@ +package monitor + +import ( + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +const retentionPurgeInterval = time.Hour + +// Service manages MCP tool execution monitor retention. +type Service struct { + db *database.DB + cfg *config.Config + logger *zap.Logger +} + +// NewService creates a monitor 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.MonitorConfig{}.RetentionDaysEffective() + } + return s.cfg.Monitor.RetentionDaysEffective() +} + +// PurgeExpired deletes tool execution rows older than retention_days when configured. +func (s *Service) PurgeExpired() { + if s == nil || s.db == nil || s.cfg == nil { + return + } + days := s.cfg.Monitor.RetentionDaysEffective() + if days <= 0 { + return + } + cutoff := time.Now().AddDate(0, 0, -days) + n, err := s.db.PurgeToolExecutionsBefore(cutoff) + if err != nil { + if s.logger != nil { + s.logger.Warn("清理过期 MCP 执行记录失败", zap.Error(err)) + } + return + } + if n > 0 && s.logger != nil { + s.logger.Info("已清理过期 MCP 执行记录", zap.Int64("deleted", n), zap.Int("retention_days", days)) + } +} + +// StartRetentionLoop periodically purges expired tool execution 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("monitor retention tick completed") + } + } + }() +} diff --git a/internal/monitor/retention_test.go b/internal/monitor/retention_test.go new file mode 100644 index 00000000..40425fd6 --- /dev/null +++ b/internal/monitor/retention_test.go @@ -0,0 +1,94 @@ +package monitor + +import ( + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +func TestServicePurgeExpired_respectsZeroRetention(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "monitor.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + + exec := &mcp.ToolExecution{ + ID: "ancient", + ToolName: "curl::get", + Arguments: map[string]interface{}{}, + Status: "completed", + StartTime: mustParseTime(t, "2020-01-01T00:00:00Z"), + } + if err := db.SaveToolExecution(exec); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + zero := 0 + svc := NewService(db, &config.Config{ + Monitor: config.MonitorConfig{RetentionDays: &zero}, + }, zap.NewNop()) + svc.PurgeExpired() + + if _, err := db.GetToolExecution("ancient"); err != nil { + t.Fatalf("record should remain when retention_days=0: %v", err) + } +} + +func TestServicePurgeExpired_deletesOldRows(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "monitor.db") + db, err := database.NewDB(dbPath, zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer db.Close() + + exec := &mcp.ToolExecution{ + ID: "ancient", + ToolName: "curl::get", + Arguments: map[string]interface{}{}, + Status: "completed", + StartTime: mustParseTime(t, "2020-01-01T00:00:00Z"), + } + if err := db.SaveToolExecution(exec); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + days := 90 + svc := NewService(db, &config.Config{ + Monitor: config.MonitorConfig{RetentionDays: &days}, + }, zap.NewNop()) + svc.PurgeExpired() + + if _, err := db.GetToolExecution("ancient"); err == nil { + t.Fatal("record should be purged when older than retention_days") + } +} + +func TestRetentionDaysEffective_defaults(t *testing.T) { + got := config.MonitorConfig{}.RetentionDaysEffective() + if got != 90 { + t.Fatalf("default = %d, want 90", got) + } + zero := 0 + cfg := config.MonitorConfig{RetentionDays: &zero} + if cfg.RetentionDaysEffective() != 0 { + t.Fatalf("zero = %d, want 0", cfg.RetentionDaysEffective()) + } +} + +func mustParseTime(t *testing.T, value string) time.Time { + t.Helper() + parsed, err := time.Parse(time.RFC3339, value) + if err != nil { + t.Fatalf("parse time: %v", err) + } + return parsed +} diff --git a/internal/robot/conn.go b/internal/robot/conn.go new file mode 100644 index 00000000..d57e361d --- /dev/null +++ b/internal/robot/conn.go @@ -0,0 +1,6 @@ +package robot + +// MessageHandler 供飞书/钉钉长连接调用的消息处理接口（由 handler.RobotHandler 实现） +type MessageHandler interface { + HandleMessage(platform, userID, text string) string +} diff --git a/internal/robot/ding.go b/internal/robot/ding.go new file mode 100644 index 00000000..7f469808 --- /dev/null +++ b/internal/robot/ding.go @@ -0,0 +1,151 @@ +package robot + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/open-dingtalk/dingtalk-stream-sdk-go/chatbot" + "github.com/open-dingtalk/dingtalk-stream-sdk-go/client" + dingutils "github.com/open-dingtalk/dingtalk-stream-sdk-go/utils" + "go.uber.org/zap" +) + +const ( + dingReconnectInitial = 5 * time.Second // 首次重连间隔 + dingReconnectMax = 60 * time.Second // 最大重连间隔 +) + +// StartDing 启动钉钉 Stream 长连接（无需公网），收到消息后调用 handler 并通过 SessionWebhook 回复。 +// 断线（如笔记本睡眠、网络中断）后会自动重连；ctx 被取消时退出，便于配置变更时重启。 +func StartDing(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.Dingtalk + if !cfg.Enabled || cfg.ClientID == "" || cfg.ClientSecret == "" { + return + } + go runDingLoop(ctx, cfg, robotsCfg.Session.StrictUserIdentityEnabled(), h, logger) +} + +// runDingLoop 循环维持钉钉长连接：断开且 ctx 未取消时按退避间隔重连。 +func runDingLoop(ctx context.Context, cfg config.RobotDingtalkConfig, strictUserIdentity bool, h MessageHandler, logger *zap.Logger) { + backoff := dingReconnectInitial + for { + streamClient := client.NewStreamClient( + client.WithAppCredential(client.NewAppCredentialConfig(cfg.ClientID, cfg.ClientSecret)), + client.WithSubscription(dingutils.SubscriptionTypeKCallback, "/v1.0/im/bot/messages/get", + chatbot.NewDefaultChatBotFrameHandler(func(ctx context.Context, msg *chatbot.BotCallbackDataModel) ([]byte, error) { + go handleDingMessage(ctx, msg, cfg, strictUserIdentity, h, logger) + return nil, nil + }).OnEventReceived), + ) + logger.Info("钉钉 Stream 正在连接…", zap.String("client_id", cfg.ClientID)) + err := streamClient.Start(ctx) + if ctx.Err() != nil { + logger.Info("钉钉 Stream 已按配置重启关闭") + return + } + if err != nil { + logger.Warn("钉钉 Stream 长连接断开（如睡眠/断网），将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + // 下次重连间隔递增，上限 60 秒，避免频繁重试 + if backoff < dingReconnectMax { + backoff *= 2 + if backoff > dingReconnectMax { + backoff = dingReconnectMax + } + } + } + } +} + +func handleDingMessage(ctx context.Context, msg *chatbot.BotCallbackDataModel, cfg config.RobotDingtalkConfig, strictUserIdentity bool, h MessageHandler, logger *zap.Logger) { + if msg == nil || msg.SessionWebhook == "" { + return + } + content := "" + if msg.Text.Content != "" { + content = strings.TrimSpace(msg.Text.Content) + } + if content == "" && msg.Msgtype == "richText" { + if cMap, ok := msg.Content.(map[string]interface{}); ok { + if rich, ok := cMap["richText"].([]interface{}); ok { + for _, c := range rich { + if m, ok := c.(map[string]interface{}); ok { + if txt, ok := m["text"].(string); ok { + content = strings.TrimSpace(txt) + break + } + } + } + } + } + } + if content == "" { + logger.Debug("钉钉消息内容为空，已忽略", zap.String("msgtype", msg.Msgtype)) + return + } + logger.Info("钉钉收到消息", zap.String("sender", msg.SenderId), zap.String("content", content)) + tenantKey := strings.TrimSpace(cfg.ClientID) + if tenantKey == "" { + tenantKey = "default" + } + userID := strings.TrimSpace(msg.SenderId) + if userID != "" { + userID = "t:" + tenantKey + "|u:" + userID + } else if cfg.AllowConversationIDFallback && !strictUserIdentity { + conversationID := strings.TrimSpace(msg.ConversationId) + if conversationID != "" { + userID = "t:" + tenantKey + "|c:" + conversationID + } + } + if userID == "" { + logger.Warn("钉钉消息缺少可用用户标识，已忽略") + return + } + reply := h.HandleMessage("dingtalk", userID, content) + // 使用 markdown 类型以便正确展示标题、列表、代码块等格式 + title := reply + if idx := strings.IndexAny(reply, "\n"); idx > 0 { + title = strings.TrimSpace(reply[:idx]) + } + if len(title) > 50 { + title = title[:50] + "…" + } + if title == "" { + title = "回复" + } + body := map[string]interface{}{ + "msgtype": "markdown", + "markdown": map[string]string{ + "title": title, + "text": reply, + }, + } + bodyBytes, _ := json.Marshal(body) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, msg.SessionWebhook, bytes.NewReader(bodyBytes)) + if err != nil { + logger.Warn("钉钉构造回复请求失败", zap.Error(err)) + return + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + logger.Warn("钉钉回复请求失败", zap.Error(err)) + return + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + logger.Warn("钉钉回复非 200", zap.Int("status", resp.StatusCode)) + return + } + logger.Debug("钉钉回复成功", zap.String("content_preview", reply)) +} diff --git a/internal/robot/discord.go b/internal/robot/discord.go new file mode 100644 index 00000000..4986e524 --- /dev/null +++ b/internal/robot/discord.go @@ -0,0 +1,121 @@ +package robot + +import ( + "context" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/bwmarrin/discordgo" + "go.uber.org/zap" +) + +const ( + discordPlatform = "discord" + discordMaxMessageRunes = 2000 +) + +// StartDiscord 启动 Discord Gateway（WebSocket，无需公网回调）。 +func StartDiscord(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.Discord + if !cfg.Enabled || strings.TrimSpace(cfg.BotToken) == "" { + return + } + go runDiscordLoop(ctx, cfg, h, logger) +} + +func runDiscordLoop(ctx context.Context, cfg config.RobotDiscordConfig, h MessageHandler, logger *zap.Logger) { + backoff := reconnectInitial + for { + err := runDiscordSession(ctx, cfg, h, logger) + if ctx.Err() != nil { + logger.Info("Discord Gateway 已按配置关闭") + return + } + if err != nil { + logger.Warn("Discord Gateway 异常，将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + if !waitReconnect(ctx, &backoff) { + return + } + } +} + +func runDiscordSession(ctx context.Context, cfg config.RobotDiscordConfig, h MessageHandler, logger *zap.Logger) error { + token := strings.TrimSpace(cfg.BotToken) + if !strings.HasPrefix(token, "Bot ") { + token = "Bot " + token + } + session, err := discordgo.New(token) + if err != nil { + return err + } + session.Identify.Intents = discordgo.IntentsGuildMessages | + discordgo.IntentsDirectMessages | + discordgo.IntentMessageContent + + session.AddHandler(func(s *discordgo.Session, m *discordgo.MessageCreate) { + if m == nil || m.Author == nil || m.Author.Bot { + return + } + text := strings.TrimSpace(m.Content) + if text == "" { + return + } + if m.GuildID != "" { + if !cfg.AllowGuildMessages { + return + } + if s.State.User == nil || !discordMentionsBot(m, s.State.User.ID) { + return + } + } + userID := discordSessionKey(m.GuildID, m.Author.ID) + logger.Info("Discord 收到消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(discordPlatform, userID, text) + discordPostReply(s, m.ChannelID, reply, logger) + }) + + if err := session.Open(); err != nil { + return err + } + logger.Info("Discord Gateway 已连接，等待收消息") + defer session.Close() + + <-ctx.Done() + return ctx.Err() +} + +func discordMentionsBot(m *discordgo.MessageCreate, botUserID string) bool { + if m == nil || botUserID == "" { + return false + } + for _, mention := range m.Mentions { + if mention != nil && mention.ID == botUserID { + return true + } + } + return strings.Contains(m.Content, "<@"+botUserID+">") || strings.Contains(m.Content, "<@!"+botUserID+">") +} + +func discordSessionKey(guildID, userID string) string { + guildID = strings.TrimSpace(guildID) + userID = strings.TrimSpace(userID) + if guildID == "" { + return "u:" + userID + } + return "g:" + guildID + "|u:" + userID +} + +func discordPostReply(s *discordgo.Session, channelID, reply string, logger *zap.Logger) { + reply = trimReply(reply) + if reply == "" { + return + } + for _, chunk := range splitTextChunks(reply, discordMaxMessageRunes) { + if _, err := s.ChannelMessageSend(channelID, chunk); err != nil { + logger.Warn("Discord 发送回复失败", zap.String("channel", channelID), zap.Error(err)) + return + } + } +} diff --git a/internal/robot/ilink/client.go b/internal/robot/ilink/client.go new file mode 100644 index 00000000..00abafdb --- /dev/null +++ b/internal/robot/ilink/client.go @@ -0,0 +1,316 @@ +package ilink + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" +) + +const ( + DefaultBaseURL = "https://ilinkai.weixin.qq.com" + DefaultBotType = "3" + DefaultBotAgent = "CyberStrikeAI/1.0" + ILinkAppID = "bot" + QRLongPollTimeout = 35 * time.Second + APIDefaultTimeout = 15 * time.Second + GetUpdatesTimeout = 35 * time.Second +) + +// Client 微信 iLink Bot HTTP 客户端（与 @tencent-weixin/openclaw-weixin 协议兼容） +type Client struct { + BaseURL string + BotToken string + BotAgent string + ClientVersion uint32 + HTTP *http.Client +} + +func NewClient(baseURL, botToken, botAgent string, clientVersion uint32) *Client { + base := strings.TrimSpace(baseURL) + if base == "" { + base = DefaultBaseURL + } + agent := strings.TrimSpace(botAgent) + if agent == "" { + agent = DefaultBotAgent + } + return &Client{ + BaseURL: strings.TrimRight(base, "/"), + BotToken: strings.TrimSpace(botToken), + BotAgent: sanitizeBotAgent(agent), + ClientVersion: clientVersion, + HTTP: &http.Client{Timeout: 0}, + } +} + +// BuildClientVersion 将 semver 编码为 iLink-App-ClientVersion（0x00MMNNPP） +func BuildClientVersion(version string) uint32 { + parts := strings.Split(version, ".") + parse := func(i int) int { + if i >= len(parts) { + return 0 + } + n, _ := strconv.Atoi(strings.TrimSpace(parts[i])) + if n < 0 { + return 0 + } + return n + } + major := parse(0) & 0xff + minor := parse(1) & 0xff + patch := parse(2) & 0xff + return uint32((major << 16) | (minor << 8) | patch) +} + +type baseInfo struct { + ChannelVersion string `json:"channel_version"` + BotAgent string `json:"bot_agent"` +} + +func (c *Client) buildBaseInfo() baseInfo { + return baseInfo{ + ChannelVersion: "1.0.0", + BotAgent: c.BotAgent, + } +} + +func randomWechatUIN() string { + var b [4]byte + _, _ = rand.Read(b[:]) + u := uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3]) + return base64.StdEncoding.EncodeToString([]byte(strconv.FormatUint(uint64(u), 10))) +} + +func (c *Client) commonHeaders() http.Header { + h := http.Header{} + h.Set("iLink-App-Id", ILinkAppID) + h.Set("iLink-App-ClientVersion", strconv.FormatUint(uint64(c.ClientVersion), 10)) + return h +} + +func (c *Client) authHeaders() http.Header { + h := c.commonHeaders() + h.Set("Content-Type", "application/json") + h.Set("AuthorizationType", "ilink_bot_token") + h.Set("X-WECHAT-UIN", randomWechatUIN()) + if c.BotToken != "" { + h.Set("Authorization", "Bearer "+c.BotToken) + } + return h +} + +func (c *Client) endpointURL(path string) (string, error) { + u, err := url.Parse(c.BaseURL + "/") + if err != nil { + return "", err + } + ref, err := url.Parse(path) + if err != nil { + return "", err + } + return u.ResolveReference(ref).String(), nil +} + +func (c *Client) doRequest(ctx context.Context, method, path string, body []byte, headers http.Header, timeout time.Duration) ([]byte, error) { + reqURL, err := c.endpointURL(path) + if err != nil { + return nil, err + } + var bodyReader io.Reader + if len(body) > 0 { + bodyReader = bytes.NewReader(body) + } + req, err := http.NewRequestWithContext(ctx, method, reqURL, bodyReader) + if err != nil { + return nil, err + } + for k, vs := range headers { + for _, v := range vs { + req.Header.Add(k, v) + } + } + client := c.HTTP + if client == nil { + client = http.DefaultClient + } + if timeout > 0 { + ctx2, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + req = req.WithContext(ctx2) + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("ilink %s %s: %d %s", method, path, resp.StatusCode, string(raw)) + } + return raw, nil +} + +// QRCodeResponse 获取二维码响应 +type QRCodeResponse struct { + QRCode string `json:"qrcode"` + QRCodeImgContent string `json:"qrcode_img_content"` +} + +// GetBotQRCode 获取绑定二维码 +func (c *Client) GetBotQRCode(ctx context.Context, botType string, localTokenList []string) (*QRCodeResponse, error) { + if strings.TrimSpace(botType) == "" { + botType = DefaultBotType + } + body, _ := json.Marshal(map[string]interface{}{ + "local_token_list": localTokenList, + }) + path := "ilink/bot/get_bot_qrcode?bot_type=" + url.QueryEscape(botType) + raw, err := c.doRequest(ctx, http.MethodPost, path, body, c.authHeaders(), APIDefaultTimeout) + if err != nil { + return nil, err + } + var out QRCodeResponse + if err := json.Unmarshal(raw, &out); err != nil { + return nil, err + } + return &out, nil +} + +// QRStatusResponse 二维码状态轮询响应 +type QRStatusResponse struct { + Status string `json:"status"` + BotToken string `json:"bot_token"` + ILinkBotID string `json:"ilink_bot_id"` + ILinkUserID string `json:"ilink_user_id"` + BaseURL string `json:"baseurl"` + RedirectHost string `json:"redirect_host"` +} + +// GetQRCodeStatus 长轮询二维码扫码状态 +func (c *Client) GetQRCodeStatus(ctx context.Context, qrcode, verifyCode string) (*QRStatusResponse, error) { + path := "ilink/bot/get_qrcode_status?qrcode=" + url.QueryEscape(qrcode) + if verifyCode != "" { + path += "&verify_code=" + url.QueryEscape(verifyCode) + } + raw, err := c.doRequest(ctx, http.MethodGet, path, nil, c.commonHeaders(), QRLongPollTimeout) + if err != nil { + if ctx.Err() != nil { + return &QRStatusResponse{Status: "wait"}, nil + } + return &QRStatusResponse{Status: "wait"}, nil + } + var out QRStatusResponse + if err := json.Unmarshal(raw, &out); err != nil { + return nil, err + } + return &out, nil +} + +// MessageItem 消息内容项 +type MessageItem struct { + Type int `json:"type"` + TextItem *struct { + Text string `json:"text"` + } `json:"text_item,omitempty"` +} + +// WeixinMessage 入站消息 +type WeixinMessage struct { + FromUserID string `json:"from_user_id"` + MessageType int `json:"message_type"` + MessageState int `json:"message_state"` + ItemList []MessageItem `json:"item_list"` + ContextToken string `json:"context_token"` +} + +// GetUpdatesResponse 长轮询消息响应 +type GetUpdatesResponse struct { + Ret int `json:"ret"` + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` + Msgs []WeixinMessage `json:"msgs"` + GetUpdatesBuf string `json:"get_updates_buf"` + LongPollingTimeoutMs int `json:"longpolling_timeout_ms"` +} + +// GetUpdates 长轮询获取新消息 +func (c *Client) GetUpdates(ctx context.Context, getUpdatesBuf string) (*GetUpdatesResponse, error) { + body, _ := json.Marshal(map[string]interface{}{ + "get_updates_buf": getUpdatesBuf, + "base_info": c.buildBaseInfo(), + }) + raw, err := c.doRequest(ctx, http.MethodPost, "ilink/bot/getupdates", body, c.authHeaders(), GetUpdatesTimeout) + if err != nil { + if ctx.Err() != nil { + return &GetUpdatesResponse{Ret: 0, GetUpdatesBuf: getUpdatesBuf}, nil + } + return &GetUpdatesResponse{Ret: 0, GetUpdatesBuf: getUpdatesBuf}, nil + } + var out GetUpdatesResponse + if err := json.Unmarshal(raw, &out); err != nil { + return nil, err + } + return &out, nil +} + +// SendTextMessage 发送文本回复 +func (c *Client) SendTextMessage(ctx context.Context, toUserID, contextToken, text, clientID string) error { + if clientID == "" { + clientID = randomClientID() + } + payload := map[string]interface{}{ + "msg": map[string]interface{}{ + "to_user_id": toUserID, + "client_id": clientID, + "message_type": 2, + "message_state": 2, + "context_token": contextToken, + "item_list": []map[string]interface{}{ + {"type": 1, "text_item": map[string]string{"text": text}}, + }, + }, + "base_info": c.buildBaseInfo(), + } + body, _ := json.Marshal(payload) + _, err := c.doRequest(ctx, http.MethodPost, "ilink/bot/sendmessage", body, c.authHeaders(), APIDefaultTimeout) + return err +} + +func randomClientID() string { + var b [8]byte + _, _ = rand.Read(b[:]) + return fmt.Sprintf("%x", b) +} + +func sanitizeBotAgent(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return DefaultBotAgent + } + if len(raw) > 256 { + return raw[:256] + } + return raw +} + +// ExtractText 从消息中提取首条文本 +func ExtractText(msg WeixinMessage) string { + for _, item := range msg.ItemList { + if item.Type == 1 && item.TextItem != nil { + return strings.TrimSpace(item.TextItem.Text) + } + } + return "" +} diff --git a/internal/robot/ilink/qrcode_image.go b/internal/robot/ilink/qrcode_image.go new file mode 100644 index 00000000..0ef6521f --- /dev/null +++ b/internal/robot/ilink/qrcode_image.go @@ -0,0 +1,26 @@ +package ilink + +import ( + "encoding/base64" + "fmt" + "strings" + + "github.com/skip2/go-qrcode" +) + +// QRCodeDataURL 将扫码内容（一般为 liteapp 链接）编码为 PNG data URL，供 Web 端展示。 +// qrcode_img_content 不是图片直链，不能用作 。 +func QRCodeDataURL(content string, size int) (string, error) { + content = strings.TrimSpace(content) + if content == "" { + return "", fmt.Errorf("empty qr content") + } + if size <= 0 { + size = 256 + } + png, err := qrcode.Encode(content, qrcode.Medium, size) + if err != nil { + return "", err + } + return "data:image/png;base64," + base64.StdEncoding.EncodeToString(png), nil +} diff --git a/internal/robot/lark.go b/internal/robot/lark.go new file mode 100644 index 00000000..2cda0601 --- /dev/null +++ b/internal/robot/lark.go @@ -0,0 +1,141 @@ +package robot + +import ( + "context" + "encoding/json" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + lark "github.com/larksuite/oapi-sdk-go/v3" + larkcore "github.com/larksuite/oapi-sdk-go/v3/core" + "github.com/larksuite/oapi-sdk-go/v3/event/dispatcher" + larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1" + larkws "github.com/larksuite/oapi-sdk-go/v3/ws" + "go.uber.org/zap" +) + +const ( + larkReconnectInitial = 5 * time.Second // 首次重连间隔 + larkReconnectMax = 60 * time.Second // 最大重连间隔 +) + +type larkTextContent struct { + Text string `json:"text"` +} + +// StartLark 启动飞书长连接（无需公网），收到消息后调用 handler 并回复。 +// 断线（如笔记本睡眠、网络中断）后会自动重连；ctx 被取消时退出，便于配置变更时重启。 +func StartLark(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.Lark + if !cfg.Enabled || cfg.AppID == "" || cfg.AppSecret == "" { + return + } + go runLarkLoop(ctx, cfg, robotsCfg.Session.StrictUserIdentityEnabled(), h, logger) +} + +// runLarkLoop 循环维持飞书长连接：断开且 ctx 未取消时按退避间隔重连。 +func runLarkLoop(ctx context.Context, cfg config.RobotLarkConfig, strictUserIdentity bool, h MessageHandler, logger *zap.Logger) { + backoff := larkReconnectInitial + for { + larkClient := lark.NewClient(cfg.AppID, cfg.AppSecret) + eventHandler := dispatcher.NewEventDispatcher("", "").OnP2MessageReceiveV1(func(ctx context.Context, event *larkim.P2MessageReceiveV1) error { + go handleLarkMessage(ctx, event, cfg, strictUserIdentity, h, larkClient, logger) + return nil + }) + wsClient := larkws.NewClient(cfg.AppID, cfg.AppSecret, + larkws.WithEventHandler(eventHandler), + larkws.WithLogLevel(larkcore.LogLevelInfo), + ) + logger.Info("飞书长连接正在连接…", zap.String("app_id", cfg.AppID)) + err := wsClient.Start(ctx) + if ctx.Err() != nil { + logger.Info("飞书长连接已按配置重启关闭") + return + } + if err != nil { + logger.Warn("飞书长连接断开（如睡眠/断网），将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + if backoff < larkReconnectMax { + backoff *= 2 + if backoff > larkReconnectMax { + backoff = larkReconnectMax + } + } + } + } +} + +func handleLarkMessage(ctx context.Context, event *larkim.P2MessageReceiveV1, cfg config.RobotLarkConfig, strictUserIdentity bool, h MessageHandler, client *lark.Client, logger *zap.Logger) { + if event == nil || event.Event == nil || event.Event.Message == nil || event.Event.Sender == nil || event.Event.Sender.SenderId == nil { + return + } + msg := event.Event.Message + msgType := larkcore.StringValue(msg.MessageType) + if msgType != larkim.MsgTypeText { + logger.Debug("飞书暂仅处理文本消息", zap.String("msg_type", msgType)) + return + } + var textBody larkTextContent + if err := json.Unmarshal([]byte(larkcore.StringValue(msg.Content)), &textBody); err != nil { + logger.Warn("飞书消息 Content 解析失败", zap.Error(err)) + return + } + text := strings.TrimSpace(textBody.Text) + if text == "" { + return + } + userID := resolveLarkUserID(event, cfg.AllowChatIDFallback && !strictUserIdentity) + if userID == "" { + logger.Warn("飞书消息缺少可用用户标识，已忽略") + return + } + messageID := larkcore.StringValue(msg.MessageId) + reply := h.HandleMessage("lark", userID, text) + contentBytes, _ := json.Marshal(larkTextContent{Text: reply}) + _, err := client.Im.Message.Reply(ctx, larkim.NewReplyMessageReqBuilder(). + MessageId(messageID). + Body(larkim.NewReplyMessageReqBodyBuilder(). + MsgType(larkim.MsgTypeText). + Content(string(contentBytes)). + Build()). + Build()) + if err != nil { + logger.Warn("飞书回复失败", zap.String("message_id", messageID), zap.Error(err)) + return + } + logger.Debug("飞书已回复", zap.String("message_id", messageID)) +} + +// resolveLarkUserID 提取飞书会话隔离键： +// tenant_key + 稳定用户标识（user_id/open_id/union_id）；按配置可选 chat_id 兜底。 +func resolveLarkUserID(event *larkim.P2MessageReceiveV1, allowChatIDFallback bool) string { + if event == nil || event.Event == nil || event.Event.Sender == nil || event.Event.Sender.SenderId == nil { + return "" + } + tenantKey := strings.TrimSpace(larkcore.StringValue(event.Event.Sender.TenantKey)) + if tenantKey == "" { + tenantKey = "default" + } + prefix := "t:" + tenantKey + "|" + if id := strings.TrimSpace(larkcore.StringValue(event.Event.Sender.SenderId.UserId)); id != "" { + return prefix + "u:" + id + } + if id := strings.TrimSpace(larkcore.StringValue(event.Event.Sender.SenderId.OpenId)); id != "" { + return prefix + "o:" + id + } + if id := strings.TrimSpace(larkcore.StringValue(event.Event.Sender.SenderId.UnionId)); id != "" { + return prefix + "n:" + id + } + if allowChatIDFallback && event.Event.Message != nil { + if id := strings.TrimSpace(larkcore.StringValue(event.Event.Message.ChatId)); id != "" { + return prefix + "c:" + id + } + } + return "" +} diff --git a/internal/robot/proactive.go b/internal/robot/proactive.go new file mode 100644 index 00000000..d600df8b --- /dev/null +++ b/internal/robot/proactive.go @@ -0,0 +1,192 @@ +package robot + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/bwmarrin/discordgo" + lark "github.com/larksuite/oapi-sdk-go/v3" + larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1" + "github.com/slack-go/slack" +) + +// SendProactive sends a message without an inbound event. Platforms whose +// reply credentials are event-scoped deliberately return an error instead of +// pretending delivery succeeded. +func SendProactive(ctx context.Context, cfg config.RobotsConfig, platform, externalUserID, message string) error { + platform = strings.ToLower(strings.TrimSpace(platform)) + userID := robotIdentityUserPart(externalUserID) + if userID == "" { + return fmt.Errorf("invalid robot recipient") + } + switch platform { + case "telegram": + if !cfg.Telegram.Enabled || strings.TrimSpace(cfg.Telegram.BotToken) == "" { + return fmt.Errorf("telegram is not configured") + } + id, err := strconv.ParseInt(userID, 10, 64) + if err != nil { + return fmt.Errorf("invalid telegram user id: %w", err) + } + return telegramSendReply(ctx, nilSafeHTTPClient(), strings.TrimSpace(cfg.Telegram.BotToken), id, message) + case "slack": + if !cfg.Slack.Enabled || strings.TrimSpace(cfg.Slack.BotToken) == "" { + return fmt.Errorf("slack is not configured") + } + api := slack.New(strings.TrimSpace(cfg.Slack.BotToken)) + channel, _, _, err := api.OpenConversationContext(ctx, &slack.OpenConversationParameters{Users: []string{userID}}) + if err != nil { + return err + } + for _, chunk := range splitTextChunks(message, slackMaxMessageRunes) { + if _, _, err = api.PostMessageContext(ctx, channel.ID, slack.MsgOptionText(chunk, false)); err != nil { + return err + } + } + return nil + case "discord": + if !cfg.Discord.Enabled || strings.TrimSpace(cfg.Discord.BotToken) == "" { + return fmt.Errorf("discord is not configured") + } + token := strings.TrimSpace(cfg.Discord.BotToken) + if !strings.HasPrefix(token, "Bot ") { + token = "Bot " + token + } + session, err := discordgo.New(token) + if err != nil { + return err + } + channel, err := session.UserChannelCreate(userID) + if err != nil { + return err + } + for _, chunk := range splitTextChunks(message, discordMaxMessageRunes) { + if _, err = session.ChannelMessageSend(channel.ID, chunk); err != nil { + return err + } + } + return nil + case "wecom": + return sendWecomProactive(ctx, cfg.Wecom, userID, message) + case "lark": + return sendLarkProactive(ctx, cfg.Lark, externalUserID, message) + default: + return fmt.Errorf("platform %s does not support proactive alerts yet", platform) + } +} + +func SupportsProactive(platform string) bool { + switch strings.ToLower(strings.TrimSpace(platform)) { + case "telegram", "slack", "discord", "wecom", "lark": + return true + default: + return false + } +} + +func nilSafeHTTPClient() *http.Client { return &http.Client{Timeout: 15 * time.Second} } + +func sendWecomProactive(ctx context.Context, cfg config.RobotWecomConfig, userID, message string) error { + if !cfg.Enabled || strings.TrimSpace(cfg.CorpID) == "" || strings.TrimSpace(cfg.Secret) == "" || cfg.AgentID == 0 { + return fmt.Errorf("wecom proactive API is not configured") + } + tokenURL := "https://qyapi.weixin.qq.com/cgi-bin/gettoken?corpid=" + cfg.CorpID + "&corpsecret=" + cfg.Secret + req, err := http.NewRequestWithContext(ctx, http.MethodGet, tokenURL, nil) + if err != nil { + return err + } + resp, err := nilSafeHTTPClient().Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return err + } + var tokenResp struct { + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` + AccessToken string `json:"access_token"` + } + if err := json.Unmarshal(body, &tokenResp); err != nil { + return err + } + if tokenResp.ErrCode != 0 || tokenResp.AccessToken == "" { + return fmt.Errorf("wecom token: %s", tokenResp.ErrMsg) + } + payload, _ := json.Marshal(map[string]interface{}{ + "touser": userID, "msgtype": "text", "agentid": cfg.AgentID, + "text": map[string]string{"content": message}, "safe": 0, + }) + sendReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token="+tokenResp.AccessToken, bytes.NewReader(payload)) + if err != nil { + return err + } + sendReq.Header.Set("Content-Type", "application/json") + sendResp, err := nilSafeHTTPClient().Do(sendReq) + if err != nil { + return err + } + defer sendResp.Body.Close() + result, _ := io.ReadAll(sendResp.Body) + var parsed struct { + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` + } + if err := json.Unmarshal(result, &parsed); err != nil { + return err + } + if parsed.ErrCode != 0 { + return fmt.Errorf("wecom send: %s", parsed.ErrMsg) + } + return nil +} + +func sendLarkProactive(ctx context.Context, cfg config.RobotLarkConfig, identity, message string) error { + if !cfg.Enabled || strings.TrimSpace(cfg.AppID) == "" || strings.TrimSpace(cfg.AppSecret) == "" { + return fmt.Errorf("lark is not configured") + } + receiveIDType, receiveID := "user_id", robotIdentityUserPart(identity) + if idx := strings.LastIndex(identity, "|o:"); idx >= 0 { + receiveIDType, receiveID = "open_id", strings.TrimSpace(identity[idx+3:]) + } + if idx := strings.LastIndex(identity, "|n:"); idx >= 0 { + receiveIDType, receiveID = "union_id", strings.TrimSpace(identity[idx+3:]) + } + if receiveID == "" { + return fmt.Errorf("invalid lark recipient") + } + content, _ := json.Marshal(larkTextContent{Text: message}) + client := lark.NewClient(cfg.AppID, cfg.AppSecret) + resp, err := client.Im.Message.Create(ctx, larkim.NewCreateMessageReqBuilder(). + ReceiveIdType(receiveIDType). + Body(larkim.NewCreateMessageReqBodyBuilder().ReceiveId(receiveID).MsgType(larkim.MsgTypeText).Content(string(content)).Build()).Build()) + if err != nil { + return err + } + if resp == nil || !resp.Success() { + return fmt.Errorf("lark send failed") + } + return nil +} + +func robotIdentityUserPart(identity string) string { + identity = strings.TrimSpace(identity) + if i := strings.LastIndex(identity, "|u:"); i >= 0 { + return strings.TrimSpace(identity[i+3:]) + } + if strings.HasPrefix(identity, "u:") { + return strings.TrimSpace(identity[2:]) + } + return identity +} diff --git a/internal/robot/qq.go b/internal/robot/qq.go new file mode 100644 index 00000000..489d8909 --- /dev/null +++ b/internal/robot/qq.go @@ -0,0 +1,209 @@ +package robot + +import ( + "context" + "strings" + "sync" + + "cyberstrike-ai/internal/config" + + "github.com/tencent-connect/botgo" + "github.com/tencent-connect/botgo/dto" + "github.com/tencent-connect/botgo/event" + "github.com/tencent-connect/botgo/openapi" + "github.com/tencent-connect/botgo/token" + "go.uber.org/zap" +) + +const ( + qqPlatform = "qq" + qqMaxMessageRunes = 3500 +) + +var ( + qqHandlerMu sync.Mutex + qqHandler MessageHandler + qqLogger *zap.Logger + qqAPI openapi.OpenAPI +) + +// StartQQ 启动 QQ 机器人 WebSocket（C2C 与群 @，出站连接，无需公网回调）。 +func StartQQ(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.QQ + if !cfg.Enabled || strings.TrimSpace(cfg.AppID) == "" || strings.TrimSpace(cfg.ClientSecret) == "" { + return + } + go runQQLoop(ctx, cfg, h, logger) +} + +func runQQLoop(ctx context.Context, cfg config.RobotQQConfig, h MessageHandler, logger *zap.Logger) { + backoff := reconnectInitial + for { + if ctx.Err() != nil { + logger.Info("QQ 机器人 WebSocket 已按配置关闭") + return + } + err := runQQSession(ctx, cfg, h, logger) + if ctx.Err() != nil { + return + } + if err != nil { + logger.Warn("QQ 机器人 WebSocket 异常，将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + if !waitReconnect(ctx, &backoff) { + return + } + } +} + +func runQQSession(ctx context.Context, cfg config.RobotQQConfig, h MessageHandler, logger *zap.Logger) error { + appID := strings.TrimSpace(cfg.AppID) + secret := strings.TrimSpace(cfg.ClientSecret) + credentials := &token.QQBotCredentials{AppID: appID, AppSecret: secret} + tokenSource := token.NewQQBotTokenSource(credentials) + + if err := token.StartRefreshAccessToken(ctx, tokenSource); err != nil { + return err + } + + var api openapi.OpenAPI + if cfg.Sandbox { + api = botgo.NewSandboxOpenAPI(appID, tokenSource) + } else { + api = botgo.NewOpenAPI(appID, tokenSource) + } + + qqHandlerMu.Lock() + qqHandler = h + qqLogger = logger + qqAPI = api + qqHandlerMu.Unlock() + defer func() { + qqHandlerMu.Lock() + qqHandler = nil + qqLogger = nil + qqAPI = nil + qqHandlerMu.Unlock() + }() + + intents := event.RegisterHandlers( + event.C2CMessageEventHandler(handleQQC2CMessage), + event.GroupATMessageEventHandler(handleQQGroupATMessage), + ) + + wsInfo, err := api.WS(ctx, nil, "") + if err != nil { + return err + } + logger.Info("QQ 机器人 WebSocket 正在连接…", zap.String("app_id", appID), zap.Bool("sandbox", cfg.Sandbox)) + + done := make(chan error, 1) + go func() { + done <- botgo.NewSessionManager().Start(wsInfo, tokenSource, &intents) + }() + + select { + case <-ctx.Done(): + return ctx.Err() + case err := <-done: + if err != nil { + return err + } + return nil + } +} + +func handleQQC2CMessage(payload *dto.WSPayload, data *dto.WSC2CMessageData) error { + if data == nil || data.Author == nil { + return nil + } + text := strings.TrimSpace(data.Content) + if text == "" { + return nil + } + userOpenID := strings.TrimSpace(data.Author.ID) + if userOpenID == "" { + return nil + } + userID := "u:" + userOpenID + qqHandlerMu.Lock() + h := qqHandler + logger := qqLogger + api := qqAPI + qqHandlerMu.Unlock() + if h == nil || api == nil { + return nil + } + logger.Info("QQ 收到 C2C 消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(qqPlatform, userID, text) + return qqPostC2CReply(context.Background(), api, userOpenID, payload, data.ID, reply, logger) +} + +func handleQQGroupATMessage(payload *dto.WSPayload, data *dto.WSGroupATMessageData) error { + if data == nil || data.Author == nil { + return nil + } + text := strings.TrimSpace(data.Content) + if text == "" { + return nil + } + userOpenID := strings.TrimSpace(data.Author.ID) + groupID := strings.TrimSpace(data.GroupID) + if userOpenID == "" { + return nil + } + userID := "g:" + groupID + "|u:" + userOpenID + qqHandlerMu.Lock() + h := qqHandler + logger := qqLogger + api := qqAPI + qqHandlerMu.Unlock() + if h == nil || api == nil { + return nil + } + logger.Info("QQ 收到群 @ 消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(qqPlatform, userID, text) + return qqPostGroupReply(context.Background(), api, groupID, payload, data.ID, reply, logger) +} + +func qqPostC2CReply(ctx context.Context, api openapi.OpenAPI, userOpenID string, payload *dto.WSPayload, msgID, reply string, logger *zap.Logger) error { + reply = trimReply(reply) + if reply == "" { + return nil + } + for _, chunk := range splitTextChunks(reply, qqMaxMessageRunes) { + msg := &dto.MessageToCreate{ + Content: chunk, + MsgID: msgID, + } + if payload != nil && payload.EventID != "" { + msg.EventID = payload.EventID + } + if _, err := api.PostC2CMessage(ctx, userOpenID, msg); err != nil { + logger.Warn("QQ 发送 C2C 回复失败", zap.String("to", userOpenID), zap.Error(err)) + return err + } + } + return nil +} + +func qqPostGroupReply(ctx context.Context, api openapi.OpenAPI, groupID string, payload *dto.WSPayload, msgID, reply string, logger *zap.Logger) error { + reply = trimReply(reply) + if reply == "" { + return nil + } + for _, chunk := range splitTextChunks(reply, qqMaxMessageRunes) { + msg := &dto.MessageToCreate{ + Content: chunk, + MsgID: msgID, + } + if payload != nil && payload.EventID != "" { + msg.EventID = payload.EventID + } + if _, err := api.PostGroupMessage(ctx, groupID, msg); err != nil { + logger.Warn("QQ 发送群消息回复失败", zap.String("group", groupID), zap.Error(err)) + return err + } + } + return nil +} diff --git a/internal/robot/reconnect.go b/internal/robot/reconnect.go new file mode 100644 index 00000000..4aee6c90 --- /dev/null +++ b/internal/robot/reconnect.go @@ -0,0 +1,38 @@ +package robot + +import ( + "context" + "time" +) + +const ( + reconnectInitial = 5 * time.Second + reconnectMax = 60 * time.Second +) + +func waitReconnect(ctx context.Context, backoff *time.Duration) bool { + if ctx.Err() != nil { + return false + } + select { + case <-ctx.Done(): + return false + case <-time.After(*backoff): + if *backoff < reconnectMax { + *backoff *= 2 + if *backoff > reconnectMax { + *backoff = reconnectMax + } + } + return true + } +} + +func bumpBackoff(backoff *time.Duration) { + if *backoff < reconnectMax { + *backoff *= 2 + if *backoff > reconnectMax { + *backoff = reconnectMax + } + } +} diff --git a/internal/robot/slack.go b/internal/robot/slack.go new file mode 100644 index 00000000..1e30a26a --- /dev/null +++ b/internal/robot/slack.go @@ -0,0 +1,135 @@ +package robot + +import ( + "context" + "strings" + + "cyberstrike-ai/internal/config" + + "github.com/slack-go/slack" + "github.com/slack-go/slack/slackevents" + "github.com/slack-go/slack/socketmode" + "go.uber.org/zap" +) + +const ( + slackPlatform = "slack" + slackMaxMessageRunes = 3900 +) + +// StartSlack 启动 Slack Socket Mode（出站 WebSocket，无需公网回调）。 +func StartSlack(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.Slack + if !cfg.Enabled || strings.TrimSpace(cfg.BotToken) == "" || strings.TrimSpace(cfg.AppToken) == "" { + return + } + go runSlackLoop(ctx, cfg, h, logger) +} + +func runSlackLoop(ctx context.Context, cfg config.RobotSlackConfig, h MessageHandler, logger *zap.Logger) { + backoff := reconnectInitial + for { + err := runSlackSocket(ctx, cfg, h, logger) + if ctx.Err() != nil { + logger.Info("Slack Socket Mode 已按配置关闭") + return + } + if err != nil { + logger.Warn("Slack Socket Mode 异常，将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + if !waitReconnect(ctx, &backoff) { + return + } + } +} + +func runSlackSocket(ctx context.Context, cfg config.RobotSlackConfig, h MessageHandler, logger *zap.Logger) error { + api := slack.New( + strings.TrimSpace(cfg.BotToken), + slack.OptionAppLevelToken(strings.TrimSpace(cfg.AppToken)), + ) + client := socketmode.New(api) + logger.Info("Slack Socket Mode 正在连接…") + + go func() { + for evt := range client.Events { + switch evt.Type { + case socketmode.EventTypeEventsAPI: + eventsAPIEvent, ok := evt.Data.(slackevents.EventsAPIEvent) + if !ok { + continue + } + client.Ack(*evt.Request) + if eventsAPIEvent.Type != slackevents.CallbackEvent { + continue + } + switch ev := eventsAPIEvent.InnerEvent.Data.(type) { + case *slackevents.MessageEvent: + handleSlackMessage(ctx, api, eventsAPIEvent.TeamID, ev, h, logger) + case *slackevents.AppMentionEvent: + handleSlackAppMention(ctx, api, eventsAPIEvent.TeamID, ev, h, logger) + } + case socketmode.EventTypeConnecting: + logger.Info("Slack Socket Mode 正在连接…") + case socketmode.EventTypeConnected: + logger.Info("Slack Socket Mode 已连接，等待收消息") + } + } + }() + + return client.RunContext(ctx) +} + +func handleSlackMessage(ctx context.Context, api *slack.Client, teamID string, ev *slackevents.MessageEvent, h MessageHandler, logger *zap.Logger) { + if ev == nil || ev.BotID != "" || ev.SubType != "" { + return + } + if ev.ChannelType != "im" { + return + } + text := strings.TrimSpace(ev.Text) + if text == "" { + return + } + userID := slackSessionKey(teamID, ev.User) + logger.Info("Slack 收到消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(slackPlatform, userID, text) + slackPostReply(ctx, api, ev.Channel, reply, logger) +} + +func handleSlackAppMention(ctx context.Context, api *slack.Client, teamID string, ev *slackevents.AppMentionEvent, h MessageHandler, logger *zap.Logger) { + if ev == nil || ev.BotID != "" { + return + } + text := strings.TrimSpace(ev.Text) + if text == "" { + return + } + userID := slackSessionKey(teamID, ev.User) + logger.Info("Slack 收到 @ 消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(slackPlatform, userID, text) + slackPostReply(ctx, api, ev.Channel, reply, logger) +} + +func slackSessionKey(teamID, userID string) string { + teamID = strings.TrimSpace(teamID) + userID = strings.TrimSpace(userID) + if teamID == "" { + teamID = "default" + } + return "t:" + teamID + "|u:" + userID +} + +func slackPostReply(ctx context.Context, api *slack.Client, channel, reply string, logger *zap.Logger) { + reply = trimReply(reply) + if reply == "" { + return + } + for _, chunk := range splitTextChunks(reply, slackMaxMessageRunes) { + _, _, err := api.PostMessageContext(ctx, channel, slack.MsgOptionText(chunk, false)) + if err != nil { + logger.Warn("Slack 发送回复失败", zap.String("channel", channel), zap.Error(err)) + return + } + } +} diff --git a/internal/robot/split.go b/internal/robot/split.go new file mode 100644 index 00000000..c5425b84 --- /dev/null +++ b/internal/robot/split.go @@ -0,0 +1,29 @@ +package robot + +import "strings" + +// splitTextChunks splits text into chunks no longer than maxRunes (rune count). +func splitTextChunks(text string, maxRunes int) []string { + text = strings.TrimSpace(text) + if text == "" || maxRunes <= 0 { + return nil + } + runes := []rune(text) + if len(runes) <= maxRunes { + return []string{text} + } + var out []string + for len(runes) > 0 { + end := maxRunes + if end > len(runes) { + end = len(runes) + } + out = append(out, string(runes[:end])) + runes = runes[end:] + } + return out +} + +func trimReply(s string) string { + return strings.TrimSpace(s) +} diff --git a/internal/robot/telegram.go b/internal/robot/telegram.go new file mode 100644 index 00000000..f8712d57 --- /dev/null +++ b/internal/robot/telegram.go @@ -0,0 +1,262 @@ +package robot + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +const ( + telegramPlatform = "telegram" + telegramAPIBase = "https://api.telegram.org" + telegramLongPollSec = 30 + telegramMaxMessageRunes = 4096 +) + +type telegramUpdate struct { + UpdateID int `json:"update_id"` + Message *telegramMessage `json:"message"` +} + +type telegramMessage struct { + MessageID int64 `json:"message_id"` + Chat telegramChat `json:"chat"` + From *telegramUser `json:"from"` + Text string `json:"text"` + Entities []telegramEntity `json:"entities"` +} + +type telegramChat struct { + ID int64 `json:"id"` + Type string `json:"type"` + Title string `json:"title"` +} + +type telegramUser struct { + ID int64 `json:"id"` + Username string `json:"username"` + IsBot bool `json:"is_bot"` +} + +type telegramEntity struct { + Type string `json:"type"` + Offset int `json:"offset"` + Length int `json:"length"` +} + +type telegramGetUpdatesResp struct { + OK bool `json:"ok"` + Result []telegramUpdate `json:"result"` + Description string `json:"description"` +} + +type telegramBotMe struct { + OK bool `json:"ok"` + Result struct { + ID int64 `json:"id"` + Username string `json:"username"` + } `json:"result"` +} + +// StartTelegram 启动 Telegram Bot 长轮询（getUpdates，无需公网回调）。 +func StartTelegram(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, logger *zap.Logger) { + cfg := robotsCfg.Telegram + if !cfg.Enabled || strings.TrimSpace(cfg.BotToken) == "" { + return + } + go runTelegramLoop(ctx, cfg, h, logger) +} + +func runTelegramLoop(ctx context.Context, cfg config.RobotTelegramConfig, h MessageHandler, logger *zap.Logger) { + backoff := reconnectInitial + for { + err := runTelegramPoll(ctx, cfg, h, logger) + if ctx.Err() != nil { + logger.Info("Telegram 长轮询已按配置关闭") + return + } + if err != nil { + logger.Warn("Telegram 长轮询异常，将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + if !waitReconnect(ctx, &backoff) { + return + } + } +} + +func runTelegramPoll(ctx context.Context, cfg config.RobotTelegramConfig, h MessageHandler, logger *zap.Logger) error { + token := strings.TrimSpace(cfg.BotToken) + botUsername := strings.TrimSpace(cfg.BotUsername) + if botUsername == "" { + if name, err := telegramGetMe(ctx, token); err != nil { + logger.Warn("Telegram getMe 失败", zap.Error(err)) + } else { + botUsername = name + } + } + offset := cfg.UpdateOffset + logger.Info("Telegram 长轮询已启动", zap.String("bot", botUsername)) + client := &http.Client{Timeout: telegramLongPollSec*time.Second + 10*time.Second} + + for { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + + updates, err := telegramGetUpdates(ctx, client, token, offset) + if err != nil { + return err + } + for _, u := range updates { + next := int64(u.UpdateID) + 1 + if next > offset { + offset = next + } + if u.Message == nil || u.Message.From == nil || u.Message.From.IsBot { + continue + } + text := strings.TrimSpace(u.Message.Text) + if text == "" { + continue + } + chatType := strings.ToLower(strings.TrimSpace(u.Message.Chat.Type)) + if chatType != "private" { + if !cfg.AllowGroupMessages { + continue + } + if botUsername != "" && !telegramMentionsBot(text, u.Message.Entities, botUsername) { + continue + } + } + userID := telegramSessionKey(chatType, u.Message.Chat.ID, u.Message.From.ID) + logger.Info("Telegram 收到消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(telegramPlatform, userID, text) + if err := telegramSendReply(ctx, client, token, u.Message.Chat.ID, reply); err != nil { + logger.Warn("Telegram 发送回复失败", zap.String("to", userID), zap.Error(err)) + } + } + } +} + +func telegramSessionKey(chatType string, chatID, fromUserID int64) string { + if chatType == "private" { + return fmt.Sprintf("u:%d", fromUserID) + } + return fmt.Sprintf("g:%d|u:%d", chatID, fromUserID) +} + +func telegramMentionsBot(text string, entities []telegramEntity, botUsername string) bool { + needle := "@" + strings.TrimPrefix(strings.ToLower(botUsername), "@") + lower := strings.ToLower(text) + if strings.Contains(lower, needle) { + return true + } + for _, e := range entities { + if e.Type != "mention" { + continue + } + if e.Offset < 0 || e.Length <= 0 || e.Offset+e.Length > len(text) { + continue + } + mention := strings.ToLower(text[e.Offset : e.Offset+e.Length]) + if mention == needle { + return true + } + } + return false +} + +func telegramAPIURL(token, method string) string { + return fmt.Sprintf("%s/bot%s/%s", telegramAPIBase, token, method) +} + +func telegramGetMe(ctx context.Context, token string) (string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, telegramAPIURL(token, "getMe"), nil) + if err != nil { + return "", err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + var parsed telegramBotMe + if err := json.Unmarshal(body, &parsed); err != nil { + return "", err + } + if !parsed.OK { + return "", fmt.Errorf("getMe failed: %s", string(body)) + } + return parsed.Result.Username, nil +} + +func telegramGetUpdates(ctx context.Context, client *http.Client, token string, offset int64) ([]telegramUpdate, error) { + url := fmt.Sprintf("%s?timeout=%d&allowed_updates=%s", telegramAPIURL(token, "getUpdates"), telegramLongPollSec, `["message"]`) + if offset > 0 { + url += fmt.Sprintf("&offset=%d", offset) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + var parsed telegramGetUpdatesResp + if err := json.Unmarshal(body, &parsed); err != nil { + return nil, err + } + if !parsed.OK { + if parsed.Description != "" { + return nil, fmt.Errorf("getUpdates: %s", parsed.Description) + } + return nil, fmt.Errorf("getUpdates failed: %s", string(body)) + } + return parsed.Result, nil +} + +func telegramSendReply(ctx context.Context, client *http.Client, token string, chatID int64, reply string) error { + reply = trimReply(reply) + if reply == "" { + return nil + } + for _, chunk := range splitTextChunks(reply, telegramMaxMessageRunes) { + payload := map[string]interface{}{ + "chat_id": chatID, + "text": chunk, + } + body, _ := json.Marshal(payload) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, telegramAPIURL(token, "sendMessage"), bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + return err + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("sendMessage status %d", resp.StatusCode) + } + } + return nil +} diff --git a/internal/robot/wechat.go b/internal/robot/wechat.go new file mode 100644 index 00000000..17d50404 --- /dev/null +++ b/internal/robot/wechat.go @@ -0,0 +1,96 @@ +package robot + +import ( + "context" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/robot/ilink" + + "go.uber.org/zap" +) + +const ( + wechatReconnectInitial = 5 * time.Second + wechatReconnectMax = 60 * time.Second + wechatPlatform = "wechat" +) + +// StartWechat 启动微信 iLink 长轮询（无需公网回调），收到消息后调用 handler 并回复。 +func StartWechat(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, appVersion string, logger *zap.Logger) { + cfg := robotsCfg.Wechat + if !cfg.Enabled || cfg.BotToken == "" { + return + } + go runWechatLoop(ctx, cfg, h, appVersion, logger) +} + +func runWechatLoop(ctx context.Context, cfg config.RobotWechatConfig, h MessageHandler, appVersion string, logger *zap.Logger) { + backoff := wechatReconnectInitial + for { + err := runWechatPoll(ctx, cfg, h, appVersion, logger) + if ctx.Err() != nil { + logger.Info("微信 iLink 长轮询已按配置关闭") + return + } + if err != nil { + logger.Warn("微信 iLink 长轮询异常，将自动重连", zap.Error(err), zap.Duration("retry_after", backoff)) + } + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + if backoff < wechatReconnectMax { + backoff *= 2 + if backoff > wechatReconnectMax { + backoff = wechatReconnectMax + } + } + } + } +} + +func runWechatPoll(ctx context.Context, cfg config.RobotWechatConfig, h MessageHandler, appVersion string, logger *zap.Logger) error { + client := ilink.NewClient(cfg.BaseURL, cfg.BotToken, cfg.BotAgent, ilink.BuildClientVersion(appVersion)) + buf := cfg.GetUpdatesBuf + logger.Info("微信 iLink 长轮询已启动", zap.String("ilink_bot_id", cfg.ILinkBotID)) + for { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + resp, err := client.GetUpdates(ctx, buf) + if err != nil { + return err + } + if resp.ErrCode != 0 && resp.Ret != 0 { + logger.Warn("微信 getUpdates 返回错误", zap.Int("errcode", resp.ErrCode), zap.String("errmsg", resp.ErrMsg)) + } + if resp.GetUpdatesBuf != "" { + buf = resp.GetUpdatesBuf + } + for _, msg := range resp.Msgs { + if msg.MessageType != 1 { + continue + } + text := ilink.ExtractText(msg) + if text == "" { + continue + } + userID := strings.TrimSpace(msg.FromUserID) + if userID == "" { + continue + } + logger.Info("微信收到消息", zap.String("from", userID), zap.String("content", text)) + reply := h.HandleMessage(wechatPlatform, userID, text) + if strings.TrimSpace(reply) == "" { + continue + } + if err := client.SendTextMessage(ctx, userID, msg.ContextToken, reply, ""); err != nil { + logger.Warn("微信发送回复失败", zap.String("to", userID), zap.Error(err)) + } + } + } +} 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 +}