feat: add new token fetching via comlink

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'
@@ -33,7 +35,12 @@ export {
 } from '../comlink/controller/comlinkControllerStore'
 export type {ComlinkNodeState} from '../comlink/node/comlinkNodeStore'
 export {getOrCreateNode, releaseNode} from '../comlink/node/comlinkNodeStore'
-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/client/clientStore.ts

@@ -38,7 +38,10 @@ const allowedKeys = Object.keys({
   apiVersion: null,
   requestTagPrefix: null,
   useProjectHostname: null,
-} satisfies Record<keyof ClientOptions, null>) as (keyof ClientOptions)[]
+} satisfies Record<keyof Omit<ClientOptions, 'middleware'>, null>) as (keyof Omit<
+  ClientOptions,
+  'middleware'
+>)[]
 
 const DEFAULT_CLIENT_CONFIG: ClientConfig = {
   apiVersion: DEFAULT_API_VERSION,
@@ -136,8 +139,8 @@ const getClientConfigKey = (options: ClientOptions) => JSON.stringify(pick(optio
  */
 export const getClient = bindActionGlobally(
   clientStore,
-  ({state, instance}, options: ClientOptions) => {
-    // Check for disallowed keys
+  ({state, instance}, options: ClientOptions, middleware?: unknown[]) => {
+    // Check for disallowed keys (excluding middleware)
     const providedKeys = Object.keys(options) as (keyof ClientOptions)[]
     const disallowedKeys = providedKeys.filter((key) => !allowedKeys.includes(key))
 
@@ -155,14 +158,16 @@ export const getClient = bindActionGlobally(
     const dataset = options.dataset ?? instance.config.dataset
     const apiHost = options.apiHost ?? instance.config.auth?.apiHost
 
-    const effectiveOptions: ClientOptions = {
+    const effectiveOptions: ClientConfig = {
       ...DEFAULT_CLIENT_CONFIG,
       ...((options.scope === 'global' || !projectId) && {useProjectHostname: false}),
       token: authMethod === 'cookie' ? undefined : (tokenFromState ?? undefined),
       ...options,
       ...(projectId && {projectId}),
       ...(dataset && {dataset}),
       ...(apiHost && {apiHost}),
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      ...(middleware && {middleware: middleware as any}),
     }
 
     if (effectiveOptions.token === null || typeof effectiveOptions.token === 'undefined') {
@@ -174,11 +179,12 @@ export const getClient = bindActionGlobally(
       delete effectiveOptions.withCredentials
     }
 
-    const key = getClientConfigKey(effectiveOptions)
+    // Key generation uses only standard options passed to getClient
+    const key = getClientConfigKey(options) // Use original options for key
 
     if (clients[key]) return clients[key]
 
-    const client = createClient(effectiveOptions)
+    const client = createClient(effectiveOptions) // Pass options including middleware here
     state.set('addClient', (prev) => ({clients: {...prev.clients, [key]: client}}))
 
     return client

packages/core/src/comlink/types.ts

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

packages/react/src/context/ComlinkTokenRefresh.tsx

@@ -0,0 +1,186 @@
+import {type Status as ComlinkStatus} from '@sanity/comlink' // Import Status as ComlinkStatus
+import {
+  type FrameMessage,
+  getIsInDashboardState, // <-- Import the new selector
+  type NewTokenResponseMessage,
+  type RequestNewTokenMessage,
+  setAuthToken,
+  type WindowMessage,
+} from '@sanity/sdk'
+import React, {
+  createContext,
+  type PropsWithChildren,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react'
+
+import {useWindowConnection} from '../hooks/comlink/useWindowConnection' // <-- Correct path to useWindowConnection
+import {useSanityInstance} from '../hooks/context/useSanityInstance' // Import the instance hook
+
+// Define specific message types extending the base types for clarity
+type SdkParentComlinkMessage = NewTokenResponseMessage | WindowMessage // Messages received by SDK
+type SdkChildComlinkMessage = RequestNewTokenMessage | FrameMessage // Messages sent by SDK
+
+export type ComlinkTokenRefreshConfig = {
+  /** Enable the Comlink token refresh feature. Defaults to false */
+  enabled?: boolean
+  /** Unique name for this SDK iframe instance's Comlink endpoint. Defaults to 'sanity-sdk-iframe' */
+  comlinkName?: string
+  /** The name of the parent window's Comlink endpoint to connect to. Defaults to 'sanity-dashboard-parent' */
+  parentName?: string
+  /** Timeout in ms for waiting for a token response from the parent. Defaults to 15000ms */
+  responseTimeout?: number
+}
+
+interface ComlinkTokenRefreshContextValue {
+  requestNewToken: () => void
+  isTokenRefreshInProgress: React.MutableRefObject<boolean>
+  comlinkStatus: ComlinkStatus
+  isEnabled: boolean
+}
+
+const ComlinkTokenRefreshContext = createContext<ComlinkTokenRefreshContextValue | null>(null)
+
+const DEFAULT_COMLINK_NAME = 'sanity-sdk-iframe'
+const DEFAULT_PARENT_NAME = 'sanity-dashboard-parent'
+const DEFAULT_RESPONSE_TIMEOUT = 15000 // 15 seconds
+
+export const ComlinkTokenRefreshProvider: React.FC<
+  PropsWithChildren<ComlinkTokenRefreshConfig>
+> = ({
+  children,
+  enabled = false, // Default to disabled
+  comlinkName = DEFAULT_COMLINK_NAME,
+  parentName = DEFAULT_PARENT_NAME,
+  responseTimeout = DEFAULT_RESPONSE_TIMEOUT,
+}) => {
+  const instance = useSanityInstance() // Get the instance
+  const [comlinkStatus, setComlinkStatus] = useState<ComlinkStatus>('idle')
+  const isTokenRefreshInProgress = useRef(false)
+  const timeoutRef = useRef<NodeJS.Timeout | null>(null)
+
+  // Get the dashboard status value once
+  const isInDashboard = useMemo(() => getIsInDashboardState(instance).getCurrent(), [instance])
+
+  // Determine if the feature should be active
+  const isDashboardComlinkEnabled = enabled && isInDashboard
+
+  const clearRefreshTimeout = useCallback(() => {
+    if (timeoutRef.current) {
+      clearTimeout(timeoutRef.current)
+      timeoutRef.current = null
+    }
+  }, [])
+
+  // Conditionally connect only if operational
+  const connectionOptions = useMemo(
+    () =>
+      isDashboardComlinkEnabled
+        ? {
+            name: comlinkName,
+            connectTo: parentName,
+            onStatus: setComlinkStatus,
+            onMessage: {
+              'dashboard/v1/auth/tokens/create': (data: NewTokenResponseMessage['payload']) => {
+                clearRefreshTimeout()
+                // Re-check isDashboardComlinkEnabled in case props/context changed
+                if (!isDashboardComlinkEnabled) return
+
+                if (data.token) {
+                  setAuthToken(instance, data.token)
+                }
+                isTokenRefreshInProgress.current = false
+              },
+            },
+          }
+        : undefined,
+    // Add isDashboardComlinkEnabled to dependency array
+    [
+      isDashboardComlinkEnabled,
+      comlinkName,
+      parentName,
+      setComlinkStatus,
+      clearRefreshTimeout,
+      instance,
+    ],
+  )
+
+  const {sendMessage} = useWindowConnection<SdkParentComlinkMessage, SdkChildComlinkMessage>(
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    connectionOptions as any,
+  )
+
+  const requestNewToken = useCallback(() => {
+    // Check if operational before proceeding
+    if (!isDashboardComlinkEnabled) {
+      return
+    }
+    if (comlinkStatus !== 'connected') {
+      return
+    }
+    if (isTokenRefreshInProgress.current) {
+      return
+    }
+
+    isTokenRefreshInProgress.current = true
+    clearRefreshTimeout()
+
+    timeoutRef.current = setTimeout(() => {
+      if (isTokenRefreshInProgress.current) {
+        isTokenRefreshInProgress.current = false
+      }
+      timeoutRef.current = null
+    }, responseTimeout)
+
+    try {
+      const messageType = 'dashboard/v1/auth/tokens/create' as const
+      sendMessage(messageType)
+    } catch {
+      isTokenRefreshInProgress.current = false
+      clearRefreshTimeout()
+    }
+  }, [isDashboardComlinkEnabled, sendMessage, comlinkStatus, responseTimeout, clearRefreshTimeout])
+
+  const contextValue = useMemo(
+    () => ({
+      requestNewToken,
+      isTokenRefreshInProgress,
+      comlinkStatus,
+      // Expose isDashboardComlinkEnabled instead of isEnabled?
+      // Let's stick to isEnabled prop for clarity, checks happen internally.
+      isEnabled: enabled,
+    }),
+    [requestNewToken, isTokenRefreshInProgress, comlinkStatus, enabled],
+  )
+
+  useEffect(() => {
+    return () => {
+      clearRefreshTimeout()
+    }
+  }, [clearRefreshTimeout])
+
+  return (
+    <ComlinkTokenRefreshContext.Provider value={contextValue}>
+      {children}
+    </ComlinkTokenRefreshContext.Provider>
+  )
+}
+
+export const useComlinkTokenRefresh = (): ComlinkTokenRefreshContextValue => {
+  const context = useContext(ComlinkTokenRefreshContext)
+  if (!context) {
+    return {
+      requestNewToken: () => {
+        /* console.warn(...) */
+      },
+      isTokenRefreshInProgress: {current: false},
+      comlinkStatus: 'idle',
+      isEnabled: false,
+    }
+  }
+  return context
+}