Simplify MCP loading and test setup without behavior changes

pkg/tools/mcp.go

@@ -38,44 +38,52 @@ func LoadMCPTools(ctx context.Context, cfg config.MCPToolsConfig, workspace stri
 	errs := make([]error, 0)
 
 	for _, serverCfg := range cfg.Servers {
-		if !serverCfg.Enabled {
+		serverTools, err := loadMCPServerTools(ctx, serverCfg, workspace, usedNames)
-			continue
+		loaded = append(loaded, serverTools...)
-		}
-
-		client := newMCPClient(serverCfg, workspace)
-		startupTimeout := durationFromMS(serverCfg.StartupTimeoutMS, defaultMCPStartupTimeout)
-
-		connectCtx, cancel := context.WithTimeout(ctx, startupTimeout)
-		remoteTools, err := client.ListTools(connectCtx)
-		cancel()
 		if err != nil {
-			errs = append(errs, fmt.Errorf("mcp server %q discovery failed: %w", serverCfg.Name, err))
+			errs = append(errs, err)
-			continue
-		}
-
-		for _, rt := range remoteTools {
-			if rt == nil || strings.TrimSpace(rt.Name) == "" {
-				continue
-			}
-
-			localName := buildLocalToolName(serverCfg, rt.Name, usedNames)
-			description := buildMCPToolDescription(serverCfg.Name, rt.Name, rt.Description)
-			parameters := normalizeMCPInputSchema(rt.InputSchema)
-
-			loaded = append(loaded, &MCPTool{
-				localName:   localName,
-				remoteName:  rt.Name,
-				description: description,
-				parameters:  parameters,
-				callTimeout: durationFromMS(serverCfg.CallTimeoutMS, defaultMCPCallTimeout),
-				client:      client,
-			})
 		}
 	}
 
 	return loaded, errors.Join(errs...)
 }
 
+func loadMCPServerTools(ctx context.Context, serverCfg config.MCPServerConfig, workspace string, usedNames map[string]int) ([]Tool, error) {
+	if !serverCfg.Enabled {
+		return nil, nil
+	}
+
+	client := newMCPClient(serverCfg, workspace)
+	startupTimeout := durationFromMS(serverCfg.StartupTimeoutMS, defaultMCPStartupTimeout)
+
+	connectCtx, cancel := context.WithTimeout(ctx, startupTimeout)
+	defer cancel()
+
+	remoteTools, err := client.ListTools(connectCtx)
+	if err != nil {
+		return nil, fmt.Errorf("mcp server %q discovery failed: %w", serverCfg.Name, err)
+	}
+
+	callTimeout := durationFromMS(serverCfg.CallTimeoutMS, defaultMCPCallTimeout)
+	loaded := make([]Tool, 0, len(remoteTools))
+	for _, rt := range remoteTools {
+		if rt == nil || strings.TrimSpace(rt.Name) == "" {
+			continue
+		}
+
+		loaded = append(loaded, &MCPTool{
+			localName:   buildLocalToolName(serverCfg, rt.Name, usedNames),
+			remoteName:  rt.Name,
+			description: buildMCPToolDescription(serverCfg.Name, rt.Name, rt.Description),
+			parameters:  normalizeMCPInputSchema(rt.InputSchema),
+			callTimeout: callTimeout,
+			client:      client,
+		})
+	}
+
+	return loaded, nil
+}
+
 type MCPTool struct {
 	localName   string
 	remoteName  string
@@ -193,45 +201,58 @@ func (c *mcpClient) buildTransport() (mcp.Transport, error) {
 
 	switch transport {
 	case "command":
-		command := strings.TrimSpace(c.cfg.Command)
+		return c.buildCommandTransport()
-		if command == "" {
-			return nil, fmt.Errorf("mcp server %q: command is required for command transport", c.cfg.Name)
-		}
-		cmd := exec.Command(command, c.cfg.Args...)
-
-		if wd := resolvePath(c.cfg.WorkingDir, c.workspace); wd != "" {
-			cmd.Dir = wd
-		}
-
-		if len(c.cfg.Env) > 0 {
-			cmd.Env = mergeEnv(os.Environ(), c.cfg.Env)
-		}
-		cmd.Stderr = os.Stderr
-
-		tr := &mcp.CommandTransport{
-			Command: cmd,
-		}
-		tr.TerminateDuration = durationFromMS(c.cfg.TerminateTimeoutMS, defaultMCPTerminateWait)
-		return tr, nil
 	case "streamable_http":
-		if strings.TrimSpace(c.cfg.URL) == "" {
+		endpoint, err := c.requiredServerURL("streamable_http")
-			return nil, fmt.Errorf("mcp server %q: url is required for streamable_http transport", c.cfg.Name)
+		if err != nil {
+			return nil, err
 		}
 		return &mcp.StreamableClientTransport{
-			Endpoint: c.cfg.URL,
+			Endpoint: endpoint,
 		}, nil
 	case "sse":
-		if strings.TrimSpace(c.cfg.URL) == "" {
+		endpoint, err := c.requiredServerURL("sse")
-			return nil, fmt.Errorf("mcp server %q: url is required for sse transport", c.cfg.Name)
+		if err != nil {
+			return nil, err
 		}
 		return &mcp.SSEClientTransport{
-			Endpoint: c.cfg.URL,
+			Endpoint: endpoint,
 		}, nil
 	default:
 		return nil, fmt.Errorf("mcp server %q: unsupported transport %q", c.cfg.Name, c.cfg.Transport)
 	}
 }
 
+func (c *mcpClient) buildCommandTransport() (mcp.Transport, error) {
+	command := strings.TrimSpace(c.cfg.Command)
+	if command == "" {
+		return nil, fmt.Errorf("mcp server %q: command is required for command transport", c.cfg.Name)
+	}
+
+	cmd := exec.Command(command, c.cfg.Args...)
+	if wd := resolvePath(c.cfg.WorkingDir, c.workspace); wd != "" {
+		cmd.Dir = wd
+	}
+	if len(c.cfg.Env) > 0 {
+		cmd.Env = mergeEnv(os.Environ(), c.cfg.Env)
+	}
+	cmd.Stderr = os.Stderr
+
+	tr := &mcp.CommandTransport{
+		Command: cmd,
+	}
+	tr.TerminateDuration = durationFromMS(c.cfg.TerminateTimeoutMS, defaultMCPTerminateWait)
+	return tr, nil
+}
+
+func (c *mcpClient) requiredServerURL(transport string) (string, error) {
+	endpoint := strings.TrimSpace(c.cfg.URL)
+	if endpoint == "" {
+		return "", fmt.Errorf("mcp server %q: url is required for %s transport", c.cfg.Name, transport)
+	}
+	return endpoint, nil
+}
+
 func formatMCPCallToolResult(result *mcp.CallToolResult) (string, error) {
 	if result == nil {
 		return "", fmt.Errorf("empty MCP response")

pkg/tools/mcp_external_integration_test.go

@@ -35,35 +35,19 @@ func TestMCPExternalPopularFilesystemCommand(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
-		Enabled: true,
+		Name:             "filesystem",
-		Servers: []config.MCPServerConfig{
+		Enabled:          true,
-			{
+		Transport:        "command",
-				Name:             "filesystem",
+		Command:          "npx",
-				Enabled:          true,
+		Args:             []string{"-y", "@modelcontextprotocol/server-filesystem", rootCanonical},
-				Transport:        "command",
+		ToolPrefix:       "mcp_fs",
-				Command:          "npx",
+		StartupTimeoutMS: 30000,
-				Args:             []string{"-y", "@modelcontextprotocol/server-filesystem", rootCanonical},
+		CallTimeoutMS:    30000,
-				ToolPrefix:       "mcp_fs",
+	})
-				StartupTimeoutMS: 30000,
-				CallTimeoutMS:    30000,
-			},
-		},
-	}
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
+	listAllowedDirs := requireToolByName(t, tools, "mcp_fs_list_allowed_directories")
-	if err != nil {
+	readFile := requireToolByName(t, tools, "mcp_fs_read_file")
-		t.Fatalf("LoadMCPTools() error: %v", err)
-	}
-
-	listAllowedDirs := findToolByName(tools, "mcp_fs_list_allowed_directories")
-	if listAllowedDirs == nil {
-		t.Fatalf("missing tool mcp_fs_list_allowed_directories; got %v", toolNames(tools))
-	}
-	readFile := findToolByName(tools, "mcp_fs_read_file")
-	if readFile == nil {
-		t.Fatalf("missing tool mcp_fs_read_file; got %v", toolNames(tools))
-	}
 
 	out, err := listAllowedDirs.Execute(ctx, map[string]interface{}{})
 	if err != nil {
@@ -89,32 +73,19 @@ func TestMCPExternalPopularMemoryCommand(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
-		Enabled: true,
+		Name:             "memory",
-		Servers: []config.MCPServerConfig{
+		Enabled:          true,
-			{
+		Transport:        "command",
-				Name:             "memory",
+		Command:          "npx",
-				Enabled:          true,
+		Args:             []string{"-y", "@modelcontextprotocol/server-memory"},
-				Transport:        "command",
+		Env:              map[string]string{"MEMORY_FILE_PATH": memoryFile},
-				Command:          "npx",
+		ToolPrefix:       "mcp_memory",
-				Args:             []string{"-y", "@modelcontextprotocol/server-memory"},
+		StartupTimeoutMS: 30000,
-				Env:              map[string]string{"MEMORY_FILE_PATH": memoryFile},
+		CallTimeoutMS:    30000,
-				ToolPrefix:       "mcp_memory",
+	})
-				StartupTimeoutMS: 30000,
-				CallTimeoutMS:    30000,
-			},
-		},
-	}
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
+	readGraph := requireToolByName(t, tools, "mcp_memory_read_graph")
-	if err != nil {
-		t.Fatalf("LoadMCPTools() error: %v", err)
-	}
-
-	readGraph := findToolByName(tools, "mcp_memory_read_graph")
-	if readGraph == nil {
-		t.Fatalf("missing tool mcp_memory_read_graph; got %v", toolNames(tools))
-	}
 
 	out, err := readGraph.Execute(ctx, map[string]interface{}{})
 	if err != nil {
@@ -135,39 +106,25 @@ func TestMCPExternalPopularEverythingSSE(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
-		Enabled: true,
+		Name:             "everything",
-		Servers: []config.MCPServerConfig{
+		Enabled:          true,
-			{
+		Transport:        "sse",
-				Name:             "everything",
+		URL:              fmt.Sprintf("http://127.0.0.1:%d/sse", port),
-				Enabled:          true,
+		ToolPrefix:       "mcp_every",
-				Transport:        "sse",
+		StartupTimeoutMS: 30000,
-				URL:              fmt.Sprintf("http://127.0.0.1:%d/sse", port),
+		CallTimeoutMS:    30000,
-				ToolPrefix:       "mcp_every",
+	})
-				StartupTimeoutMS: 30000,
-				CallTimeoutMS:    30000,
-			},
-		},
-	}
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
+	echoTool := requireToolByName(t, tools, "mcp_every_echo")
-	if err != nil {
-		t.Fatalf("LoadMCPTools() error: %v", err)
-	}
-
-	echoTool := findToolByName(tools, "mcp_every_echo")
-	if echoTool == nil {
-		t.Fatalf("missing tool mcp_every_echo; got %v", toolNames(tools))
-	}
 
 	out, err := echoTool.Execute(ctx, map[string]interface{}{"message": "hello from sse"})
 	if err != nil {
 		t.Fatalf("Execute(echo) error: %v", err)
 	}
-	if !strings.Contains(out, "Echo: hello from sse") {
+	if !strings.Contains(out, "hello from sse") {
 		t.Fatalf("unexpected echo output: %s", out)
 	}
-
 }
 
 func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) {
@@ -180,30 +137,17 @@ func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
-		Enabled: true,
+		Name:             "everything-http",
-		Servers: []config.MCPServerConfig{
+		Enabled:          true,
-			{
+		Transport:        "streamable_http",
-				Name:             "everything-http",
+		URL:              fmt.Sprintf("http://127.0.0.1:%d/mcp", port),
-				Enabled:          true,
+		ToolPrefix:       "mcp_http",
-				Transport:        "streamable_http",
+		StartupTimeoutMS: 30000,
-				URL:              fmt.Sprintf("http://127.0.0.1:%d/mcp", port),
+		CallTimeoutMS:    30000,
-				ToolPrefix:       "mcp_http",
+	})
-				StartupTimeoutMS: 30000,
-				CallTimeoutMS:    30000,
-			},
-		},
-	}
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
+	echoTool := requireToolByName(t, tools, "mcp_http_echo")
-	if err != nil {
-		t.Fatalf("LoadMCPTools() error: %v", err)
-	}
-
-	echoTool := findToolByName(tools, "mcp_http_echo")
-	if echoTool == nil {
-		t.Fatalf("missing tool mcp_http_echo; got %v", toolNames(tools))
-	}
 
 	out, err := echoTool.Execute(ctx, map[string]interface{}{"message": "hello from streamable-http"})
 	if err != nil {
@@ -212,7 +156,6 @@ func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) {
 	if !strings.Contains(out, "hello from streamable-http") {
 		t.Fatalf("unexpected echo output: %s", out)
 	}
-
 }
 
 func requireExternalMCPTests(t *testing.T) {
@@ -225,6 +168,31 @@ func requireExternalMCPTests(t *testing.T) {
 	}
 }
 
+func loadExternalMCPTools(t *testing.T, ctx context.Context, server config.MCPServerConfig) []Tool {
+	t.Helper()
+
+	cfg := config.MCPToolsConfig{
+		Enabled: true,
+		Servers: []config.MCPServerConfig{server},
+	}
+
+	tools, err := LoadMCPTools(ctx, cfg, "")
+	if err != nil {
+		t.Fatalf("LoadMCPTools() error: %v", err)
+	}
+	return tools
+}
+
+func requireToolByName(t *testing.T, tools []Tool, name string) Tool {
+	t.Helper()
+
+	tool := findToolByName(tools, name)
+	if tool == nil {
+		t.Fatalf("missing tool %s; got %v", name, toolNames(tools))
+	}
+	return tool
+}
+
 func findToolByName(tools []Tool, name string) Tool {
 	for _, tool := range tools {
 		if tool.Name() == name {
@@ -282,5 +250,4 @@ func startEverythingServer(t *testing.T, port int, mode string) {
 		_ = cmd.Process.Kill()
 		_, _ = cmd.Process.Wait()
 	})
-
 }

pkg/tools/mcp_test.go

@@ -169,36 +169,28 @@ func TestLoadMCPTools_InvalidServerAggregatesError(t *testing.T) {
 }
 
 func TestBuildTransport_CommandTerminateDefaults(t *testing.T) {
-	client := newMCPClient(config.MCPServerConfig{
+	assertCommandTransportTerminateDuration(t, config.MCPServerConfig{
 		Name:      "default-terminate",
 		Enabled:   true,
 		Transport: "command",
 		Command:   "test-command",
-	}, "")
+	}, defaultMCPTerminateWait)
-
-	tr, err := client.buildTransport()
-	if err != nil {
-		t.Fatalf("buildTransport() error: %v", err)
-	}
-
-	cmdTr, ok := tr.(*mcp.CommandTransport)
-	if !ok {
-		t.Fatalf("buildTransport() returned %T, want *mcp.CommandTransport", tr)
-	}
-	if cmdTr.TerminateDuration != defaultMCPTerminateWait {
-		t.Fatalf("TerminateDuration = %v, want %v", cmdTr.TerminateDuration, defaultMCPTerminateWait)
-	}
 }
 
 func TestBuildTransport_CommandTerminateOverride(t *testing.T) {
-	client := newMCPClient(config.MCPServerConfig{
+	assertCommandTransportTerminateDuration(t, config.MCPServerConfig{
 		Name:               "override-terminate",
 		Enabled:            true,
 		Transport:          "command",
 		Command:            "test-command",
 		TerminateTimeoutMS: 2500,
-	}, "")
+	}, 2500*time.Millisecond)
+}
 
+func assertCommandTransportTerminateDuration(t *testing.T, cfg config.MCPServerConfig, want time.Duration) {
+	t.Helper()
+
+	client := newMCPClient(cfg, "")
 	tr, err := client.buildTransport()
 	if err != nil {
 		t.Fatalf("buildTransport() error: %v", err)
@@ -208,8 +200,8 @@ func TestBuildTransport_CommandTerminateOverride(t *testing.T) {
 	if !ok {
 		t.Fatalf("buildTransport() returned %T, want *mcp.CommandTransport", tr)
 	}
-	if cmdTr.TerminateDuration != 2500*time.Millisecond {
+	if cmdTr.TerminateDuration != want {
-		t.Fatalf("TerminateDuration = %v, want %v", cmdTr.TerminateDuration, 2500*time.Millisecond)
+		t.Fatalf("TerminateDuration = %v, want %v", cmdTr.TerminateDuration, want)
 	}
 }