diff --git a/internal/agent/memory_compressor.go b/internal/agent/memory_compressor.go index edab2523..df0ce736 100644 --- a/internal/agent/memory_compressor.go +++ b/internal/agent/memory_compressor.go @@ -17,10 +17,14 @@ import ( ) const ( - DefaultMinRecentMessage = 10 - defaultChunkSize = 10 - defaultMaxImages = 3 - defaultSummaryTimeout = 10 * time.Minute + // DefaultMinRecentMessage 压缩历史消息时保留的最近消息数量,确保最近的对话上下文不被压缩 + DefaultMinRecentMessage = 5 + // defaultChunkSize 压缩历史消息时每次处理的消息块大小,将旧消息分成多个块进行摘要 + defaultChunkSize = 10 + // defaultMaxImages 压缩时最多保留的图片数量,超过此数量的图片会被移除以节省上下文空间 + defaultMaxImages = 3 + // defaultSummaryTimeout 生成消息摘要时的超时时间 + defaultSummaryTimeout = 10 * time.Minute summaryPromptTemplate = `你是一名负责为安全代理执行上下文压缩的助手,任务是在保持所有关键渗透信息完整的前提下压缩扫描数据。 diff --git a/internal/app/app.go b/internal/app/app.go index d05562ab..e66759d9 100644 --- a/internal/app/app.go +++ b/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) diff --git a/internal/handler/config.go b/internal/handler/config.go index cfea5e92..cd0e2a5b 100644 --- a/internal/handler/config.go +++ b/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 - config *config.Config - mcpServer *mcp.Server - executor *security.Executor - agent AgentUpdater // Agent接口,用于更新Agent配置 - attackChainHandler AttackChainUpdater // 攻击链处理器接口,用于更新配置 - externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器 - knowledgeToolRegistrar KnowledgeToolRegistrar // 知识库工具注册器(可选) - retrieverUpdater RetrieverUpdater // 检索器更新器(可选) - knowledgeInitializer KnowledgeInitializer // 知识库初始化器(可选) - appUpdater AppUpdater // App更新器(可选) - logger *zap.Logger - mu sync.RWMutex + configPath string + config *config.Config + mcpServer *mcp.Server + executor *security.Executor + agent AgentUpdater // Agent接口,用于更新Agent配置 + attackChainHandler AttackChainUpdater // 攻击链处理器接口,用于更新配置 + externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器 + knowledgeToolRegistrar KnowledgeToolRegistrar // 知识库工具注册器(可选) + vulnerabilityToolRegistrar VulnerabilityToolRegistrar // 漏洞工具注册器(可选) + retrieverUpdater RetrieverUpdater // 检索器更新器(可选) + knowledgeInitializer KnowledgeInitializer // 知识库初始化器(可选) + appUpdater AppUpdater // App更新器(可选) + 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("重新注册知识库工具")