fix(s3vectors): restore query pagination

src/dynavec/stores/async_s3vectors.py

@@ -257,7 +257,7 @@ async def query_pages(
         return_distance: bool = True,
         page_size: int | None = None,
     ) -> AsyncIterator[list[dict[str, Any]]]:
-        """Yield result pages from the async S3 Vectors paginator."""
+        """Yield async result pages, following S3 Vectors ``nextToken``."""
         if top_k <= 0:
             return
 
@@ -289,22 +289,13 @@ async def query_pages(
         yielded = 0
         buffer: list[dict[str, Any]] = []
 
-        with translate_client_errors():
-            paginator = client.get_paginator("query_vectors")
-            page_iterator = paginator.paginate(
-                PaginationConfig={"MaxItems": top_k},
-                **kwargs,
-            ).__aiter__()
-
         while True:
             if self._query_limiter is not None:
                 await self._query_limiter.acquire_async()
 
-            try:
-                with translate_client_errors():
-                    page = await page_iterator.__anext__()
-            except StopAsyncIteration:
-                break
+            # Older supported SDKs expose QueryVectors without a paginator.
+            with translate_client_errors():
+                page = await client.query_vectors(**kwargs)
 
             vectors = page.get("vectors", [])
 
@@ -325,25 +316,23 @@ async def query_pages(
                     buffer.extend(vectors)
 
                     while (
-                        len(buffer) >= effective_page_size
-                        and yielded < top_k
+                        yielded < top_k
+                        and len(buffer) >= min(effective_page_size, top_k - yielded)
                     ):
-                        chunk = buffer[:effective_page_size]
-                        buffer = buffer[effective_page_size:]
-
-                        remaining = top_k - yielded
-
-                        if len(chunk) > remaining:
-                            chunk = chunk[:remaining]
+                        chunk_size = min(effective_page_size, top_k - yielded)
+                        chunk = buffer[:chunk_size]
+                        buffer = buffer[chunk_size:]
 
                         yield chunk
                         yielded += len(chunk)
 
                         if yielded >= top_k:
                             return
 
-            if page.get("NextToken") is None:
+            next_token = page.get("nextToken")
+            if not next_token:
                 break
+            kwargs["nextToken"] = next_token
 
         if (
             effective_page_size is not None

src/dynavec/stores/s3vectors.py

@@ -195,7 +195,7 @@ def query_pages(
         return_distance: bool = True,
         page_size: int | None = None,
     ) -> Iterator[list[dict[str, Any]]]:
-        """Yield result pages as the S3 Vectors paginator returns them.
+        """Yield result pages from S3 Vectors, following ``nextToken``.
 
         This is what powers streaming search: an agent can begin consuming the
         first page while later pages are still in flight.
@@ -233,24 +233,13 @@ def query_pages(
         yielded = 0
         buffer: list[dict[str, Any]] = []
 
-        with translate_client_errors():
-            paginator = self._client.get_paginator("query_vectors")
-            page_iterator = iter(
-                paginator.paginate(
-                    PaginationConfig={"MaxItems": top_k},
-                    **kwargs,
-                )
-            )
-
         while True:
             if self._query_limiter is not None:
                 self._query_limiter.acquire()
 
-            try:
-                with translate_client_errors():
-                    page = next(page_iterator)
-            except StopIteration:
-                break
+            # Older supported SDKs expose QueryVectors without a paginator.
+            with translate_client_errors():
+                page = self._client.query_vectors(**kwargs)
 
             vectors = page.get("vectors", [])
 
@@ -270,23 +259,23 @@ def query_pages(
                 else:
                     buffer.extend(vectors)
 
-                    while len(buffer) >= effective_page_size and yielded < top_k:
-                        chunk = buffer[:effective_page_size]
-                        buffer = buffer[effective_page_size:]
-
-                        remaining = top_k - yielded
-
-                        if len(chunk) > remaining:
-                            chunk = chunk[:remaining]
+                    while yielded < top_k and len(buffer) >= min(
+                        effective_page_size, top_k - yielded
+                    ):
+                        chunk_size = min(effective_page_size, top_k - yielded)
+                        chunk = buffer[:chunk_size]
+                        buffer = buffer[chunk_size:]
 
                         yield chunk
                         yielded += len(chunk)
 
                         if yielded >= top_k:
                             return
 
-            if page.get("NextToken") is None:
+            next_token = page.get("nextToken")
+            if not next_token:
                 break
+            kwargs["nextToken"] = next_token
 
         if effective_page_size is not None and buffer and yielded < top_k:
             remaining = top_k - yielded

tests/test_async_s3vectors.py

@@ -6,24 +6,14 @@
 from typing import Any
 from unittest.mock import AsyncMock
 
+import aioboto3
 import pytest
+from aiobotocore.stub import AioStubber
 
 from dynavec.config import DynavecConfig
 from dynavec.stores.async_s3vectors import AsyncS3VectorsStore
 
 
-class FakeAsyncPaginator:
-    def __init__(self, pages: list[dict[str, Any]]) -> None:
-        self.pages = pages
-        self.calls: list[dict[str, Any]] = []
-
-    async def paginate(self, **kwargs: Any):
-        self.calls.append(kwargs)
-
-        for page in self.pages:
-            yield page
-
-
 def _config() -> DynavecConfig:
     return DynavecConfig(
         vector_bucket="test-bucket",
@@ -71,14 +61,15 @@ def __init__(
     ) -> None:
         self.put_calls: list[dict[str, Any]] = []
         self.get_calls: list[dict[str, Any]] = []
-        self.paginator = FakeAsyncPaginator(pages or [])
+        self.pages = iter(pages or [{"vectors": []}])
+        self.query_calls: list[dict[str, Any]] = []
 
     async def put_vectors(self, **kwargs: Any) -> None:
         self.put_calls.append(kwargs)
 
-    def get_paginator(self, name: str) -> FakeAsyncPaginator:
-        assert name == "query_vectors"
-        return self.paginator
+    async def query_vectors(self, **kwargs: Any) -> dict[str, Any]:
+        self.query_calls.append(kwargs)
+        return next(self.pages)
 
     async def get_vectors(self, **kwargs: Any) -> dict[str, Any]:
         self.get_calls.append(kwargs)
@@ -188,7 +179,7 @@ async def test_query_returns_paginated_results():
                 {"key": f"k{i}", "distance": i * 0.01}
                 for i in range(100)
             ],
-            "NextToken": "token-1",
+            "nextToken": "token-1",
         },
         {
             "vectors": [
@@ -210,6 +201,26 @@ async def test_query_returns_paginated_results():
     assert len(results) == 150
     assert results[0]["key"] == "k0"
     assert results[-1]["key"] == "k149"
+    assert "nextToken" not in client.query_calls[0]
+    assert client.query_calls[1]["nextToken"] == "token-1"
+
+
+async def test_query_with_real_async_sdk_client():
+    session = aioboto3.Session(
+        aws_access_key_id="testing", aws_secret_access_key="testing", region_name="us-east-1"
+    )
+    vectors = [{"key": "doc-1", "distance": 0.25}]
+
+    async with AsyncS3VectorsStore(_config(), session) as store:
+        with AioStubber(store._client) as stubber:
+            stubber.add_response("query_vectors", {"vectors": vectors}, {
+                "vectorBucketName": "test-bucket", "indexName": "test-index", "topK": 1,
+                "queryVector": {"float32": [1.0, 2.0, 3.0, 4.0]},
+                "returnMetadata": True, "returnDistance": True,
+            })
+
+            assert await store.query([1.0, 2.0, 3.0, 4.0], top_k=1) == vectors
+            stubber.assert_no_pending_responses()
 
 
 async def test_query_truncates_at_top_k():
@@ -219,7 +230,7 @@ async def test_query_truncates_at_top_k():
                 {"key": f"k{i}"}
                 for i in range(100)
             ],
-            "NextToken": "token-1",
+            "nextToken": "token-1",
         },
         {
             "vectors": [
@@ -283,6 +294,49 @@ async def test_query_rejects_excessive_top_k():
             )
 
 
+@pytest.mark.parametrize("page_size", [None, 2])
+async def test_query_pages_continues_after_empty_page_with_next_token(page_size):
+    client = FakeS3Client([
+        {"vectors": [{"key": "k0"}], "nextToken": "token-1"},
+        {"vectors": [], "nextToken": "token-2"},
+        {"vectors": [{"key": "k1"}, {"key": "k2"}]},
+    ])
+    limiter = AsyncMock()
+
+    async with AsyncS3VectorsStore(_config(), FakeSession(client)) as store:
+        store._query_limiter = limiter
+        pages = [
+            page async for page in store.query_pages(
+                [0.1, 0.2, 0.3, 0.4], top_k=5, page_size=page_size
+            )
+        ]
+
+    assert [hit["key"] for page in pages for hit in page] == ["k0", "k1", "k2"]
+    assert [len(page) for page in pages] == ([1, 2] if page_size is None else [2, 1])
+    assert limiter.acquire_async.await_count == 3
+
+
+async def test_query_pages_stops_at_top_k_with_large_page_size():
+    client = FakeS3Client([
+        {"vectors": [{"key": "k0"}], "nextToken": "token-1"},
+        {"vectors": [{"key": "k1"}, {"key": "k2"}], "nextToken": "unused"},
+    ])
+
+    async with AsyncS3VectorsStore(_config(), FakeSession(client)) as store:
+        pages = [page async for page in store.query_pages(
+            [1.0, 2.0, 3.0, 4.0], top_k=2, page_size=100,
+            filter={"kind": "test"}, return_metadata=False, return_distance=False,
+        )]
+
+    assert pages == [[{"key": "k0"}, {"key": "k1"}]]
+    assert len(client.query_calls) == 2
+    for call in client.query_calls:
+        assert call["topK"] == 2
+        assert call["filter"] == {"kind": "test"}
+        assert call["returnMetadata"] is False
+        assert call["returnDistance"] is False
+
+
 async def test_put_batch_uses_rate_limiter():
     config = DynavecConfig(
         vector_bucket="test-bucket",
@@ -319,11 +373,11 @@ async def test_query_pages_uses_rate_limiter_per_page():
     pages = [
         {
             "vectors": [{"key": "k0"}],
-            "NextToken": "token-1",
+            "nextToken": "token-1",
         },
         {
             "vectors": [{"key": "k1"}],
-            "NextToken": "token-2",
+            "nextToken": "token-2",
         },
         {
             "vectors": [{"key": "k2"}],

tests/test_logging.py

@@ -130,10 +130,9 @@ def test_dynamodb_store_structured_logging(caplog):
 def test_s3vectors_store_structured_logging(caplog):
     session = MagicMock()
     client = session.client.return_value
-    paginator = client.get_paginator.return_value
-    paginator.paginate.return_value = [
-        {"vectors": [{"key": "ns#1", "data": {"float32": [0.1] * 8}, "distance": 0.05}]}
-    ]
+    client.query_vectors.return_value = {
+        "vectors": [{"key": "ns#1", "data": {"float32": [0.1] * 8}, "distance": 0.05}]
+    }
 
     # 1. Opt-in: When False, no logs
     cfg_disabled = DynavecConfig(