diff --git a/mcp/client.go b/mcp/client.go index 9536bdd4..44be0ef6 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -354,10 +354,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp // ClientSession.Close cancels the listenCtx context to send notifications/cancelled. listenCtx, cancelListen := context.WithCancel(context.Background()) cs.listenCancel = cancelListen - if err := cs.subscriptionsListen(listenCtx, subscribeParams); err != nil { - cancelListen() - return nil, fmt.Errorf("opening subscriptions/listen: %w", err) - } + go cs.awaitSubscriptionsListen(listenCtx, subscribeParams, nil) } return cs, nil } @@ -496,14 +493,28 @@ type ClientSession struct { pendingElicitationsMu sync.Mutex pendingElicitations map[string]chan struct{} - // resourceSubsMu guards resourceSubs. + // resourceSubsMu guards resourceSubs and nextResourceSubGen. resourceSubsMu sync.Mutex - // resourceSubs maps a subscribed resource URI to the cancel func of the + // resourceSubs maps a subscribed resource URI to the state of the // goroutine running its dedicated subscriptions/listen stream. Populated // only under SEP-2575; the legacy protocol routes Subscribe and // Unsubscribe straight to the resources/subscribe and resources/unsubscribe // RPCs and leaves this map untouched. - resourceSubs map[string]context.CancelFunc + resourceSubs map[string]*resourceSub + // nextResourceSubGen assigns a monotonically increasing generation to each + // resourceSubs entry, so a listen goroutine only clears the entry it + // created and not one installed by a later Subscribe for the same URI. + nextResourceSubGen uint64 +} + +// resourceSub is the per-URI state of a SEP-2575 resource subscription: the +// cancel func for its subscriptions/listen stream, and the generation that +// lets a completing listen goroutine tell its own entry from one a racing +// Unsubscribe→Subscribe (or a re-subscribe inside the callback) has since +// installed for the same URI. +type resourceSub struct { + cancel context.CancelFunc + gen uint64 } type clientSessionState struct { @@ -1414,27 +1425,73 @@ func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) } uri := params.URI - var listenCtx context.Context cs.resourceSubsMu.Lock() - if _, exists := cs.resourceSubs[uri]; !exists { - var cancel context.CancelFunc - listenCtx, cancel = context.WithCancel(context.Background()) - if cs.resourceSubs == nil { - cs.resourceSubs = make(map[string]context.CancelFunc) - } - cs.resourceSubs[uri] = cancel - } - cs.resourceSubsMu.Unlock() - if listenCtx == nil { - // Already subscribed to this URI + if _, exists := cs.resourceSubs[uri]; exists { + // Already subscribed to this URI. + cs.resourceSubsMu.Unlock() return nil } + if cs.resourceSubs == nil { + cs.resourceSubs = make(map[string]*resourceSub) + } + listenCtx, cancel := context.WithCancel(context.Background()) + cs.nextResourceSubGen++ + gen := cs.nextResourceSubGen + cs.resourceSubs[uri] = &resourceSub{cancel: cancel, gen: gen} + cs.resourceSubsMu.Unlock() - return cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ + // Subscribe stays non-blocking: the listen stream is awaited on its own + // goroutine so that its completion clears the subscription. + subscribeParams := &SubscriptionsListenParams{ Notifications: &NotificationSubscriptions{ ResourceSubscriptions: []string{uri}, }, - }) + } + go cs.awaitSubscriptionsListen(listenCtx, subscribeParams, &gen) + return nil +} + +// awaitSubscriptionsListen runs a subscriptions/listen stream to completion. Connect and +// Subscribe start it on its own goroutine, so they return without waiting for +// the server to acknowledge the listen. +// +// When the stream ends for a reason other than a client-initiated Unsubscribe +// or session Close — a graceful listen result, a synthetic transport +// "terminated" error, or any jsonrpc error, all while listenCtx is not +// cancelled — the resourceSubs entries for the stream's resource URIs are +// cleared so that a later bare Subscribe re-opens the stream instead of +// no-oping. An entry is only removed if it still carries generation gen, which +// guards an Unsubscribe→Subscribe race and a re-subscribe from inside a +// callback. The SDK does not auto-resubscribe: a revoked URI would hot-loop, so +// reopening is left to the application calling Subscribe again. Connect passes +// a nil gen for its list-changed listen, which owns no resource entries, so +// its end clears nothing. +func (cs *ClientSession) awaitSubscriptionsListen(listenCtx context.Context, params *SubscriptionsListenParams, gen *uint64) { + _ = cs.subscriptionsListen(listenCtx, params) + if listenCtx.Err() != nil { + // Client-initiated teardown: Unsubscribe already removed the entry (and + // cancelAllResourceSubscriptions nils the whole map on Close), so there + // is nothing to clear. + return + } + if gen == nil || params.Notifications == nil { + return + } + for _, uri := range params.Notifications.ResourceSubscriptions { + cs.clearResourceSubIfGen(uri, *gen) + } +} + +// clearResourceSubIfGen deletes the resourceSubs entry for uri only if it is +// still present and carries the given generation, reporting whether it did. +func (cs *ClientSession) clearResourceSubIfGen(uri string, gen uint64) bool { + cs.resourceSubsMu.Lock() + defer cs.resourceSubsMu.Unlock() + if sub, ok := cs.resourceSubs[uri]; ok && sub.gen == gen { + delete(cs.resourceSubs, uri) + return true + } + return false } // Unsubscribe cancels a previous [ClientSession.Subscribe] for params.URI. @@ -1453,11 +1510,11 @@ func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribePar return fmt.Errorf("Unsubscribe: missing URI") } cs.resourceSubsMu.Lock() - cancel, ok := cs.resourceSubs[params.URI] + sub, ok := cs.resourceSubs[params.URI] delete(cs.resourceSubs, params.URI) cs.resourceSubsMu.Unlock() if ok { - cancel() + sub.cancel() } return nil } @@ -1470,8 +1527,8 @@ func (cs *ClientSession) cancelAllResourceSubscriptions() { subs := cs.resourceSubs cs.resourceSubs = nil cs.resourceSubsMu.Unlock() - for _, cancel := range subs { - cancel() + for _, sub := range subs { + sub.cancel() } } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 11486fea..6fb2d2c7 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -3602,11 +3602,9 @@ func TestSubscriptionsListen_TeardownRetiresOwnRegistration(t *testing.T) { listenCtx, cancelListen := context.WithCancel(context.Background()) defer cancelListen() - if err := cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ + go cs.subscriptionsListen(listenCtx, &SubscriptionsListenParams{ Notifications: &NotificationSubscriptions{ToolsListChanged: true, PromptsListChanged: true}, - }); err != nil { - t.Fatalf("listen: %v", err) - } + }) waitSubListenEvent(t, events, "ack") server.mu.Lock() @@ -4073,3 +4071,205 @@ func newProtocolMeta(logLevel LoggingLevel) Meta { } return m } + +// listenDropServer wraps a streamable backend handler. While armed, the first +// subscriptions/listen POST whose resource subscription list contains uri is +// served by the wrapper itself rather than the backend, then the wrapper +// disarms — so a later re-Subscribe proxies through to the real backend and +// re-fires its SubscribeHandler. An abrupt listen closes its SSE stream with +// no JSON-RPC response (the client then synthesizes a "terminated" error); a +// graceful listen writes an empty result before closing. +type listenDropServer struct { + backend http.Handler + uri string + mu sync.Mutex + armed bool + graceful bool +} + +func (d *listenDropServer) arm(graceful bool) { + d.mu.Lock() + d.armed = true + d.graceful = graceful + d.mu.Unlock() +} + +func (d *listenDropServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + body, err := io.ReadAll(r.Body) + r.Body.Close() + if err == nil { + r.Body = io.NopCloser(bytes.NewReader(body)) + if id, graceful, ok := d.matchArmed(body); ok { + serveControlledListen(w, id, graceful) + return + } + } + } + d.backend.ServeHTTP(w, r) +} + +func (d *listenDropServer) matchArmed(body []byte) (id json.RawMessage, graceful, ok bool) { + var env struct { + Method string `json:"method"` + ID json.RawMessage `json:"id"` + Params struct { + Notifications struct { + ResourceSubscriptions []string `json:"resourceSubscriptions"` + } `json:"notifications"` + } `json:"params"` + } + if json.Unmarshal(body, &env) != nil { + return nil, false, false + } + if env.Method != methodSubscriptionsListen || len(env.ID) == 0 { + return nil, false, false + } + if !slices.Contains(env.Params.Notifications.ResourceSubscriptions, d.uri) { + return nil, false, false + } + d.mu.Lock() + defer d.mu.Unlock() + if !d.armed { + return nil, false, false + } + d.armed = false + return env.ID, d.graceful, true +} + +func serveControlledListen(w http.ResponseWriter, id json.RawMessage, graceful bool) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + if flusher != nil { + flusher.Flush() + } + if graceful { + fmt.Fprintf(w, "data: {\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{}}\n\n", id) + if flusher != nil { + flusher.Flush() + } + } + // Returning closes the stream. With no response written (abrupt), the + // client synthesizes a "request terminated without response" error for the + // listen call. +} + +func newListenDropServer(t *testing.T, uri string, srv *Server) *listenDropServer { + t.Helper() + return &listenDropServer{ + backend: mustNotPanic(t, NewStreamableHTTPHandler( + func(*http.Request) *Server { return srv }, + &StreamableHTTPOptions{Stateless: true}, + )), + uri: uri, + } +} + +// hasResourceSub reports whether the session still tracks a listen for uri. +func hasResourceSub(cs *ClientSession, uri string) bool { + cs.resourceSubsMu.Lock() + defer cs.resourceSubsMu.Unlock() + _, ok := cs.resourceSubs[uri] + return ok +} + +// testResourceListenDropReopens drives the shared body for the abrupt and +// graceful cases: the first listen ends without a client Unsubscribe, the SDK +// clears the subscription entry, and a bare re-Subscribe re-opens the stream, +// re-firing the server's SubscribeHandler. +func testResourceListenDropReopens(t *testing.T, graceful bool) { + t.Helper() + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + server := resourceSubServer(t, subCh, unsubCh) + drop := newListenDropServer(t, "file:///r1", server) + drop.arm(graceful) + httpServer := httptest.NewServer(drop) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + + c := NewClient(testImpl, &ClientOptions{}) + cs, err := c.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL, DisableStandaloneSSE: true}, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatalf("connect: %v", err) + } + t.Cleanup(func() { + cs.Close() + httpServer.CloseClientConnections() + httpServer.Close() + }) + + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("subscribe: %v", err) + } + + // The listen ends without a client Unsubscribe; the SDK must clear the + // entry so the URI is no longer considered subscribed. + waitUntil(t, 10*time.Second, "resource subscription entry to clear", func() bool { + return !hasResourceSub(cs, "file:///r1") + }) + + // No SubscribeHandler yet: the first listen never reached the backend. + // A bare re-Subscribe re-opens the stream and re-fires SubscribeHandler. + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("re-subscribe: %v", err) + } + select { + case got := <-subCh: + if got != "file:///r1" { + t.Fatalf("SubscribeHandler URI = %q, want file:///r1", got) + } + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for re-subscribe SubscribeHandler") + } +} + +// TestResourceSubscriptions_AbruptDropReopens verifies that when a resource +// subscription's listen stream drops abruptly (no client Unsubscribe), the SDK +// clears the subscription so a bare re-Subscribe re-opens the stream instead of +// no-oping, and does not auto-resubscribe on its own. +func TestResourceSubscriptions_AbruptDropReopens(t *testing.T) { + testResourceListenDropReopens(t, false) +} + +// TestResourceSubscriptions_GracefulEndReopens mirrors the abrupt case for a +// listen that ends with a normal result. +func TestResourceSubscriptions_GracefulEndReopens(t *testing.T) { + testResourceListenDropReopens(t, true) +} + +// TestResourceSubscriptions_GenGuard deterministically exercises the +// generation guard that protects an Unsubscribe→Subscribe race (and a +// re-subscribe from inside a completing listen): a stale listen goroutine +// (generation 1) that completes after a newer Subscribe has installed its own +// entry (generation 2) must not clear that newer entry. +func TestResourceSubscriptions_GenGuard(t *testing.T) { + cs := &ClientSession{ + resourceSubs: map[string]*resourceSub{ + "file:///r1": {cancel: func() {}, gen: 2}, + }, + } + + if cs.clearResourceSubIfGen("file:///r1", 1) { + t.Fatal("stale generation cleared a newer subscription entry") + } + if _, ok := cs.resourceSubs["file:///r1"]; !ok { + t.Fatal("newer subscription entry was removed by a stale generation") + } + + if !cs.clearResourceSubIfGen("file:///r1", 2) { + t.Fatal("owning generation failed to clear its entry") + } + if _, ok := cs.resourceSubs["file:///r1"]; ok { + t.Fatal("entry not removed by owning generation") + } + + // Clearing an absent URI is a no-op, not a panic. + if cs.clearResourceSubIfGen("file:///gone", 1) { + t.Fatal("clearing an absent URI reported a deletion") + } +} diff --git a/mcp/shared.go b/mcp/shared.go index 8afdc20b..f3b2a348 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -158,12 +158,8 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request // Create the result to unmarshal into. // The concrete type of the result is the return type of the receiving function. res := info.newResult() - if method == methodSubscriptionsListen { - callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params) - } else { - if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil { - return nil, err - } + if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil { + return nil, err } return res, nil } diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 77ff81a2..cd141260 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -2025,7 +2025,5 @@ func TestStreamableClient_StatelessSubscriptionsListen404(t *testing.T) { if res == nil { t.Fatal("ListTools result is nil") } - if !listenServed.Load() { - t.Fatal("subscriptions/listen was not called") - } + waitUntil(t, 5*time.Second, "subscriptions/listen to be called", listenServed.Load) } diff --git a/mcp/transport.go b/mcp/transport.go index 08d3b15f..35c3e70b 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -294,24 +294,6 @@ func (e *peerCancelledError) Error() string { func (e *peerCancelledError) Unwrap() error { return context.Canceled } -// callSubscriptionsListen issues a "subscriptions/listen" call (SEP-2575) -// without awaiting its JSON-RPC response. The call's logical lifetime is the -// stream of notifications that follow on the same channel — the empty -// response, if ever delivered, only marks subscription teardown — so the -// caller has nothing useful to block on. -// -// Cancellation is driven by ctx: when it is cancelled, a background goroutine -// sends a "notifications/cancelled" notification referencing the listen's -// request ID and retires the call from the connection's outgoing-calls map. -func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) { - call := conn.Call(ctx, method, params) - - go func() { - <-ctx.Done() - _ = cancelCall(ctx, conn, call) - }() -} - // call executes and awaits a jsonrpc2 call on the given connection, // translating errors into the mcp domain. func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params, result Result) error {