fix(web): sync session state on auth events (#1835)

Author: Uarmagan

packages/web/src/auth/compass/session/SessionProvider.test.tsx

@@ -1,4 +1,4 @@
-import { waitFor } from "@testing-library/react";
+import { act, renderHook, waitFor } from "@testing-library/react";
 import { Subject } from "rxjs";
 import { authSlice } from "@web/ducks/auth/slices/auth.slice";
 import { userMetadataSlice } from "@web/ducks/auth/slices/user-metadata.slice";
@@ -101,8 +101,9 @@ const { session } = require("@web/common/classes/Session") as {
   };
 };
 
-const { sessionInit } =
+const { SessionProvider, sessionInit } =
   require("./SessionProvider") as typeof import("./SessionProvider");
+const { useSession } = require("./useSession") as typeof import("./useSession");
 
 describe("SessionProvider sessionInit", () => {
   beforeEach(() => {
@@ -164,6 +165,28 @@ describe("SessionProvider sessionInit", () => {
     );
     expect(closeStream).toHaveBeenCalledTimes(2);
   });
+
+  it("updates session consumers when SuperTokens creates a session", async () => {
+    getStream.mockReturnValue({} as EventSource);
+    doesSessionExist.mockResolvedValue(false);
+
+    const { result } = renderHook(() => useSession(), {
+      wrapper: SessionProvider,
+    });
+
+    act(() => {
+      result.current.setAuthenticated(false);
+    });
+
+    expect(result.current.authenticated).toBe(false);
+
+    sessionInit();
+    act(() => {
+      session.emit("SESSION_CREATED", { action: "SESSION_CREATED" });
+    });
+
+    expect(result.current.authenticated).toBe(true);
+  });
 });
 
 afterAll(() => {

packages/web/src/auth/compass/session/SessionProvider.tsx

@@ -60,15 +60,25 @@ export const SessionContext = createContext<CompassSession>({
 const authenticated$ = new BehaviorSubject(false);
 let isCheckingSession = false;
 let isSessionInitialized = false;
+let sessionEventVersion = 0;
 
 const $authenticated = authenticated$.pipe(skip(1), distinctUntilChanged());
 
-const handleSessionExists = async () => {
+const handleAuthenticatedSession = () => {
+  authenticated$.next(true);
   markUserAsAuthenticated(getLastKnownEmail());
-  await refreshUserMetadata();
+  void refreshUserMetadata();
+};
+
+const handleSessionExists = () => {
+  handleAuthenticatedSession();
+  if (!sse.getStream()) {
+    sse.openStream();
+  }
 };
 
 const handleSessionMissing = () => {
+  authenticated$.next(false);
   store.dispatch(authSlice.actions.resetAuth());
   store.dispatch(userMetadataSlice.actions.clear(undefined));
   clearGoogleSyncIndicatorOverride();
@@ -81,23 +91,24 @@ async function checkIfSessionExists(): Promise<boolean> {
     return false;
   }
 
-  if (isCheckingSession) return false;
+  if (isCheckingSession) return authenticated$.value;
 
   isCheckingSession = true;
+  const eventVersionAtCheckStart = sessionEventVersion;
 
   try {
     const exists = await session.doesSessionExist();
 
+    if (sessionEventVersion !== eventVersionAtCheckStart) {
+      return authenticated$.value;
+    }
+
     if (exists) {
       handleSessionExists();
-      if (!sse.getStream()) {
-        sse.openStream();
-      }
     } else {
       handleSessionMissing();
     }
 
-    authenticated$.next(exists);
     return exists;
   } catch (error) {
     console.error("Error checking auth status:", error);
@@ -118,22 +129,23 @@ export function sessionInit() {
 
   // No need to unsubscribe as this runs for the lifetime of the app
   session.events.pipe(distinctUntilKeyChanged("action")).subscribe((e) => {
-    void checkIfSessionExists();
-
     switch (e.action) {
       case "REFRESH_SESSION":
       case "SESSION_CREATED":
+        sessionEventVersion += 1;
         // Mark user as authenticated when session is created or refreshed
         // This ensures the flag is set even if markUserAsAuthenticated wasn't called during OAuth
-        markUserAsAuthenticated(getLastKnownEmail());
-        void refreshUserMetadata();
+        handleAuthenticatedSession();
         sse.closeStream();
         sse.openStream();
         break;
       case "SIGN_OUT":
-        store.dispatch(userMetadataSlice.actions.clear(undefined));
+        sessionEventVersion += 1;
+        handleSessionMissing();
         sse.closeStream();
         break;
+      default:
+        void checkIfSessionExists();
     }
   });
 }