Fix daemon shutdown deadlock and improve MCP transport observability (#206)

peteski22 · web-flow · commit 3a0e58973951 · 2025-10-09T09:58:50.000+01:00
* Bump mcp-go to v0.41.1 + go mod tidy

* Clean up temp error handling for 'method not found' on MCP servers

* Supply logging adaptor to MCP library to handle logging errors from servers

* Prevent hanging during pinging all servers when trying to quit (CTRL+C)
diff --git a/go.mod b/go.mod
@@ -8,7 +8,7 @@ require (
 	github.com/go-chi/chi/v5 v5.2.3
 	github.com/go-chi/cors v1.2.2
 	github.com/hashicorp/go-hclog v1.6.3
-	github.com/mark3labs/mcp-go v0.39.1
+	github.com/mark3labs/mcp-go v0.41.1
 	github.com/spf13/cobra v1.10.1
 	github.com/spf13/pflag v1.0.10
 	github.com/stretchr/testify v1.11.1
diff --git a/go.sum b/go.sum
@@ -40,8 +40,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/mailru/easyjson v0.9.0 h1:PrnmzHw7262yW8sTBwxi1PdJA3Iw/EKBa8psRf7d9a4=
 github.com/mailru/easyjson v0.9.0/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=
-github.com/mark3labs/mcp-go v0.39.1 h1:2oPxk7aDbQhouakkYyKl2T4hKFU1c6FDaubWyGyVE1k=
-github.com/mark3labs/mcp-go v0.39.1/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g=
+github.com/mark3labs/mcp-go v0.41.1 h1:w78eWfiQam2i8ICL7AL0WFiq7KHNJQ6UB53ZVtH4KGA=
+github.com/mark3labs/mcp-go v0.41.1/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g=
 github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
 github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
 github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
diff --git a/internal/api/mcp.go b/internal/api/mcp.go
@@ -2,12 +2,6 @@ package api
 
 import "github.com/mark3labs/mcp-go/mcp"
 
-// methodNotFoundMessage is the error message returned by MCP servers when a method is not implemented.
-// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-// See: https://github.com/mark3labs/mcp-go/issues/593
-const methodNotFoundMessage = "Method not found"
-
 // DomainMeta wraps mcp.Meta for API conversion.
 type DomainMeta mcp.Meta
 
diff --git a/internal/api/prompts.go b/internal/api/prompts.go
@@ -2,15 +2,15 @@ package api
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/danielgtaylor/huma/v2"
 	"github.com/mark3labs/mcp-go/mcp"
 
 	"github.com/mozilla-ai/mcpd/v2/internal/contracts"
-	"github.com/mozilla-ai/mcpd/v2/internal/errors"
+	errorsint "github.com/mozilla-ai/mcpd/v2/internal/errors"
 )
 
 // DomainPrompt wraps mcp.Prompt for API conversion.
@@ -147,7 +147,7 @@ func handleServerPrompts(
 ) (*PromptsListResponse, error) {
 	mcpClient, clientOk := accessor.Client(name)
 	if !clientOk {
-		return nil, fmt.Errorf("%w: %s", errors.ErrServerNotFound, name)
+		return nil, fmt.Errorf("%w: %s", errorsint.ErrServerNotFound, name)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -162,16 +162,13 @@ func handleServerPrompts(
 
 	result, err := mcpClient.ListPrompts(ctx, req)
 	if err != nil {
-		// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-		// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-		// See: https://github.com/mark3labs/mcp-go/issues/593
-		if strings.Contains(err.Error(), methodNotFoundMessage) {
-			return nil, fmt.Errorf("%w: %s", errors.ErrPromptsNotImplemented, name)
+		if errors.Is(err, mcp.ErrMethodNotFound) {
+			return nil, fmt.Errorf("%w: %s", errorsint.ErrPromptsNotImplemented, name)
 		}
-		return nil, fmt.Errorf("%w: %s: %w", errors.ErrPromptListFailed, name, err)
+		return nil, fmt.Errorf("%w: %s: %w", errorsint.ErrPromptListFailed, name, err)
 	}
 	if result == nil {
-		return nil, fmt.Errorf("%w: %s: no result", errors.ErrPromptListFailed, name)
+		return nil, fmt.Errorf("%w: %s: no result", errorsint.ErrPromptListFailed, name)
 	}
 
 	prompts := make([]Prompt, 0, len(result.Prompts))
@@ -201,7 +198,7 @@ func handleServerPromptGenerate(
 ) (*GeneratePromptResponse, error) {
 	mcpClient, clientOk := accessor.Client(serverName)
 	if !clientOk {
-		return nil, fmt.Errorf("%w: %s", errors.ErrServerNotFound, serverName)
+		return nil, fmt.Errorf("%w: %s", errorsint.ErrServerNotFound, serverName)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -214,16 +211,13 @@ func handleServerPromptGenerate(
 		},
 	})
 	if err != nil {
-		// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-		// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-		// See: https://github.com/mark3labs/mcp-go/issues/593
-		if strings.Contains(err.Error(), methodNotFoundMessage) {
-			return nil, fmt.Errorf("%w: %s", errors.ErrPromptsNotImplemented, serverName)
+		if errors.Is(err, mcp.ErrMethodNotFound) {
+			return nil, fmt.Errorf("%w: %s", errorsint.ErrPromptsNotImplemented, serverName)
 		}
-		return nil, fmt.Errorf("%w: %s: %s: %w", errors.ErrPromptGenerationFailed, serverName, promptName, err)
+		return nil, fmt.Errorf("%w: %s: %s: %w", errorsint.ErrPromptGenerationFailed, serverName, promptName, err)
 	}
 	if result == nil {
-		return nil, fmt.Errorf("%w: %s: %s: no result", errors.ErrPromptGenerationFailed, serverName, promptName)
+		return nil, fmt.Errorf("%w: %s: %s: no result", errorsint.ErrPromptGenerationFailed, serverName, promptName)
 	}
 
 	messages := make([]PromptMessage, 0, len(result.Messages))
diff --git a/internal/api/prompts_test.go b/internal/api/prompts_test.go
@@ -171,7 +171,7 @@ func TestAPI_HandleServerPrompts_MethodNotFound(t *testing.T) {
 	t.Parallel()
 
 	mockClient := &mockMCPClient{
-		listPromptsError: errors.New("Method not found"),
+		listPromptsError: mcp.ErrMethodNotFound,
 	}
 
 	accessor := newMockMCPClientAccessor()
@@ -340,7 +340,7 @@ func TestAPI_HandleServerPromptGenerate_MethodNotFound(t *testing.T) {
 	t.Parallel()
 
 	mockClient := &mockMCPClient{
-		getPromptError: errors.New("Method not found"),
+		getPromptError: mcp.ErrMethodNotFound,
 	}
 
 	accessor := newMockMCPClientAccessor()
diff --git a/internal/api/resources.go b/internal/api/resources.go
@@ -2,15 +2,15 @@ package api
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/danielgtaylor/huma/v2"
 	"github.com/mark3labs/mcp-go/mcp"
 
 	"github.com/mozilla-ai/mcpd/v2/internal/contracts"
-	"github.com/mozilla-ai/mcpd/v2/internal/errors"
+	errorsint "github.com/mozilla-ai/mcpd/v2/internal/errors"
 )
 
 // DomainResource wraps mcp.Resource for API conversion.
@@ -171,7 +171,7 @@ func handleServerResources(
 ) (*ResourcesResponse, error) {
 	mcpClient, clientOk := accessor.Client(name)
 	if !clientOk {
-		return nil, fmt.Errorf("%w: %s", errors.ErrServerNotFound, name)
+		return nil, fmt.Errorf("%w: %s", errorsint.ErrServerNotFound, name)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -186,16 +186,13 @@ func handleServerResources(
 
 	result, err := mcpClient.ListResources(ctx, req)
 	if err != nil {
-		// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-		// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-		// See: https://github.com/mark3labs/mcp-go/issues/593
-		if strings.Contains(err.Error(), methodNotFoundMessage) {
-			return nil, fmt.Errorf("%w: %s", errors.ErrResourcesNotImplemented, name)
+		if errors.Is(err, mcp.ErrMethodNotFound) {
+			return nil, fmt.Errorf("%w: %s", errorsint.ErrResourcesNotImplemented, name)
 		}
-		return nil, fmt.Errorf("%w: %s: %w", errors.ErrResourceListFailed, name, err)
+		return nil, fmt.Errorf("%w: %s: %w", errorsint.ErrResourceListFailed, name, err)
 	}
 	if result == nil {
-		return nil, fmt.Errorf("%w: %s: no result", errors.ErrResourceListFailed, name)
+		return nil, fmt.Errorf("%w: %s: no result", errorsint.ErrResourceListFailed, name)
 	}
 
 	resources := make([]Resource, 0, len(result.Resources))
@@ -224,7 +221,7 @@ func handleServerResourceTemplates(
 ) (*ResourceTemplatesResponse, error) {
 	mcpClient, clientOk := accessor.Client(name)
 	if !clientOk {
-		return nil, fmt.Errorf("%w: %s", errors.ErrServerNotFound, name)
+		return nil, fmt.Errorf("%w: %s", errorsint.ErrServerNotFound, name)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -239,16 +236,13 @@ func handleServerResourceTemplates(
 
 	result, err := mcpClient.ListResourceTemplates(ctx, req)
 	if err != nil {
-		// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-		// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-		// See: https://github.com/mark3labs/mcp-go/issues/593
-		if strings.Contains(err.Error(), methodNotFoundMessage) {
-			return nil, fmt.Errorf("%w: %s", errors.ErrResourcesNotImplemented, name)
+		if errors.Is(err, mcp.ErrMethodNotFound) {
+			return nil, fmt.Errorf("%w: %s", errorsint.ErrResourcesNotImplemented, name)
 		}
-		return nil, fmt.Errorf("%w: %s: %w", errors.ErrResourceTemplateListFailed, name, err)
+		return nil, fmt.Errorf("%w: %s: %w", errorsint.ErrResourceTemplateListFailed, name, err)
 	}
 	if result == nil {
-		return nil, fmt.Errorf("%w: %s: no result", errors.ErrResourceTemplateListFailed, name)
+		return nil, fmt.Errorf("%w: %s: no result", errorsint.ErrResourceTemplateListFailed, name)
 	}
 
 	templates := make([]ResourceTemplate, 0, len(result.ResourceTemplates))
@@ -277,7 +271,7 @@ func handleServerResourceContent(
 ) (*ResourceContentResponse, error) {
 	mcpClient, clientOk := accessor.Client(name)
 	if !clientOk {
-		return nil, fmt.Errorf("%w: %s", errors.ErrServerNotFound, name)
+		return nil, fmt.Errorf("%w: %s", errorsint.ErrServerNotFound, name)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
@@ -289,16 +283,13 @@ func handleServerResourceContent(
 		},
 	})
 	if err != nil {
-		// TODO: This string matching is fragile and should be replaced with proper JSON-RPC error code checking.
-		// Once mcp-go preserves JSON-RPC error codes, use errors.Is(err, mcp.ErrMethodNotFound) instead.
-		// See: https://github.com/mark3labs/mcp-go/issues/593
-		if strings.Contains(err.Error(), methodNotFoundMessage) {
-			return nil, fmt.Errorf("%w: %s", errors.ErrResourcesNotImplemented, name)
+		if errors.Is(err, mcp.ErrMethodNotFound) {
+			return nil, fmt.Errorf("%w: %s", errorsint.ErrResourcesNotImplemented, name)
 		}
-		return nil, fmt.Errorf("%w: %s: %s: %w", errors.ErrResourceReadFailed, name, uri, err)
+		return nil, fmt.Errorf("%w: %s: %s: %w", errorsint.ErrResourceReadFailed, name, uri, err)
 	}
 	if result == nil {
-		return nil, fmt.Errorf("%w: %s: %s: no result", errors.ErrResourceReadFailed, name, uri)
+		return nil, fmt.Errorf("%w: %s: %s: no result", errorsint.ErrResourceReadFailed, name, uri)
 	}
 
 	contents := make([]ResourceContent, 0, len(result.Contents))
diff --git a/internal/daemon/daemon.go b/internal/daemon/daemon.go
@@ -14,7 +14,9 @@ import (
 
 	"github.com/hashicorp/go-hclog"
 	"github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/client/transport"
 	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/util"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/mozilla-ai/mcpd/v2/internal/cmd"
@@ -24,6 +26,8 @@ import (
 	"github.com/mozilla-ai/mcpd/v2/internal/runtime"
 )
 
+var _ util.Logger = (*mcpLoggerAdapter)(nil)
+
 // Daemon manages MCP server lifecycles, client connections, and health monitoring.
 // It should only be created using NewDaemon to ensure proper initialization.
 type Daemon struct {
@@ -47,6 +51,11 @@ type Daemon struct {
 	clientHealthCheckInterval time.Duration
 }
 
+// mcpLoggerAdapter adapts hclog.Logger to mcp-go's util.Logger interface.
+type mcpLoggerAdapter struct {
+	logger hclog.Logger
+}
+
 // NewDaemon creates a new Daemon instance with proper initialization.
 // Use this function instead of directly creating a Daemon struct.
 func NewDaemon(deps Dependencies, opt ...Option) (*Daemon, error) {
@@ -108,6 +117,12 @@ func NewDaemon(deps Dependencies, opt ...Option) (*Daemon, error) {
 	}, nil
 }
 
+func newMCPLoggerAdapter(logger hclog.Logger) *mcpLoggerAdapter {
+	return &mcpLoggerAdapter{
+		logger: logger,
+	}
+}
+
 // StartAndManage is a long-running method that starts configured MCP servers, and the API.
 // It launches regular health checks on the MCP servers, with statuses visible via API routes.
 func (d *Daemon) StartAndManage(ctx context.Context) error {
@@ -222,7 +237,13 @@ func (d *Daemon) startMCPServer(ctx context.Context, server runtime.Server) erro
 
 	logger.Debug("attempting to start server", "binary", runtimeBinary)
 
-	stdioClient, err := client.NewStdioMCPClient(runtimeBinary, environ, args...)
+	mcpLogger := newMCPLoggerAdapter(logger.Named("transport"))
+	stdioClient, err := client.NewStdioMCPClientWithOptions(
+		runtimeBinary,
+		environ,
+		args,
+		transport.WithCommandLogger(mcpLogger),
+	)
 	if err != nil {
 		return fmt.Errorf("error starting MCP server: '%s': %w", server.Name(), err)
 	}
@@ -396,12 +417,37 @@ func (d *Daemon) pingAllServers(ctx context.Context, maxTimeout time.Duration) e
 		})
 	}
 
-	_ = g.Wait()
-	if len(errs) > 0 {
-		return errors.Join(errs...)
+	// Wait for all pings to complete, but allow interruption if parent context is cancelled.
+	// This prevents the daemon from hanging during shutdown if a ping is stuck in uninterruptible I/O.
+	done := make(chan struct{})
+	go func() {
+		_ = g.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// All pings completed normally.
+		if len(errs) > 0 {
+			return errors.Join(errs...)
+		}
+		return nil
+	case <-ctx.Done():
+		// Parent context cancelled (shutdown), return immediately without waiting for pings to complete.
+		// Any stuck pings will eventually time out or be cleaned up when the process exits.
+		d.logger.Warn("Ping operation interrupted due to context cancellation, some pings may not have completed")
+		return ctx.Err()
 	}
+}
 
-	return nil
+// Infof implements mcp-go's Logger interface.
+func (a *mcpLoggerAdapter) Infof(format string, v ...any) {
+	a.logger.Info(fmt.Sprintf(format, v...))
+}
+
+// Errorf implements mcp-go's Logger interface.
+func (a *mcpLoggerAdapter) Errorf(format string, v ...any) {
+	a.logger.Error(fmt.Sprintf(format, v...))
 }
 
 // IsValidAddr returns an error if the address is not a valid "host:port" string.
diff --git a/internal/daemon/daemon_test.go b/internal/daemon/daemon_test.go

Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ func TestAPI_HandleServerPrompts_MethodNotFound(t *testing.T) {`
`171`	`171`	`t.Parallel()`
`172`	`172`
`173`	`173`	`mockClient := &mockMCPClient{`
`174`		`- listPromptsError: errors.New("Method not found"),`
	`174`	`+ listPromptsError: mcp.ErrMethodNotFound,`
`175`	`175`	`}`
`176`	`176`
`177`	`177`	`accessor := newMockMCPClientAccessor()`
`@@ -340,7 +340,7 @@ func TestAPI_HandleServerPromptGenerate_MethodNotFound(t *testing.T) {`
`340`	`340`	`t.Parallel()`
`341`	`341`
`342`	`342`	`mockClient := &mockMCPClient{`
`343`		`- getPromptError: errors.New("Method not found"),`
	`343`	`+ getPromptError: mcp.ErrMethodNotFound,`
`344`	`344`	`}`
`345`	`345`
`346`	`346`	`accessor := newMockMCPClientAccessor()`