diff --git a/react/src/App.tsx b/react/src/App.tsx index 1a74d919d6..80990a0fd3 100644 --- a/react/src/App.tsx +++ b/react/src/App.tsx @@ -9,6 +9,7 @@ import MainLayout from './components/MainLayout/MainLayout'; import WebUINavigate from './components/WebUINavigate'; import { useSuspendedBackendaiClient } from './hooks'; import { useBAISettingUserState } from './hooks/useBAISetting'; +import { PageQueryParamAtomProvider } from './hooks/useTransitionSafeQueryParams'; // High priority to import the component import ComputeSessionListPage from './pages/ComputeSessionListPage'; import ModelStoreListPage from './pages/ModelStoreListPage'; @@ -98,7 +99,9 @@ const router = createBrowserRouter([ errorElement: , element: ( - + + + ), }, @@ -115,12 +118,14 @@ const router = createBrowserRouter([ } } > - - - - - - + + + + + + + + ), children: [ diff --git a/react/src/components/AgentSummaryList.tsx b/react/src/components/AgentSummaryList.tsx index f0abe5a1fa..f8d70050cf 100644 --- a/react/src/components/AgentSummaryList.tsx +++ b/react/src/components/AgentSummaryList.tsx @@ -42,7 +42,7 @@ import React, { import { useTranslation } from 'react-i18next'; import { graphql, FetchPolicy, useLazyLoadQuery } from 'react-relay'; import { useBAISettingUserState } from 'src/hooks/useBAISetting'; -import { useDeferredQueryParams } from 'src/hooks/useDeferredQueryParams'; +import { useTransitionSafeQueryParams } from 'src/hooks/useTransitionSafeQueryParams'; import { StringParam, withDefault } from 'use-query-params'; type AgentSummary = NonNullable< @@ -74,7 +74,7 @@ const AgentSummaryList: React.FC = ({ pageSize: 20, }); - const [queryParams, setQuery] = useDeferredQueryParams({ + const [queryParams, setQuery] = useTransitionSafeQueryParams({ order: withDefault(StringParam, undefined), filter: withDefault(StringParam, undefined), status: withDefault(StringParam, 'ALIVE'), diff --git a/react/src/hooks/reactPaginationQueryOptions.tsx b/react/src/hooks/reactPaginationQueryOptions.tsx index c2caeeb6c4..13e3fd2d26 100644 --- a/react/src/hooks/reactPaginationQueryOptions.tsx +++ b/react/src/hooks/reactPaginationQueryOptions.tsx @@ -1,6 +1,6 @@ // import { offset_to_cursor } from "../helper"; import { LazyLoadQueryOptions } from '../helper/types'; -import { useDeferredQueryParams } from './useDeferredQueryParams'; +import { useTransitionSafeQueryParams } from './useTransitionSafeQueryParams'; import { SorterResult } from 'antd/lib/table/interface'; import _ from 'lodash'; import { useMemo, useState } from 'react'; @@ -342,7 +342,7 @@ export const useBAIPaginationOptionState = ( export const useBAIPaginationOptionStateOnSearchParam = ( initialOptions: InitialPaginationOption, ): BAIPaginationOptionState => { - const [options, setOptions] = useDeferredQueryParams({ + const [options, setOptions] = useTransitionSafeQueryParams({ current: NumberParam, pageSize: NumberParam, }); diff --git a/react/src/hooks/useDeferredQueryParams.tsx b/react/src/hooks/useDeferredQueryParams.tsx deleted file mode 100644 index 75adf07cd9..0000000000 --- a/react/src/hooks/useDeferredQueryParams.tsx +++ /dev/null @@ -1,162 +0,0 @@ -import { atom, useAtomValue, useSetAtom } from 'jotai'; -import { atomWithDefault } from 'jotai/utils'; -import _ from 'lodash'; -import { useRef, useCallback, useMemo, useEffect } from 'react'; -import { - useQueryParams, - QueryParamConfigMap, - DecodedValueMap, - UrlUpdateType, -} from 'use-query-params'; - -const queryParamsAtom = atom>({}); - -// Reference counting atom for thread-safe tracking of parameter key usage -const paramRefCountAtom = atom>(new Map()); - -/** - * A custom hook that synchronizes URL search parameters with application state while handling React transitions. - * This hook solves the issue where URL parameter changes within React transitions are not properly reflected - * in the rendering cycle, as search parameter changes are detected through events rather than React's state system. - * - * @template QPCMap - Type extending QueryParamConfigMap that defines the structure of URL parameters - * @param {QPCMap} paramConfigMap - Configuration object that defines the URL parameters to be managed - * - * @returns {[ - * DecodedValueMap, - * (nextQuery: Partial> | ((prevQuery: DecodedValueMap) => Partial>), - * updateType: UrlUpdateType) => void, - * boolean - * ]} A tuple containing: - * - localQuery: The current state of the URL parameters - * - setDeferredQuery: Function to update URL parameters with transition support - * - isPending: Boolean indicating if a transition is in progress - * - * @example - * const [query, setQuery, isPending] = useDeferredQueryParams({ - * page: NumberParam, - * search: StringParam - * }); - * - * // Update URL parameters - * setQuery({ page: 2 }, 'pushIn'); - */ -export function useDeferredQueryParams( - paramConfigMap: QPCMap, -) { - const [query, setQuery] = useQueryParams(paramConfigMap); - const setSharedQuery = useSetAtom(queryParamsAtom); - const setParamRefCount = useSetAtom(paramRefCountAtom); - - const stringifiedParamConfigMap = useMemo(() => { - return JSON.stringify(paramConfigMap); - }, [paramConfigMap]); - - // Reference counting based cleanup: only remove params when last component unmounts - useEffect(() => { - const keys = Object.keys(paramConfigMap); - - // Mount: increase reference count for each key - setParamRefCount((prev) => { - const newMap = new Map(prev); - keys.forEach((key) => { - newMap.set(key, (newMap.get(key) || 0) + 1); - }); - return newMap; - }); - - return () => { - // Unmount: decrease reference count and cleanup if necessary - setParamRefCount((prev) => { - const newMap = new Map(prev); - keys.forEach((key) => { - const currentCount = (newMap.get(key) || 0) - 1; - - if (currentCount <= 0) { - // Last component using this key is unmounting, safe to cleanup - newMap.delete(key); - setSharedQuery((queryPrev) => { - const newState = { ...queryPrev }; - delete newState[key]; - return newState; - }); - } else { - // Other components still using this key - newMap.set(key, currentCount); - } - }); - return newMap; - }); - }; - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [stringifiedParamConfigMap, setParamRefCount, setSharedQuery]); - - const isBeforeInitializingRef = useRef(true); - const selectiveQueryAtom = useMemo( - () => { - const defaultValues = _.mapValues( - paramConfigMap, - (config) => config.default, - ); - return atomWithDefault((get) => { - const globalParams = get(queryParamsAtom); - const selectedParams = _.pick( - // Use query parameters from URL on initial render - isBeforeInitializingRef.current - ? query - : { - ...defaultValues, - ...globalParams, - }, - Object.keys(paramConfigMap), - ); - isBeforeInitializingRef.current = false; - return selectedParams as DecodedValueMap; - }); - }, - // eslint-disable-next-line react-hooks/exhaustive-deps - [stringifiedParamConfigMap], - ); - - let localQuery = useAtomValue(selectiveQueryAtom); - - const setDeferredQuery = useCallback( - ( - nextQuery: - | Partial> - | (( - prevQuery: DecodedValueMap, - ) => Partial>), - updateType: UrlUpdateType, - ) => { - const newQuery = - typeof nextQuery === 'function' ? nextQuery(localQuery) : nextQuery; - - // Update Jotai state - if (updateType === 'replaceIn' || updateType === 'pushIn') { - setSharedQuery((prev) => ({ - ...prev, - ...localQuery, - ...newQuery, - })); - } else { - setSharedQuery((prev) => ({ - ...prev, - ...(newQuery as DecodedValueMap), - })); - } - - // Sync all(merged) query parameters with URL - setQuery( - { - ...localQuery, - ...newQuery, - }, - updateType, - ); - }, - [localQuery, setQuery, setSharedQuery], - ); - - return [localQuery, setDeferredQuery] as const; -} diff --git a/react/src/hooks/useTransitionSafeQueryParams.tsx b/react/src/hooks/useTransitionSafeQueryParams.tsx new file mode 100644 index 0000000000..1d6be2a890 --- /dev/null +++ b/react/src/hooks/useTransitionSafeQueryParams.tsx @@ -0,0 +1,164 @@ +import { atom, PrimitiveAtom, useAtomValue, useSetAtom } from 'jotai'; +import { atomWithDefault } from 'jotai/utils'; +import _ from 'lodash'; +import { createContext, use, useCallback, useMemo } from 'react'; +import { useLocation } from 'react-router-dom'; +import { + useQueryParams, + QueryParamConfigMap, + DecodedValueMap, + UrlUpdateType, +} from 'use-query-params'; + +/** + * A custom hook that synchronizes URL search parameters with application state while handling React transitions. + * This hook solves the issue where URL parameter changes within React transitions are not properly reflected + * in the rendering cycle, as search parameter changes are detected through events rather than React's state system. + * + * @template QPCMap - Type extending QueryParamConfigMap that defines the structure of URL parameters + * @param {QPCMap} paramConfigMap - Configuration object that defines the URL parameters to be managed + * + * @returns {[ + * DecodedValueMap, + * (nextQueryParams: Partial> | ((prevQueryParams: DecodedValueMap) => Partial>), + * updateType: UrlUpdateType) => void + * ]} A tuple containing: + * - currentQueryParams: The current state of the URL parameters + * - updateQueryParams: Function to update URL parameters with transition support + * + * @example + * const [queryParams, updateQueryParams] = useTransitionSafeQueryParams({ + * page: NumberParam, + * search: StringParam + * }); + * + * // Update URL parameters + * updateQueryParams({ page: 2 }, 'pushIn'); + */ + +export function useTransitionSafeQueryParams< + QPCMap extends QueryParamConfigMap, +>(paramConfigMap: QPCMap) { + const [query, setQuery] = useQueryParams(paramConfigMap); + + const { paramKeys, defaultValues } = useMemo( + () => ({ + paramKeys: Object.keys(paramConfigMap), + defaultValues: _.mapValues(paramConfigMap, (config) => config.default), + }), + // Memoize based on the stringified version of the config map + // eslint-disable-next-line react-hooks/exhaustive-deps + [JSON.stringify(paramConfigMap)], + ); + + const pageQueryParamDeltaAtom = usePageQueryParamDeltaAtom(); + // Create page-specific atoms that reset when pagePath changes + const localQueryParamsAtom = useMemo(() => { + // Internal state atom for this specific page + let hasLoadedFromURL = false; + + // Derived atom that manages the query parameters + const queryParamsAtom = atomWithDefault((get) => { + const pageQueryParamDelta = get(pageQueryParamDeltaAtom); + + // Use URL query params on first access, then use internal state + if (!hasLoadedFromURL) { + hasLoadedFromURL = true; + return query; + } + + return _.pick( + { + ...query, + ...defaultValues, + ...pageQueryParamDelta, + }, + paramKeys, + ) as DecodedValueMap; + }); + + return queryParamsAtom; + + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [defaultValues, pageQueryParamDeltaAtom, paramKeys]); // New atoms when page changes and paramKeys change + + const localQueryParams = useAtomValue(localQueryParamsAtom); + const setPageQueryParamDelta = useSetAtom(pageQueryParamDeltaAtom); + + const updateQueryParams = useCallback( + ( + nextQueryParams: + | Partial> + | (( + prevQueryParams: DecodedValueMap, + ) => Partial>), + updateType: UrlUpdateType, + ) => { + const resolvedNextParams = + typeof nextQueryParams === 'function' + ? nextQueryParams(localQueryParams) + : nextQueryParams; + + // Update internal state atom + if (updateType === 'replaceIn' || updateType === 'pushIn') { + // do not touch + setPageQueryParamDelta((prev) => ({ + ...prev, + ...resolvedNextParams, + })); + } else { + setPageQueryParamDelta(resolvedNextParams as DecodedValueMap); + } + + // Sync all(merged) query parameters with URL + setQuery( + { + ...localQueryParams, + ...resolvedNextParams, + }, + updateType, + ); + }, + [localQueryParams, setQuery, setPageQueryParamDelta], + ); + + return [localQueryParams, updateQueryParams] as const; +} + +// Holds only the "delta" (patch) of query params modified after the initial URL load for the current page (pathname). +// It does NOT mirror the full set of current URL query params; instead it stores overrides that will be merged +// with defaults + initially decoded URL values. +const PageQueryParamDeltaAtomContext = createContext< + PrimitiveAtom> | undefined +>(undefined); + +function usePageQueryParamDeltaAtom() { + const ctx = use(PageQueryParamDeltaAtomContext); + if (!ctx) { + throw new Error( + 'usePageQueryParamDeltaAtom must be used within ', + ); + } + return ctx; +} + +export function PageQueryParamAtomProvider({ + children, +}: { + children: React.ReactNode; +}) { + const { pathname } = useLocation(); + + // When the provider re-renders and the pathname changes, a new instance is provided. + const pageQueryParamDeltaAtom = useMemo( + () => atom>({}), + // eslint-disable-next-line react-hooks/exhaustive-deps + [pathname], + ); + + return ( + + {children} + + ); +} diff --git a/react/src/pages/ComputeSessionListPage.tsx b/react/src/pages/ComputeSessionListPage.tsx index 5d4bb2c561..9548ff60e0 100644 --- a/react/src/pages/ComputeSessionListPage.tsx +++ b/react/src/pages/ComputeSessionListPage.tsx @@ -24,7 +24,7 @@ import { useCurrentProjectValue, useCurrentResourceGroupState, } from '../hooks/useCurrentProject'; -import { useDeferredQueryParams } from '../hooks/useDeferredQueryParams'; +import { useTransitionSafeQueryParams } from '../hooks/useTransitionSafeQueryParams'; import { SESSION_LAUNCHER_NOTI_PREFIX } from './SessionLauncherPage'; import { useUpdateEffect } from 'ahooks'; import { @@ -101,7 +101,7 @@ const ComputeSessionListPage = () => { pageSize: 10, }); - const [queryParams, setQuery] = useDeferredQueryParams({ + const [queryParams, setQuery] = useTransitionSafeQueryParams({ order: withDefault(StringParam, '-created_at'), filter: withDefault(StringParam, undefined), type: withDefault(StringParam, 'all'), diff --git a/react/src/pages/ServingPage.tsx b/react/src/pages/ServingPage.tsx index e921f6ea87..af4b0c9ac4 100644 --- a/react/src/pages/ServingPage.tsx +++ b/react/src/pages/ServingPage.tsx @@ -6,7 +6,7 @@ import { useUpdatableState, useWebUINavigate } from '../hooks'; import { useCurrentUserRole } from '../hooks/backendai'; import { useBAIPaginationOptionStateOnSearchParam } from '../hooks/reactPaginationQueryOptions'; import { useCurrentProjectValue } from '../hooks/useCurrentProject'; -import { useDeferredQueryParams } from '../hooks/useDeferredQueryParams'; +import { useTransitionSafeQueryParams } from '../hooks/useTransitionSafeQueryParams'; import { Button, Skeleton, theme } from 'antd'; import { filterOutEmpty, @@ -28,7 +28,7 @@ const ServingPage: React.FC = () => { const webuiNavigate = useWebUINavigate(); const currentProject = useCurrentProjectValue(); - const [queryParams, setQuery] = useDeferredQueryParams({ + const [queryParams, setQuery] = useTransitionSafeQueryParams({ order: withDefault(StringParam, '-created_at'), filter: StringParam, lifecycleStage: withDefault(StringParam, 'active'), diff --git a/react/src/pages/VFolderNodeListPage.tsx b/react/src/pages/VFolderNodeListPage.tsx index ff796f5864..f6827bc00d 100644 --- a/react/src/pages/VFolderNodeListPage.tsx +++ b/react/src/pages/VFolderNodeListPage.tsx @@ -21,7 +21,7 @@ import { } from '../hooks'; import { useBAIPaginationOptionStateOnSearchParam } from '../hooks/reactPaginationQueryOptions'; import { useCurrentProjectValue } from '../hooks/useCurrentProject'; -import { useDeferredQueryParams } from '../hooks/useDeferredQueryParams'; +import { useTransitionSafeQueryParams } from '../hooks/useTransitionSafeQueryParams'; import { useVFolderInvitationsValue } from '../hooks/useVFolderInvitations'; import { useToggle } from 'ahooks'; import { @@ -132,7 +132,7 @@ const VFolderNodeListPage: React.FC = ({ pageSize: 10, }); - const [queryParams, setQuery] = useDeferredQueryParams({ + const [queryParams, setQuery] = useTransitionSafeQueryParams({ order: withDefault(StringParam, '-created_at'), filter: withDefault(StringParam, undefined), statusCategory: withDefault(StringParam, 'active'),