fix(mcp): acquire user env-var advisory lock via execute_raw

pg_advisory_xact_lock() returns void; routing it through query_raw made
the Prisma engine try to deserialize that column and raise RawQueryError,
500ing POST /v1/mcp/server/{id}/user-env-vars. execute_raw doesn't
deserialize a result set, so the lock is taken without the void error.

Add a regression test whose fake tx.query_raw raises the real void
RawQueryError while execute_raw succeeds, and switch the concurrency
test's fake to execute_raw so it can no longer pass against the bug.
This commit is contained in:
mateo-berri 2026-06-05 00:29:34 -07:00
parent 4929f50b43
commit 5177f8f1cb
2 changed files with 62 additions and 3 deletions

View file

@ -1315,7 +1315,7 @@ async def merge_user_env_vars(
signed=True,
)
async with prisma_client.db.tx() as tx:
await tx.query_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
await tx.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
row = await tx.litellm_mcpuserenvvars.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
)

View file

@ -866,12 +866,12 @@ def _transactional_env_vars_prisma(read_delay: float = 0.0):
self._held = None
return False
async def query_raw(self, query, *args):
async def execute_raw(self, query, *args):
lock_key = args[0]
lock = self._store.locks.setdefault(lock_key, asyncio.Lock())
await lock.acquire()
self._held = lock
return [{"pg_advisory_xact_lock": None}]
return 1
class _DB:
def __init__(self, store, delay):
@ -1058,6 +1058,65 @@ async def test_merge_user_env_vars_serializes_concurrent_writes(env_vars_salt_ke
assert stored == {"TOKEN_A": "a", "TOKEN_B": "b"}
@pytest.mark.asyncio
async def test_merge_user_env_vars_acquires_lock_without_deserializing_void(
env_vars_salt_key,
):
"""``pg_advisory_xact_lock`` returns ``void``; running it through ``query_raw``
makes Prisma try to deserialize that column and raises ``RawQueryError``. The
lock must be taken via ``execute_raw`` (no result-set deserialization) so the
merge still completes."""
from unittest.mock import MagicMock
from prisma.errors import RawQueryError
from litellm.proxy._experimental.mcp_server.db import merge_user_env_vars
class _Tx:
def __init__(self):
self.stored = None
self.litellm_mcpuserenvvars = self
async def __aenter__(self):
return self
async def __aexit__(self, *exc):
return False
async def query_raw(self, query, *args):
raise RawQueryError(
{
"user_facing_error": {
"error_code": "P2010",
"meta": {
"message": "Failed to deserialize column of type 'void'."
},
}
}
)
async def execute_raw(self, query, *args):
return 1
async def find_unique(self, where):
return None
async def upsert(self, where, data):
self.stored = data["create"]["values_b64"]
tx = _Tx()
prisma = MagicMock()
prisma.db.tx = MagicMock(return_value=tx)
values = {"CORP_TOKEN": "t0ken"}
merged = await merge_user_env_vars(
prisma, "alice", "srv-1", values, allowed_names=values.keys()
)
assert merged == values
assert tx.stored is not None
@pytest.mark.asyncio
async def test_delete_mcp_server_removes_orphaned_user_env_vars():
"""Deleting a server must also drop every user's per-user env var rows for