From 9dcc21cd48aaee6650e2f8063692d5d1b78a1d41 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 16:18:25 -0700 Subject: [PATCH] refactor(mcp): batch the purge row deletion into one query The per-row delete_many loop becomes a single delete filtered to the enumerated OAuth users' (user_id IN, server_id) pairs; same rows deleted, same BYOK-sparing precision, same count-mismatch detection, one round-trip instead of N --- litellm/proxy/_experimental/mcp_server/db.py | 4 ++-- .../_experimental/mcp_server/test_db_credentials.py | 12 +++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index c6b7620b649..e4f8b0c331d 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1175,8 +1175,8 @@ async def purge_user_oauth_credentials_for_server( oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] if not oauth_rows: return 0 - deleted_count = sum( - [await repo.table.delete_many(where={"user_id": row.user_id, "server_id": server_id}) for row in oauth_rows] + deleted_count = await repo.table.delete_many( + where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}} ) if invalidate_token_cache is None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 0bdcffe1530..7269774442b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -171,7 +171,7 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _oauth_row("bob")]) - prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) invalidations = [] @@ -181,15 +181,17 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) assert purged == 2 - assert prisma.db.litellm_mcpusercredentials.delete_many.await_count == 2 + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( + where={"server_id": "srv-1", "user_id": {"in": ["alice", "bob"]}} + ) assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): """Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share - the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted, each by - its (user_id, server_id) pair, and only their users' token caches invalidated.""" + the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted (one + batched query filtered to their user_ids), and only their users' token caches invalidated.""" from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server prisma = MagicMock() @@ -205,7 +207,7 @@ async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): assert purged == 1 prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( - where={"user_id": "alice", "server_id": "srv-1"} + where={"server_id": "srv-1", "user_id": {"in": ["alice"]}} ) assert invalidations == [("alice", "srv-1")]