diff --git a/token.module b/token.module
index 97ba5c1..a5f4b9b 100644
--- a/token.module
+++ b/token.module
@@ -90,7 +90,6 @@ function token_theme() {
       'click_insert' => TRUE,
       'show_restricted' => FALSE,
       'recursion_limit' => 3,
-      'dialog' => FALSE,
     ),
     'function' => 'theme_token_tree',
     'file' => 'token.pages.inc',
@@ -621,7 +620,7 @@ function token_form_field_config_edit_form_alter(&$form, FormStateInterface $for
   $form['description'] += array('#token_types' => array());
 
   $form['token_tree'] = array(
-    '#theme' => 'token_tree',
+    '#theme' => 'token_tree_link',
     '#token_types' => array(),
     '#dialog' => TRUE,
     '#weight' => $form['description']['#weight'] + 0.5,
@@ -642,7 +641,7 @@ function token_form_system_actions_configure_alter(&$form, $form_state) {
     case 'system_send_email_action':
     case 'system_goto_action':
       $form['token_tree'] = array(
-        '#theme' => 'token_tree',
+        '#theme' => 'token_tree_link',
         '#token_types' => 'all',
         '#dialog' => TRUE,
         '#weight' => 100,
@@ -707,7 +706,7 @@ function token_form_user_admin_settings_alter(&$form, FormStateInterface $form_s
 
   // Add the token tree UI.
   $form['email']['token_tree'] = array(
-    '#theme' => 'token_tree',
+    '#theme' => 'token_tree_link',
     '#token_types' => array('user'),
     '#show_restricted' => TRUE,
     '#dialog' => TRUE,
diff --git a/token.pages.inc b/token.pages.inc
index d6fa56a..f1ec153 100644
--- a/token.pages.inc
+++ b/token.pages.inc
@@ -33,10 +33,6 @@ function theme_token_tree_link($variables) {
   }
   $variables['options']['query']['options'] += $tree_variables;
 
-  // We should never pass the dialog option to theme_token_tree(). It is only
-  // used for this function.
-  unset($variables['options']['query']['options']['dialog']);
-
   // Because PHP converts query strings with arrays into a different syntax on
   // the next request, the options have to be encoded with JSON in the query
   // string so that we can reliably decode it for token comparison.
@@ -92,10 +88,6 @@ function theme_tree_table($variables) {
  * @ingroup themeable
  */
 function theme_token_tree($variables) {
-  if (!empty($variables['dialog'])) {
-    return \Drupal::theme()->render('token_tree_link', $variables);
-  }
-
   $token_types = $variables['token_types'];
   $info = token_get_info();