diff --git a/src/Controller/TokenTreeController.php b/src/Controller/TokenTreeController.php
index 0e32798..e13d1d3 100644
--- a/src/Controller/TokenTreeController.php
+++ b/src/Controller/TokenTreeController.php
@@ -20,6 +20,13 @@ class TokenTreeController extends ControllerBase {
    * Page callback to output a token tree as an empty page.
    */
   function outputTree(Request $request) {
+    $build = [
+      '#cache' => [
+        'contexts' => [
+          'url.query_args:options',
+        ],
+      ],
+    ];
     $build['#title'] = $this->t('Available tokens');
 
     $options = $request->query->has('options') ? Json::decode($request->query->get('options')) : array();
diff --git a/src/Tests/Tree/HelpPageTest.php b/src/Tests/Tree/HelpPageTest.php
index b43f0ba..937506c 100644
--- a/src/Tests/Tree/HelpPageTest.php
+++ b/src/Tests/Tree/HelpPageTest.php
@@ -59,5 +59,9 @@ class HelpPageTest extends TokenTestBase {
     $this->assertTokenInTree('[current-page:url:unaliased]', 'current-page--url');
     $this->assertTokenInTree('[current-page:url:unaliased:args]', 'current-page--url--unaliased');
     $this->assertTokenInTree('[user:original:account-name]', 'user--original');
+
+    // Assert some of the restricted tokens to ensure they are shown.
+    $this->assertTokenInTree('[user:one-time-login-url]', 'user');
+    $this->assertTokenInTree('[user:original:cancel-url]', 'user--original');
   }
 }
diff --git a/src/Tests/Tree/TokenTreeTestTrait.php b/src/Tests/Tree/TokenTreeTestTrait.php
index 843f585..7b5a400 100644
--- a/src/Tests/Tree/TokenTreeTestTrait.php
+++ b/src/Tests/Tree/TokenTreeTestTrait.php
@@ -60,11 +60,7 @@ trait TokenTreeTestTrait {
    *   in test output.
    */
   protected function assertTokenInTree($token, $parent = '', $message = '', $group = 'Other') {
-    $xpath = "//tr";
-    if ($parent) {
-      $xpath .= '[contains(@class, "child-of-token-' . $parent . ' ")]';
-    }
-    $xpath .= '/td[contains(@class, "token-key") and text() = "' . $token . '"]';
+    $xpath = $this->getXpathForTokenInTree($token, $parent);
 
     if (!$message) {
       $message = "Token $token found.";
@@ -72,4 +68,47 @@ trait TokenTreeTestTrait {
 
     $this->assertIdentical(1, count($this->xpath($xpath)), $message, $group);
   }
+
+  /**
+   * Check to see if the specified token is present in the token browser.
+   *
+   * @param $token
+   *   The token name with the surrounding square brackets [].
+   * @param string $parent
+   *   (optional) The parent CSS identifier of this token.
+   * @param string $message
+   *   (optional) A message to display with the assertion.
+   * @param string $group
+   *   (optional) The group this message is in, which is displayed in a column
+   *   in test output.
+   */
+  protected function assertTokenNotInTree($token, $parent = '', $message = '', $group = 'Other') {
+    $xpath = $this->getXpathForTokenInTree($token, $parent);
+
+    if (!$message) {
+      $message = "Token $token not found.";
+    }
+
+    $this->assertIdentical(0, count($this->xpath($xpath)), $message, $group);
+  }
+
+  /**
+   * Get xpath to check for token in tree.
+   *
+   * @param $token
+   *   The token name with the surrounding square brackets [].
+   * @param string $parent
+   *   (optional) The parent CSS identifier of this token.
+   *
+   * @return string
+   *   The xpath to check for the token and parent.
+   */
+  protected function getXpathForTokenInTree($token, $parent = '') {
+    $xpath = "//tr";
+    if ($parent) {
+      $xpath .= '[contains(@class, "child-of-token-' . $parent . ' ")]';
+    }
+    $xpath .= '/td[contains(@class, "token-key") and text() = "' . $token . '"]';
+    return $xpath;
+  }
 }
diff --git a/src/Tests/Tree/TreeTest.php b/src/Tests/Tree/TreeTest.php
index a873029..27a1c3e 100644
--- a/src/Tests/Tree/TreeTest.php
+++ b/src/Tests/Tree/TreeTest.php
@@ -71,6 +71,17 @@ class TreeTest extends TokenTestBase {
 
     $this->assertTokenInTree('[user:account-name]', 'user');
     $this->assertTokenInTree('[user:original:account-name]', 'user--original');
+
+    // Assert some of the restricted tokens to ensure they are not shown.
+    $this->assertTokenNotInTree('[user:one-time-login-url]', 'user');
+    $this->assertTokenNotInTree('[user:original:cancel-url]', 'user--original');
+
+    // Request with show_restricted set to TRUE to show restricted tokens and
+    // check for them.
+    $this->drupalGet($this->getTokenTreeUrl(['token_types' => ['user'], 'show_restricted' => TRUE]));
+    $this->assertEqual('MISS', $this->drupalGetHeader('x-drupal-dynamic-cache'), 'Cache was not hit');
+    $this->assertTokenInTree('[user:one-time-login-url]', 'user');
+    $this->assertTokenInTree('[user:original:cancel-url]', 'user--original');
   }
 
   /**