mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 04:16:06 +00:00
feat(mcp): add dynamic tool search and shared connection pool
MCP tool search mode: when an agent has >30 MCP tools, tools are deferred and a single mcp_tool_search meta-tool is registered instead. Uses BM25 index for keyword search, activating matched tools on demand. Shared connection pool: introduces Pool that maintains one physical connection per MCP server name, shared across all agents via refcount. Eliminates N×M duplicate connections (N agents × M servers).
This commit is contained in:
1 parent
af85b0c354
commit
01e40a92b0
13 files changed
+865
-75
No files matched your search
+5
-1
@@ -574,7 +574,11 @@ func runGateway() {
|
||||
}
|
||||
}
|
||||
|
||||
contextFileInterceptor, delegateMgr = wireExtras(pgStores, agentRouter, providerRegistry, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, dynamicLoader)
|
||||
var mcpPool *mcpbridge.Pool
|
||||
contextFileInterceptor, delegateMgr, mcpPool = wireExtras(pgStores, agentRouter, providerRegistry, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, dynamicLoader)
|
||||
if mcpPool != nil {
|
||||
defer mcpPool.Stop()
|
||||
}
|
||||
gatewayAddr := loopbackAddr(cfg.Gateway.Host, cfg.Gateway.Port)
|
||||
var mcpToolLister httpapi.MCPToolLister
|
||||
if mcpMgr != nil {
|
||||
|
||||
+11
-3
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/config"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/hooks"
|
||||
httpapi "github.com/nextlevelbuilder/goclaw/internal/http"
|
||||
mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/sandbox"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/skills"
|
||||
@@ -43,7 +44,7 @@ func wireExtras(
|
||||
appCfg *config.Config,
|
||||
sandboxMgr sandbox.Manager,
|
||||
dynamicLoader *tools.DynamicToolLoader,
|
||||
) (*tools.ContextFileInterceptor, *tools.DelegateManager) {
|
||||
) (*tools.ContextFileInterceptor, *tools.DelegateManager, *mcpbridge.Pool) {
|
||||
// 1. Context file interceptor (created before resolver so callbacks can reference it)
|
||||
var contextFileInterceptor *tools.ContextFileInterceptor
|
||||
var delegateMgr *tools.DelegateManager
|
||||
@@ -82,7 +83,13 @@ func wireExtras(
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Set up agent resolver: lazy-creates Loops from DB
|
||||
// 5. Shared MCP connection pool (eliminates duplicate connections across agents)
|
||||
var mcpPool *mcpbridge.Pool
|
||||
if stores.MCP != nil {
|
||||
mcpPool = mcpbridge.NewPool()
|
||||
}
|
||||
|
||||
// 6. Set up agent resolver: lazy-creates Loops from DB
|
||||
var skillAccessStore store.SkillAccessStore
|
||||
if sas, ok := stores.Skills.(store.SkillAccessStore); ok {
|
||||
skillAccessStore = sas
|
||||
@@ -114,6 +121,7 @@ func wireExtras(
|
||||
TeamStore: stores.Teams,
|
||||
BuiltinToolStore: stores.BuiltinTools,
|
||||
MCPStore: stores.MCP,
|
||||
MCPPool: mcpPool,
|
||||
GroupWriterCache: groupWriterCache,
|
||||
OnEvent: func(event agent.AgentEvent) {
|
||||
msgBus.Broadcast(bus.Event{
|
||||
@@ -470,7 +478,7 @@ func wireExtras(
|
||||
}
|
||||
|
||||
slog.Info("resolver + interceptors + cache subscribers wired")
|
||||
return contextFileInterceptor, delegateMgr
|
||||
return contextFileInterceptor, delegateMgr, mcpPool
|
||||
}
|
||||
|
||||
// wireHTTP creates HTTP handlers (agents + skills + traces + MCP + custom tools + channel instances + providers + delegations + builtin tools).
|
||||
|
||||
@@ -29,6 +29,7 @@ func (l *Loop) buildMessages(ctx context.Context, history []providers.Message, s
|
||||
|
||||
_, hasSpawn := l.tools.Get("spawn")
|
||||
_, hasSkillSearch := l.tools.Get("skill_search")
|
||||
_, hasMCPToolSearch := l.tools.Get("mcp_tool_search")
|
||||
|
||||
// Per-user workspace: show the user's subdirectory in the system prompt.
|
||||
// Uses cached workspace from user_agent_profiles (includes channel isolation).
|
||||
@@ -77,6 +78,7 @@ func (l *Loop) buildMessages(ctx context.Context, history []providers.Message, s
|
||||
HasMemory: l.hasMemory,
|
||||
HasSpawn: l.tools != nil && hasSpawn,
|
||||
HasSkillSearch: hasSkillSearch,
|
||||
HasMCPToolSearch: hasMCPToolSearch,
|
||||
ContextFiles: contextFiles,
|
||||
AgentType: l.agentType,
|
||||
ExtraPrompt: extraSystemPrompt,
|
||||
|
||||
@@ -64,6 +64,9 @@ type ResolverDeps struct {
|
||||
// MCP server store — for per-agent MCP tool loading
|
||||
MCPStore store.MCPServerStore
|
||||
|
||||
// Shared MCP connection pool — eliminates duplicate connections across agents
|
||||
MCPPool *mcpbridge.Pool
|
||||
|
||||
// Skill access store — for per-agent skill visibility filtering
|
||||
SkillAccessStore store.SkillAccessStore
|
||||
|
||||
@@ -270,9 +273,21 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc {
|
||||
if toolsReg == deps.Tools {
|
||||
toolsReg = deps.Tools.Clone()
|
||||
}
|
||||
mcpMgr := mcpbridge.NewManager(toolsReg, mcpbridge.WithStore(deps.MCPStore))
|
||||
var mcpOpts []mcpbridge.ManagerOption
|
||||
mcpOpts = append(mcpOpts, mcpbridge.WithStore(deps.MCPStore))
|
||||
if deps.MCPPool != nil {
|
||||
mcpOpts = append(mcpOpts, mcpbridge.WithPool(deps.MCPPool))
|
||||
}
|
||||
mcpMgr := mcpbridge.NewManager(toolsReg, mcpOpts...)
|
||||
if err := mcpMgr.LoadForAgent(ctx, ag.ID, ""); err != nil {
|
||||
slog.Warn("failed to load MCP servers for agent", "agent", agentKey, "error", err)
|
||||
} else if mcpMgr.IsSearchMode() {
|
||||
// Search mode: too many tools — register mcp_tool_search meta-tool
|
||||
searchTool := mcpbridge.NewMCPToolSearchTool(mcpMgr)
|
||||
toolsReg.Register(searchTool)
|
||||
hasMCPTools = true
|
||||
slog.Info("mcp.agent.search_mode", "agent", agentKey,
|
||||
"deferred_tools", len(mcpMgr.DeferredToolInfos()))
|
||||
} else {
|
||||
toolNames := mcpMgr.ToolNames()
|
||||
if len(toolNames) > 0 {
|
||||
|
||||
@@ -35,7 +35,8 @@ type SystemPromptConfig struct {
|
||||
ExtraPrompt string // extra system prompt (subagent context, etc.)
|
||||
AgentType string // "open" or "predefined" — affects context file framing
|
||||
|
||||
HasSkillSearch bool // skill_search tool registered? (for search-mode prompt)
|
||||
HasSkillSearch bool // skill_search tool registered? (for search-mode prompt)
|
||||
HasMCPToolSearch bool // mcp_tool_search tool registered? (MCP search mode)
|
||||
|
||||
// Sandbox info — matching TS sandboxInfo in system-prompt.ts
|
||||
SandboxEnabled bool // exec tool runs inside Docker sandbox?
|
||||
@@ -56,7 +57,8 @@ var coreToolSummaries = map[string]string{
|
||||
"web_search": "Search the web",
|
||||
"web_fetch": "Fetch and extract content from a URL",
|
||||
"cron": "Manage scheduled jobs and reminders",
|
||||
"skill_search": "Search available skills by keyword (weather, translate, github, etc.)",
|
||||
"skill_search": "Search available skills by keyword (weather, translate, github, etc.)",
|
||||
"mcp_tool_search": "Search for available MCP external integration tools by keyword",
|
||||
"browser": "Browse web pages interactively",
|
||||
"tts": "Convert text to speech audio",
|
||||
"edit": "Edit a file by replacing exact text matches",
|
||||
@@ -108,6 +110,11 @@ func BuildSystemPrompt(cfg SystemPromptConfig) string {
|
||||
lines = append(lines, buildSkillsSection(cfg.SkillsSummary, cfg.HasSkillSearch)...)
|
||||
}
|
||||
|
||||
// 4.5. ## MCP Tools (full only, search mode)
|
||||
if !isMinimal && cfg.HasMCPToolSearch {
|
||||
lines = append(lines, buildMCPToolsSection()...)
|
||||
}
|
||||
|
||||
// 5. ## Memory Recall (full only)
|
||||
if !isMinimal && cfg.HasMemory {
|
||||
lines = append(lines, buildMemoryRecallSection()...)
|
||||
|
||||
@@ -9,6 +9,21 @@ import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/bootstrap"
|
||||
)
|
||||
|
||||
// buildMCPToolsSection generates the MCP tools instruction block for search mode.
|
||||
// Shown when mcp_tool_search is registered instead of individual MCP tools.
|
||||
func buildMCPToolsSection() []string {
|
||||
return []string{
|
||||
"## MCP Tools (mandatory)",
|
||||
"",
|
||||
"You have access to external tool integrations (MCP servers) with many specialized tools.",
|
||||
"Not all tools are loaded by default — use `mcp_tool_search` to discover them.",
|
||||
"1. Before performing external operations (database, API, file management, messaging), run `mcp_tool_search` with descriptive English keywords.",
|
||||
"2. Matching tools are activated immediately and can be called right away in the same turn.",
|
||||
"3. If no match found, proceed with other available tools.",
|
||||
"",
|
||||
}
|
||||
}
|
||||
|
||||
// buildSandboxSection creates the "## Sandbox" section matching TS system-prompt.ts lines 476-519.
|
||||
func buildSandboxSection(cfg SystemPromptConfig) []string {
|
||||
lines := []string{
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
// MCPToolSearchResult is a single result from a BM25 search over MCP tools.
|
||||
type MCPToolSearchResult struct {
|
||||
RegisteredName string `json:"name"`
|
||||
OriginalName string `json:"original_name"`
|
||||
ServerName string `json:"server"`
|
||||
Description string `json:"description"`
|
||||
Score float64 `json:"-"`
|
||||
}
|
||||
|
||||
// toolDoc is an internal document for BM25 scoring.
|
||||
type toolDoc struct {
|
||||
registeredName string
|
||||
originalName string
|
||||
serverName string
|
||||
description string
|
||||
tokens []string
|
||||
}
|
||||
|
||||
// mcpBM25Index is a minimal BM25 index for MCP tool search.
|
||||
type mcpBM25Index struct {
|
||||
docs []toolDoc
|
||||
df map[string]int
|
||||
avgDL float64
|
||||
k1 float64
|
||||
b float64
|
||||
}
|
||||
|
||||
func newMCPBM25Index() *mcpBM25Index {
|
||||
return &mcpBM25Index{
|
||||
df: make(map[string]int),
|
||||
k1: 1.2,
|
||||
b: 0.75,
|
||||
}
|
||||
}
|
||||
|
||||
// build indexes a list of BridgeTools for BM25 search.
|
||||
func (idx *mcpBM25Index) build(tools []*BridgeTool) {
|
||||
idx.docs = make([]toolDoc, 0, len(tools))
|
||||
idx.df = make(map[string]int)
|
||||
|
||||
totalTokens := 0
|
||||
|
||||
for _, bt := range tools {
|
||||
// Include server name + original tool name + description for broad matching
|
||||
searchText := bt.serverName + " " + bt.toolName + " " + bt.description
|
||||
tokens := tokenizeMCP(searchText)
|
||||
|
||||
idx.docs = append(idx.docs, toolDoc{
|
||||
registeredName: bt.registeredName,
|
||||
originalName: bt.toolName,
|
||||
serverName: bt.serverName,
|
||||
description: bt.description,
|
||||
tokens: tokens,
|
||||
})
|
||||
|
||||
seen := make(map[string]bool)
|
||||
for _, t := range tokens {
|
||||
if !seen[t] {
|
||||
idx.df[t]++
|
||||
seen[t] = true
|
||||
}
|
||||
}
|
||||
|
||||
totalTokens += len(tokens)
|
||||
}
|
||||
|
||||
if len(idx.docs) > 0 {
|
||||
idx.avgDL = float64(totalTokens) / float64(len(idx.docs))
|
||||
}
|
||||
}
|
||||
|
||||
// search performs a BM25 search over indexed MCP tools.
|
||||
func (idx *mcpBM25Index) search(query string, maxResults int) []MCPToolSearchResult {
|
||||
if maxResults <= 0 {
|
||||
maxResults = 5
|
||||
}
|
||||
|
||||
queryTokens := tokenizeMCP(query)
|
||||
if len(queryTokens) == 0 || len(idx.docs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
N := float64(len(idx.docs))
|
||||
|
||||
type scored struct {
|
||||
doc toolDoc
|
||||
score float64
|
||||
}
|
||||
|
||||
var results []scored
|
||||
|
||||
for _, doc := range idx.docs {
|
||||
score := 0.0
|
||||
dl := float64(len(doc.tokens))
|
||||
|
||||
tf := make(map[string]int)
|
||||
for _, t := range doc.tokens {
|
||||
tf[t]++
|
||||
}
|
||||
|
||||
for _, qt := range queryTokens {
|
||||
termFreq := float64(tf[qt])
|
||||
if termFreq == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
dfTerm := float64(idx.df[qt])
|
||||
idf := math.Log((N-dfTerm+0.5)/(dfTerm+0.5) + 1)
|
||||
|
||||
numerator := termFreq * (idx.k1 + 1)
|
||||
denominator := termFreq + idx.k1*(1-idx.b+idx.b*dl/idx.avgDL)
|
||||
score += idf * numerator / denominator
|
||||
}
|
||||
|
||||
if score > 0 {
|
||||
results = append(results, scored{doc: doc, score: score})
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by score descending (insertion sort, small N)
|
||||
for i := 1; i < len(results); i++ {
|
||||
key := results[i]
|
||||
j := i - 1
|
||||
for j >= 0 && results[j].score < key.score {
|
||||
results[j+1] = results[j]
|
||||
j--
|
||||
}
|
||||
results[j+1] = key
|
||||
}
|
||||
|
||||
if len(results) > maxResults {
|
||||
results = results[:maxResults]
|
||||
}
|
||||
|
||||
out := make([]MCPToolSearchResult, len(results))
|
||||
for i, r := range results {
|
||||
out[i] = MCPToolSearchResult{
|
||||
RegisteredName: r.doc.registeredName,
|
||||
OriginalName: r.doc.originalName,
|
||||
ServerName: r.doc.serverName,
|
||||
Description: r.doc.description,
|
||||
Score: r.score,
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// docCount returns the number of indexed documents.
|
||||
func (idx *mcpBM25Index) docCount() int { return len(idx.docs) }
|
||||
|
||||
// tokenizeMCP splits text into lowercase tokens, removing punctuation.
|
||||
// Mirrors skills.tokenize() (unexported, so duplicated here).
|
||||
func tokenizeMCP(text string) []string {
|
||||
lower := strings.ToLower(text)
|
||||
|
||||
cleaned := strings.Map(func(r rune) rune {
|
||||
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
||||
return r
|
||||
}
|
||||
return ' '
|
||||
}, lower)
|
||||
|
||||
fields := strings.Fields(cleaned)
|
||||
|
||||
var tokens []string
|
||||
for _, f := range fields {
|
||||
if len(f) > 1 {
|
||||
tokens = append(tokens, f)
|
||||
}
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
+187
-16
@@ -21,6 +21,11 @@ const (
|
||||
initialBackoff = 2 * time.Second
|
||||
maxBackoff = 60 * time.Second
|
||||
maxReconnectAttempts = 10
|
||||
|
||||
// mcpToolInlineMaxCount is the threshold above which MCP tools switch
|
||||
// to search mode (deferred loading via mcp_tool_search) instead of
|
||||
// being registered inline in the tool registry.
|
||||
mcpToolInlineMaxCount = 30
|
||||
)
|
||||
|
||||
// ServerStatus reports the connection status of an MCP server.
|
||||
@@ -51,6 +56,11 @@ type serverState struct {
|
||||
// Supports two sources:
|
||||
// - Config-based: reads from config.MCPServerConfig map (shared across all agents)
|
||||
// - DB-backed: queries MCPServerStore per agent+user for permission-filtered servers
|
||||
//
|
||||
// When total MCP tool count exceeds mcpToolInlineMaxCount, the manager
|
||||
// enters "search mode": tools are kept in deferredTools instead of the
|
||||
// registry, and only mcp_tool_search is registered. Tools are activated
|
||||
// on demand via ActivateTools().
|
||||
type Manager struct {
|
||||
mu sync.RWMutex
|
||||
servers map[string]*serverState
|
||||
@@ -61,6 +71,16 @@ type Manager struct {
|
||||
|
||||
// DB-backed servers
|
||||
store store.MCPServerStore
|
||||
|
||||
// Shared connection pool (nil = standalone mode)
|
||||
pool *Pool
|
||||
poolServers map[string]struct{} // server names acquired from pool (for cleanup)
|
||||
poolToolNames map[string][]string // per-agent tool names for pool-backed servers
|
||||
|
||||
// Search mode: deferred tools not registered in registry
|
||||
deferredTools map[string]*BridgeTool // registeredName → BridgeTool
|
||||
activatedTools map[string]struct{} // tracks activated tool names for group:mcp
|
||||
searchMode bool
|
||||
}
|
||||
|
||||
// ManagerOption configures the Manager.
|
||||
@@ -80,6 +100,14 @@ func WithStore(s store.MCPServerStore) ManagerOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithPool sets a shared connection pool for MCP servers.
|
||||
// When set, LoadForAgent uses the pool instead of creating per-agent connections.
|
||||
func WithPool(p *Pool) ManagerOption {
|
||||
return func(m *Manager) {
|
||||
m.pool = p
|
||||
}
|
||||
}
|
||||
|
||||
// NewManager creates a new MCP Manager.
|
||||
func NewManager(registry *tools.Registry, opts ...ManagerOption) *Manager {
|
||||
m := &Manager{
|
||||
@@ -139,12 +167,25 @@ func (m *Manager) LoadForAgent(ctx context.Context, agentID uuid.UUID, userID st
|
||||
continue
|
||||
}
|
||||
|
||||
if err := m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
|
||||
jsonBytesToStringSlice(srv.Args), jsonBytesToStringMap(srv.Env),
|
||||
srv.URL, jsonBytesToStringMap(srv.Headers),
|
||||
srv.ToolPrefix, srv.TimeoutSec); err != nil {
|
||||
slog.Warn("mcp.server.connect_failed", "server", srv.Name, "error", err)
|
||||
continue
|
||||
args := jsonBytesToStringSlice(srv.Args)
|
||||
env := jsonBytesToStringMap(srv.Env)
|
||||
headers := jsonBytesToStringMap(srv.Headers)
|
||||
|
||||
if m.pool != nil {
|
||||
// Pool mode: acquire shared connection, create per-agent BridgeTools
|
||||
if err := m.connectViaPool(ctx, srv.Name, srv.Transport, srv.Command,
|
||||
args, env, srv.URL, headers, srv.ToolPrefix, srv.TimeoutSec); err != nil {
|
||||
slog.Warn("mcp.server.connect_failed", "server", srv.Name, "error", err)
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
// Standalone mode: create per-agent connection
|
||||
if err := m.connectServer(ctx, srv.Name, srv.Transport, srv.Command,
|
||||
args, env, srv.URL, headers,
|
||||
srv.ToolPrefix, srv.TimeoutSec); err != nil {
|
||||
slog.Warn("mcp.server.connect_failed", "server", srv.Name, "error", err)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Apply tool filtering from grants
|
||||
@@ -153,29 +194,159 @@ func (m *Manager) LoadForAgent(ctx context.Context, agentID uuid.UUID, userID st
|
||||
}
|
||||
}
|
||||
|
||||
// Check if we should enter search mode (too many tools to inline)
|
||||
m.maybeEnterSearchMode()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// maybeEnterSearchMode moves all registered BridgeTools to deferredTools
|
||||
// if total count exceeds the inline threshold.
|
||||
func (m *Manager) maybeEnterSearchMode() {
|
||||
allNames := m.ToolNames()
|
||||
if len(allNames) <= mcpToolInlineMaxCount {
|
||||
return
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.deferredTools = make(map[string]*BridgeTool, len(allNames))
|
||||
m.activatedTools = make(map[string]struct{})
|
||||
|
||||
// Move all tools to deferred — handle both pool-backed and standalone
|
||||
for serverName := range m.servers {
|
||||
var toolNames []string
|
||||
if _, isPool := m.poolServers[serverName]; isPool {
|
||||
toolNames = m.poolToolNames[serverName]
|
||||
} else {
|
||||
toolNames = m.servers[serverName].toolNames
|
||||
}
|
||||
|
||||
for _, name := range toolNames {
|
||||
if bt, ok := m.registry.Get(name); ok {
|
||||
if bridge, ok := bt.(*BridgeTool); ok {
|
||||
m.deferredTools[name] = bridge
|
||||
m.registry.Unregister(name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clear tool names
|
||||
if _, isPool := m.poolServers[serverName]; isPool {
|
||||
m.poolToolNames[serverName] = nil
|
||||
} else {
|
||||
m.servers[serverName].toolNames = nil
|
||||
}
|
||||
}
|
||||
|
||||
tools.UnregisterToolGroup("mcp")
|
||||
m.searchMode = true
|
||||
|
||||
slog.Info("mcp.search_mode.enabled",
|
||||
"deferred_tools", len(m.deferredTools),
|
||||
"threshold", mcpToolInlineMaxCount)
|
||||
}
|
||||
|
||||
// IsSearchMode reports whether the manager is in deferred/search mode.
|
||||
func (m *Manager) IsSearchMode() bool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.searchMode
|
||||
}
|
||||
|
||||
// DeferredToolInfos returns all deferred tools for BM25 indexing.
|
||||
func (m *Manager) DeferredToolInfos() []*BridgeTool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
result := make([]*BridgeTool, 0, len(m.deferredTools))
|
||||
for _, bt := range m.deferredTools {
|
||||
result = append(result, bt)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// ActivateTools moves named deferred tools into the registry so
|
||||
// they become available on the next agent loop iteration.
|
||||
// Uses 3-phase locking to avoid deadlock with registry.mu.
|
||||
func (m *Manager) ActivateTools(names []string) {
|
||||
// Phase 1: collect tools to activate (read lock)
|
||||
m.mu.RLock()
|
||||
toActivate := make([]*BridgeTool, 0, len(names))
|
||||
for _, name := range names {
|
||||
if bt, ok := m.deferredTools[name]; ok {
|
||||
if _, exists := m.registry.Get(name); !exists {
|
||||
toActivate = append(toActivate, bt)
|
||||
}
|
||||
}
|
||||
}
|
||||
m.mu.RUnlock()
|
||||
|
||||
if len(toActivate) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Phase 2: register in registry (no Manager lock held)
|
||||
var activated []string
|
||||
for _, bt := range toActivate {
|
||||
if _, exists := m.registry.Get(bt.Name()); !exists {
|
||||
m.registry.Register(bt)
|
||||
activated = append(activated, bt.Name())
|
||||
}
|
||||
}
|
||||
|
||||
if len(activated) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Phase 3: update internal state (write lock)
|
||||
m.mu.Lock()
|
||||
for _, name := range activated {
|
||||
delete(m.deferredTools, name)
|
||||
m.activatedTools[name] = struct{}{}
|
||||
}
|
||||
activeNames := make([]string, 0, len(m.activatedTools))
|
||||
for n := range m.activatedTools {
|
||||
activeNames = append(activeNames, n)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
tools.RegisterToolGroup("mcp", activeNames)
|
||||
slog.Info("mcp.tools.activated", "tools", activated)
|
||||
}
|
||||
|
||||
// Stop shuts down all MCP server connections and unregisters tools.
|
||||
func (m *Manager) Stop() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for name, ss := range m.servers {
|
||||
if ss.cancel != nil {
|
||||
ss.cancel()
|
||||
}
|
||||
if ss.client != nil {
|
||||
if err := ss.client.Close(); err != nil {
|
||||
slog.Debug("mcp.server.close_error", "server", name, "error", err)
|
||||
if _, isPool := m.poolServers[name]; isPool {
|
||||
// Pool-backed: unregister per-agent tools, release shared connection
|
||||
for _, toolName := range m.poolToolNames[name] {
|
||||
m.registry.Unregister(toolName)
|
||||
}
|
||||
if m.pool != nil {
|
||||
m.pool.Release(name)
|
||||
}
|
||||
} else {
|
||||
// Standalone: close connection directly
|
||||
if ss.cancel != nil {
|
||||
ss.cancel()
|
||||
}
|
||||
if ss.client != nil {
|
||||
if err := ss.client.Close(); err != nil {
|
||||
slog.Debug("mcp.server.close_error", "server", name, "error", err)
|
||||
}
|
||||
}
|
||||
for _, toolName := range ss.toolNames {
|
||||
m.registry.Unregister(toolName)
|
||||
}
|
||||
}
|
||||
// Unregister tools
|
||||
for _, toolName := range ss.toolNames {
|
||||
m.registry.Unregister(toolName)
|
||||
}
|
||||
}
|
||||
m.servers = make(map[string]*serverState)
|
||||
m.poolServers = nil
|
||||
m.poolToolNames = nil
|
||||
}
|
||||
|
||||
// ServerStatus returns the status of all connected MCP servers.
|
||||
|
||||
+117
-35
@@ -13,22 +13,23 @@ import (
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
)
|
||||
|
||||
// connectServer creates a client, initializes the connection, discovers tools, and registers them.
|
||||
func (m *Manager) connectServer(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int) error {
|
||||
// connectAndDiscover creates a client, initializes the MCP handshake, and
|
||||
// discovers tools. Returns a connected serverState with discovered tool
|
||||
// definitions. The caller is responsible for registering tools and starting
|
||||
// the health loop. This function is shared by both Manager and Pool.
|
||||
func connectAndDiscover(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, timeoutSec int) (*serverState, []mcpgo.Tool, error) {
|
||||
client, err := createClient(transportType, command, args, env, url, headers)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create client: %w", err)
|
||||
return nil, nil, fmt.Errorf("create client: %w", err)
|
||||
}
|
||||
|
||||
// Start transport (SSE/streamable-http need explicit Start; stdio auto-starts)
|
||||
if transportType != "stdio" {
|
||||
if err := client.Start(ctx); err != nil {
|
||||
_ = client.Close()
|
||||
return fmt.Errorf("start transport: %w", err)
|
||||
return nil, nil, fmt.Errorf("start transport: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize MCP handshake
|
||||
initReq := mcpgo.InitializeRequest{}
|
||||
initReq.Params.ProtocolVersion = mcpgo.LATEST_PROTOCOL_VERSION
|
||||
initReq.Params.ClientInfo = mcpgo.Implementation{
|
||||
@@ -38,14 +39,13 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
|
||||
|
||||
if _, err := client.Initialize(ctx, initReq); err != nil {
|
||||
_ = client.Close()
|
||||
return fmt.Errorf("initialize: %w", err)
|
||||
return nil, nil, fmt.Errorf("initialize: %w", err)
|
||||
}
|
||||
|
||||
// Discover tools
|
||||
toolsResult, err := client.ListTools(ctx, mcpgo.ListToolsRequest{})
|
||||
if err != nil {
|
||||
_ = client.Close()
|
||||
return fmt.Errorf("list tools: %w", err)
|
||||
return nil, nil, fmt.Errorf("list tools: %w", err)
|
||||
}
|
||||
|
||||
if timeoutSec <= 0 {
|
||||
@@ -60,40 +60,29 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
|
||||
}
|
||||
ss.connected.Store(true)
|
||||
|
||||
// Register tools
|
||||
var registeredNames []string
|
||||
for _, mcpTool := range toolsResult.Tools {
|
||||
bt := NewBridgeTool(name, mcpTool, client, toolPrefix, timeoutSec, &ss.connected)
|
||||
return ss, toolsResult.Tools, nil
|
||||
}
|
||||
|
||||
// Check for name collision with existing tools
|
||||
if _, exists := m.registry.Get(bt.Name()); exists {
|
||||
slog.Warn("mcp.tool.name_collision",
|
||||
"server", name,
|
||||
"tool", bt.Name(),
|
||||
"action", "skipped",
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
m.registry.Register(bt)
|
||||
registeredNames = append(registeredNames, bt.Name())
|
||||
// connectServer creates a client, initializes the connection, discovers tools, and registers them.
|
||||
func (m *Manager) connectServer(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int) error {
|
||||
ss, mcpTools, err := connectAndDiscover(ctx, name, transportType, command, args, env, url, headers, timeoutSec)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Register tools
|
||||
registeredNames := m.registerBridgeTools(ss, mcpTools, name, toolPrefix, timeoutSec)
|
||||
ss.toolNames = registeredNames
|
||||
|
||||
// Create health monitoring context BEFORE storing server in map,
|
||||
// so DisconnectServer() can safely call ss.cancel even if called immediately.
|
||||
// Create health monitoring context
|
||||
hctx, hcancel := context.WithCancel(context.Background())
|
||||
ss.cancel = hcancel
|
||||
|
||||
// Store server state BEFORE updating MCP group, because updateMCPGroup()
|
||||
// calls ToolNames() which iterates m.servers. If we store after, the current
|
||||
// server's tools are invisible to ToolNames() and toolGroups["mcp"] ends up
|
||||
// incomplete — causing the policy engine to block MCP tools.
|
||||
// Store server state BEFORE updating MCP group
|
||||
m.mu.Lock()
|
||||
m.servers[name] = ss
|
||||
m.mu.Unlock()
|
||||
|
||||
// Register dynamic tool groups for policy filtering
|
||||
if len(registeredNames) > 0 {
|
||||
tools.RegisterToolGroup("mcp:"+name, registeredNames)
|
||||
m.updateMCPGroup()
|
||||
@@ -110,6 +99,89 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman
|
||||
return nil
|
||||
}
|
||||
|
||||
// registerBridgeTools creates BridgeTools from MCP tool definitions and
|
||||
// registers them in the Manager's registry. Returns registered tool names.
|
||||
func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int) []string {
|
||||
var registeredNames []string
|
||||
for _, mcpTool := range mcpTools {
|
||||
bt := NewBridgeTool(serverName, mcpTool, ss.client, toolPrefix, timeoutSec, &ss.connected)
|
||||
|
||||
if _, exists := m.registry.Get(bt.Name()); exists {
|
||||
slog.Warn("mcp.tool.name_collision",
|
||||
"server", serverName,
|
||||
"tool", bt.Name(),
|
||||
"action", "skipped",
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
m.registry.Register(bt)
|
||||
registeredNames = append(registeredNames, bt.Name())
|
||||
}
|
||||
return registeredNames
|
||||
}
|
||||
|
||||
// connectViaPool acquires a shared connection from the pool and creates
|
||||
// per-agent BridgeTools pointing to the shared client/connected pointers.
|
||||
func (m *Manager) connectViaPool(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, toolPrefix string, timeoutSec int) error {
|
||||
entry, err := m.pool.Acquire(ctx, name, transportType, command, args, env, url, headers, timeoutSec)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create per-agent BridgeTools from the pool's shared connection
|
||||
registeredNames := m.registerPoolBridgeTools(entry, name, toolPrefix, timeoutSec)
|
||||
|
||||
// Track server state and per-agent tool names
|
||||
m.mu.Lock()
|
||||
m.servers[name] = entry.state
|
||||
if m.poolServers == nil {
|
||||
m.poolServers = make(map[string]struct{})
|
||||
}
|
||||
m.poolServers[name] = struct{}{}
|
||||
if m.poolToolNames == nil {
|
||||
m.poolToolNames = make(map[string][]string)
|
||||
}
|
||||
m.poolToolNames[name] = registeredNames
|
||||
m.mu.Unlock()
|
||||
|
||||
if len(registeredNames) > 0 {
|
||||
tools.RegisterToolGroup("mcp:"+name, registeredNames)
|
||||
m.updateMCPGroup()
|
||||
}
|
||||
|
||||
slog.Info("mcp.server.connected_via_pool",
|
||||
"server", name,
|
||||
"transport", transportType,
|
||||
"tools", len(registeredNames),
|
||||
)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// registerPoolBridgeTools creates BridgeTools from pool entry's discovered tools,
|
||||
// pointing to the shared client/connected pointers. Returns registered tool names.
|
||||
func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int) []string {
|
||||
var registeredNames []string
|
||||
for _, mcpTool := range entry.tools {
|
||||
bt := NewBridgeTool(serverName, mcpTool, entry.state.client, toolPrefix, timeoutSec, &entry.state.connected)
|
||||
|
||||
if _, exists := m.registry.Get(bt.Name()); exists {
|
||||
slog.Warn("mcp.tool.name_collision",
|
||||
"server", serverName,
|
||||
"tool", bt.Name(),
|
||||
"action", "skipped",
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
m.registry.Register(bt)
|
||||
registeredNames = append(registeredNames, bt.Name())
|
||||
}
|
||||
|
||||
return registeredNames
|
||||
}
|
||||
|
||||
// createClient creates the appropriate MCP client based on transport type.
|
||||
func createClient(transportType, command string, args []string, env map[string]string, url string, headers map[string]string) (*mcpclient.Client, error) {
|
||||
switch transportType {
|
||||
@@ -136,9 +208,20 @@ func createClient(transportType, command string, args []string, env map[string]s
|
||||
}
|
||||
}
|
||||
|
||||
// newHealthTicker creates a ticker for health check intervals.
|
||||
func newHealthTicker() *time.Ticker {
|
||||
return time.NewTicker(healthCheckInterval)
|
||||
}
|
||||
|
||||
// isMethodNotFound returns true if the error indicates the server
|
||||
// doesn't implement the "ping" method (still considered healthy).
|
||||
func isMethodNotFound(err error) bool {
|
||||
return strings.Contains(strings.ToLower(err.Error()), "method not found")
|
||||
}
|
||||
|
||||
// healthLoop periodically pings the MCP server and attempts reconnection on failure.
|
||||
func (m *Manager) healthLoop(ctx context.Context, ss *serverState) {
|
||||
ticker := time.NewTicker(healthCheckInterval)
|
||||
ticker := newHealthTicker()
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
@@ -147,8 +230,7 @@ func (m *Manager) healthLoop(ctx context.Context, ss *serverState) {
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := ss.client.Ping(ctx); err != nil {
|
||||
// Servers that don't implement "ping" are still alive — treat as healthy.
|
||||
if strings.Contains(strings.ToLower(err.Error()), "method not found") {
|
||||
if isMethodNotFound(err) {
|
||||
ss.connected.Store(true)
|
||||
ss.mu.Lock()
|
||||
ss.reconnAttempts = 0
|
||||
|
||||
@@ -16,8 +16,12 @@ func (m *Manager) ToolNames() []string {
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
var names []string
|
||||
for _, ss := range m.servers {
|
||||
names = append(names, ss.toolNames...)
|
||||
for name, ss := range m.servers {
|
||||
if _, isPool := m.poolServers[name]; isPool {
|
||||
names = append(names, m.poolToolNames[name]...)
|
||||
} else {
|
||||
names = append(names, ss.toolNames...)
|
||||
}
|
||||
}
|
||||
return names
|
||||
}
|
||||
@@ -27,6 +31,9 @@ func (m *Manager) ServerToolNames(serverName string) []string {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
if _, isPool := m.poolServers[serverName]; isPool {
|
||||
return append([]string(nil), m.poolToolNames[serverName]...)
|
||||
}
|
||||
if ss, ok := m.servers[serverName]; ok {
|
||||
return append([]string(nil), ss.toolNames...)
|
||||
}
|
||||
@@ -49,20 +56,45 @@ func (m *Manager) unregisterAllTools() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for name, ss := range m.servers {
|
||||
if ss.cancel != nil {
|
||||
ss.cancel()
|
||||
}
|
||||
if ss.client != nil {
|
||||
_ = ss.client.Close()
|
||||
}
|
||||
for _, toolName := range ss.toolNames {
|
||||
m.registry.Unregister(toolName)
|
||||
for name := range m.servers {
|
||||
if _, isPool := m.poolServers[name]; isPool {
|
||||
// Pool-backed: unregister per-agent tools, release shared connection
|
||||
for _, toolName := range m.poolToolNames[name] {
|
||||
m.registry.Unregister(toolName)
|
||||
}
|
||||
if m.pool != nil {
|
||||
m.pool.Release(name)
|
||||
}
|
||||
} else {
|
||||
// Standalone: close connection directly
|
||||
ss := m.servers[name]
|
||||
if ss.cancel != nil {
|
||||
ss.cancel()
|
||||
}
|
||||
if ss.client != nil {
|
||||
_ = ss.client.Close()
|
||||
}
|
||||
for _, toolName := range ss.toolNames {
|
||||
m.registry.Unregister(toolName)
|
||||
}
|
||||
}
|
||||
tools.UnregisterToolGroup("mcp:" + name)
|
||||
slog.Debug("mcp.server.unregistered", "server", name, "tools", len(ss.toolNames))
|
||||
slog.Debug("mcp.server.unregistered", "server", name)
|
||||
}
|
||||
|
||||
// Clean up search mode state: unregister activated tools and clear deferred
|
||||
if m.searchMode {
|
||||
for name := range m.activatedTools {
|
||||
m.registry.Unregister(name)
|
||||
}
|
||||
m.deferredTools = nil
|
||||
m.activatedTools = nil
|
||||
m.searchMode = false
|
||||
}
|
||||
|
||||
m.servers = make(map[string]*serverState)
|
||||
m.poolServers = nil
|
||||
m.poolToolNames = nil
|
||||
tools.UnregisterToolGroup("mcp")
|
||||
}
|
||||
|
||||
@@ -114,8 +146,14 @@ func (m *Manager) filterTools(serverName string, allow, deny []string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
ss, ok := m.servers[serverName]
|
||||
if !ok {
|
||||
// Get the tool names list (pool-backed or standalone)
|
||||
var toolNames []string
|
||||
_, isPool := m.poolServers[serverName]
|
||||
if isPool {
|
||||
toolNames = m.poolToolNames[serverName]
|
||||
} else if ss, ok := m.servers[serverName]; ok {
|
||||
toolNames = ss.toolNames
|
||||
} else {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -123,7 +161,7 @@ func (m *Manager) filterTools(serverName string, allow, deny []string) {
|
||||
denySet := toSet(deny)
|
||||
|
||||
var kept []string
|
||||
for _, toolName := range ss.toolNames {
|
||||
for _, toolName := range toolNames {
|
||||
bt, ok := m.registry.Get(toolName)
|
||||
if !ok {
|
||||
continue
|
||||
@@ -151,5 +189,11 @@ func (m *Manager) filterTools(serverName string, allow, deny []string) {
|
||||
|
||||
kept = append(kept, toolName)
|
||||
}
|
||||
ss.toolNames = kept
|
||||
|
||||
// Update the correct tool names list
|
||||
if isPool {
|
||||
m.poolToolNames[serverName] = kept
|
||||
} else {
|
||||
m.servers[serverName].toolNames = kept
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
)
|
||||
|
||||
// MCPToolSearchTool searches deferred MCP tools by keyword using BM25.
|
||||
// Registered instead of individual MCP tools when the total count exceeds
|
||||
// the inline threshold (search mode). Discovered tools are activated in
|
||||
// the registry and become available on the next agent loop iteration.
|
||||
type MCPToolSearchTool struct {
|
||||
manager *Manager
|
||||
index *mcpBM25Index
|
||||
}
|
||||
|
||||
// NewMCPToolSearchTool creates an mcp_tool_search tool backed by BM25.
|
||||
func NewMCPToolSearchTool(mgr *Manager) *MCPToolSearchTool {
|
||||
t := &MCPToolSearchTool{
|
||||
manager: mgr,
|
||||
index: newMCPBM25Index(),
|
||||
}
|
||||
t.rebuildIndex()
|
||||
return t
|
||||
}
|
||||
|
||||
func (t *MCPToolSearchTool) rebuildIndex() {
|
||||
deferred := t.manager.DeferredToolInfos()
|
||||
t.index.build(deferred)
|
||||
slog.Debug("mcp_tool_search.index_built", "tools", len(deferred))
|
||||
}
|
||||
|
||||
func (t *MCPToolSearchTool) Name() string { return "mcp_tool_search" }
|
||||
|
||||
func (t *MCPToolSearchTool) Description() string {
|
||||
return "Search for available external integration tools (MCP) by keyword. " +
|
||||
"IMPORTANT: You have access to external service integrations " +
|
||||
"(databases, APIs, file systems, messaging, etc.) through MCP tools " +
|
||||
"that are NOT loaded by default. Before performing any external service " +
|
||||
"operation, you MUST search here first to discover available tools. " +
|
||||
"Use English keywords describing what you need " +
|
||||
"(e.g. 'database query', 'create issue', 'send email'). " +
|
||||
"Discovered tools become immediately available for use."
|
||||
}
|
||||
|
||||
func (t *MCPToolSearchTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "English keywords describing the operation you need (e.g. 'create github issue', 'query postgres', 'send slack message')",
|
||||
},
|
||||
"max_results": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "Maximum number of tools to return (default: 5)",
|
||||
},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *MCPToolSearchTool) Execute(ctx context.Context, args map[string]interface{}) *tools.Result {
|
||||
query, _ := args["query"].(string)
|
||||
if query == "" {
|
||||
return tools.ErrorResult("query parameter is required")
|
||||
}
|
||||
|
||||
maxResults := 5
|
||||
if mr, ok := args["max_results"].(float64); ok && int(mr) > 0 {
|
||||
maxResults = int(mr)
|
||||
}
|
||||
|
||||
results := t.index.search(query, maxResults)
|
||||
|
||||
slog.Info("mcp_tool_search", "query", query, "results", len(results))
|
||||
|
||||
if len(results) == 0 {
|
||||
return tools.NewResult(fmt.Sprintf(
|
||||
"No MCP tools found matching: %q\nProceed with other available tools.", query))
|
||||
}
|
||||
|
||||
// Activate matched tools in the registry
|
||||
names := make([]string, len(results))
|
||||
for i, r := range results {
|
||||
names[i] = r.RegisteredName
|
||||
}
|
||||
t.manager.ActivateTools(names)
|
||||
|
||||
data, _ := json.MarshalIndent(map[string]interface{}{
|
||||
"tools": results,
|
||||
"count": len(results),
|
||||
}, "", " ")
|
||||
|
||||
return tools.NewResult(string(data) +
|
||||
"\n\nThe above tools are now activated and available for use. " +
|
||||
"Call them directly by name to perform your operation.")
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
mcpgo "github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// poolEntry holds a shared connection and its discovered tools.
|
||||
type poolEntry struct {
|
||||
state *serverState // connection + health state
|
||||
tools []mcpgo.Tool // discovered MCP tool definitions
|
||||
refCount int // number of active Manager references
|
||||
}
|
||||
|
||||
// Pool manages shared MCP server connections across agents.
|
||||
// One physical connection per server name, shared by all agents
|
||||
// that have grants to that server. Each agent creates its own
|
||||
// BridgeTools pointing to the shared client/connected pointers.
|
||||
type Pool struct {
|
||||
mu sync.Mutex
|
||||
servers map[string]*poolEntry
|
||||
}
|
||||
|
||||
// NewPool creates a shared MCP connection pool.
|
||||
func NewPool() *Pool {
|
||||
return &Pool{
|
||||
servers: make(map[string]*poolEntry),
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire returns a shared connection for the named server.
|
||||
// If no connection exists, it connects using the provided config.
|
||||
// Increments the reference count.
|
||||
func (p *Pool) Acquire(ctx context.Context, name, transportType, command string, args []string, env map[string]string, url string, headers map[string]string, timeoutSec int) (*poolEntry, error) {
|
||||
p.mu.Lock()
|
||||
|
||||
if entry, ok := p.servers[name]; ok && entry.state.connected.Load() {
|
||||
entry.refCount++
|
||||
p.mu.Unlock()
|
||||
slog.Debug("mcp.pool.reuse", "server", name, "refCount", entry.refCount)
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
// If entry exists but disconnected, close old connection first
|
||||
if old, ok := p.servers[name]; ok {
|
||||
if old.state.cancel != nil {
|
||||
old.state.cancel()
|
||||
}
|
||||
if old.state.client != nil {
|
||||
_ = old.state.client.Close()
|
||||
}
|
||||
delete(p.servers, name)
|
||||
}
|
||||
|
||||
p.mu.Unlock()
|
||||
|
||||
// Connect outside the lock (may be slow)
|
||||
ss, mcpTools, err := connectAndDiscover(ctx, name, transportType, command, args, env, url, headers, timeoutSec)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Start health loop
|
||||
hctx, hcancel := context.WithCancel(context.Background())
|
||||
ss.cancel = hcancel
|
||||
go poolHealthLoop(hctx, ss)
|
||||
|
||||
entry := &poolEntry{
|
||||
state: ss,
|
||||
tools: mcpTools,
|
||||
refCount: 1,
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
// Check if another goroutine connected while we were connecting
|
||||
if existing, ok := p.servers[name]; ok && existing.state.connected.Load() {
|
||||
// Use existing, close ours
|
||||
p.mu.Unlock()
|
||||
hcancel()
|
||||
_ = ss.client.Close()
|
||||
p.mu.Lock()
|
||||
existing.refCount++
|
||||
p.mu.Unlock()
|
||||
return existing, nil
|
||||
}
|
||||
p.servers[name] = entry
|
||||
p.mu.Unlock()
|
||||
|
||||
slog.Info("mcp.pool.connected", "server", name, "tools", len(mcpTools))
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
// Release decrements the reference count for a server.
|
||||
// The connection is NOT closed when refCount reaches 0 — it stays
|
||||
// alive for future agents. Use Stop() to close all connections.
|
||||
func (p *Pool) Release(name string) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if entry, ok := p.servers[name]; ok {
|
||||
entry.refCount--
|
||||
if entry.refCount < 0 {
|
||||
entry.refCount = 0
|
||||
}
|
||||
slog.Debug("mcp.pool.release", "server", name, "refCount", entry.refCount)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop closes all pooled connections. Called on gateway shutdown.
|
||||
func (p *Pool) Stop() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
for name, entry := range p.servers {
|
||||
if entry.state.cancel != nil {
|
||||
entry.state.cancel()
|
||||
}
|
||||
if entry.state.client != nil {
|
||||
_ = entry.state.client.Close()
|
||||
}
|
||||
slog.Debug("mcp.pool.stopped", "server", name)
|
||||
}
|
||||
p.servers = make(map[string]*poolEntry)
|
||||
}
|
||||
|
||||
// poolHealthLoop is a standalone health loop for pool-managed connections.
|
||||
// Unlike Manager.healthLoop, it doesn't trigger reconnection via Manager
|
||||
// since the pool owns the connection lifecycle.
|
||||
func poolHealthLoop(ctx context.Context, ss *serverState) {
|
||||
ticker := newHealthTicker()
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := ss.client.Ping(ctx); err != nil {
|
||||
if isMethodNotFound(err) {
|
||||
ss.connected.Store(true)
|
||||
continue
|
||||
}
|
||||
ss.connected.Store(false)
|
||||
ss.mu.Lock()
|
||||
ss.lastErr = err.Error()
|
||||
ss.mu.Unlock()
|
||||
slog.Warn("mcp.pool.health_failed", "server", ss.name, "error", err)
|
||||
} else {
|
||||
ss.connected.Store(true)
|
||||
ss.mu.Lock()
|
||||
ss.reconnAttempts = 0
|
||||
ss.lastErr = ""
|
||||
ss.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -27,7 +27,7 @@ var toolGroups = map[string][]string{
|
||||
"memory_search", "memory_get",
|
||||
"sessions_list", "sessions_history", "sessions_send", "spawn", "session_status",
|
||||
"cron", "message", "create_forum_topic",
|
||||
"read_image", "create_image", "skill_search", "tts",
|
||||
"read_image", "create_image", "skill_search", "mcp_tool_search", "tts",
|
||||
"handoff", "delegate_search", "evaluate_loop",
|
||||
"team_tasks", "team_message",
|
||||
},
|
||||
|
||||
Reference in new issue
Block a user