diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py new file mode 100644 index 00000000000..02cec2475e2 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py @@ -0,0 +1,155 @@ +"""Startup backfill for oauth2 MCP server rows persisted before oauth2_flow was written. + +Rows created before the write-side stamps (DCR persist, UI create, REST create) carry a +null ``oauth2_flow`` and rely on read-time field-shape inference, which cannot tell a +DCR-registered interactive server from an M2M server unless endpoint discovery succeeds +first. This backfill classifies each null row once, at rest, using signals inference +never had, and persists the result so the read path never has to infer again. + +Signal order, strongest first: + +1. Per-user OAuth token rows exist for the server: only the interactive flow mints + per-user tokens, so this is definitive and immune to the discovery trap. BYOK API + keys share the same table (``LiteLLM_MCPUserCredentials``), so only rows whose + payload decodes as a ``type: oauth2`` token count as proof; bare keys and + undecodable rows prove nothing about the flow. +2. ``authorization_url`` persisted: interactive needs a user-facing authorization + endpoint; M2M (RFC 6749 section 4.4) never has one. +3. ``registration_url`` persisted: dynamic client registration (RFC 7591) exists to mint + clients for the interactive flow; M2M servers are configured with static credentials. +4. ``token_url`` plus decryptable ``client_id`` and ``client_secret``: ambiguous, left + unstamped. The shape is shared by M2M servers and DCR-registered interactive servers + whose authorization endpoint lives only in discovery (registered but never signed + in), so stamping client_credentials here could permanently route per-user traffic + through the proxy's stored client credential. The row keeps working through the + request-time backstop and a warning names it with the one-line fix (set oauth2_flow + via the dashboard or ``PUT /v1/mcp/server``); a completed interactive sign-in also + heals it via rule 1 at the next boot. +5. Anything else is interactive: matching how ``needs_user_oauth_token`` treats a null + flow, so the stamp never changes runtime routing for rows no rule recognizes. + +The backfill never stamps client_credentials: M2M is asserted by a human (config +requires it, the API accepts it, the dashboard sets it), mirroring the config-level +validation error. Runs before the first registry load on every boot and is idempotent: +a healed fleet has no null rows and the backfill exits after one query. +""" + +import json +from collections import Counter +from typing import Any, Literal, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials +from litellm.proxy.utils import PrismaClient +from litellm.types.mcp import MCPCredentials + +OAuth2Flow = Literal["client_credentials", "authorization_code"] +BackfillRule = Literal[ + "per_user_tokens", + "authorization_url", + "registration_url", + "ambiguous_m2m_shape", + "interactive_default", +] + +_BACKFILL_AUDIT_ACTOR = "oauth2_flow_backfill" + + +def _decrypted_credentials(raw_credentials: Any) -> Optional[MCPCredentials]: + if raw_credentials is None: + return None + if isinstance(raw_credentials, str): + try: + parsed = json.loads(raw_credentials) + except (ValueError, TypeError): + return None + else: + parsed = raw_credentials + if not isinstance(parsed, dict): + return None + return decrypt_credentials(credentials=dict(parsed)) + + +def classify_null_flow_row( + *, + has_per_user_tokens: bool, + authorization_url: Optional[str], + registration_url: Optional[str], + token_url: Optional[str], + credentials: Optional[MCPCredentials], +) -> tuple[Optional[OAuth2Flow], BackfillRule]: + if has_per_user_tokens: + return "authorization_code", "per_user_tokens" + if authorization_url: + return "authorization_code", "authorization_url" + if registration_url: + return "authorization_code", "registration_url" + if token_url and credentials and credentials.get("client_id") and credentials.get("client_secret"): + return None, "ambiguous_m2m_shape" + return "authorization_code", "interactive_default" + + +async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]: + """Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable + ones, warn on the ambiguous ones, and return counts per rule.""" + null_rows: list[Any] = await prisma_client.db.litellm_mcpservertable.find_many( + where={"auth_type": "oauth2", "oauth2_flow": None}, + ) + if not null_rows: + return {} + + server_ids = [row.server_id for row in null_rows] + token_rows: list[Any] = await prisma_client.db.litellm_mcpusercredentials.find_many( + where={"server_id": {"in": server_ids}}, + ) + server_ids_with_oauth_tokens: set[str] = { + token_row.server_id for token_row in token_rows if _decode_oauth_payload(token_row.credential_b64) is not None + } + + classified = tuple( + ( + row, + classify_null_flow_row( + has_per_user_tokens=row.server_id in server_ids_with_oauth_tokens, + authorization_url=row.authorization_url, + registration_url=row.registration_url, + token_url=row.token_url, + credentials=_decrypted_credentials(row.credentials), + ), + ) + for row in null_rows + ) + + for row, (flow, rule) in classified: + if flow is None: + verbose_proxy_logger.warning( + "oauth2_flow backfill: server_id=%s is ambiguous (client credentials + token_url, " + "no interactive signal); left unstamped. Set oauth2_flow explicitly via the " + "dashboard or PUT /v1/mcp/server: client_credentials if this server is M2M, or " + "complete an interactive sign-in and it will be stamped authorization_code at the " + "next boot.", + row.server_id, + ) + else: + verbose_proxy_logger.info( + "oauth2_flow backfill: server_id=%s stamped %s (rule=%s)", + row.server_id, + flow, + rule, + ) + + stamped_flows = {flow for _, (flow, _) in classified if flow is not None} + for stamped_flow in stamped_flows: + server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow] + await prisma_client.db.litellm_mcpservertable.update_many( + where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None}, + data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR}, + ) + + counts: dict[BackfillRule, int] = dict(Counter(rule for _, (_, rule) in classified)) + verbose_proxy_logger.info( + "oauth2_flow backfill: processed %d oauth2 server row(s): %s", + len(null_rows), + counts, + ) + return counts diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1474c15e778..c619133cebd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6366,6 +6366,17 @@ class ProxyConfig: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy._experimental.mcp_server.oauth2_flow_backfill import ( + backfill_null_oauth2_flows, + ) + + try: + if prisma_client is not None: + await backfill_null_oauth2_flows(prisma_client) + except Exception as e: # noqa: BLE001 + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db backfill - {}".format(str(e)) + ) try: await global_mcp_server_manager.reload_servers_from_database() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py new file mode 100644 index 00000000000..c1239c228aa --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py @@ -0,0 +1,234 @@ +""" +Tests for the startup oauth2_flow backfill. + +Legacy oauth2 rows with a null oauth2_flow are classified once, at rest, using +signals read-time inference never had (per-user token rows first), and the +result is persisted so the read path never infers again. The signal order is +the spec, and so is the refusal to stamp client_credentials: the M2M credential +shape is shared by DCR-registered interactive servers whose authorization +endpoint lives only in discovery, so ambiguous rows are left unstamped for a +human to assert rather than being permanently mislabeled M2M. +""" + +import base64 +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._experimental.mcp_server.oauth2_flow_backfill import ( + backfill_null_oauth2_flows, + classify_null_flow_row, +) + + +def test_classify_per_user_tokens_beat_m2m_shape(): + """The DCR trap row: creds + token_url, no authorization_url, but a user has + signed in. Tokens are definitive; the M2M shape must not win.""" + flow, rule = classify_null_flow_row( + has_per_user_tokens=True, + authorization_url=None, + registration_url=None, + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + assert flow == "authorization_code" + assert rule == "per_user_tokens" + + +def test_classify_authorization_url_beats_m2m_shape(): + flow, rule = classify_null_flow_row( + has_per_user_tokens=False, + authorization_url="https://idp.example.com/authorize", + registration_url=None, + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + assert flow == "authorization_code" + assert rule == "authorization_url" + + +def test_classify_registration_url_beats_m2m_shape(): + """A registration endpoint means DCR, and DCR exists to mint interactive + clients; an abandoned-DCR row (no sign-in yet) must not be stamped M2M.""" + flow, rule = classify_null_flow_row( + has_per_user_tokens=False, + authorization_url=None, + registration_url="https://idp.example.com/register", + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + assert flow == "authorization_code" + assert rule == "registration_url" + + +def test_classify_m2m_shape_is_ambiguous_and_unstamped(): + """The M2M shape alone must never stamp client_credentials: a DCR-registered + interactive server that nobody signed into yet has the identical shape, and a + wrong M2M stamp would permanently route its per-user traffic through the + proxy's stored client credential.""" + flow, rule = classify_null_flow_row( + has_per_user_tokens=False, + authorization_url=None, + registration_url=None, + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + assert flow is None + assert rule == "ambiguous_m2m_shape" + + +def test_classify_partial_credentials_default_interactive(): + """token_url without a full credential pair is not the M2M shape.""" + flow, rule = classify_null_flow_row( + has_per_user_tokens=False, + authorization_url=None, + registration_url=None, + token_url="https://idp.example.com/token", + credentials={"client_id": "cid"}, + ) + assert flow == "authorization_code" + assert rule == "interactive_default" + + +def test_classify_bare_row_default_interactive(): + flow, rule = classify_null_flow_row( + has_per_user_tokens=False, + authorization_url=None, + registration_url=None, + token_url=None, + credentials=None, + ) + assert flow == "authorization_code" + assert rule == "interactive_default" + + +def _row(server_id, *, authorization_url=None, registration_url=None, token_url=None, credentials=None): + return SimpleNamespace( + server_id=server_id, + authorization_url=authorization_url, + registration_url=registration_url, + token_url=token_url, + credentials=credentials, + ) + + +def _oauth_token_row(server_id): + payload = json.dumps({"type": "oauth2", "access_token": "tok", "connected_at": "2026-07-01T00:00:00Z"}) + return SimpleNamespace( + server_id=server_id, + user_id="u1", + credential_b64=base64.urlsafe_b64encode(payload.encode()).decode(), + ) + + +def _byok_key_row(server_id): + return SimpleNamespace( + server_id=server_id, + user_id="u1", + credential_b64=base64.urlsafe_b64encode(b"sk-user-supplied-upstream-key").decode(), + ) + + +def _mock_prisma(null_rows, token_rows): + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=null_rows) + mock_prisma.db.litellm_mcpservertable.update_many = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=token_rows) + return mock_prisma + + +@pytest.mark.asyncio +async def test_backfill_only_targets_null_flow_oauth2_rows(): + """The where clause is the guard that explicit and non-oauth2 rows are never touched.""" + mock_prisma = _mock_prisma([], []) + + counts = await backfill_null_oauth2_flows(mock_prisma) + + assert counts == {} + mock_prisma.db.litellm_mcpservertable.find_many.assert_awaited_once_with( + where={"auth_type": "oauth2", "oauth2_flow": None}, + ) + mock_prisma.db.litellm_mcpusercredentials.find_many.assert_not_awaited() + mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_backfill_stamps_rows_and_reports_rule_counts(): + dcr_trap_row = _row( + "signed_in_dcr", + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + m2m_row = _row( + "legacy_m2m", + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + interactive_row = _row("legacy_interactive", authorization_url="https://idp.example.com/authorize") + + mock_prisma = _mock_prisma( + [dcr_trap_row, m2m_row, interactive_row], + [_oauth_token_row("signed_in_dcr")], + ) + + counts = await backfill_null_oauth2_flows(mock_prisma) + + assert counts == {"per_user_tokens": 1, "ambiguous_m2m_shape": 1, "authorization_url": 1} + + mock_prisma.db.litellm_mcpservertable.update_many.assert_awaited_once() + call = mock_prisma.db.litellm_mcpservertable.update_many.await_args + assert sorted(call.kwargs["where"]["server_id"]["in"]) == ["legacy_interactive", "signed_in_dcr"] + assert "oauth2_flow" in call.kwargs["where"] and call.kwargs["where"]["oauth2_flow"] is None + assert call.kwargs["data"] == {"oauth2_flow": "authorization_code", "updated_by": "oauth2_flow_backfill"} + + +@pytest.mark.asyncio +async def test_backfill_handles_json_string_credentials(): + """JSON-string credential blobs must decode: the M2M shape is recognized (and + therefore deliberately left unstamped) rather than misread as credential-less.""" + m2m_row = _row( + "json_creds_m2m", + token_url="https://idp.example.com/token", + credentials='{"client_id": "cid", "client_secret": "csecret"}', + ) + mock_prisma = _mock_prisma([m2m_row], []) + + counts = await backfill_null_oauth2_flows(mock_prisma) + + assert counts == {"ambiguous_m2m_shape": 1} + mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_backfill_treats_undecodable_credentials_as_absent(): + row = _row( + "corrupt_creds", + token_url="https://idp.example.com/token", + credentials="not-json", + ) + mock_prisma = _mock_prisma([row], []) + + counts = await backfill_null_oauth2_flows(mock_prisma) + + assert counts == {"interactive_default": 1} + + +@pytest.mark.asyncio +async def test_backfill_byok_key_rows_are_not_sign_in_proof(): + """BYOK API keys live in the same table as per-user OAuth tokens; a bare key row + must not satisfy the per_user_tokens rule, or a BYOK-flavored M2M-shaped server + would be permanently stamped authorization_code. Only rows whose payload decodes + as a type oauth2 token count.""" + byok_shaped_row = _row( + "byok_m2m_shape", + token_url="https://idp.example.com/token", + credentials={"client_id": "cid", "client_secret": "csecret"}, + ) + mock_prisma = _mock_prisma([byok_shaped_row], [_byok_key_row("byok_m2m_shape")]) + + counts = await backfill_null_oauth2_flows(mock_prisma) + + assert counts == {"ambiguous_m2m_shape": 1} + mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited()