Simplify MCP loading and test setup without behavior changes

pkg/tools/mcp.go

@@ -38,42 +38,50 @@ func LoadMCPTools(ctx context.Context, cfg config.MCPToolsConfig, workspace stri
 	errs := make([]error, 0)
 
 	for _, serverCfg := range cfg.Servers {
+		serverTools, err := loadMCPServerTools(ctx, serverCfg, workspace, usedNames)
+		loaded = append(loaded, serverTools...)
+		if err != nil {
+			errs = append(errs, err)
+		}
+	}
+
+	return loaded, errors.Join(errs...)
+}
+
+func loadMCPServerTools(ctx context.Context, serverCfg config.MCPServerConfig, workspace string, usedNames map[string]int) ([]Tool, error) {
 	if !serverCfg.Enabled {
-			continue
+		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)
-		cancel()
 	if err != nil {
-			errs = append(errs, fmt.Errorf("mcp server %q discovery failed: %w", serverCfg.Name, err))
-			continue
+		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
 		}
 
-			localName := buildLocalToolName(serverCfg, rt.Name, usedNames)
-			description := buildMCPToolDescription(serverCfg.Name, rt.Name, rt.Description)
-			parameters := normalizeMCPInputSchema(rt.InputSchema)
-
 		loaded = append(loaded, &MCPTool{
-				localName:   localName,
+			localName:   buildLocalToolName(serverCfg, rt.Name, usedNames),
 			remoteName:  rt.Name,
-				description: description,
-				parameters:  parameters,
-				callTimeout: durationFromMS(serverCfg.CallTimeoutMS, defaultMCPCallTimeout),
+			description: buildMCPToolDescription(serverCfg.Name, rt.Name, rt.Description),
+			parameters:  normalizeMCPInputSchema(rt.InputSchema),
+			callTimeout: callTimeout,
 			client:      client,
 		})
 	}
-	}
 
-	return loaded, errors.Join(errs...)
+	return loaded, nil
 }
 
 type MCPTool struct {
@@ -193,16 +201,38 @@ func (c *mcpClient) buildTransport() (mcp.Transport, error) {
 
 	switch transport {
 	case "command":
+		return c.buildCommandTransport()
+	case "streamable_http":
+		endpoint, err := c.requiredServerURL("streamable_http")
+		if err != nil {
+			return nil, err
+		}
+		return &mcp.StreamableClientTransport{
+			Endpoint: endpoint,
+		}, nil
+	case "sse":
+		endpoint, err := c.requiredServerURL("sse")
+		if err != nil {
+			return nil, err
+		}
+		return &mcp.SSEClientTransport{
+			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...)
 
+	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)
 	}
@@ -213,23 +243,14 @@ func (c *mcpClient) buildTransport() (mcp.Transport, error) {
 	}
 	tr.TerminateDuration = durationFromMS(c.cfg.TerminateTimeoutMS, defaultMCPTerminateWait)
 	return tr, nil
-	case "streamable_http":
-		if strings.TrimSpace(c.cfg.URL) == "" {
-			return nil, fmt.Errorf("mcp server %q: url is required for streamable_http transport", c.cfg.Name)
-		}
-		return &mcp.StreamableClientTransport{
-			Endpoint: c.cfg.URL,
-		}, nil
-	case "sse":
-		if strings.TrimSpace(c.cfg.URL) == "" {
-			return nil, fmt.Errorf("mcp server %q: url is required for sse transport", c.cfg.Name)
-		}
-		return &mcp.SSEClientTransport{
-			Endpoint: c.cfg.URL,
-		}, nil
-	default:
-		return nil, fmt.Errorf("mcp server %q: unsupported transport %q", c.cfg.Name, c.cfg.Transport)
+}
+
+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) {

pkg/tools/mcp_external_integration_test.go

@@ -35,10 +35,7 @@ func TestMCPExternalPopularFilesystemCommand(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
-		Enabled: true,
-		Servers: []config.MCPServerConfig{
-			{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
 		Name:             "filesystem",
 		Enabled:          true,
 		Transport:        "command",
@@ -47,23 +44,10 @@ func TestMCPExternalPopularFilesystemCommand(t *testing.T) {
 		ToolPrefix:       "mcp_fs",
 		StartupTimeoutMS: 30000,
 		CallTimeoutMS:    30000,
-			},
-		},
-	}
+	})
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
-	if err != nil {
-		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))
-	}
+	listAllowedDirs := requireToolByName(t, tools, "mcp_fs_list_allowed_directories")
+	readFile := requireToolByName(t, tools, "mcp_fs_read_file")
 
 	out, err := listAllowedDirs.Execute(ctx, map[string]interface{}{})
 	if err != nil {
@@ -89,10 +73,7 @@ func TestMCPExternalPopularMemoryCommand(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
-		Enabled: true,
-		Servers: []config.MCPServerConfig{
-			{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
 		Name:             "memory",
 		Enabled:          true,
 		Transport:        "command",
@@ -102,19 +83,9 @@ func TestMCPExternalPopularMemoryCommand(t *testing.T) {
 		ToolPrefix:       "mcp_memory",
 		StartupTimeoutMS: 30000,
 		CallTimeoutMS:    30000,
-			},
-		},
-	}
+	})
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
-	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))
-	}
+	readGraph := requireToolByName(t, tools, "mcp_memory_read_graph")
 
 	out, err := readGraph.Execute(ctx, map[string]interface{}{})
 	if err != nil {
@@ -135,10 +106,7 @@ func TestMCPExternalPopularEverythingSSE(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
-		Enabled: true,
-		Servers: []config.MCPServerConfig{
-			{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
 		Name:             "everything",
 		Enabled:          true,
 		Transport:        "sse",
@@ -146,28 +114,17 @@ func TestMCPExternalPopularEverythingSSE(t *testing.T) {
 		ToolPrefix:       "mcp_every",
 		StartupTimeoutMS: 30000,
 		CallTimeoutMS:    30000,
-			},
-		},
-	}
+	})
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
-	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))
-	}
+	echoTool := requireToolByName(t, tools, "mcp_every_echo")
 
 	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,10 +137,7 @@ func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancel()
 
-	cfg := config.MCPToolsConfig{
-		Enabled: true,
-		Servers: []config.MCPServerConfig{
-			{
+	tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{
 		Name:             "everything-http",
 		Enabled:          true,
 		Transport:        "streamable_http",
@@ -191,19 +145,9 @@ func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) {
 		ToolPrefix:       "mcp_http",
 		StartupTimeoutMS: 30000,
 		CallTimeoutMS:    30000,
-			},
-		},
-	}
+	})
 
-	tools, err := LoadMCPTools(ctx, cfg, "")
-	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))
-	}
+	echoTool := requireToolByName(t, tools, "mcp_http_echo")
 
 	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",
-	}, "")
-
-	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)
-	}
+	}, 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 {
-		t.Fatalf("TerminateDuration = %v, want %v", cmdTr.TerminateDuration, 2500*time.Millisecond)
+	if cmdTr.TerminateDuration != want {
+		t.Fatalf("TerminateDuration = %v, want %v", cmdTr.TerminateDuration, want)
 	}
 }