CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-06 12:37:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2025-11-15 17:50:47 +08:00
committed by GitHub
parent f7344c0090
commit d8ed2d8cb3
18 changed files with 4749 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/agent/agent.go
+120 -15
View File
@@ -19,16 +19,18 @@ import (
// Agent AI代理
type Agent struct {
openAIClient *http.Client
config *config.OpenAIConfig
mcpServer *mcp.Server
logger *zap.Logger
maxIterations int
mu sync.RWMutex // 添加互斥锁以支持并发更新
openAIClient *http.Client
config *config.OpenAIConfig
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger
maxIterations int
mu sync.RWMutex // 添加互斥锁以支持并发更新
toolNameMapping map[string]string // 工具名称映射：OpenAI格式 -> 原始格式（用于外部MCP工具）
}
// NewAgent 创建新的Agent
func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, logger *zap.Logger, maxIterations int) *Agent {
func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger, maxIterations int) *Agent {
// 如果 maxIterations 为 0 或负数，使用默认值 30
if maxIterations <= 0 {
maxIterations = 30
@@ -55,10 +57,12 @@ func NewAgent(cfg *config.OpenAIConfig, mcpServer *mcp.Server, logger *zap.Logge
Timeout: 30 * time.Minute, // 从5分钟增加到30分钟
Transport: transport,
},
config: cfg,
mcpServer: mcpServer,
logger: logger,
maxIterations: maxIterations,
config: cfg,
mcpServer: mcpServer,
externalMCPMgr: externalMCPMgr,
logger: logger,
maxIterations: maxIterations,
toolNameMapping: make(map[string]string), // 初始化工具名称映射
}
}
@@ -578,7 +582,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
// getAvailableTools 获取可用工具
// 从MCP服务器动态获取工具列表，使用简短描述以减少token消耗
func (a *Agent) getAvailableTools() []Tool {
// 从MCP服务器获取所有已注册的工具
// 从MCP服务器获取所有已注册的内部工具
mcpTools := a.mcpServer.GetAllTools()
// 转换为OpenAI格式的工具定义
@@ -603,8 +607,91 @@ func (a *Agent) getAvailableTools() []Tool {
})
}
// 获取外部MCP工具
if a.externalMCPMgr != nil {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
externalTools, err := a.externalMCPMgr.GetAllTools(ctx)
if err != nil {
a.logger.Warn("获取外部MCP工具失败", zap.Error(err))
} else {
// 获取外部MCP配置，用于检查工具启用状态
externalMCPConfigs := a.externalMCPMgr.GetConfigs()
// 清空并重建工具名称映射
a.mu.Lock()
a.toolNameMapping = make(map[string]string)
a.mu.Unlock()
// 将外部MCP工具添加到工具列表（只添加启用的工具）
for _, externalTool := range externalTools {
// 解析工具名称：mcpName::toolName
var mcpName, actualToolName string
if idx := strings.Index(externalTool.Name, "::"); idx > 0 {
mcpName = externalTool.Name[:idx]
actualToolName = externalTool.Name[idx+2:]
} else {
continue // 跳过格式不正确的工具
}
// 检查工具是否启用
enabled := false
if cfg, exists := externalMCPConfigs[mcpName]; exists {
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
enabled = false // MCP未启用，所有工具都禁用
} else {
// MCP已启用，检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具，默认为启用（向后兼容）
if cfg.ToolEnabled == nil {
enabled = true // 未设置工具状态，默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[actualToolName]; exists {
enabled = toolEnabled // 使用配置的工具状态
} else {
enabled = true // 工具未在配置中，默认为启用
}
}
}
// 只添加启用的工具
if !enabled {
continue
}
// 使用简短描述（如果存在），否则使用详细描述
description := externalTool.ShortDescription
if description == "" {
description = externalTool.Description
}
// 转换schema中的类型为OpenAI标准类型
convertedSchema := a.convertSchemaTypes(externalTool.InputSchema)
// 将工具名称中的 "::" 替换为 "__" 以符合OpenAI命名规范
// OpenAI要求工具名称只能包含 [a-zA-Z0-9_-]
openAIName := strings.ReplaceAll(externalTool.Name, "::", "__")
// 保存名称映射关系（OpenAI格式 -> 原始格式）
a.mu.Lock()
a.toolNameMapping[openAIName] = externalTool.Name
a.mu.Unlock()
tools = append(tools, Tool{
Type: "function",
Function: FunctionDefinition{
Name: openAIName, // 使用符合OpenAI规范的名称
Description: description,
Parameters: convertedSchema,
},
})
}
}
}
a.logger.Debug("获取可用工具列表",
zap.Int("count", len(tools)),
zap.Int("internalTools", len(mcpTools)),
zap.Int("totalTools", len(tools)),
)
return tools
@@ -898,8 +985,26 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
zap.Any("args", args),
)
// 通过MCP服务器调用工具
result, executionID, err := a.mcpServer.CallTool(ctx, toolName, args)
var result *mcp.ToolResult
var executionID string
var err error
// 检查是否是外部MCP工具（通过工具名称映射）
a.mu.RLock()
originalToolName, isExternalTool := a.toolNameMapping[toolName]
a.mu.RUnlock()
if isExternalTool && a.externalMCPMgr != nil {
// 使用原始工具名称调用外部MCP工具
a.logger.Debug("调用外部MCP工具",
zap.String("openAIName", toolName),
zap.String("originalName", originalToolName),
)
result, executionID, err = a.externalMCPMgr.CallTool(ctx, originalToolName, args)
} else {
// 调用内部MCP工具
result, executionID, err = a.mcpServer.CallTool(ctx, toolName, args)
}
// 如果调用失败（如工具不存在），返回友好的错误信息而不是抛出异常
if err != nil {
internal/app/app.go
+49 -18
View File
@@ -20,14 +20,15 @@ import (
// App 应用
type App struct {
config *config.Config
logger *logger.Logger
router *gin.Engine
mcpServer *mcp.Server
agent *agent.Agent
executor *security.Executor
db *database.DB
auth *security.AuthManager
config *config.Config
logger *logger.Logger
router *gin.Engine
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager
agent *agent.Agent
executor *security.Executor
db *database.DB
auth *security.AuthManager
}
// New 创建新应用
@@ -76,12 +77,20 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
cfg.Auth.GeneratedPasswordPersistErr = ""
}
// 创建外部MCP管理器（使用与内部MCP服务器相同的存储）
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
if cfg.ExternalMCP.Servers != nil {
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
// 启动所有启用的外部MCP客户端
externalMCPMgr.StartAllEnabled()
}
// 创建Agent
maxIterations := cfg.Agent.MaxIterations
if maxIterations <= 0 {
maxIterations = 30 // 默认值
}
agent := agent.NewAgent(&cfg.OpenAI, mcpServer, log.Logger, maxIterations)
agent := agent.NewAgent(&cfg.OpenAI, mcpServer, externalMCPMgr, log.Logger, maxIterations)
// 获取配置文件路径
configPath := "config.yaml"
@@ -92,9 +101,11 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
// 创建处理器
agentHandler := handler.NewAgentHandler(agent, db, log.Logger)
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器，以便获取外部MCP执行记录
conversationHandler := handler.NewConversationHandler(db, log.Logger)
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, log.Logger)
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, externalMCPMgr, log.Logger)
externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger)
// 设置路由
setupRoutes(
@@ -104,19 +115,21 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
monitorHandler,
conversationHandler,
configHandler,
externalMCPHandler,
mcpServer,
authManager,
)
return &App{
config: cfg,
logger: log,
router: router,
mcpServer: mcpServer,
agent: agent,
executor: executor,
db: db,
auth: authManager,
config: cfg,
logger: log,
router: router,
mcpServer: mcpServer,
externalMCPMgr: externalMCPMgr,
agent: agent,
executor: executor,
db: db,
auth: authManager,
}, nil
}
@@ -144,6 +157,14 @@ func (a *App) Run() error {
return a.router.Run(addr)
}
// Shutdown 关闭应用
func (a *App) Shutdown() {
// 停止所有外部MCP客户端
if a.externalMCPMgr != nil {
a.externalMCPMgr.StopAll()
}
}
// setupRoutes 设置路由
func setupRoutes(
router *gin.Engine,
@@ -152,6 +173,7 @@ func setupRoutes(
monitorHandler *handler.MonitorHandler,
conversationHandler *handler.ConversationHandler,
configHandler *handler.ConfigHandler,
externalMCPHandler *handler.ExternalMCPHandler,
mcpServer *mcp.Server,
authManager *security.AuthManager,
) {
@@ -195,6 +217,15 @@ func setupRoutes(
protected.PUT("/config", configHandler.UpdateConfig)
protected.POST("/config/apply", configHandler.ApplyConfig)
// 外部MCP管理
protected.GET("/external-mcp", externalMCPHandler.GetExternalMCPs)
protected.GET("/external-mcp/stats", externalMCPHandler.GetExternalMCPStats)
protected.GET("/external-mcp/:name", externalMCPHandler.GetExternalMCP)
protected.PUT("/external-mcp/:name", externalMCPHandler.AddOrUpdateExternalMCP)
protected.DELETE("/external-mcp/:name", externalMCPHandler.DeleteExternalMCP)
protected.POST("/external-mcp/:name/start", externalMCPHandler.StartExternalMCP)
protected.POST("/external-mcp/:name/stop", externalMCPHandler.StopExternalMCP)
// MCP端点
protected.POST("/mcp", func(c *gin.Context) {
mcpServer.HandleHTTP(c.Writer, c.Request)
internal/config/config.go
+56 -8
View File
@@ -12,14 +12,15 @@ import (
)
type Config struct {
Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
Agent AgentConfig `yaml:"agent"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
Agent AgentConfig `yaml:"agent"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
}
type ServerConfig struct {
@@ -64,6 +65,32 @@ type AuthConfig struct {
GeneratedPasswordPersisted bool `yaml:"-" json:"-"`
GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
}
// ExternalMCPConfig 外部MCP配置
type ExternalMCPConfig struct {
Servers map[string]ExternalMCPServerConfig `yaml:"servers,omitempty" json:"servers,omitempty"`
}
// ExternalMCPServerConfig 外部MCP服务器配置
type ExternalMCPServerConfig struct {
// stdio模式配置
Command string `yaml:"command,omitempty" json:"command,omitempty"`
Args []string `yaml:"args,omitempty" json:"args,omitempty"`
// HTTP模式配置
Transport string `yaml:"transport,omitempty" json:"transport,omitempty"` // "http" 或 "stdio"
URL string `yaml:"url,omitempty" json:"url,omitempty"`
// 通用配置
Description string `yaml:"description,omitempty" json:"description,omitempty"`
Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty"` // 超时时间（秒）
ExternalMCPEnable bool `yaml:"external_mcp_enable,omitempty" json:"external_mcp_enable,omitempty"` // 是否启用外部MCP
ToolEnabled map[string]bool `yaml:"tool_enabled,omitempty" json:"tool_enabled,omitempty"` // 每个工具的启用状态（工具名称 -> 是否启用）
// 向后兼容字段（已废弃，保留用于读取旧配置）
Enabled bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` // 已废弃，使用 external_mcp_enable
Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` // 已废弃，使用 external_mcp_enable
}
type ToolConfig struct {
Name string `yaml:"name"`
Command string `yaml:"command"`
@@ -152,6 +179,27 @@ func Load(path string) (*Config, error) {
cfg.Security.Tools = tools
}
// 迁移外部MCP配置：将旧的 enabled/disabled 字段迁移到 external_mcp_enable
if cfg.ExternalMCP.Servers != nil {
for name, serverCfg := range cfg.ExternalMCP.Servers {
// 如果已经设置了 external_mcp_enable，跳过迁移
// 否则从 enabled/disabled 字段迁移
// 注意：由于 ExternalMCPEnable 是 bool 类型，零值为 false，所以需要检查是否真的设置了
// 这里我们通过检查旧的 enabled/disabled 字段来判断是否需要迁移
if serverCfg.Disabled {
// 旧配置使用 disabled，迁移到 external_mcp_enable
serverCfg.ExternalMCPEnable = false
} else if serverCfg.Enabled {
// 旧配置使用 enabled，迁移到 external_mcp_enable
serverCfg.ExternalMCPEnable = true
} else {
// 都没有设置，默认为启用
serverCfg.ExternalMCPEnable = true
}
cfg.ExternalMCP.Servers[name] = serverCfg
}
}
return &cfg, nil
}
internal/handler/config.go
+309 -22
View File
@@ -2,6 +2,7 @@ package handler
import (
"bytes"
"context"
"fmt"
"net/http"
"os"
@@ -9,6 +10,7 @@ import (
"strconv"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
@@ -20,13 +22,14 @@ import (
// ConfigHandler 配置处理器
type ConfigHandler struct {
configPath string
config *config.Config
mcpServer *mcp.Server
executor *security.Executor
agent AgentUpdater // Agent接口，用于更新Agent配置
logger *zap.Logger
mu sync.RWMutex
configPath string
config *config.Config
mcpServer *mcp.Server
executor *security.Executor
agent AgentUpdater // Agent接口，用于更新Agent配置
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
logger *zap.Logger
mu sync.RWMutex
}
// AgentUpdater Agent更新接口
@@ -36,14 +39,15 @@ type AgentUpdater interface {
}
// NewConfigHandler 创建新的配置处理器
func NewConfigHandler(configPath string, cfg *config.Config, mcpServer *mcp.Server, executor *security.Executor, agent AgentUpdater, logger *zap.Logger) *ConfigHandler {
func NewConfigHandler(configPath string, cfg *config.Config, mcpServer *mcp.Server, executor *security.Executor, agent AgentUpdater, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger) *ConfigHandler {
return &ConfigHandler{
configPath: configPath,
config: cfg,
mcpServer: mcpServer,
executor: executor,
agent: agent,
logger: logger,
configPath: configPath,
config: cfg,
mcpServer: mcpServer,
executor: executor,
agent: agent,
externalMCPMgr: externalMCPMgr,
logger: logger,
}
}
@@ -60,6 +64,8 @@ type ToolConfigInfo struct {
Name string `json:"name"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具
ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称（如果是外部工具）
}
// GetConfig 获取当前配置
@@ -67,13 +73,14 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
// 获取工具列表
// 获取工具列表（包含内部和外部工具）
tools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools))
for _, tool := range h.config.Security.Tools {
tools = append(tools, ToolConfigInfo{
Name: tool.Name,
Description: tool.ShortDescription,
Enabled: tool.Enabled,
IsExternal: false,
})
// 如果没有简短描述，使用详细描述的前100个字符
if tools[len(tools)-1].Description == "" {
@@ -85,6 +92,65 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
}
}
// 获取外部MCP工具
if h.externalMCPMgr != nil {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
externalTools, err := h.externalMCPMgr.GetAllTools(ctx)
if err == nil {
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
var mcpName, actualToolName string
if idx := strings.Index(externalTool.Name, "::"); idx > 0 {
mcpName = externalTool.Name[:idx]
actualToolName = externalTool.Name[idx+2:]
} else {
continue
}
enabled := false
if cfg, exists := externalMCPConfigs[mcpName]; exists {
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
enabled = false // MCP未启用，所有工具都禁用
} else {
// MCP已启用，检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具，默认为启用（向后兼容）
if cfg.ToolEnabled == nil {
enabled = true // 未设置工具状态，默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[actualToolName]; exists {
enabled = toolEnabled // 使用配置的工具状态
} else {
enabled = true // 工具未在配置中，默认为启用
}
}
}
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
enabled = false
}
description := externalTool.ShortDescription
if description == "" {
description = externalTool.Description
}
if len(description) > 100 {
description = description[:100] + "..."
}
tools = append(tools, ToolConfigInfo{
Name: actualToolName,
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
})
}
}
}
c.JSON(http.StatusOK, GetConfigResponse{
OpenAI: h.config.OpenAI,
MCP: h.config.MCP,
@@ -128,13 +194,14 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
searchTermLower = strings.ToLower(searchTerm)
}
// 获取所有工具并应用搜索过滤
// 获取所有内部工具并应用搜索过滤
allTools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools))
for _, tool := range h.config.Security.Tools {
toolInfo := ToolConfigInfo{
Name: tool.Name,
Description: tool.ShortDescription,
Enabled: tool.Enabled,
IsExternal: false,
}
// 如果没有简短描述，使用详细描述的前100个字符
if toolInfo.Description == "" {
@@ -157,6 +224,81 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
allTools = append(allTools, toolInfo)
}
// 获取外部MCP工具
if h.externalMCPMgr != nil {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
externalTools, err := h.externalMCPMgr.GetAllTools(ctx)
if err != nil {
h.logger.Warn("获取外部MCP工具失败", zap.Error(err))
} else {
// 获取外部MCP配置，用于判断启用状态
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
// 解析工具名称：mcpName::toolName
var mcpName, actualToolName string
if idx := strings.Index(externalTool.Name, "::"); idx > 0 {
mcpName = externalTool.Name[:idx]
actualToolName = externalTool.Name[idx+2:]
} else {
continue // 跳过格式不正确的工具
}
// 获取外部工具的启用状态
enabled := false
if cfg, exists := externalMCPConfigs[mcpName]; exists {
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
enabled = false // MCP未启用，所有工具都禁用
} else {
// MCP已启用，检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具，默认为启用（向后兼容）
if cfg.ToolEnabled == nil {
enabled = true // 未设置工具状态，默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[actualToolName]; exists {
enabled = toolEnabled // 使用配置的工具状态
} else {
enabled = true // 工具未在配置中，默认为启用
}
}
}
// 检查外部MCP是否已连接
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
enabled = false // 未连接时视为禁用
}
description := externalTool.ShortDescription
if description == "" {
description = externalTool.Description
}
if len(description) > 100 {
description = description[:100] + "..."
}
// 如果有关键词，进行搜索过滤
if searchTermLower != "" {
nameLower := strings.ToLower(actualToolName)
descLower := strings.ToLower(description)
if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) {
continue // 不匹配，跳过
}
}
allTools = append(allTools, ToolConfigInfo{
Name: actualToolName, // 显示实际工具名称，不带前缀
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
})
}
}
}
total := len(allTools)
totalPages := (total + pageSize - 1) / pageSize
if totalPages == 0 {
@@ -196,8 +338,10 @@ type UpdateConfigRequest struct {
// ToolEnableStatus 工具启用状态
type ToolEnableStatus struct {
Name string `json:"name"`
Enabled bool `json:"enabled"`
Name string `json:"name"`
Enabled bool `json:"enabled"`
IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具
ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称（如果是外部工具）
}
// UpdateConfig 更新配置
@@ -240,14 +384,28 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
// 更新工具启用状态
if req.Tools != nil {
toolMap := make(map[string]bool)
// 分离内部工具和外部工具
internalToolMap := make(map[string]bool)
// 外部工具状态：MCP名称 -> 工具名称 -> 启用状态
externalMCPToolMap := make(map[string]map[string]bool)
for _, toolStatus := range req.Tools {
toolMap[toolStatus.Name] = toolStatus.Enabled
if toolStatus.IsExternal && toolStatus.ExternalMCP != "" {
// 外部工具：保存每个工具的独立状态
mcpName := toolStatus.ExternalMCP
if externalMCPToolMap[mcpName] == nil {
externalMCPToolMap[mcpName] = make(map[string]bool)
}
externalMCPToolMap[mcpName][toolStatus.Name] = toolStatus.Enabled
} else {
// 内部工具
internalToolMap[toolStatus.Name] = toolStatus.Enabled
}
}
// 更新配置中的工具状态
// 更新内部工具状态
for i := range h.config.Security.Tools {
if enabled, ok := toolMap[h.config.Security.Tools[i].Name]; ok {
if enabled, ok := internalToolMap[h.config.Security.Tools[i].Name]; ok {
h.config.Security.Tools[i].Enabled = enabled
h.logger.Info("更新工具启用状态",
zap.String("tool", h.config.Security.Tools[i].Name),
@@ -255,6 +413,80 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
)
}
}
// 更新外部MCP工具状态
if h.externalMCPMgr != nil {
for mcpName, toolStates := range externalMCPToolMap {
// 更新配置中的工具启用状态
if h.config.ExternalMCP.Servers == nil {
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
}
cfg, exists := h.config.ExternalMCP.Servers[mcpName]
if !exists {
h.logger.Warn("外部MCP配置不存在", zap.String("mcp", mcpName))
continue
}
// 初始化ToolEnabled map
if cfg.ToolEnabled == nil {
cfg.ToolEnabled = make(map[string]bool)
}
// 更新每个工具的启用状态
for toolName, enabled := range toolStates {
cfg.ToolEnabled[toolName] = enabled
h.logger.Info("更新外部工具启用状态",
zap.String("mcp", mcpName),
zap.String("tool", toolName),
zap.Bool("enabled", enabled),
)
}
// 检查是否有任何工具启用，如果有则启用MCP
hasEnabledTool := false
for _, enabled := range cfg.ToolEnabled {
if enabled {
hasEnabledTool = true
break
}
}
// 如果MCP之前未启用，但现在有工具启用，则启用MCP
// 如果MCP之前已启用，保持启用状态（允许部分工具禁用）
if !cfg.ExternalMCPEnable && hasEnabledTool {
cfg.ExternalMCPEnable = true
h.logger.Info("自动启用外部MCP（因为有工具启用）", zap.String("mcp", mcpName))
}
h.config.ExternalMCP.Servers[mcpName] = cfg
}
// 同步更新 externalMCPMgr 中的配置，确保 GetConfigs() 返回最新配置
// 在循环外部统一更新，避免重复调用
h.externalMCPMgr.LoadConfigs(&h.config.ExternalMCP)
// 处理MCP连接状态
for mcpName := range externalMCPToolMap {
cfg := h.config.ExternalMCP.Servers[mcpName]
// 如果MCP需要启用，确保客户端已启动
if cfg.ExternalMCPEnable {
// 启动外部MCP（如果未启动）
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
if err := h.externalMCPMgr.StartClient(mcpName); err != nil {
h.logger.Warn("启动外部MCP失败",
zap.String("mcp", mcpName),
zap.Error(err),
)
} else {
h.logger.Info("启动外部MCP",
zap.String("mcp", mcpName),
)
}
}
}
}
}
}
// 保存配置到文件
@@ -318,6 +550,33 @@ func (h *ConfigHandler) saveConfig() error {
updateAgentConfig(root, h.config.Agent.MaxIterations)
updateMCPConfig(root, h.config.MCP)
updateOpenAIConfig(root, h.config.OpenAI)
// 更新外部MCP配置（使用external_mcp.go中的函数，同一包中可直接调用）
// 读取原始配置以保持向后兼容
originalConfigs := make(map[string]map[string]bool)
externalMCPNode := findMapValue(root, "external_mcp")
if externalMCPNode != nil && externalMCPNode.Kind == yaml.MappingNode {
serversNode := findMapValue(externalMCPNode, "servers")
if serversNode != nil && serversNode.Kind == yaml.MappingNode {
for i := 0; i < len(serversNode.Content); i += 2 {
if i+1 >= len(serversNode.Content) {
break
}
nameNode := serversNode.Content[i]
serverNode := serversNode.Content[i+1]
if nameNode.Kind == yaml.ScalarNode && serverNode.Kind == yaml.MappingNode {
serverName := nameNode.Value
originalConfigs[serverName] = make(map[string]bool)
if enabledVal := findBoolInMap(serverNode, "enabled"); enabledVal != nil {
originalConfigs[serverName]["enabled"] = *enabledVal
}
if disabledVal := findBoolInMap(serverNode, "disabled"); disabledVal != nil {
originalConfigs[serverName]["disabled"] = *disabledVal
}
}
}
}
}
updateExternalMCPConfig(root, h.config.ExternalMCP, originalConfigs)
if err := writeYAMLDocument(h.configPath, root); err != nil {
return fmt.Errorf("保存配置文件失败: %w", err)
@@ -504,6 +763,34 @@ func setIntInMap(mapNode *yaml.Node, key string, value int) {
valueNode.Value = fmt.Sprintf("%d", value)
}
func findBoolInMap(mapNode *yaml.Node, key string) *bool {
if mapNode == nil || mapNode.Kind != yaml.MappingNode {
return nil
}
for i := 0; i < len(mapNode.Content); i += 2 {
if i+1 >= len(mapNode.Content) {
break
}
keyNode := mapNode.Content[i]
valueNode := mapNode.Content[i+1]
if keyNode.Kind == yaml.ScalarNode && keyNode.Value == key {
if valueNode.Kind == yaml.ScalarNode {
if valueNode.Value == "true" {
result := true
return &result
} else if valueNode.Value == "false" {
result := false
return &result
}
}
return nil
}
}
return nil
}
func setBoolInMap(mapNode *yaml.Node, key string, value bool) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.ScalarNode
internal/handler/external_mcp.go
+510
View File
@@ -0,0 +1,510 @@
package handler
import (
"fmt"
"net/http"
"os"
"sync"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
"gopkg.in/yaml.v3"
)
// ExternalMCPHandler 外部MCP处理器
type ExternalMCPHandler struct {
manager *mcp.ExternalMCPManager
config *config.Config
configPath string
logger *zap.Logger
mu sync.RWMutex
}
// NewExternalMCPHandler 创建外部MCP处理器
func NewExternalMCPHandler(manager *mcp.ExternalMCPManager, cfg *config.Config, configPath string, logger *zap.Logger) *ExternalMCPHandler {
return &ExternalMCPHandler{
manager: manager,
config: cfg,
configPath: configPath,
logger: logger,
}
}
// GetExternalMCPs 获取所有外部MCP配置
func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
configs := h.manager.GetConfigs()
// 获取所有外部MCP的工具数量
toolCounts := h.manager.GetToolCounts()
// 转换为响应格式
result := make(map[string]ExternalMCPResponse)
for name, cfg := range configs {
client, exists := h.manager.GetClient(name)
status := "disconnected"
if exists {
status = client.GetStatus()
} else if h.isEnabled(cfg) {
status = "disconnected"
} else {
status = "disabled"
}
toolCount := toolCounts[name]
result[name] = ExternalMCPResponse{
Config: cfg,
Status: status,
ToolCount: toolCount,
}
}
c.JSON(http.StatusOK, gin.H{
"servers": result,
"stats": h.manager.GetStats(),
})
}
// GetExternalMCP 获取单个外部MCP配置
func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) {
name := c.Param("name")
h.mu.RLock()
defer h.mu.RUnlock()
configs := h.manager.GetConfigs()
cfg, exists := configs[name]
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "外部MCP配置不存在"})
return
}
client, clientExists := h.manager.GetClient(name)
status := "disconnected"
if clientExists {
status = client.GetStatus()
} else if h.isEnabled(cfg) {
status = "disconnected"
} else {
status = "disabled"
}
// 获取工具数量
toolCount := 0
if clientExists && client.IsConnected() {
if count, err := h.manager.GetToolCount(name); err == nil {
toolCount = count
}
}
c.JSON(http.StatusOK, ExternalMCPResponse{
Config: cfg,
Status: status,
ToolCount: toolCount,
})
}
// AddOrUpdateExternalMCP 添加或更新外部MCP配置
func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) {
var req AddOrUpdateExternalMCPRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
return
}
name := c.Param("name")
if name == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "名称不能为空"})
return
}
// 验证配置
if err := h.validateConfig(req.Config); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.mu.Lock()
defer h.mu.Unlock()
// 添加或更新配置
if err := h.manager.AddOrUpdateConfig(name, req.Config); err != nil {
h.logger.Error("添加或更新外部MCP配置失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "添加或更新配置失败: " + err.Error()})
return
}
// 更新内存中的配置
if h.config.ExternalMCP.Servers == nil {
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
}
// 如果用户提供了 disabled 或 enabled 字段，保留它们以保持向后兼容
// 同时将值迁移到 external_mcp_enable
cfg := req.Config
if req.Config.Disabled {
// 用户设置了 disabled: true
cfg.ExternalMCPEnable = false
cfg.Disabled = true
cfg.Enabled = false
} else if req.Config.Enabled {
// 用户设置了 enabled: true
cfg.ExternalMCPEnable = true
cfg.Enabled = true
cfg.Disabled = false
} else if !req.Config.ExternalMCPEnable {
// 用户没有设置任何字段，且 external_mcp_enable 为 false
// 检查现有配置是否有旧字段
if existingCfg, exists := h.config.ExternalMCP.Servers[name]; exists {
// 保留现有的旧字段
cfg.Enabled = existingCfg.Enabled
cfg.Disabled = existingCfg.Disabled
}
} else {
// 用户通过新字段启用了（external_mcp_enable: true），但没有设置旧字段
// 为了向后兼容，我们设置 enabled: true
// 这样即使原始配置中有 disabled: false，也会被转换为 enabled: true
cfg.Enabled = true
cfg.Disabled = false
}
h.config.ExternalMCP.Servers[name] = cfg
// 保存到配置文件
if err := h.saveConfig(); err != nil {
h.logger.Error("保存配置失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
return
}
h.logger.Info("外部MCP配置已更新", zap.String("name", name))
c.JSON(http.StatusOK, gin.H{"message": "配置已更新"})
}
// DeleteExternalMCP 删除外部MCP配置
func (h *ExternalMCPHandler) DeleteExternalMCP(c *gin.Context) {
name := c.Param("name")
h.mu.Lock()
defer h.mu.Unlock()
// 移除配置
if err := h.manager.RemoveConfig(name); err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "配置不存在"})
return
}
// 从内存配置中删除
if h.config.ExternalMCP.Servers != nil {
delete(h.config.ExternalMCP.Servers, name)
}
// 保存到配置文件
if err := h.saveConfig(); err != nil {
h.logger.Error("保存配置失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
return
}
h.logger.Info("外部MCP配置已删除", zap.String("name", name))
c.JSON(http.StatusOK, gin.H{"message": "配置已删除"})
}
// StartExternalMCP 启动外部MCP
func (h *ExternalMCPHandler) StartExternalMCP(c *gin.Context) {
name := c.Param("name")
h.mu.Lock()
defer h.mu.Unlock()
// 更新配置为启用
if h.config.ExternalMCP.Servers == nil {
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
}
cfg := h.config.ExternalMCP.Servers[name]
cfg.ExternalMCPEnable = true
h.config.ExternalMCP.Servers[name] = cfg
// 保存到配置文件
if err := h.saveConfig(); err != nil {
h.logger.Error("保存配置失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
return
}
// 启动客户端（这可能会花费一些时间）
h.logger.Info("开始启动外部MCP", zap.String("name", name))
if err := h.manager.StartClient(name); err != nil {
h.logger.Error("启动外部MCP失败", zap.String("name", name), zap.Error(err))
c.JSON(http.StatusBadRequest, gin.H{
"error": err.Error(),
"status": "error",
})
return
}
// 获取连接状态
client, exists := h.manager.GetClient(name)
status := "disconnected"
if exists {
status = client.GetStatus()
}
h.logger.Info("外部MCP启动完成", zap.String("name", name), zap.String("status", status))
c.JSON(http.StatusOK, gin.H{
"message": "外部MCP启动完成",
"status": status,
})
}
// StopExternalMCP 停止外部MCP
func (h *ExternalMCPHandler) StopExternalMCP(c *gin.Context) {
name := c.Param("name")
h.mu.Lock()
defer h.mu.Unlock()
// 停止客户端
if err := h.manager.StopClient(name); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
// 更新配置
if h.config.ExternalMCP.Servers == nil {
h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig)
}
cfg := h.config.ExternalMCP.Servers[name]
cfg.ExternalMCPEnable = false
h.config.ExternalMCP.Servers[name] = cfg
// 保存到配置文件
if err := h.saveConfig(); err != nil {
h.logger.Error("保存配置失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
return
}
h.logger.Info("外部MCP已停止", zap.String("name", name))
c.JSON(http.StatusOK, gin.H{"message": "外部MCP已停止"})
}
// GetExternalMCPStats 获取统计信息
func (h *ExternalMCPHandler) GetExternalMCPStats(c *gin.Context) {
stats := h.manager.GetStats()
c.JSON(http.StatusOK, stats)
}
// validateConfig 验证配置
func (h *ExternalMCPHandler) validateConfig(cfg config.ExternalMCPServerConfig) error {
transport := cfg.Transport
if transport == "" {
// 如果没有指定transport，根据是否有command或url判断
if cfg.Command != "" {
transport = "stdio"
} else if cfg.URL != "" {
transport = "http"
} else {
return fmt.Errorf("需要指定command（stdio模式）或url（http模式）")
}
}
switch transport {
case "http":
if cfg.URL == "" {
return fmt.Errorf("HTTP模式需要URL")
}
case "stdio":
if cfg.Command == "" {
return fmt.Errorf("stdio模式需要command")
}
default:
return fmt.Errorf("不支持的传输模式: %s，支持的模式: http, stdio", transport)
}
return nil
}
// isEnabled 检查是否启用
func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool {
// 优先使用 ExternalMCPEnable 字段
// 如果没有设置，检查旧的 enabled/disabled 字段（向后兼容）
if cfg.ExternalMCPEnable {
return true
}
// 向后兼容：检查旧字段
if cfg.Disabled {
return false
}
if cfg.Enabled {
return true
}
// 都没有设置，默认为启用
return true
}
// saveConfig 保存配置到文件
func (h *ExternalMCPHandler) saveConfig() error {
// 读取现有配置文件并创建备份
data, err := os.ReadFile(h.configPath)
if err != nil {
return fmt.Errorf("读取配置文件失败: %w", err)
}
if err := os.WriteFile(h.configPath+".backup", data, 0644); err != nil {
h.logger.Warn("创建配置备份失败", zap.Error(err))
}
root, err := loadYAMLDocument(h.configPath)
if err != nil {
return fmt.Errorf("解析配置文件失败: %w", err)
}
// 在更新前，读取原始配置中的 enabled/disabled 字段，以便保持向后兼容
originalConfigs := make(map[string]map[string]bool)
externalMCPNode := findMapValue(root.Content[0], "external_mcp")
if externalMCPNode != nil && externalMCPNode.Kind == yaml.MappingNode {
serversNode := findMapValue(externalMCPNode, "servers")
if serversNode != nil && serversNode.Kind == yaml.MappingNode {
// 遍历现有的服务器配置，保存 enabled/disabled 字段
for i := 0; i < len(serversNode.Content); i += 2 {
if i+1 >= len(serversNode.Content) {
break
}
nameNode := serversNode.Content[i]
serverNode := serversNode.Content[i+1]
if nameNode.Kind == yaml.ScalarNode && serverNode.Kind == yaml.MappingNode {
serverName := nameNode.Value
originalConfigs[serverName] = make(map[string]bool)
// 检查是否有 enabled 字段
if enabledVal := findBoolInMap(serverNode, "enabled"); enabledVal != nil {
originalConfigs[serverName]["enabled"] = *enabledVal
}
// 检查是否有 disabled 字段
if disabledVal := findBoolInMap(serverNode, "disabled"); disabledVal != nil {
originalConfigs[serverName]["disabled"] = *disabledVal
}
}
}
}
}
// 更新外部MCP配置
updateExternalMCPConfig(root, h.config.ExternalMCP, originalConfigs)
if err := writeYAMLDocument(h.configPath, root); err != nil {
return fmt.Errorf("保存配置文件失败: %w", err)
}
h.logger.Info("配置已保存", zap.String("path", h.configPath))
return nil
}
// updateExternalMCPConfig 更新外部MCP配置
func updateExternalMCPConfig(doc *yaml.Node, cfg config.ExternalMCPConfig, originalConfigs map[string]map[string]bool) {
root := doc.Content[0]
externalMCPNode := ensureMap(root, "external_mcp")
serversNode := ensureMap(externalMCPNode, "servers")
// 清空现有服务器配置
serversNode.Content = nil
// 添加新的服务器配置
for name, serverCfg := range cfg.Servers {
// 添加服务器名称键
nameNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: name}
serverNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
serversNode.Content = append(serversNode.Content, nameNode, serverNode)
// 设置服务器配置字段
if serverCfg.Command != "" {
setStringInMap(serverNode, "command", serverCfg.Command)
}
if len(serverCfg.Args) > 0 {
setStringArrayInMap(serverNode, "args", serverCfg.Args)
}
if serverCfg.Transport != "" {
setStringInMap(serverNode, "transport", serverCfg.Transport)
}
if serverCfg.URL != "" {
setStringInMap(serverNode, "url", serverCfg.URL)
}
if serverCfg.Description != "" {
setStringInMap(serverNode, "description", serverCfg.Description)
}
if serverCfg.Timeout > 0 {
setIntInMap(serverNode, "timeout", serverCfg.Timeout)
}
// 保存 external_mcp_enable 字段（新字段）
setBoolInMap(serverNode, "external_mcp_enable", serverCfg.ExternalMCPEnable)
// 保存 tool_enabled 字段（每个工具的启用状态）
if serverCfg.ToolEnabled != nil && len(serverCfg.ToolEnabled) > 0 {
toolEnabledNode := ensureMap(serverNode, "tool_enabled")
for toolName, enabled := range serverCfg.ToolEnabled {
setBoolInMap(toolEnabledNode, toolName, enabled)
}
}
// 保留旧的 enabled/disabled 字段以保持向后兼容
originalFields, hasOriginal := originalConfigs[name]
// 如果原始配置中有 enabled 字段，保留它
if hasOriginal {
if enabledVal, hasEnabled := originalFields["enabled"]; hasEnabled {
setBoolInMap(serverNode, "enabled", enabledVal)
}
// 如果原始配置中有 disabled 字段，保留它
// 注意：由于 omitempty，disabled: false 不会被保存，但 disabled: true 会被保存
if disabledVal, hasDisabled := originalFields["disabled"]; hasDisabled {
if disabledVal {
setBoolInMap(serverNode, "disabled", disabledVal)
} else {
// 如果原始配置中有 disabled: false，我们保存 enabled: true 来等效表示
// 因为 disabled: false 等价于 enabled: true
setBoolInMap(serverNode, "enabled", true)
}
}
}
// 如果用户在当前请求中明确设置了这些字段，也保存它们
if serverCfg.Enabled {
setBoolInMap(serverNode, "enabled", serverCfg.Enabled)
}
if serverCfg.Disabled {
setBoolInMap(serverNode, "disabled", serverCfg.Disabled)
} else if !hasOriginal && serverCfg.ExternalMCPEnable {
// 如果用户通过新字段启用了，且原始配置中没有旧字段，保存 enabled: true 以保持向后兼容
setBoolInMap(serverNode, "enabled", true)
}
}
}
// setStringArrayInMap 设置字符串数组
func setStringArrayInMap(mapNode *yaml.Node, key string, values []string) {
_, valueNode := ensureKeyValue(mapNode, key)
valueNode.Kind = yaml.SequenceNode
valueNode.Tag = "!!seq"
valueNode.Content = nil
for _, v := range values {
itemNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: v}
valueNode.Content = append(valueNode.Content, itemNode)
}
}
// AddOrUpdateExternalMCPRequest 添加或更新外部MCP请求
type AddOrUpdateExternalMCPRequest struct {
Config config.ExternalMCPServerConfig `json:"config"`
}
// ExternalMCPResponse 外部MCP响应
type ExternalMCPResponse struct {
Config config.ExternalMCPServerConfig `json:"config"`
Status string `json:"status"` // "connected", "disconnected", "disabled", "error"
ToolCount int `json:"tool_count"` // 工具数量
}
internal/handler/external_mcp_test.go
+518
View File
@@ -0,0 +1,518 @@
package handler
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func setupTestRouter() (*gin.Engine, *ExternalMCPHandler, string) {
gin.SetMode(gin.TestMode)
router := gin.New()
// 创建临时配置文件
tmpFile, err := os.CreateTemp("", "test-config-*.yaml")
if err != nil {
panic(err)
}
tmpFile.WriteString("server:\n host: 0.0.0.0\n port: 8080\n")
tmpFile.Close()
configPath := tmpFile.Name()
logger := zap.NewNop()
manager := mcp.NewExternalMCPManager(logger)
cfg := &config.Config{
ExternalMCP: config.ExternalMCPConfig{
Servers: make(map[string]config.ExternalMCPServerConfig),
},
}
handler := NewExternalMCPHandler(manager, cfg, configPath, logger)
api := router.Group("/api")
api.GET("/external-mcp", handler.GetExternalMCPs)
api.GET("/external-mcp/stats", handler.GetExternalMCPStats)
api.GET("/external-mcp/:name", handler.GetExternalMCP)
api.PUT("/external-mcp/:name", handler.AddOrUpdateExternalMCP)
api.DELETE("/external-mcp/:name", handler.DeleteExternalMCP)
api.POST("/external-mcp/:name/start", handler.StartExternalMCP)
api.POST("/external-mcp/:name/stop", handler.StopExternalMCP)
return router, handler, configPath
}
func cleanupTestConfig(configPath string) {
os.Remove(configPath)
os.Remove(configPath + ".backup")
}
func TestExternalMCPHandler_AddOrUpdateExternalMCP_Stdio(t *testing.T) {
router, _, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
// 测试添加stdio模式的配置
configJSON := `{
"command": "python3",
"args": ["/path/to/script.py", "--server", "http://example.com"],
"description": "Test stdio MCP",
"timeout": 300,
"enabled": true
}`
var configObj config.ExternalMCPServerConfig
if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil {
t.Fatalf("解析配置JSON失败: %v", err)
}
reqBody := AddOrUpdateExternalMCPRequest{
Config: configObj,
}
body, _ := json.Marshal(reqBody)
req := httptest.NewRequest("PUT", "/api/external-mcp/test-stdio", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
// 验证配置已添加
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-stdio", nil)
w2 := httptest.NewRecorder()
router.ServeHTTP(w2, req2)
if w2.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
}
var response ExternalMCPResponse
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
if response.Config.Command != "python3" {
t.Errorf("期望command为python3，实际%s", response.Config.Command)
}
if len(response.Config.Args) != 3 {
t.Errorf("期望args长度为3，实际%d", len(response.Config.Args))
}
if response.Config.Description != "Test stdio MCP" {
t.Errorf("期望description为'Test stdio MCP'，实际%s", response.Config.Description)
}
if response.Config.Timeout != 300 {
t.Errorf("期望timeout为300，实际%d", response.Config.Timeout)
}
if !response.Config.Enabled {
t.Error("期望enabled为true")
}
}
func TestExternalMCPHandler_AddOrUpdateExternalMCP_HTTP(t *testing.T) {
router, _, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
// 测试添加HTTP模式的配置
configJSON := `{
"transport": "http",
"url": "http://127.0.0.1:8081/mcp",
"enabled": true
}`
var configObj config.ExternalMCPServerConfig
if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil {
t.Fatalf("解析配置JSON失败: %v", err)
}
reqBody := AddOrUpdateExternalMCPRequest{
Config: configObj,
}
body, _ := json.Marshal(reqBody)
req := httptest.NewRequest("PUT", "/api/external-mcp/test-http", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
// 验证配置已添加
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-http", nil)
w2 := httptest.NewRecorder()
router.ServeHTTP(w2, req2)
if w2.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
}
var response ExternalMCPResponse
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
if response.Config.Transport != "http" {
t.Errorf("期望transport为http，实际%s", response.Config.Transport)
}
if response.Config.URL != "http://127.0.0.1:8081/mcp" {
t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL)
}
if !response.Config.Enabled {
t.Error("期望enabled为true")
}
}
func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidConfig(t *testing.T) {
router, _, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
testCases := []struct {
name string
configJSON string
expectedErr string
}{
{
name: "缺少command和url",
configJSON: `{"enabled": true}`,
expectedErr: "需要指定command（stdio模式）或url（http模式）",
},
{
name: "stdio模式缺少command",
configJSON: `{"args": ["test"], "enabled": true}`,
expectedErr: "stdio模式需要command",
},
{
name: "http模式缺少url",
configJSON: `{"transport": "http", "enabled": true}`,
expectedErr: "HTTP模式需要URL",
},
{
name: "无效的transport",
configJSON: `{"transport": "invalid", "enabled": true}`,
expectedErr: "不支持的传输模式",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
var configObj config.ExternalMCPServerConfig
if err := json.Unmarshal([]byte(tc.configJSON), &configObj); err != nil {
t.Fatalf("解析配置JSON失败: %v", err)
}
reqBody := AddOrUpdateExternalMCPRequest{
Config: configObj,
}
body, _ := json.Marshal(reqBody)
req := httptest.NewRequest("PUT", "/api/external-mcp/test-invalid", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String())
}
var response map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
errorMsg := response["error"].(string)
// 对于stdio模式缺少command的情况，错误信息可能略有不同
if tc.name == "stdio模式缺少command" {
if !strings.Contains(errorMsg, "stdio") && !strings.Contains(errorMsg, "command") {
t.Errorf("期望错误信息包含'stdio'或'command'，实际'%s'", errorMsg)
}
} else if !strings.Contains(errorMsg, tc.expectedErr) {
t.Errorf("期望错误信息包含'%s'，实际'%s'", tc.expectedErr, errorMsg)
}
})
}
}
func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) {
router, handler, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
// 先添加一个配置
configObj := config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
}
handler.manager.AddOrUpdateConfig("test-delete", configObj)
// 删除配置
req := httptest.NewRequest("DELETE", "/api/external-mcp/test-delete", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
// 验证配置已删除
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-delete", nil)
w2 := httptest.NewRecorder()
router.ServeHTTP(w2, req2)
if w2.Code != http.StatusNotFound {
t.Errorf("期望状态码404，实际%d: %s", w2.Code, w2.Body.String())
}
}
func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) {
router, handler, _ := setupTestRouter()
// 添加多个配置
handler.manager.AddOrUpdateConfig("test1", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
})
handler.manager.AddOrUpdateConfig("test2", config.ExternalMCPServerConfig{
URL: "http://127.0.0.1:8081/mcp",
Enabled: false,
})
req := httptest.NewRequest("GET", "/api/external-mcp", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
var response map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
servers := response["servers"].(map[string]interface{})
if len(servers) != 2 {
t.Errorf("期望2个服务器，实际%d", len(servers))
}
if _, ok := servers["test1"]; !ok {
t.Error("期望包含test1")
}
if _, ok := servers["test2"]; !ok {
t.Error("期望包含test2")
}
stats := response["stats"].(map[string]interface{})
if int(stats["total"].(float64)) != 2 {
t.Errorf("期望总数为2，实际%d", int(stats["total"].(float64)))
}
}
func TestExternalMCPHandler_GetExternalMCPStats(t *testing.T) {
router, handler, _ := setupTestRouter()
// 添加配置
handler.manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
})
handler.manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
URL: "http://127.0.0.1:8081/mcp",
Enabled: true,
})
handler.manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: false,
Disabled: true,
})
req := httptest.NewRequest("GET", "/api/external-mcp/stats", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
var stats map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &stats); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
if int(stats["total"].(float64)) != 3 {
t.Errorf("期望总数为3，实际%d", int(stats["total"].(float64)))
}
if int(stats["enabled"].(float64)) != 2 {
t.Errorf("期望启用数为2，实际%d", int(stats["enabled"].(float64)))
}
if int(stats["disabled"].(float64)) != 1 {
t.Errorf("期望停用数为1，实际%d", int(stats["disabled"].(float64)))
}
}
func TestExternalMCPHandler_StartStopExternalMCP(t *testing.T) {
router, handler, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
// 添加一个禁用的配置
handler.manager.AddOrUpdateConfig("test-start-stop", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: false,
Disabled: true,
})
// 测试启动（可能会失败，因为没有真实的服务器）
req := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/start", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
// 启动可能会失败，但应该返回合理的状态码
if w.Code != http.StatusOK {
// 如果启动失败，应该是400或500
if w.Code != http.StatusBadRequest && w.Code != http.StatusInternalServerError {
t.Errorf("期望状态码200/400/500，实际%d: %s", w.Code, w.Body.String())
}
}
// 测试停止
req2 := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/stop", nil)
w2 := httptest.NewRecorder()
router.ServeHTTP(w2, req2)
if w2.Code != http.StatusOK {
t.Errorf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
}
}
func TestExternalMCPHandler_GetExternalMCP_NotFound(t *testing.T) {
router, _, _ := setupTestRouter()
req := httptest.NewRequest("GET", "/api/external-mcp/nonexistent", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("期望状态码404，实际%d: %s", w.Code, w.Body.String())
}
}
func TestExternalMCPHandler_DeleteExternalMCP_NotFound(t *testing.T) {
router, _, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
req := httptest.NewRequest("DELETE", "/api/external-mcp/nonexistent", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
// 删除不存在的配置可能返回200（幂等操作）或404，都是合理的
if w.Code != http.StatusNotFound && w.Code != http.StatusOK {
t.Errorf("期望状态码404或200，实际%d: %s", w.Code, w.Body.String())
}
}
func TestExternalMCPHandler_AddOrUpdateExternalMCP_EmptyName(t *testing.T) {
router, _, _ := setupTestRouter()
configObj := config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
}
reqBody := AddOrUpdateExternalMCPRequest{
Config: configObj,
}
body, _ := json.Marshal(reqBody)
req := httptest.NewRequest("PUT", "/api/external-mcp/", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
// 空名称应该返回404或400
if w.Code != http.StatusNotFound && w.Code != http.StatusBadRequest {
t.Errorf("期望状态码404或400，实际%d: %s", w.Code, w.Body.String())
}
}
func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidJSON(t *testing.T) {
router, _, _ := setupTestRouter()
// 发送无效的JSON
body := []byte(`{"config": invalid json}`)
req := httptest.NewRequest("PUT", "/api/external-mcp/test", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String())
}
}
func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) {
router, handler, configPath := setupTestRouter()
defer cleanupTestConfig(configPath)
// 先添加配置
config1 := config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
}
handler.manager.AddOrUpdateConfig("test-update", config1)
// 更新配置
config2 := config.ExternalMCPServerConfig{
URL: "http://127.0.0.1:8081/mcp",
Enabled: true,
}
reqBody := AddOrUpdateExternalMCPRequest{
Config: config2,
}
body, _ := json.Marshal(reqBody)
req := httptest.NewRequest("PUT", "/api/external-mcp/test-update", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String())
}
// 验证配置已更新
req2 := httptest.NewRequest("GET", "/api/external-mcp/test-update", nil)
w2 := httptest.NewRecorder()
router.ServeHTTP(w2, req2)
if w2.Code != http.StatusOK {
t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String())
}
var response ExternalMCPResponse
if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil {
t.Fatalf("解析响应失败: %v", err)
}
if response.Config.URL != "http://127.0.0.1:8081/mcp" {
t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL)
}
if response.Config.Command != "" {
t.Errorf("期望command为空，实际%s", response.Config.Command)
}
}
internal/handler/monitor.go
+77 -16
View File
@@ -14,22 +14,29 @@ import (
// MonitorHandler 监控处理器
type MonitorHandler struct {
mcpServer *mcp.Server
executor *security.Executor
db *database.DB
logger *zap.Logger
mcpServer *mcp.Server
externalMCPMgr *mcp.ExternalMCPManager
executor *security.Executor
db *database.DB
logger *zap.Logger
}
// NewMonitorHandler 创建新的监控处理器
func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, db *database.DB, logger *zap.Logger) *MonitorHandler {
return &MonitorHandler{
mcpServer: mcpServer,
executor: executor,
db: db,
logger: logger,
mcpServer: mcpServer,
externalMCPMgr: nil, // 将在创建后设置
executor: executor,
db: db,
logger: logger,
}
}
// SetExternalMCPManager 设置外部MCP管理器
func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) {
h.externalMCPMgr = mgr
}
// MonitorResponse 监控响应
type MonitorResponse struct {
Executions []*mcp.ToolExecution `json:"executions"`
@@ -128,15 +135,49 @@ func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int) ([]*mc
}
func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
// 合并内部MCP服务器和外部MCP管理器的统计信息
stats := make(map[string]*mcp.ToolStats)
// 加载内部MCP服务器的统计信息
if h.db == nil {
return h.mcpServer.GetStats()
internalStats := h.mcpServer.GetStats()
for k, v := range internalStats {
stats[k] = v
}
} else {
dbStats, err := h.db.LoadToolStats()
if err != nil {
h.logger.Warn("从数据库加载统计信息失败，回退到内存数据", zap.Error(err))
internalStats := h.mcpServer.GetStats()
for k, v := range internalStats {
stats[k] = v
}
} else {
for k, v := range dbStats {
stats[k] = v
}
}
}
stats, err := h.db.LoadToolStats()
if err != nil {
h.logger.Warn("从数据库加载统计信息失败，回退到内存数据", zap.Error(err))
return h.mcpServer.GetStats()
// 合并外部MCP管理器的统计信息
if h.externalMCPMgr != nil {
externalStats := h.externalMCPMgr.GetToolStats()
for k, v := range externalStats {
// 如果已存在，合并统计信息
if existing, exists := stats[k]; exists {
existing.TotalCalls += v.TotalCalls
existing.SuccessCalls += v.SuccessCalls
existing.FailedCalls += v.FailedCalls
// 使用最新的调用时间
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
existing.LastCallTime = v.LastCallTime
}
} else {
stats[k] = v
}
}
}
return stats
}
@@ -145,13 +186,32 @@ func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
func (h *MonitorHandler) GetExecution(c *gin.Context) {
id := c.Param("id")
// 先从内部MCP服务器查找
exec, exists := h.mcpServer.GetExecution(id)
if !exists {
c.JSON(http.StatusNotFound, gin.H{"error": "执行记录未找到"})
if exists {
c.JSON(http.StatusOK, exec)
return
}
c.JSON(http.StatusOK, exec)
// 如果找不到，尝试从外部MCP管理器查找
if h.externalMCPMgr != nil {
exec, exists = h.externalMCPMgr.GetExecution(id)
if exists {
c.JSON(http.StatusOK, exec)
return
}
}
// 如果都找不到，尝试从数据库查找（如果使用数据库存储）
if h.db != nil {
exec, err := h.db.GetToolExecution(id)
if err == nil && exec != nil {
c.JSON(http.StatusOK, exec)
return
}
}
c.JSON(http.StatusNotFound, gin.H{"error": "执行记录未找到"})
}
// GetStats 获取统计信息
@@ -160,3 +220,4 @@ func (h *MonitorHandler) GetStats(c *gin.Context) {
c.JSON(http.StatusOK, stats)
}
internal/mcp/client.go
+474
View File
@@ -0,0 +1,474 @@
package mcp
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os/exec"
"sync"
"time"
"github.com/google/uuid"
"go.uber.org/zap"
)
// ExternalMCPClient 外部MCP客户端接口
type ExternalMCPClient interface {
// Initialize 初始化连接
Initialize(ctx context.Context) error
// ListTools 列出工具
ListTools(ctx context.Context) ([]Tool, error)
// CallTool 调用工具
CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error)
// Close 关闭连接
Close() error
// IsConnected 检查是否已连接
IsConnected() bool
// GetStatus 获取状态
GetStatus() string
}
// HTTPMCPClient HTTP模式的MCP客户端
type HTTPMCPClient struct {
url string
timeout time.Duration
client *http.Client
logger *zap.Logger
mu sync.RWMutex
status string // "disconnected", "connecting", "connected", "error"
}
// NewHTTPMCPClient 创建HTTP模式的MCP客户端
func NewHTTPMCPClient(url string, timeout time.Duration, logger *zap.Logger) *HTTPMCPClient {
if timeout <= 0 {
timeout = 30 * time.Second
}
return &HTTPMCPClient{
url: url,
timeout: timeout,
client: &http.Client{
Timeout: timeout,
},
logger: logger,
status: "disconnected",
}
}
func (c *HTTPMCPClient) setStatus(status string) {
c.mu.Lock()
defer c.mu.Unlock()
c.status = status
}
func (c *HTTPMCPClient) GetStatus() string {
c.mu.RLock()
defer c.mu.RUnlock()
return c.status
}
func (c *HTTPMCPClient) IsConnected() bool {
return c.GetStatus() == "connected"
}
func (c *HTTPMCPClient) Initialize(ctx context.Context) error {
c.setStatus("connecting")
req := Message{
ID: MessageID{value: "1"},
Method: "initialize",
Version: "2.0",
}
params := InitializeRequest{
ProtocolVersion: ProtocolVersion,
Capabilities: make(map[string]interface{}),
ClientInfo: ClientInfo{
Name: "CyberStrikeAI",
Version: "1.0.0",
},
}
paramsJSON, _ := json.Marshal(params)
req.Params = paramsJSON
_, err := c.sendRequest(ctx, &req)
if err != nil {
c.setStatus("error")
return fmt.Errorf("初始化失败: %w", err)
}
c.setStatus("connected")
return nil
}
func (c *HTTPMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
req := Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/list",
Version: "2.0",
}
req.Params = json.RawMessage("{}")
resp, err := c.sendRequest(ctx, &req)
if err != nil {
return nil, fmt.Errorf("获取工具列表失败: %w", err)
}
var listResp ListToolsResponse
if err := json.Unmarshal(resp.Result, &listResp); err != nil {
return nil, fmt.Errorf("解析工具列表失败: %w", err)
}
return listResp.Tools, nil
}
func (c *HTTPMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
req := Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/call",
Version: "2.0",
}
callReq := CallToolRequest{
Name: name,
Arguments: args,
}
paramsJSON, _ := json.Marshal(callReq)
req.Params = paramsJSON
resp, err := c.sendRequest(ctx, &req)
if err != nil {
return nil, fmt.Errorf("调用工具失败: %w", err)
}
var callResp CallToolResponse
if err := json.Unmarshal(resp.Result, &callResp); err != nil {
return nil, fmt.Errorf("解析工具调用结果失败: %w", err)
}
return &ToolResult{
Content: callResp.Content,
IsError: callResp.IsError,
}, nil
}
func (c *HTTPMCPClient) sendRequest(ctx context.Context, msg *Message) (*Message, error) {
body, err := json.Marshal(msg)
if err != nil {
return nil, fmt.Errorf("序列化请求失败: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("创建HTTP请求失败: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
resp, err := c.client.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("HTTP请求失败: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
bodyBytes, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("HTTP错误 %d: %s", resp.StatusCode, string(bodyBytes))
}
var mcpResp Message
if err := json.NewDecoder(resp.Body).Decode(&mcpResp); err != nil {
return nil, fmt.Errorf("解析响应失败: %w", err)
}
if mcpResp.Error != nil {
return nil, fmt.Errorf("MCP错误: %s (code: %d)", mcpResp.Error.Message, mcpResp.Error.Code)
}
return &mcpResp, nil
}
func (c *HTTPMCPClient) Close() error {
c.setStatus("disconnected")
return nil
}
// StdioMCPClient stdio模式的MCP客户端
type StdioMCPClient struct {
command string
args []string
timeout time.Duration
cmd *exec.Cmd
stdin io.WriteCloser
stdout io.ReadCloser
decoder *json.Decoder
encoder *json.Encoder
logger *zap.Logger
mu sync.RWMutex
status string
requestID int64
responses map[string]chan *Message
responsesMu sync.Mutex
ctx context.Context
cancel context.CancelFunc
}
// NewStdioMCPClient 创建stdio模式的MCP客户端
func NewStdioMCPClient(command string, args []string, timeout time.Duration, logger *zap.Logger) *StdioMCPClient {
if timeout <= 0 {
timeout = 30 * time.Second
}
ctx, cancel := context.WithCancel(context.Background())
return &StdioMCPClient{
command: command,
args: args,
timeout: timeout,
logger: logger,
status: "disconnected",
responses: make(map[string]chan *Message),
ctx: ctx,
cancel: cancel,
}
}
func (c *StdioMCPClient) setStatus(status string) {
c.mu.Lock()
defer c.mu.Unlock()
c.status = status
}
func (c *StdioMCPClient) GetStatus() string {
c.mu.RLock()
defer c.mu.RUnlock()
return c.status
}
func (c *StdioMCPClient) IsConnected() bool {
return c.GetStatus() == "connected"
}
func (c *StdioMCPClient) Initialize(ctx context.Context) error {
c.setStatus("connecting")
if err := c.startProcess(); err != nil {
c.setStatus("error")
return fmt.Errorf("启动进程失败: %w", err)
}
// 启动响应读取goroutine
go c.readResponses()
// 发送初始化请求
req := Message{
ID: MessageID{value: "1"},
Method: "initialize",
Version: "2.0",
}
params := InitializeRequest{
ProtocolVersion: ProtocolVersion,
Capabilities: make(map[string]interface{}),
ClientInfo: ClientInfo{
Name: "CyberStrikeAI",
Version: "1.0.0",
},
}
paramsJSON, _ := json.Marshal(params)
req.Params = paramsJSON
_, err := c.sendRequest(ctx, &req)
if err != nil {
c.setStatus("error")
c.Close()
return fmt.Errorf("初始化失败: %w", err)
}
c.setStatus("connected")
return nil
}
func (c *StdioMCPClient) startProcess() error {
cmd := exec.CommandContext(c.ctx, c.command, c.args...)
stdin, err := cmd.StdinPipe()
if err != nil {
return err
}
stdout, err := cmd.StdoutPipe()
if err != nil {
stdin.Close()
return err
}
if err := cmd.Start(); err != nil {
stdin.Close()
stdout.Close()
return err
}
c.cmd = cmd
c.stdin = stdin
c.stdout = stdout
c.decoder = json.NewDecoder(stdout)
c.encoder = json.NewEncoder(stdin)
return nil
}
func (c *StdioMCPClient) readResponses() {
defer func() {
if r := recover(); r != nil {
c.logger.Error("读取响应时发生panic", zap.Any("error", r))
}
}()
for {
var msg Message
if err := c.decoder.Decode(&msg); err != nil {
if err == io.EOF {
c.setStatus("disconnected")
break
}
c.logger.Error("读取响应失败", zap.Error(err))
break
}
// 处理响应
id := msg.ID.String()
c.responsesMu.Lock()
if ch, ok := c.responses[id]; ok {
select {
case ch <- &msg:
default:
}
delete(c.responses, id)
}
c.responsesMu.Unlock()
}
}
func (c *StdioMCPClient) sendRequest(ctx context.Context, msg *Message) (*Message, error) {
if c.encoder == nil {
return nil, fmt.Errorf("进程未启动")
}
id := msg.ID.String()
if id == "" {
c.mu.Lock()
c.requestID++
id = fmt.Sprintf("%d", c.requestID)
msg.ID = MessageID{value: id}
c.mu.Unlock()
}
// 创建响应通道
responseCh := make(chan *Message, 1)
c.responsesMu.Lock()
c.responses[id] = responseCh
c.responsesMu.Unlock()
// 发送请求
if err := c.encoder.Encode(msg); err != nil {
c.responsesMu.Lock()
delete(c.responses, id)
c.responsesMu.Unlock()
return nil, fmt.Errorf("发送请求失败: %w", err)
}
// 等待响应
select {
case resp := <-responseCh:
if resp.Error != nil {
return nil, fmt.Errorf("MCP错误: %s (code: %d)", resp.Error.Message, resp.Error.Code)
}
return resp, nil
case <-ctx.Done():
c.responsesMu.Lock()
delete(c.responses, id)
c.responsesMu.Unlock()
return nil, ctx.Err()
case <-time.After(c.timeout):
c.responsesMu.Lock()
delete(c.responses, id)
c.responsesMu.Unlock()
return nil, fmt.Errorf("请求超时")
}
}
func (c *StdioMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
req := Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/list",
Version: "2.0",
}
req.Params = json.RawMessage("{}")
resp, err := c.sendRequest(ctx, &req)
if err != nil {
return nil, fmt.Errorf("获取工具列表失败: %w", err)
}
var listResp ListToolsResponse
if err := json.Unmarshal(resp.Result, &listResp); err != nil {
return nil, fmt.Errorf("解析工具列表失败: %w", err)
}
return listResp.Tools, nil
}
func (c *StdioMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
req := Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/call",
Version: "2.0",
}
callReq := CallToolRequest{
Name: name,
Arguments: args,
}
paramsJSON, _ := json.Marshal(callReq)
req.Params = paramsJSON
resp, err := c.sendRequest(ctx, &req)
if err != nil {
return nil, fmt.Errorf("调用工具失败: %w", err)
}
var callResp CallToolResponse
if err := json.Unmarshal(resp.Result, &callResp); err != nil {
return nil, fmt.Errorf("解析工具调用结果失败: %w", err)
}
return &ToolResult{
Content: callResp.Content,
IsError: callResp.IsError,
}, nil
}
func (c *StdioMCPClient) Close() error {
c.cancel()
if c.stdin != nil {
c.stdin.Close()
}
if c.stdout != nil {
c.stdout.Close()
}
if c.cmd != nil {
c.cmd.Process.Kill()
c.cmd.Wait()
}
c.setStatus("disconnected")
return nil
}
internal/mcp/external_manager.go
+660
View File
@@ -0,0 +1,660 @@
package mcp
import (
"context"
"fmt"
"sync"
"time"
"cyberstrike-ai/internal/config"
"github.com/google/uuid"
"go.uber.org/zap"
)
// ExternalMCPManager 外部MCP管理器
type ExternalMCPManager struct {
clients map[string]ExternalMCPClient
configs map[string]config.ExternalMCPServerConfig
logger *zap.Logger
storage MonitorStorage // 可选的持久化存储
executions map[string]*ToolExecution // 执行记录
stats map[string]*ToolStats // 工具统计信息
mu sync.RWMutex
}
// NewExternalMCPManager 创建外部MCP管理器
func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager {
return NewExternalMCPManagerWithStorage(logger, nil)
}
// NewExternalMCPManagerWithStorage 创建外部MCP管理器（带持久化存储）
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
return &ExternalMCPManager{
clients: make(map[string]ExternalMCPClient),
configs: make(map[string]config.ExternalMCPServerConfig),
logger: logger,
storage: storage,
executions: make(map[string]*ToolExecution),
stats: make(map[string]*ToolStats),
}
}
// LoadConfigs 加载配置
func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) {
m.mu.Lock()
defer m.mu.Unlock()
if cfg == nil || cfg.Servers == nil {
return
}
m.configs = make(map[string]config.ExternalMCPServerConfig)
for name, serverCfg := range cfg.Servers {
m.configs[name] = serverCfg
}
}
// GetConfigs 获取所有配置
func (m *ExternalMCPManager) GetConfigs() map[string]config.ExternalMCPServerConfig {
m.mu.RLock()
defer m.mu.RUnlock()
result := make(map[string]config.ExternalMCPServerConfig)
for k, v := range m.configs {
result[k] = v
}
return result
}
// AddOrUpdateConfig 添加或更新配置
func (m *ExternalMCPManager) AddOrUpdateConfig(name string, serverCfg config.ExternalMCPServerConfig) error {
m.mu.Lock()
defer m.mu.Unlock()
// 如果已存在客户端，先关闭
if client, exists := m.clients[name]; exists {
client.Close()
delete(m.clients, name)
}
m.configs[name] = serverCfg
// 如果启用，自动连接
if m.isEnabled(serverCfg) {
go m.connectClient(name, serverCfg)
}
return nil
}
// RemoveConfig 移除配置
func (m *ExternalMCPManager) RemoveConfig(name string) error {
m.mu.Lock()
defer m.mu.Unlock()
// 关闭客户端
if client, exists := m.clients[name]; exists {
client.Close()
delete(m.clients, name)
}
delete(m.configs, name)
return nil
}
// StartClient 启动客户端
func (m *ExternalMCPManager) StartClient(name string) error {
m.mu.Lock()
serverCfg, exists := m.configs[name]
m.mu.Unlock()
if !exists {
return fmt.Errorf("配置不存在: %s", name)
}
// 检查是否已经有连接的客户端
m.mu.RLock()
_, hasClient := m.clients[name]
m.mu.RUnlock()
if hasClient {
// 检查客户端是否已连接
if client, ok := m.GetClient(name); ok && client.IsConnected() {
return fmt.Errorf("客户端已连接")
}
// 如果有客户端但未连接，先关闭
if client, ok := m.GetClient(name); ok {
client.Close()
m.mu.Lock()
delete(m.clients, name)
m.mu.Unlock()
}
}
// 更新配置为启用
m.mu.Lock()
serverCfg.ExternalMCPEnable = true
m.configs[name] = serverCfg
m.mu.Unlock()
// 连接客户端
return m.connectClient(name, serverCfg)
}
// StopClient 停止客户端
func (m *ExternalMCPManager) StopClient(name string) error {
m.mu.Lock()
defer m.mu.Unlock()
serverCfg, exists := m.configs[name]
if !exists {
return fmt.Errorf("配置不存在: %s", name)
}
// 关闭客户端
if client, exists := m.clients[name]; exists {
client.Close()
delete(m.clients, name)
}
// 更新配置为禁用
serverCfg.ExternalMCPEnable = false
m.configs[name] = serverCfg
return nil
}
// GetClient 获取客户端
func (m *ExternalMCPManager) GetClient(name string) (ExternalMCPClient, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
client, exists := m.clients[name]
return client, exists
}
// GetAllTools 获取所有外部MCP的工具
func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) {
m.mu.RLock()
clients := make(map[string]ExternalMCPClient)
for k, v := range m.clients {
clients[k] = v
}
m.mu.RUnlock()
var allTools []Tool
for name, client := range clients {
if !client.IsConnected() {
continue
}
tools, err := client.ListTools(ctx)
if err != nil {
m.logger.Warn("获取外部MCP工具列表失败",
zap.String("name", name),
zap.Error(err),
)
continue
}
// 为工具添加前缀，避免冲突
for _, tool := range tools {
tool.Name = fmt.Sprintf("%s::%s", name, tool.Name)
allTools = append(allTools, tool)
}
}
return allTools, nil
}
// CallTool 调用外部MCP工具（返回执行ID）
func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
// 解析工具名称：name::toolName
var mcpName, actualToolName string
if idx := findSubstring(toolName, "::"); idx > 0 {
mcpName = toolName[:idx]
actualToolName = toolName[idx+2:]
} else {
return nil, "", fmt.Errorf("无效的工具名称格式: %s", toolName)
}
client, exists := m.GetClient(mcpName)
if !exists {
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
}
if !client.IsConnected() {
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s", mcpName)
}
// 创建执行记录
executionID := uuid.New().String()
execution := &ToolExecution{
ID: executionID,
ToolName: toolName, // 使用完整工具名称（包含MCP名称）
Arguments: args,
Status: "running",
StartTime: time.Now(),
}
m.mu.Lock()
m.executions[executionID] = execution
// 如果内存中的执行记录超过限制，清理最旧的记录
m.cleanupOldExecutions()
m.mu.Unlock()
if m.storage != nil {
if err := m.storage.SaveToolExecution(execution); err != nil {
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
// 调用工具
result, err := client.CallTool(ctx, actualToolName, args)
// 更新执行记录
m.mu.Lock()
now := time.Now()
execution.EndTime = &now
execution.Duration = now.Sub(execution.StartTime)
if err != nil {
execution.Status = "failed"
execution.Error = err.Error()
} else if result != nil && result.IsError {
execution.Status = "failed"
if len(result.Content) > 0 {
execution.Error = result.Content[0].Text
} else {
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
} else {
execution.Status = "completed"
if result == nil {
result = &ToolResult{
Content: []Content{
{Type: "text", Text: "工具执行完成，但未返回结果"},
},
}
}
execution.Result = result
}
m.mu.Unlock()
if m.storage != nil {
if err := m.storage.SaveToolExecution(execution); err != nil {
m.logger.Warn("保存执行记录到数据库失败", zap.Error(err))
}
}
// 更新统计信息
failed := err != nil || (result != nil && result.IsError)
m.updateStats(toolName, failed)
// 如果使用存储，从内存中删除（已持久化）
if m.storage != nil {
m.mu.Lock()
delete(m.executions, executionID)
m.mu.Unlock()
}
if err != nil {
return nil, executionID, err
}
return result, executionID, nil
}
// cleanupOldExecutions 清理旧的执行记录（保持内存中的记录数量在限制内）
func (m *ExternalMCPManager) cleanupOldExecutions() {
const maxExecutionsInMemory = 1000
if len(m.executions) <= maxExecutionsInMemory {
return
}
// 按开始时间排序，删除最旧的记录
type execTime struct {
id string
startTime time.Time
}
var execs []execTime
for id, exec := range m.executions {
execs = append(execs, execTime{id: id, startTime: exec.StartTime})
}
// 按时间排序
for i := 0; i < len(execs)-1; i++ {
for j := i + 1; j < len(execs); j++ {
if execs[i].startTime.After(execs[j].startTime) {
execs[i], execs[j] = execs[j], execs[i]
}
}
}
// 删除最旧的记录
toDelete := len(m.executions) - maxExecutionsInMemory
for i := 0; i < toDelete && i < len(execs); i++ {
delete(m.executions, execs[i].id)
}
}
// GetExecution 获取执行记录（先从内存查找，再从数据库查找）
func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) {
m.mu.RLock()
exec, exists := m.executions[id]
m.mu.RUnlock()
if exists {
return exec, true
}
if m.storage != nil {
exec, err := m.storage.GetToolExecution(id)
if err == nil {
return exec, true
}
}
return nil, false
}
// updateStats 更新统计信息
func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
now := time.Now()
if m.storage != nil {
totalCalls := 1
successCalls := 0
failedCalls := 0
if failed {
failedCalls = 1
} else {
successCalls = 1
}
if err := m.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
m.logger.Warn("保存统计信息到数据库失败", zap.Error(err))
}
return
}
m.mu.Lock()
defer m.mu.Unlock()
if m.stats[toolName] == nil {
m.stats[toolName] = &ToolStats{
ToolName: toolName,
}
}
stats := m.stats[toolName]
stats.TotalCalls++
stats.LastCallTime = &now
if failed {
stats.FailedCalls++
} else {
stats.SuccessCalls++
}
}
// GetStats 获取MCP服务器统计信息
func (m *ExternalMCPManager) GetStats() map[string]interface{} {
m.mu.RLock()
defer m.mu.RUnlock()
total := len(m.configs)
enabled := 0
disabled := 0
connected := 0
for name, cfg := range m.configs {
if m.isEnabled(cfg) {
enabled++
if client, exists := m.clients[name]; exists && client.IsConnected() {
connected++
}
} else {
disabled++
}
}
return map[string]interface{}{
"total": total,
"enabled": enabled,
"disabled": disabled,
"connected": connected,
}
}
// GetToolStats 获取工具统计信息（合并内存和数据库）
// 只返回外部MCP工具的统计信息（工具名称包含 "::"）
func (m *ExternalMCPManager) GetToolStats() map[string]*ToolStats {
result := make(map[string]*ToolStats)
// 从数据库加载统计信息（如果使用数据库存储）
if m.storage != nil {
dbStats, err := m.storage.LoadToolStats()
if err == nil {
// 只保留外部MCP工具的统计信息（工具名称包含 "::"）
for k, v := range dbStats {
if findSubstring(k, "::") > 0 {
result[k] = v
}
}
} else {
m.logger.Warn("从数据库加载统计信息失败", zap.Error(err))
}
}
// 合并内存中的统计信息
m.mu.RLock()
for k, v := range m.stats {
// 如果数据库中已有该工具的统计信息，合并它们
if existing, exists := result[k]; exists {
// 创建新的统计信息对象，避免修改共享对象
merged := &ToolStats{
ToolName: k,
TotalCalls: existing.TotalCalls + v.TotalCalls,
SuccessCalls: existing.SuccessCalls + v.SuccessCalls,
FailedCalls: existing.FailedCalls + v.FailedCalls,
}
// 使用最新的调用时间
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
merged.LastCallTime = v.LastCallTime
} else if existing.LastCallTime != nil {
timeCopy := *existing.LastCallTime
merged.LastCallTime = &timeCopy
}
result[k] = merged
} else {
// 如果数据库中没有，直接使用内存中的统计信息
statCopy := *v
result[k] = &statCopy
}
}
m.mu.RUnlock()
return result
}
// GetToolCount 获取指定外部MCP的工具数量
func (m *ExternalMCPManager) GetToolCount(name string) (int, error) {
client, exists := m.GetClient(name)
if !exists {
return 0, fmt.Errorf("客户端不存在: %s", name)
}
if !client.IsConnected() {
return 0, nil
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
tools, err := client.ListTools(ctx)
if err != nil {
return 0, fmt.Errorf("获取工具列表失败: %w", err)
}
return len(tools), nil
}
// GetToolCounts 获取所有外部MCP的工具数量
func (m *ExternalMCPManager) GetToolCounts() map[string]int {
m.mu.RLock()
clients := make(map[string]ExternalMCPClient)
for k, v := range m.clients {
clients[k] = v
}
m.mu.RUnlock()
result := make(map[string]int)
for name, client := range clients {
if !client.IsConnected() {
result[name] = 0
continue
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
tools, err := client.ListTools(ctx)
cancel()
if err != nil {
m.logger.Warn("获取外部MCP工具数量失败",
zap.String("name", name),
zap.Error(err),
)
result[name] = 0
continue
}
result[name] = len(tools)
}
return result
}
// connectClient 连接客户端（异步）
func (m *ExternalMCPManager) connectClient(name string, serverCfg config.ExternalMCPServerConfig) error {
var client ExternalMCPClient
timeout := time.Duration(serverCfg.Timeout) * time.Second
if timeout <= 0 {
timeout = 30 * time.Second
}
// 根据传输模式创建客户端
transport := serverCfg.Transport
if transport == "" {
// 如果没有指定transport，根据是否有command或url判断
if serverCfg.Command != "" {
transport = "stdio"
} else if serverCfg.URL != "" {
transport = "http"
} else {
return fmt.Errorf("无法确定传输模式: 需要指定command或url")
}
}
switch transport {
case "http":
if serverCfg.URL == "" {
return fmt.Errorf("HTTP模式需要URL")
}
client = NewHTTPMCPClient(serverCfg.URL, timeout, m.logger)
case "stdio":
if serverCfg.Command == "" {
return fmt.Errorf("stdio模式需要command")
}
client = NewStdioMCPClient(serverCfg.Command, serverCfg.Args, timeout, m.logger)
default:
return fmt.Errorf("不支持的传输模式: %s", transport)
}
// 初始化连接
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
if err := client.Initialize(ctx); err != nil {
m.logger.Error("初始化外部MCP客户端失败",
zap.String("name", name),
zap.Error(err),
)
return err
}
// 保存客户端
m.mu.Lock()
m.clients[name] = client
m.mu.Unlock()
m.logger.Info("外部MCP客户端已连接",
zap.String("name", name),
zap.String("transport", transport),
)
return nil
}
// isEnabled 检查是否启用
func (m *ExternalMCPManager) isEnabled(cfg config.ExternalMCPServerConfig) bool {
// 优先使用 ExternalMCPEnable 字段
// 如果没有设置，检查旧的 enabled/disabled 字段（向后兼容）
if cfg.ExternalMCPEnable {
return true
}
// 向后兼容：检查旧字段
if cfg.Disabled {
return false
}
if cfg.Enabled {
return true
}
// 都没有设置，默认为启用
return true
}
// findSubstring 查找子字符串（简单实现）
func findSubstring(s, substr string) int {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return i
}
}
return -1
}
// StartAllEnabled 启动所有启用的客户端
func (m *ExternalMCPManager) StartAllEnabled() {
m.mu.RLock()
configs := make(map[string]config.ExternalMCPServerConfig)
for k, v := range m.configs {
configs[k] = v
}
m.mu.RUnlock()
for name, cfg := range configs {
if m.isEnabled(cfg) {
go func(n string, c config.ExternalMCPServerConfig) {
if err := m.connectClient(n, c); err != nil {
m.logger.Error("启动外部MCP客户端失败",
zap.String("name", n),
zap.Error(err),
)
}
}(name, cfg)
}
}
}
// StopAll 停止所有客户端
func (m *ExternalMCPManager) StopAll() {
m.mu.Lock()
defer m.mu.Unlock()
for name, client := range m.clients {
client.Close()
delete(m.clients, name)
}
}
internal/mcp/external_manager_test.go
+261
View File
@@ -0,0 +1,261 @@
package mcp
import (
"context"
"testing"
"time"
"cyberstrike-ai/internal/config"
"go.uber.org/zap"
)
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 测试添加stdio配置
stdioCfg := config.ExternalMCPServerConfig{
Command: "python3",
Args: []string{"/path/to/script.py"},
Transport: "stdio",
Description: "Test stdio MCP",
Timeout: 30,
Enabled: true,
}
err := manager.AddOrUpdateConfig("test-stdio", stdioCfg)
if err != nil {
t.Fatalf("添加stdio配置失败: %v", err)
}
// 测试添加HTTP配置
httpCfg := config.ExternalMCPServerConfig{
Transport: "http",
URL: "http://127.0.0.1:8081/mcp",
Description: "Test HTTP MCP",
Timeout: 30,
Enabled: false,
}
err = manager.AddOrUpdateConfig("test-http", httpCfg)
if err != nil {
t.Fatalf("添加HTTP配置失败: %v", err)
}
// 验证配置已保存
configs := manager.GetConfigs()
if len(configs) != 2 {
t.Fatalf("期望2个配置，实际%d个", len(configs))
}
if configs["test-stdio"].Command != stdioCfg.Command {
t.Errorf("stdio配置命令不匹配")
}
if configs["test-http"].URL != httpCfg.URL {
t.Errorf("HTTP配置URL不匹配")
}
}
func TestExternalMCPManager_RemoveConfig(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
cfg := config.ExternalMCPServerConfig{
Command: "python3",
Transport: "stdio",
Enabled: false,
}
manager.AddOrUpdateConfig("test-remove", cfg)
// 移除配置
err := manager.RemoveConfig("test-remove")
if err != nil {
t.Fatalf("移除配置失败: %v", err)
}
configs := manager.GetConfigs()
if _, exists := configs["test-remove"]; exists {
t.Error("配置应该已被移除")
}
}
func TestExternalMCPManager_GetStats(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 添加多个配置
manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: true,
})
manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
URL: "http://127.0.0.1:8081/mcp",
Enabled: true,
})
manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
Command: "python3",
Enabled: false,
Disabled: true, // 明确设置为禁用
})
stats := manager.GetStats()
if stats["total"].(int) != 3 {
t.Errorf("期望总数3，实际%d", stats["total"])
}
if stats["enabled"].(int) != 2 {
t.Errorf("期望启用数2，实际%d", stats["enabled"])
}
if stats["disabled"].(int) != 1 {
t.Errorf("期望停用数1，实际%d", stats["disabled"])
}
}
func TestExternalMCPManager_LoadConfigs(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
externalMCPConfig := config.ExternalMCPConfig{
Servers: map[string]config.ExternalMCPServerConfig{
"loaded1": {
Command: "python3",
Enabled: true,
},
"loaded2": {
URL: "http://127.0.0.1:8081/mcp",
Enabled: false,
},
},
}
manager.LoadConfigs(&externalMCPConfig)
configs := manager.GetConfigs()
if len(configs) != 2 {
t.Fatalf("期望2个配置，实际%d个", len(configs))
}
if configs["loaded1"].Command != "python3" {
t.Error("配置1加载失败")
}
if configs["loaded2"].URL != "http://127.0.0.1:8081/mcp" {
t.Error("配置2加载失败")
}
}
func TestHTTPMCPClient_Initialize(t *testing.T) {
// 注意：这个测试需要一个真实的HTTP MCP服务器
// 如果没有服务器，这个测试会失败
// 在实际测试中，可以使用mock服务器
logger := zap.NewNop()
client := NewHTTPMCPClient("http://127.0.0.1:8081/mcp", 5*time.Second, logger)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 这个测试可能会失败，如果没有真实的服务器
// 在实际环境中，应该使用mock服务器
err := client.Initialize(ctx)
if err != nil {
t.Logf("初始化失败（可能是没有服务器）: %v", err)
}
status := client.GetStatus()
if status == "" {
t.Error("状态不应该为空")
}
client.Close()
}
func TestStdioMCPClient_Initialize(t *testing.T) {
// 注意：这个测试需要一个真实的stdio MCP服务器
// 如果没有服务器，这个测试会失败
logger := zap.NewNop()
client := NewStdioMCPClient("echo", []string{"test"}, 5*time.Second, logger)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 这个测试可能会失败，因为echo不是MCP服务器
// 在实际环境中，应该使用真实的MCP服务器或mock
err := client.Initialize(ctx)
if err != nil {
t.Logf("初始化失败（echo不是MCP服务器）: %v", err)
}
client.Close()
}
func TestExternalMCPManager_StartStopClient(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 添加一个禁用的配置
cfg := config.ExternalMCPServerConfig{
Command: "python3",
Transport: "stdio",
Enabled: false,
}
manager.AddOrUpdateConfig("test-start-stop", cfg)
// 尝试启动（可能会失败，因为没有真实的服务器）
err := manager.StartClient("test-start-stop")
if err != nil {
t.Logf("启动失败（可能是没有服务器）: %v", err)
}
// 停止
err = manager.StopClient("test-start-stop")
if err != nil {
t.Fatalf("停止失败: %v", err)
}
// 验证配置已更新为禁用
configs := manager.GetConfigs()
if configs["test-start-stop"].Enabled {
t.Error("配置应该已被禁用")
}
}
func TestExternalMCPManager_CallTool(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 测试调用不存在的工具
_, err := manager.CallTool(context.Background(), "nonexistent::tool", map[string]interface{}{})
if err == nil {
t.Error("应该返回错误")
}
// 测试无效的工具名称格式
_, err = manager.CallTool(context.Background(), "invalid-tool-name", map[string]interface{}{})
if err == nil {
t.Error("应该返回错误（无效格式）")
}
}
func TestExternalMCPManager_GetAllTools(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
ctx := context.Background()
tools, err := manager.GetAllTools(ctx)
if err != nil {
t.Fatalf("获取工具列表失败: %v", err)
}
// 如果没有连接的客户端，应该返回空列表
if len(tools) != 0 {
t.Logf("获取到%d个工具", len(tools))
}
}