fix(web): prefetch workspace and guard routes with contract query (#36870)

web/app/(commonLayout)/__tests__/hydration-boundary.spec.tsx

@@ -7,6 +7,9 @@ const mocks = vi.hoisted(() => ({
   queryClient: undefined as QueryClient | undefined,
   profileQueryFn: vi.fn(),
   systemFeaturesQueryFn: vi.fn(),
+  workspaceQueryFn: vi.fn(),
+  workspaceQueryOptions: vi.fn(),
+  getServerConsoleClientContext: vi.fn(),
   redirect: vi.fn((url: string) => {
     throw new Error(`NEXT_REDIRECT:${url}`)
   }),
@@ -35,7 +38,17 @@ vi.mock('@/features/account-profile/server', () => ({
 }))
 
 vi.mock('@/service/server', () => ({
+  getServerConsoleClientContext: () => mocks.getServerConsoleClientContext(),
   resolveServerConsoleApiUrl: (...args: unknown[]) => mocks.resolveServerConsoleApiUrl(...args),
+  serverConsoleQuery: {
+    workspaces: {
+      current: {
+        post: {
+          queryOptions: (...args: unknown[]) => mocks.workspaceQueryOptions(...args),
+        },
+      },
+    },
+  },
 }))
 
 vi.mock('@/service/server-system-features', () => ({
@@ -70,6 +83,16 @@ describe('CommonLayoutHydrationBoundary', () => {
       },
     })
     mocks.systemFeaturesQueryFn.mockResolvedValue({ branding: { enabled: false } })
+    mocks.workspaceQueryFn.mockResolvedValue({ id: 'workspace-id', name: 'Workspace' })
+    mocks.getServerConsoleClientContext.mockResolvedValue({
+      cookie: 'session=abc',
+      csrfToken: 'csrf-token',
+    })
+    mocks.workspaceQueryOptions.mockReturnValue({
+      queryKey: ['console', 'workspaces', 'current', 'post'],
+      queryFn: mocks.workspaceQueryFn,
+      retry: false,
+    })
   })
 
   it('should hydrate common layout queries and render children', async () => {
@@ -87,6 +110,15 @@ describe('CommonLayoutHydrationBoundary', () => {
     expect(screen.getByText('Common shell')).toBeInTheDocument()
     expect(mocks.profileQueryFn).toHaveBeenCalledTimes(1)
     expect(mocks.systemFeaturesQueryFn).toHaveBeenCalledTimes(1)
+    expect(mocks.getServerConsoleClientContext).toHaveBeenCalledTimes(1)
+    expect(mocks.workspaceQueryOptions).toHaveBeenCalledWith({
+      context: {
+        cookie: 'session=abc',
+        csrfToken: 'csrf-token',
+      },
+      retry: false,
+    })
+    expect(mocks.workspaceQueryFn).toHaveBeenCalledTimes(1)
   })
 
   it('should redirect unauthorized users to the refresh route with the current path', async () => {
@@ -123,5 +155,6 @@ describe('CommonLayoutHydrationBoundary', () => {
     expect(screen.getByText('Common shell')).toBeInTheDocument()
     expect(mocks.profileQueryFn).not.toHaveBeenCalled()
     expect(mocks.systemFeaturesQueryFn).not.toHaveBeenCalled()
+    expect(mocks.workspaceQueryFn).not.toHaveBeenCalled()
   })
 })

web/app/(commonLayout)/__tests__/role-route-guard.spec.tsx

@@ -0,0 +1,113 @@
+import { useQuery } from '@tanstack/react-query'
+import { render, screen } from '@testing-library/react'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import RoleRouteGuard from '../role-route-guard'
+
+const mocks = vi.hoisted(() => ({
+  redirect: vi.fn((url: string) => {
+    throw new Error(`NEXT_REDIRECT:${url}`)
+  }),
+  currentWorkspaceQueryOptions: vi.fn(() => ({ queryKey: ['console', 'workspaces', 'current', 'post'] })),
+}))
+
+let mockPathname = '/apps'
+
+vi.mock('@/next/navigation', () => ({
+  redirect: (url: string) => mocks.redirect(url),
+  usePathname: () => mockPathname,
+}))
+
+vi.mock('@tanstack/react-query', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@tanstack/react-query')>()
+  return {
+    ...actual,
+    useQuery: vi.fn(),
+  }
+})
+
+vi.mock('@/service/client', () => ({
+  consoleQuery: {
+    workspaces: {
+      current: {
+        post: {
+          queryOptions: mocks.currentWorkspaceQueryOptions,
+        },
+      },
+    },
+  },
+}))
+
+const mockUseQuery = vi.mocked(useQuery)
+
+const setCurrentWorkspaceQuery = (overrides: { role?: string, isPending?: boolean } = {}) => {
+  mockUseQuery.mockReturnValue({
+    data: overrides.role,
+    isPending: overrides.isPending ?? false,
+  } as ReturnType<typeof useQuery>)
+}
+
+describe('RoleRouteGuard', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockPathname = '/apps'
+    setCurrentWorkspaceQuery()
+  })
+
+  it('should render loading while workspace is loading', () => {
+    setCurrentWorkspaceQuery({ isPending: true })
+
+    render((
+      <RoleRouteGuard>
+        <div>content</div>
+      </RoleRouteGuard>
+    ))
+
+    expect(screen.getByRole('status')).toBeInTheDocument()
+    expect(screen.queryByText('content')).not.toBeInTheDocument()
+    expect(mocks.redirect).not.toHaveBeenCalled()
+    expect(mocks.currentWorkspaceQueryOptions).toHaveBeenCalledWith({
+      select: expect.any(Function),
+    })
+  })
+
+  it('should redirect dataset operator on guarded routes', () => {
+    setCurrentWorkspaceQuery({ role: 'dataset_operator' })
+
+    expect(() => render((
+      <RoleRouteGuard>
+        <div>content</div>
+      </RoleRouteGuard>
+    ))).toThrow('NEXT_REDIRECT:/datasets')
+
+    expect(mocks.redirect).toHaveBeenCalledWith('/datasets')
+  })
+
+  it('should allow dataset operator on non-guarded routes', () => {
+    mockPathname = '/plugins'
+    setCurrentWorkspaceQuery({ role: 'dataset_operator' })
+
+    render((
+      <RoleRouteGuard>
+        <div>content</div>
+      </RoleRouteGuard>
+    ))
+
+    expect(screen.getByText('content')).toBeInTheDocument()
+    expect(mocks.redirect).not.toHaveBeenCalled()
+  })
+
+  it('should not block non-guarded routes while workspace is loading', () => {
+    mockPathname = '/plugins'
+    setCurrentWorkspaceQuery({ isPending: true })
+
+    render((
+      <RoleRouteGuard>
+        <div>content</div>
+      </RoleRouteGuard>
+    ))
+
+    expect(screen.getByText('content')).toBeInTheDocument()
+    expect(screen.queryByRole('status')).not.toBeInTheDocument()
+    expect(mocks.redirect).not.toHaveBeenCalled()
+  })
+})

web/app/(commonLayout)/error.tsx

@@ -13,6 +13,8 @@ type Props = {
 export default function CommonLayoutError({ error, unstable_retry }: Props) {
   const { t } = useTranslation('common')
 
+  console.error(error)
+
   // 401 already triggered jumpTo(/signin) inside service/base.ts. Render Loading
   // until the browser navigation completes, matching main's Splash behavior.
   // Showing the "Try again" button here would just flash for a few frames before

web/app/(commonLayout)/hydration-boundary.tsx

@@ -4,7 +4,7 @@ import { getQueryClientServer } from '@/context/query-client-server'
 import { serverUserProfileQueryOptions } from '@/features/account-profile/server'
 import { headers } from '@/next/headers'
 import { redirect } from '@/next/navigation'
-import { resolveServerConsoleApiUrl } from '@/service/server'
+import { getServerConsoleClientContext, resolveServerConsoleApiUrl, serverConsoleQuery } from '@/service/server'
 import { serverSystemFeaturesQueryOptions } from '@/service/server-system-features'
 import { basePath } from '@/utils/var'
 
@@ -70,9 +70,15 @@ export async function CommonLayoutHydrationBoundary({ children }: { children: Re
   }
 
   try {
+    const context = await getServerConsoleClientContext()
+
     await Promise.all([
       queryClient.fetchQuery(serverUserProfileQueryOptions()),
       queryClient.prefetchQuery(serverSystemFeaturesQueryOptions()),
+      queryClient.prefetchQuery(serverConsoleQuery.workspaces.current.post.queryOptions({
+        context,
+        retry: false,
+      })),
     ])
   }
   catch (error) {

web/app/(commonLayout)/role-route-guard.spec.tsx

@@ -1,109 +0,0 @@
-import { render, screen, waitFor } from '@testing-library/react'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
-import RoleRouteGuard from './role-route-guard'
-
-const mockReplace = vi.fn()
-const mockUseAppContext = vi.fn()
-let mockPathname = '/apps'
-
-vi.mock('@/next/navigation', () => ({
-  usePathname: () => mockPathname,
-  useRouter: () => ({
-    replace: mockReplace,
-  }),
-}))
-
-vi.mock('@/context/app-context', () => ({
-  useAppContext: () => mockUseAppContext(),
-}))
-
-type AppContextMock = {
-  isCurrentWorkspaceDatasetOperator: boolean
-  isLoadingCurrentWorkspace: boolean
-}
-
-const baseContext: AppContextMock = {
-  isCurrentWorkspaceDatasetOperator: false,
-  isLoadingCurrentWorkspace: false,
-}
-
-const setAppContext = (overrides: Partial<AppContextMock> = {}) => {
-  mockUseAppContext.mockReturnValue({
-    ...baseContext,
-    ...overrides,
-  })
-}
-
-describe('RoleRouteGuard', () => {
-  beforeEach(() => {
-    vi.clearAllMocks()
-    mockPathname = '/apps'
-    setAppContext()
-  })
-
-  it('should render loading while workspace is loading', () => {
-    setAppContext({
-      isLoadingCurrentWorkspace: true,
-    })
-
-    render((
-      <RoleRouteGuard>
-        <div>content</div>
-      </RoleRouteGuard>
-    ))
-
-    expect(screen.getByRole('status')).toBeInTheDocument()
-    expect(screen.queryByText('content')).not.toBeInTheDocument()
-    expect(mockReplace).not.toHaveBeenCalled()
-  })
-
-  it('should redirect dataset operator on guarded routes', async () => {
-    setAppContext({
-      isCurrentWorkspaceDatasetOperator: true,
-    })
-
-    render((
-      <RoleRouteGuard>
-        <div>content</div>
-      </RoleRouteGuard>
-    ))
-
-    expect(screen.queryByText('content')).not.toBeInTheDocument()
-    await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith('/datasets')
-    })
-  })
-
-  it('should allow dataset operator on non-guarded routes', () => {
-    mockPathname = '/plugins'
-    setAppContext({
-      isCurrentWorkspaceDatasetOperator: true,
-    })
-
-    render((
-      <RoleRouteGuard>
-        <div>content</div>
-      </RoleRouteGuard>
-    ))
-
-    expect(screen.getByText('content')).toBeInTheDocument()
-    expect(mockReplace).not.toHaveBeenCalled()
-  })
-
-  it('should not block non-guarded routes while workspace is loading', () => {
-    mockPathname = '/plugins'
-    setAppContext({
-      isLoadingCurrentWorkspace: true,
-    })
-
-    render((
-      <RoleRouteGuard>
-        <div>content</div>
-      </RoleRouteGuard>
-    ))
-
-    expect(screen.getByText('content')).toBeInTheDocument()
-    expect(screen.queryByRole('status')).not.toBeInTheDocument()
-    expect(mockReplace).not.toHaveBeenCalled()
-  })
-})

web/app/(commonLayout)/role-route-guard.tsx

@@ -1,33 +1,29 @@
 'use client'
 
 import type { ReactNode } from 'react'
-import { useEffect } from 'react'
+import { useQuery } from '@tanstack/react-query'
 import Loading from '@/app/components/base/loading'
-import { useAppContext } from '@/context/app-context'
-import { usePathname, useRouter } from '@/next/navigation'
+import { redirect, usePathname } from '@/next/navigation'
+import { consoleQuery } from '@/service/client'
 
 const datasetOperatorRedirectRoutes = ['/apps', '/app', '/explore', '/tools'] as const
 
 const isPathUnderRoute = (pathname: string, route: string) => pathname === route || pathname.startsWith(`${route}/`)
 
 export default function RoleRouteGuard({ children }: { children: ReactNode }) {
-  const { isCurrentWorkspaceDatasetOperator, isLoadingCurrentWorkspace } = useAppContext()
+  const currentWorkspaceRoleQuery = useQuery(consoleQuery.workspaces.current.post.queryOptions({
+    select: workspace => workspace.role,
+  }))
   const pathname = usePathname()
-  const router = useRouter()
   const shouldGuardRoute = datasetOperatorRedirectRoutes.some(route => isPathUnderRoute(pathname, route))
-  const shouldRedirect = shouldGuardRoute && !isLoadingCurrentWorkspace && isCurrentWorkspaceDatasetOperator
-
-  useEffect(() => {
-    if (shouldRedirect)
-      router.replace('/datasets')
-  }, [shouldRedirect, router])
+  const shouldRedirect = shouldGuardRoute && !currentWorkspaceRoleQuery.isPending && currentWorkspaceRoleQuery.data === 'dataset_operator'
 
   // Block rendering only for guarded routes to avoid permission flicker.
-  if (shouldGuardRoute && isLoadingCurrentWorkspace)
+  if (shouldGuardRoute && currentWorkspaceRoleQuery.isPending)
     return <Loading type="app" />
 
   if (shouldRedirect)
-    return null
+    redirect('/datasets')
 
   return <>{children}</>
 }

web/app/education-apply/education-apply-page.tsx

@@ -30,6 +30,7 @@ import {
   useEducationAdd,
   useInvalidateEducationStatus,
 } from '@/service/use-education'
+import { removeLocalStorageItem } from '@/utils/local-storage'
 import DifyLogo from '../components/base/logo/dify-logo'
 import AppliedEducationContent from './applied-education-content'
 import RoleSelector from './role-selector'
@@ -83,7 +84,7 @@ const EducationApplyAgeContent = () => {
       if (res.message === 'success') {
         onPlanInfoChanged()
         updateEducationStatus()
-        localStorage.removeItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM)
+        removeLocalStorageItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM)
         setHasSubmittedEducation(true)
       }
       else {

web/app/education-apply/hooks.ts

@@ -14,6 +14,7 @@ import { useModalContextSelector } from '@/context/modal-context'
 import { useProviderContext } from '@/context/provider-context'
 import { useRouter, useSearchParams } from '@/next/navigation'
 import { useEducationAutocomplete, useEducationVerify } from '@/service/use-education'
+import { getLocalStorageItem, setLocalStorageItem } from '@/utils/local-storage'
 import {
   EDUCATION_RE_VERIFY_ACTION,
   EDUCATION_VERIFY_URL_SEARCHPARAMS_ACTION,
@@ -133,7 +134,7 @@ const useEducationReverifyNotice = ({
 export const useEducationInit = () => {
   const setShowAccountSettingModal = useModalContextSelector(s => s.setShowAccountSettingModal)
   const setShowEducationExpireNoticeModal = useModalContextSelector(s => s.setShowEducationExpireNoticeModal)
-  const educationVerifying = localStorage.getItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM)
+  const educationVerifying = getLocalStorageItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM)
   const searchParams = useSearchParams()
   const educationVerifyAction = searchParams.get('action')
 
@@ -156,7 +157,7 @@ export const useEducationInit = () => {
       setShowAccountSettingModal({ payload: ACCOUNT_SETTING_TAB.BILLING })
 
       if (educationVerifyAction === EDUCATION_VERIFY_URL_SEARCHPARAMS_ACTION)
-        localStorage.setItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM, 'yes')
+        setLocalStorageItem(EDUCATION_VERIFYING_LOCALSTORAGE_ITEM, 'yes')
     }
     if (educationVerifyAction === EDUCATION_RE_VERIFY_ACTION)
       handleVerify()

web/app/error.tsx

@@ -15,6 +15,8 @@ export default function AppError({ error, reset, unstable_retry }: Props) {
   const { t } = useTranslation('common')
   const retry = reset ?? unstable_retry
 
+  console.error(error)
+
   if (isLegacyBase401(error))
     return <FullScreenLoading />