1package mcp
 2
 3import (
 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
12func 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}