From c187dfdc733582ca84669c62880dac8e0a09bc7b Mon Sep 17 00:00:00 2001 From: xu jiaming <78154672+xu-jia-ming@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:05:50 +0800 Subject: [PATCH] fix(mcp): read total_datasets so the fetch-all early-exit works (#19020) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 2 +- internal/handler/mcp_pagination_test.go | 95 ++++++++++++++++++++++ internal/handler/mcp_server.go | 66 ++++++++++----- mcp/server/server.py | 4 +- test/unit_test/mcp/test_server_datasets.py | 19 ++++- 5 files changed, 162 insertions(+), 24 deletions(-) create mode 100644 internal/handler/mcp_pagination_test.go diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py index fdbca04343..d2a2ef4ea1 100644 --- a/api/utils/api_utils.py +++ b/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 } """ diff --git a/internal/handler/mcp_pagination_test.go b/internal/handler/mcp_pagination_test.go new file mode 100644 index 0000000000..f856854bb6 --- /dev/null +++ b/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) + } +} diff --git a/internal/handler/mcp_server.go b/internal/handler/mcp_server.go index d7a8c509d7..721bb7e0bf 100644 --- a/internal/handler/mcp_server.go +++ b/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 +} diff --git a/mcp/server/server.py b/mcp/server/server.py index 801b42cb2b..fbb7f6051c 100644 --- a/mcp/server/server.py +++ b/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 diff --git a/test/unit_test/mcp/test_server_datasets.py b/test/unit_test/mcp/test_server_datasets.py index b7e0bc6a60..601f0e4722 100644 --- a/test/unit_test/mcp/test_server_datasets.py +++ b/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)