From 01e40a92b05bae21e498e7f4dbbe27832bf8b667 Mon Sep 17 00:00:00 2001 From: viettranx Date: Sat, 7 Mar 2026 18:52:37 +0700 Subject: [PATCH] feat(mcp): add dynamic tool search and shared connection pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- cmd/gateway.go | 6 +- cmd/gateway_managed.go | 14 +- internal/agent/loop_history.go | 2 + internal/agent/resolver.go | 17 +- internal/agent/systemprompt.go | 11 +- internal/agent/systemprompt_sections.go | 15 ++ internal/mcp/bm25_index.go | 180 +++++++++++++++++++++ internal/mcp/manager.go | 203 ++++++++++++++++++++++-- internal/mcp/manager_connect.go | 152 ++++++++++++++---- internal/mcp/manager_tools.go | 76 +++++++-- internal/mcp/mcp_tool_search.go | 102 ++++++++++++ internal/mcp/pool.go | 160 +++++++++++++++++++ internal/tools/policy.go | 2 +- 13 files changed, 865 insertions(+), 75 deletions(-) create mode 100644 internal/mcp/bm25_index.go create mode 100644 internal/mcp/mcp_tool_search.go create mode 100644 internal/mcp/pool.go diff --git a/cmd/gateway.go b/cmd/gateway.go index b8a5552b..dd2f1b2b 100644 --- a/cmd/gateway.go +++ b/cmd/gateway.go @@ -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 { diff --git a/cmd/gateway_managed.go b/cmd/gateway_managed.go index d1cd72ae..c1f7344d 100644 --- a/cmd/gateway_managed.go +++ b/cmd/gateway_managed.go @@ -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). diff --git a/internal/agent/loop_history.go b/internal/agent/loop_history.go index 525a66c8..32594575 100644 --- a/internal/agent/loop_history.go +++ b/internal/agent/loop_history.go @@ -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, diff --git a/internal/agent/resolver.go b/internal/agent/resolver.go index 7a719389..c5d8689c 100644 --- a/internal/agent/resolver.go +++ b/internal/agent/resolver.go @@ -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 { diff --git a/internal/agent/systemprompt.go b/internal/agent/systemprompt.go index 4ae2637f..d9a18ca5 100644 --- a/internal/agent/systemprompt.go +++ b/internal/agent/systemprompt.go @@ -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()...) diff --git a/internal/agent/systemprompt_sections.go b/internal/agent/systemprompt_sections.go index fd5e7cac..9cba3581 100644 --- a/internal/agent/systemprompt_sections.go +++ b/internal/agent/systemprompt_sections.go @@ -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{ diff --git a/internal/mcp/bm25_index.go b/internal/mcp/bm25_index.go new file mode 100644 index 00000000..0ae73d0a --- /dev/null +++ b/internal/mcp/bm25_index.go @@ -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 +} diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index 4f1361c2..de61e0bf 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -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. diff --git a/internal/mcp/manager_connect.go b/internal/mcp/manager_connect.go index b7606c89..1590bc97 100644 --- a/internal/mcp/manager_connect.go +++ b/internal/mcp/manager_connect.go @@ -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 diff --git a/internal/mcp/manager_tools.go b/internal/mcp/manager_tools.go index 5f8ce379..f469abf3 100644 --- a/internal/mcp/manager_tools.go +++ b/internal/mcp/manager_tools.go @@ -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 + } } diff --git a/internal/mcp/mcp_tool_search.go b/internal/mcp/mcp_tool_search.go new file mode 100644 index 00000000..e8dc5b08 --- /dev/null +++ b/internal/mcp/mcp_tool_search.go @@ -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.") +} diff --git a/internal/mcp/pool.go b/internal/mcp/pool.go new file mode 100644 index 00000000..4786051f --- /dev/null +++ b/internal/mcp/pool.go @@ -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() + } + } + } +} diff --git a/internal/tools/policy.go b/internal/tools/policy.go index 01fcb9c3..10ac7977 100644 --- a/internal/tools/policy.go +++ b/internal/tools/policy.go @@ -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", },