From 4643685e7838ce0c8fd6c0ed1f12a2a8a34f9e86 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 17:55:26 -0800 Subject: [PATCH 1/4] [Fix] /key/aliases: Add pagination and search to prevent OOMs The /key/aliases endpoint previously fetched all key aliases from the database without limit, causing OOM crashes with large key sets. Added page, size, and search query parameters with database-level filtering to enable paginated and searchable key alias retrieval. Updated the response to include pagination metadata (total_count, current_page, total_pages, size) matching the /v2/model/info pattern. Co-Authored-By: Claude Haiku 4.5 --- .../key_management_endpoints.py | 59 ++++++++++++------- .../test_key_generate_prisma.py | 34 +++++++---- 2 files changed, 61 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c1165ab26d0..8a5145bf793 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4107,13 +4107,23 @@ async def list_keys( dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def key_aliases() -> Dict[str, List[str]]: +async def key_aliases( + page: int = Query(1, ge=1, description="Page number"), + size: int = Query(50, ge=1, le=100, description="Page size"), + search: Optional[str] = Query( + None, description="Search key aliases (case-insensitive partial match)" + ), +) -> Dict[str, Any]: """ - Lists all key aliases + Lists key aliases with pagination and optional search. Returns: { - "aliases": List[str] + "aliases": List[str], + "total_count": int, + "current_page": int, + "total_pages": int, + "size": int, } """ try: @@ -4125,36 +4135,43 @@ async def key_aliases() -> Dict[str, List[str]]: verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - where: Dict[str, Any] = {} + conditions: List[Dict[str, Any]] = [{"key_alias": {"not": None}}] try: - where.update(_get_condition_to_filter_out_ui_session_tokens()) + conditions.append(_get_condition_to_filter_out_ui_session_tokens()) except NameError: # Helper may not exist in some builds; ignore if missing pass + if search: + conditions.append( + {"key_alias": {"contains": search, "mode": "insensitive"}} + ) + where: Dict[str, Any] = {"AND": conditions} + total_count = await prisma_client.db.litellm_verificationtoken.count( + where=where, + ) rows = await prisma_client.db.litellm_verificationtoken.find_many( where=where, order=[{"key_alias": "asc"}], + skip=(page - 1) * size, + take=size, ) - seen = set() - aliases: List[str] = [] - for row in rows: - alias = getattr(row, "key_alias", None) - if alias is None and isinstance(row, dict): - alias = row.get("key_alias") + aliases: List[str] = [row.key_alias for row in rows if row.key_alias] # type: ignore[misc] - if not alias: - continue + total_pages = -(-total_count // size) if total_count > 0 else 0 + verbose_proxy_logger.debug( + f"key_aliases: page={page}, size={size}, search={search!r}, " + f"total_count={total_count}, total_pages={total_pages}" + ) - alias_str = str(alias).strip() - if alias_str and alias_str not in seen: - seen.add(alias_str) - aliases.append(alias_str) - - verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") - - return {"aliases": aliases} + return { + "aliases": aliases, + "total_count": total_count, + "current_page": page, + "total_pages": total_pages, + "size": size, + } except Exception as e: verbose_proxy_logger.exception(f"Error in key_aliases: {e}") diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index c3f68762810..ed528f21e0d 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3668,9 +3668,10 @@ async def test_list_keys(prisma_client): async def test_key_aliases(prisma_client): """ Test the key_aliases function: - - Returns a list + - Returns a paginated response - Includes alias from a newly created key - - Aliases are unique and sorted + - Aliases are sorted + - Pagination and search params work correctly """ import asyncio import uuid @@ -3682,10 +3683,16 @@ async def test_key_aliases(prisma_client): setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() - # Basic call - response = await key_aliases() + # Basic call - check pagination response shape + response = await key_aliases(page=1, size=50) assert "aliases" in response assert isinstance(response["aliases"], list) + assert "total_count" in response + assert "current_page" in response + assert "total_pages" in response + assert "size" in response + assert response["current_page"] == 1 + assert response["size"] == 50 # Create a new user (and key) with a unique alias unique_id = str(uuid.uuid4()) @@ -3704,17 +3711,22 @@ async def test_key_aliases(prisma_client): # Allow async DB writes to settle await asyncio.sleep(2) - # Call again and validate - response_after = await key_aliases() + # Call again and validate alias is present + response_after = await key_aliases(page=1, size=50) aliases = response_after["aliases"] - - # Contains the new alias assert test_alias in aliases - - # Unique & sorted (endpoint dedupes and orders ascending) - assert len(aliases) == len(set(aliases)) assert aliases == sorted(aliases) + # Search by partial alias + partial = test_alias[:10] + search_response = await key_aliases(page=1, size=50, search=partial) + assert test_alias in search_response["aliases"] + + # Search with no match + no_match_response = await key_aliases(page=1, size=50, search="__no_match_xyz__") + assert len(no_match_response["aliases"]) == 0 + assert no_match_response["total_count"] == 0 + @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.asyncio From 0b0809a3d571ed2b1f713dcf5340fe826d9ec76c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 21:03:15 -0800 Subject: [PATCH 2/4] optimize key_aliases to select only key_alias column and add unit tests Add select={"key_alias": True} to the find_many call so only the alias column is fetched from the database instead of full token rows. Add five unit tests in test_key_management_endpoints.py covering response shape, pagination skip/take computation, search filter injection, absence of contains filter when no search term is given, and the select-only-alias optimization. Co-Authored-By: Claude Haiku 4.5 --- .../key_management_endpoints.py | 1 + .../test_key_management_endpoints.py | 98 +++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8a5145bf793..8ac6f282719 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4152,6 +4152,7 @@ async def key_aliases( ) rows = await prisma_client.db.litellm_verificationtoken.find_many( where=where, + select={"key_alias": True}, order=[{"key_alias": "asc"}], skip=(page - 1) * size, take=size, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 05df3c2dcbb..8580c89fdaa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -45,6 +45,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( check_team_key_model_specific_limits, delete_verification_tokens, generate_key_helper_fn, + key_aliases, list_keys, prepare_key_update_data, reset_key_spend_fn, @@ -6211,3 +6212,100 @@ async def test_generate_key_helper_fn_agent_id(): assert key_data.get("agent_id") == "test-agent-456", ( f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}" ) + + +@pytest.mark.asyncio +async def test_key_aliases_response_shape(): + """Test that key_aliases returns the correct paginated response shape.""" + mock_row1 = MagicMock() + mock_row1.key_alias = "alias-alpha" + mock_row2 = MagicMock() + mock_row2.key_alias = "alias-beta" + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=2) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[mock_row1, mock_row2] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=1, size=50, search=None) + + assert result["aliases"] == ["alias-alpha", "alias-beta"] + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + assert result["size"] == 50 + + # Both count and find_many must use the same where clause + count_where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] + find_where = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] + assert count_where == find_where + + # Non-null alias filter must be present + assert json.dumps({"key_alias": {"not": None}}) in json.dumps(count_where) + + +@pytest.mark.asyncio +async def test_key_aliases_pagination_skip_take(): + """Test that skip and take are correctly computed from page and size.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=120) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=3, size=25, search=None) + + assert result["current_page"] == 3 + assert result["size"] == 25 + assert result["total_count"] == 120 + assert result["total_pages"] == 5 # ceil(120 / 25) + + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs["skip"] == 50 # (3 - 1) * 25 + assert find_many_kwargs["take"] == 25 + + +@pytest.mark.asyncio +async def test_key_aliases_search_filter(): + """Test that the search param adds a case-insensitive contains condition.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search="my-key") + + where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] + assert ( + json.dumps({"key_alias": {"contains": "my-key", "mode": "insensitive"}}) + in json.dumps(where) + ) + + +@pytest.mark.asyncio +async def test_key_aliases_no_search_omits_contains_filter(): + """Test that without a search term no contains condition is added.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search=None) + + where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] + assert "contains" not in json.dumps(where) + + +@pytest.mark.asyncio +async def test_key_aliases_select_only_key_alias(): + """Test that find_many is called with select={key_alias: True} to avoid fetching full rows.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search=None) + + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs.get("select") == {"key_alias": True} From 06e87eea877f866af9aa3d32839cac9ffd491595 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 21:40:56 -0800 Subject: [PATCH 3/4] remove unsupported select param from find_many call LiteLLM_VerificationTokenActions.find_many() does not support the select keyword argument. Remove it and drop the corresponding test. Co-Authored-By: Claude Haiku 4.5 --- .../management_endpoints/key_management_endpoints.py | 1 - .../test_key_management_endpoints.py | 12 ------------ 2 files changed, 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8ac6f282719..8a5145bf793 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4152,7 +4152,6 @@ async def key_aliases( ) rows = await prisma_client.db.litellm_verificationtoken.find_many( where=where, - select={"key_alias": True}, order=[{"key_alias": "asc"}], skip=(page - 1) * size, take=size, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8580c89fdaa..bdfc6734664 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6297,15 +6297,3 @@ async def test_key_aliases_no_search_omits_contains_filter(): assert "contains" not in json.dumps(where) -@pytest.mark.asyncio -async def test_key_aliases_select_only_key_alias(): - """Test that find_many is called with select={key_alias: True} to avoid fetching full rows.""" - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - await key_aliases(page=1, size=50, search=None) - - find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs - assert find_many_kwargs.get("select") == {"key_alias": True} From 386c148b8db369a5c2661a17dca1e4065fb661b8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 25 Feb 2026 21:48:53 -0800 Subject: [PATCH 4/4] use raw SQL in key_aliases to avoid loading full rows into memory Replace Prisma ORM count/find_many calls with two query_raw calls that only project the key_alias column. The Prisma client wrapper does not support SELECT projection via find_many, so raw SQL is used to keep memory usage proportional to the page size rather than total key count. Update tests to mock query_raw instead of count/find_many. Co-Authored-By: Claude Haiku 4.5 --- .../key_management_endpoints.py | 52 +++++++----- .../test_key_management_endpoints.py | 79 +++++++++++-------- 2 files changed, 76 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8a5145bf793..a0442a476f9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4135,29 +4135,41 @@ async def key_aliases( verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - conditions: List[Dict[str, Any]] = [{"key_alias": {"not": None}}] - try: - conditions.append(_get_condition_to_filter_out_ui_session_tokens()) - except NameError: - # Helper may not exist in some builds; ignore if missing - pass + # Build a parameterized WHERE clause to avoid loading full rows into + # memory. Raw SQL is used because the Prisma client wrapper does not + # support column-level SELECT projection on find_many. + # + # $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens). + query_params: List[Any] = [UI_SESSION_TOKEN_TEAM_ID] + where_parts = [ + "key_alias IS NOT NULL", + "key_alias != ''", + "(team_id IS NULL OR team_id != $1)", + ] if search: - conditions.append( - {"key_alias": {"contains": search, "mode": "insensitive"}} - ) - where: Dict[str, Any] = {"AND": conditions} + query_params.append(f"%{search}%") + where_parts.append(f"key_alias ILIKE ${len(query_params)}") - total_count = await prisma_client.db.litellm_verificationtoken.count( - where=where, - ) - rows = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, - order=[{"key_alias": "asc"}], - skip=(page - 1) * size, - take=size, - ) + where_sql = " AND ".join(where_parts) - aliases: List[str] = [row.key_alias for row in rows if row.key_alias] # type: ignore[misc] + count_sql = ( + f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' + ) + count_rows = await prisma_client.db.query_raw(count_sql, *query_params) + total_count = int(count_rows[0]["count"]) if count_rows else 0 + + aliases_params = query_params + [size, (page - 1) * size] + limit_idx = len(aliases_params) - 1 + offset_idx = len(aliases_params) + aliases_sql = ( + f"SELECT key_alias" + f' FROM "LiteLLM_VerificationToken"' + f" WHERE {where_sql}" + f" ORDER BY key_alias ASC" + f" LIMIT ${limit_idx} OFFSET ${offset_idx}" + ) + alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages = -(-total_count // size) if total_count > 0 else 0 verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index bdfc6734664..75325508b01 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6217,15 +6217,12 @@ async def test_generate_key_helper_fn_agent_id(): @pytest.mark.asyncio async def test_key_aliases_response_shape(): """Test that key_aliases returns the correct paginated response shape.""" - mock_row1 = MagicMock() - mock_row1.key_alias = "alias-alpha" - mock_row2 = MagicMock() - mock_row2.key_alias = "alias-beta" - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=2) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[mock_row1, mock_row2] + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 2}], + [{"key_alias": "alias-alpha"}, {"key_alias": "alias-beta"}], + ] ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -6237,21 +6234,23 @@ async def test_key_aliases_response_shape(): assert result["total_pages"] == 1 assert result["size"] == 50 - # Both count and find_many must use the same where clause - count_where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - find_where = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] - assert count_where == find_where - - # Non-null alias filter must be present - assert json.dumps({"key_alias": {"not": None}}) in json.dumps(count_where) + # Both SQL calls must filter out null/empty aliases + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + aliases_sql = mock_prisma_client.db.query_raw.call_args_list[1].args[0] + assert "key_alias IS NOT NULL" in count_sql + assert "key_alias IS NOT NULL" in aliases_sql @pytest.mark.asyncio async def test_key_aliases_pagination_skip_take(): - """Test that skip and take are correctly computed from page and size.""" + """Test that LIMIT and OFFSET are correctly derived from page and size.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=120) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 120}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): result = await key_aliases(page=3, size=25, search=None) @@ -6261,39 +6260,49 @@ async def test_key_aliases_pagination_skip_take(): assert result["total_count"] == 120 assert result["total_pages"] == 5 # ceil(120 / 25) - find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs - assert find_many_kwargs["skip"] == 50 # (3 - 1) * 25 - assert find_many_kwargs["take"] == 25 + # aliases query params: [UI_SESSION_TOKEN_TEAM_ID, size=25, offset=50] + aliases_call_args = mock_prisma_client.db.query_raw.call_args_list[1].args + assert aliases_call_args[-2] == 25 # LIMIT = size + assert aliases_call_args[-1] == 50 # OFFSET = (3 - 1) * 25 @pytest.mark.asyncio async def test_key_aliases_search_filter(): - """Test that the search param adds a case-insensitive contains condition.""" + """Test that the search param adds a case-insensitive ILIKE condition.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await key_aliases(page=1, size=50, search="my-key") - where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - assert ( - json.dumps({"key_alias": {"contains": "my-key", "mode": "insensitive"}}) - in json.dumps(where) - ) + count_call = mock_prisma_client.db.query_raw.call_args_list[0] + count_sql = count_call.args[0] + count_params = count_call.args[1:] + + assert "ILIKE" in count_sql + assert "%my-key%" in count_params @pytest.mark.asyncio -async def test_key_aliases_no_search_omits_contains_filter(): - """Test that without a search term no contains condition is added.""" +async def test_key_aliases_no_search_omits_ilike_filter(): + """Test that without a search term no ILIKE condition is added.""" mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await key_aliases(page=1, size=50, search=None) - where = mock_prisma_client.db.litellm_verificationtoken.count.call_args.kwargs["where"] - assert "contains" not in json.dumps(where) + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "ILIKE" not in count_sql