Add files via upload

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,

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)),
 	)

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))

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))