From 0f6fadb445041aba56d04b764c6a3776b25427ff Mon Sep 17 00:00:00 2001
From: yuchou87 <yuchou87@gmail.com>
Date: Mon, 16 Feb 2026 20:00:37 +0800
Subject: [PATCH] fix(agent): register MCP tools after server initialization

Critical bug fix:
- MCP tools were never registered because servers loaded in Run()
  but tool registration happened in NewAgentLoop() with empty manager
- Move MCP tool registration from createToolRegistry to Run()
- Register MCP tools for both main agent and subag after successful server loading
- Add subagentManager field to AgentLoop for dynamic registration
- Add tool_count logging for better observability

This ensures MCP tools are properly available to both agent and subagents.
---
 pkg/agent/loop.go | 106 ++++++++++++++++++++++++++--------------------
 1 file changed, 61 insertions(+), 45 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index a18e962fa..c4d7c4ae7 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -32,22 +32,23 @@ import (
 )
 
 type AgentLoop struct {
-	bus            *bus.MessageBus
-	provider       providers.LLMProvider
-	workspace      string
-	model          string
-	contextWindow  int // Maximum context window size in tokens
-	maxIterations  int
-	sessions       *session.SessionManager
-	state          *state.Manager
-	contextBuilder *ContextBuilder
-	tools          *tools.ToolRegistry
-	mcpManager     *mcp.Manager   // MCP server manager for resource cleanup
-	mcpConfig      *config.Config // Config for lazy MCP initialization
-	mcpInitOnce    sync.Once      // Ensures MCP is initialized only once
-	running        atomic.Bool
-	summarizing    sync.Map // Tracks which sessions are currently being summarized
-	channelManager *channels.Manager
+	bus             *bus.MessageBus
+	provider        providers.LLMProvider
+	workspace       string
+	model           string
+	contextWindow   int // Maximum context window size in tokens
+	maxIterations   int
+	sessions        *session.SessionManager
+	state           *state.Manager
+	contextBuilder  *ContextBuilder
+	tools           *tools.ToolRegistry
+	mcpManager      *mcp.Manager           // MCP server manager for resource cleanup
+	mcpConfig       *config.Config         // Config for lazy MCP initialization
+	mcpInitOnce     sync.Once              // Ensures MCP is initialized only once
+	subagentManager *tools.SubagentManager // Subagent manager for MCP tool registration
+	running         atomic.Bool
+	summarizing     sync.Map // Tracks which sessions are currently being summarized
+	channelManager  *channels.Manager
 }
 
 // processOptions configures how a message is processed
@@ -105,22 +106,8 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
 	})
 	registry.Register(messageTool)
 
-	// Register MCP tools from all connected servers
-	if mcpManager != nil {
-		servers := mcpManager.GetServers()
-		for serverName, conn := range servers {
-			for _, tool := range conn.Tools {
-				mcpTool := tools.NewMCPTool(mcpManager, serverName, tool)
-				registry.Register(mcpTool)
-				logger.DebugCF("agent", "Registered MCP tool",
-					map[string]interface{}{
-						"server": serverName,
-						"tool":   tool.Name,
-						"name":   mcpTool.Name(),
-					})
-			}
-		}
-	}
+	// Note: MCP tools are registered dynamically in Run() after servers are loaded
+	// This ensures we use the agent's lifecycle context for MCP connections
 
 	return registry
 }
@@ -162,19 +149,20 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
 	contextBuilder.SetToolsRegistry(toolsRegistry)
 
 	return &AgentLoop{
-		bus:            msgBus,
-		provider:       provider,
-		workspace:      workspace,
-		model:          cfg.Agents.Defaults.Model,
-		contextWindow:  cfg.Agents.Defaults.MaxTokens, // Restore context window for summarization
-		maxIterations:  cfg.Agents.Defaults.MaxToolIterations,
-		sessions:       sessionsManager,
-		state:          stateManager,
-		contextBuilder: contextBuilder,
-		tools:          toolsRegistry,
-		mcpManager:     mcpManager,
-		mcpConfig:      cfg, // Store config for lazy initialization in Run()
-		summarizing:    sync.Map{},
+		bus:             msgBus,
+		provider:        provider,
+		workspace:       workspace,
+		model:           cfg.Agents.Defaults.Model,
+		contextWindow:   cfg.Agents.Defaults.MaxTokens, // Restore context window for summarization
+		maxIterations:   cfg.Agents.Defaults.MaxToolIterations,
+		sessions:        sessionsManager,
+		state:           stateManager,
+		contextBuilder:  contextBuilder,
+		tools:           toolsRegistry,
+		mcpManager:      mcpManager,
+		mcpConfig:       cfg, // Store config for lazy initialization in Run()
+		subagentManager: subagentManager,
+		summarizing:     sync.Map{},
 	}
 }
 
@@ -189,6 +177,34 @@ func (al *AgentLoop) Run(ctx context.Context) error {
 				map[string]interface{}{
 					"error": err.Error(),
 				})
+		} else {
+			// Register for both main agent and subagents
+			servers := al.mcpManager.GetServers()
+			toolCount := 0
+			for serverName, conn := range servers {
+				for _, tool := range conn.Tools {
+					mcpTool := tools.NewMCPTool(al.mcpManager, serverName, tool)
+					al.tools.Register(mcpTool)
+
+					// Also register for subagent
+					if al.subagentManager != nil {
+						al.subagentManager.RegisterTool(tools.NewMCPTool(al.mcpManager, serverName, tool))
+					}
+
+					toolCount++
+					logger.DebugCF("agent", "Registered MCP tool",
+						map[string]interface{}{
+							"server": serverName,
+							"tool":   tool.Name,
+							"name":   mcpTool.Name(),
+						})
+				}
+			}
+			logger.InfoCF("agent", "MCP tools registered successfully",
+				map[string]interface{}{
+					"server_count": len(servers),
+					"tool_count":   toolCount,
+				})
 		}
 	})