Add files via upload

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
+203
View File
@@ -845,6 +845,209 @@ curl -X POST http://localhost:8080/api/mcp \
}'
```
## 🔗 External MCP Integration
CyberStrikeAI supports integrating external MCP servers to extend tool capabilities. External MCP tools are automatically registered in the system, and AI can call them just like built-in tools.
### Configuration Methods
#### Method 1: Configure via Web Interface (Recommended)
1. After starting the server, access the web interface
2. Click the "Settings" button in the top-right corner
3. Find the "External MCP" configuration section in the settings
4. Click the "Add External MCP" button
5. Fill in the configuration information:
- **Name**: Unique identifier for the external MCP server (e.g., `hexstrike-ai`)
- **Transport Mode**: Choose `stdio` or `http`
- **Description**: Optional, used to describe the MCP server's functionality
- **Timeout**: Tool execution timeout in seconds, default 300 seconds
- **Enable Status**: Whether to enable this external MCP immediately
6. Fill in the corresponding configuration based on transport mode:
- **stdio Mode**:
- **Command**: Startup command for the MCP server (e.g., `python3`)
- **Args**: Startup arguments array (e.g., `["/path/to/mcp_server.py", "--arg", "value"]`)
- **HTTP Mode**:
- **URL**: HTTP endpoint address of the MCP server (e.g., `http://127.0.0.1:8888`)
7. Click the "Save" button, configuration will be automatically saved to `config.yaml`
8. The system will automatically connect to the external MCP server and load its tools
#### Method 2: Edit Configuration File Directly
Add `external_mcp` configuration in `config.yaml`:
```yaml
# External MCP Configuration
external_mcp:
servers:
# External MCP server name (unique identifier)
hexstrike-ai:
# stdio mode configuration
command: python3
args:
- /path/to/hexstrike_mcp.py
- --server
- 'http://127.0.0.1:8888'
# Or HTTP mode configuration (choose one)
# transport: http
# url: http://127.0.0.1:8888
# Common configuration
description: HexStrike AI v6.0 - Advanced Cybersecurity Automation Platform
timeout: 300 # Timeout in seconds
external_mcp_enable: true # Whether to enable
# Tool-level control (optional)
tool_enabled:
nmap_scan: true
sqlmap_scan: true
# ... other tools
```
### Transport Mode Description
#### stdio Mode
Communicates with external MCP servers via standard input/output (stdio), suitable for locally running MCP servers.
**Configuration Example:**
```yaml
external_mcp:
servers:
my-mcp-server:
command: python3
args:
- /path/to/mcp_server.py
- --config
- /path/to/config.json
description: My Custom MCP Server
timeout: 300
external_mcp_enable: true
```
**Features:**
- ✅ Local process communication, no network port required
- ✅ High security, data doesn't traverse network
- ✅ Suitable for local development and testing
#### HTTP Mode
Communicates with external MCP servers via HTTP requests, suitable for remote MCP servers or scenarios requiring cross-network access.
**Configuration Example:**
```yaml
external_mcp:
servers:
remote-mcp-server:
transport: http
url: http://192.168.1.100:8888
description: Remote MCP Server
timeout: 300
external_mcp_enable: true
```
**Features:**
- ✅ Supports remote access
- ✅ Suitable for distributed deployment
- ✅ Easy to integrate with existing HTTP services
### Tool Naming Convention
External MCP tools in the system use the naming format: `{mcp-server-name}::{tool-name}`
For example:
- External MCP server name: `hexstrike-ai`
- Tool name: `nmap_scan`
- Full tool name in system: `hexstrike-ai::nmap_scan`
AI will automatically recognize and use the full tool name when calling.
### Tool Enable/Disable Control
#### Global Enable/Disable
Control the enable status of the entire external MCP server via the `external_mcp_enable` field:
- `true`: Enabled, system will automatically connect and load tools
- `false`: Disabled, system will not connect to this server
#### Tool-Level Control
Precisely control the enable status of each tool via the `tool_enabled` field:
```yaml
external_mcp:
servers:
hexstrike-ai:
# ... other configuration
tool_enabled:
nmap_scan: true # Enable this tool
sqlmap_scan: false # Disable this tool
nuclei_scan: true # Enable this tool
```
- If a tool is not listed in `tool_enabled`, it is enabled by default
- If a tool is set to `false`, it won't appear in the tool list and AI cannot call it
### Managing External MCP
#### Via Web Interface
1. **View External MCP List**: View all configured external MCP servers in the settings interface
2. **Start/Stop**: Start or stop external MCP server connections at any time
3. **Edit Configuration**: Modify external MCP configuration information
4. **Delete Configuration**: Remove unnecessary external MCP servers
#### Via API
- **Get External MCP List**: `GET /api/external-mcp`
- **Add External MCP**: `POST /api/external-mcp`
- **Update External MCP**: `PUT /api/external-mcp/:name`
- **Delete External MCP**: `DELETE /api/external-mcp/:name`
- **Start External MCP**: `POST /api/external-mcp/:name/start`
- **Stop External MCP**: `POST /api/external-mcp/:name/stop`
### Monitoring and Statistics
External MCP tool execution records and statistics are automatically recorded in the system:
- **Execution Records**: View execution history of all external MCP tools in the "Tool Monitoring" page
- **Execution Statistics**: The execution statistics panel displays call counts and success/failure statistics for external MCP tools
- **Real-time Monitoring**: Real-time viewing of external MCP tool execution status
### Troubleshooting
**Issue: External MCP Cannot Connect**
- ✅ Check if `command` and `args` configuration are correct (stdio mode)
- ✅ Check if `url` configuration is correct and accessible (HTTP mode)
- ✅ Check if the external MCP server is running normally
- ✅ View server logs for detailed error information
- ✅ Check network connection (HTTP mode)
- ✅ Check firewall settings (HTTP mode)
**Issue: External MCP Tools Not Displayed**
- ✅ Confirm `external_mcp_enable: true`
- ✅ Check `tool_enabled` configuration to ensure tools are not disabled
- ✅ Confirm external MCP server has successfully connected
- ✅ View server logs to confirm if tools have been loaded
**Issue: External MCP Tool Execution Failed**
- ✅ Check if external MCP server is running normally
- ✅ View tool execution logs for detailed error information
- ✅ Check if timeout setting is reasonable
- ✅ Confirm external MCP server supports the tool call format
### Best Practices
1. **Naming Convention**: Use meaningful names to identify external MCP servers, avoid conflicts
2. **Timeout Settings**: Set timeout reasonably based on tool execution time
3. **Tool Control**: Use `tool_enabled` to precisely control needed tools, avoid loading too many unnecessary tools
4. **Security Considerations**: HTTP mode is recommended to use intranet addresses or configure appropriate access control
5. **Monitoring Management**: Regularly check external MCP connection status and execution statistics
## 🛠️ Security Tool Support
### Tool Overview
+203
View File
@@ -844,6 +844,209 @@ curl -X POST http://localhost:8080/api/mcp \
}'
```
## 🔗 外部 MCP 接入
CyberStrikeAI 支持接入外部 MCP 服务器,扩展工具能力。外部 MCP 工具会自动注册到系统中,AI 可以像使用内置工具一样调用它们。
### 配置方式
#### 方式一:通过 Web 界面配置(推荐)
1. 启动服务器后,访问 Web 界面
2. 点击右上角的"设置"按钮
3. 在设置界面中找到"外部 MCP"配置部分
4. 点击"添加外部 MCP"按钮
5. 填写配置信息:
- **名称**:外部 MCP 服务器的唯一标识(如:`hexstrike-ai`
- **传输模式**:选择 `stdio` 或 `http`
- **描述**:可选,用于说明该 MCP 服务器的功能
- **超时时间**:工具执行超时时间(秒),默认 300 秒
- **启用状态**:是否立即启用该外部 MCP
6. 根据传输模式填写相应配置:
- **stdio 模式**
- **命令**:MCP 服务器的启动命令(如:`python3`
- **参数**:启动参数数组(如:`["/path/to/mcp_server.py", "--arg", "value"]`
- **HTTP 模式**
- **URL**MCP 服务器的 HTTP 端点地址(如:`http://127.0.0.1:8888`
7. 点击"保存"按钮,配置会自动保存到 `config.yaml`
8. 系统会自动连接外部 MCP 服务器并加载其工具
#### 方式二:直接编辑配置文件
在 `config.yaml` 中添加 `external_mcp` 配置:
```yaml
# 外部MCP配置
external_mcp:
servers:
# 外部MCP服务器名称(唯一标识)
hexstrike-ai:
# stdio 模式配置
command: python3
args:
- /path/to/hexstrike_mcp.py
- --server
- 'http://127.0.0.1:8888'
# 或 HTTP 模式配置(二选一)
# transport: http
# url: http://127.0.0.1:8888
# 通用配置
description: HexStrike AI v6.0 - Advanced Cybersecurity Automation Platform
timeout: 300 # 超时时间(秒)
external_mcp_enable: true # 是否启用
# 工具级别控制(可选)
tool_enabled:
nmap_scan: true
sqlmap_scan: true
# ... 其他工具
```
### 传输模式说明
#### stdio 模式
通过标准输入输出(stdio)与外部 MCP 服务器通信,适合本地运行的 MCP 服务器。
**配置示例:**
```yaml
external_mcp:
servers:
my-mcp-server:
command: python3
args:
- /path/to/mcp_server.py
- --config
- /path/to/config.json
description: My Custom MCP Server
timeout: 300
external_mcp_enable: true
```
**特点:**
- ✅ 本地进程通信,无需网络端口
- ✅ 安全性高,数据不经过网络
- ✅ 适合本地开发和测试
#### HTTP 模式
通过 HTTP 请求与外部 MCP 服务器通信,适合远程 MCP 服务器或需要跨网络访问的场景。
**配置示例:**
```yaml
external_mcp:
servers:
remote-mcp-server:
transport: http
url: http://192.168.1.100:8888
description: Remote MCP Server
timeout: 300
external_mcp_enable: true
```
**特点:**
- ✅ 支持远程访问
- ✅ 适合分布式部署
- ✅ 易于集成现有 HTTP 服务
### 工具命名规则
外部 MCP 工具在系统中的名称格式为:`{mcp-server-name}::{tool-name}`
例如:
- 外部 MCP 服务器名称:`hexstrike-ai`
- 工具名称:`nmap_scan`
- 系统内完整工具名:`hexstrike-ai::nmap_scan`
AI 在调用时会自动识别并使用完整工具名。
### 工具启用控制
#### 全局启用/禁用
通过 `external_mcp_enable` 字段控制整个外部 MCP 服务器的启用状态:
- `true`:启用,系统会自动连接并加载工具
- `false`:禁用,系统不会连接该服务器
#### 工具级别控制
通过 `tool_enabled` 字段精确控制每个工具的启用状态:
```yaml
external_mcp:
servers:
hexstrike-ai:
# ... 其他配置
tool_enabled:
nmap_scan: true # 启用此工具
sqlmap_scan: false # 禁用此工具
nuclei_scan: true # 启用此工具
```
- 如果 `tool_enabled` 中未列出某个工具,默认启用
- 如果某个工具设置为 `false`,该工具不会出现在工具列表中,AI 也无法调用
### 管理外部 MCP
#### 通过 Web 界面管理
1. **查看外部 MCP 列表**:在设置界面中查看所有已配置的外部 MCP 服务器
2. **启动/停止**:可以随时启动或停止外部 MCP 服务器连接
3. **编辑配置**:修改外部 MCP 的配置信息
4. **删除配置**:移除不需要的外部 MCP 服务器
#### 通过 API 管理
- **获取外部 MCP 列表**`GET /api/external-mcp`
- **添加外部 MCP**`POST /api/external-mcp`
- **更新外部 MCP**`PUT /api/external-mcp/:name`
- **删除外部 MCP**`DELETE /api/external-mcp/:name`
- **启动外部 MCP**`POST /api/external-mcp/:name/start`
- **停止外部 MCP**`POST /api/external-mcp/:name/stop`
### 监控和统计
外部 MCP 工具的执行记录和统计信息会自动记录到系统中:
- **执行记录**:在"工具监控"页面可以查看所有外部 MCP 工具的执行历史
- **执行统计**:执行统计面板会显示外部 MCP 工具的调用次数、成功/失败统计
- **实时监控**:可以实时查看外部 MCP 工具的执行状态
### 故障排除
**问题:外部 MCP 无法连接**
- ✅ 检查 `command` 和 `args` 配置是否正确(stdio 模式)
- ✅ 检查 `url` 配置是否正确且可访问(HTTP 模式)
- ✅ 检查外部 MCP 服务器是否正常运行
- ✅ 查看服务器日志获取详细错误信息
- ✅ 检查网络连接(HTTP 模式)
- ✅ 检查防火墙设置(HTTP 模式)
**问题:外部 MCP 工具未显示**
- ✅ 确认 `external_mcp_enable: true`
- ✅ 检查 `tool_enabled` 配置,确保工具未被禁用
- ✅ 确认外部 MCP 服务器已成功连接
- ✅ 查看服务器日志确认工具是否已加载
**问题:外部 MCP 工具执行失败**
- ✅ 检查外部 MCP 服务器是否正常运行
- ✅ 查看工具执行日志获取详细错误信息
- ✅ 检查超时时间设置是否合理
- ✅ 确认外部 MCP 服务器支持的工具调用格式
### 最佳实践
1. **命名规范**:使用有意义的名称标识外部 MCP 服务器,避免冲突
2. **超时设置**:根据工具执行时间合理设置超时时间
3. **工具控制**:使用 `tool_enabled` 精确控制需要的工具,避免加载过多无用工具
4. **安全考虑**:HTTP 模式建议使用内网地址或配置适当的访问控制
5. **监控管理**:定期检查外部 MCP 的连接状态和执行统计
## 🛠️ 安全工具支持
### 工具概览
+41 -29
View File
@@ -1,45 +1,57 @@
package main
import (
"cyberstrike-ai/internal/config"
"fmt"
"os"
"cyberstrike-ai/internal/config"
)
func main() {
cfg, err := config.Load("config.yaml")
if err != nil {
fmt.Printf("❌ 加载配置失败: %v\n", err)
if len(os.Args) < 2 {
fmt.Println("Usage: go run cmd/test-config/main.go <config.yaml>")
os.Exit(1)
}
fmt.Printf("✅ 配置加载成功\n")
fmt.Printf(" 工具目录: %s\n", cfg.Security.ToolsDir)
fmt.Printf(" 工具数量: %d\n", len(cfg.Security.Tools))
configPath := os.Args[1]
cfg, err := config.Load(configPath)
if err != nil {
fmt.Printf("Error loading config: %v\n", err)
os.Exit(1)
}
if len(cfg.Security.Tools) > 0 {
fmt.Printf("\n 已加载的工具:\n")
for _, tool := range cfg.Security.Tools {
status := "❌ 禁用"
if tool.Enabled {
status = "✅ 启用"
}
shortDesc := tool.ShortDescription
if shortDesc == "" {
shortDesc = "(无简短描述,将自动提取)"
}
fmt.Printf(" %s %s\n", status, tool.Name)
fmt.Printf(" 简短描述: %s\n", shortDesc)
if len(tool.Description) > 100 {
fmt.Printf(" 详细描述: %s...\n", tool.Description[:100])
} else {
fmt.Printf(" 详细描述: %s\n", tool.Description)
}
fmt.Printf(" 参数数量: %d\n", len(tool.Parameters))
fmt.Println()
if cfg.ExternalMCP.Servers == nil {
fmt.Println("No external MCP servers configured")
os.Exit(0)
}
fmt.Printf("Found %d external MCP server(s):\n\n", len(cfg.ExternalMCP.Servers))
for name, srv := range cfg.ExternalMCP.Servers {
fmt.Printf("Name: %s\n", name)
fmt.Printf(" Transport: %s\n", getTransport(srv))
fmt.Printf(" Command: %s\n", srv.Command)
if len(srv.Args) > 0 {
fmt.Printf(" Args: %v\n", srv.Args)
}
} else {
fmt.Printf(" ⚠️ 未加载任何工具\n")
fmt.Printf(" URL: %s\n", srv.URL)
fmt.Printf(" Description: %s\n", srv.Description)
fmt.Printf(" Timeout: %d seconds\n", srv.Timeout)
fmt.Printf(" Enabled: %v\n", srv.Enabled)
fmt.Printf(" Disabled: %v\n", srv.Disabled)
fmt.Println()
}
}
func getTransport(srv config.ExternalMCPServerConfig) string {
if srv.Transport != "" {
return srv.Transport
}
if srv.Command != "" {
return "stdio"
}
if srv.URL != "" {
return "http"
}
return "unknown"
}
+145
View File
@@ -0,0 +1,145 @@
package main
import (
"context"
"fmt"
"os"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/logger"
"cyberstrike-ai/internal/mcp"
)
func main() {
if len(os.Args) < 2 {
fmt.Println("Usage: go run cmd/test-external-mcp/main.go <config.yaml>")
os.Exit(1)
}
configPath := os.Args[1]
cfg, err := config.Load(configPath)
if err != nil {
fmt.Printf("Error loading config: %v\n", err)
os.Exit(1)
}
if cfg.ExternalMCP.Servers == nil || len(cfg.ExternalMCP.Servers) == 0 {
fmt.Println("No external MCP servers configured")
os.Exit(0)
}
fmt.Printf("Found %d external MCP server(s)\n\n", len(cfg.ExternalMCP.Servers))
// 创建日志
log := logger.New("info", "stdout")
// 创建外部MCP管理器
manager := mcp.NewExternalMCPManager(log.Logger)
manager.LoadConfigs(&cfg.ExternalMCP)
// 显示配置
fmt.Println("=== 配置信息 ===")
for name, srv := range cfg.ExternalMCP.Servers {
fmt.Printf("\n%s:\n", name)
fmt.Printf(" Transport: %s\n", getTransport(srv))
if srv.Command != "" {
fmt.Printf(" Command: %s\n", srv.Command)
fmt.Printf(" Args: %v\n", srv.Args)
}
if srv.URL != "" {
fmt.Printf(" URL: %s\n", srv.URL)
}
fmt.Printf(" Description: %s\n", srv.Description)
fmt.Printf(" Timeout: %d seconds\n", srv.Timeout)
fmt.Printf(" Enabled: %v\n", srv.Enabled)
fmt.Printf(" Disabled: %v\n", srv.Disabled)
}
// 获取统计信息
fmt.Println("\n=== 统计信息 ===")
stats := manager.GetStats()
fmt.Printf("总数: %d\n", stats["total"])
fmt.Printf("已启用: %d\n", stats["enabled"])
fmt.Printf("已停用: %d\n", stats["disabled"])
fmt.Printf("已连接: %d\n", stats["connected"])
// 测试启动(仅测试启用的)
fmt.Println("\n=== 测试启动 ===")
for name, srv := range cfg.ExternalMCP.Servers {
if srv.Enabled && !srv.Disabled {
fmt.Printf("\n尝试启动 %s...\n", name)
// 注意:实际启动可能会失败,因为需要真实的MCP服务器
err := manager.StartClient(name)
if err != nil {
fmt.Printf(" 启动失败(这是正常的,如果没有真实的MCP服务器): %v\n", err)
} else {
fmt.Printf(" 启动成功\n")
// 获取客户端状态
if client, exists := manager.GetClient(name); exists {
fmt.Printf(" 状态: %s\n", client.GetStatus())
fmt.Printf(" 已连接: %v\n", client.IsConnected())
}
}
}
}
// 等待一下
time.Sleep(2 * time.Second)
// 测试获取工具列表
fmt.Println("\n=== 测试获取工具列表 ===")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
tools, err := manager.GetAllTools(ctx)
if err != nil {
fmt.Printf("获取工具列表失败: %v\n", err)
} else {
fmt.Printf("获取到 %d 个工具\n", len(tools))
for i, tool := range tools {
if i < 5 { // 只显示前5个
fmt.Printf(" - %s: %s\n", tool.Name, tool.Description)
}
}
if len(tools) > 5 {
fmt.Printf(" ... 还有 %d 个工具\n", len(tools)-5)
}
}
// 测试停止
fmt.Println("\n=== 测试停止 ===")
for name := range cfg.ExternalMCP.Servers {
fmt.Printf("\n停止 %s...\n", name)
err := manager.StopClient(name)
if err != nil {
fmt.Printf(" 停止失败: %v\n", err)
} else {
fmt.Printf(" 停止成功\n")
}
}
// 最终统计
fmt.Println("\n=== 最终统计 ===")
stats = manager.GetStats()
fmt.Printf("总数: %d\n", stats["total"])
fmt.Printf("已启用: %d\n", stats["enabled"])
fmt.Printf("已停用: %d\n", stats["disabled"])
fmt.Printf("已连接: %d\n", stats["connected"])
fmt.Println("\n=== 测试完成 ===")
}
func getTransport(srv config.ExternalMCPServerConfig) string {
if srv.Transport != "" {
return srv.Transport
}
if srv.Command != "" {
return "stdio"
}
if srv.URL != "" {
return "http"
}
return "unknown"
}
+4 -1
View File
@@ -20,7 +20,7 @@ log:
# MCP 协议配置
# MCP (Model Context Protocol) 用于工具注册和调用
mcp:
enabled: false # 是否启用 MCP 服务器
enabled: true # 是否启用 MCP 服务器
host: 0.0.0.0 # MCP 服务器监听地址
port: 8081 # MCP 服务器端口
# AI 模型配置(支持 OpenAI 兼容 API)
@@ -46,3 +46,6 @@ security:
tools_dir: tools # 工具配置文件目录(相对于配置文件所在目录)
# 系统会从该目录加载所有 .yaml 格式的工具配置文件
# 推荐方式:在 tools/ 目录下为每个工具创建独立的配置文件
# 外部MCP配置
external_mcp:
+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 {
+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)
+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
}
+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
+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("需要指定commandstdio模式)或urlhttp模式)")
}
}
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 字段,保留它
// 注意:由于 omitemptydisabled: 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"` // 工具数量
}
+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: "需要指定commandstdio模式)或urlhttp模式)",
},
{
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)
}
}
+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)
}
+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
}
+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)
}
}
+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))
}
}
+410
View File
@@ -1525,6 +1525,8 @@ header {
display: flex;
gap: 8px;
align-items: center;
flex-wrap: nowrap;
width: 100%;
}
.tools-actions button {
@@ -1557,6 +1559,8 @@ header {
display: flex;
gap: 4px;
flex: 1;
min-width: 0;
max-width: 300px;
align-items: center;
}
@@ -1639,6 +1643,23 @@ header {
color: var(--text-primary);
font-size: 0.9375rem;
margin-bottom: 4px;
display: flex;
align-items: center;
gap: 8px;
flex-wrap: wrap;
}
.external-tool-badge {
display: inline-flex;
align-items: center;
padding: 2px 8px;
background: rgba(255, 152, 0, 0.12);
border: 1px solid rgba(255, 152, 0, 0.3);
border-radius: 12px;
font-size: 0.75rem;
font-weight: 600;
color: #ff9800;
white-space: nowrap;
}
.tool-item-desc {
@@ -2007,3 +2028,392 @@ header {
background: rgba(220, 53, 69, 0.08);
border-radius: 12px;
}
/* 外部MCP配置样式 */
.external-mcp-controls {
display: flex;
flex-direction: column;
gap: 16px;
}
.external-mcp-actions {
display: flex;
align-items: center;
gap: 12px;
flex-wrap: wrap;
}
.external-mcp-stats {
display: flex;
align-items: center;
gap: 20px;
margin-left: auto;
padding: 10px 20px;
background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
border: 1px solid var(--border-color);
border-radius: 10px;
font-size: 0.875rem;
color: var(--text-secondary);
box-shadow: var(--shadow-sm);
}
.external-mcp-stats span {
display: inline-flex;
align-items: center;
gap: 8px;
padding: 4px 0;
white-space: nowrap;
}
.external-mcp-stats span strong {
color: var(--text-primary);
font-weight: 600;
margin-left: 4px;
}
.external-mcp-stats span:not(:last-child)::after {
content: '';
}
.tools-stats {
display: flex;
align-items: center;
gap: 15px;
margin-left: auto;
padding: 8px 16px;
background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
border: 1px solid var(--border-color);
border-radius: 10px;
font-size: 0.8125rem;
color: var(--text-secondary);
box-shadow: var(--shadow-sm);
flex-wrap: nowrap;
flex-shrink: 0;
white-space: nowrap;
}
.tools-stats span {
display: inline-flex;
align-items: center;
gap: 8px;
padding: 4px 0;
white-space: nowrap;
flex-shrink: 0;
}
.tools-stats span strong {
color: var(--text-primary);
font-weight: 600;
margin-left: 4px;
}
.tools-stats span:not(:last-child)::after {
content: '';
width: 1px;
height: 20px;
background: var(--border-color);
margin-left: 10px;
display: inline-block;
}
.external-mcp-list {
display: flex;
flex-direction: column;
gap: 12px;
}
.external-mcp-items {
display: flex;
flex-direction: column;
gap: 12px;
}
.external-mcp-item {
background: var(--bg-primary);
border: 1px solid var(--border-color);
border-radius: 12px;
padding: 20px;
box-shadow: var(--shadow-sm);
transition: all 0.2s ease;
display: flex;
flex-direction: column;
gap: 16px;
}
.external-mcp-item:hover {
box-shadow: var(--shadow-md);
border-color: var(--accent-color);
transform: translateY(-2px);
}
.external-mcp-item-header {
display: flex;
align-items: center;
justify-content: space-between;
gap: 16px;
flex-wrap: wrap;
}
.external-mcp-item-info {
display: flex;
align-items: center;
gap: 12px;
flex: 1;
min-width: 0;
}
.external-mcp-item-info h4 {
margin: 0;
font-size: 1.125rem;
font-weight: 600;
color: var(--text-primary);
display: flex;
align-items: center;
gap: 8px;
flex-wrap: wrap;
}
.tool-count-badge {
display: inline-flex;
align-items: center;
gap: 4px;
padding: 2px 8px;
background: rgba(0, 102, 255, 0.1);
border: 1px solid rgba(0, 102, 255, 0.3);
border-radius: 12px;
font-size: 0.75rem;
font-weight: 600;
color: var(--accent-color);
margin-left: 8px;
white-space: nowrap;
}
.external-mcp-status {
display: inline-flex;
align-items: center;
gap: 6px;
padding: 4px 12px;
border-radius: 999px;
font-size: 0.8125rem;
font-weight: 600;
white-space: nowrap;
}
.external-mcp-status.status-connected {
background: rgba(40, 167, 69, 0.12);
color: var(--success-color);
border: 1px solid rgba(40, 167, 69, 0.3);
}
.external-mcp-status.status-connected::before {
content: '';
width: 8px;
height: 8px;
border-radius: 50%;
background: var(--success-color);
display: inline-block;
animation: pulse 2s infinite;
}
@keyframes pulse {
0%, 100% {
opacity: 1;
}
50% {
opacity: 0.5;
}
}
.external-mcp-status.status-disconnected {
background: rgba(108, 117, 125, 0.12);
color: var(--text-secondary);
border: 1px solid rgba(108, 117, 125, 0.3);
}
.external-mcp-status.status-connecting {
background: rgba(0, 123, 255, 0.12);
color: var(--accent-color);
border: 1px solid rgba(0, 123, 255, 0.3);
}
.external-mcp-status.status-connecting::before {
content: '';
width: 8px;
height: 8px;
border-radius: 50%;
background: var(--accent-color);
display: inline-block;
animation: pulse 1.5s infinite;
}
.external-mcp-status.status-disabled {
background: rgba(255, 193, 7, 0.12);
color: #b8860b;
border: 1px solid rgba(255, 193, 7, 0.3);
}
.external-mcp-item-actions {
display: flex;
align-items: center;
gap: 8px;
flex-wrap: wrap;
}
.btn-small {
padding: 6px 14px;
font-size: 0.8125rem;
border-radius: 6px;
border: 1px solid var(--border-color);
background: var(--bg-secondary);
color: var(--text-primary);
cursor: pointer;
transition: all 0.2s ease;
font-weight: 500;
white-space: nowrap;
}
.btn-small:hover {
background: var(--bg-tertiary);
border-color: var(--accent-color);
color: var(--accent-color);
transform: translateY(-1px);
box-shadow: var(--shadow-sm);
}
.btn-small.btn-danger {
background: rgba(220, 53, 69, 0.08);
border-color: rgba(220, 53, 69, 0.3);
color: var(--error-color);
}
.btn-small.btn-danger:hover {
background: rgba(220, 53, 69, 0.15);
border-color: var(--error-color);
color: #c82333;
}
.external-mcp-item-details {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(220px, 1fr));
gap: 16px;
padding-top: 16px;
border-top: 1px solid var(--border-color);
}
.external-mcp-item-details > div {
display: flex;
flex-direction: column;
gap: 6px;
padding: 12px;
background: var(--bg-secondary);
border-radius: 8px;
border: 1px solid var(--border-color);
transition: all 0.2s ease;
}
.external-mcp-item-details > div:hover {
background: var(--bg-tertiary);
border-color: var(--accent-color);
transform: translateY(-1px);
box-shadow: var(--shadow-sm);
}
.external-mcp-item-details strong {
font-size: 0.75rem;
font-weight: 600;
color: var(--text-secondary);
text-transform: uppercase;
letter-spacing: 0.5px;
margin-bottom: 2px;
}
.external-mcp-item-details span {
font-size: 0.875rem;
color: var(--text-primary);
word-break: break-word;
line-height: 1.5;
}
.external-mcp-list .empty {
text-align: center;
padding: 48px 24px;
color: var(--text-muted);
font-size: 0.9375rem;
background: var(--bg-secondary);
border: 2px dashed var(--border-color);
border-radius: 12px;
}
.external-mcp-list .error {
text-align: center;
padding: 24px;
color: var(--error-color);
background: rgba(220, 53, 69, 0.08);
border: 1px solid rgba(220, 53, 69, 0.3);
border-radius: 12px;
font-size: 0.875rem;
}
.error-message {
display: block;
padding: 8px 12px;
background: rgba(220, 53, 69, 0.1);
border: 1px solid rgba(220, 53, 69, 0.3);
border-radius: 6px;
color: var(--error-color);
font-size: 0.875rem;
line-height: 1.5;
}
.form-group textarea {
padding: 12px;
border: 1px solid var(--border-color);
border-radius: 6px;
font-size: 0.9375rem;
background: var(--bg-primary);
color: var(--text-primary);
transition: border-color 0.2s;
width: 100%;
font-family: inherit;
resize: vertical;
min-height: 200px;
}
.form-group textarea:focus {
outline: none;
border-color: var(--accent-color);
box-shadow: 0 0 0 3px rgba(0, 102, 255, 0.1);
}
.form-group textarea.error {
border-color: var(--error-color);
box-shadow: 0 0 0 3px rgba(220, 53, 69, 0.1);
}
/* 响应式优化 */
@media (max-width: 768px) {
.external-mcp-actions {
flex-direction: column;
align-items: stretch;
}
.external-mcp-stats {
margin-left: 0;
width: 100%;
justify-content: space-around;
}
.external-mcp-item-header {
flex-direction: column;
align-items: flex-start;
}
.external-mcp-item-actions {
width: 100%;
justify-content: flex-end;
}
.external-mcp-item-details {
grid-template-columns: 1fr;
}
}
+650 -21
View File
@@ -1837,6 +1837,8 @@ function renderToolsList() {
if (!toolsList.contains(listContainer)) {
toolsList.appendChild(listContainer);
}
// 更新统计
updateToolsStats();
return;
}
@@ -1844,10 +1846,19 @@ function renderToolsList() {
const toolItem = document.createElement('div');
toolItem.className = 'tool-item';
toolItem.dataset.toolName = tool.name; // 保存原始工具名称
toolItem.dataset.isExternal = tool.is_external ? 'true' : 'false';
toolItem.dataset.externalMcp = tool.external_mcp || '';
// 外部工具标签
const externalBadge = tool.is_external ? '<span class="external-tool-badge" title="外部MCP工具">外部</span>' : '';
toolItem.innerHTML = `
<input type="checkbox" id="tool-${tool.name}" ${tool.enabled ? 'checked' : ''} />
<input type="checkbox" id="tool-${tool.name}" ${tool.enabled ? 'checked' : ''} ${tool.is_external ? 'data-external="true"' : ''} onchange="updateToolsStats()" />
<div class="tool-item-info">
<div class="tool-item-name">${escapeHtml(tool.name)}</div>
<div class="tool-item-name">
${escapeHtml(tool.name)}
${externalBadge}
</div>
<div class="tool-item-desc">${escapeHtml(tool.description || '无描述')}</div>
</div>
`;
@@ -1857,6 +1868,9 @@ function renderToolsList() {
if (!toolsList.contains(listContainer)) {
toolsList.appendChild(listContainer);
}
// 更新统计
updateToolsStats();
}
// 渲染工具列表分页控件
@@ -1903,6 +1917,7 @@ function selectAllTools() {
document.querySelectorAll('#tools-list input[type="checkbox"]').forEach(checkbox => {
checkbox.checked = true;
});
updateToolsStats();
}
// 全不选工具
@@ -1910,6 +1925,88 @@ function deselectAllTools() {
document.querySelectorAll('#tools-list input[type="checkbox"]').forEach(checkbox => {
checkbox.checked = false;
});
updateToolsStats();
}
// 更新工具统计信息
async function updateToolsStats() {
const statsEl = document.getElementById('tools-stats');
if (!statsEl) return;
// 计算当前页的启用工具数
const currentPageEnabled = Array.from(document.querySelectorAll('#tools-list input[type="checkbox"]:checked')).length;
const currentPageTotal = document.querySelectorAll('#tools-list input[type="checkbox"]').length;
// 计算所有工具的启用数
let totalEnabled = 0;
let totalTools = toolsPagination.total || 0;
try {
// 如果有搜索关键词,只统计搜索结果
if (toolsSearchKeyword) {
totalTools = allTools.length;
totalEnabled = allTools.filter(tool => {
const checkbox = document.getElementById(`tool-${tool.name}`);
return checkbox ? checkbox.checked : tool.enabled;
}).length;
} else {
// 没有搜索时,需要获取所有工具的状态
// 先使用当前已知的工具状态
const toolStateMap = new Map();
// 从当前页的checkbox获取状态
allTools.forEach(tool => {
const checkbox = document.getElementById(`tool-${tool.name}`);
if (checkbox) {
toolStateMap.set(tool.name, checkbox.checked);
} else {
// 如果checkbox不存在(不在当前页),使用工具原始状态
toolStateMap.set(tool.name, tool.enabled);
}
});
// 如果总工具数大于当前页,需要获取所有工具的状态
if (totalTools > allTools.length) {
// 遍历所有页面获取完整状态
let page = 1;
let hasMore = true;
const pageSize = 100; // 使用较大的页面大小以减少请求次数
while (hasMore && page <= 10) { // 限制最多10页,避免无限循环
const url = `/api/config/tools?page=${page}&page_size=${pageSize}`;
const pageResponse = await apiFetch(url);
if (!pageResponse.ok) break;
const pageResult = await pageResponse.json();
pageResult.tools.forEach(tool => {
// 如果工具不在当前页,使用服务器返回的状态
if (!toolStateMap.has(tool.name)) {
toolStateMap.set(tool.name, tool.enabled);
}
});
if (page >= pageResult.total_pages) {
hasMore = false;
} else {
page++;
}
}
}
// 计算启用的工具数
totalEnabled = Array.from(toolStateMap.values()).filter(enabled => enabled).length;
}
} catch (error) {
console.warn('获取工具统计失败,使用当前页数据', error);
// 如果获取失败,使用当前页的数据
totalTools = totalTools || currentPageTotal;
totalEnabled = currentPageEnabled;
}
statsEl.innerHTML = `
<span title="当前页启用的工具数"> 当前页已启用: <strong>${currentPageEnabled}</strong> / ${currentPageTotal}</span>
<span title="所有工具中启用的工具总数">📊 总计已启用: <strong>${totalEnabled}</strong> / ${totalTools}</span>
`;
}
// 过滤工具(已废弃,现在使用服务端搜索)
@@ -1974,39 +2071,75 @@ async function applySettings() {
document.querySelectorAll('#tools-list .tool-item').forEach(item => {
const checkbox = item.querySelector('input[type="checkbox"]');
const toolName = item.dataset.toolName;
const isExternal = item.dataset.isExternal === 'true';
const externalMcp = item.dataset.externalMcp || '';
if (toolName) {
currentPageTools.set(toolName, checkbox.checked);
currentPageTools.set(toolName, {
enabled: checkbox.checked,
is_external: isExternal,
external_mcp: externalMcp
});
}
});
// 获取所有工具列表以获取完整状态
// 获取所有工具列表以获取完整状态(遍历所有页面)
// 注意:无论是否在搜索状态下,都要获取所有工具的状态,以确保完整保存
try {
const allToolsResponse = await apiFetch(`/api/config/tools?page=1&page_size=1000`);
if (allToolsResponse.ok) {
const allToolsResult = await allToolsResponse.json();
// 使用所有工具,但用当前页的修改覆盖
allToolsResult.tools.forEach(tool => {
config.tools.push({
const allToolsMap = new Map();
let page = 1;
let hasMore = true;
const pageSize = 100; // 使用合理的页面大小
// 遍历所有页面获取所有工具(不使用搜索关键词,获取全部工具)
while (hasMore) {
const url = `/api/config/tools?page=${page}&page_size=${pageSize}`;
const pageResponse = await apiFetch(url);
if (!pageResponse.ok) {
throw new Error('获取工具列表失败');
}
const pageResult = await pageResponse.json();
// 将当前页的工具添加到映射中
// 如果工具在当前显示的页面中(匹配搜索且在当前页),使用当前页的修改
// 否则使用服务器返回的状态
pageResult.tools.forEach(tool => {
const currentPageTool = currentPageTools.get(tool.name);
allToolsMap.set(tool.name, {
name: tool.name,
enabled: currentPageTools.has(tool.name) ? currentPageTools.get(tool.name) : tool.enabled
});
});
} else {
// 如果获取失败,只使用当前页的工具
currentPageTools.forEach((enabled, toolName) => {
config.tools.push({
name: toolName,
enabled: enabled
enabled: currentPageTool ? currentPageTool.enabled : tool.enabled,
is_external: currentPageTool ? currentPageTool.is_external : (tool.is_external || false),
external_mcp: currentPageTool ? currentPageTool.external_mcp : (tool.external_mcp || '')
});
});
// 检查是否还有更多页面
if (page >= pageResult.total_pages) {
hasMore = false;
} else {
page++;
}
}
// 将所有工具添加到配置中
allToolsMap.forEach(tool => {
config.tools.push({
name: tool.name,
enabled: tool.enabled,
is_external: tool.is_external,
external_mcp: tool.external_mcp
});
});
} catch (error) {
console.warn('获取所有工具列表失败,仅使用当前页工具状态', error);
// 如果获取失败,只使用当前页的工具
currentPageTools.forEach((enabled, toolName) => {
currentPageTools.forEach((toolData, toolName) => {
config.tools.push({
name: toolName,
enabled: enabled
enabled: toolData.enabled,
is_external: toolData.is_external,
external_mcp: toolData.external_mcp
});
});
}
@@ -2432,3 +2565,499 @@ function formatExecutionDuration(start, end) {
const remainMinutes = minutes % 60;
return remainMinutes > 0 ? `${hours} 小时 ${remainMinutes}` : `${hours} 小时`;
}
// ==================== 外部MCP管理 ====================
let currentEditingMCPName = null;
// 加载外部MCP列表
async function loadExternalMCPs() {
try {
const response = await apiFetch('/api/external-mcp');
if (!response.ok) {
throw new Error('获取外部MCP列表失败');
}
const data = await response.json();
renderExternalMCPList(data.servers || {});
renderExternalMCPStats(data.stats || {});
} catch (error) {
console.error('加载外部MCP列表失败:', error);
const list = document.getElementById('external-mcp-list');
if (list) {
list.innerHTML = `<div class="error">加载失败: ${escapeHtml(error.message)}</div>`;
}
}
}
// 渲染外部MCP列表
function renderExternalMCPList(servers) {
const list = document.getElementById('external-mcp-list');
if (!list) return;
if (Object.keys(servers).length === 0) {
list.innerHTML = '<div class="empty">📋 暂无外部MCP配置<br><span style="font-size: 0.875rem; margin-top: 8px; display: block;">点击"添加外部MCP"按钮开始配置</span></div>';
return;
}
let html = '<div class="external-mcp-items">';
for (const [name, server] of Object.entries(servers)) {
const status = server.status || 'disconnected';
const statusClass = status === 'connected' ? 'status-connected' :
status === 'connecting' ? 'status-connecting' :
status === 'disabled' ? 'status-disabled' : 'status-disconnected';
const statusText = status === 'connected' ? '已连接' :
status === 'connecting' ? '连接中...' :
status === 'disabled' ? '已禁用' : '未连接';
const transport = server.config.transport || (server.config.command ? 'stdio' : 'http');
const transportIcon = transport === 'stdio' ? '⚙️' : '🌐';
html += `
<div class="external-mcp-item">
<div class="external-mcp-item-header">
<div class="external-mcp-item-info">
<h4>${transportIcon} ${escapeHtml(name)}${server.tool_count !== undefined && server.tool_count > 0 ? `<span class="tool-count-badge" title="工具数量">🔧 ${server.tool_count}</span>` : ''}</h4>
<span class="external-mcp-status ${statusClass}">${statusText}</span>
</div>
<div class="external-mcp-item-actions">
${status === 'connected' || status === 'disconnected' ?
`<button class="btn-small" id="btn-toggle-${escapeHtml(name)}" onclick="toggleExternalMCP('${escapeHtml(name)}', '${status}')" title="${status === 'connected' ? '停止连接' : '启动连接'}">
${status === 'connected' ? '⏸ 停止' : '▶ 启动'}
</button>` :
status === 'connecting' ?
`<button class="btn-small" id="btn-toggle-${escapeHtml(name)}" disabled style="opacity: 0.6; cursor: not-allowed;">
连接中...
</button>` : ''}
<button class="btn-small" onclick="editExternalMCP('${escapeHtml(name)}')" title="编辑配置" ${status === 'connecting' ? 'disabled' : ''}> 编辑</button>
<button class="btn-small btn-danger" onclick="deleteExternalMCP('${escapeHtml(name)}')" title="删除配置" ${status === 'connecting' ? 'disabled' : ''}>🗑 删除</button>
</div>
</div>
<div class="external-mcp-item-details">
<div>
<strong>传输模式</strong>
<span>${transportIcon} ${escapeHtml(transport.toUpperCase())}</span>
</div>
${server.tool_count !== undefined && server.tool_count > 0 ? `
<div>
<strong>工具数量</strong>
<span style="font-weight: 600; color: var(--accent-color);">🔧 ${server.tool_count} 个工具</span>
</div>` : server.tool_count === 0 && status === 'connected' ? `
<div>
<strong>工具数量</strong>
<span style="color: var(--text-muted);">暂无工具</span>
</div>` : ''}
${server.config.description ? `
<div>
<strong>描述</strong>
<span>${escapeHtml(server.config.description)}</span>
</div>` : ''}
${server.config.timeout ? `
<div>
<strong>超时时间</strong>
<span>${server.config.timeout} </span>
</div>` : ''}
${transport === 'stdio' && server.config.command ? `
<div>
<strong>命令</strong>
<span style="font-family: monospace; font-size: 0.8125rem;">${escapeHtml(server.config.command)}</span>
</div>` : ''}
${transport === 'http' && server.config.url ? `
<div>
<strong>URL</strong>
<span style="font-family: monospace; font-size: 0.8125rem; word-break: break-all;">${escapeHtml(server.config.url)}</span>
</div>` : ''}
</div>
</div>
`;
}
html += '</div>';
list.innerHTML = html;
}
// 渲染外部MCP统计信息
function renderExternalMCPStats(stats) {
const statsEl = document.getElementById('external-mcp-stats');
if (!statsEl) return;
const total = stats.total || 0;
const enabled = stats.enabled || 0;
const disabled = stats.disabled || 0;
const connected = stats.connected || 0;
statsEl.innerHTML = `
<span title="总配置数">📊 总数: <strong>${total}</strong></span>
<span title="已启用的配置数"> 已启用: <strong>${enabled}</strong></span>
<span title="已停用的配置数"> 已停用: <strong>${disabled}</strong></span>
<span title="当前已连接的配置数">🔗 已连接: <strong>${connected}</strong></span>
`;
}
// 显示添加外部MCP模态框
function showAddExternalMCPModal() {
currentEditingMCPName = null;
document.getElementById('external-mcp-modal-title').textContent = '添加外部MCP';
document.getElementById('external-mcp-json').value = '';
document.getElementById('external-mcp-json-error').style.display = 'none';
document.getElementById('external-mcp-json-error').textContent = '';
document.getElementById('external-mcp-json').classList.remove('error');
document.getElementById('external-mcp-modal').style.display = 'block';
}
// 关闭外部MCP模态框
function closeExternalMCPModal() {
document.getElementById('external-mcp-modal').style.display = 'none';
currentEditingMCPName = null;
}
// 编辑外部MCP
async function editExternalMCP(name) {
try {
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(name)}`);
if (!response.ok) {
throw new Error('获取外部MCP配置失败');
}
const server = await response.json();
currentEditingMCPName = name;
document.getElementById('external-mcp-modal-title').textContent = '编辑外部MCP';
// 将配置转换为对象格式(key为名称)
const config = { ...server.config };
// 移除tool_count、external_mcp_enable等前端字段,但保留enabled/disabled用于向后兼容
delete config.tool_count;
delete config.external_mcp_enable;
// 包装成对象格式:{ "name": { config } }
const configObj = {};
configObj[name] = config;
// 格式化JSON
const jsonStr = JSON.stringify(configObj, null, 2);
document.getElementById('external-mcp-json').value = jsonStr;
document.getElementById('external-mcp-json-error').style.display = 'none';
document.getElementById('external-mcp-json-error').textContent = '';
document.getElementById('external-mcp-json').classList.remove('error');
document.getElementById('external-mcp-modal').style.display = 'block';
} catch (error) {
console.error('编辑外部MCP失败:', error);
alert('编辑失败: ' + error.message);
}
}
// 格式化JSON
function formatExternalMCPJSON() {
const jsonTextarea = document.getElementById('external-mcp-json');
const errorDiv = document.getElementById('external-mcp-json-error');
try {
const jsonStr = jsonTextarea.value.trim();
if (!jsonStr) {
errorDiv.textContent = 'JSON不能为空';
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
const parsed = JSON.parse(jsonStr);
const formatted = JSON.stringify(parsed, null, 2);
jsonTextarea.value = formatted;
errorDiv.style.display = 'none';
jsonTextarea.classList.remove('error');
} catch (error) {
errorDiv.textContent = 'JSON格式错误: ' + error.message;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
}
}
// 加载示例
function loadExternalMCPExample() {
const example = {
"hexstrike-ai": {
command: "python3",
args: [
"/path/to/script.py",
"--server",
"http://example.com"
],
description: "示例描述",
timeout: 300
}
};
document.getElementById('external-mcp-json').value = JSON.stringify(example, null, 2);
document.getElementById('external-mcp-json-error').style.display = 'none';
document.getElementById('external-mcp-json').classList.remove('error');
}
// 保存外部MCP
async function saveExternalMCP() {
const jsonTextarea = document.getElementById('external-mcp-json');
const jsonStr = jsonTextarea.value.trim();
const errorDiv = document.getElementById('external-mcp-json-error');
if (!jsonStr) {
errorDiv.textContent = 'JSON配置不能为空';
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
jsonTextarea.focus();
return;
}
let configObj;
try {
configObj = JSON.parse(jsonStr);
} catch (error) {
errorDiv.textContent = 'JSON格式错误: ' + error.message;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
jsonTextarea.focus();
return;
}
// 验证必须是对象格式
if (typeof configObj !== 'object' || Array.isArray(configObj) || configObj === null) {
errorDiv.textContent = '配置错误: 必须是JSON对象格式,key为配置名称,value为配置内容';
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
// 获取所有配置名称
const names = Object.keys(configObj);
if (names.length === 0) {
errorDiv.textContent = '配置错误: 至少需要一个配置项';
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
// 验证每个配置
for (const name of names) {
if (!name || name.trim() === '') {
errorDiv.textContent = '配置错误: 配置名称不能为空';
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
const config = configObj[name];
if (typeof config !== 'object' || Array.isArray(config) || config === null) {
errorDiv.textContent = `配置错误: "${name}" 的配置必须是对象`;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
// 移除 external_mcp_enable 字段(由按钮控制,但保留 enabled/disabled 用于向后兼容)
delete config.external_mcp_enable;
// 验证配置内容
const transport = config.transport || (config.command ? 'stdio' : config.url ? 'http' : '');
if (!transport) {
errorDiv.textContent = `配置错误: "${name}" 需要指定commandstdio模式)或urlhttp模式)`;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
if (transport === 'stdio' && !config.command) {
errorDiv.textContent = `配置错误: "${name}" stdio模式需要command字段`;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
if (transport === 'http' && !config.url) {
errorDiv.textContent = `配置错误: "${name}" http模式需要url字段`;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
}
// 清除错误提示
errorDiv.style.display = 'none';
jsonTextarea.classList.remove('error');
try {
// 如果是编辑模式,只更新当前编辑的配置
if (currentEditingMCPName) {
if (!configObj[currentEditingMCPName]) {
errorDiv.textContent = `配置错误: 编辑模式下,JSON必须包含配置名称 "${currentEditingMCPName}"`;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
return;
}
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(currentEditingMCPName)}`, {
method: 'PUT',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ config: configObj[currentEditingMCPName] }),
});
if (!response.ok) {
const error = await response.json();
throw new Error(error.error || '保存失败');
}
} else {
// 添加模式:保存所有配置
for (const name of names) {
const config = configObj[name];
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(name)}`, {
method: 'PUT',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ config }),
});
if (!response.ok) {
const error = await response.json();
throw new Error(`保存 "${name}" 失败: ${error.error || '未知错误'}`);
}
}
}
closeExternalMCPModal();
await loadExternalMCPs();
alert('保存成功');
} catch (error) {
console.error('保存外部MCP失败:', error);
errorDiv.textContent = '保存失败: ' + error.message;
errorDiv.style.display = 'block';
jsonTextarea.classList.add('error');
}
}
// 删除外部MCP
async function deleteExternalMCP(name) {
if (!confirm(`确定要删除外部MCP "${name}" 吗?`)) {
return;
}
try {
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(name)}`, {
method: 'DELETE',
});
if (!response.ok) {
const error = await response.json();
throw new Error(error.error || '删除失败');
}
await loadExternalMCPs();
alert('删除成功');
} catch (error) {
console.error('删除外部MCP失败:', error);
alert('删除失败: ' + error.message);
}
}
// 切换外部MCP启停
async function toggleExternalMCP(name, currentStatus) {
const action = currentStatus === 'connected' ? 'stop' : 'start';
const buttonId = `btn-toggle-${name}`;
const button = document.getElementById(buttonId);
// 如果是启动操作,显示加载状态
if (action === 'start' && button) {
button.disabled = true;
button.style.opacity = '0.6';
button.style.cursor = 'not-allowed';
button.innerHTML = '⏳ 连接中...';
}
try {
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(name)}/${action}`, {
method: 'POST',
});
if (!response.ok) {
const error = await response.json();
throw new Error(error.error || '操作失败');
}
const result = await response.json();
// 如果是启动操作,轮询状态直到连接成功或失败
if (action === 'start') {
await pollExternalMCPStatus(name, 30); // 最多轮询30次(约30秒)
} else {
// 停止操作,直接刷新
await loadExternalMCPs();
}
} catch (error) {
console.error('切换外部MCP状态失败:', error);
alert('操作失败: ' + error.message);
// 恢复按钮状态
if (button) {
button.disabled = false;
button.style.opacity = '1';
button.style.cursor = 'pointer';
button.innerHTML = '▶ 启动';
}
// 刷新状态
await loadExternalMCPs();
}
}
// 轮询外部MCP状态
async function pollExternalMCPStatus(name, maxAttempts = 30) {
let attempts = 0;
const pollInterval = 1000; // 1秒轮询一次
while (attempts < maxAttempts) {
await new Promise(resolve => setTimeout(resolve, pollInterval));
try {
const response = await apiFetch(`/api/external-mcp/${encodeURIComponent(name)}`);
if (response.ok) {
const data = await response.json();
const status = data.status || 'disconnected';
// 更新按钮状态
const buttonId = `btn-toggle-${name}`;
const button = document.getElementById(buttonId);
if (status === 'connected') {
// 连接成功,刷新列表
await loadExternalMCPs();
return;
} else if (status === 'error' || status === 'disconnected') {
// 连接失败,刷新列表并显示错误
await loadExternalMCPs();
if (status === 'error') {
alert('连接失败,请检查配置和网络连接');
}
return;
} else if (status === 'connecting') {
// 仍在连接中,继续轮询
attempts++;
continue;
}
}
} catch (error) {
console.error('轮询状态失败:', error);
}
attempts++;
}
// 超时,刷新列表
await loadExternalMCPs();
alert('连接超时,请检查配置和网络连接');
}
// 在打开设置时加载外部MCP列表
const originalOpenSettings = openSettings;
openSettings = async function() {
await originalOpenSettings();
await loadExternalMCPs();
};
+59
View File
@@ -119,11 +119,25 @@
<input type="text" id="tools-search" placeholder="搜索工具..." onkeypress="handleSearchKeyPress(event)" oninput="if(this.value.trim() === '') clearSearch()" />
<button class="btn-search" onclick="searchTools()" title="搜索">🔍</button>
</div>
<div class="tools-stats" id="tools-stats"></div>
</div>
<div id="tools-list" class="tools-list"></div>
</div>
</div>
<!-- 外部MCP配置 -->
<div class="settings-section">
<h3>外部 MCP 配置</h3>
<div class="external-mcp-controls">
<div class="external-mcp-actions">
<button class="btn-primary" onclick="showAddExternalMCPModal()">添加外部MCP</button>
<button class="btn-secondary" onclick="loadExternalMCPs()">刷新</button>
<div class="external-mcp-stats" id="external-mcp-stats"></div>
</div>
<div id="external-mcp-list" class="external-mcp-list"></div>
</div>
</div>
<!-- Agent配置 -->
<div class="settings-section">
<h3>Agent 配置</h3>
@@ -247,6 +261,51 @@
</div>
</div>
<!-- 外部MCP配置模态框 -->
<div id="external-mcp-modal" class="modal">
<div class="modal-content" style="max-width: 900px;">
<div class="modal-header">
<h2 id="external-mcp-modal-title">添加外部MCP</h2>
<span class="modal-close" onclick="closeExternalMCPModal()">&times;</span>
</div>
<div class="modal-body">
<div class="form-group">
<label for="external-mcp-json">配置JSON <span style="color: red;">*</span></label>
<textarea id="external-mcp-json" rows="15" placeholder='{\n "hexstrike-ai": {\n "command": "python3",\n "args": ["/path/to/script.py"],\n "description": "描述",\n "timeout": 300\n }\n}' style="font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace; font-size: 0.875rem; line-height: 1.5;"></textarea>
<div class="password-hint">
<strong>配置格式:</strong>JSON对象,key为配置名称,value为配置内容。状态通过"启动/停止"按钮控制,无需在JSON中配置。<br>
<strong>配置示例:</strong><br>
<strong>stdio模式:</strong><br>
<code style="display: block; margin: 8px 0; padding: 8px; background: var(--bg-secondary); border-radius: 4px; white-space: pre-wrap;">{
"hexstrike-ai": {
"command": "python3",
"args": ["/path/to/script.py", "--server", "http://example.com"],
"description": "描述",
"timeout": 300
}
}</code>
<strong>HTTP模式:</strong><br>
<code style="display: block; margin: 8px 0; padding: 8px; background: var(--bg-secondary); border-radius: 4px; white-space: pre-wrap;">{
"cyberstrike-ai-http": {
"transport": "http",
"url": "http://127.0.0.1:8081/mcp"
}
}</code>
</div>
<div id="external-mcp-json-error" class="error-message" style="display: none; margin-top: 8px; padding: 8px; background: rgba(220, 53, 69, 0.1); border: 1px solid rgba(220, 53, 69, 0.3); border-radius: 4px; color: var(--error-color); font-size: 0.875rem;"></div>
</div>
<div class="form-group">
<button type="button" class="btn-secondary" onclick="formatExternalMCPJSON()" style="margin-top: 8px;">格式化JSON</button>
<button type="button" class="btn-secondary" onclick="loadExternalMCPExample()" style="margin-top: 8px; margin-left: 8px;">加载示例</button>
</div>
</div>
<div class="modal-footer">
<button class="btn-secondary" onclick="closeExternalMCPModal()">取消</button>
<button class="btn-primary" onclick="saveExternalMCP()">保存</button>
</div>
</div>
</div>
<!-- Marked.js for Markdown parsing -->
<script src="https://cdn.jsdelivr.net/npm/marked@11.1.1/marked.min.js"></script>
<!-- DOMPurify for HTML sanitization to prevent XSS -->