fix(web): hydrate workspace permissions before navigation render (#39524)

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

@@ -9,6 +9,8 @@ const mocks = vi.hoisted(() => ({
   profileQueryFn: vi.fn(),
   workspaceQueryFn: vi.fn(),
   workspaceQueryOptions: vi.fn(),
+  permissionQueryFn: vi.fn(),
+  permissionQueryOptions: vi.fn(),
   getServerConsoleClientContext: vi.fn(),
   redirect: vi.fn((url: string) => {
     throw new Error(`NEXT_REDIRECT:${url}`)
@@ -58,6 +60,13 @@ vi.mock('@/service/server', () => ({
         post: {
           queryOptions: (...args: unknown[]) => mocks.workspaceQueryOptions(...args),
         },
+        rbac: {
+          myPermissions: {
+            get: {
+              queryOptions: (...args: unknown[]) => mocks.permissionQueryOptions(...args),
+            },
+          },
+        },
       },
     },
   },
@@ -92,6 +101,11 @@ describe('CommonLayoutHydrationBoundary', () => {
       },
     })
     mocks.workspaceQueryFn.mockResolvedValue({ id: 'workspace-id', name: 'Workspace' })
+    mocks.permissionQueryFn.mockResolvedValue({
+      workspace: { permission_keys: ['agent.manage'] },
+      app: { default_permission_keys: [], overrides: [] },
+      dataset: { default_permission_keys: [], overrides: [] },
+    })
     mocks.getServerConsoleClientContext.mockResolvedValue({
       cookie: 'session=abc',
       csrfToken: 'csrf-token',
@@ -101,6 +115,14 @@ describe('CommonLayoutHydrationBoundary', () => {
       queryFn: mocks.workspaceQueryFn,
       retry: false,
     })
+    mocks.permissionQueryOptions.mockReturnValue({
+      queryKey: [
+        ['console', 'workspaces', 'current', 'rbac', 'myPermissions', 'get'],
+        { type: 'query' },
+      ],
+      queryFn: mocks.permissionQueryFn,
+      retry: false,
+    })
   })
 
   it('should prefetch common layout queries', async () => {
@@ -126,6 +148,14 @@ describe('CommonLayoutHydrationBoundary', () => {
       retry: false,
     })
     expect(mocks.workspaceQueryFn).toHaveBeenCalledTimes(1)
+    expect(mocks.permissionQueryOptions).toHaveBeenCalledWith({
+      context: {
+        cookie: 'session=abc',
+        csrfToken: 'csrf-token',
+      },
+      retry: false,
+    })
+    expect(mocks.permissionQueryFn).toHaveBeenCalledTimes(1)
   })
 
   it('should dehydrate only Common-owned queries', async () => {
@@ -138,11 +168,12 @@ describe('CommonLayoutHydrationBoundary', () => {
     const state = (element as ReactElement<{ state: DehydratedState }>).props.state
     const queryKeys = state.queries.map((query) => query.queryKey)
 
-    expect(queryKeys).toHaveLength(2)
+    expect(queryKeys).toHaveLength(3)
     expect(queryKeys).toEqual(
       expect.arrayContaining([
         ['common', 'user-profile'],
         ['console', 'workspaces', 'current', 'post'],
+        [['console', 'workspaces', 'current', 'rbac', 'myPermissions', 'get'], { type: 'query' }],
       ]),
     )
   })
@@ -203,5 +234,6 @@ describe('CommonLayoutHydrationBoundary', () => {
     expect(screen.getByText('Common shell')).toBeInTheDocument()
     expect(mocks.profileQueryFn).not.toHaveBeenCalled()
     expect(mocks.workspaceQueryFn).not.toHaveBeenCalled()
+    expect(mocks.permissionQueryFn).not.toHaveBeenCalled()
   })
 })

web/app/(commonLayout)/agents/__tests__/agents-access-guard.spec.tsx

@@ -22,6 +22,12 @@ vi.mock('@/context/permission-state', async () => {
 
   return createPermissionStateModuleMock(() => mockConsoleStateReader())
 })
+vi.mock('@/features/agent-v2/permissions', () => {
+  return {
+    useCanManageAgents: () =>
+      mockConsoleStateReader().workspacePermissionKeys.includes('agent.manage'),
+  }
+})
 
 type ConsoleStateFixture = {
   isLoadingCurrentWorkspace: boolean
@@ -68,8 +74,8 @@ describe('AgentsAccessGuard', () => {
     expect(mockReplace).not.toHaveBeenCalled()
   })
 
-  it('renders loading while workspace permission keys are loading', () => {
+  it('renders loading while workspace permissions are loading', () => {
-    setConsoleState({ isLoadingWorkspacePermissionKeys: true, workspacePermissionKeys: [] })
+    setConsoleState({ isLoadingWorkspacePermissionKeys: true })
 
     render(
       <AgentsAccessGuard>

web/app/(commonLayout)/agents/agents-access-guard.tsx

@@ -15,7 +15,7 @@ export function AgentsAccessGuard({ children }: { children: ReactNode }) {
   const isLoadingWorkspacePermissionKeys = useAtomValue(workspacePermissionKeysLoadingAtom)
   const canManageAgents = useCanManageAgents()
   const router = useRouter()
-  const isLoadingAccess = isLoadingCurrentWorkspace || !!isLoadingWorkspacePermissionKeys
+  const isLoadingAccess = isLoadingCurrentWorkspace || isLoadingWorkspacePermissionKeys
   const shouldRedirect = !isLoadingAccess && !!currentWorkspaceId && !canManageAgents
 
   useEffect(() => {

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

@@ -70,6 +70,12 @@ export async function CommonLayoutHydrationBoundary({ children }: { children: Re
             retry: false,
           }),
         ),
+        queryClient.prefetchQuery(
+          serverConsoleQuery.workspaces.current.rbac.myPermissions.get.queryOptions({
+            context,
+            retry: false,
+          }),
+        ),
       ])
     } catch (error) {
       await handleProfileError(error)

web/app/components/main-nav/__tests__/index.spec.tsx

@@ -228,6 +228,10 @@ vi.mock('react-i18next', async () => {
 vi.mock('@/service/client', async (importOriginal) => {
   const actual = await importOriginal<typeof import('@/service/client')>()
   const currentWorkspaceQueryKey = ['console', 'workspaces', 'current', 'post'] as const
+  const currentPermissionsQueryKey = [
+    ['console', 'workspaces', 'current', 'rbac', 'myPermissions', 'get'],
+    { type: 'query' },
+  ] as const
   const workspacesQueryKey = ['console', 'workspaces', 'get'] as const
   const consoleQuery = new Proxy(actual.consoleQuery, {
     get(target, prop, receiver) {
@@ -243,6 +247,16 @@ vi.mock('@/service/client', async (importOriginal) => {
                 ...options,
               }),
             },
+            rbac: {
+              myPermissions: {
+                get: {
+                  queryOptions: () => ({
+                    queryKey: currentPermissionsQueryKey,
+                    queryFn: () => new Promise(() => {}),
+                  }),
+                },
+              },
+            },
           },
           get: {
             queryKey: () => workspacesQueryKey,
@@ -467,7 +481,11 @@ const renderMainNav = (
       <MainNav />
       {options.extra}
     </JotaiProvider>,
-    { systemFeatures: resolvedSystemFeatures, queryClient },
+    {
+      systemFeatures: resolvedSystemFeatures,
+      workspacePermissionKeys: currentConsoleState.workspacePermissionKeys,
+      queryClient,
+    },
   )
 }
 

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

@@ -1,3 +1,4 @@
+import { toast } from '@langgenius/dify-ui/toast'
 import { cleanup, screen, waitFor } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
 import { Plan } from '@/app/components/billing/type'
@@ -9,6 +10,21 @@ let mockProviderContext: Record<string, unknown> = {}
 let mockConsoleState: Record<string, unknown> = {}
 const mockFetchSubscriptionUrls = vi.hoisted(() => vi.fn())
 const mockEducationAdd = vi.hoisted(() => vi.fn())
+const mockSwitchWorkspace = vi.hoisted(() => vi.fn())
+const mockWorkspaces = vi.hoisted(() => [
+  {
+    id: 'workspace-1',
+    name: 'Workspace One',
+    current: true,
+    plan: 'sandbox',
+  },
+  {
+    id: 'workspace-2',
+    name: 'Workspace Two',
+    current: false,
+    plan: 'sandbox',
+  },
+])
 
 vi.mock('@/context/provider-context', () => ({
   useProviderContext: () => mockProviderContext,
@@ -76,12 +92,12 @@ vi.mock('@/service/client', () => ({
       get: {
         queryOptions: () => ({
           queryKey: ['workspaces'],
-          queryFn: async () => ({ workspaces: [] }),
+          queryFn: async () => ({ workspaces: mockWorkspaces }),
         }),
       },
       switch: {
         post: {
-          mutationOptions: () => ({ mutationFn: vi.fn() }),
+          mutationOptions: () => ({ mutationFn: mockSwitchWorkspace }),
         },
       },
     },
@@ -95,7 +111,7 @@ const setupContext = (isCurrentWorkspaceManager: boolean) => {
     onPlanInfoChanged: vi.fn(),
   }
   mockConsoleState = {
-    currentWorkspace: { id: 'workspace-1', name: 'Workspace' },
+    currentWorkspace: { id: 'workspace-1', name: 'Workspace One' },
     isCurrentWorkspaceManager,
     userProfile: {
       name: 'Student',
@@ -106,7 +122,7 @@ const setupContext = (isCurrentWorkspaceManager: boolean) => {
 }
 
 const renderPage = () => {
-  const { wrapper } = createConsoleQueryWrapper()
+  const { wrapper } = createConsoleQueryWrapper({ workspacePermissionKeys: null })
   return render(<EducationApplyPage />, {
     wrapper,
   })
@@ -116,7 +132,18 @@ describe('EducationApplyPage billing boundary', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     cleanup()
+    vi.spyOn(toast, 'error').mockImplementation(() => 'toast-id')
     mockFetchSubscriptionUrls.mockResolvedValue({ url: window.location.href })
+    mockSwitchWorkspace.mockResolvedValue(undefined)
+    vi.stubGlobal('location', {
+      href: 'https://console.example.com/education-apply?token=education-token',
+      reload: vi.fn(),
+    } as unknown as Location)
+  })
+
+  afterEach(() => {
+    vi.unstubAllGlobals()
+    vi.restoreAllMocks()
   })
 
   it('lets workspace managers apply the education coupon at checkout', async () => {
@@ -142,4 +169,49 @@ describe('EducationApplyPage billing boundary', () => {
       screen.queryByRole('button', { name: 'education.useEducationDiscount' }),
     ).not.toBeInTheDocument()
   })
+
+  it('reloads the current URL after switching workspaces', async () => {
+    setupContext(true)
+    const user = userEvent.setup()
+    renderPage()
+
+    await user.click(await screen.findByRole('combobox'))
+    await user.click(await screen.findByRole('option', { name: /Workspace Two/ }))
+
+    await waitFor(() => {
+      expect(mockSwitchWorkspace.mock.calls[0]?.[0]).toEqual({
+        body: { tenant_id: 'workspace-2' },
+      })
+      expect(globalThis.location.reload).toHaveBeenCalledTimes(1)
+    })
+  })
+
+  it('disables workspace selection while switching and recovers after a failure', async () => {
+    setupContext(true)
+    let rejectSwitch!: (error: Error) => void
+    mockSwitchWorkspace.mockImplementation(
+      () =>
+        new Promise((_, reject) => {
+          rejectSwitch = reject
+        }),
+    )
+    const user = userEvent.setup()
+    renderPage()
+
+    const selector = await screen.findByRole('combobox')
+    await user.click(selector)
+    await user.click(await screen.findByRole('option', { name: /Workspace Two/ }))
+
+    await waitFor(() => {
+      expect(selector).toBeDisabled()
+    })
+
+    rejectSwitch(new Error('switch failed'))
+
+    await waitFor(() => {
+      expect(selector).toBeEnabled()
+      expect(vi.mocked(toast.error)).toHaveBeenCalledWith('common.actionMsg.modifiedUnsuccessfully')
+    })
+    expect(globalThis.location.reload).not.toHaveBeenCalled()
+  })
 })

web/app/education-apply/applied-education-content.tsx

@@ -15,6 +15,7 @@ type AppliedEducationContentProps = {
   currentWorkspace: ICurrentWorkspace
   plan: PlanType
   action: ReactNode
+  isSwitchingWorkspace: boolean
   onSwitchWorkspace: (tenantId: string) => void
 }
 
@@ -29,6 +30,7 @@ const AppliedEducationContent = ({
   currentWorkspace,
   plan,
   action,
+  isSwitchingWorkspace,
   onSwitchWorkspace,
 }: AppliedEducationContentProps) => {
   const { t } = useTranslation()
@@ -73,7 +75,10 @@ const AppliedEducationContent = ({
               if (value) onSwitchWorkspace(value)
             }}
           >
-            <SelectTrigger className="h-12! w-fit max-w-full min-w-[280px] cursor-pointer justify-between rounded-lg border-[0.5px] border-transparent bg-components-input-bg-normal px-3! py-1.5! hover:bg-state-base-hover">
+            <SelectTrigger
+              className="h-12! w-fit max-w-full min-w-[280px] cursor-pointer justify-between rounded-lg border-[0.5px] border-transparent bg-components-input-bg-normal px-3! py-1.5! hover:bg-state-base-hover"
+              disabled={isSwitchingWorkspace}
+            >
               <span className="flex min-w-0 items-center gap-3">
                 <span className="flex size-8 shrink-0 items-center justify-center rounded-lg bg-components-icon-bg-blue-solid text-[14px]">
                   <span className="bg-linear-to-r from-components-avatar-shape-fill-stop-0 to-components-avatar-shape-fill-stop-100 bg-clip-text font-semibold text-shadow-shadow-1 uppercase opacity-90">

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

@@ -6,7 +6,7 @@ import type { ICurrentWorkspace } from '@/models/common'
 import { Button } from '@langgenius/dify-ui/button'
 import { Checkbox } from '@langgenius/dify-ui/checkbox'
 import { toast } from '@langgenius/dify-ui/toast'
-import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import { useMutation, useQuery } from '@tanstack/react-query'
 import { noop } from 'es-toolkit/function'
 import { useAtomValue } from 'jotai'
 import { useState } from 'react'
@@ -50,7 +50,6 @@ const EducationApplyAgeContent = () => {
   const { handleEducationDiscount } = useEducationDiscount()
   const router = useRouter()
   const openAsyncWindow = useAsyncWindowOpen()
-  const queryClient = useQueryClient()
   const switchWorkspaceMutation = useMutation(consoleQuery.workspaces.switch.post.mutationOptions())
   const setEducationVerifying = useSetEducationVerifying()
 
@@ -115,12 +114,7 @@ const EducationApplyAgeContent = () => {
 
     try {
       await switchWorkspaceMutation.mutateAsync({ body: { tenant_id: tenantId } })
-      await Promise.all([
+      globalThis.location.reload()
-        queryClient.invalidateQueries({ queryKey: consoleQuery.workspaces.current.post.key() }),
-        queryClient.invalidateQueries({ queryKey: consoleQuery.workspaces.get.queryKey() }),
-      ])
-      onPlanInfoChanged()
-      updateEducationStatus()
     } catch {
       toast.error(t(($) => $['actionMsg.modifiedUnsuccessfully'], { ns: 'common' }))
     }
@@ -211,6 +205,7 @@ const EducationApplyAgeContent = () => {
                 currentWorkspace={currentWorkspace}
                 plan={plan.type}
                 action={renderAppliedEducationAction()}
+                isSwitchingWorkspace={switchWorkspaceMutation.isPending}
                 onSwitchWorkspace={(value) => {
                   void handleSwitchWorkspace(value)
                 }}
@@ -303,6 +298,7 @@ type AppliedEducationWorkspaceBlockProps = {
   currentWorkspace: ICurrentWorkspace
   plan: PlanType
   action: ReactNode
+  isSwitchingWorkspace: boolean
   onSwitchWorkspace: (tenantId: string) => void
 }
 
@@ -310,6 +306,7 @@ function AppliedEducationWorkspaceContent({
   currentWorkspace,
   plan,
   action,
+  isSwitchingWorkspace,
   onSwitchWorkspace,
 }: AppliedEducationWorkspaceBlockProps) {
   const { data: workspacesData } = useQuery(consoleQuery.workspaces.get.queryOptions())
@@ -321,6 +318,7 @@ function AppliedEducationWorkspaceContent({
       currentWorkspace={currentWorkspace}
       plan={plan}
       action={action}
+      isSwitchingWorkspace={isSwitchingWorkspace}
       onSwitchWorkspace={onSwitchWorkspace}
     />
   )

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

@@ -153,6 +153,32 @@ vi.mock('@/service/client', () => ({
             ...options,
           }),
         },
+        rbac: {
+          myPermissions: {
+            get: {
+              queryOptions: () => ({
+                queryKey: ['current-permissions'],
+                queryFn: async () => {
+                  if (mockPermissionKeysState.isPending) return new Promise(() => {})
+
+                  return {
+                    workspace: {
+                      permission_keys: mockPermissionKeysState.permissionKeys,
+                    },
+                    app: {
+                      default_permission_keys: [],
+                      overrides: [],
+                    },
+                    dataset: {
+                      default_permission_keys: [],
+                      overrides: [],
+                    },
+                  }
+                },
+              }),
+            },
+          },
+        },
       },
     },
     version: {
@@ -354,24 +380,6 @@ describe('Console bootstrap', () => {
       can_auto_update: false,
     }
     mockGetRequest.mockImplementation((url: string) => {
-      if (url === '/workspaces/current/rbac/my-permissions') {
-        if (mockPermissionKeysState.isPending) return new Promise(() => {})
-
-        return Promise.resolve({
-          workspace: {
-            permission_keys: mockPermissionKeysState.permissionKeys,
-          },
-          app: {
-            default_permission_keys: [],
-            overrides: [],
-          },
-          dataset: {
-            default_permission_keys: [],
-            overrides: [],
-          },
-        })
-      }
-
       if (url === '/version') return Promise.resolve(mockLangGeniusVersionState.data)
 
       return Promise.reject(new Error(`Unexpected GET ${url}`))

web/context/permission-state.ts

@@ -2,15 +2,12 @@
 
 import { atom } from 'jotai'
 import { atomWithQuery } from 'jotai-tanstack-query'
-import { workspacePermissionKeysQueryOptions } from '@/service/access-control/use-permission-keys'
+import { consoleQuery } from '@/service/client'
 import { emptyWorkspacePermissionKeys } from './app-context-normalizers'
-import { currentWorkspaceIdAtom } from './workspace-state'
 
-const workspacePermissionKeysQueryAtom = atomWithQuery((get) => {
+const workspacePermissionKeysQueryAtom = atomWithQuery(() =>
-  const workspaceId = get(currentWorkspaceIdAtom)
+  consoleQuery.workspaces.current.rbac.myPermissions.get.queryOptions(),
-
+)
-  return workspacePermissionKeysQueryOptions(workspaceId)
-})
 
 export const workspacePermissionKeysAtom = atom((get) => {
   return (

web/features/agent-v2/__tests__/permissions.spec.tsx

@@ -0,0 +1,20 @@
+import { render, screen } from '@testing-library/react'
+import { createConsoleQueryWrapper } from '@/test/console/query-data'
+import { useCanManageAgents } from '../permissions'
+
+function PermissionProbe() {
+  return <span>{String(useCanManageAgents())}</span>
+}
+
+describe('useCanManageAgents', () => {
+  it.each([
+    [['agent.manage'], 'true'],
+    [['dataset.create_and_management'], 'false'],
+  ])('resolves agent.manage from the current permission snapshot', (permissionKeys, expected) => {
+    const { wrapper } = createConsoleQueryWrapper({ workspacePermissionKeys: permissionKeys })
+
+    render(<PermissionProbe />, { wrapper })
+
+    expect(screen.getByText(expected)).toBeInTheDocument()
+  })
+})

web/models/access-control.ts

@@ -206,24 +206,6 @@ export type UpdateRolesOfMemberRequest = {
   roleIds: string[]
 }
 
-type WorkspacePermissionKeys = {
-  permission_keys: string[]
-}
-
-type ResourcePermissionKeys = {
-  default_permission_keys: string[]
-  overrides: Array<{
-    resource_id: string
-    permission_keys: string[]
-  }>
-}
-
-export type PermissionKeysResponse = {
-  workspace: WorkspacePermissionKeys
-  app: ResourcePermissionKeys
-  dataset: ResourcePermissionKeys
-}
-
 export type GetMembersOfRoleRequest = {
   roleId: string
 } & PaginationParameters

web/service/access-control/use-permission-keys.ts

@@ -1,18 +0,0 @@
-import type { PermissionKeysResponse } from '@/models/access-control'
-import { queryOptions } from '@tanstack/react-query'
-// oxlint-disable-next-line no-restricted-imports
-import { get } from '../base'
-
-const NAME_SPACE = 'workspace-permission-keys'
-
-const workspacePermissionKeysQueryKey = (workspaceId?: string) => {
-  return workspaceId ? ([NAME_SPACE, workspaceId] as const) : ([NAME_SPACE] as const)
-}
-
-export const workspacePermissionKeysQueryOptions = (workspaceId?: string) => {
-  return queryOptions<PermissionKeysResponse>({
-    queryKey: workspacePermissionKeysQueryKey(workspaceId),
-    queryFn: () => get<PermissionKeysResponse>('/workspaces/current/rbac/my-permissions'),
-    enabled: workspaceId === undefined || Boolean(workspaceId),
-  })
-}

web/test/console/query-data.ts

@@ -173,7 +173,7 @@ export const createConsoleQueryWrapper = (
   }
   if (options.workspacePermissionKeys !== null) {
     if (options.workspacePermissionKeys)
-      seedWorkspacePermissionsQuery(queryClient, 'workspace-1', options.workspacePermissionKeys)
+      seedWorkspacePermissionsQuery(queryClient, options.workspacePermissionKeys)
     else ensureWorkspacePermissionsQuery(queryClient)
   }
   const systemFeatures =

web/test/console/workspace-permissions.ts

@@ -1,10 +1,14 @@
+import type { MyPermissionsResponse } from '@dify/contracts/api/console/workspaces/types.gen'
 import type { QueryClient } from '@tanstack/react-query'
-import type { PermissionKeysResponse } from '@/models/access-control'
+
-import { workspacePermissionKeysQueryOptions } from '@/service/access-control/use-permission-keys'
+const currentWorkspacePermissionsQueryKey = [
+  ['console', 'workspaces', 'current', 'rbac', 'myPermissions', 'get'],
+  { type: 'query' },
+] as const
 
 const createWorkspacePermissionsFixture = (
   permissionKeys: readonly string[] = [],
-): PermissionKeysResponse => ({
+): MyPermissionsResponse => ({
   workspace: {
     permission_keys: [...permissionKeys],
   },
@@ -20,23 +24,22 @@ const createWorkspacePermissionsFixture = (
 
 export const seedWorkspacePermissionsQuery = (
   queryClient: QueryClient,
-  workspaceId = 'workspace-1',
   permissionKeys: readonly string[] = [],
 ) => {
   const data = createWorkspacePermissionsFixture(permissionKeys)
-  queryClient.setQueryData(workspacePermissionKeysQueryOptions(workspaceId).queryKey, data)
+  queryClient.setQueryData(currentWorkspacePermissionsQueryKey, data)
   return data
 }
 
 export const ensureWorkspacePermissionsQuery = (
   queryClient: QueryClient,
-  workspaceId = 'workspace-1',
   permissionKeys: readonly string[] = [],
 ) => {
-  const queryKey = workspacePermissionKeysQueryOptions(workspaceId).queryKey
+  const existingPermissions = queryClient.getQueryData<MyPermissionsResponse>(
-  const existingPermissions = queryClient.getQueryData<PermissionKeysResponse>(queryKey)
+    currentWorkspacePermissionsQueryKey,
+  )
   if (existingPermissions === undefined)
-    return seedWorkspacePermissionsQuery(queryClient, workspaceId, permissionKeys)
+    return seedWorkspacePermissionsQuery(queryClient, permissionKeys)
 
   return existingPermissions
 }