From cee88de684d5ce55400ef3d219d68c4a9b2dc8d2 Mon Sep 17 00:00:00 2001 From: Stephen Zhou Date: Thu, 9 Jul 2026 10:47:51 +0800 Subject: [PATCH] refactor(web): add prefetched query atom (#38572) --- .../custom-page/__tests__/index.spec.tsx | 12 ++- .../__tests__/use-web-app-brand.spec.tsx | 12 ++- .../__tests__/console-bootstrap.spec.tsx | 15 ++- web/context/app-context-defaults.ts | 10 -- web/context/app-context-state.ts | 38 +++----- web/utils/query-atoms.spec.ts | 72 ++++++++++++++ web/utils/query-atoms.ts | 55 +++++++++++ web/vitest.setup.ts | 97 +++++++++++++++++++ 8 files changed, 261 insertions(+), 50 deletions(-) create mode 100644 web/utils/query-atoms.spec.ts create mode 100644 web/utils/query-atoms.ts diff --git a/web/app/components/custom/custom-page/__tests__/index.spec.tsx b/web/app/components/custom/custom-page/__tests__/index.spec.tsx index 1cb103accaf..1ac45c685ff 100644 --- a/web/app/components/custom/custom-page/__tests__/index.spec.tsx +++ b/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, diff --git a/web/app/components/custom/custom-web-app-brand/hooks/__tests__/use-web-app-brand.spec.tsx b/web/app/components/custom/custom-web-app-brand/hooks/__tests__/use-web-app-brand.spec.tsx index 5ce08a990ab..3bd46ba9ce1 100644 --- a/web/app/components/custom/custom-web-app-brand/hooks/__tests__/use-web-app-brand.spec.tsx +++ b/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 = {} } return { - userProfile: userProfilePlaceholder, + userProfile: testUserProfile, mutateUserProfile: vi.fn(), isCurrentWorkspaceManager: true, isCurrentWorkspaceOwner: false, diff --git a/web/context/__tests__/console-bootstrap.spec.tsx b/web/context/__tests__/console-bootstrap.spec.tsx index 617ff9abdbb..9bf7b787d68 100644 --- a/web/context/__tests__/console-bootstrap.spec.tsx +++ b/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( @@ -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() diff --git a/web/context/app-context-defaults.ts b/web/context/app-context-defaults.ts index 475d7e474f5..b50f4061399 100644 --- a/web/context/app-context-defaults.ts +++ b/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: '', diff --git a/web/context/app-context-state.ts b/web/context/app-context-state.ts index a9050275a7b..e5fc2c706e5 100644 --- a/web/context/app-context-state.ts +++ b/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 = Omit, '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 - - return systemFeaturesQuery.data ?? defaultSystemFeatures +const systemFeaturesAtom = atom((get) => { + return get(systemFeaturesQueryAtom).data }) -export const userProfileAtom = atom((get): GetAccountProfileResponse => { - const accountProfileQuery = get(accountProfileQueryAtom) as SuspenseQueryResult - - 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 - - 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 - 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) }) diff --git a/web/utils/query-atoms.spec.ts b/web/utils/query-atoms.spec.ts new file mode 100644 index 00000000000..a0852bd2b9a --- /dev/null +++ b/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(() => {})) + + 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() + }) +}) diff --git a/web/utils/query-atoms.ts b/web/utils/query-atoms.ts new file mode 100644 index 00000000000..543a51499c2 --- /dev/null +++ b/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 = Awaited< + AtomWithSuspenseQueryResult +> + +/** + * 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`. + * + * 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, + getQueryClient?: (get: Getter) => QueryClient, +): WritableAtom, [], 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) + }, + ) +} diff --git a/web/vitest.setup.ts b/web/vitest.setup.ts index 60e8d7c2495..acd6dd4edb0 100644 --- a/web/vitest.setup.ts +++ b/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 +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,