refactor(web): add prefetched query atom (#38572)

web/app/components/custom/custom-page/__tests__/index.spec.tsx

@@ -10,7 +10,6 @@ import { Plan } from '@/app/components/billing/type'
 import {
   initialLangGeniusVersionInfo,
   initialWorkspaceInfo,
-  userProfilePlaceholder,
 } from '@/context/app-context-defaults'
 import { useModalContext } from '@/context/modal-context'
 import { useProviderContext } from '@/context/provider-context'
@@ -60,6 +59,15 @@ vi.mock('@langgenius/dify-ui/toast', () => ({
 const mockUseProviderContext = vi.mocked(useProviderContext)
 const mockUseModalContext = vi.mocked(useModalContext)
 
+const testUserProfile = {
+  id: '',
+  name: '',
+  email: '',
+  avatar: '',
+  avatar_url: '',
+  is_password_set: false,
+}
+
 const createProviderContext = ({
   enableBilling = false,
   planType = Plan.professional,
@@ -77,7 +85,7 @@ const createProviderContext = ({
 }
 
 const createAppContextValue = (): AppContextStateMockState => ({
-  userProfile: userProfilePlaceholder,
+  userProfile: testUserProfile,
   mutateUserProfile: vi.fn(),
   currentWorkspace: {
     ...initialWorkspaceInfo,

web/app/components/custom/custom-web-app-brand/hooks/__tests__/use-web-app-brand.spec.tsx

@@ -12,7 +12,6 @@ import { Plan } from '@/app/components/billing/type'
 import {
   initialLangGeniusVersionInfo,
   initialWorkspaceInfo,
-  userProfilePlaceholder,
 } from '@/context/app-context-defaults'
 import { useProviderContext } from '@/context/provider-context'
 import { updateCurrentWorkspace } from '@/service/common'
@@ -80,6 +79,15 @@ const mockUseProviderContext = vi.mocked(useProviderContext)
 const mockImageUpload = vi.mocked(imageUpload)
 const mockGetImageUploadErrorMessage = vi.mocked(getImageUploadErrorMessage)
 
+const testUserProfile = {
+  id: '',
+  name: '',
+  email: '',
+  avatar: '',
+  avatar_url: '',
+  is_password_set: false,
+}
+
 const createProviderContext = ({
   enableBilling = false,
   planType = Plan.professional,
@@ -110,7 +118,7 @@ const createAppContextValue = (overrides: Partial<AppContextStateMockState> = {}
   }
 
   return {
-    userProfile: userProfilePlaceholder,
+    userProfile: testUserProfile,
     mutateUserProfile: vi.fn(),
     isCurrentWorkspaceManager: true,
     isCurrentWorkspaceOwner: false,

web/context/__tests__/console-bootstrap.spec.tsx

@@ -59,7 +59,7 @@ const mockUserProfileResponseState = vi.hoisted(() => ({
       currentEnv: 'cloud',
     },
   } as {
-    profile?: {
+    profile: {
       id: string
       name: string
       email: string
@@ -259,6 +259,9 @@ function createTestQueryClient() {
 
 function renderConsoleBootstrap() {
   const queryClient = createTestQueryClient()
+  queryClient.setQueryData(['user-profile'], mockUserProfileResponseState.data)
+  queryClient.setQueryData(['system-features'], mockSystemFeaturesState.data)
+
   const view = render(
     <JotaiProvider>
       <QueryClientProvider client={queryClient}>
@@ -349,20 +352,14 @@ describe('Console bootstrap', () => {
       expect(await screen.findByText('version:1.0.0/1.0.1/cloud')).toBeInTheDocument()
     })
 
-    it('should fall back to placeholder values when profile, workspace, permission, or version data is missing', async () => {
-      mockUserProfileResponseState.data = {
-        meta: {
-          currentVersion: null,
-          currentEnv: null,
-        },
-      }
+    it('should fall back to placeholder values when workspace, permission, or version data is missing', async () => {
       mockCurrentWorkspaceQueryState.data = undefined
       mockPermissionKeysState.permissionKeys = []
       mockLangGeniusVersionState.data = undefined
 
       renderConsoleBootstrap()
 
-      expect(await screen.findByText('user:')).toBeInTheDocument()
+      expect(await screen.findByText('user:user@example.com')).toBeInTheDocument()
       expect(screen.getByText(`workspace:${initialWorkspaceInfo.name}`)).toBeInTheDocument()
       expect(screen.getByText(`role:${initialWorkspaceInfo.role}`)).toBeInTheDocument()
       expect(screen.getByText('keys:')).toBeInTheDocument()

web/context/app-context-defaults.ts

@@ -1,15 +1,5 @@
-import type { GetAccountProfileResponse } from '@dify/contracts/api/console/account/types.gen'
 import type { ICurrentWorkspace, LangGeniusVersionResponse } from '@/models/common'
 
-export const userProfilePlaceholder: GetAccountProfileResponse = {
-  id: '',
-  name: '',
-  email: '',
-  avatar: '',
-  avatar_url: '',
-  is_password_set: false,
-}
-
 export const initialLangGeniusVersionInfo: LangGeniusVersionResponse = {
   current_env: '',
   current_version: '',

web/context/app-context-state.ts

@@ -1,21 +1,16 @@
 'use client'
 
-import type { GetAccountProfileResponse } from '@dify/contracts/api/console/account/types.gen'
-import type { GetSystemFeaturesResponse } from '@dify/contracts/api/console/system-features/types.gen'
-import type { DefinedQueryObserverResult } from '@tanstack/react-query'
-import type { UserProfileWithMeta } from '@/features/account-profile/client'
 import { atom } from 'jotai'
-import { atomWithQuery, atomWithSuspenseQuery, queryClientAtom } from 'jotai-tanstack-query'
+import { atomWithQuery, queryClientAtom } from 'jotai-tanstack-query'
 import { userProfileQueryOptions } from '@/features/account-profile/client'
 import { systemFeaturesQueryOptions } from '@/features/system-features/client'
-import { defaultSystemFeatures } from '@/features/system-features/config'
 import { workspacePermissionKeysQueryOptions } from '@/service/access-control/use-permission-keys'
 import { consoleQuery } from '@/service/client'
 import { langGeniusVersionQueryOptions } from '@/service/lang-genius-version'
+import { atomWithResolvedSuspenseQuery } from '@/utils/query-atoms'
 import {
   initialLangGeniusVersionInfo,
   initialWorkspaceInfo,
-  userProfilePlaceholder,
 } from './app-context-defaults'
 import {
   emptyWorkspacePermissionKeys,
@@ -24,22 +19,16 @@ import {
   normalizeCurrentWorkspace,
 } from './app-context-normalizers'
 
-type SuspenseQueryResult<T> = Omit<DefinedQueryObserverResult<T>, 'isPlaceholderData'>
+const accountProfileQueryAtom = atomWithResolvedSuspenseQuery(() => userProfileQueryOptions())
 
-const accountProfileQueryAtom = atomWithSuspenseQuery(() => userProfileQueryOptions())
+const systemFeaturesQueryAtom = atomWithResolvedSuspenseQuery(() => systemFeaturesQueryOptions())
 
-const systemFeaturesQueryAtom = atomWithSuspenseQuery(() => systemFeaturesQueryOptions())
-
-const systemFeaturesAtom = atom((get): GetSystemFeaturesResponse => {
-  const systemFeaturesQuery = get(systemFeaturesQueryAtom) as SuspenseQueryResult<GetSystemFeaturesResponse>
-
-  return systemFeaturesQuery.data ?? defaultSystemFeatures
+const systemFeaturesAtom = atom((get) => {
+  return get(systemFeaturesQueryAtom).data
 })
 
-export const userProfileAtom = atom((get): GetAccountProfileResponse => {
-  const accountProfileQuery = get(accountProfileQueryAtom) as SuspenseQueryResult<UserProfileWithMeta>
-
-  return accountProfileQuery.data?.profile || userProfilePlaceholder
+export const userProfileAtom = atom((get) => {
+  return get(accountProfileQueryAtom).data.profile
 })
 
 export const userProfileIdAtom = atom((get) => {
@@ -51,12 +40,7 @@ export const userProfileEmailAtom = atom((get) => {
 })
 
 const profileMetaAtom = atom((get) => {
-  const accountProfileQuery = get(accountProfileQueryAtom) as SuspenseQueryResult<UserProfileWithMeta>
-
-  return accountProfileQuery.data?.meta ?? {
-    currentVersion: null,
-    currentEnv: null,
-  }
+  return get(accountProfileQueryAtom).data.meta
 })
 
 const currentWorkspaceQueryAtom = atomWithQuery(() => {
@@ -121,8 +105,8 @@ export const datasetRbacEnabledAtom = atom((get) => {
 
 const versionQueryAtom = atomWithQuery((get) => {
   const meta = get(profileMetaAtom)
-  const systemFeaturesQuery = get(systemFeaturesQueryAtom) as SuspenseQueryResult<GetSystemFeaturesResponse>
-  const enabled = Boolean(meta.currentVersion && !systemFeaturesQuery.data?.branding.enabled)
+  const systemFeaturesQuery = get(systemFeaturesQueryAtom)
+  const enabled = Boolean(meta.currentVersion && !systemFeaturesQuery.data.branding.enabled)
 
   return langGeniusVersionQueryOptions(meta.currentVersion, enabled)
 })

web/utils/query-atoms.spec.ts

@@ -0,0 +1,72 @@
+import { QueryClient } from '@tanstack/react-query'
+import { createStore } from 'jotai'
+import { describe, expect, it, vi } from 'vitest'
+import { atomWithResolvedSuspenseQuery } from './query-atoms'
+
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+      },
+    },
+  })
+
+describe('atomWithResolvedSuspenseQuery', () => {
+  it('should return the resolved suspense query result when query data is cached', () => {
+    const queryClient = createQueryClient()
+    const queryKey = ['resolved-query']
+    queryClient.setQueryData(queryKey, 'cached data')
+
+    const queryAtom = atomWithResolvedSuspenseQuery(
+      () => ({
+        queryKey,
+        queryFn: async () => 'fetched data',
+      }),
+      () => queryClient,
+    )
+    const store = createStore()
+
+    const result = store.get(queryAtom)
+
+    expect(result.data).toBe('cached data')
+    expect(result.isSuccess).toBe(true)
+  })
+
+  it('should throw when the suspense query is still pending', () => {
+    const queryClient = createQueryClient()
+    const queryKey = ['pending-query']
+    const queryFn = vi.fn(() => new Promise<string>(() => {}))
+
+    const queryAtom = atomWithResolvedSuspenseQuery(
+      () => ({
+        queryKey,
+        queryFn,
+      }),
+      () => queryClient,
+    )
+    const store = createStore()
+
+    expect(() => store.get(queryAtom)).toThrow(
+      'Suspense query must be resolved before reading: ["pending-query"]',
+    )
+    expect(queryFn).toHaveBeenCalledTimes(1)
+  })
+
+  it('should keep the wrapped suspense query atom writable', () => {
+    const queryClient = createQueryClient()
+    const queryKey = ['writable-query']
+    queryClient.setQueryData(queryKey, 'cached data')
+
+    const queryAtom = atomWithResolvedSuspenseQuery(
+      () => ({
+        queryKey,
+        queryFn: async () => 'fetched data',
+      }),
+      () => queryClient,
+    )
+    const store = createStore()
+
+    expect(() => store.set(queryAtom)).not.toThrow()
+  })
+})

web/utils/query-atoms.ts

@@ -0,0 +1,55 @@
+import type { DefaultError, QueryClient, QueryKey } from '@tanstack/react-query'
+import type { Getter, WritableAtom } from 'jotai'
+import type {
+  AtomWithSuspenseQueryOptions,
+  AtomWithSuspenseQueryResult,
+} from 'jotai-tanstack-query'
+import { atom } from 'jotai'
+import { atomWithSuspenseQuery } from 'jotai-tanstack-query'
+
+type ResolvedAtomWithSuspenseQueryResult<TData, TError = DefaultError> = Awaited<
+  AtomWithSuspenseQueryResult<TData, TError>
+>
+
+/**
+ * Creates an atomWithSuspenseQuery-compatible atom for data that is already resolved.
+ *
+ * Use this when a suspense query atom is read from another atom after an outer
+ * boundary or bootstrap step has already resolved it. The parameters match
+ * atomWithSuspenseQuery, while the returned atom exposes only the resolved query
+ * result instead of `QueryResult | Promise<QueryResult>`.
+ *
+ * This helper intentionally does not await the promise or create an async
+ * derived atom. If the suspense query is still pending, reading the atom throws
+ * to make the broken "already resolved" invariant explicit.
+ */
+export function atomWithResolvedSuspenseQuery<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+>(
+  getOptions: (
+    get: Getter,
+  ) => AtomWithSuspenseQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
+  getQueryClient?: (get: Getter) => QueryClient,
+): WritableAtom<ResolvedAtomWithSuspenseQueryResult<TData, TError>, [], void> {
+  const queryAtom = atomWithSuspenseQuery(getOptions, getQueryClient)
+
+  return atom(
+    (get) => {
+      const result = get(queryAtom)
+
+      if (result instanceof Promise) {
+        throw new TypeError(
+          `Suspense query must be resolved before reading: ${JSON.stringify(getOptions(get).queryKey)}`,
+        )
+      }
+
+      return result
+    },
+    (_get, set) => {
+      set(queryAtom)
+    },
+  )
+}

web/vitest.setup.ts

@@ -1,5 +1,9 @@
+import type { GetAccountProfileResponse } from '@dify/contracts/api/console/account/types.gen'
+import type { GetSystemFeaturesResponse } from '@dify/contracts/api/console/system-features/types.gen'
 import * as jestDomMatchers from '@testing-library/jest-dom/matchers'
 import { act, cleanup } from '@testing-library/react'
+import { getDefaultStore } from 'jotai'
+import { queryClientAtom } from 'jotai-tanstack-query'
 import * as React from 'react'
 import { afterEach, beforeEach, expect, vi } from 'vitest'
 import 'vitest-canvas-mock'
@@ -220,8 +224,101 @@ const createMockLocalStorage = () => {
 
 let mockLocalStorage: ReturnType<typeof createMockLocalStorage>
 
+const testAccountProfileQueryKey = [
+  ['console', 'account', 'profile', 'get'],
+  { type: 'query' },
+] as const
+
+const testSystemFeaturesQueryKey = [
+  ['console', 'systemFeatures', 'get'],
+  { type: 'query' },
+] as const
+
+const testAccountProfile = {
+  profile: {
+    id: 'user-1',
+    name: 'Test User',
+    email: 'test@dify.ai',
+    avatar: '',
+    avatar_url: null,
+    is_password_set: false,
+    timezone: 'UTC',
+  },
+  meta: {
+    currentVersion: null,
+    currentEnv: null,
+  },
+} satisfies {
+  profile: GetAccountProfileResponse
+  meta: {
+    currentVersion: string | null
+    currentEnv: string | null
+  }
+}
+
+const testSystemFeatures = {
+  enable_app_deploy: false,
+  sso_enforced_for_signin: false,
+  sso_enforced_for_signin_protocol: '',
+  enable_marketplace: false,
+  enable_email_code_login: false,
+  enable_email_password_login: true,
+  enable_social_oauth_login: false,
+  enable_collaboration_mode: true,
+  is_allow_create_workspace: false,
+  is_allow_register: false,
+  is_email_setup: false,
+  enable_change_email: true,
+  max_plugin_package_size: 15728640,
+  license: {
+    status: 'none',
+    expired_at: '',
+    workspaces: {
+      enabled: false,
+      size: 0,
+      limit: 0,
+    },
+  },
+  branding: {
+    enabled: false,
+    login_page_logo: '',
+    workspace_logo: '',
+    favicon: '',
+    application_title: '',
+  },
+  webapp_auth: {
+    enabled: false,
+    allow_sso: false,
+    sso_config: {
+      protocol: '',
+    },
+    allow_email_code_login: false,
+    allow_email_password_login: false,
+  },
+  plugin_installation_permission: {
+    plugin_installation_scope: 'all',
+    restrict_to_marketplace_only: false,
+  },
+  plugin_manager: {
+    enabled: false,
+  },
+  rbac_enabled: false,
+  enable_creators_platform: false,
+  enable_trial_app: false,
+  enable_explore_banner: false,
+  enable_learn_app: true,
+} satisfies GetSystemFeaturesResponse
+
+const seedResolvedAppContextQueries = () => {
+  const queryClient = getDefaultStore().get(queryClientAtom)
+
+  queryClient.setQueryData(testAccountProfileQueryKey, testAccountProfile)
+  queryClient.setQueryData(testSystemFeaturesQueryKey, testSystemFeatures)
+}
+
 beforeEach(() => {
   vi.clearAllMocks()
+  seedResolvedAppContextQueries()
   mockLocalStorage = createMockLocalStorage()
   Object.defineProperty(globalThis, 'localStorage', {
     value: mockLocalStorage,