diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 4994c63b..165482d0 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -411,6 +411,9 @@ type AsyncCall struct { // This can be used to cancel the call if needed. func (ac *AsyncCall) ID() ID { return ac.id } +// Done is closed when the call has a response or terminal error. +func (ac *AsyncCall) Done() <-chan struct{} { return ac.ready } + // retire processes the response to the call. // // It is an error to call retire more than once: retire is guarded by the diff --git a/mcp/client.go b/mcp/client.go index 74037990..54beb1d8 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -351,6 +351,7 @@ func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOp cs.listenCancel = cancelListen if err := cs.subscriptionsListen(listenCtx, subscribeParams); err != nil { cancelListen() + _ = cs.Close() return nil, fmt.Errorf("opening subscriptions/listen: %w", err) } } diff --git a/mcp/shared.go b/mcp/shared.go index 5069a470..50e8a991 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -138,7 +138,9 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request // 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) + if err := callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params); err != nil { + return nil, err + } } else { if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil { return nil, err diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index a5957bf6..48eb32f0 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -1364,6 +1364,58 @@ func TestStreamableClientConnect_DiscoverSuccess(t *testing.T) { } } +func TestStreamableClientConnect_SubscriptionsListenError(t *testing.T) { + ctx := context.Background() + + fake := &fakeStreamableServer{ + t: t, + responses: fakeResponses{ + {"POST", "", methodDiscover, ""}: { + header: header{"Content-Type": "application/json"}, + wantProtocolVersion: protocolVersion20260728, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ + ID: r.ID, + Result: mustMarshal(discoverResult), + }), http.StatusOK + }, + }, + {"POST", "", methodSubscriptionsListen, ""}: { + header: header{"Content-Type": "application/json"}, + wantProtocolVersion: protocolVersion20260728, + responseFunc: func(r *jsonrpc.Request) (string, int) { + return jsonBody(t, &jsonrpc.Response{ + ID: r.ID, + Error: &jsonrpc.Error{ + Code: jsonrpc.CodeInvalidParams, + Message: "listen rejected", + }, + }), http.StatusBadRequest + }, + }, + }, + } + + httpServer := httptest.NewServer(fake) + defer httpServer.Close() + + client := NewClient(testImpl, &ClientOptions{ + ToolListChangedHandler: func(context.Context, *ToolListChangedRequest) {}, + }) + session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err == nil { + session.Close() + t.Fatal("Connect succeeded despite rejected subscriptions/listen") + } + if !errors.Is(err, jsonrpc2.ErrRejected) { + t.Fatalf("Connect error = %v, want error wrapping jsonrpc2.ErrRejected", err) + } + if !strings.Contains(err.Error(), "opening subscriptions/listen") { + t.Fatalf("Connect error = %v, want subscriptions/listen context", err) + } +} + // TestStreamableClientConnSetMCPHeaders_ProtocolVersion covers // streamableClientConn.setMCPHeaders' selection of the Mcp-Protocol-Version // header value. diff --git a/mcp/transport.go b/mcp/transport.go index d72f9d1b..a5d5900d 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -256,22 +256,30 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a return nil, jsonrpc2.ErrNotHandled } -// 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. +// callSubscriptionsListen issues a "subscriptions/listen" call (SEP-2575). +// If the call is accepted, its 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) { +func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) error { call := conn.Call(ctx, method, params) + select { + case <-call.Done(): + return call.Await(context.Background(), nil) + case <-ctx.Done(): + _ = cancelCall(ctx, conn, call) + return nil + default: + } go func() { <-ctx.Done() _ = cancelCall(ctx, conn, call) }() + return nil } // call executes and awaits a jsonrpc2 call on the given connection,