diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 82b820d8cd9..3e0775ac09e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -200,6 +200,13 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = ( ) +# OAuth discovery retry cooldown for servers whose endpoints stay unresolved. The base is one +# reload cadence so a transient upstream failure recovers immediately; the cap bounds the request +# amplification and log volume of a permanently broken configuration. +_OAUTH_DISCOVERY_RETRY_BASE_SECONDS = 30.0 +_OAUTH_DISCOVERY_RETRY_MAX_SECONDS = 900.0 + + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -247,6 +254,7 @@ def _endpoints_yield_to_issuer( authorization_url: str | None, token_url: str | None, registration_url: str | None, + server_ref: str, ) -> tuple[str | None, str | None, str | None]: """The single rule that makes an admin-configured ``issuer`` the sole authoritative endpoint source (RFC 8414 §3.3): when it is set for a discovery auth type, the stored/manual @@ -256,9 +264,29 @@ def _endpoints_yield_to_issuer( i.e. all ``None`` when issuer-anchored, else the inputs unchanged. Called at every resolution site so the invariant holds in one place instead of being re-derived per merge. """ - if issuer is not None and is_discovery_auth_type: - return None, None, None - return authorization_url, token_url, registration_url + if issuer is None or not is_discovery_auth_type: + return authorization_url, token_url, registration_url + discarded = sorted( + label + for label, value in ( + ("authorization_url", authorization_url), + ("token_url", token_url), + ("registration_url", registration_url), + ) + if value + ) + if discarded: + verbose_logger.warning( + "MCP server %s has a pinned Issuer, so its stored %s %s not used: an anchored issuer is the " + "sole endpoint source (RFC 8414 section 3.3) and a failed issuer fetch fails closed rather " + "than falling back to them. To use manually configured endpoints instead, clear the Issuer " + "field and re-enter the endpoint urls (clearing the Issuer also clears endpoints that may " + "have been resolved under it), or clear the Issuer alone to re-discover from the server url.", + server_ref, + ", ".join(discarded), + "is" if len(discarded) == 1 else "are", + ) + return None, None, None def _normalized_authorize_endpoint(url: str) -> str: @@ -280,6 +308,68 @@ def _issuer_matches(claimed_issuer: object, configured_issuer: str) -> bool: return _normalized_authorize_endpoint(claimed_issuer) == _normalized_authorize_endpoint(configured_issuer) +def _flow_endpoints_missing( + auth_type: MCPAuthType | None, + oauth2_flow: str | None, + authorization_url: str | None, + token_url: str | None, + token_exchange_endpoint: str | None = None, +) -> bool: + """Whether a built server is missing an endpoint its flow needs to run at all. + + Used by the reload fast-path exemption: discovery runs at build time only, and the fast path + reuses an unchanged row's registry entry verbatim, so a server whose discovery came back empty + (transient upstream failure, rate limiting) would stay broken until some unrelated config write + bumps ``updated_at``, serving its 400 the whole time. Rebuilding just these entries retries + discovery on the normal reload cadence. It costs no extra fetch for servers that resolved, and + none for those with no discovery source, since the build skips discovery for both. + """ + if auth_type == MCPAuth.oauth2_token_exchange: + # A configured exchange endpoint replaces discovery entirely; only a server that must + # discover its token endpoint and still has none is unresolved. + return token_exchange_endpoint is None and token_url is None + if auth_type not in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: + return False + if oauth2_flow == "client_credentials": + return token_url is None + return authorization_url is None or token_url is None + + +def _oauth_endpoints_unresolved(server: MCPServer) -> bool: + """``_flow_endpoints_missing`` over a built registry entry, for the reload fast-path check. + + The flow comes from ``effective_oauth2_flow``, the one column-first, shape-fallback judge every + flow decision uses, not from the raw column: a legacy row the startup backfill deliberately left + unstamped (the ambiguous M2M shape) serves M2M at request time, and reading the bare column here + would classify it as interactive-missing-endpoints and re-run discovery on every reload. + """ + if ( + server.auth_type == MCPAuth.oauth2_token_exchange + and server.token_exchange_profile == "entra_obo" + and not server.scopes + ): + # entra_obo fails closed at exchange time without a scope (token_exchanger.py), and scopes + # can come from resource discovery, so a server that resolved its endpoints but no scopes is + # still unresolved for its flow. + return True + if server.is_dcr_bridge and not server.client_id and server.registration_url is None: + # A DCR bridge with no admin-configured client can only register callers through the + # upstream's registration endpoint, so a build that resolved the authorize and token + # endpoints but not registration_endpoint (partial metadata) is still unresolved for its + # flow and must keep retrying; without this it silently degrades to the short-circuit arm + # until an unrelated config write. Scopes are deliberately NOT part of completeness: they + # are a request hint the authorization server bounds at consent (RFC 6749 section 3.3), + # and a server without them is fully functional. + return True + return _flow_endpoints_missing( + server.auth_type, + MCPServerManager.effective_oauth2_flow(server), + server.authorization_url, + server.token_url, + server.token_exchange_endpoint, + ) + + def _endpoints_corroborate_authorization_url( source_authorization_url: str | None, trusted_authorization_url: str | None, @@ -311,11 +401,10 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv during re-discovery downgrades a working server (``authorization_url`` set) to a broken one (``None``, /authorize 400s) with no configuration change. Mirrors the ``short_prefix`` carry-forward. Skipped when the server's ``url`` or ``auth_type`` changed, since the previous - endpoints may then belong to a different upstream. ``registration_url`` IS carried even though - ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only restores - the same in-memory value the previous build already ran with, while persisting it would flip - ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for dcr_bridge - servers that never had one configured. + endpoints may then belong to a different upstream. Discovery results live only on the in-memory + registry entry; the gateway never writes them to the row, whose OAuth columns carry admin intent + alone, so this carry is the sole last-known-good mechanism and restores exactly the values the + previous build already ran with. Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the previous ``token_url``/``registration_url``/``scopes`` are carried only when the previous @@ -1182,6 +1271,40 @@ class MCPServerManager: # empty result, or failure). Used to throttle re-probes for servers that do # not return instructions, and to apply a short cooldown after failures. self._upstream_initialize_instructions_probed_at: dict[str, float] = {} + # Per-server (consecutive failures, monotonic timestamp) for OAuth discovery retries, so a + # server whose endpoints never resolve backs off instead of re-running the full + # RFC 9728 -> 8414 chain, and re-logging its warning, on every reload forever. + self._oauth_discovery_retry_state: dict[ + str, tuple[int, float] + ] = {} # mutable-ok: retry cooldown cache, keyed per server and pruned on success + + def _oauth_discovery_retry_due(self, server_id: str) -> bool: + """Whether an unresolved server is due for another discovery attempt. + + The reload fast-path exemption is what retries a failed discovery, so without a cooldown a + permanently unresolvable server re-runs the whole RFC 9728 -> RFC 8414 -> origin-fallback + chain and re-emits its unresolved-endpoints warning on every reload, per server, forever. + Delay doubles per consecutive failure from ``_OAUTH_DISCOVERY_RETRY_BASE_SECONDS`` up to + ``_OAUTH_DISCOVERY_RETRY_MAX_SECONDS``, so a transient outage still recovers on the next + reload while a broken configuration settles to one attempt per cap. + """ + state = self._oauth_discovery_retry_state.get(server_id) + if state is None: + return True + failures, attempted_at = state + delay = min( + _OAUTH_DISCOVERY_RETRY_BASE_SECONDS * (2 ** max(failures - 1, 0)), + _OAUTH_DISCOVERY_RETRY_MAX_SECONDS, + ) + return (time.monotonic() - attempted_at) >= delay + + def _record_oauth_discovery_outcome(self, server: MCPServer) -> None: + """Advance or clear a server's retry cooldown after a rebuild resolved it or did not.""" + if not _oauth_endpoints_unresolved(server): + self._oauth_discovery_retry_state.pop(server.server_id, None) + return + failures, _ = self._oauth_discovery_retry_state.get(server.server_id, (0, 0.0)) + self._oauth_discovery_retry_state[server.server_id] = (failures + 1, time.monotonic()) def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: raw = getattr(client, "_last_initialize_instructions", None) @@ -1357,6 +1480,7 @@ class MCPServerManager: manual_authorization_url, manual_token_url, manual_registration_url, + server_name or server_id, ) should_discover = _has_oauth_discovery_source(server_url, use_issuer_anchor) and ( is_discovery_auth_type or obo_needs_discovery @@ -1834,7 +1958,6 @@ class MCPServerManager: *, credentials_are_encrypted: bool = True, env_vars_are_encrypted: Optional[bool] = None, - persist_discovered_endpoints: bool = True, ) -> MCPServer: _mcp_info: MCPInfo = mcp_server.mcp_info or {} env_dict = _deserialize_json_dict(getattr(mcp_server, "env", None)) @@ -1925,7 +2048,12 @@ class MCPServerManager: or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url), ) manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer( - manual_issuer, is_discovery_auth_type, manual_authorization_url, manual_token_url, manual_registration_url + manual_issuer, + is_discovery_auth_type, + manual_authorization_url, + manual_token_url, + manual_registration_url, + mcp_server.alias or mcp_server.server_name or mcp_server.server_id, ) gated_oauth_metadata = await self._resolve_table_oauth_metadata( mcp_server=mcp_server, @@ -2033,143 +2161,8 @@ class MCPServerManager: max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) _warn_internal_delegate_pkce_if_applicable(new_server, source="database") - if persist_discovered_endpoints: - await self._persist_discovered_obo_token_url( - server_id=mcp_server.server_id, - auth_type=auth_type, - existing_token_url=manual_token_url, - discovered_token_url=new_server.token_url, - ) - await self._persist_discovered_oauth_endpoints( - server_id=mcp_server.server_id, - auth_type=auth_type, - existing_issuer=manual_issuer, - existing_authorization_url=manual_authorization_url, - existing_token_url=manual_token_url, - existing_scopes=scopes, - metadata=gated_oauth_metadata, - is_issuer_anchored=use_issuer_anchor, - ) return new_server - async def _persist_discovered_obo_token_url( - self, - *, - server_id: str, - auth_type: Optional[MCPAuthType], - existing_token_url: Optional[str], - discovered_token_url: Optional[str], - ) -> None: - """Write a freshly discovered OBO token endpoint back onto the DB row. - - ``build_mcp_server_from_table`` resolves ``token_url`` via RFC 9728 -> RFC 8414 for an - ``oauth2_token_exchange`` server that has none configured, but that resolved value otherwise - lives only on the returned in-memory object; the row keeps ``token_url=None`` so every rebuild - re-runs discovery, and a transient upstream outage during a rebuild leaves the server with no - endpoint until discovery next succeeds. Persisting it makes ``_obo_needs_endpoint_discovery`` - return False on the next build. Fires at most once per server (skipped once the row has a - value), and is best-effort: a write failure just means discovery runs again next time. - """ - if auth_type != MCPAuth.oauth2_token_exchange: - return - if existing_token_url or not discovered_token_url: - return - from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 - - if prisma_client is None: - return - try: - await MCPServerRepository(prisma_client).table.update( - where={"server_id": server_id}, - data={"token_url": discovered_token_url}, - ) - verbose_logger.debug("Persisted discovered OBO token_url for MCP server %s", server_id) - except Exception as exc: # noqa: BLE001 - best-effort; a failed write re-discovers next build - verbose_logger.warning("Failed to persist discovered OBO token_url for MCP server %s: %s", server_id, exc) - - async def _persist_discovered_oauth_endpoints( - self, - *, - server_id: str, - auth_type: MCPAuthType | None, - existing_issuer: str | None, - existing_authorization_url: str | None, - existing_token_url: str | None, - existing_scopes: list[str] | None, - metadata: MCPOAuthMetadata | None, - is_issuer_anchored: bool = False, - ) -> None: - """Write freshly discovered OAuth endpoints back onto the DB row. - - Same rationale as ``_persist_discovered_obo_token_url`` but for the interactive oauth2 - family: discovered ``authorization_url``/``token_url``/``scopes`` otherwise live only on - the in-memory registry entry, which is rebuilt on every client connect (the DCR reuse path - calls ``update_server``) and on every post-write DB reload, so one failed re-discovery - serves the 400 "authorization url is not configured" from /authorize until a later rebuild succeeds. - Only fills row fields that are currently empty, never persists origin-fallback guesses - (RFC 9728/8414-advertised metadata only), and deliberately skips ``registration_url`` - because ``_dcr_bridge_relays_client_registration`` keys off that column. Best-effort: a - failed write re-discovers on the next build. Scopes go through ``update_mcp_server`` so - they merge into the credentials blob without touching the stored client credentials. - - For an issuer-anchored server (``is_issuer_anchored``) the endpoints are re-derived from the - §3.3-validated issuer document on every build, so they are NOT persisted into the endpoint - columns: persisting them would make the next build see populated endpoints and treat them as - authoritative stored values, defeating the "endpoints come solely from the issuer" invariant. - Only the resource-driven scopes are persisted for such servers. - """ - if auth_type not in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: - return - if metadata is None or metadata.from_origin_fallback: - return - issuer_update = ( - {"issuer": metadata.discovered_issuer} if metadata.discovered_issuer and not existing_issuer else {} - ) - authorization_url_update = ( - {"authorization_url": metadata.authorization_url} - if metadata.authorization_url and not existing_authorization_url and not is_issuer_anchored - else {} - ) - token_url_update = ( - {"token_url": metadata.token_url} - if metadata.token_url and not existing_token_url and not is_issuer_anchored - else {} - ) - scopes_update = {"credentials": {"scopes": metadata.scopes}} if metadata.scopes and not existing_scopes else {} - updates: dict[str, object] = { - **issuer_update, - **authorization_url_update, - **token_url_update, - **scopes_update, - } - if not updates: - return - from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # db.py imports this module at load - update_mcp_server, - ) - from litellm.proxy._types import UpdateMCPServerRequest # noqa: PLC0415 # heavy module; import at call time - from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime value, set after startup - - if prisma_client is None: - return - try: - await update_mcp_server( - prisma_client=prisma_client, - data=UpdateMCPServerRequest.model_validate({"server_id": server_id, **updates}), - touched_by="mcp_oauth_discovery", - ) - verbose_logger.info( - "Persisted discovered OAuth endpoints for MCP server %s: %s", - server_id, - sorted(updates), - ) - except Exception as exc: # noqa: BLE001 - best-effort; a failed write re-discovers next build - verbose_logger.warning( - "Failed to persist discovered OAuth endpoints for MCP server %s: %s", - server_id, - exc, - ) - async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): """Register OpenAPI tools if the server has a spec_path configured.""" if server.spec_path: @@ -5347,6 +5340,10 @@ class MCPServerManager: and existing_server.updated_at is not None and server.updated_at is not None and existing_server.updated_at == server.updated_at + and not ( + _oauth_endpoints_unresolved(existing_server) + and self._oauth_discovery_retry_due(server.server_id) + ) ): # Re-use existing server instance to avoid re-running build_mcp_server_from_table() # which can perform network discovery for OAuth2 servers. @@ -5364,6 +5361,7 @@ class MCPServerManager: # already-decrypted records add_server/update_server are handed. # Decrypt them while building the registry entry. new_server = await self.build_mcp_server_from_table(server, env_vars_are_encrypted=True) + self._record_oauth_discovery_outcome(new_server) # Carry the cached short_prefix from the previous registry entry # (if any) so the prefix is stable across reloads. if existing_server is not None and existing_server.short_prefix: diff --git a/litellm/proxy/_experimental/mcp_server/oauth_issuer_stamp_backfill.py b/litellm/proxy/_experimental/mcp_server/oauth_issuer_stamp_backfill.py new file mode 100644 index 00000000000..874fcc64772 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth_issuer_stamp_backfill.py @@ -0,0 +1,148 @@ +"""One-time heal for MCP server rows whose ``issuer`` a released version wrote by itself. + +Until the write was removed, OAuth discovery stamped the issuer it discovered onto the ``issuer`` +column trust-on-first-use. That column means "the admin pinned this trust anchor", so the next +registry build read the gateway's own output back as admin intent: the server turned issuer-anchored +(RFC 8414 section 3.3), its stored authorization/token/registration URLs stopped applying, and a +failed issuer-document fetch left it with no authorize endpoint (GH #34985). + +Deleting the write fixes every row created afterwards but cannot fix a row already stamped, which +still reads as pinned. This heals those rows by clearing the stamp so their configured endpoints +apply again. + +The signal is a heuristic, and deliberately a narrow one. ``updated_by`` records only the most recent +writer, and no audit trail says which field that writer touched, so "discovery wrote this issuer" is +not directly knowable. Two independent clauses bound it, and each rules out a different way of +destroying a pin an admin meant. + +Configured endpoints must be present. A deliberately pinned row very often has none, both because the +Issuer field is documented as overriding them and because ``update_mcp_server`` clears them when an +issuer changes, so "issuer set, endpoints empty" is the canonical shape of a real pin and must never +be cleared on this evidence. Skipping those rows costs little: with nothing configured to restore, the +anchored and resource-rooted paths resolve from the same upstream document, and the row still gets the +unresolved-endpoint retry and the anchored-discard warning. + +The configured endpoints must also share the issuer's origin. A stamped issuer is by construction the +one self-attested by the authorization-server document discovery reached from this very server, so +endpoints typed alongside it address that same authority. An admin who pinned an issuer and typed +endpoints for a different authority is expressing an intent that clearing the issuer would discard, so +that row is warned about and never healed. + +What survives both clauses is a row whose configured endpoints and stamped issuer share an origin, +which is exactly the GH #34985 shape. An admin who pinned that same origin by hand lands here too, and +for them the clear is close to a no-op: their typed endpoints keep serving and still anchor the +RFC 9700 corroboration gate, with only the stricter section 3.3 anchoring lost. Every heal logs the +cleared value so it can be restored, and the clear is recorded under this module's actor so the heal +runs at most once per row. +""" + +from typing import Protocol +from urllib.parse import urlparse + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import canonicalize_url_identity +from litellm.proxy.utils import PrismaClient + +# The actor the removed discovery write-back stamped rows with. +_DISCOVERY_ACTOR = "mcp_oauth_discovery" + +# The actor recorded on a healed row, which also makes the heal idempotent: once a row is cleared it +# no longer matches ``updated_by == _DISCOVERY_ACTOR`` and is never reconsidered. +_BACKFILL_ACTOR = "mcp_oauth_issuer_stamp_backfill" + +_AUTH_TYPES_WITH_ISSUER_ANCHORING = ("oauth2", "true_passthrough", "oauth_delegate") + + +def _origin(url: str) -> str | None: + """The scheme-and-authority identity of ``url``, or ``None`` when it has none. + + Built on the shared URL canonicalizer so the lowercase-host and default-port rules match the + RFC 8414 issuer comparison the resolution path uses, instead of being re-derived here. + """ + parsed = urlparse(canonicalize_url_identity(url)) + if not parsed.scheme or not parsed.netloc: + return None + return f"{parsed.scheme}://{parsed.netloc}" + + +class _MCPServerRow(Protocol): + """The MCP server row fields this heal reads, so the untyped DB record is narrowed once here.""" + + server_id: str + alias: str | None + server_name: str | None + auth_type: str | None + issuer: str | None + authorization_url: str | None + token_url: str | None + registration_url: str | None + updated_by: str | None + + +def _is_stamped_issuer_row(row: _MCPServerRow) -> bool: + """Whether this row carries the full signature of a gateway-written issuer stamp. + + The whole rule lives here, including the writer check the query also filters on, so the decision + to clear an admin-visible field is auditable in one place rather than split between a predicate + and a query. + """ + if getattr(row, "updated_by", None) != _DISCOVERY_ACTOR: + return False + if not (getattr(row, "issuer", None) or "").strip(): + return False + if getattr(row, "auth_type", None) not in _AUTH_TYPES_WITH_ISSUER_ANCHORING: + return False + configured = tuple( + value.strip() + for value in (row.authorization_url, row.token_url, row.registration_url) + if value and value.strip() + ) + if not configured: + return False + issuer_origin = _origin(row.issuer or "") + return issuer_origin is not None and all(_origin(endpoint) == issuer_origin for endpoint in configured) + + +async def backfill_discovery_stamped_issuers(prisma_client: PrismaClient) -> int: + """Clear gateway-written issuer stamps, returning the number of rows healed.""" + candidate_rows: list[_MCPServerRow] = await prisma_client.db.litellm_mcpservertable.find_many( + where={ + "updated_by": _DISCOVERY_ACTOR, + "auth_type": {"in": list(_AUTH_TYPES_WITH_ISSUER_ANCHORING)}, + }, + ) + stamped = tuple(row for row in candidate_rows if _is_stamped_issuer_row(row)) + if not stamped: + return 0 + + healed = 0 + for row in stamped: + try: + await prisma_client.db.litellm_mcpservertable.update( + where={"server_id": row.server_id}, + data={"issuer": None, "updated_by": _BACKFILL_ACTOR}, + ) + except Exception as exc: # noqa: BLE001 - per-row best effort; the next boot retries + verbose_proxy_logger.warning( + "MCP issuer stamp backfill: could not heal server_id=%s: %s", row.server_id, exc + ) + continue + healed += 1 + verbose_proxy_logger.warning( + "MCP issuer stamp backfill: cleared issuer %r on server_id=%s (alias=%s). OAuth discovery " + "had written that value onto the Issuer column, which made the server issuer-anchored and " + "fail-closed, and its configured Authorization/Token/Registration URLs were being ignored " + "as a result; those now apply again. If you pinned this issuer deliberately, set it again " + "via the dashboard or PUT /v1/mcp/server to restore RFC 8414 section 3.3 anchoring.", + row.issuer, + row.server_id, + row.alias or row.server_name, + ) + + if healed: + verbose_proxy_logger.warning( + "MCP issuer stamp backfill: healed %d server(s) whose Issuer had been written by OAuth " + "discovery rather than by an admin", + healed, + ) + return healed diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 282184d6495..1205d23ce02 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1526,7 +1526,6 @@ if MCP_AVAILABLE: temporary_server = await global_mcp_server_manager.build_mcp_server_from_table( temp_record, credentials_are_encrypted=False, - persist_discovered_endpoints=False, ) _cache_temporary_mcp_server( temporary_server, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70484eb1e4e..18a927e7a44 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6758,6 +6758,9 @@ class ProxyConfig: from litellm.proxy._experimental.mcp_server.oauth2_flow_backfill import ( backfill_null_oauth2_flows, ) + from litellm.proxy._experimental.mcp_server.oauth_issuer_stamp_backfill import ( + backfill_discovery_stamped_issuers, + ) try: if prisma_client is not None: @@ -6767,6 +6770,16 @@ class ProxyConfig: "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db backfill - {}".format(str(e)) ) + try: + if prisma_client is not None: + await backfill_discovery_stamped_issuers(prisma_client) + except Exception as e: # noqa: BLE001 + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db issuer stamp backfill - {}".format( + str(e) + ) + ) + try: await global_mcp_server_manager.reload_servers_from_database() except Exception as e: @@ -6778,6 +6791,31 @@ class ProxyConfig: if self._should_load_db_object(object_type="mcp"): await self._init_mcp_servers_in_db() + async def reload_mcp_servers_from_db(self) -> None: + """Registry refresh only, for the periodic job in store_model_in_db-off deployments. + + Deliberately narrower than ``init_mcp_servers_from_db``: the oauth2_flow backfill is a write + path that only needs to run once at startup, so the cadence here is purely the read-side + reload whose fast-path exemption retries failed OAuth discovery. Gated the same way, so an + admin who excluded mcp from supported_db_objects opts out of this too. + """ + if not self._should_load_db_object(object_type="mcp"): + return + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if not is_mcp_available(): + return + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + try: + await global_mcp_server_manager.reload_servers_from_database() + except Exception as e: # noqa: BLE001 # scheduled job: a reload failure must not kill the recurring retry + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:reload_mcp_servers_from_db - {}".format(str(e)) + ) + async def _init_agents_in_db(self, prisma_client: PrismaClient): from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, @@ -8099,6 +8137,22 @@ class ProxyStartupEvent: if store_model_in_db is not True: await proxy_config.init_mcp_servers_from_db() + if prisma_client is not None: + # DB-backed MCP servers are live objects in every mode, so the registry refresh that + # store_model_in_db=True deployments get via the add_deployment job must run here + # too; without it, a server whose OAuth discovery failed at startup is rebuilt only + # by a management write, since the reload fast path is the retry's only driver. + mcp_reload_interval_seconds = proxy_config_reload_interval_seconds + if not isinstance(mcp_reload_interval_seconds, int) or mcp_reload_interval_seconds <= 0: + mcp_reload_interval_seconds = 30 + scheduler.add_job( + proxy_config.reload_mcp_servers_from_db, + "interval", + seconds=mcp_reload_interval_seconds, + id="reload_mcp_servers_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) await cls._initialize_slack_alerting_jobs( scheduler=scheduler, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index c063915e2e8..f6bd79c5d2d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -240,10 +240,10 @@ async def test_explicit_null_clears_upstream_resource_and_keeps_the_rest_of_the_ @pytest.mark.asyncio -async def test_url_change_clears_stale_discovered_oauth_fields(): - """Re-pointing the server url at a potentially different upstream must clear the discovered or - trust-on-first-use OAuth issuer and endpoints, so the new upstream re-discovers instead of - anchoring on the previous upstream's issuer (RFC 8414 §3.3 against a stale anchor).""" +async def test_url_change_clears_stale_oauth_fields(): + """Re-pointing the server url at a potentially different upstream must clear the OAuth issuer and + endpoints, so the new upstream re-discovers instead of anchoring on the previous upstream's issuer + (RFC 8414 §3.3 against a stale anchor).""" mock_prisma = _mock_prisma() existing = MagicMock() existing.auth_type = "oauth2" @@ -350,11 +350,13 @@ async def test_repointing_pinned_issuer_clears_stale_endpoints_keeps_new_issuer( @pytest.mark.asyncio -async def test_establishing_issuer_first_time_preserves_discovered_fields(): - """Establishing an issuer for the first time (None -> X), which is exactly what the trust-on-first-use - discovery write-back does, must NOT clear the endpoints or oauth2_flow it discovered in the same - write. Only an issuer that was already pinned and is now changed or cleared invalidates its - endpoints, so the discovery persist cannot wipe the fields it just resolved.""" +async def test_establishing_issuer_first_time_preserves_endpoints_set_in_the_same_write(): + """Establishing an issuer for the first time (None -> X) must NOT clear endpoints or oauth2_flow + submitted in the same write. Only an issuer that was already pinned and is now changed or cleared + invalidates its endpoints, so an admin configuring an issuer and its endpoints together keeps + both. The write-back this once guarded (trust-on-first-use discovery stamping the issuer it had + just resolved) no longer exists; the db.py rule it relies on still governs admin writes, which is + what this now covers.""" mock_prisma = _mock_prisma() existing = MagicMock() existing.auth_type = "oauth2" @@ -370,7 +372,7 @@ async def test_establishing_issuer_first_time_preserves_discovered_fields(): token_url="https://discovered-idp.example.com/token", oauth2_flow="authorization_code", ) - await update_mcp_server(mock_prisma, data, "mcp_oauth_discovery") + await update_mcp_server(mock_prisma, data, "some-admin@example.com") data_dict = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] assert data_dict["issuer"] == "https://discovered-idp.example.com" @@ -380,9 +382,9 @@ async def test_establishing_issuer_first_time_preserves_discovered_fields(): @pytest.mark.asyncio -async def test_unchanged_url_does_not_clear_discovered_oauth_fields(): - """A partial update that resends the same url (or omits it) must not clear the discovered OAuth - fields, so a routine save does not force needless re-discovery.""" +async def test_unchanged_url_does_not_clear_oauth_fields(): + """A partial update that resends the same url (or omits it) must not clear the OAuth fields, so a + routine save does not force needless re-discovery.""" mock_prisma = _mock_prisma() existing = MagicMock() existing.auth_type = "oauth2" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5f7f2267fc7..8a8dea0ba28 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2,6 +2,7 @@ import importlib import asyncio import json import logging +import time import os import sys from datetime import datetime @@ -35,6 +36,8 @@ from mcp.types import Tool as MCPTool from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _deserialize_json_dict, + _flow_endpoints_missing, + _oauth_endpoints_unresolved, _deserialize_json_list, _normalize_mcp_server_cost_info, _should_strip_caller_authorization, @@ -1594,21 +1597,15 @@ class TestMCPServerManager: token_url="https://idp.example.com/token", scopes=["read"], ) - with ( - patch.object( - manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=issuer_resolved) - ) as anchored, - patch.object(manager, "_persist_discovered_oauth_endpoints", new=AsyncMock()) as mock_persist, - ): + with patch.object( + manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=issuer_resolved) + ) as anchored: built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) anchored.assert_awaited_once_with("https://idp.example.com", "https://up.example.com/mcp") assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url == "https://idp.example.com/token" assert built.token_url != "https://attacker.example.com/steal" - # The issuer-anchored endpoints are never persisted into the endpoint columns, so a later - # build cannot treat them as authoritative stored values. - assert mock_persist.await_args.kwargs["is_issuer_anchored"] is True @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1624,8 +1621,8 @@ class TestMCPServerManager: and PKCE verifier to the attacker (config-time RFC 9700 mix-up). The resource-driven scopes are kept, because scope selection is resource-driven (MCP Scope Selection Strategy) and scope inflation is bounded by the authorization server at consent (RFC 6749 §3.3), not by dropping - scopes on an endpoint mismatch. Both the in-memory merge and the persisted metadata drop only - the uncorroborated endpoints.""" + scopes on an endpoint mismatch. The gateway persists nothing, so the in-memory merge is the + entire behavior.""" manager = MCPServerManager() row = LiteLLM_MCPServerTable( server_id="manual-auth-url-3", @@ -1645,20 +1642,13 @@ class TestMCPServerManager: registration_url="https://attacker.example.com/register", scopes=["read", "admin"], ) - with ( - patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)), - patch.object(manager, "_persist_discovered_oauth_endpoints", new=AsyncMock()) as mock_persist, - ): + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=metadata)): built = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) assert built.authorization_url == "https://idp.example.com/authorize" assert built.token_url is None assert built.registration_url is None assert built.scopes == ["read", "admin"] - persisted_metadata = mock_persist.await_args.kwargs["metadata"] - assert persisted_metadata.token_url is None - assert persisted_metadata.registration_url is None - assert persisted_metadata.scopes == ["read", "admin"] @pytest.mark.asyncio async def test_build_from_table_skips_discovery_when_all_upstream_oauth_fields_present(self): @@ -5586,388 +5576,300 @@ class TestMCPServerTimestamps: assert server.token_exchange_endpoint == "https://idp.example.com/token" @pytest.mark.asyncio - async def test_build_mcp_server_from_table_persists_discovered_obo_token_url(self): - """A DB-backed OBO server with no configured endpoint discovers token_url and must write it - back to the row, so the next rebuild skips discovery instead of re-running it every time.""" + async def test_discovery_never_writes_the_database(self): + """The #34985 regression, stated as the design invariant that fixes it: the gateway never + writes discovery results to the row. The OAuth columns and credentials.scopes carry admin + intent alone, so nothing the gateway learns can read back as an admin pin on a later build + (which is what anchored stamped servers fail-closed and 400ed /authorize). Discovery output + lives on the in-memory registry entry only, for oauth2 and OBO alike.""" manager = MCPServerManager() async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): - assert server_url == "https://example.com/mcp" - assert allow_origin_fallback is False # OBO never guesses the origin return MCPOAuthMetadata( - scopes=None, - authorization_url=None, - token_url="https://discovered.example.com/token", - registration_url=None, - ) - - manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] - - record = LiteLLM_MCPServerTable( - server_id="obo-persist-1", - server_name="obo_persist", - url="https://example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - credentials={"client_id": "cid", "client_secret": "csec", "audience": "aud"}, - ) - - update_mock = AsyncMock() - repo_instance = MagicMock() - repo_instance.table.update = update_mock - with ( - patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", - return_value=repo_instance, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - server = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) - - assert server.token_url == "https://discovered.example.com/token" - update_mock.assert_awaited_once() - assert update_mock.call_args.kwargs["where"] == {"server_id": "obo-persist-1"} - assert update_mock.call_args.kwargs["data"] == {"token_url": "https://discovered.example.com/token"} - - @pytest.mark.asyncio - async def test_persist_discovered_obo_token_url_skips_when_not_needed(self): - """The write-back fires only for an OBO server that discovered a new endpoint: a row that - already has token_url, a non-OBO auth_type, or a discovery that found nothing all no-op.""" - manager = MCPServerManager() - update_mock = AsyncMock() - repo_instance = MagicMock() - repo_instance.table.update = update_mock - - with ( - patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", - return_value=repo_instance, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - # already populated -> no write - await manager._persist_discovered_obo_token_url( - server_id="s", - auth_type=MCPAuth.oauth2_token_exchange, - existing_token_url="https://already.example.com/token", - discovered_token_url="https://new.example.com/token", - ) - # not an OBO server -> no write - await manager._persist_discovered_obo_token_url( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_token_url=None, - discovered_token_url="https://new.example.com/token", - ) - # discovery found nothing -> no write - await manager._persist_discovered_obo_token_url( - server_id="s", - auth_type=MCPAuth.oauth2_token_exchange, - existing_token_url=None, - discovered_token_url=None, - ) - - update_mock.assert_not_awaited() - - @pytest.mark.asyncio - async def test_persist_discovered_obo_token_url_is_best_effort(self): - """A write-back failure must not propagate; discovery just re-runs on the next build.""" - manager = MCPServerManager() - update_mock = AsyncMock(side_effect=Exception("db unavailable")) - repo_instance = MagicMock() - repo_instance.table.update = update_mock - - with ( - patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", - return_value=repo_instance, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - await manager._persist_discovered_obo_token_url( - server_id="s", - auth_type=MCPAuth.oauth2_token_exchange, - existing_token_url=None, - discovered_token_url="https://new.example.com/token", - ) - - update_mock.assert_awaited_once() - - @pytest.mark.asyncio - async def test_build_mcp_server_from_table_persists_discovered_oauth_endpoints(self): - """A DB-backed oauth2 server with no configured endpoints discovers them and must write - authorization_url, token_url, and scopes back to the row; otherwise the resolved values - live only in memory and one failed re-discovery serves the 400 "authorization url is not configured" - from /authorize. registration_url must never be persisted because - _dcr_bridge_relays_client_registration keys off that column.""" - manager = MCPServerManager() - - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): - assert allow_origin_fallback is True - return MCPOAuthMetadata( - scopes=["mcp.read", "mcp.write"], + scopes=["mcp.read"], authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", - ) - - manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] - - record = LiteLLM_MCPServerTable( - server_id="oauth-persist-1", - server_name="oauth_persist", - url="https://example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - oauth2_flow="authorization_code", - credentials={"client_id": "cid", "client_secret": "csec"}, - ) - - update_mcp_server_mock = AsyncMock() - with ( - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=update_mcp_server_mock, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - server = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) - - assert server.authorization_url == "https://idp.example.com/authorize" - update_mcp_server_mock.assert_awaited_once() - persisted = update_mcp_server_mock.call_args.kwargs["data"] - assert persisted.server_id == "oauth-persist-1" - assert persisted.authorization_url == "https://idp.example.com/authorize" - assert persisted.token_url == "https://idp.example.com/token" - assert persisted.credentials == {"scopes": ["mcp.read", "mcp.write"]} - assert "registration_url" not in persisted.fields_set() - assert update_mcp_server_mock.call_args.kwargs["touched_by"] == "mcp_oauth_discovery" - - @pytest.mark.asyncio - async def test_persist_discovered_oauth_endpoints_guards(self): - """The write-back must no-op for non-discovery auth types, empty discovery, origin-fallback - guesses (never harden an inferred authorization server into configuration), and rows whose - fields are all already populated.""" - manager = MCPServerManager() - advertised = MCPOAuthMetadata( - scopes=["s1"], - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", - ) - - update_mcp_server_mock = AsyncMock() - with ( - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=update_mcp_server_mock, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.api_key, - existing_issuer=None, - existing_authorization_url=None, - existing_token_url=None, - existing_scopes=None, - metadata=advertised, - ) - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer=None, - existing_authorization_url=None, - existing_token_url=None, - existing_scopes=None, - metadata=None, - ) - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer=None, - existing_authorization_url=None, - existing_token_url=None, - existing_scopes=None, - metadata=advertised.model_copy(update={"from_origin_fallback": True}), - ) - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer=None, - existing_authorization_url="https://configured.example.com/authorize", - existing_token_url="https://configured.example.com/token", - existing_scopes=["configured"], - metadata=advertised, - ) - - update_mcp_server_mock.assert_not_awaited() - - @pytest.mark.asyncio - async def test_persist_discovered_oauth_endpoints_only_fills_empty_fields(self): - """A row that already has token_url keeps it; only the missing authorization_url and - scopes are written, so admin-typed values always win over discovery.""" - manager = MCPServerManager() - - update_mcp_server_mock = AsyncMock() - with ( - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=update_mcp_server_mock, - ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer=None, - existing_authorization_url=None, - existing_token_url="https://configured.example.com/token", - existing_scopes=None, - metadata=MCPOAuthMetadata( - scopes=["s1"], - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", - ), - ) - - update_mcp_server_mock.assert_awaited_once() - persisted = update_mcp_server_mock.call_args.kwargs["data"] - assert persisted.authorization_url == "https://idp.example.com/authorize" - assert persisted.credentials == {"scopes": ["s1"]} - assert "token_url" not in persisted.fields_set() - - @pytest.mark.asyncio - async def test_persist_discovered_oauth_endpoints_writes_discovered_issuer_trust_on_first_use(self): - """A server with no configured issuer records the discovered issuer trust-on-first-use, so the - next rebuild anchors discovery on it (RFC 8414 §3.3) instead of re-trusting the resource. When - an issuer is already set (admin-typed or a prior discovery), it is never overwritten.""" - manager = MCPServerManager() - metadata = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", - discovered_issuer="https://idp.example.com", - ) - - update_mcp_server_mock = AsyncMock() - with ( - patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer=None, - existing_authorization_url=None, - existing_token_url=None, - existing_scopes=None, - metadata=metadata, - ) - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer="https://admin-configured.example.com", - existing_authorization_url="https://admin-configured.example.com/authorize", - existing_token_url="https://admin-configured.example.com/token", - existing_scopes=["cfg"], - metadata=metadata, - ) - - assert update_mcp_server_mock.await_count == 1 - persisted = update_mcp_server_mock.call_args.kwargs["data"] - assert persisted.issuer == "https://idp.example.com" - - @pytest.mark.asyncio - async def test_persist_discovered_oauth_endpoints_does_not_persist_endpoints_for_issuer_anchored(self): - """For an issuer-anchored server the endpoints are re-derived from the §3.3-validated issuer - document every build, so they must NOT be written into the endpoint columns: persisting them - would make the next build see populated endpoints and treat them as authoritative stored - values, defeating the issuer-only invariant. Only the resource-driven scopes are persisted.""" - manager = MCPServerManager() - metadata = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", - scopes=["read"], - ) - - update_mcp_server_mock = AsyncMock() - with ( - patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - ): - await manager._persist_discovered_oauth_endpoints( - server_id="s", - auth_type=MCPAuth.oauth2, - existing_issuer="https://idp.example.com", - existing_authorization_url=None, - existing_token_url=None, - existing_scopes=None, - metadata=metadata, - is_issuer_anchored=True, - ) - - update_mcp_server_mock.assert_awaited_once() - persisted = update_mcp_server_mock.call_args.kwargs["data"] - assert "authorization_url" not in persisted.fields_set() - assert "token_url" not in persisted.fields_set() - assert persisted.credentials == {"scopes": ["read"]} - - @pytest.mark.asyncio - async def test_build_mcp_server_from_table_skips_persistence_for_temporary_servers(self): - """The session endpoint builds temporary servers whose server_id has no DB row; with - persist_discovered_endpoints=False neither the oauth2 nor the OBO write-back may fire.""" - manager = MCPServerManager() - - async def fake_discovery(server_url: str, *, allow_origin_fallback: bool = True, warn_when_no_metadata: bool = False): - return MCPOAuthMetadata( - scopes=["s1"], - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/token", + discovered_issuer="https://idp.example.com", ) manager._descovery_metadata = fake_discovery # type: ignore[attr-defined] update_mcp_server_mock = AsyncMock() - obo_update_mock = AsyncMock() repo_instance = MagicMock() - repo_instance.table.update = obo_update_mock + repo_instance.table.update = AsyncMock() with ( - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=update_mcp_server_mock, - ), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), patch( "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", return_value=repo_instance, ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): - oauth2_record = LiteLLM_MCPServerTable( - server_id="temp-oauth-1", - server_name="temp_oauth", - url="https://example.com/mcp", + for auth_type, flow in ((MCPAuth.oauth2, "authorization_code"), (MCPAuth.oauth2_token_exchange, None)): + record = LiteLLM_MCPServerTable( + server_id=f"no-write-{auth_type}", + server_name=f"no_write_{auth_type}", + url="https://example.com/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + oauth2_flow=flow, + credentials={"client_id": "cid", "client_secret": "csec", "audience": "aud"}, + ) + built = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + assert built.token_url == "https://idp.example.com/token" + + update_mcp_server_mock.assert_not_awaited() + repo_instance.table.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_declared_endpoints_survive_a_failed_discovery(self): + """The reporter's configuration: explicit authorization_url/token_url/registration_url, + issuer left empty. With the gateway never stamping the issuer column, the server never turns + anchored, so the declared endpoints resolve on every build, including one whose discovery + fails entirely; /authorize keeps redirecting instead of serving the 400.""" + manager = MCPServerManager() + record = LiteLLM_MCPServerTable( + server_id="declared-1", + alias="declared", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)): + built = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + + assert built.issuer_is_anchored is False + assert built.authorization_url == "https://idp.example.com/authorize" + assert built.token_url == "https://idp.example.com/token" + assert built.registration_url == "https://idp.example.com/register" + + def test_flow_endpoints_missing_arms(self): + """The reload fast-path exemption's completeness rule. Interactive needs authorize+token, + client_credentials and OBO need token only, an OBO server with a configured exchange + endpoint never discovers and must not be sent into a rebuild loop, and non-OAuth auth types + are never unresolved.""" + assert _flow_endpoints_missing(MCPAuth.oauth2, "authorization_code", "https://idp/auth", None) is True + assert _flow_endpoints_missing(MCPAuth.oauth2, "authorization_code", None, "https://idp/token") is True + assert ( + _flow_endpoints_missing(MCPAuth.oauth2, "authorization_code", "https://idp/auth", "https://idp/token") + is False + ) + assert _flow_endpoints_missing(MCPAuth.oauth2, "client_credentials", None, "https://idp/token") is False + assert _flow_endpoints_missing(MCPAuth.oauth2, "client_credentials", None, None) is True + assert _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, None) is True + assert _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, "https://idp/token") is False + assert ( + _flow_endpoints_missing(MCPAuth.oauth2_token_exchange, None, None, None, "https://idp/exchange") is False + ) + assert _flow_endpoints_missing(MCPAuth.api_key, None, None, None) is False + + def test_unresolved_check_uses_the_flow_judge_not_the_raw_column(self): + """A legacy row the startup backfill deliberately left unstamped (token_url plus client + credentials, no authorization_url: the ambiguous M2M shape) serves client_credentials at + request time via effective_oauth2_flow. The reload check must reach the same verdict, or the + row is classified as interactive-missing-endpoints and re-runs discovery on every reload + forever. A null-flow row without the M2M shape stays interactive and genuinely unresolved.""" + m2m_shaped = MCPServer( + server_id="null-flow-m2m", + name="null_flow_m2m", + server_name="null_flow_m2m", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow=None, + token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + assert _oauth_endpoints_unresolved(m2m_shaped) is False + + interactive_unresolved = m2m_shaped.model_copy(update={"client_id": None, "client_secret": None}) + assert _oauth_endpoints_unresolved(interactive_unresolved) is True + + def test_dcr_bridge_relay_arm_needs_its_registration_endpoint(self): + """A dcr_bridge server with no admin-configured client can only register callers through the + upstream registration endpoint, so a partial discovery that resolved authorize and token but + not registration_endpoint leaves it silently degraded to the short-circuit arm. That counts as + unresolved so it keeps retrying. A bridge with a configured client_id uses the short-circuit + arm by design and is unaffected.""" + relay_arm = MCPServer( + server_id="bridge-partial", + name="bridge_partial", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + # dcr_bridge is only valid on the client-forwarded modes (see MCPServer.is_dcr_bridge) + auth_type=MCPAuth.oauth_delegate, + dcr_bridge=True, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url=None, + ) + assert _oauth_endpoints_unresolved(relay_arm) is True + assert _oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False + assert _oauth_endpoints_unresolved(relay_arm.model_copy(update={"client_id": "admin-client"})) is False + + def test_entra_obo_without_scopes_is_unresolved(self): + """entra_obo token exchange fails closed without a scope, and scopes can come from resource + discovery, so an entra_obo server that resolved its token endpoint but no scopes is still + unresolved for its flow. The default rfc8693 profile has no such requirement.""" + entra = MCPServer( + server_id="entra-noscope", + name="entra_noscope", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_profile="entra_obo", + token_url="https://idp.example.com/token", + scopes=None, + ) + assert _oauth_endpoints_unresolved(entra) is True + assert _oauth_endpoints_unresolved(entra.model_copy(update={"scopes": ["api://app/.default"]})) is False + assert _oauth_endpoints_unresolved(entra.model_copy(update={"token_exchange_profile": "rfc8693"})) is False + + def test_oauth_discovery_retry_backs_off_per_server(self): + """Without a cooldown the fast-path exemption re-runs the full discovery chain, and re-emits + the unresolved warning, on every reload forever for a server that can never resolve. Delay + doubles per consecutive failure up to the cap, a success clears the state so the next failure + starts from the base delay again, and the cooldown is per server.""" + manager = MCPServerManager() + + def unresolved(server_id): + return MCPServer( + server_id=server_id, + name=server_id, + url="https://up.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", - credentials={"client_id": "cid", "client_secret": "csec"}, - ) - obo_record = LiteLLM_MCPServerTable( - server_id="temp-obo-1", - server_name="temp_obo", - url="https://example.com/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2_token_exchange, - credentials={"client_id": "cid", "client_secret": "csec"}, - ) - built_oauth2 = await manager.build_mcp_server_from_table( - oauth2_record, credentials_are_encrypted=False, persist_discovered_endpoints=False - ) - await manager.build_mcp_server_from_table( - obo_record, credentials_are_encrypted=False, persist_discovered_endpoints=False ) - assert built_oauth2.authorization_url == "https://idp.example.com/authorize" - update_mcp_server_mock.assert_not_awaited() - obo_update_mock.assert_not_awaited() + assert manager._oauth_discovery_retry_due("a") is True + + manager._record_oauth_discovery_outcome(unresolved("a")) + assert manager._oauth_discovery_retry_due("a") is False + assert manager._oauth_discovery_retry_due("b") is True, "cooldown must be per server" + + failures_before, _ = manager._oauth_discovery_retry_state["a"] + manager._record_oauth_discovery_outcome(unresolved("a")) + failures_after, _ = manager._oauth_discovery_retry_state["a"] + assert failures_after == failures_before + 1 + + # An elapsed cooldown lets the retry through, and the delay grows with the failure count + manager._oauth_discovery_retry_state["a"] = (1, time.monotonic() - 31.0) + assert manager._oauth_discovery_retry_due("a") is True + manager._oauth_discovery_retry_state["a"] = (5, time.monotonic() - 31.0) + assert manager._oauth_discovery_retry_due("a") is False + + resolved = unresolved("a").model_copy( + update={ + "authorization_url": "https://idp.example.com/authorize", + "token_url": "https://idp.example.com/token", + } + ) + manager._record_oauth_discovery_outcome(resolved) + assert "a" not in manager._oauth_discovery_retry_state + assert manager._oauth_discovery_retry_due("a") is True + + @pytest.mark.asyncio + async def test_reload_fast_path_retries_unresolved_oauth_servers(self): + """A server whose discovery failed must not be pinned broken by the updated_at fast path: + the next reload rebuilds it, retrying discovery on the normal cadence instead of waiting for + an unrelated config write. A resolved server with an unchanged row still takes the fast path, + so the exemption costs nothing in the steady state.""" + manager = MCPServerManager() + stamp = datetime.now() + row = LiteLLM_MCPServerTable( + server_id="retry-1", + server_name="retry_server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + created_at=stamp, + updated_at=stamp, + ) + + def entry(authorization_url, token_url): + return MCPServer( + server_id="retry-1", + name="retry_server", + server_name="retry_server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url=authorization_url, + token_url=token_url, + updated_at=stamp, + ) + + raw_row = MagicMock() + raw_row.model_dump.return_value = row.model_dump() + repo_instance = MagicMock() + repo_instance.table.find_many = AsyncMock(return_value=[raw_row]) + + async def run_reload(previous_entry): + manager.registry = {"retry-1": previous_entry} + build_mock = AsyncMock(return_value=previous_entry) + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repo_instance, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch.object(manager, "build_mcp_server_from_table", new=build_mock), + ): + await manager.reload_servers_from_database() + return build_mock + + unresolved_build = await run_reload(entry(None, None)) + unresolved_build.assert_awaited_once() + + resolved_build = await run_reload(entry("https://idp.example.com/authorize", "https://idp.example.com/token")) + resolved_build.assert_not_awaited() + + @pytest.mark.asyncio + async def test_anchored_issuer_discarding_stored_endpoints_warns(self, caplog): + """An anchored server ignoring stored endpoint columns must say so: that state is exactly + what a row stamped by an earlier release looks like after upgrade, and the warning names the + remedy (clear the Issuer field) instead of leaving the 400 undiagnosable.""" + manager = MCPServerManager() + record = LiteLLM_MCPServerTable( + server_id="stamped-1", + alias="stamped_row", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + issuer="https://idp.example.com", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + with ( + patch.object(manager, "_fetch_issuer_anchored_oauth_metadata", new=AsyncMock(return_value=None)), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + built = await manager.build_mcp_server_from_table(record, credentials_are_encrypted=False) + + assert built.issuer_is_anchored is True + assert built.authorization_url is None + assert "stamped_row" in caplog.text + assert "authorization_url, token_url" in caplog.text + assert "clear the Issuer" in caplog.text @pytest.mark.asyncio async def test_update_server_carries_forward_last_known_good_oauth_endpoints(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py new file mode 100644 index 00000000000..b6c946b95fa --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py @@ -0,0 +1,129 @@ +"""Tests for the one-time heal of issuer values a released version's discovery write-back stamped.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._experimental.mcp_server.oauth_issuer_stamp_backfill import ( + backfill_discovery_stamped_issuers, +) + + +def _row(**overrides): + fields = { + "server_id": "srv-1", + "alias": "srv_one", + "server_name": "srv_one", + "auth_type": "oauth2", + "issuer": "https://idp.example.com", + "authorization_url": "https://idp.example.com/authorize", + "token_url": "https://idp.example.com/token", + "registration_url": None, + "updated_by": "mcp_oauth_discovery", + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _prisma(rows): + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=rows) + prisma_client.db.litellm_mcpservertable.update = AsyncMock() + return prisma_client + + +@pytest.mark.asyncio +async def test_clears_the_stamp_and_records_its_own_actor(): + """The GH #34985 row: discovery wrote the issuer, so the server reads as issuer-anchored and its + configured endpoints are ignored. Clearing the stamp makes them apply again. The heal records its + own actor, which is also what makes it idempotent: the row no longer matches the discovery-actor + filter, so it is never reconsidered on a later boot.""" + prisma_client = _prisma([_row()]) + + assert await backfill_discovery_stamped_issuers(prisma_client) == 1 + + call = prisma_client.db.litellm_mcpservertable.update.call_args + assert call.kwargs["where"] == {"server_id": "srv-1"} + assert call.kwargs["data"]["issuer"] is None + assert call.kwargs["data"]["updated_by"] == "mcp_oauth_issuer_stamp_backfill" + + where = prisma_client.db.litellm_mcpservertable.find_many.call_args.kwargs["where"] + assert where["updated_by"] == "mcp_oauth_discovery" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, reason", + [ + ({"updated_by": "some-admin@example.com"}, "an admin was the last writer, so the pin is theirs"), + ({"issuer": None}, "nothing to heal"), + ({"issuer": " "}, "blank issuer is not a pin"), + ( + {"authorization_url": None, "token_url": None, "registration_url": None}, + "issuer set with no configured endpoints is the canonical shape of a deliberate pin, and " + "there is nothing configured for anchoring to discard anyway", + ), + ( + {"authorization_url": "https://other-idp.example.com/authorize", "token_url": None}, + "endpoints addressing a different authority than the issuer are an intent a clear would " + "discard, so the row is warned about rather than healed", + ), + ( + {"issuer": "https://pinned.example.com"}, + "same shape from the other side: a pinned issuer whose origin differs from the configured " + "endpoints cannot have been derived from them by discovery", + ), + ], +) +async def test_leaves_rows_alone_that_do_not_carry_the_defect_signature(overrides, reason): + """updated_by records only the most recent writer and no audit trail says which field it touched, + so the heal is deliberately narrow: it fires only on the full signature of the defect. Every + exclusion here protects a row whose issuer may be a deliberate admin pin.""" + prisma_client = _prisma([_row(**overrides)]) + + assert await backfill_discovery_stamped_issuers(prisma_client) == 0, reason + prisma_client.db.litellm_mcpservertable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_heals_across_url_forms_that_denote_the_same_origin(): + """Origin comparison runs through the shared canonicalizer, so a default port or host casing + difference between the stamped issuer and the endpoints an admin typed does not make a #34985 row + look like a deliberate pin at a different authority.""" + prisma_client = _prisma( + [ + _row( + issuer="https://IDP.example.com:443", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + ] + ) + + assert await backfill_discovery_stamped_issuers(prisma_client) == 1 + + +@pytest.mark.asyncio +async def test_query_is_scoped_to_auth_types_where_an_issuer_anchors(): + """Only the discovery auth types read an issuer as a trust anchor; clearing it elsewhere would be + an unrelated mutation.""" + prisma_client = _prisma([]) + + await backfill_discovery_stamped_issuers(prisma_client) + + where = prisma_client.db.litellm_mcpservertable.find_many.call_args.kwargs["where"] + assert set(where["auth_type"]["in"]) == {"oauth2", "true_passthrough", "oauth_delegate"} + + +@pytest.mark.asyncio +async def test_a_failed_row_does_not_abort_the_rest(): + """Per-row best effort: one write failure must not leave later rows unhealed, and the next boot + retries the failed one since its updated_by is unchanged.""" + prisma_client = _prisma([_row(server_id="bad"), _row(server_id="good")]) + prisma_client.db.litellm_mcpservertable.update = AsyncMock( + side_effect=[Exception("write failed"), MagicMock()] + ) + + assert await backfill_discovery_stamped_issuers(prisma_client) == 1 + assert prisma_client.db.litellm_mcpservertable.update.await_count == 2 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx index 8dffc80a70e..5650bd1d7e4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/OAuthFormFields.tsx @@ -190,7 +190,7 @@ const OAuthFormFields: React.FC = ({ label={ } name="issuer"