d391ea8 · crush · git.iko.soy

internal/agent/tools/mcp/init.go+20−6

compare with main

unifiedsplit

@@ -38,8 +38,22 @@ func parseLevel(level mcp.LoggingLevel) slog.Level {

38 }

39}

40

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+

41var (

55var (

42− sessions = csync.NewMap[string, *mcp.ClientSession]()

56+ sessions = csync.NewMap[string, *ClientSession]()

43 states = csync.NewMap[string, ClientInfo]()

57 states = csync.NewMap[string, ClientInfo]()

44 broker = pubsub.NewBroker[Event]()

58 broker = pubsub.NewBroker[Event]()

45 initOnce sync.Once

59 initOnce sync.Once

@@ -102,7 +116,7 @@ type ClientInfo struct {

102 Name string

116 Name string

103 State State

117 State State

104 Error error

118 Error error

105− Client *mcp.ClientSession

119+ Client *ClientSession

106 Counts Counts

120 Counts Counts

107 ConnectedAt time.Time

121 ConnectedAt time.Time

108}

122}

@@ -239,7 +253,7 @@ func WaitForInit(ctx context.Context) error {

239 }

253 }

240}

254}

241

255

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) {

243 sess, ok := sessions.Get(name)

257 sess, ok := sessions.Get(name)

244 if !ok {

258 if !ok {

245 return nil, fmt.Errorf("mcp '%s' not available", name)

259 return nil, fmt.Errorf("mcp '%s' not available", name)

@@ -268,7 +282,7 @@ func getOrRenewClient(ctx context.Context, cfg *config.Config, name string) (*mc

268}

282}

269

283

270// updateState updates the state of an MCP client and publishes an event

284// 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) {

272 info := ClientInfo{

286 info := ClientInfo{

273 Name: name,

287 Name: name,

274 State: state,

288 State: state,

@@ -294,7 +308,7 @@ func updateState(name string, state State, err error, client *mcp.ClientSession,

294 })

308 })

295}

309}

296

310

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) {

298 timeout := mcpTimeout(m)

312 timeout := mcpTimeout(m)

299 mcpCtx, cancel := context.WithCancel(ctx)

313 mcpCtx, cancel := context.WithCancel(ctx)

300 cancelTimer := time.AfterFunc(timeout, cancel)

314 cancelTimer := time.AfterFunc(timeout, cancel)

@@ -352,7 +366,7 @@ func createSession(ctx context.Context, name string, m config.MCPConfig, resolve

352

366

353 cancelTimer.Stop()

367 cancelTimer.Stop()

354 slog.Debug("MCP client initialized", "name", name)

368 slog.Debug("MCP client initialized", "name", name)

355− return session, nil

369+ return &ClientSession{session, cancel}, nil

356}

370}

357

371

358// maybeStdioErr if a stdio mcp prints an error in non-json format, it'll fail

372// 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) {

67 updateState(name, StateConnected, nil, session, prev.Counts)

68}

69

70−func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) {

70+func getPrompts(ctx context.Context, c *ClientSession) ([]*Prompt, error) {

71 if c.InitializeResult().Capabilities.Prompts == nil {

72 return nil, nil

73 }

internal/agent/tools/mcp/resources.go+1−1

@@ -75,7 +75,7 @@ func RefreshResources(ctx context.Context, name string) {

75 updateState(name, StateConnected, nil, session, prev.Counts)

76}

77

78−func getResources(ctx context.Context, c *mcp.ClientSession) ([]*Resource, error) {

78+func getResources(ctx context.Context, c *ClientSession) ([]*Resource, error) {

79 if c.InitializeResult().Capabilities.Resources == nil {

80 return nil, nil

81 }

internal/agent/tools/mcp/tools.go+1−1

@@ -128,7 +128,7 @@ func RefreshTools(ctx context.Context, cfg *config.Config, name string) {

128 updateState(name, StateConnected, nil, session, prev.Counts)

129}

130

131−func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) {

131+func getTools(ctx context.Context, session *ClientSession) ([]*Tool, error) {

132 // Always call ListTools to get the actual available tools.

133 // The InitializeResult Capabilities.Tools field may be an empty object {},

134 // which is valid per MCP spec, but we still need to call ListTools to discover tools.