From c1d3db1ad360bf19aef92d74b9d9463d6febb878 Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Thu, 2 Apr 2026 18:42:03 +0000
Subject: [PATCH] fix(mcp): ensure CLI shutdown cleanup

---
 cmd/mcpcmd/cmd_mcp.go      | 15 +++++++++--
 cmd/mcpcmd/cmd_mcp_test.go | 52 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 cmd/mcpcmd/cmd_mcp_test.go

diff --git a/cmd/mcpcmd/cmd_mcp.go b/cmd/mcpcmd/cmd_mcp.go
index f097766..0a0fb5e 100644
--- a/cmd/mcpcmd/cmd_mcp.go
+++ b/cmd/mcpcmd/cmd_mcp.go
@@ -19,6 +19,14 @@ import (
 var workspaceFlag string
 var unrestrictedFlag bool
 
+var newMCPService = mcp.New
+var runMCPService = func(svc *mcp.Service, ctx context.Context) error {
+	return svc.Run(ctx)
+}
+var shutdownMCPService = func(svc *mcp.Service, ctx context.Context) error {
+	return svc.Shutdown(ctx)
+}
+
 var mcpCmd = &cli.Command{
 	Use:   "mcp",
 	Short: "MCP server for AI tool integration",
@@ -87,10 +95,13 @@ func runServe() error {
 	}
 
 	// Create the MCP service
-	svc, err := mcp.New(opts)
+	svc, err := newMCPService(opts)
 	if err != nil {
 		return cli.Wrap(err, "create MCP service")
 	}
+	defer func() {
+		_ = shutdownMCPService(svc, context.Background())
+	}()
 
 	// Set up signal handling for clean shutdown
 	ctx, cancel := context.WithCancel(context.Background())
@@ -105,5 +116,5 @@ func runServe() error {
 	}()
 
 	// Run the server (blocks until context cancelled or error)
-	return svc.Run(ctx)
+	return runMCPService(svc, ctx)
 }
diff --git a/cmd/mcpcmd/cmd_mcp_test.go b/cmd/mcpcmd/cmd_mcp_test.go
new file mode 100644
index 0000000..740e82a
--- /dev/null
+++ b/cmd/mcpcmd/cmd_mcp_test.go
@@ -0,0 +1,52 @@
+package mcpcmd
+
+import (
+	"context"
+	"testing"
+
+	"dappco.re/go/mcp/pkg/mcp"
+)
+
+func TestRunServe_Good_ShutsDownService(t *testing.T) {
+	oldNew := newMCPService
+	oldRun := runMCPService
+	oldShutdown := shutdownMCPService
+	oldWorkspace := workspaceFlag
+	oldUnrestricted := unrestrictedFlag
+
+	t.Cleanup(func() {
+		newMCPService = oldNew
+		runMCPService = oldRun
+		shutdownMCPService = oldShutdown
+		workspaceFlag = oldWorkspace
+		unrestrictedFlag = oldUnrestricted
+	})
+
+	workspaceFlag = ""
+	unrestrictedFlag = false
+
+	var runCalled bool
+	var shutdownCalled bool
+
+	newMCPService = func(opts mcp.Options) (*mcp.Service, error) {
+		return mcp.New(mcp.Options{})
+	}
+	runMCPService = func(svc *mcp.Service, ctx context.Context) error {
+		runCalled = true
+		return nil
+	}
+	shutdownMCPService = func(svc *mcp.Service, ctx context.Context) error {
+		shutdownCalled = true
+		return nil
+	}
+
+	if err := runServe(); err != nil {
+		t.Fatalf("runServe() returned error: %v", err)
+	}
+	if !runCalled {
+		t.Fatal("expected runMCPService to be called")
+	}
+	if !shutdownCalled {
+		t.Fatal("expected shutdownMCPService to be called")
+	}
+}