diff --git a/src/Controller/AJAXLinkController.php b/src/Controller/AJAXLinkController.php
index 81e3c4d..0a79537 100644
--- a/src/Controller/AJAXLinkController.php
+++ b/src/Controller/AJAXLinkController.php
@@ -25,17 +25,18 @@ class AJAXLinkController extends ControllerBase {
    *
    * @param \Drupal\flag\FlagInterface $flag
    *   The flag entity.
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   The flaggable entity
+   * @param int $entity_id
+   *   The flaggable ID.
    *
    * @return AjaxResponse
    *   The response object.
    *
    * @see \Drupal\flag\Plugin\ActionLink\AJAXactionLink
    */
-  public function flag(FlagInterface $flag, EntityInterface $entity) {
-
-    $flagging = \Drupal::service('flag')->flag($flag, $entity);
+  public function flag(FlagInterface $flag, $entity_id) {
+    $flag_service = \Drupal::service('flag');
+    $entity = $flag_service->getFlaggableById($flag, $entity_id);
+    $flag_service->flag($flag, $entity);
 
     return $this->generateResponse('unflag', $flag, $entity);
   }
diff --git a/src/Controller/ReloadLinkController.php b/src/Controller/ReloadLinkController.php
index 4acaec6..3088ac7 100644
--- a/src/Controller/ReloadLinkController.php
+++ b/src/Controller/ReloadLinkController.php
@@ -64,9 +64,11 @@ class ReloadLinkController extends ControllerBase implements ContainerInjectionI
    * @see \Drupal\flag\Plugin\Reload
    */
   public function flag(FlagInterface $flag, $entity_id) {
+    /* @var \Drupal\Core\Entity\EntityInterface $entity */
+    $entity = $this->flagService->getFlaggableById($flag, $entity_id);
 
     /* @var \Drupal\flag\FlaggingInterface $flagging */
-    $flagging = $this->flagService->flag($flag, $entity_id);
+    $flagging = $this->flagService->flag($flag, $entity);
 
     // Redirect back to the entity. A passed in destination query parameter
     // will automatically override this.
@@ -88,12 +90,9 @@ class ReloadLinkController extends ControllerBase implements ContainerInjectionI
    * @see \Drupal\flag\Plugin\Reload
    */
   public function unflag(FlagInterface $flag, $entity_id) {
-    $flag_id = $flag->id();
-
-    $this->flagService->unflag($flag_id, $entity_id);
-
-    $flag = $this->flagService->getFlagById($flag_id);
+    /* @var \Drupal\Core\Entity\EntityInterface $entity */
     $entity = $this->flagService->getFlaggableById($flag, $entity_id);
+    $this->flagService->unflag($flag, $entity);
 
     // Redirect back to the entity. A passed in destination query parameter
     // will automatically override this.