diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index ade30de..8e2f06d 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -3298,6 +3298,65 @@ func TestSubscriptionsListen_DisconnectScrubsMaps(t *testing.T) { } } +// TestServerSessionCloseWithActiveListen is a regression test for +// modelcontextprotocol/go-sdk#1160: ServerSession.Close must not deadlock +// when the client has an active subscriptions/listen stream. Previously, +// the server-side handler parked on ctx.Done and Close waited forever for +// the in-flight request to drain. +// +// The client's auto-listen (triggered by registering a list-changed handler) +// opens the stream on Connect — no explicit Subscribe is needed. The server +// must expose the corresponding list-changed capability, which happens +// automatically when at least one tool/prompt/resource is registered. +func TestServerSessionCloseWithActiveListen(t *testing.T) { + ctx := context.Background() + s := NewServer(&Implementation{Name: "s", Version: "0"}, nil) + AddTool(s, &Tool{Name: "t"}, sayHi) + + ct, st := NewInMemoryTransports() + if _, err := s.Connect(ctx, st, nil); err != nil { + t.Fatalf("server connect: %v", err) + } + c := NewClient(&Implementation{Name: "c", Version: "0"}, &ClientOptions{ + ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {}, + }) + cs, err := c.Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cs.Close() + + // Give the auto-listen request time to reach the server-side handler. + time.Sleep(20 * time.Millisecond) + + var ss *ServerSession + for x := range s.Sessions() { + ss = x + break + } + if ss == nil { + t.Fatal("no server session found") + } + + // Sanity check: the auto-listen must actually have registered an entry in + // listenIDs, otherwise the test below would trivially pass without + // exercising the fix. + ss.mu.Lock() + n := len(ss.listenIDs) + ss.mu.Unlock() + if n == 0 { + t.Fatal("expected auto-listen to register a request ID on the server session") + } + + done := make(chan error, 1) + go func() { done <- ss.Close() }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("ServerSession.Close deadlocked with an active subscriptions/listen") + } +} + func TestCustomMethods(t *testing.T) { type searchParams struct { ParamsBase diff --git a/mcp/server.go b/mcp/server.go index 4cdd6b9..13172a5 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1500,6 +1500,12 @@ type ServerSession struct { mu sync.Mutex state ServerSessionState + // listenIDs holds the request IDs of in-flight subscriptions/listen + // handlers on this session. subscriptions/listen parks on ctx.Done until + // cancelled, so ServerSession.Close must cancel these contexts explicitly + // via jsonrpc2.Connection.Cancel to avoid deadlocking on the jsonrpc2 + // drain. See modelcontextprotocol/go-sdk#1160. + listenIDs []jsonrpc.ID } func (ss *ServerSession) updateState(mut func(*ServerSessionState)) { @@ -1935,6 +1941,17 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, if validatedMeta.usesNewProtocol { ctx = context.WithValue(ctx, logLevelContextKey{}, validatedMeta.logLevel) } + + // subscriptions/listen parks on ctx.Done until the peer cancels it (or the + // underlying reader breaks). Track the request ID so ServerSession.Close + // can cancel the in-flight handler via jsonrpc2.Connection.Cancel and + // avoid deadlocking on the jsonrpc2 drain. + if req.Method == methodSubscriptionsListen { + ss.mu.Lock() + ss.listenIDs = append(ss.listenIDs, req.ID) + ss.mu.Unlock() + } + res, err := handleReceive(ctx, ss, req) if err != nil { return nil, err @@ -2038,6 +2055,18 @@ func (ss *ServerSession) Close() error { // Close is idempotent and conn.Close() handles concurrent calls correctly ss.keepaliveCancel() } + + // Unblock any in-flight subscriptions/listen handlers, which otherwise park + // on ctx.Done and would deadlock conn.Close (which waits for in-flight + // requests to drain). + ss.mu.Lock() + ids := ss.listenIDs + ss.listenIDs = nil + ss.mu.Unlock() + for _, id := range ids { + ss.conn.Cancel(id) + } + err := ss.conn.Close() if ss.onClose != nil && ss.calledOnClose.CompareAndSwap(false, true) {