mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(mcp): resolve admin OAuth sessions from any worker via DB-backed drafts (#36844)
The Admin UI's Authorize & Fetch Token flow stored its pending server in a module-level dict, so /register, /authorize and /token only succeeded when every leg happened to land on the process that served /session. On a proxy with NUM_WORKERS greater than 1, or more than one replica, each click was an independent draw and failed with a bare 404, which reads as intermittent. Persist the pending server as a short-lived draft row instead, so any worker resolves it. The in-memory cache is kept as the fallback for proxies with no database configured, which keeps single-process deployments working as before. A session runs under a caller-supplied id only when that id names a server that really exists, which is the edit form re-authorizing a saved server. Anything else gets a fresh id, so two concurrent sessions can never share one draft and silently adopt each other's URL or client credentials. Drafts past their lifetime are swept on each write so abandoned sessions do not accumulate, and a lost create race adopts the winner rather than failing a caller whose session is ready. Drafts are excluded from listings and never enter the runtime registry. The exclusion keeps rows whose approval status is NULL, which both short spellings of the filter drop, silently hiding every server predating the approval workflow. Measured on a two-worker proxy against the live GitHub MCP server, 120 concurrent authorize calls per leg: staging 56/120 failures, this branch 0/120, staging again 65/120 as a positive control.
This commit is contained in:
Yassin Kortam
GitHub
parent
3615cccfef
commit
8841cbc10f
4 changed files with 565 additions and 17 deletions
|
|
@ -552,13 +552,20 @@ async def get_all_mcp_servers(
|
|||
) -> list[LiteLLM_MCPServerTable]:
|
||||
"""
|
||||
Returns mcp servers from the db, optionally filtered by approval_status.
|
||||
Pass approval_status=None to return all servers regardless of approval state.
|
||||
Pass approval_status=None to return every server except drafts, which back the admin OAuth
|
||||
session flow, are addressable only by their own server_id, and must never appear in a listing.
|
||||
NULL approval_status predates the approval workflow, so those rows are kept explicitly rather
|
||||
than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them.
|
||||
"""
|
||||
try:
|
||||
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = {}
|
||||
if approval_status is not None:
|
||||
where["approval_status"] = approval_status
|
||||
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where if where else {})
|
||||
where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = (
|
||||
{"approval_status": approval_status}
|
||||
if approval_status is not None
|
||||
# mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop
|
||||
# NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts
|
||||
else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]}
|
||||
)
|
||||
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
|
||||
|
||||
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
|
||||
for table in tables:
|
||||
|
|
@ -814,6 +821,96 @@ async def create_mcp_server(
|
|||
return new_mcp_server
|
||||
|
||||
|
||||
async def create_draft_mcp_server(
|
||||
prisma_client: PrismaClient,
|
||||
data: NewMCPServerRequest,
|
||||
touched_by: str,
|
||||
ttl_seconds: int,
|
||||
server_id: str | None = None,
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
"""
|
||||
Persist a short-lived draft row backing the admin OAuth "Authorize & Fetch Token" flow.
|
||||
|
||||
The draft lives in the database rather than in process memory so that the /register,
|
||||
/authorize and /token legs resolve it whichever worker or replica accepts each request.
|
||||
|
||||
Writing is strictly create-if-absent. Any existing row for the id is returned untouched, which
|
||||
covers both a live draft for this same session and a real server the edit form is
|
||||
re-authorizing against its own id, where writing a draft would collide on the primary key.
|
||||
Each click of Authorize mints a fresh id, so nothing is lost by never overwriting, and it is
|
||||
what makes concurrent callers sharing one id safe rather than mutually destructive.
|
||||
"""
|
||||
draft_id: Final = server_id or data.server_id or str(uuid.uuid4())
|
||||
await _prune_expired_draft_mcp_servers(prisma_client, ttl_seconds)
|
||||
|
||||
existing: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
|
||||
if existing is not None:
|
||||
# Already usable by every worker, whether it is a live draft for this same session or a
|
||||
# real server the edit form is re-authorizing. Either way there is nothing to write, and
|
||||
# not writing is what keeps concurrent callers for one server_id from racing each other.
|
||||
return LiteLLM_MCPServerTable.model_validate(existing.model_dump())
|
||||
|
||||
draft_payload: Final = data.model_copy(update={"server_id": draft_id, "approval_status": MCPApprovalStatus.draft})
|
||||
try:
|
||||
return await create_mcp_server(prisma_client, draft_payload, touched_by)
|
||||
except Exception:
|
||||
# Lost the create race: the read above and this create are two statements, not one. The
|
||||
# winner wrote a draft for this same session, so adopt it rather than failing a caller
|
||||
# whose session is in fact ready. Anything else still raises.
|
||||
raced: Final = await _db_find_mcp_server_row(prisma_client, draft_id)
|
||||
if raced is None or raced.approval_status != MCPApprovalStatus.draft:
|
||||
raise
|
||||
return LiteLLM_MCPServerTable.model_validate(raced.model_dump())
|
||||
|
||||
|
||||
async def _prune_expired_draft_mcp_servers(prisma_client: PrismaClient, ttl_seconds: int) -> None:
|
||||
"""Drop drafts already past ``ttl_seconds``, so abandoned OAuth sessions do not accumulate.
|
||||
|
||||
Runs on each draft write rather than on a schedule, mirroring the in-memory cache this
|
||||
replaces, which pruned on every store. Expired drafts are unreadable by then anyway, so the
|
||||
only thing at stake is row count, and the work is bounded by how often admins authorize.
|
||||
"""
|
||||
cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
|
||||
# Age is filtered here rather than in the query: the draft set is bounded by how many OAuth
|
||||
# authorizations are in flight, so it is a handful of rows even on a busy proxy.
|
||||
drafts: Final = await _db_find_mcp_server_rows(
|
||||
prisma_client,
|
||||
where={"approval_status": MCPApprovalStatus.draft},
|
||||
)
|
||||
for row in drafts:
|
||||
# A row without a timestamp has no age to judge, so leave it rather than guess it is stale.
|
||||
# Two workers sweeping the same row is harmless: prisma's delete returns None for a row
|
||||
# that is already gone rather than raising, so the loser of that race is a no-op.
|
||||
if row.updated_at is not None and row.updated_at < cutoff:
|
||||
await delete_mcp_server(prisma_client, row.server_id)
|
||||
|
||||
|
||||
async def get_draft_mcp_server(
|
||||
prisma_client: PrismaClient, server_id: str, ttl_seconds: int
|
||||
) -> LiteLLM_MCPServerTable | None:
|
||||
"""
|
||||
Return the draft row for ``server_id`` if it has not yet aged past ``ttl_seconds``, else None.
|
||||
|
||||
Age is enforced in the query rather than by a sweeper so an expired draft is unreadable the
|
||||
moment it lapses, regardless of which process last ran a cleanup.
|
||||
"""
|
||||
cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds))
|
||||
draft_rows: Final = await _db_find_mcp_server_rows(
|
||||
prisma_client,
|
||||
where={
|
||||
"server_id": server_id,
|
||||
"approval_status": MCPApprovalStatus.draft,
|
||||
"updated_at": {"gte": cutoff},
|
||||
},
|
||||
)
|
||||
if not draft_rows:
|
||||
return None
|
||||
|
||||
table: Final = LiteLLM_MCPServerTable.model_validate(draft_rows[0].model_dump())
|
||||
decrypt_global_env_var_values(table.env_vars)
|
||||
return table
|
||||
|
||||
|
||||
async def update_mcp_server(
|
||||
prisma_client: PrismaClient,
|
||||
data: UpdateMCPServerRequest,
|
||||
|
|
|
|||
|
|
@ -1283,6 +1283,9 @@ class MCPApprovalStatus(str, enum.Enum):
|
|||
pending_review = "pending_review"
|
||||
active = "active"
|
||||
rejected = "rejected"
|
||||
# Short-lived row backing the admin OAuth "Authorize & Fetch Token" flow. Never served: the
|
||||
# registry loader and every listing exclude it, so it is reachable only by its own server_id.
|
||||
draft = "draft"
|
||||
|
||||
|
||||
from litellm.models.mcp_server import ( # noqa: E402
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ import os
|
|||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Final, Literal
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
|
|
@ -92,6 +92,9 @@ def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str)
|
|||
|
||||
DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
try:
|
||||
importlib.import_module("mcp")
|
||||
except ImportError as e:
|
||||
|
|
@ -115,11 +118,13 @@ if MCP_AVAILABLE:
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
approve_mcp_server,
|
||||
create_draft_mcp_server,
|
||||
create_mcp_server,
|
||||
delete_mcp_server,
|
||||
delete_user_credential,
|
||||
delete_user_env_vars,
|
||||
get_all_mcp_servers_for_user,
|
||||
get_draft_mcp_server,
|
||||
get_mcp_server,
|
||||
get_mcp_servers,
|
||||
get_mcp_submissions,
|
||||
|
|
@ -467,19 +472,68 @@ if MCP_AVAILABLE:
|
|||
verbose_proxy_logger.debug("Invalid temporary MCP server payload in Redis cache: %s", e)
|
||||
return None
|
||||
|
||||
def _get_prisma_client_or_none() -> "PrismaClient | None":
|
||||
"""Non-throwing counterpart to ``get_prisma_client_or_throw`` for paths that degrade
|
||||
gracefully: a proxy configured without a database keeps the in-memory OAuth session."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
return prisma_client
|
||||
|
||||
async def _persist_draft_mcp_server(
|
||||
payload: NewMCPServerRequest,
|
||||
server_id: str,
|
||||
created_by: str,
|
||||
) -> None:
|
||||
"""Write the draft row that makes the OAuth session resolvable from any worker.
|
||||
|
||||
A failure here is raised, not swallowed: without the shared row the flow degrades to
|
||||
the per-process cache and fails intermittently, which is the defect being fixed.
|
||||
"""
|
||||
prisma_client: Final = _get_prisma_client_or_none()
|
||||
if prisma_client is None:
|
||||
return
|
||||
await create_draft_mcp_server(
|
||||
prisma_client,
|
||||
payload,
|
||||
created_by,
|
||||
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
|
||||
server_id=server_id,
|
||||
)
|
||||
|
||||
async def _get_draft_mcp_server_as_mcp_server(server_id: str) -> MCPServer | None:
|
||||
"""Resolve a database-backed draft, which is the only lookup that works across workers."""
|
||||
prisma_client: Final = _get_prisma_client_or_none()
|
||||
if prisma_client is None:
|
||||
return None
|
||||
draft: Final = await get_draft_mcp_server(
|
||||
prisma_client, server_id, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS
|
||||
)
|
||||
if draft is None:
|
||||
return None
|
||||
return await global_mcp_server_manager.build_mcp_server_from_table(draft)
|
||||
|
||||
async def get_cached_temporary_mcp_server(
|
||||
server_id: str,
|
||||
) -> MCPServer | None:
|
||||
_prune_expired_temporary_mcp_servers()
|
||||
entry: Final = _temporary_mcp_servers.get(server_id)
|
||||
if entry is None:
|
||||
redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id)
|
||||
if redis_server is None:
|
||||
return None
|
||||
# Intentionally avoid repopulating local cache from Redis to prevent
|
||||
# extending effective lifetime beyond the remaining Redis TTL.
|
||||
return redis_server
|
||||
return entry.server
|
||||
if entry is not None:
|
||||
return entry.server
|
||||
|
||||
# A miss here means either an expired session or, on a multi-worker or multi-replica
|
||||
# proxy, that a different process served /session. The draft row is shared, so it
|
||||
# resolves the second case; the in-memory hit above still serves single-process
|
||||
# deployments with no database configured.
|
||||
draft_server: Final = await _get_draft_mcp_server_as_mcp_server(server_id)
|
||||
if draft_server is not None:
|
||||
return draft_server
|
||||
|
||||
redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id)
|
||||
if redis_server is None:
|
||||
return None
|
||||
# Intentionally avoid repopulating local cache from Redis to prevent
|
||||
# extending effective lifetime beyond the remaining Redis TTL.
|
||||
return redis_server
|
||||
|
||||
def _redact_mcp_credentials(
|
||||
mcp_server: LiteLLM_MCPServerTable,
|
||||
|
|
@ -709,12 +763,36 @@ if MCP_AVAILABLE:
|
|||
payload_dict["credentials"] = inherited_credentials
|
||||
return NewMCPServerRequest.model_validate(payload_dict)
|
||||
|
||||
async def _resolve_session_server_id(payload: NewMCPServerRequest) -> str:
|
||||
"""Decide the id an OAuth session runs under.
|
||||
|
||||
A caller-supplied id is honoured only when it names a server that really exists, which is
|
||||
the edit form re-authorizing a saved server against its own id. Anything else gets a fresh
|
||||
id, so two concurrent sessions can never land on one id and silently adopt each other's
|
||||
URL or client credentials. Without a database there is nothing shared to collide over, so
|
||||
the supplied id is kept and behaviour is unchanged.
|
||||
"""
|
||||
supplied: Final = payload.server_id
|
||||
if not supplied:
|
||||
return str(uuid.uuid4())
|
||||
if global_mcp_server_manager.get_mcp_server_by_id(supplied) is not None:
|
||||
return supplied
|
||||
prisma_client: Final = _get_prisma_client_or_none()
|
||||
if prisma_client is None:
|
||||
return supplied
|
||||
# A draft is another session's row, not a saved server, so re-supplying an id this
|
||||
# endpoint previously handed back must not let a later session adopt its configuration.
|
||||
existing: Final = await get_mcp_server(prisma_client, supplied)
|
||||
if existing is None or existing.approval_status == MCPApprovalStatus.draft:
|
||||
return str(uuid.uuid4())
|
||||
return supplied
|
||||
|
||||
def _build_temporary_mcp_server_record(
|
||||
payload: NewMCPServerRequest,
|
||||
created_by: str | None,
|
||||
server_id: str,
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
now: Final = datetime.utcnow()
|
||||
server_id: Final = payload.server_id or str(uuid.uuid4())
|
||||
server_name: Final = payload.server_name or payload.alias or server_id
|
||||
return LiteLLM_MCPServerTable(
|
||||
server_id=server_id,
|
||||
|
|
@ -1544,6 +1622,7 @@ if MCP_AVAILABLE:
|
|||
temp_record: Final = _build_temporary_mcp_server_record(
|
||||
payload_with_credentials,
|
||||
created_by,
|
||||
await _resolve_session_server_id(payload_with_credentials),
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -1555,6 +1634,11 @@ if MCP_AVAILABLE:
|
|||
temporary_server,
|
||||
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
|
||||
)
|
||||
await _persist_draft_mcp_server(
|
||||
payload_with_credentials,
|
||||
temp_record.server_id,
|
||||
created_by,
|
||||
)
|
||||
await _cache_temporary_mcp_server_in_redis(
|
||||
temporary_server,
|
||||
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
|
||||
|
|
|
|||
|
|
@ -1697,6 +1697,362 @@ class TestTemporaryMCPSessionEndpoints:
|
|||
assert result is None
|
||||
assert "expired" not in cache
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_resolves_draft_written_by_another_worker(self):
|
||||
"""Regression: the OAuth session must resolve on a worker that did not serve /session.
|
||||
|
||||
`_temporary_mcp_servers` is per-process, so on a multi-worker or multi-replica proxy the
|
||||
/authorize and /token legs land on a process whose dict is empty and 404. An empty dict
|
||||
here IS that other worker. Before the DB-backed draft this returned None.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
get_cached_temporary_mcp_server,
|
||||
)
|
||||
|
||||
draft_row = generate_mock_mcp_server_db_record(server_id="drafted-elsewhere")
|
||||
rebuilt_server = generate_mock_mcp_server_config_record(server_id="drafted-elsewhere")
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.build_mcp_server_from_table = AsyncMock(return_value=rebuilt_server)
|
||||
get_draft = AsyncMock(return_value=draft_row)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
|
||||
{},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server",
|
||||
get_draft,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
):
|
||||
result = await get_cached_temporary_mcp_server("drafted-elsewhere")
|
||||
|
||||
assert result is rebuilt_server
|
||||
# The shared row, not the empty per-process dict, is what answered.
|
||||
assert get_draft.await_count == 1
|
||||
assert get_draft.await_args.args[1] == "drafted-elsewhere"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_still_works_without_a_database(self):
|
||||
"""A proxy configured with no database keeps the in-memory session, rather than 404ing.
|
||||
|
||||
Pins the deliberate divergence from a DB-only design: single-process deployments with no
|
||||
DATABASE_URL must keep working exactly as before.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_TemporaryMCPServerEntry,
|
||||
get_cached_temporary_mcp_server,
|
||||
)
|
||||
|
||||
server = generate_mock_mcp_server_config_record(server_id="no-db")
|
||||
entry = _TemporaryMCPServerEntry(
|
||||
server=server,
|
||||
expires_at=datetime.utcnow() + timedelta(seconds=300),
|
||||
)
|
||||
get_draft = AsyncMock(return_value=None)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
|
||||
{"no-db": entry},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server",
|
||||
get_draft,
|
||||
),
|
||||
):
|
||||
result = await get_cached_temporary_mcp_server("no-db")
|
||||
|
||||
assert result is server
|
||||
# No database means no draft lookup is even attempted.
|
||||
assert get_draft.await_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_draft_mcp_server_never_overwrites_a_real_server(self):
|
||||
"""The edit form authorizes against a saved server's own id, so a draft write would
|
||||
collide on the primary key. That row is already visible to every worker, so it is
|
||||
returned untouched and no draft is created."""
|
||||
from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server
|
||||
|
||||
real_row = generate_mock_mcp_server_db_record(server_id="already-saved")
|
||||
real_row.approval_status = "active"
|
||||
create_call = AsyncMock()
|
||||
delete_call = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row",
|
||||
AsyncMock(return_value=real_row),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.create_mcp_server", create_call),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call),
|
||||
):
|
||||
result = await create_draft_mcp_server(
|
||||
MagicMock(),
|
||||
NewMCPServerRequest(server_id="already-saved", url="https://x.example.com/mcp"),
|
||||
"tester",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
|
||||
assert result.server_id == "already-saved"
|
||||
assert create_call.await_count == 0
|
||||
assert delete_call.await_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_draft_mcp_server_adopts_the_winner_when_it_loses_a_create_race(self):
|
||||
"""Regression: the read, delete and create are three statements, not one.
|
||||
|
||||
Two concurrent sessions for the same server_id raced and 13 of 20 returned 500 against a
|
||||
live two-worker proxy. The loser's session is in fact ready, because the winner wrote a
|
||||
draft for it, so it adopts that row instead of failing the caller.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server
|
||||
|
||||
winner_draft = generate_mock_mcp_server_db_record(server_id="raced")
|
||||
winner_draft.approval_status = "draft"
|
||||
# First lookup: nothing yet. After the losing create blows up: the winner's row.
|
||||
lookups = AsyncMock(side_effect=[None, winner_draft])
|
||||
|
||||
with (
|
||||
patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", lookups),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.create_mcp_server",
|
||||
AsyncMock(side_effect=Exception("duplicate key value violates unique constraint")),
|
||||
),
|
||||
):
|
||||
result = await create_draft_mcp_server(
|
||||
MagicMock(),
|
||||
NewMCPServerRequest(server_id="raced", url="https://x.example.com/mcp"),
|
||||
"tester",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
|
||||
assert result.server_id == "raced"
|
||||
assert lookups.await_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_draft_mcp_server_reraises_when_the_create_failure_was_not_a_race(self):
|
||||
"""A genuine database error must not be swallowed by the race-adoption path."""
|
||||
from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row",
|
||||
AsyncMock(side_effect=[None, None]),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.create_mcp_server",
|
||||
AsyncMock(side_effect=Exception("connection refused")),
|
||||
),
|
||||
pytest.raises(Exception, match="connection refused"),
|
||||
):
|
||||
await create_draft_mcp_server(
|
||||
MagicMock(),
|
||||
NewMCPServerRequest(server_id="broken", url="https://x.example.com/mcp"),
|
||||
"tester",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_draft_mcp_server_prunes_drafts_past_their_lifetime(self):
|
||||
"""Regression: abandoned OAuth sessions accumulated forever. Verified against a live
|
||||
proxy, where 12 drafts aged past the lifetime were still present and a 13th was added."""
|
||||
from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server
|
||||
|
||||
from datetime import timezone
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
stale_one = generate_mock_mcp_server_db_record(server_id="stale-1")
|
||||
stale_one.updated_at = now - timedelta(hours=1)
|
||||
stale_two = generate_mock_mcp_server_db_record(server_id="stale-2")
|
||||
stale_two.updated_at = now - timedelta(hours=1)
|
||||
# A draft still inside its lifetime must survive the sweep.
|
||||
fresh_draft = generate_mock_mcp_server_db_record(server_id="still-live")
|
||||
fresh_draft.updated_at = now
|
||||
find_rows = AsyncMock(return_value=[stale_one, stale_two, fresh_draft])
|
||||
delete_call = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", find_rows),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row",
|
||||
AsyncMock(return_value=None),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.db.create_mcp_server",
|
||||
AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id="fresh")),
|
||||
),
|
||||
):
|
||||
await create_draft_mcp_server(
|
||||
MagicMock(),
|
||||
NewMCPServerRequest(server_id="fresh", url="https://x.example.com/mcp"),
|
||||
"tester",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
|
||||
# Only drafts are considered, only the expired ones are removed, and the live one stays.
|
||||
assert find_rows.await_args.kwargs["where"]["approval_status"] == "draft"
|
||||
assert sorted(c.args[1] for c in delete_call.await_args_list) == ["stale-1", "stale-2"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_mcp_servers_hides_drafts_without_hiding_legacy_null_rows(self):
|
||||
"""Drafts are addressable only by their own id and must never appear in a listing, but a
|
||||
bare inequality would also drop pre-approval-workflow rows, since SQL evaluates
|
||||
NULL != 'draft' as NULL."""
|
||||
from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers
|
||||
|
||||
find_rows = AsyncMock(return_value=[])
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows",
|
||||
find_rows,
|
||||
):
|
||||
await get_all_mcp_servers(MagicMock())
|
||||
|
||||
where = find_rows.await_args.args[1]
|
||||
assert where == {"OR": [{"approval_status": None}, {"approval_status": {"not": "draft"}}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_session_server_id_refuses_an_unknown_caller_supplied_id(self):
|
||||
"""Regression: two concurrent sessions must never land on one id.
|
||||
|
||||
Honouring an arbitrary caller-supplied id lets a second session adopt the first's draft and
|
||||
run OAuth against its URL and client credentials, silently. An id naming no real server is
|
||||
therefore replaced with a fresh one.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_resolve_session_server_id,
|
||||
)
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.get_mcp_server_by_id.return_value = None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
|
||||
AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
resolved = await _resolve_session_server_id(
|
||||
NewMCPServerRequest(server_id="someone-elses-id", url="https://x.example.com/mcp")
|
||||
)
|
||||
|
||||
assert resolved != "someone-elses-id"
|
||||
uuid.UUID(resolved)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_session_server_id_refuses_an_id_that_names_another_sessions_draft(self):
|
||||
"""A draft row is another session's, not a saved server. Replaying an id this endpoint
|
||||
previously returned must not let a later session inherit the earlier one's config."""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_resolve_session_server_id,
|
||||
)
|
||||
|
||||
someone_elses_draft = generate_mock_mcp_server_db_record(server_id="earlier-session")
|
||||
someone_elses_draft.approval_status = "draft"
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.get_mcp_server_by_id.return_value = None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
|
||||
return_value=MagicMock(),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
|
||||
AsyncMock(return_value=someone_elses_draft),
|
||||
),
|
||||
):
|
||||
resolved = await _resolve_session_server_id(
|
||||
NewMCPServerRequest(server_id="earlier-session", url="https://x.example.com/mcp")
|
||||
)
|
||||
|
||||
assert resolved != "earlier-session"
|
||||
uuid.UUID(resolved)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_session_server_id_keeps_a_real_servers_id_for_the_edit_flow(self):
|
||||
"""The edit form re-authorizes a saved server against its own id, which must be preserved
|
||||
or the flow would authorize a throwaway id instead of the server being edited."""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_resolve_session_server_id,
|
||||
)
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.get_mcp_server_by_id.return_value = generate_mock_mcp_server_config_record(server_id="saved")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
):
|
||||
resolved = await _resolve_session_server_id(
|
||||
NewMCPServerRequest(server_id="saved", url="https://x.example.com/mcp")
|
||||
)
|
||||
|
||||
assert resolved == "saved"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_session_server_id_keeps_the_supplied_id_without_a_database(self):
|
||||
"""No database means nothing shared to collide over, so behaviour stays as it is today."""
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
_resolve_session_server_id,
|
||||
)
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.get_mcp_server_by_id.return_value = None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
|
||||
mock_manager,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
resolved = await _resolve_session_server_id(
|
||||
NewMCPServerRequest(server_id="no-db-id", url="https://x.example.com/mcp")
|
||||
)
|
||||
|
||||
assert resolved == "no-db-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_cached_temporary_mcp_server_or_404(self):
|
||||
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
|
||||
|
|
@ -5705,7 +6061,9 @@ def _edit_endpoint_patches(old_record, update_mock):
|
|||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
|
||||
AsyncMock(side_effect=old_record) if isinstance(old_record, Exception) else AsyncMock(return_value=old_record),
|
||||
AsyncMock(side_effect=old_record)
|
||||
if isinstance(old_record, Exception)
|
||||
else AsyncMock(return_value=old_record),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server",
|
||||
|
|
@ -6114,7 +6472,13 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed():
|
|||
|
||||
registry_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)),
|
||||
"..", "..", "..", "..", "litellm", "proxy", "openapi_registry.json",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"litellm",
|
||||
"proxy",
|
||||
"openapi_registry.json",
|
||||
)
|
||||
with open(registry_path) as f:
|
||||
registry = json.load(f)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue