Fix: MCP dataset discovery fails due to REST API max page size limit … (#16148)

Fix #16146

mcp/server/server.py

@@ -58,7 +58,8 @@ JSON_RESPONSE = True
 class RAGFlowConnector:
     _MAX_DATASET_CACHE = 32
     _CACHE_TTL = 300
-    _DATASET_PAGE_SIZE = 1000
+    # Keep in sync with api.utils.pagination_utils.REST_API_MAX_PAGE_SIZE.
+    _REST_API_MAX_PAGE_SIZE = 100
 
     _dataset_metadata_cache: OrderedDict[str, tuple[dict, float | int]] = OrderedDict()  # "dataset_id" -> (metadata, expiry_ts)
     _document_metadata_cache: OrderedDict[str, tuple[list[tuple[str, dict]], float | int]] = OrderedDict()  # "dataset_id" -> ([(document_id, doc_metadata)], expiry_ts)
@@ -162,11 +163,55 @@ class RAGFlowConnector:
 
         return res_json
 
-    async def list_datasets(self, *, api_key: str, page: int = 1, page_size: int = 1000, orderby: str = "create_time", desc: bool = True, id: str | None = None, name: str | None = None):
+    async def _fetch_all_datasets(
+        self,
+        *,
+        api_key: str,
+        orderby: str = "create_time",
+        desc: bool = True,
+        id: str | None = None,
+        name: str | None = None,
+    ):
+        """Fetch all accessible datasets without exceeding the REST API page-size limit."""
+        datasets = []
+        page = 1
+
+        while True:
+            logging.debug("fetching all /datasets page=%s page_size=%s", page, self._REST_API_MAX_PAGE_SIZE)
+            res_json = await self._fetch_datasets_page(
+                api_key=api_key,
+                page=page,
+                page_size=self._REST_API_MAX_PAGE_SIZE,
+                orderby=orderby,
+                desc=desc,
+                id=id,
+                name=name,
+            )
+            page_datasets = res_json.get("data", [])
+            logging.debug("received %s datasets from page=%s", len(page_datasets), page)
+            if not page_datasets:
+                break
+
+            datasets.extend(page_datasets)
+            total = res_json.get("total")
+            if total is not None and len(datasets) >= total:
+                break
+
+            page += 1
+
+        return datasets
+
+    async def list_datasets(self, *, api_key: str, page: int = 1, page_size: int = -1, orderby: str = "create_time", desc: bool = True, id: str | None = None, name: str | None = None):
         """Return accessible datasets as newline-delimited JSON for MCP tool descriptions."""
-        res_json = await self._fetch_datasets_page(api_key=api_key, page=page, page_size=page_size, orderby=orderby, desc=desc, id=id, name=name)
+        if page_size == -1:
+            datasets = await self._fetch_all_datasets(api_key=api_key, orderby=orderby, desc=desc, id=id, name=name)
+        else:
+            page_size = min(page_size, self._REST_API_MAX_PAGE_SIZE)
+            res_json = await self._fetch_datasets_page(api_key=api_key, page=page, page_size=page_size, orderby=orderby, desc=desc, id=id, name=name)
+            datasets = res_json["data"]
+
         result_list = []
-        for data in res_json["data"]:
+        for data in datasets:
             d = {"description": data["description"], "id": data["id"]}
             result_list.append(json.dumps(d, ensure_ascii=False))
         return "\n".join(result_list)
@@ -174,25 +219,13 @@ class RAGFlowConnector:
     async def resolve_dataset_ids(self, *, api_key: str):
         """Resolve all accessible dataset IDs for MCP retrieval fallback."""
         logging.info("Resolving accessible dataset IDs for MCP retrieval")
-        dataset_ids = []
-        page = 1
-
-        while True:
-            logging.debug("resolve_dataset_ids fetching /datasets page=%s page_size=%s", page, self._DATASET_PAGE_SIZE)
-            try:
-                res_json = await self._fetch_datasets_page(api_key=api_key, page=page, page_size=self._DATASET_PAGE_SIZE)
-            except Exception as exc:
-                logging.warning("resolve_dataset_ids failed to fetch /datasets page=%s error=%s", page, exc)
-                raise
-
-            datasets = res_json.get("data", [])
-            logging.debug("resolve_dataset_ids received %s datasets from page=%s", len(datasets), page)
-            dataset_ids.extend(data["id"] for data in datasets if data.get("id"))
-            total = res_json.get("total", len(dataset_ids))
-            if not datasets or len(dataset_ids) >= total:
-                break
-            page += 1
+        try:
+            datasets = await self._fetch_all_datasets(api_key=api_key)
+        except Exception as exc:
+            logging.warning("resolve_dataset_ids failed to fetch /datasets error=%s", exc)
+            raise
 
+        dataset_ids = [data["id"] for data in datasets if data.get("id")]
         resolved = list(dict.fromkeys(dataset_ids))
         logging.info("resolve_dataset_ids resolved %s accessible dataset IDs", len(resolved))
         return resolved