diff --git a/modules/message_recipient_entity/src/MessageAssociatedEntityTrait.php b/modules/message_recipient_entity/src/MessageAssociatedEntityTrait.php
index 46d7666ab231248dbfbad284d97975532f5b3c03..33344d10bc696ff042915309d46c60fbe277de0b 100644
--- a/modules/message_recipient_entity/src/MessageAssociatedEntityTrait.php
+++ b/modules/message_recipient_entity/src/MessageAssociatedEntityTrait.php
@@ -25,10 +25,10 @@ trait MessageAssociatedEntityTrait {
    * @param \Drupal\message\MessageInterface $message
    *   The message entity.
    *
-   * @return \Drupal\Core\Entity\ContentEntityInterface
-   *   The associated entity
+   * @return \Drupal\Core\Entity\ContentEntityInterface|null
+   *   The associated entity or NULL.
    */
-  public function getAssociatedEntity(MessageInterface $message): ContentEntityInterface {
+  public function getAssociatedEntity(MessageInterface $message): ?ContentEntityInterface {
     return $message->get('message_associated_entity')
       ?->first()
       ?->get('entity')