mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-05 02:18:46 +02:00
Add files via upload
This commit is contained in:
@@ -15,6 +15,26 @@ import (
|
|||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
// externalToolListCacheTTL 已连接外部 MCP 的工具列表缓存有效期,避免每次 API 请求都打远程 ListTools。
|
||||||
|
externalToolListCacheTTL = 60 * time.Second
|
||||||
|
// externalToolCountRefreshInterval 后台刷新工具数量的间隔(仅刷新缓存过期或缺失的客户端)。
|
||||||
|
externalToolCountRefreshInterval = 60 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// toolListCacheEntry 外部 MCP 工具列表缓存条目
|
||||||
|
type toolListCacheEntry struct {
|
||||||
|
tools []Tool
|
||||||
|
updatedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// listToolsInflight 合并同一 MCP 上并发的 ListTools 请求
|
||||||
|
type listToolsInflight struct {
|
||||||
|
done chan struct{}
|
||||||
|
tools []Tool
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
// ExternalMCPManager 外部MCP管理器
|
// ExternalMCPManager 外部MCP管理器
|
||||||
type ExternalMCPManager struct {
|
type ExternalMCPManager struct {
|
||||||
clients map[string]ExternalMCPClient
|
clients map[string]ExternalMCPClient
|
||||||
@@ -26,8 +46,10 @@ type ExternalMCPManager struct {
|
|||||||
errors map[string]string // 错误信息
|
errors map[string]string // 错误信息
|
||||||
toolCounts map[string]int // 工具数量缓存
|
toolCounts map[string]int // 工具数量缓存
|
||||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||||
toolCache map[string][]Tool // 工具列表缓存:MCP名称 -> 工具列表
|
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||||
|
listToolsMu sync.Mutex
|
||||||
|
listToolsInflight map[string]*listToolsInflight
|
||||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||||
@@ -51,9 +73,10 @@ func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage
|
|||||||
executions: make(map[string]*ToolExecution),
|
executions: make(map[string]*ToolExecution),
|
||||||
stats: make(map[string]*ToolStats),
|
stats: make(map[string]*ToolStats),
|
||||||
errors: make(map[string]string),
|
errors: make(map[string]string),
|
||||||
toolCounts: make(map[string]int),
|
toolCounts: make(map[string]int),
|
||||||
toolCache: make(map[string][]Tool),
|
toolCache: make(map[string]toolListCacheEntry),
|
||||||
stopRefresh: make(chan struct{}),
|
listToolsInflight: make(map[string]*listToolsInflight),
|
||||||
|
stopRefresh: make(chan struct{}),
|
||||||
runningCancels: make(map[string]context.CancelFunc),
|
runningCancels: make(map[string]context.CancelFunc),
|
||||||
abortUserNotes: make(map[string]string),
|
abortUserNotes: make(map[string]string),
|
||||||
}
|
}
|
||||||
@@ -210,15 +233,8 @@ func (m *ExternalMCPManager) StartClient(name string) error {
|
|||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
delete(m.errors, name)
|
delete(m.errors, name)
|
||||||
m.mu.Unlock()
|
m.mu.Unlock()
|
||||||
// 立即刷新工具数量和工具列表缓存
|
// 异步拉取工具列表(singleflight 去重,结果同时写入 toolCache 与 toolCounts)
|
||||||
m.triggerToolCountRefresh()
|
go m.refreshToolCache(name, client)
|
||||||
m.refreshToolCache(name, client)
|
|
||||||
// 2 秒后再刷新一次,覆盖 SSE/Streamable 等需稍等就绪的远端
|
|
||||||
go func() {
|
|
||||||
time.Sleep(2 * time.Second)
|
|
||||||
m.triggerToolCountRefresh()
|
|
||||||
m.refreshToolCache(name, client)
|
|
||||||
}()
|
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
|
|
||||||
@@ -249,6 +265,10 @@ func (m *ExternalMCPManager) StopClient(name string) error {
|
|||||||
m.toolCounts[name] = 0
|
m.toolCounts[name] = 0
|
||||||
m.toolCountsMu.Unlock()
|
m.toolCountsMu.Unlock()
|
||||||
|
|
||||||
|
m.toolCacheMu.Lock()
|
||||||
|
delete(m.toolCache, name)
|
||||||
|
m.toolCacheMu.Unlock()
|
||||||
|
|
||||||
// 更新配置为禁用
|
// 更新配置为禁用
|
||||||
serverCfg.ExternalMCPEnable = false
|
serverCfg.ExternalMCPEnable = false
|
||||||
m.configs[name] = serverCfg
|
m.configs[name] = serverCfg
|
||||||
@@ -335,16 +355,19 @@ func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPCl
|
|||||||
return nil, fmt.Errorf("外部MCP连接失败: %s", name)
|
return nil, fmt.Errorf("外部MCP连接失败: %s", name)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 已连接:尝试获取最新工具列表
|
// 已连接:缓存优先,仅在缺失或过期时打远程 ListTools
|
||||||
if client.IsConnected() {
|
if client.IsConnected() {
|
||||||
tools, err := client.ListTools(ctx)
|
if tools, ok := m.getFreshCachedTools(name); ok {
|
||||||
|
return tools, nil
|
||||||
|
}
|
||||||
|
if tools, ok := m.getAnyCachedTools(name); ok {
|
||||||
|
m.triggerToolListRefresh(name, client)
|
||||||
|
return tools, nil
|
||||||
|
}
|
||||||
|
tools, err := m.listToolsDeduped(ctx, name, client)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// 获取失败,尝试使用缓存
|
|
||||||
return m.getCachedTools(name, "连接正常但获取失败", err)
|
return m.getCachedTools(name, "连接正常但获取失败", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 获取成功,更新缓存
|
|
||||||
m.updateToolCache(name, tools)
|
|
||||||
return tools, nil
|
return tools, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -361,37 +384,126 @@ func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPCl
|
|||||||
return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status)
|
return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status)
|
||||||
}
|
}
|
||||||
|
|
||||||
// getCachedTools 获取缓存的工具列表
|
// getCachedTools 获取缓存的工具列表(含空列表缓存)
|
||||||
func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) {
|
func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) {
|
||||||
m.toolCacheMu.RLock()
|
if tools, ok := m.getAnyCachedTools(name); ok {
|
||||||
cachedTools, hasCache := m.toolCache[name]
|
|
||||||
m.toolCacheMu.RUnlock()
|
|
||||||
|
|
||||||
if hasCache && len(cachedTools) > 0 {
|
|
||||||
m.logger.Debug("使用缓存的工具列表",
|
m.logger.Debug("使用缓存的工具列表",
|
||||||
zap.String("name", name),
|
zap.String("name", name),
|
||||||
zap.String("reason", reason),
|
zap.String("reason", reason),
|
||||||
zap.Int("count", len(cachedTools)),
|
zap.Int("count", len(tools)),
|
||||||
zap.Error(originalErr),
|
zap.Error(originalErr),
|
||||||
)
|
)
|
||||||
return cachedTools, nil
|
return tools, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// 无缓存,返回错误
|
|
||||||
if originalErr != nil {
|
if originalErr != nil {
|
||||||
return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr)
|
return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr)
|
||||||
}
|
}
|
||||||
return nil, fmt.Errorf("外部MCP无缓存工具: %s", name)
|
return nil, fmt.Errorf("外部MCP无缓存工具: %s", name)
|
||||||
}
|
}
|
||||||
|
|
||||||
// updateToolCache 更新工具列表缓存
|
func (m *ExternalMCPManager) isToolCacheFresh(updatedAt time.Time) bool {
|
||||||
func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
return !updatedAt.IsZero() && time.Since(updatedAt) < externalToolListCacheTTL
|
||||||
|
}
|
||||||
|
|
||||||
|
func cloneTools(tools []Tool) []Tool {
|
||||||
|
if len(tools) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
out := make([]Tool, len(tools))
|
||||||
|
copy(out, tools)
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *ExternalMCPManager) getFreshCachedTools(name string) ([]Tool, bool) {
|
||||||
|
m.toolCacheMu.RLock()
|
||||||
|
entry, ok := m.toolCache[name]
|
||||||
|
m.toolCacheMu.RUnlock()
|
||||||
|
if !ok || !m.isToolCacheFresh(entry.updatedAt) {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
return cloneTools(entry.tools), true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *ExternalMCPManager) getAnyCachedTools(name string) ([]Tool, bool) {
|
||||||
|
m.toolCacheMu.RLock()
|
||||||
|
entry, ok := m.toolCache[name]
|
||||||
|
m.toolCacheMu.RUnlock()
|
||||||
|
if !ok {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
return cloneTools(entry.tools), true
|
||||||
|
}
|
||||||
|
|
||||||
|
// listToolsDeduped 对同一 MCP 合并并发 ListTools,并更新 toolCache / toolCounts。
|
||||||
|
func (m *ExternalMCPManager) listToolsDeduped(ctx context.Context, name string, client ExternalMCPClient) ([]Tool, error) {
|
||||||
|
m.listToolsMu.Lock()
|
||||||
|
if inflight, exists := m.listToolsInflight[name]; exists {
|
||||||
|
m.listToolsMu.Unlock()
|
||||||
|
select {
|
||||||
|
case <-inflight.done:
|
||||||
|
if inflight.err != nil {
|
||||||
|
return nil, inflight.err
|
||||||
|
}
|
||||||
|
return cloneTools(inflight.tools), nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
inflight := &listToolsInflight{done: make(chan struct{})}
|
||||||
|
m.listToolsInflight[name] = inflight
|
||||||
|
m.listToolsMu.Unlock()
|
||||||
|
|
||||||
|
inflight.tools, inflight.err = client.ListTools(ctx)
|
||||||
|
if inflight.err == nil {
|
||||||
|
m.updateToolCache(name, inflight.tools)
|
||||||
|
}
|
||||||
|
|
||||||
|
m.listToolsMu.Lock()
|
||||||
|
delete(m.listToolsInflight, name)
|
||||||
|
close(inflight.done)
|
||||||
|
m.listToolsMu.Unlock()
|
||||||
|
|
||||||
|
if inflight.err != nil {
|
||||||
|
return nil, inflight.err
|
||||||
|
}
|
||||||
|
return cloneTools(inflight.tools), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// InvalidateToolCache 清除指定外部 MCP 的工具列表缓存(手动刷新时使用)
|
||||||
|
func (m *ExternalMCPManager) InvalidateToolCache(name string) {
|
||||||
m.toolCacheMu.Lock()
|
m.toolCacheMu.Lock()
|
||||||
m.toolCache[name] = tools
|
delete(m.toolCache, name)
|
||||||
|
m.toolCacheMu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// InvalidateAllToolCaches 清除所有外部 MCP 工具列表缓存
|
||||||
|
func (m *ExternalMCPManager) InvalidateAllToolCaches() {
|
||||||
|
m.toolCacheMu.Lock()
|
||||||
|
m.toolCache = make(map[string]toolListCacheEntry)
|
||||||
|
m.toolCacheMu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *ExternalMCPManager) triggerToolListRefresh(name string, client ExternalMCPClient) {
|
||||||
|
go func() {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
_, _ = m.listToolsDeduped(ctx, name, client)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// updateToolCache 更新工具列表缓存与工具数量
|
||||||
|
func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
||||||
|
stored := cloneTools(tools)
|
||||||
|
m.toolCacheMu.Lock()
|
||||||
|
m.toolCache[name] = toolListCacheEntry{tools: stored, updatedAt: time.Now()}
|
||||||
m.toolCacheMu.Unlock()
|
m.toolCacheMu.Unlock()
|
||||||
|
|
||||||
// 如果返回空列表,记录警告
|
m.toolCountsMu.Lock()
|
||||||
if len(tools) == 0 {
|
m.toolCounts[name] = len(stored)
|
||||||
|
m.toolCountsMu.Unlock()
|
||||||
|
|
||||||
|
if len(stored) == 0 {
|
||||||
m.logger.Warn("外部MCP返回空工具列表",
|
m.logger.Warn("外部MCP返回空工具列表",
|
||||||
zap.String("name", name),
|
zap.String("name", name),
|
||||||
zap.String("hint", "服务可能暂时不可用,工具列表为空"),
|
zap.String("hint", "服务可能暂时不可用,工具列表为空"),
|
||||||
@@ -399,7 +511,7 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
|||||||
} else {
|
} else {
|
||||||
m.logger.Debug("工具列表缓存已更新",
|
m.logger.Debug("工具列表缓存已更新",
|
||||||
zap.String("name", name),
|
zap.String("name", name),
|
||||||
zap.Int("count", len(tools)),
|
zap.Int("count", len(stored)),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -854,15 +966,21 @@ func (m *ExternalMCPManager) refreshToolCounts() {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// 使用合理的超时时间(15秒),既能应对网络延迟,又不会过长阻塞
|
// 缓存仍新鲜时直接复用,避免与 GetAllTools 重复打远程
|
||||||
// 由于这是后台异步刷新,超时不会影响前端响应
|
if _, fresh := m.getFreshCachedTools(n); fresh {
|
||||||
|
m.toolCountsMu.RLock()
|
||||||
|
count := m.toolCounts[n]
|
||||||
|
m.toolCountsMu.RUnlock()
|
||||||
|
resultChan <- countResult{name: n, count: count}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||||
tools, err := c.ListTools(ctx)
|
tools, err := m.listToolsDeduped(ctx, n, c)
|
||||||
cancel()
|
cancel()
|
||||||
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
errStr := err.Error()
|
errStr := err.Error()
|
||||||
// SSE 连接 EOF:远端可能关闭了流或未按规范在流上推送响应,仅首次用 Warn 提示
|
|
||||||
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "client is closing") {
|
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "client is closing") {
|
||||||
m.logger.Warn("获取外部MCP工具数量失败(SSE 流已关闭或服务端未在流上返回 tools/list 响应)",
|
m.logger.Warn("获取外部MCP工具数量失败(SSE 流已关闭或服务端未在流上返回 tools/list 响应)",
|
||||||
zap.String("name", n),
|
zap.String("name", n),
|
||||||
@@ -875,7 +993,7 @@ func (m *ExternalMCPManager) refreshToolCounts() {
|
|||||||
zap.Error(err),
|
zap.Error(err),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
resultChan <- countResult{name: n, count: -1} // -1 表示使用旧值
|
resultChan <- countResult{name: n, count: -1}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -925,33 +1043,21 @@ func (m *ExternalMCPManager) refreshToolCache(name string, client ExternalMCPCli
|
|||||||
if !client.IsConnected() {
|
if !client.IsConnected() {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if client.GetStatus() == "error" {
|
||||||
// 检查状态,如果是error状态,不更新缓存
|
|
||||||
status := client.GetStatus()
|
|
||||||
if status == "error" {
|
|
||||||
m.logger.Debug("跳过刷新工具列表缓存(连接失败)",
|
m.logger.Debug("跳过刷新工具列表缓存(连接失败)",
|
||||||
zap.String("name", name),
|
zap.String("name", name),
|
||||||
zap.String("status", status),
|
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// 使用较短的超时时间(5秒)
|
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
if _, err := m.listToolsDeduped(ctx, name, client); err != nil {
|
||||||
tools, err := client.ListTools(ctx)
|
|
||||||
if err != nil {
|
|
||||||
m.logger.Debug("刷新工具列表缓存失败",
|
m.logger.Debug("刷新工具列表缓存失败",
|
||||||
zap.String("name", name),
|
zap.String("name", name),
|
||||||
zap.Error(err),
|
zap.Error(err),
|
||||||
)
|
)
|
||||||
// 刷新失败时不更新缓存,保留旧缓存(如果有)
|
|
||||||
return
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// 使用统一的缓存更新方法
|
|
||||||
m.updateToolCache(name, tools)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// startToolCountRefresh 启动后台刷新工具数量的goroutine
|
// startToolCountRefresh 启动后台刷新工具数量的goroutine
|
||||||
@@ -959,7 +1065,7 @@ func (m *ExternalMCPManager) startToolCountRefresh() {
|
|||||||
m.refreshWg.Add(1)
|
m.refreshWg.Add(1)
|
||||||
go func() {
|
go func() {
|
||||||
defer m.refreshWg.Done()
|
defer m.refreshWg.Done()
|
||||||
ticker := time.NewTicker(10 * time.Second) // 每10秒刷新一次
|
ticker := time.NewTicker(externalToolCountRefreshInterval)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
|
|
||||||
// 立即执行一次刷新
|
// 立即执行一次刷新
|
||||||
@@ -1168,7 +1274,7 @@ func (m *ExternalMCPManager) StopAll() {
|
|||||||
|
|
||||||
// 清理所有工具列表缓存
|
// 清理所有工具列表缓存
|
||||||
m.toolCacheMu.Lock()
|
m.toolCacheMu.Lock()
|
||||||
m.toolCache = make(map[string][]Tool)
|
m.toolCache = make(map[string]toolListCacheEntry)
|
||||||
m.toolCacheMu.Unlock()
|
m.toolCacheMu.Unlock()
|
||||||
|
|
||||||
// 停止后台刷新(使用 select 避免重复关闭 channel)
|
// 停止后台刷新(使用 select 避免重复关闭 channel)
|
||||||
|
|||||||
Reference in New Issue
Block a user