Clear access token on expiration or auth error

package.json

@@ -59,6 +59,7 @@
   },
   "dependencies": {
     "@storybook/design-system": "^7.15.15",
+    "@urql/exchange-auth": "^2.1.6",
     "chromatic": "^7.2.0",
     "date-fns": "^2.30.0",
     "filesize": "^10.0.12",

src/screens/Authentication/Authentication.tsx

@@ -7,7 +7,7 @@ import { Verify } from "./Verify";
 import { Welcome } from "./Welcome";
 
 interface AuthenticationProps {
-  setAccessToken: (token: string) => void;
+  setAccessToken: (token: string | null) => void;
   hasProjectId: boolean;
 }
 

src/screens/VisualTests/VisualTests.stories.tsx

@@ -477,7 +477,7 @@ export const AcceptingFailed = {
       ],
     },
   },
-  play: playAll(async ({ canvasElement, argsByTarget, args, argTypes }) => {
+  play: playAll(async ({ canvasElement, argsByTarget }) => {
     const button = await findByRole(canvasElement, "button", { name: "Accept" });
     await fireEvent.click(button);
     await waitFor(async () =>

src/utils/graphQLClient.tsx

@@ -1,30 +1,47 @@
 import { useAddonState } from "@storybook/manager-api";
+import { authExchange } from "@urql/exchange-auth";
 import React from "react";
-import { Client, fetchExchange, Provider } from "urql";
+import { Client, fetchExchange, mapExchange, Provider } from "urql";
 import { v4 as uuid } from "uuid";
 
 import { ACCESS_TOKEN_KEY, ADDON_ID, CHROMATIC_API_URL } from "../constants";
 
 export { Provider };
 
-let currentToken: string | null = localStorage.getItem(ACCESS_TOKEN_KEY);
-const accessTokenSharedStateKey = `${ADDON_ID}/accessToken`;
+let currentToken: string | null;
+let currentTokenExpiration: number | null;
+const setCurrentToken = (token: string | null) => {
+  try {
+    const { exp } = token ? JSON.parse(atob(token.split(".")[1])) : { exp: null };
+    currentToken = token;
+    currentTokenExpiration = exp;
+  } catch (e) {
+    currentToken = null;
+    currentTokenExpiration = null;
+  }
+  if (currentToken) {
+    localStorage.setItem(ACCESS_TOKEN_KEY, currentToken);
+  } else {
+    localStorage.removeItem(ACCESS_TOKEN_KEY);
+  }
+};
+
+setCurrentToken(localStorage.getItem(ACCESS_TOKEN_KEY));
 
 export const useAccessToken = () => {
   // We use an object rather than a straight boolean here due to https://github.com/storybookjs/storybook/pull/23991
-  const [{ token }, setToken] = useAddonState<{ token: string | null }>(accessTokenSharedStateKey, {
-    token: currentToken,
-  });
-
-  const updateToken = (newToken: string | null) => {
-    currentToken = newToken;
-    if (currentToken) {
-      localStorage.setItem(ACCESS_TOKEN_KEY, currentToken);
-    } else {
-      localStorage.removeItem(ACCESS_TOKEN_KEY);
-    }
-    setToken({ token: newToken });
-  };
+  const [{ token }, setTokenState] = useAddonState<{ token: string | null }>(
+    `${ADDON_ID}/accessToken`,
+    { token: currentToken }
+  );
+
+  const updateToken = React.useCallback(
+    (newToken: string | null) => {
+      setCurrentToken(newToken);
+      setTokenState({ token: currentToken });
+    },
+    [setTokenState]
+  );
 
   return [token, updateToken] as const;
 };
@@ -33,12 +50,55 @@ const sessionId = uuid();
 
 export const client = new Client({
   url: CHROMATIC_API_URL,
-  exchanges: [fetchExchange], // no cacheExchange to prevent sharing data between stories
+  exchanges: [
+    // We don't use cacheExchange, because it would inadvertently share data between stories.
+    mapExchange({
+      onResult(result) {
+        // Not all queries contain the `viewer` field, in which case it will be `undefined`.
+        // When we do retrieve the field but the token is invalid, it will be `null`.
+        if (result.data?.viewer === null) setCurrentToken(null);
+      },
+    }),
+    authExchange(async (utils) => {
+      return {
+        addAuthToOperation(operation) {
+          if (!currentToken) return operation;
+          return utils.appendHeaders(operation, { Authorization: `Bearer ${currentToken}` });
+        },
+
+        // Determine if the current error is an authentication error.
+        didAuthError: (error) =>
+          error.response.status === 401 ||
+          error.graphQLErrors.some((e) => e.message.includes("Must login")),
+
+        // If didAuthError returns true, clear the token. Ideally we should refresh the token here.
+        // The operation will be retried automatically.
+        async refreshAuth() {
+          setCurrentToken(null);
+        },
+
+        // Prevent making a request if we know the token is missing, invalid or expired.
+        // This handler is called repeatedly so we avoid parsing the token each time.
+        willAuthError() {
+          if (!currentToken) return true;
+          try {
+            if (!currentTokenExpiration) {
+              const { exp } = JSON.parse(atob(currentToken.split(".")[1]));
+              currentTokenExpiration = exp;
+            }
+            return Date.now() / 1000 > (currentTokenExpiration || 0);
+          } catch (e) {
+            return true;
+          }
+        },
+      };
+    }),
+    fetchExchange,
+  ],
   fetchOptions: () => ({
     headers: {
-      accept: "*/*", // workaround for https://github.com/mswjs/msw/issues/1593
-      authorization: currentToken ? `Bearer ${currentToken}` : "",
-      "x-chromatic-session-id": sessionId,
+      Accept: "*/*", // workaround for https://github.com/mswjs/msw/issues/1593
+      "X-Chromatic-Session-ID": sessionId,
     },
   }),
 });

yarn.lock

@@ -4396,14 +4396,22 @@
     "@typescript-eslint/types" "5.62.0"
     eslint-visitor-keys "^3.3.0"
 
-"@urql/core@^4.1.0":
+"@urql/core@>=4.1.0", "@urql/core@^4.1.0":
   version "4.1.2"
   resolved "https://registry.yarnpkg.com/@urql/core/-/core-4.1.2.tgz#0080796f07f34a6878a29eb0ef319ceb2083b65b"
   integrity sha512-K+JA5dxEjY7Jkt1hV8G2ShkuOscKS/r+8QnXDDxTkyMzZzviYqz5f/zxgSElObT/QSW17xCC1LFl+kwiyX5opg==
   dependencies:
     "@0no-co/graphql.web" "^1.0.1"
     wonka "^6.3.2"
 
+"@urql/exchange-auth@^2.1.6":
+  version "2.1.6"
+  resolved "https://registry.yarnpkg.com/@urql/exchange-auth/-/exchange-auth-2.1.6.tgz#8197dab76ac8109ad23c9cde57f834153055d024"
+  integrity sha512-snOlt7p5kYq0KnPDuXkKe2qW3/BucQZOElvTeo3svLQuk9JiNJVnm6ffQ6QGiGO+G3AtMrctnno1+X44fLtDuQ==
+  dependencies:
+    "@urql/core" ">=4.1.0"
+    wonka "^6.3.2"
+
 "@vitejs/plugin-react@^3.0.1", "@vitejs/plugin-react@^3.1.0":
   version "3.1.0"
   resolved "https://registry.yarnpkg.com/@vitejs/plugin-react/-/plugin-react-3.1.0.tgz#d1091f535eab8b83d6e74034d01e27d73c773240"