clean capabilities on config reload

Author: slimslenderslacks

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

@@ -44,6 +44,12 @@ type Gateway struct {
 
 	sessionCacheMu sync.RWMutex
 	sessionCache   map[*mcp.ServerSession]*ServerSessionCache
+
+	// Track registered capabilities for cleanup during reload
+	registeredToolNames            []string
+	registeredPromptNames          []string
+	registeredResourceURIs         []string
+	registeredResourceTemplateURIs []string
 }
 
 func NewGateway(config Config, docker docker.Client) *Gateway {
@@ -135,7 +141,6 @@ func (g *Gateway) Run(ctx context.Context) error {
 			log("- Client initialized: ", ss.ID())
 			g.ListRoots(ctx, ss)
 		},
-		PageSize:     100,
 		HasPrompts:   true,
 		HasResources: true,
 		HasTools:     true,
@@ -255,16 +260,40 @@ func (g *Gateway) reloadConfiguration(ctx context.Context, configuration Configu
 	// Clear existing capabilities and register new ones
 	// Note: The new SDK doesn't have bulk set methods, so we register individually
 
+	// Clear all existing capabilities by tracking them in the Gateway struct
+	if g.registeredToolNames != nil {
+		g.mcpServer.RemoveTools(g.registeredToolNames...)
+	}
+	if g.registeredPromptNames != nil {
+		g.mcpServer.RemovePrompts(g.registeredPromptNames...)
+	}
+	if g.registeredResourceURIs != nil {
+		g.mcpServer.RemoveResources(g.registeredResourceURIs...)
+	}
+	if g.registeredResourceTemplateURIs != nil {
+		g.mcpServer.RemoveResourceTemplates(g.registeredResourceTemplateURIs...)
+	}
+
+	// Reset tracking slices
+	g.registeredToolNames = nil
+	g.registeredPromptNames = nil
+	g.registeredResourceURIs = nil
+	g.registeredResourceTemplateURIs = nil
+
+	// Add new capabilities and track them
 	for _, tool := range capabilities.Tools {
 		g.mcpServer.AddTool(tool.Tool, tool.Handler)
+		g.registeredToolNames = append(g.registeredToolNames, tool.Tool.Name)
 	}
 
 	for _, prompt := range capabilities.Prompts {
 		g.mcpServer.AddPrompt(prompt.Prompt, prompt.Handler)
+		g.registeredPromptNames = append(g.registeredPromptNames, prompt.Prompt.Name)
 	}
 
 	for _, resource := range capabilities.Resources {
 		g.mcpServer.AddResource(resource.Resource, resource.Handler)
+		g.registeredResourceURIs = append(g.registeredResourceURIs, resource.Resource.URI)
 	}
 
 	// Resource templates are handled as regular resources in the new SDK
@@ -277,6 +306,7 @@ func (g *Gateway) reloadConfiguration(ctx context.Context, configuration Configu
 			MIMEType:    template.ResourceTemplate.MIMEType,
 		}
 		g.mcpServer.AddResourceTemplate(resource, template.Handler)
+		g.registeredResourceTemplateURIs = append(g.registeredResourceTemplateURIs, resource.URITemplate)
 	}
 
 	g.health.SetHealthy()

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

@@ -8,14 +8,13 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 
 	"github.com/docker/mcp-gateway/cmd/docker-mcp/internal/health"
 )
 
 func (g *Gateway) startStdioServer(ctx context.Context, _ io.Reader, _ io.Writer) error {
-	transport := newContextAwareStdioTransport(ctx)
+	transport := mcp.NewStdioTransport()
 	return g.mcpServer.Run(ctx, transport)
 }
 
@@ -130,81 +129,3 @@ func healthHandler(state *health.State) http.HandlerFunc {
 		}
 	}
 }
-
-// contextAwareStdioTransport is a custom stdio transport that handles context cancellation properly
-type contextAwareStdioTransport struct {
-	ctx context.Context
-}
-
-func newContextAwareStdioTransport(ctx context.Context) *contextAwareStdioTransport {
-	return &contextAwareStdioTransport{ctx: ctx}
-}
-
-func (t *contextAwareStdioTransport) Connect(ctx context.Context) (mcp.Connection, error) {
-	// Create the original connection once
-	transport := mcp.NewStdioTransport()
-	originalConn, err := transport.Connect(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	return newContextAwareStdioConn(t.ctx, originalConn), nil
-}
-
-// contextAwareStdioConn wraps the original connection with context-aware reading
-type contextAwareStdioConn struct {
-	ctx          context.Context
-	originalConn mcp.Connection
-}
-
-func newContextAwareStdioConn(ctx context.Context, originalConn mcp.Connection) *contextAwareStdioConn {
-	return &contextAwareStdioConn{
-		ctx:          ctx,
-		originalConn: originalConn,
-	}
-}
-
-func (c *contextAwareStdioConn) SessionID() string {
-	return c.originalConn.SessionID()
-}
-
-func (c *contextAwareStdioConn) Read(ctx context.Context) (jsonrpc.Message, error) {
-	// Create a channel to read from the original connection in a separate goroutine
-	type result struct {
-		msg jsonrpc.Message
-		err error
-	}
-
-	ch := make(chan result, 1)
-	go func() {
-		msg, err := c.originalConn.Read(context.Background())
-		ch <- result{msg, err}
-	}()
-
-	// Wait for either context cancellation or read completion
-	select {
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	case <-c.ctx.Done():
-		return nil, c.ctx.Err()
-	case res := <-ch:
-		return res.msg, res.err
-	}
-}
-
-func (c *contextAwareStdioConn) Write(ctx context.Context, msg jsonrpc.Message) error {
-	// Check context first
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-c.ctx.Done():
-		return c.ctx.Err()
-	default:
-	}
-
-	return c.originalConn.Write(ctx, msg)
-}
-
-func (c *contextAwareStdioConn) Close() error {
-	return c.originalConn.Close()
-}