Add files via upload

internal/agent/memory_compressor.go

@@ -17,10 +17,14 @@ import (
 )
 
 const (
-	DefaultMinRecentMessage = 10
+	// DefaultMinRecentMessage 压缩历史消息时保留的最近消息数量，确保最近的对话上下文不被压缩
-	defaultChunkSize        = 10
+	DefaultMinRecentMessage = 5
-	defaultMaxImages        = 3
+	// defaultChunkSize 压缩历史消息时每次处理的消息块大小，将旧消息分成多个块进行摘要
-	defaultSummaryTimeout   = 10 * time.Minute
+	defaultChunkSize = 10
+	// defaultMaxImages 压缩时最多保留的图片数量，超过此数量的图片会被移除以节省上下文空间
+	defaultMaxImages = 3
+	// defaultSummaryTimeout 生成消息摘要时的超时时间
+	defaultSummaryTimeout = 10 * time.Minute
 
 	summaryPromptTemplate = `你是一名负责为安全代理执行上下文压缩的助手，任务是在保持所有关键渗透信息完整的前提下压缩扫描数据。
 

internal/app/app.go

@@ -282,6 +282,13 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
 		agentHandler:       agentHandler,
 	}
 
+	// 设置漏洞工具注册器（内置工具，必须设置）
+	vulnerabilityRegistrar := func() error {
+		registerVulnerabilityTool(mcpServer, db, log.Logger)
+		return nil
+	}
+	configHandler.SetVulnerabilityToolRegistrar(vulnerabilityRegistrar)
+
 	// 设置知识库初始化器（用于动态初始化，需要在 App 创建后设置）
 	configHandler.SetKnowledgeInitializer(func() (*handler.KnowledgeHandler, error) {
 		knowledgeHandler, err := initializeKnowledge(cfg, db, knowledgeDBConn, mcpServer, agentHandler, app, log.Logger)

internal/handler/config.go

@@ -25,6 +25,9 @@ import (
 // KnowledgeToolRegistrar 知识库工具注册器接口
 type KnowledgeToolRegistrar func() error
 
+// VulnerabilityToolRegistrar 漏洞工具注册器接口
+type VulnerabilityToolRegistrar func() error
+
 // RetrieverUpdater 检索器更新接口
 type RetrieverUpdater interface {
 	UpdateConfig(config *knowledge.RetrievalConfig)
@@ -40,19 +43,20 @@ type AppUpdater interface {
 
 // ConfigHandler 配置处理器
 type ConfigHandler struct {
-	configPath             string
+	configPath                 string
-	config                 *config.Config
+	config                     *config.Config
-	mcpServer              *mcp.Server
+	mcpServer                  *mcp.Server
-	executor               *security.Executor
+	executor                   *security.Executor
-	agent                  AgentUpdater            // Agent接口，用于更新Agent配置
+	agent                      AgentUpdater                // Agent接口，用于更新Agent配置
-	attackChainHandler     AttackChainUpdater      // 攻击链处理器接口，用于更新配置
+	attackChainHandler         AttackChainUpdater          // 攻击链处理器接口，用于更新配置
-	externalMCPMgr         *mcp.ExternalMCPManager // 外部MCP管理器
+	externalMCPMgr             *mcp.ExternalMCPManager    // 外部MCP管理器
-	knowledgeToolRegistrar KnowledgeToolRegistrar  // 知识库工具注册器（可选）
+	knowledgeToolRegistrar     KnowledgeToolRegistrar      // 知识库工具注册器（可选）
-	retrieverUpdater       RetrieverUpdater        // 检索器更新器（可选）
+	vulnerabilityToolRegistrar VulnerabilityToolRegistrar // 漏洞工具注册器（可选）
-	knowledgeInitializer   KnowledgeInitializer    // 知识库初始化器（可选）
+	retrieverUpdater           RetrieverUpdater            // 检索器更新器（可选）
-	appUpdater             AppUpdater              // App更新器（可选）
+	knowledgeInitializer       KnowledgeInitializer         // 知识库初始化器（可选）
-	logger                 *zap.Logger
+	appUpdater                 AppUpdater                  // App更新器（可选）
-	mu                     sync.RWMutex
+	logger                     *zap.Logger
+	mu                         sync.RWMutex
 }
 
 // AttackChainUpdater 攻击链处理器更新接口
@@ -87,6 +91,13 @@ func (h *ConfigHandler) SetKnowledgeToolRegistrar(registrar KnowledgeToolRegistr
 	h.knowledgeToolRegistrar = registrar
 }
 
+// SetVulnerabilityToolRegistrar 设置漏洞工具注册器
+func (h *ConfigHandler) SetVulnerabilityToolRegistrar(registrar VulnerabilityToolRegistrar) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.vulnerabilityToolRegistrar = registrar
+}
+
 // SetRetrieverUpdater 设置检索器更新器
 func (h *ConfigHandler) SetRetrieverUpdater(updater RetrieverUpdater) {
 	h.mu.Lock()
@@ -678,6 +689,16 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
 	// 重新注册安全工具
 	h.executor.RegisterTools(h.mcpServer)
 
+	// 重新注册漏洞记录工具（内置工具，必须注册）
+	if h.vulnerabilityToolRegistrar != nil {
+		h.logger.Info("重新注册漏洞记录工具")
+		if err := h.vulnerabilityToolRegistrar(); err != nil {
+			h.logger.Error("重新注册漏洞记录工具失败", zap.Error(err))
+		} else {
+			h.logger.Info("漏洞记录工具已重新注册")
+		}
+	}
+
 	// 如果知识库启用，重新注册知识库工具
 	if h.config.Knowledge.Enabled && h.knowledgeToolRegistrar != nil {
 		h.logger.Info("重新注册知识库工具")