fix(vector-stores): paginate managed resource list by unified_resource_id cursor

`list_user_resources`, which backs the managed vector store listing, carried
the same broken cursor as the managed batch listing: `after` was applied as
`where id > after` against the table's random-uuid primary key, while the value
a client holds is the `unified_resource_id` returned as each item's `id`. That
compares a base64 cursor to unrelated uuids, so a page could re-select the rows
it just returned or skip rows that should have come next. Ordering by
non-unique `created_at` alone left the order partial as well, so rows sharing a
timestamp could shift between requests.

The listing now uses a Prisma cursor on the unique `unified_resource_id` with
`unified_resource_id` as a secondary sort key, rejects an `after` that does not
resolve to a resource the caller can list with a 400 rather than an empty page,
and derives `has_more` from an extra fetched row instead of from whether the
page came back full.

`build_list_page` read `first_id` / `last_id` through attribute access only,
which raised AttributeError on every non-empty page of this listing because it
builds plain dicts. It now reads the id from either shape.
This commit is contained in:
mateo-berri 2026-07-24 19:44:48 -07:00
parent 7b019cf152
commit f913d42abb
4 changed files with 253 additions and 21 deletions

View file

@ -542,35 +542,39 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
Returns:
Dictionary with list of resources and pagination info
"""
from fastapi import HTTPException
owner_filter = build_owner_filter(user_api_key_dict)
if owner_filter is None:
return build_list_page([])
where_clause: Dict[str, Any] = {**owner_filter}
where_clause: Dict[str, Any] = {**owner_filter, **(additional_filters or {})}
table = getattr(self.prisma_client.db, self.table_name)
if after:
where_clause["id"] = {"gt": after}
cursor_row = await table.find_first(where={**where_clause, "unified_resource_id": after})
if cursor_row is None:
raise HTTPException(
status_code=400,
detail=f"Invalid 'after' cursor: no {self.resource_type} found with id '{after}'.",
)
# Add additional filters
if additional_filters:
where_clause.update(additional_filters)
page_size = limit or 20
cursor_args: Dict[str, Any] = {"cursor": {"unified_resource_id": after}, "skip": 1} if after else {}
# Fetch resources
fetch_limit = limit or 20
table = getattr(self.prisma_client.db, self.table_name)
resources = await table.find_many(
where=where_clause,
take=fetch_limit,
order={"created_at": "desc"},
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_resource_id": "desc"}],
**cursor_args,
)
resource_objects: List[Any] = []
for resource in resources:
try:
# Stop once we have enough
if len(resource_objects) >= (limit or 20):
break
has_more = len(resources) > page_size
resource_objects: List[Any] = []
for resource in resources[:page_size]:
try:
# Parse resource object
resource_data = resource.resource_object
if isinstance(resource_data, str):
@ -590,4 +594,4 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
)
continue
return build_list_page(resource_objects, has_more=len(resource_objects) == (limit or 20))
return build_list_page(resource_objects, has_more=has_more)

View file

@ -9,6 +9,7 @@ identifying ids are denied so an empty user_id can never select an
unscoped query.
"""
from collections.abc import Mapping
from typing import Any, Dict, List, Optional
from litellm.proxy._types import (
@ -17,15 +18,23 @@ from litellm.proxy._types import (
)
def _item_id(item: Any) -> str | None:
if isinstance(item, Mapping):
return item.get("id")
return getattr(item, "id", None)
def build_list_page(items: List[Any], has_more: bool = False) -> Dict[str, Any]:
"""Build the OpenAI-style paginated list response shape used by managed
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
sourced from each item's ``.id`` attribute."""
file/batch/vector-store listings. ``first_id`` and ``last_id`` are sourced
from each item's ``id``, which is an attribute on the pydantic objects the
file and batch listings build and a key on the plain dicts the vector-store
listing builds."""
return {
"object": "list",
"data": items,
"first_id": items[0].id if items else None,
"last_id": items[-1].id if items else None,
"first_id": _item_id(items[0]) if items else None,
"last_id": _item_id(items[-1]) if items else None,
"has_more": has_more,
}

View file

@ -108,6 +108,184 @@ async def test_list_identity_less_caller_returns_empty_without_query():
resource.prisma_client.db.litellm_test_resource_table.find_many.assert_not_awaited()
def _make_row(index: int, created_by: str = "alice"):
row = MagicMock()
row.id = f"pk-{index:03d}"
row.unified_resource_id = f"unified-resource-{index:03d}"
row.created_at = 1_000_000 + index
row.created_by = created_by
row.resource_object = {
"id": f"provider-vector-store-{index:03d}",
"object": "vector_store",
"name": f"store-{index:03d}",
}
return row
def _row_matches(row, where) -> bool:
return all(
isinstance(value, dict) or getattr(row, field, None) == value
for field, value in where.items()
)
def _make_paginating_resource(rows: List) -> _StubResource:
async def find_many(where, take, order, cursor=None, skip=0):
result = sorted(
rows, key=lambda r: (r.created_at, r.unified_resource_id), reverse=True
)
id_filter = where.get("id")
if isinstance(id_filter, dict) and "gt" in id_filter:
result = [r for r in result if r.id > id_filter["gt"]]
if cursor is not None:
(cur_field, cur_val), = cursor.items()
idx = next(
(i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
None,
)
if idx is None:
return []
result = result[idx + skip:]
return result[:take]
async def find_first(where):
return next((r for r in rows if _row_matches(r, where)), None)
cache = MagicMock()
cache.async_get_cache = AsyncMock(return_value=None)
prisma = MagicMock()
table = MagicMock()
table.find_many = AsyncMock(side_effect=find_many)
table.find_first = AsyncMock(side_effect=find_first)
prisma.db = MagicMock()
setattr(prisma.db, "litellm_test_resource_table", table)
return _StubResource(internal_usage_cache=cache, prisma_client=prisma)
@pytest.mark.asyncio
async def test_list_resources_pagination_uses_unified_resource_id_cursor():
"""The ``after`` cursor a client sends back is a resource's
``unified_resource_id`` -- that is what the listing returns as each item's
``id`` and as ``last_id``. Paginating must use a Prisma cursor on that
unique column, not a ``where id > after`` filter against the random-uuid
primary key, which compares the cursor to unrelated values and so loops or
silently drops resources.
"""
rows = [_make_row(i) for i in range(3)]
resource = _make_paginating_resource(rows)
await resource.list_user_resources(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
limit=2,
after=rows[2].unified_resource_id,
)
table = resource.prisma_client.db.litellm_test_resource_table
kwargs = table.find_many.await_args.kwargs
assert kwargs["cursor"] == {"unified_resource_id": rows[2].unified_resource_id}
assert kwargs["skip"] == 1
assert kwargs["order"] == [
{"created_at": "desc"},
{"unified_resource_id": "desc"},
]
assert "id" not in kwargs["where"]
@pytest.mark.asyncio
async def test_list_resources_pagination_walks_all_pages_without_loops_or_gaps():
"""Walking pages the way a client does -- feeding ``last_id`` back as
``after`` -- must return every resource exactly once, newest first."""
rows = [_make_row(i) for i in range(5)]
resource = _make_paginating_resource(rows)
user = UserAPIKeyAuth(user_id="alice")
seen = []
after = None
for _ in range(len(rows) + 5):
page = await resource.list_user_resources(
user_api_key_dict=user, limit=2, after=after
)
seen.extend(item["id"] for item in page["data"])
if not page["has_more"]:
break
assert page["last_id"] is not None, "has_more was true but there is no cursor"
assert page["last_id"] != after, "cursor did not advance (pagination loop)"
after = page["last_id"]
assert seen == [r.unified_resource_id for r in reversed(rows)]
assert len(seen) == len(set(seen))
@pytest.mark.asyncio
async def test_list_resources_has_more_false_on_exactly_full_final_page():
"""``has_more`` must mean "another row exists", not "this page is full",
so a resource count that is an exact multiple of ``limit`` does not cost
every client an extra empty request."""
rows = [_make_row(i) for i in range(4)]
resource = _make_paginating_resource(rows)
user = UserAPIKeyAuth(user_id="alice")
first = await resource.list_user_resources(user_api_key_dict=user, limit=2)
second = await resource.list_user_resources(
user_api_key_dict=user, limit=2, after=first["last_id"]
)
assert first["has_more"] is True
assert second["has_more"] is False
assert [item["id"] for item in second["data"]] == [
rows[1].unified_resource_id,
rows[0].unified_resource_id,
]
@pytest.mark.asyncio
async def test_list_resources_rejects_unknown_after_cursor():
"""An ``after`` that does not resolve to a resource the caller can see is
a client error. Returning an empty page instead is indistinguishable from
the end of the list, so a stale cursor silently truncates the listing."""
from fastapi import HTTPException
resource = _make_paginating_resource([_make_row(0)])
with pytest.raises(HTTPException) as exc_info:
await resource.list_user_resources(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
limit=2,
after="does-not-exist-xyz",
)
assert exc_info.value.status_code == 400
assert "does-not-exist-xyz" in str(exc_info.value.detail)
resource.prisma_client.db.litellm_test_resource_table.find_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_list_resources_cursor_lookup_is_scoped_to_the_caller():
"""The cursor lookup must run against the rows the caller can list. A
Prisma cursor resolves by unique column regardless of the ``where``
filter, so an unscoped lookup would let one user anchor their page window
to another user's resource."""
from fastapi import HTTPException
resource = _make_paginating_resource([_make_row(0)])
with pytest.raises(HTTPException):
await resource.list_user_resources(
user_api_key_dict=UserAPIKeyAuth(user_id="alice"),
limit=2,
after="unified-resource-000",
additional_filters={"created_by": "bob"},
)
where = resource.prisma_client.db.litellm_test_resource_table.find_first.await_args.kwargs[
"where"
]
assert where["created_by"] == "bob"
assert where["unified_resource_id"] == "unified-resource-000"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"caller_team_id,expected",

View file

@ -5,12 +5,53 @@ Tests for managed-resource tenant isolation helpers.
import pytest
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
# ---------------------------------------------------------------------------
# build_list_page
# ---------------------------------------------------------------------------
def test_list_page_reads_ids_from_dict_items():
"""The vector-store listing builds plain dicts, so sourcing the cursor ids
through attribute access alone raised AttributeError on every non-empty
page."""
page = build_list_page(
[{"id": "resource-a"}, {"id": "resource-b"}], has_more=True
)
assert page["first_id"] == "resource-a"
assert page["last_id"] == "resource-b"
assert page["has_more"] is True
def test_list_page_reads_ids_from_object_items():
class _Item:
def __init__(self, id: str):
self.id = id
page = build_list_page([_Item("batch-a"), _Item("batch-b")])
assert page["first_id"] == "batch-a"
assert page["last_id"] == "batch-b"
assert page["has_more"] is False
def test_list_page_empty():
assert build_list_page([]) == {
"object": "list",
"data": [],
"first_id": None,
"last_id": None,
"has_more": False,
}
# ---------------------------------------------------------------------------
# build_owner_filter
# ---------------------------------------------------------------------------