fix(mcp): read total_datasets so the fetch-all early-exit works (#19020)

Fix #19019 — MCP _fetch_all_datasets() used wrong total key (total vs total_datasets), causing extra empty-page requests. Now reads total_datasets with fallback.

api/utils/api_utils.py

@@ -304,7 +304,7 @@ def get_result(code=RetCode.SUCCESS, message="", data=None, total=None):
     {
         "code": 0,
         "data": [...],        # List or object, backward compatible
-        "total": 47,          # Optional field for pagination
+        "total_datasets": 47, # Optional field for pagination
         "message": "..."      # Error or status message
     }
     """

internal/handler/mcp_pagination_test.go

@@ -0,0 +1,95 @@
+package handler
+
+import (
+	"fmt"
+	"testing"
+)
+
+func fakeDatasetPager(total int, pageSize int) (func(page, size int) ([]map[string]interface{}, int64, error), *[]int) {
+	pagesRequested := []int{}
+	page := func(p, size int) ([]map[string]interface{}, int64, error) {
+		pagesRequested = append(pagesRequested, p)
+		start := (p - 1) * size
+		end := start + size
+		if end > total {
+			end = total
+		}
+		if start >= total {
+			return nil, int64(total), nil
+		}
+		data := make([]map[string]interface{}, 0, end-start)
+		for i := start; i < end; i++ {
+			data = append(data, map[string]interface{}{"id": fmt.Sprintf("dataset-%d", i)})
+		}
+		return data, int64(total), nil
+	}
+	return page, &pagesRequested
+}
+
+func TestFetchAllDatasetIDsStopsAtTotalWithoutEmptyPage(t *testing.T) {
+	const pageSize = 100
+	tests := []struct {
+		name        string
+		total       int
+		wantPages   []int
+		wantIDCount int
+	}{
+		{"exact multiple of page size", 200, []int{1, 2}, 200},
+		{"non multiple of page size", 250, []int{1, 2, 3}, 250},
+		{"single full page", 100, []int{1}, 100},
+		{"single short page", 42, []int{1}, 42},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			page, pagesRequested := fakeDatasetPager(tt.total, pageSize)
+
+			ids, err := fetchAllDatasetIDs(page, pageSize)
+
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if len(ids) != tt.wantIDCount {
+				t.Fatalf("got %d ids, want %d", len(ids), tt.wantIDCount)
+			}
+			if len(*pagesRequested) != len(tt.wantPages) {
+				t.Fatalf("requested pages %v, want %v", *pagesRequested, tt.wantPages)
+			}
+			for i, p := range tt.wantPages {
+				if (*pagesRequested)[i] != p {
+					t.Fatalf("requested pages %v, want %v", *pagesRequested, tt.wantPages)
+				}
+			}
+		})
+	}
+}
+
+func TestFetchAllDatasetIDsPropagatesErrors(t *testing.T) {
+	wantErr := fmt.Errorf("backend unavailable")
+	page := func(page, size int) ([]map[string]interface{}, int64, error) {
+		return nil, 0, wantErr
+	}
+
+	if _, err := fetchAllDatasetIDs(page, 100); err == nil {
+		t.Fatal("expected the listPage error to propagate")
+	}
+}
+
+func TestFetchAllDatasetIDsSkipsMalformedRows(t *testing.T) {
+	page := func(p, size int) ([]map[string]interface{}, int64, error) {
+		return []map[string]interface{}{
+			{"id": "dataset-1"},
+			{"id": ""},               // empty id skipped
+			{"id": 42},               // non-string id skipped
+			{"description": "no id"}, // missing id skipped
+		}, 4, nil
+	}
+
+	ids, err := fetchAllDatasetIDs(page, 100)
+
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(ids) != 1 || ids[0] != "dataset-1" {
+		t.Fatalf("got %v, want [dataset-1]", ids)
+	}
+}

internal/handler/mcp_server.go

@@ -150,32 +150,20 @@ func MCPRetrieval(ctx context.Context, ds *dataset.DatasetService, userID string
 	datasetIDs := req.DatasetIDs
 	if len(datasetIDs) == 0 {
 		const maxPageSize = 100
-		page := 1
-		for {
-			data, _, _, err := ds.ListDatasets(ctx,
-				"", "", page, maxPageSize, "create_time", true,
+		ids, err := fetchAllDatasetIDs(func(page, pageSize int) ([]map[string]interface{}, int64, error) {
+			data, total, _, err := ds.ListDatasets(ctx,
+				"", "", page, pageSize, "create_time", true,
 				"", nil, "", userID, nil,
 			)
-			if err != nil {
-				return "", fmt.Errorf("cannot resolve accessible datasets: %w", err)
-			}
-			if len(data) == 0 {
-				break
-			}
-			for _, d := range data {
-				if id, ok := d["id"].(string); ok && id != "" {
-					datasetIDs = append(datasetIDs, id)
-				}
-			}
-			// A page smaller than maxPageSize is the last page.
-			if len(data) < maxPageSize {
-				break
-			}
-			page++
+			return data, total, err
+		}, maxPageSize)
+		if err != nil {
+			return "", fmt.Errorf("cannot resolve accessible datasets: %w", err)
 		}
-		if len(datasetIDs) == 0 {
+		if len(ids) == 0 {
 			return "", fmt.Errorf("no accessible datasets found")
 		}
+		datasetIDs = ids
 	}
 
 	searchReq := &service.SearchDatasetsRequest{
@@ -225,3 +213,39 @@ func MCPRetrieval(ctx context.Context, ds *dataset.DatasetService, userID string
 	}
 	return string(result), nil
 }
+
+// fetchAllDatasetIDs pages through listPage collecting dataset IDs until the
+// total reported by the service is reached, or a short or empty page arrives.
+// Stopping at the reported total avoids an extra empty-page request when the
+// dataset count is an exact multiple of pageSize.
+func fetchAllDatasetIDs(listPage func(page, pageSize int) ([]map[string]interface{}, int64, error), pageSize int) ([]string, error) {
+	var ids []string
+	page := 1
+	fetched := 0
+	for {
+		data, total, err := listPage(page, pageSize)
+		if err != nil {
+			return nil, err
+		}
+		if len(data) == 0 {
+			break
+		}
+		fetched += len(data)
+		for _, d := range data {
+			if id, ok := d["id"].(string); ok && id != "" {
+				ids = append(ids, id)
+			}
+		}
+		// Stop once the reported total is reached so exact multiples of
+		// pageSize do not pay an extra empty-page request.
+		if total > 0 && int64(fetched) >= total {
+			break
+		}
+		// A page smaller than pageSize is the last page.
+		if len(data) < pageSize {
+			break
+		}
+		page++
+	}
+	return ids, nil
+}

mcp/server/server.py

@@ -226,7 +226,9 @@ class RAGFlowConnector:
                 break
 
             datasets.extend(page_datasets)
-            total = res_json.get("total")
+            # The REST API reports the dataset total under "total_datasets"
+            # (see api/utils/api_utils.py get_result).
+            total = res_json.get("total_datasets", res_json.get("total"))
             if total is not None and len(datasets) >= total:
                 break
 

test/unit_test/mcp/test_server_datasets.py

@@ -63,7 +63,9 @@ def _stub_dataset_pages(monkeypatch, connector, datasets):
         page_size = params["page_size"]
         start = (page - 1) * page_size
         end = start + page_size
-        return _FakeResponse({"code": 0, "data": datasets[start:end], "total": len(datasets)})
+        # Mirrors the real REST envelope: api/utils/api_utils.py get_result() puts
+        # the dataset total under "total_datasets", not "total".
+        return _FakeResponse({"code": 0, "data": datasets[start:end], "total_datasets": len(datasets)})
 
     monkeypatch.setattr(connector, "_get", _get)
     return requests
@@ -83,6 +85,21 @@ async def test_list_datasets_default_fetches_all_with_rest_page_size_limit(monke
     assert all(request["params"]["page_size"] == 100 for request in requests)
 
 
+@pytest.mark.asyncio
+async def test_list_datasets_stops_at_total_without_empty_page_request(monkeypatch, mcp_server):
+    """When the fetched rows reach the reported total_datasets, the loop must stop
+    without issuing the extra empty-page request. Reading the wrong key ("total")
+    left this early-exit dead and always cost one more roundtrip per full listing."""
+    connector = mcp_server.RAGFlowConnector(base_url=mcp_server.BASE_URL)
+    requests = _stub_dataset_pages(monkeypatch, connector, _datasets(200))
+
+    result = await connector.list_datasets(api_key="unit-key")
+
+    rows = [json.loads(line) for line in result.splitlines()]
+    assert len(rows) == 200
+    assert [request["params"]["page"] for request in requests] == [1, 2]
+
+
 @pytest.mark.asyncio
 async def test_resolve_dataset_ids_fetches_all_pages_and_deduplicates(monkeypatch, mcp_server):
     connector = mcp_server.RAGFlowConnector(base_url=mcp_server.BASE_URL)