[ENG-2416] feat: add jwt provider session storage part 1

package.json

@@ -99,6 +99,7 @@
     "eslint-plugin-react-refresh": "^0.4.20",
     "globals": "^16.0.0",
     "immer": "^10.0.3",
+    "jose": "^6.0.12",
     "lodash.isequal": "^4.5.0",
     "react-tooltip": "^5.28.0"
   },

src/components/Connect/ConnectProvider.tsx

@@ -1,5 +1,6 @@
 import { useCallback } from "react";
 import { ConnectionsProvider } from "context/ConnectionsContextProvider";
+import { InstallationProvider } from "src/headless";
 import { useForceUpdate } from "src/hooks/useForceUpdate";
 import { Connection } from "src/services/api";
 
@@ -60,26 +61,35 @@ export function ConnectProvider({
 
   return (
     <div className={resetStyles.resetContainer} key={seed}>
-      <ConnectionsProvider groupRef={groupRef} provider={provider}>
-        <ProtectedConnectionLayout
-          resetComponent={reset}
-          provider={provider}
-          consumerRef={consumerRef}
-          consumerName={consumerName}
-          groupRef={groupRef}
-          groupName={groupName}
-          onSuccess={onSuccessFx}
-          onDisconnectSuccess={onDisconnectSuccess}
-        >
-          <RedirectHandler redirectURL={redirectUrl}>
-            <ConnectedSuccessBox
-              resetComponent={reset}
-              provider={provider}
-              onDisconnectSuccess={onDisconnectSuccess}
-            />
-          </RedirectHandler>
-        </ProtectedConnectionLayout>
-      </ConnectionsProvider>
+      {/* InstallationProvider is nested in ConnectionsProvider and API service JWT auth */}
+      <InstallationProvider
+        integration={provider}
+        consumerRef={consumerRef}
+        consumerName={consumerName}
+        groupRef={groupRef}
+        groupName={groupName}
+      >
+        <ConnectionsProvider groupRef={groupRef} provider={provider}>
+          <ProtectedConnectionLayout
+            resetComponent={reset}
+            provider={provider}
+            consumerRef={consumerRef}
+            consumerName={consumerName}
+            groupRef={groupRef}
+            groupName={groupName}
+            onSuccess={onSuccessFx}
+            onDisconnectSuccess={onDisconnectSuccess}
+          >
+            <RedirectHandler redirectURL={redirectUrl}>
+              <ConnectedSuccessBox
+                resetComponent={reset}
+                provider={provider}
+                onDisconnectSuccess={onDisconnectSuccess}
+              />
+            </RedirectHandler>
+          </ProtectedConnectionLayout>
+        </ConnectionsProvider>
+      </InstallationProvider>
     </div>
   );
 }

src/context/AmpersandContextProvider/AmpersandContextProvider.tsx

@@ -10,10 +10,11 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 
 import { ApiKeyProvider } from "../ApiKeyContextProvider";
 import { ErrorStateProvider } from "../ErrorContextProvider";
+import { JwtTokenProvider } from "../JwtTokenContextProvider";
 
 interface AmpersandProviderProps {
   options: {
-    apiKey: string;
+    apiKey?: string;
     /**
      * Use `project` instead of `projectId`.
      * @deprecated
@@ -24,6 +25,11 @@ interface AmpersandProviderProps {
      */
     project?: string;
     styles?: object;
+    /**
+     * Callback function to get a JWT token for authorization.
+     * This function should return a Promise that resolves to a JWT token string.
+     */
+    getToken?: (consumerRef: string, groupRef: string) => Promise<string>;
   };
   children: React.ReactNode;
 }
@@ -52,7 +58,7 @@ const queryClient = new QueryClient();
 
 export function AmpersandProvider(props: AmpersandProviderProps) {
   const {
-    options: { apiKey, projectId, project },
+    options: { apiKey, projectId, project, getToken },
     children,
   } = props;
   const projectIdOrName = project || projectId;
@@ -67,8 +73,16 @@ export function AmpersandProvider(props: AmpersandProviderProps) {
     );
   }
 
-  if (!apiKey) {
-    throw new Error("Cannot use AmpersandProvider without an apiKey.");
+  if (!apiKey && !getToken) {
+    throw new Error(
+      "Cannot use AmpersandProvider without an apiKey or getToken.",
+    );
+  }
+
+  if (apiKey && getToken) {
+    throw new Error(
+      "Cannot use AmpersandProvider with both apiKey and getToken.",
+    );
   }
 
   const contextValue: AmpersandContextValue = {
@@ -80,7 +94,9 @@ export function AmpersandProvider(props: AmpersandProviderProps) {
     <QueryClientProvider client={queryClient}>
       <AmpersandContext.Provider value={contextValue}>
         <ErrorStateProvider>
-          <ApiKeyProvider value={apiKey}>{children}</ApiKeyProvider>
+          <JwtTokenProvider getTokenCallback={getToken || null}>
+            <ApiKeyProvider value={apiKey || null}>{children}</ApiKeyProvider>
+          </JwtTokenProvider>
         </ErrorStateProvider>
       </AmpersandContext.Provider>
     </QueryClientProvider>

src/context/ApiKeyContextProvider.tsx

@@ -6,10 +6,5 @@ export const ApiKeyProvider = ApiKeyContext.Provider;
 
 export const useApiKey = () => {
   const apiKey = useContext(ApiKeyContext);
-
-  if (apiKey === null) {
-    throw new Error("useApiKey must be used within an ApiKeyProvider");
-  }
-
   return apiKey;
 };

src/context/JwtTokenContextProvider.tsx

@@ -0,0 +1,213 @@
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useState,
+} from "react";
+import { jwtVerify } from "jose";
+
+interface TokenCacheEntry {
+  token: string;
+  expiresAt: number;
+}
+
+const SESSION_STORAGE_PREFIX = "amp-labs_jwt_";
+
+const DEFAULT_TOKEN_EXPIRATION_TIME = 10 * 60 * 1000; // 10 minutes
+
+const createCacheKey = (consumerRef: string, groupRef: string) =>
+  `${consumerRef}:${groupRef}`;
+
+const getSessionStorageKey = (cacheKey: string) =>
+  `${SESSION_STORAGE_PREFIX}${cacheKey}`;
+
+interface JwtTokenContextValue {
+  getToken?: (consumerRef: string, groupRef: string) => Promise<string>;
+}
+
+const JwtTokenContext = createContext<JwtTokenContextValue | null>(null);
+
+interface JwtTokenProviderProps {
+  getTokenCallback:
+    | ((consumerRef: string, groupRef: string) => Promise<string>)
+    | null;
+  children: React.ReactNode;
+}
+
+/**
+ * Extract JWT token expiration time
+ */
+const getTokenExpirationTime = async (
+  token: string,
+): Promise<number | null> => {
+  try {
+    const decoded = await jwtVerify(token, new Uint8Array(0), {
+      algorithms: [], // Skip signature verification
+    });
+    const payload = decoded.payload;
+    return payload.exp && typeof payload.exp === "number"
+      ? payload.exp * 1000 // jwt expiration is in seconds, convert to milliseconds
+      : null;
+  } catch (error) {
+    console.warn("Failed to decode JWT token:", error);
+    return null;
+  }
+};
+
+/**
+ * Simplified JWT token provider with cleaner caching logic
+ */
+export function JwtTokenProvider({
+  getTokenCallback,
+  children,
+}: JwtTokenProviderProps) {
+  const [tokenCache, setTokenCache] = useState<Map<string, TokenCacheEntry>>(
+    new Map(),
+  );
+
+  // Load cached tokens from sessionStorage on mount
+  useEffect(() => {
+    try {
+      const newCache = new Map<string, TokenCacheEntry>();
+      const now = Date.now(); // current time in milliseconds
+
+      Object.keys(sessionStorage).forEach((key) => {
+        if (key.startsWith(SESSION_STORAGE_PREFIX)) {
+          const cacheKey = key.replace(SESSION_STORAGE_PREFIX, "");
+          const stored = sessionStorage.getItem(key);
+
+          if (stored) {
+            try {
+              const cacheEntry: TokenCacheEntry = JSON.parse(stored);
+              if (cacheEntry.expiresAt > now) {
+                newCache.set(cacheKey, cacheEntry);
+              } else {
+                sessionStorage.removeItem(key);
+              }
+            } catch {
+              sessionStorage.removeItem(key);
+            }
+          }
+        }
+      });
+
+      if (newCache.size > 0) {
+        setTokenCache(newCache);
+      }
+    } catch {
+      console.warn("Failed to load JWT tokens from sessionStorage");
+    }
+  }, []);
+
+  /**
+   * Get a cached token from the in-memory cache or sessionStorage
+   * @param consumerRef - The consumer reference
+   * @param groupRef - The group reference
+   * @returns The cached token or null if not found
+   */
+  const getCachedToken = useCallback(
+    (consumerRef: string, groupRef: string): string | null => {
+      const cacheKey = createCacheKey(consumerRef, groupRef);
+      const now = Date.now();
+
+      // Check in-memory cache first
+      const cached = tokenCache.get(cacheKey);
+      if (cached && cached.expiresAt > now) {
+        return cached.token;
+      } else {
+        tokenCache.delete(cacheKey);
+      }
+
+      // Check sessionStorage
+      const sessionKey = getSessionStorageKey(cacheKey);
+      const stored = sessionStorage.getItem(sessionKey);
+
+      if (stored) {
+        try {
+          const cacheEntry: TokenCacheEntry = JSON.parse(stored);
+          if (cacheEntry.expiresAt > now) {
+            // Update in-memory cache
+            setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry));
+            return cacheEntry.token;
+          } else {
+            sessionStorage.removeItem(sessionKey);
+          }
+        } catch {
+          sessionStorage.removeItem(sessionKey);
+        }
+      }
+
+      return null;
+    },
+    [tokenCache],
+  );
+
+  const setCachedToken = useCallback(
+    async (consumerRef: string, groupRef: string, token: string) => {
+      const cacheKey = createCacheKey(consumerRef, groupRef);
+      const tokenExpiration = await getTokenExpirationTime(token);
+      const expiresAt =
+        tokenExpiration || Date.now() + DEFAULT_TOKEN_EXPIRATION_TIME;
+
+      const cacheEntry: TokenCacheEntry = { token, expiresAt };
+
+      // Update both in-memory cache
+      setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry));
+
+      // Update sessionStorage
+      try {
+        sessionStorage.setItem(
+          getSessionStorageKey(cacheKey),
+          JSON.stringify(cacheEntry),
+        );
+      } catch {
+        console.warn("Failed to store JWT token in sessionStorage");
+      }
+    },
+    [],
+  );
+
+  const getToken = useCallback(
+    async (consumerRef: string, groupRef: string): Promise<string> => {
+      // Check all caches first
+      const cachedToken = getCachedToken(consumerRef, groupRef);
+      if (cachedToken) {
+        return cachedToken;
+      }
+
+      // Fetch new token if no callback provided
+      if (!getTokenCallback) {
+        throw new Error("JWT token callback not provided");
+      }
+
+      try {
+        const token = await getTokenCallback(consumerRef, groupRef);
+        await setCachedToken(consumerRef, groupRef, token);
+        return token;
+      } catch (error) {
+        console.error("Failed to get JWT token:", error);
+        throw new Error("Failed to get JWT token");
+      }
+    },
+    [getTokenCallback, getCachedToken, setCachedToken],
+  );
+
+  const contextValue: JwtTokenContextValue = {
+    getToken: getTokenCallback ? getToken : undefined,
+  };
+
+  return (
+    <JwtTokenContext.Provider value={contextValue}>
+      {children}
+    </JwtTokenContext.Provider>
+  );
+}
+
+export const useJwtToken = () => {
+  const context = useContext(JwtTokenContext);
+  if (!context) {
+    throw new Error("useJwtToken must be used within a JwtTokenProvider");
+  }
+  return context;
+};