From f913d42abba2439db273523aa75fc002d74900c4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:44:48 -0700 Subject: [PATCH] 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. --- .../base_managed_resource.py | 36 ++-- .../base_llm/managed_resources/isolation.py | 17 +- .../base_llm/test_base_managed_resource.py | 178 ++++++++++++++++++ .../test_managed_resource_isolation.py | 41 ++++ 4 files changed, 252 insertions(+), 20 deletions(-) diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 146a6aa6ae0..3a121797871 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -542,35 +542,39 @@ async def list_user_resources( 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, ) + has_more = len(resources) > page_size + resource_objects: List[Any] = [] - for resource in resources: + for resource in resources[:page_size]: try: - # Stop once we have enough - if len(resource_objects) >= (limit or 20): - break - # Parse resource object resource_data = resource.resource_object if isinstance(resource_data, str): @@ -590,4 +594,4 @@ async def list_user_resources( ) 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) diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index fd1e24f3e1d..2d3c24d12fe 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -9,6 +9,7 @@ unscoped query. """ +from collections.abc import Mapping from typing import Any, Dict, List, Optional from litellm.proxy._types import ( @@ -17,15 +18,23 @@ ) +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, } diff --git a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py b/tests/test_litellm/llms/base_llm/test_base_managed_resource.py index a6c518616f4..8583697a310 100644 --- a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py +++ b/tests/test_litellm/llms/base_llm/test_base_managed_resource.py @@ -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", diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py index b5fcd9d8219..47669e61b4a 100644 --- a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py +++ b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py @@ -5,12 +5,53 @@ 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 # ---------------------------------------------------------------------------