From 7270e3c3d1f7d29e7876ba2a7c531efb5ee3d6ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 28 Feb 2026 00:52:17 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 64 +++++++++++++++++++++++++++++++++----- internal/handler/config.go | 19 +++++++++++ internal/robot/ding.go | 11 ++++--- internal/robot/lark.go | 11 ++++--- 4 files changed, 90 insertions(+), 15 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index 90c2fb14..bbf0a5d5 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -7,6 +7,7 @@ import ( "net/http" "os" "path/filepath" + "sync" "time" "cyberstrike-ai/internal/agent" @@ -44,6 +45,10 @@ type App struct { knowledgeIndexer *knowledge.Indexer // 知识库索引器（用于动态初始化） knowledgeHandler *handler.KnowledgeHandler // 知识库处理器（用于动态初始化） agentHandler *handler.AgentHandler // Agent处理器（用于更新知识库管理器） + robotHandler *handler.RobotHandler // 机器人处理器（钉钉/飞书/企业微信） + robotMu sync.Mutex // 保护钉钉/飞书长连接的 cancel + dingCancel context.CancelFunc // 钉钉 Stream 取消函数，用于配置变更时重启 + larkCancel context.CancelFunc // 飞书长连接取消函数，用于配置变更时重启 } // New 创建新应用 @@ -327,13 +332,6 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { // 创建OpenAPI处理器 conversationHandler := handler.NewConversationHandler(db, log.Logger) robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger) - // 飞书/钉钉长连接（无需公网），启用时在后台启动 - if cfg.Robots.Lark.Enabled && cfg.Robots.Lark.AppID != "" && cfg.Robots.Lark.AppSecret != "" { - go robot.StartLark(cfg.Robots.Lark, robotHandler, log.Logger) - } - if cfg.Robots.Dingtalk.Enabled && cfg.Robots.Dingtalk.ClientID != "" && cfg.Robots.Dingtalk.ClientSecret != "" { - go robot.StartDing(cfg.Robots.Dingtalk, robotHandler, log.Logger) - } openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, resultStorage, conversationHandler, agentHandler) // 创建 App 实例（部分字段稍后填充） @@ -353,7 +351,10 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { knowledgeIndexer: knowledgeIndexer, knowledgeHandler: knowledgeHandler, agentHandler: agentHandler, + robotHandler: robotHandler, } + // 飞书/钉钉长连接（无需公网），启用时在后台启动；后续前端应用配置时会通过 RestartRobotConnections 重启 + app.startRobotConnections() // 设置漏洞工具注册器（内置工具，必须设置） vulnerabilityRegistrar := func() error { @@ -410,6 +411,9 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { configHandler.SetRetrieverUpdater(knowledgeRetriever) } + // 设置机器人连接重启器，前端应用配置后无需重启服务即可使钉钉/飞书新配置生效 + configHandler.SetRobotRestarter(app) + // 设置路由（使用 App 实例以便动态获取 handler） setupRoutes( router, @@ -462,6 +466,18 @@ func (a *App) Run() error { // Shutdown 关闭应用 func (a *App) Shutdown() { + // 停止钉钉/飞书长连接 + a.robotMu.Lock() + if a.dingCancel != nil { + a.dingCancel() + a.dingCancel = nil + } + if a.larkCancel != nil { + a.larkCancel() + a.larkCancel = nil + } + a.robotMu.Unlock() + // 停止所有外部MCP客户端 if a.externalMCPMgr != nil { a.externalMCPMgr.StopAll() @@ -475,6 +491,40 @@ func (a *App) Shutdown() { } } +// startRobotConnections 根据当前配置启动钉钉/飞书长连接（不先关闭已有连接，仅用于首次启动） +func (a *App) startRobotConnections() { + a.robotMu.Lock() + defer a.robotMu.Unlock() + cfg := a.config + if cfg.Robots.Lark.Enabled && cfg.Robots.Lark.AppID != "" && cfg.Robots.Lark.AppSecret != "" { + ctx, cancel := context.WithCancel(context.Background()) + a.larkCancel = cancel + go robot.StartLark(ctx, cfg.Robots.Lark, a.robotHandler, a.logger.Logger) + } + if cfg.Robots.Dingtalk.Enabled && cfg.Robots.Dingtalk.ClientID != "" && cfg.Robots.Dingtalk.ClientSecret != "" { + ctx, cancel := context.WithCancel(context.Background()) + a.dingCancel = cancel + go robot.StartDing(ctx, cfg.Robots.Dingtalk, a.robotHandler, a.logger.Logger) + } +} + +// RestartRobotConnections 重启钉钉/飞书长连接，使前端应用配置后立即生效（实现 handler.RobotRestarter） +func (a *App) RestartRobotConnections() { + a.robotMu.Lock() + if a.dingCancel != nil { + a.dingCancel() + a.dingCancel = nil + } + if a.larkCancel != nil { + a.larkCancel() + a.larkCancel = nil + } + a.robotMu.Unlock() + // 给旧 goroutine 一点时间退出 + time.Sleep(200 * time.Millisecond) + a.startRobotConnections() +} + // setupRoutes 设置路由 func setupRoutes( router *gin.Engine, diff --git a/internal/handler/config.go b/internal/handler/config.go index 55210567..98d56f5f 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -44,6 +44,11 @@ type AppUpdater interface { UpdateKnowledgeComponents(handler *KnowledgeHandler, manager interface{}, retriever interface{}, indexer interface{}) } +// RobotRestarter 机器人连接重启器（用于配置应用后重启钉钉/飞书长连接） +type RobotRestarter interface { + RestartRobotConnections() +} + // ConfigHandler 配置处理器 type ConfigHandler struct { configPath string @@ -59,6 +64,7 @@ type ConfigHandler struct { retrieverUpdater RetrieverUpdater // 检索器更新器（可选） knowledgeInitializer KnowledgeInitializer // 知识库初始化器（可选） appUpdater AppUpdater // App更新器（可选） + robotRestarter RobotRestarter // 机器人连接重启器（可选），ApplyConfig 时重启钉钉/飞书 logger *zap.Logger mu sync.RWMutex lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置（用于检测变更） @@ -142,6 +148,13 @@ func (h *ConfigHandler) SetAppUpdater(updater AppUpdater) { h.appUpdater = updater } +// SetRobotRestarter 设置机器人连接重启器（ApplyConfig 时用于重启钉钉/飞书长连接） +func (h *ConfigHandler) SetRobotRestarter(restarter RobotRestarter) { + h.mu.Lock() + defer h.mu.Unlock() + h.robotRestarter = restarter +} + // GetConfigResponse 获取配置响应 type GetConfigResponse struct { OpenAI config.OpenAIConfig `json:"openai"` @@ -837,6 +850,12 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) { } } + // 重启钉钉/飞书长连接，使前端修改的机器人配置立即生效（无需重启服务） + if h.robotRestarter != nil { + h.robotRestarter.RestartRobotConnections() + h.logger.Info("已触发机器人连接重启（钉钉/飞书）") + } + h.logger.Info("配置已应用", zap.Int("tools_count", len(h.config.Security.Tools)), ) diff --git a/internal/robot/ding.go b/internal/robot/ding.go index 1fe8e88c..29adca8c 100644 --- a/internal/robot/ding.go +++ b/internal/robot/ding.go @@ -15,8 +15,9 @@ import ( "go.uber.org/zap" ) -// StartDing 启动钉钉 Stream 长连接（无需公网），收到消息后调用 handler 并通过 SessionWebhook 回复 -func StartDing(cfg config.RobotDingtalkConfig, h MessageHandler, logger *zap.Logger) { +// StartDing 启动钉钉 Stream 长连接（无需公网），收到消息后调用 handler 并通过 SessionWebhook 回复。 +// ctx 被取消时长连接会退出，便于配置变更时重启。 +func StartDing(ctx context.Context, cfg config.RobotDingtalkConfig, h MessageHandler, logger *zap.Logger) { if !cfg.Enabled || cfg.ClientID == "" || cfg.ClientSecret == "" { return } @@ -30,9 +31,11 @@ func StartDing(cfg config.RobotDingtalkConfig, h MessageHandler, logger *zap.Log ) logger.Info("钉钉 Stream 正在连接…", zap.String("client_id", cfg.ClientID)) go func() { - err := streamClient.Start(context.Background()) - if err != nil { + err := streamClient.Start(ctx) + if err != nil && ctx.Err() == nil { logger.Error("钉钉 Stream 长连接退出", zap.Error(err)) + } else if ctx.Err() != nil { + logger.Info("钉钉 Stream 已按配置重启关闭") } }() logger.Info("钉钉 Stream 已启动（无需公网），等待收消息", zap.String("client_id", cfg.ClientID)) diff --git a/internal/robot/lark.go b/internal/robot/lark.go index 30c825fb..a57bc7b6 100644 --- a/internal/robot/lark.go +++ b/internal/robot/lark.go @@ -19,8 +19,9 @@ type larkTextContent struct { Text string `json:"text"` } -// StartLark 启动飞书长连接（无需公网），收到消息后调用 handler 并回复 -func StartLark(cfg config.RobotLarkConfig, h MessageHandler, logger *zap.Logger) { +// StartLark 启动飞书长连接（无需公网），收到消息后调用 handler 并回复。 +// ctx 被取消时长连接会退出，便于配置变更时重启。 +func StartLark(ctx context.Context, cfg config.RobotLarkConfig, h MessageHandler, logger *zap.Logger) { if !cfg.Enabled || cfg.AppID == "" || cfg.AppSecret == "" { return } @@ -34,9 +35,11 @@ func StartLark(cfg config.RobotLarkConfig, h MessageHandler, logger *zap.Logger) larkws.WithLogLevel(larkcore.LogLevelInfo), ) go func() { - err := wsClient.Start(context.Background()) - if err != nil { + err := wsClient.Start(ctx) + if err != nil && ctx.Err() == nil { logger.Error("飞书长连接退出", zap.Error(err)) + } else if ctx.Err() != nil { + logger.Info("飞书长连接已按配置重启关闭") } }() logger.Info("飞书长连接已启动（无需公网）", zap.String("app_id", cfg.AppID))