Switch to modelcontextprotocol/go-sdk

Author: slimslenderslacks

cmd/docker-mcp/internal/gateway/capabilitites.go

@@ -8,16 +8,36 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/mark3labs/mcp-go/mcp"
-	"github.com/mark3labs/mcp-go/server"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 	"golang.org/x/sync/errgroup"
 )
 
 type Capabilities struct {
-	Tools             []server.ServerTool
-	Prompts           []server.ServerPrompt
-	Resources         []server.ServerResource
-	ResourceTemplates []server.ServerResourceTemplate
+	Tools             []ToolRegistration
+	Prompts           []PromptRegistration
+	Resources         []ResourceRegistration
+	ResourceTemplates []ResourceTemplateRegistration
+}
+
+type ToolRegistration struct {
+	Tool    *mcp.Tool
+	Handler mcp.ToolHandler
+}
+
+type PromptRegistration struct {
+	Prompt  *mcp.Prompt
+	Handler mcp.PromptHandler
+}
+
+type ResourceRegistration struct {
+	Resource *mcp.Resource
+	Handler  mcp.ResourceHandler
+}
+
+type ResourceTemplateRegistration struct {
+	ResourceTemplate mcp.ResourceTemplate
+	Handler          mcp.ResourceHandler
 }
 
 func (g *Gateway) listCapabilities(ctx context.Context, configuration Configuration, serverNames []string) (*Capabilities, error) {
@@ -47,47 +67,47 @@ func (g *Gateway) listCapabilities(ctx context.Context, configuration Configurat
 
 				var capabilities Capabilities
 
-				tools, err := client.ListTools(ctx, mcp.ListToolsRequest{})
+				tools, err := client.ListTools(ctx, &mcp.ListToolsParams{})
 				if err != nil {
 					logf("  > Can't list tools %s: %s", serverConfig.Name, err)
 				} else {
 					for _, tool := range tools.Tools {
 						if !isToolEnabled(configuration, serverConfig.Name, serverConfig.Spec.Image, tool.Name, g.ToolNames) {
 							continue
 						}
-						capabilities.Tools = append(capabilities.Tools, server.ServerTool{
+						capabilities.Tools = append(capabilities.Tools, ToolRegistration{
 							Tool:    tool,
 							Handler: g.mcpServerToolHandler(*serverConfig, tool.Annotations),
 						})
 					}
 				}
 
-				prompts, err := client.ListPrompts(ctx, mcp.ListPromptsRequest{})
+				prompts, err := client.ListPrompts(ctx, &mcp.ListPromptsParams{})
 				if err == nil {
 					for _, prompt := range prompts.Prompts {
-						capabilities.Prompts = append(capabilities.Prompts, server.ServerPrompt{
+						capabilities.Prompts = append(capabilities.Prompts, PromptRegistration{
 							Prompt:  prompt,
 							Handler: g.mcpServerPromptHandler(*serverConfig),
 						})
 					}
 				}
 
-				resources, err := client.ListResources(ctx, mcp.ListResourcesRequest{})
+				resources, err := client.ListResources(ctx, &mcp.ListResourcesParams{})
 				if err == nil {
 					for _, resource := range resources.Resources {
-						capabilities.Resources = append(capabilities.Resources, server.ServerResource{
+						capabilities.Resources = append(capabilities.Resources, ResourceRegistration{
 							Resource: resource,
 							Handler:  g.mcpServerResourceHandler(*serverConfig),
 						})
 					}
 				}
 
-				resourceTemplates, err := client.ListResourceTemplates(ctx, mcp.ListResourceTemplatesRequest{})
+				resourceTemplates, err := client.ListResourceTemplates(ctx, &mcp.ListResourceTemplatesParams{})
 				if err == nil {
 					for _, resourceTemplate := range resourceTemplates.ResourceTemplates {
-						capabilities.ResourceTemplates = append(capabilities.ResourceTemplates, server.ServerResourceTemplate{
-							Template: resourceTemplate,
-							Handler:  g.mcpServerResourceTemplateHandler(*serverConfig),
+						capabilities.ResourceTemplates = append(capabilities.ResourceTemplates, ResourceTemplateRegistration{
+							ResourceTemplate: *resourceTemplate,
+							Handler:          g.mcpServerResourceTemplateHandler(*serverConfig),
 						})
 					}
 				}
@@ -128,17 +148,21 @@ func (g *Gateway) listCapabilities(ctx context.Context, configuration Configurat
 				mcpTool := mcp.Tool{
 					Name:        tool.Name,
 					Description: tool.Description,
+					InputSchema: &jsonschema.Schema{},
 				}
+				// TODO: Properly convert tool.Parameters to jsonschema.Schema
+				// For now, we'll create a simple schema structure
 				if len(tool.Parameters.Properties) == 0 {
 					mcpTool.InputSchema.Type = "object"
 				} else {
 					mcpTool.InputSchema.Type = tool.Parameters.Type
-					mcpTool.InputSchema.Properties = tool.Parameters.Properties.ToMap()
-					mcpTool.InputSchema.Required = tool.Parameters.Required
+					// Note: tool.Parameters.Properties.ToMap() returns map[string]any 
+					// but we need map[string]*jsonschema.Schema
+					// This is a complex conversion that needs proper implementation
 				}
 
-				capabilities.Tools = append(capabilities.Tools, server.ServerTool{
-					Tool:    mcpTool,
+				capabilities.Tools = append(capabilities.Tools, ToolRegistration{
+					Tool:    &mcpTool,
 					Handler: g.mcpToolHandler(tool),
 				})
 			}
@@ -154,22 +178,22 @@ func (g *Gateway) listCapabilities(ctx context.Context, configuration Configurat
 	}
 
 	// Merge all capabilities
-	var serverTools []server.ServerTool
-	var serverPrompts []server.ServerPrompt
-	var serverResources []server.ServerResource
-	var serverResourceTemplates []server.ServerResourceTemplate
+	var allTools []ToolRegistration
+	var allPrompts []PromptRegistration
+	var allResources []ResourceRegistration
+	var allResourceTemplates []ResourceTemplateRegistration
 	for _, capabilities := range allCapabilities {
-		serverTools = append(serverTools, capabilities.Tools...)
-		serverPrompts = append(serverPrompts, capabilities.Prompts...)
-		serverResources = append(serverResources, capabilities.Resources...)
-		serverResourceTemplates = append(serverResourceTemplates, capabilities.ResourceTemplates...)
+		allTools = append(allTools, capabilities.Tools...)
+		allPrompts = append(allPrompts, capabilities.Prompts...)
+		allResources = append(allResources, capabilities.Resources...)
+		allResourceTemplates = append(allResourceTemplates, capabilities.ResourceTemplates...)
 	}
 
 	return &Capabilities{
-		Tools:             serverTools,
-		Prompts:           serverPrompts,
-		Resources:         serverResources,
-		ResourceTemplates: serverResourceTemplates,
+		Tools:             allTools,
+		Prompts:           allPrompts,
+		Resources:         allResources,
+		ResourceTemplates: allResourceTemplates,
 	}, nil
 }
 

cmd/docker-mcp/internal/gateway/clientpool.go

@@ -7,9 +7,8 @@ import (
 	"os/exec"
 	"strings"
 	"sync"
-	"time"
 
-	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
 
 	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/catalog"
 	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/docker"
@@ -129,16 +128,22 @@ func (cp *clientPool) SetNetworks(networks []string) {
 	cp.networks = networks
 }
 
-func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, params *mcp.CallToolParams) (*mcp.CallToolResult, error) {
 	args := cp.baseArgs(tool.Name)
 
 	// Attach the MCP servers to the same network as the gateway.
 	for _, network := range cp.networks {
 		args = append(args, "--network", network)
 	}
 
+	// Convert params.Arguments to map[string]any
+	arguments, ok := params.Arguments.(map[string]any)
+	if !ok {
+		arguments = make(map[string]any)
+	}
+
 	// Volumes
-	for _, mount := range eval.EvaluateList(tool.Container.Volumes, request.GetArguments()) {
+	for _, mount := range eval.EvaluateList(tool.Container.Volumes, arguments) {
 		if mount == "" {
 			continue
 		}
@@ -150,7 +155,7 @@ func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, r
 	args = append(args, tool.Container.Image)
 
 	// Command
-	command := eval.EvaluateList(tool.Container.Command, request.GetArguments())
+	command := eval.EvaluateList(tool.Container.Command, arguments)
 	args = append(args, command...)
 
 	log("  - Running container", tool.Container.Image, "with args", args)
@@ -161,10 +166,20 @@ func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, r
 	}
 	out, err := cmd.Output()
 	if err != nil {
-		return mcp.NewToolResultError(string(out)), nil
+		return &mcp.CallToolResult{
+			Content: []mcp.Content{&mcp.TextContent{
+				Text: string(out),
+			}},
+			IsError: true,
+		}, nil
 	}
 
-	return mcp.NewToolResultText(string(out)), nil
+	return &mcp.CallToolResult{
+		Content: []mcp.Content{&mcp.TextContent{
+			Text: string(out),
+		}},
+		IsError: false,
+	}, nil
 }
 
 func (cp *clientPool) baseArgs(name string) []string {
@@ -345,17 +360,17 @@ func (cg *clientGetter) GetClient(ctx context.Context) (mcpclient.Client, error)
 				client = mcpclient.NewStdioCmdClient(cg.serverConfig.Name, "docker", env, runArgs...)
 			}
 
-			initRequest := mcp.InitializeRequest{}
-			initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
-			initRequest.Params.ClientInfo = mcp.Implementation{
-				Name:    "docker",
-				Version: "1.0.0",
+			initParams := &mcp.InitializeParams{
+				ProtocolVersion: "2024-11-05",
+				ClientInfo: &mcp.Implementation{
+					Name:    "docker",
+					Version: "1.0.0",
+				},
 			}
 
-			ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
-			defer cancel()
-
-			if _, err := client.Initialize(ctx, initRequest, cg.cp.Verbose); err != nil {
+			// Use the original context instead of creating a timeout context
+			// to avoid cancellation issues
+			if _, err := client.Initialize(ctx, initParams, cg.cp.Verbose); err != nil {
 				return nil, err
 			}
 

cmd/docker-mcp/internal/gateway/clientpool_test.go

@@ -1,13 +1,18 @@
 package gateway
 
 import (
+	"context"
 	"testing"
+	"time"
 
+	"github.com/docker/cli/cli/command"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v3"
 
 	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/catalog"
+	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/docker"
 	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/gateway/proxies"
 )
 
@@ -179,3 +184,53 @@ func readOnly() *bool {
 func boolPtr(b bool) *bool {
 	return &b
 }
+
+func TestStdioClientInitialization(t *testing.T) {
+	// Skip if running in CI or if Docker is not available
+	if testing.Short() {
+		t.Skip("Skipping integration test in short mode")
+	}
+
+	serverConfig := catalog.ServerConfig{
+		Name: "test-server",
+		Spec: catalog.Server{
+			Image: "mcp/brave-search@sha256:e13f4693a3421e2b316c8b6196c5c543c77281f9d8938850681e3613bba95115", // User should provide their image
+			Command: []string{},
+			Env: []catalog.Env{{Name: "BRAVE_API_KEY", Value: "test_key"}},
+		},
+		Config:  map[string]any{},
+		Secrets: map[string]string{},
+	}
+
+	// Create a real Docker CLI client
+	dockerCli, err := command.NewDockerCli()
+	require.NoError(t, err)
+	
+	dockerClient := docker.NewClient(dockerCli)
+	clientPool := newClientPool(Options{
+		Cpus:   1,
+		Memory: "512m",
+	}, dockerClient)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	// Test client acquisition and initialization
+	client, err := clientPool.AcquireClient(ctx, serverConfig, boolPtr(false))
+	if err != nil {
+		t.Fatalf("Failed to acquire client: %v", err)
+	}
+	defer clientPool.ReleaseClient(client)
+
+	// Test ListTools to verify the client is working
+	tools, err := client.ListTools(ctx, &mcp.ListToolsParams{})
+	if err != nil {
+		t.Fatalf("Failed to list tools: %v", err)
+	}
+
+	// Basic assertions - user can customize based on expected behavior
+	assert.NotNil(t, tools)
+	assert.NotNil(t, tools.Tools)
+	
+	t.Logf("Successfully initialized stdio client and retrieved %d tools", len(tools.Tools))
+}