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:
viettranx committed 2026-03-07 18:52:37 +07:00
1 parent af85b0c354
commit 01e40a92b0
13 files changed
+865 -75

No files matched your search

+5 -1
View File
@@ -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
View File
@@ -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).
+2
View File
@@ -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,
+16 -1
View File
@@ -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 {
+9 -2
View File
@@ -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()...)
+15
View File
@@ -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{
+180
View File
@@ -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
View File
@@ -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
View File
@@ -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
+60 -16
View File
@@ -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
}
}
+102
View File
@@ -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.")
}
+160
View File
@@ -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()
}
}
}
}
+1 -1
View File
@@ -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",
},