diff --git a/flag.module b/flag.module
index 260b417..2e87183 100644
--- a/flag.module
+++ b/flag.module
@@ -163,7 +163,7 @@ function flag_entity_view(array &$build, EntityInterface $entity, EntityViewDisp
     return;
   }
 
-  // @todo: Add caching here PLZ.
+  $build['#cache']['contexts'][] = 'user.permissions';
 
   // Get all possible flags for this entity type.
   $flag_service = \Drupal::service('flag');
@@ -175,7 +175,6 @@ function flag_entity_view(array &$build, EntityInterface $entity, EntityViewDisp
       continue;
     }
 
-    $link_type_plugin = $flag->getLinkTypePlugin();
     $flag_type_plugin = $flag->getFlagTypePlugin();
 
     // Only add cache key if flag link is displayed.
@@ -183,56 +182,53 @@ function flag_entity_view(array &$build, EntityInterface $entity, EntityViewDisp
       continue;
     }
 
-    $action = 'flag';
-    if ($flag->isFlagged($entity)) {
-      $action = 'unflag';
-    }
+    $build['flag_' . $flag->id()] = [
+      '#lazy_builder' => ['flag_lazy_link_builder', [
+        $entity->getEntityTypeId(),
+        $entity->id(),
+        $flag->id(),
+      ]],
+      '#create_placeholder' => TRUE,
+    ];
+  }
+}
 
-    // If the user does not have permission, go to the next foreach loop and
-    // don't display this flag.
-    if (!$flag->hasActionAccess($action)){
-      continue;
-    }
+/**
+ * Lazy builder callback for displaying a flag action link.
+ *
+ * @param string $entity_type_id
+ *   The entity type ID for which the link should be shown.
+ * @param string|int $entity_id
+ *   The entity ID for which the link should be shown
+ * @param string $flag_id
+ *   The flag for which the link should be shown.
+ *
+ * @return array
+ *   A render array for the action link, empty if he user does not have access.
+ */
+function flag_lazy_link_builder($entity_type_id, $entity_id, $flag_id) {
+  $entity = \Drupal::entityManager()->getStorage($entity_type_id)->load($entity_id);
+  /** @var \Drupal\flag\FlagServiceInterface $flag_service */
+  $flag_service = \Drupal::service('flag');
+  $flag = $flag_service->getFlagById($flag_id);
 
+  $action = 'flag';
+  if ($flag->isFlagged($entity)) {
+    $action = 'unflag';
+  }
+
+  // Only display if the user does have access.
+  if ($flag->hasActionAccess($action)) {
+    $link_type_plugin = $flag->getLinkTypePlugin();
     $link = $link_type_plugin->renderLink($action, $flag, $entity);
-    $build['flag_' . $flag->id()] = $link;
+    // The actual render array must be in a nested key, due to a bug in
+    // lazy builder handling that does not properly render top-level #type
+    // elements.
+    return ['link' => $link];
   }
-}
 
-/**
- * Implements hook_entity_build_defaults_alter().
- */
-function flag_entity_build_defaults_alter(array &$build, EntityInterface $entity, $view_mode = 'full', $langcode = NULL) {
-  // Add the flag ID combined with the action to the cache key and set the cache
-  // context if render caching is enabled.
-  if (isset($build['#cache']) && isset($build['#cache']['keys'])) {
-
-    // Get all possible flags for this entity type.
-    $flag_service = \Drupal::service('flag');
-    $flags = $flag_service->getFlags($entity->getEntityTypeID(),
-      $entity->bundle());
-
-    // Get the corresponding display settings.
-    $display = EntityViewDisplay::collectRenderDisplay($entity, $view_mode);
-
-    foreach ($flags as $flag) {
-      $flag_type_plugin = $flag->getFlagTypePlugin();
-
-      // Only add cache key if flag link is displayed.
-      if ($flag_type_plugin->showAsField() && !$display->getComponent('flag_' . $flag->id())) {
-        continue;
-      }
-
-      $action = 'flag';
-      if ($flag->isFlagged($entity)) {
-        $action = 'unflag';
-      }
-
-      $build['#cache']['keys'][] = $flag->id . '-' . $action;
-      $build['#cache']['contexts'][] = 'user.permissions';
-    }
-  }
-  return $build;
+  // Lazy builders must always return a render array.
+  return [];
 }
 
 /**