fix(dataset): echo back selected datasets not in the first page (#17616)

web/src/components/knowledge-base-item.tsx

@@ -1,5 +1,8 @@
 import { DocumentParserType } from '@/constants/knowledge';
-import { useFetchKnowledgeList } from '@/hooks/use-knowledge-request';
+import {
+  useFetchDatasetsByIds,
+  useFetchKnowledgeList,
+} from '@/hooks/use-knowledge-request';
 import { IDataset } from '@/interfaces/database/dataset';
 import { useBuildQueryVariableOptions } from '@/pages/agent/hooks/use-get-begin-query';
 import { useDebounce } from 'ahooks';
@@ -36,12 +39,31 @@ export function useDisableDifferenceEmbeddingDataset(name: string) {
   } = useFetchKnowledgeList(false, debouncedSearchString);
   const datasetCacheRef = useRef(new Map<string, IDataset>());
 
+  const selectedDatasetIds = useMemo(
+    () => (Array.isArray(datasetId) ? datasetId : []),
+    [datasetId],
+  );
+
+  // Selected dataset IDs that are neither in the currently loaded page nor
+  // already cached. These need to be fetched by ID so their names can be
+  // echoed back in the form field (the paginated list may not contain them).
+  const missingIds = useMemo(() => {
+    const loadedIds = new Set(datasetListOrigin.map((d) => d.id));
+    return selectedDatasetIds.filter(
+      (id) => !loadedIds.has(id) && !datasetCacheRef.current.has(id),
+    );
+  }, [datasetListOrigin, selectedDatasetIds]);
+
+  const { data: missingDatasets } = useFetchDatasetsByIds(missingIds);
+
   const datasetList = useMemo(() => {
     datasetListOrigin.forEach((dataset) => {
       datasetCacheRef.current.set(dataset.id, dataset);
     });
+    missingDatasets?.forEach((dataset) => {
+      datasetCacheRef.current.set(dataset.id, dataset);
+    });
 
-    const selectedDatasetIds = Array.isArray(datasetId) ? datasetId : [];
     const selectedDatasets = selectedDatasetIds
       .map((id) => datasetCacheRef.current.get(id))
       .filter(Boolean) as IDataset[];
@@ -54,7 +76,7 @@ export function useDisableDifferenceEmbeddingDataset(name: string) {
         ]),
       ).values(),
     );
-  }, [datasetId, datasetListOrigin]);
+  }, [datasetListOrigin, selectedDatasetIds, missingDatasets]);
 
   const selectedEmbedId = useMemo(() => {
     const data = datasetList?.find((item) => item.id === datasetId?.[0]);

web/src/hooks/use-knowledge-request.ts

@@ -41,6 +41,7 @@ import kbService, {
   listArtifacts,
   datasetFilter,
   listDataset,
+  listDatasetByIds,
   listTag,
   listWikiCommits,
   removeTag,
@@ -1003,6 +1004,8 @@ export const KnowledgeListKeys = {
       keywords,
       pageSize,
     ] as const,
+  byIds: (ids: string[]) =>
+    [KnowledgeApiAction.FetchKnowledgeList, 'byIds', ids] as const,
 };
 
 export const useFetchKnowledgeList = (
@@ -1114,6 +1117,26 @@ export const useSelectKnowledgeOptions = () => {
   return options;
 };
 
+/**
+ * Fetch datasets by a set of IDs. Used to resolve the names of
+ * already-selected datasets that are not present in the first page of
+ * the paginated list so they can be echoed back in the form field.
+ */
+export const useFetchDatasetsByIds = (ids: string[]) => {
+  const sortedIds = useMemo(() => [...ids].sort(), [ids]);
+  const { data, isFetching: loading } = useQuery<IDataset[]>({
+    queryKey: KnowledgeListKeys.byIds(sortedIds),
+    enabled: sortedIds.length > 0,
+    gcTime: 0,
+    queryFn: async () => {
+      const { data } = await listDatasetByIds(sortedIds);
+      return (data?.data ?? []) as IDataset[];
+    },
+  });
+
+  return { data, loading };
+};
+
 //#region tags
 export const useRenameTag = () => {
   const knowledgeBaseId = useKnowledgeBaseId();

web/src/services/knowledge-service.ts

@@ -268,6 +268,14 @@ export function deleteKnowledgeGraph(knowledgeId: string) {
 export const listDataset = (params?: IFetchKnowledgeListRequestParams) =>
   request.get(api.kbList, { params });
 
+// Fetch datasets by a set of IDs via the `ids` query param (comma-joined).
+// Used to echo back already-selected datasets whose names are not present
+// in the first page of the paginated list.
+export const listDatasetByIds = (ids: string[]) =>
+  request.get(api.kbList, {
+    params: { ids: ids.join(','), page_size: ids.length },
+  });
+
 export const datasetFilter = () => request.get(api.datasetFilter);
 
 export const updateKb = (datasetId: string, data: Record<string, any>) =>