git.iko.soy/crush
on main help sign in
fix(mcp): cancel context on MCP session close to prevent leak (#2157)

* fix(mcp): cancel context on MCP session close to prevent leak * refactor: rename unexported mcpSession to exported ClientSession

commitd391ea85ed2f43d8e1290a2c64269c0782a6185a authorM1xA <AnyCPU@users.noreply.github.com> dateTue 10 Feb 2026, 14:56 (8 months ago) parent8c9eb47
5 files changed+61−9
compare with main
unifiedsplit
internal/agent/tools/mcp/init.go+20−6 internal/agent/tools/mcp/init_test.go+38 internal/agent/tools/mcp/prompts.go+1−1 internal/agent/tools/mcp/resources.go+1−1 internal/agent/tools/mcp/tools.go+1−1
internal/agent/tools/mcp/init.go+20−6
@@ -38,8 +38,22 @@ func parseLevel(level mcp.LoggingLevel) slog.Level {
3838 }
3939}
4040
41+// ClientSession wraps an mcp.ClientSession with a context cancel function so
42+// that the context created during session establishment is properly cleaned up
43+// on close.
44+type ClientSession struct {
45+ *mcp.ClientSession
46+ cancel context.CancelFunc
47+}
48+
49+// Close cancels the session context and then closes the underlying session.
50+func (s *ClientSession) Close() error {
51+ s.cancel()
52+ return s.ClientSession.Close()
53+}
54+
4155var (
42− sessions = csync.NewMap[string, *mcp.ClientSession]()
56+ sessions = csync.NewMap[string, *ClientSession]()
4357 states = csync.NewMap[string, ClientInfo]()
4458 broker = pubsub.NewBroker[Event]()
4559 initOnce sync.Once
@@ -102,7 +116,7 @@ type ClientInfo struct {
102116 Name string
103117 State State
104118 Error error
105− Client *mcp.ClientSession
119+ Client *ClientSession
106120 Counts Counts
107121 ConnectedAt time.Time
108122}
@@ -239,7 +253,7 @@ func WaitForInit(ctx context.Context) error {
239253 }
240254}
241255
242−func getOrRenewClient(ctx context.Context, cfg *config.Config, name string) (*mcp.ClientSession, error) {
256+func getOrRenewClient(ctx context.Context, cfg *config.Config, name string) (*ClientSession, error) {
243257 sess, ok := sessions.Get(name)
244258 if !ok {
245259 return nil, fmt.Errorf("mcp '%s' not available", name)
@@ -268,7 +282,7 @@ func getOrRenewClient(ctx context.Context, cfg *config.Config, name string) (*mc
268282}
269283
270284// updateState updates the state of an MCP client and publishes an event
271−func updateState(name string, state State, err error, client *mcp.ClientSession, counts Counts) {
285+func updateState(name string, state State, err error, client *ClientSession, counts Counts) {
272286 info := ClientInfo{
273287 Name: name,
274288 State: state,
@@ -294,7 +308,7 @@ func updateState(name string, state State, err error, client *mcp.ClientSession,
294308 })
295309}
296310
297−func createSession(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*mcp.ClientSession, error) {
311+func createSession(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*ClientSession, error) {
298312 timeout := mcpTimeout(m)
299313 mcpCtx, cancel := context.WithCancel(ctx)
300314 cancelTimer := time.AfterFunc(timeout, cancel)
@@ -352,7 +366,7 @@ func createSession(ctx context.Context, name string, m config.MCPConfig, resolve
352366
353367 cancelTimer.Stop()
354368 slog.Debug("MCP client initialized", "name", name)
355− return session, nil
369+ return &ClientSession{session, cancel}, nil
356370}
357371
358372// maybeStdioErr if a stdio mcp prints an error in non-json format, it'll fail
internal/agent/tools/mcp/init_test.gonew file+38
@@ -0,0 +1,38 @@
1+package mcp
2+
3+import (
4+ "context"
5+ "testing"
6+
7+ "github.com/modelcontextprotocol/go-sdk/mcp"
8+ "github.com/stretchr/testify/require"
9+ "go.uber.org/goleak"
10+)
11+
12+func TestMCPSession_CancelOnClose(t *testing.T) {
13+ defer goleak.VerifyNone(t)
14+
15+ serverTransport, clientTransport := mcp.NewInMemoryTransports()
16+
17+ server := mcp.NewServer(&mcp.Implementation{Name: "test-server"}, nil)
18+ serverSession, err := server.Connect(context.Background(), serverTransport, nil)
19+ require.NoError(t, err)
20+ defer serverSession.Close()
21+
22+ ctx, cancel := context.WithCancel(context.Background())
23+
24+ client := mcp.NewClient(&mcp.Implementation{Name: "crush-test"}, nil)
25+ clientSession, err := client.Connect(ctx, clientTransport, nil)
26+ require.NoError(t, err)
27+
28+ sess := &ClientSession{clientSession, cancel}
29+
30+ // Verify the context is not cancelled before close.
31+ require.NoError(t, ctx.Err())
32+
33+ err = sess.Close()
34+ require.NoError(t, err)
35+
36+ // After Close, the context must be cancelled.
37+ require.ErrorIs(t, ctx.Err(), context.Canceled)
38+}
internal/agent/tools/mcp/prompts.go+1−1
@@ -67,7 +67,7 @@ func RefreshPrompts(ctx context.Context, name string) {
6767 updateState(name, StateConnected, nil, session, prev.Counts)
6868}
6969
70−func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) {
70+func getPrompts(ctx context.Context, c *ClientSession) ([]*Prompt, error) {
7171 if c.InitializeResult().Capabilities.Prompts == nil {
7272 return nil, nil
7373 }
internal/agent/tools/mcp/resources.go+1−1
@@ -75,7 +75,7 @@ func RefreshResources(ctx context.Context, name string) {
7575 updateState(name, StateConnected, nil, session, prev.Counts)
7676}
7777
78−func getResources(ctx context.Context, c *mcp.ClientSession) ([]*Resource, error) {
78+func getResources(ctx context.Context, c *ClientSession) ([]*Resource, error) {
7979 if c.InitializeResult().Capabilities.Resources == nil {
8080 return nil, nil
8181 }
internal/agent/tools/mcp/tools.go+1−1
@@ -128,7 +128,7 @@ func RefreshTools(ctx context.Context, cfg *config.Config, name string) {
128128 updateState(name, StateConnected, nil, session, prev.Counts)
129129}
130130
131−func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) {
131+func getTools(ctx context.Context, session *ClientSession) ([]*Tool, error) {
132132 // Always call ListTools to get the actual available tools.
133133 // The InitializeResult Capabilities.Tools field may be an empty object {},
134134 // which is valid per MCP spec, but we still need to call ListTools to discover tools.