feat: add new token fetching via comlink (#475)

* feat: add new token fetching via comlink

Author: ryanbonial

packages/core/src/_exports/index.ts

@@ -13,11 +13,13 @@ export {
   getAuthState,
   getCurrentUserState,
   getDashboardOrganizationId,
+  getIsInDashboardState,
   getLoginUrlState,
   getTokenState,
   type LoggedInAuthState,
   type LoggedOutAuthState,
   type LoggingInAuthState,
+  setAuthToken,
 } from '../auth/authStore'
 export {observeOrganizationVerificationState} from '../auth/getOrganizationVerificationState'
 export {handleAuthCallback} from '../auth/handleAuthCallback'
@@ -32,8 +34,14 @@ export {
   releaseChannel,
 } from '../comlink/controller/comlinkControllerStore'
 export type {ComlinkNodeState} from '../comlink/node/comlinkNodeStore'
+export {getOrCreateNode, releaseNode} from '../comlink/node/comlinkNodeStore'
 export {getNodeState, type NodeState} from '../comlink/node/getNodeState'
-export {type FrameMessage, type WindowMessage} from '../comlink/types'
+export {
+  type FrameMessage,
+  type NewTokenResponseMessage,
+  type RequestNewTokenMessage,
+  type WindowMessage,
+} from '../comlink/types'
 export {type AuthConfig, type AuthProvider} from '../config/authConfig'
 export {
   createDatasetHandle,

packages/core/src/auth/authStore.ts

@@ -280,3 +280,48 @@ export const getDashboardOrganizationId = bindActionGlobally(
   authStore,
   createStateSourceAction(({state: {dashboardContext}}) => dashboardContext?.orgId),
 )
+
+/**
+ * Returns a state source indicating if the SDK is running within a dashboard context.
+ * @public
+ */
+export const getIsInDashboardState = bindActionGlobally(
+  authStore,
+  createStateSourceAction(
+    ({state: {dashboardContext}}) =>
+      // Check if dashboardContext exists and is not empty
+      !!dashboardContext && Object.keys(dashboardContext).length > 0,
+  ),
+)
+
+/**
+ * Action to explicitly set the authentication token.
+ * Used internally by the Comlink token refresh.
+ * @internal
+ */
+export const setAuthToken = bindActionGlobally(authStore, ({state}, token: string | null) => {
+  const currentAuthState = state.get().authState
+  if (token) {
+    // Update state only if the new token is different or currently logged out
+    if (currentAuthState.type !== AuthStateType.LOGGED_IN || currentAuthState.token !== token) {
+      // This state update structure should trigger listeners in clientStore
+      state.set('setToken', {
+        authState: {
+          type: AuthStateType.LOGGED_IN,
+          token: token,
+          // Keep existing user or set to null? Setting to null forces refetch.
+          // Keep existing user to avoid unnecessary refetches if user data is still valid.
+          currentUser:
+            currentAuthState.type === AuthStateType.LOGGED_IN ? currentAuthState.currentUser : null,
+        },
+      })
+    }
+  } else {
+    // Handle setting token to null (logging out)
+    if (currentAuthState.type !== AuthStateType.LOGGED_OUT) {
+      state.set('setToken', {
+        authState: {type: AuthStateType.LOGGED_OUT, isDestroyingSession: false},
+      })
+    }
+  }
+})

packages/core/src/comlink/types.ts

@@ -11,3 +11,21 @@ export type FrameMessage = Message
  * @public
  */
 export type WindowMessage = Message
+
+/**
+ * Message from SDK (iframe) to Parent (dashboard) to request a new token
+ * @internal
+ */
+export type RequestNewTokenMessage = {
+  type: 'dashboard/v1/auth/tokens/create'
+  payload?: undefined
+}
+
+/**
+ * Message from Parent (dashboard) to SDK (iframe) with the new token
+ * @internal
+ */
+export type NewTokenResponseMessage = {
+  type: 'dashboard/v1/auth/tokens/create'
+  payload: {token: string | null; error?: string}
+}

packages/react/src/_exports/sdk-react.ts

@@ -4,6 +4,7 @@
 export {AuthBoundary, type AuthBoundaryProps} from '../components/auth/AuthBoundary'
 export {SanityApp, type SanityAppProps} from '../components/SanityApp'
 export {SDKProvider, type SDKProviderProps} from '../components/SDKProvider'
+export {ComlinkTokenRefreshProvider} from '../context/ComlinkTokenRefresh'
 export {ResourceProvider, type ResourceProviderProps} from '../context/ResourceProvider'
 export {useAuthState} from '../hooks/auth/useAuthState'
 export {useAuthToken} from '../hooks/auth/useAuthToken'

packages/react/src/components/auth/AuthBoundary.test.tsx

@@ -1,4 +1,4 @@
-import {AuthStateType} from '@sanity/sdk'
+import {AuthStateType, type SanityConfig} from '@sanity/sdk'
 import {render, screen, waitFor} from '@testing-library/react'
 import React from 'react'
 import {type FallbackProps} from 'react-error-boundary'
@@ -8,6 +8,7 @@ import {ResourceProvider} from '../../context/ResourceProvider'
 import {useAuthState} from '../../hooks/auth/useAuthState'
 import {useLoginUrl} from '../../hooks/auth/useLoginUrl'
 import {useVerifyOrgProjects} from '../../hooks/auth/useVerifyOrgProjects'
+import {useSanityInstance} from '../../hooks/context/useSanityInstance'
 import {AuthBoundary} from './AuthBoundary'
 
 // Mock hooks
@@ -22,6 +23,9 @@ vi.mock('../../hooks/auth/useHandleAuthCallback', () => ({
 vi.mock('../../hooks/auth/useLogOut', () => ({
   useLogOut: vi.fn(() => async () => {}),
 }))
+vi.mock('../../hooks/context/useSanityInstance', () => ({
+  useSanityInstance: vi.fn(),
+}))
 
 // Mock AuthError throwing scenario
 vi.mock('./AuthError', async (importOriginal) => {
@@ -105,8 +109,27 @@ describe('AuthBoundary', () => {
   const mockUseAuthState = vi.mocked(useAuthState)
   const mockUseLoginUrl = vi.mocked(useLoginUrl)
   const mockUseVerifyOrgProjects = vi.mocked(useVerifyOrgProjects)
+  const mockUseSanityInstance = vi.mocked(useSanityInstance)
   const testProjectIds = ['proj-test'] // Example project ID for tests
 
+  // Mock Sanity instance
+  const mockSanityInstance = {
+    instanceId: 'test-instance-id',
+    config: {
+      projectId: 'test-project',
+      dataset: 'test-dataset',
+    },
+    isDisposed: () => false,
+    dispose: () => {},
+    onDispose: () => () => {},
+    getParent: () => undefined,
+    createChild: (config: SanityConfig) => ({
+      ...mockSanityInstance,
+      config: {...mockSanityInstance.config, ...config},
+    }),
+    match: () => undefined,
+  }
+
   beforeEach(() => {
     vi.clearAllMocks()
     consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {})
@@ -116,6 +139,8 @@ describe('AuthBoundary', () => {
     mockUseLoginUrl.mockReturnValue('http://example.com/login')
     // Default mock for useVerifyOrgProjects - returns null (no error)
     mockUseVerifyOrgProjects.mockImplementation(() => null)
+    // Mock useSanityInstance to return our mock instance
+    mockUseSanityInstance.mockReturnValue(mockSanityInstance)
   })
 
   afterEach(() => {
@@ -145,9 +170,7 @@ describe('AuthBoundary', () => {
       isExchangingToken: false,
     })
     const {container} = render(
-      <ResourceProvider fallback={null}>
-        <AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>
-      </ResourceProvider>,
+      <AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>,
     )
 
     // The callback screen renders null check that it renders nothing
@@ -161,11 +184,7 @@ describe('AuthBoundary', () => {
       currentUser: null,
       token: 'exampleToken',
     })
-    render(
-      <ResourceProvider fallback={null}>
-        <AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>
-      </ResourceProvider>,
-    )
+    render(<AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>)
 
     expect(screen.getByText('Protected Content')).toBeInTheDocument()
   })
@@ -175,11 +194,7 @@ describe('AuthBoundary', () => {
       type: AuthStateType.ERROR,
       error: new Error('test error'),
     })
-    render(
-      <ResourceProvider fallback={null}>
-        <AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>
-      </ResourceProvider>,
-    )
+    render(<AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>)
 
     // The AuthBoundary should throw an AuthError internally
     // and then display the LoginError component as the fallback.
@@ -192,11 +207,7 @@ describe('AuthBoundary', () => {
   })
 
   it('renders children when logged in and org verification passes', () => {
-    render(
-      <ResourceProvider fallback={null}>
-        <AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>
-      </ResourceProvider>,
-    )
+    render(<AuthBoundary projectIds={testProjectIds}>Protected Content</AuthBoundary>)
     expect(screen.getByText('Protected Content')).toBeInTheDocument()
   })
 

packages/react/src/components/auth/AuthBoundary.tsx

@@ -2,6 +2,7 @@ import {AuthStateType} from '@sanity/sdk'
 import {useEffect, useMemo} from 'react'
 import {ErrorBoundary, type FallbackProps} from 'react-error-boundary'
 
+import {ComlinkTokenRefreshProvider} from '../../context/ComlinkTokenRefresh'
 import {useAuthState} from '../../hooks/auth/useAuthState'
 import {useLoginUrl} from '../../hooks/auth/useLoginUrl'
 import {useVerifyOrgProjects} from '../../hooks/auth/useVerifyOrgProjects'
@@ -111,9 +112,11 @@ export function AuthBoundary({
   }, [LoginErrorComponent])
 
   return (
-    <ErrorBoundary FallbackComponent={FallbackComponent}>
-      <AuthSwitch {...props} />
-    </ErrorBoundary>
+    <ComlinkTokenRefreshProvider>
+      <ErrorBoundary FallbackComponent={FallbackComponent}>
+        <AuthSwitch {...props} />
+      </ErrorBoundary>
+    </ComlinkTokenRefreshProvider>
   )
 }
 

packages/react/src/components/auth/LoginError.tsx

@@ -19,7 +19,15 @@ export type LoginErrorProps = FallbackProps
  * @alpha
  */
 export function LoginError({error, resetErrorBoundary}: LoginErrorProps): React.ReactNode {
-  if (!(error instanceof AuthError || error instanceof ConfigurationError)) throw error
+  if (
+    !(
+      error instanceof AuthError ||
+      error instanceof ConfigurationError ||
+      error instanceof ClientError
+    )
+  )
+    throw error
+
   const logout = useLogOut()
   const authState = useAuthState()
 
@@ -33,11 +41,11 @@ export function LoginError({error, resetErrorBoundary}: LoginErrorProps): React.
   }, [logout, resetErrorBoundary])
 
   useEffect(() => {
-    if (authState.type === AuthStateType.ERROR && authState.error instanceof ClientError) {
-      if (authState.error.statusCode === 401) {
+    if (error instanceof ClientError) {
+      if (error.statusCode === 401) {
         handleRetry()
-      } else if (authState.error.statusCode === 404) {
-        const errorMessage = authState.error.response.body.message || ''
+      } else if (error.statusCode === 404) {
+        const errorMessage = error.response.body.message || ''
         if (errorMessage.startsWith('Session with sid') && errorMessage.endsWith('not found')) {
           setAuthErrorMessage('The session ID is invalid or expired.')
         } else {