diff --git a/src/Entity/NotificationMessage.php b/src/Entity/NotificationMessage.php
index 26a63ad..adab32a 100644
--- a/src/Entity/NotificationMessage.php
+++ b/src/Entity/NotificationMessage.php
@@ -377,7 +377,7 @@ class NotificationMessage extends ContentEntityBase implements NotificationMessa
       if (!$definition->isRequired()) {
         continue;
       }
-      $contexts[$name] = $definition->getLabel();
+      $contexts[$definition->getDataType()] = $definition->getLabel();
     }
 
     return $contexts;
@@ -398,8 +398,13 @@ class NotificationMessage extends ContentEntityBase implements NotificationMessa
     array $contexts,
     array $required_contexts
   ) {
+    $contextsKeyedByDataType = [];
+    foreach ($contexts as $context) {
+      $contextsKeyedByDataType[$context->getContextDefinition()->getDataType()] = $context;
+    }
+
     return array_filter(
-      array_intersect_key($contexts, $required_contexts)
+      array_intersect_key($contextsKeyedByDataType, $required_contexts)
     );
   }
 
@@ -431,7 +436,9 @@ class NotificationMessage extends ContentEntityBase implements NotificationMessa
         if (!$context->hasContextValue()) {
           return FALSE;
         }
-        $instance->setContext($name, $context);
+        $contextNameParts = explode(':', $name);
+        $contextId = array_pop($contextNameParts);
+        $instance->setContext($contextId, $context);
       }
       return $instance->evaluate();
     }