[Cosmos] Honor max_item_count for query_items(feed_range=...)

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,8 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bug where `query_items(feed_range=..., max_item_count=N)` could return up to `K * N` documents per page when the supplied feed range overlapped `K` physical partition key ranges (for example, after a server-side split). The page returned to the caller is now truncated to the requested `max_item_count`.
+  * Known limitation (deferred): when a `feed_range` overlaps multiple PK ranges, only the last inner range's `x-ms-continuation` is surfaced as the page continuation token. Round-tripping that token sends it to every inner range on the next page, which is undefined server-side and can produce duplicates, missing documents, or non-terminating iteration on subsequent pages. As a safety mitigation, when the merged page is actually truncated the continuation header is suppressed so the truncated page is observed as terminal rather than producing wrong results on resume. A composite continuation token across overlapping inner ranges is tracked separately as a follow-up.
 * Fixed bug where `CosmosClient` construction with AAD credentials would crash at startup if the semantic reranking inference endpoint environment variable was not set, even when semantic reranking was not being used. The inference service is now lazily initialized on first use. See [PR 46243](https://github.com/Azure/azure-sdk-for-python/pull/46243)
 
 #### Other Changes

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -3408,6 +3408,61 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                     response_hook(last_response_headers, partial_result)
             # if the prefix partition query has results lets return it
             if results:
+                # Honor the user-requested page size (maxItemCount) across the K overlapping
+                # inner physical PK ranges. Each inner __Post above honors max_item_count per
+                # range, so the merged result can hold up to K * max_item_count documents.
+                # Truncate to the user-requested cap so a single page never exceeds it.
+                #
+                # NOTE: Use a truthy check (mirrors the contract used by _base.GetHeaders,
+                # which only emits the x-ms-max-item-count header when options['maxItemCount']
+                # is truthy). max_item_count=0/None/missing all mean "use the server default
+                # page size" and must be a no-op here, otherwise we would silently return
+                # empty pages while the server-side default page actually returned data.
+                #
+                # Known limitation (intentionally deferred — tracked separately as a
+                # follow-up: "[Cosmos] feed_range query continuation token replays documents
+                # from non-cursor PK ranges"):
+                # The continuation token surfaced to the caller (last_response_headers
+                # below is the K-th iteration's headers) is only the K-th inner range's
+                # x-ms-continuation. When the caller round-trips that token on the next
+                # page, __QueryFeed re-resolves the K overlapping ranges and sends the
+                # same caller-supplied continuation to every inner POST. Range K-1's
+                # continuation against ranges 0..K-2 is undefined server-side (may error,
+                # may restart from the beginning, may return undefined slices), so callers
+                # paginating a feed_range that overlaps multiple PK ranges may observe
+                # duplicates, missing documents, or non-terminating iteration. The correct
+                # fix is a composite continuation token spanning all K inner PK ranges.
+                # Until that lands, this branch only delivers a correct *first* page for
+                # multi-range feed_range queries.
+                #
+                # Mitigation for the deferred limitation: when truncation actually
+                # discards documents, the surfaced continuation token would describe
+                # the wrong cursor (it is the last inner range's token, but documents
+                # from earlier ranges and from the truncated tail were dropped). To
+                # avoid silent data loss, we strip the continuation header below when
+                # truncation occurs, so the truncated page is observed as terminal
+                # rather than producing wrong results on resume.
+                max_item_count = options.get("maxItemCount")
+                docs = results.get("Documents")
+                if max_item_count and isinstance(docs, list):
+                    try:
+                        cap = int(max_item_count)
+                    except (TypeError, ValueError):
+                        cap = 0
+                    if 0 < cap < len(docs):
+                        results["Documents"] = docs[:cap]
+                        # Keep the internal _count field consistent with the truncated
+                        # Documents list so any downstream consumer that introspects
+                        # the merged dict observes a coherent shape.
+                        results["_count"] = cap
+                        # The merged page was assembled from multiple inner PK ranges and
+                        # then truncated. The continuation token from the last inner range
+                        # only describes progress for that range, not for the truncated
+                        # tail or for ranges past the cursor, so resuming with it would
+                        # silently skip documents. Until a composite continuation token is
+                        # implemented, suppress the misleading token so callers see the
+                        # truncated page as terminal rather than experiencing data loss.
+                        last_response_headers.pop(http_constants.HttpHeaders.Continuation, None)
                 if last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None:
                     index_metrics_raw = last_response_headers[http_constants.HttpHeaders.IndexUtilization]
                     last_response_headers[http_constants.HttpHeaders.IndexUtilization] = (

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -3202,6 +3202,60 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                     response_hook(self.last_response_headers, partial_result)
             # if the prefix partition query has results lets return it
             if results:
+                # Honor the user-requested page size (maxItemCount) across the K overlapping
+                # inner physical PK ranges. Each inner __Post above honors max_item_count per
+                # range, so the merged result can hold up to K * max_item_count documents.
+                # Truncate to the user-requested cap so a single page never exceeds it.
+                #
+                # NOTE: Use a truthy check (mirrors the contract used by _base.GetHeaders,
+                # which only emits the x-ms-max-item-count header when options['maxItemCount']
+                # is truthy). max_item_count=0/None/missing all mean "use the server default
+                # page size" and must be a no-op here, otherwise we would silently return
+                # empty pages while the server-side default page actually returned data.
+                #
+                # Known limitation (intentionally deferred — tracked separately as a
+                # follow-up: "[Cosmos] feed_range query continuation token replays documents
+                # from non-cursor PK ranges"):
+                # The continuation token surfaced to the caller is only the K-th inner
+                # range's x-ms-continuation. When the caller round-trips that token on the
+                # next page, _CosmosClientConnection__QueryFeed re-resolves the K overlapping
+                # ranges and sends the same caller-supplied continuation to every inner POST.
+                # Range K-1's continuation against ranges 0..K-2 is undefined server-side
+                # (may error, may restart from the beginning, may return undefined slices),
+                # so callers paginating a feed_range that overlaps multiple PK ranges may
+                # observe duplicates, missing documents, or non-terminating iteration. The
+                # correct fix is a composite continuation token spanning all K inner PK
+                # ranges. Until that lands, this branch only delivers a correct *first* page
+                # for multi-range feed_range queries.
+                #
+                # Mitigation for the deferred limitation: when truncation actually
+                # discards documents, the surfaced continuation token would describe
+                # the wrong cursor (it is the last inner range's token, but documents
+                # from earlier ranges and from the truncated tail were dropped). To
+                # avoid silent data loss, we strip the continuation header below when
+                # truncation occurs, so the truncated page is observed as terminal
+                # rather than producing wrong results on resume.
+                max_item_count = options.get("maxItemCount")
+                docs = results.get("Documents")
+                if max_item_count and isinstance(docs, list):
+                    try:
+                        cap = int(max_item_count)
+                    except (TypeError, ValueError):
+                        cap = 0
+                    if 0 < cap < len(docs):
+                        results["Documents"] = docs[:cap]
+                        # Keep the internal _count field consistent with the truncated
+                        # Documents list so any downstream consumer that introspects
+                        # the merged dict observes a coherent shape.
+                        results["_count"] = cap
+                        # The merged page was assembled from multiple inner PK ranges and
+                        # then truncated. The continuation token from the last inner range
+                        # only describes progress for that range, not for the truncated
+                        # tail or for ranges past the cursor, so resuming with it would
+                        # silently skip documents. Until a composite continuation token is
+                        # implemented, suppress the misleading token so callers see the
+                        # truncated page as terminal rather than experiencing data loss.
+                        self.last_response_headers.pop(http_constants.HttpHeaders.Continuation, None)
                 if self.last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None:
                     index_metrics_raw = self.last_response_headers[http_constants.HttpHeaders.IndexUtilization]
                     self.last_response_headers[http_constants.HttpHeaders.IndexUtilization] = (