diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 711119b5ab5..08a8b1bc7b3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -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, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb330d00756..6c01575e271 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c4f5ceaba7a..997012dbc65 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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, diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index bf119c4fb2f..b6c2fd4c8e3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -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)