mcp: fix error propagation to prevent session closing (#1118)

Fixes #1117

mcp/streamable_client_test.go

@@ -1751,3 +1751,147 @@ func TestStreamableClientConnect_DiscoverUnsupportedVersionNegotiation(t *testin
 		t.Errorf("InitializeResult.ProtocolVersion = %q, want %q", got, protocolVersion20260728)
 	}
 }
+
+// TestStreamableClientHandlerErrorPropagation verifies that per-call
+// handler-level HTTP errors carrying a JSON-RPC error body do not tear down
+// the session, and that setting MCPGODEBUG=noprotocolerrorbody=1 restores the
+// pre-fix behavior in which non-transient errors permanently failed the
+// connection.
+func TestStreamableClientHandlerErrorPropagation(t *testing.T) {
+	ctx := context.Background()
+
+	// Build a JSON-RPC error response body for a given tools/call request.
+	jsonRPCErrorBody := func(id int64, code int64, msg string) string {
+		return jsonBody(t, resp(id, nil, &jsonrpc.Error{Code: code, Message: msg}))
+	}
+
+	tests := []struct {
+		name              string
+		callStatus        int    // HTTP status returned for the tools/call
+		callBody          string // response body for the tools/call ("" for empty; may be a JSON-RPC error)
+		disableBodyDecode bool   // if true, sets noprotocolerrorbody="1" for the duration of the test
+		wantErrRejected   bool   // whether the returned error should match errors.Is(err, jsonrpc2.ErrRejected)
+		wantSessionAlive  bool   // whether a subsequent call on the same session should succeed
+	}{
+		{
+			// SEP-2575 §"Missing Required Capabilities": server returns
+			// -32021 at HTTP 400. Session must survive.
+			name:             "400 with JSON-RPC error body (default)",
+			callStatus:       http.StatusBadRequest,
+			callBody:         jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"),
+			wantErrRejected:  true,
+			wantSessionAlive: true,
+		},
+		{
+			// Same server response, but the user opted out via
+			// MCPGODEBUG=noprotocolerrorbody=1. Pre-fix behavior: the
+			// session is torn down.
+			name:              "400 with JSON-RPC error body (noprotocolerrorbody=1)",
+			callStatus:        http.StatusBadRequest,
+			callBody:          jsonRPCErrorBody(1, CodeMissingRequiredClientCapabilities, "missing capability"),
+			disableBodyDecode: true,
+			wantErrRejected:   false,
+			wantSessionAlive:  false,
+		},
+		{
+			// Legacy server returns plain-text 400 with no JSON-RPC
+			// body: cannot be classified as a per-call rejection, so the
+			// session is still torn down.
+			name:             "400 with plain-text body",
+			callStatus:       http.StatusBadRequest,
+			callBody:         "Bad Request",
+			wantErrRejected:  false,
+			wantSessionAlive: false,
+		},
+		{
+			// SEP-2575: server returns -32601 for an unimplemented
+			// method at HTTP 404. Must be treated as a per-call
+			// rejection, not a terminated session.
+			name:             "404 with JSON-RPC MethodNotFound body",
+			callStatus:       http.StatusNotFound,
+			callBody:         jsonRPCErrorBody(1, jsonrpc.CodeMethodNotFound, `method not found: "tools/call"`),
+			wantErrRejected:  true,
+			wantSessionAlive: true,
+		},
+		{
+			// A bare 404 with no JSON-RPC body means the session has
+			// been terminated. Session must not survive.
+			name:             "404 with empty body (terminated session)",
+			callStatus:       http.StatusNotFound,
+			callBody:         "",
+			wantErrRejected:  false,
+			wantSessionAlive: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.disableBodyDecode {
+				prev := noprotocolerrorbody
+				noprotocolerrorbody = "1"
+				t.Cleanup(func() { noprotocolerrorbody = prev })
+			}
+
+			// Track how many tools/call requests we've served: first
+			// returns the error under test, subsequent calls succeed so
+			// we can observe whether the session is still alive.
+			var callsServed atomic.Int32
+			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, resp(r.ID.Raw().(int64), discoverResult, nil)), http.StatusOK
+						},
+					},
+					{"POST", "", "tools/call", ""}: {
+						header: header{"Content-Type": "application/json"},
+						responseFunc: func(r *jsonrpc.Request) (string, int) {
+							if callsServed.Add(1) == 1 {
+								return test.callBody, test.callStatus
+							}
+							return jsonBody(t, resp(r.ID.Raw().(int64), &CallToolResult{}, nil)), 0
+						},
+						optional: true,
+					},
+				},
+			}
+
+			httpServer := httptest.NewServer(fake)
+			defer httpServer.Close()
+
+			transport := &StreamableClientTransport{Endpoint: httpServer.URL}
+			client := NewClient(testImpl, nil)
+			session, err := client.Connect(ctx, transport, &ClientSessionOptions{ProtocolVersion: protocolVersion20260728})
+			if err != nil {
+				t.Fatalf("Connect: %v", err)
+			}
+			defer session.Close()
+
+			// First call: triggers the error under test.
+			_, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"})
+			if err == nil {
+				t.Fatal("first CallTool succeeded unexpectedly, want error")
+			}
+			if got := errors.Is(err, jsonrpc2.ErrRejected); got != test.wantErrRejected {
+				t.Errorf("errors.Is(err, ErrRejected) = %v, want %v (err = %v)", got, test.wantErrRejected, err)
+			}
+
+			// Second call: verifies whether the session survived.
+			_, err = session.CallTool(ctx, &CallToolParams{Name: "nonexistent"})
+			if test.wantSessionAlive {
+				if err != nil {
+					t.Errorf("second CallTool failed: %v (session should survive per-call rejection)", err)
+				}
+			} else {
+				if err == nil {
+					t.Error("second CallTool succeeded unexpectedly, want session torn down")
+				}
+			}
+		})
+	}
+}