fix: auth race condition when using the SDK in the Studio structure (#646)

* fix: auth race condition when using the SDK in the Studio structure
* fix(auth): fix cookie authentication in studio mode
* fix(auth): do not redirect to the login if Studio mode is enabled

Author: ryanbonial

packages/core/src/auth/authStore.test.ts

@@ -237,6 +237,7 @@ describe('authStore', () => {
     it('sets to logged in using studio token when studio mode is enabled and token exists', () => {
       const studioToken = 'studio-token'
       const projectId = 'studio-project'
+      const studioStorageKey = `__studio_auth_token_${projectId}`
       const mockStorage = {
         getItem: vi.fn(),
         setItem: vi.fn(),
@@ -252,40 +253,38 @@ describe('authStore', () => {
       })
 
       const {authState, options} = authStore.getInitialState(instance)
-      expect(getStudioTokenFromLocalStorage).toHaveBeenCalledWith(mockStorage, projectId)
+      expect(getStudioTokenFromLocalStorage).toHaveBeenCalledWith(mockStorage, studioStorageKey)
       expect(authState).toMatchObject({type: AuthStateType.LOGGED_IN, token: studioToken})
       expect(options.authMethod).toBe('localstorage')
     })
 
-    it('checks for cookie auth when studio mode is enabled and no studio token exists', async () => {
-      vi.useFakeTimers()
+    it('checks for cookie auth during initialize when studio mode is enabled and no studio token exists', () => {
       const projectId = 'studio-project'
+      const studioStorageKey = `__studio_auth_token_${projectId}`
       const mockStorage = {
         getItem: vi.fn(),
         setItem: vi.fn(),
         removeItem: vi.fn(),
       } as unknown as Storage // Mock storage
       vi.mocked(getStudioTokenFromLocalStorage).mockReturnValue(null)
-      // Mock cookie check to return true asynchronously
       vi.mocked(checkForCookieAuth).mockResolvedValue(true)
 
       instance = createSanityInstance({
         projectId,
         dataset: 'd',
         studioMode: {enabled: true},
-        auth: {storageArea: mockStorage}, // Provide mock storage
+        auth: {storageArea: mockStorage},
       })
 
-      // Initial state might be logged out before the async check completes
+      // Verify initial state without async cookie probe
       const {authState: initialAuthState} = authStore.getInitialState(instance)
-      expect(initialAuthState.type).toBe(AuthStateType.LOGGED_OUT) // Or potentially logging in depending on other factors
-      expect(getStudioTokenFromLocalStorage).toHaveBeenCalledWith(mockStorage, projectId)
-      expect(checkForCookieAuth).toHaveBeenCalledWith(projectId, expect.any(Function))
+      expect(initialAuthState.type).toBe(AuthStateType.LOGGED_OUT)
+      expect(getStudioTokenFromLocalStorage).toHaveBeenCalledWith(mockStorage, studioStorageKey)
 
-      // Wait for the promise in getInitialState to resolve
-      await vi.runAllTimersAsync()
+      // Trigger store creation + initialize
+      getAuthState(instance)
 
-      vi.useRealTimers()
+      expect(checkForCookieAuth).toHaveBeenCalledWith(projectId, expect.any(Function))
     })
 
     it('falls back to default auth (storage token) when studio mode is disabled', () => {
@@ -573,14 +572,7 @@ describe('authStore', () => {
       expect(initialOrgId.getCurrent()).toBe('initial-org-id')
 
       // Call handleCallback with the callback URL
-      await handleAuthCallback(instance, callbackUrl) // Use await as handleAuthCallback is async
-
-      // Wait for the state update to be reflected in the selector
-      await vi.waitUntil(
-        () => getDashboardOrganizationId(instance).getCurrent() === 'callback-org-id',
-      )
-      // Add a microtask yield just in case
-      await new Promise((resolve) => setTimeout(resolve, 0))
+      await handleAuthCallback(instance, callbackUrl)
 
       // Check that the orgId from the callback context is now set
       const finalOrgId = getDashboardOrganizationId(instance)

packages/core/src/auth/authStore.ts

@@ -67,7 +67,11 @@ export interface DashboardContext {
   orgId?: string
 }
 
-type AuthMethodOptions = 'localstorage' | 'cookie' | undefined
+/**
+ * The method of authentication used.
+ * @internal
+ */
+export type AuthMethodOptions = 'localstorage' | 'cookie' | undefined
 
 let tokenRefresherRunning = false
 
@@ -105,7 +109,8 @@ export const authStore = defineStore<AuthStoreState>({
     } = instance.config.auth ?? {}
     let storageArea = instance.config.auth?.storageArea
 
-    const storageKey = `__sanity_auth_token`
+    let storageKey = `__sanity_auth_token`
+    const studioModeEnabled = instance.config.studioMode?.enabled
 
     // This login URL will only be used for local development
     let loginDomain = 'https://www.sanity.io'
@@ -149,23 +154,21 @@ export const authStore = defineStore<AuthStoreState>({
       console.error('Failed to parse dashboard context from initial location:', err)
     }
 
-    if (!isInDashboard) {
+    if (!isInDashboard || studioModeEnabled) {
       // If not in dashboard, use the storage area from the config
+      // If studio mode is enabled, use the local storage area (default)
       storageArea = storageArea ?? getDefaultStorage()
     }
 
     let token: string | null
     let authMethod: AuthMethodOptions
-    if (instance.config.studioMode?.enabled) {
-      token = getStudioTokenFromLocalStorage(storageArea, instance.config.projectId)
+    if (studioModeEnabled) {
+      // In studio mode, always use the studio-specific storage key and subscribe to it
+      const studioStorageKey = `__studio_auth_token_${instance.config.projectId ?? ''}`
+      storageKey = studioStorageKey
+      token = getStudioTokenFromLocalStorage(storageArea, studioStorageKey)
       if (token) {
         authMethod = 'localstorage'
-      } else {
-        checkForCookieAuth(instance.config.projectId, clientFactory).then((isCookieAuthEnabled) => {
-          if (isCookieAuthEnabled) {
-            authMethod = 'cookie'
-          }
-        })
       }
     } else {
       token = getTokenFromStorage(storageArea, storageKey)
@@ -177,14 +180,16 @@ export const authStore = defineStore<AuthStoreState>({
     let authState: AuthState
     if (providedToken) {
       authState = {type: AuthStateType.LOGGED_IN, token: providedToken, currentUser: null}
+    } else if (token && studioModeEnabled) {
+      authState = {type: AuthStateType.LOGGED_IN, token: token ?? '', currentUser: null}
     } else if (
       getAuthCode(callbackUrl, initialLocationHref) ||
       getTokenFromLocation(initialLocationHref)
     ) {
       authState = {type: AuthStateType.LOGGING_IN, isExchangingToken: false}
       // Note: dashboardContext from the callback URL can be set later in handleAuthCallback too
-    } else if (token && !isInDashboard) {
-      // Only use token from storage if NOT running in dashboard
+    } else if (token && !isInDashboard && !studioModeEnabled) {
+      // Only use token from storage if NOT running in dashboard and studio mode is not enabled
       authState = {type: AuthStateType.LOGGED_IN, token, currentUser: null}
     } else {
       // Default to logged out if no provided token, not handling callback,
@@ -212,11 +217,37 @@ export const authStore = defineStore<AuthStoreState>({
   initialize(context) {
     const subscriptions: Subscription[] = []
     subscriptions.push(subscribeToStateAndFetchCurrentUser(context))
-
-    if (context.state.get().options?.storageArea) {
+    const storageArea = context.state.get().options?.storageArea
+    if (storageArea) {
       subscriptions.push(subscribeToStorageEventsAndSetToken(context))
     }
 
+    // If in Studio mode with no local token, resolve cookie auth asynchronously
+    try {
+      const {instance, state} = context
+      const studioModeEnabled = !!instance.config.studioMode?.enabled
+      const token: string | null =
+        state.get().authState?.type === AuthStateType.LOGGED_IN
+          ? (state.get().authState as LoggedInAuthState).token
+          : null
+      if (studioModeEnabled && !token) {
+        const projectId = instance.config.projectId
+        const clientFactory = state.get().options.clientFactory
+        checkForCookieAuth(projectId, clientFactory).then((isCookieAuthEnabled) => {
+          if (!isCookieAuthEnabled) return
+          state.set('enableCookieAuth', (prev) => ({
+            options: {...prev.options, authMethod: 'cookie'},
+            authState:
+              prev.authState.type === AuthStateType.LOGGED_IN
+                ? prev.authState
+                : {type: AuthStateType.LOGGED_IN, token: '', currentUser: null},
+          }))
+        })
+      }
+    } catch {
+      // best-effort cookie detection
+    }
+
     if (!tokenRefresherRunning) {
       tokenRefresherRunning = true
       subscriptions.push(refreshStampedToken(context))

packages/core/src/auth/studioModeAuth.test.ts

@@ -96,27 +96,27 @@ describe('getStudioTokenFromLocalStorage', () => {
     expect(getTokenFromStorageSpy).not.toHaveBeenCalled()
   })
 
-  it('should return null if projectId is undefined', () => {
+  it('should return null if storageKey is undefined', () => {
     const result = getStudioTokenFromLocalStorage(storageArea, undefined)
     expect(result).toBeNull()
     expect(getTokenFromStorageSpy).not.toHaveBeenCalled()
   })
 
   it('should call getTokenFromStorage with correct key', () => {
     getTokenFromStorageSpy.mockReturnValue(null) // Assume token not found for this test
-    getStudioTokenFromLocalStorage(storageArea, projectId)
+    getStudioTokenFromLocalStorage(storageArea, studioStorageKey)
     expect(getTokenFromStorageSpy).toHaveBeenCalledWith(storageArea, studioStorageKey)
   })
 
   it('should return the token if found in storage', () => {
     getTokenFromStorageSpy.mockReturnValue(mockToken)
-    const result = getStudioTokenFromLocalStorage(storageArea, projectId)
+    const result = getStudioTokenFromLocalStorage(storageArea, studioStorageKey)
     expect(result).toBe(mockToken)
   })
 
   it('should return null if token is not found in storage', () => {
     getTokenFromStorageSpy.mockReturnValue(null)
-    const result = getStudioTokenFromLocalStorage(storageArea, projectId)
+    const result = getStudioTokenFromLocalStorage(storageArea, studioStorageKey)
     expect(result).toBeNull()
   })
 })

packages/core/src/auth/studioModeAuth.ts

@@ -33,17 +33,16 @@ export async function checkForCookieAuth(
 /**
  * Attempts to retrieve a studio token from local storage.
  * @param storageArea - The storage area to retrieve the token from.
- * @param projectId - The project ID to retrieve the token for.
+ * @param storageKey - The storage key to retrieve the token from.
  * @returns The studio token or null if it does not exist.
  * @internal
  */
 export function getStudioTokenFromLocalStorage(
   storageArea: Storage | undefined,
-  projectId: string | undefined,
+  storageKey: string | undefined,
 ): string | null {
-  if (!storageArea || !projectId) return null
-  const studioStorageKey = `__studio_auth_token_${projectId}`
-  const token = getTokenFromStorage(storageArea, studioStorageKey)
+  if (!storageArea || !storageKey) return null
+  const token = getTokenFromStorage(storageArea, storageKey)
   if (token) {
     return token
   }

packages/core/src/auth/subscribeToStateAndFetchCurrentUser.ts

@@ -4,32 +4,43 @@ import {distinctUntilChanged, filter, map, type Subscription, switchMap} from 'r
 import {type StoreContext} from '../store/defineStore'
 import {DEFAULT_API_VERSION, REQUEST_TAG_PREFIX} from './authConstants'
 import {AuthStateType} from './authStateType'
-import {type AuthState, type AuthStoreState} from './authStore'
+import {type AuthMethodOptions, type AuthState, type AuthStoreState} from './authStore'
 
 export const subscribeToStateAndFetchCurrentUser = ({
   state,
+  instance,
 }: StoreContext<AuthStoreState>): Subscription => {
   const {clientFactory, apiHost} = state.get().options
+  const useProjectHostname = !!instance.config.studioMode?.enabled
+  const projectId = instance.config.projectId
 
   const currentUser$ = state.observable
     .pipe(
-      map(({authState}) => authState),
+      map(({authState, options}) => ({authState, authMethod: options.authMethod})),
       filter(
-        (authState): authState is Extract<AuthState, {type: AuthStateType.LOGGED_IN}> =>
-          authState.type === AuthStateType.LOGGED_IN && !authState.currentUser,
+        (
+          value,
+        ): value is {
+          authState: Extract<AuthState, {type: AuthStateType.LOGGED_IN}>
+          authMethod: AuthMethodOptions
+        } => value.authState.type === AuthStateType.LOGGED_IN && !value.authState.currentUser,
+      ),
+      map((value) => ({token: value.authState.token, authMethod: value.authMethod})),
+      distinctUntilChanged(
+        (prev, curr) => prev.token === curr.token && prev.authMethod === curr.authMethod,
       ),
-      map((authState) => authState.token),
-      distinctUntilChanged(),
     )
     .pipe(
-      map((token) =>
+      map(({token, authMethod}) =>
         clientFactory({
           apiVersion: DEFAULT_API_VERSION,
           requestTagPrefix: REQUEST_TAG_PREFIX,
-          token,
+          token: authMethod === 'cookie' ? undefined : token,
           ignoreBrowserTokenWarning: true,
-          useProjectHostname: false,
+          useProjectHostname,
           useCdn: false,
+          ...(authMethod === 'cookie' ? {withCredentials: true} : {}),
+          ...(useProjectHostname && projectId ? {projectId} : {}),
           ...(apiHost && {apiHost}),
         }),
       ),