Allow passing organizationId when refreshing the session (#2)

* Allow passing `organizationId` when refreshing * Remove accidental code style changes My editor was picking up the wrong formatter program * Read and write organization ID from session storage during refreshes * Run `npm run format` * Move `organizationId` logic under lock acquisition
workos · Aug 9, 2024 · df328df · df328df
1 parent d67d8b3
commit df328df
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/src/create-client.ts b/src/create-client.ts
@@ -12,7 +12,7 @@ import {
   removeSessionData,
   storageKeys,
 } from "./utils";
-import { getRefreshToken } from "./utils/session-data";
+import { getRefreshToken, getClaims } from "./utils/session-data";
 import { RedirectParams } from "./interfaces/create-client-options.interface";
 import Lock from "./vendor/browser-tabs-lock";
 import { RefreshError } from "./utils/authenticate-with-refresh-token";
@@ -28,6 +28,8 @@ type State = "INITIAL" | "AUTHENTICATING" | "AUTHENTICATED" | "ERROR";
 
 const DEFAULT_HOSTNAME = "api.workos.com";
 
+const ORGANIZATION_ID_SESSION_STORAGE_KEY = "workos_organization_id";
+
 export async function createClient(
   clientId: string,
   options: CreateClientOptions = {},
@@ -115,12 +117,17 @@ export async function createClient(
     return user ? (user as User) : null;
   }
 
+  function _getAccessToken() {
+    return memoryStorage.getItem(storageKeys.accessToken) as string | undefined;
+  }
+
   async function getAccessToken() {
     // TODO: should this respect onBeforeAutoRefresh ?
     if (_needsRefresh()) {
       await refreshSession();
     }
-    return memoryStorage.getItem(storageKeys.accessToken) as string | undefined;
+
+    return _getAccessToken();
   }
 
   let _refreshTimer: ReturnType<typeof setTimeout> | undefined;
@@ -209,7 +216,9 @@ export async function createClient(
 
   const REFRESH_LOCK = "WORKOS_REFRESH_SESSION";
 
-  async function refreshSession() {
+  async function refreshSession({
+    organizationId,
+  }: { organizationId?: string } = {}) {
     if (
       _authkitClientState !== "AUTHENTICATED" &&
       _authkitClientState !== "INITIAL"
@@ -220,11 +229,29 @@ export async function createClient(
     const lock = new Lock();
     try {
       _authkitClientState = "AUTHENTICATING";
+
       if (await lock.acquireLock(REFRESH_LOCK)) {
+        if (organizationId) {
+          sessionStorage.setItem(
+            ORGANIZATION_ID_SESSION_STORAGE_KEY,
+            organizationId,
+          );
+        } else {
+          const accessToken = _getAccessToken();
+          if (accessToken) {
+            organizationId = getClaims(accessToken)?.org_id;
+          } else {
+            organizationId =
+              sessionStorage.getItem(ORGANIZATION_ID_SESSION_STORAGE_KEY) ??
+              undefined;
+          }
+        }
+
         const authenticationResponse = await authenticateWithRefreshToken({
           baseUrl: _baseUrl,
           clientId: _clientId,
           refreshToken: getRefreshToken({ devMode }),
+          organizationId,
           useCookie: _useCookie,
         });
 

diff --git a/src/utils/authenticate-with-refresh-token.ts b/src/utils/authenticate-with-refresh-token.ts
@@ -5,15 +5,19 @@ interface AuthenticateWithRefreshTokenOptions {
   baseUrl: string;
   clientId: string;
   refreshToken: string | undefined;
+  organizationId?: string;
   useCookie: boolean;
 }
 
 export class RefreshError extends Error {}
 
-export async function authenticateWithRefreshToken(
-  options: AuthenticateWithRefreshTokenOptions,
-) {
-  const { baseUrl, clientId, refreshToken, useCookie } = options;
+export async function authenticateWithRefreshToken({
+  baseUrl,
+  clientId,
+  refreshToken,
+  organizationId,
+  useCookie,
+}: AuthenticateWithRefreshTokenOptions) {
   const response = await fetch(`${baseUrl}/user_management/authenticate`, {
     method: "POST",
     ...(useCookie && { credentials: "include" }),
@@ -25,6 +29,7 @@ export async function authenticateWithRefreshToken(
       client_id: clientId,
       grant_type: "refresh_token",
       ...(!useCookie && { refresh_token: refreshToken }),
+      organization_id: organizationId,
     }),
   });