diff --git a/pkg/tools/mcp.go b/pkg/tools/mcp.go index edb6f6735..c1965cca6 100644 --- a/pkg/tools/mcp.go +++ b/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 { - continue - } - - client := newMCPClient(serverCfg, workspace) - startupTimeout := durationFromMS(serverCfg.StartupTimeoutMS, defaultMCPStartupTimeout) - - connectCtx, cancel := context.WithTimeout(ctx, startupTimeout) - remoteTools, err := client.ListTools(connectCtx) - cancel() + serverTools, err := loadMCPServerTools(ctx, serverCfg, workspace, usedNames) + loaded = append(loaded, serverTools...) if err != nil { - errs = append(errs, fmt.Errorf("mcp server %q discovery failed: %w", serverCfg.Name, 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, - }) + 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 { + 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) - 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 + return c.buildCommandTransport() 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) + endpoint, err := c.requiredServerURL("streamable_http") + if err != nil { + return nil, err } return &mcp.StreamableClientTransport{ - Endpoint: c.cfg.URL, + Endpoint: endpoint, }, 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) + endpoint, err := c.requiredServerURL("sse") + 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") diff --git a/pkg/tools/mcp_external_integration_test.go b/pkg/tools/mcp_external_integration_test.go index 1b5829688..101b48f02 100644 --- a/pkg/tools/mcp_external_integration_test.go +++ b/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{ - Enabled: true, - Servers: []config.MCPServerConfig{ - { - Name: "filesystem", - Enabled: true, - Transport: "command", - Command: "npx", - Args: []string{"-y", "@modelcontextprotocol/server-filesystem", rootCanonical}, - ToolPrefix: "mcp_fs", - StartupTimeoutMS: 30000, - CallTimeoutMS: 30000, - }, - }, - } + tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{ + Name: "filesystem", + Enabled: true, + Transport: "command", + Command: "npx", + Args: []string{"-y", "@modelcontextprotocol/server-filesystem", rootCanonical}, + 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,32 +73,19 @@ func TestMCPExternalPopularMemoryCommand(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() - cfg := config.MCPToolsConfig{ - Enabled: true, - Servers: []config.MCPServerConfig{ - { - Name: "memory", - Enabled: true, - Transport: "command", - Command: "npx", - Args: []string{"-y", "@modelcontextprotocol/server-memory"}, - Env: map[string]string{"MEMORY_FILE_PATH": memoryFile}, - ToolPrefix: "mcp_memory", - StartupTimeoutMS: 30000, - CallTimeoutMS: 30000, - }, - }, - } + tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{ + Name: "memory", + Enabled: true, + Transport: "command", + Command: "npx", + Args: []string{"-y", "@modelcontextprotocol/server-memory"}, + Env: map[string]string{"MEMORY_FILE_PATH": memoryFile}, + 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,39 +106,25 @@ func TestMCPExternalPopularEverythingSSE(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() - cfg := config.MCPToolsConfig{ - Enabled: true, - Servers: []config.MCPServerConfig{ - { - Name: "everything", - Enabled: true, - Transport: "sse", - URL: fmt.Sprintf("http://127.0.0.1:%d/sse", port), - ToolPrefix: "mcp_every", - StartupTimeoutMS: 30000, - CallTimeoutMS: 30000, - }, - }, - } + tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{ + Name: "everything", + Enabled: true, + Transport: "sse", + URL: fmt.Sprintf("http://127.0.0.1:%d/sse", port), + 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,30 +137,17 @@ func TestMCPExternalPopularEverythingStreamableHTTP(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() - cfg := config.MCPToolsConfig{ - Enabled: true, - Servers: []config.MCPServerConfig{ - { - Name: "everything-http", - Enabled: true, - Transport: "streamable_http", - URL: fmt.Sprintf("http://127.0.0.1:%d/mcp", port), - ToolPrefix: "mcp_http", - StartupTimeoutMS: 30000, - CallTimeoutMS: 30000, - }, - }, - } + tools := loadExternalMCPTools(t, ctx, config.MCPServerConfig{ + Name: "everything-http", + Enabled: true, + Transport: "streamable_http", + URL: fmt.Sprintf("http://127.0.0.1:%d/mcp", port), + 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() }) - } diff --git a/pkg/tools/mcp_test.go b/pkg/tools/mcp_test.go index ae21a7883..b029e78ec 100644 --- a/pkg/tools/mcp_test.go +++ b/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) } }