Skip to content

Commit 50680db

Browse files
peteski22peteski22
committed
Improve handling for closing clients (and logging) on shutdown
* Update mcp-go to v0.37.0 * Add logging for completed shutdowns * Add timeout for client closing * Refactor MCPClientAccessor to use the MCPClient interface rather than the Client implementation * Add tests for closing the daemon's MCP servers
1 parent 1bd3061 commit 50680db

File tree

9 files changed

+784
-27
lines changed

9 files changed

+784
-27
lines changed

‎cmd/daemon.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -134,9 +134,10 @@ func (c *DaemonCmd) run(_ *cobra.Command, _ []string) error {
134134

135135
select {
136136
case <-daemonCtx.Done():
137-
logger.Info("Shutting down daemon")
137+
logger.Info("Shutting down daemon...")
138138
err := <-runErr // Wait for cleanup and deferred logging.
139-
return err // Graceful Ctrl+C / SIGTERM.
139+
logger.Info("Shutdown complete")
140+
return err // Graceful Ctrl+C / SIGTERM.
140141
case err := <-runErr:
141142
logger.Error("daemon exited with error", "error", err)
142143
return err // Propagate daemon failure.

‎go.mod

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ require (
77
github.com/danielgtaylor/huma/v2 v2.34.1
88
github.com/go-chi/chi/v5 v5.2.2
99
github.com/hashicorp/go-hclog v1.6.3
10-
github.com/mark3labs/mcp-go v0.33.0
10+
github.com/mark3labs/mcp-go v0.37.0
1111
github.com/spf13/cobra v1.9.1
1212
github.com/spf13/pflag v1.0.6
1313
github.com/stretchr/testify v1.10.0
@@ -16,17 +16,22 @@ require (
1616
)
1717

1818
require (
19+
github.com/bahlo/generic-list-go v0.2.0 // indirect
20+
github.com/buger/jsonparser v1.1.1 // indirect
1921
github.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect
2022
github.com/davecgh/go-spew v1.1.1 // indirect
2123
github.com/fatih/color v1.18.0 // indirect
2224
github.com/google/uuid v1.6.0 // indirect
2325
github.com/inconshreveable/mousetrap v1.1.0 // indirect
26+
github.com/invopop/jsonschema v0.13.0 // indirect
27+
github.com/mailru/easyjson v0.7.7 // indirect
2428
github.com/mattn/go-colorable v0.1.14 // indirect
2529
github.com/mattn/go-isatty v0.0.20 // indirect
2630
github.com/pmezard/go-difflib v1.0.0 // indirect
2731
github.com/rogpeppe/go-internal v1.11.0 // indirect
2832
github.com/russross/blackfriday/v2 v2.1.0 // indirect
2933
github.com/spf13/cast v1.9.2 // indirect
34+
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
3035
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
3136
golang.org/x/sys v0.33.0 // indirect
3237
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect

‎go.sum

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg=
22
github.com/BurntSushi/toml v1.5.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
3+
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
4+
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
5+
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
6+
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
37
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
48
github.com/cpuguy83/go-md2man/v2 v2.0.7 h1:zbFlGlXEAKlwXpmvle3d8Oe3YnkKIK4xSRTd3sHPnBo=
59
github.com/cpuguy83/go-md2man/v2 v2.0.7/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
@@ -23,15 +27,20 @@ github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB1
2327
github.com/hashicorp/go-hclog v1.6.3/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M=
2428
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
2529
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
30+
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
31+
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
32+
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
2633
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
2734
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
2835
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
2936
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
3037
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
3138
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
3239
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
33-
github.com/mark3labs/mcp-go v0.33.0 h1:naxhjnTIs/tyPZmWUZFuG0lDmdA6sUyYGGf3gsHvTCc=
34-
github.com/mark3labs/mcp-go v0.33.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
40+
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
41+
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
42+
github.com/mark3labs/mcp-go v0.37.0 h1:BywvZLPRT6Zx6mMG/MJfxLSZQkTGIcJSEGKsvr4DsoQ=
43+
github.com/mark3labs/mcp-go v0.37.0/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g=
3544
github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
3645
github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
3746
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
@@ -56,6 +65,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
5665
github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals=
5766
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
5867
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
68+
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
69+
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
5970
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
6071
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
6172
golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=

‎internal/contracts/mcp.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@ type MCPHealthMonitor interface {
2323
// MCPClientAccessor provides a way to interact with MCP servers through a client.
2424
type MCPClientAccessor interface {
2525
// Add registers a client and its tools by server name.
26-
Add(name string, c *client.Client, tools []string)
26+
Add(name string, c client.MCPClient, tools []string)
2727

2828
// Client returns the client for the given server name.
2929
// It returns a boolean to indicate whether the client was found.
30-
Client(name string) (*client.Client, bool)
30+
Client(name string) (client.MCPClient, bool)
3131

3232
// Tools returns the tools for the given server name.
3333
// It returns a boolean to indicate whether the tools were found.

‎internal/daemon/client_manager.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,21 +10,21 @@ import (
1010
// It is safe for concurrent use by multiple goroutines.
1111
type ClientManager struct {
1212
mu sync.RWMutex
13-
clients map[string]*client.Client
13+
clients map[string]client.MCPClient
1414
serverTools map[string][]string
1515
}
1616

1717
// NewClientManager creates an empty, concurrency-safe ClientManager.
1818
func NewClientManager() *ClientManager {
1919
return &ClientManager{
20-
clients: make(map[string]*client.Client),
20+
clients: make(map[string]client.MCPClient),
2121
serverTools: make(map[string][]string),
2222
}
2323
}
2424

2525
// Add registers a client and its tools by server name.
2626
// This method is safe for concurrent use.
27-
func (cm *ClientManager) Add(name string, c *client.Client, tools []string) {
27+
func (cm *ClientManager) Add(name string, c client.MCPClient, tools []string) {
2828
cm.mu.Lock()
2929
defer cm.mu.Unlock()
3030
cm.clients[name] = c
@@ -34,7 +34,7 @@ func (cm *ClientManager) Add(name string, c *client.Client, tools []string) {
3434
// Client returns the client for the given server name.
3535
// It returns a boolean to indicate whether the client was found.
3636
// This method is safe for concurrent use.
37-
func (cm *ClientManager) Client(name string) (*client.Client, bool) {
37+
func (cm *ClientManager) Client(name string) (client.MCPClient, bool) {
3838
cm.mu.RLock()
3939
defer cm.mu.RUnlock()
4040
c, ok := cm.clients[name]

‎internal/daemon/client_manager_test.go

Lines changed: 108 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,121 @@
11
package daemon
22

33
import (
4+
"context"
45
"fmt"
56
"sync"
67
"testing"
78

8-
"github.com/mark3labs/mcp-go/client"
9+
"github.com/mark3labs/mcp-go/mcp"
910
"github.com/stretchr/testify/require"
1011
)
1112

13+
// mockMCPClient is a test implementation of client.MCPClient
14+
type mockMCPClient struct{}
15+
16+
func (m *mockMCPClient) Initialize(ctx context.Context, request mcp.InitializeRequest) (*mcp.InitializeResult, error) {
17+
return nil, nil
18+
}
19+
20+
func (m *mockMCPClient) Ping(ctx context.Context) error {
21+
return nil
22+
}
23+
24+
func (m *mockMCPClient) ListResourcesByPage(
25+
ctx context.Context,
26+
request mcp.ListResourcesRequest,
27+
) (*mcp.ListResourcesResult, error) {
28+
return nil, nil
29+
}
30+
31+
func (m *mockMCPClient) ListResources(
32+
ctx context.Context,
33+
request mcp.ListResourcesRequest,
34+
) (*mcp.ListResourcesResult, error) {
35+
return nil, nil
36+
}
37+
38+
func (m *mockMCPClient) ListResourceTemplatesByPage(
39+
ctx context.Context,
40+
request mcp.ListResourceTemplatesRequest,
41+
) (*mcp.ListResourceTemplatesResult, error) {
42+
return nil, nil
43+
}
44+
45+
func (m *mockMCPClient) ListResourceTemplates(
46+
ctx context.Context,
47+
request mcp.ListResourceTemplatesRequest,
48+
) (*mcp.ListResourceTemplatesResult, error) {
49+
return nil, nil
50+
}
51+
52+
func (m *mockMCPClient) ReadResource(
53+
ctx context.Context,
54+
request mcp.ReadResourceRequest,
55+
) (*mcp.ReadResourceResult, error) {
56+
return nil, nil
57+
}
58+
59+
func (m *mockMCPClient) Subscribe(ctx context.Context, request mcp.SubscribeRequest) error {
60+
return nil
61+
}
62+
63+
func (m *mockMCPClient) Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error {
64+
return nil
65+
}
66+
67+
func (m *mockMCPClient) ListPromptsByPage(
68+
ctx context.Context,
69+
request mcp.ListPromptsRequest,
70+
) (*mcp.ListPromptsResult, error) {
71+
return nil, nil
72+
}
73+
74+
func (m *mockMCPClient) ListPrompts(
75+
ctx context.Context,
76+
request mcp.ListPromptsRequest,
77+
) (*mcp.ListPromptsResult, error) {
78+
return nil, nil
79+
}
80+
81+
func (m *mockMCPClient) GetPrompt(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
82+
return nil, nil
83+
}
84+
85+
func (m *mockMCPClient) ListToolsByPage(
86+
ctx context.Context,
87+
request mcp.ListToolsRequest,
88+
) (*mcp.ListToolsResult, error) {
89+
return nil, nil
90+
}
91+
92+
func (m *mockMCPClient) ListTools(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error) {
93+
return nil, nil
94+
}
95+
96+
func (m *mockMCPClient) CallTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
97+
return nil, nil
98+
}
99+
100+
func (m *mockMCPClient) SetLevel(ctx context.Context, request mcp.SetLevelRequest) error {
101+
return nil
102+
}
103+
104+
func (m *mockMCPClient) Complete(ctx context.Context, request mcp.CompleteRequest) (*mcp.CompleteResult, error) {
105+
return nil, nil
106+
}
107+
108+
func (m *mockMCPClient) Close() error {
109+
return nil
110+
}
111+
112+
func (m *mockMCPClient) OnNotification(handler func(notification mcp.JSONRPCNotification)) {}
113+
12114
func TestClientManager_Add_Client_Tools(t *testing.T) {
13115
t.Parallel()
14116
cm := NewClientManager()
15117

16-
c := &client.Client{}
118+
c := &mockMCPClient{}
17119
tools := []string{"tool1", "tool2"}
18120
name := "server1"
19121

@@ -34,8 +136,8 @@ func TestClientManager_List(t *testing.T) {
34136
t.Parallel()
35137
cm := NewClientManager()
36138

37-
cm.Add("server1", &client.Client{}, []string{"a"})
38-
cm.Add("server2", &client.Client{}, []string{"b"})
139+
cm.Add("server1", &mockMCPClient{}, []string{"a"})
140+
cm.Add("server2", &mockMCPClient{}, []string{"b"})
39141

40142
names := cm.List()
41143
require.Len(t, names, 2)
@@ -46,7 +148,7 @@ func TestClientManager_Remove(t *testing.T) {
46148
t.Parallel()
47149
cm := NewClientManager()
48150

49-
cm.Add("server1", &client.Client{}, []string{"tool"})
151+
cm.Add("server1", &mockMCPClient{}, []string{"tool"})
50152
cm.Remove("server1")
51153

52154
_, ok := cm.Client("server1")
@@ -83,7 +185,7 @@ func TestClientManager_ConcurrentAccess(t *testing.T) {
83185
name := fmt.Sprintf("server-%d", i)
84186
go func() {
85187
defer wg.Done()
86-
cm.Add(name, &client.Client{}, []string{"tool"})
188+
cm.Add(name, &mockMCPClient{}, []string{"tool"})
87189
}()
88190
go func() {
89191
defer wg.Done()

‎internal/daemon/daemon.go

Lines changed: 63 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import (
2727
"github.com/mozilla-ai/mcpd/v2/internal/runtime"
2828
)
2929

30+
// Daemon manages MCP server lifecycles, client connections, and health monitoring.
31+
// It should only be created using NewDaemon to ensure proper initialization.
3032
type Daemon struct {
3133
apiServer *ApiServer
3234
logger hclog.Logger
@@ -60,6 +62,8 @@ func NewDaemonOpts(logger hclog.Logger, cfgLoader config.Loader, ctxLoader confi
6062
}, nil
6163
}
6264

65+
// NewDaemon creates a new Daemon instance with proper initialization.
66+
// Use this function instead of directly creating a Daemon struct.
6367
func NewDaemon(apiAddr string, opts *Opts) (*Daemon, error) {
6468
if err := IsValidAddr(apiAddr); err != nil {
6569
return nil, fmt.Errorf("invalid API address '%s': %w", apiAddr, err)
@@ -109,15 +113,7 @@ func NewDaemon(apiAddr string, opts *Opts) (*Daemon, error) {
109113
// It launches regular health checks on the MCP servers, with statuses visible via API routes.
110114
func (d *Daemon) StartAndManage(ctx context.Context) error {
111115
// Handle clean-up.
112-
defer func() {
113-
d.logger.Info("Shutting down MCP servers and client connections")
114-
for _, n := range d.clientManager.List() {
115-
if c, ok := d.clientManager.Client(n); ok {
116-
d.logger.Info(fmt.Sprintf("Closing client %s", n))
117-
_ = c.Close()
118-
}
119-
}
120-
}()
116+
defer d.closeAllClients()
121117

122118
// Launch servers
123119
if err := d.startMCPServers(ctx); err != nil {
@@ -430,3 +426,61 @@ func loadConfig(cfgLoader config.Loader, ctxLoader configcontext.Loader) ([]runt
430426

431427
return runtime.AggregateConfigs(cfg, execCtx)
432428
}
429+
430+
// closeAllClients gracefully closes all managed clients with individual timeouts.
431+
// It closes all clients concurrently and waits for all to complete or timeout.
432+
func (d *Daemon) closeAllClients() {
433+
d.logger.Info("Shutting down MCP servers and client connections")
434+
435+
clients := d.clientManager.List()
436+
if len(clients) == 0 {
437+
return
438+
}
439+
440+
var wg sync.WaitGroup
441+
timeout := 5 * time.Second
442+
443+
// Start closing all clients concurrently
444+
for _, n := range clients {
445+
name := n
446+
c, ok := d.clientManager.Client(name)
447+
if !ok {
448+
continue
449+
}
450+
451+
wg.Add(1)
452+
go func() {
453+
defer wg.Done()
454+
d.closeClientWithTimeout(name, c, timeout)
455+
}()
456+
}
457+
458+
wg.Wait()
459+
}
460+
461+
// closeClientWithTimeout closes a single client with a timeout.
462+
func (d *Daemon) closeClientWithTimeout(name string, c client.MCPClient, timeout time.Duration) {
463+
d.logger.Info(fmt.Sprintf("Closing client %s", name))
464+
465+
done := make(chan struct{})
466+
go func() {
467+
err := c.Close()
468+
if err != nil {
469+
// 'errors' can result in things like SIGINT which returns exit code 130,
470+
// we still log the error but only for debugging purposes.
471+
d.logger.Debug("Closing client", "client", name, "error", err)
472+
}
473+
d.logger.Info(fmt.Sprintf("Closed client %s", name))
474+
close(done)
475+
}()
476+
477+
// Wait for this specific client to close or timeout.
478+
// NOTE: this could leak if we just time out clients,
479+
// but since we're exiting mcpd it isn't an issue.
480+
select {
481+
case <-done:
482+
// Closed successfully.
483+
case <-time.After(timeout):
484+
d.logger.Warn(fmt.Sprintf("Timeout (%s) closing client %s", timeout.String(), name))
485+
}
486+
}

0 commit comments

Comments
 (0)