diff --git a/menu_token.module b/menu_token.module
index 13e239f..a8097f2 100644
--- a/menu_token.module
+++ b/menu_token.module
@@ -46,8 +46,9 @@ function menu_token_menu_token_plugins() {
 
   foreach ($entities as $entity => $enabled) {
     if ($enabled) {
-      $plugins["{$entity}_context"] = array(
-        'type' => $entity,
+      $token_type = $entity_info[$entity]['token type'];
+      $plugins["{$token_type}_context"] = array(
+        'type' => $token_type,
         'label' => t('@entity_label from context', array('@entity_label' => $entity_info[$entity]['label'])),
         'description' => t('Picks a @entity_label from the current context.', array('@entity_label' => $entity_info[$entity]['label'])),
         'handler' => array(
@@ -57,8 +58,8 @@ function menu_token_menu_token_plugins() {
         ),
       );
 
-      $plugins["{$entity}_random"] = array(
-        'type' => $entity,
+      $plugins["{$token_type}_random"] = array(
+        'type' => $token_type,
         'label' => t('Random @entity_label', array('@entity_label' => $entity_info[$entity]['label'])),
         'description' => t('Picks a random @entity_label from the database.', array('@entity_label' => $entity_info[$entity]['label'])),
         'handler' => array(
@@ -68,8 +69,8 @@ function menu_token_menu_token_plugins() {
         ),
       );
 
-      $plugins["{$entity}_user_defined"] = array(
-        'type' => $entity,
+      $plugins["{$token_type}_user_defined"] = array(
+        'type' => $token_type,
         'label' => t('User-defined @entity_label', array('@entity_label' => $entity_info[$entity]['label'])),
         'description' => t('Uses a user-defined @entity_label.', array('@entity_label' => $entity_info[$entity]['label'])),
         'handler' => array(
@@ -98,12 +99,13 @@ function menu_token_translated_menu_link_alter(&$item, $map) {
   if (empty($menu_admin) && isset($item['options']['menu_token_link_path'])) {
     $info = token_get_info();
     $data = array();
+    $token_entity_mapping = token_get_entity_mapping();
 
     // Load data objects used when replacing link.
     if (isset($item['options']['menu_token_data'])) {
       foreach ($item['options']['menu_token_data'] as $type => $values) {
         if (!empty($info['types'][$type]) && $handler = menu_token_get_handler($values['plugin'])) {
-          $values['options']['_type'] = $type;
+          $values['options']['_type'] = $token_entity_mapping[$type];
           if ($object = $handler->object_load($values['options'])) {
             $data[$type] = $object;
           }
diff --git a/plugins/menu_token_entity_context.inc b/plugins/menu_token_entity_context.inc
index 6f08b9e..7fd0d73 100644
--- a/plugins/menu_token_entity_context.inc
+++ b/plugins/menu_token_entity_context.inc
@@ -21,22 +21,21 @@ class menu_token_entity_context implements menu_token_handler {
     $entity_type = $options['_type'];
     $entity_info = entity_get_info($entity_type);
 
-    // Taken from entity_uri() to determine the uri callback for this entity
-    // type.
-    // A bundle-specific callback takes precedence over the generic one for the
-    // entity type.
-    if (isset($entity_info['uri callback'])) {
-      $uri_callback = $entity_info['uri callback'];
-    }
-    else {
-      return NULL;
-    }
-
-    $path = @$uri_callback();
-    $path = current(explode('/', $path['path']));
-
-    if (arg(0) === $path && is_numeric(arg(1))) {
-      return entity_load_single($entity_type, arg(1));
+    if ($menu_item = menu_get_item()) {
+      // Use the first menu argument flagged as loadable for the entity id.
+      $id_pos = reset(array_keys($menu_item['load_functions']));
+      $entity = entity_load_single($entity_type, arg($id_pos));
+
+      if ($entity) {
+        $uri = entity_uri($entity_type, $entity);
+
+        // The first loadable argument might have resulted in an incorrect
+        // entity getting loaded. If the loaded entity's path matches at least
+        // the initial part of the current URL then it should be correct.
+        if (strpos($menu_item['href'], $uri['path']) === 0) {
+          return $entity;
+        }
+      }
     }
 
     return FALSE;