fix: retry session refresh when organizationId is missing (#2564)

When a user signs up, there's a race condition where the webhook that
creates the user's organization may not have processed yet. This results
in the JWT not having an org_id claim, causing WorkOS API calls to fail
with "organization_id must be a string".

This change adds retry logic to requireAuth() that:
- Detects when organizationId is missing from the session
- Refreshes the session up to 3 times with 1s delay between attempts
- Gives the webhook time to process and assign the organization
- Throws a clear error if org is still missing after retries

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: motatoes

ui/src/api/helpers.ts

@@ -1,4 +1,37 @@
-import { withAuth } from '@/authkit/ssr/session';
+import { withAuth, getSessionFromCookie, saveSession } from '@/authkit/ssr/session';
+import { getWorkOS } from '@/authkit/ssr/workos';
+import { getConfig } from '@/authkit/ssr/config';
+import { decodeJwt } from 'jose';
+import type { AccessToken } from '@workos-inc/node';
+
+const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+
+/**
+ * Refreshes the session and returns the new organizationId if available.
+ * This is used when the session doesn't have an org_id yet (e.g., webhook delay).
+ */
+async function refreshSessionForOrg(): Promise<string | undefined> {
+  const session = await getSessionFromCookie();
+  if (!session?.refreshToken) {
+    return undefined;
+  }
+
+  try {
+    const { accessToken, refreshToken, user, impersonator } =
+      await getWorkOS().userManagement.authenticateWithRefreshToken({
+        clientId: getConfig('clientId'),
+        refreshToken: session.refreshToken,
+      });
+
+    await saveSession({ accessToken, refreshToken, user, impersonator });
+
+    const { org_id: organizationId } = decodeJwt<AccessToken>(accessToken);
+    return organizationId;
+  } catch (error) {
+    console.error('Failed to refresh session for org:', error);
+    return undefined;
+  }
+}
 
 function redirectWithFallback(redirectUri: string, headers?: Headers) {
     const newHeaders = headers ? new Headers(headers) : new Headers();
@@ -18,16 +51,48 @@ function redirectWithFallback(redirectUri: string, headers?: Headers) {
  * Requires authentication and returns session-derived user identity.
  * Use this instead of accepting userId/organizationId from client input
  * to prevent IDOR vulnerabilities.
+ *
+ * If organizationId is missing (e.g., webhook hasn't processed yet),
+ * this will retry by refreshing the session up to 3 times.
  */
 export async function requireAuth() {
   const auth = await withAuth();
   if (!auth.user) {
     throw new Error('Unauthorized');
   }
+
+  let organizationId = auth.organizationId;
+
+  // If no organizationId, the webhook may not have processed yet.
+  // Retry by refreshing the session to get updated JWT claims.
+  if (!organizationId) {
+    const MAX_RETRIES = 3;
+    const RETRY_DELAY_MS = 1000;
+
+    for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+      console.log(`No organizationId in session, refreshing (attempt ${attempt}/${MAX_RETRIES})...`);
+
+      organizationId = await refreshSessionForOrg();
+
+      if (organizationId) {
+        console.log(`Organization found after ${attempt} refresh attempt(s): ${organizationId}`);
+        break;
+      }
+
+      if (attempt < MAX_RETRIES) {
+        await sleep(RETRY_DELAY_MS);
+      }
+    }
+  }
+
+  if (!organizationId) {
+    throw new Error('No organization available. Please try logging out and back in.');
+  }
+
   return {
     userId: auth.user.id,
     email: auth.user.email!,
-    organizationId: auth.organizationId!,
+    organizationId,
   };
 }