diff --git a/gst/vaapi/gstvaapidecode.c b/gst/vaapi/gstvaapidecode.c
index c35af825da..965bd42469 100644
--- a/gst/vaapi/gstvaapidecode.c
+++ b/gst/vaapi/gstvaapidecode.c
@@ -30,12 +30,10 @@
 
 #include "gst/vaapi/sysdeps.h"
 #include <gst/vaapi/gstvaapidisplay.h>
-#if !GST_CHECK_VERSION(1,1,0)
-#include <gst/video/videocontext.h>
-#endif
 
 #include "gstvaapidecode.h"
 #include "gstvaapipluginutil.h"
+#include "gstvaapivideocontext.h"
 #include "gstvaapivideobuffer.h"
 #if GST_CHECK_VERSION(1,0,0)
 #include "gstvaapivideobufferpool.h"
@@ -544,6 +542,20 @@ error_create_pool:
 }
 #endif
 
+#if GST_CHECK_VERSION(1,1,0)
+static void
+gst_vaapidecode_set_context(GstElement *element, GstContext *context)
+{
+    GstVaapiDecode * const decode = GST_VAAPIDECODE(element);
+    GstVaapiDisplay *display = NULL;
+
+    if (gst_vaapi_video_context_get_display(context, &display)) {
+        GST_INFO_OBJECT(element, "set display %p", display);
+        gst_vaapi_display_replace(&decode->display, display);
+    }
+}
+#endif
+
 static inline gboolean
 gst_vaapidecode_ensure_display(GstVaapiDecode *decode)
 {
@@ -777,6 +789,10 @@ gst_vaapidecode_class_init(GstVaapiDecodeClass *klass)
         GST_DEBUG_FUNCPTR(gst_vaapidecode_decide_allocation);
 #endif
 
+#if GST_CHECK_VERSION(1,1,0)
+    element_class->set_context = GST_DEBUG_FUNCPTR(gst_vaapidecode_set_context);
+#endif
+
     gst_element_class_set_static_metadata(element_class,
         "VA-API decoder",
         "Codec/Decoder/Video",
@@ -880,6 +896,24 @@ gst_vaapidecode_query(GST_PAD_QUERY_FUNCTION_ARGS)
             res = TRUE;
             break;
         }
+#endif
+#if GST_CHECK_VERSION(1,1,0)
+        case GST_QUERY_CONTEXT: {
+            const gchar *context_type = NULL;
+
+            if (gst_query_parse_context_type(query, &context_type) &&
+                !g_strcmp0(context_type, GST_VAAPI_DISPLAY_CONTEXT_TYPE_NAME) &&
+                decode->display) {
+                GstContext *context;
+
+                context = gst_vaapi_video_context_new_with_display(
+                    decode->display, FALSE);
+                gst_query_set_context(query, context);
+                gst_context_unref(context);
+                return TRUE;
+            }
+            // fall-through
+        }
 #endif
         default:
             res = GST_PAD_QUERY_FUNCTION_CALL(decode->sinkpad_query,
diff --git a/gst/vaapi/gstvaapisink.c b/gst/vaapi/gstvaapisink.c
index 6b405520ba..9cbb008e1a 100644
--- a/gst/vaapi/gstvaapisink.c
+++ b/gst/vaapi/gstvaapisink.c
@@ -32,9 +32,6 @@
 #include "gst/vaapi/sysdeps.h"
 #include <gst/gst.h>
 #include <gst/video/video.h>
-#if !GST_CHECK_VERSION(1,1,0)
-#include <gst/video/videocontext.h>
-#endif
 
 #include <gst/vaapi/gstvaapivalue.h>
 #if USE_DRM
@@ -72,6 +69,7 @@
 
 #include "gstvaapisink.h"
 #include "gstvaapipluginutil.h"
+#include "gstvaapivideocontext.h"
 #include "gstvaapivideometa.h"
 #if GST_CHECK_VERSION(1,0,0)
 #include "gstvaapivideobufferpool.h"
@@ -1330,6 +1328,20 @@ gst_vaapisink_buffer_alloc(
 }
 #endif
 
+#if GST_CHECK_VERSION(1,1,0)
+static void
+gst_vaapisink_set_context(GstElement *element, GstContext *context)
+{
+    GstVaapiSink * const sink = GST_VAAPISINK(element);
+    GstVaapiDisplay *display = NULL;
+
+    if (gst_vaapi_video_context_get_display(context, &display)) {
+        GST_INFO_OBJECT(element, "set display %p", display);
+        gst_vaapi_display_replace(&sink->display, display);
+    }
+}
+#endif
+
 static gboolean
 gst_vaapisink_query(GstBaseSink *base_sink, GstQuery *query)
 {
@@ -1339,6 +1351,30 @@ gst_vaapisink_query(GstBaseSink *base_sink, GstQuery *query)
         GST_DEBUG("sharing display %p", sink->display);
         return TRUE;
     }
+
+    switch(GST_QUERY_TYPE(query)) {
+#if GST_CHECK_VERSION(1,1,0)
+      case GST_QUERY_CONTEXT: {
+          const gchar *context_type = NULL;
+
+          if (gst_query_parse_context_type(query, &context_type) &&
+              !g_strcmp0(context_type, GST_VAAPI_DISPLAY_CONTEXT_TYPE_NAME) &&
+              sink->display) {
+              GstContext *context;
+
+              context = gst_vaapi_video_context_new_with_display(
+                  sink->display, FALSE);
+              gst_query_set_context(query, context);
+              gst_context_unref(context);
+              return TRUE;
+          }
+          // fall-through
+      }
+#endif
+      default:
+          break;
+    }
+
     return GST_BASE_SINK_CLASS(gst_vaapisink_parent_class)->query(base_sink,
         query);
 }
@@ -1455,6 +1491,10 @@ gst_vaapisink_class_init(GstVaapiSinkClass *klass)
     basesink_class->buffer_alloc = gst_vaapisink_buffer_alloc;
 #endif
 
+#if GST_CHECK_VERSION(1,1,0)
+    element_class->set_context = gst_vaapisink_set_context;
+#endif
+
     gst_element_class_set_static_metadata(element_class,
         "VA-API sink",
         "Sink/Video",